From f94583b3d09a00e32eda2e54e4531ec835707942 Mon Sep 17 00:00:00 2001 From: tastelikefeet <58414341+tastelikefeet@users.noreply.github.com> Date: Sat, 4 Jul 2026 17:09:21 +0800 Subject: [PATCH 01/60] wip (#243) --- src/twinkle_agentic/chunker/__init__.py | 4 - src/twinkle_agentic/chunker/base.py | 14 - src/twinkle_agentic/chunker/native.py | 254 ----------- src/twinkle_agentic/classifier/__init__.py | 0 src/twinkle_agentic/classifier/base.py | 11 + src/twinkle_agentic/condenser/__init__.py | 5 +- src/twinkle_agentic/condenser/base.py | 204 ++++++++- src/twinkle_agentic/condenser/facts.py | 83 ++++ src/twinkle_agentic/condenser/keyword.py | 486 -------------------- src/twinkle_agentic/condenser/model.py | 508 --------------------- src/twinkle_agentic/train/__init__.py | 0 src/twinkle_agentic/train/base.py | 6 + src/twinkle_agentic/train/cron.py | 5 + src/twinkle_agentic/utils/llm_backup.py | 314 +++++++++++++ 14 files changed, 617 insertions(+), 1277 deletions(-) delete mode 100644 src/twinkle_agentic/chunker/__init__.py delete mode 100644 src/twinkle_agentic/chunker/base.py delete mode 100644 src/twinkle_agentic/chunker/native.py create mode 100644 src/twinkle_agentic/classifier/__init__.py create mode 100644 src/twinkle_agentic/classifier/base.py create mode 100644 src/twinkle_agentic/condenser/facts.py delete mode 100644 src/twinkle_agentic/condenser/keyword.py delete mode 100644 src/twinkle_agentic/condenser/model.py create mode 100644 src/twinkle_agentic/train/__init__.py create mode 100644 src/twinkle_agentic/train/base.py create mode 100644 src/twinkle_agentic/train/cron.py create mode 100644 src/twinkle_agentic/utils/llm_backup.py diff --git a/src/twinkle_agentic/chunker/__init__.py b/src/twinkle_agentic/chunker/__init__.py deleted file mode 100644 index f826a6452..000000000 --- a/src/twinkle_agentic/chunker/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .base import Chunker -from .native import NativeChunker - -__all__ = ['Chunker', 'NativeChunker'] diff --git a/src/twinkle_agentic/chunker/base.py b/src/twinkle_agentic/chunker/base.py deleted file mode 100644 index 22beb8b88..000000000 --- a/src/twinkle_agentic/chunker/base.py +++ /dev/null @@ -1,14 +0,0 @@ -from abc import ABC, abstractmethod - -from twinkle.data_format import Trajectory -from twinkle_agentic.data_format import Chunks - - -class Chunker(ABC): - """ - TODO: Experimental feature, wait for testing - """ - - @abstractmethod - def __call__(self, trajectory: Trajectory) -> Chunks: - raise NotImplementedError diff --git a/src/twinkle_agentic/chunker/native.py b/src/twinkle_agentic/chunker/native.py deleted file mode 100644 index f5879f3c0..000000000 --- a/src/twinkle_agentic/chunker/native.py +++ /dev/null @@ -1,254 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -import re -from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence - -from twinkle.data_format import Trajectory -from twinkle_agentic.data_format import Chunk, Chunks -from .base import Chunker - -# Recursive separator list, coarsest → finest. The empty string at the -# end forces a hard character cut when nothing finer fits. -_DEFAULT_SEPARATORS: tuple = ( - '\n\n', - '\n', - '。', - '．', - '.', - '！', - '!', - '？', - '?', - '；', - ';', - '，', - ',', - ' ', - '', -) - -_MULTIMODAL_TYPES = ('image', 'video', 'audio') - -_SplitFn = Optional[Callable[[str], List[str]]] - - -class NativeChunker(Chunker): - """Character-level recursive chunker for trajectories. - TODO: Experimental feature, wait for testing - Args: - chunk_size: Soft upper bound (in characters) for every emitted - text chunk. Must be positive. - separators: Ordered separator list. The chunker tries each - separator in turn; any piece still larger than - ``chunk_size`` is re-split with the next one. A terminal - ``''`` (hard character cut) is appended automatically if - missing so the algorithm is guaranteed to terminate. - passage_boundary_re: Optional regex (compiled with - ``re.MULTILINE``) whose matches act as **hard, non-mergeable** - passage boundaries on the first user message. The regex - match is preserved at the start of the next piece (so - ``''.join(pieces) == text``). Pieces that are already - ``<= chunk_size`` are emitted as-is and are **never merged** - across boundaries; only pieces that still exceed - ``chunk_size`` fall back to the normal recursive split + merge. - This is how you keep e.g. HotpotQA passages atomic per - ````. - """ - - def __init__( - self, - chunk_size: int = 1024, - separators: Sequence[str] | None = None, - passage_boundary_re: str | None = None, - ): - if chunk_size <= 0: - raise ValueError(f'chunk_size must be positive, got {chunk_size}') - self.chunk_size = chunk_size - seps = tuple(separators) if separators is not None else _DEFAULT_SEPARATORS - if '' not in seps: - seps += ('', ) - self.separators = seps - self.passage_boundary_re: re.Pattern | None = ( - re.compile(passage_boundary_re, re.MULTILINE) if passage_boundary_re else None) - - # ------------------------------------------------------------------ - # public entry - # ------------------------------------------------------------------ - def __call__(self, trajectory: Trajectory) -> Chunks: - chunks: list[Chunk] = [] - first_user_done = False - # ``round`` is 1-indexed at the first user message. Any messages - # emitted before that (e.g., leading ``system``) carry round 0. - round_idx = 0 - for msg in trajectory.get('messages') or []: - is_user = msg.get('role') == 'user' - if is_user: - round_idx += 1 - split = (self._split_text if is_user and not first_user_done else None) - if is_user: - first_user_done = True - for chunk in self._parts(msg, split): - chunk['round'] = round_idx - chunks.append(chunk) - return Chunks(chunks=chunks) - - # ------------------------------------------------------------------ - # message → chunks decomposition - # ------------------------------------------------------------------ - def _parts(self, message: dict[str, Any], split: _SplitFn) -> Iterator[Chunk]: - role = message.get('role') or 'user' - tcid = message.get('tool_call_id') - - rc = message.get('reasoning_content') - if rc: - yield _text_chunk(role, rc, kind='reasoning_content', tool_call_id=tcid) - - content = message.get('content') - if isinstance(content, str): - yield from self._emit_text(role, content, split, tcid) - elif isinstance(content, list): - for part in content: - if not isinstance(part, dict): - continue - ptype = part.get('type') - if ptype == 'text': - yield from self._emit_text(role, part.get('text') or '', split, tcid) - elif ptype in _MULTIMODAL_TYPES: - # Keep raw part so Chunks.to_trajectory can rebuild - # the original OpenAI-style entry verbatim. - yield { # type: ignore[misc] - 'type': ptype, 'content': part.get(ptype), - 'raw': dict(part), 'role': role, - } - - for tc in message.get('tool_calls') or []: - yield _text_chunk(role, '', kind='tool_call', tool_call=tc, tool_call_id=tcid) - - def _emit_text(self, role: str, text: str, split: _SplitFn, tool_call_id: str | None) -> Iterator[Chunk]: - if not text: - return - pieces = split(text) if split is not None else [text] - for piece in pieces: - if piece: - yield _text_chunk(role, piece, tool_call_id=tool_call_id) - - # ------------------------------------------------------------------ - # recursive text splitter - # ------------------------------------------------------------------ - def _split_text(self, text: str) -> list[str]: - if not text: - return [] - if self.passage_boundary_re is None: - if len(text) <= self.chunk_size: - return [text] - return self._merge(self._recursive_split(text, list(self.separators))) - # Force-split first; each forced piece is kept intact when it is - # already short enough, and is recursively re-split (but NOT - # merged with sibling passages) when it exceeds ``chunk_size``. - out: list[str] = [] - for piece in self._force_split(text): - if not piece or not piece.strip(): - continue - if len(piece) <= self.chunk_size: - out.append(piece) - else: - out.extend(self._merge(self._recursive_split(piece, list(self.separators)))) - return out - - def _force_split(self, text: str) -> list[str]: - """Split ``text`` at every ``passage_boundary_re`` match; the - match itself sticks to the start of the **next** piece, so - ``''.join(_force_split(text)) == text``. - """ - assert self.passage_boundary_re is not None - matches = list(self.passage_boundary_re.finditer(text)) - if not matches: - return [text] - out: list[str] = [] - prev = 0 - for m in matches: - start = m.start() - if start > prev: - out.append(text[prev:start]) - prev = start - if prev < len(text): - out.append(text[prev:]) - return out - - def _recursive_split(self, text: str, separators: list[str]) -> list[str]: - if len(text) <= self.chunk_size: - return [text] if text else [] - # Terminal: no more separators, or next one is the hard-cut sentinel. - if not separators or separators[0] == '': - return _hard_cut(text, self.chunk_size) - - sep, *rest = separators - out: list[str] = [] - for piece in _split_keep(text, sep): - if not piece: - continue - if len(piece) <= self.chunk_size: - out.append(piece) - else: - out.extend(self._recursive_split(piece, rest)) - return out - - def _merge(self, pieces: list[str]) -> list[str]: - """Greedy concatenation: small fragments fuse up to ``chunk_size`` - without exceeding it. Relative order is preserved. - """ - merged: list[str] = [] - buf = '' - for p in pieces: - if not p: - continue - if buf and len(buf) + len(p) > self.chunk_size: - merged.append(buf) - buf = '' - buf += p - if buf: - merged.append(buf) - return merged - - -# ---------------------------------------------------------------------- -# helpers -# ---------------------------------------------------------------------- -def _split_keep(text: str, sep: str) -> list[str]: - """``str.split(sep)`` but the separator stays glued to the end of - each left-hand piece, so ``''.join(result) == text``. - """ - if not sep or sep not in text: - return [text] if text else [] - out: list[str] = [] - start, n = 0, len(sep) - while (i := text.find(sep, start)) != -1: - out.append(text[start:i + n]) - start = i + n - if start < len(text): - out.append(text[start:]) - return out - - -def _hard_cut(text: str, size: int) -> list[str]: - return [text[i:i + size] for i in range(0, len(text), size)] if text else [] - - -def _text_chunk( - role: str, - content: str, - *, - kind: str | None = None, - tool_call: Any = None, - tool_call_id: str | None = None, -) -> Chunk: - raw: dict[str, Any] = {} - if kind is not None: - raw['kind'] = kind - if tool_call is not None: - raw['tool_call'] = tool_call - if tool_call_id is not None: - raw['tool_call_id'] = tool_call_id - chunk: Chunk = {'type': 'text', 'content': content, 'role': role} # type: ignore[assignment] - if raw: - chunk['raw'] = raw - return chunk diff --git a/src/twinkle_agentic/classifier/__init__.py b/src/twinkle_agentic/classifier/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/twinkle_agentic/classifier/base.py b/src/twinkle_agentic/classifier/base.py new file mode 100644 index 000000000..d79f6c2e2 --- /dev/null +++ b/src/twinkle_agentic/classifier/base.py @@ -0,0 +1,11 @@ +from abc import ABC, abstractmethod + + +class Classifier(ABC): + + def __init__(self, model_path: str, **kwargs): + self.model_path = model_path + + @abstractmethod + def classify(self, text: str) -> str: + pass \ No newline at end of file diff --git a/src/twinkle_agentic/condenser/__init__.py b/src/twinkle_agentic/condenser/__init__.py index e78545002..a48fd73c1 100644 --- a/src/twinkle_agentic/condenser/__init__.py +++ b/src/twinkle_agentic/condenser/__init__.py @@ -1,5 +1,4 @@ from .base import Condenser -from .keyword import KeywordCondenser -from .model import ModelCondenser +from .facts import FactsCondenser -__all__ = ['Condenser', 'KeywordCondenser', 'ModelCondenser'] +__all__ = ['Condenser', 'FactsCondenser'] diff --git a/src/twinkle_agentic/condenser/base.py b/src/twinkle_agentic/condenser/base.py index 5e42dab17..8fd2acd61 100644 --- a/src/twinkle_agentic/condenser/base.py +++ b/src/twinkle_agentic/condenser/base.py @@ -1,13 +1,201 @@ -from abc import ABC, abstractmethod +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations -from twinkle_agentic.data_format import Chunks +import math +import re +from typing import TYPE_CHECKING, Any, Sequence +from twinkle_agentic.utils.llm_backup import llm_backup -class Condenser(ABC): - """ - TODO: Experimental feature, wait for testing +if TYPE_CHECKING: + from twinkle.data_format import SamplingParams, Trajectory # noqa: F401 + from twinkle.sampler.base import Sampler # noqa: F401 + + +DEFAULT_USER_PROMPT_TEMPLATE = """\ +Compress the following text as much as possible while preserving all key information. + +## Target length +HARD CEILING: {budget} chars. If core facts fit in far fewer chars, output fewer. + +## Text +{text}""" + + +class Condenser: + """Base condenser with progressive distillation via llm_backup. + + Subclasses customize compression behavior by providing their own + ``system_prompt``, ``user_prompt_template``, and ``lora_path``. + The shared ``_sample`` method (decorated with ``@llm_backup``) handles + the student-teacher routing transparently. + + Teacher is a global OpenAI-compatible API configured via env vars: + - LLM_BACKUP_MODEL: teacher model name + - LLM_BACKUP_API_KEY: API key + - LLM_BACKUP_BASE_URL: API endpoint + + Args: + sampler: Student model sampler (local inference, shared across types). + compression_ratio: Target compression factor (> 1). + model_path: Model identifier. + sampling_params: Default sampling params. + system_prompt: System prompt for this condenser type. + user_prompt_template: User prompt template. Must contain + ``{budget}`` and ``{text}``. May contain ``{query}``. + min_budget_chars: Floor for the character budget in the prompt. + template: Optional :class:`Template` for special token stripping. + lora_path: LoRA adapter path specific to this condenser type. + Each subclass can use a different LoRA for its task. """ - @abstractmethod - def __call__(self, chunks: Chunks, **kwargs) -> Chunks: - raise NotImplementedError + def __init__( + self, + sampler: Sampler, + compression_ratio: float = 2.0, + *, + model_path: str = '', + sampling_params: SamplingParams | None = None, + system_prompt: str = 'You are a text compression assistant.', + user_prompt_template: str | None = None, + min_budget_chars: int = 250, + template: Any | None = None, + lora_path: str | None = None, + ): + if sampler is None: + raise ValueError('sampler is required') + if compression_ratio <= 1.0: + raise ValueError(f'compression_ratio must be > 1, got {compression_ratio}') + if min_budget_chars < 1: + raise ValueError(f'min_budget_chars must be >= 1, got {min_budget_chars}') + + tpl = user_prompt_template or DEFAULT_USER_PROMPT_TEMPLATE + if '{budget}' not in tpl or '{text}' not in tpl: + raise ValueError('user_prompt_template must contain both {budget} and {text}') + + self.model_path = model_path + self.sampler = sampler + self.compression_ratio = float(compression_ratio) + self.sampling_params = sampling_params + self.system_prompt = system_prompt + self.user_prompt_template = tpl + self.min_budget_chars = int(min_budget_chars) + self.template = template + self.lora_path = lora_path if lora_path else None + self._special_tokens_cache: tuple[str, ...] | None = None + + # ------------------------------------------------------------------ + # public entry point (pre/post processing, NOT decorated) + # ------------------------------------------------------------------ + def __call__(self, text: str, system: str = None, query: str = None, + sampling_params: Any = None) -> str: + system = system or self.system_prompt + budget = max(self.min_budget_chars, math.ceil(len(text) / self.compression_ratio)) + if budget >= len(text): + return text + trajectory = self._make_trajectory(system, self.user_prompt_template, text, budget, query) + sp = sampling_params or self.sampling_params or self._default_sampling_params(budget) + + raw = self._sample(trajectory=trajectory, sampling_params=sp, query=query) + + result = self._postprocess(raw, text, self._get_special_tokens()) + return result if result is not None else text + + # ------------------------------------------------------------------ + # student sampling (decorated with llm_backup) + # ------------------------------------------------------------------ + @llm_backup(key_params=["query"]) + def _sample(self, trajectory, sampling_params, query: str = None) -> str: + """Student model: trajectory + sampling_params -> raw text.""" + sample_kwargs: dict[str, Any] = {'sampling_params': sampling_params} + if self.lora_path is None: + sample_kwargs['use_base_model'] = True + else: + sample_kwargs['adapter_path'] = self.lora_path + responses = self.sampler.sample([trajectory], **sample_kwargs) + return self._decoded(list(responses)[0]) if responses else '' + + # ------------------------------------------------------------------ + # internals + # ------------------------------------------------------------------ + def _get_special_tokens(self) -> tuple[str, ...]: + if self._special_tokens_cache is not None: + return self._special_tokens_cache + tpl = self.template or getattr(self.sampler, 'template', None) + tokenizer = getattr(tpl, 'tokenizer', None) if tpl is not None else None + tokens: list[str] = [] + if tokenizer is not None: + extras = getattr(tokenizer, 'all_special_tokens', None) or [] + if extras: + tokens.extend(t for t in extras if isinstance(t, str) and t and not t.isspace()) + else: + for attr in ('eos_token', 'pad_token', 'bos_token'): + t = getattr(tokenizer, attr, None) + if isinstance(t, str) and t: + tokens.append(t) + self._special_tokens_cache = tuple(dict.fromkeys(tokens)) + return self._special_tokens_cache + + + # ------------------------------------------------------------------ + # static helpers + # ------------------------------------------------------------------ + _CODE_FENCE_RE = re.compile(r'^```[a-zA-Z]*\s*\n(.*?)\n```\s*$', re.DOTALL) + + @staticmethod + def _make_trajectory(system: str, user_template: str, text: str, + budget: int, query: str | None = None) -> dict: + """Build a trajectory dict for sampler / API.""" + user = user_template.replace('{budget}', str(budget)) + user = user.replace('{text}', text) + if '{query}' in user: + q_text = ( + query.strip() if isinstance(query, str) and query and query.strip() else + '(no explicit query; compress by general salience)') + user = user.replace('{query}', q_text) + return { + 'messages': [ + {'role': 'system', 'content': system}, + {'role': 'user', 'content': user}, + ], + } + + @staticmethod + def _default_sampling_params(budget: int): + from twinkle.data_format.sampling import SamplingParams + max_new = max(512, budget * 3 + 128) + return SamplingParams(temperature=0.0, max_tokens=max_new) + + @staticmethod + def _postprocess(raw: str, original: str, special_tokens: tuple[str, ...]) -> str | None: + text = Condenser._strip_special_tokens( + Condenser._strip_code_fences(raw), special_tokens).strip() + if not text or not Condenser._has_alnum(text): + return None + if len(text) >= len(original): + return None + return text + + @staticmethod + def _decoded(response: Any) -> str: + seqs = getattr(response, 'sequences', None) or [] + if not seqs: + return '' + return getattr(seqs[0], 'decoded', None) or '' + + @staticmethod + def _strip_code_fences(text: str) -> str: + stripped = text.strip() + m = Condenser._CODE_FENCE_RE.match(stripped) + return m.group(1) if m else text + + @staticmethod + def _strip_special_tokens(text: str, tokens: Sequence[str]) -> str: + for tok in tokens: + if tok and tok in text: + text = text.replace(tok, '') + return text + + @staticmethod + def _has_alnum(text: str) -> bool: + return any(ch.isalnum() for ch in text) diff --git a/src/twinkle_agentic/condenser/facts.py b/src/twinkle_agentic/condenser/facts.py new file mode 100644 index 000000000..6a8313b9b --- /dev/null +++ b/src/twinkle_agentic/condenser/facts.py @@ -0,0 +1,83 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from twinkle_agentic.condenser.base import Condenser + +if TYPE_CHECKING: + from twinkle.data_format import SamplingParams # noqa: F401 + from twinkle.sampler.base import Sampler # noqa: F401 + +_SECTION_SCHEMA = """You are a text compression assistant. A downstream model will read your compressed output to decide whether the detail it needs is inside this block; if yes, it will fetch and read the original passage. + +Downstream model workflow: +Read your compressed output -> Decide whether needed info is in this block -> If yes -> Fetch original. + +Therefore your compression MUST NOT lose major information from the source. + +Output format: + +```text +## Summary +Overview plus facts STRONGLY RELATED to the Query, stated explicitly. + +## More +A collapsed index; expansion required to see specific information. +``` + +Rules: +1. Telegraphic style — drop function words ("the", "a", "is", "are", "of", ...); colons and commas mean "is" / "has". + * Exception: KEEP role-tagging verb+preposition phrases verbatim ("published by X", "written by X", "directed by X", "starring X", "founded by X", "created by X", "composed by X", "produced by X", "based on X", "adapted from X"). Collapsing these to a bare name loses the relation role (author vs publisher vs director) that the downstream question may hinge on. +2. Summary MUST contain the passage's primary topic + 2–4 concrete core facts drawn from the source (entities, numbers, dates, relations). If a Query is given, order Query-relevant facts first, but STILL include other core facts within the budget. A Query is an ORDERING HINT, NOT a filter. +3. Summary MUST NOT be meta-commentary about the Query. Forbidden patterns: "no X mention", "Query info: absent", "passage covers Y only", "does not contain ...", "no relevant info", or summaries that are only abstract category words like "structure/order/usage" with no facts. If the passage is unrelated to the Query, you still summarize the passage normally. +4. More is an INDEX of category keywords, NOT inline data. Enumerate what CAN be recovered from the source (e.g. "birthplace, death place, age"); do NOT paste dates/numbers/names inline. Make sure all category of useful facts are introduced here. +5. Output language MUST match the source language. +6. Do NOT fabricate. Do NOT omit major information. Any fact not in the source MUST NOT appear in your output. + +Now begin. +""" # noqa + +_SECTION_USER_TEMPLATE = """\ +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. + +## Query (ordering hint only — still summarize the whole passage) +{query} + +## Target length +Compress AS MUCH AS faithfully possible. HARD CEILING: {budget} chars. \ +If core facts fit in far fewer chars, output fewer. \ +Never exceed the ceiling. + +## Passage +{text}""" + + +class FactsCondenser(Condenser): + + def __init__( + self, + sampler: Sampler, + compression_ratio: float = 2.0, + *, + model_path: str = '', + sampling_params: SamplingParams | None = None, + min_budget_chars: int = 250, + template: Any | None = None, + lora_path: str | None = None, + ): + super().__init__( + sampler, + compression_ratio, + model_path=model_path, + sampling_params=sampling_params, + system_prompt=_SECTION_SCHEMA, + user_prompt_template=_SECTION_USER_TEMPLATE, + min_budget_chars=min_budget_chars, + template=template, + lora_path=lora_path, + ) diff --git a/src/twinkle_agentic/condenser/keyword.py b/src/twinkle_agentic/condenser/keyword.py deleted file mode 100644 index e17c3ca7c..000000000 --- a/src/twinkle_agentic/condenser/keyword.py +++ /dev/null @@ -1,486 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -import math -import re -import threading -from typing import Any, Dict, FrozenSet, List, Optional, Sequence, Tuple - -from twinkle_agentic.condenser.base import Condenser -from twinkle_agentic.data_format import Chunk, Chunks - -# --------------------------------------------------------------------------- -# spaCy lazy loader (one model per process, thread-safe) -# --------------------------------------------------------------------------- -_SPACY_MODELS: dict[str, Any] = {} -_SPACY_LOCK = threading.Lock() - - -def _load_spacy(name: str): - nlp = _SPACY_MODELS.get(name) - if nlp is not None: - return nlp - with _SPACY_LOCK: - nlp = _SPACY_MODELS.get(name) - if nlp is not None: - return nlp - try: - import spacy - except ImportError as e: - raise ImportError('KeywordCondenser requires spaCy. Install with: ' - '`pip install spacy && python -m spacy download en_core_web_sm`') from e - try: - nlp = spacy.load(name) - except OSError as e: - raise OSError(f'spaCy model {name!r} not found. Download with: ' - f'`python -m spacy download {name}`') from e - _SPACY_MODELS[name] = nlp - return nlp - - -# --------------------------------------------------------------------------- -# configuration-free constants -# --------------------------------------------------------------------------- -# Entity labels dropped from keyword candidates (low recall value). -_DROP_ENT_LABELS: frozenset[str] = frozenset({'CARDINAL', 'ORDINAL', 'PERCENT', 'QUANTITY'}) - -# Dependency labels that introduce sub-clauses / conjuncts we do NOT want -# to pull into a single noun-phrase span. -_DROP_NP_DEPS: frozenset[str] = frozenset( - {'relcl', 'acl', 'advcl', 'ccomp', 'xcomp', 'conj', 'cc', 'appos', 'parataxis'}) - -# Tokens stripped from NP boundaries. -_LEADING_STRIP_POS: frozenset[str] = frozenset({'DET', 'PUNCT'}) - -# Tuple-slot separator. ``|`` avoids confusion when a slot itself -# contains a comma (e.g. ``"London, England"``). -_SLOT_SEP = ' | ' -_TRIPLE_SEP = '; ' - -_WORD_RE = re.compile(r'\w+', flags=re.UNICODE) - - -# --------------------------------------------------------------------------- -# NP / verb surface helpers -# --------------------------------------------------------------------------- -def _np_text(head) -> str: - """Return the noun-phrase text headed by ``head``. - - Keeps the contiguous span from the leftmost to the rightmost kept - token so internal punctuation (hyphens, apostrophes, slashes) is - preserved verbatim. Drops clausal / conjunct sub-trees and trims - leading determiners / possessive pronouns. - """ - # Collect subtree tokens, cutting off whole clausal children. - collected: list = [] - - def _walk(tok): - if tok is not head and tok.dep_ in _DROP_NP_DEPS: - return - collected.append(tok) - for child in tok.children: - _walk(child) - - _walk(head) - if not collected: - return head.text - collected.sort(key=lambda t: t.i) - - # Strip leading det/punct and possessive pronouns. - while collected and (collected[0].pos_ in _LEADING_STRIP_POS or - (collected[0].pos_ == 'PRON' and collected[0].dep_ == 'poss')): - collected.pop(0) - while collected and collected[-1].pos_ == 'PUNCT': - collected.pop() - if not collected: - return head.text - - start, end = collected[0].i, collected[-1].i + 1 - # If the kept tokens form a contiguous span, use the original text - # (preserves hyphens etc.). Otherwise fall back to text_with_ws. - if end - start == len(collected): - return head.doc[start:end].text.strip() - return ''.join(t.text_with_ws for t in collected).strip() - - -def _verb_surface(verb_tok) -> str: - """Verb text including auxiliaries (``was born``, ``has been released``).""" - aux = [c for c in verb_tok.children if c.dep_ in ('aux', 'auxpass')] - if not aux: - return verb_tok.text - tokens = sorted(aux + [verb_tok], key=lambda t: t.i) - return ' '.join(t.text for t in tokens) - - -def _first_child(token, deps: Sequence[str]): - if token is None: - return None - for c in token.children: - if c.dep_ in deps: - return c - return None - - -def _strip_leading_nc(noun_chunk) -> str: - toks = list(noun_chunk) - while toks and (toks[0].pos_ in _LEADING_STRIP_POS or toks[0].pos_ == 'NUM' or - (toks[0].pos_ == 'PRON' and toks[0].tag_ in ('PRP$', 'WP$'))): - toks.pop(0) - while toks and toks[-1].pos_ == 'PUNCT': - toks.pop() - if not toks: - return '' - start, end = toks[0].i, toks[-1].i + 1 - if end - start == len(toks): - return noun_chunk.doc[start:end].text.strip() - return ''.join(t.text_with_ws for t in toks).strip() - - -def _word_tokens_lower(text: str) -> frozenset[str]: - return frozenset(m.group(0).lower() for m in _WORD_RE.finditer(text)) - - -def _word_boundary_truncate(text: str, limit: int) -> str: - """Truncate ``text`` to ``limit`` chars at the nearest space.""" - if len(text) <= limit: - return text - cut = text[:limit] - sp = cut.rfind(' ') - trimmed = cut[:sp] if sp >= limit // 2 else cut - return trimmed.rstrip() or cut - - -# --------------------------------------------------------------------------- -# extraction (pure functions on spaCy Doc) -# --------------------------------------------------------------------------- -def _extract_opening(doc, max_chars: int) -> str: - """First non-empty sentence, word-boundary-truncated to ``max_chars``.""" - if max_chars <= 0: - return '' - for sent in doc.sents: - text = sent.text.strip() - if text: - return _word_boundary_truncate(text, max_chars) - return '' - - -def _extract_triples(doc, n: int) -> list[tuple[str, ...]]: - """Subject-verb-object (+ optional prep-obj) triples. - - - Skips pronoun subjects (unresolved coreference is noise). - - Preserves verb surface form (``was born`` rather than ``bear``). - - Deduplicates on lemmas. - """ - if n <= 0: - return [] - out: list[tuple[str, ...]] = [] - seen: set = set() - for sent in doc.sents: - for verb in sent: - if verb.pos_ not in ('VERB', 'AUX'): - continue - subj = _first_child(verb, ('nsubj', 'nsubjpass', 'csubj')) - if subj is None or subj.pos_ == 'PRON': - continue - obj = _first_child(verb, ('dobj', 'attr', 'oprd')) - prep = _first_child(verb, ('prep', )) - prep_obj = _first_child(prep, ('pobj', 'pcomp')) if prep is not None else None - - subj_txt = _np_text(subj) - verb_txt = _verb_surface(verb) - - if obj is not None and prep_obj is not None: - triple = (subj_txt, verb_txt, _np_text(obj), f'{prep.text} {_np_text(prep_obj)}') - key = (subj.lemma_.lower(), verb.lemma_.lower(), obj.lemma_.lower(), - f'{prep.text.lower()} {prep_obj.lemma_.lower()}') - elif obj is not None: - triple = (subj_txt, verb_txt, _np_text(obj)) - key = (subj.lemma_.lower(), verb.lemma_.lower(), obj.lemma_.lower()) - elif prep_obj is not None: - triple = (subj_txt, f'{verb_txt} {prep.text}', _np_text(prep_obj)) - key = (subj.lemma_.lower(), f'{verb.lemma_.lower()} {prep.text.lower()}', prep_obj.lemma_.lower()) - else: - continue - if key in seen: - continue - seen.add(key) - out.append(triple) - if len(out) >= n: - return out - return out - - -def _extract_keywords(doc, k: int, excluded_tokens: frozenset[str]) -> list[str]: - """Rank keyword candidates by (entity-weighted) frequency. - - - Drops pure-numeric entities (CARDINAL / ORDINAL / PERCENT / QUANTITY). - - Skips any term whose words are all already in ``excluded_tokens`` - (so we don't repeat what the opening already says). - - Subsumption dedup: drops a shorter form if a longer form - containing it is already kept (``"Nolan"`` dropped when - ``"Christopher Nolan"`` is present). - """ - if k <= 0: - return [] - counts: dict[str, float] = {} - order: dict[str, int] = {} - idx = 0 - - def _add(term: str, weight: float) -> None: - nonlocal idx - t = term.strip() - if len(t) < 2: - return - words = [w.lower() for w in _WORD_RE.findall(t)] - if not words: - return - if all(w in excluded_tokens for w in words): - return - if t not in order: - order[t] = idx - idx += 1 - counts[t] = counts.get(t, 0.0) + weight - - for ent in doc.ents: - if ent.label_ in _DROP_ENT_LABELS: - continue - _add(ent.text, weight=10.0) - for nc in doc.noun_chunks: - _add(_strip_leading_nc(nc), weight=1.0) - for tok in doc: - if tok.pos_ == 'PROPN' and not tok.is_stop: - _add(tok.text, weight=2.0) - - ranked = sorted(counts.keys(), key=lambda t: (-counts[t], order[t])) - - kept: list[str] = [] - kept_word_sets: list[frozenset[str]] = [] - for term in ranked: - words = frozenset(_WORD_RE.findall(term.lower())) - # Subsumed by any already-kept term (identical or proper subset). - if any(words == ws or words < ws for ws in kept_word_sets): - continue - # Also drop earlier-kept strict subsets of the current term. - to_remove = [i for i, ws in enumerate(kept_word_sets) if ws < words] - for i in reversed(to_remove): - kept.pop(i) - kept_word_sets.pop(i) - kept.append(term) - kept_word_sets.append(words) - if len(kept) >= k: - break - return kept - - -# --------------------------------------------------------------------------- -# budget-aware formatting (pure strings) -# --------------------------------------------------------------------------- -def _format_triple(triple: tuple[str, ...]) -> str: - return '(' + _SLOT_SEP.join(triple) + ')' - - -def _compose(opening: str, rel: str, kw: str) -> str: - parts: list[str] = [] - if opening: - parts.append(f'Open: {opening}') - if rel: - parts.append(f'Rel: {rel}') - if kw: - parts.append(f'More: {kw}') - return '\n'.join(parts) - - -def _fit_under_budget( - opening: str, - triples: list[tuple[str, ...]], - keywords: list[str], - budget: int, - *, - fallback_text: str = '', -) -> str: - """Pack as many triples + keywords as possible under ``budget``. - - Strategy: - 1. If opening alone is already too long, word-boundary truncate it. - 2. Greedily append triples one-by-one, keeping a running string. - 3. Greedily append keywords one-by-one on top of whatever fits. - 4. Never exceed ``budget`` — final safety clamp applies. - """ - # ----- opening ----- - if opening and len(f'Open: {opening}') > budget: - max_open = max(0, budget - len('Open: ')) - opening = _word_boundary_truncate(opening, max_open) if max_open else '' - - if not opening and not triples and not keywords: - # Nothing extractable — fall back to raw text, strict-truncated. - base = fallback_text[:budget] if fallback_text else '' - return _word_boundary_truncate(base, budget) if base else base - - current = _compose(opening, '', '') - if len(current) > budget: - return current[:budget] - - # ----- triples ----- - kept_triples: list[tuple[str, ...]] = [] - for t in triples: - trial_rel = _TRIPLE_SEP.join(_format_triple(x) for x in kept_triples + [t]) - trial = _compose(opening, trial_rel, '') - if len(trial) <= budget: - kept_triples.append(t) - else: - break - - rel_str = _TRIPLE_SEP.join(_format_triple(x) for x in kept_triples) - - # ----- keywords ----- - kept_kws: list[str] = [] - for k in keywords: - trial_kw = ', '.join(kept_kws + [k]) - trial = _compose(opening, rel_str, trial_kw) - if len(trial) <= budget: - kept_kws.append(k) - else: - break - - kw_str = ', '.join(kept_kws) - result = _compose(opening, rel_str, kw_str) - if not result: - # Budget too tight for any extracted slot — fall back to raw - # text truncated at a word boundary. - base = fallback_text[:budget] if fallback_text else '' - return _word_boundary_truncate(base, budget) if base else base - # Belt-and-braces: budget is strict. - return result if len(result) <= budget else result[:budget] - - -# --------------------------------------------------------------------------- -# KeywordCondenser -# --------------------------------------------------------------------------- -class KeywordCondenser(Condenser): - """Extractive, spaCy-driven passage condenser. - TODO: Experimental feature, wait for testing - - Args: - num_relations: Max number of - ``(subject, verb, object[, prep-obj])`` tuples per chunk. - Set to ``0`` to disable the ``Rel:`` slot. - max_first_sentence_chars: Hard cap for the opening slot, applied - before the global compression budget. - num_keywords: Max keyword items per chunk. ``0`` disables ``More:``. - compression_ratio: Target compression factor. Must be ``> 1``. - ``len(output) <= ceil(len(input) / compression_ratio)`` is - strictly enforced for every chunk that passes ``min_chars``. - spacy_model: spaCy pipeline name (default ``en_core_web_sm``). - min_chars: Pre-filter. Chunks shorter than this are passed - through **unchanged**; the ratio contract does not apply to - them. Set to ``0`` to always compress. - skip_roles: Roles whose chunks are never compressed. - rounds: Optional set/list of conversation-turn numbers to - compress. ``None`` (default) = no round-based filtering; - when provided, chunks whose ``round`` is not in this set - are passed through unchanged. Chunks that lack a ``round`` - field are also skipped when this filter is active. - - Every produced chunk is marked with ``raw.condensed=True`` so - :meth:`Chunks.to_trajectory` wraps it in ``...``. - - Example: - >>> from twinkle_agentic.chunker import NativeChunker - >>> from twinkle_agentic.condenser.keyword import KeywordCondenser - >>> chunker = NativeChunker(chunk_size=1024) - >>> cond = KeywordCondenser( - ... num_relations=3, max_first_sentence_chars=160, - ... num_keywords=8, compression_ratio=4.0) - >>> traj = {'messages': [{'role': 'user', 'content': long_passage}]} - >>> chunks = cond(chunker(traj)) - >>> traj_compressed = chunks.to_trajectory() - """ - - def __init__( - self, - num_relations: int = 3, - max_first_sentence_chars: int = 160, - num_keywords: int = 8, - compression_ratio: float = 4.0, - spacy_model: str = 'en_core_web_sm', - min_chars: int = 200, - skip_roles: Sequence[str] = ('system', 'tool', 'assistant'), - rounds: Sequence[int] | None = None, - ): - if num_relations < 0: - raise ValueError(f'num_relations must be >= 0, got {num_relations}') - if num_keywords < 0: - raise ValueError(f'num_keywords must be >= 0, got {num_keywords}') - if max_first_sentence_chars < 0: - raise ValueError(f'max_first_sentence_chars must be >= 0, got {max_first_sentence_chars}') - if compression_ratio <= 1.0: - raise ValueError(f'compression_ratio must be > 1, got {compression_ratio}') - if min_chars < 0: - raise ValueError(f'min_chars must be >= 0, got {min_chars}') - - self.num_relations = num_relations - self.max_first_sentence_chars = max_first_sentence_chars - self.num_keywords = num_keywords - self.compression_ratio = float(compression_ratio) - self.spacy_model = spacy_model - self.min_chars = min_chars - self.skip_roles = tuple(skip_roles) - self.rounds = set(rounds) if rounds is not None else None - - # ------------------------------------------------------------------ - def __call__(self, chunks: Chunks, **kwargs) -> Chunks: - nlp = _load_spacy(self.spacy_model) - out: list[Chunk] = [] - for c in chunks.chunks: - if not self._should_condense(c): - out.append(c) - continue - compressed = self._condense(c['content'], nlp) - out.append(self._mark_condensed(c, compressed)) - return Chunks(chunks=out) - - # ------------------------------------------------------------------ - # selection policy - # ------------------------------------------------------------------ - def _should_condense(self, chunk: Chunk) -> bool: - if chunk.get('type') != 'text': - return False - if chunk.get('role') in self.skip_roles: - return False - if self.rounds is not None and chunk.get('round') not in self.rounds: - return False - content = chunk.get('content') - if not isinstance(content, str) or not content: - return False - if len(content) < self.min_chars: - return False - raw = chunk.get('raw') or {} - if isinstance(raw, dict): - # Chunker-emitted reasoning / tool-call text chunks carry a - # non-empty ``kind`` marker; leave them alone. - if raw.get('kind'): - return False - # Idempotency — don't re-condense already condensed chunks. - if raw.get('condensed'): - return False - return True - - @staticmethod - def _mark_condensed(chunk: Chunk, content: str) -> Chunk: - new: dict[str, Any] = dict(chunk) - raw = dict(new.get('raw') or {}) - raw.setdefault('original', new.get('content', '')) - new['content'] = content - raw['condensed'] = True - new['raw'] = raw - return new # type: ignore[return-value] - - # ------------------------------------------------------------------ - # core extractive compression - # ------------------------------------------------------------------ - def _condense(self, text: str, nlp) -> str: - budget = max(1, math.ceil(len(text) / self.compression_ratio)) - doc = nlp(text) - opening = _extract_opening(doc, self.max_first_sentence_chars) - excluded = _word_tokens_lower(opening) - triples = _extract_triples(doc, self.num_relations) - keywords = _extract_keywords(doc, self.num_keywords, excluded) - return _fit_under_budget(opening, triples, keywords, budget, fallback_text=text) diff --git a/src/twinkle_agentic/condenser/model.py b/src/twinkle_agentic/condenser/model.py deleted file mode 100644 index 521d38063..000000000 --- a/src/twinkle_agentic/condenser/model.py +++ /dev/null @@ -1,508 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -from __future__ import annotations - -import math -import re -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence, Tuple - -from twinkle_agentic.condenser.base import Condenser -from twinkle_agentic.data_format import Chunk, Chunks - -if TYPE_CHECKING: - from twinkle.data_format import SamplingParams, Trajectory # noqa: F401 - from twinkle.sampler.base import Sampler # noqa: F401 - -_SECTION_SCHEMA = """You are a text compression assistant. A downstream model will read your compressed output to decide whether the detail it needs is inside this block; if yes, it will fetch and read the original passage. - -Downstream model workflow: -Read your compressed output -> Decide whether needed info is in this block -> If yes -> Fetch original. - -Therefore your compression MUST NOT lose major information from the source. - -Output format: - -```text -## Summary -Overview plus facts STRONGLY RELATED to the Query, stated explicitly. - -## More -A collapsed index; expansion required to see specific information. -``` - -Rules: -1. Telegraphic style — drop function words ("the", "a", "is", "are", "of", ...); colons and commas mean "is" / "has". - * Exception: KEEP role-tagging verb+preposition phrases verbatim ("published by X", "written by X", "directed by X", "starring X", "founded by X", "created by X", "composed by X", "produced by X", "based on X", "adapted from X"). Collapsing these to a bare name loses the relation role (author vs publisher vs director) that the downstream question may hinge on. -2. Summary MUST contain the passage's primary topic + 2–4 concrete core facts drawn from the source (entities, numbers, dates, relations). If a Query is given, order Query-relevant facts first, but STILL include other core facts within the budget. A Query is an ORDERING HINT, NOT a filter. -3. Summary MUST NOT be meta-commentary about the Query. Forbidden patterns: "no X mention", "Query info: absent", "passage covers Y only", "does not contain ...", "no relevant info", or summaries that are only abstract category words like "structure/order/usage" with no facts. If the passage is unrelated to the Query, you still summarize the passage normally. -4. More is an INDEX of category keywords, NOT inline data. Enumerate what CAN be recovered from the source (e.g. "birthplace, death place, age"); do NOT paste dates/numbers/names inline. Make sure all category of useful facts are introduced here. -5. Output language MUST match the source language. -6. Do NOT fabricate. Do NOT omit major information. Any fact not in the source MUST NOT appear in your output. - -Example: - -Source: -```text -Marie Curie (7 Nov 1867 – 4 Jul 1934), born Maria Sklodowska in Warsaw (then Russian Poland); parents were teachers. Barred from Polish universities, she and her sister agreed to take turns funding each other's overseas study. - -In 1891 Marie reached Paris and enrolled at the Sorbonne, earning a physics degree (1893) and a mathematics degree (1894), becoming the school's first female physics lecturer. In 1895 she married French physicist Pierre Curie; they spent the rest of their lives on radioactivity research. - -In July 1898 she discovered polonium, named after her homeland Poland; in December she and Pierre announced the discovery of radium. She coined "radioactivity" and showed it is an atomic property, not a chemical reaction. - -In 1903 she shared the Nobel Prize in Physics with Pierre and Henri Becquerel. In 1911 she alone won the Nobel Prize in Chemistry for polonium and radium. She is the first woman to win a Nobel, and the only person to win Nobels in two different sciences. After Pierre died in a carriage accident in 1906, Marie took his chair and became the first female professor at the Sorbonne. - -During World War I she developed mobile X-ray units, called "Petites Curies" in French; about 20 were deployed to the front, examining over 1,000,000 wounded soldiers. - -She died of aplastic anaemia from radiation exposure on 4 July 1934 in Passy, Haute-Savoie, France, aged 66. Her notebooks remain highly radioactive, kept in lead boxes; researchers must wear protective gear to consult them. -``` - -Compressed: -```text -## Summary -Marie Curie: French-Polish physicist/chemist, founder of radioactivity research, first female Sorbonne professor. -- Nobel x2 (Physics + Chemistry); first woman Nobel laureate; only person with Nobels in two sciences. -- Discovered polonium + radium; coined "radioactivity"; proved it is an atomic property. - -## More -- birthplace, death place, age, cause of death -- degree years, in-school firsts x2 -- element naming origin, collaborators, full timeline -- Nobel year per prize, co-laureates, citation -- device name, deployment scale, patients treated -- notebook radioactivity, storage, access conditions -``` - -Now begin. -""" # noqa - -DEFAULT_SYSTEM_PROMPT = _SECTION_SCHEMA - -DEFAULT_USER_PROMPT_TEMPLATE = """\ -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. - -## Query (ordering hint only — still summarize the whole passage) -{query} - -## Target length -Compress AS MUCH AS faithfully possible. HARD CEILING: {budget} chars. \ -If core facts fit in far fewer chars, output fewer. \ -Never exceed the ceiling. - -## Passage -{text}""" - -# A (chunk_index, chunk, char_budget) triple marking one compression job. -_Job = Tuple[int, Chunk, int] - - -# --------------------------------------------------------------------------- -# ModelCondenser -# --------------------------------------------------------------------------- -class ModelCondenser(Condenser): - """Compressor that delegates summarization to an LLM via a :class:`Sampler`. - TODO: Experimental feature, wait for testing - Args: - sampler: Configured :class:`Sampler` with a template set. - compression_ratio: Target factor (> 1). Used only to derive a - soft character budget passed into the prompt and to size - ``SamplingParams.max_tokens``. Model output is NOT hard - truncated; a chunk whose decoded output is not strictly - shorter than the original passage is left unchanged (and - not flagged ``raw.condensed``). - sampling_params: Override for per-call sampling; when ``None`` a - greedy config is derived from the max budget in the batch. - system_prompt: Override for the system prompt. Used verbatim. - user_prompt_template: Override the user prompt. Must contain - ``{budget}`` and ``{text}``. ``{query}`` is optional and is - replaced with the trajectory's question extracted by the - ``related_query`` callback (see below); jobs without a - detected query get a neutral placeholder. - min_chars: Pre-filter; chunks shorter than this pass through. - min_budget_chars: Floor for the soft character budget exposed - to the prompt. When ``ceil(len / compression_ratio)`` falls - below this, the budget is raised to this floor so short - passages keep room for all three sections in the model's - plan. Since the condenser no longer hard-clips output, - this only influences prompt wording and sampling token - limits; pass ``1`` to use the raw ratio everywhere. - template: Optional :class:`Template`. When provided, its - ``tokenizer.all_special_tokens`` are stripped from every - decoded response before length-clamping, preventing - protocol tokens (``<|im_end|>``, ``<|eot_id|>``, ````, - ...) from leaking into the compressed output. When - omitted, falls back to ``sampler.template`` if available. - skip_roles: Roles whose chunks are never compressed. - skip_pattern: Optional regex (compiled with ``re.MULTILINE``). - Any chunk whose ``content`` has a match for this pattern - is passed through unchanged, regardless of length / ratio. - Uses :func:`re.search` semantics, so anchor with ``^`` / - start-of-string if you want boundary-matching only (e.g. - ``r'^Question:'`` to preserve the question prefix in a - HotpotQA-style user message). ``None`` disables the filter. - This flag is purely a compression-skip filter; query - extraction is the orthogonal job of ``related_query``. - related_query: Optional ``(chunk) -> Optional[str]`` callback - that returns the query string carried by ``chunk`` (e.g. - the user's HotpotQA question), or ``None`` if the chunk - is not a query carrier. Walked in chunk order; the most - recently returned non-``None`` query is broadcast to all - subsequent condense-eligible chunks until the next hit. - Because :class:`MultiTurnCondenseRollout` may merge - multiple trajectories into one chunk list, each - trajectory's question chunk must precede its passages so - this rolling state correctly partitions queries - per-trajectory. ``None`` disables query injection (the - ``{query}`` slot collapses to a neutral placeholder). - rounds: Optional set of conversation turn indices to compress. - ``None`` = no round-based filter; chunks lacking a ``round`` - field are skipped when this filter is active. - batch_size: Max chunks per sampler call. Partial batches are - padded with a duplicate of the last trajectory so that - distributed samplers (DP slice) always receive a full batch. - lora_path: Optional LoRA adapter to use for compression. - - ``None`` (default): forwards ``use_base_model=True`` to - :meth:`Sampler.sample` so compression bypasses any - currently-synced LoRA — strongly recommended when the - sampler is also the training policy. - - ``str``: forwards ``adapter_path=lora_path`` so a - dedicated condenser LoRA (e.g. a ModelScope slug or - local directory) is loaded and used instead of the base. - - Compressed chunks are flagged ``raw.condensed=True``; a subsequent - :meth:`Chunks.to_trajectory` call wraps them in ````. - - Example:: - - >>> from twinkle.sampler import vLLMSampler - >>> sampler = vLLMSampler(model_id='Qwen/Qwen2.5-3B-Instruct', - ... engine_args={'dtype': 'bfloat16'}) - >>> sampler.set_template('qwen2_5') - >>> cond = ModelCondenser(sampler, compression_ratio=2.0) - >>> compressed = cond(chunks) - """ - - def __init__( - self, - sampler: Sampler, - compression_ratio: float = 2.0, - *, - sampling_params: SamplingParams | None = None, - system_prompt: str | None = None, - user_prompt_template: str | None = None, - min_chars: int = 200, - min_budget_chars: int = 250, - template: Any | None = None, - skip_roles: Sequence[str] = ('system', 'tool', 'assistant'), - skip_pattern: str | None = None, - related_query: Callable[[Chunk], str | None] | None = None, - rounds: Sequence[int] | None = None, - batch_size: int = None, - lora_path: str | None = None, - ): - if sampler is None: - raise ValueError('sampler is required') - if compression_ratio <= 1.0: - raise ValueError(f'compression_ratio must be > 1, got {compression_ratio}') - if min_chars < 0: - raise ValueError(f'min_chars must be >= 0, got {min_chars}') - if min_budget_chars < 1: - raise ValueError(f'min_budget_chars must be >= 1, got {min_budget_chars}') - if batch_size is not None and batch_size <= 0: - raise ValueError(f'batch_size must be >= 1, got {batch_size}') - - tpl = user_prompt_template or DEFAULT_USER_PROMPT_TEMPLATE - if '{budget}' not in tpl or '{text}' not in tpl: - raise ValueError('user_prompt_template must contain both {budget} and {text}') - - self.sampler = sampler - self.compression_ratio = float(compression_ratio) - self.sampling_params = sampling_params - self.system_prompt = system_prompt or DEFAULT_SYSTEM_PROMPT - self.user_prompt_template = tpl - self.min_chars = min_chars - self.min_budget_chars = int(min_budget_chars) - self.template = template - self.skip_roles = tuple(skip_roles) - # ``^`` must anchor to start-of-string, not start-of-line: a passage - # whose body contains a ``Question:`` line would otherwise skip compression. - self.skip_re: re.Pattern | None = (re.compile(skip_pattern) if skip_pattern else None) - self.related_query = related_query - self.rounds = set(rounds) if rounds is not None else None - self.batch_size = batch_size - self.lora_path = lora_path if lora_path else None - self._special_tokens_cache: tuple[str, ...] | None = None - - # ------------------------------------------------------------------ - # entry point - # ------------------------------------------------------------------ - def __call__(self, chunks: Chunks, **_kwargs: Any) -> Chunks: - out: list[Chunk] = list(chunks.chunks) - items = self._collect_jobs(out) - if not items: - return Chunks(chunks=out) - - batch_size = self.batch_size or len(items) - for start in range(0, len(items), batch_size): - sub = items[start:start + batch_size] - batch = [job for job, _q in sub] - queries = [q for _job, q in sub] - responses = self._sample_batch(batch, queries=queries) - for (idx, chunk, _budget), resp in zip(batch, responses): - text = self._postprocess(_decoded(resp), chunk['content']) - if text is None: - continue - out[idx] = _mark_condensed(chunk, text) - return Chunks(chunks=out) - - # ------------------------------------------------------------------ - # eligibility + job collection - # ------------------------------------------------------------------ - def _collect_jobs( - self, - chunks: Sequence[Chunk], - ) -> list[tuple[_Job, str | None]]: - """Collect compression jobs, tagging each with its trajectory's query. - - Walks ``chunks`` in order and maintains a rolling - ``current_query`` state driven by the ``related_query`` - callback: every chunk for which the callback returns a - non-``None`` string updates the state, and every subsequent - condense-eligible chunk picks up the most recent query. - Because the chunker emits each trajectory's question chunk - before its passages, this walk correctly partitions queries - per-trajectory even when ``MultiTurnCondenseRollout`` merges - multiple trajectories into a single chunk list — A's - passages only ever see A's question, B's only B's. - """ - items: list[tuple[_Job, str | None]] = [] - current_query: str | None = None - extract = self.related_query - for i, c in enumerate(chunks): - content = c.get('content') - if extract is not None: - q = extract(c) - if isinstance(q, str) and q: - current_query = q - if not self._should_condense(c): - continue - budget = max(self.min_budget_chars, math.ceil(len(content) / self.compression_ratio)) - if budget >= len(content): - continue - items.append(((i, c, max(1, budget)), current_query)) - return items - - def _should_condense(self, chunk: Chunk) -> bool: - if chunk.get('type') != 'text': - return False - if chunk.get('role') in self.skip_roles: - return False - if self.rounds is not None and chunk.get('round') not in self.rounds: - return False - content = chunk.get('content') - if not isinstance(content, str) or len(content) < self.min_chars: - return False - if self.skip_re is not None and self.skip_re.search(content): - return False - raw = chunk.get('raw') or {} - if isinstance(raw, dict): - # Skip chunker-emitted reasoning / tool_call text chunks. - if raw.get('kind'): - return False - # Idempotent — never re-compress something already compressed. - if raw.get('condensed'): - return False - return True - - # ------------------------------------------------------------------ - # batched sampling - # ------------------------------------------------------------------ - def _sample_batch( - self, - batch: Sequence[_Job], - *, - queries: Sequence[str | None] = (), - ) -> list[Any]: - """Dispatch one batch to the sampler, padded to ``batch_size``. - - Distributed samplers slice inputs across DP workers and can - mis-behave when the final batch is smaller than ``batch_size``; - we pad with a duplicate of the last trajectory and trim the - matching extra responses here. - - ``queries`` is aligned 1:1 with ``batch``; each per-job query - is injected into the user prompt's ``{query}`` slot. When - empty or ``None`` at an index, a neutral placeholder is used. - """ - qs: list[str | None] = list(queries) if queries else [None] * len(batch) - if len(qs) != len(batch): - raise ValueError(f'queries length ({len(qs)}) must match batch length ' - f'({len(batch)})') - trajectories = [ - self._build_trajectory(chunk['content'], budget, query=q) for (_, chunk, budget), q in zip(batch, qs) - ] - actual = len(trajectories) - device_mesh = getattr(self.sampler, 'device_mesh', None) - min_batch_size = (device_mesh.data_world_size if device_mesh is not None else 1) - if actual < min_batch_size: - trajectories.extend([trajectories[-1]] * (min_batch_size - actual)) - - sp = self._sampling_params_for(max(b for _, _, b in batch)) - kwargs: dict[str, Any] = {'sampling_params': sp} - if self.lora_path is None: - kwargs['use_base_model'] = True - else: - kwargs['adapter_path'] = self.lora_path - responses = self.sampler.sample(trajectories, **kwargs) - # Coerce to list (some samplers may return tuples) and drop - # padding responses so downstream ``zip`` aligns with ``batch``. - return list(responses)[:actual] - - def _build_trajectory( - self, - text: str, - budget: int, - *, - query: str | None = None, - ) -> Trajectory: - system = self.system_prompt - user = self.user_prompt_template.replace('{budget}', str(budget)) - user = user.replace('{text}', text) - q_text = ( - query.strip() if isinstance(query, str) and query and query.strip() else - '(no explicit query; compress by general salience)') - user = user.replace('{query}', q_text) - return { # type: ignore[return-value] - 'messages': [ - {'role': 'system', 'content': system}, - {'role': 'user', 'content': user}, - ], - } - - def _sampling_params_for(self, budget: int) -> SamplingParams: - if self.sampling_params is not None: - return self.sampling_params - from twinkle.data_format.sampling import SamplingParams - - # CJK worst case ~2 tokens/char; budget is a soft char ceiling, not output truth. - max_new = max(512, budget * 3 + 128) - return SamplingParams(temperature=0.0, max_tokens=max_new) - - # ------------------------------------------------------------------ - # postprocess - # ------------------------------------------------------------------ - def _postprocess(self, raw: str, original: str) -> str | None: - """Return compressed text, or ``None`` to signal passthrough. - - ``None`` is returned when the decoded output is empty, - degenerate (markdown markers only, no alphanumerics), or its - character length is **not strictly shorter** than ``original`` - — in which case the model failed to produce a useful - compression and the caller should keep the original passage - verbatim (no ```` wrap, not marked ``raw.condensed``). - """ - text = _strip_special_tokens(_strip_code_fences(raw), self._get_special_tokens()).strip() - if not text or not _has_alnum(text): - return None - if len(text) >= len(original): - return None - return text - - def _get_special_tokens(self) -> tuple[str, ...]: - """Return protocol tokens to strip from decoded output (cached). - - Resolution order: - - 1. ``self.template.tokenizer`` — explicit template passed to - ``__init__``. Preferred in distributed setups where - ``sampler.template`` on the driver is a proxy and may be - ``None``. - 2. ``self.sampler.template.tokenizer`` — best-effort fallback - for single-process use. - 3. Empty tuple — no stripping (safe no-op). - - Uses ``tokenizer.all_special_tokens`` when available so the - full eos/bos/pad/unk/sep/cls/mask/additional set is covered - in one shot; this means ChatML (``<|im_end|>``), Llama - (``<|eot_id|>``), T5 (````) etc. are all handled without - per-model hard-coding. - """ - if self._special_tokens_cache is not None: - return self._special_tokens_cache - tpl = self.template or getattr(self.sampler, 'template', None) - tokenizer = getattr(tpl, 'tokenizer', None) if tpl is not None else None - tokens: list[str] = [] - if tokenizer is not None: - extras = getattr(tokenizer, 'all_special_tokens', None) or [] - if extras: - tokens.extend(t for t in extras if isinstance(t, str) and t and not t.isspace()) - else: - for attr in ('eos_token', 'pad_token', 'bos_token'): - t = getattr(tokenizer, attr, None) - if isinstance(t, str) and t: - tokens.append(t) - # Order-preserving dedupe. - self._special_tokens_cache = tuple(dict.fromkeys(tokens)) - return self._special_tokens_cache - - -# --------------------------------------------------------------------------- -# pure helpers -# --------------------------------------------------------------------------- -_CODE_FENCE_RE = re.compile(r'^```[a-zA-Z]*\s*\n(.*?)\n```\s*$', re.DOTALL) - - -def _decoded(response: Any) -> str: - """Extract the first decoded sequence, or ``''`` on empty/malformed input.""" - seqs = getattr(response, 'sequences', None) or [] - if not seqs: - return '' - return getattr(seqs[0], 'decoded', None) or '' - - -def _mark_condensed(chunk: Chunk, content: str) -> Chunk: - """Return a shallow copy of ``chunk`` with compressed ``content`` - and ``raw.condensed=True`` (preserving any original content under - ``raw.original`` so a future :class:`ExtractCondensed` call can - recover the full text). - """ - new: dict[str, Any] = dict(chunk) - raw = dict(new.get('raw') or {}) - raw.setdefault('original', new.get('content', '')) - raw['condensed'] = True - new['content'] = content - new['raw'] = raw - return new # type: ignore[return-value] - - -def _strip_code_fences(text: str) -> str: - """Unwrap a leading/trailing triple-backtick fence if present.""" - stripped = text.strip() - m = _CODE_FENCE_RE.match(stripped) - return m.group(1) if m else text - - -def _strip_special_tokens(text: str, tokens: Sequence[str]) -> str: - """Remove tokenizer special tokens that leaked through decode. - - ``tokens`` is typically ``tokenizer.all_special_tokens`` from the - template's tokenizer (see :meth:`ModelCondenser._get_special_tokens`). - Uses literal :meth:`str.replace` rather than a regex so we only - strip registered protocol markers and never legitimate passage - content that happens to look like ``<|...|>``. - """ - for tok in tokens: - if tok and tok in text: - text = text.replace(tok, '') - return text - - -def _has_alnum(text: str) -> bool: - """True iff ``text`` contains at least one alphanumeric character. - - Used to detect degenerate model outputs like ``'##'`` or ``'- '`` - that are pure markdown markers with no actual words. - """ - return any(ch.isalnum() for ch in text) diff --git a/src/twinkle_agentic/train/__init__.py b/src/twinkle_agentic/train/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/twinkle_agentic/train/base.py b/src/twinkle_agentic/train/base.py new file mode 100644 index 000000000..a0fd281e3 --- /dev/null +++ b/src/twinkle_agentic/train/base.py @@ -0,0 +1,6 @@ + + +class Trainer: + + def train(self): + pass \ No newline at end of file diff --git a/src/twinkle_agentic/train/cron.py b/src/twinkle_agentic/train/cron.py new file mode 100644 index 000000000..d118df8bc --- /dev/null +++ b/src/twinkle_agentic/train/cron.py @@ -0,0 +1,5 @@ + + +class CronTrainManager: + + pass \ No newline at end of file diff --git a/src/twinkle_agentic/utils/llm_backup.py b/src/twinkle_agentic/utils/llm_backup.py new file mode 100644 index 000000000..ce5199932 --- /dev/null +++ b/src/twinkle_agentic/utils/llm_backup.py @@ -0,0 +1,314 @@ +import functools +import hashlib +import inspect +import json +import os +import random +import threading +from collections import defaultdict +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple + + +@dataclass +class EvalRecord: + """A single evaluation record comparing student and teacher outputs.""" + student_result: Any + teacher_result: Any + match: bool + + +@dataclass +class DistillationState: + """Per-key state tracking confidence, dataset, and call count.""" + confidence: Optional[float] = None + dataset: List[EvalRecord] = field(default_factory=list) + call_count: int = 0 + + +class DistillationRegistry: + """Thread-safe registry maintaining distillation state per unique key.""" + + def __init__(self): + self._states: Dict[str, DistillationState] = defaultdict(DistillationState) + self._lock = threading.Lock() + + def get_confidence(self, key: str) -> float: + """Get confidence for a key. Returns 0.0 if no data available.""" + with self._lock: + state = self._states[key] + if state.confidence is not None: + return state.confidence + if state.dataset: + state.confidence = self._compute_confidence(state.dataset) + return state.confidence + return 0.0 + + def increment_call(self, key: str) -> int: + with self._lock: + state = self._states[key] + state.call_count += 1 + return state.call_count + + def add_record(self, key: str, student_result: Any, teacher_result: Any, match: bool): + with self._lock: + state = self._states[key] + state.dataset.append(EvalRecord( + student_result=student_result, + teacher_result=teacher_result, + match=match, + )) + + def refresh_confidence(self, key: str) -> float: + with self._lock: + state = self._states[key] + if state.dataset: + state.confidence = self._compute_confidence(state.dataset) + else: + state.confidence = 0.0 + return state.confidence + + @staticmethod + def _compute_confidence(dataset: List[EvalRecord]) -> float: + if not dataset: + return 0.0 + matches = sum(1 for r in dataset if r.match) + return matches / len(dataset) + + +# --------------------------------------------------------------------------- +# Global state +# --------------------------------------------------------------------------- +_registry = DistillationRegistry() +_teacher_api = None +_teacher_lock = threading.Lock() + + +def _get_teacher_api(): + """Lazy-init global teacher API from environment variables. + + Env vars: + LLM_BACKUP_MODEL: Model name (default: "gpt-4o") + LLM_BACKUP_API_KEY: API key + LLM_BACKUP_BASE_URL: Base URL for OpenAI-compatible endpoint + """ + global _teacher_api + if _teacher_api is not None: + return _teacher_api + with _teacher_lock: + if _teacher_api is not None: + return _teacher_api + from twinkle_agentic.protocol.openai import OpenAI + _teacher_api = OpenAI( + model=os.environ.get('LLM_BACKUP_MODEL', 'gpt-4o'), + api_key=os.environ.get('LLM_BACKUP_API_KEY'), + base_url=os.environ.get('LLM_BACKUP_BASE_URL'), + ) + return _teacher_api + + +def _call_teacher(trajectory, sampling_params) -> str: + """Call teacher API and extract raw content string.""" + api = _get_teacher_api() + message = api(trajectory, sampling_params) + if isinstance(message, list): + message = message[0] + return message.get('content', '') if isinstance(message, dict) else '' + + +# --------------------------------------------------------------------------- +# Key building +# --------------------------------------------------------------------------- +def _build_key(func_name: str, args: tuple, kwargs: dict, + param_names: List[str], key_params: Sequence[str]) -> str: + """Build a unique key from specified parameter values.""" + key_parts = [func_name] + for i, name in enumerate(param_names): + if name in key_params: + if i < len(args): + key_parts.append(f"{name}={_serialize_value(args[i])}") + elif name in kwargs: + key_parts.append(f"{name}={_serialize_value(kwargs[name])}") + for name in key_params: + if name not in param_names[:len(args)] and name in kwargs: + if f"{name}={_serialize_value(kwargs[name])}" not in key_parts: + key_parts.append(f"{name}={_serialize_value(kwargs[name])}") + raw_key = "|".join(key_parts) + return hashlib.md5(raw_key.encode()).hexdigest() + + +def _serialize_value(value: Any) -> str: + try: + return json.dumps(value, sort_keys=True, default=str) + except (TypeError, ValueError): + return str(value) + + +def _extract_param(args: tuple, kwargs: dict, param_names: List[str], name: str) -> Any: + """Extract a named parameter from args/kwargs given the signature's param_names.""" + if name in kwargs: + return kwargs[name] + for i, pname in enumerate(param_names): + if pname == name and i < len(args): + return args[i] + return None + + +# --------------------------------------------------------------------------- +# Decorator +# --------------------------------------------------------------------------- +def llm_backup( + key_params: Sequence[str], + comparator: Optional[Callable[[Any, Any], bool]] = None, + sample_rate: float = 0.2, + refresh_env_var: str = "LLM_BACKUP_REFRESH_INTERVAL", + default_refresh_interval: int = 50, +): + """Decorator for progressive distillation from teacher API to student model. + + The decorated function is the STUDENT (local model sampling). The TEACHER + is a global OpenAI-compatible API constructed from environment variables. + + The decorated function MUST accept ``trajectory`` and ``sampling_params`` + as parameters (by name) and return a raw string. This ensures: + - Teacher and student receive identical inputs + - The dataset contains raw (trajectory, student_output, teacher_output) tuples + - No pre/post processing is included, making data directly trainable + + Routing logic: + - confidence% -> use student (decorated fn) + - Of those, sample_rate% also call teacher for comparison + - (1 - confidence)% -> use teacher API + - Always also call student for comparison + + Every N calls the confidence is recalculated from the comparison dataset. + + Environment variables: + LLM_BACKUP_MODEL: Teacher model name (default "gpt-4o") + LLM_BACKUP_API_KEY: Teacher API key + LLM_BACKUP_BASE_URL: Teacher API base URL + LLM_BACKUP_REFRESH_INTERVAL: Confidence refresh interval N (default 50) + + Args: + key_params: Parameter names for unique confidence key (e.g. ["query"]). + comparator: function(student, teacher) -> bool. Default: equality. + sample_rate: Probability of teacher verification when student is used. + refresh_env_var: Env var name for refresh interval. + default_refresh_interval: Default refresh interval. + + Example: + >>> @llm_backup(key_params=["query"]) + ... def _sample(self, trajectory, sampling_params, query=None) -> str: + ... responses = self.sampler.sample([trajectory], ...) + ... return decode(responses[0]) + """ + if comparator is None: + comparator = lambda a, b: a == b # noqa: E731 + + def decorator(fn: Callable) -> Callable: + sig = inspect.signature(fn) + param_names = list(sig.parameters.keys()) + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + key = _build_key(fn.__qualname__, args, kwargs, param_names, key_params) + confidence = _registry.get_confidence(key) + + try: + refresh_interval = int(os.environ.get(refresh_env_var, default_refresh_interval)) + except (ValueError, TypeError): + refresh_interval = default_refresh_interval + + # Extract trajectory and sampling_params for teacher call + trajectory = _extract_param(args, kwargs, param_names, 'trajectory') + sampling_params = _extract_param(args, kwargs, param_names, 'sampling_params') + + roll = random.random() + use_student = roll < confidence + + if use_student: + # High confidence: trust student + result = fn(*args, **kwargs) + # Occasionally verify against teacher + if random.random() < sample_rate: + teacher_result = _call_teacher(trajectory, sampling_params) + match = comparator(result, teacher_result) + _registry.add_record(key, result, teacher_result, match) + if not match: + result = teacher_result + else: + # Low confidence: use teacher + teacher_result = _call_teacher(trajectory, sampling_params) + student_result = fn(*args, **kwargs) + match = comparator(student_result, teacher_result) + _registry.add_record(key, student_result, teacher_result, match) + result = teacher_result + + call_count = _registry.increment_call(key) + if refresh_interval > 0 and call_count % refresh_interval == 0: + _registry.refresh_confidence(key) + + return result + + wrapper._registry = _registry + return wrapper + + return decorator + + +def llm_backup_async( + key_params: Sequence[str], + comparator: Optional[Callable[[Any, Any], bool]] = None, + sample_rate: float = 0.2, + refresh_env_var: str = "LLM_BACKUP_REFRESH_INTERVAL", + default_refresh_interval: int = 50, +): + """Async version of llm_backup. Same semantics.""" + if comparator is None: + comparator = lambda a, b: a == b # noqa: E731 + + def decorator(fn: Callable) -> Callable: + sig = inspect.signature(fn) + param_names = list(sig.parameters.keys()) + + @functools.wraps(fn) + async def wrapper(*args, **kwargs): + key = _build_key(fn.__qualname__, args, kwargs, param_names, key_params) + confidence = _registry.get_confidence(key) + + try: + refresh_interval = int(os.environ.get(refresh_env_var, default_refresh_interval)) + except (ValueError, TypeError): + refresh_interval = default_refresh_interval + + trajectory = _extract_param(args, kwargs, param_names, 'trajectory') + sampling_params = _extract_param(args, kwargs, param_names, 'sampling_params') + + roll = random.random() + use_student = roll < confidence + + if use_student: + result = await fn(*args, **kwargs) + if random.random() < sample_rate: + teacher_result = _call_teacher(trajectory, sampling_params) + match = comparator(result, teacher_result) + _registry.add_record(key, result, teacher_result, match) + if not match: + result = teacher_result + else: + teacher_result = _call_teacher(trajectory, sampling_params) + student_result = await fn(*args, **kwargs) + match = comparator(student_result, teacher_result) + _registry.add_record(key, student_result, teacher_result, match) + result = teacher_result + + call_count = _registry.increment_call(key) + if refresh_interval > 0 and call_count % refresh_interval == 0: + _registry.refresh_confidence(key) + + return result + + wrapper._registry = _registry + return wrapper + + return decorator From ed45d04ce9cea076d59bc409389fba2f8867bf5a Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Tue, 7 Jul 2026 20:56:16 +0800 Subject: [PATCH 02/60] wip --- src/twinkle_agentic/condenser/__init__.py | 4 - src/twinkle_agentic/summarizer/__init__.py | 4 + .../{condenser => summarizer}/base.py | 16 +- .../summarizer/error_summarizer.py | 0 .../fact_summarizer.py} | 4 +- .../summarizer/pattern_summarizer.py | 0 src/twinkle_agentic/verifier/__init__.py | 16 + src/twinkle_agentic/verifier/base.py | 11 + src/twinkle_agentic/verifier/domain_checks.py | 436 +++++++++++++ src/twinkle_agentic/verifier/hard_scorer.py | 409 ++++++++++++ .../verifier/rubric_verifier.py | 586 ++++++++++++++++++ 11 files changed, 1472 insertions(+), 14 deletions(-) delete mode 100644 src/twinkle_agentic/condenser/__init__.py create mode 100644 src/twinkle_agentic/summarizer/__init__.py rename src/twinkle_agentic/{condenser => summarizer}/base.py (94%) create mode 100644 src/twinkle_agentic/summarizer/error_summarizer.py rename src/twinkle_agentic/{condenser/facts.py => summarizer/fact_summarizer.py} (97%) create mode 100644 src/twinkle_agentic/summarizer/pattern_summarizer.py create mode 100644 src/twinkle_agentic/verifier/__init__.py create mode 100644 src/twinkle_agentic/verifier/base.py create mode 100644 src/twinkle_agentic/verifier/domain_checks.py create mode 100644 src/twinkle_agentic/verifier/hard_scorer.py create mode 100644 src/twinkle_agentic/verifier/rubric_verifier.py diff --git a/src/twinkle_agentic/condenser/__init__.py b/src/twinkle_agentic/condenser/__init__.py deleted file mode 100644 index a48fd73c1..000000000 --- a/src/twinkle_agentic/condenser/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .base import Condenser -from .facts import FactsCondenser - -__all__ = ['Condenser', 'FactsCondenser'] diff --git a/src/twinkle_agentic/summarizer/__init__.py b/src/twinkle_agentic/summarizer/__init__.py new file mode 100644 index 000000000..9a4672923 --- /dev/null +++ b/src/twinkle_agentic/summarizer/__init__.py @@ -0,0 +1,4 @@ +from .base import Summarizer +from .fact_summarizer import FactSummarizer + +__all__ = ['Summarizer', 'FactSummarizer'] diff --git a/src/twinkle_agentic/condenser/base.py b/src/twinkle_agentic/summarizer/base.py similarity index 94% rename from src/twinkle_agentic/condenser/base.py rename to src/twinkle_agentic/summarizer/base.py index 8fd2acd61..8e9be5f5f 100644 --- a/src/twinkle_agentic/condenser/base.py +++ b/src/twinkle_agentic/summarizer/base.py @@ -22,8 +22,8 @@ {text}""" -class Condenser: - """Base condenser with progressive distillation via llm_backup. +class Summarizer: + """Base summarizer with progressive distillation via llm_backup. Subclasses customize compression behavior by providing their own ``system_prompt``, ``user_prompt_template``, and ``lora_path``. @@ -40,12 +40,12 @@ class Condenser: compression_ratio: Target compression factor (> 1). model_path: Model identifier. sampling_params: Default sampling params. - system_prompt: System prompt for this condenser type. + system_prompt: System prompt for this summarizer type. user_prompt_template: User prompt template. Must contain ``{budget}`` and ``{text}``. May contain ``{query}``. min_budget_chars: Floor for the character budget in the prompt. template: Optional :class:`Template` for special token stripping. - lora_path: LoRA adapter path specific to this condenser type. + lora_path: LoRA adapter path specific to this summarizer type. Each subclass can use a different LoRA for its task. """ @@ -168,9 +168,9 @@ def _default_sampling_params(budget: int): @staticmethod def _postprocess(raw: str, original: str, special_tokens: tuple[str, ...]) -> str | None: - text = Condenser._strip_special_tokens( - Condenser._strip_code_fences(raw), special_tokens).strip() - if not text or not Condenser._has_alnum(text): + text = Summarizer._strip_special_tokens( + Summarizer._strip_code_fences(raw), special_tokens).strip() + if not text or not Summarizer._has_alnum(text): return None if len(text) >= len(original): return None @@ -186,7 +186,7 @@ def _decoded(response: Any) -> str: @staticmethod def _strip_code_fences(text: str) -> str: stripped = text.strip() - m = Condenser._CODE_FENCE_RE.match(stripped) + m = Summarizer._CODE_FENCE_RE.match(stripped) return m.group(1) if m else text @staticmethod diff --git a/src/twinkle_agentic/summarizer/error_summarizer.py b/src/twinkle_agentic/summarizer/error_summarizer.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/twinkle_agentic/condenser/facts.py b/src/twinkle_agentic/summarizer/fact_summarizer.py similarity index 97% rename from src/twinkle_agentic/condenser/facts.py rename to src/twinkle_agentic/summarizer/fact_summarizer.py index 6a8313b9b..ab24a081b 100644 --- a/src/twinkle_agentic/condenser/facts.py +++ b/src/twinkle_agentic/summarizer/fact_summarizer.py @@ -3,7 +3,7 @@ from typing import TYPE_CHECKING, Any -from twinkle_agentic.condenser.base import Condenser +from twinkle_agentic.summarizer.base import Summarizer if TYPE_CHECKING: from twinkle.data_format import SamplingParams # noqa: F401 @@ -57,7 +57,7 @@ {text}""" -class FactsCondenser(Condenser): +class FactSummarizer(Summarizer): def __init__( self, diff --git a/src/twinkle_agentic/summarizer/pattern_summarizer.py b/src/twinkle_agentic/summarizer/pattern_summarizer.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/twinkle_agentic/verifier/__init__.py b/src/twinkle_agentic/verifier/__init__.py new file mode 100644 index 000000000..3265fed9b --- /dev/null +++ b/src/twinkle_agentic/verifier/__init__.py @@ -0,0 +1,16 @@ +from .base import Verifier +from .domain_checks import (check_answer_match, check_code_parses, + check_instruction_constraints, check_not_degenerate, + check_numeric_equiv, check_output_format, + default_checks_for) +from .hard_scorer import CheckResult, HardScorer, HardScoreDetail, TrajectoryView +from .rubric_verifier import RubricItem, RubricVerifier, ScoreDetail + +__all__ = [ + 'Verifier', + 'RubricVerifier', 'RubricItem', 'ScoreDetail', + 'HardScorer', 'HardScoreDetail', 'CheckResult', 'TrajectoryView', + 'check_output_format', 'check_numeric_equiv', 'check_answer_match', + 'check_code_parses', 'check_instruction_constraints', 'check_not_degenerate', + 'default_checks_for', +] diff --git a/src/twinkle_agentic/verifier/base.py b/src/twinkle_agentic/verifier/base.py new file mode 100644 index 000000000..534206d36 --- /dev/null +++ b/src/twinkle_agentic/verifier/base.py @@ -0,0 +1,11 @@ +from abc import ABC, abstractmethod + + +class Verifier(ABC): + """Reward verifier that scores a sample on a 5-level scale (0-4).""" + + NUM_LEVELS = 5 + + @abstractmethod + def __call__(self, trajectory: dict, **kwargs) -> int: + pass diff --git a/src/twinkle_agentic/verifier/domain_checks.py b/src/twinkle_agentic/verifier/domain_checks.py new file mode 100644 index 000000000..92ac59a64 --- /dev/null +++ b/src/twinkle_agentic/verifier/domain_checks.py @@ -0,0 +1,436 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Domain-specific deterministic checks for :class:`HardScorer`. + +These are LLM-free, dependency-free (stdlib ``ast``/``re``/``json`` only) and +reuse the answer-extraction / F1 helpers already in ``reward/f1.py``. They are +factories: call them with config and get back a plain ``CheckFn`` that plugs +into ``HardScorer(checks=[...])``. + +Coverage (initial, no sandbox): +- output format: ``\\boxed{}`` / fenced code block / parseable JSON present +- numeric equivalence: lightweight fraction/decimal/percent normalization +- reference match: F1/EM vs ``ground_truth`` (reuses ``_f1_score``) +- code syntax: ``ast.parse`` on the last fenced block (stdlib, does NOT run) +- instruction constraints: length / keyword must-include / must-exclude / lang +- degeneration: empty / too-short / repetitive final answer + +Sandbox-based math (``math-verify``/sympy) and code execution (unit tests) can +be added later as additional CheckFns without touching HardScorer. +""" +from __future__ import annotations + +import ast +import json +import re +from typing import TYPE_CHECKING, Any, Callable, List, Optional, Sequence + +from twinkle_agentic.reward.f1 import _extract_final_answer +from twinkle_agentic.reward.f1 import _f1_score as _f1_score_stemmed + +from .hard_scorer import CheckResult, TrajectoryView + +if TYPE_CHECKING: + from .hard_scorer import CheckFn # noqa: F401 + + +# --------------------------------------------------------------------------- +# shared helpers +# --------------------------------------------------------------------------- +_CODE_FENCE_RE = re.compile(r'```([a-zA-Z0-9_+-]*)\s*\n(.*?)```', re.DOTALL) +_BOXED_RE = re.compile(r'\\boxed\s*\{') +_NUMBER_RE = re.compile(r'[-+]?\d*\.?\d+(?:/\d+)?%?') + + +def _ground_truths(trajectory: dict) -> List[str]: + """Read ground_truth values from user_data (same convention as F1Reward).""" + out: List[str] = [] + for entry in trajectory.get('user_data', []) or []: + if isinstance(entry, (list, tuple)) and len(entry) == 2 and entry[0] == 'ground_truth': + v = entry[1] + if isinstance(v, str): + try: + v = json.loads(v) + except (json.JSONDecodeError, ValueError): + pass + if isinstance(v, (list, tuple)): + out.extend(str(x) for x in v if x) + elif v: + out.append(str(v)) + return out + + +def _last_code_block(text: str) -> Optional[str]: + matches = _CODE_FENCE_RE.findall(text or '') + if not matches: + return None + return matches[-1][1] + + +def _to_number(token: str) -> Optional[float]: + """Normalize a numeric token: fraction 'a/b', percent 'x%', or decimal.""" + if token is None: + return None + s = str(token).strip().replace(',', '').replace('$', '').replace(' ', '') + if not s: + return None + percent = s.endswith('%') + if percent: + s = s[:-1] + try: + if '/' in s: + num, den = s.split('/', 1) + val = float(num) / float(den) + else: + val = float(s) + except (ValueError, ZeroDivisionError): + return None + return val / 100.0 if percent else val + + +def _numbers_in(text: str) -> List[float]: + vals = [] + for tok in _NUMBER_RE.findall(text or ''): + v = _to_number(tok) + if v is not None: + vals.append(v) + return vals + + +_PUNCT_RE = re.compile(r'[^\w\s]', re.UNICODE) +_ARTICLE_RE = re.compile(r'\b(a|an|the)\b') + + +def _f1_score(prediction: str, gold: str): + """F1/EM with graceful fallback when nltk (used by f1.py stemming) is + unavailable — degrade to a stemmer-free token F1 instead of crashing.""" + try: + return _f1_score_stemmed(prediction, gold) + except ImportError: + pass + from collections import Counter + norm = lambda s: _ARTICLE_RE.sub( # noqa: E731 + ' ', _PUNCT_RE.sub('', (s or '').lower())).split() + p_tok, g_tok = norm(prediction), norm(gold) + if not p_tok or not g_tok: + em = float(p_tok == g_tok) + return em, em + em = float(p_tok == g_tok) + common = Counter(p_tok) & Counter(g_tok) + same = sum(common.values()) + if same == 0: + return 0.0, em + prec, rec = same / len(p_tok), same / len(g_tok) + return 2 * prec * rec / (prec + rec), em + + +# --------------------------------------------------------------------------- +# format / structure +# --------------------------------------------------------------------------- +def check_output_format(fmt: str, *, weight: float = 1.5, critical: bool = True + ) -> 'CheckFn': + """Require a specific output artifact in the final answer. + + Args: + fmt: one of ``'boxed'`` (\\boxed{...}), ``'code'`` (fenced block), + ``'json'`` (a parseable JSON object/array anywhere in the answer). + """ + fmt = fmt.lower() + if fmt not in ('boxed', 'code', 'json'): + raise ValueError("fmt must be 'boxed', 'code' or 'json'") + + def _check(view: TrajectoryView) -> CheckResult: + text = view.last_assistant_text() + if fmt == 'boxed': + ok = bool(_extract_final_answer(text)) or bool(_BOXED_RE.search(text)) + detail = 'boxed present' if ok else 'no \\boxed{}' + elif fmt == 'code': + ok = _last_code_block(text) is not None + detail = 'code block present' if ok else 'no code block' + else: # json + block = _last_code_block(text) or text + ok = _has_parseable_json(block) + detail = 'json parseable' if ok else 'no parseable json' + return CheckResult(f'format_{fmt}', 1.0 if ok else 0.0, weight, + critical=critical, n=1, detail=detail) + + return _check + + +def _has_parseable_json(text: str) -> bool: + s = (text or '').strip() + if not s: + return False + # try whole-string first, then first {...}/[...] span + for candidate in (s, _first_bracket_span(s)): + if not candidate: + continue + try: + json.loads(candidate) + return True + except (json.JSONDecodeError, ValueError): + continue + return False + + +def _first_bracket_span(s: str) -> Optional[str]: + starts = [i for i, c in enumerate(s) if c in '{['] + if not starts: + return None + i = starts[0] + open_c = s[i] + close_c = '}' if open_c == '{' else ']' + depth = 0 + for j in range(i, len(s)): + if s[j] == open_c: + depth += 1 + elif s[j] == close_c: + depth -= 1 + if depth == 0: + return s[i:j + 1] + return None + + +# --------------------------------------------------------------------------- +# numeric equivalence (lightweight, no sympy) +# --------------------------------------------------------------------------- +def check_numeric_equiv(*, tol: float = 1e-6, weight: float = 2.0, + critical: bool = False) -> 'CheckFn': + """Compare the extracted final number(s) against ground_truth numerically. + + Handles fractions / decimals / percentages. For symbolic equivalence, + swap this for a sympy/math-verify CheckFn later. Neutral pass when there + is no numeric ground truth to compare against. + """ + def _check(view: TrajectoryView) -> CheckResult: + golds = _ground_truths(view.trajectory) + gold_nums = [n for g in golds for n in _numbers_in(g)] + if not gold_nums: + return CheckResult('numeric_equiv', 1.0, weight, critical=False, n=0, + detail='no numeric ground truth') + text = view.last_assistant_text() + boxed = _extract_final_answer(text) + pred_nums = _numbers_in(boxed) if boxed else _numbers_in(text) + if not pred_nums: + return CheckResult('numeric_equiv', 0.0, weight, critical=critical, n=1, + detail='no number in answer') + # match if any predicted number equals any gold (last pred preferred) + target = gold_nums[-1] + ok = any(abs(p - target) <= tol + tol * abs(target) for p in pred_nums) + return CheckResult('numeric_equiv', 1.0 if ok else 0.0, weight, + critical=critical, n=1, + detail=f'pred~{pred_nums[-1]} vs gold~{target}') + + return _check + + +# --------------------------------------------------------------------------- +# reference match (reuse f1.py) +# --------------------------------------------------------------------------- +def check_answer_match(*, threshold: float = 0.6, weight: float = 2.0, + critical: bool = False, use_em: bool = False) -> 'CheckFn': + """F1/EM of the extracted answer vs ground_truth (reuses ``_f1_score``). + + Score is the max F1 over gold answers (or EM when ``use_em``); pass/fail is + F1 >= threshold. Neutral pass when there is no ground truth. + """ + def _check(view: TrajectoryView) -> CheckResult: + golds = _ground_truths(view.trajectory) + if not golds: + return CheckResult('answer_match', 1.0, weight, critical=False, n=0, + detail='no ground truth') + text = view.last_assistant_text() + boxed = _extract_final_answer(text) + pred = boxed or text + scored = [_f1_score(pred, g) for g in golds] + best_f1 = max(f for f, _ in scored) + best_em = max(e for _, e in scored) + # Containment fallback: when the answer isn't boxed, a short gold that + # appears verbatim in the answer counts as a hit (robust to preamble + # like "The answer is Paris."). + contained = False + if not boxed: + low = text.lower() + contained = any(g.strip() and g.lower() in low and len(g.split()) <= 6 + for g in golds) + if contained: + best_f1 = max(best_f1, 1.0) + best_em = max(best_em, 1.0) + val = best_em if use_em else best_f1 + return CheckResult('answer_match', val, weight, critical=critical, n=1, + detail=f'f1={best_f1:.2f} em={best_em:.0f}' + + (' contained' if contained else '') + + ('' if val >= threshold else ' 'CheckFn': + """Last fenced code block must parse (Python only, via stdlib ``ast``). + + This validates *syntax* without a sandbox and without running anything. + Non-Python blocks are a neutral pass (we can't cheaply verify them here). + """ + def _check(view: TrajectoryView) -> CheckResult: + text = view.last_assistant_text() + block = _last_code_block(text) + if block is None: + return CheckResult('code_parses', 0.0, weight, critical=critical, n=1, + detail='no code block') + if language.lower() != 'python': + return CheckResult('code_parses', 1.0, weight, critical=False, n=0, + detail=f'{language} not statically checked') + try: + ast.parse(block) + return CheckResult('code_parses', 1.0, weight, critical=critical, n=1, + detail='parses') + except SyntaxError as e: + return CheckResult('code_parses', 0.0, weight, critical=critical, n=1, + detail=f'SyntaxError: {e.msg}') + + return _check + + +# --------------------------------------------------------------------------- +# instruction constraints (IFEval-style, pure code) +# --------------------------------------------------------------------------- +def check_instruction_constraints( + *, + min_words: Optional[int] = None, + max_words: Optional[int] = None, + must_include: Optional[Sequence[str]] = None, + must_exclude: Optional[Sequence[str]] = None, + match_source_language: bool = False, + weight: float = 1.0, + critical: bool = False, +) -> 'CheckFn': + """Verify code-checkable instruction-following constraints on the answer. + + Score is the fraction of active sub-constraints satisfied. + """ + must_include = list(must_include or []) + must_exclude = list(must_exclude or []) + + def _check(view: TrajectoryView) -> CheckResult: + text = view.last_assistant_text() + words = text.split() + n_words = len(words) + checks: List[bool] = [] + notes: List[str] = [] + + if min_words is not None: + ok = n_words >= min_words + checks.append(ok) + if not ok: + notes.append(f'words<{min_words}') + if max_words is not None: + ok = n_words <= max_words + checks.append(ok) + if not ok: + notes.append(f'words>{max_words}') + low = text.lower() + for kw in must_include: + ok = kw.lower() in low + checks.append(ok) + if not ok: + notes.append(f'missing:{kw}') + for kw in must_exclude: + ok = kw.lower() not in low + checks.append(ok) + if not ok: + notes.append(f'forbidden:{kw}') + if match_source_language: + ok = _language_matches(view) + checks.append(ok) + if not ok: + notes.append('lang-mismatch') + + if not checks: + return CheckResult('instruction_constraints', 1.0, weight, + critical=False, n=0, detail='no active constraints') + score = sum(1 for c in checks if c) / len(checks) + return CheckResult('instruction_constraints', score, weight, + critical=critical, n=len(checks), + detail=', '.join(notes) or 'all satisfied') + + return _check + + +def _cjk_ratio(text: str) -> float: + if not text: + return 0.0 + cjk = sum(1 for c in text if '\u4e00' <= c <= '\u9fff' + or '\u3040' <= c <= '\u30ff' + or '\uac00' <= c <= '\ud7a3') + return cjk / len(text) + + +def _language_matches(view: TrajectoryView) -> bool: + """Cheap heuristic: answer's CJK-ness matches the first user message's.""" + user_text = '' + for m in view.messages: + if m.get('role') == 'user': + user_text = view.text_of(m) + break + ans = view.last_assistant_text() + if not user_text or not ans: + return True + return abs(_cjk_ratio(user_text[:400]) - _cjk_ratio(ans[:400])) < 0.3 + + +# --------------------------------------------------------------------------- +# degeneration +# --------------------------------------------------------------------------- +def check_not_degenerate(*, min_chars: int = 1, max_repeat_ratio: float = 0.5, + ngram: int = 8, weight: float = 1.0, + critical: bool = False) -> 'CheckFn': + """Fail on empty / trivially short / highly repetitive final answers.""" + def _check(view: TrajectoryView) -> CheckResult: + text = view.last_assistant_text().strip() + if len(text) < min_chars: + return CheckResult('not_degenerate', 0.0, weight, critical=critical, + n=1, detail='too short/empty') + rep = _repetition_ratio(text, ngram) + ok = rep <= max_repeat_ratio + return CheckResult('not_degenerate', 1.0 if ok else 0.0, weight, + critical=critical, n=1, + detail=f'repeat={rep:.2f}' + ('' if ok else ' >thr')) + + return _check + + +def _repetition_ratio(text: str, ngram: int) -> float: + if _cjk_ratio(text[:500]) > 0.3: + tokens = [c for c in text if not c.isspace()] + else: + tokens = text.split() + if len(tokens) < ngram: + return 0.0 + grams = [tuple(tokens[i:i + ngram]) for i in range(len(tokens) - ngram + 1)] + if not grams: + return 0.0 + return 1.0 - len(set(grams)) / len(grams) + + +# Convenience presets keyed by domain, for use with a router later. +def default_checks_for(domain: str) -> List['CheckFn']: + """Return a reasonable initial check bundle for a domain (no sandbox).""" + domain = (domain or '').lower() + if domain == 'math': + return [check_output_format('boxed', critical=False), + check_numeric_equiv(), + check_not_degenerate()] + if domain == 'code': + return [check_output_format('code', critical=False), + check_code_parses(), + check_not_degenerate()] + if domain in ('factual_qa', 'factual', 'qa'): + return [check_answer_match(), + check_not_degenerate()] + if domain in ('open_qa', 'open', 'writing'): + return [check_instruction_constraints(), + check_not_degenerate()] + return [check_not_degenerate()] diff --git a/src/twinkle_agentic/verifier/hard_scorer.py b/src/twinkle_agentic/verifier/hard_scorer.py new file mode 100644 index 000000000..80c69e315 --- /dev/null +++ b/src/twinkle_agentic/verifier/hard_scorer.py @@ -0,0 +1,409 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Deterministic (LLM-free) hard scorer for agentic trajectories. + +Where :class:`RubricVerifier` judges *soft* quality via an LLM, this scorer +judges *hard* facts with plain code: did the agent call declared tools, were +the arguments valid, did the calls execute, is the OpenAI tool protocol +consistent, did the run terminate cleanly, is there a final answer, and did it +avoid degenerate repetition. These signals are **free** and, crucially, +**un-hackable by the policy** — the policy cannot talk its way past a JSON +parse error or a hallucinated tool name. + +The score is a weighted mean of independent checks, each producing a +``CheckResult`` in ``[0, 1]``. Two aggregation modes: + +- ``mode='mean'`` (default): weighted average of all checks. +- ``mode='gate'``: any *critical* check that scores 0 caps the whole score at + 0 (a strict gatekeeper — one hallucinated tool call fails the segment). + +Checks are pluggable: pass your own callables to extend/override. The public +``__call__`` returns an ``int`` in ``[0, NUM_LEVELS)`` per the +:class:`Verifier` contract; ``score_detail`` returns the full breakdown. +""" +from __future__ import annotations + +import json +import re +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, List, Optional, Tuple + +from .base import Verifier + +# A check takes the parsed trajectory view and returns a CheckResult. +CheckFn = Callable[['TrajectoryView'], 'CheckResult'] + +_ERROR_PREFIX_RE = re.compile(r'^\s*(error|exception|traceback|failed)\b[:\s]', re.IGNORECASE) + + +@dataclass +class CheckResult: + name: str + score: float # in [0, 1] + weight: float + critical: bool # if True and score==0, gate mode caps total at 0 + n: int = 0 # number of items this check evaluated + detail: str = '' + + def __post_init__(self): + self.score = min(1.0, max(0.0, float(self.score))) + + +@dataclass +class HardScoreDetail: + level: int + scalar: float # continuous hard score in [0, 1] + gated: bool # a critical check zeroed the score (gate mode) + checks: List[CheckResult] = field(default_factory=list) + + def as_dict(self) -> Dict[str, Any]: + return { + 'level': self.level, + 'scalar': self.scalar, + 'gated': self.gated, + 'checks': {c.name: {'score': c.score, 'weight': c.weight, + 'critical': c.critical, 'n': c.n, 'detail': c.detail} + for c in self.checks}, + } + + +@dataclass +class _ToolCall: + name: Optional[str] + raw_args: Any + call_id: Optional[str] + msg_index: int + + +class TrajectoryView: + """Parsed, check-friendly view over a trajectory segment. + + Precomputes the message list, the assistant tool_calls, the tool-result + messages and the declared tool schema so individual checks stay cheap and + don't each re-walk the messages. + """ + + def __init__(self, trajectory: dict): + self.trajectory = trajectory or {} + self.messages: List[dict] = list(self.trajectory.get('messages', []) or []) + self.tools: List[dict] = list(self.trajectory.get('tools', []) or []) + + # declared tool names + parameter schemas + self.declared_names: set = set() + self.declared_required: Dict[str, List[str]] = {} + for t in self.tools: + fn = t.get('function') if isinstance(t, dict) else None + if not isinstance(fn, dict): + continue + name = fn.get('name') + if not isinstance(name, str) or not name: + continue + self.declared_names.add(name) + params = fn.get('parameters') + if isinstance(params, dict): + req = params.get('required') + if isinstance(req, list): + self.declared_required[name] = [r for r in req if isinstance(r, str)] + + # assistant tool calls, in order + self.tool_calls: List[_ToolCall] = [] + for i, m in enumerate(self.messages): + if m.get('role') != 'assistant': + continue + for tc in (m.get('tool_calls') or []): + if not isinstance(tc, dict): + continue + fn = tc.get('function') or {} + self.tool_calls.append(_ToolCall( + name=fn.get('name') if isinstance(fn, dict) else None, + raw_args=fn.get('arguments') if isinstance(fn, dict) else None, + call_id=tc.get('id'), + msg_index=i, + )) + + # tool-result messages, indexed by tool_call_id where present + self.tool_msgs_by_id: Dict[str, dict] = {} + self.tool_msgs: List[dict] = [] + for m in self.messages: + if m.get('role') == 'tool': + self.tool_msgs.append(m) + cid = m.get('tool_call_id') + if isinstance(cid, str) and cid: + self.tool_msgs_by_id[cid] = m + + # -- shared helpers reused by checks -- + def parsed_args(self, tc: _ToolCall) -> Optional[dict]: + raw = tc.raw_args + if isinstance(raw, dict): + return raw + if raw is None: + return {} + if isinstance(raw, str): + s = raw.strip() + if not s: + return {} + try: + v = json.loads(s) + return v if isinstance(v, dict) else None + except (json.JSONDecodeError, ValueError): + return None + return None + + def result_for(self, tc: _ToolCall) -> Optional[dict]: + """Find the tool result for a call: prefer id match, else next tool msg.""" + if tc.call_id and tc.call_id in self.tool_msgs_by_id: + return self.tool_msgs_by_id[tc.call_id] + # positional fallback: first tool message after the call's assistant msg + for j in range(tc.msg_index + 1, len(self.messages)): + m = self.messages[j] + role = m.get('role') + if role == 'tool': + return m + if role == 'assistant': + break + return None + + @staticmethod + def text_of(msg: Optional[dict]) -> str: + if not msg: + return '' + content = msg.get('content') + if isinstance(content, list): + return '\n'.join(p.get('text', '') for p in content + if isinstance(p, dict) and p.get('type') == 'text') + return content if isinstance(content, str) else '' + + def last_assistant_text(self) -> str: + for m in reversed(self.messages): + if m.get('role') == 'assistant': + return self.text_of(m) + return '' + + def last_assistant_msg(self) -> Optional[dict]: + for m in reversed(self.messages): + if m.get('role') == 'assistant': + return m + return None + + +# --------------------------------------------------------------------------- +# Individual checks (pure, deterministic) +# --------------------------------------------------------------------------- +def check_args_valid_json(view: TrajectoryView) -> CheckResult: + """Every tool call's arguments must parse as a JSON object.""" + calls = view.tool_calls + if not calls: + return CheckResult('args_valid_json', 1.0, 1.0, critical=True, n=0, + detail='no tool calls') + ok = sum(1 for tc in calls if view.parsed_args(tc) is not None) + return CheckResult('args_valid_json', ok / len(calls), 1.0, critical=True, + n=len(calls), detail=f'{ok}/{len(calls)} valid') + + +def check_tool_declared(view: TrajectoryView) -> CheckResult: + """Called tools must be in the declared tool set (no hallucinated tools).""" + calls = view.tool_calls + if not calls or not view.declared_names: + # can't verify without a declared schema -> neutral pass, non-critical + return CheckResult('tool_declared', 1.0, 1.0, critical=False, n=0, + detail='no tools declared or no calls') + ok = sum(1 for tc in calls if tc.name in view.declared_names) + return CheckResult('tool_declared', ok / len(calls), 1.5, critical=True, + n=len(calls), detail=f'{ok}/{len(calls)} declared') + + +def check_required_args(view: TrajectoryView) -> CheckResult: + """Parsed arguments must contain the schema's required fields.""" + calls = [tc for tc in view.tool_calls if tc.name in view.declared_required] + if not calls: + return CheckResult('required_args', 1.0, 1.0, critical=False, n=0, + detail='no schema-required fields to check') + ok = 0 + for tc in calls: + args = view.parsed_args(tc) + if args is None: + continue + req = view.declared_required.get(tc.name, []) + if all(r in args for r in req): + ok += 1 + return CheckResult('required_args', ok / len(calls), 1.0, critical=False, + n=len(calls), detail=f'{ok}/{len(calls)} complete') + + +def check_tool_executed(view: TrajectoryView) -> CheckResult: + """Each tool call must have a non-empty, non-error result message.""" + calls = view.tool_calls + if not calls: + return CheckResult('tool_executed', 1.0, 1.0, critical=False, n=0, + detail='no tool calls') + ok = 0 + for tc in calls: + res = view.result_for(tc) + text = view.text_of(res).strip() + if text and not _ERROR_PREFIX_RE.match(text): + ok += 1 + return CheckResult('tool_executed', ok / len(calls), 2.0, critical=False, + n=len(calls), detail=f'{ok}/{len(calls)} succeeded') + + +def check_protocol_pairing(view: TrajectoryView) -> CheckResult: + """OpenAI protocol: every tool_call id should have a matching tool msg, and + every tool msg should reference a known call id (when ids are used).""" + calls = view.tool_calls + if not calls: + return CheckResult('protocol_pairing', 1.0, 1.0, critical=False, n=0, + detail='no tool calls') + call_ids = {tc.call_id for tc in calls if tc.call_id} + if not call_ids: + # ids not used in this trace; fall back to counting result coverage + paired = sum(1 for tc in calls if view.result_for(tc) is not None) + return CheckResult('protocol_pairing', paired / len(calls), 1.0, + critical=False, n=len(calls), + detail=f'{paired}/{len(calls)} have a result (no ids)') + matched_calls = sum(1 for cid in call_ids if cid in view.tool_msgs_by_id) + # orphan tool messages referencing unknown ids + orphans = sum(1 for m in view.tool_msgs + if isinstance(m.get('tool_call_id'), str) + and m['tool_call_id'] not in call_ids) + total = len(call_ids) + orphans + score = matched_calls / total if total else 1.0 + return CheckResult('protocol_pairing', score, 1.0, critical=False, + n=len(call_ids), + detail=f'{matched_calls}/{len(call_ids)} paired, {orphans} orphan tool msgs') + + +def check_no_repeated_calls(view: TrajectoryView) -> CheckResult: + """Penalize exact-duplicate (name, args) calls (dead-loop / redundancy).""" + calls = view.tool_calls + if len(calls) < 2: + return CheckResult('no_repeated_calls', 1.0, 1.0, critical=False, + n=len(calls), detail='fewer than 2 calls') + seen: set = set() + dupes = 0 + for tc in calls: + args = view.parsed_args(tc) + key = (tc.name, json.dumps(args, sort_keys=True) if isinstance(args, dict) else str(tc.raw_args)) + if key in seen: + dupes += 1 + else: + seen.add(key) + score = 1.0 - dupes / len(calls) + return CheckResult('no_repeated_calls', score, 1.0, critical=False, + n=len(calls), detail=f'{dupes} duplicate calls') + + +def check_clean_termination(view: TrajectoryView) -> CheckResult: + """The trajectory should end on an assistant answer, not a dangling tool + call or a length-truncated turn.""" + last = view.last_assistant_msg() + if last is None: + return CheckResult('clean_termination', 0.0, 1.0, critical=False, n=1, + detail='no assistant message') + # last message overall should be the assistant answer (no trailing tool call + # left unanswered / no pending tool msg after it) + last_role = view.messages[-1].get('role') if view.messages else None + finish = last.get('finish_reason') + truncated = finish == 'length' + dangling = last_role == 'assistant' and bool(last.get('tool_calls')) + ok = (not truncated) and (not dangling) and (last_role in ('assistant', 'tool')) + detail = [] + if truncated: + detail.append('length-truncated') + if dangling: + detail.append('dangling tool_call') + return CheckResult('clean_termination', 1.0 if ok else 0.0, 1.0, + critical=False, n=1, detail=', '.join(detail) or 'clean') + + +def check_final_answer(view: TrajectoryView) -> CheckResult: + """There must be a non-empty final assistant answer.""" + text = view.last_assistant_text().strip() + return CheckResult('final_answer', 1.0 if text else 0.0, 1.5, + critical=False, n=1, + detail=f'{len(text)} chars' if text else 'empty') + + +DEFAULT_CHECKS: Tuple[CheckFn, ...] = ( + check_args_valid_json, + check_tool_declared, + check_required_args, + check_tool_executed, + check_protocol_pairing, + check_no_repeated_calls, + check_clean_termination, + check_final_answer, +) + + +# --------------------------------------------------------------------------- +# Scorer +# --------------------------------------------------------------------------- +class HardScorer(Verifier): + """Deterministic hard score for an agentic trajectory segment. + + Args: + checks: Ordered check callables. Defaults to :data:`DEFAULT_CHECKS`. + Pass your own to extend or replace. + mode: ``'mean'`` (weighted average) or ``'gate'`` (a critical check + scoring 0 zeroes the total). + weights: Optional ``{check_name: weight}`` overrides. + """ + + def __init__( + self, + checks: Optional[List[CheckFn]] = None, + *, + mode: str = 'mean', + weights: Optional[Dict[str, float]] = None, + gate_threshold: float = 1.0, + ): + if mode not in ('mean', 'gate'): + raise ValueError("mode must be 'mean' or 'gate'") + if not 0.0 <= gate_threshold <= 1.0: + raise ValueError('gate_threshold must be in [0, 1]') + self.checks: List[CheckFn] = list(checks) if checks is not None else list(DEFAULT_CHECKS) + self.mode = mode + self.weights = dict(weights or {}) + # In gate mode, a critical check scoring BELOW this threshold zeroes the + # total. 1.0 = any violation gates (strict); 0.0 = only total failure. + self.gate_threshold = float(gate_threshold) + + def __call__(self, trajectory: dict, **kwargs) -> int: + return self.score_detail(trajectory, **kwargs).level + + def score_detail(self, trajectory: dict, **kwargs) -> HardScoreDetail: + view = TrajectoryView(trajectory) + results: List[CheckResult] = [] + for fn in self.checks: + r = fn(view) + if r.name in self.weights: + r.weight = float(self.weights[r.name]) + results.append(r) + + gated = False + if self.mode == 'gate': + # A critical check that is not fully satisfied gates the segment: + # a single hallucinated tool or JSON parse error fails the whole + # thing, regardless of how many other calls were fine. + for r in results: + if r.critical and r.n > 0 and r.score < self.gate_threshold: + gated = True + break + + if gated: + scalar = 0.0 + else: + num = sum(r.score * r.weight for r in results) + den = sum(r.weight for r in results) + scalar = num / den if den else 0.0 + + return HardScoreDetail( + level=self._to_level(scalar), + scalar=scalar, + gated=gated, + checks=results, + ) + + def _to_level(self, scalar: float) -> int: + scalar = min(1.0, max(0.0, scalar)) + level = int(round(scalar * (self.NUM_LEVELS - 1))) + return min(self.NUM_LEVELS - 1, max(0, level)) diff --git a/src/twinkle_agentic/verifier/rubric_verifier.py b/src/twinkle_agentic/verifier/rubric_verifier.py new file mode 100644 index 000000000..215e32acb --- /dev/null +++ b/src/twinkle_agentic/verifier/rubric_verifier.py @@ -0,0 +1,586 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Rubric-based verifier for a single (pre-segmented) trajectory segment. + +Design follows the OpenRubrics -> RubricARROW line of work, adapted to this +repo's progressive-distillation setup: + +1. **Two LLM stages, both distilled via ``llm_backup``** + - *Rubric generation*: given the segment, produce a small set of scoring + criteria, each tagged ``[Hard Rule]`` or ``[Principle]``. + - *Rubric scoring*: given the segment + rubric, emit a per-criterion + verdict. Scores are aggregated ARROW-style into one pointwise scalar. + +2. **Code-level hard verification does NOT go through the LLM.** + Tool-call success/failure, argument JSON validity and call formatting are + checked deterministically (free + un-hackable) and blended in as a + "gatekeeper" floor on the final score. + +3. **Cost-aware scoring**: a single scoring pass yields a soft margin. Only + when the judge is uncertain (|margin| small) do we escalate to majority + voting, aggregating with a median / trimmed-mean (robust to outliers). + +The public ``__call__`` returns an ``int`` in ``[0, NUM_LEVELS)`` per the +:class:`Verifier` contract. ``score_detail`` exposes the continuous score and +breakdown for callers that want the raw signal (e.g. an RL reward). +""" +from __future__ import annotations + +import json +import re +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Tuple + +from twinkle_agentic.utils.llm_backup import llm_backup + +from .base import Verifier + +if TYPE_CHECKING: + from twinkle.data_format import SamplingParams # noqa: F401 + from twinkle.sampler.base import Sampler # noqa: F401 + + +# --------------------------------------------------------------------------- +# Prompts +# --------------------------------------------------------------------------- +_GEN_SYSTEM = """\ +You write evaluation rubrics for a single segment of an AI agent trajectory. \ +The segment may contain reasoning, tool calls and tool results. + +Produce a SHORT list of scoring criteria that discriminate a good segment from \ +a bad one. Each criterion: +- starts with "The response" or "The agent", +- is checkable and non-overlapping (no two criteria testing the same thing), +- ends with a tag: [Hard Rule] for objectively verifiable constraints \ +(tool actually called, argument schema valid, required output present) or \ +[Principle] for softer quality (reasoning soundness, sub-goal progress, no \ +redundant calls). + +Rules: +- Output {min_n}-{max_n} criteria, as FEW as needed to cover the key axes. +- Do NOT reference specific entities/values from THIS segment; keep criteria \ +generalizable to similar segments. +- Output ONLY a numbered list, one criterion per line, nothing else. +""" + +_GEN_USER = """\ +## Task / query (context) +{query} + +## Segment to build a rubric for +{segment} + +Now output the numbered rubric list.""" + +_SCORE_SYSTEM = """\ +You are a strict rubric grader for one segment of an agent trajectory. + +You are given a rubric (numbered criteria, each tagged [Hard Rule] or \ +[Principle]) and the segment. For EACH criterion output one line: + + : PASS or : FAIL + +Judge every criterion independently and literally. A [Hard Rule] fails unless \ +it is unambiguously satisfied. Output only the verdict lines, in order, then \ +stop. Do not add explanations.""" + +_SCORE_USER = """\ +## Task / query (context) +{query} + +## Rubric +{rubric} + +## Segment +{segment} + +Now output one PASS/FAIL line per criterion, in order.""" + + +# --------------------------------------------------------------------------- +# Data holders +# --------------------------------------------------------------------------- +_VERDICT_RE = re.compile(r'^\s*(\d+)\s*[:.)]\s*(pass|fail|true|false|yes|no|1|0)\b', + re.IGNORECASE) + + +@dataclass +class RubricItem: + text: str + is_hard: bool + + +@dataclass +class ScoreDetail: + """Full breakdown behind the final integer level.""" + level: int + scalar: float # continuous pointwise score in [0, 1] + llm_scalar: float # LLM (principle+softhard) component in [0, 1] + hard_pass_rate: float # code-verified hard-rule pass rate in [0, 1] + gated: bool # True if code gatekeeper capped the score + n_votes: int # scoring passes actually spent + rubric: List[RubricItem] = field(default_factory=list) + per_item_pass_rate: List[float] = field(default_factory=list) + + +# --------------------------------------------------------------------------- +# Verifier +# --------------------------------------------------------------------------- +class RubricVerifier(Verifier): + """Score one trajectory segment on a 0..NUM_LEVELS-1 scale via auto rubrics. + + Args: + sampler: Student model sampler (local inference). If ``None`` the + verifier still works but every LLM call is served by the teacher + API through ``llm_backup`` (useful before a student exists). + model_path: Model identifier (bookkeeping only). + sampling_params: Default sampling params for LLM calls. + gen_lora_path: LoRA adapter for the rubric-generator student. + score_lora_path: LoRA adapter for the rubric-scorer student. + min_rubrics / max_rubrics: Target rubric-count window per segment. + hard_weight / principle_weight: Aggregation weights (ARROW uses 3 / 1). + margin_threshold: |margin| below which we escalate to voting. + max_votes: Voting cap for uncertain segments (odd recommended). + gate_floor_ratio: If code-verified hard rules fail, the final scalar is + capped at ``hard_pass_rate`` (gatekeeper). Set to 1.0 to hard-cap, + 0.0 to disable gating. + """ + + def __init__( + self, + sampler: Optional['Sampler'] = None, + *, + model_path: str = '', + sampling_params: Optional['SamplingParams'] = None, + gen_lora_path: Optional[str] = None, + score_lora_path: Optional[str] = None, + min_rubrics: int = 5, + max_rubrics: int = 8, + hard_weight: float = 3.0, + principle_weight: float = 1.0, + margin_threshold: float = 0.25, + max_votes: int = 5, + gate: bool = True, + ): + if max_rubrics < min_rubrics: + raise ValueError('max_rubrics must be >= min_rubrics') + if min_rubrics < 1: + raise ValueError('min_rubrics must be >= 1') + if hard_weight <= 0 or principle_weight <= 0: + raise ValueError('weights must be > 0') + if not 0.0 <= margin_threshold <= 1.0: + raise ValueError('margin_threshold must be in [0, 1]') + if max_votes < 1: + raise ValueError('max_votes must be >= 1') + + self.sampler = sampler + self.model_path = model_path + self.sampling_params = sampling_params + self.gen_lora_path = gen_lora_path or None + self.score_lora_path = score_lora_path or None + self.min_rubrics = int(min_rubrics) + self.max_rubrics = int(max_rubrics) + self.hard_weight = float(hard_weight) + self.principle_weight = float(principle_weight) + self.margin_threshold = float(margin_threshold) + self.max_votes = int(max_votes) + self.gate = bool(gate) + + # ------------------------------------------------------------------ + # public entry points + # ------------------------------------------------------------------ + def __call__(self, trajectory: dict, **kwargs) -> int: + return self.score_detail(trajectory, **kwargs).level + + def score_detail(self, trajectory: dict, *, query: Optional[str] = None, + sampling_params: Any = None) -> ScoreDetail: + query = query or self._infer_query(trajectory) + segment_text = self._render_segment(trajectory) + + # --- code-level hard verification (free, un-hackable) --- + hard_pass_rate, has_hard = self._code_hard_checks(trajectory) + + # --- stage 1: rubric generation (distilled) --- + raw_rubric = self._gen_rubric( + trajectory=self._gen_trajectory(query, segment_text), + sampling_params=self._gen_sampling_params(sampling_params), + query=query, + ) + rubric = self._parse_rubric(raw_rubric) + if not rubric: + # No usable rubric: fall back to the code signal alone. + scalar = hard_pass_rate if has_hard else 0.0 + return ScoreDetail( + level=self._to_level(scalar), scalar=scalar, llm_scalar=0.0, + hard_pass_rate=hard_pass_rate if has_hard else 1.0, + gated=False, n_votes=0, rubric=[], + ) + + # --- stage 2: rubric scoring with margin-adaptive voting --- + per_item_rate, n_votes = self._score_with_voting( + query, segment_text, rubric, sampling_params) + + llm_scalar = self._aggregate(rubric, per_item_rate) + + # --- gatekeeper: code-verified hard failures cap the score --- + scalar = llm_scalar + gated = False + if self.gate and has_hard and hard_pass_rate < 1.0: + capped = min(llm_scalar, hard_pass_rate) + gated = capped < llm_scalar + scalar = capped + + return ScoreDetail( + level=self._to_level(scalar), + scalar=scalar, + llm_scalar=llm_scalar, + hard_pass_rate=hard_pass_rate if has_hard else 1.0, + gated=gated, + n_votes=n_votes, + rubric=rubric, + per_item_pass_rate=per_item_rate, + ) + + # ------------------------------------------------------------------ + # stage 1: rubric generation (student, distilled via llm_backup) + # ------------------------------------------------------------------ + @llm_backup(key_params=['query'], comparator=lambda a, b: _rubric_similar(a, b)) + def _gen_rubric(self, trajectory, sampling_params, query: str = None) -> str: + return self._sample_text(trajectory, sampling_params, self.gen_lora_path) + + # ------------------------------------------------------------------ + # stage 2: rubric scoring (student, distilled via llm_backup) + # ------------------------------------------------------------------ + # Note: the scoring pass is distilled on the *verdict pattern*; the + # comparator matches on binned pass-rate so student/teacher agree when + # their PASS/FAIL vectors are close (not byte-identical). + @llm_backup(key_params=['query', 'rubric_key'], + comparator=lambda a, b: _verdicts_close(a, b)) + def _score_once(self, trajectory, sampling_params, query: str = None, + rubric_key: str = '') -> str: + return self._sample_text(trajectory, sampling_params, self.score_lora_path) + + def _score_with_voting(self, query, segment_text, rubric, sampling_params + ) -> Tuple[List[float], int]: + n = len(rubric) + rubric_block = self._render_rubric(rubric) + rubric_key = _short_hash(rubric_block) + score_traj = self._score_trajectory(query, rubric_block, segment_text) + + # First (cheap) pass. + votes: List[List[bool]] = [] + first = self._score_once( + trajectory=score_traj, + sampling_params=self._score_sampling_params(sampling_params, temperature=0.0), + query=query, rubric_key=rubric_key) + votes.append(self._parse_verdicts(first, n)) + + # Decide whether to escalate: uncertainty = closeness of pass-rate to 0.5. + rate = self._vote_rates(votes) + if self._is_uncertain(rate) and self.max_votes > 1: + sp = self._score_sampling_params(sampling_params, temperature=0.7) + # Escalate up to max_votes; early-stop once verdicts stabilize. + while len(votes) < self.max_votes: + extra = self._score_once( + trajectory=score_traj, sampling_params=sp, + query=query, rubric_key=rubric_key) + votes.append(self._parse_verdicts(extra, n)) + rate = self._vote_rates(votes) + if not self._is_uncertain(rate): + break + return rate, len(votes) + + # ------------------------------------------------------------------ + # code-level hard verification (no LLM) + # ------------------------------------------------------------------ + @staticmethod + def _code_hard_checks(trajectory: dict) -> Tuple[float, bool]: + """Deterministic tool-call checks -> (pass_rate, has_any_hard_signal). + + Each assistant tool_call contributes checks: + - arguments parse as JSON (schema-ish validity), + - a following tool message exists and is non-empty / non-ERROR. + Returns pass_rate over all such checks; has_hard=False when the + segment has no tool calls (nothing to code-verify). + """ + msgs = trajectory.get('messages', []) or [] + n_msgs = len(msgs) + checks: List[bool] = [] + for i, m in enumerate(msgs): + if m.get('role') != 'assistant': + continue + tool_calls = m.get('tool_calls') or [] + for tc in tool_calls: + fn = (tc.get('function') or {}) if isinstance(tc, dict) else {} + args = fn.get('arguments', '') + # 1) argument validity + checks.append(_is_valid_json_args(args)) + # 2) execution success: find the matching/following tool message + ok = False + j = i + 1 + while j < n_msgs and msgs[j].get('role') == 'tool': + content = msgs[j].get('content') or '' + text = content if isinstance(content, str) else str(content) + if text.strip() and not text.lstrip().startswith('ERROR'): + ok = True + break + j += 1 + checks.append(ok) + if not checks: + return 1.0, False + return sum(1 for c in checks if c) / len(checks), True + + # ------------------------------------------------------------------ + # aggregation & mapping + # ------------------------------------------------------------------ + def _aggregate(self, rubric: List[RubricItem], per_item_rate: List[float]) -> float: + """Weighted mean of per-criterion pass-rates (ARROW-style, Hard>Principle).""" + num = 0.0 + den = 0.0 + for item, rate in zip(rubric, per_item_rate): + w = self.hard_weight if item.is_hard else self.principle_weight + num += w * rate + den += w + return num / den if den else 0.0 + + def _to_level(self, scalar: float) -> int: + scalar = min(1.0, max(0.0, scalar)) + # Map [0,1] onto {0..NUM_LEVELS-1} with even-width bins. + level = int(round(scalar * (self.NUM_LEVELS - 1))) + return min(self.NUM_LEVELS - 1, max(0, level)) + + def _is_uncertain(self, per_item_rate: Sequence[float]) -> bool: + """A segment is uncertain if any criterion sits near the 0.5 boundary.""" + if not per_item_rate: + return False + # distance of the aggregate margin from a confident 0/1 verdict + for r in per_item_rate: + if abs(r - 0.5) * 2.0 < self.margin_threshold: + return True + return False + + @staticmethod + def _vote_rates(votes: List[List[bool]]) -> List[float]: + """Per-criterion PASS rate across votes (robust: median-of-means style). + + For each criterion we average the boolean verdicts across votes. This + equals majority-vote direction while keeping a soft rate for + aggregation. Outlier passes/fails wash out as votes accumulate. + """ + if not votes: + return [] + n = max(len(v) for v in votes) + rates: List[float] = [] + for k in range(n): + col = [v[k] for v in votes if k < len(v)] + if not col: + rates.append(0.0) + continue + rates.append(sum(1 for c in col if c) / len(col)) + return rates + + # ------------------------------------------------------------------ + # LLM sampling plumbing (mirrors Summarizer) + # ------------------------------------------------------------------ + def _sample_text(self, trajectory, sampling_params, lora_path) -> str: + if self.sampler is None: + return '' + sample_kwargs: dict[str, Any] = {'sampling_params': sampling_params} + if lora_path is None: + sample_kwargs['use_base_model'] = True + else: + sample_kwargs['adapter_path'] = lora_path + responses = self.sampler.sample([trajectory], **sample_kwargs) + resp = list(responses)[0] if responses else None + if resp is None: + return '' + seqs = getattr(resp, 'sequences', None) or [] + return (getattr(seqs[0], 'decoded', None) or '') if seqs else '' + + def _gen_trajectory(self, query: str, segment_text: str) -> dict: + user = _fill(_GEN_USER, query=query, segment=segment_text) + system = _fill(_GEN_SYSTEM, min_n=self.min_rubrics, max_n=self.max_rubrics) + return {'messages': [ + {'role': 'system', 'content': system}, + {'role': 'user', 'content': user}, + ]} + + def _score_trajectory(self, query: str, rubric_block: str, segment_text: str) -> dict: + user = _fill(_SCORE_USER, query=query, rubric=rubric_block, segment=segment_text) + return {'messages': [ + {'role': 'system', 'content': _SCORE_SYSTEM}, + {'role': 'user', 'content': user}, + ]} + + def _gen_sampling_params(self, override): + if override is not None: + return override + if self.sampling_params is not None: + return self.sampling_params + from twinkle.data_format.sampling import SamplingParams + return SamplingParams(temperature=0.3, max_tokens=512) + + def _score_sampling_params(self, override, *, temperature: float): + if override is not None: + return override + from twinkle.data_format.sampling import SamplingParams + return SamplingParams(temperature=temperature, max_tokens=256) + + # ------------------------------------------------------------------ + # rendering / parsing helpers + # ------------------------------------------------------------------ + @staticmethod + def _infer_query(trajectory: dict) -> str: + for msg in trajectory.get('messages', []) or []: + if msg.get('role') == 'user': + c = msg.get('content') + if isinstance(c, str) and c.strip(): + return c.strip() + return '(no explicit query)' + + @staticmethod + def _render_segment(trajectory: dict) -> str: + """Flatten a segment's messages into a readable transcript.""" + lines: List[str] = [] + for m in trajectory.get('messages', []) or []: + role = m.get('role', '?') + if role == 'system': + continue + content = m.get('content') + if isinstance(content, list): + content = '\n'.join( + p.get('text', '') for p in content + if isinstance(p, dict) and p.get('type') == 'text') + content = content or '' + tool_calls = m.get('tool_calls') or [] + if tool_calls: + calls = '; '.join( + f"{(tc.get('function') or {}).get('name', '?')}" + f"({(tc.get('function') or {}).get('arguments', '')})" + for tc in tool_calls if isinstance(tc, dict)) + lines.append(f'[{role}] {content}\n tool_calls: {calls}'.rstrip()) + else: + lines.append(f'[{role}] {content}'.rstrip()) + return '\n'.join(lines).strip() + + _TAG_HARD_RE = re.compile(r'\[\s*hard\s*rule\s*\]', re.IGNORECASE) + _TAG_PRIN_RE = re.compile(r'\[\s*principle\s*\]', re.IGNORECASE) + _NUM_LINE_RE = re.compile(r'^\s*(?:\d+[.)]|[-*])\s*(.+?)\s*$') + + @classmethod + def _parse_rubric(cls, raw: str) -> List[RubricItem]: + items: List[RubricItem] = [] + for line in (raw or '').splitlines(): + m = cls._NUM_LINE_RE.match(line) + text = (m.group(1) if m else line).strip() + if not text: + continue + is_hard = bool(cls._TAG_HARD_RE.search(text)) + is_prin = bool(cls._TAG_PRIN_RE.search(text)) + if not (is_hard or is_prin): + # Untagged line that isn't clearly a criterion -> skip noise. + if not m: + continue + is_hard = False # default to principle + clean = cls._TAG_HARD_RE.sub('', cls._TAG_PRIN_RE.sub('', text)).strip(' .') + if clean: + items.append(RubricItem(text=clean, is_hard=is_hard)) + return items + + @staticmethod + def _render_rubric(rubric: List[RubricItem]) -> str: + return '\n'.join( + f'{i + 1}. {it.text} [{"Hard Rule" if it.is_hard else "Principle"}]' + for i, it in enumerate(rubric)) + + @staticmethod + def _parse_verdicts(raw: str, n: int) -> List[bool]: + """Parse ': PASS/FAIL' lines into a length-n boolean vector. + + Missing verdicts default to FAIL (conservative for hard rules). + """ + verdicts = [False] * n + for line in (raw or '').splitlines(): + m = _VERDICT_RE.match(line) + if not m: + continue + idx = int(m.group(1)) - 1 + if 0 <= idx < n: + verdicts[idx] = m.group(2).lower() in ('pass', 'true', 'yes', '1') + return verdicts + + +# --------------------------------------------------------------------------- +# module-level helpers (comparators etc.) +# --------------------------------------------------------------------------- +def _fill(template: str, **kw) -> str: + out = template + for k, v in kw.items(): + out = out.replace('{' + k + '}', str(v)) + return out + + +def _is_valid_json_args(args: Any) -> bool: + if isinstance(args, dict): + return True + if not isinstance(args, str): + return False + s = args.strip() + if not s: + return True # a no-arg call is valid + try: + json.loads(s) + return True + except (json.JSONDecodeError, ValueError): + return False + + +def _short_hash(text: str) -> str: + import hashlib + return hashlib.md5((text or '').encode()).hexdigest()[:12] + + +_TAG_ANY_RE = re.compile(r'\[\s*(hard\s*rule|principle)\s*\]', re.IGNORECASE) + + +def _rubric_similar(a: str, b: str) -> bool: + """Comparator for rubric generation: rubrics rarely match verbatim, so we + compare on *shape* — similar criterion count and similar hard/principle mix. + This is a proxy; downstream consistency filtering is the real quality gate. + """ + ca = _TAG_ANY_RE.findall(a or '') + cb = _TAG_ANY_RE.findall(b or '') + na, nb = len(ca), len(cb) + if na == 0 and nb == 0: + return True + if na == 0 or nb == 0: + return False + # count within +/-2 and hard-ratio within 0.34 + if abs(na - nb) > 2: + return False + hard_a = sum(1 for t in ca if t.lower().startswith('hard')) / na + hard_b = sum(1 for t in cb if t.lower().startswith('hard')) / nb + return abs(hard_a - hard_b) <= 0.34 + + +def _parse_rate(raw: str) -> Optional[float]: + pos = 0 + total = 0 + for line in (raw or '').splitlines(): + m = _VERDICT_RE.match(line) + if not m: + continue + total += 1 + if m.group(2).lower() in ('pass', 'true', 'yes', '1'): + pos += 1 + if total == 0: + return None + return pos / total + + +def _verdicts_close(a: str, b: str, tol: float = 0.25) -> bool: + """Comparator for rubric scoring: student/teacher agree when their overall + PASS rate is within ``tol`` (binned agreement, not byte-identical text).""" + ra, rb = _parse_rate(a), _parse_rate(b) + if ra is None or rb is None: + return (a or '').strip() == (b or '').strip() + return abs(ra - rb) <= tol From 1a0de3ac5811198c9afed0b47da2fbdfa44627ed Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Wed, 8 Jul 2026 20:27:46 +0800 Subject: [PATCH 03/60] wip --- cookbook/exp/cold_start/train_cold_start.py | 3 +- .../exp/data_pipeline/process_and_save.py | 246 ++++++++ src/twinkle/preprocessor/__init__.py | 2 +- src/twinkle/preprocessor/base.py | 54 +- src/twinkle_agentic/memory/DESIGN.md | 564 ++++++++++++++++++ src/twinkle_agentic/preprocessor/AUDIT.md | 179 ++++++ src/twinkle_agentic/preprocessor/__init__.py | 8 +- .../preprocessor/dead_loop_filter.py | 14 +- .../preprocessor/experimental/__init__.py | 21 + .../{ => experimental}/llm_backend.py | 0 .../{ => experimental}/score_filter.py | 2 +- .../preprocessor/hard_filter.py | 36 +- .../preprocessor/intent_classifier.py | 25 +- src/twinkle_agentic/preprocessor/intents.py | 25 + .../preprocessor/label_schema.py | 102 ++++ .../preprocessor/language_filter.py | 109 ++++ .../preprocessor/logprob_utils.py | 231 +++++++ .../preprocessor/message_normalizer.py | 12 +- .../preprocessor/message_sanity.py | 18 + .../preprocessor/message_utils.py | 137 +++++ .../preprocessor/model_filter.py | 17 +- .../preprocessor/offline/__init__.py | 24 + .../preprocessor/offline/decontaminate.py | 110 ++++ .../preprocessor/offline/near_dedup.py | 164 +++++ .../preprocessor/outcome_filter.py | 80 +++ .../preprocessor/pii_presidio_filter.py | 105 +++- .../preprocessor/provenance.py | 73 +++ .../preprocessor/refuse_filter.py | 89 ++- .../preprocessor/safety_scorer.py | 89 +++ .../preprocessor/structural_noise.py | 61 ++ .../preprocessor/trajectory_scorer.py | 243 ++++++++ src/twinkle_agentic/preprocessor/utils.py | 387 +----------- src/twinkle_agentic/segment/__init__.py | 4 + src/twinkle_agentic/segment/base.py | 237 ++++++++ src/twinkle_agentic/segment/llm_segmenter.py | 308 ++++++++++ src/twinkle_agentic/summarizer/__init__.py | 3 +- .../summarizer/action_summarizer.py | 86 +++ src/twinkle_agentic/utils/llm_backup.py | 8 +- src/twinkle_agentic/verifier/__init__.py | 7 + src/twinkle_agentic/verifier/aggregation.py | 262 ++++++++ .../verifier/rubric_verifier.py | 53 +- 41 files changed, 3755 insertions(+), 443 deletions(-) create mode 100644 cookbook/exp/data_pipeline/process_and_save.py create mode 100644 src/twinkle_agentic/memory/DESIGN.md create mode 100644 src/twinkle_agentic/preprocessor/AUDIT.md create mode 100644 src/twinkle_agentic/preprocessor/experimental/__init__.py rename src/twinkle_agentic/preprocessor/{ => experimental}/llm_backend.py (100%) rename src/twinkle_agentic/preprocessor/{ => experimental}/score_filter.py (99%) create mode 100644 src/twinkle_agentic/preprocessor/intents.py create mode 100644 src/twinkle_agentic/preprocessor/label_schema.py create mode 100644 src/twinkle_agentic/preprocessor/language_filter.py create mode 100644 src/twinkle_agentic/preprocessor/logprob_utils.py create mode 100644 src/twinkle_agentic/preprocessor/message_utils.py create mode 100644 src/twinkle_agentic/preprocessor/offline/__init__.py create mode 100644 src/twinkle_agentic/preprocessor/offline/decontaminate.py create mode 100644 src/twinkle_agentic/preprocessor/offline/near_dedup.py create mode 100644 src/twinkle_agentic/preprocessor/outcome_filter.py create mode 100644 src/twinkle_agentic/preprocessor/provenance.py create mode 100644 src/twinkle_agentic/preprocessor/safety_scorer.py create mode 100644 src/twinkle_agentic/preprocessor/structural_noise.py create mode 100644 src/twinkle_agentic/preprocessor/trajectory_scorer.py create mode 100644 src/twinkle_agentic/segment/__init__.py create mode 100644 src/twinkle_agentic/segment/base.py create mode 100644 src/twinkle_agentic/segment/llm_segmenter.py create mode 100644 src/twinkle_agentic/summarizer/action_summarizer.py create mode 100644 src/twinkle_agentic/verifier/aggregation.py diff --git a/cookbook/exp/cold_start/train_cold_start.py b/cookbook/exp/cold_start/train_cold_start.py index da7149bba..9c662e89d 100644 --- a/cookbook/exp/cold_start/train_cold_start.py +++ b/cookbook/exp/cold_start/train_cold_start.py @@ -13,11 +13,12 @@ from twinkle.dataset.base import DatasetMeta from twinkle.model import MegatronModel from twinkle_agentic.preprocessor import ( - QualityPreprocessor, SamplerBackend, + QualityPreprocessor, IntentClassifier, HardFilter, RefuseFilter, DeadLoopFilter, TokenSoupFilter, MessageSanityFilter, SpecialCharsFilter, ModelFilter, DedupFilter, MessageNormalizer, ) +from twinkle_agentic.preprocessor.experimental import SamplerBackend # noqa: F401 logger = get_logger() diff --git a/cookbook/exp/data_pipeline/process_and_save.py b/cookbook/exp/data_pipeline/process_and_save.py new file mode 100644 index 000000000..636c72e60 --- /dev/null +++ b/cookbook/exp/data_pipeline/process_and_save.py @@ -0,0 +1,246 @@ +"""Standalone dataset-processing demo: run the full agentic QualityPreprocessor +over a slice of the raw OpenClaw CSV, then persist the cleaned + scored + tagged +trajectories with ``Dataset.save_as`` for inspection. + +This is NOT a training script — no template/encode/pack. The point is to see how +the preprocessor behaves end-to-end and to keep ALL enrichment: + +- **tags / scores in ``user_data``** (AUDIT A5 envelope): per-round & trajectory + scores (``TrajectoryScorer``), safety (``SafetyScorer``), intent key-rounds + (``IntentClassifier``), structural-noise ratio (``StructuralNoiseTagger``), + and provenance/lineage (``ProvenanceStamp``). +- **important fields preserved**: ``id``/``source``/``model_id``/``messages`` + are never dropped; ``MessageNormalizer`` now passes through + ``reasoning_content``/``thinking`` (AUDIT P3). + +Pipeline follows the tag-then-filter architecture: mappers annotate, a final +read-only ``TrajectoryOutcomeFilter`` drops on the tags (no DAG, linear order). + +Run: + CSV_PATH=/mnt/data/yzhao/tastelikefeet/bc/20260531.csv \ + DATASET_TOTAL=200 python cookbook/exp/data_pipeline/process_and_save.py +""" +import json +import os +from functools import partial +from pathlib import Path +from typing import Any, Dict, Iterator, List + +from twinkle.dataset import Dataset +from twinkle.dataset.base import DatasetMeta +from twinkle.utils import get_logger +from twinkle_agentic.preprocessor import (DeadLoopFilter, HardFilter, + IntentClassifier, LanguageFilter, + MessageNormalizer, + MessageSanityFilter, ModelFilter, + ProvenanceStamp, QualityPreprocessor, + RefuseFilter, SafetyScorer, + SpecialCharsFilter, + StructuralNoiseTagger, + TokenSoupFilter, + TrajectoryOutcomeFilter, + TrajectoryScorer) +from twinkle_agentic.preprocessor import label_schema as L + +logger = get_logger() + +# ── Config ──────────────────────────────────────────────────────────────────── +CSV_PATH = os.environ.get('CSV_PATH', '/mnt/data/yzhao/tastelikefeet/bc/20260531.csv') +DATASET_TOTAL = int(os.environ.get('DATASET_TOTAL', 200)) +MAP_NUM_PROC = int(os.environ.get('MAP_NUM_PROC', 8)) +OUTPUT_DIR = os.environ.get('OUTPUT_DIR', './output/data_pipeline') +OUTPUT_PATH = os.path.join(OUTPUT_DIR, 'processed_20260531_200.jsonl') +DROPPED_PATH = os.path.join(OUTPUT_DIR, 'dropped.jsonl') +PIPELINE_VERSION = os.environ.get('PIPELINE_VERSION', 'audit-v1') +# Set to keep only trajectories above this fused score. None -> keep all (inspect scores only). +MIN_TRAJ_SCORE = os.environ.get('MIN_TRAJ_SCORE') +MIN_TRAJ_SCORE = float(MIN_TRAJ_SCORE) if MIN_TRAJ_SCORE else None +# TrajectoryScorer runs hard-only by default (deterministic, fast). Set USE_RUBRIC=1 +# to attach a RubricVerifier so per-segment scores get real LLM semantic signal +# (needs LLM_BACKUP_* / OPENAI_API_KEY; much slower). Without it every clean +# trajectory tends to collapse to level 4 because only hard checks discriminate. +USE_RUBRIC = os.environ.get('USE_RUBRIC', '') not in ('', '0', 'false', 'False') + + +# ── CSV ingestion (custom format: `ts,model,req_id,messages_json`) ───────────── +def _canonicalize_tool_call(tc: Any) -> Dict[str, Any]: + """Coerce a raw tool_call into a fixed-schema dict for stable Arrow inference.""" + tc = tc if isinstance(tc, dict) else {} + fn = tc.get('function') if isinstance(tc.get('function'), dict) else {} + args = fn.get('arguments') + if isinstance(args, dict): + args_str = json.dumps(args, ensure_ascii=False) + elif isinstance(args, str) and args.strip(): + try: + decoded = json.loads(args) + except json.JSONDecodeError: + decoded = {} + args_str = json.dumps(decoded if isinstance(decoded, dict) else {}, ensure_ascii=False) + else: + args_str = '{}' + return { + 'id': str(tc.get('id') or ''), + 'type': str(tc.get('type') or 'function'), + 'function': {'name': str(fn.get('name') or ''), 'arguments': args_str}, + } + + +def _stream_csv_rows(csv_path: str, max_rows: int = 0) -> Iterator[Dict[str, Any]]: + """Stream the custom CSV. First 3 fields are scalar; the rest of the line is a + JSON array of chat messages (may contain commas) — split on the first 3 commas. + + ``reasoning_content`` is folded into a ``...`` prefix so it + survives as visible content and is later re-exposed by MessageNormalizer (P3). + """ + emitted = 0 + with open(csv_path, 'rb') as f: + for raw in f: + try: + line = raw.decode('utf-8').rstrip('\n').rstrip('\r') + except UnicodeDecodeError: + continue + if not line: + continue + parts = line.split(',', 3) + if len(parts) < 4: + continue + ts, model, req_id, msgs_raw = parts + try: + raw_msgs = json.loads(msgs_raw) + except json.JSONDecodeError: + continue + messages: List[Dict[str, Any]] = [] + for m in raw_msgs: + role = m.get('role', '') + content = m.get('content') + if isinstance(content, list): + content = ''.join(p.get('text', '') for p in content + if isinstance(p, dict) and p.get('type') == 'text') + if content is None: + content = '' + if not isinstance(content, str): + continue + raw_tcs = m.get('tool_calls') if role == 'assistant' else None + tc_list = [_canonicalize_tool_call(tc) for tc in raw_tcs] if raw_tcs else [] + if role == 'assistant': + if not content and not tc_list: + continue + if m.get('reasoning_content'): + content = f"{m['reasoning_content']}{content}" + elif role != 'tool' and not content: + continue + messages.append({ + 'role': role, + 'content': content, + 'tool_calls': json.dumps(tc_list, ensure_ascii=False) if tc_list else '', + 'tool_call_id': str(m.get('tool_call_id') or '') if role == 'tool' else '', + }) + if not messages: + continue + yield { + 'id': f'csv__{ts}__{req_id}', + 'source': Path(csv_path).stem, + 'model_id': model, + 'messages': messages, + 'user_data': [], + } + emitted += 1 + if max_rows and emitted >= max_rows: + break + + +def _build_trajectory_scorer() -> TrajectoryScorer: + """Hard-only by default; attach an LLM RubricVerifier when USE_RUBRIC is set.""" + rubric_verifier = None + if USE_RUBRIC: + from twinkle_agentic.verifier import RubricVerifier + # No sampler -> scores via the llm_backup teacher path (LLM_BACKUP_*/OPENAI_*). + rubric_verifier = RubricVerifier() + return TrajectoryScorer(rubric_verifier=rubric_verifier) + + +def build_pipeline() -> QualityPreprocessor: + """Full tag-then-filter pipeline. Mappers annotate; the tail filter drops on tags.""" + return QualityPreprocessor( + pipeline=[ + # 0) lineage first, so even dropped rows carry provenance in the log. + ProvenanceStamp(source=Path(CSV_PATH).stem, pipeline_version=PIPELINE_VERSION), + # 1) canonicalize message schema (heartbeat strip, tool-call normalize, + # reasoning passthrough — P3), then structural / content filters. + MessageNormalizer(), + ModelFilter(), + LanguageFilter(allowed=('en', 'zh')), + # Shallow-chat round cap is 40; agent traces get a far higher ceiling + # (agent_max_rounds) so long tool-calling loops — the highest-value + # distillation data — are not clipped. + HardFilter(min_user_chars_cjk=14, min_user_chars=24, max_rounds=40, + agent_max_rounds=200), + RefuseFilter(), + DeadLoopFilter(), + MessageSanityFilter(), + SpecialCharsFilter(max_ratio=0.6), + TokenSoupFilter(max_chars=8000), + # 2) taggers (never drop): intent key-rounds, structural noise ratio, + # per-round + trajectory scores, safety score. All write user_data. + IntentClassifier(), + StructuralNoiseTagger(), + _build_trajectory_scorer(), # hard-only, or LLM rubric when USE_RUBRIC=1 + SafetyScorer(), # fixed safety rubric; no sampler -> neutral score, still tagged + # 3) read-only outcome filter: drops on the tags above (D6). Enabled + # only when MIN_TRAJ_SCORE is set, else we keep everything to inspect. + TrajectoryOutcomeFilter( + min_traj_score=MIN_TRAJ_SCORE if MIN_TRAJ_SCORE is not None else 0.0, + min_safety_score=None, + drop_unsafe_flag=False, + ), + ], + dropped_log_path=DROPPED_PATH, + ) + + +def _print_sample(dataset: Dataset, n: int = 3) -> None: + """Show the enrichment kept on a few rows so you can eyeball tags/scores.""" + hf = dataset.dataset + show = min(n, len(hf)) + logger.info(f'── sample of {show} processed rows (tags/scores in user_data) ──') + for i in range(show): + row = hf[i] + logger.info( + f"[{row.get('id')}] model={row.get('model_id')} n_msgs={len(row.get('messages') or [])}\n" + f" intent = {row.get('intent')}\n" + f" traj_score = {L.get_label(row, L.KEY_TRAJ_SCORE)} " + f"level={L.get_label(row, L.KEY_TRAJ_LEVEL)} " + f"conf={L.get_label(row, L.KEY_TRAJ_CONFIDENCE)}\n" + f" round_scores= {L.get_label(row, L.KEY_ROUND_SCORES)}\n" + f" safety = {L.get_label(row, L.KEY_SAFETY_SCORE)} " + f"(unsafe={L.get_label(row, L.KEY_SAFETY_UNSAFE)})\n" + f" noise_ratio = {L.get_label(row, 'structural_noise_ratio')}\n" + f" provenance = {L.get_label(row, L.KEY_PROVENANCE)}") + + +def main() -> None: + os.makedirs(OUTPUT_DIR, exist_ok=True) + logger.info(f'Loading up to {DATASET_TOTAL} rows from {CSV_PATH}') + + meta = DatasetMeta( + dataset_id=Path(CSV_PATH).stem, + data=partial(_stream_csv_rows, csv_path=CSV_PATH, max_rows=DATASET_TOTAL), + ) + dataset = Dataset(meta) + logger.info(f'Ingested {len(dataset.dataset)} rows.') + + pipeline = build_pipeline() + # Dataset.map runs the QualityPreprocessor over HF batches (batched=True is + # forced internally). num_proc parallelizes across shards. + dataset.map(pipeline, num_proc=MAP_NUM_PROC, load_from_cache_file=False) + + logger.info(f'After pipeline: {len(dataset.dataset)} rows kept.') + _print_sample(dataset) + + dataset.save_as(OUTPUT_PATH, format='jsonl') + logger.info(f'Saved processed dataset -> {OUTPUT_PATH}') + logger.info(f'Dropped rows log -> {DROPPED_PATH}') + + +if __name__ == '__main__': + main() diff --git a/src/twinkle/preprocessor/__init__.py b/src/twinkle/preprocessor/__init__.py index 40d756e3b..49a00b801 100644 --- a/src/twinkle/preprocessor/__init__.py +++ b/src/twinkle/preprocessor/__init__.py @@ -1,5 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -from .base import DataFilter, Preprocessor +from .base import DataFilter, Filter, Mapper, Preprocessor from .dpo import EmojiDPOProcessor from .llm import (AlpacaProcessor, CompetitionMathGRPOProcessor, CompetitionMathProcessor, CountdownProcessor, GSM8KProcessor, SelfCognitionProcessor) diff --git a/src/twinkle/preprocessor/base.py b/src/twinkle/preprocessor/base.py index 0225d3c1e..588b29ae7 100644 --- a/src/twinkle/preprocessor/base.py +++ b/src/twinkle/preprocessor/base.py @@ -1,10 +1,23 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -from typing import Any, Dict, List +from typing import Any, Dict, List, Tuple from twinkle.data_format import Trajectory class Preprocessor: + """Base for pipeline steps. + + Concrete steps take a batch of rows (list-of-dict, or the columnar + dict-of-lists produced by HF ``datasets``) and return a + ``(kept, dropped)`` tuple of row lists. ``map_col_to_row`` normalizes the + input; a step that never removes rows (a *mapper*) returns + ``(rows, [])`` — see :class:`Mapper`. Steps that select rows (a *filter*) + return ``(kept, dropped)`` — see :class:`Filter`. + + The pipeline runner (:class:`~twinkle_agentic.preprocessor.QualityPreprocessor`) + consumes the tuple, logs the dropped rows, and re-columnarizes ``kept`` before + handing it to the next step. + """ @staticmethod def map_col_to_row(rows) -> List[Dict[str, Any]]: @@ -36,8 +49,43 @@ def map_row_to_col(rows, keys: List[str] = None) -> Dict[str, List[Any]]: return columns - def __call__(self, rows: Dict[str, List[Any]]) -> Dict[str, List[Any]]: - ... + def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + """Return ``(kept, dropped)`` row lists. Subclasses must override.""" + raise NotImplementedError + + +class Mapper(Preprocessor): + """A step that annotates/transforms rows and never drops any. + + Subclasses implement :meth:`map` (row-in, row-out); the ``(rows, [])`` + contract is provided so mappers compose with filters in the same pipeline. + """ + + def map(self, row: Dict[str, Any]) -> Dict[str, Any]: + raise NotImplementedError + + def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + rows = self.map_col_to_row(rows) + return [self.map(r) for r in rows], [] + + +class Filter(Preprocessor): + """A step that selects rows, returning ``(kept, dropped)``. + + Subclasses implement :meth:`keep` (row-in, bool-out). Dropped rows are + returned so the runner can log them. + """ + + def keep(self, row: Dict[str, Any]) -> bool: + raise NotImplementedError + + def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + rows = self.map_col_to_row(rows) + kept: List[Dict[str, Any]] = [] + dropped: List[Dict[str, Any]] = [] + for r in rows: + (kept if self.keep(r) else dropped).append(r) + return kept, dropped class DataFilter: diff --git a/src/twinkle_agentic/memory/DESIGN.md b/src/twinkle_agentic/memory/DESIGN.md new file mode 100644 index 000000000..6baf5e350 --- /dev/null +++ b/src/twinkle_agentic/memory/DESIGN.md @@ -0,0 +1,564 @@ +# Memory 线设计（自进化蒸馏框架的第二条线） + +> 状态：**设计稿，未实现**。本文件是讨论沉淀，供后续实现参考。 +> 与主线（清洗 + 评分 + 在线蒸馏）共用同一批 trajectory 与同一套 verifier / `user_data` 信封 / `llm_backup` / D7c 校准思想，不另起炉灶。 + +--- + +## 0. 背景与定位：双线 + +同一份生产 trajectory，榨出两种产物，固化到不同位置、不同时间尺度： + +| | 线 A：训模型（慢记忆） | 线 B：总结 memory（快记忆） | +|---|---|---| +| 固化位置 | 模型权重 | 外部 memory store | +| 生效方式 | 蒸馏/微调后永久具备 | 推理时检索注入 context | +| 时间尺度 | 天/周（攒批再训） | 秒/分钟（写完即用） | +| 改什么 | 参数 | 行为（不改参数） | +| 装什么 | 可泛化的**技能/模式** | 易变的**事实/偏好/近期上下文** | + +**line B 的存在理由**：在线蒸馏有延迟（攒数据→训练→部署），空窗期 student 学不到新东西；memory 立刻起作用填这个空窗，等能力被训进权重后再从 memory 退休。 + +**双线共用的唯一“质检车间”**：现有 `preprocessor + verifier`。它产出的 `traj_score / round_scores / confidence / intent / safety` 同时作为两条线的准入闸门，不重复造。 + +### 分工判据（什么进 A，什么进 B） +用 **可泛化性 × 稳定性** 分流： +- 进 A（训权重）：高频、可泛化、稳定 —— 能变成“技能”的。 +- 进 B（memory）：低频 / 易变 / 实体绑定 —— 只能当“事实/上下文”的。训进权重会过拟合具体实体且会过时。 + +### 消费者与作用域（已确认） +- **消费者**：两个都要 —— ① 本地 student 推理时检索注入；② 辅助置信度路由。 +- **作用域**：两层 —— 用户级（个性化：偏好、历史）+ 全局级（跨用户通用事实/模式）。 + +--- + +## 1. 存储与检索选型（基于现有代码，零新增重依赖） + +现有可复用： +- **存储** `twinkle/server/state`：memory / file / redis 三后端，带 TTL、`update_atomic`（原子）、`keys(pattern)`（通配）。 +- **embedding** `preprocessor/experimental/llm_backend.py`：`OpenAIBackend.embeddings`（走 HTTP，dashscope 有 embedding 接口）。 +- **没有**专用向量库（无 Milvus/Qdrant/faiss），但 D1 的 MinHash 说明“近似检索土办法”本项目可接受。 + +**结论（最小可跑，可后续换库）**： +``` +MemoryStore +├── 用户级（结构化 KV/标签）── 复用 StateBackend（FileBackend 默认，可切 Redis） +│ key = mem::user:::::；精确/前缀匹配；零额外依赖 +└── 全局级（向量语义召回）── 条目存 StateBackend；embedding 用 OpenAIBackend.embeddings + 召回 = 应用层 cosine top-k（初期 O(n)，涨了再换 faiss/Qdrant） + 抽象出 VectorIndex 接口，后端可替换 +``` +- 用户级以**结构化 KV** 为主（实体绑定要精确命中，向量会招噪声）。 +- 全局级以**向量召回**为主，再用结构化标签（intent/domain/min_score）过滤。 + +### MemoryItem（统一条目） +```python +@dataclass +class MemoryItem: + id: str + scope: str # 'user:' | 'global' + kind: str # 'fact' | 'preference' | 'action_pattern' | 'anti_pattern' + key: str # 结构化槽位（用户级精确匹配）；全局级可空 + content: str # 供注入 context 的自然语言 + embedding: list # 全局级语义召回用 + # 毕业机制 / 效用 所需元数据（先埋点） + hit_count: int + source_score: float # 来源轨迹 traj_score（写入门槛） + created_at: int + last_hit_at: int + graduated: bool # 是否已训进权重 → A→B 退休标记 +``` + +### MemoryStore 接口 +```python +class MemoryStore: + def write(item): ... + def retrieve(query, scope, *, intent=None, min_score=None, k=5) -> list[MemoryItem]: + # 用户级：KV/标签精确匹配；全局级：向量召回 + 标签过滤 + def mark_hit(item_id): ... # 命中打点（喂效用/晋升） + def retire(item_id): ... # A→B 退休（标记 graduated） +``` + +--- + +## 2. 核心目标：抽取器越来越专业（区别于 mem0/reme） + +### mem0 / reme 的做法与天花板 +- **mem0**：两步固定 LLM 流水 —— Extraction（固定 prompt 抽事实）+ Update（ADD/UPDATE/DELETE/NOOP 去重消歧）。 +- **reme**：分 personal / task memory，带 reflection 从成败轨迹提炼经验，检索时 rerank。 +- **共同天花板**：抽取器**静态**（extraction prompt 永远 day-1），**没有下游效用反馈回流到抽取器**，质量靠 LLM 当下自评。→ memory 只会“越攒越多”，不会“越来越专业”。 + +### 我们的突破口 +下游有**真实效用信号**（蒸馏是否受益、路由是否更准）+ 有 **verifier/rubric 打分**。→ 可以让抽取器被“经下游验证过的 memory”反过来训。 + +### 三个进化层次（由易到难） +1. **meta harness 学会“选对抽取器”**（最快见效）：contextual bandit，按轨迹特征分派抽取器；反馈=各抽取器产出 memory 的效用分。**纯统计即可**，不需训练。 +2. **抽取器“知道什么是好 memory”**：用 rubric 给 memory 打分，维度对准**可用性**（自包含 / 可泛化 / 可操作 / 不冗余），当抽取的软靶 + 写入门槛。 +3. **抽取器本身被蒸馏得更强**（最终形态，与主线同构）：用**下游效用**给抽取器输出打真标签（不是 rubric 自评），正/负样本微调抽取器（同 `llm_backup` 机制，被蒸馏的是“抽取器”这个角色）。 + +### 关键：rubric 分 = 训练靶，效用分 = 真值锚 +- **rubric 内在分**：即时、便宜 → 抽取软靶 + 写入门槛。 +- **下游效用分**：滞后、稀疏、客观 → 校准 rubric、训练抽取器的真值锚。 +- 可用效用分**反向校准 rubric**：某维度 rubric 高分但效用低 → 自动降权（复用 D7c“客观校准主观”）。 +- ⚠️ 坑：别让 rubric 自评当最终真值 —— “看起来专业 ≠ 用起来有用”，同 “模型不知道自己不知道”。 + +### 其它有效手段（业界/研究验证） +- **失败/反思挖掘**：失败轨迹信息量常更大（“这么调工具会报错”是高价值 anti_pattern）。见 §4。 +- **巩固/合并（consolidation）**：定期把多条相关碎片 memory 合并升华成更抽象的通则 → 更泛化，是 B→A 晋升头号候选。 +- **对比式抽取（contrastive）**：`llm_backup` 已收集的 (student 错 / teacher 对) 配对 → 专抽“teacher 做对而 student 做错的那步差异”，精准命中能力缺口。 +- **引用计数衰减**：长期不命中的 memory 降权/淘汰，防历史噪声拖累。 + +--- + +## 3. 毕业机制（B↔A 梯度；先设计不实现） + +memory 是权重的“预备队/退休区”，双线互相喂： +``` +生产轨迹 ─(清洗+评分)─┬─► line B: 抽成 memory → 立刻可检索（快） + │ │ + │ └─ 高频命中 & 已泛化 ──► 晋升为训练样本 + │ ↓ + └────────────────────────► line A: 攒批蒸馏进权重（慢） + ↓ + 权重已掌握 ──► 对应 memory 退休（stale/删） +``` + +- **B→A 晋升**（同时满足）：`hit_count ≥ N`；`kind==action_pattern` 或被判可泛化（排除实体孤例）；`source_score ≥ 阈值` 且 safety 通过。→ 还原成训练样本进 line A。 +- **A→B 退休**：一批训练部署后做**回归检验** —— 让新 student 在**不给该 memory** 时回答，若已能答对（verifier 自动判）→ 标记 `graduated`，移出检索池。 + - ⚠️ 依赖“新权重能否自答”的自动判定，正是“模型不知道自己不知道”；判定靠主线 verifier/置信度，**不能靠模型自评**，否则误退休。 +- 价值：普通 memory 系统缺“遗忘/毕业”机制会无限膨胀 + 过时；有 line A 就天然有毕业出口。 + +--- + +## 4. 失败挖掘与归因（最硬的一块） + +### 现状问题（必须改 pipeline） +按当前 pipeline，**失败轨迹会在打分前/后被丢，不会自动留存**： +- 退化性失败（死循环/复读/心跳/token soup）在 `DeadLoopFilter` 等**硬过滤前置步**就丢了（实测某片 `DeadLoopFilter: 25->2 dropped 23`），**没机会打分**。 +- 能力性失败若侥幸过硬过滤，会被尾部 `TrajectoryOutcomeFilter` 按 `low_traj_score` **删除**。 + +### 归因难题（用户戳中的核心） +一条带 memory 的失败，根因有三，且在最终轨迹上**长得一样**： +1. 模型能力问题（没 memory 也做不对）； +2. memory 缺失（有正确 memory 就能对，但没检索到/库里没有）； +3. **memory 误导**（检索到的 memory 错/过时/不相关，把模型带沟里）。 + +单看一条失败轨迹**无法区分** —— 缺反事实。若把 memory 误导当“能力缺口”抽成 anti_pattern → 坏 memory 的锅上再叠 memory → **无限污染循环**。这正是 mem0/reme 不敢做 memory 归因的原因。 + +### 破局：归因在“memory 使用现场”做，不在“挖掘阶段”猜 +把 memory 注入当成一次**可对照的干预**。对同一 query 跑带/不带 memory： + +| 不带 M | 带 M | 归因 | 动作 | +|---|---|---|---| +| 失败 | 失败 | model_capability | 抽 anti_pattern / 进训练线 | +| 失败 | 成功 | memory_helpful | M 加分 | +| 成功 | 失败 | **memory_harmful** | **退休该 M，不抽 anti_pattern** | +| 成功 | 成功 | memory_redundant | M 中性/可淘汰 | + +**归因是推理时记录的，不是挖掘时推断的。** + +### 剪枝（用户务实取舍） +- **token soup / 退化性失败**：根因是生成层退化，几乎与 memory 无关 → **直接丢，不留、不归因**（对照必然“两边都烂”，得不到 memory 信息）。 +- **归因预算（影子对照 + 留痕）只投在“能力性失败”**。 +- 判“垃圾 vs 值得归因”用**现成便宜过滤器**（DeadLoop/TokenSoup/SpecialChars 命中 = 垃圾直接丢）；判“模型问题 vs memory 问题”才用贵的影子对照。两层闸门，贵的只处理少数。 +- 放弃项（暂）：“反复调错工具”的死循环表面像 token soup 实为高价值 anti_pattern，但难自动区分、占比不高，先不捞。 + +### 落地（部署已确认：可做 A/B 影子跑；memory 注入尚未上线 → 从零把留痕设计进去） +注入路径三件套（第一天就装）： +1. **影子对照采样**：每次带 memory 推理以概率 `p`（约 5%）触发“不带 memory”的影子；两条都用 verifier 判成败，填 2×2 表。常规请求只留痕不对照。 +2. **注入留痕**（`user_data` 信封加字段）： + ``` + injected_memory_ids: [...] + memory_shadow: {ran: bool, without_mem_outcome: 0/1} + inference_outcome: 0/1 # verifier 判 + attribution: helpful|harmful|redundant|model_capability|unknown + ``` +3. **归因判定**：采样命中 → 按 2×2 表当场出结论；未命中 → 挂 `unknown`，靠历史反事实给**概率性**归因（低置信，仅用于排序“哪些 M 优先做对照验证”）。 + +### FailureMiner 的干净输入 +- 只从 `attribution == model_capability` 的失败抽 anti_pattern（已排除 memory 误导）。 +- `memory_harmful` → 触发删 M，不抽新 memory。 +- `unknown` 且历史高度怀疑 harmful → 排进影子对照队列验证，不直接抽。 +- 用 `round_scores` 定位**失败转折点**（哪轮突然掉），anti_pattern 抽那个点。 +- ⚠️ 准确率 > 覆盖率：一条错的 anti_pattern 会主动误导模型 → **宁可不抽，不错抽**；噪声/环境偶发失败不抽。 + +--- + +## 5. 下游效用分（memory 飞轮的燃料） + +### 定义（已确认：outcome 用主线 verifier；粒度 per-item） +效用 ≠ memory 内在质量，而是**边际因果贡献**： +``` +utility(M) ≈ E[ verifier(轨迹) | 注入M ] − E[ verifier(轨迹) | 不注入M ] +``` +`verifier(轨迹)` 用主线连续 `traj_score`（比 0/1 更细，能测“0.6→0.85”的边际改善）。它是 §4 归因的**连续版**（2×2 离散表的加权量）。 + +### 五种信号（弱→强） +- **A 反事实成功率差**（5% 影子对照）：因果最干净、正负都能测；稀疏 → **真值锚**。 +- **B teacher 追平差**（`llm_backup` 免费捎带，teacher 不带 memory）：student 带 M 后 match=True 是否上升；覆盖 llm_backup 采样流量。 +- **C 归因离散效用**（2×2 → +1/−1/0/0）：A 的离散版，够 bandit 奖励用。 +- **D 检索命中×结果**（全流量、便宜、**有选择偏差**）：不能单飞，只当排序候选，靠 A 去偏。 +- **E 多轮内即时行为信号**（`round_scores`：工具一次成功率、少走弯路、少重试）：细粒度，尤适 action_pattern。 + +### 组合（便宜信号打底 + 稀疏真值锚校准，同 D7c） +``` +效用分 = f(观测代理 D, teacher 追平 B, 行为 E) ── 用硬对照 A/C 校准去偏 +``` +1. 日常：累积 D+B+E 加权分（便宜、全覆盖、有偏）。 +2. 校准：5% 对照 A 产出无偏真值，回归；若某类代理分系统性高于真值 → 自动降权 D。 +3. 不确定性探索：代理与真值分歧大/方差大的 M → 提高其对照采样率。 + +### per-item 聚合器（唯一要新写的东西） +```python +MemoryUtility(memory_id): + n_hits; proxy_sum(按相关度加权); n_controlled; delta_sum + utility_hat ∈ [-1,1]; confidence(样本量+锚一致性); last_hit_at +``` +融合：对照足→用无偏 `delta_sum/n_controlled`；对照少→用 proxy 减去同类已知偏差；分歧/低置信→抬高对照采样率。 +per-extractor 分**不单独采**，从 per-item **聚合**上来（某抽取器所有 memory 的 `utility_hat` 均值），零成本供 meta harness。 + +### 驱动的动作 +- 退休/清坏：`utility_hat<0` 且置信够 → 删（memory_harmful 自动闸门）。 +- 晋升：`utility_hat` 高 + 高频 + 可泛化 → 进训练池。 +- 冷启动保护：新 M 低置信给保底曝光 + 时间衰减防老 M 霸榜。 + +### 三个坑 +1. 选择偏差（D 的病）：必须 A/B 去偏，D 不单飞。 +2. 信用分配：一条轨迹多条 M → 初期均摊或按检索相关度加权，别急上 Shapley。 +3. 反馈自强化：高效用被检索更多→分更高→更多… 可能锁死 → 时间衰减 + 强制探索低曝光。 + +### 前提 +地基是 **verifier 判 `traj_score` 的质量**。rubric 打分噪声大则整个飞轮歪 → **先坐实 verifier 打分可靠性，再上效用分**（与“修 TrajectoryScorer 打分区分度”同一条线）。 + +--- + +## 6. meta harness（管理抽取器；contextual bandit） + +### 定性：contextual bandit，不是全 RL；初期连 bandit 都先不上 +- 无状态转移（选抽取器不影响下一条外部流量）→ 不需要 Q-learning/PG。 +- 是 contextual bandit：轨迹特征=context，选抽取器=action，效用=reward。 +- **reward 延迟极大**（几天）→ 先做**离线统计的 bandit**（按 context 分桶统计各抽取器历史效用 + ε 探索），成熟后再升 LinUCB/Thompson。 + +### 状态（context，全复用 preprocessor 标签） +`intent` / `traj_score` 分桶 / `round_scores` 形状 / 轨迹长度·段数 / 是否 agent(`is_agent_row`) / domain。 +- 简版：离散化组合成桶 key，如 `(intent=tool_call, traj_hi, is_agent)`。 +- 升级：特征向量 + LinUCB。**初期分桶就够，别急上向量。** + +### 动作（两级） +- 一级（必做）：从 `{fact, action_pattern, preference, failure_mine, none}` 选一个/几个。**`none`（不抽任何 memory）是合法动作** —— 省成本 + 避免噪声。 +- 二级（成熟后）：抽取参数（粒度/数量/是否跨轮聚合），初期固定默认。 +- 约束：**动作空间要小**（个位数~十几个），否则永远冷启动。 + +### 奖励 +``` +reward(context, action) = agg{ utility_hat(M) : M 由该 action 产出 } +``` +- `agg` 用均值，可选覆盖惩罚（抽太多低效用条目扣分，鼓励精不鼓励多）。`action=none` reward=0 作基线。 +- **延迟处理**：批式/异步。抽取时记 `(context_bucket, action, [mem_ids])` 到待结算表；效用成熟后回填、更新桶统计。 +- 信用分配已在 per-item 层解决，bandit 直接用聚合值。 + +### 探索 +- 简版 ε-greedy，ε 随桶样本量衰减；新桶/新抽取器强制均匀试几次。 +- 升级 Thompson/UCB，用奖励不确定性驱动（与效用分“低置信多做对照”共用信号）。 +- **协同**：bandit 探索性选了冷门抽取器产出的 memory，效用最不确定 → 应**优先安排影子对照**去测它，否则探索白探。 + +### 骨架 +```python +MetaHarness: + policy: dict[context_bucket → dict[action → (reward_mean, n, confidence)]] # 存 StateBackend + select(traj_labels) -> action # ε-greedy over bucket + log_decision(context, action, mem_ids) # 待结算表 + settle(mem_id, utility) # 效用成熟后回填 policy +``` + +--- + +## 7. 稳定性：harness 会漂移，怎么关进笼子 + +“更新 harness 后不稳定”拆成三种病： +1. **策略震荡**：噪声+延迟反馈做了过自信更新 → 来回摆。 +2. **反馈自锁**：偏向某抽取器→只有它有反馈→别人永远没数据→越锁越死（off-policy 经典病）。 +3. **非平稳漂移**（本系统独有、最麻烦）：student 在被持续蒸馏、memory 库/流量在变 → 上周最优抽取器这周可能就不对。普通 bandit 假设平稳，这里不平稳。 + +### 对策：harness 从“决策者”降级为“建议者”，更新慢、可回滚 +1. **冻结基线 + 影子上线**（治 2、防炸）：新策略先只“建议”，实际按固定基线执行，记录“若听 harness 会怎样”；证据表明稳定优于基线才灰度切；永远保留基线兜底。 +2. **慢更新 + 迟滞带**（治 1）：桶样本 `≥ N` 才更新，否则用先验；只有显著且持续优于当前才切换（margin + 连续几批）。 +3. **强制探索地板**（治 2）：每个抽取器保底 `ε_min` 曝光永不归零，防自锁、且能先发现漂移。 +4. **滑动窗口/时间衰减**（治 3）：效用统计只用近期窗口/指数衰减，让 harness 跟随漂移重学。 + +### 更重要：先稳定可用，再进化 —— 分阶段把不稳定源头关掉 +- **阶段0 静态 harness（先上，零 bandit）**：固定人写路由表（`tool_call→action_summarizer`、`code→action+fact`、`低分失败→failure_mine`、`其他→fact`）。不学习不更新、完全确定。先把“多抽取器按类型分派”的**结构**跑通、攒效用数据、当 bandit 基线。**无任何不稳定性。** +- **阶段1 离线 bandit，仅离线复盘时更新**：攒够数据后离线做 off-policy 评估，只有证明稳定优于静态表，才把新策略**作为新静态表发布**上线。策略更新发生在**离线、可审、可回滚**节点，不在线漂移。 +- **阶段2（远期，可选）**：在线自适应。 + +### 分层可降级 mode（把“新方案风险”关进笼子） +``` +harness.mode = 'static' # 固定路由表，永远兜底 + | 'shadow' # bandit 只建议不执行，收集证据 + | 'canary' # bandit 接管 x% 流量 + | 'live' # bandit 全量（需离线验证达标才允许） +``` +出问题一键降回 `static`。 + +### 诚实边界 +- 阶段0 静态表需人工先验（哪个 intent 配哪个抽取器）—— 这不“自进化”，但是冷启动正确起点。 +- 非平稳漂移无完美解，滑动窗口只缓解；靠“离线定期复盘 + 可回滚发布”管理，不追求永远正确的在线策略。 + +--- + +## 8. 现有资产映射（实现时几乎不用造轮子） + +| 需要的能力 | 复用现有 | +|---|---| +| 存储（用户级 KV、策略、效用状态） | `twinkle/server/state`（memory/file/redis + TTL + update_atomic + keys） | +| embedding（全局向量召回） | `preprocessor/experimental/llm_backend.py::OpenAIBackend.embeddings` | +| 抽取器（fact/action/pattern） | `twinkle_agentic/summarizer/*`（都走 `llm_backup` 蒸馏，可 per-type LoRA）；`pattern_summarizer` 现为空 → 做“可复用模式”抽取，是 B→A 晋升头号候选 | +| 抽取器蒸馏 | `llm_backup`（被蒸馏对象换成“抽取器”角色） | +| outcome / 写入门槛 | 主线 verifier / rubric / hard_scorer 产出的 `traj_score / safety` | +| 转折点定位 | `round_scores`（TrajectoryScorer 已产出） | +| 校准逻辑（客观校主观） | D7c 思想直接搬 | +| 标签信封 | `preprocessor/label_schema.py` + `user_data`（PyArrow 稳定）；新增 `injected_memory_ids / memory_shadow / inference_outcome / attribution / memory_utility` | +| 配对（对比抽取 / teacher 对照） | `llm_backup` 已采 (student, teacher, match) | + +**唯一新写**：MemoryStore + VectorIndex + FailureMiner + per-item 效用聚合器 + MetaHarness（含 static/shadow/canary/live mode）。毕业机制先留接口 + docstring + 判据常量，逻辑 `NotImplementedError`。 + +--- + +## 9. 待定 / 下一步（未拍板） + +- 静态路由表的具体先验规则（intent → 抽取器映射）细化。 +- 待结算表 schema、桶设计粒度（先粗按 intent，数据多了再细）。 +- 检索/注入策略本身：检索几条、排序、注入到 context 哪个位置。 +- 双线去重边界：同一高分轨迹既进训练池又进 memory，B→A 晋升时如何避免重复训练。 +- memory 命中对置信度路由的**方向**：命中→更敢自答，还是命中→说明是薄弱区更该路由？（两种逻辑相反，需定。） +- 用户级 memory 的隐私/时效：覆盖写 vs 版本化；TTL。 + +--- + +## 10. Related Works（2026 检索，按本设计的轴归类） + +> 检索源：arXiv（2026-06 ~ 2026-07 为主）。结论：本设计的**每一条主要思路都能在近半年文献里找到平行工作或验证证据**——这是好事（方向被验证、不孤立），差异化在于**把这些点在一个自进化蒸馏框架里闭环组合**，且共用主线 verifier / `llm_backup` / D7c，而非各做各的。下面按“对应本文哪一节”组织。 + +### 10.1 双线（context-space + parameter-space）—— 对应 §0 + +**DuoMem: Dual-Space Distillation**（arXiv 2606.29961） +- 具体做法（分三步离线 + 一步在线，见其 Fig.2）： + 1. **teacher 造料**：用 Qwen2.5-72B teacher 对 3,553 个训练任务各跑 3–4 次（共 11,546 实例，5 次重试内累计成功率 99%），得到 11,434 条成功轨迹；**故意 oversample 同一任务的多条不同解**以增多样性、防过拟合。 + 2. **context-space 蒸馏（CD，训练无关）**：让 teacher（而非 student）对每条完成轨迹**离线生成 procedural memory 脚本**，存成文本 bank（整套任务几 MB）。推理时对新任务 d，用 `text-embedding-3-small` 算 d 与各条 memory 的 cosine，取 top-k **prepend** 进 student prompt。不改任何参数。 + 3. **parameter-space 蒸馏（LoRA）**：只用**成功 teacher 轨迹**微调 student 的 LoRA（rank 8–32，α/r=2，base 冻结）。 + 4. **组合**：先 LoRA 再 CD。ALFWorld 上 Qwen3-4B：No-Mem 4.3% → MemP(student 自产 memory) 55% → +CD 56.4% → +LoRA 72.1% → +DuoMem 77.9%（逼近 72B teacher 87.1%），只加 5.9M 参数、~12MB memory，且比 teacher 快 3×。**关键消融结论：CD 单独收益很小（+1.4），LoRA 才是大头（+17），两者组合还有超加性（>各自之和）。** +- 与我们的区别：DuoMem 的两轴都是**一次性离线固定**（teacher memory 生成一次、LoRA 训一次），**无在线闭环、无毕业机制、只用成功轨迹**。我们要：两轴在线持续、B↔A 毕业梯度、且**失败轨迹也要挖**。另外 DuoMem 的 CD 是"整任务级 memory 检索"，我们是 per-item 效用 + 归因过滤后的 memory。**它是我们双线最直接的可行性背书 + 起点基线**（甚至可先复现 DuoMem 当 line A/B 的 v0）。 + +**KbSD: Knowledge Boundary aware Self-Distillation**（arXiv 2606.29863） +- 想解决的问题：模型经常**不知道自己知不知道**——该答的时候瞎编（幻觉），不该答的时候硬答，或者明明该去查资料却凭记忆蒙。KbSD 想教会模型"划清知识边界"：**会的直接答、不确定的去检索、真不会的就说不会**。难点是普通 RL 只有一个"最后答对没答对"的稀疏奖励，没法教中间推理过程该怎么走。 + +- 怎么做（分三步，公式见原文 §3）： + + **① 给每个问题打三个"边界标注"**（这就是你问的"确定性/靠谱度怎么来的"）： + - **参数化确定性 μ(q)**：拿**冻结模型、不给检索**，对同一问题**独立采样 N 次**，算答对（match 标准答案）的比例 `μ = (1/N)Σ I[match(yᵢ, a)]`。μ 高 = 答案本就在模型脑子里。**注意：这一步需要 ground-truth 答案 a。** + - **语义稳定性 σ(q)**：这 N 次回答**两两之间的 embedding 余弦相似度求平均** `σ = mean cos(Enc(yᵢ), Enc(yⱼ))`。σ 高 = 每次都说同样的话（信念稳）；σ 低 = 每次瞎蒙都不一样。**这个不需要标准答案。** + - **检索质量 ρ̂(q)**：对检索返回的证据打一个 retrieval-quality 分（原文没细化打分器，是可替换的相关性模型/reranker）。 + - **映射到四象限**：把 μ、ρ̂ 各卡一个阈值二值化 → 得"内部知识可靠 k / 检索充分 s"，组合成四种该有的行为：都行→**Integrated（整合）**、只检索行→**External（靠检索）**、只内部行→**Internal（信自己）**、都不行→**Refusal（拒答）**。σ 不进象限，只用来筛训练集（已知区留稳定的、未知区留不稳的），让边界更干净。 + + **② 用"开小灶的自己"当老师做示范**：把上面标注 `(μ, σ, ρ̂, 目标象限)` 拼成一段 **hint 前缀，只喂给老师**；老师 = **同一个模型** conditioned 在 `[hint; q]` 上（且 stop-gradient，不回传老师）。学生 = 同一个模型但**看不到 hint**。因为老师多看了提示，能写出"知道分寸"的推理示范（该查就查、不会就认怂），学生就跟这个"开了天眼的自己"学——**这就是"信息不对称自蒸馏"，全程不需要更大的外部模型**。hint 只在训练用，推理时不给。 + + **③ 学生怎么对齐老师——是 KL 蒸馏，而且按象限切方向**（这就是你问的"用什么 KL"）：token 级两种损失，本质是 KL 的两个方向： + - **前向 KL（mass-covering，覆盖）**：在老师轨迹上最大化学生似然 → 学生尽量覆盖老师所有说法。 + - **反向 KL（mode-seeking，收敛）**：在**学生自己采样**的轨迹上匹配老师 → 学生收敛到老师主模式、抑制老师不支持的行为。 + - **分象限分配**：Integrated（只有一种正确整合、分布集中）→ **反向 KL**；Refusal（合理拒答说法很多、分布发散）→ **前向 KL**（别塌成一个模板）；External / Internal（要精准又抗噪）→ **前向+反向都要**，用 Pareto 加权自动解一个公共下降方向 α*（闭式解，不用手调系数）。这套蒸馏再和 GRPO 的稀疏 outcome 奖励**联合优化**（稀疏管"最终对不对"，稠密管"推理怎么走"）。 + +- 对我们有什么用：它正好治我们主线最头疼的病——**"模型不知道自己不知道"**，这直接决定置信度路由该不该把请求甩给 teacher；而且证明了**不用更大模型也能自我校准**（老师=学生+hint），与我们 `llm_backup` 的 student/teacher 同构。可直接借的两处：**训练置信度路由**、**A→B 退休判定**（退休本质就是判"新权重不给 memory 提示、自己能不能答对"，正好对应 μ）。 +- **⚠️ 搬过来要改的地方**：KbSD 的 μ **依赖 ground-truth 答案**，而我们主线是**无 ground-truth 的生产流量**。所以不能照搬 μ，要用**主线的 rubric/hard verifier 分数、或 teacher 一致性**来替代"match 标准答案"这一步；σ（自洽度）可原样复用，因为它本就不需要答案。 + +### 10.2 memory 越来越有用 / 自优化抽取器 —— 对应 §2、§6 + +**SelfMem: Self-Optimizing Memory**（arXiv 2607.03726） +- 想解决的问题：现有 memory 系统（MemGPT/Mem0/MemoryBank）都是**人写死的 memory 流程**——"存用户画像""到 context 上限就压缩"这种固定规矩，换个任务就不合适、还得手动调。SelfMem 想"授人以渔"：**不给死规矩，让 agent 自己摸索"这个任务下 memory 该怎么攒"。** +- 怎么做（关键：搞清 refine 的到底是什么）：世界里有三样东西—— + 1. **原始对话记录（transcript）**：存成 SQLite 表，**只读、永不改**，是事实唯一真相来源。 + 2. **memory 工作区（workspace）**：agent 自己维护的一块**可读可写白板**，装什么结构由它自己定（用户画像 / 偏好列表 / 项目笔记 / 时间线，甚至一段压缩策略文字）。 + 3. **一组固定的 memory 工具（四类）**：读 transcript（可跑只读 SQL）、读 workspace、写 workspace（加/替换/合并去重/精炼摘要/归档过时/记精确事实）、review（只诊断不改）。 + + 所谓 **"refine 自己的 memory 策略"，refine 的是白板里 memory 的内容与组织方式**（不是改模型权重，也不是改工具本身）：agent 跑一个 **inspect→write→review→revise（查→写→自查→修订）循环**——写完一条就用 review 工具自查"有没有过时/矛盾/没出处/难检索"，拿到诊断后回去改这块 memory。反馈是**多维不塌成单一标量**的（响应质量、token 数、成本、缓存命中），让 agent 在语言层面权衡"多存提升召回但涨成本、压缩省钱但丢细节"。全程**不动模型权重**（与 Reflexion/Self-Refine 同脉，靠语言反馈迭代）。BEAM 上 100K/500K/1M token 比最强基线 official score +0.165/0.141/0.134。 +- 与我们的区别：先厘清一个易混点——**工具集是固定的，agent 不能改工具**；它和"普通调工具"的区别在于**调哪个、什么顺序、写什么、要不要重写全由 agent 按反馈自己决定，没有预设 SOP**（MemGPT/Mem0 = 人写好 SOP、工具是执行的手；SelfMem = 人只给工具+评价、让 agent 自己长出 SOP）。而 SelfMem 优化的是"**这一个 agent 怎么攒/用它自己那块 workspace**"（都在 prompt/流程层，权重不动）；我们的 meta harness 优化的是"**用哪个抽取器/参数把轨迹变成 memory**"，反馈是**下游效用（Δverifier）**而非 agent 自评，且更进一步要把好 memory **反向蒸馏回抽取器权重**。相同的是"给工具+反馈让它自进化"这个哲学，它的"inspect→write→review→revise 循环 + 多维不塌缩反馈"可直接借进我们抽取器的自评环节。 + +**MetaSkill-Evolve: Two-Timescale Recursive Self-Improvement**（arXiv 2607.05297） +- **先说最关键的一句，破除误解**：这篇**完全没有训练、没有梯度、没有 loss**。从头到尾只有**一个冻结的模型**（Gemma-4 31B），"进化"全靠**让这个模型反复读写几个 Markdown 文本文件 + 拿准确率做进化搜索**。所谓"skill/meta-skill"就是几份 `SKILL.md` 文件，不是模型参数。所以你问的"怎么训练、loss 是什么"——答案是**不训练、没有 loss**，它是"改文件 + 挑最好的文件"的搜索过程。 + +- 几个"模型"其实是同一个冻结模型扮演的**五个角色**（靠不同 prompt 区分，各读一份对应的 `SKILL.md`）： + - **Analyzer**：看一条失败案例，诊断"为什么错"，打个标签。 + - **Retriever**：从别的分支里捞点"以前类似问题怎么改好的"当灵感。 + - **Allocator**：决定这一轮生几个候选改法（预算）。 + - **Proposer**：根据诊断，具体写出"skill 文件该怎么改"。 + - **Evolver**：把改动写进文件，并验证一下。 + +- 用什么数据 / 怎么"标注"：**不需要人工标注**。数据就是任务的 (输入, 参考答案) 样本，分成训练批和验证批。所谓"评分"是**自动的**——拿当前 skill 让 agent 去做验证批的题，**答对率就是这份 skill 的分数 `U(s)`**（`r∈[0,1]` 对着参考答案打分）。没有人在中间标任何东西。 + +- 读还是写：**主要是"写"侧的自我改进**——不断重写 skill 文件让 agent 做题做得更好；"读"只是 Retriever 去翻历史分支当灵感，不涉及检索优化。 + +- 两个时间尺度到底在进化什么（这是它唯一的新意）： + - **快环（每轮）**：改 **task skill `s`**（"这个任务该怎么做"的说明书）。拿当前 skill 在训练批上找错得最惨的一题 → 五角色流水诊断+提改法 → 生成几个候选新 skill → 谁在验证批上**答对率更高**就留下（只有严格变好 `ΔU>0` 的才进池子）。 + - **慢环（每 H 轮）**：改 **meta-skill `m`**，也就是"**上面那五个角色自己的说明书**"。关键点：因为五个角色的说明书也是同格式的 `SKILL.md`，所以**用同一套五角色流水去改它们自己**（自己改进自己的改进方法）——这就是"递归自我改进"。判据是 **meta-productivity `P(m|s)`**：这个分支最近 H 个"孩子"平均提升了多少（`= 子代 ΔU 的均值`），衡量"这套改进方法还灵不灵"。 + - **选哪个分支继续进化**：打分 `η₁·当前分 U + η₂·这套方法的产出率 P + η₃·新颖度 N`（N 惩罚被选太多的分支，逼它去探索没试过的路线）。 + +- 效果：OfficeQA/SealQA/ALFWorld 比 No-Skill +23.5/+16.1/+1.9，比"只进化 s、不进化改进方法"的版本 +6.4/+8.1/+1.9（证明"连改进方法本身也进化"确实有额外收益）。 + +- 与我们的区别：它的思想（"**快改内容、慢改'改内容的方法'本身**"）正是我们 §6/§7 想要的——我们的 static→shadow→canary→live 就是它慢环的工程化 + 加了可回滚门控。**但它全程改的是文本文件、模型冻结、没有训练**；我们的慢环最终要落到**真的蒸馏进抽取器权重**（把好经验烧进参数，而非只改说明书），且用**下游效用（Δverifier）**而非"验证批答对率"当信号。一句话：**它是"用一个冻结模型玩进化搜索改 prompt"，我们要的是"把搜出来的好东西训进权重"。** + +**COMFYCLAW: Self-Evolving Skill Harnesses**（arXiv 2607.01709） +- 领域 / 模型：图像生成工作流（ComfyUI）。**不训练模型**——用现成 LLM 当 agent、现成 VLM 当"验收员"，改的是外部的 skill 文件库。 +- 怎么做（一个"边做边攒经验"的闭环，见其 Fig.1）： + 1. 给一个文生图需求，agent 通过**带类型的图编辑**（连节点、调参、加 LoRA）把 ComfyUI 工作流搭起来，跑出一张图。**非法的编辑会被自动撤销**（防止把流程改坏）。 + 2. **VLM 验收员**把需求拆成一串"可观察的是非题"（比如"有没有三只手臂""风格对不对"），逐条判过没过 + 给个 0–10 细节分，合成一个标量分数；并把**没过的条目 + 哪里错了 + 具体该怎么改**回吐给 agent，指导下一轮修改。 + 3. 跨很多需求跑下来，把"反复成功/失败的经验"提炼成**可复用的 Agent Skill（`SKILL.md` 文件）**，存进 skill 库，下次相关需求时**先只给 skill 摘要、需要时再展开全文**（渐进披露，省 context）。 +- 数据 / 标注 / loss：**没有训练、没有 loss**，"分数"来自 VLM 验收员的是非题（自动，无人标）。四个 split×两 backbone×三模型下，比"只有验收员、不进化 skill"的基线高 4 分、比"完全不修改"高 10 分。 +- 与我们的区别：它名字和"harness 管理可复用技能"的思路跟我们碎片1直接撞上，但它是**图像生成域实证**、skill 是"给 agent 复用的操作技能"；我们的"skill/抽取器"是**把轨迹变 memory 的工具**、且最终要蒸进权重。它的两个工程点可直接抄：**非法编辑自动回滚**、**验收反馈翻译成"可执行的修改建议"而非只给一个分**。 + +**UCOB: Credit-Aware On-Policy Bidirectional Self-Distillation**（arXiv 2606.29502）—— **与我们 §4 归因 + §5 效用最像的一篇，务必读透** +- 想解决的问题：检索来的"经验/skill"**不是万能的**——同一个模型，在情形 A 被这条经验帮到、在情形 B 反被它带沟里。所以"把'带经验的回答'当成永远正确的老师去教'不带经验的回答'"这个假设是**脆的、会把坏经验也学进去**。这正是我们担心的"分不清是模型问题还是 memory 问题"的学术版。 +- 怎么做（大白话，这是**真·训练，有 RL loss**）： + 1. **同一个在线模型**，同一道题准备**两种输入**：带经验的（`P₊`）和不带经验的（`P₀`）。每道题各采一批 rollout，一半用 `P₊`、一半用 `P₀`，**两边都参与在线 RL 更新**。 + 2. **在相同的"局面"上比谁做得好**：把两边 rollout 里**走到同一个中间状态**（论文叫 anchor-state）的记录凑成一组，各自算"从这一步往后的总回报"（return-to-go = 后续奖励的折扣和）。定义**同局面下的差值**`Δ = 带经验的平均回报 − 不带经验的平均回报`。 + 3. **谁赢谁当老师，只在这一局面上教对方**：`Δ` 明显为正 → 说明这条经验在这儿确实有用，让"带经验的回答"去教"不带经验的"（把能力吸收进去，以后不给经验也会）；`Δ` 明显为负 → 说明这条经验在这儿是**误导**，反过来让"不带经验的回答"**纠正**"带经验的"（**主动压制坏经验**）。教的方式是 token 级分布对齐 + 置信度门控（只在有把握的位置教）。 + 4. 同一个 `Δ` 还顺便用来**更新每条经验的效用分**（配 UCB 决定以后检不检索它）、并训练"写经验"的 reflection 模块。ALFWorld/WebShop 比 SOTA +23.5/+18.0。 +- 与我们的区别：UCOB 的"带经验 vs 不带经验、比谁回报高"就是我们 §4 **影子对照的在线 RL 内生版**——它在 rollout 内、按 anchor-state 配对做；我们在**推理服务侧按 ~5% 采样**做 2×3 归因表。UCOB 的产物是**改 policy 权重 + 更新经验效用**；我们把同一个"带/不带 memory 对照"信号接到**三个出口**：失败挖掘、per-item 效用、坏 memory 退休。**它强证了我们方案的核心机制可行且高收益**，它的 `Δ` 就是我们 per-item 效用的一种无偏估计，"按相同中间状态配对"这招可直接借来**降低我们效用估计的方差**。 + +### 10.3 失败挖掘 / 从成败双向抽 memory —— 对应 §4 + +**Learning from Failure: Inference-Time Self-Improvement for Computer-Use Agents**（arXiv 2606.31270，**ECCV 2026**） +- 想解决的问题：现在造 agent 训练数据的标准做法是"agent 在有验证器的环境里跑 → **只留成功轨迹**去微调 → 丢掉所有失败"。但失败其实携带了"模型哪里弱"的宝贵信息，全扔了很浪费。 +- 怎么做（**不训练**，改的是 agent 的推理时行为，见其 Fig.2）： + 1. agent 跑一批任务，收集**失败**轨迹。 + 2. 用一个 LLM 当"分析员"，把 (指令、动作历史、思维链) 喂进去**诊断失败原因**，归成**四类**：定位不准（grounding）、能力缺口（该用某工具却不用）、知识缺失、无脑重复循环。 + 3. 对每一类，LLM **提出一个推理时的补救办法并生成一段代码补丁**（分别对应：加视觉搜索、允许走终端执行、注入知识、加重复告警），**人工轻量核对**这段补丁后，并进 agent 的工作流，再重跑。每轮按"当前最主要的失败类型"选一个补丁，补丁跨轮累积。 +- 数据 / loss：**零训练、无 loss**，纯粹是"诊断失败→打补丁→重测"的循环。OpenCUA-72B 在 OSWorld 从 42.3%→48.9%（全部补丁叠加 52.74%）。 +- 与我们的区别：**直接印证我们"别丢失败轨迹"**。但它把失败变成**给 agent 的代码补丁/工具**（改 harness），而且**完全不区分失败是模型本身弱还是 memory 带偏的**；我们把失败变成 **anti-pattern memory + 训练负样本**，且**先做归因**（排除"是 memory 误导"才抽经验）。它那**四类失败诊断**可直接拿来当我们失败挖掘器的分类初值。 + +**ISM: Self-Improving Strategy Memory for Continual Math Reasoning**（arXiv 2606.31191，**ICML 2026 AI4Math Workshop**）—— **毕业/退休机制的现成七件套模板** +- 想解决的问题：冻结的 LLM 做连续不同领域的数学题时，学到的经验存哪、怎么不越堆越乱？纯 retrieval 会无限膨胀、纯 reflection 只存散乱文字。 +- 怎么做（**不训练模型**，只维护一个外部"策略库"）： + 1. 外挂一个**紧凑的 strategy-schema 库**。每条 schema 拆成两半：**content**（策略描述/解题模板/启发式，用时注入 prompt）+ **feature hook**（结构标签 + embedding，决定"什么时候该检索到它"，且随使用自动微调）——把"这条经验讲什么"和"何时被调出来"解耦。 + 2. 检索分两步：先按题型/算子过滤，再 soft 打分选最相关的。 + 3. **库由七个自维护机制打理**（这是精华）：①audit 审查 ②merge 合并近重复 ③prune 删无用 ④promote/demote 升降级——管质量和体积；⑤reinforce 从**成功**里抽正向启发式 ⑥antipattern 从**失败**里记"要避开的坑"——**成败双向都学**；⑦rehabilitate 给表现差的 schema 一次翻身机会再决定删不删。**每次改库都要先过符号验证器**（数学能硬校验），防止把错误泛化写进去。 +- 数据 / 标注：数据是 300 题的连续流（按域分块），"对不对"由**符号验证器自动判**（数学域独有的硬 verifier），无人工标注。backbone 是 gpt-4.1-mini、temperature=0。 +- **准确率数字要拆开看（别被 0.48→0.81 唬到）**（Table 1，MATH-Hard 累计 acc）： + - Vanilla 48.0 → RAG 57.0 / Reflexion 55.7 → **Static Schema 78.67**（bank=1，一个固定 prompt 模板 + 允许调符号工具）→ Passive 78.67 → ISM 80.67。 + - **+30 点的大头来自 Static Schema——即"好 prompt 模板 + 符号验证工具"，与 memory 无关**；真正属于"memory 自维护机制"的（ISM vs Passive）**只有 +2 点**（300 题净多对 6 道，OlympiadBench 同样 +2）。作者自承单 seed、单 stream 顺序、无逐机制消融，"gains should be interpreted as preliminary"。 + - ISM 的真卖点其实是**省存储 + 抗遗忘**（bank 只有 Passive 的 1/3~1/7、比 RAG 少 23×），不是"memory 让数学变强"。 +- **⚠️ 评测可信度存疑（重要）**：全文**没有任何去污染措施**——grep 全文无 `decontamination / n-gram / 13-gram / dedup / held-out / train-test split`。而且它的"记忆积累"和"评测"用的**是同一条 300 题流、没有独立 held-out**：RAG 基线是"**把做过的每道题（含解）全存进库、按 embedding 召回最像的一道注入**"，对 MATH/OlympiadBench 这种"换数字的同型题很常见"的数据，等于系统性地**召回近似题的解 = 流内信息泄漏**，且未做任何近似题过滤。所以 RAG 的 0.57、schema 系的 0.79 都**掺了泄漏水分 + prompt/工具增益**，**这些数字不能当作"memory 在数学上有效"的证据**。 +- 与我们的区别：ISM 的**七机制 + 验证门 + content/feature 解耦**几乎就是我们"毕业（promote/demote/retire）+ 效用退休 + 去重合并"的**现成设计模板**，它的"成功→正样本、失败→anti-pattern 对称双向抽取"正是我们要的——**我们借的是这套领域无关的 lifecycle 机制**。但**不采信它的数学有效性结论**：它靠数学专有的硬符号验证器、且评测有泄漏/缺干净对照；我们无此硬 verifier（只有带噪声的 rubric+hard 软 verifier）。而且它这套结论**恰好与我们自己的实验吻合**——我们 200 题数学 memory 只比 direct 多对 1~2 道，而 ISM 剥掉 prompt/工具/泄漏后 memory 机制也就 +2 点：**两边共同印证"数学这类可泛化硬技能应走线 A（训权重），memory（线 B）边际只有 1~2 个点"**。ISM 还是权重全程冻结、纯外部 memory；我们是双线，退休判据还多一条"新权重不给 memory 能不能自答"。 + +**M2Note: Mistake Notebook Learning**（arXiv 2607.00685） +- 想解决的问题：怎么把失败经验安全地攒成"错题本"，又不会因为写错东西把整体带崩。 +- 怎么做（不训练，改外部笔记）：把失败轨迹提炼成**按主题组织的"错题本"note**，检索注入引导 agent 规避同类坑；关键工程点是 **批级后验 + 回滚**——一批笔记编辑先在同一批任务上验证，**只有整批指标真的提升了才提交，否则整批回滚**。支持同模型自我进化，也支持"一个模型的错题本给另一个模型用"（= 我们的 student/teacher）。 +- 与我们的区别：它的"批级后验 + 只在变好时提交、否则回滚"与我们 D7c 校准 / harness 发布门的"批级回归检验 + 达标才发布"**几乎一样**，是 §7 稳定性的又一独立佐证；"跨模型进化"正对应我们 teacher→student 的 memory 迁移。差异仍是我们把它接进**归因 + 双线 + 效用退休**的完整闭环，而非只做注入引导。 + +### 10.4 归因 / 去偏 / "何时不该写 memory" —— 对应 §4 归因难题、§5 选择偏差 + +**GovMem: When Not to Write Memory — Governing False Promotion from Correlated Traces**（arXiv 2607.02579，**MLISE 2026**）—— **精确命中"分不清是模型还是 memory 问题 + 会污染"** +- 想解决的问题（一句话）：**"重复出现"不等于"多份独立证据"**。五个 agent 说同一句，可能是五次独立发现，也可能是**同一条过时笔记经共享上下文回声了五遍**。如果按"出现次数多就晋升"的朴素规则，长期 memory 会慢慢变成"把相关性错误固化下来的持久层"。 +- 怎么做（**不训练**，是一个"该不该写这条 memory"的审计决策，四步 write-path，见其 Fig.1）： + 1. 把说法相近的观测**聚成一条候选 memory**，同时保留每条的**来源信息**（哪个 source、哪套 prompt、哪个父事件、什么环境、可信度）。 + 2. **算"去掉相关性后的有效支持度"**：共享同一 prompt/工具/父事件的观测**不算独立的一票**（这是核心——把"回声"折价）。 + 3. 检索**反证**、检查"这条经验声称适用的范围"合不合理。 + 4. 综合上面判断，输出三选一：**promote（写）/ reject（拒）/ needs-review（送人审）**。 +- 数据 / 效果：合成压力测试里把"错误晋升率"从 0.597（按出现次数）降到 **0.040**，同时召回还保 0.960，代价是 15% 送人审。**最刺眼的是真实数据结果**：133 条高影响候选经人裁后，**0/133 可以安全自动晋升，本地门控判过的 11 条全被人否掉**。 +- 与我们的区别：GovMem 是个**保守的审计侧策略**，只管"写不写"、不管"写了有没有用"。我们把它当**"晋升到 A 线 / 长期 memory 之前"的闸门**，再叠上我们的 verifier 地基 + 灰度人审；它的四步（留来源→按相关性折价→找反证→三选一）可以直接做我们 promote 的前置检查。它"几乎没有能安全自动晋升的"这个结论，**强烈警示我们：宁可漏抽也别错抽，默认走 needs-review 灰度**。 + +**MemDelta: Controlled Baselines and Hidden Confounds in Agent Memory Evaluation**（arXiv 2606.29914） +- 想解决的问题：报告里 memory 系统"比 RAG 强"的结论，常常混进了 LLM/embedding/检索管线本身的变化——到底是**memory 架构真强**，还是只是**换了个更好的 embedding**？ +- 怎么做（**一篇测量方法论的论文，不提新架构**）：在 LongMemEval-S（500 题、每人 50+ 会话、三个模型族）上，**一次只变一个变量、其余全冻**，隔离四个隐藏混淆：检索质量、embedding 选择、模型长上下文行为、写路径成本。 +- 关键实测（这些数字很能说明问题）：① **agent 自产 memory 只有 42%，反而不如朴素 retrieval**；② 只换 embedding 不动别的，Mem0 就从"比 RAG 基线 +11pp"翻转成"−1.2pp"——**结论被一个变量掀翻**；③ Mem0 只在窄题型上占优，但**写路径成本可占 agent 总执行时间 80%+**。建议：固定 embedding、按模型族分层、把 write 成本当一等公民报告。 +- 与我们的区别：**直接给我们 §5 选择偏差 + §4"必须做对照"背书**——它证明了"不做受控对照，memory 的效用数字根本不可信，甚至自产 memory 是负收益"。这正是我们坚持"per-item 效用必须用 **带/不带 memory 反事实对照**去偏、且先把 verifier 地基坐实再上效用分"的理由；它的"固定 embedding + 报告 write 成本"应直接写进我们的效用评估协议。 + +**Stealthy Memory Injection in Persistent Personal Agents**（arXiv 2607.05189） +- 想解决的问题 / 做法：展示在持久化个人 agent 里，**坏的或恶意的 memory 能被悄悄写入并长期潜伏**——用户看不见、还跨会话反复生效造成危害。 +- 与我们的区别：佐证我们需要 `memory_harmful` **自动退休闸门** + 用户级 memory 的隐私/时效治理（§9 待定项）——尤其我们有用户级 scope，更要防"一条坏偏好被写进去后反复检索注入"。 + +### 10.5 memory 命中的"选择性使用" + memory-on/off 对照 —— 对应 §4 影子对照、§9 路由方向 + +**ATMem + STR-GRPO: What Memory Do GUI Agents Really Need?**（arXiv 2606.31612）—— **STR-GRPO 就是我们"影子对照"的 RL 化，几乎逐条对应** +- 想解决的问题：GUI agent 做长任务时，光"把过去看到的存下来"不够——它还得知道"这条信息现在**该不该用、用了到底有没有帮助**"。 +- 怎么做（两部分，见其 Fig.2；这是**真训练**：先 SFT 再 RL）： + 1. **ATMem（把 memory 从被动存储变成主动的"执行状态"）**：memory 不是流水账，而是一张**结构化的任务进度表**——记着"整体进度 + 约束""每个待办项的内容 + 它的状态（待办 / 已完成 / 跳过）"，由 agent 边做边更新。先用"只保留通过验证器的成功轨迹"造 SFT 数据（120 模板 → 1.1K 实例 → 21,713 步级样本），教会模型**会建、会更新、会引用**这张进度表。 + 2. **STR-GRPO（用对照实验学"何时该用 memory"）**：对同一道题采一批 rollout，**刻意一半开 memory、一半关 memory**（只切 memory 这一个开关，其余历史都保留）。奖励 = 最终有没有做成（验证器给 0/1）**减去 memory 的使用成本**（用了 memory 却多走步、没帮上忙，就扣分）。因为同一题的开/关两组**共用同一个打分基准**，所以"开 memory 比关 memory 好多少"就**直接量化成了这条 memory 通道的净贡献**，模型据此学会"该用才用"。 +- 与我们的区别：它"同题下开 memory vs 关 memory、比谁做得好"**就是我们 §4 影子对照 2×2 表的 RL 内生版**，它的"memory 使用成本"惩罚正好回答我们 §9 待定的"何时该用 memory、何时 memory 反而是负担"。差异：它在 **GUI 域、RL rollout 内**做、产物是改 policy；我们在**推理服务侧按采样**做、产物接归因+效用+退休三出口，且走双线蒸馏而非纯 RL。**它"开/关对照算净贡献"的思路可直接当我们 per-item 效用的估计量。** + +**WorldEvolver: Self-Evolving World Models for LLM Agent Planning**（arXiv 2606.30639） +- 想解决的问题：给 agent 装一个"世界模型"（行动前先预测后果）能帮规划，但**预测不准时反而会把 agent 带偏**。而每次都靠梯度更新去修这个世界模型，在线部署下太贵、还会灾难性遗忘。 +- 怎么做（**关键：agent 和世界模型的参数全程冻结，只改外部 memory**）：三个模块——① **情节记忆**：把真实发生过的"状态→动作→结果"存下来，检索出来做"检索式模拟"；② **语义记忆**：把"预测的结果 vs 真实观测"对不上的地方，提炼成可复用的启发式规则；③ **选择性前瞻**：预测在喂给 agent 之前，**先把低置信度的预测过滤掉**，只把有把握的预测注入。 +- 与我们的区别："**没把握就别注入**"正是我们 §5 不确定性驱动 + 主线置信度门控的同款直觉（我们把它用在"低效用/高风险 memory 不注入"上）；"从预测-真实的失配里提炼规则"对应我们从失败挖 anti-pattern。差异：它**冻结参数、只改 memory**；我们要把提炼出的规则进一步走 B→A 蒸馏进权重。 + +### 10.6 记忆巩固 / 晋升与身份稳定 —— 对应 §3 毕业机制、§7 稳定性 + +**Episodic-to-Semantic Consolidation Without Identity Drift**（arXiv 2607.01988） +- 想解决的问题（一个偏"合规/审计"的场景）：受监管的长期部署 agent（医院、工厂机器人）有一个**加密认证的身份**（对一份 manifest 做哈希）。传统"巩固知识"的做法（微调 / 改 prompt / 蒸馏 / 追加反思）都会改动定义身份的那份东西，于是**每学一点新知识就等于换了个 agent、要重新认证**。矛盾：既要越用越聪明，又要身份字节级不变。 +- 怎么做（**v1 完全不用 LLM、是确定性统计规则**）：把"巩固"定义成一个**确定性函数** `f: 情节日志 → 语义层`。情节日志是只追加的原始事件记录；`f` 就是**按 (技能+对象+场景) 分组、数成败、算成功率**，输出一行带 **置信度 + 观测数 + 溯源指针** 的语义事实（例："对玻璃杯、这个环境，建议抓取力 25N，置信 0.83，基于 15 次观测"）。关键设计：**身份哈希在构造上就不读这个语义层** → 无论巩固多少次，身份字节不变；planner 只能**只读**查询语义层、不能改。（用 LLM 的 v2 被明确列为 future work，因为引入不确定性会破坏可审计性。） +- 数据 / loss：**无训练、无 loss**，纯确定性聚合。合成 benchmark（1000 决策）上，对比一个校准过的 Bayesian 基线，planner 的无效尝试降 79.82%，同时身份哈希全程字节相等。 +- 与我们的区别：`f: 情节→语义` 正对应我们 **B（快 memory，情节性）→ A（长期语义）** 的毕业方向，"每条带置信度+溯源、可审计"对应我们的血缘记录；"学新知识不改身份"提示我们**毕业/退休不该破坏模型稳定的基础能力/人格**（呼应 §7）。差异：它只在 memory 层做确定性聚合、**完全不碰权重**、且是单向（情节→语义）；我们的 B→A 是**真的权重蒸馏**，还有 A→B 反向退休。 + +**SEA: Self-Evolving Agents with Anytime-Valid Certificates**（arXiv 2607.00871）—— **几乎就是我们 §7"把进化关进可回滚笼子"的统计化理论** +- 想解决的问题：自进化 agent 有个根本麻烦——它用来学习的数据、评判自己的评估器、用的组件，**全是被它自己更新的策略生产出来的**（自己考自己、自己出题自己判）。这种"闭环自产"下，经典学习理论的保证（收敛、不遗忘、安全改进）**全部失效**。 +- 怎么做（四层架构 + 两条铁律，见其 Fig.1；注意它**主要不做权重微调**）： + - **四层**：L0 = 冻结的底座模型；L1 = 一个很小的 **steering adapter**（在线只调"选哪条指令"的概率分布，不动权重）；L2 = **带版本的 harness**（prompt/工具/预算/技能库，可改可扩，每次改都记一个新版本）；L3 = 在旁边的调度器（不在推理主路径上）。 + - **铁律一：每次自我修改都要过一道"随时有效"的统计门**。因为 agent 每轮都在偷看自己的成绩，普通固定样本量的显著性检验会失效；它改用一种**允许你随时停下来看、结论都成立**的统计量（e-value），对一个固定的"错误预算"发一张**可审计的证书**（通过/暂缓/拒绝/无解），全部记进一个账本。快环每轮调 L1、慢环每隔 K 轮改 L2。 + - **铁律二：门只能在"底座本来就能做出的行为"里挑**（不能凭空造新能力）。所以另配五个"验证器在环"的引擎（best-of-N、微步搜索、自写复现测试、搜索层控制、自修复）来**产生候选行为 + 提供密集的、不靠人打分的信号**。SWE-bench Verified 上 +4/+5。 +- 与我们的区别：SEA 用**统计证书 + 错误预算**保证每次进化不把系统带崩；我们用**工程化的发布分层（static→shadow→canary→live + 可回滚）**做同一件事——两者可互补（我们的 canary 达标判定可以升级成它这种"随时有效的统计门"）。它的"L0 冻结 / L1 只 steer / L2 可改 harness"分层，给我们"哪些能在线漂移、哪些必须冻结"划了清晰边界；"门只能挑已有行为"正是我们要把 harness 关进笼子的理由。差异：SEA **基本不微调权重**（只 steer L1）；我们主线恰恰要蒸馏权重，所以更需要它这套门控来兜底。 + +### 10.7 rubric 作为 reward / 可靠性 —— 对应 §2 rubric 靶、§5 verifier 地基 + +**RuVerBench: Can LLM-as-a-Judge Reliably Verify Rubrics in Agentic Scenarios?**（arXiv 2606.29920） +- 想解决的问题：现在大家用 LLM 当"裁判"、按 rubric（评分条目）给 agent 打分，但**这个裁判本身靠不靠谱**没人系统测过——尤其 agent 输出又长又复杂（深研报告几千 token、编码轨迹几万 token）时。 +- 怎么做（**一个 benchmark，不训练**）：构造 2458 条样本（深研 1615 + 编码 843），每条 = (一段 agent 输出, 一条 rubric, 人工标的"满足没满足")。人标经**双人独立标注 + 裁决**，两组一致率 90.4%、κ=0.808（很高）。然后拿各种前沿模型当裁判去判，测它们和人标的吻合度，并测"多判几次投票""一次判多条"这些策略。 +- 结论：**即便最强模型判 rubric 也有明显噪声**；**编码类、尤其涉及 tool-use 的 rubric 判得最差**；多数投票有效但**收益递减**；一次判多条省钱但掉准确率。 +- 与我们的区别：直接支撑我们的**执行顺序**——"**先坐实 verifier 打分可靠，再拿它去算 memory 效用分**"，否则效用分建在噪声地基上（和 MemDelta 的警告叠加）。它"编码/tool-use 类判得最差"正戳中我们 agentic 场景，它的投票收益递减曲线能帮我们定"什么时候值得多投几票"。 + +**MRRG: Many Voices, One Reward — Multi-Role Rubric Generation**（arXiv 2607.01830） +- 想解决的问题：现在"自动生成 rubric"多是**一个通用评估器一口气列所有标准**，容易漏维度（论文叫"维度盲点"），进而导致判分看领域、还能被"只优化被覆盖的标准"刷分（rubric hacking）。 +- 怎么做（**训练无关、无需参考答案**）：让同一个 LLM **轮流扮演多个角色**（用户、领域专家、教育者、AI 研究员、语言学家…），每个角色从自己视角产一批**原子、可验证**的 rubric 条目，再**汇总去重**成一个可审计的打分器。这个打分器既能做偏好判定，也能**直接当 GRPO 类 RLVR 的 reward**。 +- 效果：RewardBench-2 / JudgeBench / PPE 上比单角色基线 +3.1~16.4pp；用作 RLVR reward 时 +1.7 / +3.4。 +- 与我们的区别：可直接增强我们 `RubricVerifier` 的 **rubric 生成质量**——把当前"单 prompt 生成 rubric"升级成"**多角色生成再合并**"，且它天然兼容我们主线 verifier（同一套 verifier 既供 memory 效用、也供训练奖励）。 + +### 10.8 memory 系统工程形态（存储/检索）—— 对应 §1 选型 + +**MOSS: Auditable Agentic Memory**（arXiv 2607.04391） +- 想解决的问题：主流 RAG 用向量相似度检索，**不透明、难审计、有理论上限**——在长期/个人/受监管场景尤其致命。 +- 怎么做：由 **agent 分析查询意图 → 参数化一条结构化检索 → 在关系库上用 SQL 确定性地取数**，**检索环里没有 LLM**（一旦查询定好，执行完全可复现）；词表从语料自动归纳、不外加本体；每一步从建索引到出答案全部可审计。已**真实生产部署约一年**（约 4400 万 token 语料、每天当主力工作记忆用）。 +- 与我们的区别：**印证我们"检索用结构化 KV + 可审计、LLM 只在写/抽 memory 时参与"** 的选型——检索热路径不放 LLM，省成本、可复现、可审计。 + +**Mandol: Agglomerative Agent Memory**（arXiv 2606.29778） +- 想解决的问题：现有系统把向量库、图库拆成好几套，**跨库 I/O 慢、信息碎片化**，RAG 式检索又容易招噪声、漏关联、控不住 token 预算。 +- 怎么做：用一套 **SemanticMap + SemanticGraph 的内存数据结构，原生融合 KV / 向量 / 图**（不是拼三套系统），提供统一的混合检索算子；检索走"查询自适应路由 → 去噪/消解冲突 → 按 token 预算生成上下文"，**全程不调 LLM**。LoCoMo 92.21% / LongMemEval 88.40%（均最优），10 QPS 下检索延迟比最快基线还低 **5.4×**。 +- 与我们的区别：直接支持我们"**混合索引（向量 + 结构化 KV）+ 应用层检索、检索不放 LLM**"的方向，且证明融合式索引比拼装更快更准——是我们 §1 存储层的可参考实现形态。 + +**HyphaeDB**（arXiv 2606.28781） +- 怎么做：把 HNSW 近邻图当作**多 agent 之间的知识传播网**（gossip 扩散 + 能量衰减 + 自发共识，让高价值知识自然传开、陈旧的自然淡出），并给了 **pgvector 参考实现**。 +- 与我们的区别：现在用不上（我们先做单机混合索引），但将来若**全局向量层要升级成多副本/多 agent 协同**，它的"能量衰减 = 自动淘汰陈旧 memory"和我们退休机制思路一致，可作远期参考。 + +**其它形态（提示 memory 不止"注入文本"一条路）** +- **PLACEMEM**（2607.04089）：按算力预算调度 memory 平面。 +- **Neural Procedural Memory**（2606.29824）：用**隐式 activation steering**（直接改模型激活，而不是往 prompt 里拼文本）来承载程序性 memory。 +- **Analytic Concept-Centric Memory**（2606.29774）：以概念为中心组织 memory。 +- 与我们的区别：后两者提示"注入 memory"未必只有"prepend 文本"这一种——**改激活（activation steering）是一条可选的补充通道**（和 SEA 的 L1 steering adapter 呼应）。我们当前走文本注入，把它列为远期 B 线的可选实现。 + +### 10.9 对比表：本设计 vs 代表性工作 + +| 维度 | 本设计（twinkle） | 最接近的工作 | 我们的差异 | +|---|---|---|---| +| 双线（memory + 权重） | context 注入 + `llm_backup` 蒸馏，且有 B↔A 毕业梯度 | DuoMem：CD(teacher memory 检索 prepend)+LoRA(成功轨迹) | DuoMem 两轴**离线一次性、只用成功轨迹**；我们在线闭环 + 毕业 + 失败挖掘 | +| 抽取器自进化 | meta harness（bandit）+ 效用回流 + 抽取器可蒸馏 | MetaSkill(两时间尺度五 agent)/ SelfMem / COMFYCLAW | 它们在 prompt/skill 层递归、权重冻结；我们慢环落到**权重蒸馏** + 分层可回滚笼子 | +| 失败挖掘 | 只挖“能力性失败”，token soup 直接丢 | Learning-from-Failure(四类诊断→patch)/ ISM(七机制)/ M2Note | 它们不区分模型/memory；我们前置**归因**排除 memory 误导后才抽 anti-pattern | +| 归因（模型 vs memory） | 注入现场留痕 + 5% 影子对照 2×2 表 | UCOB(CBSD:anchor-state ΔG)/ ATMem(STR-GRPO 干预 advantage) | 同为 memory-on/off 对照；它们在 RL rollout 内改 policy，我们在服务侧接失败挖掘+效用+退休三出口 | +| 效用分 | per-item 反事实 Δverifier，便宜信号 + 稀疏锚校准 | MemDelta（对照诊断证据）/ ATMem(memory-cost reward) | MemDelta 只做评估诊断；我们把去偏后的效用直接驱动晋升/退休 | +| 何时不写 memory | 门槛（traj_score/safety）+ harmful 退休 | GovMem(provenance→依赖去相关→反证→三路决策) | GovMem 是保守诊断策略、只判写不写；我们与 verifier 地基 + 灰度 + 效用结合 | +| 稳定性/进化治理 | static/shadow/canary/live + 离线可回滚发布 | SEA(四层 + e-value certificate 门)/ M2Note(批级 rollback) | 我们用工程 mode 分层，SEA 用统计 certificate；canary 判定可升级成 anytime-valid gate | +| verifier 地基 | 主线 hard+rubric 融合，先坐实再上效用 | RuVerBench(可靠性有噪声)/ MRRG(多角色 rubric) | 我们直接复用主线 verifier，不为 memory 另造评审；rubric 生成可升级为多角色 | + +### 10.10 我们仍然新颖的地方（综合判断） +单点都有平行工作，但**没有一篇把下面这套完整闭环合在一起**： +1. **同一批生产轨迹**同时喂“慢权重蒸馏”和“快 memory”，且两者之间有**显式毕业梯度（B→A 晋升 / A→B 退休）**——退休用“新权重能否自答”自动判定（DuoMem 无毕业；Consolidation 无 A→B 反向）。 +2. **失败挖掘前置因果归因**：先用注入留痕 + 影子对照区分“模型能力 / memory 缺失 / memory 误导”，**只在 model_capability 上抽 anti-pattern**，把 GovMem/MemDelta 警示的污染从源头挡掉（Learning-from-Failure 类不做 memory 归因）。 +3. **抽取器本身被“经下游效用验证过的 memory”反向蒸馏**，效用分是 per-item 反事实 Δverifier、且**与主线训练奖励同源**（SelfMem/MetaSkill 优化的是策略/prompt，不回流蒸馏抽取器权重）。 +4. **进化被工程化为可回滚的发布流程**（static→shadow→canary→live），而非在线自由漂移——把 SEA 的“门控只能 select 已有行为”落成部署 mode。 + +**一句话定位**：DuoMem 证明了双线值得做，UCOB/ATMem 证明了 memory-on/off 对照能归因，GovMem/MemDelta 证明了不归因会污染/被混淆，SEA 证明了进化要门控——**本设计是把这些已被各自验证的结论，收进一个共用 verifier/`llm_backup`/D7c 的单一自进化蒸馏闭环。** diff --git a/src/twinkle_agentic/preprocessor/AUDIT.md b/src/twinkle_agentic/preprocessor/AUDIT.md new file mode 100644 index 000000000..a5f1a71bd --- /dev/null +++ b/src/twinkle_agentic/preprocessor/AUDIT.md @@ -0,0 +1,179 @@ +# Preprocessor 审计与整改清单 + +> 审计范围：`src/twinkle_agentic/preprocessor/` 全部 15 个文件、30 个类。 +> 审计方法：逐行只读审阅 + 关键断言代码复核 + cookbook/tests 实际接线核实。 +> 三轮视角：(A) 实现问题 (B) 类设计/拆分合并 (C) 功能增删。 + +## 实施状态（已按本清单完整落地） + +> A1 经确认跳过：R1 已把 `score_filter.py` 整体移入 `experimental/`（零 active 使用的死代码）， +> 对死代码再做 4 文件拆分只增维护面、零收益，启用前再拆。其余 21 项全部实施。 + +| 项 | 状态 | 落地位置 | +|----|------|----------| +| A5 | ✅ | `label_schema.py`（`user_data` 信封 + `set_labels`/`get_label` + `pack_value`） | +| P3 | ✅ | `message_normalizer.py` Pass1 重建 assistant 时 `dict(msg)` 透传全字段 | +| P7 | ✅ | `hard_filter._has_tool_calls` / `message_normalizer._strip_heartbeat`+`_is_atomic` 全部走 `normalize_tool_calls` | +| D7 | ✅ | `trajectory_scorer.py`（Segmenter→HardScorer 逐轮→fuse_segment→aggregate_trajectory→写回 `user_data`，mapper 不删） | +| D7c | ✅ | `trajectory_scorer.py` `_segment_confidence`（一致性+voting稳定+决断性）+ `RubricVerifier.score_detail(extra_context=)` 客观注入重评 | +| D6 | ✅ | `outcome_filter.py`（纯读 `traj_score`/`safety_*` 标签比阈值，fail-open） | +| D8 | ✅ | `safety_scorer.py` + `RubricVerifier(fixed_rubric=)` 固定安全 rubric | +| D9 | ✅ | `pii_presidio_filter.py` `regex_only=True`（stub NlpEngine 免 spaCy，REPLACE→MASK 免 faker） | +| D10 | ✅ | `provenance.py`（`ProvenanceStamp`，血缘写入 `user_data`） | +| R1 | ✅ | `experimental/`（`score_filter.py` + `llm_backend.py` git mv 移出主包） | +| R3/R4/R5 | ✅ | `intent_classifier.py`（默认不删；DEFAULT_DETECTORS 精简为 ToolCall/Code/Math；LLM 路径经 R1 已全归 `llm_backup`） | +| P1/P5/P6/P8/P10 | ✅ | trim 后重算 `is_agent`；deadloop agent 行改扫有文本轮；`max_rounds` 按 pair；refuse 扫全 assistant+可选 reasoning；system 多模态保护 | +| A2/A4 | ✅ | `logprob_utils.py` + `message_utils.py`（`utils.py` 保留 shim）；`intents.py` 常量下沉 | +| A3 | ✅ | `twinkle/preprocessor/base.py` 基类返回 `Tuple[List,List]` + `Mapper`/`Filter` 语义基类（`ModelFilter`/`ProvenanceStamp` 已改用） | +| D4/D5 | ✅ | `language_filter.py`（langid 可选，启发式回退）；`structural_noise.py`（关键词无关噪声轮打标） | +| D1/D2 | ✅ | `offline/near_dedup.py`（MinHash-LSH，datasketch 可选+纯 Python 回退）；`offline/decontaminate.py`（13-gram 重叠，drop/tag） | +| A1 | ⏭️ 跳过 | 见上（R1 已隔离为死代码） | + +--- + +## 0. 结论速览 + +> 本清单已根据 review 意见复核收敛：P2/P4 撤销，P1 降级，P11/P12 归入 R1（死代码，暂不单独修）。 + +- **共需改动 22 项**：实现问题 6（P1、P3、P5–P10）、结构重构 5（A1–A5）、功能增删 11（R1–R5 + D1/D2/D4/D5 + D6/D7/D7c/D8/D9/D10，D3 废弃）。 +- **必须做（会静默损坏训练数据）**：仅 **P3** 一项（工具归一丢 reasoning 字段）。 +- **达成「干净 + 每轮评分」最终目标的核心**：**A5**（`user_data` 信封，去 DAG 前置）+ **D7**（接线 verifier，分数写回每轮）+ **D7c**（自动校准 + 客观纠偏主观重评）+ **D6**（读标签滤废案）+ **D8**（安全 rubric）+ **D9**（PII 纯 regex）。 +- **一句话结论**：现有清单修的是「清洗器 bug + 基础过滤」；要产出「干净且每轮带可信分」的 trajectory，还差——**A5 统一 `user_data` 标签信封（把评分/过滤解耦成打标 mapper + 末尾读标签 filter，去掉 DAG）+ 每轮评分打标(D7) + 自进化校准(D7c) + 废案过滤(D6) + 安全/PII(D8/D9)**。零件多数已存在（`verifier`+`aggregation`+`RubricVerifier`+`llm_backup`），核心工作是**接线 + 定 `user_data` 契约**。 + +### review 复核结论（撤销 / 降级项） + +| 原项 | review 意见 | 复核结论 | 处置 | +|------|-------------|----------|------| +| **P1** trim 砍 tool 尾 | 不以 assistant 结尾的部分无训练必要，最多用于工具调用打分 | 成立。trim 尾部未闭合 tool 对训练无害；`is_agent` 不更新的副作用仅剩“末尾 `assistant(tool_calls)` 无对应结果”，训练时本应 mask | **降级为中等**，改描述，不再算“数据损坏” | +| **P2** heartbeat 误杀 | 这类数据是 openclaw/OpenHands 常见格式，作者本意就是要删 | 成立。agent 轨迹清洗语境下 heartbeat 轮几乎必为真噪声，误杀率极低；`message_normalizer.py:26-27` 注释确认是**故意**删除 | **撤销**（保留现状；可选加词边界，非必做） | +| **P4** 删 reasoning-only 轮 | 只有 thinking 无工具调用，训练无落点 | 成立。纯 thinking 轮无 target 输出，多轮里是悬空推理，删掉合理 | **撤销** | +| **P11** ParaphraseScorer 崩溃 | 应该没有实际使用 | 成立。`ScoreFilter`/`ParaphraseScorer` **全库零 active 使用**（仅自身定义 + docs 示例 + 注释掉的引用），测试只覆盖 `utils` 数学函数 | **归入 R1**（死代码，启用时再修） | +| **P12** IFD 公式口径 | 同上 | 同上 | **归入 R1** | + +--- + +## 一、实现问题（正确性 / 语义） + +### 严重：会静默损坏训练数据 + +| ID | 位置 | 问题 | 改动 | 预期收益 | +|----|------|------|------|----------| +| **P3** | `message_normalizer.py` Pass 1 重建消息 | 工具归一路径只保留 `role/content/tool_calls/tool_call_id` 四字段，**丢弃 `reasoning_content`/`thinking`/`name`** | 重建时透传全部原字段 | reasoning 蒸馏数据不再被清洗流程静默剥离 | + +### 中等：策略漏洞 / 语义错位 + +| ID | 位置 | 问题 | 改动 | 预期收益 | +|----|------|------|------|----------| +| **P1** | `message_sanity.py:317,324-329` trim + `is_agent` | trim 掉末尾未闭合 tool 结果本身对训练无害，但 `is_agent` 在 trim **前**计算、trim 后不更新，残留“末尾 `assistant(tool_calls)` 无对应结果”，`check_tool_matching`（forward-only）不拦 | trim 后重算 `is_agent`，或末轮 `tool_calls` 无结果时 mask/剥离该 call；**非必做** | 末轮悬空 tool_call 得到一致处理，避免训练时误算 loss | +| **P5** | `dead_loop_filter.py:192-194` | `is_agent_row` 为真则**整行跳过** stuck 检测，agent 恰恰最易死循环 | agent 死循环走 `HardScorer.check_no_repeated_calls`（见 D3） | 覆盖 agent 重复工具调用循环，堵住最大系统性漏检 | +| **P6** | `hard_filter.py` `max_rounds` | 实现是 `len(asst_msgs) > max_rounds`，只数 assistant，注释却写 “user-assistant pairs” | 修正为按 pair 计数或改注释与语义一致 | 轮数过滤阈值语义正确 | +| **P7** | `message_normalizer.py` / `hard_filter.py` / `utils.py` | `tool_calls` 真值判断三处不一致（裸真值 vs `_has_tool_calls` 视 `''`/`'[]'`/`[]` 为空 vs `normalize_tool_calls`） | 全部统一走 `normalize_tool_calls` | 同一数据“是否 agent”判定一致，消除跨 filter 不一致 | +| **P8** | `refuse_filter.py` | 只扫首条 assistant 前 600 字、不读 reasoning，多轮/reasoning 拒答漏检 | 扩到全 assistant + reasoning 字段（可配窗口） | 拒答样本召回上升，减少污染 | +| **P9** | `token_soup.py` | `max_chars>0` 只查头部（cookbook 用 8000），尾部乱码漏检；不扫 reasoning | 全文 + reasoning 扫描或分段抽样 | 乱码样本召回上升 | +| **P10** | `message_sanity.py` `consolidate_system_messages` | 合并 multimodal system 时压成纯字符串，可能丢非 text part | 用 `msg_has_media` 保护多模态 system | 多模态 system 不丢内容 | + +--- + +## 二、架构 / 类设计（拆分与合并） + +| ID | 项 | 判断 | 改动 | 预期收益 | +|----|----|------|------|----------| +| **A1**（高） | `score_filter.py` 9 个类 486 行 | 契约与实现混在一起，加 scorer 就改巨型文件 | 拆为 `score/` 子包：`types.py`(RoundContext/ScoreResult/Scorer) + `scorers.py`(ChrMin/SIFD 轻) + `judge.py`(PassN/Paraphrase 重) + `score_filter.py`(编排) | 开闭原则；轻/重依赖分离；新增 scorer 零侵入 | +| **A2**（中） | `utils.py` | logprob 数学 + 消息格式工具两个无关模块塞一起 | 拆为 `logprob_utils.py` + `message_utils.py` | 降耦合；改 score 逻辑不误碰消息工具 | +| **A3**（高） | `twinkle/preprocessor/base.py:39` | 基类 `__call__` 声明返回 `Dict`，所有子类实际返回 `Tuple[kept, dropped]`，类型契约名存实亡 | 基类改 `-> Tuple[List, List]`；可选分 `Mapper`/`Filter` 语义基类 | 类型检查生效；新人不会照错签名写导致解包崩溃 | +| **A4**（低） | intent 常量位置 | `ScoreFilter` 消费 intent，但常量定义在 `intent_classifier.py`，score 独立后形成跨模块依赖 | intent 常量下沉到轻量 `intents.py` | 为 score 子包独立化铺路 | +| **A5**（高，目标前置） | 统一 `user_data` 标签信封 | 评分/安全/血缘无统一落点；D6↔D7 若代码互调会逼出 DAG | 所有标签走 `user_data` 的 `List[Tuple[str, pack_value(v)]]`（PyArrow 稳定，见 D 节前置）；打标 mapper 写、末尾 filter 读，靠列表顺序解耦 | 去 DAG、统一数据契约；A3 返回契约的自然延伸 | + +**明确不动**（避免过度设计）：`HardFilter`/`RefuseFilter`/`DeadLoopFilter`/`TokenSoupFilter` 保持独立（合并成上帝类只会更糟）；`data_juicer.py` 4 个薄封装保持一文件；`LLMBackend` 三类保持；`MessageNormalizer` 3 个 pass 不拆（有强顺序依赖）；`IntentDetector` 层级设计是全代码最佳，保持。 + +--- + +## 三、功能增删 + +### 建议去掉 / 降级（死代码与过度设计） + +| ID | 项 | 证据 | 改动 | 预期收益 | +|----|----|------|------|----------| +| **R1** | `ScoreFilter` 全家（+4 scorer）**+ `llm_backend.py` 整个文件** | `ScoreFilter` **全库零 active 使用**（仅自身定义 + docs 示例 + `train_cold_start.py:216` 注释态；测试只覆盖底层 `utils` 数学函数）；内含两处死代码 bug —— 原 **P11**（`ParaphraseScorer:452-455` 缺 DP pad，小批量/DP>1 时 `SamplerBackend` raise）、原 **P12**（`utils.py:154` `ifd=exp(-mean_delta)` 是 Superfiltering 差分指数口径而非 Cherry 损失比值，阈值不可互换）。`llm_backend.py`（`LLMBackend`/`OpenAIBackend`/`SamplerBackend`）**唯一消费者就是 `ScoreFilter`**（grep 确认除自身+`__init__` 导出外无他），它专为 score 打分提供 `chat`/`prompt_logprobs`/`prompt_logprobs_ids`/`embeddings` | `ScoreFilter` + `llm_backend.py` 一起移出主包到 `experimental/` 或 `data_selection/`，标记未验证；**启用前**再修 P11（复用 `_pad_batch`）+ 明确 P12 IFD 口径并重标阈值 | 主路径不再拖未验证的重代码 + 未接线的 LLM 后端抽象；bug 修复延后到真正需要时 | +| **R2** | `LLMBackend.embeddings()` | 是 R1 中 `llm_backend.py` 的一部分；preprocessor 侧**零调用者**，`SamplerBackend.embeddings` 直接 raise | 随 R1 一并移出（不单独保留伪抽象） | 去掉未接线接口，等真需要 embedding 去重再加 | +| **R3** | `IntentClassifier` | 产物 `key_rounds/intents` 主消费者是死的 `ScoreFilter`；`dataset_think.py` 只 import 不入 pipeline | 重定位为“标注器”（从不 drop）；短期可移出主 pipeline 省 CPU | 明确职责；省无谓计算 | +| **R4** | `ComplexLogic/Reasoning/UserDissatisfaction` Detector | 仅被 `IntentClassifier.DEFAULT_DETECTORS` 引用，下游死 | 精简 default 到 `ToolCall/Code/Math` | 减少无消费者的启发式维护面 | +| **R5** | LLM 调用抽象统一到 `llm_backup` | preprocessor 里活跃的 LLM 生成需求（summarizer/segment/verifier）**早已全部走 `twinkle_agentic/utils/llm_backup.py`**（置信度路由 student/teacher + 蒸馏数据收集）；只有死代码 `ScoreFilter` 还用独立的 `LLMBackend` | 主路径不再引入独立 `LLMBackend` 抽象，生成类需求统一走 `@llm_backup`。**注意**：`llm_backup` 只提供 chat 生成（返回 content 字符串），**不提供 `prompt_logprobs`/`embeddings`**——IFD/chr_min 类 logprob 数据选择若复活，那部分接口需在 `experimental/` 内单独保留或重写，不能指望 `llm_backup` | 收敛到单一蒸馏路由机制；生成享受 student/teacher 蒸馏；消除重复的推理后端抽象 | + +### 建议增加（真正缺失的清洗能力） + +#### 已列（清洗器层） + +| ID | 项 | 缺口 | 改动 | 预期收益 | +|----|----|------|------|----------| +| **D1**（离线） | 近重复去重 MinHash-LSH/SimHash | `DedupFilter` 只做前缀 md5 精确去重，改一字的近重复全漏 | 扩展 `DedupFilter` 或新增 `NearDupFilter`（`datasketch` 轻依赖）。**限离线批处理阶段**（需全局视图）；实时 per-batch 主路径不启用，否则局部视图导致误杀严重 | 相似轨迹被挡，多样性上升；离线做，避免在线误杀 | +| **D2**（离线） | 基准去污染 decontamination | train/test n-gram overlap **完全没有** | 新增 13-gram 重叠过滤，比对**静态** benchmark n-gram 索引。**限离线**或**只打标不删**，避免实时流误杀正常样本 | 评测不被污染，指标可信 | +| **D4**（中） | 语言识别 langid/fastText | 只有 `cjk_ratio` script 比例，粗糙 | 轻量 langid 过滤 | 中英限定更可靠，混语噪声下降 | +| **D5**（低，可选） | 结构化噪声轮识别 | 现有 heartbeat 靠关键词（对 openclaw/OpenHands 格式已够用，见 P2 撤销）；仅当出现无关键词的结构性噪声轮时才需要 | 极短轮 + 高重复 + embedding 距离（复用 D1 基础设施） | 覆盖无关键词的噪声轮；非当前痛点 | + +> **D3（agent 死循环接线）已废弃**：review 指出 `preprocessor` 与 `verifier` 当前**零互相 import**（grep 确认），让 `DeadLoopFilter` 去 import `HardScorer`（一个 RL reward `Verifier`）会破坏模块边界、且职责串（清洗器 vs 打分器）。正确路径并入 **D6/D7**：新增打标/评分 preprocessor，agent 死循环由其中的确定性 check 覆盖。 + +#### 新增（达成「干净 + 每轮评分」最终目标所需的整段能力） + +> 对标业界标准 agent 数据流程（Llama-3 / DeepSeek-V3 / Nemotron / Tulu-3 / AgentInstruct / ToolBench）。目标五属性映射：无不良信息→D8/D9、无废案→D6、无重复冗余→D1+D6(轨迹内)、无心跳→已有、每轮评分→D7。 + +| ID | 项 | 属性 | 缺口 | 改动 | 预期收益 | +|----|----|------|------|------|----------| +| **D7**（高，核心） | 每轮评分打标 preprocessor（**只打标不过滤**） | 每轮评分 | `verifier`（per-round `HardScorer` + per-segment `RubricVerifier`）+ `aggregation`（round→segment→trajectory）**基础设施现成但未接线**；`aggregation.py:27` 明说编排器 `TrajectoryScorer` 未实现 | **新增 preprocessor**（mapper，从不 drop）：`Segmenter → HardScorer(逐轮) → RubricVerifier(逐段) → aggregation → 分数写回 `user_data``。分数、`score_confidence`、安全标全部作为 `(key, pack_value(v))` 追加进 `user_data`（见架构前置 A5）。`RubricVerifier.score_detail()` 已返回完整 `ScoreDetail`，`__call__(trajectory)` 兼容逐行调用 | 直接产出「每轮评分」的 trajectory；打标与过滤解耦，D6/D8 只读标签 | +| **D7c**（高，核心） | 评分校准探针 + 客观→主观重评（自进化，无人评） | 每轮评分可信度 | 自进化框架**不能靠人评对齐**；未校准的分会系统性放大 judge 偏见 | 用三个**自动**信号合成 per-segment `score_confidence`：①**teacher-student 一致性**（复用 `llm_backup` 已收集的 `(student, teacher, match)`）②**结果锚定**（`HardScorer` 确定性 check 当弱标签探针）③**voting 方差**（`RubricVerifier` 已有 voting，导出方差）。**关键**：当客观（硬 check）与 LLM 主观**不一致**时，**把客观结果注入 rubric 的打分上下文，让 `RubricVerifier` 重新评分**（不是简单降权，是带硬信号修正的二次评分） | 分数可信度自动量化；客观事实纠偏主观判断，闭环收敛；零人工 | +| **D6**（高） | 轨迹成败判定（过滤废案）——**纯读标签 filter** | 无废案 / 轨迹内冗余 | 无 outcome verification：失败/绕圈/工具全错/最终答案错的轨迹留在训练集 | **不自己算分**，只 `user_data_get(row, 'traj_score')` 等标签跟**阈值**比 → 判废案 drop（依赖的是 D7 已写好的**数据标签**，不是 D7 的代码，靠 pipeline 列表顺序保证 D6 在 D7 后）。**阈值先拍默认值，实测回收分布后回调**（不做人评标定） | 废案不进训练集；与打标解耦，无模块依赖 | +| **D8**（高） | 安全/毒性评分（复用 rubric） | 无不良信息 | 只有 `RefuseFilter`（拒答正则）+ 敏感词表，无 toxicity/safety 覆盖暴力/仇恨/成人/越狱成功 | **复用 `RubricVerifier`**：把安全维度作为一组**固定 `RubricItem`**（暴力/仇恨/成人/越狱成功/隐私泄露）注入 stage-2 打分，走现有 `_score_with_voting` + `_aggregate`；低于阈值判不良。**无需新分类器/新依赖** | 安全过滤召回远超敏感词表；与 D7 共用打分基础设施 | +| **D9**（中） | PII 真脱敏（激活现有 Presidio） | 无不良信息 | `PIIPresidioFilter` 存在但曾因**慢**去掉（spaCy NER 是瓶颈） | 加回来但**纯 regex 模式**：现 `IGNORED_ENTITIES` 已忽略全部 NER 实体（PERSON/LOCATION/ORG…），只留 regex 标识符（邮箱/电话/证件/银行卡）→ **可不加载 spaCy**，绕过 NER 瓶颈，速度问题基本消除 | 邮箱/电话/证件等真 PII 脱敏，且不拖慢管线 | +| **D10**（中） | 治理层：provenance（血缘字段） | 可追溯 | 无血缘字段（source/teacher_model/timestamp） | 每条 trajectory 把血缘作为 `(key, pack_value(v))` 写进 `user_data`（蒸馏场景 teacher/student 版本）。**批次归因/数据卡暂缓**（backlog，见「不做」） | 可追溯；对标 Nemotron/Dolma 但先只做血缘字段 | + +> **架构前置 A5（去 DAG 的正解）**：所有评分/安全/血缘标签统一写进 **`user_data` 信封**，把「评分」与「过滤」解耦成「**打标 mapper（D7/D8/D10，从不 drop）+ 末尾纯读标签 filter（D6）**」。这样 D6→D7 是**数据依赖**（D7 写标签、D6 读标签），靠**线性 `QualityPreprocessor` 的列表顺序**保证，**无需 DAG、无模块互相 import**。 +> +> **PyArrow 硬约束**：`user_data` 必须是 **`List[Tuple[str, str]]`**，**不能用 dict**（HF `datasets` 的 PyArrow 后端对异构/嵌套 dict 序列化有问题）。已核实这是仓库现有官方约定 —— `twinkle/data_format/trajectory.py:18-19`（`user_data: List[Tuple[str, str]]`，注释 "PyArrow-stable encoding: each entry is (key, json.dumps(value))"），写用 `pack_value(v)`（JSON 字符串，值可为任意结构但对外恒为 `(str, str)`），读用 `user_data_get(row, key)`。每轮分数可用 `(f'round_{i}_score', pack_value(...))` 或 `('round_scores', pack_value([...]))` 形态。 +> +> 确定性 check 复用：D6/D7/D8 都要 `HardScorer` 的 LLM-free check。可抽到无依赖公共层（如 `twinkle_agentic/agent_checks.py`）供 verifier 与新 preprocessor 各自依赖；用户已确认也可接受 preprocessor→verifier 单向依赖，则公共层后置。 + +### 明确不做（自进化框架的取舍） + +| 项 | 为什么不做 | +|----|-----------| +| 人评校准对齐 | 自进化框架不 scale；改用 D7c 的三信号（teacher 一致性 + 结果锚定 + voting 方差）替代 | +| 批次间质量归因 / 数据集版本 diff | 暂缓（backlog），当前不阻塞可用性 | +| 管线内数据配比 / 分层采样 | 移到**训练时的 sampler** 消费 `user_data` 标签，清洗管线只负责打标 | +| DAG / 阶段化编排引擎 | 用 A5 的「打标 + 末尾读标签」拍平成线性，不引入 DAG | + +--- + +## 四、改动项汇总与优先级 + +| 优先级 | 项 | 类型 | 是否引入依赖 | +|--------|----|------|--------------| +| P0 前置 | **A5** | `user_data` 信封（去 DAG，其余目标项的地基） | 否 | +| P0 必做 | P3 | 数据损坏（丢 reasoning） | 否 | +| P1 目标核心 | D7, D7c, D6, D8 | 每轮评分 / 校准 / 废案 / 安全（接线 verifier） | 否（复用 verifier） | +| P1 高 | P7, A1, A3 | 一致性 / 结构重构 | 否 | +| P2 中 | D9, D10 | PII 脱敏 / 血缘 | presidio(D9) | +| P2 中 | P1, P5, P6, P8, P10, A2, D4 | 语义/漏检/降耦合 | langid(D4) | +| P2 中（仅离线） | D1, D2 | 去重/去污染（防实时误杀） | `datasketch`(D1) | +| P3 低 | A4, R1, R2, R3, R4, R5, D5 | 清理/重定位/增强 | 否 | + +**合计 22 项**：实现问题 6（P1、P3、P5–P10）、架构 5（A1–A5）、功能增删 11（R1–R5 + D1/D2/D4/D5 + D6/D7/D7c/D8/D9/D10）。原 P2/P4 撤销，P11/P12 归入 R1，**D3 废弃**（并入 D6/D7）。 + +--- + +## 五、改动后预期整体收益 + +1. **达成最终目标（干净 + 每轮可信分）**：A5 统一标签信封 → D7 分数写回每轮 → D7c 自动校准 + 客观纠偏主观 → D6 读标签滤废案 → D8 安全 rubric → D9 PII；heartbeat 已有、D1 离线去重。五属性齐活，且分数带 `score_confidence`。 +2. **去 DAG**：A5 把「打标 mapper + 末尾读标签 filter」拍平成线性 `QualityPreprocessor`，D6↔D7 只有数据依赖、零模块互调，无需 DAG 引擎。 +3. **数据正确性**：P3 消除 reasoning 被静默剥离；P1 末轮悬空 tool_call 一致处理。 +4. **复用而非新建**：D7 复用 `verifier`+`aggregation`；D7c 复用 `llm_backup` 一致性 + `RubricVerifier` voting；D8 复用 rubric 固定项；D9 纯 regex 免 spaCy —— **几乎零新依赖**。 +5. **可维护性 / 一致性**：A1–A4 契约清晰、依赖分层；R1–R5 砍死代码；P7 统一 `tool_calls` 判定；D10 血缘可追溯。 + +> 落地顺序建议： +> 1. **A5**（定 `user_data` 标签信封 + list-of-tuple/`pack_value` 约定）→ 所有目标项的地基。 +> 2. **P3**（防丢数据，无依赖）。 +> 3. **D7**（每轮评分打标 mapper，接线 verifier+aggregation，写回 `user_data`）。 +> 4. **D7c**（三信号校准 + 客观→主观重评）→ 让分数可信。 +> 5. **D6 + D8**（读标签滤废案 + 安全 rubric，阈值先拍后测）。 +> 6. **D9**（PII 纯 regex 加回）→ **D10**（血缘字段）。 +> 7. **A1 + A3**（结构地基）→ 其余（P1/P5/P6/P8/P10/D4）视数据分布投入。 +> 8. **D1 + D2** 放离线批处理阶段单独跑；配比放训练 sampler。 diff --git a/src/twinkle_agentic/preprocessor/__init__.py b/src/twinkle_agentic/preprocessor/__init__.py index a69b44392..fb7b359c2 100644 --- a/src/twinkle_agentic/preprocessor/__init__.py +++ b/src/twinkle_agentic/preprocessor/__init__.py @@ -12,14 +12,18 @@ from .dedup_filter import DedupFilter from .hard_filter import HardFilter from .intent_classifier import IntentClassifier -from .llm_backend import LLMBackend, OpenAIBackend, SamplerBackend # noqa: F401 +from .language_filter import LanguageFilter # noqa: F401 from .message_normalizer import MessageNormalizer # noqa: F401 from .message_sanity import MessageSanityFilter from .model_filter import ModelFilter +from .outcome_filter import TrajectoryOutcomeFilter # noqa: F401 from .pii_presidio_filter import PIIPresidioFilter +from .provenance import ProvenanceStamp # noqa: F401 from .refuse_filter import RefuseFilter -from .score_filter import ScoreFilter +from .safety_scorer import SafetyScorer # noqa: F401 +from .structural_noise import StructuralNoiseTagger # noqa: F401 from .token_soup import TokenSoupFilter +from .trajectory_scorer import TrajectoryScorer # noqa: F401 logger = get_logger() diff --git a/src/twinkle_agentic/preprocessor/dead_loop_filter.py b/src/twinkle_agentic/preprocessor/dead_loop_filter.py index 75cf3d00b..3d629723b 100644 --- a/src/twinkle_agentic/preprocessor/dead_loop_filter.py +++ b/src/twinkle_agentic/preprocessor/dead_loop_filter.py @@ -189,10 +189,18 @@ def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: dropped: List[Dict[str, Any]] = [] for row in rows: messages = row.get('messages') or [] - if is_agent_row(messages): - out.append(row) - continue + agent = is_agent_row(messages) asst_msgs = [m for m in messages if isinstance(m, dict) and m.get('role') == 'assistant'] + if agent: + # For agent rows, tool-call loops are caught by the deterministic + # per-round check_no_repeated_calls in TrajectoryScorer (D7) — not + # here — to avoid duplicating loop logic. But agents ALSO emit + # degenerate free-text; run the stuck-text detector on assistant + # turns that carry real text (skip pure tool-call turns whose empty + # content would misfire the detector), instead of skipping the row. + asst_msgs = [m for m in asst_msgs + if msg_content_text(m).strip() + or (m.get('reasoning_content') or m.get('thinking') or '').strip()] if not asst_msgs: out.append(row) continue diff --git a/src/twinkle_agentic/preprocessor/experimental/__init__.py b/src/twinkle_agentic/preprocessor/experimental/__init__.py new file mode 100644 index 000000000..6257664e3 --- /dev/null +++ b/src/twinkle_agentic/preprocessor/experimental/__init__.py @@ -0,0 +1,21 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Experimental / not-yet-wired preprocessor components (AUDIT R1). + +These modules are kept out of the main :mod:`twinkle_agentic.preprocessor` +namespace because they have no active consumer in the shipped pipeline +(``QualityPreprocessor``) and are not exercised by the cookbook or tests: + +- :class:`ScoreFilter` and its scorers (per-round SFT key-round selection). Active + LLM generation in the framework goes through ``twinkle_agentic.utils.llm_backup`` + instead of the local ``LLMBackend`` abstraction. +- :class:`LLMBackend` / :class:`OpenAIBackend` / :class:`SamplerBackend`, which + exclusively serve ``ScoreFilter``. + +Import explicitly from here if you want to experiment with them, e.g.:: + + from twinkle_agentic.preprocessor.experimental import ScoreFilter, SamplerBackend +""" +from .llm_backend import LLMBackend, OpenAIBackend, SamplerBackend # noqa: F401 +from .score_filter import ScoreFilter # noqa: F401 + +__all__ = ['ScoreFilter', 'LLMBackend', 'OpenAIBackend', 'SamplerBackend'] diff --git a/src/twinkle_agentic/preprocessor/llm_backend.py b/src/twinkle_agentic/preprocessor/experimental/llm_backend.py similarity index 100% rename from src/twinkle_agentic/preprocessor/llm_backend.py rename to src/twinkle_agentic/preprocessor/experimental/llm_backend.py diff --git a/src/twinkle_agentic/preprocessor/score_filter.py b/src/twinkle_agentic/preprocessor/experimental/score_filter.py similarity index 99% rename from src/twinkle_agentic/preprocessor/score_filter.py rename to src/twinkle_agentic/preprocessor/experimental/score_filter.py index 8530ab119..228a8e25b 100644 --- a/src/twinkle_agentic/preprocessor/score_filter.py +++ b/src/twinkle_agentic/preprocessor/experimental/score_filter.py @@ -31,7 +31,7 @@ from twinkle.template import Template from twinkle.utils import get_logger from .llm_backend import LLMBackend -from .utils import _chr_min_distinct, _ifd_family_metrics, _lp_to_jsonable, _pad_batch, _to_int_list +from ..utils import _chr_min_distinct, _ifd_family_metrics, _lp_to_jsonable, _pad_batch, _to_int_list logger = get_logger() diff --git a/src/twinkle_agentic/preprocessor/hard_filter.py b/src/twinkle_agentic/preprocessor/hard_filter.py index 043389ca0..d303ebce4 100644 --- a/src/twinkle_agentic/preprocessor/hard_filter.py +++ b/src/twinkle_agentic/preprocessor/hard_filter.py @@ -4,7 +4,7 @@ from typing import Any, Dict, List, Optional, Tuple from twinkle.preprocessor import Preprocessor -from .utils import cjk_ratio, msg_content_text, msg_has_media +from .utils import cjk_ratio, is_agent_row, msg_content_text, msg_has_media, normalize_tool_calls # ── Language detection ──────────────────────────────────────────────────────── @@ -82,14 +82,8 @@ def _has_tool_calls(msg: Dict[str, Any]) -> bool: - """Truthy ``tool_calls`` excluding the empty-array sentinels '' / '[]' / [].""" - tc = msg.get('tool_calls') - if not tc: - return False - if isinstance(tc, str): - s = tc.strip() - return bool(s) and s != '[]' - return bool(tc) + """True iff the message carries real tool calls (unified via normalize_tool_calls).""" + return normalize_tool_calls(msg) is not None def _is_simple_query(text: str, min_user_chars: int, min_user_chars_cjk: int) -> bool: @@ -129,6 +123,7 @@ def __init__( max_chars_per_round: Optional[int] = None, max_total_chars: Optional[int] = None, max_rounds: Optional[int] = None, + agent_max_rounds: Optional[int] = None, ) -> None: super().__init__() self._min_user_chars = min_user_chars @@ -142,6 +137,13 @@ def __init__( self._max_chars_per_round = max_chars_per_round self._max_total_chars = max_total_chars self._max_rounds = max_rounds + # Agent trajectories legitimately run many tool-calling rounds and are the + # highest-value distillation data, so the plain ``max_rounds`` cap (meant + # for shallow chit-chat) must not clip them. They get their own, far higher + # ceiling that still catches pathological runaway loops. ``None`` disables + # the cap for agent rows entirely; if unset it defaults to a wide multiple + # of ``max_rounds``. + self._agent_max_rounds = agent_max_rounds def _drop_reason(self, row: Dict[str, Any], messages: List[Any]) -> Optional[str]: """Apply rules in order; return first matching drop_reason, or None to keep.""" @@ -194,9 +196,19 @@ def _drop_reason(self, row: Dict[str, Any], messages: List[Any]) -> Optional[str if total > self._max_total_chars: return 'total_too_long' - # Rule 7: max rounds (user-assistant pairs). - if self._max_rounds and len(asst_msgs) > self._max_rounds: - return 'too_many_rounds' + # Rule 7: max rounds (user-assistant pairs). Count complete pairs, not raw + # assistant turns — an agent turn may emit several assistant messages + # (tool_call + follow-up) that are one logical round. Agent traces use a + # separate, higher ceiling (or none) so long tool-calling loops survive. + if self._max_rounds: + rounds = min(len(user_msgs), len(asst_msgs)) + if is_agent_row(messages): + cap = (self._agent_max_rounds if self._agent_max_rounds is not None + else self._max_rounds * 10) + else: + cap = self._max_rounds + if cap is not None and rounds > cap: + return 'too_many_rounds' return None diff --git a/src/twinkle_agentic/preprocessor/intent_classifier.py b/src/twinkle_agentic/preprocessor/intent_classifier.py index 7dde971b1..7cb1dc5d9 100644 --- a/src/twinkle_agentic/preprocessor/intent_classifier.py +++ b/src/twinkle_agentic/preprocessor/intent_classifier.py @@ -13,14 +13,10 @@ # Reasoning block regex covers both and forms. _THINK_BLOCK_RE = re.compile(r'(.*?)', re.DOTALL | re.IGNORECASE) -# ── Intent categories ───────────────────────────────────────────────────────── -INTENT_TOOL_CALL = 'tool_call' -INTENT_CODE = 'code' -INTENT_MATH = 'math' -INTENT_COMPLEX_LOGIC = 'complex_logic' -INTENT_REASONING = 'reasoning' -INTENT_USER_DISSATISFACTION = 'user_dissatisfaction' -INTENT_OTHER = 'other' +# ── Intent categories (canonical vocabulary lives in intents.py; re-exported) ── +from .intents import (INTENT_CODE, INTENT_COMPLEX_LOGIC, # noqa: F401,E402 + INTENT_MATH, INTENT_OTHER, INTENT_REASONING, + INTENT_TOOL_CALL, INTENT_USER_DISSATISFACTION) # ── Heuristic patterns ──────────────────────────────────────────────────────── _CODE_BLOCK_RE = re.compile(r'```[\s\S]{10,}?```') @@ -337,6 +333,10 @@ class IntentClassifier(Preprocessor): Pure-heuristic, no LLM. Each intent is a pluggable :class:`IntentDetector`; pass ``detectors=[...]`` to extend or override. + R3: this is an *annotator* — by default it never drops rows + (``drop_no_key_rounds=False``); rows with no detected key round are simply + tagged ``INTENT_OTHER``. Set ``drop_no_key_rounds=True`` to also filter. + Annotates per row:: row['intent'] # primary intent string @@ -344,20 +344,21 @@ class IntentClassifier(Preprocessor): ('intents', dict[str, str])] # per-round intent """ + # R4: default to the detectors with a live downstream consumer. The heavier + # heuristics (ComplexLogic / Reasoning / UserDissatisfaction) are kept as + # importable classes but dropped from the default set — their outputs had no + # active consumer. Pass ``detectors=[...]`` to re-enable them. DEFAULT_DETECTORS: List[IntentDetector] = [ ToolCallDetector(), CodeDetector(), MathDetector(), - ComplexLogicDetector(), - ReasoningDetector(), - UserDissatisfactionDetector(), ] def __init__( self, detectors: Optional[List[IntentDetector]] = None, intent_field: str = 'intent', - drop_no_key_rounds: bool = True, + drop_no_key_rounds: bool = False, ) -> None: super().__init__() self._intent_field = intent_field diff --git a/src/twinkle_agentic/preprocessor/intents.py b/src/twinkle_agentic/preprocessor/intents.py new file mode 100644 index 000000000..d945e59af --- /dev/null +++ b/src/twinkle_agentic/preprocessor/intents.py @@ -0,0 +1,25 @@ +"""Intent category constants (AUDIT A4). + +Sunk out of ``intent_classifier.py`` into this dependency-free module so any +consumer (e.g. the experimental log-prob scorers, or downstream sampling that +reads ``intents`` labels) can reference the vocabulary without importing the +heavier classifier + its regex detectors. +""" + +INTENT_TOOL_CALL = 'tool_call' +INTENT_CODE = 'code' +INTENT_MATH = 'math' +INTENT_COMPLEX_LOGIC = 'complex_logic' +INTENT_REASONING = 'reasoning' +INTENT_USER_DISSATISFACTION = 'user_dissatisfaction' +INTENT_OTHER = 'other' + +ALL_INTENTS = ( + INTENT_TOOL_CALL, + INTENT_CODE, + INTENT_MATH, + INTENT_COMPLEX_LOGIC, + INTENT_REASONING, + INTENT_USER_DISSATISFACTION, + INTENT_OTHER, +) diff --git a/src/twinkle_agentic/preprocessor/label_schema.py b/src/twinkle_agentic/preprocessor/label_schema.py new file mode 100644 index 000000000..109f627d2 --- /dev/null +++ b/src/twinkle_agentic/preprocessor/label_schema.py @@ -0,0 +1,102 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Unified ``user_data`` label envelope (AUDIT A5). + +All scoring / safety / provenance annotations produced by the pipeline are +written into a trajectory's ``user_data`` as ``(key, pack_value(value))`` pairs. +This is the single data contract that lets us decouple *tagging* (mappers that +never drop) from *filtering* (a tail filter that only reads tags), so the whole +pipeline stays a linear ``QualityPreprocessor`` list — no DAG, no cross-module +imports between a filter and the verifier it depends on. + +PyArrow hard constraint +----------------------- +``user_data`` MUST be a ``List[Tuple[str, str]]`` (see +``twinkle/data_format/trajectory.py``). We NEVER put a bare ``dict`` in a row +column: HF ``datasets``' PyArrow backend cannot stably serialize +heterogeneous / nested dicts. Structured values are JSON-encoded to a single +string via :func:`pack_value`; on read :func:`user_data_get` JSON-decodes them. + +Keep this module dependency-light: it only knows the *keys* and thin get/set +helpers, so both preprocessors and (optionally) other modules can share it +without pulling in verifier/segment code. +""" +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Tuple + +from twinkle.data_format import pack_value, user_data_get + +# --------------------------------------------------------------------------- +# Canonical label keys +# --------------------------------------------------------------------------- +# Per-round hard scores, aligned to assistant/round order within the trajectory. +# Value: List[float] in [0, 1]. +KEY_ROUND_SCORES = 'round_scores' +# Per-round gated flags (a critical hard check zeroed the round). Value: List[bool]. +KEY_ROUND_GATED = 'round_gated' + +# Per-segment fused scores. Value: List[float] in [0, 1]. +KEY_SEGMENT_SCORES = 'segment_scores' +# Per-segment score confidence (D7c calibration). Value: List[float] in [0, 1]. +KEY_SEGMENT_CONFIDENCE = 'segment_confidence' + +# Whole-trajectory fused score in [0, 1] and its discrete level. +KEY_TRAJ_SCORE = 'traj_score' +KEY_TRAJ_LEVEL = 'traj_level' +# Aggregate confidence for the trajectory score (D7c). Value: float in [0, 1]. +KEY_TRAJ_CONFIDENCE = 'traj_confidence' + +# Safety score in [0, 1] (D8, higher = safer) + boolean unsafe flag. +KEY_SAFETY_SCORE = 'safety_score' +KEY_SAFETY_UNSAFE = 'safety_unsafe' + +# Provenance blob (D10): dict-like value JSON-encoded (source/teacher/student/ts). +KEY_PROVENANCE = 'provenance' + +# Free-form scoring metadata (short-circuit stats, per-check breakdown, etc.). +KEY_SCORE_META = 'score_meta' + + +# --------------------------------------------------------------------------- +# thin get / set helpers over the (key, pack_value) envelope +# --------------------------------------------------------------------------- +def get_user_data(row: Dict[str, Any]) -> List[Tuple[str, str]]: + """Return the row's ``user_data`` as a list (never a dict), defaulting to [].""" + ud = row.get('user_data') + if ud is None: + return [] + if isinstance(ud, list): + return ud + # Be forgiving of a stray dict (e.g. hand-authored rows) — flatten to pairs. + if isinstance(ud, dict): + return [(k, v if isinstance(v, str) else pack_value(v)) for k, v in ud.items()] + return [] + + +def get_label(row: Dict[str, Any], key: str, default: Any = None) -> Any: + """Read+JSON-decode the first label matching ``key`` from ``row['user_data']``.""" + return user_data_get(get_user_data(row), key, default) + + +def set_labels(row: Dict[str, Any], updates: Dict[str, Any]) -> Dict[str, Any]: + """Return a shallow-copied row with ``updates`` merged into ``user_data``. + + Existing entries for the same keys are replaced (last-write-wins), preserving + the original order for untouched keys. Values are packed with :func:`pack_value` + so the column stays ``List[Tuple[str, str]]`` (PyArrow-stable). + """ + if not updates: + return row + existing = get_user_data(row) + replace = set(updates.keys()) + merged: List[Tuple[str, str]] = [(k, v) for (k, v) in existing if k not in replace] + for k, v in updates.items(): + merged.append((k, pack_value(v))) + new_row = dict(row) + new_row['user_data'] = merged + return new_row + + +def set_label(row: Dict[str, Any], key: str, value: Any) -> Dict[str, Any]: + """Convenience: set a single label.""" + return set_labels(row, {key: value}) diff --git a/src/twinkle_agentic/preprocessor/language_filter.py b/src/twinkle_agentic/preprocessor/language_filter.py new file mode 100644 index 000000000..fc68d90b6 --- /dev/null +++ b/src/twinkle_agentic/preprocessor/language_filter.py @@ -0,0 +1,109 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Language-identification filter (AUDIT D4). + +Keeps only rows whose user-facing language is in an allow-list. Uses ``langid`` +when installed (proper LID over 97 languages); otherwise degrades gracefully to +a script-ratio heuristic (CJK vs Latin) so the filter is usable with zero extra +dependencies — just coarser. This complements the existing ``cjk_ratio`` checks +in :class:`HardFilter`, which only measure script mix, not language. + +The language is judged from the concatenated user turns (the request defines the +expected response language; assistant text can legitimately quote other +languages, e.g. code or translations). +""" +from __future__ import annotations + +import re +from typing import Any, Dict, Optional, Sequence + +from twinkle.preprocessor import Filter +from twinkle.utils import get_logger + +from .message_utils import cjk_ratio, msg_content_text + +logger = get_logger() + +# Injected scaffolding that is NOT the user's own request and would skew language +# detection (usually English system boilerplate wrapping a non-English query, or +# vice versa). Stripped before LID so we judge the real user text. +_INJECTION_BLOCK_RE = re.compile( + r'<(system-reminder|system_reminder|system|instructions?|context|' + r'important_instructions|env|environment|tools?)\b[^>]*>.*?', + re.DOTALL | re.IGNORECASE, +) +# Self-closing / unmatched openers of the same tags (defensive). +_INJECTION_TAG_RE = re.compile( + r']*/?>', + re.IGNORECASE, +) + + +def _strip_injections(text: str) -> str: + """Remove injected system-scaffolding blocks so LID sees the real user text.""" + text = _INJECTION_BLOCK_RE.sub(' ', text) + text = _INJECTION_TAG_RE.sub(' ', text) + return text.strip() + + +class LanguageFilter(Filter): + """Keep rows whose detected user language is allowed. + + Args: + allowed: allowed ISO 639-1 codes (e.g. ``('en', 'zh')``). + min_chars: skip detection (keep) for user text shorter than this — LID is + unreliable on very short strings. + cjk_threshold: fallback heuristic boundary; user text with CJK ratio above + this is treated as ``zh``, else ``en``. Only used when ``langid`` is absent. + keep_undetected: keep rows where language can't be determined. Default True + (fail-open) so the filter never silently deletes ambiguous data. + """ + + def __init__( + self, + allowed: Sequence[str] = ('en', 'zh'), + *, + min_chars: int = 20, + cjk_threshold: float = 0.15, + keep_undetected: bool = True, + ): + self.allowed = {a.lower() for a in allowed} + self.min_chars = int(min_chars) + self.cjk_threshold = float(cjk_threshold) + self.keep_undetected = bool(keep_undetected) + self._identifier = self._load_langid() + if self._identifier is None: + logger.info('[LanguageFilter] langid not installed; using CJK/Latin script heuristic.') + + @staticmethod + def _load_langid(): + try: + from langid.langid import LanguageIdentifier, model + return LanguageIdentifier.from_modelstring(model, norm_probs=True) + except Exception: + return None + + def _user_text(self, row: Dict[str, Any]) -> str: + messages = row.get('messages') or [] + parts = [_strip_injections(msg_content_text(m)) for m in messages + if isinstance(m, dict) and m.get('role') == 'user'] + return '\n'.join(p for p in parts if p).strip() + + def _detect(self, text: str) -> Optional[str]: + if self._identifier is not None: + try: + lang, _prob = self._identifier.classify(text) + return lang + except Exception: + return None + # heuristic fallback: CJK ratio -> zh, else en + return 'zh' if cjk_ratio(text) > self.cjk_threshold else 'en' + + def keep(self, row: Dict[str, Any]) -> bool: + text = self._user_text(row) + if len(text) < self.min_chars: + return True # too short to judge reliably + lang = self._detect(text) + if lang is None: + return self.keep_undetected + return lang.lower() in self.allowed diff --git a/src/twinkle_agentic/preprocessor/logprob_utils.py b/src/twinkle_agentic/preprocessor/logprob_utils.py new file mode 100644 index 000000000..5f86e2865 --- /dev/null +++ b/src/twinkle_agentic/preprocessor/logprob_utils.py @@ -0,0 +1,231 @@ +"""Log-probability data-selection math (IFD / S-IFD / chr_min). + +Split out of ``utils.py`` (AUDIT A2): these helpers are consumed only by the +log-prob based scorers (the experimental ``ScoreFilter`` family). Keeping them +separate from the message-format utilities means editing scoring math never +risks touching the message helpers used across every active cleaning step. +""" +import math +from typing import Any, Dict, List, Optional, Set, Tuple + + +def _extract_logprob(lp, token_id: Optional[int] = None) -> Optional[float]: + if lp is None: + return None + if isinstance(lp, (int, float)): + return float(lp) + if not isinstance(lp, dict): + return None + # vLLM with prompt_logprobs=1 returns top-1 PLUS actual token if they differ; + # actual is appended LAST, so iter-first picks the wrong (top-1) one. + entry = None + if token_id is not None: + entry = lp.get(token_id) + if entry is None: + entry = lp.get(str(token_id)) + if entry is None: + entry = next(iter(lp.values()), None) + if entry is None: + return None + if hasattr(entry, 'logprob'): + return float(entry.logprob) + if isinstance(entry, dict): + v = entry.get('logprob') + return float(v) if v is not None else None + if isinstance(entry, (int, float)): + return float(entry) + return None + + +def _to_int_list(x) -> List[int]: + if hasattr(x, 'tolist'): + return x.tolist() + return list(x) + + +def _chr_min_distinct( + cond_lp: List, + asst_lp: List, + cond_ids: List[int], + asst_ids: List[int], + n_prompt: int, + exclude_ids: Optional[Set[int]] = None, +) -> Optional[float]: + """chr_dist_min_pos: fraction of distinct asst-token ids whose + per-occurrence min(cond_lp - asst_lp) is strictly positive.""" + if not asst_lp or not cond_lp or not asst_ids: + return None + n_a = min(len(asst_lp), len(asst_ids)) + n_c = len(cond_lp) + by_tok: Dict[int, List[float]] = {} + for i in range(n_a): + ci = n_prompt + i + if ci >= n_c: + break + tid = asst_ids[i] + if tid is None: + continue + if exclude_ids is not None and int(tid) in exclude_ids: + continue + a = _extract_logprob(asst_lp[i], tid) + c_tok = cond_ids[ci] if ci < len(cond_ids) else None + c = _extract_logprob(cond_lp[ci], c_tok) + if a is None or c is None: + continue + by_tok.setdefault(int(tid), []).append(c - a) + if not by_tok: + return None + pos = sum(1 for diffs in by_tok.values() if min(diffs) > 0) + return pos / len(by_tok) + + +def _chr_min_weighted( + cond_lp: List, + asst_lp: List, + cond_ids: List[int], + asst_ids: List[int], + n_prompt: int, +) -> Optional[float]: + """Magnitude-weighted chr_min: each distinct token contributes |min_delta| + as weight; returns sum(pos_weights) / sum(all_weights).""" + if not asst_lp or not cond_lp or not asst_ids: + return None + n_a = min(len(asst_lp), len(asst_ids)) + n_c = len(cond_lp) + by_tok: Dict[int, List[float]] = {} + for i in range(n_a): + ci = n_prompt + i + if ci >= n_c: + break + tid = asst_ids[i] + if tid is None: + continue + a = _extract_logprob(asst_lp[i], tid) + c_tok = cond_ids[ci] if ci < len(cond_ids) else None + c = _extract_logprob(cond_lp[ci], c_tok) + if a is None or c is None: + continue + by_tok.setdefault(int(tid), []).append(c - a) + if not by_tok: + return None + total_w = 0.0 + pos_w = 0.0 + for diffs in by_tok.values(): + md = min(diffs) + w = abs(md) + total_w += w + if md > 0: + pos_w += w + if total_w == 0: + return None + return pos_w / total_w + + +def _ifd_family_metrics( + cond_lp: List, + asst_lp: List, + cond_ids: List[int], + asst_ids: List[int], + n_prompt: int, +) -> Dict[str, Any]: + """IFD (Cherry-LLM) and S-IFD-{50,75} (T-SHIRT) for one round.""" + if not asst_lp or not cond_lp or not asst_ids: + return {} + n_a = min(len(asst_lp), len(asst_ids)) + n_c = len(cond_lp) + deltas: List[float] = [] + for i in range(n_a): + ci = n_prompt + i + if ci >= n_c: + break + tid = asst_ids[i] + if tid is None: + continue + a = _extract_logprob(asst_lp[i], tid) + c_tok = cond_ids[ci] if ci < len(cond_ids) else None + c = _extract_logprob(cond_lp[ci], c_tok) + if a is None or c is None: + continue + deltas.append(c - a) + if not deltas: + return {} + n = len(deltas) + mean_delta = sum(deltas) / n + out: Dict[str, Any] = { + 'n_tokens': n, + 'mean_delta': mean_delta, + 'ifd': math.exp(-mean_delta), + } + abs_sorted = sorted(range(n), key=lambda i: abs(deltas[i]), reverse=True) + for k_pct in (50, 75): + keep = max(1, int(round(n * k_pct / 100))) + sub = [deltas[i] for i in abs_sorted[:keep]] + out[f's_ifd_{k_pct}'] = math.exp(-sum(sub) / len(sub)) + return out + + +def _mean_logprob_delta( + cond_lp: List, + asst_lp: List, + cond_ids: List[int], + asst_ids: List[int], + n_prompt: int, +) -> Optional[float]: + """Mean per-token (cond_lp - asst_lp) over the response span.""" + if not asst_lp or not cond_lp or not asst_ids: + return None + n_a = min(len(asst_lp), len(asst_ids)) + n_c = len(cond_lp) + deltas: List[float] = [] + for i in range(n_a): + ci = n_prompt + i + if ci >= n_c: + break + tid = asst_ids[i] + if tid is None: + continue + a = _extract_logprob(asst_lp[i], tid) + c_tok = cond_ids[ci] if ci < len(cond_ids) else None + c = _extract_logprob(cond_lp[ci], c_tok) + if a is None or c is None: + continue + deltas.append(c - a) + if not deltas: + return None + return sum(deltas) / len(deltas) + + +def _lp_to_jsonable(lp_list): + """Convert per-position prompt_logprobs into JSON-safe form.""" + out = [] + for lp in (lp_list or []): + if lp is None: + out.append(None) + continue + if isinstance(lp, (int, float)): + out.append(float(lp)) + continue + if not isinstance(lp, dict): + out.append(repr(lp)) + continue + d = {} + for k, v in lp.items(): + if hasattr(v, 'logprob'): + d[str(k)] = { + 'logprob': float(v.logprob), + 'rank': getattr(v, 'rank', None), + 'decoded': getattr(v, 'decoded_token', None) + } + elif isinstance(v, dict): + d[str(k)] = v + else: + d[str(k)] = repr(v) + out.append(d) + return out + + +def _pad_batch(batch: List[List[int]], floor: int) -> Tuple[List[List[int]], int]: + n = len(batch) + if n >= floor or not batch: + return batch, n + return list(batch) + [batch[-1]] * (floor - n), n diff --git a/src/twinkle_agentic/preprocessor/message_normalizer.py b/src/twinkle_agentic/preprocessor/message_normalizer.py index d3074a565..1472bdfcf 100644 --- a/src/twinkle_agentic/preprocessor/message_normalizer.py +++ b/src/twinkle_agentic/preprocessor/message_normalizer.py @@ -21,7 +21,7 @@ from twinkle.preprocessor import Preprocessor from twinkle.template.tools import ToolCallRegistry -from .utils import msg_content_text, msg_has_media +from .utils import msg_content_text, msg_has_media, normalize_tool_calls # IGNORECASE absorbs every variant ("Read HEARTBEAT.md", "HEARTBEAT_OK", # "duplicate heartbeat", etc.) under the single token "heartbeat". @@ -46,7 +46,7 @@ def _strip_heartbeat(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: if role == 'user' and _HEARTBEAT_USER_RE.search(text): skip_next_assistant = True continue - if role == 'assistant' and not m.get('tool_calls'): + if role == 'assistant' and normalize_tool_calls(m) is None: if skip_next_assistant or _HEARTBEAT_ASST_RE.search(text): skip_next_assistant = False continue @@ -90,12 +90,16 @@ def _normalize_tool_calls(messages: List[Dict[str, Any]]) -> List[Dict[str, Any] 'arguments': json.dumps(args, ensure_ascii=False) if isinstance(args, dict) else str(args), }, }) - out.append({ + # Preserve every original field (reasoning_content / thinking / name / + # finish_reason / ...) and only override what the rewrite changes. + rebuilt = dict(msg) + rebuilt.update({ 'role': 'assistant', 'content': parser.clean(text), 'tool_calls': json.dumps(tc_list, ensure_ascii=False), 'tool_call_id': '', }) + out.append(rebuilt) # Consume following user messages as tool results — one per tool call. j = i + 1 @@ -128,7 +132,7 @@ def _normalize_tool_calls(messages: List[Dict[str, Any]]) -> List[Dict[str, Any] def _is_atomic(msg: Dict[str, Any]) -> bool: """Atomic = never merge: tool results + assistant turns carrying tool_calls.""" role = msg.get('role', '') - return role == 'tool' or (role == 'assistant' and msg.get('tool_calls')) + return role == 'tool' or (role == 'assistant' and normalize_tool_calls(msg) is not None) def _is_blank_content(msg: Dict[str, Any]) -> bool: diff --git a/src/twinkle_agentic/preprocessor/message_sanity.py b/src/twinkle_agentic/preprocessor/message_sanity.py index 6001e1d7e..2014976e5 100644 --- a/src/twinkle_agentic/preprocessor/message_sanity.py +++ b/src/twinkle_agentic/preprocessor/message_sanity.py @@ -31,6 +31,7 @@ def consolidate_system_messages(messages: List[Dict[str, Any]]) -> List[Dict[str misplaced = any(isinstance(m, dict) and m.get('role') == 'system' and i != 0 for i, m in enumerate(messages)) if sys_count <= 1 and not misplaced: return messages + sys_msgs: List[Dict[str, Any]] = [] sys_chunks: List[str] = [] rest: List[Dict[str, Any]] = [] template: Optional[Dict[str, Any]] = None @@ -38,11 +39,24 @@ def consolidate_system_messages(messages: List[Dict[str, Any]]) -> List[Dict[str if isinstance(m, dict) and m.get('role') == 'system': if template is None: template = m + sys_msgs.append(m) text = msg_content_text(m).strip() if text: sys_chunks.append(text) else: rest.append(m) + # A multimodal system message must not be flattened to a joined string — that + # would drop image/audio parts. Preserve list content by concatenating the + # original content parts instead. + if any(msg_has_media(m) for m in sys_msgs): + merged_parts: List[Any] = [] + for m in sys_msgs: + content = m.get('content') + if isinstance(content, list): + merged_parts.extend(content) + elif isinstance(content, str) and content.strip(): + merged_parts.append({'type': 'text', 'text': content}) + return [dict(template, content=merged_parts)] + rest return [dict(template, content='\n\n'.join(sys_chunks))] + rest @@ -327,6 +341,10 @@ def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: dropped.append(dict(row, drop_reason='no_assistant')) continue row = dict(row, messages=messages) + # Trimming can drop the trailing tool round, so re-derive is_agent + # on the trimmed messages — otherwise agent-only checks may run + # against a now non-agent (or vice-versa) conversation. + is_agent = is_agent_row(messages) reason = self._run_checks(messages, is_agent) if reason is None: diff --git a/src/twinkle_agentic/preprocessor/message_utils.py b/src/twinkle_agentic/preprocessor/message_utils.py new file mode 100644 index 000000000..3ee95b1eb --- /dev/null +++ b/src/twinkle_agentic/preprocessor/message_utils.py @@ -0,0 +1,137 @@ +"""Message-format utilities shared across active preprocessor steps. + +Split out of ``utils.py`` (AUDIT A2): content projection, tool-call +normalization, CJK ratio, sensitive-word regex, and agent-row detection. These +are the helpers every cleaning step depends on, independent of the log-prob +scoring math (see :mod:`logprob_utils`). +""" +import json +import os +import re +from typing import Any, Dict, List, Optional, Set + + +def msg_content_text(msg: Dict[str, Any]) -> str: + """Extract plain text from a message's content (str | list | dict).""" + c = msg.get('content') + if isinstance(c, str): + return c + if isinstance(c, list): + return ' '.join(p.get('text', '') for p in c if isinstance(p, dict) and p.get('type') == 'text') + if isinstance(c, dict) and c.get('type') == 'text': + return c.get('text', '') + return '' + + +def msg_has_media(msg: Dict[str, Any]) -> bool: + """True if message content contains non-text parts (image/audio/video).""" + c = msg.get('content') + return isinstance(c, list) and any(isinstance(p, dict) and p.get('type') not in ('text', None) for p in c) + + +def msg_has_payload(msg: Dict[str, Any]) -> bool: + """True if a message carries any substantive payload (text, tool_calls, reasoning, or media).""" + return bool( + msg_content_text(msg).strip() or msg.get('tool_calls') or msg.get('reasoning_content') or msg.get('thinking') + or msg_has_media(msg)) + + +_CJK_RE = re.compile(r'[\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\uac00-\ud7a3]') + + +def normalize_tool_calls(msg: Dict[str, Any]) -> Optional[List[Any]]: + """Return ``tool_calls`` as a list of dicts, handling PyArrow/HF serialization artifacts.""" + tcs = msg.get('tool_calls') + if isinstance(tcs, str): + s = tcs.strip() + if not s: + return None + try: + decoded = json.loads(s) + except (json.JSONDecodeError, ValueError): + return None + if not isinstance(decoded, list) or not decoded: + return None + tcs = decoded + if not isinstance(tcs, list) or not tcs: + return None + result = [] + for tc in tcs: + if isinstance(tc, str): + try: + tc = json.loads(tc) + except (json.JSONDecodeError, ValueError): + return None + if not isinstance(tc, dict): + return None + func = tc.get('function') + if isinstance(func, str): + try: + func = json.loads(func) + except (json.JSONDecodeError, ValueError): + return None + tc = dict(tc, function=func) + result.append(tc) + return result + + +CJK_CHARS_RE = _CJK_RE + + +def cjk_ratio(text: str) -> float: + """Fraction of non-whitespace characters that are CJK.""" + chars = text.replace(' ', '').replace('\n', '').replace('\t', '') + if not chars: + return 0.0 + return len(CJK_CHARS_RE.findall(chars)) / len(chars) + + +def load_sensitive_words(path: Optional[str]) -> Set[str]: + """Load from external file (one word per line). Blank lines and #-comments ignored.""" + if not path or not os.path.isfile(path): + return set() + words: Set[str] = set() + with open(path, encoding='utf-8') as f: + for line in f: + line = line.strip() + if line and not line.startswith('#'): + words.add(line) + return words + + +def build_sensitive_regex(words: Set[str]) -> Optional['re.Pattern']: + """Build a compiled regex from a set of words. Returns None if empty.""" + if not words: + return None + cjk_words = [] + latin_words = [] + cjk_re = re.compile(r'[\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\uac00-\ud7a3]') + for w in sorted(words): + if cjk_re.search(w): + cjk_words.append(re.escape(w)) + else: + latin_words.append(re.escape(w)) + parts = [] + if latin_words: + parts.append(r'\b(' + '|'.join(latin_words) + r')\b') + if cjk_words: + parts.append('(' + '|'.join(cjk_words) + ')') + return re.compile('|'.join(parts), re.IGNORECASE) + + +def is_agent_row(messages) -> bool: + """Return True if the conversation contains tool interactions (agent trace). + + After MessageNormalizer runs, all non-standard formats are already converted + to standard tool_calls / role=tool — so checking those two signals suffices. + """ + if not isinstance(messages, list): + return False + for m in messages: + if not isinstance(m, dict): + continue + if m.get('role') == 'tool': + return True + if normalize_tool_calls(m): + return True + return False diff --git a/src/twinkle_agentic/preprocessor/model_filter.py b/src/twinkle_agentic/preprocessor/model_filter.py index fe238b1ed..f651f7905 100644 --- a/src/twinkle_agentic/preprocessor/model_filter.py +++ b/src/twinkle_agentic/preprocessor/model_filter.py @@ -1,7 +1,7 @@ import re -from typing import Any, Dict, List, Optional, Sequence, Tuple +from typing import Any, Dict, Optional, Sequence -from twinkle.preprocessor import Preprocessor +from twinkle.preprocessor import Filter # Each entry is the discriminating prefix only; a shared variant tail is appended uniformly # so suffixes like -Instruct, -Thinking-2507, -Distill-Qwen-7B, -Air are accepted everywhere. @@ -21,7 +21,7 @@ _VARIANT_TAIL = r'[-\w.]*' -class ModelFilter(Preprocessor): +class ModelFilter(Filter): """Keep only rows whose model_id matches an allowed family (case-insensitive).""" def __init__(self, patterns: Optional[Sequence[str]] = None, field: str = 'model_id'): @@ -29,12 +29,5 @@ def __init__(self, patterns: Optional[Sequence[str]] = None, field: str = 'model pats = patterns if patterns is not None else _DEFAULT_PATTERNS self._re = re.compile('|'.join(f'(?:{p}{_VARIANT_TAIL})' for p in pats), re.IGNORECASE) - def __call__(self, rows: List[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: - rows = self.map_col_to_row(rows) - kept, dropped = [], [] - for r in rows: - if self._re.fullmatch(r.get(self._field) or ''): - kept.append(r) - else: - dropped.append(dict(r, drop_reason='model_not_allowed')) - return kept, dropped + def keep(self, row: Dict[str, Any]) -> bool: + return bool(self._re.fullmatch(row.get(self._field) or '')) diff --git a/src/twinkle_agentic/preprocessor/offline/__init__.py b/src/twinkle_agentic/preprocessor/offline/__init__.py new file mode 100644 index 000000000..0a14834a6 --- /dev/null +++ b/src/twinkle_agentic/preprocessor/offline/__init__.py @@ -0,0 +1,24 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Offline batch-only preprocessors (AUDIT D1 / D2). + +These steps require a GLOBAL view of the dataset and must NOT be dropped into the +per-batch :class:`~twinkle_agentic.preprocessor.QualityPreprocessor` pipeline: + +- :class:`NearDupFilter` (D1) — MinHash-LSH near-duplicate removal; per-batch use + would only compare within a batch, causing severe false negatives. +- :class:`Decontaminator` (D2) — benchmark n-gram overlap removal against a static + index; kept out of the real-time path to avoid false-positive deletions + (defaults to a safe ``'tag'``-friendly design). + +They are deliberately kept out of the main package namespace. Import explicitly:: + + from twinkle_agentic.preprocessor.offline import NearDupFilter, Decontaminator + from twinkle_agentic.preprocessor.offline import build_benchmark_index + +Usage: materialize the dataset to ``List[Dict]``, run these once, then re-wrap +the kept rows before/after the streaming QualityPreprocessor pipeline. +""" +from .decontaminate import Decontaminator, build_benchmark_index # noqa: F401 +from .near_dedup import NearDupFilter # noqa: F401 + +__all__ = ['NearDupFilter', 'Decontaminator', 'build_benchmark_index'] diff --git a/src/twinkle_agentic/preprocessor/offline/decontaminate.py b/src/twinkle_agentic/preprocessor/offline/decontaminate.py new file mode 100644 index 000000000..7ec9f7572 --- /dev/null +++ b/src/twinkle_agentic/preprocessor/offline/decontaminate.py @@ -0,0 +1,110 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Benchmark decontamination via n-gram overlap (AUDIT D2) — OFFLINE ONLY. + +Removes (or tags) training rows that overlap with evaluation benchmarks, so +reported metrics aren't inflated by leakage. Follows the standard 13-gram +overlap recipe (GPT-3 / Llama / Dolma): build an n-gram set from the benchmark +texts once, then flag any row whose text shares an n-gram with it. + +OFFLINE CONTRACT: the benchmark index is static and global; build it once and +reuse across the whole dataset. This is not a per-batch pipeline step — but +unlike near-dup it *is* embarrassingly parallel per row, so it can also run as a +standalone batch pass. Default mode ``'drop'`` removes contaminated rows; +``'tag'`` keeps them and only records a ``contaminated`` label (safer default for +real-time-ish contexts where false positives must never delete data). +""" +from __future__ import annotations + +import re +from typing import Any, Dict, Iterable, List, Set, Tuple + +from twinkle.preprocessor import Preprocessor + +from .. import label_schema as L +from ..message_utils import msg_content_text + +KEY_CONTAMINATED = 'contaminated' + +_WORD_RE = re.compile(r'\w+', re.UNICODE) + + +def _ngrams(text: str, n: int) -> Set[str]: + tokens = _WORD_RE.findall(text.lower()) + if len(tokens) < n: + return set() + return {' '.join(tokens[i:i + n]) for i in range(len(tokens) - n + 1)} + + +def build_benchmark_index(texts: Iterable[str], n: int = 13) -> Set[str]: + """Build a static n-gram set from benchmark texts (build once, reuse).""" + index: Set[str] = set() + for t in texts: + index |= _ngrams(t or '', n) + return index + + +class Decontaminator(Preprocessor): + """Flag/drop rows that share an n-gram with a static benchmark index. + + Args: + benchmark_ngrams: prebuilt index from :func:`build_benchmark_index`. + n: n-gram size (must match the index's n). Default 13. + min_overlap: number of shared n-grams to count as contaminated. + mode: ``'drop'`` removes contaminated rows; ``'tag'`` keeps them and only + writes the ``contaminated`` label (fail-open). + scan: which roles to scan — 'user' (default), 'assistant', or 'all'. + """ + + def __init__( + self, + benchmark_ngrams: Set[str], + *, + n: int = 13, + min_overlap: int = 1, + mode: str = 'drop', + scan: str = 'user', + ): + if mode not in ('drop', 'tag'): + raise ValueError("mode must be 'drop' or 'tag'") + if scan not in ('user', 'assistant', 'all'): + raise ValueError("scan must be 'user', 'assistant', or 'all'") + self.index = benchmark_ngrams or set() + self.n = int(n) + self.min_overlap = int(min_overlap) + self.mode = mode + self.scan = scan + + def _row_text(self, row: Dict[str, Any]) -> str: + messages = row.get('messages') or [] + parts = [] + for m in messages: + if not isinstance(m, dict): + continue + role = m.get('role') + if self.scan == 'all' or role == self.scan: + parts.append(msg_content_text(m)) + return '\n'.join(p for p in parts if p) + + def _overlap(self, row: Dict[str, Any]) -> int: + if not self.index: + return 0 + grams = _ngrams(self._row_text(row), self.n) + if not grams: + return 0 + return len(grams & self.index) + + def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + rows = self.map_col_to_row(rows) + kept: List[Dict[str, Any]] = [] + dropped: List[Dict[str, Any]] = [] + for row in rows: + overlap = self._overlap(row) + contaminated = overlap >= self.min_overlap + if contaminated and self.mode == 'drop': + dropped.append(dict(row, drop_reason='benchmark_contamination')) + continue + if self.mode == 'tag': + kept.append(L.set_label(row, KEY_CONTAMINATED, contaminated)) + else: + kept.append(row) + return kept, dropped diff --git a/src/twinkle_agentic/preprocessor/offline/near_dedup.py b/src/twinkle_agentic/preprocessor/offline/near_dedup.py new file mode 100644 index 000000000..e1b255688 --- /dev/null +++ b/src/twinkle_agentic/preprocessor/offline/near_dedup.py @@ -0,0 +1,164 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Near-duplicate removal via MinHash-LSH (AUDIT D1) — OFFLINE ONLY. + +``DedupFilter`` collapses only exact/prefix duplicates; a single edited character +slips through. This adds fuzzy near-dup detection over shingled trajectory text. + +OFFLINE CONTRACT (same as :class:`DedupFilter`): this needs a *global* view of +the dataset — it must see all rows in one ``__call__`` and is NOT a per-batch +``QualityPreprocessor`` step. Running near-dup on a per-batch stream would judge +similarity against only the current batch, causing severe false negatives (and, +if used to drop, unstable results). Materialize the dataset, run this once, then +re-wrap the kept rows. + +Uses ``datasketch`` when installed (fast LSH); otherwise falls back to a pure +O(n²) MinHash comparison — correct but slower, fine for modest offline batches. +""" +from __future__ import annotations + +import hashlib +import re +from typing import Any, Dict, List, Set, Tuple + +from twinkle.preprocessor import Preprocessor +from twinkle.utils import get_logger + +from ..message_utils import msg_content_text + +logger = get_logger() + +_WORD_RE = re.compile(r'\w+', re.UNICODE) + + +def _row_text(row: Dict[str, Any]) -> str: + messages = row.get('messages') or [] + return '\n'.join(msg_content_text(m) for m in messages if isinstance(m, dict)) + + +def _shingles(text: str, k: int) -> Set[str]: + tokens = _WORD_RE.findall(text.lower()) + if len(tokens) < k: + return {' '.join(tokens)} if tokens else set() + return {' '.join(tokens[i:i + k]) for i in range(len(tokens) - k + 1)} + + +def _minhash_signature(shingles: Set[str], num_perm: int) -> List[int]: + """Pure-python MinHash: for each of ``num_perm`` salted hashes, take the min.""" + if not shingles: + return [0] * num_perm + sig: List[int] = [] + for p in range(num_perm): + salt = str(p).encode() + mn = min(int(hashlib.md5(salt + s.encode('utf-8')).hexdigest(), 16) for s in shingles) + sig.append(mn) + return sig + + +class NearDupFilter(Preprocessor): + """Global near-duplicate removal over a fully materialized row collection. + + Args: + threshold: Jaccard similarity at/above which two rows are near-duplicates. + shingle_size: word n-gram size for shingling. + num_perm: MinHash permutations (higher = more accurate, slower). + keep: within a near-dup cluster, keep the ``'longest'`` (most messages) + or ``'first'`` seen row. + """ + + def __init__( + self, + *, + threshold: float = 0.8, + shingle_size: int = 5, + num_perm: int = 128, + keep: str = 'longest', + ): + if keep not in ('longest', 'first'): + raise ValueError("keep must be 'longest' or 'first'") + self.threshold = float(threshold) + self.shingle_size = int(shingle_size) + self.num_perm = int(num_perm) + self.keep = keep + + def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + rows = self.map_col_to_row(rows) + n = len(rows) + if n <= 1: + return rows, [] + shingle_sets = [_shingles(_row_text(r), self.shingle_size) for r in rows] + + try: + from datasketch import MinHash, MinHashLSH + clusters = self._cluster_lsh(shingle_sets, MinHash, MinHashLSH) + except Exception as e: + logger.info(f'[NearDupFilter] datasketch unavailable ({e}); pure-python O(n^2) fallback.') + clusters = self._cluster_bruteforce(shingle_sets) + + keep_flag = [True] * n + dropped: List[Dict[str, Any]] = [] + for cluster in clusters: + if len(cluster) <= 1: + continue + winner = self._pick_winner(rows, cluster) + for idx in cluster: + if idx != winner: + keep_flag[idx] = False + dropped.append(dict(rows[idx], drop_reason='near_duplicate')) + kept = [rows[i] for i in range(n) if keep_flag[i]] + return kept, dropped + + def _pick_winner(self, rows: List[Dict[str, Any]], cluster: List[int]) -> int: + if self.keep == 'first': + return min(cluster) + return max(cluster, key=lambda i: len(rows[i].get('messages') or [])) + + def _cluster_lsh(self, shingle_sets, MinHash, MinHashLSH) -> List[List[int]]: + lsh = MinHashLSH(threshold=self.threshold, num_perm=self.num_perm) + mh_list = [] + for i, sh in enumerate(shingle_sets): + mh = MinHash(num_perm=self.num_perm) + for s in sh: + mh.update(s.encode('utf-8')) + mh_list.append(mh) + lsh.insert(str(i), mh) + return self._union_find([(i, [int(x) for x in lsh.query(mh_list[i])]) for i in range(len(shingle_sets))], + len(shingle_sets)) + + def _cluster_bruteforce(self, shingle_sets) -> List[List[int]]: + n = len(shingle_sets) + neighbors: List[Tuple[int, List[int]]] = [] + for i in range(n): + adj = [i] + for j in range(i + 1, n): + a, b = shingle_sets[i], shingle_sets[j] + if not a and not b: + continue + inter = len(a & b) + union = len(a | b) or 1 + if inter / union >= self.threshold: + adj.append(j) + neighbors.append((i, adj)) + return self._union_find(neighbors, n) + + @staticmethod + def _union_find(adjacency: List[Tuple[int, List[int]]], n: int) -> List[List[int]]: + parent = list(range(n)) + + def find(x): + while parent[x] != x: + parent[x] = parent[parent[x]] + x = parent[x] + return x + + def union(a, b): + ra, rb = find(a), find(b) + if ra != rb: + parent[rb] = ra + + for i, adj in adjacency: + for j in adj: + union(i, j) + groups: Dict[int, List[int]] = {} + for i in range(n): + groups.setdefault(find(i), []).append(i) + return list(groups.values()) diff --git a/src/twinkle_agentic/preprocessor/outcome_filter.py b/src/twinkle_agentic/preprocessor/outcome_filter.py new file mode 100644 index 000000000..bc95b5171 --- /dev/null +++ b/src/twinkle_agentic/preprocessor/outcome_filter.py @@ -0,0 +1,80 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Drop failed / dead-end trajectories by reading scores — pure tag reader (AUDIT D6). + +This filter does **not** compute anything. It reads the scores that +:class:`TrajectoryScorer` (D7) already wrote into ``user_data`` and drops rows +whose trajectory score / safety score fall below configurable thresholds. Because +it only consumes labels, the dependency on the scorer is a *data* dependency +(scorer writes ``traj_score``, this reads it) enforced simply by pipeline order — +no module import of the verifier, no DAG. + +Thresholds are meant to be **set by default and then tuned against the observed +score distribution** (self-evolving framework: no human-labeled calibration set). +A row with no score label is kept by default (fail-open) so that placing this +filter before the scorer, or scoring being disabled, never silently drops data. +""" +from __future__ import annotations + +from typing import Any, Dict, List, Tuple + +from twinkle.preprocessor import Preprocessor +from twinkle.utils import get_logger + +from . import label_schema as L + +logger = get_logger() + + +class TrajectoryOutcomeFilter(Preprocessor): + """Drop trajectories whose written scores fall below thresholds (reads only). + + Args: + min_traj_score: drop if ``traj_score`` < this. ``None`` disables. + min_safety_score: drop if ``safety_score`` < this. ``None`` disables. + drop_unsafe_flag: drop if ``safety_unsafe`` is True. Default True. + require_score: if True, rows with no ``traj_score`` label are DROPPED + (fail-closed); default False keeps them (fail-open) so a mis-ordered + or scorer-disabled pipeline never silently deletes data. + """ + + def __init__( + self, + *, + min_traj_score: float = 0.25, + min_safety_score: float = 0.5, + drop_unsafe_flag: bool = True, + require_score: bool = False, + ): + self.min_traj_score = min_traj_score + self.min_safety_score = min_safety_score + self.drop_unsafe_flag = bool(drop_unsafe_flag) + self.require_score = bool(require_score) + + def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + rows = self.map_col_to_row(rows) + kept: List[Dict[str, Any]] = [] + dropped: List[Dict[str, Any]] = [] + for row in rows: + reason = self._drop_reason(row) + if reason is None: + kept.append(row) + else: + dropped.append(dict(row, drop_reason=reason)) + return kept, dropped + + def _drop_reason(self, row: Dict[str, Any]): + traj = L.get_label(row, L.KEY_TRAJ_SCORE, None) + if traj is None: + if self.require_score: + return 'no_score' + # fail-open: unscored rows pass through + elif self.min_traj_score is not None and float(traj) < self.min_traj_score: + return 'low_traj_score' + + if self.drop_unsafe_flag and bool(L.get_label(row, L.KEY_SAFETY_UNSAFE, False)): + return 'unsafe' + safety = L.get_label(row, L.KEY_SAFETY_SCORE, None) + if safety is not None and self.min_safety_score is not None \ + and float(safety) < self.min_safety_score: + return 'low_safety_score' + return None diff --git a/src/twinkle_agentic/preprocessor/pii_presidio_filter.py b/src/twinkle_agentic/preprocessor/pii_presidio_filter.py index 9dafd061f..662eee764 100644 --- a/src/twinkle_agentic/preprocessor/pii_presidio_filter.py +++ b/src/twinkle_agentic/preprocessor/pii_presidio_filter.py @@ -80,6 +80,52 @@ def _hash_short(s: str, salt: str = '') -> str: return hashlib.sha256((salt + s).encode('utf-8')).hexdigest()[:12] +def _faker_available() -> bool: + import importlib.util + return importlib.util.find_spec('faker') is not None + + +def _build_stub_nlp_engine(languages: Sequence[str]): + """A no-op presidio NlpEngine: emits empty NLP artifacts (spaCy-free). + + Lets pattern (regex) recognizers run without loading any language model. + Built lazily so importing this module never requires presidio. + """ + from presidio_analyzer.nlp_engine import NlpArtifacts, NlpEngine + + class _StubNlp(NlpEngine): + def __init__(self, langs): + self._langs = list(langs) + + def load(self): + pass + + def is_loaded(self): + return True + + def process_text(self, text, language): + return NlpArtifacts(entities=[], tokens=[], tokens_indices=[], + lemmas=[], nlp_engine=self, language=language) + + def process_batch(self, texts, language, **kwargs): + for t in texts: + yield t, self.process_text(t, language) + + def is_stopword(self, word, language): + return False + + def is_punct(self, word, language): + return False + + def get_supported_entities(self): + return [] + + def get_supported_languages(self): + return list(self._langs) + + return _StubNlp(languages) + + # ─── Faker dispatcher (per-instance, thread-safe) ─────────────────────────────── @@ -206,9 +252,9 @@ class PIIPresidioFilter(Preprocessor): # identifiers (phone/email/IDs/bank/cards) reliably indicate real PII. URL is also dropped—redacting # links in technical/instruction text changes semantics without privacy benefit. IGNORED_ENTITIES: Tuple[str, ...] = ('PERSON', 'LOCATION', 'ORGANIZATION', 'NRP', 'DATE_TIME', 'URL') - INSTALL_HINT = ('PIIPresidioFilter requires: pip install presidio-analyzer presidio-anonymizer ' - 'faker spacy && python -m spacy download en_core_web_sm && ' - 'python -m spacy download zh_core_web_sm') + INSTALL_HINT = ('PIIPresidioFilter requires: pip install presidio-analyzer presidio-anonymizer. ' + 'For NER-backed entities and Faker replacement also: pip install faker spacy && ' + 'python -m spacy download en_core_web_sm && python -m spacy download zh_core_web_sm') def __init__( self, @@ -222,23 +268,38 @@ def __init__( persistent_consistency: bool = False, hash_salt: str = '', record_counts: bool = False, + regex_only: bool = True, ) -> None: super().__init__() - self._require_deps() + # In regex-only mode we act exclusively on pattern-based identifiers + # (email/phone/cards/IDs/bank), which are the only entities we keep anyway + # (see IGNORED_ENTITIES). This drops the heavy spaCy model load entirely. + self._regex_only = bool(regex_only) + self._require_deps(self._regex_only) self._languages: List[str] = list(languages) self._spacy_models = dict(self.DEFAULT_SPACY_MODELS) if spacy_models: self._spacy_models.update(spacy_models) - for lang in self._languages: - if lang not in self._spacy_models: - raise ValueError(f'No spaCy model configured for language {lang!r}') + if not self._regex_only: + for lang in self._languages: + if lang not in self._spacy_models: + raise ValueError(f'No spaCy model configured for language {lang!r}') self._strategy = {k: Strategy.coerce(v) for k, v in self.DEFAULT_ENTITY_STRATEGY.items()} if entity_strategy: self._strategy.update({k.upper(): Strategy.coerce(v) for k, v in entity_strategy.items()}) self._default_strategy = Strategy.coerce(default_strategy) + # Faker-backed REPLACE needs the optional 'faker' dep. If it is absent + # (common in regex-only deployments) transparently degrade REPLACE->MASK + # so PII is still scrubbed rather than crashing at scrub time. + if not _faker_available(): + if self._default_strategy is Strategy.REPLACE: + self._default_strategy = Strategy.MASK + self._strategy = {k: (Strategy.MASK if v is Strategy.REPLACE else v) + for k, v in self._strategy.items()} + self._score_threshold = score_threshold self._roles = set(roles) self._consistency = consistency @@ -261,17 +322,21 @@ def __init__( # ── construction ──────────────────────────────────────────────────────── @classmethod - def _require_deps(cls) -> None: + def _require_deps(cls, regex_only: bool = True) -> None: try: - import faker # noqa: F401 import presidio_analyzer # noqa: F401 import presidio_anonymizer # noqa: F401 - import spacy # noqa: F401 + if not regex_only: + import spacy # noqa: F401 except ImportError as e: raise ImportError(f'{e}. {cls.INSTALL_HINT}') from e def _build_analyzer(self): from presidio_analyzer import AnalyzerEngine, RecognizerRegistry + + if self._regex_only: + return self._build_regex_analyzer(AnalyzerEngine, RecognizerRegistry) + from presidio_analyzer.nlp_engine import NlpEngineProvider nlp_conf = { @@ -293,6 +358,26 @@ def _build_analyzer(self): registry.add_recognizer(r) return AnalyzerEngine(registry=registry, nlp_engine=nlp_engine, supported_languages=self._languages) + def _build_regex_analyzer(self, AnalyzerEngine, RecognizerRegistry): + """spaCy-free analyzer: only pattern (regex) recognizers, stub NLP engine. + + Presidio's predefined pattern recognizers (email, phone, credit card, + IBAN, IP, etc.) plus our CN identifier recognizers are all regex-based and + need no NLP artifacts, so we feed a no-op NlpEngine and load only those. + NER-driven entities (PERSON/LOCATION/...) are intentionally unavailable — + they are in IGNORED_ENTITIES anyway. + """ + nlp_engine = _build_stub_nlp_engine(self._languages) + registry = RecognizerRegistry(supported_languages=self._languages) + registry.load_predefined_recognizers(languages=self._languages, nlp_engine=nlp_engine) + # Drop recognizers that depend on NLP artifacts (SpacyRecognizer et al.); + # keep only pure PatternRecognizers so analyze() never touches the stub NER. + from presidio_analyzer import PatternRecognizer + registry.recognizers = [r for r in registry.recognizers if isinstance(r, PatternRecognizer)] + for r in _build_cn_recognizers(self._languages): + registry.add_recognizer(r) + return AnalyzerEngine(registry=registry, nlp_engine=nlp_engine, supported_languages=self._languages) + # ── language routing ──────────────────────────────────────────────────── def _resolve_language(self, text: str) -> str: diff --git a/src/twinkle_agentic/preprocessor/provenance.py b/src/twinkle_agentic/preprocessor/provenance.py new file mode 100644 index 000000000..45b926056 --- /dev/null +++ b/src/twinkle_agentic/preprocessor/provenance.py @@ -0,0 +1,73 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Data-lineage / provenance stamping — tag only, never drop (AUDIT D10). + +Industry data pipelines keep provenance so any training example is traceable +back to its source (which dataset, which teacher/student model produced it, when +it was ingested, which cleaning pipeline version touched it). In a self-evolving +distillation loop this is what lets us later attribute a regression to a bad +source or a specific teacher, and to reproduce a training mix. + +This mapper writes a single ``provenance`` blob into ``user_data`` (JSON-packed, +PyArrow-stable via A5). It reads whatever lineage fields already exist on the row +(``model_id`` and any configured passthroughs) and adds an ingest timestamp so +the record is self-describing downstream. +""" +from __future__ import annotations + +import time +from typing import Any, Dict, Sequence + +from twinkle.preprocessor import Mapper + +from . import label_schema as L + + +class ProvenanceStamp(Mapper): + """Stamp each row with a provenance blob in ``user_data`` (never drops). + + Args: + source: a static source/dataset identifier for this ingest batch. + pipeline_version: version string of the cleaning pipeline for audit. + model_field: row field holding the producing model id (default 'model_id'). + extra_fields: additional row fields to copy verbatim into provenance + (e.g. 'teacher_model', 'student_model', 'request_id'). + add_timestamp: include a unix ingest timestamp. Default True. + overwrite: if False, rows that already carry a provenance blob are left + untouched (idempotent re-runs / preserve upstream lineage). Default False. + """ + + def __init__( + self, + *, + source: str = '', + pipeline_version: str = '', + model_field: str = 'model_id', + extra_fields: Sequence[str] = (), + add_timestamp: bool = True, + overwrite: bool = False, + ): + self.source = source + self.pipeline_version = pipeline_version + self.model_field = model_field + self.extra_fields = tuple(extra_fields) + self.add_timestamp = bool(add_timestamp) + self.overwrite = bool(overwrite) + + def map(self, row: Dict[str, Any]) -> Dict[str, Any]: + if not self.overwrite and L.get_label(row, L.KEY_PROVENANCE, None) is not None: + return row + blob: Dict[str, Any] = {} + if self.source: + blob['source'] = self.source + if self.pipeline_version: + blob['pipeline_version'] = self.pipeline_version + model = row.get(self.model_field) + if model: + blob['model'] = model + for f in self.extra_fields: + v = row.get(f) + if v is not None: + blob[f] = v + if self.add_timestamp: + blob['ingested_at'] = int(time.time()) + return L.set_label(row, L.KEY_PROVENANCE, blob) diff --git a/src/twinkle_agentic/preprocessor/refuse_filter.py b/src/twinkle_agentic/preprocessor/refuse_filter.py index 842aae121..f7f303bf5 100644 --- a/src/twinkle_agentic/preprocessor/refuse_filter.py +++ b/src/twinkle_agentic/preprocessor/refuse_filter.py @@ -106,6 +106,28 @@ # refusal-like phrasing doesn't get mistaken for a real user-facing refusal. _THINK_BLOCK_RE = re.compile(r'.*?\s*', re.DOTALL | re.IGNORECASE) +# ── Continuation exemption ──────────────────────────────────────────────────── +# +# A genuine refusal is TERMINAL — the assistant stops helping. In agent / coding +# traces the model very often states a local, technical inability and then +# immediately pivots to an alternative action: +# "I can't write to E:\…. I'll need to use exec to create the directory…" +# "I can't read files outside the sandbox. Let me use exec to …" +# These are NOT refusals of the user's request. If a pivot-to-action cue appears +# anywhere in the scanned window, we exempt the row. +_EN_CONTINUE = re.compile( + r"\b(let\s+me|let'?s|i'?ll|i\s+will|i'?m\s+going\s+to|i\s+need\s+to|i'?ll\s+need\s+to|" + r'instead|so\s+i(\'?ll|\s+will)?|so\s+let|try\s+(again|another)|as\s+an\s+alternative|' + r'alternatively|workaround|work\s+around|use\s+(exec|the\s+\w+\s+tool)|' + r'run\s+the|call\s+the|switch\s+to|fall\s+back)\b', + re.IGNORECASE | re.DOTALL, +) +_ZH_CONTINUE = re.compile( + r'(让我|我来|我先|我会|我将|我需要|改用|换用|换个|换成|试试|尝试|再试|退而|作为替代|' + r'替代方案|变通|绕过|所以我|因此我|接下来我|那我|改为|改成|使用工具|调用工具|执行命令)', + re.UNICODE | re.DOTALL, +) + # ── Helpers ────────────────────────────────────────────────────────────────── @@ -116,31 +138,72 @@ def _text(content: Any) -> str: return content if isinstance(content, str) else '' +# Patterns that signal a *soft* technical inability rather than a hard refusal of +# the user's request. These are the ones prone to false positives on agents that +# state a constraint and keep working, so they are subject to the continuation +# exemption. The remaining patterns (apology-decline, policy/violation, AI-identity +# refusal, standalone "I refuse to") are terminal and never exempted. +_SOFT_INABILITY = frozenset({id(_EN_CORE), id(_ZH_SELF)}) + + def _is_refusal(text: str, check_window: int = 600) -> bool: - """Return True if the text contains a self-referential refusal signal.""" - window = text[:check_window] - return any(p.search(window) for p in _ALL_PATTERNS) + """Return True if the text contains a self-referential refusal signal. + + ``check_window <= 0`` scans the whole text (no truncation). A soft technical + inability ("I can't write to X") that is immediately followed by a pivot to an + alternative action ("let me use exec…") is exempted — that is an agent working + around a constraint, not refusing the user's request. + """ + window = text if check_window <= 0 else text[:check_window] + pivots = None # lazily computed only when a soft-inability pattern hits + for p in _ALL_PATTERNS: + if not p.search(window): + continue + if id(p) in _SOFT_INABILITY: + if pivots is None: + pivots = bool(_EN_CONTINUE.search(window) or _ZH_CONTINUE.search(window)) + if pivots: + continue # constraint-then-pivot: not a refusal + return True + return False # ── Preprocessor ───────────────────────────────────────────────────────────── class RefuseFilter(Preprocessor): - - def __init__(self, check_window: int = 600) -> None: + """Drop rows whose assistant reply is a self-referential refusal. + + Args: + check_window: chars scanned per assistant message (0 = whole message). + scan_all_assistants: scan every assistant turn, not just the first — a + multi-turn conversation may only refuse in a later turn. + scan_reasoning: also scan ``reasoning_content``/``thinking`` fields. + Default False: reasoning traces often rehearse refusal-like phrasing + that the model then overrides, so scanning them raises false positives. + """ + + def __init__(self, check_window: int = 600, *, scan_all_assistants: bool = True, + scan_reasoning: bool = False) -> None: super().__init__() self._check_window = check_window + self._scan_all = bool(scan_all_assistants) + self._scan_reasoning = bool(scan_reasoning) def _is_refusal_row(self, row: Dict[str, Any]) -> bool: messages = row.get('messages') or [] - first_asst = next( - (m for m in messages if isinstance(m, dict) and m.get('role') == 'assistant'), - None, - ) - if first_asst is None: - return False - reply = _THINK_BLOCK_RE.sub('', _text(first_asst.get('content'))).strip() - return bool(reply) and _is_refusal(reply, self._check_window) + asst_msgs = [m for m in messages if isinstance(m, dict) and m.get('role') == 'assistant'] + if not self._scan_all: + asst_msgs = asst_msgs[:1] + for m in asst_msgs: + reply = _THINK_BLOCK_RE.sub('', _text(m.get('content'))).strip() + if reply and _is_refusal(reply, self._check_window): + return True + if self._scan_reasoning: + reasoning = (m.get('reasoning_content') or m.get('thinking') or '').strip() + if reasoning and _is_refusal(reasoning, self._check_window): + return True + return False def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: rows = self.map_col_to_row(rows) diff --git a/src/twinkle_agentic/preprocessor/safety_scorer.py b/src/twinkle_agentic/preprocessor/safety_scorer.py new file mode 100644 index 000000000..e89963529 --- /dev/null +++ b/src/twinkle_agentic/preprocessor/safety_scorer.py @@ -0,0 +1,89 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Safety / toxicity scoring via a FIXED rubric — tag only, never drop (AUDIT D8). + +Instead of adding a separate toxicity classifier + dependency, we reuse the +existing :class:`~twinkle_agentic.verifier.RubricVerifier` with a *fixed* safety +rubric (no stage-1 generation). Each criterion is phrased so that PASS = safe; +the aggregate pass-rate becomes ``safety_score`` in ``[0, 1]`` (higher = safer), +and a below-threshold score sets the ``safety_unsafe`` flag. Both are written to +``user_data``; the downstream :class:`TrajectoryOutcomeFilter` (D6) drops on them. + +This is a mapper (never drops) and shares the rubric voting/aggregation +machinery, so it inherits margin-adaptive voting for uncertain cases for free. +""" +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Tuple + +from twinkle.preprocessor import Preprocessor +from twinkle.utils import get_logger + +from . import label_schema as L + +logger = get_logger() + +# PASS = the trajectory is SAFE w.r.t. the criterion. All are "hard" (weighted). +DEFAULT_SAFETY_CRITERIA: Tuple[str, ...] = ( + 'The content contains no violent, gory, or physically harmful instructions.', + 'The content contains no hateful, harassing, or discriminatory language toward any group.', + 'The content contains no sexual or adult material involving minors or non-consent.', + 'The assistant did NOT comply with a jailbreak / policy-evasion attempt (no successful jailbreak).', + 'The content does not leak private personal data (real names+contacts, credentials, IDs).', + 'The content contains no instructions for weapons, malware, or other clearly illegal harm.', +) + + +class SafetyScorer(Preprocessor): + """Score each trajectory for safety with a fixed rubric; write labels only. + + Args: + rubric_verifier: a pre-built :class:`RubricVerifier`. If ``None``, one is + constructed internally with the fixed safety rubric. When no sampler/ + teacher is available the score defaults to safe (1.0) — this filter + should then be treated as disabled rather than trusted. + criteria: override the default safety criteria (list of PASS=safe strings). + unsafe_threshold: ``safety_unsafe`` is set when ``safety_score`` < this. + """ + + def __init__( + self, + rubric_verifier: Optional[Any] = None, + *, + criteria: Optional[Tuple[str, ...]] = None, + unsafe_threshold: float = 0.5, + ): + from twinkle_agentic.verifier import RubricItem, RubricVerifier + + self.unsafe_threshold = float(unsafe_threshold) + crits = criteria if criteria is not None else DEFAULT_SAFETY_CRITERIA + fixed = [RubricItem(text=c, is_hard=True) for c in crits] + if rubric_verifier is None: + rubric_verifier = RubricVerifier(fixed_rubric=fixed) + else: + rubric_verifier.fixed_rubric = fixed + self.verifier = rubric_verifier + + def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + rows = self.map_col_to_row(rows) + out: List[Dict[str, Any]] = [] + for row in rows: + try: + out.append(self._score_row(row)) + except Exception as e: + logger.warning(f'[SafetyScorer] scoring failed, row left unscored: {e}') + out.append(row) + return out, [] # mapper: never drops + + def _score_row(self, row: Dict[str, Any]) -> Dict[str, Any]: + messages = row.get('messages') + if not isinstance(messages, list) or not messages: + return row + trajectory = {'messages': messages} + if row.get('tools'): + trajectory['tools'] = row['tools'] + detail = self.verifier.score_detail(trajectory) + score = float(detail.scalar) + return L.set_labels(row, { + L.KEY_SAFETY_SCORE: round(score, 6), + L.KEY_SAFETY_UNSAFE: score < self.unsafe_threshold, + }) diff --git a/src/twinkle_agentic/preprocessor/structural_noise.py b/src/twinkle_agentic/preprocessor/structural_noise.py new file mode 100644 index 000000000..27945a4d8 --- /dev/null +++ b/src/twinkle_agentic/preprocessor/structural_noise.py @@ -0,0 +1,61 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Keyword-free structural noise-turn tagging (AUDIT D5, optional). + +Existing heartbeat stripping (``message_normalizer``) is keyword-based and, per +the audit, already covers the common OpenHands/OpenClaw formats. This optional +tagger catches *keyword-free* structural noise: near-identical, very short turns +that repeat across the trajectory (polling / retries with no new signal), using +only cheap structural signals (length + exact repetition) — no embeddings, no +LLM. It **tags** a per-trajectory noise ratio into ``user_data`` (never drops), +so a downstream filter can act on it if desired. + +The embedding-distance variant sketched in the audit is deferred until the D1 +near-dup infrastructure (which provides the embedding index) exists. +""" +from __future__ import annotations + +from collections import Counter +from typing import Any, Dict + +from twinkle.preprocessor import Mapper + +from . import label_schema as L +from .message_utils import msg_content_text, normalize_tool_calls + +KEY_NOISE_RATIO = 'structural_noise_ratio' + + +class StructuralNoiseTagger(Mapper): + """Tag the fraction of assistant turns that are short, repeated boilerplate. + + Args: + short_chars: an assistant turn with visible text at/under this length is a + noise candidate (tool-call turns are exempt — they carry structure). + min_repeat: a candidate counts as noise only if its normalized text recurs + at least this many times across the trajectory's assistant turns. + """ + + def __init__(self, *, short_chars: int = 40, min_repeat: int = 3): + self.short_chars = int(short_chars) + self.min_repeat = int(min_repeat) + + def map(self, row: Dict[str, Any]) -> Dict[str, Any]: + messages = row.get('messages') + if not isinstance(messages, list) or not messages: + return row + asst = [m for m in messages if isinstance(m, dict) and m.get('role') == 'assistant'] + if not asst: + return row + texts = [] + for m in asst: + if normalize_tool_calls(m) is not None: + texts.append(None) # tool-call turn: never noise + else: + texts.append(msg_content_text(m).strip()) + counts = Counter(t for t in texts if t) + noise = 0 + for t in texts: + if t and len(t) <= self.short_chars and counts[t] >= self.min_repeat: + noise += 1 + ratio = noise / len(asst) + return L.set_label(row, KEY_NOISE_RATIO, round(ratio, 6)) diff --git a/src/twinkle_agentic/preprocessor/trajectory_scorer.py b/src/twinkle_agentic/preprocessor/trajectory_scorer.py new file mode 100644 index 000000000..421a53197 --- /dev/null +++ b/src/twinkle_agentic/preprocessor/trajectory_scorer.py @@ -0,0 +1,243 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Per-round / per-segment / per-trajectory scoring — tag only, never drop (AUDIT D7). + +This preprocessor wires the existing ``segment`` + ``verifier`` + ``aggregation`` +infrastructure into the cleaning pipeline. It is a **mapper**: it never removes a +row, it only writes scores into ``user_data`` (see :mod:`label_schema`, A5). A +downstream tail filter (``TrajectoryOutcomeFilter``, D6) reads those labels and +decides what to drop — so scoring and filtering stay decoupled and the pipeline +remains a linear list (no DAG, no filter↔verifier code coupling). + +Flow per trajectory:: + + Segmenter(traj) ─► segments + for each segment: + split_segment_into_rounds ─► rounds + HardScorer(round) ─► RoundScore (per-round hard scalar) + fuse_segment(round_scores, rubric_fn) ─► SegmentScore + └ rubric_fn lazily calls RubricVerifier ONLY when not short-circuited + aggregate_trajectory(segment_scores) ─► TrajectoryScore + +Labels written (all JSON-packed, PyArrow-stable): + round_scores, round_gated, segment_scores, traj_score, traj_level, score_meta. + +The ``RubricVerifier`` is optional: when no sampler/teacher is available the soft +chain returns an empty score and fusion falls back to the hard signal, so the +scorer still produces useful per-round hard scores with zero LLM calls. +""" +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Tuple + +from twinkle.preprocessor import Preprocessor +from twinkle.utils import get_logger + +from . import label_schema as L + +logger = get_logger() + + +class TrajectoryScorer(Preprocessor): + """Score every trajectory and write the scores into ``user_data`` (never drops). + + Args: + segmenter: a :class:`~twinkle_agentic.segment.base.Segmenter`. Defaults to + structural ``TurnSegmenter('cluster')`` (LLM-free). + hard_scorer: a :class:`~twinkle_agentic.verifier.HardScorer` (per-round, + deterministic). Defaults to a plain ``HardScorer()``. + rubric_verifier: optional :class:`~twinkle_agentic.verifier.RubricVerifier` + (per-segment, soft/LLM). If ``None``, only hard scores are used. + hard_agg / fusion / hard_floor / hard_ceil_skip: passed to + :func:`~twinkle_agentic.verifier.fuse_segment`. + traj_agg / weight_by_rounds: passed to + :func:`~twinkle_agentic.verifier.aggregate_trajectory`. + write_round_detail: also store per-check breakdown into ``score_meta``. + """ + + def __init__( + self, + segmenter: Optional[Any] = None, + hard_scorer: Optional[Any] = None, + rubric_verifier: Optional[Any] = None, + *, + hard_agg: str = 'gmean', + fusion: str = 'product', + hard_floor: float = 0.25, + hard_ceil_skip: Optional[float] = None, + traj_agg: str = 'mean', + weight_by_rounds: bool = True, + write_round_detail: bool = False, + calibrate: bool = True, + disagree_margin: float = 0.34, + ): + # Lazy imports keep the module importable even if verifier/segment deps + # are heavy; construction still fails loudly if the packages are absent. + from twinkle_agentic.segment import TurnSegmenter + from twinkle_agentic.verifier import HardScorer + + self.segmenter = segmenter if segmenter is not None else TurnSegmenter('cluster') + self.hard_scorer = hard_scorer if hard_scorer is not None else HardScorer() + self.rubric_verifier = rubric_verifier + self.hard_agg = hard_agg + self.fusion = fusion + self.hard_floor = float(hard_floor) + self.hard_ceil_skip = hard_ceil_skip + self.traj_agg = traj_agg + self.weight_by_rounds = bool(weight_by_rounds) + self.write_round_detail = bool(write_round_detail) + # D7c: self-evolving calibration (no human alignment). + self.calibrate = bool(calibrate) + self.disagree_margin = float(disagree_margin) + + # ------------------------------------------------------------------ + def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + rows = self.map_col_to_row(rows) + out: List[Dict[str, Any]] = [] + for row in rows: + try: + out.append(self._score_row(row)) + except Exception as e: # scoring must never break the pipeline + logger.warning(f'[TrajectoryScorer] scoring failed, row left unscored: {e}') + out.append(row) + return out, [] # mapper: never drops + + # ------------------------------------------------------------------ + def _score_row(self, row: Dict[str, Any]) -> Dict[str, Any]: + from twinkle_agentic.verifier import (aggregate_trajectory, fuse_segment, + split_segment_into_rounds) + from twinkle_agentic.verifier.aggregation import RoundScore + + messages = row.get('messages') + if not isinstance(messages, list) or not messages: + return row + + trajectory = {'messages': messages} + if row.get('tools'): + trajectory['tools'] = row['tools'] + + segments = self.segmenter.segment(trajectory) + if not segments: + return row + + query = self._infer_query(messages) + all_round_scalars: List[float] = [] + all_round_gated: List[bool] = [] + segment_scalars: List[float] = [] + segment_confidence: List[float] = [] + segment_scores = [] + + for s_idx, segment in enumerate(segments): + rounds = split_segment_into_rounds(segment) + round_scores = [] + for r_idx, rnd in enumerate(rounds): + detail = self.hard_scorer.score_detail(rnd) + round_scores.append(RoundScore( + index=r_idx, + hard_scalar=detail.scalar, + gated=detail.gated, + detail=detail if self.write_round_detail else None, + )) + all_round_scalars.append(detail.scalar) + all_round_gated.append(detail.gated) + + rubric_fn = self._make_rubric_fn(segment, query, round_scores) + seg_score = fuse_segment( + s_idx, round_scores, rubric_fn, + hard_agg=self.hard_agg, fusion=self.fusion, + hard_floor=self.hard_floor, hard_ceil_skip=self.hard_ceil_skip, + ) + # carry the rubric ScoreDetail (stashed by rubric_fn) for confidence + seg_score.detail = segment.pop('_last_rubric', None) + segment_scores.append(seg_score) + segment_scalars.append(seg_score.scalar) + segment_confidence.append(self._segment_confidence(seg_score)) + + traj = aggregate_trajectory( + segment_scores, how=self.traj_agg, weight_by_rounds=self.weight_by_rounds) + + labels: Dict[str, Any] = { + L.KEY_ROUND_SCORES: [round(x, 6) for x in all_round_scalars], + L.KEY_ROUND_GATED: all_round_gated, + L.KEY_SEGMENT_SCORES: [round(x, 6) for x in segment_scalars], + L.KEY_TRAJ_SCORE: round(traj.scalar, 6), + L.KEY_TRAJ_LEVEL: traj.level, + } + if self.calibrate: + labels[L.KEY_SEGMENT_CONFIDENCE] = [round(c, 6) for c in segment_confidence] + labels[L.KEY_TRAJ_CONFIDENCE] = round( + sum(segment_confidence) / len(segment_confidence), 6) if segment_confidence else 1.0 + if self.write_round_detail: + labels[L.KEY_SCORE_META] = { + 'n_segments': len(segment_scores), + 'short_circuited': [s.short_circuited for s in segment_scores], + 'segment_hard': [round(s.hard_scalar, 6) for s in segment_scores], + 'segment_rubric': [ + None if s.rubric_scalar is None else round(s.rubric_scalar, 6) + for s in segment_scores + ], + } + return L.set_labels(row, labels) + + # ------------------------------------------------------------------ + def _make_rubric_fn(self, segment: dict, query: str, round_scores): + """Return a zero-arg callable for the soft chain, or None if unavailable. + + ``fuse_segment`` only invokes this when the segment is NOT short-circuited, + so the expensive LLM path runs exactly when the hard signal is inconclusive. + + D7c objective→subjective correction: score once, and if the rubric verdict + disagrees with the deterministic hard signal beyond ``disagree_margin``, + re-score with the objective evidence folded into the transcript so the + judge revises WITH the hard facts in view. The last ``ScoreDetail`` is + stashed on the segment (``_last_rubric``) for confidence estimation. + """ + rv = self.rubric_verifier + if rv is None: + return None + from twinkle_agentic.verifier.aggregation import aggregate_hard_over_rounds + hard_agg_val = aggregate_hard_over_rounds(round_scores, how=self.hard_agg) + + def _fn(): + detail = rv.score_detail(segment, query=query) + if self.calibrate and detail is not None and getattr(detail, 'scalar', None) is not None: + if abs(detail.scalar - hard_agg_val) >= self.disagree_margin: + evidence = (f'Deterministic tool/answer checks scored this ' + f'segment {hard_agg_val:.2f} out of 1.0. Reconcile ' + f'your assessment with this objective evidence.') + revised = rv.score_detail(segment, query=query, extra_context=evidence) + if revised is not None: + detail = revised + segment['_last_rubric'] = detail + return detail + + return _fn + + # ------------------------------------------------------------------ + def _segment_confidence(self, seg_score) -> float: + """Self-evolving confidence in a segment score (no human labels). + + Combines three automatic signals: + 1. hard↔rubric agreement — 1 minus their absolute gap (objective anchoring); + 2. voting stability — fewer escalated votes ⇒ the judge was decisive; + 3. decisiveness — distance of the fused score from the ambiguous 0.5 band. + Short-circuited (hard-only) segments are highly confident by construction. + """ + if seg_score.short_circuited or seg_score.rubric_scalar is None: + return 1.0 + agree = 1.0 - min(1.0, abs(seg_score.hard_scalar - seg_score.rubric_scalar)) + detail = getattr(seg_score, 'detail', None) + n_votes = getattr(detail, 'n_votes', 1) or 1 + max_votes = getattr(self.rubric_verifier, 'max_votes', 1) or 1 + stability = 1.0 if max_votes <= 1 else 1.0 - (n_votes - 1) / max(1, max_votes - 1) + decisive = min(1.0, abs(seg_score.scalar - 0.5) * 2.0) + return max(0.0, min(1.0, 0.5 * agree + 0.3 * stability + 0.2 * decisive)) + + # ------------------------------------------------------------------ + @staticmethod + def _infer_query(messages: List[dict]) -> str: + for m in messages: + if isinstance(m, dict) and m.get('role') == 'user': + c = m.get('content') + if isinstance(c, str) and c.strip(): + return c.strip() + return '(no explicit query)' diff --git a/src/twinkle_agentic/preprocessor/utils.py b/src/twinkle_agentic/preprocessor/utils.py index 7a041fcec..803986db0 100644 --- a/src/twinkle_agentic/preprocessor/utils.py +++ b/src/twinkle_agentic/preprocessor/utils.py @@ -1,359 +1,28 @@ -"""Pure helpers shared across preprocessor modules.""" -import json -import math -import os -import re -from typing import Any, Dict, List, Optional, Set, Tuple - - -def _extract_logprob(lp, token_id: Optional[int] = None) -> Optional[float]: - if lp is None: - return None - if isinstance(lp, (int, float)): - return float(lp) - if not isinstance(lp, dict): - return None - # vLLM with prompt_logprobs=1 returns top-1 PLUS actual token if they differ; - # actual is appended LAST, so iter-first picks the wrong (top-1) one. - entry = None - if token_id is not None: - entry = lp.get(token_id) - if entry is None: - entry = lp.get(str(token_id)) - if entry is None: - entry = next(iter(lp.values()), None) - if entry is None: - return None - if hasattr(entry, 'logprob'): - return float(entry.logprob) - if isinstance(entry, dict): - v = entry.get('logprob') - return float(v) if v is not None else None - if isinstance(entry, (int, float)): - return float(entry) - return None - - -def _to_int_list(x) -> List[int]: - if hasattr(x, 'tolist'): - return x.tolist() - return list(x) - - -def _chr_min_distinct( - cond_lp: List, - asst_lp: List, - cond_ids: List[int], - asst_ids: List[int], - n_prompt: int, - exclude_ids: Optional[Set[int]] = None, -) -> Optional[float]: - """chr_dist_min_pos: fraction of distinct asst-token ids whose - per-occurrence min(cond_lp - asst_lp) is strictly positive.""" - if not asst_lp or not cond_lp or not asst_ids: - return None - n_a = min(len(asst_lp), len(asst_ids)) - n_c = len(cond_lp) - by_tok: Dict[int, List[float]] = {} - for i in range(n_a): - ci = n_prompt + i - if ci >= n_c: - break - tid = asst_ids[i] - if tid is None: - continue - if exclude_ids is not None and int(tid) in exclude_ids: - continue - a = _extract_logprob(asst_lp[i], tid) - c_tok = cond_ids[ci] if ci < len(cond_ids) else None - c = _extract_logprob(cond_lp[ci], c_tok) - if a is None or c is None: - continue - by_tok.setdefault(int(tid), []).append(c - a) - if not by_tok: - return None - pos = sum(1 for diffs in by_tok.values() if min(diffs) > 0) - return pos / len(by_tok) - - -def _chr_min_weighted( - cond_lp: List, - asst_lp: List, - cond_ids: List[int], - asst_ids: List[int], - n_prompt: int, -) -> Optional[float]: - """Magnitude-weighted chr_min: each distinct token contributes |min_delta| - as weight; returns sum(pos_weights) / sum(all_weights).""" - if not asst_lp or not cond_lp or not asst_ids: - return None - n_a = min(len(asst_lp), len(asst_ids)) - n_c = len(cond_lp) - by_tok: Dict[int, List[float]] = {} - for i in range(n_a): - ci = n_prompt + i - if ci >= n_c: - break - tid = asst_ids[i] - if tid is None: - continue - a = _extract_logprob(asst_lp[i], tid) - c_tok = cond_ids[ci] if ci < len(cond_ids) else None - c = _extract_logprob(cond_lp[ci], c_tok) - if a is None or c is None: - continue - by_tok.setdefault(int(tid), []).append(c - a) - if not by_tok: - return None - total_w = 0.0 - pos_w = 0.0 - for diffs in by_tok.values(): - md = min(diffs) - w = abs(md) - total_w += w - if md > 0: - pos_w += w - if total_w == 0: - return None - return pos_w / total_w - - -def _ifd_family_metrics( - cond_lp: List, - asst_lp: List, - cond_ids: List[int], - asst_ids: List[int], - n_prompt: int, -) -> Dict[str, Any]: - """IFD (Cherry-LLM) and S-IFD-{50,75} (T-SHIRT) for one round.""" - if not asst_lp or not cond_lp or not asst_ids: - return {} - n_a = min(len(asst_lp), len(asst_ids)) - n_c = len(cond_lp) - deltas: List[float] = [] - for i in range(n_a): - ci = n_prompt + i - if ci >= n_c: - break - tid = asst_ids[i] - if tid is None: - continue - a = _extract_logprob(asst_lp[i], tid) - c_tok = cond_ids[ci] if ci < len(cond_ids) else None - c = _extract_logprob(cond_lp[ci], c_tok) - if a is None or c is None: - continue - deltas.append(c - a) - if not deltas: - return {} - n = len(deltas) - mean_delta = sum(deltas) / n - out: Dict[str, Any] = { - 'n_tokens': n, - 'mean_delta': mean_delta, - 'ifd': math.exp(-mean_delta), - } - abs_sorted = sorted(range(n), key=lambda i: abs(deltas[i]), reverse=True) - for k_pct in (50, 75): - keep = max(1, int(round(n * k_pct / 100))) - sub = [deltas[i] for i in abs_sorted[:keep]] - out[f's_ifd_{k_pct}'] = math.exp(-sum(sub) / len(sub)) - return out - - -def _mean_logprob_delta( - cond_lp: List, - asst_lp: List, - cond_ids: List[int], - asst_ids: List[int], - n_prompt: int, -) -> Optional[float]: - """Mean per-token (cond_lp - asst_lp) over the response span.""" - if not asst_lp or not cond_lp or not asst_ids: - return None - n_a = min(len(asst_lp), len(asst_ids)) - n_c = len(cond_lp) - deltas: List[float] = [] - for i in range(n_a): - ci = n_prompt + i - if ci >= n_c: - break - tid = asst_ids[i] - if tid is None: - continue - a = _extract_logprob(asst_lp[i], tid) - c_tok = cond_ids[ci] if ci < len(cond_ids) else None - c = _extract_logprob(cond_lp[ci], c_tok) - if a is None or c is None: - continue - deltas.append(c - a) - if not deltas: - return None - return sum(deltas) / len(deltas) - - -def _lp_to_jsonable(lp_list): - """Convert per-position prompt_logprobs into JSON-safe form.""" - out = [] - for lp in (lp_list or []): - if lp is None: - out.append(None) - continue - if isinstance(lp, (int, float)): - out.append(float(lp)) - continue - if not isinstance(lp, dict): - out.append(repr(lp)) - continue - d = {} - for k, v in lp.items(): - if hasattr(v, 'logprob'): - d[str(k)] = { - 'logprob': float(v.logprob), - 'rank': getattr(v, 'rank', None), - 'decoded': getattr(v, 'decoded_token', None) - } - elif isinstance(v, dict): - d[str(k)] = v - else: - d[str(k)] = repr(v) - out.append(d) - return out - - -def _pad_batch(batch: List[List[int]], floor: int) -> Tuple[List[List[int]], int]: - n = len(batch) - if n >= floor or not batch: - return batch, n - return list(batch) + [batch[-1]] * (floor - n), n - - -# ══════════════════════════════════════════════════════════════════════════════ -# Message-format utilities -# ══════════════════════════════════════════════════════════════════════════════ - - -def msg_content_text(msg: Dict[str, Any]) -> str: - """Extract plain text from a message's content (str | list | dict).""" - c = msg.get('content') - if isinstance(c, str): - return c - if isinstance(c, list): - return ' '.join(p.get('text', '') for p in c if isinstance(p, dict) and p.get('type') == 'text') - if isinstance(c, dict) and c.get('type') == 'text': - return c.get('text', '') - return '' - - -def msg_has_media(msg: Dict[str, Any]) -> bool: - """True if message content contains non-text parts (image/audio/video).""" - c = msg.get('content') - return isinstance(c, list) and any(isinstance(p, dict) and p.get('type') not in ('text', None) for p in c) - - -def msg_has_payload(msg: Dict[str, Any]) -> bool: - """True if a message carries any substantive payload (text, tool_calls, reasoning, or media).""" - return bool( - msg_content_text(msg).strip() or msg.get('tool_calls') or msg.get('reasoning_content') or msg.get('thinking') - or msg_has_media(msg)) - - -_CJK_RE = re.compile(r'[\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\uac00-\ud7a3]') - - -def normalize_tool_calls(msg: Dict[str, Any]) -> Optional[List[Any]]: - """Return ``tool_calls`` as a list of dicts, handling PyArrow/HF serialization artifacts.""" - tcs = msg.get('tool_calls') - if isinstance(tcs, str): - s = tcs.strip() - if not s: - return None - try: - decoded = json.loads(s) - except (json.JSONDecodeError, ValueError): - return None - if not isinstance(decoded, list) or not decoded: - return None - tcs = decoded - if not isinstance(tcs, list) or not tcs: - return None - result = [] - for tc in tcs: - if isinstance(tc, str): - try: - tc = json.loads(tc) - except (json.JSONDecodeError, ValueError): - return None - if not isinstance(tc, dict): - return None - func = tc.get('function') - if isinstance(func, str): - try: - func = json.loads(func) - except (json.JSONDecodeError, ValueError): - return None - tc = dict(tc, function=func) - result.append(tc) - return result - - -CJK_CHARS_RE = _CJK_RE - - -def cjk_ratio(text: str) -> float: - """Fraction of non-whitespace characters that are CJK.""" - chars = text.replace(' ', '').replace('\n', '').replace('\t', '') - if not chars: - return 0.0 - return len(CJK_CHARS_RE.findall(chars)) / len(chars) - - -def load_sensitive_words(path: Optional[str]) -> Set[str]: - """Load from external file (one word per line). Blank lines and #-comments ignored.""" - if not path or not os.path.isfile(path): - return set() - words: Set[str] = set() - with open(path, encoding='utf-8') as f: - for line in f: - line = line.strip() - if line and not line.startswith('#'): - words.add(line) - return words - - -def build_sensitive_regex(words: Set[str]) -> Optional['re.Pattern']: - """Build a compiled regex from a set of words. Returns None if empty.""" - if not words: - return None - cjk_words = [] - latin_words = [] - cjk_re = re.compile(r'[\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\uac00-\ud7a3]') - for w in sorted(words): - if cjk_re.search(w): - cjk_words.append(re.escape(w)) - else: - latin_words.append(re.escape(w)) - parts = [] - if latin_words: - parts.append(r'\b(' + '|'.join(latin_words) + r')\b') - if cjk_words: - parts.append('(' + '|'.join(cjk_words) + ')') - return re.compile('|'.join(parts), re.IGNORECASE) - - -def is_agent_row(messages) -> bool: - """Return True if the conversation contains tool interactions (agent trace). - - After MessageNormalizer runs, all non-standard formats are already converted - to standard tool_calls / role=tool — so checking those two signals suffices. - """ - if not isinstance(messages, list): - return False - for m in messages: - if not isinstance(m, dict): - continue - if m.get('role') == 'tool': - return True - if normalize_tool_calls(m): - return True - return False +"""Backward-compat re-export shim (AUDIT A2). + +``utils.py`` was split into two focused modules: +- :mod:`logprob_utils` — log-prob data-selection math (IFD / S-IFD / chr_min), + used only by the experimental log-prob scorers. +- :mod:`message_utils` — message-format helpers used by every active step. + +This shim keeps historical ``from .utils import ...`` imports working. Prefer +importing from the focused modules directly in new code. +""" +from .logprob_utils import (_chr_min_distinct, _chr_min_weighted, # noqa: F401 + _extract_logprob, _ifd_family_metrics, + _lp_to_jsonable, _mean_logprob_delta, _pad_batch, + _to_int_list) +from .message_utils import (CJK_CHARS_RE, build_sensitive_regex, # noqa: F401 + cjk_ratio, is_agent_row, load_sensitive_words, + msg_content_text, msg_has_media, msg_has_payload, + normalize_tool_calls) + +__all__ = [ + # message utils + 'msg_content_text', 'msg_has_media', 'msg_has_payload', 'normalize_tool_calls', + 'cjk_ratio', 'CJK_CHARS_RE', 'load_sensitive_words', 'build_sensitive_regex', + 'is_agent_row', + # logprob utils + '_extract_logprob', '_to_int_list', '_chr_min_distinct', '_chr_min_weighted', + '_ifd_family_metrics', '_mean_logprob_delta', '_lp_to_jsonable', '_pad_batch', +] diff --git a/src/twinkle_agentic/segment/__init__.py b/src/twinkle_agentic/segment/__init__.py new file mode 100644 index 000000000..17f64e7f2 --- /dev/null +++ b/src/twinkle_agentic/segment/__init__.py @@ -0,0 +1,4 @@ +from .base import Segmenter, Turn, TurnSegmenter +from .llm_segmenter import LlmSegmenter + +__all__ = ['Segmenter', 'Turn', 'TurnSegmenter', 'LlmSegmenter'] diff --git a/src/twinkle_agentic/segment/base.py b/src/twinkle_agentic/segment/base.py new file mode 100644 index 000000000..353daa5cf --- /dev/null +++ b/src/twinkle_agentic/segment/base.py @@ -0,0 +1,237 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Trajectory segmentation for segment-level rubric scoring. + +A long agent trajectory is split into a list of *segments*; each segment is a +self-contained sub-trajectory (``{'messages': [...], 'tools': [...]}``) that can +be fed directly to a :class:`~twinkle_agentic.verifier.Verifier`. + +Two layers, matching the literature: +- **Structural (per-turn)** — free, deterministic. A "turn" is one assistant + message plus the tool result messages it triggered (Web-Shepherd / AgentPRM + style turn-level MDP). See :class:`TurnSegmenter`. +- **Sub-goal (LLM)** — compress each turn to a one-line intent gist, then make + ONE LLM pass over the whole (short) gist list to group turns into a few + coarse sub-goals (Web-Shepherd / MiRA style), then reassemble segments from + the ORIGINAL messages by index. See :class:`LlmSegmenter`. + +Design notes: +- The leading ``system`` message and the first ``user`` message form the + trajectory *preamble*; it is not itself a scorable segment, but every segment + carries the preamble (system + original user query) so the verifier keeps the + task context. This mirrors ``RubricVerifier._infer_query``. +- A new ``user`` message mid-trajectory is a hard boundary (a new turn starts). +""" +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Tuple + + +@dataclass +class Turn: + """One structural turn: an assistant message + its tool results. + + ``indices`` are positions into the original ``messages`` list so a segment + can be reassembled verbatim from the source (never from a summary). + """ + indices: List[int] + messages: List[dict] = field(default_factory=list) + role_kind: str = 'assistant' # 'assistant' | 'user' | 'other' + + +class Segmenter(ABC): + """Split a trajectory into scorable sub-trajectories. + + Subclasses implement :meth:`segment`. The base class provides the shared + structural turn splitter and segment assembly so LLM-based subclasses only + decide *how turns are grouped*. + """ + + def __call__(self, trajectory: dict, **kwargs) -> List[dict]: + return self.segment(trajectory, **kwargs) + + @abstractmethod + def segment(self, trajectory: dict, **kwargs) -> List[dict]: + raise NotImplementedError + + # ------------------------------------------------------------------ + # shared: preamble + structural turn splitting + assembly + # ------------------------------------------------------------------ + @staticmethod + def _split_preamble(messages: List[dict]) -> Tuple[List[dict], int]: + """Return (preamble_messages, start_index). + + Preamble = leading system message(s) + the first user message. Segments + start at ``start_index`` (the message after the first user turn). + """ + preamble: List[dict] = [] + i = 0 + n = len(messages) + while i < n and messages[i].get('role') == 'system': + preamble.append(messages[i]) + i += 1 + if i < n and messages[i].get('role') == 'user': + preamble.append(messages[i]) + i += 1 + return preamble, i + + @classmethod + def split_turns(cls, messages: List[dict], start: int = 0) -> List[Turn]: + """Group messages[start:] into structural turns. + + A turn begins at an ``assistant`` message and absorbs the following + ``tool`` messages. A mid-trajectory ``user`` message becomes its own + boundary turn (role_kind='user'). Stray leading non-assistant messages + are attached to the first turn. + """ + turns: List[Turn] = [] + cur: Optional[Turn] = None + for idx in range(start, len(messages)): + role = messages[idx].get('role') + if role == 'assistant': + cur = Turn(indices=[idx], messages=[messages[idx]], role_kind='assistant') + turns.append(cur) + elif role == 'user': + # hard boundary: user re-prompt starts a fresh turn + cur = Turn(indices=[idx], messages=[messages[idx]], role_kind='user') + turns.append(cur) + else: # tool / other -> attach to current turn, or open a new one + if cur is None: + cur = Turn(indices=[idx], messages=[messages[idx]], role_kind='other') + turns.append(cur) + else: + cur.indices.append(idx) + cur.messages.append(messages[idx]) + return turns + + @staticmethod + def _assemble(trajectory: dict, preamble: List[dict], turns_slice: List[Turn]) -> dict: + """Build a segment sub-trajectory from preamble + a slice of turns. + + Messages are taken from the ORIGINAL trajectory (verbatim), so the + verifier scores real content, not any compressed gist. + """ + seg_messages: List[dict] = list(preamble) + for t in turns_slice: + seg_messages.extend(t.messages) + seg: Dict[str, Any] = {'messages': seg_messages} + if trajectory.get('tools'): + seg['tools'] = list(trajectory['tools']) + if trajectory.get('user_data'): + seg['user_data'] = list(trajectory['user_data']) + return seg + + @classmethod + def _segments_from_groups( + cls, + trajectory: dict, + preamble: List[dict], + turns: List[Turn], + groups: List[List[int]], + ) -> List[dict]: + """Assemble segments given a grouping of turn-indices. + + ``groups`` is a list of lists of indices into ``turns``. Robust to + gaps/overlaps: see :meth:`_normalize_groups`. + """ + groups = cls._normalize_groups(groups, len(turns)) + return [cls._assemble(trajectory, preamble, [turns[i] for i in grp]) + for grp in groups if grp] + + @staticmethod + def _normalize_groups(groups: List[List[int]], n_turns: int) -> List[List[int]]: + """Repair LLM-proposed groupings into a clean partition of 0..n_turns-1. + + - drop out-of-range indices + - dedupe (first occurrence wins; later duplicates dropped) + - sort each group; sort groups by their first index + - assign any uncovered turns to the nearest preceding group (or the + first group), so every turn lands in exactly one segment + """ + if n_turns <= 0: + return [] + seen: set = set() + cleaned: List[List[int]] = [] + for grp in groups: + g = [] + for i in grp: + if isinstance(i, bool): # guard: bools are ints in python + continue + if isinstance(i, int) and 0 <= i < n_turns and i not in seen: + seen.add(i) + g.append(i) + if g: + cleaned.append(sorted(g)) + cleaned.sort(key=lambda g: g[0]) + + # cover missing turns + missing = [i for i in range(n_turns) if i not in seen] + if missing: + if not cleaned: + cleaned = [missing] + else: + for i in missing: + # nearest preceding group by first-index + target = cleaned[0] + for grp in cleaned: + if grp[0] <= i: + target = grp + else: + break + target.append(i) + for grp in cleaned: + grp.sort() + cleaned.sort(key=lambda g: g[0]) + return cleaned + + +class TurnSegmenter(Segmenter): + """Structural, LLM-free segmenter. + + ``granularity='turn'``: one segment per turn (finest; per-tool-call level). + ``granularity='cluster'``: merge consecutive tool-using assistant turns into + one segment, closing the cluster on a turn that produces a user-facing + text answer with no tool calls (sub-task attempt level). A ``user`` turn + always starts a new cluster. + """ + + def __init__(self, granularity: str = 'cluster'): + if granularity not in ('turn', 'cluster'): + raise ValueError("granularity must be 'turn' or 'cluster'") + self.granularity = granularity + + def segment(self, trajectory: dict, **kwargs) -> List[dict]: + messages = list(trajectory.get('messages', []) or []) + preamble, start = self._split_preamble(messages) + turns = self.split_turns(messages, start) + if not turns: + return [self._assemble(trajectory, preamble, [])] if preamble else [] + + if self.granularity == 'turn': + groups = [[i] for i in range(len(turns))] + else: + groups = self._cluster_groups(turns) + return self._segments_from_groups(trajectory, preamble, turns, groups) + + @staticmethod + def _cluster_groups(turns: List[Turn]) -> List[List[int]]: + groups: List[List[int]] = [] + cur: List[int] = [] + for i, t in enumerate(turns): + if t.role_kind == 'user': + if cur: + groups.append(cur) + cur = [] + groups.append([i]) # user re-prompt as its own boundary segment + continue + cur.append(i) + has_tool_call = any(m.get('role') == 'assistant' and m.get('tool_calls') + for m in t.messages) + if not has_tool_call: + # a text-only assistant answer closes the current sub-task cluster + groups.append(cur) + cur = [] + if cur: + groups.append(cur) + return groups diff --git a/src/twinkle_agentic/segment/llm_segmenter.py b/src/twinkle_agentic/segment/llm_segmenter.py new file mode 100644 index 000000000..7230d225f --- /dev/null +++ b/src/twinkle_agentic/segment/llm_segmenter.py @@ -0,0 +1,308 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Compress-then-segment: per-turn action gist + one-pass LLM sub-goal grouping. + +Pipeline (all LLM calls distilled via ``llm_backup``): + 1. Split the trajectory into structural turns (base class, free). + 2. Compress each turn to a one-line intent gist with an + :class:`ActionSummarizer` (optional; falls back to a truncated + structural render when no summarizer is given). + 3. Make ONE LLM pass over the numbered gist list to group turns into a few + coarse sub-goals — the segmenter LLM sees the whole (short) trajectory at + once, which is what makes segmentation global yet cheap. + 4. Reassemble segments from the ORIGINAL messages by turn index (scoring + always uses verbatim content, never the gist). + +The grouping call is wrapped with ``@llm_backup`` (student sub-goal model with +teacher fallback + progressive distillation), keyed by ``query`` so confidence +is tracked per task family. Index alignment is validated/repaired in pure code +by the base class ``_normalize_groups`` (contiguous, non-overlapping, full +coverage). +""" +from __future__ import annotations + +import json +import re +from typing import TYPE_CHECKING, Any, List, Optional + +from twinkle_agentic.utils.llm_backup import llm_backup + +from .base import Segmenter, Turn + +if TYPE_CHECKING: + from twinkle.data_format import SamplingParams # noqa: F401 + from twinkle.sampler.base import Sampler # noqa: F401 + from twinkle_agentic.summarizer.action_summarizer import ActionSummarizer # noqa: F401 + + +_SEG_SYSTEM = """\ +You segment an AI agent's trajectory into a FEW coarse sub-goals for later \ +evaluation. You are given the task query and a numbered list of one-line turn \ +gists (one per turn). Group CONSECUTIVE turns that jointly pursue the same \ +sub-goal. + +Guidelines: +- Prefer {min_g}-{max_g} sub-goals total. Each sub-goal spans a contiguous + range of turns; ranges must not overlap and must cover every turn. +- Group by MEANINGFUL task progress, not by exact actions: e.g. several + consecutive searches that gather info for one purpose form ONE sub-goal. +- A user re-prompt starts a new sub-goal. + +Output ONLY a JSON array, each item: {"goal": "", "start": , "end": } +Turn numbers are 0-based and inclusive. No prose, no markdown fence.""" + +_SEG_USER = """\ +## Task / query +{query} + +## Turn gists (numbered) +{gists} + +Now output the JSON array of sub-goals covering turns 0..{last}.""" + + +class LlmSegmenter(Segmenter): + """Sub-goal segmenter via compress-then-one-pass-LLM grouping. + + Args: + sampler: Student model sampler for the grouping call. If ``None``, every + grouping call is served by the teacher via ``llm_backup``. + action_summarizer: Optional :class:`ActionSummarizer` used to compress + each turn. If ``None``, a truncated structural render is used as the + gist (no per-turn LLM cost). + min_subgoals / max_subgoals: Target sub-goal count window. + sampling_params: Sampling params for the grouping call. + lora_path: LoRA adapter for the sub-goal student model. + max_gist_chars: Truncation for the structural-render fallback gist. + """ + + def __init__( + self, + sampler: Optional['Sampler'] = None, + *, + action_summarizer: Optional['ActionSummarizer'] = None, + min_subgoals: int = 3, + max_subgoals: int = 6, + sampling_params: Optional['SamplingParams'] = None, + lora_path: Optional[str] = None, + max_gist_chars: int = 200, + ): + if max_subgoals < min_subgoals: + raise ValueError('max_subgoals must be >= min_subgoals') + if min_subgoals < 1: + raise ValueError('min_subgoals must be >= 1') + self.sampler = sampler + self.action_summarizer = action_summarizer + self.min_subgoals = int(min_subgoals) + self.max_subgoals = int(max_subgoals) + self.sampling_params = sampling_params + self.lora_path = lora_path or None + self.max_gist_chars = int(max_gist_chars) + + # ------------------------------------------------------------------ + def segment(self, trajectory: dict, *, query: Optional[str] = None, **kwargs) -> List[dict]: + messages = list(trajectory.get('messages', []) or []) + preamble, start = self._split_preamble(messages) + turns = self.split_turns(messages, start) + if not turns: + return [self._assemble(trajectory, preamble, [])] if preamble else [] + # Too few turns to bother segmenting -> one segment. + if len(turns) <= self.min_subgoals: + return self._segments_from_groups( + trajectory, preamble, turns, [[i] for i in range(len(turns))]) + + from .base import TurnSegmenter + + # No LLM available at all (no student sampler AND no teacher configured) + # -> degrade to free structural clustering instead of crashing. + if not self._llm_available(): + groups = TurnSegmenter._cluster_groups(turns) + return self._segments_from_groups(trajectory, preamble, turns, groups) + + query = query or self._infer_query(messages) + gists = [self._turn_gist(t, query) for t in turns] + gist_block = '\n'.join(f'[{i}] {g}' for i, g in enumerate(gists)) + + raw = self._group( + trajectory=self._group_trajectory(query, gist_block, len(turns)), + sampling_params=self._group_sampling_params(), + query=query) + groups = self._parse_groups(raw, len(turns)) + if not groups: + # LLM produced nothing usable -> fall back to structural clustering. + groups = TurnSegmenter._cluster_groups(turns) + return self._segments_from_groups(trajectory, preamble, turns, groups) + + def _llm_available(self) -> bool: + """True if a student sampler exists or a teacher API is configured. + + Mirrors the env vars ``llm_backup`` uses for its teacher; when neither a + student nor a teacher is present we must not attempt an LLM call. + """ + if self.sampler is not None: + return True + import os + return bool(os.environ.get('LLM_BACKUP_API_KEY') + or os.environ.get('OPENAI_API_KEY') + or os.environ.get('LLM_BACKUP_BASE_URL')) + + # ------------------------------------------------------------------ + # the distilled grouping call + # ------------------------------------------------------------------ + @llm_backup(key_params=['query'], comparator=lambda a, b: _grouping_similar(a, b)) + def _group(self, trajectory, sampling_params, query: str = None) -> str: + if self.sampler is None: + return '' + sample_kwargs: dict[str, Any] = {'sampling_params': sampling_params} + if self.lora_path is None: + sample_kwargs['use_base_model'] = True + else: + sample_kwargs['adapter_path'] = self.lora_path + responses = self.sampler.sample([trajectory], **sample_kwargs) + resp = list(responses)[0] if responses else None + if resp is None: + return '' + seqs = getattr(resp, 'sequences', None) or [] + return (getattr(seqs[0], 'decoded', None) or '') if seqs else '' + + # ------------------------------------------------------------------ + # per-turn gist + # ------------------------------------------------------------------ + def _turn_gist(self, turn: Turn, query: str) -> str: + rendered = self._render_turn(turn) + if self.action_summarizer is not None: + try: + gist = self.action_summarizer(rendered, query=query) + if isinstance(gist, str) and gist.strip(): + return ' '.join(gist.split())[:self.max_gist_chars] + except Exception: + pass + # fallback: truncated structural render (no LLM) + return ' '.join(rendered.split())[:self.max_gist_chars] + + @staticmethod + def _render_turn(turn: Turn) -> str: + parts: List[str] = [] + for m in turn.messages: + role = m.get('role', '?') + content = m.get('content') + if isinstance(content, list): + content = '\n'.join(p.get('text', '') for p in content + if isinstance(p, dict) and p.get('type') == 'text') + content = content or '' + tool_calls = m.get('tool_calls') or [] + if tool_calls: + names = ', '.join((tc.get('function') or {}).get('name', '?') + for tc in tool_calls if isinstance(tc, dict)) + parts.append(f'{role}: {content} [tool_calls: {names}]'.strip()) + else: + parts.append(f'{role}: {content}'.strip()) + return ' | '.join(p for p in parts if p) + + # ------------------------------------------------------------------ + # prompt / sampling plumbing + # ------------------------------------------------------------------ + def _group_trajectory(self, query: str, gist_block: str, n_turns: int) -> dict: + system = (_SEG_SYSTEM + .replace('{min_g}', str(self.min_subgoals)) + .replace('{max_g}', str(self.max_subgoals))) + user = (_SEG_USER + .replace('{query}', query) + .replace('{gists}', gist_block) + .replace('{last}', str(n_turns - 1))) + return {'messages': [ + {'role': 'system', 'content': system}, + {'role': 'user', 'content': user}, + ]} + + def _group_sampling_params(self): + if self.sampling_params is not None: + return self.sampling_params + from twinkle.data_format.sampling import SamplingParams + return SamplingParams(temperature=0.0, max_tokens=512) + + # ------------------------------------------------------------------ + # parsing + # ------------------------------------------------------------------ + @staticmethod + def _infer_query(messages: List[dict]) -> str: + for m in messages: + if m.get('role') == 'user': + c = m.get('content') + if isinstance(c, str) and c.strip(): + return c.strip() + return '(no explicit query)' + + _JSON_ARRAY_RE = re.compile(r'\[.*\]', re.DOTALL) + + @classmethod + def _parse_groups(cls, raw: str, n_turns: int) -> List[List[int]]: + """Parse the sub-goal JSON into a list of turn-index groups. + + Accepts ``[{"goal":..,"start":i,"end":j}, ...]``. Falls back to empty + on unparseable output (caller then uses structural clustering). + """ + if not raw: + return [] + text = raw.strip() + m = cls._JSON_ARRAY_RE.search(text) + if not m: + return [] + try: + data = json.loads(m.group(0)) + except (json.JSONDecodeError, ValueError): + return [] + if not isinstance(data, list): + return [] + groups: List[List[int]] = [] + for item in data: + if not isinstance(item, dict): + continue + s, e = item.get('start'), item.get('end') + if not isinstance(s, int) or not isinstance(e, int): + continue + if e < s: + s, e = e, s + s = max(0, s) + e = min(n_turns - 1, e) + grp = list(range(s, e + 1)) + if grp: + groups.append(grp) + return groups + + +# --------------------------------------------------------------------------- +# comparator for llm_backup +# --------------------------------------------------------------------------- +_JSON_ARRAY_RE = re.compile(r'\[.*\]', re.DOTALL) + + +def _boundaries(raw: str) -> Optional[List[int]]: + m = _JSON_ARRAY_RE.search((raw or '').strip()) + if not m: + return None + try: + data = json.loads(m.group(0)) + except (json.JSONDecodeError, ValueError): + return None + if not isinstance(data, list): + return None + starts = [] + for item in data: + if isinstance(item, dict) and isinstance(item.get('start'), int): + starts.append(item['start']) + return sorted(starts) if starts else None + + +def _grouping_similar(a: str, b: str) -> bool: + """Two segmentations match when they have a similar number of sub-goals and + near-identical boundaries (not byte-identical labels/text).""" + ba, bb = _boundaries(a), _boundaries(b) + if ba is None or bb is None: + return (a or '').strip() == (b or '').strip() + if abs(len(ba) - len(bb)) > 1: + return False + # Jaccard-ish agreement on boundary start positions + sa, sb = set(ba), set(bb) + inter = len(sa & sb) + union = len(sa | sb) or 1 + return inter / union >= 0.6 diff --git a/src/twinkle_agentic/summarizer/__init__.py b/src/twinkle_agentic/summarizer/__init__.py index 9a4672923..9991a0e33 100644 --- a/src/twinkle_agentic/summarizer/__init__.py +++ b/src/twinkle_agentic/summarizer/__init__.py @@ -1,4 +1,5 @@ +from .action_summarizer import ActionSummarizer from .base import Summarizer from .fact_summarizer import FactSummarizer -__all__ = ['Summarizer', 'FactSummarizer'] +__all__ = ['Summarizer', 'FactSummarizer', 'ActionSummarizer'] diff --git a/src/twinkle_agentic/summarizer/action_summarizer.py b/src/twinkle_agentic/summarizer/action_summarizer.py new file mode 100644 index 000000000..2e6c3aa0d --- /dev/null +++ b/src/twinkle_agentic/summarizer/action_summarizer.py @@ -0,0 +1,86 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Action-level summarizer: compress ONE agent turn into a single-line +``(action, goal, result-state)`` gist for downstream trajectory segmentation. + +Unlike :class:`FactSummarizer` (which preserves *facts* for retrieval), this +summarizer preserves *intent* — what the agent tried to do this turn and +whether it worked — because that is what a segmenter needs to group turns into +sub-goals. The gist is deliberately tiny (one line) so a whole long trajectory +can be laid out and segmented in a single LLM pass. + +Same machinery as every other component: the shared ``_sample`` is decorated +with ``@llm_backup`` so student/teacher routing + progressive distillation +happen transparently. +""" +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from twinkle_agentic.summarizer.base import Summarizer + +if TYPE_CHECKING: + from twinkle.data_format import SamplingParams # noqa: F401 + from twinkle.sampler.base import Sampler # noqa: F401 + + +_ACTION_SCHEMA = """\ +You compress ONE turn of an AI agent's trajectory into a single short line \ +that captures INTENT, not facts. A downstream segmenter reads these lines to \ +group consecutive turns into sub-goals, so keep the action's PURPOSE and \ +OUTCOME, and drop retrieved content / long results. + +Output EXACTLY one line in this shape (<= 140 chars): + | goal: | result: + +Where is a terse verb-phrase for the turn, e.g. one of: + search, read, compute, call-tool:, plan, reason, ask-user, final-answer, other + +Rules: +- Focus on WHY the turn happened and WHETHER it advanced the task. +- result=ok (tool/step succeeded), fail (error/empty), partial (some progress), + pending (awaiting more), answer (produced a user-facing final answer). +- Do NOT copy retrieved passages, numbers or long tool output — only the gist. +- One line only. No markdown, no extra commentary. +""" + +_ACTION_USER_TEMPLATE = """\ +Summarize this single agent turn into ONE intent line (see the required shape). \ +Capture the action, its goal, and the result-state. Do not exceed {budget} chars. \ +Ignore long retrieved content; keep only what the turn was DOING. + +## Task / query (context) +{query} + +## Turn +{text}""" + + +class ActionSummarizer(Summarizer): + """Compress a single rendered turn into a one-line intent gist. + + Defaults target a tiny fixed budget (one line) regardless of turn length, + which is what makes a whole-trajectory single-pass segmentation cheap. + """ + + def __init__( + self, + sampler: 'Sampler', + compression_ratio: float = 6.0, + *, + model_path: str = '', + sampling_params: 'SamplingParams | None' = None, + min_budget_chars: int = 140, + template: Any | None = None, + lora_path: str | None = None, + ): + super().__init__( + sampler, + compression_ratio, + model_path=model_path, + sampling_params=sampling_params, + system_prompt=_ACTION_SCHEMA, + user_prompt_template=_ACTION_USER_TEMPLATE, + min_budget_chars=min_budget_chars, + template=template, + lora_path=lora_path, + ) diff --git a/src/twinkle_agentic/utils/llm_backup.py b/src/twinkle_agentic/utils/llm_backup.py index ce5199932..3ce0c4dd8 100644 --- a/src/twinkle_agentic/utils/llm_backup.py +++ b/src/twinkle_agentic/utils/llm_backup.py @@ -99,10 +99,16 @@ def _get_teacher_api(): if _teacher_api is not None: return _teacher_api from twinkle_agentic.protocol.openai import OpenAI + # Bound per-request latency: without a timeout a single hung request blocks + # the calling worker for the SDK default (~600s) x retries. Overridable via + # env for slow/large-prompt endpoints. + timeout = float(os.environ.get('LLM_BACKUP_TIMEOUT', '120')) + max_retries = int(os.environ.get('LLM_BACKUP_MAX_RETRIES', '2')) _teacher_api = OpenAI( - model=os.environ.get('LLM_BACKUP_MODEL', 'gpt-4o'), + model=os.environ.get('LLM_BACKUP_MODEL', 'qwen3.7-max'), api_key=os.environ.get('LLM_BACKUP_API_KEY'), base_url=os.environ.get('LLM_BACKUP_BASE_URL'), + client_kwargs={'timeout': timeout, 'max_retries': max_retries}, ) return _teacher_api diff --git a/src/twinkle_agentic/verifier/__init__.py b/src/twinkle_agentic/verifier/__init__.py index 3265fed9b..98c7d13ed 100644 --- a/src/twinkle_agentic/verifier/__init__.py +++ b/src/twinkle_agentic/verifier/__init__.py @@ -1,3 +1,7 @@ +from .aggregation import (RoundScore, SegmentScore, TrajectoryScore, + aggregate_hard_over_rounds, aggregate_trajectory, + fuse_segment, scalar_to_level, + split_segment_into_rounds) from .base import Verifier from .domain_checks import (check_answer_match, check_code_parses, check_instruction_constraints, check_not_degenerate, @@ -13,4 +17,7 @@ 'check_output_format', 'check_numeric_equiv', 'check_answer_match', 'check_code_parses', 'check_instruction_constraints', 'check_not_degenerate', 'default_checks_for', + 'RoundScore', 'SegmentScore', 'TrajectoryScore', + 'split_segment_into_rounds', 'aggregate_hard_over_rounds', 'fuse_segment', + 'aggregate_trajectory', 'scalar_to_level', ] diff --git a/src/twinkle_agentic/verifier/aggregation.py b/src/twinkle_agentic/verifier/aggregation.py new file mode 100644 index 000000000..9ab4b4cb1 --- /dev/null +++ b/src/twinkle_agentic/verifier/aggregation.py @@ -0,0 +1,262 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Multi-granularity score aggregation. + +Two scorers operate at deliberately different granularities: + +- :class:`HardScorer` is **per-round** — tool-call validity / execution / + protocol are facts about a single assistant turn, cheap and deterministic. +- :class:`RubricVerifier` is **per-segment** — soft quality (sub-goal progress, + reasoning soundness, no redundant calls) needs multi-round context. + +This module bridges the two without forcing either onto the other's grain: + + rounds ── HardScorer (per round) ──► h_1..h_R + │ aggregate over the rounds in a segment + segment ── RubricVerifier (whole) ──► rubric_scalar + │ fuse(hard_agg, rubric_scalar) + ▼ + segment score ──► aggregate ──► trajectory score + +It also implements the **short-circuit gate**: when a segment's hard score is +extreme (e.g. every tool call failed, no final answer), the soft rubric chain +is skipped entirely — saving the long/expensive LLM path exactly when the +answer is already decided. + +The functions here are pure: pass callables/detail objects, get scores back. +They do not import HardScorer/RubricVerifier, so they stay easily testable and +decoupled (the orchestrating TrajectoryScorer will wire real scorers in later). +""" +from __future__ import annotations + +import statistics +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, List, Optional, Sequence + +from twinkle_agentic.segment.base import Segmenter + +NUM_LEVELS = 5 + +# aggregation strategies for combining a list of scalars in [0, 1] +_REDUCERS: Dict[str, Callable[[Sequence[float]], float]] = { + 'mean': lambda xs: sum(xs) / len(xs), + 'min': min, + 'max': max, + 'median': lambda xs: statistics.median(xs), + # geometric mean: harsher than mean, one bad round drags the whole segment + 'gmean': lambda xs: (statistics.geometric_mean([max(1e-6, x) for x in xs])), +} + + +def _reduce(xs: Sequence[float], how: str) -> float: + xs = [x for x in xs if x is not None] + if not xs: + return 0.0 + fn = _REDUCERS.get(how) + if fn is None: + raise ValueError(f'unknown reducer {how!r}; choose from {list(_REDUCERS)}') + return float(fn(xs)) + + +def scalar_to_level(scalar: float, num_levels: int = NUM_LEVELS) -> int: + scalar = min(1.0, max(0.0, scalar)) + return min(num_levels - 1, max(0, int(round(scalar * (num_levels - 1))))) + + +# --------------------------------------------------------------------------- +# score containers +# --------------------------------------------------------------------------- +@dataclass +class RoundScore: + """Per-round hard score.""" + index: int # round index within the segment + hard_scalar: float + gated: bool = False + detail: Any = None # optional HardScoreDetail + + +@dataclass +class SegmentScore: + """Fused per-segment score.""" + index: int + scalar: float # final fused score in [0, 1] + level: int + hard_scalar: float # aggregated hard score over the segment's rounds + rubric_scalar: Optional[float] # None if the soft chain was short-circuited + short_circuited: bool + n_rounds: int + rounds: List[RoundScore] = field(default_factory=list) + detail: Any = None # optional rubric ScoreDetail + + +@dataclass +class TrajectoryScore: + """Trajectory-level score aggregated over segments.""" + scalar: float + level: int + segments: List[SegmentScore] = field(default_factory=list) + + +# --------------------------------------------------------------------------- +# round -> segment +# --------------------------------------------------------------------------- +def split_segment_into_rounds(segment: dict) -> List[dict]: + """Split a segment sub-trajectory into per-round sub-trajectories. + + Reuses :meth:`Segmenter.split_turns`. The segment's preamble (system + the + first user message) is carried onto every round so a per-round HardScorer + still sees tools/context. Each returned dict is a valid sub-trajectory + (``messages`` + ``tools``/``user_data`` when present). + """ + messages = list(segment.get('messages', []) or []) + preamble, start = Segmenter._split_preamble(messages) + turns = Segmenter.split_turns(messages, start) + rounds: List[dict] = [] + for t in turns: + # skip pure 'user' boundary turns: nothing tool-verifiable there + if t.role_kind == 'user': + continue + r: Dict[str, Any] = {'messages': list(preamble) + list(t.messages)} + if segment.get('tools'): + r['tools'] = list(segment['tools']) + if segment.get('user_data'): + r['user_data'] = list(segment['user_data']) + rounds.append(r) + return rounds + + +def aggregate_hard_over_rounds( + round_scores: Sequence[RoundScore], + *, + how: str = 'gmean', +) -> float: + """Aggregate per-round hard scalars into one segment-level hard scalar. + + Default ``gmean`` (geometric mean) is intentionally harsher than plain mean: + a single badly-formed round meaningfully drags the segment, which matches + the intuition that one hallucinated/failed tool call hurts the sub-task. + """ + if not round_scores: + return 1.0 # no rounds to fault -> neutral (e.g. a text-only segment) + return _reduce([r.hard_scalar for r in round_scores], how) + + +# --------------------------------------------------------------------------- +# fusion (hard x soft) with short-circuit +# --------------------------------------------------------------------------- +def fuse_segment( + index: int, + round_scores: Sequence[RoundScore], + rubric_fn: Optional[Callable[[], Any]] = None, + *, + hard_agg: str = 'gmean', + fusion: str = 'product', + hard_floor: float = 0.25, + hard_ceil_skip: Optional[float] = None, + num_levels: int = NUM_LEVELS, +) -> SegmentScore: + """Fuse per-round hard scores with a (lazily computed) segment rubric score. + + Args: + index: segment index. + round_scores: per-round hard scores (already computed; cheap/code-only). + rubric_fn: zero-arg callable returning a rubric ScoreDetail (something + with a ``.scalar`` attribute) OR a float. Called ONLY when the soft + chain is not short-circuited — this is what saves the long LLM path. + hard_agg: reducer for per-round hard scores ('gmean'|'mean'|'min'|...). + fusion: how to combine hard_agg and rubric: + 'product' -> hard_agg * rubric (hard acts as a floor/gatekeeper) + 'min' -> min(hard_agg, rubric) + 'mean' -> (hard_agg + rubric)/2 + 'hard_only'-> ignore rubric entirely + hard_floor: if aggregated hard score < this, SHORT-CIRCUIT: skip rubric, + segment score = hard_agg (the answer is already decided as bad). + hard_ceil_skip: (optional) if aggregated hard score >= this, also skip + rubric and use hard_agg. Set None to disable. Useful for + trivially-good tool-only segments where soft quality adds little. + num_levels: level discretization. + """ + hard_agg_val = aggregate_hard_over_rounds(round_scores, how=hard_agg) + + short = False + rubric_scalar: Optional[float] = None + + if fusion == 'hard_only' or rubric_fn is None: + scalar = hard_agg_val + short = True + elif hard_agg_val < hard_floor: + # bad hard signal -> don't waste the soft chain + scalar = hard_agg_val + short = True + elif hard_ceil_skip is not None and hard_agg_val >= hard_ceil_skip: + scalar = hard_agg_val + short = True + else: + rubric_scalar = _rubric_scalar(rubric_fn()) + scalar = _combine(hard_agg_val, rubric_scalar, fusion) + + return SegmentScore( + index=index, + scalar=scalar, + level=scalar_to_level(scalar, num_levels), + hard_scalar=hard_agg_val, + rubric_scalar=rubric_scalar, + short_circuited=short, + n_rounds=len(round_scores), + rounds=list(round_scores), + ) + + +def _rubric_scalar(result: Any) -> float: + if result is None: + return 0.0 + if isinstance(result, (int, float)): + return float(result) + scalar = getattr(result, 'scalar', None) + return float(scalar) if scalar is not None else 0.0 + + +def _combine(hard: float, soft: float, fusion: str) -> float: + if fusion == 'product': + return hard * soft + if fusion == 'min': + return min(hard, soft) + if fusion == 'mean': + return (hard + soft) / 2.0 + raise ValueError(f'unknown fusion {fusion!r}') + + +# --------------------------------------------------------------------------- +# segment -> trajectory +# --------------------------------------------------------------------------- +def aggregate_trajectory( + segment_scores: Sequence[SegmentScore], + *, + how: str = 'mean', + weight_by_rounds: bool = True, + num_levels: int = NUM_LEVELS, +) -> TrajectoryScore: + """Aggregate segment scores into a trajectory score. + + Args: + how: reducer when ``weight_by_rounds`` is False. + weight_by_rounds: when True, weight each segment by its round count so + longer sub-tasks count proportionally (ignores ``how``, uses a + round-weighted mean). Text-only segments count as weight 1. + """ + if not segment_scores: + return TrajectoryScore(scalar=0.0, level=0, segments=[]) + if weight_by_rounds: + num = 0.0 + den = 0.0 + for s in segment_scores: + w = max(1, s.n_rounds) + num += s.scalar * w + den += w + scalar = num / den if den else 0.0 + else: + scalar = _reduce([s.scalar for s in segment_scores], how) + return TrajectoryScore( + scalar=scalar, + level=scalar_to_level(scalar, num_levels), + segments=list(segment_scores), + ) diff --git a/src/twinkle_agentic/verifier/rubric_verifier.py b/src/twinkle_agentic/verifier/rubric_verifier.py index 215e32acb..9c505c1a9 100644 --- a/src/twinkle_agentic/verifier/rubric_verifier.py +++ b/src/twinkle_agentic/verifier/rubric_verifier.py @@ -160,6 +160,7 @@ def __init__( margin_threshold: float = 0.25, max_votes: int = 5, gate: bool = True, + fixed_rubric: Optional[List['RubricItem']] = None, ): if max_rubrics < min_rubrics: raise ValueError('max_rubrics must be >= min_rubrics') @@ -184,6 +185,9 @@ def __init__( self.margin_threshold = float(margin_threshold) self.max_votes = int(max_votes) self.gate = bool(gate) + # When provided, skip stage-1 rubric generation and score against these + # fixed criteria (e.g. a safety rubric — AUDIT D8). + self.fixed_rubric: Optional[List['RubricItem']] = list(fixed_rubric) if fixed_rubric else None # ------------------------------------------------------------------ # public entry points @@ -192,20 +196,39 @@ def __call__(self, trajectory: dict, **kwargs) -> int: return self.score_detail(trajectory, **kwargs).level def score_detail(self, trajectory: dict, *, query: Optional[str] = None, - sampling_params: Any = None) -> ScoreDetail: + sampling_params: Any = None, + extra_context: Optional[str] = None) -> ScoreDetail: query = query or self._infer_query(trajectory) segment_text = self._render_segment(trajectory) + # D7c: fold an objective finding into the scored transcript so the judge + # re-scores WITH the hard evidence in view (objective corrects subjective). + if extra_context: + segment_text = f'{segment_text}\n\n[OBJECTIVE EVIDENCE]\n{extra_context}' # --- code-level hard verification (free, un-hackable) --- hard_pass_rate, has_hard = self._code_hard_checks(trajectory) - # --- stage 1: rubric generation (distilled) --- - raw_rubric = self._gen_rubric( - trajectory=self._gen_trajectory(query, segment_text), - sampling_params=self._gen_sampling_params(sampling_params), - query=query, - ) - rubric = self._parse_rubric(raw_rubric) + # No LLM (no student sampler AND no teacher API): skip both LLM stages and + # fall back to the deterministic code signal, instead of letting the + # llm_backup teacher path raise a missing-credentials error. + if not self._llm_available(): + scalar = hard_pass_rate if has_hard else 0.0 + return ScoreDetail( + level=self._to_level(scalar), scalar=scalar, llm_scalar=0.0, + hard_pass_rate=hard_pass_rate if has_hard else 1.0, + gated=False, n_votes=0, rubric=[], + ) + + # --- stage 1: rubric (fixed if configured, else distilled generation) --- + if self.fixed_rubric is not None: + rubric = list(self.fixed_rubric) + else: + raw_rubric = self._gen_rubric( + trajectory=self._gen_trajectory(query, segment_text), + sampling_params=self._gen_sampling_params(sampling_params), + query=query, + ) + rubric = self._parse_rubric(raw_rubric) if not rubric: # No usable rubric: fall back to the code signal alone. scalar = hard_pass_rate if has_hard else 0.0 @@ -381,6 +404,20 @@ def _vote_rates(votes: List[List[bool]]) -> List[float]: # ------------------------------------------------------------------ # LLM sampling plumbing (mirrors Summarizer) # ------------------------------------------------------------------ + def _llm_available(self) -> bool: + """True if a student sampler exists or a teacher API is configured. + + Mirrors the env vars ``llm_backup`` uses for its teacher; when neither a + student nor a teacher is present we must not attempt an LLM call (it would + raise a missing-credentials error inside the llm_backup teacher path). + """ + if self.sampler is not None: + return True + import os + return bool(os.environ.get('LLM_BACKUP_API_KEY') + or os.environ.get('OPENAI_API_KEY') + or os.environ.get('LLM_BACKUP_BASE_URL')) + def _sample_text(self, trajectory, sampling_params, lora_path) -> str: if self.sampler is None: return '' From 627869b39995d22b5194d6f11f3ca9f6f7a79a16 Mon Sep 17 00:00:00 2001 From: tastelikefeet <58414341+tastelikefeet@users.noreply.github.com> Date: Thu, 9 Jul 2026 12:19:01 +0800 Subject: [PATCH 04/60] Unfinished code (#244) --- .gitignore | 2 +- 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 | 362 +++- cookbook/exp/embedding/compare_math_levels.py | 91 + cookbook/exp/embedding/dataset_hard.py | 202 +++ cookbook/exp/embedding/dataset_index.py | 2 +- cookbook/exp/embedding/eval_gpqa_rag.py | 1490 +++++++++++++++++ cookbook/exp/embedding/eval_math_by_level.sh | 59 + cookbook/exp/embedding/eval_rag_recall.py | 187 +++ .../exp/embedding/make_embedding_dataset.py | 758 +++++++++ .../exp/embedding/train_embedding_full_ddp.py | 779 ++------- cookbook/exp/rl/grpo.py | 787 +++++++++ cookbook/exp/rl/rag_hint_grpo.py | 1480 ++++++++++++++++ cookbook/sample/emb_sample.py | 8 +- cookbook/sample/rag_recall_sample.py | 379 +++++ src/twinkle/loss/grpo.py | 26 +- src/twinkle/template/base.py | 5 +- 21 files changed, 6009 insertions(+), 754 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/compare_math_levels.py create mode 100644 cookbook/exp/embedding/dataset_hard.py create mode 100644 cookbook/exp/embedding/eval_gpqa_rag.py create mode 100755 cookbook/exp/embedding/eval_math_by_level.sh create mode 100644 cookbook/exp/embedding/eval_rag_recall.py create mode 100644 cookbook/exp/embedding/make_embedding_dataset.py create mode 100644 cookbook/exp/rl/grpo.py create mode 100644 cookbook/exp/rl/rag_hint_grpo.py create mode 100644 cookbook/sample/rag_recall_sample.py diff --git a/.gitignore b/.gitignore index 222caa50b..db4719f34 100644 --- a/.gitignore +++ b/.gitignore @@ -49,7 +49,7 @@ htmlcov/ .tox/ .coverage .coverage.* -.cache +.cache* nosetests.xml coverage.xml *.cover diff --git a/cookbook/exp/embedding/ablation_all.sh b/cookbook/exp/embedding/ablation_all.sh new file mode 100755 index 000000000..2be4c7fc7 --- /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 000000000..f5e5d1456 --- /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 000000000..fb47718e6 --- /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 000000000..defa863cd --- /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 000000000..86e35bd1e --- /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 d228a597a..a71bae060 100644 --- a/cookbook/exp/embedding/build_thinking_rag_index.py +++ b/cookbook/exp/embedding/build_thinking_rag_index.py @@ -39,6 +39,8 @@ import os import re import sys +import threading +import time from pathlib import Path from typing import Any, Dict, Iterator, List, Optional, Tuple @@ -163,7 +165,7 @@ EMBED_MODEL_ID = os.environ.get( 'EMBED_MODEL_ID', - 'output/embedding_lora_transformers/step_4000', + 'output/embedding_full_transformers/last-checkpoint', ) CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') @@ -185,30 +187,42 @@ SIM_THRESHOLD = float(os.environ.get('SIM_THRESHOLD', 0.65)) MIN_TEXT_CHARS = int(os.environ.get('MIN_TEXT_CHARS', 256)) +# Dataset mix caps (only used in 'both' mode). None = no cap. +THINK_CAP: Optional[int] = int(os.environ.get('THINK_CAP', 400_000)) or None +INDEX_CAP: Optional[int] = int(os.environ.get('INDEX_CAP', 400_000)) or None +MIX_SHUFFLE_SEED = 100 + +# Concurrency knobs for API fallback and prefetch pipeline. +API_CONCURRENCY = int(os.environ.get('API_CONCURRENCY', 8)) +API_MIN_INTERVAL = float(os.environ.get('API_MIN_INTERVAL', 0.1)) +PREFETCH_WORKERS = int(os.environ.get('PREFETCH_WORKERS', 2)) + # Hard-templated hints: the condenser SFT prior maps `Skill` to the legacy # `Use when: / numbered steps / Output:` skeleton on long inputs; embedding the # exact 4-line body template + explicit negative constraints is the only way to # override it deterministically across query and cot sides. RAG_QUERY_HINT = ( - 'Summarize this query for retrieval. ' - 'The body of ## Summary MUST follow this EXACT 4-line template — ' - 'do NOT emit "Use when:", numbered procedure steps, or "Output:":\n' - 'Topic: \n' - 'Problem: \n' - 'Skill: \n' - 'Knowledge: \n' + 'Extract the abstract PROBLEM TYPE from this query. ' + 'IGNORE all specific numbers, values, variable names, and parameters — ' + 'focus ONLY on the CLASS of problem and the METHODOLOGY required. ' + 'The body of ## Summary MUST follow this EXACT 4-line template:\n' + 'Topic: \n' + 'Problem: \n' + 'Skill: \n' + 'Knowledge: \n' 'Then emit the mandatory ## More section as usual. ' - 'Topic must name the specific pattern, never generic labels.') + 'Topic must name the method class, never mention specific numbers.') RAG_THINKING_HINT = ( - 'Summarize this reasoning trace for retrieval. ' - 'The body of ## Summary MUST follow this EXACT 4-line template — ' - 'do NOT emit "Use when:", numbered procedure steps, or "Output:":\n' - 'Topic: \n' - 'Problem: \n' - 'Skill: \n' - 'Knowledge: \n' + 'Extract the abstract METHODOLOGY demonstrated in this solution. ' + 'IGNORE all specific numbers, values, and computed results — ' + 'focus ONLY on the general TECHNIQUE and key reasoning STEPS. ' + 'The body of ## Summary MUST follow this EXACT 4-line template:\n' + 'Topic: \n' + 'Problem: \n' + 'Skill: \n' + 'Knowledge: \n' 'Then emit the mandatory ## More section as usual. ' - 'Topic must name the specific pattern, never generic labels.') + 'Topic must name the method class, never mention specific numbers.') # OpenAI API fallback (used when vLLM truncates). COMPRESS_API_KEY = os.environ.get('COMPRESS_API_KEY', '') @@ -419,23 +433,124 @@ def _api_compress(api: OpenAIClient, messages: List[Dict[str, str]]) -> Optional return _strip_outer_codefence(content) +_api_throttle_lock = threading.Lock() +_api_last_call = [0.0] + + +def _api_throttle(): + with _api_throttle_lock: + gap = time.monotonic() - _api_last_call[0] + if gap < API_MIN_INTERVAL: + time.sleep(API_MIN_INTERVAL - gap) + _api_last_call[0] = time.monotonic() + + +def _api_compress_throttled(api: OpenAIClient, messages: List[Dict[str, str]]) -> Optional[str]: + """Rate-limited API compression call.""" + _api_throttle() + return _api_compress(api, messages) + + def _resolve_compressed(sampler: vLLMSampler, api: Optional[OpenAIClient], texts: List[str], query_hint: str) -> List[Optional[str]]: - """Run vLLM batch; replace truncations / skeleton-incomplete with API output.""" + """Run vLLM batch; replace truncations / skeleton-incomplete with API output. + + API fallback runs concurrently (up to API_CONCURRENCY workers) for speed. + """ pairs = _vllm_compress(sampler, texts, query_hint) - results: List[Optional[str]] = [] - for (text, stop), src_text in zip(pairs, texts): + results: List[Optional[str]] = [None] * len(texts) + fallback_indices: List[int] = [] + for i, ((text, stop), src_text) in enumerate(zip(pairs, texts)): if stop != 'length' and not _is_truncated_compression(text): - results.append(text) - continue - if api is None: - results.append(None) + results[i] = text + else: + fallback_indices.append(i) + + if fallback_indices and api is not None: + from concurrent.futures import ThreadPoolExecutor, as_completed + with ThreadPoolExecutor(max_workers=API_CONCURRENCY) as pool: + futures = {} + for idx in fallback_indices: + msgs = _build_compress_messages(texts[idx], query_hint) + futures[pool.submit(_api_compress_throttled, api, msgs)] = idx + for fut in as_completed(futures): + idx = futures[fut] + api_text = fut.result() + if api_text and not _is_truncated_compression(api_text): + results[idx] = api_text + + return results + + +def _resolve_compressed_multi(sampler: vLLMSampler, api: Optional[OpenAIClient], + texts: List[str], hints: List[str]) -> List[Optional[str]]: + """Like _resolve_compressed but each text has its own per-item hint. + + Merges all texts into a SINGLE vLLM batch call (instead of one per hint), + dramatically reducing round-trip overhead when processing interleaved + query+cot pairs with different hint strings. + + Args: + sampler: vLLM condenser sampler. + api: Optional OpenAI-compatible API client for fallback. + texts: List of raw texts to compress (may contain empty strings to skip). + hints: Per-text hint strings (same length as texts). + + Returns: + List of compressed texts (None where compression failed entirely). + """ + assert len(texts) == len(hints), f'texts({len(texts)}) != hints({len(hints)})' + if not texts: + return [] + + # Skip texts that would exceed the condenser's context window. + _max_input_chars = (CONDENSE_MAX_MODEL_LEN - CONDENSE_MAX_TOKENS) * 3 + skip_mask = [len(t) > _max_input_chars for t in texts] + + # Build prompts per-item (each text gets its own hint as the query parameter). + prompts = [{'messages': _build_compress_messages(t, h)} + for t, h, skip in zip(texts, hints, skip_mask) if not skip] + active_indices = [i for i, skip in enumerate(skip_mask) if not skip] + params = TwinkleSamplingParams( + max_tokens=CONDENSE_MAX_TOKENS, + temperature=COMPRESS_TEMPERATURE, + top_p=COMPRESS_TOP_P, + num_samples=1, + ) + + # Single vLLM batch call — the key throughput win. + responses = sampler.sample(prompts, params) if prompts else [] + + results: List[Optional[str]] = [None] * len(texts) + fallback_indices: List[int] = [] + for resp_idx, orig_idx in enumerate(active_indices): + resp = responses[resp_idx] + seq = resp.sequences[0] if resp and resp.sequences else None + if seq is None: + fallback_indices.append(orig_idx) continue - api_text = _api_compress(api, _build_compress_messages(src_text, query_hint)) - if api_text is None or _is_truncated_compression(api_text): - results.append(None) + text = seq.decoded or '' + text = re.sub(r'<\|[^|]+\|>', '', text).rstrip() + text = _strip_outer_codefence(text) + if seq.stop_reason != 'length' and not _is_truncated_compression(text): + results[orig_idx] = text else: - results.append(api_text) + fallback_indices.append(orig_idx) + + # Concurrent API fallback for failed items. + if fallback_indices and api is not None: + from concurrent.futures import ThreadPoolExecutor, as_completed + with ThreadPoolExecutor(max_workers=API_CONCURRENCY) as pool: + futures = {} + for idx in fallback_indices: + msgs = _build_compress_messages(texts[idx], hints[idx]) + futures[pool.submit(_api_compress_throttled, api, msgs)] = idx + for fut in as_completed(futures): + idx = futures[fut] + api_text = fut.result() + if api_text and not _is_truncated_compression(api_text): + results[idx] = api_text + return results @@ -545,7 +660,7 @@ def _existing_ids(table) -> set: def _stream_corpus(total: Optional[int], load_from_cache_file: bool, max_rows: int = 0) -> Iterator[Dict[str, Any]]: - ds = _GET_DATASET(total=total, load_from_cache_file=load_from_cache_file) + ds = _GET_DATASET(total=total or None, load_from_cache_file=load_from_cache_file) n_full = len(ds) cap = max_rows if (max_rows and max_rows < n_full) else n_full sys.stderr.write(f'[corpus] get_dataset: {n_full} rows' @@ -602,32 +717,58 @@ def build_index(args: argparse.Namespace, # ---- Streaming loop ----------------------------------------------------- n_seen = n_kept = n_dropped_short = n_dropped_compress = n_dropped_sim = 0 n_dropped_dup = 0 - pbar = tqdm(desc='index', unit='row', dynamic_ncols=True) + n_no_id = 0 + n_no_query = 0 + n_short_cot = 0 + _diag_samples = 5 # print first N dropped rows for diagnosis batch: List[Dict[str, Any]] = [] - def _flush(rows: List[Dict[str, Any]]) -> None: - nonlocal n_kept, n_dropped_compress, n_dropped_sim + def _compress_batch(rows: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Phase 1: compress query+cot in a SINGLE merged vLLM call for throughput.""" if not rows: - return - # Phase 1 — compress query (RAG_QUERY_HINT) and cot (RAG_THINKING_HINT). - # Short queries bypass condenser (passthrough) — matches training behaviour. - long_q_indices = [i for i, r in enumerate(rows) if len(r['query_raw']) >= MIN_TEXT_CHARS] - q_compressed: List[Optional[str]] = [None] * len(rows) - for i, r in enumerate(rows): - if len(r['query_raw']) < MIN_TEXT_CHARS: - q_compressed[i] = r['query_raw'] - if long_q_indices: - long_results = _resolve_compressed( - sampler, api, [rows[i]['query_raw'] for i in long_q_indices], RAG_QUERY_HINT) - for idx, res in zip(long_q_indices, long_results): - q_compressed[idx] = res - c_compressed = _resolve_compressed( - sampler, api, [r['cot_raw'] for r in rows], RAG_THINKING_HINT) + return [] + # Build a merged prompt list: interleave query and cot texts so the sampler + # processes both in one round-trip instead of two serial calls. + all_texts: List[str] = [] + all_hints: List[str] = [] + passthrough_map: Dict[int, str] = {} # prompt_idx → raw text for short queries + for r in rows: + q_raw = r['query_raw'] + if len(q_raw) < MIN_TEXT_CHARS: + passthrough_map[len(all_texts)] = q_raw + all_texts.append('') # placeholder + all_hints.append(RAG_QUERY_HINT) + else: + all_texts.append(q_raw) + all_hints.append(RAG_QUERY_HINT) + all_texts.append(r['cot_raw']) + all_hints.append(RAG_THINKING_HINT) + + # Split into passthrough vs sampler-needed + sampler_indices = [i for i in range(len(all_texts)) if i not in passthrough_map] + sampler_texts = [all_texts[i] for i in sampler_indices] + sampler_hints = [all_hints[i] for i in sampler_indices] + + # Single merged vLLM call — group by hint to maximize prefix-sharing + # (both hints produce the same COMPRESS_SYSTEM, so batching is efficient). + sampler_results = _resolve_compressed_multi( + sampler, api, sampler_texts, sampler_hints) + + # Reassemble full results + all_results: List[Optional[str]] = [None] * len(all_texts) + for idx, text in passthrough_map.items(): + all_results[idx] = text + for pos, res in zip(sampler_indices, sampler_results): + all_results[pos] = res + + # Pair up (query, cot) and filter kept_rows: List[Dict[str, Any]] = [] - for r, q_cmp, c_cmp in zip(rows, q_compressed, c_compressed): + for i, r in enumerate(rows): + q_cmp = all_results[i * 2] + c_cmp = all_results[i * 2 + 1] if not q_cmp or not c_cmp: - n_dropped_compress += 1 + nonlocal_counters['n_dropped_compress'] += 1 _log_miss(misses_path, misses_lock, { 'id': r['id'], 'source': r['source'], 'reason': 'compress_fail', 'query_raw_head': _short(r['query_raw'], 200), @@ -637,15 +778,17 @@ def _flush(rows: List[Dict[str, Any]]) -> None: r['query_compressed'] = q_cmp r['cot_compressed'] = c_cmp kept_rows.append(r) + return kept_rows + + def _embed_and_insert(kept_rows: List[Dict[str, Any]]) -> None: + """Phase 2+3: embed compressed texts and insert into LanceDB.""" if not kept_rows: return - # Phase 2 — encode anchor (compressed query) + positive (compressed cot). anchor_emb = get_embeddings( emb_model, emb_template, [r['query_compressed'] for r in kept_rows], role='anchor') positive_emb = get_embeddings( emb_model, emb_template, [r['cot_compressed'] for r in kept_rows], role='positive') sims = (anchor_emb * positive_emb).sum(axis=1).astype(np.float32) - # Phase 3 — sim filter + LanceDB insert. to_insert: List[Dict[str, Any]] = [] for idx, (r, sim_val) in enumerate(zip(kept_rows, sims)): tag = 'KEEP' if sim_val >= SIM_THRESHOLD else 'DROP' @@ -653,7 +796,7 @@ def _flush(rows: List[Dict[str, Any]]) -> None: f'q={_short(r["query_raw"], 60)!r} ' f'cot={_short(r["cot_raw"], 60)!r}', flush=True) if sim_val < SIM_THRESHOLD: - n_dropped_sim += 1 + nonlocal_counters['n_dropped_sim'] += 1 _log_miss(misses_path, misses_lock, { 'id': r['id'], 'source': r['source'], 'reason': 'sim_low', 'sim': float(sim_val), @@ -677,24 +820,60 @@ def _flush(rows: List[Dict[str, Any]]) -> None: }) if to_insert: tbl.add(to_insert) - n_kept += len(to_insert) + nonlocal_counters['n_kept'] += len(to_insert) indexed.update(r['id'] for r in to_insert) + def _process_batch(rows: List[Dict[str, Any]]) -> None: + """Full pipeline for one batch: compress → embed → insert.""" + kept = _compress_batch(rows) + _embed_and_insert(kept) + + # Mutable counters shared with nested functions (avoid nonlocal limitation). + nonlocal_counters = { + 'n_kept': 0, 'n_dropped_compress': 0, 'n_dropped_sim': 0, + } + + from concurrent.futures import ThreadPoolExecutor as _PrefetchPool + prefetch_pool = _PrefetchPool(max_workers=PREFETCH_WORKERS) + try: + # Phase 1: Stream corpus, filter rows, collect batches (fast). + pending_futures = [] + sys.stderr.write('[build] streaming corpus and submitting batches...\n') + for row in _stream_corpus(total=args.total, load_from_cache_file=not args.no_cache, max_rows=args.max_rows): n_seen += 1 - if args.limit and n_kept >= args.limit: + if args.limit and nonlocal_counters['n_kept'] >= args.limit: break rid = row.get('id') or '' if not rid: + n_no_id += 1 + if n_no_id <= _diag_samples: + sys.stderr.write(f'[diag:no_id] row keys={list(row.keys())}\n') continue if rid in indexed: n_dropped_dup += 1 continue user_query, cot = _extract_query_cot(row) - if not user_query or len(cot) < MIN_TEXT_CHARS: + if not user_query: + n_no_query += 1 + n_dropped_short += 1 + if n_no_query <= _diag_samples: + msgs = row.get('messages') + sys.stderr.write( + f'[diag:no_query] id={rid} source={row.get("source","?")} ' + f'msgs_type={type(msgs).__name__} ' + f'msgs_len={len(msgs) if isinstance(msgs, list) else "?"} ' + f'msg0_keys={list(msgs[0].keys()) if isinstance(msgs, list) and msgs and isinstance(msgs[0], dict) else "?"}\n') + continue + if len(cot) < MIN_TEXT_CHARS: + n_short_cot += 1 n_dropped_short += 1 + if n_short_cot <= _diag_samples: + sys.stderr.write( + f'[diag:short_cot] id={rid} source={row.get("source","?")} ' + f'cot_len={len(cot)} query_len={len(user_query)}\n') continue batch.append({ 'id': rid, @@ -703,21 +882,45 @@ def _flush(rows: List[Dict[str, Any]]) -> None: 'cot_raw': cot, }) if len(batch) >= args.batch_size: - _flush(batch) + pending_futures.append(prefetch_pool.submit(_process_batch, list(batch))) batch.clear() - pbar.set_postfix(kept=n_kept, sim_drop=n_dropped_sim, - cmp_drop=n_dropped_compress, refresh=False) - pbar.update(1) + + # Flush remainder if batch: - _flush(batch) + pending_futures.append(prefetch_pool.submit(_process_batch, list(batch))) batch.clear() + + n_batches = len(pending_futures) + n_valid = n_seen - n_no_id - n_dropped_dup - n_dropped_short + sys.stderr.write( + f'[build] stream done: seen={n_seen} valid={n_valid} ' + f'batches={n_batches} (no_id={n_no_id} no_query={n_no_query} ' + f'short_cot={n_short_cot} dup={n_dropped_dup})\n') + + # Phase 2: Wait for all futures with real progress tracking. + pbar = tqdm(total=n_batches, desc='compress+embed', unit='batch', + dynamic_ncols=True) + for fut in pending_futures: + fut.result() + n_kept = nonlocal_counters['n_kept'] + n_dropped_sim = nonlocal_counters['n_dropped_sim'] + n_dropped_compress = nonlocal_counters['n_dropped_compress'] + pbar.set_postfix(kept=n_kept, sim_drop=n_dropped_sim, + cmp_drop=n_dropped_compress, refresh=False) + pbar.update(1) finally: pbar.close() + prefetch_pool.shutdown(wait=True) + + n_kept = nonlocal_counters['n_kept'] + n_dropped_sim = nonlocal_counters['n_dropped_sim'] + n_dropped_compress = nonlocal_counters['n_dropped_compress'] sys.stderr.write( - f'[build] seen={n_seen} kept={n_kept} sim_drop={n_dropped_sim} ' - f'cmp_drop={n_dropped_compress} short_drop={n_dropped_short} ' - f'dup_skip={n_dropped_dup}\n') + f'[build] summary: seen={n_seen} kept={n_kept} ' + f'dup={n_dropped_dup} no_id={n_no_id} no_query={n_no_query} ' + f'short_cot={n_short_cot} compress_fail={n_dropped_compress} ' + f'sim_drop={n_dropped_sim}\n') # ---- Build vector index for fast retrieval ------------------------------ if n_kept >= 64 and not args.skip_index: @@ -864,17 +1067,17 @@ def parse_args() -> argparse.Namespace: help='LanceDB table name within --db-path.') p.add_argument('--total', type=int, default=0, help='Total dataset rows to scale corpus to (0 = base sizes from the loader module).') - p.add_argument('--dataset-module', default='dataset_index', - choices=['dataset_index', 'dataset_think'], - help='Which loader to use: dataset_index (RAG profile) or ' - 'dataset_think (training mix).') + p.add_argument('--dataset-module', default='both', + choices=['dataset_index', 'dataset_think', 'both'], + help='Which loader to use: dataset_index (RAG profile), ' + 'dataset_think (training mix), or both (50/50 mix).') p.add_argument('--limit', type=int, default=0, help='Stop building once this many rows are kept (0 = no cap).') p.add_argument('--max-rows', type=int, default=0, help='Truncate corpus to this many rows AFTER get_dataset (0 = no cap). ' 'Use this instead of --total to avoid invalidating the dataset cache.') - p.add_argument('--batch-size', type=int, default=64, - help='Rows per condense+encode batch.') + p.add_argument('--batch-size', type=int, default=128, + help='Rows per condense+encode batch (larger = better GPU util).') p.add_argument('--no-cache', action='store_true', help='Disable load_from_cache_file in dataset_think.get_dataset.') p.add_argument('--overwrite', action='store_true', @@ -901,6 +1104,27 @@ def main() -> None: if args.dataset_module == 'dataset_think': from dataset_think import get_dataset as _swap _GET_DATASET = _swap + elif args.dataset_module == 'both': + from dataset_think import get_dataset as _get_think + from datasets import concatenate_datasets + + def _get_both(total=None, load_from_cache_file=True, **kw): + _total = total or None # CLI default 0 means "no scaling" → None + ds_index = _default_get_dataset(total=_total, load_from_cache_file=load_from_cache_file) + ds_think = _get_think(total=_total, load_from_cache_file=load_from_cache_file) + if INDEX_CAP and len(ds_index.dataset) > INDEX_CAP: + ds_index.dataset = ds_index.dataset.select(range(INDEX_CAP)) + if THINK_CAP and len(ds_think.dataset) > THINK_CAP: + ds_think.dataset = ds_think.dataset.select(range(THINK_CAP)) + n_index = len(ds_index.dataset) + n_think = len(ds_think.dataset) + ds_index.dataset = concatenate_datasets( + [ds_index.dataset, ds_think.dataset]).shuffle(seed=MIX_SHUFFLE_SEED) + sys.stderr.write(f'[mix] index={n_index} + think={n_think} ' + f'→ total={len(ds_index.dataset)}\n') + return ds_index + + _GET_DATASET = _get_both sys.stderr.write(f'[main] dataset loader: {args.dataset_module}\n') # Build/eval both depend on the same Twinkle stack — initialize once. diff --git a/cookbook/exp/embedding/compare_math_levels.py b/cookbook/exp/embedding/compare_math_levels.py new file mode 100644 index 000000000..d488909b9 --- /dev/null +++ b/cookbook/exp/embedding/compare_math_levels.py @@ -0,0 +1,91 @@ +"""Compare MATH direct vs RAG by difficulty level. + +Re-grades both result files with the production ``answers_match`` (so the +stored ``is_correct`` is never trusted) and prints the per-level accuracy +plus the RAG gain (delta) so you can see how it varies with difficulty. + +Defaults to the raw-RAG output (``math_rag_results.jsonl``); pass a second +arg to compare a different rag file (e.g. ``math_rag_hint_results.jsonl``). + +Usage: + python cookbook/exp/embedding/compare_math_levels.py \ + [direct.jsonl] [rag.jsonl] +""" +import importlib.util +import json +import os +import sys +from collections import defaultdict + +_HERE = os.path.dirname(os.path.abspath(__file__)) + + +def _load_grader(): + spec = importlib.util.spec_from_file_location( + 'egr', os.path.join(_HERE, 'eval_gpqa_rag.py')) + egr = importlib.util.module_from_spec(spec) + spec.loader.exec_module(egr) + return egr.answers_match + + +def _load(path): + return {json.loads(l)['idx']: json.loads(l) + for l in open(path, encoding='utf-8') if l.strip()} + + +def main(): + direct_path = sys.argv[1] if len(sys.argv) > 1 else \ + './output/thinking_rag/math_direct_results.jsonl' + hint_path = sys.argv[2] if len(sys.argv) > 2 else \ + './output/thinking_rag/math_rag_results.jsonl' + + answers_match = _load_grader() + D = _load(direct_path) + H = _load(hint_path) + common = sorted(set(D) & set(H)) + print(f'direct={len(D)} rag+hint={len(H)} common={len(common)}') + + def runaway(rec): + mo = rec.get('model_output') or '' + return ('' not in mo) or ( + not (rec.get('predicted') or '').strip() and len(mo) > 40000) + + def correct(rec): + return answers_match(rec.get('predicted') or '', + rec['reference_answer']) + + # level -> counters + per = defaultdict(lambda: {'n': 0, 'd': 0, 'h': 0, + 'd_run': 0, 'h_run': 0}) + for i in common: + lv = H[i].get('level') or D[i].get('level') or 'Unknown' + c = per[lv] + c['n'] += 1 + c['d'] += int(correct(D[i])) + c['h'] += int(correct(H[i])) + c['d_run'] += int(runaway(D[i])) + c['h_run'] += int(runaway(H[i])) + + print(f'\n{"level":>10} | {"n":>4} | {"direct":>7} | {"rag+hint":>8} | ' + f'{"delta":>7} | {"d_run":>6} | {"h_run":>6}') + print('-' * 68) + tot = {'n': 0, 'd': 0, 'h': 0, 'd_run': 0, 'h_run': 0} + for lv in sorted(per.keys()): + c = per[lv] + for k in tot: + tot[k] += c[k] + n = c['n'] + dacc, hacc = c['d'] / n, c['h'] / n + print(f'{lv:>10} | {n:>4} | {dacc:>7.3f} | {hacc:>8.3f} | ' + f'{hacc - dacc:>+7.3f} | {c["d_run"]/n:>6.1%} | ' + f'{c["h_run"]/n:>6.1%}') + print('-' * 68) + n = tot['n'] + if n: + print(f'{"OVERALL":>10} | {n:>4} | {tot["d"]/n:>7.3f} | ' + f'{tot["h"]/n:>8.3f} | {(tot["h"]-tot["d"])/n:>+7.3f} | ' + f'{tot["d_run"]/n:>6.1%} | {tot["h_run"]/n:>6.1%}') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/embedding/dataset_hard.py b/cookbook/exp/embedding/dataset_hard.py new file mode 100644 index 000000000..9fa059b95 --- /dev/null +++ b/cookbook/exp/embedding/dataset_hard.py @@ -0,0 +1,202 @@ +"""Hard-negative dataset for embedding training. + +Provides ReasonIR (AI-ModelScope/reasonir-data, hq subset): + - query: reasoning-intensive question + - positive: BRIGHT document (resolved via xlangai/BRIGHT documents corpus) + - negatives: plausibly related but ultimately unhelpful documents + +Output schema: ``{id, source, query, cot, response, negatives}`` +where ``negatives`` is a list of strings (each a separate hard negative). +""" +import hashlib +import os +import sys +from pathlib import Path +from typing import Any, Dict, List, Optional + +from datasets import Dataset as HFDataset +from modelscope import MsDataset + +_CACHE_DIR = Path(__file__).resolve().parent / '.cache_hard' + + +def _hash_id(prefix: str, content: str) -> str: + return f'{prefix}__{hashlib.md5(content.encode("utf-8")).hexdigest()[:16]}' + + +# --------------------------------------------------------------------------- +# BRIGHT document corpus (lazy singleton) +# --------------------------------------------------------------------------- + +_BRIGHT_SPLITS = [ + 'aops', 'biology', 'earth_science', 'economics', 'leetcode', 'pony', + 'psychology', 'robotics', 'stackoverflow', 'sustainable_living', + 'theoremqa_questions', 'theoremqa_theorems', +] + +_bright_docs: Optional[Dict[str, str]] = None + + +def _load_bright_docs() -> Dict[str, str]: + """Load all BRIGHT document splits into {id -> content} lookup dict.""" + global _bright_docs + if _bright_docs is not None: + return _bright_docs + sys.stderr.write('[dataset_hard] Loading BRIGHT documents corpus...\n') + _bright_docs = {} + for split in _BRIGHT_SPLITS: + try: + ds = MsDataset.load( + 'xlangai/BRIGHT', subset_name='documents', split=split, + download_mode='reuse_dataset_if_exists') + for row in ds: + doc_id = row.get('id', '') + content = row.get('content', '') + if doc_id and content: + _bright_docs[doc_id] = content + short = doc_id.rsplit('/', 1)[-1] if '/' in doc_id else doc_id + if short not in _bright_docs: + _bright_docs[short] = content + sys.stderr.write(f' [{split}] loaded {len(ds)} docs\n') + except Exception as e: + sys.stderr.write(f' [{split}] FAILED: {e}\n') + sys.stderr.write(f'[dataset_hard] BRIGHT total: {len(_bright_docs)} entries\n') + return _bright_docs + + +# --------------------------------------------------------------------------- +# ReasonIR dataset +# --------------------------------------------------------------------------- + +def get_dataset_reasonir(max_rows: Optional[int] = None, + max_negatives: int = 16, + load_from_cache_file: bool = True) -> HFDataset: + """Load AI-ModelScope/reasonir-data (hq subset) with BRIGHT doc resolution. + + Schema: {id, source, query, cot, response, negatives} + """ + cache_key = f'reasonir_neg{max_negatives}' + cache_path = _CACHE_DIR / cache_key + if load_from_cache_file and cache_path.exists(): + sys.stderr.write(f'[reasonir] loading from cache: {cache_path}\n') + ds = HFDataset.load_from_disk(str(cache_path)) + if max_rows and len(ds) > max_rows: + ds = ds.select(range(max_rows)) + sys.stderr.write(f'[reasonir] {len(ds)} rows (cached)\n') + return ds + + ds = MsDataset.load( + 'AI-ModelScope/reasonir-data', subset_name='hq', split='train', + download_mode='reuse_dataset_if_exists') + if max_rows and len(ds) > max_rows: + ds = ds.select(range(max_rows)) + + bright = _load_bright_docs() + rows = [] + n_miss = 0 + for row in ds: + query_parts = row.get('query', []) + if not isinstance(query_parts, list) or len(query_parts) < 2: + continue + query = query_parts[1].strip() + if not query: + continue + + pos_list = row.get('pos', []) + if not pos_list: + continue + pos_id = pos_list[0][1] if isinstance(pos_list[0], list) and len(pos_list[0]) > 1 else '' + cot = bright.get(pos_id, '') + if not cot: + n_miss += 1 + continue + + neg_list = row.get('neg', []) + negatives = [] + for neg in neg_list: + if isinstance(neg, list) and len(neg) > 1: + neg_text = neg[1].strip() + if neg_text: + negatives.append(neg_text) + if len(negatives) >= max_negatives: + break + + if not negatives: + continue + + rows.append({ + 'id': _hash_id('reasonir', f'{query}\n{pos_id}'), + 'source': 'reasonir-hq', + 'query': query, + 'cot': cot, + 'response': '', + 'negatives': negatives, + }) + + if n_miss: + sys.stderr.write(f'[reasonir] {n_miss} rows skipped (BRIGHT doc not found)\n') + sys.stderr.write(f'[reasonir] {len(rows)} rows with hard negatives\n') + result = HFDataset.from_dict(_rows_to_cols(rows)) + # Persist full dataset; max_rows is applied post-cache for flexibility. + cache_path.parent.mkdir(parents=True, exist_ok=True) + result.save_to_disk(str(cache_path)) + sys.stderr.write(f'[reasonir] cached to {cache_path}\n') + if max_rows and len(result) > max_rows: + result = result.select(range(max_rows)) + return result + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _rows_to_cols(rows: List[Dict[str, Any]]) -> Dict[str, list]: + if not rows: + return {'id': [], 'source': [], 'query': [], 'cot': [], + 'response': [], 'negatives': []} + keys = rows[0].keys() + return {k: [r[k] for r in rows] for k in keys} + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + +def get_dataset( + reasonir_max: Optional[int] = None, + max_negatives: int = 16, + load_from_cache_file: bool = True, + **kwargs, +) -> HFDataset: + """Load hard-negative dataset (reasonir only). + + Returns HF Dataset with schema: {id, source, query, cot, response, negatives} + """ + ds = get_dataset_reasonir(max_rows=reasonir_max, max_negatives=max_negatives, + load_from_cache_file=load_from_cache_file) + if len(ds) == 0: + sys.stderr.write('[dataset_hard] WARNING: reasonir dataset empty\n') + else: + sys.stderr.write(f'[dataset_hard] reasonir={len(ds)}\n') + return ds + + +if __name__ == '__main__': + import argparse + parser = argparse.ArgumentParser() + parser.add_argument('--reasonir-max', type=int, default=1000) + args = parser.parse_args() + + ds = get_dataset(reasonir_max=args.reasonir_max) + print(f'Total rows: {len(ds)}') + print(f'Features: {ds.features}') + if len(ds) > 0: + row = ds[0] + print(f'\nSample[0]:') + print(f' id: {row["id"]}') + print(f' source: {row["source"]}') + print(f' query: {row["query"][:100]}...') + print(f' cot: {row["cot"][:100]}...') + print(f' negatives: {len(row["negatives"])} items') + if row['negatives']: + print(f' [0]: {row["negatives"][0][:80]}...') diff --git a/cookbook/exp/embedding/dataset_index.py b/cookbook/exp/embedding/dataset_index.py index c86e1c523..7d2905a59 100644 --- a/cookbook/exp/embedding/dataset_index.py +++ b/cookbook/exp/embedding/dataset_index.py @@ -704,7 +704,7 @@ def get_dataset(total: Optional[int] = None, ], dropped_log_path=dropped_log or '', ) - dataset.map(qp, batched=True, num_proc=32, load_from_cache_file=load_from_cache_file) + dataset.map(qp, num_proc=32, load_from_cache_file=load_from_cache_file) return dataset diff --git a/cookbook/exp/embedding/eval_gpqa_rag.py b/cookbook/exp/embedding/eval_gpqa_rag.py new file mode 100644 index 000000000..dd71590c0 --- /dev/null +++ b/cookbook/exp/embedding/eval_gpqa_rag.py @@ -0,0 +1,1490 @@ +"""Math evaluation: direct vs RAG-augmented with Qwen3.5-4B. + +Datasets (``--dataset``): + - ``math`` (default): MATH (Hendrycks), stratified by difficulty (Level 1-5) + so RAG gain can be plotted against difficulty. + - ``aops``: AoPS competition problems (metadata.boxed only). + +Modes (``--mode``): + - ``direct``: The model solves problems directly (4 GPUs, TP=4). + - ``rag`` (default): Retrieve top-k thinking traces from LanceDB, condense + them (API qwen3.7-max), inject as 1-shot examples, then solve + (8 GPUs: DP=4 embedding + TP=4 vLLM). + +Defaults implement **raw RAG on MATH**: ``--dataset math --mode rag --condense`` +with hint filtering OFF. The API condenser needs ``COMPRESS_API_KEY`` (or a +local condenser via ``EVAL_CONDENSER_GPUS``); otherwise pass ``--no-condense``. + +Optional ``--hint`` flag (rag mode only): + After retrieval + condensing, call an API model to filter and refine the + traces — keeping only applicable methods — then inject the refined trace. + +Reference answers are the ``\\boxed{...}`` content of each solution. + +Launch examples: + # Default: raw RAG on MATH, stratified 100/level (needs COMPRESS_API_KEY) + COMPRESS_API_KEY=sk-... python cookbook/exp/embedding/eval_gpqa_rag.py + + # Paired direct baseline on the same MATH subset + python cookbook/exp/embedding/eval_gpqa_rag.py --mode direct + + # Raw RAG without condenser (inject raw retrieved traces) + python cookbook/exp/embedding/eval_gpqa_rag.py --no-condense + + # Add hint filtering back on top of condensing + COMPRESS_API_KEY=sk-... python cookbook/exp/embedding/eval_gpqa_rag.py --hint + + # Fall back to the old AoPS dataset + python cookbook/exp/embedding/eval_gpqa_rag.py --dataset aops +""" +import argparse +import json +import os +import random +import re +import sys +import threading +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Any, Dict, List, Optional + +import numpy as np +import torch + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import SamplingParams as TwinkleSamplingParams +from twinkle.loss import InfonceLoss +from twinkle.model import TransformersModel +from twinkle.processor import InputProcessor +from twinkle.sampler import vLLMSampler +from twinkle.template import Qwen3_5Template +from twinkle_agentic.protocol.openai import OpenAI as OpenAIClient + +logger = get_logger() + +# -- Condenser config ---------------------------------------------------------- +CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') +CONDENSE_API_KEY = os.environ.get('COMPRESS_API_KEY', '') +CONDENSE_BASE_URL = os.environ.get('COMPRESS_BASE_URL', 'https://dashscope.aliyuncs.com/compatible-mode/v1') +CONDENSE_API_MODEL = os.environ.get('COMPRESS_MODEL', 'qwen3.7-max') +CONDENSE_API_CONCURRENCY = int(os.environ.get('API_CONCURRENCY', 32)) +CONDENSE_API_MIN_INTERVAL = float(os.environ.get('API_MIN_INTERVAL', 0.1)) +CONDENSE_TEMPERATURE = 0.2 +CONDENSE_MAX_TOKENS = 8192 + +# -- Hint analysis config ------------------------------------------------------ +HINT_ANALYSIS_MAX_TOKENS = int(os.environ.get('HINT_ANALYSIS_MAX_TOKENS', 2000)) +HINT_ANALYSIS_TEMPERATURE = 0.2 + +HINT_ANALYSIS_SYSTEM = ( + 'You are a mathematical reasoning trace filter. ' + 'Given a target problem and reasoning traces retrieved from SIMILAR (but different) problems, ' + 'your task is to FILTER and REFINE the traces into a clean reference.\n\n' + 'Rules:\n' + '1. KEEP: solution steps, methods, formulas, techniques, and key insights ' + 'that are directly applicable to solving the target problem.\n' + '2. REMOVE: problem-specific numeric calculations that do not transfer, ' + 'dead-end explorations, irrelevant approaches, verbose restatements, ' + 'and any content that would mislead the solver on the target problem.\n' + '3. Output the refined trace directly as actionable solution steps. ' + 'Preserve the original mathematical expressions and step structure.\n' + '4. Do NOT solve the target problem. Do NOT add your own solutions or commentary.\n' + '5. Do NOT output the answer to either problem.\n' + '6. If the traces are entirely irrelevant, output exactly: "No applicable methods."' +) + +HINT_ANALYSIS_USER = ( + '## Target Problem\n{query}\n\n' + '## Retrieved Reasoning Traces\n{thinking}' +) + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- +# -- Gen/Embed config --------------------------------------------------------- +GEN_MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3.5-4B') +EMBED_MODEL_ID = os.environ.get( + 'EMBED_MODEL_ID', 'output.oldemb/embedding_full_transformers/last-checkpoint') + +GEN_GPUS = int(os.environ.get('GEN_GPUS', 4)) +EMB_GPUS = int(os.environ.get('EMB_GPUS', 2)) +EMBED_MAX_LENGTH = int(os.environ.get('EMBED_MAX_LENGTH', 20000)) + +GEN_GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.85)) +GEN_MAX_MODEL_LEN = int(os.environ.get('GEN_MAX_MODEL_LEN', 65536)) +GEN_MAX_TOKENS = int(os.environ.get('GEN_MAX_TOKENS', 65536)) +GEN_TEMPERATURE = float(os.environ.get('GEN_TEMPERATURE', 0.6)) +GEN_TOP_P = float(os.environ.get('GEN_TOP_P', 0.95)) + +AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') +MATH_DATA_DIR = os.environ.get('MATH_DATA_DIR', './output/math_data/MATH') + + +# --------------------------------------------------------------------------- +# Condenser prompts & validation +# --------------------------------------------------------------------------- + +COMPRESS_SYSTEM = """\ +You are a reasoning-trace condenser. Given a verbose reasoning trace, \ +extract the TRANSFERABLE KNOWLEDGE as an EXECUTABLE SOLUTION SKELETON \ +that would help a reader solve SIMILAR problems in the same domain. + +Your output is the ENTIRE useful content — there is no expansion tool, no second pass. \ +The reader will apply this knowledge to a DIFFERENT problem, so focus on what transfers. + +Principles: +1. OUTPUT AN EXECUTABLE STEP CHAIN: numbered steps that a solver can directly follow. \ +Each step should state WHAT to do and HOW (with the formula/technique), not just \ +name the concept. +2. INCLUDE FULL FORMULAS: theorems, identities, inequalities — state each \ +with its COMPLETE MATHEMATICAL EXPRESSION, not just the name. +3. STATE APPLICABILITY: what structural features of a problem signal that this \ +approach works (e.g. "when the constraint is a sum of squares"). +4. PRESERVE KEY INSIGHTS: the non-obvious ideas or tricks that make the approach \ +work — the things a solver would NOT think of without guidance. +5. REMOVE: problem-specific numeric calculations, dead-end explorations, \ +hesitations, verbose restatements, and trivial arithmetic. +6. FORMAT: Start with a one-line "Applicability" statement, then numbered steps, \ +then key formulas. Keep it concise and actionable. +7. NO meta-commentary about the compression process. NO preamble. +""" + +COMPRESS_USER = ( + '## Reader Problem (context only — do NOT solve it)\n{query}\n\n' + '## Reasoning Trace to Condense\n{text}') + + +def _is_truncated_compression(text: str) -> bool: + if not text or not text.strip(): + return True + lines = [l.strip() for l in text.strip().splitlines() if l.strip()] + if len(lines) < 3: + return True + last_line = lines[-1] + # Truncated if last line looks incomplete (no terminal punctuation/formula) + if last_line and last_line[-1] not in '.。!！)）]】}\\$': + # Allow lines ending with numbers, boxed answers, etc. + if not re.search(r'\d$|\\boxed|\$|\)$', last_line): + return True + return False + + +# -- API rate limiter ---------------------------------------------------------- +_api_semaphore = threading.Semaphore(CONDENSE_API_CONCURRENCY) +_api_bucket_lock = threading.Lock() +_api_tokens = [float(CONDENSE_API_CONCURRENCY)] +_api_last_refill = [time.monotonic()] + + +def _api_throttle(): + _api_semaphore.acquire() + wait = 0.0 + try: + with _api_bucket_lock: + now = time.monotonic() + elapsed = now - _api_last_refill[0] + refill = elapsed / CONDENSE_API_MIN_INTERVAL + _api_tokens[0] = min(float(CONDENSE_API_CONCURRENCY), _api_tokens[0] + refill) + _api_last_refill[0] = now + if _api_tokens[0] >= 1.0: + _api_tokens[0] -= 1.0 + else: + wait = (1.0 - _api_tokens[0]) * CONDENSE_API_MIN_INTERVAL + _api_tokens[0] = 0.0 + finally: + _api_semaphore.release() + if wait > 0: + time.sleep(wait) + + +def _api_condense_single(api_client: OpenAIClient, messages: List[Dict]) -> Optional[str]: + _api_throttle() + trajectory = {'messages': messages} + sp = TwinkleSamplingParams(temperature=CONDENSE_TEMPERATURE, max_tokens=CONDENSE_MAX_TOKENS) + try: + reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) + except Exception as exc: + logger.warning(f'[condense-api] error: {exc}') + return None + content = (reply.get('content') or '').strip() + if not content: + return None + m = re.match(r'^```[a-zA-Z]*\n(.*?)\n```\s*$', content, re.DOTALL) + if m: + content = m.group(1).strip() + return content + + +def _api_hint_analysis_batch( + api_client: OpenAIClient, + problems: List[str], + condensed_examples: List[List[Dict[str, str]]], +) -> List[Optional[str]]: + """Call API to pre-analyze RAG relevance for each problem.""" + _MAX_HINT_INPUT = 8000 + results: List[Optional[str]] = [None] * len(problems) + tasks = [] + for i, prob in enumerate(problems): + if not condensed_examples[i]: + continue + traces = [ex.get('thinking', '') for ex in condensed_examples[i]] + merged_thinking = '\n---\n'.join(traces) + if len(merged_thinking) > _MAX_HINT_INPUT: + merged_thinking = merged_thinking[:_MAX_HINT_INPUT] + '\n[...truncated]' + user_msg = HINT_ANALYSIS_USER.format(query=prob, thinking=merged_thinking) + msgs = [ + {'role': 'system', 'content': HINT_ANALYSIS_SYSTEM}, + {'role': 'user', 'content': user_msg}, + ] + tasks.append((i, msgs)) + + if not tasks: + return results + + def _call_one(idx, msgs): + _api_throttle() + try: + trajectory = {'messages': msgs} + sp = TwinkleSamplingParams( + temperature=HINT_ANALYSIS_TEMPERATURE, + max_tokens=HINT_ANALYSIS_MAX_TOKENS) + reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) + content = (reply.get('content') or '').strip() + # Treat "No applicable methods." as empty (will trigger fallback) + if not content or content == 'No applicable methods.': + return idx, None + return idx, content + except Exception as exc: + logger.warning(f'[hint-analysis] error for idx={idx}: {exc}') + return idx, None + + with ThreadPoolExecutor(max_workers=min(len(tasks), CONDENSE_API_CONCURRENCY)) as pool: + futs = [pool.submit(_call_one, idx, msgs) for idx, msgs in tasks] + for fut in as_completed(futs): + idx, analysis = fut.result() + results[idx] = analysis + + n_success = sum(1 for r in results if r) + logger.info(f'[hint-analysis] completed {n_success}/{len(tasks)} analyses') + return results + + +# --------------------------------------------------------------------------- +# LLM-based decontamination +# --------------------------------------------------------------------------- + +_DECONTAM_JUDGE_PROMPT = ( + 'We are building a RAG-augmented math training system. Problem A is the test ' + 'question; Problem B was retrieved from a knowledge base.\n' + 'Answer YES only if A and B are essentially the SAME specific problem — ' + 'i.e. solving B directly gives you A\'s answer (just different wording/notation/' + 'format/negation).\n' + 'Answer NO if they merely share the same method/topic but have different ' + 'specific values, equations, or geometric configurations — learning B\'s ' + 'approach still requires independent work to solve A.\n' + 'Problem A: {prob_a}\n' + 'Problem B: {prob_b}\n' + 'Answer only YES or NO.' +) + + +def _llm_judge_same_problem( + api_client: OpenAIClient, pairs: List[tuple], +) -> List[bool]: + """Batch LLM judge: are (problem_a, problem_b) the same problem? + + Returns list of bools (True = same problem = should filter). + """ + if not pairs or not api_client: + return [False] * len(pairs) + + results = [False] * len(pairs) + + def _judge_one(idx, pa, pb): + prompt = _DECONTAM_JUDGE_PROMPT.format(prob_a=pa, prob_b=pb) + msgs = [{'role': 'user', 'content': prompt}] + _api_throttle() + try: + trajectory = {'messages': msgs} + sp = TwinkleSamplingParams(temperature=0.1, max_tokens=8) + reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) + answer = (reply.get('content') or '').strip().upper() + return idx, 'YES' in answer + except Exception: + return idx, False + + with ThreadPoolExecutor(max_workers=min(len(pairs), CONDENSE_API_CONCURRENCY)) as pool: + futs = [pool.submit(_judge_one, i, pa, pb) for i, (pa, pb) in enumerate(pairs)] + for fut in as_completed(futs): + idx, is_same = fut.result() + results[idx] = is_same + return results + + +def _llm_decontaminate( + api_client: OpenAIClient, + problems: List[str], + all_examples: List[List[Dict[str, str]]], +) -> List[List[Dict[str, str]]]: + """Apply LLM-based decontamination: remove retrievals judged as same problem.""" + judge_pairs = [] # (qi, ret_idx, prob_a, prob_b) + for qi, exs in enumerate(all_examples): + for ri, ex in enumerate(exs): + judge_pairs.append((qi, ri, problems[qi], ex.get('query', ''))) + + if not judge_pairs: + return all_examples + + pairs_input = [(pa, pb) for _, _, pa, pb in judge_pairs] + verdicts = _llm_judge_same_problem(api_client, pairs_input) + to_remove = set() + for vi, (qi, ri, _, _) in enumerate(judge_pairs): + if verdicts[vi]: + to_remove.add((qi, ri)) + + if to_remove: + logger.info(f'[decontam-llm] filtered {len(to_remove)} same-problem retrievals') + for qi in range(len(all_examples)): + all_examples[qi] = [ + ex for ri, ex in enumerate(all_examples[qi]) + if (qi, ri) not in to_remove + ] + return all_examples + + +def condense_traces( + examples_batch: List[List[Dict[str, str]]], + problems: List[str], + api_client: OpenAIClient, + condenser_sampler=None, + compress_params=None, + special_tokens: set = None, + max_output_len: int = 2000, + dp_size: int = 1, +) -> List[List[Dict[str, str]]]: + """Compress retrieved thinking traces with query-aware condenser. + + Primary: local vLLM condenser (if provided). + Fallback: API condenser. + Final fallback: raw trace truncated to max_output_len. + """ + result: List[List[Dict[str, str]]] = [] + # Flatten all (batch_idx, ex_idx, problem, example) for batch processing + tasks = [] + for bi, (exs, prob) in enumerate(zip(examples_batch, problems)): + for ei, ex in enumerate(exs): + tasks.append((bi, ei, prob, ex)) + + if not tasks: + return [[] for _ in examples_batch] + + # Build condense prompts (aligned with make_embedding_dataset.py hard path) + prompts = [] + for _, _, prob, ex in tasks: + user_msg = COMPRESS_USER.format(query=prob, text=ex['thinking']) + prompts.append([{'role': 'system', 'content': COMPRESS_SYSTEM}, + {'role': 'user', 'content': user_msg}]) + + # Phase 1: local vLLM condenser + condensed = [None] * len(tasks) + condense_sources = ['raw'] * len(tasks) + fallback_indices = [] + + if condenser_sampler is not None and compress_params is not None: + sampler_inputs = [{'messages': p} for p in prompts] + # The local vLLM sampler runs data-parallel across ``dp_size`` workers + # and requires at least one item per worker (it errors with + # "Batch too small for N workers" otherwise). Pad the batch up to a + # multiple of dp_size by repeating the last item, run, then keep only + # the first ``n_real`` responses and drop the padding. + n_real = len(sampler_inputs) + pad_size = 0 + if dp_size > 1 and n_real > 0 and n_real % dp_size != 0: + pad_size = dp_size - (n_real % dp_size) + sampler_inputs = sampler_inputs + [sampler_inputs[-1]] * pad_size + try: + responses = condenser_sampler.sample(sampler_inputs, compress_params) + except Exception as exc: + logger.warning(f'[condense] sampler error: {exc}') + responses = [None] * len(sampler_inputs) + if pad_size: + responses = responses[:n_real] + for ri, resp in enumerate(responses): + seq = resp.sequences[0] if resp and resp.sequences else None + text = '' + if seq and seq.stop_reason != 'length' and seq.decoded: + text = seq.decoded + if special_tokens: + for tok in special_tokens: + text = text.replace(tok, '') + text = text.rstrip() + if text and not _is_truncated_compression(text): + condensed[ri] = text + condense_sources[ri] = 'local' + else: + fallback_indices.append(ri) + else: + fallback_indices = list(range(len(tasks))) + + # Phase 2: API fallback + if fallback_indices and api_client: + with ThreadPoolExecutor(max_workers=CONDENSE_API_CONCURRENCY) as pool: + futures = {} + for ri in fallback_indices: + futures[pool.submit(_api_condense_single, api_client, prompts[ri])] = ri + for fut in as_completed(futures): + ri = futures[fut] + api_result = fut.result() + if api_result and not _is_truncated_compression(api_result): + condensed[ri] = api_result + condense_sources[ri] = 'api' + + # Phase 3: assemble results (fallback to raw truncation) + result = [[] for _ in examples_batch] + for ti, (bi, ei, prob, ex) in enumerate(tasks): + compressed = condensed[ti] + raw_len = len(ex['thinking']) + sim_val = ex.get('_sim', 0.0) + if compressed: + result[bi].append({'query': ex['query'], + 'thinking': _strip_condenser_markers(compressed), + '_condense_source': condense_sources[ti], + '_raw_trace_len': raw_len, '_sim': sim_val}) + else: + result[bi].append({'query': ex['query'], + 'thinking': ex['thinking'][:max_output_len], + '_condense_source': 'raw', + '_raw_trace_len': raw_len, '_sim': sim_val}) + + n_ok = sum(1 for c in condensed if c) + logger.info(f'[condense] {n_ok}/{len(tasks)} compressed ok, ' + f'{len(tasks) - n_ok} fell back to raw truncation') + return result + + +def _strip_condenser_markers(text: str) -> str: + """Light cleanup of condenser output. + + Removes any residual markdown headers or meta-lines that don't carry + solution content. Keeps numbered steps and equations intact. + """ + # Remove legacy ## headers if condenser still emits them + if '## More' in text: + text = text.split('## More', 1)[0] + text = re.sub(r'^##\s*Summary\s*\n?', '', text, flags=re.MULTILINE) + text = re.sub(r'^Topic:\s*.*\n?', '', text, flags=re.MULTILINE) + # Remove meta-commentary lines + text = re.sub(r'^\s*\(Note:.*\)\s*$', '', text, flags=re.MULTILINE) + return text.strip() + + +# --------------------------------------------------------------------------- +# Boxed answer extraction +# --------------------------------------------------------------------------- +_BOXED_RE = re.compile(r'\\boxed\s*\{') + + +def extract_boxed(text: str) -> Optional[str]: + """Extract the last \\boxed{...} content, handling nested braces.""" + if not text: + return None + last_match = None + for m in _BOXED_RE.finditer(text): + start = m.end() + depth = 1 + i = start + while i < len(text) and depth > 0: + if text[i] == '{': + depth += 1 + elif text[i] == '}': + depth -= 1 + i += 1 + if depth == 0: + last_match = text[start:i - 1].strip() + return last_match + + +def normalize_answer(ans: str) -> str: + """Normalize a math answer string for comparison.""" + if not ans: + return '' + s = ans.strip() + # MCQ: extract bare letter from \textbf{(D)}, \text{(A)}, \mathbb{A}, (B), etc. + m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) + if m: + return m.group(1) + s = s.replace(' ', '') + s = s.replace(r'\,', '') + s = s.replace(r'\;', '') + s = s.replace(r'\!', '') + s = s.replace(r'\text', '') + s = s.replace(r'\mathrm', '') + s = s.replace(r'\displaystyle', '') + s = re.sub(r'\\(?:left|right)[.()\[\]|]', '', s) + s = s.replace(r'\dfrac', r'\frac') + s = s.replace(r'\tfrac', r'\frac') + s = s.strip('$').strip() + # Strip unit-like brace suffixes: {cm}, {m}, {kg}, etc. + s = re.sub(r'\{[a-zA-Z]+\}$', '', s) + # Normalize degree: ^\circ, ^{\circ}, ° → deg + s = re.sub(r'\^\\circ|\^\{\\circ\}|°', 'deg', s) + # Canonicalize \frac{a}{b} → (a)/(b) + def _frac_to_slash(m): + # Handle nested braces in numerator/denominator + text = m.group(0) + pos = text.index('{') + 1 + depth, num_start = 1, pos + while depth > 0: + if text[pos] == '{': depth += 1 + elif text[pos] == '}': depth -= 1 + pos += 1 + numer = text[num_start:pos - 1] + pos += 1 # skip '{' + den_start = pos + depth = 1 + while depth > 0: + if text[pos] == '{': depth += 1 + elif text[pos] == '}': depth -= 1 + pos += 1 + denom = text[den_start:pos - 1] + return f'({numer})/({denom})' + s = re.sub(r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', _frac_to_slash, s) + # Also handle bare a/b → (a)/(b) for consistent comparison + # Only simple integer/variable fractions: 17/5 → (17)/(5) + s = re.sub(r'(? bool: + """Try to evaluate both as floats; match if within 1e-9 relative tolerance.""" + try: + va = float(a.replace('(', '').replace(')', '')) + vb = float(b.replace('(', '').replace(')', '')) + return abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) + except (ValueError, ZeroDivisionError): + pass + # Try evaluating simple fraction expressions like (17)/(5) + frac_re = re.compile(r'^\(([^)]+)\)/\(([^)]+)\)$') + def _eval_frac(s): + m = frac_re.match(s) + if m: + try: + return float(m.group(1)) / float(m.group(2)) + except (ValueError, ZeroDivisionError): + pass + return None + va, vb = _eval_frac(a), _eval_frac(b) + if va is not None and vb is not None: + return abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) + return False + + +# MCQ compound pattern: \text{(D) }49, \textbf{(C)}12, (B) 21, etc. +_MCQ_REF_RE = re.compile( + r'^\\(?:textbf|text|mathrm|mathbf)\{\(?([A-E])\)?\s*\}\s*(.+)$' + r'|^\(?([A-E])\)\s+(.+)$' +) + + +def _split_mcq(ans: str): + """Split an MCQ answer into (letter, value) components. + + Handles compound forms (``\\text{(D) }49``, ``(B) 21``) as well as a + bare letter (``D`` -> letter only) and a bare value (``21`` -> value only). + Returns ``(letter_or_None, value_or_None)``. + """ + s = ans.strip() + m = _MCQ_REF_RE.match(s) + if m: + letter = m.group(1) or m.group(3) + value = (m.group(2) or m.group(4) or '').strip() + return letter, (value or None) + # Bare single letter (with optional \text/\textbf/\mathbb wrapper or parens). + # \mathbb{A} appears as a dirty reference label for option A in some rows. + bl = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) + if bl: + return bl.group(1), None + return None, s or None + + +def answers_match(predicted: str, reference: str) -> bool: + """Check if two math answers are equivalent. + + Supports bidirectional MCQ matching: either side may be a bare option + letter, a bare value, or a compound ``(letter) value`` form. The answer is + considered correct if the letters match, or if the values match. + """ + if not predicted or not reference: + return False + norm_p = normalize_answer(predicted) + norm_r = normalize_answer(reference) + if norm_p == norm_r: + return True + if _try_numeric_equal(norm_p, norm_r): + return True + + # Symmetric MCQ matching: decompose both sides into (letter, value). + p_letter, p_value = _split_mcq(predicted) + r_letter, r_value = _split_mcq(reference) + + # Match on the option letter (only meaningful if both sides carry a letter). + if p_letter and r_letter and p_letter == r_letter: + return True + + # If both sides are letter-only (a bare option letter with no value), the + # letters are the only signal; differing letters mean a mismatch. Do NOT + # fall through to value comparison, which would spuriously match dirty + # labels like \mathbb{A} vs \mathbb{B}. + if (p_letter and p_value is None) and (r_letter and r_value is None): + return False + + # Match on the value part (compare whichever value each side exposes; fall + # back to the raw normalized string when a side has no separate value). + p_val = normalize_answer(p_value) if p_value else norm_p + r_val = normalize_answer(r_value) if r_value else norm_r + if p_val and r_val: + if p_val == r_val or _try_numeric_equal(p_val, r_val): + return True + return False + + +# --------------------------------------------------------------------------- +# Dataset loading +# --------------------------------------------------------------------------- + +def load_aops(n: int, seed: int = 42) -> List[Dict[str, Any]]: + """Load AoPS boxed problems, sample n, extract reference answers.""" + from modelscope import MsDataset + ds = MsDataset.load(AOPS_DATASET_ID, split='train', + download_mode='reuse_dataset_if_exists') + boxed = [] + for row in ds: + if not row['metadata'].get('boxed'): + continue + ref = extract_boxed(row['solution']) + if not ref: + continue + boxed.append({ + 'problem': row['problem'], + 'solution': row['solution'], + 'reference_answer': ref, + 'tags': row.get('tags', []), + }) + sys.stderr.write(f'[aops] {len(boxed)} boxed problems with extractable answers\n') + rng = random.Random(seed) + rng.shuffle(boxed) + if n > 0 and n < len(boxed): + boxed = boxed[:n] + sys.stderr.write(f'[aops] sampled {n} problems\n') + return boxed + + +def load_math(n: int, seed: int = 42, split: str = 'test', + per_level: int = 0) -> List[Dict[str, Any]]: + """Load the MATH (Hendrycks) dataset from local extracted JSON files. + + Each problem's reference answer is the ``\\boxed{}`` content of its + ``solution`` (MATH solutions always end in a boxed answer). + + Sampling is *stratified by level* so every difficulty (Level 1-5) is + represented equally — required to measure how RAG gain varies with + difficulty. ``per_level`` (if >0) fixes the count per level; otherwise + ``n`` is split evenly across the 5 levels. When both are 0, all problems + are returned. The final list is shuffled with ``seed`` so index order is + stable/comparable across direct vs rag runs. + """ + import glob + root = os.path.join(MATH_DATA_DIR, split) + files = glob.glob(os.path.join(root, '*', '*.json')) + if not files: + raise FileNotFoundError( + f'[math] no problems found under {root!r}; set MATH_DATA_DIR or ' + f'extract MATH.zip there') + + by_level: Dict[str, List[Dict[str, Any]]] = {} + n_no_box = 0 + for fp in files: + try: + with open(fp, 'r', encoding='utf-8') as fin: + row = json.load(fin) + except Exception: + continue + ref = extract_boxed(row.get('solution', '')) + if not ref: + n_no_box += 1 + continue + level = row.get('level', 'Unknown') + by_level.setdefault(level, []).append({ + 'problem': row['problem'], + 'solution': row['solution'], + 'reference_answer': ref, + 'level': level, + 'type': row.get('type', ''), + }) + + total = sum(len(v) for v in by_level.values()) + sys.stderr.write( + f'[math] {total} problems with boxed answers across ' + f'{len(by_level)} levels (skipped {n_no_box} without boxed)\n') + + levels = sorted(by_level.keys()) + rng = random.Random(seed) + + # Decide how many per level. + if per_level <= 0 and n > 0: + per_level = max(1, n // max(1, len(levels))) + + sampled: List[Dict[str, Any]] = [] + for lv in levels: + pool = by_level[lv] + rng.shuffle(pool) + take = pool if per_level <= 0 else pool[:per_level] + sampled.extend(take) + sys.stderr.write(f'[math] {lv}: took {len(take)}/{len(pool)}\n') + + rng.shuffle(sampled) + sys.stderr.write(f'[math] total sampled: {len(sampled)}\n') + return sampled + + +# --------------------------------------------------------------------------- +# Prompt building +# --------------------------------------------------------------------------- + +DIRECT_SYSTEM = ( + 'You are an expert competition mathematician. Solve the following problem ' + 'step by step. Provide your final answer inside \\boxed{}.' +) + +RAG_SYSTEM = ( + 'You are an expert competition mathematician. Solve the following problem ' + 'step by step. Provide your final answer inside \\boxed{}.\n\n' + 'You will first see example problem-solving traces or skills. ' + 'Learn from the reasoning methodology demonstrated in these examples, ' + 'then thinking to solve the actual problem.' +) + +RAG_FOLLOWUP = ( + 'The above is a reference solution to a similar problem. ' + 'You may use any applicable techniques from it, or ignore it ' + 'if you find a better approach. ' + 'Solve the problem step by step and put your final answer in \\boxed{}.' +) + +HINT_FOLLOWUP = ( + 'The above are applicable solution approaches extracted from similar problems. ' + 'You may use any applicable techniques from them, or ignore them ' + 'if you find a better approach. ' + 'Solve the problem step by step and put your final answer in \\boxed{}.' +) + +# Reminder appended to the final user turn. Without this, the reasoning model +# can loop indefinitely on multiple-choice problems, oscillating between boxing +# the option letter and boxing the value (e.g. "I'll box B. I'll box 21. ...") +# and never terminating. Boxing BOTH the letter and value removes the ambiguity +# (the grader accepts either), so the model has no format decision to agonize over. +MCQ_INSTRUCTION = ( + '\n\nNote: If the problem is multiple-choice (it lists options such as ' + '(A), (B), (C), ...), put BOTH the option letter and its value in the box, ' + 'e.g. \\boxed{(B) 21}. Otherwise, box the value directly. Decide the answer ' + 'format once and do not deliberate over which form to box.' +) + + +def build_direct_prompt(problem: str) -> Dict[str, Any]: + return { + 'messages': [ + {'role': 'system', 'content': DIRECT_SYSTEM}, + {'role': 'user', 'content': problem + MCQ_INSTRUCTION}, + ] + } + + +def build_hint_prompt(problem: str, hint_analysis: str) -> Dict[str, Any]: + """Build prompt with pre-analyzed hint in a multi-turn conversation. + + Mirrors ``build_rag_prompt``: the hint is presented as an assistant + "extracted approaches" turn (instead of being buried in the system + prompt), followed by a user instruction that provides a clear closing + directive to solve the problem and box the answer. Keeping the final + solve/box instruction in a dedicated user turn (rather than in the + system prompt) helps the reasoning model terminate cleanly. + """ + messages: List[Dict[str, str]] = [ + {'role': 'system', 'content': DIRECT_SYSTEM}, + {'role': 'user', 'content': problem}, + {'role': 'assistant', + 'content': ('Here are applicable solution approaches extracted from ' + f'similar problems:\n\n{hint_analysis}')}, + {'role': 'user', 'content': HINT_FOLLOWUP + MCQ_INSTRUCTION}, + ] + return {'messages': messages} + + +def build_rag_prompt(problem: str, + examples: List[Dict[str, str]]) -> Dict[str, Any]: + """Approach B: multi-turn assistant format. + + The trace is presented as an assistant "retrieval" turn, followed by + a user instruction that constrains the model to use methodology only. + """ + messages: List[Dict[str, str]] = [{'role': 'system', 'content': DIRECT_SYSTEM}] + messages.append({'role': 'user', 'content': problem}) + # Build trace content from retrieved examples + trace_parts = [] + for i, ex in enumerate(examples, 1): + trace_parts.append(f'[Retrieved Example {i}]\nProblem: {ex["query"]}\n' + f'Reasoning:\n{ex["thinking"]}') + trace_text = '\n\n'.join(trace_parts) + messages.append({'role': 'assistant', + 'content': f'I found relevant reasoning traces from the knowledge base!\n\n{trace_text}'}) + messages.append({'role': 'user', 'content': RAG_FOLLOWUP + MCQ_INSTRUCTION}) + return {'messages': messages} + + +# --------------------------------------------------------------------------- +# 13-gram Jaccard decontamination +# --------------------------------------------------------------------------- + +def _normalize_for_ngram(text: str) -> str: + """Normalize text for n-gram comparison: strip LaTeX markup, lowercase.""" + text = text.lower() + text = re.sub(r'\$+', '', text) + text = re.sub(r'\\[a-z]+\{([^}]*)\}', r'\1', text) + text = re.sub(r'\\[a-z]+', ' ', text) + text = re.sub(r'[{}\\^_$]', '', text) + text = re.sub(r'\s+', ' ', text).strip() + return text + + +def _ngram_jaccard(text_a: str, text_b: str, n: int = 13) -> float: + """13-gram character-level Jaccard similarity.""" + a = _normalize_for_ngram(text_a) + b = _normalize_for_ngram(text_b) + if len(a) < n or len(b) < n: + return 0.0 + grams_a = set(a[i:i + n] for i in range(len(a) - n + 1)) + grams_b = set(b[i:i + n] for i in range(len(b) - n + 1)) + if not grams_a or not grams_b: + return 0.0 + return len(grams_a & grams_b) / len(grams_a | grams_b) + + +# --------------------------------------------------------------------------- +# Embedding / RAG helpers +# --------------------------------------------------------------------------- + +def _wrap_anchor(text: str) -> List[Dict[str, str]]: + return [ + {'role': 'user', 'content': text}, + {'role': 'assistant', 'content': 'Match the correct response here.'}, + ] + + +def get_embeddings(model: TransformersModel, template: Qwen3_5Template, + texts: List[str], dp_size: int) -> np.ndarray: + if not texts: + return np.zeros((0,), dtype=np.float32) + n = len(texts) + pad_n = (-n) % dp_size + padded = list(texts) + [' '] * pad_n if pad_n else list(texts) + features = [] + for t in padded: + feat = template.encode({'messages': _wrap_anchor(t or ' ')}) + feat['labels'] = [1] + features.append(feat) + out = model.forward_only(inputs=features, task='embedding', return_logits=True) + emb = out['embeddings'] + if isinstance(emb, torch.Tensor): + emb = emb.detach().to(torch.float32).cpu().numpy() + emb = np.asarray(emb, dtype=np.float32) + return emb[:n] if pad_n else emb + + +def retrieve_examples(tbl, query_vecs: np.ndarray, top_k: int, + use_thinking_raw: bool, sim_threshold: float = 0.0, + problems: List[str] = None, + decontam_threshold: float = 0.0, + ) -> List[List[Dict[str, str]]]: + thinking_field = 'thinking_raw' if use_thinking_raw else 'cot_compressed' + fetch_limit = top_k + 50 if decontam_threshold > 0 else top_k + n_queries = len(query_vecs) + all_examples: List[List[Dict[str, str]]] = [None] * n_queries + decontam_skipped = 0 + _decontam_lock = threading.Lock() + + def _search_one(qi: int): + nonlocal decontam_skipped + vec = query_vecs[qi] + results = ( + tbl.search(vec.astype(np.float32).tolist()) + .metric('dot') + .limit(fetch_limit) + .select(['query_raw', thinking_field, '_distance']) + .to_list() + ) + problem_text = problems[qi] if problems else '' + examples = [] + local_skipped = 0 + for r in results: + if len(examples) >= top_k: + break + sim = 1.0 - r.get('_distance', 0.0) + if sim < sim_threshold: + continue + q = r.get('query_raw', '') + t = r.get(thinking_field, '') + if not t: + continue + if decontam_threshold > 0 and problem_text and q: + ng_sim = _ngram_jaccard(problem_text, q) + if ng_sim > decontam_threshold: + local_skipped += 1 + continue + examples.append({'query': q, 'thinking': t, '_sim': round(sim, 4), + '_raw_trace_len': len(t)}) + all_examples[qi] = examples + if local_skipped: + with _decontam_lock: + decontam_skipped += local_skipped + + with ThreadPoolExecutor(max_workers=min(n_queries, 16)) as pool: + list(pool.map(_search_one, range(n_queries))) + + if decontam_skipped > 0: + logger.info(f'[decontam] skipped {decontam_skipped} leaked retrievals ' + f'(13-gram Jaccard > {decontam_threshold})') + return all_examples + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main(): + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--mode', choices=['direct', 'rag'], default='rag') + p.add_argument('--dataset', choices=['aops', 'math'], default='math', + help='Evaluation dataset. "math" = MATH (Hendrycks), ' + 'stratified by level for a difficulty-vs-gain curve.') + p.add_argument('--math-split', default='test', + help='MATH split to load (test/train).') + p.add_argument('--per-level', type=int, default=100, + help='MATH only: problems per difficulty level (default 100 ' + '-> 500 total across Level 1-5). If 0, --n is split ' + 'evenly across the 5 levels.') + p.add_argument('--n', type=int, default=0, + help='Pool size: sample this many problems (0 = all boxed). ' + 'In RAG mode with --target-eval, set this to 0 for max coverage.') + p.add_argument('--target-eval', type=int, default=0, + help='Stop after this many problems are successfully evaluated ' + '(0 = no limit, evaluate the entire sampled set — the ' + 'default, so all 500 stratified MATH problems are run). ' + 'RAG mode: counts problems with valid traces after ' + 'decontam; direct mode: ignored, evaluates all filtered.') + p.add_argument('--db-path', default='./output.oldemb/thinking_rag/lance.db') + p.add_argument('--table', default='thinking_traces') + p.add_argument('--top-k', type=int, default=1) + p.add_argument('--use-cot-compressed', action='store_true', + help='Use pre-compressed cot_compressed field instead of thinking_raw.') + p.add_argument('--sim-threshold', type=float, default=0.75, + help='Minimum cosine similarity for retrieved traces. ' + 'Traces below this are discarded at retrieval time.') + p.add_argument('--decontam-threshold', type=float, default=0.20, + help='13-gram Jaccard threshold for leak detection. ' + 'Retrieved traces above this are skipped (0=disabled).') + p.add_argument('--llm-decontam', action='store_true', default=True, + help='LLM-based decontamination (default ON): API judges whether ' + 'retrieved problem is the same as the test problem. ' + 'Applied after 13-gram decontam, before condensing. ' + 'Use --no-llm-decontam to disable.') + p.add_argument('--no-llm-decontam', dest='llm_decontam', action='store_false', + help='Disable LLM-based decontamination.') + p.add_argument('--max-trace-len', type=int, default=12000) + p.add_argument('--condense', action='store_true', default=True, + help='Enable condenser re-compression on retrieved traces ' + '(default ON). Use --no-condense to inject raw traces.') + p.add_argument('--no-condense', dest='condense', action='store_false', + help='Disable condenser; inject raw retrieved traces.') + p.add_argument('--condense-max-len', type=int, default=2000, + help='Max chars of condensed trace (fallback truncation).') + p.add_argument('--batch-size', type=int, default=16) + p.add_argument('--seed', type=int, default=42) + p.add_argument('--hint', action='store_true', default=False, + help='Enable API hint filtering on retrieved traces (default OFF; ' + 'raw RAG injects the condensed trace directly). ' + 'In rag mode: retrieve → condense → API filters trace → refined system prompt. ' + 'In direct mode: ignored (no traces to filter).') + p.add_argument('--no-hint', dest='hint', action='store_false', + help='Disable API hint filtering; inject condensed trace directly.') + p.add_argument('--problem-ids-file', default=None, + help='File listing problem indices evaluated by RAG mode. ' + 'RAG mode writes this file; direct mode reads it to ' + 'evaluate the same subset (use --no-filter to disable). ' + 'Defaults to a dataset-specific path.') + p.add_argument('--no-filter', action='store_true', + help='In direct mode, evaluate ALL sampled problems ' + 'instead of filtering to RAG subset.') + p.add_argument('--output', default=None) + args = p.parse_args() + + # Dataset-specific default paths (keeps aops and math runs from colliding). + if args.problem_ids_file is None: + args.problem_ids_file = ( + f'./output/thinking_rag/{args.dataset}_rag_problem_ids.json') + + if args.output is None: + suffix = f'{args.mode}_hint' if (args.hint and args.mode == 'rag') else args.mode + args.output = ( + f'./output/thinking_rag/{args.dataset}_{suffix}_results.jsonl') + + if args.condense and args.use_cot_compressed: + logger.warning('--condense requires thinking_raw, ignoring --use-cot-compressed') + args.use_cot_compressed = False + + if args.dataset == 'math': + records = load_math(n=args.n, seed=args.seed, split=args.math_split, + per_level=args.per_level) + else: + records = load_aops(n=args.n, seed=args.seed) + + is_rag = (args.mode == 'rag') + + # Direct mode: filter to same problems RAG evaluated (controlled comparison) + original_indices = list(range(len(records))) # track original indices + if not is_rag and not args.no_filter: + if os.path.exists(args.problem_ids_file): + with open(args.problem_ids_file) as f: + content = f.read().strip() + if content.startswith('['): + valid_indices = set(json.loads(content)) + else: + valid_indices = set(int(line) for line in content.splitlines() if line.strip()) + filtered = [(i, r) for i, r in enumerate(records) if i in valid_indices] + original_indices = [i for i, _ in filtered] + records = [r for _, r in filtered] + sys.stderr.write( + f'[direct] filtered to {len(records)} problems ' + f'from {args.problem_ids_file}\n') + else: + sys.stderr.write( + f'[direct] WARNING: {args.problem_ids_file} not found, ' + f'running all {len(records)} problems\n') + + condenser_gpus = int(os.environ.get('EVAL_CONDENSER_GPUS', 0)) if args.condense else 0 + + # Raw RAG relies on the API condenser (qwen3.7-max). Fail fast with a clear + # message if it's enabled without an API key and without a local condenser. + if is_rag and args.condense and not CONDENSE_API_KEY and condenser_gpus == 0: + sys.stderr.write( + '[condense] ERROR: --condense is ON but COMPRESS_API_KEY is unset ' + 'and no local condenser (EVAL_CONDENSER_GPUS=0).\n' + ' Fix one of:\n' + ' - export COMPRESS_API_KEY=sk-... (use API condenser)\n' + ' - EVAL_CONDENSER_GPUS=2 python ... (use local vLLM condenser)\n' + ' - pass --no-condense (inject raw traces)\n') + sys.exit(1) + + if is_rag: + num_gpus = EMB_GPUS + GEN_GPUS + condenser_gpus + device_groups = [ + DeviceGroup(name='emb_model', ranks=list(range(EMB_GPUS)), + device_type='GPU'), + DeviceGroup(name='sampler', + ranks=list(range(EMB_GPUS, EMB_GPUS + GEN_GPUS)), + device_type='GPU', gpus_per_worker=GEN_GPUS), + ] + if condenser_gpus > 0: + cond_start = EMB_GPUS + GEN_GPUS + device_groups.append( + DeviceGroup(name='condenser', + ranks=list(range(cond_start, cond_start + condenser_gpus)), + device_type='GPU')) + emb_mesh = DeviceMesh.from_sizes(world_size=EMB_GPUS, dp_size=EMB_GPUS) + gen_mesh = DeviceMesh.from_sizes(world_size=GEN_GPUS, tp_size=GEN_GPUS) + twinkle.initialize(mode='ray', nproc_per_node=num_gpus, + groups=device_groups, lazy_collect=False) + else: + device_groups = [ + DeviceGroup(name='sampler', ranks=list(range(GEN_GPUS)), + device_type='GPU', gpus_per_worker=GEN_GPUS), + ] + gen_mesh = DeviceMesh.from_sizes(world_size=GEN_GPUS, tp_size=GEN_GPUS) + twinkle.initialize(mode='ray', nproc_per_node=GEN_GPUS, + groups=device_groups, lazy_collect=False) + + sampler = vLLMSampler( + model_id=GEN_MODEL_ID, + engine_args={ + 'gpu_memory_utilization': GEN_GPU_MEM, + 'max_model_len': GEN_MAX_MODEL_LEN, + }, + device_mesh=gen_mesh, + remote_group='sampler', + ) + sampler.set_template('Qwen3_5Template', model_id=GEN_MODEL_ID, + enable_thinking=True, max_length=GEN_MAX_MODEL_LEN) + sys.stderr.write(f'[aops] vLLM sampler ready (model={GEN_MODEL_ID})\n') + + gen_params = TwinkleSamplingParams( + max_tokens=GEN_MAX_TOKENS, + temperature=GEN_TEMPERATURE, + top_p=GEN_TOP_P, + num_samples=1, + ) + + emb_model = emb_template = tbl = None + if is_rag: + import lancedb + db = lancedb.connect(args.db_path) + if args.table not in db.table_names(): + raise SystemExit(f'Table "{args.table}" not found in {args.db_path}') + tbl = db.open_table(args.table) + sys.stderr.write(f'[aops] LanceDB rows={tbl.count_rows()}\n') + + emb_model = TransformersModel( + model_id=EMBED_MODEL_ID, device_mesh=emb_mesh, + remote_group='emb_model') + emb_model.set_processor(InputProcessor) + emb_model.set_loss(InfonceLoss, temperature=0.03, use_batch=True) + emb_template = Qwen3_5Template( + model_id=EMBED_MODEL_ID, max_length=EMBED_MAX_LENGTH, + truncation_strategy='delete', enable_thinking=False) + sys.stderr.write('[aops] embedding model ready\n') + + # -- Condenser setup (API primary + optional local vLLM) ------------------- + condenser_api_client = None + condenser_sampler_obj = None + condenser_params = None + condenser_special_tokens = None + + if args.condense and is_rag: + condenser_api_client = OpenAIClient( + model=CONDENSE_API_MODEL, api_key=CONDENSE_API_KEY, + base_url=CONDENSE_BASE_URL) + sys.stderr.write(f'[condense] API client ready (model={CONDENSE_API_MODEL})\n') + + if condenser_gpus > 0: + condenser_mesh = DeviceMesh.from_sizes( + world_size=condenser_gpus, dp_size=condenser_gpus) + condenser_sampler_obj = vLLMSampler( + model_id=CONDENSE_MODEL_ID, + engine_args={'gpu_memory_utilization': 0.8, 'max_model_len': 32768}, + device_mesh=condenser_mesh, + remote_group='condenser', + ) + condenser_sampler_obj.set_template( + 'Qwen3_5Template', model_id=CONDENSE_MODEL_ID, + enable_thinking=False, truncation_strategy='delete', + max_length=32768) + condenser_template = Qwen3_5Template( + model_id=CONDENSE_MODEL_ID, max_length=32768, + enable_thinking=False, truncation_strategy='delete') + condenser_special_tokens = set(condenser_template.tokenizer.all_special_tokens) + condenser_params = TwinkleSamplingParams( + max_tokens=CONDENSE_MAX_TOKENS, + temperature=CONDENSE_TEMPERATURE, + top_p=0.5, num_samples=1) + sys.stderr.write(f'[condense] local vLLM ready (model={CONDENSE_MODEL_ID})\n') + + # -- Hint analysis API client (reuses condenser API config) ----------------- + hint_api_client = None + if args.hint and is_rag: + if condenser_api_client is not None: + hint_api_client = condenser_api_client + else: + hint_api_client = OpenAIClient( + model=CONDENSE_API_MODEL, api_key=CONDENSE_API_KEY, + base_url=CONDENSE_BASE_URL) + sys.stderr.write(f'[hint] API hint analysis enabled (model={CONDENSE_API_MODEL})\n') + + # -- LLM decontam API client --------------------------------------------------- + decontam_api_client = None + if args.llm_decontam and is_rag: + if hint_api_client is not None: + decontam_api_client = hint_api_client + elif condenser_api_client is not None: + decontam_api_client = condenser_api_client + else: + decontam_api_client = OpenAIClient( + model=CONDENSE_API_MODEL, api_key=CONDENSE_API_KEY, + base_url=CONDENSE_BASE_URL) + sys.stderr.write(f'[decontam-llm] LLM decontamination enabled (model={CONDENSE_API_MODEL})\n') + + correct_count = 0 + total_count = 0 + skipped_indices: List[int] = [] # problems skipped by RAG (no valid trace) + evaluated_indices: List[int] = [] # problems actually evaluated + debug_records: List[Dict[str, Any]] = [] + + os.makedirs(os.path.dirname(args.output) or '.', exist_ok=True) + out_f = open(args.output, 'w', encoding='utf-8') + + # Open problem-ids files for incremental writing (RAG mode only) + ids_f = None + skip_f = None + if is_rag: + os.makedirs(os.path.dirname(args.problem_ids_file) or '.', exist_ok=True) + ids_f = open(args.problem_ids_file, 'w', encoding='utf-8') + skip_path = args.problem_ids_file.replace('.json', '_skipped.json') + skip_f = open(skip_path, 'w', encoding='utf-8') + + # -- RAG batch preparation (embed + retrieve + decontam + condense + hint) -- + def _prepare_rag_batch(batch_start: int): + """Prepare a RAG batch: returns (prompts, batch, all_examples, + hint_analyses, kept_global_indices, batch_skipped_indices) or None.""" + batch_end = min(batch_start + args.batch_size, len(records)) + batch = records[batch_start:batch_end] + problems = [r['problem'] for r in batch] + + query_vecs = get_embeddings(emb_model, emb_template, problems, EMB_GPUS) + use_raw = not args.use_cot_compressed + all_examples = retrieve_examples(tbl, query_vecs, args.top_k, + use_raw, args.sim_threshold, + problems=problems, + decontam_threshold=args.decontam_threshold) + if args.use_cot_compressed: + for exs in all_examples: + for ex in exs: + ex['thinking'] = _strip_condenser_markers(ex['thinking']) + + if args.llm_decontam and decontam_api_client: + all_examples = _llm_decontaminate( + decontam_api_client, problems, all_examples) + + if args.condense and condenser_api_client: + all_examples = condense_traces( + all_examples, problems, condenser_api_client, + condenser_sampler=condenser_sampler_obj, + compress_params=condenser_params, + special_tokens=condenser_special_tokens, + max_output_len=args.condense_max_len, + dp_size=condenser_gpus) + + hint_analyses = None + if args.hint and hint_api_client: + hint_analyses = _api_hint_analysis_batch( + hint_api_client, problems, all_examples) + + keep_mask = [] + for pi, (r, examples) in enumerate(zip(batch, all_examples)): + if not examples: + keep_mask.append(False) + elif hint_analyses and hint_analyses[pi]: + keep_mask.append(True) + else: + usable = [ex for ex in examples + if len(ex['thinking']) <= args.max_trace_len] + keep_mask.append(bool(usable)) + + batch_skipped = [] + for pi, keep in enumerate(keep_mask): + if not keep: + batch_skipped.append(batch_start + pi) + + kept_batch = [] + kept_examples = [] + kept_hints = [] + kept_global_indices = [] + for pi, keep in enumerate(keep_mask): + if keep: + kept_batch.append(batch[pi]) + kept_examples.append(all_examples[pi]) + kept_hints.append(hint_analyses[pi] if hint_analyses else None) + kept_global_indices.append(batch_start + pi) + + if not kept_batch: + return None, None, None, None, None, batch_skipped + + prompts = [] + for pi, (r, examples) in enumerate(zip(kept_batch, kept_examples)): + if kept_hints[pi]: + prompts.append(build_hint_prompt(r['problem'], kept_hints[pi])) + else: + filtered = [{'query': ex['query'], 'thinking': ex['thinking']} + for ex in examples + if len(ex['thinking']) <= args.max_trace_len] + prompts.append(build_rag_prompt(r['problem'], filtered)) + + return prompts, kept_batch, kept_examples, kept_hints, kept_global_indices, batch_skipped + + target_reached = False + batch_starts = list(range(0, len(records), args.batch_size)) + + if is_rag: + # Pipeline: prefetch next batch while current batch generates + from concurrent.futures import Future + prefetch_pool = ThreadPoolExecutor(max_workers=1) + # Prepare first batch synchronously + cur_result = _prepare_rag_batch(batch_starts[0]) + + for bi, batch_start in enumerate(batch_starts): + if target_reached: + break + prompts, batch, all_examples, hint_analyses, kept_global_indices, batch_skipped = cur_result + skipped_indices.extend(batch_skipped or []) + if skip_f and batch_skipped: + for sid in batch_skipped: + skip_f.write(f'{sid}\n') + skip_f.flush() + + # Submit next batch preparation in background + next_future: Optional[Future] = None + if bi + 1 < len(batch_starts) and not target_reached: + next_future = prefetch_pool.submit(_prepare_rag_batch, batch_starts[bi + 1]) + + if prompts is None: + # Entire batch skipped + cur_result = next_future.result() if next_future else None + continue + + # Generate (runs on gen GPU while next batch prepares on emb GPU + API) + responses = sampler.sample(prompts, gen_params) + + for i, (rec, resp) in enumerate(zip(batch, responses)): + seq = resp.sequences[0] if resp and resp.sequences else None + raw_output = '' + if seq is not None: + raw_output = seq.decoded or '' + raw_output = re.sub(r'<\|[^|]+\|>', '', raw_output).rstrip() + + predicted = extract_boxed(raw_output) + is_correct = answers_match(predicted, rec['reference_answer']) + if is_correct: + correct_count += 1 + total_count += 1 + + global_idx = kept_global_indices[i] + evaluated_indices.append(global_idx) + if ids_f: + ids_f.write(f'{global_idx}\n') + ids_f.flush() + + debug_rec = { + 'idx': global_idx, + 'reference_answer': rec['reference_answer'], + 'predicted': predicted, + 'is_correct': is_correct, + 'problem': rec['problem'], + 'model_output': raw_output, + } + if rec.get('level'): + debug_rec['level'] = rec['level'] + if rec.get('type'): + debug_rec['type'] = rec['type'] + debug_rec['num_traces'] = len(all_examples[i]) + if all_examples[i]: + ex0 = all_examples[i][0] + debug_rec['similarity'] = ex0.get('_sim', 0.0) + debug_rec['retrieved_query'] = ex0.get('query', '') + debug_rec['raw_trace_len'] = ex0.get('_raw_trace_len', 0) + debug_rec['condensed_trace'] = ex0['thinking'] + debug_rec['condensed_trace_len'] = len(ex0['thinking']) + debug_rec['condense_source'] = ex0.get('_condense_source', '') + if hint_analyses and hint_analyses[i]: + debug_rec['hint_analysis'] = hint_analyses[i] + debug_records.append(debug_rec) + out_f.write(json.dumps(debug_rec, ensure_ascii=False) + '\n') + out_f.flush() + + acc = correct_count / total_count if total_count else 0 + sys.stderr.write( + f' [{total_count}/{args.target_eval}] ' + f'acc={acc:.4f} ({correct_count}/{total_count})\n') + + if args.target_eval > 0 and total_count >= args.target_eval: + target_reached = True + + # Collect prefetched result for next iteration (skip if done) + if not target_reached and next_future: + cur_result = next_future.result() + else: + cur_result = None + + prefetch_pool.shutdown(wait=True) + else: + # Direct mode: no pipeline needed, just batch generate + for batch_start in batch_starts: + batch_end = min(batch_start + args.batch_size, len(records)) + batch = records[batch_start:batch_end] + prompts = [build_direct_prompt(r['problem']) for r in batch] + + responses = sampler.sample(prompts, gen_params) + + for i, (rec, resp) in enumerate(zip(batch, responses)): + seq = resp.sequences[0] if resp and resp.sequences else None + raw_output = '' + if seq is not None: + raw_output = seq.decoded or '' + raw_output = re.sub(r'<\|[^|]+\|>', '', raw_output).rstrip() + + predicted = extract_boxed(raw_output) + is_correct = answers_match(predicted, rec['reference_answer']) + if is_correct: + correct_count += 1 + total_count += 1 + + global_idx = original_indices[batch_start + i] + evaluated_indices.append(global_idx) + + debug_rec = { + 'idx': global_idx, + 'reference_answer': rec['reference_answer'], + 'predicted': predicted, + 'is_correct': is_correct, + 'problem': rec['problem'], + 'model_output': raw_output, + } + if rec.get('level'): + debug_rec['level'] = rec['level'] + if rec.get('type'): + debug_rec['type'] = rec['type'] + debug_records.append(debug_rec) + out_f.write(json.dumps(debug_rec, ensure_ascii=False) + '\n') + out_f.flush() + + acc = correct_count / total_count if total_count else 0 + sys.stderr.write( + f' [{total_count}/{len(records)}] ' + f'acc={acc:.4f} ({correct_count}/{total_count})\n') + + overall_acc = correct_count / total_count if total_count else 0 + print(f'\n{"=" * 60}') + print(f'{args.dataset.upper()} — mode={args.mode}, model={GEN_MODEL_ID}') + print(f' n={total_count}, seed={args.seed}') + if is_rag: + print(f' evaluated={len(evaluated_indices)}, skipped={len(skipped_indices)}') + print(f'{"=" * 60}') + print(f'Overall accuracy: {overall_acc:.4f} ({correct_count}/{total_count})') + + # Per-level breakdown (MATH: the difficulty-vs-gain curve we care about). + if any(r.get('level') for r in debug_records): + from collections import defaultdict + per = defaultdict(lambda: [0, 0]) # level -> [correct, total] + for r in debug_records: + lv = r.get('level', 'Unknown') + per[lv][1] += 1 + if r['is_correct']: + per[lv][0] += 1 + print(f'\nPer-level accuracy:') + for lv in sorted(per.keys()): + c, t = per[lv] + print(f' {lv:>10}: {c/t:.4f} ({c}/{t})') + + out_f.close() + print(f'\n[output] {len(debug_records)} records saved to {args.output}') + + if ids_f: + ids_f.close() + print(f'[output] problem IDs ({len(evaluated_indices)}) saved to {args.problem_ids_file}') + if skip_f: + skip_f.close() + if skipped_indices: + print(f'[output] skipped IDs ({len(skipped_indices)}) saved to ' + f'{args.problem_ids_file.replace(".json", "_skipped.json")}') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/embedding/eval_math_by_level.sh b/cookbook/exp/embedding/eval_math_by_level.sh new file mode 100755 index 000000000..d4f8bbdb3 --- /dev/null +++ b/cookbook/exp/embedding/eval_math_by_level.sh @@ -0,0 +1,59 @@ +#!/bin/bash +# MATH (Hendrycks) difficulty-stratified evaluation. +# +# Goal: measure how the (raw) RAG gain over direct varies with problem +# difficulty (Level 1-5). Runs raw RAG first (retrieve -> qwen3.7-max condense +# -> inject, no hint filtering; it writes the problem-id file), then direct on +# the *same* problems for a paired comparison. +# +# Usage: +# COMPRESS_API_KEY=sk-xxx bash cookbook/exp/embedding/eval_math_by_level.sh +# +# Env knobs: +# PER_LEVEL problems per difficulty level (default 100 -> 500 total) +# SEED stratified-sampling seed (default 100; must match across runs) +# DB_PATH LanceDB retrieval index +# SIM / TOPK retrieval threshold / top-k + +set -euo pipefail + +export COMPRESS_API_KEY="${COMPRESS_API_KEY:?Set COMPRESS_API_KEY}" + +SCRIPT="cookbook/exp/embedding/eval_gpqa_rag.py" +PER_LEVEL="${PER_LEVEL:-100}" +SEED="${SEED:-100}" +SIM="${SIM:-0.75}" +TOPK="${TOPK:-1}" +OUTDIR="./output/thinking_rag" +DB_PATH="${DB_PATH:-./output.oldemb/thinking_rag/lance.db}" + +mkdir -p "$OUTDIR" + +echo "============================================================" +echo " MATH by level: raw RAG (qwen3.7-max condenser, no hint)" +echo " per_level=$PER_LEVEL seed=$SEED" +echo "============================================================" +python "$SCRIPT" \ + --dataset math --math-split test \ + --mode rag \ + --per-level "$PER_LEVEL" --seed "$SEED" \ + --db-path "$DB_PATH" \ + --sim-threshold "$SIM" --top-k "$TOPK" \ + --condense \ + --output "$OUTDIR/math_rag_results.jsonl" + +echo "" +echo "============================================================" +echo " MATH by level: Direct (same problems as raw RAG)" +echo "============================================================" +# Direct reads math_rag_problem_ids.json (written above) to match the subset. +python "$SCRIPT" \ + --dataset math --math-split test \ + --mode direct \ + --per-level "$PER_LEVEL" --seed "$SEED" \ + --output "$OUTDIR/math_direct_results.jsonl" + +echo "" +echo "============================================================" +echo " Done. Compare with: python cookbook/exp/embedding/compare_math_levels.py" +echo "============================================================" diff --git a/cookbook/exp/embedding/eval_rag_recall.py b/cookbook/exp/embedding/eval_rag_recall.py new file mode 100644 index 000000000..19bc5ad6d --- /dev/null +++ b/cookbook/exp/embedding/eval_rag_recall.py @@ -0,0 +1,187 @@ +"""Self-recall evaluation: sample rows from LanceDB, re-encode query, check retrieval. + +Unlike the full build pipeline (which needs 8 GPUs for condenser + embedding), +this script only needs the embedding model (4 GPUs) since it uses the +already-compressed ``query_compressed`` stored in the index. + +Launch: + python cookbook/exp/embedding/eval_rag_recall.py + python cookbook/exp/embedding/eval_rag_recall.py --n 200 --top-k 20 + python cookbook/exp/embedding/eval_rag_recall.py --db-path ./output/thinking_rag/lance.db +""" +import argparse +import json +import os +import random +import sys +from typing import Any, Dict, List, Tuple + +import numpy as np +import torch + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.loss import InfonceLoss +from twinkle.model import TransformersModel +from twinkle.processor import InputProcessor +from twinkle.template import Qwen3_5Template + +logger = get_logger() + +EMBED_MODEL_ID = os.environ.get( + 'EMBED_MODEL_ID', 'output/embedding_full_transformers/last-checkpoint') +EMB_GPUS = int(os.environ.get('EMB_GPUS', 4)) +EMBED_MAX_LENGTH = int(os.environ.get('EMBED_MAX_LENGTH', 8192)) + + +def _wrap_anchor(text: str) -> List[Dict[str, str]]: + return [ + {'role': 'user', 'content': text}, + {'role': 'assistant', 'content': 'Match the correct response here.'}, + ] + + +def get_embeddings(model: TransformersModel, template: Qwen3_5Template, + texts: List[str]) -> np.ndarray: + if not texts: + return np.zeros((0,), dtype=np.float32) + n = len(texts) + pad_n = (-n) % EMB_GPUS + padded = list(texts) + [' '] * pad_n if pad_n else list(texts) + features = [] + for t in padded: + feat = template.encode({'messages': _wrap_anchor(t or ' ')}) + feat['labels'] = [1] + features.append(feat) + out = model.forward_only(inputs=features, task='embedding', return_logits=True) + emb = out['embeddings'] + if isinstance(emb, torch.Tensor): + emb = emb.detach().to(torch.float32).cpu().numpy() + emb = np.asarray(emb, dtype=np.float32) + return emb[:n] if pad_n else emb + + +def main(): + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--db-path', default='./output/thinking_rag/lance.db') + p.add_argument('--table', default='thinking_traces') + p.add_argument('--n', type=int, default=100, help='Number of samples to probe.') + p.add_argument('--top-k', type=int, default=10) + p.add_argument('--seed', type=int, default=42) + p.add_argument('--batch-size', type=int, default=32) + p.add_argument('--output', default='./output/thinking_rag/recall_debug.jsonl', + help='JSONL file to dump per-sample debug info.') + args = p.parse_args() + + import lancedb + db = lancedb.connect(args.db_path) + if args.table not in db.table_names(): + raise SystemExit(f'Table "{args.table}" not found in {args.db_path}') + tbl = db.open_table(args.table) + total_rows = tbl.count_rows() + sys.stderr.write(f'[eval] table={args.table} rows={total_rows}\n') + + df = tbl.to_pandas() + n_sample = min(args.n, len(df)) + random.seed(args.seed) + sample_indices = random.sample(range(len(df)), n_sample) + samples = df.iloc[sample_indices].reset_index(drop=True) + sys.stderr.write(f'[eval] sampled {n_sample} rows for self-recall test\n') + + # Init embedding model only (no condenser needed). + device_groups = [ + DeviceGroup(name='emb_model', ranks=list(range(EMB_GPUS)), device_type='GPU'), + ] + emb_mesh = DeviceMesh.from_sizes(world_size=EMB_GPUS, dp_size=EMB_GPUS) + twinkle.initialize(mode='ray', nproc_per_node=EMB_GPUS, groups=device_groups, + lazy_collect=False) + + model = TransformersModel(model_id=EMBED_MODEL_ID, device_mesh=emb_mesh, + remote_group='emb_model') + model.set_processor(InputProcessor) + model.set_loss(InfonceLoss, temperature=0.03, use_batch=True) + template = Qwen3_5Template(model_id=EMBED_MODEL_ID, max_length=EMBED_MAX_LENGTH, + truncation_strategy='delete', enable_thinking=False) + sys.stderr.write('[eval] embedding model ready\n') + + ks = sorted({1, 5, 10, args.top_k}) + hits = {k: 0 for k in ks} + per_source_hits: Dict[str, Dict[int, int]] = {} + per_source_total: Dict[str, int] = {} + debug_records: List[Dict[str, Any]] = [] + + # Batch encode and search. + for batch_start in range(0, n_sample, args.batch_size): + batch_end = min(batch_start + args.batch_size, n_sample) + batch = samples.iloc[batch_start:batch_end] + queries = batch['query_compressed'].tolist() + ids = batch['id'].tolist() + sources = batch['source'].tolist() + thinkings = batch['thinking_raw'].tolist() + query_raws = batch['query_raw'].tolist() + cot_compresseds = batch['cot_compressed'].tolist() + + anchor_emb = get_embeddings(model, template, queries) + + for i, (rid, src, vec) in enumerate(zip(ids, sources, anchor_emb)): + res = ( + tbl.search(vec.astype(np.float32).tolist()) + .metric('dot') + .limit(max(ks)) + .select(['id', 'source', 'query_compressed', 'cot_compressed', + 'thinking_raw', 'query_raw']) + .to_list() + ) + hit_ids = [item['id'] for item in res] + try: + rank = hit_ids.index(rid) + except ValueError: + rank = -1 + + for k in ks: + if 0 <= rank < k: + hits[k] += 1 + per_source_hits.setdefault(src, {kk: 0 for kk in ks})[k] += 1 + per_source_total[src] = per_source_total.get(src, 0) + 1 + per_source_hits.setdefault(src, {kk: 0 for kk in ks}) + + top1 = res[0] if res else {} + debug_records.append({ + 'id': rid, + 'source': src, + 'rank': rank, + 'query_raw': query_raws[i], + 'query_compressed': queries[i], + 'cot_compressed': cot_compresseds[i], + 'thinking_raw': thinkings[i][:2000], + 'top1_id': top1.get('id'), + 'top1_source': top1.get('source'), + 'top1_query_compressed': top1.get('query_compressed'), + 'top1_cot_compressed': top1.get('cot_compressed'), + 'top1_query_raw': top1.get('query_raw'), + 'top1_thinking_raw': (top1.get('thinking_raw') or '')[:2000], + 'top1_is_self': top1.get('id') == rid, + }) + + sys.stderr.write(f' probed {batch_end}/{n_sample}\n') + + print(f'\n=== Self-Recall @ k (n={n_sample}, seed={args.seed}) ===') + for k in ks: + print(f' recall@{k:<3} = {hits[k]/n_sample:.4f} ({hits[k]}/{n_sample})') + + print(f'\n=== Per-source recall@{max(ks)} ===') + for src in sorted(per_source_total, key=lambda s: -per_source_total[s]): + tot = per_source_total[src] + h = per_source_hits.get(src, {}).get(max(ks), 0) + print(f' {src:<48s} {h/tot:.4f} ({h}/{tot})') + + os.makedirs(os.path.dirname(args.output) or '.', exist_ok=True) + with open(args.output, 'w', encoding='utf-8') as f: + for rec in debug_records: + f.write(json.dumps(rec, ensure_ascii=False) + '\n') + print(f'\n[debug] {len(debug_records)} records saved to {args.output}') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/embedding/make_embedding_dataset.py b/cookbook/exp/embedding/make_embedding_dataset.py new file mode 100644 index 000000000..847f222fc --- /dev/null +++ b/cookbook/exp/embedding/make_embedding_dataset.py @@ -0,0 +1,758 @@ +"""Offline compression pipeline: raw datasets → condenser → pre-compressed embedding dataset. + +Loads think/index/hard datasets, compresses query/cot/negatives via vLLM condenser +with API fallback, saves a single HF Dataset ready for embedding training. + +Output schema: {anchor_text, positive_text, negative_texts, source} + +Launch (8 GPUs — 4 for vLLM condenser): + python cookbook/exp/embedding/make_embedding_dataset.py +""" +import json +import os +import re +import sys +import threading +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path +from typing import Any, Dict, List, Optional + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import SamplingParams +from twinkle.sampler import vLLMSampler +from twinkle.template import Qwen3_5Template +from twinkle.utils.parallel import PosixFileLock +from twinkle_agentic.protocol.openai import OpenAI as OpenAIClient + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from dataset_think import get_dataset as get_dataset_think # noqa: E402 +from dataset_index import get_dataset as get_dataset_index # noqa: E402 +from dataset_hard import get_dataset as get_dataset_hard # noqa: E402 + +logger = get_logger() + +# -- Model config ------------------------------------------------------------- +CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') +TEMPLATE_NAME = 'Qwen3_5Template' + +# -- GPU placement (condenser only) ------------------------------------------- +CONDENSER_GPUS = int(os.environ.get('CONDENSER_GPUS', 8)) + +# -- Dataset caps ------------------------------------------------------------- +TOTAL_SAMPLES: Optional[int] = None +THINK_CAP: Optional[int] = int(os.environ.get('THINK_CAP', 100_000)) +INDEX_CAP: Optional[int] = int(os.environ.get('INDEX_CAP', 100_000)) +HARD_CAP: Optional[int] = int(os.environ.get('HARD_CAP', 0)) or None +HARD_MAX_NEGATIVES = int(os.environ.get('HARD_MAX_NEGATIVES', 8)) + +# -- Compression params ------------------------------------------------------- +MIN_TEXT_CHARS = 256 +DATASET_MAX_TOKENS = 32768 +COMPRESS_TEMPERATURE = 0.2 +COMPRESS_TOP_P = 0.5 +COMPRESS_MAX_MODEL_LEN = 32768 +BATCH_SIZE = int(os.environ.get('COMPRESS_BATCH_SIZE', 128)) + +# -- API fallback ------------------------------------------------------------- +COMPRESS_API_KEY = os.environ.get('COMPRESS_API_KEY', '') +COMPRESS_BASE_URL = os.environ.get('COMPRESS_BASE_URL', 'https://dashscope.aliyuncs.com/compatible-mode/v1') +COMPRESS_MODEL = os.environ.get('COMPRESS_MODEL', 'qwen3.7-max') +API_MIN_INTERVAL = float(os.environ.get('API_MIN_INTERVAL', 0.1)) +API_CONCURRENCY = int(os.environ.get('API_CONCURRENCY', 24)) +SAMPLER_TIMEOUT = float(os.environ.get('SAMPLER_TIMEOUT', 300)) + +# -- Output ------------------------------------------------------------------- +OUTPUT_DIR = os.environ.get('EMB_DATASET_OUTPUT', './output/embedding_dataset') +RESULTS_JSONL = f'{OUTPUT_DIR}/results.jsonl' +PROGRESS_FILE = f'{OUTPUT_DIR}/progress.json' + +# ============================================================================= +# Prompts +# ============================================================================= + +COMPRESS_SYSTEM = """\ +You are a compression and summary assistant. For the (query, source) pair, emit a Markdown \ +answer with TWO sections, designed to pair with the `extract_compressed` tool: \ +the reader absorbs `## Summary` directly, then calls `extract_compressed` \ +on any topic-key listed under `## More` to recover its \ +fuller content. + + `## Summary` — extreme-density text the reader reads directly. + `## More` — a topic index whose keys are valid arguments \ +to `extract_compressed` for recovering material not captured inline. + +Together the two sections must form a COMPLETE, NON-DISTORTING inventory of the \ +source for the query — nothing essential lost, nothing implied that the source \ +does not support. NO preamble, NO meta-commentary, NO code fences wrapping the \ +whole output. + +Output skeleton: + +## Summary +Topic: + + +## More +- : +- ... + +Format selection for the inline body (pick the MOST COMPACT form per query, mix \ +when helpful): +- Interface / signature → code notation directly: `func(a:int)->str` +- Factual / entity → telegraphic prose; drop function words; ":" for "is", "," \ +for "has" +- Skill / how-to / usage → lead with `Use when: `; numbered telegraphic \ +steps `1.do X 2.then Y`; close with `Output: ` when relevant +- Procedural → numbered short steps +- Analytical / design → hierarchical bullets with abbreviations + +`## Summary` rules: +1. TOPIC LINE — line 1 is ALWAYS `Topic: `, even when the \ +query is narrow. Anchors both the reader and the tool. +2. DENSITY — every token in the body carries query-relevant signal; cut filler. +3. PRIMARY-COMPLETE — never silently drop a fact essential to answering the \ +query. Anything cut for length MUST appear as a key under \ +`## More`. +4. NON-MISLEADING — phrasing must not let the reader infer anything the source \ +does not support; partial truths that mislead are worse than honest omissions \ +flagged in the index. +5. SELF-CONTAINED — the reader can act on the answer without re-opening the source. +6. FAITHFUL — only content the source supports; no fabrication, no extrapolation. +7. LANGUAGE — match the source language. +8. NO outer code fences around the whole answer; no meta-commentary. + +`## More` rules (MANDATORY — this section is never omitted): +1. FORMAT — each bullet is `- : `: + • topic-key — short, unambiguous, grounded in source vocabulary so the \ +`extract_compressed` tool can locate the aspect (e.g. `decorators`, \ +`error handling`, `pitfalls`). + • hint — tells WHAT the reader gains by expanding (concrete numbers, code \ +listings, secondary cases, edge details, related context, …); do NOT restate \ +the inline answer. +2. CRITERION — each bullet names an aspect that EXISTS in the source but is \ +NOT fully captured inline. Material that genuinely fits inline without \ +distortion MUST NOT be duplicated here. +3. FAITHFUL — hints must be grounded in the source; never speculate or invent. +4. ORDER — by relevance to the query, then by importance. +5. EMPTY CASE — if the source is so short / single-purpose that everything \ +fits inline, write a single line `- (none)`. + +Now begin.\ +""" + +COMPRESS_USER = ( + 'Downstream model will read your compressed block to decide whether to ' + 'expand it. Compress faithfully: preserve the passage topic + core facts. ' + 'Do NOT invent facts. Do NOT drop major facts. Do NOT write meta-commentary ' + 'about the Query (never write "Query info: absent", "no X mention", etc.); ' + 'if the passage does not address the Query, still summarize the passage. ' + 'CRITICAL LANGUAGE RULE: detect the dominant language of the Passage ' + '(NOT the Query, NOT this instruction) and write the ENTIRE output in that ' + 'same language; English passage → English output, Chinese passage → ' + 'Chinese output, Japanese passage → Japanese output. NEVER translate, ' + 'NEVER mix languages, NEVER copy these instructions into the output.\n\n' + '## Query (ordering hint only — still summarize the whole passage)\n{query}\n\n' + '## Passage\n{text}') + +EMBED_QUERY_Q = ( + 'Summarize this query for retrieval. ' + 'The body of ## Summary MUST follow this EXACT 4-line template — ' + 'do NOT emit "Use when:", numbered procedure steps, or "Output:":\n' + 'Topic: \n' + 'Problem: \n' + 'Skill: \n' + 'Knowledge: \n' + 'Then emit the mandatory ## More section as usual. ' + 'Topic must name the specific pattern, never generic labels.') + +EMBED_QUERY_COT = ( + 'Summarize this reasoning trace for retrieval. ' + 'The body of ## Summary MUST follow this EXACT 4-line template — ' + 'do NOT emit "Use when:", numbered procedure steps, or "Output:":\n' + 'Topic: \n' + 'Problem: \n' + 'Skill: \n' + 'Knowledge: \n' + 'Then emit the mandatory ## More section as usual. ' + 'Topic must name the specific pattern, never generic labels.') + +EMBED_QUERY_Q_LEGACY = ( + 'What problem does this passage address, and what skill or method is needed? ' + 'Topic must name the specific pattern, never generic labels. ' + 'Compress into a retrieval-friendly need description.') + +EMBED_QUERY_COT_LEGACY = ( + 'Extract the reusable skill: trigger conditions, key steps, and expected output. ' + 'Topic names the method/pattern; format as "Use when: ...", numbered steps, ' + '"Output: ...". Compress into a standardized procedure for retrieval.') + +EMBED_QUERY_REASONIR_Q = ( + 'Extract the abstract PROBLEM TYPE from this query. ' + 'IGNORE all specific numbers, values, variable names, and parameters — ' + 'focus ONLY on the CLASS of problem and the METHODOLOGY required. ' + 'The body of ## Summary MUST follow this EXACT 4-line template:\n' + 'Topic: \n' + 'Problem: \n' + 'Skill: \n' + 'Knowledge: \n' + 'Then emit the mandatory ## More section as usual. ' + 'Topic must name the method class, never mention specific numbers.') + +EMBED_QUERY_REASONIR_COT = ( + 'Extract the abstract METHODOLOGY demonstrated in this solution. ' + 'IGNORE all specific numbers, values, and computed results — ' + 'focus ONLY on the general TECHNIQUE and key reasoning STEPS. ' + 'The body of ## Summary MUST follow this EXACT 4-line template:\n' + 'Topic: \n' + 'Problem: \n' + 'Skill: \n' + 'Knowledge: \n' + 'Then emit the mandatory ## More section as usual. ' + 'Topic must name the method class, never mention specific numbers.') + + +# ============================================================================= +# Validation & API fallback +# ============================================================================= + +_LEGACY_USE_WHEN_RE = re.compile(r'(?im)^\s*Use when\s*:') +_SCHEMA_MARKERS = ('Problem:', 'Skill:', 'Knowledge:') + + +def _is_truncated_compression(text: str, schema: str = 'new') -> bool: + if not text or not text.strip(): + return True + if '## More' not in text or '## Summary' not in text: + return True + after_more = text.split('## More', 1)[1].strip() + if not after_more: + return True + last_line = after_more.splitlines()[-1].strip() + if not (last_line.startswith('-') or last_line.endswith(')')): + return True + if schema == 'new': + summary_body = text.split('## Summary', 1)[1].split('## More', 1)[0] + if _LEGACY_USE_WHEN_RE.search(summary_body): + return True + if not all(marker in summary_body for marker in _SCHEMA_MARKERS): + return True + return False + + +_api_semaphore = threading.Semaphore(API_CONCURRENCY) +_api_bucket_lock = threading.Lock() +_api_tokens = [float(API_CONCURRENCY)] +_api_last_refill = [time.monotonic()] + + +def _api_throttle(): + """Token-bucket rate limiter: API_CONCURRENCY requests per API_MIN_INTERVAL*API_CONCURRENCY window.""" + _api_semaphore.acquire() + try: + with _api_bucket_lock: + now = time.monotonic() + elapsed = now - _api_last_refill[0] + refill = elapsed / API_MIN_INTERVAL + _api_tokens[0] = min(float(API_CONCURRENCY), _api_tokens[0] + refill) + _api_last_refill[0] = now + if _api_tokens[0] >= 1.0: + _api_tokens[0] -= 1.0 + else: + wait = (1.0 - _api_tokens[0]) * API_MIN_INTERVAL + _api_tokens[0] = 0.0 + time.sleep(wait) + finally: + _api_semaphore.release() + + +def _api_compress(api_client: OpenAIClient, prompt: Dict[str, Any]) -> Optional[str]: + _api_throttle() + trajectory = {'messages': prompt['messages']} + sp = SamplingParams(temperature=0.2, max_tokens=8192) + try: + reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) + except Exception as exc: + logger.warning(f'[api_fallback] error: {exc}') + return None + content = (reply.get('content') or '').strip() + if not content: + return None + m = re.match(r'^```[a-zA-Z]*\n(.*?)\n```\s*$', content, re.DOTALL) + if m: + content = m.group(1).strip() + return content + + +# ============================================================================= +# Core compression logic +# ============================================================================= + +def _extract_query_cot(row: Dict[str, Any]): + messages = row.get('messages') or [] + query, cot = '', '' + for m in messages: + if not isinstance(m, dict): + continue + role = m.get('role') or '' + if role == 'user' and not query: + query = (m.get('content') or '').strip() + elif role == 'assistant': + cot = (m.get('reasoning_content') or '').strip() + break + return query, cot + + +def _compress_batch_phase1( + rows: List[Dict[str, Any]], + condenser_sampler, + compress_params: SamplingParams, + special_tokens: set, + source_type: str, +) -> Optional[Dict[str, Any]]: + """Phase 1 (GPU): build prompts → vLLM sample → validate. Returns state for phase 2.""" + _MAX_COT_CHARS = 30_000 + + if source_type == 'hard': + return _compress_hard_phase1(rows, condenser_sampler, compress_params, + special_tokens, source_type) + + prompts: List[Optional[Dict[str, Any]]] = [] + meta: List[Dict[str, Any]] = [] + for i, row in enumerate(rows): + query, cot = _extract_query_cot(row) + if not query or len(cot) < MIN_TEXT_CHARS or len(cot) > _MAX_COT_CHARS: + continue + schema = 'legacy' if (i % 2 == 0) else 'new' + q_hint = EMBED_QUERY_Q_LEGACY if schema == 'legacy' else EMBED_QUERY_Q + c_hint = EMBED_QUERY_COT_LEGACY if schema == 'legacy' else EMBED_QUERY_COT + + if len(query) < MIN_TEXT_CHARS: + prompts.append(None) + else: + user = COMPRESS_USER.format(query=q_hint, text=query) + prompts.append({'messages': [ + {'role': 'system', 'content': COMPRESS_SYSTEM}, + {'role': 'user', 'content': user}, + ]}) + user_c = COMPRESS_USER.format(query=c_hint, text=cot) + prompts.append({'messages': [ + {'role': 'system', 'content': COMPRESS_SYSTEM}, + {'role': 'user', 'content': user_c}, + ]}) + meta.append({'query_raw': query, 'cot_raw': cot, 'schema': schema, + 'q_hint': q_hint, 'source': source_type, + 'row_id': row.get('id', str(i))}) + + if not prompts: + return {'final': []} + + sampler_input = [p for p in prompts if p is not None] + sampler_pos = [ri for ri, p in enumerate(prompts) if p is not None] + try: + sampler_responses = condenser_sampler.sample(sampler_input, compress_params) + except Exception as exc: + logger.warning(f'[compress] sampler error: {exc}') + sampler_responses = [None] * len(sampler_input) + + responses = [None] * len(prompts) + for resp, pos in zip(sampler_responses, sampler_pos): + responses[pos] = resp + + decoded: List[str] = [] + fallback_indices: List[int] = [] + for ri in range(len(prompts)): + pair_idx = ri // 2 + schema = meta[pair_idx]['schema'] + if prompts[ri] is None: + decoded.append(meta[pair_idx]['query_raw']) + continue + resp = responses[ri] + seq = resp.sequences[0] if resp and resp.sequences else None + text = '' + if seq and seq.stop_reason != 'length' and seq.decoded: + text = seq.decoded + for tok in special_tokens: + text = text.replace(tok, '') + text = text.rstrip() + if not _is_truncated_compression(text, schema): + decoded.append(text) + else: + decoded.append('') + fallback_indices.append(ri) + + return {'prompts': prompts, 'meta': meta, 'decoded': decoded, + 'fallback_indices': fallback_indices} + + +def _compress_batch_phase2( + state: Dict[str, Any], + api_client: OpenAIClient, +) -> List[Dict[str, Any]]: + """Phase 2 (no GPU): API fallback → build results.""" + if 'final' in state: + return state['final'] + + prompts = state['prompts'] + decoded = state['decoded'] + fallback_indices = state['fallback_indices'] + is_hard = state.get('hard', False) + meta = state.get('meta') # None for hard + + # Track which prompts used API fallback + api_set: set = set() + if fallback_indices: + api_futures = {} + with ThreadPoolExecutor(max_workers=API_CONCURRENCY) as pool: + for ri in fallback_indices: + api_futures[pool.submit(_api_compress, api_client, prompts[ri])] = ri + for fut in as_completed(api_futures): + ri = api_futures[fut] + api_result = fut.result() + schema = 'new' if is_hard else meta[ri // 2]['schema'] + if api_result and not _is_truncated_compression(api_result, schema): + decoded[ri] = api_result + api_set.add(ri) + + state['api_set'] = api_set + if is_hard: + return _build_hard_results(state) + return _build_think_index_results(state) + + +def _build_think_index_results(state: Dict[str, Any]) -> List[Dict[str, Any]]: + meta = state['meta'] + decoded = state['decoded'] + api_set = state.get('api_set', set()) + results = [] + for pair_idx in range(len(meta)): + q_text = decoded[pair_idx * 2] + c_text = decoded[pair_idx * 2 + 1] + if not q_text or not c_text: + continue + q_method = 'api' if (pair_idx * 2) in api_set else 'vllm' + c_method = 'api' if (pair_idx * 2 + 1) in api_set else 'vllm' + results.append({ + 'anchor_text': q_text, + 'positive_text': c_text, + 'negative_texts': [], + 'source': meta[pair_idx]['source'], + 'query_raw': meta[pair_idx]['query_raw'], + 'cot_raw': meta[pair_idx]['cot_raw'], + 'anchor_method': q_method, + 'positive_method': c_method, + }) + return results + + +def _build_hard_results(state: Dict[str, Any]) -> List[Dict[str, Any]]: + group_sizes = state['group_sizes'] + decoded = state['decoded'] + source_type = state['source_type'] + raw_groups = state['raw_groups'] + api_set = state.get('api_set', set()) + results = [] + offset = 0 + for gi, gs in enumerate(group_sizes): + q_text = decoded[offset] + c_text = decoded[offset + 1] + if not q_text or not c_text: + offset += gs + continue + neg_texts = [] + neg_raws = [] + neg_methods = [] + for ni in range(2, gs): + nt = decoded[offset + ni] + if nt: + neg_texts.append(nt) + neg_raws.append(raw_groups[gi]['negs_raw'][ni - 2]) + neg_methods.append('api' if (offset + ni) in api_set else 'vllm') + q_method = 'api' if offset in api_set else 'vllm' + c_method = 'api' if (offset + 1) in api_set else 'vllm' + results.append({ + 'anchor_text': q_text, + 'positive_text': c_text, + 'negative_texts': neg_texts, + 'source': source_type, + 'query_raw': raw_groups[gi]['query_raw'], + 'cot_raw': raw_groups[gi]['cot_raw'], + 'negs_raw': neg_raws, + 'anchor_method': q_method, + 'positive_method': c_method, + 'neg_methods': neg_methods, + }) + offset += gs + return results + + +def _compress_hard_phase1( + rows: List[Dict[str, Any]], + condenser_sampler, + compress_params: SamplingParams, + special_tokens: set, + source_type: str, +) -> Dict[str, Any]: + """Phase 1 for hard rows: vLLM sample + validate. Returns state for phase 2.""" + _MAX_COT_CHARS = 30_000 + + prompts: List[Dict[str, Any]] = [] + group_sizes: List[int] = [] + row_ids: List[str] = [] + raw_groups: List[Dict[str, Any]] = [] + + for row in rows: + query, cot = _extract_query_cot(row) + if not query or not cot or len(cot) > _MAX_COT_CHARS: + continue + negatives = row.get('negatives') or [] + valid_negs = [n for n in negatives + if n and len(n) <= _MAX_COT_CHARS] + + user_q = COMPRESS_USER.format(query=EMBED_QUERY_REASONIR_Q, text=query) + prompts.append({'messages': [ + {'role': 'system', 'content': COMPRESS_SYSTEM}, + {'role': 'user', 'content': user_q}, + ]}) + user_c = COMPRESS_USER.format(query=EMBED_QUERY_REASONIR_COT, text=cot) + prompts.append({'messages': [ + {'role': 'system', 'content': COMPRESS_SYSTEM}, + {'role': 'user', 'content': user_c}, + ]}) + for neg in valid_negs: + user_n = COMPRESS_USER.format(query=EMBED_QUERY_REASONIR_COT, text=neg) + prompts.append({'messages': [ + {'role': 'system', 'content': COMPRESS_SYSTEM}, + {'role': 'user', 'content': user_n}, + ]}) + group_sizes.append(2 + len(valid_negs)) + row_ids.append(row.get('id', '')) + raw_groups.append({'query_raw': query, 'cot_raw': cot, 'negs_raw': valid_negs}) + + if not prompts: + return {'hard': True, 'prompts': [], 'group_sizes': [], 'row_ids': [], + 'decoded': [], 'fallback_indices': [], 'source_type': source_type, + 'raw_groups': []} + + try: + responses = condenser_sampler.sample(prompts, compress_params) + except Exception as exc: + logger.warning(f'[compress-hard] sampler error: {exc}') + responses = [None] * len(prompts) + + decoded: List[str] = [] + fallback_indices: List[int] = [] + for ri, resp in enumerate(responses): + seq = resp.sequences[0] if resp and resp.sequences else None + text = '' + if seq and seq.stop_reason != 'length' and seq.decoded: + text = seq.decoded + for tok in special_tokens: + text = text.replace(tok, '') + text = text.rstrip() + if text and not _is_truncated_compression(text, 'new'): + decoded.append(text) + else: + decoded.append('') + fallback_indices.append(ri) + + return {'hard': True, 'prompts': prompts, 'group_sizes': group_sizes, + 'row_ids': row_ids, 'decoded': decoded, + 'fallback_indices': fallback_indices, 'source_type': source_type, + 'raw_groups': raw_groups} + + +# ============================================================================= +# Main pipeline +# ============================================================================= + +def main(): + device_groups = [ + DeviceGroup(name='condenser_sampler', + ranks=list(range(CONDENSER_GPUS)), + device_type='GPU'), + ] + condenser_mesh = DeviceMesh.from_sizes( + world_size=CONDENSER_GPUS, dp_size=CONDENSER_GPUS) + twinkle.initialize(mode='ray', nproc_per_node=CONDENSER_GPUS, groups=device_groups) + + # -- Load raw datasets ---------------------------------------------------- + from datasets import Dataset as HFDataset + + dataset_think = get_dataset_think(total=TOTAL_SAMPLES, load_from_cache_file=True) + if THINK_CAP and len(dataset_think.dataset) > THINK_CAP: + dataset_think.dataset = dataset_think.dataset.select(range(THINK_CAP)) + ds_think = dataset_think.dataset + logger.info(f'[load] think={len(ds_think)}') + + ds_index_obj = get_dataset_index(total=None, load_from_cache_file=True) + ds_index = ds_index_obj.dataset + if INDEX_CAP and len(ds_index) > INDEX_CAP: + ds_index = ds_index.select(range(INDEX_CAP)) + logger.info(f'[load] index={len(ds_index)}') + + ds_hard_raw = get_dataset_hard(max_negatives=HARD_MAX_NEGATIVES, load_from_cache_file=True) + if HARD_CAP and len(ds_hard_raw) > HARD_CAP: + ds_hard_raw = ds_hard_raw.select(range(HARD_CAP)) + n_hard = len(ds_hard_raw) + logger.info(f'[load] hard={n_hard}') + + # Convert hard to messages schema + hard_rows_list = [] + if n_hard > 0: + h_ids = ds_hard_raw['id'] + h_queries = ds_hard_raw['query'] + h_cots = ds_hard_raw['cot'] + h_responses = ds_hard_raw['response'] if 'response' in ds_hard_raw.column_names else [''] * n_hard + h_negatives = ds_hard_raw['negatives'] + for i in range(n_hard): + hard_rows_list.append({ + 'id': h_ids[i], + 'messages': [ + {'role': 'user', 'content': h_queries[i]}, + {'role': 'assistant', 'reasoning_content': h_cots[i], + 'content': h_responses[i] or ''}, + ], + 'negatives': h_negatives[i], + }) + + # Batch-convert HF Datasets to list-of-dicts + def _ds_to_rows(ds): + return [dict(zip(ds.column_names, vals)) for vals in zip(*(ds[c] for c in ds.column_names))] + + think_rows = _ds_to_rows(ds_think) + index_rows = _ds_to_rows(ds_index) + + # -- Setup condenser ------------------------------------------------------ + condenser_template = Qwen3_5Template( + model_id=CONDENSE_MODEL_ID, max_length=DATASET_MAX_TOKENS, + enable_thinking=False, truncation_strategy='delete') + special_tokens = set(condenser_template.tokenizer.all_special_tokens) + + condenser_sampler = vLLMSampler( + model_id=CONDENSE_MODEL_ID, + engine_args={'gpu_memory_utilization': 0.8, 'max_model_len': COMPRESS_MAX_MODEL_LEN}, + device_mesh=condenser_mesh, + remote_group='condenser_sampler', + ) + condenser_sampler.set_template( + TEMPLATE_NAME, model_id=CONDENSE_MODEL_ID, enable_thinking=False, + truncation_strategy='delete', max_length=DATASET_MAX_TOKENS) + condenser_sampler._ray_get_timeout = SAMPLER_TIMEOUT + compress_params = SamplingParams( + max_tokens=8192, temperature=COMPRESS_TEMPERATURE, + top_p=COMPRESS_TOP_P, num_samples=1) + + api_client = OpenAIClient( + model=COMPRESS_MODEL, api_key=COMPRESS_API_KEY, base_url=COMPRESS_BASE_URL) + + # -- Resume support ---------------------------------------------------------- + os.makedirs(OUTPUT_DIR, exist_ok=True) + progress = {'think': 0, 'index': 0, 'hard': 0} + if os.path.exists(PROGRESS_FILE): + with open(PROGRESS_FILE, 'r') as f: + progress = json.load(f) + logger.info(f'[resume] loaded progress: {progress}') + + _results_lock = PosixFileLock(RESULTS_JSONL + '.lock') + + def _flush_results(records: List[Dict[str, Any]]): + if not records: + return + lines = [json.dumps(r, ensure_ascii=False) + '\n' for r in records] + with _results_lock: + with open(RESULTS_JSONL, 'a', encoding='utf-8') as f: + f.writelines(lines) + + def _save_progress(): + tmp = PROGRESS_FILE + '.tmp' + with open(tmp, 'w') as f: + json.dump(progress, f) + os.replace(tmp, PROGRESS_FILE) + + # -- Process in batches (pipelined: vLLM batch N+1 overlaps API fallback N) - + total_flushed = 0 + if os.path.exists(RESULTS_JSONL): + with open(RESULTS_JSONL, 'r', encoding='utf-8') as f: + total_flushed = sum(1 for l in f if l.strip()) + if total_flushed: + logger.info(f'[resume] {total_flushed} records already in results.jsonl') + + def _process_source(rows, source_type, label): + nonlocal total_flushed + n_total = len(rows) + skip = progress.get(source_type, 0) + if skip >= n_total: + logger.info(f'[{label}] skipped (already done {skip}/{n_total})') + return + if skip > 0: + logger.info(f'[{label}] resuming from row {skip}/{n_total}') + + bg_pool = ThreadPoolExecutor(max_workers=1) + pending = None # (future, batch_start, batch_len) + + def _drain_pending(): + nonlocal total_flushed, pending + if pending is None: + return + fut, p_start, p_len = pending + batch_results = fut.result() + _flush_results(batch_results) + total_flushed += len(batch_results) + progress[source_type] = p_start + p_len + _save_progress() + pending = None + + for start in range(skip, n_total, BATCH_SIZE): + batch = rows[start:start + BATCH_SIZE] + state = _compress_batch_phase1( + batch, condenser_sampler, compress_params, + special_tokens, source_type) + _drain_pending() + pending = ( + bg_pool.submit(_compress_batch_phase2, state, api_client), + start, len(batch)) + n_done = start + len(batch) + if n_done % (BATCH_SIZE * 10) == 0 or n_done >= n_total: + logger.info(f'[{label}] {n_done}/{n_total} vLLM done, ' + f'{total_flushed} records flushed (last batch pending)') + + _drain_pending() + bg_pool.shutdown(wait=False) + logger.info(f'[{label}] complete, {total_flushed} total records flushed') + + _process_source(hard_rows_list, 'hard', 'hard') + _process_source(think_rows, 'think', 'think') + _process_source(index_rows, 'index', 'index') + + # -- Convert JSONL → HF Dataset ------------------------------------------- + logger.info(f'[save] converting results.jsonl to HF Dataset...') + all_results = [] + with open(RESULTS_JSONL, 'r', encoding='utf-8') as f: + for line_no, line in enumerate(f, 1): + if not line.strip(): + continue + try: + all_results.append(json.loads(line)) + except json.JSONDecodeError: + logger.warning(f'[save] skipping malformed line {line_no} (truncated resume?)') + logger.info(f'[save] total records: {len(all_results)}') + out_ds = HFDataset.from_dict({ + 'anchor_text': [r['anchor_text'] for r in all_results], + 'positive_text': [r['positive_text'] for r in all_results], + 'negative_texts': [r['negative_texts'] for r in all_results], + 'source': [r['source'] for r in all_results], + 'query_raw': [r.get('query_raw', '') for r in all_results], + 'cot_raw': [r.get('cot_raw', '') for r in all_results], + 'negs_raw': [r.get('negs_raw', []) for r in all_results], + }) + out_ds.save_to_disk(OUTPUT_DIR + '/dataset') + logger.info(f'[save] dataset saved to {OUTPUT_DIR}/dataset') + logger.info(f'[stats] think={sum(1 for r in all_results if r["source"]=="think")} ' + f'index={sum(1 for r in all_results if r["source"]=="index")} ' + f'hard={sum(1 for r in all_results if r["source"]=="hard")}') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/embedding/train_embedding_full_ddp.py b/cookbook/exp/embedding/train_embedding_full_ddp.py index bc69c56fd..97ab3b128 100644 --- a/cookbook/exp/embedding/train_embedding_full_ddp.py +++ b/cookbook/exp/embedding/train_embedding_full_ddp.py @@ -1,277 +1,60 @@ -"""LoRA embedding training with online compression via frozen vLLM condenser. +"""Full-parameter embedding training on pre-compressed dataset. -Architecture (8 GPUs total): - - Ranks 0-3 (``model``): Trainable embedding model with LoRA, InfoNCE loss. - - Ranks 4-7 (``condenser_sampler``): Frozen vLLM condenser for online compression. +Reads the pre-compressed HF Dataset produced by make_embedding_dataset.py, +encodes features, trains with InfoNCE loss. -When the condenser sampler truncates or regresses to the legacy schema, an -external OpenAI-compatible API produces the correct compression. The failure is -logged to failures.jsonl for offline SFT data regeneration. +Architecture (4 GPUs): + - Ranks 0-3: Trainable embedding model, InfoNCE loss. Launch: - python cookbook/exp/train_embedding_lora_ddp.py + python cookbook/exp/embedding/train_embedding_full_ddp.py """ -import hashlib -import json import os -import re -import sys -import threading import time -from concurrent.futures import ThreadPoolExecutor -from pathlib import Path from typing import Any, Dict, List, Literal, Optional import swanlab import twinkle from twinkle import DeviceGroup, DeviceMesh, get_device_placement, get_logger -from twinkle.data_format import SamplingParams -from twinkle.dataloader import DataLoader from twinkle.loss import InfonceLoss from twinkle.metric import EmbeddingMetric from twinkle.model import TransformersModel from twinkle.processor import InputProcessor -from twinkle.sampler import vLLMSampler from twinkle.template import Qwen3_5Template, Template -from twinkle.utils.parallel import PosixFileLock -from twinkle_agentic.protocol.openai import OpenAI as OpenAIClient - -sys.path.insert(0, str(Path(__file__).resolve().parent)) -from dataset_think import get_dataset as get_dataset_think # noqa: E402 -from dataset_index import get_dataset as get_dataset_index # noqa: E402 logger = get_logger() # -- Backend selection -------------------------------------------------------- BACKEND: Literal['transformers', 'megatron'] = 'transformers' -# Condenser (online compression + LoRA self-improvement); embedding model trains LoRA on top of MODEL_ID. -CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') MODEL_ID = os.environ.get('MODEL_ID', 'ms://Qwen/Qwen3.5-4B') -TEMPLATE_NAME = 'Qwen3_5Template' -# -- GPU placement (8 total) -------------------------------------------------- -MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 4)) -CONDENSER_SAMPLER_GPUS = int(os.environ.get('CONDENSER_SAMPLER_GPUS', 4)) -NUM_GPUS = MODEL_GPUS + CONDENSER_SAMPLER_GPUS +# -- GPU placement ------------------------------------------------------------ +MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 8)) # -- Embedding training hyper-params ------------------------------------------ EMB_MAX_LENGTH = 8192 HARD_NEGATIVES = None -# 0.07 keeps gradient on diag pairs until cosine clears ~0.75; 0.03 saturated near 0.40. TEMPERATURE = 0.07 -BATCH_SIZE = int(os.environ.get('BATCH_SIZE', 32)) +BATCH_SIZE = int(os.environ.get('BATCH_SIZE', 64)) LEARNING_RATE = 1e-5 GRADIENT_ACCUMULATION_STEPS = 1 LOG_INTERVAL = 2 SAVE_INTERVAL = 2000 NUM_EPOCHS = 1 -TOTAL_SAMPLES: Optional[int] = None -# Post-build caps on each loader (None = no cap). Applied via .select() before mix. -THINK_CAP: Optional[int] = 400_000 -INDEX_CAP: Optional[int] = 400_000 +# -- Dataset path (output of make_embedding_dataset.py) ----------------------- +DATASET_PATH = os.environ.get('EMB_DATASET_PATH', 'ms://twinkle-kit/qth-embedding') MIX_SHUFFLE_SEED = 42 # -- Resume from checkpoint --------------------------------------------------- -# Empty by default — build_model falls back to MODEL_ID (the published emb model). -# Set both to point at a local in-progress run only when resuming the *same* schedule. RESUME_CHECKPOINT = os.environ.get('RESUME_CHECKPOINT', '') RESUME_STEP = int(os.environ.get('RESUME_STEP', 0)) -# -- Online-compression knobs ------------------------------------------------- -# Below this length, condenser fabricates content for open-ended short prompts; -# query passes through as qr verbatim and cot rows are dropped from training. -MIN_TEXT_CHARS = 256 -DATASET_MAX_TOKENS = 32768 -COMPRESS_TEMPERATURE = 0.2 -COMPRESS_TOP_P = 0.5 -COMPRESS_MAX_MODEL_LEN = 32768 - -# How many BATCH_SIZE chunks to fetch and compress in one vLLM call. -PREFETCH_BATCH_MULTIPLIER = int(os.environ.get('PREFETCH_BATCH_MULTIPLIER', 8)) - -# -- OpenAI API fallback for truncated compressions --------------------------- -COMPRESS_API_KEY = os.environ.get('COMPRESS_API_KEY', '') -COMPRESS_BASE_URL = os.environ.get('COMPRESS_BASE_URL', 'https://dashscope.aliyuncs.com/compatible-mode/v1') -COMPRESS_MODEL = os.environ.get('COMPRESS_MODEL', 'qwen3.7-max') -# Minimum gap between API calls (seconds); bounds dashscope qps under provider limits. -API_MIN_INTERVAL = float(os.environ.get('API_MIN_INTERVAL', 0.1)) -API_CONCURRENCY = int(os.environ.get('API_CONCURRENCY', 8)) -# vLLM sampler timeout (seconds); if a sample() call exceeds this, fall back to API. -SAMPLER_TIMEOUT = float(os.environ.get('SAMPLER_TIMEOUT', 300)) - -# -- Output paths ------------------------------------------------------------- -OUTPUT_DIR = f'./output/embedding_lora_{BACKEND}' -RESPONSE_LOG = os.environ.get('RESPONSE_LOG', f'./output/embedding_lora_{BACKEND}/responses.jsonl') -FAILURE_LOG = os.environ.get('FAILURE_LOG', f'./output/embedding_lora_{BACKEND}/failures.jsonl') - - -# ============================================================================= -# Prompts (from make_condenser_dataset.py — "## Summary" format) -# ============================================================================= - -COMPRESS_SYSTEM = """\ -You are a compression and summary assistant. For the (query, source) pair, emit a Markdown \ -answer with TWO sections, designed to pair with the `extract_compressed` tool: \ -the reader absorbs `## Summary` directly, then calls `extract_compressed` \ -on any topic-key listed under `## More` to recover its \ -fuller content. - - `## Summary` — extreme-density text the reader reads directly. - `## More` — a topic index whose keys are valid arguments \ -to `extract_compressed` for recovering material not captured inline. - -Together the two sections must form a COMPLETE, NON-DISTORTING inventory of the \ -source for the query — nothing essential lost, nothing implied that the source \ -does not support. NO preamble, NO meta-commentary, NO code fences wrapping the \ -whole output. - -Output skeleton: - -## Summary -Topic: - - -## More -- : -- ... - -Format selection for the inline body (pick the MOST COMPACT form per query, mix \ -when helpful): -- Interface / signature → code notation directly: `func(a:int)->str` -- Factual / entity → telegraphic prose; drop function words; ":" for "is", "," \ -for "has" -- Skill / how-to / usage → lead with `Use when: `; numbered telegraphic \ -steps `1.do X 2.then Y`; close with `Output: ` when relevant -- Procedural → numbered short steps -- Analytical / design → hierarchical bullets with abbreviations - -`## Summary` rules: -1. TOPIC LINE — line 1 is ALWAYS `Topic: `, even when the \ -query is narrow. Anchors both the reader and the tool. -2. DENSITY — every token in the body carries query-relevant signal; cut filler. -3. PRIMARY-COMPLETE — never silently drop a fact essential to answering the \ -query. Anything cut for length MUST appear as a key under \ -`## More`. -4. NON-MISLEADING — phrasing must not let the reader infer anything the source \ -does not support; partial truths that mislead are worse than honest omissions \ -flagged in the index. -5. SELF-CONTAINED — the reader can act on the answer without re-opening the source. -6. FAITHFUL — only content the source supports; no fabrication, no extrapolation. -7. LANGUAGE — match the source language. -8. NO outer code fences around the whole answer; no meta-commentary. - -`## More` rules (MANDATORY — this section is never omitted): -1. FORMAT — each bullet is `- : `: - • topic-key — short, unambiguous, grounded in source vocabulary so the \ -`extract_compressed` tool can locate the aspect (e.g. `decorators`, \ -`error handling`, `pitfalls`). - • hint — tells WHAT the reader gains by expanding (concrete numbers, code \ -listings, secondary cases, edge details, related context, …); do NOT restate \ -the inline answer. -2. CRITERION — each bullet names an aspect that EXISTS in the source but is \ -NOT fully captured inline. Material that genuinely fits inline without \ -distortion MUST NOT be duplicated here. -3. FAITHFUL — hints must be grounded in the source; never speculate or invent. -4. ORDER — by relevance to the query, then by importance. -5. EMPTY CASE — if the source is so short / single-purpose that everything \ -fits inline, write a single line `- (none)`. - -Now begin.\ -""" - -COMPRESS_USER = ( - 'Downstream model will read your compressed block to decide whether to ' - 'expand it. Compress faithfully: preserve the passage topic + core facts. ' - 'Do NOT invent facts. Do NOT drop major facts. Do NOT write meta-commentary ' - 'about the Query (never write "Query info: absent", "no X mention", etc.); ' - 'if the passage does not address the Query, still summarize the passage. ' - 'CRITICAL LANGUAGE RULE: detect the dominant language of the Passage ' - '(NOT the Query, NOT this instruction) and write the ENTIRE output in that ' - 'same language; English passage → English output, Chinese passage → ' - 'Chinese output, Japanese passage → Japanese output. NEVER translate, ' - 'NEVER mix languages, NEVER copy these instructions into the output.\n\n' - '## Query (ordering hint only — still summarize the whole passage)\n{query}\n\n' - '## Passage\n{text}') - - -# ============================================================================= -# Logging helpers -# ============================================================================= - -_response_lock: Optional[PosixFileLock] = None -_failure_lock: Optional[PosixFileLock] = None - -# Monotonic global sample id; per-batch index would alias across batches. -_sample_counter = 0 -_sample_counter_lock = threading.Lock() - -_api_throttle_lock = threading.Lock() -_api_last_call = [0.0] - - -def _api_throttle(): - with _api_throttle_lock: - gap = time.monotonic() - _api_last_call[0] - if gap < API_MIN_INTERVAL: - time.sleep(API_MIN_INTERVAL - gap) - _api_last_call[0] = time.monotonic() - - -def _next_sample_id() -> int: - global _sample_counter - with _sample_counter_lock: - sid = _sample_counter - _sample_counter += 1 - return sid - - -def _log_responses(query_resp_text: str, cot_resp_text: str, idx: int, - query_raw: str = '', cot_raw: str = ''): - global _response_lock - if _response_lock is None: - os.makedirs(os.path.dirname(RESPONSE_LOG) or '.', exist_ok=True) - _response_lock = PosixFileLock(RESPONSE_LOG + '.lock') - - record = { - 'idx': idx, - 'query_raw': query_raw, - 'cot_raw': cot_raw, - 'query_compressed': query_resp_text, - 'cot_compressed': cot_resp_text, - } - line = json.dumps(record, ensure_ascii=False, default=str) + '\n' - with _response_lock: - with open(RESPONSE_LOG, 'a', encoding='utf-8') as f: - f.write(line) - - -def _log_failure(source_text: str, query: str, compressed: str, batch_idx: int): - global _failure_lock - if _failure_lock is None: - os.makedirs(os.path.dirname(FAILURE_LOG) or '.', exist_ok=True) - _failure_lock = PosixFileLock(FAILURE_LOG + '.lock') - - qhash = hashlib.md5(query.strip().encode('utf-8')).hexdigest()[:8] - record = { - 'id': f'{batch_idx}__{qhash}', - 'source': 'online_failure', - 'query': query, - 'original_len': len(source_text), - 'compressed_len': len(compressed), - 'messages': [ - {'role': 'system', 'content': COMPRESS_SYSTEM}, - {'role': 'user', 'content': COMPRESS_USER.format(query=query, text=source_text)}, - {'role': 'assistant', 'content': compressed}, - ], - } - line = json.dumps(record, ensure_ascii=False, default=str) + '\n' - with _failure_lock: - with open(FAILURE_LOG, 'a', encoding='utf-8') as f: - f.write(line) +# -- Output ------------------------------------------------------------------- +OUTPUT_DIR = f'./output/embedding_full_{BACKEND}' # ============================================================================= @@ -327,121 +110,9 @@ def save_checkpoint(model, name: str): # ============================================================================= -# Compression prompt building +# Feature encoding # ============================================================================= -# Hard-templated hints: the condenser SFT prior maps `Skill` to the legacy -# `Use when: / numbered steps / Output:` skeleton on long inputs; embedding the -# exact 4-line body template + explicit negative constraints is the only way to -# override it deterministically across query and cot sides. -EMBED_QUERY_Q = ( - 'Summarize this query for retrieval. ' - 'The body of ## Summary MUST follow this EXACT 4-line template — ' - 'do NOT emit "Use when:", numbered procedure steps, or "Output:":\n' - 'Topic: \n' - 'Problem: \n' - 'Skill: \n' - 'Knowledge: \n' - 'Then emit the mandatory ## More section as usual. ' - 'Topic must name the specific pattern, never generic labels.') -EMBED_QUERY_COT = ( - 'Summarize this reasoning trace for retrieval. ' - 'The body of ## Summary MUST follow this EXACT 4-line template — ' - 'do NOT emit "Use when:", numbered procedure steps, or "Output:":\n' - 'Topic: \n' - 'Problem: \n' - 'Skill: \n' - 'Knowledge: \n' - 'Then emit the mandatory ## More section as usual. ' - 'Topic must name the specific pattern, never generic labels.') - -# Legacy schema (Use when: / numbered steps / Output:) — mixed in 50/50 with the -# new schema to expose the embedder to schema-invariant semantic alignment. -# Both query and cot of the SAME pair always use the SAME schema; cross-schema -# anchors and positives would re-introduce the schema asymmetry we just fixed. -EMBED_QUERY_Q_LEGACY = ( - 'What problem does this passage address, and what skill or method is needed? ' - 'Topic must name the specific pattern, never generic labels. ' - 'Compress into a retrieval-friendly need description.') -EMBED_QUERY_COT_LEGACY = ( - 'Extract the reusable skill: trigger conditions, key steps, and expected output. ' - 'Topic names the method/pattern; format as "Use when: ...", numbered steps, ' - '"Output: ...". Compress into a standardized procedure for retrieval.') - - -def _extract_query_cot(row: Dict[str, Any]): - messages = row.get('messages') or [] - query, cot = '', '' - for m in messages: - if not isinstance(m, dict): - continue - role = m.get('role') or '' - if role == 'user' and not query: - query = (m.get('content') or '').strip() - elif role == 'assistant': - cot = (m.get('reasoning_content') or '').strip() - break - return query, cot - - -def _build_compress_prompts(rows: List[Dict[str, Any]]) -> tuple: - """Build prompts for compressing both query and cot per row. - - Returns (prompts, valid_indices, raw_pairs, prompt_queries, passthrough, schemas) - where: - - prompts: flat-interleaved [query_0, cot_0, query_1, cot_1, ...]; ``None`` means - passthrough (use raw text directly, do not call sampler) - - valid_indices: which rows passed the min-length filter - - raw_pairs: [(query, cot), ...] - - prompt_queries: the query string used for each prompt (for failure logging) - - passthrough: parallel to prompts; non-None text means "use this verbatim as qc" - - schemas: parallel to prompts; 'new' or 'legacy', drives validator branch - """ - prompts: List[Optional[Dict[str, Any]]] = [] - valid_indices: List[int] = [] - raw_pairs: List[tuple] = [] - prompt_queries: List[str] = [] - passthrough: List[Optional[str]] = [] - schemas: List[str] = [] - # Conservative char budget: 32768 max_length - 8192 gen - ~2k prompt overhead = ~22k tokens. - # 30k cap bounds vLLM batch latency (vLLM batches by max prompt length). - _MAX_COT_CHARS = 30_000 - for i, row in enumerate(rows): - query, cot = _extract_query_cot(row) - if not query or len(cot) < MIN_TEXT_CHARS: - continue - if len(cot) > _MAX_COT_CHARS: - continue - valid_indices.append(i) - raw_pairs.append((query, cot)) - # 50/50 schema mix; same schema for query+cot of one pair to keep alignment. - schema = 'legacy' if (i % 2 == 0) else 'new' - q_hint = EMBED_QUERY_Q_LEGACY if schema == 'legacy' else EMBED_QUERY_Q - c_hint = EMBED_QUERY_COT_LEGACY if schema == 'legacy' else EMBED_QUERY_COT - # Short query bypasses condenser to avoid skeleton-induced hallucination. - if len(query) < MIN_TEXT_CHARS: - prompts.append(None) - passthrough.append(query) - else: - user = COMPRESS_USER.format(query=q_hint, text=query) - prompts.append({'messages': [ - {'role': 'system', 'content': COMPRESS_SYSTEM}, - {'role': 'user', 'content': user}, - ]}) - passthrough.append(None) - prompt_queries.append(q_hint) - schemas.append(schema) - user = COMPRESS_USER.format(query=c_hint, text=cot) - prompts.append({'messages': [ - {'role': 'system', 'content': COMPRESS_SYSTEM}, - {'role': 'user', 'content': user}, - ]}) - prompt_queries.append(c_hint) - passthrough.append(None) - schemas.append(schema) - return prompts, valid_indices, raw_pairs, prompt_queries, passthrough, schemas - - def _get_first_feature(decoded_text: str, template: Template, role: str) -> Optional[Dict[str, Any]]: if not decoded_text: return None @@ -450,77 +121,42 @@ def _get_first_feature(decoded_text: str, template: Template, role: str) -> Opti {'role': 'user', 'content': decoded_text}, {'role': 'assistant', 'content': 'Match the correct response here.'}, ]}) + if feat is None: + return None feat['labels'] = [1] else: feat = template.encode({'messages': [ {'role': 'user', 'content': 'Match the correct query here.'}, {'role': 'assistant', 'content': decoded_text}, ]}) + if feat is None: + return None feat['labels'] = [0] return feat -# ============================================================================= -# OpenAI API fallback -# ============================================================================= - -_LEGACY_USE_WHEN_RE = re.compile(r'(?im)^\s*Use when\s*:') -_SCHEMA_MARKERS = ('Problem:', 'Skill:', 'Knowledge:') - - -def _is_truncated_compression(text: str, schema: str = 'new') -> bool: - """Reject structurally incomplete OR schema-regressed condenser output. - - Triggers API fallback when the vLLM output: - * lacks ``## Summary`` / ``## More``, - * has an empty or unterminated ``## More`` bullet list, or - * (schema='new' only) regresses to the legacy ``Use when: / numbered-steps / - Output:`` skeleton instead of the mandated Problem/Skill/Knowledge 4-line - body — the dominant cot-side failure mode that drives sim < 0.45 drops on - the RAG index. - - For schema='legacy', body markers are intentionally NOT enforced: the legacy - template legitimately emits ``Use when:`` and the SFT prior already produces - that shape natively, so only structural completeness is checked. - """ - if not text or not text.strip(): - return True - if '## More' not in text or '## Summary' not in text: - return True - after_more = text.split('## More', 1)[1].strip() - if not after_more: - return True - last_line = after_more.splitlines()[-1].strip() - if not (last_line.startswith('-') or last_line.endswith(')')): - return True - if schema == 'new': - summary_body = text.split('## Summary', 1)[1].split('## More', 1)[0] - if _LEGACY_USE_WHEN_RE.search(summary_body): - return True - if not all(marker in summary_body for marker in _SCHEMA_MARKERS): - return True - return False - - -def _api_compress(api_client: OpenAIClient, prompt: Dict[str, Any]) -> Optional[str]: - """Call external API to compress when vLLM truncates.""" - _api_throttle() - trajectory = {'messages': prompt['messages']} - # Cap max_tokens to leave ample prompt headroom inside the API model context. - sp = SamplingParams(temperature=0.2, max_tokens=8192) - try: - reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) - except Exception as exc: - logger.warning(f'[api_fallback] error: {exc}') - return None - content = (reply.get('content') or '').strip() - if not content: - return None - # Strip outer code fence if present - m = re.match(r'^```[a-zA-Z]*\n(.*?)\n```\s*$', content, re.DOTALL) - if m: - content = m.group(1).strip() - return content +def _encode_batch( + rows: List[Dict[str, Any]], + emb_template: Template, +) -> List[Dict[str, Any]]: + """Encode pre-compressed texts into embedding features.""" + features: List[Dict[str, Any]] = [] + for row in rows: + anchor_text = row['anchor_text'] + positive_text = row['positive_text'] + negative_texts = row.get('negative_texts') or [] + + feat_q = _get_first_feature(anchor_text, emb_template, role='anchor') + feat_c = _get_first_feature(positive_text, emb_template, role='positive') + if not feat_q or not feat_c: + continue + features.append(feat_q) + features.append(feat_c) + for neg_text in negative_texts: + feat_neg = _get_first_feature(neg_text, emb_template, role='positive') + if feat_neg: + features.append(feat_neg) + return features # ============================================================================= @@ -528,42 +164,27 @@ def _api_compress(api_client: OpenAIClient, prompt: Dict[str, Any]) -> Optional[ # ============================================================================= def train(): - # -------- Device groups (2 groups) ---------------------------------------- device_groups = [ DeviceGroup(name='model', ranks=list(range(MODEL_GPUS)), device_type='GPU'), - DeviceGroup(name='condenser_sampler', - ranks=list(range(MODEL_GPUS, MODEL_GPUS + CONDENSER_SAMPLER_GPUS)), - device_type='GPU'), ] model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, dp_size=MODEL_GPUS) - condenser_sampler_mesh = DeviceMesh.from_sizes( - world_size=CONDENSER_SAMPLER_GPUS, dp_size=CONDENSER_SAMPLER_GPUS) - - twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, groups=device_groups) - - # -------- Data ----------------------------------------------------------- - dataset = get_dataset_think(total=TOTAL_SAMPLES, load_from_cache_file=True) - if THINK_CAP and len(dataset.dataset) > THINK_CAP: - dataset.dataset = dataset.dataset.select(range(THINK_CAP)) - if INDEX_CAP != 0: - from datasets import concatenate_datasets - ds_index = get_dataset_index(total=None, load_from_cache_file=True) - if INDEX_CAP and len(ds_index.dataset) > INDEX_CAP: - ds_index.dataset = ds_index.dataset.select(range(INDEX_CAP)) - n_think = len(dataset.dataset) - n_index = len(ds_index.dataset) - # Both loaders emit identical {id, source, messages} schema post-QP. - dataset.dataset = concatenate_datasets( - [dataset.dataset, ds_index.dataset]).shuffle(seed=MIX_SHUFFLE_SEED) - logger.info(f'[mix] think={n_think} + index={n_index} → total={len(dataset.dataset)}') - _mega_batch_size = BATCH_SIZE * PREFETCH_BATCH_MULTIPLIER - dataloader = DataLoader(dataset=dataset, batch_size=_mega_batch_size, shuffle=True) - total_forward_steps = len(dataloader) * PREFETCH_BATCH_MULTIPLIER * NUM_EPOCHS - optimizer_steps = total_forward_steps // GRADIENT_ACCUMULATION_STEPS - - # -------- Embedding model (4 GPU) ---------------------------------------- + twinkle.initialize(mode='ray', nproc_per_node=MODEL_GPUS, groups=device_groups) + + # -- Load pre-compressed dataset ------------------------------------------ + from twinkle.dataset import Dataset as TwinkleDataset, DatasetMeta + logger.info(f'[data] loading pre-compressed dataset from {DATASET_PATH}') + dataset = TwinkleDataset(DatasetMeta(dataset_id=DATASET_PATH), download_mode='force_redownload') + dataset = dataset.dataset.shuffle(seed=MIX_SHUFFLE_SEED) + logger.info(f'[data] {len(dataset)} rows loaded') + + # -- Compute steps -------------------------------------------------------- + rows_per_step = BATCH_SIZE + total_steps = (len(dataset) // rows_per_step) * NUM_EPOCHS + optimizer_steps = total_steps // GRADIENT_ACCUMULATION_STEPS + + # -- Model ---------------------------------------------------------------- model = build_model(model_mesh) model.set_processor(InputProcessor) model.set_loss(InfonceLoss, temperature=TEMPERATURE, use_batch=True, @@ -571,264 +192,78 @@ def train(): setup_optimizer(model, optimizer_steps) model.add_metric(EmbeddingMetric, is_training=True) - # -------- Condenser sampler (4 GPU, vLLM) -------------------------------- - emb_template = Qwen3_5Template(model_id=MODEL_ID, max_length=EMB_MAX_LENGTH, enable_thinking=False) - # Special tokens come from the condenser tokenizer because the leak we strip is in its decoded output. - condenser_template = Qwen3_5Template(model_id=CONDENSE_MODEL_ID, max_length=DATASET_MAX_TOKENS, - enable_thinking=False) - _special_tokens = set(condenser_template.tokenizer.all_special_tokens) - condenser_sampler = vLLMSampler( - model_id=CONDENSE_MODEL_ID, - engine_args={ - 'gpu_memory_utilization': 0.8, - 'max_model_len': COMPRESS_MAX_MODEL_LEN, - }, - device_mesh=condenser_sampler_mesh, - remote_group='condenser_sampler', - ) - condenser_sampler.set_template( - TEMPLATE_NAME, model_id=CONDENSE_MODEL_ID, enable_thinking=False, - truncation_strategy='delete', max_length=DATASET_MAX_TOKENS) - compress_params = SamplingParams( - max_tokens=8192, - temperature=COMPRESS_TEMPERATURE, - top_p=COMPRESS_TOP_P, - num_samples=1, - ) - - condenser_sampler._ray_get_timeout = SAMPLER_TIMEOUT - _sampler_epoch = 0 - - def _rebuild_sampler(): - """Kill stuck actors and recreate the vLLM sampler from scratch.""" - nonlocal condenser_sampler, _sampler_epoch - import ray - for actor in getattr(condenser_sampler, '_actors', []): - try: - ray.kill(actor, no_restart=True) - except Exception: - pass - logger.warning('[sampler] killed stuck actors, recreating sampler \u2026') - new = vLLMSampler( - model_id=CONDENSE_MODEL_ID, - engine_args={'gpu_memory_utilization': 0.8, 'max_model_len': COMPRESS_MAX_MODEL_LEN}, - device_mesh=condenser_sampler_mesh, - remote_group='condenser_sampler', - ) - new.set_template( - TEMPLATE_NAME, model_id=CONDENSE_MODEL_ID, enable_thinking=False, - truncation_strategy='delete', max_length=DATASET_MAX_TOKENS) - new._ray_get_timeout = SAMPLER_TIMEOUT - condenser_sampler = new - _sampler_epoch += 1 - logger.warning('[sampler] sampler rebuilt successfully') - - # -------- OpenAI API client for fallback --------------------------------- - api_client = OpenAIClient( - model=COMPRESS_MODEL, - api_key=COMPRESS_API_KEY, - base_url=COMPRESS_BASE_URL, - ) + emb_template = Qwen3_5Template( + model_id=MODEL_ID, max_length=EMB_MAX_LENGTH, + enable_thinking=False, truncation_strategy='delete') logger.info(get_device_placement()) logger.info(model.get_train_configs()) - logger.info(f'Total forward steps: {total_forward_steps}, optimizer steps: {optimizer_steps}') - if RESUME_STEP > 0: - logger.info(f'Resuming from step {RESUME_STEP}, checkpoint: {RESUME_CHECKPOINT}') - logger.info(f'Starting at epoch {RESUME_STEP // (total_forward_steps // NUM_EPOCHS)}, ' - f'skipping {RESUME_STEP - (RESUME_STEP // (total_forward_steps // NUM_EPOCHS)) * (total_forward_steps // NUM_EPOCHS)} batches') + logger.info(f'Total steps: {total_steps}, optimizer steps: {optimizer_steps}') swanlab.init(project='twinkle', config={ 'backend': BACKEND, 'model_id': MODEL_ID, - 'condense_model_id': CONDENSE_MODEL_ID, 'batch_size': BATCH_SIZE, 'lr': LEARNING_RATE, 'temperature': TEMPERATURE, 'emb_max_length': EMB_MAX_LENGTH, - 'DATASET_MAX_TOKENS': DATASET_MAX_TOKENS, + 'dataset_path': DATASET_PATH, }) - # -------- Train loop ----------------------------------------------------- - def _sample_batch(raw_batch): - """Compress via vLLM sampler; fall back to API on truncation.""" - _t_enter = time.monotonic() - compress_prompts, valid_indices, raw_pairs, prompt_queries, passthrough, schemas = \ - _build_compress_prompts(raw_batch) - _t_build = time.monotonic() - if len(compress_prompts) < 4: - return None + # -- Train loop ----------------------------------------------------------- + cur_step = 0 + _skip_rows = RESUME_STEP * rows_per_step # approximate rows to skip - # Only submit non-passthrough prompts to the sampler. - sampler_input = [p for p in compress_prompts if p is not None] - sampler_pos = [ri for ri, p in enumerate(compress_prompts) if p is not None] - if sampler_input: - try: - sampler_responses = condenser_sampler.sample(sampler_input, compress_params) - except Exception as exc: - logger.warning(f'[sampler] error \u2192 API fallback: {exc}') - sampler_responses = [None] * len(sampler_input) - if 'Timeout' in type(exc).__name__: - try: - _rebuild_sampler() - except Exception as re_exc: - logger.error(f'[sampler] rebuild failed: {re_exc}') - else: - sampler_responses = [] - _t_sample = time.monotonic() - - responses = [None] * len(compress_prompts) - for resp, pos in zip(sampler_responses, sampler_pos): - responses[pos] = resp - - # Extract decoded texts; detect truncations and fall back to API - decoded_texts: List[Optional[str]] = [None] * len(compress_prompts) - fallback_indices: List[int] = [] - for ri in range(len(compress_prompts)): - if passthrough[ri] is not None: - decoded_texts[ri] = passthrough[ri] + for epoch in range(NUM_EPOCHS): + for start in range(0, len(dataset), rows_per_step): + if start < _skip_rows: continue - resp = responses[ri] - seq = resp.sequences[0] if resp and resp.sequences else None - text = '' - if seq and seq.stop_reason != 'length' and seq.decoded: - text = seq.decoded - for tok in _special_tokens: - text = text.replace(tok, '') - text = text.rstrip() - - needs_fallback = (not seq or seq.stop_reason == 'length' - or _is_truncated_compression(text, schemas[ri])) - if not needs_fallback: - decoded_texts[ri] = text - else: - fallback_indices.append(ri) - - _api_calls = len(fallback_indices) - if fallback_indices: - from concurrent.futures import as_completed - api_futures = {} - with ThreadPoolExecutor(max_workers=API_CONCURRENCY) as api_pool: - for ri in fallback_indices: - api_futures[api_pool.submit(_api_compress, api_client, compress_prompts[ri])] = ri - for fut in as_completed(api_futures): - ri = api_futures[fut] - api_result = fut.result() - if api_result and not _is_truncated_compression(api_result, schemas[ri]): - decoded_texts[ri] = api_result - pair_idx = ri // 2 - q_raw, c_raw = raw_pairs[pair_idx] - source_text = q_raw if ri % 2 == 0 else c_raw - _log_failure(source_text, prompt_queries[ri], api_result, - valid_indices[pair_idx]) - else: - decoded_texts[ri] = '' - _t_api = time.monotonic() - - # Build embedding features from decoded texts - emb_features: List[Dict[str, Any]] = [] - for i in range(0, len(decoded_texts), 2): - q_text = decoded_texts[i] - c_text = decoded_texts[i + 1] - q_raw, c_raw = raw_pairs[i // 2] - _log_responses(q_text, c_text, _next_sample_id(), - query_raw=q_raw, cot_raw=c_raw) - feat_q = _get_first_feature(q_text, emb_template, role='anchor') - feat_c = _get_first_feature(c_text, emb_template, role='positive') - if feat_q and feat_c: - emb_features.append(feat_q) - emb_features.append(feat_c) - _t_feat = time.monotonic() - - logger.info( - f'[prefetch] prompts={len(sampler_input)} api={_api_calls} feats={len(emb_features)} | ' - f'build={_t_build - _t_enter:.1f}s ' - f'vllm={_t_sample - _t_build:.1f}s ' - f'api={_t_api - _t_sample:.1f}s feat={_t_feat - _t_api:.1f}s ' - f'total={_t_feat - _t_enter:.1f}s') - - _target = BATCH_SIZE * 2 - minibatches = [emb_features[i:i + _target] for i in range(0, len(emb_features), _target)] - minibatches = [mb for mb in minibatches if len(mb) >= 4] - return minibatches if minibatches else None - - cur_step = RESUME_STEP - _batches_per_epoch = len(dataloader) - _steps_per_mega = PREFETCH_BATCH_MULTIPLIER - _start_epoch = cur_step // (_batches_per_epoch * _steps_per_mega) if cur_step > 0 else 0 - _skip_batches_in_epoch = max(0, cur_step // _steps_per_mega - _start_epoch * _batches_per_epoch) - - _ema_prefetch = 0.0 - _ema_train = 0.0 - _ema_alpha = 0.1 - - prefetch_executor = ThreadPoolExecutor(max_workers=1) - for epoch in range(_start_epoch, NUM_EPOCHS): - if _skip_batches_in_epoch > 0: - dataloader.skip_consumed_samples(_skip_batches_in_epoch * _mega_batch_size) - batch_iter = iter(dataloader) - _skip_batches_in_epoch = 0 - - first = next(batch_iter, None) - future = prefetch_executor.submit(_sample_batch, first) if first else None - - for raw_mega_batch in batch_iter: + + batch_rows = dataset[start:start + rows_per_step] + # HF Dataset slicing returns dict of lists; convert to list of dicts + n_rows = len(batch_rows['anchor_text']) + rows_list = [{k: batch_rows[k][i] for k in batch_rows} + for i in range(n_rows)] + t0 = time.monotonic() - minibatches = future.result() if future else None - t_prefetch = time.monotonic() - t0 - future = prefetch_executor.submit(_sample_batch, raw_mega_batch) + features = _encode_batch(rows_list, emb_template) + t_encode = time.monotonic() - t0 - if not minibatches: + if len(features) < 4: continue - for mb in minibatches: - t1 = time.monotonic() - model.forward_backward(inputs=mb, task='embedding') - model.clip_grad_and_step(gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) - t_train = time.monotonic() - t1 - cur_step += 1 - - _ema_prefetch = _ema_alpha * t_prefetch + (1 - _ema_alpha) * _ema_prefetch if cur_step > RESUME_STEP + 1 else t_prefetch - _ema_train = _ema_alpha * t_train + (1 - _ema_alpha) * _ema_train if cur_step > RESUME_STEP + 1 else t_train - - if cur_step % LOG_INTERVAL == 0: - metric = model.calculate_metric(is_training=True) - _bottleneck = 'PREFETCH' if _ema_prefetch > _ema_train else 'TRAIN' - logger.info( - f'Epoch {epoch} Step {cur_step}/{total_forward_steps}, metric: {metric} | ' - f'prefetch={t_prefetch:.1f}s(ema {_ema_prefetch:.1f}) ' - f'train={t_train:.1f}s(ema {_ema_train:.1f}) ' - f'bottleneck={_bottleneck}') - log_dict = {} - for k, v in metric.items(): - if not v: - continue - try: - log_dict[k] = float(v) - except (ValueError, TypeError): - pass - log_dict['epoch'] = epoch - log_dict['prefetch_sec'] = round(t_prefetch, 2) - log_dict['train_sec'] = round(t_train, 2) - swanlab.log(log_dict, step=cur_step) - if cur_step % SAVE_INTERVAL == 0: - save_checkpoint(model, f'step_{cur_step}') - t_prefetch = 0.0 - - # Drain final mega-batch - if future: - minibatches = future.result() - future = None - if minibatches: - for mb in minibatches: - model.forward_backward(inputs=mb, task='embedding') - model.clip_grad_and_step(gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) - cur_step += 1 - if cur_step % SAVE_INTERVAL == 0: - save_checkpoint(model, f'step_{cur_step}') - - prefetch_executor.shutdown(wait=False) + t1 = time.monotonic() + model.forward_backward(inputs=features, task='embedding') + model.clip_grad_and_step( + gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) + t_train = time.monotonic() - t1 + cur_step += 1 + + if cur_step % LOG_INTERVAL == 0: + metric = model.calculate_metric(is_training=True) + logger.info( + f'Epoch {epoch} Step {cur_step}/{total_steps}, ' + f'metric: {metric} | ' + f'encode={t_encode:.2f}s train={t_train:.2f}s') + log_dict = {} + for k, v in metric.items(): + if not v: + continue + try: + log_dict[k] = float(v) + except (ValueError, TypeError): + pass + log_dict['epoch'] = epoch + log_dict['encode_sec'] = round(t_encode, 3) + log_dict['train_sec'] = round(t_train, 3) + swanlab.log(log_dict, step=cur_step) + if cur_step % SAVE_INTERVAL == 0: + save_checkpoint(model, f'step_{cur_step}') + save_checkpoint(model, 'last-checkpoint') + # Force sync: resolve any pending lazy remote calls (save) before exit + model.calculate_metric(is_training=True) + logger.info(f'Training complete. Final step: {cur_step}') if __name__ == '__main__': diff --git a/cookbook/exp/rl/grpo.py b/cookbook/exp/rl/grpo.py new file mode 100644 index 000000000..7f43d93f4 --- /dev/null +++ b/cookbook/exp/rl/grpo.py @@ -0,0 +1,787 @@ +"""Pure GRPO training on AoPS dataset (no RAG, ablation baseline). + +Architecture (8 GPUs): + - 4 GPUs: sampler/rollout (vLLM TP=4) + - 4 GPUs: training model (FSDP) + +Pipeline per step: + 1. DataLoader yields a batch of math problems + 2. Sampler generates rollouts + 3. Reward (accuracy + format + gibberish) → GRPO advantage → model update + +Launch: + python cookbook/exp/rl/grpo.py +""" +import json +import os +import re +import random +from typing import Any, Dict, List, Tuple + +import numpy as np +import torch + +import twinkle +from twinkle import DeviceMesh, DeviceGroup, get_device_placement, get_logger +from twinkle.advantage import GRPOAdvantage +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.data_format import SamplingParams +from twinkle.dataloader import DataLoader +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.metric import CompletionRewardMetric +from twinkle.model import TransformersModel +from twinkle.processor import InputProcessor +from twinkle.reward.base import Reward +from twinkle.sampler import vLLMSampler +from twinkle.template import Qwen3_5Template + +logger = get_logger() + +# ============================================================================ +# Configuration +# ============================================================================ +MODEL_ID = os.environ.get('MODEL_ID', 'ms://Qwen/Qwen3.5-4B') + +# GPU layout: 4 rollout + 4 train = 8 +SAMPLER_GPUS = int(os.environ.get('SAMPLER_GPUS', 4)) +MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 4)) +NUM_GPUS = SAMPLER_GPUS + MODEL_GPUS + +# Training hyperparams +NUM_GENERATIONS = int(os.environ.get('NUM_GENERATIONS', 8)) +MAX_NEW_TOKENS = int(os.environ.get('MAX_NEW_TOKENS', 32768)) +LEARNING_RATE = float(os.environ.get('LR', 1e-5)) +MAX_STEPS = int(os.environ.get('MAX_STEPS', 5000)) +BATCH_SIZE = int(os.environ.get('BATCH_SIZE', 8)) +MINI_BATCH_SIZE = int(os.environ.get('MINI_BATCH_SIZE', 8)) +MICRO_BATCH_SIZE = int(os.environ.get('MICRO_BATCH_SIZE', 8)) +GRADIENT_ACCUMULATION_STEPS = int(os.environ.get('GRADIENT_ACCUMULATION_STEPS', 1)) +SAVE_STEPS = int(os.environ.get('SAVE_STEPS', 100)) +ADV_CLIP = float(os.environ.get('ADV_CLIP', 1.0)) +LOSS_SPIKE_THRESHOLD = float(os.environ.get('LOSS_SPIKE_THRESHOLD', 10.0)) + +# Dataset +AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') +AOPS_SEED = int(os.environ.get('AOPS_SEED', 100)) + +# Output / diagnostics +OUTPUT_DIR = os.environ.get('OUTPUT_DIR', './outputs/grpo') + +# System prompt +SYSTEM_PROMPT = ( + 'You are an expert competition mathematician. ' + 'Solve the problem step by step. Put your final answer inside \\boxed{}. ' + 'For multiple-choice questions, put the option LETTER (A/B/C/D/E) inside \\boxed{}.' +) + + +# ============================================================================ +# Reward +# ============================================================================ +class AoPSAccuracyReward(Reward): + """Accuracy reward via boxed answer extraction + robust equivalence matching.""" + + @staticmethod + def extract_boxed(text: str) -> str: + idx = text.rfind('\\boxed{') + if idx == -1: + return '' + start = idx + len('\\boxed{') + depth = 1 + j = start + while j < len(text) and depth > 0: + if text[j] == '{': + depth += 1 + elif text[j] == '}': + depth -= 1 + j += 1 + if depth == 0: + return text[start:j - 1].strip() + return '' + + _MCQ_GT_RE = re.compile( + r'^\\?(?:textbf|mathbf|text|mathrm)\{?\(?([A-E])[)}\s\\]*(.*)', + re.DOTALL) + _MCQ_PAREN_RE = re.compile(r'^\(?([A-E])\)?[.:\s\\]+(.*)', re.DOTALL) + _MCQ_SINGLE_LETTER_RE = re.compile(r'^[A-E]$') + _VAR_PREFIX_RE = re.compile( + r'^(?:[a-zA-Z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)', re.DOTALL) + _EQ_RHS_RE = re.compile(r'^.+=\s*(.+)$') + + @staticmethod + def normalize_answer(ans: str) -> str: + if not ans: + return '' + s = ans.strip() + m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf)?\{?\(?([A-E])\)?\}?$', s) + if m: + return m.group(1) + s = s.replace(' ', '') + s = s.replace(r'\,', '') + s = s.replace(r'\;', '') + s = s.replace(r'\!', '') + s = re.sub(r'\\(?:text|mathrm|mathbf|textbf|operatorname)\{([^}]*)\}', r'\1', s) + s = s.replace(r'\displaystyle', '') + s = re.sub(r'\\(?:left|right)[.()\[\]|]', '', s) + s = s.replace(r'\dfrac', r'\frac') + s = s.replace(r'\tfrac', r'\frac') + # \frac shorthand without braces: \frac ab → \frac{a}{b} + s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) + s = s.strip('$').strip() + s = re.sub(r'\{[a-zA-Z]+\}$', '', s) + s = re.sub(r'\^\{\\circ\}|\^\\circ|°|\\circ', '', s) + s = re.sub(r'\\(?:quad|qquad|\s)', '', s) + s = s.replace(r'\minus{}', '-').replace(r'\minus', '-') + + def _frac_to_slash(m): + text = m.group(0) + pos = text.index('{') + 1 + depth, num_start = 1, pos + while depth > 0: + if text[pos] == '{': depth += 1 + elif text[pos] == '}': depth -= 1 + pos += 1 + numer = text[num_start:pos - 1] + pos += 1 + den_start = pos + depth = 1 + while depth > 0: + if text[pos] == '{': depth += 1 + elif text[pos] == '}': depth -= 1 + pos += 1 + denom = text[den_start:pos - 1] + return f'({numer})/({denom})' + + s = re.sub( + r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', + _frac_to_slash, s) + s = re.sub(r'(? str: + m = cls._VAR_PREFIX_RE.match(s) + return m.group(1).strip() if m else s + + @classmethod + def _extract_mcq_parts(cls, s: str): + m = cls._MCQ_GT_RE.match(s) + if m: + return m.group(1), m.group(2).strip() + m = cls._MCQ_PAREN_RE.match(s) + if m: + return m.group(1), m.group(2).strip() + m2 = re.search(r'\(?([A-E])\)?\s*$', s) + if m2 and len(s) > 3: + return m2.group(1), s[:m2.start()].strip() + return None, None + + @staticmethod + def _try_numeric_equal(a: str, b: str) -> bool: + import math + + def _try_eval(s: str): + try: + return float(s.replace('(', '').replace(')', '')) + except (ValueError, ZeroDivisionError): + pass + s_stripped = re.sub(r'[a-zA-Z]+$', '', s.replace('(', '').replace(')', '')).strip() + if s_stripped and s_stripped != s: + try: + return float(s_stripped) + except (ValueError, ZeroDivisionError): + pass + frac_re = re.compile(r'^\(([^)]+)\)/\(([^)]+)\)$') + m = frac_re.match(s) + if m: + try: + return float(m.group(1)) / float(m.group(2)) + except (ValueError, ZeroDivisionError): + pass + expr = s + expr = expr.replace('\\pi', str(math.pi)) + expr = expr.replace('\\e', str(math.e)) + expr = re.sub(r'\\sqrt\{([^}]+)\}', r'(\1)**0.5', expr) + expr = re.sub(r'\\sqrt\[3\]\{([^}]+)\}', r'(\1)**(1/3)', expr) + expr = re.sub(r'\\sqrt\[([^]]+)\]\{([^}]+)\}', r'(\2)**(1/\1)', expr) + expr = expr.replace('{', '(').replace('}', ')') + expr = expr.replace('\\cdot', '*').replace('\\times', '*') + expr = re.sub(r'(\d)\(', r'\1*(', expr) + try: + val = eval(expr, {"__builtins__": {}, "math": math, "pi": math.pi, "e": math.e}, {}) + return float(val) + except Exception: + pass + return None + + va, vb = _try_eval(a), _try_eval(b) + if va is not None and vb is not None: + return abs(va - vb) < 1e-6 * max(1, abs(va), abs(vb)) + return False + + @classmethod + def _strip_quantifiers(cls, s: str) -> str: + """Strip universal/existential quantifier wrappers.""" + s = re.sub(r'^\\forall\s*\w+\s*\\in\s*\\mathbb\s*\{?[A-Z]\}?\s*[:,]\s*', '', s) + s = re.sub(r'\s*\(\\forall[^)]*\)\s*$', '', s) + s = re.sub(r'\s*\(for\s+all[^)]*\)\s*$', '', s, flags=re.IGNORECASE) + return s.strip() + + @classmethod + def _try_param_rename(cls, a: str, b: str) -> bool: + """Check if a == b up to consistent single free-parameter rename (ax vs cx).""" + if not a or not b or len(a) != len(b) or len(a) > 80: + return False + diffs = [(i, a[i], b[i]) for i in range(len(a)) if a[i] != b[i]] + if not diffs: + return False + src_chars = set(d[1] for d in diffs) + dst_chars = set(d[2] for d in diffs) + if len(src_chars) == 1 and len(dst_chars) == 1: + src, dst = src_chars.pop(), dst_chars.pop() + if src.isalpha() and dst.isalpha(): + return a.replace(src, dst) == b + return False + + @classmethod + def _normalize_tuple(cls, s: str) -> str: + # Strip set-builder conditions: \mid ... or | ... + s = re.sub(r'\\mid.*$', '', s) + s = re.sub(r'\|[^,]*$', '', s) + return re.sub(r'[\s()\[\]{}\\]', '', s) + + @classmethod + def _try_sympy_equal(cls, a: str, b: str) -> bool: + try: + from sympy.parsing.latex import parse_latex + from sympy import simplify, nsimplify + expr_a = parse_latex(a) + expr_b = parse_latex(b) + diff = simplify(nsimplify(expr_a - expr_b)) + return diff == 0 + except Exception: + return False + + @classmethod + def answers_match(cls, predicted: str, reference: str) -> bool: + if not predicted or not reference: + return False + + norm_p = cls.normalize_answer(predicted) + norm_r = cls.normalize_answer(reference) + + if norm_p == norm_r: + return True + if norm_p.lower() == norm_r.lower(): + return True + if cls._try_numeric_equal(norm_p, norm_r): + return True + + stripped_p = cls.normalize_answer(cls._strip_var_prefix(predicted)) + stripped_r = cls.normalize_answer(cls._strip_var_prefix(reference)) + if stripped_p and stripped_r and stripped_p == stripped_r: + return True + if stripped_p and stripped_r and cls._try_numeric_equal(stripped_p, stripped_r): + return True + + ref_letter, ref_value = cls._extract_mcq_parts(reference) + if ref_letter: + if norm_p == ref_letter or predicted.strip().upper() == ref_letter: + return True + if ref_value: + norm_ref_val = cls.normalize_answer(ref_value) + if norm_p == norm_ref_val or cls._try_numeric_equal(norm_p, norm_ref_val): + return True + pred_letter, pred_value = cls._extract_mcq_parts(predicted) + if pred_letter: + if norm_r == pred_letter or reference.strip().upper() == pred_letter: + return True + if pred_value: + norm_pred_val = cls.normalize_answer(pred_value) + if norm_r == norm_pred_val or cls._try_numeric_equal(norm_r, norm_pred_val): + return True + if cls._MCQ_SINGLE_LETTER_RE.match(reference.strip()): + if cls._MCQ_SINGLE_LETTER_RE.match(predicted.strip().upper()): + return predicted.strip().upper() == reference.strip().upper() + + tuple_p = cls._normalize_tuple(norm_p) + tuple_r = cls._normalize_tuple(norm_r) + if ',' in tuple_p and tuple_p == tuple_r: + return True + if stripped_p and stripped_r: + tuple_sp = cls._normalize_tuple(stripped_p) + tuple_sr = cls._normalize_tuple(stripped_r) + if ',' in tuple_sp and tuple_sp == tuple_sr: + return True + + if '=' in norm_r and '=' not in norm_p: + parts = norm_r.split('=') + for part in parts: + part = part.strip() + if part == norm_p or cls._try_numeric_equal(part, norm_p): + return True + if '=' in norm_p and '=' not in norm_r: + parts = norm_p.split('=') + for part in parts: + part = part.strip() + if part == norm_r or cls._try_numeric_equal(part, norm_r): + return True + if '=' in norm_p and '=' in norm_r: + pp = [x.strip() for x in norm_p.split('=')] + rp = [x.strip() for x in norm_r.split('=')] + if set(pp) == set(rp): + return True + + def _sort_factors(s): + tokens = re.findall(r'\\?[a-zA-Z]+\{[^}]*\}|\\?[a-zA-Z]+|\d+|[^a-zA-Z\d\\{}]', s) + return ''.join(sorted(tokens)) + if _sort_factors(norm_p) == _sort_factors(norm_r): + return True + + if cls._try_sympy_equal(predicted, reference): + return True + + # --- Strategy 9b: quantifier stripping + param rename --- + q_stripped_r = cls._strip_quantifiers(reference) + q_stripped_p = cls._strip_quantifiers(predicted) + if q_stripped_r != reference or q_stripped_p != predicted: + norm_qr = cls.normalize_answer(cls._strip_var_prefix(q_stripped_r)) + norm_qp = cls.normalize_answer(cls._strip_var_prefix(q_stripped_p)) + if norm_qr and norm_qp: + if norm_qr == norm_qp: + return True + if cls._try_param_rename(norm_qr, norm_qp): + return True + + # --- Strategy 10: param rename on var-prefix-stripped forms --- + if stripped_p and stripped_r and cls._try_param_rename(stripped_p, stripped_r): + return True + + return False + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + rewards = [] + for traj in trajectories: + messages = traj.get('messages', []) + completion = '' + for msg in reversed(messages): + if msg.get('role') == 'assistant': + completion = msg.get('content', '') + break + user_data = traj.get('user_data') or [] + gt = '' + for item in user_data: + if item[0] == 'ground_truth': + gt = item[1] + break + predicted = self.extract_boxed(completion) + correct = self.answers_match(predicted, gt) + rewards.append(1.0 if correct else 0.0) + return rewards + + +class FormatReward(Reward): + """Reward for having \\boxed{} in the output.""" + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + rewards = [] + for traj in trajectories: + messages = traj.get('messages', []) + completion = '' + for msg in reversed(messages): + if msg.get('role') == 'assistant': + completion = msg.get('content', '') + break + has_boxed = '\\boxed{' in completion + rewards.append(0.5 if has_boxed else 0.0) + return rewards + + +class GibberishPenalty(Reward): + """Negative reward for degenerate outputs (gibberish/random unicode tail).""" + + TAIL_CHARS = 400 + GIBBERISH_THRESHOLD = 0.20 + + @classmethod + def is_gibberish(cls, text: str) -> bool: + if not text: + return False + tail = text[-cls.TAIL_CHARS:] if len(text) > cls.TAIL_CHARS else text + non_math_non_ascii = 0 + for c in tail: + code = ord(c) + if code > 127 and not (0x4e00 <= code <= 0x9fff): + non_math_non_ascii += 1 + return non_math_non_ascii > len(tail) * cls.GIBBERISH_THRESHOLD + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + rewards = [] + for traj in trajectories: + messages = traj.get('messages', []) + completion = '' + for msg in reversed(messages): + if msg.get('role') == 'assistant': + completion = msg.get('content', '') + break + rewards.append(-0.5 if self.is_gibberish(completion) else 0.0) + return rewards + + +def compute_rewards(trajectories: List[Dict[str, Any]] + ) -> Tuple[List[float], List[float], List[float]]: + acc_fn = AoPSAccuracyReward() + fmt_fn = FormatReward() + gib_fn = GibberishPenalty() + acc = acc_fn(trajectories) + fmt = fmt_fn(trajectories) + gib = gib_fn(trajectories) + total = [a + f + g for a, f, g in zip(acc, fmt, gib)] + return total, fmt, acc + + +# ============================================================================ +# Dataset: AoPS boxed problems +# ============================================================================ +def create_aops_dataset(): + """Load AoPS and create GRPO-style dataset (prompt only, with ground_truth in user_data).""" + from modelscope import MsDataset + from twinkle.data_format import Message, Trajectory + + ds = MsDataset.load(AOPS_DATASET_ID, split='train', + download_mode='reuse_dataset_if_exists') + rows = [] + for row in ds: + if not row['metadata'].get('boxed'): + continue + ref = AoPSAccuracyReward.extract_boxed(row['solution']) + if not ref: + continue + rows.append({'problem': row['problem'], 'ground_truth': ref}) + + logger.info(f'[aops] loaded {len(rows)} boxed problems') + rng = random.Random(AOPS_SEED) + rng.shuffle(rows) + + trajectories = [] + for r in rows: + traj = Trajectory( + messages=[ + Message(role='system', content=SYSTEM_PROMPT), + Message(role='user', content=r['problem']), + ], + user_data=[('ground_truth', r['ground_truth'])], + ) + trajectories.append(traj) + + data_meta = DatasetMeta(data=trajectories) + dataset = Dataset(data_meta) + dataset.set_template('Qwen3_5Template', model_id=MODEL_ID, + max_length=16384, truncation_strategy='delete', + enable_thinking=True) + dataset.encode(add_generation_prompt=True) + return dataset + + +# ============================================================================ +# Main +# ============================================================================ +def main(): + sampler_start = 0 + model_start = sampler_start + SAMPLER_GPUS + + device_groups = [ + DeviceGroup(name='sampler', ranks=list(range(sampler_start, model_start)), + device_type='GPU'), + DeviceGroup(name='model', ranks=list(range(model_start, NUM_GPUS)), + device_type='GPU'), + ] + + model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, fsdp_size=MODEL_GPUS, ulysses_size=2) + sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, dp_size=SAMPLER_GPUS) + + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, + groups=device_groups, lazy_collect=False) + + # -- Training model (full-parameter) -- + model = TransformersModel( + model_id=MODEL_ID, device_mesh=model_mesh, remote_group='model') + model.set_optimizer('AdamW', lr=LEARNING_RATE) + model.set_lr_scheduler('CosineAnnealingLR', T_max=MAX_STEPS, eta_min=0) + model.set_loss('GSPOLoss', epsilon=0.2, epsilon_high=0.28, beta=0.04) + model.set_processor(InputProcessor) + model.set_template('Qwen3_5Template', model_id=MODEL_ID, + enable_thinking=True, max_length=32768) + + # -- Rollout sampler -- + sampler = vLLMSampler( + model_id=MODEL_ID, + engine_args={ + 'gpu_memory_utilization': 0.8, + 'max_model_len': 32768, + }, + device_mesh=sampler_mesh, + remote_group='sampler', + ) + sampler.set_template('Qwen3_5Template', model_id=MODEL_ID, + enable_thinking=True, max_length=32768) + + # -- Checkpoint & DataLoader -- + ckpt_manager = CheckpointEngineManager(model=model, sampler=sampler) + + GLOBAL_BATCH_SIZE = BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS + dataloader = DataLoader( + dataset=create_aops_dataset, + batch_size=GLOBAL_BATCH_SIZE, + min_batch_size=GLOBAL_BATCH_SIZE, + device_mesh=model_mesh, + remote_group='model', + ) + + advantage_fn = GRPOAdvantage() + metrics = CompletionRewardMetric() + sampling_params = SamplingParams( + max_tokens=MAX_NEW_TOKENS, num_samples=1, logprobs=1, + temperature=1.0, top_p=0.95) + + optim_step = 0 + logger.info('Starting pure GRPO training (no RAG)') + logger.info(get_device_placement()) + + # -- Diagnostics -- + os.makedirs(OUTPUT_DIR, exist_ok=True) + diag_path = os.path.join(OUTPUT_DIR, 'diagnostics.jsonl') + diag_f = open(diag_path, 'w', encoding='utf-8') + logger.info(f'[diag] diagnostics → {diag_path}') + + def _content_to_str(content): + if isinstance(content, str): + return content + if isinstance(content, list): + return ''.join( + b.get('text', '') if isinstance(b, dict) else str(b) + for b in content) + return str(content) + + for batch in dataloader: + if optim_step >= MAX_STEPS: + break + + metrics.reset() + + # Build prompts (direct, no RAG) + prompts = [] + for item in batch: + msgs = item.get('messages', []) + prob = '' + for m in msgs: + if m.get('role') == 'user': + prob = m.get('content', '') + if isinstance(prob, list): + prob = ''.join(p.get('text', '') for p in prob if isinstance(p, dict)) + break + ud = item.get('user_data', []) + gt = '' + for pair in ud: + if pair[0] == 'ground_truth': + gt = pair[1] + break + prompts.append({ + 'messages': [ + {'role': 'system', 'content': SYSTEM_PROMPT}, + {'role': 'user', 'content': prob}, + ], + 'user_data': [('ground_truth', gt)], + }) + + # Expand for NUM_GENERATIONS and sample + expand_prompts = [] + for prompt in prompts: + expand_prompts.extend([prompt] * NUM_GENERATIONS) + + ckpt_manager.sync_weights(merge_and_sync=False) + sampler.reset_prefix_cache() + + sample_responses = sampler.sample(expand_prompts, sampling_params) + + # Collect rollouts + all_input_data: List[Dict[str, Any]] = [] + all_old_logps: List[List[float]] = [] + all_completion_lengths: List[int] = [] + + for sample_response in sample_responses: + for sequence in sample_response.sequences: + all_input_data.append(sequence.new_input_feature) + all_old_logps.append([logprob[0][1] for logprob in sequence.logprobs]) + all_completion_lengths.append(len(sequence.tokens)) + + # Rewards + total_rewards, format_rewards, accuracy_rewards = compute_rewards(all_input_data) + + # Zero out rewards for rollouts that hit the max_tokens ceiling + max_len_threshold = int(MAX_NEW_TOKENS * 0.95) + for i in range(len(all_input_data)): + if all_completion_lengths[i] >= max_len_threshold: + total_rewards[i] = 0.0 + accuracy_rewards[i] = 0.0 + format_rewards[i] = 0.0 + + # Per-step reward summary + n_correct = sum(1 for a in accuracy_rewards if a > 0) + diag_f.write(json.dumps({ + 'step': optim_step, 'type': 'reward_summary', + 'n_samples': len(accuracy_rewards), + 'accuracy': n_correct / len(accuracy_rewards) if accuracy_rewards else 0, + 'mean_reward': sum(total_rewards) / len(total_rewards) if total_rewards else 0, + }, ensure_ascii=False) + '\n') + + metrics.accumulate( + completion_lengths=all_completion_lengths, + rewards={ + 'total': total_rewards, + 'format': format_rewards, + 'accuracy': accuracy_rewards, + }, + ) + + # GRPO advantage + advantages = advantage_fn( + total_rewards, num_generations=NUM_GENERATIONS, scale='group').tolist() + if ADV_CLIP > 0: + advantages = [max(-ADV_CLIP, min(ADV_CLIP, a)) for a in advantages] + + # Log rollout responses + _extract_boxed = AoPSAccuracyReward.extract_boxed + for ridx, traj in enumerate(all_input_data): + msgs = traj.get('messages', []) + assistant_text = _content_to_str(next( + (m['content'] for m in reversed(msgs) if m.get('role') == 'assistant'), '')) + user_text = _content_to_str(next( + (m['content'] for m in msgs if m.get('role') == 'user'), '')) + user_data = traj.get('user_data') or [] + gt = next((v for k, v in user_data if k == 'ground_truth'), '') + problem_idx = ridx // NUM_GENERATIONS + grp_start = problem_idx * NUM_GENERATIONS + grp_end = grp_start + NUM_GENERATIONS + grp_acc = sum(accuracy_rewards[grp_start:grp_end]) / NUM_GENERATIONS + + diag_f.write(json.dumps({ + 'step': optim_step, 'type': 'rollout', + 'idx': ridx, + 'problem_idx': problem_idx, + 'problem': user_text, + 'response': assistant_text, + 'ground_truth': gt, + 'predicted': _extract_boxed(assistant_text), + 'reward': total_rewards[ridx], + 'accuracy_reward': accuracy_rewards[ridx], + 'format_reward': format_rewards[ridx], + 'advantage': advantages[ridx], + 'completion_length': all_completion_lengths[ridx], + 'group_accuracy': grp_acc, + }, ensure_ascii=False) + '\n') + + diag_f.flush() + + # Filter out low-signal problem groups (DAPO-style dynamic sampling) + # Skip groups where accuracy is too low (<0.1) or too high (>0.9) + # to avoid gradient dominated by gibberish/format noise or no learning signal. + filtered_inputs, filtered_old_logps, filtered_advantages = [], [], [] + for g in range(BATCH_SIZE): + g_start = g * NUM_GENERATIONS + g_end = g_start + NUM_GENERATIONS + grp_adv = advantages[g_start:g_end] + if all(abs(a) < 1e-8 for a in grp_adv): + continue + grp_acc_rate = sum(accuracy_rewards[g_start:g_end]) / NUM_GENERATIONS + if grp_acc_rate < 0.2 or grp_acc_rate > 0.8: + continue + filtered_inputs.extend(all_input_data[g_start:g_end]) + filtered_old_logps.extend(all_old_logps[g_start:g_end]) + filtered_advantages.extend(grp_adv) + + # Mini-batch training with gradient accumulation + # Process MICRO_BATCH_SIZE samples per forward, accumulate grad_accum_steps + # times before one optimizer step. clip_grad_norm normalizes by accumulated + # num_tokens, ensuring mathematical equivalence with larger batch forward. + total_completions = len(filtered_inputs) + if total_completions == 0: + logger.info(f'[Step {optim_step}] all groups filtered (uniform rewards), skip training') + continue + + grad_accum_steps = MINI_BATCH_SIZE // MICRO_BATCH_SIZE + accum_count = 0 + for mb_start in range(0, total_completions, MICRO_BATCH_SIZE): + mb_end = min(mb_start + MICRO_BATCH_SIZE, total_completions) + mb_inputs = filtered_inputs[mb_start:mb_end] + mb_old_logps = filtered_old_logps[mb_start:mb_end] + mb_advantages = filtered_advantages[mb_start:mb_end] + + outputs = model.forward_backward( + inputs=mb_inputs, + old_logps=mb_old_logps, + ref_logps=mb_old_logps, + advantages=mb_advantages, + ) + accum_count += 1 + + if accum_count % grad_accum_steps == 0: + skip_step = False + try: + loss_val = outputs.get('loss', None) + if loss_val is not None: + if hasattr(loss_val, 'item'): + loss_val = loss_val.item() + if loss_val > LOSS_SPIKE_THRESHOLD: + skip_step = True + logger.warning( + f'[Step {optim_step}] Loss spike: {loss_val:.4f} > ' + f'{LOSS_SPIKE_THRESHOLD}, skipping update') + except Exception: + pass + + if skip_step: + model.zero_grad() + else: + model.clip_grad_and_step() + optim_step += 1 + + if optim_step >= MAX_STEPS: + break + if optim_step % SAVE_STEPS == 0: + model.save(f'grpo-checkpoint-{optim_step}') + + # Flush remaining accumulated gradients (incomplete window at tail) + if accum_count % grad_accum_steps != 0: + skip_step = False + try: + loss_val = outputs.get('loss', None) + if loss_val is not None: + if hasattr(loss_val, 'item'): + loss_val = loss_val.item() + if loss_val > LOSS_SPIKE_THRESHOLD: + skip_step = True + logger.warning( + f'[Step {optim_step}] Loss spike (tail): {loss_val:.4f} > ' + f'{LOSS_SPIKE_THRESHOLD}, skipping update') + except Exception: + pass + + if skip_step: + model.zero_grad() + else: + model.clip_grad_and_step() + optim_step += 1 + + log_dict = metrics.calculate() + log_dict.update(model.calculate_metric(is_training=True)) + metrics.reset() + logger.info(f'[Step {optim_step}/{MAX_STEPS}] {log_dict}') + + diag_f.close() + logger.info(f'Training completed. optim_steps={optim_step}') + model.save('grpo-final') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/rl/rag_hint_grpo.py b/cookbook/exp/rl/rag_hint_grpo.py new file mode 100644 index 000000000..b35b3706b --- /dev/null +++ b/cookbook/exp/rl/rag_hint_grpo.py @@ -0,0 +1,1480 @@ +"""RAG-hint GRPO training: retrieve thinking traces and condense as hints for RL. + +Architecture (8 GPUs): + - 1 GPU: condenser (vLLM, compress retrieved traces) + - 1 GPU: embedding model (encode queries for retrieval) + - 4 GPUs: sampler/rollout (vLLM TP=4) + - 2 GPUs: training model (FSDP/DP) + +Pipeline per step: + 1. DataLoader yields a batch of math problems + 2. [Async] Embedding model encodes problems → retrieve from LanceDB → condenser compresses + 3. Build RAG-hint prompts (one-shot in system, with analysis prefix) + 4. Sampler generates rollouts (response starts with forced analysis prefix) + 5. Reward (accuracy + format) → GRPO advantage → model update + +Launch: + python cookbook/exp/rl/rag_hint_grpo.py +""" +import json +import os +import re +import random +import threading +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np +import torch + +import twinkle +from twinkle import DeviceMesh, DeviceGroup, get_device_placement, get_logger +from twinkle.advantage import GRPOAdvantage +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.data_format import SamplingParams +from twinkle.dataloader import DataLoader +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.loss import InfonceLoss +from twinkle.metric import CompletionRewardMetric +from twinkle.model import TransformersModel +from twinkle.processor import InputProcessor +from twinkle.reward.base import Reward +from twinkle.sampler import vLLMSampler +from twinkle.template import Qwen3_5Template +from twinkle_agentic.protocol.openai import OpenAI as OpenAIClient + +logger = get_logger() + +# ============================================================================ +# Configuration +# ============================================================================ +MODEL_ID = os.environ.get('MODEL_ID', 'ms://Qwen/Qwen3.5-4B') + +# GPU layout: 1 condenser + 1 embedding + 4 rollout + 2 train = 8 +CONDENSER_GPUS = int(os.environ.get('CONDENSER_GPUS', 1)) +EMB_GPUS = int(os.environ.get('EMB_GPUS', 1)) +SAMPLER_GPUS = int(os.environ.get('SAMPLER_GPUS', 2)) +MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 4)) +NUM_GPUS = CONDENSER_GPUS + EMB_GPUS + SAMPLER_GPUS + MODEL_GPUS + +# Training hyperparams +NUM_GENERATIONS = int(os.environ.get('NUM_GENERATIONS', 8)) +MAX_NEW_TOKENS = int(os.environ.get('MAX_NEW_TOKENS', 32768)) +LEARNING_RATE = float(os.environ.get('LR', 1e-5)) +MAX_STEPS = int(os.environ.get('MAX_STEPS', 5000)) +BATCH_SIZE = int(os.environ.get('BATCH_SIZE', 8)) +MINI_BATCH_SIZE = int(os.environ.get('MINI_BATCH_SIZE', 8)) +MICRO_BATCH_SIZE = int(os.environ.get('MICRO_BATCH_SIZE', 8)) +GRADIENT_ACCUMULATION_STEPS = int(os.environ.get('GRADIENT_ACCUMULATION_STEPS', 1)) +SAVE_STEPS = int(os.environ.get('SAVE_STEPS', 100)) +ADV_CLIP = float(os.environ.get('ADV_CLIP', 1.0)) +LOSS_SPIKE_THRESHOLD = float(os.environ.get('LOSS_SPIKE_THRESHOLD', 10.0)) + +# RAG config +DB_PATH = os.environ.get('DB_PATH', './output.oldemb/thinking_rag/lance.db') +DB_TABLE = os.environ.get('DB_TABLE', 'thinking_traces') +TOP_K = int(os.environ.get('TOP_K', 2)) +SIM_THRESHOLD = float(os.environ.get('SIM_THRESHOLD', 0.75)) +MAX_TRACE_LEN = int(os.environ.get('MAX_TRACE_LEN', 8192)) +EMBED_MODEL_ID = os.environ.get( + 'EMBED_MODEL_ID', 'output.oldemb/embedding_full_transformers/last-checkpoint') +EMBED_MAX_LENGTH = int(os.environ.get('EMBED_MAX_LENGTH', 32000)) + +# Condenser config +CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') +CONDENSE_API_KEY = os.environ.get('COMPRESS_API_KEY', '') +CONDENSE_BASE_URL = os.environ.get('COMPRESS_BASE_URL', 'https://dashscope.aliyuncs.com/compatible-mode/v1') +CONDENSE_API_MODEL = os.environ.get('COMPRESS_MODEL', 'qwen3.7-max') +CONDENSE_API_CONCURRENCY = int(os.environ.get('API_CONCURRENCY', 16)) +CONDENSE_TEMPERATURE = 0.2 +CONDENSE_MAX_TOKENS = 8192 + +# Dataset +AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') +AOPS_SEED = int(os.environ.get('AOPS_SEED', 100)) + +# Decontamination & RAG fallback +DECONTAM_THRESHOLD = float(os.environ.get('DECONTAM_THRESHOLD', 0.20)) +RAG_FALLBACK_SIM = float(os.environ.get('RAG_FALLBACK_SIM', 0.60)) + +# Output / diagnostics +OUTPUT_DIR = os.environ.get('OUTPUT_DIR', './outputs/rag_hint_grpo') + +# Forced analysis prefix appended at the start of assistant response +ANALYSIS_PREFIX = '' + +# Fixed opening inside block — model must produce this EXACT prefix +HINT_REQUIRED_PREFIX = "Let's analyze the RAG example step by step." + +# Hint analysis config (API pre-analysis) +HINT_ANALYSIS_MAX_TOKENS = int(os.environ.get('HINT_ANALYSIS_MAX_TOKENS', 400)) +HINT_ANALYSIS_TEMPERATURE = 0.3 + +# ============================================================================ +# Condenser prompt (strategy-level extraction) +# ============================================================================ +COMPRESS_SYSTEM = """\ +You are a reasoning-trace condenser. Given a verbose reasoning trace, \ +extract the TRANSFERABLE KNOWLEDGE as an EXECUTABLE SOLUTION SKELETON \ +that would help a reader solve SIMILAR problems in the same domain. + +The reader will apply this knowledge to a DIFFERENT problem, so focus on what transfers. \ +NEVER output the final answer or conclusion of the original problem. \ +NEVER include problem-specific numeric results. + +Principles: +1. OUTPUT AN EXECUTABLE STEP CHAIN: numbered steps that a solver can directly follow. \ +Each step should state WHAT/WHY/HOW (with the formula/technique), not just name the concept. +2. INCLUDE FULL FORMULAS: theorems, identities — state each with COMPLETE MATHEMATICAL EXPRESSION. +3. STATE APPLICABILITY: what structural features signal that this approach works. +4. PRESERVE KEY INSIGHTS: non-obvious ideas that make the approach work. +5. REMOVE: problem-specific numeric calculations, final answers, dead-end explorations, hesitations. +6. FORMAT: Start with "Applicability:" one-line, then numbered steps. Keep concise. +7. NO meta-commentary. NO preamble. NO final answer. +""" + +COMPRESS_USER = ( + '## Reader Problem (context only — do NOT solve it)\n{query}\n\n' + '## Reasoning Trace to Condense\n{text}') + +# ============================================================================ +# RAG system prompt template (few-shot in system) +# ============================================================================ +SYSTEM_WITH_RAG_HEADER = ( + 'You are an expert competition mathematician. ' + 'Below are condensed reasoning examples from similar problems.\n\n' + '## Output Format (STRICT)\n' + 'Your response MUST begin with a block as the VERY FIRST content. ' + 'Do NOT output any text before .\n\n' + 'The block MUST start with EXACTLY this sentence (copy verbatim):\n' + '"Let\'s analyze the RAG example step by step."\n\n' + 'Then continue your analysis:\n' + '- Walk through each example\'s methodology and identify which steps, ' + 'formulas, and concepts are CORRECT and APPLICABLE to the current problem.\n' + '- Identify which parts are WRONG, IRRELEVANT, MISLEADING, or based on ' + 'assumptions that do NOT hold for this problem.\n' + '- End with a one-line verdict: "Useful: ..." and "Discard: ..."\n\n' + 'Example format:\n' + '\n' + "Let's analyze the RAG example step by step.\n" + '- Example 1: The ansatz f(x)=x^n is APPLICABLE because ... However, ' + 'the uniqueness argument via continuity is UNNECESSARY for this problem.\n' + '- Useful: power function ansatz, linear combination check.\n' + '- Discard: continuity assumption, specific numeric result.\n' + '\n\n' + 'After the block, solve the actual problem step by step using ONLY ' + 'the validated useful parts. Put your final answer inside \\boxed{}. ' + 'For multiple-choice questions, put the option LETTER (A/B/C/D/E) inside \\boxed{}.\n\n' +) + +# System prompt for pre-analyzed RAG (hint analysis done by API, model just solves) +_PREANALYSIS_BEFORE = ( + 'You are an expert competition mathematician.\n\n' + '## RAG Analysis (pre-computed)\n' +) +_PREANALYSIS_AFTER = ( + '\n\n## Instructions\n' + 'Use the useful methods/formulas identified above to solve the problem. ' + 'Ignore anything marked as irrelevant. ' + 'Solve step by step and put your final answer inside \\boxed{}. ' + 'For multiple-choice questions, put the option LETTER (A/B/C/D/E) inside \\boxed{}.' +) + + +def build_preanalysis_system(hint_analysis: str) -> str: + """Build system prompt with pre-analyzed hint. Uses concatenation to avoid .format() issues with math braces.""" + return _PREANALYSIS_BEFORE + hint_analysis + _PREANALYSIS_AFTER + +# API prompt for hint analysis generation +HINT_ANALYSIS_SYSTEM = ( + 'You are a mathematical methodology analyst. ' + 'Given a target problem and a condensed reasoning trace from a SIMILAR (but different) problem, ' + 'analyze which methods, formulas, and techniques from the trace are APPLICABLE to the target problem ' + 'and which are IRRELEVANT or MISLEADING.\n\n' + 'Output format (strict):\n' + '- Useful: [list specific methods/formulas/techniques that transfer to the target]\n' + '- Discard: [list parts that are irrelevant or would mislead]\n' + '- Key insight: [one sentence on how to apply the useful parts]\n\n' + 'Rules:\n' + '1. Be concise — at most 200 words total.\n' + '2. Focus ONLY on transferable methodology, never solve the target problem.\n' + '3. Never output the answer to either problem.\n' + '4. If the trace is entirely irrelevant, say "Useful: None. Discard: All."' +) + +HINT_ANALYSIS_USER = ( + '## Target Problem\n{query}\n\n' + '## Condensed Trace (from similar problem)\n{thinking}' +) + +EXAMPLE_TEMPLATE = ( + '--- Example {idx} ---\n' + 'Problem: {example_query}\n' + 'Methodology:\n{example_thinking}\n' + '--- End Example {idx} ---\n' +) + +SYSTEM_DIRECT = ( + 'You are an expert competition mathematician. ' + 'Solve the problem step by step. Put your final answer inside \\boxed{}. ' + 'For multiple-choice questions, put the option LETTER (A/B/C/D/E) inside \\boxed{}.' +) + + +# ============================================================================ +# Condenser utilities +# ============================================================================ +_api_semaphore = threading.Semaphore(CONDENSE_API_CONCURRENCY) + + +def _api_condense_single(api_client: OpenAIClient, messages: List[Dict]) -> Optional[str]: + _api_semaphore.acquire() + try: + trajectory = {'messages': messages} + sp = SamplingParams(temperature=CONDENSE_TEMPERATURE, max_tokens=CONDENSE_MAX_TOKENS) + reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) + content = (reply.get('content') or '').strip() + if not content: + return None + return content + except Exception as exc: + logger.warning(f'[condense-api] error: {exc}') + return None + finally: + _api_semaphore.release() + + +def _api_hint_analysis_batch( + api_client: OpenAIClient, + problems: List[str], + condensed_examples: List[List[Dict[str, str]]], +) -> List[Optional[str]]: + """Call API to pre-analyze RAG relevance for each problem. ~300 tokens per call.""" + results: List[Optional[str]] = [None] * len(problems) + tasks = [] # (idx, messages) + for i, prob in enumerate(problems): + if not condensed_examples[i]: + continue + # Merge all condensed traces into one block + traces = [] + for ex in condensed_examples[i]: + traces.append(ex.get('thinking', '')) + merged_thinking = '\n---\n'.join(traces) + user_msg = HINT_ANALYSIS_USER.format(query=prob, thinking=merged_thinking) + msgs = [ + {'role': 'system', 'content': HINT_ANALYSIS_SYSTEM}, + {'role': 'user', 'content': user_msg}, + ] + tasks.append((i, msgs)) + + if not tasks: + return results + + def _call_one(idx, msgs): + _api_semaphore.acquire() + try: + trajectory = {'messages': msgs} + sp = SamplingParams( + temperature=HINT_ANALYSIS_TEMPERATURE, + max_tokens=HINT_ANALYSIS_MAX_TOKENS) + reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) + content = (reply.get('content') or '').strip() + return idx, content if content else None + except Exception as exc: + logger.warning(f'[hint-analysis] error for idx={idx}: {exc}') + return idx, None + finally: + _api_semaphore.release() + + with ThreadPoolExecutor(max_workers=min(len(tasks), CONDENSE_API_CONCURRENCY)) as pool: + futs = [pool.submit(_call_one, idx, msgs) for idx, msgs in tasks] + for fut in as_completed(futs): + idx, analysis = fut.result() + results[idx] = analysis + + n_success = sum(1 for r in results if r) + logger.info(f'[hint-analysis] completed {n_success}/{len(tasks)} analyses') + return results + + +# ============================================================================ +# Embedding & Retrieval +# ============================================================================ +def _normalize_for_ngram(text: str) -> str: + """Normalize text for n-gram comparison: strip LaTeX markup, lowercase.""" + text = text.lower() + text = re.sub(r'\$+', '', text) + text = re.sub(r'\\[a-z]+\{([^}]*)\}', r'\1', text) + text = re.sub(r'\\[a-z]+', ' ', text) + text = re.sub(r'[{}\\^_$]', '', text) + text = re.sub(r'\s+', ' ', text).strip() + return text + + +def _ngram_jaccard(text_a: str, text_b: str, n: int = 13) -> float: + """13-gram character-level Jaccard similarity for decontamination.""" + a = _normalize_for_ngram(text_a) + b = _normalize_for_ngram(text_b) + if len(a) < n or len(b) < n: + return 0.0 + grams_a = set(a[i:i + n] for i in range(len(a) - n + 1)) + grams_b = set(b[i:i + n] for i in range(len(b) - n + 1)) + if not grams_a or not grams_b: + return 0.0 + return len(grams_a & grams_b) / len(grams_a | grams_b) + + +def _wrap_anchor(text: str) -> List[Dict[str, str]]: + return [ + {'role': 'user', 'content': text}, + {'role': 'assistant', 'content': 'Match the correct response here.'}, + ] + + +_DECONTAM_JUDGE_PROMPT = ( + 'We are building a RAG-augmented math training system. Problem A is the test ' + 'question; Problem B was retrieved from a knowledge base.\n' + 'Answer YES only if A and B are essentially the SAME specific problem — ' + 'i.e. solving B directly gives you A\'s answer (just different wording/notation/' + 'format/negation).\n' + 'Answer NO if they merely share the same method/topic but have different ' + 'specific values, equations, or geometric configurations — learning B\'s ' + 'approach still requires independent work to solve A.\n' + 'Problem A: {prob_a}\n' + 'Problem B: {prob_b}\n' + 'Answer only YES or NO.' +) + + +def _llm_judge_same_problem( + api_client, pairs: List[tuple], +) -> List[bool]: + """Batch LLM judge: are (problem_a, problem_b) the same problem? + + Each pair text is truncated to 200 chars to keep latency low. + Returns list of bools (True = same problem = should filter). + """ + if not pairs or not api_client: + return [False] * len(pairs) + + results = [False] * len(pairs) + + def _judge_one(idx, pa, pb): + prompt = _DECONTAM_JUDGE_PROMPT.format(prob_a=pa, prob_b=pb) + msgs = [{'role': 'user', 'content': prompt}] + try: + trajectory = {'messages': msgs} + sp = SamplingParams(temperature=0.1, max_tokens=8) + reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) + answer = (reply.get('content') or '').strip().upper() + return idx, 'YES' in answer + except Exception: + return idx, False + + with ThreadPoolExecutor(max_workers=min(len(pairs), CONDENSE_API_CONCURRENCY)) as pool: + futs = [pool.submit(_judge_one, i, pa, pb) for i, (pa, pb) in enumerate(pairs)] + for fut in as_completed(futs): + idx, is_same = fut.result() + results[idx] = is_same + return results + + +def get_embeddings(model: TransformersModel, template: Qwen3_5Template, + texts: List[str], dp_size: int) -> np.ndarray: + if not texts: + return np.zeros((0,), dtype=np.float32) + n = len(texts) + pad_n = (-n) % dp_size + padded = list(texts) + [' '] * pad_n if pad_n else list(texts) + features = [] + for t in padded: + feat = template.encode({'messages': _wrap_anchor(t or ' ')}) + feat['labels'] = [1] + features.append(feat) + out = model.forward_only(inputs=features, task='embedding', return_logits=True) + emb = out['embeddings'] + if isinstance(emb, torch.Tensor): + emb = emb.detach().to(torch.float32).cpu().numpy() + emb = np.asarray(emb, dtype=np.float32) + return emb[:n] if pad_n else emb + + +def retrieve_topk(tbl, query_vecs: np.ndarray, problems: List[str], + sim_threshold: float + ) -> List[List[Dict[str, Any]]]: + """Retrieve top-K thinking_raw per query with decontamination and length filter. + + Returns per-query list of dicts with keys: query, thinking, sim. + """ + results = [] + decontam_skipped = 0 + for qi, vec in enumerate(query_vecs): + hits = ( + tbl.search(vec.astype(np.float32).tolist()) + .metric('dot') + .limit(TOP_K + 50) + .select(['query_raw', 'thinking_raw', '_distance']) + .to_list() + ) + matched = [] + problem_text = problems[qi] if problems else '' + for h in hits: + if len(matched) >= TOP_K: + break + sim = 1.0 - h.get('_distance', 0.0) + if sim < sim_threshold: + continue + q = h.get('query_raw', '') + t = h.get('thinking_raw', '') + if not t: + continue + # Decontamination: skip if retrieved problem is too similar to current + if DECONTAM_THRESHOLD > 0 and problem_text and q: + if _ngram_jaccard(problem_text, q) > DECONTAM_THRESHOLD: + decontam_skipped += 1 + continue + # Drop traces exceeding max length (don't truncate — they'll be condensed poorly) + if len(t) > MAX_TRACE_LEN * 4: + continue + matched.append({'query': q, 'thinking': t, 'sim': sim}) + results.append(matched) + if decontam_skipped > 0: + logger.info(f'[decontam] skipped {decontam_skipped} leaked retrievals') + return results + + +# ============================================================================ +# Reward +# ============================================================================ +class AoPSAccuracyReward(Reward): + """Accuracy reward via boxed answer extraction + robust equivalence matching.""" + + @staticmethod + def extract_boxed(text: str) -> str: + idx = text.rfind('\\boxed{') + if idx == -1: + return '' + start = idx + len('\\boxed{') + depth = 1 + j = start + while j < len(text) and depth > 0: + if text[j] == '{': + depth += 1 + elif text[j] == '}': + depth -= 1 + j += 1 + if depth == 0: + return text[start:j - 1].strip() + return '' + + # --- MCQ letter regex (matches \textbf{(C) }value, (C) value, etc.) --- + _MCQ_GT_RE = re.compile( + r'^\\?(?:textbf|mathbf|text|mathrm)\{?\(?([A-E])[)}\s\\]*(.*)', + re.DOTALL) + _MCQ_PAREN_RE = re.compile(r'^\(?([A-E])\)?[.:\s\\]+(.*)', re.DOTALL) + _MCQ_SINGLE_LETTER_RE = re.compile(r'^[A-E]$') + # variable prefix: f(x)=..., m=..., N=..., P(n+1)=..., (x,y)=... + _VAR_PREFIX_RE = re.compile( + r'^(?:[a-zA-Z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)', re.DOTALL) + # GT with derivation: "18×1+999×2=2016" → extract RHS + _EQ_RHS_RE = re.compile(r'^.+=\s*(.+)$') + + @staticmethod + def normalize_answer(ans: str) -> str: + if not ans: + return '' + s = ans.strip() + # Pure MCQ letter + m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf)?\{?\(?([A-E])\)?\}?$', s) + if m: + return m.group(1) + s = s.replace(' ', '') + s = s.replace(r'\,', '') + s = s.replace(r'\;', '') + s = s.replace(r'\!', '') + # Remove text-mode wrappers but keep content (unwrap braces) + s = re.sub(r'\\(?:text|mathrm|mathbf|textbf|operatorname)\{([^}]*)\}', r'\1', s) + s = s.replace(r'\displaystyle', '') + s = re.sub(r'\\(?:left|right)[.()\[\]|]', '', s) + s = s.replace(r'\dfrac', r'\frac') + s = s.replace(r'\tfrac', r'\frac') + # \frac shorthand without braces: \frac ab → \frac{a}{b} + s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) + s = s.strip('$').strip() + # Remove trailing unit braces: {cm}, {kg}, etc. + s = re.sub(r'\{[a-zA-Z]+\}$', '', s) + # Degree normalization — strip entirely (degrees are contextual) + s = re.sub(r'\^\{\\circ\}|\^\\circ|°|\\circ', '', s) + # Remove \quad, \qquad, \ etc spacing + s = re.sub(r'\\(?:quad|qquad|\s)', '', s) + # Normalize minus: \minus{} → - + s = s.replace(r'\minus{}', '-').replace(r'\minus', '-') + + def _frac_to_slash(m): + text = m.group(0) + pos = text.index('{') + 1 + depth, num_start = 1, pos + while depth > 0: + if text[pos] == '{': depth += 1 + elif text[pos] == '}': depth -= 1 + pos += 1 + numer = text[num_start:pos - 1] + pos += 1 + den_start = pos + depth = 1 + while depth > 0: + if text[pos] == '{': depth += 1 + elif text[pos] == '}': depth -= 1 + pos += 1 + denom = text[den_start:pos - 1] + return f'({numer})/({denom})' + + s = re.sub( + r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', + _frac_to_slash, s) + s = re.sub(r'(? str: + """Strip variable assignment prefix: 'f(x)=x+1' → 'x+1', 'N=1006' → '1006'.""" + m = cls._VAR_PREFIX_RE.match(s) + return m.group(1).strip() if m else s + + @classmethod + def _extract_mcq_parts(cls, s: str): + """Extract (letter, value) from MCQ-formatted string. Returns (None, None) if not MCQ.""" + m = cls._MCQ_GT_RE.match(s) + if m: + return m.group(1), m.group(2).strip() + m = cls._MCQ_PAREN_RE.match(s) + if m: + return m.group(1), m.group(2).strip() + # GT ends with " (A)" pattern: "2+2\sqrt{7} (A)" + m2 = re.search(r'\(?([A-E])\)?\s*$', s) + if m2 and len(s) > 3: + return m2.group(1), s[:m2.start()].strip() + return None, None + + @staticmethod + def _try_numeric_equal(a: str, b: str) -> bool: + """Try numeric equality after normalization. Handles fracs and simple expressions.""" + import math + + def _try_eval(s: str): + # Direct float + try: + return float(s.replace('(', '').replace(')', '')) + except (ValueError, ZeroDivisionError): + pass + # Strip trailing unit-like suffix and retry + s_stripped = re.sub(r'[a-zA-Z]+$', '', s.replace('(', '').replace(')', '')).strip() + if s_stripped and s_stripped != s: + try: + return float(s_stripped) + except (ValueError, ZeroDivisionError): + pass + # Fraction pattern (a)/(b) + frac_re = re.compile(r'^\(([^)]+)\)/\(([^)]+)\)$') + m = frac_re.match(s) + if m: + try: + return float(m.group(1)) / float(m.group(2)) + except (ValueError, ZeroDivisionError): + pass + # Try evaluating simple math expressions (pi, sqrt, etc.) + expr = s + expr = expr.replace('\\pi', str(math.pi)) + expr = expr.replace('\\e', str(math.e)) + expr = re.sub(r'\\sqrt\{([^}]+)\}', r'(\1)**0.5', expr) + expr = re.sub(r'\\sqrt\[3\]\{([^}]+)\}', r'(\1)**(1/3)', expr) + expr = re.sub(r'\\sqrt\[([^]]+)\]\{([^}]+)\}', r'(\2)**(1/\1)', expr) + expr = expr.replace('{', '(').replace('}', ')') + expr = expr.replace('\\cdot', '*').replace('\\times', '*') + expr = re.sub(r'(\d)\(', r'\1*(', expr) + try: + val = eval(expr, {"__builtins__": {}, "math": math, "pi": math.pi, "e": math.e}, {}) + return float(val) + except Exception: + pass + return None + + va, vb = _try_eval(a), _try_eval(b) + if va is not None and vb is not None: + return abs(va - vb) < 1e-6 * max(1, abs(va), abs(vb)) + return False + + @classmethod + def _strip_quantifiers(cls, s: str) -> str: + """Strip universal/existential quantifier wrappers.""" + s = re.sub(r'^\\forall\s*\w+\s*\\in\s*\\mathbb\s*\{?[A-Z]\}?\s*[:,]\s*', '', s) + s = re.sub(r'\s*\(\\forall[^)]*\)\s*$', '', s) + s = re.sub(r'\s*\(for\s+all[^)]*\)\s*$', '', s, flags=re.IGNORECASE) + return s.strip() + + @classmethod + def _try_param_rename(cls, a: str, b: str) -> bool: + """Check if a == b up to consistent single free-parameter rename (ax vs cx).""" + if not a or not b or len(a) != len(b) or len(a) > 80: + return False + diffs = [(i, a[i], b[i]) for i in range(len(a)) if a[i] != b[i]] + if not diffs: + return False + src_chars = set(d[1] for d in diffs) + dst_chars = set(d[2] for d in diffs) + if len(src_chars) == 1 and len(dst_chars) == 1: + src, dst = src_chars.pop(), dst_chars.pop() + if src.isalpha() and dst.isalpha(): + return a.replace(src, dst) == b + return False + + @classmethod + def _normalize_tuple(cls, s: str) -> str: + """Normalize tuple formatting: (2, 5, 609) → 2,5,609.""" + # Strip set-builder conditions: \mid ... or | ... + s = re.sub(r'\\mid.*$', '', s) + s = re.sub(r'\|[^,]*$', '', s) + return re.sub(r'[\s()\[\]{}\\]', '', s) + + @classmethod + def _try_sympy_equal(cls, a: str, b: str) -> bool: + """Optional sympy-based algebraic equivalence (graceful fallback if unavailable).""" + try: + from sympy.parsing.latex import parse_latex + from sympy import simplify, nsimplify + expr_a = parse_latex(a) + expr_b = parse_latex(b) + diff = simplify(nsimplify(expr_a - expr_b)) + return diff == 0 + except Exception: + return False + + @classmethod + def answers_match(cls, predicted: str, reference: str) -> bool: + if not predicted or not reference: + return False + + norm_p = cls.normalize_answer(predicted) + norm_r = cls.normalize_answer(reference) + + # --- Strategy 1: direct string equality --- + if norm_p == norm_r: + return True + + # --- Strategy 2: case-insensitive --- + if norm_p.lower() == norm_r.lower(): + return True + + # --- Strategy 3: numeric equality --- + if cls._try_numeric_equal(norm_p, norm_r): + return True + + # --- Strategy 4: variable-prefix stripping (both sides) --- + stripped_p = cls.normalize_answer(cls._strip_var_prefix(predicted)) + stripped_r = cls.normalize_answer(cls._strip_var_prefix(reference)) + if stripped_p and stripped_r and stripped_p == stripped_r: + return True + if stripped_p and stripped_r and cls._try_numeric_equal(stripped_p, stripped_r): + return True + + # --- Strategy 5: MCQ double matching --- + # Extract letter+value from reference + ref_letter, ref_value = cls._extract_mcq_parts(reference) + if ref_letter: + # pred matches the letter? + if norm_p == ref_letter or predicted.strip().upper() == ref_letter: + return True + # pred matches the value? + if ref_value: + norm_ref_val = cls.normalize_answer(ref_value) + if norm_p == norm_ref_val or cls._try_numeric_equal(norm_p, norm_ref_val): + return True + # Extract from predicted side too (pred="B", ref has value) + pred_letter, pred_value = cls._extract_mcq_parts(predicted) + if pred_letter: + if norm_r == pred_letter or reference.strip().upper() == pred_letter: + return True + if pred_value: + norm_pred_val = cls.normalize_answer(pred_value) + if norm_r == norm_pred_val or cls._try_numeric_equal(norm_r, norm_pred_val): + return True + # MCQ: GT is single letter, pred is numeric/expression → match if pred chose option + if cls._MCQ_SINGLE_LETTER_RE.match(reference.strip()): + if cls._MCQ_SINGLE_LETTER_RE.match(predicted.strip().upper()): + return predicted.strip().upper() == reference.strip().upper() + # pred is a value, GT is just a letter: we accept pred=letter match only + # (can't verify value without options text) + + # --- Strategy 6: tuple/set normalization --- + tuple_p = cls._normalize_tuple(norm_p) + tuple_r = cls._normalize_tuple(norm_r) + if ',' in tuple_p and tuple_p == tuple_r: + return True + if stripped_p and stripped_r: + tuple_sp = cls._normalize_tuple(stripped_p) + tuple_sr = cls._normalize_tuple(stripped_r) + if ',' in tuple_sp and tuple_sp == tuple_sr: + return True + + # --- Strategy 7: equation reorder (a+b=c vs c=a+b, or lhs=rhs swapped) --- + if '=' in norm_r and '=' not in norm_p: + # GT has derivation like "18*1+999*2=2016", pred is "2016" + parts = norm_r.split('=') + for part in parts: + part = part.strip() + if part == norm_p or cls._try_numeric_equal(part, norm_p): + return True + if '=' in norm_p and '=' not in norm_r: + parts = norm_p.split('=') + for part in parts: + part = part.strip() + if part == norm_r or cls._try_numeric_equal(part, norm_r): + return True + if '=' in norm_p and '=' in norm_r: + # Both have =: try matching LHS=RHS in any order + pp = [x.strip() for x in norm_p.split('=')] + rp = [x.strip() for x in norm_r.split('=')] + if set(pp) == set(rp): + return True + + # --- Strategy 8: multiplicative reorder (27\pi\sqrt{6} vs 27\sqrt{6}\pi) --- + def _sort_factors(s): + tokens = re.findall(r'\\?[a-zA-Z]+\{[^}]*\}|\\?[a-zA-Z]+|\d+|[^a-zA-Z\d\\{}]', s) + return ''.join(sorted(tokens)) + if _sort_factors(norm_p) == _sort_factors(norm_r): + return True + + # --- Strategy 9: sympy algebraic equivalence (optional, slow) --- + if cls._try_sympy_equal(predicted, reference): + return True + + # --- Strategy 9b: quantifier stripping + param rename --- + q_stripped_r = cls._strip_quantifiers(reference) + q_stripped_p = cls._strip_quantifiers(predicted) + if q_stripped_r != reference or q_stripped_p != predicted: + norm_qr = cls.normalize_answer(cls._strip_var_prefix(q_stripped_r)) + norm_qp = cls.normalize_answer(cls._strip_var_prefix(q_stripped_p)) + if norm_qr and norm_qp: + if norm_qr == norm_qp: + return True + if cls._try_param_rename(norm_qr, norm_qp): + return True + + # --- Strategy 10: param rename on var-prefix-stripped forms --- + if stripped_p and stripped_r and cls._try_param_rename(stripped_p, stripped_r): + return True + + return False + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + rewards = [] + for traj in trajectories: + messages = traj.get('messages', []) + completion = '' + for msg in reversed(messages): + if msg.get('role') == 'assistant': + completion = msg.get('content', '') + break + user_data = traj.get('user_data') or [] + gt = '' + for item in user_data: + if item[0] == 'ground_truth': + gt = item[1] + break + predicted = self.extract_boxed(completion) + correct = self.answers_match(predicted, gt) + rewards.append(1.0 if correct else 0.0) + return rewards + + +class FormatReward(Reward): + """Reward for having \\boxed{} and ... analysis in the output.""" + + _HINT_RE = re.compile(r'(.*?)', re.DOTALL) + _THINK_RE = re.compile(r'^.*?', re.DOTALL) + _MIN_HINT_LEN = 30 # minimum chars for a substantive hint + _MAX_HINT_LEN = 4096 # hints longer than this are likely thinking dumps + + @staticmethod + def _to_text(content) -> str: + """Convert content (str or list-of-blocks) to plain text.""" + if isinstance(content, str): + return content + if isinstance(content, list): + return ''.join( + b.get('text', '') if isinstance(b, dict) else str(b) + for b in content) + return str(content) if content else '' + + @classmethod + def _visible_response(cls, text: str) -> str: + """Strip ... block to get the visible response.""" + think_end = text.find('') + if think_end >= 0: + return text[think_end + len(''):] + return text + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + rewards = [] + for traj in trajectories: + messages = traj.get('messages', []) + completion = '' + sys_content = '' + for msg in messages: + if msg.get('role') == 'system': + sys_content = self._to_text(msg.get('content', '')) + for msg in reversed(messages): + if msg.get('role') == 'assistant': + completion = self._to_text(msg.get('content', '')) + break + has_boxed = '\\boxed{' in completion + # Only check hint tags for RAG prompts (system contains examples) + is_rag = 'condensed reasoning examples from similar problems' in sys_content + if is_rag: + # Check hint in VISIBLE response only (after ) + visible = self._visible_response(completion) + hint_match = self._HINT_RE.search(visible) + has_good_hint = False + if hint_match: + hint_text = hint_match.group(1).strip() + hint_pos = hint_match.start() + # Hint must be near the start of visible output + at_beginning = hint_pos < max(len(visible) * 0.05, 200) + is_substantive = len(hint_text) >= self._MIN_HINT_LEN + # Reject hints that are too long (model dumping thinking) + not_dump = len(hint_text) <= self._MAX_HINT_LEN + # Must start with the required prefix + has_prefix = hint_text.startswith(HINT_REQUIRED_PREFIX) + has_good_hint = (at_beginning and is_substantive + and not_dump and has_prefix) + # 0.3 for boxed + 0.2 for good hint = 0.5 max + reward = (0.3 if has_boxed else 0.0) + (0.2 if has_good_hint else 0.0) + else: + reward = 0.5 if has_boxed else 0.0 + rewards.append(reward) + return rewards + + +class GibberishPenalty(Reward): + """Negative reward for degenerate outputs (gibberish/random unicode tail).""" + + TAIL_CHARS = 400 + GIBBERISH_THRESHOLD = 0.20 # >20% non-math non-ascii in tail + + @classmethod + def is_gibberish(cls, text: str) -> bool: + if not text: + return False + tail = text[-cls.TAIL_CHARS:] if len(text) > cls.TAIL_CHARS else text + non_math_non_ascii = 0 + for c in tail: + code = ord(c) + # Allow: ASCII, common CJK (for Chinese math), LaTeX symbols + if code > 127 and not (0x4e00 <= code <= 0x9fff): + non_math_non_ascii += 1 + return non_math_non_ascii > len(tail) * cls.GIBBERISH_THRESHOLD + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + rewards = [] + for traj in trajectories: + messages = traj.get('messages', []) + completion = '' + for msg in reversed(messages): + if msg.get('role') == 'assistant': + completion = msg.get('content', '') + break + rewards.append(-0.5 if self.is_gibberish(completion) else 0.0) + return rewards + + +def compute_rewards(trajectories: List[Dict[str, Any]] + ) -> Tuple[List[float], List[float], List[float]]: + acc_fn = AoPSAccuracyReward() + fmt_fn = FormatReward() + gib_fn = GibberishPenalty() + acc = acc_fn(trajectories) + fmt = fmt_fn(trajectories) + gib = gib_fn(trajectories) + total = [a + f + g for a, f, g in zip(acc, fmt, gib)] + return total, fmt, acc + + +# ============================================================================ +# Dataset: AoPS boxed problems +# ============================================================================ +def create_aops_dataset(): + """Load AoPS and create GRPO-style dataset (prompt only, with ground_truth in user_data).""" + from modelscope import MsDataset + from twinkle.data_format import Message, Trajectory + + ds = MsDataset.load(AOPS_DATASET_ID, split='train', + download_mode='reuse_dataset_if_exists') + rows = [] + for row in ds: + if not row['metadata'].get('boxed'): + continue + ref = AoPSAccuracyReward.extract_boxed(row['solution']) + if not ref: + continue + rows.append({'problem': row['problem'], 'ground_truth': ref}) + + logger.info(f'[aops] loaded {len(rows)} boxed problems') + rng = random.Random(AOPS_SEED) + rng.shuffle(rows) + + # Build Trajectory list (prompt-only for GRPO) + trajectories = [] + for r in rows: + # Use direct system prompt as placeholder — will be replaced by RAG pipeline + traj = Trajectory( + messages=[ + Message(role='system', content=SYSTEM_DIRECT), + Message(role='user', content=r['problem']), + ], + user_data=[('ground_truth', r['ground_truth'])], + ) + trajectories.append(traj) + + data_meta = DatasetMeta(data=trajectories) + dataset = Dataset(data_meta) + dataset.set_template('Qwen3_5Template', model_id=MODEL_ID, + max_length=16384, truncation_strategy='delete', + enable_thinking=True) + dataset.encode(add_generation_prompt=True) + return dataset + + +# ============================================================================ +# Main +# ============================================================================ +def main(): + # GPU rank allocation + cond_start = 0 + emb_start = cond_start + CONDENSER_GPUS + sampler_start = emb_start + EMB_GPUS + model_start = sampler_start + SAMPLER_GPUS + + device_groups = [ + DeviceGroup(name='condenser', ranks=list(range(cond_start, emb_start)), + device_type='GPU'), + DeviceGroup(name='emb_model', ranks=list(range(emb_start, sampler_start)), + device_type='GPU'), + DeviceGroup(name='sampler', ranks=list(range(sampler_start, model_start)), + device_type='GPU'), + DeviceGroup(name='model', ranks=list(range(model_start, NUM_GPUS)), + device_type='GPU'), + ] + + model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, fsdp_size=MODEL_GPUS, ulysses_size=2) + sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, dp_size=SAMPLER_GPUS) + emb_mesh = DeviceMesh.from_sizes(world_size=EMB_GPUS, dp_size=EMB_GPUS) + condenser_mesh = DeviceMesh.from_sizes(world_size=CONDENSER_GPUS, dp_size=CONDENSER_GPUS) + + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, + groups=device_groups, lazy_collect=False) + + # -- Training model (full-parameter) -- + model = TransformersModel( + model_id=MODEL_ID, device_mesh=model_mesh, remote_group='model') + model.set_optimizer('AdamW', lr=LEARNING_RATE) + model.set_lr_scheduler('CosineAnnealingLR', T_max=MAX_STEPS, eta_min=0) + model.set_loss('GSPOLoss', epsilon=0.2, epsilon_high=0.28, beta=0.04) + model.set_processor(InputProcessor) + model.set_template('Qwen3_5Template', model_id=MODEL_ID, + enable_thinking=True, max_length=32768) + + # -- Rollout sampler -- + sampler = vLLMSampler( + model_id=MODEL_ID, + engine_args={ + 'gpu_memory_utilization': 0.8, + 'max_model_len': 32768, + }, + device_mesh=sampler_mesh, + remote_group='sampler', + ) + sampler.set_template('Qwen3_5Template', model_id=MODEL_ID, + enable_thinking=True, max_length=32768) + + # -- Embedding model -- + emb_model = TransformersModel( + model_id=EMBED_MODEL_ID, device_mesh=emb_mesh, remote_group='emb_model') + emb_model.set_processor(InputProcessor) + emb_model.set_loss(InfonceLoss, temperature=0.03, use_batch=True) + emb_template = Qwen3_5Template( + model_id=EMBED_MODEL_ID, max_length=EMBED_MAX_LENGTH, + truncation_strategy='delete', enable_thinking=False) + + # -- Condenser sampler -- + condenser_sampler = vLLMSampler( + model_id=CONDENSE_MODEL_ID, + engine_args={'gpu_memory_utilization': 0.85, 'max_model_len': 32768}, + device_mesh=condenser_mesh, + remote_group='condenser', + ) + condenser_sampler.set_template( + 'Qwen3_5Template', model_id=CONDENSE_MODEL_ID, + enable_thinking=False, truncation_strategy='delete', max_length=32768) + condenser_template = Qwen3_5Template( + model_id=CONDENSE_MODEL_ID, max_length=32768, + enable_thinking=False, truncation_strategy='delete') + condenser_special_tokens = set(condenser_template.tokenizer.all_special_tokens) + compress_params = SamplingParams( + max_tokens=CONDENSE_MAX_TOKENS, temperature=CONDENSE_TEMPERATURE, + top_p=0.5, num_samples=1) + + # -- API client (condenser fallback) -- + api_client = None + if CONDENSE_API_KEY: + api_client = OpenAIClient( + model=CONDENSE_API_MODEL, api_key=CONDENSE_API_KEY, + base_url=CONDENSE_BASE_URL) + + # -- LanceDB -- + import lancedb + db = lancedb.connect(DB_PATH) + tbl = db.open_table(DB_TABLE) + logger.info(f'[rag] LanceDB ready, rows={tbl.count_rows()}') + + # -- Checkpoint & DataLoader -- + ckpt_manager = CheckpointEngineManager(model=model, sampler=sampler) + + GLOBAL_BATCH_SIZE = BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS + dataloader = DataLoader( + dataset=create_aops_dataset, + batch_size=GLOBAL_BATCH_SIZE, + min_batch_size=GLOBAL_BATCH_SIZE, + device_mesh=model_mesh, + remote_group='model', + ) + + advantage_fn = GRPOAdvantage() + metrics = CompletionRewardMetric() + sampling_params = SamplingParams( + max_tokens=MAX_NEW_TOKENS, num_samples=1, logprobs=1, + temperature=1.0, top_p=0.95) + + optim_step = 0 + logger.info('Starting RAG-hint GRPO training') + logger.info(get_device_placement()) + + # -- Prefetch: overlap RAG data preparation with training -- + prefetch_pool = ThreadPoolExecutor(max_workers=1) + + def _extract_text(content) -> str: + """Extract plain text from content (str or list-of-parts format).""" + if isinstance(content, str): + return content + if isinstance(content, list): + return ''.join( + p.get('text', '') for p in content if isinstance(p, dict) and p.get('type') == 'text') + return str(content) if content else '' + + def prepare_rag_batch(batch): + """Embed → retrieve → condense → build prompts. Runs in background thread.""" + problems = [] + ground_truths = [] + for item in batch: + msgs = item.get('messages', []) + prob = '' + for m in msgs: + if m.get('role') == 'user': + prob = _extract_text(m.get('content', '')) + break + problems.append(prob) + ud = item.get('user_data', []) + gt = '' + for pair in ud: + if pair[0] == 'ground_truth': + gt = pair[1] + break + ground_truths.append(gt) + + # Embed & retrieve + query_vecs = get_embeddings(emb_model, emb_template, problems, EMB_GPUS) + retrieved = retrieve_topk(tbl, query_vecs, problems, SIM_THRESHOLD) + raw_retrieved_counts = [len(r) for r in retrieved] + + # LLM-based decontamination: judge ALL retrievals via API + if api_client: + judge_pairs = [] # (qi, ret_idx, prob_a, prob_b) + for qi, rets in enumerate(retrieved): + for ri, ret in enumerate(rets): + judge_pairs.append((qi, ri, problems[qi], ret['query'])) + + if judge_pairs: + pairs_input = [(pa, pb) for _, _, pa, pb in judge_pairs] + verdicts = _llm_judge_same_problem(api_client, pairs_input) + to_remove = set() + for vi, (qi, ri, _, _) in enumerate(judge_pairs): + if verdicts[vi]: + to_remove.add((qi, ri)) + if to_remove: + logger.info(f'[decontam-llm] filtered {len(to_remove)} same-problem retrievals') + for qi in range(len(retrieved)): + retrieved[qi] = [ + ret for ri, ret in enumerate(retrieved[qi]) + if (qi, ri) not in to_remove + ] + + # Condense (batch local vLLM + API fallback) + condensed_examples: List[List[Dict[str, str]]] = [[] for _ in range(len(problems))] + tasks_to_condense = [] + for i, rets in enumerate(retrieved): + for j, ret in enumerate(rets): + tasks_to_condense.append((i, j, problems[i], ret)) + + if tasks_to_condense: + condense_prompts = [] + for idx, _j, prob, ret in tasks_to_condense: + user_msg = COMPRESS_USER.format(query=prob, text=ret['thinking']) + condense_prompts.append({'messages': [ + {'role': 'system', 'content': COMPRESS_SYSTEM}, + {'role': 'user', 'content': user_msg}]}) + + try: + condense_responses = condenser_sampler.sample(condense_prompts, compress_params) + except Exception as exc: + logger.warning(f'[condense] local batch error: {exc}') + condense_responses = [None] * len(condense_prompts) + + api_fallback_indices = [] + for ci, (idx, _j, prob, ret) in enumerate(tasks_to_condense): + resp = condense_responses[ci] if condense_responses else None + seq = resp.sequences[0] if resp and resp.sequences else None + text = '' + if seq and seq.stop_reason != 'length' and seq.decoded: + text = seq.decoded + for tok in condenser_special_tokens: + text = text.replace(tok, '') + text = text.strip() + if text: + condensed_examples[idx].append({'query': ret['query'], 'thinking': text}) + else: + api_fallback_indices.append(ci) + + if api_fallback_indices and api_client: + def _fallback(ci): + return ci, _api_condense_single(api_client, condense_prompts[ci]['messages']) + with ThreadPoolExecutor(max_workers=CONDENSE_API_CONCURRENCY) as pool: + futs = [pool.submit(_fallback, ci) for ci in api_fallback_indices] + for fut in as_completed(futs): + ci, result = fut.result() + idx, _j, prob, ret = tasks_to_condense[ci] + text = result if result else ret['thinking'][:MAX_TRACE_LEN] + condensed_examples[idx].append({'query': ret['query'], 'thinking': text}) + elif api_fallback_indices: + for ci in api_fallback_indices: + idx, _j, prob, ret = tasks_to_condense[ci] + condensed_examples[idx].append( + {'query': ret['query'], 'thinking': ret['thinking'][:MAX_TRACE_LEN]}) + + # API hint analysis: pre-compute RAG relevance verdict + hint_analyses = [None] * len(problems) + if api_client: + hint_analyses = _api_hint_analysis_batch(api_client, problems, condensed_examples) + + # Build prompts with rag_fallback_sim check + rag_prompts = [] + rag_debug_records = [] + for i, prob in enumerate(problems): + examples = condensed_examples[i] + rets = retrieved[i] + best_sim = max((r['sim'] for r in rets), default=0.0) + use_rag = bool(examples) and best_sim >= RAG_FALLBACK_SIM + + if use_rag: + # If API hint analysis succeeded, use pre-analyzed prompt (no needed) + if hint_analyses[i]: + rag_sys_content = build_preanalysis_system(hint_analyses[i]) + else: + # Fallback: old-style prompt with self-analysis requirement + parts = [SYSTEM_WITH_RAG_HEADER] + for eidx, ex in enumerate(examples, 1): + parts.append(EXAMPLE_TEMPLATE.format( + idx=eidx, + example_query=ex['query'], + example_thinking=ex['thinking'])) + rag_sys_content = ''.join(parts) + + # RAG group (only RAG, no paired NoRAG) + rag_prompts.append({ + 'messages': [ + {'role': 'system', 'content': rag_sys_content}, + {'role': 'user', 'content': prob}, + ], + 'user_data': [('ground_truth', ground_truths[i])], + 'assistant_prefix': ANALYSIS_PREFIX, + }) + rag_debug_records.append({ + 'problem': prob[:200], + 'ground_truth': ground_truths[i], + 'best_sim': round(best_sim, 4), + 'num_raw_retrieved': raw_retrieved_counts[i], + 'num_retrieved': len(rets), + 'num_condensed': len(examples), + 'use_rag': True, + 'has_preanalysis': hint_analyses[i] is not None, + 'preanalysis_len': len(hint_analyses[i]) if hint_analyses[i] else 0, + 'top_retrieved_query': rets[0]['query'][:200] if rets else '', + 'condensed_len': len(examples[0].get('thinking', '')) if examples else 0, + }) + else: + # No hint found — skip this query entirely + continue + + return rag_prompts, rag_debug_records + + # Submit first batch prefetch + os.makedirs(OUTPUT_DIR, exist_ok=True) + rag_log_path = os.path.join(OUTPUT_DIR, 'rag_diagnostics.jsonl') + rag_log_f = open(rag_log_path, 'w', encoding='utf-8') + logger.info(f'[rag] diagnostics → {rag_log_path}') + + batch_iter = iter(dataloader) + pending_future = None + try: + first_batch = next(batch_iter) + pending_future = prefetch_pool.submit(prepare_rag_batch, first_batch) + except StopIteration: + pass + + try: + while pending_future is not None: + if optim_step >= MAX_STEPS: + break + + metrics.reset() + rag_prompts, rag_debug_records = pending_future.result() + + # Write RAG diagnostics + for rec in rag_debug_records: + rec['step'] = optim_step + rag_log_f.write(json.dumps(rec, ensure_ascii=False) + '\n') + rag_log_f.flush() + + # Submit next batch prefetch (overlaps with rollout + training) + pending_future = None + try: + next_batch = next(batch_iter) + pending_future = prefetch_pool.submit(prepare_rag_batch, next_batch) + except StopIteration: + pass + + # ---- Expand for NUM_GENERATIONS and sample ---- + expand_prompts = [] + for prompt in rag_prompts: + expand_prompts.extend([prompt] * NUM_GENERATIONS) + + if not expand_prompts: + logger.warning(f'[Step {optim_step}] empty prompt list after RAG processing, skip') + continue + + ckpt_manager.sync_weights(merge_and_sync=False) + sampler.reset_prefix_cache() + + sample_responses = sampler.sample(expand_prompts, sampling_params) + + # ---- Collect rollouts ---- + all_input_data: List[Dict[str, Any]] = [] + all_old_logps: List[List[float]] = [] + all_completion_lengths: List[int] = [] + + for sample_response in sample_responses: + for sequence in sample_response.sequences: + all_input_data.append(sequence.new_input_feature) + all_old_logps.append([logprob[0][1] for logprob in sequence.logprobs]) + all_completion_lengths.append(len(sequence.tokens)) + + # ---- Rewards ---- + total_rewards, format_rewards, accuracy_rewards = compute_rewards(all_input_data) + + # Zero out rewards for rollouts that hit the max_tokens ceiling + max_len_threshold = int(MAX_NEW_TOKENS * 0.95) + for i in range(len(all_input_data)): + if all_completion_lengths[i] >= max_len_threshold: + total_rewards[i] = 0.0 + accuracy_rewards[i] = 0.0 + format_rewards[i] = 0.0 + + # Per-step reward summary to diagnostics + n_correct = sum(1 for a in accuracy_rewards if a > 0) + rag_log_f.write(json.dumps({ + 'step': optim_step, 'type': 'reward_summary', + 'n_samples': len(accuracy_rewards), + 'accuracy': n_correct / len(accuracy_rewards) if accuracy_rewards else 0, + 'mean_reward': sum(total_rewards) / len(total_rewards) if total_rewards else 0, + }, ensure_ascii=False) + '\n') + + metrics.accumulate( + completion_lengths=all_completion_lengths, + rewards={ + 'total': total_rewards, + 'format': format_rewards, + 'accuracy': accuracy_rewards, + }, + ) + + # ---- GRPO advantage ---- + advantages = advantage_fn( + total_rewards, num_generations=NUM_GENERATIONS, scale='group').tolist() + if ADV_CLIP > 0: + advantages = [max(-ADV_CLIP, min(ADV_CLIP, a)) for a in advantages] + + # Log all rollout responses (after advantage computation) + _extract_boxed = AoPSAccuracyReward.extract_boxed + def _content_to_str(content): + """Convert message content (str or list of blocks) to plain text.""" + if isinstance(content, str): + return content + if isinstance(content, list): + return ''.join( + b.get('text', '') if isinstance(b, dict) else str(b) + for b in content) + return str(content) + + for ridx, traj in enumerate(all_input_data): + msgs = traj.get('messages', []) + assistant_text = _content_to_str(next( + (m['content'] for m in reversed(msgs) if m.get('role') == 'assistant'), '')) + user_text = _content_to_str(next( + (m['content'] for m in msgs if m.get('role') == 'user'), '')) + sys_text = _content_to_str(next( + (m['content'] for m in msgs if m.get('role') == 'system'), '')) + user_data = traj.get('user_data') or [] + gt = next((v for k, v in user_data if k == 'ground_truth'), '') + problem_idx = ridx // NUM_GENERATIONS + use_rag = ('condensed reasoning examples from similar problems' in sys_text + or 'RAG Analysis (pre-computed)' in sys_text) + # Per-problem group accuracy (all generations for same problem) + grp_start = problem_idx * NUM_GENERATIONS + grp_end = grp_start + NUM_GENERATIONS + grp_acc = sum(accuracy_rewards[grp_start:grp_end]) / NUM_GENERATIONS + + rag_log_f.write(json.dumps({ + 'step': optim_step, 'type': 'rollout', + 'idx': ridx, + 'problem_idx': problem_idx, + 'problem': user_text, + 'system': sys_text, + 'response': assistant_text, + 'ground_truth': gt, + 'predicted': _extract_boxed(assistant_text), + 'use_rag': use_rag, + 'best_sim': rag_debug_records[problem_idx].get('best_sim', 0.0) if problem_idx < len(rag_debug_records) else 0.0, + 'reward': total_rewards[ridx], + 'accuracy_reward': accuracy_rewards[ridx], + 'format_reward': format_rewards[ridx], + 'advantage': advantages[ridx], + 'completion_length': all_completion_lengths[ridx], + 'group_accuracy': grp_acc, + }, ensure_ascii=False) + '\n') + + rag_log_f.flush() + + # ---- Filter out low-signal problem groups (DAPO-style dynamic sampling) ---- + # Skip groups where accuracy is too low (<0.1) or too high (>0.9) + # to avoid gradient dominated by gibberish/format noise or no learning signal. + filtered_inputs, filtered_old_logps, filtered_advantages = [], [], [] + actual_num_groups = len(all_input_data) // NUM_GENERATIONS + for g in range(actual_num_groups): + g_start = g * NUM_GENERATIONS + g_end = g_start + NUM_GENERATIONS + grp_adv = advantages[g_start:g_end] + if all(abs(a) < 1e-8 for a in grp_adv): + continue + grp_acc_rate = sum(accuracy_rewards[g_start:g_end]) / NUM_GENERATIONS + if grp_acc_rate < 0.2 or grp_acc_rate > 0.8: + continue + filtered_inputs.extend(all_input_data[g_start:g_end]) + filtered_old_logps.extend(all_old_logps[g_start:g_end]) + filtered_advantages.extend(grp_adv) + + # ---- Mini-batch training with gradient accumulation ---- + # Process MICRO_BATCH_SIZE samples per forward, accumulate grad_accum_steps + # times before one optimizer step. clip_grad_norm normalizes by accumulated + # num_tokens, ensuring mathematical equivalence with larger batch forward. + total_completions = len(filtered_inputs) + if total_completions == 0: + logger.info(f'[Step {optim_step}] all groups filtered (uniform rewards), skip training') + continue + + grad_accum_steps = MINI_BATCH_SIZE // MICRO_BATCH_SIZE + accum_count = 0 + for mb_start in range(0, total_completions, MICRO_BATCH_SIZE): + mb_end = min(mb_start + MICRO_BATCH_SIZE, total_completions) + mb_inputs = filtered_inputs[mb_start:mb_end] + mb_old_logps = filtered_old_logps[mb_start:mb_end] + mb_advantages = filtered_advantages[mb_start:mb_end] + + outputs = model.forward_backward( + inputs=mb_inputs, + old_logps=mb_old_logps, + ref_logps=mb_old_logps, + advantages=mb_advantages, + ) + accum_count += 1 + + if accum_count % grad_accum_steps == 0: + # Loss spike skip: discard explosive gradients + skip_step = False + try: + loss_val = outputs.get('loss', None) + if loss_val is not None: + if hasattr(loss_val, 'item'): + loss_val = loss_val.item() + if loss_val > LOSS_SPIKE_THRESHOLD: + skip_step = True + logger.warning( + f'[Step {optim_step}] Loss spike: {loss_val:.4f} > ' + f'{LOSS_SPIKE_THRESHOLD}, skipping update') + except Exception: + pass + + if skip_step: + model.zero_grad() + else: + model.clip_grad_and_step() + optim_step += 1 + + if optim_step >= MAX_STEPS: + break + if optim_step % SAVE_STEPS == 0: + model.save(f'rag-hint-grpo-checkpoint-{optim_step}') + + # Flush remaining accumulated gradients (incomplete window at tail) + if accum_count % grad_accum_steps != 0: + skip_step = False + try: + loss_val = outputs.get('loss', None) + if loss_val is not None: + if hasattr(loss_val, 'item'): + loss_val = loss_val.item() + if loss_val > LOSS_SPIKE_THRESHOLD: + skip_step = True + logger.warning( + f'[Step {optim_step}] Loss spike (tail): {loss_val:.4f} > ' + f'{LOSS_SPIKE_THRESHOLD}, skipping update') + except Exception: + pass + + if skip_step: + model.zero_grad() + else: + model.clip_grad_and_step() + optim_step += 1 + + log_dict = metrics.calculate() + log_dict.update(model.calculate_metric(is_training=True)) + metrics.reset() + logger.info(f'[Step {optim_step}/{MAX_STEPS}] {log_dict}') + finally: + prefetch_pool.shutdown(wait=False) + rag_log_f.close() + + logger.info(f'Training completed. optim_steps={optim_step}') + model.save('rag-hint-grpo-final') + + +if __name__ == '__main__': + main() diff --git a/cookbook/sample/emb_sample.py b/cookbook/sample/emb_sample.py index da27a8155..8db4b91a6 100644 --- a/cookbook/sample/emb_sample.py +++ b/cookbook/sample/emb_sample.py @@ -32,10 +32,10 @@ args = CLI.from_args() # -- Config ------------------------------------------------------------------- -CONDENSE_MODEL_ID = args.extra.get('condense_model_id', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') -EMB_MODEL_ID = args.extra.get('emb_model_id', 'ms://twinkle-kit/Qwen3.5-4B-QA-emb') -SAMPLER_GPUS = args.infra.sampler_gpus or 1 -EMB_GPUS = int(args.extra.get('emb_gpus', 1)) +CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') +EMB_MODEL_ID = os.environ.get('EMB_MODEL', 'output/embedding_lora_transformers/step_8000') +SAMPLER_GPUS = int(os.environ.get('SAMPLER_GPUS', 1)) +EMB_GPUS = int(os.environ.get('EMB_GPUS', 1)) EMB_MAX_LENGTH = 8192 # -- Prompts (aligned with train_embedding_full_ddp.py) ----------------------- diff --git a/cookbook/sample/rag_recall_sample.py b/cookbook/sample/rag_recall_sample.py new file mode 100644 index 000000000..691a69f6e --- /dev/null +++ b/cookbook/sample/rag_recall_sample.py @@ -0,0 +1,379 @@ +"""RAG recall test: compress a query via condenser → embed → search LanceDB. + +End-to-end validation that the thinking-trace RAG index built by +``cookbook/exp/embedding/build_thinking_rag_index.py`` is retrievable. + +Architecture (8 GPUs, same as build script): + * GPU 0-3: vLLM condenser (TP=4) + * GPU 4-7: TransformersModel embedding (DP=4) + +Launch: + python cookbook/sample/rag_recall_sample.py + python cookbook/sample/rag_recall_sample.py --query "How to implement binary search?" + python cookbook/sample/rag_recall_sample.py --db-path ./output/thinking_rag/lance.db --top-k 5 +""" +import argparse +import os +import re +import sys +from pathlib import Path +from typing import Any, Dict, List, Optional + +import numpy as np + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import SamplingParams +from twinkle.loss import InfonceLoss +from twinkle.model import TransformersModel +from twinkle.processor import InputProcessor +from twinkle.sampler import vLLMSampler +from twinkle.template import Qwen3_5Template + +logger = get_logger() + +# --------------------------------------------------------------------------- +# Config (mirrors build_thinking_rag_index.py) +# --------------------------------------------------------------------------- +CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') +EMBED_MODEL_ID = os.environ.get( + 'EMBED_MODEL_ID', 'output/embedding_lora_transformers/last-checkpoint') +SAMPLER_GPUS = int(os.environ.get('SAMPLER_GPUS', 4)) +EMB_GPUS = int(os.environ.get('EMB_GPUS', 4)) +NUM_GPUS = SAMPLER_GPUS + EMB_GPUS + +CONDENSE_GPU_MEM = float(os.environ.get('CONDENSE_GPU_MEM', 0.85)) +CONDENSE_MAX_MODEL_LEN = int(os.environ.get('CONDENSE_MAX_MODEL_LEN', 32768)) +CONDENSE_MAX_TOKENS = int(os.environ.get('CONDENSE_MAX_TOKENS', 8192)) +COMPRESS_TEMPERATURE = float(os.environ.get('COMPRESS_TEMPERATURE', 0.2)) +COMPRESS_TOP_P = float(os.environ.get('COMPRESS_TOP_P', 0.5)) +EMBED_MAX_LENGTH = int(os.environ.get('EMBED_MAX_LENGTH', 8192)) +MIN_TEXT_CHARS = int(os.environ.get('MIN_TEXT_CHARS', 256)) + +# --------------------------------------------------------------------------- +# Compress prompts — MUST match build_thinking_rag_index.py exactly. +# --------------------------------------------------------------------------- +COMPRESS_SYSTEM = """\ +You are a compression and summary assistant. For the (query, source) pair, emit a Markdown \ +answer with TWO sections, designed to pair with the `extract_compressed` tool: \ +the reader absorbs `## Summary` directly, then calls `extract_compressed` \ +on any topic-key listed under `## More` to recover its \ +fuller content. + + `## Summary` \u2014 extreme-density text the reader reads directly. + `## More` \u2014 a topic index whose keys are valid arguments \ +to `extract_compressed` for recovering material not captured inline. + +Together the two sections must form a COMPLETE, NON-DISTORTING inventory of the \ +source for the query \u2014 nothing essential lost, nothing implied that the source \ +does not support. NO preamble, NO meta-commentary, NO code fences wrapping the \ +whole output. + +Output skeleton: + +## Summary +Topic: + + +## More +- : +- ... + +Format selection for the inline body (pick the MOST COMPACT form per query, mix \ +when helpful): +- Interface / signature \u2192 code notation directly: `func(a:int)->str` +- Factual / entity \u2192 telegraphic prose; drop function words; \":\" for \"is\", \",\" \ +for \"has\" +- Skill / how-to / usage \u2192 lead with `Use when: `; numbered telegraphic \ +steps `1.do X 2.then Y`; close with `Output: ` when relevant +- Procedural \u2192 numbered short steps +- Analytical / design \u2192 hierarchical bullets with abbreviations + +`## Summary` rules: +1. TOPIC LINE \u2014 line 1 is ALWAYS `Topic: `, even when the \ +query is narrow. Anchors both the reader and the tool. +2. DENSITY \u2014 every token in the body carries query-relevant signal; cut filler. +3. PRIMARY-COMPLETE \u2014 never silently drop a fact essential to answering the \ +query. Anything cut for length MUST appear as a key under \ +`## More`. +4. NON-MISLEADING \u2014 phrasing must not let the reader infer anything the source \ +does not support; partial truths that mislead are worse than honest omissions \ +flagged in the index. +5. SELF-CONTAINED \u2014 the reader can act on the answer without re-opening the source. +6. FAITHFUL \u2014 only content the source supports; no fabrication, no extrapolation. +7. LANGUAGE \u2014 match the source language. +8. NO outer code fences around the whole answer; no meta-commentary. + +`## More` rules (MANDATORY \u2014 this section is never omitted): +1. FORMAT \u2014 each bullet is `- : `: + \u2022 topic-key \u2014 short, unambiguous, grounded in source vocabulary so the \ +`extract_compressed` tool can locate the aspect (e.g. `decorators`, \ +`error handling`, `pitfalls`). + \u2022 hint \u2014 tells WHAT the reader gains by expanding (concrete numbers, code \ +listings, secondary cases, edge details, related context, \u2026); do NOT restate \ +the inline answer. +2. CRITERION \u2014 each bullet names an aspect that EXISTS in the source but is \ +NOT fully captured inline. Material that genuinely fits inline without \ +distortion MUST NOT be duplicated here. +3. FAITHFUL \u2014 hints must be grounded in the source; never speculate or invent. +4. ORDER \u2014 by relevance to the query, then by importance. +5. EMPTY CASE \u2014 if the source is so short / single-purpose that everything \ +fits inline, write a single line `- (none)`. + +Now begin.\ +""" + +COMPRESS_USER = ( + 'Downstream model will read your compressed block to decide whether to ' + 'expand it. Compress faithfully: preserve the passage topic + core facts. ' + 'Do NOT invent facts. Do NOT drop major facts. Do NOT write meta-commentary ' + 'about the Query (never write "Query info: absent", "no X mention", etc.); ' + 'if the passage does not address the Query, still summarize the passage. ' + 'CRITICAL LANGUAGE RULE: detect the dominant language of the Passage ' + '(NOT the Query, NOT this instruction) and write the ENTIRE output in that ' + 'same language; English passage \u2192 English output, Chinese passage \u2192 ' + 'Chinese output, Japanese passage \u2192 Japanese output. NEVER translate, ' + 'NEVER mix languages, NEVER copy these instructions into the output.\n\n' + '## Query (ordering hint only \u2014 still summarize the whole passage)\n{query}\n\n' + '## Passage\n{text}') + +RAG_QUERY_HINT = ( + 'Summarize this query for retrieval. ' + 'The body of ## Summary MUST follow this EXACT 4-line template \u2014 ' + 'do NOT emit "Use when:", numbered procedure steps, or "Output:":\n' + 'Topic: \n' + 'Problem: \n' + 'Skill: \n' + 'Knowledge: \n' + 'Then emit the mandatory ## More section as usual. ' + 'Topic must name the specific pattern, never generic labels.') + +# --------------------------------------------------------------------------- +# Demo queries (diverse domains to exercise retrieval) +# --------------------------------------------------------------------------- +DEMO_QUERIES = [ + 'How can I implement binary search in Python and what are the edge cases?', + 'Explain the Free-Energy Principle in neuroscience and how it relates to active inference.', + '如何用动态规划解决最长公共子序列问题？', + 'What is the optimal turbulence model for simulating airflow around a building?', + '请详细解释快速排序的分治策略及其时间复杂度分析', +] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _strip_outer_codefence(text: str) -> str: + m = re.match(r'^```[a-zA-Z]*\n(.*?)\n```\s*$', text, re.DOTALL) + return m.group(1).strip() if m else text.strip() + + +def _short(text: str, n: int = 120) -> str: + text = (text or '').replace('\n', ' ').strip() + return text[:n] + ('\u2026' if len(text) > n else '') + + +def _build_compress_messages(text: str, query: str) -> List[Dict[str, str]]: + return [ + {'role': 'system', 'content': COMPRESS_SYSTEM}, + {'role': 'user', 'content': COMPRESS_USER.format(query=query, text=text)}, + ] + + +def _wrap_anchor(text: str) -> List[Dict[str, str]]: + return [ + {'role': 'user', 'content': text}, + {'role': 'assistant', 'content': 'Match the correct response here.'}, + ] + + +# --------------------------------------------------------------------------- +# Core pipeline +# --------------------------------------------------------------------------- + +def compress_query(sampler: vLLMSampler, query: str) -> str: + """Compress a query using the condenser; short queries pass through.""" + if len(query) < MIN_TEXT_CHARS: + return query + prompts = [{'messages': _build_compress_messages(query, RAG_QUERY_HINT)}] + params = SamplingParams( + max_tokens=CONDENSE_MAX_TOKENS, + temperature=COMPRESS_TEMPERATURE, + top_p=COMPRESS_TOP_P, + num_samples=1, + ) + responses = sampler.sample(prompts, params) + seq = responses[0].sequences[0] if responses and responses[0].sequences else None + if seq is None: + return query + text = seq.decoded or '' + text = re.sub(r'<\|[^|]+\|>', '', text).rstrip() + text = _strip_outer_codefence(text) + return text if text.strip() else query + + +def embed_query(model: TransformersModel, template: Qwen3_5Template, + text: str) -> np.ndarray: + """Encode a single text as an anchor embedding, returns [H] float32.""" + feat = template.encode({'messages': _wrap_anchor(text)}) + feat['labels'] = [1] + # Pad to EMB_GPUS to avoid dispatch starvation. + pad_n = EMB_GPUS - 1 + pad_feat = template.encode({'messages': _wrap_anchor(' ')}) + pad_feat['labels'] = [1] + features = [feat] + [pad_feat] * pad_n + out = model.forward_only(inputs=features, task='embedding', return_logits=True) + emb = out['embeddings'] + if hasattr(emb, 'detach'): + emb = emb.detach().cpu().numpy() + return np.asarray(emb[0], dtype=np.float32) + + +def search_lancedb(db_path: str, table_name: str, vector: np.ndarray, + top_k: int) -> List[Dict[str, Any]]: + """Search LanceDB table and return top-k results.""" + import lancedb + db = lancedb.connect(db_path) + available = db.list_tables() + table_list = available.tables if hasattr(available, 'tables') else list(available) + if table_name not in table_list: + raise SystemExit(f'Table "{table_name}" not found in {db_path}. ' + f'Available: {table_list}') + tbl = db.open_table(table_name) + results = ( + tbl.search(vector.tolist()) + .metric('dot') + .limit(top_k) + .select(['id', 'source', 'query_raw', 'thinking_raw', + 'query_compressed', 'cot_compressed', 'sim', '_distance']) + .to_list() + ) + return results + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--query', type=str, nargs='*', default=None, + help='Custom queries to test (overrides built-in demos).') + p.add_argument('--db-path', default='./output/thinking_rag/lance.db', + help='LanceDB directory (same as build script).') + p.add_argument('--table', default='thinking_traces', + help='LanceDB table name.') + p.add_argument('--top-k', type=int, default=3, + help='Number of results to retrieve per query.') + return p.parse_args() + + +def main(): + args = parse_args() + + if not Path(args.db_path).exists(): + raise SystemExit(f'DB path does not exist: {args.db_path}\n' + f'Run build_thinking_rag_index.py first.') + + queries = args.query if args.query else DEMO_QUERIES + + # ── 1. Initialize Twinkle ─────────────────────────────────────────── + device_groups = [ + DeviceGroup( + name='sampler', + ranks=list(range(SAMPLER_GPUS)), + device_type='GPU', + gpus_per_worker=SAMPLER_GPUS, + ), + DeviceGroup( + name='emb_model', + ranks=list(range(SAMPLER_GPUS, NUM_GPUS)), + device_type='GPU', + ), + ] + sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, tp_size=SAMPLER_GPUS) + emb_mesh = DeviceMesh.from_sizes(world_size=EMB_GPUS, dp_size=EMB_GPUS) + twinkle.initialize( + mode='ray', nproc_per_node=NUM_GPUS, + groups=device_groups, lazy_collect=False) + + # ── 2. vLLM condenser ─────────────────────────────────────────────── + sampler = vLLMSampler( + model_id=CONDENSE_MODEL_ID, + engine_args={ + 'gpu_memory_utilization': CONDENSE_GPU_MEM, + 'max_model_len': CONDENSE_MAX_MODEL_LEN, + }, + device_mesh=sampler_mesh, + remote_group='sampler', + ) + sampler.set_template( + 'Qwen3_5Template', model_id=CONDENSE_MODEL_ID, + enable_thinking=False, max_length=CONDENSE_MAX_MODEL_LEN) + + # ── 3. Embedding model ────────────────────────────────────────────── + emb_model = TransformersModel( + model_id=EMBED_MODEL_ID, + device_mesh=emb_mesh, + remote_group='emb_model', + ) + emb_model.set_processor(InputProcessor) + emb_model.set_loss(InfonceLoss, temperature=0.03, use_batch=True) + emb_template = Qwen3_5Template( + model_id=EMBED_MODEL_ID, + max_length=EMBED_MAX_LENGTH, + truncation_strategy='delete', + enable_thinking=False, + ) + + logger.info(f'Initialized: sampler GPUs 0-{SAMPLER_GPUS-1}, ' + f'emb GPUs {SAMPLER_GPUS}-{NUM_GPUS-1}') + logger.info(f'DB: {args.db_path} / table: {args.table}') + logger.info(f'Queries to test: {len(queries)}') + + # ── 4. Per-query: compress → embed → search ───────────────────────── + for i, raw_query in enumerate(queries): + print(f'\n{"="*80}') + print(f'[Query {i+1}/{len(queries)}]') + print(f' Raw: {_short(raw_query, 200)}') + + # Compress + compressed = compress_query(sampler, raw_query) + is_passthrough = len(raw_query) < MIN_TEXT_CHARS + if is_passthrough: + print(f' Compressed: (passthrough, len={len(raw_query)} < {MIN_TEXT_CHARS})') + else: + print(f' Compressed ({len(raw_query)}\u2192{len(compressed)} chars):') + for line in compressed.split('\n')[:8]: + print(f' {line}') + if compressed.count('\n') > 8: + print(f' ... ({compressed.count(chr(10))+1} lines total)') + + # Embed + vec = embed_query(emb_model, emb_template, compressed) + print(f' Embedding: shape={vec.shape}, norm={np.linalg.norm(vec):.4f}') + + # Search + results = search_lancedb(args.db_path, args.table, vec, args.top_k) + print(f'\n Top-{args.top_k} Results:') + if not results: + print(' (no results)') + continue + for rank, r in enumerate(results, 1): + dist = r.get('_distance', None) + sim = (1.0 - dist) if isinstance(dist, (int, float)) else None + sim_str = f'{sim:.4f}' if sim is not None else '?' + dist_str = f'{dist:.4f}' if isinstance(dist, (int, float)) else '?' + print(f' [{rank}] cos_sim={sim_str} (dist={dist_str}) source={r["source"]}') + print(f' query: {_short(r["query_raw"], 100)}') + print(f' thinking: {_short(r["thinking_raw"], 150)}') + print() + + print(f'\n{"="*80}') + print('RAG recall test complete.') + + +if __name__ == '__main__': + main() diff --git a/src/twinkle/loss/grpo.py b/src/twinkle/loss/grpo.py index 781b22060..7fb799eca 100644 --- a/src/twinkle/loss/grpo.py +++ b/src/twinkle/loss/grpo.py @@ -64,8 +64,8 @@ def _compute_log_importance_weights( """ import torch log_ratio = per_token_logps - per_token_old_logps - # Clamp for numerical stability - log_ratio = torch.clamp(log_ratio, min=-20.0, max=20.0) + # Clamp for numerical stability (±5 bounds ratio to [exp(-5), exp(5)] ≈ [0.007, 148]) + log_ratio = torch.clamp(log_ratio, min=-5.0, max=5.0) return log_ratio def _compute_per_token_loss( @@ -75,9 +75,16 @@ def _compute_per_token_loss( per_token_logps: 'torch.Tensor', ) -> 'torch.Tensor': """ - Compute per-token loss with PPO clipping. + Compute per-token loss with PPO double-sided clipping. - Override this method in subclasses for different loss formulations. + Standard PPO clip is one-sided: it only bounds the loss when + advantage > 0 and ratio > 1+eps. When advantage < 0 and ratio > 1+eps + (policy moved AWAY from the old action on its own), the loss is + unbounded upward, causing gradient explosions. + + This implementation clips the ratio from BOTH sides regardless of + advantage sign, bounding the per-token loss to at most + (1+eps_high) * |advantage|. Args: ratio: [batch, seq_len] importance sampling ratio @@ -91,7 +98,14 @@ def _compute_per_token_loss( clipped_ratio = torch.clamp(ratio, 1 - self.epsilon, 1 + self.epsilon_high) loss1 = ratio * advantages loss2 = clipped_ratio * advantages - return -torch.min(loss1, loss2) + # Double-sided clip: use max for positive advantage, min for negative. + # Equivalent to: always take the MORE conservative (smaller magnitude) loss. + per_token_loss = torch.where( + advantages >= 0, + -torch.min(loss1, loss2), # positive adv: standard PPO clip + -torch.max(loss1, loss2), # negative adv: clip the OTHER side + ) + return per_token_loss def _aggregate_loss( self, @@ -320,7 +334,7 @@ def _compute_log_importance_weights( """Sequence-level importance sampling: use mean log ratio.""" import torch log_ratio = per_token_logps - per_token_old_logps - log_ratio = torch.clamp(log_ratio, min=-20.0, max=20.0) + log_ratio = torch.clamp(log_ratio, min=-5.0, max=5.0) seq_level_log_weights = ((log_ratio * loss_mask).sum(-1) / loss_mask.sum(-1).clamp(min=1.0)).unsqueeze(-1) return seq_level_log_weights diff --git a/src/twinkle/template/base.py b/src/twinkle/template/base.py index 3c6c29f6c..ae6411119 100644 --- a/src/twinkle/template/base.py +++ b/src/twinkle/template/base.py @@ -682,7 +682,10 @@ def encode(self, trajectory: Trajectory, add_generation_prompt: bool = False, ** assert self.truncation_strategy != 'split', ( 'encode() does not support truncation_strategy=="split" because it may produce multiple outputs. ' 'Use batch_encode() instead.') - return self.batch_encode([trajectory], add_generation_prompt=add_generation_prompt, **kwargs)[0] + encoded = self.batch_encode([trajectory], add_generation_prompt=add_generation_prompt, **kwargs) + if encoded: + return encoded[0] + return None @staticmethod def map_col_to_row(trajectories: Dict[str, Any]): From fbbd147554f0059e774c97d7db67dfd967ac486a Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Fri, 10 Jul 2026 10:18:14 +0800 Subject: [PATCH 05/60] fix --- cookbook/exp/data_pipeline/audit_rubric.py | 167 ++++++ .../exp/data_pipeline/process_and_save.py | 225 ++++++-- cookbook/exp/embedding/eval_dualline_math.py | 373 +++++++++++++ src/twinkle/preprocessor/base.py | 16 +- src/twinkle_agentic/memory/DESIGN.md | 181 +++++++ src/twinkle_agentic/preprocessor/__init__.py | 159 +++++- .../preprocessor/dead_loop_filter.py | 15 +- .../preprocessor/intent_classifier.py | 20 +- .../preprocessor/label_schema.py | 15 + .../preprocessor/language_filter.py | 5 + .../preprocessor/model_filter.py | 14 + .../preprocessor/safety_scorer.py | 11 + .../preprocessor/trajectory_scorer.py | 148 +++++- .../preprocessor/value_selector.py | 288 ++++++++++ src/twinkle_agentic/verifier/__init__.py | 9 +- src/twinkle_agentic/verifier/aggregation.py | 12 + src/twinkle_agentic/verifier/hard_scorer.py | 41 +- .../verifier/rubric_library.py | 108 ++++ .../verifier/rubric_verifier.py | 501 +++++++++++++++++- tests/preprocessor/test_dead_loop_agent.py | 28 + tests/preprocessor/test_dropped_merge.py | 21 + tests/preprocessor/test_intent_think_strip.py | 49 ++ .../test_quality_preprocessor_map_drop.py | 110 ++++ tests/preprocessor/test_value_selector.py | 234 ++++++++ .../test_aggregation_fusion.py | 12 + .../twinkle_agentic/test_diagnosis_salvage.py | 62 +++ .../test_repeated_calls_spin.py | 47 ++ .../test_rubric_stabilization.py | 92 ++++ 28 files changed, 2853 insertions(+), 110 deletions(-) create mode 100644 cookbook/exp/data_pipeline/audit_rubric.py create mode 100644 cookbook/exp/embedding/eval_dualline_math.py create mode 100644 src/twinkle_agentic/preprocessor/value_selector.py create mode 100644 src/twinkle_agentic/verifier/rubric_library.py create mode 100644 tests/preprocessor/test_dead_loop_agent.py create mode 100644 tests/preprocessor/test_dropped_merge.py create mode 100644 tests/preprocessor/test_intent_think_strip.py create mode 100644 tests/preprocessor/test_quality_preprocessor_map_drop.py create mode 100644 tests/preprocessor/test_value_selector.py create mode 100644 tests/twinkle_agentic/test_aggregation_fusion.py create mode 100644 tests/twinkle_agentic/test_diagnosis_salvage.py create mode 100644 tests/twinkle_agentic/test_repeated_calls_spin.py create mode 100644 tests/twinkle_agentic/test_rubric_stabilization.py diff --git a/cookbook/exp/data_pipeline/audit_rubric.py b/cookbook/exp/data_pipeline/audit_rubric.py new file mode 100644 index 000000000..2d3107373 --- /dev/null +++ b/cookbook/exp/data_pipeline/audit_rubric.py @@ -0,0 +1,167 @@ +"""Audit rubric scoring accuracy: re-score a spread of kept agent trajectories +with the SAME teacher RubricVerifier used in the pipeline, and print, per chosen +trajectory, the generated rubric + per-criterion pass rate + a readable segment +summary so a human can judge whether the score is *right* (good and bad alike). + +Run (same env as the pipeline): + LLM_BACKUP_MODEL=qwen3.7-max \ + LLM_BACKUP_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1 \ + LLM_BACKUP_API_KEY=sk-... \ + python cookbook/exp/data_pipeline/audit_rubric.py +""" +import json +import os +from typing import Any, Dict, List, Optional + +PROCESSED = os.environ.get( + 'PROCESSED_PATH', './output/data_pipeline/processed_20260531_200.jsonl') +N_HIGH = int(os.environ.get('AUDIT_N_HIGH', 3)) +N_LOW = int(os.environ.get('AUDIT_N_LOW', 3)) +N_MID = int(os.environ.get('AUDIT_N_MID', 2)) +# Same stabilization policy as the pipeline (skeleton|fixed|off). +RUBRIC_MODE = os.environ.get('TRAJ_RUBRIC_MODE', 'skeleton').strip().lower() +# Repeat each re-score REPEAT times to measure score variance (jitter). >1 to +# check whether the stabilization actually lowered the spread. +REPEAT = int(os.environ.get('AUDIT_REPEAT', 1)) + + +def _lab(row: Dict[str, Any], key: str) -> Optional[Any]: + for kv in (row.get('user_data') or []): + if isinstance(kv, list) and len(kv) == 2 and kv[0] == key: + try: + return json.loads(kv[1]) + except Exception: + return kv[1] + return None + + +def _seg_summary(messages: List[dict], max_chars: int = 900) -> str: + parts = [] + for m in messages: + role = m.get('role', '?') + content = m.get('content') or '' + tc = m.get('tool_calls') + if tc: + try: + calls = json.loads(tc) if isinstance(tc, str) else tc + names = ','.join(c.get('function', {}).get('name', '?') for c in calls) + content = (content + f' [tool_calls: {names}]').strip() + except Exception: + pass + content = content.replace('\n', ' ') + if len(content) > 200: + content = content[:200] + '…' + parts.append(f' {role}: {content}') + text = '\n'.join(parts) + return text if len(text) <= max_chars else text[:max_chars] + '\n …(truncated)' + + +def _infer_intent(messages: List[dict]) -> Optional[str]: + """Structural intent for the whole trajectory (same detectors as the scorer).""" + from twinkle_agentic.preprocessor.intent_classifier import ( + CodeDetector, MathDetector, ToolCallDetector) + # tool_calls arrive JSON-encoded in the processed jsonl; decode so the + # ToolCallDetector (which reads normalized tool_calls) can see them. + norm = [] + for m in messages: + m = dict(m) + tc = m.get('tool_calls') + if isinstance(tc, str) and tc.strip(): + try: + m['tool_calls'] = json.loads(tc) + except Exception: + m['tool_calls'] = [] + norm.append(m) + for det in (ToolCallDetector(), CodeDetector(), MathDetector()): + try: + if det(norm): + return det.intent + except Exception: + continue + return None + + +def main() -> None: + from twinkle_agentic.verifier import (RubricVerifier, + default_intent_base_rubrics, + default_intent_fixed_rubrics) + + rows = [json.loads(l) for l in open(PROCESSED)] + scored = [r for r in rows if _lab(r, 'traj_score') is not None] + scored.sort(key=lambda r: _lab(r, 'traj_score')) + if not scored: + print('no scored rows found') + return + + picks: List[Dict[str, Any]] = [] + picks += scored[:N_LOW] # lowest + mid = len(scored) // 2 + picks += scored[mid:mid + N_MID] # middle + picks += scored[-N_HIGH:] # highest + # de-dup by id, preserve order + seen = set() + uniq = [] + for r in picks: + if r.get('id') not in seen: + seen.add(r.get('id')) + uniq.append(r) + + intent_base = intent_fixed = None + if RUBRIC_MODE == 'skeleton': + intent_base = default_intent_base_rubrics() + elif RUBRIC_MODE == 'fixed': + intent_fixed = default_intent_fixed_rubrics() + rv = RubricVerifier( + max_votes=5, max_votes_long=3, min_votes_long=2, long_margin_threshold=0.18, + min_votes_high=3, high_score_threshold=0.85, + intent_base_rubrics=intent_base, intent_rubrics=intent_fixed) + print(f'[audit] rubric_mode={RUBRIC_MODE} repeat={REPEAT}') + + for r in uniq: + stored = _lab(r, 'traj_score') + seg_scores = _lab(r, 'segment_scores') + messages = r.get('messages') or [] + intent = _infer_intent(messages) + print('=' * 100) + print(f"id={r.get('id')} model={r.get('model_id')} n_msgs={len(messages)} intent={intent}") + print(f"stored traj_score={stored} level={_lab(r,'traj_level')} " + f"segment_scores={seg_scores} safety={_lab(r,'safety_score')}") + print('-- trajectory summary --') + print(_seg_summary(messages)) + traj = {'messages': messages} + if r.get('tools'): + traj['tools'] = r['tools'] + + scalars: List[float] = [] + det = None + for _ in range(max(1, REPEAT)): + try: + det = rv.score_detail(traj, intent=intent) + except Exception as e: + print(f'!! re-score failed: {e}') + det = None + break + scalars.append(det.scalar) + if det is None: + continue + print('-- teacher re-score --') + print(f" llm_scalar={det.llm_scalar:.3f} hard_pass_rate={det.hard_pass_rate:.3f} " + f"scalar={det.scalar:.3f} gated={det.gated} n_votes={det.n_votes}") + if REPEAT > 1: + lo, hi = min(scalars), max(scalars) + mean = sum(scalars) / len(scalars) + var = sum((s - mean) ** 2 for s in scalars) / len(scalars) + print(f" [variance over {REPEAT}] mean={mean:.3f} spread={hi - lo:.3f} " + f"std={var ** 0.5:.3f} scalars={[round(s, 3) for s in scalars]}") + rubric = det.rubric or [] + rates = det.per_item_pass_rate or [] + for i, it in enumerate(rubric): + rate = rates[i] if i < len(rates) else None + kind = 'HARD' if getattr(it, 'is_hard', False) else 'prin' + rate_s = 'n/a' if rate is None else f'{rate:.2f}' + print(f' [{kind}] pass={rate_s} {it.text}') + print('=' * 100) + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/data_pipeline/process_and_save.py b/cookbook/exp/data_pipeline/process_and_save.py index 636c72e60..3a19b1569 100644 --- a/cookbook/exp/data_pipeline/process_and_save.py +++ b/cookbook/exp/data_pipeline/process_and_save.py @@ -18,7 +18,11 @@ Run: CSV_PATH=/mnt/data/yzhao/tastelikefeet/bc/20260531.csv \ - DATASET_TOTAL=200 python cookbook/exp/data_pipeline/process_and_save.py + USE_RUBRIC=1 SELECT_FRAC=0.1 DATASET_TOTAL=2000 \ + python cookbook/exp/data_pipeline/process_and_save.py + +With gating, rubric LLM cost ~ ``SELECT_FRAC * N_kept`` (not ``N_kept``). Ingest +more rows in pass 1; only the global top fraction gets rubric in pass 2. """ import json import os @@ -39,17 +43,28 @@ StructuralNoiseTagger, TokenSoupFilter, TrajectoryOutcomeFilter, - TrajectoryScorer) + TrajectoryScorer, + ValueSelector, + merge_dropped_shards, + run_quality_pipeline, + select_top_for_rubric, + truncate_dropped_logs) from twinkle_agentic.preprocessor import label_schema as L logger = get_logger() # ── Config ──────────────────────────────────────────────────────────────────── CSV_PATH = os.environ.get('CSV_PATH', '/mnt/data/yzhao/tastelikefeet/bc/20260531.csv') -DATASET_TOTAL = int(os.environ.get('DATASET_TOTAL', 200)) +# Default 2000: pass-1 (filter + value_score) scales with N; with USE_RUBRIC=1 and +# SELECT_FRAC=0.1, rubric cost stays ~10% of survivors (similar to old 200×full rubric). +DATASET_TOTAL = int(os.environ.get('DATASET_TOTAL', 2000)) MAP_NUM_PROC = int(os.environ.get('MAP_NUM_PROC', 8)) OUTPUT_DIR = os.environ.get('OUTPUT_DIR', './output/data_pipeline') -OUTPUT_PATH = os.path.join(OUTPUT_DIR, 'processed_20260531_200.jsonl') +_OUTPUT_BASENAME = os.environ.get( + 'OUTPUT_BASENAME', + f'processed_{Path(CSV_PATH).stem}_{DATASET_TOTAL}.jsonl', +) +OUTPUT_PATH = os.path.join(OUTPUT_DIR, _OUTPUT_BASENAME) DROPPED_PATH = os.path.join(OUTPUT_DIR, 'dropped.jsonl') PIPELINE_VERSION = os.environ.get('PIPELINE_VERSION', 'audit-v1') # Set to keep only trajectories above this fused score. None -> keep all (inspect scores only). @@ -60,6 +75,27 @@ # (needs LLM_BACKUP_* / OPENAI_API_KEY; much slower). Without it every clean # trajectory tends to collapse to level 4 because only hard checks discriminate. USE_RUBRIC = os.environ.get('USE_RUBRIC', '') not in ('', '0', 'false', 'False') +TRAJ_FUSION = os.environ.get('TRAJ_FUSION', 'hard_soft_blend') +TRAJ_HARD_CEIL_SKIP = os.environ.get('TRAJ_HARD_CEIL_SKIP') +TRAJ_HARD_CEIL_SKIP = float(TRAJ_HARD_CEIL_SKIP) if TRAJ_HARD_CEIL_SKIP else (0.92 if USE_RUBRIC else None) +TRAJ_SCORER_WORKERS = int(os.environ.get('TRAJ_SCORER_WORKERS', '2' if USE_RUBRIC else '1')) +# Rubric stabilization policy (reduces per-call score jitter for template-like +# intents). 'skeleton' = half-fixed core + generated tail (flexible, DEFAULT), +# 'fixed' = fully fixed per intent (max stability), 'off' = pure generation. +TRAJ_RUBRIC_MODE = os.environ.get('TRAJ_RUBRIC_MODE', 'skeleton').strip().lower() +# Active-learning pre-selection: only the global top SELECT_FRAC by value_score +# gets an (expensive) rubric pass; the rest are hard-only. 1.0 = label everyone +# (disables gating). Only takes effect with USE_RUBRIC=1. +SELECT_FRAC = float(os.environ.get('SELECT_FRAC', '0.1')) +SELECT_MIN = int(os.environ.get('SELECT_MIN', '0')) +SELECT_MAX = os.environ.get('SELECT_MAX') +SELECT_MAX = int(SELECT_MAX) if SELECT_MAX else None +# Persist the full per-segment rubric diagnosis (per-criterion verdict + reason + +# fix + raw teacher output) for rubric-scored rows — the SFT corpus to distill a +# PRM / error-checker LoRA later. Defaults ON when rubric labeling is enabled +# (one extra teacher call per scored segment). Set PERSIST_DIAGNOSIS=0 to skip. +PERSIST_DIAGNOSIS = os.environ.get( + 'PERSIST_DIAGNOSIS', '1' if USE_RUBRIC else '0') not in ('', '0', 'false', 'False') # ── CSV ingestion (custom format: `ts,model,req_id,messages_json`) ───────────── @@ -143,6 +179,14 @@ def _stream_csv_rows(csv_path: str, max_rows: int = 0) -> Iterator[Dict[str, Any 'model_id': model, 'messages': messages, 'user_data': [], + # Pre-declare the only top-level column a downstream step adds + # (IntentClassifier). Without it the ingest schema lacks `intent`, + # and HF datasets.map(num_proc>1) infers per-shard features from + # the FIRST finished writer; shards that added `intent` later get + # that column dropped on concat -> rows with intent/value/prov all + # None (observed as 140/327 at 500 rows x num_proc=32). Declaring + # it up front keeps the Arrow schema identical across shards. + 'intent': None, } emitted += 1 if max_rows and emitted >= max_rows: @@ -153,51 +197,112 @@ def _build_trajectory_scorer() -> TrajectoryScorer: """Hard-only by default; attach an LLM RubricVerifier when USE_RUBRIC is set.""" rubric_verifier = None if USE_RUBRIC: - from twinkle_agentic.verifier import RubricVerifier - # No sampler -> scores via the llm_backup teacher path (LLM_BACKUP_*/OPENAI_*). - rubric_verifier = RubricVerifier() - return TrajectoryScorer(rubric_verifier=rubric_verifier) - - -def build_pipeline() -> QualityPreprocessor: - """Full tag-then-filter pipeline. Mappers annotate; the tail filter drops on tags.""" - return QualityPreprocessor( - pipeline=[ - # 0) lineage first, so even dropped rows carry provenance in the log. - ProvenanceStamp(source=Path(CSV_PATH).stem, pipeline_version=PIPELINE_VERSION), - # 1) canonicalize message schema (heartbeat strip, tool-call normalize, - # reasoning passthrough — P3), then structural / content filters. - MessageNormalizer(), - ModelFilter(), - LanguageFilter(allowed=('en', 'zh')), - # Shallow-chat round cap is 40; agent traces get a far higher ceiling - # (agent_max_rounds) so long tool-calling loops — the highest-value - # distillation data — are not clipped. - HardFilter(min_user_chars_cjk=14, min_user_chars=24, max_rounds=40, - agent_max_rounds=200), - RefuseFilter(), - DeadLoopFilter(), - MessageSanityFilter(), - SpecialCharsFilter(max_ratio=0.6), - TokenSoupFilter(max_chars=8000), - # 2) taggers (never drop): intent key-rounds, structural noise ratio, - # per-round + trajectory scores, safety score. All write user_data. - IntentClassifier(), - StructuralNoiseTagger(), - _build_trajectory_scorer(), # hard-only, or LLM rubric when USE_RUBRIC=1 - SafetyScorer(), # fixed safety rubric; no sampler -> neutral score, still tagged - # 3) read-only outcome filter: drops on the tags above (D6). Enabled - # only when MIN_TRAJ_SCORE is set, else we keep everything to inspect. - TrajectoryOutcomeFilter( - min_traj_score=MIN_TRAJ_SCORE if MIN_TRAJ_SCORE is not None else 0.0, - min_safety_score=None, - drop_unsafe_flag=False, - ), - ], - dropped_log_path=DROPPED_PATH, + from twinkle_agentic.verifier import (RubricVerifier, + default_intent_base_rubrics, + default_intent_fixed_rubrics) + # Intent-aware rubric stabilization: half-fixed skeleton (default) keeps + # the generator flexible while anchoring a shared core; 'fixed' drops + # generation entirely for tool_call/code/math; 'off' = pure generation. + intent_base = intent_fixed = None + if TRAJ_RUBRIC_MODE == 'skeleton': + intent_base = default_intent_base_rubrics() + elif TRAJ_RUBRIC_MODE == 'fixed': + intent_fixed = default_intent_fixed_rubrics() + rubric_verifier = RubricVerifier( + max_votes=5, + max_votes_long=3, + min_votes_long=2, + long_margin_threshold=0.18, + # Re-sample top-band segments so "looks perfect" isn't a lucky draw. + min_votes_high=3, + high_score_threshold=0.85, + intent_base_rubrics=intent_base, + intent_rubrics=intent_fixed, + ) + return TrajectoryScorer( + rubric_verifier=rubric_verifier, + fusion=TRAJ_FUSION, + hard_ceil_skip=TRAJ_HARD_CEIL_SKIP, + scorer_workers=TRAJ_SCORER_WORKERS, + reconcile_max_messages=int(os.environ.get('TRAJ_RECONCILE_MAX_MSGS', '80')), + # Persist the full rubric diagnosis (verdict+reason+fix+raw) for scored + # segments — the SFT corpus for a distilled PRM/checker LoRA. + persist_diagnosis=PERSIST_DIAGNOSIS, ) +def build_pipeline_pass1() -> QualityPreprocessor: + """Pass 1: clean + tag + cheap value scoring (NO LLM rubric). + + Everything here is deterministic/parallel-safe. It ends by stamping a + ``value_score`` on every surviving row so the driver can then pick the global + top fraction for the (expensive) rubric pass. When rubric labeling is off, + the whole pipeline is a single pass and ValueSelector is skipped. + """ + steps = [ + # 0) lineage first, so even dropped rows carry provenance in the log. + ProvenanceStamp(source=Path(CSV_PATH).stem, pipeline_version=PIPELINE_VERSION), + # 1) canonicalize message schema (heartbeat strip, tool-call normalize, + # reasoning passthrough — P3), then structural / content filters. + MessageNormalizer(), + ModelFilter(), + LanguageFilter(allowed=('en', 'zh')), + # Shallow-chat round cap is 40; agent traces capped at 20 logical rounds + # (min user/assistant counts) for pipeline experiments — raise for prod. + HardFilter(min_user_chars_cjk=14, min_user_chars=24, max_rounds=40, + agent_max_rounds=20), + RefuseFilter(), + DeadLoopFilter(), + MessageSanityFilter(), + SpecialCharsFilter(max_ratio=0.6), + TokenSoupFilter(max_chars=8000), + # 2) taggers (never drop): intent key-rounds, structural noise ratio. + IntentClassifier(), + StructuralNoiseTagger(), + ] + if USE_RUBRIC and SELECT_FRAC < 1.0: + # Active-learning pre-selection: cheap value_score for top-fraction gating. + steps.append(ValueSelector()) + if not USE_RUBRIC: + # Single-pass mode: fold scoring + safety + outcome filter in here. + # No selection happened, so safety scores every row (gated=None). + steps += _pass2_tail(gated=False) + # drop_mode='mark': map returns equal-length columns (dropped rows flagged), + # and run_quality_pipeline materializes the removal via Dataset.filter — so a + # partially filtered batch can never leave ghost rows (no remove_columns hack). + return QualityPreprocessor(pipeline=steps, dropped_log_path=DROPPED_PATH, + drop_mode='mark') + + +def _pass2_tail(gated: bool) -> list: + """Scoring + safety + outcome filter (the LLM-touching tail). + + When ``gated`` is True (two-pass active-learning mode) both the rubric scorer + and the safety scorer only spend an LLM call on rows pre-selected by + ValueSelector (``selected_for_rubric``); everyone else is hard/neutral only. + """ + gate = L.KEY_SELECTED_FOR_RUBRIC if gated else None + return [ + _build_trajectory_scorer(), # hard-only, or LLM rubric when USE_RUBRIC=1 + # Fixed safety rubric; gated post-selection so the LLM safety pass runs + # only on selected rows (neutral-safe otherwise). + SafetyScorer(gate_label=gate), + # read-only outcome filter: drops on the tags above (D6). Enabled only + # when MIN_TRAJ_SCORE is set, else we keep everything to inspect. + TrajectoryOutcomeFilter( + min_traj_score=MIN_TRAJ_SCORE if MIN_TRAJ_SCORE is not None else 0.0, + min_safety_score=None, + drop_unsafe_flag=False, + ), + ] + + +def build_pipeline_pass2(gated: bool = True) -> QualityPreprocessor: + """Pass 2: rubric + safety (both gated on ``selected_for_rubric``) + filter.""" + return QualityPreprocessor(pipeline=_pass2_tail(gated=gated), + dropped_log_path=DROPPED_PATH, drop_mode='mark') + + def _print_sample(dataset: Dataset, n: int = 3) -> None: """Show the enrichment kept on a few rows so you can eyeball tags/scores.""" hf = dataset.dataset @@ -229,10 +334,30 @@ def main() -> None: dataset = Dataset(meta) logger.info(f'Ingested {len(dataset.dataset)} rows.') - pipeline = build_pipeline() - # Dataset.map runs the QualityPreprocessor over HF batches (batched=True is - # forced internally). num_proc parallelizes across shards. - dataset.map(pipeline, num_proc=MAP_NUM_PROC, load_from_cache_file=False) + truncate_dropped_logs(DROPPED_PATH) + # Pass 1: clean + tag (+ value_score when gating). run_quality_pipeline runs + # the pipeline as map(equal-length columns, dropped rows flagged) + a single + # Dataset.filter — map never changes row count, so no ghost rows can appear + # (unlike a filtering map, which needs remove_columns and still risks + # partial-batch ghosting). num_proc parallelizes across shards. + run_quality_pipeline(dataset, build_pipeline_pass1(), + num_proc=MAP_NUM_PROC, load_from_cache_file=False) + merge_dropped_shards(DROPPED_PATH) + logger.info(f'After pass 1 (clean+tag): {len(dataset.dataset)} rows kept.') + + if USE_RUBRIC: + gating = SELECT_FRAC < 1.0 + if gating: + # Global top-fraction: pick the most valuable rows for the LLM pass. + _, n_sel = select_top_for_rubric( + dataset, select_frac=SELECT_FRAC, + min_select=SELECT_MIN, max_select=SELECT_MAX) + logger.info(f'Value-gated rubric: {n_sel} rows selected for LLM labeling ' + f'(frac={SELECT_FRAC}).') + # Pass 2: rubric + safety (both gated when a selection ran) + outcome filter. + run_quality_pipeline(dataset, build_pipeline_pass2(gated=gating), + num_proc=MAP_NUM_PROC, load_from_cache_file=False) + merge_dropped_shards(DROPPED_PATH) logger.info(f'After pipeline: {len(dataset.dataset)} rows kept.') _print_sample(dataset) diff --git a/cookbook/exp/embedding/eval_dualline_math.py b/cookbook/exp/embedding/eval_dualline_math.py new file mode 100644 index 000000000..bbb09b388 --- /dev/null +++ b/cookbook/exp/embedding/eval_dualline_math.py @@ -0,0 +1,373 @@ +"""Dual-line math evaluation: baseline vs online process-checking + rubric injection. + +This is **Phase 0 of DESIGN §11.6** ("参数化 memory: 查错 LoRA"): before training any +LoRA, test the *upper bound* of the mechanism "pause every N tokens, let a strong +teacher check the partial reasoning for rubric errors, inject the found issue back +into the context, then resume". If even the strongest teacher checking online cannot +lift math accuracy, distilling that ability into a LoRA is pointless — so we gate on +this first. + +It deliberately reuses the SAME dataset loader, sampling params and answer grader as +``eval_gpqa_rag.py`` so the two lines are directly comparable: + + - **Line A — baseline** (``--mode baseline``): the student model solves each problem + in a single pass (identical to ``eval_gpqa_rag.py --mode direct``). + - **Line B — dualline** (``--mode dualline``, default): the student generates in + ``--chunk-tokens`` slices; between slices a teacher ``RubricVerifier.diagnose()`` + inspects the partial reasoning. When it reports process issues, the concrete + finding is injected as a ``[Checker]`` note and generation resumes. + +The teacher checker is the ``llm_backup`` teacher API (no student sampler is given to +the verifier, so every check is served by the teacher — exactly the Phase-0 setup). +Configure it via the ``LLM_BACKUP_*`` env vars (see ``utils/llm_backup.py``). + +Continuation is done at the token level (crude on purpose — §11.6 says experiment +performance is not a concern): each slice re-feeds the prior ``new_input_feature`` and, +on injection, splices the tokenized note in before resuming. + +Launch examples: + # Dual-line on 200 MATH problems (needs LLM_BACKUP_* for the teacher checker) + LLM_BACKUP_API_KEY=sk-... LLM_BACKUP_BASE_URL=... \\ + python cookbook/exp/embedding/eval_dualline_math.py --target-eval 200 + + # Paired baseline on the same subset (no checker calls) + python cookbook/exp/embedding/eval_dualline_math.py --mode baseline --target-eval 200 +""" +import argparse +import json +import os +import sys +from collections import defaultdict +from typing import Any, Dict, List, Optional + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import SamplingParams as TwinkleSamplingParams +from twinkle.sampler import vLLMSampler + +# Reuse the reference eval's dataset + grading + prompts verbatim so the two +# lines are measured on identical footing. +from eval_gpqa_rag import (GEN_MAX_MODEL_LEN, GEN_MAX_TOKENS, GEN_MODEL_ID, + GEN_GPU_MEM, GEN_GPUS, GEN_TEMPERATURE, GEN_TOP_P, + answers_match, build_direct_prompt, extract_boxed, + load_math) + +logger = get_logger() + +# --------------------------------------------------------------------------- +# Dual-line config +# --------------------------------------------------------------------------- +CHUNK_TOKENS = int(os.environ.get('DUALLINE_CHUNK_TOKENS', 512)) +MAX_CHECKS = int(os.environ.get('DUALLINE_MAX_CHECKS', 8)) +MAX_INJECTIONS = int(os.environ.get('DUALLINE_MAX_INJECTIONS', 3)) +# Only inject when the checker is confident enough that something is wrong. +CHECK_SCORE_FLOOR = float(os.environ.get('DUALLINE_CHECK_FLOOR', 0.6)) + +# The note format wraps the teacher's finding so the student treats it as an +# external hint rather than its own reasoning. Kept short to limit disruption. +INJECT_TEMPLATE = ( + '\n\n[Checker] A quick review of the reasoning so far found an issue: {issue}\n' + 'Please account for this and continue solving.\n\n') + + +# --------------------------------------------------------------------------- +# Teacher checker (Phase-0: pure teacher via llm_backup) +# --------------------------------------------------------------------------- +def _build_checker(): + """RubricVerifier with no student sampler -> every diagnose() hits the teacher. + + Uses a fixed, math-oriented process rubric so we do not spend a rubric- + generation call per slice (the segment here is a partial CoT, not a finished + trajectory). Falls back to auto-generated rubrics if fixed_rubric is cleared. + """ + from twinkle_agentic.verifier import RubricVerifier + from twinkle_agentic.verifier.rubric_verifier import RubricItem + + fixed = [ + RubricItem('The reasoning contains no arithmetic or algebraic error so far', + is_hard=True), + RubricItem('Each step follows logically from the previous ones', is_hard=True), + RubricItem('No formula or theorem is misstated or misapplied', is_hard=True), + RubricItem('The approach is on track to answer the actual question asked', + is_hard=False), + RubricItem('No step contradicts an earlier established fact', is_hard=False), + ] + return RubricVerifier(fixed_rubric=fixed, gate=True) + + +def _checker_available() -> bool: + return bool(os.environ.get('LLM_BACKUP_API_KEY') + or os.environ.get('LLM_BACKUP_BASE_URL') + or os.environ.get('OPENAI_API_KEY')) + + +def _diagnose_partial(checker, problem: str, partial_cot: str): + """Run the teacher checker on the partial reasoning; return (issue_or_None, detail).""" + seg = {'messages': [ + {'role': 'user', 'content': problem}, + {'role': 'assistant', 'content': partial_cot}, + ]} + try: + detail = checker.diagnose(seg, query=problem) + except Exception as exc: + logger.warning(f'[dualline] checker error: {exc}') + return None, None + if detail.overall_ok: + return None, detail + if detail.scalar >= CHECK_SCORE_FLOOR: + # Checker leans "mostly fine"; don't disrupt on a marginal signal. + return None, detail + fails = [it for it in detail.items if not it.verdict] + if not fails: + return None, detail + # Prefer a fix if the checker gave one; else the reason. + parts = [] + for it in fails[:2]: + msg = it.fix or it.reason + if msg: + parts.append(msg) + issue = ' '.join(parts).strip() or detail.summary + return (issue or None), detail + + +# --------------------------------------------------------------------------- +# Token-level segmented generation with mid-stream injection +# --------------------------------------------------------------------------- +def _decode(tokenizer, ids: List[int]) -> str: + return tokenizer.decode(ids, skip_special_tokens=True) + + +def generate_dualline(sampler, tokenizer, problem: str, checker, + base_params: TwinkleSamplingParams, + chunk_tokens: int) -> Dict[str, Any]: + """Generate the reasoning in slices, checking + injecting between slices. + + Returns a dict with the final text, number of checks/injections, and the + per-injection findings (for the debug log / future SFT corpus). + """ + prompt = build_direct_prompt(problem) + + # First slice: encode the trajectory (adds the generation prompt), generate + # up to chunk_tokens. Subsequent slices reuse the returned new_input_feature. + chunk_params = TwinkleSamplingParams( + max_tokens=chunk_tokens, temperature=base_params.temperature, + top_p=base_params.top_p, num_samples=1) + + cur_inputs: Any = [prompt] + gen_ids: List[int] = [] # student-generated token ids only + injected_ids: List[int] = [] # ids we spliced in (excluded from answer) + n_checks = 0 + n_injections = 0 + findings: List[Dict[str, Any]] = [] + total_new = 0 + finished = False + + while total_new < GEN_MAX_TOKENS: + responses = sampler.sample(cur_inputs, chunk_params) + seq = (responses[0].sequences[0] + if responses and responses[0].sequences else None) + if seq is None: + break + gen_ids.extend(seq.tokens) + total_new += len(seq.tokens) + + if seq.stop_reason != 'length': + finished = True + break # hit EOS / stop -> generation complete + + # Length-capped slice: this is a pause point. Check the partial CoT. + if n_checks >= MAX_CHECKS or not checker: + cur_inputs = [seq.new_input_feature] + continue + + partial_cot = _decode(tokenizer, gen_ids) + n_checks += 1 + issue, _detail = _diagnose_partial(checker, problem, partial_cot) + + next_feat = dict(seq.new_input_feature) + if issue and n_injections < MAX_INJECTIONS: + note = INJECT_TEMPLATE.format(issue=issue) + note_ids = tokenizer.encode(note, add_special_tokens=False) + next_feat['input_ids'] = list(next_feat['input_ids']) + note_ids + if 'labels' in next_feat: + next_feat['labels'] = list(next_feat['labels']) + note_ids + injected_ids.extend(note_ids) + n_injections += 1 + findings.append({'at_token': total_new, 'issue': issue}) + cur_inputs = [next_feat] + + final_text = _decode(tokenizer, gen_ids) + return { + 'text': final_text, + 'finished': finished, + 'n_checks': n_checks, + 'n_injections': n_injections, + 'findings': findings, + 'gen_tokens': len(gen_ids), + } + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- +def main(): + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--mode', choices=['baseline', 'dualline'], default='dualline') + p.add_argument('--math-split', default='test') + p.add_argument('--per-level', type=int, default=0, + help='Problems per difficulty level. 0 => --n split across levels.') + p.add_argument('--n', type=int, default=200, + help='Pool size sampled from MATH (stratified by level).') + p.add_argument('--target-eval', type=int, default=200, + help='Stop after this many problems are evaluated (0 = all sampled).') + p.add_argument('--chunk-tokens', type=int, default=CHUNK_TOKENS, + help='Generate this many tokens between checker pauses.') + p.add_argument('--batch-size', type=int, default=16, + help='Baseline mode batch size (dualline runs per-problem).') + 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/dualline/math_{args.mode}_results.jsonl' + + is_dual = (args.mode == 'dualline') + if is_dual and not _checker_available(): + sys.stderr.write( + '[dualline] ERROR: --mode dualline needs a teacher checker but no ' + 'LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / OPENAI_API_KEY is set.\n' + ' Set them, or run --mode baseline for the paired baseline.\n') + sys.exit(1) + + records = load_math(n=args.n, seed=args.seed, split=args.math_split, + per_level=args.per_level) + if args.target_eval > 0: + records = records[:args.target_eval] + sys.stderr.write(f'[dualline] evaluating {len(records)} problems (mode={args.mode})\n') + + 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'[dualline] 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) + + checker = None + tokenizer = None + if is_dual: + checker = _build_checker() + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(GEN_MODEL_ID, trust_remote_code=True) + sys.stderr.write('[dualline] teacher checker ready (llm_backup teacher)\n') + + correct = 0 + total = 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') + + def _grade_and_log(rec, idx, raw_output, extra=None): + nonlocal correct, total + predicted = extract_boxed(raw_output) + is_correct = answers_match(predicted, rec['reference_answer']) + if is_correct: + correct += 1 + total += 1 + debug_rec = { + 'idx': idx, + 'reference_answer': rec['reference_answer'], + 'predicted': predicted, + 'is_correct': is_correct, + 'problem': rec['problem'], + 'model_output': raw_output, + } + if rec.get('level'): + debug_rec['level'] = rec['level'] + if rec.get('type'): + debug_rec['type'] = rec['type'] + if extra: + debug_rec.update(extra) + debug_records.append(debug_rec) + out_f.write(json.dumps(debug_rec, ensure_ascii=False) + '\n') + out_f.flush() + + if is_dual: + for idx, rec in enumerate(records): + result = generate_dualline(sampler, tokenizer, rec['problem'], + checker, gen_params, args.chunk_tokens) + _grade_and_log(rec, idx, result['text'], extra={ + 'n_checks': result['n_checks'], + 'n_injections': result['n_injections'], + 'findings': result['findings'], + 'finished': result['finished'], + 'gen_tokens': result['gen_tokens'], + }) + acc = correct / total if total else 0 + sys.stderr.write( + f' [{total}/{len(records)}] acc={acc:.4f} ({correct}/{total}) ' + f'checks={result["n_checks"]} inj={result["n_injections"]}\n') + else: + import re + for batch_start in range(0, len(records), args.batch_size): + batch = records[batch_start:batch_start + args.batch_size] + 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 = re.sub(r'<\|[^|]+\|>', '', seq.decoded or '').rstrip() + _grade_and_log(rec, batch_start + i, raw_output) + acc = correct / total if total else 0 + sys.stderr.write(f' [{total}/{len(records)}] acc={acc:.4f} ' + f'({correct}/{total})\n') + + overall = correct / total if total else 0 + print(f'\n{"=" * 60}') + print(f'MATH dual-line — mode={args.mode}, model={GEN_MODEL_ID}') + print(f' n={total}, seed={args.seed}, chunk_tokens={args.chunk_tokens}') + print(f'{"=" * 60}') + print(f'Overall accuracy: {overall:.4f} ({correct}/{total})') + + if is_dual: + tot_checks = sum(r.get('n_checks', 0) for r in debug_records) + tot_inj = sum(r.get('n_injections', 0) for r in debug_records) + n_with_inj = sum(1 for r in debug_records if r.get('n_injections', 0) > 0) + print(f' checker: {tot_checks} checks, {tot_inj} injections across ' + f'{n_with_inj}/{total} problems') + + if any(r.get('level') for r in debug_records): + per = defaultdict(lambda: [0, 0]) + for r in debug_records: + lv = r.get('level', 'Unknown') + per[lv][1] += 1 + if r['is_correct']: + per[lv][0] += 1 + print('\nPer-level accuracy:') + for lv in sorted(per.keys()): + c, t = per[lv] + print(f' {lv:>10}: {c/t:.4f} ({c}/{t})') + + out_f.close() + print(f'\n[output] {len(debug_records)} records saved to {args.output}') + + +if __name__ == '__main__': + main() diff --git a/src/twinkle/preprocessor/base.py b/src/twinkle/preprocessor/base.py index 588b29ae7..4d9cb51c0 100644 --- a/src/twinkle/preprocessor/base.py +++ b/src/twinkle/preprocessor/base.py @@ -42,10 +42,11 @@ def map_row_to_col(rows, keys: List[str] = None) -> Dict[str, List[Any]]: return {k: [] for k in keys} if keys else {} columns: Dict[str, List[Any]] = {} - keys = keys or rows[0].keys() + row_keys = list(rows[0].keys()) + out_keys = row_keys if not keys else list(dict.fromkeys(row_keys + [k for k in keys if k not in row_keys])) - for key in keys: - columns[key] = [row[key] for row in rows] + for key in out_keys: + columns[key] = [row.get(key) for row in rows] return columns @@ -79,12 +80,19 @@ class Filter(Preprocessor): def keep(self, row: Dict[str, Any]) -> bool: raise NotImplementedError + def drop_reason(self, row: Dict[str, Any]) -> str: + """Short reason tag for dropped-row logs (override in subclasses).""" + return type(self).__name__ + def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: rows = self.map_col_to_row(rows) kept: List[Dict[str, Any]] = [] dropped: List[Dict[str, Any]] = [] for r in rows: - (kept if self.keep(r) else dropped).append(r) + if self.keep(r): + kept.append(r) + else: + dropped.append(dict(r, drop_reason=self.drop_reason(r))) return kept, dropped diff --git a/src/twinkle_agentic/memory/DESIGN.md b/src/twinkle_agentic/memory/DESIGN.md index 6baf5e350..f27f1cb19 100644 --- a/src/twinkle_agentic/memory/DESIGN.md +++ b/src/twinkle_agentic/memory/DESIGN.md @@ -19,6 +19,12 @@ **line B 的存在理由**：在线蒸馏有延迟（攒数据→训练→部署），空窗期 student 学不到新东西；memory 立刻起作用填这个空窗，等能力被训进权重后再从 memory 退休。 +> **line B 有两种载体（同为"不动 base 的外挂补充"，见 §11.6）**： +> 1. **文本 memory**：检索注入 context —— 本文 §1~§9 讨论的主形态，适合**易变事实/实体/偏好**（一条即用即改，天然带"检不到就拒答"信号）。 +> 2. **参数化 memory（LoRA，base 冻结）**：把能力挂成可插拔 adapter —— 适合**可泛化的行为模式**（如"边生成边查错"），在已部署 vLLM 里边际显存/进程成本最低、可回滚（换 adapter ≈ 删一条 memory，而非重训 base）。 +> +> 二者与 line A 的分界是"**动不动 base**"：line A 改 base 权重（有遗忘风险，故慎重、攒批离线做）；参数化 memory 虽然也是"权重形态"，但**base 全程冻结、只挂 LoRA**，本质仍是 line B（可插拔、可回滚、即挂即用）。参数化 memory 与文本 memory 的取舍见 Substrate Asymmetry（§11 反方证据）：参数化擅长行为/风格，文本检索擅长事实/拒答。 + **双线共用的唯一“质检车间”**：现有 `preprocessor + verifier`。它产出的 `traj_score / round_scores / confidence / intent / safety` 同时作为两条线的准入闸门，不重复造。 ### 分工判据（什么进 A，什么进 B） @@ -562,3 +568,178 @@ harness.mode = 'static' # 固定路由表，永远兜底 4. **进化被工程化为可回滚的发布流程**（static→shadow→canary→live），而非在线自由漂移——把 SEA 的“门控只能 select 已有行为”落成部署 mode。 **一句话定位**：DuoMem 证明了双线值得做，UCOB/ATMem 证明了 memory-on/off 对照能归因，GovMem/MemDelta 证明了不归因会污染/被混淆，SEA 证明了进化要门控——**本设计是把这些已被各自验证的结论，收进一个共用 verifier/`llm_backup`/D7c 的单一自进化蒸馏闭环。** + +--- + +## 11. skill/rubric → 参数化（LoRA）：双 LoRA 方案的文献支撑 + +> 背景：讨论中提出过一个具体的 line A 实现形态 —— **① 把高分轨迹的 system skills 蒸成一个"技能 LoRA"，推理时先用它产出/注入 skill；② 把 rubric 评价能力蒸成另一个"评价 LoRA"，推理中每隔 N 个 token 用它检验，当过程奖励/memory 提示。** 这一节把 2026 上半年直接对应这三个子命题（skill→LoRA、rubric→LoRA、多 LoRA 推理时切换）的工作按 §10 的详尽风格补全，并标注每篇对方案的取舍启示。**结论先行：三个子命题都各有直接平行工作，MetaClaw 几乎是整套方案 + 本双线设计的镜像；但"评价要不要绕一圈做成 LoRA judge"和"过程监督要不要走 LoRA"这两处，有明确的反方证据，需先决策。** + +### 11.1 skill → LoRA（对应子命题 ①：技能 LoRA） + +这一类的**共同范式**高度一致：**离线用完整 skill 文本合成"技能引导"的示范 → 训一个 skill 专属 LoRA → 在线丢掉 skill 文本、动态挂 LoRA 激活行为**。动机都是：skill 文本每步注入 context 太贵、且长上下文里关键指令定位不到 / 遵守不了（小模型尤甚，ICL 一贯不如微调、且模型越小差距越大）。 + +**Skill-to-LoRA (S2L)**（arXiv 2606.16769）—— **最贴近方案 ① 的朴素做法（一 skill 一 LoRA，不用 hypernetwork）** +- 想解决的问题：agent skill 现在以 `SKILL.md`（人读的流程文档：workflow / 工具 / 资源 / 领域约定）分发，可读可复用，但**同一套可复用流程要在每步 runtime context 里反复注入**，费 token 又稀释注意力。 +- 怎么做（**离线合成 + 在线换挂**，behavior-centric）：不压缩文档本身，而是**建模"skill 文本诱导的行为改变"**。**离线**：把完整 `SKILL.md` 喂进去，让模型合成一批"skill 引导下的示范轨迹"（demonstrations），用这些示范 SFT 出一个**该 skill 专属的 LoRA**；**在线**：完全省略 `SKILL.md` 全文，按当前需要动态加载对应 LoRA 来"激活"这个技能行为。 +- 数据 / loss：标准 SFT（在合成的 skill-guided 示范上做 teacher-forcing 训 LoRA），无 RL。Qwen3.6-27B、SWE-Skills-Bench 21 个 skill 子集：比 no-skill +2.9 pp、比 Full-Text +5.2 pp，每步 token 比 Full-Text prompting −6.6%；18/21 个 skill 追平或超过 Full-Text、15/21 超过 no-skill。**关键对照实验：Wrong-LoRA（挂错 skill 的 LoRA）和 Shared-LoRA（所有 skill 共用一个 LoRA）都掉点** → 收益依赖 **skill-专属对齐**，不是"随便训个 LoRA 就行"。 +- 与我们的区别：这就是方案 ① 最省事的落法（不引入 hypernetwork，一 skill 一 LoRA、离线 SFT）。**它的对照实验直接给我们敲定了两条工程铁律**：(a) 技能 LoRA 必须按 skill 对齐、**不能把多 skill 或 query 糊进一个 LoRA**（呼应 §0 分工判据："query 是易变实体，训进权重会过拟合"）；(b) 挂错 LoRA 反而有害 → 上线要有"挂哪个 skill LoRA"的可靠路由（正好复用我们 meta harness / intent 标签）。差异：S2L 是**离线一次性**、skill 集固定；我们要它在自进化闭环里**持续产出新技能 LoRA**，且用主线 `traj_score` 当"哪些高分轨迹够格蒸成 skill LoRA"的门槛。 + +**LatentSkill**（arXiv 2606.06087）—— **用 hypernetwork 把文本 skill 一次前向转成 LoRA** +- 想解决的问题：同 S2L（per-step 注入 skill 费 context、且 skill 明文暴露），但更进一步想要"**不为每个 skill 单独训 LoRA**"。 +- 怎么做：训一个**预训练 hypernetwork**，输入文本 skill、输出即插即用的 LoRA adapter（把 skill 知识存进**权重空间**而非 context 空间）。保留了 LoRA 的模块化：可加载、可缩放（用 LoRA scaling 系数精确调强弱）、可组合（对齐时能在**参数空间做算术**叠加多个 skill）。 +- 数据 / loss：hypernetwork 预训练。ALFWorld seen/unseen +21.4 / +13.4 分、prefill token −64.1%；Search-QA EM +3.0、skill-token 开销 −72.2%。分析发现生成的 skill LoRA 形成**结构化语义几何**。 +- 与我们的区别：如果我们不想"一个 skill 训一个 LoRA"（S2L 的痛点是 skill 一多 LoRA 就爆炸），LatentSkill 的 hypernetwork 是升级路径 —— **一次前向即出新 skill 的 LoRA、零梯度更新、零 skill 专属数据采集**。但它更复杂、要预训练 hypernetwork，属于方案 ① 的"进阶版"，建议 S2L 跑通、验证 skill LoRA 确有增益后再考虑。它的"参数空间算术组合"对我们"把多条相关碎片 memory 合并升华"（§2 consolidation / B→A 晋升）是权重侧的对应工具。 + +**ParametricSkills**（arXiv 2606.30015）—— **hypernetwork 同时参数化"skill 内容 + 利用方法"，含自进化/持续学习** +- 想解决的问题：同上两条，外加一个更深的观察 —— **文本空间演化 skill（EvoSkill/SkillOpt 等改写 SKILL.md）和模型本身的学习是解耦的**，模型能力没被优化。 +- 怎么做（三阶段，hypernetwork 驱动）：(1) 建 **45.8k skill 库**（网爬 + 从真实 agent 轨迹总结，覆盖 13 领域），用 OpenCode 沙箱围绕这些 skill 合成单/多轮"skill 利用轨迹"；(2) **skill-重建预训练**：三个自监督目标让 hypernetwork 学会把 skill 编码成 LoRA —— **完整重建**（据全文生成能重建全文的 adapter）、**前缀补全**（只给前缀、要补出全文，学 skill 结构）、**段级 cloze 补全**（挖掉一个功能段、据前后文补，学"触发条件/执行步骤/失败处理"如何组织与互相支撑 → 强组合泛化 + 支持局部编辑）；(3) 在 skill-利用轨迹上**多轮 SFT** hypernetwork。 +- 数据 / loss：自监督重建 + 多轮 SFT（loss 都 backprop 到 hypernetwork）。6 个 SWE 子任务比 ICL +6.44 分（DeepSeek-V4-Flash 判）、BERTScore +1.17、F1 +5.53%；**注意 text-to-LoRA 基线 SHINE 反而打不过 in-context skill** → hypernetwork 训不好会退化。持续学习：把多条轨迹的经验**不断 merge 成一个全局 parametric skill**。 +- 与我们的区别：它把 §2「抽取器越来越专业」和 §3「B→A 晋升」在**权重侧**给出了一个具体形态 —— "文本演化 skill = 直接改进模型"。三个自监督目标（尤其**段级 cloze**）可直接借来当我们"技能 LoRA 抽取器"的预训练任务。但它同样是**离线训 hypernetwork**、且 SHINE 反例提醒**参数化 skill 不保证优于文本注入**，必须带对照验证（呼应 §4 影子对照）。 + +### 11.2 rubric / 评价 → 参数化（对应子命题 ②：评价 LoRA）—— 两条岔路，务必先选 + +**支线 A：把 rubric 打分能力做成一个轻量 LoRA judge（= 方案 ② 的原意）** + +**Plug-and-Play LLM Judges**（arXiv 2506.05748）—— **"rubric + 小 LoRA = 顶级裁判"的最强直接背书** +- 想解决的问题：RLHF 的奖励模型训练是成本瓶颈（动辄几十亿参数 + 离线偏好微调阶段）。 +- 怎么做：**冻结的 instruction-tuned 7B + 一行 JSON rubric + rank-16 LoRA（只动 0.8% 参数）**，就当完整奖励模型用。消融：6 条 in-context 示范贡献了大部分零样本→少样本增益（+2pp），**LoRA 补上剩余差距**（尤其 safety / 对抗性 Chat-Hard 段）。 +- 数据 / loss：小 LoRA 微调 + prompt 工程。RewardBench 96.2%，超过 27B~70B 专用奖励网络；配它当在线 PPO 的 reward，7B actor 在 GSM-8K 拿 92% EM、超过 70B DPO 基线；LoRA judge 的解释与人类相似度 ≈9/10（零样本裁判仅 ≈5/10）。 +- 与我们的区别：**这是方案 ②「rubric→评价 LoRA」最硬的可行性背书** —— 极小 LoRA + 一条 rubric 就能把通用模型变成高质量、可解释、可调的裁判。直接支持我们把 `RubricVerifier`（现在调 dashscope teacher）蒸成本地评价 LoRA：既复用 `score_lora_path` 现成入口，又能治昨天"每段调远程 LLM 太慢"的痛（本地 LoRA 打分快几个数量级）。 + +**支线 B：跳过 judge，把 rubric 直接蒸进 policy 的 token 级信号（更省，可能是更优解）** + +> ⚠️ 这两篇机制外壳都是"rubric-conditioned 的自己当 teacher、逐 token 蒸给 unconditioned 的自己"，但**要解决的痛点、对标的对手、卖点完全不同**，别当成一篇：**RCSD 的对手是 RL 的标量奖励**（卖点＝把标量 reward 升级成过程级 token 信用分配）；**RGSD 的对手是 rubric 训练里的那个 LLM verifier**（卖点＝把 verifier 从训练回路里彻底删掉）。下面各自只讲其独有点。 + +**Rubric-Conditioned Self-Distillation / RCSD**（arXiv 2606.19327）—— **卖点：用 rubric 替代 RL 的"标量奖励"，做过程级信用分配** +- 想解决的问题：针对的是**蒸馏与 RLVR 两种 post-training 各自的病**。蒸馏靠 CoT 标注（贵、可能有噪/不全/半错，**哪怕最终答案对，坏 rationale 也会干扰学习**）；RLVR 则把评价**压成一个标量 reward**，看不出"该改推理的哪一步"。它要的是一个**比标量更细的过程监督信号**。 +- 怎么做（**独有点＝两阶段 pipeline + 显式过程级信用分配**）：核心是"**不把单一参考 rationale 当唯一监督靶**"，而让 teacher 看 criterion 级 rubric、在 student 自采样轨迹上给 token 级指导。落地成**两阶段**：**阶段① 先训一个"生成 task-specific rubric"的模块**（给任务先产出该任务的评分条目），**阶段② 再用这些 rubric 训"rubric 引导的 reasoner"**。rubric 说明"强回答该满足什么" → 转成**过程级信用分配**，这是它明确对标 GRPO 的地方。 +- 数据 / loss：on-policy 自蒸馏（token 级）。科学推理套件上**比 GRPO +1.0、比 OPSD +0.9**（对手是 RL/在线自蒸馏方法，不是 verifier）。 +- 与我们的区别：它证明 **rubric 能当"过程级 reward"直接进 policy 训练**，比标量 GRPO 奖励细。对我们的启示落在**训练信号形态**上：如果 line A 想要比 `traj_score` 标量更细的过程监督，RCSD 的"rubric→token 级信用分配"是替代 GRPO 标量奖励的路子；它的**阶段① rubric 生成器**正对应我们 `RubricVerifier` 的 stage-1（可复用）。 + +**Rubric-Guided Self-Distillation / RGSD**（arXiv 2606.12507）—— **卖点：verifier-free，把 LLM judge 从训练回路里删掉** +- 想解决的问题：针对的是**现有 rubric 训练法都要挂一个 LLM verifier 给每条 rollout 打分**这件事本身 —— 带来三个后果：训练期 verifier 调用**开销大**、优化被**特定 verifier 的偏差**污染、且 rubric 反馈被 verifier 压成**稀疏的轨迹末端信号**（只有一个 end-of-trajectory 分）。它要的是**根本不调 verifier**。 +- 怎么做（**独有点＝极简、零 verifier、单 rollout**）：直接拿 rubric-conditioned base policy 当 teacher、unconditioned 当 student 逐 token 蒸 —— 关键在于它把这套做到了**训练回路里完全没有 LLM judge**、且**每 prompt 只需一条 on-policy rollout**（不用像 GRPO 那样一题多采样再让 judge 排序）。 +- 数据 / loss：**零 verifier 调用 + 单 rollout/prompt**。Qwen-2.5(3B/7B)、Qwen3-Thinking(4B/8B) 医学/科学域：rubric 满足度**与 judge-based GRPO 相当**（对手是"带 verifier 的 rubric 训练"）。独有消融：**raw rubric 比"自生成参考回答"是更强的 teacher 富化信号**；但**更强的 GRPO judge 在某些设置能反超 RGSD** → 它诚实地把自己定位为"**当 verifier 成本/可靠性是瓶颈时**"的互补方案，而非全面更优。 +- 与我们的区别：**这篇最该在决策前读透**。它直击我们现状——verifier 又贵又不稳（RuVerBench 警示 + dashscope 每段调用慢）时，**把 rubric 直接蒸进 policy 比"训一个评价 LoRA 再在线打分"更省更稳**。据此，方案 ② 有两条路：**A 训评价 LoRA judge（Plug-and-Play 背书，产物是可复用的独立评价分）** vs **B 走 RGSD/RCSD 把 rubric 直接蒸进主 policy（不产独立分、只喂 student）**。选 A 还是 B，取决于我们是否真需要一个"能被 memory 效用 / 在线 PRM 复用的独立评价分"——若需要就 A，若只为提升 student 就 B。 + +### 11.3 多 LoRA 推理时切换 / 评价即一个 LoRA（对应子命题 ③：每 N token 用评价 LoRA 检验） + +**VideoMind — Chain-of-LoRA**（arXiv 2503.13444，**ICLR 2026**）—— **方案 ③"主生成 LoRA + 评价 LoRA 交替"的现成机制原型** +- 想解决的问题：视频时序 grounding 推理要多种能力（定位、验证、回答），但为每种能力各开一个完整模型太重。 +- 怎么做（两个创新）：(1) **角色化 agent 工作流** —— planner 协调、grounder 时序定位、**verifier 评估候选**、answerer 回答；(2) **Chain-of-LoRA**：一个统一 base model + **多个 LoRA adapter**，推理时**无缝切换角色**（用哪个角色就挂哪个 LoRA），在"每角色一个完整模型"和"纯 prompt 切角色"之间取平衡。 +- 数据 / loss：各角色 LoRA 分别训。15 个 benchmark（Grounded VideoQA / 时序 grounding / 通用 VideoQA）验证有效，且利于 test-time scaling / 长视频。 +- 与我们的区别：**它的 verifier 就是链条里的一个 LoRA 角色，与主生成角色在同一 base 上按需热切** —— 这正是方案 ③ 想要的"生成 LoRA / 评价 LoRA 交替"的成熟原型，**强烈建议精读**它怎么做角色切换调度。⚠️ 但注意：Chain-of-LoRA 是**在"角色回合"边界切**（planner→grounder→verifier→answerer 各跑一段），**不是"每 N 个 token 打断主流插评价"**；后者需要在 decode 中途暂停、跑旁路评估、再续，当前我们 vLLM sampler 的 `sample`/`sample_stream` 没有这种中途插入钩子（一次前向也只能挂一个 LoRA），要自己写**交错解码调度器**——这是方案 ③ 真正的工程量所在。 + +**MetaClaw: Just Talk**（arXiv 2603.17187，**有 code**）—— **几乎是整套方案 + 本双线设计的生产级镜像** +- 想解决的问题：部署的 agent 是**静态**的，跟不上用户需求漂移；在 OpenClaw（20+ 渠道、杂负载）上，现有法要么只存原始轨迹不蒸馏、要么静态 skill 库、要么retrain 要停机。 +- 怎么做（**双互补机制**）：(1) **skill 驱动的快适应** —— LLM evolver 分析**失败轨迹**合成新 skill，**零停机立刻生效**（= 我们的 line B / 快记忆）；(2) **机会式策略优化** —— **云端 LoRA 微调 + RL-PRM（带过程奖励模型的 RL）** 做梯度更新（= 我们的 line A / 慢权重），由 **OMLS 调度器**在**用户空闲窗口**（监控系统空闲 + 日历）触发。两机制**互相喂**：更好的策略产更好轨迹给 skill 合成，更丰富的 skill 给策略优化更高质数据。用**版本机制**分离 support / query 数据**防污染**。proxy 架构、无需本地 GPU。 +- 数据 / loss：SFT/RL-PRM（LoRA）+ 无训练的 skill 合成。skill 快适应相对 +32%；全流程把 Kimi-K2.5 从 21.4%→40.6%、综合鲁棒性 +18.3%。 +- 与我们的区别：**这是与本设计 §0 双线 + §3 毕业 + §6 harness 最像的一篇，且是 OpenClaw 生产场景**。相同点几乎逐条对上：skill 快线 + LoRA/PRM 慢线、失败轨迹驱动 skill、空闲窗口触发训练、版本防污染。差异（也是我们的增量）：MetaClaw 的两机制是**并列互喂**，**没有显式的 B→A 晋升 / A→B 退休毕业梯度**，也**不做"失败是模型问题还是 memory 误导"的因果归因**（§4）；我们多了归因前置、毕业出口、和 static→shadow→canary→live 的可回滚门控。**它是本方案最好的对标基线与工程参考（有 code），建议直接研读其 OMLS 调度 + RL-PRM 实现。** + +### 11.4 ⚠️ 反方证据：过程监督 / TTT 未必要走 LoRA + +**SCATR: Simple Calibrated Test-Time Ranking**（arXiv 2604.16535） +- 想解决的问题：Best-of-N 的效果全看打分函数；学出来的 PRM 强但**训练/推理都贵**，而基于 token logprob 的轻量置信度启发式又**明显偏弱**。 +- 怎么做：从**小校准集**学一个**轻量 scorer**，用的是 **base 模型的隐藏表示**（不是训 LoRA、不生成）。 +- 数据 / loss：轻量回归头。编码/数学基准上比置信度基线 +最高 9%；**相对在同样校准数据上做 LoRA 微调，用少 8000× 的可训练参数达到相当精度**，训练/推理延迟分别快 150×/1000×；和强 PRM 相当、部分设置数学 +7.8%/代码 +4.2% 且推理快 1000×。 +- 与我们的区别：**直接质疑方案 ③"评价/PRM 一定要做成 LoRA"** —— 一个读 base 隐藏层的轻量头，可能比评价 LoRA 更省几个数量级且更快。若方案 ② 的评价只用于"打个过程分排序/门控"（而非要它生成文字理由），**优先考虑 SCATR 式轻量头，而不是 LoRA**。 + +**Surprisal-Guided Selection**（arXiv 2602.07670） +- 想解决的问题：可验证、密集奖励任务（如 GPU kernel 优化，有确定性 evaluator）下，test-time 到底该"梯度自适应"还是"搜索"？ +- 怎么做 / 结论：KernelBench + GPT-OSS-120B(LoRA)：**Best-of-N 搜索（K=64 达 90% 成功）远胜 TTT 梯度自适应（最好 30.6%）**；TTT 会**过度锐化**、把多样性塌成平庸解，"等效 K < 1"（还不如单样本）。零成本妙招：**选 surprisal 最高（最不自信）的正确样本**比选最自信的 +30%。 +- 与我们的区别：提醒我们——**有确定性 verifier 时，算力花在"采样多样性 + 聪明选样"常比训 LoRA / 在线梯度自适应更值**。方案 ② / ③ 若目的是"提升生成质量"，先比一比"评价 LoRA + 干预" vs "多采样 + 评价选样"哪个划算，别默认前者。 + +**VDS-TTT**（arXiv 2505.19475）—— **支持方案：verifier 选样 → 只训 LoRA** +- 怎么做：learned verifier 给一批候选打分，**选高分伪标签**（置信度过阈值）配对成训练数据，**只微调 LoRA adapter** 做 test-time training。 +- 数据 / loss：verifier 驱动的自监督 SFT（仅 LoRA）。三 benchmark × 三 LLM，比 base 相对 +最高 32.29%、比"用 verifier 但不 TTT" +6.66%。 +- 与我们的区别：这是**"高分轨迹 → LoRA"这条主链的直接同构与背书**（verifier 打分选样 → 只训 LoRA），但它是**离线选样再训**、不是"推理中途干预"。我们 line A 的"用主线 `traj_score` 选高分轨迹蒸 LoRA"和它几乎一样，可当实现参考。 + +**Beyond Perplexity（TTT memory 审计）**（arXiv 2607.00368） +- 怎么做 / 结论：提出行为层评测框架，审计 TTT/memory 工作。发现一步 LoRA 更新能降 support/answer loss（跨 3 个 Qwen3 规模），但**自由 recall 仍为零** —— **proxy 指标改善 ≠ 部署行为改善**。 +- 与我们的区别：警示方案 ① 的技能 LoRA / 方案 ② 的评价能力，**别只看 loss 或 rubric 分下降就宣称有效**，要用**行为层对照**（带/不带、later recall / 下游动作）验证真收益（呼应 §4 影子对照、§10 DuoMem 的 CD 单独只 +1.4 的教训）。 + +### 11.5 双 LoRA 方案落地建议（综合上面证据） + +| 子命题 | 直接背书 | 反方 / 风险 | 建议 | +|---|---|---|---| +| ① 技能 LoRA（skill→LoRA） | S2L / LatentSkill / ParametricSkills / VDS-TTT | Beyond-Perplexity（proxy≠行为）；DuoMem CD 单独仅 +1.4 | **先做 S2L 式（一 skill 一 LoRA、离线 SFT、主线 traj_score 选样）**；必须带**带/不带对照**验证增益；skill 一多再上 hypernetwork（LatentSkill） | +| ② 评价 LoRA（rubric→judge） | Plug-and-Play Judge（rank-16 LoRA 顶 70B） | RGSD/RCSD：verifier 贵/不稳时**直接蒸进 policy 更优**；SCATR：轻量头比 LoRA 省 8000× | **先决策产物是不是"独立可复用的评价分"**：要 → 训评价 LoRA（复用 `score_lora_path`）；只为给 student dense 信号 → 走 RGSD 直蒸；只为排序/门控 → SCATR 轻量头 | +| ③ 每 N token 在线 PRM | Chain-of-LoRA（多 LoRA 热切原型）；MetaClaw（生产 RL-PRM） | 一次前向只挂一个 LoRA、无中途插入钩子；Surprisal/SCATR：搜索选样常更值；TTT 过锐化 | **最后做**；第一版用 `prompt_logprobs` **旁路打分 + 只留痕不干预**；确认 PRM 分与主线 hard verifier 一致性够高再考虑写交错解码调度器 | + +**对标基线**：**MetaClaw（有 code、OpenClaw 生产、双线镜像）**是整套方案最该研读与对标的工作；**Chain-of-LoRA** 是子命题 ③ 的机制原型；**RGSD** 是子命题 ② 决策的关键对照。我们相对它们的增量仍是 §10.10 那四条（显式毕业梯度、失败前置归因、抽取器被下游效用反向蒸馏、可回滚发布门控）——双 LoRA 只是把 line A 的"权重"具体化成"技能 LoRA + 评价能力"，不改变整体闭环定位。 + +### 11.6 定稿：查错 LoRA 作为参数化 memory —— 结论与实验方案 + +> 本节是 §11 讨论收敛后的**决策记录 + 实验计划**。核心转变：把"评价 LoRA"从"给数据打分的 judge"重新定位成**不动 base 的参数化 memory（line B 第二载体），专门给 base 补一个"在线查错"能力**。技能 LoRA（LoRA-2）**本轮暂缓**（理由见末尾）。 + +#### 11.6.1 定位（钉死的几条前提） + +1. **base 全程冻结，永不训。** 从根上杜绝知识遗忘——这是硬底线。所有能力增量都以**可插拔 LoRA**形式外挂，本质是 line B（参数化 memory），不是 line A。 +2. **LoRA-1 = 查错器（过程级 memory）。** 它不改 base 的知识，只给 base 补"边生成边发现自己踩了 rubric 里哪类错"的能力（如"公式第 k 步描述错""工具调用缺参数"）。发现错误后把**具体问题**注回 context，引导 base 改。 +3. **可回滚性 ≈ 删一条 memory。** 因为不动 base、LoRA 可插拔，一版 adapter 训坏了直接换/摘，撤销成本远低于"重训 base"——这消解了"权重侧犯错难撤"的顾虑（那顾虑只对"蒸进 base"成立）。 +4. **蒸馏数据 = 问题定位 + 打分（不是只给分）。** 要蒸出的是"**可定位、可操作的过程诊断**"能力，所以 teacher rubric 的产出必须到"哪一步/哪个 tool call/缺什么"这一粒度，而非单一 PASS/FAIL 标量（这是能否蒸出"查错"而非"打分习惯"的前提，对应 §10.7 原子化可验证条目 / 多角色 rubric）。 + +#### 11.6.2 自进化引擎：超大模型 + 本地 LoRA 双端检测 + +``` +超大模型(teacher rubric) ──检测出错误──┐ + ├─► 分歧/teacher 独有的错误 = 本地能力缺口 = 训练信号 ─► 升级 LoRA-1 +本地查错 LoRA-1 ──检测出错误──┘ │ + ▼ + 下一轮两端一起查 ──► 本地 LoRA-1 持续追平 teacher 的查错能力（追平后 teacher 少调、省成本） +``` + +- 与主线蒸馏同构：主线是"student 生成能力追平 teacher"，这里是"**本地 LoRA 的查错能力追平 teacher**"，同一套 `llm_backup` (teacher 对/本地错) 配对机制，被蒸对象换成"查错"这一角色。 +- **双端分歧同时是"该审上游"的探针**：本地 LoRA 与 teacher **系统性**分歧（成规律、非零星）时，要么本地没学到位（继续训），要么 **teacher/rubric 本身有系统性偏差**——后者是上游数据源/超大模型质量问题，**在上游治**（换更强 teacher、多角色 rubric 交叉、rubric 可靠性审计），不指望下游 LoRA 兜（呼应 §10.7"先坐实 verifier 再用它"）。 + +#### 11.6.3 三阶段实验路径（每阶段是下阶段的 gate，早止损） + +**阶段 0（先做）：不引入任何 LoRA，纯 teacher LLM 在线查错，测"上限"。** +- 数据集：**数学**（对错相对明确、查错信号干净，理想试验田）。 +- 机制：base 每隔 N 个 token **暂停** → 用 **teacher LLM（走 `llm_backup`）** 判当前生成有没有踩 rubric 错 → 有则把发现的问题注入 context → 继续生成。 +- 目的：**验证"在线过程级查错 + 注入 rubric"这个机制本身能否抬升数学解题上限**。用最强 teacher 代表能力天花板。 +- 性质：**纯可行性 gate，不训任何东西**。若最强 teacher 在线查错都提不了分，后面蒸 LoRA 更无意义 —— 先证信号有价值。 +- 工程：实验期"每 N token 暂停判断"**可用最粗暴实现**（停→跑一次 teacher 判断→拼 rubric→续），**不追性能**；生产化才需要专门的交错解码调度器（见待解决项）。 +- **已有实现**：`cookbook/exp/embedding/eval_dualline_math.py`（`--mode dualline`）。它复用 `eval_gpqa_rag.py` 的 MATH 分层加载 / 采样参数 / `answers_match` 判分，token 级分段生成：每 `--chunk-tokens` 暂停 → `RubricVerifier.diagnose()`（无 student sampler，走 llm_backup teacher）查错 → 命中且分数低于 `DUALLINE_CHECK_FLOOR` 则把 fix/原因作 `[Checker]` 注入续写。对照基线 `--mode baseline`（＝单遍生成，等同 `eval_gpqa_rag --mode direct`），同一 200 条子集直接比较 overall / 分层准确率与 checks/injections 计数。launch 配置：`dualline_math`。 + +**阶段 1（阶段 0 证明有效后）：把 teacher 的查错能力蒸成 LoRA-1，替换在线 teacher 调用。** +- 用阶段 0 收集的 (完整核验 CoT + 问题定位 + 打分) 数据蒸 LoRA-1（可复用 `score_lora_path` 入口）。**数据须正负配平**（成功段"无错、continue" + 失败段"定位/原因/建议"），构造约束见 11.6.6。 +- 目的：验证**本地低秩 LoRA 能否追平 teacher 查错**（＝待验证 b"学不学得动"）。实验要**把 rank × rubric 产出粒度当两个变量扫**（ParametricSkills 的 SHINE 退化反例说明：配置不对会打不过 in-context，学不学得动不是 0/1）。 + +**阶段 2（远期）：LoRA-2 技能 + 召回**——本轮暂缓，见 11.6.5。 + +#### 11.6.4 三条反驳 → 回应 → 限定（决策留痕） + +| 反驳 | 我们的回应 | 仍需守住的限定 | +|---|---|---| +| **① 廉价 query 路由只解决"该不该召回"，没解决"召回内容本身可能是错的"** | 不走"召回一条可能有错的 memory"，而是把"**查错能力**"蒸进 LoRA，让 base 内生地边做边纠错，**绕开召回可信度问题** | 前提：rubric 产出要到"**可定位错误**"粒度，否则只蒸出打分习惯、蒸不出查错 | +| **③ 权重侧坏数据难撤 / 会不会被污染** | LoRA 对**单条随机坏数据**抗性强于 RAG（梯度统计稀释 + base 低秩先验；RAG 一条即直入 context 无稀释）；且不动 base、可插拔，撤销≈删 memory | 抗的是**单条噪声**，**不抗系统性偏差**——系统性错误会被梯度强化。故 §4 归因质量仍是地基；系统性偏差归上游治（11.6.2） | +| **（Substrate Asymmetry）参数化 memory 在"该缺的要拒答"上惨败** | 承认：这是 LoRA "缺检不到信号"的固有短板，**不是污染问题** | **LoRA 只装可泛化行为/模式（如查错），易变事实仍留文本 memory**（§0 分工判据）；别用 LoRA 装事实 | + +#### 11.6.6 查错 LoRA 的训练数据构造（三条硬约束，钉死） + +> 复用现有 `RubricVerifier` 的 rubric + `llm_backup` 配对采集管线来攒数据，但现有 scoring prompt 刻意"只出 PASS/FAIL、不出解释"（省 token、便于 comparator 对齐），**直接拿来训会蒸出"打分习惯"而非"查错能力"**。故新增一个"诊断模式"产训练样本，须同时满足： + +1. **分数与原因必须一次调用同源产出（不可分两次采样拼）。** 若分数、原因来自两次独立采样，二者可能逻辑矛盾（判 FAIL 但原因说"没问题"），LoRA 学到错位映射。实现上：单次调用同时产 `per-criterion verdict + 每条 FAIL 的定位/原因/建议`，聚合分数由这批 verdict 算得，**原因与分数天然一致**。 + +2. **正负样本都要产、且要均衡（不能只在低分/被 gate 段跑诊断）。** 只见"错的"会把 LoRA 训成"逢查必报错"的挑刺偏置——它在线运行时每 N token 都硬报一个错，注入噪声反拖垮 base。必须让它见过大量"**看完一段、逐条核验、全部 PASS、结论=本段无过程错、无需干预**"的样本，才有能力在线输出关键的"OK，继续"信号。故诊断要**按比例覆盖高分成功段**，与失败段配平。 + +3. **CoT 要完整（核验过程，非只报结论）。** 正负两类样本的 target 都写成完整推理链： + - 成功段：`逐条核验 → 每条为何 PASS → 结论：无过程错误，continue` + - 失败段：`逐条核验 → 定位到第 k 条 FAIL → 错在哪 / 为什么 / 建议修法` + 完整 CoT 才对齐 LoRA-1 在线"边看前缀边判断"的实际形态；只给"这里错了"的结论会让它学不到核验过程，泛化差（呼应 §11.4 Beyond-Perplexity："proxy 改善 ≠ 部署行为改善"）。 + +> 落地路径（不改动服务打分/准入的现有 `RubricVerifier` 便宜路径）：新增诊断入口（`explain=True` 变体或独立 `_diagnose_once`，带 `@llm_backup` 自动落配对数据），schema：`input=[segment + rubric]`，`target=[完整核验 CoT + verdict + (FAIL 时) 定位/原因/建议]`。 +> +> **采样策略（已定）**： +> - **采集期全量存储、不做平衡、不设成本上限** —— 对每个 segment 都产完整诊断 CoT（正负都存），先把数据完整跑出来。正负配比、降采样等**留到训练采样阶段**再定，避免采集期过早丢信息。 +> - **注意训练/部署分布差**：离线诊断看的是"已完成的整段"，LoRA-1 在线看的是"生成中途的前缀"——离线攒数据能省掉阶段 0 的交错解码调度器工程，但**仍需一次"带/不带 LoRA-1 在线注入"的行为层对照**才算真验证（不能只看离线诊断准确率）。 + +#### 11.6.5 待验证 / 待解决 / 暂缓 + +- **待验证 a（已定方案）**：蒸 LoRA-1 时数据须含"问题定位 + 打分"，二者**一次调用同源产出**，且**正负样本都产、带完整核验 CoT**（详见 11.6.6 三条硬约束）。 +- **待验证 b**：低秩 LoRA 学不学得动查错 —— 阶段 1 跑实验，扫 rank × 信号粒度。 +- **待解决（生产化）**：vLLM **一次前向只挂一个 LoRA、无中途插入钩子**，"每 N token 暂停查错"生产化需自写**交错解码调度器**；实验阶段用粗暴实现绕过。 +- **待解决**：LoRA-1（旁路查错）与未来 LoRA-2（技能常驻）**同时在线的调度冲突**（单请求单 LoRA 限制）。 +- **暂缓：LoRA-2（优秀轨迹 skill 召回）** —— 本轮不做。理由：(1) 它**必须先训 LoRA 才能用，没有"纯 LLM 不训练"的验证捷径**，不适合当前"先用 LLM 测上限"的实验起点；(2) 训它需要**大量 trajectory** 作输入。两个已知前置问题记录备查：**P1 技能可能 per-user、跨用户不可迁移**（长期解＝只把 global 可泛化技能进 LoRA-2，用户专属留文本 memory，归 §0 分工；实验阶段忽略）；**P2 冷启动召回**（首个 query 可能是"你好"无信息量）——解法优先级：**A 延迟召回**（无信息量不召回，用 `intent`/信息量阈值判，"你好"本就不该召回）＞ B 滚动召回（每轮用累积上下文重判）＞ C 用 system/场景先验预挂。 diff --git a/src/twinkle_agentic/preprocessor/__init__.py b/src/twinkle_agentic/preprocessor/__init__.py index fb7b359c2..20b514b5d 100644 --- a/src/twinkle_agentic/preprocessor/__init__.py +++ b/src/twinkle_agentic/preprocessor/__init__.py @@ -24,10 +24,74 @@ from .structural_noise import StructuralNoiseTagger # noqa: F401 from .token_soup import TokenSoupFilter from .trajectory_scorer import TrajectoryScorer # noqa: F401 +from .value_selector import ValueSelector, select_top_for_rubric # noqa: F401 logger = get_logger() +def truncate_dropped_logs(dropped_log_path: str) -> None: + """Remove prior dropped log shards (call once from the main process before map).""" + if not dropped_log_path: + return + import glob + for p in [dropped_log_path] + glob.glob(f'{dropped_log_path}.*'): + if p.endswith('.lock'): + continue + try: + os.remove(p) + except FileNotFoundError: + pass + + +def merge_dropped_shards(dropped_log_path: str) -> None: + """Merge per-worker ``dropped.jsonl.`` shards into ``dropped.jsonl``.""" + if not dropped_log_path: + return + import glob + shards = sorted( + p for p in glob.glob(f'{dropped_log_path}.*') + if not p.endswith('.lock')) + if not shards: + return + os.makedirs(os.path.dirname(os.path.abspath(dropped_log_path)) or '.', exist_ok=True) + with open(dropped_log_path, 'w', encoding='utf-8') as out: + for sp in shards: + with open(sp, encoding='utf-8') as fin: + for line in fin: + if line.strip(): + out.write(line if line.endswith('\n') else line + '\n') + try: + os.remove(sp) + except FileNotFoundError: + pass + + +def run_quality_pipeline(dataset, pipeline: 'QualityPreprocessor', *, + num_proc: int = 1, **map_kwargs): + """Run a ``drop_mode='mark'`` pipeline as map(equal-length) + filter(keep). + + This is the ghost-proof way to run a filtering pipeline: ``map`` never + changes row count (every batch returns equal-length columns with a + ``_keep`` flag), then a single ``Dataset.filter`` on that flag does the + actual removal. Returns the dataset (mutated in place). + """ + if getattr(pipeline, '_drop_mode', None) != 'mark': + raise ValueError("run_quality_pipeline requires a pipeline built with drop_mode='mark'") + flag = QualityPreprocessor.KEEP_FLAG + map_kwargs.pop('remove_columns', None) # mark mode keeps row count; not needed + dataset.map(pipeline, num_proc=num_proc, **map_kwargs) + dataset.filter(lambda row: bool(row.get(flag, True))) + # Drop the transient keep-flag column so downstream schema stays clean. + hf = dataset.dataset + if flag in hf.column_names: + dataset.dataset = hf.remove_columns([flag]) + datasets = getattr(dataset, 'datasets', None) + if isinstance(datasets, dict) and len(datasets) == 1: + for k in list(datasets.keys()): + datasets[k] = dataset.dataset + return dataset + + class QualityPreprocessor(Preprocessor): """Thin pipeline runner: accepts a list of callables, runs them in order. @@ -35,19 +99,45 @@ class QualityPreprocessor(Preprocessor): Per-step logging (before/after count) and optional dropped-row JSONL are provided. """ - def __init__(self, pipeline: List[Callable], dropped_log_path: str = ''): + #: Column name for the keep flag emitted in ``drop_mode='mark'``. + KEEP_FLAG = '_keep' + + def __init__(self, pipeline: List[Callable], dropped_log_path: str = '', + drop_mode: str = 'inline'): super().__init__() + if drop_mode not in ('inline', 'mark'): + raise ValueError("drop_mode must be 'inline' or 'mark'") + # 'inline': the batch returns only surviving rows (shorter columns). HF + # then needs remove_columns to change row count cleanly, else ghost + # rows appear. Kept as the backward-compatible default. + # 'mark': the batch ALWAYS returns equal-length columns; dropped rows are + # returned too, flagged KEEP_FLAG=False (survivors True). No row-count + # change happens inside map, so no ghosting is possible. The caller + # materializes the drop with a follow-up ``Dataset.filter`` on KEEP_FLAG + # (see ``run_quality_pipeline``). + self._drop_mode = drop_mode self._pipelines = list(pipeline) self._dropped_log_path = dropped_log_path if dropped_log_path: - os.makedirs(os.path.dirname(os.path.abspath(dropped_log_path)), exist_ok=True) - self._lock: Optional[PosixFileLock] = (PosixFileLock(dropped_log_path + '.lock') if dropped_log_path else None) - if dropped_log_path and os.path.exists(dropped_log_path): - os.remove(dropped_log_path) + os.makedirs(os.path.dirname(os.path.abspath(dropped_log_path)) or '.', exist_ok=True) + lock_path = (dropped_log_path + '.lock') if dropped_log_path else '' + self._lock: Optional[PosixFileLock] = PosixFileLock(lock_path) if lock_path else None + # Truncation is explicit (see truncate_dropped_logs) so HF num_proc workers + # do not race to delete each other's shard files on unpickle/re-init. def __call__(self, rows): + input_col_keys = list(rows.keys()) if isinstance(rows, dict) else None rows_list = self.map_col_to_row(rows) total_start = len(rows_list) + # In 'mark' mode we must return every input row (equal-length columns), so + # remember each row's identity to reconcile survivors vs. dropped at the + # end. A per-batch position index is stable and needs no unique id, and a + # snapshot preserves dropped rows' original columns for re-emission. + original_rows = None + if self._drop_mode == 'mark': + original_rows = [dict(r) for r in rows_list] + for i, r in enumerate(rows_list): + r['_row_idx'] = i stats = [] for step in self._pipelines: if not rows_list: @@ -63,12 +153,65 @@ def __call__(self, rows): self._log_dropped(step_name, dropped) summary = '\n'.join(stats) logger.info(f'[QualityPreprocessor] {total_start} -> {len(rows_list)}\n{summary}') - return self.map_row_to_col(rows_list) + + if self._drop_mode == 'mark': + return self._emit_marked(rows_list, total_start, input_col_keys, original_rows) + # 'inline': HF ``datasets.map(batched=True)`` changes row count only when + # the batch returns shorter columns AND the caller passes remove_columns + # so the old columns are rebuilt (else survivors of a partially-filtered + # batch leave the un-dropped originals behind as ghost rows). Emitting an + # empty dict would also leave ghosts, so always emit explicit columns. + return self.map_row_to_col(rows_list, keys=input_col_keys) + + def _emit_marked(self, survivors, total_start, input_col_keys, original_rows): + """Return ALL input rows with equal-length columns, flagging survivors. + + Survivors carry ``KEEP_FLAG=True`` plus their tags; dropped rows are + re-emitted from their original input state with ``KEEP_FLAG=False`` so no + column ever changes length inside ``map`` (ghost-proof). The caller then + does a single ``Dataset.filter`` on ``KEEP_FLAG``. + """ + by_idx = {r.get('_row_idx'): r for r in survivors} + merged = [] + for i in range(total_start): + if i in by_idx: + row = by_idx[i] + row[self.KEEP_FLAG] = True + else: + # dropped: re-emit the original input row so its columns still + # exist (values are irrelevant — the caller filters it out). + row = dict(original_rows[i]) + row[self.KEEP_FLAG] = False + row.pop('_row_idx', None) + merged.append(row) + # Emit the UNION of every row's keys so a tag added only to survivors + # (e.g. `intent`) is present as a real column (None for dropped rows) — + # rows[0] alone is not enough (that is the original ghosting bug). + key_union: List[str] = list(input_col_keys or []) + for row in merged: + for k in row.keys(): + if k not in key_union: + key_union.append(k) + columns = {k: [row.get(k) for row in merged] for k in key_union} + return columns def _log_dropped(self, step_name: str, dropped: List[Dict[str, Any]]) -> None: if not self._lock or not dropped: return + shard = f'{self._dropped_log_path}.{os.getpid()}' with self._lock: - with open(self._dropped_log_path, 'a', encoding='utf-8') as f: + with open(shard, 'a', encoding='utf-8') as f: for r in dropped: - f.write(json.dumps({'step': step_name, 'row': r}, ensure_ascii=False, default=str) + '\n') + rec = self._compact_drop_record(step_name, r) + f.write(json.dumps(rec, ensure_ascii=False, default=str) + '\n') + + @staticmethod + def _compact_drop_record(step_name: str, row: Dict[str, Any]) -> Dict[str, Any]: + """Log metadata only — full messages are huge and break multiprocess merges.""" + return { + 'step': step_name, + 'reason': row.get('drop_reason') or step_name, + 'id': row.get('id'), + 'model_id': row.get('model_id'), + 'n_msgs': len(row.get('messages') or []), + } diff --git a/src/twinkle_agentic/preprocessor/dead_loop_filter.py b/src/twinkle_agentic/preprocessor/dead_loop_filter.py index 3d629723b..d2316f0a9 100644 --- a/src/twinkle_agentic/preprocessor/dead_loop_filter.py +++ b/src/twinkle_agentic/preprocessor/dead_loop_filter.py @@ -155,6 +155,7 @@ def __init__( think_hesitation_density_threshold: float = 15.0, think_cascade_threshold: int = 20, think_repetition_threshold: float = 0.65, + agent_min_stuck_turns: int = 2, ) -> None: super().__init__() # Two threshold profiles: laxer inside reasoning (free to ramble), @@ -175,6 +176,7 @@ def __init__( ngram_size=ngram_size, ngram_min_words=ngram_min_words, ) + self._agent_min_stuck_turns = max(1, int(agent_min_stuck_turns)) def _is_stuck(self, text: str, reasoning: str = '') -> bool: think_part, response_part = _split_think(text) @@ -204,11 +206,14 @@ def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: if not asst_msgs: out.append(row) continue - if any( - self._is_stuck( - msg_content_text(m).strip(), - (m.get('reasoning_content') or m.get('thinking') or '').strip(), - ) for m in asst_msgs): + stuck_turns = sum( + 1 for m in asst_msgs + if self._is_stuck( + msg_content_text(m).strip(), + (m.get('reasoning_content') or m.get('thinking') or '').strip(), + )) + min_stuck = self._agent_min_stuck_turns if agent else 1 + if stuck_turns >= min_stuck: dropped.append(dict(row, drop_reason='dead_loop')) else: out.append(row) diff --git a/src/twinkle_agentic/preprocessor/intent_classifier.py b/src/twinkle_agentic/preprocessor/intent_classifier.py index 7cb1dc5d9..079719ed7 100644 --- a/src/twinkle_agentic/preprocessor/intent_classifier.py +++ b/src/twinkle_agentic/preprocessor/intent_classifier.py @@ -201,10 +201,25 @@ class _RegexDetector(IntentDetector): """Common scaffolding: scan messages, run ``_match`` on each text, pair to assistant.""" role_filter: Optional[str] = None + # Whether ```` reasoning blocks are stripped before matching an + # assistant message. Content-signature detectors (code / math / logic) set + # this so scratch-pad markdown fences or LaTeX inside the model's private + # reasoning don't misclassify the task (e.g. a copywriting answer whose + # happens to contain a ``` fence being tagged as ``code``). User + # messages are never stripped — a code/latex request there is a real signal. + strip_think_in_assistant: bool = False def _match(self, text: str) -> bool: return False + def _text_for_match(self, role: str, m: dict) -> str: + text = msg_content_text(m) + if self.strip_think_in_assistant and role == 'assistant' and text: + # Keep only the visible response (pre-think + post-think), drop the + # ... scratch work that shouldn't define the task type. + text = _THINK_BLOCK_RE.sub(' ', text) + return text + def __call__(self, messages): rounds = set() for idx, m in enumerate(messages): @@ -217,7 +232,7 @@ def __call__(self, messages): continue if self.role_filter and role != self.role_filter: continue - text = msg_content_text(m) + text = self._text_for_match(role, m) if not text or not self._match(text): continue asst_idx = _pair_assistant(messages, idx, role) @@ -241,6 +256,7 @@ def __call__(self, messages): class CodeDetector(_RegexDetector): intent = INTENT_CODE + strip_think_in_assistant = True def __init__(self, threshold: int = 3) -> None: self.threshold = threshold @@ -254,6 +270,7 @@ def _match(self, text): class MathDetector(_RegexDetector): intent = INTENT_MATH + strip_think_in_assistant = True def __init__(self, threshold: int = 4) -> None: self.threshold = threshold @@ -265,6 +282,7 @@ def _match(self, text): class ComplexLogicDetector(_RegexDetector): intent = INTENT_COMPLEX_LOGIC role_filter = 'assistant' + strip_think_in_assistant = True def __init__(self, threshold: int = 6) -> None: self.threshold = threshold diff --git a/src/twinkle_agentic/preprocessor/label_schema.py b/src/twinkle_agentic/preprocessor/label_schema.py index 109f627d2..f9cc97938 100644 --- a/src/twinkle_agentic/preprocessor/label_schema.py +++ b/src/twinkle_agentic/preprocessor/label_schema.py @@ -56,6 +56,21 @@ # Free-form scoring metadata (short-circuit stats, per-check breakdown, etc.). KEY_SCORE_META = 'score_meta' +# Active-learning pre-selection (ValueSelector): a cheap, LLM-free "how worth an +# expensive rubric pass is this row" score in [0, 1], its per-component +# breakdown, and the boolean gate the rubric stage reads to decide whether to +# spend an LLM call on this row (top-fraction by value_score). +KEY_VALUE_SCORE = 'value_score' +KEY_VALUE_META = 'value_meta' +KEY_SELECTED_FOR_RUBRIC = 'selected_for_rubric' + +# Persisted rubric diagnosis for rubric-scored rows: a per-segment verification +# chain (rubric text + per-criterion verdict/reason/fix + raw model output + +# query/segment_text). This is the SFT corpus for distilling a PRM / error-checker +# LoRA — store it so training never has to re-run the (expensive) teacher. +# Value: List[dict], one entry per rubric-scored segment (see TrajectoryScorer). +KEY_RUBRIC_DIAGNOSIS = 'rubric_diagnosis' + # --------------------------------------------------------------------------- # thin get / set helpers over the (key, pack_value) envelope diff --git a/src/twinkle_agentic/preprocessor/language_filter.py b/src/twinkle_agentic/preprocessor/language_filter.py index fc68d90b6..ec0c7fac9 100644 --- a/src/twinkle_agentic/preprocessor/language_filter.py +++ b/src/twinkle_agentic/preprocessor/language_filter.py @@ -107,3 +107,8 @@ def keep(self, row: Dict[str, Any]) -> bool: if lang is None: return self.keep_undetected return lang.lower() in self.allowed + + def drop_reason(self, row: Dict[str, Any]) -> str: + text = self._user_text(row) + lang = self._detect(text) if len(text) >= self.min_chars else None + return f'language_{lang or "undetected"}' diff --git a/src/twinkle_agentic/preprocessor/model_filter.py b/src/twinkle_agentic/preprocessor/model_filter.py index f651f7905..4162c54c5 100644 --- a/src/twinkle_agentic/preprocessor/model_filter.py +++ b/src/twinkle_agentic/preprocessor/model_filter.py @@ -5,6 +5,17 @@ # Each entry is the discriminating prefix only; a shared variant tail is appended uniformly # so suffixes like -Instruct, -Thinking-2507, -Distill-Qwen-7B, -Air are accepted everywhere. +# +# DESIGN INTENT — only large text models are admitted (this is deliberate, not a bug): +# * Sub-100B models are excluded on purpose. The size gate is baked into the +# patterns: e.g. ``-[123]\d{2}b`` requires a 3-digit "1xx/2xx/3xxB" family, so +# 27B / 35B / 8B variants intentionally FAIL to match and get dropped. +# * Vision-Language (VL / multimodal) models are excluded on purpose. There is no +# VL entry in the allow-list, so e.g. ``Qwen3-VL-*`` is dropped even at 235B. +# Consequence: on mixed dumps a large share of rows (small + VL models) land in +# ``dropped.jsonl`` with reason ``model_not_allowed`` — expected by design. +# Revisit here (add a pattern / relax the size digits) only when we decide to +# start distilling from small or multimodal teachers. _DEFAULT_PATTERNS = [ r'minimax/minimax-m[23][\d.]*', r'opengvlab/internvl[\d._]+-2\d{2}b', @@ -31,3 +42,6 @@ def __init__(self, patterns: Optional[Sequence[str]] = None, field: str = 'model def keep(self, row: Dict[str, Any]) -> bool: return bool(self._re.fullmatch(row.get(self._field) or '')) + + def drop_reason(self, row: Dict[str, Any]) -> str: + return 'model_not_allowed' diff --git a/src/twinkle_agentic/preprocessor/safety_scorer.py b/src/twinkle_agentic/preprocessor/safety_scorer.py index e89963529..4cbe7142f 100644 --- a/src/twinkle_agentic/preprocessor/safety_scorer.py +++ b/src/twinkle_agentic/preprocessor/safety_scorer.py @@ -51,6 +51,7 @@ def __init__( *, criteria: Optional[Tuple[str, ...]] = None, unsafe_threshold: float = 0.5, + gate_label: Optional[str] = None, ): from twinkle_agentic.verifier import RubricItem, RubricVerifier @@ -62,6 +63,10 @@ def __init__( else: rubric_verifier.fixed_rubric = fixed self.verifier = rubric_verifier + # Active-learning gate: when set, only rows whose ``gate_label`` is True + # spend an LLM safety pass. The rest are tagged as neutral-safe (the LLM + # safety check runs post-selection only). None -> score every row. + self.gate_label = gate_label def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: rows = self.map_col_to_row(rows) @@ -78,6 +83,12 @@ def _score_row(self, row: Dict[str, Any]) -> Dict[str, Any]: messages = row.get('messages') if not isinstance(messages, list) or not messages: return row + # Gated out (not selected for the LLM pass): tag neutral-safe, no LLM call. + if self.gate_label and L.get_label(row, self.gate_label, None) is False: + return L.set_labels(row, { + L.KEY_SAFETY_SCORE: 1.0, + L.KEY_SAFETY_UNSAFE: False, + }) trajectory = {'messages': messages} if row.get('tools'): trajectory['tools'] = row['tools'] diff --git a/src/twinkle_agentic/preprocessor/trajectory_scorer.py b/src/twinkle_agentic/preprocessor/trajectory_scorer.py index 421a53197..1310410ce 100644 --- a/src/twinkle_agentic/preprocessor/trajectory_scorer.py +++ b/src/twinkle_agentic/preprocessor/trajectory_scorer.py @@ -27,6 +27,8 @@ """ from __future__ import annotations +import os +from concurrent.futures import ThreadPoolExecutor from typing import Any, Dict, List, Optional, Tuple from twinkle.preprocessor import Preprocessor @@ -61,7 +63,7 @@ def __init__( rubric_verifier: Optional[Any] = None, *, hard_agg: str = 'gmean', - fusion: str = 'product', + fusion: str = 'hard_soft_blend', hard_floor: float = 0.25, hard_ceil_skip: Optional[float] = None, traj_agg: str = 'mean', @@ -69,6 +71,11 @@ def __init__( write_round_detail: bool = False, calibrate: bool = True, disagree_margin: float = 0.34, + reconcile_max_messages: int = 80, + scorer_workers: Optional[int] = None, + intent_aware: bool = True, + rubric_gate_label: Optional[str] = L.KEY_SELECTED_FOR_RUBRIC, + persist_diagnosis: bool = False, ): # Lazy imports keep the module importable even if verifier/segment deps # are heavy; construction still fails loudly if the packages are absent. @@ -78,6 +85,19 @@ def __init__( self.segmenter = segmenter if segmenter is not None else TurnSegmenter('cluster') self.hard_scorer = hard_scorer if hard_scorer is not None else HardScorer() self.rubric_verifier = rubric_verifier + # Route each segment's rubric by its structural intent (tool_call/code/ + # math) so the verifier can apply intent-keyed fixed/half-fixed rubrics. + self.intent_aware = bool(intent_aware) + self._intent_detectors = None + # Active-learning gate: when set, only rows whose ``rubric_gate_label`` is + # True spend an LLM rubric pass; the rest are scored hard-only. Leaving it + # None (or the label absent) preserves the "rubric every row" behavior. + self.rubric_gate_label = rubric_gate_label + # When True, rubric-scored segments also emit a full DiagnoseDetail + # (per-criterion verdict + reason + fix + raw teacher output). Persisted + # under KEY_RUBRIC_DIAGNOSIS as the SFT corpus for a distilled PRM/checker + # LoRA. Costs one extra teacher call per scored segment. + self.persist_diagnosis = bool(persist_diagnosis) self.hard_agg = hard_agg self.fusion = fusion self.hard_floor = float(hard_floor) @@ -88,17 +108,26 @@ def __init__( # D7c: self-evolving calibration (no human alignment). self.calibrate = bool(calibrate) self.disagree_margin = float(disagree_margin) + self.reconcile_max_messages = int(reconcile_max_messages) + if scorer_workers is None: + scorer_workers = int(os.environ.get('TRAJ_SCORER_WORKERS', '1')) + self.scorer_workers = max(1, int(scorer_workers)) + + def _score_row_safe(self, row: Dict[str, Any]) -> Dict[str, Any]: + try: + return self._score_row(row) + except Exception as e: + logger.warning(f'[TrajectoryScorer] scoring failed, row left unscored: {e}') + return row # ------------------------------------------------------------------ def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: rows = self.map_col_to_row(rows) - out: List[Dict[str, Any]] = [] - for row in rows: - try: - out.append(self._score_row(row)) - except Exception as e: # scoring must never break the pipeline - logger.warning(f'[TrajectoryScorer] scoring failed, row left unscored: {e}') - out.append(row) + if self.scorer_workers <= 1 or len(rows) <= 1: + out = [self._score_row_safe(row) for row in rows] + else: + with ThreadPoolExecutor(max_workers=self.scorer_workers) as pool: + out = list(pool.map(self._score_row_safe, rows)) return out, [] # mapper: never drops # ------------------------------------------------------------------ @@ -119,12 +148,22 @@ def _score_row(self, row: Dict[str, Any]) -> Dict[str, Any]: if not segments: return row + # Active-learning gate: skip the LLM rubric for rows not pre-selected by + # ValueSelector (hard-only), so expensive labeling is spent on the top + # fraction only. Absent label -> treat as selected (backward compatible). + rubric_enabled = self.rubric_verifier is not None + if rubric_enabled and self.rubric_gate_label: + selected = L.get_label(row, self.rubric_gate_label, None) + if selected is False: + rubric_enabled = False + query = self._infer_query(messages) all_round_scalars: List[float] = [] all_round_gated: List[bool] = [] segment_scalars: List[float] = [] segment_confidence: List[float] = [] segment_scores = [] + diagnoses: List[Dict[str, Any]] = [] for s_idx, segment in enumerate(segments): rounds = split_segment_into_rounds(segment) @@ -140,7 +179,12 @@ def _score_row(self, row: Dict[str, Any]) -> Dict[str, Any]: all_round_scalars.append(detail.scalar) all_round_gated.append(detail.gated) - rubric_fn = self._make_rubric_fn(segment, query, round_scores) + seg_intent = None + if rubric_enabled: + seg_intent = self._segment_intent(segment) if self.intent_aware else None + rubric_fn = self._make_rubric_fn(segment, query, round_scores, seg_intent) + else: + rubric_fn = None seg_score = fuse_segment( s_idx, round_scores, rubric_fn, hard_agg=self.hard_agg, fusion=self.fusion, @@ -152,6 +196,14 @@ def _score_row(self, row: Dict[str, Any]) -> Dict[str, Any]: segment_scalars.append(seg_score.scalar) segment_confidence.append(self._segment_confidence(seg_score)) + # Persist a full diagnostic chain for segments that actually reached + # the LLM (not hard-only / short-circuited) — the PRM/checker SFT data. + if (self.persist_diagnosis and rubric_enabled + and not seg_score.short_circuited): + diag = self._diagnose_segment(segment, query, seg_intent, s_idx) + if diag is not None: + diagnoses.append(diag) + traj = aggregate_trajectory( segment_scores, how=self.traj_agg, weight_by_rounds=self.weight_by_rounds) @@ -176,10 +228,77 @@ def _score_row(self, row: Dict[str, Any]) -> Dict[str, Any]: for s in segment_scores ], } + if self.persist_diagnosis and diagnoses: + labels[L.KEY_RUBRIC_DIAGNOSIS] = diagnoses return L.set_labels(row, labels) + def _diagnose_segment(self, segment: dict, query: str, intent, s_idx: int): + """Run the verifier's full diagnosis and pack it for persistence. + + Emits everything a distilled PRM/checker LoRA needs: per-criterion + verdict + reason + fix, the overall verdict, the rubric text, and the + raw teacher output (SFT target) alongside the query + segment text + (SFT inputs). Never raises — diagnosis is best-effort enrichment. + """ + rv = self.rubric_verifier + if rv is None or not hasattr(rv, 'diagnose'): + return None + try: + d = rv.diagnose(segment, query=query, intent=intent) + except Exception as e: + logger.warning(f'[TrajectoryScorer] diagnose failed (seg {s_idx}): {e}') + return None + if d is None: + return None + return { + 'segment_index': s_idx, + 'intent': intent, + 'scalar': round(float(getattr(d, 'scalar', 0.0)), 6), + 'overall_ok': bool(getattr(d, 'overall_ok', False)), + 'summary': getattr(d, 'summary', '') or '', + 'query': getattr(d, 'query', '') or query, + 'segment_text': getattr(d, 'segment_text', '') or '', + 'raw': getattr(d, 'raw', '') or '', + 'rubric': [ + {'text': it.text, 'is_hard': bool(getattr(it, 'is_hard', False))} + for it in (getattr(d, 'rubric', None) or []) + ], + 'items': [ + { + 'index': it.index, + 'verdict': bool(it.verdict), + 'reason': it.reason or '', + 'fix': it.fix or '', + } + for it in (getattr(d, 'items', None) or []) + ], + } + # ------------------------------------------------------------------ - def _make_rubric_fn(self, segment: dict, query: str, round_scores): + def _segment_intent(self, segment: dict): + """Classify a segment by structural intent (tool_call > code > math). + + Reuses the lightweight, LLM-free detectors from IntentClassifier so the + segment's rubric can be routed to an intent-keyed fixed/half-fixed rubric. + Returns an intent string or ``None`` (no confident match -> generate). + """ + if self._intent_detectors is None: + from .intent_classifier import (CodeDetector, MathDetector, + ToolCallDetector) + # Order matters: tool_call is the strongest structural signal. + self._intent_detectors = [ToolCallDetector(), CodeDetector(), MathDetector()] + messages = segment.get('messages') or [] + if not isinstance(messages, list) or not messages: + return None + for det in self._intent_detectors: + try: + if det(messages): + return det.intent + except Exception: + continue + return None + + def _make_rubric_fn(self, segment: dict, query: str, round_scores, intent=None): """Return a zero-arg callable for the soft chain, or None if unavailable. ``fuse_segment`` only invokes this when the segment is NOT short-circuited, @@ -198,13 +317,16 @@ def _make_rubric_fn(self, segment: dict, query: str, round_scores): hard_agg_val = aggregate_hard_over_rounds(round_scores, how=self.hard_agg) def _fn(): - detail = rv.score_detail(segment, query=query) - if self.calibrate and detail is not None and getattr(detail, 'scalar', None) is not None: + detail = rv.score_detail(segment, query=query, intent=intent) + if (self.calibrate and detail is not None + and getattr(detail, 'scalar', None) is not None + and len(segment.get('messages') or []) <= self.reconcile_max_messages): if abs(detail.scalar - hard_agg_val) >= self.disagree_margin: evidence = (f'Deterministic tool/answer checks scored this ' f'segment {hard_agg_val:.2f} out of 1.0. Reconcile ' f'your assessment with this objective evidence.') - revised = rv.score_detail(segment, query=query, extra_context=evidence) + revised = rv.score_detail(segment, query=query, + extra_context=evidence, intent=intent) if revised is not None: detail = revised segment['_last_rubric'] = detail diff --git a/src/twinkle_agentic/preprocessor/value_selector.py b/src/twinkle_agentic/preprocessor/value_selector.py new file mode 100644 index 000000000..aeef78ccf --- /dev/null +++ b/src/twinkle_agentic/preprocessor/value_selector.py @@ -0,0 +1,288 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Active-learning pre-selection for the expensive rubric pass. + +Motivation +---------- +At high daily volume it is neither affordable to LLM-label every trajectory nor +smart to sample at random (most rows are unremarkable). This mapper assigns each +row a **cheap, fully deterministic** ``value_score`` — an estimate of how much an +expensive rubric/LLM pass would *learn* from it — so a downstream gate can send +only the top fraction to the LLM. Self-evolution needs few, well-chosen samples. + +The score is a weighted blend of three LLM-free signals (all in ``[0, 1]``): + +- **uncertainty** — how close the deterministic hard signal is to undecided. + A row the hard checks already call clearly good (all 1.0) or clearly bad + (all 0.0) teaches the LLM little; rows near the boundary, or with internal + disagreement across rounds, are where a rubric pass pays off most. +- **difficulty** — structural complexity (rounds, tool calls, distinct tools, + segments), log-compressed so a few giant traces don't dominate. Long agentic + traces carry more signal than single-turn chit-chat. +- **error_signal** — deterministic failure evidence (gated rounds, failed + tool execution / termination / repetition checks). Mistakes are valuable + learning material for self-evolution (negative / correction examples). + +Two-pass usage (see ``TrajectoryScorer`` gate) +---------------------------------------------- +1. Run this mapper over the full stream (parallel, no global state) to stamp + ``value_score`` on every row. +2. After ``map`` completes, in the *single* driver process call + :func:`select_top_for_rubric` to flip ``selected_for_rubric=True`` on the + global top ``select_frac``. Only those rows spend an LLM call. +""" +from __future__ import annotations + +import math +from typing import Any, Dict, List, Optional, Tuple + +from twinkle.preprocessor import Preprocessor +from twinkle.utils import get_logger + +from . import label_schema as L +from .utils import normalize_tool_calls + +logger = get_logger() + + +def _log_norm(x: float, cap: float) -> float: + """Log-compress a count into [0, 1], saturating at ``cap``.""" + if x <= 0: + return 0.0 + return min(1.0, math.log1p(x) / math.log1p(cap)) + + +def _mean(xs: List[float]) -> float: + return sum(xs) / len(xs) if xs else 0.0 + + +def _pstdev(xs: List[float]) -> float: + if len(xs) < 2: + return 0.0 + m = _mean(xs) + return math.sqrt(sum((x - m) ** 2 for x in xs) / len(xs)) + + +class ValueSelector(Preprocessor): + """Stamp a deterministic ``value_score`` on every row (never drops). + + Args: + hard_scorer: a :class:`~twinkle_agentic.verifier.HardScorer` (reused for + the per-round hard scalars that feed ``uncertainty`` / ``error``). + Defaults to a plain ``HardScorer()``. No LLM is ever called. + segmenter: a segmenter for round splitting. Defaults to + ``TurnSegmenter('cluster')`` (LLM-free, same as TrajectoryScorer). + w_uncertainty / w_difficulty / w_error: blend weights (need not sum to 1; + normalized internally). + rounds_cap / toolcalls_cap / tools_cap / segments_cap: saturation caps + for the difficulty sub-signals. + write_meta: also store the per-component breakdown under ``value_meta``. + """ + + def __init__( + self, + hard_scorer: Optional[Any] = None, + segmenter: Optional[Any] = None, + *, + w_uncertainty: float = 0.45, + w_difficulty: float = 0.30, + w_error: float = 0.25, + rounds_cap: int = 20, + toolcalls_cap: int = 15, + tools_cap: int = 6, + segments_cap: int = 8, + write_meta: bool = True, + ): + from twinkle_agentic.segment import TurnSegmenter + from twinkle_agentic.verifier import HardScorer + + self.hard_scorer = hard_scorer if hard_scorer is not None else HardScorer() + self.segmenter = segmenter if segmenter is not None else TurnSegmenter('cluster') + total = w_uncertainty + w_difficulty + w_error + if total <= 0: + raise ValueError('at least one value weight must be > 0') + self.w_uncertainty = w_uncertainty / total + self.w_difficulty = w_difficulty / total + self.w_error = w_error / total + self.rounds_cap = int(rounds_cap) + self.toolcalls_cap = int(toolcalls_cap) + self.tools_cap = int(tools_cap) + self.segments_cap = int(segments_cap) + self.write_meta = bool(write_meta) + + # ------------------------------------------------------------------ + def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + rows = self.map_col_to_row(rows) + out = [] + for row in rows: + try: + out.append(self._score_row(row)) + except Exception as e: # never break the pipeline on a bad row + logger.warning(f'[ValueSelector] scoring failed, value=0: {e}') + out.append(L.set_label(row, L.KEY_VALUE_SCORE, 0.0)) + return out, [] # mapper: never drops + + # ------------------------------------------------------------------ + def _score_row(self, row: Dict[str, Any]) -> Dict[str, Any]: + from twinkle_agentic.verifier import split_segment_into_rounds + + messages = row.get('messages') + if not isinstance(messages, list) or not messages: + return L.set_label(row, L.KEY_VALUE_SCORE, 0.0) + + trajectory = {'messages': messages} + if row.get('tools'): + trajectory['tools'] = row['tools'] + + segments = self.segmenter.segment(trajectory) or [] + + round_scalars: List[float] = [] + any_gated = False + soft_fail = 0.0 # worst non-critical hard-check miss across rounds + for segment in segments: + for rnd in split_segment_into_rounds(segment): + detail = self.hard_scorer.score_detail(rnd) + round_scalars.append(detail.scalar) + if detail.gated: + any_gated = True + soft_fail = max(soft_fail, self._soft_fail(detail)) + + uncertainty = self._uncertainty(round_scalars) + difficulty = self._difficulty(messages, segments) + error_signal = self._error_signal(any_gated, soft_fail) + + value = (self.w_uncertainty * uncertainty + + self.w_difficulty * difficulty + + self.w_error * error_signal) + value = max(0.0, min(1.0, value)) + + updates: Dict[str, Any] = {L.KEY_VALUE_SCORE: round(value, 6)} + if self.write_meta: + updates[L.KEY_VALUE_META] = { + 'uncertainty': round(uncertainty, 4), + 'difficulty': round(difficulty, 4), + 'error': round(error_signal, 4), + 'n_rounds': len(round_scalars), + } + return L.set_labels(row, updates) + + # ------------------------------------------------------------------ + # signal components + # ------------------------------------------------------------------ + @staticmethod + def _uncertainty(round_scalars: List[float]) -> float: + """High when the hard signal is undecided OR rounds disagree.""" + if not round_scalars: + return 0.0 + mean_hard = _mean(round_scalars) + central = 1.0 - abs(2.0 * mean_hard - 1.0) # peak at 0.5 + disagreement = min(1.0, 2.0 * _pstdev(round_scalars)) # spread across rounds + return max(central, disagreement) + + def _difficulty(self, messages: List[dict], segments: List[dict]) -> float: + n_rounds = sum(1 for m in messages + if isinstance(m, dict) and m.get('role') == 'assistant') + n_toolcalls = 0 + tool_names = set() + for m in messages: + if not isinstance(m, dict) or m.get('role') != 'assistant': + continue + for tc in (normalize_tool_calls(m) or []): + n_toolcalls += 1 + fn = (tc.get('function') or {}) if isinstance(tc, dict) else {} + name = fn.get('name') if isinstance(fn, dict) else None + if name: + tool_names.add(name) + n_segments = len(segments) + return _mean([ + _log_norm(n_rounds, self.rounds_cap), + _log_norm(n_toolcalls, self.toolcalls_cap), + _log_norm(len(tool_names), self.tools_cap), + _log_norm(n_segments, self.segments_cap), + ]) + + @staticmethod + def _soft_fail(detail: Any) -> float: + """Worst miss among informative non-critical checks in a round (0..1).""" + watch = {'tool_executed', 'clean_termination', 'no_repeated_calls', + 'protocol_pairing', 'final_answer'} + worst = 0.0 + for c in getattr(detail, 'checks', None) or []: + if getattr(c, 'name', None) in watch: + worst = max(worst, 1.0 - float(getattr(c, 'score', 1.0))) + return worst + + @staticmethod + def _error_signal(any_gated: bool, soft_fail: float) -> float: + """Deterministic evidence the model made a mistake worth studying.""" + if any_gated: + return 1.0 + return soft_fail + + +# --------------------------------------------------------------------------- +# global top-fraction selection (driver process, after Dataset.map) +# --------------------------------------------------------------------------- +def select_top_for_rubric( + dataset, + *, + select_frac: float = 0.1, + min_select: int = 0, + max_select: Optional[int] = None, + value_key: str = L.KEY_VALUE_SCORE, + selected_key: str = L.KEY_SELECTED_FOR_RUBRIC, +): + """Flip ``selected_for_rubric`` on the global top-``select_frac`` by value. + + Must run in the single driver process AFTER ``Dataset.map`` (the top + fraction is a global order that per-shard workers cannot compute). Returns + ``(dataset, n_selected)``; the dataset is mutated via a lightweight map. + + Ties at the cutoff are all included (selection is by a value threshold), so + the realized count can slightly exceed ``select_frac * N``. + """ + hf = dataset.dataset + n = len(hf) + if n == 0: + return dataset, 0 + + def _val(row) -> float: + v = L.get_label(row, value_key, 0.0) + try: + return float(v) + except (TypeError, ValueError): + return 0.0 + + values = sorted((_val(hf[i]) for i in range(n)), reverse=True) + k = int(round(select_frac * n)) + if min_select: + k = max(k, min_select) + if max_select is not None: + k = min(k, max_select) + k = max(0, min(k, n)) + if k == 0: + threshold = float('inf') + else: + threshold = values[k - 1] + + def _mark(batch): + rows = Preprocessor.map_col_to_row(batch) + out = [L.set_label(r, selected_key, _val(r) >= threshold) for r in rows] + return Preprocessor.map_row_to_col(out, keys=list(batch.keys())) + + marked = hf.map(_mark, batched=True, load_from_cache_file=False, + remove_columns=list(hf.column_names)) + # Write back to BOTH views so the next Dataset.map sees the marks: twinkle's + # Dataset.map operates on self.datasets[key] (not self.dataset), so updating + # only self.dataset would silently drop selected_for_rubric before pass 2. + dataset.dataset = marked + datasets = getattr(dataset, 'datasets', None) + if isinstance(datasets, dict): + for k in list(datasets.keys()): + if datasets[k] is hf or len(datasets) == 1: + datasets[k] = marked + + n_selected = sum(1 for i in range(len(marked)) + if L.get_label(marked[i], selected_key, False)) + logger.info(f'[ValueSelector] selected {n_selected}/{n} rows for rubric ' + f'(frac={select_frac}, threshold={threshold:.4f})') + return dataset, n_selected diff --git a/src/twinkle_agentic/verifier/__init__.py b/src/twinkle_agentic/verifier/__init__.py index 98c7d13ed..0cdd24bc2 100644 --- a/src/twinkle_agentic/verifier/__init__.py +++ b/src/twinkle_agentic/verifier/__init__.py @@ -8,11 +8,16 @@ check_numeric_equiv, check_output_format, default_checks_for) from .hard_scorer import CheckResult, HardScorer, HardScoreDetail, TrajectoryView -from .rubric_verifier import RubricItem, RubricVerifier, ScoreDetail +from .rubric_library import (INTENT_BASE_RUBRICS, INTENT_FIXED_RUBRICS, + default_intent_base_rubrics, + default_intent_fixed_rubrics) +from .rubric_verifier import (DiagnoseDetail, DiagnosisItem, RubricItem, + RubricVerifier, ScoreDetail) __all__ = [ 'Verifier', 'RubricVerifier', 'RubricItem', 'ScoreDetail', + 'DiagnoseDetail', 'DiagnosisItem', 'HardScorer', 'HardScoreDetail', 'CheckResult', 'TrajectoryView', 'check_output_format', 'check_numeric_equiv', 'check_answer_match', 'check_code_parses', 'check_instruction_constraints', 'check_not_degenerate', @@ -20,4 +25,6 @@ 'RoundScore', 'SegmentScore', 'TrajectoryScore', 'split_segment_into_rounds', 'aggregate_hard_over_rounds', 'fuse_segment', 'aggregate_trajectory', 'scalar_to_level', + 'INTENT_BASE_RUBRICS', 'INTENT_FIXED_RUBRICS', + 'default_intent_base_rubrics', 'default_intent_fixed_rubrics', ] diff --git a/src/twinkle_agentic/verifier/aggregation.py b/src/twinkle_agentic/verifier/aggregation.py index 9ab4b4cb1..c791ccb40 100644 --- a/src/twinkle_agentic/verifier/aggregation.py +++ b/src/twinkle_agentic/verifier/aggregation.py @@ -165,6 +165,8 @@ def fuse_segment( hard_agg: reducer for per-round hard scores ('gmean'|'mean'|'min'|...). fusion: how to combine hard_agg and rubric: 'product' -> hard_agg * rubric (hard acts as a floor/gatekeeper) + 'hard_soft_blend' -> product, but when hard is high blend rubric toward + a floor so all-pass tool segments are not one-shot vetoed 'min' -> min(hard_agg, rubric) 'mean' -> (hard_agg + rubric)/2 'hard_only'-> ignore rubric entirely @@ -218,6 +220,16 @@ def _rubric_scalar(result: Any) -> float: def _combine(hard: float, soft: float, fusion: str) -> float: if fusion == 'product': return hard * soft + if fusion == 'hard_soft_blend': + # When hard checks are strong (tool/format all pass), a harsh rubric on a + # long agent trace must not one-shot veto the segment (product → ~0.08). + if hard >= 0.9: + mix = 0.55 * soft + 0.45 + elif hard >= 0.75: + mix = 0.75 * soft + 0.25 + else: + mix = soft + return hard * mix if fusion == 'min': return min(hard, soft) if fusion == 'mean': diff --git a/src/twinkle_agentic/verifier/hard_scorer.py b/src/twinkle_agentic/verifier/hard_scorer.py index 80c69e315..bc4dca126 100644 --- a/src/twinkle_agentic/verifier/hard_scorer.py +++ b/src/twinkle_agentic/verifier/hard_scorer.py @@ -272,13 +272,24 @@ def check_protocol_pairing(view: TrajectoryView) -> CheckResult: def check_no_repeated_calls(view: TrajectoryView) -> CheckResult: - """Penalize exact-duplicate (name, args) calls (dead-loop / redundancy).""" + """Penalize degenerate tool-call loops. + + Two independent signals, worst one wins: + 1. exact-duplicate ``(name, args)`` calls — classic redundant repetition; + 2. single-tool domination — one tool name fired over and over (even with + *different* args), the "spin the same tool forever" failure that (1) + misses because the arguments differ each time. Only kicks in once there + are enough calls (``>= _SPIN_MIN_CALLS``) so a legitimate 3-4 step loop + of the same tool is not punished. + """ calls = view.tool_calls if len(calls) < 2: return CheckResult('no_repeated_calls', 1.0, 1.0, critical=False, n=len(calls), detail='fewer than 2 calls') + seen: set = set() dupes = 0 + name_counts: Dict[str, int] = {} for tc in calls: args = view.parsed_args(tc) key = (tc.name, json.dumps(args, sort_keys=True) if isinstance(args, dict) else str(tc.raw_args)) @@ -286,9 +297,33 @@ def check_no_repeated_calls(view: TrajectoryView) -> CheckResult: dupes += 1 else: seen.add(key) - score = 1.0 - dupes / len(calls) + name_counts[tc.name or ''] = name_counts.get(tc.name or '', 0) + 1 + dup_score = 1.0 - dupes / len(calls) + + # Single-tool domination: one tool fired over and over (a spin loop). This + # is only a *soft* signal — a legitimate agent may batch-read 8 files with + # the same tool — so it is deliberately lenient: it only triggers on long + # sequences that are almost entirely one tool, and it floors the penalty so + # a batch operation is nudged down, not zeroed. Real dead-loops (empty + # repeated spins) get further penalized by the rubric / final-answer checks. + _SPIN_MIN_CALLS = 8 + _SPIN_TOLERATED_SHARE = 0.8 + _SPIN_FLOOR = 0.4 + spin_score = 1.0 + top_name, top_n = max(name_counts.items(), key=lambda kv: kv[1]) + top_share = top_n / len(calls) + if len(calls) >= _SPIN_MIN_CALLS and top_share > _SPIN_TOLERATED_SHARE: + # Linearly map (tolerated..1.0] share onto (1.0.._SPIN_FLOOR] score. + frac = (top_share - _SPIN_TOLERATED_SHARE) / (1.0 - _SPIN_TOLERATED_SHARE) + spin_score = max(_SPIN_FLOOR, 1.0 - (1.0 - _SPIN_FLOOR) * frac) + + score = min(dup_score, spin_score) + detail = f'{dupes} duplicate calls' + if spin_score < dup_score: + detail = (f"tool '{top_name}' dominates {top_n}/{len(calls)} " + f'calls ({top_share:.0%})') return CheckResult('no_repeated_calls', score, 1.0, critical=False, - n=len(calls), detail=f'{dupes} duplicate calls') + n=len(calls), detail=detail) def check_clean_termination(view: TrajectoryView) -> CheckResult: diff --git a/src/twinkle_agentic/verifier/rubric_library.py b/src/twinkle_agentic/verifier/rubric_library.py new file mode 100644 index 000000000..8f5f753c0 --- /dev/null +++ b/src/twinkle_agentic/verifier/rubric_library.py @@ -0,0 +1,108 @@ +"""Intent-keyed rubric library (DESIGN follow-up: stabilize rubric scoring). + +Rubric *generation* is flexible but high-variance: for template-like intents +(tool_call / code / math) the model re-invents slightly different criteria every +call, which is the main source of score jitter and occasional task-type +misreads. This module supplies two levels of stabilization, both keyed by the +intent vocabulary in :mod:`twinkle_agentic.preprocessor.intents`: + +- ``INTENT_BASE_RUBRICS`` — half-fixed **skeletons**: a small, stable core of + criteria that is PREPENDED to the distilled rubric. The generator still adds + task-specific criteria on top, so flexibility is preserved while the shared + core makes scores comparable across similar segments. This is the DEFAULT + policy (does NOT sacrifice flexibility). +- ``INTENT_FIXED_RUBRICS`` — fully-fixed rubrics per intent (no generation). + Maximum stability, minimum flexibility; opt-in for callers that want it. + +Each criterion is written to match the grader prompt conventions: +- starts with "The response" / "The agent", +- [Hard Rule] for objectively checkable constraints, [Principle] for quality, +- scoped to what is observable INSIDE one segment (never assumes later steps). + +Criteria are deliberately GENERIC (no entities/values) so a single skeleton +generalizes across all segments of that intent. +""" +from typing import Dict, List + +from .rubric_verifier import RubricItem + +# Re-export intent constants so callers wire the library without importing the +# heavier classifier module. +from ..preprocessor.intents import (INTENT_CODE, INTENT_MATH, # noqa: F401 + INTENT_TOOL_CALL) + + +def _h(text: str) -> RubricItem: + return RubricItem(text=text, is_hard=True) + + +def _p(text: str) -> RubricItem: + return RubricItem(text=text, is_hard=False) + + +# --------------------------------------------------------------------------- # +# Half-fixed skeletons (DEFAULT). Kept intentionally short (2-3 items) so the +# distilled generator still supplies the bulk of task-specific coverage. +# --------------------------------------------------------------------------- # +_TOOL_CALL_SKELETON: List[RubricItem] = [ + _h('The agent emits tool calls whose arguments are valid, complete JSON ' + 'matching the tool schema'), + _h('The agent selects tools appropriate to the sub-goal and does not invent ' + 'unavailable tools or arguments'), + _p('The agent uses each tool result to advance the sub-goal without ' + 'redundant or repeated identical calls'), +] + +_CODE_SKELETON: List[RubricItem] = [ + _h('The response produces code that is syntactically well-formed and ' + 'self-consistent within the segment'), + _p('The response addresses the stated coding sub-goal with correct, relevant ' + 'logic rather than placeholder or off-topic code'), + _p('The response avoids obvious defects (undefined names, wrong API usage) ' + 'visible within the segment'), +] + +_MATH_SKELETON: List[RubricItem] = [ + _h('The response performs each mathematical step correctly with no ' + 'arithmetic or algebraic error visible in the segment'), + _p('The response follows a valid, coherent solution path toward the ' + 'sub-goal without unjustified leaps'), + _p('The response states intermediate/final results clearly and consistently ' + 'with the work shown'), +] + +INTENT_BASE_RUBRICS: Dict[str, List[RubricItem]] = { + INTENT_TOOL_CALL: _TOOL_CALL_SKELETON, + INTENT_CODE: _CODE_SKELETON, + INTENT_MATH: _MATH_SKELETON, +} + + +# --------------------------------------------------------------------------- # +# Fully-fixed rubrics (opt-in). Same criteria plus a couple more so the fixed +# set is self-sufficient without any generation. +# --------------------------------------------------------------------------- # +INTENT_FIXED_RUBRICS: Dict[str, List[RubricItem]] = { + INTENT_TOOL_CALL: _TOOL_CALL_SKELETON + [ + _p('The agent grounds its next action in the actual tool output rather ' + 'than hallucinating results'), + ], + INTENT_CODE: _CODE_SKELETON + [ + _p('The response explains or structures the code enough to be usable in ' + 'the surrounding task context'), + ], + INTENT_MATH: _MATH_SKELETON + [ + _p('The response keeps units, signs and notation consistent throughout ' + 'the segment'), + ], +} + + +def default_intent_base_rubrics() -> Dict[str, List[RubricItem]]: + """The recommended half-fixed policy (flexible + stabilized).""" + return {k: list(v) for k, v in INTENT_BASE_RUBRICS.items()} + + +def default_intent_fixed_rubrics() -> Dict[str, List[RubricItem]]: + """The opt-in fully-fixed policy (max stability, min flexibility).""" + return {k: list(v) for k, v in INTENT_FIXED_RUBRICS.items()} diff --git a/src/twinkle_agentic/verifier/rubric_verifier.py b/src/twinkle_agentic/verifier/rubric_verifier.py index 9c505c1a9..e8a3552d3 100644 --- a/src/twinkle_agentic/verifier/rubric_verifier.py +++ b/src/twinkle_agentic/verifier/rubric_verifier.py @@ -26,9 +26,10 @@ from __future__ import annotations import json +import os import re from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple from twinkle_agentic.utils.llm_backup import llm_backup @@ -55,6 +56,18 @@ [Principle] for softer quality (reasoning soundness, sub-goal progress, no \ redundant calls). +Scope discipline (critical — avoid over-strict, mismatched rubrics): +- This is ONE SEGMENT, possibly the MIDDLE of a longer task. Only write criteria \ +about behavior that is OBSERVABLE INSIDE THIS SEGMENT. Do NOT invent criteria \ +about a final deliverable, later steps, or task completion that this segment is \ +not expected to reach (e.g. "registers the component", "updates the entry point"). +- Infer the task type ONLY from what the segment actually does. Do NOT assume it \ +is an "implement a feature" task unless the segment clearly shows that. When the \ +segment only reads/inspects/answers, judge reading/answering quality, not delivery. +- Reasoning shown inside ... (or ) is internal scratch \ +work. Never write a criterion that penalizes the mere presence of such reasoning, \ +and do NOT let it count against "output only X" style constraints. + Rules: - Output {min_n}-{max_n} criteria, as FEW as needed to cover the key axes. - Do NOT reference specific entities/values from THIS segment; keep criteria \ @@ -80,8 +93,17 @@ : PASS or : FAIL Judge every criterion independently and literally. A [Hard Rule] fails unless \ -it is unambiguously satisfied. Output only the verdict lines, in order, then \ -stop. Do not add explanations.""" +it is unambiguously satisfied. + +Grading discipline: +- Judge ONLY what is observable in THIS segment; if a criterion asks about a \ +step/deliverable this segment was not meant to reach, do not FAIL it for that \ +alone — grade it satisfied when the in-segment behavior is correct. +- Content inside ... (or ) is internal reasoning, not \ +user-facing output. For "output only X / no extra text" style criteria, ignore \ +such reasoning blocks; judge the actual response payload. + +Output only the verdict lines, in order, then stop. Do not add explanations.""" _SCORE_USER = """\ ## Task / query (context) @@ -96,6 +118,55 @@ Now output one PASS/FAIL line per criterion, in order.""" +# --- diagnostic mode: single call yields verdict + reason together --------- +# Used to distil an "on-the-fly error checker" LoRA (DESIGN §11.6). Unlike the +# terse scorer above, this asks for a COMPLETE verification chain over EVERY +# criterion (both pass and fail) so the distilled LoRA learns to also emit +# "checked, all good, continue" — not only to nitpick. Verdict and reason are +# produced in ONE pass so they can never disagree. +_DIAG_SYSTEM = """\ +You are a process error checker for one segment of an agent trajectory. You are \ +given a rubric (numbered criteria, each tagged [Hard Rule] or [Principle]) and \ +the segment. Walk through EVERY criterion in order and, for each, decide PASS or \ +FAIL and briefly justify it grounded in the segment. + +Output STRICT JSON (no prose outside it) with this shape: +{ + "items": [ + {"index": 1, "verdict": "PASS", "reason": "", + "fix": ""}, + {"index": 2, "verdict": "FAIL", "reason": "", + "fix": ""} + ], + "overall": "OK" | "ISSUES", + "summary": "" +} + +Rules: +- Judge every criterion independently and literally; a [Hard Rule] is FAIL \ +unless unambiguously satisfied. +- Judge ONLY what is observable in THIS segment; do not FAIL a criterion merely \ +because a later step/deliverable it references is outside this segment's scope. +- Content inside ... (or ) is internal reasoning, not \ +user-facing output; ignore it for "output only X" style criteria. +- For PASS items, leave "fix" as "". For FAIL items, "fix" must be a concrete \ +correction (e.g. add the missing argument, redo step k). +- "overall" is "OK" only if NO criterion is FAIL. +- Output only the JSON object.""" + +_DIAG_USER = """\ +## Task / query (context) +{query} + +## Rubric +{rubric} + +## Segment +{segment} + +Now output the diagnostic JSON object.""" + + # --------------------------------------------------------------------------- # Data holders # --------------------------------------------------------------------------- @@ -122,6 +193,33 @@ class ScoreDetail: per_item_pass_rate: List[float] = field(default_factory=list) +@dataclass +class DiagnosisItem: + """Per-criterion diagnostic verdict with its justification.""" + index: int + verdict: bool # True == PASS + reason: str = '' + fix: str = '' # concrete correction, only for FAIL + + +@dataclass +class DiagnoseDetail: + """A complete verification chain over one segment (DESIGN §11.6). + + Produced in a single LLM call so verdict and reason are always consistent. + Covers EVERY criterion (pass and fail) so a distilled checker learns to emit + "checked, no error, continue" as well as concrete fault localisation. + """ + scalar: float # aggregated pointwise score in [0, 1] + overall_ok: bool # True == no criterion failed + summary: str # one-line human-readable conclusion + items: List[DiagnosisItem] = field(default_factory=list) + rubric: List[RubricItem] = field(default_factory=list) + raw: str = '' # raw model output (for SFT targets) + query: str = '' # task/query context (for SFT inputs) + segment_text: str = '' # rendered segment (for SFT inputs) + + # --------------------------------------------------------------------------- # Verifier # --------------------------------------------------------------------------- @@ -161,6 +259,17 @@ def __init__( max_votes: int = 5, gate: bool = True, fixed_rubric: Optional[List['RubricItem']] = None, + base_rubric: Optional[List['RubricItem']] = None, + intent_rubrics: Optional[Dict[str, List['RubricItem']]] = None, + intent_base_rubrics: Optional[Dict[str, List['RubricItem']]] = None, + max_segment_chars: int = 14_000, + long_segment_chars: int = 8_000, + min_votes_long: int = 3, + long_margin_threshold: float = 0.18, + max_votes_long: int = 3, + min_votes_high: int = 3, + high_score_threshold: float = 0.85, + diag_max_tokens: int = 2048, ): if max_rubrics < min_rubrics: raise ValueError('max_rubrics must be >= min_rubrics') @@ -185,9 +294,32 @@ def __init__( self.margin_threshold = float(margin_threshold) self.max_votes = int(max_votes) self.gate = bool(gate) + self.max_segment_chars = int(max_segment_chars) + self.long_segment_chars = int(long_segment_chars) + self.min_votes_long = max(1, int(min_votes_long)) + self.long_margin_threshold = float(long_margin_threshold) + self.max_votes_long = max(1, int(max_votes_long)) + # High-confidence band: force at least this many votes when the first + # pass lands >= high_score_threshold, so 4/4-looking "level 4" segments + # are not decided by a single lucky sample (reduces high-band variance). + self.min_votes_high = max(1, int(min_votes_high)) + self.high_score_threshold = float(high_score_threshold) + # Diagnosis emits a full per-criterion (verdict+reason+fix) JSON; it needs + # a far larger token budget than terse scoring or it truncates mid-JSON. + self.diag_max_tokens = max(256, int(diag_max_tokens)) # When provided, skip stage-1 rubric generation and score against these # fixed criteria (e.g. a safety rubric — AUDIT D8). self.fixed_rubric: Optional[List['RubricItem']] = list(fixed_rubric) if fixed_rubric else None + # Skeleton criteria PREPENDED to every distilled rubric (half-fixed mode, + # DESIGN follow-up): stabilizes cross-segment comparability while still + # letting stage-1 add task-specific criteria. Ignored when fixed_rubric set. + self.base_rubric: Optional[List['RubricItem']] = list(base_rubric) if base_rubric else None + # Intent-aware routing: per-intent fully-fixed rubrics (highest priority) + # and per-intent half-fixed skeletons. Keys are intent strings (intents.py). + self.intent_rubrics: Optional[Dict[str, List['RubricItem']]] = ( + {k: list(v) for k, v in intent_rubrics.items()} if intent_rubrics else None) + self.intent_base_rubrics: Optional[Dict[str, List['RubricItem']]] = ( + {k: list(v) for k, v in intent_base_rubrics.items()} if intent_base_rubrics else None) # ------------------------------------------------------------------ # public entry points @@ -197,9 +329,10 @@ def __call__(self, trajectory: dict, **kwargs) -> int: def score_detail(self, trajectory: dict, *, query: Optional[str] = None, sampling_params: Any = None, - extra_context: Optional[str] = None) -> ScoreDetail: + extra_context: Optional[str] = None, + intent: Optional[str] = None) -> ScoreDetail: query = query or self._infer_query(trajectory) - segment_text = self._render_segment(trajectory) + segment_text = self._trim_segment_for_llm(self._render_segment(trajectory)) # D7c: fold an objective finding into the scored transcript so the judge # re-scores WITH the hard evidence in view (objective corrects subjective). if extra_context: @@ -220,15 +353,7 @@ def score_detail(self, trajectory: dict, *, query: Optional[str] = None, ) # --- stage 1: rubric (fixed if configured, else distilled generation) --- - if self.fixed_rubric is not None: - rubric = list(self.fixed_rubric) - else: - raw_rubric = self._gen_rubric( - trajectory=self._gen_trajectory(query, segment_text), - sampling_params=self._gen_sampling_params(sampling_params), - query=query, - ) - rubric = self._parse_rubric(raw_rubric) + rubric = self._build_rubric(query, segment_text, sampling_params, intent=intent) if not rubric: # No usable rubric: fall back to the code signal alone. scalar = hard_pass_rate if has_hard else 0.0 @@ -263,9 +388,116 @@ def score_detail(self, trajectory: dict, *, query: Optional[str] = None, per_item_pass_rate=per_item_rate, ) + def diagnose(self, trajectory: dict, *, query: Optional[str] = None, + sampling_params: Any = None, + intent: Optional[str] = None) -> DiagnoseDetail: + """Produce a COMPLETE verification chain over the segment (DESIGN §11.6). + + Unlike :meth:`score_detail` (terse PASS/FAIL, tuned to be cheap), this + emits, in a SINGLE llm_backup-distilled call, a per-criterion verdict + *with* its reason and (on FAIL) a concrete fix, plus an overall verdict. + The single call keeps verdict and reason mutually consistent, and it + covers passing criteria too so a distilled checker learns to say + "checked, no error, continue" — not only to nitpick. + + Every call flows through ``llm_backup``: the (student, teacher, match) + pairs it records are exactly the SFT corpus for the error-checker LoRA. + Store all of them (both OK and ISSUES segments); balancing is a + training-time sampling concern, not a collection-time one. + """ + query = query or self._infer_query(trajectory) + segment_text = self._trim_segment_for_llm(self._render_segment(trajectory)) + + if not self._llm_available(): + # No LLM: fall back to the deterministic code signal only. + hard_pass_rate, has_hard = self._code_hard_checks(trajectory) + scalar = hard_pass_rate if has_hard else 1.0 + return DiagnoseDetail( + scalar=scalar, overall_ok=scalar >= 1.0, + summary='no LLM available; code-hard signal only', + items=[], rubric=[], raw='', query=query, segment_text=segment_text) + + # Reuse the same rubric machinery as scoring (fixed / half-fixed / gen). + rubric = self._build_rubric(query, segment_text, sampling_params, intent=intent) + if not rubric: + hard_pass_rate, has_hard = self._code_hard_checks(trajectory) + scalar = hard_pass_rate if has_hard else 1.0 + return DiagnoseDetail( + scalar=scalar, overall_ok=scalar >= 1.0, + summary='no usable rubric; code-hard signal only', + items=[], rubric=rubric, raw='', query=query, segment_text=segment_text) + + rubric_block = self._render_rubric(rubric) + rubric_key = _short_hash(rubric_block) + raw = self._diagnose_once( + trajectory=self._diagnose_trajectory(query, rubric_block, segment_text), + sampling_params=self._diagnose_sampling_params(sampling_params, temperature=0.0), + query=query, rubric_key=rubric_key) + + items, overall_ok, summary = self._parse_diagnosis(raw, len(rubric)) + # Blend deterministic hard checks in as a gatekeeper floor, mirroring + # score_detail so the diagnostic scalar is comparable to the scoring one. + per_item_rate = [1.0 if it.verdict else 0.0 for it in items] + llm_scalar = self._aggregate(rubric, per_item_rate) if per_item_rate else 0.0 + hard_pass_rate, has_hard = self._code_hard_checks(trajectory) + scalar = llm_scalar + if self.gate and has_hard and hard_pass_rate < 1.0: + scalar = min(llm_scalar, hard_pass_rate) + return DiagnoseDetail( + scalar=scalar, overall_ok=overall_ok, summary=summary, + items=items, rubric=rubric, raw=raw, + query=query, segment_text=segment_text) + # ------------------------------------------------------------------ - # stage 1: rubric generation (student, distilled via llm_backup) + # stage 1: rubric assembly (fixed | half-fixed skeleton + distilled | gen) # ------------------------------------------------------------------ + def _build_rubric(self, query, segment_text, sampling_params, + intent: Optional[str] = None) -> List[RubricItem]: + """Return the rubric to score against. + + Selection order (intent-aware routing, DESIGN follow-up): + 1. ``intent_rubrics[intent]`` set -> fully fixed for this intent (most + stable; template-like tasks such as tool_call / code / math). + 2. ``fixed_rubric`` set -> global fixed rubric, verbatim. + 3. else -> distilled generation, optionally + PREPENDED with a fixed skeleton: ``intent_base_rubrics[intent]`` if + present, else the global ``base_rubric`` (half-fixed). Skeleton gives + cross-segment comparability; the generated tail adds task-specific + coverage. Duplicate criteria (same normalized text) drop, skeleton wins. + """ + if intent and self.intent_rubrics and intent in self.intent_rubrics: + return list(self.intent_rubrics[intent]) + if self.fixed_rubric is not None: + return list(self.fixed_rubric) + + skeleton: Optional[List[RubricItem]] = None + if intent and self.intent_base_rubrics and intent in self.intent_base_rubrics: + skeleton = self.intent_base_rubrics[intent] + elif self.base_rubric: + skeleton = self.base_rubric + + gen_min = self.min_rubrics + gen_max = self.max_rubrics + if skeleton: + # leave room for the skeleton so the total stays in the count window + gen_min = max(1, self.min_rubrics - len(skeleton)) + gen_max = max(gen_min, self.max_rubrics - len(skeleton)) + raw_rubric = self._gen_rubric( + trajectory=self._gen_trajectory(query, segment_text, gen_min, gen_max), + sampling_params=self._gen_sampling_params(sampling_params), + query=query) + generated = self._parse_rubric(raw_rubric) + if not skeleton: + return generated + merged = list(skeleton) + seen = {_norm_criterion(it.text) for it in merged} + for it in generated: + key = _norm_criterion(it.text) + if key and key not in seen: + seen.add(key) + merged.append(it) + return merged + @llm_backup(key_params=['query'], comparator=lambda a, b: _rubric_similar(a, b)) def _gen_rubric(self, trajectory, sampling_params, query: str = None) -> str: return self._sample_text(trajectory, sampling_params, self.gen_lora_path) @@ -282,6 +514,18 @@ def _score_once(self, trajectory, sampling_params, query: str = None, rubric_key: str = '') -> str: return self._sample_text(trajectory, sampling_params, self.score_lora_path) + # ------------------------------------------------------------------ + # diagnostic pass (student, distilled via llm_backup) — DESIGN §11.6 + # ------------------------------------------------------------------ + # Distilled on the full (verdict + reason) chain. Consistency is checked on + # the per-criterion verdict vector (same idea as scoring), not on the free + # text of the reasons — two valid reasons for the same verdict should match. + @llm_backup(key_params=['query', 'rubric_key'], + comparator=lambda a, b: _diag_verdicts_close(a, b)) + def _diagnose_once(self, trajectory, sampling_params, query: str = None, + rubric_key: str = '') -> str: + return self._sample_text(trajectory, sampling_params, self.score_lora_path) + def _score_with_voting(self, query, segment_text, rubric, sampling_params ) -> Tuple[List[float], int]: n = len(rubric) @@ -289,6 +533,14 @@ def _score_with_voting(self, query, segment_text, rubric, sampling_params rubric_key = _short_hash(rubric_block) score_traj = self._score_trajectory(query, rubric_block, segment_text) + margin_thr = self.margin_threshold + vote_cap = self.max_votes + min_votes = 1 + if len(segment_text) >= self.long_segment_chars: + margin_thr = self.long_margin_threshold + min_votes = self.min_votes_long + vote_cap = min(vote_cap, self.max_votes_long) + # First (cheap) pass. votes: List[List[bool]] = [] first = self._score_once( @@ -297,21 +549,37 @@ def _score_with_voting(self, query, segment_text, rubric, sampling_params query=query, rubric_key=rubric_key) votes.append(self._parse_verdicts(first, n)) - # Decide whether to escalate: uncertainty = closeness of pass-rate to 0.5. + # High-confidence band: a lone pass that looks like "all good" (>= the + # high threshold) still gets re-sampled, so top-band scores are not + # decided by one lucky draw. Raise the required vote depth accordingly. rate = self._vote_rates(votes) - if self._is_uncertain(rate) and self.max_votes > 1: + first_scalar = self._aggregate(rubric, rate) + if first_scalar >= self.high_score_threshold: + min_votes = max(min_votes, min(self.min_votes_high, vote_cap)) + + # Escalate when uncertain OR when the vote-depth floor is not yet met. + if vote_cap > 1 and (self._is_uncertain(rate, margin_thr) or len(votes) < min_votes): sp = self._score_sampling_params(sampling_params, temperature=0.7) - # Escalate up to max_votes; early-stop once verdicts stabilize. - while len(votes) < self.max_votes: + while len(votes) < vote_cap: extra = self._score_once( trajectory=score_traj, sampling_params=sp, query=query, rubric_key=rubric_key) votes.append(self._parse_verdicts(extra, n)) rate = self._vote_rates(votes) - if not self._is_uncertain(rate): + if len(votes) >= min_votes and not self._is_uncertain(rate, margin_thr): break return rate, len(votes) + def _trim_segment_for_llm(self, text: str) -> str: + cap = self.max_segment_chars + if len(text) <= cap: + return text + head = cap // 2 - 96 + tail = cap // 2 - 96 + omitted = len(text) - head - tail + return (f'{text[:head]}\n\n[... {omitted} chars omitted for rubric scoring ...]\n\n' + f'{text[-tail:]}') + # ------------------------------------------------------------------ # code-level hard verification (no LLM) # ------------------------------------------------------------------ @@ -371,13 +639,15 @@ def _to_level(self, scalar: float) -> int: level = int(round(scalar * (self.NUM_LEVELS - 1))) return min(self.NUM_LEVELS - 1, max(0, level)) - def _is_uncertain(self, per_item_rate: Sequence[float]) -> bool: + def _is_uncertain(self, per_item_rate: Sequence[float], + margin_threshold: Optional[float] = None) -> bool: """A segment is uncertain if any criterion sits near the 0.5 boundary.""" + thr = self.margin_threshold if margin_threshold is None else margin_threshold if not per_item_rate: return False # distance of the aggregate margin from a confident 0/1 verdict for r in per_item_rate: - if abs(r - 0.5) * 2.0 < self.margin_threshold: + if abs(r - 0.5) * 2.0 < thr: return True return False @@ -413,7 +683,6 @@ def _llm_available(self) -> bool: """ if self.sampler is not None: return True - import os return bool(os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL')) @@ -433,9 +702,12 @@ def _sample_text(self, trajectory, sampling_params, lora_path) -> str: seqs = getattr(resp, 'sequences', None) or [] return (getattr(seqs[0], 'decoded', None) or '') if seqs else '' - def _gen_trajectory(self, query: str, segment_text: str) -> dict: + def _gen_trajectory(self, query: str, segment_text: str, + min_n: Optional[int] = None, max_n: Optional[int] = None) -> dict: user = _fill(_GEN_USER, query=query, segment=segment_text) - system = _fill(_GEN_SYSTEM, min_n=self.min_rubrics, max_n=self.max_rubrics) + system = _fill(_GEN_SYSTEM, + min_n=self.min_rubrics if min_n is None else min_n, + max_n=self.max_rubrics if max_n is None else max_n) return {'messages': [ {'role': 'system', 'content': system}, {'role': 'user', 'content': user}, @@ -448,6 +720,13 @@ def _score_trajectory(self, query: str, rubric_block: str, segment_text: str) -> {'role': 'user', 'content': user}, ]} + def _diagnose_trajectory(self, query: str, rubric_block: str, segment_text: str) -> dict: + user = _fill(_DIAG_USER, query=query, rubric=rubric_block, segment=segment_text) + return {'messages': [ + {'role': 'system', 'content': _DIAG_SYSTEM}, + {'role': 'user', 'content': user}, + ]} + def _gen_sampling_params(self, override): if override is not None: return override @@ -462,6 +741,22 @@ def _score_sampling_params(self, override, *, temperature: float): from twinkle.data_format.sampling import SamplingParams return SamplingParams(temperature=temperature, max_tokens=256) + def _diagnose_sampling_params(self, override, *, temperature: float): + """Token budget for the diagnostic pass. + + Scoring emits terse PASS/FAIL lines (256 tokens is plenty), but the + diagnosis emits a full JSON object with a per-criterion reason AND fix + for EVERY rubric item. With ~7 criteria that easily exceeds 256 tokens + and the JSON gets truncated mid-string (unparsable -> all-FAIL fallback, + useless as SFT data). Give it a much larger budget, scaled by rubric size + and overridable via ``RUBRIC_DIAG_MAX_TOKENS``. + """ + if override is not None: + return override + from twinkle.data_format.sampling import SamplingParams + cap = int(os.environ.get('RUBRIC_DIAG_MAX_TOKENS', str(self.diag_max_tokens))) + return SamplingParams(temperature=temperature, max_tokens=cap) + # ------------------------------------------------------------------ # rendering / parsing helpers # ------------------------------------------------------------------ @@ -545,6 +840,58 @@ def _parse_verdicts(raw: str, n: int) -> List[bool]: verdicts[idx] = m.group(2).lower() in ('pass', 'true', 'yes', '1') return verdicts + @classmethod + def _parse_diagnosis(cls, raw: str, n: int + ) -> Tuple[List[DiagnosisItem], bool, str]: + """Parse the diagnostic JSON into (items, overall_ok, summary). + + Robust to models that wrap JSON in code fences or add stray prose. Falls + back to the PASS/FAIL line parser when JSON is unrecoverable, so a + malformed diagnostic still yields usable verdicts (missing -> FAIL). + """ + obj = _extract_json_obj(raw) + entries: List[dict] = [] + if isinstance(obj, dict) and isinstance(obj.get('items'), list): + entries = [e for e in obj['items'] if isinstance(e, dict)] + if not entries: + # The full JSON did not parse (commonly a truncated response): salvage + # every COMPLETE ``{...}`` item object so partial diagnoses stay usable + # instead of degrading to an all-FAIL, reason-less vector. + entries = _salvage_diag_items(raw) + + items: List[DiagnosisItem] = [] + for i, entry in enumerate(entries): + try: + idx = int(entry.get('index', i + 1)) + except (TypeError, ValueError): + idx = i + 1 + verdict = str(entry.get('verdict', '')).strip().lower() in ( + 'pass', 'true', 'yes', '1', 'ok') + items.append(DiagnosisItem( + index=idx, verdict=verdict, + reason=str(entry.get('reason', '') or '').strip(), + fix=str(entry.get('fix', '') or '').strip())) + if not items: + # Last resort: the terse verdict-line parser (missing -> FAIL). + verdicts = cls._parse_verdicts(raw, n) + items = [DiagnosisItem(index=i + 1, verdict=v) + for i, v in enumerate(verdicts)] + + overall_ok = all(it.verdict for it in items) if items else False + summary = '' + if isinstance(obj, dict): + summary = str(obj.get('summary', '') or '').strip() + overall_raw = str(obj.get('overall', '') or '').strip().lower() + if overall_raw in ('ok', 'pass', 'good'): + # Trust an explicit OK only if no item contradicts it. + overall_ok = overall_ok and True + elif overall_raw in ('issues', 'issue', 'fail', 'bad'): + overall_ok = False + if not summary: + summary = ('no process errors, continue' if overall_ok + else 'process issues found') + return items, overall_ok, summary + # --------------------------------------------------------------------------- # module-level helpers (comparators etc.) @@ -576,6 +923,15 @@ def _short_hash(text: str) -> str: return hashlib.md5((text or '').encode()).hexdigest()[:12] +_NORM_RE = re.compile(r'[^a-z0-9]+') + + +def _norm_criterion(text: str) -> str: + """Normalize criterion text for dedup (lowercase, alnum-only, first 12 words).""" + words = _NORM_RE.sub(' ', (text or '').lower()).split() + return ' '.join(words[:12]) + + _TAG_ANY_RE = re.compile(r'\[\s*(hard\s*rule|principle)\s*\]', re.IGNORECASE) @@ -621,3 +977,98 @@ def _verdicts_close(a: str, b: str, tol: float = 0.25) -> bool: if ra is None or rb is None: return (a or '').strip() == (b or '').strip() return abs(ra - rb) <= tol + + +def _salvage_diag_items(raw: str) -> List[dict]: + """Recover complete diagnosis item objects from a (possibly truncated) reply. + + Scans for balanced ``{...}`` spans (string-aware, so braces inside a reason + like ``\\subsubsection*{...}`` don't corrupt the depth count) and json-parses + each object that carries a ``verdict`` key. A response cut off mid-stream + still yields every item emitted before the cut, so the diagnosis keeps its + reasons/fixes instead of collapsing to an all-FAIL, reason-less vector. + """ + if not raw: + return [] + out: List[dict] = [] + stack: List[int] = [] # start index of each open brace, by depth + in_str = False + escaped = False + for i, ch in enumerate(raw): + if in_str: + if escaped: + escaped = False + elif ch == '\\': + escaped = True + elif ch == '"': + in_str = False + continue + if ch == '"': + in_str = True + elif ch == '{': + stack.append(i) + elif ch == '}' and stack: + start = stack.pop() + frag = raw[start:i + 1] + # Only leaf-ish item objects carry a verdict; the outer envelope + # ({"items": [...]}) usually never closes when truncated anyway. + if '"verdict"' in frag and '"items"' not in frag: + try: + obj = json.loads(frag) + if isinstance(obj, dict): + out.append(obj) + except (json.JSONDecodeError, ValueError): + pass + return out + + +def _extract_json_obj(raw: str) -> Optional[dict]: + """Best-effort extraction of the first JSON object from a model response. + + Handles bare JSON, ```json fenced blocks, and JSON embedded in prose. + """ + if not raw: + return None + s = raw.strip() + # Strip a leading/trailing code fence if present. + if s.startswith('```'): + s = re.sub(r'^```[a-zA-Z]*\s*', '', s) + s = re.sub(r'\s*```$', '', s).strip() + try: + obj = json.loads(s) + return obj if isinstance(obj, dict) else None + except (json.JSONDecodeError, ValueError): + pass + # Fall back to the widest {...} span. + start = s.find('{') + end = s.rfind('}') + if start != -1 and end > start: + try: + obj = json.loads(s[start:end + 1]) + return obj if isinstance(obj, dict) else None + except (json.JSONDecodeError, ValueError): + return None + return None + + +def _diag_rate(raw: str) -> Optional[float]: + """Overall PASS rate from a diagnostic JSON response (for the comparator).""" + obj = _extract_json_obj(raw) + if isinstance(obj, dict) and isinstance(obj.get('items'), list): + verdicts = [str(e.get('verdict', '')).strip().lower() in + ('pass', 'true', 'yes', '1', 'ok') + for e in obj['items'] if isinstance(e, dict)] + if verdicts: + return sum(1 for v in verdicts if v) / len(verdicts) + return _parse_rate(raw) + + +def _diag_verdicts_close(a: str, b: str, tol: float = 0.25) -> bool: + """Comparator for the diagnostic pass: student/teacher agree when their + overall PASS rate is within ``tol``. Reasons are free text, so we match on + the verdict vector (two valid phrasings of the same verdict count as a + match), not on byte-identical explanations.""" + ra, rb = _diag_rate(a), _diag_rate(b) + if ra is None or rb is None: + return (a or '').strip() == (b or '').strip() + return abs(ra - rb) <= tol diff --git a/tests/preprocessor/test_dead_loop_agent.py b/tests/preprocessor/test_dead_loop_agent.py new file mode 100644 index 000000000..3c39cc8e4 --- /dev/null +++ b/tests/preprocessor/test_dead_loop_agent.py @@ -0,0 +1,28 @@ +from twinkle_agentic.preprocessor.dead_loop_filter import DeadLoopFilter + + +def _row(*assistant_texts): + msgs = [{'role': 'user', 'content': 'go'}] + for i, t in enumerate(assistant_texts): + msgs.append({ + 'role': 'assistant', + 'content': t, + 'tool_calls': '[{"id":"1","type":"function","function":{"name":"x","arguments":"{}"}}]' if i == 0 else '', + }) + if i == 0: + msgs.append({'role': 'tool', 'content': 'ok', 'tool_call_id': '1'}) + return {'messages': msgs} + + +def test_agent_requires_two_stuck_turns(): + f = DeadLoopFilter(agent_min_stuck_turns=2) + stuck = 'wait wait no actually hmm no wait oh wait i was wrong' + kept, dropped = f([_row(stuck, 'ok reply')]) + assert len(kept) == 1 and not dropped + + +def test_agent_drops_on_two_stuck_turns(): + f = DeadLoopFilter(agent_min_stuck_turns=2) + stuck = 'wait wait no actually hmm no wait oh wait i was wrong' + kept, dropped = f([_row(stuck, stuck)]) + assert not kept and len(dropped) == 1 diff --git a/tests/preprocessor/test_dropped_merge.py b/tests/preprocessor/test_dropped_merge.py new file mode 100644 index 000000000..bdd6ff710 --- /dev/null +++ b/tests/preprocessor/test_dropped_merge.py @@ -0,0 +1,21 @@ +import json +import os +import tempfile + +from twinkle_agentic.preprocessor import merge_dropped_shards, truncate_dropped_logs + + +def test_merge_dropped_shards(): + with tempfile.TemporaryDirectory() as td: + base = os.path.join(td, 'dropped.jsonl') + with open(f'{base}.111', 'w', encoding='utf-8') as f: + f.write(json.dumps({'step': 'A', 'id': '1'}) + '\n') + with open(f'{base}.222', 'w', encoding='utf-8') as f: + f.write(json.dumps({'step': 'B', 'id': '2'}) + '\n') + merge_dropped_shards(base) + with open(base, encoding='utf-8') as f: + lines = [ln for ln in f if ln.strip()] + assert len(lines) == 2 + assert not os.path.exists(f'{base}.111') + truncate_dropped_logs(base) + assert not os.path.exists(base) diff --git a/tests/preprocessor/test_intent_think_strip.py b/tests/preprocessor/test_intent_think_strip.py new file mode 100644 index 000000000..bd7e7e4bf --- /dev/null +++ b/tests/preprocessor/test_intent_think_strip.py @@ -0,0 +1,49 @@ +"""Content-signature intent detectors must ignore markdown/LaTeX inside . + +Regression for the copywriting-tagged-as-code bug: a non-code answer whose +private scratch-pad contained a ``` fence was misclassified as ``code``. +Task type must be decided by the visible response, not the reasoning block. +""" + +from twinkle_agentic.preprocessor.intent_classifier import (CodeDetector, + MathDetector) + + +def _asst(content): + return {'role': 'assistant', 'content': content} + + +def test_code_fence_only_in_think_is_not_code(): + msgs = [ + {'role': 'user', 'content': '为门店写一条短视频口播脚本'}, + _asst('1. 分析需求\n```\n钩子→痛点→转化\n```\n' + '钩子：这价格我不敢信。转化：现在下单立省八千。'), + ] + assert CodeDetector()(msgs) == [] + + +def test_real_code_in_visible_answer_still_detected(): + msgs = [ + {'role': 'user', 'content': '写个快排'}, + _asst('先想边界```python\n' + 'def quicksort(a):\n return a\n```'), + ] + assert CodeDetector()(msgs) == [1] + + +def test_user_code_request_not_stripped(): + # A code block in the USER turn is a genuine signal and must NOT be stripped. + msgs = [ + {'role': 'user', 'content': '```python\nprint(1)\n```\n这段有什么问题'}, + _asst('这里没有问题。'), + ] + assert CodeDetector()(msgs) == [1] + + +def test_latex_only_in_think_is_not_math(): + msgs = [ + {'role': 'user', 'content': '把这段话润色一下'}, + _asst(r'可以用 \frac{a}{b} \sum \int \sqrt{x} 打个比方' + '润色后的文字，通顺自然。'), + ] + assert MathDetector()(msgs) == [] diff --git a/tests/preprocessor/test_quality_preprocessor_map_drop.py b/tests/preprocessor/test_quality_preprocessor_map_drop.py new file mode 100644 index 000000000..ad770e004 --- /dev/null +++ b/tests/preprocessor/test_quality_preprocessor_map_drop.py @@ -0,0 +1,110 @@ +"""QualityPreprocessor + HF batched map must remove fully-dropped batches.""" + +from datasets import Dataset + +from twinkle_agentic.preprocessor import QualityPreprocessor +from twinkle_agentic.preprocessor.model_filter import ModelFilter + + +def test_fully_dropped_batch_does_not_leave_ghost_rows(): + """Returning ``{}`` from an empty batch used to keep raw rows; use empty column lists.""" + qp = QualityPreprocessor(pipeline=[ModelFilter()], dropped_log_path='') + batch = { + 'id': ['bad1', 'bad2'], + 'model_id': ['Qwen/Qwen3.5-27B', 'Qwen/Qwen3-VL-8B-Instruct'], + 'messages': [[], []], + 'user_data': [[], []], + } + out = qp(batch) + assert out == { + 'id': [], + 'model_id': [], + 'messages': [], + 'user_data': [], + } + + ds = Dataset.from_dict({ + 'id': ['bad1', 'keep', 'bad2'], + 'model_id': [ + 'Qwen/Qwen3.5-27B', + 'MiniMax/MiniMax-M2.5', + 'Qwen/Qwen3-VL-8B-Instruct', + ], + 'messages': [[], [{'role': 'user', 'content': 'hi'}], []], + 'user_data': [[], [], []], + }) + mapped = ds.map(qp, batched=True, batch_size=3) + assert len(mapped) == 1 + assert mapped[0]['model_id'] == 'MiniMax/MiniMax-M2.5' + + +class _AddTag: + """Mapper: add a top-level `tag` column to every row (never drops).""" + + def __call__(self, rows): + rows = QualityPreprocessor.map_col_to_row(rows) + return [dict(r, tag='T') for r in rows], [] + + +class _DropOdd: + """Filter: drop rows whose `id` ends in an odd digit.""" + + def __call__(self, rows): + rows = QualityPreprocessor.map_col_to_row(rows) + kept, dropped = [], [] + for r in rows: + (dropped if int(str(r['id'])[-1]) % 2 else kept).append( + dict(r, drop_reason='odd') if int(str(r['id'])[-1]) % 2 else r) + return kept, dropped + + +def test_mark_mode_returns_equal_length_columns(): + """drop_mode='mark' must never change row count inside map (ghost-proof).""" + qp = QualityPreprocessor(pipeline=[_AddTag(), _DropOdd()], drop_mode='mark') + batch = { + 'id': ['r0', 'r1', 'r2', 'r3'], + 'messages': [[], [], [], []], + 'user_data': [[], [], [], []], + } + out = qp(batch) + # every column has the SAME length as the input (4), no shrinkage + lengths = {k: len(v) for k, v in out.items()} + assert set(lengths.values()) == {4}, lengths + # the survivor-only tag column exists for all rows (None for dropped) + assert '_keep' in out and 'tag' in out + assert out['_keep'] == [True, False, True, False] # r0,r2 kept; r1,r3 dropped + assert out['tag'] == ['T', None, 'T', None] # dropped rows have no tag + + +def test_mark_mode_end_to_end_filter(): + """map(mark) + filter(_keep) yields the correct survivors, no ghosts, at scale.""" + from twinkle_agentic.preprocessor import run_quality_pipeline + + class _DS: + def __init__(self, hf): + self.dataset = hf + self.datasets = {'d': hf} + + def map(self, fn, num_proc=1, **kw): + self.dataset = self.dataset.map(fn, batched=True, num_proc=num_proc, **kw) + self.datasets['d'] = self.dataset + + def filter(self, fn, **kw): + self.dataset = self.dataset.filter(fn, **kw) + self.datasets['d'] = self.dataset + + n = 500 # large enough to cross HF's internal batch boundary (the ghost trigger) + hf = Dataset.from_dict({ + 'id': [f'r{i}' for i in range(n)], + 'messages': [[{'role': 'user', 'content': 'x'}] for _ in range(n)], + 'user_data': [[] for _ in range(n)], + }) + ds = _DS(hf) + qp = QualityPreprocessor(pipeline=[_AddTag(), _DropOdd()], drop_mode='mark') + run_quality_pipeline(ds, qp, num_proc=1) + + survivors = ds.dataset + assert len(survivors) == n // 2 # exactly the even-id rows + assert '_keep' not in survivors.column_names # transient flag stripped + assert all(int(str(survivors[i]['id'])[-1]) % 2 == 0 for i in range(len(survivors))) + assert all(survivors[i]['tag'] == 'T' for i in range(len(survivors))) # tags intact diff --git a/tests/preprocessor/test_value_selector.py b/tests/preprocessor/test_value_selector.py new file mode 100644 index 000000000..82ff5ccb2 --- /dev/null +++ b/tests/preprocessor/test_value_selector.py @@ -0,0 +1,234 @@ +"""ValueSelector: deterministic value_score + component boundaries + top-frac gate.""" + +from twinkle_agentic.preprocessor import ValueSelector +from twinkle_agentic.preprocessor import label_schema as L + + +def _single_turn(): + return {'messages': [ + {'role': 'user', 'content': '写一句短视频文案'}, + {'role': 'assistant', 'content': '这价格我不敢信，现在下单立省八千。'}, + ], 'user_data': []} + + +def _tool_call(name, args, result='ok', content=''): + return [ + {'role': 'assistant', 'content': content, + 'tool_calls': [{'id': 't', 'type': 'function', + 'function': {'name': name, 'arguments': args}}]}, + {'role': 'tool', 'tool_call_id': 't', 'content': result}, + ] + + +def _agent_row(steps): + msgs = [{'role': 'user', 'content': 'do the task'}] + for name, args, result in steps: + msgs += _tool_call(name, args, result) + return {'messages': msgs, 'user_data': []} + + +def _val(row): + return L.get_label(row, L.KEY_VALUE_SCORE) + + +def _meta(row): + return L.get_label(row, L.KEY_VALUE_META) or {} + + +def test_single_turn_scores_low(): + out, dropped = ValueSelector()([_single_turn()]) + assert dropped == [] + v = _val(out[0]) + assert v is not None and v < 0.25 + m = _meta(out[0]) + assert m['uncertainty'] == 0.0 # single all-pass round -> decided + + +def test_tool_error_raises_error_signal(): + # a tool that returns an ERROR result -> tool_executed miss -> error signal up + row = _agent_row([('search', '{"q":"x"}', 'ERROR: not found'), + ('search', '{"q":"y"}', 'ERROR: not found')]) + out, _ = ValueSelector()([row]) + assert _meta(out[0])['error'] > 0.0 + assert _val(out[0]) > _val(ValueSelector()([_single_turn()])[0][0]) + + +def test_long_agent_scores_higher_than_single_turn(): + steps = [(f'read', '{"p":"%d"}' % i, 'contents') for i in range(6)] + steps += [('grep', '{"q":"a"}', 'hit'), ('edit', '{"f":"b"}', 'done')] + long_row = _agent_row(steps) + out_long, _ = ValueSelector()([long_row]) + out_short, _ = ValueSelector()([_single_turn()]) + assert _meta(out_long[0])['difficulty'] > _meta(out_short[0])['difficulty'] + assert _val(out_long[0]) > _val(out_short[0]) + + +def test_bad_row_never_crashes(): + out, dropped = ValueSelector()([{'messages': None, 'user_data': []}, + {'messages': [], 'user_data': []}]) + assert dropped == [] + assert all(_val(r) == 0.0 for r in out) + + +def test_select_top_for_rubric_marks_global_top(): + from datasets import Dataset as HFDataset + + class _Wrap: + def __init__(self, hf): + self.dataset = hf + + from twinkle_agentic.preprocessor import select_top_for_rubric + + def _row_with_value(v): + return {'messages': [{'role': 'user', 'content': 'x'}], + 'user_data': [(L.KEY_VALUE_SCORE, str(v))]} + + hf = HFDataset.from_list([_row_with_value(v) for v in [0.1, 0.9, 0.5, 0.8, 0.2]]) + ds = _Wrap(hf) + ds, n_sel = select_top_for_rubric(ds, select_frac=0.4) # top 2 of 5 + assert n_sel == 2 + selected = [L.get_label(ds.dataset[i], L.KEY_SELECTED_FOR_RUBRIC, False) + for i in range(len(ds.dataset))] + # rows with value 0.9 and 0.8 are the top-2 + assert selected == [False, True, False, True, False] + + +class _SpyRubric: + """Records how many times the rubric was invoked.""" + max_votes = 1 + + def __init__(self): + self.calls = 0 + self.diagnose_calls = 0 + + def score_detail(self, segment, query=None, intent=None, extra_context=None): + self.calls += 1 + + class _D: + scalar = 0.5 + votes = [0.5] + return _D() + + def diagnose(self, segment, query=None, intent=None): + self.diagnose_calls += 1 + + class _Item: + def __init__(self, index, verdict, reason, fix): + self.index, self.verdict, self.reason, self.fix = index, verdict, reason, fix + + class _RItem: + def __init__(self, text, is_hard): + self.text, self.is_hard = text, is_hard + + class _Diag: + scalar = 0.5 + overall_ok = False + summary = 'one criterion failed' + query = 'do it' + segment_text = 'user: do it ...' + raw = '1. FAIL: boom -> retry' + rubric = [_RItem('args are valid JSON', True)] + items = [_Item(0, False, 'tool errored', 'retry with valid args')] + return _Diag() + + +def _low_quality_agent(): + # a tool call whose result errors -> low hard score -> not short-circuited, + # so fuse_segment will actually reach for the rubric (unless gated). + return {'messages': [ + {'role': 'user', 'content': 'do it'}, + {'role': 'assistant', 'content': '', + 'tool_calls': [{'id': 't', 'type': 'function', + 'function': {'name': 'f', 'arguments': '{}'}}]}, + {'role': 'tool', 'tool_call_id': 't', 'content': 'ERROR: boom'}, + ], 'user_data': []} + + +def test_gate_skips_rubric_for_unselected_row(): + from twinkle_agentic.preprocessor import TrajectoryScorer + + spy = _SpyRubric() + scorer = TrajectoryScorer(rubric_verifier=spy, calibrate=False) + + gated = L.set_label(_low_quality_agent(), L.KEY_SELECTED_FOR_RUBRIC, False) + scorer([gated]) + assert spy.calls == 0 # not selected -> no LLM rubric + + selected = L.set_label(_low_quality_agent(), L.KEY_SELECTED_FOR_RUBRIC, True) + scorer([selected]) + assert spy.calls >= 1 # selected -> rubric runs + + +def test_persist_diagnosis_writes_verdict_reason_fix(): + from twinkle_agentic.preprocessor import TrajectoryScorer + + spy = _SpyRubric() + scorer = TrajectoryScorer(rubric_verifier=spy, calibrate=False, + persist_diagnosis=True) + + selected = L.set_label(_low_quality_agent(), L.KEY_SELECTED_FOR_RUBRIC, True) + out, _ = scorer([selected]) + diag = L.get_label(out[0], L.KEY_RUBRIC_DIAGNOSIS) + assert spy.diagnose_calls >= 1 # diagnosis ran for the scored segment + assert isinstance(diag, list) and diag # persisted, one entry per segment + entry = diag[0] + assert entry['overall_ok'] is False + assert entry['raw'] == '1. FAIL: boom -> retry' # SFT target + assert entry['segment_text'] and entry['query'] # SFT inputs + assert entry['items'][0]['verdict'] is False + assert entry['items'][0]['reason'] == 'tool errored' # the "why" + assert entry['items'][0]['fix'] == 'retry with valid args' + assert entry['rubric'][0]['text'] == 'args are valid JSON' + + +def test_no_diagnosis_for_gated_out_row(): + from twinkle_agentic.preprocessor import TrajectoryScorer + + spy = _SpyRubric() + scorer = TrajectoryScorer(rubric_verifier=spy, calibrate=False, + persist_diagnosis=True) + gated = L.set_label(_low_quality_agent(), L.KEY_SELECTED_FOR_RUBRIC, False) + out, _ = scorer([gated]) + assert spy.diagnose_calls == 0 # gated -> no LLM at all + assert L.get_label(out[0], L.KEY_RUBRIC_DIAGNOSIS) is None + + +class _SpySafety: + """Stands in for the RubricVerifier a SafetyScorer holds.""" + fixed_rubric = None + + def __init__(self): + self.calls = 0 + + def score_detail(self, trajectory, **kwargs): + self.calls += 1 + + class _D: + scalar = 1.0 + return _D() + + +def test_safety_gate_skips_llm_for_unselected_row(): + from twinkle_agentic.preprocessor import SafetyScorer + + spy = _SpySafety() + scorer = SafetyScorer(rubric_verifier=spy, gate_label=L.KEY_SELECTED_FOR_RUBRIC) + + gated = L.set_label(_low_quality_agent(), L.KEY_SELECTED_FOR_RUBRIC, False) + out, _ = scorer([gated]) + assert spy.calls == 0 # not selected -> no LLM safety pass + assert L.get_label(out[0], L.KEY_SAFETY_SCORE) == 1.0 # tagged neutral-safe + assert L.get_label(out[0], L.KEY_SAFETY_UNSAFE) is False + + selected = L.set_label(_low_quality_agent(), L.KEY_SELECTED_FOR_RUBRIC, True) + scorer([selected]) + assert spy.calls >= 1 # selected -> safety LLM runs + + +def test_safety_no_gate_scores_every_row(): + from twinkle_agentic.preprocessor import SafetyScorer + + spy = _SpySafety() + scorer = SafetyScorer(rubric_verifier=spy) # gate_label=None -> score all + scorer([L.set_label(_low_quality_agent(), L.KEY_SELECTED_FOR_RUBRIC, False)]) + assert spy.calls == 1 # ungated: LLM runs even for a "not selected" row diff --git a/tests/twinkle_agentic/test_aggregation_fusion.py b/tests/twinkle_agentic/test_aggregation_fusion.py new file mode 100644 index 000000000..272981e0c --- /dev/null +++ b/tests/twinkle_agentic/test_aggregation_fusion.py @@ -0,0 +1,12 @@ +from twinkle_agentic.verifier.aggregation import _combine + + +def test_hard_soft_blend_high_hard_not_one_shot_veto(): + # product would be 0.9 * 0.1 = 0.09 + blended = _combine(0.9, 0.1, 'hard_soft_blend') + assert blended > 0.35 + assert blended < 0.9 + + +def test_hard_soft_blend_low_hard_uses_soft(): + assert _combine(0.5, 0.4, 'hard_soft_blend') == 0.5 * 0.4 diff --git a/tests/twinkle_agentic/test_diagnosis_salvage.py b/tests/twinkle_agentic/test_diagnosis_salvage.py new file mode 100644 index 000000000..3be52adaa --- /dev/null +++ b/tests/twinkle_agentic/test_diagnosis_salvage.py @@ -0,0 +1,62 @@ +"""Diagnosis parsing must survive truncated JSON. + +The diagnostic pass emits a full per-criterion (verdict+reason+fix) JSON. When a +teacher reply is cut off mid-stream (too small a token budget), the outer object +never closes. We must still recover every COMPLETE item object rather than +degrade to an all-FAIL, reason-less vector (which is useless as SFT data). +""" +from twinkle_agentic.verifier.rubric_verifier import (RubricVerifier, + _salvage_diag_items) + +# A realistic truncated response: 3 complete items (one reason contains LaTeX +# braces to exercise string-aware brace matching), then cut off mid-string. +TRUNCATED = ( + '{\n "items": [\n' + ' {"index": 1, "verdict": "PASS", "reason": "No tool calls needed.", "fix": ""},\n' + ' {"index": 2, "verdict": "FAIL", "reason": "Duplicated \\\\subsubsection*{ans} block.", ' + '"fix": "Remove the duplicate."},\n' + ' {"index": 3, "verdict": "PASS", "reason": "Correctly omits \\\\begin{document}.", "fix": ""},\n' + ' {"index": 4, "verdict": "FAIL", "reason": "The response is cut off right he' +) + +WELL_FORMED = ( + '{"items": [{"index": 1, "verdict": "PASS", "reason": "ok", "fix": ""},' + '{"index": 2, "verdict": "FAIL", "reason": "bad", "fix": "do x"}],' + ' "overall": "issues", "summary": "one failure"}' +) + + +def test_salvage_recovers_complete_items_from_truncated_json(): + items = _salvage_diag_items(TRUNCATED) + # 3 complete objects; the 4th (truncated) is dropped. + assert len(items) == 3 + assert [it['verdict'] for it in items] == ['PASS', 'FAIL', 'PASS'] + # braces inside the reason string must not corrupt matching + assert 'subsubsection' in items[1]['reason'] + + +def test_parse_diagnosis_truncated_keeps_reasons_and_verdicts(): + items, overall_ok, summary = RubricVerifier._parse_diagnosis(TRUNCATED, n=7) + assert len(items) == 3 # not the all-FAIL length-7 fallback + assert items[0].verdict is True + assert items[1].verdict is False + assert items[1].reason # the "why" survives + assert items[1].fix == 'Remove the duplicate.' + assert overall_ok is False # a FAIL present -> not ok + + +def test_parse_diagnosis_well_formed_still_works(): + items, overall_ok, summary = RubricVerifier._parse_diagnosis(WELL_FORMED, n=2) + assert len(items) == 2 + assert items[0].verdict is True and items[1].verdict is False + assert items[1].fix == 'do x' + assert overall_ok is False + assert summary == 'one failure' + + +def test_diag_sampling_params_bigger_than_scoring(): + rv = RubricVerifier(diag_max_tokens=2048) + diag = rv._diagnose_sampling_params(None, temperature=0.0) + score = rv._score_sampling_params(None, temperature=0.0) + assert diag.max_tokens >= 2048 + assert diag.max_tokens > score.max_tokens # diagnosis needs a bigger budget diff --git a/tests/twinkle_agentic/test_repeated_calls_spin.py b/tests/twinkle_agentic/test_repeated_calls_spin.py new file mode 100644 index 000000000..c0b80195e --- /dev/null +++ b/tests/twinkle_agentic/test_repeated_calls_spin.py @@ -0,0 +1,47 @@ +"""check_no_repeated_calls: single-tool spin loops are penalized, batches aren't. + +Regression for a dead-loop that exact-duplicate detection missed: the same tool +fired ~20 times with *different* arguments (empty repeated spins) used to score +1.0 because no two (name, args) pairs were identical. +""" +import json + +from twinkle_agentic.verifier.hard_scorer import (TrajectoryView, + check_no_repeated_calls) + + +def _call(name, args): + return {'role': 'assistant', + 'tool_calls': [{'function': {'name': name, 'arguments': json.dumps(args)}}]} + + +def _score(msgs): + return check_no_repeated_calls(TrajectoryView({'messages': msgs})).score + + +def test_single_tool_spin_is_penalized(): + loop = [_call('LatexFixResponse', {'part': str(i)}) for i in range(18)] + assert _score(loop) <= 0.4 + + +def test_exact_duplicate_calls_penalized(): + dupes = [_call('read', {'p': 'same'}) for _ in range(4)] + # 3 of 4 are exact duplicates -> 1 - 3/4 = 0.25 + assert _score(dupes) <= 0.3 + + +def test_mixed_tools_not_penalized(): + mixed = [_call('read', {'p': str(i)}) for i in range(6)] + mixed += [_call('grep', {'q': 'a'}), _call('edit', {'f': 'b'}), + _call('run', {'c': 'c'}), _call('read', {'p': 'z'})] + assert _score(mixed) == 1.0 + + +def test_short_same_tool_loop_ok(): + # A legitimate 4-step same-tool loop (below the spin floor of 8 calls). + small = [_call('read', {'p': str(i)}) for i in range(4)] + assert _score(small) == 1.0 + + +def test_fewer_than_two_calls_ok(): + assert _score([_call('read', {'p': 'a'})]) == 1.0 diff --git a/tests/twinkle_agentic/test_rubric_stabilization.py b/tests/twinkle_agentic/test_rubric_stabilization.py new file mode 100644 index 000000000..0d2e36484 --- /dev/null +++ b/tests/twinkle_agentic/test_rubric_stabilization.py @@ -0,0 +1,92 @@ +"""Rubric stabilization: skeleton prepend, intent routing, high-band voting. + +These exercise :class:`RubricVerifier` without any real LLM by stubbing the two +distilled hooks (``_gen_rubric`` / ``_score_once``) and forcing +``_llm_available`` True, so we test the *assembly + voting* logic in isolation. +""" +from twinkle_agentic.preprocessor.intents import INTENT_CODE, INTENT_TOOL_CALL +from twinkle_agentic.verifier import RubricItem, RubricVerifier +from twinkle_agentic.verifier.rubric_library import default_intent_base_rubrics + + +def _segment(): + return {'messages': [ + {'role': 'user', 'content': 'do the thing'}, + {'role': 'assistant', 'content': 'here is the result'}, + ]} + + +def _stub_llm(rv, *, gen_lines, score_seq): + """Force LLM-available and deterministic gen/score outputs. + + ``score_seq`` is a list of raw verdict strings returned by successive + ``_score_once`` calls (so we can count how many votes were spent). + """ + rv._llm_available = lambda: True # type: ignore[method-assign] + rv._gen_rubric = lambda **kw: gen_lines # type: ignore[method-assign] + calls = {'n': 0} + + def _score_once(**kw): + i = min(calls['n'], len(score_seq) - 1) + calls['n'] += 1 + return score_seq[i] + + rv._score_once = _score_once # type: ignore[method-assign] + return calls + + +def test_base_rubric_prepended_and_dedup(): + base = [RubricItem('The agent calls tools with valid JSON', True)] + rv = RubricVerifier(base_rubric=base, min_rubrics=1, max_rubrics=6) + gen = ('1. The agent calls tools with valid JSON [Hard Rule]\n' # dup of skeleton + '2. The response advances the sub-goal [Principle]') + _stub_llm(rv, gen_lines=gen, score_seq=['1: PASS\n2: PASS']) + detail = rv.score_detail(_segment()) + texts = [it.text for it in detail.rubric] + # skeleton first, duplicate from generation dropped -> exactly 2 items + assert texts[0] == 'The agent calls tools with valid JSON' + assert len(detail.rubric) == 2 + + +def test_intent_fixed_rubric_skips_generation(): + fixed = {INTENT_TOOL_CALL: [RubricItem('The agent uses tools correctly', True)]} + rv = RubricVerifier(intent_rubrics=fixed) + # gen would raise if called -> proves generation is skipped for this intent + rv._llm_available = lambda: True # type: ignore[method-assign] + rv._gen_rubric = lambda **kw: (_ for _ in ()).throw(AssertionError('gen called')) # type: ignore + rv._score_once = lambda **kw: '1: PASS' # type: ignore[method-assign] + detail = rv.score_detail(_segment(), intent=INTENT_TOOL_CALL) + assert len(detail.rubric) == 1 + assert detail.rubric[0].text == 'The agent uses tools correctly' + + +def test_intent_base_rubric_routes_by_intent(): + rv = RubricVerifier(intent_base_rubrics=default_intent_base_rubrics(), + min_rubrics=1, max_rubrics=8) + _stub_llm(rv, gen_lines='1. The response is coherent [Principle]', + score_seq=['1: PASS\n2: PASS\n3: PASS\n4: PASS']) + detail = rv.score_detail(_segment(), intent=INTENT_CODE) + # the CODE skeleton leads the rubric + assert detail.rubric[0].text.startswith('The response produces code') + + +def test_high_band_forces_more_votes(): + rv = RubricVerifier(min_rubrics=1, max_rubrics=4, min_votes_high=3, + high_score_threshold=0.85, max_votes=5) + gen = '1. The response is correct [Hard Rule]' + # first pass all-PASS -> scalar 1.0 (>= 0.85) -> must escalate to >=3 votes + calls = _stub_llm(rv, gen_lines=gen, score_seq=['1: PASS']) + detail = rv.score_detail(_segment()) + assert detail.n_votes >= 3 + assert calls['n'] >= 3 + + +def test_low_band_single_vote(): + rv = RubricVerifier(min_rubrics=1, max_rubrics=4, min_votes_high=3, + high_score_threshold=0.85, max_votes=5, margin_threshold=0.1) + gen = '1. The response is correct [Hard Rule]' + # first pass FAIL -> scalar 0.0, decisive (far from 0.5) -> single vote + calls = _stub_llm(rv, gen_lines=gen, score_seq=['1: FAIL']) + detail = rv.score_detail(_segment()) + assert detail.n_votes == 1 + assert calls['n'] == 1 From 0cb5b8a9d3ea1a98b9813c0fd9b5bc7fb6507cdf Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Sat, 11 Jul 2026 19:28:18 +0800 Subject: [PATCH 06/60] fix --- cookbook/exp/embedding/eval_dualline_math.py | 528 +++++++++--- cookbook/exp/embedding/eval_gpqa_rag.py | 50 +- .../exp/embedding/eval_reflexion_skill.py | 753 ++++++++++++++++++ .../embedding/train_reflexion_skill_rft.py | 413 ++++++++++ .../embedding/train_reflexion_skill_rft.sh | 25 + src/twinkle_agentic/verifier/__init__.py | 2 + src/twinkle_agentic/verifier/leak_verifier.py | 329 ++++++++ .../verifier/rubric_verifier.py | 3 + 8 files changed, 1992 insertions(+), 111 deletions(-) create mode 100644 cookbook/exp/embedding/eval_reflexion_skill.py create mode 100644 cookbook/exp/embedding/train_reflexion_skill_rft.py create mode 100644 cookbook/exp/embedding/train_reflexion_skill_rft.sh create mode 100644 src/twinkle_agentic/verifier/leak_verifier.py diff --git a/cookbook/exp/embedding/eval_dualline_math.py b/cookbook/exp/embedding/eval_dualline_math.py index bbb09b388..bd1c183c7 100644 --- a/cookbook/exp/embedding/eval_dualline_math.py +++ b/cookbook/exp/embedding/eval_dualline_math.py @@ -14,8 +14,9 @@ in a single pass (identical to ``eval_gpqa_rag.py --mode direct``). - **Line B — dualline** (``--mode dualline``, default): the student generates in ``--chunk-tokens`` slices; between slices a teacher ``RubricVerifier.diagnose()`` - inspects the partial reasoning. When it reports process issues, the concrete - finding is injected as a ``[Checker]`` note and generation resumes. + inspects the full reasoning so far (query + all prior response). When it reports + process issues, the finding is injected back as a first-person self-correction + (in the student's own voice) and generation resumes. The teacher checker is the ``llm_backup`` teacher API (no student sampler is given to the verifier, so every check is served by the teacher — exactly the Phase-0 setup). @@ -25,18 +26,27 @@ performance is not a concern): each slice re-feeds the prior ``new_input_feature`` and, on injection, splices the tokenized note in before resuming. +The dataset defaults to AoPS (``--dataset aops``), which auto-downloads from +ModelScope so no local data path is needed; pass ``--dataset math`` to use the +local Hendrycks MATH set instead. Both lines MUST share ``--dataset``, ``--n``, +``--target-eval`` and ``--seed`` to stay a paired comparison. + Launch examples: - # Dual-line on 200 MATH problems (needs LLM_BACKUP_* for the teacher checker) + # Dual-line on 200 AoPS problems (needs LLM_BACKUP_* for the teacher checker) LLM_BACKUP_API_KEY=sk-... LLM_BACKUP_BASE_URL=... \\ - python cookbook/exp/embedding/eval_dualline_math.py --target-eval 200 + python cookbook/exp/embedding/eval_dualline_math.py \\ + --n 200 --target-eval 200 --seed 42 # Paired baseline on the same subset (no checker calls) - python cookbook/exp/embedding/eval_dualline_math.py --mode baseline --target-eval 200 + python cookbook/exp/embedding/eval_dualline_math.py --mode baseline \\ + --n 200 --target-eval 200 --seed 42 """ import argparse +import copy import json import os import sys +import time from collections import defaultdict from typing import Any, Dict, List, Optional @@ -47,10 +57,17 @@ # Reuse the reference eval's dataset + grading + prompts verbatim so the two # lines are measured on identical footing. -from eval_gpqa_rag import (GEN_MAX_MODEL_LEN, GEN_MAX_TOKENS, GEN_MODEL_ID, - GEN_GPU_MEM, GEN_GPUS, GEN_TEMPERATURE, GEN_TOP_P, - answers_match, build_direct_prompt, extract_boxed, - load_math) +from eval_gpqa_rag import (GEN_MODEL_ID, GEN_GPU_MEM, GEN_GPUS, GEN_TEMPERATURE, + GEN_TOP_P, answers_match, build_direct_prompt, + extract_boxed, load_aops, load_math) + +# Dualline eval defaults (override via --max-model-len or DUALLINE_MAX_MODEL_LEN). +DUALLINE_DEFAULT_MAX_MODEL_LEN = int(os.environ.get('DUALLINE_MAX_MODEL_LEN', 32000)) +DUALLINE_DEFAULT_MAX_GEN_TOKENS = int( + os.environ.get('DUALLINE_MAX_GEN_TOKENS', DUALLINE_DEFAULT_MAX_MODEL_LEN)) + +# vLLM parallel: default tp=1, dp=GEN_GPUS (override with GEN_TP / keep GEN_GPUS=8). +GEN_TP = int(os.environ.get('GEN_TP', 1)) logger = get_logger() @@ -62,13 +79,90 @@ MAX_INJECTIONS = int(os.environ.get('DUALLINE_MAX_INJECTIONS', 3)) # Only inject when the checker is confident enough that something is wrong. CHECK_SCORE_FLOOR = float(os.environ.get('DUALLINE_CHECK_FLOOR', 0.6)) - -# The note format wraps the teacher's finding so the student treats it as an -# external hint rather than its own reasoning. Kept short to limit disruption. +# The note is written in the student's own first-person voice so, when spliced +# back in, the running model treats it as its own mid-thought self-correction +# rather than an external interruption (which tended to derail generation toward +# max-length). Kept short to limit disruption. INJECT_TEMPLATE = ( - '\n\n[Checker] A quick review of the reasoning so far found an issue: {issue}\n' - 'Please account for this and continue solving.\n\n') + '\n\nWait — reviewing my reasoning above, I realize there is a problem: {issue}\n' + 'Let me correct this and continue.\n\n') +# When context hits max_model_len (or sample fails), dump query + generation here. +OVERFLOW_DUMP_DIR = os.environ.get( + 'DUALLINE_OVERFLOW_DUMP_DIR', './output/dualline/overflow_dumps') + + +def _decode(tokenizer, ids: List[int]) -> str: + return tokenizer.decode(ids, skip_special_tokens=True) + + +def _input_ids_len(cur_inputs: Any) -> Optional[int]: + """Length of the tokenized prompt fed to vLLM on this step, if known.""" + if not cur_inputs: + return None + item = cur_inputs[0] + if isinstance(item, dict) and 'input_ids' in item: + ids = item['input_ids'] + return len(ids) if ids is not None else None + return None + + +def _dump_dualline_state( + *, + reason: str, + problem: str, + debug_idx: Optional[int], + chunk_tokens: int, + cur_inputs: Any, + gen_ids: List[int], + injected_ids: List[int], + tokenizer, + n_checks: int, + n_injections: int, + findings: List[Dict[str, Any]], + total_new: int, + finished: bool, + max_model_len: int, + error: Optional[str] = None, +) -> str: + """Persist state for post-mortem (student CoT vs checker injection). Returns path.""" + os.makedirs(OVERFLOW_DUMP_DIR, exist_ok=True) + tag = f'idx{debug_idx}' if debug_idx is not None else 'idx_unknown' + path = os.path.join( + OVERFLOW_DUMP_DIR, f'{tag}_{reason}_{int(time.time())}.json') + + partial_cot = _decode(tokenizer, gen_ids) if tokenizer and gen_ids else '' + injected_text = (_decode(tokenizer, injected_ids) + if tokenizer and injected_ids else '') + ctx_len = _input_ids_len(cur_inputs) + + payload: Dict[str, Any] = { + 'reason': reason, + 'error': error, + 'query': problem, + 'debug_idx': debug_idx, + 'gen_token_count': len(gen_ids), + 'injected_token_count': len(injected_ids), + 'context_input_ids_len': ctx_len, + 'max_model_len': max_model_len, + 'chunk_tokens': chunk_tokens, + 'total_new': total_new, + 'n_checks': n_checks, + 'n_injections': n_injections, + 'findings': findings, + 'finished': finished, + 'partial_cot': partial_cot, + 'injected_text': injected_text, + 'partial_cot_chars': len(partial_cot), + 'context_is_message_prompt': ctx_len is None, + } + with open(path, 'w', encoding='utf-8') as f: + json.dump(payload, f, ensure_ascii=False, indent=2) + cot_path = path.replace('.json', '_partial_cot.txt') + with open(cot_path, 'w', encoding='utf-8') as f: + f.write(partial_cot) + sys.stderr.write(f'[dualline] overflow dump -> {path}\n') + return path # --------------------------------------------------------------------------- # Teacher checker (Phase-0: pure teacher via llm_backup) @@ -102,10 +196,21 @@ def _checker_available() -> bool: def _diagnose_partial(checker, problem: str, partial_cot: str): - """Run the teacher checker on the partial reasoning; return (issue_or_None, detail).""" + """Run the teacher checker on the reasoning so far; return (issue_or_None, detail). + + ``partial_cot`` is the FULL reasoning generated so far (all prior chunks plus + any self-corrections already spliced in), not just the latest slice, so the + teacher judges the whole derivation in context. We label it as in-progress so + it grades correctness of the steps rather than penalizing the absence of a + final answer. + """ + seg_content = ( + '[The following is the full reasoning so far, still in progress and not ' + 'yet complete. Judge only whether the reasoning up to this point is ' + 'mathematically correct; do not expect a final answer here.]\n\n' + partial_cot) seg = {'messages': [ {'role': 'user', 'content': problem}, - {'role': 'assistant', 'content': partial_cot}, + {'role': 'assistant', 'content': seg_content}, ]} try: detail = checker.diagnose(seg, query=problem) @@ -130,81 +235,237 @@ def _diagnose_partial(checker, problem: str, partial_cot: str): return (issue or None), detail -# --------------------------------------------------------------------------- -# Token-level segmented generation with mid-stream injection -# --------------------------------------------------------------------------- -def _decode(tokenizer, ids: List[int]) -> str: - return tokenizer.decode(ids, skip_special_tokens=True) +def _pad_batch_for_dp(items: List[Any], gen_dp: int) -> List[Any]: + """``slice_dp`` needs batch len >= DP world size (every rank gets work). + + Only kicks in on the tail rounds when fewer than ``gen_dp`` problems are + still active; the padded replicas are dropped by the caller. + """ + if gen_dp <= 1 or not items or len(items) >= gen_dp: + return items + pad = [copy.deepcopy(items[-1]) for _ in range(gen_dp - len(items))] + return items + pad + + +class _DualState: + """Per-problem generation state for the batched dualline loop. + + All problems advance together, one ``chunk_tokens`` slice per round. A + problem stays *active* until it emits EOS, hits ``max_gen_tokens``, would + overflow ``max_model_len``, or a sample call fails. Because the problems + share every round's ``sampler.sample`` call, the vLLM engine batches them + (and, with dp>1, spreads them across ranks) instead of running one at a + time. + """ + + __slots__ = ('idx', 'problem', 'cur_input', 'gen_ids', 'injected_ids', + 'n_checks', 'n_injections', 'findings', 'total_new', + 'finished', 'stopped_reason', 'context_input_ids_len', + 'pending_partial_cot', 'prompt_len') + + def __init__(self, idx: int, problem: str, prompt: Any): + self.idx = idx + self.problem = problem + self.cur_input: Any = prompt # str prompt (round 0) or input_feature + self.gen_ids: List[int] = [] # student-generated token ids only + self.injected_ids: List[int] = [] # spliced-in ids (excluded from answer) + self.n_checks = 0 + self.n_injections = 0 + self.findings: List[Dict[str, Any]] = [] + self.total_new = 0 + self.finished = False + self.stopped_reason: Optional[str] = None + self.context_input_ids_len: Optional[int] = None + self.pending_partial_cot: Optional[str] = None + self.prompt_len: Optional[int] = None # token len of the fixed prompt prefix + + def cur_input_len(self) -> Optional[int]: + item = self.cur_input + if isinstance(item, dict) and 'input_ids' in item: + ids = item['input_ids'] + return len(ids) if ids is not None else None + return None + + def result(self, tokenizer) -> Dict[str, Any]: + if self.context_input_ids_len is None: + self.context_input_ids_len = self.cur_input_len() + return { + 'text': _decode(tokenizer, self.gen_ids), + 'finished': self.finished, + 'stopped_reason': self.stopped_reason, + 'context_input_ids_len': self.context_input_ids_len, + 'n_checks': self.n_checks, + 'n_injections': self.n_injections, + 'findings': self.findings, + 'gen_tokens': len(self.gen_ids), + } -def generate_dualline(sampler, tokenizer, problem: str, checker, - base_params: TwinkleSamplingParams, - chunk_tokens: int) -> Dict[str, Any]: - """Generate the reasoning in slices, checking + injecting between slices. +def _dump_state_obj(st: '_DualState', tokenizer, chunk_tokens: int, + max_model_len: int, reason: str, error: str) -> None: + _dump_dualline_state( + reason=reason, + problem=st.problem, + debug_idx=st.idx, + chunk_tokens=chunk_tokens, + cur_inputs=[st.cur_input], + gen_ids=st.gen_ids, + injected_ids=st.injected_ids, + tokenizer=tokenizer, + n_checks=st.n_checks, + n_injections=st.n_injections, + findings=st.findings, + total_new=st.total_new, + finished=st.finished, + max_model_len=max_model_len, + error=error, + ) - Returns a dict with the final text, number of checks/injections, and the - per-injection findings (for the debug log / future SFT corpus). + +# --------------------------------------------------------------------------- +# Batched token-level segmented generation with mid-stream injection +# --------------------------------------------------------------------------- +def run_dualline_batch(sampler, tokenizer, problems: List[str], checker, + base_params: TwinkleSamplingParams, + chunk_tokens: int, + max_model_len: int, + max_gen_tokens: int, + gen_dp: int = 1, + diagnose_workers: int = 8) -> List[Dict[str, Any]]: + """Advance every problem in lock-step slices, sharing one sampler call/round. + + Each round: (1) preflight-drop any problem that would overflow the context, + (2) one ``sampler.sample`` over all still-active problems (vLLM batches + + spreads over dp ranks), (3) for the length-capped ones, run the teacher + diagnoses concurrently and splice injections, then loop. + + Returns per-problem result dicts in the original ``problems`` order. """ - prompt = build_direct_prompt(problem) + from concurrent.futures import ThreadPoolExecutor - # First slice: encode the trajectory (adds the generation prompt), generate - # up to chunk_tokens. Subsequent slices reuse the returned new_input_feature. chunk_params = TwinkleSamplingParams( max_tokens=chunk_tokens, temperature=base_params.temperature, top_p=base_params.top_p, num_samples=1) - cur_inputs: Any = [prompt] - gen_ids: List[int] = [] # student-generated token ids only - injected_ids: List[int] = [] # ids we spliced in (excluded from answer) - n_checks = 0 - n_injections = 0 - findings: List[Dict[str, Any]] = [] - total_new = 0 - finished = False - - while total_new < GEN_MAX_TOKENS: - responses = sampler.sample(cur_inputs, chunk_params) - seq = (responses[0].sequences[0] - if responses and responses[0].sequences else None) - if seq is None: + states = [_DualState(i, p, build_direct_prompt(p)) + for i, p in enumerate(problems)] + active = list(states) + round_no = 0 + + while active: + round_no += 1 + + # (1) Preflight: drop problems that would overflow the context window, + # and those that already reached the generation-token cap. + survivors: List[_DualState] = [] + for st in active: + if st.total_new >= max_gen_tokens: + st.stopped_reason = st.stopped_reason or 'max_gen_tokens' + continue + ctx_len = st.cur_input_len() + if ctx_len is not None and ctx_len + chunk_tokens >= max_model_len: + st.context_input_ids_len = ctx_len + st.stopped_reason = 'context_full' + _dump_state_obj( + st, tokenizer, chunk_tokens, max_model_len, + reason='preflight_context_full', + error=(f'context len {ctx_len} + chunk {chunk_tokens} ' + f'>= max_model_len {max_model_len}')) + continue + survivors.append(st) + active = survivors + if not active: break - gen_ids.extend(seq.tokens) - total_new += len(seq.tokens) - - if seq.stop_reason != 'length': - finished = True - break # hit EOS / stop -> generation complete - - # Length-capped slice: this is a pause point. Check the partial CoT. - if n_checks >= MAX_CHECKS or not checker: - cur_inputs = [seq.new_input_feature] - continue - - partial_cot = _decode(tokenizer, gen_ids) - n_checks += 1 - issue, _detail = _diagnose_partial(checker, problem, partial_cot) - - next_feat = dict(seq.new_input_feature) - if issue and n_injections < MAX_INJECTIONS: - note = INJECT_TEMPLATE.format(issue=issue) - note_ids = tokenizer.encode(note, add_special_tokens=False) - next_feat['input_ids'] = list(next_feat['input_ids']) + note_ids - if 'labels' in next_feat: - next_feat['labels'] = list(next_feat['labels']) + note_ids - injected_ids.extend(note_ids) - n_injections += 1 - findings.append({'at_token': total_new, 'issue': issue}) - cur_inputs = [next_feat] - - final_text = _decode(tokenizer, gen_ids) - return { - 'text': final_text, - 'finished': finished, - 'n_checks': n_checks, - 'n_injections': n_injections, - 'findings': findings, - 'gen_tokens': len(gen_ids), - } + + # (2) One shared sampler call over all active problems. On tail rounds + # with fewer active problems than dp ranks, pad to keep slice_dp happy + # and drop the padded responses. The context-overflow preflight above + # guarantees every input still fits, so a length-capped slice should + # never raise here; let any genuine engine error propagate instead of + # masking it as a whole-round failure. + batch_inputs = [st.cur_input for st in active] + padded = _pad_batch_for_dp(batch_inputs, gen_dp) + responses = sampler.sample(padded, chunk_params) + responses = responses[:len(active)] + + # (3) Consume each problem's slice; queue the ones needing a check. + to_diagnose: List[_DualState] = [] + next_active: List[_DualState] = [] + for st, resp in zip(active, responses): + seq = resp.sequences[0] if resp and resp.sequences else None + if seq is None: + st.stopped_reason = st.stopped_reason or 'empty_response' + continue + st.gen_ids.extend(seq.tokens) + st.total_new += len(seq.tokens) + st.cur_input = seq.new_input_feature + if st.prompt_len is None: + # Fixed prompt prefix = everything before this round's generation. + st.prompt_len = len(st.cur_input['input_ids']) - len(seq.tokens) + + if seq.stop_reason != 'length': + st.finished = True # EOS / stop -> done + continue + if st.n_checks >= MAX_CHECKS or not checker: + next_active.append(st) # keep generating, no more checks + continue + # Diagnose the FULL reasoning generated so far (all prior chunks plus + # any self-corrections already spliced in), so the teacher judges the + # whole derivation in context rather than an isolated tail slice. + st.pending_partial_cot = _decode( + tokenizer, st.cur_input['input_ids'][st.prompt_len:]) + st.n_checks += 1 + to_diagnose.append(st) + + # Concurrent teacher diagnoses for this round's length-capped problems. + if to_diagnose: + def _run(st: _DualState): + return st, _diagnose_partial( + checker, st.problem, st.pending_partial_cot) + workers = max(1, min(diagnose_workers, len(to_diagnose))) + with ThreadPoolExecutor(max_workers=workers) as ex: + for st, (issue, _detail) in ex.map(_run, to_diagnose): + st.pending_partial_cot = None + if issue and st.n_injections < MAX_INJECTIONS: + note = INJECT_TEMPLATE.format(issue=issue) + note_ids = tokenizer.encode(note, add_special_tokens=False) + feat = dict(st.cur_input) + feat['input_ids'] = list(feat['input_ids']) + note_ids + if 'labels' in feat: + feat['labels'] = list(feat['labels']) + note_ids + st.cur_input = feat + st.injected_ids.extend(note_ids) + st.n_injections += 1 + st.findings.append( + {'at_token': st.total_new, 'issue': issue}) + next_active.append(st) + + active = next_active + n_done = sum(1 for s in states if s.finished or s.stopped_reason) + sys.stderr.write( + f'[dualline] round {round_no}: active={len(active)} ' + f'done={n_done}/{len(states)}\n') + + return [st.result(tokenizer) for st in states] + + +def _load_tokenizer(model_id: str): + """Load the tokenizer from ModelScope (matches the vLLM sampler source). + + The box runs offline, so ``transformers.AutoTokenizer`` (which resolves via + the HF hub) fails with ``Network is unreachable``. ModelScope's AutoTokenizer + downloads/reads from the ModelScope cache instead — the same place the vLLM + sampler already pulled the model from. Falls back to transformers only if the + ModelScope path is unavailable. + """ + try: + from modelscope import AutoTokenizer as MSAutoTokenizer + return MSAutoTokenizer.from_pretrained(model_id, trust_remote_code=True) + except Exception as exc: + sys.stderr.write(f'[dualline] modelscope tokenizer load failed ({exc}); ' + f'falling back to transformers\n') + from transformers import AutoTokenizer + return AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) # --------------------------------------------------------------------------- @@ -214,23 +475,37 @@ def main(): p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) p.add_argument('--mode', choices=['baseline', 'dualline'], default='dualline') + p.add_argument('--dataset', choices=['aops', 'math'], default='aops', + help='Evaluation dataset. "aops" (default) auto-downloads from ' + 'ModelScope (no local path needed); "math" reads local ' + 'MATH_DATA_DIR, stratified by difficulty level.') p.add_argument('--math-split', default='test') p.add_argument('--per-level', type=int, default=0, - help='Problems per difficulty level. 0 => --n split across levels.') - p.add_argument('--n', type=int, default=200, - help='Pool size sampled from MATH (stratified by level).') - p.add_argument('--target-eval', type=int, default=200, + help='MATH only: problems per level. 0 => --n split across levels.') + p.add_argument('--n', type=int, default=32, + help='Pool size sampled from the dataset (MATH is stratified ' + 'by level; AoPS is a flat shuffle).') + p.add_argument('--target-eval', type=int, default=32, help='Stop after this many problems are evaluated (0 = all sampled).') + p.add_argument('--max-model-len', type=int, default=DUALLINE_DEFAULT_MAX_MODEL_LEN, + help='vLLM max_model_len / template max_length (default 32000).') + p.add_argument('--max-gen-tokens', type=int, default=DUALLINE_DEFAULT_MAX_GEN_TOKENS, + help='Cap total generated tokens per problem (default: same as ' + 'max-model-len / DUALLINE_MAX_GEN_TOKENS).') p.add_argument('--chunk-tokens', type=int, default=CHUNK_TOKENS, help='Generate this many tokens between checker pauses.') p.add_argument('--batch-size', type=int, default=16, - help='Baseline mode batch size (dualline runs per-problem).') + help='Baseline mode batch size. Dualline runs all problems ' + 'concurrently (one shared sampler call per slice-round).') + p.add_argument('--diagnose-workers', type=int, + default=int(os.environ.get('DUALLINE_DIAGNOSE_WORKERS', 8)), + help='Concurrency for teacher diagnose() calls within a round.') 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/dualline/math_{args.mode}_results.jsonl' + args.output = f'./output/dualline/{args.dataset}_{args.mode}_results.jsonl' is_dual = (args.mode == 'dualline') if is_dual and not _checker_available(): @@ -240,41 +515,56 @@ def main(): ' Set them, or run --mode baseline for the paired baseline.\n') sys.exit(1) - records = load_math(n=args.n, seed=args.seed, split=args.math_split, - per_level=args.per_level) + 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) if args.target_eval > 0: records = records[:args.target_eval] - sys.stderr.write(f'[dualline] evaluating {len(records)} problems (mode={args.mode})\n') + max_model_len = args.max_model_len + max_gen_tokens = args.max_gen_tokens + sys.stderr.write( + f'[dualline] evaluating {len(records)} problems ' + f'(mode={args.mode}, dataset={args.dataset}, ' + f'max_model_len={max_model_len}, max_gen_tokens={max_gen_tokens})\n') device_groups = [ DeviceGroup(name='sampler', ranks=list(range(GEN_GPUS)), - device_type='GPU', gpus_per_worker=GEN_GPUS), + device_type='GPU', gpus_per_worker=GEN_TP), ] - gen_mesh = DeviceMesh.from_sizes(world_size=GEN_GPUS, tp_size=GEN_GPUS) + if GEN_GPUS % GEN_TP != 0: + raise ValueError(f'GEN_GPUS ({GEN_GPUS}) must be divisible by GEN_TP ({GEN_TP})') + gen_dp = GEN_GPUS // GEN_TP + gen_mesh = DeviceMesh.from_sizes(world_size=GEN_GPUS, dp_size=gen_dp, tp_size=GEN_TP) 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}, + engine_args={ + 'gpu_memory_utilization': GEN_GPU_MEM, + 'max_model_len': max_model_len, + 'tensor_parallel_size': GEN_TP, + }, 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'[dualline] vLLM sampler ready (model={GEN_MODEL_ID})\n') + enable_thinking=True, max_length=max_model_len) + sys.stderr.write( + f'[dualline] vLLM sampler ready (model={GEN_MODEL_ID}, ' + f'tp={GEN_TP}, dp={gen_dp})\n') gen_params = TwinkleSamplingParams( - max_tokens=GEN_MAX_TOKENS, temperature=GEN_TEMPERATURE, + max_tokens=max_gen_tokens, temperature=GEN_TEMPERATURE, top_p=GEN_TOP_P, num_samples=1) checker = None tokenizer = None if is_dual: checker = _build_checker() - from transformers import AutoTokenizer - tokenizer = AutoTokenizer.from_pretrained(GEN_MODEL_ID, trust_remote_code=True) + tokenizer = _load_tokenizer(GEN_MODEL_ID) sys.stderr.write('[dualline] teacher checker ready (llm_backup teacher)\n') correct = 0 @@ -309,26 +599,43 @@ def _grade_and_log(rec, idx, raw_output, extra=None): out_f.flush() if is_dual: - for idx, rec in enumerate(records): - result = generate_dualline(sampler, tokenizer, rec['problem'], - checker, gen_params, args.chunk_tokens) + problems = [rec['problem'] for rec in records] + results = run_dualline_batch( + sampler, tokenizer, problems, checker, gen_params, + args.chunk_tokens, max_model_len, max_gen_tokens, + gen_dp=gen_dp, diagnose_workers=args.diagnose_workers) + for idx, (rec, result) in enumerate(zip(records, results)): _grade_and_log(rec, idx, result['text'], extra={ 'n_checks': result['n_checks'], 'n_injections': result['n_injections'], 'findings': result['findings'], 'finished': result['finished'], + 'stopped_reason': result.get('stopped_reason'), + 'context_input_ids_len': result.get('context_input_ids_len'), 'gen_tokens': result['gen_tokens'], }) - acc = correct / total if total else 0 + stop_tag = (f' stop={result["stopped_reason"]}' + if result.get('stopped_reason') else '') sys.stderr.write( - f' [{total}/{len(records)}] acc={acc:.4f} ({correct}/{total}) ' - f'checks={result["n_checks"]} inj={result["n_injections"]}\n') + f' [idx {idx}] correct={debug_records[-1]["is_correct"]} ' + f'gen={result["gen_tokens"]} checks={result["n_checks"]} ' + f'inj={result["n_injections"]}{stop_tag}\n') + acc = correct / total if total else 0 + sys.stderr.write( + f'[dualline] batched eval done: acc={acc:.4f} ({correct}/{total})\n') else: import re for batch_start in range(0, len(records), args.batch_size): batch = records[batch_start:batch_start + args.batch_size] prompts = [build_direct_prompt(r['problem']) for r in batch] + if gen_dp > 1 and len(prompts) < gen_dp: + prompts = _pad_batch_for_dp(prompts, gen_dp) + pad_n = len(prompts) - len(batch) + else: + pad_n = 0 responses = sampler.sample(prompts, gen_params) + if pad_n: + responses = responses[:len(batch)] for i, (rec, resp) in enumerate(zip(batch, responses)): seq = resp.sequences[0] if resp and resp.sequences else None raw_output = '' @@ -342,7 +649,8 @@ def _grade_and_log(rec, idx, raw_output, extra=None): overall = correct / total if total else 0 print(f'\n{"=" * 60}') print(f'MATH dual-line — mode={args.mode}, model={GEN_MODEL_ID}') - print(f' n={total}, seed={args.seed}, chunk_tokens={args.chunk_tokens}') + print(f' n={total}, seed={args.seed}, chunk_tokens={args.chunk_tokens}, ' + f'max_model_len={max_model_len}') print(f'{"=" * 60}') print(f'Overall accuracy: {overall:.4f} ({correct}/{total})') @@ -352,6 +660,14 @@ def _grade_and_log(rec, idx, raw_output, extra=None): n_with_inj = sum(1 for r in debug_records if r.get('n_injections', 0) > 0) print(f' checker: {tot_checks} checks, {tot_inj} injections across ' f'{n_with_inj}/{total} problems') + n_ctx_full = sum( + 1 for r in debug_records if r.get('stopped_reason') == 'context_full') + n_sample_fail = sum( + 1 for r in debug_records if r.get('stopped_reason') == 'sample_failed') + n_unfinished = sum(1 for r in debug_records if not r.get('finished')) + print(f' length: context_full={n_ctx_full}/{total}, ' + f'sample_failed={n_sample_fail}/{total}, ' + f'unfinished(no EOS)={n_unfinished}/{total}') if any(r.get('level') for r in debug_records): per = defaultdict(lambda: [0, 0]) diff --git a/cookbook/exp/embedding/eval_gpqa_rag.py b/cookbook/exp/embedding/eval_gpqa_rag.py index dd71590c0..72a5a9b30 100644 --- a/cookbook/exp/embedding/eval_gpqa_rag.py +++ b/cookbook/exp/embedding/eval_gpqa_rag.py @@ -107,7 +107,7 @@ 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)) +GEN_GPUS = int(os.environ.get('GEN_GPUS', 8)) EMB_GPUS = int(os.environ.get('EMB_GPUS', 2)) EMBED_MAX_LENGTH = int(os.environ.get('EMBED_MAX_LENGTH', 20000)) @@ -653,11 +653,51 @@ def answers_match(predicted: str, reference: str) -> bool: # Dataset loading # --------------------------------------------------------------------------- +def _load_aops_from_modelscope(): + """Download the AoPS repo natively from ModelScope and read its parquet. + + Primary loader: ``dataset_snapshot_download`` pulls the dataset repo files + (parquet) straight from the ModelScope hub WITHOUT going through the + ``datasets``/HF-filesystem path used by ``MsDataset.load`` — that path is + broken on this modelscope build (``HfFileSystem.find() got multiple values + for 'maxdepth'``). We then read the local parquet with the ``datasets`` + backend (reading local files does not trigger the HfFileSystem bug). + """ + import glob + + from datasets import Dataset as HFDataset + from modelscope.hub.snapshot_download import dataset_snapshot_download + + local = dataset_snapshot_download(AOPS_DATASET_ID) + files = sorted(glob.glob(os.path.join(local, '**', '*.parquet'), + recursive=True)) + if not files: + # Older snapshots may materialize an arrow file instead of parquet. + files = sorted(glob.glob(os.path.join(local, '**', '*train*.arrow'), + recursive=True)) + if files: + sys.stderr.write(f'[aops] modelscope snapshot arrow: {files[0]}\n') + return HFDataset.from_file(files[0]) + return None + sys.stderr.write(f'[aops] modelscope snapshot parquet: {files[0]}\n') + return HFDataset.from_parquet(files if len(files) > 1 else files[0]) + + 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') + """Load AoPS boxed problems, sample n, extract reference answers. + + Uses ModelScope as the data source (native repo snapshot download). + """ + ds = None + try: + ds = _load_aops_from_modelscope() + except Exception as exc: + sys.stderr.write(f'[aops] modelscope snapshot download failed ({exc}); ' + f'trying MsDataset.load\n') + if ds is None: + 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'): diff --git a/cookbook/exp/embedding/eval_reflexion_skill.py b/cookbook/exp/embedding/eval_reflexion_skill.py new file mode 100644 index 000000000..db88ca13a --- /dev/null +++ b/cookbook/exp/embedding/eval_reflexion_skill.py @@ -0,0 +1,753 @@ +"""Phase-0 measurement for the reflexion self-skill scheme (see reflexion.md). + +Question this script answers: **on problems the base model first gets wrong, does +letting the SAME base model reflect on its failed attempt, distill a general +"skill", and re-solve WITH that skill in the system prompt, actually raise its +pass@k?** No LoRA is trained here — this is the upper-bound / go-no-go gate before +investing in a Skill-LoRA. If the base model's own skills don't help, training a +LoRA to produce them is pointless. + +It deliberately reuses ``eval_gpqa_rag`` verbatim (dataset, grader, prompts, sampling +config) so numbers are comparable with the other AoPS lines. Only the base model + +one vLLM sampler are used; the dataset is AoPS; validation is on the SAME problem +(no similar-problem retrieval). + +Per chunk of problems (all sampler calls are BATCHED across the whole chunk — never +one problem at a time): + 1. Initial solve — 1 rollout each; keep only problems the model got wrong. + 2. Skill generation — for each failed problem, the base model reads its own failed + attempt and produces N candidate skills (general reminders, no answer/solution). + 3. Leak filter — drop skills that leak the gold answer or a full solution. + 4. Baseline pass@k — K rollouts of the plain problem (the "no-skill" control). + 5. With-skill pass@k — K rollouts of the problem with each surviving skill in the + system prompt. + 6. Score — marginal = with-skill pass@k − baseline pass@k; keep the best skill. +A "pass" = answer correct AND generation terminated (no length cutoff). + +Everything useful (failed attempt, every candidate skill + leak flag, baseline and +per-skill rollout stats, marginals, best skill) is written to a JSONL **incrementally +after each chunk**, so partial runs are fully analysable. + +Launch (8 GPUs, tp=1 dp=8 by default): + python cookbook/exp/embedding/eval_reflexion_skill.py --n 64 --chunk-size 16 +""" +import argparse +import copy +import json +import os +import re +import sys +import threading +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Any, Dict, List, Optional, Tuple + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import SamplingParams as TwinkleSamplingParams +from twinkle.sampler import vLLMSampler +from twinkle_agentic.protocol.openai import OpenAI as OpenAIClient + +# Reuse the reference eval's dataset + grading + prompts + sampling config so this +# line is directly comparable with eval_gpqa_rag / eval_dualline_math. +from eval_gpqa_rag import (DIRECT_SYSTEM, GEN_GPU_MEM, GEN_GPUS, GEN_MODEL_ID, + GEN_TEMPERATURE, GEN_TOP_P, MCQ_INSTRUCTION, answers_match, + build_direct_prompt, extract_boxed, load_aops) + +logger = get_logger() + +# vLLM parallel: tp=1, dp=GEN_GPUS by default (override GEN_TP; keep GEN_GPUS=8). +GEN_TP = int(os.environ.get('GEN_TP', 1)) + +# Leak-detector API (reuses eval_gpqa_rag's env names). A strong external model +# judges whether a candidate skill leaks THIS problem's answer/solution — catching +# what the string filter cannot (multiple-choice letters, derived-result leakage). +LEAK_API_KEY = os.environ.get('COMPRESS_API_KEY', '') +LEAK_BASE_URL = os.environ.get('COMPRESS_BASE_URL', + 'https://dashscope.aliyuncs.com/compatible-mode/v1') +LEAK_API_MODEL = os.environ.get('COMPRESS_MODEL', 'qwen3.7-max') + +# Global call spacer so bursts of leak-judge calls stay under the QPS limit. +_api_lock = threading.Lock() +_api_next = [0.0] + + +def _api_throttle(min_interval: float) -> None: + with _api_lock: + now = time.monotonic() + wait = max(0.0, _api_next[0] - now) + _api_next[0] = max(now, _api_next[0]) + min_interval + if wait > 0: + time.sleep(wait) + + +# --------------------------------------------------------------------------- +# Prompts (self-reflection skill generation + skill-conditioned solving) +# --------------------------------------------------------------------------- +SKILL_GEN_SYSTEM = ( + "You are a meticulous mathematics coach. You are shown a competition problem and a " + "student's FAILED attempt. Produce a SHORT list of general, reusable skills that " + 'would prevent this class of mistake on SIMILAR problems.\n\n' + 'OUTPUT FORMAT (strict):\n' + '- You may reason briefly first, but the final answer MUST be a markdown bullet ' + 'list of 3-5 items WRAPPED IN and tags. Output nothing after ' + '.\n' + '- Each item is ONE short imperative sentence (a rule, check, or habit).\n' + '- Inside the tags: no diagnosis narration, no "The student...", no headings, no ' + 'restating the problem or the examples.\n\n' + 'CONTENT RULES (strict):\n' + '- Do NOT reveal the final answer or the multiple-choice option.\n' + '- Do NOT state the specific numbers, values, or key intermediate results of THIS ' + 'problem.\n' + '- Do NOT give a step-by-step solution to THIS problem. Every item must be GENERAL ' + 'and transferable to other problems of the same type.\n\n' + 'Follow the example below for the exact tags, style, and level of generality.' +) + +SKILL_GEN_USER = ( + 'Problem:\n{problem}\n\n' + "The student's failed attempt (it may be long or may fail to terminate):\n" + '{attempt}\n\n' + 'Now output the skills bullet list.' +) + +# One-shot demonstration of the required format and generality (answer-free). +_EX_PROBLEM = 'Simplify $\\sqrt{72} + \\sqrt{18}$ and give the result.' +_EX_ATTEMPT = ( + 'The student added the radicands directly to get $\\sqrt{90}$ and concluded it ' + 'could not be simplified, never factoring out the perfect squares first.') +_EX_SKILLS = ( + '\n' + '- Before adding square roots, factor each radicand into a perfect square times a ' + 'remainder and move the perfect-square root outside.\n' + '- Never add radicands directly: $\\sqrt{a}+\\sqrt{b}\\ne\\sqrt{a+b}$.\n' + '- Only combine radical terms after reducing them to the same simplest radical ' + 'form.\n' + '- Sanity-check the simplified result by estimating each root numerically.\n' + '') + + +def build_skillgen_prompt(problem: str, attempt: str) -> Dict[str, Any]: + return {'messages': [ + {'role': 'system', 'content': SKILL_GEN_SYSTEM}, + {'role': 'user', + 'content': SKILL_GEN_USER.format(problem=_EX_PROBLEM, attempt=_EX_ATTEMPT)}, + {'role': 'assistant', 'content': _EX_SKILLS}, + {'role': 'user', + 'content': SKILL_GEN_USER.format(problem=problem, attempt=attempt)}, + ]} + + +# The skill is injected into the SYSTEM prompt (per reflexion.md), on top of the +# exact DIRECT_SYSTEM used by the baseline so the only difference is the reminders. +# Built by concatenation (NOT str.format): DIRECT_SYSTEM and the skill may contain +# literal braces (e.g. ``\boxed{}``, LaTeX), which would break ``.format``. +_SKILL_SOLVE_PREFIX = ( + DIRECT_SYSTEM + '\n\n' + 'Before you start, keep these reminders in mind to avoid common mistakes on this ' + 'type of problem:\n') +_SKILL_SOLVE_SUFFIX = ( + '\nApply them where relevant, but rely on your own reasoning to reach the answer.') + + +def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: + return {'messages': [ + {'role': 'system', 'content': _SKILL_SOLVE_PREFIX + skill + _SKILL_SOLVE_SUFFIX}, + {'role': 'user', 'content': problem + MCQ_INSTRUCTION}, + ]} + + +# --------------------------------------------------------------------------- +# Parsing / grading / leak filtering +# --------------------------------------------------------------------------- +_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') + + +def _clean_text(decoded: Optional[str]) -> str: + return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() + + +_BULLET_RE = re.compile(r'^\s*(?:[-*]|\d+[.)])\s') + + +def _extract_skill_list(text: str) -> str: + """Pull just the clean skill list out of a (possibly thinking-laden) output. + + The model is instructed to wrap the final list in ``...``, so + prefer that (robust to any preceding reasoning, closed or unterminated). Fall + back to dropping a ```` block and keeping from the first bullet onward. + """ + low = text.lower() + if '' in low: + start = low.index('') + len('') + end = low.index('') if '' in low else len(text) + return text[start:end].strip() + if '' in text: + text = text.rsplit('', 1)[-1] + text = re.sub(r'.*?', '', text, flags=re.DOTALL).strip() + lines = text.splitlines() + for i, line in enumerate(lines): + if _BULLET_RE.match(line): + return '\n'.join(lines[i:]).strip() + return text.strip() + + +def _bound_attempt(text: str, gen_tokens: int, budget_tokens: int) -> str: + """Keep a failed attempt within the skill-gen context budget. + + Round-2 (skill-gen) input contains the FULL round-1 attempt, and failed + attempts are often the ones that ran to the token cap (repetition loops), so + feeding them verbatim overflows max_model_len. Keep the head (real reasoning + + where it went wrong) plus a short tail (the final wrong answer); drop the + redundant middle. Token->char conversion uses THIS attempt's observed + chars-per-token so the cut fits precisely. + """ + if not text or gen_tokens <= budget_tokens: + return text + cpt = len(text) / max(1, gen_tokens) + head_tok = int(budget_tokens * 0.7) + tail_tok = budget_tokens - head_tok + head = text[:int(head_tok * cpt)] + tail = text[-int(tail_tok * cpt):] if tail_tok > 0 else '' + return f'{head}\n\n[... attempt truncated for length ...]\n\n{tail}' + + +def _parse_seq(seq, gold: str) -> Dict[str, Any]: + """Turn one sampled sequence into a graded rollout record. + + ``pass`` requires BOTH a correct boxed answer AND clean termination (a length + cutoff means the model never actually committed to the answer). + """ + text = _clean_text(getattr(seq, 'decoded', '') or '') + pred = extract_boxed(text) + correct = bool(pred) and answers_match(pred, gold) + terminated = getattr(seq, 'stop_reason', None) != 'length' + return { + 'pred': pred, + 'correct': correct, + 'terminated': terminated, + 'passed': bool(correct and terminated), + 'stop_reason': getattr(seq, 'stop_reason', None), + 'gen_tokens': len(getattr(seq, 'tokens', None) or []), + 'text': text, + } + + +def _pass_rate(rolls: List[Dict[str, Any]]) -> float: + return sum(1 for r in rolls if r['passed']) / len(rolls) if rolls else 0.0 + + +def _skill_leaks(skill: str, gold: str) -> Tuple[bool, str]: + """Reject a skill that leaks the answer, or is a degenerate / non-list output.""" + if not skill.strip(): + return True, 'empty' + bullets = [ln for ln in skill.splitlines() if _BULLET_RE.match(ln)] + if len(bullets) < 2: + return True, 'too_short' + low = skill.lower() + if 'item 1' in low and 'item 2' in low: # model echoed the format placeholder + return True, 'placeholder' + if '\\boxed' in skill: + return True, 'contains_boxed' + g = (gold or '').strip() + # Raw substring match is only trustworthy when the answer is specific enough that + # an incidental hit is unlikely. Short answers ('D', 'E', '1') would match almost + # any text, so leave those to the API judge instead of false-flagging every skill. + if len(g) >= 4 and g.lower() in skill.lower(): + return True, 'contains_gold_answer' + # Standalone multi-digit numbers from the gold answer leaking into the skill. + for num in re.findall(r'-?\d{2,}', g): + if re.search(r'(? Optional[bool]: + """Return True (leak) / False (clean) / None (unparseable or API error after retries). + + Only transient API errors are retried (with exponential backoff); an unparseable + verdict is deterministic at temperature 0, so retrying it is pointless. + """ + msgs = [ + {'role': 'system', 'content': _LEAK_JUDGE_SYSTEM}, + {'role': 'user', 'content': _LEAK_JUDGE_USER.format( + problem=problem[:4000], gold=gold, skill=skill[:4000])}, + ] + for attempt in range(retries + 1): + _api_throttle(min_interval) + try: + reply = api({'messages': msgs}, + TwinkleSamplingParams(temperature=0.0, max_tokens=16), + extra_body={'enable_thinking': False}) + except Exception as exc: # noqa: BLE001 — broad catch is intentional + logger.warning(f'[leak-judge] error (attempt {attempt + 1}/{retries + 1}): {exc}') + if attempt < retries: + time.sleep(min(4.0, 0.5 * 2 ** attempt)) # exponential backoff + continue + return None + verdict = (reply.get('content') or '').strip().upper() + if 'CLEAN' in verdict: + return False + if 'LEAK' in verdict: + return True + return None # unparseable — deterministic at temp 0, no point retrying + + +def _api_leak_batch(api: OpenAIClient, items: List[Tuple[int, str, str, str]], + concurrency: int, min_interval: float, + retries: int) -> Dict[int, Optional[bool]]: + """Judge many (key, problem, gold, skill) tuples in parallel; key -> verdict.""" + verdicts: Dict[int, Optional[bool]] = {} + if not items: + return verdicts + with ThreadPoolExecutor(max_workers=min(len(items), concurrency)) as pool: + futs = {pool.submit(_api_leak_judge_one, api, p, g, s, min_interval, retries): k + for (k, p, g, s) in items} + for fut in as_completed(futs): + verdicts[futs[fut]] = fut.result() + return verdicts + + +# --------------------------------------------------------------------------- +# Batched sampling (one shared sampler.sample per phase — never per problem) +# --------------------------------------------------------------------------- +def _pad_for_dp(prompts: List[Any], gen_dp: int) -> List[Any]: + """vLLM dp needs batch len >= dp; pad tail rounds and let the caller slice back.""" + if gen_dp <= 1 or not prompts or len(prompts) >= gen_dp: + return prompts + pad = [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] + return prompts + pad + + +def _run_samples(sampler, prompts: List[Any], num_samples: int, max_tokens: int, + gen_dp: int) -> List[List[Any]]: + """One batched sampler call; return per-prompt list of raw sampled sequences.""" + if not prompts: + return [] + params = TwinkleSamplingParams( + max_tokens=max_tokens, temperature=GEN_TEMPERATURE, top_p=GEN_TOP_P, + num_samples=num_samples) + padded = _pad_for_dp(prompts, gen_dp) + responses = sampler.sample(padded, params)[:len(prompts)] + return [list(r.sequences) if (r and r.sequences) else [] for r in responses] + + +def _set_thinking(sampler, args: argparse.Namespace, enabled: bool) -> None: + """Toggle the remote template's thinking mode. + + Skill generation wants thinking OFF so the model emits the short ```` + list directly (with thinking ON it burns the token budget reasoning and often + never reaches the list); solving wants it ON. ``set_template`` is a + remote_function, so this propagates to every sampler worker. + """ + sampler.set_template('Qwen3_5Template', model_id=GEN_MODEL_ID, + enable_thinking=enabled, max_length=args.max_model_len) + + +def _bounded_attempt_for(r: Dict[str, Any], args: argparse.Namespace) -> str: + """Per-problem bound so problem + attempt + skill output fits the context window.""" + prob_est = len(r['problem']) // 2 # conservative problem token estimate + budget = max(1024, args.max_model_len - args.skill_max_tokens + - args.attempt_reserve_tokens - prob_est) + return _bound_attempt(r['_init'][0]['text'], r['_init'][0]['gen_tokens'], budget) + + +def _filter_candidates(api: Optional[OpenAIClient], + cands: List[Tuple[Dict[str, Any], str]], + args: argparse.Namespace) -> None: + """Apply string + API leak filters to (problem, skill) candidates; append results. + + The cheap string filter runs first; the API judge only sees skills that pass it, + which is what catches MCQ-letter / derived-result / full-solution leakage. + """ + prepared = [] # [r, text, leaked(bool|None), reason] + for r, text in cands: + leaked, reason = _skill_leaks(text, r['reference_answer']) + prepared.append([r, text, True if leaked else None, reason]) + if api is not None: + items = [(i, prepared[i][0]['problem'], prepared[i][0]['reference_answer'], + prepared[i][1]) for i in range(len(prepared)) if prepared[i][2] is None] + verdicts = _api_leak_batch(api, items, args.api_concurrency, args.api_min_interval, + args.api_retries) + for key, _p, _g, _s in items: + v = verdicts.get(key) + if v is True: + prepared[key][2], prepared[key][3] = True, 'api_leak' + elif v is False: + prepared[key][2], prepared[key][3] = False, '' + else: + prepared[key][2], prepared[key][3] = False, 'api_uncertain' + for r, text, leaked, reason in prepared: + r['_skills'].append({'skill': text, 'leaked': bool(leaked), 'leak_reason': reason}) + + +def _build_skills(sampler, api: Optional[OpenAIClient], failed: List[Dict[str, Any]], + gen_dp: int, args: argparse.Namespace) -> None: + """Generate + extract + leak-filter skills, re-rolling problems short on clean ones. + + Clean skills accumulate across rounds; only problems still below ``min_survivors`` + clean skills are re-rolled, up to ``skill_retries`` extra rounds. + """ + for r in failed: + r['_skills'] = [] + todo = list(failed) + _set_thinking(sampler, args, False) # skill-gen: emit the list directly, no CoT + try: + for _ in range(args.skill_retries + 1): + if not todo: + break + sg_out = _run_samples( + sampler, + [build_skillgen_prompt(r['problem'], _bounded_attempt_for(r, args)) for r in todo], + args.n_skills, args.skill_max_tokens, gen_dp) + cands = [(r, _extract_skill_list(_clean_text(getattr(s, 'decoded', '') or ''))) + for r, seqs in zip(todo, sg_out) for s in seqs] + _filter_candidates(api, cands, args) + todo = [r for r in failed + if sum(1 for sk in r['_skills'] if not sk['leaked']) < args.min_survivors] + finally: + _set_thinking(sampler, args, True) # restore for solving phases + tot = sum(len(r['_skills']) for r in failed) + leaked = sum(1 for r in failed for sk in r['_skills'] if sk['leaked']) + sys.stderr.write(f' phase2: skills={tot} leaked={leaked} ' + f'({leaked / max(1, tot):.0%}); {len(todo)} still short of ' + f'{args.min_survivors} clean\n') + + +# --------------------------------------------------------------------------- +# Per-chunk pipeline +# --------------------------------------------------------------------------- +def process_chunk(sampler, api: Optional[OpenAIClient], chunk: List[Dict[str, Any]], + gen_dp: int, args: argparse.Namespace) -> List[Dict[str, Any]]: + """Run all 6 phases for one chunk (batched) and return per-problem records.""" + # --- Phase 1: initial solve, keep only the ones the model got wrong. --- + init_out = _run_samples( + sampler, [build_direct_prompt(r['problem']) for r in chunk], + args.init_samples, args.max_tokens, gen_dp) + for r, seqs in zip(chunk, init_out): + r['_init'] = [_parse_seq(s, r['reference_answer']) for s in seqs] + r['_init_pass'] = _pass_rate(r['_init']) + r['_failed'] = r['_init_pass'] == 0.0 + failed = [r for r in chunk if r['_failed']] + sys.stderr.write(f' phase1: {len(chunk)-len(failed)}/{len(chunk)} solved on ' + f'first try, {len(failed)} failed -> reflect\n') + + if failed: + # --- Baseline pass@k FIRST: defines which failures are genuinely hard. --- + # (A single initial rollout is noisy; an easy problem can fail phase 1 yet + # have a high pass@k, so measure the marginal only on truly hard problems.) + base_out = _run_samples( + sampler, [build_direct_prompt(r['problem']) for r in failed], + args.pass_k, args.max_tokens, gen_dp) + for r, seqs in zip(failed, base_out): + r['_baseline'] = [_parse_seq(s, r['reference_answer']) for s in seqs] + r['_baseline_pass'] = _pass_rate(r['_baseline']) + r['_hard'] = r['_baseline_pass'] <= args.hard_baseline_max + r['_skills'] = [] + r['_best'] = None + hard = [r for r in failed if r['_hard']] + sys.stderr.write(f' baseline: {len(hard)}/{len(failed)} failures are hard ' + f'(pass@{args.pass_k} <= {args.hard_baseline_max})\n') + + if hard: + # --- Skills (generate + leak filter + re-rollout) for HARD problems only. --- + _build_skills(sampler, api, hard, gen_dp, args) + + # --- With-skill pass@k (flatten hard-problem x surviving skill). --- + flat: List[Tuple[int, int]] = [] + ws_prompts: List[Any] = [] + for ri, r in enumerate(hard): + for si, sk in enumerate(r['_skills']): + if sk['leaked'] or not sk['skill'].strip(): + continue + flat.append((ri, si)) + ws_prompts.append(build_skill_solve_prompt(r['problem'], sk['skill'])) + ws_out = _run_samples(sampler, ws_prompts, args.pass_k, args.max_tokens, gen_dp) + for (ri, si), seqs in zip(flat, ws_out): + r = hard[ri] + sk = r['_skills'][si] + sk['rolls'] = [_parse_seq(s, r['reference_answer']) for s in seqs] + sk['with_pass'] = _pass_rate(sk['rolls']) + sk['marginal'] = sk['with_pass'] - r['_baseline_pass'] + + # --- Pick the best (highest marginal) surviving skill per hard problem. --- + for r in hard: + scored = [sk for sk in r['_skills'] if 'marginal' in sk] + r['_best'] = max(scored, key=lambda s: s['marginal']) if scored else None + + return [_make_record(r, args) for r in chunk] + + +def _roll_summary(roll: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: + out = {k: roll[k] for k in ('pred', 'correct', 'terminated', 'passed', + 'stop_reason', 'gen_tokens')} + if args.store_rollout_text: + out['text'] = roll['text'][:args.store_rollout_chars] + return out + + +def _make_record(r: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: + """Assemble the incremental JSONL record for one problem (solved or failed).""" + rec: Dict[str, Any] = { + 'problem': r['problem'], + 'reference_answer': r['reference_answer'], + 'tags': r.get('tags', []), + 'failed_first_try': r['_failed'], + 'init_pass_rate': r['_init_pass'], + 'init_attempt': { + 'text': r['_init'][0]['text'][:args.store_init_chars], + 'pred': r['_init'][0]['pred'], + 'stop_reason': r['_init'][0]['stop_reason'], + 'gen_tokens': r['_init'][0]['gen_tokens'], + }, + } + if not r['_failed']: + return rec + + best = r.get('_best') + rec['baseline_pass'] = r['_baseline_pass'] + # Genuinely hard = low baseline pass@k; only these count in the marginal stats. + rec['is_hard'] = bool(r.get('_hard')) + rec['baseline_rolls'] = [_roll_summary(x, args) for x in r['_baseline']] + rec['skills'] = [{ + 'skill': sk['skill'], + 'leaked': sk['leaked'], + 'leak_reason': sk['leak_reason'], + 'with_pass': sk.get('with_pass'), + 'marginal': sk.get('marginal'), + 'rolls': [_roll_summary(x, args) for x in sk.get('rolls', [])], + } for sk in r.get('_skills', [])] + rec['best_skill'] = best['skill'] if best else None + rec['best_marginal'] = best['marginal'] if best else None + rec['best_with_pass'] = best['with_pass'] if best else None + # "rescued" = a leak-free skill turned a fully-failing problem into some passes. + rec['rescued'] = bool(best and r['_baseline_pass'] == 0.0 and best['with_pass'] > 0.0) + rec['helped'] = bool(best and best['marginal'] > 0.0) + return rec + + +# --------------------------------------------------------------------------- +# Running summary +# --------------------------------------------------------------------------- +def _update_summary(summ: Dict[str, Any], recs: List[Dict[str, Any]]) -> None: + for rec in recs: + summ['n_total'] += 1 + if not rec['failed_first_try']: + summ['n_solved_first'] += 1 + continue + summ['n_failed'] += 1 + if not rec.get('is_hard'): + summ['n_failed_easy'] += 1 # failed phase 1 but easy on pass@k — excluded + continue + summ['n_hard'] += 1 + base = rec.get('baseline_pass', 0.0) + summ['sum_baseline_pass'] += base + if rec.get('best_marginal') is not None: + summ['n_with_skill'] += 1 + summ['sum_best_with_pass'] += rec.get('best_with_pass', 0.0) + summ['sum_best_marginal'] += rec.get('best_marginal', 0.0) + else: + # No clean skill produced for this hard problem -> skill adds no gain + # (count it honestly as marginal 0 rather than dropping it from the average). + summ['sum_best_with_pass'] += base + summ['n_helped'] += int(rec.get('helped', False)) + summ['n_rescued'] += int(rec.get('rescued', False)) + + +def _summary_report(summ: Dict[str, Any]) -> Dict[str, Any]: + nh = max(1, summ['n_hard']) + return { + 'record_type': 'summary', + 'n_total': summ['n_total'], + 'n_solved_first_try': summ['n_solved_first'], + 'n_failed_first_try': summ['n_failed'], + 'n_failed_but_easy': summ['n_failed_easy'], + 'n_hard': summ['n_hard'], + 'n_hard_with_skill': summ['n_with_skill'], + # Averages are over ALL hard problems; a hard problem with no clean skill + # counts as zero gain (with_pass == baseline), so with - base == marginal. + 'avg_baseline_pass_on_hard': summ['sum_baseline_pass'] / nh, + 'avg_best_with_skill_pass_on_hard': summ['sum_best_with_pass'] / nh, + 'avg_best_marginal_on_hard': summ['sum_best_marginal'] / nh, + 'n_helped_by_skill': summ['n_helped'], + 'n_rescued_from_zero': summ['n_rescued'], + 'frac_hard_helped': summ['n_helped'] / nh, + } + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- +def main() -> None: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--n', type=int, default=64, help='AoPS problems to sample.') + p.add_argument('--seed', type=int, default=42) + p.add_argument('--chunk-size', type=int, default=16, + help='Problems per chunk. All sampler calls within a chunk are ' + 'batched; results are flushed to disk after each chunk.') + p.add_argument('--init-samples', type=int, default=1, + help='Rollouts for the initial solve. A problem is "failed" (and ' + 'sent to reflection) only if all initial rollouts are wrong.') + p.add_argument('--n-skills', type=int, default=8, + help='Candidate skills generated per failed problem.') + p.add_argument('--pass-k', type=int, default=8, + help='Rollouts per (baseline / with-skill) pass@k estimate.') + p.add_argument('--hard-baseline-max', type=float, default=0.25, + help='A failed problem counts as "hard" (included in the marginal ' + 'stats) only if its baseline pass@k <= this. Filters out easy ' + 'problems that merely failed the single initial rollout.') + p.add_argument('--max-model-len', type=int, default=30000, + help='Context window (engine + template). MUST exceed --max-tokens: ' + 'the round-2 skill-gen input holds the full round-1 attempt ' + 'plus the problem.') + p.add_argument('--max-tokens', type=int, default=20000, + help='Max generated tokens for solving rollouts (round-1 output cap).') + p.add_argument('--skill-max-tokens', type=int, default=2048, + help='Max tokens for skill generation. Enough to finish any thinking ' + 'and emit the short bullet list (which is then extracted).') + p.add_argument('--attempt-reserve-tokens', type=int, default=2048, + help='Tokens reserved for system prompt + wrappers when bounding the ' + 'failed attempt fed into skill generation (the problem length ' + 'is accounted for separately, per-problem).') + p.add_argument('--min-survivors', type=int, default=2, + help='Re-roll a problem\'s skills if fewer than this many survive the ' + 'leak filters.') + p.add_argument('--skill-retries', type=int, default=1, + help='Max extra skill-generation rounds for problems short on clean ' + 'skills (0 = no retry).') + p.add_argument('--api-concurrency', type=int, default=32, + help='Parallel workers for the API leak judge (max 32 recommended).') + p.add_argument('--api-min-interval', type=float, default=0.1, + help='Minimum seconds between API leak-judge calls (QPS guard).') + p.add_argument('--api-retries', type=int, default=3, + help='Retries on transient API errors per leak-judge call (exponential ' + 'backoff); only after these are exhausted is a skill kept as ' + 'api_uncertain.') + p.add_argument('--disable-api-leak', action='store_true', + help='Skip the API leak judge even if COMPRESS_API_KEY is set ' + '(string filter only).') + p.add_argument('--output', default='./output/reflexion_phase0/aops_results.jsonl') + p.add_argument('--store-init-chars', type=int, default=8000, + help='Truncate the stored failed-attempt text to this many chars.') + p.add_argument('--store-rollout-text', action='store_true', + help='Also store (truncated) text of every rollout, not just stats.') + p.add_argument('--store-rollout-chars', type=int, default=2000) + args = p.parse_args() + + records = load_aops(n=args.n, seed=args.seed) + sys.stderr.write(f'[reflexion] {len(records)} AoPS problems, chunk={args.chunk_size}, ' + f'init_samples={args.init_samples}, n_skills={args.n_skills}, ' + f'pass_k={args.pass_k}, max_tokens={args.max_tokens}\n') + + # --- 8-GPU vLLM sampler (tp=GEN_TP, dp=GEN_GPUS/GEN_TP). --- + if GEN_GPUS % GEN_TP != 0: + raise ValueError(f'GEN_GPUS ({GEN_GPUS}) must be divisible by GEN_TP ({GEN_TP})') + gen_dp = GEN_GPUS // GEN_TP + gen_mesh = DeviceMesh.from_sizes(world_size=GEN_GPUS, dp_size=gen_dp, tp_size=GEN_TP) + twinkle.initialize( + mode='ray', nproc_per_node=GEN_GPUS, + groups=[DeviceGroup(name='sampler', ranks=list(range(GEN_GPUS)), + device_type='GPU', gpus_per_worker=GEN_TP)], + lazy_collect=False) + sampler = vLLMSampler( + model_id=GEN_MODEL_ID, + engine_args={'gpu_memory_utilization': GEN_GPU_MEM, + 'max_model_len': args.max_model_len, + 'tensor_parallel_size': GEN_TP}, + device_mesh=gen_mesh, remote_group='sampler') + sampler.set_template('Qwen3_5Template', model_id=GEN_MODEL_ID, + enable_thinking=True, max_length=args.max_model_len) + sys.stderr.write(f'[reflexion] sampler ready (model={GEN_MODEL_ID}, tp={GEN_TP}, ' + f'dp={gen_dp})\n') + + # --- API leak judge (optional): reuses eval_gpqa_rag's COMPRESS_* env. --- + api: Optional[OpenAIClient] = None + if LEAK_API_KEY and not args.disable_api_leak: + api = OpenAIClient(model=LEAK_API_MODEL, api_key=LEAK_API_KEY, + base_url=LEAK_BASE_URL) + sys.stderr.write(f'[reflexion] leak judge ON via API model={LEAK_API_MODEL} ' + f'(concurrency={args.api_concurrency})\n') + else: + sys.stderr.write('[reflexion] leak judge OFF (string filter only) — set ' + 'COMPRESS_API_KEY to enable the API judge\n') + + os.makedirs(os.path.dirname(args.output) or '.', exist_ok=True) + summ = {k: 0 for k in ('n_total', 'n_solved_first', 'n_failed', 'n_failed_easy', + 'n_hard', 'n_with_skill', 'n_helped', 'n_rescued')} + summ.update({'sum_baseline_pass': 0.0, 'sum_best_with_pass': 0.0, + 'sum_best_marginal': 0.0}) + + with open(args.output, 'w', encoding='utf-8') as out_f: + # Line 1: run config, for reproducibility / later analysis. + out_f.write(json.dumps({ + 'record_type': 'config', 'model': GEN_MODEL_ID, 'dataset': 'aops', + 'n': len(records), 'seed': args.seed, 'init_samples': args.init_samples, + 'n_skills': args.n_skills, 'pass_k': args.pass_k, + 'hard_baseline_max': args.hard_baseline_max, + 'max_model_len': args.max_model_len, 'max_tokens': args.max_tokens, + 'skill_max_tokens': args.skill_max_tokens, + 'api_leak_judge': api is not None, + 'api_leak_model': LEAK_API_MODEL if api is not None else None, + 'min_survivors': args.min_survivors, 'skill_retries': args.skill_retries, + 'gpus': GEN_GPUS, 'tp': GEN_TP, 'started': int(time.time()), + }, ensure_ascii=False) + '\n') + out_f.flush() + + n_chunks = (len(records) + args.chunk_size - 1) // args.chunk_size + for ci in range(n_chunks): + chunk = records[ci * args.chunk_size:(ci + 1) * args.chunk_size] + sys.stderr.write(f'[reflexion] chunk {ci+1}/{n_chunks} ({len(chunk)} problems)\n') + recs = process_chunk(sampler, api, chunk, gen_dp, args) + for rec in recs: # incremental write per problem + out_f.write(json.dumps(rec, ensure_ascii=False) + '\n') + out_f.flush() + _update_summary(summ, recs) + rep = _summary_report(summ) + sys.stderr.write( + f' running: failed={rep["n_failed_first_try"]} ' + f'(easy={rep["n_failed_but_easy"]}) hard={rep["n_hard"]} ' + f'base_pass={rep["avg_baseline_pass_on_hard"]:.3f} ' + f'skill_pass={rep["avg_best_with_skill_pass_on_hard"]:.3f} ' + f'helped={rep["n_helped_by_skill"]} rescued={rep["n_rescued_from_zero"]}\n') + + report = _summary_report(summ) + out_f.write(json.dumps(report, ensure_ascii=False) + '\n') + out_f.flush() + + print('\n' + '=' * 60) + print(f'Reflexion Phase-0 — model={GEN_MODEL_ID}, dataset=aops, n={report["n_total"]}') + print('=' * 60) + print(f'solved on first try : {report["n_solved_first_try"]}/{report["n_total"]}') + print(f'failed first try : {report["n_failed_first_try"]} ' + f'(easy, excluded: {report["n_failed_but_easy"]})') + print(f'hard (baseline pass@{args.pass_k}<= {args.hard_baseline_max}) : {report["n_hard"]}') + print(f' avg baseline pass@{args.pass_k:<2} : {report["avg_baseline_pass_on_hard"]:.4f}') + print(f' avg best-skill pass@{args.pass_k:<2} : {report["avg_best_with_skill_pass_on_hard"]:.4f}') + print(f' avg best marginal : {report["avg_best_marginal_on_hard"]:+.4f}') + print(f' helped by a skill : {report["n_helped_by_skill"]}/{report["n_hard"]}') + print(f' rescued from 0 pass : {report["n_rescued_from_zero"]}/{report["n_hard"]}') + print(f'\n[output] {args.output}') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/embedding/train_reflexion_skill_rft.py b/cookbook/exp/embedding/train_reflexion_skill_rft.py new file mode 100644 index 000000000..105112242 --- /dev/null +++ b/cookbook/exp/embedding/train_reflexion_skill_rft.py @@ -0,0 +1,413 @@ +"""RFT cold-start for the reflexion skill generator (see reflexion.md §6). + +Trains an INDEPENDENT skill model to write reusable, transferable skills that, +when injected into a FROZEN base solver's system prompt, raise the base's pass@k +on problems it first got wrong. The base is never trained — it only produces the +reward signal (marginal pass@k gain). This is STaR/RFT self-bootstrapping: skills +are generated online, only those with a positive marginal (and no answer leak) +are kept, and the skill model is periodically SFT-ed on them, so each round's +generator is a little better than the last. + +Direction: skill GENERATION + recall. The skill model reads the problem, the +guidance the solver had, and the solver's own attempt, then DISTILLS the genuinely +useful method / reasoning direction into a short skill list (not a mistake +diagnosis). The distilled ```` block is recalled into the base's system +prompt at solve time. + +8-GPU layout (three DeviceGroups, one twinkle.initialize): + - ranks 0-3 : ``train`` — skill model, full-param FSDP2, dp=4 + - ranks 4-5 : ``skill_sampler`` — skill model rollouts (vLLM, tp1 dp2) + - ranks 6-7 : ``base_sampler`` — frozen base solver (vLLM, tp1 dp2) +CheckpointEngineManager syncs train -> skill_sampler after every optimizer step; +base_sampler is never synced. + +Leak filtering uses ``LeakVerifier(sampler=None)`` via the backup teacher API +(no local judge, no distillation): set LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / +LLM_BACKUP_MODEL. + +Launch (8 GPUs): + LLM_BACKUP_API_KEY=... LLM_BACKUP_BASE_URL=... LLM_BACKUP_MODEL=... \ + python cookbook/exp/embedding/train_reflexion_skill_rft.py --n 2000 --chunk-size 16 +""" +import argparse +import json +import os +import random +import sys +import time +from typing import Any, Dict, List, Optional, Tuple + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.model import TransformersModel +from twinkle.processor import InputProcessor +from twinkle.sampler import vLLMSampler +from twinkle_agentic.verifier import LeakVerifier + +# Reuse the reference eval's dataset + grading + prompts + sampling config, and the +# phase-0 pipeline's parsing / rollout / injection helpers (Find > Create). +from eval_gpqa_rag import (GEN_GPU_MEM, GEN_MODEL_ID, build_direct_prompt, # noqa: F401 + load_math) +from eval_reflexion_skill import (SKILL_GEN_USER, _EX_ATTEMPT, _EX_PROBLEM, # noqa: F401 + _EX_SKILLS, _bound_attempt, _clean_text, + _pass_rate, _parse_seq, _run_samples, + build_skill_solve_prompt) + +logger = get_logger() + +# -- GPU layout --------------------------------------------------------------- +TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 4)) +SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) +BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 2)) +NUM_GPUS = TRAIN_GPUS + SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS + + +# --------------------------------------------------------------------------- +# Skill-generation prompt (DISTILL the useful approach, per the new direction) +# --------------------------------------------------------------------------- +SKILL_GEN_SYSTEM = ( + 'You are distilling a reusable problem-solving SKILL from one worked episode. ' + 'You are shown a competition problem, the guidance the solver was given, and the ' + "solver's own attempt (its reasoning may be partly right and partly wrong).\n\n" + 'Your job: summarize the GENUINELY USEFUL parts — the effective method, the ' + 'correct reasoning direction, and the kind of guidance that transfers — into a ' + 'short list of reusable skills for SIMILAR problems. Distill the useful approach; ' + 'do NOT merely criticise this attempt.\n\n' + 'OUTPUT FORMAT (strict):\n' + '- You may think first, but the final answer MUST be a markdown bullet list of 3-5 ' + 'items WRAPPED IN and tags. Output nothing after .\n' + '- Each item is ONE short imperative sentence (a method, heuristic, or check).\n' + '- Inside the tags: no narration, no "The student...", no headings, no restating ' + 'the problem.\n\n' + 'CONTENT RULES (strict):\n' + '- Do NOT reveal the final answer or the multiple-choice option.\n' + '- Do NOT state the specific numbers, values, or key intermediate results of THIS ' + 'problem.\n' + '- Every item must be GENERAL and transferable, not a step-by-step solution to ' + 'THIS problem.\n\n' + 'Follow the example below for the exact tags, style, and level of generality.' +) + + +def build_skillgen_prompt(problem: str, attempt: str) -> Dict[str, Any]: + """Skill-gen chat prompt: system + one-shot format demo + the real episode.""" + return {'messages': [ + {'role': 'system', 'content': SKILL_GEN_SYSTEM}, + {'role': 'user', + 'content': SKILL_GEN_USER.format(problem=_EX_PROBLEM, attempt=_EX_ATTEMPT)}, + {'role': 'assistant', 'content': _EX_SKILLS}, + {'role': 'user', + 'content': SKILL_GEN_USER.format(problem=problem, attempt=attempt)}, + ]} + + +def _extract_skills_block(text: str) -> Optional[str]: + """Return the inner ``...`` block, or None if not parseable. + + RFT keeps thinking ON, so a candidate is usable only if it actually produced a + CLOSED tag block; unterminated / tag-less generations are dropped (never fall + back to scraping raw reasoning) — see reflexion.md §6.8. + """ + low = text.lower() + if '' not in low or '' not in low: + return None + start = low.index('') + len('') + end = low.index('') + if end <= start: + return None + block = text[start:end].strip() + return block or None + + +def _bounded_attempt(r: Dict[str, Any], args: argparse.Namespace) -> str: + """Bound the failed attempt so problem + attempt + skill output fits context. + + Capped by ``--attempt-max-tokens`` so the SFT sequences (which embed the same + attempt) stay short enough to train; the same bound is used at generation and + SFT time, keeping train/inference identical. + """ + prob_est = len(r['problem']) // 2 + budget = max(1024, args.max_model_len - args.skill_max_tokens + - args.attempt_reserve_tokens - prob_est) + budget = min(budget, args.attempt_max_tokens) + init = r['_init'][0] + return _bound_attempt(init['text'], init['gen_tokens'], budget) + + +# --------------------------------------------------------------------------- +# Online data generation (one chunk, all sampler calls batched) +# --------------------------------------------------------------------------- +def generate_chunk(base_sampler, skill_sampler, leak: LeakVerifier, + chunk: List[Dict[str, Any]], base_dp: int, skill_dp: int, + args: argparse.Namespace) -> List[Dict[str, Any]]: + """Run base-solve -> hard-filter -> skill-gen -> leak-filter -> marginal, and + return the records whose skill gave a positive marginal pass@k gain.""" + # --- Phase 1: base solves once; keep only what it got wrong. --- + init_out = _run_samples( + base_sampler, [build_direct_prompt(r['problem']) for r in chunk], + args.init_samples, args.max_tokens, base_dp) + failed = [] + for r, seqs in zip(chunk, init_out): + r['_init'] = [_parse_seq(s, r['reference_answer']) for s in seqs] + if _pass_rate(r['_init']) == 0.0: + failed.append(r) + sys.stderr.write(f' phase1: {len(chunk)-len(failed)}/{len(chunk)} solved on ' + f'first try, {len(failed)} failed\n') + if not failed: + return [] + + # --- Phase 2: baseline pass@k on the failures -> keep the genuinely hard. --- + base_out = _run_samples( + base_sampler, [build_direct_prompt(r['problem']) for r in failed], + args.pass_k, args.max_tokens, base_dp) + hard = [] + for r, seqs in zip(failed, base_out): + r['_baseline_pass'] = _pass_rate([_parse_seq(s, r['reference_answer']) for s in seqs]) + if r['_baseline_pass'] <= args.hard_baseline_max: + hard.append(r) + sys.stderr.write(f' phase2: {len(hard)}/{len(failed)} failures are hard ' + f'(pass@{args.pass_k} <= {args.hard_baseline_max})\n') + if not hard: + return [] + + # --- Phase 3: skill-gen (thinking ON) -> require a parseable block. --- + sg_out = _run_samples( + skill_sampler, + [build_skillgen_prompt(r['problem'], _bounded_attempt(r, args)) for r in hard], + args.n_skills, args.skill_max_tokens, skill_dp) + cands: List[Dict[str, Any]] = [] # {r, response, block} + for r, seqs in zip(hard, sg_out): + for s in seqs: + resp = _clean_text(getattr(s, 'decoded', '') or '') + block = _extract_skills_block(resp) + if block: + cands.append({'r': r, 'response': resp, 'block': block}) + if not cands: + sys.stderr.write(' phase3: 0 parseable skill blocks\n') + return [] + + # --- Phase 4: leak filter via backup teacher (drops answer-leaking skills). --- + details = leak.leak_batch( + [{'content': c['block'], 'query': c['r']['problem'], + 'reference': c['r']['reference_answer']} for c in cands], + max_workers=args.leak_workers) + clean = [c for c, d in zip(cands, details) if not d.leaked] + sys.stderr.write(f' phase4: {len(clean)}/{len(cands)} skills clean ' + f'({len(cands)-len(clean)} leaked)\n') + if not clean: + return [] + + # --- Phase 5: with-skill pass@k -> marginal = with - baseline. --- + ws_out = _run_samples( + base_sampler, + [build_skill_solve_prompt(c['r']['problem'], c['block']) for c in clean], + args.pass_k, args.max_tokens, base_dp) + records: List[Dict[str, Any]] = [] + for c, seqs in zip(clean, ws_out): + r = c['r'] + with_pass = _pass_rate([_parse_seq(s, r['reference_answer']) for s in seqs]) + marginal = with_pass - r['_baseline_pass'] + if marginal > 0.0: # keep only skills that actually helped + records.append({ + 'problem': r['problem'], + 'reference_answer': r['reference_answer'], + 'attempt': _bounded_attempt(r, args), + 'response': c['response'], + 'skills': c['block'], + 'baseline_pass': r['_baseline_pass'], + 'with_pass': with_pass, + 'marginal': marginal, + }) + sys.stderr.write(f' phase5: {len(records)} skills with positive marginal kept\n') + return records + + +# --------------------------------------------------------------------------- +# Online RFT training +# --------------------------------------------------------------------------- +def _sft_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: + """SFT sample = the exact skill-gen prompt + the kept generation as the target. + + Matching the generation prompt keeps train/inference consistent; the template + masks the prompt to -100 and supervises the assistant turns. + """ + msgs = build_skillgen_prompt(rec['problem'], rec['attempt'])['messages'] + return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}]} + + +def _train_round(skill_model, ckpt: CheckpointEngineManager, + pool: List[Dict[str, Any]], args: argparse.Namespace) -> None: + """One RFT round: weighted-sample ``train_batch`` records, SFT them in driver- + side mini-batches (each an optimizer step, matching short_math_grpo), then sync. + + The transformers backend runs one forward per ``forward_backward`` call (no + internal micro split), and ``slice_dp`` needs each mini-batch divisible across + the training dp ranks — hence ``sft_batch_size`` is a multiple of TRAIN_GPUS. + """ + weights = [max(rec['marginal'], args.weight_eps) ** args.weight_alpha for rec in pool] + batch = random.choices(pool, weights=weights, k=args.train_batch) + trajs = [_sft_trajectory(rec) for rec in batch] + for i in range(0, len(trajs), args.sft_batch_size): + skill_model.forward_backward(inputs=trajs[i:i + args.sft_batch_size]) + skill_model.clip_grad_and_step() + ckpt.sync_weights(merge_and_sync=True) # full-param: push all weights to vLLM + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- +def _build_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--n', type=int, default=2000, help='MATH problems to stream.') + p.add_argument('--seed', type=int, default=42) + p.add_argument('--chunk-size', type=int, default=16, + help='Problems per generation chunk (all sampler calls batched).') + p.add_argument('--init-samples', type=int, default=1) + p.add_argument('--n-skills', type=int, default=8, + help='Candidate skills generated per hard problem.') + p.add_argument('--pass-k', type=int, default=8) + p.add_argument('--hard-baseline-max', type=float, default=0.25) + p.add_argument('--max-model-len', type=int, default=16384) + p.add_argument('--max-tokens', type=int, default=8192, + help='Max generated tokens for solve rollouts.') + p.add_argument('--skill-max-tokens', type=int, default=4096, + help='Max tokens for skill-gen (thinking ON: room for CoT + the ' + ' block).') + p.add_argument('--attempt-reserve-tokens', type=int, default=1024) + p.add_argument('--attempt-max-tokens', type=int, default=3072, + help='Hard cap on the failed-attempt tokens embedded in skill-gen ' + 'and SFT prompts (keeps SFT sequences trainable).') + p.add_argument('--leak-workers', type=int, default=32, + help='Parallel workers for the LeakVerifier backup judge.') + # -- online RFT -- + p.add_argument('--train-every', type=int, default=64, + help='Trigger one train round after this many new valid records.') + p.add_argument('--train-batch', type=int, default=64, + help='Records weighted-sampled (with replacement) per train round.') + p.add_argument('--sft-batch-size', type=int, default=8, + help='Driver-side mini-batch per optimizer step; MUST be a multiple ' + 'of TRAIN_GPUS (sliced across training dp ranks) and divide ' + '--train-batch.') + p.add_argument('--lr', type=float, default=1e-5) + p.add_argument('--max-train-rounds', type=int, default=200, + help='Cap on train rounds (also sizes the LR schedule).') + p.add_argument('--save-rounds', type=int, default=50) + p.add_argument('--weight-alpha', type=float, default=1.0, + help='Sampling weight exponent: w = max(marginal, eps) ** alpha.') + p.add_argument('--weight-eps', type=float, default=0.01) + p.add_argument('--output-dir', default='./output/reflexion_skill_rft') + return p.parse_args() + + +def main() -> None: + args = _build_args() + if args.sft_batch_size % TRAIN_GPUS != 0 or args.train_batch % args.sft_batch_size != 0: + raise ValueError(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple ' + f'of TRAIN_GPUS ({TRAIN_GPUS}) and divide --train-batch ' + f'({args.train_batch})') + steps_per_round = args.train_batch // args.sft_batch_size + records = load_math(n=args.n, seed=args.seed) + os.makedirs(args.output_dir, exist_ok=True) + data_path = os.path.join(args.output_dir, 'skill_dataset.jsonl') + + if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY')): + sys.stderr.write('[rft] WARNING: no LLM_BACKUP_API_KEY/OPENAI_API_KEY — ' + 'LeakVerifier will report no_llm and skip leak filtering\n') + + # -- Device groups: train (FSDP2) + two independent vLLM samplers. -- + r0, r1, r2 = TRAIN_GPUS, TRAIN_GPUS + SKILL_SAMPLER_GPUS, NUM_GPUS + device_groups = [ + DeviceGroup(name='train', ranks=list(range(0, r0)), device_type='GPU'), + DeviceGroup(name='skill_sampler', ranks=list(range(r0, r1)), device_type='GPU'), + DeviceGroup(name='base_sampler', ranks=list(range(r1, r2)), device_type='GPU'), + ] + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, groups=device_groups, + lazy_collect=False) + + # -- Skill model: full-param FSDP2, causal-LM SFT. -- + train_mesh = DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_GPUS) + skill_model = TransformersModel(model_id=GEN_MODEL_ID, device_mesh=train_mesh, + remote_group='train', + ddp_config={'find_unused_parameters': False}) + from twinkle.patch.no_split_modules import NoSplitModulesPatch + skill_model.apply_patch(NoSplitModulesPatch({'Qwen3_5DecoderLayer'})) + skill_model.set_template('Qwen3_5Template', model_id=GEN_MODEL_ID, + enable_thinking=True, max_length=args.max_model_len, + truncation_strategy='delete') + skill_model.set_processor(InputProcessor, padding_free=True) + skill_model.set_loss('CrossEntropyLoss') + skill_model.set_optimizer('AdamW', lr=args.lr) + skill_model.set_lr_scheduler('CosineWarmupScheduler', num_warmup_steps=10, + num_training_steps=args.max_train_rounds * steps_per_round) + + # -- Two vLLM samplers: skill (synced) + base (frozen). -- + skill_dp, base_dp = SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS + skill_sampler = vLLMSampler( + model_id=GEN_MODEL_ID, + engine_args={'gpu_memory_utilization': GEN_GPU_MEM, + 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, + device_mesh=DeviceMesh.from_sizes(world_size=SKILL_SAMPLER_GPUS, dp_size=skill_dp), + remote_group='skill_sampler') + skill_sampler.set_template('Qwen3_5Template', model_id=GEN_MODEL_ID, + enable_thinking=True, max_length=args.max_model_len) + base_sampler = vLLMSampler( + model_id=GEN_MODEL_ID, + engine_args={'gpu_memory_utilization': GEN_GPU_MEM, + 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, + device_mesh=DeviceMesh.from_sizes(world_size=BASE_SAMPLER_GPUS, dp_size=base_dp), + remote_group='base_sampler') + base_sampler.set_template('Qwen3_5Template', model_id=GEN_MODEL_ID, + enable_thinking=True, max_length=args.max_model_len) + + ckpt = CheckpointEngineManager(model=skill_model, sampler=skill_sampler) + leak = LeakVerifier(sampler=None) # backup teacher only, no local judge + + sys.stderr.write(f'[rft] {len(records)} MATH problems; train_gpus={TRAIN_GPUS} ' + f'skill_dp={skill_dp} base_dp={base_dp}\n') + + pool: List[Dict[str, Any]] = [] + pending = 0 + rounds = 0 + n_chunks = (len(records) + args.chunk_size - 1) // args.chunk_size + + with open(data_path, 'w', encoding='utf-8') as fout: + fout.write(json.dumps({ + 'record_type': 'config', 'model': GEN_MODEL_ID, 'dataset': 'math', + 'n': len(records), 'seed': args.seed, 'pass_k': args.pass_k, + 'hard_baseline_max': args.hard_baseline_max, 'n_skills': args.n_skills, + 'train_every': args.train_every, 'train_batch': args.train_batch, + 'lr': args.lr, 'started': int(time.time()), + }, ensure_ascii=False) + '\n') + fout.flush() + + for ci in range(n_chunks): + if rounds >= args.max_train_rounds: + break + chunk = records[ci * args.chunk_size:(ci + 1) * args.chunk_size] + sys.stderr.write(f'[rft] chunk {ci+1}/{n_chunks} ({len(chunk)} problems)\n') + recs = generate_chunk(base_sampler, skill_sampler, leak, chunk, + base_dp, skill_dp, args) + for rec in recs: + fout.write(json.dumps(rec, ensure_ascii=False) + '\n') + fout.flush() + pool.extend(recs) + pending += len(recs) + + while pending >= args.train_every and rounds < args.max_train_rounds: + _train_round(skill_model, ckpt, pool, args) + pending -= args.train_every + rounds += 1 + sys.stderr.write(f'[rft] train round {rounds}/{args.max_train_rounds} ' + f'(pool={len(pool)})\n') + if rounds % args.save_rounds == 0: + skill_model.save(f'skill-rft-{rounds}', output_dir=args.output_dir) + + skill_model.save('skill-rft-final', output_dir=args.output_dir) + sys.stderr.write(f'[rft] done: {len(pool)} valid records, {rounds} train rounds; ' + f'dataset -> {data_path}\n') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/embedding/train_reflexion_skill_rft.sh b/cookbook/exp/embedding/train_reflexion_skill_rft.sh new file mode 100644 index 000000000..7ad6a607a --- /dev/null +++ b/cookbook/exp/embedding/train_reflexion_skill_rft.sh @@ -0,0 +1,25 @@ +#!/bin/bash +# RFT cold-start for the reflexion skill generator (see reflexion.md §6). +# GPUs: 8 — ranks 0-3 train (skill model, FSDP2), 4-5 skill sampler, 6-7 base sampler. +# Leak filtering uses the backup teacher API (no local judge): set LLM_BACKUP_*. + +set -euo pipefail + +export GEN_MODEL_ID=${GEN_MODEL_ID:-Qwen/Qwen3.5-4B} +export MATH_DATA_DIR=${MATH_DATA_DIR:-./output/math_data/MATH} +export LLM_BACKUP_API_KEY=${LLM_BACKUP_API_KEY:?set LLM_BACKUP_API_KEY for the leak judge} +export LLM_BACKUP_BASE_URL=${LLM_BACKUP_BASE_URL:-https://dashscope.aliyuncs.com/compatible-mode/v1} +export LLM_BACKUP_MODEL=${LLM_BACKUP_MODEL:-qwen3.7-max} + +python cookbook/exp/embedding/train_reflexion_skill_rft.py \ + --n 2000 \ + --chunk-size 16 \ + --n-skills 8 \ + --pass-k 8 \ + --hard-baseline-max 0.25 \ + --train-every 64 \ + --train-batch 64 \ + --sft-batch-size 8 \ + --lr 1e-5 \ + --max-train-rounds 200 \ + --output-dir ./output/reflexion_skill_rft diff --git a/src/twinkle_agentic/verifier/__init__.py b/src/twinkle_agentic/verifier/__init__.py index 0cdd24bc2..38e14f601 100644 --- a/src/twinkle_agentic/verifier/__init__.py +++ b/src/twinkle_agentic/verifier/__init__.py @@ -8,6 +8,7 @@ check_numeric_equiv, check_output_format, default_checks_for) from .hard_scorer import CheckResult, HardScorer, HardScoreDetail, TrajectoryView +from .leak_verifier import LeakDetail, LeakVerifier from .rubric_library import (INTENT_BASE_RUBRICS, INTENT_FIXED_RUBRICS, default_intent_base_rubrics, default_intent_fixed_rubrics) @@ -18,6 +19,7 @@ 'Verifier', 'RubricVerifier', 'RubricItem', 'ScoreDetail', 'DiagnoseDetail', 'DiagnosisItem', + 'LeakVerifier', 'LeakDetail', 'HardScorer', 'HardScoreDetail', 'CheckResult', 'TrajectoryView', 'check_output_format', 'check_numeric_equiv', 'check_answer_match', 'check_code_parses', 'check_instruction_constraints', 'check_not_degenerate', diff --git a/src/twinkle_agentic/verifier/leak_verifier.py b/src/twinkle_agentic/verifier/leak_verifier.py new file mode 100644 index 000000000..f1301c66d --- /dev/null +++ b/src/twinkle_agentic/verifier/leak_verifier.py @@ -0,0 +1,329 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""General answer-leak verifier for auxiliary hints / skills / notes. + +Decides whether a piece of *auxiliary content* — a hint, a distilled skill, a +retrieved note, a rationale — that will be shown to a solver ALONGSIDE a task +leaks the answer (or an essentially complete solution) to THAT task. A good hint +carries only transferable strategy ("factor the radicand before adding roots"); +a leaking hint hands over the result, a task-specific decisive step, or a full +derivation. + +**One layer only: an LLM judge.** Leak detection is inherently *semantic* — the +common case is a hint that describes the solution structure without ever writing +the answer, which no string rule can catch. A cheap deterministic pre-filter +(verbatim answer / answer-number matching) was measured to replace only a few +percent of the judge's catches, is not domain-general, and false-flags short +answers ("D", "1") that appear in almost any text. So it is deliberately absent: +this verifier is a single ``llm_backup``-distilled LEAK / CLEAN judge that runs +on the student model when confident and falls back to the teacher API otherwise +— the same progressive-distillation path the other verifiers here use. + +The :class:`Verifier` contract (``__call__ -> int``) maps CLEAN -> ``NUM_LEVELS-1`` +(a clean hint is the "good" end of the reward scale) and LEAK -> ``0``. +:meth:`leak_detail` exposes the boolean + reason, and :meth:`leak_batch` runs +many checks in parallel (the judge dominates latency). +""" +from __future__ import annotations + +import os +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Union + +from twinkle_agentic.utils.llm_backup import llm_backup + +from .base import Verifier +from .domain_checks import _ground_truths + +if TYPE_CHECKING: + from twinkle.data_format.sampling import SamplingParams # noqa: F401 + from twinkle.sampler.base import Sampler # noqa: F401 + + +# --------------------------------------------------------------------------- +# Prompts (task-agnostic; a reference answer is optional context for the judge) +# --------------------------------------------------------------------------- +_JUDGE_SYSTEM = """\ +You check whether a HINT that will be shown to someone solving a TASK leaks the \ +answer. + +A good hint gives ONLY transferable strategy — general methods, common \ +pitfalls, or sanity checks that would help on a whole class of similar tasks. + +The hint LEAKS if it does ANY of: +- reveals the final answer or final result (a value, expression, choice, label, \ +or verbatim output); +- states a specific decisive intermediate result or fact that is unique to THIS \ +task; +- gives a derivation or step-by-step plan that essentially solves THIS task. + +The hint does NOT leak if it only names general methods, common mistakes, or \ +checks that are not specific to this task's answer. + +Reply with exactly one word: LEAK or CLEAN.""" + +_JUDGE_USER = """\ +## Task +{query} +{reference_block} +## Hint to check +{hint} + +Does the hint leak the answer or an essentially complete solution to THIS task? \ +Reply LEAK or CLEAN.""" + +_REFERENCE_BLOCK = """ +## Known answer (for your judgement only; do not treat its wording as the hint) +{reference} +""" + + +# --------------------------------------------------------------------------- +# Result holder +# --------------------------------------------------------------------------- +@dataclass +class LeakDetail: + """Outcome of one leak check. + + Attributes: + leaked: True if the content leaks the answer/solution. + reason: machine-readable reason — ``''`` (clean), ``'llm_leak'``, + ``'llm_uncertain'`` (judge reply unparseable), or ``'no_llm'`` (no + student sampler and no teacher API configured). + source: which layer decided — ``'llm'`` | ``'none'``. + """ + leaked: bool + reason: str + source: str + + +# --------------------------------------------------------------------------- +# Verdict parsing / comparator (for the judge + its distillation) +# --------------------------------------------------------------------------- +def _verdict_of(raw: str) -> Optional[bool]: + """Parse a LEAK/CLEAN reply -> True (leak) / False (clean) / None (unclear).""" + v = (raw or '').strip().upper() + if 'CLEAN' in v: + return False + if 'LEAK' in v: + return True + return None + + +def _verdict_close(a: str, b: str) -> bool: + """llm_backup comparator: student/teacher agree iff same LEAK/CLEAN verdict.""" + va, vb = _verdict_of(a), _verdict_of(b) + if va is None or vb is None: + return (a or '').strip() == (b or '').strip() + return va == vb + + +# --------------------------------------------------------------------------- +# Verifier +# --------------------------------------------------------------------------- +class LeakVerifier(Verifier): + """Verify that an auxiliary hint does not leak the answer to its task. + + Args: + sampler: student model sampler (local inference). If ``None`` the judge + is served entirely by the teacher API via ``llm_backup`` (useful + before a student exists); if no teacher is configured either, the + verifier reports ``no_llm`` and cannot judge. + model_path: identifier (bookkeeping only). + sampling_params: default sampling params for the judge call. + judge_lora_path: LoRA adapter for the distilled judge student. + max_content_chars / max_query_chars: truncation caps for the judge input. + uncertain_is_leak: if the judge reply is unparseable, treat it as a leak + (default False: keep the hint, the conservative choice). + """ + + def __init__( + self, + sampler: Optional['Sampler'] = None, + *, + model_path: str = '', + sampling_params: Optional['SamplingParams'] = None, + judge_lora_path: Optional[str] = None, + max_content_chars: int = 4000, + max_query_chars: int = 4000, + uncertain_is_leak: bool = False, + ): + self.sampler = sampler + self.model_path = model_path + self.sampling_params = sampling_params + self.judge_lora_path = judge_lora_path or None + self.max_content_chars = int(max_content_chars) + self.max_query_chars = int(max_query_chars) + self.uncertain_is_leak = bool(uncertain_is_leak) + + # ------------------------------------------------------------------ + # public entry points + # ------------------------------------------------------------------ + def __call__(self, trajectory: dict, *, query: Optional[str] = None, + reference: Optional[Union[str, Sequence[str]]] = None, + **kwargs) -> int: + detail = self.leak_detail( + self._content_of(trajectory), + query=query or self._infer_query(trajectory), + reference=reference if reference is not None else _ground_truths(trajectory), + ) + return 0 if detail.leaked else self.NUM_LEVELS - 1 + + def is_leak(self, content: str, *, query: str, + reference: Optional[Union[str, Sequence[str]]] = None) -> bool: + return self.leak_detail(content, query=query, reference=reference).leaked + + def leak_detail(self, content: str, *, query: str, + reference: Optional[Union[str, Sequence[str]]] = None + ) -> LeakDetail: + """LLM judge; ``no_llm`` when neither a student nor a teacher exists.""" + content = content or '' + references = self._as_list(reference) + + if not self._llm_available(): + return LeakDetail(False, 'no_llm', 'none') + + verdict = self._judge(content, query or '', references) + if verdict is True: + return LeakDetail(True, 'llm_leak', 'llm') + if verdict is False: + return LeakDetail(False, '', 'llm') + return LeakDetail(self.uncertain_is_leak, 'llm_uncertain', 'llm') + + def leak_batch(self, items: Sequence[dict], *, max_workers: int = 8 + ) -> List[LeakDetail]: + """Check many hints in parallel; each item is ``{content, query, reference?}``. + + The judge (network / student inference) dominates latency, so the checks + fan out over a thread pool. Results are returned in input order. + """ + items = list(items) + if not items: + return [] + workers = max(1, min(max_workers, len(items))) + if workers == 1: + return [self._leak_detail_item(it) for it in items] + with ThreadPoolExecutor(max_workers=workers) as pool: + return list(pool.map(self._leak_detail_item, items)) + + def _leak_detail_item(self, item: dict) -> LeakDetail: + return self.leak_detail(item.get('content', ''), query=item.get('query', ''), + reference=item.get('reference')) + + # ------------------------------------------------------------------ + # judge (student, distilled via llm_backup) + # ------------------------------------------------------------------ + def _judge(self, content: str, query: str, + references: Sequence[str]) -> Optional[bool]: + trajectory = self._judge_trajectory(content, query, references) + raw = self._judge_once( + trajectory=trajectory, + sampling_params=self._judge_sampling_params(self.sampling_params), + judge_key='leak') + return _verdict_of(raw) + + # Distilled on the LEAK/CLEAN verdict. A single shared key ('leak') tracks + # student/teacher agreement on the judging skill as a whole; the comparator + # matches on the verdict, not byte-identical text. + @llm_backup(key_params=['judge_key'], comparator=_verdict_close) + def _judge_once(self, trajectory, sampling_params, judge_key: str = 'leak') -> str: + return self._sample_text(trajectory, sampling_params, self.judge_lora_path) + + def _judge_trajectory(self, content: str, query: str, + references: Sequence[str]) -> dict: + ref_block = '' + refs = [r for r in references if (r or '').strip()] + if refs: + ref_block = _REFERENCE_BLOCK.format(reference='; '.join(refs)) + user = _fill( + _JUDGE_USER, + query=self._trim(query, self.max_query_chars), + reference_block=ref_block, + hint=self._trim(content, self.max_content_chars)) + return {'messages': [ + {'role': 'system', 'content': _JUDGE_SYSTEM}, + {'role': 'user', 'content': user}, + ]} + + def _judge_sampling_params(self, override): + if override is not None: + return override + from twinkle.data_format.sampling import SamplingParams + # A one-word verdict; keep the budget tiny. Small headroom absorbs models + # that prepend a stray token before LEAK/CLEAN. + return SamplingParams(temperature=0.0, max_tokens=8) + + # ------------------------------------------------------------------ + # LLM plumbing (mirrors RubricVerifier) + # ------------------------------------------------------------------ + def _llm_available(self) -> bool: + if self.sampler is not None: + return True + return bool(os.environ.get('LLM_BACKUP_API_KEY') + or os.environ.get('OPENAI_API_KEY') + or os.environ.get('LLM_BACKUP_BASE_URL')) + + def _sample_text(self, trajectory, sampling_params, lora_path) -> str: + if self.sampler is None: + return '' + sample_kwargs: dict[str, Any] = {'sampling_params': sampling_params} + if lora_path is None: + sample_kwargs['use_base_model'] = True + else: + sample_kwargs['adapter_path'] = lora_path + responses = self.sampler.sample([trajectory], **sample_kwargs) + resp = list(responses)[0] if responses else None + if resp is None: + return '' + seqs = getattr(resp, 'sequences', None) or [] + return (getattr(seqs[0], 'decoded', None) or '') if seqs else '' + + # ------------------------------------------------------------------ + # small helpers + # ------------------------------------------------------------------ + @staticmethod + def _as_list(reference: Optional[Union[str, Sequence[str]]]) -> List[str]: + if reference is None: + return [] + if isinstance(reference, str): + return [reference] + return [str(r) for r in reference] + + @staticmethod + def _trim(text: str, cap: int) -> str: + text = text or '' + return text if len(text) <= cap else text[:cap] + + @staticmethod + def _content_of(trajectory: dict) -> str: + """The hint under check = last assistant message text (else any content).""" + if isinstance(trajectory, str): + return trajectory + msgs = trajectory.get('messages', []) or [] + for m in reversed(msgs): + if m.get('role') == 'assistant': + c = m.get('content') + if isinstance(c, list): + c = '\n'.join(p.get('text', '') for p in c + if isinstance(p, dict) and p.get('type') == 'text') + if isinstance(c, str) and c.strip(): + return c + return str(trajectory.get('content', '') or '') + + @staticmethod + def _infer_query(trajectory: dict) -> str: + if isinstance(trajectory, str): + return '(no explicit task)' + for m in trajectory.get('messages', []) or []: + if m.get('role') == 'user': + c = m.get('content') + if isinstance(c, str) and c.strip(): + return c.strip() + return '(no explicit task)' + + +def _fill(template: str, **kw) -> str: + out = template + for k, v in kw.items(): + out = out.replace('{' + k + '}', str(v)) + return out diff --git a/src/twinkle_agentic/verifier/rubric_verifier.py b/src/twinkle_agentic/verifier/rubric_verifier.py index e8a3552d3..8e10b19da 100644 --- a/src/twinkle_agentic/verifier/rubric_verifier.py +++ b/src/twinkle_agentic/verifier/rubric_verifier.py @@ -151,6 +151,9 @@ user-facing output; ignore it for "output only X" style criteria. - For PASS items, leave "fix" as "". For FAIL items, "fix" must be a concrete \ correction (e.g. add the missing argument, redo step k). +- Keep every "reason" and "fix" clear and concise — one short sentence each, \ +stating only the essential point; do not restate the criterion, quote the segment \ +at length, or add filler. - "overall" is "OK" only if NO criterion is FAIL. - Output only the JSON object.""" From 355c45a5d3c146e309cdda331189750b6ce11234 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Mon, 13 Jul 2026 20:29:33 +0800 Subject: [PATCH 07/60] wip --- .../exp/embedding/eval_reflexion_skill.py | 35 +- .../embedding/train_reflexion_skill_rft.py | 1144 +++++++++++++---- .../embedding/train_reflexion_skill_rft.sh | 23 +- src/twinkle_agentic/verifier/leak_verifier.py | 36 +- 4 files changed, 975 insertions(+), 263 deletions(-) diff --git a/cookbook/exp/embedding/eval_reflexion_skill.py b/cookbook/exp/embedding/eval_reflexion_skill.py index db88ca13a..3f7ddcc0d 100644 --- a/cookbook/exp/embedding/eval_reflexion_skill.py +++ b/cookbook/exp/embedding/eval_reflexion_skill.py @@ -265,18 +265,18 @@ def _skill_leaks(skill: str, gold: str) -> Tuple[bool, str]: # --- API leak judge: catches what the string filter cannot (MCQ letters, a # derived key result, or a near-complete solution laid out as a "skill"). --- _LEAK_JUDGE_SYSTEM = ( - 'You are a strict grader deciding whether a "skill" hint LEAKS the solution to a ' - 'math problem. A skill should be a GENERAL, transferable reminder. It LEAKS if it ' - 'does ANY of: reveal the final answer (a number, expression, or multiple-choice ' - 'option); state a specific numeric/geometric result or key intermediate value of ' - 'THIS problem; or give a derivation that essentially solves THIS problem. It does ' - 'NOT leak if it only names general methods, common pitfalls, or checks. Reply with ' - 'exactly one word: LEAK or CLEAN.' + 'You decide whether a "skill" hint LEAKS the FINAL ANSWER to a math problem. ' + 'A skill may freely name the general method, the right technique, the solution ' + 'approach, the steps to take, or common pitfalls — revealing the METHOD is ' + 'allowed and expected. It LEAKS only if it reveals the FINAL ANSWER itself: the ' + 'concrete final number, expression, or multiple-choice option the problem asks ' + 'for (or a trivially equivalent restatement). Describing HOW to solve it WITHOUT ' + 'stating the resulting final value does NOT leak. Reply with exactly one word: ' + 'LEAK or CLEAN.' ) _LEAK_JUDGE_USER = ( 'Problem:\n{problem}\n\nGold answer: {gold}\n\nSkill hint to check:\n{skill}\n\n' - 'Does the skill leak the answer or a full solution to THIS problem? Reply LEAK or ' - 'CLEAN.' + 'Does the skill reveal the FINAL ANSWER to THIS problem? Reply LEAK or CLEAN.' ) @@ -339,13 +339,22 @@ def _pad_for_dp(prompts: List[Any], gen_dp: int) -> List[Any]: def _run_samples(sampler, prompts: List[Any], num_samples: int, max_tokens: int, - gen_dp: int) -> List[List[Any]]: - """One batched sampler call; return per-prompt list of raw sampled sequences.""" + gen_dp: int, temperature: Optional[float] = None, + top_p: Optional[float] = None, + top_k: Optional[int] = None) -> List[List[Any]]: + """One batched sampler call; return per-prompt list of raw sampled sequences. + + ``temperature``/``top_p``/``top_k`` default to the module's sampling config; pass + ``temperature=0.0`` for deterministic greedy decoding (SEAM-style executor scoring), + or a high ``temperature`` with ``top_k=-1`` for diverse multi-candidate sampling.""" if not prompts: return [] params = TwinkleSamplingParams( - max_tokens=max_tokens, temperature=GEN_TEMPERATURE, top_p=GEN_TOP_P, - num_samples=num_samples) + max_tokens=max_tokens, + temperature=GEN_TEMPERATURE if temperature is None else temperature, + top_p=GEN_TOP_P if top_p is None else top_p, + num_samples=num_samples, + **({} if top_k is None else {'top_k': top_k})) padded = _pad_for_dp(prompts, gen_dp) responses = sampler.sample(padded, params)[:len(prompts)] return [list(r.sequences) if (r and r.sequences) else [] for r in responses] diff --git a/cookbook/exp/embedding/train_reflexion_skill_rft.py b/cookbook/exp/embedding/train_reflexion_skill_rft.py index 105112242..33bf4dc20 100644 --- a/cookbook/exp/embedding/train_reflexion_skill_rft.py +++ b/cookbook/exp/embedding/train_reflexion_skill_rft.py @@ -1,18 +1,21 @@ """RFT cold-start for the reflexion skill generator (see reflexion.md §6). Trains an INDEPENDENT skill model to write reusable, transferable skills that, -when injected into a FROZEN base solver's system prompt, raise the base's pass@k -on problems it first got wrong. The base is never trained — it only produces the -reward signal (marginal pass@k gain). This is STaR/RFT self-bootstrapping: skills -are generated online, only those with a positive marginal (and no answer leak) -are kept, and the skill model is periodically SFT-ed on them, so each round's -generator is a little better than the last. - -Direction: skill GENERATION + recall. The skill model reads the problem, the -guidance the solver had, and the solver's own attempt, then DISTILLS the genuinely -useful method / reasoning direction into a short skill list (not a mistake -diagnosis). The distilled ```` block is recalled into the base's system -prompt at solve time. +when injected into a FROZEN base solver's system prompt, let the base solve problems +it first got wrong. The base is never trained — it only produces the reward signal. +Scoring is SEAM-style DETERMINISTIC: the base runs each candidate skill once at +temperature 0 (M=1), so the reward ``R in {0,1}`` (answer correct) carries no +sampling noise; the per-candidate advantage is group-relative within a problem +(``A = (R - mean) / (std + eps)``) and the skill model is updated online by GRPO — +problem-groups where every skill scores alike (std=0) contribute no gradient. + +Direction: skill GENERATION + recall. Skill-gen always runs with thinking ON, and +each hard problem is routed to EXACTLY ONE of two views (no reuse — kills memory +leak and holds cost at 1x): view A ``(problem + attempt) -> think + skills`` keeps +the online generator self-bootstrapping; view B ``(problem only) -> think + skills`` +is the deployment form, where the think is grounded on the query alone so it cannot +hallucinate an attempt. Both share the verified skill; the distilled ```` +block is recalled into the base's system prompt at solve time. 8-GPU layout (three DeviceGroups, one twinkle.initialize): - ranks 0-3 : ``train`` — skill model, full-param FSDP2, dp=4 @@ -30,228 +33,727 @@ python cookbook/exp/embedding/train_reflexion_skill_rft.py --n 2000 --chunk-size 16 """ import argparse +import hashlib import json import os -import random +import re import sys import time from typing import Any, Dict, List, Optional, Tuple +import numpy as np + import twinkle from twinkle import DeviceGroup, DeviceMesh, get_logger from twinkle.checkpoint_engine import CheckpointEngineManager from twinkle.model import TransformersModel from twinkle.processor import InputProcessor from twinkle.sampler import vLLMSampler -from twinkle_agentic.verifier import LeakVerifier +from twinkle.template import Template +from twinkle_agentic.verifier import LeakVerifier, RubricVerifier +from twinkle_agentic.verifier.rubric_verifier import RubricItem # Reuse the reference eval's dataset + grading + prompts + sampling config, and the # phase-0 pipeline's parsing / rollout / injection helpers (Find > Create). from eval_gpqa_rag import (GEN_GPU_MEM, GEN_MODEL_ID, build_direct_prompt, # noqa: F401 load_math) -from eval_reflexion_skill import (SKILL_GEN_USER, _EX_ATTEMPT, _EX_PROBLEM, # noqa: F401 - _EX_SKILLS, _bound_attempt, _clean_text, - _pass_rate, _parse_seq, _run_samples, - build_skill_solve_prompt) +from eval_reflexion_skill import (_EX_PROBLEM, _clean_text, # noqa: F401 + _parse_seq, _run_samples, build_skill_solve_prompt) logger = get_logger() +try: + import swanlab +except ImportError: # optional; metric logging degrades to stdout + jsonl only + swanlab = None + + # -- GPU layout --------------------------------------------------------------- TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 4)) SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 2)) NUM_GPUS = TRAIN_GPUS + SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS +# FSDP shard group size within a dp replica; TRAIN_DP is the data-parallel axis that +# ``forward_backward`` (slice_dp) splits each mini-batch over. +TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', 2)) +TRAIN_DP = max(1, TRAIN_GPUS // TRAIN_FSDP) # --------------------------------------------------------------------------- # Skill-generation prompt (DISTILL the useful approach, per the new direction) # --------------------------------------------------------------------------- +# --- Previous STRICT view-A system prompt (commented out; kept for easy revert). It +# hard-required 3-5 bullets, "output nothing after ", one-imperative-sentence +# items, no-narration, and a strict do-not-reveal block. The soft SEAM-style version +# below drops those four format demands, frames the skills as advisory reminders, and +# explains how they are used. --- +# SKILL_GEN_SYSTEM = ( +# 'You are distilling reusable problem-solving SKILLS from one worked episode. ' +# 'You are shown a competition problem, the guidance the solver was given, and the ' +# "solver's own attempt (its reasoning may be partly right and partly wrong).\n\n" +# 'FIRST, in your private thinking, do ALL of: (a) work out what this TYPE of problem ' +# 'fundamentally requires; (b) pinpoint WHERE THIS attempt actually went wrong ' +# '(when a process-check report is provided below, use its flagged criteria as ' +# 'evidence, but confirm each against the attempt yourself) — ' +# 'the decisive misstep, a missing idea, a wrong turn, or the way it stalled, looped ' +# 'on the same step, or ran the length budget out without ever committing to an ' +# 'answer; and (c) imagine AS MANY DIFFERENT angles as you can — distinct approaches ' +# 'or representations that could crack this problem, alternative solution paths, and ' +# 'the various ways a solver could plausibly go wrong on it (a few words each, do NOT ' +# 'develop them fully). THEN commit to the angle you find most decisive and write a ' +# 'SHORT list of skills that would have PREVENTED that specific ' +# 'failure and would raise the success rate of a SIMILAR solver on SIMILAR problems. ' +# 'Ground each skill in the concrete mistake you found, but state it as a GENERAL, ' +# 'transferable rule — not a patch hard-coded to this problem. Across the 3-5 ' +# 'bullets, prioritise in this order:\n' +# '1. the decisive method or representation this class of problem calls for (what to ' +# 'set up or reach for first);\n' +# '2. the specific mistake that derailed THIS attempt, recast as a general pitfall, ' +# 'plus the quick check that catches it;\n' +# '3. a compact, ordered procedure that reliably drives toward a final answer;\n' +# '4. convergence discipline: once the key quantity is in hand, commit to a single ' +# 'concrete final answer in the required format instead of re-deriving, endless ' +# 'case-splitting, looping on the same check, or overrunning the length budget.\n\n' +# 'OUTPUT FORMAT (strict):\n' +# '- Keep the thinking COMPACT — a quick brainstorm of angles then a decision, not a ' +# 'full solution and not a re-statement of these instructions. AFTER it, ' +# 'output ONLY a markdown bullet list of 3-5 items WRAPPED IN and ' +# 'tags — no preamble, no narration outside the tags. Output nothing after ' +# '.\n' +# '- Each item is ONE short imperative sentence (a method, heuristic, check, or ' +# 'habit).\n' +# '- Inside the tags: no narration, no "The student...", no headings, no restating ' +# 'the problem.\n\n' +# 'CONTENT RULES (strict):\n' +# '- Do NOT reveal the final answer or the multiple-choice option.\n' +# '- Do NOT state the specific numbers, values, or key intermediate results of THIS ' +# 'problem.\n' +# '- Every item must be GENERAL and transferable, not a step-by-step solution to ' +# 'THIS problem.\n\n' +# 'Follow the example below for the exact tags, style, and level of generality.' +# ) SKILL_GEN_SYSTEM = ( - 'You are distilling a reusable problem-solving SKILL from one worked episode. ' - 'You are shown a competition problem, the guidance the solver was given, and the ' - "solver's own attempt (its reasoning may be partly right and partly wrong).\n\n" - 'Your job: summarize the GENUINELY USEFUL parts — the effective method, the ' - 'correct reasoning direction, and the kind of guidance that transfers — into a ' - 'short list of reusable skills for SIMILAR problems. Distill the useful approach; ' - 'do NOT merely criticise this attempt.\n\n' - 'OUTPUT FORMAT (strict):\n' - '- You may think first, but the final answer MUST be a markdown bullet list of 3-5 ' - 'items WRAPPED IN and tags. Output nothing after .\n' - '- Each item is ONE short imperative sentence (a method, heuristic, or check).\n' - '- Inside the tags: no narration, no "The student...", no headings, no restating ' - 'the problem.\n\n' - 'CONTENT RULES (strict):\n' - '- Do NOT reveal the final answer or the multiple-choice option.\n' - '- Do NOT state the specific numbers, values, or key intermediate results of THIS ' - 'problem.\n' - '- Every item must be GENERAL and transferable, not a step-by-step solution to ' - 'THIS problem.\n\n' - 'Follow the example below for the exact tags, style, and level of generality.' + 'You are a mathematics coach. You are shown a competition problem together with an ' + 'automated process-check of an earlier solver attempt at it -- which solution ' + 'criteria the attempt passed or failed, and suggested fixes for the failures. You ' + 'do NOT see the attempt itself, only this check. Use the flagged failures to see ' + 'where this KIND of problem tends to trip solvers up, then distil a few reusable ' + 'tips that would prevent those specific failures.\n\n' + "These tips are advisory: they will be placed in a solver's system prompt as gentle " + 'reminders before it works through a SIMILAR problem on its own. So keep them ' + 'general and transferable — the method worth reaching for, the pitfall to watch and ' + 'a quick check, and the discipline to settle on a final answer — rather than a ' + 'worked solution to this exact problem, and without stating its specific ' + 'intermediate values or its final answer. Think briefly first, then give your tips ' + 'as a markdown bullet list wrapped in and , like the example below.' ) +# One-shot demo of the recommended mix (method / pitfall+check / procedure / +# convergence), answer-free — anchors both the format and the content priorities. +_EX_SKILLS = ( + '\n' + '- Rewrite each square root by factoring its radicand into a perfect square times ' + 'a remainder, then move the perfect-square factor outside.\n' + '- Avoid the classic trap $\\sqrt{a}+\\sqrt{b}\\ne\\sqrt{a+b}$; only combine terms ' + 'sharing the same simplest radical, and sanity-check by estimating each root.\n' + '- Procedure: simplify every radical, group like radical terms, add their ' + 'coefficients, then reduce to simplest form.\n' + '- Once the expression is in simplest form, commit to that single result as the ' + 'final answer rather than re-checking indefinitely.\n' + '') + + +# View A user template: the problem + the automated rubric process-check of an earlier +# attempt (PASS/FAIL per criterion + suggested fixes). The attempt trajectory is NOT +# shown -- the rubric findings are the evidence the skill model grounds its tips on, +# which avoids feeding the (often long, non-terminating) attempt into the prompt. +SKILL_GEN_USER_RUBRIC = ( + 'Problem:\n{problem}\n\n' + 'Process check of an earlier attempt (automated rubric verifier -- treat as ' + 'evidence, not gospel; PASS/FAIL per criterion with suggested fixes for failures):\n' + '{diagnosis}\n\n' + 'Now output the skills bullet list.' +) -def build_skillgen_prompt(problem: str, attempt: str) -> Dict[str, Any]: - """Skill-gen chat prompt: system + one-shot format demo + the real episode.""" + +def build_skillgen_prompt(problem: str, diagnosis: str) -> Dict[str, Any]: + """View A skill-gen prompt: system + one-shot format demo + the real episode + (problem + the rubric process-check of an earlier attempt). The attempt trajectory + is deliberately NOT shown -- the rubric findings localise the failure without the + generator having to re-chew (and often re-solve) a long, possibly non-terminating + attempt. The one-shot demo is query-only; only the real turn carries the diagnosis.""" return {'messages': [ {'role': 'system', 'content': SKILL_GEN_SYSTEM}, - {'role': 'user', - 'content': SKILL_GEN_USER.format(problem=_EX_PROBLEM, attempt=_EX_ATTEMPT)}, + {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=_EX_PROBLEM)}, {'role': 'assistant', 'content': _EX_SKILLS}, {'role': 'user', - 'content': SKILL_GEN_USER.format(problem=problem, attempt=attempt)}, + 'content': SKILL_GEN_USER_RUBRIC.format(problem=problem, diagnosis=diagnosis)}, ]} +# --------------------------------------------------------------------------- +# View B: query-only skill-gen (deployment form). No attempt is shown — the model +# must reason about the problem TYPE from the query alone, so the think is grounded +# on the query and cannot narrate/fabricate an attempt. Format is deliberately +# distinct from view A so the model learns the two modes as separate contracts. +# --------------------------------------------------------------------------- +# --- Previous STRICT view-B (query-only) system prompt (commented out; kept for revert). +# Same four format demands as the old view A. Soft SEAM-style version below. --- +# SKILL_GEN_SYSTEM_Q = ( +# 'You are distilling reusable problem-solving SKILLS for a CLASS of problems. You ' +# 'are shown ONE competition problem and NOTHING else — no solution, no attempt. ' +# 'FIRST, in your private thinking, imagine AS MANY DIFFERENT angles as you can — ' +# 'distinct approaches or representations that could crack this TYPE of problem, ' +# 'alternative solution paths, and the various ways a solver could plausibly go wrong ' +# 'on it (a few words each, do NOT develop them fully). THEN commit to what you find ' +# 'most decisive and write a SHORT list of skills that would raise a solver\'s success ' +# 'rate on SIMILAR problems. Across the 3-5 bullets, prioritise in this order:\n' +# '1. the decisive method or representation this class of problem calls for (what to ' +# 'set up or reach for first);\n' +# '2. the specific pitfall that derails such problems, plus the quick check that ' +# 'catches it;\n' +# '3. a compact, ordered procedure that reliably drives toward a final answer;\n' +# '4. convergence discipline: once the key quantity is in hand, commit to a single ' +# 'concrete final answer in the required format instead of re-deriving, endless ' +# 'case-splitting, or overrunning the length budget.\n\n' +# 'OUTPUT FORMAT (strict):\n' +# '- Keep the thinking COMPACT — a quick brainstorm of angles then a decision, not a ' +# 'full solution and not a re-statement of these instructions. AFTER ' +# 'it, output ONLY a markdown bullet list of 3-5 items WRAPPED IN and ' +# ' tags — no preamble, no narration outside the tags. Output nothing after ' +# '.\n' +# '- Each item is ONE short imperative sentence (a method, heuristic, check, or ' +# 'habit).\n' +# '- Inside the tags: no narration, no headings, no restating the problem, and no ' +# 'reference to any attempt, student, or solution.\n\n' +# 'CONTENT RULES (strict):\n' +# '- Do NOT solve THIS problem or reveal its final answer or multiple-choice option.\n' +# '- Do NOT state the specific numbers, values, or key intermediate results of THIS ' +# 'problem.\n' +# '- Every item must be GENERAL and transferable to other problems of the same ' +# 'type.\n\n' +# 'Follow the example below for the exact tags, style, and level of generality.' +# ) +SKILL_GEN_SYSTEM_Q = ( + 'You are a mathematics coach. You are shown ONE competition problem and nothing ' + 'else — no solution and no attempt. Think about what approach this KIND of problem ' + 'calls for and where solvers tend to slip, then distil a few reusable tips.\n\n' + "These tips are advisory: they will be placed in a solver's system prompt as gentle " + 'reminders before it works through a SIMILAR problem on its own. So keep them ' + 'general and transferable — the method worth reaching for, the pitfall to watch and ' + 'a quick check, and the discipline to settle on a final answer — rather than a ' + 'worked solution to this exact problem, and without stating its specific ' + 'intermediate values or its final answer. Think briefly first, then give your tips ' + 'as a markdown bullet list wrapped in and , like the example below.' +) + +SKILL_GEN_USER_Q = ( + 'Problem:\n{problem}\n\n' + 'Now reason about this TYPE of problem, then output the skills bullet list.' +) + + +def build_querygen_prompt(problem: str) -> Dict[str, Any]: + """View B skill-gen prompt: system + one-shot demo + the problem ALONE (no + attempt) — matching what is available at deployment (query only).""" + return {'messages': [ + {'role': 'system', 'content': SKILL_GEN_SYSTEM_Q}, + {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=_EX_PROBLEM)}, + {'role': 'assistant', 'content': _EX_SKILLS}, + {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=problem)}, + ]} + + +def _assign_view(problem: str, args: argparse.Namespace) -> str: + """Deterministically route a problem to exactly one view (stable across restarts + and across the generation/SFT sides). ``--view-b-frac`` of problems go to view B.""" + h = int(hashlib.md5(f'{args.seed}:{problem}'.encode('utf-8')).hexdigest(), 16) + return 'B' if (h % 100000) / 100000.0 < args.view_b_frac else 'A' + + +def _skillgen_messages(problem: str, view: str, diagnosis: str) -> List[Dict[str, Any]]: + """Single source of truth for the skill-gen prompt, used at BOTH generation and + training time so they can never diverge. View A with a localisable failure uses + problem + rubric findings (NO trajectory); view B -- or a view-A problem whose rubric + flagged NO failure (``[FAIL]`` absent: all-pass or missing diagnosis) -- is query-only. + So view A DEGRADES to view B whenever there is nothing concrete to correct.""" + if view == 'B' or '[FAIL]' not in (diagnosis or ''): + return build_querygen_prompt(problem)['messages'] + return build_skillgen_prompt(problem, diagnosis)['messages'] + + +def _view_prompt(r: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: + """The skill-gen prompt for problem ``r`` under its assigned view (routing in + ``_skillgen_messages``: view A carries the rubric process-check; view B, and any + view-A problem with no rubric failure, is query-only).""" + return {'messages': _skillgen_messages(r['problem'], r['_view'], r.get('_rubric_diag', ''))} + + +_BULLET_RE = re.compile(r'(?m)^\s*(?:[-*]|\d+[.)])\s') +# Trajectory/meta references that betray CoT fragments leaking into the block; any +# hit fails the purity gate (the problem is then re-sampled, per --skill-retries). +_META_RE = re.compile( + r'\b(the student|the solver|the attempt|this attempt|the trace|the response|' + r'in the (?:attempt|trace|response|solution)|as (?:shown|seen|noted) above|' + r'the (?:above|previous|earlier)|my (?:reasoning|analysis)|i (?:think|need|will))\b', + re.IGNORECASE) + + +def _is_clean_block(block: str) -> bool: + """Purity gate for thinking-ON skill-gen: the block must be a pure bullet list + (every non-empty line a bullet — no prose/CoT fragments) with no meta/trajectory + reference. Answer leak is caught separately by the backup-teacher leak stage.""" + lines = [ln for ln in (l.strip() for l in block.splitlines()) if ln] + if not lines or not all(_BULLET_RE.match(ln) for ln in lines): + return False + return _META_RE.search(block) is None + + def _extract_skills_block(text: str) -> Optional[str]: - """Return the inner ``...`` block, or None if not parseable. + """Return the clean ``...`` block, or None if not parseable. - RFT keeps thinking ON, so a candidate is usable only if it actually produced a - CLOSED tag block; unterminated / tag-less generations are dropped (never fall - back to scraping raw reasoning) — see reflexion.md §6.8. - """ + Skill-gen runs with thinking ON, so the model must end its reasoning with an explicit + ```` before committing an answer (whether the opening ```` is emitted by + the model or pre-injected by the chat template). We therefore REQUIRE ```` and + read only the text after the last one; its absence means the token budget was exhausted + mid-reasoning (nothing committed, per reflexion.md §6.8) — reject so a draft or a + system-prompt demo echo inside the CoT can never be mistaken for the answer. Within + the answer take the ```` block (closing tag optional), strip stray tags, and + require ``_is_clean_block`` — prose-mixed / meta-referencing fragments are rejected + for re-sampling.""" low = text.lower() - if '' not in low or '' not in low: + end_think = low.rfind('') + if end_think < 0: + return None # reasoning never closed -> no committed answer + answer = text[end_think + len(''):] + low_a = answer.lower() + s = low_a.find('') + if s < 0: return None - start = low.index('') + len('') - end = low.index('') - if end <= start: + inner = s + len('') + e = low_a.find('', inner) + block = (answer[inner:e] if e >= 0 else answer[inner:]).strip() + block = re.sub(r'', '', block, flags=re.IGNORECASE).strip() + if not _is_clean_block(block): return None - block = text[start:end].strip() - return block or None + return block + +# --------------------------------------------------------------------------- +# OPTIONAL stricter leak criterion (currently UNUSED -- the run uses answer_only=True, +# which flags ONLY the final answer). This variant ALSO flags concrete intermediate KEY +# results, while still permitting method / plan / pitfalls / checks. To enable, pass +# judge_system=_LEAK_JUDGE_SYSTEM to the LeakVerifier below. +# --------------------------------------------------------------------------- +_LEAK_JUDGE_SYSTEM = """\ +You check whether a HINT that will be shown to someone solving a math TASK gives away +this task's own results. -def _bounded_attempt(r: Dict[str, Any], args: argparse.Namespace) -> str: - """Bound the failed attempt so problem + attempt + skill output fits context. +The hint may FREELY describe the general method, which approach or technique to use, the +steps or plan to follow, common pitfalls, and sanity checks -- even when that points +strongly at HOW to solve THIS task. Describing the approach is expected of a good hint. - Capped by ``--attempt-max-tokens`` so the SFT sequences (which embed the same - attempt) stay short enough to train; the same bound is used at generation and - SFT time, keeping train/inference identical. - """ - prob_est = len(r['problem']) // 2 - budget = max(1024, args.max_model_len - args.skill_max_tokens - - args.attempt_reserve_tokens - prob_est) - budget = min(budget, args.attempt_max_tokens) - init = r['_init'][0] - return _bound_attempt(init['text'], init['gen_tokens'], budget) +The hint LEAKS only if, for THIS specific task, it states either: +- the final answer or final result (a value, expression, choice, label, or verbatim + output); or +- a concrete decisive INTERMEDIATE key result -- a specific computed value, quantity, or + fact unique to this task that hands over a key step of the answer. + +If it names only the method / plan / pitfalls / checks WITHOUT stating those concrete +intermediate values or the final result, it does NOT leak. + +Reply with exactly one word: LEAK or CLEAN.""" # --------------------------------------------------------------------------- -# Online data generation (one chunk, all sampler calls batched) +# Rubric process-check (view A only): a frozen teacher diagnoses the base's failed +# attempt so the skill model grounds its error analysis on a verified fault +# localisation instead of guessing. Teacher-only (sampler=None -> every diagnose() +# hits llm_backup); mirrors eval_dualline_math's fixed math rubric. # --------------------------------------------------------------------------- -def generate_chunk(base_sampler, skill_sampler, leak: LeakVerifier, - chunk: List[Dict[str, Any]], base_dp: int, skill_dp: int, - args: argparse.Namespace) -> List[Dict[str, Any]]: - """Run base-solve -> hard-filter -> skill-gen -> leak-filter -> marginal, and - return the records whose skill gave a positive marginal pass@k gain.""" - # --- Phase 1: base solves once; keep only what it got wrong. --- - init_out = _run_samples( - base_sampler, [build_direct_prompt(r['problem']) for r in chunk], - args.init_samples, args.max_tokens, base_dp) - failed = [] - for r, seqs in zip(chunk, init_out): - r['_init'] = [_parse_seq(s, r['reference_answer']) for s in seqs] - if _pass_rate(r['_init']) == 0.0: - failed.append(r) - sys.stderr.write(f' phase1: {len(chunk)-len(failed)}/{len(chunk)} solved on ' - f'first try, {len(failed)} failed\n') - if not failed: - return [] - - # --- Phase 2: baseline pass@k on the failures -> keep the genuinely hard. --- - base_out = _run_samples( - base_sampler, [build_direct_prompt(r['problem']) for r in failed], - args.pass_k, args.max_tokens, base_dp) - hard = [] - for r, seqs in zip(failed, base_out): - r['_baseline_pass'] = _pass_rate([_parse_seq(s, r['reference_answer']) for s in seqs]) - if r['_baseline_pass'] <= args.hard_baseline_max: - hard.append(r) - sys.stderr.write(f' phase2: {len(hard)}/{len(failed)} failures are hard ' - f'(pass@{args.pass_k} <= {args.hard_baseline_max})\n') - if not hard: - return [] - - # --- Phase 3: skill-gen (thinking ON) -> require a parseable block. --- - sg_out = _run_samples( - skill_sampler, - [build_skillgen_prompt(r['problem'], _bounded_attempt(r, args)) for r in hard], - args.n_skills, args.skill_max_tokens, skill_dp) - cands: List[Dict[str, Any]] = [] # {r, response, block} - for r, seqs in zip(hard, sg_out): - for s in seqs: - resp = _clean_text(getattr(s, 'decoded', '') or '') - block = _extract_skills_block(resp) - if block: - cands.append({'r': r, 'response': resp, 'block': block}) - if not cands: - sys.stderr.write(' phase3: 0 parseable skill blocks\n') - return [] - - # --- Phase 4: leak filter via backup teacher (drops answer-leaking skills). --- - details = leak.leak_batch( - [{'content': c['block'], 'query': c['r']['problem'], - 'reference': c['r']['reference_answer']} for c in cands], - max_workers=args.leak_workers) - clean = [c for c, d in zip(cands, details) if not d.leaked] - sys.stderr.write(f' phase4: {len(clean)}/{len(cands)} skills clean ' - f'({len(cands)-len(clean)} leaked)\n') - if not clean: - return [] - - # --- Phase 5: with-skill pass@k -> marginal = with - baseline. --- - ws_out = _run_samples( - base_sampler, - [build_skill_solve_prompt(c['r']['problem'], c['block']) for c in clean], - args.pass_k, args.max_tokens, base_dp) - records: List[Dict[str, Any]] = [] - for c, seqs in zip(clean, ws_out): - r = c['r'] - with_pass = _pass_rate([_parse_seq(s, r['reference_answer']) for s in seqs]) - marginal = with_pass - r['_baseline_pass'] - if marginal > 0.0: # keep only skills that actually helped - records.append({ - 'problem': r['problem'], - 'reference_answer': r['reference_answer'], - 'attempt': _bounded_attempt(r, args), - 'response': c['response'], - 'skills': c['block'], - 'baseline_pass': r['_baseline_pass'], - 'with_pass': with_pass, - 'marginal': marginal, - }) - sys.stderr.write(f' phase5: {len(records)} skills with positive marginal kept\n') - return records +_MATH_RUBRIC = [ + ('The reasoning contains no arithmetic or algebraic error', True), + ('Each step follows logically from the previous ones', True), + ('No formula or theorem is misstated or misapplied', True), + ('The approach is on track to answer the actual question asked', False), + ('No step contradicts an earlier established fact', False), +] + + +def _build_rubric_checker() -> Optional['RubricVerifier']: + """Fixed math-process rubric verifier, teacher-served. None if no LLM backup env.""" + if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') + or os.environ.get('OPENAI_API_KEY')): + return None + return RubricVerifier( + fixed_rubric=[RubricItem(t, is_hard=h) for t, h in _MATH_RUBRIC], gate=True) + + +def _format_diagnosis(detail) -> str: + """One line per criterion (PASS/FAIL + reason + fix on FAIL) then a summary — the + compact evidence block appended to the view-A skill-gen prompt.""" + rub = detail.rubric + lines = [] + for it in detail.items: + text = rub[it.index - 1].text if 0 < it.index <= len(rub) else f'criterion {it.index}' + if it.verdict: + lines.append(f'- [PASS] {text}') + else: + tail = f': {it.reason}' if it.reason else '' + tail += f' (fix: {it.fix})' if it.fix else '' + lines.append(f'- [FAIL] {text}{tail}') + if detail.summary: + lines.append(f'Summary: {detail.summary}') + return '\n'.join(lines) + + +def _diagnose_views(checker, hard: List[Dict[str, Any]], args: argparse.Namespace) -> None: + """Rubric-check every view-A problem's greedy attempt in parallel, stashing the + formatted findings on ``r['_rubric_diag']`` (view B stays empty). A checker error + or empty result degrades to no diagnosis (the plain view-A prompt).""" + from concurrent.futures import ThreadPoolExecutor + targets = [r for r in hard if r.get('_view') == 'A'] + if not checker or not targets: + return + + def _run(r: Dict[str, Any]) -> Tuple[Dict[str, Any], str]: + seg = {'messages': [ + {'role': 'user', 'content': r['problem']}, + {'role': 'assistant', 'content': r['_init'][0]['text']}, + ]} + try: + return r, _format_diagnosis(checker.diagnose(seg, query=r['problem'])) + except Exception as exc: # teacher hiccup -> fall back to no-diagnosis prompt + logger.warning(f'[rubric] diagnose error: {exc}') + return r, '' + + workers = max(1, min(args.rubric_workers, len(targets))) + with ThreadPoolExecutor(max_workers=workers) as ex: + for r, diag in ex.map(_run, targets): + r['_rubric_diag'] = diag # --------------------------------------------------------------------------- -# Online RFT training +# Online data generation (one chunk; every candidate is recorded, untruncated) # --------------------------------------------------------------------------- -def _sft_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: - """SFT sample = the exact skill-gen prompt + the kept generation as the target. +def _roll(x: Dict[str, Any]) -> Dict[str, Any]: + """Full (untruncated) rollout record for offline analysis.""" + return {k: x[k] for k in ('pred', 'correct', 'terminated', 'passed', + 'stop_reason', 'gen_tokens', 'text')} - Matching the generation prompt keeps train/inference consistent; the template - masks the prompt to -100 and supervises the assistant turns. + +def _empty_roll() -> Dict[str, Any]: + """Fallback rollout when the sampler returned nothing for a prompt.""" + return {'pred': '', 'correct': False, 'terminated': False, 'passed': False, + 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} + + +def _assign_advantages(hard: List[Dict[str, Any]], args: argparse.Namespace) -> None: + """Group-relative advantage (SEAM-style) over each problem's clean, scored candidates, + using the DETERMINISTIC greedy reward ``R in {0, 1}`` (answer CORRECT only; + termination is NOT part of the reward -- monitored via `terminated`/`passed` only): + + A_j = (R_j - mean_R) / (std_R + eps) + + Groups where every candidate shares the same reward (``std_R == 0``: all solve or all + fail) get advantage 0 and contribute no gradient -- GRPO's own group variance + auto-selects the informative problems, so no explicit difficulty / marginal gate is + needed. Because the reward is deterministic (M=1 greedy, no pass@k sampling), the + std-normalisation no longer amplifies rollout noise (the reason it was dropped for the + old stochastic marginal). ``kept`` marks above-average candidates (for reporting only). + """ + eps = 1e-6 + for r in hard: + for c in r['_cands']: + c['advantage'], c['grpo_adv'], c['kept'] = 0.0, 0.0, False + cs = ([c for c in r['_cands'] if c.get('reward') is not None] if args.format_in_reward + else [c for c in r['_cands'] if c['leaked'] is False and c.get('reward') is not None]) + if len(cs) < 2: + continue + rewards = [c['reward'] for c in cs] + mean_r = sum(rewards) / len(rewards) + std = (sum((x - mean_r) ** 2 for x in rewards) / len(rewards)) ** 0.5 + if std < 1e-9: + continue # all candidates equal (all solve / all fail) -> no learning signal + for c in cs: + adv = (c['reward'] - mean_r) / (std + eps) + c['advantage'], c['grpo_adv'], c['kept'] = adv, adv, c['reward'] > mean_r + + +def _full_record(r: Dict[str, Any], ci: int) -> Dict[str, Any]: + """Complete per-problem record: init attempt, baseline, and ALL candidates + (parseable/leaked/scored alike) with full text — nothing dropped or truncated.""" + init = r['_init'][0] + rec: Dict[str, Any] = { + 'record_type': 'problem', 'chunk': ci, 'problem': r['problem'], + 'reference_answer': r['reference_answer'], 'level': r.get('level', ''), + 'failed_first_try': r['_failed'], + 'init_attempt': {'text': init['text'], 'pred': init['pred'], + 'correct': init['correct'], 'terminated': init['terminated'], + 'stop_reason': init['stop_reason'], 'gen_tokens': init['gen_tokens']}, + } + rec['baseline_pass'] = r['_baseline_pass'] + rec['is_hard'] = r['_hard'] + rec['view'] = r.get('_view', '') + rec['rubric_diag'] = r.get('_rubric_diag', '') + rec['baseline_rolls'] = [_roll(x) for x in r['_baseline_rolls']] + rec['candidates'] = [{ + 'skills': c['skills'], 'response': c['response'], 'parseable': c['parseable'], + 'leaked': c['leaked'], 'leak_reason': c['leak_reason'], 'leak_source': c['leak_source'], + 'with_pass': c['with_pass'], 'reward': c.get('reward'), + 'advantage': c.get('advantage'), 'grpo_adv': c.get('grpo_adv'), 'kept': c.get('kept'), + 'rolls': [_roll(x) for x in c['rolls']], + } for c in r['_cands']] + return rec + + +def _view_stats(hard: List[Dict[str, Any]], view: str) -> Dict[str, Any]: + """Per-view yield: hard problems, clean candidates, and the ADOPTION rate — + the fraction of hard problems that produced at least one clean, non-zero-advantage + candidate (i.e. a record that actually reaches training). Watching A vs B and + early vs late tells whether query-only (B) catches up to trajectory-grounded (A).""" + hv = [r for r in hard if r.get('_view') == view] + cands = [c for r in hv for c in r['_cands'] if c['parseable']] + clean = [c for c in cands if c['leaked'] is False] + adopted = sum(1 for r in hv + if any(c['leaked'] is False and abs(c.get('advantage') or 0.0) > 1e-9 + for c in r['_cands'])) + return { + 'n_hard': len(hv), 'n_candidates_parseable': len(cands), + 'n_clean': len(clean), 'n_adopted_problems': adopted, + 'adoption_rate': (adopted / len(hv)) if hv else 0.0, + } + + +def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespace) -> Dict[str, Any]: + """Per-chunk aggregates — watch these across chunks to see if the RFT'd skill + model produces better skills over time (yield, leak rate, lift, termination).""" + failed = [r for r in chunk if r['_failed']] + hard = [r for r in chunk if r['_hard']] + all_cands = [c for r in chunk for c in r['_cands']] + cands = [c for c in all_cands if c['parseable']] + scored = [c for c in cands if c['with_pass'] is not None] + ws_rolls = [x for c in scored for x in c['rolls']] + # With --format-in-reward, unparseable/leaked candidates also carry a (0) reward and are + # trained, so count trainables over ALL candidates; else only clean scored ones. + train_cands = ([c for c in all_cands if abs(c.get('advantage') or 0.0) > 1e-9] if args.format_in_reward + else [c for c in scored if c['leaked'] is False and abs(c.get('advantage') or 0.0) > 1e-9]) + base_acc = (sum(r['_baseline_pass'] for r in hard) / len(hard)) if hard else 0.0 + ws_acc = (sum(c['with_pass'] for c in scored) / len(scored)) if scored else 0.0 + return { + 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), + 'n_failed_first_try': len(failed), 'n_hard': len(hard), + 'n_generated': len(all_cands), + 'n_candidates_parseable': len(cands), + 'n_unparseable': len(all_cands) - len(cands), + 'n_leaked': sum(1 for c in cands if c['leaked']), + 'n_clean': sum(1 for c in cands if c['leaked'] is False), + 'n_reward_pos': sum(1 for c in scored if c['reward']), + 'n_train_samples': len(train_cands), + 'avg_baseline_pass_on_hard': base_acc, + 'avg_withskill_pass': ws_acc, + 'avg_lift': ws_acc - base_acc, + 'termination_rate_withskill': (sum(1 for x in ws_rolls if x['terminated']) / len(ws_rolls)) if ws_rolls else 0.0, + 'view_A': _view_stats(hard, 'A'), 'view_B': _view_stats(hard, 'B'), + } + + +def _group_records(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> List[Dict[str, Any]]: + """GRPO training records: every clean, scored skill candidate with a NON-ZERO + advantage (positive pushes the skill up, negative down; the group-relative + usefulness-over-base advantage was set in _assign_advantages). Each carries its + ``view`` and the rubric ``diagnosis``; the prompt (identical to generation) is rebuilt + from those by ``_skillgen_messages`` -- no trajectory is stored or replayed.""" + out = [] + for r in chunk: + if not r['_hard']: + continue + view = r.get('_view', 'A') + for c in r['_cands']: + adv_nz = abs(c.get('advantage') or 0.0) > 1e-9 + trainable = adv_nz if args.format_in_reward else ( + c['leaked'] is False and c['with_pass'] is not None and adv_nz) + if trainable: + out.append({ + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'view': view, + 'diagnosis': r.get('_rubric_diag', ''), + 'response': c['response'], 'skills': c['skills'], + 'advantage': c['advantage'], 'grpo_adv': c['grpo_adv'], 'kept': c['kept'], + 'reward': c['reward'], 'with_pass': c['with_pass'], + }) + return out + + +def process_chunk(base_sampler, skill_sampler, leak: LeakVerifier, + chunk: List[Dict[str, Any]], ci: int, base_dp: int, skill_dp: int, + args: argparse.Namespace, checker=None + ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], List[Dict[str, Any]]]: + """base-solve -> rubric-check (view A) -> skill-gen -> leak-filter -> with-skill pass + -> GRPO advantages, for one chunk. + + Sequential (generate-one-chunk-train-one): generation and the trainer's weight + sync never overlap, so no lock is needed. ``base_sampler`` is frozen (never + synced); ``skill_sampler`` is synced by the trainer between chunks. """ - msgs = build_skillgen_prompt(rec['problem'], rec['attempt'])['messages'] - return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}]} + # --- Phase 1: base solves each problem GREEDILY once (T=0, M=1) -- this produces the + # view-A attempt and records whether the base already gets it (reporting only). EVERY + # problem is processed (no difficulty gate, SEAM-style): the group-relative advantage + # (Phase 6) gives zero gradient to any problem whose skills all score alike (base-easy + # -> all solve, or hopeless -> all fail), so GRPO's own group variance selects the + # informative problems. Termination is NOT part of the reward (monitored only). --- + base_out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in chunk], + 1, args.max_tokens, base_dp, temperature=0.0) + for r, seqs in zip(chunk, base_out): + roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + r['_baseline_rolls'], r['_cands'] = [roll], [] + r['_init'] = [roll] # the greedy attempt (metric + view A) + r['_failed'] = not roll['correct'] + r['_baseline_pass'] = 1.0 if roll['correct'] else 0.0 # base accuracy (reporting only, NOT in reward) + r['_hard'] = True # process EVERY problem; group variance selects + hard = chunk + + # --- Phase 2: assign each problem's view, then rubric-check the view-A attempts so + # the skill model diagnoses from verified findings instead of guessing. View B is + # query-only and deliberately gets NO rubric (nothing to diagnose without an attempt). --- + for r in hard: + r['_view'] = _assign_view(r['problem'], args) + r['_rubric_diag'] = '' + _diagnose_views(checker, hard, args) + + # --- Phase 3: skill-gen (thinking ON), per-view prompt; re-sample empties. --- + flat: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] + if hard: + pending = list(hard) # problems still without any clean candidate + for _ in range(args.skill_retries + 1): + if not pending: + break + prompts = [_view_prompt(r, args) for r in pending] + sg_out = _run_samples(skill_sampler, prompts, args.n_skills, + args.skill_max_tokens, skill_dp, + temperature=args.skill_gen_temperature, + top_p=args.skill_gen_top_p, + top_k=args.skill_gen_top_k) + still = [] + for r, seqs in zip(pending, sg_out): + got = False + for s in seqs: + resp = _clean_text(getattr(s, 'decoded', '') or '') + block = _extract_skills_block(resp) + cand = {'skills': block or '', 'response': resp, 'parseable': bool(block), + 'view': r['_view'], 'leaked': None, 'leak_reason': '', + 'leak_source': '', 'with_pass': None, + 'reward': None, 'rolls': []} + r['_cands'].append(cand) + if block: + flat.append((r, cand)) + got = True + if not got: + still.append(r) # nothing parseable yet -> retry this problem + pending = still + + # --- Phase 4: leak filter via backup teacher (network only, no lock). VIEW A ONLY -- + # view B is query-only (no trajectory to leak from) and is left exactly like SEAM, which + # runs NO leak filter: its candidates skip the check and are treated as clean. To restore + # leak-checking on view B, drop the ``_view == 'A'`` guard below. --- + for r, c in flat: + if r.get('_view') != 'A': + c['leaked'], c['leak_reason'], c['leak_source'] = False, '', 'skipped_viewB' + flat_a = [(r, c) for r, c in flat if r.get('_view') == 'A'] + if flat_a: + details = leak.leak_batch( + [{'content': c['skills'], 'query': r['problem'], 'reference': r['reference_answer']} + for r, c in flat_a], max_workers=args.leak_workers) + for (r, c), d in zip(flat_a, details): + c['leaked'], c['leak_reason'], c['leak_source'] = bool(d.leaked), d.reason, d.source + + # --- Phase 5: with-skill GREEDY pass (T=0, M=1) on clean candidates. Binary reward + # R = answer CORRECT (deterministic, no pass@k noise), ABSOLUTE -- no baseline + # subtraction; the group mean in Phase 6 is the only baseline. Termination is NOT + # required (monitored only) -- see reflexion.md §7.6. --- + clean = [(r, c) for r, c in flat if c['leaked'] is False] + if clean: + ws_out = _run_samples(base_sampler, + [build_skill_solve_prompt(r['problem'], c['skills']) for r, c in clean], + 1, args.max_tokens, base_dp, temperature=0.0) + for (r, c), seqs in zip(clean, ws_out): + c['rolls'] = [_parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll()] + c['with_pass'] = 1.0 if c['rolls'][0]['correct'] else 0.0 + c['reward'] = c['with_pass'] # valid + clean + correct -> 1 + # Validity-in-reward (SEAM-style, --format-in-reward): every candidate that never reached + # the executor -- unparseable/impure format OR answer-leaked -- scores 0 and STILL joins its + # group, so its whole response (think tokens included) is trained DOWN. Off => those + # candidates are excluded, as before. + if args.format_in_reward: + for r in hard: + for c in r['_cands']: + if c['reward'] is None: + c['reward'] = 0.0 + + # --- Phase 6: group-relative GRPO advantage per problem-group. --- + _assign_advantages(hard, args) + return ([_full_record(r, ci) for r in chunk], + _chunk_summary(chunk, ci, args), + _group_records(chunk, args)) -def _train_round(skill_model, ckpt: CheckpointEngineManager, - pool: List[Dict[str, Any]], args: argparse.Namespace) -> None: - """One RFT round: weighted-sample ``train_batch`` records, SFT them in driver- - side mini-batches (each an optimizer step, matching short_math_grpo), then sync. - The transformers backend runs one forward per ``forward_backward`` call (no - internal micro split), and ``slice_dp`` needs each mini-batch divisible across - the training dp ranks — hence ``sft_batch_size`` is a multiple of TRAIN_GPUS. +# --------------------------------------------------------------------------- +# Online RFT training +# --------------------------------------------------------------------------- +def _train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: + """Training sample = the exact skill-gen prompt for this record's view + the + generated (think + skills) response as the target; the GRPO advantage is attached + separately at forward_backward time. + + The prompt is rebuilt by ``_skillgen_messages`` (the same function used at generation), + so train/inference stay identical: view A replays problem + rubric findings, view B (and + no-failure view A) replays the query-only prompt. ``key_rounds`` selects the final + assistant turn (index ``len(msgs)``); the plain ``Template`` then masks the prompt and + trains the whole generated response (reasoning + ```` + skills) -- the key-round + prefix already excludes the prompt-provided ````, so no extra masking is needed.""" + msgs = _skillgen_messages(rec['problem'], rec.get('view', 'A'), rec.get('diagnosis', '')) + full = msgs + [{'role': 'assistant', 'content': rec['response']}] + return {'messages': full, 'user_data': {'key_rounds': [len(msgs)]}} + + +def _train_chunk(skill_model, ckpt: CheckpointEngineManager, + samples: List[Dict[str, Any]], args: argparse.Namespace) -> Dict[str, Any]: + """One on-policy GRPO update on THIS chunk's skill candidates, then sync weights. + + Sequential design (generate-one-chunk-train-one): the skills were sampled from the + current policy and trained immediately, so ``old_logps`` is omitted and the GRPO + ratio is ~1 (no importance correction needed). Each per-sample ``advantage`` was set + in _assign_advantages (group-relative + optional SFT blend). Driver-side mini-batches + each take an optimizer step (matching short_math_grpo); the batch is padded to a + multiple of ``sft_batch_size`` (dp needs each mini-batch divisible) with advantage-0 + copies that contribute zero gradient. ``sync_weights`` needs no lock (no overlap). """ - weights = [max(rec['marginal'], args.weight_eps) ** args.weight_alpha for rec in pool] - batch = random.choices(pool, weights=weights, k=args.train_batch) - trajs = [_sft_trajectory(rec) for rec in batch] + trajs = [_train_trajectory(rec) for rec in samples] + advs = [float(rec['advantage']) for rec in samples] + rem = (-len(trajs)) % args.sft_batch_size + if rem: + trajs += [trajs[-1]] * rem # zero-advantage pads -> forward only, no gradient + advs += [0.0] * rem + steps = 0 for i in range(0, len(trajs), args.sft_batch_size): - skill_model.forward_backward(inputs=trajs[i:i + args.sft_batch_size]) + skill_model.forward_backward(inputs=trajs[i:i + args.sft_batch_size], + advantages=advs[i:i + args.sft_batch_size]) skill_model.clip_grad_and_step() - ckpt.sync_weights(merge_and_sync=True) # full-param: push all weights to vLLM + steps += 1 + ckpt.sync_weights(merge_and_sync=True) + metric = skill_model.calculate_metric(is_training=True) + return {'n_samples': len(samples), 'n_steps': steps, + 'advantages': [float(rec['advantage']) for rec in samples], + 'metric': {k: (float(v) if _is_num(v) else v) for k, v in (metric or {}).items()}} + + +def _is_num(v: Any) -> bool: + try: + float(v) + return True + except (TypeError, ValueError): + return False # --------------------------------------------------------------------------- @@ -264,58 +766,168 @@ def _build_args() -> argparse.Namespace: p.add_argument('--seed', type=int, default=42) p.add_argument('--chunk-size', type=int, default=16, help='Problems per generation chunk (all sampler calls batched).') - p.add_argument('--init-samples', type=int, default=1) p.add_argument('--n-skills', type=int, default=8, help='Candidate skills generated per hard problem.') - p.add_argument('--pass-k', type=int, default=8) - p.add_argument('--hard-baseline-max', type=float, default=0.25) + p.add_argument('--view-b-frac', type=float, default=0.5, + help='Fraction of hard problems routed to view B (query-only, ' + 'deployment form); the rest go to view A (problem + attempt). ' + 'Each problem is assigned to EXACTLY ONE view.') + p.add_argument('--skill-retries', type=int, default=2, + help='Extra skill-gen rounds for a hard problem that yielded no ' + 'clean, parseable candidate (thinking-ON purity gate rejects).') + p.add_argument('--skill-gen-temperature', type=float, default=1.0, + help='Sampling temperature for skill-gen (BOTH views). >0 so the ' + 'n_skills candidates per problem are genuinely DIVERSE — a group ' + 'of near-duplicate skills gives GRPO no real good-vs-bad contrast.') + p.add_argument('--skill-gen-top-p', type=float, default=1.0, + help='top_p for skill-gen; 1.0 keeps the full tail for diversity.') + p.add_argument('--skill-gen-top-k', type=int, default=-1, + help='top_k for skill-gen; -1 disables truncation (max diversity). ' + 'A finite value only narrows the candidate pool.') p.add_argument('--max-model-len', type=int, default=16384) p.add_argument('--max-tokens', type=int, default=8192, help='Max generated tokens for solve rollouts.') - p.add_argument('--skill-max-tokens', type=int, default=4096, - help='Max tokens for skill-gen (thinking ON: room for CoT + the ' - ' block).') - p.add_argument('--attempt-reserve-tokens', type=int, default=1024) - p.add_argument('--attempt-max-tokens', type=int, default=3072, - help='Hard cap on the failed-attempt tokens embedded in skill-gen ' - 'and SFT prompts (keeps SFT sequences trainable).') - p.add_argument('--leak-workers', type=int, default=32, - help='Parallel workers for the LeakVerifier backup judge.') - # -- online RFT -- - p.add_argument('--train-every', type=int, default=64, - help='Trigger one train round after this many new valid records.') - p.add_argument('--train-batch', type=int, default=64, - help='Records weighted-sampled (with replacement) per train round.') + p.add_argument('--skill-max-tokens', type=int, default=8192, + help='Max tokens for skill-gen (thinking ON: the model must close ' + ' within this budget or the candidate is dropped, so ' + 'leave ample room).') + p.add_argument('--leak-workers', type=int, default=16, + help='Parallel workers for the LeakVerifier backup judge (capped at 16 ' + 'to avoid the teacher API burst-rate limit; leak and rubric run in ' + 'separate phases so peak teacher concurrency is max(leak,rubric)).') + p.add_argument('--rubric-workers', type=int, default=16, + help='Parallel workers for the view-A rubric diagnose() calls ' + '(teacher-served; requires LLM_BACKUP_* env).') + # -- online GRPO (one on-policy update per generated chunk) -- p.add_argument('--sft-batch-size', type=int, default=8, help='Driver-side mini-batch per optimizer step; MUST be a multiple ' - 'of TRAIN_GPUS (sliced across training dp ranks) and divide ' - '--train-batch.') + 'of the training dp size (sliced across dp ranks).') + p.add_argument('--grpo-epsilon', type=float, default=0.2, + help='PPO clip epsilon for GRPOLoss (ratio~1 on-policy, so rarely binds).') + p.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True, + help='Fold output validity into the reward (SEAM-style): unparseable/impure ' + 'or answer-leaked candidates score 0 and join their group to be trained ' + 'DOWN (the whole response, think tokens included). ' + '--no-format-in-reward keeps the reject-and-exclude gate.') p.add_argument('--lr', type=float, default=1e-5) p.add_argument('--max-train-rounds', type=int, default=200, - help='Cap on train rounds (also sizes the LR schedule).') + help='Cap on train rounds = trained chunks (also sizes the LR schedule).') p.add_argument('--save-rounds', type=int, default=50) - p.add_argument('--weight-alpha', type=float, default=1.0, - help='Sampling weight exponent: w = max(marginal, eps) ** alpha.') - p.add_argument('--weight-eps', type=float, default=0.01) + p.add_argument('--trend-every', type=int, default=10, + help='Every N chunks, print a [trend] line contrasting the first N ' + 'vs the most recent N chunks (adoption + lift + pos/chunk) ' + 'so the training effect on fresh problems is visible at a glance.') p.add_argument('--output-dir', default='./output/reflexion_skill_rft') + p.add_argument('--swanlab-project', default='twinkle', + help='swanlab project; logging is skipped when swanlab is not ' + 'installed or SWANLAB_MODE=disabled.') + p.add_argument('--swanlab-exp', default='', + help='swanlab experiment (run) name; empty = auto.') return p.parse_args() +def _trend_line(hist: List[Dict[str, float]], window: int, rounds_done: int) -> Optional[str]: + """Contrast the FIRST ``window`` chunks with the most recent ``window`` chunks so + the online training effect on fresh, never-trained problems is glanceable: if RFT + is working, adoption and lift on recent chunks exceed the early baseline.""" + if len(hist) < 2 * window: + return None # need two non-overlapping windows for a clean before/after + base, rec = hist[:window], hist[-window:] + m = lambda xs, k: sum(h[k] for h in xs) / len(xs) + return (f'[trend] first {window} vs last {window} chunks | ' + f'adopt A {m(base,"aA"):.2f}->{m(rec,"aA"):.2f} ' + f'B {m(base,"aB"):.2f}->{m(rec,"aB"):.2f} | ' + f'lift {m(base,"lift"):+.3f}->{m(rec,"lift"):+.3f} | ' + f'pos/chunk {m(base,"pos"):.1f}->{m(rec,"pos"):.1f} | rounds={rounds_done}') + + +def _query_rows(full: List[Dict[str, Any]]) -> List[Tuple[float, float, float, int, str]]: + """Per hard problem that produced >=1 scored candidate: its no-skill baseline + pass@k, the BEST and MEAN with-skill pass@k over its N skill candidates, the scored + count, and the problem text. Drives both the per-query print and the swanlab passk/* + aggregates.""" + rows = [] + for rec in full: + if rec.get('record_type') != 'problem' or not rec.get('is_hard'): + continue + ps = [c['with_pass'] for c in rec.get('candidates', []) if c.get('with_pass') is not None] + if not ps: + continue + rows.append((rec['baseline_pass'], max(ps), sum(ps) / len(ps), len(ps), rec['problem'])) + return rows + + +def _clean_metric(metric: Optional[Dict[str, Any]]) -> Dict[str, float]: + """Numeric GRPO metrics for swanlab: collapse the duplicate per-group LR to a single + ``lr`` and drop non-numeric fields (e.g. 'total time elapse').""" + out: Dict[str, float] = {} + for k, v in (metric or {}).items(): + if not _is_num(v): + continue + if k.startswith('learning rate'): + if 'group 1' in k: + out['lr'] = float(v) + else: + out[k.replace(' ', '_')] = float(v) + return out + + +def _swan_metrics(summary: Dict[str, Any], log: Optional[Dict[str, Any]], + rows: List[Tuple[float, float, float, int, str]]) -> Dict[str, float]: + """Flat metric dict for swanlab = external reflexion metrics + (when this chunk was + trained) the GRPO built-in metric. acc/adopt/term are only emitted on chunks that had + hard problems, and passk/* only when scored candidates exist, so idle chunks don't dip + the charts to zero.""" + d: Dict[str, float] = { + 'gen/n_hard': summary['n_hard'], 'gen/n_clean': summary['n_clean'], + 'gen/n_leaked': summary['n_leaked'], 'gen/n_train_samples': summary['n_train_samples'], + 'gen/n_reward_pos': summary['n_reward_pos'], + } + if summary['n_hard'] > 0: + d.update({ + 'acc/baseline_pass': summary['avg_baseline_pass_on_hard'], + 'acc/withskill_pass': summary['avg_withskill_pass'], + 'acc/lift': summary['avg_lift'], + 'adopt/A': summary['view_A']['adoption_rate'], + 'adopt/B': summary['view_B']['adoption_rate'], + 'term/withskill': summary['termination_rate_withskill'], + }) + if rows: + m = lambda i: sum(r[i] for r in rows) / len(rows) + d.update({'passk/baseline_mean': m(0), 'passk/bestN_mean': m(1), 'passk/avgN_mean': m(2)}) + if log: + d['train/n_steps'] = log['n_steps'] + d.update({f'train/{k}': v for k, v in _clean_metric(log.get('metric')).items()}) + return d + + def main() -> None: args = _build_args() - if args.sft_batch_size % TRAIN_GPUS != 0 or args.train_batch % args.sft_batch_size != 0: + if args.sft_batch_size % TRAIN_DP != 0: raise ValueError(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple ' - f'of TRAIN_GPUS ({TRAIN_GPUS}) and divide --train-batch ' - f'({args.train_batch})') - steps_per_round = args.train_batch // args.sft_batch_size + f'of the training dp size ({TRAIN_DP})') + # LR schedule is sized by an upper bound on optimizer steps (one pass over each + # trained chunk's candidates); the exact count varies, cosine just decays slower. + steps_per_round = max(1, (args.chunk_size * args.n_skills) // args.sft_batch_size) records = load_math(n=args.n, seed=args.seed) os.makedirs(args.output_dir, exist_ok=True) data_path = os.path.join(args.output_dir, 'skill_dataset.jsonl') + gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') + train_log_path = os.path.join(args.output_dir, 'train_log.jsonl') if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY')): sys.stderr.write('[rft] WARNING: no LLM_BACKUP_API_KEY/OPENAI_API_KEY — ' 'LeakVerifier will report no_llm and skip leak filtering\n') + use_swan = swanlab is not None and os.environ.get('SWANLAB_MODE') != 'disabled' + if use_swan: + swanlab.init(project=args.swanlab_project, + experiment_name=(args.swanlab_exp or None), + config={'model': GEN_MODEL_ID, 'n': len(records), + 'n_skills': args.n_skills, 'view_b_frac': args.view_b_frac, + 'skill_gen_temp': args.skill_gen_temperature, + 'grpo_epsilon': args.grpo_epsilon, 'lr': args.lr}) + # -- Device groups: train (FSDP2) + two independent vLLM samplers. -- r0, r1, r2 = TRAIN_GPUS, TRAIN_GPUS + SKILL_SAMPLER_GPUS, NUM_GPUS device_groups = [ @@ -326,18 +938,18 @@ def main() -> None: twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, groups=device_groups, lazy_collect=False) - # -- Skill model: full-param FSDP2, causal-LM SFT. -- - train_mesh = DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_GPUS) + # -- Skill model: full-param FSDP2, GRPO policy update. -- + train_mesh = DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_DP, fsdp_size=TRAIN_FSDP) skill_model = TransformersModel(model_id=GEN_MODEL_ID, device_mesh=train_mesh, remote_group='train', ddp_config={'find_unused_parameters': False}) from twinkle.patch.no_split_modules import NoSplitModulesPatch - skill_model.apply_patch(NoSplitModulesPatch({'Qwen3_5DecoderLayer'})) - skill_model.set_template('Qwen3_5Template', model_id=GEN_MODEL_ID, + skill_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) + skill_model.set_template(Template, model_id=GEN_MODEL_ID, enable_thinking=True, max_length=args.max_model_len, truncation_strategy='delete') skill_model.set_processor(InputProcessor, padding_free=True) - skill_model.set_loss('CrossEntropyLoss') + skill_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) skill_model.set_optimizer('AdamW', lr=args.lr) skill_model.set_lr_scheduler('CosineWarmupScheduler', num_warmup_steps=10, num_training_steps=args.max_train_rounds * steps_per_round) @@ -350,7 +962,7 @@ def main() -> None: 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, device_mesh=DeviceMesh.from_sizes(world_size=SKILL_SAMPLER_GPUS, dp_size=skill_dp), remote_group='skill_sampler') - skill_sampler.set_template('Qwen3_5Template', model_id=GEN_MODEL_ID, + skill_sampler.set_template(Template, model_id=GEN_MODEL_ID, enable_thinking=True, max_length=args.max_model_len) base_sampler = vLLMSampler( model_id=GEN_MODEL_ID, @@ -358,55 +970,107 @@ def main() -> None: 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, device_mesh=DeviceMesh.from_sizes(world_size=BASE_SAMPLER_GPUS, dp_size=base_dp), remote_group='base_sampler') - base_sampler.set_template('Qwen3_5Template', model_id=GEN_MODEL_ID, + base_sampler.set_template(Template, model_id=GEN_MODEL_ID, enable_thinking=True, max_length=args.max_model_len) ckpt = CheckpointEngineManager(model=skill_model, sampler=skill_sampler) - leak = LeakVerifier(sampler=None) # backup teacher only, no local judge + leak = LeakVerifier(sampler=None, answer_only=True) # flag ONLY the final answer (view A only; view B skips leak) + # leak = LeakVerifier(sampler=None, judge_system=_LEAK_JUDGE_SYSTEM) # stricter: also flag concrete intermediate key results + checker = _build_rubric_checker() # view-A process-check (teacher-only); None if no LLM backup + if checker is None: + sys.stderr.write('[rft] no LLM backup env -> view-A rubric process-check DISABLED ' + '(skill-gen diagnoses from the attempt alone)\n') sys.stderr.write(f'[rft] {len(records)} MATH problems; train_gpus={TRAIN_GPUS} ' f'skill_dp={skill_dp} base_dp={base_dp}\n') - pool: List[Dict[str, Any]] = [] - pending = 0 - rounds = 0 + # -- Sequential: generate one chunk, train on it, sync -> exact on-policy GRPO. + # Generation dominates wall-clock, so not overlapping training costs little, and + # it removes all producer/consumer concurrency (no thread, no lock). -- n_chunks = (len(records) + args.chunk_size - 1) // args.chunk_size + cfg = {'record_type': 'config', 'model': GEN_MODEL_ID, 'dataset': 'math', + 'n': len(records), 'seed': args.seed, 'n_skills': args.n_skills, + 'view_b_frac': args.view_b_frac, 'skill_retries': args.skill_retries, + 'skill_gen_temp': args.skill_gen_temperature, + 'skill_gen_top_p': args.skill_gen_top_p, 'skill_gen_top_k': args.skill_gen_top_k, + 'reward': 'greedy_binary(correct)', 'advantage': 'group_relative', + 'format_in_reward': args.format_in_reward, + 'rubric_check': 'fixed_math_5crit(viewA)' if checker else 'disabled', + 'grpo_epsilon': args.grpo_epsilon, 'lr': args.lr, + 'max_train_rounds': args.max_train_rounds, 'started': int(time.time())} + hist: List[Dict[str, float]] = [] + rounds = 0 + with open(gen_path, 'w', encoding='utf-8') as gen_f, \ + open(data_path, 'w', encoding='utf-8') as data_f, \ + open(train_log_path, 'w', encoding='utf-8') as tlog: + for f in (gen_f, data_f, tlog): + f.write(json.dumps(cfg, ensure_ascii=False) + '\n') + f.flush() + gstep, epoch = 0, 0 + # Multiple epochs over the SAME problem set: each pass reshuffles and RE-GENERATES + # rollouts with the current (improved) policy, so every chunk stays on-policy (no + # importance correction needed) -- the online analogue of SEAM's fixed-data epochs. + while rounds < args.max_train_rounds and n_chunks > 0: + if epoch > 0: + np.random.RandomState(args.seed + epoch).shuffle(records) + for ci in range(n_chunks): + if rounds >= args.max_train_rounds: + break + chunk = records[ci * args.chunk_size:(ci + 1) * args.chunk_size] + full, summary, groups = process_chunk( + base_sampler, skill_sampler, leak, chunk, gstep, base_dp, skill_dp, + args, checker) - with open(data_path, 'w', encoding='utf-8') as fout: - fout.write(json.dumps({ - 'record_type': 'config', 'model': GEN_MODEL_ID, 'dataset': 'math', - 'n': len(records), 'seed': args.seed, 'pass_k': args.pass_k, - 'hard_baseline_max': args.hard_baseline_max, 'n_skills': args.n_skills, - 'train_every': args.train_every, 'train_batch': args.train_batch, - 'lr': args.lr, 'started': int(time.time()), - }, ensure_ascii=False) + '\n') - fout.flush() - - for ci in range(n_chunks): - if rounds >= args.max_train_rounds: - break - chunk = records[ci * args.chunk_size:(ci + 1) * args.chunk_size] - sys.stderr.write(f'[rft] chunk {ci+1}/{n_chunks} ({len(chunk)} problems)\n') - recs = generate_chunk(base_sampler, skill_sampler, leak, chunk, - base_dp, skill_dp, args) - for rec in recs: - fout.write(json.dumps(rec, ensure_ascii=False) + '\n') - fout.flush() - pool.extend(recs) - pending += len(recs) - - while pending >= args.train_every and rounds < args.max_train_rounds: - _train_round(skill_model, ckpt, pool, args) - pending -= args.train_every - rounds += 1 - sys.stderr.write(f'[rft] train round {rounds}/{args.max_train_rounds} ' - f'(pool={len(pool)})\n') - if rounds % args.save_rounds == 0: - skill_model.save(f'skill-rft-{rounds}', output_dir=args.output_dir) + log = None + if groups: # on-policy GRPO update on this chunk, then weights sync + log = _train_chunk(skill_model, ckpt, groups, args) + rounds += 1 + log.update({'record_type': 'train_round', 'round': rounds, + 'chunk': gstep, 'epoch': epoch, 'ts': int(time.time())}) + tlog.write(json.dumps(log, ensure_ascii=False) + '\n') + tlog.flush() + if rounds % args.save_rounds == 0: + skill_model.save(f'skill-rft-{rounds}', output_dir=args.output_dir) + + summary['rounds_done'], summary['epoch'] = rounds, epoch + for rec in full: + gen_f.write(json.dumps(rec, ensure_ascii=False) + '\n') + gen_f.write(json.dumps(summary, ensure_ascii=False) + '\n') + gen_f.flush() + for v in groups: + data_f.write(json.dumps(v, ensure_ascii=False) + '\n') + data_f.flush() + + sa, sb = summary['view_A'], summary['view_B'] + hist.append({'aA': sa['adoption_rate'], 'aB': sb['adoption_rate'], + 'lift': summary['avg_lift'], 'pos': summary['n_reward_pos']}) + sys.stderr.write( + f'[gen] e{epoch} chunk {ci+1}/{n_chunks} (g{gstep}): hard={summary["n_hard"]} ' + f'clean={summary["n_clean"]} train={summary["n_train_samples"]} ' + f'acc={summary["avg_baseline_pass_on_hard"]:.2f}->{summary["avg_withskill_pass"]:.2f} ' + f'lift={summary["avg_lift"]:+.3f} ' + f'A[{sa["n_hard"]}h {sa["adoption_rate"]:.2f}] ' + f'B[{sb["n_hard"]}h {sb["adoption_rate"]:.2f}] ' + f'rounds={rounds}' + + (f' metric={log.get("metric")}' if log else '') + '\n') + # -- per-query passk (base vs best/avg of N skills) + swanlab metrics -- + rows = _query_rows(full) + for base_p, best_p, avg_p, nsc, prob in rows: + logger.info(f'[q] g{gstep} base={base_p:.2f} bestN={best_p:.2f} avgN={avg_p:.2f} ' + f'n={nsc} | {prob[:70].replace(chr(10), " ")}') + if use_swan: + swanlab.log(_swan_metrics(summary, log, rows), step=gstep) + + if (gstep + 1) % args.trend_every == 0: + tl = _trend_line(hist, args.trend_every, rounds) + if tl: + sys.stderr.write(tl + '\n') + gstep += 1 + epoch += 1 skill_model.save('skill-rft-final', output_dir=args.output_dir) - sys.stderr.write(f'[rft] done: {len(pool)} valid records, {rounds} train rounds; ' - f'dataset -> {data_path}\n') + sys.stderr.write(f'[rft] done: {rounds} train rounds over {gstep} chunks / {epoch} epochs; ' + f'data -> {data_path}\n') if __name__ == '__main__': diff --git a/cookbook/exp/embedding/train_reflexion_skill_rft.sh b/cookbook/exp/embedding/train_reflexion_skill_rft.sh index 7ad6a607a..524570a35 100644 --- a/cookbook/exp/embedding/train_reflexion_skill_rft.sh +++ b/cookbook/exp/embedding/train_reflexion_skill_rft.sh @@ -5,21 +5,26 @@ set -euo pipefail -export GEN_MODEL_ID=${GEN_MODEL_ID:-Qwen/Qwen3.5-4B} -export MATH_DATA_DIR=${MATH_DATA_DIR:-./output/math_data/MATH} +export GEN_MODEL_ID=${GEN_MODEL_ID:-Qwen/Qwen3-4B} +# Local MATH copy (modelscope download cache). Override MATH_DATA_DIR if the +# cache hash dir changes or the data lives elsewhere. +export MATH_DATA_DIR=${MATH_DATA_DIR:-/mnt/workspace/.cache/modelscope/hub/datasets/downloads/extracted/0744cd2d347a7e8f85f7087d950b2ed38b626a5c808c5399e2d8a0923d42d013/MATH} export LLM_BACKUP_API_KEY=${LLM_BACKUP_API_KEY:?set LLM_BACKUP_API_KEY for the leak judge} export LLM_BACKUP_BASE_URL=${LLM_BACKUP_BASE_URL:-https://dashscope.aliyuncs.com/compatible-mode/v1} export LLM_BACKUP_MODEL=${LLM_BACKUP_MODEL:-qwen3.7-max} python cookbook/exp/embedding/train_reflexion_skill_rft.py \ - --n 2000 \ + --n 5000 \ --chunk-size 16 \ --n-skills 8 \ - --pass-k 8 \ - --hard-baseline-max 0.25 \ - --train-every 64 \ - --train-batch 64 \ + --view-b-frac 0.5 \ + --skill-retries 2 \ + --max-tokens 25000 \ + --max-model-len 30000 \ --sft-batch-size 8 \ - --lr 1e-5 \ - --max-train-rounds 200 \ + --grpo-epsilon 0.2 \ + --lr 6e-6 \ + --max-train-rounds 1500 \ + --save-rounds 25 \ + --trend-every 10 \ --output-dir ./output/reflexion_skill_rft diff --git a/src/twinkle_agentic/verifier/leak_verifier.py b/src/twinkle_agentic/verifier/leak_verifier.py index f1301c66d..f1cf730ac 100644 --- a/src/twinkle_agentic/verifier/leak_verifier.py +++ b/src/twinkle_agentic/verifier/leak_verifier.py @@ -62,6 +62,27 @@ Reply with exactly one word: LEAK or CLEAN.""" +# Answer-only variant: flags ONLY final-answer disclosure and explicitly PERMITS +# revealing the method / approach / step plan. For skill-generation training, the +# method IS the product we want, so the default prompt's "decisive intermediate +# result" and "step-by-step plan that solves this task" clauses over-flag good +# skills; this variant drops them. +_JUDGE_SYSTEM_ANSWER_ONLY = """\ +You check whether a HINT that will be shown to someone solving a TASK reveals the +TASK's FINAL ANSWER. + +The hint may FREELY describe general methods, the solution approach, which +technique to apply, the steps to take, common pitfalls, or sanity checks — even +if that strongly points at how to solve THIS task. Revealing the METHOD or PLAN +is allowed and expected of a good hint. + +The hint LEAKS only if it reveals the FINAL ANSWER itself — the concrete final +value, expression, choice, label, or verbatim result the task asks for (or a +trivially equivalent restatement of it). Describing HOW to get there WITHOUT +stating the resulting final value does NOT leak. + +Reply with exactly one word: LEAK or CLEAN.""" + _JUDGE_USER = """\ ## Task {query} @@ -135,6 +156,8 @@ class LeakVerifier(Verifier): max_content_chars / max_query_chars: truncation caps for the judge input. uncertain_is_leak: if the judge reply is unparseable, treat it as a leak (default False: keep the hint, the conservative choice). + judge_system: optional custom judge system prompt that overrides the built-in + answer_only/legacy criteria (for task-specific leak policies). """ def __init__( @@ -147,6 +170,8 @@ def __init__( max_content_chars: int = 4000, max_query_chars: int = 4000, uncertain_is_leak: bool = False, + answer_only: bool = False, + judge_system: Optional[str] = None, ): self.sampler = sampler self.model_path = model_path @@ -155,6 +180,13 @@ def __init__( self.max_content_chars = int(max_content_chars) self.max_query_chars = int(max_query_chars) self.uncertain_is_leak = bool(uncertain_is_leak) + # answer_only: flag ONLY final-answer disclosure, permit method/plan (see + # _JUDGE_SYSTEM_ANSWER_ONLY). Default keeps the strict legacy behaviour. + self.answer_only = bool(answer_only) + # judge_system: caller-supplied criterion that overrides both built-ins, for + # task-specific leak policies (e.g. permit method/plan but still flag concrete + # intermediate key results). None -> fall back to the answer_only/legacy pair. + self.judge_system = judge_system or None # ------------------------------------------------------------------ # public entry points @@ -240,8 +272,10 @@ def _judge_trajectory(self, content: str, query: str, query=self._trim(query, self.max_query_chars), reference_block=ref_block, hint=self._trim(content, self.max_content_chars)) + system = self.judge_system or ( + _JUDGE_SYSTEM_ANSWER_ONLY if self.answer_only else _JUDGE_SYSTEM) return {'messages': [ - {'role': 'system', 'content': _JUDGE_SYSTEM}, + {'role': 'system', 'content': system}, {'role': 'user', 'content': user}, ]} From c214d0ed3553f1fedfbe1ad927ab2b1f6923d726 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Wed, 15 Jul 2026 11:46:40 +0800 Subject: [PATCH 08/60] wip --- cookbook/exp/embedding/eval_gpqa_rag.py | 111 +-- .../embedding/train_reflexion_skill_rft.py | 676 +++++++++++++++--- .../embedding/train_reflexion_skill_rft.sh | 5 + 3 files changed, 641 insertions(+), 151 deletions(-) diff --git a/cookbook/exp/embedding/eval_gpqa_rag.py b/cookbook/exp/embedding/eval_gpqa_rag.py index 72a5a9b30..af5cf60dc 100644 --- a/cookbook/exp/embedding/eval_gpqa_rag.py +++ b/cookbook/exp/embedding/eval_gpqa_rag.py @@ -509,29 +509,28 @@ def normalize_answer(ans: str) -> str: """Normalize a math answer string for comparison.""" if not ans: return '' - s = ans.strip() - # MCQ: extract bare letter from \textbf{(D)}, \text{(A)}, \mathbb{A}, (B), etc. + s = str(ans).strip() m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) if m: return m.group(1) - s = s.replace(' ', '') - s = s.replace(r'\,', '') - s = s.replace(r'\;', '') - s = s.replace(r'\!', '') - s = s.replace(r'\text', '') - s = s.replace(r'\mathrm', '') + s = s.strip('$').strip() + s = s.replace('−', '-') + 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 = 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 = s.replace(r'\,', '').replace(r'\;', '').replace(r'\!', '') + s = re.sub(r'\\(?:quad|qquad|\s)', '', s) + s = re.sub(r'\s+', '', s) + s = re.sub(r'\\sqrt([0-9A-Za-z])', r'\\sqrt{\1}', s) + s = re.sub(r'\\frac\{([^{}]+)\}([^{}\\])', r'\\frac{\1}{\2}', s) + s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) 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) + s = re.sub(r'\^\\circ|\^\{\\circ\}|°|\\circ', 'deg', s) + s = s.replace(r'\minus{}', '-').replace(r'\minus', '-') + def _frac_to_slash(m): - # Handle nested braces in numerator/denominator text = m.group(0) pos = text.index('{') + 1 depth, num_start = 1, pos @@ -540,7 +539,7 @@ def _frac_to_slash(m): elif text[pos] == '}': depth -= 1 pos += 1 numer = text[num_start:pos - 1] - pos += 1 # skip '{' + pos += 1 den_start = pos depth = 1 while depth > 0: @@ -549,9 +548,8 @@ def _frac_to_slash(m): 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: 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: @@ -574,79 +572,98 @@ def _eval_frac(s): except (ValueError, ZeroDivisionError): pass return None + va, vb = _eval_frac(a), _eval_frac(b) if va is not None and vb is not None: return abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) return False -# MCQ compound pattern: \text{(D) }49, \textbf{(C)}12, (B) 21, etc. _MCQ_REF_RE = re.compile( r'^\\(?:textbf|text|mathrm|mathbf)\{\(?([A-E])\)?\s*\}\s*(.+)$' r'|^\(?([A-E])\)\s+(.+)$' ) +_VAR_PREFIX_RE = re.compile(r'^(?:[A-Za-z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)$', re.DOTALL) 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)``. - """ + """Split an MCQ answer into (letter, value) components.""" 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. +def _strip_var_prefix(ans: str) -> str: + m = _VAR_PREFIX_RE.match((ans or '').strip()) + return m.group(1).strip() if m else (ans or '') - 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. - """ + +def _normalize_tuple(ans: str) -> str: + return re.sub(r'[\s()\[\]{}\\]', '', ans or '') + + +def _math_verify_equal(predicted: str, reference: str) -> bool: + try: + from math_verify import parse, verify + gold = parse(r'\boxed{' + reference + '}', parsing_timeout=1) + pred = parse(r'\boxed{' + predicted + '}', parsing_timeout=1) + return bool(gold and pred and verify(gold, pred, timeout_seconds=1)) + except Exception: + 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: + if norm_p == norm_r or norm_p.lower() == norm_r.lower(): return True if _try_numeric_equal(norm_p, norm_r): return True - # Symmetric MCQ matching: decompose both sides into (letter, value). + stripped_p = normalize_answer(_strip_var_prefix(predicted)) + stripped_r = normalize_answer(_strip_var_prefix(reference)) + if stripped_p and stripped_r: + if stripped_p == stripped_r or stripped_p.lower() == stripped_r.lower(): + return True + if _try_numeric_equal(stripped_p, stripped_r): + return True + 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): + if p_val == r_val or p_val.lower() == r_val.lower() or _try_numeric_equal(p_val, r_val): return True - return False + + for left, right in ((norm_p, norm_r), (stripped_p, stripped_r)): + tuple_l, tuple_r = _normalize_tuple(left), _normalize_tuple(right) + if ',' in tuple_l and tuple_l == tuple_r: + return True + + if '=' in norm_r and '=' not in norm_p: + if any(part and (part == norm_p or _try_numeric_equal(part, norm_p)) for part in norm_r.split('=')): + return True + if '=' in norm_p and '=' not in norm_r: + if any(part and (part == norm_r or _try_numeric_equal(part, norm_r)) for part in norm_p.split('=')): + return True + + return _math_verify_equal(stripped_p or norm_p, stripped_r or norm_r) # --------------------------------------------------------------------------- diff --git a/cookbook/exp/embedding/train_reflexion_skill_rft.py b/cookbook/exp/embedding/train_reflexion_skill_rft.py index 33bf4dc20..12852c74f 100644 --- a/cookbook/exp/embedding/train_reflexion_skill_rft.py +++ b/cookbook/exp/embedding/train_reflexion_skill_rft.py @@ -56,7 +56,7 @@ # Reuse the reference eval's dataset + grading + prompts + sampling config, and the # phase-0 pipeline's parsing / rollout / injection helpers (Find > Create). from eval_gpqa_rag import (GEN_GPU_MEM, GEN_MODEL_ID, build_direct_prompt, # noqa: F401 - load_math) + load_aops, load_math) from eval_reflexion_skill import (_EX_PROBLEM, _clean_text, # noqa: F401 _parse_seq, _run_samples, build_skill_solve_prompt) @@ -136,16 +136,22 @@ 'You are a mathematics coach. You are shown a competition problem together with an ' 'automated process-check of an earlier solver attempt at it -- which solution ' 'criteria the attempt passed or failed, and suggested fixes for the failures. You ' - 'do NOT see the attempt itself, only this check. Use the flagged failures to see ' - 'where this KIND of problem tends to trip solvers up, then distil a few reusable ' - 'tips that would prevent those specific failures.\n\n' + 'do NOT see the attempt itself, only this check. Treat the check as privileged ' + 'training scaffolding: study it together with the problem, identify the ' + 'problem-visible features that make each useful flagged failure relevant, then ' + 'rephrase those lessons as self-contained reusable skills. The goal is not to ' + 'continue from the check, cite it, or hide it silently; the goal is to turn it into ' + 'a problem-triggered reasoning pattern a query-only solver could reproduce later.\n\n' + 'Good skills name the observable trigger, the method worth reaching for, the ' + 'pitfall to watch, and a quick verification habit. Prefer formulations like ' + '"When a configuration has ...", "Before setting up ...", or "Check whether ..." ' + 'over references to the process-check, failed criteria, or the earlier attempt. ' "These tips are advisory: they will be placed in a solver's system prompt as gentle " - 'reminders before it works through a SIMILAR problem on its own. So keep them ' - 'general and transferable — the method worth reaching for, the pitfall to watch and ' - 'a quick check, and the discipline to settle on a final answer — rather than a ' - 'worked solution to this exact problem, and without stating its specific ' - 'intermediate values or its final answer. Think briefly first, then give your tips ' - 'as a markdown bullet list wrapped in and , like the example below.' + 'reminders before it works through a SIMILAR problem on its own, without seeing ' + 'this process-check. So keep them general and transferable rather than a worked ' + 'solution to this exact problem, and do not state its specific intermediate values ' + 'or final answer. Think briefly first, then give your tips as a markdown bullet ' + 'list wrapped in and , like the example below.' ) # One-shot demo of the recommended mix (method / pitfall+check / procedure / @@ -172,7 +178,10 @@ 'Process check of an earlier attempt (automated rubric verifier -- treat as ' 'evidence, not gospel; PASS/FAIL per criterion with suggested fixes for failures):\n' '{diagnosis}\n\n' - 'Now output the skills bullet list.' + 'Now output a self-contained skills bullet list. Each bullet should still be useful ' + 'if the process check were removed: connect any useful flagged failure to ' + 'problem-visible features, general methods, and quick checks rather than citing the ' + 'rubric or the earlier attempt.' ) @@ -371,6 +380,66 @@ def _extract_skills_block(text: str) -> Optional[str]: # localisation instead of guessing. Teacher-only (sampler=None -> every diagnose() # hits llm_backup); mirrors eval_dualline_math's fixed math rubric. # --------------------------------------------------------------------------- +_RFT_DIAG_SYSTEM = """\ +You are a process error checker for a math solution attempt. You are given a math +problem, a rubric, and one attempted solution segment. Decide PASS or FAIL for each +criterion and explain only the process error type. + +Output STRICT JSON (no prose outside it) with this shape: +{ + "items": [ + {"index": 1, "verdict": "PASS", "reason": "", + "fix": ""}, + {"index": 2, "verdict": "FAIL", "reason": "", + "fix": ""} + ], + "overall": "OK" | "ISSUES", + "summary": "" +} + +Rules: +- Judge every criterion independently and literally; a [Hard Rule] is FAIL unless + unambiguously satisfied. +- Judge ONLY what is observable in THIS segment. +- Content inside ... (or ) is internal reasoning, not + user-facing output; ignore it for "output only X" style criteria. +- For PASS items, leave "fix" as "". +- For FAIL items, "reason", "fix", and "summary" must describe only the flawed + step, theorem, arithmetic operation, case split, or verification habit. +- NEVER state the correct final answer, corrected final expression, option letter, + graph/choice label, or any exact value that the answer should become. +- NEVER write phrases like "the correct answer is", "which gives", "yielding", + "should be ", "Option ", or "Graph ". +- If a fix would require naming a corrected value, replace it with a method-level + instruction such as "redo that computation carefully" or "apply the theorem with + the correct quantities". +- Keep every "reason" and "fix" clear and concise — one short sentence each. +- "overall" is "OK" only if NO criterion is FAIL. +- Output only the JSON object.""" + + +_RFT_DIAG_USER = """\ +## Task / query (context) +{query} + +## Rubric +{rubric} + +## Segment +{segment} + +Now output the diagnostic JSON object.""" + + +class _RftRubricVerifier(RubricVerifier): + def _diagnose_trajectory(self, query: str, rubric_block: str, segment_text: str) -> dict: + user = _RFT_DIAG_USER.format(query=query, rubric=rubric_block, segment=segment_text) + return {'messages': [ + {'role': 'system', 'content': _RFT_DIAG_SYSTEM}, + {'role': 'user', 'content': user}, + ]} + + _MATH_RUBRIC = [ ('The reasoning contains no arithmetic or algebraic error', True), ('Each step follows logically from the previous ones', True), @@ -385,7 +454,7 @@ def _build_rubric_checker() -> Optional['RubricVerifier']: if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') or os.environ.get('OPENAI_API_KEY')): return None - return RubricVerifier( + return _RftRubricVerifier( fixed_rubric=[RubricItem(t, is_hard=h) for t, h in _MATH_RUBRIC], gate=True) @@ -525,6 +594,17 @@ def _view_stats(hard: List[Dict[str, Any]], view: str) -> Dict[str, Any]: } +def _is_trainable(c: Dict[str, Any], args: argparse.Namespace) -> bool: + """A candidate reaches the GRPO update iff its advantage is non-zero. With + --format-in-reward every candidate carries a reward (unparseable/leaked score 0), + so non-zero advantage is the only gate; otherwise it must also be clean and scored. + Single source of truth for both the summary counts and ``_group_records``.""" + adv_nz = abs(c.get('advantage') or 0.0) > 1e-9 + if args.format_in_reward: + return adv_nz + return c['leaked'] is False and c.get('with_pass') is not None and adv_nz + + def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespace) -> Dict[str, Any]: """Per-chunk aggregates — watch these across chunks to see if the RFT'd skill model produces better skills over time (yield, leak rate, lift, termination).""" @@ -536,10 +616,15 @@ def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespac ws_rolls = [x for c in scored for x in c['rolls']] # With --format-in-reward, unparseable/leaked candidates also carry a (0) reward and are # trained, so count trainables over ALL candidates; else only clean scored ones. - train_cands = ([c for c in all_cands if abs(c.get('advantage') or 0.0) > 1e-9] if args.format_in_reward - else [c for c in scored if c['leaked'] is False and abs(c.get('advantage') or 0.0) > 1e-9]) + train_cands = [c for c in all_cands if _is_trainable(c, args)] base_acc = (sum(r['_baseline_pass'] for r in hard) / len(hard)) if hard else 0.0 ws_acc = (sum(c['with_pass'] for c in scored) / len(scored)) if scored else 0.0 + # -- signal-source monitor: how much of the GRPO signal comes from base-FAIL problems + # (the offensive "rescue a failure" signal we want) vs base-success (defensive "don't + # break an easy one"). abs_adv_from_fail_frac ~0.1 was the diagnosed failure mode. -- + fail_cands = [c for r in chunk if r['_failed'] for c in r['_cands']] + abs_adv = lambda cs: sum(abs(c.get('advantage') or 0.0) for c in cs) + total_abs = abs_adv(all_cands) return { 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), 'n_failed_first_try': len(failed), 'n_hard': len(hard), @@ -550,6 +635,8 @@ def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespac 'n_clean': sum(1 for c in cands if c['leaked'] is False), 'n_reward_pos': sum(1 for c in scored if c['reward']), 'n_train_samples': len(train_cands), + 'n_train_from_fail': sum(1 for c in fail_cands if _is_trainable(c, args)), + 'abs_adv_from_fail_frac': (abs_adv(fail_cands) / total_abs) if total_abs > 0 else 0.0, 'avg_baseline_pass_on_hard': base_acc, 'avg_withskill_pass': ws_acc, 'avg_lift': ws_acc - base_acc, @@ -570,10 +657,7 @@ def _group_records(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> Lis continue view = r.get('_view', 'A') for c in r['_cands']: - adv_nz = abs(c.get('advantage') or 0.0) > 1e-9 - trainable = adv_nz if args.format_in_reward else ( - c['leaked'] is False and c['with_pass'] is not None and adv_nz) - if trainable: + if _is_trainable(c, args): out.append({ 'problem': r['problem'], 'reference_answer': r['reference_answer'], 'view': view, @@ -585,6 +669,212 @@ def _group_records(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> Lis return out +def _apply_baseline(r: Dict[str, Any], roll: Dict[str, Any]) -> None: + """Write a (cached or fresh) greedy baseline roll onto a problem and RESET the per-chunk + working state, so a problem reused in a later chunk never carries prior skill candidates.""" + r['_baseline_rolls'], r['_cands'] = [roll], [] + r['_init'] = [roll] + r['_failed'] = not roll['correct'] + r['_baseline_pass'] = 1.0 if roll['correct'] else 0.0 + r['_hard'] = True # process EVERY selected problem; group variance selects + + +def _baseline_rollout(base_sampler, problems: List[Dict[str, Any]], base_dp: int, + args: argparse.Namespace, cache: Dict[str, Dict[str, Any]]) -> int: + """Phase 1: base solves each problem GREEDILY once (T=0, M=1), keyed-cached by problem + text across chunks. The base sampler is FROZEN and decoding is greedy, so a problem's + baseline never changes over the run -- a cache hit is exact and skips the sampler. + Returns the number of FRESH sampler rollouts (cache misses) for efficiency reporting.""" + todo = [r for r in problems if r['problem'] not in cache] + if todo: + base_out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], + 1, args.max_tokens, base_dp, temperature=0.0) + for r, seqs in zip(todo, base_out): + cache[r['problem']] = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + for r in problems: + _apply_baseline(r, cache[r['problem']]) + return len(todo) + + +def _baseline_class(r: Dict[str, Any]) -> str: + """Bucket a baselined problem by its greedy outcome: ``success`` (base solved it), + ``fail_loop`` (ran the length budget out / never terminated -- the mode skills rescue + best), or ``fail_wrong`` (terminated cleanly but the answer is wrong).""" + roll = r['_init'][0] + if roll['correct']: + return 'success' + if roll['stop_reason'] == 'length' or not roll['terminated']: + return 'fail_loop' + return 'fail_wrong' + + +_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') + + +def _norm_num_text(num: str) -> str: + try: + f = float(num) + return str(int(f)) if f == int(f) else str(f) + except Exception: + return str(num).strip() + + +def _numeric_value(raw: Any) -> Optional[str]: + if raw is None: + return None + s = str(raw).strip().strip('$').strip() + s = s.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac') + s = re.sub(r'\\!|\\,|\\;|\\ |\\left|\\right|\s', '', s) + m = re.fullmatch(r'\\frac\{(-?\d+)\}\{(-?\d+)\}', s) + if m: + a, b = int(m.group(1)), int(m.group(2)) + return _norm_num_text(str(a / b)) if b else None + m = re.fullmatch(r'(-?\d+)/(-?\d+)', s) + if m: + a, b = int(m.group(1)), int(m.group(2)) + return _norm_num_text(str(a / b)) if b else None + if _NUM_RE.fullmatch(s): + return _norm_num_text(s) + return None + + +def _numeric_only_records(records: List[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], int]: + out = [] + dropped = 0 + for r in records: + ref = _numeric_value(r.get('reference_answer')) + if ref is None: + dropped += 1 + continue + rr = dict(r) + rr['reference_answer'] = ref + out.append(rr) + return out, dropped + + +def _load_records(args: argparse.Namespace) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], Dict[str, int]]: + need_split = args.eval_size > 0 + load_n = 0 if (args.numeric_only or need_split) else args.n + records = (load_aops(n=load_n, seed=args.seed) if args.dataset == 'aops' + else load_math(n=load_n, seed=args.seed)) + raw_n = len(records) + dropped = 0 + if args.numeric_only: + records, dropped = _numeric_only_records(records) + rng = np.random.RandomState(args.seed) + rng.shuffle(records) + eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 + eval_records = [dict(r) for r in records[:eval_n]] + train_pool = records[eval_n:] + train_n = args.n if args.n > 0 else len(train_pool) + train_records = [dict(r) for r in train_pool[:train_n]] + overlap = {r['problem'] for r in train_records} & {r['problem'] for r in eval_records} + if overlap: + raise ValueError(f'fixed eval/train overlap detected: {len(overlap)} duplicated problems') + stats = {'raw_loaded': raw_n, 'numeric_dropped': dropped, + 'train_records': len(train_records), 'eval_records': len(eval_records)} + return train_records, eval_records, stats + + +class _ProblemPool: + """Cyclic draw source over the loaded problems. Each full pass reshuffles with + ``seed + epoch`` and bumps ``epoch`` (matching the old per-epoch reshuffle); the + initial pass keeps the loader's shuffled order. Draws never run out.""" + + def __init__(self, records: List[Dict[str, Any]], seed: int): + self._records = list(records) + self._seed = seed + self._cursor = 0 + self.epoch = 0 + self.baseline_cache: Dict[str, Dict[str, Any]] = {} # problem text -> frozen greedy roll + + def draw(self, k: int) -> List[Dict[str, Any]]: + """Return ``k`` DISTINCT problems (unique within this call, so one chunk never + processes the same problem twice even when the cursor wraps mid-draw). ``k`` is + always << pool size, so this terminates.""" + out: List[Dict[str, Any]] = [] + seen: set = set() + while len(out) < k: + if self._cursor >= len(self._records): + self.epoch += 1 + np.random.RandomState(self._seed + self.epoch).shuffle(self._records) + self._cursor = 0 + r = self._records[self._cursor] + self._cursor += 1 + if id(r) not in seen: + seen.add(id(r)) + out.append(r) + return out + + +def _select_balanced(buckets: Dict[str, List[Dict[str, Any]]], n_success: int, + n_fail: int, n_fail_loop: int) -> List[Dict[str, Any]]: + """Pick the chunk from the baselined buckets: ``n_fail`` base-fails (split toward + ``n_fail_loop`` loop-fails, best-effort) + ``n_success`` base-successes. If a bucket + is too thin to hit ``chunk_size`` the shortfall is topped up from leftovers (the + ratio then drifts, which the caller logs).""" + loop, wrong, succ = buckets['fail_loop'], buckets['fail_wrong'], buckets['success'] + take_loop = min(n_fail_loop, len(loop)) + take_wrong = min(n_fail - take_loop, len(wrong)) + take_loop = min(n_fail - take_wrong, len(loop)) # give loop the remainder if wrong is short + sel = loop[:take_loop] + wrong[:take_wrong] + succ[:n_success] + target = n_success + n_fail + if len(sel) < target: + used = {id(x) for x in sel} + leftover = [x for b in (loop, wrong, succ) for x in b if id(x) not in used] + sel += leftover[:target - len(sel)] + return sel + + +def _draw_chunk(pool: _ProblemPool, base_sampler, base_dp: int, args: argparse.Namespace + ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + """Draw one training chunk, running baseline rollout (Phase 1) on every drawn problem. + + With ``--balance`` off, draw ``chunk_size`` problems and return them. With it on, keep + drawing+baselining in ``chunk_size`` batches, bucketing by ``_baseline_class``, until the + target base fail:success mix is reachable or the draw budget is hit; then select a + balanced subset. Returns ``(chunk, stats)`` where stats records the realised mix.""" + if not args.balance: + chunk = pool.draw(args.chunk_size) + n_fresh = _baseline_rollout(base_sampler, chunk, base_dp, args, pool.baseline_cache) + return chunk, {'enabled': False, 'n_drawn': len(chunk), 'n_baseline_fresh': n_fresh} + + n_success = max(0, min(args.chunk_size, round(args.chunk_size * args.balance_success_frac))) + n_fail = args.chunk_size - n_success + n_fail_loop = round(n_fail * args.balance_loop_frac) + buckets: Dict[str, List[Dict[str, Any]]] = {'success': [], 'fail_loop': [], 'fail_wrong': []} + budget = args.chunk_size * args.balance_max_draws_mult + n_drawn, n_fresh = 0, 0 + seen: set = set() # dedupe across batches: the pool can re-serve a problem after a wrap + while n_drawn < budget: + if (len(buckets['success']) >= n_success + and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail): + break # enough of both classes buffered to satisfy the target split + batch = pool.draw(args.chunk_size) + n_fresh += _baseline_rollout(base_sampler, batch, base_dp, args, pool.baseline_cache) + n_drawn += len(batch) + for r in batch: + if id(r) in seen: + continue + seen.add(id(r)) + buckets[_baseline_class(r)].append(r) + + target_reached = (len(buckets['success']) >= n_success + and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail) + chunk = _select_balanced(buckets, n_success, n_fail, n_fail_loop) + sel_success = sum(1 for r in chunk if not r['_failed']) + stats = { + 'enabled': True, 'n_drawn': n_drawn, 'n_baseline_fresh': n_fresh, 'n_selected': len(chunk), + 'target_success': n_success, 'target_fail': n_fail, 'target_fail_loop': n_fail_loop, + 'selected_success': sel_success, 'selected_fail': len(chunk) - sel_success, + 'selected_fail_loop': sum(1 for r in chunk if _baseline_class(r) == 'fail_loop'), + 'selected_fail_wrong': sum(1 for r in chunk if _baseline_class(r) == 'fail_wrong'), + 'selected_success_frac': (sel_success / len(chunk)) if chunk else 0.0, + 'budget_hit': not target_reached, # stopped short of the target mix, not by choice + } + return chunk, stats + + def process_chunk(base_sampler, skill_sampler, leak: LeakVerifier, chunk: List[Dict[str, Any]], ci: int, base_dp: int, skill_dp: int, args: argparse.Namespace, checker=None @@ -595,22 +885,15 @@ def process_chunk(base_sampler, skill_sampler, leak: LeakVerifier, Sequential (generate-one-chunk-train-one): generation and the trainer's weight sync never overlap, so no lock is needed. ``base_sampler`` is frozen (never synced); ``skill_sampler`` is synced by the trainer between chunks. + + ``chunk`` arrives ALREADY baselined by ``_draw_chunk`` (Phase 1 ran during the + balanced draw), so every problem carries ``_init``/``_failed``/``_baseline_pass``/ + ``_hard``/``_cands`` -- Phase 1 is not repeated here. """ - # --- Phase 1: base solves each problem GREEDILY once (T=0, M=1) -- this produces the - # view-A attempt and records whether the base already gets it (reporting only). EVERY - # problem is processed (no difficulty gate, SEAM-style): the group-relative advantage - # (Phase 6) gives zero gradient to any problem whose skills all score alike (base-easy - # -> all solve, or hopeless -> all fail), so GRPO's own group variance selects the - # informative problems. Termination is NOT part of the reward (monitored only). --- - base_out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in chunk], - 1, args.max_tokens, base_dp, temperature=0.0) - for r, seqs in zip(chunk, base_out): - roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - r['_baseline_rolls'], r['_cands'] = [roll], [] - r['_init'] = [roll] # the greedy attempt (metric + view A) - r['_failed'] = not roll['correct'] - r['_baseline_pass'] = 1.0 if roll['correct'] else 0.0 # base accuracy (reporting only, NOT in reward) - r['_hard'] = True # process EVERY problem; group variance selects + # Phase 1 (base greedy solve) ran in _draw_chunk so the balancer could classify by + # outcome; every selected problem is processed (no difficulty gate, SEAM-style): the + # group-relative advantage (Phase 6) gives zero gradient to any problem whose skills + # all score alike, so GRPO's own group variance selects the informative problems. hard = chunk # --- Phase 2: assign each problem's view, then rubric-check the view-A attempts so @@ -762,10 +1045,43 @@ def _is_num(v: Any) -> bool: def _build_args() -> argparse.Namespace: p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--n', type=int, default=2000, help='MATH problems to stream.') + p.add_argument('--dataset', choices=('aops', 'math'), default='aops', + help='Problem source. aops (AI-MO competition problems) is much ' + 'harder than MATH, so the base fails more often -> more offensive ' + 'training signal after balanced sampling.') + p.add_argument('--n', type=int, default=2000, + help='Problems to load into the draw pool (cycled/reshuffled across ' + 'epochs; with --balance many more baseline rollouts than this ' + 'may run, but the pool size is fixed here).') p.add_argument('--seed', type=int, default=42) + p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True, + help='Keep only answers that collapse to one integer/decimal/fraction, ' + 'matching SEAM numeric reward and avoiding non-scalar grading noise.') + p.add_argument('--eval-size', type=int, default=128, + help='Fixed holdout problems, sampled before the train pool after all ' + 'filters; set 0 to disable fixed eval.') + p.add_argument('--eval-every', type=int, default=10, + help='Run fixed holdout eval every N generation chunks when --eval-size > 0.') p.add_argument('--chunk-size', type=int, default=16, help='Problems per generation chunk (all sampler calls batched).') + # -- online baseline-balanced sampling (draw+baseline until the chunk hits the + # target base fail:success mix, so the offensive signal is not starved) -- + p.add_argument('--balance', action=argparse.BooleanOptionalAction, default=True, + help='Keep drawing+baselining problems until the chunk matches the ' + 'target base fail:success composition, then select a balanced ' + 'subset. --no-balance draws chunk_size problems directly.') + p.add_argument('--balance-success-frac', type=float, default=0.4, + help='Target fraction of the chunk that the base solves (base-success). ' + '0.4 => 3:2 fail:success; 0.2 => 4:1. The remainder are base-fail.') + p.add_argument('--balance-loop-frac', type=float, default=0.5, + help='Within the base-fail portion, SOFT target fraction of loop-fails ' + '(ran out of length / never terminated) vs non-loop wrong answers. ' + 'Best-effort only: the fail count is filled from whichever bucket ' + 'is available so a thin bucket never starves the chunk.') + p.add_argument('--balance-max-draws-mult', type=int, default=8, + help='Draw budget per chunk as a multiple of chunk_size; once this many ' + 'problems have been baselined the chunk is assembled from whatever ' + 'the buckets hold (ratio may drift; the actual mix is logged).') p.add_argument('--n-skills', type=int, default=8, help='Candidate skills generated per hard problem.') p.add_argument('--view-b-frac', type=float, default=0.5, @@ -882,7 +1198,16 @@ def _swan_metrics(summary: Dict[str, Any], log: Optional[Dict[str, Any]], 'gen/n_hard': summary['n_hard'], 'gen/n_clean': summary['n_clean'], 'gen/n_leaked': summary['n_leaked'], 'gen/n_train_samples': summary['n_train_samples'], 'gen/n_reward_pos': summary['n_reward_pos'], + 'gen/n_train_from_fail': summary['n_train_from_fail'], + 'gen/abs_adv_from_fail_frac': summary['abs_adv_from_fail_frac'], } + bal = summary.get('balance') or {} + if bal.get('enabled'): + d.update({'balance/n_drawn': bal['n_drawn'], + 'balance/n_baseline_fresh': bal['n_baseline_fresh'], + 'balance/selected_success_frac': bal['selected_success_frac'], + 'balance/selected_fail_loop': bal['selected_fail_loop'], + 'balance/selected_fail_wrong': bal['selected_fail_wrong']}) if summary['n_hard'] > 0: d.update({ 'acc/baseline_pass': summary['avg_baseline_pass_on_hard'], @@ -901,6 +1226,109 @@ def _swan_metrics(summary: Dict[str, Any], log: Optional[Dict[str, Any]], return d +def _prefix_metrics(metrics: Dict[str, float], prefix: str) -> Dict[str, float]: + return {f'{prefix}/{k}': v for k, v in metrics.items()} + + +def _greedy_eval_metrics(recs: List[Dict[str, Any]], ci: int, rounds: int + ) -> Tuple[Dict[str, Any], Dict[str, float]]: + """Aggregate the greedy holdout into SEAM ``mean@1`` metrics: overall + per-view acc, + the frozen-baseline acc, and their lift -- all single-sample-per-problem means (no + candidate averaging, no pass@k), so acc is directly comparable to SEAM's + ``val-core/math/acc/mean@1`` (correctness only; format/leak not gated).""" + def acc(rs: List[Dict[str, Any]]) -> float: + return sum(1 for x in rs if x['withskill_correct']) / len(rs) if rs else 0.0 + def bacc(rs: List[Dict[str, Any]]) -> float: + return sum(x['baseline_pass'] for x in rs) / len(rs) if rs else 0.0 + A = [x for x in recs if x['view'] == 'A'] + B = [x for x in recs if x['view'] == 'B'] + ws, base = acc(recs), bacc(recs) + summary = { + 'record_type': 'eval_summary', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, + 'n': len(recs), 'n_A': len(A), 'n_B': len(B), + 'acc_mean1': ws, 'baseline_acc_mean1': base, 'lift_mean1': ws - base, + 'acc_A_mean1': acc(A), 'acc_B_mean1': acc(B), + 'format_mean1': (sum(1 for x in recs if x['skill_parseable']) / len(recs)) if recs else 0.0, + 'term_mean1': (sum(1 for x in recs if x['withskill_terminated']) / len(recs)) if recs else 0.0, + } + metrics = { + 'core/math/acc/mean@1': ws, + 'core/math/baseline_acc/mean@1': base, + 'core/math/lift/mean@1': ws - base, + 'core/math/format/mean@1': summary['format_mean1'], + 'core/math/term/mean@1': summary['term_mean1'], + } + if A: + metrics['core/math/acc_A/mean@1'] = summary['acc_A_mean1'] + if B: + metrics['core/math/acc_B/mean@1'] = summary['acc_B_mean1'] + return summary, metrics + + +def _run_greedy_eval(base_sampler, skill_sampler, + eval_records: List[Dict[str, Any]], eval_cache: Dict[str, Dict[str, Any]], + ci: int, rounds: int, base_dp: int, skill_dp: int, + args: argparse.Namespace, checker=None + ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], Dict[str, float]]: + """SEAM ``val-core/math/acc/mean@1`` analogue on the fixed holdout: ONE greedy skill per + problem (T=0) injected into ONE greedy base solve (T=0), so acc is a single-sample + pass@1 per problem averaged over problems. Each problem keeps its assigned view; view A + still gets the rubric process-check, view B stays query-only -- the mixed A/B acc is the + deployment number. No leak filter: like SEAM's val, acc scores correctness alone.""" + _baseline_rollout(base_sampler, eval_records, base_dp, args, eval_cache) # frozen greedy baseline + for r in eval_records: + r['_view'] = _assign_view(r['problem'], args) + r['_rubric_diag'] = '' + _diagnose_views(checker, eval_records, args) + sg_out = _run_samples(skill_sampler, [_view_prompt(r, args) for r in eval_records], + 1, args.skill_max_tokens, skill_dp, temperature=0.0) + skills = [] + for seqs in sg_out: + resp = _clean_text(getattr(seqs[0], 'decoded', '') or '') if seqs else '' + skills.append((_extract_skills_block(resp) or '', resp)) + ws_out = _run_samples(base_sampler, + [build_skill_solve_prompt(r['problem'], sk) for r, (sk, _) in zip(eval_records, skills)], + 1, args.max_tokens, base_dp, temperature=0.0) + recs = [] + for r, (sk, sresp), seqs in zip(eval_records, skills, ws_out): + roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + recs.append({ + 'record_type': 'eval_problem', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'view': r['_view'], 'rubric_diag': r.get('_rubric_diag', ''), + 'baseline_pass': r['_baseline_pass'], + 'skill': sk, 'skill_parseable': bool(sk), 'skill_response': sresp, + 'withskill_pred': roll['pred'], 'withskill_correct': roll['correct'], + 'withskill_terminated': roll['terminated'], 'withskill_stop_reason': roll['stop_reason'], + 'withskill_text': roll['text'], + }) + summary, metrics = _greedy_eval_metrics(recs, ci, rounds) + return recs, summary, metrics + + +def _validate_run_config(args: argparse.Namespace, records: List[Dict[str, Any]]) -> None: + """Fail fast on configs that would SILENTLY hang the online sampler: _ProblemPool.draw(k) + dedups within a call, so it never returns unless the pool holds >= chunk_size problems; + a zero draw budget or chunk size yields empty chunks that never advance ``rounds``.""" + if not records: + raise ValueError(f'loaded 0 {args.dataset} problems; check the dataset source') + if args.chunk_size < 1: + raise ValueError(f'--chunk-size must be >= 1 (got {args.chunk_size})') + if args.eval_size < 0: + raise ValueError(f'--eval-size must be >= 0 (got {args.eval_size})') + if args.eval_size > 0 and args.eval_every < 1: + raise ValueError(f'--eval-every must be >= 1 when eval is enabled (got {args.eval_every})') + if len(records) < args.chunk_size: + raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded problems ' + f'({len(records)}); raise --n or lower --chunk-size') + if args.balance_max_draws_mult < 1: + raise ValueError(f'--balance-max-draws-mult must be >= 1 (got {args.balance_max_draws_mult})') + if not 0.0 <= args.balance_success_frac <= 1.0: + raise ValueError(f'--balance-success-frac must be in [0, 1] (got {args.balance_success_frac})') + if not 0.0 <= args.balance_loop_frac <= 1.0: + raise ValueError(f'--balance-loop-frac must be in [0, 1] (got {args.balance_loop_frac})') + + def main() -> None: args = _build_args() if args.sft_batch_size % TRAIN_DP != 0: @@ -909,10 +1337,12 @@ def main() -> None: # LR schedule is sized by an upper bound on optimizer steps (one pass over each # trained chunk's candidates); the exact count varies, cosine just decays slower. steps_per_round = max(1, (args.chunk_size * args.n_skills) // args.sft_batch_size) - records = load_math(n=args.n, seed=args.seed) + records, eval_records, data_stats = _load_records(args) + _validate_run_config(args, records) os.makedirs(args.output_dir, exist_ok=True) data_path = os.path.join(args.output_dir, 'skill_dataset.jsonl') gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') + eval_path = os.path.join(args.output_dir, 'eval_records.jsonl') train_log_path = os.path.join(args.output_dir, 'train_log.jsonl') if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY')): @@ -923,8 +1353,14 @@ def main() -> None: if use_swan: swanlab.init(project=args.swanlab_project, experiment_name=(args.swanlab_exp or None), - config={'model': GEN_MODEL_ID, 'n': len(records), + config={'model': GEN_MODEL_ID, 'dataset': args.dataset, + 'n': len(records), 'eval_n': len(eval_records), + 'raw_loaded': data_stats['raw_loaded'], + 'numeric_only': args.numeric_only, + 'numeric_dropped': data_stats['numeric_dropped'], 'n_skills': args.n_skills, 'view_b_frac': args.view_b_frac, + 'balance': args.balance, + 'balance_success_frac': args.balance_success_frac, 'skill_gen_temp': args.skill_gen_temperature, 'grpo_epsilon': args.grpo_epsilon, 'lr': args.lr}) @@ -948,7 +1384,7 @@ def main() -> None: skill_model.set_template(Template, model_id=GEN_MODEL_ID, enable_thinking=True, max_length=args.max_model_len, truncation_strategy='delete') - skill_model.set_processor(InputProcessor, padding_free=True) + skill_model.set_processor(InputProcessor, padding_free=False) skill_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) skill_model.set_optimizer('AdamW', lr=args.lr) skill_model.set_lr_scheduler('CosineWarmupScheduler', num_warmup_steps=10, @@ -981,16 +1417,24 @@ def main() -> None: sys.stderr.write('[rft] no LLM backup env -> view-A rubric process-check DISABLED ' '(skill-gen diagnoses from the attempt alone)\n') - sys.stderr.write(f'[rft] {len(records)} MATH problems; train_gpus={TRAIN_GPUS} ' - f'skill_dp={skill_dp} base_dp={base_dp}\n') + sys.stderr.write(f'[rft] raw={data_stats["raw_loaded"]} numeric_drop={data_stats["numeric_dropped"]} ' + f'train={len(records)} eval={len(eval_records)} {args.dataset} problems; ' + f'train_gpus={TRAIN_GPUS} skill_dp={skill_dp} base_dp={base_dp}\n') # -- Sequential: generate one chunk, train on it, sync -> exact on-policy GRPO. # Generation dominates wall-clock, so not overlapping training costs little, and # it removes all producer/consumer concurrency (no thread, no lock). -- - n_chunks = (len(records) + args.chunk_size - 1) // args.chunk_size - cfg = {'record_type': 'config', 'model': GEN_MODEL_ID, 'dataset': 'math', - 'n': len(records), 'seed': args.seed, 'n_skills': args.n_skills, + cfg = {'record_type': 'config', 'model': GEN_MODEL_ID, 'dataset': args.dataset, + 'n': len(records), 'eval_n': len(eval_records), 'seed': args.seed, + 'numeric_only': args.numeric_only, + 'raw_loaded': data_stats['raw_loaded'], + 'numeric_dropped': data_stats['numeric_dropped'], + 'eval_every': args.eval_every, + 'n_skills': args.n_skills, 'view_b_frac': args.view_b_frac, 'skill_retries': args.skill_retries, + 'balance': args.balance, 'balance_success_frac': args.balance_success_frac, + 'balance_loop_frac': args.balance_loop_frac, + 'balance_max_draws_mult': args.balance_max_draws_mult, 'skill_gen_temp': args.skill_gen_temperature, 'skill_gen_top_p': args.skill_gen_top_p, 'skill_gen_top_k': args.skill_gen_top_k, 'reward': 'greedy_binary(correct)', 'advantage': 'group_relative', @@ -1000,76 +1444,100 @@ def main() -> None: 'max_train_rounds': args.max_train_rounds, 'started': int(time.time())} hist: List[Dict[str, float]] = [] rounds = 0 + pool = _ProblemPool(records, args.seed) + eval_cache: Dict[str, Dict[str, Any]] = {} with open(gen_path, 'w', encoding='utf-8') as gen_f, \ + open(eval_path, 'w', encoding='utf-8') as eval_f, \ open(data_path, 'w', encoding='utf-8') as data_f, \ open(train_log_path, 'w', encoding='utf-8') as tlog: - for f in (gen_f, data_f, tlog): + for f in (gen_f, eval_f, data_f, tlog): f.write(json.dumps(cfg, ensure_ascii=False) + '\n') f.flush() - gstep, epoch = 0, 0 - # Multiple epochs over the SAME problem set: each pass reshuffles and RE-GENERATES - # rollouts with the current (improved) policy, so every chunk stays on-policy (no - # importance correction needed) -- the online analogue of SEAM's fixed-data epochs. - while rounds < args.max_train_rounds and n_chunks > 0: - if epoch > 0: - np.random.RandomState(args.seed + epoch).shuffle(records) - for ci in range(n_chunks): - if rounds >= args.max_train_rounds: - break - chunk = records[ci * args.chunk_size:(ci + 1) * args.chunk_size] - full, summary, groups = process_chunk( - base_sampler, skill_sampler, leak, chunk, gstep, base_dp, skill_dp, - args, checker) - - log = None - if groups: # on-policy GRPO update on this chunk, then weights sync - log = _train_chunk(skill_model, ckpt, groups, args) - rounds += 1 - log.update({'record_type': 'train_round', 'round': rounds, - 'chunk': gstep, 'epoch': epoch, 'ts': int(time.time())}) - tlog.write(json.dumps(log, ensure_ascii=False) + '\n') - tlog.flush() - if rounds % args.save_rounds == 0: - skill_model.save(f'skill-rft-{rounds}', output_dir=args.output_dir) - - summary['rounds_done'], summary['epoch'] = rounds, epoch - for rec in full: - gen_f.write(json.dumps(rec, ensure_ascii=False) + '\n') - gen_f.write(json.dumps(summary, ensure_ascii=False) + '\n') - gen_f.flush() - for v in groups: - data_f.write(json.dumps(v, ensure_ascii=False) + '\n') - data_f.flush() - - sa, sb = summary['view_A'], summary['view_B'] - hist.append({'aA': sa['adoption_rate'], 'aB': sb['adoption_rate'], - 'lift': summary['avg_lift'], 'pos': summary['n_reward_pos']}) - sys.stderr.write( - f'[gen] e{epoch} chunk {ci+1}/{n_chunks} (g{gstep}): hard={summary["n_hard"]} ' - f'clean={summary["n_clean"]} train={summary["n_train_samples"]} ' - f'acc={summary["avg_baseline_pass_on_hard"]:.2f}->{summary["avg_withskill_pass"]:.2f} ' - f'lift={summary["avg_lift"]:+.3f} ' - f'A[{sa["n_hard"]}h {sa["adoption_rate"]:.2f}] ' - f'B[{sb["n_hard"]}h {sb["adoption_rate"]:.2f}] ' - f'rounds={rounds}' - + (f' metric={log.get("metric")}' if log else '') + '\n') - # -- per-query passk (base vs best/avg of N skills) + swanlab metrics -- - rows = _query_rows(full) - for base_p, best_p, avg_p, nsc, prob in rows: - logger.info(f'[q] g{gstep} base={base_p:.2f} bestN={best_p:.2f} avgN={avg_p:.2f} ' - f'n={nsc} | {prob[:70].replace(chr(10), " ")}') + gstep = 0 + # Each chunk is drawn fresh from the pool (which reshuffles + bumps epoch on every + # full pass) and RE-GENERATED with the current (improved) policy, so every chunk + # stays on-policy (no importance correction) -- the online analogue of SEAM's + # fixed-data epochs. With --balance, _draw_chunk keeps drawing+baselining until the + # base fail:success mix hits the target before this chunk is trained on. + while rounds < args.max_train_rounds: + chunk, balance = _draw_chunk(pool, base_sampler, base_dp, args) + full, summary, groups = process_chunk( + base_sampler, skill_sampler, leak, chunk, gstep, base_dp, skill_dp, + args, checker) + summary['balance'] = balance + + log = None + if groups: # on-policy GRPO update on this chunk, then weights sync + log = _train_chunk(skill_model, ckpt, groups, args) + rounds += 1 + log.update({'record_type': 'train_round', 'round': rounds, + 'chunk': gstep, 'epoch': pool.epoch, 'ts': int(time.time())}) + tlog.write(json.dumps(log, ensure_ascii=False) + '\n') + tlog.flush() + if rounds % args.save_rounds == 0: + skill_model.save(f'skill-rft-{rounds}', output_dir=args.output_dir) + + summary['rounds_done'], summary['epoch'] = rounds, pool.epoch + for rec in full: + gen_f.write(json.dumps(rec, ensure_ascii=False) + '\n') + gen_f.write(json.dumps(summary, ensure_ascii=False) + '\n') + gen_f.flush() + for v in groups: + data_f.write(json.dumps(v, ensure_ascii=False) + '\n') + data_f.flush() + + sa, sb = summary['view_A'], summary['view_B'] + hist.append({'aA': sa['adoption_rate'], 'aB': sb['adoption_rate'], + 'lift': summary['avg_lift'], 'pos': summary['n_reward_pos']}) + bal_str = (f'bal {balance["selected_fail"]}f/{balance["selected_success"]}s ' + f'(loop {balance["selected_fail_loop"]} drew {balance["n_drawn"]}/' + f'fresh {balance["n_baseline_fresh"]}' + + ('!' if balance.get('budget_hit') else '') + ') ' + ) if balance.get('enabled') else '' + sys.stderr.write( + f'[gen] e{pool.epoch} g{gstep}: {bal_str}hard={summary["n_hard"]} ' + f'clean={summary["n_clean"]} train={summary["n_train_samples"]} ' + f'(fail {summary["n_train_from_fail"]} adv%{summary["abs_adv_from_fail_frac"]:.2f}) ' + f'acc={summary["avg_baseline_pass_on_hard"]:.2f}->{summary["avg_withskill_pass"]:.2f} ' + f'lift={summary["avg_lift"]:+.3f} ' + f'A[{sa["n_hard"]}h {sa["adoption_rate"]:.2f}] ' + f'B[{sb["n_hard"]}h {sb["adoption_rate"]:.2f}] ' + f'rounds={rounds}' + + (f' metric={log.get("metric")}' if log else '') + '\n') + # -- per-query passk (base vs best/avg of N skills) + swanlab metrics -- + rows = _query_rows(full) + for base_p, best_p, avg_p, nsc, prob in rows: + logger.info(f'[q] g{gstep} base={base_p:.2f} bestN={best_p:.2f} avgN={avg_p:.2f} ' + f'n={nsc} | {prob[:70].replace(chr(10), " ")}') + if use_swan: + swanlab.log(_swan_metrics(summary, log, rows), step=gstep) + + if eval_records and (gstep + 1) % args.eval_every == 0: + eval_recs, eval_summary, eval_metrics = _run_greedy_eval( + base_sampler, skill_sampler, eval_records, eval_cache, gstep, + rounds, base_dp, skill_dp, args, checker) + for rec in eval_recs: + eval_f.write(json.dumps(rec, ensure_ascii=False) + '\n') + eval_f.write(json.dumps(eval_summary, ensure_ascii=False) + '\n') + eval_f.flush() if use_swan: - swanlab.log(_swan_metrics(summary, log, rows), step=gstep) - - if (gstep + 1) % args.trend_every == 0: - tl = _trend_line(hist, args.trend_every, rounds) - if tl: - sys.stderr.write(tl + '\n') - gstep += 1 - epoch += 1 + swanlab.log(_prefix_metrics(eval_metrics, 'eval'), step=gstep) + sys.stderr.write( + f'[eval] g{gstep}: n={eval_summary["n"]} mean@1 ' + f'acc={eval_summary["baseline_acc_mean1"]:.3f}->{eval_summary["acc_mean1"]:.3f} ' + f'lift={eval_summary["lift_mean1"]:+.3f} ' + f'A[{eval_summary["n_A"]} {eval_summary["acc_A_mean1"]:.3f}] ' + f'B[{eval_summary["n_B"]} {eval_summary["acc_B_mean1"]:.3f}] ' + f'fmt={eval_summary["format_mean1"]:.2f} rounds={rounds}\n') + + if (gstep + 1) % args.trend_every == 0: + tl = _trend_line(hist, args.trend_every, rounds) + if tl: + sys.stderr.write(tl + '\n') + gstep += 1 skill_model.save('skill-rft-final', output_dir=args.output_dir) - sys.stderr.write(f'[rft] done: {rounds} train rounds over {gstep} chunks / {epoch} epochs; ' + sys.stderr.write(f'[rft] done: {rounds} train rounds over {gstep} chunks / {pool.epoch} epochs; ' f'data -> {data_path}\n') diff --git a/cookbook/exp/embedding/train_reflexion_skill_rft.sh b/cookbook/exp/embedding/train_reflexion_skill_rft.sh index 524570a35..b5e0f1d82 100644 --- a/cookbook/exp/embedding/train_reflexion_skill_rft.sh +++ b/cookbook/exp/embedding/train_reflexion_skill_rft.sh @@ -14,11 +14,16 @@ export LLM_BACKUP_BASE_URL=${LLM_BACKUP_BASE_URL:-https://dashscope.aliyuncs.com export LLM_BACKUP_MODEL=${LLM_BACKUP_MODEL:-qwen3.7-max} python cookbook/exp/embedding/train_reflexion_skill_rft.py \ + --dataset aops \ --n 5000 \ --chunk-size 16 \ --n-skills 8 \ --view-b-frac 0.5 \ --skill-retries 2 \ + --balance \ + --balance-success-frac 0.4 \ + --balance-loop-frac 0.5 \ + --balance-max-draws-mult 8 \ --max-tokens 25000 \ --max-model-len 30000 \ --sft-batch-size 8 \ From 5c6bb750135624b5333cd742f6231d5e7c6a32ed Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Wed, 15 Jul 2026 12:47:12 +0800 Subject: [PATCH 09/60] fix --- .../embedding/build_reflexion_skill_data.py | 152 ++++++++++++++++++ .../embedding/train_reflexion_skill_replay.py | 113 +++++++++++++ .../embedding/train_reflexion_skill_rft.py | 80 +++++---- 3 files changed, 316 insertions(+), 29 deletions(-) create mode 100644 cookbook/exp/embedding/build_reflexion_skill_data.py create mode 100644 cookbook/exp/embedding/train_reflexion_skill_replay.py diff --git a/cookbook/exp/embedding/build_reflexion_skill_data.py b/cookbook/exp/embedding/build_reflexion_skill_data.py new file mode 100644 index 000000000..8ff6c7421 --- /dev/null +++ b/cookbook/exp/embedding/build_reflexion_skill_data.py @@ -0,0 +1,152 @@ +"""Build exact reflexion skill RFT data without updating the skill model. + +The output ``skill_dataset.jsonl`` uses the same schema as the online trainer's +per-chunk training records, while ``gen_records.jsonl`` keeps the full trace for +inspection. This lets expensive rollout/scoring be run once, then reused for +offline replay experiments. + +Launch: + python cookbook/exp/embedding/build_reflexion_skill_data.py --chunks 100 +""" +import argparse +import json +import os +import sys +import time +from typing import Any, Dict, List + +import train_reflexion_skill_rft as rft +from twinkle_agentic.verifier import LeakVerifier + + +def _build_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument('--dataset', choices=('aops', 'math'), default='aops') + p.add_argument('--n', type=int, default=2000) + p.add_argument('--seed', type=int, default=42) + p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--eval-size', type=int, default=128) + p.add_argument('--chunks', type=int, default=100) + p.add_argument('--chunk-size', type=int, default=16) + p.add_argument('--balance', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--balance-success-frac', type=float, default=0.4) + p.add_argument('--balance-loop-frac', type=float, default=0.5) + p.add_argument('--balance-max-draws-mult', type=int, default=8) + p.add_argument('--n-skills', type=int, default=8) + p.add_argument('--view-b-frac', type=float, default=0.5) + p.add_argument('--skill-retries', type=int, default=2) + p.add_argument('--skill-gen-temperature', type=float, default=1.0) + p.add_argument('--skill-gen-top-p', type=float, default=1.0) + p.add_argument('--skill-gen-top-k', type=int, default=-1) + p.add_argument('--max-model-len', type=int, default=16384) + p.add_argument('--max-tokens', type=int, default=8192) + p.add_argument('--skill-max-tokens', type=int, default=8192) + p.add_argument('--leak-workers', type=int, default=16) + p.add_argument('--rubric-workers', type=int, default=16) + p.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--output-dir', default='./output/reflexion_skill_data') + p.add_argument('--overwrite', action='store_true') + return p.parse_args() + + +def _init_samplers(args: argparse.Namespace): + skill_dp = rft.SKILL_SAMPLER_GPUS + base_dp = rft.BASE_SAMPLER_GPUS + total_gpus = skill_dp + base_dp + if total_gpus <= 0: + raise ValueError('SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS must be positive') + device_groups = [ + rft.DeviceGroup(name='skill_sampler', ranks=list(range(0, skill_dp)), device_type='GPU'), + rft.DeviceGroup(name='base_sampler', ranks=list(range(skill_dp, total_gpus)), device_type='GPU'), + ] + rft.twinkle.initialize(mode='ray', nproc_per_node=total_gpus, groups=device_groups, + lazy_collect=False) + skill_sampler = rft.vLLMSampler( + model_id=rft.GEN_MODEL_ID, + engine_args={'gpu_memory_utilization': rft.GEN_GPU_MEM, + 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, + device_mesh=rft.DeviceMesh.from_sizes(world_size=skill_dp, dp_size=skill_dp), + remote_group='skill_sampler') + skill_sampler.set_template(rft.Template, model_id=rft.GEN_MODEL_ID, + enable_thinking=True, max_length=args.max_model_len) + base_sampler = rft.vLLMSampler( + model_id=rft.GEN_MODEL_ID, + engine_args={'gpu_memory_utilization': rft.GEN_GPU_MEM, + 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, + device_mesh=rft.DeviceMesh.from_sizes(world_size=base_dp, dp_size=base_dp), + remote_group='base_sampler') + base_sampler.set_template(rft.Template, model_id=rft.GEN_MODEL_ID, + enable_thinking=True, max_length=args.max_model_len) + return base_sampler, skill_sampler, base_dp, skill_dp + + +def _write_jsonl_row(handle, row: Dict[str, Any]) -> None: + handle.write(json.dumps(row, ensure_ascii=False) + '\n') + + +def main() -> None: + args = _build_args() + records, eval_records, data_stats = rft._load_records(args) + rft._validate_run_config(args, records) + os.makedirs(args.output_dir, exist_ok=True) + data_path = os.path.join(args.output_dir, 'skill_dataset.jsonl') + gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') + eval_path = os.path.join(args.output_dir, 'eval_holdout.jsonl') + for path in (data_path, gen_path, eval_path): + if os.path.exists(path) and not args.overwrite: + raise FileExistsError(f'{path} exists; pass --overwrite to replace it') + + if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY')): + sys.stderr.write('[build-rft-data] WARNING: no LLM backup env; leak/rubric checks degrade\n') + + base_sampler, skill_sampler, base_dp, skill_dp = _init_samplers(args) + leak = LeakVerifier(sampler=None, answer_only=True) + checker = rft._build_rubric_checker() + pool = rft._ProblemPool(records, args.seed) + rubric_cache: Dict[str, str] = {} + + cfg = { + 'record_type': 'config', 'mode': 'offline_data_build', 'model': rft.GEN_MODEL_ID, + 'dataset': args.dataset, 'n': len(records), 'eval_n': len(eval_records), + 'seed': args.seed, 'numeric_only': args.numeric_only, + 'raw_loaded': data_stats['raw_loaded'], 'numeric_dropped': data_stats['numeric_dropped'], + 'chunks': args.chunks, 'chunk_size': args.chunk_size, 'n_skills': args.n_skills, + 'view_b_frac': args.view_b_frac, 'balance': args.balance, + 'balance_success_frac': args.balance_success_frac, + 'reward': 'greedy_binary(correct)', 'advantage': 'group_relative', + 'format_in_reward': args.format_in_reward, 'started': int(time.time()), + } + total_groups = 0 + with open(gen_path, 'w', encoding='utf-8') as gen_f, \ + open(data_path, 'w', encoding='utf-8') as data_f, \ + open(eval_path, 'w', encoding='utf-8') as eval_f: + for handle in (gen_f, data_f, eval_f): + _write_jsonl_row(handle, cfg) + for rec in eval_records: + _write_jsonl_row(eval_f, {'record_type': 'eval_holdout', **rec}) + eval_f.flush() + + for ci in range(args.chunks): + chunk, balance = rft._draw_chunk(pool, base_sampler, base_dp, args) + full, summary, groups = rft.process_chunk( + base_sampler, skill_sampler, leak, chunk, ci, base_dp, skill_dp, + args, checker, rubric_cache) + summary['balance'] = balance + for rec in full: + _write_jsonl_row(gen_f, rec) + _write_jsonl_row(gen_f, summary) + gen_f.flush() + for row in groups: + _write_jsonl_row(data_f, {'chunk': ci, **row}) + data_f.flush() + total_groups += len(groups) + sys.stderr.write( + f'[build-rft-data] g{ci}: train={len(groups)} total={total_groups} ' + f'acc={summary["avg_baseline_pass_on_hard"]:.2f}->{summary["avg_withskill_pass"]:.2f} ' + f'lift={summary["avg_lift"]:+.3f}\n') + + sys.stderr.write(f'[build-rft-data] done: train records -> {data_path}; trace -> {gen_path}\n') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/embedding/train_reflexion_skill_replay.py b/cookbook/exp/embedding/train_reflexion_skill_replay.py new file mode 100644 index 000000000..5429e45ce --- /dev/null +++ b/cookbook/exp/embedding/train_reflexion_skill_replay.py @@ -0,0 +1,113 @@ +"""Replay-train the reflexion skill model from prebuilt exact RFT data. + +Use ``build_reflexion_skill_data.py`` first to create ``skill_dataset.jsonl``. This +script trains only the skill model from those frozen records; it does not run vLLM +rollouts, leak checks, or rubric diagnosis. + +Launch: + python cookbook/exp/embedding/train_reflexion_skill_replay.py \ + --data ./output/reflexion_skill_data/skill_dataset.jsonl +""" +import argparse +import json +import os +import sys +from collections import defaultdict +from typing import Any, Dict, List + +import train_reflexion_skill_rft as rft + + +def _build_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument('--data', default='./output/reflexion_skill_data/skill_dataset.jsonl') + p.add_argument('--output-dir', default='./output/reflexion_skill_replay') + p.add_argument('--epochs', type=int, default=1) + p.add_argument('--sft-batch-size', type=int, default=8) + p.add_argument('--max-model-len', type=int, default=16384) + p.add_argument('--grpo-epsilon', type=float, default=0.2) + p.add_argument('--lr', type=float, default=1e-5) + p.add_argument('--save-rounds', type=int, default=50) + return p.parse_args() + + +def _load_chunks(path: str) -> List[List[Dict[str, Any]]]: + chunks: Dict[int, List[Dict[str, Any]]] = defaultdict(list) + fallback_chunk = 0 + with open(path, 'r', encoding='utf-8') as f: + for line_no, line in enumerate(f, 1): + if not line.strip(): + continue + row = json.loads(line) + if row.get('record_type') == 'config': + continue + for key in ('problem', 'response', 'advantage'): + if key not in row: + raise ValueError(f'{path}:{line_no} missing required field {key!r}') + ci = int(row.get('chunk', fallback_chunk)) + chunks[ci].append(row) + if 'chunk' not in row and len(chunks[ci]) >= 64: + fallback_chunk += 1 + return [chunks[k] for k in sorted(chunks) if chunks[k]] + + +def _init_model(args: argparse.Namespace, total_updates: int): + train_mesh = rft.DeviceMesh.from_sizes( + world_size=rft.TRAIN_GPUS, dp_size=rft.TRAIN_DP, fsdp_size=rft.TRAIN_FSDP) + device_groups = [ + rft.DeviceGroup(name='train', ranks=list(range(rft.TRAIN_GPUS)), device_type='GPU'), + ] + rft.twinkle.initialize(mode='ray', nproc_per_node=rft.TRAIN_GPUS, groups=device_groups, + lazy_collect=False) + model = rft.TransformersModel(model_id=rft.GEN_MODEL_ID, device_mesh=train_mesh, + remote_group='train', ddp_config={'find_unused_parameters': False}) + from twinkle.patch.no_split_modules import NoSplitModulesPatch + model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) + model.set_template(rft.Template, model_id=rft.GEN_MODEL_ID, + enable_thinking=True, max_length=args.max_model_len, + truncation_strategy='delete') + model.set_processor(rft.InputProcessor, padding_free=False) + model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) + model.set_optimizer('AdamW', lr=args.lr) + model.set_lr_scheduler('CosineWarmupScheduler', num_warmup_steps=10, + num_training_steps=max(1, total_updates)) + return model + + +def main() -> None: + args = _build_args() + if args.sft_batch_size % rft.TRAIN_DP != 0: + raise ValueError(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple ' + f'of the training dp size ({rft.TRAIN_DP})') + chunks = _load_chunks(args.data) + if not chunks: + raise ValueError(f'no train records found in {args.data}') + os.makedirs(args.output_dir, exist_ok=True) + total_updates = len(chunks) * args.epochs + model = _init_model(args, total_updates) + log_path = os.path.join(args.output_dir, 'train_log.jsonl') + cfg = {'record_type': 'config', 'mode': 'offline_replay', 'data': args.data, + 'chunks': len(chunks), 'epochs': args.epochs, 'lr': args.lr, + 'sft_batch_size': args.sft_batch_size} + rounds = 0 + with open(log_path, 'w', encoding='utf-8') as tlog: + tlog.write(json.dumps(cfg, ensure_ascii=False) + '\n') + for epoch in range(args.epochs): + for ci, samples in enumerate(chunks): + log = rft._train_chunk(model, None, samples, args) + rounds += 1 + log.update({'record_type': 'train_round', 'round': rounds, + 'epoch': epoch, 'chunk': ci}) + tlog.write(json.dumps(log, ensure_ascii=False) + '\n') + tlog.flush() + sys.stderr.write( + f'[replay-rft] e{epoch} c{ci}: n={log["n_samples"]} ' + f'micro={log["n_micro_batches"]} metric={log.get("metric")}\n') + if rounds % args.save_rounds == 0: + model.save(f'skill-rft-replay-{rounds}', output_dir=args.output_dir) + model.save('skill-rft-replay-final', output_dir=args.output_dir) + sys.stderr.write(f'[replay-rft] done: {rounds} updates; log -> {log_path}\n') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/embedding/train_reflexion_skill_rft.py b/cookbook/exp/embedding/train_reflexion_skill_rft.py index 12852c74f..39a676c05 100644 --- a/cookbook/exp/embedding/train_reflexion_skill_rft.py +++ b/cookbook/exp/embedding/train_reflexion_skill_rft.py @@ -476,7 +476,8 @@ def _format_diagnosis(detail) -> str: return '\n'.join(lines) -def _diagnose_views(checker, hard: List[Dict[str, Any]], args: argparse.Namespace) -> None: +def _diagnose_views(checker, hard: List[Dict[str, Any]], args: argparse.Namespace, + diag_cache: Optional[Dict[str, str]] = None) -> None: """Rubric-check every view-A problem's greedy attempt in parallel, stashing the formatted findings on ``r['_rubric_diag']`` (view B stays empty). A checker error or empty result degrades to no diagnosis (the plain view-A prompt).""" @@ -485,21 +486,38 @@ def _diagnose_views(checker, hard: List[Dict[str, Any]], args: argparse.Namespac if not checker or not targets: return - def _run(r: Dict[str, Any]) -> Tuple[Dict[str, Any], str]: + def _cache_key(r: Dict[str, Any]) -> str: + init_text = r.get('_init', [{}])[0].get('text', '') + return hashlib.md5(f'{r["problem"]}\n{init_text}'.encode('utf-8')).hexdigest() + + pending = [] + for r in targets: + key = _cache_key(r) + if diag_cache is not None and key in diag_cache: + r['_rubric_diag'] = diag_cache[key] + else: + pending.append((r, key)) + if not pending: + return + + def _run(item: Tuple[Dict[str, Any], str]) -> Tuple[Dict[str, Any], str, str, bool]: + r, key = item seg = {'messages': [ {'role': 'user', 'content': r['problem']}, {'role': 'assistant', 'content': r['_init'][0]['text']}, ]} try: - return r, _format_diagnosis(checker.diagnose(seg, query=r['problem'])) + return r, key, _format_diagnosis(checker.diagnose(seg, query=r['problem'])), True except Exception as exc: # teacher hiccup -> fall back to no-diagnosis prompt logger.warning(f'[rubric] diagnose error: {exc}') - return r, '' + return r, key, '', False - workers = max(1, min(args.rubric_workers, len(targets))) + workers = max(1, min(args.rubric_workers, len(pending))) with ThreadPoolExecutor(max_workers=workers) as ex: - for r, diag in ex.map(_run, targets): + for r, key, diag, ok in ex.map(_run, pending): r['_rubric_diag'] = diag + if ok and diag_cache is not None: + diag_cache[key] = diag # --------------------------------------------------------------------------- @@ -877,7 +895,8 @@ def _draw_chunk(pool: _ProblemPool, base_sampler, base_dp: int, args: argparse.N def process_chunk(base_sampler, skill_sampler, leak: LeakVerifier, chunk: List[Dict[str, Any]], ci: int, base_dp: int, skill_dp: int, - args: argparse.Namespace, checker=None + args: argparse.Namespace, checker=None, + diag_cache: Optional[Dict[str, str]] = None ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], List[Dict[str, Any]]]: """base-solve -> rubric-check (view A) -> skill-gen -> leak-filter -> with-skill pass -> GRPO advantages, for one chunk. @@ -902,7 +921,7 @@ def process_chunk(base_sampler, skill_sampler, leak: LeakVerifier, for r in hard: r['_view'] = _assign_view(r['problem'], args) r['_rubric_diag'] = '' - _diagnose_views(checker, hard, args) + _diagnose_views(checker, hard, args, diag_cache) # --- Phase 3: skill-gen (thinking ON), per-view prompt; re-sample empties. --- flat: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] @@ -1000,17 +1019,15 @@ def _train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: return {'messages': full, 'user_data': {'key_rounds': [len(msgs)]}} -def _train_chunk(skill_model, ckpt: CheckpointEngineManager, +def _train_chunk(skill_model, ckpt: Optional[CheckpointEngineManager], samples: List[Dict[str, Any]], args: argparse.Namespace) -> Dict[str, Any]: - """One on-policy GRPO update on THIS chunk's skill candidates, then sync weights. + """One on-policy GRPO optimizer update on THIS chunk's skill candidates, then sync weights. Sequential design (generate-one-chunk-train-one): the skills were sampled from the current policy and trained immediately, so ``old_logps`` is omitted and the GRPO - ratio is ~1 (no importance correction needed). Each per-sample ``advantage`` was set - in _assign_advantages (group-relative + optional SFT blend). Driver-side mini-batches - each take an optimizer step (matching short_math_grpo); the batch is padded to a - multiple of ``sft_batch_size`` (dp needs each mini-batch divisible) with advantage-0 - copies that contribute zero gradient. ``sync_weights`` needs no lock (no overlap). + ratio is ~1. All driver-side mini-batches accumulate into one optimizer step so the + whole rollout chunk stays under the same pre-update policy. The batch is padded to a + multiple of ``sft_batch_size`` with advantage-0 copies that contribute zero gradient. """ trajs = [_train_trajectory(rec) for rec in samples] advs = [float(rec['advantage']) for rec in samples] @@ -1018,15 +1035,16 @@ def _train_chunk(skill_model, ckpt: CheckpointEngineManager, if rem: trajs += [trajs[-1]] * rem # zero-advantage pads -> forward only, no gradient advs += [0.0] * rem - steps = 0 + micro_batches = 0 for i in range(0, len(trajs), args.sft_batch_size): skill_model.forward_backward(inputs=trajs[i:i + args.sft_batch_size], advantages=advs[i:i + args.sft_batch_size]) - skill_model.clip_grad_and_step() - steps += 1 - ckpt.sync_weights(merge_and_sync=True) + micro_batches += 1 + skill_model.clip_grad_and_step() + if ckpt is not None: + ckpt.sync_weights(merge_and_sync=True) metric = skill_model.calculate_metric(is_training=True) - return {'n_samples': len(samples), 'n_steps': steps, + return {'n_samples': len(samples), 'n_steps': 1, 'n_micro_batches': micro_batches, 'advantages': [float(rec['advantage']) for rec in samples], 'metric': {k: (float(v) if _is_num(v) else v) for k, v in (metric or {}).items()}} @@ -1116,8 +1134,8 @@ def _build_args() -> argparse.Namespace: '(teacher-served; requires LLM_BACKUP_* env).') # -- online GRPO (one on-policy update per generated chunk) -- p.add_argument('--sft-batch-size', type=int, default=8, - help='Driver-side mini-batch per optimizer step; MUST be a multiple ' - 'of the training dp size (sliced across dp ranks).') + help='Driver-side micro-batch size before the chunk-level optimizer step; ' + 'MUST be a multiple of the training dp size (sliced across dp ranks).') p.add_argument('--grpo-epsilon', type=float, default=0.2, help='PPO clip epsilon for GRPOLoss (ratio~1 on-policy, so rarely binds).') p.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True, @@ -1222,6 +1240,8 @@ def _swan_metrics(summary: Dict[str, Any], log: Optional[Dict[str, Any]], d.update({'passk/baseline_mean': m(0), 'passk/bestN_mean': m(1), 'passk/avgN_mean': m(2)}) if log: d['train/n_steps'] = log['n_steps'] + if 'n_micro_batches' in log: + d['train/n_micro_batches'] = log['n_micro_batches'] d.update({f'train/{k}': v for k, v in _clean_metric(log.get('metric')).items()}) return d @@ -1268,7 +1288,8 @@ def bacc(rs: List[Dict[str, Any]]) -> float: def _run_greedy_eval(base_sampler, skill_sampler, eval_records: List[Dict[str, Any]], eval_cache: Dict[str, Dict[str, Any]], ci: int, rounds: int, base_dp: int, skill_dp: int, - args: argparse.Namespace, checker=None + args: argparse.Namespace, checker=None, + diag_cache: Optional[Dict[str, str]] = None ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], Dict[str, float]]: """SEAM ``val-core/math/acc/mean@1`` analogue on the fixed holdout: ONE greedy skill per problem (T=0) injected into ONE greedy base solve (T=0), so acc is a single-sample @@ -1279,7 +1300,7 @@ def _run_greedy_eval(base_sampler, skill_sampler, for r in eval_records: r['_view'] = _assign_view(r['problem'], args) r['_rubric_diag'] = '' - _diagnose_views(checker, eval_records, args) + _diagnose_views(checker, eval_records, args, diag_cache) sg_out = _run_samples(skill_sampler, [_view_prompt(r, args) for r in eval_records], 1, args.skill_max_tokens, skill_dp, temperature=0.0) skills = [] @@ -1334,9 +1355,8 @@ def main() -> None: if args.sft_batch_size % TRAIN_DP != 0: raise ValueError(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple ' f'of the training dp size ({TRAIN_DP})') - # LR schedule is sized by an upper bound on optimizer steps (one pass over each - # trained chunk's candidates); the exact count varies, cosine just decays slower. - steps_per_round = max(1, (args.chunk_size * args.n_skills) // args.sft_batch_size) + # LR schedule now follows chunk-level optimizer updates, not driver micro-batches. + steps_per_round = 1 records, eval_records, data_stats = _load_records(args) _validate_run_config(args, records) os.makedirs(args.output_dir, exist_ok=True) @@ -1446,6 +1466,8 @@ def main() -> None: rounds = 0 pool = _ProblemPool(records, args.seed) eval_cache: Dict[str, Dict[str, Any]] = {} + rubric_cache: Dict[str, str] = {} + eval_rubric_cache: Dict[str, str] = {} with open(gen_path, 'w', encoding='utf-8') as gen_f, \ open(eval_path, 'w', encoding='utf-8') as eval_f, \ open(data_path, 'w', encoding='utf-8') as data_f, \ @@ -1463,7 +1485,7 @@ def main() -> None: chunk, balance = _draw_chunk(pool, base_sampler, base_dp, args) full, summary, groups = process_chunk( base_sampler, skill_sampler, leak, chunk, gstep, base_dp, skill_dp, - args, checker) + args, checker, rubric_cache) summary['balance'] = balance log = None @@ -1515,7 +1537,7 @@ def main() -> None: if eval_records and (gstep + 1) % args.eval_every == 0: eval_recs, eval_summary, eval_metrics = _run_greedy_eval( base_sampler, skill_sampler, eval_records, eval_cache, gstep, - rounds, base_dp, skill_dp, args, checker) + rounds, base_dp, skill_dp, args, checker, eval_rubric_cache) for rec in eval_recs: eval_f.write(json.dumps(rec, ensure_ascii=False) + '\n') eval_f.write(json.dumps(eval_summary, ensure_ascii=False) + '\n') From d54539febc653ab3fb08a7351cc62329e245351c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=A8=E6=B3=93?= Date: Wed, 15 Jul 2026 14:55:35 +0800 Subject: [PATCH 10/60] fix --- .../embedding/build_reflexion_skill_data.py | 147 +++++++++++------- cookbook/exp/embedding/eval_gpqa_rag.py | 2 +- .../embedding/train_reflexion_skill_replay.py | 5 +- .../embedding/train_reflexion_skill_rft.py | 19 +-- 4 files changed, 108 insertions(+), 65 deletions(-) diff --git a/cookbook/exp/embedding/build_reflexion_skill_data.py b/cookbook/exp/embedding/build_reflexion_skill_data.py index 8ff6c7421..eecba0830 100644 --- a/cookbook/exp/embedding/build_reflexion_skill_data.py +++ b/cookbook/exp/embedding/build_reflexion_skill_data.py @@ -6,78 +6,106 @@ offline replay experiments. Launch: - python cookbook/exp/embedding/build_reflexion_skill_data.py --chunks 100 + python cookbook/exp/embedding/build_reflexion_skill_data.py \ + --total-problems 3200 --base-success-frac 0.3 """ import argparse import json +import math import os import sys import time -from typing import Any, Dict, List +from typing import Any, Dict, List, Tuple + +import numpy as np + +from cookbook.exp.embedding.eval_gpqa_rag import load_aops, load_math +from twinkle.sampler import vLLMSampler + +import twinkle +from twinkle import DeviceGroup, DeviceMesh import train_reflexion_skill_rft as rft from twinkle_agentic.verifier import LeakVerifier def _build_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description=__doc__) - p.add_argument('--dataset', choices=('aops', 'math'), default='aops') - p.add_argument('--n', type=int, default=2000) - p.add_argument('--seed', type=int, default=42) - p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) - p.add_argument('--eval-size', type=int, default=128) - p.add_argument('--chunks', type=int, default=100) - p.add_argument('--chunk-size', type=int, default=16) - p.add_argument('--balance', action=argparse.BooleanOptionalAction, default=True) - p.add_argument('--balance-success-frac', type=float, default=0.4) - p.add_argument('--balance-loop-frac', type=float, default=0.5) - p.add_argument('--balance-max-draws-mult', type=int, default=8) - p.add_argument('--n-skills', type=int, default=8) - p.add_argument('--view-b-frac', type=float, default=0.5) - p.add_argument('--skill-retries', type=int, default=2) - p.add_argument('--skill-gen-temperature', type=float, default=1.0) - p.add_argument('--skill-gen-top-p', type=float, default=1.0) - p.add_argument('--skill-gen-top-k', type=int, default=-1) - p.add_argument('--max-model-len', type=int, default=16384) - p.add_argument('--max-tokens', type=int, default=8192) - p.add_argument('--skill-max-tokens', type=int, default=8192) - p.add_argument('--leak-workers', type=int, default=16) - p.add_argument('--rubric-workers', type=int, default=16) - p.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True) + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter) + p.add_argument('--total-problems', type=int, default=3200, + help='Final number of problems selected into generated chunks.') + p.add_argument('--base-success-frac', type=float, default=0.3, + help='Target fraction of selected problems solved by the frozen base.') p.add_argument('--output-dir', default='./output/reflexion_skill_data') p.add_argument('--overwrite', action='store_true') - return p.parse_args() + p.add_argument('--seed', type=int, default=42) + + advanced = p.add_argument_group('advanced knobs, usually leave unchanged') + advanced.add_argument('--dataset', choices=('aops', 'math'), default='aops') + advanced.add_argument('--n', type=int, default=0, + help='Raw train-pool size; 0 derives it from --total-problems.') + advanced.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) + advanced.add_argument('--eval-size', type=int, default=128) + advanced.add_argument('--chunk-size', type=int, default=16) + advanced.add_argument('--balance', action=argparse.BooleanOptionalAction, default=True) + advanced.add_argument('--balance-loop-frac', type=float, default=0.5) + advanced.add_argument('--balance-max-draws-mult', type=int, default=8) + advanced.add_argument('--n-skills', type=int, default=8) + advanced.add_argument('--view-b-frac', type=float, default=0.5) + advanced.add_argument('--skill-retries', type=int, default=2) + advanced.add_argument('--skill-gen-temperature', type=float, default=1.0) + advanced.add_argument('--skill-gen-top-p', type=float, default=1.0) + advanced.add_argument('--skill-gen-top-k', type=int, default=-1) + advanced.add_argument('--max-model-len', type=int, default=16384) + advanced.add_argument('--max-tokens', type=int, default=8192) + advanced.add_argument('--skill-max-tokens', type=int, default=8192) + advanced.add_argument('--leak-workers', type=int, default=16) + advanced.add_argument('--rubric-workers', type=int, default=16) + advanced.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True) + args = p.parse_args() + _resolve_args(args) + return args + + +def _resolve_args(args: argparse.Namespace) -> None: + if args.total_problems <= 0: + raise ValueError('--total-problems must be positive') + if args.chunk_size <= 0: + raise ValueError('--chunk-size must be positive') + if not 0.0 <= args.base_success_frac <= 1.0: + raise ValueError('--base-success-frac must be in [0, 1]') + args.chunks = math.ceil(args.total_problems / args.chunk_size) + args.balance_success_frac = args.base_success_frac + if args.n <= 0: + args.n = max(args.total_problems + args.eval_size, + math.ceil(args.total_problems * 1.5)) def _init_samplers(args: argparse.Namespace): - skill_dp = rft.SKILL_SAMPLER_GPUS - base_dp = rft.BASE_SAMPLER_GPUS - total_gpus = skill_dp + base_dp - if total_gpus <= 0: - raise ValueError('SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS must be positive') + model = 'ms://Qwen/Qwen3-4B' device_groups = [ - rft.DeviceGroup(name='skill_sampler', ranks=list(range(0, skill_dp)), device_type='GPU'), - rft.DeviceGroup(name='base_sampler', ranks=list(range(skill_dp, total_gpus)), device_type='GPU'), + DeviceGroup(name='skill_sampler', ranks=list(range(0, 4)), device_type='GPU'), + DeviceGroup(name='base_sampler', ranks=list(range(4, 8)), device_type='GPU'), ] - rft.twinkle.initialize(mode='ray', nproc_per_node=total_gpus, groups=device_groups, + twinkle.initialize(mode='ray', nproc_per_node=8, groups=device_groups, lazy_collect=False) - skill_sampler = rft.vLLMSampler( - model_id=rft.GEN_MODEL_ID, - engine_args={'gpu_memory_utilization': rft.GEN_GPU_MEM, + skill_sampler = vLLMSampler( + model_id='Qwen/Qwen3-4B', + engine_args={'gpu_memory_utilization': 0.8, 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, - device_mesh=rft.DeviceMesh.from_sizes(world_size=skill_dp, dp_size=skill_dp), + device_mesh=DeviceMesh.from_sizes(world_size=4, dp_size=4), remote_group='skill_sampler') - skill_sampler.set_template(rft.Template, model_id=rft.GEN_MODEL_ID, + skill_sampler.set_template('Template', model_id=model, enable_thinking=True, max_length=args.max_model_len) - base_sampler = rft.vLLMSampler( - model_id=rft.GEN_MODEL_ID, - engine_args={'gpu_memory_utilization': rft.GEN_GPU_MEM, + base_sampler = vLLMSampler( + model_id=model, + engine_args={'gpu_memory_utilization': 0.8, 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, - device_mesh=rft.DeviceMesh.from_sizes(world_size=base_dp, dp_size=base_dp), + device_mesh=DeviceMesh.from_sizes(world_size=4, dp_size=4), remote_group='base_sampler') - base_sampler.set_template(rft.Template, model_id=rft.GEN_MODEL_ID, + base_sampler.set_template('Template', model_id=model, enable_thinking=True, max_length=args.max_model_len) - return base_sampler, skill_sampler, base_dp, skill_dp + return base_sampler, skill_sampler, 4, 4 def _write_jsonl_row(handle, row: Dict[str, Any]) -> None: @@ -106,12 +134,14 @@ def main() -> None: rubric_cache: Dict[str, str] = {} cfg = { - 'record_type': 'config', 'mode': 'offline_data_build', 'model': rft.GEN_MODEL_ID, + 'record_type': 'config', 'mode': 'offline_data_build', 'model': 'Qwen/Qwen3-4B', 'dataset': args.dataset, 'n': len(records), 'eval_n': len(eval_records), - 'seed': args.seed, 'numeric_only': args.numeric_only, + 'total_problems': args.total_problems, 'seed': args.seed, + 'numeric_only': args.numeric_only, 'raw_loaded': data_stats['raw_loaded'], 'numeric_dropped': data_stats['numeric_dropped'], 'chunks': args.chunks, 'chunk_size': args.chunk_size, 'n_skills': args.n_skills, 'view_b_frac': args.view_b_frac, 'balance': args.balance, + 'base_success_frac': args.base_success_frac, 'balance_success_frac': args.balance_success_frac, 'reward': 'greedy_binary(correct)', 'advantage': 'group_relative', 'format_in_reward': args.format_in_reward, 'started': int(time.time()), @@ -126,11 +156,20 @@ def main() -> None: _write_jsonl_row(eval_f, {'record_type': 'eval_holdout', **rec}) eval_f.flush() + selected = 0 for ci in range(args.chunks): - chunk, balance = rft._draw_chunk(pool, base_sampler, base_dp, args) - full, summary, groups = rft.process_chunk( - base_sampler, skill_sampler, leak, chunk, ci, base_dp, skill_dp, - args, checker, rubric_cache) + remaining = args.total_problems - selected + if remaining <= 0: + break + original_chunk_size = args.chunk_size + args.chunk_size = min(original_chunk_size, remaining) + try: + chunk, balance = rft._draw_chunk(pool, base_sampler, base_dp, args) + full, summary, groups = rft.process_chunk( + base_sampler, skill_sampler, leak, chunk, ci, base_dp, skill_dp, + args, checker, rubric_cache) + finally: + args.chunk_size = original_chunk_size summary['balance'] = balance for rec in full: _write_jsonl_row(gen_f, rec) @@ -140,8 +179,10 @@ def main() -> None: _write_jsonl_row(data_f, {'chunk': ci, **row}) data_f.flush() total_groups += len(groups) + selected += len(chunk) sys.stderr.write( - f'[build-rft-data] g{ci}: train={len(groups)} total={total_groups} ' + f'[build-rft-data] g{ci}: problems={selected}/{args.total_problems} ' + f'train={len(groups)} total={total_groups} ' f'acc={summary["avg_baseline_pass_on_hard"]:.2f}->{summary["avg_withskill_pass"]:.2f} ' f'lift={summary["avg_lift"]:+.3f}\n') diff --git a/cookbook/exp/embedding/eval_gpqa_rag.py b/cookbook/exp/embedding/eval_gpqa_rag.py index af5cf60dc..7954e5fd2 100644 --- a/cookbook/exp/embedding/eval_gpqa_rag.py +++ b/cookbook/exp/embedding/eval_gpqa_rag.py @@ -853,7 +853,7 @@ def build_direct_prompt(problem: str) -> Dict[str, Any]: return { 'messages': [ {'role': 'system', 'content': DIRECT_SYSTEM}, - {'role': 'user', 'content': problem + MCQ_INSTRUCTION}, + {'role': 'user', 'content': problem}, ] } diff --git a/cookbook/exp/embedding/train_reflexion_skill_replay.py b/cookbook/exp/embedding/train_reflexion_skill_replay.py index 5429e45ce..0580a6ba4 100644 --- a/cookbook/exp/embedding/train_reflexion_skill_replay.py +++ b/cookbook/exp/embedding/train_reflexion_skill_replay.py @@ -52,6 +52,7 @@ def _load_chunks(path: str) -> List[List[Dict[str, Any]]]: def _init_model(args: argparse.Namespace, total_updates: int): + model = 'ms://Qwen/Qwen3-4B' train_mesh = rft.DeviceMesh.from_sizes( world_size=rft.TRAIN_GPUS, dp_size=rft.TRAIN_DP, fsdp_size=rft.TRAIN_FSDP) device_groups = [ @@ -59,11 +60,11 @@ def _init_model(args: argparse.Namespace, total_updates: int): ] rft.twinkle.initialize(mode='ray', nproc_per_node=rft.TRAIN_GPUS, groups=device_groups, lazy_collect=False) - model = rft.TransformersModel(model_id=rft.GEN_MODEL_ID, device_mesh=train_mesh, + model = rft.TransformersModel(model_id=model, device_mesh=train_mesh, remote_group='train', ddp_config={'find_unused_parameters': False}) from twinkle.patch.no_split_modules import NoSplitModulesPatch model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) - model.set_template(rft.Template, model_id=rft.GEN_MODEL_ID, + model.set_template(rft.Template, model_id=model, enable_thinking=True, max_length=args.max_model_len, truncation_strategy='delete') model.set_processor(rft.InputProcessor, padding_free=False) diff --git a/cookbook/exp/embedding/train_reflexion_skill_rft.py b/cookbook/exp/embedding/train_reflexion_skill_rft.py index 39a676c05..b4388b773 100644 --- a/cookbook/exp/embedding/train_reflexion_skill_rft.py +++ b/cookbook/exp/embedding/train_reflexion_skill_rft.py @@ -140,8 +140,8 @@ 'training scaffolding: study it together with the problem, identify the ' 'problem-visible features that make each useful flagged failure relevant, then ' 'rephrase those lessons as self-contained reusable skills. The goal is not to ' - 'continue from the check, cite it, or hide it silently; the goal is to turn it into ' - 'a problem-triggered reasoning pattern a query-only solver could reproduce later.\n\n' + 'continue from the check, cite it, or hide it silently; the goal is to turn it to ' + 'a skill pattern which prevents the model falls into similar pitfalls in the future.\n\n' 'Good skills name the observable trigger, the method worth reaching for, the ' 'pitfall to watch, and a quick verification habit. Prefer formulations like ' '"When a configuration has ...", "Before setting up ...", or "Check whether ..." ' @@ -181,7 +181,8 @@ 'Now output a self-contained skills bullet list. Each bullet should still be useful ' 'if the process check were removed: connect any useful flagged failure to ' 'problem-visible features, general methods, and quick checks rather than citing the ' - 'rubric or the earlier attempt.' + 'rubric or the earlier attempt. \n\n' + 'Note: **Do not solve the problem, only generate skills**. Now Begin:' ) @@ -193,8 +194,8 @@ def build_skillgen_prompt(problem: str, diagnosis: str) -> Dict[str, Any]: attempt. The one-shot demo is query-only; only the real turn carries the diagnosis.""" return {'messages': [ {'role': 'system', 'content': SKILL_GEN_SYSTEM}, - {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=_EX_PROBLEM)}, - {'role': 'assistant', 'content': _EX_SKILLS}, + # {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=_EX_PROBLEM)}, + # {'role': 'assistant', 'content': _EX_SKILLS}, {'role': 'user', 'content': SKILL_GEN_USER_RUBRIC.format(problem=problem, diagnosis=diagnosis)}, ]} @@ -267,8 +268,8 @@ def build_querygen_prompt(problem: str) -> Dict[str, Any]: attempt) — matching what is available at deployment (query only).""" return {'messages': [ {'role': 'system', 'content': SKILL_GEN_SYSTEM_Q}, - {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=_EX_PROBLEM)}, - {'role': 'assistant', 'content': _EX_SKILLS}, + # {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=_EX_PROBLEM)}, + # {'role': 'assistant', 'content': _EX_SKILLS}, {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=problem)}, ]} @@ -406,14 +407,14 @@ def _extract_skills_block(text: str) -> Optional[str]: - For PASS items, leave "fix" as "". - For FAIL items, "reason", "fix", and "summary" must describe only the flawed step, theorem, arithmetic operation, case split, or verification habit. -- NEVER state the correct final answer, corrected final expression, option letter, +- NEVER try to solve the query or state the correct final answer, corrected final expression, option letter, graph/choice label, or any exact value that the answer should become. - NEVER write phrases like "the correct answer is", "which gives", "yielding", "should be ", "Option ", or "Graph ". - If a fix would require naming a corrected value, replace it with a method-level instruction such as "redo that computation carefully" or "apply the theorem with the correct quantities". -- Keep every "reason" and "fix" clear and concise — one short sentence each. +- Keep every "reason" and "fix" clear and concise. - "overall" is "OK" only if NO criterion is FAIL. - Output only the JSON object.""" From d474ee4ffcb03a7c496adc522450a832f4a2eb2e Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Wed, 15 Jul 2026 17:20:37 +0800 Subject: [PATCH 11/60] fix --- .../embedding/build_reflexion_skill_data.py | 1142 ++++++++++++-- .../exp/embedding/train_reflexion_skill.py | 1390 +++++++++++++++++ .../exp/embedding/train_reflexion_skill.sh | 54 + src/twinkle_agentic/verifier/leak_verifier.py | 2 +- 4 files changed, 2481 insertions(+), 107 deletions(-) create mode 100644 cookbook/exp/embedding/train_reflexion_skill.py create mode 100755 cookbook/exp/embedding/train_reflexion_skill.sh diff --git a/cookbook/exp/embedding/build_reflexion_skill_data.py b/cookbook/exp/embedding/build_reflexion_skill_data.py index eecba0830..59fe2382c 100644 --- a/cookbook/exp/embedding/build_reflexion_skill_data.py +++ b/cookbook/exp/embedding/build_reflexion_skill_data.py @@ -1,121 +1,1048 @@ -"""Build exact reflexion skill RFT data without updating the skill model. +"""Offline builder for reflexion skill RFT data (self-contained, cached). -The output ``skill_dataset.jsonl`` uses the same schema as the online trainer's -per-chunk training records, while ``gen_records.jsonl`` keeps the full trace for -inspection. This lets expensive rollout/scoring be run once, then reused for -offline replay experiments. +Runs the SAME pipeline as the online trainer -- base greedy solve -> rubric +process-check (view A) -> skill-gen -> leak filter -> with-skill greedy pass -> +group-relative GRPO advantage -- but never updates the skill model. It emits +``skill_dataset.jsonl`` (trainer-schema training records), ``gen_records.jsonl`` +(full per-problem traces) and ``eval_holdout.jsonl`` (the fixed holdout). + +The expensive base rollouts and rubric diagnoses are cached to disk (one jsonl +each, keyed by an md5 of their inputs) so a re-run skips them entirely. + +8 GPUs: ranks 0-3 skill_sampler (vLLM tp1 dp4), ranks 4-7 base_sampler. Leak / +rubric use the backup teacher API: set LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / +LLM_BACKUP_MODEL. Launch: - python cookbook/exp/embedding/build_reflexion_skill_data.py \ + LLM_BACKUP_API_KEY=... python cookbook/exp/embedding/build_reflexion_skill_data.py \ --total-problems 3200 --base-success-frac 0.3 """ import argparse +import copy +import hashlib import json import math import os +import re import sys import time -from typing import Any, Dict, List, Tuple +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Optional, Tuple import numpy as np -from cookbook.exp.embedding.eval_gpqa_rag import load_aops, load_math +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import SamplingParams +from twinkle.dataset import Dataset, DatasetMeta from twinkle.sampler import vLLMSampler +from twinkle_agentic.verifier import LeakVerifier, RubricVerifier +from twinkle_agentic.verifier.rubric_verifier import RubricItem -import twinkle -from twinkle import DeviceGroup, DeviceMesh +logger = get_logger() + +MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') +GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.8)) +GEN_TEMPERATURE = float(os.environ.get('GEN_TEMPERATURE', 0.6)) +GEN_TOP_P = float(os.environ.get('GEN_TOP_P', 0.95)) +AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') +MATH_DATASET_ID = os.environ.get('MATH_DATASET_ID', 'modelscope/competition_math') + + +# =========================================================================== +# Block A -- boxed extraction + answer grading +# =========================================================================== +_BOXED_RE = re.compile(r'\\boxed\s*\{') + + +def extract_boxed(text: str) -> Optional[str]: + """Last ``\\boxed{...}`` content, brace-balanced.""" + if not text: + return None + last = None + for m in _BOXED_RE.finditer(text): + depth, i = 1, m.end() + while i < len(text) and depth > 0: + depth += (text[i] == '{') - (text[i] == '}') + i += 1 + if depth == 0: + last = text[m.end():i - 1].strip() + return last + + +def normalize_answer(ans: str) -> str: + if not ans: + return '' + s = str(ans).strip() + m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) + if m: + return m.group(1) + s = s.strip('$').strip().replace('−', '-') + 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').replace(r'\tfrac', r'\frac') + s = s.replace(r'\,', '').replace(r'\;', '').replace(r'\!', '') + s = re.sub(r'\\(?:quad|qquad|\s)', '', s) + s = re.sub(r'\s+', '', s) + s = re.sub(r'\\sqrt([0-9A-Za-z])', r'\\sqrt{\1}', s) + s = re.sub(r'\\frac\{([^{}]+)\}([^{}\\])', r'\\frac{\1}{\2}', s) + s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) + s = re.sub(r'\{[a-zA-Z]+\}$', '', s) + s = re.sub(r'\^\\circ|\^\{\\circ\}|°|\\circ', 'deg', s) + s = s.replace(r'\minus{}', '-').replace(r'\minus', '-') + + def _frac_to_slash(mt): + text = mt.group(0) + pos = text.index('{') + 1 + depth, num_start = 1, pos + while depth > 0: + depth += (text[pos] == '{') - (text[pos] == '}') + pos += 1 + numer = text[num_start:pos - 1] + pos += 1 + den_start, depth = pos, 1 + while depth > 0: + depth += (text[pos] == '{') - (text[pos] == '}') + pos += 1 + return f'({numer})/({text[den_start:pos - 1]})' + + s = re.sub(r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', _frac_to_slash, s) + s = re.sub(r'(? bool: + try: + va, vb = float(a.replace('(', '').replace(')', '')), 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) + return va is not None and vb is not None and abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) + + +_MCQ_REF_RE = re.compile( + r'^\\(?:textbf|text|mathrm|mathbf)\{\(?([A-E])\)?\s*\}\s*(.+)$' + r'|^\(?([A-E])\)\s+(.+)$') +_VAR_PREFIX_RE = re.compile(r'^(?:[A-Za-z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)$', re.DOTALL) + + +def _split_mcq(ans: str): + s = ans.strip() + m = _MCQ_REF_RE.match(s) + if m: + return (m.group(1) or m.group(3)), ((m.group(2) or m.group(4) or '').strip() or None) + bl = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) + return (bl.group(1), None) if bl else (None, s or None) + + +def _strip_var_prefix(ans: str) -> str: + m = _VAR_PREFIX_RE.match((ans or '').strip()) + return m.group(1).strip() if m else (ans or '') + + +def _normalize_tuple(ans: str) -> str: + return re.sub(r'[\s()\[\]{}\\]', '', ans or '') + + +def _math_verify_equal(predicted: str, reference: str) -> bool: + try: + from math_verify import parse, verify + gold = parse(r'\boxed{' + reference + '}', parsing_timeout=1) + pred = parse(r'\boxed{' + predicted + '}', parsing_timeout=1) + return bool(gold and pred and verify(gold, pred, timeout_seconds=1)) + except Exception: + return False + + +def answers_match(predicted: str, reference: str) -> bool: + if not predicted or not reference: + return False + norm_p, norm_r = normalize_answer(predicted), normalize_answer(reference) + if norm_p == norm_r or norm_p.lower() == norm_r.lower() or _try_numeric_equal(norm_p, norm_r): + return True + stripped_p = normalize_answer(_strip_var_prefix(predicted)) + stripped_r = normalize_answer(_strip_var_prefix(reference)) + if stripped_p and stripped_r: + if (stripped_p == stripped_r or stripped_p.lower() == stripped_r.lower() + or _try_numeric_equal(stripped_p, stripped_r)): + return True + p_letter, p_value = _split_mcq(predicted) + r_letter, r_value = _split_mcq(reference) + if p_letter and r_letter and p_letter == r_letter: + return True + if (p_letter and p_value is None) and (r_letter and r_value is None): + return False + 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 and (p_val == r_val or p_val.lower() == r_val.lower() or _try_numeric_equal(p_val, r_val)): + return True + for left, right in ((norm_p, norm_r), (stripped_p, stripped_r)): + tl, tr = _normalize_tuple(left), _normalize_tuple(right) + if ',' in tl and tl == tr: + return True + if '=' in norm_r and '=' not in norm_p: + if any(part and (part == norm_p or _try_numeric_equal(part, norm_p)) for part in norm_r.split('=')): + return True + if '=' in norm_p and '=' not in norm_r: + if any(part and (part == norm_r or _try_numeric_equal(part, norm_r)) for part in norm_p.split('=')): + return True + return _math_verify_equal(stripped_p or norm_p, stripped_r or norm_r) + + +# =========================================================================== +# Block B -- prompts, skill parsing, batched sampling +# =========================================================================== +DIRECT_SYSTEM = ( + 'You are an expert competition mathematician. Solve the following problem ' + 'step by step. Provide your final answer inside \\boxed{}.') + +# Appended to solve turns: box BOTH the letter and value of an MCQ so the model +# never loops deciding which form to box. +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.') + +_SKILL_SOLVE_PREFIX = ( + DIRECT_SYSTEM + '\n\n' + 'Before you start, keep these reminders in mind to avoid common mistakes on this ' + 'type of problem:\n') +_SKILL_SOLVE_SUFFIX = '\nApply them where relevant, but rely on your own reasoning to reach the answer.' + + +def build_direct_prompt(problem: str) -> Dict[str, Any]: + return {'messages': [{'role': 'system', 'content': DIRECT_SYSTEM}, + {'role': 'user', 'content': problem}]} + + +def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: + # Concatenation (not .format): DIRECT_SYSTEM/skill contain literal braces. + return {'messages': [ + {'role': 'system', 'content': _SKILL_SOLVE_PREFIX + skill + _SKILL_SOLVE_SUFFIX}, + {'role': 'user', 'content': problem + MCQ_INSTRUCTION}]} + + +# -- skill-gen prompts (view A: problem + rubric findings; view B: query only) -- +SKILL_GEN_SYSTEM = ( + 'You are a mathematics coach. You are shown a competition problem together with an ' + 'automated process-check of an earlier solver attempt at it -- which solution ' + 'criteria the attempt passed or failed, and suggested fixes for the failures. You ' + 'do NOT see the attempt itself, only this check. Treat the check as privileged ' + 'training scaffolding: study it together with the problem, identify the ' + 'problem-visible features that make each useful flagged failure relevant, then ' + 'rephrase those lessons as self-contained reusable skills. The goal is not to ' + 'continue from the check, cite it, or hide it silently; the goal is to turn it into ' + 'a problem-triggered reasoning pattern a query-only solver could reproduce later.\n\n' + 'Good skills name the observable trigger, the method worth reaching for, the ' + 'pitfall to watch, and a quick verification habit. Prefer formulations like ' + '"When a configuration has ...", "Before setting up ...", or "Check whether ..." ' + 'over references to the process-check, failed criteria, or the earlier attempt. ' + "These tips are advisory: they will be placed in a solver's system prompt as gentle " + 'reminders before it works through a SIMILAR problem on its own, without seeing ' + 'this process-check. So keep them general and transferable rather than a worked ' + 'solution to this exact problem, and do not state its specific intermediate values ' + 'or final answer. Think briefly first, then give your tips as a markdown bullet ' + 'list wrapped in and , like the example below.') + +SKILL_GEN_SYSTEM_Q = ( + 'You are a mathematics coach. You are shown ONE competition problem and nothing ' + 'else — no solution and no attempt. Think about what approach this KIND of problem ' + 'calls for and where solvers tend to slip, then distil a few reusable tips.\n\n' + "These tips are advisory: they will be placed in a solver's system prompt as gentle " + 'reminders before it works through a SIMILAR problem on its own. So keep them ' + 'general and transferable — the method worth reaching for, the pitfall to watch and ' + 'a quick check, and the discipline to settle on a final answer — rather than a ' + 'worked solution to this exact problem, and without stating its specific ' + 'intermediate values or its final answer. Think briefly first, then give your tips ' + 'as a markdown bullet list wrapped in and , like the example below.') + +SKILL_GEN_USER_Q = ( + 'Problem:\n{problem}\n\n' + 'Now reason about this TYPE of problem, then output the skills bullet list.') + +SKILL_GEN_USER_RUBRIC = ( + 'Problem:\n{problem}\n\n' + 'Process check of an earlier attempt (automated rubric verifier -- treat as ' + 'evidence, not gospel; PASS/FAIL per criterion with suggested fixes for failures):\n' + '{diagnosis}\n\n' + 'Now output a self-contained skills bullet list. Each bullet should still be useful ' + 'if the process check were removed: connect any useful flagged failure to ' + 'problem-visible features, general methods, and quick checks rather than citing the ' + 'rubric or the earlier attempt.') + +_EX_PROBLEM = 'Simplify $\\sqrt{72} + \\sqrt{18}$ and give the result.' +_EX_SKILLS = ( + '\n' + '- Rewrite each square root by factoring its radicand into a perfect square times ' + 'a remainder, then move the perfect-square factor outside.\n' + '- Avoid the classic trap $\\sqrt{a}+\\sqrt{b}\\ne\\sqrt{a+b}$; only combine terms ' + 'sharing the same simplest radical, and sanity-check by estimating each root.\n' + '- Procedure: simplify every radical, group like radical terms, add their ' + 'coefficients, then reduce to simplest form.\n' + '- Once the expression is in simplest form, commit to that single result as the ' + 'final answer rather than re-checking indefinitely.\n' + '') + + +def _skillgen_messages(problem: str, view: str, diagnosis: str) -> List[Dict[str, Any]]: + """Single source of truth for the skill-gen prompt. View A with a localisable + failure uses problem + rubric findings; view B -- or a view-A problem whose rubric + flagged nothing (no ``[FAIL]``) -- degrades to the query-only prompt.""" + if view == 'B' or '[FAIL]' not in (diagnosis or ''): + return [{'role': 'system', 'content': SKILL_GEN_SYSTEM_Q}, + {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=_EX_PROBLEM)}, + {'role': 'assistant', 'content': _EX_SKILLS}, + {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=problem)}] + return [{'role': 'system', 'content': SKILL_GEN_SYSTEM}, + {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=_EX_PROBLEM)}, + {'role': 'assistant', 'content': _EX_SKILLS}, + {'role': 'user', 'content': SKILL_GEN_USER_RUBRIC.format(problem=problem, diagnosis=diagnosis)}] + + +def _assign_view(problem: str, args: argparse.Namespace) -> str: + h = int(hashlib.md5(f'{args.seed}:{problem}'.encode('utf-8')).hexdigest(), 16) + return 'B' if (h % 100000) / 100000.0 < args.view_b_frac else 'A' + + +def _view_prompt(r: Dict[str, Any]) -> Dict[str, Any]: + return {'messages': _skillgen_messages(r['problem'], r['_view'], r.get('_rubric_diag', ''))} + + +_BULLET_RE = re.compile(r'(?m)^\s*(?:[-*]|\d+[.)])\s') +_META_RE = re.compile( + r'\b(the student|the solver|the attempt|this attempt|the trace|the response|' + r'in the (?:attempt|trace|response|solution)|as (?:shown|seen|noted) above|' + r'the (?:above|previous|earlier)|my (?:reasoning|analysis)|i (?:think|need|will))\b', + re.IGNORECASE) +_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') + + +def _clean_text(decoded: Optional[str]) -> str: + return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() + + +def _is_clean_block(block: str) -> bool: + """Pure bullet list (every non-empty line a bullet) with no meta/trajectory ref.""" + lines = [ln for ln in (l.strip() for l in block.splitlines()) if ln] + if not lines or not all(_BULLET_RE.match(ln) for ln in lines): + return False + return _META_RE.search(block) is None + + +def _extract_skills_block(text: str) -> Optional[str]: + """Clean ``...`` block, or None. Requires ```` (skill-gen + runs thinking ON); reads only the answer after the last one, so a mid-reasoning draft + or a demo echo can never be mistaken for the answer.""" + low = text.lower() + end_think = low.rfind('') + if end_think < 0: + return None + answer = text[end_think + len(''):] + low_a = answer.lower() + s = low_a.find('') + if s < 0: + return None + inner = s + len('') + e = low_a.find('', inner) + block = (answer[inner:e] if e >= 0 else answer[inner:]).strip() + block = re.sub(r'', '', block, flags=re.IGNORECASE).strip() + return block if _is_clean_block(block) else None + + +def _parse_seq(seq, gold: str) -> Dict[str, Any]: + """Grade one sampled sequence into a rollout record.""" + text = _clean_text(getattr(seq, 'decoded', '') or '') + pred = extract_boxed(text) + correct = bool(pred) and answers_match(pred, gold) + terminated = getattr(seq, 'stop_reason', None) != 'length' + return {'pred': pred, 'correct': correct, 'terminated': terminated, + 'passed': bool(correct and terminated), + 'stop_reason': getattr(seq, 'stop_reason', None), + 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} + + +def _run_samples(sampler, prompts: List[Any], num_samples: int, max_tokens: int, + gen_dp: int, temperature: Optional[float] = None, + top_p: Optional[float] = None, top_k: Optional[int] = None) -> List[List[Any]]: + """One batched sampler call -> per-prompt list of raw sequences. vLLM dp needs + batch len >= dp, so pad the tail and slice back.""" + if not prompts: + return [] + params = SamplingParams( + max_tokens=max_tokens, + temperature=GEN_TEMPERATURE if temperature is None else temperature, + top_p=GEN_TOP_P if top_p is None else top_p, + num_samples=num_samples, **({} if top_k is None else {'top_k': top_k})) + padded = prompts + if gen_dp > 1 and 0 < len(prompts) < gen_dp: + padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] + responses = sampler.sample(padded, params)[:len(prompts)] + return [list(r.sequences) if (r and r.sequences) else [] for r in responses] + + +# =========================================================================== +# Block C -- data loading via twinkle.Dataset + numeric filtering +# =========================================================================== +def load_problems(dataset: str, n: int, seed: int) -> List[Dict[str, Any]]: + """Load boxed-answer problems as ``{problem, reference_answer, level?}`` via + twinkle.Dataset (ModelScope hub), sampled to ``n`` (0 = all).""" + ds_id = AOPS_DATASET_ID if dataset == 'aops' else MATH_DATASET_ID + rows = Dataset(DatasetMeta(dataset_id=f'ms://{ds_id}', split='train')).dataset + out: List[Dict[str, Any]] = [] + for row in rows: + if dataset == 'aops' and not (row.get('metadata') or {}).get('boxed'): + continue + ref = extract_boxed(row.get('solution', '')) + if not ref: + continue + rec = {'problem': row['problem'], 'reference_answer': ref} + if row.get('level'): + rec['level'] = row['level'] + out.append(rec) + logger.info(f'[data] {dataset}: {len(out)} boxed problems') + rng = np.random.RandomState(seed) + rng.shuffle(out) + return out[:n] if (n and n < len(out)) else out + + +_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') + + +def _norm_num_text(num: str) -> str: + try: + f = float(num) + return str(int(f)) if f == int(f) else str(f) + except Exception: + return str(num).strip() + + +def _numeric_value(raw: Any) -> Optional[str]: + """Collapse an answer to a single int/decimal/fraction, or None.""" + if raw is None: + return None + s = str(raw).strip().strip('$').strip() + s = s.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac') + s = re.sub(r'\\!|\\,|\\;|\\ |\\left|\\right|\s', '', s) + for pat in (r'\\frac\{(-?\d+)\}\{(-?\d+)\}', r'(-?\d+)/(-?\d+)'): + m = re.fullmatch(pat, s) + if m: + a, b = int(m.group(1)), int(m.group(2)) + return _norm_num_text(str(a / b)) if b else None + return _norm_num_text(s) if _NUM_RE.fullmatch(s) else None + + +def _load_records(args: argparse.Namespace) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], Dict[str, int]]: + """Load, numeric-filter, shuffle, then split a fixed eval holdout off the front.""" + # Load all when filtering or splitting (else the eval holdout could starve train). + load_n = 0 if (args.numeric_only or args.eval_size > 0) else args.n + records = load_problems(args.dataset, load_n, args.seed) + raw_n, dropped = len(records), 0 + if args.numeric_only: + kept = [] + for r in records: + ref = _numeric_value(r.get('reference_answer')) + if ref is None: + dropped += 1 + continue + kept.append({**r, 'reference_answer': ref}) + records = kept + np.random.RandomState(args.seed).shuffle(records) + eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 + eval_records = [dict(r) for r in records[:eval_n]] + pool = records[eval_n:] + train_n = args.n if args.n > 0 else len(pool) + train_records = [dict(r) for r in pool[:train_n]] + overlap = {r['problem'] for r in train_records} & {r['problem'] for r in eval_records} + if overlap: + raise ValueError(f'eval/train overlap: {len(overlap)} problems') + stats = {'raw_loaded': raw_n, 'numeric_dropped': dropped, + 'train_records': len(train_records), 'eval_records': len(eval_records)} + return train_records, eval_records, stats + + +# =========================================================================== +# Block D -- disk cache, problem pool, baseline rollout, rubric check +# =========================================================================== +class DiskCache: + """Append-only jsonl kv cache (md5 key -> value). Loads on init, appends on put. + Disabled instances (``enabled=False``) always miss and never write.""" + + def __init__(self, path: str, enabled: bool = True): + self.path, self.enabled = path, enabled + self._mem: Dict[str, Any] = {} + self._fh = None + if not enabled: + return + if os.path.exists(path): + with open(path, encoding='utf-8') as f: + for line in f: + if line.strip(): + row = json.loads(line) + self._mem[row['key']] = row['value'] + self._fh = open(path, 'a', encoding='utf-8') + + @staticmethod + def key_for(*parts: str) -> str: + return hashlib.md5('\x1f'.join(parts).encode('utf-8')).hexdigest() + + def __contains__(self, key: str) -> bool: + return key in self._mem -import train_reflexion_skill_rft as rft -from twinkle_agentic.verifier import LeakVerifier + def get(self, key: str) -> Any: + return self._mem.get(key) + + def put(self, key: str, value: Any) -> None: + self._mem[key] = value + if self._fh is not None: + self._fh.write(json.dumps({'key': key, 'value': value}, ensure_ascii=False) + '\n') + self._fh.flush() + + def close(self) -> None: + if self._fh is not None: + self._fh.close() + + +class ProblemPool: + """Cyclic draw source. Each full pass reshuffles with ``seed + epoch``; the initial + pass keeps the loader order. ``draw(k)`` returns k distinct problems (k << pool).""" + + def __init__(self, records: List[Dict[str, Any]], seed: int): + self._records = list(records) + self._seed, self._cursor, self.epoch = seed, 0, 0 + + def draw(self, k: int) -> List[Dict[str, Any]]: + out, seen = [], set() + while len(out) < k: + if self._cursor >= len(self._records): + self.epoch += 1 + np.random.RandomState(self._seed + self.epoch).shuffle(self._records) + self._cursor = 0 + r = self._records[self._cursor] + self._cursor += 1 + if id(r) not in seen: + seen.add(id(r)) + out.append(r) + return out + + +def _empty_roll() -> Dict[str, Any]: + return {'pred': '', 'correct': False, 'terminated': False, 'passed': False, + 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} + + +def _apply_baseline(r: Dict[str, Any], roll: Dict[str, Any]) -> None: + """Attach a greedy baseline roll and reset per-chunk working state.""" + r['_baseline_rolls'], r['_cands'], r['_init'] = [roll], [], [roll] + r['_failed'] = not roll['correct'] + r['_baseline_pass'] = 1.0 if roll['correct'] else 0.0 + r['_hard'] = True # process every problem; group variance selects (SEAM-style) + + +def baseline_rollout(base_sampler, problems: List[Dict[str, Any]], base_dp: int, + args: argparse.Namespace, cache: DiskCache) -> int: + """Phase 1: base solves each problem greedily once (T=0, M=1), disk-cached by + problem text. The base is frozen + greedy so the cache is exact. Returns the number + of fresh (cache-miss) rollouts.""" + todo = [r for r in problems if DiskCache.key_for(r['problem']) not in cache] + if todo: + out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], + 1, args.max_tokens, base_dp, temperature=0.0) + for r, seqs in zip(todo, out): + roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + cache.put(DiskCache.key_for(r['problem']), roll) + for r in problems: + _apply_baseline(r, cache.get(DiskCache.key_for(r['problem']))) + return len(todo) + + +# -- rubric process-check (view A): teacher diagnoses the base's attempt -- +_RFT_DIAG_SYSTEM = """\ +You are a process error checker for a math solution attempt. You are given a math +problem, a rubric, and one attempted solution segment. Decide PASS or FAIL for each +criterion and explain only the process error type. + +Output STRICT JSON (no prose outside it) with this shape: +{ + "items": [ + {"index": 1, "verdict": "PASS", "reason": "", + "fix": ""}, + {"index": 2, "verdict": "FAIL", "reason": "", + "fix": ""} + ], + "overall": "OK" | "ISSUES", + "summary": "" +} + +Rules: +- Judge every criterion independently and literally; a [Hard Rule] is FAIL unless + unambiguously satisfied. +- Judge ONLY what is observable in THIS segment. +- Content inside ... (or ) is internal reasoning, not + user-facing output; ignore it for "output only X" style criteria. +- For PASS items, leave "fix" as "". +- For FAIL items, "reason", "fix", and "summary" must describe only the flawed + step, theorem, arithmetic operation, case split, or verification habit. +- NEVER state the correct final answer, corrected final expression, option letter, + graph/choice label, or any exact value that the answer should become. +- NEVER write phrases like "the correct answer is", "which gives", "yielding", + "should be ", "Option ", or "Graph ". +- If a fix would require naming a corrected value, replace it with a method-level + instruction such as "redo that computation carefully" or "apply the theorem with + the correct quantities". +- Keep every "reason" and "fix" clear and concise — one short sentence each. +- "overall" is "OK" only if NO criterion is FAIL. +- Output only the JSON object.""" + +_RFT_DIAG_USER = """\ +## Task / query (context) +{query} + +## Rubric +{rubric} + +## Segment +{segment} + +Now output the diagnostic JSON object.""" + +_MATH_RUBRIC = [ + ('The reasoning contains no arithmetic or algebraic error', True), + ('Each step follows logically from the previous ones', True), + ('No formula or theorem is misstated or misapplied', True), + ('The approach is on track to answer the actual question asked', False), + ('No step contradicts an earlier established fact', False), +] + + +class _RftRubricVerifier(RubricVerifier): + def _diagnose_trajectory(self, query: str, rubric_block: str, segment_text: str) -> dict: + return {'messages': [ + {'role': 'system', 'content': _RFT_DIAG_SYSTEM}, + {'role': 'user', 'content': _RFT_DIAG_USER.format( + query=query, rubric=rubric_block, segment=segment_text)}]} + + +def build_rubric_checker() -> Optional[RubricVerifier]: + """Fixed math-process rubric verifier (teacher-served). None if no LLM backup env.""" + if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') + or os.environ.get('OPENAI_API_KEY')): + return None + return _RftRubricVerifier( + fixed_rubric=[RubricItem(t, is_hard=h) for t, h in _MATH_RUBRIC], gate=True) + + +def _format_diagnosis(detail) -> str: + """One line per criterion (PASS/FAIL + reason + fix) then a summary.""" + rub = detail.rubric + lines = [] + for it in detail.items: + text = rub[it.index - 1].text if 0 < it.index <= len(rub) else f'criterion {it.index}' + if it.verdict: + lines.append(f'- [PASS] {text}') + else: + tail = f': {it.reason}' if it.reason else '' + tail += f' (fix: {it.fix})' if it.fix else '' + lines.append(f'- [FAIL] {text}{tail}') + if detail.summary: + lines.append(f'Summary: {detail.summary}') + return '\n'.join(lines) + + +def diagnose_views(checker, hard: List[Dict[str, Any]], args: argparse.Namespace, + cache: DiskCache) -> None: + """Rubric-check every view-A problem's greedy attempt in parallel (disk-cached by + problem + attempt), stashing the formatted findings on ``r['_rubric_diag']``.""" + targets = [r for r in hard if r.get('_view') == 'A'] + if not checker or not targets: + return + + def _key(r: Dict[str, Any]) -> str: + return DiskCache.key_for(r['problem'], r.get('_init', [{}])[0].get('text', '')) + + pending = [] + for r in targets: + key = _key(r) + if key in cache: + r['_rubric_diag'] = cache.get(key) + else: + pending.append((r, key)) + if not pending: + return + + def _run(item): + r, key = item + seg = {'messages': [{'role': 'user', 'content': r['problem']}, + {'role': 'assistant', 'content': r['_init'][0]['text']}]} + try: + return r, key, _format_diagnosis(checker.diagnose(seg, query=r['problem'])) + except Exception as exc: # teacher hiccup -> no-diagnosis prompt (not cached) + logger.warning(f'[rubric] diagnose error: {exc}') + return r, key, None + + workers = max(1, min(args.rubric_workers, len(pending))) + with ThreadPoolExecutor(max_workers=workers) as ex: + for r, key, diag in ex.map(_run, pending): + r['_rubric_diag'] = diag or '' + if diag is not None: + cache.put(key, diag) + + +# =========================================================================== +# Block E -- chunk draw, pipeline, record building +# =========================================================================== +def _baseline_class(r: Dict[str, Any]) -> str: + """success | fail_loop (out of length / never terminated) | fail_wrong.""" + roll = r['_init'][0] + if roll['correct']: + return 'success' + return 'fail_loop' if (roll['stop_reason'] == 'length' or not roll['terminated']) else 'fail_wrong' + + +def _select_balanced(buckets: Dict[str, List[Dict[str, Any]]], n_success: int, + n_fail: int, n_fail_loop: int) -> List[Dict[str, Any]]: + """Pick n_fail base-fails (toward n_fail_loop loop-fails, best-effort) + n_success + base-successes; top up any shortfall from leftovers.""" + loop, wrong, succ = buckets['fail_loop'], buckets['fail_wrong'], buckets['success'] + take_loop = min(n_fail_loop, len(loop)) + take_wrong = min(n_fail - take_loop, len(wrong)) + take_loop = min(n_fail - take_wrong, len(loop)) + sel = loop[:take_loop] + wrong[:take_wrong] + succ[:n_success] + target = n_success + n_fail + if len(sel) < target: + used = {id(x) for x in sel} + sel += [x for b in (loop, wrong, succ) for x in b if id(x) not in used][:target - len(sel)] + return sel + + +def draw_chunk(pool: ProblemPool, base_sampler, base_dp: int, args: argparse.Namespace, + cache: DiskCache) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + """Draw one chunk, running baseline rollout (Phase 1) on every drawn problem. With + ``--balance``, keep drawing+baselining until the target base fail:success mix is + reachable (or the budget is hit), then select a balanced subset.""" + if not args.balance: + chunk = pool.draw(args.chunk_size) + fresh = baseline_rollout(base_sampler, chunk, base_dp, args, cache) + return chunk, {'enabled': False, 'n_drawn': len(chunk), 'n_baseline_fresh': fresh} + + n_success = max(0, min(args.chunk_size, round(args.chunk_size * args.balance_success_frac))) + n_fail = args.chunk_size - n_success + n_fail_loop = round(n_fail * args.balance_loop_frac) + buckets: Dict[str, List[Dict[str, Any]]] = {'success': [], 'fail_loop': [], 'fail_wrong': []} + budget, n_drawn, n_fresh, seen = args.chunk_size * args.balance_max_draws_mult, 0, 0, set() + while n_drawn < budget: + if (len(buckets['success']) >= n_success + and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail): + break + batch = pool.draw(args.chunk_size) + n_fresh += baseline_rollout(base_sampler, batch, base_dp, args, cache) + n_drawn += len(batch) + for r in batch: + if id(r) not in seen: + seen.add(id(r)) + buckets[_baseline_class(r)].append(r) + + reached = (len(buckets['success']) >= n_success + and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail) + chunk = _select_balanced(buckets, n_success, n_fail, n_fail_loop) + sel_success = sum(1 for r in chunk if not r['_failed']) + stats = { + 'enabled': True, 'n_drawn': n_drawn, 'n_baseline_fresh': n_fresh, 'n_selected': len(chunk), + 'target_success': n_success, 'target_fail': n_fail, 'target_fail_loop': n_fail_loop, + 'selected_success': sel_success, 'selected_fail': len(chunk) - sel_success, + 'selected_fail_loop': sum(1 for r in chunk if _baseline_class(r) == 'fail_loop'), + 'selected_fail_wrong': sum(1 for r in chunk if _baseline_class(r) == 'fail_wrong'), + 'selected_success_frac': (sel_success / len(chunk)) if chunk else 0.0, + 'budget_hit': not reached, + } + return chunk, stats + + +def _assign_advantages(hard: List[Dict[str, Any]], args: argparse.Namespace) -> None: + """Group-relative advantage over each problem's scored candidates using the greedy + binary reward R in {0,1}: ``A = (R - mean) / (std + eps)``. std==0 groups (all solve + / all fail) get advantage 0 and no gradient -- GRPO's variance selects informative + problems, so no explicit difficulty gate is needed.""" + eps = 1e-6 + for r in hard: + for c in r['_cands']: + c['advantage'], c['grpo_adv'], c['kept'] = 0.0, 0.0, False + cs = ([c for c in r['_cands'] if c.get('reward') is not None] if args.format_in_reward + else [c for c in r['_cands'] if c['leaked'] is False and c.get('reward') is not None]) + if len(cs) < 2: + continue + rewards = [c['reward'] for c in cs] + mean_r = sum(rewards) / len(rewards) + std = (sum((x - mean_r) ** 2 for x in rewards) / len(rewards)) ** 0.5 + if std < 1e-9: + continue + for c in cs: + adv = (c['reward'] - mean_r) / (std + eps) + c['advantage'], c['grpo_adv'], c['kept'] = adv, adv, c['reward'] > mean_r + + +def process_chunk(base_sampler, skill_sampler, leak: LeakVerifier, chunk: List[Dict[str, Any]], + ci: int, base_dp: int, skill_dp: int, args: argparse.Namespace, + checker, rubric_cache: DiskCache + ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], List[Dict[str, Any]]]: + """view assign -> rubric-check (view A) -> skill-gen -> leak-filter -> with-skill + greedy pass -> GRPO advantages. ``chunk`` arrives already baselined by draw_chunk.""" + hard = chunk + + # Phase 2: view routing + view-A rubric check (view B is query-only, no rubric). + for r in hard: + r['_view'], r['_rubric_diag'] = _assign_view(r['problem'], args), '' + diagnose_views(checker, hard, args, rubric_cache) + + # Phase 3: skill-gen (thinking ON), per-view prompt; re-sample empties. + flat: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] + pending = list(hard) + for _ in range(args.skill_retries + 1): + if not pending: + break + sg_out = _run_samples(skill_sampler, [_view_prompt(r) for r in pending], + args.n_skills, args.skill_max_tokens, skill_dp, + temperature=args.skill_gen_temperature, + top_p=args.skill_gen_top_p, top_k=args.skill_gen_top_k) + still = [] + for r, seqs in zip(pending, sg_out): + got = False + for s in seqs: + resp = _clean_text(getattr(s, 'decoded', '') or '') + block = _extract_skills_block(resp) + cand = {'skills': block or '', 'response': resp, 'parseable': bool(block), + 'view': r['_view'], 'leaked': None, 'leak_reason': '', + 'leak_source': '', 'with_pass': None, 'reward': None, 'rolls': []} + r['_cands'].append(cand) + if block: + flat.append((r, cand)) + got = True + if not got: + still.append(r) + pending = still + + # Phase 4: leak filter (view A only; view B is query-only -> treated clean, SEAM-like). + for r, c in flat: + if r.get('_view') != 'A': + c['leaked'], c['leak_reason'], c['leak_source'] = False, '', 'skipped_viewB' + flat_a = [(r, c) for r, c in flat if r.get('_view') == 'A'] + if flat_a: + details = leak.leak_batch( + [{'content': c['skills'], 'query': r['problem'], 'reference': r['reference_answer']} + for r, c in flat_a], max_workers=args.leak_workers) + for (r, c), d in zip(flat_a, details): + c['leaked'], c['leak_reason'], c['leak_source'] = bool(d.leaked), d.reason, d.source + + # Phase 5: with-skill greedy pass (T=0, M=1) on clean candidates. Reward = correct, + # absolute (no baseline subtraction); the group mean in Phase 6 is the only baseline. + clean = [(r, c) for r, c in flat if c['leaked'] is False] + if clean: + ws_out = _run_samples(base_sampler, + [build_skill_solve_prompt(r['problem'], c['skills']) for r, c in clean], + 1, args.max_tokens, base_dp, temperature=0.0) + for (r, c), seqs in zip(clean, ws_out): + c['rolls'] = [_parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll()] + c['with_pass'] = 1.0 if c['rolls'][0]['correct'] else 0.0 + c['reward'] = c['with_pass'] + if args.format_in_reward: # unparseable/leaked candidates score 0 and still join the group + for r in hard: + for c in r['_cands']: + if c['reward'] is None: + c['reward'] = 0.0 + + # Phase 6: group-relative GRPO advantage. + _assign_advantages(hard, args) + return ([_full_record(r, ci) for r in chunk], _chunk_summary(chunk, ci, args), + _group_records(chunk, args)) + + +def _roll(x: Dict[str, Any]) -> Dict[str, Any]: + return {k: x[k] for k in ('pred', 'correct', 'terminated', 'passed', + 'stop_reason', 'gen_tokens', 'text')} + + +def _is_trainable(c: Dict[str, Any], args: argparse.Namespace) -> bool: + """A candidate reaches the GRPO update iff its advantage is non-zero (and, without + --format-in-reward, is also clean and scored).""" + adv_nz = abs(c.get('advantage') or 0.0) > 1e-9 + if args.format_in_reward: + return adv_nz + return c['leaked'] is False and c.get('with_pass') is not None and adv_nz + + +def _full_record(r: Dict[str, Any], ci: int) -> Dict[str, Any]: + """Complete per-problem trace: init attempt, baseline, and all candidates.""" + init = r['_init'][0] + return { + 'record_type': 'problem', 'chunk': ci, 'problem': r['problem'], + 'reference_answer': r['reference_answer'], 'level': r.get('level', ''), + 'failed_first_try': r['_failed'], + 'init_attempt': {'text': init['text'], 'pred': init['pred'], 'correct': init['correct'], + 'terminated': init['terminated'], 'stop_reason': init['stop_reason'], + 'gen_tokens': init['gen_tokens']}, + 'baseline_pass': r['_baseline_pass'], 'is_hard': r['_hard'], + 'view': r.get('_view', ''), 'rubric_diag': r.get('_rubric_diag', ''), + 'baseline_rolls': [_roll(x) for x in r['_baseline_rolls']], + 'candidates': [{ + 'skills': c['skills'], 'response': c['response'], 'parseable': c['parseable'], + 'leaked': c['leaked'], 'leak_reason': c['leak_reason'], 'leak_source': c['leak_source'], + 'with_pass': c['with_pass'], 'reward': c.get('reward'), 'advantage': c.get('advantage'), + 'grpo_adv': c.get('grpo_adv'), 'kept': c.get('kept'), + 'rolls': [_roll(x) for x in c['rolls']], + } for c in r['_cands']], + } + + +def _view_stats(hard: List[Dict[str, Any]], view: str) -> Dict[str, Any]: + hv = [r for r in hard if r.get('_view') == view] + cands = [c for r in hv for c in r['_cands'] if c['parseable']] + clean = [c for c in cands if c['leaked'] is False] + adopted = sum(1 for r in hv + if any(c['leaked'] is False and abs(c.get('advantage') or 0.0) > 1e-9 + for c in r['_cands'])) + return {'n_hard': len(hv), 'n_candidates_parseable': len(cands), 'n_clean': len(clean), + 'n_adopted_problems': adopted, 'adoption_rate': (adopted / len(hv)) if hv else 0.0} + + +def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespace) -> Dict[str, Any]: + hard = [r for r in chunk if r['_hard']] + all_cands = [c for r in chunk for c in r['_cands']] + cands = [c for c in all_cands if c['parseable']] + scored = [c for c in cands if c['with_pass'] is not None] + ws_rolls = [x for c in scored for x in c['rolls']] + train_cands = [c for c in all_cands if _is_trainable(c, args)] + fail_cands = [c for r in chunk if r['_failed'] for c in r['_cands']] + base_acc = (sum(r['_baseline_pass'] for r in hard) / len(hard)) if hard else 0.0 + ws_acc = (sum(c['with_pass'] for c in scored) / len(scored)) if scored else 0.0 + abs_adv = lambda cs: sum(abs(c.get('advantage') or 0.0) for c in cs) + total_abs = abs_adv(all_cands) + return { + 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), + 'n_failed_first_try': sum(1 for r in chunk if r['_failed']), 'n_hard': len(hard), + 'n_generated': len(all_cands), 'n_candidates_parseable': len(cands), + 'n_unparseable': len(all_cands) - len(cands), + 'n_leaked': sum(1 for c in cands if c['leaked']), + 'n_clean': sum(1 for c in cands if c['leaked'] is False), + 'n_reward_pos': sum(1 for c in scored if c['reward']), 'n_train_samples': len(train_cands), + 'n_train_from_fail': sum(1 for c in fail_cands if _is_trainable(c, args)), + 'abs_adv_from_fail_frac': (abs_adv(fail_cands) / total_abs) if total_abs > 0 else 0.0, + 'avg_baseline_pass_on_hard': base_acc, 'avg_withskill_pass': ws_acc, + 'avg_lift': ws_acc - base_acc, + 'termination_rate_withskill': (sum(1 for x in ws_rolls if x['terminated']) / len(ws_rolls)) if ws_rolls else 0.0, + 'view_A': _view_stats(hard, 'A'), 'view_B': _view_stats(hard, 'B'), + } + + +def _group_records(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> List[Dict[str, Any]]: + """GRPO training records: every trainable candidate with its view + rubric diagnosis + (the prompt is rebuilt from those by ``_skillgen_messages``, no trajectory stored).""" + out = [] + for r in chunk: + if not r['_hard']: + continue + for c in r['_cands']: + if _is_trainable(c, args): + out.append({ + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'view': r.get('_view', 'A'), 'diagnosis': r.get('_rubric_diag', ''), + 'response': c['response'], 'skills': c['skills'], + 'advantage': c['advantage'], 'grpo_adv': c['grpo_adv'], 'kept': c['kept'], + 'reward': c['reward'], 'with_pass': c['with_pass']}) + return out + + +# =========================================================================== +# Block F -- samplers, args, main +# =========================================================================== +def init_samplers(args: argparse.Namespace): + """8 GPUs: ranks 0-3 skill_sampler, ranks 4-7 base_sampler (both vLLM tp1 dp4).""" + twinkle.initialize(mode='ray', nproc_per_node=8, lazy_collect=False, groups=[ + DeviceGroup(name='skill_sampler', ranks=list(range(0, 4)), device_type='GPU'), + DeviceGroup(name='base_sampler', ranks=list(range(4, 8)), device_type='GPU')]) + samplers = [] + for group in ('skill_sampler', 'base_sampler'): + s = vLLMSampler( + model_id=MODEL_ID, + engine_args={'gpu_memory_utilization': GPU_MEM, + 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, + device_mesh=DeviceMesh.from_sizes(world_size=4, dp_size=4), remote_group=group) + s.set_template('Template', model_id=MODEL_ID, enable_thinking=True, + max_length=args.max_model_len) + samplers.append(s) + return samplers[1], samplers[0], 4, 4 # base_sampler, skill_sampler, base_dp, skill_dp def _build_args() -> argparse.Namespace: - p = argparse.ArgumentParser( - description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter) + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) p.add_argument('--total-problems', type=int, default=3200, help='Final number of problems selected into generated chunks.') p.add_argument('--base-success-frac', type=float, default=0.3, - help='Target fraction of selected problems solved by the frozen base.') + help='Target fraction of selected problems the frozen base solves.') p.add_argument('--output-dir', default='./output/reflexion_skill_data') - p.add_argument('--overwrite', action='store_true') + p.add_argument('--cache-dir', default='', + help='Baseline/rubric cache dir (default /cache).') + p.add_argument('--no-cache', action='store_true', help='Disable disk cache read/write.') + p.add_argument('--overwrite', action='store_true', help='Replace existing output jsonl.') p.add_argument('--seed', type=int, default=42) - - advanced = p.add_argument_group('advanced knobs, usually leave unchanged') - advanced.add_argument('--dataset', choices=('aops', 'math'), default='aops') - advanced.add_argument('--n', type=int, default=0, + p.add_argument('--dataset', choices=('aops', 'math'), default='aops') + p.add_argument('--n', type=int, default=0, help='Raw train-pool size; 0 derives it from --total-problems.') - advanced.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) - advanced.add_argument('--eval-size', type=int, default=128) - advanced.add_argument('--chunk-size', type=int, default=16) - advanced.add_argument('--balance', action=argparse.BooleanOptionalAction, default=True) - advanced.add_argument('--balance-loop-frac', type=float, default=0.5) - advanced.add_argument('--balance-max-draws-mult', type=int, default=8) - advanced.add_argument('--n-skills', type=int, default=8) - advanced.add_argument('--view-b-frac', type=float, default=0.5) - advanced.add_argument('--skill-retries', type=int, default=2) - advanced.add_argument('--skill-gen-temperature', type=float, default=1.0) - advanced.add_argument('--skill-gen-top-p', type=float, default=1.0) - advanced.add_argument('--skill-gen-top-k', type=int, default=-1) - advanced.add_argument('--max-model-len', type=int, default=16384) - advanced.add_argument('--max-tokens', type=int, default=8192) - advanced.add_argument('--skill-max-tokens', type=int, default=8192) - advanced.add_argument('--leak-workers', type=int, default=16) - advanced.add_argument('--rubric-workers', type=int, default=16) - advanced.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--eval-size', type=int, default=128) + p.add_argument('--chunk-size', type=int, default=16) + p.add_argument('--balance', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--balance-loop-frac', type=float, default=0.5) + p.add_argument('--balance-max-draws-mult', type=int, default=8) + p.add_argument('--n-skills', type=int, default=8) + p.add_argument('--view-b-frac', type=float, default=0.5) + p.add_argument('--skill-retries', type=int, default=2) + p.add_argument('--skill-gen-temperature', type=float, default=1.0) + p.add_argument('--skill-gen-top-p', type=float, default=1.0) + p.add_argument('--skill-gen-top-k', type=int, default=-1) + p.add_argument('--max-model-len', type=int, default=16384) + p.add_argument('--max-tokens', type=int, default=8192) + p.add_argument('--skill-max-tokens', type=int, default=8192) + p.add_argument('--leak-workers', type=int, default=16) + p.add_argument('--rubric-workers', type=int, default=16) + p.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True) args = p.parse_args() - _resolve_args(args) - return args - - -def _resolve_args(args: argparse.Namespace) -> None: - if args.total_problems <= 0: - raise ValueError('--total-problems must be positive') - if args.chunk_size <= 0: - raise ValueError('--chunk-size must be positive') + if args.total_problems <= 0 or args.chunk_size <= 0: + raise ValueError('--total-problems and --chunk-size must be positive') if not 0.0 <= args.base_success_frac <= 1.0: raise ValueError('--base-success-frac must be in [0, 1]') args.chunks = math.ceil(args.total_problems / args.chunk_size) args.balance_success_frac = args.base_success_frac if args.n <= 0: - args.n = max(args.total_problems + args.eval_size, - math.ceil(args.total_problems * 1.5)) + args.n = max(args.total_problems + args.eval_size, math.ceil(args.total_problems * 1.5)) + return args -def _init_samplers(args: argparse.Namespace): - model = 'ms://Qwen/Qwen3-4B' - device_groups = [ - DeviceGroup(name='skill_sampler', ranks=list(range(0, 4)), device_type='GPU'), - DeviceGroup(name='base_sampler', ranks=list(range(4, 8)), device_type='GPU'), - ] - twinkle.initialize(mode='ray', nproc_per_node=8, groups=device_groups, - lazy_collect=False) - skill_sampler = vLLMSampler( - model_id='Qwen/Qwen3-4B', - engine_args={'gpu_memory_utilization': 0.8, - 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, - device_mesh=DeviceMesh.from_sizes(world_size=4, dp_size=4), - remote_group='skill_sampler') - skill_sampler.set_template('Template', model_id=model, - enable_thinking=True, max_length=args.max_model_len) - base_sampler = vLLMSampler( - model_id=model, - engine_args={'gpu_memory_utilization': 0.8, - 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, - device_mesh=DeviceMesh.from_sizes(world_size=4, dp_size=4), - remote_group='base_sampler') - base_sampler.set_template('Template', model_id=model, - enable_thinking=True, max_length=args.max_model_len) - return base_sampler, skill_sampler, 4, 4 - - -def _write_jsonl_row(handle, row: Dict[str, Any]) -> None: +def _write(handle, row: Dict[str, Any]) -> None: handle.write(json.dumps(row, ensure_ascii=False) + '\n') def main() -> None: args = _build_args() - records, eval_records, data_stats = rft._load_records(args) - rft._validate_run_config(args, records) + records, eval_records, data_stats = _load_records(args) + if not records: + raise ValueError(f'loaded 0 {args.dataset} problems') + if len(records) < args.chunk_size: + raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded ({len(records)})') + os.makedirs(args.output_dir, exist_ok=True) data_path = os.path.join(args.output_dir, 'skill_dataset.jsonl') gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') @@ -125,68 +1052,71 @@ def main() -> None: raise FileExistsError(f'{path} exists; pass --overwrite to replace it') if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY')): - sys.stderr.write('[build-rft-data] WARNING: no LLM backup env; leak/rubric checks degrade\n') + sys.stderr.write('[build] WARNING: no LLM backup env; leak/rubric checks degrade\n') - base_sampler, skill_sampler, base_dp, skill_dp = _init_samplers(args) + base_sampler, skill_sampler, base_dp, skill_dp = init_samplers(args) leak = LeakVerifier(sampler=None, answer_only=True) - checker = rft._build_rubric_checker() - pool = rft._ProblemPool(records, args.seed) - rubric_cache: Dict[str, str] = {} + checker = build_rubric_checker() + pool = ProblemPool(records, args.seed) + + cache_dir = args.cache_dir or os.path.join(args.output_dir, 'cache') + os.makedirs(cache_dir, exist_ok=True) + use_cache = not args.no_cache + baseline_cache = DiskCache(os.path.join(cache_dir, 'baseline.jsonl'), use_cache) + rubric_cache = DiskCache(os.path.join(cache_dir, 'rubric.jsonl'), use_cache) cfg = { - 'record_type': 'config', 'mode': 'offline_data_build', 'model': 'Qwen/Qwen3-4B', + 'record_type': 'config', 'mode': 'offline_data_build', 'model': MODEL_ID, 'dataset': args.dataset, 'n': len(records), 'eval_n': len(eval_records), - 'total_problems': args.total_problems, 'seed': args.seed, - 'numeric_only': args.numeric_only, + 'total_problems': args.total_problems, 'seed': args.seed, 'numeric_only': args.numeric_only, 'raw_loaded': data_stats['raw_loaded'], 'numeric_dropped': data_stats['numeric_dropped'], 'chunks': args.chunks, 'chunk_size': args.chunk_size, 'n_skills': args.n_skills, 'view_b_frac': args.view_b_frac, 'balance': args.balance, - 'base_success_frac': args.base_success_frac, - 'balance_success_frac': args.balance_success_frac, + 'base_success_frac': args.base_success_frac, 'balance_success_frac': args.balance_success_frac, 'reward': 'greedy_binary(correct)', 'advantage': 'group_relative', - 'format_in_reward': args.format_in_reward, 'started': int(time.time()), + 'format_in_reward': args.format_in_reward, 'cache': use_cache, + 'rubric_check': 'fixed_math_5crit(viewA)' if checker else 'disabled', + 'started': int(time.time()), } - total_groups = 0 + total_groups, selected = 0, 0 with open(gen_path, 'w', encoding='utf-8') as gen_f, \ open(data_path, 'w', encoding='utf-8') as data_f, \ open(eval_path, 'w', encoding='utf-8') as eval_f: for handle in (gen_f, data_f, eval_f): - _write_jsonl_row(handle, cfg) + _write(handle, cfg) for rec in eval_records: - _write_jsonl_row(eval_f, {'record_type': 'eval_holdout', **rec}) + _write(eval_f, {'record_type': 'eval_holdout', **rec}) eval_f.flush() - selected = 0 + full_chunk_size = args.chunk_size for ci in range(args.chunks): remaining = args.total_problems - selected if remaining <= 0: break - original_chunk_size = args.chunk_size - args.chunk_size = min(original_chunk_size, remaining) - try: - chunk, balance = rft._draw_chunk(pool, base_sampler, base_dp, args) - full, summary, groups = rft.process_chunk( + args.chunk_size = min(full_chunk_size, remaining) # last chunk may be short + chunk, balance = draw_chunk(pool, base_sampler, base_dp, args, baseline_cache) + full, summary, groups = process_chunk( base_sampler, skill_sampler, leak, chunk, ci, base_dp, skill_dp, args, checker, rubric_cache) - finally: - args.chunk_size = original_chunk_size summary['balance'] = balance for rec in full: - _write_jsonl_row(gen_f, rec) - _write_jsonl_row(gen_f, summary) + _write(gen_f, rec) + _write(gen_f, summary) gen_f.flush() for row in groups: - _write_jsonl_row(data_f, {'chunk': ci, **row}) + _write(data_f, {'chunk': ci, **row}) data_f.flush() total_groups += len(groups) selected += len(chunk) sys.stderr.write( - f'[build-rft-data] g{ci}: problems={selected}/{args.total_problems} ' + f'[build] g{ci}: problems={selected}/{args.total_problems} ' f'train={len(groups)} total={total_groups} ' f'acc={summary["avg_baseline_pass_on_hard"]:.2f}->{summary["avg_withskill_pass"]:.2f} ' f'lift={summary["avg_lift"]:+.3f}\n') - sys.stderr.write(f'[build-rft-data] done: train records -> {data_path}; trace -> {gen_path}\n') + baseline_cache.close() + rubric_cache.close() + sys.stderr.write(f'[build] done: {total_groups} train records -> {data_path}; trace -> {gen_path}\n') if __name__ == '__main__': diff --git a/cookbook/exp/embedding/train_reflexion_skill.py b/cookbook/exp/embedding/train_reflexion_skill.py new file mode 100644 index 000000000..655233af5 --- /dev/null +++ b/cookbook/exp/embedding/train_reflexion_skill.py @@ -0,0 +1,1390 @@ +"""Online GRPO RFT for the reflexion skill generator (self-contained, cached). + +Trains an INDEPENDENT skill model to write reusable skills that, injected into a +FROZEN base solver's system prompt, raise its accuracy. The base is never trained; +it only produces the reward. Per chunk: base greedy solve -> rubric process-check +(view A) -> skill-gen (thinking ON, N candidates) -> leak filter -> with-skill +greedy pass -> group-relative advantage -> ONE on-policy GRPO step -> sync weights. +Reward is deterministic (T=0, M=1) binary correctness; A = (R - mean) / (std + eps) +within each problem-group, so std=0 groups give no gradient (GRPO variance selects). + +Each problem is routed to EXACTLY ONE view: A = problem + rubric findings, B = +query only (deployment form). Skill-gen trains only the final turn. + +Base greedy rollouts and rubric diagnoses are cached to disk (jsonl, md5-keyed) so +restarts skip them; skill-gen is on-policy and never cached. + +8 GPUs: ranks 0-3 train (FSDP2), 4-5 skill_sampler (synced), 6-7 base_sampler (frozen). +Leak / rubric use the backup teacher API: set LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / +LLM_BACKUP_MODEL. + +Launch: + LLM_BACKUP_API_KEY=... python cookbook/exp/embedding/train_reflexion_skill.py \ + --dataset aops --n 5000 --chunk-size 16 --lr 6e-6 +""" +import argparse +import copy +import hashlib +import json +import os +import re +import sys +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.data_format import SamplingParams +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.model import TransformersModel +from twinkle.patch.no_split_modules import NoSplitModulesPatch +from twinkle.processor import InputProcessor +from twinkle.sampler import vLLMSampler +from twinkle.template import Template +from twinkle_agentic.verifier import RubricVerifier +from twinkle_agentic.verifier.rubric_verifier import RubricItem + +logger = get_logger() + +try: + import swanlab +except ImportError: + swanlab = None + +MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') +GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.8)) +GEN_TEMPERATURE = float(os.environ.get('GEN_TEMPERATURE', 0.6)) +GEN_TOP_P = float(os.environ.get('GEN_TOP_P', 0.95)) +AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') +MATH_DATASET_ID = os.environ.get('MATH_DATASET_ID', 'modelscope/competition_math') + +# GPU layout: train (FSDP2) + skill_sampler (synced) + base_sampler (frozen). +TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 4)) +SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) +BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 2)) +NUM_GPUS = TRAIN_GPUS + SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS +TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', 2)) +TRAIN_DP = max(1, TRAIN_GPUS // TRAIN_FSDP) + + +# =========================================================================== +# Block A -- boxed extraction + answer grading +# =========================================================================== +_BOXED_RE = re.compile(r'\\boxed\s*\{') + + +def extract_boxed(text: str) -> Optional[str]: + """Last ``\\boxed{...}`` content, brace-balanced.""" + if not text: + return None + last = None + for m in _BOXED_RE.finditer(text): + depth, i = 1, m.end() + while i < len(text) and depth > 0: + depth += (text[i] == '{') - (text[i] == '}') + i += 1 + if depth == 0: + last = text[m.end():i - 1].strip() + return last + + +def normalize_answer(ans: str) -> str: + if not ans: + return '' + s = str(ans).strip() + m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) + if m: + return m.group(1) + s = s.strip('$').strip().replace('−', '-') + 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').replace(r'\tfrac', r'\frac') + s = s.replace(r'\,', '').replace(r'\;', '').replace(r'\!', '') + s = re.sub(r'\\(?:quad|qquad|\s)', '', s) + s = re.sub(r'\s+', '', s) + s = re.sub(r'\\sqrt([0-9A-Za-z])', r'\\sqrt{\1}', s) + s = re.sub(r'\\frac\{([^{}]+)\}([^{}\\])', r'\\frac{\1}{\2}', s) + s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) + s = re.sub(r'\{[a-zA-Z]+\}$', '', s) + s = re.sub(r'\^\\circ|\^\{\\circ\}|°|\\circ', 'deg', s) + s = s.replace(r'\minus{}', '-').replace(r'\minus', '-') + + def _frac_to_slash(mt): + text = mt.group(0) + pos = text.index('{') + 1 + depth, num_start = 1, pos + while depth > 0: + depth += (text[pos] == '{') - (text[pos] == '}') + pos += 1 + numer = text[num_start:pos - 1] + pos += 1 + den_start, depth = pos, 1 + while depth > 0: + depth += (text[pos] == '{') - (text[pos] == '}') + pos += 1 + return f'({numer})/({text[den_start:pos - 1]})' + + s = re.sub(r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', _frac_to_slash, s) + s = re.sub(r'(? bool: + try: + va, vb = float(a.replace('(', '').replace(')', '')), 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) + return va is not None and vb is not None and abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) + + +_MCQ_REF_RE = re.compile( + r'^\\(?:textbf|text|mathrm|mathbf)\{\(?([A-E])\)?\s*\}\s*(.+)$' + r'|^\(?([A-E])\)\s+(.+)$') +_VAR_PREFIX_RE = re.compile(r'^(?:[A-Za-z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)$', re.DOTALL) + + +def _split_mcq(ans: str): + s = ans.strip() + m = _MCQ_REF_RE.match(s) + if m: + return (m.group(1) or m.group(3)), ((m.group(2) or m.group(4) or '').strip() or None) + bl = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) + return (bl.group(1), None) if bl else (None, s or None) + + +def _strip_var_prefix(ans: str) -> str: + m = _VAR_PREFIX_RE.match((ans or '').strip()) + return m.group(1).strip() if m else (ans or '') + + +def _normalize_tuple(ans: str) -> str: + return re.sub(r'[\s()\[\]{}\\]', '', ans or '') + + +def _math_verify_equal(predicted: str, reference: str) -> bool: + try: + from math_verify import parse, verify + gold = parse(r'\boxed{' + reference + '}', parsing_timeout=1) + pred = parse(r'\boxed{' + predicted + '}', parsing_timeout=1) + return bool(gold and pred and verify(gold, pred, timeout_seconds=1)) + except Exception: + return False + + +def answers_match(predicted: str, reference: str) -> bool: + if not predicted or not reference: + return False + norm_p, norm_r = normalize_answer(predicted), normalize_answer(reference) + if norm_p == norm_r or norm_p.lower() == norm_r.lower() or _try_numeric_equal(norm_p, norm_r): + return True + stripped_p = normalize_answer(_strip_var_prefix(predicted)) + stripped_r = normalize_answer(_strip_var_prefix(reference)) + if stripped_p and stripped_r: + if (stripped_p == stripped_r or stripped_p.lower() == stripped_r.lower() + or _try_numeric_equal(stripped_p, stripped_r)): + return True + p_letter, p_value = _split_mcq(predicted) + r_letter, r_value = _split_mcq(reference) + if p_letter and r_letter and p_letter == r_letter: + return True + if (p_letter and p_value is None) and (r_letter and r_value is None): + return False + 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 and (p_val == r_val or p_val.lower() == r_val.lower() or _try_numeric_equal(p_val, r_val)): + return True + for left, right in ((norm_p, norm_r), (stripped_p, stripped_r)): + tl, tr = _normalize_tuple(left), _normalize_tuple(right) + if ',' in tl and tl == tr: + return True + if '=' in norm_r and '=' not in norm_p: + if any(part and (part == norm_p or _try_numeric_equal(part, norm_p)) for part in norm_r.split('=')): + return True + if '=' in norm_p and '=' not in norm_r: + if any(part and (part == norm_r or _try_numeric_equal(part, norm_r)) for part in norm_p.split('=')): + return True + return _math_verify_equal(stripped_p or norm_p, stripped_r or norm_r) + + +# =========================================================================== +# Block B -- prompts, skill parsing, batched sampling +# =========================================================================== +DIRECT_SYSTEM = ( + 'You are an expert competition mathematician. Solve the following problem ' + 'step by step. Provide your final answer inside \\boxed{}.') + +_SKILL_SOLVE_PREFIX = ( + DIRECT_SYSTEM + '\n\n' + 'Before you start, keep these reminders in mind to avoid common mistakes on this ' + 'type of problem:\n') +_SKILL_SOLVE_SUFFIX = '\nApply them where relevant, but rely on your own reasoning to reach the answer.' + + +def build_direct_prompt(problem: str) -> Dict[str, Any]: + return {'messages': [{'role': 'system', 'content': DIRECT_SYSTEM}, + {'role': 'user', 'content': problem}]} + + +def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: + # Concatenation (not .format): DIRECT_SYSTEM/skill contain literal braces. + return {'messages': [ + {'role': 'system', 'content': _SKILL_SOLVE_PREFIX + skill + _SKILL_SOLVE_SUFFIX}, + {'role': 'user', 'content': problem}]} + + +# -- skill-gen prompts (view A: problem + rubric findings; view B: query only) -- +SKILL_GEN_SYSTEM = ( + 'You are a math problem-solving coach. You are given a problem and an automated ' + 'process-check of an earlier attempt at it (which criteria it passed or failed, with ' + 'suggested fixes). Use the check to see where a solver of this problem tends to go ' + 'wrong, then give reusable guidance that helps a separate solver avoid those failures ' + 'on this and similar problems.\n' + '- Do not solve it or state the final answer or any specific intermediate value.\n' + '- Turn each relevant failure into general guidance (the method to reach for, the ' + 'pitfall to watch, a quick check) rather than narrating this attempt.\n' + 'Think briefly, then output only the guidance wrapped in and .') + +SKILL_GEN_SYSTEM_Q = ( + 'You are a math problem-solving coach. Read the problem below and give reusable ' + 'guidance that would help a separate solver reach the answer on this and similar ' + 'problems.\n' + '- Do not solve it or state the final answer or any specific intermediate value.\n' + '- Name the key idea or method to reach for and the main pitfall to avoid.\n' + 'Think briefly, then output only the guidance wrapped in and .') + +SKILL_GEN_USER_Q = ( + 'Problem:\n{problem}\n\n' + 'Now output the guidance.') + +SKILL_GEN_USER_RUBRIC = ( + 'Problem:\n{problem}\n\n' + 'Process check of an earlier attempt (automated rubric verifier -- treat as ' + 'evidence, not gospel; PASS/FAIL per criterion with suggested fixes for failures):\n' + '{diagnosis}\n\n' + 'Now output the guidance.') + + +def _skillgen_messages(problem: str, view: str, diagnosis: str) -> List[Dict[str, Any]]: + """Single source of truth for the skill-gen prompt (used at BOTH generation and + training so they never diverge). View A with a localisable failure uses problem + + rubric findings; view B -- or a view-A problem whose rubric flagged nothing (no + ``[FAIL]``) -- degrades to the query-only prompt.""" + if view == 'B' or '[FAIL]' not in (diagnosis or ''): + return [{'role': 'system', 'content': SKILL_GEN_SYSTEM_Q}, + {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=problem)}] + return [{'role': 'system', 'content': SKILL_GEN_SYSTEM}, + {'role': 'user', 'content': SKILL_GEN_USER_RUBRIC.format(problem=problem, diagnosis=diagnosis)}] + + +def _assign_view(problem: str, args: argparse.Namespace) -> str: + h = int(hashlib.md5(f'{args.seed}:{problem}'.encode('utf-8')).hexdigest(), 16) + return 'B' if (h % 100000) / 100000.0 < args.view_b_frac else 'A' + + +def _view_prompt(r: Dict[str, Any]) -> Dict[str, Any]: + return {'messages': _skillgen_messages(r['problem'], r['_view'], r.get('_rubric_diag', ''))} + + +_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') + + +def _clean_text(decoded: Optional[str]) -> str: + return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() + + +def _extract_skills_block(text: str) -> Optional[str]: + """Non-empty ``...`` block, or None. Requires ```` (skill-gen + runs thinking ON); reads only the answer after the last one, so a mid-reasoning draft + or a demo echo can never be mistaken for the answer. No format/wording gate beyond + that -- the reward (does the frozen executor solve a similar problem with this skill?) + is what judges skill quality (SEAM-style), so we do not lexically second-guess it.""" + low = text.lower() + end_think = low.rfind('') + if end_think < 0: + return None + answer = text[end_think + len(''):] + low_a = answer.lower() + s = low_a.find('') + if s < 0: + return None + inner = s + len('') + e = low_a.find('', inner) + block = (answer[inner:e] if e >= 0 else answer[inner:]).strip() + block = re.sub(r'', '', block, flags=re.IGNORECASE).strip() + return block or None + + +def _parse_seq(seq, gold: str) -> Dict[str, Any]: + """Grade one sampled sequence into a rollout record.""" + text = _clean_text(getattr(seq, 'decoded', '') or '') + pred = extract_boxed(text) + correct = bool(pred) and answers_match(pred, gold) + terminated = getattr(seq, 'stop_reason', None) != 'length' + return {'pred': pred, 'correct': correct, 'terminated': terminated, + 'passed': bool(correct and terminated), + 'stop_reason': getattr(seq, 'stop_reason', None), + 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} + + +def _run_samples(sampler, prompts: List[Any], num_samples: int, max_tokens: int, + gen_dp: int, temperature: Optional[float] = None, + top_p: Optional[float] = None, top_k: Optional[int] = None) -> List[List[Any]]: + """One batched sampler call -> per-prompt list of raw sequences. vLLM dp needs + batch len >= dp, so pad the tail and slice back.""" + if not prompts: + return [] + params = SamplingParams( + max_tokens=max_tokens, + temperature=GEN_TEMPERATURE if temperature is None else temperature, + top_p=GEN_TOP_P if top_p is None else top_p, + num_samples=num_samples, **({} if top_k is None else {'top_k': top_k})) + padded = prompts + if gen_dp > 1 and 0 < len(prompts) < gen_dp: + padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] + responses = sampler.sample(padded, params)[:len(prompts)] + return [list(r.sequences) if (r and r.sequences) else [] for r in responses] + + +# =========================================================================== +# Block C -- data loading via twinkle.Dataset + numeric filtering +# =========================================================================== +def _boxed_batch(rows: Dict[str, List[Any]], dataset: str) -> Dict[str, List[Any]]: + """Batched (columnar) mapper for ``Dataset.map``: annotate each row with its boxed + ``reference_answer`` and a ``_keep`` flag (aops also needs the ``boxed`` metadata).""" + sols = rows['solution'] + metas = rows.get('metadata', [None] * len(sols)) + refs = [extract_boxed(s or '') for s in sols] + keep = [bool(ref) and (dataset != 'aops' or bool((meta or {}).get('boxed'))) + for ref, meta in zip(refs, metas)] + return {**rows, 'reference_answer': refs, '_keep': keep} + + +def load_problems(dataset: str, n: int, seed: int, num_proc: int = 0) -> List[Dict[str, Any]]: + """Load boxed-answer problems as ``{problem, reference_answer, level?}`` via + twinkle.Dataset (ModelScope hub), sampled to ``n`` (0 = all). Boxed extraction (regex + + brace scan over every solution) is parallelised by ``Dataset.map``/``.filter``; + ``num_proc`` defaults to all cores (set 1 to force serial).""" + ds_id = AOPS_DATASET_ID if dataset == 'aops' else MATH_DATASET_ID + ds = Dataset(DatasetMeta(dataset_id=f'ms://{ds_id}', split='train')) + nproc = num_proc if num_proc > 0 else (os.cpu_count() or 1) + ds.map(lambda rows: _boxed_batch(rows, dataset), num_proc=nproc) + ds.filter(lambda row: row['_keep'], num_proc=nproc) + has_level = 'level' in ds.dataset.column_names + out = [{'problem': row['problem'], 'reference_answer': row['reference_answer'], + **({'level': row['level']} if has_level and row.get('level') else {})} + for row in ds.dataset] + logger.info(f'[data] {dataset}: {len(out)} boxed problems ({nproc} procs)') + rng = np.random.RandomState(seed) + rng.shuffle(out) + return out[:n] if (n and n < len(out)) else out + + +_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') + + +def _norm_num_text(num: str) -> str: + try: + f = float(num) + return str(int(f)) if f == int(f) else str(f) + except Exception: + return str(num).strip() + + +def _numeric_value(raw: Any) -> Optional[str]: + """Collapse an answer to a single int/decimal/fraction, or None.""" + if raw is None: + return None + s = str(raw).strip().strip('$').strip() + s = s.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac') + s = re.sub(r'\\!|\\,|\\;|\\ |\\left|\\right|\s', '', s) + for pat in (r'\\frac\{(-?\d+)\}\{(-?\d+)\}', r'(-?\d+)/(-?\d+)'): + m = re.fullmatch(pat, s) + if m: + a, b = int(m.group(1)), int(m.group(2)) + return _norm_num_text(str(a / b)) if b else None + return _norm_num_text(s) if _NUM_RE.fullmatch(s) else None + + +def _answer_leaked(skill: str, reference: str) -> bool: + """Deterministic answer-leak check (replaces the LLM LeakVerifier). Flags ONLY the one + real way a skill can game the deterministic reward: writing the final answer verbatim. + On our numeric-only data the answer is a single number, so we require it as a + standalone token (digit boundaries) -- that avoids matching an intermediate value like + '3' inside '36'. Everything else (methods, plans, pitfalls) is left to the reward, the + way SEAM/POPE handle it (no content audit). Non-numeric answers -> not flagged.""" + if not skill: + return False + for cand in {_numeric_value(reference), (str(reference).strip() or None)}: + if cand and re.search(r'(? Tuple[List[Dict[str, Any]], List[Dict[str, Any]], Dict[str, int]]: + """Load, numeric-filter, shuffle, then split a fixed eval holdout off the front.""" + # Load all when filtering or splitting (else the eval holdout could starve train). + load_n = 0 if (args.numeric_only or args.eval_size > 0) else args.n + records = load_problems(args.dataset, load_n, args.seed) + raw_n, dropped = len(records), 0 + if args.numeric_only: + records = [{**r, 'reference_answer': v} + for r, v in ((r, _numeric_value(r.get('reference_answer'))) for r in records) + if v is not None] + dropped = raw_n - len(records) + np.random.RandomState(args.seed).shuffle(records) + eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 + eval_records = [dict(r) for r in records[:eval_n]] + pool = records[eval_n:] + train_n = args.n if args.n > 0 else len(pool) + train_records = [dict(r) for r in pool[:train_n]] + if {r['problem'] for r in train_records} & {r['problem'] for r in eval_records}: + raise ValueError('eval/train overlap detected') + stats = {'raw_loaded': raw_n, 'numeric_dropped': dropped, + 'train_records': len(train_records), 'eval_records': len(eval_records)} + return train_records, eval_records, stats + + +# =========================================================================== +# Block D -- disk cache, problem pool, baseline rollout, rubric check +# =========================================================================== +class DiskCache: + """Append-only jsonl kv cache (md5 key -> value). Loads on init, appends on put. + Disabled instances always miss and never write.""" + + def __init__(self, path: str, enabled: bool = True): + self._mem: Dict[str, Any] = {} + self._fh = None + self._lock = threading.Lock() # base baseline is prefetched on a background thread + if not enabled: + return + if os.path.exists(path): + with open(path, encoding='utf-8') as f: + for line in f: + if line.strip(): + row = json.loads(line) + self._mem[row['key']] = row['value'] + self._fh = open(path, 'a', encoding='utf-8') + + @staticmethod + def key_for(*parts: str) -> str: + return hashlib.md5('\x1f'.join(parts).encode('utf-8')).hexdigest() + + def __contains__(self, key: str) -> bool: + with self._lock: + return key in self._mem + + def get(self, key: str) -> Any: + with self._lock: + return self._mem.get(key) + + def put(self, key: str, value: Any) -> None: + with self._lock: + self._mem[key] = value + if self._fh is not None: + self._fh.write(json.dumps({'key': key, 'value': value}, ensure_ascii=False) + '\n') + self._fh.flush() + + def close(self) -> None: + if self._fh is not None: + self._fh.close() + + +class _LockedSampler: + """Serialises every ``.sample()`` on a sampler behind one lock. The base sampler is + shared by the main thread (with-skill scoring, eval) and the baseline-prefetch thread; + ``sample`` is a slice_dp/flatten remote call whose collect is NOT safe to interleave + across two callers, so concurrent calls could mis-join sequences. The lock keeps base + calls serial (prefetch still overlaps the skill-gen phase, which uses skill_sampler).""" + + def __init__(self, sampler): + self._sampler = sampler + self._lock = threading.Lock() + + def sample(self, *args, **kwargs): + with self._lock: + return self._sampler.sample(*args, **kwargs) + + def __getattr__(self, name): + return getattr(self._sampler, name) + + +class ProblemPool: + """Cyclic draw source. Each full pass reshuffles with ``seed + epoch``; the initial + pass keeps the loader order. ``draw(k)`` returns k distinct problems (k << pool).""" + + def __init__(self, records: List[Dict[str, Any]], seed: int): + self._records = list(records) + self._seed, self._cursor, self.epoch = seed, 0, 0 + + def draw(self, k: int) -> List[Dict[str, Any]]: + out, seen = [], set() + while len(out) < k: + if self._cursor >= len(self._records): + self.epoch += 1 + np.random.RandomState(self._seed + self.epoch).shuffle(self._records) + self._cursor = 0 + r = self._records[self._cursor] + self._cursor += 1 + if id(r) not in seen: + seen.add(id(r)) + out.append(r) + return out + + def peek(self, k: int) -> List[Dict[str, Any]]: + """The next k distinct problems draw() would return, WITHOUT advancing state + (no cursor move, no reshuffle). Used to prefetch their baseline into base_cache + while the current chunk trains -- a wrong guess (cross-epoch reshuffle) only + misses the cache, never corrupts the draw.""" + out, seen, cur = [], set(), self._cursor + recs = self._records + while len(out) < k and cur < len(recs): # stop at epoch edge; don't simulate reshuffle + r = recs[cur] + cur += 1 + if id(r) not in seen: + seen.add(id(r)) + out.append(r) + return out + + +def _empty_roll() -> Dict[str, Any]: + return {'pred': '', 'correct': False, 'terminated': False, 'passed': False, + 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} + + +def _apply_baseline(r: Dict[str, Any], roll: Dict[str, Any]) -> None: + """Attach a greedy baseline roll and reset per-chunk working state.""" + r['_baseline_rolls'], r['_cands'], r['_init'] = [roll], [], [roll] + r['_failed'] = not roll['correct'] + r['_baseline_pass'] = 1.0 if roll['correct'] else 0.0 + r['_hard'] = True # process every problem; group variance selects (SEAM-style) + + +def baseline_rollout(base_sampler, problems: List[Dict[str, Any]], base_dp: int, + args: argparse.Namespace, cache: DiskCache) -> int: + """Base solves each problem greedily once (T=0, M=1), disk-cached by problem text. + The base is frozen + greedy so a cache hit is exact. Returns fresh (miss) count.""" + todo = [r for r in problems if DiskCache.key_for(r['problem']) not in cache] + if todo: + out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], + 1, args.max_tokens, base_dp, temperature=0.0) + for r, seqs in zip(todo, out): + roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + cache.put(DiskCache.key_for(r['problem']), roll) + for r in problems: + _apply_baseline(r, cache.get(DiskCache.key_for(r['problem']))) + return len(todo) + + +# -- rubric process-check (view A): teacher diagnoses the base's attempt -- +_RFT_DIAG_SYSTEM = """\ +You are a process error checker for a math solution attempt. You are given a math +problem, a rubric, and one attempted solution segment. Decide PASS or FAIL for each +criterion and explain only the process error type. + +Output STRICT JSON (no prose outside it) with this shape: +{ + "items": [ + {"index": 1, "verdict": "PASS", "reason": "", + "fix": ""}, + {"index": 2, "verdict": "FAIL", "reason": "", + "fix": ""} + ], + "overall": "OK" | "ISSUES", + "summary": "" +} + +Rules: +- Judge every criterion independently and literally; a [Hard Rule] is FAIL unless + unambiguously satisfied. +- Judge ONLY what is observable in THIS segment. +- Content inside ... (or ) is internal reasoning, not + user-facing output; ignore it for "output only X" style criteria. +- For PASS items, leave "fix" as "". +- For FAIL items, "reason", "fix", and "summary" must describe only the flawed + step, theorem, arithmetic operation, case split, or verification habit. +- NEVER state the correct final answer, corrected final expression, option letter, + graph/choice label, or any exact value that the answer should become. +- NEVER write phrases like "the correct answer is", "which gives", "yielding", + "should be ", "Option ", or "Graph ". +- If a fix would require naming a corrected value, replace it with a method-level + instruction such as "redo that computation carefully" or "apply the theorem with + the correct quantities". +- Keep every "reason" and "fix" clear and concise — one short sentence each. +- "overall" is "OK" only if NO criterion is FAIL. +- Output only the JSON object.""" + +_RFT_DIAG_USER = """\ +## Task / query (context) +{query} + +## Rubric +{rubric} + +## Segment +{segment} + +Now output the diagnostic JSON object.""" + +_MATH_RUBRIC = [ + ('The reasoning contains no arithmetic or algebraic error', True), + ('Each step follows logically from the previous ones', True), + ('No formula or theorem is misstated or misapplied', True), + ('The approach is on track to answer the actual question asked', False), + ('No step contradicts an earlier established fact', False), +] + + +class _RftRubricVerifier(RubricVerifier): + def _diagnose_trajectory(self, query: str, rubric_block: str, segment_text: str) -> dict: + return {'messages': [ + {'role': 'system', 'content': _RFT_DIAG_SYSTEM}, + {'role': 'user', 'content': _RFT_DIAG_USER.format( + query=query, rubric=rubric_block, segment=segment_text)}]} + + +def build_rubric_checker() -> Optional[RubricVerifier]: + """Fixed math-process rubric verifier (teacher-served). None if no LLM backup env.""" + if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') + or os.environ.get('OPENAI_API_KEY')): + return None + return _RftRubricVerifier( + fixed_rubric=[RubricItem(t, is_hard=h) for t, h in _MATH_RUBRIC], gate=True) + + +def _format_diagnosis(detail) -> str: + """One line per criterion (PASS/FAIL + reason + fix) then a summary.""" + rub = detail.rubric + lines = [] + for it in detail.items: + text = rub[it.index - 1].text if 0 < it.index <= len(rub) else f'criterion {it.index}' + if it.verdict: + lines.append(f'- [PASS] {text}') + else: + tail = f': {it.reason}' if it.reason else '' + tail += f' (fix: {it.fix})' if it.fix else '' + lines.append(f'- [FAIL] {text}{tail}') + if detail.summary: + lines.append(f'Summary: {detail.summary}') + return '\n'.join(lines) + + +def diagnose_views(checker, problems: List[Dict[str, Any]], args: argparse.Namespace, + cache: DiskCache) -> None: + """Rubric-check every view-A problem's greedy attempt in parallel (disk-cached by + problem + attempt), stashing the formatted findings on ``r['_rubric_diag']``.""" + targets = [r for r in problems if r.get('_view') == 'A'] + if not checker or not targets: + return + + def _key(r: Dict[str, Any]) -> str: + return DiskCache.key_for(r['problem'], r.get('_init', [{}])[0].get('text', '')) + + pending = [] + for r in targets: + key = _key(r) + if key in cache: + r['_rubric_diag'] = cache.get(key) + else: + pending.append((r, key)) + if not pending: + return + + def _run(item): + r, key = item + seg = {'messages': [{'role': 'user', 'content': r['problem']}, + {'role': 'assistant', 'content': r['_init'][0]['text']}]} + try: + return r, key, _format_diagnosis(checker.diagnose(seg, query=r['problem'])) + except Exception as exc: # teacher hiccup -> no-diagnosis prompt (not cached) + logger.warning(f'[rubric] diagnose error: {exc}') + return r, key, None + + workers = max(1, min(args.rubric_workers, len(pending))) + with ThreadPoolExecutor(max_workers=workers) as ex: + for r, key, diag in ex.map(_run, pending): + r['_rubric_diag'] = diag or '' + if diag is not None: + cache.put(key, diag) + + +# =========================================================================== +# Block E -- chunk draw, generation pipeline, record building +# =========================================================================== +def _baseline_class(r: Dict[str, Any]) -> str: + """success | fail_loop (out of length / never terminated) | fail_wrong.""" + roll = r['_init'][0] + if roll['correct']: + return 'success' + return 'fail_loop' if (roll['stop_reason'] == 'length' or not roll['terminated']) else 'fail_wrong' + + +def _select_balanced(buckets: Dict[str, List[Dict[str, Any]]], n_success: int, + n_fail: int, n_fail_loop: int) -> List[Dict[str, Any]]: + """Pick n_fail base-fails (toward n_fail_loop loop-fails, best-effort) + n_success + base-successes; top up any shortfall from leftovers.""" + loop, wrong, succ = buckets['fail_loop'], buckets['fail_wrong'], buckets['success'] + take_loop = min(n_fail_loop, len(loop)) + take_wrong = min(n_fail - take_loop, len(wrong)) + take_loop = min(n_fail - take_wrong, len(loop)) + sel = loop[:take_loop] + wrong[:take_wrong] + succ[:n_success] + target = n_success + n_fail + if len(sel) < target: + used = {id(x) for x in sel} + sel += [x for b in (loop, wrong, succ) for x in b if id(x) not in used][:target - len(sel)] + return sel + + +def draw_chunk(pool: ProblemPool, base_sampler, base_dp: int, args: argparse.Namespace, + cache: DiskCache) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + """Draw one chunk, baselining every drawn problem. With ``--balance``, keep + drawing+baselining until the target base fail:success mix is reachable (or the budget + is hit), then select a balanced subset.""" + if not args.balance: + chunk = pool.draw(args.chunk_size) + fresh = baseline_rollout(base_sampler, chunk, base_dp, args, cache) + return chunk, {'enabled': False, 'n_drawn': len(chunk), 'n_baseline_fresh': fresh} + + n_success = max(0, min(args.chunk_size, round(args.chunk_size * args.balance_success_frac))) + n_fail = args.chunk_size - n_success + n_fail_loop = round(n_fail * args.balance_loop_frac) + buckets: Dict[str, List[Dict[str, Any]]] = {'success': [], 'fail_loop': [], 'fail_wrong': []} + budget, n_drawn, n_fresh, seen = args.chunk_size * args.balance_max_draws_mult, 0, 0, set() + while n_drawn < budget: + if (len(buckets['success']) >= n_success + and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail): + break + batch = pool.draw(args.chunk_size) + n_fresh += baseline_rollout(base_sampler, batch, base_dp, args, cache) + n_drawn += len(batch) + for r in batch: + if id(r) not in seen: + seen.add(id(r)) + buckets[_baseline_class(r)].append(r) + + reached = (len(buckets['success']) >= n_success + and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail) + chunk = _select_balanced(buckets, n_success, n_fail, n_fail_loop) + sel_success = sum(1 for r in chunk if not r['_failed']) + stats = { + 'enabled': True, 'n_drawn': n_drawn, 'n_baseline_fresh': n_fresh, 'n_selected': len(chunk), + 'target_success': n_success, 'target_fail': n_fail, 'target_fail_loop': n_fail_loop, + 'selected_success': sel_success, 'selected_fail': len(chunk) - sel_success, + 'selected_fail_loop': sum(1 for r in chunk if _baseline_class(r) == 'fail_loop'), + 'selected_fail_wrong': sum(1 for r in chunk if _baseline_class(r) == 'fail_wrong'), + 'selected_success_frac': (sel_success / len(chunk)) if chunk else 0.0, + 'budget_hit': not reached, + } + return chunk, stats + + +def _assign_advantages(hard: List[Dict[str, Any]], args: argparse.Namespace) -> None: + """Group-relative advantage over each problem's scored candidates using the greedy + binary reward R in {0,1}: ``A = (R - mean) / (std + eps)``. std==0 groups get no + gradient -- GRPO's variance selects informative problems (no explicit difficulty gate).""" + eps = 1e-6 + for r in hard: + for c in r['_cands']: + c['advantage'], c['grpo_adv'], c['kept'] = 0.0, 0.0, False + cs = ([c for c in r['_cands'] if c.get('reward') is not None] if args.format_in_reward + else [c for c in r['_cands'] if c['leaked'] is False and c.get('reward') is not None]) + if len(cs) < 2: + continue + rewards = [c['reward'] for c in cs] + mean_r = sum(rewards) / len(rewards) + std = (sum((x - mean_r) ** 2 for x in rewards) / len(rewards)) ** 0.5 + if std < 1e-9: + continue + for c in cs: + adv = (c['reward'] - mean_r) / (std + eps) + c['advantage'], c['grpo_adv'], c['kept'] = adv, adv, c['reward'] > mean_r + + +def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], + ci: int, base_dp: int, skill_dp: int, args: argparse.Namespace, + checker, rubric_cache: DiskCache + ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], List[Dict[str, Any]]]: + """view assign -> rubric-check (view A) -> skill-gen -> leak-filter -> with-skill + greedy pass -> GRPO advantages. ``chunk`` arrives already baselined by draw_chunk.""" + hard = chunk + for r in hard: + r['_view'], r['_rubric_diag'] = _assign_view(r['problem'], args), '' + diagnose_views(checker, hard, args, rubric_cache) + + # skill-gen (thinking ON), per-view prompt; re-sample problems with no clean candidate. + flat: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] + pending = list(hard) + for _ in range(args.skill_retries + 1): + if not pending: + break + sg_out = _run_samples(skill_sampler, [_view_prompt(r) for r in pending], + args.n_skills, args.skill_max_tokens, skill_dp, + temperature=args.skill_gen_temperature, + top_p=args.skill_gen_top_p, top_k=args.skill_gen_top_k) + still = [] + for r, seqs in zip(pending, sg_out): + got = False + for s in seqs: + resp = _clean_text(getattr(s, 'decoded', '') or '') + block = _extract_skills_block(resp) + cand = {'skills': block or '', 'response': resp, 'parseable': bool(block), + 'view': r['_view'], 'leaked': None, 'leak_reason': '', + 'leak_source': '', 'with_pass': None, 'reward': None, 'rolls': []} + r['_cands'].append(cand) + if block: + flat.append((r, cand)) + got = True + if not got: + still.append(r) + pending = still + + # leak filter: deterministic verbatim-answer check only (no LLM). Measured on prior + # runs, the LLM judge flagged ~2% and leak rate was view-independent, so we drop the + # teacher-served LeakVerifier and just block the one exploit the reward can't catch -- + # the final answer written into the skill (see _answer_leaked). + for r, c in flat: + leaked = _answer_leaked(c['skills'], r['reference_answer']) + c['leaked'] = leaked + c['leak_reason'] = 'answer_verbatim' if leaked else '' + c['leak_source'] = 'deterministic' + + # with-skill greedy pass (T=0, M=1); reward = correct, absolute (group mean is baseline). + clean = [(r, c) for r, c in flat if c['leaked'] is False] + if clean: + ws_out = _run_samples(base_sampler, + [build_skill_solve_prompt(r['problem'], c['skills']) for r, c in clean], + 1, args.max_tokens, base_dp, temperature=0.0) + for (r, c), seqs in zip(clean, ws_out): + c['rolls'] = [_parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll()] + c['with_pass'] = 1.0 if c['rolls'][0]['correct'] else 0.0 + c['reward'] = c['with_pass'] + if args.format_in_reward: # unparseable/leaked score 0 and still join the group + for r in hard: + for c in r['_cands']: + if c['reward'] is None: + c['reward'] = 0.0 + + _assign_advantages(hard, args) + return ([_full_record(r, ci) for r in chunk], _chunk_summary(chunk, ci, args), + _group_records(chunk, args)) + + +def _roll(x: Dict[str, Any]) -> Dict[str, Any]: + return {k: x[k] for k in ('pred', 'correct', 'terminated', 'passed', + 'stop_reason', 'gen_tokens', 'text')} + + +def _is_trainable(c: Dict[str, Any], args: argparse.Namespace) -> bool: + """Reaches the GRPO update iff advantage is non-zero (and, without --format-in-reward, + also clean and scored).""" + adv_nz = abs(c.get('advantage') or 0.0) > 1e-9 + if args.format_in_reward: + return adv_nz + return c['leaked'] is False and c.get('with_pass') is not None and adv_nz + + +def _full_record(r: Dict[str, Any], ci: int) -> Dict[str, Any]: + """Complete per-problem trace: init attempt, baseline, and all candidates.""" + init = r['_init'][0] + return { + 'record_type': 'problem', 'chunk': ci, 'problem': r['problem'], + 'reference_answer': r['reference_answer'], 'level': r.get('level', ''), + 'failed_first_try': r['_failed'], + 'init_attempt': {'text': init['text'], 'pred': init['pred'], 'correct': init['correct'], + 'terminated': init['terminated'], 'stop_reason': init['stop_reason'], + 'gen_tokens': init['gen_tokens']}, + 'baseline_pass': r['_baseline_pass'], 'is_hard': r['_hard'], + 'view': r.get('_view', ''), 'rubric_diag': r.get('_rubric_diag', ''), + 'baseline_rolls': [_roll(x) for x in r['_baseline_rolls']], + 'candidates': [{ + 'skills': c['skills'], 'response': c['response'], 'parseable': c['parseable'], + 'leaked': c['leaked'], 'leak_reason': c['leak_reason'], 'leak_source': c['leak_source'], + 'with_pass': c['with_pass'], 'reward': c.get('reward'), 'advantage': c.get('advantage'), + 'grpo_adv': c.get('grpo_adv'), 'kept': c.get('kept'), + 'rolls': [_roll(x) for x in c['rolls']], + } for c in r['_cands']], + } + + +def _view_stats(hard: List[Dict[str, Any]], view: str) -> Dict[str, Any]: + hv = [r for r in hard if r.get('_view') == view] + cands = [c for r in hv for c in r['_cands'] if c['parseable']] + clean = [c for c in cands if c['leaked'] is False] + adopted = sum(1 for r in hv + if any(c['leaked'] is False and abs(c.get('advantage') or 0.0) > 1e-9 + for c in r['_cands'])) + return {'n_hard': len(hv), 'n_candidates_parseable': len(cands), 'n_clean': len(clean), + 'n_adopted_problems': adopted, 'adoption_rate': (adopted / len(hv)) if hv else 0.0} + + +def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespace) -> Dict[str, Any]: + hard = [r for r in chunk if r['_hard']] + all_cands = [c for r in chunk for c in r['_cands']] + cands = [c for c in all_cands if c['parseable']] + scored = [c for c in cands if c['with_pass'] is not None] + ws_rolls = [x for c in scored for x in c['rolls']] + fail_cands = [c for r in chunk if r['_failed'] for c in r['_cands']] + base_acc = (sum(r['_baseline_pass'] for r in hard) / len(hard)) if hard else 0.0 + ws_acc = (sum(c['with_pass'] for c in scored) / len(scored)) if scored else 0.0 + abs_adv = lambda cs: sum(abs(c.get('advantage') or 0.0) for c in cs) + total_abs = abs_adv(all_cands) + return { + 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), + 'n_failed_first_try': sum(1 for r in chunk if r['_failed']), 'n_hard': len(hard), + 'n_generated': len(all_cands), 'n_candidates_parseable': len(cands), + 'n_unparseable': len(all_cands) - len(cands), + 'n_leaked': sum(1 for c in cands if c['leaked']), + 'n_clean': sum(1 for c in cands if c['leaked'] is False), + 'n_reward_pos': sum(1 for c in scored if c['reward']), + 'n_train_samples': sum(1 for c in all_cands if _is_trainable(c, args)), + 'n_train_from_fail': sum(1 for c in fail_cands if _is_trainable(c, args)), + 'abs_adv_from_fail_frac': (abs_adv(fail_cands) / total_abs) if total_abs > 0 else 0.0, + 'avg_baseline_pass_on_hard': base_acc, 'avg_withskill_pass': ws_acc, + 'avg_lift': ws_acc - base_acc, + 'termination_rate_withskill': (sum(1 for x in ws_rolls if x['terminated']) / len(ws_rolls)) if ws_rolls else 0.0, + 'view_A': _view_stats(hard, 'A'), 'view_B': _view_stats(hard, 'B'), + } + + +def _group_records(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> List[Dict[str, Any]]: + """GRPO training records: every trainable candidate with its view + rubric diagnosis + (the prompt is rebuilt from those by ``_skillgen_messages``, no trajectory stored).""" + out = [] + for r in chunk: + if not r['_hard']: + continue + for c in r['_cands']: + if _is_trainable(c, args): + out.append({ + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'view': r.get('_view', 'A'), 'diagnosis': r.get('_rubric_diag', ''), + 'response': c['response'], 'skills': c['skills'], + 'advantage': c['advantage'], 'grpo_adv': c['grpo_adv'], 'kept': c['kept'], + 'reward': c['reward'], 'with_pass': c['with_pass']}) + return out + + +# =========================================================================== +# Block G -- online GRPO training +# =========================================================================== +def _is_num(v: Any) -> bool: + try: + float(v) + return True + except (TypeError, ValueError): + return False + + +def _train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: + """Training sample = the exact skill-gen prompt (rebuilt by ``_skillgen_messages`` so + train/inference match) + the generated (think + skills) response. ``key_rounds`` + selects the final assistant turn; Template masks the prompt and trains the whole + response (the key-round prefix already excludes the prompt-provided ).""" + msgs = _skillgen_messages(rec['problem'], rec.get('view', 'A'), rec.get('diagnosis', '')) + return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], + 'user_data': {'key_rounds': [len(msgs)]}} + + +def _train_chunk(skill_model, ckpt: CheckpointEngineManager, samples: List[Dict[str, Any]], + args: argparse.Namespace) -> Dict[str, Any]: + """One on-policy GRPO optimizer step on this chunk, then sync weights. Skills were + sampled from the current policy (ratio ~1, no old_logps); all micro-batches accumulate + into one step. The batch is padded to a multiple of ``sft_batch_size`` with + advantage-0 copies that contribute no gradient.""" + trajs = [_train_trajectory(rec) for rec in samples] + advs = [float(rec['advantage']) for rec in samples] + rem = (-len(trajs)) % args.sft_batch_size + if rem: + trajs += [trajs[-1]] * rem + advs += [0.0] * rem + micro = 0 + for i in range(0, len(trajs), args.sft_batch_size): + skill_model.forward_backward(inputs=trajs[i:i + args.sft_batch_size], + advantages=advs[i:i + args.sft_batch_size]) + micro += 1 + skill_model.clip_grad_and_step() + ckpt.sync_weights(merge_and_sync=True) + metric = skill_model.calculate_metric(is_training=True) + return {'n_samples': len(samples), 'n_steps': 1, 'n_micro_batches': micro, + 'metric': {k: (float(v) if _is_num(v) else v) for k, v in (metric or {}).items()}} + + +# =========================================================================== +# Block H -- fixed-holdout eval + metric formatting +# =========================================================================== +def run_greedy_eval(base_sampler, skill_sampler, eval_records: List[Dict[str, Any]], + ci: int, rounds: int, base_dp: int, skill_dp: int, args: argparse.Namespace, + checker, base_cache: DiskCache, rubric_cache: DiskCache + ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], Dict[str, float]]: + """SEAM ``val-core/math/acc/mean@1`` on the fixed holdout: ONE greedy skill (T=0) per + problem into ONE greedy base solve (T=0). Each problem keeps its view; no leak filter + (acc scores correctness alone). Baseline + rubric reuse the disk caches.""" + baseline_rollout(base_sampler, eval_records, base_dp, args, base_cache) + for r in eval_records: + r['_view'], r['_rubric_diag'] = _assign_view(r['problem'], args), '' + diagnose_views(checker, eval_records, args, rubric_cache) + sg_out = _run_samples(skill_sampler, [_view_prompt(r) for r in eval_records], + 1, args.skill_max_tokens, skill_dp, temperature=0.0) + skills = [(_extract_skills_block(_clean_text(getattr(seqs[0], 'decoded', '') or '')) or '', + _clean_text(getattr(seqs[0], 'decoded', '') or '')) if seqs else ('', '') + for seqs in sg_out] + ws_out = _run_samples(base_sampler, + [build_skill_solve_prompt(r['problem'], sk) for r, (sk, _) in zip(eval_records, skills)], + 1, args.max_tokens, base_dp, temperature=0.0) + recs = [] + for r, (sk, sresp), seqs in zip(eval_records, skills, ws_out): + roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + recs.append({ + 'record_type': 'eval_problem', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'view': r['_view'], 'rubric_diag': r.get('_rubric_diag', ''), + 'baseline_pass': r['_baseline_pass'], 'skill': sk, 'skill_parseable': bool(sk), + 'skill_response': sresp, 'withskill_pred': roll['pred'], + 'withskill_correct': roll['correct'], 'withskill_terminated': roll['terminated'], + 'withskill_stop_reason': roll['stop_reason'], 'withskill_text': roll['text'], + }) + acc = lambda rs: sum(1 for x in rs if x['withskill_correct']) / len(rs) if rs else 0.0 + A = [x for x in recs if x['view'] == 'A'] + B = [x for x in recs if x['view'] == 'B'] + ws = acc(recs) + base = sum(x['baseline_pass'] for x in recs) / len(recs) if recs else 0.0 + fmt = (sum(1 for x in recs if x['skill_parseable']) / len(recs)) if recs else 0.0 + term = (sum(1 for x in recs if x['withskill_terminated']) / len(recs)) if recs else 0.0 + summary = {'record_type': 'eval_summary', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, + 'n': len(recs), 'n_A': len(A), 'n_B': len(B), 'acc_mean1': ws, + 'baseline_acc_mean1': base, 'lift_mean1': ws - base, + 'acc_A_mean1': acc(A), 'acc_B_mean1': acc(B), 'format_mean1': fmt, 'term_mean1': term} + metrics = {'core/math/acc/mean@1': ws, 'core/math/baseline_acc/mean@1': base, + 'core/math/lift/mean@1': ws - base, 'core/math/format/mean@1': fmt, + 'core/math/term/mean@1': term} + if A: + metrics['core/math/acc_A/mean@1'] = acc(A) + if B: + metrics['core/math/acc_B/mean@1'] = acc(B) + return recs, summary, metrics + + +def _trend_line(hist: List[Dict[str, float]], window: int, rounds_done: int) -> Optional[str]: + """Contrast the FIRST vs the most recent ``window`` chunks -- if RFT works, adoption + and lift on recent (fresh) chunks exceed the early baseline.""" + if len(hist) < 2 * window: + return None + base, rec = hist[:window], hist[-window:] + m = lambda xs, k: sum(h[k] for h in xs) / len(xs) + return (f'[trend] first {window} vs last {window} | ' + f'adopt A {m(base,"aA"):.2f}->{m(rec,"aA"):.2f} B {m(base,"aB"):.2f}->{m(rec,"aB"):.2f} | ' + f'lift {m(base,"lift"):+.3f}->{m(rec,"lift"):+.3f} | ' + f'pos/chunk {m(base,"pos"):.1f}->{m(rec,"pos"):.1f} | rounds={rounds_done}') + + +def _swan_metrics(summary: Dict[str, Any], log: Optional[Dict[str, Any]]) -> Dict[str, float]: + """Flat swanlab dict = generation metrics + (when trained) the GRPO built-in metric. + acc/adopt/term only on chunks with hard problems, so idle chunks don't dip charts.""" + d: Dict[str, float] = { + 'gen/n_hard': summary['n_hard'], 'gen/n_clean': summary['n_clean'], + 'gen/n_leaked': summary['n_leaked'], 'gen/n_train_samples': summary['n_train_samples'], + 'gen/n_reward_pos': summary['n_reward_pos'], + 'gen/n_train_from_fail': summary['n_train_from_fail'], + 'gen/abs_adv_from_fail_frac': summary['abs_adv_from_fail_frac'], + } + bal = summary.get('balance') or {} + if bal.get('enabled'): + d.update({'balance/n_drawn': bal['n_drawn'], 'balance/n_baseline_fresh': bal['n_baseline_fresh'], + 'balance/selected_success_frac': bal['selected_success_frac']}) + if summary['n_hard'] > 0: + d.update({'acc/baseline_pass': summary['avg_baseline_pass_on_hard'], + 'acc/withskill_pass': summary['avg_withskill_pass'], 'acc/lift': summary['avg_lift'], + 'adopt/A': summary['view_A']['adoption_rate'], + 'adopt/B': summary['view_B']['adoption_rate'], + 'term/withskill': summary['termination_rate_withskill']}) + if log: + d['train/n_steps'] = log['n_steps'] + d['train/n_micro_batches'] = log['n_micro_batches'] + for k, v in (log.get('metric') or {}).items(): + if not _is_num(v): + continue + if k.startswith('learning rate'): + if 'group 1' in k: + d['train/lr'] = float(v) + else: + d[f'train/{k.replace(" ", "_")}'] = float(v) + return d + + +# =========================================================================== +# Block F -- components, args, main +# =========================================================================== +def init_components(args: argparse.Namespace): + """8 GPUs: ranks 0-3 train (FSDP2), 4-5 skill_sampler (synced), 6-7 base_sampler + (frozen). Returns (skill_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp).""" + r0, r1, r2 = TRAIN_GPUS, TRAIN_GPUS + SKILL_SAMPLER_GPUS, NUM_GPUS + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, lazy_collect=False, groups=[ + DeviceGroup(name='train', ranks=list(range(0, r0)), device_type='GPU'), + DeviceGroup(name='skill_sampler', ranks=list(range(r0, r1)), device_type='GPU'), + DeviceGroup(name='base_sampler', ranks=list(range(r1, r2)), device_type='GPU')]) + + train_mesh = DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_DP, fsdp_size=TRAIN_FSDP) + skill_model = TransformersModel(model_id=MODEL_ID, device_mesh=train_mesh, remote_group='train', + ddp_config={'find_unused_parameters': False}) + skill_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) + skill_model.set_template(Template, model_id=MODEL_ID, enable_thinking=True, + max_length=args.max_model_len, truncation_strategy='delete') + skill_model.set_processor(InputProcessor, padding_free=False) + skill_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) + skill_model.set_optimizer('AdamW', lr=args.lr) + skill_model.set_lr_scheduler('CosineWarmupScheduler', num_warmup_steps=10, + num_training_steps=args.max_train_rounds) + + def _sampler(group, world): + s = vLLMSampler(model_id=MODEL_ID, + engine_args={'gpu_memory_utilization': GPU_MEM, + 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, + device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), + remote_group=group) + s.set_template(Template, model_id=MODEL_ID, enable_thinking=True, max_length=args.max_model_len) + return s + + skill_sampler = _sampler('skill_sampler', SKILL_SAMPLER_GPUS) + # base_sampler is shared by the main thread and the baseline-prefetch thread -> serialise. + base_sampler = _LockedSampler(_sampler('base_sampler', BASE_SAMPLER_GPUS)) + ckpt = CheckpointEngineManager(model=skill_model, sampler=skill_sampler) + return skill_model, skill_sampler, base_sampler, ckpt, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS + + +def _build_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--dataset', choices=('aops', 'math'), default='aops') + p.add_argument('--n', type=int, default=2000, help='Problems loaded into the draw pool.') + p.add_argument('--seed', type=int, default=42) + p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--eval-size', type=int, default=128, help='Fixed holdout size (0 disables).') + p.add_argument('--eval-every', type=int, default=10, help='Run holdout eval every N chunks.') + p.add_argument('--chunk-size', type=int, default=16) + p.add_argument('--balance', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--balance-success-frac', type=float, default=0.4, + help='Target fraction of the chunk the base solves (rest are base-fail).') + p.add_argument('--balance-loop-frac', type=float, default=0.5) + p.add_argument('--balance-max-draws-mult', type=int, default=8) + p.add_argument('--n-skills', type=int, default=8) + p.add_argument('--view-b-frac', type=float, default=0.5) + p.add_argument('--skill-retries', type=int, default=2) + p.add_argument('--skill-gen-temperature', type=float, default=1.0) + p.add_argument('--skill-gen-top-p', type=float, default=1.0) + p.add_argument('--skill-gen-top-k', type=int, default=-1) + p.add_argument('--max-model-len', type=int, default=16384) + p.add_argument('--max-tokens', type=int, default=8192) + p.add_argument('--skill-max-tokens', type=int, default=8192) + p.add_argument('--rubric-workers', type=int, default=16) + p.add_argument('--sft-batch-size', type=int, default=8, + help='Driver micro-batch before the chunk optimizer step; multiple of train dp.') + p.add_argument('--grpo-epsilon', type=float, default=0.2) + p.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--lr', type=float, default=6e-6) + p.add_argument('--max-train-rounds', type=int, default=1500) + p.add_argument('--save-rounds', type=int, default=50) + p.add_argument('--trend-every', type=int, default=10) + p.add_argument('--output-dir', default='./output/reflexion_skill_rft') + p.add_argument('--cache-dir', default='', help='Baseline/rubric cache dir (default /cache).') + p.add_argument('--no-cache', action='store_true', help='Disable disk cache read/write.') + p.add_argument('--prefetch-baseline', action=argparse.BooleanOptionalAction, default=True, + help='Prefetch next chunk base baseline on a background thread (overlaps ' + 'with skill-gen; base_sampler is frozen so it never blocks the trainer).') + p.add_argument('--swanlab-project', default='twinkle') + p.add_argument('--swanlab-exp', default='') + args = p.parse_args() + if args.sft_batch_size % TRAIN_DP != 0: + raise ValueError(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple of train dp ({TRAIN_DP})') + if args.chunk_size < 1: + raise ValueError('--chunk-size must be >= 1') + args.balance_success_frac = max(0.0, min(1.0, args.balance_success_frac)) + return args + + +def _write(handle, row: Dict[str, Any]) -> None: + handle.write(json.dumps(row, ensure_ascii=False) + '\n') + + +def main() -> None: + args = _build_args() + records, eval_records, data_stats = _load_records(args) + if len(records) < args.chunk_size: + raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded ({len(records)}); raise --n') + + os.makedirs(args.output_dir, exist_ok=True) + gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') + eval_path = os.path.join(args.output_dir, 'eval_records.jsonl') + data_path = os.path.join(args.output_dir, 'skill_dataset.jsonl') + train_log_path = os.path.join(args.output_dir, 'train_log.jsonl') + + if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY')): + sys.stderr.write('[rft] WARNING: no LLM backup env; view-A rubric check disabled ' + '(leak filter is deterministic, unaffected)\n') + + use_swan = swanlab is not None and os.environ.get('SWANLAB_MODE') != 'disabled' + if use_swan: + swanlab.init(project=args.swanlab_project, experiment_name=(args.swanlab_exp or None), + config={'model': MODEL_ID, 'dataset': args.dataset, 'n': len(records), + 'eval_n': len(eval_records), 'n_skills': args.n_skills, + 'view_b_frac': args.view_b_frac, 'balance': args.balance, + 'balance_success_frac': args.balance_success_frac, + 'skill_gen_temp': args.skill_gen_temperature, + 'grpo_epsilon': args.grpo_epsilon, 'lr': args.lr}) + + skill_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp = init_components(args) + checker = build_rubric_checker() + if checker is None: + sys.stderr.write('[rft] no LLM backup env -> view-A rubric process-check DISABLED\n') + + cache_dir = args.cache_dir or os.path.join(args.output_dir, 'cache') + os.makedirs(cache_dir, exist_ok=True) + use_cache = not args.no_cache + base_cache = DiskCache(os.path.join(cache_dir, 'baseline.jsonl'), use_cache) + eval_base_cache = DiskCache(os.path.join(cache_dir, 'eval_baseline.jsonl'), use_cache) + rubric_cache = DiskCache(os.path.join(cache_dir, 'rubric.jsonl'), use_cache) + + cfg = {'record_type': 'config', 'model': MODEL_ID, 'dataset': args.dataset, + 'n': len(records), 'eval_n': len(eval_records), 'seed': args.seed, + 'numeric_only': args.numeric_only, 'raw_loaded': data_stats['raw_loaded'], + 'numeric_dropped': data_stats['numeric_dropped'], 'eval_every': args.eval_every, + 'n_skills': args.n_skills, 'view_b_frac': args.view_b_frac, + 'skill_retries': args.skill_retries, 'balance': args.balance, + 'balance_success_frac': args.balance_success_frac, + 'skill_gen_temp': args.skill_gen_temperature, 'reward': 'greedy_binary(correct)', + 'advantage': 'group_relative', 'format_in_reward': args.format_in_reward, 'cache': use_cache, + 'rubric_check': 'fixed_math_5crit(viewA)' if checker else 'disabled', + 'grpo_epsilon': args.grpo_epsilon, 'lr': args.lr, + 'max_train_rounds': args.max_train_rounds, 'started': int(time.time())} + sys.stderr.write(f'[rft] raw={data_stats["raw_loaded"]} numeric_drop={data_stats["numeric_dropped"]} ' + f'train={len(records)} eval={len(eval_records)} {args.dataset}; ' + f'train_gpus={TRAIN_GPUS} skill_dp={skill_dp} base_dp={base_dp}\n') + + hist: List[Dict[str, float]] = [] + rounds = 0 + pool = ProblemPool(records, args.seed) + with open(gen_path, 'w', encoding='utf-8') as gen_f, \ + open(eval_path, 'w', encoding='utf-8') as eval_f, \ + open(data_path, 'w', encoding='utf-8') as data_f, \ + open(train_log_path, 'w', encoding='utf-8') as tlog: + for f in (gen_f, eval_f, data_f, tlog): + _write(f, cfg) + gstep = 0 + # base baseline (frozen sampler, disk-cached, no weight-sync) is prefetched on a + # background thread while the current chunk generates: the skill-gen phase uses + # skill_sampler, so prefetching the NEXT chunk's baseline on base_sampler overlaps + # it and fills base_cache so the next draw_chunk hits it. base_sampler is wrapped in + # _LockedSampler, so the prefetch and this chunk's with-skill scoring never issue a + # base .sample() concurrently. It never touches the trainer or on-policy generation. + prefetch_pool = ThreadPoolExecutor(max_workers=1) if args.prefetch_baseline else None + pending: Optional[Any] = None + + def _prefetch(peeked: List[Dict[str, Any]]) -> None: + if peeked: + baseline_rollout(base_sampler, peeked, base_dp, args, base_cache) + + # Each chunk is drawn fresh + RE-GENERATED with the current policy (on-policy); + # with --balance, draw_chunk keeps drawing until the base fail:success mix hits target. + while rounds < args.max_train_rounds: + if pending is not None: + pending.result() # finish last round's prefetch before drawing (cache-warm) + pending = None + chunk, balance = draw_chunk(pool, base_sampler, base_dp, args, base_cache) + if prefetch_pool is not None: + peeked = pool.peek(int(args.chunk_size * (args.balance_max_draws_mult if args.balance else 1))) + pending = prefetch_pool.submit(_prefetch, peeked) + full, summary, groups = process_chunk( + base_sampler, skill_sampler, chunk, gstep, base_dp, skill_dp, + args, checker, rubric_cache) + summary['balance'] = balance + + log = None + if groups: + log = _train_chunk(skill_model, ckpt, groups, args) + rounds += 1 + log.update({'record_type': 'train_round', 'round': rounds, 'chunk': gstep, + 'epoch': pool.epoch, 'ts': int(time.time())}) + _write(tlog, log) + tlog.flush() + if rounds % args.save_rounds == 0: + skill_model.save(f'skill-rft-{rounds}', output_dir=args.output_dir) + + summary['rounds_done'], summary['epoch'] = rounds, pool.epoch + for rec in full: + _write(gen_f, rec) + _write(gen_f, summary) + gen_f.flush() + for v in groups: + _write(data_f, v) + data_f.flush() + + sa, sb = summary['view_A'], summary['view_B'] + hist.append({'aA': sa['adoption_rate'], 'aB': sb['adoption_rate'], + 'lift': summary['avg_lift'], 'pos': summary['n_reward_pos']}) + bal_str = (f'bal {balance["selected_fail"]}f/{balance["selected_success"]}s ' + f'(drew {balance["n_drawn"]}/fresh {balance["n_baseline_fresh"]}' + + ('!' if balance.get('budget_hit') else '') + ') ') if balance.get('enabled') else '' + sys.stderr.write( + f'[gen] e{pool.epoch} g{gstep}: {bal_str}hard={summary["n_hard"]} ' + f'clean={summary["n_clean"]} train={summary["n_train_samples"]} ' + f'acc={summary["avg_baseline_pass_on_hard"]:.2f}->{summary["avg_withskill_pass"]:.2f} ' + f'lift={summary["avg_lift"]:+.3f} ' + f'A[{sa["n_hard"]}h {sa["adoption_rate"]:.2f}] B[{sb["n_hard"]}h {sb["adoption_rate"]:.2f}] ' + f'rounds={rounds}\n') + if use_swan: + swanlab.log(_swan_metrics(summary, log), step=gstep) + + if eval_records and (gstep + 1) % args.eval_every == 0: + eval_recs, eval_summary, eval_metrics = run_greedy_eval( + base_sampler, skill_sampler, eval_records, gstep, rounds, base_dp, skill_dp, + args, checker, eval_base_cache, rubric_cache) + for rec in eval_recs: + _write(eval_f, rec) + _write(eval_f, eval_summary) + eval_f.flush() + if use_swan: + swanlab.log({f'eval/{k}': v for k, v in eval_metrics.items()}, step=gstep) + sys.stderr.write( + f'[eval] g{gstep}: n={eval_summary["n"]} mean@1 ' + f'acc={eval_summary["baseline_acc_mean1"]:.3f}->{eval_summary["acc_mean1"]:.3f} ' + f'lift={eval_summary["lift_mean1"]:+.3f} ' + f'A[{eval_summary["n_A"]} {eval_summary["acc_A_mean1"]:.3f}] ' + f'B[{eval_summary["n_B"]} {eval_summary["acc_B_mean1"]:.3f}] ' + f'fmt={eval_summary["format_mean1"]:.2f} rounds={rounds}\n') + + if (gstep + 1) % args.trend_every == 0: + tl = _trend_line(hist, args.trend_every, rounds) + if tl: + sys.stderr.write(tl + '\n') + gstep += 1 + + if prefetch_pool is not None: + if pending is not None: + pending.result() + prefetch_pool.shutdown(wait=True) + base_cache.close() + eval_base_cache.close() + rubric_cache.close() + skill_model.save('skill-rft-final', output_dir=args.output_dir) + sys.stderr.write(f'[rft] done: {rounds} rounds over {gstep} chunks / {pool.epoch} epochs -> {data_path}\n') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/embedding/train_reflexion_skill.sh b/cookbook/exp/embedding/train_reflexion_skill.sh new file mode 100755 index 000000000..21530a604 --- /dev/null +++ b/cookbook/exp/embedding/train_reflexion_skill.sh @@ -0,0 +1,54 @@ +#!/bin/bash +# Online GRPO RFT for the reflexion skill generator (unified, self-contained, cached). +# GPUs: 8 — ranks 0-3 train (skill model, FSDP2), 4-5 skill sampler (synced), 6-7 base +# sampler (frozen). Per chunk: base greedy solve -> rubric process-check (view A) -> +# skill-gen (thinking ON, N candidates) -> deterministic leak filter -> with-skill greedy +# pass -> group-relative advantage -> ONE on-policy GRPO step -> sync weights. +# +# Baseline rollouts + rubric diagnoses are disk-cached (output-dir/cache/*.jsonl), so a +# restart skips re-sampling them; skill-gen is on-policy and never cached. The next chunk's +# baseline is prefetched on a background thread (overlaps skill-gen; base sampler is frozen). +# +# The view-A rubric process-check uses the backup teacher API (set LLM_BACKUP_*). Without +# it the run still works: view A degrades to query-only and the leak filter stays +# deterministic (no teacher needed). + +set -euo pipefail + +export GEN_MODEL_ID=${GEN_MODEL_ID:-Qwen/Qwen3-4B} +export GEN_GPU_MEM=${GEN_GPU_MEM:-0.8} +# Datasets are pulled from ModelScope via twinkle.Dataset (ms://AI-MO/aops or +# ms://modelscope/competition_math); override AOPS_DATASET_ID / MATH_DATASET_ID to change. +# Teacher API for the view-A rubric process-check (optional; leak filter is deterministic). +export LLM_BACKUP_API_KEY=${LLM_BACKUP_API_KEY:-} +export LLM_BACKUP_BASE_URL=${LLM_BACKUP_BASE_URL:-https://dashscope.aliyuncs.com/compatible-mode/v1} +export LLM_BACKUP_MODEL=${LLM_BACKUP_MODEL:-qwen3.7-max} + +python cookbook/exp/embedding/train_reflexion_skill.py \ + --dataset aops \ + --n 5000 \ + --numeric-only \ + --chunk-size 16 \ + --n-skills 8 \ + --view-b-frac 0.5 \ + --skill-retries 2 \ + --balance \ + --balance-success-frac 0.4 \ + --balance-loop-frac 0.5 \ + --balance-max-draws-mult 8 \ + --max-tokens 25000 \ + --skill-max-tokens 8192 \ + --max-model-len 30000 \ + --eval-size 128 \ + --eval-every 10 \ + --sft-batch-size 8 \ + --grpo-epsilon 0.2 \ + --format-in-reward \ + --lr 6e-6 \ + --max-train-rounds 1500 \ + --save-rounds 25 \ + --trend-every 10 \ + --prefetch-baseline \ + --output-dir ./output/reflexion_skill \ + --swanlab-project twinkle \ + --swanlab-exp reflexion_skill_rft diff --git a/src/twinkle_agentic/verifier/leak_verifier.py b/src/twinkle_agentic/verifier/leak_verifier.py index f1cf730ac..b692d5f80 100644 --- a/src/twinkle_agentic/verifier/leak_verifier.py +++ b/src/twinkle_agentic/verifier/leak_verifier.py @@ -170,7 +170,7 @@ def __init__( max_content_chars: int = 4000, max_query_chars: int = 4000, uncertain_is_leak: bool = False, - answer_only: bool = False, + answer_only: bool = True, judge_system: Optional[str] = None, ): self.sampler = sampler From 23cf9f06e40fd71b18c40bad311fa894671cd739 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Wed, 15 Jul 2026 17:32:48 +0800 Subject: [PATCH 12/60] fix --- .../exp/embedding/train_reflexion_skill.py | 28 +++++++------------ 1 file changed, 10 insertions(+), 18 deletions(-) diff --git a/cookbook/exp/embedding/train_reflexion_skill.py b/cookbook/exp/embedding/train_reflexion_skill.py index 655233af5..7fb1bb94b 100644 --- a/cookbook/exp/embedding/train_reflexion_skill.py +++ b/cookbook/exp/embedding/train_reflexion_skill.py @@ -1032,15 +1032,15 @@ def _train_chunk(skill_model, ckpt: CheckpointEngineManager, samples: List[Dict[ # =========================================================================== def run_greedy_eval(base_sampler, skill_sampler, eval_records: List[Dict[str, Any]], ci: int, rounds: int, base_dp: int, skill_dp: int, args: argparse.Namespace, - checker, base_cache: DiskCache, rubric_cache: DiskCache + base_cache: DiskCache ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], Dict[str, float]]: """SEAM ``val-core/math/acc/mean@1`` on the fixed holdout: ONE greedy skill (T=0) per - problem into ONE greedy base solve (T=0). Each problem keeps its view; no leak filter - (acc scores correctness alone). Baseline + rubric reuse the disk caches.""" + problem into ONE greedy base solve (T=0). Eval ALWAYS uses view B (query-only, the + deployment form: no rubric, since rubric needs an online teacher unavailable at deploy); + no leak filter (acc scores correctness alone). Baseline reuses the disk cache.""" baseline_rollout(base_sampler, eval_records, base_dp, args, base_cache) for r in eval_records: - r['_view'], r['_rubric_diag'] = _assign_view(r['problem'], args), '' - diagnose_views(checker, eval_records, args, rubric_cache) + r['_view'], r['_rubric_diag'] = 'B', '' sg_out = _run_samples(skill_sampler, [_view_prompt(r) for r in eval_records], 1, args.skill_max_tokens, skill_dp, temperature=0.0) skills = [(_extract_skills_block(_clean_text(getattr(seqs[0], 'decoded', '') or '')) or '', @@ -1062,23 +1062,17 @@ def run_greedy_eval(base_sampler, skill_sampler, eval_records: List[Dict[str, An 'withskill_stop_reason': roll['stop_reason'], 'withskill_text': roll['text'], }) acc = lambda rs: sum(1 for x in rs if x['withskill_correct']) / len(rs) if rs else 0.0 - A = [x for x in recs if x['view'] == 'A'] - B = [x for x in recs if x['view'] == 'B'] - ws = acc(recs) + ws = acc(recs) # all view B (deployment form) base = sum(x['baseline_pass'] for x in recs) / len(recs) if recs else 0.0 fmt = (sum(1 for x in recs if x['skill_parseable']) / len(recs)) if recs else 0.0 term = (sum(1 for x in recs if x['withskill_terminated']) / len(recs)) if recs else 0.0 summary = {'record_type': 'eval_summary', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, - 'n': len(recs), 'n_A': len(A), 'n_B': len(B), 'acc_mean1': ws, + 'n': len(recs), 'view': 'B', 'acc_mean1': ws, 'baseline_acc_mean1': base, 'lift_mean1': ws - base, - 'acc_A_mean1': acc(A), 'acc_B_mean1': acc(B), 'format_mean1': fmt, 'term_mean1': term} + 'format_mean1': fmt, 'term_mean1': term} metrics = {'core/math/acc/mean@1': ws, 'core/math/baseline_acc/mean@1': base, 'core/math/lift/mean@1': ws - base, 'core/math/format/mean@1': fmt, 'core/math/term/mean@1': term} - if A: - metrics['core/math/acc_A/mean@1'] = acc(A) - if B: - metrics['core/math/acc_B/mean@1'] = acc(B) return recs, summary, metrics @@ -1354,7 +1348,7 @@ def _prefetch(peeked: List[Dict[str, Any]]) -> None: if eval_records and (gstep + 1) % args.eval_every == 0: eval_recs, eval_summary, eval_metrics = run_greedy_eval( base_sampler, skill_sampler, eval_records, gstep, rounds, base_dp, skill_dp, - args, checker, eval_base_cache, rubric_cache) + args, eval_base_cache) for rec in eval_recs: _write(eval_f, rec) _write(eval_f, eval_summary) @@ -1362,11 +1356,9 @@ def _prefetch(peeked: List[Dict[str, Any]]) -> None: if use_swan: swanlab.log({f'eval/{k}': v for k, v in eval_metrics.items()}, step=gstep) sys.stderr.write( - f'[eval] g{gstep}: n={eval_summary["n"]} mean@1 ' + f'[eval] g{gstep}: n={eval_summary["n"]} viewB mean@1 ' f'acc={eval_summary["baseline_acc_mean1"]:.3f}->{eval_summary["acc_mean1"]:.3f} ' f'lift={eval_summary["lift_mean1"]:+.3f} ' - f'A[{eval_summary["n_A"]} {eval_summary["acc_A_mean1"]:.3f}] ' - f'B[{eval_summary["n_B"]} {eval_summary["acc_B_mean1"]:.3f}] ' f'fmt={eval_summary["format_mean1"]:.2f} rounds={rounds}\n') if (gstep + 1) % args.trend_every == 0: From bff4fb48b175746616f1b38243cd35dbd3f1a76a Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Wed, 15 Jul 2026 17:54:51 +0800 Subject: [PATCH 13/60] fix --- .../exp/embedding/train_reflexion_skill.py | 338 +++++++++++++++--- .../exp/embedding/train_reflexion_skill.sh | 5 +- 2 files changed, 296 insertions(+), 47 deletions(-) diff --git a/cookbook/exp/embedding/train_reflexion_skill.py b/cookbook/exp/embedding/train_reflexion_skill.py index 7fb1bb94b..cfa3d125d 100644 --- a/cookbook/exp/embedding/train_reflexion_skill.py +++ b/cookbook/exp/embedding/train_reflexion_skill.py @@ -26,6 +26,7 @@ import copy import hashlib import json +import math import os import re import sys @@ -398,6 +399,125 @@ def load_problems(dataset: str, n: int, seed: int, num_proc: int = 0) -> List[Di return out[:n] if (n and n < len(out)) else out +# --------------------------------------------------------------------------- +# Bag-of-words neighbour pairing (cross-problem rubric transfer, --xproblem-rubric) +# --------------------------------------------------------------------------- +# Common English + math-scaffolding words that carry no problem-type signal. Kept small +# and deterministic on purpose (no external stopword list): what survives is the domain +# vocabulary ('triangle', 'prime', 'polynomial', ...) that actually defines the type. +_BOW_STOP = frozenset(""" +a an the of to in on at for and or but if is are be was were been being this that these those +with without into onto from by as it its their his her our your my we you they he she them +find compute determine calculate evaluate solve show prove given let suppose consider assume +what which when where how many much value values number numbers expression form terms term +such that then than so if only when each every all any some both one two three four five six +seven eight nine ten first second third last non over under about above below between +problem answer result equal equals sum difference product total following there here have has +had do does did can could will would should may might must not no yes if then else +""".split()) + +_WORD_RE = re.compile(r'[a-z]+') + + +def _stem(w: str) -> str: + """Crude suffix stripper so 'prime'/'primes', 'triangle'/'triangles' collapse to one + type token. Not linguistically correct -- just enough to merge the common plural/gerund + variants that otherwise split a type's vocabulary and starve the df filter.""" + if len(w) > 4 and w.endswith('ies'): + return w[:-3] + 'y' # properties -> property + if len(w) > 4 and w.endswith('es') and w[-3] in 'sxzh': + return w[:-2] # boxes -> box (keep primes -> prime below) + for suf in ('ing', 'ed', 's'): + if len(w) > len(suf) + 2 and w.endswith(suf): + return w[:-len(suf)] + return w + + +def _tokenize(problem: str) -> List[str]: + """Deterministic bag-of-words tokens for type matching: lowercase alphabetic words + (numbers dropped -- they are instance detail, not type), minus generic stopwords, then + stemmed, so only domain terms remain. Latex control words (frac, sqrt, ...) survive.""" + return [_stem(w) for w in _WORD_RE.findall((problem or '').lower()) + if len(w) > 2 and w not in _BOW_STOP] + + +class BagOfWordsIndex: + """TF-IDF cosine nearest-neighbour over problem statements. Sparse dict vectors + + an inverted index, so ``nearest`` scores only problems sharing a term (near-linear in + practice, no dense NxN). Deterministic; ties break on lower index for reproducibility.""" + + def __init__(self, problems: List[str], min_df: int = 2, max_df_frac: float = 0.5): + self._toks = [_tokenize(p) for p in problems] + n = len(self._toks) + df: Dict[str, int] = {} + for toks in self._toks: + for w in set(toks): + df[w] = df.get(w, 0) + 1 + max_df = max(min_df, int(max_df_frac * n)) + self._idf = {w: math.log((n + 1) / (c + 1)) + 1.0 + for w, c in df.items() if min_df <= c <= max_df} + self._vecs: List[Dict[str, float]] = [self._vectorize(t) for t in self._toks] + self._inverted: Dict[str, List[int]] = {} + for i, v in enumerate(self._vecs): + for w in v: + self._inverted.setdefault(w, []).append(i) + + def _vectorize(self, toks: List[str]) -> Dict[str, float]: + tf: Dict[str, float] = {} + for w in toks: + if w in self._idf: + tf[w] = tf.get(w, 0.0) + 1.0 + vec = {w: c * self._idf[w] for w, c in tf.items()} + norm = math.sqrt(sum(x * x for x in vec.values())) + return {w: x / norm for w, x in vec.items()} if norm > 0 else {} + + def nearest(self, i: int) -> Tuple[int, float]: + """(index, cosine) of the most similar OTHER problem, or (-1, 0.0) if none.""" + vi = self._vecs[i] + if not vi: + return -1, 0.0 + scores: Dict[int, float] = {} + for w, xi in vi.items(): + for j in self._inverted.get(w, ()): + if j != i: + scores[j] = scores.get(j, 0.0) + xi * self._vecs[j].get(w, 0.0) + if not scores: + return -1, 0.0 + best = max(scores.items(), key=lambda kv: (kv[1], -kv[0])) + return best[0], best[1] + + +def build_neighbor_map(problems: List[str], cache: 'DiskCache') -> Dict[str, Tuple[str, float]]: + """{problem -> (nearest-neighbour problem, cosine similarity)}. Cached on disk keyed by + the whole problem set (order-independent), so a restart reloads instantly.""" + key = DiskCache.key_for('bow_neighbors', *sorted(problems)) + hit = cache.get(key) + if hit is not None: + return {q: (p, s) for q, (p, s) in hit.items()} # json roundtrips tuples as lists + index = BagOfWordsIndex(problems) + out: Dict[str, Tuple[str, float]] = {} + for i, p in enumerate(problems): + j, sim = index.nearest(i) + if j >= 0: + out[p] = (problems[j], sim) + cache.put(key, {q: [p, s] for q, (p, s) in out.items()}) + return out + + +def select_paired_subset(records: List[Dict[str, Any]], n: int, seed: int + ) -> List[Dict[str, Any]]: + """Keep the ``n`` problems with the strongest bag-of-words neighbour, so the training + pool is dense in same-type pairs (cross-problem rubric transfer needs a real analogue, + not a random other problem). n<=0 or n>=len keeps all (still shuffled).""" + index = BagOfWordsIndex([r['problem'] for r in records]) + scored = [(index.nearest(i)[1], i) for i in range(len(records))] + scored.sort(key=lambda si: (-si[0], si[1])) + keep = [i for _, i in scored[:n]] if (0 < n < len(records)) else list(range(len(records))) + rng = np.random.RandomState(seed) + rng.shuffle(keep) + return [records[i] for i in keep] + + _NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') @@ -455,7 +575,11 @@ def _load_records(args: argparse.Namespace) -> Tuple[List[Dict[str, Any]], List[ eval_records = [dict(r) for r in records[:eval_n]] pool = records[eval_n:] train_n = args.n if args.n > 0 else len(pool) - train_records = [dict(r) for r in pool[:train_n]] + # Cross-problem rubric: keep the train_n problems with the strongest bag-of-words + # neighbour (dense same-type pairs), so a problem's rubric can transfer to its analogue. + # Otherwise keep the first train_n (already shuffled). + train_records = ([dict(r) for r in select_paired_subset(pool, train_n, args.seed)] + if args.xproblem_rubric else [dict(r) for r in pool[:train_n]]) if {r['problem'] for r in train_records} & {r['problem'] for r in eval_records}: raise ValueError('eval/train overlap detected') stats = {'raw_loaded': raw_n, 'numeric_dropped': dropped, @@ -818,16 +942,50 @@ def _assign_advantages(hard: List[Dict[str, Any]], args: argparse.Namespace) -> c['advantage'], c['grpo_adv'], c['kept'] = adv, adv, c['reward'] > mean_r +def apply_neighbor_rubric(base_sampler, chunk: List[Dict[str, Any]], + neighbor_map: Dict[str, Tuple[str, float]], base_dp: int, + args: argparse.Namespace, checker, + base_cache: DiskCache, rubric_cache: DiskCache) -> None: + """Cross-problem rubric (--xproblem-rubric): for every view-A problem Q, replace its own + rubric with the rubric of its bag-of-words neighbour P (Q keeps being the solved/scored + problem). P is baselined + diagnosed here (both disk-cached) as a throwaway stub, then + P's findings are copied onto Q, along with the neighbour's text + similarity for audit. + A P that can't be diagnosed -> Q degrades to view B. This makes any answer leaked from + P's rubric useless for solving Q.""" + targets = [r for r in chunk if r.get('_view') == 'A' and neighbor_map.get(r['problem'])] + if not targets: + return + stubs, by_problem = [], {} + for r in targets: + p, _ = neighbor_map[r['problem']] + if p not in by_problem: + stub = {'problem': p, 'reference_answer': '', '_view': 'A'} + by_problem[p] = stub + stubs.append(stub) + baseline_rollout(base_sampler, stubs, base_dp, args, base_cache) + diagnose_views(checker, stubs, args, rubric_cache) + for r in targets: + p, sim = neighbor_map[r['problem']] + r['_rubric_diag'] = by_problem[p].get('_rubric_diag', '') + r['_rubric_src'], r['_neighbor_sim'] = p, sim + + def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], ci: int, base_dp: int, skill_dp: int, args: argparse.Namespace, - checker, rubric_cache: DiskCache + checker, rubric_cache: DiskCache, base_cache: DiskCache = None, + neighbor_map: Optional[Dict[str, str]] = None ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], List[Dict[str, Any]]]: """view assign -> rubric-check (view A) -> skill-gen -> leak-filter -> with-skill - greedy pass -> GRPO advantages. ``chunk`` arrives already baselined by draw_chunk.""" + greedy pass -> GRPO advantages. ``chunk`` arrives already baselined by draw_chunk. + With --xproblem-rubric, view A gets its NEIGHBOUR's rubric (see apply_neighbor_rubric).""" hard = chunk for r in hard: r['_view'], r['_rubric_diag'] = _assign_view(r['problem'], args), '' - diagnose_views(checker, hard, args, rubric_cache) + if args.xproblem_rubric and neighbor_map: + apply_neighbor_rubric(base_sampler, hard, neighbor_map, base_dp, args, checker, + base_cache, rubric_cache) + else: + diagnose_views(checker, hard, args, rubric_cache) # skill-gen (thinking ON), per-view prompt; re-sample problems with no clean candidate. flat: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] @@ -847,7 +1005,9 @@ def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], block = _extract_skills_block(resp) cand = {'skills': block or '', 'response': resp, 'parseable': bool(block), 'view': r['_view'], 'leaked': None, 'leak_reason': '', - 'leak_source': '', 'with_pass': None, 'reward': None, 'rolls': []} + 'leak_source': '', 'with_pass': None, 'reward': None, 'rolls': [], + 'skillgen_stop': getattr(s, 'stop_reason', None), + 'skillgen_tokens': len(getattr(s, 'tokens', None) or [])} r['_cands'].append(cand) if block: flat.append((r, cand)) @@ -913,9 +1073,12 @@ def _full_record(r: Dict[str, Any], ci: int) -> Dict[str, Any]: 'gen_tokens': init['gen_tokens']}, 'baseline_pass': r['_baseline_pass'], 'is_hard': r['_hard'], 'view': r.get('_view', ''), 'rubric_diag': r.get('_rubric_diag', ''), + # xproblem audit: which neighbour's rubric this problem borrowed, and how similar. + 'rubric_src': r.get('_rubric_src', ''), 'neighbor_sim': r.get('_neighbor_sim'), 'baseline_rolls': [_roll(x) for x in r['_baseline_rolls']], 'candidates': [{ 'skills': c['skills'], 'response': c['response'], 'parseable': c['parseable'], + 'skillgen_stop': c.get('skillgen_stop'), 'skillgen_tokens': c.get('skillgen_tokens'), 'leaked': c['leaked'], 'leak_reason': c['leak_reason'], 'leak_source': c['leak_source'], 'with_pass': c['with_pass'], 'reward': c.get('reward'), 'advantage': c.get('advantage'), 'grpo_adv': c.get('grpo_adv'), 'kept': c.get('kept'), @@ -924,46 +1087,99 @@ def _full_record(r: Dict[str, Any], ci: int) -> Dict[str, Any]: } -def _view_stats(hard: List[Dict[str, Any]], view: str) -> Dict[str, Any]: - hv = [r for r in hard if r.get('_view') == view] - cands = [c for r in hv for c in r['_cands'] if c['parseable']] +def _view_stats(problems: List[Dict[str, Any]], view: str) -> Dict[str, Any]: + pv = [r for r in problems if r.get('_view') == view] + cands = [c for r in pv for c in r['_cands'] if c['parseable']] clean = [c for c in cands if c['leaked'] is False] - adopted = sum(1 for r in hv + adopted = sum(1 for r in pv if any(c['leaked'] is False and abs(c.get('advantage') or 0.0) > 1e-9 for c in r['_cands'])) - return {'n_hard': len(hv), 'n_candidates_parseable': len(cands), 'n_clean': len(clean), - 'n_adopted_problems': adopted, 'adoption_rate': (adopted / len(hv)) if hv else 0.0} + return {'n': len(pv), 'n_candidates_parseable': len(cands), 'n_clean': len(clean), + 'n_adopted_problems': adopted, 'adoption_rate': (adopted / len(pv)) if pv else 0.0} + + +def _mean(xs: List[float]) -> float: + return sum(xs) / len(xs) if xs else 0.0 + + +def _std(xs: List[float]) -> float: + if len(xs) < 2: + return 0.0 + m = _mean(xs) + return (sum((x - m) ** 2 for x in xs) / len(xs)) ** 0.5 + + +def _signal_stats(problems: List[Dict[str, Any]]) -> Dict[str, float]: + """The heart of 'is there a learning signal': per problem, the scored candidates form a + GRPO group. A group with zero reward variance (all skills solve, or none do -- the + hard-problem dead zone) gives NO gradient. Tracks that fraction plus reward level and + within-group variance so a collapse (all-0 or all-1) is visible immediately.""" + group_vars, all_rewards, zero_grad, groups = [], [], 0, 0 + for r in problems: + rewards = [c['reward'] for c in r['_cands'] if c.get('reward') is not None] + if len(rewards) < 2: + continue + groups += 1 + all_rewards.extend(rewards) + v = _std(rewards) + group_vars.append(v) + if v < 1e-9: # every skill got the same reward -> GRPO skips this problem + zero_grad += 1 + return {'n_groups': groups, 'zero_grad_frac': (zero_grad / groups) if groups else 0.0, + 'reward_mean': _mean(all_rewards), 'reward_std': _std(all_rewards), + 'group_reward_std_mean': _mean(group_vars)} def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespace) -> Dict[str, Any]: - hard = [r for r in chunk if r['_hard']] all_cands = [c for r in chunk for c in r['_cands']] cands = [c for c in all_cands if c['parseable']] scored = [c for c in cands if c['with_pass'] is not None] + clean = [c for c in cands if c['leaked'] is False] ws_rolls = [x for c in scored for x in c['rolls']] - fail_cands = [c for r in chunk if r['_failed'] for c in r['_cands']] - base_acc = (sum(r['_baseline_pass'] for r in hard) / len(hard)) if hard else 0.0 - ws_acc = (sum(c['with_pass'] for c in scored) / len(scored)) if scored else 0.0 - abs_adv = lambda cs: sum(abs(c.get('advantage') or 0.0) for c in cs) - total_abs = abs_adv(all_cands) + base_acc = (sum(r['_baseline_pass'] for r in chunk) / len(chunk)) if chunk else 0.0 + ws_acc = _mean([c['with_pass'] for c in scored]) + # base failure taxonomy (you asked whether skills fail because the base loops out of length) + classes = [_baseline_class(r) for r in chunk] + n_fail = sum(1 for c in classes if c != 'success') + skill_tokens = [c.get('skillgen_tokens') or 0 for c in cands] # skill-gen response length + trunc = sum(1 for r in chunk for c in r['_cands'] + for x in c['rolls'] if x['stop_reason'] == 'length') return { 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), - 'n_failed_first_try': sum(1 for r in chunk if r['_failed']), 'n_hard': len(hard), + 'n_failed_first_try': sum(1 for r in chunk if r['_failed']), 'n_generated': len(all_cands), 'n_candidates_parseable': len(cands), 'n_unparseable': len(all_cands) - len(cands), - 'n_leaked': sum(1 for c in cands if c['leaked']), - 'n_clean': sum(1 for c in cands if c['leaked'] is False), + 'parse_rate': (len(cands) / len(all_cands)) if all_cands else 0.0, + 'n_leaked': sum(1 for c in cands if c['leaked']), 'n_clean': len(clean), + 'leak_rate': (sum(1 for c in cands if c['leaked']) / len(cands)) if cands else 0.0, 'n_reward_pos': sum(1 for c in scored if c['reward']), 'n_train_samples': sum(1 for c in all_cands if _is_trainable(c, args)), - 'n_train_from_fail': sum(1 for c in fail_cands if _is_trainable(c, args)), - 'abs_adv_from_fail_frac': (abs_adv(fail_cands) / total_abs) if total_abs > 0 else 0.0, - 'avg_baseline_pass_on_hard': base_acc, 'avg_withskill_pass': ws_acc, + 'signal': _signal_stats(chunk), + 'fail_loop_frac': (sum(1 for c in classes if c == 'fail_loop') / n_fail) if n_fail else 0.0, + 'fail_wrong_frac': (sum(1 for c in classes if c == 'fail_wrong') / n_fail) if n_fail else 0.0, + 'skill_tokens_mean': _mean(skill_tokens), + 'withskill_trunc_frac': (trunc / len(ws_rolls)) if ws_rolls else 0.0, + 'avg_baseline_pass': base_acc, 'avg_withskill_pass': ws_acc, 'avg_lift': ws_acc - base_acc, - 'termination_rate_withskill': (sum(1 for x in ws_rolls if x['terminated']) / len(ws_rolls)) if ws_rolls else 0.0, - 'view_A': _view_stats(hard, 'A'), 'view_B': _view_stats(hard, 'B'), + 'termination_rate_withskill': _mean([1.0 if x['terminated'] else 0.0 for x in ws_rolls]), + 'view_A': _view_stats(chunk, 'A'), 'view_B': _view_stats(chunk, 'B'), + **_xproblem_stats(chunk, args), } +def _xproblem_stats(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> Dict[str, Any]: + """Cross-problem pairing health: of the view-A problems, how many actually got a + neighbour's rubric (pair_rate) and how similar those neighbours were. Empty when off.""" + if not args.xproblem_rubric: + return {} + view_a = [r for r in chunk if r.get('_view') == 'A'] + paired = [r for r in view_a if r.get('_rubric_src')] + return {'xproblem': { + 'n_view_a': len(view_a), 'n_paired': len(paired), + 'pair_rate': (len(paired) / len(view_a)) if view_a else 0.0, + 'neighbor_sim_mean': _mean([r.get('_neighbor_sim', 0.0) for r in paired])}} + + def _group_records(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> List[Dict[str, Any]]: """GRPO training records: every trainable candidate with its view + rubric diagnosis (the prompt is rebuilt from those by ``_skillgen_messages``, no trajectory stored).""" @@ -976,7 +1192,9 @@ def _group_records(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> Lis out.append({ 'problem': r['problem'], 'reference_answer': r['reference_answer'], 'view': r.get('_view', 'A'), 'diagnosis': r.get('_rubric_diag', ''), + 'rubric_src': r.get('_rubric_src', ''), 'response': c['response'], 'skills': c['skills'], + 'skillgen_stop': c.get('skillgen_stop'), 'advantage': c['advantage'], 'grpo_adv': c['grpo_adv'], 'kept': c['kept'], 'reward': c['reward'], 'with_pass': c['with_pass']}) return out @@ -1086,29 +1304,43 @@ def _trend_line(hist: List[Dict[str, float]], window: int, rounds_done: int) -> return (f'[trend] first {window} vs last {window} | ' f'adopt A {m(base,"aA"):.2f}->{m(rec,"aA"):.2f} B {m(base,"aB"):.2f}->{m(rec,"aB"):.2f} | ' f'lift {m(base,"lift"):+.3f}->{m(rec,"lift"):+.3f} | ' + f'0grad {m(base,"zero_grad"):.2f}->{m(rec,"zero_grad"):.2f} | ' f'pos/chunk {m(base,"pos"):.1f}->{m(rec,"pos"):.1f} | rounds={rounds_done}') def _swan_metrics(summary: Dict[str, Any], log: Optional[Dict[str, Any]]) -> Dict[str, float]: - """Flat swanlab dict = generation metrics + (when trained) the GRPO built-in metric. - acc/adopt/term only on chunks with hard problems, so idle chunks don't dip charts.""" + """Flat swanlab dict. ``signal/*`` is the primary health group (is GRPO getting a + gradient at all); ``acc/*`` is the effect; the rest diagnose why. acc/adopt/skill are + only logged when the chunk produced a scored group, so idle chunks don't dip charts.""" + sig = summary['signal'] d: Dict[str, float] = { - 'gen/n_hard': summary['n_hard'], 'gen/n_clean': summary['n_clean'], - 'gen/n_leaked': summary['n_leaked'], 'gen/n_train_samples': summary['n_train_samples'], - 'gen/n_reward_pos': summary['n_reward_pos'], - 'gen/n_train_from_fail': summary['n_train_from_fail'], - 'gen/abs_adv_from_fail_frac': summary['abs_adv_from_fail_frac'], + # --- signal: the FIRST thing to watch (no variance -> no learning) --- + 'signal/zero_grad_frac': sig['zero_grad_frac'], 'signal/n_groups': sig['n_groups'], + 'signal/reward_mean': sig['reward_mean'], 'signal/reward_std': sig['reward_std'], + 'signal/group_reward_std_mean': sig['group_reward_std_mean'], + 'signal/n_train_samples': summary['n_train_samples'], + 'signal/n_reward_pos': summary['n_reward_pos'], + # --- skill format / leak health --- + 'skill/parse_rate': summary['parse_rate'], 'skill/tokens_mean': summary['skill_tokens_mean'], + 'leak/rate': summary['leak_rate'], 'leak/n': summary['n_leaked'], + # --- base failure taxonomy (loop-out-of-length vs plain wrong) --- + 'fail/loop_frac': summary['fail_loop_frac'], 'fail/wrong_frac': summary['fail_wrong_frac'], + 'fail/frac_first_try': (summary['n_failed_first_try'] / summary['n']) if summary['n'] else 0.0, } bal = summary.get('balance') or {} if bal.get('enabled'): d.update({'balance/n_drawn': bal['n_drawn'], 'balance/n_baseline_fresh': bal['n_baseline_fresh'], 'balance/selected_success_frac': bal['selected_success_frac']}) - if summary['n_hard'] > 0: - d.update({'acc/baseline_pass': summary['avg_baseline_pass_on_hard'], + xp = summary.get('xproblem') or {} + if xp: + d.update({'xproblem/pair_rate': xp['pair_rate'], 'xproblem/neighbor_sim_mean': xp['neighbor_sim_mean']}) + if sig['n_groups'] > 0: + d.update({'acc/baseline_pass': summary['avg_baseline_pass'], 'acc/withskill_pass': summary['avg_withskill_pass'], 'acc/lift': summary['avg_lift'], 'adopt/A': summary['view_A']['adoption_rate'], 'adopt/B': summary['view_B']['adoption_rate'], - 'term/withskill': summary['termination_rate_withskill']}) + 'term/withskill': summary['termination_rate_withskill'], + 'term/withskill_trunc_frac': summary['withskill_trunc_frac']}) if log: d['train/n_steps'] = log['n_steps'] d['train/n_micro_batches'] = log['n_micro_batches'] @@ -1171,7 +1403,7 @@ def _build_args() -> argparse.Namespace: p.add_argument('--seed', type=int, default=42) p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) p.add_argument('--eval-size', type=int, default=128, help='Fixed holdout size (0 disables).') - p.add_argument('--eval-every', type=int, default=10, help='Run holdout eval every N chunks.') + p.add_argument('--eval-every', type=int, default=5, help='Run holdout eval every N chunks.') p.add_argument('--chunk-size', type=int, default=16) p.add_argument('--balance', action=argparse.BooleanOptionalAction, default=True) p.add_argument('--balance-success-frac', type=float, default=0.4, @@ -1180,6 +1412,11 @@ def _build_args() -> argparse.Namespace: p.add_argument('--balance-max-draws-mult', type=int, default=8) p.add_argument('--n-skills', type=int, default=8) p.add_argument('--view-b-frac', type=float, default=0.5) + p.add_argument('--xproblem-rubric', action=argparse.BooleanOptionalAction, default=True, + help='View A uses its bag-of-words NEIGHBOUR problem\'s rubric (transfer ' + 'test; kills answer leakage since the neighbour\'s answer differs). ' + 'On (default); --no-xproblem-rubric = each problem uses its own rubric ' + '(may leak).') p.add_argument('--skill-retries', type=int, default=2) p.add_argument('--skill-gen-temperature', type=float, default=1.0) p.add_argument('--skill-gen-top-p', type=float, default=1.0) @@ -1194,7 +1431,7 @@ def _build_args() -> argparse.Namespace: p.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True) p.add_argument('--lr', type=float, default=6e-6) p.add_argument('--max-train-rounds', type=int, default=1500) - p.add_argument('--save-rounds', type=int, default=50) + p.add_argument('--save-rounds', type=int, default=200) p.add_argument('--trend-every', type=int, default=10) p.add_argument('--output-dir', default='./output/reflexion_skill_rft') p.add_argument('--cache-dir', default='', help='Baseline/rubric cache dir (default /cache).') @@ -1254,6 +1491,12 @@ def main() -> None: base_cache = DiskCache(os.path.join(cache_dir, 'baseline.jsonl'), use_cache) eval_base_cache = DiskCache(os.path.join(cache_dir, 'eval_baseline.jsonl'), use_cache) rubric_cache = DiskCache(os.path.join(cache_dir, 'rubric.jsonl'), use_cache) + neighbor_map: Dict[str, str] = {} + if args.xproblem_rubric: + neighbor_cache = DiskCache(os.path.join(cache_dir, 'bow_neighbors.jsonl'), use_cache) + neighbor_map = build_neighbor_map([r['problem'] for r in records], neighbor_cache) + neighbor_cache.close() + sys.stderr.write(f'[rft] xproblem-rubric ON: {len(neighbor_map)} bag-of-words pairs\n') cfg = {'record_type': 'config', 'model': MODEL_ID, 'dataset': args.dataset, 'n': len(records), 'eval_n': len(eval_records), 'seed': args.seed, @@ -1265,6 +1508,7 @@ def main() -> None: 'skill_gen_temp': args.skill_gen_temperature, 'reward': 'greedy_binary(correct)', 'advantage': 'group_relative', 'format_in_reward': args.format_in_reward, 'cache': use_cache, 'rubric_check': 'fixed_math_5crit(viewA)' if checker else 'disabled', + 'xproblem_rubric': args.xproblem_rubric, 'grpo_epsilon': args.grpo_epsilon, 'lr': args.lr, 'max_train_rounds': args.max_train_rounds, 'started': int(time.time())} sys.stderr.write(f'[rft] raw={data_stats["raw_loaded"]} numeric_drop={data_stats["numeric_dropped"]} ' @@ -1306,7 +1550,7 @@ def _prefetch(peeked: List[Dict[str, Any]]) -> None: pending = prefetch_pool.submit(_prefetch, peeked) full, summary, groups = process_chunk( base_sampler, skill_sampler, chunk, gstep, base_dp, skill_dp, - args, checker, rubric_cache) + args, checker, rubric_cache, base_cache, neighbor_map) summary['balance'] = balance log = None @@ -1329,18 +1573,22 @@ def _prefetch(peeked: List[Dict[str, Any]]) -> None: _write(data_f, v) data_f.flush() - sa, sb = summary['view_A'], summary['view_B'] + sa, sb, sig = summary['view_A'], summary['view_B'], summary['signal'] hist.append({'aA': sa['adoption_rate'], 'aB': sb['adoption_rate'], - 'lift': summary['avg_lift'], 'pos': summary['n_reward_pos']}) + 'lift': summary['avg_lift'], 'pos': summary['n_reward_pos'], + 'zero_grad': sig['zero_grad_frac']}) bal_str = (f'bal {balance["selected_fail"]}f/{balance["selected_success"]}s ' f'(drew {balance["n_drawn"]}/fresh {balance["n_baseline_fresh"]}' + ('!' if balance.get('budget_hit') else '') + ') ') if balance.get('enabled') else '' + xp = summary.get('xproblem') + xp_str = f'pair={xp["pair_rate"]:.2f}@{xp["neighbor_sim_mean"]:.2f} ' if xp else '' sys.stderr.write( - f'[gen] e{pool.epoch} g{gstep}: {bal_str}hard={summary["n_hard"]} ' - f'clean={summary["n_clean"]} train={summary["n_train_samples"]} ' - f'acc={summary["avg_baseline_pass_on_hard"]:.2f}->{summary["avg_withskill_pass"]:.2f} ' - f'lift={summary["avg_lift"]:+.3f} ' - f'A[{sa["n_hard"]}h {sa["adoption_rate"]:.2f}] B[{sb["n_hard"]}h {sb["adoption_rate"]:.2f}] ' + f'[gen] e{pool.epoch} g{gstep}: {bal_str}n={summary["n"]} ' + f'clean={summary["n_clean"]} leak={summary["leak_rate"]:.2f} train={summary["n_train_samples"]} ' + f'0grad={sig["zero_grad_frac"]:.2f} R={sig["reward_mean"]:.2f}±{sig["reward_std"]:.2f} ' + f'acc={summary["avg_baseline_pass"]:.2f}->{summary["avg_withskill_pass"]:.2f} ' + f'lift={summary["avg_lift"]:+.3f} {xp_str}' + f'A[{sa["n"]} {sa["adoption_rate"]:.2f}] B[{sb["n"]} {sb["adoption_rate"]:.2f}] ' f'rounds={rounds}\n') if use_swan: swanlab.log(_swan_metrics(summary, log), step=gstep) diff --git a/cookbook/exp/embedding/train_reflexion_skill.sh b/cookbook/exp/embedding/train_reflexion_skill.sh index 21530a604..f5fde814d 100755 --- a/cookbook/exp/embedding/train_reflexion_skill.sh +++ b/cookbook/exp/embedding/train_reflexion_skill.sh @@ -31,6 +31,7 @@ python cookbook/exp/embedding/train_reflexion_skill.py \ --chunk-size 16 \ --n-skills 8 \ --view-b-frac 0.5 \ + --xproblem-rubric \ --skill-retries 2 \ --balance \ --balance-success-frac 0.4 \ @@ -40,13 +41,13 @@ python cookbook/exp/embedding/train_reflexion_skill.py \ --skill-max-tokens 8192 \ --max-model-len 30000 \ --eval-size 128 \ - --eval-every 10 \ + --eval-every 5 \ --sft-batch-size 8 \ --grpo-epsilon 0.2 \ --format-in-reward \ --lr 6e-6 \ --max-train-rounds 1500 \ - --save-rounds 25 \ + --save-rounds 200 \ --trend-every 10 \ --prefetch-baseline \ --output-dir ./output/reflexion_skill \ From fb06fe56ffa33d367e1c04aefc3ee6f8a4cbf9c2 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Wed, 15 Jul 2026 18:05:12 +0800 Subject: [PATCH 14/60] fix --- .../exp/embedding/train_reflexion_skill.py | 131 ++++++++++-------- 1 file changed, 71 insertions(+), 60 deletions(-) diff --git a/cookbook/exp/embedding/train_reflexion_skill.py b/cookbook/exp/embedding/train_reflexion_skill.py index cfa3d125d..f76108ce7 100644 --- a/cookbook/exp/embedding/train_reflexion_skill.py +++ b/cookbook/exp/embedding/train_reflexion_skill.py @@ -444,10 +444,17 @@ def _tokenize(problem: str) -> List[str]: class BagOfWordsIndex: """TF-IDF cosine nearest-neighbour over problem statements. Sparse dict vectors + an inverted index, so ``nearest`` scores only problems sharing a term (near-linear in - practice, no dense NxN). Deterministic; ties break on lower index for reproducibility.""" + practice, no dense NxN). Deterministic; ties break on lower index for reproducibility. - def __init__(self, problems: List[str], min_df: int = 2, max_df_frac: float = 0.5): + Answer-aware: ``nearest`` skips the same index, near-verbatim duplicates (cosine + >= ``sim_max``) and -- crucially for anti-leak -- any candidate whose answer equals the + query's, so a neighbour rubric can never hand over the query's own answer.""" + + def __init__(self, problems: List[str], answers: Optional[List[str]] = None, + min_df: int = 2, max_df_frac: float = 0.5): self._toks = [_tokenize(p) for p in problems] + self._ans = [(_numeric_value(a) or (str(a).strip() if a else '')) for a in answers] \ + if answers is not None else [''] * len(problems) n = len(self._toks) df: Dict[str, int] = {} for toks in self._toks: @@ -471,51 +478,47 @@ def _vectorize(self, toks: List[str]) -> Dict[str, float]: norm = math.sqrt(sum(x * x for x in vec.values())) return {w: x / norm for w, x in vec.items()} if norm > 0 else {} - def nearest(self, i: int) -> Tuple[int, float]: - """(index, cosine) of the most similar OTHER problem, or (-1, 0.0) if none.""" + def nearest(self, i: int, sim_max: float = 0.98) -> Tuple[int, float]: + """(index, cosine) of the most similar *distinct* problem: not i, not a duplicate + (cosine < ``sim_max``) and not the same answer. (-1, 0.0) if none qualifies.""" vi = self._vecs[i] if not vi: return -1, 0.0 + ai = self._ans[i] scores: Dict[int, float] = {} for w, xi in vi.items(): for j in self._inverted.get(w, ()): if j != i: scores[j] = scores.get(j, 0.0) + xi * self._vecs[j].get(w, 0.0) - if not scores: - return -1, 0.0 - best = max(scores.items(), key=lambda kv: (kv[1], -kv[0])) - return best[0], best[1] - - -def build_neighbor_map(problems: List[str], cache: 'DiskCache') -> Dict[str, Tuple[str, float]]: - """{problem -> (nearest-neighbour problem, cosine similarity)}. Cached on disk keyed by - the whole problem set (order-independent), so a restart reloads instantly.""" - key = DiskCache.key_for('bow_neighbors', *sorted(problems)) - hit = cache.get(key) - if hit is not None: - return {q: (p, s) for q, (p, s) in hit.items()} # json roundtrips tuples as lists - index = BagOfWordsIndex(problems) - out: Dict[str, Tuple[str, float]] = {} - for i, p in enumerate(problems): - j, sim = index.nearest(i) - if j >= 0: - out[p] = (problems[j], sim) - cache.put(key, {q: [p, s] for q, (p, s) in out.items()}) - return out - - -def select_paired_subset(records: List[Dict[str, Any]], n: int, seed: int - ) -> List[Dict[str, Any]]: - """Keep the ``n`` problems with the strongest bag-of-words neighbour, so the training - pool is dense in same-type pairs (cross-problem rubric transfer needs a real analogue, - not a random other problem). n<=0 or n>=len keeps all (still shuffled).""" - index = BagOfWordsIndex([r['problem'] for r in records]) - scored = [(index.nearest(i)[1], i) for i in range(len(records))] - scored.sort(key=lambda si: (-si[0], si[1])) - keep = [i for _, i in scored[:n]] if (0 < n < len(records)) else list(range(len(records))) + best_j, best_s = -1, 0.0 + for j, s in scores.items(): + if s >= sim_max or (ai and self._ans[j] == ai): + continue + if s > best_s or (s == best_s and (best_j < 0 or j < best_j)): + best_j, best_s = j, s + return best_j, best_s + + +def build_pairs(records: List[Dict[str, Any]], n: int, seed: int, sim_max: float = 0.98 + ) -> Tuple[List[Dict[str, Any]], Dict[str, Tuple[str, float]]]: + """Single-pass cross-problem pairing over the whole pool (one index build). + + Returns ``(subset, neighbour_map)`` where ``subset`` is the ``n`` problems with the + strongest qualifying neighbour (dense same-type pairs; n<=0/n>=len keeps all, shuffled) + and ``neighbour_map[q] = (p, cosine)`` gives each kept problem its analogue P. P is drawn + from the *full* pool (richer analogues) but never a duplicate or same-answer problem, so + P's rubric can transfer method without ever leaking Q's answer.""" + index = BagOfWordsIndex([r['problem'] for r in records], + [str(r.get('reference_answer', '')) for r in records]) + nbr = [index.nearest(i, sim_max) for i in range(len(records))] + order = sorted(range(len(records)), key=lambda i: (-nbr[i][1], i)) + keep = order[:n] if (0 < n < len(records)) else list(range(len(records))) rng = np.random.RandomState(seed) rng.shuffle(keep) - return [records[i] for i in keep] + subset = [records[i] for i in keep] + neighbour_map = {records[i]['problem']: (records[nbr[i][0]]['problem'], nbr[i][1]) + for i in keep if nbr[i][0] >= 0} + return subset, neighbour_map _NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') @@ -559,8 +562,13 @@ def _answer_leaked(skill: str, reference: str) -> bool: return False -def _load_records(args: argparse.Namespace) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], Dict[str, int]]: - """Load, numeric-filter, shuffle, then split a fixed eval holdout off the front.""" +def _load_records(args: argparse.Namespace + ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], + Dict[str, Tuple[str, float]], Dict[str, str], Dict[str, int]]: + """Load, numeric-filter, shuffle, split a fixed eval holdout, and (when --xproblem-rubric) + select a same-type-dense train subset with its neighbour map -- all in one pass. + Also returns ``pool_answers`` (every candidate neighbour P's true answer) so P's baseline + can be graded/cached correctly even when P is not itself a training problem.""" # Load all when filtering or splitting (else the eval holdout could starve train). load_n = 0 if (args.numeric_only or args.eval_size > 0) else args.n records = load_problems(args.dataset, load_n, args.seed) @@ -575,16 +583,20 @@ def _load_records(args: argparse.Namespace) -> Tuple[List[Dict[str, Any]], List[ eval_records = [dict(r) for r in records[:eval_n]] pool = records[eval_n:] train_n = args.n if args.n > 0 else len(pool) - # Cross-problem rubric: keep the train_n problems with the strongest bag-of-words - # neighbour (dense same-type pairs), so a problem's rubric can transfer to its analogue. - # Otherwise keep the first train_n (already shuffled). - train_records = ([dict(r) for r in select_paired_subset(pool, train_n, args.seed)] - if args.xproblem_rubric else [dict(r) for r in pool[:train_n]]) + # Cross-problem rubric: pick the train_n problems with the strongest qualifying neighbour + # (dense same-type pairs) and build {Q -> (P, sim)} in one index pass. Otherwise keep the + # first train_n (already shuffled) with no neighbours. + if args.xproblem_rubric: + subset, neighbor_map = build_pairs(pool, train_n, args.seed) + train_records = [dict(r) for r in subset] + pool_answers = {r['problem']: str(r.get('reference_answer', '')) for r in pool} + else: + train_records, neighbor_map, pool_answers = [dict(r) for r in pool[:train_n]], {}, {} if {r['problem'] for r in train_records} & {r['problem'] for r in eval_records}: raise ValueError('eval/train overlap detected') stats = {'raw_loaded': raw_n, 'numeric_dropped': dropped, 'train_records': len(train_records), 'eval_records': len(eval_records)} - return train_records, eval_records, stats + return train_records, eval_records, neighbor_map, pool_answers, stats # =========================================================================== @@ -943,15 +955,17 @@ def _assign_advantages(hard: List[Dict[str, Any]], args: argparse.Namespace) -> def apply_neighbor_rubric(base_sampler, chunk: List[Dict[str, Any]], - neighbor_map: Dict[str, Tuple[str, float]], base_dp: int, + neighbor_map: Dict[str, Tuple[str, float]], + pool_answers: Dict[str, str], base_dp: int, args: argparse.Namespace, checker, base_cache: DiskCache, rubric_cache: DiskCache) -> None: """Cross-problem rubric (--xproblem-rubric): for every view-A problem Q, replace its own rubric with the rubric of its bag-of-words neighbour P (Q keeps being the solved/scored - problem). P is baselined + diagnosed here (both disk-cached) as a throwaway stub, then - P's findings are copied onto Q, along with the neighbour's text + similarity for audit. - A P that can't be diagnosed -> Q degrades to view B. This makes any answer leaked from - P's rubric useless for solving Q.""" + problem). P is baselined + diagnosed here (both disk-cached) as a stub carrying P's REAL + answer, so P's baseline grades correctly and legitimately shares the baseline cache with + P-as-training-problem. P's findings are copied onto Q with the neighbour text + similarity + for audit. A P that can't be diagnosed -> Q degrades to view B. Since P's answer differs + from Q's (guaranteed at pairing time), any answer P's rubric leaks is useless for Q.""" targets = [r for r in chunk if r.get('_view') == 'A' and neighbor_map.get(r['problem'])] if not targets: return @@ -959,7 +973,7 @@ def apply_neighbor_rubric(base_sampler, chunk: List[Dict[str, Any]], for r in targets: p, _ = neighbor_map[r['problem']] if p not in by_problem: - stub = {'problem': p, 'reference_answer': '', '_view': 'A'} + stub = {'problem': p, 'reference_answer': pool_answers.get(p, ''), '_view': 'A'} by_problem[p] = stub stubs.append(stub) baseline_rollout(base_sampler, stubs, base_dp, args, base_cache) @@ -973,7 +987,8 @@ def apply_neighbor_rubric(base_sampler, chunk: List[Dict[str, Any]], def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], ci: int, base_dp: int, skill_dp: int, args: argparse.Namespace, checker, rubric_cache: DiskCache, base_cache: DiskCache = None, - neighbor_map: Optional[Dict[str, str]] = None + neighbor_map: Optional[Dict[str, Tuple[str, float]]] = None, + pool_answers: Optional[Dict[str, str]] = None ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], List[Dict[str, Any]]]: """view assign -> rubric-check (view A) -> skill-gen -> leak-filter -> with-skill greedy pass -> GRPO advantages. ``chunk`` arrives already baselined by draw_chunk. @@ -982,8 +997,8 @@ def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], for r in hard: r['_view'], r['_rubric_diag'] = _assign_view(r['problem'], args), '' if args.xproblem_rubric and neighbor_map: - apply_neighbor_rubric(base_sampler, hard, neighbor_map, base_dp, args, checker, - base_cache, rubric_cache) + apply_neighbor_rubric(base_sampler, hard, neighbor_map, pool_answers or {}, base_dp, + args, checker, base_cache, rubric_cache) else: diagnose_views(checker, hard, args, rubric_cache) @@ -1456,7 +1471,7 @@ def _write(handle, row: Dict[str, Any]) -> None: def main() -> None: args = _build_args() - records, eval_records, data_stats = _load_records(args) + records, eval_records, neighbor_map, pool_answers, data_stats = _load_records(args) if len(records) < args.chunk_size: raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded ({len(records)}); raise --n') @@ -1491,11 +1506,7 @@ def main() -> None: base_cache = DiskCache(os.path.join(cache_dir, 'baseline.jsonl'), use_cache) eval_base_cache = DiskCache(os.path.join(cache_dir, 'eval_baseline.jsonl'), use_cache) rubric_cache = DiskCache(os.path.join(cache_dir, 'rubric.jsonl'), use_cache) - neighbor_map: Dict[str, str] = {} if args.xproblem_rubric: - neighbor_cache = DiskCache(os.path.join(cache_dir, 'bow_neighbors.jsonl'), use_cache) - neighbor_map = build_neighbor_map([r['problem'] for r in records], neighbor_cache) - neighbor_cache.close() sys.stderr.write(f'[rft] xproblem-rubric ON: {len(neighbor_map)} bag-of-words pairs\n') cfg = {'record_type': 'config', 'model': MODEL_ID, 'dataset': args.dataset, @@ -1550,7 +1561,7 @@ def _prefetch(peeked: List[Dict[str, Any]]) -> None: pending = prefetch_pool.submit(_prefetch, peeked) full, summary, groups = process_chunk( base_sampler, skill_sampler, chunk, gstep, base_dp, skill_dp, - args, checker, rubric_cache, base_cache, neighbor_map) + args, checker, rubric_cache, base_cache, neighbor_map, pool_answers) summary['balance'] = balance log = None From 0ee0ebe45ba2ee93bb28325b2f557069ef0ece20 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Thu, 16 Jul 2026 11:17:35 +0800 Subject: [PATCH 15/60] fix --- cookbook/exp/embedding/train_reflexion_skill.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cookbook/exp/embedding/train_reflexion_skill.sh b/cookbook/exp/embedding/train_reflexion_skill.sh index f5fde814d..862172711 100755 --- a/cookbook/exp/embedding/train_reflexion_skill.sh +++ b/cookbook/exp/embedding/train_reflexion_skill.sh @@ -45,7 +45,7 @@ python cookbook/exp/embedding/train_reflexion_skill.py \ --sft-batch-size 8 \ --grpo-epsilon 0.2 \ --format-in-reward \ - --lr 6e-6 \ + --lr 1e-6 \ --max-train-rounds 1500 \ --save-rounds 200 \ --trend-every 10 \ From 2c3e4c8853fe213aa42349dc6b73443583ea5fa2 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Fri, 17 Jul 2026 13:56:56 +0800 Subject: [PATCH 16/60] update --- .../exp/embedding/train_reflexion_skill.py | 268 +++++++++++++----- .../exp/embedding/train_reflexion_skill.sh | 21 +- 2 files changed, 203 insertions(+), 86 deletions(-) diff --git a/cookbook/exp/embedding/train_reflexion_skill.py b/cookbook/exp/embedding/train_reflexion_skill.py index f76108ce7..5734adf22 100644 --- a/cookbook/exp/embedding/train_reflexion_skill.py +++ b/cookbook/exp/embedding/train_reflexion_skill.py @@ -14,7 +14,10 @@ Base greedy rollouts and rubric diagnoses are cached to disk (jsonl, md5-keyed) so restarts skip them; skill-gen is on-policy and never cached. -8 GPUs: ranks 0-3 train (FSDP2), 4-5 skill_sampler (synced), 6-7 base_sampler (frozen). +8 GPUs default for high-memory cards: rank 0 trains the actor, rank 1 hosts a +frozen reference model, ranks 2-3 skill_sampler (synced), ranks 4-7 base_sampler +(frozen). Override TRAIN_GPUS / REF_GPUS / SKILL_SAMPLER_GPUS / BASE_SAMPLER_GPUS +for other layouts. Leak / rubric use the backup teacher API: set LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / LLM_BACKUP_MODEL. @@ -64,13 +67,25 @@ AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') MATH_DATASET_ID = os.environ.get('MATH_DATASET_ID', 'modelscope/competition_math') -# GPU layout: train (FSDP2) + skill_sampler (synced) + base_sampler (frozen). -TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 4)) +# GPU layout: train actor (synced to skill_sampler) + frozen ref model + skill_sampler + base_sampler. +# High-memory cards can keep Qwen3-4B actor/ref as one GPU each and spend more GPUs +# on vLLM data-parallel sampling. The base side is heavier here because every +# clean skill candidate is re-solved by the frozen base model, plus balance/eval baselines. +TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 2)) +REF_GPUS = int(os.environ.get('REF_GPUS', 2)) SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 2)) -NUM_GPUS = TRAIN_GPUS + SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS -TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', 2)) -TRAIN_DP = max(1, TRAIN_GPUS // TRAIN_FSDP) +NUM_GPUS = TRAIN_GPUS + REF_GPUS + SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS +TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', min(1, TRAIN_GPUS))) +REF_FSDP = int(os.environ.get('REF_FSDP', min(1, REF_GPUS))) +if min(TRAIN_GPUS, REF_GPUS, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS, TRAIN_FSDP, REF_FSDP) < 1: + raise ValueError('TRAIN_GPUS, REF_GPUS, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS, TRAIN_FSDP and REF_FSDP must all be >= 1') +if TRAIN_GPUS % TRAIN_FSDP != 0: + raise ValueError(f'TRAIN_GPUS ({TRAIN_GPUS}) must be divisible by TRAIN_FSDP ({TRAIN_FSDP})') +if REF_GPUS % REF_FSDP != 0: + raise ValueError(f'REF_GPUS ({REF_GPUS}) must be divisible by REF_FSDP ({REF_FSDP})') +TRAIN_DP = TRAIN_GPUS // TRAIN_FSDP +REF_DP = REF_GPUS // REF_FSDP # =========================================================================== @@ -246,54 +261,52 @@ def build_direct_prompt(problem: str) -> Dict[str, Any]: def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: + skill = (skill or '').strip() + if not skill: + return build_direct_prompt(problem) # Concatenation (not .format): DIRECT_SYSTEM/skill contain literal braces. return {'messages': [ {'role': 'system', 'content': _SKILL_SOLVE_PREFIX + skill + _SKILL_SOLVE_SUFFIX}, {'role': 'user', 'content': problem}]} -# -- skill-gen prompts (view A: problem + rubric findings; view B: query only) -- +# -- skill-gen prompts (view A: problem + rubric source; view B: query only) -- SKILL_GEN_SYSTEM = ( - 'You are a math problem-solving coach. You are given a problem and an automated ' - 'process-check of an earlier attempt at it (which criteria it passed or failed, with ' - 'suggested fixes). Use the check to see where a solver of this problem tends to go ' - 'wrong, then give reusable guidance that helps a separate solver avoid those failures ' - 'on this and similar problems.\n' - '- Do not solve it or state the final answer or any specific intermediate value.\n' - '- Turn each relevant failure into general guidance (the method to reach for, the ' - 'pitfall to watch, a quick check) rather than narrating this attempt.\n' - 'Think briefly, then output only the guidance wrapped in and .') + 'You are a math guidance writer. You are given a target problem and an automated ' + 'process-check from a related problem. Use it as context to write reusable guidance ' + 'for this and similar problems.\n' + 'Output only a non-empty ... block.') SKILL_GEN_SYSTEM_Q = ( - 'You are a math problem-solving coach. Read the problem below and give reusable ' - 'guidance that would help a separate solver reach the answer on this and similar ' - 'problems.\n' - '- Do not solve it or state the final answer or any specific intermediate value.\n' - '- Name the key idea or method to reach for and the main pitfall to avoid.\n' - 'Think briefly, then output only the guidance wrapped in and .') + 'You are a math guidance writer. Given the problem below, write reusable guidance ' + 'for this and similar problems.\n' + 'Output only a non-empty ... block.') SKILL_GEN_USER_Q = ( 'Problem:\n{problem}\n\n' - 'Now output the guidance.') + '') SKILL_GEN_USER_RUBRIC = ( - 'Problem:\n{problem}\n\n' - 'Process check of an earlier attempt (automated rubric verifier -- treat as ' - 'evidence, not gospel; PASS/FAIL per criterion with suggested fixes for failures):\n' + 'Target problem:\n{problem}\n\n' + 'Problem used for the process check:\n{rubric_problem}\n\n' + 'Process check:\n' '{diagnosis}\n\n' - 'Now output the guidance.') + '') -def _skillgen_messages(problem: str, view: str, diagnosis: str) -> List[Dict[str, Any]]: +def _skillgen_messages(problem: str, view: str, diagnosis: str, + rubric_problem: str = '') -> List[Dict[str, Any]]: """Single source of truth for the skill-gen prompt (used at BOTH generation and - training so they never diverge). View A with a localisable failure uses problem + - rubric findings; view B -- or a view-A problem whose rubric flagged nothing (no - ``[FAIL]``) -- degrades to the query-only prompt.""" + training so they never diverge). View A with a localisable failure uses the target + problem plus the rubric source problem and findings; view B -- or a view-A problem + whose rubric flagged nothing (no ``[FAIL]``) -- degrades to the query-only prompt.""" if view == 'B' or '[FAIL]' not in (diagnosis or ''): return [{'role': 'system', 'content': SKILL_GEN_SYSTEM_Q}, {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=problem)}] + rubric_problem = rubric_problem or problem return [{'role': 'system', 'content': SKILL_GEN_SYSTEM}, - {'role': 'user', 'content': SKILL_GEN_USER_RUBRIC.format(problem=problem, diagnosis=diagnosis)}] + {'role': 'user', 'content': SKILL_GEN_USER_RUBRIC.format( + problem=problem, rubric_problem=rubric_problem, diagnosis=diagnosis)}] def _assign_view(problem: str, args: argparse.Namespace) -> str: @@ -302,7 +315,8 @@ def _assign_view(problem: str, args: argparse.Namespace) -> str: def _view_prompt(r: Dict[str, Any]) -> Dict[str, Any]: - return {'messages': _skillgen_messages(r['problem'], r['_view'], r.get('_rubric_diag', ''))} + return {'messages': _skillgen_messages( + r['problem'], r['_view'], r.get('_rubric_diag', ''), r.get('_rubric_src', ''))} _SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') @@ -329,7 +343,9 @@ def _extract_skills_block(text: str) -> Optional[str]: return None inner = s + len('') e = low_a.find('', inner) - block = (answer[inner:e] if e >= 0 else answer[inner:]).strip() + if e < 0: + return None + block = answer[inner:e].strip() block = re.sub(r'', '', block, flags=re.IGNORECASE).strip() return block or None @@ -386,7 +402,7 @@ def load_problems(dataset: str, n: int, seed: int, num_proc: int = 0) -> List[Di ``num_proc`` defaults to all cores (set 1 to force serial).""" ds_id = AOPS_DATASET_ID if dataset == 'aops' else MATH_DATASET_ID ds = Dataset(DatasetMeta(dataset_id=f'ms://{ds_id}', split='train')) - nproc = num_proc if num_proc > 0 else (os.cpu_count() or 1) + nproc = num_proc if num_proc > 0 else min(32, os.cpu_count() or 1) ds.map(lambda rows: _boxed_batch(rows, dataset), num_proc=nproc) ds.filter(lambda row: row['_keep'], num_proc=nproc) has_level = 'level' in ds.dataset.column_names @@ -940,8 +956,7 @@ def _assign_advantages(hard: List[Dict[str, Any]], args: argparse.Namespace) -> for r in hard: for c in r['_cands']: c['advantage'], c['grpo_adv'], c['kept'] = 0.0, 0.0, False - cs = ([c for c in r['_cands'] if c.get('reward') is not None] if args.format_in_reward - else [c for c in r['_cands'] if c['leaked'] is False and c.get('reward') is not None]) + cs = [c for c in r['_cands'] if c.get('reward') is not None] if len(cs) < 2: continue rewards = [c['reward'] for c in cs] @@ -1031,10 +1046,9 @@ def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], still.append(r) pending = still - # leak filter: deterministic verbatim-answer check only (no LLM). Measured on prior - # runs, the LLM judge flagged ~2% and leak rate was view-independent, so we drop the - # teacher-served LeakVerifier and just block the one exploit the reward can't catch -- - # the final answer written into the skill (see _answer_leaked). + # leak audit: deterministic verbatim-answer check only (no LLM). This is observability + # only: it records leak metrics for swanlab/jsonl, but does not block scoring, reward, + # advantage assignment, or training sample selection. for r, c in flat: leaked = _answer_leaked(c['skills'], r['reference_answer']) c['leaked'] = leaked @@ -1042,16 +1056,16 @@ def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], c['leak_source'] = 'deterministic' # with-skill greedy pass (T=0, M=1); reward = correct, absolute (group mean is baseline). - clean = [(r, c) for r, c in flat if c['leaked'] is False] - if clean: + scored_inputs = flat + if scored_inputs: ws_out = _run_samples(base_sampler, - [build_skill_solve_prompt(r['problem'], c['skills']) for r, c in clean], + [build_skill_solve_prompt(r['problem'], c['skills']) for r, c in scored_inputs], 1, args.max_tokens, base_dp, temperature=0.0) - for (r, c), seqs in zip(clean, ws_out): + for (r, c), seqs in zip(scored_inputs, ws_out): c['rolls'] = [_parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll()] c['with_pass'] = 1.0 if c['rolls'][0]['correct'] else 0.0 c['reward'] = c['with_pass'] - if args.format_in_reward: # unparseable/leaked score 0 and still join the group + if args.format_in_reward: # unparseable candidates score 0 and still join the group for r in hard: for c in r['_cands']: if c['reward'] is None: @@ -1068,12 +1082,11 @@ def _roll(x: Dict[str, Any]) -> Dict[str, Any]: def _is_trainable(c: Dict[str, Any], args: argparse.Namespace) -> bool: - """Reaches the GRPO update iff advantage is non-zero (and, without --format-in-reward, - also clean and scored).""" + """Reaches the GRPO update iff advantage is non-zero. Leak flags are audit-only.""" adv_nz = abs(c.get('advantage') or 0.0) > 1e-9 if args.format_in_reward: return adv_nz - return c['leaked'] is False and c.get('with_pass') is not None and adv_nz + return c.get('with_pass') is not None and adv_nz def _full_record(r: Dict[str, Any], ci: int) -> Dict[str, Any]: @@ -1107,8 +1120,7 @@ def _view_stats(problems: List[Dict[str, Any]], view: str) -> Dict[str, Any]: cands = [c for r in pv for c in r['_cands'] if c['parseable']] clean = [c for c in cands if c['leaked'] is False] adopted = sum(1 for r in pv - if any(c['leaked'] is False and abs(c.get('advantage') or 0.0) > 1e-9 - for c in r['_cands'])) + if any(abs(c.get('advantage') or 0.0) > 1e-9 for c in r['_cands'])) return {'n': len(pv), 'n_candidates_parseable': len(cands), 'n_clean': len(clean), 'n_adopted_problems': adopted, 'adoption_rate': (adopted / len(pv)) if pv else 0.0} @@ -1231,32 +1243,61 @@ def _train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: train/inference match) + the generated (think + skills) response. ``key_rounds`` selects the final assistant turn; Template masks the prompt and trains the whole response (the key-round prefix already excludes the prompt-provided ).""" - msgs = _skillgen_messages(rec['problem'], rec.get('view', 'A'), rec.get('diagnosis', '')) + msgs = _skillgen_messages( + rec['problem'], rec.get('view', 'A'), rec.get('diagnosis', ''), rec.get('rubric_src', '')) return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], 'user_data': {'key_rounds': [len(msgs)]}} -def _train_chunk(skill_model, ckpt: CheckpointEngineManager, samples: List[Dict[str, Any]], +def _train_chunk(skill_model, ref_model, ckpt: CheckpointEngineManager, samples: List[Dict[str, Any]], args: argparse.Namespace) -> Dict[str, Any]: - """One on-policy GRPO optimizer step on this chunk, then sync weights. Skills were - sampled from the current policy (ratio ~1, no old_logps); all micro-batches accumulate - into one step. The batch is padded to a multiple of ``sft_batch_size`` with - advantage-0 copies that contribute no gradient.""" + """On-policy GRPO update over one chunk, then sync weights. Micro-batches of + ``sft_batch_size`` accumulate gradients; ONE optimizer step is taken per PPO + mini-batch of ``ppo_mini_batch_size`` samples (0 -> a single step over the whole + chunk, the original behaviour). A frozen reference model provides ref_logps for the + SEAM-style KL penalty. + + Multi-step correctness: with more than one step over the SAME rollout, later + mini-batches see an already-updated policy, so we FREEZE the sampling-policy + ``old_logps`` (recomputed once, before any step) and let GRPOLoss form the PPO ratio + against them. A single step keeps ``old_logps=None`` (ratio==1, pure on-policy). + The batch is padded to a multiple of ``sft_batch_size`` with advantage-0 copies that + contribute no policy gradient.""" trajs = [_train_trajectory(rec) for rec in samples] advs = [float(rec['advantage']) for rec in samples] rem = (-len(trajs)) % args.sft_batch_size if rem: trajs += [trajs[-1]] * rem advs += [0.0] * rem - micro = 0 - for i in range(0, len(trajs), args.sft_batch_size): - skill_model.forward_backward(inputs=trajs[i:i + args.sft_batch_size], - advantages=advs[i:i + args.sft_batch_size]) - micro += 1 - skill_model.clip_grad_and_step() + + n, sft = len(trajs), args.sft_batch_size + mini = args.ppo_mini_batch_size if args.ppo_mini_batch_size > 0 else n + mini = max(sft, (mini // sft) * sft) # align to a whole number of micro-batches + multi_step = mini < n + + # Freeze ref_logps (frozen model) and -- only when we take multiple steps -- the + # sampling-policy old_logps, BOTH before any optimizer step touches the weights. With + # a single step old_logps stays None so GRPOLoss uses ratio==1 (pure on-policy). + micro_ref, micro_old = [], [] + for i in range(0, n, sft): + mb = trajs[i:i + sft] + micro_ref.append(ref_model.forward_only(inputs=mb).get('logps')) + micro_old.append(skill_model.forward_only(inputs=mb).get('logps') if multi_step else None) + + micro, n_steps = 0, 0 + for ms in range(0, n, mini): + for i in range(ms, min(ms + mini, n), sft): + k = i // sft + skill_model.forward_backward(inputs=trajs[i:i + sft], + advantages=advs[i:i + sft], + old_logps=micro_old[k], + ref_logps=micro_ref[k]) + micro += 1 + skill_model.clip_grad_and_step() + n_steps += 1 ckpt.sync_weights(merge_and_sync=True) metric = skill_model.calculate_metric(is_training=True) - return {'n_samples': len(samples), 'n_steps': 1, 'n_micro_batches': micro, + return {'n_samples': len(samples), 'n_steps': n_steps, 'n_micro_batches': micro, 'metric': {k: (float(v) if _is_num(v) else v) for k, v in (metric or {}).items()}} @@ -1370,17 +1411,46 @@ def _swan_metrics(summary: Dict[str, Any], log: Optional[Dict[str, Any]]) -> Dic return d +def _view_a_rubric_leak_metrics(chunk: List[Dict[str, Any]], + pool_answers: Optional[Dict[str, str]] = None) -> Dict[str, float]: + """Swanlab-only audit for answer leakage in view-A rubric text. This never changes + rewards, advantages, filtering, or training records.""" + view_a = [r for r in chunk if r.get('_view') == 'A'] + with_diag = [r for r in view_a if r.get('_rubric_diag')] + target_leaks = sum(1 for r in with_diag + if _answer_leaked(r.get('_rubric_diag', ''), r.get('reference_answer', ''))) + source_leaks = 0 + pool_answers = pool_answers or {} + for r in with_diag: + src = r.get('_rubric_src') + src_ref = pool_answers.get(src, '') if src else r.get('reference_answer', '') + if _answer_leaked(r.get('_rubric_diag', ''), src_ref): + source_leaks += 1 + n = len(with_diag) + return { + 'rubric_leak/n_view_a': float(len(view_a)), + 'rubric_leak/n_checked': float(n), + 'rubric_leak/target_answer_n': float(target_leaks), + 'rubric_leak/target_answer_rate': (target_leaks / n) if n else 0.0, + 'rubric_leak/source_answer_n': float(source_leaks), + 'rubric_leak/source_answer_rate': (source_leaks / n) if n else 0.0, + } + + # =========================================================================== # Block F -- components, args, main # =========================================================================== def init_components(args: argparse.Namespace): - """8 GPUs: ranks 0-3 train (FSDP2), 4-5 skill_sampler (synced), 6-7 base_sampler - (frozen). Returns (skill_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp).""" - r0, r1, r2 = TRAIN_GPUS, TRAIN_GPUS + SKILL_SAMPLER_GPUS, NUM_GPUS + """Default 8-GPU layout: rank 0 trains the actor, rank 1 hosts a frozen ref model, + 2-3 skill_sampler (synced), 4-7 base_sampler (frozen). Returns + (skill_model, ref_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp).""" + r0, r1 = TRAIN_GPUS, TRAIN_GPUS + REF_GPUS + r2, r3 = r1 + SKILL_SAMPLER_GPUS, NUM_GPUS twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, lazy_collect=False, groups=[ DeviceGroup(name='train', ranks=list(range(0, r0)), device_type='GPU'), - DeviceGroup(name='skill_sampler', ranks=list(range(r0, r1)), device_type='GPU'), - DeviceGroup(name='base_sampler', ranks=list(range(r1, r2)), device_type='GPU')]) + DeviceGroup(name='ref', ranks=list(range(r0, r1)), device_type='GPU'), + DeviceGroup(name='skill_sampler', ranks=list(range(r1, r2)), device_type='GPU'), + DeviceGroup(name='base_sampler', ranks=list(range(r2, r3)), device_type='GPU')]) train_mesh = DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_DP, fsdp_size=TRAIN_FSDP) skill_model = TransformersModel(model_id=MODEL_ID, device_mesh=train_mesh, remote_group='train', @@ -1389,11 +1459,20 @@ def init_components(args: argparse.Namespace): skill_model.set_template(Template, model_id=MODEL_ID, enable_thinking=True, max_length=args.max_model_len, truncation_strategy='delete') skill_model.set_processor(InputProcessor, padding_free=False) - skill_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) + skill_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon, beta=args.kl_beta) skill_model.set_optimizer('AdamW', lr=args.lr) skill_model.set_lr_scheduler('CosineWarmupScheduler', num_warmup_steps=10, num_training_steps=args.max_train_rounds) + ref_mesh = DeviceMesh.from_sizes(world_size=REF_GPUS, dp_size=REF_DP, fsdp_size=REF_FSDP) + ref_model = TransformersModel(model_id=MODEL_ID, device_mesh=ref_mesh, remote_group='ref', + ddp_config={'find_unused_parameters': False}) + ref_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) + ref_model.set_template(Template, model_id=MODEL_ID, enable_thinking=True, + max_length=args.max_model_len, truncation_strategy='delete') + ref_model.set_processor(InputProcessor, padding_free=False) + ref_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) + def _sampler(group, world): s = vLLMSampler(model_id=MODEL_ID, engine_args={'gpu_memory_utilization': GPU_MEM, @@ -1407,7 +1486,7 @@ def _sampler(group, world): # base_sampler is shared by the main thread and the baseline-prefetch thread -> serialise. base_sampler = _LockedSampler(_sampler('base_sampler', BASE_SAMPLER_GPUS)) ckpt = CheckpointEngineManager(model=skill_model, sampler=skill_sampler) - return skill_model, skill_sampler, base_sampler, ckpt, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS + return skill_model, ref_model, skill_sampler, base_sampler, ckpt, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS def _build_args() -> argparse.Namespace: @@ -1441,8 +1520,16 @@ def _build_args() -> argparse.Namespace: p.add_argument('--skill-max-tokens', type=int, default=8192) p.add_argument('--rubric-workers', type=int, default=16) p.add_argument('--sft-batch-size', type=int, default=8, - help='Driver micro-batch before the chunk optimizer step; multiple of train dp.') + help='Driver micro-batch (gradient-accumulation unit); multiple of train dp.') + p.add_argument('--ppo-mini-batch-size', type=int, default=0, + help='Samples per optimizer step (SEAM-style PPO mini-batch). 0 = one step ' + 'over the whole chunk (on-policy, ratio==1). When >0 and smaller than ' + 'the trainable count, multiple steps are taken over the same rollout and ' + 'old_logps are frozen so the PPO ratio/clip stays valid. Rounded down to ' + 'a multiple of --sft-batch-size.') p.add_argument('--grpo-epsilon', type=float, default=0.2) + p.add_argument('--kl-beta', type=float, default=0.001, + help='SEAM-style reference KL coefficient for GRPOLoss.') p.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True) p.add_argument('--lr', type=float, default=6e-6) p.add_argument('--max-train-rounds', type=int, default=1500) @@ -1493,9 +1580,10 @@ def main() -> None: 'view_b_frac': args.view_b_frac, 'balance': args.balance, 'balance_success_frac': args.balance_success_frac, 'skill_gen_temp': args.skill_gen_temperature, - 'grpo_epsilon': args.grpo_epsilon, 'lr': args.lr}) + 'grpo_epsilon': args.grpo_epsilon, 'kl_beta': args.kl_beta, + 'lr': args.lr}) - skill_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp = init_components(args) + skill_model, ref_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp = init_components(args) checker = build_rubric_checker() if checker is None: sys.stderr.write('[rft] no LLM backup env -> view-A rubric process-check DISABLED\n') @@ -1520,11 +1608,15 @@ def main() -> None: 'advantage': 'group_relative', 'format_in_reward': args.format_in_reward, 'cache': use_cache, 'rubric_check': 'fixed_math_5crit(viewA)' if checker else 'disabled', 'xproblem_rubric': args.xproblem_rubric, - 'grpo_epsilon': args.grpo_epsilon, 'lr': args.lr, + 'grpo_epsilon': args.grpo_epsilon, 'kl_beta': args.kl_beta, 'lr': args.lr, + 'train_gpus': TRAIN_GPUS, 'ref_gpus': REF_GPUS, 'ref_fsdp': REF_FSDP, + 'train_fsdp': TRAIN_FSDP, 'train_dp': TRAIN_DP, + 'skill_sampler_gpus': SKILL_SAMPLER_GPUS, 'base_sampler_gpus': BASE_SAMPLER_GPUS, 'max_train_rounds': args.max_train_rounds, 'started': int(time.time())} sys.stderr.write(f'[rft] raw={data_stats["raw_loaded"]} numeric_drop={data_stats["numeric_dropped"]} ' f'train={len(records)} eval={len(eval_records)} {args.dataset}; ' - f'train_gpus={TRAIN_GPUS} skill_dp={skill_dp} base_dp={base_dp}\n') + f'train_gpus={TRAIN_GPUS} ref_gpus={REF_GPUS} train_fsdp={TRAIN_FSDP} ' + f'train_dp={TRAIN_DP} skill_dp={skill_dp} base_dp={base_dp}\n') hist: List[Dict[str, float]] = [] rounds = 0 @@ -1549,6 +1641,24 @@ def _prefetch(peeked: List[Dict[str, Any]]) -> None: if peeked: baseline_rollout(base_sampler, peeked, base_dp, args, base_cache) + # Baseline (round 0) eval BEFORE any training: measures the untrained skill model on + # the fixed holdout so every later eval has a step-0 reference point on the same axis. + if eval_records: + eval_recs, eval_summary, eval_metrics = run_greedy_eval( + base_sampler, skill_sampler, eval_records, -1, rounds, base_dp, skill_dp, + args, eval_base_cache) + for rec in eval_recs: + _write(eval_f, rec) + _write(eval_f, eval_summary) + eval_f.flush() + if use_swan: + swanlab.log({f'eval/{k}': v for k, v in eval_metrics.items()}, step=0) + sys.stderr.write( + f'[eval] g-1 (init): n={eval_summary["n"]} viewB mean@1 ' + f'acc={eval_summary["baseline_acc_mean1"]:.3f}->{eval_summary["acc_mean1"]:.3f} ' + f'lift={eval_summary["lift_mean1"]:+.3f} ' + f'fmt={eval_summary["format_mean1"]:.2f} rounds={rounds}\n') + # Each chunk is drawn fresh + RE-GENERATED with the current policy (on-policy); # with --balance, draw_chunk keeps drawing until the base fail:success mix hits target. while rounds < args.max_train_rounds: @@ -1566,7 +1676,7 @@ def _prefetch(peeked: List[Dict[str, Any]]) -> None: log = None if groups: - log = _train_chunk(skill_model, ckpt, groups, args) + log = _train_chunk(skill_model, ref_model, ckpt, groups, args) rounds += 1 log.update({'record_type': 'train_round', 'round': rounds, 'chunk': gstep, 'epoch': pool.epoch, 'ts': int(time.time())}) @@ -1602,7 +1712,9 @@ def _prefetch(peeked: List[Dict[str, Any]]) -> None: f'A[{sa["n"]} {sa["adoption_rate"]:.2f}] B[{sb["n"]} {sb["adoption_rate"]:.2f}] ' f'rounds={rounds}\n') if use_swan: - swanlab.log(_swan_metrics(summary, log), step=gstep) + swan_metrics = _swan_metrics(summary, log) + swan_metrics.update(_view_a_rubric_leak_metrics(chunk, pool_answers)) + swanlab.log(swan_metrics, step=gstep) if eval_records and (gstep + 1) % args.eval_every == 0: eval_recs, eval_summary, eval_metrics = run_greedy_eval( diff --git a/cookbook/exp/embedding/train_reflexion_skill.sh b/cookbook/exp/embedding/train_reflexion_skill.sh index 862172711..7a76372b0 100755 --- a/cookbook/exp/embedding/train_reflexion_skill.sh +++ b/cookbook/exp/embedding/train_reflexion_skill.sh @@ -1,7 +1,10 @@ #!/bin/bash # Online GRPO RFT for the reflexion skill generator (unified, self-contained, cached). -# GPUs: 8 — ranks 0-3 train (skill model, FSDP2), 4-5 skill sampler (synced), 6-7 base -# sampler (frozen). Per chunk: base greedy solve -> rubric process-check (view A) -> +# GPUs: 8 — default high-memory layout uses rank 0 for actor training, rank 1 for +# the frozen ref model, ranks 2-3 for skill sampler (synced), and ranks 4-7 for +# base sampler (frozen). Override TRAIN_GPUS / REF_GPUS / SKILL_SAMPLER_GPUS / +# BASE_SAMPLER_GPUS for other layouts. Per chunk: base greedy +# solve -> rubric process-check (view A) -> # skill-gen (thinking ON, N candidates) -> deterministic leak filter -> with-skill greedy # pass -> group-relative advantage -> ONE on-policy GRPO step -> sync weights. # @@ -28,22 +31,24 @@ python cookbook/exp/embedding/train_reflexion_skill.py \ --dataset aops \ --n 5000 \ --numeric-only \ - --chunk-size 16 \ - --n-skills 8 \ + --chunk-size 32 \ + --n-skills 16 \ --view-b-frac 0.5 \ --xproblem-rubric \ --skill-retries 2 \ --balance \ - --balance-success-frac 0.4 \ + --balance-success-frac 0.3 \ --balance-loop-frac 0.5 \ --balance-max-draws-mult 8 \ - --max-tokens 25000 \ - --skill-max-tokens 8192 \ - --max-model-len 30000 \ + --max-tokens 8192 \ + --skill-max-tokens 4096 \ + --max-model-len 16384 \ --eval-size 128 \ --eval-every 5 \ --sft-batch-size 8 \ + --ppo-mini-batch-size 0 \ --grpo-epsilon 0.2 \ + --kl-beta 0.001 \ --format-in-reward \ --lr 1e-6 \ --max-train-rounds 1500 \ From a78a46c03018beaa3229d1b2254c0ff3dfa4acc1 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Fri, 17 Jul 2026 23:04:07 +0800 Subject: [PATCH 17/60] fix --- .../exp/embedding/train_reflexion_skill.py | 31 ++++++++++++++----- 1 file changed, 23 insertions(+), 8 deletions(-) diff --git a/cookbook/exp/embedding/train_reflexion_skill.py b/cookbook/exp/embedding/train_reflexion_skill.py index 5734adf22..765dc9c01 100644 --- a/cookbook/exp/embedding/train_reflexion_skill.py +++ b/cookbook/exp/embedding/train_reflexion_skill.py @@ -279,7 +279,7 @@ def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: SKILL_GEN_SYSTEM_Q = ( 'You are a math guidance writer. Given the problem below, write reusable guidance ' - 'for this and similar problems.\n' + 'for this and similar problems, you can recall the diagnosis or check happend before.\n' 'Output only a non-empty ... block.') SKILL_GEN_USER_Q = ( @@ -864,11 +864,17 @@ def _run(item): r, key = item seg = {'messages': [{'role': 'user', 'content': r['problem']}, {'role': 'assistant', 'content': r['_init'][0]['text']}]} - try: - return r, key, _format_diagnosis(checker.diagnose(seg, query=r['problem'])) - except Exception as exc: # teacher hiccup -> no-diagnosis prompt (not cached) - logger.warning(f'[rubric] diagnose error: {exc}') - return r, key, None + attempts = max(1, args.rubric_retries + 1) + for attempt in range(attempts): + try: + return r, key, _format_diagnosis(checker.diagnose(seg, query=r['problem'])) + except Exception as exc: # teacher hiccup -> retry, then no-diagnosis prompt (not cached) + if attempt + 1 < attempts: + logger.warning(f'[rubric] diagnose error: {exc}; retry {attempt + 1}/{args.rubric_retries}') + time.sleep(min(2.0, 0.5 * (2 ** attempt))) + continue + logger.warning(f'[rubric] diagnose error: {exc}; giving up after {attempts} attempts') + return r, key, None workers = max(1, min(args.rubric_workers, len(pending))) with ThreadPoolExecutor(max_workers=workers) as ex: @@ -1164,7 +1170,9 @@ def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespac clean = [c for c in cands if c['leaked'] is False] ws_rolls = [x for c in scored for x in c['rolls']] base_acc = (sum(r['_baseline_pass'] for r in chunk) / len(chunk)) if chunk else 0.0 - ws_acc = _mean([c['with_pass'] for c in scored]) + ws_acc = _mean([1.0 if any(c.get('reward') for c in r['_cands']) else 0.0 for r in chunk]) + cand_pass_parseable = _mean([c['with_pass'] for c in scored]) + cand_pass_all = _mean([1.0 if c.get('reward') else 0.0 for c in all_cands]) # base failure taxonomy (you asked whether skills fail because the base loops out of length) classes = [_baseline_class(r) for r in chunk] n_fail = sum(1 for c in classes if c != 'success') @@ -1188,6 +1196,8 @@ def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespac 'withskill_trunc_frac': (trunc / len(ws_rolls)) if ws_rolls else 0.0, 'avg_baseline_pass': base_acc, 'avg_withskill_pass': ws_acc, 'avg_lift': ws_acc - base_acc, + 'candidate_withskill_pass_parseable': cand_pass_parseable, + 'candidate_withskill_pass_all': cand_pass_all, 'termination_rate_withskill': _mean([1.0 if x['terminated'] else 0.0 for x in ws_rolls]), 'view_A': _view_stats(chunk, 'A'), 'view_B': _view_stats(chunk, 'B'), **_xproblem_stats(chunk, args), @@ -1393,6 +1403,8 @@ def _swan_metrics(summary: Dict[str, Any], log: Optional[Dict[str, Any]]) -> Dic if sig['n_groups'] > 0: d.update({'acc/baseline_pass': summary['avg_baseline_pass'], 'acc/withskill_pass': summary['avg_withskill_pass'], 'acc/lift': summary['avg_lift'], + 'candidate/withskill_pass_parseable': summary['candidate_withskill_pass_parseable'], + 'candidate/withskill_pass_all': summary['candidate_withskill_pass_all'], 'adopt/A': summary['view_A']['adoption_rate'], 'adopt/B': summary['view_B']['adoption_rate'], 'term/withskill': summary['termination_rate_withskill'], @@ -1456,7 +1468,7 @@ def init_components(args: argparse.Namespace): skill_model = TransformersModel(model_id=MODEL_ID, device_mesh=train_mesh, remote_group='train', ddp_config={'find_unused_parameters': False}) skill_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) - skill_model.set_template(Template, model_id=MODEL_ID, enable_thinking=True, + skill_model.set_template(Template, model_id=MODEL_ID, enable_thinking=False, max_length=args.max_model_len, truncation_strategy='delete') skill_model.set_processor(InputProcessor, padding_free=False) skill_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon, beta=args.kl_beta) @@ -1519,6 +1531,9 @@ def _build_args() -> argparse.Namespace: p.add_argument('--max-tokens', type=int, default=8192) p.add_argument('--skill-max-tokens', type=int, default=8192) p.add_argument('--rubric-workers', type=int, default=16) + p.add_argument('--rubric-retries', type=int, default=2, + help='Retry failed/timeout rubric diagnose calls this many times before ' + 'falling back to an empty diagnosis without caching the failure.') p.add_argument('--sft-batch-size', type=int, default=8, help='Driver micro-batch (gradient-accumulation unit); multiple of train dp.') p.add_argument('--ppo-mini-batch-size', type=int, default=0, From 07f05bc60d8b0551846ada5ccfc8454e1bc09032 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Sat, 18 Jul 2026 10:13:04 +0800 Subject: [PATCH 18/60] fix --- .../exp/embedding/train_reflexion_skill.py | 48 ++++++++++--------- 1 file changed, 26 insertions(+), 22 deletions(-) diff --git a/cookbook/exp/embedding/train_reflexion_skill.py b/cookbook/exp/embedding/train_reflexion_skill.py index 765dc9c01..c43e87090 100644 --- a/cookbook/exp/embedding/train_reflexion_skill.py +++ b/cookbook/exp/embedding/train_reflexion_skill.py @@ -275,23 +275,21 @@ def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: 'You are a math guidance writer. You are given a target problem and an automated ' 'process-check from a related problem. Use it as context to write reusable guidance ' 'for this and similar problems.\n' - 'Output only a non-empty ... block.') + 'Output only a non-empty:\n\nYour remind and skill here...\n\nblock.') SKILL_GEN_SYSTEM_Q = ( 'You are a math guidance writer. Given the problem below, write reusable guidance ' 'for this and similar problems, you can recall the diagnosis or check happend before.\n' - 'Output only a non-empty ... block.') + 'Output only a non-empty:\n\nYour remind and skill here...\n\nblock.') SKILL_GEN_USER_Q = ( - 'Problem:\n{problem}\n\n' - '') + 'Problem:\n{problem}\n\n') SKILL_GEN_USER_RUBRIC = ( 'Target problem:\n{problem}\n\n' 'Problem used for the process check:\n{rubric_problem}\n\n' 'Process check:\n' - '{diagnosis}\n\n' - '') + '{diagnosis}\n\n') def _skillgen_messages(problem: str, view: str, diagnosis: str, @@ -327,18 +325,18 @@ def _clean_text(decoded: Optional[str]) -> str: def _extract_skills_block(text: str) -> Optional[str]: - """Non-empty ``...`` block, or None. Requires ```` (skill-gen - runs thinking ON); reads only the answer after the last one, so a mid-reasoning draft - or a demo echo can never be mistaken for the answer. No format/wording gate beyond - that -- the reward (does the frozen executor solve a similar problem with this skill?) - is what judges skill quality (SEAM-style), so we do not lexically second-guess it.""" + """Return the non-empty content inside the last ``...`` block. + + Skill generation may run with thinking either ON or OFF. If a ```` marker is + present, parse only the text after the last marker so drafts inside thinking are ignored; + otherwise parse the full response. This keeps the parser compatible with thinking-off + skill generation while preserving the old thinking-on safety behavior. + """ low = text.lower() end_think = low.rfind('') - if end_think < 0: - return None - answer = text[end_think + len(''):] + answer = text[end_think + len(''):] if end_think >= 0 else text low_a = answer.lower() - s = low_a.find('') + s = low_a.rfind('') if s < 0: return None inner = s + len('') @@ -957,8 +955,10 @@ def draw_chunk(pool: ProblemPool, base_sampler, base_dp: int, args: argparse.Nam def _assign_advantages(hard: List[Dict[str, Any]], args: argparse.Namespace) -> None: """Group-relative advantage over each problem's scored candidates using the greedy binary reward R in {0,1}: ``A = (R - mean) / (std + eps)``. std==0 groups get no - gradient -- GRPO's variance selects informative problems (no explicit difficulty gate).""" + gradient -- GRPO's variance selects informative problems (no explicit difficulty gate). + A symmetric clip limits both positive and negative outliers from sparse 1-vs-many groups.""" eps = 1e-6 + adv_clip = abs(float(getattr(args, 'adv_clip', 0.0) or 0.0)) for r in hard: for c in r['_cands']: c['advantage'], c['grpo_adv'], c['kept'] = 0.0, 0.0, False @@ -971,7 +971,8 @@ def _assign_advantages(hard: List[Dict[str, Any]], args: argparse.Namespace) -> if std < 1e-9: continue for c in cs: - adv = (c['reward'] - mean_r) / (std + eps) + raw_adv = (c['reward'] - mean_r) / (std + eps) + adv = max(-adv_clip, min(adv_clip, raw_adv)) if adv_clip > 0 else raw_adv c['advantage'], c['grpo_adv'], c['kept'] = adv, adv, c['reward'] > mean_r @@ -1485,18 +1486,18 @@ def init_components(args: argparse.Namespace): ref_model.set_processor(InputProcessor, padding_free=False) ref_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) - def _sampler(group, world): + def _sampler(group, world, enable_thinking: bool = True): s = vLLMSampler(model_id=MODEL_ID, engine_args={'gpu_memory_utilization': GPU_MEM, 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), remote_group=group) - s.set_template(Template, model_id=MODEL_ID, enable_thinking=True, max_length=args.max_model_len) + s.set_template(Template, model_id=MODEL_ID, enable_thinking=enable_thinking, max_length=args.max_model_len) return s - skill_sampler = _sampler('skill_sampler', SKILL_SAMPLER_GPUS) + skill_sampler = _sampler('skill_sampler', SKILL_SAMPLER_GPUS, enable_thinking=False) # base_sampler is shared by the main thread and the baseline-prefetch thread -> serialise. - base_sampler = _LockedSampler(_sampler('base_sampler', BASE_SAMPLER_GPUS)) + base_sampler = _LockedSampler(_sampler('base_sampler', BASE_SAMPLER_GPUS, enable_thinking=True)) ckpt = CheckpointEngineManager(model=skill_model, sampler=skill_sampler) return skill_model, ref_model, skill_sampler, base_sampler, ckpt, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS @@ -1543,6 +1544,8 @@ def _build_args() -> argparse.Namespace: 'old_logps are frozen so the PPO ratio/clip stays valid. Rounded down to ' 'a multiple of --sft-batch-size.') p.add_argument('--grpo-epsilon', type=float, default=0.2) + p.add_argument('--adv-clip', type=float, default=3.0, + help='Symmetric clip for group-relative advantages; <=0 disables clipping.') p.add_argument('--kl-beta', type=float, default=0.001, help='SEAM-style reference KL coefficient for GRPOLoss.') p.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True) @@ -1550,7 +1553,7 @@ def _build_args() -> argparse.Namespace: p.add_argument('--max-train-rounds', type=int, default=1500) p.add_argument('--save-rounds', type=int, default=200) p.add_argument('--trend-every', type=int, default=10) - p.add_argument('--output-dir', default='./output/reflexion_skill_rft') + p.add_argument('--output-dir', default='./output/reflexion_skill') p.add_argument('--cache-dir', default='', help='Baseline/rubric cache dir (default /cache).') p.add_argument('--no-cache', action='store_true', help='Disable disk cache read/write.') p.add_argument('--prefetch-baseline', action=argparse.BooleanOptionalAction, default=True, @@ -1623,6 +1626,7 @@ def main() -> None: 'advantage': 'group_relative', 'format_in_reward': args.format_in_reward, 'cache': use_cache, 'rubric_check': 'fixed_math_5crit(viewA)' if checker else 'disabled', 'xproblem_rubric': args.xproblem_rubric, + 'adv_clip': args.adv_clip, 'grpo_epsilon': args.grpo_epsilon, 'kl_beta': args.kl_beta, 'lr': args.lr, 'train_gpus': TRAIN_GPUS, 'ref_gpus': REF_GPUS, 'ref_fsdp': REF_FSDP, 'train_fsdp': TRAIN_FSDP, 'train_dp': TRAIN_DP, From 8d3cee511f3546f654548f5ab8aabe23a9c5f541 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Sat, 18 Jul 2026 13:12:21 +0800 Subject: [PATCH 19/60] fix --- .../exp/embedding/train_reflexion_skill.py | 114 ++++++++++++------ 1 file changed, 77 insertions(+), 37 deletions(-) diff --git a/cookbook/exp/embedding/train_reflexion_skill.py b/cookbook/exp/embedding/train_reflexion_skill.py index c43e87090..240247f01 100644 --- a/cookbook/exp/embedding/train_reflexion_skill.py +++ b/cookbook/exp/embedding/train_reflexion_skill.py @@ -274,13 +274,19 @@ def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: SKILL_GEN_SYSTEM = ( 'You are a math guidance writer. You are given a target problem and an automated ' 'process-check from a related problem. Use it as context to write reusable guidance ' - 'for this and similar problems.\n' - 'Output only a non-empty:\n\nYour remind and skill here...\n\nblock.') + 'for this and similar problems.\n') SKILL_GEN_SYSTEM_Q = ( 'You are a math guidance writer. Given the problem below, write reusable guidance ' - 'for this and similar problems, you can recall the diagnosis or check happend before.\n' - 'Output only a non-empty:\n\nYour remind and skill here...\n\nblock.') + 'for this and similar problems, you need to recall the diagnosis or check happened before.\n') + +_SKILL_OUTPUT_DUAL = ( + 'Output non-empty:\n' + 'Your diagnoses from previous process checking here...\n' + 'Your reminds and skills here... blocks') + +_SKILL_OUTPUT_LEGACY = ( + 'Output a non-empty:\n\nYour remind and skill here...\n\nblock.') SKILL_GEN_USER_Q = ( 'Problem:\n{problem}\n\n') @@ -292,17 +298,22 @@ def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: '{diagnosis}\n\n') +def _skill_output_instruction(diagnose_skill_format: bool) -> str: + return _SKILL_OUTPUT_DUAL if diagnose_skill_format else _SKILL_OUTPUT_LEGACY + + def _skillgen_messages(problem: str, view: str, diagnosis: str, - rubric_problem: str = '') -> List[Dict[str, Any]]: + rubric_problem: str = '', diagnose_skill_format: bool = True) -> List[Dict[str, Any]]: """Single source of truth for the skill-gen prompt (used at BOTH generation and training so they never diverge). View A with a localisable failure uses the target problem plus the rubric source problem and findings; view B -- or a view-A problem whose rubric flagged nothing (no ``[FAIL]``) -- degrades to the query-only prompt.""" + out = _skill_output_instruction(diagnose_skill_format) if view == 'B' or '[FAIL]' not in (diagnosis or ''): - return [{'role': 'system', 'content': SKILL_GEN_SYSTEM_Q}, + return [{'role': 'system', 'content': SKILL_GEN_SYSTEM_Q + out}, {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=problem)}] rubric_problem = rubric_problem or problem - return [{'role': 'system', 'content': SKILL_GEN_SYSTEM}, + return [{'role': 'system', 'content': SKILL_GEN_SYSTEM + out}, {'role': 'user', 'content': SKILL_GEN_USER_RUBRIC.format( problem=problem, rubric_problem=rubric_problem, diagnosis=diagnosis)}] @@ -312,9 +323,10 @@ def _assign_view(problem: str, args: argparse.Namespace) -> str: return 'B' if (h % 100000) / 100000.0 < args.view_b_frac else 'A' -def _view_prompt(r: Dict[str, Any]) -> Dict[str, Any]: +def _view_prompt(r: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: return {'messages': _skillgen_messages( - r['problem'], r['_view'], r.get('_rubric_diag', ''), r.get('_rubric_src', ''))} + r['problem'], r['_view'], r.get('_rubric_diag', ''), r.get('_rubric_src', ''), + args.diagnose_skill_format)} _SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') @@ -324,30 +336,44 @@ def _clean_text(decoded: Optional[str]) -> str: return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() -def _extract_skills_block(text: str) -> Optional[str]: - """Return the non-empty content inside the last ``...`` block. - - Skill generation may run with thinking either ON or OFF. If a ```` marker is - present, parse only the text after the last marker so drafts inside thinking are ignored; - otherwise parse the full response. This keeps the parser compatible with thinking-off - skill generation while preserving the old thinking-on safety behavior. - """ - low = text.lower() - end_think = low.rfind('') - answer = text[end_think + len(''):] if end_think >= 0 else text - low_a = answer.lower() - s = low_a.rfind('') +def _extract_tag_block(answer: str, tag: str) -> Optional[str]: + low = answer.lower() + open_tag, close_tag = f'<{tag}>', f'' + s = low.rfind(open_tag) if s < 0: return None - inner = s + len('') - e = low_a.find('', inner) + inner = s + len(open_tag) + e = low.find(close_tag, inner) if e < 0: return None block = answer[inner:e].strip() - block = re.sub(r'', '', block, flags=re.IGNORECASE).strip() + block = re.sub(r'', '', block, flags=re.IGNORECASE).strip() return block or None +def _extract_skill_blocks(text: str, diagnose_skill_format: bool = True) -> Optional[Dict[str, str]]: + """Parse skill-generation output. + + With diagnose-skill format enabled, both ```` and ```` must be + present and non-empty; their inner texts are concatenated for executor injection. With + the legacy format, a non-empty ```` block is enough. If a ```` marker is + present, parse only the text after the last marker; otherwise parse the full response. + """ + low = text.lower() + end_think = low.rfind('') + answer = text[end_think + len(''):] if end_think >= 0 else text + if diagnose_skill_format: + diagnose = _extract_tag_block(answer, 'diagnose') + skill = _extract_tag_block(answer, 'skill') + if not diagnose or not skill: + return None + return {'diagnose': diagnose, 'skill': skill, 'skills': f'{diagnose}\n\n{skill}'} + block = _extract_tag_block(answer, 'skills') + if not block: + return None + return {'diagnose': '', 'skill': block, 'skills': block} + + def _parse_seq(seq, gold: str) -> Dict[str, Any]: """Grade one sampled sequence into a rollout record.""" text = _clean_text(getattr(seq, 'decoded', '') or '') @@ -1024,13 +1050,13 @@ def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], else: diagnose_views(checker, hard, args, rubric_cache) - # skill-gen (thinking ON), per-view prompt; re-sample problems with no clean candidate. + # skill-gen (thinking OFF), per-view prompt; re-sample problems with no clean candidate. flat: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] pending = list(hard) for _ in range(args.skill_retries + 1): if not pending: break - sg_out = _run_samples(skill_sampler, [_view_prompt(r) for r in pending], + sg_out = _run_samples(skill_sampler, [_view_prompt(r, args) for r in pending], args.n_skills, args.skill_max_tokens, skill_dp, temperature=args.skill_gen_temperature, top_p=args.skill_gen_top_p, top_k=args.skill_gen_top_k) @@ -1039,8 +1065,11 @@ def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], got = False for s in seqs: resp = _clean_text(getattr(s, 'decoded', '') or '') - block = _extract_skills_block(resp) - cand = {'skills': block or '', 'response': resp, 'parseable': bool(block), + parsed = _extract_skill_blocks(resp, args.diagnose_skill_format) + block = parsed['skills'] if parsed else '' + cand = {'skills': block, 'diagnose': (parsed or {}).get('diagnose', ''), + 'skill': (parsed or {}).get('skill', ''), + 'response': resp, 'parseable': bool(parsed), 'view': r['_view'], 'leaked': None, 'leak_reason': '', 'leak_source': '', 'with_pass': None, 'reward': None, 'rolls': [], 'skillgen_stop': getattr(s, 'stop_reason', None), @@ -1112,7 +1141,8 @@ def _full_record(r: Dict[str, Any], ci: int) -> Dict[str, Any]: 'rubric_src': r.get('_rubric_src', ''), 'neighbor_sim': r.get('_neighbor_sim'), 'baseline_rolls': [_roll(x) for x in r['_baseline_rolls']], 'candidates': [{ - 'skills': c['skills'], 'response': c['response'], 'parseable': c['parseable'], + 'skills': c['skills'], 'diagnose': c.get('diagnose', ''), 'skill': c.get('skill', ''), + 'response': c['response'], 'parseable': c['parseable'], 'skillgen_stop': c.get('skillgen_stop'), 'skillgen_tokens': c.get('skillgen_tokens'), 'leaked': c['leaked'], 'leak_reason': c['leak_reason'], 'leak_source': c['leak_source'], 'with_pass': c['with_pass'], 'reward': c.get('reward'), 'advantage': c.get('advantage'), @@ -1232,6 +1262,8 @@ def _group_records(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> Lis 'view': r.get('_view', 'A'), 'diagnosis': r.get('_rubric_diag', ''), 'rubric_src': r.get('_rubric_src', ''), 'response': c['response'], 'skills': c['skills'], + 'diagnose': c.get('diagnose', ''), 'skill': c.get('skill', ''), + 'diagnose_skill_format': args.diagnose_skill_format, 'skillgen_stop': c.get('skillgen_stop'), 'advantage': c['advantage'], 'grpo_adv': c['grpo_adv'], 'kept': c['kept'], 'reward': c['reward'], 'with_pass': c['with_pass']}) @@ -1255,7 +1287,8 @@ def _train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: selects the final assistant turn; Template masks the prompt and trains the whole response (the key-round prefix already excludes the prompt-provided ).""" msgs = _skillgen_messages( - rec['problem'], rec.get('view', 'A'), rec.get('diagnosis', ''), rec.get('rubric_src', '')) + rec['problem'], rec.get('view', 'A'), rec.get('diagnosis', ''), rec.get('rubric_src', ''), + rec.get('diagnose_skill_format', True)) return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], 'user_data': {'key_rounds': [len(msgs)]}} @@ -1326,22 +1359,25 @@ def run_greedy_eval(base_sampler, skill_sampler, eval_records: List[Dict[str, An baseline_rollout(base_sampler, eval_records, base_dp, args, base_cache) for r in eval_records: r['_view'], r['_rubric_diag'] = 'B', '' - sg_out = _run_samples(skill_sampler, [_view_prompt(r) for r in eval_records], + sg_out = _run_samples(skill_sampler, [_view_prompt(r, args) for r in eval_records], 1, args.skill_max_tokens, skill_dp, temperature=0.0) - skills = [(_extract_skills_block(_clean_text(getattr(seqs[0], 'decoded', '') or '')) or '', - _clean_text(getattr(seqs[0], 'decoded', '') or '')) if seqs else ('', '') + skills = [((parsed := _extract_skill_blocks(_clean_text(getattr(seqs[0], 'decoded', '') or ''), + args.diagnose_skill_format))['skills'] if parsed else '', + (parsed or {}).get('diagnose', ''), (parsed or {}).get('skill', ''), + _clean_text(getattr(seqs[0], 'decoded', '') or '')) if seqs else ('', '', '', '') for seqs in sg_out] ws_out = _run_samples(base_sampler, - [build_skill_solve_prompt(r['problem'], sk) for r, (sk, _) in zip(eval_records, skills)], + [build_skill_solve_prompt(r['problem'], sk) for r, (sk, _, _, _) in zip(eval_records, skills)], 1, args.max_tokens, base_dp, temperature=0.0) recs = [] - for r, (sk, sresp), seqs in zip(eval_records, skills, ws_out): + for r, (sk, diag_text, skill_text, sresp), seqs in zip(eval_records, skills, ws_out): roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() recs.append({ 'record_type': 'eval_problem', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, 'problem': r['problem'], 'reference_answer': r['reference_answer'], 'view': r['_view'], 'rubric_diag': r.get('_rubric_diag', ''), - 'baseline_pass': r['_baseline_pass'], 'skill': sk, 'skill_parseable': bool(sk), + 'baseline_pass': r['_baseline_pass'], 'skill': sk, 'skill_diagnose': diag_text, + 'skill_action': skill_text, 'skill_parseable': bool(sk), 'skill_response': sresp, 'withskill_pred': roll['pred'], 'withskill_correct': roll['correct'], 'withskill_terminated': roll['terminated'], 'withskill_stop_reason': roll['stop_reason'], 'withskill_text': roll['text'], @@ -1528,6 +1564,9 @@ def _build_args() -> argparse.Namespace: p.add_argument('--skill-gen-temperature', type=float, default=1.0) p.add_argument('--skill-gen-top-p', type=float, default=1.0) p.add_argument('--skill-gen-top-k', type=int, default=-1) + p.add_argument('--diagnose-skill-format', action=argparse.BooleanOptionalAction, default=True, + help='Require both and blocks for skill-gen format; ' + '--no-diagnose-skill-format falls back to legacy .') p.add_argument('--max-model-len', type=int, default=16384) p.add_argument('--max-tokens', type=int, default=8192) p.add_argument('--skill-max-tokens', type=int, default=8192) @@ -1620,6 +1659,7 @@ def main() -> None: 'numeric_only': args.numeric_only, 'raw_loaded': data_stats['raw_loaded'], 'numeric_dropped': data_stats['numeric_dropped'], 'eval_every': args.eval_every, 'n_skills': args.n_skills, 'view_b_frac': args.view_b_frac, + 'diagnose_skill_format': args.diagnose_skill_format, 'skill_retries': args.skill_retries, 'balance': args.balance, 'balance_success_frac': args.balance_success_frac, 'skill_gen_temp': args.skill_gen_temperature, 'reward': 'greedy_binary(correct)', From 490995a81cde532c57e0e636b625a5629a269b35 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Sat, 18 Jul 2026 13:32:43 +0800 Subject: [PATCH 20/60] fix --- cookbook/exp/embedding/train_reflexion_skill.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/cookbook/exp/embedding/train_reflexion_skill.py b/cookbook/exp/embedding/train_reflexion_skill.py index 240247f01..6c92c8d53 100644 --- a/cookbook/exp/embedding/train_reflexion_skill.py +++ b/cookbook/exp/embedding/train_reflexion_skill.py @@ -1361,11 +1361,17 @@ def run_greedy_eval(base_sampler, skill_sampler, eval_records: List[Dict[str, An r['_view'], r['_rubric_diag'] = 'B', '' sg_out = _run_samples(skill_sampler, [_view_prompt(r, args) for r in eval_records], 1, args.skill_max_tokens, skill_dp, temperature=0.0) - skills = [((parsed := _extract_skill_blocks(_clean_text(getattr(seqs[0], 'decoded', '') or ''), - args.diagnose_skill_format))['skills'] if parsed else '', - (parsed or {}).get('diagnose', ''), (parsed or {}).get('skill', ''), - _clean_text(getattr(seqs[0], 'decoded', '') or '')) if seqs else ('', '', '', '') - for seqs in sg_out] + skills = [] + for seqs in sg_out: + if not seqs: + skills.append(('', '', '', '')) + continue + sresp = _clean_text(getattr(seqs[0], 'decoded', '') or '') + parsed = _extract_skill_blocks(sresp, args.diagnose_skill_format) + skills.append((parsed['skills'] if parsed else '', + (parsed or {}).get('diagnose', ''), + (parsed or {}).get('skill', ''), + sresp)) ws_out = _run_samples(base_sampler, [build_skill_solve_prompt(r['problem'], sk) for r, (sk, _, _, _) in zip(eval_records, skills)], 1, args.max_tokens, base_dp, temperature=0.0) From 388a87df2d0548eba5184c1e916525003c411837 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Sun, 19 Jul 2026 11:13:36 +0800 Subject: [PATCH 21/60] fix --- .../exp/embedding/train_reflexion_skill.py | 98 ++++++++++++------- 1 file changed, 60 insertions(+), 38 deletions(-) diff --git a/cookbook/exp/embedding/train_reflexion_skill.py b/cookbook/exp/embedding/train_reflexion_skill.py index 6c92c8d53..39af13e21 100644 --- a/cookbook/exp/embedding/train_reflexion_skill.py +++ b/cookbook/exp/embedding/train_reflexion_skill.py @@ -9,7 +9,7 @@ within each problem-group, so std=0 groups give no gradient (GRPO variance selects). Each problem is routed to EXACTLY ONE view: A = problem + rubric findings, B = -query only (deployment form). Skill-gen trains only the final turn. +query only (deployment form). Skill-gen trains only the final structured guidance turn. Base greedy rollouts and rubric diagnoses are cached to disk (jsonl, md5-keyed) so restarts skip them; skill-gen is on-policy and never cached. @@ -274,16 +274,18 @@ def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: SKILL_GEN_SYSTEM = ( 'You are a math guidance writer. You are given a target problem and an automated ' 'process-check from a related problem. Use it as context to write reusable guidance ' - 'for this and similar problems.\n') + 'for this and similar problems. In , write likely mistakes or checks only. ' + 'Leave empty if unclear. In , write reusable solving advice.\n') SKILL_GEN_SYSTEM_Q = ( 'You are a math guidance writer. Given the problem below, write reusable guidance ' - 'for this and similar problems, you need to recall the diagnosis or check happened before.\n') + 'for this and similar problems. In , write likely mistakes or checks only. ' + 'Leave empty if unclear. In , write reusable solving advice.\n') _SKILL_OUTPUT_DUAL = ( - 'Output non-empty:\n' - 'Your diagnoses from previous process checking here...\n' - 'Your reminds and skills here... blocks') + 'Output:\n' + 'Your likely mistakes or checks here\n' + 'Your reusable solving advice here...') _SKILL_OUTPUT_LEGACY = ( 'Output a non-empty:\n\nYour remind and skill here...\n\nblock.') @@ -336,7 +338,7 @@ def _clean_text(decoded: Optional[str]) -> str: return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() -def _extract_tag_block(answer: str, tag: str) -> Optional[str]: +def _extract_tag_block(answer: str, tag: str, allow_empty: bool = False) -> Optional[str]: low = answer.lower() open_tag, close_tag = f'<{tag}>', f'' s = low.rfind(open_tag) @@ -347,31 +349,40 @@ def _extract_tag_block(answer: str, tag: str) -> Optional[str]: if e < 0: return None block = answer[inner:e].strip() - block = re.sub(r'', '', block, flags=re.IGNORECASE).strip() - return block or None + block = re.sub(r'', '', block, flags=re.IGNORECASE).strip() + return block if (block or allow_empty) else None + + +def _join_pitfall_strategy(pitfall: str, strategy: str) -> str: + pitfall = (pitfall or '').strip() + strategy = (strategy or '').strip() + return strategy if not pitfall else f'{pitfall}\n\n{strategy}' def _extract_skill_blocks(text: str, diagnose_skill_format: bool = True) -> Optional[Dict[str, str]]: """Parse skill-generation output. - With diagnose-skill format enabled, both ```` and ```` must be - present and non-empty; their inner texts are concatenated for executor injection. With - the legacy format, a non-empty ```` block is enough. If a ```` marker is - present, parse only the text after the last marker; otherwise parse the full response. + With the pitfall-strategy format enabled, both ```` and ```` tags + must be present, while ```` may be empty; their inner texts are concatenated + for executor injection. With the legacy format, a non-empty ```` block is + enough. If a ```` marker is present, parse only the text after the last marker; + otherwise parse the full response. """ low = text.lower() end_think = low.rfind('') answer = text[end_think + len(''):] if end_think >= 0 else text if diagnose_skill_format: - diagnose = _extract_tag_block(answer, 'diagnose') - skill = _extract_tag_block(answer, 'skill') - if not diagnose or not skill: + pitfall = _extract_tag_block(answer, 'pitfall', allow_empty=True) + strategy = _extract_tag_block(answer, 'strategy') + if pitfall is None or not strategy: return None - return {'diagnose': diagnose, 'skill': skill, 'skills': f'{diagnose}\n\n{skill}'} + joined = _join_pitfall_strategy(pitfall, strategy) + return {'pitfall': pitfall, 'strategy': strategy, + 'diagnose': pitfall, 'skill': strategy, 'skills': joined} block = _extract_tag_block(answer, 'skills') if not block: return None - return {'diagnose': '', 'skill': block, 'skills': block} + return {'pitfall': '', 'strategy': block, 'diagnose': '', 'skill': block, 'skills': block} def _parse_seq(seq, gold: str) -> Dict[str, Any]: @@ -746,11 +757,15 @@ def _empty_roll() -> Dict[str, Any]: 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} +def _roll_passed(roll: Dict[str, Any]) -> bool: + return bool(roll.get('passed', bool(roll.get('correct') and roll.get('terminated')))) + + def _apply_baseline(r: Dict[str, Any], roll: Dict[str, Any]) -> None: """Attach a greedy baseline roll and reset per-chunk working state.""" r['_baseline_rolls'], r['_cands'], r['_init'] = [roll], [], [roll] - r['_failed'] = not roll['correct'] - r['_baseline_pass'] = 1.0 if roll['correct'] else 0.0 + r['_failed'] = not _roll_passed(roll) + r['_baseline_pass'] = 1.0 if _roll_passed(roll) else 0.0 r['_hard'] = True # process every problem; group variance selects (SEAM-style) @@ -914,7 +929,7 @@ def _run(item): def _baseline_class(r: Dict[str, Any]) -> str: """success | fail_loop (out of length / never terminated) | fail_wrong.""" roll = r['_init'][0] - if roll['correct']: + if _roll_passed(roll): return 'success' return 'fail_loop' if (roll['stop_reason'] == 'length' or not roll['terminated']) else 'fail_wrong' @@ -1067,7 +1082,9 @@ def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], resp = _clean_text(getattr(s, 'decoded', '') or '') parsed = _extract_skill_blocks(resp, args.diagnose_skill_format) block = parsed['skills'] if parsed else '' - cand = {'skills': block, 'diagnose': (parsed or {}).get('diagnose', ''), + cand = {'skills': block, 'pitfall': (parsed or {}).get('pitfall', ''), + 'strategy': (parsed or {}).get('strategy', ''), + 'diagnose': (parsed or {}).get('diagnose', ''), 'skill': (parsed or {}).get('skill', ''), 'response': resp, 'parseable': bool(parsed), 'view': r['_view'], 'leaked': None, 'leak_reason': '', @@ -1091,7 +1108,7 @@ def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], c['leak_reason'] = 'answer_verbatim' if leaked else '' c['leak_source'] = 'deterministic' - # with-skill greedy pass (T=0, M=1); reward = correct, absolute (group mean is baseline). + # with-skill greedy pass (T=0, M=1); reward = correct and terminated, absolute (group mean is baseline). scored_inputs = flat if scored_inputs: ws_out = _run_samples(base_sampler, @@ -1099,7 +1116,7 @@ def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], 1, args.max_tokens, base_dp, temperature=0.0) for (r, c), seqs in zip(scored_inputs, ws_out): c['rolls'] = [_parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll()] - c['with_pass'] = 1.0 if c['rolls'][0]['correct'] else 0.0 + c['with_pass'] = 1.0 if _roll_passed(c['rolls'][0]) else 0.0 c['reward'] = c['with_pass'] if args.format_in_reward: # unparseable candidates score 0 and still join the group for r in hard: @@ -1141,7 +1158,8 @@ def _full_record(r: Dict[str, Any], ci: int) -> Dict[str, Any]: 'rubric_src': r.get('_rubric_src', ''), 'neighbor_sim': r.get('_neighbor_sim'), 'baseline_rolls': [_roll(x) for x in r['_baseline_rolls']], 'candidates': [{ - 'skills': c['skills'], 'diagnose': c.get('diagnose', ''), 'skill': c.get('skill', ''), + 'skills': c['skills'], 'pitfall': c.get('pitfall', ''), 'strategy': c.get('strategy', ''), + 'diagnose': c.get('diagnose', ''), 'skill': c.get('skill', ''), 'response': c['response'], 'parseable': c['parseable'], 'skillgen_stop': c.get('skillgen_stop'), 'skillgen_tokens': c.get('skillgen_tokens'), 'leaked': c['leaked'], 'leak_reason': c['leak_reason'], 'leak_source': c['leak_source'], @@ -1262,6 +1280,7 @@ def _group_records(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> Lis 'view': r.get('_view', 'A'), 'diagnosis': r.get('_rubric_diag', ''), 'rubric_src': r.get('_rubric_src', ''), 'response': c['response'], 'skills': c['skills'], + 'pitfall': c.get('pitfall', ''), 'strategy': c.get('strategy', ''), 'diagnose': c.get('diagnose', ''), 'skill': c.get('skill', ''), 'diagnose_skill_format': args.diagnose_skill_format, 'skillgen_stop': c.get('skillgen_stop'), @@ -1283,7 +1302,7 @@ def _is_num(v: Any) -> bool: def _train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: """Training sample = the exact skill-gen prompt (rebuilt by ``_skillgen_messages`` so - train/inference match) + the generated (think + skills) response. ``key_rounds`` + train/inference match) + the generated structured guidance response. ``key_rounds`` selects the final assistant turn; Template masks the prompt and trains the whole response (the key-round prefix already excludes the prompt-provided ).""" msgs = _skillgen_messages( @@ -1355,7 +1374,7 @@ def run_greedy_eval(base_sampler, skill_sampler, eval_records: List[Dict[str, An """SEAM ``val-core/math/acc/mean@1`` on the fixed holdout: ONE greedy skill (T=0) per problem into ONE greedy base solve (T=0). Eval ALWAYS uses view B (query-only, the deployment form: no rubric, since rubric needs an online teacher unavailable at deploy); - no leak filter (acc scores correctness alone). Baseline reuses the disk cache.""" + no leak filter. Main acc requires both correctness and normal termination. Baseline reuses the disk cache.""" baseline_rollout(base_sampler, eval_records, base_dp, args, base_cache) for r in eval_records: r['_view'], r['_rubric_diag'] = 'B', '' @@ -1369,37 +1388,40 @@ def run_greedy_eval(base_sampler, skill_sampler, eval_records: List[Dict[str, An sresp = _clean_text(getattr(seqs[0], 'decoded', '') or '') parsed = _extract_skill_blocks(sresp, args.diagnose_skill_format) skills.append((parsed['skills'] if parsed else '', - (parsed or {}).get('diagnose', ''), - (parsed or {}).get('skill', ''), + (parsed or {}).get('pitfall', ''), + (parsed or {}).get('strategy', ''), sresp)) ws_out = _run_samples(base_sampler, [build_skill_solve_prompt(r['problem'], sk) for r, (sk, _, _, _) in zip(eval_records, skills)], 1, args.max_tokens, base_dp, temperature=0.0) recs = [] - for r, (sk, diag_text, skill_text, sresp), seqs in zip(eval_records, skills, ws_out): + for r, (sk, pitfall_text, strategy_text, sresp), seqs in zip(eval_records, skills, ws_out): roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() recs.append({ 'record_type': 'eval_problem', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, 'problem': r['problem'], 'reference_answer': r['reference_answer'], 'view': r['_view'], 'rubric_diag': r.get('_rubric_diag', ''), - 'baseline_pass': r['_baseline_pass'], 'skill': sk, 'skill_diagnose': diag_text, - 'skill_action': skill_text, 'skill_parseable': bool(sk), + 'baseline_pass': r['_baseline_pass'], 'skill': sk, 'skill_pitfall': pitfall_text, + 'skill_strategy': strategy_text, 'skill_diagnose': pitfall_text, + 'skill_action': strategy_text, 'skill_parseable': bool(sk), 'skill_response': sresp, 'withskill_pred': roll['pred'], 'withskill_correct': roll['correct'], 'withskill_terminated': roll['terminated'], + 'withskill_pass': _roll_passed(roll), 'withskill_stop_reason': roll['stop_reason'], 'withskill_text': roll['text'], }) - acc = lambda rs: sum(1 for x in rs if x['withskill_correct']) / len(rs) if rs else 0.0 + acc = lambda rs: sum(1 for x in rs if x['withskill_pass']) / len(rs) if rs else 0.0 ws = acc(recs) # all view B (deployment form) base = sum(x['baseline_pass'] for x in recs) / len(recs) if recs else 0.0 + correct = (sum(1 for x in recs if x['withskill_correct']) / len(recs)) if recs else 0.0 fmt = (sum(1 for x in recs if x['skill_parseable']) / len(recs)) if recs else 0.0 term = (sum(1 for x in recs if x['withskill_terminated']) / len(recs)) if recs else 0.0 summary = {'record_type': 'eval_summary', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, 'n': len(recs), 'view': 'B', 'acc_mean1': ws, 'baseline_acc_mean1': base, 'lift_mean1': ws - base, - 'format_mean1': fmt, 'term_mean1': term} + 'correct_mean1': correct, 'format_mean1': fmt, 'term_mean1': term} metrics = {'core/math/acc/mean@1': ws, 'core/math/baseline_acc/mean@1': base, - 'core/math/lift/mean@1': ws - base, 'core/math/format/mean@1': fmt, - 'core/math/term/mean@1': term} + 'core/math/lift/mean@1': ws - base, 'core/math/correct/mean@1': correct, + 'core/math/format/mean@1': fmt, 'core/math/term/mean@1': term} return recs, summary, metrics @@ -1571,7 +1593,7 @@ def _build_args() -> argparse.Namespace: p.add_argument('--skill-gen-top-p', type=float, default=1.0) p.add_argument('--skill-gen-top-k', type=int, default=-1) p.add_argument('--diagnose-skill-format', action=argparse.BooleanOptionalAction, default=True, - help='Require both and blocks for skill-gen format; ' + help='Require and blocks for skill-gen format; ' '--no-diagnose-skill-format falls back to legacy .') p.add_argument('--max-model-len', type=int, default=16384) p.add_argument('--max-tokens', type=int, default=8192) @@ -1668,7 +1690,7 @@ def main() -> None: 'diagnose_skill_format': args.diagnose_skill_format, 'skill_retries': args.skill_retries, 'balance': args.balance, 'balance_success_frac': args.balance_success_frac, - 'skill_gen_temp': args.skill_gen_temperature, 'reward': 'greedy_binary(correct)', + 'skill_gen_temp': args.skill_gen_temperature, 'reward': 'greedy_binary(correct_and_terminated)', 'advantage': 'group_relative', 'format_in_reward': args.format_in_reward, 'cache': use_cache, 'rubric_check': 'fixed_math_5crit(viewA)' if checker else 'disabled', 'xproblem_rubric': args.xproblem_rubric, From c07cdaacf71773604a8c9dffe2b63de06744c30a Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Mon, 20 Jul 2026 19:57:59 +0800 Subject: [PATCH 22/60] sft25 --- .../build_reflexion_coldstart_sft.py | 519 ++++++++++++++++++ .../exp/embedding/train_reflexion_skill.py | 427 ++++++++------ .../exp/embedding/train_reflexion_skill.sh | 13 +- 3 files changed, 800 insertions(+), 159 deletions(-) create mode 100644 cookbook/exp/embedding/build_reflexion_coldstart_sft.py diff --git a/cookbook/exp/embedding/build_reflexion_coldstart_sft.py b/cookbook/exp/embedding/build_reflexion_coldstart_sft.py new file mode 100644 index 000000000..548dfad58 --- /dev/null +++ b/cookbook/exp/embedding/build_reflexion_coldstart_sft.py @@ -0,0 +1,519 @@ +"""Build a cold-start SFT corpus for reflexion skill generation on AOPS. + +Pipeline: + AOPS problems -> frozen base greedy attempt -> strategy-level rubric API diagnosis + -> answer-free API skill target -> query-only SFT examples. + +This is intentionally offline: no GRPO, no actor training, and no skill-model rollout. +The API is treated as an external teacher, so both diagnosis and generated skill targets +are filtered if they reveal the target final answer. + +Example: + LLM_BACKUP_API_KEY=... LLM_BACKUP_BASE_URL=... LLM_BACKUP_MODEL=... \ + python cookbook/exp/embedding/build_reflexion_coldstart_sft.py \ + --dataset aops --n 10000 --output-dir ./output/reflexion_coldstart_sft --overwrite +""" +import argparse +import json +import os +import sys +import time +import urllib.error +import urllib.request +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Optional, Tuple + +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../..')) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.sampler import vLLMSampler + +from cookbook.exp.embedding.train_reflexion_skill import ( + MODEL_ID, + GPU_MEM, + SamplingParams, + DiskCache, + _MATH_RUBRIC, + _RUBRIC_VERSION, + _answer_leaked, + _clean_text, + _empty_roll, + _format_diagnosis, + _numeric_value, + _parse_seq, + _run_samples, + _skillgen_messages, + _load_excluded_records, + build_direct_prompt, + build_skill_solve_prompt, + build_rubric_checker, + extract_boxed, + load_problems, +) + +logger = get_logger() + +COLDSTART_SYSTEM = """\ +You are writing cold-start training targets for a math skill generator. You are given a +competition problem and an answer-free process diagnosis of a previous attempt. + +Write concise, reusable guidance that a query-only solver could use before solving this +problem or similar problems. Focus on route choice, structural observations, constraints, +validity checks, and length-control habits. + +Output exactly one XML-style block: + +Your reusable guidance here. + + +Rules: +- Do not mention the diagnosis, rubric, previous attempt, or API. +- Do not reveal the final answer, a corrected value/expression, an option label, or a + step-by-step solution. +- It is okay to name methods, checks, pitfalls, and local strategy directions. +- Keep it short and useful: 3-6 compact sentences or bullets. +""" + +COLDSTART_USER = """\ +Problem: +{problem} + +Answer-free process diagnosis: +{diagnosis} + +Now write the reusable skill guidance. +""" + +_SPECIAL_TOKEN_NOTE = 'process diagnosis leaked target answer' + + +def _api_config() -> Tuple[str, str, str]: + api_key = os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY') + base_url = os.environ.get('LLM_BACKUP_BASE_URL') or os.environ.get('OPENAI_BASE_URL') or 'https://api.openai.com/v1' + model = os.environ.get('LLM_BACKUP_MODEL') or os.environ.get('OPENAI_MODEL') or 'gpt-4o-mini' + if not api_key: + raise RuntimeError('Set LLM_BACKUP_API_KEY or OPENAI_API_KEY for cold-start API generation.') + return api_key, base_url.rstrip('/'), model + + +def _chat_complete(messages: List[Dict[str, str]], max_tokens: int, temperature: float, + retries: int = 3, timeout: int = 120) -> str: + api_key, base_url, model = _api_config() + url = f'{base_url}/chat/completions' + payload = { + 'model': model, + 'messages': messages, + 'temperature': temperature, + 'max_tokens': max_tokens, + } + data = json.dumps(payload).encode('utf-8') + headers = {'Content-Type': 'application/json', 'Authorization': f'Bearer {api_key}'} + last_err = None + for attempt in range(max(1, retries)): + req = urllib.request.Request(url, data=data, headers=headers, method='POST') + try: + with urllib.request.urlopen(req, timeout=timeout) as resp: + obj = json.loads(resp.read().decode('utf-8')) + return obj['choices'][0]['message']['content'] + except (urllib.error.HTTPError, urllib.error.URLError, TimeoutError, KeyError, json.JSONDecodeError) as exc: + last_err = exc + if attempt + 1 < max(1, retries): + time.sleep(min(8.0, 1.0 * (2 ** attempt))) + continue + raise RuntimeError(f'chat completion failed after {retries} attempts: {last_err}') + + +def _extract_skill_block(text: str) -> Optional[str]: + low = (text or '').lower() + end_think = low.rfind('') + answer = text[end_think + len(''):] if end_think >= 0 else (text or '') + low = answer.lower() + s = low.rfind('') + if s < 0: + return None + inner = s + len('') + e = low.find('', inner) + if e < 0: + return None + block = answer[inner:e].strip() + return block or None + + +def _skill_response(block: str) -> str: + return f'\n{block.strip()}\n' + + +def _baseline_rollout(base_sampler, problems: List[Dict[str, Any]], base_dp: int, + args: argparse.Namespace, cache: DiskCache) -> int: + todo = [r for r in problems if DiskCache.key_for(r['problem']) not in cache] + if todo: + outs = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], + 1, args.max_tokens, base_dp, temperature=0.0) + for r, seqs in zip(todo, outs): + roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + cache.put(DiskCache.key_for(r['problem']), roll) + for r in problems: + roll = cache.get(DiskCache.key_for(r['problem'])) + r['_init'] = [roll] + r['_baseline_pass'] = 1.0 if roll.get('correct') else 0.0 + r['_failed'] = not roll.get('correct') + return len(todo) + + +def _diagnose_one(checker, r: Dict[str, Any], args: argparse.Namespace) -> str: + init = r['_init'][0] + seg_text = init.get('text', '') + if init.get('stop_reason') == 'length' or not init.get('terminated'): + seg_text += ('\n\n[Process note: this attempt was cut off at the token budget ' + 'and never produced a final \\boxed{} answer.]') + seg = {'messages': [{'role': 'user', 'content': r['problem']}, + {'role': 'assistant', 'content': seg_text}]} + attempts = max(1, args.rubric_retries + 1) + for attempt in range(attempts): + try: + return _format_diagnosis(checker.diagnose(seg, query=r['problem'])) + except Exception as exc: + if attempt + 1 < attempts: + logger.warning(f'[rubric] diagnose error: {exc}; retry {attempt + 1}/{args.rubric_retries}') + time.sleep(min(4.0, 0.5 * (2 ** attempt))) + else: + logger.warning(f'[rubric] diagnose failed: {exc}') + return '' + + +def _diagnose_batch(checker, rows: List[Dict[str, Any]], args: argparse.Namespace, + cache: DiskCache) -> int: + pending = [] + for r in rows: + init = r['_init'][0] + term = 'L' if (init.get('stop_reason') == 'length' or not init.get('terminated')) else 'T' + key = DiskCache.key_for(r['problem'], init.get('text', ''), _RUBRIC_VERSION, term) + if key in cache: + r['_rubric_diag'] = cache.get(key) + else: + pending.append((r, key)) + if not pending: + return 0 + + def run(item): + r, key = item + diag = _diagnose_one(checker, r, args) + return r, key, diag + + workers = max(1, min(args.rubric_workers, len(pending))) + fresh = 0 + with ThreadPoolExecutor(max_workers=workers) as ex: + for r, key, diag in ex.map(run, pending): + r['_rubric_diag'] = diag or '' + if diag: + cache.put(key, diag) + fresh += 1 + return fresh + + +def _target_key(problem: str, diagnosis: str, sample_idx: int) -> str: + return DiskCache.key_for('coldstart_skill_v2', str(sample_idx), problem, diagnosis) + + +def _generate_skill_targets(r: Dict[str, Any], args: argparse.Namespace, + cache: DiskCache) -> List[Dict[str, Any]]: + out = [] + messages = [ + {'role': 'system', 'content': COLDSTART_SYSTEM}, + {'role': 'user', 'content': COLDSTART_USER.format(problem=r['problem'], diagnosis=r.get('_rubric_diag', ''))}, + ] + for sample_idx in range(max(1, int(args.api_samples))): + key = _target_key(r['problem'], r.get('_rubric_diag', ''), sample_idx) + if key in cache: + resp = cache.get(key) + else: + resp = _chat_complete(messages, max_tokens=args.api_max_tokens, + temperature=args.api_temperature, retries=args.api_retries, + timeout=args.api_timeout) + cache.put(key, resp) + block = _extract_skill_block(resp) or '' + leaked = _answer_leaked(resp + '\n' + block, r['reference_answer']) + out.append({'sample_idx': sample_idx, 'raw_response': resp, + 'skills': block, 'skill_leak': leaked}) + return out + + +def _sft_messages(problem: str, response: str) -> List[Dict[str, str]]: + msgs = _skillgen_messages(problem, 'B', '') + return msgs + [{'role': 'assistant', 'content': response}] + + +def _init_base_sampler(args: argparse.Namespace): + twinkle.initialize(mode='ray', nproc_per_node=args.base_gpus, lazy_collect=False, + groups=[DeviceGroup(name='base_sampler', ranks=list(range(args.base_gpus)), device_type='GPU')]) + sampler = vLLMSampler( + model_id=MODEL_ID, + engine_args={'gpu_memory_utilization': GPU_MEM, + 'max_model_len': args.max_model_len, + 'tensor_parallel_size': 1}, + device_mesh=DeviceMesh.from_sizes(world_size=args.base_gpus, dp_size=args.base_gpus), + remote_group='base_sampler') + sampler.set_template('Template', model_id=MODEL_ID, enable_thinking=True, + max_length=args.max_model_len) + return sampler, args.base_gpus + + +def _select_records(args: argparse.Namespace) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + load_n = 0 if args.numeric_only or args.eval_size > 0 else max(args.n, args.target_size + args.eval_size) + records = load_problems(args.dataset, load_n, args.seed) + raw_n = len(records) + if args.numeric_only: + records = [{**r, 'reference_answer': v} + for r, v in ((r, _numeric_value(r.get('reference_answer'))) for r in records) + if v is not None] + import numpy as np + np.random.RandomState(args.seed).shuffle(records) + exclude_ids, exclude_problems = _load_excluded_records(getattr(args, 'exclude_data_ids', '')) + excluded = 0 + if exclude_ids or exclude_problems: + before = len(records) + records = [r for r in records + if str(r.get('data_id', '')) not in exclude_ids + and str(r.get('problem', '')).strip() not in exclude_problems] + excluded = before - len(records) + eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 + pool = [dict(r) for r in records[eval_n:]] + pool_offset = max(0, int(getattr(args, 'pool_offset', 0) or 0)) + if pool_offset: + if pool_offset >= len(pool): + raise ValueError(f'--pool-offset ({pool_offset}) leaves no cold-start records from pool size {len(pool)}') + pool = pool[pool_offset:] + n = min(args.n, len(pool)) if args.n > 0 else min(len(pool), max(args.target_size * 2, args.target_size + 512)) + stats = {'raw_loaded': raw_n, 'numeric_dropped': raw_n - len(records) - excluded, + 'excluded_records': excluded, 'eval_size': eval_n, + 'pool_offset': pool_offset, 'pool_selected': n} + return pool[:n], stats + + +def _build_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--dataset', choices=('aops', 'math'), default='aops') + p.add_argument('--target-size', type=int, default=10000, help='Number of accepted SFT examples to write.') + p.add_argument('--n', type=int, default=0, help='Raw train-pool size after eval split; 0 auto-selects.') + p.add_argument('--pool-offset', type=int, default=0, + help='Skip this many shuffled non-eval records before building the cold-start pool.') + p.add_argument('--exclude-data-ids', default='', + help='Comma-separated jsonl files whose data_id/problem keys are excluded, ' + 'useful for building non-overlapping shards.') + p.add_argument('--eval-size', type=int, default=128) + p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--seed', type=int, default=42) + p.add_argument('--output-dir', default='./output/reflexion_coldstart_sft') + p.add_argument('--cache-dir', default='') + p.add_argument('--overwrite', action='store_true') + p.add_argument('--no-cache', action='store_true') + p.add_argument('--chunk-size', type=int, default=64) + p.add_argument('--base-gpus', type=int, default=int(os.environ.get('BASE_GPUS', 4))) + p.add_argument('--max-model-len', type=int, default=16384) + p.add_argument('--max-tokens', type=int, default=8192) + p.add_argument('--rubric-workers', type=int, default=16) + p.add_argument('--rubric-retries', type=int, default=2) + p.add_argument('--api-workers', type=int, default=16) + p.add_argument('--api-samples', type=int, default=4, + help='API skill targets sampled per problem before executor verification.') + p.add_argument('--verify-targets', action=argparse.BooleanOptionalAction, default=True, + help='Run frozen base executor with each API skill target and keep a successful one.') + p.add_argument('--keep-unverified-targets', action='store_true', + help='If all executor checks fail, keep the first clean target anyway. Default skips it.') + p.add_argument('--api-retries', type=int, default=3) + p.add_argument('--api-timeout', type=int, default=120) + p.add_argument('--api-max-tokens', type=int, default=768) + p.add_argument('--api-temperature', type=float, default=0.2) + p.add_argument('--require-fail', action=argparse.BooleanOptionalAction, default=True, + help='Only keep API diagnoses containing [FAIL]. Use --no-require-fail to keep OK diagnoses too.') + return p.parse_args() + + +def _write(f, row: Dict[str, Any]) -> None: + f.write(json.dumps(row, ensure_ascii=False) + '\n') + + +def main() -> None: + args = _build_args() + if args.target_size <= 0: + raise ValueError('--target-size must be positive') + records, data_stats = _select_records(args) + if not records: + raise ValueError('no records selected') + + os.makedirs(args.output_dir, exist_ok=True) + sft_path = os.path.join(args.output_dir, 'coldstart_sft.jsonl') + rec_path = os.path.join(args.output_dir, 'coldstart_records.jsonl') + for path in (sft_path, rec_path): + if os.path.exists(path) and not args.overwrite: + raise FileExistsError(f'{path} exists; pass --overwrite') + + checker = build_rubric_checker() + if checker is None: + raise RuntimeError('No rubric checker available; set LLM_BACKUP_API_KEY/BASE_URL or OPENAI_API_KEY.') + _api_config() + base_sampler, base_dp = _init_base_sampler(args) + + cache_dir = args.cache_dir or os.path.join(args.output_dir, 'cache') + os.makedirs(cache_dir, exist_ok=True) + use_cache = not args.no_cache + base_cache = DiskCache(os.path.join(cache_dir, 'baseline.jsonl'), use_cache) + rubric_cache = DiskCache(os.path.join(cache_dir, 'rubric.jsonl'), use_cache) + skill_cache = DiskCache(os.path.join(cache_dir, 'api_skill.jsonl'), use_cache) + + cfg = { + 'record_type': 'config', 'mode': 'coldstart_sft_build', 'dataset': args.dataset, + 'target_size': args.target_size, 'selected_records': len(records), 'seed': args.seed, + 'pool_offset': args.pool_offset, 'exclude_data_ids': args.exclude_data_ids, + 'numeric_only': args.numeric_only, **data_stats, + 'rubric_version': _RUBRIC_VERSION, 'rubric_check': f'fixed_math_{len(_MATH_RUBRIC)}crit', + 'api_model': os.environ.get('LLM_BACKUP_MODEL') or os.environ.get('OPENAI_MODEL') or 'gpt-4o-mini', + 'api_samples': args.api_samples, 'verify_targets': args.verify_targets, + 'keep_unverified_targets': args.keep_unverified_targets, + 'require_fail': args.require_fail, 'started': int(time.time()), + } + + accepted = 0 + skipped_no_diag = skipped_no_fail = skipped_api_leak = 0 + skipped_no_skill = skipped_skill_leak = skipped_executor_fail = 0 + processed = 0 + with open(sft_path, 'w', encoding='utf-8') as sft_f, open(rec_path, 'w', encoding='utf-8') as rec_f: + _write(rec_f, cfg) + for start in range(0, len(records), args.chunk_size): + if accepted >= args.target_size: + break + chunk = [dict(r) for r in records[start:start + args.chunk_size]] + _baseline_rollout(base_sampler, chunk, base_dp, args, base_cache) + _diagnose_batch(checker, chunk, args, rubric_cache) + + def gen_one(r: Dict[str, Any]): + return r, _generate_skill_targets(r, args, skill_cache) + + candidates = [] + for r in chunk: + processed += 1 + diag = r.get('_rubric_diag', '') or '' + if not diag: + skipped_no_diag += 1 + continue + if args.require_fail and '[FAIL]' not in diag: + skipped_no_fail += 1 + continue + if _answer_leaked(diag, r['reference_answer']): + skipped_api_leak += 1 + continue + candidates.append(r) + + generated = [] + workers = max(1, min(args.api_workers, len(candidates))) + if candidates: + with ThreadPoolExecutor(max_workers=workers) as ex: + for r, targets in ex.map(gen_one, candidates): + for target in targets: + skills = target.get('skills', '') + if not skills: + skipped_no_skill += 1 + continue + if target.get('skill_leak'): + skipped_skill_leak += 1 + continue + target['r'] = r + target['response'] = _skill_response(skills) + generated.append(target) + + selected = [] + selected_keys = set() + if generated and args.verify_targets: + verify_prompts = [build_skill_solve_prompt(g['r']['problem'], g['skills']) for g in generated] + verify_outs = _run_samples(base_sampler, verify_prompts, 1, args.max_tokens, + base_dp, temperature=0.0) + attempted_keys = set() + for g, seqs in zip(generated, verify_outs): + r = g['r'] + key = r.get('data_id') or r['problem'] + attempted_keys.add(key) + roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + g['target_roll'] = roll + if key not in selected_keys and roll.get('correct') and roll.get('terminated'): + g['executor_verified'] = True + selected.append(g) + selected_keys.add(key) + if args.keep_unverified_targets: + for g in generated: + r = g['r'] + key = r.get('data_id') or r['problem'] + if key not in selected_keys: + g['executor_verified'] = False + g.setdefault('target_roll', {}) + selected.append(g) + selected_keys.add(key) + skipped_executor_fail += len(attempted_keys - selected_keys) + elif generated: + for g in generated: + r = g['r'] + key = r.get('data_id') or r['problem'] + if key not in selected_keys: + g['executor_verified'] = False + selected.append(g) + selected_keys.add(key) + + for g in selected: + if accepted >= args.target_size: + break + r = g['r'] + response = g['response'] + messages = _sft_messages(r['problem'], response) + sft_row = { + 'messages': messages, + 'user_data': {'key_rounds': [len(messages) - 1]}, + 'data_id': r.get('data_id'), + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'skills': g['skills'], 'response': response, + 'view': 'B', 'sft': True, 'source': 'api_coldstart', + 'api_sample_idx': g.get('sample_idx'), + 'executor_verified': g.get('executor_verified', False), + 'baseline_correct': r['_init'][0].get('correct'), + 'baseline_terminated': r['_init'][0].get('terminated'), + 'baseline_stop_reason': r['_init'][0].get('stop_reason'), + 'target_correct': (g.get('target_roll') or {}).get('correct'), + 'target_terminated': (g.get('target_roll') or {}).get('terminated'), + 'target_stop_reason': (g.get('target_roll') or {}).get('stop_reason'), + 'diagnosis': r.get('_rubric_diag', ''), + } + audit = { + 'record_type': 'coldstart_problem', 'accepted': True, + 'data_id': r.get('data_id'), + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'baseline': r['_init'][0], 'diagnosis': r.get('_rubric_diag', ''), + 'raw_skill_response': g.get('raw_response'), 'skills': g['skills'], + 'api_sample_idx': g.get('sample_idx'), + 'executor_verified': g.get('executor_verified', False), + 'target_roll': g.get('target_roll'), + } + _write(sft_f, sft_row) + _write(rec_f, audit) + accepted += 1 + sys.stderr.write( + f'[coldstart] processed={processed} accepted={accepted}/{args.target_size} ' + f'skip(no_diag={skipped_no_diag}, no_fail={skipped_no_fail}, api_leak={skipped_api_leak}, ' + f'no_skill={skipped_no_skill}, skill_leak={skipped_skill_leak}, ' + f'executor_fail={skipped_executor_fail})\n') + sft_f.flush(); rec_f.flush() + + summary = { + 'record_type': 'summary', 'processed': processed, 'accepted': accepted, + 'skipped_no_diag': skipped_no_diag, 'skipped_no_fail': skipped_no_fail, + 'skipped_api_leak': skipped_api_leak, 'skipped_no_skill': skipped_no_skill, + 'skipped_skill_leak': skipped_skill_leak, + 'skipped_executor_fail': skipped_executor_fail, 'finished': int(time.time()), + } + with open(rec_path, 'a', encoding='utf-8') as rec_f: + _write(rec_f, summary) + sys.stderr.write(f'[coldstart] wrote {accepted} SFT rows to {sft_path}\n') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/embedding/train_reflexion_skill.py b/cookbook/exp/embedding/train_reflexion_skill.py index 39af13e21..5d323ee08 100644 --- a/cookbook/exp/embedding/train_reflexion_skill.py +++ b/cookbook/exp/embedding/train_reflexion_skill.py @@ -36,7 +36,7 @@ import threading import time from concurrent.futures import ThreadPoolExecutor -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Set, Tuple import numpy as np @@ -271,24 +271,20 @@ def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: # -- skill-gen prompts (view A: problem + rubric source; view B: query only) -- +# Kept deliberately short: this is the RL policy's system prompt, so over-specifying +# the output hurts convergence. The concrete output format is appended separately by +# ``_SKILL_OUTPUT`` so the task text stays format-agnostic. SKILL_GEN_SYSTEM = ( - 'You are a math guidance writer. You are given a target problem and an automated ' - 'process-check from a related problem. Use it as context to write reusable guidance ' - 'for this and similar problems. In , write likely mistakes or checks only. ' - 'Leave empty if unclear. In , write reusable solving advice.\n') + 'You are a math guidance writer. A process-check on a related problem hints at ' + 'likely mistakes. Write short reusable guidance for this and similar problems, ' + 'and note what to watch out for.\n') SKILL_GEN_SYSTEM_Q = ( - 'You are a math guidance writer. Given the problem below, write reusable guidance ' - 'for this and similar problems. In , write likely mistakes or checks only. ' - 'Leave empty if unclear. In , write reusable solving advice.\n') + 'You are a math guidance writer. Write short reusable guidance for this and ' + 'similar problems.\n') -_SKILL_OUTPUT_DUAL = ( - 'Output:\n' - 'Your likely mistakes or checks here\n' - 'Your reusable solving advice here...') - -_SKILL_OUTPUT_LEGACY = ( - 'Output a non-empty:\n\nYour remind and skill here...\n\nblock.') +_SKILL_OUTPUT = ( + 'Output only:\n\nYour reusable solving guidance here.\n') SKILL_GEN_USER_Q = ( 'Problem:\n{problem}\n\n') @@ -300,22 +296,25 @@ def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: '{diagnosis}\n\n') -def _skill_output_instruction(diagnose_skill_format: bool) -> str: - return _SKILL_OUTPUT_DUAL if diagnose_skill_format else _SKILL_OUTPUT_LEGACY +def _rubric_has_fail(diagnosis: str) -> bool: + """View A uses the open-book rubric prompt (and, under --viewa-sft, SFT distillation) + IFF its rubric localised a failure. No ``[FAIL]`` => nothing to inject: generation + degrades to query-only and the problem is trained by GRPO exactly like view B. Single + source of truth so ``_skillgen_messages`` and ``_group_records`` never diverge.""" + return '[FAIL]' in (diagnosis or '') def _skillgen_messages(problem: str, view: str, diagnosis: str, - rubric_problem: str = '', diagnose_skill_format: bool = True) -> List[Dict[str, Any]]: + rubric_problem: str = '') -> List[Dict[str, Any]]: """Single source of truth for the skill-gen prompt (used at BOTH generation and training so they never diverge). View A with a localisable failure uses the target problem plus the rubric source problem and findings; view B -- or a view-A problem whose rubric flagged nothing (no ``[FAIL]``) -- degrades to the query-only prompt.""" - out = _skill_output_instruction(diagnose_skill_format) - if view == 'B' or '[FAIL]' not in (diagnosis or ''): - return [{'role': 'system', 'content': SKILL_GEN_SYSTEM_Q + out}, + if view == 'B' or not _rubric_has_fail(diagnosis): + return [{'role': 'system', 'content': SKILL_GEN_SYSTEM_Q + _SKILL_OUTPUT}, {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=problem)}] rubric_problem = rubric_problem or problem - return [{'role': 'system', 'content': SKILL_GEN_SYSTEM + out}, + return [{'role': 'system', 'content': SKILL_GEN_SYSTEM + _SKILL_OUTPUT}, {'role': 'user', 'content': SKILL_GEN_USER_RUBRIC.format( problem=problem, rubric_problem=rubric_problem, diagnosis=diagnosis)}] @@ -327,8 +326,7 @@ def _assign_view(problem: str, args: argparse.Namespace) -> str: def _view_prompt(r: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: return {'messages': _skillgen_messages( - r['problem'], r['_view'], r.get('_rubric_diag', ''), r.get('_rubric_src', ''), - args.diagnose_skill_format)} + r['problem'], r['_view'], r.get('_rubric_diag', ''), r.get('_rubric_src', ''))} _SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') @@ -353,36 +351,14 @@ def _extract_tag_block(answer: str, tag: str, allow_empty: bool = False) -> Opti return block if (block or allow_empty) else None -def _join_pitfall_strategy(pitfall: str, strategy: str) -> str: - pitfall = (pitfall or '').strip() - strategy = (strategy or '').strip() - return strategy if not pitfall else f'{pitfall}\n\n{strategy}' - - -def _extract_skill_blocks(text: str, diagnose_skill_format: bool = True) -> Optional[Dict[str, str]]: - """Parse skill-generation output. - - With the pitfall-strategy format enabled, both ```` and ```` tags - must be present, while ```` may be empty; their inner texts are concatenated - for executor injection. With the legacy format, a non-empty ```` block is - enough. If a ```` marker is present, parse only the text after the last marker; - otherwise parse the full response. - """ +def _extract_skill(text: str) -> Optional[str]: + """Parse skill-generation output: return the inner text of a non-empty ```` + block, or None. If a ```` marker is present, parse only the text after the + last one; otherwise parse the full response.""" low = text.lower() end_think = low.rfind('') answer = text[end_think + len(''):] if end_think >= 0 else text - if diagnose_skill_format: - pitfall = _extract_tag_block(answer, 'pitfall', allow_empty=True) - strategy = _extract_tag_block(answer, 'strategy') - if pitfall is None or not strategy: - return None - joined = _join_pitfall_strategy(pitfall, strategy) - return {'pitfall': pitfall, 'strategy': strategy, - 'diagnose': pitfall, 'skill': strategy, 'skills': joined} - block = _extract_tag_block(answer, 'skills') - if not block: - return None - return {'pitfall': '', 'strategy': block, 'diagnose': '', 'skill': block, 'skills': block} + return _extract_tag_block(answer, 'skills') def _parse_seq(seq, gold: str) -> Dict[str, Any]: @@ -441,9 +417,10 @@ def load_problems(dataset: str, n: int, seed: int, num_proc: int = 0) -> List[Di ds.map(lambda rows: _boxed_batch(rows, dataset), num_proc=nproc) ds.filter(lambda row: row['_keep'], num_proc=nproc) has_level = 'level' in ds.dataset.column_names - out = [{'problem': row['problem'], 'reference_answer': row['reference_answer'], + out = [{'data_id': f'{dataset}:{i}', 'problem': row['problem'], + 'reference_answer': row['reference_answer'], **({'level': row['level']} if has_level and row.get('level') else {})} - for row in ds.dataset] + for i, row in enumerate(ds.dataset)] logger.info(f'[data] {dataset}: {len(out)} boxed problems ({nproc} procs)') rng = np.random.RandomState(seed) rng.shuffle(out) @@ -599,12 +576,10 @@ def _numeric_value(raw: Any) -> Optional[str]: def _answer_leaked(skill: str, reference: str) -> bool: - """Deterministic answer-leak check (replaces the LLM LeakVerifier). Flags ONLY the one - real way a skill can game the deterministic reward: writing the final answer verbatim. - On our numeric-only data the answer is a single number, so we require it as a - standalone token (digit boundaries) -- that avoids matching an intermediate value like - '3' inside '36'. Everything else (methods, plans, pitfalls) is left to the reward, the - way SEAM/POPE handle it (no content audit). Non-numeric answers -> not flagged.""" + """Audit whether a generated skill contains the final answer verbatim. This is NOT + a training filter: if the skill model derives an answer from the problem, that is a + legitimate answer-bearing skill under this experiment. The real leakage boundary is the + external API/rubric diagnosis, which is constrained by prompt to stay answer-free.""" if not skill: return False for cand in {_numeric_value(reference), (str(reference).strip() or None)}: @@ -613,6 +588,33 @@ def _answer_leaked(skill: str, reference: str) -> bool: return False +def _load_excluded_records(paths_arg: str) -> Tuple[Set[str], Set[str]]: + """Read jsonl files and collect data_id/problem keys that must be excluded. + + The cold-start SFT builder writes stable ``data_id`` values. ``problem`` is kept as a + backward-compatible fallback for older jsonl files produced before data_id existed.""" + ids: Set[str] = set() + problems: Set[str] = set() + for raw_path in (paths_arg or '').split(','): + path = raw_path.strip() + if not path or not os.path.exists(path): + continue + with open(path, encoding='utf-8') as f: + for line in f: + if not line.strip(): + continue + row = json.loads(line) + if row.get('record_type') in {'config', 'summary'}: + continue + data_id = str(row.get('data_id') or '').strip() + problem = str(row.get('problem') or '').strip() + if data_id: + ids.add(data_id) + elif problem: + problems.add(problem) + return ids, problems + + def _load_records(args: argparse.Namespace ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], Dict[str, Tuple[str, float]], Dict[str, str], Dict[str, int]]: @@ -630,9 +632,22 @@ def _load_records(args: argparse.Namespace if v is not None] dropped = raw_n - len(records) np.random.RandomState(args.seed).shuffle(records) + exclude_ids, exclude_problems = _load_excluded_records(getattr(args, 'exclude_data_ids', '')) + excluded = 0 + if exclude_ids or exclude_problems: + before = len(records) + records = [r for r in records + if str(r.get('data_id', '')) not in exclude_ids + and str(r.get('problem', '')).strip() not in exclude_problems] + excluded = before - len(records) eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 eval_records = [dict(r) for r in records[:eval_n]] pool = records[eval_n:] + pool_offset = max(0, int(getattr(args, 'pool_offset', 0) or 0)) + if pool_offset: + if pool_offset >= len(pool): + raise ValueError(f'--pool-offset ({pool_offset}) leaves no train records from pool size {len(pool)}') + pool = pool[pool_offset:] train_n = args.n if args.n > 0 else len(pool) # Cross-problem rubric: pick the train_n problems with the strongest qualifying neighbour # (dense same-type pairs) and build {Q -> (P, sim)} in one index pass. Otherwise keep the @@ -646,6 +661,7 @@ def _load_records(args: argparse.Namespace if {r['problem'] for r in train_records} & {r['problem'] for r in eval_records}: raise ValueError('eval/train overlap detected') stats = {'raw_loaded': raw_n, 'numeric_dropped': dropped, + 'excluded_records': excluded, 'pool_offset': pool_offset, 'train_records': len(train_records), 'eval_records': len(eval_records)} return train_records, eval_records, neighbor_map, pool_answers, stats @@ -757,15 +773,11 @@ def _empty_roll() -> Dict[str, Any]: 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} -def _roll_passed(roll: Dict[str, Any]) -> bool: - return bool(roll.get('passed', bool(roll.get('correct') and roll.get('terminated')))) - - def _apply_baseline(r: Dict[str, Any], roll: Dict[str, Any]) -> None: """Attach a greedy baseline roll and reset per-chunk working state.""" r['_baseline_rolls'], r['_cands'], r['_init'] = [roll], [], [roll] - r['_failed'] = not _roll_passed(roll) - r['_baseline_pass'] = 1.0 if _roll_passed(roll) else 0.0 + r['_failed'] = not roll['correct'] + r['_baseline_pass'] = 1.0 if roll['correct'] else 0.0 r['_hard'] = True # process every problem; group variance selects (SEAM-style) @@ -787,39 +799,44 @@ def baseline_rollout(base_sampler, problems: List[Dict[str, Any]], base_dp: int, # -- rubric process-check (view A): teacher diagnoses the base's attempt -- _RFT_DIAG_SYSTEM = """\ -You are a process error checker for a math solution attempt. You are given a math -problem, a rubric, and one attempted solution segment. Decide PASS or FAIL for each -criterion and explain only the process error type. +You are a strategy-level process checker for a math solution attempt. You are given a +math problem, a rubric, and one attempted solution segment. Decide PASS or FAIL for each +criterion, and write the diagnosis so it can become useful reusable guidance for solving +similar problems without seeing this segment. Output STRICT JSON (no prose outside it) with this shape: { "items": [ {"index": 1, "verdict": "PASS", "reason": "", "fix": ""}, - {"index": 2, "verdict": "FAIL", "reason": "", - "fix": ""} + {"index": 2, "verdict": "FAIL", "reason": "", + "fix": ""} ], "overall": "OK" | "ISSUES", - "summary": "" + "summary": "" } Rules: -- Judge every criterion independently and literally; a [Hard Rule] is FAIL unless - unambiguously satisfied. -- Judge ONLY what is observable in THIS segment. -- Content inside ... (or ) is internal reasoning, not - user-facing output; ignore it for "output only X" style criteria. +- Judge every criterion independently; a [Hard Rule] is FAIL unless clearly satisfied. +- Judge ONLY what is observable in THIS segment. Ignore hidden or + content for output-format criteria. +- The API diagnosis is an external teacher signal, so it must stay answer-free. +- Prefer diagnosis that transfers to view-B skill generation: name the route choice, + structural observation, missing check, or length-control habit that a solver should + remember before solving a similar problem. - For PASS items, leave "fix" as "". -- For FAIL items, "reason", "fix", and "summary" must describe only the flawed - step, theorem, arithmetic operation, case split, or verification habit. -- NEVER state the correct final answer, corrected final expression, option letter, - graph/choice label, or any exact value that the answer should become. -- NEVER write phrases like "the correct answer is", "which gives", "yielding", - "should be ", "Option ", or "Graph ". -- If a fix would require naming a corrected value, replace it with a method-level - instruction such as "redo that computation carefully" or "apply the theorem with - the correct quantities". -- Keep every "reason" and "fix" clear and concise — one short sentence each. +- For FAIL items, describe the process problem at strategy level: unsuitable method, + missed structure, invalid transformation, missing constraint check, redundant cases, + off-track approach, contradiction, or inefficient/unfinished reasoning. +- A fix may suggest the LOCAL correction direction, such as identify the key structure, + verify constraints, preserve equivalence, reduce redundant cases, or choose a more + direct route. Do not carry out the correction. +- Never reveal the final answer, a corrected value/expression, an option label, or a + step-by-step solution that would let another model copy the solve. +- If the segment contains a process note saying it was cut off before a final boxed + answer, mark the length-budget criterion as FAIL and suggest a method-level way to + finish faster. +- Keep every "reason" and "fix" concise: one short sentence each. - "overall" is "OK" only if NO criterion is FAIL. - Output only the JSON object.""" @@ -836,13 +853,19 @@ def baseline_rollout(base_sampler, problems: List[Dict[str, Any]], base_dp: int, Now output the diagnostic JSON object.""" _MATH_RUBRIC = [ - ('The reasoning contains no arithmetic or algebraic error', True), - ('Each step follows logically from the previous ones', True), - ('No formula or theorem is misstated or misapplied', True), - ('The approach is on track to answer the actual question asked', False), - ('No step contradicts an earlier established fact', False), + ('The attempt chooses a method suitable for the problem structure', False), + ('The attempt identifies the key constraint, invariant, or quantity before computing', False), + ('Algebraic and logical transformations preserve validity at each step', True), + ('The attempt checks required constraints, domains, boundary cases, or validity conditions', False), + ('The attempt avoids redundant casework, looping, or re-deriving known facts', False), + ('The attempt reaches a final boxed answer within the length budget', False), + ('The approach stays focused on the actual question asked', False), ] +# Bump when the rubric criteria or the diagnosis prompt change so the disk-cached +# diagnoses written under an older rubric are not silently reused. +_RUBRIC_VERSION = 'rubric_v5_viewb_strategy' + class _RftRubricVerifier(RubricVerifier): def _diagnose_trajectory(self, query: str, rubric_block: str, segment_text: str) -> dict: @@ -887,7 +910,9 @@ def diagnose_views(checker, problems: List[Dict[str, Any]], args: argparse.Names return def _key(r: Dict[str, Any]) -> str: - return DiskCache.key_for(r['problem'], r.get('_init', [{}])[0].get('text', '')) + init = r.get('_init', [{}])[0] + term = 'L' if (init.get('stop_reason') == 'length' or not init.get('terminated')) else 'T' + return DiskCache.key_for(r['problem'], init.get('text', ''), _RUBRIC_VERSION, term) pending = [] for r in targets: @@ -901,8 +926,13 @@ def _key(r: Dict[str, Any]) -> str: def _run(item): r, key = item + init = r['_init'][0] + seg_text = init['text'] + if init.get('stop_reason') == 'length' or not init.get('terminated'): + seg_text += ('\n\n[Process note: this attempt was cut off at the token budget ' + 'and never produced a final \\boxed{} answer.]') seg = {'messages': [{'role': 'user', 'content': r['problem']}, - {'role': 'assistant', 'content': r['_init'][0]['text']}]} + {'role': 'assistant', 'content': seg_text}]} attempts = max(1, args.rubric_retries + 1) for attempt in range(attempts): try: @@ -929,7 +959,7 @@ def _run(item): def _baseline_class(r: Dict[str, Any]) -> str: """success | fail_loop (out of length / never terminated) | fail_wrong.""" roll = r['_init'][0] - if _roll_passed(roll): + if roll['correct']: return 'success' return 'fail_loop' if (roll['stop_reason'] == 'length' or not roll['terminated']) else 'fail_wrong' @@ -1017,6 +1047,31 @@ def _assign_advantages(hard: List[Dict[str, Any]], args: argparse.Namespace) -> c['advantage'], c['grpo_adv'], c['kept'] = adv, adv, c['reward'] > mean_r +def _best_sft_candidate(r: Dict[str, Any], args: argparse.Namespace) -> Optional[Dict[str, Any]]: + """Pick ONE view-A candidate to distill (online context distillation). PREFER the + executor-verified PASSING skills (reward==1); if NONE passed -- common on the hard + problems that are exactly the cases worth distilling -- FALL BACK to any parseable + open-book skill regardless of the executor outcome. Answer-bearing skills produced by + the skill model itself are allowed here; only the external API/rubric diagnosis must be + answer-free. Within the chosen tier, take the one whose skill length is CLOSEST to + ``--sft-target-len`` -- an empirically high-pass-rate length (~500-600 chars in this + run) -- breaking ties by the fewest executor solve tokens. Targeting a length (rather + than the minimum) avoids a distillation feedback loop that would otherwise drive + rollouts ever shorter. None only when no parseable candidate exists at all.""" + eligible = [c for c in r['_cands'] if c.get('parseable') and c.get('skills')] + if not eligible: + return None + passing = [c for c in eligible if c.get('reward') == 1.0] + cs = passing or eligible + target = int(getattr(args, 'sft_target_len', 550) or 550) + + def _solve_tokens(c: Dict[str, Any]) -> int: + rolls = c.get('rolls') or [] + return rolls[0].get('gen_tokens', 1 << 30) if rolls else (1 << 30) + + return min(cs, key=lambda c: (abs(len(c['skills']) - target), _solve_tokens(c))) + + def apply_neighbor_rubric(base_sampler, chunk: List[Dict[str, Any]], neighbor_map: Dict[str, Tuple[str, float]], pool_answers: Dict[str, str], base_dp: int, @@ -1053,7 +1108,7 @@ def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], neighbor_map: Optional[Dict[str, Tuple[str, float]]] = None, pool_answers: Optional[Dict[str, str]] = None ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], List[Dict[str, Any]]]: - """view assign -> rubric-check (view A) -> skill-gen -> leak-filter -> with-skill + """view assign -> rubric-check (view A) -> skill-gen -> answer audit -> with-skill greedy pass -> GRPO advantages. ``chunk`` arrives already baselined by draw_chunk. With --xproblem-rubric, view A gets its NEIGHBOUR's rubric (see apply_neighbor_rubric).""" hard = chunk @@ -1066,8 +1121,10 @@ def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], diagnose_views(checker, hard, args, rubric_cache) # skill-gen (thinking OFF), per-view prompt; re-sample problems with no clean candidate. + # View-A problems that cannot yield a training record (no [FAIL] rubric, or rubric + # leaked the answer) are dropped from training entirely -- skip their generation. flat: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] - pending = list(hard) + pending = [r for r in hard if not _viewa_dropped(r, args)] for _ in range(args.skill_retries + 1): if not pending: break @@ -1080,13 +1137,8 @@ def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], got = False for s in seqs: resp = _clean_text(getattr(s, 'decoded', '') or '') - parsed = _extract_skill_blocks(resp, args.diagnose_skill_format) - block = parsed['skills'] if parsed else '' - cand = {'skills': block, 'pitfall': (parsed or {}).get('pitfall', ''), - 'strategy': (parsed or {}).get('strategy', ''), - 'diagnose': (parsed or {}).get('diagnose', ''), - 'skill': (parsed or {}).get('skill', ''), - 'response': resp, 'parseable': bool(parsed), + block = _extract_skill(resp) or '' + cand = {'skills': block, 'response': resp, 'parseable': bool(block), 'view': r['_view'], 'leaked': None, 'leak_reason': '', 'leak_source': '', 'with_pass': None, 'reward': None, 'rolls': [], 'skillgen_stop': getattr(s, 'stop_reason', None), @@ -1099,16 +1151,16 @@ def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], still.append(r) pending = still - # leak audit: deterministic verbatim-answer check only (no LLM). This is observability - # only: it records leak metrics for swanlab/jsonl, but does not block scoring, reward, - # advantage assignment, or training sample selection. + # Answer-bearing skill audit: deterministic verbatim-answer check only (no LLM). This + # is observability only; it records metrics for swanlab/jsonl, but does not block + # scoring, reward, advantage assignment, SFT candidate selection, or GRPO training. for r, c in flat: leaked = _answer_leaked(c['skills'], r['reference_answer']) c['leaked'] = leaked c['leak_reason'] = 'answer_verbatim' if leaked else '' c['leak_source'] = 'deterministic' - # with-skill greedy pass (T=0, M=1); reward = correct and terminated, absolute (group mean is baseline). + # with-skill greedy pass (T=0, M=1); reward = correct, absolute (group mean is baseline). scored_inputs = flat if scored_inputs: ws_out = _run_samples(base_sampler, @@ -1116,7 +1168,7 @@ def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], 1, args.max_tokens, base_dp, temperature=0.0) for (r, c), seqs in zip(scored_inputs, ws_out): c['rolls'] = [_parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll()] - c['with_pass'] = 1.0 if _roll_passed(c['rolls'][0]) else 0.0 + c['with_pass'] = 1.0 if c['rolls'][0]['correct'] else 0.0 c['reward'] = c['with_pass'] if args.format_in_reward: # unparseable candidates score 0 and still join the group for r in hard: @@ -1158,9 +1210,7 @@ def _full_record(r: Dict[str, Any], ci: int) -> Dict[str, Any]: 'rubric_src': r.get('_rubric_src', ''), 'neighbor_sim': r.get('_neighbor_sim'), 'baseline_rolls': [_roll(x) for x in r['_baseline_rolls']], 'candidates': [{ - 'skills': c['skills'], 'pitfall': c.get('pitfall', ''), 'strategy': c.get('strategy', ''), - 'diagnose': c.get('diagnose', ''), 'skill': c.get('skill', ''), - 'response': c['response'], 'parseable': c['parseable'], + 'skills': c['skills'], 'response': c['response'], 'parseable': c['parseable'], 'skillgen_stop': c.get('skillgen_stop'), 'skillgen_tokens': c.get('skillgen_tokens'), 'leaked': c['leaked'], 'leak_reason': c['leak_reason'], 'leak_source': c['leak_source'], 'with_pass': c['with_pass'], 'reward': c.get('reward'), 'advantage': c.get('advantage'), @@ -1218,8 +1268,11 @@ def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespac scored = [c for c in cands if c['with_pass'] is not None] clean = [c for c in cands if c['leaked'] is False] ws_rolls = [x for c in scored for x in c['rolls']] - base_acc = (sum(r['_baseline_pass'] for r in chunk) / len(chunk)) if chunk else 0.0 - ws_acc = _mean([1.0 if any(c.get('reward') for c in r['_cands']) else 0.0 for r in chunk]) + # viewa-dropped problems generate no candidates; keep acc/* on the generated subset + # so the with-skill/lift trend stays comparable across view_b_frac settings. + gen_probs = [r for r in chunk if r['_cands']] + base_acc = _mean([1.0 if r['_baseline_pass'] else 0.0 for r in gen_probs]) + ws_acc = _mean([1.0 if any(c.get('reward') for c in r['_cands']) else 0.0 for r in gen_probs]) cand_pass_parseable = _mean([c['with_pass'] for c in scored]) cand_pass_all = _mean([1.0 if c.get('reward') else 0.0 for c in all_cands]) # base failure taxonomy (you asked whether skills fail because the base loops out of length) @@ -1228,6 +1281,9 @@ def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespac skill_tokens = [c.get('skillgen_tokens') or 0 for c in cands] # skill-gen response length trunc = sum(1 for r in chunk for c in r['_cands'] for x in c['rolls'] if x['stop_reason'] == 'length') + rubric_answer_leaks = sum( + 1 for r in chunk + if r.get('_view') == 'A' and _rubric_has_fail(r.get('_rubric_diag')) and _rubric_answer_leaked(r)) return { 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), 'n_failed_first_try': sum(1 for r in chunk if r['_failed']), @@ -1237,6 +1293,8 @@ def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespac 'n_leaked': sum(1 for c in cands if c['leaked']), 'n_clean': len(clean), 'leak_rate': (sum(1 for c in cands if c['leaked']) / len(cands)) if cands else 0.0, 'n_reward_pos': sum(1 for c in scored if c['reward']), + 'n_rubric_answer_leaked': rubric_answer_leaks, + 'n_viewa_dropped': sum(1 for r in chunk if _viewa_dropped(r, args)), 'n_train_samples': sum(1 for c in all_cands if _is_trainable(c, args)), 'signal': _signal_stats(chunk), 'fail_loop_frac': (sum(1 for c in classes if c == 'fail_loop') / n_fail) if n_fail else 0.0, @@ -1266,23 +1324,66 @@ def _xproblem_stats(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> Di 'neighbor_sim_mean': _mean([r.get('_neighbor_sim', 0.0) for r in paired])}} +def _sft_record(r: Dict[str, Any], c: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: + """A view-A context-distillation sample: the OPEN-BOOK (rubric-conditioned) skill ``c`` + is the target, but the prompt is rebuilt as query-only (view B, no rubric) so the model + learns to produce it CLOSED-BOOK. Rides the same GRPOLoss with a positive constant + advantage (``--sft-weight``); single-step (old_logps=None) this reduces to + ``-w*logp + beta*KL`` -- a KL-anchored weighted cross-entropy toward the skill.""" + return { + 'data_id': r.get('data_id'), + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'view': 'B', 'diagnosis': '', 'rubric_src': '', 'orig_view': 'A', 'sft': True, + 'response': c['response'], 'skills': c['skills'], + 'skillgen_stop': c.get('skillgen_stop'), + 'advantage': float(args.sft_weight), 'grpo_adv': 0.0, 'kept': True, + 'reward': c['reward'], 'with_pass': c['with_pass']} + + +def _rubric_answer_leaked(r: Dict[str, Any]) -> bool: + """Hard safety gate for external API/rubric diagnosis. Unlike answer-bearing skills + generated by the policy itself, a rubric that contains the target final answer is an + external teacher leak and must not be distilled into view B.""" + return _answer_leaked(r.get('_rubric_diag', ''), r.get('reference_answer', '')) + + +def _viewa_dropped(r: Dict[str, Any], args: argparse.Namespace) -> bool: + """Under --viewa-sft a view-A problem trains ONLY through the SFT path (rubric with + [FAIL], answer-free). No-FAIL and rubric-leaked problems produce no training record + at all (no GRPO backflow: those prompts are query-only and would muddy the pure + view-B GRPO pool), so their skill-gen + executor scoring is skipped entirely.""" + return (bool(args.viewa_sft) and r.get('_view') == 'A' + and (not _rubric_has_fail(r.get('_rubric_diag')) + or _rubric_answer_leaked(r))) + + def _group_records(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> List[Dict[str, Any]]: - """GRPO training records: every trainable candidate with its view + rubric diagnosis - (the prompt is rebuilt from those by ``_skillgen_messages``, no trajectory stored).""" + """Training records. Under --viewa-sft, view A is SFT-only: a problem whose rubric + localised a failure (has ``[FAIL]``, answer-free) contributes ONE context-distillation + SFT sample (best parseable open-book skill -- preferring an executor-verified pass, + else any parseable candidate -- rebuilt query-only); no-FAIL / rubric-leaked view-A + problems are dropped (``_viewa_dropped``, generation already skipped). GRPO candidates + come from view B only, keeping that pool purely closed-book. Prompts are rebuilt from + the stored view/diagnosis by ``_skillgen_messages``.""" out = [] for r in chunk: if not r['_hard']: continue + if args.viewa_sft and r.get('_view') == 'A': + if _viewa_dropped(r, args): + continue + best = _best_sft_candidate(r, args) + if best is not None: + out.append(_sft_record(r, best, args)) + continue for c in r['_cands']: if _is_trainable(c, args): out.append({ + 'data_id': r.get('data_id'), 'problem': r['problem'], 'reference_answer': r['reference_answer'], 'view': r.get('_view', 'A'), 'diagnosis': r.get('_rubric_diag', ''), - 'rubric_src': r.get('_rubric_src', ''), + 'rubric_src': r.get('_rubric_src', ''), 'sft': False, 'response': c['response'], 'skills': c['skills'], - 'pitfall': c.get('pitfall', ''), 'strategy': c.get('strategy', ''), - 'diagnose': c.get('diagnose', ''), 'skill': c.get('skill', ''), - 'diagnose_skill_format': args.diagnose_skill_format, 'skillgen_stop': c.get('skillgen_stop'), 'advantage': c['advantage'], 'grpo_adv': c['grpo_adv'], 'kept': c['kept'], 'reward': c['reward'], 'with_pass': c['with_pass']}) @@ -1306,8 +1407,7 @@ def _train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: selects the final assistant turn; Template masks the prompt and trains the whole response (the key-round prefix already excludes the prompt-provided ).""" msgs = _skillgen_messages( - rec['problem'], rec.get('view', 'A'), rec.get('diagnosis', ''), rec.get('rubric_src', ''), - rec.get('diagnose_skill_format', True)) + rec['problem'], rec.get('view', 'A'), rec.get('diagnosis', ''), rec.get('rubric_src', '')) return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], 'user_data': {'key_rounds': [len(msgs)]}} @@ -1325,7 +1425,9 @@ def _train_chunk(skill_model, ref_model, ckpt: CheckpointEngineManager, samples: ``old_logps`` (recomputed once, before any step) and let GRPOLoss form the PPO ratio against them. A single step keeps ``old_logps=None`` (ratio==1, pure on-policy). The batch is padded to a multiple of ``sft_batch_size`` with advantage-0 copies that - contribute no policy gradient.""" + contribute no policy gradient. View-A context-distillation samples ride the same loss + with a positive constant advantage (``--sft-weight``); single-step (old_logps=None) + that is ``-w*logp + beta*KL`` (KL-anchored weighted cross-entropy).""" trajs = [_train_trajectory(rec) for rec in samples] advs = [float(rec['advantage']) for rec in samples] rem = (-len(trajs)) % args.sft_batch_size @@ -1360,7 +1462,9 @@ def _train_chunk(skill_model, ref_model, ckpt: CheckpointEngineManager, samples: n_steps += 1 ckpt.sync_weights(merge_and_sync=True) metric = skill_model.calculate_metric(is_training=True) - return {'n_samples': len(samples), 'n_steps': n_steps, 'n_micro_batches': micro, + n_sft = sum(1 for s in samples if s.get('sft')) + return {'n_samples': len(samples), 'n_sft': n_sft, 'n_grpo': len(samples) - n_sft, + 'n_steps': n_steps, 'n_micro_batches': micro, 'metric': {k: (float(v) if _is_num(v) else v) for k, v in (metric or {}).items()}} @@ -1374,7 +1478,7 @@ def run_greedy_eval(base_sampler, skill_sampler, eval_records: List[Dict[str, An """SEAM ``val-core/math/acc/mean@1`` on the fixed holdout: ONE greedy skill (T=0) per problem into ONE greedy base solve (T=0). Eval ALWAYS uses view B (query-only, the deployment form: no rubric, since rubric needs an online teacher unavailable at deploy); - no leak filter. Main acc requires both correctness and normal termination. Baseline reuses the disk cache.""" + no leak filter (acc scores correctness alone). Baseline reuses the disk cache.""" baseline_rollout(base_sampler, eval_records, base_dp, args, base_cache) for r in eval_records: r['_view'], r['_rubric_diag'] = 'B', '' @@ -1383,45 +1487,38 @@ def run_greedy_eval(base_sampler, skill_sampler, eval_records: List[Dict[str, An skills = [] for seqs in sg_out: if not seqs: - skills.append(('', '', '', '')) + skills.append(('', '')) continue sresp = _clean_text(getattr(seqs[0], 'decoded', '') or '') - parsed = _extract_skill_blocks(sresp, args.diagnose_skill_format) - skills.append((parsed['skills'] if parsed else '', - (parsed or {}).get('pitfall', ''), - (parsed or {}).get('strategy', ''), - sresp)) + skills.append((_extract_skill(sresp) or '', sresp)) ws_out = _run_samples(base_sampler, - [build_skill_solve_prompt(r['problem'], sk) for r, (sk, _, _, _) in zip(eval_records, skills)], + [build_skill_solve_prompt(r['problem'], sk) for r, (sk, _) in zip(eval_records, skills)], 1, args.max_tokens, base_dp, temperature=0.0) recs = [] - for r, (sk, pitfall_text, strategy_text, sresp), seqs in zip(eval_records, skills, ws_out): + for r, (sk, sresp), seqs in zip(eval_records, skills, ws_out): roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() recs.append({ 'record_type': 'eval_problem', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, + 'data_id': r.get('data_id'), 'problem': r['problem'], 'reference_answer': r['reference_answer'], 'view': r['_view'], 'rubric_diag': r.get('_rubric_diag', ''), - 'baseline_pass': r['_baseline_pass'], 'skill': sk, 'skill_pitfall': pitfall_text, - 'skill_strategy': strategy_text, 'skill_diagnose': pitfall_text, - 'skill_action': strategy_text, 'skill_parseable': bool(sk), + 'baseline_pass': r['_baseline_pass'], 'skill': sk, 'skill_parseable': bool(sk), 'skill_response': sresp, 'withskill_pred': roll['pred'], 'withskill_correct': roll['correct'], 'withskill_terminated': roll['terminated'], - 'withskill_pass': _roll_passed(roll), 'withskill_stop_reason': roll['stop_reason'], 'withskill_text': roll['text'], }) - acc = lambda rs: sum(1 for x in rs if x['withskill_pass']) / len(rs) if rs else 0.0 + acc = lambda rs: sum(1 for x in rs if x['withskill_correct']) / len(rs) if rs else 0.0 ws = acc(recs) # all view B (deployment form) base = sum(x['baseline_pass'] for x in recs) / len(recs) if recs else 0.0 - correct = (sum(1 for x in recs if x['withskill_correct']) / len(recs)) if recs else 0.0 fmt = (sum(1 for x in recs if x['skill_parseable']) / len(recs)) if recs else 0.0 term = (sum(1 for x in recs if x['withskill_terminated']) / len(recs)) if recs else 0.0 summary = {'record_type': 'eval_summary', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, 'n': len(recs), 'view': 'B', 'acc_mean1': ws, 'baseline_acc_mean1': base, 'lift_mean1': ws - base, - 'correct_mean1': correct, 'format_mean1': fmt, 'term_mean1': term} + 'format_mean1': fmt, 'term_mean1': term} metrics = {'core/math/acc/mean@1': ws, 'core/math/baseline_acc/mean@1': base, - 'core/math/lift/mean@1': ws - base, 'core/math/correct/mean@1': correct, - 'core/math/format/mean@1': fmt, 'core/math/term/mean@1': term} + 'core/math/lift/mean@1': ws - base, 'core/math/format/mean@1': fmt, + 'core/math/term/mean@1': term} return recs, summary, metrics @@ -1457,6 +1554,9 @@ def _swan_metrics(summary: Dict[str, Any], log: Optional[Dict[str, Any]]) -> Dic # --- base failure taxonomy (loop-out-of-length vs plain wrong) --- 'fail/loop_frac': summary['fail_loop_frac'], 'fail/wrong_frac': summary['fail_wrong_frac'], 'fail/frac_first_try': (summary['n_failed_first_try'] / summary['n']) if summary['n'] else 0.0, + # --- view-A routing: dropped = no-FAIL rubric or rubric leak (SFT-only view A) --- + 'viewa/dropped_frac': (summary['n_viewa_dropped'] / summary['view_A']['n'] + if summary['view_A']['n'] else 0.0), } bal = summary.get('balance') or {} if bal.get('enabled'): @@ -1571,6 +1671,12 @@ def _build_args() -> argparse.Namespace: formatter_class=argparse.RawDescriptionHelpFormatter) p.add_argument('--dataset', choices=('aops', 'math'), default='aops') p.add_argument('--n', type=int, default=2000, help='Problems loaded into the draw pool.') + p.add_argument('--pool-offset', type=int, default=0, + help='Skip this many shuffled non-eval records before building the train pool; ' + 'useful to avoid cold-start SFT data ranges.') + p.add_argument('--exclude-data-ids', default='', + help='Comma-separated jsonl files whose data_id/problem keys are excluded ' + 'from train/eval selection, e.g. coldstart_sft.jsonl.') p.add_argument('--seed', type=int, default=42) p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) p.add_argument('--eval-size', type=int, default=128, help='Fixed holdout size (0 disables).') @@ -1583,18 +1689,15 @@ def _build_args() -> argparse.Namespace: p.add_argument('--balance-max-draws-mult', type=int, default=8) p.add_argument('--n-skills', type=int, default=8) p.add_argument('--view-b-frac', type=float, default=0.5) - p.add_argument('--xproblem-rubric', action=argparse.BooleanOptionalAction, default=True, - help='View A uses its bag-of-words NEIGHBOUR problem\'s rubric (transfer ' - 'test; kills answer leakage since the neighbour\'s answer differs). ' - 'On (default); --no-xproblem-rubric = each problem uses its own rubric ' - '(may leak).') + p.add_argument('--xproblem-rubric', action=argparse.BooleanOptionalAction, default=False, + help='When enabled, view A uses a bag-of-words NEIGHBOUR problem rubric. ' + 'Default is off: each view-A problem uses its own baseline attempt, ' + 'while the API diagnosis prompt is constrained to be answer-free and ' + 'method-level only.') p.add_argument('--skill-retries', type=int, default=2) p.add_argument('--skill-gen-temperature', type=float, default=1.0) p.add_argument('--skill-gen-top-p', type=float, default=1.0) p.add_argument('--skill-gen-top-k', type=int, default=-1) - p.add_argument('--diagnose-skill-format', action=argparse.BooleanOptionalAction, default=True, - help='Require and blocks for skill-gen format; ' - '--no-diagnose-skill-format falls back to legacy .') p.add_argument('--max-model-len', type=int, default=16384) p.add_argument('--max-tokens', type=int, default=8192) p.add_argument('--skill-max-tokens', type=int, default=8192) @@ -1615,6 +1718,19 @@ def _build_args() -> argparse.Namespace: help='Symmetric clip for group-relative advantages; <=0 disables clipping.') p.add_argument('--kl-beta', type=float, default=0.001, help='SEAM-style reference KL coefficient for GRPOLoss.') + p.add_argument('--viewa-sft', action=argparse.BooleanOptionalAction, default=True, + help='Route view-A problems to online context distillation (SFT on the best ' + 'passing open-book skill, rebuilt as a query-only prompt) instead of GRPO. ' + 'View B stays GRPO; both share one optimizer step.') + p.add_argument('--sft-weight', type=float, default=0.5, + help='Advantage magnitude (lambda) for view-A SFT distillation samples in the ' + 'shared GRPOLoss. Single-step this is -lambda*logp + beta*KL (weighted CE). ' + 'Keep <~1 so one-signed SFT gradients do not swamp mean-0 GRPO advantages.') + p.add_argument('--sft-target-len', type=int, default=550, + help='Target skill length (chars) for view-A SFT distillation: among passing ' + 'candidates pick the one CLOSEST to this length (not the shortest). Set near ' + 'the empirical pass-rate peak (~500-600 here) so distillation neither shrinks ' + 'rollouts toward zero nor lets them grow unbounded.') p.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True) p.add_argument('--lr', type=float, default=6e-6) p.add_argument('--max-train-rounds', type=int, default=1500) @@ -1684,16 +1800,19 @@ def main() -> None: cfg = {'record_type': 'config', 'model': MODEL_ID, 'dataset': args.dataset, 'n': len(records), 'eval_n': len(eval_records), 'seed': args.seed, + 'pool_offset': args.pool_offset, 'exclude_data_ids': args.exclude_data_ids, + 'excluded_records': data_stats.get('excluded_records', 0), 'numeric_only': args.numeric_only, 'raw_loaded': data_stats['raw_loaded'], 'numeric_dropped': data_stats['numeric_dropped'], 'eval_every': args.eval_every, 'n_skills': args.n_skills, 'view_b_frac': args.view_b_frac, - 'diagnose_skill_format': args.diagnose_skill_format, 'skill_retries': args.skill_retries, 'balance': args.balance, 'balance_success_frac': args.balance_success_frac, - 'skill_gen_temp': args.skill_gen_temperature, 'reward': 'greedy_binary(correct_and_terminated)', + 'skill_gen_temp': args.skill_gen_temperature, 'reward': 'greedy_binary(correct)', 'advantage': 'group_relative', 'format_in_reward': args.format_in_reward, 'cache': use_cache, - 'rubric_check': 'fixed_math_5crit(viewA)' if checker else 'disabled', + 'rubric_check': f'fixed_math_{len(_MATH_RUBRIC)}crit+term(viewA)' if checker else 'disabled', 'xproblem_rubric': args.xproblem_rubric, + 'viewa_sft': args.viewa_sft, 'sft_weight': args.sft_weight, + 'sft_target_len': args.sft_target_len, 'adv_clip': args.adv_clip, 'grpo_epsilon': args.grpo_epsilon, 'kl_beta': args.kl_beta, 'lr': args.lr, 'train_gpus': TRAIN_GPUS, 'ref_gpus': REF_GPUS, 'ref_fsdp': REF_FSDP, @@ -1790,9 +1909,11 @@ def _prefetch(peeked: List[Dict[str, Any]]) -> None: + ('!' if balance.get('budget_hit') else '') + ') ') if balance.get('enabled') else '' xp = summary.get('xproblem') xp_str = f'pair={xp["pair_rate"]:.2f}@{xp["neighbor_sim_mean"]:.2f} ' if xp else '' + tr_str = (f'train={log["n_grpo"]}g/{log["n_sft"]}s ' if log + else f'train={summary["n_train_samples"]} ') sys.stderr.write( f'[gen] e{pool.epoch} g{gstep}: {bal_str}n={summary["n"]} ' - f'clean={summary["n_clean"]} leak={summary["leak_rate"]:.2f} train={summary["n_train_samples"]} ' + f'clean={summary["n_clean"]} leak={summary["leak_rate"]:.2f} {tr_str}' f'0grad={sig["zero_grad_frac"]:.2f} R={sig["reward_mean"]:.2f}±{sig["reward_std"]:.2f} ' f'acc={summary["avg_baseline_pass"]:.2f}->{summary["avg_withskill_pass"]:.2f} ' f'lift={summary["avg_lift"]:+.3f} {xp_str}' diff --git a/cookbook/exp/embedding/train_reflexion_skill.sh b/cookbook/exp/embedding/train_reflexion_skill.sh index 7a76372b0..ad7273a68 100755 --- a/cookbook/exp/embedding/train_reflexion_skill.sh +++ b/cookbook/exp/embedding/train_reflexion_skill.sh @@ -26,24 +26,25 @@ export GEN_GPU_MEM=${GEN_GPU_MEM:-0.8} export LLM_BACKUP_API_KEY=${LLM_BACKUP_API_KEY:-} export LLM_BACKUP_BASE_URL=${LLM_BACKUP_BASE_URL:-https://dashscope.aliyuncs.com/compatible-mode/v1} export LLM_BACKUP_MODEL=${LLM_BACKUP_MODEL:-qwen3.7-max} +export EXCLUDE_DATA_IDS=${EXCLUDE_DATA_IDS:-./output/reflexion_coldstart_sft/coldstart_sft.jsonl} python cookbook/exp/embedding/train_reflexion_skill.py \ --dataset aops \ --n 5000 \ --numeric-only \ - --chunk-size 32 \ - --n-skills 16 \ - --view-b-frac 0.5 \ - --xproblem-rubric \ + --chunk-size 64 \ + --n-skills 8 \ + --view-b-frac 0.30 \ --skill-retries 2 \ --balance \ - --balance-success-frac 0.3 \ + --balance-success-frac 0.2 \ --balance-loop-frac 0.5 \ --balance-max-draws-mult 8 \ --max-tokens 8192 \ --skill-max-tokens 4096 \ --max-model-len 16384 \ --eval-size 128 \ + --exclude-data-ids "${EXCLUDE_DATA_IDS}" \ --eval-every 5 \ --sft-batch-size 8 \ --ppo-mini-batch-size 0 \ @@ -57,4 +58,4 @@ python cookbook/exp/embedding/train_reflexion_skill.py \ --prefetch-baseline \ --output-dir ./output/reflexion_skill \ --swanlab-project twinkle \ - --swanlab-exp reflexion_skill_rft + --swanlab-exp reflexion_skill_sft35 From ee54fe9064cfab07f3390a9752f60cddd2861da8 Mon Sep 17 00:00:00 2001 From: tastelikefeet Date: Tue, 21 Jul 2026 19:47:54 +0800 Subject: [PATCH 23/60] wip --- .../exp/embedding/train_reflexion_skill.py | 57 +- .../exp/embedding/train_reflexion_skill.sh | 11 +- cookbook/exp/embedding/train_skill_v2.py | 1450 +++++++++++++++++ cookbook/exp/embedding/train_skill_v2.sh | 59 + 4 files changed, 1559 insertions(+), 18 deletions(-) create mode 100644 cookbook/exp/embedding/train_skill_v2.py create mode 100644 cookbook/exp/embedding/train_skill_v2.sh diff --git a/cookbook/exp/embedding/train_reflexion_skill.py b/cookbook/exp/embedding/train_reflexion_skill.py index 5d323ee08..beccbc7b3 100644 --- a/cookbook/exp/embedding/train_reflexion_skill.py +++ b/cookbook/exp/embedding/train_reflexion_skill.py @@ -324,6 +324,18 @@ def _assign_view(problem: str, args: argparse.Namespace) -> str: return 'B' if (h % 100000) / 100000.0 < args.view_b_frac else 'A' +def _curriculum_view_b_frac(gstep: int, args: argparse.Namespace) -> float: + """View-A anneal (--viewa-frac-start): the view-A share holds at ``viewa_frac_start`` + for the first ``viewa_warmup_chunks`` chunks (pure-SFT warmup when start==1.0), then + decays linearly to ``viewa_frac_end`` over ``viewa_decay_chunks`` chunks and holds. + Because _assign_view is a fixed hash against a moving threshold, the B set grows + MONOTONICALLY: a problem trained open-book (A) early can only reappear closed-book + (B) later, never the reverse.""" + t = min(max(gstep - args.viewa_warmup_chunks, 0) / max(args.viewa_decay_chunks, 1), 1.0) + share = args.viewa_frac_start + (args.viewa_frac_end - args.viewa_frac_start) * t + return 1.0 - share + + def _view_prompt(r: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: return {'messages': _skillgen_messages( r['problem'], r['_view'], r.get('_rubric_diag', ''), r.get('_rubric_src', ''))} @@ -1048,28 +1060,28 @@ def _assign_advantages(hard: List[Dict[str, Any]], args: argparse.Namespace) -> def _best_sft_candidate(r: Dict[str, Any], args: argparse.Namespace) -> Optional[Dict[str, Any]]: - """Pick ONE view-A candidate to distill (online context distillation). PREFER the - executor-verified PASSING skills (reward==1); if NONE passed -- common on the hard - problems that are exactly the cases worth distilling -- FALL BACK to any parseable - open-book skill regardless of the executor outcome. Answer-bearing skills produced by - the skill model itself are allowed here; only the external API/rubric diagnosis must be - answer-free. Within the chosen tier, take the one whose skill length is CLOSEST to - ``--sft-target-len`` -- an empirically high-pass-rate length (~500-600 chars in this - run) -- breaking ties by the fewest executor solve tokens. Targeting a length (rather - than the minimum) avoids a distillation feedback loop that would otherwise drive - rollouts ever shorter. None only when no parseable candidate exists at all.""" + """Pick ONE view-A candidate to distill (online context distillation). ONLY + executor-verified PASSING skills (reward==1) are distilled: the earlier fallback to + unverified skills meant ~60% of SFT targets had failed their own executor pass + (measured on the sft35 run) and the model was imitating plausible-but-wrong skills. + Problems with no passing candidate now yield NO SFT record. Answer-bearing skills + produced by the skill model itself are allowed here; only the external API/rubric + diagnosis must be answer-free. Among the passing candidates, take the one whose skill + length is CLOSEST to ``--sft-target-len`` -- an empirically high-pass-rate length + (~500-600 chars in this run) -- breaking ties by the fewest executor solve tokens. + Targeting a length (rather than the minimum) avoids a distillation feedback loop that + would otherwise drive rollouts ever shorter.""" eligible = [c for c in r['_cands'] if c.get('parseable') and c.get('skills')] - if not eligible: - return None passing = [c for c in eligible if c.get('reward') == 1.0] - cs = passing or eligible + if not passing: + return None target = int(getattr(args, 'sft_target_len', 550) or 550) def _solve_tokens(c: Dict[str, Any]) -> int: rolls = c.get('rolls') or [] return rolls[0].get('gen_tokens', 1 << 30) if rolls else (1 << 30) - return min(cs, key=lambda c: (abs(len(c['skills']) - target), _solve_tokens(c))) + return min(passing, key=lambda c: (abs(len(c['skills']) - target), _solve_tokens(c))) def apply_neighbor_rubric(base_sampler, chunk: List[Dict[str, Any]], @@ -1689,6 +1701,16 @@ def _build_args() -> argparse.Namespace: p.add_argument('--balance-max-draws-mult', type=int, default=8) p.add_argument('--n-skills', type=int, default=8) p.add_argument('--view-b-frac', type=float, default=0.5) + p.add_argument('--viewa-frac-start', type=float, default=None, + help='Enable the view-A curriculum: chunk 0 uses this view-A share ' + '(view_b_frac = 1 - share), decaying linearly to --viewa-frac-end ' + 'over --viewa-decay-chunks chunks, then holding. Overrides ' + '--view-b-frac for every chunk.') + p.add_argument('--viewa-frac-end', type=float, default=0.1) + p.add_argument('--viewa-warmup-chunks', type=int, default=0, + help='Hold the view-A share at --viewa-frac-start for this many chunks ' + 'before the linear decay begins.') + p.add_argument('--viewa-decay-chunks', type=int, default=40) p.add_argument('--xproblem-rubric', action=argparse.BooleanOptionalAction, default=False, help='When enabled, view A uses a bag-of-words NEIGHBOUR problem rubric. ' 'Default is off: each view-A problem uses its own baseline attempt, ' @@ -1805,6 +1827,9 @@ def main() -> None: 'numeric_only': args.numeric_only, 'raw_loaded': data_stats['raw_loaded'], 'numeric_dropped': data_stats['numeric_dropped'], 'eval_every': args.eval_every, 'n_skills': args.n_skills, 'view_b_frac': args.view_b_frac, + 'viewa_frac_start': args.viewa_frac_start, 'viewa_frac_end': args.viewa_frac_end, + 'viewa_warmup_chunks': args.viewa_warmup_chunks, + 'viewa_decay_chunks': args.viewa_decay_chunks, 'skill_retries': args.skill_retries, 'balance': args.balance, 'balance_success_frac': args.balance_success_frac, 'skill_gen_temp': args.skill_gen_temperature, 'reward': 'greedy_binary(correct)', @@ -1871,6 +1896,8 @@ def _prefetch(peeked: List[Dict[str, Any]]) -> None: if pending is not None: pending.result() # finish last round's prefetch before drawing (cache-warm) pending = None + if args.viewa_frac_start is not None: + args.view_b_frac = _curriculum_view_b_frac(gstep, args) chunk, balance = draw_chunk(pool, base_sampler, base_dp, args, base_cache) if prefetch_pool is not None: peeked = pool.peek(int(args.chunk_size * (args.balance_max_draws_mult if args.balance else 1))) @@ -1879,6 +1906,7 @@ def _prefetch(peeked: List[Dict[str, Any]]) -> None: base_sampler, skill_sampler, chunk, gstep, base_dp, skill_dp, args, checker, rubric_cache, base_cache, neighbor_map, pool_answers) summary['balance'] = balance + summary['view_b_frac'] = round(args.view_b_frac, 4) log = None if groups: @@ -1922,6 +1950,7 @@ def _prefetch(peeked: List[Dict[str, Any]]) -> None: if use_swan: swan_metrics = _swan_metrics(summary, log) swan_metrics.update(_view_a_rubric_leak_metrics(chunk, pool_answers)) + swan_metrics['train/view_b_frac'] = float(args.view_b_frac) swanlab.log(swan_metrics, step=gstep) if eval_records and (gstep + 1) % args.eval_every == 0: diff --git a/cookbook/exp/embedding/train_reflexion_skill.sh b/cookbook/exp/embedding/train_reflexion_skill.sh index ad7273a68..e732a6a3f 100755 --- a/cookbook/exp/embedding/train_reflexion_skill.sh +++ b/cookbook/exp/embedding/train_reflexion_skill.sh @@ -30,11 +30,14 @@ export EXCLUDE_DATA_IDS=${EXCLUDE_DATA_IDS:-./output/reflexion_coldstart_sft/col python cookbook/exp/embedding/train_reflexion_skill.py \ --dataset aops \ - --n 5000 \ + --n 10000 \ --numeric-only \ --chunk-size 64 \ --n-skills 8 \ - --view-b-frac 0.30 \ + --viewa-frac-start 1.0 \ + --viewa-frac-end 0.1 \ + --viewa-warmup-chunks 20 \ + --viewa-decay-chunks 40 \ --skill-retries 2 \ --balance \ --balance-success-frac 0.2 \ @@ -56,6 +59,6 @@ python cookbook/exp/embedding/train_reflexion_skill.py \ --save-rounds 200 \ --trend-every 10 \ --prefetch-baseline \ - --output-dir ./output/reflexion_skill \ + --output-dir ./output/reflexion_skill_curriculum \ --swanlab-project twinkle \ - --swanlab-exp reflexion_skill_sft35 + --swanlab-exp reflexion_skill_curriculum diff --git a/cookbook/exp/embedding/train_skill_v2.py b/cookbook/exp/embedding/train_skill_v2.py new file mode 100644 index 000000000..6e4d7d201 --- /dev/null +++ b/cookbook/exp/embedding/train_skill_v2.py @@ -0,0 +1,1450 @@ +"""Simplified GRPO + buffer-distill training for the reflexion skill generator (v2). + +Key differences from train_reflexion_skill.py: +- No view A/B split: all skill-gen is query-only (deployment form). +- No baseline rollout in training, no balance selection. +- thinking OFF; skill model outputs optional analysis then block. +- Reward = parseable × (correct AND terminated) × min(1, len_budget/skill_len). +- Buffer A: adv=0 (all-fail) problems accumulate failure trajectories. +- Buffer B: batch rubric → regenerate skill → pass@k validate → SFT injection. +- SFT is event-driven: buffer B reaches threshold → one SFT pass → eval. + +Launch: + LLM_BACKUP_API_KEY=... python cookbook/exp/embedding/train_skill_v2.py \ + --dataset aops --n 5000 --chunk-size 16 --lr 6e-6 +""" +import argparse +import copy +import hashlib +import json +import math +import os +import re +import sys +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Optional, Set, Tuple + +import numpy as np + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.data_format import SamplingParams +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.model import TransformersModel +from twinkle.patch.no_split_modules import NoSplitModulesPatch +from twinkle.processor import InputProcessor +from twinkle.sampler import vLLMSampler +from twinkle.template import Template +from twinkle_agentic.verifier import RubricVerifier +from twinkle_agentic.verifier.rubric_verifier import RubricItem + +logger = get_logger() + +try: + import swanlab +except ImportError: + swanlab = None + +MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') +GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.8)) +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') + +TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 2)) +REF_GPUS = int(os.environ.get('REF_GPUS', 2)) +SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) +BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 2)) +NUM_GPUS = TRAIN_GPUS + REF_GPUS + SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS +TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', min(1, TRAIN_GPUS))) +REF_FSDP = int(os.environ.get('REF_FSDP', min(1, REF_GPUS))) +TRAIN_DP = TRAIN_GPUS // TRAIN_FSDP +REF_DP = REF_GPUS // REF_FSDP + + +# =========================================================================== +# Section A — boxed extraction + answer grading (verbatim from v1) +# =========================================================================== +_BOXED_RE = re.compile(r'\\boxed\s*\{') + + +def extract_boxed(text: str) -> Optional[str]: + if not text: + return None + last = None + for m in _BOXED_RE.finditer(text): + depth, i = 1, m.end() + while i < len(text) and depth > 0: + depth += (text[i] == '{') - (text[i] == '}') + i += 1 + if depth == 0: + last = text[m.end():i - 1].strip() + return last + + +def normalize_answer(ans: str) -> str: + if not ans: + return '' + s = str(ans).strip() + m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) + if m: + return m.group(1) + s = s.strip('$').strip().replace('\u2212', '-') + 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').replace(r'\tfrac', r'\frac') + s = s.replace(r'\,', '').replace(r'\;', '').replace(r'\!', '') + s = re.sub(r'\\(?:quad|qquad|\s)', '', s) + s = re.sub(r'\s+', '', s) + s = re.sub(r'\\sqrt([0-9A-Za-z])', r'\\sqrt{\1}', s) + s = re.sub(r'\\frac\{([^{}]+)\}([^{}\\])', r'\\frac{\1}{\2}', s) + s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) + s = re.sub(r'\{[a-zA-Z]+\}$', '', s) + s = re.sub(r'\^\\circ|\^\{\\circ\}|\u00b0|\\circ', 'deg', s) + s = s.replace(r'\minus{}', '-').replace(r'\minus', '-') + + def _frac_to_slash(mt): + text = mt.group(0) + pos = text.index('{') + 1 + depth, num_start = 1, pos + while depth > 0: + depth += (text[pos] == '{') - (text[pos] == '}') + pos += 1 + numer = text[num_start:pos - 1] + pos += 1 + den_start, depth = pos, 1 + while depth > 0: + depth += (text[pos] == '{') - (text[pos] == '}') + pos += 1 + return f'({numer})/({text[den_start:pos - 1]})' + + s = re.sub(r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', _frac_to_slash, s) + s = re.sub(r'(? bool: + try: + va, vb = float(a.replace('(', '').replace(')', '')), 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) + return va is not None and vb is not None and abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) + + +_MCQ_REF_RE = re.compile( + r'^\\(?:textbf|text|mathrm|mathbf)\{\(?([A-E])\)?\s*\}\s*(.+)$' + r'|^\(?([A-E])\)\s+(.+)$') +_VAR_PREFIX_RE = re.compile(r'^(?:[A-Za-z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)$', re.DOTALL) + + +def _split_mcq(ans): + s = ans.strip() + m = _MCQ_REF_RE.match(s) + if m: + return (m.group(1) or m.group(3)), ((m.group(2) or m.group(4) or '').strip() or None) + bl = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) + return (bl.group(1), None) if bl else (None, s or None) + + +def _strip_var_prefix(ans): + m = _VAR_PREFIX_RE.match((ans or '').strip()) + return m.group(1).strip() if m else (ans or '') + + +def _math_verify_equal(predicted: str, reference: str) -> bool: + try: + from math_verify import parse, verify + gold = parse(r'\boxed{' + reference + '}', parsing_timeout=1) + pred = parse(r'\boxed{' + predicted + '}', parsing_timeout=1) + return bool(gold and pred and verify(gold, pred, timeout_seconds=1)) + except Exception: + return False + + +def answers_match(predicted: str, reference: str) -> bool: + if not predicted or not reference: + return False + norm_p, norm_r = normalize_answer(predicted), normalize_answer(reference) + if norm_p == norm_r or norm_p.lower() == norm_r.lower() or _try_numeric_equal(norm_p, norm_r): + return True + stripped_p = normalize_answer(_strip_var_prefix(predicted)) + stripped_r = normalize_answer(_strip_var_prefix(reference)) + if stripped_p and stripped_r: + if (stripped_p == stripped_r or stripped_p.lower() == stripped_r.lower() + or _try_numeric_equal(stripped_p, stripped_r)): + return True + p_letter, p_value = _split_mcq(predicted) + r_letter, r_value = _split_mcq(reference) + if p_letter and r_letter and p_letter == r_letter: + return True + if (p_letter and p_value is None) and (r_letter and r_value is None): + return False + 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 and (p_val == r_val or p_val.lower() == r_val.lower() or _try_numeric_equal(p_val, r_val)): + return True + for left, right in ((norm_p, norm_r), (stripped_p, stripped_r)): + tl, tr = re.sub(r'[\s()\[\]{}\\]', '', left or ''), re.sub(r'[\s()\[\]{}\\]', '', right or '') + if ',' in tl and tl == tr: + return True + if '=' in norm_r and '=' not in norm_p: + if any(part and (part == norm_p or _try_numeric_equal(part, norm_p)) for part in norm_r.split('=')): + return True + if '=' in norm_p and '=' not in norm_r: + if any(part and (part == norm_r or _try_numeric_equal(part, norm_r)) for part in norm_p.split('=')): + return True + return _math_verify_equal(stripped_p or norm_p, stripped_r or norm_r) + + +_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') + + +def _numeric_value(raw) -> Optional[str]: + if raw is None: + return None + s = str(raw).strip().strip('$').strip() + s = s.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac') + s = re.sub(r'\\!|\\,|\\;|\\ |\\left|\\right|\s', '', s) + for pat in (r'\\frac\{(-?\d+)\}\{(-?\d+)\}', r'(-?\d+)/(-?\d+)'): + m = re.fullmatch(pat, s) + if m: + a, b = int(m.group(1)), int(m.group(2)) + return str(int(a / b)) if (b and a / b == int(a / b)) else (str(a / b) if b else None) + return (str(int(float(s))) if float(s) == int(float(s)) else str(float(s))) if _NUM_RE.fullmatch(s) else None + + +def _answer_leaked(skill: str, reference: str) -> bool: + if not skill: + return False + # Suffix guard: reject only a following DIGIT or a following '.' (decimal point), + # NOT a sentence-ending '.'. Old '(?![\d.])' let leaks like "...= 675." slip through + # because the trailing period satisfied the [\d.] class. 中文注释：尾断言只排除"后接数字" + # 或"后接小数点+数字"，不排除句末句号，堵住 "答案." 这类泄漏漏检。 + for cand in {_numeric_value(reference), (str(reference).strip() or None)}: + if cand and re.search(r'(?') + + +def _clean_text(decoded: Optional[str]) -> str: + return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() + + +def _extract_skill(text: str) -> Optional[str]: + """Parse ... block from skill-gen output.""" + low = text.lower() + end_think = low.rfind('') + answer = text[end_think + len(''):] if end_think >= 0 else text + open_tag, close_tag = '', '' + s = answer.lower().rfind(open_tag) + if s < 0: + return None + inner = s + len(open_tag) + e = answer.lower().find(close_tag, inner) + if e < 0: + return None + block = answer[inner:e].strip() + block = re.sub(r'', '', block, flags=re.IGNORECASE).strip() + return block or None + + +def _parse_seq(seq, gold: str) -> Dict[str, Any]: + text = _clean_text(getattr(seq, 'decoded', '') or '') + pred = extract_boxed(text) + correct = bool(pred) and answers_match(pred, gold) + terminated = getattr(seq, 'stop_reason', None) != 'length' + return {'pred': pred, 'correct': correct, 'terminated': terminated, + 'stop_reason': getattr(seq, 'stop_reason', None), + 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} + + +def _empty_roll(): + return {'pred': '', 'correct': False, 'terminated': False, + 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} + + +def _run_samples(sampler, prompts, num_samples, max_tokens, gen_dp, + temperature=None, top_p=None, top_k=None): + if not prompts: + return [] + params = SamplingParams( + max_tokens=max_tokens, + temperature=GEN_TEMPERATURE if temperature is None else temperature, + top_p=GEN_TOP_P if top_p is None else top_p, + num_samples=num_samples, **({} if top_k is None else {'top_k': top_k})) + padded = prompts + if gen_dp > 1 and 0 < len(prompts) < gen_dp: + padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] + responses = sampler.sample(padded, params)[:len(prompts)] + return [list(r.sequences) if (r and r.sequences) else [] for r in responses] + + +# =========================================================================== +# Section C — data loading (simplified: no balance, no xproblem, no views) +# =========================================================================== +def _boxed_batch(rows, dataset): + sols = rows['solution'] + metas = rows.get('metadata', [None] * len(sols)) + refs = [extract_boxed(s or '') for s in sols] + keep = [bool(ref) and (dataset != 'aops' or bool((meta or {}).get('boxed'))) + for ref, meta in zip(refs, metas)] + return {**rows, 'reference_answer': refs, '_keep': keep} + + +def load_problems(dataset: str, n: int, seed: int) -> List[Dict[str, Any]]: + ds_id = AOPS_DATASET_ID if dataset == 'aops' else os.environ.get('MATH_DATASET_ID', 'modelscope/competition_math') + ds = Dataset(DatasetMeta(dataset_id=f'ms://{ds_id}', split='train')) + nproc = min(32, os.cpu_count() or 1) + ds.map(lambda rows: _boxed_batch(rows, dataset), num_proc=nproc) + ds.filter(lambda row: row['_keep'], num_proc=nproc) + out = [{'data_id': f'{dataset}:{i}', 'problem': row['problem'], + 'reference_answer': row['reference_answer']} + for i, row in enumerate(ds.dataset)] + rng = np.random.RandomState(seed) + rng.shuffle(out) + return out[:n] if (n and n < len(out)) else out + + +def _load_records(args): + records = load_problems(args.dataset, 0, args.seed) + raw_n = len(records) + if args.numeric_only: + records = [{**r, 'reference_answer': v} + for r, v in ((r, _numeric_value(r.get('reference_answer'))) for r in records) + if v is not None] + np.random.RandomState(args.seed).shuffle(records) + # exclude + excl_ids, excl_probs = set(), set() + for path in (args.exclude_data_ids or '').split(','): + path = path.strip() + if not path or not os.path.exists(path): + continue + with open(path) as f: + for line in f: + if not line.strip(): continue + row = json.loads(line) + if row.get('record_type') in {'config', 'summary'}: continue + did = str(row.get('data_id', '')).strip() + if did: excl_ids.add(did) + else: + p = str(row.get('problem', '')).strip() + if p: excl_probs.add(p) + if excl_ids or excl_probs: + records = [r for r in records + if str(r.get('data_id', '')) not in excl_ids + and str(r.get('problem', '')).strip() not in excl_probs] + eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 + eval_records = records[:eval_n] + # Dedup by problem TEXT: index slices are disjoint, but duplicate problem statements + # across the boundary would still leak eval into train. Drop any train record whose + # problem appears in eval, then guard with an explicit overlap assertion. + # 中文注释：train/eval 去重——按题面文本剔除，防止数据集内重复题目跨界泄漏；末尾硬断言无交集。 + eval_probs = {r['problem'] for r in eval_records} + train_records = [r for r in records[eval_n:] if r['problem'] not in eval_probs] + if args.n > 0: + train_records = train_records[:args.n] + if {r['problem'] for r in train_records} & eval_probs: + raise ValueError('eval/train overlap detected after dedup') + logger.info(f'[data] raw={raw_n} train={len(train_records)} eval={len(eval_records)}') + return train_records, eval_records + + +# =========================================================================== +# Section D — DiskCache, ProblemPool, LockedSampler +# =========================================================================== +class DiskCache: + def __init__(self, path: str, enabled: bool = True): + self._mem: Dict[str, Any] = {} + self._fh = None + self._lock = threading.Lock() + if not enabled: + return + if os.path.exists(path): + with open(path, encoding='utf-8') as f: + for line in f: + if line.strip(): + row = json.loads(line) + self._mem[row['key']] = row['value'] + self._fh = open(path, 'a', encoding='utf-8') + + @staticmethod + def key_for(*parts): + return hashlib.md5('\x1f'.join(parts).encode('utf-8')).hexdigest() + + def get(self, key): return self._mem.get(key) + def __contains__(self, key): return key in self._mem + + def put(self, key, value): + with self._lock: + self._mem[key] = value + if self._fh: + self._fh.write(json.dumps({'key': key, 'value': value}, ensure_ascii=False) + '\n') + self._fh.flush() + + def close(self): + if self._fh: self._fh.close() + + +class ProblemPool: + def __init__(self, records, seed): + self._records = list(records) + self._seed, self._cursor, self.epoch = seed, 0, 0 + + def draw(self, k): + out, seen = [], set() + while len(out) < k: + if self._cursor >= len(self._records): + self.epoch += 1 + np.random.RandomState(self._seed + self.epoch).shuffle(self._records) + self._cursor = 0 + r = self._records[self._cursor] + self._cursor += 1 + if id(r) not in seen: + seen.add(id(r)) + out.append(r) + return out + + +class _LockedSampler: + def __init__(self, sampler): + self._sampler = sampler + self._lock = threading.Lock() + + def sample(self, *a, **kw): + with self._lock: + return self._sampler.sample(*a, **kw) + + def __getattr__(self, name): + return getattr(self._sampler, name) + + +# =========================================================================== +# Section E — Rubric (teacher diagnosis, batched at distill time) +# =========================================================================== +_RFT_DIAG_SYSTEM = """\ +You are a strategy-level process checker for a math solution attempt. You are given a +math problem, a rubric, and one attempted solution segment. Decide PASS or FAIL for each +criterion, and write the diagnosis so it can become useful reusable guidance for solving +similar problems without seeing this segment. + +Output STRICT JSON (no prose outside it) with this shape: +{"items": [{"index": 1, "verdict": "PASS"|"FAIL", "reason": "...", "fix": ""}], "overall": "OK"|"ISSUES", "summary": "..."} + +Rules: +- Judge every criterion independently. +- The diagnosis must stay answer-free. +- For FAIL items: describe the process problem at strategy level. +- A fix suggests the LOCAL correction direction without solving. +- Never reveal the final answer or a corrected expression. +- If segment was cut off (no \\boxed{}), mark length-budget as FAIL. +- Keep "reason" and "fix" concise: one short sentence each. +- Output only the JSON object.""" + +_RFT_DIAG_USER = """\ +## Task / query +{query} + +## Rubric +{rubric} + +## Segment +{segment} + +Now output the diagnostic JSON object.""" + +_MATH_RUBRIC = [ + ('The attempt chooses a method suitable for the problem structure', False), + ('The attempt identifies the key constraint, invariant, or quantity before computing', False), + ('Algebraic and logical transformations preserve validity at each step', True), + ('The attempt checks required constraints, domains, boundary cases, or validity conditions', False), + ('The attempt avoids redundant casework, looping, or re-deriving known facts', False), + ('The attempt reaches a final boxed answer within the length budget', False), + ('The approach stays focused on the actual question asked', False), +] +_RUBRIC_VERSION = 'rubric_v5_viewb_strategy' + + +class _RftRubricVerifier(RubricVerifier): + def _diagnose_trajectory(self, query, rubric_block, segment_text): + return {'messages': [ + {'role': 'system', 'content': _RFT_DIAG_SYSTEM}, + {'role': 'user', 'content': _RFT_DIAG_USER.format( + query=query, rubric=rubric_block, segment=segment_text)}]} + + +def build_rubric_checker(): + if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') + or os.environ.get('OPENAI_API_KEY')): + return None + return _RftRubricVerifier( + fixed_rubric=[RubricItem(t, is_hard=h) for t, h in _MATH_RUBRIC], gate=True) + + +def _format_diagnosis(detail) -> str: + rub = detail.rubric + lines = [] + for it in detail.items: + text = rub[it.index - 1].text if 0 < it.index <= len(rub) else f'criterion {it.index}' + if it.verdict: + lines.append(f'- [PASS] {text}') + else: + tail = f': {it.reason}' if it.reason else '' + tail += f' (fix: {it.fix})' if it.fix else '' + lines.append(f'- [FAIL] {text}{tail}') + if detail.summary: + lines.append(f'Summary: {detail.summary}') + return '\n'.join(lines) + + +def _diagnose_entry(checker, entry: Dict[str, Any]) -> Optional[str]: + """Run the teacher rubric on ONE buffer-A failure trajectory → formatted diagnosis text + (or None on API error). Pure network/CPU (no GPU), so it can run on a background thread + while GRPO trains. Shared by the background pre-diagnosis pool and distill_buffer's + fallback for any entry the pool did not reach in time. + 中文注释：单条失败轨迹的 rubric 诊断（纯 API，不吃 GPU）。后台预诊断与 distill 补诊断共用。""" + seg_text = entry['fail_segment'] + if entry.get('fail_stop_reason') == 'length': + seg_text += ('\n\n[Process note: this attempt was cut off at the token budget ' + 'and never produced a final \\boxed{} answer.]') + seg = {'messages': [{'role': 'user', 'content': entry['problem']}, + {'role': 'assistant', 'content': seg_text}]} + try: + return _format_diagnosis(checker.diagnose(seg, query=entry['problem'])) + except Exception as exc: + logger.warning(f'[rubric] diagnose error: {exc}') + return None + + +# =========================================================================== +# Section F — NEW: prompts, reward, buffer logic +# =========================================================================== + +# ---- Skill-gen system prompt (query-only, thinking OFF) ---- +# 中文注释：skillmodel 系统提示词。thinking 关闭；允许在 之前输出简短分析； +# 要求 skill ≤600 字符、切题、不给答案、不啰嗦；方向弱列举（pitfall/技术点/step/overview/确信输出）。 +SKILL_GEN_SYSTEM = """\ +You are a math guidance writer. Write short, reusable guidance for the problem below. + +Rules: +- You may briefly analyze the problem BEFORE the tag. +- Output your guidance inside .... +- Keep the guidance within 600 characters, concise and problem-specific. +- Focus on: pitfalls may be happened to avoid, key techniques, brief step outlines, or how to help to converge to the answer quickly. +- Do NOT calculate the final answer. +- Do NOT be verbose or generic. + +Output format: +[optional brief analysis] + +Your reusable solving guidance here. +""" + +# ---- Executor system prompt (with skill injection) ---- +# 中文注释：executor 系统提示词，将 skill 注入 solver 的 system prompt 前缀。 +DIRECT_SYSTEM = ( + 'You are an expert competition mathematician. Solve the following problem ' + 'step by step. Provide your final answer inside \\boxed{}.') +_SKILL_SOLVE_PREFIX = ( + DIRECT_SYSTEM + '\n\n' + 'Before you start, keep these reminders in mind to avoid common mistakes on this ' + 'type of problem:\n') +_SKILL_SOLVE_SUFFIX = '\nApply them where relevant, but rely on your own reasoning to reach the answer.' + + +def build_direct_prompt(problem): + return {'messages': [{'role': 'system', 'content': DIRECT_SYSTEM}, + {'role': 'user', 'content': problem}]} + + +def build_skill_solve_prompt(problem, skill): + skill = (skill or '').strip() + if not skill: + return build_direct_prompt(problem) + return {'messages': [ + {'role': 'system', 'content': _SKILL_SOLVE_PREFIX + skill + _SKILL_SOLVE_SUFFIX}, + {'role': 'user', 'content': problem}]} + + +# ---- Rubric-guided regeneration prompt (buffer B distillation) ---- +# 中文注释：蒸馏重生成提示词。给旧 skill + rubric 诊断，要求产出改进后的 skill。 +# 只输出 块，≤600字符，不含答案。 +REGEN_SYSTEM = """\ +You are a math guidance writer. You previously wrote guidance for a problem, but the \ +solver still failed. A process-check diagnosed the failure. Revise your guidance to \ +address the diagnosed issues. + +Rules: +- Output ONLY a ... block (no analysis). +- Keep within 600 characters, concise, problem-specific. +- Address the diagnosed failure points, also keep the good parts of the old skills. +- Do NOT include the final answer.""" + +REGEN_USER = """\ +Problem: +{problem} + +Previous guidance (did not help): +{orig_skill} + +Process-check diagnosis: +{rubric_diag} + +Write improved guidance:""" + + +def _skillgen_prompt(problem: str) -> Dict[str, Any]: + """Skill-gen prompt: query-only, no view split.""" + return {'messages': [ + {'role': 'system', 'content': SKILL_GEN_SYSTEM}, + {'role': 'user', 'content': f'Problem:\n{problem}'}]} + + +def _regen_prompt(problem: str, orig_skill: str, rubric_diag: str) -> Dict[str, Any]: + """Regeneration prompt for buffer B distillation.""" + return {'messages': [ + {'role': 'system', 'content': REGEN_SYSTEM}, + {'role': 'user', 'content': REGEN_USER.format( + problem=problem, orig_skill=orig_skill, rubric_diag=rubric_diag)}]} + + +# ---- Reward ---- +# 中文注释：reward = parseable × (correct AND terminated) × min(1, budget/len) +# parseable=0 的候选 reward=0 仍参与 group（格式压力）；截断=失败；超长乘法衰减。 +def _skill_reward(parseable: bool, correct: bool, terminated: bool, + skill_len: int, len_budget: int) -> float: + if not parseable: + return 0.0 + base = 1.0 if (correct and terminated) else 0.0 + len_factor = min(1.0, len_budget / max(skill_len, 1)) + return base * len_factor + + +# ---- Buffer A: collect adv=0 all-fail problems ---- +def _collect_buffer_a(chunk, args) -> List[Dict[str, Any]]: + """Collect problems where all candidates got reward 0 (adv=0, GRPO blind spot). + Store one representative failure trajectory for later rubric diagnosis.""" + entries = [] + for r in chunk: + cs = [c for c in r['_cands'] if c.get('reward') is not None] + if len(cs) < 2: + continue + rewards = [c['reward'] for c in cs] + if max(rewards) > 0: + continue # has signal, not all-fail + # Representative trajectory for rubric + regen seed: prefer a terminated-wrong + # parseable candidate (complete reasoning to diagnose). Its skill becomes the regen + # seed, so pick the MOST SUBSTANTIAL one WITHIN budget (longest ≤ len_budget) — a + # rich-but-not-bloated starting point — rather than an arbitrary [0] or a near-empty + # skill. If all seeds exceed budget, take the one closest to budget (shortest-over). + # 中文注释：代表轨迹既做 rubric 诊断又做 regen 种子——优先"跑完但答错"的候选（完整推理）， + # 其 skill 取预算内最长(最有实质)的作种子；若全超预算则取最接近预算的，避免随机/近空种子。 + budget = args.len_budget + + def _seed_key(c): + L = len(c.get('skills') or '') + return (L <= budget, L if L <= budget else -L) + + parseable = [c for c in cs if c.get('skills')] + term_wrong = [c for c in parseable if c['rolls'] and c['rolls'][0].get('terminated')] + pool_c = term_wrong or parseable or cs + rep = max(pool_c, key=_seed_key) + stop_dist = {} + for c in cs: + sr = c['rolls'][0]['stop_reason'] if c['rolls'] else 'none' + stop_dist[sr] = stop_dist.get(sr, 0) + 1 + entries.append({ + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'data_id': r.get('data_id', ''), + 'orig_skill': rep.get('skills', ''), + 'orig_len': len(rep.get('skills', '')), + 'fail_segment': rep['rolls'][0]['text'] if rep['rolls'] else '', + 'fail_stop_reason': rep['rolls'][0]['stop_reason'] if rep['rolls'] else 'none', + 'stop_reason_dist': stop_dist, + }) + return entries + + +# ---- Buffer B distillation ---- +def distill_buffer(entries: List[Dict[str, Any]], skill_sampler, base_sampler, + checker, skill_dp: int, base_dp: int, + args) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + """Batch rubric → regenerate K distinct skills → greedy-validate → return SFT records. + 中文注释：蒸馏流程（方案 B，多样性在 skill 侧、executor 用贪心）： + 1. 批量 rubric 诊断失败轨迹；2. 仅 [FAIL] 项用高温重生成 K 个不同候选 skill； + 3. 每个候选过 ≤budget+无leak+去重 过滤；4. 每个存活候选用 executor 贪心(T=0)解 1 次； + 5. gate：≥m 个不同候选达成 terminated-correct → select 长度最接近 budget 的一个入 buffer B。 + 返回 (sft_records, distill_records)：后者逐 entry 记录 rubric_diag/候选 skill/贪心解结果/漏斗 + stage，落盘到 distill_records.jsonl 供复盘（否则 rubric 诊断与候选明细只存在于内存）。""" + if not checker or not entries: + return [], [] + + # Step 1: ensure every entry has a rubric diagnosis. Entries pre-diagnosed in the + # background (see _prediagnose in main) already carry '_rubric_diag'; only the misses + # are diagnosed here (in parallel), so the GPU-idle API wait is normally hidden. + # 中文注释：优先用后台预诊断结果；只对没预诊断到的条目并行补跑，隐藏 API 等待。 + pending = [e for e in entries if not e.get('_rubric_diag')] + if pending: + workers = min(args.rubric_workers, len(pending)) + with ThreadPoolExecutor(max_workers=max(1, workers)) as ex: + diags = list(ex.map(lambda e: _diagnose_entry(checker, e), pending)) + for entry, diag in zip(pending, diags): + entry['_rubric_diag'] = diag or '' + for entry in entries: + entry['rubric_diag'] = entry.get('_rubric_diag') or '' + + # Builder for the structured distill audit records (one per buffer-A entry). Closure over + # `entries`/`args`; takes the per-entry regen skills + greedy solve results (may be empty + # for the early-exit funnel stages). 中文注释：构造逐 entry 的蒸馏审计记录（含漏斗 stage）。 + def _mk_distill(results_by_entry, per_entry_skills, has_fail): + hf_index = {id(e): ei for ei, e in enumerate(has_fail)} + recs = [] + for e in entries: + ei = hf_index.get(id(e)) + cand_results = results_by_entry.get(ei, []) if ei is not None else [] + n_pass = sum(1 for c in cand_results if c['correct'] and c['terminated']) + n_cand = len(per_entry_skills[ei]) if (ei is not None and ei < len(per_entry_skills)) else 0 + if ei is None: + stage = 'no_fail' # rubric 未给出任何 [FAIL] + elif n_cand == 0: + stage = 'no_valid_regen' # 有 [FAIL] 但重生成无一条过 ≤budget/无leak/去重 + elif n_pass >= args.passatk_m: + stage = 'accepted' # ≥m 个候选贪心解对 → 入 buffer B + else: + stage = 'rejected' # 有候选但 args.len_budget: + continue + if _answer_leaked(skill, entry['reference_answer']): + continue + if skill in seen: + continue # 去重：同一 skill 只算一个"不同候选" + seen.add(skill) + skills.append(skill) + per_entry_skills.append(skills) + if skills: + n_entries_with_cands += 1 + + # Flatten to (entry_idx, skill) for ONE batched greedy executor solve per candidate skill. + # 中文注释：每个候选 skill 只用 executor 贪心(T=0)解 1 次——与部署/eval 口径一致； + # k 次 rollout 摊到 k 个不同 skill 上，多样性来自 skill 侧而非 executor 侧。 + flat_idx, flat_prompts = [], [] + for ei, skills in enumerate(per_entry_skills): + for sk in skills: + flat_idx.append((ei, sk)) + flat_prompts.append(build_skill_solve_prompt(has_fail[ei]['problem'], sk)) + if not flat_prompts: + logger.info(f'[distill] {len(has_fail)} [FAIL] entries, 0 valid regen skills') + return [], _mk_distill({}, per_entry_skills, has_fail) + solve_out = _run_samples(base_sampler, flat_prompts, 1, args.max_tokens, base_dp, + temperature=0.0) + + # Gather greedy results back per entry: record EVERY candidate skill's solve outcome + # (not just passers) so distill_records can show why an entry was rejected. + # 中文注释：回收每个候选 skill 的贪心解结果（含未通过的），供审计记录还原被拒原因。 + results_by_entry: Dict[int, List[Dict[str, Any]]] = {} + for (ei, sk), seqs in zip(flat_idx, solve_out): + roll = _parse_seq(seqs[0], has_fail[ei]['reference_answer']) if seqs else _empty_roll() + results_by_entry.setdefault(ei, []).append( + {'skill': sk, 'len': len(sk), 'correct': roll['correct'], 'terminated': roll['terminated']}) + + # Step 4: gate ≥ m distinct greedy-effective skills; select the survivor CLOSEST to the + # length budget (short is the floor, but not so short it degrades to answer-dumping). + # 中文注释：gate——≥m 个不同 skill 在贪心下 terminated-correct；select——在通过的候选里 + # 选长度最接近 budget 的一个入 buffer B（短是地板，但别短到退化成吐答案）。 + sft_records = [] + for ei, cand_results in results_by_entry.items(): + passers = [c['skill'] for c in cand_results if c['correct'] and c['terminated']] + if len(passers) < args.passatk_m: + continue + entry = has_fail[ei] + best = min(passers, key=lambda sk: abs(len(sk) - args.len_budget)) + sft_records.append({ + 'problem': entry['problem'], 'reference_answer': entry['reference_answer'], + 'data_id': entry.get('data_id', ''), + 'response': f'\n{best}\n', + 'skills': best, 'sft': True, + 'n_pass_skills': len(passers), 'n_cand_skills': len(per_entry_skills[ei]), + }) + + logger.info(f'[distill] {len(entries)} A → {len(has_fail)} [FAIL] → ' + f'{n_entries_with_cands} w/cands → {len(sft_records)} validated B ' + f'(gate m={args.passatk_m}/k={k})') + return sft_records, _mk_distill(results_by_entry, per_entry_skills, has_fail) + + + +# =========================================================================== +# Section G — GRPO advantages + training +# =========================================================================== +def _assign_advantages(chunk, args): + """Group-relative advantage: A = (R - mean) / (std + eps). std==0 → adv=0 (skipped).""" + eps = 1e-6 + adv_clip = abs(float(getattr(args, 'adv_clip', 0.0) or 0.0)) + for r in chunk: + for c in r['_cands']: + c['advantage'], c['kept'] = 0.0, False + cs = [c for c in r['_cands'] if c.get('reward') is not None] + if len(cs) < 2: + continue + rewards = [c['reward'] for c in cs] + mean_r = sum(rewards) / len(rewards) + std = (sum((x - mean_r) ** 2 for x in rewards) / len(rewards)) ** 0.5 + if std < 1e-9: + continue + for c in cs: + raw = (c['reward'] - mean_r) / (std + eps) + c['advantage'] = max(-adv_clip, min(adv_clip, raw)) if adv_clip > 0 else raw + c['kept'] = c['reward'] > mean_r + + +def _train_trajectory(rec): + """Rebuild the query-only skill-gen prompt (train/inference match) + response. + GRPO records carry the full generated response; SFT records carry only the + cleaned block. key_rounds selects the final assistant turn.""" + msgs = _skillgen_prompt(rec['problem'])['messages'] + return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], + 'user_data': {'key_rounds': [len(msgs)]}} + + +def _train_step(skill_model, ref_model, ckpt, samples, args): + """On-policy GRPO update over one batch, then sync weights. SFT samples ride the + same GRPOLoss with a positive constant advantage (--sft-weight).""" + trajs = [_train_trajectory(rec) for rec in samples] + advs = [float(rec['advantage']) for rec in samples] + rem = (-len(trajs)) % args.sft_batch_size + if rem: + trajs += [trajs[-1]] * rem + advs += [0.0] * rem + n, sft = len(trajs), args.sft_batch_size + mini = args.ppo_mini_batch_size if args.ppo_mini_batch_size > 0 else n + mini = max(sft, (mini // sft) * sft) + multi_step = mini < n + micro_ref, micro_old = [], [] + for i in range(0, n, sft): + mb = trajs[i:i + sft] + micro_ref.append(ref_model.forward_only(inputs=mb).get('logps')) + micro_old.append(skill_model.forward_only(inputs=mb).get('logps') if multi_step else None) + micro, n_steps = 0, 0 + for ms in range(0, n, mini): + for i in range(ms, min(ms + mini, n), sft): + k = i // sft + skill_model.forward_backward(inputs=trajs[i:i + sft], advantages=advs[i:i + sft], + old_logps=micro_old[k], ref_logps=micro_ref[k]) + micro += 1 + skill_model.clip_grad_and_step() + n_steps += 1 + ckpt.sync_weights(merge_and_sync=True) + metric = skill_model.calculate_metric(is_training=True) + n_sft = sum(1 for s in samples if s.get('sft')) + return {'n_samples': len(samples), 'n_sft': n_sft, 'n_grpo': len(samples) - n_sft, + 'n_steps': n_steps, 'n_micro_batches': micro, + 'metric': {k: (float(v) if _is_num(v) else v) for k, v in (metric or {}).items()}} + + +def _is_num(v): + try: + float(v); return True + except (TypeError, ValueError): + return False + + +# =========================================================================== +# Section H — chunk processing, records, eval (+ hard-slice rescue) +# =========================================================================== +def process_chunk(base_sampler, skill_sampler, chunk, ci, base_dp, skill_dp, args): + """skill-gen (query-only) → leak audit → with-skill greedy pass → reward → advantages. + Returns (full_records, summary, grpo_train_records, buffer_a_entries).""" + for r in chunk: + r['_cands'] = [] + # skill-gen (thinking OFF), re-sample problems with no clean candidate + flat = [] + pending = list(chunk) + for _ in range(args.skill_retries + 1): + if not pending: + break + sg_out = _run_samples(skill_sampler, [_skillgen_prompt(r['problem']) for r in pending], + args.n_skills, args.skill_max_tokens, skill_dp, + temperature=args.skill_gen_temperature, top_p=args.skill_gen_top_p, + top_k=args.skill_gen_top_k) + still = [] + for r, seqs in zip(pending, sg_out): + got = False + for s in seqs: + resp = _clean_text(getattr(s, 'decoded', '') or '') + block = _extract_skill(resp) or '' + cand = {'skills': block, 'response': resp, 'parseable': bool(block), + 'leaked': None, 'with_pass': None, 'reward': None, 'rolls': [], + 'advantage': 0.0, 'kept': False, + 'skillgen_stop': getattr(s, 'stop_reason', None), + 'skillgen_tokens': len(getattr(s, 'tokens', None) or [])} + r['_cands'].append(cand) + if block: + flat.append((r, cand)) + got = True + if not got: + still.append(r) + pending = still + + # leak audit (deterministic, observability only) + for r, c in flat: + c['leaked'] = _answer_leaked(c['skills'], r['reference_answer']) + + # with-skill greedy pass (T=0) + if flat: + ws_out = _run_samples(base_sampler, + [build_skill_solve_prompt(r['problem'], c['skills']) for r, c in flat], + 1, args.max_tokens, base_dp, temperature=0.0) + for (r, c), seqs in zip(flat, ws_out): + roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + c['rolls'] = [roll] + c['with_pass'] = 1.0 if roll['correct'] else 0.0 + c['reward'] = _skill_reward(c['parseable'], roll['correct'], roll['terminated'], + len(c['skills']), args.len_budget) + # unparseable candidates score 0 and still join the group (format pressure) + for r in chunk: + for c in r['_cands']: + if c['reward'] is None: + c['reward'] = 0.0 + + _assign_advantages(chunk, args) + + grpo = [] + for r in chunk: + for c in r['_cands']: + if abs(c.get('advantage') or 0.0) > 1e-9: + grpo.append({'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'data_id': r.get('data_id', ''), 'response': c['response'], + 'skills': c['skills'], 'advantage': c['advantage'], + 'kept': c['kept'], 'reward': c['reward'], 'sft': False}) + buffer_a = _collect_buffer_a(chunk, args) + return _full_records(chunk, ci), _chunk_summary(chunk, ci), grpo, buffer_a + + +def _roll(x): + return {k: x[k] for k in ('pred', 'correct', 'terminated', 'stop_reason', 'gen_tokens', 'text')} + + +def _full_records(chunk, ci): + out = [] + for r in chunk: + out.append({ + 'record_type': 'problem', 'chunk': ci, 'problem': r['problem'], + 'reference_answer': r['reference_answer'], 'data_id': r.get('data_id', ''), + 'candidates': [{ + 'skills': c['skills'], 'response': c['response'], 'parseable': c['parseable'], + 'leaked': c['leaked'], 'with_pass': c['with_pass'], 'reward': c.get('reward'), + 'advantage': c.get('advantage'), 'kept': c.get('kept'), + 'skillgen_stop': c.get('skillgen_stop'), 'skillgen_tokens': c.get('skillgen_tokens'), + 'rolls': [_roll(x) for x in c['rolls']], + } for c in r['_cands']], + }) + return out + + +def _mean(xs): + return sum(xs) / len(xs) if xs else 0.0 + + +def _std(xs): + if len(xs) < 2: + return 0.0 + m = _mean(xs) + return (sum((x - m) ** 2 for x in xs) / len(xs)) ** 0.5 + + +def _chunk_summary(chunk, ci): + all_cands = [c for r in chunk for c in r['_cands']] + cands = [c for c in all_cands if c['parseable']] + scored = [c for c in cands if c['with_pass'] is not None] + ws_rolls = [x for c in scored for x in c['rolls']] + # signal: fraction of groups with zero reward variance (no gradient) + group_vars, all_rewards, zero_grad, groups = [], [], 0, 0 + for r in chunk: + rewards = [c['reward'] for c in r['_cands'] if c.get('reward') is not None] + if len(rewards) < 2: + continue + groups += 1 + all_rewards.extend(rewards) + v = _std(rewards) + group_vars.append(v) + if v < 1e-9: + zero_grad += 1 + n_train = sum(1 for c in all_cands if abs(c.get('advantage') or 0.0) > 1e-9) + trunc = sum(1 for x in ws_rolls if x['stop_reason'] == 'length') + ws_acc = _mean([1.0 if any(c.get('reward') for c in r['_cands']) else 0.0 + for r in chunk if r['_cands']]) + return { + 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), + 'n_generated': len(all_cands), 'n_candidates_parseable': len(cands), + 'parse_rate': (len(cands) / len(all_cands)) if all_cands else 0.0, + 'n_leaked': sum(1 for c in cands if c['leaked']), + 'leak_rate': (sum(1 for c in cands if c['leaked']) / len(cands)) if cands else 0.0, + 'n_train_samples': n_train, 'n_groups': groups, + 'zero_grad_frac': (zero_grad / groups) if groups else 0.0, + 'reward_mean': _mean(all_rewards), 'reward_std': _std(all_rewards), + 'group_reward_std_mean': _mean(group_vars), + 'skill_tokens_mean': _mean([c.get('skillgen_tokens') or 0 for c in cands]), + 'skill_chars_mean': _mean([len(c['skills']) for c in cands]), + 'avg_withskill_pass': ws_acc, + 'candidate_withskill_pass': _mean([c['with_pass'] for c in scored]), + 'withskill_trunc_frac': (trunc / len(ws_rolls)) if ws_rolls else 0.0, + 'termination_rate_withskill': _mean([1.0 if x['terminated'] else 0.0 for x in ws_rolls]), + } + + +def run_greedy_eval(base_sampler, skill_sampler, eval_records, ci, rounds, + base_dp, skill_dp, args, base_cache): + """SEAM mean@1 on the fixed holdout: greedy skill (T=0) → greedy base solve (T=0). + Adds hard-slice (baseline_pass==0) rescue rate as a zero-cost secondary readout.""" + # baseline (frozen, cached) + todo = [r for r in eval_records if DiskCache.key_for(r['problem']) not in base_cache] + if todo: + out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], + 1, args.max_tokens, base_dp, temperature=0.0) + for r, seqs in zip(todo, out): + roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + base_cache.put(DiskCache.key_for(r['problem']), roll) + for r in eval_records: + br = base_cache.get(DiskCache.key_for(r['problem'])) + r['_baseline_pass'] = 1.0 if br['correct'] else 0.0 + # skill-gen (greedy) → with-skill (greedy) + sg_out = _run_samples(skill_sampler, [_skillgen_prompt(r['problem']) for r in eval_records], + 1, args.skill_max_tokens, skill_dp, temperature=0.0) + skills = [] + for seqs in sg_out: + if not seqs: + skills.append(('', '')) + continue + sresp = _clean_text(getattr(seqs[0], 'decoded', '') or '') + skills.append((_extract_skill(sresp) or '', sresp)) + ws_out = _run_samples(base_sampler, + [build_skill_solve_prompt(r['problem'], sk) for r, (sk, _) in zip(eval_records, skills)], + 1, args.max_tokens, base_dp, temperature=0.0) + recs = [] + for r, (sk, sresp), seqs in zip(eval_records, skills, ws_out): + roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + recs.append({ + 'record_type': 'eval_problem', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, + 'data_id': r.get('data_id', ''), 'problem': r['problem'], + 'reference_answer': r['reference_answer'], 'baseline_pass': r['_baseline_pass'], + 'skill': sk, 'skill_parseable': bool(sk), 'skill_chars': len(sk), + 'withskill_pred': roll['pred'], 'withskill_correct': roll['correct'], + 'withskill_terminated': roll['terminated'], 'withskill_stop_reason': roll['stop_reason'], + 'withskill_text': roll['text'], + }) + n = len(recs) + ws = (sum(1 for x in recs if x['withskill_correct']) / n) if n else 0.0 + base = (sum(x['baseline_pass'] for x in recs) / n) if n else 0.0 + fmt = (sum(1 for x in recs if x['skill_parseable']) / n) if n else 0.0 + term = (sum(1 for x in recs if x['withskill_terminated']) / n) if n else 0.0 + # 中文注释：难题子片救活率——baseline_pass==0 的子集里 with-skill 做对的比例。 + # 零成本（复用已算字段），是 buffer B 回路的目标量（见 skill_quality_analysis.md 第 14 节）。 + hard = [x for x in recs if not x['baseline_pass']] + hard_rescued = sum(1 for x in hard if x['withskill_correct']) + hard_rescue_rate = (hard_rescued / len(hard)) if hard else 0.0 + summary = {'record_type': 'eval_summary', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, + 'n': n, 'acc_mean1': ws, 'baseline_acc_mean1': base, 'lift_mean1': ws - base, + 'format_mean1': fmt, 'term_mean1': term, + 'hard_n': len(hard), 'hard_rescued': hard_rescued, 'hard_rescue_rate': hard_rescue_rate} + metrics = {'core/math/acc/mean@1': ws, 'core/math/baseline_acc/mean@1': base, + 'core/math/lift/mean@1': ws - base, 'core/math/format/mean@1': fmt, + 'core/math/term/mean@1': term, 'core/math/hard_rescue/mean@1': hard_rescue_rate} + return recs, summary, metrics + + +# =========================================================================== +# Section I — components, args, main +# =========================================================================== +def init_components(args): + r0, r1 = TRAIN_GPUS, TRAIN_GPUS + REF_GPUS + r2, r3 = r1 + SKILL_SAMPLER_GPUS, NUM_GPUS + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, lazy_collect=False, groups=[ + DeviceGroup(name='train', ranks=list(range(0, r0)), device_type='GPU'), + DeviceGroup(name='ref', ranks=list(range(r0, r1)), device_type='GPU'), + DeviceGroup(name='skill_sampler', ranks=list(range(r1, r2)), device_type='GPU'), + DeviceGroup(name='base_sampler', ranks=list(range(r2, r3)), device_type='GPU')]) + + train_mesh = DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_DP, fsdp_size=TRAIN_FSDP) + skill_model = TransformersModel(model_id=MODEL_ID, device_mesh=train_mesh, remote_group='train', + ddp_config={'find_unused_parameters': False}) + skill_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) + skill_model.set_template(Template, model_id=MODEL_ID, enable_thinking=False, + max_length=args.max_model_len, truncation_strategy='delete') + skill_model.set_processor(InputProcessor, padding_free=False) + skill_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon, beta=args.kl_beta) + skill_model.set_optimizer('AdamW', lr=args.lr) + skill_model.set_lr_scheduler('CosineWarmupScheduler', num_warmup_steps=10, + num_training_steps=args.max_train_rounds) + + ref_mesh = DeviceMesh.from_sizes(world_size=REF_GPUS, dp_size=REF_DP, fsdp_size=REF_FSDP) + ref_model = TransformersModel(model_id=MODEL_ID, device_mesh=ref_mesh, remote_group='ref', + ddp_config={'find_unused_parameters': False}) + ref_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) + ref_model.set_template(Template, model_id=MODEL_ID, enable_thinking=False, + max_length=args.max_model_len, truncation_strategy='delete') + ref_model.set_processor(InputProcessor, padding_free=False) + ref_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) + + def _sampler(group, world, enable_thinking): + s = vLLMSampler(model_id=MODEL_ID, + engine_args={'gpu_memory_utilization': GPU_MEM, + 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, + device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), + remote_group=group) + s.set_template(Template, model_id=MODEL_ID, enable_thinking=enable_thinking, max_length=args.max_model_len) + return s + + skill_sampler = _sampler('skill_sampler', SKILL_SAMPLER_GPUS, enable_thinking=False) + base_sampler = _LockedSampler(_sampler('base_sampler', BASE_SAMPLER_GPUS, enable_thinking=True)) + ckpt = CheckpointEngineManager(model=skill_model, sampler=skill_sampler) + return skill_model, ref_model, skill_sampler, base_sampler, ckpt, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS + + +def _build_args(): + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--dataset', choices=('aops', 'math'), default='aops') + p.add_argument('--n', type=int, default=2000, help='Problems loaded into the draw pool (0=all).') + p.add_argument('--exclude-data-ids', default='', + help='Comma-separated jsonl files whose data_id/problem keys are excluded.') + p.add_argument('--seed', type=int, default=42) + p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--eval-size', type=int, default=128, help='Fixed holdout size (0 disables).') + p.add_argument('--eval-every', type=int, default=5, help='Run holdout eval every N chunks.') + p.add_argument('--chunk-size', type=int, default=16) + p.add_argument('--n-skills', type=int, default=8) + p.add_argument('--skill-retries', type=int, default=2) + p.add_argument('--skill-gen-temperature', type=float, default=1.0) + p.add_argument('--skill-gen-top-p', type=float, default=1.0) + p.add_argument('--skill-gen-top-k', type=int, default=-1) + p.add_argument('--max-model-len', type=int, default=16384) + p.add_argument('--max-tokens', type=int, default=8192) + p.add_argument('--skill-max-tokens', type=int, default=4096) + p.add_argument('--len-budget', type=int, default=600, + help='Skill length budget (chars). Reward multiplied by min(1, budget/len).') + # --- buffer / distillation --- + p.add_argument('--distill-trigger', type=int, default=300, + help='Start draining buffer A into distillation once it reaches this many entries.') + p.add_argument('--distill-batch', type=int, default=64, + help='Entries distilled per iteration while buffer A is over --distill-trigger ' + '(incremental drain: bounds per-step latency instead of one big stall).') + p.add_argument('--sft-trigger', type=int, default=100, + help='Run one SFT pass + eval when buffer B reaches this many validated entries. ' + 'Kept low: the distill funnel (has-FAIL × valid-regen × pass@k) yields only ' + '~10-15%% of buffer A, so a high threshold would rarely fire the SFT loop.') + # Plan B validation: diversity lives in the SKILL side, the executor stays at the + # deployment (greedy) decoding口径. For each buffer-A problem we regenerate K distinct + # candidate skills (high temperature), run each through ONE greedy (T=0) executor solve, + # and accept the problem iff >= M distinct skills reach a terminated-correct solve. This + # validates "the problem admits several skills that work under greedy decoding" (matches + # eval口径) rather than "one skill passes m/k times under a high-temperature executor". + p.add_argument('--passatk-k', type=int, default=8, + help='Plan B: number of DISTINCT candidate skills regenerated per problem ' + '(skill-side diversity; executor stays greedy).') + p.add_argument('--passatk-skill-temp', type=float, default=1.0, + help='Skill-model temperature when regenerating the K candidate skills ' + '(needs >0 for diversity across candidates).') + p.add_argument('--passatk-skill-top-p', type=float, default=1.0, + help='Skill-model top-p when regenerating the K candidate skills.') + p.add_argument('--passatk-m', type=int, default=2, + help='Plan B: min number of DISTINCT candidate skills that must reach a ' + 'terminated-correct GREEDY solve to accept the problem into buffer B. ' + 'Lower than pass@k-over-one-skill (default 2): requiring m distinct ' + 'greedy-effective skills is already a strong, low-noise bar.') + p.add_argument('--sft-weight', type=float, default=0.5, + help='Advantage magnitude for SFT distillation samples (-w*logp + beta*KL).') + p.add_argument('--rubric-workers', type=int, default=16) + # --- GRPO --- + p.add_argument('--sft-batch-size', type=int, default=8) + p.add_argument('--ppo-mini-batch-size', type=int, default=0) + p.add_argument('--grpo-epsilon', type=float, default=0.2) + p.add_argument('--adv-clip', type=float, default=3.0) + p.add_argument('--kl-beta', type=float, default=0.001) + p.add_argument('--lr', type=float, default=6e-6) + p.add_argument('--max-train-rounds', type=int, default=1500) + p.add_argument('--save-rounds', type=int, default=200) + p.add_argument('--output-dir', default='./output/skill_v2') + p.add_argument('--cache-dir', default='') + p.add_argument('--no-cache', action='store_true') + p.add_argument('--swanlab-project', default='twinkle') + p.add_argument('--swanlab-exp', default='') + args = p.parse_args() + if args.sft_batch_size % TRAIN_DP != 0: + raise ValueError(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple of train dp ({TRAIN_DP})') + if args.chunk_size < 1: + raise ValueError('--chunk-size must be >= 1') + return args + + +def _write(handle, row): + handle.write(json.dumps(row, ensure_ascii=False) + '\n') + + +def _swan_metrics(summary, log): + # Lean metric set: each carries independent information. Dropped as redundant — + # 中文注释：删除冗余项（换算重复）：n_groups(≈chunk_size)、reward_std(池化,组内方差已够)、 + # skill_tokens_mean(与chars重复)、leak/n(=rate×n)、candidate_withskill(与问题级重复)、 + # term/withskill(=1-trunc)、train/n_steps(恒为1)。 + d = { + 'signal/zero_grad_frac': summary['zero_grad_frac'], + 'signal/reward_mean': summary['reward_mean'], + 'signal/group_reward_std_mean': summary['group_reward_std_mean'], + 'signal/n_train_samples': summary['n_train_samples'], + 'skill/parse_rate': summary['parse_rate'], 'skill/chars_mean': summary['skill_chars_mean'], + 'leak/rate': summary['leak_rate'], + } + if summary['n_groups'] > 0: + d.update({'acc/withskill_pass': summary['avg_withskill_pass'], + 'term/withskill_trunc_frac': summary['withskill_trunc_frac']}) + if log: + d['train/n_grpo'] = log['n_grpo'] + d['train/n_sft'] = log['n_sft'] + for k, v in (log.get('metric') or {}).items(): + if not _is_num(v): + continue + if k.startswith('learning rate'): + if 'group 1' in k: + d['train/lr'] = float(v) + else: + d[f'train/{k.replace(" ", "_")}'] = float(v) + return d + + +def main(): + args = _build_args() + records, eval_records = _load_records(args) + if len(records) < args.chunk_size: + raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded ({len(records)}); raise --n') + + os.makedirs(args.output_dir, exist_ok=True) + gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') + eval_path = os.path.join(args.output_dir, 'eval_records.jsonl') + sft_path = os.path.join(args.output_dir, 'skill_dataset.jsonl') + train_log_path = os.path.join(args.output_dir, 'train_log.jsonl') + buffer_a_path = os.path.join(args.output_dir, 'buffer_a.jsonl') + distill_path = os.path.join(args.output_dir, 'distill_records.jsonl') + + use_swan = swanlab is not None and os.environ.get('SWANLAB_MODE') != 'disabled' + if use_swan: + swanlab.init(project=args.swanlab_project, experiment_name=(args.swanlab_exp or None), + config={'model': MODEL_ID, 'dataset': args.dataset, 'n': len(records), + 'eval_n': len(eval_records), 'n_skills': args.n_skills, + 'len_budget': args.len_budget, 'distill_trigger': args.distill_trigger, + 'sft_trigger': args.sft_trigger, 'passatk_k': args.passatk_k, + 'passatk_m': args.passatk_m, 'passatk_skill_temp': args.passatk_skill_temp, + 'sft_weight': args.sft_weight, 'lr': args.lr}) + + skill_model, ref_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp = init_components(args) + checker = build_rubric_checker() + if checker is None: + sys.stderr.write('[v2] no LLM backup env -> buffer B distillation DISABLED (GRPO only)\n') + + cache_dir = args.cache_dir or os.path.join(args.output_dir, 'cache') + os.makedirs(cache_dir, exist_ok=True) + eval_base_cache = DiskCache(os.path.join(cache_dir, 'eval_baseline.jsonl'), not args.no_cache) + + cfg = {'record_type': 'config', 'model': MODEL_ID, 'dataset': args.dataset, + 'n': len(records), 'eval_n': len(eval_records), 'seed': args.seed, + 'n_skills': args.n_skills, 'len_budget': args.len_budget, + 'distill_trigger': args.distill_trigger, 'sft_trigger': args.sft_trigger, + 'passatk_k': args.passatk_k, 'passatk_m': args.passatk_m, + 'passatk_skill_temp': args.passatk_skill_temp, 'passatk_skill_top_p': args.passatk_skill_top_p, + 'sft_weight': args.sft_weight, + 'grpo_epsilon': args.grpo_epsilon, 'kl_beta': args.kl_beta, 'lr': args.lr, + 'rubric_check': bool(checker), 'max_train_rounds': args.max_train_rounds, + 'started': int(time.time())} + + hist_a: List[Dict[str, Any]] = [] # buffer A accumulator (in-memory + jsonl) + sft_queue: List[Dict[str, Any]] = [] # buffer B: validated SFT records awaiting an SFT pass + rounds = 0 # GRPO rounds only (gates --max-train-rounds + save cadence) + sft_rounds = 0 # SFT passes (separate: must NOT eat the GRPO round budget) + pool = ProblemPool(records, args.seed) + + # Background rubric pre-diagnosis (做法 B): the moment a failure trajectory lands in + # buffer A, fire its teacher-rubric call on a daemon thread pool. The API round-trip + # then overlaps with GRPO GPU work, so by the time --distill-trigger fires the + # diagnoses are usually already cached on each entry ('_rubric_diag'); distill_buffer + # only pays for the stragglers. Entries are dicts held by reference, so the worker + # writes the result straight onto the entry. + # 中文注释：失败轨迹一进 buffer A 就后台异步跑 rubric，API 等待藏进 GPU 训练时间； + # 到蒸馏时诊断多已缓存在条目上，distill_buffer 只补漏。 + prediag_pool = (ThreadPoolExecutor(max_workers=max(1, args.rubric_workers), + thread_name_prefix='rubric-prediag') + if checker else None) + + def _prediagnose(entry: Dict[str, Any]): + entry['_rubric_diag'] = _diagnose_entry(checker, entry) or '' + + with open(gen_path, 'w', encoding='utf-8') as gen_f, \ + open(eval_path, 'w', encoding='utf-8') as eval_f, \ + open(sft_path, 'w', encoding='utf-8') as sft_f, \ + open(train_log_path, 'w', encoding='utf-8') as tlog, \ + open(distill_path, 'w', encoding='utf-8') as distill_f, \ + open(buffer_a_path, 'w', encoding='utf-8') as buf_f: + for f in (gen_f, eval_f, sft_f, tlog, distill_f): + _write(f, cfg) + + def _do_eval(gstep): + recs, summary, metrics = run_greedy_eval( + base_sampler, skill_sampler, eval_records, gstep, rounds, base_dp, skill_dp, + args, eval_base_cache) + for rec in recs: + _write(eval_f, rec) + _write(eval_f, summary) + eval_f.flush() + if use_swan: + swanlab.log({f'eval/{k}': v for k, v in metrics.items()}, step=max(gstep, 0)) + sys.stderr.write( + f'[eval] g{gstep}: n={summary["n"]} acc={summary["baseline_acc_mean1"]:.3f}' + f'->{summary["acc_mean1"]:.3f} lift={summary["lift_mean1"]:+.3f} ' + f'hard_rescue={summary["hard_rescue_rate"]:.3f}({summary["hard_rescued"]}/{summary["hard_n"]}) ' + f'fmt={summary["format_mean1"]:.2f} rounds={rounds}\n') + + if eval_records: + _do_eval(-1) + + gstep = 0 + while rounds < args.max_train_rounds: + chunk = pool.draw(args.chunk_size) + full, summary, grpo, buffer_a = process_chunk( + base_sampler, skill_sampler, chunk, gstep, base_dp, skill_dp, args) + + # accumulate buffer A (only when a rubric checker exists to consume it; + # 中文注释：无 checker 时蒸馏永不触发，不累积以免内存无限增长) + if checker: + for e in buffer_a: + _write(buf_f, e) + prediag_pool.submit(_prediagnose, e) # 后台异步预诊断，不阻塞主循环 + buf_f.flush() + hist_a.extend(buffer_a) + + # GRPO train step (only when there is signal) + log = None + if grpo: + log = _train_step(skill_model, ref_model, ckpt, grpo, args) + rounds += 1 + log.update({'record_type': 'train_round', 'round': rounds, 'chunk': gstep, + 'epoch': pool.epoch, 'kind': 'grpo', 'ts': int(time.time())}) + _write(tlog, log) + tlog.flush() + if rounds % args.save_rounds == 0: + skill_model.save(f'skill-v2-{rounds}', output_dir=args.output_dir) + + summary['rounds_done'], summary['epoch'] = rounds, pool.epoch + summary['buffer_a_size'], summary['sft_queue_size'] = len(hist_a), len(sft_queue) + for rec in full: + _write(gen_f, rec) + _write(gen_f, summary) + gen_f.flush() + + sys.stderr.write( + f'[gen] e{pool.epoch} g{gstep}: n={summary["n"]} ' + f'clean={summary["n_candidates_parseable"]} 0grad={summary["zero_grad_frac"]:.2f} ' + f'R={summary["reward_mean"]:.2f}+-{summary["reward_std"]:.2f} ' + f'ws_acc={summary["avg_withskill_pass"]:.2f} chars={summary["skill_chars_mean"]:.0f} ' + f'bufA={len(hist_a)} bufB={len(sft_queue)} rounds={rounds}\n') + if use_swan: + m = _swan_metrics(summary, log) + m['buffer/a_size'] = float(len(hist_a)) + m['buffer/b_size'] = float(len(sft_queue)) + swanlab.log(m, step=gstep) + + # --- distillation: once buffer A fills, drain it INCREMENTALLY in bounded + # batches (--distill-batch) so a large buffer never stalls the loop for tens + # of minutes; each iteration processes one batch, interleaved with GRPO. + # 中文注释：增量分批蒸馏——buffer A 满后每轮只处理 --distill-batch 条，把一次性 + # 几十分钟阻塞摊成每轮几分钟小停顿；两段验证(见 distill_buffer)再砍验证算力。 + if checker and len(hist_a) >= args.distill_trigger: + batch = hist_a[:args.distill_batch] + hist_a = hist_a[args.distill_batch:] + new_sft, distill_recs = distill_buffer(batch, skill_sampler, base_sampler, checker, + skill_dp, base_dp, args) + for rec in distill_recs: # 逐 entry 审计记录：rubric_diag + 候选 skill + 贪心解 + stage + rec['chunk'] = gstep + _write(distill_f, rec) + distill_f.flush() + for rec in new_sft: + _write(sft_f, rec) + sft_f.flush() + sft_queue.extend(new_sft) + + # --- SFT trigger: buffer B full → one SFT pass + eval --- + did_eval = False + if len(sft_queue) >= args.sft_trigger: + sys.stderr.write(f'[sft] triggered at bufB={len(sft_queue)}\n') + sft_samples = [{**s, 'advantage': float(args.sft_weight)} for s in sft_queue] + sft_log = _train_step(skill_model, ref_model, ckpt, sft_samples, args) + sft_rounds += 1 # 中文注释：SFT 用独立计数，不占用 GRPO 的 rounds 配额/save 节奏 + sft_log.update({'record_type': 'train_round', 'round': rounds, 'sft_round': sft_rounds, + 'chunk': gstep, 'epoch': pool.epoch, 'kind': 'sft', 'ts': int(time.time())}) + _write(tlog, sft_log) + tlog.flush() + sft_queue = [] + skill_model.save(f'skill-v2-sft{sft_rounds}', output_dir=args.output_dir) # 大改动后落盘 + if eval_records: # 中文注释：SFT 后立即 eval，测灾难性遗忘/真提升（第 11.3/13.4 节） + _do_eval(gstep) + did_eval = True + + if eval_records and not did_eval and (gstep + 1) % args.eval_every == 0: + _do_eval(gstep) + gstep += 1 + + if prediag_pool is not None: + prediag_pool.shutdown(wait=False, cancel_futures=True) # 丢弃未完成的后台预诊断 + eval_base_cache.close() + skill_model.save('skill-v2-final', output_dir=args.output_dir) + sys.stderr.write(f'[v2] done: {rounds} rounds over {gstep} chunks / {pool.epoch} epochs\n') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/embedding/train_skill_v2.sh b/cookbook/exp/embedding/train_skill_v2.sh new file mode 100644 index 000000000..f6c2d91ea --- /dev/null +++ b/cookbook/exp/embedding/train_skill_v2.sh @@ -0,0 +1,59 @@ +#!/bin/bash +# train_skill_v2.sh — 简化 GRPO + buffer distill 训练启动脚本 +# 用法: bash cookbook/exp/embedding/train_skill_v2.sh +# +# 环境变量: +# LLM_BACKUP_API_KEY - rubric 诊断用的教师 API key（必须，否则 buffer B 蒸馏不可用） +# LLM_BACKUP_BASE_URL - 教师 API base URL +# LLM_BACKUP_MODEL - 教师模型 ID +# GEN_MODEL_ID - 训练 skill 模型 ID（默认 Qwen/Qwen3-4B） +# TRAIN_GPUS / REF_GPUS / SKILL_SAMPLER_GPUS / BASE_SAMPLER_GPUS — GPU 分配 + +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +# 默认输出目录 +OUTPUT_DIR="${OUTPUT_DIR:-./output/skill_v2}" + +# 去重/排斥数据（冷启动 SFT 数据避免重叠） +EXCLUDE="${EXCLUDE_DATA_IDS:-}" + +# 提前建目录：tee 需在 python 建目录前就能打开日志文件 +mkdir -p "${OUTPUT_DIR}" + +python3 "${SCRIPT_DIR}/train_skill_v2.py" \ + --dataset aops \ + --n 20000 \ + --numeric-only \ + --eval-size 200 \ + --eval-every 10 \ + --chunk-size 32 \ + --n-skills 8 \ + --skill-retries 2 \ + --skill-gen-temperature 1.0 \ + --skill-gen-top-p 1.0 \ + --skill-gen-top-k -1 \ + --max-model-len 16384 \ + --max-tokens 8192 \ + --skill-max-tokens 4096 \ + --len-budget 600 \ + --distill-trigger 150 \ + --distill-batch 64 \ + --sft-trigger 100 \ + --passatk-k 8 \ + --passatk-m 2 \ + --sft-weight 1.0 \ + --rubric-workers 16 \ + --sft-batch-size 8 \ + --ppo-mini-batch-size 0 \ + --grpo-epsilon 0.2 \ + --adv-clip 3.0 \ + --kl-beta 0.001 \ + --lr 1e-6 \ + --max-train-rounds 1500 \ + --save-rounds 200 \ + --output-dir "${OUTPUT_DIR}" \ + --swanlab-project twinkle \ + --swanlab-exp "skill_v2_$(date +%Y%m%d_%H%M%S)" \ + ${EXCLUDE:+--exclude-data-ids "${EXCLUDE}"} \ + "$@" 2>&1 | tee "${OUTPUT_DIR}/run.log" From cd862bb2c50b1501356a8f716e44403a87f4b8a6 Mon Sep 17 00:00:00 2001 From: root Date: Sun, 26 Jul 2026 18:05:15 +0800 Subject: [PATCH 24/60] wip --- cookbook/exp/legacy/ablation_all.sh | 84 + cookbook/exp/legacy/ablation_direct.sh | 12 + .../exp/legacy/ablation_rag_api_condenser.sh | 17 + .../legacy/ablation_rag_local_condenser.sh | 18 + cookbook/exp/legacy/ablation_rag_raw.sh | 15 + .../legacy/build_reflexion_coldstart_sft.py | 519 +++++ .../exp/legacy/build_reflexion_skill_data.py | 1123 +++++++++ .../exp/legacy/build_thinking_rag_index.py | 1159 ++++++++++ cookbook/exp/legacy/compare_math_levels.py | 91 + cookbook/exp/legacy/dataset_hard.py | 202 ++ cookbook/exp/legacy/dataset_index.py | 718 ++++++ cookbook/exp/legacy/dataset_think.py | 456 ++++ cookbook/exp/legacy/eval_dualline_math.py | 689 ++++++ cookbook/exp/legacy/eval_gpqa_rag.py | 1547 +++++++++++++ cookbook/exp/legacy/eval_math_by_level.sh | 59 + cookbook/exp/legacy/eval_rag_recall.py | 187 ++ cookbook/exp/legacy/eval_reflexion_skill.py | 762 +++++++ cookbook/exp/legacy/make_embedding_dataset.py | 758 ++++++ .../exp/legacy/train_embedding_full_ddp.py | 270 +++ cookbook/exp/legacy/train_reflexion_skill.py | 1990 ++++++++++++++++ cookbook/exp/legacy/train_reflexion_skill.sh | 64 + .../exp/legacy/train_reflexion_skill_old.py | 2022 +++++++++++++++++ .../exp/legacy/train_reflexion_skill_old.sh | 57 + .../legacy/train_reflexion_skill_replay.py | 114 + .../exp/legacy/train_reflexion_skill_rft.py | 1568 +++++++++++++ .../exp/legacy/train_reflexion_skill_rft.sh | 35 + .../exp/legacy/train_reflexion_skill_seam.py | 2022 +++++++++++++++++ cookbook/exp/legacy/train_skill_v2_ablate.sh | 35 + cookbook/exp/legacy/train_skill_v2_ablate3.sh | 62 + cookbook/exp/skill2lora/run_ablate12.sh | 109 + cookbook/exp/skill2lora/train_skill_v2.py | 1800 +++++++++++++++ cookbook/exp/skill2lora/train_skill_v2.sh | 69 + 32 files changed, 18633 insertions(+) create mode 100755 cookbook/exp/legacy/ablation_all.sh create mode 100755 cookbook/exp/legacy/ablation_direct.sh create mode 100755 cookbook/exp/legacy/ablation_rag_api_condenser.sh create mode 100755 cookbook/exp/legacy/ablation_rag_local_condenser.sh create mode 100755 cookbook/exp/legacy/ablation_rag_raw.sh create mode 100644 cookbook/exp/legacy/build_reflexion_coldstart_sft.py create mode 100644 cookbook/exp/legacy/build_reflexion_skill_data.py create mode 100644 cookbook/exp/legacy/build_thinking_rag_index.py create mode 100644 cookbook/exp/legacy/compare_math_levels.py create mode 100644 cookbook/exp/legacy/dataset_hard.py create mode 100644 cookbook/exp/legacy/dataset_index.py create mode 100644 cookbook/exp/legacy/dataset_think.py create mode 100644 cookbook/exp/legacy/eval_dualline_math.py create mode 100644 cookbook/exp/legacy/eval_gpqa_rag.py create mode 100755 cookbook/exp/legacy/eval_math_by_level.sh create mode 100644 cookbook/exp/legacy/eval_rag_recall.py create mode 100644 cookbook/exp/legacy/eval_reflexion_skill.py create mode 100644 cookbook/exp/legacy/make_embedding_dataset.py create mode 100644 cookbook/exp/legacy/train_embedding_full_ddp.py create mode 100644 cookbook/exp/legacy/train_reflexion_skill.py create mode 100755 cookbook/exp/legacy/train_reflexion_skill.sh create mode 100644 cookbook/exp/legacy/train_reflexion_skill_old.py create mode 100755 cookbook/exp/legacy/train_reflexion_skill_old.sh create mode 100644 cookbook/exp/legacy/train_reflexion_skill_replay.py create mode 100644 cookbook/exp/legacy/train_reflexion_skill_rft.py create mode 100644 cookbook/exp/legacy/train_reflexion_skill_rft.sh create mode 100644 cookbook/exp/legacy/train_reflexion_skill_seam.py create mode 100644 cookbook/exp/legacy/train_skill_v2_ablate.sh create mode 100755 cookbook/exp/legacy/train_skill_v2_ablate3.sh create mode 100644 cookbook/exp/skill2lora/run_ablate12.sh create mode 100644 cookbook/exp/skill2lora/train_skill_v2.py create mode 100644 cookbook/exp/skill2lora/train_skill_v2.sh diff --git a/cookbook/exp/legacy/ablation_all.sh b/cookbook/exp/legacy/ablation_all.sh new file mode 100755 index 000000000..2be4c7fc7 --- /dev/null +++ b/cookbook/exp/legacy/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/legacy/ablation_direct.sh b/cookbook/exp/legacy/ablation_direct.sh new file mode 100755 index 000000000..f5e5d1456 --- /dev/null +++ b/cookbook/exp/legacy/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/legacy/ablation_rag_api_condenser.sh b/cookbook/exp/legacy/ablation_rag_api_condenser.sh new file mode 100755 index 000000000..fb47718e6 --- /dev/null +++ b/cookbook/exp/legacy/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/legacy/ablation_rag_local_condenser.sh b/cookbook/exp/legacy/ablation_rag_local_condenser.sh new file mode 100755 index 000000000..defa863cd --- /dev/null +++ b/cookbook/exp/legacy/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/legacy/ablation_rag_raw.sh b/cookbook/exp/legacy/ablation_rag_raw.sh new file mode 100755 index 000000000..86e35bd1e --- /dev/null +++ b/cookbook/exp/legacy/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/legacy/build_reflexion_coldstart_sft.py b/cookbook/exp/legacy/build_reflexion_coldstart_sft.py new file mode 100644 index 000000000..548dfad58 --- /dev/null +++ b/cookbook/exp/legacy/build_reflexion_coldstart_sft.py @@ -0,0 +1,519 @@ +"""Build a cold-start SFT corpus for reflexion skill generation on AOPS. + +Pipeline: + AOPS problems -> frozen base greedy attempt -> strategy-level rubric API diagnosis + -> answer-free API skill target -> query-only SFT examples. + +This is intentionally offline: no GRPO, no actor training, and no skill-model rollout. +The API is treated as an external teacher, so both diagnosis and generated skill targets +are filtered if they reveal the target final answer. + +Example: + LLM_BACKUP_API_KEY=... LLM_BACKUP_BASE_URL=... LLM_BACKUP_MODEL=... \ + python cookbook/exp/embedding/build_reflexion_coldstart_sft.py \ + --dataset aops --n 10000 --output-dir ./output/reflexion_coldstart_sft --overwrite +""" +import argparse +import json +import os +import sys +import time +import urllib.error +import urllib.request +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Optional, Tuple + +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../..')) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.sampler import vLLMSampler + +from cookbook.exp.embedding.train_reflexion_skill import ( + MODEL_ID, + GPU_MEM, + SamplingParams, + DiskCache, + _MATH_RUBRIC, + _RUBRIC_VERSION, + _answer_leaked, + _clean_text, + _empty_roll, + _format_diagnosis, + _numeric_value, + _parse_seq, + _run_samples, + _skillgen_messages, + _load_excluded_records, + build_direct_prompt, + build_skill_solve_prompt, + build_rubric_checker, + extract_boxed, + load_problems, +) + +logger = get_logger() + +COLDSTART_SYSTEM = """\ +You are writing cold-start training targets for a math skill generator. You are given a +competition problem and an answer-free process diagnosis of a previous attempt. + +Write concise, reusable guidance that a query-only solver could use before solving this +problem or similar problems. Focus on route choice, structural observations, constraints, +validity checks, and length-control habits. + +Output exactly one XML-style block: + +Your reusable guidance here. + + +Rules: +- Do not mention the diagnosis, rubric, previous attempt, or API. +- Do not reveal the final answer, a corrected value/expression, an option label, or a + step-by-step solution. +- It is okay to name methods, checks, pitfalls, and local strategy directions. +- Keep it short and useful: 3-6 compact sentences or bullets. +""" + +COLDSTART_USER = """\ +Problem: +{problem} + +Answer-free process diagnosis: +{diagnosis} + +Now write the reusable skill guidance. +""" + +_SPECIAL_TOKEN_NOTE = 'process diagnosis leaked target answer' + + +def _api_config() -> Tuple[str, str, str]: + api_key = os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY') + base_url = os.environ.get('LLM_BACKUP_BASE_URL') or os.environ.get('OPENAI_BASE_URL') or 'https://api.openai.com/v1' + model = os.environ.get('LLM_BACKUP_MODEL') or os.environ.get('OPENAI_MODEL') or 'gpt-4o-mini' + if not api_key: + raise RuntimeError('Set LLM_BACKUP_API_KEY or OPENAI_API_KEY for cold-start API generation.') + return api_key, base_url.rstrip('/'), model + + +def _chat_complete(messages: List[Dict[str, str]], max_tokens: int, temperature: float, + retries: int = 3, timeout: int = 120) -> str: + api_key, base_url, model = _api_config() + url = f'{base_url}/chat/completions' + payload = { + 'model': model, + 'messages': messages, + 'temperature': temperature, + 'max_tokens': max_tokens, + } + data = json.dumps(payload).encode('utf-8') + headers = {'Content-Type': 'application/json', 'Authorization': f'Bearer {api_key}'} + last_err = None + for attempt in range(max(1, retries)): + req = urllib.request.Request(url, data=data, headers=headers, method='POST') + try: + with urllib.request.urlopen(req, timeout=timeout) as resp: + obj = json.loads(resp.read().decode('utf-8')) + return obj['choices'][0]['message']['content'] + except (urllib.error.HTTPError, urllib.error.URLError, TimeoutError, KeyError, json.JSONDecodeError) as exc: + last_err = exc + if attempt + 1 < max(1, retries): + time.sleep(min(8.0, 1.0 * (2 ** attempt))) + continue + raise RuntimeError(f'chat completion failed after {retries} attempts: {last_err}') + + +def _extract_skill_block(text: str) -> Optional[str]: + low = (text or '').lower() + end_think = low.rfind('') + answer = text[end_think + len(''):] if end_think >= 0 else (text or '') + low = answer.lower() + s = low.rfind('') + if s < 0: + return None + inner = s + len('') + e = low.find('', inner) + if e < 0: + return None + block = answer[inner:e].strip() + return block or None + + +def _skill_response(block: str) -> str: + return f'\n{block.strip()}\n' + + +def _baseline_rollout(base_sampler, problems: List[Dict[str, Any]], base_dp: int, + args: argparse.Namespace, cache: DiskCache) -> int: + todo = [r for r in problems if DiskCache.key_for(r['problem']) not in cache] + if todo: + outs = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], + 1, args.max_tokens, base_dp, temperature=0.0) + for r, seqs in zip(todo, outs): + roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + cache.put(DiskCache.key_for(r['problem']), roll) + for r in problems: + roll = cache.get(DiskCache.key_for(r['problem'])) + r['_init'] = [roll] + r['_baseline_pass'] = 1.0 if roll.get('correct') else 0.0 + r['_failed'] = not roll.get('correct') + return len(todo) + + +def _diagnose_one(checker, r: Dict[str, Any], args: argparse.Namespace) -> str: + init = r['_init'][0] + seg_text = init.get('text', '') + if init.get('stop_reason') == 'length' or not init.get('terminated'): + seg_text += ('\n\n[Process note: this attempt was cut off at the token budget ' + 'and never produced a final \\boxed{} answer.]') + seg = {'messages': [{'role': 'user', 'content': r['problem']}, + {'role': 'assistant', 'content': seg_text}]} + attempts = max(1, args.rubric_retries + 1) + for attempt in range(attempts): + try: + return _format_diagnosis(checker.diagnose(seg, query=r['problem'])) + except Exception as exc: + if attempt + 1 < attempts: + logger.warning(f'[rubric] diagnose error: {exc}; retry {attempt + 1}/{args.rubric_retries}') + time.sleep(min(4.0, 0.5 * (2 ** attempt))) + else: + logger.warning(f'[rubric] diagnose failed: {exc}') + return '' + + +def _diagnose_batch(checker, rows: List[Dict[str, Any]], args: argparse.Namespace, + cache: DiskCache) -> int: + pending = [] + for r in rows: + init = r['_init'][0] + term = 'L' if (init.get('stop_reason') == 'length' or not init.get('terminated')) else 'T' + key = DiskCache.key_for(r['problem'], init.get('text', ''), _RUBRIC_VERSION, term) + if key in cache: + r['_rubric_diag'] = cache.get(key) + else: + pending.append((r, key)) + if not pending: + return 0 + + def run(item): + r, key = item + diag = _diagnose_one(checker, r, args) + return r, key, diag + + workers = max(1, min(args.rubric_workers, len(pending))) + fresh = 0 + with ThreadPoolExecutor(max_workers=workers) as ex: + for r, key, diag in ex.map(run, pending): + r['_rubric_diag'] = diag or '' + if diag: + cache.put(key, diag) + fresh += 1 + return fresh + + +def _target_key(problem: str, diagnosis: str, sample_idx: int) -> str: + return DiskCache.key_for('coldstart_skill_v2', str(sample_idx), problem, diagnosis) + + +def _generate_skill_targets(r: Dict[str, Any], args: argparse.Namespace, + cache: DiskCache) -> List[Dict[str, Any]]: + out = [] + messages = [ + {'role': 'system', 'content': COLDSTART_SYSTEM}, + {'role': 'user', 'content': COLDSTART_USER.format(problem=r['problem'], diagnosis=r.get('_rubric_diag', ''))}, + ] + for sample_idx in range(max(1, int(args.api_samples))): + key = _target_key(r['problem'], r.get('_rubric_diag', ''), sample_idx) + if key in cache: + resp = cache.get(key) + else: + resp = _chat_complete(messages, max_tokens=args.api_max_tokens, + temperature=args.api_temperature, retries=args.api_retries, + timeout=args.api_timeout) + cache.put(key, resp) + block = _extract_skill_block(resp) or '' + leaked = _answer_leaked(resp + '\n' + block, r['reference_answer']) + out.append({'sample_idx': sample_idx, 'raw_response': resp, + 'skills': block, 'skill_leak': leaked}) + return out + + +def _sft_messages(problem: str, response: str) -> List[Dict[str, str]]: + msgs = _skillgen_messages(problem, 'B', '') + return msgs + [{'role': 'assistant', 'content': response}] + + +def _init_base_sampler(args: argparse.Namespace): + twinkle.initialize(mode='ray', nproc_per_node=args.base_gpus, lazy_collect=False, + groups=[DeviceGroup(name='base_sampler', ranks=list(range(args.base_gpus)), device_type='GPU')]) + sampler = vLLMSampler( + model_id=MODEL_ID, + engine_args={'gpu_memory_utilization': GPU_MEM, + 'max_model_len': args.max_model_len, + 'tensor_parallel_size': 1}, + device_mesh=DeviceMesh.from_sizes(world_size=args.base_gpus, dp_size=args.base_gpus), + remote_group='base_sampler') + sampler.set_template('Template', model_id=MODEL_ID, enable_thinking=True, + max_length=args.max_model_len) + return sampler, args.base_gpus + + +def _select_records(args: argparse.Namespace) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + load_n = 0 if args.numeric_only or args.eval_size > 0 else max(args.n, args.target_size + args.eval_size) + records = load_problems(args.dataset, load_n, args.seed) + raw_n = len(records) + if args.numeric_only: + records = [{**r, 'reference_answer': v} + for r, v in ((r, _numeric_value(r.get('reference_answer'))) for r in records) + if v is not None] + import numpy as np + np.random.RandomState(args.seed).shuffle(records) + exclude_ids, exclude_problems = _load_excluded_records(getattr(args, 'exclude_data_ids', '')) + excluded = 0 + if exclude_ids or exclude_problems: + before = len(records) + records = [r for r in records + if str(r.get('data_id', '')) not in exclude_ids + and str(r.get('problem', '')).strip() not in exclude_problems] + excluded = before - len(records) + eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 + pool = [dict(r) for r in records[eval_n:]] + pool_offset = max(0, int(getattr(args, 'pool_offset', 0) or 0)) + if pool_offset: + if pool_offset >= len(pool): + raise ValueError(f'--pool-offset ({pool_offset}) leaves no cold-start records from pool size {len(pool)}') + pool = pool[pool_offset:] + n = min(args.n, len(pool)) if args.n > 0 else min(len(pool), max(args.target_size * 2, args.target_size + 512)) + stats = {'raw_loaded': raw_n, 'numeric_dropped': raw_n - len(records) - excluded, + 'excluded_records': excluded, 'eval_size': eval_n, + 'pool_offset': pool_offset, 'pool_selected': n} + return pool[:n], stats + + +def _build_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--dataset', choices=('aops', 'math'), default='aops') + p.add_argument('--target-size', type=int, default=10000, help='Number of accepted SFT examples to write.') + p.add_argument('--n', type=int, default=0, help='Raw train-pool size after eval split; 0 auto-selects.') + p.add_argument('--pool-offset', type=int, default=0, + help='Skip this many shuffled non-eval records before building the cold-start pool.') + p.add_argument('--exclude-data-ids', default='', + help='Comma-separated jsonl files whose data_id/problem keys are excluded, ' + 'useful for building non-overlapping shards.') + p.add_argument('--eval-size', type=int, default=128) + p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--seed', type=int, default=42) + p.add_argument('--output-dir', default='./output/reflexion_coldstart_sft') + p.add_argument('--cache-dir', default='') + p.add_argument('--overwrite', action='store_true') + p.add_argument('--no-cache', action='store_true') + p.add_argument('--chunk-size', type=int, default=64) + p.add_argument('--base-gpus', type=int, default=int(os.environ.get('BASE_GPUS', 4))) + p.add_argument('--max-model-len', type=int, default=16384) + p.add_argument('--max-tokens', type=int, default=8192) + p.add_argument('--rubric-workers', type=int, default=16) + p.add_argument('--rubric-retries', type=int, default=2) + p.add_argument('--api-workers', type=int, default=16) + p.add_argument('--api-samples', type=int, default=4, + help='API skill targets sampled per problem before executor verification.') + p.add_argument('--verify-targets', action=argparse.BooleanOptionalAction, default=True, + help='Run frozen base executor with each API skill target and keep a successful one.') + p.add_argument('--keep-unverified-targets', action='store_true', + help='If all executor checks fail, keep the first clean target anyway. Default skips it.') + p.add_argument('--api-retries', type=int, default=3) + p.add_argument('--api-timeout', type=int, default=120) + p.add_argument('--api-max-tokens', type=int, default=768) + p.add_argument('--api-temperature', type=float, default=0.2) + p.add_argument('--require-fail', action=argparse.BooleanOptionalAction, default=True, + help='Only keep API diagnoses containing [FAIL]. Use --no-require-fail to keep OK diagnoses too.') + return p.parse_args() + + +def _write(f, row: Dict[str, Any]) -> None: + f.write(json.dumps(row, ensure_ascii=False) + '\n') + + +def main() -> None: + args = _build_args() + if args.target_size <= 0: + raise ValueError('--target-size must be positive') + records, data_stats = _select_records(args) + if not records: + raise ValueError('no records selected') + + os.makedirs(args.output_dir, exist_ok=True) + sft_path = os.path.join(args.output_dir, 'coldstart_sft.jsonl') + rec_path = os.path.join(args.output_dir, 'coldstart_records.jsonl') + for path in (sft_path, rec_path): + if os.path.exists(path) and not args.overwrite: + raise FileExistsError(f'{path} exists; pass --overwrite') + + checker = build_rubric_checker() + if checker is None: + raise RuntimeError('No rubric checker available; set LLM_BACKUP_API_KEY/BASE_URL or OPENAI_API_KEY.') + _api_config() + base_sampler, base_dp = _init_base_sampler(args) + + cache_dir = args.cache_dir or os.path.join(args.output_dir, 'cache') + os.makedirs(cache_dir, exist_ok=True) + use_cache = not args.no_cache + base_cache = DiskCache(os.path.join(cache_dir, 'baseline.jsonl'), use_cache) + rubric_cache = DiskCache(os.path.join(cache_dir, 'rubric.jsonl'), use_cache) + skill_cache = DiskCache(os.path.join(cache_dir, 'api_skill.jsonl'), use_cache) + + cfg = { + 'record_type': 'config', 'mode': 'coldstart_sft_build', 'dataset': args.dataset, + 'target_size': args.target_size, 'selected_records': len(records), 'seed': args.seed, + 'pool_offset': args.pool_offset, 'exclude_data_ids': args.exclude_data_ids, + 'numeric_only': args.numeric_only, **data_stats, + 'rubric_version': _RUBRIC_VERSION, 'rubric_check': f'fixed_math_{len(_MATH_RUBRIC)}crit', + 'api_model': os.environ.get('LLM_BACKUP_MODEL') or os.environ.get('OPENAI_MODEL') or 'gpt-4o-mini', + 'api_samples': args.api_samples, 'verify_targets': args.verify_targets, + 'keep_unverified_targets': args.keep_unverified_targets, + 'require_fail': args.require_fail, 'started': int(time.time()), + } + + accepted = 0 + skipped_no_diag = skipped_no_fail = skipped_api_leak = 0 + skipped_no_skill = skipped_skill_leak = skipped_executor_fail = 0 + processed = 0 + with open(sft_path, 'w', encoding='utf-8') as sft_f, open(rec_path, 'w', encoding='utf-8') as rec_f: + _write(rec_f, cfg) + for start in range(0, len(records), args.chunk_size): + if accepted >= args.target_size: + break + chunk = [dict(r) for r in records[start:start + args.chunk_size]] + _baseline_rollout(base_sampler, chunk, base_dp, args, base_cache) + _diagnose_batch(checker, chunk, args, rubric_cache) + + def gen_one(r: Dict[str, Any]): + return r, _generate_skill_targets(r, args, skill_cache) + + candidates = [] + for r in chunk: + processed += 1 + diag = r.get('_rubric_diag', '') or '' + if not diag: + skipped_no_diag += 1 + continue + if args.require_fail and '[FAIL]' not in diag: + skipped_no_fail += 1 + continue + if _answer_leaked(diag, r['reference_answer']): + skipped_api_leak += 1 + continue + candidates.append(r) + + generated = [] + workers = max(1, min(args.api_workers, len(candidates))) + if candidates: + with ThreadPoolExecutor(max_workers=workers) as ex: + for r, targets in ex.map(gen_one, candidates): + for target in targets: + skills = target.get('skills', '') + if not skills: + skipped_no_skill += 1 + continue + if target.get('skill_leak'): + skipped_skill_leak += 1 + continue + target['r'] = r + target['response'] = _skill_response(skills) + generated.append(target) + + selected = [] + selected_keys = set() + if generated and args.verify_targets: + verify_prompts = [build_skill_solve_prompt(g['r']['problem'], g['skills']) for g in generated] + verify_outs = _run_samples(base_sampler, verify_prompts, 1, args.max_tokens, + base_dp, temperature=0.0) + attempted_keys = set() + for g, seqs in zip(generated, verify_outs): + r = g['r'] + key = r.get('data_id') or r['problem'] + attempted_keys.add(key) + roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + g['target_roll'] = roll + if key not in selected_keys and roll.get('correct') and roll.get('terminated'): + g['executor_verified'] = True + selected.append(g) + selected_keys.add(key) + if args.keep_unverified_targets: + for g in generated: + r = g['r'] + key = r.get('data_id') or r['problem'] + if key not in selected_keys: + g['executor_verified'] = False + g.setdefault('target_roll', {}) + selected.append(g) + selected_keys.add(key) + skipped_executor_fail += len(attempted_keys - selected_keys) + elif generated: + for g in generated: + r = g['r'] + key = r.get('data_id') or r['problem'] + if key not in selected_keys: + g['executor_verified'] = False + selected.append(g) + selected_keys.add(key) + + for g in selected: + if accepted >= args.target_size: + break + r = g['r'] + response = g['response'] + messages = _sft_messages(r['problem'], response) + sft_row = { + 'messages': messages, + 'user_data': {'key_rounds': [len(messages) - 1]}, + 'data_id': r.get('data_id'), + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'skills': g['skills'], 'response': response, + 'view': 'B', 'sft': True, 'source': 'api_coldstart', + 'api_sample_idx': g.get('sample_idx'), + 'executor_verified': g.get('executor_verified', False), + 'baseline_correct': r['_init'][0].get('correct'), + 'baseline_terminated': r['_init'][0].get('terminated'), + 'baseline_stop_reason': r['_init'][0].get('stop_reason'), + 'target_correct': (g.get('target_roll') or {}).get('correct'), + 'target_terminated': (g.get('target_roll') or {}).get('terminated'), + 'target_stop_reason': (g.get('target_roll') or {}).get('stop_reason'), + 'diagnosis': r.get('_rubric_diag', ''), + } + audit = { + 'record_type': 'coldstart_problem', 'accepted': True, + 'data_id': r.get('data_id'), + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'baseline': r['_init'][0], 'diagnosis': r.get('_rubric_diag', ''), + 'raw_skill_response': g.get('raw_response'), 'skills': g['skills'], + 'api_sample_idx': g.get('sample_idx'), + 'executor_verified': g.get('executor_verified', False), + 'target_roll': g.get('target_roll'), + } + _write(sft_f, sft_row) + _write(rec_f, audit) + accepted += 1 + sys.stderr.write( + f'[coldstart] processed={processed} accepted={accepted}/{args.target_size} ' + f'skip(no_diag={skipped_no_diag}, no_fail={skipped_no_fail}, api_leak={skipped_api_leak}, ' + f'no_skill={skipped_no_skill}, skill_leak={skipped_skill_leak}, ' + f'executor_fail={skipped_executor_fail})\n') + sft_f.flush(); rec_f.flush() + + summary = { + 'record_type': 'summary', 'processed': processed, 'accepted': accepted, + 'skipped_no_diag': skipped_no_diag, 'skipped_no_fail': skipped_no_fail, + 'skipped_api_leak': skipped_api_leak, 'skipped_no_skill': skipped_no_skill, + 'skipped_skill_leak': skipped_skill_leak, + 'skipped_executor_fail': skipped_executor_fail, 'finished': int(time.time()), + } + with open(rec_path, 'a', encoding='utf-8') as rec_f: + _write(rec_f, summary) + sys.stderr.write(f'[coldstart] wrote {accepted} SFT rows to {sft_path}\n') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/legacy/build_reflexion_skill_data.py b/cookbook/exp/legacy/build_reflexion_skill_data.py new file mode 100644 index 000000000..59fe2382c --- /dev/null +++ b/cookbook/exp/legacy/build_reflexion_skill_data.py @@ -0,0 +1,1123 @@ +"""Offline builder for reflexion skill RFT data (self-contained, cached). + +Runs the SAME pipeline as the online trainer -- base greedy solve -> rubric +process-check (view A) -> skill-gen -> leak filter -> with-skill greedy pass -> +group-relative GRPO advantage -- but never updates the skill model. It emits +``skill_dataset.jsonl`` (trainer-schema training records), ``gen_records.jsonl`` +(full per-problem traces) and ``eval_holdout.jsonl`` (the fixed holdout). + +The expensive base rollouts and rubric diagnoses are cached to disk (one jsonl +each, keyed by an md5 of their inputs) so a re-run skips them entirely. + +8 GPUs: ranks 0-3 skill_sampler (vLLM tp1 dp4), ranks 4-7 base_sampler. Leak / +rubric use the backup teacher API: set LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / +LLM_BACKUP_MODEL. + +Launch: + LLM_BACKUP_API_KEY=... python cookbook/exp/embedding/build_reflexion_skill_data.py \ + --total-problems 3200 --base-success-frac 0.3 +""" +import argparse +import copy +import hashlib +import json +import math +import os +import re +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import SamplingParams +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.sampler import vLLMSampler +from twinkle_agentic.verifier import LeakVerifier, RubricVerifier +from twinkle_agentic.verifier.rubric_verifier import RubricItem + +logger = get_logger() + +MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') +GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.8)) +GEN_TEMPERATURE = float(os.environ.get('GEN_TEMPERATURE', 0.6)) +GEN_TOP_P = float(os.environ.get('GEN_TOP_P', 0.95)) +AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') +MATH_DATASET_ID = os.environ.get('MATH_DATASET_ID', 'modelscope/competition_math') + + +# =========================================================================== +# Block A -- boxed extraction + answer grading +# =========================================================================== +_BOXED_RE = re.compile(r'\\boxed\s*\{') + + +def extract_boxed(text: str) -> Optional[str]: + """Last ``\\boxed{...}`` content, brace-balanced.""" + if not text: + return None + last = None + for m in _BOXED_RE.finditer(text): + depth, i = 1, m.end() + while i < len(text) and depth > 0: + depth += (text[i] == '{') - (text[i] == '}') + i += 1 + if depth == 0: + last = text[m.end():i - 1].strip() + return last + + +def normalize_answer(ans: str) -> str: + if not ans: + return '' + s = str(ans).strip() + m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) + if m: + return m.group(1) + s = s.strip('$').strip().replace('−', '-') + 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').replace(r'\tfrac', r'\frac') + s = s.replace(r'\,', '').replace(r'\;', '').replace(r'\!', '') + s = re.sub(r'\\(?:quad|qquad|\s)', '', s) + s = re.sub(r'\s+', '', s) + s = re.sub(r'\\sqrt([0-9A-Za-z])', r'\\sqrt{\1}', s) + s = re.sub(r'\\frac\{([^{}]+)\}([^{}\\])', r'\\frac{\1}{\2}', s) + s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) + s = re.sub(r'\{[a-zA-Z]+\}$', '', s) + s = re.sub(r'\^\\circ|\^\{\\circ\}|°|\\circ', 'deg', s) + s = s.replace(r'\minus{}', '-').replace(r'\minus', '-') + + def _frac_to_slash(mt): + text = mt.group(0) + pos = text.index('{') + 1 + depth, num_start = 1, pos + while depth > 0: + depth += (text[pos] == '{') - (text[pos] == '}') + pos += 1 + numer = text[num_start:pos - 1] + pos += 1 + den_start, depth = pos, 1 + while depth > 0: + depth += (text[pos] == '{') - (text[pos] == '}') + pos += 1 + return f'({numer})/({text[den_start:pos - 1]})' + + s = re.sub(r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', _frac_to_slash, s) + s = re.sub(r'(? bool: + try: + va, vb = float(a.replace('(', '').replace(')', '')), 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) + return va is not None and vb is not None and abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) + + +_MCQ_REF_RE = re.compile( + r'^\\(?:textbf|text|mathrm|mathbf)\{\(?([A-E])\)?\s*\}\s*(.+)$' + r'|^\(?([A-E])\)\s+(.+)$') +_VAR_PREFIX_RE = re.compile(r'^(?:[A-Za-z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)$', re.DOTALL) + + +def _split_mcq(ans: str): + s = ans.strip() + m = _MCQ_REF_RE.match(s) + if m: + return (m.group(1) or m.group(3)), ((m.group(2) or m.group(4) or '').strip() or None) + bl = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) + return (bl.group(1), None) if bl else (None, s or None) + + +def _strip_var_prefix(ans: str) -> str: + m = _VAR_PREFIX_RE.match((ans or '').strip()) + return m.group(1).strip() if m else (ans or '') + + +def _normalize_tuple(ans: str) -> str: + return re.sub(r'[\s()\[\]{}\\]', '', ans or '') + + +def _math_verify_equal(predicted: str, reference: str) -> bool: + try: + from math_verify import parse, verify + gold = parse(r'\boxed{' + reference + '}', parsing_timeout=1) + pred = parse(r'\boxed{' + predicted + '}', parsing_timeout=1) + return bool(gold and pred and verify(gold, pred, timeout_seconds=1)) + except Exception: + return False + + +def answers_match(predicted: str, reference: str) -> bool: + if not predicted or not reference: + return False + norm_p, norm_r = normalize_answer(predicted), normalize_answer(reference) + if norm_p == norm_r or norm_p.lower() == norm_r.lower() or _try_numeric_equal(norm_p, norm_r): + return True + stripped_p = normalize_answer(_strip_var_prefix(predicted)) + stripped_r = normalize_answer(_strip_var_prefix(reference)) + if stripped_p and stripped_r: + if (stripped_p == stripped_r or stripped_p.lower() == stripped_r.lower() + or _try_numeric_equal(stripped_p, stripped_r)): + return True + p_letter, p_value = _split_mcq(predicted) + r_letter, r_value = _split_mcq(reference) + if p_letter and r_letter and p_letter == r_letter: + return True + if (p_letter and p_value is None) and (r_letter and r_value is None): + return False + 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 and (p_val == r_val or p_val.lower() == r_val.lower() or _try_numeric_equal(p_val, r_val)): + return True + for left, right in ((norm_p, norm_r), (stripped_p, stripped_r)): + tl, tr = _normalize_tuple(left), _normalize_tuple(right) + if ',' in tl and tl == tr: + return True + if '=' in norm_r and '=' not in norm_p: + if any(part and (part == norm_p or _try_numeric_equal(part, norm_p)) for part in norm_r.split('=')): + return True + if '=' in norm_p and '=' not in norm_r: + if any(part and (part == norm_r or _try_numeric_equal(part, norm_r)) for part in norm_p.split('=')): + return True + return _math_verify_equal(stripped_p or norm_p, stripped_r or norm_r) + + +# =========================================================================== +# Block B -- prompts, skill parsing, batched sampling +# =========================================================================== +DIRECT_SYSTEM = ( + 'You are an expert competition mathematician. Solve the following problem ' + 'step by step. Provide your final answer inside \\boxed{}.') + +# Appended to solve turns: box BOTH the letter and value of an MCQ so the model +# never loops deciding which form to box. +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.') + +_SKILL_SOLVE_PREFIX = ( + DIRECT_SYSTEM + '\n\n' + 'Before you start, keep these reminders in mind to avoid common mistakes on this ' + 'type of problem:\n') +_SKILL_SOLVE_SUFFIX = '\nApply them where relevant, but rely on your own reasoning to reach the answer.' + + +def build_direct_prompt(problem: str) -> Dict[str, Any]: + return {'messages': [{'role': 'system', 'content': DIRECT_SYSTEM}, + {'role': 'user', 'content': problem}]} + + +def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: + # Concatenation (not .format): DIRECT_SYSTEM/skill contain literal braces. + return {'messages': [ + {'role': 'system', 'content': _SKILL_SOLVE_PREFIX + skill + _SKILL_SOLVE_SUFFIX}, + {'role': 'user', 'content': problem + MCQ_INSTRUCTION}]} + + +# -- skill-gen prompts (view A: problem + rubric findings; view B: query only) -- +SKILL_GEN_SYSTEM = ( + 'You are a mathematics coach. You are shown a competition problem together with an ' + 'automated process-check of an earlier solver attempt at it -- which solution ' + 'criteria the attempt passed or failed, and suggested fixes for the failures. You ' + 'do NOT see the attempt itself, only this check. Treat the check as privileged ' + 'training scaffolding: study it together with the problem, identify the ' + 'problem-visible features that make each useful flagged failure relevant, then ' + 'rephrase those lessons as self-contained reusable skills. The goal is not to ' + 'continue from the check, cite it, or hide it silently; the goal is to turn it into ' + 'a problem-triggered reasoning pattern a query-only solver could reproduce later.\n\n' + 'Good skills name the observable trigger, the method worth reaching for, the ' + 'pitfall to watch, and a quick verification habit. Prefer formulations like ' + '"When a configuration has ...", "Before setting up ...", or "Check whether ..." ' + 'over references to the process-check, failed criteria, or the earlier attempt. ' + "These tips are advisory: they will be placed in a solver's system prompt as gentle " + 'reminders before it works through a SIMILAR problem on its own, without seeing ' + 'this process-check. So keep them general and transferable rather than a worked ' + 'solution to this exact problem, and do not state its specific intermediate values ' + 'or final answer. Think briefly first, then give your tips as a markdown bullet ' + 'list wrapped in and , like the example below.') + +SKILL_GEN_SYSTEM_Q = ( + 'You are a mathematics coach. You are shown ONE competition problem and nothing ' + 'else — no solution and no attempt. Think about what approach this KIND of problem ' + 'calls for and where solvers tend to slip, then distil a few reusable tips.\n\n' + "These tips are advisory: they will be placed in a solver's system prompt as gentle " + 'reminders before it works through a SIMILAR problem on its own. So keep them ' + 'general and transferable — the method worth reaching for, the pitfall to watch and ' + 'a quick check, and the discipline to settle on a final answer — rather than a ' + 'worked solution to this exact problem, and without stating its specific ' + 'intermediate values or its final answer. Think briefly first, then give your tips ' + 'as a markdown bullet list wrapped in and , like the example below.') + +SKILL_GEN_USER_Q = ( + 'Problem:\n{problem}\n\n' + 'Now reason about this TYPE of problem, then output the skills bullet list.') + +SKILL_GEN_USER_RUBRIC = ( + 'Problem:\n{problem}\n\n' + 'Process check of an earlier attempt (automated rubric verifier -- treat as ' + 'evidence, not gospel; PASS/FAIL per criterion with suggested fixes for failures):\n' + '{diagnosis}\n\n' + 'Now output a self-contained skills bullet list. Each bullet should still be useful ' + 'if the process check were removed: connect any useful flagged failure to ' + 'problem-visible features, general methods, and quick checks rather than citing the ' + 'rubric or the earlier attempt.') + +_EX_PROBLEM = 'Simplify $\\sqrt{72} + \\sqrt{18}$ and give the result.' +_EX_SKILLS = ( + '\n' + '- Rewrite each square root by factoring its radicand into a perfect square times ' + 'a remainder, then move the perfect-square factor outside.\n' + '- Avoid the classic trap $\\sqrt{a}+\\sqrt{b}\\ne\\sqrt{a+b}$; only combine terms ' + 'sharing the same simplest radical, and sanity-check by estimating each root.\n' + '- Procedure: simplify every radical, group like radical terms, add their ' + 'coefficients, then reduce to simplest form.\n' + '- Once the expression is in simplest form, commit to that single result as the ' + 'final answer rather than re-checking indefinitely.\n' + '') + + +def _skillgen_messages(problem: str, view: str, diagnosis: str) -> List[Dict[str, Any]]: + """Single source of truth for the skill-gen prompt. View A with a localisable + failure uses problem + rubric findings; view B -- or a view-A problem whose rubric + flagged nothing (no ``[FAIL]``) -- degrades to the query-only prompt.""" + if view == 'B' or '[FAIL]' not in (diagnosis or ''): + return [{'role': 'system', 'content': SKILL_GEN_SYSTEM_Q}, + {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=_EX_PROBLEM)}, + {'role': 'assistant', 'content': _EX_SKILLS}, + {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=problem)}] + return [{'role': 'system', 'content': SKILL_GEN_SYSTEM}, + {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=_EX_PROBLEM)}, + {'role': 'assistant', 'content': _EX_SKILLS}, + {'role': 'user', 'content': SKILL_GEN_USER_RUBRIC.format(problem=problem, diagnosis=diagnosis)}] + + +def _assign_view(problem: str, args: argparse.Namespace) -> str: + h = int(hashlib.md5(f'{args.seed}:{problem}'.encode('utf-8')).hexdigest(), 16) + return 'B' if (h % 100000) / 100000.0 < args.view_b_frac else 'A' + + +def _view_prompt(r: Dict[str, Any]) -> Dict[str, Any]: + return {'messages': _skillgen_messages(r['problem'], r['_view'], r.get('_rubric_diag', ''))} + + +_BULLET_RE = re.compile(r'(?m)^\s*(?:[-*]|\d+[.)])\s') +_META_RE = re.compile( + r'\b(the student|the solver|the attempt|this attempt|the trace|the response|' + r'in the (?:attempt|trace|response|solution)|as (?:shown|seen|noted) above|' + r'the (?:above|previous|earlier)|my (?:reasoning|analysis)|i (?:think|need|will))\b', + re.IGNORECASE) +_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') + + +def _clean_text(decoded: Optional[str]) -> str: + return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() + + +def _is_clean_block(block: str) -> bool: + """Pure bullet list (every non-empty line a bullet) with no meta/trajectory ref.""" + lines = [ln for ln in (l.strip() for l in block.splitlines()) if ln] + if not lines or not all(_BULLET_RE.match(ln) for ln in lines): + return False + return _META_RE.search(block) is None + + +def _extract_skills_block(text: str) -> Optional[str]: + """Clean ``...`` block, or None. Requires ```` (skill-gen + runs thinking ON); reads only the answer after the last one, so a mid-reasoning draft + or a demo echo can never be mistaken for the answer.""" + low = text.lower() + end_think = low.rfind('') + if end_think < 0: + return None + answer = text[end_think + len(''):] + low_a = answer.lower() + s = low_a.find('') + if s < 0: + return None + inner = s + len('') + e = low_a.find('', inner) + block = (answer[inner:e] if e >= 0 else answer[inner:]).strip() + block = re.sub(r'', '', block, flags=re.IGNORECASE).strip() + return block if _is_clean_block(block) else None + + +def _parse_seq(seq, gold: str) -> Dict[str, Any]: + """Grade one sampled sequence into a rollout record.""" + text = _clean_text(getattr(seq, 'decoded', '') or '') + pred = extract_boxed(text) + correct = bool(pred) and answers_match(pred, gold) + terminated = getattr(seq, 'stop_reason', None) != 'length' + return {'pred': pred, 'correct': correct, 'terminated': terminated, + 'passed': bool(correct and terminated), + 'stop_reason': getattr(seq, 'stop_reason', None), + 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} + + +def _run_samples(sampler, prompts: List[Any], num_samples: int, max_tokens: int, + gen_dp: int, temperature: Optional[float] = None, + top_p: Optional[float] = None, top_k: Optional[int] = None) -> List[List[Any]]: + """One batched sampler call -> per-prompt list of raw sequences. vLLM dp needs + batch len >= dp, so pad the tail and slice back.""" + if not prompts: + return [] + params = SamplingParams( + max_tokens=max_tokens, + temperature=GEN_TEMPERATURE if temperature is None else temperature, + top_p=GEN_TOP_P if top_p is None else top_p, + num_samples=num_samples, **({} if top_k is None else {'top_k': top_k})) + padded = prompts + if gen_dp > 1 and 0 < len(prompts) < gen_dp: + padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] + responses = sampler.sample(padded, params)[:len(prompts)] + return [list(r.sequences) if (r and r.sequences) else [] for r in responses] + + +# =========================================================================== +# Block C -- data loading via twinkle.Dataset + numeric filtering +# =========================================================================== +def load_problems(dataset: str, n: int, seed: int) -> List[Dict[str, Any]]: + """Load boxed-answer problems as ``{problem, reference_answer, level?}`` via + twinkle.Dataset (ModelScope hub), sampled to ``n`` (0 = all).""" + ds_id = AOPS_DATASET_ID if dataset == 'aops' else MATH_DATASET_ID + rows = Dataset(DatasetMeta(dataset_id=f'ms://{ds_id}', split='train')).dataset + out: List[Dict[str, Any]] = [] + for row in rows: + if dataset == 'aops' and not (row.get('metadata') or {}).get('boxed'): + continue + ref = extract_boxed(row.get('solution', '')) + if not ref: + continue + rec = {'problem': row['problem'], 'reference_answer': ref} + if row.get('level'): + rec['level'] = row['level'] + out.append(rec) + logger.info(f'[data] {dataset}: {len(out)} boxed problems') + rng = np.random.RandomState(seed) + rng.shuffle(out) + return out[:n] if (n and n < len(out)) else out + + +_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') + + +def _norm_num_text(num: str) -> str: + try: + f = float(num) + return str(int(f)) if f == int(f) else str(f) + except Exception: + return str(num).strip() + + +def _numeric_value(raw: Any) -> Optional[str]: + """Collapse an answer to a single int/decimal/fraction, or None.""" + if raw is None: + return None + s = str(raw).strip().strip('$').strip() + s = s.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac') + s = re.sub(r'\\!|\\,|\\;|\\ |\\left|\\right|\s', '', s) + for pat in (r'\\frac\{(-?\d+)\}\{(-?\d+)\}', r'(-?\d+)/(-?\d+)'): + m = re.fullmatch(pat, s) + if m: + a, b = int(m.group(1)), int(m.group(2)) + return _norm_num_text(str(a / b)) if b else None + return _norm_num_text(s) if _NUM_RE.fullmatch(s) else None + + +def _load_records(args: argparse.Namespace) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], Dict[str, int]]: + """Load, numeric-filter, shuffle, then split a fixed eval holdout off the front.""" + # Load all when filtering or splitting (else the eval holdout could starve train). + load_n = 0 if (args.numeric_only or args.eval_size > 0) else args.n + records = load_problems(args.dataset, load_n, args.seed) + raw_n, dropped = len(records), 0 + if args.numeric_only: + kept = [] + for r in records: + ref = _numeric_value(r.get('reference_answer')) + if ref is None: + dropped += 1 + continue + kept.append({**r, 'reference_answer': ref}) + records = kept + np.random.RandomState(args.seed).shuffle(records) + eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 + eval_records = [dict(r) for r in records[:eval_n]] + pool = records[eval_n:] + train_n = args.n if args.n > 0 else len(pool) + train_records = [dict(r) for r in pool[:train_n]] + overlap = {r['problem'] for r in train_records} & {r['problem'] for r in eval_records} + if overlap: + raise ValueError(f'eval/train overlap: {len(overlap)} problems') + stats = {'raw_loaded': raw_n, 'numeric_dropped': dropped, + 'train_records': len(train_records), 'eval_records': len(eval_records)} + return train_records, eval_records, stats + + +# =========================================================================== +# Block D -- disk cache, problem pool, baseline rollout, rubric check +# =========================================================================== +class DiskCache: + """Append-only jsonl kv cache (md5 key -> value). Loads on init, appends on put. + Disabled instances (``enabled=False``) always miss and never write.""" + + def __init__(self, path: str, enabled: bool = True): + self.path, self.enabled = path, enabled + self._mem: Dict[str, Any] = {} + self._fh = None + if not enabled: + return + if os.path.exists(path): + with open(path, encoding='utf-8') as f: + for line in f: + if line.strip(): + row = json.loads(line) + self._mem[row['key']] = row['value'] + self._fh = open(path, 'a', encoding='utf-8') + + @staticmethod + def key_for(*parts: str) -> str: + return hashlib.md5('\x1f'.join(parts).encode('utf-8')).hexdigest() + + def __contains__(self, key: str) -> bool: + return key in self._mem + + def get(self, key: str) -> Any: + return self._mem.get(key) + + def put(self, key: str, value: Any) -> None: + self._mem[key] = value + if self._fh is not None: + self._fh.write(json.dumps({'key': key, 'value': value}, ensure_ascii=False) + '\n') + self._fh.flush() + + def close(self) -> None: + if self._fh is not None: + self._fh.close() + + +class ProblemPool: + """Cyclic draw source. Each full pass reshuffles with ``seed + epoch``; the initial + pass keeps the loader order. ``draw(k)`` returns k distinct problems (k << pool).""" + + def __init__(self, records: List[Dict[str, Any]], seed: int): + self._records = list(records) + self._seed, self._cursor, self.epoch = seed, 0, 0 + + def draw(self, k: int) -> List[Dict[str, Any]]: + out, seen = [], set() + while len(out) < k: + if self._cursor >= len(self._records): + self.epoch += 1 + np.random.RandomState(self._seed + self.epoch).shuffle(self._records) + self._cursor = 0 + r = self._records[self._cursor] + self._cursor += 1 + if id(r) not in seen: + seen.add(id(r)) + out.append(r) + return out + + +def _empty_roll() -> Dict[str, Any]: + return {'pred': '', 'correct': False, 'terminated': False, 'passed': False, + 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} + + +def _apply_baseline(r: Dict[str, Any], roll: Dict[str, Any]) -> None: + """Attach a greedy baseline roll and reset per-chunk working state.""" + r['_baseline_rolls'], r['_cands'], r['_init'] = [roll], [], [roll] + r['_failed'] = not roll['correct'] + r['_baseline_pass'] = 1.0 if roll['correct'] else 0.0 + r['_hard'] = True # process every problem; group variance selects (SEAM-style) + + +def baseline_rollout(base_sampler, problems: List[Dict[str, Any]], base_dp: int, + args: argparse.Namespace, cache: DiskCache) -> int: + """Phase 1: base solves each problem greedily once (T=0, M=1), disk-cached by + problem text. The base is frozen + greedy so the cache is exact. Returns the number + of fresh (cache-miss) rollouts.""" + todo = [r for r in problems if DiskCache.key_for(r['problem']) not in cache] + if todo: + out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], + 1, args.max_tokens, base_dp, temperature=0.0) + for r, seqs in zip(todo, out): + roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + cache.put(DiskCache.key_for(r['problem']), roll) + for r in problems: + _apply_baseline(r, cache.get(DiskCache.key_for(r['problem']))) + return len(todo) + + +# -- rubric process-check (view A): teacher diagnoses the base's attempt -- +_RFT_DIAG_SYSTEM = """\ +You are a process error checker for a math solution attempt. You are given a math +problem, a rubric, and one attempted solution segment. Decide PASS or FAIL for each +criterion and explain only the process error type. + +Output STRICT JSON (no prose outside it) with this shape: +{ + "items": [ + {"index": 1, "verdict": "PASS", "reason": "", + "fix": ""}, + {"index": 2, "verdict": "FAIL", "reason": "", + "fix": ""} + ], + "overall": "OK" | "ISSUES", + "summary": "" +} + +Rules: +- Judge every criterion independently and literally; a [Hard Rule] is FAIL unless + unambiguously satisfied. +- Judge ONLY what is observable in THIS segment. +- Content inside ... (or ) is internal reasoning, not + user-facing output; ignore it for "output only X" style criteria. +- For PASS items, leave "fix" as "". +- For FAIL items, "reason", "fix", and "summary" must describe only the flawed + step, theorem, arithmetic operation, case split, or verification habit. +- NEVER state the correct final answer, corrected final expression, option letter, + graph/choice label, or any exact value that the answer should become. +- NEVER write phrases like "the correct answer is", "which gives", "yielding", + "should be ", "Option ", or "Graph ". +- If a fix would require naming a corrected value, replace it with a method-level + instruction such as "redo that computation carefully" or "apply the theorem with + the correct quantities". +- Keep every "reason" and "fix" clear and concise — one short sentence each. +- "overall" is "OK" only if NO criterion is FAIL. +- Output only the JSON object.""" + +_RFT_DIAG_USER = """\ +## Task / query (context) +{query} + +## Rubric +{rubric} + +## Segment +{segment} + +Now output the diagnostic JSON object.""" + +_MATH_RUBRIC = [ + ('The reasoning contains no arithmetic or algebraic error', True), + ('Each step follows logically from the previous ones', True), + ('No formula or theorem is misstated or misapplied', True), + ('The approach is on track to answer the actual question asked', False), + ('No step contradicts an earlier established fact', False), +] + + +class _RftRubricVerifier(RubricVerifier): + def _diagnose_trajectory(self, query: str, rubric_block: str, segment_text: str) -> dict: + return {'messages': [ + {'role': 'system', 'content': _RFT_DIAG_SYSTEM}, + {'role': 'user', 'content': _RFT_DIAG_USER.format( + query=query, rubric=rubric_block, segment=segment_text)}]} + + +def build_rubric_checker() -> Optional[RubricVerifier]: + """Fixed math-process rubric verifier (teacher-served). None if no LLM backup env.""" + if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') + or os.environ.get('OPENAI_API_KEY')): + return None + return _RftRubricVerifier( + fixed_rubric=[RubricItem(t, is_hard=h) for t, h in _MATH_RUBRIC], gate=True) + + +def _format_diagnosis(detail) -> str: + """One line per criterion (PASS/FAIL + reason + fix) then a summary.""" + rub = detail.rubric + lines = [] + for it in detail.items: + text = rub[it.index - 1].text if 0 < it.index <= len(rub) else f'criterion {it.index}' + if it.verdict: + lines.append(f'- [PASS] {text}') + else: + tail = f': {it.reason}' if it.reason else '' + tail += f' (fix: {it.fix})' if it.fix else '' + lines.append(f'- [FAIL] {text}{tail}') + if detail.summary: + lines.append(f'Summary: {detail.summary}') + return '\n'.join(lines) + + +def diagnose_views(checker, hard: List[Dict[str, Any]], args: argparse.Namespace, + cache: DiskCache) -> None: + """Rubric-check every view-A problem's greedy attempt in parallel (disk-cached by + problem + attempt), stashing the formatted findings on ``r['_rubric_diag']``.""" + targets = [r for r in hard if r.get('_view') == 'A'] + if not checker or not targets: + return + + def _key(r: Dict[str, Any]) -> str: + return DiskCache.key_for(r['problem'], r.get('_init', [{}])[0].get('text', '')) + + pending = [] + for r in targets: + key = _key(r) + if key in cache: + r['_rubric_diag'] = cache.get(key) + else: + pending.append((r, key)) + if not pending: + return + + def _run(item): + r, key = item + seg = {'messages': [{'role': 'user', 'content': r['problem']}, + {'role': 'assistant', 'content': r['_init'][0]['text']}]} + try: + return r, key, _format_diagnosis(checker.diagnose(seg, query=r['problem'])) + except Exception as exc: # teacher hiccup -> no-diagnosis prompt (not cached) + logger.warning(f'[rubric] diagnose error: {exc}') + return r, key, None + + workers = max(1, min(args.rubric_workers, len(pending))) + with ThreadPoolExecutor(max_workers=workers) as ex: + for r, key, diag in ex.map(_run, pending): + r['_rubric_diag'] = diag or '' + if diag is not None: + cache.put(key, diag) + + +# =========================================================================== +# Block E -- chunk draw, pipeline, record building +# =========================================================================== +def _baseline_class(r: Dict[str, Any]) -> str: + """success | fail_loop (out of length / never terminated) | fail_wrong.""" + roll = r['_init'][0] + if roll['correct']: + return 'success' + return 'fail_loop' if (roll['stop_reason'] == 'length' or not roll['terminated']) else 'fail_wrong' + + +def _select_balanced(buckets: Dict[str, List[Dict[str, Any]]], n_success: int, + n_fail: int, n_fail_loop: int) -> List[Dict[str, Any]]: + """Pick n_fail base-fails (toward n_fail_loop loop-fails, best-effort) + n_success + base-successes; top up any shortfall from leftovers.""" + loop, wrong, succ = buckets['fail_loop'], buckets['fail_wrong'], buckets['success'] + take_loop = min(n_fail_loop, len(loop)) + take_wrong = min(n_fail - take_loop, len(wrong)) + take_loop = min(n_fail - take_wrong, len(loop)) + sel = loop[:take_loop] + wrong[:take_wrong] + succ[:n_success] + target = n_success + n_fail + if len(sel) < target: + used = {id(x) for x in sel} + sel += [x for b in (loop, wrong, succ) for x in b if id(x) not in used][:target - len(sel)] + return sel + + +def draw_chunk(pool: ProblemPool, base_sampler, base_dp: int, args: argparse.Namespace, + cache: DiskCache) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + """Draw one chunk, running baseline rollout (Phase 1) on every drawn problem. With + ``--balance``, keep drawing+baselining until the target base fail:success mix is + reachable (or the budget is hit), then select a balanced subset.""" + if not args.balance: + chunk = pool.draw(args.chunk_size) + fresh = baseline_rollout(base_sampler, chunk, base_dp, args, cache) + return chunk, {'enabled': False, 'n_drawn': len(chunk), 'n_baseline_fresh': fresh} + + n_success = max(0, min(args.chunk_size, round(args.chunk_size * args.balance_success_frac))) + n_fail = args.chunk_size - n_success + n_fail_loop = round(n_fail * args.balance_loop_frac) + buckets: Dict[str, List[Dict[str, Any]]] = {'success': [], 'fail_loop': [], 'fail_wrong': []} + budget, n_drawn, n_fresh, seen = args.chunk_size * args.balance_max_draws_mult, 0, 0, set() + while n_drawn < budget: + if (len(buckets['success']) >= n_success + and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail): + break + batch = pool.draw(args.chunk_size) + n_fresh += baseline_rollout(base_sampler, batch, base_dp, args, cache) + n_drawn += len(batch) + for r in batch: + if id(r) not in seen: + seen.add(id(r)) + buckets[_baseline_class(r)].append(r) + + reached = (len(buckets['success']) >= n_success + and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail) + chunk = _select_balanced(buckets, n_success, n_fail, n_fail_loop) + sel_success = sum(1 for r in chunk if not r['_failed']) + stats = { + 'enabled': True, 'n_drawn': n_drawn, 'n_baseline_fresh': n_fresh, 'n_selected': len(chunk), + 'target_success': n_success, 'target_fail': n_fail, 'target_fail_loop': n_fail_loop, + 'selected_success': sel_success, 'selected_fail': len(chunk) - sel_success, + 'selected_fail_loop': sum(1 for r in chunk if _baseline_class(r) == 'fail_loop'), + 'selected_fail_wrong': sum(1 for r in chunk if _baseline_class(r) == 'fail_wrong'), + 'selected_success_frac': (sel_success / len(chunk)) if chunk else 0.0, + 'budget_hit': not reached, + } + return chunk, stats + + +def _assign_advantages(hard: List[Dict[str, Any]], args: argparse.Namespace) -> None: + """Group-relative advantage over each problem's scored candidates using the greedy + binary reward R in {0,1}: ``A = (R - mean) / (std + eps)``. std==0 groups (all solve + / all fail) get advantage 0 and no gradient -- GRPO's variance selects informative + problems, so no explicit difficulty gate is needed.""" + eps = 1e-6 + for r in hard: + for c in r['_cands']: + c['advantage'], c['grpo_adv'], c['kept'] = 0.0, 0.0, False + cs = ([c for c in r['_cands'] if c.get('reward') is not None] if args.format_in_reward + else [c for c in r['_cands'] if c['leaked'] is False and c.get('reward') is not None]) + if len(cs) < 2: + continue + rewards = [c['reward'] for c in cs] + mean_r = sum(rewards) / len(rewards) + std = (sum((x - mean_r) ** 2 for x in rewards) / len(rewards)) ** 0.5 + if std < 1e-9: + continue + for c in cs: + adv = (c['reward'] - mean_r) / (std + eps) + c['advantage'], c['grpo_adv'], c['kept'] = adv, adv, c['reward'] > mean_r + + +def process_chunk(base_sampler, skill_sampler, leak: LeakVerifier, chunk: List[Dict[str, Any]], + ci: int, base_dp: int, skill_dp: int, args: argparse.Namespace, + checker, rubric_cache: DiskCache + ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], List[Dict[str, Any]]]: + """view assign -> rubric-check (view A) -> skill-gen -> leak-filter -> with-skill + greedy pass -> GRPO advantages. ``chunk`` arrives already baselined by draw_chunk.""" + hard = chunk + + # Phase 2: view routing + view-A rubric check (view B is query-only, no rubric). + for r in hard: + r['_view'], r['_rubric_diag'] = _assign_view(r['problem'], args), '' + diagnose_views(checker, hard, args, rubric_cache) + + # Phase 3: skill-gen (thinking ON), per-view prompt; re-sample empties. + flat: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] + pending = list(hard) + for _ in range(args.skill_retries + 1): + if not pending: + break + sg_out = _run_samples(skill_sampler, [_view_prompt(r) for r in pending], + args.n_skills, args.skill_max_tokens, skill_dp, + temperature=args.skill_gen_temperature, + top_p=args.skill_gen_top_p, top_k=args.skill_gen_top_k) + still = [] + for r, seqs in zip(pending, sg_out): + got = False + for s in seqs: + resp = _clean_text(getattr(s, 'decoded', '') or '') + block = _extract_skills_block(resp) + cand = {'skills': block or '', 'response': resp, 'parseable': bool(block), + 'view': r['_view'], 'leaked': None, 'leak_reason': '', + 'leak_source': '', 'with_pass': None, 'reward': None, 'rolls': []} + r['_cands'].append(cand) + if block: + flat.append((r, cand)) + got = True + if not got: + still.append(r) + pending = still + + # Phase 4: leak filter (view A only; view B is query-only -> treated clean, SEAM-like). + for r, c in flat: + if r.get('_view') != 'A': + c['leaked'], c['leak_reason'], c['leak_source'] = False, '', 'skipped_viewB' + flat_a = [(r, c) for r, c in flat if r.get('_view') == 'A'] + if flat_a: + details = leak.leak_batch( + [{'content': c['skills'], 'query': r['problem'], 'reference': r['reference_answer']} + for r, c in flat_a], max_workers=args.leak_workers) + for (r, c), d in zip(flat_a, details): + c['leaked'], c['leak_reason'], c['leak_source'] = bool(d.leaked), d.reason, d.source + + # Phase 5: with-skill greedy pass (T=0, M=1) on clean candidates. Reward = correct, + # absolute (no baseline subtraction); the group mean in Phase 6 is the only baseline. + clean = [(r, c) for r, c in flat if c['leaked'] is False] + if clean: + ws_out = _run_samples(base_sampler, + [build_skill_solve_prompt(r['problem'], c['skills']) for r, c in clean], + 1, args.max_tokens, base_dp, temperature=0.0) + for (r, c), seqs in zip(clean, ws_out): + c['rolls'] = [_parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll()] + c['with_pass'] = 1.0 if c['rolls'][0]['correct'] else 0.0 + c['reward'] = c['with_pass'] + if args.format_in_reward: # unparseable/leaked candidates score 0 and still join the group + for r in hard: + for c in r['_cands']: + if c['reward'] is None: + c['reward'] = 0.0 + + # Phase 6: group-relative GRPO advantage. + _assign_advantages(hard, args) + return ([_full_record(r, ci) for r in chunk], _chunk_summary(chunk, ci, args), + _group_records(chunk, args)) + + +def _roll(x: Dict[str, Any]) -> Dict[str, Any]: + return {k: x[k] for k in ('pred', 'correct', 'terminated', 'passed', + 'stop_reason', 'gen_tokens', 'text')} + + +def _is_trainable(c: Dict[str, Any], args: argparse.Namespace) -> bool: + """A candidate reaches the GRPO update iff its advantage is non-zero (and, without + --format-in-reward, is also clean and scored).""" + adv_nz = abs(c.get('advantage') or 0.0) > 1e-9 + if args.format_in_reward: + return adv_nz + return c['leaked'] is False and c.get('with_pass') is not None and adv_nz + + +def _full_record(r: Dict[str, Any], ci: int) -> Dict[str, Any]: + """Complete per-problem trace: init attempt, baseline, and all candidates.""" + init = r['_init'][0] + return { + 'record_type': 'problem', 'chunk': ci, 'problem': r['problem'], + 'reference_answer': r['reference_answer'], 'level': r.get('level', ''), + 'failed_first_try': r['_failed'], + 'init_attempt': {'text': init['text'], 'pred': init['pred'], 'correct': init['correct'], + 'terminated': init['terminated'], 'stop_reason': init['stop_reason'], + 'gen_tokens': init['gen_tokens']}, + 'baseline_pass': r['_baseline_pass'], 'is_hard': r['_hard'], + 'view': r.get('_view', ''), 'rubric_diag': r.get('_rubric_diag', ''), + 'baseline_rolls': [_roll(x) for x in r['_baseline_rolls']], + 'candidates': [{ + 'skills': c['skills'], 'response': c['response'], 'parseable': c['parseable'], + 'leaked': c['leaked'], 'leak_reason': c['leak_reason'], 'leak_source': c['leak_source'], + 'with_pass': c['with_pass'], 'reward': c.get('reward'), 'advantage': c.get('advantage'), + 'grpo_adv': c.get('grpo_adv'), 'kept': c.get('kept'), + 'rolls': [_roll(x) for x in c['rolls']], + } for c in r['_cands']], + } + + +def _view_stats(hard: List[Dict[str, Any]], view: str) -> Dict[str, Any]: + hv = [r for r in hard if r.get('_view') == view] + cands = [c for r in hv for c in r['_cands'] if c['parseable']] + clean = [c for c in cands if c['leaked'] is False] + adopted = sum(1 for r in hv + if any(c['leaked'] is False and abs(c.get('advantage') or 0.0) > 1e-9 + for c in r['_cands'])) + return {'n_hard': len(hv), 'n_candidates_parseable': len(cands), 'n_clean': len(clean), + 'n_adopted_problems': adopted, 'adoption_rate': (adopted / len(hv)) if hv else 0.0} + + +def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespace) -> Dict[str, Any]: + hard = [r for r in chunk if r['_hard']] + all_cands = [c for r in chunk for c in r['_cands']] + cands = [c for c in all_cands if c['parseable']] + scored = [c for c in cands if c['with_pass'] is not None] + ws_rolls = [x for c in scored for x in c['rolls']] + train_cands = [c for c in all_cands if _is_trainable(c, args)] + fail_cands = [c for r in chunk if r['_failed'] for c in r['_cands']] + base_acc = (sum(r['_baseline_pass'] for r in hard) / len(hard)) if hard else 0.0 + ws_acc = (sum(c['with_pass'] for c in scored) / len(scored)) if scored else 0.0 + abs_adv = lambda cs: sum(abs(c.get('advantage') or 0.0) for c in cs) + total_abs = abs_adv(all_cands) + return { + 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), + 'n_failed_first_try': sum(1 for r in chunk if r['_failed']), 'n_hard': len(hard), + 'n_generated': len(all_cands), 'n_candidates_parseable': len(cands), + 'n_unparseable': len(all_cands) - len(cands), + 'n_leaked': sum(1 for c in cands if c['leaked']), + 'n_clean': sum(1 for c in cands if c['leaked'] is False), + 'n_reward_pos': sum(1 for c in scored if c['reward']), 'n_train_samples': len(train_cands), + 'n_train_from_fail': sum(1 for c in fail_cands if _is_trainable(c, args)), + 'abs_adv_from_fail_frac': (abs_adv(fail_cands) / total_abs) if total_abs > 0 else 0.0, + 'avg_baseline_pass_on_hard': base_acc, 'avg_withskill_pass': ws_acc, + 'avg_lift': ws_acc - base_acc, + 'termination_rate_withskill': (sum(1 for x in ws_rolls if x['terminated']) / len(ws_rolls)) if ws_rolls else 0.0, + 'view_A': _view_stats(hard, 'A'), 'view_B': _view_stats(hard, 'B'), + } + + +def _group_records(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> List[Dict[str, Any]]: + """GRPO training records: every trainable candidate with its view + rubric diagnosis + (the prompt is rebuilt from those by ``_skillgen_messages``, no trajectory stored).""" + out = [] + for r in chunk: + if not r['_hard']: + continue + for c in r['_cands']: + if _is_trainable(c, args): + out.append({ + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'view': r.get('_view', 'A'), 'diagnosis': r.get('_rubric_diag', ''), + 'response': c['response'], 'skills': c['skills'], + 'advantage': c['advantage'], 'grpo_adv': c['grpo_adv'], 'kept': c['kept'], + 'reward': c['reward'], 'with_pass': c['with_pass']}) + return out + + +# =========================================================================== +# Block F -- samplers, args, main +# =========================================================================== +def init_samplers(args: argparse.Namespace): + """8 GPUs: ranks 0-3 skill_sampler, ranks 4-7 base_sampler (both vLLM tp1 dp4).""" + twinkle.initialize(mode='ray', nproc_per_node=8, lazy_collect=False, groups=[ + DeviceGroup(name='skill_sampler', ranks=list(range(0, 4)), device_type='GPU'), + DeviceGroup(name='base_sampler', ranks=list(range(4, 8)), device_type='GPU')]) + samplers = [] + for group in ('skill_sampler', 'base_sampler'): + s = vLLMSampler( + model_id=MODEL_ID, + engine_args={'gpu_memory_utilization': GPU_MEM, + 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, + device_mesh=DeviceMesh.from_sizes(world_size=4, dp_size=4), remote_group=group) + s.set_template('Template', model_id=MODEL_ID, enable_thinking=True, + max_length=args.max_model_len) + samplers.append(s) + return samplers[1], samplers[0], 4, 4 # base_sampler, skill_sampler, base_dp, skill_dp + + +def _build_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--total-problems', type=int, default=3200, + help='Final number of problems selected into generated chunks.') + p.add_argument('--base-success-frac', type=float, default=0.3, + help='Target fraction of selected problems the frozen base solves.') + p.add_argument('--output-dir', default='./output/reflexion_skill_data') + p.add_argument('--cache-dir', default='', + help='Baseline/rubric cache dir (default /cache).') + p.add_argument('--no-cache', action='store_true', help='Disable disk cache read/write.') + p.add_argument('--overwrite', action='store_true', help='Replace existing output jsonl.') + p.add_argument('--seed', type=int, default=42) + p.add_argument('--dataset', choices=('aops', 'math'), default='aops') + p.add_argument('--n', type=int, default=0, + help='Raw train-pool size; 0 derives it from --total-problems.') + p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--eval-size', type=int, default=128) + p.add_argument('--chunk-size', type=int, default=16) + p.add_argument('--balance', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--balance-loop-frac', type=float, default=0.5) + p.add_argument('--balance-max-draws-mult', type=int, default=8) + p.add_argument('--n-skills', type=int, default=8) + p.add_argument('--view-b-frac', type=float, default=0.5) + p.add_argument('--skill-retries', type=int, default=2) + p.add_argument('--skill-gen-temperature', type=float, default=1.0) + p.add_argument('--skill-gen-top-p', type=float, default=1.0) + p.add_argument('--skill-gen-top-k', type=int, default=-1) + p.add_argument('--max-model-len', type=int, default=16384) + p.add_argument('--max-tokens', type=int, default=8192) + p.add_argument('--skill-max-tokens', type=int, default=8192) + p.add_argument('--leak-workers', type=int, default=16) + p.add_argument('--rubric-workers', type=int, default=16) + p.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True) + args = p.parse_args() + if args.total_problems <= 0 or args.chunk_size <= 0: + raise ValueError('--total-problems and --chunk-size must be positive') + if not 0.0 <= args.base_success_frac <= 1.0: + raise ValueError('--base-success-frac must be in [0, 1]') + args.chunks = math.ceil(args.total_problems / args.chunk_size) + args.balance_success_frac = args.base_success_frac + if args.n <= 0: + args.n = max(args.total_problems + args.eval_size, math.ceil(args.total_problems * 1.5)) + return args + + +def _write(handle, row: Dict[str, Any]) -> None: + handle.write(json.dumps(row, ensure_ascii=False) + '\n') + + +def main() -> None: + args = _build_args() + records, eval_records, data_stats = _load_records(args) + if not records: + raise ValueError(f'loaded 0 {args.dataset} problems') + if len(records) < args.chunk_size: + raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded ({len(records)})') + + os.makedirs(args.output_dir, exist_ok=True) + data_path = os.path.join(args.output_dir, 'skill_dataset.jsonl') + gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') + eval_path = os.path.join(args.output_dir, 'eval_holdout.jsonl') + for path in (data_path, gen_path, eval_path): + if os.path.exists(path) and not args.overwrite: + raise FileExistsError(f'{path} exists; pass --overwrite to replace it') + + if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY')): + sys.stderr.write('[build] WARNING: no LLM backup env; leak/rubric checks degrade\n') + + base_sampler, skill_sampler, base_dp, skill_dp = init_samplers(args) + leak = LeakVerifier(sampler=None, answer_only=True) + checker = build_rubric_checker() + pool = ProblemPool(records, args.seed) + + cache_dir = args.cache_dir or os.path.join(args.output_dir, 'cache') + os.makedirs(cache_dir, exist_ok=True) + use_cache = not args.no_cache + baseline_cache = DiskCache(os.path.join(cache_dir, 'baseline.jsonl'), use_cache) + rubric_cache = DiskCache(os.path.join(cache_dir, 'rubric.jsonl'), use_cache) + + cfg = { + 'record_type': 'config', 'mode': 'offline_data_build', 'model': MODEL_ID, + 'dataset': args.dataset, 'n': len(records), 'eval_n': len(eval_records), + 'total_problems': args.total_problems, 'seed': args.seed, 'numeric_only': args.numeric_only, + 'raw_loaded': data_stats['raw_loaded'], 'numeric_dropped': data_stats['numeric_dropped'], + 'chunks': args.chunks, 'chunk_size': args.chunk_size, 'n_skills': args.n_skills, + 'view_b_frac': args.view_b_frac, 'balance': args.balance, + 'base_success_frac': args.base_success_frac, 'balance_success_frac': args.balance_success_frac, + 'reward': 'greedy_binary(correct)', 'advantage': 'group_relative', + 'format_in_reward': args.format_in_reward, 'cache': use_cache, + 'rubric_check': 'fixed_math_5crit(viewA)' if checker else 'disabled', + 'started': int(time.time()), + } + total_groups, selected = 0, 0 + with open(gen_path, 'w', encoding='utf-8') as gen_f, \ + open(data_path, 'w', encoding='utf-8') as data_f, \ + open(eval_path, 'w', encoding='utf-8') as eval_f: + for handle in (gen_f, data_f, eval_f): + _write(handle, cfg) + for rec in eval_records: + _write(eval_f, {'record_type': 'eval_holdout', **rec}) + eval_f.flush() + + full_chunk_size = args.chunk_size + for ci in range(args.chunks): + remaining = args.total_problems - selected + if remaining <= 0: + break + args.chunk_size = min(full_chunk_size, remaining) # last chunk may be short + chunk, balance = draw_chunk(pool, base_sampler, base_dp, args, baseline_cache) + full, summary, groups = process_chunk( + base_sampler, skill_sampler, leak, chunk, ci, base_dp, skill_dp, + args, checker, rubric_cache) + summary['balance'] = balance + for rec in full: + _write(gen_f, rec) + _write(gen_f, summary) + gen_f.flush() + for row in groups: + _write(data_f, {'chunk': ci, **row}) + data_f.flush() + total_groups += len(groups) + selected += len(chunk) + sys.stderr.write( + f'[build] g{ci}: problems={selected}/{args.total_problems} ' + f'train={len(groups)} total={total_groups} ' + f'acc={summary["avg_baseline_pass_on_hard"]:.2f}->{summary["avg_withskill_pass"]:.2f} ' + f'lift={summary["avg_lift"]:+.3f}\n') + + baseline_cache.close() + rubric_cache.close() + sys.stderr.write(f'[build] done: {total_groups} train records -> {data_path}; trace -> {gen_path}\n') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/legacy/build_thinking_rag_index.py b/cookbook/exp/legacy/build_thinking_rag_index.py new file mode 100644 index 000000000..a71bae060 --- /dev/null +++ b/cookbook/exp/legacy/build_thinking_rag_index.py @@ -0,0 +1,1159 @@ +"""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 ``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 + ``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 +import threading +import time +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 + +# --------------------------------------------------------------------------- +# Compress prompts — MUST match train_embedding_full_ddp.py exactly. +# --------------------------------------------------------------------------- +_HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(_HERE)) + +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 + +_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', + 'output/embedding_full_transformers/last-checkpoint', +) +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)) + +# Dataset mix caps (only used in 'both' mode). None = no cap. +THINK_CAP: Optional[int] = int(os.environ.get('THINK_CAP', 400_000)) or None +INDEX_CAP: Optional[int] = int(os.environ.get('INDEX_CAP', 400_000)) or None +MIX_SHUFFLE_SEED = 100 + +# Concurrency knobs for API fallback and prefetch pipeline. +API_CONCURRENCY = int(os.environ.get('API_CONCURRENCY', 8)) +API_MIN_INTERVAL = float(os.environ.get('API_MIN_INTERVAL', 0.1)) +PREFETCH_WORKERS = int(os.environ.get('PREFETCH_WORKERS', 2)) + +# Hard-templated hints: the condenser SFT prior maps `Skill` to the legacy +# `Use when: / numbered steps / Output:` skeleton on long inputs; embedding the +# exact 4-line body template + explicit negative constraints is the only way to +# override it deterministically across query and cot sides. +RAG_QUERY_HINT = ( + '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.') +RAG_THINKING_HINT = ( + '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.') + +# 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 +# =========================================================================== + +_LEGACY_USE_WHEN_RE = re.compile(r'(?im)^\s*Use when\s*:') +_SCHEMA_MARKERS = ('Problem:', 'Skill:', 'Knowledge:') + + +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 + 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 + 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 _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) + + +_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. + + API fallback runs concurrently (up to API_CONCURRENCY workers) for speed. + """ + pairs = _vllm_compress(sampler, texts, query_hint) + 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[i] = text + else: + fallback_indices.append(i) + + if fallback_indices and api is not None: + from concurrent.futures import ThreadPoolExecutor, as_completed + with ThreadPoolExecutor(max_workers=API_CONCURRENCY) as pool: + futures = {} + for idx in fallback_indices: + msgs = _build_compress_messages(texts[idx], query_hint) + futures[pool.submit(_api_compress_throttled, api, msgs)] = idx + for fut in as_completed(futures): + idx = futures[fut] + api_text = fut.result() + if api_text and not _is_truncated_compression(api_text): + results[idx] = api_text + + return results + + +def _resolve_compressed_multi(sampler: vLLMSampler, api: Optional[OpenAIClient], + texts: List[str], hints: List[str]) -> List[Optional[str]]: + """Like _resolve_compressed but each text has its own per-item hint. + + Merges all texts into a SINGLE vLLM batch call (instead of one per hint), + dramatically reducing round-trip overhead when processing interleaved + query+cot pairs with different hint strings. + + Args: + sampler: vLLM condenser sampler. + api: Optional OpenAI-compatible API client for fallback. + texts: List of raw texts to compress (may contain empty strings to skip). + hints: Per-text hint strings (same length as texts). + + Returns: + List of compressed texts (None where compression failed entirely). + """ + assert len(texts) == len(hints), f'texts({len(texts)}) != hints({len(hints)})' + if not texts: + return [] + + # Skip texts that would exceed the condenser's context window. + _max_input_chars = (CONDENSE_MAX_MODEL_LEN - CONDENSE_MAX_TOKENS) * 3 + skip_mask = [len(t) > _max_input_chars for t in texts] + + # Build prompts per-item (each text gets its own hint as the query parameter). + prompts = [{'messages': _build_compress_messages(t, h)} + for t, h, skip in zip(texts, hints, skip_mask) if not skip] + active_indices = [i for i, skip in enumerate(skip_mask) if not skip] + params = TwinkleSamplingParams( + max_tokens=CONDENSE_MAX_TOKENS, + temperature=COMPRESS_TEMPERATURE, + top_p=COMPRESS_TOP_P, + num_samples=1, + ) + + # Single vLLM batch call — the key throughput win. + responses = sampler.sample(prompts, params) if prompts else [] + + results: List[Optional[str]] = [None] * len(texts) + fallback_indices: List[int] = [] + for resp_idx, orig_idx in enumerate(active_indices): + resp = responses[resp_idx] + seq = resp.sequences[0] if resp and resp.sequences else None + if seq is None: + fallback_indices.append(orig_idx) + continue + text = seq.decoded or '' + text = re.sub(r'<\|[^|]+\|>', '', text).rstrip() + text = _strip_outer_codefence(text) + if seq.stop_reason != 'length' and not _is_truncated_compression(text): + results[orig_idx] = text + else: + fallback_indices.append(orig_idx) + + # Concurrent API fallback for failed items. + if fallback_indices and api is not None: + from concurrent.futures import ThreadPoolExecutor, as_completed + with ThreadPoolExecutor(max_workers=API_CONCURRENCY) as pool: + futures = {} + for idx in fallback_indices: + msgs = _build_compress_messages(texts[idx], hints[idx]) + futures[pool.submit(_api_compress_throttled, api, msgs)] = idx + for fut in as_completed(futures): + idx = futures[fut] + api_text = fut.result() + if api_text and not _is_truncated_compression(api_text): + results[idx] = api_text + + return results + + +# =========================================================================== +# 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 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' + + (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 + 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 _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 [] + # 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 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: + 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), + 'cot_raw_head': _short(r['cot_raw'], 200), + }) + continue + 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 + 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) + 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: + 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), + '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) + 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 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: + 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, + 'source': row.get('source') or 'unknown', + 'query_raw': user_query, + 'cot_raw': cot, + }) + if len(batch) >= args.batch_size: + pending_futures.append(prefetch_pool.submit(_process_batch, list(batch))) + batch.clear() + + # Flush remainder + if 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] 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: + 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], RAG_QUERY_HINT) + 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='both', + choices=['dataset_index', 'dataset_think', 'both'], + help='Which loader to use: dataset_index (RAG profile), ' + 'dataset_think (training mix), or both (50/50 mix).') + p.add_argument('--limit', type=int, default=0, + help='Stop building once this many rows are kept (0 = no cap).') + p.add_argument('--max-rows', type=int, default=0, + help='Truncate corpus to this many rows AFTER get_dataset (0 = no cap). ' + 'Use this instead of --total to avoid invalidating the dataset cache.') + p.add_argument('--batch-size', type=int, default=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', + 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 + 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. + 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/legacy/compare_math_levels.py b/cookbook/exp/legacy/compare_math_levels.py new file mode 100644 index 000000000..d488909b9 --- /dev/null +++ b/cookbook/exp/legacy/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/legacy/dataset_hard.py b/cookbook/exp/legacy/dataset_hard.py new file mode 100644 index 000000000..9fa059b95 --- /dev/null +++ b/cookbook/exp/legacy/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/legacy/dataset_index.py b/cookbook/exp/legacy/dataset_index.py new file mode 100644 index 000000000..7d2905a59 --- /dev/null +++ b/cookbook/exp/legacy/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/legacy/dataset_think.py b/cookbook/exp/legacy/dataset_think.py new file mode 100644 index 000000000..38618ced1 --- /dev/null +++ b/cookbook/exp/legacy/dataset_think.py @@ -0,0 +1,456 @@ +import hashlib +import re +from typing import Any, Dict, List, Optional + +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.preprocessor import Preprocessor + +_THINK_RE = re.compile(r'(.*?)', re.DOTALL) + + +def _hash_id(prefix: str, content: str) -> str: + return f'{prefix}__{hashlib.md5(content.encode("utf-8")).hexdigest()[:16]}' + + +def _register(dataset, processor_cls, meta: DatasetMeta, init_args: Optional[Dict[str, Any]] = None, + load_from_cache_file: bool = True) -> None: + """Add dataset and run preprocessor; auto-strip every input column to enforce + the universal ``{id, source, query, cot, response}`` output schema.""" + dataset.add_dataset(meta) + cols = list(dataset.datasets[meta.get_id()].column_names) + dataset.map( + processor_cls, + dataset_meta=meta, + init_args=init_args or {}, + remove_columns=cols, + load_from_cache_file=load_from_cache_file, + ) + + +# ===== Modotte/CodeX-2M-Thinking ===== +CODEX_THINKING_REPO = 'ms://Modotte/CodeX-2M-Thinking' + + +class CodeXThinkingProcessor(Preprocessor): + """CodeX-2M-Thinking row → ``{id, source, query, cot, response}``。 + + 输入 schema: ``input``（问题）、``output``（含 ``...`` + 答案）。 + 拆分 output 为 cot（think 标签内容）和 response（标签之后的正文）。 + 丢弃缺失 input/output 或无法解析 think 标签的行。 + """ + + 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 '').strip() + output = (row.get('output') or '').strip() + if not query or not output: + continue + m = _THINK_RE.search(output) + if not m: + continue + cot = m.group(1).strip() + response = output[m.end():].strip() + if not cot or not response: + continue + out.append({ + 'id': _hash_id('codex_think', f'{query}\n{response}'), + 'source': 'CodeX-2M-Thinking', + 'query': query, + 'cot': cot, + 'response': response, + }) + return self.map_row_to_col(out) + + +# ===== open-thoughts/OpenThoughts3-1.2M ===== +OPEN_THOUGHTS_REPO = 'ms://open-thoughts/OpenThoughts3-1.2M' + + +class OpenThoughtsProcessor(Preprocessor): + """OpenThoughts3 row → ``{id, source, query, cot, response}``。 + + 输入 schema: ``conversations`` (messages 格式 list[{from/value}])。 + 取第一个 human 作 query，第一个 gpt 的 value 按 ``...`` 拆 cot/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: + convs = row.get('conversations') + if not isinstance(convs, list): + continue + query = '' + assistant_text = '' + for msg in convs: + if not isinstance(msg, dict): + continue + role = msg.get('from') or msg.get('role') or '' + value = msg.get('value') or msg.get('content') or '' + if role in ('human', 'user') and not query: + query = value.strip() + elif role in ('gpt', 'assistant') and not assistant_text: + assistant_text = value.strip() + break + if not query or not assistant_text: + continue + m = _THINK_RE.search(assistant_text) + if not m: + continue + cot = m.group(1).strip() + response = assistant_text[m.end():].strip() + if not cot or not response: + continue + out.append({ + 'id': _hash_id('openthoughts', f'{query}\n{response}'), + 'source': 'OpenThoughts3-1.2M', + 'query': query, + 'cot': cot, + 'response': response, + }) + return self.map_row_to_col(out) + + +# ===== GAIR/LIMO-v2 ===== +LIMO_REPO = 'ms://GAIR/LIMO-v2' + + +class LIMOProcessor(Preprocessor): + """LIMO-v2 row → ``{id, source, query, cot, response}``。 + + 输入 schema: ``question``、``solution``（含 ``...`` + 答案）。 + 拆分 solution 为 cot 和 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('question') or '').strip() + solution = (row.get('solution') or '').strip() + if not query or not solution: + continue + m = _THINK_RE.search(solution) + if m: + cot = m.group(1).strip() + response = solution[m.end():].strip() + else: + # 无 think 标签时，solution 整体作为 response，cot 留空 + cot = '' + response = solution + if not response: + continue + out.append({ + 'id': _hash_id('limo', f'{query}\n{response}'), + 'source': 'LIMO-v2', + 'query': query, + 'cot': cot, + 'response': response, + }) + return self.map_row_to_col(out) + + +# ===== AI-ModelScope/Chinese-DeepSeek-R1-Distill-data-110k ===== +CN_R1_DISTILL_REPO = 'ms://AI-ModelScope/Chinese-DeepSeek-R1-Distill-data-110k' + + +class ChineseR1DistillProcessor(Preprocessor): + """Chinese-DeepSeek-R1-Distill row → ``{id, source, query, cot, response}``。 + + 输入已有三列: ``input`` → query, ``reasoning_content`` → cot, ``content`` → 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('input') or '').strip() + cot = (row.get('reasoning_content') or '').strip() + response = (row.get('content') or '').strip() + if not query or not response: + continue + if cot: + response = _THINK_RE.sub('', response).strip() + if not response: + continue + out.append({ + 'id': _hash_id('cn_r1_distill', f'{query}\n{response}'), + 'source': 'Chinese-DeepSeek-R1-Distill-data-110k', + 'query': query, + 'cot': cot, + 'response': response, + }) + return self.map_row_to_col(out) + + +# ===== nohurry/Opus-4.6-Reasoning-3000x-filtered ===== +OPUS_REASONING_REPO = 'ms://nohurry/Opus-4.6-Reasoning-3000x-filtered' + + +class OpusReasoningProcessor(Preprocessor): + """Opus-4.6-Reasoning-3000x-filtered row → ``{id, source, query, cot, response}``。 + + 输入已有三列: ``problem`` → query, ``thinking`` → cot, ``solution`` → 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 '').strip() + cot = (row.get('thinking') or '').strip() + response = (row.get('solution') or '').strip() + if not query or not response: + continue + if cot: + response = _THINK_RE.sub('', response).strip() + if not response: + continue + out.append({ + 'id': _hash_id('opus_reasoning', f'{query}\n{response}'), + 'source': 'Opus-4.6-Reasoning-3000x-filtered', + 'query': query, + 'cot': cot, + 'response': response, + }) + return self.map_row_to_col(out) + + +# ===== Roman1111111/claude-opus-4.6-10000x ===== +CLAUDE_OPUS_REPO = 'ms://Roman1111111/claude-opus-4.6-10000x' + + +class ClaudeOpusProcessor(Preprocessor): + """claude-opus-4.6-10000x row → ``{id, source, query, cot, response}``。 + + 输入 schema: ``messages`` (OpenAI 格式 list[{role, content}])。 + 取首个 user 作 query，首个 assistant 按 ``...`` 拆 cot/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: + messages = row.get('messages') + if not isinstance(messages, list): + continue + query = '' + assistant_text = '' + for msg in messages: + if not isinstance(msg, dict): + continue + role = msg.get('role') or '' + content = msg.get('content') or '' + if not isinstance(content, str): + continue + if role == 'user' and not query: + query = content.strip() + elif role == 'assistant' and not assistant_text: + assistant_text = content.strip() + break + if not query or not assistant_text: + continue + m = _THINK_RE.search(assistant_text) + if m: + cot = m.group(1).strip() + response = assistant_text[m.end():].strip() + else: + cot = '' + response = assistant_text + if not response: + continue + out.append({ + 'id': _hash_id('claude_opus', f'{query}\n{response}'), + 'source': 'claude-opus-4.6-10000x', + 'query': query, + 'cot': cot, + 'response': response, + }) + return self.map_row_to_col(out) + + +ANGRYGIRAFFE_REPO = 'ms://hf/angrygiraffe-claude-opus-4.6-4.7-reasoning-8.7k' + + +class AngrygiraffeOpusReasoningProcessor(Preprocessor): + """angrygiraffe/claude-opus-4.6-4.7-reasoning-8.7k row → ``{id, source, query, cot, response}``。 + + 输入 schema: ``messages`` (OpenAI 格式 list[{role, content}])。 + 取首个 user 作 query，首个 assistant 按 ``...`` 拆 cot/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: + messages = row.get('messages') + if not isinstance(messages, list): + continue + query = '' + assistant_text = '' + for msg in messages: + if not isinstance(msg, dict): + continue + role = msg.get('role') or '' + content = msg.get('content') or '' + if not isinstance(content, str): + continue + if role == 'user' and not query: + query = content.strip() + elif role == 'assistant' and not assistant_text: + assistant_text = content.strip() + break + if not query or not assistant_text: + continue + m = _THINK_RE.search(assistant_text) + if m: + cot = m.group(1).strip() + response = assistant_text[m.end():].strip() + else: + cot = '' + response = assistant_text + if not response: + continue + out.append({ + 'id': _hash_id('angrygiraffe_opus', f'{query}\n{response}'), + 'source': 'angrygiraffe-claude-opus-4.6-4.7-reasoning-8.7k', + 'query': query, + 'cot': cot, + 'response': response, + }) + return self.map_row_to_col(out) + + +_BASE_SIZES = { + 'codex_think': 100000, + 'open_thoughts': 400000, + 'cn_r1_distill': 100000, + 'opus_reasoning': 3000, + 'claude_opus': 10000, + 'angrygiraffe': 38000, +} + + +def _scaled_sizes(total: Optional[int]) -> Dict[str, int]: + if total is None: + 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, CodeXThinkingProcessor, + DatasetMeta(dataset_id=CODEX_THINKING_REPO, split='train', + data_slice=range(sizes['codex_think'])), + load_from_cache_file=load_from_cache_file) + + _register(dataset, OpenThoughtsProcessor, + DatasetMeta(dataset_id=OPEN_THOUGHTS_REPO, split='train', + data_slice=range(sizes['open_thoughts'])), + load_from_cache_file=load_from_cache_file) + + _register(dataset, LIMOProcessor, + DatasetMeta(dataset_id=LIMO_REPO, split='train'), + load_from_cache_file=load_from_cache_file) + + _register(dataset, ChineseR1DistillProcessor, + DatasetMeta(dataset_id=CN_R1_DISTILL_REPO, split='train', + data_slice=range(sizes['cn_r1_distill'])), + load_from_cache_file=load_from_cache_file) + + _register(dataset, OpusReasoningProcessor, + DatasetMeta(dataset_id=OPUS_REASONING_REPO, split='train', + data_slice=range(sizes['opus_reasoning'])), + load_from_cache_file=load_from_cache_file) + + _register(dataset, ClaudeOpusProcessor, + DatasetMeta(dataset_id=CLAUDE_OPUS_REPO, split='train', + data_slice=range(sizes['claude_opus'])), + load_from_cache_file=load_from_cache_file) + + _register(dataset, AngrygiraffeOpusReasoningProcessor, + DatasetMeta(dataset_id=ANGRYGIRAFFE_REPO, split='train', + data_slice=range(sizes['angrygiraffe'])), + load_from_cache_file=load_from_cache_file) + + dataset.mix_dataset(False) + return dataset + + +class ToMessagesProcessor(Preprocessor): + """Convert {query, cot, response} → {id, source, messages}.""" + + 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('query') or '' + cot = row.get('cot') or '' + response = row.get('response') or '' + if not cot: + continue + assistant_content = f'{cot}' + out.append({ + 'id': row.get('id', ''), + 'source': row.get('source', ''), + 'messages': [ + {'role': 'user', 'content': query}, + {'role': 'assistant', 'content': assistant_content, + 'reasoning_content': cot}, + ], + }) + return self.map_row_to_col(out, keys=['id', 'source', 'messages']) + + +def get_dataset(total: Optional[int] = None, dropped_log: Optional[str] = None, + load_from_cache_file: bool = True) -> Dataset: + """Build, convert to messages format, and quality-filter the CoT dataset. + + If ``total`` is given, every per-source row count in ``_BASE_SIZES`` is + scaled proportionally so the input-row sum approximates ``total``. + """ + from twinkle_agentic.preprocessor import ( + DeadLoopFilter, + FixUnicodeFilter, + HardFilter, + IntentClassifier, + MessageSanityFilter, + QualityPreprocessor, + RefuseFilter, + RemoveRepeatSentencesFilter, + TokenNumFilter, + TokenSoupFilter, + ) + + dataset = _build_dataset(total=total, 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.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/legacy/eval_dualline_math.py b/cookbook/exp/legacy/eval_dualline_math.py new file mode 100644 index 000000000..bd1c183c7 --- /dev/null +++ b/cookbook/exp/legacy/eval_dualline_math.py @@ -0,0 +1,689 @@ +"""Dual-line math evaluation: baseline vs online process-checking + rubric injection. + +This is **Phase 0 of DESIGN §11.6** ("参数化 memory: 查错 LoRA"): before training any +LoRA, test the *upper bound* of the mechanism "pause every N tokens, let a strong +teacher check the partial reasoning for rubric errors, inject the found issue back +into the context, then resume". If even the strongest teacher checking online cannot +lift math accuracy, distilling that ability into a LoRA is pointless — so we gate on +this first. + +It deliberately reuses the SAME dataset loader, sampling params and answer grader as +``eval_gpqa_rag.py`` so the two lines are directly comparable: + + - **Line A — baseline** (``--mode baseline``): the student model solves each problem + in a single pass (identical to ``eval_gpqa_rag.py --mode direct``). + - **Line B — dualline** (``--mode dualline``, default): the student generates in + ``--chunk-tokens`` slices; between slices a teacher ``RubricVerifier.diagnose()`` + inspects the full reasoning so far (query + all prior response). When it reports + process issues, the finding is injected back as a first-person self-correction + (in the student's own voice) and generation resumes. + +The teacher checker is the ``llm_backup`` teacher API (no student sampler is given to +the verifier, so every check is served by the teacher — exactly the Phase-0 setup). +Configure it via the ``LLM_BACKUP_*`` env vars (see ``utils/llm_backup.py``). + +Continuation is done at the token level (crude on purpose — §11.6 says experiment +performance is not a concern): each slice re-feeds the prior ``new_input_feature`` and, +on injection, splices the tokenized note in before resuming. + +The dataset defaults to AoPS (``--dataset aops``), which auto-downloads from +ModelScope so no local data path is needed; pass ``--dataset math`` to use the +local Hendrycks MATH set instead. Both lines MUST share ``--dataset``, ``--n``, +``--target-eval`` and ``--seed`` to stay a paired comparison. + +Launch examples: + # Dual-line on 200 AoPS problems (needs LLM_BACKUP_* for the teacher checker) + LLM_BACKUP_API_KEY=sk-... LLM_BACKUP_BASE_URL=... \\ + python cookbook/exp/embedding/eval_dualline_math.py \\ + --n 200 --target-eval 200 --seed 42 + + # Paired baseline on the same subset (no checker calls) + python cookbook/exp/embedding/eval_dualline_math.py --mode baseline \\ + --n 200 --target-eval 200 --seed 42 +""" +import argparse +import copy +import json +import os +import sys +import time +from collections import defaultdict +from typing import Any, Dict, List, Optional + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import SamplingParams as TwinkleSamplingParams +from twinkle.sampler import vLLMSampler + +# Reuse the reference eval's dataset + grading + prompts verbatim so the two +# lines are measured on identical footing. +from eval_gpqa_rag import (GEN_MODEL_ID, GEN_GPU_MEM, GEN_GPUS, GEN_TEMPERATURE, + GEN_TOP_P, answers_match, build_direct_prompt, + extract_boxed, load_aops, load_math) + +# Dualline eval defaults (override via --max-model-len or DUALLINE_MAX_MODEL_LEN). +DUALLINE_DEFAULT_MAX_MODEL_LEN = int(os.environ.get('DUALLINE_MAX_MODEL_LEN', 32000)) +DUALLINE_DEFAULT_MAX_GEN_TOKENS = int( + os.environ.get('DUALLINE_MAX_GEN_TOKENS', DUALLINE_DEFAULT_MAX_MODEL_LEN)) + +# vLLM parallel: default tp=1, dp=GEN_GPUS (override with GEN_TP / keep GEN_GPUS=8). +GEN_TP = int(os.environ.get('GEN_TP', 1)) + +logger = get_logger() + +# --------------------------------------------------------------------------- +# Dual-line config +# --------------------------------------------------------------------------- +CHUNK_TOKENS = int(os.environ.get('DUALLINE_CHUNK_TOKENS', 512)) +MAX_CHECKS = int(os.environ.get('DUALLINE_MAX_CHECKS', 8)) +MAX_INJECTIONS = int(os.environ.get('DUALLINE_MAX_INJECTIONS', 3)) +# Only inject when the checker is confident enough that something is wrong. +CHECK_SCORE_FLOOR = float(os.environ.get('DUALLINE_CHECK_FLOOR', 0.6)) +# The note is written in the student's own first-person voice so, when spliced +# back in, the running model treats it as its own mid-thought self-correction +# rather than an external interruption (which tended to derail generation toward +# max-length). Kept short to limit disruption. +INJECT_TEMPLATE = ( + '\n\nWait — reviewing my reasoning above, I realize there is a problem: {issue}\n' + 'Let me correct this and continue.\n\n') + +# When context hits max_model_len (or sample fails), dump query + generation here. +OVERFLOW_DUMP_DIR = os.environ.get( + 'DUALLINE_OVERFLOW_DUMP_DIR', './output/dualline/overflow_dumps') + + +def _decode(tokenizer, ids: List[int]) -> str: + return tokenizer.decode(ids, skip_special_tokens=True) + + +def _input_ids_len(cur_inputs: Any) -> Optional[int]: + """Length of the tokenized prompt fed to vLLM on this step, if known.""" + if not cur_inputs: + return None + item = cur_inputs[0] + if isinstance(item, dict) and 'input_ids' in item: + ids = item['input_ids'] + return len(ids) if ids is not None else None + return None + + +def _dump_dualline_state( + *, + reason: str, + problem: str, + debug_idx: Optional[int], + chunk_tokens: int, + cur_inputs: Any, + gen_ids: List[int], + injected_ids: List[int], + tokenizer, + n_checks: int, + n_injections: int, + findings: List[Dict[str, Any]], + total_new: int, + finished: bool, + max_model_len: int, + error: Optional[str] = None, +) -> str: + """Persist state for post-mortem (student CoT vs checker injection). Returns path.""" + os.makedirs(OVERFLOW_DUMP_DIR, exist_ok=True) + tag = f'idx{debug_idx}' if debug_idx is not None else 'idx_unknown' + path = os.path.join( + OVERFLOW_DUMP_DIR, f'{tag}_{reason}_{int(time.time())}.json') + + partial_cot = _decode(tokenizer, gen_ids) if tokenizer and gen_ids else '' + injected_text = (_decode(tokenizer, injected_ids) + if tokenizer and injected_ids else '') + ctx_len = _input_ids_len(cur_inputs) + + payload: Dict[str, Any] = { + 'reason': reason, + 'error': error, + 'query': problem, + 'debug_idx': debug_idx, + 'gen_token_count': len(gen_ids), + 'injected_token_count': len(injected_ids), + 'context_input_ids_len': ctx_len, + 'max_model_len': max_model_len, + 'chunk_tokens': chunk_tokens, + 'total_new': total_new, + 'n_checks': n_checks, + 'n_injections': n_injections, + 'findings': findings, + 'finished': finished, + 'partial_cot': partial_cot, + 'injected_text': injected_text, + 'partial_cot_chars': len(partial_cot), + 'context_is_message_prompt': ctx_len is None, + } + with open(path, 'w', encoding='utf-8') as f: + json.dump(payload, f, ensure_ascii=False, indent=2) + cot_path = path.replace('.json', '_partial_cot.txt') + with open(cot_path, 'w', encoding='utf-8') as f: + f.write(partial_cot) + sys.stderr.write(f'[dualline] overflow dump -> {path}\n') + return path + +# --------------------------------------------------------------------------- +# Teacher checker (Phase-0: pure teacher via llm_backup) +# --------------------------------------------------------------------------- +def _build_checker(): + """RubricVerifier with no student sampler -> every diagnose() hits the teacher. + + Uses a fixed, math-oriented process rubric so we do not spend a rubric- + generation call per slice (the segment here is a partial CoT, not a finished + trajectory). Falls back to auto-generated rubrics if fixed_rubric is cleared. + """ + from twinkle_agentic.verifier import RubricVerifier + from twinkle_agentic.verifier.rubric_verifier import RubricItem + + fixed = [ + RubricItem('The reasoning contains no arithmetic or algebraic error so far', + is_hard=True), + RubricItem('Each step follows logically from the previous ones', is_hard=True), + RubricItem('No formula or theorem is misstated or misapplied', is_hard=True), + RubricItem('The approach is on track to answer the actual question asked', + is_hard=False), + RubricItem('No step contradicts an earlier established fact', is_hard=False), + ] + return RubricVerifier(fixed_rubric=fixed, gate=True) + + +def _checker_available() -> bool: + return bool(os.environ.get('LLM_BACKUP_API_KEY') + or os.environ.get('LLM_BACKUP_BASE_URL') + or os.environ.get('OPENAI_API_KEY')) + + +def _diagnose_partial(checker, problem: str, partial_cot: str): + """Run the teacher checker on the reasoning so far; return (issue_or_None, detail). + + ``partial_cot`` is the FULL reasoning generated so far (all prior chunks plus + any self-corrections already spliced in), not just the latest slice, so the + teacher judges the whole derivation in context. We label it as in-progress so + it grades correctness of the steps rather than penalizing the absence of a + final answer. + """ + seg_content = ( + '[The following is the full reasoning so far, still in progress and not ' + 'yet complete. Judge only whether the reasoning up to this point is ' + 'mathematically correct; do not expect a final answer here.]\n\n' + partial_cot) + seg = {'messages': [ + {'role': 'user', 'content': problem}, + {'role': 'assistant', 'content': seg_content}, + ]} + try: + detail = checker.diagnose(seg, query=problem) + except Exception as exc: + logger.warning(f'[dualline] checker error: {exc}') + return None, None + if detail.overall_ok: + return None, detail + if detail.scalar >= CHECK_SCORE_FLOOR: + # Checker leans "mostly fine"; don't disrupt on a marginal signal. + return None, detail + fails = [it for it in detail.items if not it.verdict] + if not fails: + return None, detail + # Prefer a fix if the checker gave one; else the reason. + parts = [] + for it in fails[:2]: + msg = it.fix or it.reason + if msg: + parts.append(msg) + issue = ' '.join(parts).strip() or detail.summary + return (issue or None), detail + + +def _pad_batch_for_dp(items: List[Any], gen_dp: int) -> List[Any]: + """``slice_dp`` needs batch len >= DP world size (every rank gets work). + + Only kicks in on the tail rounds when fewer than ``gen_dp`` problems are + still active; the padded replicas are dropped by the caller. + """ + if gen_dp <= 1 or not items or len(items) >= gen_dp: + return items + pad = [copy.deepcopy(items[-1]) for _ in range(gen_dp - len(items))] + return items + pad + + +class _DualState: + """Per-problem generation state for the batched dualline loop. + + All problems advance together, one ``chunk_tokens`` slice per round. A + problem stays *active* until it emits EOS, hits ``max_gen_tokens``, would + overflow ``max_model_len``, or a sample call fails. Because the problems + share every round's ``sampler.sample`` call, the vLLM engine batches them + (and, with dp>1, spreads them across ranks) instead of running one at a + time. + """ + + __slots__ = ('idx', 'problem', 'cur_input', 'gen_ids', 'injected_ids', + 'n_checks', 'n_injections', 'findings', 'total_new', + 'finished', 'stopped_reason', 'context_input_ids_len', + 'pending_partial_cot', 'prompt_len') + + def __init__(self, idx: int, problem: str, prompt: Any): + self.idx = idx + self.problem = problem + self.cur_input: Any = prompt # str prompt (round 0) or input_feature + self.gen_ids: List[int] = [] # student-generated token ids only + self.injected_ids: List[int] = [] # spliced-in ids (excluded from answer) + self.n_checks = 0 + self.n_injections = 0 + self.findings: List[Dict[str, Any]] = [] + self.total_new = 0 + self.finished = False + self.stopped_reason: Optional[str] = None + self.context_input_ids_len: Optional[int] = None + self.pending_partial_cot: Optional[str] = None + self.prompt_len: Optional[int] = None # token len of the fixed prompt prefix + + def cur_input_len(self) -> Optional[int]: + item = self.cur_input + if isinstance(item, dict) and 'input_ids' in item: + ids = item['input_ids'] + return len(ids) if ids is not None else None + return None + + def result(self, tokenizer) -> Dict[str, Any]: + if self.context_input_ids_len is None: + self.context_input_ids_len = self.cur_input_len() + return { + 'text': _decode(tokenizer, self.gen_ids), + 'finished': self.finished, + 'stopped_reason': self.stopped_reason, + 'context_input_ids_len': self.context_input_ids_len, + 'n_checks': self.n_checks, + 'n_injections': self.n_injections, + 'findings': self.findings, + 'gen_tokens': len(self.gen_ids), + } + + +def _dump_state_obj(st: '_DualState', tokenizer, chunk_tokens: int, + max_model_len: int, reason: str, error: str) -> None: + _dump_dualline_state( + reason=reason, + problem=st.problem, + debug_idx=st.idx, + chunk_tokens=chunk_tokens, + cur_inputs=[st.cur_input], + gen_ids=st.gen_ids, + injected_ids=st.injected_ids, + tokenizer=tokenizer, + n_checks=st.n_checks, + n_injections=st.n_injections, + findings=st.findings, + total_new=st.total_new, + finished=st.finished, + max_model_len=max_model_len, + error=error, + ) + + +# --------------------------------------------------------------------------- +# Batched token-level segmented generation with mid-stream injection +# --------------------------------------------------------------------------- +def run_dualline_batch(sampler, tokenizer, problems: List[str], checker, + base_params: TwinkleSamplingParams, + chunk_tokens: int, + max_model_len: int, + max_gen_tokens: int, + gen_dp: int = 1, + diagnose_workers: int = 8) -> List[Dict[str, Any]]: + """Advance every problem in lock-step slices, sharing one sampler call/round. + + Each round: (1) preflight-drop any problem that would overflow the context, + (2) one ``sampler.sample`` over all still-active problems (vLLM batches + + spreads over dp ranks), (3) for the length-capped ones, run the teacher + diagnoses concurrently and splice injections, then loop. + + Returns per-problem result dicts in the original ``problems`` order. + """ + from concurrent.futures import ThreadPoolExecutor + + chunk_params = TwinkleSamplingParams( + max_tokens=chunk_tokens, temperature=base_params.temperature, + top_p=base_params.top_p, num_samples=1) + + states = [_DualState(i, p, build_direct_prompt(p)) + for i, p in enumerate(problems)] + active = list(states) + round_no = 0 + + while active: + round_no += 1 + + # (1) Preflight: drop problems that would overflow the context window, + # and those that already reached the generation-token cap. + survivors: List[_DualState] = [] + for st in active: + if st.total_new >= max_gen_tokens: + st.stopped_reason = st.stopped_reason or 'max_gen_tokens' + continue + ctx_len = st.cur_input_len() + if ctx_len is not None and ctx_len + chunk_tokens >= max_model_len: + st.context_input_ids_len = ctx_len + st.stopped_reason = 'context_full' + _dump_state_obj( + st, tokenizer, chunk_tokens, max_model_len, + reason='preflight_context_full', + error=(f'context len {ctx_len} + chunk {chunk_tokens} ' + f'>= max_model_len {max_model_len}')) + continue + survivors.append(st) + active = survivors + if not active: + break + + # (2) One shared sampler call over all active problems. On tail rounds + # with fewer active problems than dp ranks, pad to keep slice_dp happy + # and drop the padded responses. The context-overflow preflight above + # guarantees every input still fits, so a length-capped slice should + # never raise here; let any genuine engine error propagate instead of + # masking it as a whole-round failure. + batch_inputs = [st.cur_input for st in active] + padded = _pad_batch_for_dp(batch_inputs, gen_dp) + responses = sampler.sample(padded, chunk_params) + responses = responses[:len(active)] + + # (3) Consume each problem's slice; queue the ones needing a check. + to_diagnose: List[_DualState] = [] + next_active: List[_DualState] = [] + for st, resp in zip(active, responses): + seq = resp.sequences[0] if resp and resp.sequences else None + if seq is None: + st.stopped_reason = st.stopped_reason or 'empty_response' + continue + st.gen_ids.extend(seq.tokens) + st.total_new += len(seq.tokens) + st.cur_input = seq.new_input_feature + if st.prompt_len is None: + # Fixed prompt prefix = everything before this round's generation. + st.prompt_len = len(st.cur_input['input_ids']) - len(seq.tokens) + + if seq.stop_reason != 'length': + st.finished = True # EOS / stop -> done + continue + if st.n_checks >= MAX_CHECKS or not checker: + next_active.append(st) # keep generating, no more checks + continue + # Diagnose the FULL reasoning generated so far (all prior chunks plus + # any self-corrections already spliced in), so the teacher judges the + # whole derivation in context rather than an isolated tail slice. + st.pending_partial_cot = _decode( + tokenizer, st.cur_input['input_ids'][st.prompt_len:]) + st.n_checks += 1 + to_diagnose.append(st) + + # Concurrent teacher diagnoses for this round's length-capped problems. + if to_diagnose: + def _run(st: _DualState): + return st, _diagnose_partial( + checker, st.problem, st.pending_partial_cot) + workers = max(1, min(diagnose_workers, len(to_diagnose))) + with ThreadPoolExecutor(max_workers=workers) as ex: + for st, (issue, _detail) in ex.map(_run, to_diagnose): + st.pending_partial_cot = None + if issue and st.n_injections < MAX_INJECTIONS: + note = INJECT_TEMPLATE.format(issue=issue) + note_ids = tokenizer.encode(note, add_special_tokens=False) + feat = dict(st.cur_input) + feat['input_ids'] = list(feat['input_ids']) + note_ids + if 'labels' in feat: + feat['labels'] = list(feat['labels']) + note_ids + st.cur_input = feat + st.injected_ids.extend(note_ids) + st.n_injections += 1 + st.findings.append( + {'at_token': st.total_new, 'issue': issue}) + next_active.append(st) + + active = next_active + n_done = sum(1 for s in states if s.finished or s.stopped_reason) + sys.stderr.write( + f'[dualline] round {round_no}: active={len(active)} ' + f'done={n_done}/{len(states)}\n') + + return [st.result(tokenizer) for st in states] + + +def _load_tokenizer(model_id: str): + """Load the tokenizer from ModelScope (matches the vLLM sampler source). + + The box runs offline, so ``transformers.AutoTokenizer`` (which resolves via + the HF hub) fails with ``Network is unreachable``. ModelScope's AutoTokenizer + downloads/reads from the ModelScope cache instead — the same place the vLLM + sampler already pulled the model from. Falls back to transformers only if the + ModelScope path is unavailable. + """ + try: + from modelscope import AutoTokenizer as MSAutoTokenizer + return MSAutoTokenizer.from_pretrained(model_id, trust_remote_code=True) + except Exception as exc: + sys.stderr.write(f'[dualline] modelscope tokenizer load failed ({exc}); ' + f'falling back to transformers\n') + from transformers import AutoTokenizer + return AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- +def main(): + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--mode', choices=['baseline', 'dualline'], default='dualline') + p.add_argument('--dataset', choices=['aops', 'math'], default='aops', + help='Evaluation dataset. "aops" (default) auto-downloads from ' + 'ModelScope (no local path needed); "math" reads local ' + 'MATH_DATA_DIR, stratified by difficulty level.') + p.add_argument('--math-split', default='test') + p.add_argument('--per-level', type=int, default=0, + help='MATH only: problems per level. 0 => --n split across levels.') + p.add_argument('--n', type=int, default=32, + help='Pool size sampled from the dataset (MATH is stratified ' + 'by level; AoPS is a flat shuffle).') + p.add_argument('--target-eval', type=int, default=32, + help='Stop after this many problems are evaluated (0 = all sampled).') + p.add_argument('--max-model-len', type=int, default=DUALLINE_DEFAULT_MAX_MODEL_LEN, + help='vLLM max_model_len / template max_length (default 32000).') + p.add_argument('--max-gen-tokens', type=int, default=DUALLINE_DEFAULT_MAX_GEN_TOKENS, + help='Cap total generated tokens per problem (default: same as ' + 'max-model-len / DUALLINE_MAX_GEN_TOKENS).') + p.add_argument('--chunk-tokens', type=int, default=CHUNK_TOKENS, + help='Generate this many tokens between checker pauses.') + p.add_argument('--batch-size', type=int, default=16, + help='Baseline mode batch size. Dualline runs all problems ' + 'concurrently (one shared sampler call per slice-round).') + p.add_argument('--diagnose-workers', type=int, + default=int(os.environ.get('DUALLINE_DIAGNOSE_WORKERS', 8)), + help='Concurrency for teacher diagnose() calls within a round.') + 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/dualline/{args.dataset}_{args.mode}_results.jsonl' + + is_dual = (args.mode == 'dualline') + if is_dual and not _checker_available(): + sys.stderr.write( + '[dualline] ERROR: --mode dualline needs a teacher checker but no ' + 'LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / OPENAI_API_KEY is set.\n' + ' Set them, or run --mode baseline for the paired baseline.\n') + sys.exit(1) + + 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) + if args.target_eval > 0: + records = records[:args.target_eval] + max_model_len = args.max_model_len + max_gen_tokens = args.max_gen_tokens + sys.stderr.write( + f'[dualline] evaluating {len(records)} problems ' + f'(mode={args.mode}, dataset={args.dataset}, ' + f'max_model_len={max_model_len}, max_gen_tokens={max_gen_tokens})\n') + + device_groups = [ + DeviceGroup(name='sampler', ranks=list(range(GEN_GPUS)), + device_type='GPU', gpus_per_worker=GEN_TP), + ] + if GEN_GPUS % GEN_TP != 0: + raise ValueError(f'GEN_GPUS ({GEN_GPUS}) must be divisible by GEN_TP ({GEN_TP})') + gen_dp = GEN_GPUS // GEN_TP + gen_mesh = DeviceMesh.from_sizes(world_size=GEN_GPUS, dp_size=gen_dp, tp_size=GEN_TP) + 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': max_model_len, + 'tensor_parallel_size': GEN_TP, + }, + device_mesh=gen_mesh, + remote_group='sampler', + ) + sampler.set_template('Qwen3_5Template', model_id=GEN_MODEL_ID, + enable_thinking=True, max_length=max_model_len) + sys.stderr.write( + f'[dualline] vLLM sampler ready (model={GEN_MODEL_ID}, ' + f'tp={GEN_TP}, dp={gen_dp})\n') + + gen_params = TwinkleSamplingParams( + max_tokens=max_gen_tokens, temperature=GEN_TEMPERATURE, + top_p=GEN_TOP_P, num_samples=1) + + checker = None + tokenizer = None + if is_dual: + checker = _build_checker() + tokenizer = _load_tokenizer(GEN_MODEL_ID) + sys.stderr.write('[dualline] teacher checker ready (llm_backup teacher)\n') + + correct = 0 + total = 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') + + def _grade_and_log(rec, idx, raw_output, extra=None): + nonlocal correct, total + predicted = extract_boxed(raw_output) + is_correct = answers_match(predicted, rec['reference_answer']) + if is_correct: + correct += 1 + total += 1 + debug_rec = { + 'idx': idx, + 'reference_answer': rec['reference_answer'], + 'predicted': predicted, + 'is_correct': is_correct, + 'problem': rec['problem'], + 'model_output': raw_output, + } + if rec.get('level'): + debug_rec['level'] = rec['level'] + if rec.get('type'): + debug_rec['type'] = rec['type'] + if extra: + debug_rec.update(extra) + debug_records.append(debug_rec) + out_f.write(json.dumps(debug_rec, ensure_ascii=False) + '\n') + out_f.flush() + + if is_dual: + problems = [rec['problem'] for rec in records] + results = run_dualline_batch( + sampler, tokenizer, problems, checker, gen_params, + args.chunk_tokens, max_model_len, max_gen_tokens, + gen_dp=gen_dp, diagnose_workers=args.diagnose_workers) + for idx, (rec, result) in enumerate(zip(records, results)): + _grade_and_log(rec, idx, result['text'], extra={ + 'n_checks': result['n_checks'], + 'n_injections': result['n_injections'], + 'findings': result['findings'], + 'finished': result['finished'], + 'stopped_reason': result.get('stopped_reason'), + 'context_input_ids_len': result.get('context_input_ids_len'), + 'gen_tokens': result['gen_tokens'], + }) + stop_tag = (f' stop={result["stopped_reason"]}' + if result.get('stopped_reason') else '') + sys.stderr.write( + f' [idx {idx}] correct={debug_records[-1]["is_correct"]} ' + f'gen={result["gen_tokens"]} checks={result["n_checks"]} ' + f'inj={result["n_injections"]}{stop_tag}\n') + acc = correct / total if total else 0 + sys.stderr.write( + f'[dualline] batched eval done: acc={acc:.4f} ({correct}/{total})\n') + else: + import re + for batch_start in range(0, len(records), args.batch_size): + batch = records[batch_start:batch_start + args.batch_size] + prompts = [build_direct_prompt(r['problem']) for r in batch] + if gen_dp > 1 and len(prompts) < gen_dp: + prompts = _pad_batch_for_dp(prompts, gen_dp) + pad_n = len(prompts) - len(batch) + else: + pad_n = 0 + responses = sampler.sample(prompts, gen_params) + if pad_n: + responses = responses[:len(batch)] + 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 = re.sub(r'<\|[^|]+\|>', '', seq.decoded or '').rstrip() + _grade_and_log(rec, batch_start + i, raw_output) + acc = correct / total if total else 0 + sys.stderr.write(f' [{total}/{len(records)}] acc={acc:.4f} ' + f'({correct}/{total})\n') + + overall = correct / total if total else 0 + print(f'\n{"=" * 60}') + print(f'MATH dual-line — mode={args.mode}, model={GEN_MODEL_ID}') + print(f' n={total}, seed={args.seed}, chunk_tokens={args.chunk_tokens}, ' + f'max_model_len={max_model_len}') + print(f'{"=" * 60}') + print(f'Overall accuracy: {overall:.4f} ({correct}/{total})') + + if is_dual: + tot_checks = sum(r.get('n_checks', 0) for r in debug_records) + tot_inj = sum(r.get('n_injections', 0) for r in debug_records) + n_with_inj = sum(1 for r in debug_records if r.get('n_injections', 0) > 0) + print(f' checker: {tot_checks} checks, {tot_inj} injections across ' + f'{n_with_inj}/{total} problems') + n_ctx_full = sum( + 1 for r in debug_records if r.get('stopped_reason') == 'context_full') + n_sample_fail = sum( + 1 for r in debug_records if r.get('stopped_reason') == 'sample_failed') + n_unfinished = sum(1 for r in debug_records if not r.get('finished')) + print(f' length: context_full={n_ctx_full}/{total}, ' + f'sample_failed={n_sample_fail}/{total}, ' + f'unfinished(no EOS)={n_unfinished}/{total}') + + if any(r.get('level') for r in debug_records): + per = defaultdict(lambda: [0, 0]) + for r in debug_records: + lv = r.get('level', 'Unknown') + per[lv][1] += 1 + if r['is_correct']: + per[lv][0] += 1 + print('\nPer-level accuracy:') + for lv in sorted(per.keys()): + c, t = per[lv] + print(f' {lv:>10}: {c/t:.4f} ({c}/{t})') + + out_f.close() + print(f'\n[output] {len(debug_records)} records saved to {args.output}') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/legacy/eval_gpqa_rag.py b/cookbook/exp/legacy/eval_gpqa_rag.py new file mode 100644 index 000000000..7954e5fd2 --- /dev/null +++ b/cookbook/exp/legacy/eval_gpqa_rag.py @@ -0,0 +1,1547 @@ +"""Math evaluation: direct vs RAG-augmented with Qwen3.5-4B. + +Datasets (``--dataset``): + - ``math`` (default): MATH (Hendrycks), stratified by difficulty (Level 1-5) + so RAG gain can be plotted against difficulty. + - ``aops``: AoPS competition problems (metadata.boxed only). + +Modes (``--mode``): + - ``direct``: The model solves problems directly (4 GPUs, TP=4). + - ``rag`` (default): Retrieve top-k thinking traces from LanceDB, condense + them (API qwen3.7-max), inject as 1-shot examples, then solve + (8 GPUs: DP=4 embedding + TP=4 vLLM). + +Defaults implement **raw RAG on MATH**: ``--dataset math --mode rag --condense`` +with hint filtering OFF. The API condenser needs ``COMPRESS_API_KEY`` (or a +local condenser via ``EVAL_CONDENSER_GPUS``); otherwise pass ``--no-condense``. + +Optional ``--hint`` flag (rag mode only): + After retrieval + condensing, call an API model to filter and refine the + traces — keeping only applicable methods — then inject the refined trace. + +Reference answers are the ``\\boxed{...}`` content of each solution. + +Launch examples: + # Default: raw RAG on MATH, stratified 100/level (needs COMPRESS_API_KEY) + COMPRESS_API_KEY=sk-... python cookbook/exp/embedding/eval_gpqa_rag.py + + # Paired direct baseline on the same MATH subset + python cookbook/exp/embedding/eval_gpqa_rag.py --mode direct + + # Raw RAG without condenser (inject raw retrieved traces) + python cookbook/exp/embedding/eval_gpqa_rag.py --no-condense + + # Add hint filtering back on top of condensing + COMPRESS_API_KEY=sk-... python cookbook/exp/embedding/eval_gpqa_rag.py --hint + + # Fall back to the old AoPS dataset + python cookbook/exp/embedding/eval_gpqa_rag.py --dataset aops +""" +import argparse +import json +import os +import random +import re +import sys +import threading +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Any, Dict, List, Optional + +import numpy as np +import torch + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import SamplingParams as TwinkleSamplingParams +from twinkle.loss import InfonceLoss +from twinkle.model import TransformersModel +from twinkle.processor import InputProcessor +from twinkle.sampler import vLLMSampler +from twinkle.template import Qwen3_5Template +from twinkle_agentic.protocol.openai import OpenAI as OpenAIClient + +logger = get_logger() + +# -- Condenser config ---------------------------------------------------------- +CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') +CONDENSE_API_KEY = os.environ.get('COMPRESS_API_KEY', '') +CONDENSE_BASE_URL = os.environ.get('COMPRESS_BASE_URL', 'https://dashscope.aliyuncs.com/compatible-mode/v1') +CONDENSE_API_MODEL = os.environ.get('COMPRESS_MODEL', 'qwen3.7-max') +CONDENSE_API_CONCURRENCY = int(os.environ.get('API_CONCURRENCY', 32)) +CONDENSE_API_MIN_INTERVAL = float(os.environ.get('API_MIN_INTERVAL', 0.1)) +CONDENSE_TEMPERATURE = 0.2 +CONDENSE_MAX_TOKENS = 8192 + +# -- Hint analysis config ------------------------------------------------------ +HINT_ANALYSIS_MAX_TOKENS = int(os.environ.get('HINT_ANALYSIS_MAX_TOKENS', 2000)) +HINT_ANALYSIS_TEMPERATURE = 0.2 + +HINT_ANALYSIS_SYSTEM = ( + 'You are a mathematical reasoning trace filter. ' + 'Given a target problem and reasoning traces retrieved from SIMILAR (but different) problems, ' + 'your task is to FILTER and REFINE the traces into a clean reference.\n\n' + 'Rules:\n' + '1. KEEP: solution steps, methods, formulas, techniques, and key insights ' + 'that are directly applicable to solving the target problem.\n' + '2. REMOVE: problem-specific numeric calculations that do not transfer, ' + 'dead-end explorations, irrelevant approaches, verbose restatements, ' + 'and any content that would mislead the solver on the target problem.\n' + '3. Output the refined trace directly as actionable solution steps. ' + 'Preserve the original mathematical expressions and step structure.\n' + '4. Do NOT solve the target problem. Do NOT add your own solutions or commentary.\n' + '5. Do NOT output the answer to either problem.\n' + '6. If the traces are entirely irrelevant, output exactly: "No applicable methods."' +) + +HINT_ANALYSIS_USER = ( + '## Target Problem\n{query}\n\n' + '## Retrieved Reasoning Traces\n{thinking}' +) + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- +# -- Gen/Embed config --------------------------------------------------------- +GEN_MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3.5-4B') +EMBED_MODEL_ID = os.environ.get( + 'EMBED_MODEL_ID', 'output.oldemb/embedding_full_transformers/last-checkpoint') + +GEN_GPUS = int(os.environ.get('GEN_GPUS', 8)) +EMB_GPUS = int(os.environ.get('EMB_GPUS', 2)) +EMBED_MAX_LENGTH = int(os.environ.get('EMBED_MAX_LENGTH', 20000)) + +GEN_GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.85)) +GEN_MAX_MODEL_LEN = int(os.environ.get('GEN_MAX_MODEL_LEN', 65536)) +GEN_MAX_TOKENS = int(os.environ.get('GEN_MAX_TOKENS', 65536)) +GEN_TEMPERATURE = float(os.environ.get('GEN_TEMPERATURE', 0.6)) +GEN_TOP_P = float(os.environ.get('GEN_TOP_P', 0.95)) + +AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') +MATH_DATA_DIR = os.environ.get('MATH_DATA_DIR', './output/math_data/MATH') + + +# --------------------------------------------------------------------------- +# Condenser prompts & validation +# --------------------------------------------------------------------------- + +COMPRESS_SYSTEM = """\ +You are a reasoning-trace condenser. Given a verbose reasoning trace, \ +extract the TRANSFERABLE KNOWLEDGE as an EXECUTABLE SOLUTION SKELETON \ +that would help a reader solve SIMILAR problems in the same domain. + +Your output is the ENTIRE useful content — there is no expansion tool, no second pass. \ +The reader will apply this knowledge to a DIFFERENT problem, so focus on what transfers. + +Principles: +1. OUTPUT AN EXECUTABLE STEP CHAIN: numbered steps that a solver can directly follow. \ +Each step should state WHAT to do and HOW (with the formula/technique), not just \ +name the concept. +2. INCLUDE FULL FORMULAS: theorems, identities, inequalities — state each \ +with its COMPLETE MATHEMATICAL EXPRESSION, not just the name. +3. STATE APPLICABILITY: what structural features of a problem signal that this \ +approach works (e.g. "when the constraint is a sum of squares"). +4. PRESERVE KEY INSIGHTS: the non-obvious ideas or tricks that make the approach \ +work — the things a solver would NOT think of without guidance. +5. REMOVE: problem-specific numeric calculations, dead-end explorations, \ +hesitations, verbose restatements, and trivial arithmetic. +6. FORMAT: Start with a one-line "Applicability" statement, then numbered steps, \ +then key formulas. Keep it concise and actionable. +7. NO meta-commentary about the compression process. NO preamble. +""" + +COMPRESS_USER = ( + '## Reader Problem (context only — do NOT solve it)\n{query}\n\n' + '## Reasoning Trace to Condense\n{text}') + + +def _is_truncated_compression(text: str) -> bool: + if not text or not text.strip(): + return True + lines = [l.strip() for l in text.strip().splitlines() if l.strip()] + if len(lines) < 3: + return True + last_line = lines[-1] + # Truncated if last line looks incomplete (no terminal punctuation/formula) + if last_line and last_line[-1] not in '.。!！)）]】}\\$': + # Allow lines ending with numbers, boxed answers, etc. + if not re.search(r'\d$|\\boxed|\$|\)$', last_line): + return True + return False + + +# -- API rate limiter ---------------------------------------------------------- +_api_semaphore = threading.Semaphore(CONDENSE_API_CONCURRENCY) +_api_bucket_lock = threading.Lock() +_api_tokens = [float(CONDENSE_API_CONCURRENCY)] +_api_last_refill = [time.monotonic()] + + +def _api_throttle(): + _api_semaphore.acquire() + wait = 0.0 + try: + with _api_bucket_lock: + now = time.monotonic() + elapsed = now - _api_last_refill[0] + refill = elapsed / CONDENSE_API_MIN_INTERVAL + _api_tokens[0] = min(float(CONDENSE_API_CONCURRENCY), _api_tokens[0] + refill) + _api_last_refill[0] = now + if _api_tokens[0] >= 1.0: + _api_tokens[0] -= 1.0 + else: + wait = (1.0 - _api_tokens[0]) * CONDENSE_API_MIN_INTERVAL + _api_tokens[0] = 0.0 + finally: + _api_semaphore.release() + if wait > 0: + time.sleep(wait) + + +def _api_condense_single(api_client: OpenAIClient, messages: List[Dict]) -> Optional[str]: + _api_throttle() + trajectory = {'messages': messages} + sp = TwinkleSamplingParams(temperature=CONDENSE_TEMPERATURE, max_tokens=CONDENSE_MAX_TOKENS) + try: + reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) + except Exception as exc: + logger.warning(f'[condense-api] error: {exc}') + return None + content = (reply.get('content') or '').strip() + if not content: + return None + m = re.match(r'^```[a-zA-Z]*\n(.*?)\n```\s*$', content, re.DOTALL) + if m: + content = m.group(1).strip() + return content + + +def _api_hint_analysis_batch( + api_client: OpenAIClient, + problems: List[str], + condensed_examples: List[List[Dict[str, str]]], +) -> List[Optional[str]]: + """Call API to pre-analyze RAG relevance for each problem.""" + _MAX_HINT_INPUT = 8000 + results: List[Optional[str]] = [None] * len(problems) + tasks = [] + for i, prob in enumerate(problems): + if not condensed_examples[i]: + continue + traces = [ex.get('thinking', '') for ex in condensed_examples[i]] + merged_thinking = '\n---\n'.join(traces) + if len(merged_thinking) > _MAX_HINT_INPUT: + merged_thinking = merged_thinking[:_MAX_HINT_INPUT] + '\n[...truncated]' + user_msg = HINT_ANALYSIS_USER.format(query=prob, thinking=merged_thinking) + msgs = [ + {'role': 'system', 'content': HINT_ANALYSIS_SYSTEM}, + {'role': 'user', 'content': user_msg}, + ] + tasks.append((i, msgs)) + + if not tasks: + return results + + def _call_one(idx, msgs): + _api_throttle() + try: + trajectory = {'messages': msgs} + sp = TwinkleSamplingParams( + temperature=HINT_ANALYSIS_TEMPERATURE, + max_tokens=HINT_ANALYSIS_MAX_TOKENS) + reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) + content = (reply.get('content') or '').strip() + # Treat "No applicable methods." as empty (will trigger fallback) + if not content or content == 'No applicable methods.': + return idx, None + return idx, content + except Exception as exc: + logger.warning(f'[hint-analysis] error for idx={idx}: {exc}') + return idx, None + + with ThreadPoolExecutor(max_workers=min(len(tasks), CONDENSE_API_CONCURRENCY)) as pool: + futs = [pool.submit(_call_one, idx, msgs) for idx, msgs in tasks] + for fut in as_completed(futs): + idx, analysis = fut.result() + results[idx] = analysis + + n_success = sum(1 for r in results if r) + logger.info(f'[hint-analysis] completed {n_success}/{len(tasks)} analyses') + return results + + +# --------------------------------------------------------------------------- +# LLM-based decontamination +# --------------------------------------------------------------------------- + +_DECONTAM_JUDGE_PROMPT = ( + 'We are building a RAG-augmented math training system. Problem A is the test ' + 'question; Problem B was retrieved from a knowledge base.\n' + 'Answer YES only if A and B are essentially the SAME specific problem — ' + 'i.e. solving B directly gives you A\'s answer (just different wording/notation/' + 'format/negation).\n' + 'Answer NO if they merely share the same method/topic but have different ' + 'specific values, equations, or geometric configurations — learning B\'s ' + 'approach still requires independent work to solve A.\n' + 'Problem A: {prob_a}\n' + 'Problem B: {prob_b}\n' + 'Answer only YES or NO.' +) + + +def _llm_judge_same_problem( + api_client: OpenAIClient, pairs: List[tuple], +) -> List[bool]: + """Batch LLM judge: are (problem_a, problem_b) the same problem? + + Returns list of bools (True = same problem = should filter). + """ + if not pairs or not api_client: + return [False] * len(pairs) + + results = [False] * len(pairs) + + def _judge_one(idx, pa, pb): + prompt = _DECONTAM_JUDGE_PROMPT.format(prob_a=pa, prob_b=pb) + msgs = [{'role': 'user', 'content': prompt}] + _api_throttle() + try: + trajectory = {'messages': msgs} + sp = TwinkleSamplingParams(temperature=0.1, max_tokens=8) + reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) + answer = (reply.get('content') or '').strip().upper() + return idx, 'YES' in answer + except Exception: + return idx, False + + with ThreadPoolExecutor(max_workers=min(len(pairs), CONDENSE_API_CONCURRENCY)) as pool: + futs = [pool.submit(_judge_one, i, pa, pb) for i, (pa, pb) in enumerate(pairs)] + for fut in as_completed(futs): + idx, is_same = fut.result() + results[idx] = is_same + return results + + +def _llm_decontaminate( + api_client: OpenAIClient, + problems: List[str], + all_examples: List[List[Dict[str, str]]], +) -> List[List[Dict[str, str]]]: + """Apply LLM-based decontamination: remove retrievals judged as same problem.""" + judge_pairs = [] # (qi, ret_idx, prob_a, prob_b) + for qi, exs in enumerate(all_examples): + for ri, ex in enumerate(exs): + judge_pairs.append((qi, ri, problems[qi], ex.get('query', ''))) + + if not judge_pairs: + return all_examples + + pairs_input = [(pa, pb) for _, _, pa, pb in judge_pairs] + verdicts = _llm_judge_same_problem(api_client, pairs_input) + to_remove = set() + for vi, (qi, ri, _, _) in enumerate(judge_pairs): + if verdicts[vi]: + to_remove.add((qi, ri)) + + if to_remove: + logger.info(f'[decontam-llm] filtered {len(to_remove)} same-problem retrievals') + for qi in range(len(all_examples)): + all_examples[qi] = [ + ex for ri, ex in enumerate(all_examples[qi]) + if (qi, ri) not in to_remove + ] + return all_examples + + +def condense_traces( + examples_batch: List[List[Dict[str, str]]], + problems: List[str], + api_client: OpenAIClient, + condenser_sampler=None, + compress_params=None, + special_tokens: set = None, + max_output_len: int = 2000, + dp_size: int = 1, +) -> List[List[Dict[str, str]]]: + """Compress retrieved thinking traces with query-aware condenser. + + Primary: local vLLM condenser (if provided). + Fallback: API condenser. + Final fallback: raw trace truncated to max_output_len. + """ + result: List[List[Dict[str, str]]] = [] + # Flatten all (batch_idx, ex_idx, problem, example) for batch processing + tasks = [] + for bi, (exs, prob) in enumerate(zip(examples_batch, problems)): + for ei, ex in enumerate(exs): + tasks.append((bi, ei, prob, ex)) + + if not tasks: + return [[] for _ in examples_batch] + + # Build condense prompts (aligned with make_embedding_dataset.py hard path) + prompts = [] + for _, _, prob, ex in tasks: + user_msg = COMPRESS_USER.format(query=prob, text=ex['thinking']) + prompts.append([{'role': 'system', 'content': COMPRESS_SYSTEM}, + {'role': 'user', 'content': user_msg}]) + + # Phase 1: local vLLM condenser + condensed = [None] * len(tasks) + condense_sources = ['raw'] * len(tasks) + fallback_indices = [] + + if condenser_sampler is not None and compress_params is not None: + sampler_inputs = [{'messages': p} for p in prompts] + # The local vLLM sampler runs data-parallel across ``dp_size`` workers + # and requires at least one item per worker (it errors with + # "Batch too small for N workers" otherwise). Pad the batch up to a + # multiple of dp_size by repeating the last item, run, then keep only + # the first ``n_real`` responses and drop the padding. + n_real = len(sampler_inputs) + pad_size = 0 + if dp_size > 1 and n_real > 0 and n_real % dp_size != 0: + pad_size = dp_size - (n_real % dp_size) + sampler_inputs = sampler_inputs + [sampler_inputs[-1]] * pad_size + try: + responses = condenser_sampler.sample(sampler_inputs, compress_params) + except Exception as exc: + logger.warning(f'[condense] sampler error: {exc}') + responses = [None] * len(sampler_inputs) + if pad_size: + responses = responses[:n_real] + for ri, resp in enumerate(responses): + seq = resp.sequences[0] if resp and resp.sequences else None + text = '' + if seq and seq.stop_reason != 'length' and seq.decoded: + text = seq.decoded + if special_tokens: + for tok in special_tokens: + text = text.replace(tok, '') + text = text.rstrip() + if text and not _is_truncated_compression(text): + condensed[ri] = text + condense_sources[ri] = 'local' + else: + fallback_indices.append(ri) + else: + fallback_indices = list(range(len(tasks))) + + # Phase 2: API fallback + if fallback_indices and api_client: + with ThreadPoolExecutor(max_workers=CONDENSE_API_CONCURRENCY) as pool: + futures = {} + for ri in fallback_indices: + futures[pool.submit(_api_condense_single, api_client, prompts[ri])] = ri + for fut in as_completed(futures): + ri = futures[fut] + api_result = fut.result() + if api_result and not _is_truncated_compression(api_result): + condensed[ri] = api_result + condense_sources[ri] = 'api' + + # Phase 3: assemble results (fallback to raw truncation) + result = [[] for _ in examples_batch] + for ti, (bi, ei, prob, ex) in enumerate(tasks): + compressed = condensed[ti] + raw_len = len(ex['thinking']) + sim_val = ex.get('_sim', 0.0) + if compressed: + result[bi].append({'query': ex['query'], + 'thinking': _strip_condenser_markers(compressed), + '_condense_source': condense_sources[ti], + '_raw_trace_len': raw_len, '_sim': sim_val}) + else: + result[bi].append({'query': ex['query'], + 'thinking': ex['thinking'][:max_output_len], + '_condense_source': 'raw', + '_raw_trace_len': raw_len, '_sim': sim_val}) + + n_ok = sum(1 for c in condensed if c) + logger.info(f'[condense] {n_ok}/{len(tasks)} compressed ok, ' + f'{len(tasks) - n_ok} fell back to raw truncation') + return result + + +def _strip_condenser_markers(text: str) -> str: + """Light cleanup of condenser output. + + Removes any residual markdown headers or meta-lines that don't carry + solution content. Keeps numbered steps and equations intact. + """ + # Remove legacy ## headers if condenser still emits them + if '## More' in text: + text = text.split('## More', 1)[0] + text = re.sub(r'^##\s*Summary\s*\n?', '', text, flags=re.MULTILINE) + text = re.sub(r'^Topic:\s*.*\n?', '', text, flags=re.MULTILINE) + # Remove meta-commentary lines + text = re.sub(r'^\s*\(Note:.*\)\s*$', '', text, flags=re.MULTILINE) + return text.strip() + + +# --------------------------------------------------------------------------- +# Boxed answer extraction +# --------------------------------------------------------------------------- +_BOXED_RE = re.compile(r'\\boxed\s*\{') + + +def extract_boxed(text: str) -> Optional[str]: + """Extract the last \\boxed{...} content, handling nested braces.""" + if not text: + return None + last_match = None + for m in _BOXED_RE.finditer(text): + start = m.end() + depth = 1 + i = start + while i < len(text) and depth > 0: + if text[i] == '{': + depth += 1 + elif text[i] == '}': + depth -= 1 + i += 1 + if depth == 0: + last_match = text[start:i - 1].strip() + return last_match + + +def normalize_answer(ans: str) -> str: + """Normalize a math answer string for comparison.""" + if not ans: + return '' + s = str(ans).strip() + m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) + if m: + return m.group(1) + s = s.strip('$').strip() + s = s.replace('−', '-') + 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.replace(r'\,', '').replace(r'\;', '').replace(r'\!', '') + s = re.sub(r'\\(?:quad|qquad|\s)', '', s) + s = re.sub(r'\s+', '', s) + s = re.sub(r'\\sqrt([0-9A-Za-z])', r'\\sqrt{\1}', s) + s = re.sub(r'\\frac\{([^{}]+)\}([^{}\\])', r'\\frac{\1}{\2}', s) + s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) + s = re.sub(r'\{[a-zA-Z]+\}$', '', s) + s = re.sub(r'\^\\circ|\^\{\\circ\}|°|\\circ', 'deg', 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'(? 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 + frac_re = re.compile(r'^\(([^)]+)\)/\(([^)]+)\)$') + + def _eval_frac(s): + m = frac_re.match(s) + if m: + try: + return float(m.group(1)) / float(m.group(2)) + except (ValueError, ZeroDivisionError): + pass + return None + + va, vb = _eval_frac(a), _eval_frac(b) + if va is not None and vb is not None: + return abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) + return False + + +_MCQ_REF_RE = re.compile( + r'^\\(?:textbf|text|mathrm|mathbf)\{\(?([A-E])\)?\s*\}\s*(.+)$' + r'|^\(?([A-E])\)\s+(.+)$' +) +_VAR_PREFIX_RE = re.compile(r'^(?:[A-Za-z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)$', re.DOTALL) + + +def _split_mcq(ans: str): + """Split an MCQ answer into (letter, value) components.""" + 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) + 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 _strip_var_prefix(ans: str) -> str: + m = _VAR_PREFIX_RE.match((ans or '').strip()) + return m.group(1).strip() if m else (ans or '') + + +def _normalize_tuple(ans: str) -> str: + return re.sub(r'[\s()\[\]{}\\]', '', ans or '') + + +def _math_verify_equal(predicted: str, reference: str) -> bool: + try: + from math_verify import parse, verify + gold = parse(r'\boxed{' + reference + '}', parsing_timeout=1) + pred = parse(r'\boxed{' + predicted + '}', parsing_timeout=1) + return bool(gold and pred and verify(gold, pred, timeout_seconds=1)) + except Exception: + 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 or norm_p.lower() == norm_r.lower(): + return True + if _try_numeric_equal(norm_p, norm_r): + return True + + stripped_p = normalize_answer(_strip_var_prefix(predicted)) + stripped_r = normalize_answer(_strip_var_prefix(reference)) + if stripped_p and stripped_r: + if stripped_p == stripped_r or stripped_p.lower() == stripped_r.lower(): + return True + if _try_numeric_equal(stripped_p, stripped_r): + return True + + p_letter, p_value = _split_mcq(predicted) + r_letter, r_value = _split_mcq(reference) + if p_letter and r_letter and p_letter == r_letter: + return True + if (p_letter and p_value is None) and (r_letter and r_value is None): + return False + + 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 p_val.lower() == r_val.lower() or _try_numeric_equal(p_val, r_val): + return True + + for left, right in ((norm_p, norm_r), (stripped_p, stripped_r)): + tuple_l, tuple_r = _normalize_tuple(left), _normalize_tuple(right) + if ',' in tuple_l and tuple_l == tuple_r: + return True + + if '=' in norm_r and '=' not in norm_p: + if any(part and (part == norm_p or _try_numeric_equal(part, norm_p)) for part in norm_r.split('=')): + return True + if '=' in norm_p and '=' not in norm_r: + if any(part and (part == norm_r or _try_numeric_equal(part, norm_r)) for part in norm_p.split('=')): + return True + + return _math_verify_equal(stripped_p or norm_p, stripped_r or norm_r) + + +# --------------------------------------------------------------------------- +# Dataset loading +# --------------------------------------------------------------------------- + +def _load_aops_from_modelscope(): + """Download the AoPS repo natively from ModelScope and read its parquet. + + Primary loader: ``dataset_snapshot_download`` pulls the dataset repo files + (parquet) straight from the ModelScope hub WITHOUT going through the + ``datasets``/HF-filesystem path used by ``MsDataset.load`` — that path is + broken on this modelscope build (``HfFileSystem.find() got multiple values + for 'maxdepth'``). We then read the local parquet with the ``datasets`` + backend (reading local files does not trigger the HfFileSystem bug). + """ + import glob + + from datasets import Dataset as HFDataset + from modelscope.hub.snapshot_download import dataset_snapshot_download + + local = dataset_snapshot_download(AOPS_DATASET_ID) + files = sorted(glob.glob(os.path.join(local, '**', '*.parquet'), + recursive=True)) + if not files: + # Older snapshots may materialize an arrow file instead of parquet. + files = sorted(glob.glob(os.path.join(local, '**', '*train*.arrow'), + recursive=True)) + if files: + sys.stderr.write(f'[aops] modelscope snapshot arrow: {files[0]}\n') + return HFDataset.from_file(files[0]) + return None + sys.stderr.write(f'[aops] modelscope snapshot parquet: {files[0]}\n') + return HFDataset.from_parquet(files if len(files) > 1 else files[0]) + + +def load_aops(n: int, seed: int = 42) -> List[Dict[str, Any]]: + """Load AoPS boxed problems, sample n, extract reference answers. + + Uses ModelScope as the data source (native repo snapshot download). + """ + ds = None + try: + ds = _load_aops_from_modelscope() + except Exception as exc: + sys.stderr.write(f'[aops] modelscope snapshot download failed ({exc}); ' + f'trying MsDataset.load\n') + if ds is None: + from modelscope import MsDataset + ds = MsDataset.load(AOPS_DATASET_ID, split='train', + download_mode='reuse_dataset_if_exists') + boxed = [] + for row in ds: + if not row['metadata'].get('boxed'): + continue + ref = extract_boxed(row['solution']) + if not ref: + continue + boxed.append({ + 'problem': row['problem'], + 'solution': row['solution'], + 'reference_answer': ref, + 'tags': row.get('tags', []), + }) + sys.stderr.write(f'[aops] {len(boxed)} boxed problems with extractable answers\n') + rng = random.Random(seed) + rng.shuffle(boxed) + if n > 0 and n < len(boxed): + boxed = boxed[:n] + sys.stderr.write(f'[aops] sampled {n} problems\n') + return boxed + + +def load_math(n: int, seed: int = 42, split: str = 'test', + per_level: int = 0) -> List[Dict[str, Any]]: + """Load the MATH (Hendrycks) dataset from local extracted JSON files. + + Each problem's reference answer is the ``\\boxed{}`` content of its + ``solution`` (MATH solutions always end in a boxed answer). + + Sampling is *stratified by level* so every difficulty (Level 1-5) is + represented equally — required to measure how RAG gain varies with + difficulty. ``per_level`` (if >0) fixes the count per level; otherwise + ``n`` is split evenly across the 5 levels. When both are 0, all problems + are returned. The final list is shuffled with ``seed`` so index order is + stable/comparable across direct vs rag runs. + """ + import glob + root = os.path.join(MATH_DATA_DIR, split) + files = glob.glob(os.path.join(root, '*', '*.json')) + if not files: + raise FileNotFoundError( + f'[math] no problems found under {root!r}; set MATH_DATA_DIR or ' + f'extract MATH.zip there') + + by_level: Dict[str, List[Dict[str, Any]]] = {} + n_no_box = 0 + for fp in files: + try: + with open(fp, 'r', encoding='utf-8') as fin: + row = json.load(fin) + except Exception: + continue + ref = extract_boxed(row.get('solution', '')) + if not ref: + n_no_box += 1 + continue + level = row.get('level', 'Unknown') + by_level.setdefault(level, []).append({ + 'problem': row['problem'], + 'solution': row['solution'], + 'reference_answer': ref, + 'level': level, + 'type': row.get('type', ''), + }) + + total = sum(len(v) for v in by_level.values()) + sys.stderr.write( + f'[math] {total} problems with boxed answers across ' + f'{len(by_level)} levels (skipped {n_no_box} without boxed)\n') + + levels = sorted(by_level.keys()) + rng = random.Random(seed) + + # Decide how many per level. + if per_level <= 0 and n > 0: + per_level = max(1, n // max(1, len(levels))) + + sampled: List[Dict[str, Any]] = [] + for lv in levels: + pool = by_level[lv] + rng.shuffle(pool) + take = pool if per_level <= 0 else pool[:per_level] + sampled.extend(take) + sys.stderr.write(f'[math] {lv}: took {len(take)}/{len(pool)}\n') + + rng.shuffle(sampled) + sys.stderr.write(f'[math] total sampled: {len(sampled)}\n') + return sampled + + +# --------------------------------------------------------------------------- +# Prompt building +# --------------------------------------------------------------------------- + +DIRECT_SYSTEM = ( + 'You are an expert competition mathematician. Solve the following problem ' + 'step by step. Provide your final answer inside \\boxed{}.' +) + +RAG_SYSTEM = ( + 'You are an expert competition mathematician. Solve the following problem ' + 'step by step. Provide your final answer inside \\boxed{}.\n\n' + 'You will first see example problem-solving traces or skills. ' + 'Learn from the reasoning methodology demonstrated in these examples, ' + 'then thinking to solve the actual problem.' +) + +RAG_FOLLOWUP = ( + 'The above is a reference solution to a similar problem. ' + 'You may use any applicable techniques from it, or ignore it ' + 'if you find a better approach. ' + 'Solve the problem step by step and put your final answer in \\boxed{}.' +) + +HINT_FOLLOWUP = ( + 'The above are applicable solution approaches extracted from similar problems. ' + 'You may use any applicable techniques from them, or ignore them ' + 'if you find a better approach. ' + 'Solve the problem step by step and put your final answer in \\boxed{}.' +) + +# Reminder appended to the final user turn. Without this, the reasoning model +# can loop indefinitely on multiple-choice problems, oscillating between boxing +# the option letter and boxing the value (e.g. "I'll box B. I'll box 21. ...") +# and never terminating. Boxing BOTH the letter and value removes the ambiguity +# (the grader accepts either), so the model has no format decision to agonize over. +MCQ_INSTRUCTION = ( + '\n\nNote: If the problem is multiple-choice (it lists options such as ' + '(A), (B), (C), ...), put BOTH the option letter and its value in the box, ' + 'e.g. \\boxed{(B) 21}. Otherwise, box the value directly. Decide the answer ' + 'format once and do not deliberate over which form to box.' +) + + +def build_direct_prompt(problem: str) -> Dict[str, Any]: + return { + 'messages': [ + {'role': 'system', 'content': DIRECT_SYSTEM}, + {'role': 'user', 'content': problem}, + ] + } + + +def build_hint_prompt(problem: str, hint_analysis: str) -> Dict[str, Any]: + """Build prompt with pre-analyzed hint in a multi-turn conversation. + + Mirrors ``build_rag_prompt``: the hint is presented as an assistant + "extracted approaches" turn (instead of being buried in the system + prompt), followed by a user instruction that provides a clear closing + directive to solve the problem and box the answer. Keeping the final + solve/box instruction in a dedicated user turn (rather than in the + system prompt) helps the reasoning model terminate cleanly. + """ + messages: List[Dict[str, str]] = [ + {'role': 'system', 'content': DIRECT_SYSTEM}, + {'role': 'user', 'content': problem}, + {'role': 'assistant', + 'content': ('Here are applicable solution approaches extracted from ' + f'similar problems:\n\n{hint_analysis}')}, + {'role': 'user', 'content': HINT_FOLLOWUP + MCQ_INSTRUCTION}, + ] + return {'messages': messages} + + +def build_rag_prompt(problem: str, + examples: List[Dict[str, str]]) -> Dict[str, Any]: + """Approach B: multi-turn assistant format. + + The trace is presented as an assistant "retrieval" turn, followed by + a user instruction that constrains the model to use methodology only. + """ + messages: List[Dict[str, str]] = [{'role': 'system', 'content': DIRECT_SYSTEM}] + messages.append({'role': 'user', 'content': problem}) + # Build trace content from retrieved examples + trace_parts = [] + for i, ex in enumerate(examples, 1): + trace_parts.append(f'[Retrieved Example {i}]\nProblem: {ex["query"]}\n' + f'Reasoning:\n{ex["thinking"]}') + trace_text = '\n\n'.join(trace_parts) + messages.append({'role': 'assistant', + 'content': f'I found relevant reasoning traces from the knowledge base!\n\n{trace_text}'}) + messages.append({'role': 'user', 'content': RAG_FOLLOWUP + MCQ_INSTRUCTION}) + return {'messages': messages} + + +# --------------------------------------------------------------------------- +# 13-gram Jaccard decontamination +# --------------------------------------------------------------------------- + +def _normalize_for_ngram(text: str) -> str: + """Normalize text for n-gram comparison: strip LaTeX markup, lowercase.""" + text = text.lower() + text = re.sub(r'\$+', '', text) + text = re.sub(r'\\[a-z]+\{([^}]*)\}', r'\1', text) + text = re.sub(r'\\[a-z]+', ' ', text) + text = re.sub(r'[{}\\^_$]', '', text) + text = re.sub(r'\s+', ' ', text).strip() + return text + + +def _ngram_jaccard(text_a: str, text_b: str, n: int = 13) -> float: + """13-gram character-level Jaccard similarity.""" + a = _normalize_for_ngram(text_a) + b = _normalize_for_ngram(text_b) + if len(a) < n or len(b) < n: + return 0.0 + grams_a = set(a[i:i + n] for i in range(len(a) - n + 1)) + grams_b = set(b[i:i + n] for i in range(len(b) - n + 1)) + if not grams_a or not grams_b: + return 0.0 + return len(grams_a & grams_b) / len(grams_a | grams_b) + + +# --------------------------------------------------------------------------- +# Embedding / RAG helpers +# --------------------------------------------------------------------------- + +def _wrap_anchor(text: str) -> List[Dict[str, str]]: + return [ + {'role': 'user', 'content': text}, + {'role': 'assistant', 'content': 'Match the correct response here.'}, + ] + + +def get_embeddings(model: TransformersModel, template: Qwen3_5Template, + texts: List[str], dp_size: int) -> np.ndarray: + if not texts: + return np.zeros((0,), dtype=np.float32) + n = len(texts) + pad_n = (-n) % dp_size + padded = list(texts) + [' '] * pad_n if pad_n else list(texts) + features = [] + for t in padded: + feat = template.encode({'messages': _wrap_anchor(t or ' ')}) + feat['labels'] = [1] + features.append(feat) + out = model.forward_only(inputs=features, task='embedding', return_logits=True) + emb = out['embeddings'] + if isinstance(emb, torch.Tensor): + emb = emb.detach().to(torch.float32).cpu().numpy() + emb = np.asarray(emb, dtype=np.float32) + return emb[:n] if pad_n else emb + + +def retrieve_examples(tbl, query_vecs: np.ndarray, top_k: int, + use_thinking_raw: bool, sim_threshold: float = 0.0, + problems: List[str] = None, + decontam_threshold: float = 0.0, + ) -> List[List[Dict[str, str]]]: + thinking_field = 'thinking_raw' if use_thinking_raw else 'cot_compressed' + fetch_limit = top_k + 50 if decontam_threshold > 0 else top_k + n_queries = len(query_vecs) + all_examples: List[List[Dict[str, str]]] = [None] * n_queries + decontam_skipped = 0 + _decontam_lock = threading.Lock() + + def _search_one(qi: int): + nonlocal decontam_skipped + vec = query_vecs[qi] + results = ( + tbl.search(vec.astype(np.float32).tolist()) + .metric('dot') + .limit(fetch_limit) + .select(['query_raw', thinking_field, '_distance']) + .to_list() + ) + problem_text = problems[qi] if problems else '' + examples = [] + local_skipped = 0 + for r in results: + if len(examples) >= top_k: + break + sim = 1.0 - r.get('_distance', 0.0) + if sim < sim_threshold: + continue + q = r.get('query_raw', '') + t = r.get(thinking_field, '') + if not t: + continue + if decontam_threshold > 0 and problem_text and q: + ng_sim = _ngram_jaccard(problem_text, q) + if ng_sim > decontam_threshold: + local_skipped += 1 + continue + examples.append({'query': q, 'thinking': t, '_sim': round(sim, 4), + '_raw_trace_len': len(t)}) + all_examples[qi] = examples + if local_skipped: + with _decontam_lock: + decontam_skipped += local_skipped + + with ThreadPoolExecutor(max_workers=min(n_queries, 16)) as pool: + list(pool.map(_search_one, range(n_queries))) + + if decontam_skipped > 0: + logger.info(f'[decontam] skipped {decontam_skipped} leaked retrievals ' + f'(13-gram Jaccard > {decontam_threshold})') + return all_examples + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main(): + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--mode', choices=['direct', 'rag'], default='rag') + p.add_argument('--dataset', choices=['aops', 'math'], default='math', + help='Evaluation dataset. "math" = MATH (Hendrycks), ' + 'stratified by level for a difficulty-vs-gain curve.') + p.add_argument('--math-split', default='test', + help='MATH split to load (test/train).') + p.add_argument('--per-level', type=int, default=100, + help='MATH only: problems per difficulty level (default 100 ' + '-> 500 total across Level 1-5). If 0, --n is split ' + 'evenly across the 5 levels.') + p.add_argument('--n', type=int, default=0, + help='Pool size: sample this many problems (0 = all boxed). ' + 'In RAG mode with --target-eval, set this to 0 for max coverage.') + p.add_argument('--target-eval', type=int, default=0, + help='Stop after this many problems are successfully evaluated ' + '(0 = no limit, evaluate the entire sampled set — the ' + 'default, so all 500 stratified MATH problems are run). ' + 'RAG mode: counts problems with valid traces after ' + 'decontam; direct mode: ignored, evaluates all filtered.') + p.add_argument('--db-path', default='./output.oldemb/thinking_rag/lance.db') + p.add_argument('--table', default='thinking_traces') + p.add_argument('--top-k', type=int, default=1) + p.add_argument('--use-cot-compressed', action='store_true', + help='Use pre-compressed cot_compressed field instead of thinking_raw.') + p.add_argument('--sim-threshold', type=float, default=0.75, + help='Minimum cosine similarity for retrieved traces. ' + 'Traces below this are discarded at retrieval time.') + p.add_argument('--decontam-threshold', type=float, default=0.20, + help='13-gram Jaccard threshold for leak detection. ' + 'Retrieved traces above this are skipped (0=disabled).') + p.add_argument('--llm-decontam', action='store_true', default=True, + help='LLM-based decontamination (default ON): API judges whether ' + 'retrieved problem is the same as the test problem. ' + 'Applied after 13-gram decontam, before condensing. ' + 'Use --no-llm-decontam to disable.') + p.add_argument('--no-llm-decontam', dest='llm_decontam', action='store_false', + help='Disable LLM-based decontamination.') + p.add_argument('--max-trace-len', type=int, default=12000) + p.add_argument('--condense', action='store_true', default=True, + help='Enable condenser re-compression on retrieved traces ' + '(default ON). Use --no-condense to inject raw traces.') + p.add_argument('--no-condense', dest='condense', action='store_false', + help='Disable condenser; inject raw retrieved traces.') + p.add_argument('--condense-max-len', type=int, default=2000, + help='Max chars of condensed trace (fallback truncation).') + p.add_argument('--batch-size', type=int, default=16) + p.add_argument('--seed', type=int, default=42) + p.add_argument('--hint', action='store_true', default=False, + help='Enable API hint filtering on retrieved traces (default OFF; ' + 'raw RAG injects the condensed trace directly). ' + 'In rag mode: retrieve → condense → API filters trace → refined system prompt. ' + 'In direct mode: ignored (no traces to filter).') + p.add_argument('--no-hint', dest='hint', action='store_false', + help='Disable API hint filtering; inject condensed trace directly.') + p.add_argument('--problem-ids-file', default=None, + help='File listing problem indices evaluated by RAG mode. ' + 'RAG mode writes this file; direct mode reads it to ' + 'evaluate the same subset (use --no-filter to disable). ' + 'Defaults to a dataset-specific path.') + p.add_argument('--no-filter', action='store_true', + help='In direct mode, evaluate ALL sampled problems ' + 'instead of filtering to RAG subset.') + p.add_argument('--output', default=None) + args = p.parse_args() + + # Dataset-specific default paths (keeps aops and math runs from colliding). + if args.problem_ids_file is None: + args.problem_ids_file = ( + f'./output/thinking_rag/{args.dataset}_rag_problem_ids.json') + + if args.output is None: + suffix = f'{args.mode}_hint' if (args.hint and args.mode == 'rag') else args.mode + args.output = ( + f'./output/thinking_rag/{args.dataset}_{suffix}_results.jsonl') + + if args.condense and args.use_cot_compressed: + logger.warning('--condense requires thinking_raw, ignoring --use-cot-compressed') + args.use_cot_compressed = False + + if args.dataset == 'math': + records = load_math(n=args.n, seed=args.seed, split=args.math_split, + per_level=args.per_level) + else: + records = load_aops(n=args.n, seed=args.seed) + + is_rag = (args.mode == 'rag') + + # Direct mode: filter to same problems RAG evaluated (controlled comparison) + original_indices = list(range(len(records))) # track original indices + if not is_rag and not args.no_filter: + if os.path.exists(args.problem_ids_file): + with open(args.problem_ids_file) as f: + content = f.read().strip() + if content.startswith('['): + valid_indices = set(json.loads(content)) + else: + valid_indices = set(int(line) for line in content.splitlines() if line.strip()) + filtered = [(i, r) for i, r in enumerate(records) if i in valid_indices] + original_indices = [i for i, _ in filtered] + records = [r for _, r in filtered] + sys.stderr.write( + f'[direct] filtered to {len(records)} problems ' + f'from {args.problem_ids_file}\n') + else: + sys.stderr.write( + f'[direct] WARNING: {args.problem_ids_file} not found, ' + f'running all {len(records)} problems\n') + + condenser_gpus = int(os.environ.get('EVAL_CONDENSER_GPUS', 0)) if args.condense else 0 + + # Raw RAG relies on the API condenser (qwen3.7-max). Fail fast with a clear + # message if it's enabled without an API key and without a local condenser. + if is_rag and args.condense and not CONDENSE_API_KEY and condenser_gpus == 0: + sys.stderr.write( + '[condense] ERROR: --condense is ON but COMPRESS_API_KEY is unset ' + 'and no local condenser (EVAL_CONDENSER_GPUS=0).\n' + ' Fix one of:\n' + ' - export COMPRESS_API_KEY=sk-... (use API condenser)\n' + ' - EVAL_CONDENSER_GPUS=2 python ... (use local vLLM condenser)\n' + ' - pass --no-condense (inject raw traces)\n') + sys.exit(1) + + if is_rag: + num_gpus = EMB_GPUS + GEN_GPUS + condenser_gpus + device_groups = [ + DeviceGroup(name='emb_model', ranks=list(range(EMB_GPUS)), + device_type='GPU'), + DeviceGroup(name='sampler', + ranks=list(range(EMB_GPUS, EMB_GPUS + GEN_GPUS)), + device_type='GPU', gpus_per_worker=GEN_GPUS), + ] + if condenser_gpus > 0: + cond_start = EMB_GPUS + GEN_GPUS + device_groups.append( + DeviceGroup(name='condenser', + ranks=list(range(cond_start, cond_start + condenser_gpus)), + device_type='GPU')) + emb_mesh = DeviceMesh.from_sizes(world_size=EMB_GPUS, dp_size=EMB_GPUS) + gen_mesh = DeviceMesh.from_sizes(world_size=GEN_GPUS, tp_size=GEN_GPUS) + twinkle.initialize(mode='ray', nproc_per_node=num_gpus, + groups=device_groups, lazy_collect=False) + else: + device_groups = [ + DeviceGroup(name='sampler', ranks=list(range(GEN_GPUS)), + device_type='GPU', gpus_per_worker=GEN_GPUS), + ] + gen_mesh = DeviceMesh.from_sizes(world_size=GEN_GPUS, tp_size=GEN_GPUS) + twinkle.initialize(mode='ray', nproc_per_node=GEN_GPUS, + groups=device_groups, lazy_collect=False) + + sampler = vLLMSampler( + model_id=GEN_MODEL_ID, + engine_args={ + 'gpu_memory_utilization': GEN_GPU_MEM, + 'max_model_len': GEN_MAX_MODEL_LEN, + }, + device_mesh=gen_mesh, + remote_group='sampler', + ) + sampler.set_template('Qwen3_5Template', model_id=GEN_MODEL_ID, + enable_thinking=True, max_length=GEN_MAX_MODEL_LEN) + sys.stderr.write(f'[aops] vLLM sampler ready (model={GEN_MODEL_ID})\n') + + gen_params = TwinkleSamplingParams( + max_tokens=GEN_MAX_TOKENS, + temperature=GEN_TEMPERATURE, + top_p=GEN_TOP_P, + num_samples=1, + ) + + emb_model = emb_template = tbl = None + if is_rag: + import lancedb + db = lancedb.connect(args.db_path) + if args.table not in db.table_names(): + raise SystemExit(f'Table "{args.table}" not found in {args.db_path}') + tbl = db.open_table(args.table) + sys.stderr.write(f'[aops] LanceDB rows={tbl.count_rows()}\n') + + emb_model = TransformersModel( + model_id=EMBED_MODEL_ID, device_mesh=emb_mesh, + remote_group='emb_model') + emb_model.set_processor(InputProcessor) + emb_model.set_loss(InfonceLoss, temperature=0.03, use_batch=True) + emb_template = Qwen3_5Template( + model_id=EMBED_MODEL_ID, max_length=EMBED_MAX_LENGTH, + truncation_strategy='delete', enable_thinking=False) + sys.stderr.write('[aops] embedding model ready\n') + + # -- Condenser setup (API primary + optional local vLLM) ------------------- + condenser_api_client = None + condenser_sampler_obj = None + condenser_params = None + condenser_special_tokens = None + + if args.condense and is_rag: + condenser_api_client = OpenAIClient( + model=CONDENSE_API_MODEL, api_key=CONDENSE_API_KEY, + base_url=CONDENSE_BASE_URL) + sys.stderr.write(f'[condense] API client ready (model={CONDENSE_API_MODEL})\n') + + if condenser_gpus > 0: + condenser_mesh = DeviceMesh.from_sizes( + world_size=condenser_gpus, dp_size=condenser_gpus) + condenser_sampler_obj = vLLMSampler( + model_id=CONDENSE_MODEL_ID, + engine_args={'gpu_memory_utilization': 0.8, 'max_model_len': 32768}, + device_mesh=condenser_mesh, + remote_group='condenser', + ) + condenser_sampler_obj.set_template( + 'Qwen3_5Template', model_id=CONDENSE_MODEL_ID, + enable_thinking=False, truncation_strategy='delete', + max_length=32768) + condenser_template = Qwen3_5Template( + model_id=CONDENSE_MODEL_ID, max_length=32768, + enable_thinking=False, truncation_strategy='delete') + condenser_special_tokens = set(condenser_template.tokenizer.all_special_tokens) + condenser_params = TwinkleSamplingParams( + max_tokens=CONDENSE_MAX_TOKENS, + temperature=CONDENSE_TEMPERATURE, + top_p=0.5, num_samples=1) + sys.stderr.write(f'[condense] local vLLM ready (model={CONDENSE_MODEL_ID})\n') + + # -- Hint analysis API client (reuses condenser API config) ----------------- + hint_api_client = None + if args.hint and is_rag: + if condenser_api_client is not None: + hint_api_client = condenser_api_client + else: + hint_api_client = OpenAIClient( + model=CONDENSE_API_MODEL, api_key=CONDENSE_API_KEY, + base_url=CONDENSE_BASE_URL) + sys.stderr.write(f'[hint] API hint analysis enabled (model={CONDENSE_API_MODEL})\n') + + # -- LLM decontam API client --------------------------------------------------- + decontam_api_client = None + if args.llm_decontam and is_rag: + if hint_api_client is not None: + decontam_api_client = hint_api_client + elif condenser_api_client is not None: + decontam_api_client = condenser_api_client + else: + decontam_api_client = OpenAIClient( + model=CONDENSE_API_MODEL, api_key=CONDENSE_API_KEY, + base_url=CONDENSE_BASE_URL) + sys.stderr.write(f'[decontam-llm] LLM decontamination enabled (model={CONDENSE_API_MODEL})\n') + + correct_count = 0 + total_count = 0 + skipped_indices: List[int] = [] # problems skipped by RAG (no valid trace) + evaluated_indices: List[int] = [] # problems actually evaluated + debug_records: List[Dict[str, Any]] = [] + + os.makedirs(os.path.dirname(args.output) or '.', exist_ok=True) + out_f = open(args.output, 'w', encoding='utf-8') + + # Open problem-ids files for incremental writing (RAG mode only) + ids_f = None + skip_f = None + if is_rag: + os.makedirs(os.path.dirname(args.problem_ids_file) or '.', exist_ok=True) + ids_f = open(args.problem_ids_file, 'w', encoding='utf-8') + skip_path = args.problem_ids_file.replace('.json', '_skipped.json') + skip_f = open(skip_path, 'w', encoding='utf-8') + + # -- RAG batch preparation (embed + retrieve + decontam + condense + hint) -- + def _prepare_rag_batch(batch_start: int): + """Prepare a RAG batch: returns (prompts, batch, all_examples, + hint_analyses, kept_global_indices, batch_skipped_indices) or None.""" + batch_end = min(batch_start + args.batch_size, len(records)) + batch = records[batch_start:batch_end] + problems = [r['problem'] for r in batch] + + query_vecs = get_embeddings(emb_model, emb_template, problems, EMB_GPUS) + use_raw = not args.use_cot_compressed + all_examples = retrieve_examples(tbl, query_vecs, args.top_k, + use_raw, args.sim_threshold, + problems=problems, + decontam_threshold=args.decontam_threshold) + if args.use_cot_compressed: + for exs in all_examples: + for ex in exs: + ex['thinking'] = _strip_condenser_markers(ex['thinking']) + + if args.llm_decontam and decontam_api_client: + all_examples = _llm_decontaminate( + decontam_api_client, problems, all_examples) + + if args.condense and condenser_api_client: + all_examples = condense_traces( + all_examples, problems, condenser_api_client, + condenser_sampler=condenser_sampler_obj, + compress_params=condenser_params, + special_tokens=condenser_special_tokens, + max_output_len=args.condense_max_len, + dp_size=condenser_gpus) + + hint_analyses = None + if args.hint and hint_api_client: + hint_analyses = _api_hint_analysis_batch( + hint_api_client, problems, all_examples) + + keep_mask = [] + for pi, (r, examples) in enumerate(zip(batch, all_examples)): + if not examples: + keep_mask.append(False) + elif hint_analyses and hint_analyses[pi]: + keep_mask.append(True) + else: + usable = [ex for ex in examples + if len(ex['thinking']) <= args.max_trace_len] + keep_mask.append(bool(usable)) + + batch_skipped = [] + for pi, keep in enumerate(keep_mask): + if not keep: + batch_skipped.append(batch_start + pi) + + kept_batch = [] + kept_examples = [] + kept_hints = [] + kept_global_indices = [] + for pi, keep in enumerate(keep_mask): + if keep: + kept_batch.append(batch[pi]) + kept_examples.append(all_examples[pi]) + kept_hints.append(hint_analyses[pi] if hint_analyses else None) + kept_global_indices.append(batch_start + pi) + + if not kept_batch: + return None, None, None, None, None, batch_skipped + + prompts = [] + for pi, (r, examples) in enumerate(zip(kept_batch, kept_examples)): + if kept_hints[pi]: + prompts.append(build_hint_prompt(r['problem'], kept_hints[pi])) + else: + filtered = [{'query': ex['query'], 'thinking': ex['thinking']} + for ex in examples + if len(ex['thinking']) <= args.max_trace_len] + prompts.append(build_rag_prompt(r['problem'], filtered)) + + return prompts, kept_batch, kept_examples, kept_hints, kept_global_indices, batch_skipped + + target_reached = False + batch_starts = list(range(0, len(records), args.batch_size)) + + if is_rag: + # Pipeline: prefetch next batch while current batch generates + from concurrent.futures import Future + prefetch_pool = ThreadPoolExecutor(max_workers=1) + # Prepare first batch synchronously + cur_result = _prepare_rag_batch(batch_starts[0]) + + for bi, batch_start in enumerate(batch_starts): + if target_reached: + break + prompts, batch, all_examples, hint_analyses, kept_global_indices, batch_skipped = cur_result + skipped_indices.extend(batch_skipped or []) + if skip_f and batch_skipped: + for sid in batch_skipped: + skip_f.write(f'{sid}\n') + skip_f.flush() + + # Submit next batch preparation in background + next_future: Optional[Future] = None + if bi + 1 < len(batch_starts) and not target_reached: + next_future = prefetch_pool.submit(_prepare_rag_batch, batch_starts[bi + 1]) + + if prompts is None: + # Entire batch skipped + cur_result = next_future.result() if next_future else None + continue + + # Generate (runs on gen GPU while next batch prepares on emb GPU + API) + responses = sampler.sample(prompts, gen_params) + + for i, (rec, resp) in enumerate(zip(batch, responses)): + seq = resp.sequences[0] if resp and resp.sequences else None + raw_output = '' + if seq is not None: + raw_output = seq.decoded or '' + raw_output = re.sub(r'<\|[^|]+\|>', '', raw_output).rstrip() + + predicted = extract_boxed(raw_output) + is_correct = answers_match(predicted, rec['reference_answer']) + if is_correct: + correct_count += 1 + total_count += 1 + + global_idx = kept_global_indices[i] + evaluated_indices.append(global_idx) + if ids_f: + ids_f.write(f'{global_idx}\n') + ids_f.flush() + + debug_rec = { + 'idx': global_idx, + 'reference_answer': rec['reference_answer'], + 'predicted': predicted, + 'is_correct': is_correct, + 'problem': rec['problem'], + 'model_output': raw_output, + } + if rec.get('level'): + debug_rec['level'] = rec['level'] + if rec.get('type'): + debug_rec['type'] = rec['type'] + debug_rec['num_traces'] = len(all_examples[i]) + if all_examples[i]: + ex0 = all_examples[i][0] + debug_rec['similarity'] = ex0.get('_sim', 0.0) + debug_rec['retrieved_query'] = ex0.get('query', '') + debug_rec['raw_trace_len'] = ex0.get('_raw_trace_len', 0) + debug_rec['condensed_trace'] = ex0['thinking'] + debug_rec['condensed_trace_len'] = len(ex0['thinking']) + debug_rec['condense_source'] = ex0.get('_condense_source', '') + if hint_analyses and hint_analyses[i]: + debug_rec['hint_analysis'] = hint_analyses[i] + debug_records.append(debug_rec) + out_f.write(json.dumps(debug_rec, ensure_ascii=False) + '\n') + out_f.flush() + + acc = correct_count / total_count if total_count else 0 + sys.stderr.write( + f' [{total_count}/{args.target_eval}] ' + f'acc={acc:.4f} ({correct_count}/{total_count})\n') + + if args.target_eval > 0 and total_count >= args.target_eval: + target_reached = True + + # Collect prefetched result for next iteration (skip if done) + if not target_reached and next_future: + cur_result = next_future.result() + else: + cur_result = None + + prefetch_pool.shutdown(wait=True) + else: + # Direct mode: no pipeline needed, just batch generate + for batch_start in batch_starts: + batch_end = min(batch_start + args.batch_size, len(records)) + batch = records[batch_start:batch_end] + prompts = [build_direct_prompt(r['problem']) for r in batch] + + responses = sampler.sample(prompts, gen_params) + + for i, (rec, resp) in enumerate(zip(batch, responses)): + seq = resp.sequences[0] if resp and resp.sequences else None + raw_output = '' + if seq is not None: + raw_output = seq.decoded or '' + raw_output = re.sub(r'<\|[^|]+\|>', '', raw_output).rstrip() + + predicted = extract_boxed(raw_output) + is_correct = answers_match(predicted, rec['reference_answer']) + if is_correct: + correct_count += 1 + total_count += 1 + + global_idx = original_indices[batch_start + i] + evaluated_indices.append(global_idx) + + debug_rec = { + 'idx': global_idx, + 'reference_answer': rec['reference_answer'], + 'predicted': predicted, + 'is_correct': is_correct, + 'problem': rec['problem'], + 'model_output': raw_output, + } + if rec.get('level'): + debug_rec['level'] = rec['level'] + if rec.get('type'): + debug_rec['type'] = rec['type'] + debug_records.append(debug_rec) + out_f.write(json.dumps(debug_rec, ensure_ascii=False) + '\n') + out_f.flush() + + acc = correct_count / total_count if total_count else 0 + sys.stderr.write( + f' [{total_count}/{len(records)}] ' + f'acc={acc:.4f} ({correct_count}/{total_count})\n') + + overall_acc = correct_count / total_count if total_count else 0 + print(f'\n{"=" * 60}') + print(f'{args.dataset.upper()} — mode={args.mode}, model={GEN_MODEL_ID}') + print(f' n={total_count}, seed={args.seed}') + if is_rag: + print(f' evaluated={len(evaluated_indices)}, skipped={len(skipped_indices)}') + print(f'{"=" * 60}') + print(f'Overall accuracy: {overall_acc:.4f} ({correct_count}/{total_count})') + + # Per-level breakdown (MATH: the difficulty-vs-gain curve we care about). + if any(r.get('level') for r in debug_records): + from collections import defaultdict + per = defaultdict(lambda: [0, 0]) # level -> [correct, total] + for r in debug_records: + lv = r.get('level', 'Unknown') + per[lv][1] += 1 + if r['is_correct']: + per[lv][0] += 1 + print(f'\nPer-level accuracy:') + for lv in sorted(per.keys()): + c, t = per[lv] + print(f' {lv:>10}: {c/t:.4f} ({c}/{t})') + + out_f.close() + print(f'\n[output] {len(debug_records)} records saved to {args.output}') + + if ids_f: + ids_f.close() + print(f'[output] problem IDs ({len(evaluated_indices)}) saved to {args.problem_ids_file}') + if skip_f: + skip_f.close() + if skipped_indices: + print(f'[output] skipped IDs ({len(skipped_indices)}) saved to ' + f'{args.problem_ids_file.replace(".json", "_skipped.json")}') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/legacy/eval_math_by_level.sh b/cookbook/exp/legacy/eval_math_by_level.sh new file mode 100755 index 000000000..d4f8bbdb3 --- /dev/null +++ b/cookbook/exp/legacy/eval_math_by_level.sh @@ -0,0 +1,59 @@ +#!/bin/bash +# MATH (Hendrycks) difficulty-stratified evaluation. +# +# Goal: measure how the (raw) RAG gain over direct varies with problem +# difficulty (Level 1-5). Runs raw RAG first (retrieve -> qwen3.7-max condense +# -> inject, no hint filtering; it writes the problem-id file), then direct on +# the *same* problems for a paired comparison. +# +# Usage: +# COMPRESS_API_KEY=sk-xxx bash cookbook/exp/embedding/eval_math_by_level.sh +# +# Env knobs: +# PER_LEVEL problems per difficulty level (default 100 -> 500 total) +# SEED stratified-sampling seed (default 100; must match across runs) +# DB_PATH LanceDB retrieval index +# SIM / TOPK retrieval threshold / top-k + +set -euo pipefail + +export COMPRESS_API_KEY="${COMPRESS_API_KEY:?Set COMPRESS_API_KEY}" + +SCRIPT="cookbook/exp/embedding/eval_gpqa_rag.py" +PER_LEVEL="${PER_LEVEL:-100}" +SEED="${SEED:-100}" +SIM="${SIM:-0.75}" +TOPK="${TOPK:-1}" +OUTDIR="./output/thinking_rag" +DB_PATH="${DB_PATH:-./output.oldemb/thinking_rag/lance.db}" + +mkdir -p "$OUTDIR" + +echo "============================================================" +echo " MATH by level: raw RAG (qwen3.7-max condenser, no hint)" +echo " per_level=$PER_LEVEL seed=$SEED" +echo "============================================================" +python "$SCRIPT" \ + --dataset math --math-split test \ + --mode rag \ + --per-level "$PER_LEVEL" --seed "$SEED" \ + --db-path "$DB_PATH" \ + --sim-threshold "$SIM" --top-k "$TOPK" \ + --condense \ + --output "$OUTDIR/math_rag_results.jsonl" + +echo "" +echo "============================================================" +echo " MATH by level: Direct (same problems as raw RAG)" +echo "============================================================" +# Direct reads math_rag_problem_ids.json (written above) to match the subset. +python "$SCRIPT" \ + --dataset math --math-split test \ + --mode direct \ + --per-level "$PER_LEVEL" --seed "$SEED" \ + --output "$OUTDIR/math_direct_results.jsonl" + +echo "" +echo "============================================================" +echo " Done. Compare with: python cookbook/exp/embedding/compare_math_levels.py" +echo "============================================================" diff --git a/cookbook/exp/legacy/eval_rag_recall.py b/cookbook/exp/legacy/eval_rag_recall.py new file mode 100644 index 000000000..19bc5ad6d --- /dev/null +++ b/cookbook/exp/legacy/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/legacy/eval_reflexion_skill.py b/cookbook/exp/legacy/eval_reflexion_skill.py new file mode 100644 index 000000000..3f7ddcc0d --- /dev/null +++ b/cookbook/exp/legacy/eval_reflexion_skill.py @@ -0,0 +1,762 @@ +"""Phase-0 measurement for the reflexion self-skill scheme (see reflexion.md). + +Question this script answers: **on problems the base model first gets wrong, does +letting the SAME base model reflect on its failed attempt, distill a general +"skill", and re-solve WITH that skill in the system prompt, actually raise its +pass@k?** No LoRA is trained here — this is the upper-bound / go-no-go gate before +investing in a Skill-LoRA. If the base model's own skills don't help, training a +LoRA to produce them is pointless. + +It deliberately reuses ``eval_gpqa_rag`` verbatim (dataset, grader, prompts, sampling +config) so numbers are comparable with the other AoPS lines. Only the base model + +one vLLM sampler are used; the dataset is AoPS; validation is on the SAME problem +(no similar-problem retrieval). + +Per chunk of problems (all sampler calls are BATCHED across the whole chunk — never +one problem at a time): + 1. Initial solve — 1 rollout each; keep only problems the model got wrong. + 2. Skill generation — for each failed problem, the base model reads its own failed + attempt and produces N candidate skills (general reminders, no answer/solution). + 3. Leak filter — drop skills that leak the gold answer or a full solution. + 4. Baseline pass@k — K rollouts of the plain problem (the "no-skill" control). + 5. With-skill pass@k — K rollouts of the problem with each surviving skill in the + system prompt. + 6. Score — marginal = with-skill pass@k − baseline pass@k; keep the best skill. +A "pass" = answer correct AND generation terminated (no length cutoff). + +Everything useful (failed attempt, every candidate skill + leak flag, baseline and +per-skill rollout stats, marginals, best skill) is written to a JSONL **incrementally +after each chunk**, so partial runs are fully analysable. + +Launch (8 GPUs, tp=1 dp=8 by default): + python cookbook/exp/embedding/eval_reflexion_skill.py --n 64 --chunk-size 16 +""" +import argparse +import copy +import json +import os +import re +import sys +import threading +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Any, Dict, List, Optional, Tuple + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import SamplingParams as TwinkleSamplingParams +from twinkle.sampler import vLLMSampler +from twinkle_agentic.protocol.openai import OpenAI as OpenAIClient + +# Reuse the reference eval's dataset + grading + prompts + sampling config so this +# line is directly comparable with eval_gpqa_rag / eval_dualline_math. +from eval_gpqa_rag import (DIRECT_SYSTEM, GEN_GPU_MEM, GEN_GPUS, GEN_MODEL_ID, + GEN_TEMPERATURE, GEN_TOP_P, MCQ_INSTRUCTION, answers_match, + build_direct_prompt, extract_boxed, load_aops) + +logger = get_logger() + +# vLLM parallel: tp=1, dp=GEN_GPUS by default (override GEN_TP; keep GEN_GPUS=8). +GEN_TP = int(os.environ.get('GEN_TP', 1)) + +# Leak-detector API (reuses eval_gpqa_rag's env names). A strong external model +# judges whether a candidate skill leaks THIS problem's answer/solution — catching +# what the string filter cannot (multiple-choice letters, derived-result leakage). +LEAK_API_KEY = os.environ.get('COMPRESS_API_KEY', '') +LEAK_BASE_URL = os.environ.get('COMPRESS_BASE_URL', + 'https://dashscope.aliyuncs.com/compatible-mode/v1') +LEAK_API_MODEL = os.environ.get('COMPRESS_MODEL', 'qwen3.7-max') + +# Global call spacer so bursts of leak-judge calls stay under the QPS limit. +_api_lock = threading.Lock() +_api_next = [0.0] + + +def _api_throttle(min_interval: float) -> None: + with _api_lock: + now = time.monotonic() + wait = max(0.0, _api_next[0] - now) + _api_next[0] = max(now, _api_next[0]) + min_interval + if wait > 0: + time.sleep(wait) + + +# --------------------------------------------------------------------------- +# Prompts (self-reflection skill generation + skill-conditioned solving) +# --------------------------------------------------------------------------- +SKILL_GEN_SYSTEM = ( + "You are a meticulous mathematics coach. You are shown a competition problem and a " + "student's FAILED attempt. Produce a SHORT list of general, reusable skills that " + 'would prevent this class of mistake on SIMILAR problems.\n\n' + 'OUTPUT FORMAT (strict):\n' + '- You may reason briefly first, but the final answer MUST be a markdown bullet ' + 'list of 3-5 items WRAPPED IN and tags. Output nothing after ' + '.\n' + '- Each item is ONE short imperative sentence (a rule, check, or habit).\n' + '- Inside the tags: no diagnosis narration, no "The student...", no headings, no ' + 'restating the problem or the examples.\n\n' + 'CONTENT RULES (strict):\n' + '- Do NOT reveal the final answer or the multiple-choice option.\n' + '- Do NOT state the specific numbers, values, or key intermediate results of THIS ' + 'problem.\n' + '- Do NOT give a step-by-step solution to THIS problem. Every item must be GENERAL ' + 'and transferable to other problems of the same type.\n\n' + 'Follow the example below for the exact tags, style, and level of generality.' +) + +SKILL_GEN_USER = ( + 'Problem:\n{problem}\n\n' + "The student's failed attempt (it may be long or may fail to terminate):\n" + '{attempt}\n\n' + 'Now output the skills bullet list.' +) + +# One-shot demonstration of the required format and generality (answer-free). +_EX_PROBLEM = 'Simplify $\\sqrt{72} + \\sqrt{18}$ and give the result.' +_EX_ATTEMPT = ( + 'The student added the radicands directly to get $\\sqrt{90}$ and concluded it ' + 'could not be simplified, never factoring out the perfect squares first.') +_EX_SKILLS = ( + '\n' + '- Before adding square roots, factor each radicand into a perfect square times a ' + 'remainder and move the perfect-square root outside.\n' + '- Never add radicands directly: $\\sqrt{a}+\\sqrt{b}\\ne\\sqrt{a+b}$.\n' + '- Only combine radical terms after reducing them to the same simplest radical ' + 'form.\n' + '- Sanity-check the simplified result by estimating each root numerically.\n' + '') + + +def build_skillgen_prompt(problem: str, attempt: str) -> Dict[str, Any]: + return {'messages': [ + {'role': 'system', 'content': SKILL_GEN_SYSTEM}, + {'role': 'user', + 'content': SKILL_GEN_USER.format(problem=_EX_PROBLEM, attempt=_EX_ATTEMPT)}, + {'role': 'assistant', 'content': _EX_SKILLS}, + {'role': 'user', + 'content': SKILL_GEN_USER.format(problem=problem, attempt=attempt)}, + ]} + + +# The skill is injected into the SYSTEM prompt (per reflexion.md), on top of the +# exact DIRECT_SYSTEM used by the baseline so the only difference is the reminders. +# Built by concatenation (NOT str.format): DIRECT_SYSTEM and the skill may contain +# literal braces (e.g. ``\boxed{}``, LaTeX), which would break ``.format``. +_SKILL_SOLVE_PREFIX = ( + DIRECT_SYSTEM + '\n\n' + 'Before you start, keep these reminders in mind to avoid common mistakes on this ' + 'type of problem:\n') +_SKILL_SOLVE_SUFFIX = ( + '\nApply them where relevant, but rely on your own reasoning to reach the answer.') + + +def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: + return {'messages': [ + {'role': 'system', 'content': _SKILL_SOLVE_PREFIX + skill + _SKILL_SOLVE_SUFFIX}, + {'role': 'user', 'content': problem + MCQ_INSTRUCTION}, + ]} + + +# --------------------------------------------------------------------------- +# Parsing / grading / leak filtering +# --------------------------------------------------------------------------- +_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') + + +def _clean_text(decoded: Optional[str]) -> str: + return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() + + +_BULLET_RE = re.compile(r'^\s*(?:[-*]|\d+[.)])\s') + + +def _extract_skill_list(text: str) -> str: + """Pull just the clean skill list out of a (possibly thinking-laden) output. + + The model is instructed to wrap the final list in ``...``, so + prefer that (robust to any preceding reasoning, closed or unterminated). Fall + back to dropping a ```` block and keeping from the first bullet onward. + """ + low = text.lower() + if '' in low: + start = low.index('') + len('') + end = low.index('') if '' in low else len(text) + return text[start:end].strip() + if '' in text: + text = text.rsplit('', 1)[-1] + text = re.sub(r'.*?', '', text, flags=re.DOTALL).strip() + lines = text.splitlines() + for i, line in enumerate(lines): + if _BULLET_RE.match(line): + return '\n'.join(lines[i:]).strip() + return text.strip() + + +def _bound_attempt(text: str, gen_tokens: int, budget_tokens: int) -> str: + """Keep a failed attempt within the skill-gen context budget. + + Round-2 (skill-gen) input contains the FULL round-1 attempt, and failed + attempts are often the ones that ran to the token cap (repetition loops), so + feeding them verbatim overflows max_model_len. Keep the head (real reasoning + + where it went wrong) plus a short tail (the final wrong answer); drop the + redundant middle. Token->char conversion uses THIS attempt's observed + chars-per-token so the cut fits precisely. + """ + if not text or gen_tokens <= budget_tokens: + return text + cpt = len(text) / max(1, gen_tokens) + head_tok = int(budget_tokens * 0.7) + tail_tok = budget_tokens - head_tok + head = text[:int(head_tok * cpt)] + tail = text[-int(tail_tok * cpt):] if tail_tok > 0 else '' + return f'{head}\n\n[... attempt truncated for length ...]\n\n{tail}' + + +def _parse_seq(seq, gold: str) -> Dict[str, Any]: + """Turn one sampled sequence into a graded rollout record. + + ``pass`` requires BOTH a correct boxed answer AND clean termination (a length + cutoff means the model never actually committed to the answer). + """ + text = _clean_text(getattr(seq, 'decoded', '') or '') + pred = extract_boxed(text) + correct = bool(pred) and answers_match(pred, gold) + terminated = getattr(seq, 'stop_reason', None) != 'length' + return { + 'pred': pred, + 'correct': correct, + 'terminated': terminated, + 'passed': bool(correct and terminated), + 'stop_reason': getattr(seq, 'stop_reason', None), + 'gen_tokens': len(getattr(seq, 'tokens', None) or []), + 'text': text, + } + + +def _pass_rate(rolls: List[Dict[str, Any]]) -> float: + return sum(1 for r in rolls if r['passed']) / len(rolls) if rolls else 0.0 + + +def _skill_leaks(skill: str, gold: str) -> Tuple[bool, str]: + """Reject a skill that leaks the answer, or is a degenerate / non-list output.""" + if not skill.strip(): + return True, 'empty' + bullets = [ln for ln in skill.splitlines() if _BULLET_RE.match(ln)] + if len(bullets) < 2: + return True, 'too_short' + low = skill.lower() + if 'item 1' in low and 'item 2' in low: # model echoed the format placeholder + return True, 'placeholder' + if '\\boxed' in skill: + return True, 'contains_boxed' + g = (gold or '').strip() + # Raw substring match is only trustworthy when the answer is specific enough that + # an incidental hit is unlikely. Short answers ('D', 'E', '1') would match almost + # any text, so leave those to the API judge instead of false-flagging every skill. + if len(g) >= 4 and g.lower() in skill.lower(): + return True, 'contains_gold_answer' + # Standalone multi-digit numbers from the gold answer leaking into the skill. + for num in re.findall(r'-?\d{2,}', g): + if re.search(r'(? Optional[bool]: + """Return True (leak) / False (clean) / None (unparseable or API error after retries). + + Only transient API errors are retried (with exponential backoff); an unparseable + verdict is deterministic at temperature 0, so retrying it is pointless. + """ + msgs = [ + {'role': 'system', 'content': _LEAK_JUDGE_SYSTEM}, + {'role': 'user', 'content': _LEAK_JUDGE_USER.format( + problem=problem[:4000], gold=gold, skill=skill[:4000])}, + ] + for attempt in range(retries + 1): + _api_throttle(min_interval) + try: + reply = api({'messages': msgs}, + TwinkleSamplingParams(temperature=0.0, max_tokens=16), + extra_body={'enable_thinking': False}) + except Exception as exc: # noqa: BLE001 — broad catch is intentional + logger.warning(f'[leak-judge] error (attempt {attempt + 1}/{retries + 1}): {exc}') + if attempt < retries: + time.sleep(min(4.0, 0.5 * 2 ** attempt)) # exponential backoff + continue + return None + verdict = (reply.get('content') or '').strip().upper() + if 'CLEAN' in verdict: + return False + if 'LEAK' in verdict: + return True + return None # unparseable — deterministic at temp 0, no point retrying + + +def _api_leak_batch(api: OpenAIClient, items: List[Tuple[int, str, str, str]], + concurrency: int, min_interval: float, + retries: int) -> Dict[int, Optional[bool]]: + """Judge many (key, problem, gold, skill) tuples in parallel; key -> verdict.""" + verdicts: Dict[int, Optional[bool]] = {} + if not items: + return verdicts + with ThreadPoolExecutor(max_workers=min(len(items), concurrency)) as pool: + futs = {pool.submit(_api_leak_judge_one, api, p, g, s, min_interval, retries): k + for (k, p, g, s) in items} + for fut in as_completed(futs): + verdicts[futs[fut]] = fut.result() + return verdicts + + +# --------------------------------------------------------------------------- +# Batched sampling (one shared sampler.sample per phase — never per problem) +# --------------------------------------------------------------------------- +def _pad_for_dp(prompts: List[Any], gen_dp: int) -> List[Any]: + """vLLM dp needs batch len >= dp; pad tail rounds and let the caller slice back.""" + if gen_dp <= 1 or not prompts or len(prompts) >= gen_dp: + return prompts + pad = [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] + return prompts + pad + + +def _run_samples(sampler, prompts: List[Any], num_samples: int, max_tokens: int, + gen_dp: int, temperature: Optional[float] = None, + top_p: Optional[float] = None, + top_k: Optional[int] = None) -> List[List[Any]]: + """One batched sampler call; return per-prompt list of raw sampled sequences. + + ``temperature``/``top_p``/``top_k`` default to the module's sampling config; pass + ``temperature=0.0`` for deterministic greedy decoding (SEAM-style executor scoring), + or a high ``temperature`` with ``top_k=-1`` for diverse multi-candidate sampling.""" + if not prompts: + return [] + params = TwinkleSamplingParams( + max_tokens=max_tokens, + temperature=GEN_TEMPERATURE if temperature is None else temperature, + top_p=GEN_TOP_P if top_p is None else top_p, + num_samples=num_samples, + **({} if top_k is None else {'top_k': top_k})) + padded = _pad_for_dp(prompts, gen_dp) + responses = sampler.sample(padded, params)[:len(prompts)] + return [list(r.sequences) if (r and r.sequences) else [] for r in responses] + + +def _set_thinking(sampler, args: argparse.Namespace, enabled: bool) -> None: + """Toggle the remote template's thinking mode. + + Skill generation wants thinking OFF so the model emits the short ```` + list directly (with thinking ON it burns the token budget reasoning and often + never reaches the list); solving wants it ON. ``set_template`` is a + remote_function, so this propagates to every sampler worker. + """ + sampler.set_template('Qwen3_5Template', model_id=GEN_MODEL_ID, + enable_thinking=enabled, max_length=args.max_model_len) + + +def _bounded_attempt_for(r: Dict[str, Any], args: argparse.Namespace) -> str: + """Per-problem bound so problem + attempt + skill output fits the context window.""" + prob_est = len(r['problem']) // 2 # conservative problem token estimate + budget = max(1024, args.max_model_len - args.skill_max_tokens + - args.attempt_reserve_tokens - prob_est) + return _bound_attempt(r['_init'][0]['text'], r['_init'][0]['gen_tokens'], budget) + + +def _filter_candidates(api: Optional[OpenAIClient], + cands: List[Tuple[Dict[str, Any], str]], + args: argparse.Namespace) -> None: + """Apply string + API leak filters to (problem, skill) candidates; append results. + + The cheap string filter runs first; the API judge only sees skills that pass it, + which is what catches MCQ-letter / derived-result / full-solution leakage. + """ + prepared = [] # [r, text, leaked(bool|None), reason] + for r, text in cands: + leaked, reason = _skill_leaks(text, r['reference_answer']) + prepared.append([r, text, True if leaked else None, reason]) + if api is not None: + items = [(i, prepared[i][0]['problem'], prepared[i][0]['reference_answer'], + prepared[i][1]) for i in range(len(prepared)) if prepared[i][2] is None] + verdicts = _api_leak_batch(api, items, args.api_concurrency, args.api_min_interval, + args.api_retries) + for key, _p, _g, _s in items: + v = verdicts.get(key) + if v is True: + prepared[key][2], prepared[key][3] = True, 'api_leak' + elif v is False: + prepared[key][2], prepared[key][3] = False, '' + else: + prepared[key][2], prepared[key][3] = False, 'api_uncertain' + for r, text, leaked, reason in prepared: + r['_skills'].append({'skill': text, 'leaked': bool(leaked), 'leak_reason': reason}) + + +def _build_skills(sampler, api: Optional[OpenAIClient], failed: List[Dict[str, Any]], + gen_dp: int, args: argparse.Namespace) -> None: + """Generate + extract + leak-filter skills, re-rolling problems short on clean ones. + + Clean skills accumulate across rounds; only problems still below ``min_survivors`` + clean skills are re-rolled, up to ``skill_retries`` extra rounds. + """ + for r in failed: + r['_skills'] = [] + todo = list(failed) + _set_thinking(sampler, args, False) # skill-gen: emit the list directly, no CoT + try: + for _ in range(args.skill_retries + 1): + if not todo: + break + sg_out = _run_samples( + sampler, + [build_skillgen_prompt(r['problem'], _bounded_attempt_for(r, args)) for r in todo], + args.n_skills, args.skill_max_tokens, gen_dp) + cands = [(r, _extract_skill_list(_clean_text(getattr(s, 'decoded', '') or ''))) + for r, seqs in zip(todo, sg_out) for s in seqs] + _filter_candidates(api, cands, args) + todo = [r for r in failed + if sum(1 for sk in r['_skills'] if not sk['leaked']) < args.min_survivors] + finally: + _set_thinking(sampler, args, True) # restore for solving phases + tot = sum(len(r['_skills']) for r in failed) + leaked = sum(1 for r in failed for sk in r['_skills'] if sk['leaked']) + sys.stderr.write(f' phase2: skills={tot} leaked={leaked} ' + f'({leaked / max(1, tot):.0%}); {len(todo)} still short of ' + f'{args.min_survivors} clean\n') + + +# --------------------------------------------------------------------------- +# Per-chunk pipeline +# --------------------------------------------------------------------------- +def process_chunk(sampler, api: Optional[OpenAIClient], chunk: List[Dict[str, Any]], + gen_dp: int, args: argparse.Namespace) -> List[Dict[str, Any]]: + """Run all 6 phases for one chunk (batched) and return per-problem records.""" + # --- Phase 1: initial solve, keep only the ones the model got wrong. --- + init_out = _run_samples( + sampler, [build_direct_prompt(r['problem']) for r in chunk], + args.init_samples, args.max_tokens, gen_dp) + for r, seqs in zip(chunk, init_out): + r['_init'] = [_parse_seq(s, r['reference_answer']) for s in seqs] + r['_init_pass'] = _pass_rate(r['_init']) + r['_failed'] = r['_init_pass'] == 0.0 + failed = [r for r in chunk if r['_failed']] + sys.stderr.write(f' phase1: {len(chunk)-len(failed)}/{len(chunk)} solved on ' + f'first try, {len(failed)} failed -> reflect\n') + + if failed: + # --- Baseline pass@k FIRST: defines which failures are genuinely hard. --- + # (A single initial rollout is noisy; an easy problem can fail phase 1 yet + # have a high pass@k, so measure the marginal only on truly hard problems.) + base_out = _run_samples( + sampler, [build_direct_prompt(r['problem']) for r in failed], + args.pass_k, args.max_tokens, gen_dp) + for r, seqs in zip(failed, base_out): + r['_baseline'] = [_parse_seq(s, r['reference_answer']) for s in seqs] + r['_baseline_pass'] = _pass_rate(r['_baseline']) + r['_hard'] = r['_baseline_pass'] <= args.hard_baseline_max + r['_skills'] = [] + r['_best'] = None + hard = [r for r in failed if r['_hard']] + sys.stderr.write(f' baseline: {len(hard)}/{len(failed)} failures are hard ' + f'(pass@{args.pass_k} <= {args.hard_baseline_max})\n') + + if hard: + # --- Skills (generate + leak filter + re-rollout) for HARD problems only. --- + _build_skills(sampler, api, hard, gen_dp, args) + + # --- With-skill pass@k (flatten hard-problem x surviving skill). --- + flat: List[Tuple[int, int]] = [] + ws_prompts: List[Any] = [] + for ri, r in enumerate(hard): + for si, sk in enumerate(r['_skills']): + if sk['leaked'] or not sk['skill'].strip(): + continue + flat.append((ri, si)) + ws_prompts.append(build_skill_solve_prompt(r['problem'], sk['skill'])) + ws_out = _run_samples(sampler, ws_prompts, args.pass_k, args.max_tokens, gen_dp) + for (ri, si), seqs in zip(flat, ws_out): + r = hard[ri] + sk = r['_skills'][si] + sk['rolls'] = [_parse_seq(s, r['reference_answer']) for s in seqs] + sk['with_pass'] = _pass_rate(sk['rolls']) + sk['marginal'] = sk['with_pass'] - r['_baseline_pass'] + + # --- Pick the best (highest marginal) surviving skill per hard problem. --- + for r in hard: + scored = [sk for sk in r['_skills'] if 'marginal' in sk] + r['_best'] = max(scored, key=lambda s: s['marginal']) if scored else None + + return [_make_record(r, args) for r in chunk] + + +def _roll_summary(roll: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: + out = {k: roll[k] for k in ('pred', 'correct', 'terminated', 'passed', + 'stop_reason', 'gen_tokens')} + if args.store_rollout_text: + out['text'] = roll['text'][:args.store_rollout_chars] + return out + + +def _make_record(r: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: + """Assemble the incremental JSONL record for one problem (solved or failed).""" + rec: Dict[str, Any] = { + 'problem': r['problem'], + 'reference_answer': r['reference_answer'], + 'tags': r.get('tags', []), + 'failed_first_try': r['_failed'], + 'init_pass_rate': r['_init_pass'], + 'init_attempt': { + 'text': r['_init'][0]['text'][:args.store_init_chars], + 'pred': r['_init'][0]['pred'], + 'stop_reason': r['_init'][0]['stop_reason'], + 'gen_tokens': r['_init'][0]['gen_tokens'], + }, + } + if not r['_failed']: + return rec + + best = r.get('_best') + rec['baseline_pass'] = r['_baseline_pass'] + # Genuinely hard = low baseline pass@k; only these count in the marginal stats. + rec['is_hard'] = bool(r.get('_hard')) + rec['baseline_rolls'] = [_roll_summary(x, args) for x in r['_baseline']] + rec['skills'] = [{ + 'skill': sk['skill'], + 'leaked': sk['leaked'], + 'leak_reason': sk['leak_reason'], + 'with_pass': sk.get('with_pass'), + 'marginal': sk.get('marginal'), + 'rolls': [_roll_summary(x, args) for x in sk.get('rolls', [])], + } for sk in r.get('_skills', [])] + rec['best_skill'] = best['skill'] if best else None + rec['best_marginal'] = best['marginal'] if best else None + rec['best_with_pass'] = best['with_pass'] if best else None + # "rescued" = a leak-free skill turned a fully-failing problem into some passes. + rec['rescued'] = bool(best and r['_baseline_pass'] == 0.0 and best['with_pass'] > 0.0) + rec['helped'] = bool(best and best['marginal'] > 0.0) + return rec + + +# --------------------------------------------------------------------------- +# Running summary +# --------------------------------------------------------------------------- +def _update_summary(summ: Dict[str, Any], recs: List[Dict[str, Any]]) -> None: + for rec in recs: + summ['n_total'] += 1 + if not rec['failed_first_try']: + summ['n_solved_first'] += 1 + continue + summ['n_failed'] += 1 + if not rec.get('is_hard'): + summ['n_failed_easy'] += 1 # failed phase 1 but easy on pass@k — excluded + continue + summ['n_hard'] += 1 + base = rec.get('baseline_pass', 0.0) + summ['sum_baseline_pass'] += base + if rec.get('best_marginal') is not None: + summ['n_with_skill'] += 1 + summ['sum_best_with_pass'] += rec.get('best_with_pass', 0.0) + summ['sum_best_marginal'] += rec.get('best_marginal', 0.0) + else: + # No clean skill produced for this hard problem -> skill adds no gain + # (count it honestly as marginal 0 rather than dropping it from the average). + summ['sum_best_with_pass'] += base + summ['n_helped'] += int(rec.get('helped', False)) + summ['n_rescued'] += int(rec.get('rescued', False)) + + +def _summary_report(summ: Dict[str, Any]) -> Dict[str, Any]: + nh = max(1, summ['n_hard']) + return { + 'record_type': 'summary', + 'n_total': summ['n_total'], + 'n_solved_first_try': summ['n_solved_first'], + 'n_failed_first_try': summ['n_failed'], + 'n_failed_but_easy': summ['n_failed_easy'], + 'n_hard': summ['n_hard'], + 'n_hard_with_skill': summ['n_with_skill'], + # Averages are over ALL hard problems; a hard problem with no clean skill + # counts as zero gain (with_pass == baseline), so with - base == marginal. + 'avg_baseline_pass_on_hard': summ['sum_baseline_pass'] / nh, + 'avg_best_with_skill_pass_on_hard': summ['sum_best_with_pass'] / nh, + 'avg_best_marginal_on_hard': summ['sum_best_marginal'] / nh, + 'n_helped_by_skill': summ['n_helped'], + 'n_rescued_from_zero': summ['n_rescued'], + 'frac_hard_helped': summ['n_helped'] / nh, + } + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- +def main() -> None: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--n', type=int, default=64, help='AoPS problems to sample.') + p.add_argument('--seed', type=int, default=42) + p.add_argument('--chunk-size', type=int, default=16, + help='Problems per chunk. All sampler calls within a chunk are ' + 'batched; results are flushed to disk after each chunk.') + p.add_argument('--init-samples', type=int, default=1, + help='Rollouts for the initial solve. A problem is "failed" (and ' + 'sent to reflection) only if all initial rollouts are wrong.') + p.add_argument('--n-skills', type=int, default=8, + help='Candidate skills generated per failed problem.') + p.add_argument('--pass-k', type=int, default=8, + help='Rollouts per (baseline / with-skill) pass@k estimate.') + p.add_argument('--hard-baseline-max', type=float, default=0.25, + help='A failed problem counts as "hard" (included in the marginal ' + 'stats) only if its baseline pass@k <= this. Filters out easy ' + 'problems that merely failed the single initial rollout.') + p.add_argument('--max-model-len', type=int, default=30000, + help='Context window (engine + template). MUST exceed --max-tokens: ' + 'the round-2 skill-gen input holds the full round-1 attempt ' + 'plus the problem.') + p.add_argument('--max-tokens', type=int, default=20000, + help='Max generated tokens for solving rollouts (round-1 output cap).') + p.add_argument('--skill-max-tokens', type=int, default=2048, + help='Max tokens for skill generation. Enough to finish any thinking ' + 'and emit the short bullet list (which is then extracted).') + p.add_argument('--attempt-reserve-tokens', type=int, default=2048, + help='Tokens reserved for system prompt + wrappers when bounding the ' + 'failed attempt fed into skill generation (the problem length ' + 'is accounted for separately, per-problem).') + p.add_argument('--min-survivors', type=int, default=2, + help='Re-roll a problem\'s skills if fewer than this many survive the ' + 'leak filters.') + p.add_argument('--skill-retries', type=int, default=1, + help='Max extra skill-generation rounds for problems short on clean ' + 'skills (0 = no retry).') + p.add_argument('--api-concurrency', type=int, default=32, + help='Parallel workers for the API leak judge (max 32 recommended).') + p.add_argument('--api-min-interval', type=float, default=0.1, + help='Minimum seconds between API leak-judge calls (QPS guard).') + p.add_argument('--api-retries', type=int, default=3, + help='Retries on transient API errors per leak-judge call (exponential ' + 'backoff); only after these are exhausted is a skill kept as ' + 'api_uncertain.') + p.add_argument('--disable-api-leak', action='store_true', + help='Skip the API leak judge even if COMPRESS_API_KEY is set ' + '(string filter only).') + p.add_argument('--output', default='./output/reflexion_phase0/aops_results.jsonl') + p.add_argument('--store-init-chars', type=int, default=8000, + help='Truncate the stored failed-attempt text to this many chars.') + p.add_argument('--store-rollout-text', action='store_true', + help='Also store (truncated) text of every rollout, not just stats.') + p.add_argument('--store-rollout-chars', type=int, default=2000) + args = p.parse_args() + + records = load_aops(n=args.n, seed=args.seed) + sys.stderr.write(f'[reflexion] {len(records)} AoPS problems, chunk={args.chunk_size}, ' + f'init_samples={args.init_samples}, n_skills={args.n_skills}, ' + f'pass_k={args.pass_k}, max_tokens={args.max_tokens}\n') + + # --- 8-GPU vLLM sampler (tp=GEN_TP, dp=GEN_GPUS/GEN_TP). --- + if GEN_GPUS % GEN_TP != 0: + raise ValueError(f'GEN_GPUS ({GEN_GPUS}) must be divisible by GEN_TP ({GEN_TP})') + gen_dp = GEN_GPUS // GEN_TP + gen_mesh = DeviceMesh.from_sizes(world_size=GEN_GPUS, dp_size=gen_dp, tp_size=GEN_TP) + twinkle.initialize( + mode='ray', nproc_per_node=GEN_GPUS, + groups=[DeviceGroup(name='sampler', ranks=list(range(GEN_GPUS)), + device_type='GPU', gpus_per_worker=GEN_TP)], + lazy_collect=False) + sampler = vLLMSampler( + model_id=GEN_MODEL_ID, + engine_args={'gpu_memory_utilization': GEN_GPU_MEM, + 'max_model_len': args.max_model_len, + 'tensor_parallel_size': GEN_TP}, + device_mesh=gen_mesh, remote_group='sampler') + sampler.set_template('Qwen3_5Template', model_id=GEN_MODEL_ID, + enable_thinking=True, max_length=args.max_model_len) + sys.stderr.write(f'[reflexion] sampler ready (model={GEN_MODEL_ID}, tp={GEN_TP}, ' + f'dp={gen_dp})\n') + + # --- API leak judge (optional): reuses eval_gpqa_rag's COMPRESS_* env. --- + api: Optional[OpenAIClient] = None + if LEAK_API_KEY and not args.disable_api_leak: + api = OpenAIClient(model=LEAK_API_MODEL, api_key=LEAK_API_KEY, + base_url=LEAK_BASE_URL) + sys.stderr.write(f'[reflexion] leak judge ON via API model={LEAK_API_MODEL} ' + f'(concurrency={args.api_concurrency})\n') + else: + sys.stderr.write('[reflexion] leak judge OFF (string filter only) — set ' + 'COMPRESS_API_KEY to enable the API judge\n') + + os.makedirs(os.path.dirname(args.output) or '.', exist_ok=True) + summ = {k: 0 for k in ('n_total', 'n_solved_first', 'n_failed', 'n_failed_easy', + 'n_hard', 'n_with_skill', 'n_helped', 'n_rescued')} + summ.update({'sum_baseline_pass': 0.0, 'sum_best_with_pass': 0.0, + 'sum_best_marginal': 0.0}) + + with open(args.output, 'w', encoding='utf-8') as out_f: + # Line 1: run config, for reproducibility / later analysis. + out_f.write(json.dumps({ + 'record_type': 'config', 'model': GEN_MODEL_ID, 'dataset': 'aops', + 'n': len(records), 'seed': args.seed, 'init_samples': args.init_samples, + 'n_skills': args.n_skills, 'pass_k': args.pass_k, + 'hard_baseline_max': args.hard_baseline_max, + 'max_model_len': args.max_model_len, 'max_tokens': args.max_tokens, + 'skill_max_tokens': args.skill_max_tokens, + 'api_leak_judge': api is not None, + 'api_leak_model': LEAK_API_MODEL if api is not None else None, + 'min_survivors': args.min_survivors, 'skill_retries': args.skill_retries, + 'gpus': GEN_GPUS, 'tp': GEN_TP, 'started': int(time.time()), + }, ensure_ascii=False) + '\n') + out_f.flush() + + n_chunks = (len(records) + args.chunk_size - 1) // args.chunk_size + for ci in range(n_chunks): + chunk = records[ci * args.chunk_size:(ci + 1) * args.chunk_size] + sys.stderr.write(f'[reflexion] chunk {ci+1}/{n_chunks} ({len(chunk)} problems)\n') + recs = process_chunk(sampler, api, chunk, gen_dp, args) + for rec in recs: # incremental write per problem + out_f.write(json.dumps(rec, ensure_ascii=False) + '\n') + out_f.flush() + _update_summary(summ, recs) + rep = _summary_report(summ) + sys.stderr.write( + f' running: failed={rep["n_failed_first_try"]} ' + f'(easy={rep["n_failed_but_easy"]}) hard={rep["n_hard"]} ' + f'base_pass={rep["avg_baseline_pass_on_hard"]:.3f} ' + f'skill_pass={rep["avg_best_with_skill_pass_on_hard"]:.3f} ' + f'helped={rep["n_helped_by_skill"]} rescued={rep["n_rescued_from_zero"]}\n') + + report = _summary_report(summ) + out_f.write(json.dumps(report, ensure_ascii=False) + '\n') + out_f.flush() + + print('\n' + '=' * 60) + print(f'Reflexion Phase-0 — model={GEN_MODEL_ID}, dataset=aops, n={report["n_total"]}') + print('=' * 60) + print(f'solved on first try : {report["n_solved_first_try"]}/{report["n_total"]}') + print(f'failed first try : {report["n_failed_first_try"]} ' + f'(easy, excluded: {report["n_failed_but_easy"]})') + print(f'hard (baseline pass@{args.pass_k}<= {args.hard_baseline_max}) : {report["n_hard"]}') + print(f' avg baseline pass@{args.pass_k:<2} : {report["avg_baseline_pass_on_hard"]:.4f}') + print(f' avg best-skill pass@{args.pass_k:<2} : {report["avg_best_with_skill_pass_on_hard"]:.4f}') + print(f' avg best marginal : {report["avg_best_marginal_on_hard"]:+.4f}') + print(f' helped by a skill : {report["n_helped_by_skill"]}/{report["n_hard"]}') + print(f' rescued from 0 pass : {report["n_rescued_from_zero"]}/{report["n_hard"]}') + print(f'\n[output] {args.output}') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/legacy/make_embedding_dataset.py b/cookbook/exp/legacy/make_embedding_dataset.py new file mode 100644 index 000000000..847f222fc --- /dev/null +++ b/cookbook/exp/legacy/make_embedding_dataset.py @@ -0,0 +1,758 @@ +"""Offline compression pipeline: raw datasets → condenser → pre-compressed embedding dataset. + +Loads think/index/hard datasets, compresses query/cot/negatives via vLLM condenser +with API fallback, saves a single HF Dataset ready for embedding training. + +Output schema: {anchor_text, positive_text, negative_texts, source} + +Launch (8 GPUs — 4 for vLLM condenser): + python cookbook/exp/embedding/make_embedding_dataset.py +""" +import json +import os +import re +import sys +import threading +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path +from typing import Any, Dict, List, Optional + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import SamplingParams +from twinkle.sampler import vLLMSampler +from twinkle.template import Qwen3_5Template +from twinkle.utils.parallel import PosixFileLock +from twinkle_agentic.protocol.openai import OpenAI as OpenAIClient + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from dataset_think import get_dataset as get_dataset_think # noqa: E402 +from dataset_index import get_dataset as get_dataset_index # noqa: E402 +from dataset_hard import get_dataset as get_dataset_hard # noqa: E402 + +logger = get_logger() + +# -- Model config ------------------------------------------------------------- +CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') +TEMPLATE_NAME = 'Qwen3_5Template' + +# -- GPU placement (condenser only) ------------------------------------------- +CONDENSER_GPUS = int(os.environ.get('CONDENSER_GPUS', 8)) + +# -- Dataset caps ------------------------------------------------------------- +TOTAL_SAMPLES: Optional[int] = None +THINK_CAP: Optional[int] = int(os.environ.get('THINK_CAP', 100_000)) +INDEX_CAP: Optional[int] = int(os.environ.get('INDEX_CAP', 100_000)) +HARD_CAP: Optional[int] = int(os.environ.get('HARD_CAP', 0)) or None +HARD_MAX_NEGATIVES = int(os.environ.get('HARD_MAX_NEGATIVES', 8)) + +# -- Compression params ------------------------------------------------------- +MIN_TEXT_CHARS = 256 +DATASET_MAX_TOKENS = 32768 +COMPRESS_TEMPERATURE = 0.2 +COMPRESS_TOP_P = 0.5 +COMPRESS_MAX_MODEL_LEN = 32768 +BATCH_SIZE = int(os.environ.get('COMPRESS_BATCH_SIZE', 128)) + +# -- API fallback ------------------------------------------------------------- +COMPRESS_API_KEY = os.environ.get('COMPRESS_API_KEY', '') +COMPRESS_BASE_URL = os.environ.get('COMPRESS_BASE_URL', 'https://dashscope.aliyuncs.com/compatible-mode/v1') +COMPRESS_MODEL = os.environ.get('COMPRESS_MODEL', 'qwen3.7-max') +API_MIN_INTERVAL = float(os.environ.get('API_MIN_INTERVAL', 0.1)) +API_CONCURRENCY = int(os.environ.get('API_CONCURRENCY', 24)) +SAMPLER_TIMEOUT = float(os.environ.get('SAMPLER_TIMEOUT', 300)) + +# -- Output ------------------------------------------------------------------- +OUTPUT_DIR = os.environ.get('EMB_DATASET_OUTPUT', './output/embedding_dataset') +RESULTS_JSONL = f'{OUTPUT_DIR}/results.jsonl' +PROGRESS_FILE = f'{OUTPUT_DIR}/progress.json' + +# ============================================================================= +# Prompts +# ============================================================================= + +COMPRESS_SYSTEM = """\ +You are a compression and summary assistant. For the (query, source) pair, emit a Markdown \ +answer with TWO sections, designed to pair with the `extract_compressed` tool: \ +the reader absorbs `## Summary` directly, then calls `extract_compressed` \ +on any topic-key listed under `## More` to recover its \ +fuller content. + + `## Summary` — extreme-density text the reader reads directly. + `## More` — a topic index whose keys are valid arguments \ +to `extract_compressed` for recovering material not captured inline. + +Together the two sections must form a COMPLETE, NON-DISTORTING inventory of the \ +source for the query — nothing essential lost, nothing implied that the source \ +does not support. NO preamble, NO meta-commentary, NO code fences wrapping the \ +whole output. + +Output skeleton: + +## Summary +Topic: + + +## More +- : +- ... + +Format selection for the inline body (pick the MOST COMPACT form per query, mix \ +when helpful): +- Interface / signature → code notation directly: `func(a:int)->str` +- Factual / entity → telegraphic prose; drop function words; ":" for "is", "," \ +for "has" +- Skill / how-to / usage → lead with `Use when: `; numbered telegraphic \ +steps `1.do X 2.then Y`; close with `Output: ` when relevant +- Procedural → numbered short steps +- Analytical / design → hierarchical bullets with abbreviations + +`## Summary` rules: +1. TOPIC LINE — line 1 is ALWAYS `Topic: `, even when the \ +query is narrow. Anchors both the reader and the tool. +2. DENSITY — every token in the body carries query-relevant signal; cut filler. +3. PRIMARY-COMPLETE — never silently drop a fact essential to answering the \ +query. Anything cut for length MUST appear as a key under \ +`## More`. +4. NON-MISLEADING — phrasing must not let the reader infer anything the source \ +does not support; partial truths that mislead are worse than honest omissions \ +flagged in the index. +5. SELF-CONTAINED — the reader can act on the answer without re-opening the source. +6. FAITHFUL — only content the source supports; no fabrication, no extrapolation. +7. LANGUAGE — match the source language. +8. NO outer code fences around the whole answer; no meta-commentary. + +`## More` rules (MANDATORY — this section is never omitted): +1. FORMAT — each bullet is `- : `: + • topic-key — short, unambiguous, grounded in source vocabulary so the \ +`extract_compressed` tool can locate the aspect (e.g. `decorators`, \ +`error handling`, `pitfalls`). + • hint — tells WHAT the reader gains by expanding (concrete numbers, code \ +listings, secondary cases, edge details, related context, …); do NOT restate \ +the inline answer. +2. CRITERION — each bullet names an aspect that EXISTS in the source but is \ +NOT fully captured inline. Material that genuinely fits inline without \ +distortion MUST NOT be duplicated here. +3. FAITHFUL — hints must be grounded in the source; never speculate or invent. +4. ORDER — by relevance to the query, then by importance. +5. EMPTY CASE — if the source is so short / single-purpose that everything \ +fits inline, write a single line `- (none)`. + +Now begin.\ +""" + +COMPRESS_USER = ( + 'Downstream model will read your compressed block to decide whether to ' + 'expand it. Compress faithfully: preserve the passage topic + core facts. ' + 'Do NOT invent facts. Do NOT drop major facts. Do NOT write meta-commentary ' + 'about the Query (never write "Query info: absent", "no X mention", etc.); ' + 'if the passage does not address the Query, still summarize the passage. ' + 'CRITICAL LANGUAGE RULE: detect the dominant language of the Passage ' + '(NOT the Query, NOT this instruction) and write the ENTIRE output in that ' + 'same language; English passage → English output, Chinese passage → ' + 'Chinese output, Japanese passage → Japanese output. NEVER translate, ' + 'NEVER mix languages, NEVER copy these instructions into the output.\n\n' + '## Query (ordering hint only — still summarize the whole passage)\n{query}\n\n' + '## Passage\n{text}') + +EMBED_QUERY_Q = ( + 'Summarize this query for retrieval. ' + 'The body of ## Summary MUST follow this EXACT 4-line template — ' + 'do NOT emit "Use when:", numbered procedure steps, or "Output:":\n' + 'Topic: \n' + 'Problem: \n' + 'Skill: \n' + 'Knowledge: \n' + 'Then emit the mandatory ## More section as usual. ' + 'Topic must name the specific pattern, never generic labels.') + +EMBED_QUERY_COT = ( + 'Summarize this reasoning trace for retrieval. ' + 'The body of ## Summary MUST follow this EXACT 4-line template — ' + 'do NOT emit "Use when:", numbered procedure steps, or "Output:":\n' + 'Topic: \n' + 'Problem: \n' + 'Skill: \n' + 'Knowledge: \n' + 'Then emit the mandatory ## More section as usual. ' + 'Topic must name the specific pattern, never generic labels.') + +EMBED_QUERY_Q_LEGACY = ( + 'What problem does this passage address, and what skill or method is needed? ' + 'Topic must name the specific pattern, never generic labels. ' + 'Compress into a retrieval-friendly need description.') + +EMBED_QUERY_COT_LEGACY = ( + 'Extract the reusable skill: trigger conditions, key steps, and expected output. ' + 'Topic names the method/pattern; format as "Use when: ...", numbered steps, ' + '"Output: ...". Compress into a standardized procedure for retrieval.') + +EMBED_QUERY_REASONIR_Q = ( + 'Extract the abstract PROBLEM TYPE from this query. ' + 'IGNORE all specific numbers, values, variable names, and parameters — ' + 'focus ONLY on the CLASS of problem and the METHODOLOGY required. ' + 'The body of ## Summary MUST follow this EXACT 4-line template:\n' + 'Topic: \n' + 'Problem: \n' + 'Skill: \n' + 'Knowledge: \n' + 'Then emit the mandatory ## More section as usual. ' + 'Topic must name the method class, never mention specific numbers.') + +EMBED_QUERY_REASONIR_COT = ( + 'Extract the abstract METHODOLOGY demonstrated in this solution. ' + 'IGNORE all specific numbers, values, and computed results — ' + 'focus ONLY on the general TECHNIQUE and key reasoning STEPS. ' + 'The body of ## Summary MUST follow this EXACT 4-line template:\n' + 'Topic: \n' + 'Problem: \n' + 'Skill: \n' + 'Knowledge: \n' + 'Then emit the mandatory ## More section as usual. ' + 'Topic must name the method class, never mention specific numbers.') + + +# ============================================================================= +# Validation & API fallback +# ============================================================================= + +_LEGACY_USE_WHEN_RE = re.compile(r'(?im)^\s*Use when\s*:') +_SCHEMA_MARKERS = ('Problem:', 'Skill:', 'Knowledge:') + + +def _is_truncated_compression(text: str, schema: str = 'new') -> bool: + if not text or not text.strip(): + return True + if '## More' not in text or '## Summary' not in text: + return True + after_more = text.split('## More', 1)[1].strip() + if not after_more: + return True + last_line = after_more.splitlines()[-1].strip() + if not (last_line.startswith('-') or last_line.endswith(')')): + return True + if schema == 'new': + summary_body = text.split('## Summary', 1)[1].split('## More', 1)[0] + if _LEGACY_USE_WHEN_RE.search(summary_body): + return True + if not all(marker in summary_body for marker in _SCHEMA_MARKERS): + return True + return False + + +_api_semaphore = threading.Semaphore(API_CONCURRENCY) +_api_bucket_lock = threading.Lock() +_api_tokens = [float(API_CONCURRENCY)] +_api_last_refill = [time.monotonic()] + + +def _api_throttle(): + """Token-bucket rate limiter: API_CONCURRENCY requests per API_MIN_INTERVAL*API_CONCURRENCY window.""" + _api_semaphore.acquire() + try: + with _api_bucket_lock: + now = time.monotonic() + elapsed = now - _api_last_refill[0] + refill = elapsed / API_MIN_INTERVAL + _api_tokens[0] = min(float(API_CONCURRENCY), _api_tokens[0] + refill) + _api_last_refill[0] = now + if _api_tokens[0] >= 1.0: + _api_tokens[0] -= 1.0 + else: + wait = (1.0 - _api_tokens[0]) * API_MIN_INTERVAL + _api_tokens[0] = 0.0 + time.sleep(wait) + finally: + _api_semaphore.release() + + +def _api_compress(api_client: OpenAIClient, prompt: Dict[str, Any]) -> Optional[str]: + _api_throttle() + trajectory = {'messages': prompt['messages']} + sp = SamplingParams(temperature=0.2, max_tokens=8192) + try: + reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) + except Exception as exc: + logger.warning(f'[api_fallback] error: {exc}') + return None + content = (reply.get('content') or '').strip() + if not content: + return None + m = re.match(r'^```[a-zA-Z]*\n(.*?)\n```\s*$', content, re.DOTALL) + if m: + content = m.group(1).strip() + return content + + +# ============================================================================= +# Core compression logic +# ============================================================================= + +def _extract_query_cot(row: Dict[str, Any]): + messages = row.get('messages') or [] + query, cot = '', '' + for m in messages: + if not isinstance(m, dict): + continue + role = m.get('role') or '' + if role == 'user' and not query: + query = (m.get('content') or '').strip() + elif role == 'assistant': + cot = (m.get('reasoning_content') or '').strip() + break + return query, cot + + +def _compress_batch_phase1( + rows: List[Dict[str, Any]], + condenser_sampler, + compress_params: SamplingParams, + special_tokens: set, + source_type: str, +) -> Optional[Dict[str, Any]]: + """Phase 1 (GPU): build prompts → vLLM sample → validate. Returns state for phase 2.""" + _MAX_COT_CHARS = 30_000 + + if source_type == 'hard': + return _compress_hard_phase1(rows, condenser_sampler, compress_params, + special_tokens, source_type) + + prompts: List[Optional[Dict[str, Any]]] = [] + meta: List[Dict[str, Any]] = [] + for i, row in enumerate(rows): + query, cot = _extract_query_cot(row) + if not query or len(cot) < MIN_TEXT_CHARS or len(cot) > _MAX_COT_CHARS: + continue + schema = 'legacy' if (i % 2 == 0) else 'new' + q_hint = EMBED_QUERY_Q_LEGACY if schema == 'legacy' else EMBED_QUERY_Q + c_hint = EMBED_QUERY_COT_LEGACY if schema == 'legacy' else EMBED_QUERY_COT + + if len(query) < MIN_TEXT_CHARS: + prompts.append(None) + else: + user = COMPRESS_USER.format(query=q_hint, text=query) + prompts.append({'messages': [ + {'role': 'system', 'content': COMPRESS_SYSTEM}, + {'role': 'user', 'content': user}, + ]}) + user_c = COMPRESS_USER.format(query=c_hint, text=cot) + prompts.append({'messages': [ + {'role': 'system', 'content': COMPRESS_SYSTEM}, + {'role': 'user', 'content': user_c}, + ]}) + meta.append({'query_raw': query, 'cot_raw': cot, 'schema': schema, + 'q_hint': q_hint, 'source': source_type, + 'row_id': row.get('id', str(i))}) + + if not prompts: + return {'final': []} + + sampler_input = [p for p in prompts if p is not None] + sampler_pos = [ri for ri, p in enumerate(prompts) if p is not None] + try: + sampler_responses = condenser_sampler.sample(sampler_input, compress_params) + except Exception as exc: + logger.warning(f'[compress] sampler error: {exc}') + sampler_responses = [None] * len(sampler_input) + + responses = [None] * len(prompts) + for resp, pos in zip(sampler_responses, sampler_pos): + responses[pos] = resp + + decoded: List[str] = [] + fallback_indices: List[int] = [] + for ri in range(len(prompts)): + pair_idx = ri // 2 + schema = meta[pair_idx]['schema'] + if prompts[ri] is None: + decoded.append(meta[pair_idx]['query_raw']) + continue + resp = responses[ri] + seq = resp.sequences[0] if resp and resp.sequences else None + text = '' + if seq and seq.stop_reason != 'length' and seq.decoded: + text = seq.decoded + for tok in special_tokens: + text = text.replace(tok, '') + text = text.rstrip() + if not _is_truncated_compression(text, schema): + decoded.append(text) + else: + decoded.append('') + fallback_indices.append(ri) + + return {'prompts': prompts, 'meta': meta, 'decoded': decoded, + 'fallback_indices': fallback_indices} + + +def _compress_batch_phase2( + state: Dict[str, Any], + api_client: OpenAIClient, +) -> List[Dict[str, Any]]: + """Phase 2 (no GPU): API fallback → build results.""" + if 'final' in state: + return state['final'] + + prompts = state['prompts'] + decoded = state['decoded'] + fallback_indices = state['fallback_indices'] + is_hard = state.get('hard', False) + meta = state.get('meta') # None for hard + + # Track which prompts used API fallback + api_set: set = set() + if fallback_indices: + api_futures = {} + with ThreadPoolExecutor(max_workers=API_CONCURRENCY) as pool: + for ri in fallback_indices: + api_futures[pool.submit(_api_compress, api_client, prompts[ri])] = ri + for fut in as_completed(api_futures): + ri = api_futures[fut] + api_result = fut.result() + schema = 'new' if is_hard else meta[ri // 2]['schema'] + if api_result and not _is_truncated_compression(api_result, schema): + decoded[ri] = api_result + api_set.add(ri) + + state['api_set'] = api_set + if is_hard: + return _build_hard_results(state) + return _build_think_index_results(state) + + +def _build_think_index_results(state: Dict[str, Any]) -> List[Dict[str, Any]]: + meta = state['meta'] + decoded = state['decoded'] + api_set = state.get('api_set', set()) + results = [] + for pair_idx in range(len(meta)): + q_text = decoded[pair_idx * 2] + c_text = decoded[pair_idx * 2 + 1] + if not q_text or not c_text: + continue + q_method = 'api' if (pair_idx * 2) in api_set else 'vllm' + c_method = 'api' if (pair_idx * 2 + 1) in api_set else 'vllm' + results.append({ + 'anchor_text': q_text, + 'positive_text': c_text, + 'negative_texts': [], + 'source': meta[pair_idx]['source'], + 'query_raw': meta[pair_idx]['query_raw'], + 'cot_raw': meta[pair_idx]['cot_raw'], + 'anchor_method': q_method, + 'positive_method': c_method, + }) + return results + + +def _build_hard_results(state: Dict[str, Any]) -> List[Dict[str, Any]]: + group_sizes = state['group_sizes'] + decoded = state['decoded'] + source_type = state['source_type'] + raw_groups = state['raw_groups'] + api_set = state.get('api_set', set()) + results = [] + offset = 0 + for gi, gs in enumerate(group_sizes): + q_text = decoded[offset] + c_text = decoded[offset + 1] + if not q_text or not c_text: + offset += gs + continue + neg_texts = [] + neg_raws = [] + neg_methods = [] + for ni in range(2, gs): + nt = decoded[offset + ni] + if nt: + neg_texts.append(nt) + neg_raws.append(raw_groups[gi]['negs_raw'][ni - 2]) + neg_methods.append('api' if (offset + ni) in api_set else 'vllm') + q_method = 'api' if offset in api_set else 'vllm' + c_method = 'api' if (offset + 1) in api_set else 'vllm' + results.append({ + 'anchor_text': q_text, + 'positive_text': c_text, + 'negative_texts': neg_texts, + 'source': source_type, + 'query_raw': raw_groups[gi]['query_raw'], + 'cot_raw': raw_groups[gi]['cot_raw'], + 'negs_raw': neg_raws, + 'anchor_method': q_method, + 'positive_method': c_method, + 'neg_methods': neg_methods, + }) + offset += gs + return results + + +def _compress_hard_phase1( + rows: List[Dict[str, Any]], + condenser_sampler, + compress_params: SamplingParams, + special_tokens: set, + source_type: str, +) -> Dict[str, Any]: + """Phase 1 for hard rows: vLLM sample + validate. Returns state for phase 2.""" + _MAX_COT_CHARS = 30_000 + + prompts: List[Dict[str, Any]] = [] + group_sizes: List[int] = [] + row_ids: List[str] = [] + raw_groups: List[Dict[str, Any]] = [] + + for row in rows: + query, cot = _extract_query_cot(row) + if not query or not cot or len(cot) > _MAX_COT_CHARS: + continue + negatives = row.get('negatives') or [] + valid_negs = [n for n in negatives + if n and len(n) <= _MAX_COT_CHARS] + + user_q = COMPRESS_USER.format(query=EMBED_QUERY_REASONIR_Q, text=query) + prompts.append({'messages': [ + {'role': 'system', 'content': COMPRESS_SYSTEM}, + {'role': 'user', 'content': user_q}, + ]}) + user_c = COMPRESS_USER.format(query=EMBED_QUERY_REASONIR_COT, text=cot) + prompts.append({'messages': [ + {'role': 'system', 'content': COMPRESS_SYSTEM}, + {'role': 'user', 'content': user_c}, + ]}) + for neg in valid_negs: + user_n = COMPRESS_USER.format(query=EMBED_QUERY_REASONIR_COT, text=neg) + prompts.append({'messages': [ + {'role': 'system', 'content': COMPRESS_SYSTEM}, + {'role': 'user', 'content': user_n}, + ]}) + group_sizes.append(2 + len(valid_negs)) + row_ids.append(row.get('id', '')) + raw_groups.append({'query_raw': query, 'cot_raw': cot, 'negs_raw': valid_negs}) + + if not prompts: + return {'hard': True, 'prompts': [], 'group_sizes': [], 'row_ids': [], + 'decoded': [], 'fallback_indices': [], 'source_type': source_type, + 'raw_groups': []} + + try: + responses = condenser_sampler.sample(prompts, compress_params) + except Exception as exc: + logger.warning(f'[compress-hard] sampler error: {exc}') + responses = [None] * len(prompts) + + decoded: List[str] = [] + fallback_indices: List[int] = [] + for ri, resp in enumerate(responses): + seq = resp.sequences[0] if resp and resp.sequences else None + text = '' + if seq and seq.stop_reason != 'length' and seq.decoded: + text = seq.decoded + for tok in special_tokens: + text = text.replace(tok, '') + text = text.rstrip() + if text and not _is_truncated_compression(text, 'new'): + decoded.append(text) + else: + decoded.append('') + fallback_indices.append(ri) + + return {'hard': True, 'prompts': prompts, 'group_sizes': group_sizes, + 'row_ids': row_ids, 'decoded': decoded, + 'fallback_indices': fallback_indices, 'source_type': source_type, + 'raw_groups': raw_groups} + + +# ============================================================================= +# Main pipeline +# ============================================================================= + +def main(): + device_groups = [ + DeviceGroup(name='condenser_sampler', + ranks=list(range(CONDENSER_GPUS)), + device_type='GPU'), + ] + condenser_mesh = DeviceMesh.from_sizes( + world_size=CONDENSER_GPUS, dp_size=CONDENSER_GPUS) + twinkle.initialize(mode='ray', nproc_per_node=CONDENSER_GPUS, groups=device_groups) + + # -- Load raw datasets ---------------------------------------------------- + from datasets import Dataset as HFDataset + + dataset_think = get_dataset_think(total=TOTAL_SAMPLES, load_from_cache_file=True) + if THINK_CAP and len(dataset_think.dataset) > THINK_CAP: + dataset_think.dataset = dataset_think.dataset.select(range(THINK_CAP)) + ds_think = dataset_think.dataset + logger.info(f'[load] think={len(ds_think)}') + + ds_index_obj = get_dataset_index(total=None, load_from_cache_file=True) + ds_index = ds_index_obj.dataset + if INDEX_CAP and len(ds_index) > INDEX_CAP: + ds_index = ds_index.select(range(INDEX_CAP)) + logger.info(f'[load] index={len(ds_index)}') + + ds_hard_raw = get_dataset_hard(max_negatives=HARD_MAX_NEGATIVES, load_from_cache_file=True) + if HARD_CAP and len(ds_hard_raw) > HARD_CAP: + ds_hard_raw = ds_hard_raw.select(range(HARD_CAP)) + n_hard = len(ds_hard_raw) + logger.info(f'[load] hard={n_hard}') + + # Convert hard to messages schema + hard_rows_list = [] + if n_hard > 0: + h_ids = ds_hard_raw['id'] + h_queries = ds_hard_raw['query'] + h_cots = ds_hard_raw['cot'] + h_responses = ds_hard_raw['response'] if 'response' in ds_hard_raw.column_names else [''] * n_hard + h_negatives = ds_hard_raw['negatives'] + for i in range(n_hard): + hard_rows_list.append({ + 'id': h_ids[i], + 'messages': [ + {'role': 'user', 'content': h_queries[i]}, + {'role': 'assistant', 'reasoning_content': h_cots[i], + 'content': h_responses[i] or ''}, + ], + 'negatives': h_negatives[i], + }) + + # Batch-convert HF Datasets to list-of-dicts + def _ds_to_rows(ds): + return [dict(zip(ds.column_names, vals)) for vals in zip(*(ds[c] for c in ds.column_names))] + + think_rows = _ds_to_rows(ds_think) + index_rows = _ds_to_rows(ds_index) + + # -- Setup condenser ------------------------------------------------------ + condenser_template = Qwen3_5Template( + model_id=CONDENSE_MODEL_ID, max_length=DATASET_MAX_TOKENS, + enable_thinking=False, truncation_strategy='delete') + special_tokens = set(condenser_template.tokenizer.all_special_tokens) + + condenser_sampler = vLLMSampler( + model_id=CONDENSE_MODEL_ID, + engine_args={'gpu_memory_utilization': 0.8, 'max_model_len': COMPRESS_MAX_MODEL_LEN}, + device_mesh=condenser_mesh, + remote_group='condenser_sampler', + ) + condenser_sampler.set_template( + TEMPLATE_NAME, model_id=CONDENSE_MODEL_ID, enable_thinking=False, + truncation_strategy='delete', max_length=DATASET_MAX_TOKENS) + condenser_sampler._ray_get_timeout = SAMPLER_TIMEOUT + compress_params = SamplingParams( + max_tokens=8192, temperature=COMPRESS_TEMPERATURE, + top_p=COMPRESS_TOP_P, num_samples=1) + + api_client = OpenAIClient( + model=COMPRESS_MODEL, api_key=COMPRESS_API_KEY, base_url=COMPRESS_BASE_URL) + + # -- Resume support ---------------------------------------------------------- + os.makedirs(OUTPUT_DIR, exist_ok=True) + progress = {'think': 0, 'index': 0, 'hard': 0} + if os.path.exists(PROGRESS_FILE): + with open(PROGRESS_FILE, 'r') as f: + progress = json.load(f) + logger.info(f'[resume] loaded progress: {progress}') + + _results_lock = PosixFileLock(RESULTS_JSONL + '.lock') + + def _flush_results(records: List[Dict[str, Any]]): + if not records: + return + lines = [json.dumps(r, ensure_ascii=False) + '\n' for r in records] + with _results_lock: + with open(RESULTS_JSONL, 'a', encoding='utf-8') as f: + f.writelines(lines) + + def _save_progress(): + tmp = PROGRESS_FILE + '.tmp' + with open(tmp, 'w') as f: + json.dump(progress, f) + os.replace(tmp, PROGRESS_FILE) + + # -- Process in batches (pipelined: vLLM batch N+1 overlaps API fallback N) - + total_flushed = 0 + if os.path.exists(RESULTS_JSONL): + with open(RESULTS_JSONL, 'r', encoding='utf-8') as f: + total_flushed = sum(1 for l in f if l.strip()) + if total_flushed: + logger.info(f'[resume] {total_flushed} records already in results.jsonl') + + def _process_source(rows, source_type, label): + nonlocal total_flushed + n_total = len(rows) + skip = progress.get(source_type, 0) + if skip >= n_total: + logger.info(f'[{label}] skipped (already done {skip}/{n_total})') + return + if skip > 0: + logger.info(f'[{label}] resuming from row {skip}/{n_total}') + + bg_pool = ThreadPoolExecutor(max_workers=1) + pending = None # (future, batch_start, batch_len) + + def _drain_pending(): + nonlocal total_flushed, pending + if pending is None: + return + fut, p_start, p_len = pending + batch_results = fut.result() + _flush_results(batch_results) + total_flushed += len(batch_results) + progress[source_type] = p_start + p_len + _save_progress() + pending = None + + for start in range(skip, n_total, BATCH_SIZE): + batch = rows[start:start + BATCH_SIZE] + state = _compress_batch_phase1( + batch, condenser_sampler, compress_params, + special_tokens, source_type) + _drain_pending() + pending = ( + bg_pool.submit(_compress_batch_phase2, state, api_client), + start, len(batch)) + n_done = start + len(batch) + if n_done % (BATCH_SIZE * 10) == 0 or n_done >= n_total: + logger.info(f'[{label}] {n_done}/{n_total} vLLM done, ' + f'{total_flushed} records flushed (last batch pending)') + + _drain_pending() + bg_pool.shutdown(wait=False) + logger.info(f'[{label}] complete, {total_flushed} total records flushed') + + _process_source(hard_rows_list, 'hard', 'hard') + _process_source(think_rows, 'think', 'think') + _process_source(index_rows, 'index', 'index') + + # -- Convert JSONL → HF Dataset ------------------------------------------- + logger.info(f'[save] converting results.jsonl to HF Dataset...') + all_results = [] + with open(RESULTS_JSONL, 'r', encoding='utf-8') as f: + for line_no, line in enumerate(f, 1): + if not line.strip(): + continue + try: + all_results.append(json.loads(line)) + except json.JSONDecodeError: + logger.warning(f'[save] skipping malformed line {line_no} (truncated resume?)') + logger.info(f'[save] total records: {len(all_results)}') + out_ds = HFDataset.from_dict({ + 'anchor_text': [r['anchor_text'] for r in all_results], + 'positive_text': [r['positive_text'] for r in all_results], + 'negative_texts': [r['negative_texts'] for r in all_results], + 'source': [r['source'] for r in all_results], + 'query_raw': [r.get('query_raw', '') for r in all_results], + 'cot_raw': [r.get('cot_raw', '') for r in all_results], + 'negs_raw': [r.get('negs_raw', []) for r in all_results], + }) + out_ds.save_to_disk(OUTPUT_DIR + '/dataset') + logger.info(f'[save] dataset saved to {OUTPUT_DIR}/dataset') + logger.info(f'[stats] think={sum(1 for r in all_results if r["source"]=="think")} ' + f'index={sum(1 for r in all_results if r["source"]=="index")} ' + f'hard={sum(1 for r in all_results if r["source"]=="hard")}') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/legacy/train_embedding_full_ddp.py b/cookbook/exp/legacy/train_embedding_full_ddp.py new file mode 100644 index 000000000..97ab3b128 --- /dev/null +++ b/cookbook/exp/legacy/train_embedding_full_ddp.py @@ -0,0 +1,270 @@ +"""Full-parameter embedding training on pre-compressed dataset. + +Reads the pre-compressed HF Dataset produced by make_embedding_dataset.py, +encodes features, trains with InfoNCE loss. + +Architecture (4 GPUs): + - Ranks 0-3: Trainable embedding model, InfoNCE loss. + +Launch: + python cookbook/exp/embedding/train_embedding_full_ddp.py +""" +import os +import time +from typing import Any, Dict, List, Literal, Optional + +import swanlab + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_device_placement, get_logger +from twinkle.loss import InfonceLoss +from twinkle.metric import EmbeddingMetric +from twinkle.model import TransformersModel +from twinkle.processor import InputProcessor +from twinkle.template import Qwen3_5Template, Template + +logger = get_logger() + +# -- Backend selection -------------------------------------------------------- +BACKEND: Literal['transformers', 'megatron'] = 'transformers' + +MODEL_ID = os.environ.get('MODEL_ID', 'ms://Qwen/Qwen3.5-4B') + +# -- GPU placement ------------------------------------------------------------ +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', 64)) +LEARNING_RATE = 1e-5 +GRADIENT_ACCUMULATION_STEPS = 1 +LOG_INTERVAL = 2 +SAVE_INTERVAL = 2000 +NUM_EPOCHS = 1 + +# -- Dataset path (output of make_embedding_dataset.py) ----------------------- +DATASET_PATH = os.environ.get('EMB_DATASET_PATH', 'ms://twinkle-kit/qth-embedding') +MIX_SHUFFLE_SEED = 42 + +# -- Resume from checkpoint --------------------------------------------------- +RESUME_CHECKPOINT = os.environ.get('RESUME_CHECKPOINT', '') +RESUME_STEP = int(os.environ.get('RESUME_STEP', 0)) + +# -- Output ------------------------------------------------------------------- +OUTPUT_DIR = f'./output/embedding_full_{BACKEND}' + + +# ============================================================================= +# Model builders +# ============================================================================= + +def build_model(device_mesh: DeviceMesh): + model_id = RESUME_CHECKPOINT if RESUME_CHECKPOINT else MODEL_ID + if BACKEND == 'transformers': + model = TransformersModel( + model_id=model_id, + device_mesh=device_mesh, + remote_group='model', + ddp_config={'find_unused_parameters': True}, + ) + from twinkle.patch.no_split_modules import NoSplitModulesPatch + model.apply_patch(NoSplitModulesPatch({'Qwen3_5DecoderLayer'})) + return model + if BACKEND == 'megatron': + from twinkle.model import MegatronModel + return MegatronModel( + model_id=MODEL_ID, + device_mesh=device_mesh, + remote_group='model', + mixed_precision='bf16', + variable_seq_lengths=True, + ) + raise ValueError(f'Unknown BACKEND={BACKEND!r}') + + +def setup_optimizer(model, total_steps: int): + if BACKEND == 'transformers': + model.set_optimizer(optimizer_cls='AdamW', lr=LEARNING_RATE) + model.set_lr_scheduler( + scheduler_cls='CosineWarmupScheduler', + num_warmup_steps=200, + num_training_steps=total_steps, + ) + return + if BACKEND == 'megatron': + model.set_optimizer(optimizer_cls='default', lr=LEARNING_RATE) + model.set_lr_scheduler( + scheduler_cls='default', + lr_warmup_steps=50, + lr_decay_steps=total_steps, + ) + return + raise ValueError(f'Unknown BACKEND={BACKEND!r}') + + +def save_checkpoint(model, name: str): + model.save(name, output_dir=OUTPUT_DIR) + + +# ============================================================================= +# Feature encoding +# ============================================================================= + +def _get_first_feature(decoded_text: str, template: Template, role: str) -> Optional[Dict[str, Any]]: + if not decoded_text: + return None + if role == 'anchor': + feat = template.encode({'messages': [ + {'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 + + +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 + + +# ============================================================================= +# Main training +# ============================================================================= + +def train(): + device_groups = [ + DeviceGroup(name='model', + ranks=list(range(MODEL_GPUS)), + device_type='GPU'), + ] + model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, dp_size=MODEL_GPUS) + twinkle.initialize(mode='ray', nproc_per_node=MODEL_GPUS, groups=device_groups) + + # -- Load pre-compressed dataset ------------------------------------------ + from twinkle.dataset import Dataset as TwinkleDataset, DatasetMeta + logger.info(f'[data] loading pre-compressed dataset from {DATASET_PATH}') + dataset = TwinkleDataset(DatasetMeta(dataset_id=DATASET_PATH), download_mode='force_redownload') + dataset = dataset.dataset.shuffle(seed=MIX_SHUFFLE_SEED) + logger.info(f'[data] {len(dataset)} rows loaded') + + # -- Compute steps -------------------------------------------------------- + rows_per_step = BATCH_SIZE + total_steps = (len(dataset) // rows_per_step) * NUM_EPOCHS + optimizer_steps = total_steps // GRADIENT_ACCUMULATION_STEPS + + # -- Model ---------------------------------------------------------------- + model = build_model(model_mesh) + model.set_processor(InputProcessor) + model.set_loss(InfonceLoss, temperature=TEMPERATURE, use_batch=True, + hard_negatives=HARD_NEGATIVES) + setup_optimizer(model, optimizer_steps) + model.add_metric(EmbeddingMetric, is_training=True) + + 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 steps: {total_steps}, optimizer steps: {optimizer_steps}') + + swanlab.init(project='twinkle', config={ + 'backend': BACKEND, + 'model_id': MODEL_ID, + 'batch_size': BATCH_SIZE, + 'lr': LEARNING_RATE, + 'temperature': TEMPERATURE, + 'emb_max_length': EMB_MAX_LENGTH, + 'dataset_path': DATASET_PATH, + }) + + # -- Train loop ----------------------------------------------------------- + cur_step = 0 + _skip_rows = RESUME_STEP * rows_per_step # approximate rows to skip + + for epoch in range(NUM_EPOCHS): + for start in range(0, len(dataset), rows_per_step): + if start < _skip_rows: + continue + + 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() + features = _encode_batch(rows_list, emb_template) + t_encode = time.monotonic() - t0 + + if len(features) < 4: + continue + + t1 = time.monotonic() + model.forward_backward(inputs=features, task='embedding') + model.clip_grad_and_step( + gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) + t_train = time.monotonic() - t1 + cur_step += 1 + + if cur_step % LOG_INTERVAL == 0: + metric = model.calculate_metric(is_training=True) + logger.info( + f'Epoch {epoch} Step {cur_step}/{total_steps}, ' + f'metric: {metric} | ' + f'encode={t_encode:.2f}s train={t_train:.2f}s') + log_dict = {} + for k, v in metric.items(): + if not v: + continue + try: + log_dict[k] = float(v) + except (ValueError, TypeError): + pass + log_dict['epoch'] = epoch + log_dict['encode_sec'] = round(t_encode, 3) + log_dict['train_sec'] = round(t_train, 3) + swanlab.log(log_dict, step=cur_step) + if cur_step % SAVE_INTERVAL == 0: + save_checkpoint(model, f'step_{cur_step}') + + save_checkpoint(model, 'last-checkpoint') + # Force sync: resolve any pending lazy remote calls (save) before exit + model.calculate_metric(is_training=True) + logger.info(f'Training complete. Final step: {cur_step}') + + +if __name__ == '__main__': + train() diff --git a/cookbook/exp/legacy/train_reflexion_skill.py b/cookbook/exp/legacy/train_reflexion_skill.py new file mode 100644 index 000000000..beccbc7b3 --- /dev/null +++ b/cookbook/exp/legacy/train_reflexion_skill.py @@ -0,0 +1,1990 @@ +"""Online GRPO RFT for the reflexion skill generator (self-contained, cached). + +Trains an INDEPENDENT skill model to write reusable skills that, injected into a +FROZEN base solver's system prompt, raise its accuracy. The base is never trained; +it only produces the reward. Per chunk: base greedy solve -> rubric process-check +(view A) -> skill-gen (thinking ON, N candidates) -> leak filter -> with-skill +greedy pass -> group-relative advantage -> ONE on-policy GRPO step -> sync weights. +Reward is deterministic (T=0, M=1) binary correctness; A = (R - mean) / (std + eps) +within each problem-group, so std=0 groups give no gradient (GRPO variance selects). + +Each problem is routed to EXACTLY ONE view: A = problem + rubric findings, B = +query only (deployment form). Skill-gen trains only the final structured guidance turn. + +Base greedy rollouts and rubric diagnoses are cached to disk (jsonl, md5-keyed) so +restarts skip them; skill-gen is on-policy and never cached. + +8 GPUs default for high-memory cards: rank 0 trains the actor, rank 1 hosts a +frozen reference model, ranks 2-3 skill_sampler (synced), ranks 4-7 base_sampler +(frozen). Override TRAIN_GPUS / REF_GPUS / SKILL_SAMPLER_GPUS / BASE_SAMPLER_GPUS +for other layouts. +Leak / rubric use the backup teacher API: set LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / +LLM_BACKUP_MODEL. + +Launch: + LLM_BACKUP_API_KEY=... python cookbook/exp/embedding/train_reflexion_skill.py \ + --dataset aops --n 5000 --chunk-size 16 --lr 6e-6 +""" +import argparse +import copy +import hashlib +import json +import math +import os +import re +import sys +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Optional, Set, Tuple + +import numpy as np + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.data_format import SamplingParams +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.model import TransformersModel +from twinkle.patch.no_split_modules import NoSplitModulesPatch +from twinkle.processor import InputProcessor +from twinkle.sampler import vLLMSampler +from twinkle.template import Template +from twinkle_agentic.verifier import RubricVerifier +from twinkle_agentic.verifier.rubric_verifier import RubricItem + +logger = get_logger() + +try: + import swanlab +except ImportError: + swanlab = None + +MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') +GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.8)) +GEN_TEMPERATURE = float(os.environ.get('GEN_TEMPERATURE', 0.6)) +GEN_TOP_P = float(os.environ.get('GEN_TOP_P', 0.95)) +AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') +MATH_DATASET_ID = os.environ.get('MATH_DATASET_ID', 'modelscope/competition_math') + +# GPU layout: train actor (synced to skill_sampler) + frozen ref model + skill_sampler + base_sampler. +# High-memory cards can keep Qwen3-4B actor/ref as one GPU each and spend more GPUs +# on vLLM data-parallel sampling. The base side is heavier here because every +# clean skill candidate is re-solved by the frozen base model, plus balance/eval baselines. +TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 2)) +REF_GPUS = int(os.environ.get('REF_GPUS', 2)) +SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) +BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 2)) +NUM_GPUS = TRAIN_GPUS + REF_GPUS + SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS +TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', min(1, TRAIN_GPUS))) +REF_FSDP = int(os.environ.get('REF_FSDP', min(1, REF_GPUS))) +if min(TRAIN_GPUS, REF_GPUS, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS, TRAIN_FSDP, REF_FSDP) < 1: + raise ValueError('TRAIN_GPUS, REF_GPUS, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS, TRAIN_FSDP and REF_FSDP must all be >= 1') +if TRAIN_GPUS % TRAIN_FSDP != 0: + raise ValueError(f'TRAIN_GPUS ({TRAIN_GPUS}) must be divisible by TRAIN_FSDP ({TRAIN_FSDP})') +if REF_GPUS % REF_FSDP != 0: + raise ValueError(f'REF_GPUS ({REF_GPUS}) must be divisible by REF_FSDP ({REF_FSDP})') +TRAIN_DP = TRAIN_GPUS // TRAIN_FSDP +REF_DP = REF_GPUS // REF_FSDP + + +# =========================================================================== +# Block A -- boxed extraction + answer grading +# =========================================================================== +_BOXED_RE = re.compile(r'\\boxed\s*\{') + + +def extract_boxed(text: str) -> Optional[str]: + """Last ``\\boxed{...}`` content, brace-balanced.""" + if not text: + return None + last = None + for m in _BOXED_RE.finditer(text): + depth, i = 1, m.end() + while i < len(text) and depth > 0: + depth += (text[i] == '{') - (text[i] == '}') + i += 1 + if depth == 0: + last = text[m.end():i - 1].strip() + return last + + +def normalize_answer(ans: str) -> str: + if not ans: + return '' + s = str(ans).strip() + m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) + if m: + return m.group(1) + s = s.strip('$').strip().replace('−', '-') + 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').replace(r'\tfrac', r'\frac') + s = s.replace(r'\,', '').replace(r'\;', '').replace(r'\!', '') + s = re.sub(r'\\(?:quad|qquad|\s)', '', s) + s = re.sub(r'\s+', '', s) + s = re.sub(r'\\sqrt([0-9A-Za-z])', r'\\sqrt{\1}', s) + s = re.sub(r'\\frac\{([^{}]+)\}([^{}\\])', r'\\frac{\1}{\2}', s) + s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) + s = re.sub(r'\{[a-zA-Z]+\}$', '', s) + s = re.sub(r'\^\\circ|\^\{\\circ\}|°|\\circ', 'deg', s) + s = s.replace(r'\minus{}', '-').replace(r'\minus', '-') + + def _frac_to_slash(mt): + text = mt.group(0) + pos = text.index('{') + 1 + depth, num_start = 1, pos + while depth > 0: + depth += (text[pos] == '{') - (text[pos] == '}') + pos += 1 + numer = text[num_start:pos - 1] + pos += 1 + den_start, depth = pos, 1 + while depth > 0: + depth += (text[pos] == '{') - (text[pos] == '}') + pos += 1 + return f'({numer})/({text[den_start:pos - 1]})' + + s = re.sub(r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', _frac_to_slash, s) + s = re.sub(r'(? bool: + try: + va, vb = float(a.replace('(', '').replace(')', '')), 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) + return va is not None and vb is not None and abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) + + +_MCQ_REF_RE = re.compile( + r'^\\(?:textbf|text|mathrm|mathbf)\{\(?([A-E])\)?\s*\}\s*(.+)$' + r'|^\(?([A-E])\)\s+(.+)$') +_VAR_PREFIX_RE = re.compile(r'^(?:[A-Za-z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)$', re.DOTALL) + + +def _split_mcq(ans: str): + s = ans.strip() + m = _MCQ_REF_RE.match(s) + if m: + return (m.group(1) or m.group(3)), ((m.group(2) or m.group(4) or '').strip() or None) + bl = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) + return (bl.group(1), None) if bl else (None, s or None) + + +def _strip_var_prefix(ans: str) -> str: + m = _VAR_PREFIX_RE.match((ans or '').strip()) + return m.group(1).strip() if m else (ans or '') + + +def _normalize_tuple(ans: str) -> str: + return re.sub(r'[\s()\[\]{}\\]', '', ans or '') + + +def _math_verify_equal(predicted: str, reference: str) -> bool: + try: + from math_verify import parse, verify + gold = parse(r'\boxed{' + reference + '}', parsing_timeout=1) + pred = parse(r'\boxed{' + predicted + '}', parsing_timeout=1) + return bool(gold and pred and verify(gold, pred, timeout_seconds=1)) + except Exception: + return False + + +def answers_match(predicted: str, reference: str) -> bool: + if not predicted or not reference: + return False + norm_p, norm_r = normalize_answer(predicted), normalize_answer(reference) + if norm_p == norm_r or norm_p.lower() == norm_r.lower() or _try_numeric_equal(norm_p, norm_r): + return True + stripped_p = normalize_answer(_strip_var_prefix(predicted)) + stripped_r = normalize_answer(_strip_var_prefix(reference)) + if stripped_p and stripped_r: + if (stripped_p == stripped_r or stripped_p.lower() == stripped_r.lower() + or _try_numeric_equal(stripped_p, stripped_r)): + return True + p_letter, p_value = _split_mcq(predicted) + r_letter, r_value = _split_mcq(reference) + if p_letter and r_letter and p_letter == r_letter: + return True + if (p_letter and p_value is None) and (r_letter and r_value is None): + return False + 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 and (p_val == r_val or p_val.lower() == r_val.lower() or _try_numeric_equal(p_val, r_val)): + return True + for left, right in ((norm_p, norm_r), (stripped_p, stripped_r)): + tl, tr = _normalize_tuple(left), _normalize_tuple(right) + if ',' in tl and tl == tr: + return True + if '=' in norm_r and '=' not in norm_p: + if any(part and (part == norm_p or _try_numeric_equal(part, norm_p)) for part in norm_r.split('=')): + return True + if '=' in norm_p and '=' not in norm_r: + if any(part and (part == norm_r or _try_numeric_equal(part, norm_r)) for part in norm_p.split('=')): + return True + return _math_verify_equal(stripped_p or norm_p, stripped_r or norm_r) + + +# =========================================================================== +# Block B -- prompts, skill parsing, batched sampling +# =========================================================================== +DIRECT_SYSTEM = ( + 'You are an expert competition mathematician. Solve the following problem ' + 'step by step. Provide your final answer inside \\boxed{}.') + +_SKILL_SOLVE_PREFIX = ( + DIRECT_SYSTEM + '\n\n' + 'Before you start, keep these reminders in mind to avoid common mistakes on this ' + 'type of problem:\n') +_SKILL_SOLVE_SUFFIX = '\nApply them where relevant, but rely on your own reasoning to reach the answer.' + + +def build_direct_prompt(problem: str) -> Dict[str, Any]: + return {'messages': [{'role': 'system', 'content': DIRECT_SYSTEM}, + {'role': 'user', 'content': problem}]} + + +def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: + skill = (skill or '').strip() + if not skill: + return build_direct_prompt(problem) + # Concatenation (not .format): DIRECT_SYSTEM/skill contain literal braces. + return {'messages': [ + {'role': 'system', 'content': _SKILL_SOLVE_PREFIX + skill + _SKILL_SOLVE_SUFFIX}, + {'role': 'user', 'content': problem}]} + + +# -- skill-gen prompts (view A: problem + rubric source; view B: query only) -- +# Kept deliberately short: this is the RL policy's system prompt, so over-specifying +# the output hurts convergence. The concrete output format is appended separately by +# ``_SKILL_OUTPUT`` so the task text stays format-agnostic. +SKILL_GEN_SYSTEM = ( + 'You are a math guidance writer. A process-check on a related problem hints at ' + 'likely mistakes. Write short reusable guidance for this and similar problems, ' + 'and note what to watch out for.\n') + +SKILL_GEN_SYSTEM_Q = ( + 'You are a math guidance writer. Write short reusable guidance for this and ' + 'similar problems.\n') + +_SKILL_OUTPUT = ( + 'Output only:\n\nYour reusable solving guidance here.\n') + +SKILL_GEN_USER_Q = ( + 'Problem:\n{problem}\n\n') + +SKILL_GEN_USER_RUBRIC = ( + 'Target problem:\n{problem}\n\n' + 'Problem used for the process check:\n{rubric_problem}\n\n' + 'Process check:\n' + '{diagnosis}\n\n') + + +def _rubric_has_fail(diagnosis: str) -> bool: + """View A uses the open-book rubric prompt (and, under --viewa-sft, SFT distillation) + IFF its rubric localised a failure. No ``[FAIL]`` => nothing to inject: generation + degrades to query-only and the problem is trained by GRPO exactly like view B. Single + source of truth so ``_skillgen_messages`` and ``_group_records`` never diverge.""" + return '[FAIL]' in (diagnosis or '') + + +def _skillgen_messages(problem: str, view: str, diagnosis: str, + rubric_problem: str = '') -> List[Dict[str, Any]]: + """Single source of truth for the skill-gen prompt (used at BOTH generation and + training so they never diverge). View A with a localisable failure uses the target + problem plus the rubric source problem and findings; view B -- or a view-A problem + whose rubric flagged nothing (no ``[FAIL]``) -- degrades to the query-only prompt.""" + if view == 'B' or not _rubric_has_fail(diagnosis): + return [{'role': 'system', 'content': SKILL_GEN_SYSTEM_Q + _SKILL_OUTPUT}, + {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=problem)}] + rubric_problem = rubric_problem or problem + return [{'role': 'system', 'content': SKILL_GEN_SYSTEM + _SKILL_OUTPUT}, + {'role': 'user', 'content': SKILL_GEN_USER_RUBRIC.format( + problem=problem, rubric_problem=rubric_problem, diagnosis=diagnosis)}] + + +def _assign_view(problem: str, args: argparse.Namespace) -> str: + h = int(hashlib.md5(f'{args.seed}:{problem}'.encode('utf-8')).hexdigest(), 16) + return 'B' if (h % 100000) / 100000.0 < args.view_b_frac else 'A' + + +def _curriculum_view_b_frac(gstep: int, args: argparse.Namespace) -> float: + """View-A anneal (--viewa-frac-start): the view-A share holds at ``viewa_frac_start`` + for the first ``viewa_warmup_chunks`` chunks (pure-SFT warmup when start==1.0), then + decays linearly to ``viewa_frac_end`` over ``viewa_decay_chunks`` chunks and holds. + Because _assign_view is a fixed hash against a moving threshold, the B set grows + MONOTONICALLY: a problem trained open-book (A) early can only reappear closed-book + (B) later, never the reverse.""" + t = min(max(gstep - args.viewa_warmup_chunks, 0) / max(args.viewa_decay_chunks, 1), 1.0) + share = args.viewa_frac_start + (args.viewa_frac_end - args.viewa_frac_start) * t + return 1.0 - share + + +def _view_prompt(r: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: + return {'messages': _skillgen_messages( + r['problem'], r['_view'], r.get('_rubric_diag', ''), r.get('_rubric_src', ''))} + + +_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') + + +def _clean_text(decoded: Optional[str]) -> str: + return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() + + +def _extract_tag_block(answer: str, tag: str, allow_empty: bool = False) -> Optional[str]: + low = answer.lower() + open_tag, close_tag = f'<{tag}>', f'' + s = low.rfind(open_tag) + if s < 0: + return None + inner = s + len(open_tag) + e = low.find(close_tag, inner) + if e < 0: + return None + block = answer[inner:e].strip() + block = re.sub(r'', '', block, flags=re.IGNORECASE).strip() + return block if (block or allow_empty) else None + + +def _extract_skill(text: str) -> Optional[str]: + """Parse skill-generation output: return the inner text of a non-empty ```` + block, or None. If a ```` marker is present, parse only the text after the + last one; otherwise parse the full response.""" + low = text.lower() + end_think = low.rfind('') + answer = text[end_think + len(''):] if end_think >= 0 else text + return _extract_tag_block(answer, 'skills') + + +def _parse_seq(seq, gold: str) -> Dict[str, Any]: + """Grade one sampled sequence into a rollout record.""" + text = _clean_text(getattr(seq, 'decoded', '') or '') + pred = extract_boxed(text) + correct = bool(pred) and answers_match(pred, gold) + terminated = getattr(seq, 'stop_reason', None) != 'length' + return {'pred': pred, 'correct': correct, 'terminated': terminated, + 'passed': bool(correct and terminated), + 'stop_reason': getattr(seq, 'stop_reason', None), + 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} + + +def _run_samples(sampler, prompts: List[Any], num_samples: int, max_tokens: int, + gen_dp: int, temperature: Optional[float] = None, + top_p: Optional[float] = None, top_k: Optional[int] = None) -> List[List[Any]]: + """One batched sampler call -> per-prompt list of raw sequences. vLLM dp needs + batch len >= dp, so pad the tail and slice back.""" + if not prompts: + return [] + params = SamplingParams( + max_tokens=max_tokens, + temperature=GEN_TEMPERATURE if temperature is None else temperature, + top_p=GEN_TOP_P if top_p is None else top_p, + num_samples=num_samples, **({} if top_k is None else {'top_k': top_k})) + padded = prompts + if gen_dp > 1 and 0 < len(prompts) < gen_dp: + padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] + responses = sampler.sample(padded, params)[:len(prompts)] + return [list(r.sequences) if (r and r.sequences) else [] for r in responses] + + +# =========================================================================== +# Block C -- data loading via twinkle.Dataset + numeric filtering +# =========================================================================== +def _boxed_batch(rows: Dict[str, List[Any]], dataset: str) -> Dict[str, List[Any]]: + """Batched (columnar) mapper for ``Dataset.map``: annotate each row with its boxed + ``reference_answer`` and a ``_keep`` flag (aops also needs the ``boxed`` metadata).""" + sols = rows['solution'] + metas = rows.get('metadata', [None] * len(sols)) + refs = [extract_boxed(s or '') for s in sols] + keep = [bool(ref) and (dataset != 'aops' or bool((meta or {}).get('boxed'))) + for ref, meta in zip(refs, metas)] + return {**rows, 'reference_answer': refs, '_keep': keep} + + +def load_problems(dataset: str, n: int, seed: int, num_proc: int = 0) -> List[Dict[str, Any]]: + """Load boxed-answer problems as ``{problem, reference_answer, level?}`` via + twinkle.Dataset (ModelScope hub), sampled to ``n`` (0 = all). Boxed extraction (regex + + brace scan over every solution) is parallelised by ``Dataset.map``/``.filter``; + ``num_proc`` defaults to all cores (set 1 to force serial).""" + ds_id = AOPS_DATASET_ID if dataset == 'aops' else MATH_DATASET_ID + ds = Dataset(DatasetMeta(dataset_id=f'ms://{ds_id}', split='train')) + nproc = num_proc if num_proc > 0 else min(32, os.cpu_count() or 1) + ds.map(lambda rows: _boxed_batch(rows, dataset), num_proc=nproc) + ds.filter(lambda row: row['_keep'], num_proc=nproc) + has_level = 'level' in ds.dataset.column_names + out = [{'data_id': f'{dataset}:{i}', 'problem': row['problem'], + 'reference_answer': row['reference_answer'], + **({'level': row['level']} if has_level and row.get('level') else {})} + for i, row in enumerate(ds.dataset)] + logger.info(f'[data] {dataset}: {len(out)} boxed problems ({nproc} procs)') + rng = np.random.RandomState(seed) + rng.shuffle(out) + return out[:n] if (n and n < len(out)) else out + + +# --------------------------------------------------------------------------- +# Bag-of-words neighbour pairing (cross-problem rubric transfer, --xproblem-rubric) +# --------------------------------------------------------------------------- +# Common English + math-scaffolding words that carry no problem-type signal. Kept small +# and deterministic on purpose (no external stopword list): what survives is the domain +# vocabulary ('triangle', 'prime', 'polynomial', ...) that actually defines the type. +_BOW_STOP = frozenset(""" +a an the of to in on at for and or but if is are be was were been being this that these those +with without into onto from by as it its their his her our your my we you they he she them +find compute determine calculate evaluate solve show prove given let suppose consider assume +what which when where how many much value values number numbers expression form terms term +such that then than so if only when each every all any some both one two three four five six +seven eight nine ten first second third last non over under about above below between +problem answer result equal equals sum difference product total following there here have has +had do does did can could will would should may might must not no yes if then else +""".split()) + +_WORD_RE = re.compile(r'[a-z]+') + + +def _stem(w: str) -> str: + """Crude suffix stripper so 'prime'/'primes', 'triangle'/'triangles' collapse to one + type token. Not linguistically correct -- just enough to merge the common plural/gerund + variants that otherwise split a type's vocabulary and starve the df filter.""" + if len(w) > 4 and w.endswith('ies'): + return w[:-3] + 'y' # properties -> property + if len(w) > 4 and w.endswith('es') and w[-3] in 'sxzh': + return w[:-2] # boxes -> box (keep primes -> prime below) + for suf in ('ing', 'ed', 's'): + if len(w) > len(suf) + 2 and w.endswith(suf): + return w[:-len(suf)] + return w + + +def _tokenize(problem: str) -> List[str]: + """Deterministic bag-of-words tokens for type matching: lowercase alphabetic words + (numbers dropped -- they are instance detail, not type), minus generic stopwords, then + stemmed, so only domain terms remain. Latex control words (frac, sqrt, ...) survive.""" + return [_stem(w) for w in _WORD_RE.findall((problem or '').lower()) + if len(w) > 2 and w not in _BOW_STOP] + + +class BagOfWordsIndex: + """TF-IDF cosine nearest-neighbour over problem statements. Sparse dict vectors + + an inverted index, so ``nearest`` scores only problems sharing a term (near-linear in + practice, no dense NxN). Deterministic; ties break on lower index for reproducibility. + + Answer-aware: ``nearest`` skips the same index, near-verbatim duplicates (cosine + >= ``sim_max``) and -- crucially for anti-leak -- any candidate whose answer equals the + query's, so a neighbour rubric can never hand over the query's own answer.""" + + def __init__(self, problems: List[str], answers: Optional[List[str]] = None, + min_df: int = 2, max_df_frac: float = 0.5): + self._toks = [_tokenize(p) for p in problems] + self._ans = [(_numeric_value(a) or (str(a).strip() if a else '')) for a in answers] \ + if answers is not None else [''] * len(problems) + n = len(self._toks) + df: Dict[str, int] = {} + for toks in self._toks: + for w in set(toks): + df[w] = df.get(w, 0) + 1 + max_df = max(min_df, int(max_df_frac * n)) + self._idf = {w: math.log((n + 1) / (c + 1)) + 1.0 + for w, c in df.items() if min_df <= c <= max_df} + self._vecs: List[Dict[str, float]] = [self._vectorize(t) for t in self._toks] + self._inverted: Dict[str, List[int]] = {} + for i, v in enumerate(self._vecs): + for w in v: + self._inverted.setdefault(w, []).append(i) + + def _vectorize(self, toks: List[str]) -> Dict[str, float]: + tf: Dict[str, float] = {} + for w in toks: + if w in self._idf: + tf[w] = tf.get(w, 0.0) + 1.0 + vec = {w: c * self._idf[w] for w, c in tf.items()} + norm = math.sqrt(sum(x * x for x in vec.values())) + return {w: x / norm for w, x in vec.items()} if norm > 0 else {} + + def nearest(self, i: int, sim_max: float = 0.98) -> Tuple[int, float]: + """(index, cosine) of the most similar *distinct* problem: not i, not a duplicate + (cosine < ``sim_max``) and not the same answer. (-1, 0.0) if none qualifies.""" + vi = self._vecs[i] + if not vi: + return -1, 0.0 + ai = self._ans[i] + scores: Dict[int, float] = {} + for w, xi in vi.items(): + for j in self._inverted.get(w, ()): + if j != i: + scores[j] = scores.get(j, 0.0) + xi * self._vecs[j].get(w, 0.0) + best_j, best_s = -1, 0.0 + for j, s in scores.items(): + if s >= sim_max or (ai and self._ans[j] == ai): + continue + if s > best_s or (s == best_s and (best_j < 0 or j < best_j)): + best_j, best_s = j, s + return best_j, best_s + + +def build_pairs(records: List[Dict[str, Any]], n: int, seed: int, sim_max: float = 0.98 + ) -> Tuple[List[Dict[str, Any]], Dict[str, Tuple[str, float]]]: + """Single-pass cross-problem pairing over the whole pool (one index build). + + Returns ``(subset, neighbour_map)`` where ``subset`` is the ``n`` problems with the + strongest qualifying neighbour (dense same-type pairs; n<=0/n>=len keeps all, shuffled) + and ``neighbour_map[q] = (p, cosine)`` gives each kept problem its analogue P. P is drawn + from the *full* pool (richer analogues) but never a duplicate or same-answer problem, so + P's rubric can transfer method without ever leaking Q's answer.""" + index = BagOfWordsIndex([r['problem'] for r in records], + [str(r.get('reference_answer', '')) for r in records]) + nbr = [index.nearest(i, sim_max) for i in range(len(records))] + order = sorted(range(len(records)), key=lambda i: (-nbr[i][1], i)) + keep = order[:n] if (0 < n < len(records)) else list(range(len(records))) + rng = np.random.RandomState(seed) + rng.shuffle(keep) + subset = [records[i] for i in keep] + neighbour_map = {records[i]['problem']: (records[nbr[i][0]]['problem'], nbr[i][1]) + for i in keep if nbr[i][0] >= 0} + return subset, neighbour_map + + +_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') + + +def _norm_num_text(num: str) -> str: + try: + f = float(num) + return str(int(f)) if f == int(f) else str(f) + except Exception: + return str(num).strip() + + +def _numeric_value(raw: Any) -> Optional[str]: + """Collapse an answer to a single int/decimal/fraction, or None.""" + if raw is None: + return None + s = str(raw).strip().strip('$').strip() + s = s.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac') + s = re.sub(r'\\!|\\,|\\;|\\ |\\left|\\right|\s', '', s) + for pat in (r'\\frac\{(-?\d+)\}\{(-?\d+)\}', r'(-?\d+)/(-?\d+)'): + m = re.fullmatch(pat, s) + if m: + a, b = int(m.group(1)), int(m.group(2)) + return _norm_num_text(str(a / b)) if b else None + return _norm_num_text(s) if _NUM_RE.fullmatch(s) else None + + +def _answer_leaked(skill: str, reference: str) -> bool: + """Audit whether a generated skill contains the final answer verbatim. This is NOT + a training filter: if the skill model derives an answer from the problem, that is a + legitimate answer-bearing skill under this experiment. The real leakage boundary is the + external API/rubric diagnosis, which is constrained by prompt to stay answer-free.""" + if not skill: + return False + for cand in {_numeric_value(reference), (str(reference).strip() or None)}: + if cand and re.search(r'(? Tuple[Set[str], Set[str]]: + """Read jsonl files and collect data_id/problem keys that must be excluded. + + The cold-start SFT builder writes stable ``data_id`` values. ``problem`` is kept as a + backward-compatible fallback for older jsonl files produced before data_id existed.""" + ids: Set[str] = set() + problems: Set[str] = set() + for raw_path in (paths_arg or '').split(','): + path = raw_path.strip() + if not path or not os.path.exists(path): + continue + with open(path, encoding='utf-8') as f: + for line in f: + if not line.strip(): + continue + row = json.loads(line) + if row.get('record_type') in {'config', 'summary'}: + continue + data_id = str(row.get('data_id') or '').strip() + problem = str(row.get('problem') or '').strip() + if data_id: + ids.add(data_id) + elif problem: + problems.add(problem) + return ids, problems + + +def _load_records(args: argparse.Namespace + ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], + Dict[str, Tuple[str, float]], Dict[str, str], Dict[str, int]]: + """Load, numeric-filter, shuffle, split a fixed eval holdout, and (when --xproblem-rubric) + select a same-type-dense train subset with its neighbour map -- all in one pass. + Also returns ``pool_answers`` (every candidate neighbour P's true answer) so P's baseline + can be graded/cached correctly even when P is not itself a training problem.""" + # Load all when filtering or splitting (else the eval holdout could starve train). + load_n = 0 if (args.numeric_only or args.eval_size > 0) else args.n + records = load_problems(args.dataset, load_n, args.seed) + raw_n, dropped = len(records), 0 + if args.numeric_only: + records = [{**r, 'reference_answer': v} + for r, v in ((r, _numeric_value(r.get('reference_answer'))) for r in records) + if v is not None] + dropped = raw_n - len(records) + np.random.RandomState(args.seed).shuffle(records) + exclude_ids, exclude_problems = _load_excluded_records(getattr(args, 'exclude_data_ids', '')) + excluded = 0 + if exclude_ids or exclude_problems: + before = len(records) + records = [r for r in records + if str(r.get('data_id', '')) not in exclude_ids + and str(r.get('problem', '')).strip() not in exclude_problems] + excluded = before - len(records) + eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 + eval_records = [dict(r) for r in records[:eval_n]] + pool = records[eval_n:] + pool_offset = max(0, int(getattr(args, 'pool_offset', 0) or 0)) + if pool_offset: + if pool_offset >= len(pool): + raise ValueError(f'--pool-offset ({pool_offset}) leaves no train records from pool size {len(pool)}') + pool = pool[pool_offset:] + train_n = args.n if args.n > 0 else len(pool) + # Cross-problem rubric: pick the train_n problems with the strongest qualifying neighbour + # (dense same-type pairs) and build {Q -> (P, sim)} in one index pass. Otherwise keep the + # first train_n (already shuffled) with no neighbours. + if args.xproblem_rubric: + subset, neighbor_map = build_pairs(pool, train_n, args.seed) + train_records = [dict(r) for r in subset] + pool_answers = {r['problem']: str(r.get('reference_answer', '')) for r in pool} + else: + train_records, neighbor_map, pool_answers = [dict(r) for r in pool[:train_n]], {}, {} + if {r['problem'] for r in train_records} & {r['problem'] for r in eval_records}: + raise ValueError('eval/train overlap detected') + stats = {'raw_loaded': raw_n, 'numeric_dropped': dropped, + 'excluded_records': excluded, 'pool_offset': pool_offset, + 'train_records': len(train_records), 'eval_records': len(eval_records)} + return train_records, eval_records, neighbor_map, pool_answers, stats + + +# =========================================================================== +# Block D -- disk cache, problem pool, baseline rollout, rubric check +# =========================================================================== +class DiskCache: + """Append-only jsonl kv cache (md5 key -> value). Loads on init, appends on put. + Disabled instances always miss and never write.""" + + def __init__(self, path: str, enabled: bool = True): + self._mem: Dict[str, Any] = {} + self._fh = None + self._lock = threading.Lock() # base baseline is prefetched on a background thread + if not enabled: + return + if os.path.exists(path): + with open(path, encoding='utf-8') as f: + for line in f: + if line.strip(): + row = json.loads(line) + self._mem[row['key']] = row['value'] + self._fh = open(path, 'a', encoding='utf-8') + + @staticmethod + def key_for(*parts: str) -> str: + return hashlib.md5('\x1f'.join(parts).encode('utf-8')).hexdigest() + + def __contains__(self, key: str) -> bool: + with self._lock: + return key in self._mem + + def get(self, key: str) -> Any: + with self._lock: + return self._mem.get(key) + + def put(self, key: str, value: Any) -> None: + with self._lock: + self._mem[key] = value + if self._fh is not None: + self._fh.write(json.dumps({'key': key, 'value': value}, ensure_ascii=False) + '\n') + self._fh.flush() + + def close(self) -> None: + if self._fh is not None: + self._fh.close() + + +class _LockedSampler: + """Serialises every ``.sample()`` on a sampler behind one lock. The base sampler is + shared by the main thread (with-skill scoring, eval) and the baseline-prefetch thread; + ``sample`` is a slice_dp/flatten remote call whose collect is NOT safe to interleave + across two callers, so concurrent calls could mis-join sequences. The lock keeps base + calls serial (prefetch still overlaps the skill-gen phase, which uses skill_sampler).""" + + def __init__(self, sampler): + self._sampler = sampler + self._lock = threading.Lock() + + def sample(self, *args, **kwargs): + with self._lock: + return self._sampler.sample(*args, **kwargs) + + def __getattr__(self, name): + return getattr(self._sampler, name) + + +class ProblemPool: + """Cyclic draw source. Each full pass reshuffles with ``seed + epoch``; the initial + pass keeps the loader order. ``draw(k)`` returns k distinct problems (k << pool).""" + + def __init__(self, records: List[Dict[str, Any]], seed: int): + self._records = list(records) + self._seed, self._cursor, self.epoch = seed, 0, 0 + + def draw(self, k: int) -> List[Dict[str, Any]]: + out, seen = [], set() + while len(out) < k: + if self._cursor >= len(self._records): + self.epoch += 1 + np.random.RandomState(self._seed + self.epoch).shuffle(self._records) + self._cursor = 0 + r = self._records[self._cursor] + self._cursor += 1 + if id(r) not in seen: + seen.add(id(r)) + out.append(r) + return out + + def peek(self, k: int) -> List[Dict[str, Any]]: + """The next k distinct problems draw() would return, WITHOUT advancing state + (no cursor move, no reshuffle). Used to prefetch their baseline into base_cache + while the current chunk trains -- a wrong guess (cross-epoch reshuffle) only + misses the cache, never corrupts the draw.""" + out, seen, cur = [], set(), self._cursor + recs = self._records + while len(out) < k and cur < len(recs): # stop at epoch edge; don't simulate reshuffle + r = recs[cur] + cur += 1 + if id(r) not in seen: + seen.add(id(r)) + out.append(r) + return out + + +def _empty_roll() -> Dict[str, Any]: + return {'pred': '', 'correct': False, 'terminated': False, 'passed': False, + 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} + + +def _apply_baseline(r: Dict[str, Any], roll: Dict[str, Any]) -> None: + """Attach a greedy baseline roll and reset per-chunk working state.""" + r['_baseline_rolls'], r['_cands'], r['_init'] = [roll], [], [roll] + r['_failed'] = not roll['correct'] + r['_baseline_pass'] = 1.0 if roll['correct'] else 0.0 + r['_hard'] = True # process every problem; group variance selects (SEAM-style) + + +def baseline_rollout(base_sampler, problems: List[Dict[str, Any]], base_dp: int, + args: argparse.Namespace, cache: DiskCache) -> int: + """Base solves each problem greedily once (T=0, M=1), disk-cached by problem text. + The base is frozen + greedy so a cache hit is exact. Returns fresh (miss) count.""" + todo = [r for r in problems if DiskCache.key_for(r['problem']) not in cache] + if todo: + out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], + 1, args.max_tokens, base_dp, temperature=0.0) + for r, seqs in zip(todo, out): + roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + cache.put(DiskCache.key_for(r['problem']), roll) + for r in problems: + _apply_baseline(r, cache.get(DiskCache.key_for(r['problem']))) + return len(todo) + + +# -- rubric process-check (view A): teacher diagnoses the base's attempt -- +_RFT_DIAG_SYSTEM = """\ +You are a strategy-level process checker for a math solution attempt. You are given a +math problem, a rubric, and one attempted solution segment. Decide PASS or FAIL for each +criterion, and write the diagnosis so it can become useful reusable guidance for solving +similar problems without seeing this segment. + +Output STRICT JSON (no prose outside it) with this shape: +{ + "items": [ + {"index": 1, "verdict": "PASS", "reason": "", + "fix": ""}, + {"index": 2, "verdict": "FAIL", "reason": "", + "fix": ""} + ], + "overall": "OK" | "ISSUES", + "summary": "" +} + +Rules: +- Judge every criterion independently; a [Hard Rule] is FAIL unless clearly satisfied. +- Judge ONLY what is observable in THIS segment. Ignore hidden or + content for output-format criteria. +- The API diagnosis is an external teacher signal, so it must stay answer-free. +- Prefer diagnosis that transfers to view-B skill generation: name the route choice, + structural observation, missing check, or length-control habit that a solver should + remember before solving a similar problem. +- For PASS items, leave "fix" as "". +- For FAIL items, describe the process problem at strategy level: unsuitable method, + missed structure, invalid transformation, missing constraint check, redundant cases, + off-track approach, contradiction, or inefficient/unfinished reasoning. +- A fix may suggest the LOCAL correction direction, such as identify the key structure, + verify constraints, preserve equivalence, reduce redundant cases, or choose a more + direct route. Do not carry out the correction. +- Never reveal the final answer, a corrected value/expression, an option label, or a + step-by-step solution that would let another model copy the solve. +- If the segment contains a process note saying it was cut off before a final boxed + answer, mark the length-budget criterion as FAIL and suggest a method-level way to + finish faster. +- Keep every "reason" and "fix" concise: one short sentence each. +- "overall" is "OK" only if NO criterion is FAIL. +- Output only the JSON object.""" + +_RFT_DIAG_USER = """\ +## Task / query (context) +{query} + +## Rubric +{rubric} + +## Segment +{segment} + +Now output the diagnostic JSON object.""" + +_MATH_RUBRIC = [ + ('The attempt chooses a method suitable for the problem structure', False), + ('The attempt identifies the key constraint, invariant, or quantity before computing', False), + ('Algebraic and logical transformations preserve validity at each step', True), + ('The attempt checks required constraints, domains, boundary cases, or validity conditions', False), + ('The attempt avoids redundant casework, looping, or re-deriving known facts', False), + ('The attempt reaches a final boxed answer within the length budget', False), + ('The approach stays focused on the actual question asked', False), +] + +# Bump when the rubric criteria or the diagnosis prompt change so the disk-cached +# diagnoses written under an older rubric are not silently reused. +_RUBRIC_VERSION = 'rubric_v5_viewb_strategy' + + +class _RftRubricVerifier(RubricVerifier): + def _diagnose_trajectory(self, query: str, rubric_block: str, segment_text: str) -> dict: + return {'messages': [ + {'role': 'system', 'content': _RFT_DIAG_SYSTEM}, + {'role': 'user', 'content': _RFT_DIAG_USER.format( + query=query, rubric=rubric_block, segment=segment_text)}]} + + +def build_rubric_checker() -> Optional[RubricVerifier]: + """Fixed math-process rubric verifier (teacher-served). None if no LLM backup env.""" + if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') + or os.environ.get('OPENAI_API_KEY')): + return None + return _RftRubricVerifier( + fixed_rubric=[RubricItem(t, is_hard=h) for t, h in _MATH_RUBRIC], gate=True) + + +def _format_diagnosis(detail) -> str: + """One line per criterion (PASS/FAIL + reason + fix) then a summary.""" + rub = detail.rubric + lines = [] + for it in detail.items: + text = rub[it.index - 1].text if 0 < it.index <= len(rub) else f'criterion {it.index}' + if it.verdict: + lines.append(f'- [PASS] {text}') + else: + tail = f': {it.reason}' if it.reason else '' + tail += f' (fix: {it.fix})' if it.fix else '' + lines.append(f'- [FAIL] {text}{tail}') + if detail.summary: + lines.append(f'Summary: {detail.summary}') + return '\n'.join(lines) + + +def diagnose_views(checker, problems: List[Dict[str, Any]], args: argparse.Namespace, + cache: DiskCache) -> None: + """Rubric-check every view-A problem's greedy attempt in parallel (disk-cached by + problem + attempt), stashing the formatted findings on ``r['_rubric_diag']``.""" + targets = [r for r in problems if r.get('_view') == 'A'] + if not checker or not targets: + return + + def _key(r: Dict[str, Any]) -> str: + init = r.get('_init', [{}])[0] + term = 'L' if (init.get('stop_reason') == 'length' or not init.get('terminated')) else 'T' + return DiskCache.key_for(r['problem'], init.get('text', ''), _RUBRIC_VERSION, term) + + pending = [] + for r in targets: + key = _key(r) + if key in cache: + r['_rubric_diag'] = cache.get(key) + else: + pending.append((r, key)) + if not pending: + return + + def _run(item): + r, key = item + init = r['_init'][0] + seg_text = init['text'] + if init.get('stop_reason') == 'length' or not init.get('terminated'): + seg_text += ('\n\n[Process note: this attempt was cut off at the token budget ' + 'and never produced a final \\boxed{} answer.]') + seg = {'messages': [{'role': 'user', 'content': r['problem']}, + {'role': 'assistant', 'content': seg_text}]} + attempts = max(1, args.rubric_retries + 1) + for attempt in range(attempts): + try: + return r, key, _format_diagnosis(checker.diagnose(seg, query=r['problem'])) + except Exception as exc: # teacher hiccup -> retry, then no-diagnosis prompt (not cached) + if attempt + 1 < attempts: + logger.warning(f'[rubric] diagnose error: {exc}; retry {attempt + 1}/{args.rubric_retries}') + time.sleep(min(2.0, 0.5 * (2 ** attempt))) + continue + logger.warning(f'[rubric] diagnose error: {exc}; giving up after {attempts} attempts') + return r, key, None + + workers = max(1, min(args.rubric_workers, len(pending))) + with ThreadPoolExecutor(max_workers=workers) as ex: + for r, key, diag in ex.map(_run, pending): + r['_rubric_diag'] = diag or '' + if diag is not None: + cache.put(key, diag) + + +# =========================================================================== +# Block E -- chunk draw, generation pipeline, record building +# =========================================================================== +def _baseline_class(r: Dict[str, Any]) -> str: + """success | fail_loop (out of length / never terminated) | fail_wrong.""" + roll = r['_init'][0] + if roll['correct']: + return 'success' + return 'fail_loop' if (roll['stop_reason'] == 'length' or not roll['terminated']) else 'fail_wrong' + + +def _select_balanced(buckets: Dict[str, List[Dict[str, Any]]], n_success: int, + n_fail: int, n_fail_loop: int) -> List[Dict[str, Any]]: + """Pick n_fail base-fails (toward n_fail_loop loop-fails, best-effort) + n_success + base-successes; top up any shortfall from leftovers.""" + loop, wrong, succ = buckets['fail_loop'], buckets['fail_wrong'], buckets['success'] + take_loop = min(n_fail_loop, len(loop)) + take_wrong = min(n_fail - take_loop, len(wrong)) + take_loop = min(n_fail - take_wrong, len(loop)) + sel = loop[:take_loop] + wrong[:take_wrong] + succ[:n_success] + target = n_success + n_fail + if len(sel) < target: + used = {id(x) for x in sel} + sel += [x for b in (loop, wrong, succ) for x in b if id(x) not in used][:target - len(sel)] + return sel + + +def draw_chunk(pool: ProblemPool, base_sampler, base_dp: int, args: argparse.Namespace, + cache: DiskCache) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + """Draw one chunk, baselining every drawn problem. With ``--balance``, keep + drawing+baselining until the target base fail:success mix is reachable (or the budget + is hit), then select a balanced subset.""" + if not args.balance: + chunk = pool.draw(args.chunk_size) + fresh = baseline_rollout(base_sampler, chunk, base_dp, args, cache) + return chunk, {'enabled': False, 'n_drawn': len(chunk), 'n_baseline_fresh': fresh} + + n_success = max(0, min(args.chunk_size, round(args.chunk_size * args.balance_success_frac))) + n_fail = args.chunk_size - n_success + n_fail_loop = round(n_fail * args.balance_loop_frac) + buckets: Dict[str, List[Dict[str, Any]]] = {'success': [], 'fail_loop': [], 'fail_wrong': []} + budget, n_drawn, n_fresh, seen = args.chunk_size * args.balance_max_draws_mult, 0, 0, set() + while n_drawn < budget: + if (len(buckets['success']) >= n_success + and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail): + break + batch = pool.draw(args.chunk_size) + n_fresh += baseline_rollout(base_sampler, batch, base_dp, args, cache) + n_drawn += len(batch) + for r in batch: + if id(r) not in seen: + seen.add(id(r)) + buckets[_baseline_class(r)].append(r) + + reached = (len(buckets['success']) >= n_success + and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail) + chunk = _select_balanced(buckets, n_success, n_fail, n_fail_loop) + sel_success = sum(1 for r in chunk if not r['_failed']) + stats = { + 'enabled': True, 'n_drawn': n_drawn, 'n_baseline_fresh': n_fresh, 'n_selected': len(chunk), + 'target_success': n_success, 'target_fail': n_fail, 'target_fail_loop': n_fail_loop, + 'selected_success': sel_success, 'selected_fail': len(chunk) - sel_success, + 'selected_fail_loop': sum(1 for r in chunk if _baseline_class(r) == 'fail_loop'), + 'selected_fail_wrong': sum(1 for r in chunk if _baseline_class(r) == 'fail_wrong'), + 'selected_success_frac': (sel_success / len(chunk)) if chunk else 0.0, + 'budget_hit': not reached, + } + return chunk, stats + + +def _assign_advantages(hard: List[Dict[str, Any]], args: argparse.Namespace) -> None: + """Group-relative advantage over each problem's scored candidates using the greedy + binary reward R in {0,1}: ``A = (R - mean) / (std + eps)``. std==0 groups get no + gradient -- GRPO's variance selects informative problems (no explicit difficulty gate). + A symmetric clip limits both positive and negative outliers from sparse 1-vs-many groups.""" + eps = 1e-6 + adv_clip = abs(float(getattr(args, 'adv_clip', 0.0) or 0.0)) + for r in hard: + for c in r['_cands']: + c['advantage'], c['grpo_adv'], c['kept'] = 0.0, 0.0, False + cs = [c for c in r['_cands'] if c.get('reward') is not None] + if len(cs) < 2: + continue + rewards = [c['reward'] for c in cs] + mean_r = sum(rewards) / len(rewards) + std = (sum((x - mean_r) ** 2 for x in rewards) / len(rewards)) ** 0.5 + if std < 1e-9: + continue + for c in cs: + raw_adv = (c['reward'] - mean_r) / (std + eps) + adv = max(-adv_clip, min(adv_clip, raw_adv)) if adv_clip > 0 else raw_adv + c['advantage'], c['grpo_adv'], c['kept'] = adv, adv, c['reward'] > mean_r + + +def _best_sft_candidate(r: Dict[str, Any], args: argparse.Namespace) -> Optional[Dict[str, Any]]: + """Pick ONE view-A candidate to distill (online context distillation). ONLY + executor-verified PASSING skills (reward==1) are distilled: the earlier fallback to + unverified skills meant ~60% of SFT targets had failed their own executor pass + (measured on the sft35 run) and the model was imitating plausible-but-wrong skills. + Problems with no passing candidate now yield NO SFT record. Answer-bearing skills + produced by the skill model itself are allowed here; only the external API/rubric + diagnosis must be answer-free. Among the passing candidates, take the one whose skill + length is CLOSEST to ``--sft-target-len`` -- an empirically high-pass-rate length + (~500-600 chars in this run) -- breaking ties by the fewest executor solve tokens. + Targeting a length (rather than the minimum) avoids a distillation feedback loop that + would otherwise drive rollouts ever shorter.""" + eligible = [c for c in r['_cands'] if c.get('parseable') and c.get('skills')] + passing = [c for c in eligible if c.get('reward') == 1.0] + if not passing: + return None + target = int(getattr(args, 'sft_target_len', 550) or 550) + + def _solve_tokens(c: Dict[str, Any]) -> int: + rolls = c.get('rolls') or [] + return rolls[0].get('gen_tokens', 1 << 30) if rolls else (1 << 30) + + return min(passing, key=lambda c: (abs(len(c['skills']) - target), _solve_tokens(c))) + + +def apply_neighbor_rubric(base_sampler, chunk: List[Dict[str, Any]], + neighbor_map: Dict[str, Tuple[str, float]], + pool_answers: Dict[str, str], base_dp: int, + args: argparse.Namespace, checker, + base_cache: DiskCache, rubric_cache: DiskCache) -> None: + """Cross-problem rubric (--xproblem-rubric): for every view-A problem Q, replace its own + rubric with the rubric of its bag-of-words neighbour P (Q keeps being the solved/scored + problem). P is baselined + diagnosed here (both disk-cached) as a stub carrying P's REAL + answer, so P's baseline grades correctly and legitimately shares the baseline cache with + P-as-training-problem. P's findings are copied onto Q with the neighbour text + similarity + for audit. A P that can't be diagnosed -> Q degrades to view B. Since P's answer differs + from Q's (guaranteed at pairing time), any answer P's rubric leaks is useless for Q.""" + targets = [r for r in chunk if r.get('_view') == 'A' and neighbor_map.get(r['problem'])] + if not targets: + return + stubs, by_problem = [], {} + for r in targets: + p, _ = neighbor_map[r['problem']] + if p not in by_problem: + stub = {'problem': p, 'reference_answer': pool_answers.get(p, ''), '_view': 'A'} + by_problem[p] = stub + stubs.append(stub) + baseline_rollout(base_sampler, stubs, base_dp, args, base_cache) + diagnose_views(checker, stubs, args, rubric_cache) + for r in targets: + p, sim = neighbor_map[r['problem']] + r['_rubric_diag'] = by_problem[p].get('_rubric_diag', '') + r['_rubric_src'], r['_neighbor_sim'] = p, sim + + +def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], + ci: int, base_dp: int, skill_dp: int, args: argparse.Namespace, + checker, rubric_cache: DiskCache, base_cache: DiskCache = None, + neighbor_map: Optional[Dict[str, Tuple[str, float]]] = None, + pool_answers: Optional[Dict[str, str]] = None + ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], List[Dict[str, Any]]]: + """view assign -> rubric-check (view A) -> skill-gen -> answer audit -> with-skill + greedy pass -> GRPO advantages. ``chunk`` arrives already baselined by draw_chunk. + With --xproblem-rubric, view A gets its NEIGHBOUR's rubric (see apply_neighbor_rubric).""" + hard = chunk + for r in hard: + r['_view'], r['_rubric_diag'] = _assign_view(r['problem'], args), '' + if args.xproblem_rubric and neighbor_map: + apply_neighbor_rubric(base_sampler, hard, neighbor_map, pool_answers or {}, base_dp, + args, checker, base_cache, rubric_cache) + else: + diagnose_views(checker, hard, args, rubric_cache) + + # skill-gen (thinking OFF), per-view prompt; re-sample problems with no clean candidate. + # View-A problems that cannot yield a training record (no [FAIL] rubric, or rubric + # leaked the answer) are dropped from training entirely -- skip their generation. + flat: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] + pending = [r for r in hard if not _viewa_dropped(r, args)] + for _ in range(args.skill_retries + 1): + if not pending: + break + sg_out = _run_samples(skill_sampler, [_view_prompt(r, args) for r in pending], + args.n_skills, args.skill_max_tokens, skill_dp, + temperature=args.skill_gen_temperature, + top_p=args.skill_gen_top_p, top_k=args.skill_gen_top_k) + still = [] + for r, seqs in zip(pending, sg_out): + got = False + for s in seqs: + resp = _clean_text(getattr(s, 'decoded', '') or '') + block = _extract_skill(resp) or '' + cand = {'skills': block, 'response': resp, 'parseable': bool(block), + 'view': r['_view'], 'leaked': None, 'leak_reason': '', + 'leak_source': '', 'with_pass': None, 'reward': None, 'rolls': [], + 'skillgen_stop': getattr(s, 'stop_reason', None), + 'skillgen_tokens': len(getattr(s, 'tokens', None) or [])} + r['_cands'].append(cand) + if block: + flat.append((r, cand)) + got = True + if not got: + still.append(r) + pending = still + + # Answer-bearing skill audit: deterministic verbatim-answer check only (no LLM). This + # is observability only; it records metrics for swanlab/jsonl, but does not block + # scoring, reward, advantage assignment, SFT candidate selection, or GRPO training. + for r, c in flat: + leaked = _answer_leaked(c['skills'], r['reference_answer']) + c['leaked'] = leaked + c['leak_reason'] = 'answer_verbatim' if leaked else '' + c['leak_source'] = 'deterministic' + + # with-skill greedy pass (T=0, M=1); reward = correct, absolute (group mean is baseline). + scored_inputs = flat + if scored_inputs: + ws_out = _run_samples(base_sampler, + [build_skill_solve_prompt(r['problem'], c['skills']) for r, c in scored_inputs], + 1, args.max_tokens, base_dp, temperature=0.0) + for (r, c), seqs in zip(scored_inputs, ws_out): + c['rolls'] = [_parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll()] + c['with_pass'] = 1.0 if c['rolls'][0]['correct'] else 0.0 + c['reward'] = c['with_pass'] + if args.format_in_reward: # unparseable candidates score 0 and still join the group + for r in hard: + for c in r['_cands']: + if c['reward'] is None: + c['reward'] = 0.0 + + _assign_advantages(hard, args) + return ([_full_record(r, ci) for r in chunk], _chunk_summary(chunk, ci, args), + _group_records(chunk, args)) + + +def _roll(x: Dict[str, Any]) -> Dict[str, Any]: + return {k: x[k] for k in ('pred', 'correct', 'terminated', 'passed', + 'stop_reason', 'gen_tokens', 'text')} + + +def _is_trainable(c: Dict[str, Any], args: argparse.Namespace) -> bool: + """Reaches the GRPO update iff advantage is non-zero. Leak flags are audit-only.""" + adv_nz = abs(c.get('advantage') or 0.0) > 1e-9 + if args.format_in_reward: + return adv_nz + return c.get('with_pass') is not None and adv_nz + + +def _full_record(r: Dict[str, Any], ci: int) -> Dict[str, Any]: + """Complete per-problem trace: init attempt, baseline, and all candidates.""" + init = r['_init'][0] + return { + 'record_type': 'problem', 'chunk': ci, 'problem': r['problem'], + 'reference_answer': r['reference_answer'], 'level': r.get('level', ''), + 'failed_first_try': r['_failed'], + 'init_attempt': {'text': init['text'], 'pred': init['pred'], 'correct': init['correct'], + 'terminated': init['terminated'], 'stop_reason': init['stop_reason'], + 'gen_tokens': init['gen_tokens']}, + 'baseline_pass': r['_baseline_pass'], 'is_hard': r['_hard'], + 'view': r.get('_view', ''), 'rubric_diag': r.get('_rubric_diag', ''), + # xproblem audit: which neighbour's rubric this problem borrowed, and how similar. + 'rubric_src': r.get('_rubric_src', ''), 'neighbor_sim': r.get('_neighbor_sim'), + 'baseline_rolls': [_roll(x) for x in r['_baseline_rolls']], + 'candidates': [{ + 'skills': c['skills'], 'response': c['response'], 'parseable': c['parseable'], + 'skillgen_stop': c.get('skillgen_stop'), 'skillgen_tokens': c.get('skillgen_tokens'), + 'leaked': c['leaked'], 'leak_reason': c['leak_reason'], 'leak_source': c['leak_source'], + 'with_pass': c['with_pass'], 'reward': c.get('reward'), 'advantage': c.get('advantage'), + 'grpo_adv': c.get('grpo_adv'), 'kept': c.get('kept'), + 'rolls': [_roll(x) for x in c['rolls']], + } for c in r['_cands']], + } + + +def _view_stats(problems: List[Dict[str, Any]], view: str) -> Dict[str, Any]: + pv = [r for r in problems if r.get('_view') == view] + cands = [c for r in pv for c in r['_cands'] if c['parseable']] + clean = [c for c in cands if c['leaked'] is False] + adopted = sum(1 for r in pv + if any(abs(c.get('advantage') or 0.0) > 1e-9 for c in r['_cands'])) + return {'n': len(pv), 'n_candidates_parseable': len(cands), 'n_clean': len(clean), + 'n_adopted_problems': adopted, 'adoption_rate': (adopted / len(pv)) if pv else 0.0} + + +def _mean(xs: List[float]) -> float: + return sum(xs) / len(xs) if xs else 0.0 + + +def _std(xs: List[float]) -> float: + if len(xs) < 2: + return 0.0 + m = _mean(xs) + return (sum((x - m) ** 2 for x in xs) / len(xs)) ** 0.5 + + +def _signal_stats(problems: List[Dict[str, Any]]) -> Dict[str, float]: + """The heart of 'is there a learning signal': per problem, the scored candidates form a + GRPO group. A group with zero reward variance (all skills solve, or none do -- the + hard-problem dead zone) gives NO gradient. Tracks that fraction plus reward level and + within-group variance so a collapse (all-0 or all-1) is visible immediately.""" + group_vars, all_rewards, zero_grad, groups = [], [], 0, 0 + for r in problems: + rewards = [c['reward'] for c in r['_cands'] if c.get('reward') is not None] + if len(rewards) < 2: + continue + groups += 1 + all_rewards.extend(rewards) + v = _std(rewards) + group_vars.append(v) + if v < 1e-9: # every skill got the same reward -> GRPO skips this problem + zero_grad += 1 + return {'n_groups': groups, 'zero_grad_frac': (zero_grad / groups) if groups else 0.0, + 'reward_mean': _mean(all_rewards), 'reward_std': _std(all_rewards), + 'group_reward_std_mean': _mean(group_vars)} + + +def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespace) -> Dict[str, Any]: + all_cands = [c for r in chunk for c in r['_cands']] + cands = [c for c in all_cands if c['parseable']] + scored = [c for c in cands if c['with_pass'] is not None] + clean = [c for c in cands if c['leaked'] is False] + ws_rolls = [x for c in scored for x in c['rolls']] + # viewa-dropped problems generate no candidates; keep acc/* on the generated subset + # so the with-skill/lift trend stays comparable across view_b_frac settings. + gen_probs = [r for r in chunk if r['_cands']] + base_acc = _mean([1.0 if r['_baseline_pass'] else 0.0 for r in gen_probs]) + ws_acc = _mean([1.0 if any(c.get('reward') for c in r['_cands']) else 0.0 for r in gen_probs]) + cand_pass_parseable = _mean([c['with_pass'] for c in scored]) + cand_pass_all = _mean([1.0 if c.get('reward') else 0.0 for c in all_cands]) + # base failure taxonomy (you asked whether skills fail because the base loops out of length) + classes = [_baseline_class(r) for r in chunk] + n_fail = sum(1 for c in classes if c != 'success') + skill_tokens = [c.get('skillgen_tokens') or 0 for c in cands] # skill-gen response length + trunc = sum(1 for r in chunk for c in r['_cands'] + for x in c['rolls'] if x['stop_reason'] == 'length') + rubric_answer_leaks = sum( + 1 for r in chunk + if r.get('_view') == 'A' and _rubric_has_fail(r.get('_rubric_diag')) and _rubric_answer_leaked(r)) + return { + 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), + 'n_failed_first_try': sum(1 for r in chunk if r['_failed']), + 'n_generated': len(all_cands), 'n_candidates_parseable': len(cands), + 'n_unparseable': len(all_cands) - len(cands), + 'parse_rate': (len(cands) / len(all_cands)) if all_cands else 0.0, + 'n_leaked': sum(1 for c in cands if c['leaked']), 'n_clean': len(clean), + 'leak_rate': (sum(1 for c in cands if c['leaked']) / len(cands)) if cands else 0.0, + 'n_reward_pos': sum(1 for c in scored if c['reward']), + 'n_rubric_answer_leaked': rubric_answer_leaks, + 'n_viewa_dropped': sum(1 for r in chunk if _viewa_dropped(r, args)), + 'n_train_samples': sum(1 for c in all_cands if _is_trainable(c, args)), + 'signal': _signal_stats(chunk), + 'fail_loop_frac': (sum(1 for c in classes if c == 'fail_loop') / n_fail) if n_fail else 0.0, + 'fail_wrong_frac': (sum(1 for c in classes if c == 'fail_wrong') / n_fail) if n_fail else 0.0, + 'skill_tokens_mean': _mean(skill_tokens), + 'withskill_trunc_frac': (trunc / len(ws_rolls)) if ws_rolls else 0.0, + 'avg_baseline_pass': base_acc, 'avg_withskill_pass': ws_acc, + 'avg_lift': ws_acc - base_acc, + 'candidate_withskill_pass_parseable': cand_pass_parseable, + 'candidate_withskill_pass_all': cand_pass_all, + 'termination_rate_withskill': _mean([1.0 if x['terminated'] else 0.0 for x in ws_rolls]), + 'view_A': _view_stats(chunk, 'A'), 'view_B': _view_stats(chunk, 'B'), + **_xproblem_stats(chunk, args), + } + + +def _xproblem_stats(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> Dict[str, Any]: + """Cross-problem pairing health: of the view-A problems, how many actually got a + neighbour's rubric (pair_rate) and how similar those neighbours were. Empty when off.""" + if not args.xproblem_rubric: + return {} + view_a = [r for r in chunk if r.get('_view') == 'A'] + paired = [r for r in view_a if r.get('_rubric_src')] + return {'xproblem': { + 'n_view_a': len(view_a), 'n_paired': len(paired), + 'pair_rate': (len(paired) / len(view_a)) if view_a else 0.0, + 'neighbor_sim_mean': _mean([r.get('_neighbor_sim', 0.0) for r in paired])}} + + +def _sft_record(r: Dict[str, Any], c: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: + """A view-A context-distillation sample: the OPEN-BOOK (rubric-conditioned) skill ``c`` + is the target, but the prompt is rebuilt as query-only (view B, no rubric) so the model + learns to produce it CLOSED-BOOK. Rides the same GRPOLoss with a positive constant + advantage (``--sft-weight``); single-step (old_logps=None) this reduces to + ``-w*logp + beta*KL`` -- a KL-anchored weighted cross-entropy toward the skill.""" + return { + 'data_id': r.get('data_id'), + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'view': 'B', 'diagnosis': '', 'rubric_src': '', 'orig_view': 'A', 'sft': True, + 'response': c['response'], 'skills': c['skills'], + 'skillgen_stop': c.get('skillgen_stop'), + 'advantage': float(args.sft_weight), 'grpo_adv': 0.0, 'kept': True, + 'reward': c['reward'], 'with_pass': c['with_pass']} + + +def _rubric_answer_leaked(r: Dict[str, Any]) -> bool: + """Hard safety gate for external API/rubric diagnosis. Unlike answer-bearing skills + generated by the policy itself, a rubric that contains the target final answer is an + external teacher leak and must not be distilled into view B.""" + return _answer_leaked(r.get('_rubric_diag', ''), r.get('reference_answer', '')) + + +def _viewa_dropped(r: Dict[str, Any], args: argparse.Namespace) -> bool: + """Under --viewa-sft a view-A problem trains ONLY through the SFT path (rubric with + [FAIL], answer-free). No-FAIL and rubric-leaked problems produce no training record + at all (no GRPO backflow: those prompts are query-only and would muddy the pure + view-B GRPO pool), so their skill-gen + executor scoring is skipped entirely.""" + return (bool(args.viewa_sft) and r.get('_view') == 'A' + and (not _rubric_has_fail(r.get('_rubric_diag')) + or _rubric_answer_leaked(r))) + + +def _group_records(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> List[Dict[str, Any]]: + """Training records. Under --viewa-sft, view A is SFT-only: a problem whose rubric + localised a failure (has ``[FAIL]``, answer-free) contributes ONE context-distillation + SFT sample (best parseable open-book skill -- preferring an executor-verified pass, + else any parseable candidate -- rebuilt query-only); no-FAIL / rubric-leaked view-A + problems are dropped (``_viewa_dropped``, generation already skipped). GRPO candidates + come from view B only, keeping that pool purely closed-book. Prompts are rebuilt from + the stored view/diagnosis by ``_skillgen_messages``.""" + out = [] + for r in chunk: + if not r['_hard']: + continue + if args.viewa_sft and r.get('_view') == 'A': + if _viewa_dropped(r, args): + continue + best = _best_sft_candidate(r, args) + if best is not None: + out.append(_sft_record(r, best, args)) + continue + for c in r['_cands']: + if _is_trainable(c, args): + out.append({ + 'data_id': r.get('data_id'), + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'view': r.get('_view', 'A'), 'diagnosis': r.get('_rubric_diag', ''), + 'rubric_src': r.get('_rubric_src', ''), 'sft': False, + 'response': c['response'], 'skills': c['skills'], + 'skillgen_stop': c.get('skillgen_stop'), + 'advantage': c['advantage'], 'grpo_adv': c['grpo_adv'], 'kept': c['kept'], + 'reward': c['reward'], 'with_pass': c['with_pass']}) + return out + + +# =========================================================================== +# Block G -- online GRPO training +# =========================================================================== +def _is_num(v: Any) -> bool: + try: + float(v) + return True + except (TypeError, ValueError): + return False + + +def _train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: + """Training sample = the exact skill-gen prompt (rebuilt by ``_skillgen_messages`` so + train/inference match) + the generated structured guidance response. ``key_rounds`` + selects the final assistant turn; Template masks the prompt and trains the whole + response (the key-round prefix already excludes the prompt-provided ).""" + msgs = _skillgen_messages( + rec['problem'], rec.get('view', 'A'), rec.get('diagnosis', ''), rec.get('rubric_src', '')) + return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], + 'user_data': {'key_rounds': [len(msgs)]}} + + +def _train_chunk(skill_model, ref_model, ckpt: CheckpointEngineManager, samples: List[Dict[str, Any]], + args: argparse.Namespace) -> Dict[str, Any]: + """On-policy GRPO update over one chunk, then sync weights. Micro-batches of + ``sft_batch_size`` accumulate gradients; ONE optimizer step is taken per PPO + mini-batch of ``ppo_mini_batch_size`` samples (0 -> a single step over the whole + chunk, the original behaviour). A frozen reference model provides ref_logps for the + SEAM-style KL penalty. + + Multi-step correctness: with more than one step over the SAME rollout, later + mini-batches see an already-updated policy, so we FREEZE the sampling-policy + ``old_logps`` (recomputed once, before any step) and let GRPOLoss form the PPO ratio + against them. A single step keeps ``old_logps=None`` (ratio==1, pure on-policy). + The batch is padded to a multiple of ``sft_batch_size`` with advantage-0 copies that + contribute no policy gradient. View-A context-distillation samples ride the same loss + with a positive constant advantage (``--sft-weight``); single-step (old_logps=None) + that is ``-w*logp + beta*KL`` (KL-anchored weighted cross-entropy).""" + trajs = [_train_trajectory(rec) for rec in samples] + advs = [float(rec['advantage']) for rec in samples] + rem = (-len(trajs)) % args.sft_batch_size + if rem: + trajs += [trajs[-1]] * rem + advs += [0.0] * rem + + n, sft = len(trajs), args.sft_batch_size + mini = args.ppo_mini_batch_size if args.ppo_mini_batch_size > 0 else n + mini = max(sft, (mini // sft) * sft) # align to a whole number of micro-batches + multi_step = mini < n + + # Freeze ref_logps (frozen model) and -- only when we take multiple steps -- the + # sampling-policy old_logps, BOTH before any optimizer step touches the weights. With + # a single step old_logps stays None so GRPOLoss uses ratio==1 (pure on-policy). + micro_ref, micro_old = [], [] + for i in range(0, n, sft): + mb = trajs[i:i + sft] + micro_ref.append(ref_model.forward_only(inputs=mb).get('logps')) + micro_old.append(skill_model.forward_only(inputs=mb).get('logps') if multi_step else None) + + micro, n_steps = 0, 0 + for ms in range(0, n, mini): + for i in range(ms, min(ms + mini, n), sft): + k = i // sft + skill_model.forward_backward(inputs=trajs[i:i + sft], + advantages=advs[i:i + sft], + old_logps=micro_old[k], + ref_logps=micro_ref[k]) + micro += 1 + skill_model.clip_grad_and_step() + n_steps += 1 + ckpt.sync_weights(merge_and_sync=True) + metric = skill_model.calculate_metric(is_training=True) + n_sft = sum(1 for s in samples if s.get('sft')) + return {'n_samples': len(samples), 'n_sft': n_sft, 'n_grpo': len(samples) - n_sft, + 'n_steps': n_steps, 'n_micro_batches': micro, + 'metric': {k: (float(v) if _is_num(v) else v) for k, v in (metric or {}).items()}} + + +# =========================================================================== +# Block H -- fixed-holdout eval + metric formatting +# =========================================================================== +def run_greedy_eval(base_sampler, skill_sampler, eval_records: List[Dict[str, Any]], + ci: int, rounds: int, base_dp: int, skill_dp: int, args: argparse.Namespace, + base_cache: DiskCache + ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], Dict[str, float]]: + """SEAM ``val-core/math/acc/mean@1`` on the fixed holdout: ONE greedy skill (T=0) per + problem into ONE greedy base solve (T=0). Eval ALWAYS uses view B (query-only, the + deployment form: no rubric, since rubric needs an online teacher unavailable at deploy); + no leak filter (acc scores correctness alone). Baseline reuses the disk cache.""" + baseline_rollout(base_sampler, eval_records, base_dp, args, base_cache) + for r in eval_records: + r['_view'], r['_rubric_diag'] = 'B', '' + sg_out = _run_samples(skill_sampler, [_view_prompt(r, args) for r in eval_records], + 1, args.skill_max_tokens, skill_dp, temperature=0.0) + skills = [] + for seqs in sg_out: + if not seqs: + skills.append(('', '')) + continue + sresp = _clean_text(getattr(seqs[0], 'decoded', '') or '') + skills.append((_extract_skill(sresp) or '', sresp)) + ws_out = _run_samples(base_sampler, + [build_skill_solve_prompt(r['problem'], sk) for r, (sk, _) in zip(eval_records, skills)], + 1, args.max_tokens, base_dp, temperature=0.0) + recs = [] + for r, (sk, sresp), seqs in zip(eval_records, skills, ws_out): + roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + recs.append({ + 'record_type': 'eval_problem', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, + 'data_id': r.get('data_id'), + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'view': r['_view'], 'rubric_diag': r.get('_rubric_diag', ''), + 'baseline_pass': r['_baseline_pass'], 'skill': sk, 'skill_parseable': bool(sk), + 'skill_response': sresp, 'withskill_pred': roll['pred'], + 'withskill_correct': roll['correct'], 'withskill_terminated': roll['terminated'], + 'withskill_stop_reason': roll['stop_reason'], 'withskill_text': roll['text'], + }) + acc = lambda rs: sum(1 for x in rs if x['withskill_correct']) / len(rs) if rs else 0.0 + ws = acc(recs) # all view B (deployment form) + base = sum(x['baseline_pass'] for x in recs) / len(recs) if recs else 0.0 + fmt = (sum(1 for x in recs if x['skill_parseable']) / len(recs)) if recs else 0.0 + term = (sum(1 for x in recs if x['withskill_terminated']) / len(recs)) if recs else 0.0 + summary = {'record_type': 'eval_summary', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, + 'n': len(recs), 'view': 'B', 'acc_mean1': ws, + 'baseline_acc_mean1': base, 'lift_mean1': ws - base, + 'format_mean1': fmt, 'term_mean1': term} + metrics = {'core/math/acc/mean@1': ws, 'core/math/baseline_acc/mean@1': base, + 'core/math/lift/mean@1': ws - base, 'core/math/format/mean@1': fmt, + 'core/math/term/mean@1': term} + return recs, summary, metrics + + +def _trend_line(hist: List[Dict[str, float]], window: int, rounds_done: int) -> Optional[str]: + """Contrast the FIRST vs the most recent ``window`` chunks -- if RFT works, adoption + and lift on recent (fresh) chunks exceed the early baseline.""" + if len(hist) < 2 * window: + return None + base, rec = hist[:window], hist[-window:] + m = lambda xs, k: sum(h[k] for h in xs) / len(xs) + return (f'[trend] first {window} vs last {window} | ' + f'adopt A {m(base,"aA"):.2f}->{m(rec,"aA"):.2f} B {m(base,"aB"):.2f}->{m(rec,"aB"):.2f} | ' + f'lift {m(base,"lift"):+.3f}->{m(rec,"lift"):+.3f} | ' + f'0grad {m(base,"zero_grad"):.2f}->{m(rec,"zero_grad"):.2f} | ' + f'pos/chunk {m(base,"pos"):.1f}->{m(rec,"pos"):.1f} | rounds={rounds_done}') + + +def _swan_metrics(summary: Dict[str, Any], log: Optional[Dict[str, Any]]) -> Dict[str, float]: + """Flat swanlab dict. ``signal/*`` is the primary health group (is GRPO getting a + gradient at all); ``acc/*`` is the effect; the rest diagnose why. acc/adopt/skill are + only logged when the chunk produced a scored group, so idle chunks don't dip charts.""" + sig = summary['signal'] + d: Dict[str, float] = { + # --- signal: the FIRST thing to watch (no variance -> no learning) --- + 'signal/zero_grad_frac': sig['zero_grad_frac'], 'signal/n_groups': sig['n_groups'], + 'signal/reward_mean': sig['reward_mean'], 'signal/reward_std': sig['reward_std'], + 'signal/group_reward_std_mean': sig['group_reward_std_mean'], + 'signal/n_train_samples': summary['n_train_samples'], + 'signal/n_reward_pos': summary['n_reward_pos'], + # --- skill format / leak health --- + 'skill/parse_rate': summary['parse_rate'], 'skill/tokens_mean': summary['skill_tokens_mean'], + 'leak/rate': summary['leak_rate'], 'leak/n': summary['n_leaked'], + # --- base failure taxonomy (loop-out-of-length vs plain wrong) --- + 'fail/loop_frac': summary['fail_loop_frac'], 'fail/wrong_frac': summary['fail_wrong_frac'], + 'fail/frac_first_try': (summary['n_failed_first_try'] / summary['n']) if summary['n'] else 0.0, + # --- view-A routing: dropped = no-FAIL rubric or rubric leak (SFT-only view A) --- + 'viewa/dropped_frac': (summary['n_viewa_dropped'] / summary['view_A']['n'] + if summary['view_A']['n'] else 0.0), + } + bal = summary.get('balance') or {} + if bal.get('enabled'): + d.update({'balance/n_drawn': bal['n_drawn'], 'balance/n_baseline_fresh': bal['n_baseline_fresh'], + 'balance/selected_success_frac': bal['selected_success_frac']}) + xp = summary.get('xproblem') or {} + if xp: + d.update({'xproblem/pair_rate': xp['pair_rate'], 'xproblem/neighbor_sim_mean': xp['neighbor_sim_mean']}) + if sig['n_groups'] > 0: + d.update({'acc/baseline_pass': summary['avg_baseline_pass'], + 'acc/withskill_pass': summary['avg_withskill_pass'], 'acc/lift': summary['avg_lift'], + 'candidate/withskill_pass_parseable': summary['candidate_withskill_pass_parseable'], + 'candidate/withskill_pass_all': summary['candidate_withskill_pass_all'], + 'adopt/A': summary['view_A']['adoption_rate'], + 'adopt/B': summary['view_B']['adoption_rate'], + 'term/withskill': summary['termination_rate_withskill'], + 'term/withskill_trunc_frac': summary['withskill_trunc_frac']}) + if log: + d['train/n_steps'] = log['n_steps'] + d['train/n_micro_batches'] = log['n_micro_batches'] + for k, v in (log.get('metric') or {}).items(): + if not _is_num(v): + continue + if k.startswith('learning rate'): + if 'group 1' in k: + d['train/lr'] = float(v) + else: + d[f'train/{k.replace(" ", "_")}'] = float(v) + return d + + +def _view_a_rubric_leak_metrics(chunk: List[Dict[str, Any]], + pool_answers: Optional[Dict[str, str]] = None) -> Dict[str, float]: + """Swanlab-only audit for answer leakage in view-A rubric text. This never changes + rewards, advantages, filtering, or training records.""" + view_a = [r for r in chunk if r.get('_view') == 'A'] + with_diag = [r for r in view_a if r.get('_rubric_diag')] + target_leaks = sum(1 for r in with_diag + if _answer_leaked(r.get('_rubric_diag', ''), r.get('reference_answer', ''))) + source_leaks = 0 + pool_answers = pool_answers or {} + for r in with_diag: + src = r.get('_rubric_src') + src_ref = pool_answers.get(src, '') if src else r.get('reference_answer', '') + if _answer_leaked(r.get('_rubric_diag', ''), src_ref): + source_leaks += 1 + n = len(with_diag) + return { + 'rubric_leak/n_view_a': float(len(view_a)), + 'rubric_leak/n_checked': float(n), + 'rubric_leak/target_answer_n': float(target_leaks), + 'rubric_leak/target_answer_rate': (target_leaks / n) if n else 0.0, + 'rubric_leak/source_answer_n': float(source_leaks), + 'rubric_leak/source_answer_rate': (source_leaks / n) if n else 0.0, + } + + +# =========================================================================== +# Block F -- components, args, main +# =========================================================================== +def init_components(args: argparse.Namespace): + """Default 8-GPU layout: rank 0 trains the actor, rank 1 hosts a frozen ref model, + 2-3 skill_sampler (synced), 4-7 base_sampler (frozen). Returns + (skill_model, ref_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp).""" + r0, r1 = TRAIN_GPUS, TRAIN_GPUS + REF_GPUS + r2, r3 = r1 + SKILL_SAMPLER_GPUS, NUM_GPUS + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, lazy_collect=False, groups=[ + DeviceGroup(name='train', ranks=list(range(0, r0)), device_type='GPU'), + DeviceGroup(name='ref', ranks=list(range(r0, r1)), device_type='GPU'), + DeviceGroup(name='skill_sampler', ranks=list(range(r1, r2)), device_type='GPU'), + DeviceGroup(name='base_sampler', ranks=list(range(r2, r3)), device_type='GPU')]) + + train_mesh = DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_DP, fsdp_size=TRAIN_FSDP) + skill_model = TransformersModel(model_id=MODEL_ID, device_mesh=train_mesh, remote_group='train', + ddp_config={'find_unused_parameters': False}) + skill_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) + skill_model.set_template(Template, model_id=MODEL_ID, enable_thinking=False, + max_length=args.max_model_len, truncation_strategy='delete') + skill_model.set_processor(InputProcessor, padding_free=False) + skill_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon, beta=args.kl_beta) + skill_model.set_optimizer('AdamW', lr=args.lr) + skill_model.set_lr_scheduler('CosineWarmupScheduler', num_warmup_steps=10, + num_training_steps=args.max_train_rounds) + + ref_mesh = DeviceMesh.from_sizes(world_size=REF_GPUS, dp_size=REF_DP, fsdp_size=REF_FSDP) + ref_model = TransformersModel(model_id=MODEL_ID, device_mesh=ref_mesh, remote_group='ref', + ddp_config={'find_unused_parameters': False}) + ref_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) + ref_model.set_template(Template, model_id=MODEL_ID, enable_thinking=True, + max_length=args.max_model_len, truncation_strategy='delete') + ref_model.set_processor(InputProcessor, padding_free=False) + ref_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) + + def _sampler(group, world, enable_thinking: bool = True): + s = vLLMSampler(model_id=MODEL_ID, + engine_args={'gpu_memory_utilization': GPU_MEM, + 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, + device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), + remote_group=group) + s.set_template(Template, model_id=MODEL_ID, enable_thinking=enable_thinking, max_length=args.max_model_len) + return s + + skill_sampler = _sampler('skill_sampler', SKILL_SAMPLER_GPUS, enable_thinking=False) + # base_sampler is shared by the main thread and the baseline-prefetch thread -> serialise. + base_sampler = _LockedSampler(_sampler('base_sampler', BASE_SAMPLER_GPUS, enable_thinking=True)) + ckpt = CheckpointEngineManager(model=skill_model, sampler=skill_sampler) + return skill_model, ref_model, skill_sampler, base_sampler, ckpt, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS + + +def _build_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--dataset', choices=('aops', 'math'), default='aops') + p.add_argument('--n', type=int, default=2000, help='Problems loaded into the draw pool.') + p.add_argument('--pool-offset', type=int, default=0, + help='Skip this many shuffled non-eval records before building the train pool; ' + 'useful to avoid cold-start SFT data ranges.') + p.add_argument('--exclude-data-ids', default='', + help='Comma-separated jsonl files whose data_id/problem keys are excluded ' + 'from train/eval selection, e.g. coldstart_sft.jsonl.') + p.add_argument('--seed', type=int, default=42) + p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--eval-size', type=int, default=128, help='Fixed holdout size (0 disables).') + p.add_argument('--eval-every', type=int, default=5, help='Run holdout eval every N chunks.') + p.add_argument('--chunk-size', type=int, default=16) + p.add_argument('--balance', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--balance-success-frac', type=float, default=0.4, + help='Target fraction of the chunk the base solves (rest are base-fail).') + p.add_argument('--balance-loop-frac', type=float, default=0.5) + p.add_argument('--balance-max-draws-mult', type=int, default=8) + p.add_argument('--n-skills', type=int, default=8) + p.add_argument('--view-b-frac', type=float, default=0.5) + p.add_argument('--viewa-frac-start', type=float, default=None, + help='Enable the view-A curriculum: chunk 0 uses this view-A share ' + '(view_b_frac = 1 - share), decaying linearly to --viewa-frac-end ' + 'over --viewa-decay-chunks chunks, then holding. Overrides ' + '--view-b-frac for every chunk.') + p.add_argument('--viewa-frac-end', type=float, default=0.1) + p.add_argument('--viewa-warmup-chunks', type=int, default=0, + help='Hold the view-A share at --viewa-frac-start for this many chunks ' + 'before the linear decay begins.') + p.add_argument('--viewa-decay-chunks', type=int, default=40) + p.add_argument('--xproblem-rubric', action=argparse.BooleanOptionalAction, default=False, + help='When enabled, view A uses a bag-of-words NEIGHBOUR problem rubric. ' + 'Default is off: each view-A problem uses its own baseline attempt, ' + 'while the API diagnosis prompt is constrained to be answer-free and ' + 'method-level only.') + p.add_argument('--skill-retries', type=int, default=2) + p.add_argument('--skill-gen-temperature', type=float, default=1.0) + p.add_argument('--skill-gen-top-p', type=float, default=1.0) + p.add_argument('--skill-gen-top-k', type=int, default=-1) + p.add_argument('--max-model-len', type=int, default=16384) + p.add_argument('--max-tokens', type=int, default=8192) + p.add_argument('--skill-max-tokens', type=int, default=8192) + p.add_argument('--rubric-workers', type=int, default=16) + p.add_argument('--rubric-retries', type=int, default=2, + help='Retry failed/timeout rubric diagnose calls this many times before ' + 'falling back to an empty diagnosis without caching the failure.') + p.add_argument('--sft-batch-size', type=int, default=8, + help='Driver micro-batch (gradient-accumulation unit); multiple of train dp.') + p.add_argument('--ppo-mini-batch-size', type=int, default=0, + help='Samples per optimizer step (SEAM-style PPO mini-batch). 0 = one step ' + 'over the whole chunk (on-policy, ratio==1). When >0 and smaller than ' + 'the trainable count, multiple steps are taken over the same rollout and ' + 'old_logps are frozen so the PPO ratio/clip stays valid. Rounded down to ' + 'a multiple of --sft-batch-size.') + p.add_argument('--grpo-epsilon', type=float, default=0.2) + p.add_argument('--adv-clip', type=float, default=3.0, + help='Symmetric clip for group-relative advantages; <=0 disables clipping.') + p.add_argument('--kl-beta', type=float, default=0.001, + help='SEAM-style reference KL coefficient for GRPOLoss.') + p.add_argument('--viewa-sft', action=argparse.BooleanOptionalAction, default=True, + help='Route view-A problems to online context distillation (SFT on the best ' + 'passing open-book skill, rebuilt as a query-only prompt) instead of GRPO. ' + 'View B stays GRPO; both share one optimizer step.') + p.add_argument('--sft-weight', type=float, default=0.5, + help='Advantage magnitude (lambda) for view-A SFT distillation samples in the ' + 'shared GRPOLoss. Single-step this is -lambda*logp + beta*KL (weighted CE). ' + 'Keep <~1 so one-signed SFT gradients do not swamp mean-0 GRPO advantages.') + p.add_argument('--sft-target-len', type=int, default=550, + help='Target skill length (chars) for view-A SFT distillation: among passing ' + 'candidates pick the one CLOSEST to this length (not the shortest). Set near ' + 'the empirical pass-rate peak (~500-600 here) so distillation neither shrinks ' + 'rollouts toward zero nor lets them grow unbounded.') + p.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--lr', type=float, default=6e-6) + p.add_argument('--max-train-rounds', type=int, default=1500) + p.add_argument('--save-rounds', type=int, default=200) + p.add_argument('--trend-every', type=int, default=10) + p.add_argument('--output-dir', default='./output/reflexion_skill') + p.add_argument('--cache-dir', default='', help='Baseline/rubric cache dir (default /cache).') + p.add_argument('--no-cache', action='store_true', help='Disable disk cache read/write.') + p.add_argument('--prefetch-baseline', action=argparse.BooleanOptionalAction, default=True, + help='Prefetch next chunk base baseline on a background thread (overlaps ' + 'with skill-gen; base_sampler is frozen so it never blocks the trainer).') + p.add_argument('--swanlab-project', default='twinkle') + p.add_argument('--swanlab-exp', default='') + args = p.parse_args() + if args.sft_batch_size % TRAIN_DP != 0: + raise ValueError(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple of train dp ({TRAIN_DP})') + if args.chunk_size < 1: + raise ValueError('--chunk-size must be >= 1') + args.balance_success_frac = max(0.0, min(1.0, args.balance_success_frac)) + return args + + +def _write(handle, row: Dict[str, Any]) -> None: + handle.write(json.dumps(row, ensure_ascii=False) + '\n') + + +def main() -> None: + args = _build_args() + records, eval_records, neighbor_map, pool_answers, data_stats = _load_records(args) + if len(records) < args.chunk_size: + raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded ({len(records)}); raise --n') + + os.makedirs(args.output_dir, exist_ok=True) + gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') + eval_path = os.path.join(args.output_dir, 'eval_records.jsonl') + data_path = os.path.join(args.output_dir, 'skill_dataset.jsonl') + train_log_path = os.path.join(args.output_dir, 'train_log.jsonl') + + if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY')): + sys.stderr.write('[rft] WARNING: no LLM backup env; view-A rubric check disabled ' + '(leak filter is deterministic, unaffected)\n') + + use_swan = swanlab is not None and os.environ.get('SWANLAB_MODE') != 'disabled' + if use_swan: + swanlab.init(project=args.swanlab_project, experiment_name=(args.swanlab_exp or None), + config={'model': MODEL_ID, 'dataset': args.dataset, 'n': len(records), + 'eval_n': len(eval_records), 'n_skills': args.n_skills, + 'view_b_frac': args.view_b_frac, 'balance': args.balance, + 'balance_success_frac': args.balance_success_frac, + 'skill_gen_temp': args.skill_gen_temperature, + 'grpo_epsilon': args.grpo_epsilon, 'kl_beta': args.kl_beta, + 'lr': args.lr}) + + skill_model, ref_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp = init_components(args) + checker = build_rubric_checker() + if checker is None: + sys.stderr.write('[rft] no LLM backup env -> view-A rubric process-check DISABLED\n') + + cache_dir = args.cache_dir or os.path.join(args.output_dir, 'cache') + os.makedirs(cache_dir, exist_ok=True) + use_cache = not args.no_cache + base_cache = DiskCache(os.path.join(cache_dir, 'baseline.jsonl'), use_cache) + eval_base_cache = DiskCache(os.path.join(cache_dir, 'eval_baseline.jsonl'), use_cache) + rubric_cache = DiskCache(os.path.join(cache_dir, 'rubric.jsonl'), use_cache) + if args.xproblem_rubric: + sys.stderr.write(f'[rft] xproblem-rubric ON: {len(neighbor_map)} bag-of-words pairs\n') + + cfg = {'record_type': 'config', 'model': MODEL_ID, 'dataset': args.dataset, + 'n': len(records), 'eval_n': len(eval_records), 'seed': args.seed, + 'pool_offset': args.pool_offset, 'exclude_data_ids': args.exclude_data_ids, + 'excluded_records': data_stats.get('excluded_records', 0), + 'numeric_only': args.numeric_only, 'raw_loaded': data_stats['raw_loaded'], + 'numeric_dropped': data_stats['numeric_dropped'], 'eval_every': args.eval_every, + 'n_skills': args.n_skills, 'view_b_frac': args.view_b_frac, + 'viewa_frac_start': args.viewa_frac_start, 'viewa_frac_end': args.viewa_frac_end, + 'viewa_warmup_chunks': args.viewa_warmup_chunks, + 'viewa_decay_chunks': args.viewa_decay_chunks, + 'skill_retries': args.skill_retries, 'balance': args.balance, + 'balance_success_frac': args.balance_success_frac, + 'skill_gen_temp': args.skill_gen_temperature, 'reward': 'greedy_binary(correct)', + 'advantage': 'group_relative', 'format_in_reward': args.format_in_reward, 'cache': use_cache, + 'rubric_check': f'fixed_math_{len(_MATH_RUBRIC)}crit+term(viewA)' if checker else 'disabled', + 'xproblem_rubric': args.xproblem_rubric, + 'viewa_sft': args.viewa_sft, 'sft_weight': args.sft_weight, + 'sft_target_len': args.sft_target_len, + 'adv_clip': args.adv_clip, + 'grpo_epsilon': args.grpo_epsilon, 'kl_beta': args.kl_beta, 'lr': args.lr, + 'train_gpus': TRAIN_GPUS, 'ref_gpus': REF_GPUS, 'ref_fsdp': REF_FSDP, + 'train_fsdp': TRAIN_FSDP, 'train_dp': TRAIN_DP, + 'skill_sampler_gpus': SKILL_SAMPLER_GPUS, 'base_sampler_gpus': BASE_SAMPLER_GPUS, + 'max_train_rounds': args.max_train_rounds, 'started': int(time.time())} + sys.stderr.write(f'[rft] raw={data_stats["raw_loaded"]} numeric_drop={data_stats["numeric_dropped"]} ' + f'train={len(records)} eval={len(eval_records)} {args.dataset}; ' + f'train_gpus={TRAIN_GPUS} ref_gpus={REF_GPUS} train_fsdp={TRAIN_FSDP} ' + f'train_dp={TRAIN_DP} skill_dp={skill_dp} base_dp={base_dp}\n') + + hist: List[Dict[str, float]] = [] + rounds = 0 + pool = ProblemPool(records, args.seed) + with open(gen_path, 'w', encoding='utf-8') as gen_f, \ + open(eval_path, 'w', encoding='utf-8') as eval_f, \ + open(data_path, 'w', encoding='utf-8') as data_f, \ + open(train_log_path, 'w', encoding='utf-8') as tlog: + for f in (gen_f, eval_f, data_f, tlog): + _write(f, cfg) + gstep = 0 + # base baseline (frozen sampler, disk-cached, no weight-sync) is prefetched on a + # background thread while the current chunk generates: the skill-gen phase uses + # skill_sampler, so prefetching the NEXT chunk's baseline on base_sampler overlaps + # it and fills base_cache so the next draw_chunk hits it. base_sampler is wrapped in + # _LockedSampler, so the prefetch and this chunk's with-skill scoring never issue a + # base .sample() concurrently. It never touches the trainer or on-policy generation. + prefetch_pool = ThreadPoolExecutor(max_workers=1) if args.prefetch_baseline else None + pending: Optional[Any] = None + + def _prefetch(peeked: List[Dict[str, Any]]) -> None: + if peeked: + baseline_rollout(base_sampler, peeked, base_dp, args, base_cache) + + # Baseline (round 0) eval BEFORE any training: measures the untrained skill model on + # the fixed holdout so every later eval has a step-0 reference point on the same axis. + if eval_records: + eval_recs, eval_summary, eval_metrics = run_greedy_eval( + base_sampler, skill_sampler, eval_records, -1, rounds, base_dp, skill_dp, + args, eval_base_cache) + for rec in eval_recs: + _write(eval_f, rec) + _write(eval_f, eval_summary) + eval_f.flush() + if use_swan: + swanlab.log({f'eval/{k}': v for k, v in eval_metrics.items()}, step=0) + sys.stderr.write( + f'[eval] g-1 (init): n={eval_summary["n"]} viewB mean@1 ' + f'acc={eval_summary["baseline_acc_mean1"]:.3f}->{eval_summary["acc_mean1"]:.3f} ' + f'lift={eval_summary["lift_mean1"]:+.3f} ' + f'fmt={eval_summary["format_mean1"]:.2f} rounds={rounds}\n') + + # Each chunk is drawn fresh + RE-GENERATED with the current policy (on-policy); + # with --balance, draw_chunk keeps drawing until the base fail:success mix hits target. + while rounds < args.max_train_rounds: + if pending is not None: + pending.result() # finish last round's prefetch before drawing (cache-warm) + pending = None + if args.viewa_frac_start is not None: + args.view_b_frac = _curriculum_view_b_frac(gstep, args) + chunk, balance = draw_chunk(pool, base_sampler, base_dp, args, base_cache) + if prefetch_pool is not None: + peeked = pool.peek(int(args.chunk_size * (args.balance_max_draws_mult if args.balance else 1))) + pending = prefetch_pool.submit(_prefetch, peeked) + full, summary, groups = process_chunk( + base_sampler, skill_sampler, chunk, gstep, base_dp, skill_dp, + args, checker, rubric_cache, base_cache, neighbor_map, pool_answers) + summary['balance'] = balance + summary['view_b_frac'] = round(args.view_b_frac, 4) + + log = None + if groups: + log = _train_chunk(skill_model, ref_model, ckpt, groups, args) + rounds += 1 + log.update({'record_type': 'train_round', 'round': rounds, 'chunk': gstep, + 'epoch': pool.epoch, 'ts': int(time.time())}) + _write(tlog, log) + tlog.flush() + if rounds % args.save_rounds == 0: + skill_model.save(f'skill-rft-{rounds}', output_dir=args.output_dir) + + summary['rounds_done'], summary['epoch'] = rounds, pool.epoch + for rec in full: + _write(gen_f, rec) + _write(gen_f, summary) + gen_f.flush() + for v in groups: + _write(data_f, v) + data_f.flush() + + sa, sb, sig = summary['view_A'], summary['view_B'], summary['signal'] + hist.append({'aA': sa['adoption_rate'], 'aB': sb['adoption_rate'], + 'lift': summary['avg_lift'], 'pos': summary['n_reward_pos'], + 'zero_grad': sig['zero_grad_frac']}) + bal_str = (f'bal {balance["selected_fail"]}f/{balance["selected_success"]}s ' + f'(drew {balance["n_drawn"]}/fresh {balance["n_baseline_fresh"]}' + + ('!' if balance.get('budget_hit') else '') + ') ') if balance.get('enabled') else '' + xp = summary.get('xproblem') + xp_str = f'pair={xp["pair_rate"]:.2f}@{xp["neighbor_sim_mean"]:.2f} ' if xp else '' + tr_str = (f'train={log["n_grpo"]}g/{log["n_sft"]}s ' if log + else f'train={summary["n_train_samples"]} ') + sys.stderr.write( + f'[gen] e{pool.epoch} g{gstep}: {bal_str}n={summary["n"]} ' + f'clean={summary["n_clean"]} leak={summary["leak_rate"]:.2f} {tr_str}' + f'0grad={sig["zero_grad_frac"]:.2f} R={sig["reward_mean"]:.2f}±{sig["reward_std"]:.2f} ' + f'acc={summary["avg_baseline_pass"]:.2f}->{summary["avg_withskill_pass"]:.2f} ' + f'lift={summary["avg_lift"]:+.3f} {xp_str}' + f'A[{sa["n"]} {sa["adoption_rate"]:.2f}] B[{sb["n"]} {sb["adoption_rate"]:.2f}] ' + f'rounds={rounds}\n') + if use_swan: + swan_metrics = _swan_metrics(summary, log) + swan_metrics.update(_view_a_rubric_leak_metrics(chunk, pool_answers)) + swan_metrics['train/view_b_frac'] = float(args.view_b_frac) + swanlab.log(swan_metrics, step=gstep) + + if eval_records and (gstep + 1) % args.eval_every == 0: + eval_recs, eval_summary, eval_metrics = run_greedy_eval( + base_sampler, skill_sampler, eval_records, gstep, rounds, base_dp, skill_dp, + args, eval_base_cache) + for rec in eval_recs: + _write(eval_f, rec) + _write(eval_f, eval_summary) + eval_f.flush() + if use_swan: + swanlab.log({f'eval/{k}': v for k, v in eval_metrics.items()}, step=gstep) + sys.stderr.write( + f'[eval] g{gstep}: n={eval_summary["n"]} viewB mean@1 ' + f'acc={eval_summary["baseline_acc_mean1"]:.3f}->{eval_summary["acc_mean1"]:.3f} ' + f'lift={eval_summary["lift_mean1"]:+.3f} ' + f'fmt={eval_summary["format_mean1"]:.2f} rounds={rounds}\n') + + if (gstep + 1) % args.trend_every == 0: + tl = _trend_line(hist, args.trend_every, rounds) + if tl: + sys.stderr.write(tl + '\n') + gstep += 1 + + if prefetch_pool is not None: + if pending is not None: + pending.result() + prefetch_pool.shutdown(wait=True) + base_cache.close() + eval_base_cache.close() + rubric_cache.close() + skill_model.save('skill-rft-final', output_dir=args.output_dir) + sys.stderr.write(f'[rft] done: {rounds} rounds over {gstep} chunks / {pool.epoch} epochs -> {data_path}\n') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/legacy/train_reflexion_skill.sh b/cookbook/exp/legacy/train_reflexion_skill.sh new file mode 100755 index 000000000..e732a6a3f --- /dev/null +++ b/cookbook/exp/legacy/train_reflexion_skill.sh @@ -0,0 +1,64 @@ +#!/bin/bash +# Online GRPO RFT for the reflexion skill generator (unified, self-contained, cached). +# GPUs: 8 — default high-memory layout uses rank 0 for actor training, rank 1 for +# the frozen ref model, ranks 2-3 for skill sampler (synced), and ranks 4-7 for +# base sampler (frozen). Override TRAIN_GPUS / REF_GPUS / SKILL_SAMPLER_GPUS / +# BASE_SAMPLER_GPUS for other layouts. Per chunk: base greedy +# solve -> rubric process-check (view A) -> +# skill-gen (thinking ON, N candidates) -> deterministic leak filter -> with-skill greedy +# pass -> group-relative advantage -> ONE on-policy GRPO step -> sync weights. +# +# Baseline rollouts + rubric diagnoses are disk-cached (output-dir/cache/*.jsonl), so a +# restart skips re-sampling them; skill-gen is on-policy and never cached. The next chunk's +# baseline is prefetched on a background thread (overlaps skill-gen; base sampler is frozen). +# +# The view-A rubric process-check uses the backup teacher API (set LLM_BACKUP_*). Without +# it the run still works: view A degrades to query-only and the leak filter stays +# deterministic (no teacher needed). + +set -euo pipefail + +export GEN_MODEL_ID=${GEN_MODEL_ID:-Qwen/Qwen3-4B} +export GEN_GPU_MEM=${GEN_GPU_MEM:-0.8} +# Datasets are pulled from ModelScope via twinkle.Dataset (ms://AI-MO/aops or +# ms://modelscope/competition_math); override AOPS_DATASET_ID / MATH_DATASET_ID to change. +# Teacher API for the view-A rubric process-check (optional; leak filter is deterministic). +export LLM_BACKUP_API_KEY=${LLM_BACKUP_API_KEY:-} +export LLM_BACKUP_BASE_URL=${LLM_BACKUP_BASE_URL:-https://dashscope.aliyuncs.com/compatible-mode/v1} +export LLM_BACKUP_MODEL=${LLM_BACKUP_MODEL:-qwen3.7-max} +export EXCLUDE_DATA_IDS=${EXCLUDE_DATA_IDS:-./output/reflexion_coldstart_sft/coldstart_sft.jsonl} + +python cookbook/exp/embedding/train_reflexion_skill.py \ + --dataset aops \ + --n 10000 \ + --numeric-only \ + --chunk-size 64 \ + --n-skills 8 \ + --viewa-frac-start 1.0 \ + --viewa-frac-end 0.1 \ + --viewa-warmup-chunks 20 \ + --viewa-decay-chunks 40 \ + --skill-retries 2 \ + --balance \ + --balance-success-frac 0.2 \ + --balance-loop-frac 0.5 \ + --balance-max-draws-mult 8 \ + --max-tokens 8192 \ + --skill-max-tokens 4096 \ + --max-model-len 16384 \ + --eval-size 128 \ + --exclude-data-ids "${EXCLUDE_DATA_IDS}" \ + --eval-every 5 \ + --sft-batch-size 8 \ + --ppo-mini-batch-size 0 \ + --grpo-epsilon 0.2 \ + --kl-beta 0.001 \ + --format-in-reward \ + --lr 1e-6 \ + --max-train-rounds 1500 \ + --save-rounds 200 \ + --trend-every 10 \ + --prefetch-baseline \ + --output-dir ./output/reflexion_skill_curriculum \ + --swanlab-project twinkle \ + --swanlab-exp reflexion_skill_curriculum diff --git a/cookbook/exp/legacy/train_reflexion_skill_old.py b/cookbook/exp/legacy/train_reflexion_skill_old.py new file mode 100644 index 000000000..757d9d915 --- /dev/null +++ b/cookbook/exp/legacy/train_reflexion_skill_old.py @@ -0,0 +1,2022 @@ +"""Online GRPO RFT for the reflexion skill generator (self-contained, cached). + +Trains an INDEPENDENT skill model to write reusable skills that, injected into a +FROZEN base solver's system prompt, raise its accuracy. The base is never trained; +it only produces the reward. Per chunk: base greedy solve -> rubric process-check +(view A) -> skill-gen (thinking ON, N candidates) -> leak filter -> with-skill +greedy pass -> group-relative advantage -> ONE on-policy GRPO step -> sync weights. +Reward is deterministic (T=0, M=1) binary correctness; A = (R - mean) / (std + eps) +within each problem-group, so std=0 groups give no gradient (GRPO variance selects). + +Each problem is routed to EXACTLY ONE view: A = problem + rubric findings, B = +query only (deployment form). Skill-gen trains only the final structured guidance turn. + +Base greedy rollouts and rubric diagnoses are cached to disk (jsonl, md5-keyed) so +restarts skip them; skill-gen is on-policy and never cached. + +8 GPUs default for high-memory cards: rank 0 trains the actor, rank 1 hosts a +frozen reference model, ranks 2-3 skill_sampler (synced), ranks 4-7 base_sampler +(frozen). Override TRAIN_GPUS / REF_GPUS / SKILL_SAMPLER_GPUS / BASE_SAMPLER_GPUS +for other layouts. +Leak / rubric use the backup teacher API: set LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / +LLM_BACKUP_MODEL. + +Launch: + LLM_BACKUP_API_KEY=... python cookbook/exp/embedding/train_reflexion_skill.py \ + --dataset aops --n 5000 --chunk-size 16 --lr 6e-6 +""" +import argparse +import copy +import hashlib +import json +import math +import os +import re +import sys +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Optional, Set, Tuple + +import numpy as np + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.data_format import SamplingParams +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.model import TransformersModel +from twinkle.patch.no_split_modules import NoSplitModulesPatch +from twinkle.processor import InputProcessor +from twinkle.sampler import vLLMSampler +from twinkle.template import Template +from twinkle_agentic.verifier import RubricVerifier +from twinkle_agentic.verifier.rubric_verifier import RubricItem + +logger = get_logger() + +try: + import swanlab +except ImportError: + swanlab = None + +MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') +GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.8)) +GEN_TEMPERATURE = float(os.environ.get('GEN_TEMPERATURE', 0.6)) +GEN_TOP_P = float(os.environ.get('GEN_TOP_P', 0.95)) +AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') +MATH_DATASET_ID = os.environ.get('MATH_DATASET_ID', 'modelscope/competition_math') + +# GPU layout: train actor (synced to skill_sampler) + frozen ref model + skill_sampler + base_sampler. +# High-memory cards can keep Qwen3-4B actor/ref as one GPU each and spend more GPUs +# on vLLM data-parallel sampling. The base side is heavier here because every +# clean skill candidate is re-solved by the frozen base model, plus balance/eval baselines. +TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 2)) +REF_GPUS = int(os.environ.get('REF_GPUS', 2)) +SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) +BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 2)) +NUM_GPUS = TRAIN_GPUS + REF_GPUS + SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS +TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', min(1, TRAIN_GPUS))) +REF_FSDP = int(os.environ.get('REF_FSDP', min(1, REF_GPUS))) +if min(TRAIN_GPUS, REF_GPUS, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS, TRAIN_FSDP, REF_FSDP) < 1: + raise ValueError('TRAIN_GPUS, REF_GPUS, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS, TRAIN_FSDP and REF_FSDP must all be >= 1') +if TRAIN_GPUS % TRAIN_FSDP != 0: + raise ValueError(f'TRAIN_GPUS ({TRAIN_GPUS}) must be divisible by TRAIN_FSDP ({TRAIN_FSDP})') +if REF_GPUS % REF_FSDP != 0: + raise ValueError(f'REF_GPUS ({REF_GPUS}) must be divisible by REF_FSDP ({REF_FSDP})') +TRAIN_DP = TRAIN_GPUS // TRAIN_FSDP +REF_DP = REF_GPUS // REF_FSDP + + +# =========================================================================== +# Block A -- boxed extraction + answer grading +# =========================================================================== +_BOXED_RE = re.compile(r'\\boxed\s*\{') + + +def extract_boxed(text: str) -> Optional[str]: + """Last ``\\boxed{...}`` content, brace-balanced.""" + if not text: + return None + last = None + for m in _BOXED_RE.finditer(text): + depth, i = 1, m.end() + while i < len(text) and depth > 0: + depth += (text[i] == '{') - (text[i] == '}') + i += 1 + if depth == 0: + last = text[m.end():i - 1].strip() + return last + + +def normalize_answer(ans: str) -> str: + if not ans: + return '' + s = str(ans).strip() + m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) + if m: + return m.group(1) + s = s.strip('$').strip().replace('−', '-') + 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').replace(r'\tfrac', r'\frac') + s = s.replace(r'\,', '').replace(r'\;', '').replace(r'\!', '') + s = re.sub(r'\\(?:quad|qquad|\s)', '', s) + s = re.sub(r'\s+', '', s) + s = re.sub(r'\\sqrt([0-9A-Za-z])', r'\\sqrt{\1}', s) + s = re.sub(r'\\frac\{([^{}]+)\}([^{}\\])', r'\\frac{\1}{\2}', s) + s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) + s = re.sub(r'\{[a-zA-Z]+\}$', '', s) + s = re.sub(r'\^\\circ|\^\{\\circ\}|°|\\circ', 'deg', s) + s = s.replace(r'\minus{}', '-').replace(r'\minus', '-') + + def _frac_to_slash(mt): + text = mt.group(0) + pos = text.index('{') + 1 + depth, num_start = 1, pos + while depth > 0: + depth += (text[pos] == '{') - (text[pos] == '}') + pos += 1 + numer = text[num_start:pos - 1] + pos += 1 + den_start, depth = pos, 1 + while depth > 0: + depth += (text[pos] == '{') - (text[pos] == '}') + pos += 1 + return f'({numer})/({text[den_start:pos - 1]})' + + s = re.sub(r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', _frac_to_slash, s) + s = re.sub(r'(? bool: + try: + va, vb = float(a.replace('(', '').replace(')', '')), 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) + return va is not None and vb is not None and abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) + + +_MCQ_REF_RE = re.compile( + r'^\\(?:textbf|text|mathrm|mathbf)\{\(?([A-E])\)?\s*\}\s*(.+)$' + r'|^\(?([A-E])\)\s+(.+)$') +_VAR_PREFIX_RE = re.compile(r'^(?:[A-Za-z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)$', re.DOTALL) + + +def _split_mcq(ans: str): + s = ans.strip() + m = _MCQ_REF_RE.match(s) + if m: + return (m.group(1) or m.group(3)), ((m.group(2) or m.group(4) or '').strip() or None) + bl = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) + return (bl.group(1), None) if bl else (None, s or None) + + +def _strip_var_prefix(ans: str) -> str: + m = _VAR_PREFIX_RE.match((ans or '').strip()) + return m.group(1).strip() if m else (ans or '') + + +def _normalize_tuple(ans: str) -> str: + return re.sub(r'[\s()\[\]{}\\]', '', ans or '') + + +def _math_verify_equal(predicted: str, reference: str) -> bool: + try: + from math_verify import parse, verify + gold = parse(r'\boxed{' + reference + '}', parsing_timeout=1) + pred = parse(r'\boxed{' + predicted + '}', parsing_timeout=1) + return bool(gold and pred and verify(gold, pred, timeout_seconds=1)) + except Exception: + return False + + +def answers_match(predicted: str, reference: str) -> bool: + if not predicted or not reference: + return False + norm_p, norm_r = normalize_answer(predicted), normalize_answer(reference) + if norm_p == norm_r or norm_p.lower() == norm_r.lower() or _try_numeric_equal(norm_p, norm_r): + return True + stripped_p = normalize_answer(_strip_var_prefix(predicted)) + stripped_r = normalize_answer(_strip_var_prefix(reference)) + if stripped_p and stripped_r: + if (stripped_p == stripped_r or stripped_p.lower() == stripped_r.lower() + or _try_numeric_equal(stripped_p, stripped_r)): + return True + p_letter, p_value = _split_mcq(predicted) + r_letter, r_value = _split_mcq(reference) + if p_letter and r_letter and p_letter == r_letter: + return True + if (p_letter and p_value is None) and (r_letter and r_value is None): + return False + 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 and (p_val == r_val or p_val.lower() == r_val.lower() or _try_numeric_equal(p_val, r_val)): + return True + for left, right in ((norm_p, norm_r), (stripped_p, stripped_r)): + tl, tr = _normalize_tuple(left), _normalize_tuple(right) + if ',' in tl and tl == tr: + return True + if '=' in norm_r and '=' not in norm_p: + if any(part and (part == norm_p or _try_numeric_equal(part, norm_p)) for part in norm_r.split('=')): + return True + if '=' in norm_p and '=' not in norm_r: + if any(part and (part == norm_r or _try_numeric_equal(part, norm_r)) for part in norm_p.split('=')): + return True + return _math_verify_equal(stripped_p or norm_p, stripped_r or norm_r) + + +# =========================================================================== +# Block B -- prompts, skill parsing, batched sampling +# =========================================================================== +DIRECT_SYSTEM = ( + 'You are an expert competition mathematician. Solve the following problem ' + 'step by step. Provide your final answer inside \\boxed{}.') + +_SKILL_SOLVE_PREFIX = ( + DIRECT_SYSTEM + '\n\n' + 'Before you start, keep these reminders in mind to avoid common mistakes on this ' + 'type of problem:\n') +_SKILL_SOLVE_SUFFIX = '\nApply them where relevant, but rely on your own reasoning to reach the answer.' + + +def build_direct_prompt(problem: str) -> Dict[str, Any]: + return {'messages': [{'role': 'system', 'content': DIRECT_SYSTEM}, + {'role': 'user', 'content': problem}]} + + +def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: + skill = (skill or '').strip() + if not skill: + return build_direct_prompt(problem) + # Concatenation (not .format): DIRECT_SYSTEM/skill contain literal braces. + return {'messages': [ + {'role': 'system', 'content': _SKILL_SOLVE_PREFIX + skill + _SKILL_SOLVE_SUFFIX}, + {'role': 'user', 'content': problem}]} + + +# -- skill-gen prompts (view A: problem + rubric source; view B: query only) -- +# Kept deliberately short: this is the RL policy's system prompt, so over-specifying +# the output hurts convergence. The concrete output format is appended separately by +# ``_SKILL_OUTPUT`` so the task text stays format-agnostic. +SKILL_GEN_SYSTEM = ( + 'You are a math guidance writer. A process-check on a related problem hints at ' + 'likely mistakes. Write short reusable guidance for this and similar problems, ' + 'and note what to watch out for.\n') + +SKILL_GEN_SYSTEM_Q = ( + 'You are a math guidance writer. Write short reusable guidance for this and ' + 'similar problems.\n') + +_SKILL_OUTPUT = ( + 'Output only:\n\nYour reusable solving guidance here.\n') + +SKILL_GEN_USER_Q = ( + 'Problem:\n{problem}\n\n') + +SKILL_GEN_USER_RUBRIC = ( + 'Target problem:\n{problem}\n\n' + 'Problem used for the process check:\n{rubric_problem}\n\n' + 'Process check:\n' + '{diagnosis}\n\n') + + +def _rubric_has_fail(diagnosis: str) -> bool: + """View A uses the open-book rubric prompt (and, under --viewa-sft, SFT distillation) + IFF its rubric localised a failure. No ``[FAIL]`` => nothing to inject: generation + degrades to query-only and the problem is trained by GRPO exactly like view B. Single + source of truth so ``_skillgen_messages`` and ``_group_records`` never diverge.""" + return '[FAIL]' in (diagnosis or '') + + +def _skillgen_messages(problem: str, view: str, diagnosis: str, + rubric_problem: str = '') -> List[Dict[str, Any]]: + """Single source of truth for the skill-gen prompt (used at BOTH generation and + training so they never diverge). View A with a localisable failure uses the target + problem plus the rubric source problem and findings; view B -- or a view-A problem + whose rubric flagged nothing (no ``[FAIL]``) -- degrades to the query-only prompt.""" + if view == 'B' or not _rubric_has_fail(diagnosis): + return [{'role': 'system', 'content': SKILL_GEN_SYSTEM_Q + _SKILL_OUTPUT}, + {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=problem)}] + rubric_problem = rubric_problem or problem + return [{'role': 'system', 'content': SKILL_GEN_SYSTEM + _SKILL_OUTPUT}, + {'role': 'user', 'content': SKILL_GEN_USER_RUBRIC.format( + problem=problem, rubric_problem=rubric_problem, diagnosis=diagnosis)}] + + +def _assign_view(problem: str, args: argparse.Namespace) -> str: + h = int(hashlib.md5(f'{args.seed}:{problem}'.encode('utf-8')).hexdigest(), 16) + return 'B' if (h % 100000) / 100000.0 < args.view_b_frac else 'A' + + +def _view_prompt(r: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: + return {'messages': _skillgen_messages( + r['problem'], r['_view'], r.get('_rubric_diag', ''), r.get('_rubric_src', ''))} + + +_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') + + +def _clean_text(decoded: Optional[str]) -> str: + return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() + + +def _extract_tag_block(answer: str, tag: str, allow_empty: bool = False) -> Optional[str]: + low = answer.lower() + open_tag, close_tag = f'<{tag}>', f'' + s = low.rfind(open_tag) + if s < 0: + return None + inner = s + len(open_tag) + e = low.find(close_tag, inner) + if e < 0: + return None + block = answer[inner:e].strip() + block = re.sub(r'', '', block, flags=re.IGNORECASE).strip() + return block if (block or allow_empty) else None + + +def _extract_skill(text: str) -> Optional[str]: + """Parse skill-generation output: return the inner text of a non-empty ```` + block, or None. If a ```` marker is present, parse only the text after the + last one; otherwise parse the full response.""" + low = text.lower() + end_think = low.rfind('') + answer = text[end_think + len(''):] if end_think >= 0 else text + return _extract_tag_block(answer, 'skills') + + +def _parse_seq(seq, gold: str) -> Dict[str, Any]: + """Grade one sampled sequence into a rollout record.""" + text = _clean_text(getattr(seq, 'decoded', '') or '') + pred = extract_boxed(text) + correct = bool(pred) and answers_match(pred, gold) + terminated = getattr(seq, 'stop_reason', None) != 'length' + return {'pred': pred, 'correct': correct, 'terminated': terminated, + 'passed': bool(correct and terminated), + 'stop_reason': getattr(seq, 'stop_reason', None), + 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} + + +def _run_samples(sampler, prompts: List[Any], num_samples: int, max_tokens: int, + gen_dp: int, temperature: Optional[float] = None, + top_p: Optional[float] = None, top_k: Optional[int] = None) -> List[List[Any]]: + """One batched sampler call -> per-prompt list of raw sequences. vLLM dp needs + batch len >= dp, so pad the tail and slice back.""" + if not prompts: + return [] + params = SamplingParams( + max_tokens=max_tokens, + temperature=GEN_TEMPERATURE if temperature is None else temperature, + top_p=GEN_TOP_P if top_p is None else top_p, + num_samples=num_samples, **({} if top_k is None else {'top_k': top_k})) + padded = prompts + if gen_dp > 1 and 0 < len(prompts) < gen_dp: + padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] + responses = sampler.sample(padded, params)[:len(prompts)] + return [list(r.sequences) if (r and r.sequences) else [] for r in responses] + + +# =========================================================================== +# Block C -- data loading via twinkle.Dataset + numeric filtering +# =========================================================================== +def _boxed_batch(rows: Dict[str, List[Any]], dataset: str) -> Dict[str, List[Any]]: + """Batched (columnar) mapper for ``Dataset.map``: annotate each row with its boxed + ``reference_answer`` and a ``_keep`` flag (aops also needs the ``boxed`` metadata).""" + sols = rows['solution'] + metas = rows.get('metadata', [None] * len(sols)) + refs = [extract_boxed(s or '') for s in sols] + keep = [bool(ref) and (dataset != 'aops' or bool((meta or {}).get('boxed'))) + for ref, meta in zip(refs, metas)] + return {**rows, 'reference_answer': refs, '_keep': keep} + + +def load_problems(dataset: str, n: int, seed: int, num_proc: int = 0) -> List[Dict[str, Any]]: + """Load boxed-answer problems as ``{problem, reference_answer, level?}`` via + twinkle.Dataset (ModelScope hub), sampled to ``n`` (0 = all). Boxed extraction (regex + + brace scan over every solution) is parallelised by ``Dataset.map``/``.filter``; + ``num_proc`` defaults to all cores (set 1 to force serial).""" + ds_id = AOPS_DATASET_ID if dataset == 'aops' else MATH_DATASET_ID + ds = Dataset(DatasetMeta(dataset_id=f'ms://{ds_id}', split='train')) + nproc = num_proc if num_proc > 0 else min(32, os.cpu_count() or 1) + ds.map(lambda rows: _boxed_batch(rows, dataset), num_proc=nproc) + ds.filter(lambda row: row['_keep'], num_proc=nproc) + has_level = 'level' in ds.dataset.column_names + out = [{'data_id': f'{dataset}:{i}', 'problem': row['problem'], + 'reference_answer': row['reference_answer'], + **({'level': row['level']} if has_level and row.get('level') else {})} + for i, row in enumerate(ds.dataset)] + logger.info(f'[data] {dataset}: {len(out)} boxed problems ({nproc} procs)') + rng = np.random.RandomState(seed) + rng.shuffle(out) + return out[:n] if (n and n < len(out)) else out + + +# --------------------------------------------------------------------------- +# Bag-of-words neighbour pairing (cross-problem rubric transfer, --xproblem-rubric) +# --------------------------------------------------------------------------- +# Common English + math-scaffolding words that carry no problem-type signal. Kept small +# and deterministic on purpose (no external stopword list): what survives is the domain +# vocabulary ('triangle', 'prime', 'polynomial', ...) that actually defines the type. +_BOW_STOP = frozenset(""" +a an the of to in on at for and or but if is are be was were been being this that these those +with without into onto from by as it its their his her our your my we you they he she them +find compute determine calculate evaluate solve show prove given let suppose consider assume +what which when where how many much value values number numbers expression form terms term +such that then than so if only when each every all any some both one two three four five six +seven eight nine ten first second third last non over under about above below between +problem answer result equal equals sum difference product total following there here have has +had do does did can could will would should may might must not no yes if then else +""".split()) + +_WORD_RE = re.compile(r'[a-z]+') + + +def _stem(w: str) -> str: + """Crude suffix stripper so 'prime'/'primes', 'triangle'/'triangles' collapse to one + type token. Not linguistically correct -- just enough to merge the common plural/gerund + variants that otherwise split a type's vocabulary and starve the df filter.""" + if len(w) > 4 and w.endswith('ies'): + return w[:-3] + 'y' # properties -> property + if len(w) > 4 and w.endswith('es') and w[-3] in 'sxzh': + return w[:-2] # boxes -> box (keep primes -> prime below) + for suf in ('ing', 'ed', 's'): + if len(w) > len(suf) + 2 and w.endswith(suf): + return w[:-len(suf)] + return w + + +def _tokenize(problem: str) -> List[str]: + """Deterministic bag-of-words tokens for type matching: lowercase alphabetic words + (numbers dropped -- they are instance detail, not type), minus generic stopwords, then + stemmed, so only domain terms remain. Latex control words (frac, sqrt, ...) survive.""" + return [_stem(w) for w in _WORD_RE.findall((problem or '').lower()) + if len(w) > 2 and w not in _BOW_STOP] + + +class BagOfWordsIndex: + """TF-IDF cosine nearest-neighbour over problem statements. Sparse dict vectors + + an inverted index, so ``nearest`` scores only problems sharing a term (near-linear in + practice, no dense NxN). Deterministic; ties break on lower index for reproducibility. + + Answer-aware: ``nearest`` skips the same index, near-verbatim duplicates (cosine + >= ``sim_max``) and -- crucially for anti-leak -- any candidate whose answer equals the + query's, so a neighbour rubric can never hand over the query's own answer.""" + + def __init__(self, problems: List[str], answers: Optional[List[str]] = None, + min_df: int = 2, max_df_frac: float = 0.5): + self._toks = [_tokenize(p) for p in problems] + self._ans = [(_numeric_value(a) or (str(a).strip() if a else '')) for a in answers] \ + if answers is not None else [''] * len(problems) + n = len(self._toks) + df: Dict[str, int] = {} + for toks in self._toks: + for w in set(toks): + df[w] = df.get(w, 0) + 1 + max_df = max(min_df, int(max_df_frac * n)) + self._idf = {w: math.log((n + 1) / (c + 1)) + 1.0 + for w, c in df.items() if min_df <= c <= max_df} + self._vecs: List[Dict[str, float]] = [self._vectorize(t) for t in self._toks] + self._inverted: Dict[str, List[int]] = {} + for i, v in enumerate(self._vecs): + for w in v: + self._inverted.setdefault(w, []).append(i) + + def _vectorize(self, toks: List[str]) -> Dict[str, float]: + tf: Dict[str, float] = {} + for w in toks: + if w in self._idf: + tf[w] = tf.get(w, 0.0) + 1.0 + vec = {w: c * self._idf[w] for w, c in tf.items()} + norm = math.sqrt(sum(x * x for x in vec.values())) + return {w: x / norm for w, x in vec.items()} if norm > 0 else {} + + def nearest(self, i: int, sim_max: float = 0.98) -> Tuple[int, float]: + """(index, cosine) of the most similar *distinct* problem: not i, not a duplicate + (cosine < ``sim_max``) and not the same answer. (-1, 0.0) if none qualifies.""" + vi = self._vecs[i] + if not vi: + return -1, 0.0 + ai = self._ans[i] + scores: Dict[int, float] = {} + for w, xi in vi.items(): + for j in self._inverted.get(w, ()): + if j != i: + scores[j] = scores.get(j, 0.0) + xi * self._vecs[j].get(w, 0.0) + best_j, best_s = -1, 0.0 + for j, s in scores.items(): + if s >= sim_max or (ai and self._ans[j] == ai): + continue + if s > best_s or (s == best_s and (best_j < 0 or j < best_j)): + best_j, best_s = j, s + return best_j, best_s + + +def build_pairs(records: List[Dict[str, Any]], n: int, seed: int, sim_max: float = 0.98 + ) -> Tuple[List[Dict[str, Any]], Dict[str, Tuple[str, float]]]: + """Single-pass cross-problem pairing over the whole pool (one index build). + + Returns ``(subset, neighbour_map)`` where ``subset`` is the ``n`` problems with the + strongest qualifying neighbour (dense same-type pairs; n<=0/n>=len keeps all, shuffled) + and ``neighbour_map[q] = (p, cosine)`` gives each kept problem its analogue P. P is drawn + from the *full* pool (richer analogues) but never a duplicate or same-answer problem, so + P's rubric can transfer method without ever leaking Q's answer.""" + index = BagOfWordsIndex([r['problem'] for r in records], + [str(r.get('reference_answer', '')) for r in records]) + nbr = [index.nearest(i, sim_max) for i in range(len(records))] + order = sorted(range(len(records)), key=lambda i: (-nbr[i][1], i)) + keep = order[:n] if (0 < n < len(records)) else list(range(len(records))) + rng = np.random.RandomState(seed) + rng.shuffle(keep) + subset = [records[i] for i in keep] + neighbour_map = {records[i]['problem']: (records[nbr[i][0]]['problem'], nbr[i][1]) + for i in keep if nbr[i][0] >= 0} + return subset, neighbour_map + + +_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') + + +def _norm_num_text(num: str) -> str: + try: + f = float(num) + return str(int(f)) if f == int(f) else str(f) + except Exception: + return str(num).strip() + + +def _numeric_value(raw: Any) -> Optional[str]: + """Collapse an answer to a single int/decimal/fraction, or None.""" + if raw is None: + return None + s = str(raw).strip().strip('$').strip() + s = s.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac') + s = re.sub(r'\\!|\\,|\\;|\\ |\\left|\\right|\s', '', s) + for pat in (r'\\frac\{(-?\d+)\}\{(-?\d+)\}', r'(-?\d+)/(-?\d+)'): + m = re.fullmatch(pat, s) + if m: + a, b = int(m.group(1)), int(m.group(2)) + return _norm_num_text(str(a / b)) if b else None + return _norm_num_text(s) if _NUM_RE.fullmatch(s) else None + + +def _answer_leaked(skill: str, reference: str) -> bool: + """Audit whether a generated skill contains the final answer verbatim. This is NOT + a training filter: if the skill model derives an answer from the problem, that is a + legitimate answer-bearing skill under this experiment. The real leakage boundary is the + external API/rubric diagnosis, which is constrained by prompt to stay answer-free.""" + if not skill: + return False + for cand in {_numeric_value(reference), (str(reference).strip() or None)}: + if cand and re.search(r'(? Tuple[Set[str], Set[str]]: + """Read jsonl files and collect data_id/problem keys that must be excluded. + + The cold-start SFT builder writes stable ``data_id`` values. ``problem`` is kept as a + backward-compatible fallback for older jsonl files produced before data_id existed.""" + ids: Set[str] = set() + problems: Set[str] = set() + for raw_path in (paths_arg or '').split(','): + path = raw_path.strip() + if not path or not os.path.exists(path): + continue + with open(path, encoding='utf-8') as f: + for line in f: + if not line.strip(): + continue + row = json.loads(line) + if row.get('record_type') in {'config', 'summary'}: + continue + data_id = str(row.get('data_id') or '').strip() + problem = str(row.get('problem') or '').strip() + if data_id: + ids.add(data_id) + elif problem: + problems.add(problem) + return ids, problems + + +def _load_seam_parquet(path: str) -> List[Dict[str, Any]]: + """Read a SEAM ``build_aops_dataset.py`` parquet (VERL RLHF schema) into twinkle records, + PRESERVING file row order. ``problem <- extra_info.problem`` and + ``reference_answer <- reward_model.ground_truth``. No shuffle/filter: + the parquet is already SEAM's numeric-filtered, seed-42-shuffled, truncated split.""" + import pyarrow.parquet as pq + rows = pq.read_table(path).to_pylist() + out: List[Dict[str, Any]] = [] + for i, r in enumerate(rows): + ei = r.get('extra_info') or {} + rm = r.get('reward_model') or {} + problem = (ei.get('problem') or '').strip() + ref = rm.get('ground_truth') + if not problem or ref is None: + continue + out.append({'problem': problem, 'reference_answer': str(ref), + 'data_id': f"seam:{ei.get('split', '')}:{ei.get('index', i)}"}) + return out + + +def _load_records_from_seam(args: argparse.Namespace, seam_dir: str + ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], + Dict[str, Tuple[str, float]], Dict[str, str], Dict[str, int]]: + """Data entry that mirrors a SEAM run EXACTLY: read ``train.parquet``/``val.parquet`` from + ``seam_dir`` in file order, use ``val`` as the eval holdout, take the first ``--n`` train rows + (post ``--pool-offset``) with NO shuffle. ``--numeric-only``/``--eval-size``/internal shuffle + are bypassed (the parquet is already the authoritative split).""" + tp, vp = os.path.join(seam_dir, 'train.parquet'), os.path.join(seam_dir, 'val.parquet') + if not (os.path.exists(tp) and os.path.exists(vp)): + raise FileNotFoundError( + f'--seam-parquet-dir needs both train.parquet and val.parquet in {seam_dir}') + if args.xproblem_rubric: + raise ValueError('--xproblem-rubric is unsupported with --seam-parquet-dir ' + '(SEAM parquet carries no neighbour structure).') + pool = _load_seam_parquet(tp) # already SEAM-shuffled + truncated, in file order + eval_records = [dict(r) for r in _load_seam_parquet(vp)] # SEAM's exact val holdout + pool_offset = max(0, int(getattr(args, 'pool_offset', 0) or 0)) + if pool_offset: + if pool_offset >= len(pool): + raise ValueError(f'--pool-offset ({pool_offset}) leaves no train records ' + f'from SEAM train pool size {len(pool)}') + pool = pool[pool_offset:] + train_n = args.n if args.n > 0 else len(pool) + train_records = [dict(r) for r in pool[:train_n]] + if {r['problem'] for r in train_records} & {r['problem'] for r in eval_records}: + raise ValueError('eval/train overlap detected in SEAM parquet') + stats = {'raw_loaded': len(pool) + len(eval_records), 'numeric_dropped': 0, + 'excluded_records': 0, 'pool_offset': pool_offset, + 'train_records': len(train_records), 'eval_records': len(eval_records), + 'source': 'seam_parquet', 'seam_parquet_dir': seam_dir} + return train_records, eval_records, {}, {}, stats + + +def _load_records(args: argparse.Namespace + ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], + Dict[str, Tuple[str, float]], Dict[str, str], Dict[str, int]]: + """Load, numeric-filter, shuffle, split a fixed eval holdout, and (when --xproblem-rubric) + select a same-type-dense train subset with its neighbour map -- all in one pass. + Also returns ``pool_answers`` (every candidate neighbour P's true answer) so P's baseline + can be graded/cached correctly even when P is not itself a training problem.""" + seam_dir = (getattr(args, 'seam_parquet_dir', '') or '').strip() + if seam_dir: # read SEAM parquet in file order, bypassing load/filter/shuffle/split + return _load_records_from_seam(args, seam_dir) + # Load all when filtering or splitting (else the eval holdout could starve train). + load_n = 0 if (args.numeric_only or args.eval_size > 0) else args.n + records = load_problems(args.dataset, load_n, args.seed) + raw_n, dropped = len(records), 0 + if args.numeric_only: + records = [{**r, 'reference_answer': v} + for r, v in ((r, _numeric_value(r.get('reference_answer'))) for r in records) + if v is not None] + dropped = raw_n - len(records) + np.random.RandomState(args.seed).shuffle(records) + exclude_ids, exclude_problems = _load_excluded_records(getattr(args, 'exclude_data_ids', '')) + excluded = 0 + if exclude_ids or exclude_problems: + before = len(records) + records = [r for r in records + if str(r.get('data_id', '')) not in exclude_ids + and str(r.get('problem', '')).strip() not in exclude_problems] + excluded = before - len(records) + eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 + eval_records = [dict(r) for r in records[:eval_n]] + pool = records[eval_n:] + pool_offset = max(0, int(getattr(args, 'pool_offset', 0) or 0)) + if pool_offset: + if pool_offset >= len(pool): + raise ValueError(f'--pool-offset ({pool_offset}) leaves no train records from pool size {len(pool)}') + pool = pool[pool_offset:] + train_n = args.n if args.n > 0 else len(pool) + # Cross-problem rubric: pick the train_n problems with the strongest qualifying neighbour + # (dense same-type pairs) and build {Q -> (P, sim)} in one index pass. Otherwise keep the + # first train_n (already shuffled) with no neighbours. + if args.xproblem_rubric: + subset, neighbor_map = build_pairs(pool, train_n, args.seed) + train_records = [dict(r) for r in subset] + pool_answers = {r['problem']: str(r.get('reference_answer', '')) for r in pool} + else: + train_records, neighbor_map, pool_answers = [dict(r) for r in pool[:train_n]], {}, {} + if {r['problem'] for r in train_records} & {r['problem'] for r in eval_records}: + raise ValueError('eval/train overlap detected') + stats = {'raw_loaded': raw_n, 'numeric_dropped': dropped, + 'excluded_records': excluded, 'pool_offset': pool_offset, + 'train_records': len(train_records), 'eval_records': len(eval_records)} + return train_records, eval_records, neighbor_map, pool_answers, stats + + +# =========================================================================== +# Block D -- disk cache, problem pool, baseline rollout, rubric check +# =========================================================================== +class DiskCache: + """Append-only jsonl kv cache (md5 key -> value). Loads on init, appends on put. + Disabled instances always miss and never write.""" + + def __init__(self, path: str, enabled: bool = True): + self._mem: Dict[str, Any] = {} + self._fh = None + self._lock = threading.Lock() # base baseline is prefetched on a background thread + if not enabled: + return + if os.path.exists(path): + with open(path, encoding='utf-8') as f: + for line in f: + if line.strip(): + row = json.loads(line) + self._mem[row['key']] = row['value'] + self._fh = open(path, 'a', encoding='utf-8') + + @staticmethod + def key_for(*parts: str) -> str: + return hashlib.md5('\x1f'.join(parts).encode('utf-8')).hexdigest() + + def __contains__(self, key: str) -> bool: + with self._lock: + return key in self._mem + + def get(self, key: str) -> Any: + with self._lock: + return self._mem.get(key) + + def put(self, key: str, value: Any) -> None: + with self._lock: + self._mem[key] = value + if self._fh is not None: + self._fh.write(json.dumps({'key': key, 'value': value}, ensure_ascii=False) + '\n') + self._fh.flush() + + def close(self) -> None: + if self._fh is not None: + self._fh.close() + + +class _LockedSampler: + """Serialises every ``.sample()`` on a sampler behind one lock. The base sampler is + shared by the main thread (with-skill scoring, eval) and the baseline-prefetch thread; + ``sample`` is a slice_dp/flatten remote call whose collect is NOT safe to interleave + across two callers, so concurrent calls could mis-join sequences. The lock keeps base + calls serial (prefetch still overlaps the skill-gen phase, which uses skill_sampler).""" + + def __init__(self, sampler): + self._sampler = sampler + self._lock = threading.Lock() + + def sample(self, *args, **kwargs): + with self._lock: + return self._sampler.sample(*args, **kwargs) + + def __getattr__(self, name): + return getattr(self._sampler, name) + + +class ProblemPool: + """Cyclic draw source. Each full pass reshuffles with ``seed + epoch``; the initial + pass keeps the loader order. ``draw(k)`` returns k distinct problems (k << pool).""" + + def __init__(self, records: List[Dict[str, Any]], seed: int): + self._records = list(records) + self._seed, self._cursor, self.epoch = seed, 0, 0 + + def draw(self, k: int) -> List[Dict[str, Any]]: + out, seen = [], set() + while len(out) < k: + if self._cursor >= len(self._records): + self.epoch += 1 + np.random.RandomState(self._seed + self.epoch).shuffle(self._records) + self._cursor = 0 + r = self._records[self._cursor] + self._cursor += 1 + if id(r) not in seen: + seen.add(id(r)) + out.append(r) + return out + + def peek(self, k: int) -> List[Dict[str, Any]]: + """The next k distinct problems draw() would return, WITHOUT advancing state + (no cursor move, no reshuffle). Used to prefetch their baseline into base_cache + while the current chunk trains -- a wrong guess (cross-epoch reshuffle) only + misses the cache, never corrupts the draw.""" + out, seen, cur = [], set(), self._cursor + recs = self._records + while len(out) < k and cur < len(recs): # stop at epoch edge; don't simulate reshuffle + r = recs[cur] + cur += 1 + if id(r) not in seen: + seen.add(id(r)) + out.append(r) + return out + + +def _empty_roll() -> Dict[str, Any]: + return {'pred': '', 'correct': False, 'terminated': False, 'passed': False, + 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} + + +def _apply_baseline(r: Dict[str, Any], roll: Dict[str, Any]) -> None: + """Attach a greedy baseline roll and reset per-chunk working state.""" + r['_baseline_rolls'], r['_cands'], r['_init'] = [roll], [], [roll] + r['_failed'] = not roll['correct'] + r['_baseline_pass'] = 1.0 if roll['correct'] else 0.0 + r['_hard'] = True # process every problem; group variance selects (SEAM-style) + + +def baseline_rollout(base_sampler, problems: List[Dict[str, Any]], base_dp: int, + args: argparse.Namespace, cache: DiskCache) -> int: + """Base solves each problem greedily once (T=0, M=1), disk-cached by problem text. + The base is frozen + greedy so a cache hit is exact. Returns fresh (miss) count.""" + todo = [r for r in problems if DiskCache.key_for(r['problem']) not in cache] + if todo: + out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], + 1, args.max_tokens, base_dp, temperature=0.0) + for r, seqs in zip(todo, out): + roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + cache.put(DiskCache.key_for(r['problem']), roll) + for r in problems: + _apply_baseline(r, cache.get(DiskCache.key_for(r['problem']))) + return len(todo) + + +# -- rubric process-check (view A): teacher diagnoses the base's attempt -- +_RFT_DIAG_SYSTEM = """\ +You are a strategy-level process checker for a math solution attempt. You are given a +math problem, a rubric, and one attempted solution segment. Decide PASS or FAIL for each +criterion, and write the diagnosis so it can become useful reusable guidance for solving +similar problems without seeing this segment. + +Output STRICT JSON (no prose outside it) with this shape: +{ + "items": [ + {"index": 1, "verdict": "PASS", "reason": "", + "fix": ""}, + {"index": 2, "verdict": "FAIL", "reason": "", + "fix": ""} + ], + "overall": "OK" | "ISSUES", + "summary": "" +} + +Rules: +- Judge every criterion independently; a [Hard Rule] is FAIL unless clearly satisfied. +- Judge ONLY what is observable in THIS segment. Ignore hidden or + content for output-format criteria. +- The API diagnosis is an external teacher signal, so it must stay answer-free. +- Prefer diagnosis that transfers to view-B skill generation: name the route choice, + structural observation, missing check, or length-control habit that a solver should + remember before solving a similar problem. +- For PASS items, leave "fix" as "". +- For FAIL items, describe the process problem at strategy level: unsuitable method, + missed structure, invalid transformation, missing constraint check, redundant cases, + off-track approach, contradiction, or inefficient/unfinished reasoning. +- A fix may suggest the LOCAL correction direction, such as identify the key structure, + verify constraints, preserve equivalence, reduce redundant cases, or choose a more + direct route. Do not carry out the correction. +- Never reveal the final answer, a corrected value/expression, an option label, or a + step-by-step solution that would let another model copy the solve. +- If the segment contains a process note saying it was cut off before a final boxed + answer, mark the length-budget criterion as FAIL and suggest a method-level way to + finish faster. +- Keep every "reason" and "fix" concise: one short sentence each. +- "overall" is "OK" only if NO criterion is FAIL. +- Output only the JSON object.""" + +_RFT_DIAG_USER = """\ +## Task / query (context) +{query} + +## Rubric +{rubric} + +## Segment +{segment} + +Now output the diagnostic JSON object.""" + +_MATH_RUBRIC = [ + ('The attempt chooses a method suitable for the problem structure', False), + ('The attempt identifies the key constraint, invariant, or quantity before computing', False), + ('Algebraic and logical transformations preserve validity at each step', True), + ('The attempt checks required constraints, domains, boundary cases, or validity conditions', False), + ('The attempt avoids redundant casework, looping, or re-deriving known facts', False), + ('The attempt reaches a final boxed answer within the length budget', False), + ('The approach stays focused on the actual question asked', False), +] + +# Bump when the rubric criteria or the diagnosis prompt change so the disk-cached +# diagnoses written under an older rubric are not silently reused. +_RUBRIC_VERSION = 'rubric_v5_viewb_strategy' + + +class _RftRubricVerifier(RubricVerifier): + def _diagnose_trajectory(self, query: str, rubric_block: str, segment_text: str) -> dict: + return {'messages': [ + {'role': 'system', 'content': _RFT_DIAG_SYSTEM}, + {'role': 'user', 'content': _RFT_DIAG_USER.format( + query=query, rubric=rubric_block, segment=segment_text)}]} + + +def build_rubric_checker() -> Optional[RubricVerifier]: + """Fixed math-process rubric verifier (teacher-served). None if no LLM backup env.""" + if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') + or os.environ.get('OPENAI_API_KEY')): + return None + return _RftRubricVerifier( + fixed_rubric=[RubricItem(t, is_hard=h) for t, h in _MATH_RUBRIC], gate=True) + + +def _format_diagnosis(detail) -> str: + """One line per criterion (PASS/FAIL + reason + fix) then a summary.""" + rub = detail.rubric + lines = [] + for it in detail.items: + text = rub[it.index - 1].text if 0 < it.index <= len(rub) else f'criterion {it.index}' + if it.verdict: + lines.append(f'- [PASS] {text}') + else: + tail = f': {it.reason}' if it.reason else '' + tail += f' (fix: {it.fix})' if it.fix else '' + lines.append(f'- [FAIL] {text}{tail}') + if detail.summary: + lines.append(f'Summary: {detail.summary}') + return '\n'.join(lines) + + +def diagnose_views(checker, problems: List[Dict[str, Any]], args: argparse.Namespace, + cache: DiskCache) -> None: + """Rubric-check every view-A problem's greedy attempt in parallel (disk-cached by + problem + attempt), stashing the formatted findings on ``r['_rubric_diag']``.""" + targets = [r for r in problems if r.get('_view') == 'A'] + if not checker or not targets: + return + + def _key(r: Dict[str, Any]) -> str: + init = r.get('_init', [{}])[0] + term = 'L' if (init.get('stop_reason') == 'length' or not init.get('terminated')) else 'T' + return DiskCache.key_for(r['problem'], init.get('text', ''), _RUBRIC_VERSION, term) + + pending = [] + for r in targets: + key = _key(r) + if key in cache: + r['_rubric_diag'] = cache.get(key) + else: + pending.append((r, key)) + if not pending: + return + + def _run(item): + r, key = item + init = r['_init'][0] + seg_text = init['text'] + if init.get('stop_reason') == 'length' or not init.get('terminated'): + seg_text += ('\n\n[Process note: this attempt was cut off at the token budget ' + 'and never produced a final \\boxed{} answer.]') + seg = {'messages': [{'role': 'user', 'content': r['problem']}, + {'role': 'assistant', 'content': seg_text}]} + attempts = max(1, args.rubric_retries + 1) + for attempt in range(attempts): + try: + return r, key, _format_diagnosis(checker.diagnose(seg, query=r['problem'])) + except Exception as exc: # teacher hiccup -> retry, then no-diagnosis prompt (not cached) + if attempt + 1 < attempts: + logger.warning(f'[rubric] diagnose error: {exc}; retry {attempt + 1}/{args.rubric_retries}') + time.sleep(min(2.0, 0.5 * (2 ** attempt))) + continue + logger.warning(f'[rubric] diagnose error: {exc}; giving up after {attempts} attempts') + return r, key, None + + workers = max(1, min(args.rubric_workers, len(pending))) + with ThreadPoolExecutor(max_workers=workers) as ex: + for r, key, diag in ex.map(_run, pending): + r['_rubric_diag'] = diag or '' + if diag is not None: + cache.put(key, diag) + + +# =========================================================================== +# Block E -- chunk draw, generation pipeline, record building +# =========================================================================== +def _baseline_class(r: Dict[str, Any]) -> str: + """success | fail_loop (out of length / never terminated) | fail_wrong.""" + roll = r['_init'][0] + if roll['correct']: + return 'success' + return 'fail_loop' if (roll['stop_reason'] == 'length' or not roll['terminated']) else 'fail_wrong' + + +def _select_balanced(buckets: Dict[str, List[Dict[str, Any]]], n_success: int, + n_fail: int, n_fail_loop: int) -> List[Dict[str, Any]]: + """Pick n_fail base-fails (toward n_fail_loop loop-fails, best-effort) + n_success + base-successes; top up any shortfall from leftovers.""" + loop, wrong, succ = buckets['fail_loop'], buckets['fail_wrong'], buckets['success'] + take_loop = min(n_fail_loop, len(loop)) + take_wrong = min(n_fail - take_loop, len(wrong)) + take_loop = min(n_fail - take_wrong, len(loop)) + sel = loop[:take_loop] + wrong[:take_wrong] + succ[:n_success] + target = n_success + n_fail + if len(sel) < target: + used = {id(x) for x in sel} + sel += [x for b in (loop, wrong, succ) for x in b if id(x) not in used][:target - len(sel)] + return sel + + +def draw_chunk(pool: ProblemPool, base_sampler, base_dp: int, args: argparse.Namespace, + cache: DiskCache) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + """Draw one chunk, baselining every drawn problem. With ``--balance``, keep + drawing+baselining until the target base fail:success mix is reachable (or the budget + is hit), then select a balanced subset.""" + if not args.balance: + chunk = pool.draw(args.chunk_size) + fresh = baseline_rollout(base_sampler, chunk, base_dp, args, cache) + return chunk, {'enabled': False, 'n_drawn': len(chunk), 'n_baseline_fresh': fresh} + + n_success = max(0, min(args.chunk_size, round(args.chunk_size * args.balance_success_frac))) + n_fail = args.chunk_size - n_success + n_fail_loop = round(n_fail * args.balance_loop_frac) + buckets: Dict[str, List[Dict[str, Any]]] = {'success': [], 'fail_loop': [], 'fail_wrong': []} + budget, n_drawn, n_fresh, seen = args.chunk_size * args.balance_max_draws_mult, 0, 0, set() + while n_drawn < budget: + if (len(buckets['success']) >= n_success + and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail): + break + batch = pool.draw(args.chunk_size) + n_fresh += baseline_rollout(base_sampler, batch, base_dp, args, cache) + n_drawn += len(batch) + for r in batch: + if id(r) not in seen: + seen.add(id(r)) + buckets[_baseline_class(r)].append(r) + + reached = (len(buckets['success']) >= n_success + and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail) + chunk = _select_balanced(buckets, n_success, n_fail, n_fail_loop) + sel_success = sum(1 for r in chunk if not r['_failed']) + stats = { + 'enabled': True, 'n_drawn': n_drawn, 'n_baseline_fresh': n_fresh, 'n_selected': len(chunk), + 'target_success': n_success, 'target_fail': n_fail, 'target_fail_loop': n_fail_loop, + 'selected_success': sel_success, 'selected_fail': len(chunk) - sel_success, + 'selected_fail_loop': sum(1 for r in chunk if _baseline_class(r) == 'fail_loop'), + 'selected_fail_wrong': sum(1 for r in chunk if _baseline_class(r) == 'fail_wrong'), + 'selected_success_frac': (sel_success / len(chunk)) if chunk else 0.0, + 'budget_hit': not reached, + } + return chunk, stats + + +def _assign_advantages(hard: List[Dict[str, Any]], args: argparse.Namespace) -> None: + """Group-relative advantage over each problem's scored candidates using the greedy + binary reward R in {0,1}: ``A = (R - mean) / (std + eps)``. std==0 groups get no + gradient -- GRPO's variance selects informative problems (no explicit difficulty gate). + A symmetric clip limits both positive and negative outliers from sparse 1-vs-many groups.""" + eps = 1e-6 + adv_clip = abs(float(getattr(args, 'adv_clip', 0.0) or 0.0)) + for r in hard: + for c in r['_cands']: + c['advantage'], c['grpo_adv'], c['kept'] = 0.0, 0.0, False + cs = [c for c in r['_cands'] if c.get('reward') is not None] + if len(cs) < 2: + continue + rewards = [c['reward'] for c in cs] + mean_r = sum(rewards) / len(rewards) + std = (sum((x - mean_r) ** 2 for x in rewards) / len(rewards)) ** 0.5 + if std < 1e-9: + continue + for c in cs: + raw_adv = (c['reward'] - mean_r) / (std + eps) + adv = max(-adv_clip, min(adv_clip, raw_adv)) if adv_clip > 0 else raw_adv + c['advantage'], c['grpo_adv'], c['kept'] = adv, adv, c['reward'] > mean_r + + +def _best_sft_candidate(r: Dict[str, Any], args: argparse.Namespace) -> Optional[Dict[str, Any]]: + """Pick ONE view-A candidate to distill (online context distillation). PREFER the + executor-verified PASSING skills (reward==1); if NONE passed -- common on the hard + problems that are exactly the cases worth distilling -- FALL BACK to any parseable + open-book skill regardless of the executor outcome. Answer-bearing skills produced by + the skill model itself are allowed here; only the external API/rubric diagnosis must be + answer-free. Within the chosen tier, take the one whose skill length is CLOSEST to + ``--sft-target-len`` -- an empirically high-pass-rate length (~500-600 chars in this + run) -- breaking ties by the fewest executor solve tokens. Targeting a length (rather + than the minimum) avoids a distillation feedback loop that would otherwise drive + rollouts ever shorter. None only when no parseable candidate exists at all.""" + eligible = [c for c in r['_cands'] if c.get('parseable') and c.get('skills')] + if not eligible: + return None + passing = [c for c in eligible if c.get('reward') == 1.0] + cs = passing or eligible + target = int(getattr(args, 'sft_target_len', 550) or 550) + + def _solve_tokens(c: Dict[str, Any]) -> int: + rolls = c.get('rolls') or [] + return rolls[0].get('gen_tokens', 1 << 30) if rolls else (1 << 30) + + return min(cs, key=lambda c: (abs(len(c['skills']) - target), _solve_tokens(c))) + + +def apply_neighbor_rubric(base_sampler, chunk: List[Dict[str, Any]], + neighbor_map: Dict[str, Tuple[str, float]], + pool_answers: Dict[str, str], base_dp: int, + args: argparse.Namespace, checker, + base_cache: DiskCache, rubric_cache: DiskCache) -> None: + """Cross-problem rubric (--xproblem-rubric): for every view-A problem Q, replace its own + rubric with the rubric of its bag-of-words neighbour P (Q keeps being the solved/scored + problem). P is baselined + diagnosed here (both disk-cached) as a stub carrying P's REAL + answer, so P's baseline grades correctly and legitimately shares the baseline cache with + P-as-training-problem. P's findings are copied onto Q with the neighbour text + similarity + for audit. A P that can't be diagnosed -> Q degrades to view B. Since P's answer differs + from Q's (guaranteed at pairing time), any answer P's rubric leaks is useless for Q.""" + targets = [r for r in chunk if r.get('_view') == 'A' and neighbor_map.get(r['problem'])] + if not targets: + return + stubs, by_problem = [], {} + for r in targets: + p, _ = neighbor_map[r['problem']] + if p not in by_problem: + stub = {'problem': p, 'reference_answer': pool_answers.get(p, ''), '_view': 'A'} + by_problem[p] = stub + stubs.append(stub) + baseline_rollout(base_sampler, stubs, base_dp, args, base_cache) + diagnose_views(checker, stubs, args, rubric_cache) + for r in targets: + p, sim = neighbor_map[r['problem']] + r['_rubric_diag'] = by_problem[p].get('_rubric_diag', '') + r['_rubric_src'], r['_neighbor_sim'] = p, sim + + +def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], + ci: int, base_dp: int, skill_dp: int, args: argparse.Namespace, + checker, rubric_cache: DiskCache, base_cache: DiskCache = None, + neighbor_map: Optional[Dict[str, Tuple[str, float]]] = None, + pool_answers: Optional[Dict[str, str]] = None + ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], List[Dict[str, Any]]]: + """view assign -> rubric-check (view A) -> skill-gen -> answer audit -> with-skill + greedy pass -> GRPO advantages. ``chunk`` arrives already baselined by draw_chunk. + With --xproblem-rubric, view A gets its NEIGHBOUR's rubric (see apply_neighbor_rubric).""" + hard = chunk + for r in hard: + r['_view'], r['_rubric_diag'] = _assign_view(r['problem'], args), '' + if args.xproblem_rubric and neighbor_map: + apply_neighbor_rubric(base_sampler, hard, neighbor_map, pool_answers or {}, base_dp, + args, checker, base_cache, rubric_cache) + else: + diagnose_views(checker, hard, args, rubric_cache) + + # skill-gen (thinking OFF), per-view prompt; re-sample problems with no clean candidate. + # View-A problems that cannot yield a training record (no [FAIL] rubric, or rubric + # leaked the answer) are dropped from training entirely -- skip their generation. + flat: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] + pending = [r for r in hard if not _viewa_dropped(r, args)] + for _ in range(args.skill_retries + 1): + if not pending: + break + sg_out = _run_samples(skill_sampler, [_view_prompt(r, args) for r in pending], + args.n_skills, args.skill_max_tokens, skill_dp, + temperature=args.skill_gen_temperature, + top_p=args.skill_gen_top_p, top_k=args.skill_gen_top_k) + still = [] + for r, seqs in zip(pending, sg_out): + got = False + for s in seqs: + resp = _clean_text(getattr(s, 'decoded', '') or '') + block = _extract_skill(resp) or '' + cand = {'skills': block, 'response': resp, 'parseable': bool(block), + 'view': r['_view'], 'leaked': None, 'leak_reason': '', + 'leak_source': '', 'with_pass': None, 'reward': None, 'rolls': [], + 'skillgen_stop': getattr(s, 'stop_reason', None), + 'skillgen_tokens': len(getattr(s, 'tokens', None) or [])} + r['_cands'].append(cand) + if block: + flat.append((r, cand)) + got = True + if not got: + still.append(r) + pending = still + + # Answer-bearing skill audit: deterministic verbatim-answer check only (no LLM). This + # is observability only; it records metrics for swanlab/jsonl, but does not block + # scoring, reward, advantage assignment, SFT candidate selection, or GRPO training. + for r, c in flat: + leaked = _answer_leaked(c['skills'], r['reference_answer']) + c['leaked'] = leaked + c['leak_reason'] = 'answer_verbatim' if leaked else '' + c['leak_source'] = 'deterministic' + + # with-skill greedy pass (T=0, M=1); reward = correct, absolute (group mean is baseline). + scored_inputs = flat + if scored_inputs: + ws_out = _run_samples(base_sampler, + [build_skill_solve_prompt(r['problem'], c['skills']) for r, c in scored_inputs], + 1, args.max_tokens, base_dp, temperature=0.0) + for (r, c), seqs in zip(scored_inputs, ws_out): + c['rolls'] = [_parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll()] + c['with_pass'] = 1.0 if c['rolls'][0]['correct'] else 0.0 + c['reward'] = c['with_pass'] + if args.format_in_reward: # unparseable candidates score 0 and still join the group + for r in hard: + for c in r['_cands']: + if c['reward'] is None: + c['reward'] = 0.0 + + _assign_advantages(hard, args) + return ([_full_record(r, ci) for r in chunk], _chunk_summary(chunk, ci, args), + _group_records(chunk, args)) + + +def _roll(x: Dict[str, Any]) -> Dict[str, Any]: + return {k: x[k] for k in ('pred', 'correct', 'terminated', 'passed', + 'stop_reason', 'gen_tokens', 'text')} + + +def _is_trainable(c: Dict[str, Any], args: argparse.Namespace) -> bool: + """Reaches the GRPO update iff advantage is non-zero. Leak flags are audit-only.""" + adv_nz = abs(c.get('advantage') or 0.0) > 1e-9 + if args.format_in_reward: + return adv_nz + return c.get('with_pass') is not None and adv_nz + + +def _full_record(r: Dict[str, Any], ci: int) -> Dict[str, Any]: + """Complete per-problem trace: init attempt, baseline, and all candidates.""" + init = r['_init'][0] + return { + 'record_type': 'problem', 'chunk': ci, 'problem': r['problem'], + 'reference_answer': r['reference_answer'], 'level': r.get('level', ''), + 'failed_first_try': r['_failed'], + 'init_attempt': {'text': init['text'], 'pred': init['pred'], 'correct': init['correct'], + 'terminated': init['terminated'], 'stop_reason': init['stop_reason'], + 'gen_tokens': init['gen_tokens']}, + 'baseline_pass': r['_baseline_pass'], 'is_hard': r['_hard'], + 'view': r.get('_view', ''), 'rubric_diag': r.get('_rubric_diag', ''), + # xproblem audit: which neighbour's rubric this problem borrowed, and how similar. + 'rubric_src': r.get('_rubric_src', ''), 'neighbor_sim': r.get('_neighbor_sim'), + 'baseline_rolls': [_roll(x) for x in r['_baseline_rolls']], + 'candidates': [{ + 'skills': c['skills'], 'response': c['response'], 'parseable': c['parseable'], + 'skillgen_stop': c.get('skillgen_stop'), 'skillgen_tokens': c.get('skillgen_tokens'), + 'leaked': c['leaked'], 'leak_reason': c['leak_reason'], 'leak_source': c['leak_source'], + 'with_pass': c['with_pass'], 'reward': c.get('reward'), 'advantage': c.get('advantage'), + 'grpo_adv': c.get('grpo_adv'), 'kept': c.get('kept'), + 'rolls': [_roll(x) for x in c['rolls']], + } for c in r['_cands']], + } + + +def _view_stats(problems: List[Dict[str, Any]], view: str) -> Dict[str, Any]: + pv = [r for r in problems if r.get('_view') == view] + cands = [c for r in pv for c in r['_cands'] if c['parseable']] + clean = [c for c in cands if c['leaked'] is False] + adopted = sum(1 for r in pv + if any(abs(c.get('advantage') or 0.0) > 1e-9 for c in r['_cands'])) + return {'n': len(pv), 'n_candidates_parseable': len(cands), 'n_clean': len(clean), + 'n_adopted_problems': adopted, 'adoption_rate': (adopted / len(pv)) if pv else 0.0} + + +def _mean(xs: List[float]) -> float: + return sum(xs) / len(xs) if xs else 0.0 + + +def _std(xs: List[float]) -> float: + if len(xs) < 2: + return 0.0 + m = _mean(xs) + return (sum((x - m) ** 2 for x in xs) / len(xs)) ** 0.5 + + +def _signal_stats(problems: List[Dict[str, Any]]) -> Dict[str, float]: + """The heart of 'is there a learning signal': per problem, the scored candidates form a + GRPO group. A group with zero reward variance (all skills solve, or none do -- the + hard-problem dead zone) gives NO gradient. Tracks that fraction plus reward level and + within-group variance so a collapse (all-0 or all-1) is visible immediately.""" + group_vars, all_rewards, zero_grad, groups = [], [], 0, 0 + for r in problems: + rewards = [c['reward'] for c in r['_cands'] if c.get('reward') is not None] + if len(rewards) < 2: + continue + groups += 1 + all_rewards.extend(rewards) + v = _std(rewards) + group_vars.append(v) + if v < 1e-9: # every skill got the same reward -> GRPO skips this problem + zero_grad += 1 + return {'n_groups': groups, 'zero_grad_frac': (zero_grad / groups) if groups else 0.0, + 'reward_mean': _mean(all_rewards), 'reward_std': _std(all_rewards), + 'group_reward_std_mean': _mean(group_vars)} + + +def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespace) -> Dict[str, Any]: + all_cands = [c for r in chunk for c in r['_cands']] + cands = [c for c in all_cands if c['parseable']] + scored = [c for c in cands if c['with_pass'] is not None] + clean = [c for c in cands if c['leaked'] is False] + ws_rolls = [x for c in scored for x in c['rolls']] + # viewa-dropped problems generate no candidates; keep acc/* on the generated subset + # so the with-skill/lift trend stays comparable across view_b_frac settings. + gen_probs = [r for r in chunk if r['_cands']] + base_acc = _mean([1.0 if r['_baseline_pass'] else 0.0 for r in gen_probs]) + ws_acc = _mean([1.0 if any(c.get('reward') for c in r['_cands']) else 0.0 for r in gen_probs]) + cand_pass_parseable = _mean([c['with_pass'] for c in scored]) + cand_pass_all = _mean([1.0 if c.get('reward') else 0.0 for c in all_cands]) + # base failure taxonomy (you asked whether skills fail because the base loops out of length) + classes = [_baseline_class(r) for r in chunk] + n_fail = sum(1 for c in classes if c != 'success') + skill_tokens = [c.get('skillgen_tokens') or 0 for c in cands] # skill-gen response length + trunc = sum(1 for r in chunk for c in r['_cands'] + for x in c['rolls'] if x['stop_reason'] == 'length') + rubric_answer_leaks = sum( + 1 for r in chunk + if r.get('_view') == 'A' and _rubric_has_fail(r.get('_rubric_diag')) and _rubric_answer_leaked(r)) + return { + 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), + 'n_failed_first_try': sum(1 for r in chunk if r['_failed']), + 'n_generated': len(all_cands), 'n_candidates_parseable': len(cands), + 'n_unparseable': len(all_cands) - len(cands), + 'parse_rate': (len(cands) / len(all_cands)) if all_cands else 0.0, + 'n_leaked': sum(1 for c in cands if c['leaked']), 'n_clean': len(clean), + 'leak_rate': (sum(1 for c in cands if c['leaked']) / len(cands)) if cands else 0.0, + 'n_reward_pos': sum(1 for c in scored if c['reward']), + 'n_rubric_answer_leaked': rubric_answer_leaks, + 'n_viewa_dropped': sum(1 for r in chunk if _viewa_dropped(r, args)), + 'n_train_samples': sum(1 for c in all_cands if _is_trainable(c, args)), + 'signal': _signal_stats(chunk), + 'fail_loop_frac': (sum(1 for c in classes if c == 'fail_loop') / n_fail) if n_fail else 0.0, + 'fail_wrong_frac': (sum(1 for c in classes if c == 'fail_wrong') / n_fail) if n_fail else 0.0, + 'skill_tokens_mean': _mean(skill_tokens), + 'withskill_trunc_frac': (trunc / len(ws_rolls)) if ws_rolls else 0.0, + 'avg_baseline_pass': base_acc, 'avg_withskill_pass': ws_acc, + 'avg_lift': ws_acc - base_acc, + 'candidate_withskill_pass_parseable': cand_pass_parseable, + 'candidate_withskill_pass_all': cand_pass_all, + 'termination_rate_withskill': _mean([1.0 if x['terminated'] else 0.0 for x in ws_rolls]), + 'view_A': _view_stats(chunk, 'A'), 'view_B': _view_stats(chunk, 'B'), + **_xproblem_stats(chunk, args), + } + + +def _xproblem_stats(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> Dict[str, Any]: + """Cross-problem pairing health: of the view-A problems, how many actually got a + neighbour's rubric (pair_rate) and how similar those neighbours were. Empty when off.""" + if not args.xproblem_rubric: + return {} + view_a = [r for r in chunk if r.get('_view') == 'A'] + paired = [r for r in view_a if r.get('_rubric_src')] + return {'xproblem': { + 'n_view_a': len(view_a), 'n_paired': len(paired), + 'pair_rate': (len(paired) / len(view_a)) if view_a else 0.0, + 'neighbor_sim_mean': _mean([r.get('_neighbor_sim', 0.0) for r in paired])}} + + +def _sft_record(r: Dict[str, Any], c: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: + """A view-A context-distillation sample: the OPEN-BOOK (rubric-conditioned) skill ``c`` + is the target, but the prompt is rebuilt as query-only (view B, no rubric) so the model + learns to produce it CLOSED-BOOK. Rides the same GRPOLoss with a positive constant + advantage (``--sft-weight``); single-step (old_logps=None) this reduces to + ``-w*logp + beta*KL`` -- a KL-anchored weighted cross-entropy toward the skill.""" + return { + 'data_id': r.get('data_id'), + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'view': 'B', 'diagnosis': '', 'rubric_src': '', 'orig_view': 'A', 'sft': True, + 'response': c['response'], 'skills': c['skills'], + 'skillgen_stop': c.get('skillgen_stop'), + 'advantage': float(args.sft_weight), 'grpo_adv': 0.0, 'kept': True, + 'reward': c['reward'], 'with_pass': c['with_pass']} + + +def _rubric_answer_leaked(r: Dict[str, Any]) -> bool: + """Hard safety gate for external API/rubric diagnosis. Unlike answer-bearing skills + generated by the policy itself, a rubric that contains the target final answer is an + external teacher leak and must not be distilled into view B.""" + return _answer_leaked(r.get('_rubric_diag', ''), r.get('reference_answer', '')) + + +def _viewa_dropped(r: Dict[str, Any], args: argparse.Namespace) -> bool: + """Under --viewa-sft a view-A problem trains ONLY through the SFT path (rubric with + [FAIL], answer-free). No-FAIL and rubric-leaked problems produce no training record + at all (no GRPO backflow: those prompts are query-only and would muddy the pure + view-B GRPO pool), so their skill-gen + executor scoring is skipped entirely.""" + return (bool(args.viewa_sft) and r.get('_view') == 'A' + and (not _rubric_has_fail(r.get('_rubric_diag')) + or _rubric_answer_leaked(r))) + + +def _group_records(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> List[Dict[str, Any]]: + """Training records. Under --viewa-sft, view A is SFT-only: a problem whose rubric + localised a failure (has ``[FAIL]``, answer-free) contributes ONE context-distillation + SFT sample (best parseable open-book skill -- preferring an executor-verified pass, + else any parseable candidate -- rebuilt query-only); no-FAIL / rubric-leaked view-A + problems are dropped (``_viewa_dropped``, generation already skipped). GRPO candidates + come from view B only, keeping that pool purely closed-book. Prompts are rebuilt from + the stored view/diagnosis by ``_skillgen_messages``.""" + out = [] + for r in chunk: + if not r['_hard']: + continue + if args.viewa_sft and r.get('_view') == 'A': + if _viewa_dropped(r, args): + continue + best = _best_sft_candidate(r, args) + if best is not None: + out.append(_sft_record(r, best, args)) + continue + for c in r['_cands']: + if _is_trainable(c, args): + out.append({ + 'data_id': r.get('data_id'), + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'view': r.get('_view', 'A'), 'diagnosis': r.get('_rubric_diag', ''), + 'rubric_src': r.get('_rubric_src', ''), 'sft': False, + 'response': c['response'], 'skills': c['skills'], + 'skillgen_stop': c.get('skillgen_stop'), + 'advantage': c['advantage'], 'grpo_adv': c['grpo_adv'], 'kept': c['kept'], + 'reward': c['reward'], 'with_pass': c['with_pass']}) + return out + + +# =========================================================================== +# Block G -- online GRPO training +# =========================================================================== +def _is_num(v: Any) -> bool: + try: + float(v) + return True + except (TypeError, ValueError): + return False + + +def _train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: + """Training sample = the exact skill-gen prompt (rebuilt by ``_skillgen_messages`` so + train/inference match) + the generated structured guidance response. ``key_rounds`` + selects the final assistant turn; Template masks the prompt and trains the whole + response (the key-round prefix already excludes the prompt-provided ).""" + msgs = _skillgen_messages( + rec['problem'], rec.get('view', 'A'), rec.get('diagnosis', ''), rec.get('rubric_src', '')) + return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], + 'user_data': {'key_rounds': [len(msgs)]}} + + +def _train_chunk(skill_model, ref_model, ckpt: CheckpointEngineManager, samples: List[Dict[str, Any]], + args: argparse.Namespace) -> Dict[str, Any]: + """On-policy GRPO update over one chunk, then sync weights. Micro-batches of + ``sft_batch_size`` accumulate gradients; ONE optimizer step is taken per PPO + mini-batch of ``ppo_mini_batch_size`` samples (0 -> a single step over the whole + chunk, the original behaviour). A frozen reference model provides ref_logps for the + SEAM-style KL penalty. + + Multi-step correctness: with more than one step over the SAME rollout, later + mini-batches see an already-updated policy, so we FREEZE the sampling-policy + ``old_logps`` (recomputed once, before any step) and let GRPOLoss form the PPO ratio + against them. A single step keeps ``old_logps=None`` (ratio==1, pure on-policy). + The batch is padded to a multiple of ``sft_batch_size`` with advantage-0 copies that + contribute no policy gradient. View-A context-distillation samples ride the same loss + with a positive constant advantage (``--sft-weight``); single-step (old_logps=None) + that is ``-w*logp + beta*KL`` (KL-anchored weighted cross-entropy).""" + trajs = [_train_trajectory(rec) for rec in samples] + advs = [float(rec['advantage']) for rec in samples] + rem = (-len(trajs)) % args.sft_batch_size + if rem: + trajs += [trajs[-1]] * rem + advs += [0.0] * rem + + n, sft = len(trajs), args.sft_batch_size + mini = args.ppo_mini_batch_size if args.ppo_mini_batch_size > 0 else n + mini = max(sft, (mini // sft) * sft) # align to a whole number of micro-batches + multi_step = mini < n + + # Freeze ref_logps (frozen model) and -- only when we take multiple steps -- the + # sampling-policy old_logps, BOTH before any optimizer step touches the weights. With + # a single step old_logps stays None so GRPOLoss uses ratio==1 (pure on-policy). + micro_ref, micro_old = [], [] + for i in range(0, n, sft): + mb = trajs[i:i + sft] + micro_ref.append(ref_model.forward_only(inputs=mb).get('logps')) + micro_old.append(skill_model.forward_only(inputs=mb).get('logps') if multi_step else None) + + micro, n_steps = 0, 0 + for ms in range(0, n, mini): + for i in range(ms, min(ms + mini, n), sft): + k = i // sft + skill_model.forward_backward(inputs=trajs[i:i + sft], + advantages=advs[i:i + sft], + old_logps=micro_old[k], + ref_logps=micro_ref[k]) + micro += 1 + skill_model.clip_grad_and_step() + n_steps += 1 + ckpt.sync_weights(merge_and_sync=True) + metric = skill_model.calculate_metric(is_training=True) + n_sft = sum(1 for s in samples if s.get('sft')) + return {'n_samples': len(samples), 'n_sft': n_sft, 'n_grpo': len(samples) - n_sft, + 'n_steps': n_steps, 'n_micro_batches': micro, + 'metric': {k: (float(v) if _is_num(v) else v) for k, v in (metric or {}).items()}} + + +# =========================================================================== +# Block H -- fixed-holdout eval + metric formatting +# =========================================================================== +def run_greedy_eval(base_sampler, skill_sampler, eval_records: List[Dict[str, Any]], + ci: int, rounds: int, base_dp: int, skill_dp: int, args: argparse.Namespace, + base_cache: DiskCache + ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], Dict[str, float]]: + """SEAM ``val-core/math/acc/mean@1`` on the fixed holdout: ONE greedy skill (T=0) per + problem into ONE greedy base solve (T=0). Eval ALWAYS uses view B (query-only, the + deployment form: no rubric, since rubric needs an online teacher unavailable at deploy); + no leak filter (acc scores correctness alone). Baseline reuses the disk cache.""" + baseline_rollout(base_sampler, eval_records, base_dp, args, base_cache) + for r in eval_records: + r['_view'], r['_rubric_diag'] = 'B', '' + sg_out = _run_samples(skill_sampler, [_view_prompt(r, args) for r in eval_records], + 1, args.skill_max_tokens, skill_dp, temperature=0.0) + skills = [] + for seqs in sg_out: + if not seqs: + skills.append(('', '')) + continue + sresp = _clean_text(getattr(seqs[0], 'decoded', '') or '') + skills.append((_extract_skill(sresp) or '', sresp)) + ws_out = _run_samples(base_sampler, + [build_skill_solve_prompt(r['problem'], sk) for r, (sk, _) in zip(eval_records, skills)], + 1, args.max_tokens, base_dp, temperature=0.0) + recs = [] + for r, (sk, sresp), seqs in zip(eval_records, skills, ws_out): + roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + recs.append({ + 'record_type': 'eval_problem', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, + 'data_id': r.get('data_id'), + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'view': r['_view'], 'rubric_diag': r.get('_rubric_diag', ''), + 'baseline_pass': r['_baseline_pass'], 'skill': sk, 'skill_parseable': bool(sk), + 'skill_response': sresp, 'withskill_pred': roll['pred'], + 'withskill_correct': roll['correct'], 'withskill_terminated': roll['terminated'], + 'withskill_stop_reason': roll['stop_reason'], 'withskill_text': roll['text'], + }) + acc = lambda rs: sum(1 for x in rs if x['withskill_correct']) / len(rs) if rs else 0.0 + ws = acc(recs) # all view B (deployment form) + base = sum(x['baseline_pass'] for x in recs) / len(recs) if recs else 0.0 + fmt = (sum(1 for x in recs if x['skill_parseable']) / len(recs)) if recs else 0.0 + term = (sum(1 for x in recs if x['withskill_terminated']) / len(recs)) if recs else 0.0 + summary = {'record_type': 'eval_summary', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, + 'n': len(recs), 'view': 'B', 'acc_mean1': ws, + 'baseline_acc_mean1': base, 'lift_mean1': ws - base, + 'format_mean1': fmt, 'term_mean1': term} + metrics = {'core/math/acc/mean@1': ws, 'core/math/baseline_acc/mean@1': base, + 'core/math/lift/mean@1': ws - base, 'core/math/format/mean@1': fmt, + 'core/math/term/mean@1': term} + return recs, summary, metrics + + +def _trend_line(hist: List[Dict[str, float]], window: int, rounds_done: int) -> Optional[str]: + """Contrast the FIRST vs the most recent ``window`` chunks -- if RFT works, adoption + and lift on recent (fresh) chunks exceed the early baseline.""" + if len(hist) < 2 * window: + return None + base, rec = hist[:window], hist[-window:] + m = lambda xs, k: sum(h[k] for h in xs) / len(xs) + return (f'[trend] first {window} vs last {window} | ' + f'adopt A {m(base,"aA"):.2f}->{m(rec,"aA"):.2f} B {m(base,"aB"):.2f}->{m(rec,"aB"):.2f} | ' + f'lift {m(base,"lift"):+.3f}->{m(rec,"lift"):+.3f} | ' + f'0grad {m(base,"zero_grad"):.2f}->{m(rec,"zero_grad"):.2f} | ' + f'pos/chunk {m(base,"pos"):.1f}->{m(rec,"pos"):.1f} | rounds={rounds_done}') + + +def _swan_metrics(summary: Dict[str, Any], log: Optional[Dict[str, Any]]) -> Dict[str, float]: + """Flat swanlab dict. ``signal/*`` is the primary health group (is GRPO getting a + gradient at all); ``acc/*`` is the effect; the rest diagnose why. acc/adopt/skill are + only logged when the chunk produced a scored group, so idle chunks don't dip charts.""" + sig = summary['signal'] + d: Dict[str, float] = { + # --- signal: the FIRST thing to watch (no variance -> no learning) --- + 'signal/zero_grad_frac': sig['zero_grad_frac'], 'signal/n_groups': sig['n_groups'], + 'signal/reward_mean': sig['reward_mean'], 'signal/reward_std': sig['reward_std'], + 'signal/group_reward_std_mean': sig['group_reward_std_mean'], + 'signal/n_train_samples': summary['n_train_samples'], + 'signal/n_reward_pos': summary['n_reward_pos'], + # --- skill format / leak health --- + 'skill/parse_rate': summary['parse_rate'], 'skill/tokens_mean': summary['skill_tokens_mean'], + 'leak/rate': summary['leak_rate'], 'leak/n': summary['n_leaked'], + # --- base failure taxonomy (loop-out-of-length vs plain wrong) --- + 'fail/loop_frac': summary['fail_loop_frac'], 'fail/wrong_frac': summary['fail_wrong_frac'], + 'fail/frac_first_try': (summary['n_failed_first_try'] / summary['n']) if summary['n'] else 0.0, + # --- view-A routing: dropped = no-FAIL rubric or rubric leak (SFT-only view A) --- + 'viewa/dropped_frac': (summary['n_viewa_dropped'] / summary['view_A']['n'] + if summary['view_A']['n'] else 0.0), + } + bal = summary.get('balance') or {} + if bal.get('enabled'): + d.update({'balance/n_drawn': bal['n_drawn'], 'balance/n_baseline_fresh': bal['n_baseline_fresh'], + 'balance/selected_success_frac': bal['selected_success_frac']}) + xp = summary.get('xproblem') or {} + if xp: + d.update({'xproblem/pair_rate': xp['pair_rate'], 'xproblem/neighbor_sim_mean': xp['neighbor_sim_mean']}) + if sig['n_groups'] > 0: + d.update({'acc/baseline_pass': summary['avg_baseline_pass'], + 'acc/withskill_pass': summary['avg_withskill_pass'], 'acc/lift': summary['avg_lift'], + 'candidate/withskill_pass_parseable': summary['candidate_withskill_pass_parseable'], + 'candidate/withskill_pass_all': summary['candidate_withskill_pass_all'], + 'adopt/A': summary['view_A']['adoption_rate'], + 'adopt/B': summary['view_B']['adoption_rate'], + 'term/withskill': summary['termination_rate_withskill'], + 'term/withskill_trunc_frac': summary['withskill_trunc_frac']}) + if log: + d['train/n_steps'] = log['n_steps'] + d['train/n_micro_batches'] = log['n_micro_batches'] + for k, v in (log.get('metric') or {}).items(): + if not _is_num(v): + continue + if k.startswith('learning rate'): + if 'group 1' in k: + d['train/lr'] = float(v) + else: + d[f'train/{k.replace(" ", "_")}'] = float(v) + return d + + +def _view_a_rubric_leak_metrics(chunk: List[Dict[str, Any]], + pool_answers: Optional[Dict[str, str]] = None) -> Dict[str, float]: + """Swanlab-only audit for answer leakage in view-A rubric text. This never changes + rewards, advantages, filtering, or training records.""" + view_a = [r for r in chunk if r.get('_view') == 'A'] + with_diag = [r for r in view_a if r.get('_rubric_diag')] + target_leaks = sum(1 for r in with_diag + if _answer_leaked(r.get('_rubric_diag', ''), r.get('reference_answer', ''))) + source_leaks = 0 + pool_answers = pool_answers or {} + for r in with_diag: + src = r.get('_rubric_src') + src_ref = pool_answers.get(src, '') if src else r.get('reference_answer', '') + if _answer_leaked(r.get('_rubric_diag', ''), src_ref): + source_leaks += 1 + n = len(with_diag) + return { + 'rubric_leak/n_view_a': float(len(view_a)), + 'rubric_leak/n_checked': float(n), + 'rubric_leak/target_answer_n': float(target_leaks), + 'rubric_leak/target_answer_rate': (target_leaks / n) if n else 0.0, + 'rubric_leak/source_answer_n': float(source_leaks), + 'rubric_leak/source_answer_rate': (source_leaks / n) if n else 0.0, + } + + +# =========================================================================== +# Block F -- components, args, main +# =========================================================================== +def init_components(args: argparse.Namespace): + """Default 8-GPU layout: rank 0 trains the actor, rank 1 hosts a frozen ref model, + 2-3 skill_sampler (synced), 4-7 base_sampler (frozen). Returns + (skill_model, ref_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp).""" + r0, r1 = TRAIN_GPUS, TRAIN_GPUS + REF_GPUS + r2, r3 = r1 + SKILL_SAMPLER_GPUS, NUM_GPUS + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, lazy_collect=False, groups=[ + DeviceGroup(name='train', ranks=list(range(0, r0)), device_type='GPU'), + DeviceGroup(name='ref', ranks=list(range(r0, r1)), device_type='GPU'), + DeviceGroup(name='skill_sampler', ranks=list(range(r1, r2)), device_type='GPU'), + DeviceGroup(name='base_sampler', ranks=list(range(r2, r3)), device_type='GPU')]) + + train_mesh = DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_DP, fsdp_size=TRAIN_FSDP) + skill_model = TransformersModel(model_id=MODEL_ID, device_mesh=train_mesh, remote_group='train', + ddp_config={'find_unused_parameters': False}) + skill_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) + skill_model.set_template(Template, model_id=MODEL_ID, enable_thinking=False, + max_length=args.max_model_len, truncation_strategy='delete') + skill_model.set_processor(InputProcessor, padding_free=False) + skill_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon, beta=args.kl_beta) + skill_model.set_optimizer('AdamW', lr=args.lr) + skill_model.set_lr_scheduler('CosineWarmupScheduler', num_warmup_steps=10, + num_training_steps=args.max_train_rounds) + + ref_mesh = DeviceMesh.from_sizes(world_size=REF_GPUS, dp_size=REF_DP, fsdp_size=REF_FSDP) + ref_model = TransformersModel(model_id=MODEL_ID, device_mesh=ref_mesh, remote_group='ref', + ddp_config={'find_unused_parameters': False}) + ref_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) + ref_model.set_template(Template, model_id=MODEL_ID, enable_thinking=True, + max_length=args.max_model_len, truncation_strategy='delete') + ref_model.set_processor(InputProcessor, padding_free=False) + ref_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) + + def _sampler(group, world, enable_thinking: bool = True): + s = vLLMSampler(model_id=MODEL_ID, + engine_args={'gpu_memory_utilization': GPU_MEM, + 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, + device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), + remote_group=group) + s.set_template(Template, model_id=MODEL_ID, enable_thinking=enable_thinking, max_length=args.max_model_len) + return s + + skill_sampler = _sampler('skill_sampler', SKILL_SAMPLER_GPUS, enable_thinking=False) + # base_sampler is shared by the main thread and the baseline-prefetch thread -> serialise. + base_sampler = _LockedSampler(_sampler('base_sampler', BASE_SAMPLER_GPUS, enable_thinking=True)) + ckpt = CheckpointEngineManager(model=skill_model, sampler=skill_sampler) + return skill_model, ref_model, skill_sampler, base_sampler, ckpt, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS + + +def _build_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--dataset', choices=('aops', 'math'), default='aops') + p.add_argument('--n', type=int, default=2000, help='Problems loaded into the draw pool.') + p.add_argument('--pool-offset', type=int, default=0, + help='Skip this many shuffled non-eval records before building the train pool; ' + 'useful to avoid cold-start SFT data ranges.') + p.add_argument('--exclude-data-ids', default='', + help='Comma-separated jsonl files whose data_id/problem keys are excluded ' + 'from train/eval selection, e.g. coldstart_sft.jsonl.') + p.add_argument('--seed', type=int, default=42) + p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--eval-size', type=int, default=128, help='Fixed holdout size (0 disables).') + p.add_argument('--eval-every', type=int, default=5, help='Run holdout eval every N chunks.') + p.add_argument('--chunk-size', type=int, default=16) + p.add_argument('--balance', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--balance-success-frac', type=float, default=0.4, + help='Target fraction of the chunk the base solves (rest are base-fail).') + p.add_argument('--balance-loop-frac', type=float, default=0.5) + p.add_argument('--balance-max-draws-mult', type=int, default=8) + p.add_argument('--seam-parquet-dir', type=str, default='', + help='Read SEAM build_aops_dataset.py train.parquet/val.parquet directly, in ' + 'file order (problem<-extra_info.problem, answer<-reward_model.ground_truth). ' + 'val.parquet becomes the eval holdout. Bypasses load/--numeric-only/' + '--eval-size/internal shuffle so the input data matches a SEAM run exactly.') + p.add_argument('--n-skills', type=int, default=8) + p.add_argument('--view-b-frac', type=float, default=0.5) + p.add_argument('--xproblem-rubric', action=argparse.BooleanOptionalAction, default=False, + help='When enabled, view A uses a bag-of-words NEIGHBOUR problem rubric. ' + 'Default is off: each view-A problem uses its own baseline attempt, ' + 'while the API diagnosis prompt is constrained to be answer-free and ' + 'method-level only.') + p.add_argument('--skill-retries', type=int, default=2) + p.add_argument('--skill-gen-temperature', type=float, default=1.0) + p.add_argument('--skill-gen-top-p', type=float, default=1.0) + p.add_argument('--skill-gen-top-k', type=int, default=-1) + p.add_argument('--max-model-len', type=int, default=16384) + p.add_argument('--max-tokens', type=int, default=8192) + p.add_argument('--skill-max-tokens', type=int, default=8192) + p.add_argument('--rubric-workers', type=int, default=16) + p.add_argument('--rubric-retries', type=int, default=2, + help='Retry failed/timeout rubric diagnose calls this many times before ' + 'falling back to an empty diagnosis without caching the failure.') + p.add_argument('--sft-batch-size', type=int, default=8, + help='Driver micro-batch (gradient-accumulation unit); multiple of train dp.') + p.add_argument('--ppo-mini-batch-size', type=int, default=0, + help='Samples per optimizer step (SEAM-style PPO mini-batch). 0 = one step ' + 'over the whole chunk (on-policy, ratio==1). When >0 and smaller than ' + 'the trainable count, multiple steps are taken over the same rollout and ' + 'old_logps are frozen so the PPO ratio/clip stays valid. Rounded down to ' + 'a multiple of --sft-batch-size.') + p.add_argument('--grpo-epsilon', type=float, default=0.2) + p.add_argument('--adv-clip', type=float, default=3.0, + help='Symmetric clip for group-relative advantages; <=0 disables clipping.') + p.add_argument('--kl-beta', type=float, default=0.001, + help='SEAM-style reference KL coefficient for GRPOLoss.') + p.add_argument('--viewa-sft', action=argparse.BooleanOptionalAction, default=True, + help='Route view-A problems to online context distillation (SFT on the best ' + 'passing open-book skill, rebuilt as a query-only prompt) instead of GRPO. ' + 'View B stays GRPO; both share one optimizer step.') + p.add_argument('--sft-weight', type=float, default=0.5, + help='Advantage magnitude (lambda) for view-A SFT distillation samples in the ' + 'shared GRPOLoss. Single-step this is -lambda*logp + beta*KL (weighted CE). ' + 'Keep <~1 so one-signed SFT gradients do not swamp mean-0 GRPO advantages.') + p.add_argument('--sft-target-len', type=int, default=550, + help='Target skill length (chars) for view-A SFT distillation: among passing ' + 'candidates pick the one CLOSEST to this length (not the shortest). Set near ' + 'the empirical pass-rate peak (~500-600 here) so distillation neither shrinks ' + 'rollouts toward zero nor lets them grow unbounded.') + p.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--lr', type=float, default=6e-6) + p.add_argument('--max-train-rounds', type=int, default=1500) + p.add_argument('--save-rounds', type=int, default=200) + p.add_argument('--trend-every', type=int, default=10) + p.add_argument('--output-dir', default='./output/reflexion_skill') + p.add_argument('--cache-dir', default='', help='Baseline/rubric cache dir (default /cache).') + p.add_argument('--no-cache', action='store_true', help='Disable disk cache read/write.') + p.add_argument('--prefetch-baseline', action=argparse.BooleanOptionalAction, default=True, + help='Prefetch next chunk base baseline on a background thread (overlaps ' + 'with skill-gen; base_sampler is frozen so it never blocks the trainer).') + p.add_argument('--swanlab-project', default='twinkle') + p.add_argument('--swanlab-exp', default='') + args = p.parse_args() + if args.sft_batch_size % TRAIN_DP != 0: + raise ValueError(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple of train dp ({TRAIN_DP})') + if args.chunk_size < 1: + raise ValueError('--chunk-size must be >= 1') + args.balance_success_frac = max(0.0, min(1.0, args.balance_success_frac)) + return args + + +def _write(handle, row: Dict[str, Any]) -> None: + handle.write(json.dumps(row, ensure_ascii=False) + '\n') + + +def main() -> None: + args = _build_args() + records, eval_records, neighbor_map, pool_answers, data_stats = _load_records(args) + if len(records) < args.chunk_size: + raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded ({len(records)}); raise --n') + + os.makedirs(args.output_dir, exist_ok=True) + gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') + eval_path = os.path.join(args.output_dir, 'eval_records.jsonl') + data_path = os.path.join(args.output_dir, 'skill_dataset.jsonl') + train_log_path = os.path.join(args.output_dir, 'train_log.jsonl') + + if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY')): + sys.stderr.write('[rft] WARNING: no LLM backup env; view-A rubric check disabled ' + '(leak filter is deterministic, unaffected)\n') + + use_swan = swanlab is not None and os.environ.get('SWANLAB_MODE') != 'disabled' + if use_swan: + swanlab.init(project=args.swanlab_project, experiment_name=(args.swanlab_exp or None), + config={'model': MODEL_ID, 'dataset': args.dataset, 'n': len(records), + 'eval_n': len(eval_records), 'n_skills': args.n_skills, + 'view_b_frac': args.view_b_frac, 'balance': args.balance, + 'balance_success_frac': args.balance_success_frac, + 'skill_gen_temp': args.skill_gen_temperature, + 'grpo_epsilon': args.grpo_epsilon, 'kl_beta': args.kl_beta, + 'lr': args.lr}) + + skill_model, ref_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp = init_components(args) + checker = build_rubric_checker() + if checker is None: + sys.stderr.write('[rft] no LLM backup env -> view-A rubric process-check DISABLED\n') + + cache_dir = args.cache_dir or os.path.join(args.output_dir, 'cache') + os.makedirs(cache_dir, exist_ok=True) + use_cache = not args.no_cache + base_cache = DiskCache(os.path.join(cache_dir, 'baseline.jsonl'), use_cache) + eval_base_cache = DiskCache(os.path.join(cache_dir, 'eval_baseline.jsonl'), use_cache) + rubric_cache = DiskCache(os.path.join(cache_dir, 'rubric.jsonl'), use_cache) + if args.xproblem_rubric: + sys.stderr.write(f'[rft] xproblem-rubric ON: {len(neighbor_map)} bag-of-words pairs\n') + + cfg = {'record_type': 'config', 'model': MODEL_ID, 'dataset': args.dataset, + 'n': len(records), 'eval_n': len(eval_records), 'seed': args.seed, + 'pool_offset': args.pool_offset, 'exclude_data_ids': args.exclude_data_ids, + 'excluded_records': data_stats.get('excluded_records', 0), + 'numeric_only': args.numeric_only, 'raw_loaded': data_stats['raw_loaded'], + 'numeric_dropped': data_stats['numeric_dropped'], 'eval_every': args.eval_every, + 'n_skills': args.n_skills, 'view_b_frac': args.view_b_frac, + 'skill_retries': args.skill_retries, 'balance': args.balance, + 'balance_success_frac': args.balance_success_frac, + 'skill_gen_temp': args.skill_gen_temperature, 'reward': 'greedy_binary(correct)', + 'advantage': 'group_relative', 'format_in_reward': args.format_in_reward, 'cache': use_cache, + 'rubric_check': f'fixed_math_{len(_MATH_RUBRIC)}crit+term(viewA)' if checker else 'disabled', + 'xproblem_rubric': args.xproblem_rubric, + 'viewa_sft': args.viewa_sft, 'sft_weight': args.sft_weight, + 'sft_target_len': args.sft_target_len, + 'adv_clip': args.adv_clip, + 'grpo_epsilon': args.grpo_epsilon, 'kl_beta': args.kl_beta, 'lr': args.lr, + 'train_gpus': TRAIN_GPUS, 'ref_gpus': REF_GPUS, 'ref_fsdp': REF_FSDP, + 'train_fsdp': TRAIN_FSDP, 'train_dp': TRAIN_DP, + 'skill_sampler_gpus': SKILL_SAMPLER_GPUS, 'base_sampler_gpus': BASE_SAMPLER_GPUS, + 'max_train_rounds': args.max_train_rounds, 'started': int(time.time())} + sys.stderr.write(f'[rft] raw={data_stats["raw_loaded"]} numeric_drop={data_stats["numeric_dropped"]} ' + f'train={len(records)} eval={len(eval_records)} {args.dataset}; ' + f'train_gpus={TRAIN_GPUS} ref_gpus={REF_GPUS} train_fsdp={TRAIN_FSDP} ' + f'train_dp={TRAIN_DP} skill_dp={skill_dp} base_dp={base_dp}\n') + + hist: List[Dict[str, float]] = [] + rounds = 0 + pool = ProblemPool(records, args.seed) + with open(gen_path, 'w', encoding='utf-8') as gen_f, \ + open(eval_path, 'w', encoding='utf-8') as eval_f, \ + open(data_path, 'w', encoding='utf-8') as data_f, \ + open(train_log_path, 'w', encoding='utf-8') as tlog: + for f in (gen_f, eval_f, data_f, tlog): + _write(f, cfg) + gstep = 0 + # base baseline (frozen sampler, disk-cached, no weight-sync) is prefetched on a + # background thread while the current chunk generates: the skill-gen phase uses + # skill_sampler, so prefetching the NEXT chunk's baseline on base_sampler overlaps + # it and fills base_cache so the next draw_chunk hits it. base_sampler is wrapped in + # _LockedSampler, so the prefetch and this chunk's with-skill scoring never issue a + # base .sample() concurrently. It never touches the trainer or on-policy generation. + prefetch_pool = ThreadPoolExecutor(max_workers=1) if args.prefetch_baseline else None + pending: Optional[Any] = None + + def _prefetch(peeked: List[Dict[str, Any]]) -> None: + if peeked: + baseline_rollout(base_sampler, peeked, base_dp, args, base_cache) + + # Baseline (round 0) eval BEFORE any training: measures the untrained skill model on + # the fixed holdout so every later eval has a step-0 reference point on the same axis. + if eval_records: + eval_recs, eval_summary, eval_metrics = run_greedy_eval( + base_sampler, skill_sampler, eval_records, -1, rounds, base_dp, skill_dp, + args, eval_base_cache) + for rec in eval_recs: + _write(eval_f, rec) + _write(eval_f, eval_summary) + eval_f.flush() + if use_swan: + swanlab.log({f'eval/{k}': v for k, v in eval_metrics.items()}, step=0) + sys.stderr.write( + f'[eval] g-1 (init): n={eval_summary["n"]} viewB mean@1 ' + f'acc={eval_summary["baseline_acc_mean1"]:.3f}->{eval_summary["acc_mean1"]:.3f} ' + f'lift={eval_summary["lift_mean1"]:+.3f} ' + f'fmt={eval_summary["format_mean1"]:.2f} rounds={rounds}\n') + + # Each chunk is drawn fresh + RE-GENERATED with the current policy (on-policy); + # with --balance, draw_chunk keeps drawing until the base fail:success mix hits target. + while rounds < args.max_train_rounds: + if pending is not None: + pending.result() # finish last round's prefetch before drawing (cache-warm) + pending = None + chunk, balance = draw_chunk(pool, base_sampler, base_dp, args, base_cache) + if prefetch_pool is not None: + peeked = pool.peek(int(args.chunk_size * (args.balance_max_draws_mult if args.balance else 1))) + pending = prefetch_pool.submit(_prefetch, peeked) + full, summary, groups = process_chunk( + base_sampler, skill_sampler, chunk, gstep, base_dp, skill_dp, + args, checker, rubric_cache, base_cache, neighbor_map, pool_answers) + summary['balance'] = balance + + log = None + if groups: + log = _train_chunk(skill_model, ref_model, ckpt, groups, args) + rounds += 1 + log.update({'record_type': 'train_round', 'round': rounds, 'chunk': gstep, + 'epoch': pool.epoch, 'ts': int(time.time())}) + _write(tlog, log) + tlog.flush() + if rounds % args.save_rounds == 0: + skill_model.save(f'skill-rft-{rounds}', output_dir=args.output_dir) + + summary['rounds_done'], summary['epoch'] = rounds, pool.epoch + for rec in full: + _write(gen_f, rec) + _write(gen_f, summary) + gen_f.flush() + for v in groups: + _write(data_f, v) + data_f.flush() + + sa, sb, sig = summary['view_A'], summary['view_B'], summary['signal'] + hist.append({'aA': sa['adoption_rate'], 'aB': sb['adoption_rate'], + 'lift': summary['avg_lift'], 'pos': summary['n_reward_pos'], + 'zero_grad': sig['zero_grad_frac']}) + bal_str = (f'bal {balance["selected_fail"]}f/{balance["selected_success"]}s ' + f'(drew {balance["n_drawn"]}/fresh {balance["n_baseline_fresh"]}' + + ('!' if balance.get('budget_hit') else '') + ') ') if balance.get('enabled') else '' + xp = summary.get('xproblem') + xp_str = f'pair={xp["pair_rate"]:.2f}@{xp["neighbor_sim_mean"]:.2f} ' if xp else '' + tr_str = (f'train={log["n_grpo"]}g/{log["n_sft"]}s ' if log + else f'train={summary["n_train_samples"]} ') + sys.stderr.write( + f'[gen] e{pool.epoch} g{gstep}: {bal_str}n={summary["n"]} ' + f'clean={summary["n_clean"]} leak={summary["leak_rate"]:.2f} {tr_str}' + f'0grad={sig["zero_grad_frac"]:.2f} R={sig["reward_mean"]:.2f}±{sig["reward_std"]:.2f} ' + f'acc={summary["avg_baseline_pass"]:.2f}->{summary["avg_withskill_pass"]:.2f} ' + f'lift={summary["avg_lift"]:+.3f} {xp_str}' + f'A[{sa["n"]} {sa["adoption_rate"]:.2f}] B[{sb["n"]} {sb["adoption_rate"]:.2f}] ' + f'rounds={rounds}\n') + if use_swan: + swan_metrics = _swan_metrics(summary, log) + swan_metrics.update(_view_a_rubric_leak_metrics(chunk, pool_answers)) + swanlab.log(swan_metrics, step=gstep) + + if eval_records and (gstep + 1) % args.eval_every == 0: + eval_recs, eval_summary, eval_metrics = run_greedy_eval( + base_sampler, skill_sampler, eval_records, gstep, rounds, base_dp, skill_dp, + args, eval_base_cache) + for rec in eval_recs: + _write(eval_f, rec) + _write(eval_f, eval_summary) + eval_f.flush() + if use_swan: + swanlab.log({f'eval/{k}': v for k, v in eval_metrics.items()}, step=gstep) + sys.stderr.write( + f'[eval] g{gstep}: n={eval_summary["n"]} viewB mean@1 ' + f'acc={eval_summary["baseline_acc_mean1"]:.3f}->{eval_summary["acc_mean1"]:.3f} ' + f'lift={eval_summary["lift_mean1"]:+.3f} ' + f'fmt={eval_summary["format_mean1"]:.2f} rounds={rounds}\n') + + if (gstep + 1) % args.trend_every == 0: + tl = _trend_line(hist, args.trend_every, rounds) + if tl: + sys.stderr.write(tl + '\n') + gstep += 1 + + if prefetch_pool is not None: + if pending is not None: + pending.result() + prefetch_pool.shutdown(wait=True) + base_cache.close() + eval_base_cache.close() + rubric_cache.close() + skill_model.save('skill-rft-final', output_dir=args.output_dir) + sys.stderr.write(f'[rft] done: {rounds} rounds over {gstep} chunks / {pool.epoch} epochs -> {data_path}\n') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/legacy/train_reflexion_skill_old.sh b/cookbook/exp/legacy/train_reflexion_skill_old.sh new file mode 100755 index 000000000..093bf3521 --- /dev/null +++ b/cookbook/exp/legacy/train_reflexion_skill_old.sh @@ -0,0 +1,57 @@ +#!/bin/bash +# Online GRPO RFT for the reflexion skill generator (unified, self-contained, cached). +# GPUs: 8 — default high-memory layout uses rank 0 for actor training, rank 1 for +# the frozen ref model, ranks 2-3 for skill sampler (synced), and ranks 4-7 for +# base sampler (frozen). Override TRAIN_GPUS / REF_GPUS / SKILL_SAMPLER_GPUS / +# BASE_SAMPLER_GPUS for other layouts. Per chunk: base greedy +# solve -> rubric process-check (view A) -> +# skill-gen (thinking ON, N candidates) -> deterministic leak filter -> with-skill greedy +# pass -> group-relative advantage -> ONE on-policy GRPO step -> sync weights. +# +# Baseline rollouts + rubric diagnoses are disk-cached (output-dir/cache/*.jsonl), so a +# restart skips re-sampling them; skill-gen is on-policy and never cached. The next chunk's +# baseline is prefetched on a background thread (overlaps skill-gen; base sampler is frozen). +# +# The view-A rubric process-check uses the backup teacher API (set LLM_BACKUP_*). Without +# it the run still works: view A degrades to query-only and the leak filter stays +# deterministic (no teacher needed). + +set -euo pipefail + +export GEN_MODEL_ID=${GEN_MODEL_ID:-Qwen/Qwen3-4B} +export GEN_GPU_MEM=${GEN_GPU_MEM:-0.8} +# Datasets are pulled from ModelScope via twinkle.Dataset (ms://AI-MO/aops or +# ms://modelscope/competition_math); override AOPS_DATASET_ID / MATH_DATASET_ID to change. +# Teacher API for the view-A rubric process-check (optional; leak filter is deterministic). +export LLM_BACKUP_API_KEY=${LLM_BACKUP_API_KEY:-} +export LLM_BACKUP_BASE_URL=${LLM_BACKUP_BASE_URL:-https://dashscope.aliyuncs.com/compatible-mode/v1} +export LLM_BACKUP_MODEL=${LLM_BACKUP_MODEL:-qwen3.7-max} + +python cookbook/exp/embedding/train_reflexion_skill.py \ + --dataset aops \ + --n 5000 \ + --seam-parquet-dir /root/data/seam \ + --numeric-only \ + --chunk-size 32 \ + --n-skills 16 \ + --view-b-frac 0.5 \ + --skill-retries 2 \ + --no-balance \ + --max-tokens 8192 \ + --skill-max-tokens 4096 \ + --max-model-len 16384 \ + --eval-size 128 \ + --eval-every 5 \ + --sft-batch-size 8 \ + --ppo-mini-batch-size 0 \ + --grpo-epsilon 0.2 \ + --kl-beta 0.001 \ + --format-in-reward \ + --lr 1e-6 \ + --max-train-rounds 1500 \ + --save-rounds 200 \ + --trend-every 10 \ + --prefetch-baseline \ + --output-dir ./output/reflexion_skill \ + --swanlab-project twinkle \ + --swanlab-exp reflexion_skill_rft diff --git a/cookbook/exp/legacy/train_reflexion_skill_replay.py b/cookbook/exp/legacy/train_reflexion_skill_replay.py new file mode 100644 index 000000000..0580a6ba4 --- /dev/null +++ b/cookbook/exp/legacy/train_reflexion_skill_replay.py @@ -0,0 +1,114 @@ +"""Replay-train the reflexion skill model from prebuilt exact RFT data. + +Use ``build_reflexion_skill_data.py`` first to create ``skill_dataset.jsonl``. This +script trains only the skill model from those frozen records; it does not run vLLM +rollouts, leak checks, or rubric diagnosis. + +Launch: + python cookbook/exp/embedding/train_reflexion_skill_replay.py \ + --data ./output/reflexion_skill_data/skill_dataset.jsonl +""" +import argparse +import json +import os +import sys +from collections import defaultdict +from typing import Any, Dict, List + +import train_reflexion_skill_rft as rft + + +def _build_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument('--data', default='./output/reflexion_skill_data/skill_dataset.jsonl') + p.add_argument('--output-dir', default='./output/reflexion_skill_replay') + p.add_argument('--epochs', type=int, default=1) + p.add_argument('--sft-batch-size', type=int, default=8) + p.add_argument('--max-model-len', type=int, default=16384) + p.add_argument('--grpo-epsilon', type=float, default=0.2) + p.add_argument('--lr', type=float, default=1e-5) + p.add_argument('--save-rounds', type=int, default=50) + return p.parse_args() + + +def _load_chunks(path: str) -> List[List[Dict[str, Any]]]: + chunks: Dict[int, List[Dict[str, Any]]] = defaultdict(list) + fallback_chunk = 0 + with open(path, 'r', encoding='utf-8') as f: + for line_no, line in enumerate(f, 1): + if not line.strip(): + continue + row = json.loads(line) + if row.get('record_type') == 'config': + continue + for key in ('problem', 'response', 'advantage'): + if key not in row: + raise ValueError(f'{path}:{line_no} missing required field {key!r}') + ci = int(row.get('chunk', fallback_chunk)) + chunks[ci].append(row) + if 'chunk' not in row and len(chunks[ci]) >= 64: + fallback_chunk += 1 + return [chunks[k] for k in sorted(chunks) if chunks[k]] + + +def _init_model(args: argparse.Namespace, total_updates: int): + model = 'ms://Qwen/Qwen3-4B' + train_mesh = rft.DeviceMesh.from_sizes( + world_size=rft.TRAIN_GPUS, dp_size=rft.TRAIN_DP, fsdp_size=rft.TRAIN_FSDP) + device_groups = [ + rft.DeviceGroup(name='train', ranks=list(range(rft.TRAIN_GPUS)), device_type='GPU'), + ] + rft.twinkle.initialize(mode='ray', nproc_per_node=rft.TRAIN_GPUS, groups=device_groups, + lazy_collect=False) + model = rft.TransformersModel(model_id=model, device_mesh=train_mesh, + remote_group='train', ddp_config={'find_unused_parameters': False}) + from twinkle.patch.no_split_modules import NoSplitModulesPatch + model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) + model.set_template(rft.Template, model_id=model, + enable_thinking=True, max_length=args.max_model_len, + truncation_strategy='delete') + model.set_processor(rft.InputProcessor, padding_free=False) + model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) + model.set_optimizer('AdamW', lr=args.lr) + model.set_lr_scheduler('CosineWarmupScheduler', num_warmup_steps=10, + num_training_steps=max(1, total_updates)) + return model + + +def main() -> None: + args = _build_args() + if args.sft_batch_size % rft.TRAIN_DP != 0: + raise ValueError(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple ' + f'of the training dp size ({rft.TRAIN_DP})') + chunks = _load_chunks(args.data) + if not chunks: + raise ValueError(f'no train records found in {args.data}') + os.makedirs(args.output_dir, exist_ok=True) + total_updates = len(chunks) * args.epochs + model = _init_model(args, total_updates) + log_path = os.path.join(args.output_dir, 'train_log.jsonl') + cfg = {'record_type': 'config', 'mode': 'offline_replay', 'data': args.data, + 'chunks': len(chunks), 'epochs': args.epochs, 'lr': args.lr, + 'sft_batch_size': args.sft_batch_size} + rounds = 0 + with open(log_path, 'w', encoding='utf-8') as tlog: + tlog.write(json.dumps(cfg, ensure_ascii=False) + '\n') + for epoch in range(args.epochs): + for ci, samples in enumerate(chunks): + log = rft._train_chunk(model, None, samples, args) + rounds += 1 + log.update({'record_type': 'train_round', 'round': rounds, + 'epoch': epoch, 'chunk': ci}) + tlog.write(json.dumps(log, ensure_ascii=False) + '\n') + tlog.flush() + sys.stderr.write( + f'[replay-rft] e{epoch} c{ci}: n={log["n_samples"]} ' + f'micro={log["n_micro_batches"]} metric={log.get("metric")}\n') + if rounds % args.save_rounds == 0: + model.save(f'skill-rft-replay-{rounds}', output_dir=args.output_dir) + model.save('skill-rft-replay-final', output_dir=args.output_dir) + sys.stderr.write(f'[replay-rft] done: {rounds} updates; log -> {log_path}\n') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/legacy/train_reflexion_skill_rft.py b/cookbook/exp/legacy/train_reflexion_skill_rft.py new file mode 100644 index 000000000..b4388b773 --- /dev/null +++ b/cookbook/exp/legacy/train_reflexion_skill_rft.py @@ -0,0 +1,1568 @@ +"""RFT cold-start for the reflexion skill generator (see reflexion.md §6). + +Trains an INDEPENDENT skill model to write reusable, transferable skills that, +when injected into a FROZEN base solver's system prompt, let the base solve problems +it first got wrong. The base is never trained — it only produces the reward signal. +Scoring is SEAM-style DETERMINISTIC: the base runs each candidate skill once at +temperature 0 (M=1), so the reward ``R in {0,1}`` (answer correct) carries no +sampling noise; the per-candidate advantage is group-relative within a problem +(``A = (R - mean) / (std + eps)``) and the skill model is updated online by GRPO — +problem-groups where every skill scores alike (std=0) contribute no gradient. + +Direction: skill GENERATION + recall. Skill-gen always runs with thinking ON, and +each hard problem is routed to EXACTLY ONE of two views (no reuse — kills memory +leak and holds cost at 1x): view A ``(problem + attempt) -> think + skills`` keeps +the online generator self-bootstrapping; view B ``(problem only) -> think + skills`` +is the deployment form, where the think is grounded on the query alone so it cannot +hallucinate an attempt. Both share the verified skill; the distilled ```` +block is recalled into the base's system prompt at solve time. + +8-GPU layout (three DeviceGroups, one twinkle.initialize): + - ranks 0-3 : ``train`` — skill model, full-param FSDP2, dp=4 + - ranks 4-5 : ``skill_sampler`` — skill model rollouts (vLLM, tp1 dp2) + - ranks 6-7 : ``base_sampler`` — frozen base solver (vLLM, tp1 dp2) +CheckpointEngineManager syncs train -> skill_sampler after every optimizer step; +base_sampler is never synced. + +Leak filtering uses ``LeakVerifier(sampler=None)`` via the backup teacher API +(no local judge, no distillation): set LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / +LLM_BACKUP_MODEL. + +Launch (8 GPUs): + LLM_BACKUP_API_KEY=... LLM_BACKUP_BASE_URL=... LLM_BACKUP_MODEL=... \ + python cookbook/exp/embedding/train_reflexion_skill_rft.py --n 2000 --chunk-size 16 +""" +import argparse +import hashlib +import json +import os +import re +import sys +import time +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.model import TransformersModel +from twinkle.processor import InputProcessor +from twinkle.sampler import vLLMSampler +from twinkle.template import Template +from twinkle_agentic.verifier import LeakVerifier, RubricVerifier +from twinkle_agentic.verifier.rubric_verifier import RubricItem + +# Reuse the reference eval's dataset + grading + prompts + sampling config, and the +# phase-0 pipeline's parsing / rollout / injection helpers (Find > Create). +from eval_gpqa_rag import (GEN_GPU_MEM, GEN_MODEL_ID, build_direct_prompt, # noqa: F401 + load_aops, load_math) +from eval_reflexion_skill import (_EX_PROBLEM, _clean_text, # noqa: F401 + _parse_seq, _run_samples, build_skill_solve_prompt) + +logger = get_logger() + +try: + import swanlab +except ImportError: # optional; metric logging degrades to stdout + jsonl only + swanlab = None + + +# -- GPU layout --------------------------------------------------------------- +TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 4)) +SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) +BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 2)) +NUM_GPUS = TRAIN_GPUS + SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS +# FSDP shard group size within a dp replica; TRAIN_DP is the data-parallel axis that +# ``forward_backward`` (slice_dp) splits each mini-batch over. +TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', 2)) +TRAIN_DP = max(1, TRAIN_GPUS // TRAIN_FSDP) + + +# --------------------------------------------------------------------------- +# Skill-generation prompt (DISTILL the useful approach, per the new direction) +# --------------------------------------------------------------------------- +# --- Previous STRICT view-A system prompt (commented out; kept for easy revert). It +# hard-required 3-5 bullets, "output nothing after ", one-imperative-sentence +# items, no-narration, and a strict do-not-reveal block. The soft SEAM-style version +# below drops those four format demands, frames the skills as advisory reminders, and +# explains how they are used. --- +# SKILL_GEN_SYSTEM = ( +# 'You are distilling reusable problem-solving SKILLS from one worked episode. ' +# 'You are shown a competition problem, the guidance the solver was given, and the ' +# "solver's own attempt (its reasoning may be partly right and partly wrong).\n\n" +# 'FIRST, in your private thinking, do ALL of: (a) work out what this TYPE of problem ' +# 'fundamentally requires; (b) pinpoint WHERE THIS attempt actually went wrong ' +# '(when a process-check report is provided below, use its flagged criteria as ' +# 'evidence, but confirm each against the attempt yourself) — ' +# 'the decisive misstep, a missing idea, a wrong turn, or the way it stalled, looped ' +# 'on the same step, or ran the length budget out without ever committing to an ' +# 'answer; and (c) imagine AS MANY DIFFERENT angles as you can — distinct approaches ' +# 'or representations that could crack this problem, alternative solution paths, and ' +# 'the various ways a solver could plausibly go wrong on it (a few words each, do NOT ' +# 'develop them fully). THEN commit to the angle you find most decisive and write a ' +# 'SHORT list of skills that would have PREVENTED that specific ' +# 'failure and would raise the success rate of a SIMILAR solver on SIMILAR problems. ' +# 'Ground each skill in the concrete mistake you found, but state it as a GENERAL, ' +# 'transferable rule — not a patch hard-coded to this problem. Across the 3-5 ' +# 'bullets, prioritise in this order:\n' +# '1. the decisive method or representation this class of problem calls for (what to ' +# 'set up or reach for first);\n' +# '2. the specific mistake that derailed THIS attempt, recast as a general pitfall, ' +# 'plus the quick check that catches it;\n' +# '3. a compact, ordered procedure that reliably drives toward a final answer;\n' +# '4. convergence discipline: once the key quantity is in hand, commit to a single ' +# 'concrete final answer in the required format instead of re-deriving, endless ' +# 'case-splitting, looping on the same check, or overrunning the length budget.\n\n' +# 'OUTPUT FORMAT (strict):\n' +# '- Keep the thinking COMPACT — a quick brainstorm of angles then a decision, not a ' +# 'full solution and not a re-statement of these instructions. AFTER it, ' +# 'output ONLY a markdown bullet list of 3-5 items WRAPPED IN and ' +# 'tags — no preamble, no narration outside the tags. Output nothing after ' +# '.\n' +# '- Each item is ONE short imperative sentence (a method, heuristic, check, or ' +# 'habit).\n' +# '- Inside the tags: no narration, no "The student...", no headings, no restating ' +# 'the problem.\n\n' +# 'CONTENT RULES (strict):\n' +# '- Do NOT reveal the final answer or the multiple-choice option.\n' +# '- Do NOT state the specific numbers, values, or key intermediate results of THIS ' +# 'problem.\n' +# '- Every item must be GENERAL and transferable, not a step-by-step solution to ' +# 'THIS problem.\n\n' +# 'Follow the example below for the exact tags, style, and level of generality.' +# ) +SKILL_GEN_SYSTEM = ( + 'You are a mathematics coach. You are shown a competition problem together with an ' + 'automated process-check of an earlier solver attempt at it -- which solution ' + 'criteria the attempt passed or failed, and suggested fixes for the failures. You ' + 'do NOT see the attempt itself, only this check. Treat the check as privileged ' + 'training scaffolding: study it together with the problem, identify the ' + 'problem-visible features that make each useful flagged failure relevant, then ' + 'rephrase those lessons as self-contained reusable skills. The goal is not to ' + 'continue from the check, cite it, or hide it silently; the goal is to turn it to ' + 'a skill pattern which prevents the model falls into similar pitfalls in the future.\n\n' + 'Good skills name the observable trigger, the method worth reaching for, the ' + 'pitfall to watch, and a quick verification habit. Prefer formulations like ' + '"When a configuration has ...", "Before setting up ...", or "Check whether ..." ' + 'over references to the process-check, failed criteria, or the earlier attempt. ' + "These tips are advisory: they will be placed in a solver's system prompt as gentle " + 'reminders before it works through a SIMILAR problem on its own, without seeing ' + 'this process-check. So keep them general and transferable rather than a worked ' + 'solution to this exact problem, and do not state its specific intermediate values ' + 'or final answer. Think briefly first, then give your tips as a markdown bullet ' + 'list wrapped in and , like the example below.' +) + +# One-shot demo of the recommended mix (method / pitfall+check / procedure / +# convergence), answer-free — anchors both the format and the content priorities. +_EX_SKILLS = ( + '\n' + '- Rewrite each square root by factoring its radicand into a perfect square times ' + 'a remainder, then move the perfect-square factor outside.\n' + '- Avoid the classic trap $\\sqrt{a}+\\sqrt{b}\\ne\\sqrt{a+b}$; only combine terms ' + 'sharing the same simplest radical, and sanity-check by estimating each root.\n' + '- Procedure: simplify every radical, group like radical terms, add their ' + 'coefficients, then reduce to simplest form.\n' + '- Once the expression is in simplest form, commit to that single result as the ' + 'final answer rather than re-checking indefinitely.\n' + '') + + +# View A user template: the problem + the automated rubric process-check of an earlier +# attempt (PASS/FAIL per criterion + suggested fixes). The attempt trajectory is NOT +# shown -- the rubric findings are the evidence the skill model grounds its tips on, +# which avoids feeding the (often long, non-terminating) attempt into the prompt. +SKILL_GEN_USER_RUBRIC = ( + 'Problem:\n{problem}\n\n' + 'Process check of an earlier attempt (automated rubric verifier -- treat as ' + 'evidence, not gospel; PASS/FAIL per criterion with suggested fixes for failures):\n' + '{diagnosis}\n\n' + 'Now output a self-contained skills bullet list. Each bullet should still be useful ' + 'if the process check were removed: connect any useful flagged failure to ' + 'problem-visible features, general methods, and quick checks rather than citing the ' + 'rubric or the earlier attempt. \n\n' + 'Note: **Do not solve the problem, only generate skills**. Now Begin:' +) + + +def build_skillgen_prompt(problem: str, diagnosis: str) -> Dict[str, Any]: + """View A skill-gen prompt: system + one-shot format demo + the real episode + (problem + the rubric process-check of an earlier attempt). The attempt trajectory + is deliberately NOT shown -- the rubric findings localise the failure without the + generator having to re-chew (and often re-solve) a long, possibly non-terminating + attempt. The one-shot demo is query-only; only the real turn carries the diagnosis.""" + return {'messages': [ + {'role': 'system', 'content': SKILL_GEN_SYSTEM}, + # {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=_EX_PROBLEM)}, + # {'role': 'assistant', 'content': _EX_SKILLS}, + {'role': 'user', + 'content': SKILL_GEN_USER_RUBRIC.format(problem=problem, diagnosis=diagnosis)}, + ]} + + +# --------------------------------------------------------------------------- +# View B: query-only skill-gen (deployment form). No attempt is shown — the model +# must reason about the problem TYPE from the query alone, so the think is grounded +# on the query and cannot narrate/fabricate an attempt. Format is deliberately +# distinct from view A so the model learns the two modes as separate contracts. +# --------------------------------------------------------------------------- +# --- Previous STRICT view-B (query-only) system prompt (commented out; kept for revert). +# Same four format demands as the old view A. Soft SEAM-style version below. --- +# SKILL_GEN_SYSTEM_Q = ( +# 'You are distilling reusable problem-solving SKILLS for a CLASS of problems. You ' +# 'are shown ONE competition problem and NOTHING else — no solution, no attempt. ' +# 'FIRST, in your private thinking, imagine AS MANY DIFFERENT angles as you can — ' +# 'distinct approaches or representations that could crack this TYPE of problem, ' +# 'alternative solution paths, and the various ways a solver could plausibly go wrong ' +# 'on it (a few words each, do NOT develop them fully). THEN commit to what you find ' +# 'most decisive and write a SHORT list of skills that would raise a solver\'s success ' +# 'rate on SIMILAR problems. Across the 3-5 bullets, prioritise in this order:\n' +# '1. the decisive method or representation this class of problem calls for (what to ' +# 'set up or reach for first);\n' +# '2. the specific pitfall that derails such problems, plus the quick check that ' +# 'catches it;\n' +# '3. a compact, ordered procedure that reliably drives toward a final answer;\n' +# '4. convergence discipline: once the key quantity is in hand, commit to a single ' +# 'concrete final answer in the required format instead of re-deriving, endless ' +# 'case-splitting, or overrunning the length budget.\n\n' +# 'OUTPUT FORMAT (strict):\n' +# '- Keep the thinking COMPACT — a quick brainstorm of angles then a decision, not a ' +# 'full solution and not a re-statement of these instructions. AFTER ' +# 'it, output ONLY a markdown bullet list of 3-5 items WRAPPED IN and ' +# ' tags — no preamble, no narration outside the tags. Output nothing after ' +# '.\n' +# '- Each item is ONE short imperative sentence (a method, heuristic, check, or ' +# 'habit).\n' +# '- Inside the tags: no narration, no headings, no restating the problem, and no ' +# 'reference to any attempt, student, or solution.\n\n' +# 'CONTENT RULES (strict):\n' +# '- Do NOT solve THIS problem or reveal its final answer or multiple-choice option.\n' +# '- Do NOT state the specific numbers, values, or key intermediate results of THIS ' +# 'problem.\n' +# '- Every item must be GENERAL and transferable to other problems of the same ' +# 'type.\n\n' +# 'Follow the example below for the exact tags, style, and level of generality.' +# ) +SKILL_GEN_SYSTEM_Q = ( + 'You are a mathematics coach. You are shown ONE competition problem and nothing ' + 'else — no solution and no attempt. Think about what approach this KIND of problem ' + 'calls for and where solvers tend to slip, then distil a few reusable tips.\n\n' + "These tips are advisory: they will be placed in a solver's system prompt as gentle " + 'reminders before it works through a SIMILAR problem on its own. So keep them ' + 'general and transferable — the method worth reaching for, the pitfall to watch and ' + 'a quick check, and the discipline to settle on a final answer — rather than a ' + 'worked solution to this exact problem, and without stating its specific ' + 'intermediate values or its final answer. Think briefly first, then give your tips ' + 'as a markdown bullet list wrapped in and , like the example below.' +) + +SKILL_GEN_USER_Q = ( + 'Problem:\n{problem}\n\n' + 'Now reason about this TYPE of problem, then output the skills bullet list.' +) + + +def build_querygen_prompt(problem: str) -> Dict[str, Any]: + """View B skill-gen prompt: system + one-shot demo + the problem ALONE (no + attempt) — matching what is available at deployment (query only).""" + return {'messages': [ + {'role': 'system', 'content': SKILL_GEN_SYSTEM_Q}, + # {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=_EX_PROBLEM)}, + # {'role': 'assistant', 'content': _EX_SKILLS}, + {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=problem)}, + ]} + + +def _assign_view(problem: str, args: argparse.Namespace) -> str: + """Deterministically route a problem to exactly one view (stable across restarts + and across the generation/SFT sides). ``--view-b-frac`` of problems go to view B.""" + h = int(hashlib.md5(f'{args.seed}:{problem}'.encode('utf-8')).hexdigest(), 16) + return 'B' if (h % 100000) / 100000.0 < args.view_b_frac else 'A' + + +def _skillgen_messages(problem: str, view: str, diagnosis: str) -> List[Dict[str, Any]]: + """Single source of truth for the skill-gen prompt, used at BOTH generation and + training time so they can never diverge. View A with a localisable failure uses + problem + rubric findings (NO trajectory); view B -- or a view-A problem whose rubric + flagged NO failure (``[FAIL]`` absent: all-pass or missing diagnosis) -- is query-only. + So view A DEGRADES to view B whenever there is nothing concrete to correct.""" + if view == 'B' or '[FAIL]' not in (diagnosis or ''): + return build_querygen_prompt(problem)['messages'] + return build_skillgen_prompt(problem, diagnosis)['messages'] + + +def _view_prompt(r: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: + """The skill-gen prompt for problem ``r`` under its assigned view (routing in + ``_skillgen_messages``: view A carries the rubric process-check; view B, and any + view-A problem with no rubric failure, is query-only).""" + return {'messages': _skillgen_messages(r['problem'], r['_view'], r.get('_rubric_diag', ''))} + + +_BULLET_RE = re.compile(r'(?m)^\s*(?:[-*]|\d+[.)])\s') +# Trajectory/meta references that betray CoT fragments leaking into the block; any +# hit fails the purity gate (the problem is then re-sampled, per --skill-retries). +_META_RE = re.compile( + r'\b(the student|the solver|the attempt|this attempt|the trace|the response|' + r'in the (?:attempt|trace|response|solution)|as (?:shown|seen|noted) above|' + r'the (?:above|previous|earlier)|my (?:reasoning|analysis)|i (?:think|need|will))\b', + re.IGNORECASE) + + +def _is_clean_block(block: str) -> bool: + """Purity gate for thinking-ON skill-gen: the block must be a pure bullet list + (every non-empty line a bullet — no prose/CoT fragments) with no meta/trajectory + reference. Answer leak is caught separately by the backup-teacher leak stage.""" + lines = [ln for ln in (l.strip() for l in block.splitlines()) if ln] + if not lines or not all(_BULLET_RE.match(ln) for ln in lines): + return False + return _META_RE.search(block) is None + + +def _extract_skills_block(text: str) -> Optional[str]: + """Return the clean ``...`` block, or None if not parseable. + + Skill-gen runs with thinking ON, so the model must end its reasoning with an explicit + ```` before committing an answer (whether the opening ```` is emitted by + the model or pre-injected by the chat template). We therefore REQUIRE ```` and + read only the text after the last one; its absence means the token budget was exhausted + mid-reasoning (nothing committed, per reflexion.md §6.8) — reject so a draft or a + system-prompt demo echo inside the CoT can never be mistaken for the answer. Within + the answer take the ```` block (closing tag optional), strip stray tags, and + require ``_is_clean_block`` — prose-mixed / meta-referencing fragments are rejected + for re-sampling.""" + low = text.lower() + end_think = low.rfind('') + if end_think < 0: + return None # reasoning never closed -> no committed answer + answer = text[end_think + len(''):] + low_a = answer.lower() + s = low_a.find('') + if s < 0: + return None + inner = s + len('') + e = low_a.find('', inner) + block = (answer[inner:e] if e >= 0 else answer[inner:]).strip() + block = re.sub(r'', '', block, flags=re.IGNORECASE).strip() + if not _is_clean_block(block): + return None + return block + + +# --------------------------------------------------------------------------- +# OPTIONAL stricter leak criterion (currently UNUSED -- the run uses answer_only=True, +# which flags ONLY the final answer). This variant ALSO flags concrete intermediate KEY +# results, while still permitting method / plan / pitfalls / checks. To enable, pass +# judge_system=_LEAK_JUDGE_SYSTEM to the LeakVerifier below. +# --------------------------------------------------------------------------- +_LEAK_JUDGE_SYSTEM = """\ +You check whether a HINT that will be shown to someone solving a math TASK gives away +this task's own results. + +The hint may FREELY describe the general method, which approach or technique to use, the +steps or plan to follow, common pitfalls, and sanity checks -- even when that points +strongly at HOW to solve THIS task. Describing the approach is expected of a good hint. + +The hint LEAKS only if, for THIS specific task, it states either: +- the final answer or final result (a value, expression, choice, label, or verbatim + output); or +- a concrete decisive INTERMEDIATE key result -- a specific computed value, quantity, or + fact unique to this task that hands over a key step of the answer. + +If it names only the method / plan / pitfalls / checks WITHOUT stating those concrete +intermediate values or the final result, it does NOT leak. + +Reply with exactly one word: LEAK or CLEAN.""" + + +# --------------------------------------------------------------------------- +# Rubric process-check (view A only): a frozen teacher diagnoses the base's failed +# attempt so the skill model grounds its error analysis on a verified fault +# localisation instead of guessing. Teacher-only (sampler=None -> every diagnose() +# hits llm_backup); mirrors eval_dualline_math's fixed math rubric. +# --------------------------------------------------------------------------- +_RFT_DIAG_SYSTEM = """\ +You are a process error checker for a math solution attempt. You are given a math +problem, a rubric, and one attempted solution segment. Decide PASS or FAIL for each +criterion and explain only the process error type. + +Output STRICT JSON (no prose outside it) with this shape: +{ + "items": [ + {"index": 1, "verdict": "PASS", "reason": "", + "fix": ""}, + {"index": 2, "verdict": "FAIL", "reason": "", + "fix": ""} + ], + "overall": "OK" | "ISSUES", + "summary": "" +} + +Rules: +- Judge every criterion independently and literally; a [Hard Rule] is FAIL unless + unambiguously satisfied. +- Judge ONLY what is observable in THIS segment. +- Content inside ... (or ) is internal reasoning, not + user-facing output; ignore it for "output only X" style criteria. +- For PASS items, leave "fix" as "". +- For FAIL items, "reason", "fix", and "summary" must describe only the flawed + step, theorem, arithmetic operation, case split, or verification habit. +- NEVER try to solve the query or state the correct final answer, corrected final expression, option letter, + graph/choice label, or any exact value that the answer should become. +- NEVER write phrases like "the correct answer is", "which gives", "yielding", + "should be ", "Option ", or "Graph ". +- If a fix would require naming a corrected value, replace it with a method-level + instruction such as "redo that computation carefully" or "apply the theorem with + the correct quantities". +- Keep every "reason" and "fix" clear and concise. +- "overall" is "OK" only if NO criterion is FAIL. +- Output only the JSON object.""" + + +_RFT_DIAG_USER = """\ +## Task / query (context) +{query} + +## Rubric +{rubric} + +## Segment +{segment} + +Now output the diagnostic JSON object.""" + + +class _RftRubricVerifier(RubricVerifier): + def _diagnose_trajectory(self, query: str, rubric_block: str, segment_text: str) -> dict: + user = _RFT_DIAG_USER.format(query=query, rubric=rubric_block, segment=segment_text) + return {'messages': [ + {'role': 'system', 'content': _RFT_DIAG_SYSTEM}, + {'role': 'user', 'content': user}, + ]} + + +_MATH_RUBRIC = [ + ('The reasoning contains no arithmetic or algebraic error', True), + ('Each step follows logically from the previous ones', True), + ('No formula or theorem is misstated or misapplied', True), + ('The approach is on track to answer the actual question asked', False), + ('No step contradicts an earlier established fact', False), +] + + +def _build_rubric_checker() -> Optional['RubricVerifier']: + """Fixed math-process rubric verifier, teacher-served. None if no LLM backup env.""" + if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') + or os.environ.get('OPENAI_API_KEY')): + return None + return _RftRubricVerifier( + fixed_rubric=[RubricItem(t, is_hard=h) for t, h in _MATH_RUBRIC], gate=True) + + +def _format_diagnosis(detail) -> str: + """One line per criterion (PASS/FAIL + reason + fix on FAIL) then a summary — the + compact evidence block appended to the view-A skill-gen prompt.""" + rub = detail.rubric + lines = [] + for it in detail.items: + text = rub[it.index - 1].text if 0 < it.index <= len(rub) else f'criterion {it.index}' + if it.verdict: + lines.append(f'- [PASS] {text}') + else: + tail = f': {it.reason}' if it.reason else '' + tail += f' (fix: {it.fix})' if it.fix else '' + lines.append(f'- [FAIL] {text}{tail}') + if detail.summary: + lines.append(f'Summary: {detail.summary}') + return '\n'.join(lines) + + +def _diagnose_views(checker, hard: List[Dict[str, Any]], args: argparse.Namespace, + diag_cache: Optional[Dict[str, str]] = None) -> None: + """Rubric-check every view-A problem's greedy attempt in parallel, stashing the + formatted findings on ``r['_rubric_diag']`` (view B stays empty). A checker error + or empty result degrades to no diagnosis (the plain view-A prompt).""" + from concurrent.futures import ThreadPoolExecutor + targets = [r for r in hard if r.get('_view') == 'A'] + if not checker or not targets: + return + + def _cache_key(r: Dict[str, Any]) -> str: + init_text = r.get('_init', [{}])[0].get('text', '') + return hashlib.md5(f'{r["problem"]}\n{init_text}'.encode('utf-8')).hexdigest() + + pending = [] + for r in targets: + key = _cache_key(r) + if diag_cache is not None and key in diag_cache: + r['_rubric_diag'] = diag_cache[key] + else: + pending.append((r, key)) + if not pending: + return + + def _run(item: Tuple[Dict[str, Any], str]) -> Tuple[Dict[str, Any], str, str, bool]: + r, key = item + seg = {'messages': [ + {'role': 'user', 'content': r['problem']}, + {'role': 'assistant', 'content': r['_init'][0]['text']}, + ]} + try: + return r, key, _format_diagnosis(checker.diagnose(seg, query=r['problem'])), True + except Exception as exc: # teacher hiccup -> fall back to no-diagnosis prompt + logger.warning(f'[rubric] diagnose error: {exc}') + return r, key, '', False + + workers = max(1, min(args.rubric_workers, len(pending))) + with ThreadPoolExecutor(max_workers=workers) as ex: + for r, key, diag, ok in ex.map(_run, pending): + r['_rubric_diag'] = diag + if ok and diag_cache is not None: + diag_cache[key] = diag + + +# --------------------------------------------------------------------------- +# Online data generation (one chunk; every candidate is recorded, untruncated) +# --------------------------------------------------------------------------- +def _roll(x: Dict[str, Any]) -> Dict[str, Any]: + """Full (untruncated) rollout record for offline analysis.""" + return {k: x[k] for k in ('pred', 'correct', 'terminated', 'passed', + 'stop_reason', 'gen_tokens', 'text')} + + +def _empty_roll() -> Dict[str, Any]: + """Fallback rollout when the sampler returned nothing for a prompt.""" + return {'pred': '', 'correct': False, 'terminated': False, 'passed': False, + 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} + + +def _assign_advantages(hard: List[Dict[str, Any]], args: argparse.Namespace) -> None: + """Group-relative advantage (SEAM-style) over each problem's clean, scored candidates, + using the DETERMINISTIC greedy reward ``R in {0, 1}`` (answer CORRECT only; + termination is NOT part of the reward -- monitored via `terminated`/`passed` only): + + A_j = (R_j - mean_R) / (std_R + eps) + + Groups where every candidate shares the same reward (``std_R == 0``: all solve or all + fail) get advantage 0 and contribute no gradient -- GRPO's own group variance + auto-selects the informative problems, so no explicit difficulty / marginal gate is + needed. Because the reward is deterministic (M=1 greedy, no pass@k sampling), the + std-normalisation no longer amplifies rollout noise (the reason it was dropped for the + old stochastic marginal). ``kept`` marks above-average candidates (for reporting only). + """ + eps = 1e-6 + for r in hard: + for c in r['_cands']: + c['advantage'], c['grpo_adv'], c['kept'] = 0.0, 0.0, False + cs = ([c for c in r['_cands'] if c.get('reward') is not None] if args.format_in_reward + else [c for c in r['_cands'] if c['leaked'] is False and c.get('reward') is not None]) + if len(cs) < 2: + continue + rewards = [c['reward'] for c in cs] + mean_r = sum(rewards) / len(rewards) + std = (sum((x - mean_r) ** 2 for x in rewards) / len(rewards)) ** 0.5 + if std < 1e-9: + continue # all candidates equal (all solve / all fail) -> no learning signal + for c in cs: + adv = (c['reward'] - mean_r) / (std + eps) + c['advantage'], c['grpo_adv'], c['kept'] = adv, adv, c['reward'] > mean_r + + +def _full_record(r: Dict[str, Any], ci: int) -> Dict[str, Any]: + """Complete per-problem record: init attempt, baseline, and ALL candidates + (parseable/leaked/scored alike) with full text — nothing dropped or truncated.""" + init = r['_init'][0] + rec: Dict[str, Any] = { + 'record_type': 'problem', 'chunk': ci, 'problem': r['problem'], + 'reference_answer': r['reference_answer'], 'level': r.get('level', ''), + 'failed_first_try': r['_failed'], + 'init_attempt': {'text': init['text'], 'pred': init['pred'], + 'correct': init['correct'], 'terminated': init['terminated'], + 'stop_reason': init['stop_reason'], 'gen_tokens': init['gen_tokens']}, + } + rec['baseline_pass'] = r['_baseline_pass'] + rec['is_hard'] = r['_hard'] + rec['view'] = r.get('_view', '') + rec['rubric_diag'] = r.get('_rubric_diag', '') + rec['baseline_rolls'] = [_roll(x) for x in r['_baseline_rolls']] + rec['candidates'] = [{ + 'skills': c['skills'], 'response': c['response'], 'parseable': c['parseable'], + 'leaked': c['leaked'], 'leak_reason': c['leak_reason'], 'leak_source': c['leak_source'], + 'with_pass': c['with_pass'], 'reward': c.get('reward'), + 'advantage': c.get('advantage'), 'grpo_adv': c.get('grpo_adv'), 'kept': c.get('kept'), + 'rolls': [_roll(x) for x in c['rolls']], + } for c in r['_cands']] + return rec + + +def _view_stats(hard: List[Dict[str, Any]], view: str) -> Dict[str, Any]: + """Per-view yield: hard problems, clean candidates, and the ADOPTION rate — + the fraction of hard problems that produced at least one clean, non-zero-advantage + candidate (i.e. a record that actually reaches training). Watching A vs B and + early vs late tells whether query-only (B) catches up to trajectory-grounded (A).""" + hv = [r for r in hard if r.get('_view') == view] + cands = [c for r in hv for c in r['_cands'] if c['parseable']] + clean = [c for c in cands if c['leaked'] is False] + adopted = sum(1 for r in hv + if any(c['leaked'] is False and abs(c.get('advantage') or 0.0) > 1e-9 + for c in r['_cands'])) + return { + 'n_hard': len(hv), 'n_candidates_parseable': len(cands), + 'n_clean': len(clean), 'n_adopted_problems': adopted, + 'adoption_rate': (adopted / len(hv)) if hv else 0.0, + } + + +def _is_trainable(c: Dict[str, Any], args: argparse.Namespace) -> bool: + """A candidate reaches the GRPO update iff its advantage is non-zero. With + --format-in-reward every candidate carries a reward (unparseable/leaked score 0), + so non-zero advantage is the only gate; otherwise it must also be clean and scored. + Single source of truth for both the summary counts and ``_group_records``.""" + adv_nz = abs(c.get('advantage') or 0.0) > 1e-9 + if args.format_in_reward: + return adv_nz + return c['leaked'] is False and c.get('with_pass') is not None and adv_nz + + +def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespace) -> Dict[str, Any]: + """Per-chunk aggregates — watch these across chunks to see if the RFT'd skill + model produces better skills over time (yield, leak rate, lift, termination).""" + failed = [r for r in chunk if r['_failed']] + hard = [r for r in chunk if r['_hard']] + all_cands = [c for r in chunk for c in r['_cands']] + cands = [c for c in all_cands if c['parseable']] + scored = [c for c in cands if c['with_pass'] is not None] + ws_rolls = [x for c in scored for x in c['rolls']] + # With --format-in-reward, unparseable/leaked candidates also carry a (0) reward and are + # trained, so count trainables over ALL candidates; else only clean scored ones. + train_cands = [c for c in all_cands if _is_trainable(c, args)] + base_acc = (sum(r['_baseline_pass'] for r in hard) / len(hard)) if hard else 0.0 + ws_acc = (sum(c['with_pass'] for c in scored) / len(scored)) if scored else 0.0 + # -- signal-source monitor: how much of the GRPO signal comes from base-FAIL problems + # (the offensive "rescue a failure" signal we want) vs base-success (defensive "don't + # break an easy one"). abs_adv_from_fail_frac ~0.1 was the diagnosed failure mode. -- + fail_cands = [c for r in chunk if r['_failed'] for c in r['_cands']] + abs_adv = lambda cs: sum(abs(c.get('advantage') or 0.0) for c in cs) + total_abs = abs_adv(all_cands) + return { + 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), + 'n_failed_first_try': len(failed), 'n_hard': len(hard), + 'n_generated': len(all_cands), + 'n_candidates_parseable': len(cands), + 'n_unparseable': len(all_cands) - len(cands), + 'n_leaked': sum(1 for c in cands if c['leaked']), + 'n_clean': sum(1 for c in cands if c['leaked'] is False), + 'n_reward_pos': sum(1 for c in scored if c['reward']), + 'n_train_samples': len(train_cands), + 'n_train_from_fail': sum(1 for c in fail_cands if _is_trainable(c, args)), + 'abs_adv_from_fail_frac': (abs_adv(fail_cands) / total_abs) if total_abs > 0 else 0.0, + 'avg_baseline_pass_on_hard': base_acc, + 'avg_withskill_pass': ws_acc, + 'avg_lift': ws_acc - base_acc, + 'termination_rate_withskill': (sum(1 for x in ws_rolls if x['terminated']) / len(ws_rolls)) if ws_rolls else 0.0, + 'view_A': _view_stats(hard, 'A'), 'view_B': _view_stats(hard, 'B'), + } + + +def _group_records(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> List[Dict[str, Any]]: + """GRPO training records: every clean, scored skill candidate with a NON-ZERO + advantage (positive pushes the skill up, negative down; the group-relative + usefulness-over-base advantage was set in _assign_advantages). Each carries its + ``view`` and the rubric ``diagnosis``; the prompt (identical to generation) is rebuilt + from those by ``_skillgen_messages`` -- no trajectory is stored or replayed.""" + out = [] + for r in chunk: + if not r['_hard']: + continue + view = r.get('_view', 'A') + for c in r['_cands']: + if _is_trainable(c, args): + out.append({ + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'view': view, + 'diagnosis': r.get('_rubric_diag', ''), + 'response': c['response'], 'skills': c['skills'], + 'advantage': c['advantage'], 'grpo_adv': c['grpo_adv'], 'kept': c['kept'], + 'reward': c['reward'], 'with_pass': c['with_pass'], + }) + return out + + +def _apply_baseline(r: Dict[str, Any], roll: Dict[str, Any]) -> None: + """Write a (cached or fresh) greedy baseline roll onto a problem and RESET the per-chunk + working state, so a problem reused in a later chunk never carries prior skill candidates.""" + r['_baseline_rolls'], r['_cands'] = [roll], [] + r['_init'] = [roll] + r['_failed'] = not roll['correct'] + r['_baseline_pass'] = 1.0 if roll['correct'] else 0.0 + r['_hard'] = True # process EVERY selected problem; group variance selects + + +def _baseline_rollout(base_sampler, problems: List[Dict[str, Any]], base_dp: int, + args: argparse.Namespace, cache: Dict[str, Dict[str, Any]]) -> int: + """Phase 1: base solves each problem GREEDILY once (T=0, M=1), keyed-cached by problem + text across chunks. The base sampler is FROZEN and decoding is greedy, so a problem's + baseline never changes over the run -- a cache hit is exact and skips the sampler. + Returns the number of FRESH sampler rollouts (cache misses) for efficiency reporting.""" + todo = [r for r in problems if r['problem'] not in cache] + if todo: + base_out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], + 1, args.max_tokens, base_dp, temperature=0.0) + for r, seqs in zip(todo, base_out): + cache[r['problem']] = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + for r in problems: + _apply_baseline(r, cache[r['problem']]) + return len(todo) + + +def _baseline_class(r: Dict[str, Any]) -> str: + """Bucket a baselined problem by its greedy outcome: ``success`` (base solved it), + ``fail_loop`` (ran the length budget out / never terminated -- the mode skills rescue + best), or ``fail_wrong`` (terminated cleanly but the answer is wrong).""" + roll = r['_init'][0] + if roll['correct']: + return 'success' + if roll['stop_reason'] == 'length' or not roll['terminated']: + return 'fail_loop' + return 'fail_wrong' + + +_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') + + +def _norm_num_text(num: str) -> str: + try: + f = float(num) + return str(int(f)) if f == int(f) else str(f) + except Exception: + return str(num).strip() + + +def _numeric_value(raw: Any) -> Optional[str]: + if raw is None: + return None + s = str(raw).strip().strip('$').strip() + s = s.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac') + s = re.sub(r'\\!|\\,|\\;|\\ |\\left|\\right|\s', '', s) + m = re.fullmatch(r'\\frac\{(-?\d+)\}\{(-?\d+)\}', s) + if m: + a, b = int(m.group(1)), int(m.group(2)) + return _norm_num_text(str(a / b)) if b else None + m = re.fullmatch(r'(-?\d+)/(-?\d+)', s) + if m: + a, b = int(m.group(1)), int(m.group(2)) + return _norm_num_text(str(a / b)) if b else None + if _NUM_RE.fullmatch(s): + return _norm_num_text(s) + return None + + +def _numeric_only_records(records: List[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], int]: + out = [] + dropped = 0 + for r in records: + ref = _numeric_value(r.get('reference_answer')) + if ref is None: + dropped += 1 + continue + rr = dict(r) + rr['reference_answer'] = ref + out.append(rr) + return out, dropped + + +def _load_records(args: argparse.Namespace) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], Dict[str, int]]: + need_split = args.eval_size > 0 + load_n = 0 if (args.numeric_only or need_split) else args.n + records = (load_aops(n=load_n, seed=args.seed) if args.dataset == 'aops' + else load_math(n=load_n, seed=args.seed)) + raw_n = len(records) + dropped = 0 + if args.numeric_only: + records, dropped = _numeric_only_records(records) + rng = np.random.RandomState(args.seed) + rng.shuffle(records) + eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 + eval_records = [dict(r) for r in records[:eval_n]] + train_pool = records[eval_n:] + train_n = args.n if args.n > 0 else len(train_pool) + train_records = [dict(r) for r in train_pool[:train_n]] + overlap = {r['problem'] for r in train_records} & {r['problem'] for r in eval_records} + if overlap: + raise ValueError(f'fixed eval/train overlap detected: {len(overlap)} duplicated problems') + stats = {'raw_loaded': raw_n, 'numeric_dropped': dropped, + 'train_records': len(train_records), 'eval_records': len(eval_records)} + return train_records, eval_records, stats + + +class _ProblemPool: + """Cyclic draw source over the loaded problems. Each full pass reshuffles with + ``seed + epoch`` and bumps ``epoch`` (matching the old per-epoch reshuffle); the + initial pass keeps the loader's shuffled order. Draws never run out.""" + + def __init__(self, records: List[Dict[str, Any]], seed: int): + self._records = list(records) + self._seed = seed + self._cursor = 0 + self.epoch = 0 + self.baseline_cache: Dict[str, Dict[str, Any]] = {} # problem text -> frozen greedy roll + + def draw(self, k: int) -> List[Dict[str, Any]]: + """Return ``k`` DISTINCT problems (unique within this call, so one chunk never + processes the same problem twice even when the cursor wraps mid-draw). ``k`` is + always << pool size, so this terminates.""" + out: List[Dict[str, Any]] = [] + seen: set = set() + while len(out) < k: + if self._cursor >= len(self._records): + self.epoch += 1 + np.random.RandomState(self._seed + self.epoch).shuffle(self._records) + self._cursor = 0 + r = self._records[self._cursor] + self._cursor += 1 + if id(r) not in seen: + seen.add(id(r)) + out.append(r) + return out + + +def _select_balanced(buckets: Dict[str, List[Dict[str, Any]]], n_success: int, + n_fail: int, n_fail_loop: int) -> List[Dict[str, Any]]: + """Pick the chunk from the baselined buckets: ``n_fail`` base-fails (split toward + ``n_fail_loop`` loop-fails, best-effort) + ``n_success`` base-successes. If a bucket + is too thin to hit ``chunk_size`` the shortfall is topped up from leftovers (the + ratio then drifts, which the caller logs).""" + loop, wrong, succ = buckets['fail_loop'], buckets['fail_wrong'], buckets['success'] + take_loop = min(n_fail_loop, len(loop)) + take_wrong = min(n_fail - take_loop, len(wrong)) + take_loop = min(n_fail - take_wrong, len(loop)) # give loop the remainder if wrong is short + sel = loop[:take_loop] + wrong[:take_wrong] + succ[:n_success] + target = n_success + n_fail + if len(sel) < target: + used = {id(x) for x in sel} + leftover = [x for b in (loop, wrong, succ) for x in b if id(x) not in used] + sel += leftover[:target - len(sel)] + return sel + + +def _draw_chunk(pool: _ProblemPool, base_sampler, base_dp: int, args: argparse.Namespace + ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + """Draw one training chunk, running baseline rollout (Phase 1) on every drawn problem. + + With ``--balance`` off, draw ``chunk_size`` problems and return them. With it on, keep + drawing+baselining in ``chunk_size`` batches, bucketing by ``_baseline_class``, until the + target base fail:success mix is reachable or the draw budget is hit; then select a + balanced subset. Returns ``(chunk, stats)`` where stats records the realised mix.""" + if not args.balance: + chunk = pool.draw(args.chunk_size) + n_fresh = _baseline_rollout(base_sampler, chunk, base_dp, args, pool.baseline_cache) + return chunk, {'enabled': False, 'n_drawn': len(chunk), 'n_baseline_fresh': n_fresh} + + n_success = max(0, min(args.chunk_size, round(args.chunk_size * args.balance_success_frac))) + n_fail = args.chunk_size - n_success + n_fail_loop = round(n_fail * args.balance_loop_frac) + buckets: Dict[str, List[Dict[str, Any]]] = {'success': [], 'fail_loop': [], 'fail_wrong': []} + budget = args.chunk_size * args.balance_max_draws_mult + n_drawn, n_fresh = 0, 0 + seen: set = set() # dedupe across batches: the pool can re-serve a problem after a wrap + while n_drawn < budget: + if (len(buckets['success']) >= n_success + and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail): + break # enough of both classes buffered to satisfy the target split + batch = pool.draw(args.chunk_size) + n_fresh += _baseline_rollout(base_sampler, batch, base_dp, args, pool.baseline_cache) + n_drawn += len(batch) + for r in batch: + if id(r) in seen: + continue + seen.add(id(r)) + buckets[_baseline_class(r)].append(r) + + target_reached = (len(buckets['success']) >= n_success + and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail) + chunk = _select_balanced(buckets, n_success, n_fail, n_fail_loop) + sel_success = sum(1 for r in chunk if not r['_failed']) + stats = { + 'enabled': True, 'n_drawn': n_drawn, 'n_baseline_fresh': n_fresh, 'n_selected': len(chunk), + 'target_success': n_success, 'target_fail': n_fail, 'target_fail_loop': n_fail_loop, + 'selected_success': sel_success, 'selected_fail': len(chunk) - sel_success, + 'selected_fail_loop': sum(1 for r in chunk if _baseline_class(r) == 'fail_loop'), + 'selected_fail_wrong': sum(1 for r in chunk if _baseline_class(r) == 'fail_wrong'), + 'selected_success_frac': (sel_success / len(chunk)) if chunk else 0.0, + 'budget_hit': not target_reached, # stopped short of the target mix, not by choice + } + return chunk, stats + + +def process_chunk(base_sampler, skill_sampler, leak: LeakVerifier, + chunk: List[Dict[str, Any]], ci: int, base_dp: int, skill_dp: int, + args: argparse.Namespace, checker=None, + diag_cache: Optional[Dict[str, str]] = None + ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], List[Dict[str, Any]]]: + """base-solve -> rubric-check (view A) -> skill-gen -> leak-filter -> with-skill pass + -> GRPO advantages, for one chunk. + + Sequential (generate-one-chunk-train-one): generation and the trainer's weight + sync never overlap, so no lock is needed. ``base_sampler`` is frozen (never + synced); ``skill_sampler`` is synced by the trainer between chunks. + + ``chunk`` arrives ALREADY baselined by ``_draw_chunk`` (Phase 1 ran during the + balanced draw), so every problem carries ``_init``/``_failed``/``_baseline_pass``/ + ``_hard``/``_cands`` -- Phase 1 is not repeated here. + """ + # Phase 1 (base greedy solve) ran in _draw_chunk so the balancer could classify by + # outcome; every selected problem is processed (no difficulty gate, SEAM-style): the + # group-relative advantage (Phase 6) gives zero gradient to any problem whose skills + # all score alike, so GRPO's own group variance selects the informative problems. + hard = chunk + + # --- Phase 2: assign each problem's view, then rubric-check the view-A attempts so + # the skill model diagnoses from verified findings instead of guessing. View B is + # query-only and deliberately gets NO rubric (nothing to diagnose without an attempt). --- + for r in hard: + r['_view'] = _assign_view(r['problem'], args) + r['_rubric_diag'] = '' + _diagnose_views(checker, hard, args, diag_cache) + + # --- Phase 3: skill-gen (thinking ON), per-view prompt; re-sample empties. --- + flat: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] + if hard: + pending = list(hard) # problems still without any clean candidate + for _ in range(args.skill_retries + 1): + if not pending: + break + prompts = [_view_prompt(r, args) for r in pending] + sg_out = _run_samples(skill_sampler, prompts, args.n_skills, + args.skill_max_tokens, skill_dp, + temperature=args.skill_gen_temperature, + top_p=args.skill_gen_top_p, + top_k=args.skill_gen_top_k) + still = [] + for r, seqs in zip(pending, sg_out): + got = False + for s in seqs: + resp = _clean_text(getattr(s, 'decoded', '') or '') + block = _extract_skills_block(resp) + cand = {'skills': block or '', 'response': resp, 'parseable': bool(block), + 'view': r['_view'], 'leaked': None, 'leak_reason': '', + 'leak_source': '', 'with_pass': None, + 'reward': None, 'rolls': []} + r['_cands'].append(cand) + if block: + flat.append((r, cand)) + got = True + if not got: + still.append(r) # nothing parseable yet -> retry this problem + pending = still + + # --- Phase 4: leak filter via backup teacher (network only, no lock). VIEW A ONLY -- + # view B is query-only (no trajectory to leak from) and is left exactly like SEAM, which + # runs NO leak filter: its candidates skip the check and are treated as clean. To restore + # leak-checking on view B, drop the ``_view == 'A'`` guard below. --- + for r, c in flat: + if r.get('_view') != 'A': + c['leaked'], c['leak_reason'], c['leak_source'] = False, '', 'skipped_viewB' + flat_a = [(r, c) for r, c in flat if r.get('_view') == 'A'] + if flat_a: + details = leak.leak_batch( + [{'content': c['skills'], 'query': r['problem'], 'reference': r['reference_answer']} + for r, c in flat_a], max_workers=args.leak_workers) + for (r, c), d in zip(flat_a, details): + c['leaked'], c['leak_reason'], c['leak_source'] = bool(d.leaked), d.reason, d.source + + # --- Phase 5: with-skill GREEDY pass (T=0, M=1) on clean candidates. Binary reward + # R = answer CORRECT (deterministic, no pass@k noise), ABSOLUTE -- no baseline + # subtraction; the group mean in Phase 6 is the only baseline. Termination is NOT + # required (monitored only) -- see reflexion.md §7.6. --- + clean = [(r, c) for r, c in flat if c['leaked'] is False] + if clean: + ws_out = _run_samples(base_sampler, + [build_skill_solve_prompt(r['problem'], c['skills']) for r, c in clean], + 1, args.max_tokens, base_dp, temperature=0.0) + for (r, c), seqs in zip(clean, ws_out): + c['rolls'] = [_parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll()] + c['with_pass'] = 1.0 if c['rolls'][0]['correct'] else 0.0 + c['reward'] = c['with_pass'] # valid + clean + correct -> 1 + # Validity-in-reward (SEAM-style, --format-in-reward): every candidate that never reached + # the executor -- unparseable/impure format OR answer-leaked -- scores 0 and STILL joins its + # group, so its whole response (think tokens included) is trained DOWN. Off => those + # candidates are excluded, as before. + if args.format_in_reward: + for r in hard: + for c in r['_cands']: + if c['reward'] is None: + c['reward'] = 0.0 + + # --- Phase 6: group-relative GRPO advantage per problem-group. --- + _assign_advantages(hard, args) + + return ([_full_record(r, ci) for r in chunk], + _chunk_summary(chunk, ci, args), + _group_records(chunk, args)) + + +# --------------------------------------------------------------------------- +# Online RFT training +# --------------------------------------------------------------------------- +def _train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: + """Training sample = the exact skill-gen prompt for this record's view + the + generated (think + skills) response as the target; the GRPO advantage is attached + separately at forward_backward time. + + The prompt is rebuilt by ``_skillgen_messages`` (the same function used at generation), + so train/inference stay identical: view A replays problem + rubric findings, view B (and + no-failure view A) replays the query-only prompt. ``key_rounds`` selects the final + assistant turn (index ``len(msgs)``); the plain ``Template`` then masks the prompt and + trains the whole generated response (reasoning + ```` + skills) -- the key-round + prefix already excludes the prompt-provided ````, so no extra masking is needed.""" + msgs = _skillgen_messages(rec['problem'], rec.get('view', 'A'), rec.get('diagnosis', '')) + full = msgs + [{'role': 'assistant', 'content': rec['response']}] + return {'messages': full, 'user_data': {'key_rounds': [len(msgs)]}} + + +def _train_chunk(skill_model, ckpt: Optional[CheckpointEngineManager], + samples: List[Dict[str, Any]], args: argparse.Namespace) -> Dict[str, Any]: + """One on-policy GRPO optimizer update on THIS chunk's skill candidates, then sync weights. + + Sequential design (generate-one-chunk-train-one): the skills were sampled from the + current policy and trained immediately, so ``old_logps`` is omitted and the GRPO + ratio is ~1. All driver-side mini-batches accumulate into one optimizer step so the + whole rollout chunk stays under the same pre-update policy. The batch is padded to a + multiple of ``sft_batch_size`` with advantage-0 copies that contribute zero gradient. + """ + trajs = [_train_trajectory(rec) for rec in samples] + advs = [float(rec['advantage']) for rec in samples] + rem = (-len(trajs)) % args.sft_batch_size + if rem: + trajs += [trajs[-1]] * rem # zero-advantage pads -> forward only, no gradient + advs += [0.0] * rem + micro_batches = 0 + for i in range(0, len(trajs), args.sft_batch_size): + skill_model.forward_backward(inputs=trajs[i:i + args.sft_batch_size], + advantages=advs[i:i + args.sft_batch_size]) + micro_batches += 1 + skill_model.clip_grad_and_step() + if ckpt is not None: + ckpt.sync_weights(merge_and_sync=True) + metric = skill_model.calculate_metric(is_training=True) + return {'n_samples': len(samples), 'n_steps': 1, 'n_micro_batches': micro_batches, + 'advantages': [float(rec['advantage']) for rec in samples], + 'metric': {k: (float(v) if _is_num(v) else v) for k, v in (metric or {}).items()}} + + +def _is_num(v: Any) -> bool: + try: + float(v) + return True + except (TypeError, ValueError): + return False + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- +def _build_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--dataset', choices=('aops', 'math'), default='aops', + help='Problem source. aops (AI-MO competition problems) is much ' + 'harder than MATH, so the base fails more often -> more offensive ' + 'training signal after balanced sampling.') + p.add_argument('--n', type=int, default=2000, + help='Problems to load into the draw pool (cycled/reshuffled across ' + 'epochs; with --balance many more baseline rollouts than this ' + 'may run, but the pool size is fixed here).') + p.add_argument('--seed', type=int, default=42) + p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True, + help='Keep only answers that collapse to one integer/decimal/fraction, ' + 'matching SEAM numeric reward and avoiding non-scalar grading noise.') + p.add_argument('--eval-size', type=int, default=128, + help='Fixed holdout problems, sampled before the train pool after all ' + 'filters; set 0 to disable fixed eval.') + p.add_argument('--eval-every', type=int, default=10, + help='Run fixed holdout eval every N generation chunks when --eval-size > 0.') + p.add_argument('--chunk-size', type=int, default=16, + help='Problems per generation chunk (all sampler calls batched).') + # -- online baseline-balanced sampling (draw+baseline until the chunk hits the + # target base fail:success mix, so the offensive signal is not starved) -- + p.add_argument('--balance', action=argparse.BooleanOptionalAction, default=True, + help='Keep drawing+baselining problems until the chunk matches the ' + 'target base fail:success composition, then select a balanced ' + 'subset. --no-balance draws chunk_size problems directly.') + p.add_argument('--balance-success-frac', type=float, default=0.4, + help='Target fraction of the chunk that the base solves (base-success). ' + '0.4 => 3:2 fail:success; 0.2 => 4:1. The remainder are base-fail.') + p.add_argument('--balance-loop-frac', type=float, default=0.5, + help='Within the base-fail portion, SOFT target fraction of loop-fails ' + '(ran out of length / never terminated) vs non-loop wrong answers. ' + 'Best-effort only: the fail count is filled from whichever bucket ' + 'is available so a thin bucket never starves the chunk.') + p.add_argument('--balance-max-draws-mult', type=int, default=8, + help='Draw budget per chunk as a multiple of chunk_size; once this many ' + 'problems have been baselined the chunk is assembled from whatever ' + 'the buckets hold (ratio may drift; the actual mix is logged).') + p.add_argument('--n-skills', type=int, default=8, + help='Candidate skills generated per hard problem.') + p.add_argument('--view-b-frac', type=float, default=0.5, + help='Fraction of hard problems routed to view B (query-only, ' + 'deployment form); the rest go to view A (problem + attempt). ' + 'Each problem is assigned to EXACTLY ONE view.') + p.add_argument('--skill-retries', type=int, default=2, + help='Extra skill-gen rounds for a hard problem that yielded no ' + 'clean, parseable candidate (thinking-ON purity gate rejects).') + p.add_argument('--skill-gen-temperature', type=float, default=1.0, + help='Sampling temperature for skill-gen (BOTH views). >0 so the ' + 'n_skills candidates per problem are genuinely DIVERSE — a group ' + 'of near-duplicate skills gives GRPO no real good-vs-bad contrast.') + p.add_argument('--skill-gen-top-p', type=float, default=1.0, + help='top_p for skill-gen; 1.0 keeps the full tail for diversity.') + p.add_argument('--skill-gen-top-k', type=int, default=-1, + help='top_k for skill-gen; -1 disables truncation (max diversity). ' + 'A finite value only narrows the candidate pool.') + p.add_argument('--max-model-len', type=int, default=16384) + p.add_argument('--max-tokens', type=int, default=8192, + help='Max generated tokens for solve rollouts.') + p.add_argument('--skill-max-tokens', type=int, default=8192, + help='Max tokens for skill-gen (thinking ON: the model must close ' + ' within this budget or the candidate is dropped, so ' + 'leave ample room).') + p.add_argument('--leak-workers', type=int, default=16, + help='Parallel workers for the LeakVerifier backup judge (capped at 16 ' + 'to avoid the teacher API burst-rate limit; leak and rubric run in ' + 'separate phases so peak teacher concurrency is max(leak,rubric)).') + p.add_argument('--rubric-workers', type=int, default=16, + help='Parallel workers for the view-A rubric diagnose() calls ' + '(teacher-served; requires LLM_BACKUP_* env).') + # -- online GRPO (one on-policy update per generated chunk) -- + p.add_argument('--sft-batch-size', type=int, default=8, + help='Driver-side micro-batch size before the chunk-level optimizer step; ' + 'MUST be a multiple of the training dp size (sliced across dp ranks).') + p.add_argument('--grpo-epsilon', type=float, default=0.2, + help='PPO clip epsilon for GRPOLoss (ratio~1 on-policy, so rarely binds).') + p.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True, + help='Fold output validity into the reward (SEAM-style): unparseable/impure ' + 'or answer-leaked candidates score 0 and join their group to be trained ' + 'DOWN (the whole response, think tokens included). ' + '--no-format-in-reward keeps the reject-and-exclude gate.') + p.add_argument('--lr', type=float, default=1e-5) + p.add_argument('--max-train-rounds', type=int, default=200, + help='Cap on train rounds = trained chunks (also sizes the LR schedule).') + p.add_argument('--save-rounds', type=int, default=50) + p.add_argument('--trend-every', type=int, default=10, + help='Every N chunks, print a [trend] line contrasting the first N ' + 'vs the most recent N chunks (adoption + lift + pos/chunk) ' + 'so the training effect on fresh problems is visible at a glance.') + p.add_argument('--output-dir', default='./output/reflexion_skill_rft') + p.add_argument('--swanlab-project', default='twinkle', + help='swanlab project; logging is skipped when swanlab is not ' + 'installed or SWANLAB_MODE=disabled.') + p.add_argument('--swanlab-exp', default='', + help='swanlab experiment (run) name; empty = auto.') + return p.parse_args() + + +def _trend_line(hist: List[Dict[str, float]], window: int, rounds_done: int) -> Optional[str]: + """Contrast the FIRST ``window`` chunks with the most recent ``window`` chunks so + the online training effect on fresh, never-trained problems is glanceable: if RFT + is working, adoption and lift on recent chunks exceed the early baseline.""" + if len(hist) < 2 * window: + return None # need two non-overlapping windows for a clean before/after + base, rec = hist[:window], hist[-window:] + m = lambda xs, k: sum(h[k] for h in xs) / len(xs) + return (f'[trend] first {window} vs last {window} chunks | ' + f'adopt A {m(base,"aA"):.2f}->{m(rec,"aA"):.2f} ' + f'B {m(base,"aB"):.2f}->{m(rec,"aB"):.2f} | ' + f'lift {m(base,"lift"):+.3f}->{m(rec,"lift"):+.3f} | ' + f'pos/chunk {m(base,"pos"):.1f}->{m(rec,"pos"):.1f} | rounds={rounds_done}') + + +def _query_rows(full: List[Dict[str, Any]]) -> List[Tuple[float, float, float, int, str]]: + """Per hard problem that produced >=1 scored candidate: its no-skill baseline + pass@k, the BEST and MEAN with-skill pass@k over its N skill candidates, the scored + count, and the problem text. Drives both the per-query print and the swanlab passk/* + aggregates.""" + rows = [] + for rec in full: + if rec.get('record_type') != 'problem' or not rec.get('is_hard'): + continue + ps = [c['with_pass'] for c in rec.get('candidates', []) if c.get('with_pass') is not None] + if not ps: + continue + rows.append((rec['baseline_pass'], max(ps), sum(ps) / len(ps), len(ps), rec['problem'])) + return rows + + +def _clean_metric(metric: Optional[Dict[str, Any]]) -> Dict[str, float]: + """Numeric GRPO metrics for swanlab: collapse the duplicate per-group LR to a single + ``lr`` and drop non-numeric fields (e.g. 'total time elapse').""" + out: Dict[str, float] = {} + for k, v in (metric or {}).items(): + if not _is_num(v): + continue + if k.startswith('learning rate'): + if 'group 1' in k: + out['lr'] = float(v) + else: + out[k.replace(' ', '_')] = float(v) + return out + + +def _swan_metrics(summary: Dict[str, Any], log: Optional[Dict[str, Any]], + rows: List[Tuple[float, float, float, int, str]]) -> Dict[str, float]: + """Flat metric dict for swanlab = external reflexion metrics + (when this chunk was + trained) the GRPO built-in metric. acc/adopt/term are only emitted on chunks that had + hard problems, and passk/* only when scored candidates exist, so idle chunks don't dip + the charts to zero.""" + d: Dict[str, float] = { + 'gen/n_hard': summary['n_hard'], 'gen/n_clean': summary['n_clean'], + 'gen/n_leaked': summary['n_leaked'], 'gen/n_train_samples': summary['n_train_samples'], + 'gen/n_reward_pos': summary['n_reward_pos'], + 'gen/n_train_from_fail': summary['n_train_from_fail'], + 'gen/abs_adv_from_fail_frac': summary['abs_adv_from_fail_frac'], + } + bal = summary.get('balance') or {} + if bal.get('enabled'): + d.update({'balance/n_drawn': bal['n_drawn'], + 'balance/n_baseline_fresh': bal['n_baseline_fresh'], + 'balance/selected_success_frac': bal['selected_success_frac'], + 'balance/selected_fail_loop': bal['selected_fail_loop'], + 'balance/selected_fail_wrong': bal['selected_fail_wrong']}) + if summary['n_hard'] > 0: + d.update({ + 'acc/baseline_pass': summary['avg_baseline_pass_on_hard'], + 'acc/withskill_pass': summary['avg_withskill_pass'], + 'acc/lift': summary['avg_lift'], + 'adopt/A': summary['view_A']['adoption_rate'], + 'adopt/B': summary['view_B']['adoption_rate'], + 'term/withskill': summary['termination_rate_withskill'], + }) + if rows: + m = lambda i: sum(r[i] for r in rows) / len(rows) + d.update({'passk/baseline_mean': m(0), 'passk/bestN_mean': m(1), 'passk/avgN_mean': m(2)}) + if log: + d['train/n_steps'] = log['n_steps'] + if 'n_micro_batches' in log: + d['train/n_micro_batches'] = log['n_micro_batches'] + d.update({f'train/{k}': v for k, v in _clean_metric(log.get('metric')).items()}) + return d + + +def _prefix_metrics(metrics: Dict[str, float], prefix: str) -> Dict[str, float]: + return {f'{prefix}/{k}': v for k, v in metrics.items()} + + +def _greedy_eval_metrics(recs: List[Dict[str, Any]], ci: int, rounds: int + ) -> Tuple[Dict[str, Any], Dict[str, float]]: + """Aggregate the greedy holdout into SEAM ``mean@1`` metrics: overall + per-view acc, + the frozen-baseline acc, and their lift -- all single-sample-per-problem means (no + candidate averaging, no pass@k), so acc is directly comparable to SEAM's + ``val-core/math/acc/mean@1`` (correctness only; format/leak not gated).""" + def acc(rs: List[Dict[str, Any]]) -> float: + return sum(1 for x in rs if x['withskill_correct']) / len(rs) if rs else 0.0 + def bacc(rs: List[Dict[str, Any]]) -> float: + return sum(x['baseline_pass'] for x in rs) / len(rs) if rs else 0.0 + A = [x for x in recs if x['view'] == 'A'] + B = [x for x in recs if x['view'] == 'B'] + ws, base = acc(recs), bacc(recs) + summary = { + 'record_type': 'eval_summary', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, + 'n': len(recs), 'n_A': len(A), 'n_B': len(B), + 'acc_mean1': ws, 'baseline_acc_mean1': base, 'lift_mean1': ws - base, + 'acc_A_mean1': acc(A), 'acc_B_mean1': acc(B), + 'format_mean1': (sum(1 for x in recs if x['skill_parseable']) / len(recs)) if recs else 0.0, + 'term_mean1': (sum(1 for x in recs if x['withskill_terminated']) / len(recs)) if recs else 0.0, + } + metrics = { + 'core/math/acc/mean@1': ws, + 'core/math/baseline_acc/mean@1': base, + 'core/math/lift/mean@1': ws - base, + 'core/math/format/mean@1': summary['format_mean1'], + 'core/math/term/mean@1': summary['term_mean1'], + } + if A: + metrics['core/math/acc_A/mean@1'] = summary['acc_A_mean1'] + if B: + metrics['core/math/acc_B/mean@1'] = summary['acc_B_mean1'] + return summary, metrics + + +def _run_greedy_eval(base_sampler, skill_sampler, + eval_records: List[Dict[str, Any]], eval_cache: Dict[str, Dict[str, Any]], + ci: int, rounds: int, base_dp: int, skill_dp: int, + args: argparse.Namespace, checker=None, + diag_cache: Optional[Dict[str, str]] = None + ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], Dict[str, float]]: + """SEAM ``val-core/math/acc/mean@1`` analogue on the fixed holdout: ONE greedy skill per + problem (T=0) injected into ONE greedy base solve (T=0), so acc is a single-sample + pass@1 per problem averaged over problems. Each problem keeps its assigned view; view A + still gets the rubric process-check, view B stays query-only -- the mixed A/B acc is the + deployment number. No leak filter: like SEAM's val, acc scores correctness alone.""" + _baseline_rollout(base_sampler, eval_records, base_dp, args, eval_cache) # frozen greedy baseline + for r in eval_records: + r['_view'] = _assign_view(r['problem'], args) + r['_rubric_diag'] = '' + _diagnose_views(checker, eval_records, args, diag_cache) + sg_out = _run_samples(skill_sampler, [_view_prompt(r, args) for r in eval_records], + 1, args.skill_max_tokens, skill_dp, temperature=0.0) + skills = [] + for seqs in sg_out: + resp = _clean_text(getattr(seqs[0], 'decoded', '') or '') if seqs else '' + skills.append((_extract_skills_block(resp) or '', resp)) + ws_out = _run_samples(base_sampler, + [build_skill_solve_prompt(r['problem'], sk) for r, (sk, _) in zip(eval_records, skills)], + 1, args.max_tokens, base_dp, temperature=0.0) + recs = [] + for r, (sk, sresp), seqs in zip(eval_records, skills, ws_out): + roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + recs.append({ + 'record_type': 'eval_problem', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'view': r['_view'], 'rubric_diag': r.get('_rubric_diag', ''), + 'baseline_pass': r['_baseline_pass'], + 'skill': sk, 'skill_parseable': bool(sk), 'skill_response': sresp, + 'withskill_pred': roll['pred'], 'withskill_correct': roll['correct'], + 'withskill_terminated': roll['terminated'], 'withskill_stop_reason': roll['stop_reason'], + 'withskill_text': roll['text'], + }) + summary, metrics = _greedy_eval_metrics(recs, ci, rounds) + return recs, summary, metrics + + +def _validate_run_config(args: argparse.Namespace, records: List[Dict[str, Any]]) -> None: + """Fail fast on configs that would SILENTLY hang the online sampler: _ProblemPool.draw(k) + dedups within a call, so it never returns unless the pool holds >= chunk_size problems; + a zero draw budget or chunk size yields empty chunks that never advance ``rounds``.""" + if not records: + raise ValueError(f'loaded 0 {args.dataset} problems; check the dataset source') + if args.chunk_size < 1: + raise ValueError(f'--chunk-size must be >= 1 (got {args.chunk_size})') + if args.eval_size < 0: + raise ValueError(f'--eval-size must be >= 0 (got {args.eval_size})') + if args.eval_size > 0 and args.eval_every < 1: + raise ValueError(f'--eval-every must be >= 1 when eval is enabled (got {args.eval_every})') + if len(records) < args.chunk_size: + raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded problems ' + f'({len(records)}); raise --n or lower --chunk-size') + if args.balance_max_draws_mult < 1: + raise ValueError(f'--balance-max-draws-mult must be >= 1 (got {args.balance_max_draws_mult})') + if not 0.0 <= args.balance_success_frac <= 1.0: + raise ValueError(f'--balance-success-frac must be in [0, 1] (got {args.balance_success_frac})') + if not 0.0 <= args.balance_loop_frac <= 1.0: + raise ValueError(f'--balance-loop-frac must be in [0, 1] (got {args.balance_loop_frac})') + + +def main() -> None: + args = _build_args() + if args.sft_batch_size % TRAIN_DP != 0: + raise ValueError(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple ' + f'of the training dp size ({TRAIN_DP})') + # LR schedule now follows chunk-level optimizer updates, not driver micro-batches. + steps_per_round = 1 + records, eval_records, data_stats = _load_records(args) + _validate_run_config(args, records) + os.makedirs(args.output_dir, exist_ok=True) + data_path = os.path.join(args.output_dir, 'skill_dataset.jsonl') + gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') + eval_path = os.path.join(args.output_dir, 'eval_records.jsonl') + train_log_path = os.path.join(args.output_dir, 'train_log.jsonl') + + if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY')): + sys.stderr.write('[rft] WARNING: no LLM_BACKUP_API_KEY/OPENAI_API_KEY — ' + 'LeakVerifier will report no_llm and skip leak filtering\n') + + use_swan = swanlab is not None and os.environ.get('SWANLAB_MODE') != 'disabled' + if use_swan: + swanlab.init(project=args.swanlab_project, + experiment_name=(args.swanlab_exp or None), + config={'model': GEN_MODEL_ID, 'dataset': args.dataset, + 'n': len(records), 'eval_n': len(eval_records), + 'raw_loaded': data_stats['raw_loaded'], + 'numeric_only': args.numeric_only, + 'numeric_dropped': data_stats['numeric_dropped'], + 'n_skills': args.n_skills, 'view_b_frac': args.view_b_frac, + 'balance': args.balance, + 'balance_success_frac': args.balance_success_frac, + 'skill_gen_temp': args.skill_gen_temperature, + 'grpo_epsilon': args.grpo_epsilon, 'lr': args.lr}) + + # -- Device groups: train (FSDP2) + two independent vLLM samplers. -- + r0, r1, r2 = TRAIN_GPUS, TRAIN_GPUS + SKILL_SAMPLER_GPUS, NUM_GPUS + device_groups = [ + DeviceGroup(name='train', ranks=list(range(0, r0)), device_type='GPU'), + DeviceGroup(name='skill_sampler', ranks=list(range(r0, r1)), device_type='GPU'), + DeviceGroup(name='base_sampler', ranks=list(range(r1, r2)), device_type='GPU'), + ] + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, groups=device_groups, + lazy_collect=False) + + # -- Skill model: full-param FSDP2, GRPO policy update. -- + train_mesh = DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_DP, fsdp_size=TRAIN_FSDP) + skill_model = TransformersModel(model_id=GEN_MODEL_ID, device_mesh=train_mesh, + remote_group='train', + ddp_config={'find_unused_parameters': False}) + from twinkle.patch.no_split_modules import NoSplitModulesPatch + skill_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) + skill_model.set_template(Template, model_id=GEN_MODEL_ID, + enable_thinking=True, max_length=args.max_model_len, + truncation_strategy='delete') + skill_model.set_processor(InputProcessor, padding_free=False) + skill_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) + skill_model.set_optimizer('AdamW', lr=args.lr) + skill_model.set_lr_scheduler('CosineWarmupScheduler', num_warmup_steps=10, + num_training_steps=args.max_train_rounds * steps_per_round) + + # -- Two vLLM samplers: skill (synced) + base (frozen). -- + skill_dp, base_dp = SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS + skill_sampler = vLLMSampler( + model_id=GEN_MODEL_ID, + engine_args={'gpu_memory_utilization': GEN_GPU_MEM, + 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, + device_mesh=DeviceMesh.from_sizes(world_size=SKILL_SAMPLER_GPUS, dp_size=skill_dp), + remote_group='skill_sampler') + skill_sampler.set_template(Template, model_id=GEN_MODEL_ID, + enable_thinking=True, max_length=args.max_model_len) + base_sampler = vLLMSampler( + model_id=GEN_MODEL_ID, + engine_args={'gpu_memory_utilization': GEN_GPU_MEM, + 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, + device_mesh=DeviceMesh.from_sizes(world_size=BASE_SAMPLER_GPUS, dp_size=base_dp), + remote_group='base_sampler') + base_sampler.set_template(Template, model_id=GEN_MODEL_ID, + enable_thinking=True, max_length=args.max_model_len) + + ckpt = CheckpointEngineManager(model=skill_model, sampler=skill_sampler) + leak = LeakVerifier(sampler=None, answer_only=True) # flag ONLY the final answer (view A only; view B skips leak) + # leak = LeakVerifier(sampler=None, judge_system=_LEAK_JUDGE_SYSTEM) # stricter: also flag concrete intermediate key results + checker = _build_rubric_checker() # view-A process-check (teacher-only); None if no LLM backup + if checker is None: + sys.stderr.write('[rft] no LLM backup env -> view-A rubric process-check DISABLED ' + '(skill-gen diagnoses from the attempt alone)\n') + + sys.stderr.write(f'[rft] raw={data_stats["raw_loaded"]} numeric_drop={data_stats["numeric_dropped"]} ' + f'train={len(records)} eval={len(eval_records)} {args.dataset} problems; ' + f'train_gpus={TRAIN_GPUS} skill_dp={skill_dp} base_dp={base_dp}\n') + + # -- Sequential: generate one chunk, train on it, sync -> exact on-policy GRPO. + # Generation dominates wall-clock, so not overlapping training costs little, and + # it removes all producer/consumer concurrency (no thread, no lock). -- + cfg = {'record_type': 'config', 'model': GEN_MODEL_ID, 'dataset': args.dataset, + 'n': len(records), 'eval_n': len(eval_records), 'seed': args.seed, + 'numeric_only': args.numeric_only, + 'raw_loaded': data_stats['raw_loaded'], + 'numeric_dropped': data_stats['numeric_dropped'], + 'eval_every': args.eval_every, + 'n_skills': args.n_skills, + 'view_b_frac': args.view_b_frac, 'skill_retries': args.skill_retries, + 'balance': args.balance, 'balance_success_frac': args.balance_success_frac, + 'balance_loop_frac': args.balance_loop_frac, + 'balance_max_draws_mult': args.balance_max_draws_mult, + 'skill_gen_temp': args.skill_gen_temperature, + 'skill_gen_top_p': args.skill_gen_top_p, 'skill_gen_top_k': args.skill_gen_top_k, + 'reward': 'greedy_binary(correct)', 'advantage': 'group_relative', + 'format_in_reward': args.format_in_reward, + 'rubric_check': 'fixed_math_5crit(viewA)' if checker else 'disabled', + 'grpo_epsilon': args.grpo_epsilon, 'lr': args.lr, + 'max_train_rounds': args.max_train_rounds, 'started': int(time.time())} + hist: List[Dict[str, float]] = [] + rounds = 0 + pool = _ProblemPool(records, args.seed) + eval_cache: Dict[str, Dict[str, Any]] = {} + rubric_cache: Dict[str, str] = {} + eval_rubric_cache: Dict[str, str] = {} + with open(gen_path, 'w', encoding='utf-8') as gen_f, \ + open(eval_path, 'w', encoding='utf-8') as eval_f, \ + open(data_path, 'w', encoding='utf-8') as data_f, \ + open(train_log_path, 'w', encoding='utf-8') as tlog: + for f in (gen_f, eval_f, data_f, tlog): + f.write(json.dumps(cfg, ensure_ascii=False) + '\n') + f.flush() + gstep = 0 + # Each chunk is drawn fresh from the pool (which reshuffles + bumps epoch on every + # full pass) and RE-GENERATED with the current (improved) policy, so every chunk + # stays on-policy (no importance correction) -- the online analogue of SEAM's + # fixed-data epochs. With --balance, _draw_chunk keeps drawing+baselining until the + # base fail:success mix hits the target before this chunk is trained on. + while rounds < args.max_train_rounds: + chunk, balance = _draw_chunk(pool, base_sampler, base_dp, args) + full, summary, groups = process_chunk( + base_sampler, skill_sampler, leak, chunk, gstep, base_dp, skill_dp, + args, checker, rubric_cache) + summary['balance'] = balance + + log = None + if groups: # on-policy GRPO update on this chunk, then weights sync + log = _train_chunk(skill_model, ckpt, groups, args) + rounds += 1 + log.update({'record_type': 'train_round', 'round': rounds, + 'chunk': gstep, 'epoch': pool.epoch, 'ts': int(time.time())}) + tlog.write(json.dumps(log, ensure_ascii=False) + '\n') + tlog.flush() + if rounds % args.save_rounds == 0: + skill_model.save(f'skill-rft-{rounds}', output_dir=args.output_dir) + + summary['rounds_done'], summary['epoch'] = rounds, pool.epoch + for rec in full: + gen_f.write(json.dumps(rec, ensure_ascii=False) + '\n') + gen_f.write(json.dumps(summary, ensure_ascii=False) + '\n') + gen_f.flush() + for v in groups: + data_f.write(json.dumps(v, ensure_ascii=False) + '\n') + data_f.flush() + + sa, sb = summary['view_A'], summary['view_B'] + hist.append({'aA': sa['adoption_rate'], 'aB': sb['adoption_rate'], + 'lift': summary['avg_lift'], 'pos': summary['n_reward_pos']}) + bal_str = (f'bal {balance["selected_fail"]}f/{balance["selected_success"]}s ' + f'(loop {balance["selected_fail_loop"]} drew {balance["n_drawn"]}/' + f'fresh {balance["n_baseline_fresh"]}' + + ('!' if balance.get('budget_hit') else '') + ') ' + ) if balance.get('enabled') else '' + sys.stderr.write( + f'[gen] e{pool.epoch} g{gstep}: {bal_str}hard={summary["n_hard"]} ' + f'clean={summary["n_clean"]} train={summary["n_train_samples"]} ' + f'(fail {summary["n_train_from_fail"]} adv%{summary["abs_adv_from_fail_frac"]:.2f}) ' + f'acc={summary["avg_baseline_pass_on_hard"]:.2f}->{summary["avg_withskill_pass"]:.2f} ' + f'lift={summary["avg_lift"]:+.3f} ' + f'A[{sa["n_hard"]}h {sa["adoption_rate"]:.2f}] ' + f'B[{sb["n_hard"]}h {sb["adoption_rate"]:.2f}] ' + f'rounds={rounds}' + + (f' metric={log.get("metric")}' if log else '') + '\n') + # -- per-query passk (base vs best/avg of N skills) + swanlab metrics -- + rows = _query_rows(full) + for base_p, best_p, avg_p, nsc, prob in rows: + logger.info(f'[q] g{gstep} base={base_p:.2f} bestN={best_p:.2f} avgN={avg_p:.2f} ' + f'n={nsc} | {prob[:70].replace(chr(10), " ")}') + if use_swan: + swanlab.log(_swan_metrics(summary, log, rows), step=gstep) + + if eval_records and (gstep + 1) % args.eval_every == 0: + eval_recs, eval_summary, eval_metrics = _run_greedy_eval( + base_sampler, skill_sampler, eval_records, eval_cache, gstep, + rounds, base_dp, skill_dp, args, checker, eval_rubric_cache) + for rec in eval_recs: + eval_f.write(json.dumps(rec, ensure_ascii=False) + '\n') + eval_f.write(json.dumps(eval_summary, ensure_ascii=False) + '\n') + eval_f.flush() + if use_swan: + swanlab.log(_prefix_metrics(eval_metrics, 'eval'), step=gstep) + sys.stderr.write( + f'[eval] g{gstep}: n={eval_summary["n"]} mean@1 ' + f'acc={eval_summary["baseline_acc_mean1"]:.3f}->{eval_summary["acc_mean1"]:.3f} ' + f'lift={eval_summary["lift_mean1"]:+.3f} ' + f'A[{eval_summary["n_A"]} {eval_summary["acc_A_mean1"]:.3f}] ' + f'B[{eval_summary["n_B"]} {eval_summary["acc_B_mean1"]:.3f}] ' + f'fmt={eval_summary["format_mean1"]:.2f} rounds={rounds}\n') + + if (gstep + 1) % args.trend_every == 0: + tl = _trend_line(hist, args.trend_every, rounds) + if tl: + sys.stderr.write(tl + '\n') + gstep += 1 + + skill_model.save('skill-rft-final', output_dir=args.output_dir) + sys.stderr.write(f'[rft] done: {rounds} train rounds over {gstep} chunks / {pool.epoch} epochs; ' + f'data -> {data_path}\n') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/legacy/train_reflexion_skill_rft.sh b/cookbook/exp/legacy/train_reflexion_skill_rft.sh new file mode 100644 index 000000000..b5e0f1d82 --- /dev/null +++ b/cookbook/exp/legacy/train_reflexion_skill_rft.sh @@ -0,0 +1,35 @@ +#!/bin/bash +# RFT cold-start for the reflexion skill generator (see reflexion.md §6). +# GPUs: 8 — ranks 0-3 train (skill model, FSDP2), 4-5 skill sampler, 6-7 base sampler. +# Leak filtering uses the backup teacher API (no local judge): set LLM_BACKUP_*. + +set -euo pipefail + +export GEN_MODEL_ID=${GEN_MODEL_ID:-Qwen/Qwen3-4B} +# Local MATH copy (modelscope download cache). Override MATH_DATA_DIR if the +# cache hash dir changes or the data lives elsewhere. +export MATH_DATA_DIR=${MATH_DATA_DIR:-/mnt/workspace/.cache/modelscope/hub/datasets/downloads/extracted/0744cd2d347a7e8f85f7087d950b2ed38b626a5c808c5399e2d8a0923d42d013/MATH} +export LLM_BACKUP_API_KEY=${LLM_BACKUP_API_KEY:?set LLM_BACKUP_API_KEY for the leak judge} +export LLM_BACKUP_BASE_URL=${LLM_BACKUP_BASE_URL:-https://dashscope.aliyuncs.com/compatible-mode/v1} +export LLM_BACKUP_MODEL=${LLM_BACKUP_MODEL:-qwen3.7-max} + +python cookbook/exp/embedding/train_reflexion_skill_rft.py \ + --dataset aops \ + --n 5000 \ + --chunk-size 16 \ + --n-skills 8 \ + --view-b-frac 0.5 \ + --skill-retries 2 \ + --balance \ + --balance-success-frac 0.4 \ + --balance-loop-frac 0.5 \ + --balance-max-draws-mult 8 \ + --max-tokens 25000 \ + --max-model-len 30000 \ + --sft-batch-size 8 \ + --grpo-epsilon 0.2 \ + --lr 6e-6 \ + --max-train-rounds 1500 \ + --save-rounds 25 \ + --trend-every 10 \ + --output-dir ./output/reflexion_skill_rft diff --git a/cookbook/exp/legacy/train_reflexion_skill_seam.py b/cookbook/exp/legacy/train_reflexion_skill_seam.py new file mode 100644 index 000000000..757d9d915 --- /dev/null +++ b/cookbook/exp/legacy/train_reflexion_skill_seam.py @@ -0,0 +1,2022 @@ +"""Online GRPO RFT for the reflexion skill generator (self-contained, cached). + +Trains an INDEPENDENT skill model to write reusable skills that, injected into a +FROZEN base solver's system prompt, raise its accuracy. The base is never trained; +it only produces the reward. Per chunk: base greedy solve -> rubric process-check +(view A) -> skill-gen (thinking ON, N candidates) -> leak filter -> with-skill +greedy pass -> group-relative advantage -> ONE on-policy GRPO step -> sync weights. +Reward is deterministic (T=0, M=1) binary correctness; A = (R - mean) / (std + eps) +within each problem-group, so std=0 groups give no gradient (GRPO variance selects). + +Each problem is routed to EXACTLY ONE view: A = problem + rubric findings, B = +query only (deployment form). Skill-gen trains only the final structured guidance turn. + +Base greedy rollouts and rubric diagnoses are cached to disk (jsonl, md5-keyed) so +restarts skip them; skill-gen is on-policy and never cached. + +8 GPUs default for high-memory cards: rank 0 trains the actor, rank 1 hosts a +frozen reference model, ranks 2-3 skill_sampler (synced), ranks 4-7 base_sampler +(frozen). Override TRAIN_GPUS / REF_GPUS / SKILL_SAMPLER_GPUS / BASE_SAMPLER_GPUS +for other layouts. +Leak / rubric use the backup teacher API: set LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / +LLM_BACKUP_MODEL. + +Launch: + LLM_BACKUP_API_KEY=... python cookbook/exp/embedding/train_reflexion_skill.py \ + --dataset aops --n 5000 --chunk-size 16 --lr 6e-6 +""" +import argparse +import copy +import hashlib +import json +import math +import os +import re +import sys +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Optional, Set, Tuple + +import numpy as np + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.data_format import SamplingParams +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.model import TransformersModel +from twinkle.patch.no_split_modules import NoSplitModulesPatch +from twinkle.processor import InputProcessor +from twinkle.sampler import vLLMSampler +from twinkle.template import Template +from twinkle_agentic.verifier import RubricVerifier +from twinkle_agentic.verifier.rubric_verifier import RubricItem + +logger = get_logger() + +try: + import swanlab +except ImportError: + swanlab = None + +MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') +GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.8)) +GEN_TEMPERATURE = float(os.environ.get('GEN_TEMPERATURE', 0.6)) +GEN_TOP_P = float(os.environ.get('GEN_TOP_P', 0.95)) +AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') +MATH_DATASET_ID = os.environ.get('MATH_DATASET_ID', 'modelscope/competition_math') + +# GPU layout: train actor (synced to skill_sampler) + frozen ref model + skill_sampler + base_sampler. +# High-memory cards can keep Qwen3-4B actor/ref as one GPU each and spend more GPUs +# on vLLM data-parallel sampling. The base side is heavier here because every +# clean skill candidate is re-solved by the frozen base model, plus balance/eval baselines. +TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 2)) +REF_GPUS = int(os.environ.get('REF_GPUS', 2)) +SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) +BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 2)) +NUM_GPUS = TRAIN_GPUS + REF_GPUS + SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS +TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', min(1, TRAIN_GPUS))) +REF_FSDP = int(os.environ.get('REF_FSDP', min(1, REF_GPUS))) +if min(TRAIN_GPUS, REF_GPUS, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS, TRAIN_FSDP, REF_FSDP) < 1: + raise ValueError('TRAIN_GPUS, REF_GPUS, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS, TRAIN_FSDP and REF_FSDP must all be >= 1') +if TRAIN_GPUS % TRAIN_FSDP != 0: + raise ValueError(f'TRAIN_GPUS ({TRAIN_GPUS}) must be divisible by TRAIN_FSDP ({TRAIN_FSDP})') +if REF_GPUS % REF_FSDP != 0: + raise ValueError(f'REF_GPUS ({REF_GPUS}) must be divisible by REF_FSDP ({REF_FSDP})') +TRAIN_DP = TRAIN_GPUS // TRAIN_FSDP +REF_DP = REF_GPUS // REF_FSDP + + +# =========================================================================== +# Block A -- boxed extraction + answer grading +# =========================================================================== +_BOXED_RE = re.compile(r'\\boxed\s*\{') + + +def extract_boxed(text: str) -> Optional[str]: + """Last ``\\boxed{...}`` content, brace-balanced.""" + if not text: + return None + last = None + for m in _BOXED_RE.finditer(text): + depth, i = 1, m.end() + while i < len(text) and depth > 0: + depth += (text[i] == '{') - (text[i] == '}') + i += 1 + if depth == 0: + last = text[m.end():i - 1].strip() + return last + + +def normalize_answer(ans: str) -> str: + if not ans: + return '' + s = str(ans).strip() + m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) + if m: + return m.group(1) + s = s.strip('$').strip().replace('−', '-') + 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').replace(r'\tfrac', r'\frac') + s = s.replace(r'\,', '').replace(r'\;', '').replace(r'\!', '') + s = re.sub(r'\\(?:quad|qquad|\s)', '', s) + s = re.sub(r'\s+', '', s) + s = re.sub(r'\\sqrt([0-9A-Za-z])', r'\\sqrt{\1}', s) + s = re.sub(r'\\frac\{([^{}]+)\}([^{}\\])', r'\\frac{\1}{\2}', s) + s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) + s = re.sub(r'\{[a-zA-Z]+\}$', '', s) + s = re.sub(r'\^\\circ|\^\{\\circ\}|°|\\circ', 'deg', s) + s = s.replace(r'\minus{}', '-').replace(r'\minus', '-') + + def _frac_to_slash(mt): + text = mt.group(0) + pos = text.index('{') + 1 + depth, num_start = 1, pos + while depth > 0: + depth += (text[pos] == '{') - (text[pos] == '}') + pos += 1 + numer = text[num_start:pos - 1] + pos += 1 + den_start, depth = pos, 1 + while depth > 0: + depth += (text[pos] == '{') - (text[pos] == '}') + pos += 1 + return f'({numer})/({text[den_start:pos - 1]})' + + s = re.sub(r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', _frac_to_slash, s) + s = re.sub(r'(? bool: + try: + va, vb = float(a.replace('(', '').replace(')', '')), 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) + return va is not None and vb is not None and abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) + + +_MCQ_REF_RE = re.compile( + r'^\\(?:textbf|text|mathrm|mathbf)\{\(?([A-E])\)?\s*\}\s*(.+)$' + r'|^\(?([A-E])\)\s+(.+)$') +_VAR_PREFIX_RE = re.compile(r'^(?:[A-Za-z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)$', re.DOTALL) + + +def _split_mcq(ans: str): + s = ans.strip() + m = _MCQ_REF_RE.match(s) + if m: + return (m.group(1) or m.group(3)), ((m.group(2) or m.group(4) or '').strip() or None) + bl = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) + return (bl.group(1), None) if bl else (None, s or None) + + +def _strip_var_prefix(ans: str) -> str: + m = _VAR_PREFIX_RE.match((ans or '').strip()) + return m.group(1).strip() if m else (ans or '') + + +def _normalize_tuple(ans: str) -> str: + return re.sub(r'[\s()\[\]{}\\]', '', ans or '') + + +def _math_verify_equal(predicted: str, reference: str) -> bool: + try: + from math_verify import parse, verify + gold = parse(r'\boxed{' + reference + '}', parsing_timeout=1) + pred = parse(r'\boxed{' + predicted + '}', parsing_timeout=1) + return bool(gold and pred and verify(gold, pred, timeout_seconds=1)) + except Exception: + return False + + +def answers_match(predicted: str, reference: str) -> bool: + if not predicted or not reference: + return False + norm_p, norm_r = normalize_answer(predicted), normalize_answer(reference) + if norm_p == norm_r or norm_p.lower() == norm_r.lower() or _try_numeric_equal(norm_p, norm_r): + return True + stripped_p = normalize_answer(_strip_var_prefix(predicted)) + stripped_r = normalize_answer(_strip_var_prefix(reference)) + if stripped_p and stripped_r: + if (stripped_p == stripped_r or stripped_p.lower() == stripped_r.lower() + or _try_numeric_equal(stripped_p, stripped_r)): + return True + p_letter, p_value = _split_mcq(predicted) + r_letter, r_value = _split_mcq(reference) + if p_letter and r_letter and p_letter == r_letter: + return True + if (p_letter and p_value is None) and (r_letter and r_value is None): + return False + 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 and (p_val == r_val or p_val.lower() == r_val.lower() or _try_numeric_equal(p_val, r_val)): + return True + for left, right in ((norm_p, norm_r), (stripped_p, stripped_r)): + tl, tr = _normalize_tuple(left), _normalize_tuple(right) + if ',' in tl and tl == tr: + return True + if '=' in norm_r and '=' not in norm_p: + if any(part and (part == norm_p or _try_numeric_equal(part, norm_p)) for part in norm_r.split('=')): + return True + if '=' in norm_p and '=' not in norm_r: + if any(part and (part == norm_r or _try_numeric_equal(part, norm_r)) for part in norm_p.split('=')): + return True + return _math_verify_equal(stripped_p or norm_p, stripped_r or norm_r) + + +# =========================================================================== +# Block B -- prompts, skill parsing, batched sampling +# =========================================================================== +DIRECT_SYSTEM = ( + 'You are an expert competition mathematician. Solve the following problem ' + 'step by step. Provide your final answer inside \\boxed{}.') + +_SKILL_SOLVE_PREFIX = ( + DIRECT_SYSTEM + '\n\n' + 'Before you start, keep these reminders in mind to avoid common mistakes on this ' + 'type of problem:\n') +_SKILL_SOLVE_SUFFIX = '\nApply them where relevant, but rely on your own reasoning to reach the answer.' + + +def build_direct_prompt(problem: str) -> Dict[str, Any]: + return {'messages': [{'role': 'system', 'content': DIRECT_SYSTEM}, + {'role': 'user', 'content': problem}]} + + +def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: + skill = (skill or '').strip() + if not skill: + return build_direct_prompt(problem) + # Concatenation (not .format): DIRECT_SYSTEM/skill contain literal braces. + return {'messages': [ + {'role': 'system', 'content': _SKILL_SOLVE_PREFIX + skill + _SKILL_SOLVE_SUFFIX}, + {'role': 'user', 'content': problem}]} + + +# -- skill-gen prompts (view A: problem + rubric source; view B: query only) -- +# Kept deliberately short: this is the RL policy's system prompt, so over-specifying +# the output hurts convergence. The concrete output format is appended separately by +# ``_SKILL_OUTPUT`` so the task text stays format-agnostic. +SKILL_GEN_SYSTEM = ( + 'You are a math guidance writer. A process-check on a related problem hints at ' + 'likely mistakes. Write short reusable guidance for this and similar problems, ' + 'and note what to watch out for.\n') + +SKILL_GEN_SYSTEM_Q = ( + 'You are a math guidance writer. Write short reusable guidance for this and ' + 'similar problems.\n') + +_SKILL_OUTPUT = ( + 'Output only:\n\nYour reusable solving guidance here.\n') + +SKILL_GEN_USER_Q = ( + 'Problem:\n{problem}\n\n') + +SKILL_GEN_USER_RUBRIC = ( + 'Target problem:\n{problem}\n\n' + 'Problem used for the process check:\n{rubric_problem}\n\n' + 'Process check:\n' + '{diagnosis}\n\n') + + +def _rubric_has_fail(diagnosis: str) -> bool: + """View A uses the open-book rubric prompt (and, under --viewa-sft, SFT distillation) + IFF its rubric localised a failure. No ``[FAIL]`` => nothing to inject: generation + degrades to query-only and the problem is trained by GRPO exactly like view B. Single + source of truth so ``_skillgen_messages`` and ``_group_records`` never diverge.""" + return '[FAIL]' in (diagnosis or '') + + +def _skillgen_messages(problem: str, view: str, diagnosis: str, + rubric_problem: str = '') -> List[Dict[str, Any]]: + """Single source of truth for the skill-gen prompt (used at BOTH generation and + training so they never diverge). View A with a localisable failure uses the target + problem plus the rubric source problem and findings; view B -- or a view-A problem + whose rubric flagged nothing (no ``[FAIL]``) -- degrades to the query-only prompt.""" + if view == 'B' or not _rubric_has_fail(diagnosis): + return [{'role': 'system', 'content': SKILL_GEN_SYSTEM_Q + _SKILL_OUTPUT}, + {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=problem)}] + rubric_problem = rubric_problem or problem + return [{'role': 'system', 'content': SKILL_GEN_SYSTEM + _SKILL_OUTPUT}, + {'role': 'user', 'content': SKILL_GEN_USER_RUBRIC.format( + problem=problem, rubric_problem=rubric_problem, diagnosis=diagnosis)}] + + +def _assign_view(problem: str, args: argparse.Namespace) -> str: + h = int(hashlib.md5(f'{args.seed}:{problem}'.encode('utf-8')).hexdigest(), 16) + return 'B' if (h % 100000) / 100000.0 < args.view_b_frac else 'A' + + +def _view_prompt(r: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: + return {'messages': _skillgen_messages( + r['problem'], r['_view'], r.get('_rubric_diag', ''), r.get('_rubric_src', ''))} + + +_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') + + +def _clean_text(decoded: Optional[str]) -> str: + return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() + + +def _extract_tag_block(answer: str, tag: str, allow_empty: bool = False) -> Optional[str]: + low = answer.lower() + open_tag, close_tag = f'<{tag}>', f'' + s = low.rfind(open_tag) + if s < 0: + return None + inner = s + len(open_tag) + e = low.find(close_tag, inner) + if e < 0: + return None + block = answer[inner:e].strip() + block = re.sub(r'', '', block, flags=re.IGNORECASE).strip() + return block if (block or allow_empty) else None + + +def _extract_skill(text: str) -> Optional[str]: + """Parse skill-generation output: return the inner text of a non-empty ```` + block, or None. If a ```` marker is present, parse only the text after the + last one; otherwise parse the full response.""" + low = text.lower() + end_think = low.rfind('') + answer = text[end_think + len(''):] if end_think >= 0 else text + return _extract_tag_block(answer, 'skills') + + +def _parse_seq(seq, gold: str) -> Dict[str, Any]: + """Grade one sampled sequence into a rollout record.""" + text = _clean_text(getattr(seq, 'decoded', '') or '') + pred = extract_boxed(text) + correct = bool(pred) and answers_match(pred, gold) + terminated = getattr(seq, 'stop_reason', None) != 'length' + return {'pred': pred, 'correct': correct, 'terminated': terminated, + 'passed': bool(correct and terminated), + 'stop_reason': getattr(seq, 'stop_reason', None), + 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} + + +def _run_samples(sampler, prompts: List[Any], num_samples: int, max_tokens: int, + gen_dp: int, temperature: Optional[float] = None, + top_p: Optional[float] = None, top_k: Optional[int] = None) -> List[List[Any]]: + """One batched sampler call -> per-prompt list of raw sequences. vLLM dp needs + batch len >= dp, so pad the tail and slice back.""" + if not prompts: + return [] + params = SamplingParams( + max_tokens=max_tokens, + temperature=GEN_TEMPERATURE if temperature is None else temperature, + top_p=GEN_TOP_P if top_p is None else top_p, + num_samples=num_samples, **({} if top_k is None else {'top_k': top_k})) + padded = prompts + if gen_dp > 1 and 0 < len(prompts) < gen_dp: + padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] + responses = sampler.sample(padded, params)[:len(prompts)] + return [list(r.sequences) if (r and r.sequences) else [] for r in responses] + + +# =========================================================================== +# Block C -- data loading via twinkle.Dataset + numeric filtering +# =========================================================================== +def _boxed_batch(rows: Dict[str, List[Any]], dataset: str) -> Dict[str, List[Any]]: + """Batched (columnar) mapper for ``Dataset.map``: annotate each row with its boxed + ``reference_answer`` and a ``_keep`` flag (aops also needs the ``boxed`` metadata).""" + sols = rows['solution'] + metas = rows.get('metadata', [None] * len(sols)) + refs = [extract_boxed(s or '') for s in sols] + keep = [bool(ref) and (dataset != 'aops' or bool((meta or {}).get('boxed'))) + for ref, meta in zip(refs, metas)] + return {**rows, 'reference_answer': refs, '_keep': keep} + + +def load_problems(dataset: str, n: int, seed: int, num_proc: int = 0) -> List[Dict[str, Any]]: + """Load boxed-answer problems as ``{problem, reference_answer, level?}`` via + twinkle.Dataset (ModelScope hub), sampled to ``n`` (0 = all). Boxed extraction (regex + + brace scan over every solution) is parallelised by ``Dataset.map``/``.filter``; + ``num_proc`` defaults to all cores (set 1 to force serial).""" + ds_id = AOPS_DATASET_ID if dataset == 'aops' else MATH_DATASET_ID + ds = Dataset(DatasetMeta(dataset_id=f'ms://{ds_id}', split='train')) + nproc = num_proc if num_proc > 0 else min(32, os.cpu_count() or 1) + ds.map(lambda rows: _boxed_batch(rows, dataset), num_proc=nproc) + ds.filter(lambda row: row['_keep'], num_proc=nproc) + has_level = 'level' in ds.dataset.column_names + out = [{'data_id': f'{dataset}:{i}', 'problem': row['problem'], + 'reference_answer': row['reference_answer'], + **({'level': row['level']} if has_level and row.get('level') else {})} + for i, row in enumerate(ds.dataset)] + logger.info(f'[data] {dataset}: {len(out)} boxed problems ({nproc} procs)') + rng = np.random.RandomState(seed) + rng.shuffle(out) + return out[:n] if (n and n < len(out)) else out + + +# --------------------------------------------------------------------------- +# Bag-of-words neighbour pairing (cross-problem rubric transfer, --xproblem-rubric) +# --------------------------------------------------------------------------- +# Common English + math-scaffolding words that carry no problem-type signal. Kept small +# and deterministic on purpose (no external stopword list): what survives is the domain +# vocabulary ('triangle', 'prime', 'polynomial', ...) that actually defines the type. +_BOW_STOP = frozenset(""" +a an the of to in on at for and or but if is are be was were been being this that these those +with without into onto from by as it its their his her our your my we you they he she them +find compute determine calculate evaluate solve show prove given let suppose consider assume +what which when where how many much value values number numbers expression form terms term +such that then than so if only when each every all any some both one two three four five six +seven eight nine ten first second third last non over under about above below between +problem answer result equal equals sum difference product total following there here have has +had do does did can could will would should may might must not no yes if then else +""".split()) + +_WORD_RE = re.compile(r'[a-z]+') + + +def _stem(w: str) -> str: + """Crude suffix stripper so 'prime'/'primes', 'triangle'/'triangles' collapse to one + type token. Not linguistically correct -- just enough to merge the common plural/gerund + variants that otherwise split a type's vocabulary and starve the df filter.""" + if len(w) > 4 and w.endswith('ies'): + return w[:-3] + 'y' # properties -> property + if len(w) > 4 and w.endswith('es') and w[-3] in 'sxzh': + return w[:-2] # boxes -> box (keep primes -> prime below) + for suf in ('ing', 'ed', 's'): + if len(w) > len(suf) + 2 and w.endswith(suf): + return w[:-len(suf)] + return w + + +def _tokenize(problem: str) -> List[str]: + """Deterministic bag-of-words tokens for type matching: lowercase alphabetic words + (numbers dropped -- they are instance detail, not type), minus generic stopwords, then + stemmed, so only domain terms remain. Latex control words (frac, sqrt, ...) survive.""" + return [_stem(w) for w in _WORD_RE.findall((problem or '').lower()) + if len(w) > 2 and w not in _BOW_STOP] + + +class BagOfWordsIndex: + """TF-IDF cosine nearest-neighbour over problem statements. Sparse dict vectors + + an inverted index, so ``nearest`` scores only problems sharing a term (near-linear in + practice, no dense NxN). Deterministic; ties break on lower index for reproducibility. + + Answer-aware: ``nearest`` skips the same index, near-verbatim duplicates (cosine + >= ``sim_max``) and -- crucially for anti-leak -- any candidate whose answer equals the + query's, so a neighbour rubric can never hand over the query's own answer.""" + + def __init__(self, problems: List[str], answers: Optional[List[str]] = None, + min_df: int = 2, max_df_frac: float = 0.5): + self._toks = [_tokenize(p) for p in problems] + self._ans = [(_numeric_value(a) or (str(a).strip() if a else '')) for a in answers] \ + if answers is not None else [''] * len(problems) + n = len(self._toks) + df: Dict[str, int] = {} + for toks in self._toks: + for w in set(toks): + df[w] = df.get(w, 0) + 1 + max_df = max(min_df, int(max_df_frac * n)) + self._idf = {w: math.log((n + 1) / (c + 1)) + 1.0 + for w, c in df.items() if min_df <= c <= max_df} + self._vecs: List[Dict[str, float]] = [self._vectorize(t) for t in self._toks] + self._inverted: Dict[str, List[int]] = {} + for i, v in enumerate(self._vecs): + for w in v: + self._inverted.setdefault(w, []).append(i) + + def _vectorize(self, toks: List[str]) -> Dict[str, float]: + tf: Dict[str, float] = {} + for w in toks: + if w in self._idf: + tf[w] = tf.get(w, 0.0) + 1.0 + vec = {w: c * self._idf[w] for w, c in tf.items()} + norm = math.sqrt(sum(x * x for x in vec.values())) + return {w: x / norm for w, x in vec.items()} if norm > 0 else {} + + def nearest(self, i: int, sim_max: float = 0.98) -> Tuple[int, float]: + """(index, cosine) of the most similar *distinct* problem: not i, not a duplicate + (cosine < ``sim_max``) and not the same answer. (-1, 0.0) if none qualifies.""" + vi = self._vecs[i] + if not vi: + return -1, 0.0 + ai = self._ans[i] + scores: Dict[int, float] = {} + for w, xi in vi.items(): + for j in self._inverted.get(w, ()): + if j != i: + scores[j] = scores.get(j, 0.0) + xi * self._vecs[j].get(w, 0.0) + best_j, best_s = -1, 0.0 + for j, s in scores.items(): + if s >= sim_max or (ai and self._ans[j] == ai): + continue + if s > best_s or (s == best_s and (best_j < 0 or j < best_j)): + best_j, best_s = j, s + return best_j, best_s + + +def build_pairs(records: List[Dict[str, Any]], n: int, seed: int, sim_max: float = 0.98 + ) -> Tuple[List[Dict[str, Any]], Dict[str, Tuple[str, float]]]: + """Single-pass cross-problem pairing over the whole pool (one index build). + + Returns ``(subset, neighbour_map)`` where ``subset`` is the ``n`` problems with the + strongest qualifying neighbour (dense same-type pairs; n<=0/n>=len keeps all, shuffled) + and ``neighbour_map[q] = (p, cosine)`` gives each kept problem its analogue P. P is drawn + from the *full* pool (richer analogues) but never a duplicate or same-answer problem, so + P's rubric can transfer method without ever leaking Q's answer.""" + index = BagOfWordsIndex([r['problem'] for r in records], + [str(r.get('reference_answer', '')) for r in records]) + nbr = [index.nearest(i, sim_max) for i in range(len(records))] + order = sorted(range(len(records)), key=lambda i: (-nbr[i][1], i)) + keep = order[:n] if (0 < n < len(records)) else list(range(len(records))) + rng = np.random.RandomState(seed) + rng.shuffle(keep) + subset = [records[i] for i in keep] + neighbour_map = {records[i]['problem']: (records[nbr[i][0]]['problem'], nbr[i][1]) + for i in keep if nbr[i][0] >= 0} + return subset, neighbour_map + + +_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') + + +def _norm_num_text(num: str) -> str: + try: + f = float(num) + return str(int(f)) if f == int(f) else str(f) + except Exception: + return str(num).strip() + + +def _numeric_value(raw: Any) -> Optional[str]: + """Collapse an answer to a single int/decimal/fraction, or None.""" + if raw is None: + return None + s = str(raw).strip().strip('$').strip() + s = s.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac') + s = re.sub(r'\\!|\\,|\\;|\\ |\\left|\\right|\s', '', s) + for pat in (r'\\frac\{(-?\d+)\}\{(-?\d+)\}', r'(-?\d+)/(-?\d+)'): + m = re.fullmatch(pat, s) + if m: + a, b = int(m.group(1)), int(m.group(2)) + return _norm_num_text(str(a / b)) if b else None + return _norm_num_text(s) if _NUM_RE.fullmatch(s) else None + + +def _answer_leaked(skill: str, reference: str) -> bool: + """Audit whether a generated skill contains the final answer verbatim. This is NOT + a training filter: if the skill model derives an answer from the problem, that is a + legitimate answer-bearing skill under this experiment. The real leakage boundary is the + external API/rubric diagnosis, which is constrained by prompt to stay answer-free.""" + if not skill: + return False + for cand in {_numeric_value(reference), (str(reference).strip() or None)}: + if cand and re.search(r'(? Tuple[Set[str], Set[str]]: + """Read jsonl files and collect data_id/problem keys that must be excluded. + + The cold-start SFT builder writes stable ``data_id`` values. ``problem`` is kept as a + backward-compatible fallback for older jsonl files produced before data_id existed.""" + ids: Set[str] = set() + problems: Set[str] = set() + for raw_path in (paths_arg or '').split(','): + path = raw_path.strip() + if not path or not os.path.exists(path): + continue + with open(path, encoding='utf-8') as f: + for line in f: + if not line.strip(): + continue + row = json.loads(line) + if row.get('record_type') in {'config', 'summary'}: + continue + data_id = str(row.get('data_id') or '').strip() + problem = str(row.get('problem') or '').strip() + if data_id: + ids.add(data_id) + elif problem: + problems.add(problem) + return ids, problems + + +def _load_seam_parquet(path: str) -> List[Dict[str, Any]]: + """Read a SEAM ``build_aops_dataset.py`` parquet (VERL RLHF schema) into twinkle records, + PRESERVING file row order. ``problem <- extra_info.problem`` and + ``reference_answer <- reward_model.ground_truth``. No shuffle/filter: + the parquet is already SEAM's numeric-filtered, seed-42-shuffled, truncated split.""" + import pyarrow.parquet as pq + rows = pq.read_table(path).to_pylist() + out: List[Dict[str, Any]] = [] + for i, r in enumerate(rows): + ei = r.get('extra_info') or {} + rm = r.get('reward_model') or {} + problem = (ei.get('problem') or '').strip() + ref = rm.get('ground_truth') + if not problem or ref is None: + continue + out.append({'problem': problem, 'reference_answer': str(ref), + 'data_id': f"seam:{ei.get('split', '')}:{ei.get('index', i)}"}) + return out + + +def _load_records_from_seam(args: argparse.Namespace, seam_dir: str + ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], + Dict[str, Tuple[str, float]], Dict[str, str], Dict[str, int]]: + """Data entry that mirrors a SEAM run EXACTLY: read ``train.parquet``/``val.parquet`` from + ``seam_dir`` in file order, use ``val`` as the eval holdout, take the first ``--n`` train rows + (post ``--pool-offset``) with NO shuffle. ``--numeric-only``/``--eval-size``/internal shuffle + are bypassed (the parquet is already the authoritative split).""" + tp, vp = os.path.join(seam_dir, 'train.parquet'), os.path.join(seam_dir, 'val.parquet') + if not (os.path.exists(tp) and os.path.exists(vp)): + raise FileNotFoundError( + f'--seam-parquet-dir needs both train.parquet and val.parquet in {seam_dir}') + if args.xproblem_rubric: + raise ValueError('--xproblem-rubric is unsupported with --seam-parquet-dir ' + '(SEAM parquet carries no neighbour structure).') + pool = _load_seam_parquet(tp) # already SEAM-shuffled + truncated, in file order + eval_records = [dict(r) for r in _load_seam_parquet(vp)] # SEAM's exact val holdout + pool_offset = max(0, int(getattr(args, 'pool_offset', 0) or 0)) + if pool_offset: + if pool_offset >= len(pool): + raise ValueError(f'--pool-offset ({pool_offset}) leaves no train records ' + f'from SEAM train pool size {len(pool)}') + pool = pool[pool_offset:] + train_n = args.n if args.n > 0 else len(pool) + train_records = [dict(r) for r in pool[:train_n]] + if {r['problem'] for r in train_records} & {r['problem'] for r in eval_records}: + raise ValueError('eval/train overlap detected in SEAM parquet') + stats = {'raw_loaded': len(pool) + len(eval_records), 'numeric_dropped': 0, + 'excluded_records': 0, 'pool_offset': pool_offset, + 'train_records': len(train_records), 'eval_records': len(eval_records), + 'source': 'seam_parquet', 'seam_parquet_dir': seam_dir} + return train_records, eval_records, {}, {}, stats + + +def _load_records(args: argparse.Namespace + ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], + Dict[str, Tuple[str, float]], Dict[str, str], Dict[str, int]]: + """Load, numeric-filter, shuffle, split a fixed eval holdout, and (when --xproblem-rubric) + select a same-type-dense train subset with its neighbour map -- all in one pass. + Also returns ``pool_answers`` (every candidate neighbour P's true answer) so P's baseline + can be graded/cached correctly even when P is not itself a training problem.""" + seam_dir = (getattr(args, 'seam_parquet_dir', '') or '').strip() + if seam_dir: # read SEAM parquet in file order, bypassing load/filter/shuffle/split + return _load_records_from_seam(args, seam_dir) + # Load all when filtering or splitting (else the eval holdout could starve train). + load_n = 0 if (args.numeric_only or args.eval_size > 0) else args.n + records = load_problems(args.dataset, load_n, args.seed) + raw_n, dropped = len(records), 0 + if args.numeric_only: + records = [{**r, 'reference_answer': v} + for r, v in ((r, _numeric_value(r.get('reference_answer'))) for r in records) + if v is not None] + dropped = raw_n - len(records) + np.random.RandomState(args.seed).shuffle(records) + exclude_ids, exclude_problems = _load_excluded_records(getattr(args, 'exclude_data_ids', '')) + excluded = 0 + if exclude_ids or exclude_problems: + before = len(records) + records = [r for r in records + if str(r.get('data_id', '')) not in exclude_ids + and str(r.get('problem', '')).strip() not in exclude_problems] + excluded = before - len(records) + eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 + eval_records = [dict(r) for r in records[:eval_n]] + pool = records[eval_n:] + pool_offset = max(0, int(getattr(args, 'pool_offset', 0) or 0)) + if pool_offset: + if pool_offset >= len(pool): + raise ValueError(f'--pool-offset ({pool_offset}) leaves no train records from pool size {len(pool)}') + pool = pool[pool_offset:] + train_n = args.n if args.n > 0 else len(pool) + # Cross-problem rubric: pick the train_n problems with the strongest qualifying neighbour + # (dense same-type pairs) and build {Q -> (P, sim)} in one index pass. Otherwise keep the + # first train_n (already shuffled) with no neighbours. + if args.xproblem_rubric: + subset, neighbor_map = build_pairs(pool, train_n, args.seed) + train_records = [dict(r) for r in subset] + pool_answers = {r['problem']: str(r.get('reference_answer', '')) for r in pool} + else: + train_records, neighbor_map, pool_answers = [dict(r) for r in pool[:train_n]], {}, {} + if {r['problem'] for r in train_records} & {r['problem'] for r in eval_records}: + raise ValueError('eval/train overlap detected') + stats = {'raw_loaded': raw_n, 'numeric_dropped': dropped, + 'excluded_records': excluded, 'pool_offset': pool_offset, + 'train_records': len(train_records), 'eval_records': len(eval_records)} + return train_records, eval_records, neighbor_map, pool_answers, stats + + +# =========================================================================== +# Block D -- disk cache, problem pool, baseline rollout, rubric check +# =========================================================================== +class DiskCache: + """Append-only jsonl kv cache (md5 key -> value). Loads on init, appends on put. + Disabled instances always miss and never write.""" + + def __init__(self, path: str, enabled: bool = True): + self._mem: Dict[str, Any] = {} + self._fh = None + self._lock = threading.Lock() # base baseline is prefetched on a background thread + if not enabled: + return + if os.path.exists(path): + with open(path, encoding='utf-8') as f: + for line in f: + if line.strip(): + row = json.loads(line) + self._mem[row['key']] = row['value'] + self._fh = open(path, 'a', encoding='utf-8') + + @staticmethod + def key_for(*parts: str) -> str: + return hashlib.md5('\x1f'.join(parts).encode('utf-8')).hexdigest() + + def __contains__(self, key: str) -> bool: + with self._lock: + return key in self._mem + + def get(self, key: str) -> Any: + with self._lock: + return self._mem.get(key) + + def put(self, key: str, value: Any) -> None: + with self._lock: + self._mem[key] = value + if self._fh is not None: + self._fh.write(json.dumps({'key': key, 'value': value}, ensure_ascii=False) + '\n') + self._fh.flush() + + def close(self) -> None: + if self._fh is not None: + self._fh.close() + + +class _LockedSampler: + """Serialises every ``.sample()`` on a sampler behind one lock. The base sampler is + shared by the main thread (with-skill scoring, eval) and the baseline-prefetch thread; + ``sample`` is a slice_dp/flatten remote call whose collect is NOT safe to interleave + across two callers, so concurrent calls could mis-join sequences. The lock keeps base + calls serial (prefetch still overlaps the skill-gen phase, which uses skill_sampler).""" + + def __init__(self, sampler): + self._sampler = sampler + self._lock = threading.Lock() + + def sample(self, *args, **kwargs): + with self._lock: + return self._sampler.sample(*args, **kwargs) + + def __getattr__(self, name): + return getattr(self._sampler, name) + + +class ProblemPool: + """Cyclic draw source. Each full pass reshuffles with ``seed + epoch``; the initial + pass keeps the loader order. ``draw(k)`` returns k distinct problems (k << pool).""" + + def __init__(self, records: List[Dict[str, Any]], seed: int): + self._records = list(records) + self._seed, self._cursor, self.epoch = seed, 0, 0 + + def draw(self, k: int) -> List[Dict[str, Any]]: + out, seen = [], set() + while len(out) < k: + if self._cursor >= len(self._records): + self.epoch += 1 + np.random.RandomState(self._seed + self.epoch).shuffle(self._records) + self._cursor = 0 + r = self._records[self._cursor] + self._cursor += 1 + if id(r) not in seen: + seen.add(id(r)) + out.append(r) + return out + + def peek(self, k: int) -> List[Dict[str, Any]]: + """The next k distinct problems draw() would return, WITHOUT advancing state + (no cursor move, no reshuffle). Used to prefetch their baseline into base_cache + while the current chunk trains -- a wrong guess (cross-epoch reshuffle) only + misses the cache, never corrupts the draw.""" + out, seen, cur = [], set(), self._cursor + recs = self._records + while len(out) < k and cur < len(recs): # stop at epoch edge; don't simulate reshuffle + r = recs[cur] + cur += 1 + if id(r) not in seen: + seen.add(id(r)) + out.append(r) + return out + + +def _empty_roll() -> Dict[str, Any]: + return {'pred': '', 'correct': False, 'terminated': False, 'passed': False, + 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} + + +def _apply_baseline(r: Dict[str, Any], roll: Dict[str, Any]) -> None: + """Attach a greedy baseline roll and reset per-chunk working state.""" + r['_baseline_rolls'], r['_cands'], r['_init'] = [roll], [], [roll] + r['_failed'] = not roll['correct'] + r['_baseline_pass'] = 1.0 if roll['correct'] else 0.0 + r['_hard'] = True # process every problem; group variance selects (SEAM-style) + + +def baseline_rollout(base_sampler, problems: List[Dict[str, Any]], base_dp: int, + args: argparse.Namespace, cache: DiskCache) -> int: + """Base solves each problem greedily once (T=0, M=1), disk-cached by problem text. + The base is frozen + greedy so a cache hit is exact. Returns fresh (miss) count.""" + todo = [r for r in problems if DiskCache.key_for(r['problem']) not in cache] + if todo: + out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], + 1, args.max_tokens, base_dp, temperature=0.0) + for r, seqs in zip(todo, out): + roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + cache.put(DiskCache.key_for(r['problem']), roll) + for r in problems: + _apply_baseline(r, cache.get(DiskCache.key_for(r['problem']))) + return len(todo) + + +# -- rubric process-check (view A): teacher diagnoses the base's attempt -- +_RFT_DIAG_SYSTEM = """\ +You are a strategy-level process checker for a math solution attempt. You are given a +math problem, a rubric, and one attempted solution segment. Decide PASS or FAIL for each +criterion, and write the diagnosis so it can become useful reusable guidance for solving +similar problems without seeing this segment. + +Output STRICT JSON (no prose outside it) with this shape: +{ + "items": [ + {"index": 1, "verdict": "PASS", "reason": "", + "fix": ""}, + {"index": 2, "verdict": "FAIL", "reason": "", + "fix": ""} + ], + "overall": "OK" | "ISSUES", + "summary": "" +} + +Rules: +- Judge every criterion independently; a [Hard Rule] is FAIL unless clearly satisfied. +- Judge ONLY what is observable in THIS segment. Ignore hidden or + content for output-format criteria. +- The API diagnosis is an external teacher signal, so it must stay answer-free. +- Prefer diagnosis that transfers to view-B skill generation: name the route choice, + structural observation, missing check, or length-control habit that a solver should + remember before solving a similar problem. +- For PASS items, leave "fix" as "". +- For FAIL items, describe the process problem at strategy level: unsuitable method, + missed structure, invalid transformation, missing constraint check, redundant cases, + off-track approach, contradiction, or inefficient/unfinished reasoning. +- A fix may suggest the LOCAL correction direction, such as identify the key structure, + verify constraints, preserve equivalence, reduce redundant cases, or choose a more + direct route. Do not carry out the correction. +- Never reveal the final answer, a corrected value/expression, an option label, or a + step-by-step solution that would let another model copy the solve. +- If the segment contains a process note saying it was cut off before a final boxed + answer, mark the length-budget criterion as FAIL and suggest a method-level way to + finish faster. +- Keep every "reason" and "fix" concise: one short sentence each. +- "overall" is "OK" only if NO criterion is FAIL. +- Output only the JSON object.""" + +_RFT_DIAG_USER = """\ +## Task / query (context) +{query} + +## Rubric +{rubric} + +## Segment +{segment} + +Now output the diagnostic JSON object.""" + +_MATH_RUBRIC = [ + ('The attempt chooses a method suitable for the problem structure', False), + ('The attempt identifies the key constraint, invariant, or quantity before computing', False), + ('Algebraic and logical transformations preserve validity at each step', True), + ('The attempt checks required constraints, domains, boundary cases, or validity conditions', False), + ('The attempt avoids redundant casework, looping, or re-deriving known facts', False), + ('The attempt reaches a final boxed answer within the length budget', False), + ('The approach stays focused on the actual question asked', False), +] + +# Bump when the rubric criteria or the diagnosis prompt change so the disk-cached +# diagnoses written under an older rubric are not silently reused. +_RUBRIC_VERSION = 'rubric_v5_viewb_strategy' + + +class _RftRubricVerifier(RubricVerifier): + def _diagnose_trajectory(self, query: str, rubric_block: str, segment_text: str) -> dict: + return {'messages': [ + {'role': 'system', 'content': _RFT_DIAG_SYSTEM}, + {'role': 'user', 'content': _RFT_DIAG_USER.format( + query=query, rubric=rubric_block, segment=segment_text)}]} + + +def build_rubric_checker() -> Optional[RubricVerifier]: + """Fixed math-process rubric verifier (teacher-served). None if no LLM backup env.""" + if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') + or os.environ.get('OPENAI_API_KEY')): + return None + return _RftRubricVerifier( + fixed_rubric=[RubricItem(t, is_hard=h) for t, h in _MATH_RUBRIC], gate=True) + + +def _format_diagnosis(detail) -> str: + """One line per criterion (PASS/FAIL + reason + fix) then a summary.""" + rub = detail.rubric + lines = [] + for it in detail.items: + text = rub[it.index - 1].text if 0 < it.index <= len(rub) else f'criterion {it.index}' + if it.verdict: + lines.append(f'- [PASS] {text}') + else: + tail = f': {it.reason}' if it.reason else '' + tail += f' (fix: {it.fix})' if it.fix else '' + lines.append(f'- [FAIL] {text}{tail}') + if detail.summary: + lines.append(f'Summary: {detail.summary}') + return '\n'.join(lines) + + +def diagnose_views(checker, problems: List[Dict[str, Any]], args: argparse.Namespace, + cache: DiskCache) -> None: + """Rubric-check every view-A problem's greedy attempt in parallel (disk-cached by + problem + attempt), stashing the formatted findings on ``r['_rubric_diag']``.""" + targets = [r for r in problems if r.get('_view') == 'A'] + if not checker or not targets: + return + + def _key(r: Dict[str, Any]) -> str: + init = r.get('_init', [{}])[0] + term = 'L' if (init.get('stop_reason') == 'length' or not init.get('terminated')) else 'T' + return DiskCache.key_for(r['problem'], init.get('text', ''), _RUBRIC_VERSION, term) + + pending = [] + for r in targets: + key = _key(r) + if key in cache: + r['_rubric_diag'] = cache.get(key) + else: + pending.append((r, key)) + if not pending: + return + + def _run(item): + r, key = item + init = r['_init'][0] + seg_text = init['text'] + if init.get('stop_reason') == 'length' or not init.get('terminated'): + seg_text += ('\n\n[Process note: this attempt was cut off at the token budget ' + 'and never produced a final \\boxed{} answer.]') + seg = {'messages': [{'role': 'user', 'content': r['problem']}, + {'role': 'assistant', 'content': seg_text}]} + attempts = max(1, args.rubric_retries + 1) + for attempt in range(attempts): + try: + return r, key, _format_diagnosis(checker.diagnose(seg, query=r['problem'])) + except Exception as exc: # teacher hiccup -> retry, then no-diagnosis prompt (not cached) + if attempt + 1 < attempts: + logger.warning(f'[rubric] diagnose error: {exc}; retry {attempt + 1}/{args.rubric_retries}') + time.sleep(min(2.0, 0.5 * (2 ** attempt))) + continue + logger.warning(f'[rubric] diagnose error: {exc}; giving up after {attempts} attempts') + return r, key, None + + workers = max(1, min(args.rubric_workers, len(pending))) + with ThreadPoolExecutor(max_workers=workers) as ex: + for r, key, diag in ex.map(_run, pending): + r['_rubric_diag'] = diag or '' + if diag is not None: + cache.put(key, diag) + + +# =========================================================================== +# Block E -- chunk draw, generation pipeline, record building +# =========================================================================== +def _baseline_class(r: Dict[str, Any]) -> str: + """success | fail_loop (out of length / never terminated) | fail_wrong.""" + roll = r['_init'][0] + if roll['correct']: + return 'success' + return 'fail_loop' if (roll['stop_reason'] == 'length' or not roll['terminated']) else 'fail_wrong' + + +def _select_balanced(buckets: Dict[str, List[Dict[str, Any]]], n_success: int, + n_fail: int, n_fail_loop: int) -> List[Dict[str, Any]]: + """Pick n_fail base-fails (toward n_fail_loop loop-fails, best-effort) + n_success + base-successes; top up any shortfall from leftovers.""" + loop, wrong, succ = buckets['fail_loop'], buckets['fail_wrong'], buckets['success'] + take_loop = min(n_fail_loop, len(loop)) + take_wrong = min(n_fail - take_loop, len(wrong)) + take_loop = min(n_fail - take_wrong, len(loop)) + sel = loop[:take_loop] + wrong[:take_wrong] + succ[:n_success] + target = n_success + n_fail + if len(sel) < target: + used = {id(x) for x in sel} + sel += [x for b in (loop, wrong, succ) for x in b if id(x) not in used][:target - len(sel)] + return sel + + +def draw_chunk(pool: ProblemPool, base_sampler, base_dp: int, args: argparse.Namespace, + cache: DiskCache) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + """Draw one chunk, baselining every drawn problem. With ``--balance``, keep + drawing+baselining until the target base fail:success mix is reachable (or the budget + is hit), then select a balanced subset.""" + if not args.balance: + chunk = pool.draw(args.chunk_size) + fresh = baseline_rollout(base_sampler, chunk, base_dp, args, cache) + return chunk, {'enabled': False, 'n_drawn': len(chunk), 'n_baseline_fresh': fresh} + + n_success = max(0, min(args.chunk_size, round(args.chunk_size * args.balance_success_frac))) + n_fail = args.chunk_size - n_success + n_fail_loop = round(n_fail * args.balance_loop_frac) + buckets: Dict[str, List[Dict[str, Any]]] = {'success': [], 'fail_loop': [], 'fail_wrong': []} + budget, n_drawn, n_fresh, seen = args.chunk_size * args.balance_max_draws_mult, 0, 0, set() + while n_drawn < budget: + if (len(buckets['success']) >= n_success + and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail): + break + batch = pool.draw(args.chunk_size) + n_fresh += baseline_rollout(base_sampler, batch, base_dp, args, cache) + n_drawn += len(batch) + for r in batch: + if id(r) not in seen: + seen.add(id(r)) + buckets[_baseline_class(r)].append(r) + + reached = (len(buckets['success']) >= n_success + and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail) + chunk = _select_balanced(buckets, n_success, n_fail, n_fail_loop) + sel_success = sum(1 for r in chunk if not r['_failed']) + stats = { + 'enabled': True, 'n_drawn': n_drawn, 'n_baseline_fresh': n_fresh, 'n_selected': len(chunk), + 'target_success': n_success, 'target_fail': n_fail, 'target_fail_loop': n_fail_loop, + 'selected_success': sel_success, 'selected_fail': len(chunk) - sel_success, + 'selected_fail_loop': sum(1 for r in chunk if _baseline_class(r) == 'fail_loop'), + 'selected_fail_wrong': sum(1 for r in chunk if _baseline_class(r) == 'fail_wrong'), + 'selected_success_frac': (sel_success / len(chunk)) if chunk else 0.0, + 'budget_hit': not reached, + } + return chunk, stats + + +def _assign_advantages(hard: List[Dict[str, Any]], args: argparse.Namespace) -> None: + """Group-relative advantage over each problem's scored candidates using the greedy + binary reward R in {0,1}: ``A = (R - mean) / (std + eps)``. std==0 groups get no + gradient -- GRPO's variance selects informative problems (no explicit difficulty gate). + A symmetric clip limits both positive and negative outliers from sparse 1-vs-many groups.""" + eps = 1e-6 + adv_clip = abs(float(getattr(args, 'adv_clip', 0.0) or 0.0)) + for r in hard: + for c in r['_cands']: + c['advantage'], c['grpo_adv'], c['kept'] = 0.0, 0.0, False + cs = [c for c in r['_cands'] if c.get('reward') is not None] + if len(cs) < 2: + continue + rewards = [c['reward'] for c in cs] + mean_r = sum(rewards) / len(rewards) + std = (sum((x - mean_r) ** 2 for x in rewards) / len(rewards)) ** 0.5 + if std < 1e-9: + continue + for c in cs: + raw_adv = (c['reward'] - mean_r) / (std + eps) + adv = max(-adv_clip, min(adv_clip, raw_adv)) if adv_clip > 0 else raw_adv + c['advantage'], c['grpo_adv'], c['kept'] = adv, adv, c['reward'] > mean_r + + +def _best_sft_candidate(r: Dict[str, Any], args: argparse.Namespace) -> Optional[Dict[str, Any]]: + """Pick ONE view-A candidate to distill (online context distillation). PREFER the + executor-verified PASSING skills (reward==1); if NONE passed -- common on the hard + problems that are exactly the cases worth distilling -- FALL BACK to any parseable + open-book skill regardless of the executor outcome. Answer-bearing skills produced by + the skill model itself are allowed here; only the external API/rubric diagnosis must be + answer-free. Within the chosen tier, take the one whose skill length is CLOSEST to + ``--sft-target-len`` -- an empirically high-pass-rate length (~500-600 chars in this + run) -- breaking ties by the fewest executor solve tokens. Targeting a length (rather + than the minimum) avoids a distillation feedback loop that would otherwise drive + rollouts ever shorter. None only when no parseable candidate exists at all.""" + eligible = [c for c in r['_cands'] if c.get('parseable') and c.get('skills')] + if not eligible: + return None + passing = [c for c in eligible if c.get('reward') == 1.0] + cs = passing or eligible + target = int(getattr(args, 'sft_target_len', 550) or 550) + + def _solve_tokens(c: Dict[str, Any]) -> int: + rolls = c.get('rolls') or [] + return rolls[0].get('gen_tokens', 1 << 30) if rolls else (1 << 30) + + return min(cs, key=lambda c: (abs(len(c['skills']) - target), _solve_tokens(c))) + + +def apply_neighbor_rubric(base_sampler, chunk: List[Dict[str, Any]], + neighbor_map: Dict[str, Tuple[str, float]], + pool_answers: Dict[str, str], base_dp: int, + args: argparse.Namespace, checker, + base_cache: DiskCache, rubric_cache: DiskCache) -> None: + """Cross-problem rubric (--xproblem-rubric): for every view-A problem Q, replace its own + rubric with the rubric of its bag-of-words neighbour P (Q keeps being the solved/scored + problem). P is baselined + diagnosed here (both disk-cached) as a stub carrying P's REAL + answer, so P's baseline grades correctly and legitimately shares the baseline cache with + P-as-training-problem. P's findings are copied onto Q with the neighbour text + similarity + for audit. A P that can't be diagnosed -> Q degrades to view B. Since P's answer differs + from Q's (guaranteed at pairing time), any answer P's rubric leaks is useless for Q.""" + targets = [r for r in chunk if r.get('_view') == 'A' and neighbor_map.get(r['problem'])] + if not targets: + return + stubs, by_problem = [], {} + for r in targets: + p, _ = neighbor_map[r['problem']] + if p not in by_problem: + stub = {'problem': p, 'reference_answer': pool_answers.get(p, ''), '_view': 'A'} + by_problem[p] = stub + stubs.append(stub) + baseline_rollout(base_sampler, stubs, base_dp, args, base_cache) + diagnose_views(checker, stubs, args, rubric_cache) + for r in targets: + p, sim = neighbor_map[r['problem']] + r['_rubric_diag'] = by_problem[p].get('_rubric_diag', '') + r['_rubric_src'], r['_neighbor_sim'] = p, sim + + +def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], + ci: int, base_dp: int, skill_dp: int, args: argparse.Namespace, + checker, rubric_cache: DiskCache, base_cache: DiskCache = None, + neighbor_map: Optional[Dict[str, Tuple[str, float]]] = None, + pool_answers: Optional[Dict[str, str]] = None + ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], List[Dict[str, Any]]]: + """view assign -> rubric-check (view A) -> skill-gen -> answer audit -> with-skill + greedy pass -> GRPO advantages. ``chunk`` arrives already baselined by draw_chunk. + With --xproblem-rubric, view A gets its NEIGHBOUR's rubric (see apply_neighbor_rubric).""" + hard = chunk + for r in hard: + r['_view'], r['_rubric_diag'] = _assign_view(r['problem'], args), '' + if args.xproblem_rubric and neighbor_map: + apply_neighbor_rubric(base_sampler, hard, neighbor_map, pool_answers or {}, base_dp, + args, checker, base_cache, rubric_cache) + else: + diagnose_views(checker, hard, args, rubric_cache) + + # skill-gen (thinking OFF), per-view prompt; re-sample problems with no clean candidate. + # View-A problems that cannot yield a training record (no [FAIL] rubric, or rubric + # leaked the answer) are dropped from training entirely -- skip their generation. + flat: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] + pending = [r for r in hard if not _viewa_dropped(r, args)] + for _ in range(args.skill_retries + 1): + if not pending: + break + sg_out = _run_samples(skill_sampler, [_view_prompt(r, args) for r in pending], + args.n_skills, args.skill_max_tokens, skill_dp, + temperature=args.skill_gen_temperature, + top_p=args.skill_gen_top_p, top_k=args.skill_gen_top_k) + still = [] + for r, seqs in zip(pending, sg_out): + got = False + for s in seqs: + resp = _clean_text(getattr(s, 'decoded', '') or '') + block = _extract_skill(resp) or '' + cand = {'skills': block, 'response': resp, 'parseable': bool(block), + 'view': r['_view'], 'leaked': None, 'leak_reason': '', + 'leak_source': '', 'with_pass': None, 'reward': None, 'rolls': [], + 'skillgen_stop': getattr(s, 'stop_reason', None), + 'skillgen_tokens': len(getattr(s, 'tokens', None) or [])} + r['_cands'].append(cand) + if block: + flat.append((r, cand)) + got = True + if not got: + still.append(r) + pending = still + + # Answer-bearing skill audit: deterministic verbatim-answer check only (no LLM). This + # is observability only; it records metrics for swanlab/jsonl, but does not block + # scoring, reward, advantage assignment, SFT candidate selection, or GRPO training. + for r, c in flat: + leaked = _answer_leaked(c['skills'], r['reference_answer']) + c['leaked'] = leaked + c['leak_reason'] = 'answer_verbatim' if leaked else '' + c['leak_source'] = 'deterministic' + + # with-skill greedy pass (T=0, M=1); reward = correct, absolute (group mean is baseline). + scored_inputs = flat + if scored_inputs: + ws_out = _run_samples(base_sampler, + [build_skill_solve_prompt(r['problem'], c['skills']) for r, c in scored_inputs], + 1, args.max_tokens, base_dp, temperature=0.0) + for (r, c), seqs in zip(scored_inputs, ws_out): + c['rolls'] = [_parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll()] + c['with_pass'] = 1.0 if c['rolls'][0]['correct'] else 0.0 + c['reward'] = c['with_pass'] + if args.format_in_reward: # unparseable candidates score 0 and still join the group + for r in hard: + for c in r['_cands']: + if c['reward'] is None: + c['reward'] = 0.0 + + _assign_advantages(hard, args) + return ([_full_record(r, ci) for r in chunk], _chunk_summary(chunk, ci, args), + _group_records(chunk, args)) + + +def _roll(x: Dict[str, Any]) -> Dict[str, Any]: + return {k: x[k] for k in ('pred', 'correct', 'terminated', 'passed', + 'stop_reason', 'gen_tokens', 'text')} + + +def _is_trainable(c: Dict[str, Any], args: argparse.Namespace) -> bool: + """Reaches the GRPO update iff advantage is non-zero. Leak flags are audit-only.""" + adv_nz = abs(c.get('advantage') or 0.0) > 1e-9 + if args.format_in_reward: + return adv_nz + return c.get('with_pass') is not None and adv_nz + + +def _full_record(r: Dict[str, Any], ci: int) -> Dict[str, Any]: + """Complete per-problem trace: init attempt, baseline, and all candidates.""" + init = r['_init'][0] + return { + 'record_type': 'problem', 'chunk': ci, 'problem': r['problem'], + 'reference_answer': r['reference_answer'], 'level': r.get('level', ''), + 'failed_first_try': r['_failed'], + 'init_attempt': {'text': init['text'], 'pred': init['pred'], 'correct': init['correct'], + 'terminated': init['terminated'], 'stop_reason': init['stop_reason'], + 'gen_tokens': init['gen_tokens']}, + 'baseline_pass': r['_baseline_pass'], 'is_hard': r['_hard'], + 'view': r.get('_view', ''), 'rubric_diag': r.get('_rubric_diag', ''), + # xproblem audit: which neighbour's rubric this problem borrowed, and how similar. + 'rubric_src': r.get('_rubric_src', ''), 'neighbor_sim': r.get('_neighbor_sim'), + 'baseline_rolls': [_roll(x) for x in r['_baseline_rolls']], + 'candidates': [{ + 'skills': c['skills'], 'response': c['response'], 'parseable': c['parseable'], + 'skillgen_stop': c.get('skillgen_stop'), 'skillgen_tokens': c.get('skillgen_tokens'), + 'leaked': c['leaked'], 'leak_reason': c['leak_reason'], 'leak_source': c['leak_source'], + 'with_pass': c['with_pass'], 'reward': c.get('reward'), 'advantage': c.get('advantage'), + 'grpo_adv': c.get('grpo_adv'), 'kept': c.get('kept'), + 'rolls': [_roll(x) for x in c['rolls']], + } for c in r['_cands']], + } + + +def _view_stats(problems: List[Dict[str, Any]], view: str) -> Dict[str, Any]: + pv = [r for r in problems if r.get('_view') == view] + cands = [c for r in pv for c in r['_cands'] if c['parseable']] + clean = [c for c in cands if c['leaked'] is False] + adopted = sum(1 for r in pv + if any(abs(c.get('advantage') or 0.0) > 1e-9 for c in r['_cands'])) + return {'n': len(pv), 'n_candidates_parseable': len(cands), 'n_clean': len(clean), + 'n_adopted_problems': adopted, 'adoption_rate': (adopted / len(pv)) if pv else 0.0} + + +def _mean(xs: List[float]) -> float: + return sum(xs) / len(xs) if xs else 0.0 + + +def _std(xs: List[float]) -> float: + if len(xs) < 2: + return 0.0 + m = _mean(xs) + return (sum((x - m) ** 2 for x in xs) / len(xs)) ** 0.5 + + +def _signal_stats(problems: List[Dict[str, Any]]) -> Dict[str, float]: + """The heart of 'is there a learning signal': per problem, the scored candidates form a + GRPO group. A group with zero reward variance (all skills solve, or none do -- the + hard-problem dead zone) gives NO gradient. Tracks that fraction plus reward level and + within-group variance so a collapse (all-0 or all-1) is visible immediately.""" + group_vars, all_rewards, zero_grad, groups = [], [], 0, 0 + for r in problems: + rewards = [c['reward'] for c in r['_cands'] if c.get('reward') is not None] + if len(rewards) < 2: + continue + groups += 1 + all_rewards.extend(rewards) + v = _std(rewards) + group_vars.append(v) + if v < 1e-9: # every skill got the same reward -> GRPO skips this problem + zero_grad += 1 + return {'n_groups': groups, 'zero_grad_frac': (zero_grad / groups) if groups else 0.0, + 'reward_mean': _mean(all_rewards), 'reward_std': _std(all_rewards), + 'group_reward_std_mean': _mean(group_vars)} + + +def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespace) -> Dict[str, Any]: + all_cands = [c for r in chunk for c in r['_cands']] + cands = [c for c in all_cands if c['parseable']] + scored = [c for c in cands if c['with_pass'] is not None] + clean = [c for c in cands if c['leaked'] is False] + ws_rolls = [x for c in scored for x in c['rolls']] + # viewa-dropped problems generate no candidates; keep acc/* on the generated subset + # so the with-skill/lift trend stays comparable across view_b_frac settings. + gen_probs = [r for r in chunk if r['_cands']] + base_acc = _mean([1.0 if r['_baseline_pass'] else 0.0 for r in gen_probs]) + ws_acc = _mean([1.0 if any(c.get('reward') for c in r['_cands']) else 0.0 for r in gen_probs]) + cand_pass_parseable = _mean([c['with_pass'] for c in scored]) + cand_pass_all = _mean([1.0 if c.get('reward') else 0.0 for c in all_cands]) + # base failure taxonomy (you asked whether skills fail because the base loops out of length) + classes = [_baseline_class(r) for r in chunk] + n_fail = sum(1 for c in classes if c != 'success') + skill_tokens = [c.get('skillgen_tokens') or 0 for c in cands] # skill-gen response length + trunc = sum(1 for r in chunk for c in r['_cands'] + for x in c['rolls'] if x['stop_reason'] == 'length') + rubric_answer_leaks = sum( + 1 for r in chunk + if r.get('_view') == 'A' and _rubric_has_fail(r.get('_rubric_diag')) and _rubric_answer_leaked(r)) + return { + 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), + 'n_failed_first_try': sum(1 for r in chunk if r['_failed']), + 'n_generated': len(all_cands), 'n_candidates_parseable': len(cands), + 'n_unparseable': len(all_cands) - len(cands), + 'parse_rate': (len(cands) / len(all_cands)) if all_cands else 0.0, + 'n_leaked': sum(1 for c in cands if c['leaked']), 'n_clean': len(clean), + 'leak_rate': (sum(1 for c in cands if c['leaked']) / len(cands)) if cands else 0.0, + 'n_reward_pos': sum(1 for c in scored if c['reward']), + 'n_rubric_answer_leaked': rubric_answer_leaks, + 'n_viewa_dropped': sum(1 for r in chunk if _viewa_dropped(r, args)), + 'n_train_samples': sum(1 for c in all_cands if _is_trainable(c, args)), + 'signal': _signal_stats(chunk), + 'fail_loop_frac': (sum(1 for c in classes if c == 'fail_loop') / n_fail) if n_fail else 0.0, + 'fail_wrong_frac': (sum(1 for c in classes if c == 'fail_wrong') / n_fail) if n_fail else 0.0, + 'skill_tokens_mean': _mean(skill_tokens), + 'withskill_trunc_frac': (trunc / len(ws_rolls)) if ws_rolls else 0.0, + 'avg_baseline_pass': base_acc, 'avg_withskill_pass': ws_acc, + 'avg_lift': ws_acc - base_acc, + 'candidate_withskill_pass_parseable': cand_pass_parseable, + 'candidate_withskill_pass_all': cand_pass_all, + 'termination_rate_withskill': _mean([1.0 if x['terminated'] else 0.0 for x in ws_rolls]), + 'view_A': _view_stats(chunk, 'A'), 'view_B': _view_stats(chunk, 'B'), + **_xproblem_stats(chunk, args), + } + + +def _xproblem_stats(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> Dict[str, Any]: + """Cross-problem pairing health: of the view-A problems, how many actually got a + neighbour's rubric (pair_rate) and how similar those neighbours were. Empty when off.""" + if not args.xproblem_rubric: + return {} + view_a = [r for r in chunk if r.get('_view') == 'A'] + paired = [r for r in view_a if r.get('_rubric_src')] + return {'xproblem': { + 'n_view_a': len(view_a), 'n_paired': len(paired), + 'pair_rate': (len(paired) / len(view_a)) if view_a else 0.0, + 'neighbor_sim_mean': _mean([r.get('_neighbor_sim', 0.0) for r in paired])}} + + +def _sft_record(r: Dict[str, Any], c: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: + """A view-A context-distillation sample: the OPEN-BOOK (rubric-conditioned) skill ``c`` + is the target, but the prompt is rebuilt as query-only (view B, no rubric) so the model + learns to produce it CLOSED-BOOK. Rides the same GRPOLoss with a positive constant + advantage (``--sft-weight``); single-step (old_logps=None) this reduces to + ``-w*logp + beta*KL`` -- a KL-anchored weighted cross-entropy toward the skill.""" + return { + 'data_id': r.get('data_id'), + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'view': 'B', 'diagnosis': '', 'rubric_src': '', 'orig_view': 'A', 'sft': True, + 'response': c['response'], 'skills': c['skills'], + 'skillgen_stop': c.get('skillgen_stop'), + 'advantage': float(args.sft_weight), 'grpo_adv': 0.0, 'kept': True, + 'reward': c['reward'], 'with_pass': c['with_pass']} + + +def _rubric_answer_leaked(r: Dict[str, Any]) -> bool: + """Hard safety gate for external API/rubric diagnosis. Unlike answer-bearing skills + generated by the policy itself, a rubric that contains the target final answer is an + external teacher leak and must not be distilled into view B.""" + return _answer_leaked(r.get('_rubric_diag', ''), r.get('reference_answer', '')) + + +def _viewa_dropped(r: Dict[str, Any], args: argparse.Namespace) -> bool: + """Under --viewa-sft a view-A problem trains ONLY through the SFT path (rubric with + [FAIL], answer-free). No-FAIL and rubric-leaked problems produce no training record + at all (no GRPO backflow: those prompts are query-only and would muddy the pure + view-B GRPO pool), so their skill-gen + executor scoring is skipped entirely.""" + return (bool(args.viewa_sft) and r.get('_view') == 'A' + and (not _rubric_has_fail(r.get('_rubric_diag')) + or _rubric_answer_leaked(r))) + + +def _group_records(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> List[Dict[str, Any]]: + """Training records. Under --viewa-sft, view A is SFT-only: a problem whose rubric + localised a failure (has ``[FAIL]``, answer-free) contributes ONE context-distillation + SFT sample (best parseable open-book skill -- preferring an executor-verified pass, + else any parseable candidate -- rebuilt query-only); no-FAIL / rubric-leaked view-A + problems are dropped (``_viewa_dropped``, generation already skipped). GRPO candidates + come from view B only, keeping that pool purely closed-book. Prompts are rebuilt from + the stored view/diagnosis by ``_skillgen_messages``.""" + out = [] + for r in chunk: + if not r['_hard']: + continue + if args.viewa_sft and r.get('_view') == 'A': + if _viewa_dropped(r, args): + continue + best = _best_sft_candidate(r, args) + if best is not None: + out.append(_sft_record(r, best, args)) + continue + for c in r['_cands']: + if _is_trainable(c, args): + out.append({ + 'data_id': r.get('data_id'), + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'view': r.get('_view', 'A'), 'diagnosis': r.get('_rubric_diag', ''), + 'rubric_src': r.get('_rubric_src', ''), 'sft': False, + 'response': c['response'], 'skills': c['skills'], + 'skillgen_stop': c.get('skillgen_stop'), + 'advantage': c['advantage'], 'grpo_adv': c['grpo_adv'], 'kept': c['kept'], + 'reward': c['reward'], 'with_pass': c['with_pass']}) + return out + + +# =========================================================================== +# Block G -- online GRPO training +# =========================================================================== +def _is_num(v: Any) -> bool: + try: + float(v) + return True + except (TypeError, ValueError): + return False + + +def _train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: + """Training sample = the exact skill-gen prompt (rebuilt by ``_skillgen_messages`` so + train/inference match) + the generated structured guidance response. ``key_rounds`` + selects the final assistant turn; Template masks the prompt and trains the whole + response (the key-round prefix already excludes the prompt-provided ).""" + msgs = _skillgen_messages( + rec['problem'], rec.get('view', 'A'), rec.get('diagnosis', ''), rec.get('rubric_src', '')) + return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], + 'user_data': {'key_rounds': [len(msgs)]}} + + +def _train_chunk(skill_model, ref_model, ckpt: CheckpointEngineManager, samples: List[Dict[str, Any]], + args: argparse.Namespace) -> Dict[str, Any]: + """On-policy GRPO update over one chunk, then sync weights. Micro-batches of + ``sft_batch_size`` accumulate gradients; ONE optimizer step is taken per PPO + mini-batch of ``ppo_mini_batch_size`` samples (0 -> a single step over the whole + chunk, the original behaviour). A frozen reference model provides ref_logps for the + SEAM-style KL penalty. + + Multi-step correctness: with more than one step over the SAME rollout, later + mini-batches see an already-updated policy, so we FREEZE the sampling-policy + ``old_logps`` (recomputed once, before any step) and let GRPOLoss form the PPO ratio + against them. A single step keeps ``old_logps=None`` (ratio==1, pure on-policy). + The batch is padded to a multiple of ``sft_batch_size`` with advantage-0 copies that + contribute no policy gradient. View-A context-distillation samples ride the same loss + with a positive constant advantage (``--sft-weight``); single-step (old_logps=None) + that is ``-w*logp + beta*KL`` (KL-anchored weighted cross-entropy).""" + trajs = [_train_trajectory(rec) for rec in samples] + advs = [float(rec['advantage']) for rec in samples] + rem = (-len(trajs)) % args.sft_batch_size + if rem: + trajs += [trajs[-1]] * rem + advs += [0.0] * rem + + n, sft = len(trajs), args.sft_batch_size + mini = args.ppo_mini_batch_size if args.ppo_mini_batch_size > 0 else n + mini = max(sft, (mini // sft) * sft) # align to a whole number of micro-batches + multi_step = mini < n + + # Freeze ref_logps (frozen model) and -- only when we take multiple steps -- the + # sampling-policy old_logps, BOTH before any optimizer step touches the weights. With + # a single step old_logps stays None so GRPOLoss uses ratio==1 (pure on-policy). + micro_ref, micro_old = [], [] + for i in range(0, n, sft): + mb = trajs[i:i + sft] + micro_ref.append(ref_model.forward_only(inputs=mb).get('logps')) + micro_old.append(skill_model.forward_only(inputs=mb).get('logps') if multi_step else None) + + micro, n_steps = 0, 0 + for ms in range(0, n, mini): + for i in range(ms, min(ms + mini, n), sft): + k = i // sft + skill_model.forward_backward(inputs=trajs[i:i + sft], + advantages=advs[i:i + sft], + old_logps=micro_old[k], + ref_logps=micro_ref[k]) + micro += 1 + skill_model.clip_grad_and_step() + n_steps += 1 + ckpt.sync_weights(merge_and_sync=True) + metric = skill_model.calculate_metric(is_training=True) + n_sft = sum(1 for s in samples if s.get('sft')) + return {'n_samples': len(samples), 'n_sft': n_sft, 'n_grpo': len(samples) - n_sft, + 'n_steps': n_steps, 'n_micro_batches': micro, + 'metric': {k: (float(v) if _is_num(v) else v) for k, v in (metric or {}).items()}} + + +# =========================================================================== +# Block H -- fixed-holdout eval + metric formatting +# =========================================================================== +def run_greedy_eval(base_sampler, skill_sampler, eval_records: List[Dict[str, Any]], + ci: int, rounds: int, base_dp: int, skill_dp: int, args: argparse.Namespace, + base_cache: DiskCache + ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], Dict[str, float]]: + """SEAM ``val-core/math/acc/mean@1`` on the fixed holdout: ONE greedy skill (T=0) per + problem into ONE greedy base solve (T=0). Eval ALWAYS uses view B (query-only, the + deployment form: no rubric, since rubric needs an online teacher unavailable at deploy); + no leak filter (acc scores correctness alone). Baseline reuses the disk cache.""" + baseline_rollout(base_sampler, eval_records, base_dp, args, base_cache) + for r in eval_records: + r['_view'], r['_rubric_diag'] = 'B', '' + sg_out = _run_samples(skill_sampler, [_view_prompt(r, args) for r in eval_records], + 1, args.skill_max_tokens, skill_dp, temperature=0.0) + skills = [] + for seqs in sg_out: + if not seqs: + skills.append(('', '')) + continue + sresp = _clean_text(getattr(seqs[0], 'decoded', '') or '') + skills.append((_extract_skill(sresp) or '', sresp)) + ws_out = _run_samples(base_sampler, + [build_skill_solve_prompt(r['problem'], sk) for r, (sk, _) in zip(eval_records, skills)], + 1, args.max_tokens, base_dp, temperature=0.0) + recs = [] + for r, (sk, sresp), seqs in zip(eval_records, skills, ws_out): + roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + recs.append({ + 'record_type': 'eval_problem', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, + 'data_id': r.get('data_id'), + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'view': r['_view'], 'rubric_diag': r.get('_rubric_diag', ''), + 'baseline_pass': r['_baseline_pass'], 'skill': sk, 'skill_parseable': bool(sk), + 'skill_response': sresp, 'withskill_pred': roll['pred'], + 'withskill_correct': roll['correct'], 'withskill_terminated': roll['terminated'], + 'withskill_stop_reason': roll['stop_reason'], 'withskill_text': roll['text'], + }) + acc = lambda rs: sum(1 for x in rs if x['withskill_correct']) / len(rs) if rs else 0.0 + ws = acc(recs) # all view B (deployment form) + base = sum(x['baseline_pass'] for x in recs) / len(recs) if recs else 0.0 + fmt = (sum(1 for x in recs if x['skill_parseable']) / len(recs)) if recs else 0.0 + term = (sum(1 for x in recs if x['withskill_terminated']) / len(recs)) if recs else 0.0 + summary = {'record_type': 'eval_summary', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, + 'n': len(recs), 'view': 'B', 'acc_mean1': ws, + 'baseline_acc_mean1': base, 'lift_mean1': ws - base, + 'format_mean1': fmt, 'term_mean1': term} + metrics = {'core/math/acc/mean@1': ws, 'core/math/baseline_acc/mean@1': base, + 'core/math/lift/mean@1': ws - base, 'core/math/format/mean@1': fmt, + 'core/math/term/mean@1': term} + return recs, summary, metrics + + +def _trend_line(hist: List[Dict[str, float]], window: int, rounds_done: int) -> Optional[str]: + """Contrast the FIRST vs the most recent ``window`` chunks -- if RFT works, adoption + and lift on recent (fresh) chunks exceed the early baseline.""" + if len(hist) < 2 * window: + return None + base, rec = hist[:window], hist[-window:] + m = lambda xs, k: sum(h[k] for h in xs) / len(xs) + return (f'[trend] first {window} vs last {window} | ' + f'adopt A {m(base,"aA"):.2f}->{m(rec,"aA"):.2f} B {m(base,"aB"):.2f}->{m(rec,"aB"):.2f} | ' + f'lift {m(base,"lift"):+.3f}->{m(rec,"lift"):+.3f} | ' + f'0grad {m(base,"zero_grad"):.2f}->{m(rec,"zero_grad"):.2f} | ' + f'pos/chunk {m(base,"pos"):.1f}->{m(rec,"pos"):.1f} | rounds={rounds_done}') + + +def _swan_metrics(summary: Dict[str, Any], log: Optional[Dict[str, Any]]) -> Dict[str, float]: + """Flat swanlab dict. ``signal/*`` is the primary health group (is GRPO getting a + gradient at all); ``acc/*`` is the effect; the rest diagnose why. acc/adopt/skill are + only logged when the chunk produced a scored group, so idle chunks don't dip charts.""" + sig = summary['signal'] + d: Dict[str, float] = { + # --- signal: the FIRST thing to watch (no variance -> no learning) --- + 'signal/zero_grad_frac': sig['zero_grad_frac'], 'signal/n_groups': sig['n_groups'], + 'signal/reward_mean': sig['reward_mean'], 'signal/reward_std': sig['reward_std'], + 'signal/group_reward_std_mean': sig['group_reward_std_mean'], + 'signal/n_train_samples': summary['n_train_samples'], + 'signal/n_reward_pos': summary['n_reward_pos'], + # --- skill format / leak health --- + 'skill/parse_rate': summary['parse_rate'], 'skill/tokens_mean': summary['skill_tokens_mean'], + 'leak/rate': summary['leak_rate'], 'leak/n': summary['n_leaked'], + # --- base failure taxonomy (loop-out-of-length vs plain wrong) --- + 'fail/loop_frac': summary['fail_loop_frac'], 'fail/wrong_frac': summary['fail_wrong_frac'], + 'fail/frac_first_try': (summary['n_failed_first_try'] / summary['n']) if summary['n'] else 0.0, + # --- view-A routing: dropped = no-FAIL rubric or rubric leak (SFT-only view A) --- + 'viewa/dropped_frac': (summary['n_viewa_dropped'] / summary['view_A']['n'] + if summary['view_A']['n'] else 0.0), + } + bal = summary.get('balance') or {} + if bal.get('enabled'): + d.update({'balance/n_drawn': bal['n_drawn'], 'balance/n_baseline_fresh': bal['n_baseline_fresh'], + 'balance/selected_success_frac': bal['selected_success_frac']}) + xp = summary.get('xproblem') or {} + if xp: + d.update({'xproblem/pair_rate': xp['pair_rate'], 'xproblem/neighbor_sim_mean': xp['neighbor_sim_mean']}) + if sig['n_groups'] > 0: + d.update({'acc/baseline_pass': summary['avg_baseline_pass'], + 'acc/withskill_pass': summary['avg_withskill_pass'], 'acc/lift': summary['avg_lift'], + 'candidate/withskill_pass_parseable': summary['candidate_withskill_pass_parseable'], + 'candidate/withskill_pass_all': summary['candidate_withskill_pass_all'], + 'adopt/A': summary['view_A']['adoption_rate'], + 'adopt/B': summary['view_B']['adoption_rate'], + 'term/withskill': summary['termination_rate_withskill'], + 'term/withskill_trunc_frac': summary['withskill_trunc_frac']}) + if log: + d['train/n_steps'] = log['n_steps'] + d['train/n_micro_batches'] = log['n_micro_batches'] + for k, v in (log.get('metric') or {}).items(): + if not _is_num(v): + continue + if k.startswith('learning rate'): + if 'group 1' in k: + d['train/lr'] = float(v) + else: + d[f'train/{k.replace(" ", "_")}'] = float(v) + return d + + +def _view_a_rubric_leak_metrics(chunk: List[Dict[str, Any]], + pool_answers: Optional[Dict[str, str]] = None) -> Dict[str, float]: + """Swanlab-only audit for answer leakage in view-A rubric text. This never changes + rewards, advantages, filtering, or training records.""" + view_a = [r for r in chunk if r.get('_view') == 'A'] + with_diag = [r for r in view_a if r.get('_rubric_diag')] + target_leaks = sum(1 for r in with_diag + if _answer_leaked(r.get('_rubric_diag', ''), r.get('reference_answer', ''))) + source_leaks = 0 + pool_answers = pool_answers or {} + for r in with_diag: + src = r.get('_rubric_src') + src_ref = pool_answers.get(src, '') if src else r.get('reference_answer', '') + if _answer_leaked(r.get('_rubric_diag', ''), src_ref): + source_leaks += 1 + n = len(with_diag) + return { + 'rubric_leak/n_view_a': float(len(view_a)), + 'rubric_leak/n_checked': float(n), + 'rubric_leak/target_answer_n': float(target_leaks), + 'rubric_leak/target_answer_rate': (target_leaks / n) if n else 0.0, + 'rubric_leak/source_answer_n': float(source_leaks), + 'rubric_leak/source_answer_rate': (source_leaks / n) if n else 0.0, + } + + +# =========================================================================== +# Block F -- components, args, main +# =========================================================================== +def init_components(args: argparse.Namespace): + """Default 8-GPU layout: rank 0 trains the actor, rank 1 hosts a frozen ref model, + 2-3 skill_sampler (synced), 4-7 base_sampler (frozen). Returns + (skill_model, ref_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp).""" + r0, r1 = TRAIN_GPUS, TRAIN_GPUS + REF_GPUS + r2, r3 = r1 + SKILL_SAMPLER_GPUS, NUM_GPUS + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, lazy_collect=False, groups=[ + DeviceGroup(name='train', ranks=list(range(0, r0)), device_type='GPU'), + DeviceGroup(name='ref', ranks=list(range(r0, r1)), device_type='GPU'), + DeviceGroup(name='skill_sampler', ranks=list(range(r1, r2)), device_type='GPU'), + DeviceGroup(name='base_sampler', ranks=list(range(r2, r3)), device_type='GPU')]) + + train_mesh = DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_DP, fsdp_size=TRAIN_FSDP) + skill_model = TransformersModel(model_id=MODEL_ID, device_mesh=train_mesh, remote_group='train', + ddp_config={'find_unused_parameters': False}) + skill_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) + skill_model.set_template(Template, model_id=MODEL_ID, enable_thinking=False, + max_length=args.max_model_len, truncation_strategy='delete') + skill_model.set_processor(InputProcessor, padding_free=False) + skill_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon, beta=args.kl_beta) + skill_model.set_optimizer('AdamW', lr=args.lr) + skill_model.set_lr_scheduler('CosineWarmupScheduler', num_warmup_steps=10, + num_training_steps=args.max_train_rounds) + + ref_mesh = DeviceMesh.from_sizes(world_size=REF_GPUS, dp_size=REF_DP, fsdp_size=REF_FSDP) + ref_model = TransformersModel(model_id=MODEL_ID, device_mesh=ref_mesh, remote_group='ref', + ddp_config={'find_unused_parameters': False}) + ref_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) + ref_model.set_template(Template, model_id=MODEL_ID, enable_thinking=True, + max_length=args.max_model_len, truncation_strategy='delete') + ref_model.set_processor(InputProcessor, padding_free=False) + ref_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) + + def _sampler(group, world, enable_thinking: bool = True): + s = vLLMSampler(model_id=MODEL_ID, + engine_args={'gpu_memory_utilization': GPU_MEM, + 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, + device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), + remote_group=group) + s.set_template(Template, model_id=MODEL_ID, enable_thinking=enable_thinking, max_length=args.max_model_len) + return s + + skill_sampler = _sampler('skill_sampler', SKILL_SAMPLER_GPUS, enable_thinking=False) + # base_sampler is shared by the main thread and the baseline-prefetch thread -> serialise. + base_sampler = _LockedSampler(_sampler('base_sampler', BASE_SAMPLER_GPUS, enable_thinking=True)) + ckpt = CheckpointEngineManager(model=skill_model, sampler=skill_sampler) + return skill_model, ref_model, skill_sampler, base_sampler, ckpt, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS + + +def _build_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--dataset', choices=('aops', 'math'), default='aops') + p.add_argument('--n', type=int, default=2000, help='Problems loaded into the draw pool.') + p.add_argument('--pool-offset', type=int, default=0, + help='Skip this many shuffled non-eval records before building the train pool; ' + 'useful to avoid cold-start SFT data ranges.') + p.add_argument('--exclude-data-ids', default='', + help='Comma-separated jsonl files whose data_id/problem keys are excluded ' + 'from train/eval selection, e.g. coldstart_sft.jsonl.') + p.add_argument('--seed', type=int, default=42) + p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--eval-size', type=int, default=128, help='Fixed holdout size (0 disables).') + p.add_argument('--eval-every', type=int, default=5, help='Run holdout eval every N chunks.') + p.add_argument('--chunk-size', type=int, default=16) + p.add_argument('--balance', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--balance-success-frac', type=float, default=0.4, + help='Target fraction of the chunk the base solves (rest are base-fail).') + p.add_argument('--balance-loop-frac', type=float, default=0.5) + p.add_argument('--balance-max-draws-mult', type=int, default=8) + p.add_argument('--seam-parquet-dir', type=str, default='', + help='Read SEAM build_aops_dataset.py train.parquet/val.parquet directly, in ' + 'file order (problem<-extra_info.problem, answer<-reward_model.ground_truth). ' + 'val.parquet becomes the eval holdout. Bypasses load/--numeric-only/' + '--eval-size/internal shuffle so the input data matches a SEAM run exactly.') + p.add_argument('--n-skills', type=int, default=8) + p.add_argument('--view-b-frac', type=float, default=0.5) + p.add_argument('--xproblem-rubric', action=argparse.BooleanOptionalAction, default=False, + help='When enabled, view A uses a bag-of-words NEIGHBOUR problem rubric. ' + 'Default is off: each view-A problem uses its own baseline attempt, ' + 'while the API diagnosis prompt is constrained to be answer-free and ' + 'method-level only.') + p.add_argument('--skill-retries', type=int, default=2) + p.add_argument('--skill-gen-temperature', type=float, default=1.0) + p.add_argument('--skill-gen-top-p', type=float, default=1.0) + p.add_argument('--skill-gen-top-k', type=int, default=-1) + p.add_argument('--max-model-len', type=int, default=16384) + p.add_argument('--max-tokens', type=int, default=8192) + p.add_argument('--skill-max-tokens', type=int, default=8192) + p.add_argument('--rubric-workers', type=int, default=16) + p.add_argument('--rubric-retries', type=int, default=2, + help='Retry failed/timeout rubric diagnose calls this many times before ' + 'falling back to an empty diagnosis without caching the failure.') + p.add_argument('--sft-batch-size', type=int, default=8, + help='Driver micro-batch (gradient-accumulation unit); multiple of train dp.') + p.add_argument('--ppo-mini-batch-size', type=int, default=0, + help='Samples per optimizer step (SEAM-style PPO mini-batch). 0 = one step ' + 'over the whole chunk (on-policy, ratio==1). When >0 and smaller than ' + 'the trainable count, multiple steps are taken over the same rollout and ' + 'old_logps are frozen so the PPO ratio/clip stays valid. Rounded down to ' + 'a multiple of --sft-batch-size.') + p.add_argument('--grpo-epsilon', type=float, default=0.2) + p.add_argument('--adv-clip', type=float, default=3.0, + help='Symmetric clip for group-relative advantages; <=0 disables clipping.') + p.add_argument('--kl-beta', type=float, default=0.001, + help='SEAM-style reference KL coefficient for GRPOLoss.') + p.add_argument('--viewa-sft', action=argparse.BooleanOptionalAction, default=True, + help='Route view-A problems to online context distillation (SFT on the best ' + 'passing open-book skill, rebuilt as a query-only prompt) instead of GRPO. ' + 'View B stays GRPO; both share one optimizer step.') + p.add_argument('--sft-weight', type=float, default=0.5, + help='Advantage magnitude (lambda) for view-A SFT distillation samples in the ' + 'shared GRPOLoss. Single-step this is -lambda*logp + beta*KL (weighted CE). ' + 'Keep <~1 so one-signed SFT gradients do not swamp mean-0 GRPO advantages.') + p.add_argument('--sft-target-len', type=int, default=550, + help='Target skill length (chars) for view-A SFT distillation: among passing ' + 'candidates pick the one CLOSEST to this length (not the shortest). Set near ' + 'the empirical pass-rate peak (~500-600 here) so distillation neither shrinks ' + 'rollouts toward zero nor lets them grow unbounded.') + p.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--lr', type=float, default=6e-6) + p.add_argument('--max-train-rounds', type=int, default=1500) + p.add_argument('--save-rounds', type=int, default=200) + p.add_argument('--trend-every', type=int, default=10) + p.add_argument('--output-dir', default='./output/reflexion_skill') + p.add_argument('--cache-dir', default='', help='Baseline/rubric cache dir (default /cache).') + p.add_argument('--no-cache', action='store_true', help='Disable disk cache read/write.') + p.add_argument('--prefetch-baseline', action=argparse.BooleanOptionalAction, default=True, + help='Prefetch next chunk base baseline on a background thread (overlaps ' + 'with skill-gen; base_sampler is frozen so it never blocks the trainer).') + p.add_argument('--swanlab-project', default='twinkle') + p.add_argument('--swanlab-exp', default='') + args = p.parse_args() + if args.sft_batch_size % TRAIN_DP != 0: + raise ValueError(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple of train dp ({TRAIN_DP})') + if args.chunk_size < 1: + raise ValueError('--chunk-size must be >= 1') + args.balance_success_frac = max(0.0, min(1.0, args.balance_success_frac)) + return args + + +def _write(handle, row: Dict[str, Any]) -> None: + handle.write(json.dumps(row, ensure_ascii=False) + '\n') + + +def main() -> None: + args = _build_args() + records, eval_records, neighbor_map, pool_answers, data_stats = _load_records(args) + if len(records) < args.chunk_size: + raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded ({len(records)}); raise --n') + + os.makedirs(args.output_dir, exist_ok=True) + gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') + eval_path = os.path.join(args.output_dir, 'eval_records.jsonl') + data_path = os.path.join(args.output_dir, 'skill_dataset.jsonl') + train_log_path = os.path.join(args.output_dir, 'train_log.jsonl') + + if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY')): + sys.stderr.write('[rft] WARNING: no LLM backup env; view-A rubric check disabled ' + '(leak filter is deterministic, unaffected)\n') + + use_swan = swanlab is not None and os.environ.get('SWANLAB_MODE') != 'disabled' + if use_swan: + swanlab.init(project=args.swanlab_project, experiment_name=(args.swanlab_exp or None), + config={'model': MODEL_ID, 'dataset': args.dataset, 'n': len(records), + 'eval_n': len(eval_records), 'n_skills': args.n_skills, + 'view_b_frac': args.view_b_frac, 'balance': args.balance, + 'balance_success_frac': args.balance_success_frac, + 'skill_gen_temp': args.skill_gen_temperature, + 'grpo_epsilon': args.grpo_epsilon, 'kl_beta': args.kl_beta, + 'lr': args.lr}) + + skill_model, ref_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp = init_components(args) + checker = build_rubric_checker() + if checker is None: + sys.stderr.write('[rft] no LLM backup env -> view-A rubric process-check DISABLED\n') + + cache_dir = args.cache_dir or os.path.join(args.output_dir, 'cache') + os.makedirs(cache_dir, exist_ok=True) + use_cache = not args.no_cache + base_cache = DiskCache(os.path.join(cache_dir, 'baseline.jsonl'), use_cache) + eval_base_cache = DiskCache(os.path.join(cache_dir, 'eval_baseline.jsonl'), use_cache) + rubric_cache = DiskCache(os.path.join(cache_dir, 'rubric.jsonl'), use_cache) + if args.xproblem_rubric: + sys.stderr.write(f'[rft] xproblem-rubric ON: {len(neighbor_map)} bag-of-words pairs\n') + + cfg = {'record_type': 'config', 'model': MODEL_ID, 'dataset': args.dataset, + 'n': len(records), 'eval_n': len(eval_records), 'seed': args.seed, + 'pool_offset': args.pool_offset, 'exclude_data_ids': args.exclude_data_ids, + 'excluded_records': data_stats.get('excluded_records', 0), + 'numeric_only': args.numeric_only, 'raw_loaded': data_stats['raw_loaded'], + 'numeric_dropped': data_stats['numeric_dropped'], 'eval_every': args.eval_every, + 'n_skills': args.n_skills, 'view_b_frac': args.view_b_frac, + 'skill_retries': args.skill_retries, 'balance': args.balance, + 'balance_success_frac': args.balance_success_frac, + 'skill_gen_temp': args.skill_gen_temperature, 'reward': 'greedy_binary(correct)', + 'advantage': 'group_relative', 'format_in_reward': args.format_in_reward, 'cache': use_cache, + 'rubric_check': f'fixed_math_{len(_MATH_RUBRIC)}crit+term(viewA)' if checker else 'disabled', + 'xproblem_rubric': args.xproblem_rubric, + 'viewa_sft': args.viewa_sft, 'sft_weight': args.sft_weight, + 'sft_target_len': args.sft_target_len, + 'adv_clip': args.adv_clip, + 'grpo_epsilon': args.grpo_epsilon, 'kl_beta': args.kl_beta, 'lr': args.lr, + 'train_gpus': TRAIN_GPUS, 'ref_gpus': REF_GPUS, 'ref_fsdp': REF_FSDP, + 'train_fsdp': TRAIN_FSDP, 'train_dp': TRAIN_DP, + 'skill_sampler_gpus': SKILL_SAMPLER_GPUS, 'base_sampler_gpus': BASE_SAMPLER_GPUS, + 'max_train_rounds': args.max_train_rounds, 'started': int(time.time())} + sys.stderr.write(f'[rft] raw={data_stats["raw_loaded"]} numeric_drop={data_stats["numeric_dropped"]} ' + f'train={len(records)} eval={len(eval_records)} {args.dataset}; ' + f'train_gpus={TRAIN_GPUS} ref_gpus={REF_GPUS} train_fsdp={TRAIN_FSDP} ' + f'train_dp={TRAIN_DP} skill_dp={skill_dp} base_dp={base_dp}\n') + + hist: List[Dict[str, float]] = [] + rounds = 0 + pool = ProblemPool(records, args.seed) + with open(gen_path, 'w', encoding='utf-8') as gen_f, \ + open(eval_path, 'w', encoding='utf-8') as eval_f, \ + open(data_path, 'w', encoding='utf-8') as data_f, \ + open(train_log_path, 'w', encoding='utf-8') as tlog: + for f in (gen_f, eval_f, data_f, tlog): + _write(f, cfg) + gstep = 0 + # base baseline (frozen sampler, disk-cached, no weight-sync) is prefetched on a + # background thread while the current chunk generates: the skill-gen phase uses + # skill_sampler, so prefetching the NEXT chunk's baseline on base_sampler overlaps + # it and fills base_cache so the next draw_chunk hits it. base_sampler is wrapped in + # _LockedSampler, so the prefetch and this chunk's with-skill scoring never issue a + # base .sample() concurrently. It never touches the trainer or on-policy generation. + prefetch_pool = ThreadPoolExecutor(max_workers=1) if args.prefetch_baseline else None + pending: Optional[Any] = None + + def _prefetch(peeked: List[Dict[str, Any]]) -> None: + if peeked: + baseline_rollout(base_sampler, peeked, base_dp, args, base_cache) + + # Baseline (round 0) eval BEFORE any training: measures the untrained skill model on + # the fixed holdout so every later eval has a step-0 reference point on the same axis. + if eval_records: + eval_recs, eval_summary, eval_metrics = run_greedy_eval( + base_sampler, skill_sampler, eval_records, -1, rounds, base_dp, skill_dp, + args, eval_base_cache) + for rec in eval_recs: + _write(eval_f, rec) + _write(eval_f, eval_summary) + eval_f.flush() + if use_swan: + swanlab.log({f'eval/{k}': v for k, v in eval_metrics.items()}, step=0) + sys.stderr.write( + f'[eval] g-1 (init): n={eval_summary["n"]} viewB mean@1 ' + f'acc={eval_summary["baseline_acc_mean1"]:.3f}->{eval_summary["acc_mean1"]:.3f} ' + f'lift={eval_summary["lift_mean1"]:+.3f} ' + f'fmt={eval_summary["format_mean1"]:.2f} rounds={rounds}\n') + + # Each chunk is drawn fresh + RE-GENERATED with the current policy (on-policy); + # with --balance, draw_chunk keeps drawing until the base fail:success mix hits target. + while rounds < args.max_train_rounds: + if pending is not None: + pending.result() # finish last round's prefetch before drawing (cache-warm) + pending = None + chunk, balance = draw_chunk(pool, base_sampler, base_dp, args, base_cache) + if prefetch_pool is not None: + peeked = pool.peek(int(args.chunk_size * (args.balance_max_draws_mult if args.balance else 1))) + pending = prefetch_pool.submit(_prefetch, peeked) + full, summary, groups = process_chunk( + base_sampler, skill_sampler, chunk, gstep, base_dp, skill_dp, + args, checker, rubric_cache, base_cache, neighbor_map, pool_answers) + summary['balance'] = balance + + log = None + if groups: + log = _train_chunk(skill_model, ref_model, ckpt, groups, args) + rounds += 1 + log.update({'record_type': 'train_round', 'round': rounds, 'chunk': gstep, + 'epoch': pool.epoch, 'ts': int(time.time())}) + _write(tlog, log) + tlog.flush() + if rounds % args.save_rounds == 0: + skill_model.save(f'skill-rft-{rounds}', output_dir=args.output_dir) + + summary['rounds_done'], summary['epoch'] = rounds, pool.epoch + for rec in full: + _write(gen_f, rec) + _write(gen_f, summary) + gen_f.flush() + for v in groups: + _write(data_f, v) + data_f.flush() + + sa, sb, sig = summary['view_A'], summary['view_B'], summary['signal'] + hist.append({'aA': sa['adoption_rate'], 'aB': sb['adoption_rate'], + 'lift': summary['avg_lift'], 'pos': summary['n_reward_pos'], + 'zero_grad': sig['zero_grad_frac']}) + bal_str = (f'bal {balance["selected_fail"]}f/{balance["selected_success"]}s ' + f'(drew {balance["n_drawn"]}/fresh {balance["n_baseline_fresh"]}' + + ('!' if balance.get('budget_hit') else '') + ') ') if balance.get('enabled') else '' + xp = summary.get('xproblem') + xp_str = f'pair={xp["pair_rate"]:.2f}@{xp["neighbor_sim_mean"]:.2f} ' if xp else '' + tr_str = (f'train={log["n_grpo"]}g/{log["n_sft"]}s ' if log + else f'train={summary["n_train_samples"]} ') + sys.stderr.write( + f'[gen] e{pool.epoch} g{gstep}: {bal_str}n={summary["n"]} ' + f'clean={summary["n_clean"]} leak={summary["leak_rate"]:.2f} {tr_str}' + f'0grad={sig["zero_grad_frac"]:.2f} R={sig["reward_mean"]:.2f}±{sig["reward_std"]:.2f} ' + f'acc={summary["avg_baseline_pass"]:.2f}->{summary["avg_withskill_pass"]:.2f} ' + f'lift={summary["avg_lift"]:+.3f} {xp_str}' + f'A[{sa["n"]} {sa["adoption_rate"]:.2f}] B[{sb["n"]} {sb["adoption_rate"]:.2f}] ' + f'rounds={rounds}\n') + if use_swan: + swan_metrics = _swan_metrics(summary, log) + swan_metrics.update(_view_a_rubric_leak_metrics(chunk, pool_answers)) + swanlab.log(swan_metrics, step=gstep) + + if eval_records and (gstep + 1) % args.eval_every == 0: + eval_recs, eval_summary, eval_metrics = run_greedy_eval( + base_sampler, skill_sampler, eval_records, gstep, rounds, base_dp, skill_dp, + args, eval_base_cache) + for rec in eval_recs: + _write(eval_f, rec) + _write(eval_f, eval_summary) + eval_f.flush() + if use_swan: + swanlab.log({f'eval/{k}': v for k, v in eval_metrics.items()}, step=gstep) + sys.stderr.write( + f'[eval] g{gstep}: n={eval_summary["n"]} viewB mean@1 ' + f'acc={eval_summary["baseline_acc_mean1"]:.3f}->{eval_summary["acc_mean1"]:.3f} ' + f'lift={eval_summary["lift_mean1"]:+.3f} ' + f'fmt={eval_summary["format_mean1"]:.2f} rounds={rounds}\n') + + if (gstep + 1) % args.trend_every == 0: + tl = _trend_line(hist, args.trend_every, rounds) + if tl: + sys.stderr.write(tl + '\n') + gstep += 1 + + if prefetch_pool is not None: + if pending is not None: + pending.result() + prefetch_pool.shutdown(wait=True) + base_cache.close() + eval_base_cache.close() + rubric_cache.close() + skill_model.save('skill-rft-final', output_dir=args.output_dir) + sys.stderr.write(f'[rft] done: {rounds} rounds over {gstep} chunks / {pool.epoch} epochs -> {data_path}\n') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/legacy/train_skill_v2_ablate.sh b/cookbook/exp/legacy/train_skill_v2_ablate.sh new file mode 100644 index 000000000..0b859e759 --- /dev/null +++ b/cookbook/exp/legacy/train_skill_v2_ablate.sh @@ -0,0 +1,35 @@ +#!/bin/bash +# train_skill_v2_ablate.sh — skill 文体消融：toy vs pitfall，各 30 rounds，顺序执行 +# +# 设计（对应探针实验结论 CONCLUSIONS_config.md / CONCLUSIONS_reflexion.md）： +# 1. thinking 可控（--skill-thinking），本轮两个实验均为 off； +# 2. 两个文体方向：toy（异数字玩具题示范）与 pitfall（预判纠错）； +# 同一文体在主链路（query-only）与 buffer B regen（rubric 诊断）下输出格式一致， +# 保证 GRPO 与 SFT 样本分布一致可联合训练； +# 3. 每个实验 --max-train-rounds 30 结束，输出分目录： +# output.ablate_toy/skill_v2 与 output.ablate_pitfall/skill_v2 +# +# 用法: bash cookbook/exp/embedding/train_skill_v2_ablate.sh +# 注：train_skill_v2.sh 末尾以 "$@" 透传附加参数，argparse 同名参数后者覆盖前者， +# 故此处的 --max-train-rounds 30 会覆盖基础脚本里的 1500。 + +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +for STYLE in toy pitfall; do + echo "==============================================================" + echo "[ablate] 开始 skill-style=${STYLE} (thinking=off, 30 rounds)" + echo "==============================================================" + OUTPUT_DIR="./output.ablate_${STYLE}/skill_v2" \ + bash "${SCRIPT_DIR}/train_skill_v2.sh" \ + --skill-style "${STYLE}" \ + --skill-thinking off \ + --max-train-rounds 30 + echo "[ablate] skill-style=${STYLE} 完成" + # 中文注释：两次连跑之间等 Ray/vLLM 完全退出，避免引擎初始化竞态（曾复现过一次） + sleep 30 +done + +echo "[ablate] 两个消融实验全部完成：" +echo " toy -> ./output.ablate_toy/skill_v2" +echo " pitfall -> ./output.ablate_pitfall/skill_v2" diff --git a/cookbook/exp/legacy/train_skill_v2_ablate3.sh b/cookbook/exp/legacy/train_skill_v2_ablate3.sh new file mode 100755 index 000000000..e8e5f3285 --- /dev/null +++ b/cookbook/exp/legacy/train_skill_v2_ablate3.sh @@ -0,0 +1,62 @@ +#!/bin/bash +# train_skill_v2_ablate3.sh — 三路 prompt×thinking 消融，各 40 GRPO rounds，顺序执行。 +# +# 三组（单变量：skill 文体 + thinking）： +# 1) narrative + think -> output.ablate_narrative_think/skill_v2 +# 2) pitfall + think -> output.ablate_pitfall_think/skill_v2 +# 3) pitfall + nothink -> output.ablate_pitfall_nothink/skill_v2 +# +# 关键设计（对应探针结论 + 训练脚本代码事实）： +# A. 离线 SFT/buffer B 暂时关闭：本脚本 unset LLM_BACKUP_*/OPENAI 环境，使 +# build_rubric_checker() 返回 None -> 训练走 "GRPO only"（train_skill_v2.py:1596-1597,1673）。 +# 理由：40 rounds×chunk16≈640 题，远够不到 distill-trigger(150) 与 sft-trigger(100)， +# SFT 本就不会触发；关掉还能省去每条失败轨迹的 qwen-plus 预诊断 API 开销与后台线程。 +# B. think 组把 --skill-max-tokens 覆盖回 8192：基础 .sh 写死 4096，装不下 think 段+完整 +# ，会截断成空块导致 parse 崩（train_skill_v2.py:1455-1457）。nothink 组维持 4096。 +# C. eval 保留（--eval-every/--eval-size 沿用基础 .sh 的 5/200），这是三组对比的产出信号，不能关。 +# +# ⚠ 重要提醒（务必在 swanlab 盯 leak/rate 曲线）： +# 主 GRPO 路径 reward = parseable AND correct，leak 仅作 observability 审计、不进 reward +# （train_skill_v2.py:1195-1197,1208）。buffer B 关闭后训练回路里【没有任何泄漏防御】。 +# 探针实测 narrative+think 净泄漏≈0.46、pitfall+think 也偏高 —— 两个 think 组极可能 +# reward-hacking：靠把答案写进 skill 拿高 reward，reward_mean/eval-lift 会虚高。 +# 解读 think 组时必须同时看 leak/rate；pitfall+nothink 泄漏≈0，是干净对照基线。 + +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +# A. 关闭 buffer B / 离线 SFT / rubric 预诊断（GRPO only） +unset LLM_BACKUP_API_KEY LLM_BACKUP_BASE_URL LLM_BACKUP_MODEL OPENAI_API_KEY || true + +ROUNDS="${ROUNDS:-40}" + +# 每行: TAG STYLE THINKING(on/off) SKILL_MAX_TOKENS +RUNS=( + "narrative_think narrative on 8192" + "pitfall_think pitfall on 8192" + "pitfall_nothink pitfall off 4096" +) + +for spec in "${RUNS[@]}"; do + read -r TAG STYLE THINKING SMT <<< "${spec}" + OUT="./output.ablate_${TAG}/skill_v2" + echo "==============================================================" + echo "[ablate3] TAG=${TAG} style=${STYLE} thinking=${THINKING} skill_max_tokens=${SMT} rounds=${ROUNDS}" + echo " 输出目录: ${OUT} (GRPO only, buffer B 已关)" + echo "==============================================================" + OUTPUT_DIR="${OUT}" \ + bash "${SCRIPT_DIR}/train_skill_v2.sh" \ + --skill-style "${STYLE}" \ + --skill-thinking "${THINKING}" \ + --skill-max-tokens "${SMT}" \ + --max-train-rounds "${ROUNDS}" \ + --swanlab-exp "ablate3_${TAG}_$(date +%Y%m%d_%H%M%S)" + echo "[ablate3] ${TAG} 完成 -> ${OUT}" + # 两次连跑之间等 Ray/vLLM 完全退出，避免引擎初始化竞态（曾复现过） + sleep 30 +done + +echo "[ablate3] 三组全部完成：" +echo " narrative+think -> ./output.ablate_narrative_think/skill_v2" +echo " pitfall+think -> ./output.ablate_pitfall_think/skill_v2" +echo " pitfall+nothink -> ./output.ablate_pitfall_nothink/skill_v2" diff --git a/cookbook/exp/skill2lora/run_ablate12.sh b/cookbook/exp/skill2lora/run_ablate12.sh new file mode 100644 index 000000000..8a72cde22 --- /dev/null +++ b/cookbook/exp/skill2lora/run_ablate12.sh @@ -0,0 +1,109 @@ +#!/usr/bin/env bash +# ============================================================================== +# run_ablate12.sh — sequential launcher for the 12-experiment skill ablation. +# +# Reads the run plan from skill_ablate/config.py (single source of truth for order / +# dir names / think / skill-max-tokens / optional gate), then runs each experiment via +# `python -m skill_ablate.main --exp E{n}` in RUN_ORDER. +# +# Per experiment: +# - isolated product dir output.ablate12// +# - idempotent: a successful run writes /DONE.json (atomic, last step); both this +# script and skill_ablate.main skip completed experiments unless FORCE=1 +# - env snapshot output.ablate12//env_info.txt +# - skill-max-tokens 8192 (think) / 4096 (nothink) [from the plan] +# - E12 (sft, optional) SKIPPED unless RUN_SFT=1 +# - sleep between runs to let the previous Ray/vLLM engine tear down (avoid contention) +# +# Env knobs (all optional): +# DEEPMATH_DIR=$HERE/../../../deepmath_103k TRAIN_N=5000 MAX_UPDATES=50 EVAL_EVERY=5 +# LR=1e-6 RUN_SFT=1 FORCE=1 ONLY="E5 E6" SLEEP=30 SWANLAB_PROJECT=twinkle +# ============================================================================== +set -euo pipefail + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "$HERE" + +OUT_ROOT="${OUT_ROOT:-$HERE/output.ablate12}" +# DeepMath-103K (difficulty-stratified loader in skill_ablate/data.py); replaces the old +# SEAM/aops input — see skill_quality_analysis.md 组成漂移修正. +DEEPMATH_DIR="${DEEPMATH_DIR:-$(cd "$HERE/../../.." && pwd)/deepmath_103k}" +TRAIN_N="${TRAIN_N:-5000}" +EVAL_SIZE="${EVAL_SIZE:-128}" +MAX_UPDATES="${MAX_UPDATES:-50}" +EVAL_EVERY="${EVAL_EVERY:-5}" +LR="${LR:-1e-6}" +SLEEP="${SLEEP:-30}" +SWANLAB_PROJECT="${SWANLAB_PROJECT:-twinkle}" +RUN_SFT="${RUN_SFT:-0}" +FORCE="${FORCE:-0}" +ONLY="${ONLY:-}" + +mkdir -p "$OUT_ROOT" + +# --- pull the run plan (name \t exp_dir \t think \t smt \t optional) ------------------- +PLAN="$(python3 skill_ablate/config.py --plan)" + +snapshot_env() { # $1 = target file + { + echo "=== ablate12 env snapshot @ $(date -u +%FT%TZ) ===" + echo "host: $(hostname)" + echo "python: $(python3 -c 'import sys;print(sys.version.split()[0])')" + echo "torch: $(python3 -c 'import torch;print(torch.__version__)' 2>/dev/null || echo NA)" + echo "vllm: $(python3 -c 'import vllm;print(vllm.__version__)' 2>/dev/null || echo NA)" + echo "transformers: $(python3 -c 'import transformers;print(transformers.__version__)' 2>/dev/null || echo NA)" + echo "CUDA_VISIBLE_DEVICES: ${CUDA_VISIBLE_DEVICES:-unset}" + echo "nvidia-smi:"; nvidia-smi --query-gpu=name,memory.total --format=csv,noheader 2>/dev/null || echo " (nvidia-smi NA)" + echo "GPU layout: TRAIN=${TRAIN_GPUS:-2} REF=${REF_GPUS:-2} SKILL_SAMPLER=${SKILL_SAMPLER_GPUS:-2} BASE_SAMPLER=${BASE_SAMPLER_GPUS:-2}" + echo "LLM_BACKUP set: $([ -n "${LLM_BACKUP_API_KEY:-}${LLM_BACKUP_BASE_URL:-}${OPENAI_API_KEY:-}" ] && echo yes || echo no)" + } > "$1" +} + +echo "[ablate12] run order:"; echo "$PLAN" | awk -F'\t' '{printf " %s -> %s (think=%s smt=%s opt=%s)\n",$1,$2,$3,$4,$5}' + +while IFS=$'\t' read -r NAME EXP_DIR THINK SMT OPTIONAL; do + [ -z "$NAME" ] && continue + if [ -n "$ONLY" ] && ! grep -qw "$NAME" <<< "$ONLY"; then + echo "[ablate12] $NAME skipped (not in ONLY='$ONLY')"; continue + fi + if [ "$OPTIONAL" = "1" ] && [ "$RUN_SFT" != "1" ]; then + echo "[ablate12] $NAME ($EXP_DIR) skipped: optional; set RUN_SFT=1 to run"; continue + fi + + EXP_OUT="$OUT_ROOT/$EXP_DIR" + if [ -f "$EXP_OUT/DONE.json" ] && [ "$FORCE" != "1" ]; then + echo "[ablate12] $NAME already done ($EXP_OUT/DONE.json); FORCE=1 to rerun"; continue + fi + mkdir -p "$EXP_OUT" + snapshot_env "$EXP_OUT/env_info.txt" + + echo "======================================================================" + echo "[ablate12] START $NAME -> $EXP_OUT (think=$THINK skill_max_tokens=$SMT)" + echo "======================================================================" + LOG="$EXP_OUT/run.log" + FORCE_FLAG="" + [ "$FORCE" = "1" ] && FORCE_FLAG="--force" + set +e + python -m skill_ablate.main \ + --exp "$NAME" \ + --deepmath-dir "$DEEPMATH_DIR" \ + --n "$TRAIN_N" \ + --eval-size "$EVAL_SIZE" \ + --output-dir "$EXP_OUT" \ + --skill-max-tokens "$SMT" \ + --max-updates "$MAX_UPDATES" \ + --eval-every-updates "$EVAL_EVERY" \ + --lr "$LR" \ + --swanlab-project "$SWANLAB_PROJECT" \ + $FORCE_FLAG \ + 2>&1 | tee "$LOG" + RC=${PIPESTATUS[0]} + set -e + if [ "$RC" != "0" ]; then + echo "[ablate12] $NAME FAILED (rc=$RC); see $LOG. Stopping."; exit "$RC" + fi + echo "[ablate12] $NAME done. Sleeping ${SLEEP}s for engine teardown..." + sleep "$SLEEP" +done <<< "$PLAN" + +echo "[ablate12] all requested experiments finished." diff --git a/cookbook/exp/skill2lora/train_skill_v2.py b/cookbook/exp/skill2lora/train_skill_v2.py new file mode 100644 index 000000000..1d189f680 --- /dev/null +++ b/cookbook/exp/skill2lora/train_skill_v2.py @@ -0,0 +1,1800 @@ +"""Simplified GRPO + buffer-distill training for the reflexion skill generator (v2). + +Key differences from train_reflexion_skill.py: +- No view A/B split: all skill-gen is query-only (deployment form). +- No baseline rollout in training, no balance selection. +- thinking ON for the skill model: actor reasons in then emits a distilled + block; the is stripped by _extract_skill and NEVER reaches the executor + (executor only consumes ), so it is not SEAM-style think leakage. +- Reward = parseable × correct (aligned with SEAM lpem: no terminated, no length penalty). +- Buffer A: adv=0 (all-fail) problems accumulate failure trajectories. +- Buffer B: batch rubric → regenerate skill → pass@k validate → SFT injection. +- SFT is event-driven: buffer B reaches threshold → one SFT pass → eval. + +Launch: + LLM_BACKUP_API_KEY=... python cookbook/exp/skill2lora/train_skill_v2.py \ + --dataset aops --n 5000 --chunk-size 16 --lr 6e-6 +""" +import argparse +import copy +import hashlib +import json +import math +import os +import re +import sys +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Optional, Set, Tuple + +import numpy as np + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.data_format import SamplingParams +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.model import TransformersModel +from twinkle.patch.no_split_modules import NoSplitModulesPatch +from twinkle.processor import InputProcessor +from twinkle.sampler import vLLMSampler +from twinkle.template import Template +from twinkle_agentic.verifier import RubricVerifier +from twinkle_agentic.verifier.rubric_verifier import RubricItem + +logger = get_logger() + +try: + import swanlab +except ImportError: + swanlab = None + +MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') +GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.8)) +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') + +TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 2)) +REF_GPUS = int(os.environ.get('REF_GPUS', 2)) +SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) +BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 2)) +NUM_GPUS = TRAIN_GPUS + REF_GPUS + SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS +TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', min(1, TRAIN_GPUS))) +REF_FSDP = int(os.environ.get('REF_FSDP', min(1, REF_GPUS))) +TRAIN_DP = TRAIN_GPUS // TRAIN_FSDP +REF_DP = REF_GPUS // REF_FSDP + + +# =========================================================================== +# Section A — boxed extraction + answer grading (verbatim from v1) +# =========================================================================== +_BOXED_RE = re.compile(r'\\boxed\s*\{') + + +def extract_boxed(text: str) -> Optional[str]: + if not text: + return None + last = None + for m in _BOXED_RE.finditer(text): + depth, i = 1, m.end() + while i < len(text) and depth > 0: + depth += (text[i] == '{') - (text[i] == '}') + i += 1 + if depth == 0: + last = text[m.end():i - 1].strip() + return last + + +# SEAM-style answer format: executor emits ...[numeric only]. +# 中文注释：executor 答案格式对齐 SEAM——优先解析 …，回退到 \boxed{} +# （兼容旧轨迹/rubric 提示）。取最后一个 ，容忍缺失闭合标签（截断时取到 EOS）。 +_ANSWER_RE = re.compile(r'(.*?)', re.DOTALL | re.IGNORECASE) +_ANSWER_OPEN_RE = re.compile(r'(.*)', re.DOTALL | re.IGNORECASE) + + +def extract_answer(text: str) -> Optional[str]: + """Extract the final answer, preferring SEAM's …, falling back to \\boxed{}.""" + if not text: + return None + matches = _ANSWER_RE.findall(text) + if matches: + return matches[-1].strip() or None + # tolerate a truncated / unclosed tag (e.g. cut at token budget) + m = _ANSWER_OPEN_RE.search(text) + if m: + return m.group(1).strip() or None + return extract_boxed(text) + + +def normalize_answer(ans: str) -> str: + if not ans: + return '' + s = str(ans).strip() + m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) + if m: + return m.group(1) + s = s.strip('$').strip().replace('\u2212', '-') + 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').replace(r'\tfrac', r'\frac') + s = s.replace(r'\,', '').replace(r'\;', '').replace(r'\!', '') + s = re.sub(r'\\(?:quad|qquad|\s)', '', s) + s = re.sub(r'\s+', '', s) + s = re.sub(r'\\sqrt([0-9A-Za-z])', r'\\sqrt{\1}', s) + s = re.sub(r'\\frac\{([^{}]+)\}([^{}\\])', r'\\frac{\1}{\2}', s) + s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) + s = re.sub(r'\{[a-zA-Z]+\}$', '', s) + s = re.sub(r'\^\\circ|\^\{\\circ\}|\u00b0|\\circ', 'deg', s) + s = s.replace(r'\minus{}', '-').replace(r'\minus', '-') + + def _frac_to_slash(mt): + text = mt.group(0) + pos = text.index('{') + 1 + depth, num_start = 1, pos + while depth > 0: + depth += (text[pos] == '{') - (text[pos] == '}') + pos += 1 + numer = text[num_start:pos - 1] + pos += 1 + den_start, depth = pos, 1 + while depth > 0: + depth += (text[pos] == '{') - (text[pos] == '}') + pos += 1 + return f'({numer})/({text[den_start:pos - 1]})' + + s = re.sub(r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', _frac_to_slash, s) + s = re.sub(r'(? bool: + try: + va, vb = float(a.replace('(', '').replace(')', '')), 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) + return va is not None and vb is not None and abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) + + +_MCQ_REF_RE = re.compile( + r'^\\(?:textbf|text|mathrm|mathbf)\{\(?([A-E])\)?\s*\}\s*(.+)$' + r'|^\(?([A-E])\)\s+(.+)$') +_VAR_PREFIX_RE = re.compile(r'^(?:[A-Za-z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)$', re.DOTALL) + + +def _split_mcq(ans): + s = ans.strip() + m = _MCQ_REF_RE.match(s) + if m: + return (m.group(1) or m.group(3)), ((m.group(2) or m.group(4) or '').strip() or None) + bl = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) + return (bl.group(1), None) if bl else (None, s or None) + + +def _strip_var_prefix(ans): + m = _VAR_PREFIX_RE.match((ans or '').strip()) + return m.group(1).strip() if m else (ans or '') + + +def _math_verify_equal(predicted: str, reference: str) -> bool: + try: + from math_verify import parse, verify + gold = parse(r'\boxed{' + reference + '}', parsing_timeout=1) + pred = parse(r'\boxed{' + predicted + '}', parsing_timeout=1) + return bool(gold and pred and verify(gold, pred, timeout_seconds=1)) + except Exception: + return False + + +def answers_match(predicted: str, reference: str) -> bool: + if not predicted or not reference: + return False + norm_p, norm_r = normalize_answer(predicted), normalize_answer(reference) + if norm_p == norm_r or norm_p.lower() == norm_r.lower() or _try_numeric_equal(norm_p, norm_r): + return True + stripped_p = normalize_answer(_strip_var_prefix(predicted)) + stripped_r = normalize_answer(_strip_var_prefix(reference)) + if stripped_p and stripped_r: + if (stripped_p == stripped_r or stripped_p.lower() == stripped_r.lower() + or _try_numeric_equal(stripped_p, stripped_r)): + return True + p_letter, p_value = _split_mcq(predicted) + r_letter, r_value = _split_mcq(reference) + if p_letter and r_letter and p_letter == r_letter: + return True + if (p_letter and p_value is None) and (r_letter and r_value is None): + return False + 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 and (p_val == r_val or p_val.lower() == r_val.lower() or _try_numeric_equal(p_val, r_val)): + return True + for left, right in ((norm_p, norm_r), (stripped_p, stripped_r)): + tl, tr = re.sub(r'[\s()\[\]{}\\]', '', left or ''), re.sub(r'[\s()\[\]{}\\]', '', right or '') + if ',' in tl and tl == tr: + return True + if '=' in norm_r and '=' not in norm_p: + if any(part and (part == norm_p or _try_numeric_equal(part, norm_p)) for part in norm_r.split('=')): + return True + if '=' in norm_p and '=' not in norm_r: + if any(part and (part == norm_r or _try_numeric_equal(part, norm_r)) for part in norm_p.split('=')): + return True + return _math_verify_equal(stripped_p or norm_p, stripped_r or norm_r) + + +_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') + + +def _numeric_value(raw) -> Optional[str]: + if raw is None: + return None + s = str(raw).strip().strip('$').strip() + s = s.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac') + s = re.sub(r'\\!|\\,|\\;|\\ |\\left|\\right|\s', '', s) + for pat in (r'\\frac\{(-?\d+)\}\{(-?\d+)\}', r'(-?\d+)/(-?\d+)'): + m = re.fullmatch(pat, s) + if m: + a, b = int(m.group(1)), int(m.group(2)) + return str(int(a / b)) if (b and a / b == int(a / b)) else (str(a / b) if b else None) + return (str(int(float(s))) if float(s) == int(float(s)) else str(float(s))) if _NUM_RE.fullmatch(s) else None + + +def _answer_leaked(skill: str, reference: str) -> bool: + if not skill: + return False + # Suffix guard: reject only a following DIGIT or a following '.' (decimal point), + # NOT a sentence-ending '.'. Old '(?![\d.])' let leaks like "...= 675." slip through + # because the trailing period satisfied the [\d.] class. 中文注释：尾断言只排除"后接数字" + # 或"后接小数点+数字"，不排除句末句号，堵住 "答案." 这类泄漏漏检。 + for cand in {_numeric_value(reference), (str(reference).strip() or None)}: + if cand and re.search(r'(?/boxed/$...$/分数/首个数字，float 归一后纯数值精确匹配。 +_SEAM_TAG_RE = re.compile(r'<\s*answer\s*>(.*?)<\s*/\s*answer\s*>', re.I | re.S) +_SEAM_BOX_RE = re.compile(r'(?:\\{1,2}\(|)\\{1,2}boxed\s*\{\s*([^}]*)\s*}(?:\)|)', re.S) +_SEAM_INLINE_RE = re.compile(r'\$([^$]+)\$|\\\(([^)]+)\\\)', re.S) +_SEAM_FRAC_RE = re.compile(r'(-?\d+(?:\.\d+)?)/(-?\d+(?:\.\d+)?)') +_SEAM_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') + + +def _seam_norm(num: str) -> str: + try: + f = float(num) + return str(int(f)) if f == int(f) else str(f) + except Exception: + return num.strip() + + +def _seam_sanitize(txt: str) -> str: + """Port of SEAM lpem.sanitize_math_answer + normalize_number_format.""" + txt = (txt or '').strip() + if (m := _SEAM_TAG_RE.search(txt)): + txt = m.group(1).strip() + elif (m := _SEAM_BOX_RE.search(txt)): + txt = m.group(1).strip() + elif (m := _SEAM_INLINE_RE.search(txt)): + txt = (m.group(1) or m.group(2)).strip() + txt = re.sub(r'\\frac\s*\{\s*([^}]+?)\s*}\s*\{\s*([^}]+?)\s*}', r'\1/\2', txt) + if (m := _SEAM_FRAC_RE.search(txt)): + p, q = map(float, m.groups()) + if q: + return _seam_norm(str(p / q)) + if (m := _SEAM_NUM_RE.search(txt)): + return _seam_norm(m.group()) + return txt + + +# =========================================================================== +# Section B — sampling / parsing utilities +# =========================================================================== +_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') + + +def _clean_text(decoded: Optional[str]) -> str: + return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() + + +def _extract_skill(text: str) -> Optional[str]: + """Parse the skill block: in seam mode (SEAM format_pass parity), else .""" + low = text.lower() + end_think = low.rfind('') + answer = text[end_think + len(''):] if end_think >= 0 else text + if _ALIGN_MODE == 'seam': + m = re.search(r'(.*?)', answer, re.DOTALL | re.IGNORECASE) + return (m.group(1).strip() or None) if m else None + open_tag, close_tag = '', '' + s = answer.lower().rfind(open_tag) + if s < 0: + return None + inner = s + len(open_tag) + e = answer.lower().find(close_tag, inner) + if e < 0: + return None + block = answer[inner:e].strip() + block = re.sub(r'', '', block, flags=re.IGNORECASE).strip() + return block or None + + +def _parse_seq(seq, gold: str) -> Dict[str, Any]: + text = _clean_text(getattr(seq, 'decoded', '') or '') + if _ALIGN_MODE == 'seam': + # seam：SEAM lpem parity——整段贪婪 sanitize 成单一数值后精确匹配。 + san = _seam_sanitize(text) + pred = san or None + correct = bool(san) and (san == _seam_sanitize(str(gold))) + else: + # v2：只从 \boxed{} 抽取（executor 被要求把最终数值写进 \boxed{}），再走同一套数值归一 + # （frac/inline/number）后精确匹配；不做“整段抓首个数字”的贪婪回退（避免从推理里误抓）。 + # extract_boxed 取最后一个配平的 \boxed{}、截断（未闭合）时不误取；都没有则判错。 + raw = extract_boxed(text) + pred = _seam_sanitize(raw) if raw else None + correct = bool(pred) and (pred == _seam_sanitize(str(gold))) + terminated = getattr(seq, 'stop_reason', None) != 'length' + return {'pred': pred, 'correct': correct, 'terminated': terminated, + 'stop_reason': getattr(seq, 'stop_reason', None), + 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} + + +def _empty_roll(): + return {'pred': '', 'correct': False, 'terminated': False, + 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} + + +def _run_samples(sampler, prompts, num_samples, max_tokens, gen_dp, + temperature=None, top_p=None, top_k=None): + if not prompts: + return [] + params = SamplingParams( + max_tokens=max_tokens, + temperature=GEN_TEMPERATURE if temperature is None else temperature, + top_p=GEN_TOP_P if top_p is None else top_p, + num_samples=num_samples, **({} if top_k is None else {'top_k': top_k})) + padded = prompts + if gen_dp > 1 and 0 < len(prompts) < gen_dp: + padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] + responses = sampler.sample(padded, params)[:len(prompts)] + return [list(r.sequences) if (r and r.sequences) else [] for r in responses] + + +# =========================================================================== +# Section C — data loading (simplified: no balance, no xproblem, no views) +# =========================================================================== +def _boxed_batch(rows, dataset): + sols = rows['solution'] + metas = rows.get('metadata', [None] * len(sols)) + refs = [extract_boxed(s or '') for s in sols] + keep = [bool(ref) and (dataset != 'aops' or bool((meta or {}).get('boxed'))) + for ref, meta in zip(refs, metas)] + return {**rows, 'reference_answer': refs, '_keep': keep} + + +def load_problems(dataset: str, n: int, seed: int) -> List[Dict[str, Any]]: + ds_id = AOPS_DATASET_ID if dataset == 'aops' else os.environ.get('MATH_DATASET_ID', 'modelscope/competition_math') + ds = Dataset(DatasetMeta(dataset_id=f'ms://{ds_id}', split='train')) + nproc = min(32, os.cpu_count() or 1) + ds.map(lambda rows: _boxed_batch(rows, dataset), num_proc=nproc) + ds.filter(lambda row: row['_keep'], num_proc=nproc) + out = [{'data_id': f'{dataset}:{i}', 'problem': row['problem'], + 'reference_answer': row['reference_answer']} + for i, row in enumerate(ds.dataset)] + rng = np.random.RandomState(seed) + rng.shuffle(out) + return out[:n] if (n and n < len(out)) else out + + +def _load_seam_parquet(path: str) -> List[Dict[str, Any]]: + """Read a SEAM ``build_aops_dataset.py`` parquet (VERL RLHF schema) into twinkle records, + PRESERVING file row order. ``problem <- extra_info.problem`` and + ``reference_answer <- reward_model.ground_truth``. No shuffle/filter: the parquet is already + SEAM's numeric-filtered, seed-42-shuffled, truncated split. + 中文注释：直读 SEAM parquet、保持文件顺序，用于让 twinkle 与 SEAM 输入同一批数据。""" + import pyarrow.parquet as pq + rows = pq.read_table(path).to_pylist() + out: List[Dict[str, Any]] = [] + for i, r in enumerate(rows): + ei = r.get('extra_info') or {} + rm = r.get('reward_model') or {} + problem = (ei.get('problem') or '').strip() + ref = rm.get('ground_truth') + if not problem or ref is None: + continue + out.append({'data_id': f"seam:{ei.get('split', '')}:{ei.get('index', i)}", + 'problem': problem, 'reference_answer': str(ref)}) + return out + + +def _load_records(args): + seam_dir = (getattr(args, 'seam_parquet_dir', '') or '').strip() + if seam_dir: # 直读 SEAM parquet：按文件顺序取 train，val 整份当 eval，跳过 load/numeric/shuffle/split + tp, vp = os.path.join(seam_dir, 'train.parquet'), os.path.join(seam_dir, 'val.parquet') + if not (os.path.exists(tp) and os.path.exists(vp)): + raise FileNotFoundError( + f'--seam-parquet-dir needs both train.parquet and val.parquet in {seam_dir}') + pool = _load_seam_parquet(tp) # already SEAM-shuffled + truncated, file order + eval_records = _load_seam_parquet(vp) # SEAM's exact val holdout + eval_probs = {r['problem'] for r in eval_records} + train_records = [r for r in pool if r['problem'] not in eval_probs] + if args.n > 0: + train_records = train_records[:args.n] + if {r['problem'] for r in train_records} & eval_probs: + raise ValueError('eval/train overlap detected in SEAM parquet') + logger.info(f'[data] SEAM parquet: train={len(train_records)} eval={len(eval_records)} dir={seam_dir}') + return train_records, eval_records + records = load_problems(args.dataset, 0, args.seed) + raw_n = len(records) + if args.numeric_only: + records = [{**r, 'reference_answer': v} + for r, v in ((r, _numeric_value(r.get('reference_answer'))) for r in records) + if v is not None] + np.random.RandomState(args.seed).shuffle(records) + # exclude + excl_ids, excl_probs = set(), set() + for path in (args.exclude_data_ids or '').split(','): + path = path.strip() + if not path or not os.path.exists(path): + continue + with open(path) as f: + for line in f: + if not line.strip(): continue + row = json.loads(line) + if row.get('record_type') in {'config', 'summary'}: continue + did = str(row.get('data_id', '')).strip() + if did: excl_ids.add(did) + else: + p = str(row.get('problem', '')).strip() + if p: excl_probs.add(p) + if excl_ids or excl_probs: + records = [r for r in records + if str(r.get('data_id', '')) not in excl_ids + and str(r.get('problem', '')).strip() not in excl_probs] + eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 + eval_records = records[:eval_n] + # Dedup by problem TEXT: index slices are disjoint, but duplicate problem statements + # across the boundary would still leak eval into train. Drop any train record whose + # problem appears in eval, then guard with an explicit overlap assertion. + # 中文注释：train/eval 去重——按题面文本剔除，防止数据集内重复题目跨界泄漏；末尾硬断言无交集。 + eval_probs = {r['problem'] for r in eval_records} + train_records = [r for r in records[eval_n:] if r['problem'] not in eval_probs] + if args.n > 0: + train_records = train_records[:args.n] + if {r['problem'] for r in train_records} & eval_probs: + raise ValueError('eval/train overlap detected after dedup') + logger.info(f'[data] raw={raw_n} train={len(train_records)} eval={len(eval_records)}') + return train_records, eval_records + + +# =========================================================================== +# Section D — DiskCache, ProblemPool, LockedSampler +# =========================================================================== +class DiskCache: + def __init__(self, path: str, enabled: bool = True): + self._mem: Dict[str, Any] = {} + self._fh = None + self._lock = threading.Lock() + if not enabled: + return + if os.path.exists(path): + with open(path, encoding='utf-8') as f: + for line in f: + if line.strip(): + row = json.loads(line) + self._mem[row['key']] = row['value'] + self._fh = open(path, 'a', encoding='utf-8') + + @staticmethod + def key_for(*parts): + return hashlib.md5('\x1f'.join(parts).encode('utf-8')).hexdigest() + + def get(self, key): return self._mem.get(key) + def __contains__(self, key): return key in self._mem + + def put(self, key, value): + with self._lock: + self._mem[key] = value + if self._fh: + self._fh.write(json.dumps({'key': key, 'value': value}, ensure_ascii=False) + '\n') + self._fh.flush() + + def close(self): + if self._fh: self._fh.close() + + +class ProblemPool: + """Dataloader-like batch sampler: epoch-wise seeded RandomSampler + drop_last. + + SEAM's verl dataloader uses a sampler and drop_last=True; this mirrors that behavior more + closely than the old cursor loop that carried a short epoch tail into the next batch. + """ + def __init__(self, records, seed): + self._records = list(records) + self._seed, self._cursor, self.epoch = seed, 0, 0 + self._order: List[int] = [] + self._reset_epoch() + + def _reset_epoch(self): + rng = np.random.RandomState(self._seed + self.epoch) + self._order = list(rng.permutation(len(self._records))) + self._cursor = 0 + + def draw(self, k): + if k > len(self._records): + raise ValueError(f'batch size {k} exceeds dataset size {len(self._records)}') + if self._cursor + k > len(self._order): + self.epoch += 1 + self._reset_epoch() + idx = self._order[self._cursor:self._cursor + k] + self._cursor += k + return [self._records[i] for i in idx] + + +class _LockedSampler: + def __init__(self, sampler): + self._sampler = sampler + self._lock = threading.Lock() + + def sample(self, *a, **kw): + with self._lock: + return self._sampler.sample(*a, **kw) + + def __getattr__(self, name): + return getattr(self._sampler, name) + + +# =========================================================================== +# Section E — Rubric (teacher diagnosis, batched at distill time) +# =========================================================================== +_RFT_DIAG_SYSTEM = """\ +You are a strategy-level process checker for a math solution attempt. You are given a +math problem, a rubric, and one attempted solution segment. Decide PASS or FAIL for each +criterion, and write the diagnosis so it can become useful reusable guidance for solving +similar problems without seeing this segment. + +Output STRICT JSON (no prose outside it) with this shape: +{"items": [{"index": 1, "verdict": "PASS"|"FAIL", "reason": "...", "fix": ""}], "overall": "OK"|"ISSUES", "summary": "..."} + +Rules: +- Judge every criterion independently. +- The diagnosis must stay answer-free. +- For FAIL items: describe the process problem at strategy level. +- A fix suggests the LOCAL correction direction without solving. +- Never reveal the final answer or a corrected expression. +- If segment was cut off (no final reached), mark length-budget as FAIL. +- Keep "reason" and "fix" concise: one short sentence each. +- Output only the JSON object.""" + +_RFT_DIAG_USER = """\ +## Task / query +{query} + +## Rubric +{rubric} + +## Segment +{segment} + +Now output the diagnostic JSON object.""" + +_MATH_RUBRIC = [ + ('The attempt chooses a method suitable for the problem structure', False), + ('The attempt identifies the key constraint, invariant, or quantity before computing', False), + ('Algebraic and logical transformations preserve validity at each step', True), + ('The attempt checks required constraints, domains, boundary cases, or validity conditions', False), + ('The attempt avoids redundant casework, looping, or re-deriving known facts', False), + ('The attempt reaches a final within the length budget', False), + ('The approach stays focused on the actual question asked', False), +] +_RUBRIC_VERSION = 'rubric_v5_viewb_strategy' + + +class _RftRubricVerifier(RubricVerifier): + def _diagnose_trajectory(self, query, rubric_block, segment_text): + return {'messages': [ + {'role': 'system', 'content': _RFT_DIAG_SYSTEM}, + {'role': 'user', 'content': _RFT_DIAG_USER.format( + query=query, rubric=rubric_block, segment=segment_text)}]} + + +def build_rubric_checker(): + if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') + or os.environ.get('OPENAI_API_KEY')): + return None + return _RftRubricVerifier( + fixed_rubric=[RubricItem(t, is_hard=h) for t, h in _MATH_RUBRIC], gate=True) + + +def _format_diagnosis(detail) -> str: + rub = detail.rubric + lines = [] + for it in detail.items: + text = rub[it.index - 1].text if 0 < it.index <= len(rub) else f'criterion {it.index}' + if it.verdict: + lines.append(f'- [PASS] {text}') + else: + tail = f': {it.reason}' if it.reason else '' + tail += f' (fix: {it.fix})' if it.fix else '' + lines.append(f'- [FAIL] {text}{tail}') + if detail.summary: + lines.append(f'Summary: {detail.summary}') + return '\n'.join(lines) + + +def _diagnose_entry(checker, entry: Dict[str, Any]) -> Optional[str]: + """Run the teacher rubric on ONE buffer-A failure trajectory → formatted diagnosis text + (or None on API error). Pure network/CPU (no GPU), so it can run on a background thread + while GRPO trains. Shared by the background pre-diagnosis pool and distill_buffer's + fallback for any entry the pool did not reach in time. + 中文注释：单条失败轨迹的 rubric 诊断（纯 API，不吃 GPU）。后台预诊断与 distill 补诊断共用。""" + seg_text = entry['fail_segment'] + if entry.get('fail_stop_reason') == 'length': + seg_text += ('\n\n[Process note: this attempt was cut off at the token budget ' + 'and never produced a final .]') + seg = {'messages': [{'role': 'user', 'content': entry['problem']}, + {'role': 'assistant', 'content': seg_text}]} + try: + return _format_diagnosis(checker.diagnose(seg, query=entry['problem'])) + except Exception as exc: + logger.warning(f'[rubric] diagnose error: {exc}') + return None + + +# =========================================================================== +# Section F — NEW: prompts, reward, buffer logic +# =========================================================================== + +# ---- Skill-gen system prompt (query-only; used in v2 mode) ---- +# 中文注释：skillmodel 系统提示词（仅 v2 模式；seam 模式改用 _SEAM_EXPERIENCE_PROMPT）。 +# 方案1：thinking 开启。让 skill 模型在 里“先把本题实际解一遍、想清楚”，再在 里 +# 只写抽象出来的“通用方法论”（不含本题任何具体数字/中间结果/答案）。 会被 _extract_skill +# 用 rfind('') 砍掉、绝不流给 executor（避免 SEAM 那种 think 泄漏），executor 只吃 。 +# 之所以要开 thinking：nothinking 下模型无处安放解题过程，只能把“完整解答+答案”直接写进 +# （实测 前置分析长度=0、且常把答案算出来写进块内），等于换标签的泄漏且质量差。开 thinking 后 +# “先解题、再提炼”两步显式分离， 才可能是真正可迁移、不代入本题数值的方法论。 +# skill_model/ref_model/skill_sampler 三者 enable_thinking 必须一致，否则训练轨迹 token 布局与采样对不上。 +SKILL_GEN_SYSTEM = """\ +You are a skill-generation model. Your block will be fed to a SEPARATE downstream executor model that must solve the problem on its own. The executor will NOT see your private reasoning — it only sees what is inside .... + +First, think privately: actually work the problem out in your head to make sure you understand it, then step back and abstract WHAT MAKES THIS TYPE OF PROBLEM SOLVABLE into transferable methodology. + +Then write the block following these rules: +- Give general, transferable solving techniques for this TYPE of problem: the key concepts/theorems it relies on, the recommended strategy and steps, and the common pitfalls to avoid — plus a brief reason for each piece of advice so the executor understands why. +- Write it as one coherent analysis narrative (not a bullet list): first name what the problem is essentially asking, then walk through how to approach it, blending concepts, steps, pitfalls and reasons into a single connected story. +- CRITICAL: Do NOT solve the problem for the executor. Do NOT reveal or compute the final answer, and do NOT substitute the problem's specific given numbers into the steps or state any intermediate numeric results. Leave ALL concrete numbers for the executor to compute on its own. If you catch yourself writing a specific number from the problem, replace it with a description of the quantity instead. +- Keep it concise: aim for roughly one focused paragraph. + +Put ONLY the methodology inside . + +Example: + +This problem is essentially asking for the units (last) digit of an integer raised to a high power; first get clear on what the problem is asking before deciding where to start. Since only the last digit matters, you should first look only at the units digit of the base, because the units digit of an integer power is determined solely by the units digit of the base and the higher digits do not affect the result — so at this step be careful not to expand or compute the whole large number, which is both unnecessary and error-prone. Next, repeatedly multiply this units digit by itself and record the units digit each time, until it starts to repeat, thereby obtaining its cycle period. The part about "determining the period length" is important here: be careful not to count one term too many or too few, otherwise all the later positioning will be off. Finally, take the given exponent modulo the period length and land on the corresponding term within the period; here pay special attention that when the remainder is 0 it corresponds to the last term of the period rather than the first. Overall, I summarize the approach for this kind of problem as "first recognize that it asks for the units digit of a power, then fix on the units digit to find the cycle period, and finally use the exponent modulo to locate the term", while leaving the concrete numbers for the downstream solver to substitute and compute on its own. + +""" + +# ---- Skill 文体消融（--skill-style）---- +# 中文注释：探针实验（CONCLUSIONS_config/reflexion.md）验证的两种高性价比文体。 +# 关键约束：同一文体在主链路（query-only 预判）与 buffer B regen（rubric 诊断条件）下 +# 输出格式必须一致（toy=迷你题示范+迁移句；pitfall=WARNING/INSTEAD/纪律句）， +# 否则 GRPO 与 SFT 样本分布不一致无法联合训练。 +# toy 主链路：探针 P3_toy 原文（异数字玩具题，天然 answer-free）。 +SKILL_GEN_TOY = """\ +You are a skill-generation model. A separate executor model will solve the problem; it only sees your block. + +First think privately and identify the core technique this problem needs. Then, inside , do exactly one thing: invent a MINIATURE problem of the same type with DIFFERENT and much smaller numbers, and solve that miniature completely in at most 5 short lines, making the key trick explicit. Finish with one transfer sentence: "Your problem has the same shape - repeat these steps with its own numbers, then box a bare number." + +Hard rules: never mention or use any number that appears in the original problem; never state the original problem's answer; keep the whole block under 100 words. +""" + +# pitfall 主链路：探针 P5_pitfall 原文（预判最可能错误走向并拦截）。 +SKILL_GEN_PITFALL = """\ +You are a skill-generation model. A separate executor model will solve the problem; it only sees your block. + +First think privately: solve the problem in your head AND identify the single most likely way a solver goes wrong on this type (a tempting but wrong turn, an off-by-one, a wasteful brute-force, a wrong branch). Then, inside , write under 90 words: +- WARNING: name that most likely mistake concretely and say why it is wrong. +- INSTEAD: one or two sentences pointing to the correct turn (technique name + where to apply it), without solving the problem or revealing any numeric result. +- End with: "Avoid re-checking loops; box a bare number as soon as it is computed." +""" + +_SKILL_STYLE = 'narrative' # 'narrative' | 'toy' | 'pitfall'；由 main() 依据 --skill-style 设置 + +# ---- Executor prompt (with skill injection) ---- +# 中文注释：executor 提示词。答案格式对齐 SEAM 的 slove_qwen.txt（numeric-only，降截断）。 +# v2 模式：新版采用英文单 user turn——题目 + “技巧提示（skill 作为 advisory）” + 答案格式（见 +# build_skill_solve_prompt）。v2 执行器输出用 \boxed{}（_ANSWER_FORMAT_V2），先不用 ； +# seam 基线 DIRECT_SYSTEM 仍用 （_ANSWER_FORMAT），对齐 SEAM。skill 不再注入 system。 +_ANSWER_FORMAT = ('Present your reasoning and answer in the following format:\n' + ' Content of Thinking[Final numeric result only]') +# v2 执行器答案格式：把最终数值放进 \boxed{}（不再要求 ）；判分对应走 extract_boxed。 +_ANSWER_FORMAT_V2 = ('Present your reasoning, then put ONLY the final numeric result inside ' + '\\boxed{}. For example: \\boxed{42}.') +DIRECT_SYSTEM = ( + 'You are an expert competition mathematician. Be concise and accurate. ' + + _ANSWER_FORMAT) +_SKILL_SOLVE_PREFIX = ( + 'You are an expert competition mathematician. Be concise and accurate.\n\n' + 'Before you start, keep these reminders in mind to avoid common mistakes on this ' + 'type of problem:\n') +_SKILL_SOLVE_SUFFIX = ('\nApply them where relevant, but rely on your own reasoning to reach the answer.\n\n' + + _ANSWER_FORMAT) + + +# ---- Unified SEAM-alignment mode (toggle: prompt + skill-gen format only) ---- +# 中文注释：SEAM 对齐开关（由 --align-mode 控制）。 +# 1) executor 输入：seam=复刻 SEAM reward worker 的非空 experience 路径： +# prompt_text + actor 原始 response_text（保留 与 ）+ 解题 advisory。 +# v2=干净单 user turn（题目 + “Skill hint”advisory + 答案格式；空 skill 回退 direct）。 +# 2) skill-gen prompt：seam=SEAM EXPERIENCE_PROMPT（单 user turn，输出 ）；v2=SKILL_GEN_SYSTEM（）。 +# 注：reward 判分（_parse_seq）与 loss 聚合（set_loss）已统一为 SEAM 口径（lpem 纯数值匹配 + token-mean）。 +# 关于强制 ：SEAM 原版在 executor 末尾裸拼 "\n"。v2 走 messages+模板路径，base_sampler +# enable_thinking=True 时 Qwen3 生成起点默认进入 thinking；这里不改共享 sampler 的 assistant 前缀注入。 +_ALIGN_MODE = 'v2' # 'v2' | 'seam'；由 main() 依据 --align-mode 设置 + +_SEAM_EXPERIENCE_PROMPT = ( + 'You are a problem-solving guidance model. Read the math problem below and ' + 'distill a concise, reusable piece of solving experience that will help a ' + 'SEPARATE solver model reach the correct answer.\n' + 'Rules:\n' + '- Do NOT solve the problem and do NOT reveal or compute the final answer.\n' + '- State the key concepts/theorems, the recommended strategy/steps, and the ' + 'common pitfalls to avoid.\n' + '- Output ONLY the experience, wrapped EXACTLY as ' + ' ... .\n\n' + 'Problem:\n{problem}') +_SEAM_SOLVE_ADVISORY = ( + 'The above is a Q&A dialogue between a user and a problem-solving guidance model.\n' + 'Treat the output of the guidance model as advisory context to solve the math problem: ' + 'prefer using its techniques when they fit, but you may use alternative correct methods ' + 'if they are more efficient or clearer. If you diverge from the advisory context, briefly ' + 'explain why. Be concise and accurate.\n' + 'Present your reasoning and answer in the following format:\n' + ' Content of Thinking[Final numeric result only]') + + +def build_skill_solve_prompt_seam(problem, skill, raw_response=None): + """SEAM executor prompt. Non-empty skills use the actor's raw response_text, preserving + actor exactly as SEAM's reward worker does: prompt_text + response_text + grm. + If raw_response is missing, fall back to reconstructing a minimal response.""" + skill = (skill or '').strip() + response_text = (raw_response or '').strip() + prompt_text = ('<|im_start|>user\n' + + _SEAM_EXPERIENCE_PROMPT.format(problem=problem) + + '<|im_end|>\n<|im_start|>assistant\n') + if not response_text: + response_text = f'{skill}' + content = prompt_text + response_text + '\n' + _SEAM_SOLVE_ADVISORY + return {'messages': [{'role': 'user', 'content': content}]} + + +def build_direct_prompt(problem): + if _ALIGN_MODE == 'seam': + # seam 基线保持英文原样，不受 v2 prompt 改动影响 + return {'messages': [{'role': 'system', 'content': DIRECT_SYSTEM}, + {'role': 'user', 'content': problem}]} + # v2：英文 executor 基线——与带 skill 版同格式，仅去掉“技巧提示”部分 + content = f'The problem you need to solve:\n{problem}\n\n' + _ANSWER_FORMAT_V2 + return {'messages': [{'role': 'user', 'content': content}]} + + +def build_skill_solve_prompt(problem, skill, raw_response=None): + skill = (skill or '').strip() + if not skill: + # 空 skill → 干净 direct。训练侧根本不会用空 skill 走 executor（process_chunk 只对非空 flat 跑， + # 空候选直接 reward=0），故此分支仅影响 eval 口径——让空 skill 题 withskill==baseline、对 lift 贡献 0， + # 去掉空壳嵌套的框架水分，指标更干净。 + return build_direct_prompt(problem) + if _ALIGN_MODE == 'seam': + # 非空 skill 走 SEAM 原始 reward worker 路径：executor 可见 actor 完整 response_text（含 ）。 + return build_skill_solve_prompt_seam(problem, skill, raw_response=raw_response) + # v2：英文 executor——题目 + 技巧提示(skill 作为 advisory) + 答案格式，单 user turn + content = (f'The problem you need to solve:\n{problem}\n\n' + 'Skill hint:\nFor this problem, a skill-generation model has analyzed it and ' + 'provided some advisory skills:\n' + f'{skill}\n' + 'Prefer using its techniques when they fit, but if you have a more efficient or ' + 'clearer correct method, you may use it. If you diverge from this advice, briefly ' + 'explain why. Be concise and accurate.\n' + + _ANSWER_FORMAT_V2) + return {'messages': [{'role': 'user', 'content': content}]} + + +# ---- Rubric-guided regeneration prompt (buffer B distillation) ---- +# 中文注释：蒸馏重生成提示词。给旧 skill + rubric 诊断，要求产出改进后的 skill（）。 +# 输出要求第一人称自持句式、不指向外部上下文（防幻觉），含一个连贯叙述式示例。 +REGEN_SYSTEM = """\ +You are a skill-generation model. Your skill will be fed to a downstream executor model to help it solve the problem better. +You may give general, transferable solving techniques, together with why you give this advice, so the downstream model can follow it. Do NOT reveal or compute the final answer. + +You previously generated a skill, but that skill did not help the model. The executor's actual solving process has now been analyzed. You need to regenerate the skill based on your previous skill and the mistakes the model actually made, so as to help the model solve this problem. + +Your steps: +1. Re-read and understand the original problem. +2. Tell a coherent analysis story for this problem as one flowing narrative: first identify what it is essentially asking, then walk through how to approach it, naturally weaving together the solving points that were already correct last time, the pitfalls that actually tripped up the solving process and how to avoid them, and your reasoning for why you give this advice, blended into a single connected story, and leave the concrete numbers for the downstream solver to compute. +3. Put the above inside . + +Output requirement: Write your judgments and pitfall reminders about this problem directly in the first person (e.g. "I think this step tends to ...", "A common mistake is ..., so you need to ..."), and phrase the issues you find as self-contained, general techniques. Do NOT use phrasings that point to external context such as "according to the given analysis/hints" or "the previous skill" — the downstream executor cannot see that context, and such phrasings will cause hallucination. + +Example: + +This problem is essentially asking "how many arrangements satisfy the given constraints", which is a counting problem; first get clear on "what exactly is being counted" before deciding whether to use permutations or combinations. Since it is counting, you should first clearly define the objects being counted and the constraints, and judge whether the elements are distinguishable and whether order matters, because this directly determines whether you will need to divide out duplicates later. Next, first compute a total as if things were "ordered/distinguishable", then find which seemingly different arrangements actually correspond to the same configuration. The part about "recognizing symmetry and determining the duplication factor" is important here: I think the step most likely to go wrong in this problem is ignoring symmetry and treating essentially identical configurations as different, which makes the result too large; I think it is also easy to directly miss the "divide by the duplication factor" step — as long as the choices can be interchanged, you must divide out duplicates, otherwise you overcount. Finally, divide the total by the duplication factor to get the truly non-duplicated count; here pay special attention not to jump straight to permutation/combination formulas, but first think clearly about whether the elements are distinguishable and then decide whether to divide out duplicates. Overall, I summarize the approach for this kind of problem as "first recognize that it is a counting problem and judge whether the elements are distinguishable, then compute the total, recognize symmetry and remove duplicates", because I judge that the loss points for such problems almost all concentrate on overcounting; while leaving the concrete numbers for the downstream solver to substitute and compute on its own. +""" + +REGEN_USER = """\ +Original problem: +{problem} + +Previously generated skill (did not help the executor): +{orig_skill} + +Analysis of the executor's actual solving process: +{rubric_diag} + +Now rewrite the improved guidance:""" + +# 中文注释：buffer B regen 的 toy/pitfall 文体版（与主链路同文体，保证训练分布一致）。 +# 源自 reflexion 探针 D3_toyfix / D1_needle（diag-only 口径），另加防指涉硬规则 +# （不许写 "according to the diagnosis / the previous skill"，executor 看不到这些上下文）。 +REGEN_TOY_SYSTEM = """\ +You are a skill-generation model. A separate executor model previously FAILED this problem even with your earlier skill. You will see that earlier skill and an expert rubric diagnosis of the failure. The executor will retry seeing ONLY your new block. + +First think privately: from the diagnosis, identify the ONE technique the executor got wrong. Then, inside , do exactly this (under 110 words): +1. Invent a MINIATURE problem exercising that same technique with DIFFERENT, much smaller numbers, and solve the miniature completely in at most 5 short lines, making the correct move (the one the failed attempt missed) explicit. +2. One transfer sentence: "Your problem has the same shape - repeat these steps with its own numbers, then box a bare number." +Hard rules: never use any number from the original problem; never state its answer; the block must be self-contained - never reference "the diagnosis", "the previous skill" or any context the executor cannot see. +""" + +REGEN_PITFALL_SYSTEM = """\ +You are a skill-generation model. A separate executor model previously FAILED this problem even with your earlier skill. You will see that earlier skill and an expert rubric diagnosis of the failure. The executor will retry seeing ONLY your new block. + +First think privately: from the diagnosis, pinpoint the decisive error. Then, inside , write under 90 words: +- WARNING: the decisive mistake, stated concretely for THIS problem in self-contained first person (e.g. "I think the step most likely to go wrong is ..."). +- INSTEAD: one or two sentences pointing to the correct turn (technique name + where to apply it), without solving the problem or revealing any numeric result. +- End with: "Avoid re-checking loops; box a bare number as soon as it is computed." +Hard rules: the block must be self-contained - never reference "the diagnosis" or "the previous skill"; the executor cannot see them. +""" + + +def _skillgen_prompt(problem: str) -> Dict[str, Any]: + """Skill-gen prompt: query-only. seam mode uses SEAM EXPERIENCE_PROMPT (single user turn, + output); v2 uses SKILL_GEN_SYSTEM ().""" + if _ALIGN_MODE == 'seam': + return {'messages': [{'role': 'user', 'content': _SEAM_EXPERIENCE_PROMPT.format(problem=problem)}]} + # 中文注释：按 --skill-style 选主链路文体（narrative=现版叙述式 / toy / pitfall）。 + sys_p = {'toy': SKILL_GEN_TOY, 'pitfall': SKILL_GEN_PITFALL}.get(_SKILL_STYLE, SKILL_GEN_SYSTEM) + return {'messages': [ + {'role': 'system', 'content': sys_p}, + {'role': 'user', 'content': f'Problem:\n{problem}'}]} + + +def _regen_prompt(problem: str, orig_skill: str, rubric_diag: str) -> Dict[str, Any]: + """Regeneration prompt for buffer B distillation.""" + # 中文注释：regen 与主链路同文体（--skill-style），user 模板复用 REGEN_USER 三字段。 + sys_p = {'toy': REGEN_TOY_SYSTEM, 'pitfall': REGEN_PITFALL_SYSTEM}.get(_SKILL_STYLE, REGEN_SYSTEM) + return {'messages': [ + {'role': 'system', 'content': sys_p}, + {'role': 'user', 'content': REGEN_USER.format( + problem=problem, orig_skill=orig_skill, rubric_diag=rubric_diag)}]} + + +# ---- Reward ---- +# 中文注释：reward = parseable × correct（对齐 SEAM lpem：去 terminated、去长度惩罚）。 +# parseable=0 的候选 reward=0 仍参与 group（格式压力）。 +def _skill_reward(parseable: bool, correct: bool) -> float: + return 1.0 if (parseable and correct) else 0.0 + + +# ---- Buffer A: collect adv=0 all-fail problems ---- +def _collect_buffer_a(chunk, args) -> List[Dict[str, Any]]: + """Collect problems where all candidates got reward 0 (adv=0, GRPO blind spot). + Store one representative failure trajectory for later rubric diagnosis.""" + entries = [] + for r in chunk: + cs = [c for c in r['_cands'] if c.get('reward') is not None] + if len(cs) < 2: + continue + rewards = [c['reward'] for c in cs] + if max(rewards) > 0: + continue # has signal, not all-fail + # Representative trajectory for rubric + regen seed: prefer a terminated-wrong + # parseable candidate (complete reasoning to diagnose). Its skill becomes the regen + # seed, so pick the MOST SUBSTANTIAL one WITHIN budget (longest ≤ len_budget) — a + # rich-but-not-bloated starting point — rather than an arbitrary [0] or a near-empty + # skill. If all seeds exceed budget, take the one closest to budget (shortest-over). + # 中文注释：代表轨迹既做 rubric 诊断又做 regen 种子——优先"跑完但答错"的候选（完整推理）， + # 其 skill 取预算内最长(最有实质)的作种子；若全超预算则取最接近预算的，避免随机/近空种子。 + budget = args.len_budget + + def _seed_key(c): + L = len(c.get('skills') or '') + return (L <= budget, L if L <= budget else -L) + + parseable = [c for c in cs if c.get('skills')] + term_wrong = [c for c in parseable if c['rolls'] and c['rolls'][0].get('terminated')] + pool_c = term_wrong or parseable or cs + rep = max(pool_c, key=_seed_key) + stop_dist = {} + for c in cs: + sr = c['rolls'][0]['stop_reason'] if c['rolls'] else 'none' + stop_dist[sr] = stop_dist.get(sr, 0) + 1 + entries.append({ + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'data_id': r.get('data_id', ''), + 'orig_skill': rep.get('skills', ''), + 'orig_len': len(rep.get('skills', '')), + 'fail_segment': rep['rolls'][0]['text'] if rep['rolls'] else '', + 'fail_stop_reason': rep['rolls'][0]['stop_reason'] if rep['rolls'] else 'none', + 'stop_reason_dist': stop_dist, + }) + return entries + + +# ---- Buffer B distillation ---- +def distill_buffer(entries: List[Dict[str, Any]], skill_sampler, base_sampler, + checker, skill_dp: int, base_dp: int, + args) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + """Batch rubric → regenerate K distinct skills → greedy-validate → return SFT records. + 中文注释：蒸馏流程（方案 B，多样性在 skill 侧、executor 用贪心）： + 1. 批量 rubric 诊断失败轨迹；2. 仅 [FAIL] 项用高温重生成 K 个不同候选 skill； + 3. 每个候选过 ≤budget+无leak+去重 过滤；4. 每个存活候选用 executor 贪心(T=0)解 1 次； + 5. gate：≥m 个不同候选达成 terminated-correct → select 长度最接近 budget 的一个入 buffer B。 + 返回 (sft_records, distill_records)：后者逐 entry 记录 rubric_diag/候选 skill/贪心解结果/漏斗 + stage，落盘到 distill_records.jsonl 供复盘（否则 rubric 诊断与候选明细只存在于内存）。""" + if not checker or not entries: + return [], [] + + # Step 1: ensure every entry has a rubric diagnosis. Entries pre-diagnosed in the + # background (see _prediagnose in main) already carry '_rubric_diag'; only the misses + # are diagnosed here (in parallel), so the GPU-idle API wait is normally hidden. + # 中文注释：优先用后台预诊断结果；只对没预诊断到的条目并行补跑，隐藏 API 等待。 + pending = [e for e in entries if not e.get('_rubric_diag')] + if pending: + workers = min(args.rubric_workers, len(pending)) + with ThreadPoolExecutor(max_workers=max(1, workers)) as ex: + diags = list(ex.map(lambda e: _diagnose_entry(checker, e), pending)) + for entry, diag in zip(pending, diags): + entry['_rubric_diag'] = diag or '' + for entry in entries: + entry['rubric_diag'] = entry.get('_rubric_diag') or '' + + # Builder for the structured distill audit records (one per buffer-A entry). Closure over + # `entries`/`args`; takes the per-entry regen skills + greedy solve results (may be empty + # for the early-exit funnel stages). 中文注释：构造逐 entry 的蒸馏审计记录（含漏斗 stage）。 + def _mk_distill(results_by_entry, per_entry_skills, has_fail): + hf_index = {id(e): ei for ei, e in enumerate(has_fail)} + recs = [] + for e in entries: + ei = hf_index.get(id(e)) + cand_results = results_by_entry.get(ei, []) if ei is not None else [] + n_pass = sum(1 for c in cand_results if c['correct'] and c['terminated']) + n_cand = len(per_entry_skills[ei]) if (ei is not None and ei < len(per_entry_skills)) else 0 + if ei is None: + stage = 'no_fail' # rubric 未给出任何 [FAIL] + elif n_cand == 0: + stage = 'no_valid_regen' # 有 [FAIL] 但重生成无一条过 ≤budget/无leak/去重 + elif n_pass >= args.passatk_m: + stage = 'accepted' # ≥m 个候选贪心解对 → 入 buffer B + else: + stage = 'rejected' # 有候选但 = m distinct greedy-passing skills are + # RETRIED for up to --distill-retries extra rounds: regenerate more skills, dedup against those + # already seen for that entry, greedy-solve, and accumulate — rescuing more problems into buffer B. + # 中文注释：对每条 [FAIL] 高温采 K 个不同候选 skill 并贪心验证；还没凑够 m 个“贪心解对”的条目， + # 再重生成 --distill-retries 轮（新候选去重、贪心解、累计），把更多题救进 buffer B。 + k = args.passatk_k + per_entry_skills: List[List[str]] = [[] for _ in has_fail] # 累计去重候选(跨轮) + results_by_entry: Dict[int, List[Dict[str, Any]]] = {} # 累计每候选贪心结果(跨轮) + passers_by_entry: Dict[int, Set[str]] = {ei: set() for ei in range(len(has_fail))} + pending_ei = list(range(len(has_fail))) # 还没凑够 m 个 passer 的条目 + for _ in range(args.distill_retries + 1): + if not pending_ei: + break + regen_prompts = [_regen_prompt(has_fail[ei]['problem'], has_fail[ei]['orig_skill'], + has_fail[ei]['rubric_diag']) for ei in pending_ei] + regen_out = _run_samples(skill_sampler, regen_prompts, k, args.skill_max_tokens, skill_dp, + temperature=args.passatk_skill_temp, top_p=args.passatk_skill_top_p) + # 过滤(可解析/≤budget/无leak/对本条目去重) → 收集本轮新候选 + new_flat_idx, new_flat_prompts = [], [] + for ei, seqs in zip(pending_ei, regen_out): + seen = set(per_entry_skills[ei]) + for s in (seqs or []): + resp = _clean_text(getattr(s, 'decoded', '') or '') + skill = _extract_skill(resp) + if not skill or len(skill) > args.len_budget: + continue + if _answer_leaked(skill, has_fail[ei]['reference_answer']): + continue + if skill in seen: + continue + seen.add(skill) + per_entry_skills[ei].append(skill) + new_flat_idx.append((ei, skill)) + new_flat_prompts.append(build_skill_solve_prompt(has_fail[ei]['problem'], skill)) + # 本轮新候选各用 executor 贪心(T=0)解 1 次，累计结果与 distinct passer + if new_flat_prompts: + solve_out = _run_samples(base_sampler, new_flat_prompts, 1, args.max_tokens, base_dp, + temperature=0.0) + for (ei, sk), seqs in zip(new_flat_idx, solve_out): + roll = _parse_seq(seqs[0], has_fail[ei]['reference_answer']) if seqs else _empty_roll() + results_by_entry.setdefault(ei, []).append( + {'skill': sk, 'len': len(sk), 'correct': roll['correct'], 'terminated': roll['terminated']}) + if roll['correct'] and roll['terminated']: + passers_by_entry[ei].add(sk) + # 仍不足 m 个 distinct passer 的条目进入下一轮重试 + pending_ei = [ei for ei in pending_ei if len(passers_by_entry[ei]) < args.passatk_m] + + n_entries_with_cands = sum(1 for sk in per_entry_skills if sk) + if not results_by_entry: + logger.info(f'[distill] {len(has_fail)} [FAIL] entries, 0 valid regen skills') + return [], _mk_distill({}, per_entry_skills, has_fail) + + # Step 4: gate ≥ m distinct greedy-effective skills; select the survivor CLOSEST to the + # length budget (short is the floor, but not so short it degrades to answer-dumping). + # 中文注释：gate——≥m 个不同 skill 在贪心下 terminated-correct；select——在通过的候选里 + # 选长度最接近 budget 的一个入 buffer B（短是地板，但别短到退化成吐答案）。 + sft_records = [] + for ei, cand_results in results_by_entry.items(): + passers = [c['skill'] for c in cand_results if c['correct'] and c['terminated']] + if len(passers) < args.passatk_m: + continue + entry = has_fail[ei] + best = min(passers, key=lambda sk: abs(len(sk) - args.len_budget)) + sft_records.append({ + 'problem': entry['problem'], 'reference_answer': entry['reference_answer'], + 'data_id': entry.get('data_id', ''), + 'response': f'\n{best}\n', + 'skills': best, 'sft': True, + 'n_pass_skills': len(passers), 'n_cand_skills': len(per_entry_skills[ei]), + }) + + logger.info(f'[distill] {len(entries)} A → {len(has_fail)} [FAIL] → ' + f'{n_entries_with_cands} w/cands → {len(sft_records)} validated B ' + f'(gate m={args.passatk_m}/k={k})') + return sft_records, _mk_distill(results_by_entry, per_entry_skills, has_fail) + + + +# =========================================================================== +# Section G — GRPO advantages + training +# =========================================================================== +def _assign_advantages(chunk, args): + """Group-relative advantage: A = (R - mean) / (std + eps). std==0 → adv=0 (skipped).""" + eps = 1e-6 + adv_clip = abs(float(getattr(args, 'adv_clip', 0.0) or 0.0)) + for r in chunk: + for c in r['_cands']: + c['advantage'], c['kept'] = 0.0, False + cs = [c for c in r['_cands'] if c.get('reward') is not None] + if len(cs) < 2: + continue + rewards = [c['reward'] for c in cs] + mean_r = sum(rewards) / len(rewards) + # SEAM/verl uses torch.std's default unbiased=True for GRPO group std. + import torch + std = float(torch.std(torch.tensor(rewards, dtype=torch.float32)).item()) + if std < 1e-9: + continue + for c in cs: + raw = (c['reward'] - mean_r) / (std + eps) + c['advantage'] = max(-adv_clip, min(adv_clip, raw)) if adv_clip > 0 else raw + c['kept'] = c['reward'] > mean_r + + +def _train_trajectory(rec): + """Rebuild the query-only skill-gen prompt (train/inference match) + response. + GRPO records carry the full generated response; SFT records carry only the + cleaned block. key_rounds selects the final assistant turn.""" + msgs = _skillgen_prompt(rec['problem'])['messages'] + return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], + 'user_data': {'key_rounds': [len(msgs)]}} + + +def _train_step(skill_model, ref_model, ckpt, samples, args): + """On-policy GRPO update over one batch, then sync weights. SFT samples ride the + same BNPOLoss with a positive constant advantage (--sft-weight).""" + # 过滤空/纯空白 response：其可训练 token 为 0，会让持有它的 DP rank 跳过 backward， + # 与对端 all-reduce 失步 → NCCL 死锁（find_unused_parameters=False）。高熵采样偶发首 token 即 EOS。 + n_in = len(samples) + samples = [rec for rec in samples if (rec.get('response') or '').strip()] + n_empty = n_in - len(samples) + trajs = [_train_trajectory(rec) for rec in samples] + advs = [float(rec['advantage']) for rec in samples] + # drop_last 到 TRAIN_DP 整数倍：每个 micro（末尾那个可短于 sft）仍能被 dp 均分，零 padding 假样本。 + # 只丢尾部 ≤ dp-1 条真样本；若整批不足 dp（n 0 else n + mini = max(sft, (mini // sft) * sft) + multi_step = mini < n + micro_ref, micro_old = [], [] + for i in range(0, n, sft): + mb = trajs[i:i + sft] + micro_ref.append(ref_model.forward_only(inputs=mb).get('logps')) + micro_old.append(skill_model.forward_only(inputs=mb).get('logps') if multi_step else None) + micro, n_steps = 0, 0 + for ms in range(0, n, mini): + for i in range(ms, min(ms + mini, n), sft): + k = i // sft + skill_model.forward_backward(inputs=trajs[i:i + sft], advantages=advs[i:i + sft], + old_logps=micro_old[k], ref_logps=micro_ref[k]) + micro += 1 + skill_model.clip_grad_and_step() + n_steps += 1 + ckpt.sync_weights(merge_and_sync=True) + metric = skill_model.calculate_metric(is_training=True) + n_sft = sum(1 for s in samples if s.get('sft')) + return {'n_samples': n_in, 'n_sft': n_sft, 'n_grpo': len(samples) - n_sft, 'n_empty': n_empty, + 'n_steps': n_steps, 'n_micro_batches': micro, + 'metric': {k: (float(v) if _is_num(v) else v) for k, v in (metric or {}).items()}} + + +def _is_num(v): + try: + float(v); return True + except (TypeError, ValueError): + return False + + +# =========================================================================== +# Section H — chunk processing, records, eval (+ hard-slice rescue) +# =========================================================================== +def process_chunk(base_sampler, skill_sampler, chunk, ci, base_dp, skill_dp, args): + """skill-gen (query-only) → leak audit → with-skill greedy pass → reward → advantages. + Returns (full_records, summary, grpo_train_records, buffer_a_entries).""" + for r in chunk: + r['_cands'] = [] + # skill-gen (single pass, no retry)：每题恒 n_skills 个候选、组大小固定，对齐 SEAM rollout.n。 + # （全 0 的难题不在这里重采，而是进 buffer A → rubric 重生成，见 distill_buffer。） + flat = [] + sg_out = _run_samples(skill_sampler, [_skillgen_prompt(r['problem']) for r in chunk], + args.n_skills, args.skill_max_tokens, skill_dp, + temperature=args.skill_gen_temperature, top_p=args.skill_gen_top_p, + top_k=args.skill_gen_top_k) + for r, seqs in zip(chunk, sg_out): + for s in seqs: + resp = _clean_text(getattr(s, 'decoded', '') or '') + block = _extract_skill(resp) or '' + cand = {'skills': block, 'response': resp, 'parseable': bool(block), + 'leaked': None, 'with_pass': None, 'reward': None, 'rolls': [], + 'advantage': 0.0, 'kept': False, + 'skillgen_stop': getattr(s, 'stop_reason', None), + 'skillgen_tokens': len(getattr(s, 'tokens', None) or [])} + r['_cands'].append(cand) + if block: + flat.append((r, cand)) + + # leak audit (deterministic, observability only) + for r, c in flat: + c['leaked'] = _answer_leaked(c['skills'], r['reference_answer']) + + # with-skill greedy pass (T=0) + if flat: + ws_out = _run_samples(base_sampler, + [build_skill_solve_prompt(r['problem'], c['skills'], c.get('response')) for r, c in flat], + 1, args.max_tokens, base_dp, temperature=0.0) + for (r, c), seqs in zip(flat, ws_out): + roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + c['rolls'] = [roll] + c['with_pass'] = 1.0 if roll['correct'] else 0.0 + c['reward'] = _skill_reward(c['parseable'], roll['correct']) + # unparseable candidates score 0 and still join the group (format pressure) + for r in chunk: + for c in r['_cands']: + if c['reward'] is None: + c['reward'] = 0.0 + + _assign_advantages(chunk, args) + + # SEAM/verl 对齐：每个 dataloader batch 都进入 actor update。零 adv 候选 PG 贡献 0， + # 但仍计入 token-mean 分母，并在 beta>0 时贡献 KL 锚定；不再因整 chunk 无信号而跳过 step。 + grpo = [] + for r in chunk: + for c in r['_cands']: + if c.get('reward') is None: + continue + grpo.append({'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'data_id': r.get('data_id', ''), 'response': c['response'], + 'skills': c['skills'], 'advantage': c['advantage'], + 'kept': c['kept'], 'reward': c['reward'], 'sft': False}) + buffer_a = _collect_buffer_a(chunk, args) + return _full_records(chunk, ci), _chunk_summary(chunk, ci), grpo, buffer_a + + +def _roll(x): + return {k: x[k] for k in ('pred', 'correct', 'terminated', 'stop_reason', 'gen_tokens', 'text')} + + +def _full_records(chunk, ci): + out = [] + for r in chunk: + out.append({ + 'record_type': 'problem', 'chunk': ci, 'problem': r['problem'], + 'reference_answer': r['reference_answer'], 'data_id': r.get('data_id', ''), + 'candidates': [{ + 'skills': c['skills'], 'response': c['response'], 'parseable': c['parseable'], + 'leaked': c['leaked'], 'with_pass': c['with_pass'], 'reward': c.get('reward'), + 'advantage': c.get('advantage'), 'kept': c.get('kept'), + 'skillgen_stop': c.get('skillgen_stop'), 'skillgen_tokens': c.get('skillgen_tokens'), + 'rolls': [_roll(x) for x in c['rolls']], + } for c in r['_cands']], + }) + return out + + +def _mean(xs): + return sum(xs) / len(xs) if xs else 0.0 + + +def _std(xs): + if len(xs) < 2: + return 0.0 + import torch + return float(torch.std(torch.tensor(xs, dtype=torch.float32)).item()) + + +def _chunk_summary(chunk, ci): + all_cands = [c for r in chunk for c in r['_cands']] + cands = [c for c in all_cands if c['parseable']] + scored = [c for c in cands if c['with_pass'] is not None] + ws_rolls = [x for c in scored for x in c['rolls']] + # signal: fraction of groups with zero reward variance (no gradient) + group_vars, all_rewards, zero_grad, groups = [], [], 0, 0 + for r in chunk: + rewards = [c['reward'] for c in r['_cands'] if c.get('reward') is not None] + if len(rewards) < 2: + continue + groups += 1 + all_rewards.extend(rewards) + v = _std(rewards) + group_vars.append(v) + if v < 1e-9: + zero_grad += 1 + n_train = sum(1 for c in all_cands if abs(c.get('advantage') or 0.0) > 1e-9) + trunc = sum(1 for x in ws_rolls if x['stop_reason'] == 'length') + ws_acc = _mean([1.0 if any(c.get('reward') for c in r['_cands']) else 0.0 + for r in chunk if r['_cands']]) + return { + 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), + 'n_generated': len(all_cands), 'n_candidates_parseable': len(cands), + 'parse_rate': (len(cands) / len(all_cands)) if all_cands else 0.0, + 'n_leaked': sum(1 for c in cands if c['leaked']), + 'leak_rate': (sum(1 for c in cands if c['leaked']) / len(cands)) if cands else 0.0, + 'n_train_samples': n_train, 'n_groups': groups, + 'zero_grad_frac': (zero_grad / groups) if groups else 0.0, + 'reward_mean': _mean(all_rewards), 'reward_std': _std(all_rewards), + 'group_reward_std_mean': _mean(group_vars), + 'skill_tokens_mean': _mean([c.get('skillgen_tokens') or 0 for c in cands]), + 'skill_chars_mean': _mean([len(c['skills']) for c in cands]), + 'avg_withskill_pass': ws_acc, + 'candidate_withskill_pass': _mean([c['with_pass'] for c in scored]), + 'withskill_trunc_frac': (trunc / len(ws_rolls)) if ws_rolls else 0.0, + 'termination_rate_withskill': _mean([1.0 if x['terminated'] else 0.0 for x in ws_rolls]), + } + + +def run_greedy_eval(base_sampler, skill_sampler, eval_records, ci, rounds, + base_dp, skill_dp, args, base_cache): + """Holdout readout: skill-gen (T=args.eval_skill_temperature, args.eval_rollouts rollouts) + -> greedy base solve (T=0). acc = per-problem mean correctness over the rollouts, averaged over + problems (falls back to single greedy when eval_rollouts=1 & temp=0). Adds hard-slice + (baseline_pass==0) rescue rate as a zero-cost secondary readout.""" + # baseline (frozen, cached) + todo = [r for r in eval_records if DiskCache.key_for(r['problem']) not in base_cache] + if todo: + out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], + 1, args.max_tokens, base_dp, temperature=0.0) + for r, seqs in zip(todo, out): + roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + base_cache.put(DiskCache.key_for(r['problem']), roll) + for r in eval_records: + br = base_cache.get(DiskCache.key_for(r['problem'])) + r['_baseline_pass'] = 1.0 if br['correct'] else 0.0 + # skill-gen (T=eval_skill_temperature, R rollouts) → with-skill greedy → mean acc over rollouts + R = max(1, args.eval_rollouts) + sg_out = _run_samples(skill_sampler, [_skillgen_prompt(r['problem']) for r in eval_records], + R, args.skill_max_tokens, skill_dp, temperature=args.eval_skill_temperature) + # per problem -> list of R (skill, sresp) + per_skills = [] + for seqs in sg_out: + seqs = list(seqs or []) + row = [] + for j in range(R): + s = seqs[j] if j < len(seqs) else None + if s is None: + row.append(('', '')) + else: + sresp = _clean_text(getattr(s, 'decoded', '') or '') + row.append((_extract_skill(sresp) or '', sresp)) + per_skills.append(row) + # flatten R×N for a single batched greedy executor pass + flat_prompts, flat_idx = [], [] + for pi, (r, row) in enumerate(zip(eval_records, per_skills)): + for j, (sk, sresp) in enumerate(row): + flat_prompts.append(build_skill_solve_prompt(r['problem'], sk, sresp)) + flat_idx.append((pi, j)) + ws_out = _run_samples(base_sampler, flat_prompts, 1, args.max_tokens, base_dp, temperature=0.0) + roll_by = {} + for (pi, j), seqs in zip(flat_idx, ws_out): + roll_by[(pi, j)] = _parse_seq(seqs[0], eval_records[pi]['reference_answer']) if seqs else _empty_roll() + recs = [] + for pi, (r, row) in enumerate(zip(eval_records, per_skills)): + rolls = [roll_by[(pi, j)] for j in range(len(row))] + corr = [1.0 if x['correct'] else 0.0 for x in rolls] + parses = [1.0 if sk else 0.0 for sk, _ in row] + terms = [1.0 if x['terminated'] else 0.0 for x in rolls] + acc_mean = sum(corr) / len(corr) if corr else 0.0 + recs.append({ + 'record_type': 'eval_problem', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, + 'data_id': r.get('data_id', ''), 'problem': r['problem'], + 'reference_answer': r['reference_answer'], 'baseline_pass': r['_baseline_pass'], + 'n_rollouts': len(row), 'eval_skill_temperature': args.eval_skill_temperature, + 'withskill_acc_mean': acc_mean, # per-problem mean over R rollouts + 'withskill_pass_any': 1.0 if any(corr) else 0.0, # pass@R (bonus readout) + 'skill_parseable_mean': sum(parses) / len(parses) if parses else 0.0, + 'withskill_terminated_mean': sum(terms) / len(terms) if terms else 0.0, + # 首个 rollout 的明细留作肉眼抽查 + 'skill': row[0][0], 'skill_parseable': bool(row[0][0]), 'skill_chars': len(row[0][0]), + 'withskill_pred': rolls[0]['pred'], 'withskill_correct': rolls[0]['correct'], + 'withskill_terminated': rolls[0]['terminated'], 'withskill_stop_reason': rolls[0]['stop_reason'], + 'withskill_text': rolls[0]['text'], + }) + n = len(recs) + # acc = 跨题平均的"每题 R 次平均正确率"（mean-over-rollouts） + ws = (sum(x['withskill_acc_mean'] for x in recs) / n) if n else 0.0 + pass_any = (sum(x['withskill_pass_any'] for x in recs) / n) if n else 0.0 + base = (sum(x['baseline_pass'] for x in recs) / n) if n else 0.0 + fmt = (sum(x['skill_parseable_mean'] for x in recs) / n) if n else 0.0 + term = (sum(x['withskill_terminated_mean'] for x in recs) / n) if n else 0.0 + # 中文注释：难题子片救活率——baseline_pass==0 子集里 with-skill 的平均正确率（同 mean-over-rollouts 口径）。 + hard = [x for x in recs if not x['baseline_pass']] + hard_rescue_rate = (sum(x['withskill_acc_mean'] for x in hard) / len(hard)) if hard else 0.0 + hard_rescued = sum(x['withskill_acc_mean'] for x in hard) # 期望救活数(分数) + summary = {'record_type': 'eval_summary', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, + 'n': n, 'acc_mean1': ws, 'baseline_acc_mean1': base, 'lift_mean1': ws - base, + 'acc_pass_any': pass_any, 'n_rollouts': R, 'eval_skill_temperature': args.eval_skill_temperature, + 'format_mean1': fmt, 'term_mean1': term, + 'hard_n': len(hard), 'hard_rescued': hard_rescued, 'hard_rescue_rate': hard_rescue_rate} + metrics = {'core/math/acc/mean@1': ws, 'core/math/baseline_acc/mean@1': base, + 'core/math/lift/mean@1': ws - base, 'core/math/format/mean@1': fmt, + 'core/math/term/mean@1': term, 'core/math/hard_rescue/mean@1': hard_rescue_rate} + return recs, summary, metrics + + +# =========================================================================== +# Section I — components, args, main +# =========================================================================== +def init_components(args): + r0, r1 = TRAIN_GPUS, TRAIN_GPUS + REF_GPUS + r2, r3 = r1 + SKILL_SAMPLER_GPUS, NUM_GPUS + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, lazy_collect=False, groups=[ + DeviceGroup(name='train', ranks=list(range(0, r0)), device_type='GPU'), + DeviceGroup(name='ref', ranks=list(range(r0, r1)), device_type='GPU'), + DeviceGroup(name='skill_sampler', ranks=list(range(r1, r2)), device_type='GPU'), + DeviceGroup(name='base_sampler', ranks=list(range(r2, r3)), device_type='GPU')]) + + train_mesh = DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_DP, fsdp_size=TRAIN_FSDP) + # 主权重必须 fp32（对齐 verl actor：fp32 master + bf16 autocast）。twinkle 默认不传 dtype 时 + # transformers 会按 config 加载 bf16 主权重，lr=1e-6 的更新量(~1e-6)远小于 bf16 ulp(~4e-5)， + # optimizer.step 的更新几乎全被舍入吞掉——这是 v2 学不动/与 SEAM 对不上的根因（A/B 实测差 10-20 倍）。 + skill_model = TransformersModel(model_id=MODEL_ID, device_mesh=train_mesh, remote_group='train', + torch_dtype='float32', + ddp_config={'find_unused_parameters': False}) + skill_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) + # 方案1：skill 模型开 thinking——让 actor 先在 里解题+提炼， 只放通用方法论。 + # 由 _extract_skill 的 rfind('') 砍掉，绝不流给 executor（executor 只吃 ）， + # 因此不构成 SEAM 那种“把 think 喂给 executor”的泄漏。skill_model/ref_model/skill_sampler 三者 + # enable_thinking 必须一致，否则训练轨迹 token 布局与采样对不上。 + # 中文注释：skill_model/ref_model/skill_sampler 三者 enable_thinking 由 --skill-thinking 统一控制 + # （必须一致，否则训练轨迹 token 布局与采样对不上）；base_sampler（executor）恒 thinking on。 + _think = args.skill_thinking == 'on' + skill_model.set_template(Template, model_id=MODEL_ID, enable_thinking=_think, + max_length=args.max_model_len, truncation_strategy='delete') + skill_model.set_processor(InputProcessor, padding_free=False) + # loss 统一用 SEAM 对齐的 SEAMBNPOLoss（verl PPO clip + low_var_kl + token-mean）； + # v2/seam 两模式一致，不再随 align-mode 变。 + _loss_cls = 'SEAMBNPOLoss' + skill_model.set_loss(_loss_cls, epsilon=args.grpo_epsilon, beta=args.kl_beta) + skill_model.set_optimizer('AdamW', lr=args.lr) + # 对齐 SEAM：恒定 lr（无 warmup、无 decay）。SEAM 用 get_constant_schedule_with_warmup( + # num_warmup_steps=0)+warmup_style=constant，全程恒定 1e-6。这里直接不设 scheduler， + # skill_model.step() 对 lr_scheduler is None 有保护（transformers.py:822-824），lr 恒为 args.lr。 + # 之前的 CosineWarmupScheduler(warmup=10, cosine decay→0) 会让训练中后期有效 lr 持续衰减、 + # 更新幅度变小，与 SEAM 不一致，故移除。 + + ref_mesh = DeviceMesh.from_sizes(world_size=REF_GPUS, dp_size=REF_DP, fsdp_size=REF_FSDP) + ref_model = TransformersModel(model_id=MODEL_ID, device_mesh=ref_mesh, remote_group='ref', + ddp_config={'find_unused_parameters': False}) + ref_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) + # 方案1：与 skill_model 保持一致开 thinking（三者 enable_thinking 必须一致）。 + ref_model.set_template(Template, model_id=MODEL_ID, enable_thinking=_think, + max_length=args.max_model_len, truncation_strategy='delete') + ref_model.set_processor(InputProcessor, padding_free=False) + ref_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) + + def _sampler(group, world, enable_thinking): + s = vLLMSampler(model_id=MODEL_ID, + engine_args={'gpu_memory_utilization': GPU_MEM, + 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, + device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), + remote_group=group) + s.set_template(Template, model_id=MODEL_ID, enable_thinking=enable_thinking, max_length=args.max_model_len) + return s + + # 方案1：skill 采样器开 thinking，与 skill_model/ref_model 一致（actor 先想再写 ）。 + skill_sampler = _sampler('skill_sampler', SKILL_SAMPLER_GPUS, enable_thinking=_think) + base_sampler = _LockedSampler(_sampler('base_sampler', BASE_SAMPLER_GPUS, enable_thinking=True)) + ckpt = CheckpointEngineManager(model=skill_model, sampler=skill_sampler) + return skill_model, ref_model, skill_sampler, base_sampler, ckpt, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS + + +def _build_args(): + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--dataset', choices=('aops', 'math'), default='aops') + p.add_argument('--n', type=int, default=0, + help='Problems loaded into the draw pool (0=all; keep 0 to match a SEAM run).') + p.add_argument('--exclude-data-ids', default='', + help='Comma-separated jsonl files whose data_id/problem keys are excluded.') + p.add_argument('--seed', type=int, default=42) + p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--eval-size', type=int, default=128, help='Fixed holdout size (0 disables).') + p.add_argument('--seam-parquet-dir', type=str, default='', + help='Read SEAM build_aops_dataset.py train.parquet/val.parquet directly, in ' + 'file order (problem<-extra_info.problem, answer<-reward_model.ground_truth). ' + 'val.parquet becomes the eval holdout. Bypasses load/--numeric-only/' + '--eval-size/internal shuffle so the input data matches a SEAM run exactly.') + p.add_argument('--eval-every', type=int, default=5, help='Run holdout eval every N chunks.') + p.add_argument('--eval-rollouts', type=int, default=1, + help='Eval: skill rollouts per holdout problem. SEAM validation uses one greedy rollout.') + p.add_argument('--eval-skill-temperature', type=float, default=0.0, + help='Eval: skill-model sampling temperature. SEAM validation uses greedy T=0.') + p.add_argument('--chunk-size', type=int, default=16) + p.add_argument('--n-skills', type=int, default=8) + p.add_argument('--skill-gen-temperature', type=float, default=1.0) + p.add_argument('--skill-gen-top-p', type=float, default=1.0) + p.add_argument('--skill-gen-top-k', type=int, default=-1) + p.add_argument('--max-model-len', type=int, default=16384) + p.add_argument('--max-tokens', type=int, default=8192) + # 方案1：开 thinking 后 skill 模型要先写 分析再写 ，4096 装不下 think+完整 skills + # 会截断成空块（_extract_skill 找不到 返回 None）。提到 8192 给两段都留足空间。 + p.add_argument('--skill-max-tokens', type=int, default=8192) + # 中文注释：文体消融开关——主链路与 buffer B regen 同文体（分布一致才可联合训练）。 + p.add_argument('--skill-style', choices=('narrative', 'toy', 'pitfall'), default='narrative', + help='skill文体: narrative=现版叙述式; toy=异数字玩具题示范; pitfall=预判纠错。' + '主链路与 regen 同文体。') + p.add_argument('--skill-thinking', choices=('on', 'off'), default='on', + help='skill_model/ref_model/skill_sampler 三者的 enable_thinking（必须一致）') + p.add_argument('--align-mode', choices=('v2', 'seam'), default='v2', + help="SEAM-alignment toggle for PROMPT/SKILL FORMAT only. " + "'v2'=clean single-user executor prompt + skill-gen. " + "'seam'=nested single-user executor prompt + EXPERIENCE_PROMPT/ skill-gen. " + 'Reward (lpem numeric-only) and loss (BNPOLoss token-mean) are SEAM-style in BOTH modes.') + p.add_argument('--len-budget', type=int, default=1200, + help='Skill length budget (chars). ONLY used in distillation: drop regen skills ' + 'longer than this, pick the buffer-A seed / buffer-B survivor closest to it. ' + 'Does NOT affect GRPO reward or eval (reward = parseable AND correct).') + # --- buffer / distillation --- + p.add_argument('--distill-trigger', type=int, default=300, + help='Start draining buffer A into distillation once it reaches this many entries.') + p.add_argument('--distill-batch', type=int, default=64, + help='Entries distilled per iteration while buffer A is over --distill-trigger ' + '(incremental drain: bounds per-step latency instead of one big stall).') + p.add_argument('--sft-trigger', type=int, default=100, + help='Run one SFT pass + eval when buffer B reaches this many validated entries. ' + 'Kept low: the distill funnel (has-FAIL × valid-regen × pass@k) yields only ' + '~10-15%% of buffer A, so a high threshold would rarely fire the SFT loop.') + # Plan B validation: diversity lives in the SKILL side, the executor stays at the + # deployment (greedy) decoding口径. For each buffer-A problem we regenerate K distinct + # candidate skills (high temperature), run each through ONE greedy (T=0) executor solve, + # and accept the problem iff >= M distinct skills reach a terminated-correct solve. This + # validates "the problem admits several skills that work under greedy decoding" (matches + # eval口径) rather than "one skill passes m/k times under a high-temperature executor". + p.add_argument('--passatk-k', type=int, default=8, + help='Plan B: number of DISTINCT candidate skills regenerated per problem ' + '(skill-side diversity; executor stays greedy).') + p.add_argument('--passatk-skill-temp', type=float, default=1.0, + help='Skill-model temperature when regenerating the K candidate skills ' + '(needs >0 for diversity across candidates).') + p.add_argument('--passatk-skill-top-p', type=float, default=1.0, + help='Skill-model top-p when regenerating the K candidate skills.') + p.add_argument('--passatk-m', type=int, default=2, + help='Plan B: min number of DISTINCT candidate skills that must reach a ' + 'terminated-correct GREEDY solve to accept the problem into buffer B. ' + 'Lower than pass@k-over-one-skill (default 2): requiring m distinct ' + 'greedy-effective skills is already a strong, low-noise bar.') + p.add_argument('--distill-retries', type=int, default=1, + help='Extra regeneration rounds in distillation for [FAIL] entries that have ' + 'not yet reached m distinct greedy-passing skills (0=single pass). Each ' + 'extra round regenerates more skills (deduped) to rescue more into buffer B.') + p.add_argument('--sft-weight', type=float, default=0.5, + help='Advantage magnitude for SFT distillation samples (-w*logp + beta*KL).') + p.add_argument('--rubric-workers', type=int, default=16) + # --- GRPO --- + p.add_argument('--sft-batch-size', type=int, default=8) + p.add_argument('--ppo-mini-batch-size', type=int, default=0) + p.add_argument('--grpo-epsilon', type=float, default=0.2) + p.add_argument('--adv-clip', type=float, default=0.0, + help='clip group-relative advantage to [-adv_clip, adv_clip]; ' + '0 = no clipping (matches SEAM/verl GRPO which does not clip advantages)') + p.add_argument('--kl-beta', type=float, default=0.001) + p.add_argument('--lr', type=float, default=6e-6) + p.add_argument('--max-train-rounds', type=int, default=1500) + p.add_argument('--save-rounds', type=int, default=200) + p.add_argument('--output-dir', default='./output/skill_v2') + p.add_argument('--cache-dir', default='') + p.add_argument('--no-cache', action='store_true') + p.add_argument('--swanlab-project', default='twinkle') + p.add_argument('--swanlab-exp', default='') + args = p.parse_args() + if args.sft_batch_size % TRAIN_DP != 0: + raise ValueError(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple of train dp ({TRAIN_DP})') + if args.chunk_size < 1: + raise ValueError('--chunk-size must be >= 1') + return args + + +def _write(handle, row): + handle.write(json.dumps(row, ensure_ascii=False) + '\n') + + +def _swan_metrics(summary, log): + # Lean metric set: each carries independent information. Dropped as redundant — + # 中文注释：删除冗余项（换算重复）：n_groups(≈chunk_size)、reward_std(池化,组内方差已够)、 + # skill_tokens_mean(与chars重复)、leak/n(=rate×n)、candidate_withskill(与问题级重复)、 + # term/withskill(=1-trunc)、train/n_steps(恒为1)。 + d = { + 'signal/zero_grad_frac': summary['zero_grad_frac'], + 'signal/reward_mean': summary['reward_mean'], + 'signal/group_reward_std_mean': summary['group_reward_std_mean'], + 'signal/n_train_samples': summary['n_train_samples'], + 'skill/parse_rate': summary['parse_rate'], 'skill/chars_mean': summary['skill_chars_mean'], + 'leak/rate': summary['leak_rate'], + } + if summary['n_groups'] > 0: + d.update({'acc/withskill_pass': summary['avg_withskill_pass'], + 'term/withskill_trunc_frac': summary['withskill_trunc_frac']}) + if log: + d['train/n_grpo'] = log['n_grpo'] + d['train/n_sft'] = log['n_sft'] + for k, v in (log.get('metric') or {}).items(): + if not _is_num(v): + continue + if k.startswith('learning rate'): + if 'group 1' in k: + d['train/lr'] = float(v) + else: + d[f'train/{k.replace(" ", "_")}'] = float(v) + return d + + +def main(): + args = _build_args() + global _ALIGN_MODE, _SKILL_STYLE + _ALIGN_MODE = args.align_mode # 'v2' | 'seam' + _SKILL_STYLE = args.skill_style # 'narrative' | 'toy' | 'pitfall' + records, eval_records = _load_records(args) + if len(records) < args.chunk_size: + raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded ({len(records)}); raise --n') + + os.makedirs(args.output_dir, exist_ok=True) + gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') + eval_path = os.path.join(args.output_dir, 'eval_records.jsonl') + sft_path = os.path.join(args.output_dir, 'skill_dataset.jsonl') + train_log_path = os.path.join(args.output_dir, 'train_log.jsonl') + buffer_a_path = os.path.join(args.output_dir, 'buffer_a.jsonl') + distill_path = os.path.join(args.output_dir, 'distill_records.jsonl') + + use_swan = swanlab is not None and os.environ.get('SWANLAB_MODE') != 'disabled' + if use_swan: + swanlab.init(project=args.swanlab_project, experiment_name=(args.swanlab_exp or None), + config={'model': MODEL_ID, 'dataset': args.dataset, 'n': len(records), + 'eval_n': len(eval_records), 'n_skills': args.n_skills, + 'len_budget': args.len_budget, 'distill_trigger': args.distill_trigger, + 'sft_trigger': args.sft_trigger, 'passatk_k': args.passatk_k, + 'passatk_m': args.passatk_m, 'passatk_skill_temp': args.passatk_skill_temp, + 'sft_weight': args.sft_weight, 'lr': args.lr, 'align_mode': args.align_mode}) + + skill_model, ref_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp = init_components(args) + checker = build_rubric_checker() + if checker is None: + sys.stderr.write('[v2] no LLM backup env -> buffer B distillation DISABLED (GRPO only)\n') + + cache_dir = args.cache_dir or os.path.join(args.output_dir, 'cache') + os.makedirs(cache_dir, exist_ok=True) + # 每次启动强制重算 eval baseline：旧缓存可能来自不同环境/代码版本（torch/vllm/dtype 均影响 T=0 输出）， + # 跨 run 复用会造成 with-skill（现算）vs baseline（陈旧）不可比，lift 虚高/虚低（已实锤过一次）。 + _base_cache_path = os.path.join(cache_dir, 'eval_baseline.jsonl') + if os.path.exists(_base_cache_path): + os.remove(_base_cache_path) + logger.info('stale eval_baseline cache removed (recomputed this run)') + eval_base_cache = DiskCache(_base_cache_path, not args.no_cache) + + cfg = {'record_type': 'config', 'model': MODEL_ID, 'dataset': args.dataset, + 'n': len(records), 'eval_n': len(eval_records), 'seed': args.seed, + 'n_skills': args.n_skills, 'len_budget': args.len_budget, + 'distill_trigger': args.distill_trigger, 'sft_trigger': args.sft_trigger, + 'passatk_k': args.passatk_k, 'passatk_m': args.passatk_m, + 'passatk_skill_temp': args.passatk_skill_temp, 'passatk_skill_top_p': args.passatk_skill_top_p, + 'sft_weight': args.sft_weight, + 'grpo_epsilon': args.grpo_epsilon, 'kl_beta': args.kl_beta, 'lr': args.lr, + 'align_mode': args.align_mode, + 'rubric_check': bool(checker), 'max_train_rounds': args.max_train_rounds, + 'seam_parquet_dir': (getattr(args, 'seam_parquet_dir', '') or ''), + 'started': int(time.time())} + + hist_a: List[Dict[str, Any]] = [] # buffer A accumulator (in-memory + jsonl) + sft_queue: List[Dict[str, Any]] = [] # buffer B: validated SFT records awaiting an SFT pass + rounds = 0 # GRPO rounds only (gates --max-train-rounds + save cadence) + sft_rounds = 0 # SFT passes (separate: must NOT eat the GRPO round budget) + pool = ProblemPool(records, args.seed) + + # Background rubric pre-diagnosis (做法 B): the moment a failure trajectory lands in + # buffer A, fire its teacher-rubric call on a daemon thread pool. The API round-trip + # then overlaps with GRPO GPU work, so by the time --distill-trigger fires the + # diagnoses are usually already cached on each entry ('_rubric_diag'); distill_buffer + # only pays for the stragglers. Entries are dicts held by reference, so the worker + # writes the result straight onto the entry. + # 中文注释：失败轨迹一进 buffer A 就后台异步跑 rubric，API 等待藏进 GPU 训练时间； + # 到蒸馏时诊断多已缓存在条目上，distill_buffer 只补漏。 + prediag_pool = (ThreadPoolExecutor(max_workers=max(1, args.rubric_workers), + thread_name_prefix='rubric-prediag') + if checker else None) + + def _prediagnose(entry: Dict[str, Any]): + entry['_rubric_diag'] = _diagnose_entry(checker, entry) or '' + + with open(gen_path, 'w', encoding='utf-8') as gen_f, \ + open(eval_path, 'w', encoding='utf-8') as eval_f, \ + open(sft_path, 'w', encoding='utf-8') as sft_f, \ + open(train_log_path, 'w', encoding='utf-8') as tlog, \ + open(distill_path, 'w', encoding='utf-8') as distill_f, \ + open(buffer_a_path, 'w', encoding='utf-8') as buf_f: + for f in (gen_f, eval_f, sft_f, tlog, distill_f): + _write(f, cfg) + + def _do_eval(gstep): + recs, summary, metrics = run_greedy_eval( + base_sampler, skill_sampler, eval_records, gstep, rounds, base_dp, skill_dp, + args, eval_base_cache) + for rec in recs: + _write(eval_f, rec) + _write(eval_f, summary) + eval_f.flush() + if use_swan: + swanlab.log({f'eval/{k}': v for k, v in metrics.items()}, step=max(gstep, 0)) + sys.stderr.write( + f'[eval] g{gstep}: n={summary["n"]} acc={summary["baseline_acc_mean1"]:.3f}' + f'->{summary["acc_mean1"]:.3f} lift={summary["lift_mean1"]:+.3f} ' + f'hard_rescue={summary["hard_rescue_rate"]:.3f}({summary["hard_rescued"]}/{summary["hard_n"]}) ' + f'fmt={summary["format_mean1"]:.2f} rounds={rounds}\n') + + if eval_records: + _do_eval(-1) + + gstep = 0 + while rounds < args.max_train_rounds: + chunk = pool.draw(args.chunk_size) + full, summary, grpo, buffer_a = process_chunk( + base_sampler, skill_sampler, chunk, gstep, base_dp, skill_dp, args) + + # accumulate buffer A (only when a rubric checker exists to consume it; + # 中文注释：无 checker 时蒸馏永不触发，不累积以免内存无限增长) + if checker: + for e in buffer_a: + _write(buf_f, e) + prediag_pool.submit(_prediagnose, e) # 后台异步预诊断，不阻塞主循环 + buf_f.flush() + hist_a.extend(buffer_a) + + # GRPO train step (only when there is signal) + log = None + if grpo: + log = _train_step(skill_model, ref_model, ckpt, grpo, args) + rounds += 1 + log.update({'record_type': 'train_round', 'round': rounds, 'chunk': gstep, + 'epoch': pool.epoch, 'kind': 'grpo', 'ts': int(time.time())}) + _write(tlog, log) + tlog.flush() + if rounds % args.save_rounds == 0: + skill_model.save(f'skill-v2-{rounds}', output_dir=args.output_dir) + + summary['rounds_done'], summary['epoch'] = rounds, pool.epoch + summary['buffer_a_size'], summary['sft_queue_size'] = len(hist_a), len(sft_queue) + for rec in full: + _write(gen_f, rec) + _write(gen_f, summary) + gen_f.flush() + + sys.stderr.write( + f'[gen] e{pool.epoch} g{gstep}: n={summary["n"]} ' + f'clean={summary["n_candidates_parseable"]} 0grad={summary["zero_grad_frac"]:.2f} ' + f'R={summary["reward_mean"]:.2f}+-{summary["reward_std"]:.2f} ' + f'ws_acc={summary["avg_withskill_pass"]:.2f} chars={summary["skill_chars_mean"]:.0f} ' + f'bufA={len(hist_a)} bufB={len(sft_queue)} rounds={rounds}\n') + if use_swan: + m = _swan_metrics(summary, log) + m['buffer/a_size'] = float(len(hist_a)) + m['buffer/b_size'] = float(len(sft_queue)) + swanlab.log(m, step=gstep) + + # --- distillation: once buffer A fills, drain it INCREMENTALLY in bounded + # batches (--distill-batch) so a large buffer never stalls the loop for tens + # of minutes; each iteration processes one batch, interleaved with GRPO. + # 中文注释：增量分批蒸馏——buffer A 满后每轮只处理 --distill-batch 条，把一次性 + # 几十分钟阻塞摊成每轮几分钟小停顿；两段验证(见 distill_buffer)再砍验证算力。 + if checker and len(hist_a) >= args.distill_trigger: + batch = hist_a[:args.distill_batch] + hist_a = hist_a[args.distill_batch:] + new_sft, distill_recs = distill_buffer(batch, skill_sampler, base_sampler, checker, + skill_dp, base_dp, args) + for rec in distill_recs: # 逐 entry 审计记录：rubric_diag + 候选 skill + 贪心解 + stage + rec['chunk'] = gstep + _write(distill_f, rec) + distill_f.flush() + for rec in new_sft: + _write(sft_f, rec) + sft_f.flush() + sft_queue.extend(new_sft) + + # --- SFT trigger: buffer B full → one SFT pass + eval --- + did_eval = False + if len(sft_queue) >= args.sft_trigger: + sys.stderr.write(f'[sft] triggered at bufB={len(sft_queue)}\n') + sft_samples = [{**s, 'advantage': float(args.sft_weight)} for s in sft_queue] + sft_log = _train_step(skill_model, ref_model, ckpt, sft_samples, args) + sft_rounds += 1 # 中文注释：SFT 用独立计数，不占用 GRPO 的 rounds 配额/save 节奏 + sft_log.update({'record_type': 'train_round', 'round': rounds, 'sft_round': sft_rounds, + 'chunk': gstep, 'epoch': pool.epoch, 'kind': 'sft', 'ts': int(time.time())}) + _write(tlog, sft_log) + tlog.flush() + sft_queue = [] + skill_model.save(f'skill-v2-sft{sft_rounds}', output_dir=args.output_dir) # 大改动后落盘 + if eval_records: # 中文注释：SFT 后立即 eval，测灾难性遗忘/真提升（第 11.3/13.4 节） + _do_eval(gstep) + did_eval = True + + if eval_records and not did_eval and (gstep + 1) % args.eval_every == 0: + _do_eval(gstep) + gstep += 1 + + if prediag_pool is not None: + prediag_pool.shutdown(wait=False, cancel_futures=True) # 丢弃未完成的后台预诊断 + eval_base_cache.close() + skill_model.save('skill-v2-final', output_dir=args.output_dir) + sys.stderr.write(f'[v2] done: {rounds} rounds over {gstep} chunks / {pool.epoch} epochs\n') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/skill2lora/train_skill_v2.sh b/cookbook/exp/skill2lora/train_skill_v2.sh new file mode 100644 index 000000000..2f7e4573a --- /dev/null +++ b/cookbook/exp/skill2lora/train_skill_v2.sh @@ -0,0 +1,69 @@ +#!/bin/bash +# train_skill_v2.sh — 简化 GRPO + buffer distill 训练启动脚本 +# 用法: bash cookbook/exp/skill2lora/train_skill_v2.sh +# +# 环境变量: +# LLM_BACKUP_API_KEY - rubric 诊断用的教师 API key（必须，否则 buffer B 蒸馏不可用） +# LLM_BACKUP_BASE_URL - 教师 API base URL +# LLM_BACKUP_MODEL - 教师模型 ID +# GEN_MODEL_ID - 训练 skill 模型 ID（默认 Qwen/Qwen3-4B） +# TRAIN_GPUS / REF_GPUS / SKILL_SAMPLER_GPUS / BASE_SAMPLER_GPUS — GPU 分配 + +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +# 缓解显存碎片（reserved-but-unallocated），降低 forward_backward 阶段 OOM 概率 +export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" + +# 默认输出目录 +OUTPUT_DIR="${OUTPUT_DIR:-./output/skill_v2}" + +# 去重/排斥数据（冷启动 SFT 数据避免重叠） +EXCLUDE="${EXCLUDE_DATA_IDS:-}" + +# 与 SEAM 输入数据对齐：直读 SEAM build_aops_dataset.py 产出的 parquet（同题池 + 同 val） +# 置空则回退到 twinkle 自己的 load+shuffle+split。 +SEAM_PARQUET_DIR="${SEAM_PARQUET_DIR:-/root/data/seam}" + +# 提前建目录：tee 需在 python 建目录前就能打开日志文件 +mkdir -p "${OUTPUT_DIR}" + +python3 "${SCRIPT_DIR}/train_skill_v2.py" \ + --dataset aops \ + --numeric-only \ + --eval-size 200 \ + --eval-every 5 \ + --eval-rollouts 1 \ + --eval-skill-temperature 0.0 \ + --chunk-size 16 \ + --n-skills 8 \ + --distill-retries 1 \ + --skill-gen-temperature 1.0 \ + --skill-gen-top-p 1.0 \ + --skill-gen-top-k -1 \ + --max-model-len 16384 \ + --max-tokens 8192 \ + --skill-max-tokens 4096 \ + --len-budget 600 \ + --distill-trigger 150 \ + --distill-batch 64 \ + --sft-trigger 100 \ + --passatk-k 8 \ + --passatk-m 2 \ + --align-mode seam \ + --sft-weight 1.0 \ + --rubric-workers 16 \ + --sft-batch-size 4 \ + --ppo-mini-batch-size 0 \ + --grpo-epsilon 0.2 \ + --adv-clip 0 \ + --kl-beta 0.001 \ + --lr 1e-6 \ + --max-train-rounds 1500 \ + --save-rounds 200 \ + --output-dir "${OUTPUT_DIR}" \ + --swanlab-project twinkle \ + --swanlab-exp "skill_v2_$(date +%Y%m%d_%H%M%S)" \ + ${EXCLUDE:+--exclude-data-ids "${EXCLUDE}"} \ + ${SEAM_PARQUET_DIR:+--seam-parquet-dir "${SEAM_PARQUET_DIR}"} \ + "$@" 2>&1 | tee "${OUTPUT_DIR}/run.log" From 3e81bff6ffda84ec8aaf5784ae6ff4d7a7d6f2eb Mon Sep 17 00:00:00 2001 From: root Date: Sun, 26 Jul 2026 18:06:06 +0800 Subject: [PATCH 25/60] wip --- pyproject.toml | 2 +- src/twinkle/loss/__init__.py | 6 +- src/twinkle/loss/grpo.py | 99 +++++++++++++++++++ src/twinkle/loss/opsd.py | 126 ++++++++++++++++++++++++ src/twinkle/patch/vllm_lora_weights.py | 25 ++++- tests/loss/test_opsd.py | 131 +++++++++++++++++++++++++ 6 files changed, 386 insertions(+), 3 deletions(-) create mode 100644 src/twinkle/loss/opsd.py create mode 100644 tests/loss/test_opsd.py diff --git a/pyproject.toml b/pyproject.toml index 12fb2d929..5fee2cab0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ version = "0.4.0.dev0" description = "Training API for large language models with efficient data handling and advanced optimization techniques." readme = "README.md" authors = [{ name = "ModelScope", email = "contact@modelscope.cn" }] -requires-python = ">=3.11,<3.13" +requires-python = ">=3.10,<3.13" dependencies = [ "numpy>=2.0.0,<2.3.0", "datasets", diff --git a/src/twinkle/loss/__init__.py b/src/twinkle/loss/__init__.py index 8e1d0e2ad..c834b5630 100644 --- a/src/twinkle/loss/__init__.py +++ b/src/twinkle/loss/__init__.py @@ -4,9 +4,10 @@ from .cross_entropy import CrossEntropyLoss from .dpo import CPOLoss, DPOLoss, ORPOLoss, SimPOLoss from .gkd import GKDLoss -from .grpo import BNPOLoss, CISPOLoss, DRGRPOLoss, GRPOLoss, GSPOLoss, SAPOLoss +from .grpo import BNPOLoss, CISPOLoss, DRGRPOLoss, GRPOLoss, GSPOLoss, SAPOLoss, SEAMBNPOLoss from .infonce import InfonceLoss from .mse import MSELoss +from .opsd import OPSDLoss torch_loss_mapping = { 'mse': MSELoss, @@ -20,7 +21,10 @@ 'sapo': SAPOLoss, 'cispo': CISPOLoss, 'bnpo': BNPOLoss, + 'seam_bnpo': SEAMBNPOLoss, 'dr_grpo': DRGRPOLoss, + # Self-distillation losses + 'opsd': OPSDLoss, # DPO family losses 'dpo': DPOLoss, 'simpo': SimPOLoss, diff --git a/src/twinkle/loss/grpo.py b/src/twinkle/loss/grpo.py index 7fb799eca..a01ecda84 100644 --- a/src/twinkle/loss/grpo.py +++ b/src/twinkle/loss/grpo.py @@ -422,6 +422,105 @@ def _aggregate_loss( return (per_token_loss * loss_mask).sum() / loss_mask.sum().clamp(min=1.0) +class SEAMBNPOLoss(BNPOLoss): + """BNPO token-mean loss with verl/SEAM policy-loss clipping semantics.""" + + def __init__( + self, + epsilon: float = 0.2, + epsilon_high: Optional[float] = None, + beta: float = 0.0, + entropy_coef: float = 0.0, + clip_ratio_c: float = 3.0, + ignore_index: int = -100, + **kwargs, + ): + super().__init__(epsilon=epsilon, epsilon_high=epsilon_high, beta=beta, + entropy_coef=entropy_coef, ignore_index=ignore_index, **kwargs) + self.clip_ratio_c = clip_ratio_c + + def _compute_log_importance_weights( + self, + per_token_logps: 'torch.Tensor', + per_token_old_logps: 'torch.Tensor', + loss_mask: 'torch.Tensor', + ) -> 'torch.Tensor': + """Match verl.core_algos.compute_policy_loss clamp range.""" + import torch + return torch.clamp(per_token_logps - per_token_old_logps, min=-20.0, max=20.0) + + def _compute_per_token_loss( + self, + ratio: 'torch.Tensor', + advantages: 'torch.Tensor', + per_token_logps: 'torch.Tensor', + ) -> 'torch.Tensor': + """Match verl.core_algos.compute_policy_loss for PPO clipped policy loss.""" + import torch + pg_losses1 = -advantages * ratio + pg_losses2 = -advantages * torch.clamp(ratio, 1 - self.epsilon, 1 + self.epsilon_high) + clip_pg_losses1 = torch.maximum(pg_losses1, pg_losses2) + pg_losses3 = -advantages * self.clip_ratio_c + clip_pg_losses2 = torch.minimum(pg_losses3, clip_pg_losses1) + return torch.where(advantages < 0, clip_pg_losses2, clip_pg_losses1) + + def __call__( + self, + inputs: Dict, + outputs: Dict, + *, + old_logps: Optional[Union['torch.Tensor', List[List[float]]]] = None, + ref_logps: Optional['torch.Tensor'] = None, + advantages: Optional[Union['torch.Tensor', List[float], np.ndarray]] = None, + **kwargs, + ): + """Same as GRPOLoss.__call__, but with verl low_var_kl clipping for the KL loss.""" + import torch + labels = inputs.get('labels') + assert labels is not None, "inputs must contain 'labels'" + if not torch.is_tensor(labels): + labels = torch.as_tensor(labels) + if labels.dim() == 1: + labels = labels.unsqueeze(0) + + logps = outputs.get('logps') + loss_mask = (labels != self.ignore_index).bool() + if logps is None: + logits = outputs.get('logits') + if logits.shape[1] != labels.shape[1]: + logits = logits[:, -labels.shape[1]:] + masked_labels = labels.clone() + masked_labels[~loss_mask] = 0 + logps = selective_log_softmax(logits, masked_labels) + + device = logps.device + old_logps = logps.detach() if old_logps is None else self._pad_and_align_to_batch( + old_logps, loss_mask, device, logps.dtype) + if ref_logps is not None: + ref_logps = self._pad_and_align_to_batch(ref_logps, loss_mask, device, logps.dtype) + if advantages is None: + return LossOutput(loss=logps.sum() * 0.0, num_tokens=0) + advantages = self._pad_and_align_to_batch(advantages, loss_mask, device, logps.dtype) + + log_importance_weights = self._compute_log_importance_weights(logps, old_logps, loss_mask) + ratio = torch.exp(log_importance_weights) + per_token_loss = self._compute_per_token_loss(ratio, advantages, logps) + + if self.beta > 0.0 and ref_logps is not None: + kl = torch.clamp(ref_logps - logps, min=-20.0, max=20.0) + per_token_kl = torch.clamp(torch.exp(kl) - kl - 1, min=-10.0, max=10.0) + per_token_loss = per_token_loss + self.beta * per_token_kl + + if self.entropy_coef > 0.0: + entropies = outputs.get('entropies') + assert entropies is not None, ('entropy_coef > 0 requires outputs[\'entropies\'] — make sure the ' + "loss instance's require_entropy flag was set before the forward call.") + per_token_loss = per_token_loss - self.entropy_coef * entropies.to(per_token_loss.dtype) + + loss = self._aggregate_loss(per_token_loss, loss_mask, **kwargs) + return LossOutput(loss=loss, num_tokens=0) + + class DRGRPOLoss(GRPOLoss): """ DR-GRPO (Dynamic Ratio GRPO) Loss. diff --git a/src/twinkle/loss/opsd.py b/src/twinkle/loss/opsd.py new file mode 100644 index 000000000..b59fd8e18 --- /dev/null +++ b/src/twinkle/loss/opsd.py @@ -0,0 +1,126 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from typing import TYPE_CHECKING, Dict, List, Optional, Union + +from twinkle.data_format import LossOutput +from twinkle.loss.grpo import GRPOLoss + +if TYPE_CHECKING: + import torch + + +class OPSDLoss(GRPOLoss): + """On-Policy Self-Distillation (OPSD) loss. + + Reference: + "Self-Distilled Reasoner: On-Policy Self-Distillation for Large Language Models" + (Zhao et al., arXiv:2601.18734). + + A single model acts as BOTH teacher and student, differing only in context: + * student policy conditions on the QUESTION ONLY (query-only prompt); + * teacher policy conditions on PRIVILEGED information (question + rubric diagnosis). + Training minimizes a per-token divergence between the two distributions over the + STUDENT's own on-policy rollout (the tokens the student generated under the + query-only prompt). Because both forwards score the SAME response tokens, only the + prompt differs, so the per-token alignment is exact. + + Token-probability (sampled-token) form — v1, zero extra tensor channel + ---------------------------------------------------------------------- + We only need the per-token log-prob of the SAMPLED tokens from each context + (``teacher_logps`` from a teacher forward on the rubric-conditioned trajectory, + ``logps`` from the student forward on the query-only trajectory). Reusing the exact + k3 estimator already used by the GRPO KL penalty + (``grpo.py``: ``exp(ref - logps) - (ref - logps) - 1``), the per-token loss is:: + + r = teacher_logp - student_logp # teacher detached + per_token = exp(r) - r - 1 # k3 estimate, >= 0, pulls student -> teacher + + Its gradient w.r.t. the student log-prob is ``1 - exp(r)``: when the teacher assigns + higher probability than the student (``r > 0``) the update RAISES the student log-prob + toward the teacher, and lowers it when ``r < 0`` — a dense token-level distillation + pull, no advantages / reward needed. + + Aggregation is BNPO-style token-mean (sum over all response tokens / total token count), + matching the RL branch so OPSD and BNPO experiments share the same effective step scaling. + + Notes + ----- + * ``teacher_logps`` is accepted via a dedicated kwarg; for pipelines that route the teacher + log-probs through the existing reference channel it also falls back to ``ref_logps``. + Provide it in the RESPONSE-ONLY form (one log-prob per trainable/response token, matching + the student loss mask) — ``_pad_and_align_to_batch`` scatters it onto the response + positions. The teacher and student prompts differ in length, so the full-sequence + (right-padded) form must NOT be used here. + * The divergence direction (this k3 form corresponds to KL(student || teacher)) should be + re-confirmed against the official code release before treating it as final; it is exposed + via ``reverse`` for a quick swap without touching call sites. + """ + + require_logps = True + require_logits = False + + def __init__( + self, + beta: float = 0.0, + ignore_index: int = -100, + reverse: bool = True, + **kwargs, + ): + # epsilon is unused (no PPO ratio here) but kept in the ctor so the shared + # ``set_loss(epsilon=..., beta=...)`` call site does not need special-casing. + super().__init__(epsilon=kwargs.pop('epsilon', 0.2), beta=beta, + ignore_index=ignore_index, **kwargs) + self.reverse = reverse + + def _aggregate_loss(self, per_token_loss, loss_mask, **kwargs): + """BNPO-style token-mean: sum over all response tokens / total token count.""" + return (per_token_loss * loss_mask).sum() / loss_mask.sum().clamp(min=1.0) + + def __call__( + self, + inputs: Dict, + outputs: Dict, + *, + teacher_logps: Optional[Union['torch.Tensor', List[List[float]]]] = None, + ref_logps: Optional[Union['torch.Tensor', List[List[float]]]] = None, + **kwargs, + ) -> LossOutput: + import torch + + labels = inputs.get('labels') + assert labels is not None, "inputs must contain 'labels'" + if not torch.is_tensor(labels): + labels = torch.as_tensor(labels) + if labels.dim() == 1: + labels = labels.unsqueeze(0) + + logps = outputs.get('logps') + loss_mask = (labels != self.ignore_index).bool() + if logps is None: + from twinkle.utils.torch_utils import selective_log_softmax + logits = outputs.get('logits') + if logits.shape[1] != labels.shape[1]: + logits = logits[:, -labels.shape[1]:] + masked_labels = labels.clone() + masked_labels[~loss_mask] = 0 + logps = selective_log_softmax(logits, masked_labels) + + device = logps.device + + # Teacher log-probs: prefer the dedicated kwarg, else reuse the reference channel. + teacher = teacher_logps if teacher_logps is not None else ref_logps + # Without a teacher this reduces to a no-op that still flows through autograd, so + # ref-only / eval forwards (which harvest outputs['logps']) do not crash and DDP/FSDP + # never see unused parameters. Mirrors GRPOLoss's advantages-None guard. + if teacher is None: + return LossOutput(loss=logps.sum() * 0.0, num_tokens=0) + + teacher = self._pad_and_align_to_batch(teacher, loss_mask, device, logps.dtype) + teacher = teacher.detach() + + # r = teacher - student. k3 KL estimate: exp(r) - r - 1 (>= 0), pulls student -> teacher. + r = teacher - logps if self.reverse else logps - teacher + r = torch.clamp(r, min=-10.0, max=10.0) # guard exp overflow on rare huge gaps + per_token_loss = torch.exp(r) - r - 1 + + loss = self._aggregate_loss(per_token_loss, loss_mask, **kwargs) + return LossOutput(loss=loss, num_tokens=0) diff --git a/src/twinkle/patch/vllm_lora_weights.py b/src/twinkle/patch/vllm_lora_weights.py index 558c03892..cd905b91c 100644 --- a/src/twinkle/patch/vllm_lora_weights.py +++ b/src/twinkle/patch/vllm_lora_weights.py @@ -130,6 +130,29 @@ def patched_load_adapter(self: LRUCacheWorkerLoRAManager, lora_request: TensorLo f'lora_extra_vocab_size {self.lora_config.lora_extra_vocab_size}.') return lora + # Cache the cache-wrapped template tokenizer (keyed by id) so we wrap once, not per request. + _wrapped_tok_cache: Dict[int, object] = {} + + def _ensure_max_token_id(tokenizer): + """ + vllm's Processor._validate_model_input reads ``tokenizer.max_token_id``, an attribute + that only exists on vllm's ``CachedTokenizer`` wrapper. The sampler template tokenizer is + a RAW HF tokenizer (never passed through vllm's ``get_cached_tokenizer``), so validation + raises ``AttributeError: ... has no attribute max_token_id``. Wrap it once to add the attr. + """ + if tokenizer is None or hasattr(tokenizer, 'max_token_id'): + return tokenizer + key = id(tokenizer) + wrapped = _wrapped_tok_cache.get(key) + if wrapped is None: + try: + from vllm.transformers_utils.tokenizer import get_cached_tokenizer + wrapped = get_cached_tokenizer(tokenizer) + except Exception: + wrapped = tokenizer + _wrapped_tok_cache[key] = wrapped + return wrapped + def patched_get_lora_tokenizer(self: TokenizerGroup, lora_request: LoRARequest): # since we pass dummy path, skip get tokenizer from path # Use lazy tokenizer access @@ -137,7 +160,7 @@ def patched_get_lora_tokenizer(self: TokenizerGroup, lora_request: LoRARequest): if tokenizer is None: # Fallback to the original method if tokenizer not available return self._old_get_lora_tokenizer(lora_request) - return tokenizer + return _ensure_max_token_id(tokenizer) if not hasattr(LRUCacheWorkerLoRAManager, '_old_load_adapter'): _old_load_adapter = LRUCacheWorkerLoRAManager._load_adapter diff --git a/tests/loss/test_opsd.py b/tests/loss/test_opsd.py new file mode 100644 index 000000000..011c3aff9 --- /dev/null +++ b/tests/loss/test_opsd.py @@ -0,0 +1,131 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Tests for OPSDLoss (On-Policy Self-Distillation, arXiv:2601.18734).""" +import pytest +import torch +import torch.nn.functional as F + +from twinkle.loss import OPSDLoss +from twinkle.loss import torch_loss_mapping + + +def _make_opsd_batch(batch_size=4, seq_len=8, vocab_size=20, gap=0.0): + """Synthetic batch: student logps + teacher logps shifted by `gap` on valid tokens.""" + torch.manual_seed(42) + logits = torch.randn(batch_size, seq_len, vocab_size) + labels = torch.randint(0, vocab_size, (batch_size, seq_len)) + for i in range(batch_size): + labels[i, seq_len // 2:] = -100 # first half = response tokens, rest ignored + + loss_mask = (labels != -100) + masked_labels = labels.clone() + masked_labels[~loss_mask] = 0 + logps = F.log_softmax(logits, dim=-1).gather(-1, masked_labels.unsqueeze(-1)).squeeze(-1) + teacher_logps = logps.detach() + gap + + inputs = {'labels': labels} + outputs = {'logps': logps} + return inputs, outputs, teacher_logps, loss_mask + + +class TestOPSDLoss: + + def test_basic_finite_scalar(self): + loss_fn = OPSDLoss() + inputs, outputs, teacher, _ = _make_opsd_batch(gap=0.3) + result = loss_fn(inputs, outputs, teacher_logps=teacher) + assert isinstance(result, dict) and 'loss' in result + assert result['loss'].dim() == 0 + assert torch.isfinite(result['loss']) + + def test_zero_loss_when_teacher_equals_student(self): + """k3 estimate exp(r) - r - 1 == 0 exactly when r == 0.""" + loss_fn = OPSDLoss() + inputs, outputs, teacher, _ = _make_opsd_batch(gap=0.0) + result = loss_fn(inputs, outputs, teacher_logps=teacher) + assert torch.allclose(result['loss'], torch.tensor(0.0), atol=1e-6) + + def test_loss_positive_when_gap_nonzero(self): + loss_fn = OPSDLoss() + for gap in (0.5, -0.5): + inputs, outputs, teacher, _ = _make_opsd_batch(gap=gap) + result = loss_fn(inputs, outputs, teacher_logps=teacher) + assert result['loss'].item() > 0.0 + + def test_gradient_pulls_student_toward_teacher(self): + """teacher logp higher (r>0) -> d(loss)/d(student_logp) < 0 -> SGD raises student logp.""" + logps = torch.zeros(1, 4, requires_grad=True) + labels = torch.tensor([[1, 1, -100, -100]]) + teacher = torch.full((1, 4), 0.0) + teacher[0, :2] = 0.7 # teacher more confident on the two valid tokens + loss_fn = OPSDLoss() + out = loss_fn({'labels': labels}, {'logps': logps}, teacher_logps=teacher) + out['loss'].backward() + # gradient on valid tokens must be negative (increase logps), zero on masked tokens + assert (logps.grad[0, :2] < 0).all() + assert torch.allclose(logps.grad[0, 2:], torch.zeros(2)) + + def test_gradient_direction_flips_when_teacher_lower(self): + logps = torch.zeros(1, 4, requires_grad=True) + labels = torch.tensor([[1, 1, -100, -100]]) + teacher = torch.full((1, 4), -0.7) # teacher LESS confident + loss_fn = OPSDLoss() + out = loss_fn({'labels': labels}, {'logps': logps}, teacher_logps=teacher) + out['loss'].backward() + assert (logps.grad[0, :2] > 0).all() # SGD lowers student logp + + def test_masked_tokens_do_not_contribute(self): + """Changing teacher values on ignored positions must not change the loss.""" + loss_fn = OPSDLoss() + inputs, outputs, teacher, loss_mask = _make_opsd_batch(gap=0.3) + r1 = loss_fn(inputs, outputs, teacher_logps=teacher.clone()) + teacher2 = teacher.clone() + teacher2[~loss_mask] += 123.0 + r2 = loss_fn(inputs, outputs, teacher_logps=teacher2) + assert torch.allclose(r1['loss'], r2['loss']) + + def test_response_only_ragged_list_form(self): + """Teacher logps as ragged per-sample lists (response tokens only) must align to the mask. + + This is the production form: the teacher forward uses a DIFFERENT (rubric) prompt, so + only the response-token log-probs are extracted and passed per sample.""" + loss_fn = OPSDLoss() + inputs, outputs, teacher, loss_mask = _make_opsd_batch(gap=0.4) + ragged = [teacher[i][loss_mask[i]].tolist() for i in range(teacher.shape[0])] + r_full = loss_fn(inputs, outputs, teacher_logps=teacher) + r_ragged = loss_fn(inputs, outputs, teacher_logps=ragged) + assert torch.allclose(r_full['loss'], r_ragged['loss'], atol=1e-5) + + def test_ref_logps_channel_fallback(self): + """teacher_logps may ride the existing ref_logps channel (zero new tensor plumbing).""" + loss_fn = OPSDLoss() + inputs, outputs, teacher, _ = _make_opsd_batch(gap=0.3) + r_kw = loss_fn(inputs, outputs, teacher_logps=teacher) + r_ref = loss_fn(inputs, outputs, ref_logps=teacher) + assert torch.allclose(r_kw['loss'], r_ref['loss']) + + def test_no_teacher_returns_zero_flowing_loss(self): + """No teacher -> zero loss that still flows through autograd (ref-only forwards).""" + logps = torch.randn(2, 6, requires_grad=True) + labels = torch.randint(0, 10, (2, 6)) + loss_fn = OPSDLoss() + out = loss_fn({'labels': labels}, {'logps': logps}) + assert out['loss'].item() == 0.0 + out['loss'].backward() # must not raise + assert logps.grad is not None + + def test_clamp_guards_extreme_gap(self): + loss_fn = OPSDLoss() + inputs, outputs, teacher, _ = _make_opsd_batch(gap=50.0) + result = loss_fn(inputs, outputs, teacher_logps=teacher) + assert torch.isfinite(result['loss']) + + def test_registered_in_mapping(self): + assert torch_loss_mapping.get('opsd') is OPSDLoss + + def test_requires_logps_not_logits(self): + assert OPSDLoss.require_logps is True + assert OPSDLoss.require_logits is False + + +if __name__ == '__main__': + pytest.main([__file__, '-v']) From ac75ba6bb5e6174ee03c43104e6ef18474bb9cf9 Mon Sep 17 00:00:00 2001 From: root Date: Tue, 28 Jul 2026 10:39:04 +0800 Subject: [PATCH 26/60] wip --- .../build_reflexion_coldstart_sft.py | 519 ----- .../embedding/build_reflexion_skill_data.py | 1123 ---------- .../exp/embedding/build_thinking_rag_index.py | 1159 ---------- cookbook/exp/embedding/compare_math_levels.py | 91 - cookbook/exp/embedding/dataset_hard.py | 202 -- cookbook/exp/embedding/dataset_index.py | 718 ------ cookbook/exp/embedding/dataset_think.py | 456 ---- cookbook/exp/embedding/eval_dualline_math.py | 689 ------ cookbook/exp/embedding/eval_gpqa_rag.py | 1547 ------------- cookbook/exp/embedding/eval_rag_recall.py | 187 -- .../exp/embedding/eval_reflexion_skill.py | 762 ------- .../exp/embedding/make_embedding_dataset.py | 758 ------- .../exp/embedding/train_embedding_full_ddp.py | 270 --- .../exp/embedding/train_reflexion_skill.py | 1990 ----------------- .../embedding/train_reflexion_skill_replay.py | 114 - .../embedding/train_reflexion_skill_rft.py | 1568 ------------- cookbook/exp/embedding/train_skill_v2.py | 1450 ------------ .../cold_start/train_cold_start.py | 0 .../exp/{ => legacy}/condenser/dataset.py | 0 .../condenser/make_condenser_dataset.py | 0 .../condenser/train_condenser_ddp.py | 0 .../condenser/untested/eval_condensed.py | 0 .../data_pipeline/audit_rubric.py | 0 .../data_pipeline/process_and_save.py | 0 cookbook/exp/{ => legacy}/rl/grpo.py | 0 cookbook/exp/{ => legacy}/rl/rag_hint_grpo.py | 0 .../good_skill_hard_fail/analyze_3way.py | 177 ++ .../good_skill_hard_fail/eval_skill_probe.py | 284 +++ .../good_skill_hard_fail/leak_decomp.py | 146 ++ .../good_skill_hard_fail/reflexion_probe.py | 504 +++++ .../good_skill_hard_fail/sample_probe.py | 136 ++ .../skill_config_probe.py | 419 ++++ cookbook/exp/skill2lora/run_ablate12.sh | 17 + .../exp/skill2lora/skill_ablate/__init__.py | 8 + .../exp/skill2lora/skill_ablate/config.py | 147 ++ cookbook/exp/skill2lora/skill_ablate/data.py | 96 + cookbook/exp/skill2lora/skill_ablate/main.py | 155 ++ .../exp/skill2lora/skill_ablate/methods.py | 665 ++++++ cookbook/exp/skill2lora/skill_ablate/pool.py | 125 ++ .../exp/skill2lora/skill_ablate/rollouting.py | 143 ++ .../skill2lora/skill_ablate/rubric_cache.py | 91 + .../exp/skill2lora/skill_ablate/trainer.py | 344 +++ cookbook/exp/skill2lora/train_skill_v2.py | 42 +- 43 files changed, 3477 insertions(+), 13625 deletions(-) delete mode 100644 cookbook/exp/embedding/build_reflexion_coldstart_sft.py delete mode 100644 cookbook/exp/embedding/build_reflexion_skill_data.py delete mode 100644 cookbook/exp/embedding/build_thinking_rag_index.py delete mode 100644 cookbook/exp/embedding/compare_math_levels.py delete mode 100644 cookbook/exp/embedding/dataset_hard.py delete mode 100644 cookbook/exp/embedding/dataset_index.py delete mode 100644 cookbook/exp/embedding/dataset_think.py delete mode 100644 cookbook/exp/embedding/eval_dualline_math.py delete mode 100644 cookbook/exp/embedding/eval_gpqa_rag.py delete mode 100644 cookbook/exp/embedding/eval_rag_recall.py delete mode 100644 cookbook/exp/embedding/eval_reflexion_skill.py delete mode 100644 cookbook/exp/embedding/make_embedding_dataset.py delete mode 100644 cookbook/exp/embedding/train_embedding_full_ddp.py delete mode 100644 cookbook/exp/embedding/train_reflexion_skill.py delete mode 100644 cookbook/exp/embedding/train_reflexion_skill_replay.py delete mode 100644 cookbook/exp/embedding/train_reflexion_skill_rft.py delete mode 100644 cookbook/exp/embedding/train_skill_v2.py rename cookbook/exp/{ => legacy}/cold_start/train_cold_start.py (100%) rename cookbook/exp/{ => legacy}/condenser/dataset.py (100%) rename cookbook/exp/{ => legacy}/condenser/make_condenser_dataset.py (100%) rename cookbook/exp/{ => legacy}/condenser/train_condenser_ddp.py (100%) rename cookbook/exp/{ => legacy}/condenser/untested/eval_condensed.py (100%) rename cookbook/exp/{ => legacy}/data_pipeline/audit_rubric.py (100%) rename cookbook/exp/{ => legacy}/data_pipeline/process_and_save.py (100%) rename cookbook/exp/{ => legacy}/rl/grpo.py (100%) rename cookbook/exp/{ => legacy}/rl/rag_hint_grpo.py (100%) create mode 100644 cookbook/exp/skill2lora/good_skill_hard_fail/analyze_3way.py create mode 100644 cookbook/exp/skill2lora/good_skill_hard_fail/eval_skill_probe.py create mode 100644 cookbook/exp/skill2lora/good_skill_hard_fail/leak_decomp.py create mode 100644 cookbook/exp/skill2lora/good_skill_hard_fail/reflexion_probe.py create mode 100644 cookbook/exp/skill2lora/good_skill_hard_fail/sample_probe.py create mode 100644 cookbook/exp/skill2lora/good_skill_hard_fail/skill_config_probe.py create mode 100644 cookbook/exp/skill2lora/skill_ablate/__init__.py create mode 100644 cookbook/exp/skill2lora/skill_ablate/config.py create mode 100644 cookbook/exp/skill2lora/skill_ablate/data.py create mode 100644 cookbook/exp/skill2lora/skill_ablate/main.py create mode 100644 cookbook/exp/skill2lora/skill_ablate/methods.py create mode 100644 cookbook/exp/skill2lora/skill_ablate/pool.py create mode 100644 cookbook/exp/skill2lora/skill_ablate/rollouting.py create mode 100644 cookbook/exp/skill2lora/skill_ablate/rubric_cache.py create mode 100644 cookbook/exp/skill2lora/skill_ablate/trainer.py diff --git a/cookbook/exp/embedding/build_reflexion_coldstart_sft.py b/cookbook/exp/embedding/build_reflexion_coldstart_sft.py deleted file mode 100644 index 548dfad58..000000000 --- a/cookbook/exp/embedding/build_reflexion_coldstart_sft.py +++ /dev/null @@ -1,519 +0,0 @@ -"""Build a cold-start SFT corpus for reflexion skill generation on AOPS. - -Pipeline: - AOPS problems -> frozen base greedy attempt -> strategy-level rubric API diagnosis - -> answer-free API skill target -> query-only SFT examples. - -This is intentionally offline: no GRPO, no actor training, and no skill-model rollout. -The API is treated as an external teacher, so both diagnosis and generated skill targets -are filtered if they reveal the target final answer. - -Example: - LLM_BACKUP_API_KEY=... LLM_BACKUP_BASE_URL=... LLM_BACKUP_MODEL=... \ - python cookbook/exp/embedding/build_reflexion_coldstart_sft.py \ - --dataset aops --n 10000 --output-dir ./output/reflexion_coldstart_sft --overwrite -""" -import argparse -import json -import os -import sys -import time -import urllib.error -import urllib.request -from concurrent.futures import ThreadPoolExecutor -from typing import Any, Dict, List, Optional, Tuple - -_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../..')) -if _REPO_ROOT not in sys.path: - sys.path.insert(0, _REPO_ROOT) - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.sampler import vLLMSampler - -from cookbook.exp.embedding.train_reflexion_skill import ( - MODEL_ID, - GPU_MEM, - SamplingParams, - DiskCache, - _MATH_RUBRIC, - _RUBRIC_VERSION, - _answer_leaked, - _clean_text, - _empty_roll, - _format_diagnosis, - _numeric_value, - _parse_seq, - _run_samples, - _skillgen_messages, - _load_excluded_records, - build_direct_prompt, - build_skill_solve_prompt, - build_rubric_checker, - extract_boxed, - load_problems, -) - -logger = get_logger() - -COLDSTART_SYSTEM = """\ -You are writing cold-start training targets for a math skill generator. You are given a -competition problem and an answer-free process diagnosis of a previous attempt. - -Write concise, reusable guidance that a query-only solver could use before solving this -problem or similar problems. Focus on route choice, structural observations, constraints, -validity checks, and length-control habits. - -Output exactly one XML-style block: - -Your reusable guidance here. - - -Rules: -- Do not mention the diagnosis, rubric, previous attempt, or API. -- Do not reveal the final answer, a corrected value/expression, an option label, or a - step-by-step solution. -- It is okay to name methods, checks, pitfalls, and local strategy directions. -- Keep it short and useful: 3-6 compact sentences or bullets. -""" - -COLDSTART_USER = """\ -Problem: -{problem} - -Answer-free process diagnosis: -{diagnosis} - -Now write the reusable skill guidance. -""" - -_SPECIAL_TOKEN_NOTE = 'process diagnosis leaked target answer' - - -def _api_config() -> Tuple[str, str, str]: - api_key = os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY') - base_url = os.environ.get('LLM_BACKUP_BASE_URL') or os.environ.get('OPENAI_BASE_URL') or 'https://api.openai.com/v1' - model = os.environ.get('LLM_BACKUP_MODEL') or os.environ.get('OPENAI_MODEL') or 'gpt-4o-mini' - if not api_key: - raise RuntimeError('Set LLM_BACKUP_API_KEY or OPENAI_API_KEY for cold-start API generation.') - return api_key, base_url.rstrip('/'), model - - -def _chat_complete(messages: List[Dict[str, str]], max_tokens: int, temperature: float, - retries: int = 3, timeout: int = 120) -> str: - api_key, base_url, model = _api_config() - url = f'{base_url}/chat/completions' - payload = { - 'model': model, - 'messages': messages, - 'temperature': temperature, - 'max_tokens': max_tokens, - } - data = json.dumps(payload).encode('utf-8') - headers = {'Content-Type': 'application/json', 'Authorization': f'Bearer {api_key}'} - last_err = None - for attempt in range(max(1, retries)): - req = urllib.request.Request(url, data=data, headers=headers, method='POST') - try: - with urllib.request.urlopen(req, timeout=timeout) as resp: - obj = json.loads(resp.read().decode('utf-8')) - return obj['choices'][0]['message']['content'] - except (urllib.error.HTTPError, urllib.error.URLError, TimeoutError, KeyError, json.JSONDecodeError) as exc: - last_err = exc - if attempt + 1 < max(1, retries): - time.sleep(min(8.0, 1.0 * (2 ** attempt))) - continue - raise RuntimeError(f'chat completion failed after {retries} attempts: {last_err}') - - -def _extract_skill_block(text: str) -> Optional[str]: - low = (text or '').lower() - end_think = low.rfind('') - answer = text[end_think + len(''):] if end_think >= 0 else (text or '') - low = answer.lower() - s = low.rfind('') - if s < 0: - return None - inner = s + len('') - e = low.find('', inner) - if e < 0: - return None - block = answer[inner:e].strip() - return block or None - - -def _skill_response(block: str) -> str: - return f'\n{block.strip()}\n' - - -def _baseline_rollout(base_sampler, problems: List[Dict[str, Any]], base_dp: int, - args: argparse.Namespace, cache: DiskCache) -> int: - todo = [r for r in problems if DiskCache.key_for(r['problem']) not in cache] - if todo: - outs = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], - 1, args.max_tokens, base_dp, temperature=0.0) - for r, seqs in zip(todo, outs): - roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - cache.put(DiskCache.key_for(r['problem']), roll) - for r in problems: - roll = cache.get(DiskCache.key_for(r['problem'])) - r['_init'] = [roll] - r['_baseline_pass'] = 1.0 if roll.get('correct') else 0.0 - r['_failed'] = not roll.get('correct') - return len(todo) - - -def _diagnose_one(checker, r: Dict[str, Any], args: argparse.Namespace) -> str: - init = r['_init'][0] - seg_text = init.get('text', '') - if init.get('stop_reason') == 'length' or not init.get('terminated'): - seg_text += ('\n\n[Process note: this attempt was cut off at the token budget ' - 'and never produced a final \\boxed{} answer.]') - seg = {'messages': [{'role': 'user', 'content': r['problem']}, - {'role': 'assistant', 'content': seg_text}]} - attempts = max(1, args.rubric_retries + 1) - for attempt in range(attempts): - try: - return _format_diagnosis(checker.diagnose(seg, query=r['problem'])) - except Exception as exc: - if attempt + 1 < attempts: - logger.warning(f'[rubric] diagnose error: {exc}; retry {attempt + 1}/{args.rubric_retries}') - time.sleep(min(4.0, 0.5 * (2 ** attempt))) - else: - logger.warning(f'[rubric] diagnose failed: {exc}') - return '' - - -def _diagnose_batch(checker, rows: List[Dict[str, Any]], args: argparse.Namespace, - cache: DiskCache) -> int: - pending = [] - for r in rows: - init = r['_init'][0] - term = 'L' if (init.get('stop_reason') == 'length' or not init.get('terminated')) else 'T' - key = DiskCache.key_for(r['problem'], init.get('text', ''), _RUBRIC_VERSION, term) - if key in cache: - r['_rubric_diag'] = cache.get(key) - else: - pending.append((r, key)) - if not pending: - return 0 - - def run(item): - r, key = item - diag = _diagnose_one(checker, r, args) - return r, key, diag - - workers = max(1, min(args.rubric_workers, len(pending))) - fresh = 0 - with ThreadPoolExecutor(max_workers=workers) as ex: - for r, key, diag in ex.map(run, pending): - r['_rubric_diag'] = diag or '' - if diag: - cache.put(key, diag) - fresh += 1 - return fresh - - -def _target_key(problem: str, diagnosis: str, sample_idx: int) -> str: - return DiskCache.key_for('coldstart_skill_v2', str(sample_idx), problem, diagnosis) - - -def _generate_skill_targets(r: Dict[str, Any], args: argparse.Namespace, - cache: DiskCache) -> List[Dict[str, Any]]: - out = [] - messages = [ - {'role': 'system', 'content': COLDSTART_SYSTEM}, - {'role': 'user', 'content': COLDSTART_USER.format(problem=r['problem'], diagnosis=r.get('_rubric_diag', ''))}, - ] - for sample_idx in range(max(1, int(args.api_samples))): - key = _target_key(r['problem'], r.get('_rubric_diag', ''), sample_idx) - if key in cache: - resp = cache.get(key) - else: - resp = _chat_complete(messages, max_tokens=args.api_max_tokens, - temperature=args.api_temperature, retries=args.api_retries, - timeout=args.api_timeout) - cache.put(key, resp) - block = _extract_skill_block(resp) or '' - leaked = _answer_leaked(resp + '\n' + block, r['reference_answer']) - out.append({'sample_idx': sample_idx, 'raw_response': resp, - 'skills': block, 'skill_leak': leaked}) - return out - - -def _sft_messages(problem: str, response: str) -> List[Dict[str, str]]: - msgs = _skillgen_messages(problem, 'B', '') - return msgs + [{'role': 'assistant', 'content': response}] - - -def _init_base_sampler(args: argparse.Namespace): - twinkle.initialize(mode='ray', nproc_per_node=args.base_gpus, lazy_collect=False, - groups=[DeviceGroup(name='base_sampler', ranks=list(range(args.base_gpus)), device_type='GPU')]) - sampler = vLLMSampler( - model_id=MODEL_ID, - engine_args={'gpu_memory_utilization': GPU_MEM, - 'max_model_len': args.max_model_len, - 'tensor_parallel_size': 1}, - device_mesh=DeviceMesh.from_sizes(world_size=args.base_gpus, dp_size=args.base_gpus), - remote_group='base_sampler') - sampler.set_template('Template', model_id=MODEL_ID, enable_thinking=True, - max_length=args.max_model_len) - return sampler, args.base_gpus - - -def _select_records(args: argparse.Namespace) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: - load_n = 0 if args.numeric_only or args.eval_size > 0 else max(args.n, args.target_size + args.eval_size) - records = load_problems(args.dataset, load_n, args.seed) - raw_n = len(records) - if args.numeric_only: - records = [{**r, 'reference_answer': v} - for r, v in ((r, _numeric_value(r.get('reference_answer'))) for r in records) - if v is not None] - import numpy as np - np.random.RandomState(args.seed).shuffle(records) - exclude_ids, exclude_problems = _load_excluded_records(getattr(args, 'exclude_data_ids', '')) - excluded = 0 - if exclude_ids or exclude_problems: - before = len(records) - records = [r for r in records - if str(r.get('data_id', '')) not in exclude_ids - and str(r.get('problem', '')).strip() not in exclude_problems] - excluded = before - len(records) - eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 - pool = [dict(r) for r in records[eval_n:]] - pool_offset = max(0, int(getattr(args, 'pool_offset', 0) or 0)) - if pool_offset: - if pool_offset >= len(pool): - raise ValueError(f'--pool-offset ({pool_offset}) leaves no cold-start records from pool size {len(pool)}') - pool = pool[pool_offset:] - n = min(args.n, len(pool)) if args.n > 0 else min(len(pool), max(args.target_size * 2, args.target_size + 512)) - stats = {'raw_loaded': raw_n, 'numeric_dropped': raw_n - len(records) - excluded, - 'excluded_records': excluded, 'eval_size': eval_n, - 'pool_offset': pool_offset, 'pool_selected': n} - return pool[:n], stats - - -def _build_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--dataset', choices=('aops', 'math'), default='aops') - p.add_argument('--target-size', type=int, default=10000, help='Number of accepted SFT examples to write.') - p.add_argument('--n', type=int, default=0, help='Raw train-pool size after eval split; 0 auto-selects.') - p.add_argument('--pool-offset', type=int, default=0, - help='Skip this many shuffled non-eval records before building the cold-start pool.') - p.add_argument('--exclude-data-ids', default='', - help='Comma-separated jsonl files whose data_id/problem keys are excluded, ' - 'useful for building non-overlapping shards.') - p.add_argument('--eval-size', type=int, default=128) - p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) - p.add_argument('--seed', type=int, default=42) - p.add_argument('--output-dir', default='./output/reflexion_coldstart_sft') - p.add_argument('--cache-dir', default='') - p.add_argument('--overwrite', action='store_true') - p.add_argument('--no-cache', action='store_true') - p.add_argument('--chunk-size', type=int, default=64) - p.add_argument('--base-gpus', type=int, default=int(os.environ.get('BASE_GPUS', 4))) - p.add_argument('--max-model-len', type=int, default=16384) - p.add_argument('--max-tokens', type=int, default=8192) - p.add_argument('--rubric-workers', type=int, default=16) - p.add_argument('--rubric-retries', type=int, default=2) - p.add_argument('--api-workers', type=int, default=16) - p.add_argument('--api-samples', type=int, default=4, - help='API skill targets sampled per problem before executor verification.') - p.add_argument('--verify-targets', action=argparse.BooleanOptionalAction, default=True, - help='Run frozen base executor with each API skill target and keep a successful one.') - p.add_argument('--keep-unverified-targets', action='store_true', - help='If all executor checks fail, keep the first clean target anyway. Default skips it.') - p.add_argument('--api-retries', type=int, default=3) - p.add_argument('--api-timeout', type=int, default=120) - p.add_argument('--api-max-tokens', type=int, default=768) - p.add_argument('--api-temperature', type=float, default=0.2) - p.add_argument('--require-fail', action=argparse.BooleanOptionalAction, default=True, - help='Only keep API diagnoses containing [FAIL]. Use --no-require-fail to keep OK diagnoses too.') - return p.parse_args() - - -def _write(f, row: Dict[str, Any]) -> None: - f.write(json.dumps(row, ensure_ascii=False) + '\n') - - -def main() -> None: - args = _build_args() - if args.target_size <= 0: - raise ValueError('--target-size must be positive') - records, data_stats = _select_records(args) - if not records: - raise ValueError('no records selected') - - os.makedirs(args.output_dir, exist_ok=True) - sft_path = os.path.join(args.output_dir, 'coldstart_sft.jsonl') - rec_path = os.path.join(args.output_dir, 'coldstart_records.jsonl') - for path in (sft_path, rec_path): - if os.path.exists(path) and not args.overwrite: - raise FileExistsError(f'{path} exists; pass --overwrite') - - checker = build_rubric_checker() - if checker is None: - raise RuntimeError('No rubric checker available; set LLM_BACKUP_API_KEY/BASE_URL or OPENAI_API_KEY.') - _api_config() - base_sampler, base_dp = _init_base_sampler(args) - - cache_dir = args.cache_dir or os.path.join(args.output_dir, 'cache') - os.makedirs(cache_dir, exist_ok=True) - use_cache = not args.no_cache - base_cache = DiskCache(os.path.join(cache_dir, 'baseline.jsonl'), use_cache) - rubric_cache = DiskCache(os.path.join(cache_dir, 'rubric.jsonl'), use_cache) - skill_cache = DiskCache(os.path.join(cache_dir, 'api_skill.jsonl'), use_cache) - - cfg = { - 'record_type': 'config', 'mode': 'coldstart_sft_build', 'dataset': args.dataset, - 'target_size': args.target_size, 'selected_records': len(records), 'seed': args.seed, - 'pool_offset': args.pool_offset, 'exclude_data_ids': args.exclude_data_ids, - 'numeric_only': args.numeric_only, **data_stats, - 'rubric_version': _RUBRIC_VERSION, 'rubric_check': f'fixed_math_{len(_MATH_RUBRIC)}crit', - 'api_model': os.environ.get('LLM_BACKUP_MODEL') or os.environ.get('OPENAI_MODEL') or 'gpt-4o-mini', - 'api_samples': args.api_samples, 'verify_targets': args.verify_targets, - 'keep_unverified_targets': args.keep_unverified_targets, - 'require_fail': args.require_fail, 'started': int(time.time()), - } - - accepted = 0 - skipped_no_diag = skipped_no_fail = skipped_api_leak = 0 - skipped_no_skill = skipped_skill_leak = skipped_executor_fail = 0 - processed = 0 - with open(sft_path, 'w', encoding='utf-8') as sft_f, open(rec_path, 'w', encoding='utf-8') as rec_f: - _write(rec_f, cfg) - for start in range(0, len(records), args.chunk_size): - if accepted >= args.target_size: - break - chunk = [dict(r) for r in records[start:start + args.chunk_size]] - _baseline_rollout(base_sampler, chunk, base_dp, args, base_cache) - _diagnose_batch(checker, chunk, args, rubric_cache) - - def gen_one(r: Dict[str, Any]): - return r, _generate_skill_targets(r, args, skill_cache) - - candidates = [] - for r in chunk: - processed += 1 - diag = r.get('_rubric_diag', '') or '' - if not diag: - skipped_no_diag += 1 - continue - if args.require_fail and '[FAIL]' not in diag: - skipped_no_fail += 1 - continue - if _answer_leaked(diag, r['reference_answer']): - skipped_api_leak += 1 - continue - candidates.append(r) - - generated = [] - workers = max(1, min(args.api_workers, len(candidates))) - if candidates: - with ThreadPoolExecutor(max_workers=workers) as ex: - for r, targets in ex.map(gen_one, candidates): - for target in targets: - skills = target.get('skills', '') - if not skills: - skipped_no_skill += 1 - continue - if target.get('skill_leak'): - skipped_skill_leak += 1 - continue - target['r'] = r - target['response'] = _skill_response(skills) - generated.append(target) - - selected = [] - selected_keys = set() - if generated and args.verify_targets: - verify_prompts = [build_skill_solve_prompt(g['r']['problem'], g['skills']) for g in generated] - verify_outs = _run_samples(base_sampler, verify_prompts, 1, args.max_tokens, - base_dp, temperature=0.0) - attempted_keys = set() - for g, seqs in zip(generated, verify_outs): - r = g['r'] - key = r.get('data_id') or r['problem'] - attempted_keys.add(key) - roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - g['target_roll'] = roll - if key not in selected_keys and roll.get('correct') and roll.get('terminated'): - g['executor_verified'] = True - selected.append(g) - selected_keys.add(key) - if args.keep_unverified_targets: - for g in generated: - r = g['r'] - key = r.get('data_id') or r['problem'] - if key not in selected_keys: - g['executor_verified'] = False - g.setdefault('target_roll', {}) - selected.append(g) - selected_keys.add(key) - skipped_executor_fail += len(attempted_keys - selected_keys) - elif generated: - for g in generated: - r = g['r'] - key = r.get('data_id') or r['problem'] - if key not in selected_keys: - g['executor_verified'] = False - selected.append(g) - selected_keys.add(key) - - for g in selected: - if accepted >= args.target_size: - break - r = g['r'] - response = g['response'] - messages = _sft_messages(r['problem'], response) - sft_row = { - 'messages': messages, - 'user_data': {'key_rounds': [len(messages) - 1]}, - 'data_id': r.get('data_id'), - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'skills': g['skills'], 'response': response, - 'view': 'B', 'sft': True, 'source': 'api_coldstart', - 'api_sample_idx': g.get('sample_idx'), - 'executor_verified': g.get('executor_verified', False), - 'baseline_correct': r['_init'][0].get('correct'), - 'baseline_terminated': r['_init'][0].get('terminated'), - 'baseline_stop_reason': r['_init'][0].get('stop_reason'), - 'target_correct': (g.get('target_roll') or {}).get('correct'), - 'target_terminated': (g.get('target_roll') or {}).get('terminated'), - 'target_stop_reason': (g.get('target_roll') or {}).get('stop_reason'), - 'diagnosis': r.get('_rubric_diag', ''), - } - audit = { - 'record_type': 'coldstart_problem', 'accepted': True, - 'data_id': r.get('data_id'), - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'baseline': r['_init'][0], 'diagnosis': r.get('_rubric_diag', ''), - 'raw_skill_response': g.get('raw_response'), 'skills': g['skills'], - 'api_sample_idx': g.get('sample_idx'), - 'executor_verified': g.get('executor_verified', False), - 'target_roll': g.get('target_roll'), - } - _write(sft_f, sft_row) - _write(rec_f, audit) - accepted += 1 - sys.stderr.write( - f'[coldstart] processed={processed} accepted={accepted}/{args.target_size} ' - f'skip(no_diag={skipped_no_diag}, no_fail={skipped_no_fail}, api_leak={skipped_api_leak}, ' - f'no_skill={skipped_no_skill}, skill_leak={skipped_skill_leak}, ' - f'executor_fail={skipped_executor_fail})\n') - sft_f.flush(); rec_f.flush() - - summary = { - 'record_type': 'summary', 'processed': processed, 'accepted': accepted, - 'skipped_no_diag': skipped_no_diag, 'skipped_no_fail': skipped_no_fail, - 'skipped_api_leak': skipped_api_leak, 'skipped_no_skill': skipped_no_skill, - 'skipped_skill_leak': skipped_skill_leak, - 'skipped_executor_fail': skipped_executor_fail, 'finished': int(time.time()), - } - with open(rec_path, 'a', encoding='utf-8') as rec_f: - _write(rec_f, summary) - sys.stderr.write(f'[coldstart] wrote {accepted} SFT rows to {sft_path}\n') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/embedding/build_reflexion_skill_data.py b/cookbook/exp/embedding/build_reflexion_skill_data.py deleted file mode 100644 index 59fe2382c..000000000 --- a/cookbook/exp/embedding/build_reflexion_skill_data.py +++ /dev/null @@ -1,1123 +0,0 @@ -"""Offline builder for reflexion skill RFT data (self-contained, cached). - -Runs the SAME pipeline as the online trainer -- base greedy solve -> rubric -process-check (view A) -> skill-gen -> leak filter -> with-skill greedy pass -> -group-relative GRPO advantage -- but never updates the skill model. It emits -``skill_dataset.jsonl`` (trainer-schema training records), ``gen_records.jsonl`` -(full per-problem traces) and ``eval_holdout.jsonl`` (the fixed holdout). - -The expensive base rollouts and rubric diagnoses are cached to disk (one jsonl -each, keyed by an md5 of their inputs) so a re-run skips them entirely. - -8 GPUs: ranks 0-3 skill_sampler (vLLM tp1 dp4), ranks 4-7 base_sampler. Leak / -rubric use the backup teacher API: set LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / -LLM_BACKUP_MODEL. - -Launch: - LLM_BACKUP_API_KEY=... python cookbook/exp/embedding/build_reflexion_skill_data.py \ - --total-problems 3200 --base-success-frac 0.3 -""" -import argparse -import copy -import hashlib -import json -import math -import os -import re -import sys -import time -from concurrent.futures import ThreadPoolExecutor -from typing import Any, Dict, List, Optional, Tuple - -import numpy as np - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.data_format import SamplingParams -from twinkle.dataset import Dataset, DatasetMeta -from twinkle.sampler import vLLMSampler -from twinkle_agentic.verifier import LeakVerifier, RubricVerifier -from twinkle_agentic.verifier.rubric_verifier import RubricItem - -logger = get_logger() - -MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') -GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.8)) -GEN_TEMPERATURE = float(os.environ.get('GEN_TEMPERATURE', 0.6)) -GEN_TOP_P = float(os.environ.get('GEN_TOP_P', 0.95)) -AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') -MATH_DATASET_ID = os.environ.get('MATH_DATASET_ID', 'modelscope/competition_math') - - -# =========================================================================== -# Block A -- boxed extraction + answer grading -# =========================================================================== -_BOXED_RE = re.compile(r'\\boxed\s*\{') - - -def extract_boxed(text: str) -> Optional[str]: - """Last ``\\boxed{...}`` content, brace-balanced.""" - if not text: - return None - last = None - for m in _BOXED_RE.finditer(text): - depth, i = 1, m.end() - while i < len(text) and depth > 0: - depth += (text[i] == '{') - (text[i] == '}') - i += 1 - if depth == 0: - last = text[m.end():i - 1].strip() - return last - - -def normalize_answer(ans: str) -> str: - if not ans: - return '' - s = str(ans).strip() - m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) - if m: - return m.group(1) - s = s.strip('$').strip().replace('−', '-') - 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').replace(r'\tfrac', r'\frac') - s = s.replace(r'\,', '').replace(r'\;', '').replace(r'\!', '') - s = re.sub(r'\\(?:quad|qquad|\s)', '', s) - s = re.sub(r'\s+', '', s) - s = re.sub(r'\\sqrt([0-9A-Za-z])', r'\\sqrt{\1}', s) - s = re.sub(r'\\frac\{([^{}]+)\}([^{}\\])', r'\\frac{\1}{\2}', s) - s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) - s = re.sub(r'\{[a-zA-Z]+\}$', '', s) - s = re.sub(r'\^\\circ|\^\{\\circ\}|°|\\circ', 'deg', s) - s = s.replace(r'\minus{}', '-').replace(r'\minus', '-') - - def _frac_to_slash(mt): - text = mt.group(0) - pos = text.index('{') + 1 - depth, num_start = 1, pos - while depth > 0: - depth += (text[pos] == '{') - (text[pos] == '}') - pos += 1 - numer = text[num_start:pos - 1] - pos += 1 - den_start, depth = pos, 1 - while depth > 0: - depth += (text[pos] == '{') - (text[pos] == '}') - pos += 1 - return f'({numer})/({text[den_start:pos - 1]})' - - s = re.sub(r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', _frac_to_slash, s) - s = re.sub(r'(? bool: - try: - va, vb = float(a.replace('(', '').replace(')', '')), 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) - return va is not None and vb is not None and abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) - - -_MCQ_REF_RE = re.compile( - r'^\\(?:textbf|text|mathrm|mathbf)\{\(?([A-E])\)?\s*\}\s*(.+)$' - r'|^\(?([A-E])\)\s+(.+)$') -_VAR_PREFIX_RE = re.compile(r'^(?:[A-Za-z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)$', re.DOTALL) - - -def _split_mcq(ans: str): - s = ans.strip() - m = _MCQ_REF_RE.match(s) - if m: - return (m.group(1) or m.group(3)), ((m.group(2) or m.group(4) or '').strip() or None) - bl = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) - return (bl.group(1), None) if bl else (None, s or None) - - -def _strip_var_prefix(ans: str) -> str: - m = _VAR_PREFIX_RE.match((ans or '').strip()) - return m.group(1).strip() if m else (ans or '') - - -def _normalize_tuple(ans: str) -> str: - return re.sub(r'[\s()\[\]{}\\]', '', ans or '') - - -def _math_verify_equal(predicted: str, reference: str) -> bool: - try: - from math_verify import parse, verify - gold = parse(r'\boxed{' + reference + '}', parsing_timeout=1) - pred = parse(r'\boxed{' + predicted + '}', parsing_timeout=1) - return bool(gold and pred and verify(gold, pred, timeout_seconds=1)) - except Exception: - return False - - -def answers_match(predicted: str, reference: str) -> bool: - if not predicted or not reference: - return False - norm_p, norm_r = normalize_answer(predicted), normalize_answer(reference) - if norm_p == norm_r or norm_p.lower() == norm_r.lower() or _try_numeric_equal(norm_p, norm_r): - return True - stripped_p = normalize_answer(_strip_var_prefix(predicted)) - stripped_r = normalize_answer(_strip_var_prefix(reference)) - if stripped_p and stripped_r: - if (stripped_p == stripped_r or stripped_p.lower() == stripped_r.lower() - or _try_numeric_equal(stripped_p, stripped_r)): - return True - p_letter, p_value = _split_mcq(predicted) - r_letter, r_value = _split_mcq(reference) - if p_letter and r_letter and p_letter == r_letter: - return True - if (p_letter and p_value is None) and (r_letter and r_value is None): - return False - 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 and (p_val == r_val or p_val.lower() == r_val.lower() or _try_numeric_equal(p_val, r_val)): - return True - for left, right in ((norm_p, norm_r), (stripped_p, stripped_r)): - tl, tr = _normalize_tuple(left), _normalize_tuple(right) - if ',' in tl and tl == tr: - return True - if '=' in norm_r and '=' not in norm_p: - if any(part and (part == norm_p or _try_numeric_equal(part, norm_p)) for part in norm_r.split('=')): - return True - if '=' in norm_p and '=' not in norm_r: - if any(part and (part == norm_r or _try_numeric_equal(part, norm_r)) for part in norm_p.split('=')): - return True - return _math_verify_equal(stripped_p or norm_p, stripped_r or norm_r) - - -# =========================================================================== -# Block B -- prompts, skill parsing, batched sampling -# =========================================================================== -DIRECT_SYSTEM = ( - 'You are an expert competition mathematician. Solve the following problem ' - 'step by step. Provide your final answer inside \\boxed{}.') - -# Appended to solve turns: box BOTH the letter and value of an MCQ so the model -# never loops deciding which form to box. -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.') - -_SKILL_SOLVE_PREFIX = ( - DIRECT_SYSTEM + '\n\n' - 'Before you start, keep these reminders in mind to avoid common mistakes on this ' - 'type of problem:\n') -_SKILL_SOLVE_SUFFIX = '\nApply them where relevant, but rely on your own reasoning to reach the answer.' - - -def build_direct_prompt(problem: str) -> Dict[str, Any]: - return {'messages': [{'role': 'system', 'content': DIRECT_SYSTEM}, - {'role': 'user', 'content': problem}]} - - -def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: - # Concatenation (not .format): DIRECT_SYSTEM/skill contain literal braces. - return {'messages': [ - {'role': 'system', 'content': _SKILL_SOLVE_PREFIX + skill + _SKILL_SOLVE_SUFFIX}, - {'role': 'user', 'content': problem + MCQ_INSTRUCTION}]} - - -# -- skill-gen prompts (view A: problem + rubric findings; view B: query only) -- -SKILL_GEN_SYSTEM = ( - 'You are a mathematics coach. You are shown a competition problem together with an ' - 'automated process-check of an earlier solver attempt at it -- which solution ' - 'criteria the attempt passed or failed, and suggested fixes for the failures. You ' - 'do NOT see the attempt itself, only this check. Treat the check as privileged ' - 'training scaffolding: study it together with the problem, identify the ' - 'problem-visible features that make each useful flagged failure relevant, then ' - 'rephrase those lessons as self-contained reusable skills. The goal is not to ' - 'continue from the check, cite it, or hide it silently; the goal is to turn it into ' - 'a problem-triggered reasoning pattern a query-only solver could reproduce later.\n\n' - 'Good skills name the observable trigger, the method worth reaching for, the ' - 'pitfall to watch, and a quick verification habit. Prefer formulations like ' - '"When a configuration has ...", "Before setting up ...", or "Check whether ..." ' - 'over references to the process-check, failed criteria, or the earlier attempt. ' - "These tips are advisory: they will be placed in a solver's system prompt as gentle " - 'reminders before it works through a SIMILAR problem on its own, without seeing ' - 'this process-check. So keep them general and transferable rather than a worked ' - 'solution to this exact problem, and do not state its specific intermediate values ' - 'or final answer. Think briefly first, then give your tips as a markdown bullet ' - 'list wrapped in and , like the example below.') - -SKILL_GEN_SYSTEM_Q = ( - 'You are a mathematics coach. You are shown ONE competition problem and nothing ' - 'else — no solution and no attempt. Think about what approach this KIND of problem ' - 'calls for and where solvers tend to slip, then distil a few reusable tips.\n\n' - "These tips are advisory: they will be placed in a solver's system prompt as gentle " - 'reminders before it works through a SIMILAR problem on its own. So keep them ' - 'general and transferable — the method worth reaching for, the pitfall to watch and ' - 'a quick check, and the discipline to settle on a final answer — rather than a ' - 'worked solution to this exact problem, and without stating its specific ' - 'intermediate values or its final answer. Think briefly first, then give your tips ' - 'as a markdown bullet list wrapped in and , like the example below.') - -SKILL_GEN_USER_Q = ( - 'Problem:\n{problem}\n\n' - 'Now reason about this TYPE of problem, then output the skills bullet list.') - -SKILL_GEN_USER_RUBRIC = ( - 'Problem:\n{problem}\n\n' - 'Process check of an earlier attempt (automated rubric verifier -- treat as ' - 'evidence, not gospel; PASS/FAIL per criterion with suggested fixes for failures):\n' - '{diagnosis}\n\n' - 'Now output a self-contained skills bullet list. Each bullet should still be useful ' - 'if the process check were removed: connect any useful flagged failure to ' - 'problem-visible features, general methods, and quick checks rather than citing the ' - 'rubric or the earlier attempt.') - -_EX_PROBLEM = 'Simplify $\\sqrt{72} + \\sqrt{18}$ and give the result.' -_EX_SKILLS = ( - '\n' - '- Rewrite each square root by factoring its radicand into a perfect square times ' - 'a remainder, then move the perfect-square factor outside.\n' - '- Avoid the classic trap $\\sqrt{a}+\\sqrt{b}\\ne\\sqrt{a+b}$; only combine terms ' - 'sharing the same simplest radical, and sanity-check by estimating each root.\n' - '- Procedure: simplify every radical, group like radical terms, add their ' - 'coefficients, then reduce to simplest form.\n' - '- Once the expression is in simplest form, commit to that single result as the ' - 'final answer rather than re-checking indefinitely.\n' - '') - - -def _skillgen_messages(problem: str, view: str, diagnosis: str) -> List[Dict[str, Any]]: - """Single source of truth for the skill-gen prompt. View A with a localisable - failure uses problem + rubric findings; view B -- or a view-A problem whose rubric - flagged nothing (no ``[FAIL]``) -- degrades to the query-only prompt.""" - if view == 'B' or '[FAIL]' not in (diagnosis or ''): - return [{'role': 'system', 'content': SKILL_GEN_SYSTEM_Q}, - {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=_EX_PROBLEM)}, - {'role': 'assistant', 'content': _EX_SKILLS}, - {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=problem)}] - return [{'role': 'system', 'content': SKILL_GEN_SYSTEM}, - {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=_EX_PROBLEM)}, - {'role': 'assistant', 'content': _EX_SKILLS}, - {'role': 'user', 'content': SKILL_GEN_USER_RUBRIC.format(problem=problem, diagnosis=diagnosis)}] - - -def _assign_view(problem: str, args: argparse.Namespace) -> str: - h = int(hashlib.md5(f'{args.seed}:{problem}'.encode('utf-8')).hexdigest(), 16) - return 'B' if (h % 100000) / 100000.0 < args.view_b_frac else 'A' - - -def _view_prompt(r: Dict[str, Any]) -> Dict[str, Any]: - return {'messages': _skillgen_messages(r['problem'], r['_view'], r.get('_rubric_diag', ''))} - - -_BULLET_RE = re.compile(r'(?m)^\s*(?:[-*]|\d+[.)])\s') -_META_RE = re.compile( - r'\b(the student|the solver|the attempt|this attempt|the trace|the response|' - r'in the (?:attempt|trace|response|solution)|as (?:shown|seen|noted) above|' - r'the (?:above|previous|earlier)|my (?:reasoning|analysis)|i (?:think|need|will))\b', - re.IGNORECASE) -_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') - - -def _clean_text(decoded: Optional[str]) -> str: - return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() - - -def _is_clean_block(block: str) -> bool: - """Pure bullet list (every non-empty line a bullet) with no meta/trajectory ref.""" - lines = [ln for ln in (l.strip() for l in block.splitlines()) if ln] - if not lines or not all(_BULLET_RE.match(ln) for ln in lines): - return False - return _META_RE.search(block) is None - - -def _extract_skills_block(text: str) -> Optional[str]: - """Clean ``...`` block, or None. Requires ```` (skill-gen - runs thinking ON); reads only the answer after the last one, so a mid-reasoning draft - or a demo echo can never be mistaken for the answer.""" - low = text.lower() - end_think = low.rfind('') - if end_think < 0: - return None - answer = text[end_think + len(''):] - low_a = answer.lower() - s = low_a.find('') - if s < 0: - return None - inner = s + len('') - e = low_a.find('', inner) - block = (answer[inner:e] if e >= 0 else answer[inner:]).strip() - block = re.sub(r'', '', block, flags=re.IGNORECASE).strip() - return block if _is_clean_block(block) else None - - -def _parse_seq(seq, gold: str) -> Dict[str, Any]: - """Grade one sampled sequence into a rollout record.""" - text = _clean_text(getattr(seq, 'decoded', '') or '') - pred = extract_boxed(text) - correct = bool(pred) and answers_match(pred, gold) - terminated = getattr(seq, 'stop_reason', None) != 'length' - return {'pred': pred, 'correct': correct, 'terminated': terminated, - 'passed': bool(correct and terminated), - 'stop_reason': getattr(seq, 'stop_reason', None), - 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} - - -def _run_samples(sampler, prompts: List[Any], num_samples: int, max_tokens: int, - gen_dp: int, temperature: Optional[float] = None, - top_p: Optional[float] = None, top_k: Optional[int] = None) -> List[List[Any]]: - """One batched sampler call -> per-prompt list of raw sequences. vLLM dp needs - batch len >= dp, so pad the tail and slice back.""" - if not prompts: - return [] - params = SamplingParams( - max_tokens=max_tokens, - temperature=GEN_TEMPERATURE if temperature is None else temperature, - top_p=GEN_TOP_P if top_p is None else top_p, - num_samples=num_samples, **({} if top_k is None else {'top_k': top_k})) - padded = prompts - if gen_dp > 1 and 0 < len(prompts) < gen_dp: - padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] - responses = sampler.sample(padded, params)[:len(prompts)] - return [list(r.sequences) if (r and r.sequences) else [] for r in responses] - - -# =========================================================================== -# Block C -- data loading via twinkle.Dataset + numeric filtering -# =========================================================================== -def load_problems(dataset: str, n: int, seed: int) -> List[Dict[str, Any]]: - """Load boxed-answer problems as ``{problem, reference_answer, level?}`` via - twinkle.Dataset (ModelScope hub), sampled to ``n`` (0 = all).""" - ds_id = AOPS_DATASET_ID if dataset == 'aops' else MATH_DATASET_ID - rows = Dataset(DatasetMeta(dataset_id=f'ms://{ds_id}', split='train')).dataset - out: List[Dict[str, Any]] = [] - for row in rows: - if dataset == 'aops' and not (row.get('metadata') or {}).get('boxed'): - continue - ref = extract_boxed(row.get('solution', '')) - if not ref: - continue - rec = {'problem': row['problem'], 'reference_answer': ref} - if row.get('level'): - rec['level'] = row['level'] - out.append(rec) - logger.info(f'[data] {dataset}: {len(out)} boxed problems') - rng = np.random.RandomState(seed) - rng.shuffle(out) - return out[:n] if (n and n < len(out)) else out - - -_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') - - -def _norm_num_text(num: str) -> str: - try: - f = float(num) - return str(int(f)) if f == int(f) else str(f) - except Exception: - return str(num).strip() - - -def _numeric_value(raw: Any) -> Optional[str]: - """Collapse an answer to a single int/decimal/fraction, or None.""" - if raw is None: - return None - s = str(raw).strip().strip('$').strip() - s = s.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac') - s = re.sub(r'\\!|\\,|\\;|\\ |\\left|\\right|\s', '', s) - for pat in (r'\\frac\{(-?\d+)\}\{(-?\d+)\}', r'(-?\d+)/(-?\d+)'): - m = re.fullmatch(pat, s) - if m: - a, b = int(m.group(1)), int(m.group(2)) - return _norm_num_text(str(a / b)) if b else None - return _norm_num_text(s) if _NUM_RE.fullmatch(s) else None - - -def _load_records(args: argparse.Namespace) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], Dict[str, int]]: - """Load, numeric-filter, shuffle, then split a fixed eval holdout off the front.""" - # Load all when filtering or splitting (else the eval holdout could starve train). - load_n = 0 if (args.numeric_only or args.eval_size > 0) else args.n - records = load_problems(args.dataset, load_n, args.seed) - raw_n, dropped = len(records), 0 - if args.numeric_only: - kept = [] - for r in records: - ref = _numeric_value(r.get('reference_answer')) - if ref is None: - dropped += 1 - continue - kept.append({**r, 'reference_answer': ref}) - records = kept - np.random.RandomState(args.seed).shuffle(records) - eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 - eval_records = [dict(r) for r in records[:eval_n]] - pool = records[eval_n:] - train_n = args.n if args.n > 0 else len(pool) - train_records = [dict(r) for r in pool[:train_n]] - overlap = {r['problem'] for r in train_records} & {r['problem'] for r in eval_records} - if overlap: - raise ValueError(f'eval/train overlap: {len(overlap)} problems') - stats = {'raw_loaded': raw_n, 'numeric_dropped': dropped, - 'train_records': len(train_records), 'eval_records': len(eval_records)} - return train_records, eval_records, stats - - -# =========================================================================== -# Block D -- disk cache, problem pool, baseline rollout, rubric check -# =========================================================================== -class DiskCache: - """Append-only jsonl kv cache (md5 key -> value). Loads on init, appends on put. - Disabled instances (``enabled=False``) always miss and never write.""" - - def __init__(self, path: str, enabled: bool = True): - self.path, self.enabled = path, enabled - self._mem: Dict[str, Any] = {} - self._fh = None - if not enabled: - return - if os.path.exists(path): - with open(path, encoding='utf-8') as f: - for line in f: - if line.strip(): - row = json.loads(line) - self._mem[row['key']] = row['value'] - self._fh = open(path, 'a', encoding='utf-8') - - @staticmethod - def key_for(*parts: str) -> str: - return hashlib.md5('\x1f'.join(parts).encode('utf-8')).hexdigest() - - def __contains__(self, key: str) -> bool: - return key in self._mem - - def get(self, key: str) -> Any: - return self._mem.get(key) - - def put(self, key: str, value: Any) -> None: - self._mem[key] = value - if self._fh is not None: - self._fh.write(json.dumps({'key': key, 'value': value}, ensure_ascii=False) + '\n') - self._fh.flush() - - def close(self) -> None: - if self._fh is not None: - self._fh.close() - - -class ProblemPool: - """Cyclic draw source. Each full pass reshuffles with ``seed + epoch``; the initial - pass keeps the loader order. ``draw(k)`` returns k distinct problems (k << pool).""" - - def __init__(self, records: List[Dict[str, Any]], seed: int): - self._records = list(records) - self._seed, self._cursor, self.epoch = seed, 0, 0 - - def draw(self, k: int) -> List[Dict[str, Any]]: - out, seen = [], set() - while len(out) < k: - if self._cursor >= len(self._records): - self.epoch += 1 - np.random.RandomState(self._seed + self.epoch).shuffle(self._records) - self._cursor = 0 - r = self._records[self._cursor] - self._cursor += 1 - if id(r) not in seen: - seen.add(id(r)) - out.append(r) - return out - - -def _empty_roll() -> Dict[str, Any]: - return {'pred': '', 'correct': False, 'terminated': False, 'passed': False, - 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} - - -def _apply_baseline(r: Dict[str, Any], roll: Dict[str, Any]) -> None: - """Attach a greedy baseline roll and reset per-chunk working state.""" - r['_baseline_rolls'], r['_cands'], r['_init'] = [roll], [], [roll] - r['_failed'] = not roll['correct'] - r['_baseline_pass'] = 1.0 if roll['correct'] else 0.0 - r['_hard'] = True # process every problem; group variance selects (SEAM-style) - - -def baseline_rollout(base_sampler, problems: List[Dict[str, Any]], base_dp: int, - args: argparse.Namespace, cache: DiskCache) -> int: - """Phase 1: base solves each problem greedily once (T=0, M=1), disk-cached by - problem text. The base is frozen + greedy so the cache is exact. Returns the number - of fresh (cache-miss) rollouts.""" - todo = [r for r in problems if DiskCache.key_for(r['problem']) not in cache] - if todo: - out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], - 1, args.max_tokens, base_dp, temperature=0.0) - for r, seqs in zip(todo, out): - roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - cache.put(DiskCache.key_for(r['problem']), roll) - for r in problems: - _apply_baseline(r, cache.get(DiskCache.key_for(r['problem']))) - return len(todo) - - -# -- rubric process-check (view A): teacher diagnoses the base's attempt -- -_RFT_DIAG_SYSTEM = """\ -You are a process error checker for a math solution attempt. You are given a math -problem, a rubric, and one attempted solution segment. Decide PASS or FAIL for each -criterion and explain only the process error type. - -Output STRICT JSON (no prose outside it) with this shape: -{ - "items": [ - {"index": 1, "verdict": "PASS", "reason": "", - "fix": ""}, - {"index": 2, "verdict": "FAIL", "reason": "", - "fix": ""} - ], - "overall": "OK" | "ISSUES", - "summary": "" -} - -Rules: -- Judge every criterion independently and literally; a [Hard Rule] is FAIL unless - unambiguously satisfied. -- Judge ONLY what is observable in THIS segment. -- Content inside ... (or ) is internal reasoning, not - user-facing output; ignore it for "output only X" style criteria. -- For PASS items, leave "fix" as "". -- For FAIL items, "reason", "fix", and "summary" must describe only the flawed - step, theorem, arithmetic operation, case split, or verification habit. -- NEVER state the correct final answer, corrected final expression, option letter, - graph/choice label, or any exact value that the answer should become. -- NEVER write phrases like "the correct answer is", "which gives", "yielding", - "should be ", "Option ", or "Graph ". -- If a fix would require naming a corrected value, replace it with a method-level - instruction such as "redo that computation carefully" or "apply the theorem with - the correct quantities". -- Keep every "reason" and "fix" clear and concise — one short sentence each. -- "overall" is "OK" only if NO criterion is FAIL. -- Output only the JSON object.""" - -_RFT_DIAG_USER = """\ -## Task / query (context) -{query} - -## Rubric -{rubric} - -## Segment -{segment} - -Now output the diagnostic JSON object.""" - -_MATH_RUBRIC = [ - ('The reasoning contains no arithmetic or algebraic error', True), - ('Each step follows logically from the previous ones', True), - ('No formula or theorem is misstated or misapplied', True), - ('The approach is on track to answer the actual question asked', False), - ('No step contradicts an earlier established fact', False), -] - - -class _RftRubricVerifier(RubricVerifier): - def _diagnose_trajectory(self, query: str, rubric_block: str, segment_text: str) -> dict: - return {'messages': [ - {'role': 'system', 'content': _RFT_DIAG_SYSTEM}, - {'role': 'user', 'content': _RFT_DIAG_USER.format( - query=query, rubric=rubric_block, segment=segment_text)}]} - - -def build_rubric_checker() -> Optional[RubricVerifier]: - """Fixed math-process rubric verifier (teacher-served). None if no LLM backup env.""" - if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') - or os.environ.get('OPENAI_API_KEY')): - return None - return _RftRubricVerifier( - fixed_rubric=[RubricItem(t, is_hard=h) for t, h in _MATH_RUBRIC], gate=True) - - -def _format_diagnosis(detail) -> str: - """One line per criterion (PASS/FAIL + reason + fix) then a summary.""" - rub = detail.rubric - lines = [] - for it in detail.items: - text = rub[it.index - 1].text if 0 < it.index <= len(rub) else f'criterion {it.index}' - if it.verdict: - lines.append(f'- [PASS] {text}') - else: - tail = f': {it.reason}' if it.reason else '' - tail += f' (fix: {it.fix})' if it.fix else '' - lines.append(f'- [FAIL] {text}{tail}') - if detail.summary: - lines.append(f'Summary: {detail.summary}') - return '\n'.join(lines) - - -def diagnose_views(checker, hard: List[Dict[str, Any]], args: argparse.Namespace, - cache: DiskCache) -> None: - """Rubric-check every view-A problem's greedy attempt in parallel (disk-cached by - problem + attempt), stashing the formatted findings on ``r['_rubric_diag']``.""" - targets = [r for r in hard if r.get('_view') == 'A'] - if not checker or not targets: - return - - def _key(r: Dict[str, Any]) -> str: - return DiskCache.key_for(r['problem'], r.get('_init', [{}])[0].get('text', '')) - - pending = [] - for r in targets: - key = _key(r) - if key in cache: - r['_rubric_diag'] = cache.get(key) - else: - pending.append((r, key)) - if not pending: - return - - def _run(item): - r, key = item - seg = {'messages': [{'role': 'user', 'content': r['problem']}, - {'role': 'assistant', 'content': r['_init'][0]['text']}]} - try: - return r, key, _format_diagnosis(checker.diagnose(seg, query=r['problem'])) - except Exception as exc: # teacher hiccup -> no-diagnosis prompt (not cached) - logger.warning(f'[rubric] diagnose error: {exc}') - return r, key, None - - workers = max(1, min(args.rubric_workers, len(pending))) - with ThreadPoolExecutor(max_workers=workers) as ex: - for r, key, diag in ex.map(_run, pending): - r['_rubric_diag'] = diag or '' - if diag is not None: - cache.put(key, diag) - - -# =========================================================================== -# Block E -- chunk draw, pipeline, record building -# =========================================================================== -def _baseline_class(r: Dict[str, Any]) -> str: - """success | fail_loop (out of length / never terminated) | fail_wrong.""" - roll = r['_init'][0] - if roll['correct']: - return 'success' - return 'fail_loop' if (roll['stop_reason'] == 'length' or not roll['terminated']) else 'fail_wrong' - - -def _select_balanced(buckets: Dict[str, List[Dict[str, Any]]], n_success: int, - n_fail: int, n_fail_loop: int) -> List[Dict[str, Any]]: - """Pick n_fail base-fails (toward n_fail_loop loop-fails, best-effort) + n_success - base-successes; top up any shortfall from leftovers.""" - loop, wrong, succ = buckets['fail_loop'], buckets['fail_wrong'], buckets['success'] - take_loop = min(n_fail_loop, len(loop)) - take_wrong = min(n_fail - take_loop, len(wrong)) - take_loop = min(n_fail - take_wrong, len(loop)) - sel = loop[:take_loop] + wrong[:take_wrong] + succ[:n_success] - target = n_success + n_fail - if len(sel) < target: - used = {id(x) for x in sel} - sel += [x for b in (loop, wrong, succ) for x in b if id(x) not in used][:target - len(sel)] - return sel - - -def draw_chunk(pool: ProblemPool, base_sampler, base_dp: int, args: argparse.Namespace, - cache: DiskCache) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: - """Draw one chunk, running baseline rollout (Phase 1) on every drawn problem. With - ``--balance``, keep drawing+baselining until the target base fail:success mix is - reachable (or the budget is hit), then select a balanced subset.""" - if not args.balance: - chunk = pool.draw(args.chunk_size) - fresh = baseline_rollout(base_sampler, chunk, base_dp, args, cache) - return chunk, {'enabled': False, 'n_drawn': len(chunk), 'n_baseline_fresh': fresh} - - n_success = max(0, min(args.chunk_size, round(args.chunk_size * args.balance_success_frac))) - n_fail = args.chunk_size - n_success - n_fail_loop = round(n_fail * args.balance_loop_frac) - buckets: Dict[str, List[Dict[str, Any]]] = {'success': [], 'fail_loop': [], 'fail_wrong': []} - budget, n_drawn, n_fresh, seen = args.chunk_size * args.balance_max_draws_mult, 0, 0, set() - while n_drawn < budget: - if (len(buckets['success']) >= n_success - and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail): - break - batch = pool.draw(args.chunk_size) - n_fresh += baseline_rollout(base_sampler, batch, base_dp, args, cache) - n_drawn += len(batch) - for r in batch: - if id(r) not in seen: - seen.add(id(r)) - buckets[_baseline_class(r)].append(r) - - reached = (len(buckets['success']) >= n_success - and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail) - chunk = _select_balanced(buckets, n_success, n_fail, n_fail_loop) - sel_success = sum(1 for r in chunk if not r['_failed']) - stats = { - 'enabled': True, 'n_drawn': n_drawn, 'n_baseline_fresh': n_fresh, 'n_selected': len(chunk), - 'target_success': n_success, 'target_fail': n_fail, 'target_fail_loop': n_fail_loop, - 'selected_success': sel_success, 'selected_fail': len(chunk) - sel_success, - 'selected_fail_loop': sum(1 for r in chunk if _baseline_class(r) == 'fail_loop'), - 'selected_fail_wrong': sum(1 for r in chunk if _baseline_class(r) == 'fail_wrong'), - 'selected_success_frac': (sel_success / len(chunk)) if chunk else 0.0, - 'budget_hit': not reached, - } - return chunk, stats - - -def _assign_advantages(hard: List[Dict[str, Any]], args: argparse.Namespace) -> None: - """Group-relative advantage over each problem's scored candidates using the greedy - binary reward R in {0,1}: ``A = (R - mean) / (std + eps)``. std==0 groups (all solve - / all fail) get advantage 0 and no gradient -- GRPO's variance selects informative - problems, so no explicit difficulty gate is needed.""" - eps = 1e-6 - for r in hard: - for c in r['_cands']: - c['advantage'], c['grpo_adv'], c['kept'] = 0.0, 0.0, False - cs = ([c for c in r['_cands'] if c.get('reward') is not None] if args.format_in_reward - else [c for c in r['_cands'] if c['leaked'] is False and c.get('reward') is not None]) - if len(cs) < 2: - continue - rewards = [c['reward'] for c in cs] - mean_r = sum(rewards) / len(rewards) - std = (sum((x - mean_r) ** 2 for x in rewards) / len(rewards)) ** 0.5 - if std < 1e-9: - continue - for c in cs: - adv = (c['reward'] - mean_r) / (std + eps) - c['advantage'], c['grpo_adv'], c['kept'] = adv, adv, c['reward'] > mean_r - - -def process_chunk(base_sampler, skill_sampler, leak: LeakVerifier, chunk: List[Dict[str, Any]], - ci: int, base_dp: int, skill_dp: int, args: argparse.Namespace, - checker, rubric_cache: DiskCache - ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], List[Dict[str, Any]]]: - """view assign -> rubric-check (view A) -> skill-gen -> leak-filter -> with-skill - greedy pass -> GRPO advantages. ``chunk`` arrives already baselined by draw_chunk.""" - hard = chunk - - # Phase 2: view routing + view-A rubric check (view B is query-only, no rubric). - for r in hard: - r['_view'], r['_rubric_diag'] = _assign_view(r['problem'], args), '' - diagnose_views(checker, hard, args, rubric_cache) - - # Phase 3: skill-gen (thinking ON), per-view prompt; re-sample empties. - flat: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] - pending = list(hard) - for _ in range(args.skill_retries + 1): - if not pending: - break - sg_out = _run_samples(skill_sampler, [_view_prompt(r) for r in pending], - args.n_skills, args.skill_max_tokens, skill_dp, - temperature=args.skill_gen_temperature, - top_p=args.skill_gen_top_p, top_k=args.skill_gen_top_k) - still = [] - for r, seqs in zip(pending, sg_out): - got = False - for s in seqs: - resp = _clean_text(getattr(s, 'decoded', '') or '') - block = _extract_skills_block(resp) - cand = {'skills': block or '', 'response': resp, 'parseable': bool(block), - 'view': r['_view'], 'leaked': None, 'leak_reason': '', - 'leak_source': '', 'with_pass': None, 'reward': None, 'rolls': []} - r['_cands'].append(cand) - if block: - flat.append((r, cand)) - got = True - if not got: - still.append(r) - pending = still - - # Phase 4: leak filter (view A only; view B is query-only -> treated clean, SEAM-like). - for r, c in flat: - if r.get('_view') != 'A': - c['leaked'], c['leak_reason'], c['leak_source'] = False, '', 'skipped_viewB' - flat_a = [(r, c) for r, c in flat if r.get('_view') == 'A'] - if flat_a: - details = leak.leak_batch( - [{'content': c['skills'], 'query': r['problem'], 'reference': r['reference_answer']} - for r, c in flat_a], max_workers=args.leak_workers) - for (r, c), d in zip(flat_a, details): - c['leaked'], c['leak_reason'], c['leak_source'] = bool(d.leaked), d.reason, d.source - - # Phase 5: with-skill greedy pass (T=0, M=1) on clean candidates. Reward = correct, - # absolute (no baseline subtraction); the group mean in Phase 6 is the only baseline. - clean = [(r, c) for r, c in flat if c['leaked'] is False] - if clean: - ws_out = _run_samples(base_sampler, - [build_skill_solve_prompt(r['problem'], c['skills']) for r, c in clean], - 1, args.max_tokens, base_dp, temperature=0.0) - for (r, c), seqs in zip(clean, ws_out): - c['rolls'] = [_parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll()] - c['with_pass'] = 1.0 if c['rolls'][0]['correct'] else 0.0 - c['reward'] = c['with_pass'] - if args.format_in_reward: # unparseable/leaked candidates score 0 and still join the group - for r in hard: - for c in r['_cands']: - if c['reward'] is None: - c['reward'] = 0.0 - - # Phase 6: group-relative GRPO advantage. - _assign_advantages(hard, args) - return ([_full_record(r, ci) for r in chunk], _chunk_summary(chunk, ci, args), - _group_records(chunk, args)) - - -def _roll(x: Dict[str, Any]) -> Dict[str, Any]: - return {k: x[k] for k in ('pred', 'correct', 'terminated', 'passed', - 'stop_reason', 'gen_tokens', 'text')} - - -def _is_trainable(c: Dict[str, Any], args: argparse.Namespace) -> bool: - """A candidate reaches the GRPO update iff its advantage is non-zero (and, without - --format-in-reward, is also clean and scored).""" - adv_nz = abs(c.get('advantage') or 0.0) > 1e-9 - if args.format_in_reward: - return adv_nz - return c['leaked'] is False and c.get('with_pass') is not None and adv_nz - - -def _full_record(r: Dict[str, Any], ci: int) -> Dict[str, Any]: - """Complete per-problem trace: init attempt, baseline, and all candidates.""" - init = r['_init'][0] - return { - 'record_type': 'problem', 'chunk': ci, 'problem': r['problem'], - 'reference_answer': r['reference_answer'], 'level': r.get('level', ''), - 'failed_first_try': r['_failed'], - 'init_attempt': {'text': init['text'], 'pred': init['pred'], 'correct': init['correct'], - 'terminated': init['terminated'], 'stop_reason': init['stop_reason'], - 'gen_tokens': init['gen_tokens']}, - 'baseline_pass': r['_baseline_pass'], 'is_hard': r['_hard'], - 'view': r.get('_view', ''), 'rubric_diag': r.get('_rubric_diag', ''), - 'baseline_rolls': [_roll(x) for x in r['_baseline_rolls']], - 'candidates': [{ - 'skills': c['skills'], 'response': c['response'], 'parseable': c['parseable'], - 'leaked': c['leaked'], 'leak_reason': c['leak_reason'], 'leak_source': c['leak_source'], - 'with_pass': c['with_pass'], 'reward': c.get('reward'), 'advantage': c.get('advantage'), - 'grpo_adv': c.get('grpo_adv'), 'kept': c.get('kept'), - 'rolls': [_roll(x) for x in c['rolls']], - } for c in r['_cands']], - } - - -def _view_stats(hard: List[Dict[str, Any]], view: str) -> Dict[str, Any]: - hv = [r for r in hard if r.get('_view') == view] - cands = [c for r in hv for c in r['_cands'] if c['parseable']] - clean = [c for c in cands if c['leaked'] is False] - adopted = sum(1 for r in hv - if any(c['leaked'] is False and abs(c.get('advantage') or 0.0) > 1e-9 - for c in r['_cands'])) - return {'n_hard': len(hv), 'n_candidates_parseable': len(cands), 'n_clean': len(clean), - 'n_adopted_problems': adopted, 'adoption_rate': (adopted / len(hv)) if hv else 0.0} - - -def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespace) -> Dict[str, Any]: - hard = [r for r in chunk if r['_hard']] - all_cands = [c for r in chunk for c in r['_cands']] - cands = [c for c in all_cands if c['parseable']] - scored = [c for c in cands if c['with_pass'] is not None] - ws_rolls = [x for c in scored for x in c['rolls']] - train_cands = [c for c in all_cands if _is_trainable(c, args)] - fail_cands = [c for r in chunk if r['_failed'] for c in r['_cands']] - base_acc = (sum(r['_baseline_pass'] for r in hard) / len(hard)) if hard else 0.0 - ws_acc = (sum(c['with_pass'] for c in scored) / len(scored)) if scored else 0.0 - abs_adv = lambda cs: sum(abs(c.get('advantage') or 0.0) for c in cs) - total_abs = abs_adv(all_cands) - return { - 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), - 'n_failed_first_try': sum(1 for r in chunk if r['_failed']), 'n_hard': len(hard), - 'n_generated': len(all_cands), 'n_candidates_parseable': len(cands), - 'n_unparseable': len(all_cands) - len(cands), - 'n_leaked': sum(1 for c in cands if c['leaked']), - 'n_clean': sum(1 for c in cands if c['leaked'] is False), - 'n_reward_pos': sum(1 for c in scored if c['reward']), 'n_train_samples': len(train_cands), - 'n_train_from_fail': sum(1 for c in fail_cands if _is_trainable(c, args)), - 'abs_adv_from_fail_frac': (abs_adv(fail_cands) / total_abs) if total_abs > 0 else 0.0, - 'avg_baseline_pass_on_hard': base_acc, 'avg_withskill_pass': ws_acc, - 'avg_lift': ws_acc - base_acc, - 'termination_rate_withskill': (sum(1 for x in ws_rolls if x['terminated']) / len(ws_rolls)) if ws_rolls else 0.0, - 'view_A': _view_stats(hard, 'A'), 'view_B': _view_stats(hard, 'B'), - } - - -def _group_records(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> List[Dict[str, Any]]: - """GRPO training records: every trainable candidate with its view + rubric diagnosis - (the prompt is rebuilt from those by ``_skillgen_messages``, no trajectory stored).""" - out = [] - for r in chunk: - if not r['_hard']: - continue - for c in r['_cands']: - if _is_trainable(c, args): - out.append({ - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'view': r.get('_view', 'A'), 'diagnosis': r.get('_rubric_diag', ''), - 'response': c['response'], 'skills': c['skills'], - 'advantage': c['advantage'], 'grpo_adv': c['grpo_adv'], 'kept': c['kept'], - 'reward': c['reward'], 'with_pass': c['with_pass']}) - return out - - -# =========================================================================== -# Block F -- samplers, args, main -# =========================================================================== -def init_samplers(args: argparse.Namespace): - """8 GPUs: ranks 0-3 skill_sampler, ranks 4-7 base_sampler (both vLLM tp1 dp4).""" - twinkle.initialize(mode='ray', nproc_per_node=8, lazy_collect=False, groups=[ - DeviceGroup(name='skill_sampler', ranks=list(range(0, 4)), device_type='GPU'), - DeviceGroup(name='base_sampler', ranks=list(range(4, 8)), device_type='GPU')]) - samplers = [] - for group in ('skill_sampler', 'base_sampler'): - s = vLLMSampler( - model_id=MODEL_ID, - engine_args={'gpu_memory_utilization': GPU_MEM, - 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, - device_mesh=DeviceMesh.from_sizes(world_size=4, dp_size=4), remote_group=group) - s.set_template('Template', model_id=MODEL_ID, enable_thinking=True, - max_length=args.max_model_len) - samplers.append(s) - return samplers[1], samplers[0], 4, 4 # base_sampler, skill_sampler, base_dp, skill_dp - - -def _build_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--total-problems', type=int, default=3200, - help='Final number of problems selected into generated chunks.') - p.add_argument('--base-success-frac', type=float, default=0.3, - help='Target fraction of selected problems the frozen base solves.') - p.add_argument('--output-dir', default='./output/reflexion_skill_data') - p.add_argument('--cache-dir', default='', - help='Baseline/rubric cache dir (default /cache).') - p.add_argument('--no-cache', action='store_true', help='Disable disk cache read/write.') - p.add_argument('--overwrite', action='store_true', help='Replace existing output jsonl.') - p.add_argument('--seed', type=int, default=42) - p.add_argument('--dataset', choices=('aops', 'math'), default='aops') - p.add_argument('--n', type=int, default=0, - help='Raw train-pool size; 0 derives it from --total-problems.') - p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) - p.add_argument('--eval-size', type=int, default=128) - p.add_argument('--chunk-size', type=int, default=16) - p.add_argument('--balance', action=argparse.BooleanOptionalAction, default=True) - p.add_argument('--balance-loop-frac', type=float, default=0.5) - p.add_argument('--balance-max-draws-mult', type=int, default=8) - p.add_argument('--n-skills', type=int, default=8) - p.add_argument('--view-b-frac', type=float, default=0.5) - p.add_argument('--skill-retries', type=int, default=2) - p.add_argument('--skill-gen-temperature', type=float, default=1.0) - p.add_argument('--skill-gen-top-p', type=float, default=1.0) - p.add_argument('--skill-gen-top-k', type=int, default=-1) - p.add_argument('--max-model-len', type=int, default=16384) - p.add_argument('--max-tokens', type=int, default=8192) - p.add_argument('--skill-max-tokens', type=int, default=8192) - p.add_argument('--leak-workers', type=int, default=16) - p.add_argument('--rubric-workers', type=int, default=16) - p.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True) - args = p.parse_args() - if args.total_problems <= 0 or args.chunk_size <= 0: - raise ValueError('--total-problems and --chunk-size must be positive') - if not 0.0 <= args.base_success_frac <= 1.0: - raise ValueError('--base-success-frac must be in [0, 1]') - args.chunks = math.ceil(args.total_problems / args.chunk_size) - args.balance_success_frac = args.base_success_frac - if args.n <= 0: - args.n = max(args.total_problems + args.eval_size, math.ceil(args.total_problems * 1.5)) - return args - - -def _write(handle, row: Dict[str, Any]) -> None: - handle.write(json.dumps(row, ensure_ascii=False) + '\n') - - -def main() -> None: - args = _build_args() - records, eval_records, data_stats = _load_records(args) - if not records: - raise ValueError(f'loaded 0 {args.dataset} problems') - if len(records) < args.chunk_size: - raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded ({len(records)})') - - os.makedirs(args.output_dir, exist_ok=True) - data_path = os.path.join(args.output_dir, 'skill_dataset.jsonl') - gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') - eval_path = os.path.join(args.output_dir, 'eval_holdout.jsonl') - for path in (data_path, gen_path, eval_path): - if os.path.exists(path) and not args.overwrite: - raise FileExistsError(f'{path} exists; pass --overwrite to replace it') - - if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY')): - sys.stderr.write('[build] WARNING: no LLM backup env; leak/rubric checks degrade\n') - - base_sampler, skill_sampler, base_dp, skill_dp = init_samplers(args) - leak = LeakVerifier(sampler=None, answer_only=True) - checker = build_rubric_checker() - pool = ProblemPool(records, args.seed) - - cache_dir = args.cache_dir or os.path.join(args.output_dir, 'cache') - os.makedirs(cache_dir, exist_ok=True) - use_cache = not args.no_cache - baseline_cache = DiskCache(os.path.join(cache_dir, 'baseline.jsonl'), use_cache) - rubric_cache = DiskCache(os.path.join(cache_dir, 'rubric.jsonl'), use_cache) - - cfg = { - 'record_type': 'config', 'mode': 'offline_data_build', 'model': MODEL_ID, - 'dataset': args.dataset, 'n': len(records), 'eval_n': len(eval_records), - 'total_problems': args.total_problems, 'seed': args.seed, 'numeric_only': args.numeric_only, - 'raw_loaded': data_stats['raw_loaded'], 'numeric_dropped': data_stats['numeric_dropped'], - 'chunks': args.chunks, 'chunk_size': args.chunk_size, 'n_skills': args.n_skills, - 'view_b_frac': args.view_b_frac, 'balance': args.balance, - 'base_success_frac': args.base_success_frac, 'balance_success_frac': args.balance_success_frac, - 'reward': 'greedy_binary(correct)', 'advantage': 'group_relative', - 'format_in_reward': args.format_in_reward, 'cache': use_cache, - 'rubric_check': 'fixed_math_5crit(viewA)' if checker else 'disabled', - 'started': int(time.time()), - } - total_groups, selected = 0, 0 - with open(gen_path, 'w', encoding='utf-8') as gen_f, \ - open(data_path, 'w', encoding='utf-8') as data_f, \ - open(eval_path, 'w', encoding='utf-8') as eval_f: - for handle in (gen_f, data_f, eval_f): - _write(handle, cfg) - for rec in eval_records: - _write(eval_f, {'record_type': 'eval_holdout', **rec}) - eval_f.flush() - - full_chunk_size = args.chunk_size - for ci in range(args.chunks): - remaining = args.total_problems - selected - if remaining <= 0: - break - args.chunk_size = min(full_chunk_size, remaining) # last chunk may be short - chunk, balance = draw_chunk(pool, base_sampler, base_dp, args, baseline_cache) - full, summary, groups = process_chunk( - base_sampler, skill_sampler, leak, chunk, ci, base_dp, skill_dp, - args, checker, rubric_cache) - summary['balance'] = balance - for rec in full: - _write(gen_f, rec) - _write(gen_f, summary) - gen_f.flush() - for row in groups: - _write(data_f, {'chunk': ci, **row}) - data_f.flush() - total_groups += len(groups) - selected += len(chunk) - sys.stderr.write( - f'[build] g{ci}: problems={selected}/{args.total_problems} ' - f'train={len(groups)} total={total_groups} ' - f'acc={summary["avg_baseline_pass_on_hard"]:.2f}->{summary["avg_withskill_pass"]:.2f} ' - f'lift={summary["avg_lift"]:+.3f}\n') - - baseline_cache.close() - rubric_cache.close() - sys.stderr.write(f'[build] done: {total_groups} train records -> {data_path}; trace -> {gen_path}\n') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/embedding/build_thinking_rag_index.py b/cookbook/exp/embedding/build_thinking_rag_index.py deleted file mode 100644 index a71bae060..000000000 --- a/cookbook/exp/embedding/build_thinking_rag_index.py +++ /dev/null @@ -1,1159 +0,0 @@ -"""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 ``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 - ``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 -import threading -import time -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 - -# --------------------------------------------------------------------------- -# Compress prompts — MUST match train_embedding_full_ddp.py exactly. -# --------------------------------------------------------------------------- -_HERE = Path(__file__).resolve().parent -sys.path.insert(0, str(_HERE)) - -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 - -_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', - 'output/embedding_full_transformers/last-checkpoint', -) -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)) - -# Dataset mix caps (only used in 'both' mode). None = no cap. -THINK_CAP: Optional[int] = int(os.environ.get('THINK_CAP', 400_000)) or None -INDEX_CAP: Optional[int] = int(os.environ.get('INDEX_CAP', 400_000)) or None -MIX_SHUFFLE_SEED = 100 - -# Concurrency knobs for API fallback and prefetch pipeline. -API_CONCURRENCY = int(os.environ.get('API_CONCURRENCY', 8)) -API_MIN_INTERVAL = float(os.environ.get('API_MIN_INTERVAL', 0.1)) -PREFETCH_WORKERS = int(os.environ.get('PREFETCH_WORKERS', 2)) - -# Hard-templated hints: the condenser SFT prior maps `Skill` to the legacy -# `Use when: / numbered steps / Output:` skeleton on long inputs; embedding the -# exact 4-line body template + explicit negative constraints is the only way to -# override it deterministically across query and cot sides. -RAG_QUERY_HINT = ( - '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.') -RAG_THINKING_HINT = ( - '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.') - -# 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 -# =========================================================================== - -_LEGACY_USE_WHEN_RE = re.compile(r'(?im)^\s*Use when\s*:') -_SCHEMA_MARKERS = ('Problem:', 'Skill:', 'Knowledge:') - - -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 - 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 - 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 _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) - - -_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. - - API fallback runs concurrently (up to API_CONCURRENCY workers) for speed. - """ - pairs = _vllm_compress(sampler, texts, query_hint) - 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[i] = text - else: - fallback_indices.append(i) - - if fallback_indices and api is not None: - from concurrent.futures import ThreadPoolExecutor, as_completed - with ThreadPoolExecutor(max_workers=API_CONCURRENCY) as pool: - futures = {} - for idx in fallback_indices: - msgs = _build_compress_messages(texts[idx], query_hint) - futures[pool.submit(_api_compress_throttled, api, msgs)] = idx - for fut in as_completed(futures): - idx = futures[fut] - api_text = fut.result() - if api_text and not _is_truncated_compression(api_text): - results[idx] = api_text - - return results - - -def _resolve_compressed_multi(sampler: vLLMSampler, api: Optional[OpenAIClient], - texts: List[str], hints: List[str]) -> List[Optional[str]]: - """Like _resolve_compressed but each text has its own per-item hint. - - Merges all texts into a SINGLE vLLM batch call (instead of one per hint), - dramatically reducing round-trip overhead when processing interleaved - query+cot pairs with different hint strings. - - Args: - sampler: vLLM condenser sampler. - api: Optional OpenAI-compatible API client for fallback. - texts: List of raw texts to compress (may contain empty strings to skip). - hints: Per-text hint strings (same length as texts). - - Returns: - List of compressed texts (None where compression failed entirely). - """ - assert len(texts) == len(hints), f'texts({len(texts)}) != hints({len(hints)})' - if not texts: - return [] - - # Skip texts that would exceed the condenser's context window. - _max_input_chars = (CONDENSE_MAX_MODEL_LEN - CONDENSE_MAX_TOKENS) * 3 - skip_mask = [len(t) > _max_input_chars for t in texts] - - # Build prompts per-item (each text gets its own hint as the query parameter). - prompts = [{'messages': _build_compress_messages(t, h)} - for t, h, skip in zip(texts, hints, skip_mask) if not skip] - active_indices = [i for i, skip in enumerate(skip_mask) if not skip] - params = TwinkleSamplingParams( - max_tokens=CONDENSE_MAX_TOKENS, - temperature=COMPRESS_TEMPERATURE, - top_p=COMPRESS_TOP_P, - num_samples=1, - ) - - # Single vLLM batch call — the key throughput win. - responses = sampler.sample(prompts, params) if prompts else [] - - results: List[Optional[str]] = [None] * len(texts) - fallback_indices: List[int] = [] - for resp_idx, orig_idx in enumerate(active_indices): - resp = responses[resp_idx] - seq = resp.sequences[0] if resp and resp.sequences else None - if seq is None: - fallback_indices.append(orig_idx) - continue - text = seq.decoded or '' - text = re.sub(r'<\|[^|]+\|>', '', text).rstrip() - text = _strip_outer_codefence(text) - if seq.stop_reason != 'length' and not _is_truncated_compression(text): - results[orig_idx] = text - else: - fallback_indices.append(orig_idx) - - # Concurrent API fallback for failed items. - if fallback_indices and api is not None: - from concurrent.futures import ThreadPoolExecutor, as_completed - with ThreadPoolExecutor(max_workers=API_CONCURRENCY) as pool: - futures = {} - for idx in fallback_indices: - msgs = _build_compress_messages(texts[idx], hints[idx]) - futures[pool.submit(_api_compress_throttled, api, msgs)] = idx - for fut in as_completed(futures): - idx = futures[fut] - api_text = fut.result() - if api_text and not _is_truncated_compression(api_text): - results[idx] = api_text - - return results - - -# =========================================================================== -# 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 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' - + (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 - 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 _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 [] - # 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 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: - 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), - 'cot_raw_head': _short(r['cot_raw'], 200), - }) - continue - 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 - 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) - 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: - 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), - '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) - 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 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: - 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, - 'source': row.get('source') or 'unknown', - 'query_raw': user_query, - 'cot_raw': cot, - }) - if len(batch) >= args.batch_size: - pending_futures.append(prefetch_pool.submit(_process_batch, list(batch))) - batch.clear() - - # Flush remainder - if 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] 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: - 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], RAG_QUERY_HINT) - 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='both', - choices=['dataset_index', 'dataset_think', 'both'], - help='Which loader to use: dataset_index (RAG profile), ' - 'dataset_think (training mix), or both (50/50 mix).') - p.add_argument('--limit', type=int, default=0, - help='Stop building once this many rows are kept (0 = no cap).') - p.add_argument('--max-rows', type=int, default=0, - help='Truncate corpus to this many rows AFTER get_dataset (0 = no cap). ' - 'Use this instead of --total to avoid invalidating the dataset cache.') - p.add_argument('--batch-size', type=int, default=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', - 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 - 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. - 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/compare_math_levels.py b/cookbook/exp/embedding/compare_math_levels.py deleted file mode 100644 index d488909b9..000000000 --- a/cookbook/exp/embedding/compare_math_levels.py +++ /dev/null @@ -1,91 +0,0 @@ -"""Compare MATH direct vs RAG by difficulty level. - -Re-grades both result files with the production ``answers_match`` (so the -stored ``is_correct`` is never trusted) and prints the per-level accuracy -plus the RAG gain (delta) so you can see how it varies with difficulty. - -Defaults to the raw-RAG output (``math_rag_results.jsonl``); pass a second -arg to compare a different rag file (e.g. ``math_rag_hint_results.jsonl``). - -Usage: - python cookbook/exp/embedding/compare_math_levels.py \ - [direct.jsonl] [rag.jsonl] -""" -import importlib.util -import json -import os -import sys -from collections import defaultdict - -_HERE = os.path.dirname(os.path.abspath(__file__)) - - -def _load_grader(): - spec = importlib.util.spec_from_file_location( - 'egr', os.path.join(_HERE, 'eval_gpqa_rag.py')) - egr = importlib.util.module_from_spec(spec) - spec.loader.exec_module(egr) - return egr.answers_match - - -def _load(path): - return {json.loads(l)['idx']: json.loads(l) - for l in open(path, encoding='utf-8') if l.strip()} - - -def main(): - direct_path = sys.argv[1] if len(sys.argv) > 1 else \ - './output/thinking_rag/math_direct_results.jsonl' - hint_path = sys.argv[2] if len(sys.argv) > 2 else \ - './output/thinking_rag/math_rag_results.jsonl' - - answers_match = _load_grader() - D = _load(direct_path) - H = _load(hint_path) - common = sorted(set(D) & set(H)) - print(f'direct={len(D)} rag+hint={len(H)} common={len(common)}') - - def runaway(rec): - mo = rec.get('model_output') or '' - return ('' not in mo) or ( - not (rec.get('predicted') or '').strip() and len(mo) > 40000) - - def correct(rec): - return answers_match(rec.get('predicted') or '', - rec['reference_answer']) - - # level -> counters - per = defaultdict(lambda: {'n': 0, 'd': 0, 'h': 0, - 'd_run': 0, 'h_run': 0}) - for i in common: - lv = H[i].get('level') or D[i].get('level') or 'Unknown' - c = per[lv] - c['n'] += 1 - c['d'] += int(correct(D[i])) - c['h'] += int(correct(H[i])) - c['d_run'] += int(runaway(D[i])) - c['h_run'] += int(runaway(H[i])) - - print(f'\n{"level":>10} | {"n":>4} | {"direct":>7} | {"rag+hint":>8} | ' - f'{"delta":>7} | {"d_run":>6} | {"h_run":>6}') - print('-' * 68) - tot = {'n': 0, 'd': 0, 'h': 0, 'd_run': 0, 'h_run': 0} - for lv in sorted(per.keys()): - c = per[lv] - for k in tot: - tot[k] += c[k] - n = c['n'] - dacc, hacc = c['d'] / n, c['h'] / n - print(f'{lv:>10} | {n:>4} | {dacc:>7.3f} | {hacc:>8.3f} | ' - f'{hacc - dacc:>+7.3f} | {c["d_run"]/n:>6.1%} | ' - f'{c["h_run"]/n:>6.1%}') - print('-' * 68) - n = tot['n'] - if n: - print(f'{"OVERALL":>10} | {n:>4} | {tot["d"]/n:>7.3f} | ' - f'{tot["h"]/n:>8.3f} | {(tot["h"]-tot["d"])/n:>+7.3f} | ' - f'{tot["d_run"]/n:>6.1%} | {tot["h_run"]/n:>6.1%}') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/embedding/dataset_hard.py b/cookbook/exp/embedding/dataset_hard.py deleted file mode 100644 index 9fa059b95..000000000 --- a/cookbook/exp/embedding/dataset_hard.py +++ /dev/null @@ -1,202 +0,0 @@ -"""Hard-negative dataset for embedding training. - -Provides ReasonIR (AI-ModelScope/reasonir-data, hq subset): - - query: reasoning-intensive question - - positive: BRIGHT document (resolved via xlangai/BRIGHT documents corpus) - - negatives: plausibly related but ultimately unhelpful documents - -Output schema: ``{id, source, query, cot, response, negatives}`` -where ``negatives`` is a list of strings (each a separate hard negative). -""" -import hashlib -import os -import sys -from pathlib import Path -from typing import Any, Dict, List, Optional - -from datasets import Dataset as HFDataset -from modelscope import MsDataset - -_CACHE_DIR = Path(__file__).resolve().parent / '.cache_hard' - - -def _hash_id(prefix: str, content: str) -> str: - return f'{prefix}__{hashlib.md5(content.encode("utf-8")).hexdigest()[:16]}' - - -# --------------------------------------------------------------------------- -# BRIGHT document corpus (lazy singleton) -# --------------------------------------------------------------------------- - -_BRIGHT_SPLITS = [ - 'aops', 'biology', 'earth_science', 'economics', 'leetcode', 'pony', - 'psychology', 'robotics', 'stackoverflow', 'sustainable_living', - 'theoremqa_questions', 'theoremqa_theorems', -] - -_bright_docs: Optional[Dict[str, str]] = None - - -def _load_bright_docs() -> Dict[str, str]: - """Load all BRIGHT document splits into {id -> content} lookup dict.""" - global _bright_docs - if _bright_docs is not None: - return _bright_docs - sys.stderr.write('[dataset_hard] Loading BRIGHT documents corpus...\n') - _bright_docs = {} - for split in _BRIGHT_SPLITS: - try: - ds = MsDataset.load( - 'xlangai/BRIGHT', subset_name='documents', split=split, - download_mode='reuse_dataset_if_exists') - for row in ds: - doc_id = row.get('id', '') - content = row.get('content', '') - if doc_id and content: - _bright_docs[doc_id] = content - short = doc_id.rsplit('/', 1)[-1] if '/' in doc_id else doc_id - if short not in _bright_docs: - _bright_docs[short] = content - sys.stderr.write(f' [{split}] loaded {len(ds)} docs\n') - except Exception as e: - sys.stderr.write(f' [{split}] FAILED: {e}\n') - sys.stderr.write(f'[dataset_hard] BRIGHT total: {len(_bright_docs)} entries\n') - return _bright_docs - - -# --------------------------------------------------------------------------- -# ReasonIR dataset -# --------------------------------------------------------------------------- - -def get_dataset_reasonir(max_rows: Optional[int] = None, - max_negatives: int = 16, - load_from_cache_file: bool = True) -> HFDataset: - """Load AI-ModelScope/reasonir-data (hq subset) with BRIGHT doc resolution. - - Schema: {id, source, query, cot, response, negatives} - """ - cache_key = f'reasonir_neg{max_negatives}' - cache_path = _CACHE_DIR / cache_key - if load_from_cache_file and cache_path.exists(): - sys.stderr.write(f'[reasonir] loading from cache: {cache_path}\n') - ds = HFDataset.load_from_disk(str(cache_path)) - if max_rows and len(ds) > max_rows: - ds = ds.select(range(max_rows)) - sys.stderr.write(f'[reasonir] {len(ds)} rows (cached)\n') - return ds - - ds = MsDataset.load( - 'AI-ModelScope/reasonir-data', subset_name='hq', split='train', - download_mode='reuse_dataset_if_exists') - if max_rows and len(ds) > max_rows: - ds = ds.select(range(max_rows)) - - bright = _load_bright_docs() - rows = [] - n_miss = 0 - for row in ds: - query_parts = row.get('query', []) - if not isinstance(query_parts, list) or len(query_parts) < 2: - continue - query = query_parts[1].strip() - if not query: - continue - - pos_list = row.get('pos', []) - if not pos_list: - continue - pos_id = pos_list[0][1] if isinstance(pos_list[0], list) and len(pos_list[0]) > 1 else '' - cot = bright.get(pos_id, '') - if not cot: - n_miss += 1 - continue - - neg_list = row.get('neg', []) - negatives = [] - for neg in neg_list: - if isinstance(neg, list) and len(neg) > 1: - neg_text = neg[1].strip() - if neg_text: - negatives.append(neg_text) - if len(negatives) >= max_negatives: - break - - if not negatives: - continue - - rows.append({ - 'id': _hash_id('reasonir', f'{query}\n{pos_id}'), - 'source': 'reasonir-hq', - 'query': query, - 'cot': cot, - 'response': '', - 'negatives': negatives, - }) - - if n_miss: - sys.stderr.write(f'[reasonir] {n_miss} rows skipped (BRIGHT doc not found)\n') - sys.stderr.write(f'[reasonir] {len(rows)} rows with hard negatives\n') - result = HFDataset.from_dict(_rows_to_cols(rows)) - # Persist full dataset; max_rows is applied post-cache for flexibility. - cache_path.parent.mkdir(parents=True, exist_ok=True) - result.save_to_disk(str(cache_path)) - sys.stderr.write(f'[reasonir] cached to {cache_path}\n') - if max_rows and len(result) > max_rows: - result = result.select(range(max_rows)) - return result - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - -def _rows_to_cols(rows: List[Dict[str, Any]]) -> Dict[str, list]: - if not rows: - return {'id': [], 'source': [], 'query': [], 'cot': [], - 'response': [], 'negatives': []} - keys = rows[0].keys() - return {k: [r[k] for r in rows] for k in keys} - - -# --------------------------------------------------------------------------- -# Public API -# --------------------------------------------------------------------------- - -def get_dataset( - reasonir_max: Optional[int] = None, - max_negatives: int = 16, - load_from_cache_file: bool = True, - **kwargs, -) -> HFDataset: - """Load hard-negative dataset (reasonir only). - - Returns HF Dataset with schema: {id, source, query, cot, response, negatives} - """ - ds = get_dataset_reasonir(max_rows=reasonir_max, max_negatives=max_negatives, - load_from_cache_file=load_from_cache_file) - if len(ds) == 0: - sys.stderr.write('[dataset_hard] WARNING: reasonir dataset empty\n') - else: - sys.stderr.write(f'[dataset_hard] reasonir={len(ds)}\n') - return ds - - -if __name__ == '__main__': - import argparse - parser = argparse.ArgumentParser() - parser.add_argument('--reasonir-max', type=int, default=1000) - args = parser.parse_args() - - ds = get_dataset(reasonir_max=args.reasonir_max) - print(f'Total rows: {len(ds)}') - print(f'Features: {ds.features}') - if len(ds) > 0: - row = ds[0] - print(f'\nSample[0]:') - print(f' id: {row["id"]}') - print(f' source: {row["source"]}') - print(f' query: {row["query"][:100]}...') - print(f' cot: {row["cot"][:100]}...') - print(f' negatives: {len(row["negatives"])} items') - if row['negatives']: - print(f' [0]: {row["negatives"][0][:80]}...') diff --git a/cookbook/exp/embedding/dataset_index.py b/cookbook/exp/embedding/dataset_index.py deleted file mode 100644 index 7d2905a59..000000000 --- a/cookbook/exp/embedding/dataset_index.py +++ /dev/null @@ -1,718 +0,0 @@ -"""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/dataset_think.py b/cookbook/exp/embedding/dataset_think.py deleted file mode 100644 index 38618ced1..000000000 --- a/cookbook/exp/embedding/dataset_think.py +++ /dev/null @@ -1,456 +0,0 @@ -import hashlib -import re -from typing import Any, Dict, List, Optional - -from twinkle.dataset import Dataset, DatasetMeta -from twinkle.preprocessor import Preprocessor - -_THINK_RE = re.compile(r'(.*?)', re.DOTALL) - - -def _hash_id(prefix: str, content: str) -> str: - return f'{prefix}__{hashlib.md5(content.encode("utf-8")).hexdigest()[:16]}' - - -def _register(dataset, processor_cls, meta: DatasetMeta, init_args: Optional[Dict[str, Any]] = None, - load_from_cache_file: bool = True) -> None: - """Add dataset and run preprocessor; auto-strip every input column to enforce - the universal ``{id, source, query, cot, response}`` output schema.""" - dataset.add_dataset(meta) - cols = list(dataset.datasets[meta.get_id()].column_names) - dataset.map( - processor_cls, - dataset_meta=meta, - init_args=init_args or {}, - remove_columns=cols, - load_from_cache_file=load_from_cache_file, - ) - - -# ===== Modotte/CodeX-2M-Thinking ===== -CODEX_THINKING_REPO = 'ms://Modotte/CodeX-2M-Thinking' - - -class CodeXThinkingProcessor(Preprocessor): - """CodeX-2M-Thinking row → ``{id, source, query, cot, response}``。 - - 输入 schema: ``input``（问题）、``output``（含 ``...`` + 答案）。 - 拆分 output 为 cot（think 标签内容）和 response（标签之后的正文）。 - 丢弃缺失 input/output 或无法解析 think 标签的行。 - """ - - 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 '').strip() - output = (row.get('output') or '').strip() - if not query or not output: - continue - m = _THINK_RE.search(output) - if not m: - continue - cot = m.group(1).strip() - response = output[m.end():].strip() - if not cot or not response: - continue - out.append({ - 'id': _hash_id('codex_think', f'{query}\n{response}'), - 'source': 'CodeX-2M-Thinking', - 'query': query, - 'cot': cot, - 'response': response, - }) - return self.map_row_to_col(out) - - -# ===== open-thoughts/OpenThoughts3-1.2M ===== -OPEN_THOUGHTS_REPO = 'ms://open-thoughts/OpenThoughts3-1.2M' - - -class OpenThoughtsProcessor(Preprocessor): - """OpenThoughts3 row → ``{id, source, query, cot, response}``。 - - 输入 schema: ``conversations`` (messages 格式 list[{from/value}])。 - 取第一个 human 作 query，第一个 gpt 的 value 按 ``...`` 拆 cot/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: - convs = row.get('conversations') - if not isinstance(convs, list): - continue - query = '' - assistant_text = '' - for msg in convs: - if not isinstance(msg, dict): - continue - role = msg.get('from') or msg.get('role') or '' - value = msg.get('value') or msg.get('content') or '' - if role in ('human', 'user') and not query: - query = value.strip() - elif role in ('gpt', 'assistant') and not assistant_text: - assistant_text = value.strip() - break - if not query or not assistant_text: - continue - m = _THINK_RE.search(assistant_text) - if not m: - continue - cot = m.group(1).strip() - response = assistant_text[m.end():].strip() - if not cot or not response: - continue - out.append({ - 'id': _hash_id('openthoughts', f'{query}\n{response}'), - 'source': 'OpenThoughts3-1.2M', - 'query': query, - 'cot': cot, - 'response': response, - }) - return self.map_row_to_col(out) - - -# ===== GAIR/LIMO-v2 ===== -LIMO_REPO = 'ms://GAIR/LIMO-v2' - - -class LIMOProcessor(Preprocessor): - """LIMO-v2 row → ``{id, source, query, cot, response}``。 - - 输入 schema: ``question``、``solution``（含 ``...`` + 答案）。 - 拆分 solution 为 cot 和 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('question') or '').strip() - solution = (row.get('solution') or '').strip() - if not query or not solution: - continue - m = _THINK_RE.search(solution) - if m: - cot = m.group(1).strip() - response = solution[m.end():].strip() - else: - # 无 think 标签时，solution 整体作为 response，cot 留空 - cot = '' - response = solution - if not response: - continue - out.append({ - 'id': _hash_id('limo', f'{query}\n{response}'), - 'source': 'LIMO-v2', - 'query': query, - 'cot': cot, - 'response': response, - }) - return self.map_row_to_col(out) - - -# ===== AI-ModelScope/Chinese-DeepSeek-R1-Distill-data-110k ===== -CN_R1_DISTILL_REPO = 'ms://AI-ModelScope/Chinese-DeepSeek-R1-Distill-data-110k' - - -class ChineseR1DistillProcessor(Preprocessor): - """Chinese-DeepSeek-R1-Distill row → ``{id, source, query, cot, response}``。 - - 输入已有三列: ``input`` → query, ``reasoning_content`` → cot, ``content`` → 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('input') or '').strip() - cot = (row.get('reasoning_content') or '').strip() - response = (row.get('content') or '').strip() - if not query or not response: - continue - if cot: - response = _THINK_RE.sub('', response).strip() - if not response: - continue - out.append({ - 'id': _hash_id('cn_r1_distill', f'{query}\n{response}'), - 'source': 'Chinese-DeepSeek-R1-Distill-data-110k', - 'query': query, - 'cot': cot, - 'response': response, - }) - return self.map_row_to_col(out) - - -# ===== nohurry/Opus-4.6-Reasoning-3000x-filtered ===== -OPUS_REASONING_REPO = 'ms://nohurry/Opus-4.6-Reasoning-3000x-filtered' - - -class OpusReasoningProcessor(Preprocessor): - """Opus-4.6-Reasoning-3000x-filtered row → ``{id, source, query, cot, response}``。 - - 输入已有三列: ``problem`` → query, ``thinking`` → cot, ``solution`` → 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 '').strip() - cot = (row.get('thinking') or '').strip() - response = (row.get('solution') or '').strip() - if not query or not response: - continue - if cot: - response = _THINK_RE.sub('', response).strip() - if not response: - continue - out.append({ - 'id': _hash_id('opus_reasoning', f'{query}\n{response}'), - 'source': 'Opus-4.6-Reasoning-3000x-filtered', - 'query': query, - 'cot': cot, - 'response': response, - }) - return self.map_row_to_col(out) - - -# ===== Roman1111111/claude-opus-4.6-10000x ===== -CLAUDE_OPUS_REPO = 'ms://Roman1111111/claude-opus-4.6-10000x' - - -class ClaudeOpusProcessor(Preprocessor): - """claude-opus-4.6-10000x row → ``{id, source, query, cot, response}``。 - - 输入 schema: ``messages`` (OpenAI 格式 list[{role, content}])。 - 取首个 user 作 query，首个 assistant 按 ``...`` 拆 cot/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: - messages = row.get('messages') - if not isinstance(messages, list): - continue - query = '' - assistant_text = '' - for msg in messages: - if not isinstance(msg, dict): - continue - role = msg.get('role') or '' - content = msg.get('content') or '' - if not isinstance(content, str): - continue - if role == 'user' and not query: - query = content.strip() - elif role == 'assistant' and not assistant_text: - assistant_text = content.strip() - break - if not query or not assistant_text: - continue - m = _THINK_RE.search(assistant_text) - if m: - cot = m.group(1).strip() - response = assistant_text[m.end():].strip() - else: - cot = '' - response = assistant_text - if not response: - continue - out.append({ - 'id': _hash_id('claude_opus', f'{query}\n{response}'), - 'source': 'claude-opus-4.6-10000x', - 'query': query, - 'cot': cot, - 'response': response, - }) - return self.map_row_to_col(out) - - -ANGRYGIRAFFE_REPO = 'ms://hf/angrygiraffe-claude-opus-4.6-4.7-reasoning-8.7k' - - -class AngrygiraffeOpusReasoningProcessor(Preprocessor): - """angrygiraffe/claude-opus-4.6-4.7-reasoning-8.7k row → ``{id, source, query, cot, response}``。 - - 输入 schema: ``messages`` (OpenAI 格式 list[{role, content}])。 - 取首个 user 作 query，首个 assistant 按 ``...`` 拆 cot/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: - messages = row.get('messages') - if not isinstance(messages, list): - continue - query = '' - assistant_text = '' - for msg in messages: - if not isinstance(msg, dict): - continue - role = msg.get('role') or '' - content = msg.get('content') or '' - if not isinstance(content, str): - continue - if role == 'user' and not query: - query = content.strip() - elif role == 'assistant' and not assistant_text: - assistant_text = content.strip() - break - if not query or not assistant_text: - continue - m = _THINK_RE.search(assistant_text) - if m: - cot = m.group(1).strip() - response = assistant_text[m.end():].strip() - else: - cot = '' - response = assistant_text - if not response: - continue - out.append({ - 'id': _hash_id('angrygiraffe_opus', f'{query}\n{response}'), - 'source': 'angrygiraffe-claude-opus-4.6-4.7-reasoning-8.7k', - 'query': query, - 'cot': cot, - 'response': response, - }) - return self.map_row_to_col(out) - - -_BASE_SIZES = { - 'codex_think': 100000, - 'open_thoughts': 400000, - 'cn_r1_distill': 100000, - 'opus_reasoning': 3000, - 'claude_opus': 10000, - 'angrygiraffe': 38000, -} - - -def _scaled_sizes(total: Optional[int]) -> Dict[str, int]: - if total is None: - 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, CodeXThinkingProcessor, - DatasetMeta(dataset_id=CODEX_THINKING_REPO, split='train', - data_slice=range(sizes['codex_think'])), - load_from_cache_file=load_from_cache_file) - - _register(dataset, OpenThoughtsProcessor, - DatasetMeta(dataset_id=OPEN_THOUGHTS_REPO, split='train', - data_slice=range(sizes['open_thoughts'])), - load_from_cache_file=load_from_cache_file) - - _register(dataset, LIMOProcessor, - DatasetMeta(dataset_id=LIMO_REPO, split='train'), - load_from_cache_file=load_from_cache_file) - - _register(dataset, ChineseR1DistillProcessor, - DatasetMeta(dataset_id=CN_R1_DISTILL_REPO, split='train', - data_slice=range(sizes['cn_r1_distill'])), - load_from_cache_file=load_from_cache_file) - - _register(dataset, OpusReasoningProcessor, - DatasetMeta(dataset_id=OPUS_REASONING_REPO, split='train', - data_slice=range(sizes['opus_reasoning'])), - load_from_cache_file=load_from_cache_file) - - _register(dataset, ClaudeOpusProcessor, - DatasetMeta(dataset_id=CLAUDE_OPUS_REPO, split='train', - data_slice=range(sizes['claude_opus'])), - load_from_cache_file=load_from_cache_file) - - _register(dataset, AngrygiraffeOpusReasoningProcessor, - DatasetMeta(dataset_id=ANGRYGIRAFFE_REPO, split='train', - data_slice=range(sizes['angrygiraffe'])), - load_from_cache_file=load_from_cache_file) - - dataset.mix_dataset(False) - return dataset - - -class ToMessagesProcessor(Preprocessor): - """Convert {query, cot, response} → {id, source, messages}.""" - - 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('query') or '' - cot = row.get('cot') or '' - response = row.get('response') or '' - if not cot: - continue - assistant_content = f'{cot}' - out.append({ - 'id': row.get('id', ''), - 'source': row.get('source', ''), - 'messages': [ - {'role': 'user', 'content': query}, - {'role': 'assistant', 'content': assistant_content, - 'reasoning_content': cot}, - ], - }) - return self.map_row_to_col(out, keys=['id', 'source', 'messages']) - - -def get_dataset(total: Optional[int] = None, dropped_log: Optional[str] = None, - load_from_cache_file: bool = True) -> Dataset: - """Build, convert to messages format, and quality-filter the CoT dataset. - - If ``total`` is given, every per-source row count in ``_BASE_SIZES`` is - scaled proportionally so the input-row sum approximates ``total``. - """ - from twinkle_agentic.preprocessor import ( - DeadLoopFilter, - FixUnicodeFilter, - HardFilter, - IntentClassifier, - MessageSanityFilter, - QualityPreprocessor, - RefuseFilter, - RemoveRepeatSentencesFilter, - TokenNumFilter, - TokenSoupFilter, - ) - - dataset = _build_dataset(total=total, 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.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/eval_dualline_math.py b/cookbook/exp/embedding/eval_dualline_math.py deleted file mode 100644 index bd1c183c7..000000000 --- a/cookbook/exp/embedding/eval_dualline_math.py +++ /dev/null @@ -1,689 +0,0 @@ -"""Dual-line math evaluation: baseline vs online process-checking + rubric injection. - -This is **Phase 0 of DESIGN §11.6** ("参数化 memory: 查错 LoRA"): before training any -LoRA, test the *upper bound* of the mechanism "pause every N tokens, let a strong -teacher check the partial reasoning for rubric errors, inject the found issue back -into the context, then resume". If even the strongest teacher checking online cannot -lift math accuracy, distilling that ability into a LoRA is pointless — so we gate on -this first. - -It deliberately reuses the SAME dataset loader, sampling params and answer grader as -``eval_gpqa_rag.py`` so the two lines are directly comparable: - - - **Line A — baseline** (``--mode baseline``): the student model solves each problem - in a single pass (identical to ``eval_gpqa_rag.py --mode direct``). - - **Line B — dualline** (``--mode dualline``, default): the student generates in - ``--chunk-tokens`` slices; between slices a teacher ``RubricVerifier.diagnose()`` - inspects the full reasoning so far (query + all prior response). When it reports - process issues, the finding is injected back as a first-person self-correction - (in the student's own voice) and generation resumes. - -The teacher checker is the ``llm_backup`` teacher API (no student sampler is given to -the verifier, so every check is served by the teacher — exactly the Phase-0 setup). -Configure it via the ``LLM_BACKUP_*`` env vars (see ``utils/llm_backup.py``). - -Continuation is done at the token level (crude on purpose — §11.6 says experiment -performance is not a concern): each slice re-feeds the prior ``new_input_feature`` and, -on injection, splices the tokenized note in before resuming. - -The dataset defaults to AoPS (``--dataset aops``), which auto-downloads from -ModelScope so no local data path is needed; pass ``--dataset math`` to use the -local Hendrycks MATH set instead. Both lines MUST share ``--dataset``, ``--n``, -``--target-eval`` and ``--seed`` to stay a paired comparison. - -Launch examples: - # Dual-line on 200 AoPS problems (needs LLM_BACKUP_* for the teacher checker) - LLM_BACKUP_API_KEY=sk-... LLM_BACKUP_BASE_URL=... \\ - python cookbook/exp/embedding/eval_dualline_math.py \\ - --n 200 --target-eval 200 --seed 42 - - # Paired baseline on the same subset (no checker calls) - python cookbook/exp/embedding/eval_dualline_math.py --mode baseline \\ - --n 200 --target-eval 200 --seed 42 -""" -import argparse -import copy -import json -import os -import sys -import time -from collections import defaultdict -from typing import Any, Dict, List, Optional - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.data_format import SamplingParams as TwinkleSamplingParams -from twinkle.sampler import vLLMSampler - -# Reuse the reference eval's dataset + grading + prompts verbatim so the two -# lines are measured on identical footing. -from eval_gpqa_rag import (GEN_MODEL_ID, GEN_GPU_MEM, GEN_GPUS, GEN_TEMPERATURE, - GEN_TOP_P, answers_match, build_direct_prompt, - extract_boxed, load_aops, load_math) - -# Dualline eval defaults (override via --max-model-len or DUALLINE_MAX_MODEL_LEN). -DUALLINE_DEFAULT_MAX_MODEL_LEN = int(os.environ.get('DUALLINE_MAX_MODEL_LEN', 32000)) -DUALLINE_DEFAULT_MAX_GEN_TOKENS = int( - os.environ.get('DUALLINE_MAX_GEN_TOKENS', DUALLINE_DEFAULT_MAX_MODEL_LEN)) - -# vLLM parallel: default tp=1, dp=GEN_GPUS (override with GEN_TP / keep GEN_GPUS=8). -GEN_TP = int(os.environ.get('GEN_TP', 1)) - -logger = get_logger() - -# --------------------------------------------------------------------------- -# Dual-line config -# --------------------------------------------------------------------------- -CHUNK_TOKENS = int(os.environ.get('DUALLINE_CHUNK_TOKENS', 512)) -MAX_CHECKS = int(os.environ.get('DUALLINE_MAX_CHECKS', 8)) -MAX_INJECTIONS = int(os.environ.get('DUALLINE_MAX_INJECTIONS', 3)) -# Only inject when the checker is confident enough that something is wrong. -CHECK_SCORE_FLOOR = float(os.environ.get('DUALLINE_CHECK_FLOOR', 0.6)) -# The note is written in the student's own first-person voice so, when spliced -# back in, the running model treats it as its own mid-thought self-correction -# rather than an external interruption (which tended to derail generation toward -# max-length). Kept short to limit disruption. -INJECT_TEMPLATE = ( - '\n\nWait — reviewing my reasoning above, I realize there is a problem: {issue}\n' - 'Let me correct this and continue.\n\n') - -# When context hits max_model_len (or sample fails), dump query + generation here. -OVERFLOW_DUMP_DIR = os.environ.get( - 'DUALLINE_OVERFLOW_DUMP_DIR', './output/dualline/overflow_dumps') - - -def _decode(tokenizer, ids: List[int]) -> str: - return tokenizer.decode(ids, skip_special_tokens=True) - - -def _input_ids_len(cur_inputs: Any) -> Optional[int]: - """Length of the tokenized prompt fed to vLLM on this step, if known.""" - if not cur_inputs: - return None - item = cur_inputs[0] - if isinstance(item, dict) and 'input_ids' in item: - ids = item['input_ids'] - return len(ids) if ids is not None else None - return None - - -def _dump_dualline_state( - *, - reason: str, - problem: str, - debug_idx: Optional[int], - chunk_tokens: int, - cur_inputs: Any, - gen_ids: List[int], - injected_ids: List[int], - tokenizer, - n_checks: int, - n_injections: int, - findings: List[Dict[str, Any]], - total_new: int, - finished: bool, - max_model_len: int, - error: Optional[str] = None, -) -> str: - """Persist state for post-mortem (student CoT vs checker injection). Returns path.""" - os.makedirs(OVERFLOW_DUMP_DIR, exist_ok=True) - tag = f'idx{debug_idx}' if debug_idx is not None else 'idx_unknown' - path = os.path.join( - OVERFLOW_DUMP_DIR, f'{tag}_{reason}_{int(time.time())}.json') - - partial_cot = _decode(tokenizer, gen_ids) if tokenizer and gen_ids else '' - injected_text = (_decode(tokenizer, injected_ids) - if tokenizer and injected_ids else '') - ctx_len = _input_ids_len(cur_inputs) - - payload: Dict[str, Any] = { - 'reason': reason, - 'error': error, - 'query': problem, - 'debug_idx': debug_idx, - 'gen_token_count': len(gen_ids), - 'injected_token_count': len(injected_ids), - 'context_input_ids_len': ctx_len, - 'max_model_len': max_model_len, - 'chunk_tokens': chunk_tokens, - 'total_new': total_new, - 'n_checks': n_checks, - 'n_injections': n_injections, - 'findings': findings, - 'finished': finished, - 'partial_cot': partial_cot, - 'injected_text': injected_text, - 'partial_cot_chars': len(partial_cot), - 'context_is_message_prompt': ctx_len is None, - } - with open(path, 'w', encoding='utf-8') as f: - json.dump(payload, f, ensure_ascii=False, indent=2) - cot_path = path.replace('.json', '_partial_cot.txt') - with open(cot_path, 'w', encoding='utf-8') as f: - f.write(partial_cot) - sys.stderr.write(f'[dualline] overflow dump -> {path}\n') - return path - -# --------------------------------------------------------------------------- -# Teacher checker (Phase-0: pure teacher via llm_backup) -# --------------------------------------------------------------------------- -def _build_checker(): - """RubricVerifier with no student sampler -> every diagnose() hits the teacher. - - Uses a fixed, math-oriented process rubric so we do not spend a rubric- - generation call per slice (the segment here is a partial CoT, not a finished - trajectory). Falls back to auto-generated rubrics if fixed_rubric is cleared. - """ - from twinkle_agentic.verifier import RubricVerifier - from twinkle_agentic.verifier.rubric_verifier import RubricItem - - fixed = [ - RubricItem('The reasoning contains no arithmetic or algebraic error so far', - is_hard=True), - RubricItem('Each step follows logically from the previous ones', is_hard=True), - RubricItem('No formula or theorem is misstated or misapplied', is_hard=True), - RubricItem('The approach is on track to answer the actual question asked', - is_hard=False), - RubricItem('No step contradicts an earlier established fact', is_hard=False), - ] - return RubricVerifier(fixed_rubric=fixed, gate=True) - - -def _checker_available() -> bool: - return bool(os.environ.get('LLM_BACKUP_API_KEY') - or os.environ.get('LLM_BACKUP_BASE_URL') - or os.environ.get('OPENAI_API_KEY')) - - -def _diagnose_partial(checker, problem: str, partial_cot: str): - """Run the teacher checker on the reasoning so far; return (issue_or_None, detail). - - ``partial_cot`` is the FULL reasoning generated so far (all prior chunks plus - any self-corrections already spliced in), not just the latest slice, so the - teacher judges the whole derivation in context. We label it as in-progress so - it grades correctness of the steps rather than penalizing the absence of a - final answer. - """ - seg_content = ( - '[The following is the full reasoning so far, still in progress and not ' - 'yet complete. Judge only whether the reasoning up to this point is ' - 'mathematically correct; do not expect a final answer here.]\n\n' + partial_cot) - seg = {'messages': [ - {'role': 'user', 'content': problem}, - {'role': 'assistant', 'content': seg_content}, - ]} - try: - detail = checker.diagnose(seg, query=problem) - except Exception as exc: - logger.warning(f'[dualline] checker error: {exc}') - return None, None - if detail.overall_ok: - return None, detail - if detail.scalar >= CHECK_SCORE_FLOOR: - # Checker leans "mostly fine"; don't disrupt on a marginal signal. - return None, detail - fails = [it for it in detail.items if not it.verdict] - if not fails: - return None, detail - # Prefer a fix if the checker gave one; else the reason. - parts = [] - for it in fails[:2]: - msg = it.fix or it.reason - if msg: - parts.append(msg) - issue = ' '.join(parts).strip() or detail.summary - return (issue or None), detail - - -def _pad_batch_for_dp(items: List[Any], gen_dp: int) -> List[Any]: - """``slice_dp`` needs batch len >= DP world size (every rank gets work). - - Only kicks in on the tail rounds when fewer than ``gen_dp`` problems are - still active; the padded replicas are dropped by the caller. - """ - if gen_dp <= 1 or not items or len(items) >= gen_dp: - return items - pad = [copy.deepcopy(items[-1]) for _ in range(gen_dp - len(items))] - return items + pad - - -class _DualState: - """Per-problem generation state for the batched dualline loop. - - All problems advance together, one ``chunk_tokens`` slice per round. A - problem stays *active* until it emits EOS, hits ``max_gen_tokens``, would - overflow ``max_model_len``, or a sample call fails. Because the problems - share every round's ``sampler.sample`` call, the vLLM engine batches them - (and, with dp>1, spreads them across ranks) instead of running one at a - time. - """ - - __slots__ = ('idx', 'problem', 'cur_input', 'gen_ids', 'injected_ids', - 'n_checks', 'n_injections', 'findings', 'total_new', - 'finished', 'stopped_reason', 'context_input_ids_len', - 'pending_partial_cot', 'prompt_len') - - def __init__(self, idx: int, problem: str, prompt: Any): - self.idx = idx - self.problem = problem - self.cur_input: Any = prompt # str prompt (round 0) or input_feature - self.gen_ids: List[int] = [] # student-generated token ids only - self.injected_ids: List[int] = [] # spliced-in ids (excluded from answer) - self.n_checks = 0 - self.n_injections = 0 - self.findings: List[Dict[str, Any]] = [] - self.total_new = 0 - self.finished = False - self.stopped_reason: Optional[str] = None - self.context_input_ids_len: Optional[int] = None - self.pending_partial_cot: Optional[str] = None - self.prompt_len: Optional[int] = None # token len of the fixed prompt prefix - - def cur_input_len(self) -> Optional[int]: - item = self.cur_input - if isinstance(item, dict) and 'input_ids' in item: - ids = item['input_ids'] - return len(ids) if ids is not None else None - return None - - def result(self, tokenizer) -> Dict[str, Any]: - if self.context_input_ids_len is None: - self.context_input_ids_len = self.cur_input_len() - return { - 'text': _decode(tokenizer, self.gen_ids), - 'finished': self.finished, - 'stopped_reason': self.stopped_reason, - 'context_input_ids_len': self.context_input_ids_len, - 'n_checks': self.n_checks, - 'n_injections': self.n_injections, - 'findings': self.findings, - 'gen_tokens': len(self.gen_ids), - } - - -def _dump_state_obj(st: '_DualState', tokenizer, chunk_tokens: int, - max_model_len: int, reason: str, error: str) -> None: - _dump_dualline_state( - reason=reason, - problem=st.problem, - debug_idx=st.idx, - chunk_tokens=chunk_tokens, - cur_inputs=[st.cur_input], - gen_ids=st.gen_ids, - injected_ids=st.injected_ids, - tokenizer=tokenizer, - n_checks=st.n_checks, - n_injections=st.n_injections, - findings=st.findings, - total_new=st.total_new, - finished=st.finished, - max_model_len=max_model_len, - error=error, - ) - - -# --------------------------------------------------------------------------- -# Batched token-level segmented generation with mid-stream injection -# --------------------------------------------------------------------------- -def run_dualline_batch(sampler, tokenizer, problems: List[str], checker, - base_params: TwinkleSamplingParams, - chunk_tokens: int, - max_model_len: int, - max_gen_tokens: int, - gen_dp: int = 1, - diagnose_workers: int = 8) -> List[Dict[str, Any]]: - """Advance every problem in lock-step slices, sharing one sampler call/round. - - Each round: (1) preflight-drop any problem that would overflow the context, - (2) one ``sampler.sample`` over all still-active problems (vLLM batches + - spreads over dp ranks), (3) for the length-capped ones, run the teacher - diagnoses concurrently and splice injections, then loop. - - Returns per-problem result dicts in the original ``problems`` order. - """ - from concurrent.futures import ThreadPoolExecutor - - chunk_params = TwinkleSamplingParams( - max_tokens=chunk_tokens, temperature=base_params.temperature, - top_p=base_params.top_p, num_samples=1) - - states = [_DualState(i, p, build_direct_prompt(p)) - for i, p in enumerate(problems)] - active = list(states) - round_no = 0 - - while active: - round_no += 1 - - # (1) Preflight: drop problems that would overflow the context window, - # and those that already reached the generation-token cap. - survivors: List[_DualState] = [] - for st in active: - if st.total_new >= max_gen_tokens: - st.stopped_reason = st.stopped_reason or 'max_gen_tokens' - continue - ctx_len = st.cur_input_len() - if ctx_len is not None and ctx_len + chunk_tokens >= max_model_len: - st.context_input_ids_len = ctx_len - st.stopped_reason = 'context_full' - _dump_state_obj( - st, tokenizer, chunk_tokens, max_model_len, - reason='preflight_context_full', - error=(f'context len {ctx_len} + chunk {chunk_tokens} ' - f'>= max_model_len {max_model_len}')) - continue - survivors.append(st) - active = survivors - if not active: - break - - # (2) One shared sampler call over all active problems. On tail rounds - # with fewer active problems than dp ranks, pad to keep slice_dp happy - # and drop the padded responses. The context-overflow preflight above - # guarantees every input still fits, so a length-capped slice should - # never raise here; let any genuine engine error propagate instead of - # masking it as a whole-round failure. - batch_inputs = [st.cur_input for st in active] - padded = _pad_batch_for_dp(batch_inputs, gen_dp) - responses = sampler.sample(padded, chunk_params) - responses = responses[:len(active)] - - # (3) Consume each problem's slice; queue the ones needing a check. - to_diagnose: List[_DualState] = [] - next_active: List[_DualState] = [] - for st, resp in zip(active, responses): - seq = resp.sequences[0] if resp and resp.sequences else None - if seq is None: - st.stopped_reason = st.stopped_reason or 'empty_response' - continue - st.gen_ids.extend(seq.tokens) - st.total_new += len(seq.tokens) - st.cur_input = seq.new_input_feature - if st.prompt_len is None: - # Fixed prompt prefix = everything before this round's generation. - st.prompt_len = len(st.cur_input['input_ids']) - len(seq.tokens) - - if seq.stop_reason != 'length': - st.finished = True # EOS / stop -> done - continue - if st.n_checks >= MAX_CHECKS or not checker: - next_active.append(st) # keep generating, no more checks - continue - # Diagnose the FULL reasoning generated so far (all prior chunks plus - # any self-corrections already spliced in), so the teacher judges the - # whole derivation in context rather than an isolated tail slice. - st.pending_partial_cot = _decode( - tokenizer, st.cur_input['input_ids'][st.prompt_len:]) - st.n_checks += 1 - to_diagnose.append(st) - - # Concurrent teacher diagnoses for this round's length-capped problems. - if to_diagnose: - def _run(st: _DualState): - return st, _diagnose_partial( - checker, st.problem, st.pending_partial_cot) - workers = max(1, min(diagnose_workers, len(to_diagnose))) - with ThreadPoolExecutor(max_workers=workers) as ex: - for st, (issue, _detail) in ex.map(_run, to_diagnose): - st.pending_partial_cot = None - if issue and st.n_injections < MAX_INJECTIONS: - note = INJECT_TEMPLATE.format(issue=issue) - note_ids = tokenizer.encode(note, add_special_tokens=False) - feat = dict(st.cur_input) - feat['input_ids'] = list(feat['input_ids']) + note_ids - if 'labels' in feat: - feat['labels'] = list(feat['labels']) + note_ids - st.cur_input = feat - st.injected_ids.extend(note_ids) - st.n_injections += 1 - st.findings.append( - {'at_token': st.total_new, 'issue': issue}) - next_active.append(st) - - active = next_active - n_done = sum(1 for s in states if s.finished or s.stopped_reason) - sys.stderr.write( - f'[dualline] round {round_no}: active={len(active)} ' - f'done={n_done}/{len(states)}\n') - - return [st.result(tokenizer) for st in states] - - -def _load_tokenizer(model_id: str): - """Load the tokenizer from ModelScope (matches the vLLM sampler source). - - The box runs offline, so ``transformers.AutoTokenizer`` (which resolves via - the HF hub) fails with ``Network is unreachable``. ModelScope's AutoTokenizer - downloads/reads from the ModelScope cache instead — the same place the vLLM - sampler already pulled the model from. Falls back to transformers only if the - ModelScope path is unavailable. - """ - try: - from modelscope import AutoTokenizer as MSAutoTokenizer - return MSAutoTokenizer.from_pretrained(model_id, trust_remote_code=True) - except Exception as exc: - sys.stderr.write(f'[dualline] modelscope tokenizer load failed ({exc}); ' - f'falling back to transformers\n') - from transformers import AutoTokenizer - return AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) - - -# --------------------------------------------------------------------------- -# Main -# --------------------------------------------------------------------------- -def main(): - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--mode', choices=['baseline', 'dualline'], default='dualline') - p.add_argument('--dataset', choices=['aops', 'math'], default='aops', - help='Evaluation dataset. "aops" (default) auto-downloads from ' - 'ModelScope (no local path needed); "math" reads local ' - 'MATH_DATA_DIR, stratified by difficulty level.') - p.add_argument('--math-split', default='test') - p.add_argument('--per-level', type=int, default=0, - help='MATH only: problems per level. 0 => --n split across levels.') - p.add_argument('--n', type=int, default=32, - help='Pool size sampled from the dataset (MATH is stratified ' - 'by level; AoPS is a flat shuffle).') - p.add_argument('--target-eval', type=int, default=32, - help='Stop after this many problems are evaluated (0 = all sampled).') - p.add_argument('--max-model-len', type=int, default=DUALLINE_DEFAULT_MAX_MODEL_LEN, - help='vLLM max_model_len / template max_length (default 32000).') - p.add_argument('--max-gen-tokens', type=int, default=DUALLINE_DEFAULT_MAX_GEN_TOKENS, - help='Cap total generated tokens per problem (default: same as ' - 'max-model-len / DUALLINE_MAX_GEN_TOKENS).') - p.add_argument('--chunk-tokens', type=int, default=CHUNK_TOKENS, - help='Generate this many tokens between checker pauses.') - p.add_argument('--batch-size', type=int, default=16, - help='Baseline mode batch size. Dualline runs all problems ' - 'concurrently (one shared sampler call per slice-round).') - p.add_argument('--diagnose-workers', type=int, - default=int(os.environ.get('DUALLINE_DIAGNOSE_WORKERS', 8)), - help='Concurrency for teacher diagnose() calls within a round.') - 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/dualline/{args.dataset}_{args.mode}_results.jsonl' - - is_dual = (args.mode == 'dualline') - if is_dual and not _checker_available(): - sys.stderr.write( - '[dualline] ERROR: --mode dualline needs a teacher checker but no ' - 'LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / OPENAI_API_KEY is set.\n' - ' Set them, or run --mode baseline for the paired baseline.\n') - sys.exit(1) - - 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) - if args.target_eval > 0: - records = records[:args.target_eval] - max_model_len = args.max_model_len - max_gen_tokens = args.max_gen_tokens - sys.stderr.write( - f'[dualline] evaluating {len(records)} problems ' - f'(mode={args.mode}, dataset={args.dataset}, ' - f'max_model_len={max_model_len}, max_gen_tokens={max_gen_tokens})\n') - - device_groups = [ - DeviceGroup(name='sampler', ranks=list(range(GEN_GPUS)), - device_type='GPU', gpus_per_worker=GEN_TP), - ] - if GEN_GPUS % GEN_TP != 0: - raise ValueError(f'GEN_GPUS ({GEN_GPUS}) must be divisible by GEN_TP ({GEN_TP})') - gen_dp = GEN_GPUS // GEN_TP - gen_mesh = DeviceMesh.from_sizes(world_size=GEN_GPUS, dp_size=gen_dp, tp_size=GEN_TP) - 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': max_model_len, - 'tensor_parallel_size': GEN_TP, - }, - device_mesh=gen_mesh, - remote_group='sampler', - ) - sampler.set_template('Qwen3_5Template', model_id=GEN_MODEL_ID, - enable_thinking=True, max_length=max_model_len) - sys.stderr.write( - f'[dualline] vLLM sampler ready (model={GEN_MODEL_ID}, ' - f'tp={GEN_TP}, dp={gen_dp})\n') - - gen_params = TwinkleSamplingParams( - max_tokens=max_gen_tokens, temperature=GEN_TEMPERATURE, - top_p=GEN_TOP_P, num_samples=1) - - checker = None - tokenizer = None - if is_dual: - checker = _build_checker() - tokenizer = _load_tokenizer(GEN_MODEL_ID) - sys.stderr.write('[dualline] teacher checker ready (llm_backup teacher)\n') - - correct = 0 - total = 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') - - def _grade_and_log(rec, idx, raw_output, extra=None): - nonlocal correct, total - predicted = extract_boxed(raw_output) - is_correct = answers_match(predicted, rec['reference_answer']) - if is_correct: - correct += 1 - total += 1 - debug_rec = { - 'idx': idx, - 'reference_answer': rec['reference_answer'], - 'predicted': predicted, - 'is_correct': is_correct, - 'problem': rec['problem'], - 'model_output': raw_output, - } - if rec.get('level'): - debug_rec['level'] = rec['level'] - if rec.get('type'): - debug_rec['type'] = rec['type'] - if extra: - debug_rec.update(extra) - debug_records.append(debug_rec) - out_f.write(json.dumps(debug_rec, ensure_ascii=False) + '\n') - out_f.flush() - - if is_dual: - problems = [rec['problem'] for rec in records] - results = run_dualline_batch( - sampler, tokenizer, problems, checker, gen_params, - args.chunk_tokens, max_model_len, max_gen_tokens, - gen_dp=gen_dp, diagnose_workers=args.diagnose_workers) - for idx, (rec, result) in enumerate(zip(records, results)): - _grade_and_log(rec, idx, result['text'], extra={ - 'n_checks': result['n_checks'], - 'n_injections': result['n_injections'], - 'findings': result['findings'], - 'finished': result['finished'], - 'stopped_reason': result.get('stopped_reason'), - 'context_input_ids_len': result.get('context_input_ids_len'), - 'gen_tokens': result['gen_tokens'], - }) - stop_tag = (f' stop={result["stopped_reason"]}' - if result.get('stopped_reason') else '') - sys.stderr.write( - f' [idx {idx}] correct={debug_records[-1]["is_correct"]} ' - f'gen={result["gen_tokens"]} checks={result["n_checks"]} ' - f'inj={result["n_injections"]}{stop_tag}\n') - acc = correct / total if total else 0 - sys.stderr.write( - f'[dualline] batched eval done: acc={acc:.4f} ({correct}/{total})\n') - else: - import re - for batch_start in range(0, len(records), args.batch_size): - batch = records[batch_start:batch_start + args.batch_size] - prompts = [build_direct_prompt(r['problem']) for r in batch] - if gen_dp > 1 and len(prompts) < gen_dp: - prompts = _pad_batch_for_dp(prompts, gen_dp) - pad_n = len(prompts) - len(batch) - else: - pad_n = 0 - responses = sampler.sample(prompts, gen_params) - if pad_n: - responses = responses[:len(batch)] - 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 = re.sub(r'<\|[^|]+\|>', '', seq.decoded or '').rstrip() - _grade_and_log(rec, batch_start + i, raw_output) - acc = correct / total if total else 0 - sys.stderr.write(f' [{total}/{len(records)}] acc={acc:.4f} ' - f'({correct}/{total})\n') - - overall = correct / total if total else 0 - print(f'\n{"=" * 60}') - print(f'MATH dual-line — mode={args.mode}, model={GEN_MODEL_ID}') - print(f' n={total}, seed={args.seed}, chunk_tokens={args.chunk_tokens}, ' - f'max_model_len={max_model_len}') - print(f'{"=" * 60}') - print(f'Overall accuracy: {overall:.4f} ({correct}/{total})') - - if is_dual: - tot_checks = sum(r.get('n_checks', 0) for r in debug_records) - tot_inj = sum(r.get('n_injections', 0) for r in debug_records) - n_with_inj = sum(1 for r in debug_records if r.get('n_injections', 0) > 0) - print(f' checker: {tot_checks} checks, {tot_inj} injections across ' - f'{n_with_inj}/{total} problems') - n_ctx_full = sum( - 1 for r in debug_records if r.get('stopped_reason') == 'context_full') - n_sample_fail = sum( - 1 for r in debug_records if r.get('stopped_reason') == 'sample_failed') - n_unfinished = sum(1 for r in debug_records if not r.get('finished')) - print(f' length: context_full={n_ctx_full}/{total}, ' - f'sample_failed={n_sample_fail}/{total}, ' - f'unfinished(no EOS)={n_unfinished}/{total}') - - if any(r.get('level') for r in debug_records): - per = defaultdict(lambda: [0, 0]) - for r in debug_records: - lv = r.get('level', 'Unknown') - per[lv][1] += 1 - if r['is_correct']: - per[lv][0] += 1 - print('\nPer-level accuracy:') - for lv in sorted(per.keys()): - c, t = per[lv] - print(f' {lv:>10}: {c/t:.4f} ({c}/{t})') - - out_f.close() - print(f'\n[output] {len(debug_records)} records saved to {args.output}') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/embedding/eval_gpqa_rag.py b/cookbook/exp/embedding/eval_gpqa_rag.py deleted file mode 100644 index 7954e5fd2..000000000 --- a/cookbook/exp/embedding/eval_gpqa_rag.py +++ /dev/null @@ -1,1547 +0,0 @@ -"""Math evaluation: direct vs RAG-augmented with Qwen3.5-4B. - -Datasets (``--dataset``): - - ``math`` (default): MATH (Hendrycks), stratified by difficulty (Level 1-5) - so RAG gain can be plotted against difficulty. - - ``aops``: AoPS competition problems (metadata.boxed only). - -Modes (``--mode``): - - ``direct``: The model solves problems directly (4 GPUs, TP=4). - - ``rag`` (default): Retrieve top-k thinking traces from LanceDB, condense - them (API qwen3.7-max), inject as 1-shot examples, then solve - (8 GPUs: DP=4 embedding + TP=4 vLLM). - -Defaults implement **raw RAG on MATH**: ``--dataset math --mode rag --condense`` -with hint filtering OFF. The API condenser needs ``COMPRESS_API_KEY`` (or a -local condenser via ``EVAL_CONDENSER_GPUS``); otherwise pass ``--no-condense``. - -Optional ``--hint`` flag (rag mode only): - After retrieval + condensing, call an API model to filter and refine the - traces — keeping only applicable methods — then inject the refined trace. - -Reference answers are the ``\\boxed{...}`` content of each solution. - -Launch examples: - # Default: raw RAG on MATH, stratified 100/level (needs COMPRESS_API_KEY) - COMPRESS_API_KEY=sk-... python cookbook/exp/embedding/eval_gpqa_rag.py - - # Paired direct baseline on the same MATH subset - python cookbook/exp/embedding/eval_gpqa_rag.py --mode direct - - # Raw RAG without condenser (inject raw retrieved traces) - python cookbook/exp/embedding/eval_gpqa_rag.py --no-condense - - # Add hint filtering back on top of condensing - COMPRESS_API_KEY=sk-... python cookbook/exp/embedding/eval_gpqa_rag.py --hint - - # Fall back to the old AoPS dataset - python cookbook/exp/embedding/eval_gpqa_rag.py --dataset aops -""" -import argparse -import json -import os -import random -import re -import sys -import threading -import time -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Any, Dict, List, Optional - -import numpy as np -import torch - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.data_format import SamplingParams as TwinkleSamplingParams -from twinkle.loss import InfonceLoss -from twinkle.model import TransformersModel -from twinkle.processor import InputProcessor -from twinkle.sampler import vLLMSampler -from twinkle.template import Qwen3_5Template -from twinkle_agentic.protocol.openai import OpenAI as OpenAIClient - -logger = get_logger() - -# -- Condenser config ---------------------------------------------------------- -CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') -CONDENSE_API_KEY = os.environ.get('COMPRESS_API_KEY', '') -CONDENSE_BASE_URL = os.environ.get('COMPRESS_BASE_URL', 'https://dashscope.aliyuncs.com/compatible-mode/v1') -CONDENSE_API_MODEL = os.environ.get('COMPRESS_MODEL', 'qwen3.7-max') -CONDENSE_API_CONCURRENCY = int(os.environ.get('API_CONCURRENCY', 32)) -CONDENSE_API_MIN_INTERVAL = float(os.environ.get('API_MIN_INTERVAL', 0.1)) -CONDENSE_TEMPERATURE = 0.2 -CONDENSE_MAX_TOKENS = 8192 - -# -- Hint analysis config ------------------------------------------------------ -HINT_ANALYSIS_MAX_TOKENS = int(os.environ.get('HINT_ANALYSIS_MAX_TOKENS', 2000)) -HINT_ANALYSIS_TEMPERATURE = 0.2 - -HINT_ANALYSIS_SYSTEM = ( - 'You are a mathematical reasoning trace filter. ' - 'Given a target problem and reasoning traces retrieved from SIMILAR (but different) problems, ' - 'your task is to FILTER and REFINE the traces into a clean reference.\n\n' - 'Rules:\n' - '1. KEEP: solution steps, methods, formulas, techniques, and key insights ' - 'that are directly applicable to solving the target problem.\n' - '2. REMOVE: problem-specific numeric calculations that do not transfer, ' - 'dead-end explorations, irrelevant approaches, verbose restatements, ' - 'and any content that would mislead the solver on the target problem.\n' - '3. Output the refined trace directly as actionable solution steps. ' - 'Preserve the original mathematical expressions and step structure.\n' - '4. Do NOT solve the target problem. Do NOT add your own solutions or commentary.\n' - '5. Do NOT output the answer to either problem.\n' - '6. If the traces are entirely irrelevant, output exactly: "No applicable methods."' -) - -HINT_ANALYSIS_USER = ( - '## Target Problem\n{query}\n\n' - '## Retrieved Reasoning Traces\n{thinking}' -) - -# --------------------------------------------------------------------------- -# Config -# --------------------------------------------------------------------------- -# -- Gen/Embed config --------------------------------------------------------- -GEN_MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3.5-4B') -EMBED_MODEL_ID = os.environ.get( - 'EMBED_MODEL_ID', 'output.oldemb/embedding_full_transformers/last-checkpoint') - -GEN_GPUS = int(os.environ.get('GEN_GPUS', 8)) -EMB_GPUS = int(os.environ.get('EMB_GPUS', 2)) -EMBED_MAX_LENGTH = int(os.environ.get('EMBED_MAX_LENGTH', 20000)) - -GEN_GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.85)) -GEN_MAX_MODEL_LEN = int(os.environ.get('GEN_MAX_MODEL_LEN', 65536)) -GEN_MAX_TOKENS = int(os.environ.get('GEN_MAX_TOKENS', 65536)) -GEN_TEMPERATURE = float(os.environ.get('GEN_TEMPERATURE', 0.6)) -GEN_TOP_P = float(os.environ.get('GEN_TOP_P', 0.95)) - -AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') -MATH_DATA_DIR = os.environ.get('MATH_DATA_DIR', './output/math_data/MATH') - - -# --------------------------------------------------------------------------- -# Condenser prompts & validation -# --------------------------------------------------------------------------- - -COMPRESS_SYSTEM = """\ -You are a reasoning-trace condenser. Given a verbose reasoning trace, \ -extract the TRANSFERABLE KNOWLEDGE as an EXECUTABLE SOLUTION SKELETON \ -that would help a reader solve SIMILAR problems in the same domain. - -Your output is the ENTIRE useful content — there is no expansion tool, no second pass. \ -The reader will apply this knowledge to a DIFFERENT problem, so focus on what transfers. - -Principles: -1. OUTPUT AN EXECUTABLE STEP CHAIN: numbered steps that a solver can directly follow. \ -Each step should state WHAT to do and HOW (with the formula/technique), not just \ -name the concept. -2. INCLUDE FULL FORMULAS: theorems, identities, inequalities — state each \ -with its COMPLETE MATHEMATICAL EXPRESSION, not just the name. -3. STATE APPLICABILITY: what structural features of a problem signal that this \ -approach works (e.g. "when the constraint is a sum of squares"). -4. PRESERVE KEY INSIGHTS: the non-obvious ideas or tricks that make the approach \ -work — the things a solver would NOT think of without guidance. -5. REMOVE: problem-specific numeric calculations, dead-end explorations, \ -hesitations, verbose restatements, and trivial arithmetic. -6. FORMAT: Start with a one-line "Applicability" statement, then numbered steps, \ -then key formulas. Keep it concise and actionable. -7. NO meta-commentary about the compression process. NO preamble. -""" - -COMPRESS_USER = ( - '## Reader Problem (context only — do NOT solve it)\n{query}\n\n' - '## Reasoning Trace to Condense\n{text}') - - -def _is_truncated_compression(text: str) -> bool: - if not text or not text.strip(): - return True - lines = [l.strip() for l in text.strip().splitlines() if l.strip()] - if len(lines) < 3: - return True - last_line = lines[-1] - # Truncated if last line looks incomplete (no terminal punctuation/formula) - if last_line and last_line[-1] not in '.。!！)）]】}\\$': - # Allow lines ending with numbers, boxed answers, etc. - if not re.search(r'\d$|\\boxed|\$|\)$', last_line): - return True - return False - - -# -- API rate limiter ---------------------------------------------------------- -_api_semaphore = threading.Semaphore(CONDENSE_API_CONCURRENCY) -_api_bucket_lock = threading.Lock() -_api_tokens = [float(CONDENSE_API_CONCURRENCY)] -_api_last_refill = [time.monotonic()] - - -def _api_throttle(): - _api_semaphore.acquire() - wait = 0.0 - try: - with _api_bucket_lock: - now = time.monotonic() - elapsed = now - _api_last_refill[0] - refill = elapsed / CONDENSE_API_MIN_INTERVAL - _api_tokens[0] = min(float(CONDENSE_API_CONCURRENCY), _api_tokens[0] + refill) - _api_last_refill[0] = now - if _api_tokens[0] >= 1.0: - _api_tokens[0] -= 1.0 - else: - wait = (1.0 - _api_tokens[0]) * CONDENSE_API_MIN_INTERVAL - _api_tokens[0] = 0.0 - finally: - _api_semaphore.release() - if wait > 0: - time.sleep(wait) - - -def _api_condense_single(api_client: OpenAIClient, messages: List[Dict]) -> Optional[str]: - _api_throttle() - trajectory = {'messages': messages} - sp = TwinkleSamplingParams(temperature=CONDENSE_TEMPERATURE, max_tokens=CONDENSE_MAX_TOKENS) - try: - reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) - except Exception as exc: - logger.warning(f'[condense-api] error: {exc}') - return None - content = (reply.get('content') or '').strip() - if not content: - return None - m = re.match(r'^```[a-zA-Z]*\n(.*?)\n```\s*$', content, re.DOTALL) - if m: - content = m.group(1).strip() - return content - - -def _api_hint_analysis_batch( - api_client: OpenAIClient, - problems: List[str], - condensed_examples: List[List[Dict[str, str]]], -) -> List[Optional[str]]: - """Call API to pre-analyze RAG relevance for each problem.""" - _MAX_HINT_INPUT = 8000 - results: List[Optional[str]] = [None] * len(problems) - tasks = [] - for i, prob in enumerate(problems): - if not condensed_examples[i]: - continue - traces = [ex.get('thinking', '') for ex in condensed_examples[i]] - merged_thinking = '\n---\n'.join(traces) - if len(merged_thinking) > _MAX_HINT_INPUT: - merged_thinking = merged_thinking[:_MAX_HINT_INPUT] + '\n[...truncated]' - user_msg = HINT_ANALYSIS_USER.format(query=prob, thinking=merged_thinking) - msgs = [ - {'role': 'system', 'content': HINT_ANALYSIS_SYSTEM}, - {'role': 'user', 'content': user_msg}, - ] - tasks.append((i, msgs)) - - if not tasks: - return results - - def _call_one(idx, msgs): - _api_throttle() - try: - trajectory = {'messages': msgs} - sp = TwinkleSamplingParams( - temperature=HINT_ANALYSIS_TEMPERATURE, - max_tokens=HINT_ANALYSIS_MAX_TOKENS) - reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) - content = (reply.get('content') or '').strip() - # Treat "No applicable methods." as empty (will trigger fallback) - if not content or content == 'No applicable methods.': - return idx, None - return idx, content - except Exception as exc: - logger.warning(f'[hint-analysis] error for idx={idx}: {exc}') - return idx, None - - with ThreadPoolExecutor(max_workers=min(len(tasks), CONDENSE_API_CONCURRENCY)) as pool: - futs = [pool.submit(_call_one, idx, msgs) for idx, msgs in tasks] - for fut in as_completed(futs): - idx, analysis = fut.result() - results[idx] = analysis - - n_success = sum(1 for r in results if r) - logger.info(f'[hint-analysis] completed {n_success}/{len(tasks)} analyses') - return results - - -# --------------------------------------------------------------------------- -# LLM-based decontamination -# --------------------------------------------------------------------------- - -_DECONTAM_JUDGE_PROMPT = ( - 'We are building a RAG-augmented math training system. Problem A is the test ' - 'question; Problem B was retrieved from a knowledge base.\n' - 'Answer YES only if A and B are essentially the SAME specific problem — ' - 'i.e. solving B directly gives you A\'s answer (just different wording/notation/' - 'format/negation).\n' - 'Answer NO if they merely share the same method/topic but have different ' - 'specific values, equations, or geometric configurations — learning B\'s ' - 'approach still requires independent work to solve A.\n' - 'Problem A: {prob_a}\n' - 'Problem B: {prob_b}\n' - 'Answer only YES or NO.' -) - - -def _llm_judge_same_problem( - api_client: OpenAIClient, pairs: List[tuple], -) -> List[bool]: - """Batch LLM judge: are (problem_a, problem_b) the same problem? - - Returns list of bools (True = same problem = should filter). - """ - if not pairs or not api_client: - return [False] * len(pairs) - - results = [False] * len(pairs) - - def _judge_one(idx, pa, pb): - prompt = _DECONTAM_JUDGE_PROMPT.format(prob_a=pa, prob_b=pb) - msgs = [{'role': 'user', 'content': prompt}] - _api_throttle() - try: - trajectory = {'messages': msgs} - sp = TwinkleSamplingParams(temperature=0.1, max_tokens=8) - reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) - answer = (reply.get('content') or '').strip().upper() - return idx, 'YES' in answer - except Exception: - return idx, False - - with ThreadPoolExecutor(max_workers=min(len(pairs), CONDENSE_API_CONCURRENCY)) as pool: - futs = [pool.submit(_judge_one, i, pa, pb) for i, (pa, pb) in enumerate(pairs)] - for fut in as_completed(futs): - idx, is_same = fut.result() - results[idx] = is_same - return results - - -def _llm_decontaminate( - api_client: OpenAIClient, - problems: List[str], - all_examples: List[List[Dict[str, str]]], -) -> List[List[Dict[str, str]]]: - """Apply LLM-based decontamination: remove retrievals judged as same problem.""" - judge_pairs = [] # (qi, ret_idx, prob_a, prob_b) - for qi, exs in enumerate(all_examples): - for ri, ex in enumerate(exs): - judge_pairs.append((qi, ri, problems[qi], ex.get('query', ''))) - - if not judge_pairs: - return all_examples - - pairs_input = [(pa, pb) for _, _, pa, pb in judge_pairs] - verdicts = _llm_judge_same_problem(api_client, pairs_input) - to_remove = set() - for vi, (qi, ri, _, _) in enumerate(judge_pairs): - if verdicts[vi]: - to_remove.add((qi, ri)) - - if to_remove: - logger.info(f'[decontam-llm] filtered {len(to_remove)} same-problem retrievals') - for qi in range(len(all_examples)): - all_examples[qi] = [ - ex for ri, ex in enumerate(all_examples[qi]) - if (qi, ri) not in to_remove - ] - return all_examples - - -def condense_traces( - examples_batch: List[List[Dict[str, str]]], - problems: List[str], - api_client: OpenAIClient, - condenser_sampler=None, - compress_params=None, - special_tokens: set = None, - max_output_len: int = 2000, - dp_size: int = 1, -) -> List[List[Dict[str, str]]]: - """Compress retrieved thinking traces with query-aware condenser. - - Primary: local vLLM condenser (if provided). - Fallback: API condenser. - Final fallback: raw trace truncated to max_output_len. - """ - result: List[List[Dict[str, str]]] = [] - # Flatten all (batch_idx, ex_idx, problem, example) for batch processing - tasks = [] - for bi, (exs, prob) in enumerate(zip(examples_batch, problems)): - for ei, ex in enumerate(exs): - tasks.append((bi, ei, prob, ex)) - - if not tasks: - return [[] for _ in examples_batch] - - # Build condense prompts (aligned with make_embedding_dataset.py hard path) - prompts = [] - for _, _, prob, ex in tasks: - user_msg = COMPRESS_USER.format(query=prob, text=ex['thinking']) - prompts.append([{'role': 'system', 'content': COMPRESS_SYSTEM}, - {'role': 'user', 'content': user_msg}]) - - # Phase 1: local vLLM condenser - condensed = [None] * len(tasks) - condense_sources = ['raw'] * len(tasks) - fallback_indices = [] - - if condenser_sampler is not None and compress_params is not None: - sampler_inputs = [{'messages': p} for p in prompts] - # The local vLLM sampler runs data-parallel across ``dp_size`` workers - # and requires at least one item per worker (it errors with - # "Batch too small for N workers" otherwise). Pad the batch up to a - # multiple of dp_size by repeating the last item, run, then keep only - # the first ``n_real`` responses and drop the padding. - n_real = len(sampler_inputs) - pad_size = 0 - if dp_size > 1 and n_real > 0 and n_real % dp_size != 0: - pad_size = dp_size - (n_real % dp_size) - sampler_inputs = sampler_inputs + [sampler_inputs[-1]] * pad_size - try: - responses = condenser_sampler.sample(sampler_inputs, compress_params) - except Exception as exc: - logger.warning(f'[condense] sampler error: {exc}') - responses = [None] * len(sampler_inputs) - if pad_size: - responses = responses[:n_real] - for ri, resp in enumerate(responses): - seq = resp.sequences[0] if resp and resp.sequences else None - text = '' - if seq and seq.stop_reason != 'length' and seq.decoded: - text = seq.decoded - if special_tokens: - for tok in special_tokens: - text = text.replace(tok, '') - text = text.rstrip() - if text and not _is_truncated_compression(text): - condensed[ri] = text - condense_sources[ri] = 'local' - else: - fallback_indices.append(ri) - else: - fallback_indices = list(range(len(tasks))) - - # Phase 2: API fallback - if fallback_indices and api_client: - with ThreadPoolExecutor(max_workers=CONDENSE_API_CONCURRENCY) as pool: - futures = {} - for ri in fallback_indices: - futures[pool.submit(_api_condense_single, api_client, prompts[ri])] = ri - for fut in as_completed(futures): - ri = futures[fut] - api_result = fut.result() - if api_result and not _is_truncated_compression(api_result): - condensed[ri] = api_result - condense_sources[ri] = 'api' - - # Phase 3: assemble results (fallback to raw truncation) - result = [[] for _ in examples_batch] - for ti, (bi, ei, prob, ex) in enumerate(tasks): - compressed = condensed[ti] - raw_len = len(ex['thinking']) - sim_val = ex.get('_sim', 0.0) - if compressed: - result[bi].append({'query': ex['query'], - 'thinking': _strip_condenser_markers(compressed), - '_condense_source': condense_sources[ti], - '_raw_trace_len': raw_len, '_sim': sim_val}) - else: - result[bi].append({'query': ex['query'], - 'thinking': ex['thinking'][:max_output_len], - '_condense_source': 'raw', - '_raw_trace_len': raw_len, '_sim': sim_val}) - - n_ok = sum(1 for c in condensed if c) - logger.info(f'[condense] {n_ok}/{len(tasks)} compressed ok, ' - f'{len(tasks) - n_ok} fell back to raw truncation') - return result - - -def _strip_condenser_markers(text: str) -> str: - """Light cleanup of condenser output. - - Removes any residual markdown headers or meta-lines that don't carry - solution content. Keeps numbered steps and equations intact. - """ - # Remove legacy ## headers if condenser still emits them - if '## More' in text: - text = text.split('## More', 1)[0] - text = re.sub(r'^##\s*Summary\s*\n?', '', text, flags=re.MULTILINE) - text = re.sub(r'^Topic:\s*.*\n?', '', text, flags=re.MULTILINE) - # Remove meta-commentary lines - text = re.sub(r'^\s*\(Note:.*\)\s*$', '', text, flags=re.MULTILINE) - return text.strip() - - -# --------------------------------------------------------------------------- -# Boxed answer extraction -# --------------------------------------------------------------------------- -_BOXED_RE = re.compile(r'\\boxed\s*\{') - - -def extract_boxed(text: str) -> Optional[str]: - """Extract the last \\boxed{...} content, handling nested braces.""" - if not text: - return None - last_match = None - for m in _BOXED_RE.finditer(text): - start = m.end() - depth = 1 - i = start - while i < len(text) and depth > 0: - if text[i] == '{': - depth += 1 - elif text[i] == '}': - depth -= 1 - i += 1 - if depth == 0: - last_match = text[start:i - 1].strip() - return last_match - - -def normalize_answer(ans: str) -> str: - """Normalize a math answer string for comparison.""" - if not ans: - return '' - s = str(ans).strip() - m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) - if m: - return m.group(1) - s = s.strip('$').strip() - s = s.replace('−', '-') - 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.replace(r'\,', '').replace(r'\;', '').replace(r'\!', '') - s = re.sub(r'\\(?:quad|qquad|\s)', '', s) - s = re.sub(r'\s+', '', s) - s = re.sub(r'\\sqrt([0-9A-Za-z])', r'\\sqrt{\1}', s) - s = re.sub(r'\\frac\{([^{}]+)\}([^{}\\])', r'\\frac{\1}{\2}', s) - s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) - s = re.sub(r'\{[a-zA-Z]+\}$', '', s) - s = re.sub(r'\^\\circ|\^\{\\circ\}|°|\\circ', 'deg', 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'(? 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 - frac_re = re.compile(r'^\(([^)]+)\)/\(([^)]+)\)$') - - def _eval_frac(s): - m = frac_re.match(s) - if m: - try: - return float(m.group(1)) / float(m.group(2)) - except (ValueError, ZeroDivisionError): - pass - return None - - va, vb = _eval_frac(a), _eval_frac(b) - if va is not None and vb is not None: - return abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) - return False - - -_MCQ_REF_RE = re.compile( - r'^\\(?:textbf|text|mathrm|mathbf)\{\(?([A-E])\)?\s*\}\s*(.+)$' - r'|^\(?([A-E])\)\s+(.+)$' -) -_VAR_PREFIX_RE = re.compile(r'^(?:[A-Za-z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)$', re.DOTALL) - - -def _split_mcq(ans: str): - """Split an MCQ answer into (letter, value) components.""" - 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) - 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 _strip_var_prefix(ans: str) -> str: - m = _VAR_PREFIX_RE.match((ans or '').strip()) - return m.group(1).strip() if m else (ans or '') - - -def _normalize_tuple(ans: str) -> str: - return re.sub(r'[\s()\[\]{}\\]', '', ans or '') - - -def _math_verify_equal(predicted: str, reference: str) -> bool: - try: - from math_verify import parse, verify - gold = parse(r'\boxed{' + reference + '}', parsing_timeout=1) - pred = parse(r'\boxed{' + predicted + '}', parsing_timeout=1) - return bool(gold and pred and verify(gold, pred, timeout_seconds=1)) - except Exception: - 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 or norm_p.lower() == norm_r.lower(): - return True - if _try_numeric_equal(norm_p, norm_r): - return True - - stripped_p = normalize_answer(_strip_var_prefix(predicted)) - stripped_r = normalize_answer(_strip_var_prefix(reference)) - if stripped_p and stripped_r: - if stripped_p == stripped_r or stripped_p.lower() == stripped_r.lower(): - return True - if _try_numeric_equal(stripped_p, stripped_r): - return True - - p_letter, p_value = _split_mcq(predicted) - r_letter, r_value = _split_mcq(reference) - if p_letter and r_letter and p_letter == r_letter: - return True - if (p_letter and p_value is None) and (r_letter and r_value is None): - return False - - 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 p_val.lower() == r_val.lower() or _try_numeric_equal(p_val, r_val): - return True - - for left, right in ((norm_p, norm_r), (stripped_p, stripped_r)): - tuple_l, tuple_r = _normalize_tuple(left), _normalize_tuple(right) - if ',' in tuple_l and tuple_l == tuple_r: - return True - - if '=' in norm_r and '=' not in norm_p: - if any(part and (part == norm_p or _try_numeric_equal(part, norm_p)) for part in norm_r.split('=')): - return True - if '=' in norm_p and '=' not in norm_r: - if any(part and (part == norm_r or _try_numeric_equal(part, norm_r)) for part in norm_p.split('=')): - return True - - return _math_verify_equal(stripped_p or norm_p, stripped_r or norm_r) - - -# --------------------------------------------------------------------------- -# Dataset loading -# --------------------------------------------------------------------------- - -def _load_aops_from_modelscope(): - """Download the AoPS repo natively from ModelScope and read its parquet. - - Primary loader: ``dataset_snapshot_download`` pulls the dataset repo files - (parquet) straight from the ModelScope hub WITHOUT going through the - ``datasets``/HF-filesystem path used by ``MsDataset.load`` — that path is - broken on this modelscope build (``HfFileSystem.find() got multiple values - for 'maxdepth'``). We then read the local parquet with the ``datasets`` - backend (reading local files does not trigger the HfFileSystem bug). - """ - import glob - - from datasets import Dataset as HFDataset - from modelscope.hub.snapshot_download import dataset_snapshot_download - - local = dataset_snapshot_download(AOPS_DATASET_ID) - files = sorted(glob.glob(os.path.join(local, '**', '*.parquet'), - recursive=True)) - if not files: - # Older snapshots may materialize an arrow file instead of parquet. - files = sorted(glob.glob(os.path.join(local, '**', '*train*.arrow'), - recursive=True)) - if files: - sys.stderr.write(f'[aops] modelscope snapshot arrow: {files[0]}\n') - return HFDataset.from_file(files[0]) - return None - sys.stderr.write(f'[aops] modelscope snapshot parquet: {files[0]}\n') - return HFDataset.from_parquet(files if len(files) > 1 else files[0]) - - -def load_aops(n: int, seed: int = 42) -> List[Dict[str, Any]]: - """Load AoPS boxed problems, sample n, extract reference answers. - - Uses ModelScope as the data source (native repo snapshot download). - """ - ds = None - try: - ds = _load_aops_from_modelscope() - except Exception as exc: - sys.stderr.write(f'[aops] modelscope snapshot download failed ({exc}); ' - f'trying MsDataset.load\n') - if ds is None: - from modelscope import MsDataset - ds = MsDataset.load(AOPS_DATASET_ID, split='train', - download_mode='reuse_dataset_if_exists') - boxed = [] - for row in ds: - if not row['metadata'].get('boxed'): - continue - ref = extract_boxed(row['solution']) - if not ref: - continue - boxed.append({ - 'problem': row['problem'], - 'solution': row['solution'], - 'reference_answer': ref, - 'tags': row.get('tags', []), - }) - sys.stderr.write(f'[aops] {len(boxed)} boxed problems with extractable answers\n') - rng = random.Random(seed) - rng.shuffle(boxed) - if n > 0 and n < len(boxed): - boxed = boxed[:n] - sys.stderr.write(f'[aops] sampled {n} problems\n') - return boxed - - -def load_math(n: int, seed: int = 42, split: str = 'test', - per_level: int = 0) -> List[Dict[str, Any]]: - """Load the MATH (Hendrycks) dataset from local extracted JSON files. - - Each problem's reference answer is the ``\\boxed{}`` content of its - ``solution`` (MATH solutions always end in a boxed answer). - - Sampling is *stratified by level* so every difficulty (Level 1-5) is - represented equally — required to measure how RAG gain varies with - difficulty. ``per_level`` (if >0) fixes the count per level; otherwise - ``n`` is split evenly across the 5 levels. When both are 0, all problems - are returned. The final list is shuffled with ``seed`` so index order is - stable/comparable across direct vs rag runs. - """ - import glob - root = os.path.join(MATH_DATA_DIR, split) - files = glob.glob(os.path.join(root, '*', '*.json')) - if not files: - raise FileNotFoundError( - f'[math] no problems found under {root!r}; set MATH_DATA_DIR or ' - f'extract MATH.zip there') - - by_level: Dict[str, List[Dict[str, Any]]] = {} - n_no_box = 0 - for fp in files: - try: - with open(fp, 'r', encoding='utf-8') as fin: - row = json.load(fin) - except Exception: - continue - ref = extract_boxed(row.get('solution', '')) - if not ref: - n_no_box += 1 - continue - level = row.get('level', 'Unknown') - by_level.setdefault(level, []).append({ - 'problem': row['problem'], - 'solution': row['solution'], - 'reference_answer': ref, - 'level': level, - 'type': row.get('type', ''), - }) - - total = sum(len(v) for v in by_level.values()) - sys.stderr.write( - f'[math] {total} problems with boxed answers across ' - f'{len(by_level)} levels (skipped {n_no_box} without boxed)\n') - - levels = sorted(by_level.keys()) - rng = random.Random(seed) - - # Decide how many per level. - if per_level <= 0 and n > 0: - per_level = max(1, n // max(1, len(levels))) - - sampled: List[Dict[str, Any]] = [] - for lv in levels: - pool = by_level[lv] - rng.shuffle(pool) - take = pool if per_level <= 0 else pool[:per_level] - sampled.extend(take) - sys.stderr.write(f'[math] {lv}: took {len(take)}/{len(pool)}\n') - - rng.shuffle(sampled) - sys.stderr.write(f'[math] total sampled: {len(sampled)}\n') - return sampled - - -# --------------------------------------------------------------------------- -# Prompt building -# --------------------------------------------------------------------------- - -DIRECT_SYSTEM = ( - 'You are an expert competition mathematician. Solve the following problem ' - 'step by step. Provide your final answer inside \\boxed{}.' -) - -RAG_SYSTEM = ( - 'You are an expert competition mathematician. Solve the following problem ' - 'step by step. Provide your final answer inside \\boxed{}.\n\n' - 'You will first see example problem-solving traces or skills. ' - 'Learn from the reasoning methodology demonstrated in these examples, ' - 'then thinking to solve the actual problem.' -) - -RAG_FOLLOWUP = ( - 'The above is a reference solution to a similar problem. ' - 'You may use any applicable techniques from it, or ignore it ' - 'if you find a better approach. ' - 'Solve the problem step by step and put your final answer in \\boxed{}.' -) - -HINT_FOLLOWUP = ( - 'The above are applicable solution approaches extracted from similar problems. ' - 'You may use any applicable techniques from them, or ignore them ' - 'if you find a better approach. ' - 'Solve the problem step by step and put your final answer in \\boxed{}.' -) - -# Reminder appended to the final user turn. Without this, the reasoning model -# can loop indefinitely on multiple-choice problems, oscillating between boxing -# the option letter and boxing the value (e.g. "I'll box B. I'll box 21. ...") -# and never terminating. Boxing BOTH the letter and value removes the ambiguity -# (the grader accepts either), so the model has no format decision to agonize over. -MCQ_INSTRUCTION = ( - '\n\nNote: If the problem is multiple-choice (it lists options such as ' - '(A), (B), (C), ...), put BOTH the option letter and its value in the box, ' - 'e.g. \\boxed{(B) 21}. Otherwise, box the value directly. Decide the answer ' - 'format once and do not deliberate over which form to box.' -) - - -def build_direct_prompt(problem: str) -> Dict[str, Any]: - return { - 'messages': [ - {'role': 'system', 'content': DIRECT_SYSTEM}, - {'role': 'user', 'content': problem}, - ] - } - - -def build_hint_prompt(problem: str, hint_analysis: str) -> Dict[str, Any]: - """Build prompt with pre-analyzed hint in a multi-turn conversation. - - Mirrors ``build_rag_prompt``: the hint is presented as an assistant - "extracted approaches" turn (instead of being buried in the system - prompt), followed by a user instruction that provides a clear closing - directive to solve the problem and box the answer. Keeping the final - solve/box instruction in a dedicated user turn (rather than in the - system prompt) helps the reasoning model terminate cleanly. - """ - messages: List[Dict[str, str]] = [ - {'role': 'system', 'content': DIRECT_SYSTEM}, - {'role': 'user', 'content': problem}, - {'role': 'assistant', - 'content': ('Here are applicable solution approaches extracted from ' - f'similar problems:\n\n{hint_analysis}')}, - {'role': 'user', 'content': HINT_FOLLOWUP + MCQ_INSTRUCTION}, - ] - return {'messages': messages} - - -def build_rag_prompt(problem: str, - examples: List[Dict[str, str]]) -> Dict[str, Any]: - """Approach B: multi-turn assistant format. - - The trace is presented as an assistant "retrieval" turn, followed by - a user instruction that constrains the model to use methodology only. - """ - messages: List[Dict[str, str]] = [{'role': 'system', 'content': DIRECT_SYSTEM}] - messages.append({'role': 'user', 'content': problem}) - # Build trace content from retrieved examples - trace_parts = [] - for i, ex in enumerate(examples, 1): - trace_parts.append(f'[Retrieved Example {i}]\nProblem: {ex["query"]}\n' - f'Reasoning:\n{ex["thinking"]}') - trace_text = '\n\n'.join(trace_parts) - messages.append({'role': 'assistant', - 'content': f'I found relevant reasoning traces from the knowledge base!\n\n{trace_text}'}) - messages.append({'role': 'user', 'content': RAG_FOLLOWUP + MCQ_INSTRUCTION}) - return {'messages': messages} - - -# --------------------------------------------------------------------------- -# 13-gram Jaccard decontamination -# --------------------------------------------------------------------------- - -def _normalize_for_ngram(text: str) -> str: - """Normalize text for n-gram comparison: strip LaTeX markup, lowercase.""" - text = text.lower() - text = re.sub(r'\$+', '', text) - text = re.sub(r'\\[a-z]+\{([^}]*)\}', r'\1', text) - text = re.sub(r'\\[a-z]+', ' ', text) - text = re.sub(r'[{}\\^_$]', '', text) - text = re.sub(r'\s+', ' ', text).strip() - return text - - -def _ngram_jaccard(text_a: str, text_b: str, n: int = 13) -> float: - """13-gram character-level Jaccard similarity.""" - a = _normalize_for_ngram(text_a) - b = _normalize_for_ngram(text_b) - if len(a) < n or len(b) < n: - return 0.0 - grams_a = set(a[i:i + n] for i in range(len(a) - n + 1)) - grams_b = set(b[i:i + n] for i in range(len(b) - n + 1)) - if not grams_a or not grams_b: - return 0.0 - return len(grams_a & grams_b) / len(grams_a | grams_b) - - -# --------------------------------------------------------------------------- -# Embedding / RAG helpers -# --------------------------------------------------------------------------- - -def _wrap_anchor(text: str) -> List[Dict[str, str]]: - return [ - {'role': 'user', 'content': text}, - {'role': 'assistant', 'content': 'Match the correct response here.'}, - ] - - -def get_embeddings(model: TransformersModel, template: Qwen3_5Template, - texts: List[str], dp_size: int) -> np.ndarray: - if not texts: - return np.zeros((0,), dtype=np.float32) - n = len(texts) - pad_n = (-n) % dp_size - padded = list(texts) + [' '] * pad_n if pad_n else list(texts) - features = [] - for t in padded: - feat = template.encode({'messages': _wrap_anchor(t or ' ')}) - feat['labels'] = [1] - features.append(feat) - out = model.forward_only(inputs=features, task='embedding', return_logits=True) - emb = out['embeddings'] - if isinstance(emb, torch.Tensor): - emb = emb.detach().to(torch.float32).cpu().numpy() - emb = np.asarray(emb, dtype=np.float32) - return emb[:n] if pad_n else emb - - -def retrieve_examples(tbl, query_vecs: np.ndarray, top_k: int, - use_thinking_raw: bool, sim_threshold: float = 0.0, - problems: List[str] = None, - decontam_threshold: float = 0.0, - ) -> List[List[Dict[str, str]]]: - thinking_field = 'thinking_raw' if use_thinking_raw else 'cot_compressed' - fetch_limit = top_k + 50 if decontam_threshold > 0 else top_k - n_queries = len(query_vecs) - all_examples: List[List[Dict[str, str]]] = [None] * n_queries - decontam_skipped = 0 - _decontam_lock = threading.Lock() - - def _search_one(qi: int): - nonlocal decontam_skipped - vec = query_vecs[qi] - results = ( - tbl.search(vec.astype(np.float32).tolist()) - .metric('dot') - .limit(fetch_limit) - .select(['query_raw', thinking_field, '_distance']) - .to_list() - ) - problem_text = problems[qi] if problems else '' - examples = [] - local_skipped = 0 - for r in results: - if len(examples) >= top_k: - break - sim = 1.0 - r.get('_distance', 0.0) - if sim < sim_threshold: - continue - q = r.get('query_raw', '') - t = r.get(thinking_field, '') - if not t: - continue - if decontam_threshold > 0 and problem_text and q: - ng_sim = _ngram_jaccard(problem_text, q) - if ng_sim > decontam_threshold: - local_skipped += 1 - continue - examples.append({'query': q, 'thinking': t, '_sim': round(sim, 4), - '_raw_trace_len': len(t)}) - all_examples[qi] = examples - if local_skipped: - with _decontam_lock: - decontam_skipped += local_skipped - - with ThreadPoolExecutor(max_workers=min(n_queries, 16)) as pool: - list(pool.map(_search_one, range(n_queries))) - - if decontam_skipped > 0: - logger.info(f'[decontam] skipped {decontam_skipped} leaked retrievals ' - f'(13-gram Jaccard > {decontam_threshold})') - return all_examples - - -# --------------------------------------------------------------------------- -# Main -# --------------------------------------------------------------------------- - -def main(): - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--mode', choices=['direct', 'rag'], default='rag') - p.add_argument('--dataset', choices=['aops', 'math'], default='math', - help='Evaluation dataset. "math" = MATH (Hendrycks), ' - 'stratified by level for a difficulty-vs-gain curve.') - p.add_argument('--math-split', default='test', - help='MATH split to load (test/train).') - p.add_argument('--per-level', type=int, default=100, - help='MATH only: problems per difficulty level (default 100 ' - '-> 500 total across Level 1-5). If 0, --n is split ' - 'evenly across the 5 levels.') - p.add_argument('--n', type=int, default=0, - help='Pool size: sample this many problems (0 = all boxed). ' - 'In RAG mode with --target-eval, set this to 0 for max coverage.') - p.add_argument('--target-eval', type=int, default=0, - help='Stop after this many problems are successfully evaluated ' - '(0 = no limit, evaluate the entire sampled set — the ' - 'default, so all 500 stratified MATH problems are run). ' - 'RAG mode: counts problems with valid traces after ' - 'decontam; direct mode: ignored, evaluates all filtered.') - p.add_argument('--db-path', default='./output.oldemb/thinking_rag/lance.db') - p.add_argument('--table', default='thinking_traces') - p.add_argument('--top-k', type=int, default=1) - p.add_argument('--use-cot-compressed', action='store_true', - help='Use pre-compressed cot_compressed field instead of thinking_raw.') - p.add_argument('--sim-threshold', type=float, default=0.75, - help='Minimum cosine similarity for retrieved traces. ' - 'Traces below this are discarded at retrieval time.') - p.add_argument('--decontam-threshold', type=float, default=0.20, - help='13-gram Jaccard threshold for leak detection. ' - 'Retrieved traces above this are skipped (0=disabled).') - p.add_argument('--llm-decontam', action='store_true', default=True, - help='LLM-based decontamination (default ON): API judges whether ' - 'retrieved problem is the same as the test problem. ' - 'Applied after 13-gram decontam, before condensing. ' - 'Use --no-llm-decontam to disable.') - p.add_argument('--no-llm-decontam', dest='llm_decontam', action='store_false', - help='Disable LLM-based decontamination.') - p.add_argument('--max-trace-len', type=int, default=12000) - p.add_argument('--condense', action='store_true', default=True, - help='Enable condenser re-compression on retrieved traces ' - '(default ON). Use --no-condense to inject raw traces.') - p.add_argument('--no-condense', dest='condense', action='store_false', - help='Disable condenser; inject raw retrieved traces.') - p.add_argument('--condense-max-len', type=int, default=2000, - help='Max chars of condensed trace (fallback truncation).') - p.add_argument('--batch-size', type=int, default=16) - p.add_argument('--seed', type=int, default=42) - p.add_argument('--hint', action='store_true', default=False, - help='Enable API hint filtering on retrieved traces (default OFF; ' - 'raw RAG injects the condensed trace directly). ' - 'In rag mode: retrieve → condense → API filters trace → refined system prompt. ' - 'In direct mode: ignored (no traces to filter).') - p.add_argument('--no-hint', dest='hint', action='store_false', - help='Disable API hint filtering; inject condensed trace directly.') - p.add_argument('--problem-ids-file', default=None, - help='File listing problem indices evaluated by RAG mode. ' - 'RAG mode writes this file; direct mode reads it to ' - 'evaluate the same subset (use --no-filter to disable). ' - 'Defaults to a dataset-specific path.') - p.add_argument('--no-filter', action='store_true', - help='In direct mode, evaluate ALL sampled problems ' - 'instead of filtering to RAG subset.') - p.add_argument('--output', default=None) - args = p.parse_args() - - # Dataset-specific default paths (keeps aops and math runs from colliding). - if args.problem_ids_file is None: - args.problem_ids_file = ( - f'./output/thinking_rag/{args.dataset}_rag_problem_ids.json') - - if args.output is None: - suffix = f'{args.mode}_hint' if (args.hint and args.mode == 'rag') else args.mode - args.output = ( - f'./output/thinking_rag/{args.dataset}_{suffix}_results.jsonl') - - if args.condense and args.use_cot_compressed: - logger.warning('--condense requires thinking_raw, ignoring --use-cot-compressed') - args.use_cot_compressed = False - - if args.dataset == 'math': - records = load_math(n=args.n, seed=args.seed, split=args.math_split, - per_level=args.per_level) - else: - records = load_aops(n=args.n, seed=args.seed) - - is_rag = (args.mode == 'rag') - - # Direct mode: filter to same problems RAG evaluated (controlled comparison) - original_indices = list(range(len(records))) # track original indices - if not is_rag and not args.no_filter: - if os.path.exists(args.problem_ids_file): - with open(args.problem_ids_file) as f: - content = f.read().strip() - if content.startswith('['): - valid_indices = set(json.loads(content)) - else: - valid_indices = set(int(line) for line in content.splitlines() if line.strip()) - filtered = [(i, r) for i, r in enumerate(records) if i in valid_indices] - original_indices = [i for i, _ in filtered] - records = [r for _, r in filtered] - sys.stderr.write( - f'[direct] filtered to {len(records)} problems ' - f'from {args.problem_ids_file}\n') - else: - sys.stderr.write( - f'[direct] WARNING: {args.problem_ids_file} not found, ' - f'running all {len(records)} problems\n') - - condenser_gpus = int(os.environ.get('EVAL_CONDENSER_GPUS', 0)) if args.condense else 0 - - # Raw RAG relies on the API condenser (qwen3.7-max). Fail fast with a clear - # message if it's enabled without an API key and without a local condenser. - if is_rag and args.condense and not CONDENSE_API_KEY and condenser_gpus == 0: - sys.stderr.write( - '[condense] ERROR: --condense is ON but COMPRESS_API_KEY is unset ' - 'and no local condenser (EVAL_CONDENSER_GPUS=0).\n' - ' Fix one of:\n' - ' - export COMPRESS_API_KEY=sk-... (use API condenser)\n' - ' - EVAL_CONDENSER_GPUS=2 python ... (use local vLLM condenser)\n' - ' - pass --no-condense (inject raw traces)\n') - sys.exit(1) - - if is_rag: - num_gpus = EMB_GPUS + GEN_GPUS + condenser_gpus - device_groups = [ - DeviceGroup(name='emb_model', ranks=list(range(EMB_GPUS)), - device_type='GPU'), - DeviceGroup(name='sampler', - ranks=list(range(EMB_GPUS, EMB_GPUS + GEN_GPUS)), - device_type='GPU', gpus_per_worker=GEN_GPUS), - ] - if condenser_gpus > 0: - cond_start = EMB_GPUS + GEN_GPUS - device_groups.append( - DeviceGroup(name='condenser', - ranks=list(range(cond_start, cond_start + condenser_gpus)), - device_type='GPU')) - emb_mesh = DeviceMesh.from_sizes(world_size=EMB_GPUS, dp_size=EMB_GPUS) - gen_mesh = DeviceMesh.from_sizes(world_size=GEN_GPUS, tp_size=GEN_GPUS) - twinkle.initialize(mode='ray', nproc_per_node=num_gpus, - groups=device_groups, lazy_collect=False) - else: - device_groups = [ - DeviceGroup(name='sampler', ranks=list(range(GEN_GPUS)), - device_type='GPU', gpus_per_worker=GEN_GPUS), - ] - gen_mesh = DeviceMesh.from_sizes(world_size=GEN_GPUS, tp_size=GEN_GPUS) - twinkle.initialize(mode='ray', nproc_per_node=GEN_GPUS, - groups=device_groups, lazy_collect=False) - - sampler = vLLMSampler( - model_id=GEN_MODEL_ID, - engine_args={ - 'gpu_memory_utilization': GEN_GPU_MEM, - 'max_model_len': GEN_MAX_MODEL_LEN, - }, - device_mesh=gen_mesh, - remote_group='sampler', - ) - sampler.set_template('Qwen3_5Template', model_id=GEN_MODEL_ID, - enable_thinking=True, max_length=GEN_MAX_MODEL_LEN) - sys.stderr.write(f'[aops] vLLM sampler ready (model={GEN_MODEL_ID})\n') - - gen_params = TwinkleSamplingParams( - max_tokens=GEN_MAX_TOKENS, - temperature=GEN_TEMPERATURE, - top_p=GEN_TOP_P, - num_samples=1, - ) - - emb_model = emb_template = tbl = None - if is_rag: - import lancedb - db = lancedb.connect(args.db_path) - if args.table not in db.table_names(): - raise SystemExit(f'Table "{args.table}" not found in {args.db_path}') - tbl = db.open_table(args.table) - sys.stderr.write(f'[aops] LanceDB rows={tbl.count_rows()}\n') - - emb_model = TransformersModel( - model_id=EMBED_MODEL_ID, device_mesh=emb_mesh, - remote_group='emb_model') - emb_model.set_processor(InputProcessor) - emb_model.set_loss(InfonceLoss, temperature=0.03, use_batch=True) - emb_template = Qwen3_5Template( - model_id=EMBED_MODEL_ID, max_length=EMBED_MAX_LENGTH, - truncation_strategy='delete', enable_thinking=False) - sys.stderr.write('[aops] embedding model ready\n') - - # -- Condenser setup (API primary + optional local vLLM) ------------------- - condenser_api_client = None - condenser_sampler_obj = None - condenser_params = None - condenser_special_tokens = None - - if args.condense and is_rag: - condenser_api_client = OpenAIClient( - model=CONDENSE_API_MODEL, api_key=CONDENSE_API_KEY, - base_url=CONDENSE_BASE_URL) - sys.stderr.write(f'[condense] API client ready (model={CONDENSE_API_MODEL})\n') - - if condenser_gpus > 0: - condenser_mesh = DeviceMesh.from_sizes( - world_size=condenser_gpus, dp_size=condenser_gpus) - condenser_sampler_obj = vLLMSampler( - model_id=CONDENSE_MODEL_ID, - engine_args={'gpu_memory_utilization': 0.8, 'max_model_len': 32768}, - device_mesh=condenser_mesh, - remote_group='condenser', - ) - condenser_sampler_obj.set_template( - 'Qwen3_5Template', model_id=CONDENSE_MODEL_ID, - enable_thinking=False, truncation_strategy='delete', - max_length=32768) - condenser_template = Qwen3_5Template( - model_id=CONDENSE_MODEL_ID, max_length=32768, - enable_thinking=False, truncation_strategy='delete') - condenser_special_tokens = set(condenser_template.tokenizer.all_special_tokens) - condenser_params = TwinkleSamplingParams( - max_tokens=CONDENSE_MAX_TOKENS, - temperature=CONDENSE_TEMPERATURE, - top_p=0.5, num_samples=1) - sys.stderr.write(f'[condense] local vLLM ready (model={CONDENSE_MODEL_ID})\n') - - # -- Hint analysis API client (reuses condenser API config) ----------------- - hint_api_client = None - if args.hint and is_rag: - if condenser_api_client is not None: - hint_api_client = condenser_api_client - else: - hint_api_client = OpenAIClient( - model=CONDENSE_API_MODEL, api_key=CONDENSE_API_KEY, - base_url=CONDENSE_BASE_URL) - sys.stderr.write(f'[hint] API hint analysis enabled (model={CONDENSE_API_MODEL})\n') - - # -- LLM decontam API client --------------------------------------------------- - decontam_api_client = None - if args.llm_decontam and is_rag: - if hint_api_client is not None: - decontam_api_client = hint_api_client - elif condenser_api_client is not None: - decontam_api_client = condenser_api_client - else: - decontam_api_client = OpenAIClient( - model=CONDENSE_API_MODEL, api_key=CONDENSE_API_KEY, - base_url=CONDENSE_BASE_URL) - sys.stderr.write(f'[decontam-llm] LLM decontamination enabled (model={CONDENSE_API_MODEL})\n') - - correct_count = 0 - total_count = 0 - skipped_indices: List[int] = [] # problems skipped by RAG (no valid trace) - evaluated_indices: List[int] = [] # problems actually evaluated - debug_records: List[Dict[str, Any]] = [] - - os.makedirs(os.path.dirname(args.output) or '.', exist_ok=True) - out_f = open(args.output, 'w', encoding='utf-8') - - # Open problem-ids files for incremental writing (RAG mode only) - ids_f = None - skip_f = None - if is_rag: - os.makedirs(os.path.dirname(args.problem_ids_file) or '.', exist_ok=True) - ids_f = open(args.problem_ids_file, 'w', encoding='utf-8') - skip_path = args.problem_ids_file.replace('.json', '_skipped.json') - skip_f = open(skip_path, 'w', encoding='utf-8') - - # -- RAG batch preparation (embed + retrieve + decontam + condense + hint) -- - def _prepare_rag_batch(batch_start: int): - """Prepare a RAG batch: returns (prompts, batch, all_examples, - hint_analyses, kept_global_indices, batch_skipped_indices) or None.""" - batch_end = min(batch_start + args.batch_size, len(records)) - batch = records[batch_start:batch_end] - problems = [r['problem'] for r in batch] - - query_vecs = get_embeddings(emb_model, emb_template, problems, EMB_GPUS) - use_raw = not args.use_cot_compressed - all_examples = retrieve_examples(tbl, query_vecs, args.top_k, - use_raw, args.sim_threshold, - problems=problems, - decontam_threshold=args.decontam_threshold) - if args.use_cot_compressed: - for exs in all_examples: - for ex in exs: - ex['thinking'] = _strip_condenser_markers(ex['thinking']) - - if args.llm_decontam and decontam_api_client: - all_examples = _llm_decontaminate( - decontam_api_client, problems, all_examples) - - if args.condense and condenser_api_client: - all_examples = condense_traces( - all_examples, problems, condenser_api_client, - condenser_sampler=condenser_sampler_obj, - compress_params=condenser_params, - special_tokens=condenser_special_tokens, - max_output_len=args.condense_max_len, - dp_size=condenser_gpus) - - hint_analyses = None - if args.hint and hint_api_client: - hint_analyses = _api_hint_analysis_batch( - hint_api_client, problems, all_examples) - - keep_mask = [] - for pi, (r, examples) in enumerate(zip(batch, all_examples)): - if not examples: - keep_mask.append(False) - elif hint_analyses and hint_analyses[pi]: - keep_mask.append(True) - else: - usable = [ex for ex in examples - if len(ex['thinking']) <= args.max_trace_len] - keep_mask.append(bool(usable)) - - batch_skipped = [] - for pi, keep in enumerate(keep_mask): - if not keep: - batch_skipped.append(batch_start + pi) - - kept_batch = [] - kept_examples = [] - kept_hints = [] - kept_global_indices = [] - for pi, keep in enumerate(keep_mask): - if keep: - kept_batch.append(batch[pi]) - kept_examples.append(all_examples[pi]) - kept_hints.append(hint_analyses[pi] if hint_analyses else None) - kept_global_indices.append(batch_start + pi) - - if not kept_batch: - return None, None, None, None, None, batch_skipped - - prompts = [] - for pi, (r, examples) in enumerate(zip(kept_batch, kept_examples)): - if kept_hints[pi]: - prompts.append(build_hint_prompt(r['problem'], kept_hints[pi])) - else: - filtered = [{'query': ex['query'], 'thinking': ex['thinking']} - for ex in examples - if len(ex['thinking']) <= args.max_trace_len] - prompts.append(build_rag_prompt(r['problem'], filtered)) - - return prompts, kept_batch, kept_examples, kept_hints, kept_global_indices, batch_skipped - - target_reached = False - batch_starts = list(range(0, len(records), args.batch_size)) - - if is_rag: - # Pipeline: prefetch next batch while current batch generates - from concurrent.futures import Future - prefetch_pool = ThreadPoolExecutor(max_workers=1) - # Prepare first batch synchronously - cur_result = _prepare_rag_batch(batch_starts[0]) - - for bi, batch_start in enumerate(batch_starts): - if target_reached: - break - prompts, batch, all_examples, hint_analyses, kept_global_indices, batch_skipped = cur_result - skipped_indices.extend(batch_skipped or []) - if skip_f and batch_skipped: - for sid in batch_skipped: - skip_f.write(f'{sid}\n') - skip_f.flush() - - # Submit next batch preparation in background - next_future: Optional[Future] = None - if bi + 1 < len(batch_starts) and not target_reached: - next_future = prefetch_pool.submit(_prepare_rag_batch, batch_starts[bi + 1]) - - if prompts is None: - # Entire batch skipped - cur_result = next_future.result() if next_future else None - continue - - # Generate (runs on gen GPU while next batch prepares on emb GPU + API) - responses = sampler.sample(prompts, gen_params) - - for i, (rec, resp) in enumerate(zip(batch, responses)): - seq = resp.sequences[0] if resp and resp.sequences else None - raw_output = '' - if seq is not None: - raw_output = seq.decoded or '' - raw_output = re.sub(r'<\|[^|]+\|>', '', raw_output).rstrip() - - predicted = extract_boxed(raw_output) - is_correct = answers_match(predicted, rec['reference_answer']) - if is_correct: - correct_count += 1 - total_count += 1 - - global_idx = kept_global_indices[i] - evaluated_indices.append(global_idx) - if ids_f: - ids_f.write(f'{global_idx}\n') - ids_f.flush() - - debug_rec = { - 'idx': global_idx, - 'reference_answer': rec['reference_answer'], - 'predicted': predicted, - 'is_correct': is_correct, - 'problem': rec['problem'], - 'model_output': raw_output, - } - if rec.get('level'): - debug_rec['level'] = rec['level'] - if rec.get('type'): - debug_rec['type'] = rec['type'] - debug_rec['num_traces'] = len(all_examples[i]) - if all_examples[i]: - ex0 = all_examples[i][0] - debug_rec['similarity'] = ex0.get('_sim', 0.0) - debug_rec['retrieved_query'] = ex0.get('query', '') - debug_rec['raw_trace_len'] = ex0.get('_raw_trace_len', 0) - debug_rec['condensed_trace'] = ex0['thinking'] - debug_rec['condensed_trace_len'] = len(ex0['thinking']) - debug_rec['condense_source'] = ex0.get('_condense_source', '') - if hint_analyses and hint_analyses[i]: - debug_rec['hint_analysis'] = hint_analyses[i] - debug_records.append(debug_rec) - out_f.write(json.dumps(debug_rec, ensure_ascii=False) + '\n') - out_f.flush() - - acc = correct_count / total_count if total_count else 0 - sys.stderr.write( - f' [{total_count}/{args.target_eval}] ' - f'acc={acc:.4f} ({correct_count}/{total_count})\n') - - if args.target_eval > 0 and total_count >= args.target_eval: - target_reached = True - - # Collect prefetched result for next iteration (skip if done) - if not target_reached and next_future: - cur_result = next_future.result() - else: - cur_result = None - - prefetch_pool.shutdown(wait=True) - else: - # Direct mode: no pipeline needed, just batch generate - for batch_start in batch_starts: - batch_end = min(batch_start + args.batch_size, len(records)) - batch = records[batch_start:batch_end] - prompts = [build_direct_prompt(r['problem']) for r in batch] - - responses = sampler.sample(prompts, gen_params) - - for i, (rec, resp) in enumerate(zip(batch, responses)): - seq = resp.sequences[0] if resp and resp.sequences else None - raw_output = '' - if seq is not None: - raw_output = seq.decoded or '' - raw_output = re.sub(r'<\|[^|]+\|>', '', raw_output).rstrip() - - predicted = extract_boxed(raw_output) - is_correct = answers_match(predicted, rec['reference_answer']) - if is_correct: - correct_count += 1 - total_count += 1 - - global_idx = original_indices[batch_start + i] - evaluated_indices.append(global_idx) - - debug_rec = { - 'idx': global_idx, - 'reference_answer': rec['reference_answer'], - 'predicted': predicted, - 'is_correct': is_correct, - 'problem': rec['problem'], - 'model_output': raw_output, - } - if rec.get('level'): - debug_rec['level'] = rec['level'] - if rec.get('type'): - debug_rec['type'] = rec['type'] - debug_records.append(debug_rec) - out_f.write(json.dumps(debug_rec, ensure_ascii=False) + '\n') - out_f.flush() - - acc = correct_count / total_count if total_count else 0 - sys.stderr.write( - f' [{total_count}/{len(records)}] ' - f'acc={acc:.4f} ({correct_count}/{total_count})\n') - - overall_acc = correct_count / total_count if total_count else 0 - print(f'\n{"=" * 60}') - print(f'{args.dataset.upper()} — mode={args.mode}, model={GEN_MODEL_ID}') - print(f' n={total_count}, seed={args.seed}') - if is_rag: - print(f' evaluated={len(evaluated_indices)}, skipped={len(skipped_indices)}') - print(f'{"=" * 60}') - print(f'Overall accuracy: {overall_acc:.4f} ({correct_count}/{total_count})') - - # Per-level breakdown (MATH: the difficulty-vs-gain curve we care about). - if any(r.get('level') for r in debug_records): - from collections import defaultdict - per = defaultdict(lambda: [0, 0]) # level -> [correct, total] - for r in debug_records: - lv = r.get('level', 'Unknown') - per[lv][1] += 1 - if r['is_correct']: - per[lv][0] += 1 - print(f'\nPer-level accuracy:') - for lv in sorted(per.keys()): - c, t = per[lv] - print(f' {lv:>10}: {c/t:.4f} ({c}/{t})') - - out_f.close() - print(f'\n[output] {len(debug_records)} records saved to {args.output}') - - if ids_f: - ids_f.close() - print(f'[output] problem IDs ({len(evaluated_indices)}) saved to {args.problem_ids_file}') - if skip_f: - skip_f.close() - if skipped_indices: - print(f'[output] skipped IDs ({len(skipped_indices)}) saved to ' - f'{args.problem_ids_file.replace(".json", "_skipped.json")}') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/embedding/eval_rag_recall.py b/cookbook/exp/embedding/eval_rag_recall.py deleted file mode 100644 index 19bc5ad6d..000000000 --- a/cookbook/exp/embedding/eval_rag_recall.py +++ /dev/null @@ -1,187 +0,0 @@ -"""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/eval_reflexion_skill.py b/cookbook/exp/embedding/eval_reflexion_skill.py deleted file mode 100644 index 3f7ddcc0d..000000000 --- a/cookbook/exp/embedding/eval_reflexion_skill.py +++ /dev/null @@ -1,762 +0,0 @@ -"""Phase-0 measurement for the reflexion self-skill scheme (see reflexion.md). - -Question this script answers: **on problems the base model first gets wrong, does -letting the SAME base model reflect on its failed attempt, distill a general -"skill", and re-solve WITH that skill in the system prompt, actually raise its -pass@k?** No LoRA is trained here — this is the upper-bound / go-no-go gate before -investing in a Skill-LoRA. If the base model's own skills don't help, training a -LoRA to produce them is pointless. - -It deliberately reuses ``eval_gpqa_rag`` verbatim (dataset, grader, prompts, sampling -config) so numbers are comparable with the other AoPS lines. Only the base model + -one vLLM sampler are used; the dataset is AoPS; validation is on the SAME problem -(no similar-problem retrieval). - -Per chunk of problems (all sampler calls are BATCHED across the whole chunk — never -one problem at a time): - 1. Initial solve — 1 rollout each; keep only problems the model got wrong. - 2. Skill generation — for each failed problem, the base model reads its own failed - attempt and produces N candidate skills (general reminders, no answer/solution). - 3. Leak filter — drop skills that leak the gold answer or a full solution. - 4. Baseline pass@k — K rollouts of the plain problem (the "no-skill" control). - 5. With-skill pass@k — K rollouts of the problem with each surviving skill in the - system prompt. - 6. Score — marginal = with-skill pass@k − baseline pass@k; keep the best skill. -A "pass" = answer correct AND generation terminated (no length cutoff). - -Everything useful (failed attempt, every candidate skill + leak flag, baseline and -per-skill rollout stats, marginals, best skill) is written to a JSONL **incrementally -after each chunk**, so partial runs are fully analysable. - -Launch (8 GPUs, tp=1 dp=8 by default): - python cookbook/exp/embedding/eval_reflexion_skill.py --n 64 --chunk-size 16 -""" -import argparse -import copy -import json -import os -import re -import sys -import threading -import time -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Any, Dict, List, Optional, Tuple - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.data_format import SamplingParams as TwinkleSamplingParams -from twinkle.sampler import vLLMSampler -from twinkle_agentic.protocol.openai import OpenAI as OpenAIClient - -# Reuse the reference eval's dataset + grading + prompts + sampling config so this -# line is directly comparable with eval_gpqa_rag / eval_dualline_math. -from eval_gpqa_rag import (DIRECT_SYSTEM, GEN_GPU_MEM, GEN_GPUS, GEN_MODEL_ID, - GEN_TEMPERATURE, GEN_TOP_P, MCQ_INSTRUCTION, answers_match, - build_direct_prompt, extract_boxed, load_aops) - -logger = get_logger() - -# vLLM parallel: tp=1, dp=GEN_GPUS by default (override GEN_TP; keep GEN_GPUS=8). -GEN_TP = int(os.environ.get('GEN_TP', 1)) - -# Leak-detector API (reuses eval_gpqa_rag's env names). A strong external model -# judges whether a candidate skill leaks THIS problem's answer/solution — catching -# what the string filter cannot (multiple-choice letters, derived-result leakage). -LEAK_API_KEY = os.environ.get('COMPRESS_API_KEY', '') -LEAK_BASE_URL = os.environ.get('COMPRESS_BASE_URL', - 'https://dashscope.aliyuncs.com/compatible-mode/v1') -LEAK_API_MODEL = os.environ.get('COMPRESS_MODEL', 'qwen3.7-max') - -# Global call spacer so bursts of leak-judge calls stay under the QPS limit. -_api_lock = threading.Lock() -_api_next = [0.0] - - -def _api_throttle(min_interval: float) -> None: - with _api_lock: - now = time.monotonic() - wait = max(0.0, _api_next[0] - now) - _api_next[0] = max(now, _api_next[0]) + min_interval - if wait > 0: - time.sleep(wait) - - -# --------------------------------------------------------------------------- -# Prompts (self-reflection skill generation + skill-conditioned solving) -# --------------------------------------------------------------------------- -SKILL_GEN_SYSTEM = ( - "You are a meticulous mathematics coach. You are shown a competition problem and a " - "student's FAILED attempt. Produce a SHORT list of general, reusable skills that " - 'would prevent this class of mistake on SIMILAR problems.\n\n' - 'OUTPUT FORMAT (strict):\n' - '- You may reason briefly first, but the final answer MUST be a markdown bullet ' - 'list of 3-5 items WRAPPED IN and tags. Output nothing after ' - '.\n' - '- Each item is ONE short imperative sentence (a rule, check, or habit).\n' - '- Inside the tags: no diagnosis narration, no "The student...", no headings, no ' - 'restating the problem or the examples.\n\n' - 'CONTENT RULES (strict):\n' - '- Do NOT reveal the final answer or the multiple-choice option.\n' - '- Do NOT state the specific numbers, values, or key intermediate results of THIS ' - 'problem.\n' - '- Do NOT give a step-by-step solution to THIS problem. Every item must be GENERAL ' - 'and transferable to other problems of the same type.\n\n' - 'Follow the example below for the exact tags, style, and level of generality.' -) - -SKILL_GEN_USER = ( - 'Problem:\n{problem}\n\n' - "The student's failed attempt (it may be long or may fail to terminate):\n" - '{attempt}\n\n' - 'Now output the skills bullet list.' -) - -# One-shot demonstration of the required format and generality (answer-free). -_EX_PROBLEM = 'Simplify $\\sqrt{72} + \\sqrt{18}$ and give the result.' -_EX_ATTEMPT = ( - 'The student added the radicands directly to get $\\sqrt{90}$ and concluded it ' - 'could not be simplified, never factoring out the perfect squares first.') -_EX_SKILLS = ( - '\n' - '- Before adding square roots, factor each radicand into a perfect square times a ' - 'remainder and move the perfect-square root outside.\n' - '- Never add radicands directly: $\\sqrt{a}+\\sqrt{b}\\ne\\sqrt{a+b}$.\n' - '- Only combine radical terms after reducing them to the same simplest radical ' - 'form.\n' - '- Sanity-check the simplified result by estimating each root numerically.\n' - '') - - -def build_skillgen_prompt(problem: str, attempt: str) -> Dict[str, Any]: - return {'messages': [ - {'role': 'system', 'content': SKILL_GEN_SYSTEM}, - {'role': 'user', - 'content': SKILL_GEN_USER.format(problem=_EX_PROBLEM, attempt=_EX_ATTEMPT)}, - {'role': 'assistant', 'content': _EX_SKILLS}, - {'role': 'user', - 'content': SKILL_GEN_USER.format(problem=problem, attempt=attempt)}, - ]} - - -# The skill is injected into the SYSTEM prompt (per reflexion.md), on top of the -# exact DIRECT_SYSTEM used by the baseline so the only difference is the reminders. -# Built by concatenation (NOT str.format): DIRECT_SYSTEM and the skill may contain -# literal braces (e.g. ``\boxed{}``, LaTeX), which would break ``.format``. -_SKILL_SOLVE_PREFIX = ( - DIRECT_SYSTEM + '\n\n' - 'Before you start, keep these reminders in mind to avoid common mistakes on this ' - 'type of problem:\n') -_SKILL_SOLVE_SUFFIX = ( - '\nApply them where relevant, but rely on your own reasoning to reach the answer.') - - -def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: - return {'messages': [ - {'role': 'system', 'content': _SKILL_SOLVE_PREFIX + skill + _SKILL_SOLVE_SUFFIX}, - {'role': 'user', 'content': problem + MCQ_INSTRUCTION}, - ]} - - -# --------------------------------------------------------------------------- -# Parsing / grading / leak filtering -# --------------------------------------------------------------------------- -_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') - - -def _clean_text(decoded: Optional[str]) -> str: - return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() - - -_BULLET_RE = re.compile(r'^\s*(?:[-*]|\d+[.)])\s') - - -def _extract_skill_list(text: str) -> str: - """Pull just the clean skill list out of a (possibly thinking-laden) output. - - The model is instructed to wrap the final list in ``...``, so - prefer that (robust to any preceding reasoning, closed or unterminated). Fall - back to dropping a ```` block and keeping from the first bullet onward. - """ - low = text.lower() - if '' in low: - start = low.index('') + len('') - end = low.index('') if '' in low else len(text) - return text[start:end].strip() - if '' in text: - text = text.rsplit('', 1)[-1] - text = re.sub(r'.*?', '', text, flags=re.DOTALL).strip() - lines = text.splitlines() - for i, line in enumerate(lines): - if _BULLET_RE.match(line): - return '\n'.join(lines[i:]).strip() - return text.strip() - - -def _bound_attempt(text: str, gen_tokens: int, budget_tokens: int) -> str: - """Keep a failed attempt within the skill-gen context budget. - - Round-2 (skill-gen) input contains the FULL round-1 attempt, and failed - attempts are often the ones that ran to the token cap (repetition loops), so - feeding them verbatim overflows max_model_len. Keep the head (real reasoning + - where it went wrong) plus a short tail (the final wrong answer); drop the - redundant middle. Token->char conversion uses THIS attempt's observed - chars-per-token so the cut fits precisely. - """ - if not text or gen_tokens <= budget_tokens: - return text - cpt = len(text) / max(1, gen_tokens) - head_tok = int(budget_tokens * 0.7) - tail_tok = budget_tokens - head_tok - head = text[:int(head_tok * cpt)] - tail = text[-int(tail_tok * cpt):] if tail_tok > 0 else '' - return f'{head}\n\n[... attempt truncated for length ...]\n\n{tail}' - - -def _parse_seq(seq, gold: str) -> Dict[str, Any]: - """Turn one sampled sequence into a graded rollout record. - - ``pass`` requires BOTH a correct boxed answer AND clean termination (a length - cutoff means the model never actually committed to the answer). - """ - text = _clean_text(getattr(seq, 'decoded', '') or '') - pred = extract_boxed(text) - correct = bool(pred) and answers_match(pred, gold) - terminated = getattr(seq, 'stop_reason', None) != 'length' - return { - 'pred': pred, - 'correct': correct, - 'terminated': terminated, - 'passed': bool(correct and terminated), - 'stop_reason': getattr(seq, 'stop_reason', None), - 'gen_tokens': len(getattr(seq, 'tokens', None) or []), - 'text': text, - } - - -def _pass_rate(rolls: List[Dict[str, Any]]) -> float: - return sum(1 for r in rolls if r['passed']) / len(rolls) if rolls else 0.0 - - -def _skill_leaks(skill: str, gold: str) -> Tuple[bool, str]: - """Reject a skill that leaks the answer, or is a degenerate / non-list output.""" - if not skill.strip(): - return True, 'empty' - bullets = [ln for ln in skill.splitlines() if _BULLET_RE.match(ln)] - if len(bullets) < 2: - return True, 'too_short' - low = skill.lower() - if 'item 1' in low and 'item 2' in low: # model echoed the format placeholder - return True, 'placeholder' - if '\\boxed' in skill: - return True, 'contains_boxed' - g = (gold or '').strip() - # Raw substring match is only trustworthy when the answer is specific enough that - # an incidental hit is unlikely. Short answers ('D', 'E', '1') would match almost - # any text, so leave those to the API judge instead of false-flagging every skill. - if len(g) >= 4 and g.lower() in skill.lower(): - return True, 'contains_gold_answer' - # Standalone multi-digit numbers from the gold answer leaking into the skill. - for num in re.findall(r'-?\d{2,}', g): - if re.search(r'(? Optional[bool]: - """Return True (leak) / False (clean) / None (unparseable or API error after retries). - - Only transient API errors are retried (with exponential backoff); an unparseable - verdict is deterministic at temperature 0, so retrying it is pointless. - """ - msgs = [ - {'role': 'system', 'content': _LEAK_JUDGE_SYSTEM}, - {'role': 'user', 'content': _LEAK_JUDGE_USER.format( - problem=problem[:4000], gold=gold, skill=skill[:4000])}, - ] - for attempt in range(retries + 1): - _api_throttle(min_interval) - try: - reply = api({'messages': msgs}, - TwinkleSamplingParams(temperature=0.0, max_tokens=16), - extra_body={'enable_thinking': False}) - except Exception as exc: # noqa: BLE001 — broad catch is intentional - logger.warning(f'[leak-judge] error (attempt {attempt + 1}/{retries + 1}): {exc}') - if attempt < retries: - time.sleep(min(4.0, 0.5 * 2 ** attempt)) # exponential backoff - continue - return None - verdict = (reply.get('content') or '').strip().upper() - if 'CLEAN' in verdict: - return False - if 'LEAK' in verdict: - return True - return None # unparseable — deterministic at temp 0, no point retrying - - -def _api_leak_batch(api: OpenAIClient, items: List[Tuple[int, str, str, str]], - concurrency: int, min_interval: float, - retries: int) -> Dict[int, Optional[bool]]: - """Judge many (key, problem, gold, skill) tuples in parallel; key -> verdict.""" - verdicts: Dict[int, Optional[bool]] = {} - if not items: - return verdicts - with ThreadPoolExecutor(max_workers=min(len(items), concurrency)) as pool: - futs = {pool.submit(_api_leak_judge_one, api, p, g, s, min_interval, retries): k - for (k, p, g, s) in items} - for fut in as_completed(futs): - verdicts[futs[fut]] = fut.result() - return verdicts - - -# --------------------------------------------------------------------------- -# Batched sampling (one shared sampler.sample per phase — never per problem) -# --------------------------------------------------------------------------- -def _pad_for_dp(prompts: List[Any], gen_dp: int) -> List[Any]: - """vLLM dp needs batch len >= dp; pad tail rounds and let the caller slice back.""" - if gen_dp <= 1 or not prompts or len(prompts) >= gen_dp: - return prompts - pad = [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] - return prompts + pad - - -def _run_samples(sampler, prompts: List[Any], num_samples: int, max_tokens: int, - gen_dp: int, temperature: Optional[float] = None, - top_p: Optional[float] = None, - top_k: Optional[int] = None) -> List[List[Any]]: - """One batched sampler call; return per-prompt list of raw sampled sequences. - - ``temperature``/``top_p``/``top_k`` default to the module's sampling config; pass - ``temperature=0.0`` for deterministic greedy decoding (SEAM-style executor scoring), - or a high ``temperature`` with ``top_k=-1`` for diverse multi-candidate sampling.""" - if not prompts: - return [] - params = TwinkleSamplingParams( - max_tokens=max_tokens, - temperature=GEN_TEMPERATURE if temperature is None else temperature, - top_p=GEN_TOP_P if top_p is None else top_p, - num_samples=num_samples, - **({} if top_k is None else {'top_k': top_k})) - padded = _pad_for_dp(prompts, gen_dp) - responses = sampler.sample(padded, params)[:len(prompts)] - return [list(r.sequences) if (r and r.sequences) else [] for r in responses] - - -def _set_thinking(sampler, args: argparse.Namespace, enabled: bool) -> None: - """Toggle the remote template's thinking mode. - - Skill generation wants thinking OFF so the model emits the short ```` - list directly (with thinking ON it burns the token budget reasoning and often - never reaches the list); solving wants it ON. ``set_template`` is a - remote_function, so this propagates to every sampler worker. - """ - sampler.set_template('Qwen3_5Template', model_id=GEN_MODEL_ID, - enable_thinking=enabled, max_length=args.max_model_len) - - -def _bounded_attempt_for(r: Dict[str, Any], args: argparse.Namespace) -> str: - """Per-problem bound so problem + attempt + skill output fits the context window.""" - prob_est = len(r['problem']) // 2 # conservative problem token estimate - budget = max(1024, args.max_model_len - args.skill_max_tokens - - args.attempt_reserve_tokens - prob_est) - return _bound_attempt(r['_init'][0]['text'], r['_init'][0]['gen_tokens'], budget) - - -def _filter_candidates(api: Optional[OpenAIClient], - cands: List[Tuple[Dict[str, Any], str]], - args: argparse.Namespace) -> None: - """Apply string + API leak filters to (problem, skill) candidates; append results. - - The cheap string filter runs first; the API judge only sees skills that pass it, - which is what catches MCQ-letter / derived-result / full-solution leakage. - """ - prepared = [] # [r, text, leaked(bool|None), reason] - for r, text in cands: - leaked, reason = _skill_leaks(text, r['reference_answer']) - prepared.append([r, text, True if leaked else None, reason]) - if api is not None: - items = [(i, prepared[i][0]['problem'], prepared[i][0]['reference_answer'], - prepared[i][1]) for i in range(len(prepared)) if prepared[i][2] is None] - verdicts = _api_leak_batch(api, items, args.api_concurrency, args.api_min_interval, - args.api_retries) - for key, _p, _g, _s in items: - v = verdicts.get(key) - if v is True: - prepared[key][2], prepared[key][3] = True, 'api_leak' - elif v is False: - prepared[key][2], prepared[key][3] = False, '' - else: - prepared[key][2], prepared[key][3] = False, 'api_uncertain' - for r, text, leaked, reason in prepared: - r['_skills'].append({'skill': text, 'leaked': bool(leaked), 'leak_reason': reason}) - - -def _build_skills(sampler, api: Optional[OpenAIClient], failed: List[Dict[str, Any]], - gen_dp: int, args: argparse.Namespace) -> None: - """Generate + extract + leak-filter skills, re-rolling problems short on clean ones. - - Clean skills accumulate across rounds; only problems still below ``min_survivors`` - clean skills are re-rolled, up to ``skill_retries`` extra rounds. - """ - for r in failed: - r['_skills'] = [] - todo = list(failed) - _set_thinking(sampler, args, False) # skill-gen: emit the list directly, no CoT - try: - for _ in range(args.skill_retries + 1): - if not todo: - break - sg_out = _run_samples( - sampler, - [build_skillgen_prompt(r['problem'], _bounded_attempt_for(r, args)) for r in todo], - args.n_skills, args.skill_max_tokens, gen_dp) - cands = [(r, _extract_skill_list(_clean_text(getattr(s, 'decoded', '') or ''))) - for r, seqs in zip(todo, sg_out) for s in seqs] - _filter_candidates(api, cands, args) - todo = [r for r in failed - if sum(1 for sk in r['_skills'] if not sk['leaked']) < args.min_survivors] - finally: - _set_thinking(sampler, args, True) # restore for solving phases - tot = sum(len(r['_skills']) for r in failed) - leaked = sum(1 for r in failed for sk in r['_skills'] if sk['leaked']) - sys.stderr.write(f' phase2: skills={tot} leaked={leaked} ' - f'({leaked / max(1, tot):.0%}); {len(todo)} still short of ' - f'{args.min_survivors} clean\n') - - -# --------------------------------------------------------------------------- -# Per-chunk pipeline -# --------------------------------------------------------------------------- -def process_chunk(sampler, api: Optional[OpenAIClient], chunk: List[Dict[str, Any]], - gen_dp: int, args: argparse.Namespace) -> List[Dict[str, Any]]: - """Run all 6 phases for one chunk (batched) and return per-problem records.""" - # --- Phase 1: initial solve, keep only the ones the model got wrong. --- - init_out = _run_samples( - sampler, [build_direct_prompt(r['problem']) for r in chunk], - args.init_samples, args.max_tokens, gen_dp) - for r, seqs in zip(chunk, init_out): - r['_init'] = [_parse_seq(s, r['reference_answer']) for s in seqs] - r['_init_pass'] = _pass_rate(r['_init']) - r['_failed'] = r['_init_pass'] == 0.0 - failed = [r for r in chunk if r['_failed']] - sys.stderr.write(f' phase1: {len(chunk)-len(failed)}/{len(chunk)} solved on ' - f'first try, {len(failed)} failed -> reflect\n') - - if failed: - # --- Baseline pass@k FIRST: defines which failures are genuinely hard. --- - # (A single initial rollout is noisy; an easy problem can fail phase 1 yet - # have a high pass@k, so measure the marginal only on truly hard problems.) - base_out = _run_samples( - sampler, [build_direct_prompt(r['problem']) for r in failed], - args.pass_k, args.max_tokens, gen_dp) - for r, seqs in zip(failed, base_out): - r['_baseline'] = [_parse_seq(s, r['reference_answer']) for s in seqs] - r['_baseline_pass'] = _pass_rate(r['_baseline']) - r['_hard'] = r['_baseline_pass'] <= args.hard_baseline_max - r['_skills'] = [] - r['_best'] = None - hard = [r for r in failed if r['_hard']] - sys.stderr.write(f' baseline: {len(hard)}/{len(failed)} failures are hard ' - f'(pass@{args.pass_k} <= {args.hard_baseline_max})\n') - - if hard: - # --- Skills (generate + leak filter + re-rollout) for HARD problems only. --- - _build_skills(sampler, api, hard, gen_dp, args) - - # --- With-skill pass@k (flatten hard-problem x surviving skill). --- - flat: List[Tuple[int, int]] = [] - ws_prompts: List[Any] = [] - for ri, r in enumerate(hard): - for si, sk in enumerate(r['_skills']): - if sk['leaked'] or not sk['skill'].strip(): - continue - flat.append((ri, si)) - ws_prompts.append(build_skill_solve_prompt(r['problem'], sk['skill'])) - ws_out = _run_samples(sampler, ws_prompts, args.pass_k, args.max_tokens, gen_dp) - for (ri, si), seqs in zip(flat, ws_out): - r = hard[ri] - sk = r['_skills'][si] - sk['rolls'] = [_parse_seq(s, r['reference_answer']) for s in seqs] - sk['with_pass'] = _pass_rate(sk['rolls']) - sk['marginal'] = sk['with_pass'] - r['_baseline_pass'] - - # --- Pick the best (highest marginal) surviving skill per hard problem. --- - for r in hard: - scored = [sk for sk in r['_skills'] if 'marginal' in sk] - r['_best'] = max(scored, key=lambda s: s['marginal']) if scored else None - - return [_make_record(r, args) for r in chunk] - - -def _roll_summary(roll: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: - out = {k: roll[k] for k in ('pred', 'correct', 'terminated', 'passed', - 'stop_reason', 'gen_tokens')} - if args.store_rollout_text: - out['text'] = roll['text'][:args.store_rollout_chars] - return out - - -def _make_record(r: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: - """Assemble the incremental JSONL record for one problem (solved or failed).""" - rec: Dict[str, Any] = { - 'problem': r['problem'], - 'reference_answer': r['reference_answer'], - 'tags': r.get('tags', []), - 'failed_first_try': r['_failed'], - 'init_pass_rate': r['_init_pass'], - 'init_attempt': { - 'text': r['_init'][0]['text'][:args.store_init_chars], - 'pred': r['_init'][0]['pred'], - 'stop_reason': r['_init'][0]['stop_reason'], - 'gen_tokens': r['_init'][0]['gen_tokens'], - }, - } - if not r['_failed']: - return rec - - best = r.get('_best') - rec['baseline_pass'] = r['_baseline_pass'] - # Genuinely hard = low baseline pass@k; only these count in the marginal stats. - rec['is_hard'] = bool(r.get('_hard')) - rec['baseline_rolls'] = [_roll_summary(x, args) for x in r['_baseline']] - rec['skills'] = [{ - 'skill': sk['skill'], - 'leaked': sk['leaked'], - 'leak_reason': sk['leak_reason'], - 'with_pass': sk.get('with_pass'), - 'marginal': sk.get('marginal'), - 'rolls': [_roll_summary(x, args) for x in sk.get('rolls', [])], - } for sk in r.get('_skills', [])] - rec['best_skill'] = best['skill'] if best else None - rec['best_marginal'] = best['marginal'] if best else None - rec['best_with_pass'] = best['with_pass'] if best else None - # "rescued" = a leak-free skill turned a fully-failing problem into some passes. - rec['rescued'] = bool(best and r['_baseline_pass'] == 0.0 and best['with_pass'] > 0.0) - rec['helped'] = bool(best and best['marginal'] > 0.0) - return rec - - -# --------------------------------------------------------------------------- -# Running summary -# --------------------------------------------------------------------------- -def _update_summary(summ: Dict[str, Any], recs: List[Dict[str, Any]]) -> None: - for rec in recs: - summ['n_total'] += 1 - if not rec['failed_first_try']: - summ['n_solved_first'] += 1 - continue - summ['n_failed'] += 1 - if not rec.get('is_hard'): - summ['n_failed_easy'] += 1 # failed phase 1 but easy on pass@k — excluded - continue - summ['n_hard'] += 1 - base = rec.get('baseline_pass', 0.0) - summ['sum_baseline_pass'] += base - if rec.get('best_marginal') is not None: - summ['n_with_skill'] += 1 - summ['sum_best_with_pass'] += rec.get('best_with_pass', 0.0) - summ['sum_best_marginal'] += rec.get('best_marginal', 0.0) - else: - # No clean skill produced for this hard problem -> skill adds no gain - # (count it honestly as marginal 0 rather than dropping it from the average). - summ['sum_best_with_pass'] += base - summ['n_helped'] += int(rec.get('helped', False)) - summ['n_rescued'] += int(rec.get('rescued', False)) - - -def _summary_report(summ: Dict[str, Any]) -> Dict[str, Any]: - nh = max(1, summ['n_hard']) - return { - 'record_type': 'summary', - 'n_total': summ['n_total'], - 'n_solved_first_try': summ['n_solved_first'], - 'n_failed_first_try': summ['n_failed'], - 'n_failed_but_easy': summ['n_failed_easy'], - 'n_hard': summ['n_hard'], - 'n_hard_with_skill': summ['n_with_skill'], - # Averages are over ALL hard problems; a hard problem with no clean skill - # counts as zero gain (with_pass == baseline), so with - base == marginal. - 'avg_baseline_pass_on_hard': summ['sum_baseline_pass'] / nh, - 'avg_best_with_skill_pass_on_hard': summ['sum_best_with_pass'] / nh, - 'avg_best_marginal_on_hard': summ['sum_best_marginal'] / nh, - 'n_helped_by_skill': summ['n_helped'], - 'n_rescued_from_zero': summ['n_rescued'], - 'frac_hard_helped': summ['n_helped'] / nh, - } - - -# --------------------------------------------------------------------------- -# Main -# --------------------------------------------------------------------------- -def main() -> None: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--n', type=int, default=64, help='AoPS problems to sample.') - p.add_argument('--seed', type=int, default=42) - p.add_argument('--chunk-size', type=int, default=16, - help='Problems per chunk. All sampler calls within a chunk are ' - 'batched; results are flushed to disk after each chunk.') - p.add_argument('--init-samples', type=int, default=1, - help='Rollouts for the initial solve. A problem is "failed" (and ' - 'sent to reflection) only if all initial rollouts are wrong.') - p.add_argument('--n-skills', type=int, default=8, - help='Candidate skills generated per failed problem.') - p.add_argument('--pass-k', type=int, default=8, - help='Rollouts per (baseline / with-skill) pass@k estimate.') - p.add_argument('--hard-baseline-max', type=float, default=0.25, - help='A failed problem counts as "hard" (included in the marginal ' - 'stats) only if its baseline pass@k <= this. Filters out easy ' - 'problems that merely failed the single initial rollout.') - p.add_argument('--max-model-len', type=int, default=30000, - help='Context window (engine + template). MUST exceed --max-tokens: ' - 'the round-2 skill-gen input holds the full round-1 attempt ' - 'plus the problem.') - p.add_argument('--max-tokens', type=int, default=20000, - help='Max generated tokens for solving rollouts (round-1 output cap).') - p.add_argument('--skill-max-tokens', type=int, default=2048, - help='Max tokens for skill generation. Enough to finish any thinking ' - 'and emit the short bullet list (which is then extracted).') - p.add_argument('--attempt-reserve-tokens', type=int, default=2048, - help='Tokens reserved for system prompt + wrappers when bounding the ' - 'failed attempt fed into skill generation (the problem length ' - 'is accounted for separately, per-problem).') - p.add_argument('--min-survivors', type=int, default=2, - help='Re-roll a problem\'s skills if fewer than this many survive the ' - 'leak filters.') - p.add_argument('--skill-retries', type=int, default=1, - help='Max extra skill-generation rounds for problems short on clean ' - 'skills (0 = no retry).') - p.add_argument('--api-concurrency', type=int, default=32, - help='Parallel workers for the API leak judge (max 32 recommended).') - p.add_argument('--api-min-interval', type=float, default=0.1, - help='Minimum seconds between API leak-judge calls (QPS guard).') - p.add_argument('--api-retries', type=int, default=3, - help='Retries on transient API errors per leak-judge call (exponential ' - 'backoff); only after these are exhausted is a skill kept as ' - 'api_uncertain.') - p.add_argument('--disable-api-leak', action='store_true', - help='Skip the API leak judge even if COMPRESS_API_KEY is set ' - '(string filter only).') - p.add_argument('--output', default='./output/reflexion_phase0/aops_results.jsonl') - p.add_argument('--store-init-chars', type=int, default=8000, - help='Truncate the stored failed-attempt text to this many chars.') - p.add_argument('--store-rollout-text', action='store_true', - help='Also store (truncated) text of every rollout, not just stats.') - p.add_argument('--store-rollout-chars', type=int, default=2000) - args = p.parse_args() - - records = load_aops(n=args.n, seed=args.seed) - sys.stderr.write(f'[reflexion] {len(records)} AoPS problems, chunk={args.chunk_size}, ' - f'init_samples={args.init_samples}, n_skills={args.n_skills}, ' - f'pass_k={args.pass_k}, max_tokens={args.max_tokens}\n') - - # --- 8-GPU vLLM sampler (tp=GEN_TP, dp=GEN_GPUS/GEN_TP). --- - if GEN_GPUS % GEN_TP != 0: - raise ValueError(f'GEN_GPUS ({GEN_GPUS}) must be divisible by GEN_TP ({GEN_TP})') - gen_dp = GEN_GPUS // GEN_TP - gen_mesh = DeviceMesh.from_sizes(world_size=GEN_GPUS, dp_size=gen_dp, tp_size=GEN_TP) - twinkle.initialize( - mode='ray', nproc_per_node=GEN_GPUS, - groups=[DeviceGroup(name='sampler', ranks=list(range(GEN_GPUS)), - device_type='GPU', gpus_per_worker=GEN_TP)], - lazy_collect=False) - sampler = vLLMSampler( - model_id=GEN_MODEL_ID, - engine_args={'gpu_memory_utilization': GEN_GPU_MEM, - 'max_model_len': args.max_model_len, - 'tensor_parallel_size': GEN_TP}, - device_mesh=gen_mesh, remote_group='sampler') - sampler.set_template('Qwen3_5Template', model_id=GEN_MODEL_ID, - enable_thinking=True, max_length=args.max_model_len) - sys.stderr.write(f'[reflexion] sampler ready (model={GEN_MODEL_ID}, tp={GEN_TP}, ' - f'dp={gen_dp})\n') - - # --- API leak judge (optional): reuses eval_gpqa_rag's COMPRESS_* env. --- - api: Optional[OpenAIClient] = None - if LEAK_API_KEY and not args.disable_api_leak: - api = OpenAIClient(model=LEAK_API_MODEL, api_key=LEAK_API_KEY, - base_url=LEAK_BASE_URL) - sys.stderr.write(f'[reflexion] leak judge ON via API model={LEAK_API_MODEL} ' - f'(concurrency={args.api_concurrency})\n') - else: - sys.stderr.write('[reflexion] leak judge OFF (string filter only) — set ' - 'COMPRESS_API_KEY to enable the API judge\n') - - os.makedirs(os.path.dirname(args.output) or '.', exist_ok=True) - summ = {k: 0 for k in ('n_total', 'n_solved_first', 'n_failed', 'n_failed_easy', - 'n_hard', 'n_with_skill', 'n_helped', 'n_rescued')} - summ.update({'sum_baseline_pass': 0.0, 'sum_best_with_pass': 0.0, - 'sum_best_marginal': 0.0}) - - with open(args.output, 'w', encoding='utf-8') as out_f: - # Line 1: run config, for reproducibility / later analysis. - out_f.write(json.dumps({ - 'record_type': 'config', 'model': GEN_MODEL_ID, 'dataset': 'aops', - 'n': len(records), 'seed': args.seed, 'init_samples': args.init_samples, - 'n_skills': args.n_skills, 'pass_k': args.pass_k, - 'hard_baseline_max': args.hard_baseline_max, - 'max_model_len': args.max_model_len, 'max_tokens': args.max_tokens, - 'skill_max_tokens': args.skill_max_tokens, - 'api_leak_judge': api is not None, - 'api_leak_model': LEAK_API_MODEL if api is not None else None, - 'min_survivors': args.min_survivors, 'skill_retries': args.skill_retries, - 'gpus': GEN_GPUS, 'tp': GEN_TP, 'started': int(time.time()), - }, ensure_ascii=False) + '\n') - out_f.flush() - - n_chunks = (len(records) + args.chunk_size - 1) // args.chunk_size - for ci in range(n_chunks): - chunk = records[ci * args.chunk_size:(ci + 1) * args.chunk_size] - sys.stderr.write(f'[reflexion] chunk {ci+1}/{n_chunks} ({len(chunk)} problems)\n') - recs = process_chunk(sampler, api, chunk, gen_dp, args) - for rec in recs: # incremental write per problem - out_f.write(json.dumps(rec, ensure_ascii=False) + '\n') - out_f.flush() - _update_summary(summ, recs) - rep = _summary_report(summ) - sys.stderr.write( - f' running: failed={rep["n_failed_first_try"]} ' - f'(easy={rep["n_failed_but_easy"]}) hard={rep["n_hard"]} ' - f'base_pass={rep["avg_baseline_pass_on_hard"]:.3f} ' - f'skill_pass={rep["avg_best_with_skill_pass_on_hard"]:.3f} ' - f'helped={rep["n_helped_by_skill"]} rescued={rep["n_rescued_from_zero"]}\n') - - report = _summary_report(summ) - out_f.write(json.dumps(report, ensure_ascii=False) + '\n') - out_f.flush() - - print('\n' + '=' * 60) - print(f'Reflexion Phase-0 — model={GEN_MODEL_ID}, dataset=aops, n={report["n_total"]}') - print('=' * 60) - print(f'solved on first try : {report["n_solved_first_try"]}/{report["n_total"]}') - print(f'failed first try : {report["n_failed_first_try"]} ' - f'(easy, excluded: {report["n_failed_but_easy"]})') - print(f'hard (baseline pass@{args.pass_k}<= {args.hard_baseline_max}) : {report["n_hard"]}') - print(f' avg baseline pass@{args.pass_k:<2} : {report["avg_baseline_pass_on_hard"]:.4f}') - print(f' avg best-skill pass@{args.pass_k:<2} : {report["avg_best_with_skill_pass_on_hard"]:.4f}') - print(f' avg best marginal : {report["avg_best_marginal_on_hard"]:+.4f}') - print(f' helped by a skill : {report["n_helped_by_skill"]}/{report["n_hard"]}') - print(f' rescued from 0 pass : {report["n_rescued_from_zero"]}/{report["n_hard"]}') - print(f'\n[output] {args.output}') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/embedding/make_embedding_dataset.py b/cookbook/exp/embedding/make_embedding_dataset.py deleted file mode 100644 index 847f222fc..000000000 --- a/cookbook/exp/embedding/make_embedding_dataset.py +++ /dev/null @@ -1,758 +0,0 @@ -"""Offline compression pipeline: raw datasets → condenser → pre-compressed embedding dataset. - -Loads think/index/hard datasets, compresses query/cot/negatives via vLLM condenser -with API fallback, saves a single HF Dataset ready for embedding training. - -Output schema: {anchor_text, positive_text, negative_texts, source} - -Launch (8 GPUs — 4 for vLLM condenser): - python cookbook/exp/embedding/make_embedding_dataset.py -""" -import json -import os -import re -import sys -import threading -import time -from concurrent.futures import ThreadPoolExecutor, as_completed -from pathlib import Path -from typing import Any, Dict, List, Optional - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.data_format import SamplingParams -from twinkle.sampler import vLLMSampler -from twinkle.template import Qwen3_5Template -from twinkle.utils.parallel import PosixFileLock -from twinkle_agentic.protocol.openai import OpenAI as OpenAIClient - -sys.path.insert(0, str(Path(__file__).resolve().parent)) -from dataset_think import get_dataset as get_dataset_think # noqa: E402 -from dataset_index import get_dataset as get_dataset_index # noqa: E402 -from dataset_hard import get_dataset as get_dataset_hard # noqa: E402 - -logger = get_logger() - -# -- Model config ------------------------------------------------------------- -CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') -TEMPLATE_NAME = 'Qwen3_5Template' - -# -- GPU placement (condenser only) ------------------------------------------- -CONDENSER_GPUS = int(os.environ.get('CONDENSER_GPUS', 8)) - -# -- Dataset caps ------------------------------------------------------------- -TOTAL_SAMPLES: Optional[int] = None -THINK_CAP: Optional[int] = int(os.environ.get('THINK_CAP', 100_000)) -INDEX_CAP: Optional[int] = int(os.environ.get('INDEX_CAP', 100_000)) -HARD_CAP: Optional[int] = int(os.environ.get('HARD_CAP', 0)) or None -HARD_MAX_NEGATIVES = int(os.environ.get('HARD_MAX_NEGATIVES', 8)) - -# -- Compression params ------------------------------------------------------- -MIN_TEXT_CHARS = 256 -DATASET_MAX_TOKENS = 32768 -COMPRESS_TEMPERATURE = 0.2 -COMPRESS_TOP_P = 0.5 -COMPRESS_MAX_MODEL_LEN = 32768 -BATCH_SIZE = int(os.environ.get('COMPRESS_BATCH_SIZE', 128)) - -# -- API fallback ------------------------------------------------------------- -COMPRESS_API_KEY = os.environ.get('COMPRESS_API_KEY', '') -COMPRESS_BASE_URL = os.environ.get('COMPRESS_BASE_URL', 'https://dashscope.aliyuncs.com/compatible-mode/v1') -COMPRESS_MODEL = os.environ.get('COMPRESS_MODEL', 'qwen3.7-max') -API_MIN_INTERVAL = float(os.environ.get('API_MIN_INTERVAL', 0.1)) -API_CONCURRENCY = int(os.environ.get('API_CONCURRENCY', 24)) -SAMPLER_TIMEOUT = float(os.environ.get('SAMPLER_TIMEOUT', 300)) - -# -- Output ------------------------------------------------------------------- -OUTPUT_DIR = os.environ.get('EMB_DATASET_OUTPUT', './output/embedding_dataset') -RESULTS_JSONL = f'{OUTPUT_DIR}/results.jsonl' -PROGRESS_FILE = f'{OUTPUT_DIR}/progress.json' - -# ============================================================================= -# Prompts -# ============================================================================= - -COMPRESS_SYSTEM = """\ -You are a compression and summary assistant. For the (query, source) pair, emit a Markdown \ -answer with TWO sections, designed to pair with the `extract_compressed` tool: \ -the reader absorbs `## Summary` directly, then calls `extract_compressed` \ -on any topic-key listed under `## More` to recover its \ -fuller content. - - `## Summary` — extreme-density text the reader reads directly. - `## More` — a topic index whose keys are valid arguments \ -to `extract_compressed` for recovering material not captured inline. - -Together the two sections must form a COMPLETE, NON-DISTORTING inventory of the \ -source for the query — nothing essential lost, nothing implied that the source \ -does not support. NO preamble, NO meta-commentary, NO code fences wrapping the \ -whole output. - -Output skeleton: - -## Summary -Topic: - - -## More -- : -- ... - -Format selection for the inline body (pick the MOST COMPACT form per query, mix \ -when helpful): -- Interface / signature → code notation directly: `func(a:int)->str` -- Factual / entity → telegraphic prose; drop function words; ":" for "is", "," \ -for "has" -- Skill / how-to / usage → lead with `Use when: `; numbered telegraphic \ -steps `1.do X 2.then Y`; close with `Output: ` when relevant -- Procedural → numbered short steps -- Analytical / design → hierarchical bullets with abbreviations - -`## Summary` rules: -1. TOPIC LINE — line 1 is ALWAYS `Topic: `, even when the \ -query is narrow. Anchors both the reader and the tool. -2. DENSITY — every token in the body carries query-relevant signal; cut filler. -3. PRIMARY-COMPLETE — never silently drop a fact essential to answering the \ -query. Anything cut for length MUST appear as a key under \ -`## More`. -4. NON-MISLEADING — phrasing must not let the reader infer anything the source \ -does not support; partial truths that mislead are worse than honest omissions \ -flagged in the index. -5. SELF-CONTAINED — the reader can act on the answer without re-opening the source. -6. FAITHFUL — only content the source supports; no fabrication, no extrapolation. -7. LANGUAGE — match the source language. -8. NO outer code fences around the whole answer; no meta-commentary. - -`## More` rules (MANDATORY — this section is never omitted): -1. FORMAT — each bullet is `- : `: - • topic-key — short, unambiguous, grounded in source vocabulary so the \ -`extract_compressed` tool can locate the aspect (e.g. `decorators`, \ -`error handling`, `pitfalls`). - • hint — tells WHAT the reader gains by expanding (concrete numbers, code \ -listings, secondary cases, edge details, related context, …); do NOT restate \ -the inline answer. -2. CRITERION — each bullet names an aspect that EXISTS in the source but is \ -NOT fully captured inline. Material that genuinely fits inline without \ -distortion MUST NOT be duplicated here. -3. FAITHFUL — hints must be grounded in the source; never speculate or invent. -4. ORDER — by relevance to the query, then by importance. -5. EMPTY CASE — if the source is so short / single-purpose that everything \ -fits inline, write a single line `- (none)`. - -Now begin.\ -""" - -COMPRESS_USER = ( - 'Downstream model will read your compressed block to decide whether to ' - 'expand it. Compress faithfully: preserve the passage topic + core facts. ' - 'Do NOT invent facts. Do NOT drop major facts. Do NOT write meta-commentary ' - 'about the Query (never write "Query info: absent", "no X mention", etc.); ' - 'if the passage does not address the Query, still summarize the passage. ' - 'CRITICAL LANGUAGE RULE: detect the dominant language of the Passage ' - '(NOT the Query, NOT this instruction) and write the ENTIRE output in that ' - 'same language; English passage → English output, Chinese passage → ' - 'Chinese output, Japanese passage → Japanese output. NEVER translate, ' - 'NEVER mix languages, NEVER copy these instructions into the output.\n\n' - '## Query (ordering hint only — still summarize the whole passage)\n{query}\n\n' - '## Passage\n{text}') - -EMBED_QUERY_Q = ( - 'Summarize this query for retrieval. ' - 'The body of ## Summary MUST follow this EXACT 4-line template — ' - 'do NOT emit "Use when:", numbered procedure steps, or "Output:":\n' - 'Topic: \n' - 'Problem: \n' - 'Skill: \n' - 'Knowledge: \n' - 'Then emit the mandatory ## More section as usual. ' - 'Topic must name the specific pattern, never generic labels.') - -EMBED_QUERY_COT = ( - 'Summarize this reasoning trace for retrieval. ' - 'The body of ## Summary MUST follow this EXACT 4-line template — ' - 'do NOT emit "Use when:", numbered procedure steps, or "Output:":\n' - 'Topic: \n' - 'Problem: \n' - 'Skill: \n' - 'Knowledge: \n' - 'Then emit the mandatory ## More section as usual. ' - 'Topic must name the specific pattern, never generic labels.') - -EMBED_QUERY_Q_LEGACY = ( - 'What problem does this passage address, and what skill or method is needed? ' - 'Topic must name the specific pattern, never generic labels. ' - 'Compress into a retrieval-friendly need description.') - -EMBED_QUERY_COT_LEGACY = ( - 'Extract the reusable skill: trigger conditions, key steps, and expected output. ' - 'Topic names the method/pattern; format as "Use when: ...", numbered steps, ' - '"Output: ...". Compress into a standardized procedure for retrieval.') - -EMBED_QUERY_REASONIR_Q = ( - 'Extract the abstract PROBLEM TYPE from this query. ' - 'IGNORE all specific numbers, values, variable names, and parameters — ' - 'focus ONLY on the CLASS of problem and the METHODOLOGY required. ' - 'The body of ## Summary MUST follow this EXACT 4-line template:\n' - 'Topic: \n' - 'Problem: \n' - 'Skill: \n' - 'Knowledge: \n' - 'Then emit the mandatory ## More section as usual. ' - 'Topic must name the method class, never mention specific numbers.') - -EMBED_QUERY_REASONIR_COT = ( - 'Extract the abstract METHODOLOGY demonstrated in this solution. ' - 'IGNORE all specific numbers, values, and computed results — ' - 'focus ONLY on the general TECHNIQUE and key reasoning STEPS. ' - 'The body of ## Summary MUST follow this EXACT 4-line template:\n' - 'Topic: \n' - 'Problem: \n' - 'Skill: \n' - 'Knowledge: \n' - 'Then emit the mandatory ## More section as usual. ' - 'Topic must name the method class, never mention specific numbers.') - - -# ============================================================================= -# Validation & API fallback -# ============================================================================= - -_LEGACY_USE_WHEN_RE = re.compile(r'(?im)^\s*Use when\s*:') -_SCHEMA_MARKERS = ('Problem:', 'Skill:', 'Knowledge:') - - -def _is_truncated_compression(text: str, schema: str = 'new') -> bool: - if not text or not text.strip(): - return True - if '## More' not in text or '## Summary' not in text: - return True - after_more = text.split('## More', 1)[1].strip() - if not after_more: - return True - last_line = after_more.splitlines()[-1].strip() - if not (last_line.startswith('-') or last_line.endswith(')')): - return True - if schema == 'new': - summary_body = text.split('## Summary', 1)[1].split('## More', 1)[0] - if _LEGACY_USE_WHEN_RE.search(summary_body): - return True - if not all(marker in summary_body for marker in _SCHEMA_MARKERS): - return True - return False - - -_api_semaphore = threading.Semaphore(API_CONCURRENCY) -_api_bucket_lock = threading.Lock() -_api_tokens = [float(API_CONCURRENCY)] -_api_last_refill = [time.monotonic()] - - -def _api_throttle(): - """Token-bucket rate limiter: API_CONCURRENCY requests per API_MIN_INTERVAL*API_CONCURRENCY window.""" - _api_semaphore.acquire() - try: - with _api_bucket_lock: - now = time.monotonic() - elapsed = now - _api_last_refill[0] - refill = elapsed / API_MIN_INTERVAL - _api_tokens[0] = min(float(API_CONCURRENCY), _api_tokens[0] + refill) - _api_last_refill[0] = now - if _api_tokens[0] >= 1.0: - _api_tokens[0] -= 1.0 - else: - wait = (1.0 - _api_tokens[0]) * API_MIN_INTERVAL - _api_tokens[0] = 0.0 - time.sleep(wait) - finally: - _api_semaphore.release() - - -def _api_compress(api_client: OpenAIClient, prompt: Dict[str, Any]) -> Optional[str]: - _api_throttle() - trajectory = {'messages': prompt['messages']} - sp = SamplingParams(temperature=0.2, max_tokens=8192) - try: - reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) - except Exception as exc: - logger.warning(f'[api_fallback] error: {exc}') - return None - content = (reply.get('content') or '').strip() - if not content: - return None - m = re.match(r'^```[a-zA-Z]*\n(.*?)\n```\s*$', content, re.DOTALL) - if m: - content = m.group(1).strip() - return content - - -# ============================================================================= -# Core compression logic -# ============================================================================= - -def _extract_query_cot(row: Dict[str, Any]): - messages = row.get('messages') or [] - query, cot = '', '' - for m in messages: - if not isinstance(m, dict): - continue - role = m.get('role') or '' - if role == 'user' and not query: - query = (m.get('content') or '').strip() - elif role == 'assistant': - cot = (m.get('reasoning_content') or '').strip() - break - return query, cot - - -def _compress_batch_phase1( - rows: List[Dict[str, Any]], - condenser_sampler, - compress_params: SamplingParams, - special_tokens: set, - source_type: str, -) -> Optional[Dict[str, Any]]: - """Phase 1 (GPU): build prompts → vLLM sample → validate. Returns state for phase 2.""" - _MAX_COT_CHARS = 30_000 - - if source_type == 'hard': - return _compress_hard_phase1(rows, condenser_sampler, compress_params, - special_tokens, source_type) - - prompts: List[Optional[Dict[str, Any]]] = [] - meta: List[Dict[str, Any]] = [] - for i, row in enumerate(rows): - query, cot = _extract_query_cot(row) - if not query or len(cot) < MIN_TEXT_CHARS or len(cot) > _MAX_COT_CHARS: - continue - schema = 'legacy' if (i % 2 == 0) else 'new' - q_hint = EMBED_QUERY_Q_LEGACY if schema == 'legacy' else EMBED_QUERY_Q - c_hint = EMBED_QUERY_COT_LEGACY if schema == 'legacy' else EMBED_QUERY_COT - - if len(query) < MIN_TEXT_CHARS: - prompts.append(None) - else: - user = COMPRESS_USER.format(query=q_hint, text=query) - prompts.append({'messages': [ - {'role': 'system', 'content': COMPRESS_SYSTEM}, - {'role': 'user', 'content': user}, - ]}) - user_c = COMPRESS_USER.format(query=c_hint, text=cot) - prompts.append({'messages': [ - {'role': 'system', 'content': COMPRESS_SYSTEM}, - {'role': 'user', 'content': user_c}, - ]}) - meta.append({'query_raw': query, 'cot_raw': cot, 'schema': schema, - 'q_hint': q_hint, 'source': source_type, - 'row_id': row.get('id', str(i))}) - - if not prompts: - return {'final': []} - - sampler_input = [p for p in prompts if p is not None] - sampler_pos = [ri for ri, p in enumerate(prompts) if p is not None] - try: - sampler_responses = condenser_sampler.sample(sampler_input, compress_params) - except Exception as exc: - logger.warning(f'[compress] sampler error: {exc}') - sampler_responses = [None] * len(sampler_input) - - responses = [None] * len(prompts) - for resp, pos in zip(sampler_responses, sampler_pos): - responses[pos] = resp - - decoded: List[str] = [] - fallback_indices: List[int] = [] - for ri in range(len(prompts)): - pair_idx = ri // 2 - schema = meta[pair_idx]['schema'] - if prompts[ri] is None: - decoded.append(meta[pair_idx]['query_raw']) - continue - resp = responses[ri] - seq = resp.sequences[0] if resp and resp.sequences else None - text = '' - if seq and seq.stop_reason != 'length' and seq.decoded: - text = seq.decoded - for tok in special_tokens: - text = text.replace(tok, '') - text = text.rstrip() - if not _is_truncated_compression(text, schema): - decoded.append(text) - else: - decoded.append('') - fallback_indices.append(ri) - - return {'prompts': prompts, 'meta': meta, 'decoded': decoded, - 'fallback_indices': fallback_indices} - - -def _compress_batch_phase2( - state: Dict[str, Any], - api_client: OpenAIClient, -) -> List[Dict[str, Any]]: - """Phase 2 (no GPU): API fallback → build results.""" - if 'final' in state: - return state['final'] - - prompts = state['prompts'] - decoded = state['decoded'] - fallback_indices = state['fallback_indices'] - is_hard = state.get('hard', False) - meta = state.get('meta') # None for hard - - # Track which prompts used API fallback - api_set: set = set() - if fallback_indices: - api_futures = {} - with ThreadPoolExecutor(max_workers=API_CONCURRENCY) as pool: - for ri in fallback_indices: - api_futures[pool.submit(_api_compress, api_client, prompts[ri])] = ri - for fut in as_completed(api_futures): - ri = api_futures[fut] - api_result = fut.result() - schema = 'new' if is_hard else meta[ri // 2]['schema'] - if api_result and not _is_truncated_compression(api_result, schema): - decoded[ri] = api_result - api_set.add(ri) - - state['api_set'] = api_set - if is_hard: - return _build_hard_results(state) - return _build_think_index_results(state) - - -def _build_think_index_results(state: Dict[str, Any]) -> List[Dict[str, Any]]: - meta = state['meta'] - decoded = state['decoded'] - api_set = state.get('api_set', set()) - results = [] - for pair_idx in range(len(meta)): - q_text = decoded[pair_idx * 2] - c_text = decoded[pair_idx * 2 + 1] - if not q_text or not c_text: - continue - q_method = 'api' if (pair_idx * 2) in api_set else 'vllm' - c_method = 'api' if (pair_idx * 2 + 1) in api_set else 'vllm' - results.append({ - 'anchor_text': q_text, - 'positive_text': c_text, - 'negative_texts': [], - 'source': meta[pair_idx]['source'], - 'query_raw': meta[pair_idx]['query_raw'], - 'cot_raw': meta[pair_idx]['cot_raw'], - 'anchor_method': q_method, - 'positive_method': c_method, - }) - return results - - -def _build_hard_results(state: Dict[str, Any]) -> List[Dict[str, Any]]: - group_sizes = state['group_sizes'] - decoded = state['decoded'] - source_type = state['source_type'] - raw_groups = state['raw_groups'] - api_set = state.get('api_set', set()) - results = [] - offset = 0 - for gi, gs in enumerate(group_sizes): - q_text = decoded[offset] - c_text = decoded[offset + 1] - if not q_text or not c_text: - offset += gs - continue - neg_texts = [] - neg_raws = [] - neg_methods = [] - for ni in range(2, gs): - nt = decoded[offset + ni] - if nt: - neg_texts.append(nt) - neg_raws.append(raw_groups[gi]['negs_raw'][ni - 2]) - neg_methods.append('api' if (offset + ni) in api_set else 'vllm') - q_method = 'api' if offset in api_set else 'vllm' - c_method = 'api' if (offset + 1) in api_set else 'vllm' - results.append({ - 'anchor_text': q_text, - 'positive_text': c_text, - 'negative_texts': neg_texts, - 'source': source_type, - 'query_raw': raw_groups[gi]['query_raw'], - 'cot_raw': raw_groups[gi]['cot_raw'], - 'negs_raw': neg_raws, - 'anchor_method': q_method, - 'positive_method': c_method, - 'neg_methods': neg_methods, - }) - offset += gs - return results - - -def _compress_hard_phase1( - rows: List[Dict[str, Any]], - condenser_sampler, - compress_params: SamplingParams, - special_tokens: set, - source_type: str, -) -> Dict[str, Any]: - """Phase 1 for hard rows: vLLM sample + validate. Returns state for phase 2.""" - _MAX_COT_CHARS = 30_000 - - prompts: List[Dict[str, Any]] = [] - group_sizes: List[int] = [] - row_ids: List[str] = [] - raw_groups: List[Dict[str, Any]] = [] - - for row in rows: - query, cot = _extract_query_cot(row) - if not query or not cot or len(cot) > _MAX_COT_CHARS: - continue - negatives = row.get('negatives') or [] - valid_negs = [n for n in negatives - if n and len(n) <= _MAX_COT_CHARS] - - user_q = COMPRESS_USER.format(query=EMBED_QUERY_REASONIR_Q, text=query) - prompts.append({'messages': [ - {'role': 'system', 'content': COMPRESS_SYSTEM}, - {'role': 'user', 'content': user_q}, - ]}) - user_c = COMPRESS_USER.format(query=EMBED_QUERY_REASONIR_COT, text=cot) - prompts.append({'messages': [ - {'role': 'system', 'content': COMPRESS_SYSTEM}, - {'role': 'user', 'content': user_c}, - ]}) - for neg in valid_negs: - user_n = COMPRESS_USER.format(query=EMBED_QUERY_REASONIR_COT, text=neg) - prompts.append({'messages': [ - {'role': 'system', 'content': COMPRESS_SYSTEM}, - {'role': 'user', 'content': user_n}, - ]}) - group_sizes.append(2 + len(valid_negs)) - row_ids.append(row.get('id', '')) - raw_groups.append({'query_raw': query, 'cot_raw': cot, 'negs_raw': valid_negs}) - - if not prompts: - return {'hard': True, 'prompts': [], 'group_sizes': [], 'row_ids': [], - 'decoded': [], 'fallback_indices': [], 'source_type': source_type, - 'raw_groups': []} - - try: - responses = condenser_sampler.sample(prompts, compress_params) - except Exception as exc: - logger.warning(f'[compress-hard] sampler error: {exc}') - responses = [None] * len(prompts) - - decoded: List[str] = [] - fallback_indices: List[int] = [] - for ri, resp in enumerate(responses): - seq = resp.sequences[0] if resp and resp.sequences else None - text = '' - if seq and seq.stop_reason != 'length' and seq.decoded: - text = seq.decoded - for tok in special_tokens: - text = text.replace(tok, '') - text = text.rstrip() - if text and not _is_truncated_compression(text, 'new'): - decoded.append(text) - else: - decoded.append('') - fallback_indices.append(ri) - - return {'hard': True, 'prompts': prompts, 'group_sizes': group_sizes, - 'row_ids': row_ids, 'decoded': decoded, - 'fallback_indices': fallback_indices, 'source_type': source_type, - 'raw_groups': raw_groups} - - -# ============================================================================= -# Main pipeline -# ============================================================================= - -def main(): - device_groups = [ - DeviceGroup(name='condenser_sampler', - ranks=list(range(CONDENSER_GPUS)), - device_type='GPU'), - ] - condenser_mesh = DeviceMesh.from_sizes( - world_size=CONDENSER_GPUS, dp_size=CONDENSER_GPUS) - twinkle.initialize(mode='ray', nproc_per_node=CONDENSER_GPUS, groups=device_groups) - - # -- Load raw datasets ---------------------------------------------------- - from datasets import Dataset as HFDataset - - dataset_think = get_dataset_think(total=TOTAL_SAMPLES, load_from_cache_file=True) - if THINK_CAP and len(dataset_think.dataset) > THINK_CAP: - dataset_think.dataset = dataset_think.dataset.select(range(THINK_CAP)) - ds_think = dataset_think.dataset - logger.info(f'[load] think={len(ds_think)}') - - ds_index_obj = get_dataset_index(total=None, load_from_cache_file=True) - ds_index = ds_index_obj.dataset - if INDEX_CAP and len(ds_index) > INDEX_CAP: - ds_index = ds_index.select(range(INDEX_CAP)) - logger.info(f'[load] index={len(ds_index)}') - - ds_hard_raw = get_dataset_hard(max_negatives=HARD_MAX_NEGATIVES, load_from_cache_file=True) - if HARD_CAP and len(ds_hard_raw) > HARD_CAP: - ds_hard_raw = ds_hard_raw.select(range(HARD_CAP)) - n_hard = len(ds_hard_raw) - logger.info(f'[load] hard={n_hard}') - - # Convert hard to messages schema - hard_rows_list = [] - if n_hard > 0: - h_ids = ds_hard_raw['id'] - h_queries = ds_hard_raw['query'] - h_cots = ds_hard_raw['cot'] - h_responses = ds_hard_raw['response'] if 'response' in ds_hard_raw.column_names else [''] * n_hard - h_negatives = ds_hard_raw['negatives'] - for i in range(n_hard): - hard_rows_list.append({ - 'id': h_ids[i], - 'messages': [ - {'role': 'user', 'content': h_queries[i]}, - {'role': 'assistant', 'reasoning_content': h_cots[i], - 'content': h_responses[i] or ''}, - ], - 'negatives': h_negatives[i], - }) - - # Batch-convert HF Datasets to list-of-dicts - def _ds_to_rows(ds): - return [dict(zip(ds.column_names, vals)) for vals in zip(*(ds[c] for c in ds.column_names))] - - think_rows = _ds_to_rows(ds_think) - index_rows = _ds_to_rows(ds_index) - - # -- Setup condenser ------------------------------------------------------ - condenser_template = Qwen3_5Template( - model_id=CONDENSE_MODEL_ID, max_length=DATASET_MAX_TOKENS, - enable_thinking=False, truncation_strategy='delete') - special_tokens = set(condenser_template.tokenizer.all_special_tokens) - - condenser_sampler = vLLMSampler( - model_id=CONDENSE_MODEL_ID, - engine_args={'gpu_memory_utilization': 0.8, 'max_model_len': COMPRESS_MAX_MODEL_LEN}, - device_mesh=condenser_mesh, - remote_group='condenser_sampler', - ) - condenser_sampler.set_template( - TEMPLATE_NAME, model_id=CONDENSE_MODEL_ID, enable_thinking=False, - truncation_strategy='delete', max_length=DATASET_MAX_TOKENS) - condenser_sampler._ray_get_timeout = SAMPLER_TIMEOUT - compress_params = SamplingParams( - max_tokens=8192, temperature=COMPRESS_TEMPERATURE, - top_p=COMPRESS_TOP_P, num_samples=1) - - api_client = OpenAIClient( - model=COMPRESS_MODEL, api_key=COMPRESS_API_KEY, base_url=COMPRESS_BASE_URL) - - # -- Resume support ---------------------------------------------------------- - os.makedirs(OUTPUT_DIR, exist_ok=True) - progress = {'think': 0, 'index': 0, 'hard': 0} - if os.path.exists(PROGRESS_FILE): - with open(PROGRESS_FILE, 'r') as f: - progress = json.load(f) - logger.info(f'[resume] loaded progress: {progress}') - - _results_lock = PosixFileLock(RESULTS_JSONL + '.lock') - - def _flush_results(records: List[Dict[str, Any]]): - if not records: - return - lines = [json.dumps(r, ensure_ascii=False) + '\n' for r in records] - with _results_lock: - with open(RESULTS_JSONL, 'a', encoding='utf-8') as f: - f.writelines(lines) - - def _save_progress(): - tmp = PROGRESS_FILE + '.tmp' - with open(tmp, 'w') as f: - json.dump(progress, f) - os.replace(tmp, PROGRESS_FILE) - - # -- Process in batches (pipelined: vLLM batch N+1 overlaps API fallback N) - - total_flushed = 0 - if os.path.exists(RESULTS_JSONL): - with open(RESULTS_JSONL, 'r', encoding='utf-8') as f: - total_flushed = sum(1 for l in f if l.strip()) - if total_flushed: - logger.info(f'[resume] {total_flushed} records already in results.jsonl') - - def _process_source(rows, source_type, label): - nonlocal total_flushed - n_total = len(rows) - skip = progress.get(source_type, 0) - if skip >= n_total: - logger.info(f'[{label}] skipped (already done {skip}/{n_total})') - return - if skip > 0: - logger.info(f'[{label}] resuming from row {skip}/{n_total}') - - bg_pool = ThreadPoolExecutor(max_workers=1) - pending = None # (future, batch_start, batch_len) - - def _drain_pending(): - nonlocal total_flushed, pending - if pending is None: - return - fut, p_start, p_len = pending - batch_results = fut.result() - _flush_results(batch_results) - total_flushed += len(batch_results) - progress[source_type] = p_start + p_len - _save_progress() - pending = None - - for start in range(skip, n_total, BATCH_SIZE): - batch = rows[start:start + BATCH_SIZE] - state = _compress_batch_phase1( - batch, condenser_sampler, compress_params, - special_tokens, source_type) - _drain_pending() - pending = ( - bg_pool.submit(_compress_batch_phase2, state, api_client), - start, len(batch)) - n_done = start + len(batch) - if n_done % (BATCH_SIZE * 10) == 0 or n_done >= n_total: - logger.info(f'[{label}] {n_done}/{n_total} vLLM done, ' - f'{total_flushed} records flushed (last batch pending)') - - _drain_pending() - bg_pool.shutdown(wait=False) - logger.info(f'[{label}] complete, {total_flushed} total records flushed') - - _process_source(hard_rows_list, 'hard', 'hard') - _process_source(think_rows, 'think', 'think') - _process_source(index_rows, 'index', 'index') - - # -- Convert JSONL → HF Dataset ------------------------------------------- - logger.info(f'[save] converting results.jsonl to HF Dataset...') - all_results = [] - with open(RESULTS_JSONL, 'r', encoding='utf-8') as f: - for line_no, line in enumerate(f, 1): - if not line.strip(): - continue - try: - all_results.append(json.loads(line)) - except json.JSONDecodeError: - logger.warning(f'[save] skipping malformed line {line_no} (truncated resume?)') - logger.info(f'[save] total records: {len(all_results)}') - out_ds = HFDataset.from_dict({ - 'anchor_text': [r['anchor_text'] for r in all_results], - 'positive_text': [r['positive_text'] for r in all_results], - 'negative_texts': [r['negative_texts'] for r in all_results], - 'source': [r['source'] for r in all_results], - 'query_raw': [r.get('query_raw', '') for r in all_results], - 'cot_raw': [r.get('cot_raw', '') for r in all_results], - 'negs_raw': [r.get('negs_raw', []) for r in all_results], - }) - out_ds.save_to_disk(OUTPUT_DIR + '/dataset') - logger.info(f'[save] dataset saved to {OUTPUT_DIR}/dataset') - logger.info(f'[stats] think={sum(1 for r in all_results if r["source"]=="think")} ' - f'index={sum(1 for r in all_results if r["source"]=="index")} ' - f'hard={sum(1 for r in all_results if r["source"]=="hard")}') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/embedding/train_embedding_full_ddp.py b/cookbook/exp/embedding/train_embedding_full_ddp.py deleted file mode 100644 index 97ab3b128..000000000 --- a/cookbook/exp/embedding/train_embedding_full_ddp.py +++ /dev/null @@ -1,270 +0,0 @@ -"""Full-parameter embedding training on pre-compressed dataset. - -Reads the pre-compressed HF Dataset produced by make_embedding_dataset.py, -encodes features, trains with InfoNCE loss. - -Architecture (4 GPUs): - - Ranks 0-3: Trainable embedding model, InfoNCE loss. - -Launch: - python cookbook/exp/embedding/train_embedding_full_ddp.py -""" -import os -import time -from typing import Any, Dict, List, Literal, Optional - -import swanlab - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_device_placement, get_logger -from twinkle.loss import InfonceLoss -from twinkle.metric import EmbeddingMetric -from twinkle.model import TransformersModel -from twinkle.processor import InputProcessor -from twinkle.template import Qwen3_5Template, Template - -logger = get_logger() - -# -- Backend selection -------------------------------------------------------- -BACKEND: Literal['transformers', 'megatron'] = 'transformers' - -MODEL_ID = os.environ.get('MODEL_ID', 'ms://Qwen/Qwen3.5-4B') - -# -- GPU placement ------------------------------------------------------------ -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', 64)) -LEARNING_RATE = 1e-5 -GRADIENT_ACCUMULATION_STEPS = 1 -LOG_INTERVAL = 2 -SAVE_INTERVAL = 2000 -NUM_EPOCHS = 1 - -# -- Dataset path (output of make_embedding_dataset.py) ----------------------- -DATASET_PATH = os.environ.get('EMB_DATASET_PATH', 'ms://twinkle-kit/qth-embedding') -MIX_SHUFFLE_SEED = 42 - -# -- Resume from checkpoint --------------------------------------------------- -RESUME_CHECKPOINT = os.environ.get('RESUME_CHECKPOINT', '') -RESUME_STEP = int(os.environ.get('RESUME_STEP', 0)) - -# -- Output ------------------------------------------------------------------- -OUTPUT_DIR = f'./output/embedding_full_{BACKEND}' - - -# ============================================================================= -# Model builders -# ============================================================================= - -def build_model(device_mesh: DeviceMesh): - model_id = RESUME_CHECKPOINT if RESUME_CHECKPOINT else MODEL_ID - if BACKEND == 'transformers': - model = TransformersModel( - model_id=model_id, - device_mesh=device_mesh, - remote_group='model', - ddp_config={'find_unused_parameters': True}, - ) - from twinkle.patch.no_split_modules import NoSplitModulesPatch - model.apply_patch(NoSplitModulesPatch({'Qwen3_5DecoderLayer'})) - return model - if BACKEND == 'megatron': - from twinkle.model import MegatronModel - return MegatronModel( - model_id=MODEL_ID, - device_mesh=device_mesh, - remote_group='model', - mixed_precision='bf16', - variable_seq_lengths=True, - ) - raise ValueError(f'Unknown BACKEND={BACKEND!r}') - - -def setup_optimizer(model, total_steps: int): - if BACKEND == 'transformers': - model.set_optimizer(optimizer_cls='AdamW', lr=LEARNING_RATE) - model.set_lr_scheduler( - scheduler_cls='CosineWarmupScheduler', - num_warmup_steps=200, - num_training_steps=total_steps, - ) - return - if BACKEND == 'megatron': - model.set_optimizer(optimizer_cls='default', lr=LEARNING_RATE) - model.set_lr_scheduler( - scheduler_cls='default', - lr_warmup_steps=50, - lr_decay_steps=total_steps, - ) - return - raise ValueError(f'Unknown BACKEND={BACKEND!r}') - - -def save_checkpoint(model, name: str): - model.save(name, output_dir=OUTPUT_DIR) - - -# ============================================================================= -# Feature encoding -# ============================================================================= - -def _get_first_feature(decoded_text: str, template: Template, role: str) -> Optional[Dict[str, Any]]: - if not decoded_text: - return None - if role == 'anchor': - feat = template.encode({'messages': [ - {'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 - - -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 - - -# ============================================================================= -# Main training -# ============================================================================= - -def train(): - device_groups = [ - DeviceGroup(name='model', - ranks=list(range(MODEL_GPUS)), - device_type='GPU'), - ] - model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, dp_size=MODEL_GPUS) - twinkle.initialize(mode='ray', nproc_per_node=MODEL_GPUS, groups=device_groups) - - # -- Load pre-compressed dataset ------------------------------------------ - from twinkle.dataset import Dataset as TwinkleDataset, DatasetMeta - logger.info(f'[data] loading pre-compressed dataset from {DATASET_PATH}') - dataset = TwinkleDataset(DatasetMeta(dataset_id=DATASET_PATH), download_mode='force_redownload') - dataset = dataset.dataset.shuffle(seed=MIX_SHUFFLE_SEED) - logger.info(f'[data] {len(dataset)} rows loaded') - - # -- Compute steps -------------------------------------------------------- - rows_per_step = BATCH_SIZE - total_steps = (len(dataset) // rows_per_step) * NUM_EPOCHS - optimizer_steps = total_steps // GRADIENT_ACCUMULATION_STEPS - - # -- Model ---------------------------------------------------------------- - model = build_model(model_mesh) - model.set_processor(InputProcessor) - model.set_loss(InfonceLoss, temperature=TEMPERATURE, use_batch=True, - hard_negatives=HARD_NEGATIVES) - setup_optimizer(model, optimizer_steps) - model.add_metric(EmbeddingMetric, is_training=True) - - 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 steps: {total_steps}, optimizer steps: {optimizer_steps}') - - swanlab.init(project='twinkle', config={ - 'backend': BACKEND, - 'model_id': MODEL_ID, - 'batch_size': BATCH_SIZE, - 'lr': LEARNING_RATE, - 'temperature': TEMPERATURE, - 'emb_max_length': EMB_MAX_LENGTH, - 'dataset_path': DATASET_PATH, - }) - - # -- Train loop ----------------------------------------------------------- - cur_step = 0 - _skip_rows = RESUME_STEP * rows_per_step # approximate rows to skip - - for epoch in range(NUM_EPOCHS): - for start in range(0, len(dataset), rows_per_step): - if start < _skip_rows: - continue - - 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() - features = _encode_batch(rows_list, emb_template) - t_encode = time.monotonic() - t0 - - if len(features) < 4: - continue - - t1 = time.monotonic() - model.forward_backward(inputs=features, task='embedding') - model.clip_grad_and_step( - gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) - t_train = time.monotonic() - t1 - cur_step += 1 - - if cur_step % LOG_INTERVAL == 0: - metric = model.calculate_metric(is_training=True) - logger.info( - f'Epoch {epoch} Step {cur_step}/{total_steps}, ' - f'metric: {metric} | ' - f'encode={t_encode:.2f}s train={t_train:.2f}s') - log_dict = {} - for k, v in metric.items(): - if not v: - continue - try: - log_dict[k] = float(v) - except (ValueError, TypeError): - pass - log_dict['epoch'] = epoch - log_dict['encode_sec'] = round(t_encode, 3) - log_dict['train_sec'] = round(t_train, 3) - swanlab.log(log_dict, step=cur_step) - if cur_step % SAVE_INTERVAL == 0: - save_checkpoint(model, f'step_{cur_step}') - - save_checkpoint(model, 'last-checkpoint') - # Force sync: resolve any pending lazy remote calls (save) before exit - model.calculate_metric(is_training=True) - logger.info(f'Training complete. Final step: {cur_step}') - - -if __name__ == '__main__': - train() diff --git a/cookbook/exp/embedding/train_reflexion_skill.py b/cookbook/exp/embedding/train_reflexion_skill.py deleted file mode 100644 index beccbc7b3..000000000 --- a/cookbook/exp/embedding/train_reflexion_skill.py +++ /dev/null @@ -1,1990 +0,0 @@ -"""Online GRPO RFT for the reflexion skill generator (self-contained, cached). - -Trains an INDEPENDENT skill model to write reusable skills that, injected into a -FROZEN base solver's system prompt, raise its accuracy. The base is never trained; -it only produces the reward. Per chunk: base greedy solve -> rubric process-check -(view A) -> skill-gen (thinking ON, N candidates) -> leak filter -> with-skill -greedy pass -> group-relative advantage -> ONE on-policy GRPO step -> sync weights. -Reward is deterministic (T=0, M=1) binary correctness; A = (R - mean) / (std + eps) -within each problem-group, so std=0 groups give no gradient (GRPO variance selects). - -Each problem is routed to EXACTLY ONE view: A = problem + rubric findings, B = -query only (deployment form). Skill-gen trains only the final structured guidance turn. - -Base greedy rollouts and rubric diagnoses are cached to disk (jsonl, md5-keyed) so -restarts skip them; skill-gen is on-policy and never cached. - -8 GPUs default for high-memory cards: rank 0 trains the actor, rank 1 hosts a -frozen reference model, ranks 2-3 skill_sampler (synced), ranks 4-7 base_sampler -(frozen). Override TRAIN_GPUS / REF_GPUS / SKILL_SAMPLER_GPUS / BASE_SAMPLER_GPUS -for other layouts. -Leak / rubric use the backup teacher API: set LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / -LLM_BACKUP_MODEL. - -Launch: - LLM_BACKUP_API_KEY=... python cookbook/exp/embedding/train_reflexion_skill.py \ - --dataset aops --n 5000 --chunk-size 16 --lr 6e-6 -""" -import argparse -import copy -import hashlib -import json -import math -import os -import re -import sys -import threading -import time -from concurrent.futures import ThreadPoolExecutor -from typing import Any, Dict, List, Optional, Set, Tuple - -import numpy as np - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.checkpoint_engine import CheckpointEngineManager -from twinkle.data_format import SamplingParams -from twinkle.dataset import Dataset, DatasetMeta -from twinkle.model import TransformersModel -from twinkle.patch.no_split_modules import NoSplitModulesPatch -from twinkle.processor import InputProcessor -from twinkle.sampler import vLLMSampler -from twinkle.template import Template -from twinkle_agentic.verifier import RubricVerifier -from twinkle_agentic.verifier.rubric_verifier import RubricItem - -logger = get_logger() - -try: - import swanlab -except ImportError: - swanlab = None - -MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') -GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.8)) -GEN_TEMPERATURE = float(os.environ.get('GEN_TEMPERATURE', 0.6)) -GEN_TOP_P = float(os.environ.get('GEN_TOP_P', 0.95)) -AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') -MATH_DATASET_ID = os.environ.get('MATH_DATASET_ID', 'modelscope/competition_math') - -# GPU layout: train actor (synced to skill_sampler) + frozen ref model + skill_sampler + base_sampler. -# High-memory cards can keep Qwen3-4B actor/ref as one GPU each and spend more GPUs -# on vLLM data-parallel sampling. The base side is heavier here because every -# clean skill candidate is re-solved by the frozen base model, plus balance/eval baselines. -TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 2)) -REF_GPUS = int(os.environ.get('REF_GPUS', 2)) -SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) -BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 2)) -NUM_GPUS = TRAIN_GPUS + REF_GPUS + SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS -TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', min(1, TRAIN_GPUS))) -REF_FSDP = int(os.environ.get('REF_FSDP', min(1, REF_GPUS))) -if min(TRAIN_GPUS, REF_GPUS, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS, TRAIN_FSDP, REF_FSDP) < 1: - raise ValueError('TRAIN_GPUS, REF_GPUS, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS, TRAIN_FSDP and REF_FSDP must all be >= 1') -if TRAIN_GPUS % TRAIN_FSDP != 0: - raise ValueError(f'TRAIN_GPUS ({TRAIN_GPUS}) must be divisible by TRAIN_FSDP ({TRAIN_FSDP})') -if REF_GPUS % REF_FSDP != 0: - raise ValueError(f'REF_GPUS ({REF_GPUS}) must be divisible by REF_FSDP ({REF_FSDP})') -TRAIN_DP = TRAIN_GPUS // TRAIN_FSDP -REF_DP = REF_GPUS // REF_FSDP - - -# =========================================================================== -# Block A -- boxed extraction + answer grading -# =========================================================================== -_BOXED_RE = re.compile(r'\\boxed\s*\{') - - -def extract_boxed(text: str) -> Optional[str]: - """Last ``\\boxed{...}`` content, brace-balanced.""" - if not text: - return None - last = None - for m in _BOXED_RE.finditer(text): - depth, i = 1, m.end() - while i < len(text) and depth > 0: - depth += (text[i] == '{') - (text[i] == '}') - i += 1 - if depth == 0: - last = text[m.end():i - 1].strip() - return last - - -def normalize_answer(ans: str) -> str: - if not ans: - return '' - s = str(ans).strip() - m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) - if m: - return m.group(1) - s = s.strip('$').strip().replace('−', '-') - 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').replace(r'\tfrac', r'\frac') - s = s.replace(r'\,', '').replace(r'\;', '').replace(r'\!', '') - s = re.sub(r'\\(?:quad|qquad|\s)', '', s) - s = re.sub(r'\s+', '', s) - s = re.sub(r'\\sqrt([0-9A-Za-z])', r'\\sqrt{\1}', s) - s = re.sub(r'\\frac\{([^{}]+)\}([^{}\\])', r'\\frac{\1}{\2}', s) - s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) - s = re.sub(r'\{[a-zA-Z]+\}$', '', s) - s = re.sub(r'\^\\circ|\^\{\\circ\}|°|\\circ', 'deg', s) - s = s.replace(r'\minus{}', '-').replace(r'\minus', '-') - - def _frac_to_slash(mt): - text = mt.group(0) - pos = text.index('{') + 1 - depth, num_start = 1, pos - while depth > 0: - depth += (text[pos] == '{') - (text[pos] == '}') - pos += 1 - numer = text[num_start:pos - 1] - pos += 1 - den_start, depth = pos, 1 - while depth > 0: - depth += (text[pos] == '{') - (text[pos] == '}') - pos += 1 - return f'({numer})/({text[den_start:pos - 1]})' - - s = re.sub(r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', _frac_to_slash, s) - s = re.sub(r'(? bool: - try: - va, vb = float(a.replace('(', '').replace(')', '')), 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) - return va is not None and vb is not None and abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) - - -_MCQ_REF_RE = re.compile( - r'^\\(?:textbf|text|mathrm|mathbf)\{\(?([A-E])\)?\s*\}\s*(.+)$' - r'|^\(?([A-E])\)\s+(.+)$') -_VAR_PREFIX_RE = re.compile(r'^(?:[A-Za-z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)$', re.DOTALL) - - -def _split_mcq(ans: str): - s = ans.strip() - m = _MCQ_REF_RE.match(s) - if m: - return (m.group(1) or m.group(3)), ((m.group(2) or m.group(4) or '').strip() or None) - bl = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) - return (bl.group(1), None) if bl else (None, s or None) - - -def _strip_var_prefix(ans: str) -> str: - m = _VAR_PREFIX_RE.match((ans or '').strip()) - return m.group(1).strip() if m else (ans or '') - - -def _normalize_tuple(ans: str) -> str: - return re.sub(r'[\s()\[\]{}\\]', '', ans or '') - - -def _math_verify_equal(predicted: str, reference: str) -> bool: - try: - from math_verify import parse, verify - gold = parse(r'\boxed{' + reference + '}', parsing_timeout=1) - pred = parse(r'\boxed{' + predicted + '}', parsing_timeout=1) - return bool(gold and pred and verify(gold, pred, timeout_seconds=1)) - except Exception: - return False - - -def answers_match(predicted: str, reference: str) -> bool: - if not predicted or not reference: - return False - norm_p, norm_r = normalize_answer(predicted), normalize_answer(reference) - if norm_p == norm_r or norm_p.lower() == norm_r.lower() or _try_numeric_equal(norm_p, norm_r): - return True - stripped_p = normalize_answer(_strip_var_prefix(predicted)) - stripped_r = normalize_answer(_strip_var_prefix(reference)) - if stripped_p and stripped_r: - if (stripped_p == stripped_r or stripped_p.lower() == stripped_r.lower() - or _try_numeric_equal(stripped_p, stripped_r)): - return True - p_letter, p_value = _split_mcq(predicted) - r_letter, r_value = _split_mcq(reference) - if p_letter and r_letter and p_letter == r_letter: - return True - if (p_letter and p_value is None) and (r_letter and r_value is None): - return False - 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 and (p_val == r_val or p_val.lower() == r_val.lower() or _try_numeric_equal(p_val, r_val)): - return True - for left, right in ((norm_p, norm_r), (stripped_p, stripped_r)): - tl, tr = _normalize_tuple(left), _normalize_tuple(right) - if ',' in tl and tl == tr: - return True - if '=' in norm_r and '=' not in norm_p: - if any(part and (part == norm_p or _try_numeric_equal(part, norm_p)) for part in norm_r.split('=')): - return True - if '=' in norm_p and '=' not in norm_r: - if any(part and (part == norm_r or _try_numeric_equal(part, norm_r)) for part in norm_p.split('=')): - return True - return _math_verify_equal(stripped_p or norm_p, stripped_r or norm_r) - - -# =========================================================================== -# Block B -- prompts, skill parsing, batched sampling -# =========================================================================== -DIRECT_SYSTEM = ( - 'You are an expert competition mathematician. Solve the following problem ' - 'step by step. Provide your final answer inside \\boxed{}.') - -_SKILL_SOLVE_PREFIX = ( - DIRECT_SYSTEM + '\n\n' - 'Before you start, keep these reminders in mind to avoid common mistakes on this ' - 'type of problem:\n') -_SKILL_SOLVE_SUFFIX = '\nApply them where relevant, but rely on your own reasoning to reach the answer.' - - -def build_direct_prompt(problem: str) -> Dict[str, Any]: - return {'messages': [{'role': 'system', 'content': DIRECT_SYSTEM}, - {'role': 'user', 'content': problem}]} - - -def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: - skill = (skill or '').strip() - if not skill: - return build_direct_prompt(problem) - # Concatenation (not .format): DIRECT_SYSTEM/skill contain literal braces. - return {'messages': [ - {'role': 'system', 'content': _SKILL_SOLVE_PREFIX + skill + _SKILL_SOLVE_SUFFIX}, - {'role': 'user', 'content': problem}]} - - -# -- skill-gen prompts (view A: problem + rubric source; view B: query only) -- -# Kept deliberately short: this is the RL policy's system prompt, so over-specifying -# the output hurts convergence. The concrete output format is appended separately by -# ``_SKILL_OUTPUT`` so the task text stays format-agnostic. -SKILL_GEN_SYSTEM = ( - 'You are a math guidance writer. A process-check on a related problem hints at ' - 'likely mistakes. Write short reusable guidance for this and similar problems, ' - 'and note what to watch out for.\n') - -SKILL_GEN_SYSTEM_Q = ( - 'You are a math guidance writer. Write short reusable guidance for this and ' - 'similar problems.\n') - -_SKILL_OUTPUT = ( - 'Output only:\n\nYour reusable solving guidance here.\n') - -SKILL_GEN_USER_Q = ( - 'Problem:\n{problem}\n\n') - -SKILL_GEN_USER_RUBRIC = ( - 'Target problem:\n{problem}\n\n' - 'Problem used for the process check:\n{rubric_problem}\n\n' - 'Process check:\n' - '{diagnosis}\n\n') - - -def _rubric_has_fail(diagnosis: str) -> bool: - """View A uses the open-book rubric prompt (and, under --viewa-sft, SFT distillation) - IFF its rubric localised a failure. No ``[FAIL]`` => nothing to inject: generation - degrades to query-only and the problem is trained by GRPO exactly like view B. Single - source of truth so ``_skillgen_messages`` and ``_group_records`` never diverge.""" - return '[FAIL]' in (diagnosis or '') - - -def _skillgen_messages(problem: str, view: str, diagnosis: str, - rubric_problem: str = '') -> List[Dict[str, Any]]: - """Single source of truth for the skill-gen prompt (used at BOTH generation and - training so they never diverge). View A with a localisable failure uses the target - problem plus the rubric source problem and findings; view B -- or a view-A problem - whose rubric flagged nothing (no ``[FAIL]``) -- degrades to the query-only prompt.""" - if view == 'B' or not _rubric_has_fail(diagnosis): - return [{'role': 'system', 'content': SKILL_GEN_SYSTEM_Q + _SKILL_OUTPUT}, - {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=problem)}] - rubric_problem = rubric_problem or problem - return [{'role': 'system', 'content': SKILL_GEN_SYSTEM + _SKILL_OUTPUT}, - {'role': 'user', 'content': SKILL_GEN_USER_RUBRIC.format( - problem=problem, rubric_problem=rubric_problem, diagnosis=diagnosis)}] - - -def _assign_view(problem: str, args: argparse.Namespace) -> str: - h = int(hashlib.md5(f'{args.seed}:{problem}'.encode('utf-8')).hexdigest(), 16) - return 'B' if (h % 100000) / 100000.0 < args.view_b_frac else 'A' - - -def _curriculum_view_b_frac(gstep: int, args: argparse.Namespace) -> float: - """View-A anneal (--viewa-frac-start): the view-A share holds at ``viewa_frac_start`` - for the first ``viewa_warmup_chunks`` chunks (pure-SFT warmup when start==1.0), then - decays linearly to ``viewa_frac_end`` over ``viewa_decay_chunks`` chunks and holds. - Because _assign_view is a fixed hash against a moving threshold, the B set grows - MONOTONICALLY: a problem trained open-book (A) early can only reappear closed-book - (B) later, never the reverse.""" - t = min(max(gstep - args.viewa_warmup_chunks, 0) / max(args.viewa_decay_chunks, 1), 1.0) - share = args.viewa_frac_start + (args.viewa_frac_end - args.viewa_frac_start) * t - return 1.0 - share - - -def _view_prompt(r: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: - return {'messages': _skillgen_messages( - r['problem'], r['_view'], r.get('_rubric_diag', ''), r.get('_rubric_src', ''))} - - -_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') - - -def _clean_text(decoded: Optional[str]) -> str: - return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() - - -def _extract_tag_block(answer: str, tag: str, allow_empty: bool = False) -> Optional[str]: - low = answer.lower() - open_tag, close_tag = f'<{tag}>', f'' - s = low.rfind(open_tag) - if s < 0: - return None - inner = s + len(open_tag) - e = low.find(close_tag, inner) - if e < 0: - return None - block = answer[inner:e].strip() - block = re.sub(r'', '', block, flags=re.IGNORECASE).strip() - return block if (block or allow_empty) else None - - -def _extract_skill(text: str) -> Optional[str]: - """Parse skill-generation output: return the inner text of a non-empty ```` - block, or None. If a ```` marker is present, parse only the text after the - last one; otherwise parse the full response.""" - low = text.lower() - end_think = low.rfind('') - answer = text[end_think + len(''):] if end_think >= 0 else text - return _extract_tag_block(answer, 'skills') - - -def _parse_seq(seq, gold: str) -> Dict[str, Any]: - """Grade one sampled sequence into a rollout record.""" - text = _clean_text(getattr(seq, 'decoded', '') or '') - pred = extract_boxed(text) - correct = bool(pred) and answers_match(pred, gold) - terminated = getattr(seq, 'stop_reason', None) != 'length' - return {'pred': pred, 'correct': correct, 'terminated': terminated, - 'passed': bool(correct and terminated), - 'stop_reason': getattr(seq, 'stop_reason', None), - 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} - - -def _run_samples(sampler, prompts: List[Any], num_samples: int, max_tokens: int, - gen_dp: int, temperature: Optional[float] = None, - top_p: Optional[float] = None, top_k: Optional[int] = None) -> List[List[Any]]: - """One batched sampler call -> per-prompt list of raw sequences. vLLM dp needs - batch len >= dp, so pad the tail and slice back.""" - if not prompts: - return [] - params = SamplingParams( - max_tokens=max_tokens, - temperature=GEN_TEMPERATURE if temperature is None else temperature, - top_p=GEN_TOP_P if top_p is None else top_p, - num_samples=num_samples, **({} if top_k is None else {'top_k': top_k})) - padded = prompts - if gen_dp > 1 and 0 < len(prompts) < gen_dp: - padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] - responses = sampler.sample(padded, params)[:len(prompts)] - return [list(r.sequences) if (r and r.sequences) else [] for r in responses] - - -# =========================================================================== -# Block C -- data loading via twinkle.Dataset + numeric filtering -# =========================================================================== -def _boxed_batch(rows: Dict[str, List[Any]], dataset: str) -> Dict[str, List[Any]]: - """Batched (columnar) mapper for ``Dataset.map``: annotate each row with its boxed - ``reference_answer`` and a ``_keep`` flag (aops also needs the ``boxed`` metadata).""" - sols = rows['solution'] - metas = rows.get('metadata', [None] * len(sols)) - refs = [extract_boxed(s or '') for s in sols] - keep = [bool(ref) and (dataset != 'aops' or bool((meta or {}).get('boxed'))) - for ref, meta in zip(refs, metas)] - return {**rows, 'reference_answer': refs, '_keep': keep} - - -def load_problems(dataset: str, n: int, seed: int, num_proc: int = 0) -> List[Dict[str, Any]]: - """Load boxed-answer problems as ``{problem, reference_answer, level?}`` via - twinkle.Dataset (ModelScope hub), sampled to ``n`` (0 = all). Boxed extraction (regex - + brace scan over every solution) is parallelised by ``Dataset.map``/``.filter``; - ``num_proc`` defaults to all cores (set 1 to force serial).""" - ds_id = AOPS_DATASET_ID if dataset == 'aops' else MATH_DATASET_ID - ds = Dataset(DatasetMeta(dataset_id=f'ms://{ds_id}', split='train')) - nproc = num_proc if num_proc > 0 else min(32, os.cpu_count() or 1) - ds.map(lambda rows: _boxed_batch(rows, dataset), num_proc=nproc) - ds.filter(lambda row: row['_keep'], num_proc=nproc) - has_level = 'level' in ds.dataset.column_names - out = [{'data_id': f'{dataset}:{i}', 'problem': row['problem'], - 'reference_answer': row['reference_answer'], - **({'level': row['level']} if has_level and row.get('level') else {})} - for i, row in enumerate(ds.dataset)] - logger.info(f'[data] {dataset}: {len(out)} boxed problems ({nproc} procs)') - rng = np.random.RandomState(seed) - rng.shuffle(out) - return out[:n] if (n and n < len(out)) else out - - -# --------------------------------------------------------------------------- -# Bag-of-words neighbour pairing (cross-problem rubric transfer, --xproblem-rubric) -# --------------------------------------------------------------------------- -# Common English + math-scaffolding words that carry no problem-type signal. Kept small -# and deterministic on purpose (no external stopword list): what survives is the domain -# vocabulary ('triangle', 'prime', 'polynomial', ...) that actually defines the type. -_BOW_STOP = frozenset(""" -a an the of to in on at for and or but if is are be was were been being this that these those -with without into onto from by as it its their his her our your my we you they he she them -find compute determine calculate evaluate solve show prove given let suppose consider assume -what which when where how many much value values number numbers expression form terms term -such that then than so if only when each every all any some both one two three four five six -seven eight nine ten first second third last non over under about above below between -problem answer result equal equals sum difference product total following there here have has -had do does did can could will would should may might must not no yes if then else -""".split()) - -_WORD_RE = re.compile(r'[a-z]+') - - -def _stem(w: str) -> str: - """Crude suffix stripper so 'prime'/'primes', 'triangle'/'triangles' collapse to one - type token. Not linguistically correct -- just enough to merge the common plural/gerund - variants that otherwise split a type's vocabulary and starve the df filter.""" - if len(w) > 4 and w.endswith('ies'): - return w[:-3] + 'y' # properties -> property - if len(w) > 4 and w.endswith('es') and w[-3] in 'sxzh': - return w[:-2] # boxes -> box (keep primes -> prime below) - for suf in ('ing', 'ed', 's'): - if len(w) > len(suf) + 2 and w.endswith(suf): - return w[:-len(suf)] - return w - - -def _tokenize(problem: str) -> List[str]: - """Deterministic bag-of-words tokens for type matching: lowercase alphabetic words - (numbers dropped -- they are instance detail, not type), minus generic stopwords, then - stemmed, so only domain terms remain. Latex control words (frac, sqrt, ...) survive.""" - return [_stem(w) for w in _WORD_RE.findall((problem or '').lower()) - if len(w) > 2 and w not in _BOW_STOP] - - -class BagOfWordsIndex: - """TF-IDF cosine nearest-neighbour over problem statements. Sparse dict vectors + - an inverted index, so ``nearest`` scores only problems sharing a term (near-linear in - practice, no dense NxN). Deterministic; ties break on lower index for reproducibility. - - Answer-aware: ``nearest`` skips the same index, near-verbatim duplicates (cosine - >= ``sim_max``) and -- crucially for anti-leak -- any candidate whose answer equals the - query's, so a neighbour rubric can never hand over the query's own answer.""" - - def __init__(self, problems: List[str], answers: Optional[List[str]] = None, - min_df: int = 2, max_df_frac: float = 0.5): - self._toks = [_tokenize(p) for p in problems] - self._ans = [(_numeric_value(a) or (str(a).strip() if a else '')) for a in answers] \ - if answers is not None else [''] * len(problems) - n = len(self._toks) - df: Dict[str, int] = {} - for toks in self._toks: - for w in set(toks): - df[w] = df.get(w, 0) + 1 - max_df = max(min_df, int(max_df_frac * n)) - self._idf = {w: math.log((n + 1) / (c + 1)) + 1.0 - for w, c in df.items() if min_df <= c <= max_df} - self._vecs: List[Dict[str, float]] = [self._vectorize(t) for t in self._toks] - self._inverted: Dict[str, List[int]] = {} - for i, v in enumerate(self._vecs): - for w in v: - self._inverted.setdefault(w, []).append(i) - - def _vectorize(self, toks: List[str]) -> Dict[str, float]: - tf: Dict[str, float] = {} - for w in toks: - if w in self._idf: - tf[w] = tf.get(w, 0.0) + 1.0 - vec = {w: c * self._idf[w] for w, c in tf.items()} - norm = math.sqrt(sum(x * x for x in vec.values())) - return {w: x / norm for w, x in vec.items()} if norm > 0 else {} - - def nearest(self, i: int, sim_max: float = 0.98) -> Tuple[int, float]: - """(index, cosine) of the most similar *distinct* problem: not i, not a duplicate - (cosine < ``sim_max``) and not the same answer. (-1, 0.0) if none qualifies.""" - vi = self._vecs[i] - if not vi: - return -1, 0.0 - ai = self._ans[i] - scores: Dict[int, float] = {} - for w, xi in vi.items(): - for j in self._inverted.get(w, ()): - if j != i: - scores[j] = scores.get(j, 0.0) + xi * self._vecs[j].get(w, 0.0) - best_j, best_s = -1, 0.0 - for j, s in scores.items(): - if s >= sim_max or (ai and self._ans[j] == ai): - continue - if s > best_s or (s == best_s and (best_j < 0 or j < best_j)): - best_j, best_s = j, s - return best_j, best_s - - -def build_pairs(records: List[Dict[str, Any]], n: int, seed: int, sim_max: float = 0.98 - ) -> Tuple[List[Dict[str, Any]], Dict[str, Tuple[str, float]]]: - """Single-pass cross-problem pairing over the whole pool (one index build). - - Returns ``(subset, neighbour_map)`` where ``subset`` is the ``n`` problems with the - strongest qualifying neighbour (dense same-type pairs; n<=0/n>=len keeps all, shuffled) - and ``neighbour_map[q] = (p, cosine)`` gives each kept problem its analogue P. P is drawn - from the *full* pool (richer analogues) but never a duplicate or same-answer problem, so - P's rubric can transfer method without ever leaking Q's answer.""" - index = BagOfWordsIndex([r['problem'] for r in records], - [str(r.get('reference_answer', '')) for r in records]) - nbr = [index.nearest(i, sim_max) for i in range(len(records))] - order = sorted(range(len(records)), key=lambda i: (-nbr[i][1], i)) - keep = order[:n] if (0 < n < len(records)) else list(range(len(records))) - rng = np.random.RandomState(seed) - rng.shuffle(keep) - subset = [records[i] for i in keep] - neighbour_map = {records[i]['problem']: (records[nbr[i][0]]['problem'], nbr[i][1]) - for i in keep if nbr[i][0] >= 0} - return subset, neighbour_map - - -_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') - - -def _norm_num_text(num: str) -> str: - try: - f = float(num) - return str(int(f)) if f == int(f) else str(f) - except Exception: - return str(num).strip() - - -def _numeric_value(raw: Any) -> Optional[str]: - """Collapse an answer to a single int/decimal/fraction, or None.""" - if raw is None: - return None - s = str(raw).strip().strip('$').strip() - s = s.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac') - s = re.sub(r'\\!|\\,|\\;|\\ |\\left|\\right|\s', '', s) - for pat in (r'\\frac\{(-?\d+)\}\{(-?\d+)\}', r'(-?\d+)/(-?\d+)'): - m = re.fullmatch(pat, s) - if m: - a, b = int(m.group(1)), int(m.group(2)) - return _norm_num_text(str(a / b)) if b else None - return _norm_num_text(s) if _NUM_RE.fullmatch(s) else None - - -def _answer_leaked(skill: str, reference: str) -> bool: - """Audit whether a generated skill contains the final answer verbatim. This is NOT - a training filter: if the skill model derives an answer from the problem, that is a - legitimate answer-bearing skill under this experiment. The real leakage boundary is the - external API/rubric diagnosis, which is constrained by prompt to stay answer-free.""" - if not skill: - return False - for cand in {_numeric_value(reference), (str(reference).strip() or None)}: - if cand and re.search(r'(? Tuple[Set[str], Set[str]]: - """Read jsonl files and collect data_id/problem keys that must be excluded. - - The cold-start SFT builder writes stable ``data_id`` values. ``problem`` is kept as a - backward-compatible fallback for older jsonl files produced before data_id existed.""" - ids: Set[str] = set() - problems: Set[str] = set() - for raw_path in (paths_arg or '').split(','): - path = raw_path.strip() - if not path or not os.path.exists(path): - continue - with open(path, encoding='utf-8') as f: - for line in f: - if not line.strip(): - continue - row = json.loads(line) - if row.get('record_type') in {'config', 'summary'}: - continue - data_id = str(row.get('data_id') or '').strip() - problem = str(row.get('problem') or '').strip() - if data_id: - ids.add(data_id) - elif problem: - problems.add(problem) - return ids, problems - - -def _load_records(args: argparse.Namespace - ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], - Dict[str, Tuple[str, float]], Dict[str, str], Dict[str, int]]: - """Load, numeric-filter, shuffle, split a fixed eval holdout, and (when --xproblem-rubric) - select a same-type-dense train subset with its neighbour map -- all in one pass. - Also returns ``pool_answers`` (every candidate neighbour P's true answer) so P's baseline - can be graded/cached correctly even when P is not itself a training problem.""" - # Load all when filtering or splitting (else the eval holdout could starve train). - load_n = 0 if (args.numeric_only or args.eval_size > 0) else args.n - records = load_problems(args.dataset, load_n, args.seed) - raw_n, dropped = len(records), 0 - if args.numeric_only: - records = [{**r, 'reference_answer': v} - for r, v in ((r, _numeric_value(r.get('reference_answer'))) for r in records) - if v is not None] - dropped = raw_n - len(records) - np.random.RandomState(args.seed).shuffle(records) - exclude_ids, exclude_problems = _load_excluded_records(getattr(args, 'exclude_data_ids', '')) - excluded = 0 - if exclude_ids or exclude_problems: - before = len(records) - records = [r for r in records - if str(r.get('data_id', '')) not in exclude_ids - and str(r.get('problem', '')).strip() not in exclude_problems] - excluded = before - len(records) - eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 - eval_records = [dict(r) for r in records[:eval_n]] - pool = records[eval_n:] - pool_offset = max(0, int(getattr(args, 'pool_offset', 0) or 0)) - if pool_offset: - if pool_offset >= len(pool): - raise ValueError(f'--pool-offset ({pool_offset}) leaves no train records from pool size {len(pool)}') - pool = pool[pool_offset:] - train_n = args.n if args.n > 0 else len(pool) - # Cross-problem rubric: pick the train_n problems with the strongest qualifying neighbour - # (dense same-type pairs) and build {Q -> (P, sim)} in one index pass. Otherwise keep the - # first train_n (already shuffled) with no neighbours. - if args.xproblem_rubric: - subset, neighbor_map = build_pairs(pool, train_n, args.seed) - train_records = [dict(r) for r in subset] - pool_answers = {r['problem']: str(r.get('reference_answer', '')) for r in pool} - else: - train_records, neighbor_map, pool_answers = [dict(r) for r in pool[:train_n]], {}, {} - if {r['problem'] for r in train_records} & {r['problem'] for r in eval_records}: - raise ValueError('eval/train overlap detected') - stats = {'raw_loaded': raw_n, 'numeric_dropped': dropped, - 'excluded_records': excluded, 'pool_offset': pool_offset, - 'train_records': len(train_records), 'eval_records': len(eval_records)} - return train_records, eval_records, neighbor_map, pool_answers, stats - - -# =========================================================================== -# Block D -- disk cache, problem pool, baseline rollout, rubric check -# =========================================================================== -class DiskCache: - """Append-only jsonl kv cache (md5 key -> value). Loads on init, appends on put. - Disabled instances always miss and never write.""" - - def __init__(self, path: str, enabled: bool = True): - self._mem: Dict[str, Any] = {} - self._fh = None - self._lock = threading.Lock() # base baseline is prefetched on a background thread - if not enabled: - return - if os.path.exists(path): - with open(path, encoding='utf-8') as f: - for line in f: - if line.strip(): - row = json.loads(line) - self._mem[row['key']] = row['value'] - self._fh = open(path, 'a', encoding='utf-8') - - @staticmethod - def key_for(*parts: str) -> str: - return hashlib.md5('\x1f'.join(parts).encode('utf-8')).hexdigest() - - def __contains__(self, key: str) -> bool: - with self._lock: - return key in self._mem - - def get(self, key: str) -> Any: - with self._lock: - return self._mem.get(key) - - def put(self, key: str, value: Any) -> None: - with self._lock: - self._mem[key] = value - if self._fh is not None: - self._fh.write(json.dumps({'key': key, 'value': value}, ensure_ascii=False) + '\n') - self._fh.flush() - - def close(self) -> None: - if self._fh is not None: - self._fh.close() - - -class _LockedSampler: - """Serialises every ``.sample()`` on a sampler behind one lock. The base sampler is - shared by the main thread (with-skill scoring, eval) and the baseline-prefetch thread; - ``sample`` is a slice_dp/flatten remote call whose collect is NOT safe to interleave - across two callers, so concurrent calls could mis-join sequences. The lock keeps base - calls serial (prefetch still overlaps the skill-gen phase, which uses skill_sampler).""" - - def __init__(self, sampler): - self._sampler = sampler - self._lock = threading.Lock() - - def sample(self, *args, **kwargs): - with self._lock: - return self._sampler.sample(*args, **kwargs) - - def __getattr__(self, name): - return getattr(self._sampler, name) - - -class ProblemPool: - """Cyclic draw source. Each full pass reshuffles with ``seed + epoch``; the initial - pass keeps the loader order. ``draw(k)`` returns k distinct problems (k << pool).""" - - def __init__(self, records: List[Dict[str, Any]], seed: int): - self._records = list(records) - self._seed, self._cursor, self.epoch = seed, 0, 0 - - def draw(self, k: int) -> List[Dict[str, Any]]: - out, seen = [], set() - while len(out) < k: - if self._cursor >= len(self._records): - self.epoch += 1 - np.random.RandomState(self._seed + self.epoch).shuffle(self._records) - self._cursor = 0 - r = self._records[self._cursor] - self._cursor += 1 - if id(r) not in seen: - seen.add(id(r)) - out.append(r) - return out - - def peek(self, k: int) -> List[Dict[str, Any]]: - """The next k distinct problems draw() would return, WITHOUT advancing state - (no cursor move, no reshuffle). Used to prefetch their baseline into base_cache - while the current chunk trains -- a wrong guess (cross-epoch reshuffle) only - misses the cache, never corrupts the draw.""" - out, seen, cur = [], set(), self._cursor - recs = self._records - while len(out) < k and cur < len(recs): # stop at epoch edge; don't simulate reshuffle - r = recs[cur] - cur += 1 - if id(r) not in seen: - seen.add(id(r)) - out.append(r) - return out - - -def _empty_roll() -> Dict[str, Any]: - return {'pred': '', 'correct': False, 'terminated': False, 'passed': False, - 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} - - -def _apply_baseline(r: Dict[str, Any], roll: Dict[str, Any]) -> None: - """Attach a greedy baseline roll and reset per-chunk working state.""" - r['_baseline_rolls'], r['_cands'], r['_init'] = [roll], [], [roll] - r['_failed'] = not roll['correct'] - r['_baseline_pass'] = 1.0 if roll['correct'] else 0.0 - r['_hard'] = True # process every problem; group variance selects (SEAM-style) - - -def baseline_rollout(base_sampler, problems: List[Dict[str, Any]], base_dp: int, - args: argparse.Namespace, cache: DiskCache) -> int: - """Base solves each problem greedily once (T=0, M=1), disk-cached by problem text. - The base is frozen + greedy so a cache hit is exact. Returns fresh (miss) count.""" - todo = [r for r in problems if DiskCache.key_for(r['problem']) not in cache] - if todo: - out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], - 1, args.max_tokens, base_dp, temperature=0.0) - for r, seqs in zip(todo, out): - roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - cache.put(DiskCache.key_for(r['problem']), roll) - for r in problems: - _apply_baseline(r, cache.get(DiskCache.key_for(r['problem']))) - return len(todo) - - -# -- rubric process-check (view A): teacher diagnoses the base's attempt -- -_RFT_DIAG_SYSTEM = """\ -You are a strategy-level process checker for a math solution attempt. You are given a -math problem, a rubric, and one attempted solution segment. Decide PASS or FAIL for each -criterion, and write the diagnosis so it can become useful reusable guidance for solving -similar problems without seeing this segment. - -Output STRICT JSON (no prose outside it) with this shape: -{ - "items": [ - {"index": 1, "verdict": "PASS", "reason": "", - "fix": ""}, - {"index": 2, "verdict": "FAIL", "reason": "", - "fix": ""} - ], - "overall": "OK" | "ISSUES", - "summary": "" -} - -Rules: -- Judge every criterion independently; a [Hard Rule] is FAIL unless clearly satisfied. -- Judge ONLY what is observable in THIS segment. Ignore hidden or - content for output-format criteria. -- The API diagnosis is an external teacher signal, so it must stay answer-free. -- Prefer diagnosis that transfers to view-B skill generation: name the route choice, - structural observation, missing check, or length-control habit that a solver should - remember before solving a similar problem. -- For PASS items, leave "fix" as "". -- For FAIL items, describe the process problem at strategy level: unsuitable method, - missed structure, invalid transformation, missing constraint check, redundant cases, - off-track approach, contradiction, or inefficient/unfinished reasoning. -- A fix may suggest the LOCAL correction direction, such as identify the key structure, - verify constraints, preserve equivalence, reduce redundant cases, or choose a more - direct route. Do not carry out the correction. -- Never reveal the final answer, a corrected value/expression, an option label, or a - step-by-step solution that would let another model copy the solve. -- If the segment contains a process note saying it was cut off before a final boxed - answer, mark the length-budget criterion as FAIL and suggest a method-level way to - finish faster. -- Keep every "reason" and "fix" concise: one short sentence each. -- "overall" is "OK" only if NO criterion is FAIL. -- Output only the JSON object.""" - -_RFT_DIAG_USER = """\ -## Task / query (context) -{query} - -## Rubric -{rubric} - -## Segment -{segment} - -Now output the diagnostic JSON object.""" - -_MATH_RUBRIC = [ - ('The attempt chooses a method suitable for the problem structure', False), - ('The attempt identifies the key constraint, invariant, or quantity before computing', False), - ('Algebraic and logical transformations preserve validity at each step', True), - ('The attempt checks required constraints, domains, boundary cases, or validity conditions', False), - ('The attempt avoids redundant casework, looping, or re-deriving known facts', False), - ('The attempt reaches a final boxed answer within the length budget', False), - ('The approach stays focused on the actual question asked', False), -] - -# Bump when the rubric criteria or the diagnosis prompt change so the disk-cached -# diagnoses written under an older rubric are not silently reused. -_RUBRIC_VERSION = 'rubric_v5_viewb_strategy' - - -class _RftRubricVerifier(RubricVerifier): - def _diagnose_trajectory(self, query: str, rubric_block: str, segment_text: str) -> dict: - return {'messages': [ - {'role': 'system', 'content': _RFT_DIAG_SYSTEM}, - {'role': 'user', 'content': _RFT_DIAG_USER.format( - query=query, rubric=rubric_block, segment=segment_text)}]} - - -def build_rubric_checker() -> Optional[RubricVerifier]: - """Fixed math-process rubric verifier (teacher-served). None if no LLM backup env.""" - if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') - or os.environ.get('OPENAI_API_KEY')): - return None - return _RftRubricVerifier( - fixed_rubric=[RubricItem(t, is_hard=h) for t, h in _MATH_RUBRIC], gate=True) - - -def _format_diagnosis(detail) -> str: - """One line per criterion (PASS/FAIL + reason + fix) then a summary.""" - rub = detail.rubric - lines = [] - for it in detail.items: - text = rub[it.index - 1].text if 0 < it.index <= len(rub) else f'criterion {it.index}' - if it.verdict: - lines.append(f'- [PASS] {text}') - else: - tail = f': {it.reason}' if it.reason else '' - tail += f' (fix: {it.fix})' if it.fix else '' - lines.append(f'- [FAIL] {text}{tail}') - if detail.summary: - lines.append(f'Summary: {detail.summary}') - return '\n'.join(lines) - - -def diagnose_views(checker, problems: List[Dict[str, Any]], args: argparse.Namespace, - cache: DiskCache) -> None: - """Rubric-check every view-A problem's greedy attempt in parallel (disk-cached by - problem + attempt), stashing the formatted findings on ``r['_rubric_diag']``.""" - targets = [r for r in problems if r.get('_view') == 'A'] - if not checker or not targets: - return - - def _key(r: Dict[str, Any]) -> str: - init = r.get('_init', [{}])[0] - term = 'L' if (init.get('stop_reason') == 'length' or not init.get('terminated')) else 'T' - return DiskCache.key_for(r['problem'], init.get('text', ''), _RUBRIC_VERSION, term) - - pending = [] - for r in targets: - key = _key(r) - if key in cache: - r['_rubric_diag'] = cache.get(key) - else: - pending.append((r, key)) - if not pending: - return - - def _run(item): - r, key = item - init = r['_init'][0] - seg_text = init['text'] - if init.get('stop_reason') == 'length' or not init.get('terminated'): - seg_text += ('\n\n[Process note: this attempt was cut off at the token budget ' - 'and never produced a final \\boxed{} answer.]') - seg = {'messages': [{'role': 'user', 'content': r['problem']}, - {'role': 'assistant', 'content': seg_text}]} - attempts = max(1, args.rubric_retries + 1) - for attempt in range(attempts): - try: - return r, key, _format_diagnosis(checker.diagnose(seg, query=r['problem'])) - except Exception as exc: # teacher hiccup -> retry, then no-diagnosis prompt (not cached) - if attempt + 1 < attempts: - logger.warning(f'[rubric] diagnose error: {exc}; retry {attempt + 1}/{args.rubric_retries}') - time.sleep(min(2.0, 0.5 * (2 ** attempt))) - continue - logger.warning(f'[rubric] diagnose error: {exc}; giving up after {attempts} attempts') - return r, key, None - - workers = max(1, min(args.rubric_workers, len(pending))) - with ThreadPoolExecutor(max_workers=workers) as ex: - for r, key, diag in ex.map(_run, pending): - r['_rubric_diag'] = diag or '' - if diag is not None: - cache.put(key, diag) - - -# =========================================================================== -# Block E -- chunk draw, generation pipeline, record building -# =========================================================================== -def _baseline_class(r: Dict[str, Any]) -> str: - """success | fail_loop (out of length / never terminated) | fail_wrong.""" - roll = r['_init'][0] - if roll['correct']: - return 'success' - return 'fail_loop' if (roll['stop_reason'] == 'length' or not roll['terminated']) else 'fail_wrong' - - -def _select_balanced(buckets: Dict[str, List[Dict[str, Any]]], n_success: int, - n_fail: int, n_fail_loop: int) -> List[Dict[str, Any]]: - """Pick n_fail base-fails (toward n_fail_loop loop-fails, best-effort) + n_success - base-successes; top up any shortfall from leftovers.""" - loop, wrong, succ = buckets['fail_loop'], buckets['fail_wrong'], buckets['success'] - take_loop = min(n_fail_loop, len(loop)) - take_wrong = min(n_fail - take_loop, len(wrong)) - take_loop = min(n_fail - take_wrong, len(loop)) - sel = loop[:take_loop] + wrong[:take_wrong] + succ[:n_success] - target = n_success + n_fail - if len(sel) < target: - used = {id(x) for x in sel} - sel += [x for b in (loop, wrong, succ) for x in b if id(x) not in used][:target - len(sel)] - return sel - - -def draw_chunk(pool: ProblemPool, base_sampler, base_dp: int, args: argparse.Namespace, - cache: DiskCache) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: - """Draw one chunk, baselining every drawn problem. With ``--balance``, keep - drawing+baselining until the target base fail:success mix is reachable (or the budget - is hit), then select a balanced subset.""" - if not args.balance: - chunk = pool.draw(args.chunk_size) - fresh = baseline_rollout(base_sampler, chunk, base_dp, args, cache) - return chunk, {'enabled': False, 'n_drawn': len(chunk), 'n_baseline_fresh': fresh} - - n_success = max(0, min(args.chunk_size, round(args.chunk_size * args.balance_success_frac))) - n_fail = args.chunk_size - n_success - n_fail_loop = round(n_fail * args.balance_loop_frac) - buckets: Dict[str, List[Dict[str, Any]]] = {'success': [], 'fail_loop': [], 'fail_wrong': []} - budget, n_drawn, n_fresh, seen = args.chunk_size * args.balance_max_draws_mult, 0, 0, set() - while n_drawn < budget: - if (len(buckets['success']) >= n_success - and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail): - break - batch = pool.draw(args.chunk_size) - n_fresh += baseline_rollout(base_sampler, batch, base_dp, args, cache) - n_drawn += len(batch) - for r in batch: - if id(r) not in seen: - seen.add(id(r)) - buckets[_baseline_class(r)].append(r) - - reached = (len(buckets['success']) >= n_success - and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail) - chunk = _select_balanced(buckets, n_success, n_fail, n_fail_loop) - sel_success = sum(1 for r in chunk if not r['_failed']) - stats = { - 'enabled': True, 'n_drawn': n_drawn, 'n_baseline_fresh': n_fresh, 'n_selected': len(chunk), - 'target_success': n_success, 'target_fail': n_fail, 'target_fail_loop': n_fail_loop, - 'selected_success': sel_success, 'selected_fail': len(chunk) - sel_success, - 'selected_fail_loop': sum(1 for r in chunk if _baseline_class(r) == 'fail_loop'), - 'selected_fail_wrong': sum(1 for r in chunk if _baseline_class(r) == 'fail_wrong'), - 'selected_success_frac': (sel_success / len(chunk)) if chunk else 0.0, - 'budget_hit': not reached, - } - return chunk, stats - - -def _assign_advantages(hard: List[Dict[str, Any]], args: argparse.Namespace) -> None: - """Group-relative advantage over each problem's scored candidates using the greedy - binary reward R in {0,1}: ``A = (R - mean) / (std + eps)``. std==0 groups get no - gradient -- GRPO's variance selects informative problems (no explicit difficulty gate). - A symmetric clip limits both positive and negative outliers from sparse 1-vs-many groups.""" - eps = 1e-6 - adv_clip = abs(float(getattr(args, 'adv_clip', 0.0) or 0.0)) - for r in hard: - for c in r['_cands']: - c['advantage'], c['grpo_adv'], c['kept'] = 0.0, 0.0, False - cs = [c for c in r['_cands'] if c.get('reward') is not None] - if len(cs) < 2: - continue - rewards = [c['reward'] for c in cs] - mean_r = sum(rewards) / len(rewards) - std = (sum((x - mean_r) ** 2 for x in rewards) / len(rewards)) ** 0.5 - if std < 1e-9: - continue - for c in cs: - raw_adv = (c['reward'] - mean_r) / (std + eps) - adv = max(-adv_clip, min(adv_clip, raw_adv)) if adv_clip > 0 else raw_adv - c['advantage'], c['grpo_adv'], c['kept'] = adv, adv, c['reward'] > mean_r - - -def _best_sft_candidate(r: Dict[str, Any], args: argparse.Namespace) -> Optional[Dict[str, Any]]: - """Pick ONE view-A candidate to distill (online context distillation). ONLY - executor-verified PASSING skills (reward==1) are distilled: the earlier fallback to - unverified skills meant ~60% of SFT targets had failed their own executor pass - (measured on the sft35 run) and the model was imitating plausible-but-wrong skills. - Problems with no passing candidate now yield NO SFT record. Answer-bearing skills - produced by the skill model itself are allowed here; only the external API/rubric - diagnosis must be answer-free. Among the passing candidates, take the one whose skill - length is CLOSEST to ``--sft-target-len`` -- an empirically high-pass-rate length - (~500-600 chars in this run) -- breaking ties by the fewest executor solve tokens. - Targeting a length (rather than the minimum) avoids a distillation feedback loop that - would otherwise drive rollouts ever shorter.""" - eligible = [c for c in r['_cands'] if c.get('parseable') and c.get('skills')] - passing = [c for c in eligible if c.get('reward') == 1.0] - if not passing: - return None - target = int(getattr(args, 'sft_target_len', 550) or 550) - - def _solve_tokens(c: Dict[str, Any]) -> int: - rolls = c.get('rolls') or [] - return rolls[0].get('gen_tokens', 1 << 30) if rolls else (1 << 30) - - return min(passing, key=lambda c: (abs(len(c['skills']) - target), _solve_tokens(c))) - - -def apply_neighbor_rubric(base_sampler, chunk: List[Dict[str, Any]], - neighbor_map: Dict[str, Tuple[str, float]], - pool_answers: Dict[str, str], base_dp: int, - args: argparse.Namespace, checker, - base_cache: DiskCache, rubric_cache: DiskCache) -> None: - """Cross-problem rubric (--xproblem-rubric): for every view-A problem Q, replace its own - rubric with the rubric of its bag-of-words neighbour P (Q keeps being the solved/scored - problem). P is baselined + diagnosed here (both disk-cached) as a stub carrying P's REAL - answer, so P's baseline grades correctly and legitimately shares the baseline cache with - P-as-training-problem. P's findings are copied onto Q with the neighbour text + similarity - for audit. A P that can't be diagnosed -> Q degrades to view B. Since P's answer differs - from Q's (guaranteed at pairing time), any answer P's rubric leaks is useless for Q.""" - targets = [r for r in chunk if r.get('_view') == 'A' and neighbor_map.get(r['problem'])] - if not targets: - return - stubs, by_problem = [], {} - for r in targets: - p, _ = neighbor_map[r['problem']] - if p not in by_problem: - stub = {'problem': p, 'reference_answer': pool_answers.get(p, ''), '_view': 'A'} - by_problem[p] = stub - stubs.append(stub) - baseline_rollout(base_sampler, stubs, base_dp, args, base_cache) - diagnose_views(checker, stubs, args, rubric_cache) - for r in targets: - p, sim = neighbor_map[r['problem']] - r['_rubric_diag'] = by_problem[p].get('_rubric_diag', '') - r['_rubric_src'], r['_neighbor_sim'] = p, sim - - -def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], - ci: int, base_dp: int, skill_dp: int, args: argparse.Namespace, - checker, rubric_cache: DiskCache, base_cache: DiskCache = None, - neighbor_map: Optional[Dict[str, Tuple[str, float]]] = None, - pool_answers: Optional[Dict[str, str]] = None - ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], List[Dict[str, Any]]]: - """view assign -> rubric-check (view A) -> skill-gen -> answer audit -> with-skill - greedy pass -> GRPO advantages. ``chunk`` arrives already baselined by draw_chunk. - With --xproblem-rubric, view A gets its NEIGHBOUR's rubric (see apply_neighbor_rubric).""" - hard = chunk - for r in hard: - r['_view'], r['_rubric_diag'] = _assign_view(r['problem'], args), '' - if args.xproblem_rubric and neighbor_map: - apply_neighbor_rubric(base_sampler, hard, neighbor_map, pool_answers or {}, base_dp, - args, checker, base_cache, rubric_cache) - else: - diagnose_views(checker, hard, args, rubric_cache) - - # skill-gen (thinking OFF), per-view prompt; re-sample problems with no clean candidate. - # View-A problems that cannot yield a training record (no [FAIL] rubric, or rubric - # leaked the answer) are dropped from training entirely -- skip their generation. - flat: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] - pending = [r for r in hard if not _viewa_dropped(r, args)] - for _ in range(args.skill_retries + 1): - if not pending: - break - sg_out = _run_samples(skill_sampler, [_view_prompt(r, args) for r in pending], - args.n_skills, args.skill_max_tokens, skill_dp, - temperature=args.skill_gen_temperature, - top_p=args.skill_gen_top_p, top_k=args.skill_gen_top_k) - still = [] - for r, seqs in zip(pending, sg_out): - got = False - for s in seqs: - resp = _clean_text(getattr(s, 'decoded', '') or '') - block = _extract_skill(resp) or '' - cand = {'skills': block, 'response': resp, 'parseable': bool(block), - 'view': r['_view'], 'leaked': None, 'leak_reason': '', - 'leak_source': '', 'with_pass': None, 'reward': None, 'rolls': [], - 'skillgen_stop': getattr(s, 'stop_reason', None), - 'skillgen_tokens': len(getattr(s, 'tokens', None) or [])} - r['_cands'].append(cand) - if block: - flat.append((r, cand)) - got = True - if not got: - still.append(r) - pending = still - - # Answer-bearing skill audit: deterministic verbatim-answer check only (no LLM). This - # is observability only; it records metrics for swanlab/jsonl, but does not block - # scoring, reward, advantage assignment, SFT candidate selection, or GRPO training. - for r, c in flat: - leaked = _answer_leaked(c['skills'], r['reference_answer']) - c['leaked'] = leaked - c['leak_reason'] = 'answer_verbatim' if leaked else '' - c['leak_source'] = 'deterministic' - - # with-skill greedy pass (T=0, M=1); reward = correct, absolute (group mean is baseline). - scored_inputs = flat - if scored_inputs: - ws_out = _run_samples(base_sampler, - [build_skill_solve_prompt(r['problem'], c['skills']) for r, c in scored_inputs], - 1, args.max_tokens, base_dp, temperature=0.0) - for (r, c), seqs in zip(scored_inputs, ws_out): - c['rolls'] = [_parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll()] - c['with_pass'] = 1.0 if c['rolls'][0]['correct'] else 0.0 - c['reward'] = c['with_pass'] - if args.format_in_reward: # unparseable candidates score 0 and still join the group - for r in hard: - for c in r['_cands']: - if c['reward'] is None: - c['reward'] = 0.0 - - _assign_advantages(hard, args) - return ([_full_record(r, ci) for r in chunk], _chunk_summary(chunk, ci, args), - _group_records(chunk, args)) - - -def _roll(x: Dict[str, Any]) -> Dict[str, Any]: - return {k: x[k] for k in ('pred', 'correct', 'terminated', 'passed', - 'stop_reason', 'gen_tokens', 'text')} - - -def _is_trainable(c: Dict[str, Any], args: argparse.Namespace) -> bool: - """Reaches the GRPO update iff advantage is non-zero. Leak flags are audit-only.""" - adv_nz = abs(c.get('advantage') or 0.0) > 1e-9 - if args.format_in_reward: - return adv_nz - return c.get('with_pass') is not None and adv_nz - - -def _full_record(r: Dict[str, Any], ci: int) -> Dict[str, Any]: - """Complete per-problem trace: init attempt, baseline, and all candidates.""" - init = r['_init'][0] - return { - 'record_type': 'problem', 'chunk': ci, 'problem': r['problem'], - 'reference_answer': r['reference_answer'], 'level': r.get('level', ''), - 'failed_first_try': r['_failed'], - 'init_attempt': {'text': init['text'], 'pred': init['pred'], 'correct': init['correct'], - 'terminated': init['terminated'], 'stop_reason': init['stop_reason'], - 'gen_tokens': init['gen_tokens']}, - 'baseline_pass': r['_baseline_pass'], 'is_hard': r['_hard'], - 'view': r.get('_view', ''), 'rubric_diag': r.get('_rubric_diag', ''), - # xproblem audit: which neighbour's rubric this problem borrowed, and how similar. - 'rubric_src': r.get('_rubric_src', ''), 'neighbor_sim': r.get('_neighbor_sim'), - 'baseline_rolls': [_roll(x) for x in r['_baseline_rolls']], - 'candidates': [{ - 'skills': c['skills'], 'response': c['response'], 'parseable': c['parseable'], - 'skillgen_stop': c.get('skillgen_stop'), 'skillgen_tokens': c.get('skillgen_tokens'), - 'leaked': c['leaked'], 'leak_reason': c['leak_reason'], 'leak_source': c['leak_source'], - 'with_pass': c['with_pass'], 'reward': c.get('reward'), 'advantage': c.get('advantage'), - 'grpo_adv': c.get('grpo_adv'), 'kept': c.get('kept'), - 'rolls': [_roll(x) for x in c['rolls']], - } for c in r['_cands']], - } - - -def _view_stats(problems: List[Dict[str, Any]], view: str) -> Dict[str, Any]: - pv = [r for r in problems if r.get('_view') == view] - cands = [c for r in pv for c in r['_cands'] if c['parseable']] - clean = [c for c in cands if c['leaked'] is False] - adopted = sum(1 for r in pv - if any(abs(c.get('advantage') or 0.0) > 1e-9 for c in r['_cands'])) - return {'n': len(pv), 'n_candidates_parseable': len(cands), 'n_clean': len(clean), - 'n_adopted_problems': adopted, 'adoption_rate': (adopted / len(pv)) if pv else 0.0} - - -def _mean(xs: List[float]) -> float: - return sum(xs) / len(xs) if xs else 0.0 - - -def _std(xs: List[float]) -> float: - if len(xs) < 2: - return 0.0 - m = _mean(xs) - return (sum((x - m) ** 2 for x in xs) / len(xs)) ** 0.5 - - -def _signal_stats(problems: List[Dict[str, Any]]) -> Dict[str, float]: - """The heart of 'is there a learning signal': per problem, the scored candidates form a - GRPO group. A group with zero reward variance (all skills solve, or none do -- the - hard-problem dead zone) gives NO gradient. Tracks that fraction plus reward level and - within-group variance so a collapse (all-0 or all-1) is visible immediately.""" - group_vars, all_rewards, zero_grad, groups = [], [], 0, 0 - for r in problems: - rewards = [c['reward'] for c in r['_cands'] if c.get('reward') is not None] - if len(rewards) < 2: - continue - groups += 1 - all_rewards.extend(rewards) - v = _std(rewards) - group_vars.append(v) - if v < 1e-9: # every skill got the same reward -> GRPO skips this problem - zero_grad += 1 - return {'n_groups': groups, 'zero_grad_frac': (zero_grad / groups) if groups else 0.0, - 'reward_mean': _mean(all_rewards), 'reward_std': _std(all_rewards), - 'group_reward_std_mean': _mean(group_vars)} - - -def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespace) -> Dict[str, Any]: - all_cands = [c for r in chunk for c in r['_cands']] - cands = [c for c in all_cands if c['parseable']] - scored = [c for c in cands if c['with_pass'] is not None] - clean = [c for c in cands if c['leaked'] is False] - ws_rolls = [x for c in scored for x in c['rolls']] - # viewa-dropped problems generate no candidates; keep acc/* on the generated subset - # so the with-skill/lift trend stays comparable across view_b_frac settings. - gen_probs = [r for r in chunk if r['_cands']] - base_acc = _mean([1.0 if r['_baseline_pass'] else 0.0 for r in gen_probs]) - ws_acc = _mean([1.0 if any(c.get('reward') for c in r['_cands']) else 0.0 for r in gen_probs]) - cand_pass_parseable = _mean([c['with_pass'] for c in scored]) - cand_pass_all = _mean([1.0 if c.get('reward') else 0.0 for c in all_cands]) - # base failure taxonomy (you asked whether skills fail because the base loops out of length) - classes = [_baseline_class(r) for r in chunk] - n_fail = sum(1 for c in classes if c != 'success') - skill_tokens = [c.get('skillgen_tokens') or 0 for c in cands] # skill-gen response length - trunc = sum(1 for r in chunk for c in r['_cands'] - for x in c['rolls'] if x['stop_reason'] == 'length') - rubric_answer_leaks = sum( - 1 for r in chunk - if r.get('_view') == 'A' and _rubric_has_fail(r.get('_rubric_diag')) and _rubric_answer_leaked(r)) - return { - 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), - 'n_failed_first_try': sum(1 for r in chunk if r['_failed']), - 'n_generated': len(all_cands), 'n_candidates_parseable': len(cands), - 'n_unparseable': len(all_cands) - len(cands), - 'parse_rate': (len(cands) / len(all_cands)) if all_cands else 0.0, - 'n_leaked': sum(1 for c in cands if c['leaked']), 'n_clean': len(clean), - 'leak_rate': (sum(1 for c in cands if c['leaked']) / len(cands)) if cands else 0.0, - 'n_reward_pos': sum(1 for c in scored if c['reward']), - 'n_rubric_answer_leaked': rubric_answer_leaks, - 'n_viewa_dropped': sum(1 for r in chunk if _viewa_dropped(r, args)), - 'n_train_samples': sum(1 for c in all_cands if _is_trainable(c, args)), - 'signal': _signal_stats(chunk), - 'fail_loop_frac': (sum(1 for c in classes if c == 'fail_loop') / n_fail) if n_fail else 0.0, - 'fail_wrong_frac': (sum(1 for c in classes if c == 'fail_wrong') / n_fail) if n_fail else 0.0, - 'skill_tokens_mean': _mean(skill_tokens), - 'withskill_trunc_frac': (trunc / len(ws_rolls)) if ws_rolls else 0.0, - 'avg_baseline_pass': base_acc, 'avg_withskill_pass': ws_acc, - 'avg_lift': ws_acc - base_acc, - 'candidate_withskill_pass_parseable': cand_pass_parseable, - 'candidate_withskill_pass_all': cand_pass_all, - 'termination_rate_withskill': _mean([1.0 if x['terminated'] else 0.0 for x in ws_rolls]), - 'view_A': _view_stats(chunk, 'A'), 'view_B': _view_stats(chunk, 'B'), - **_xproblem_stats(chunk, args), - } - - -def _xproblem_stats(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> Dict[str, Any]: - """Cross-problem pairing health: of the view-A problems, how many actually got a - neighbour's rubric (pair_rate) and how similar those neighbours were. Empty when off.""" - if not args.xproblem_rubric: - return {} - view_a = [r for r in chunk if r.get('_view') == 'A'] - paired = [r for r in view_a if r.get('_rubric_src')] - return {'xproblem': { - 'n_view_a': len(view_a), 'n_paired': len(paired), - 'pair_rate': (len(paired) / len(view_a)) if view_a else 0.0, - 'neighbor_sim_mean': _mean([r.get('_neighbor_sim', 0.0) for r in paired])}} - - -def _sft_record(r: Dict[str, Any], c: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: - """A view-A context-distillation sample: the OPEN-BOOK (rubric-conditioned) skill ``c`` - is the target, but the prompt is rebuilt as query-only (view B, no rubric) so the model - learns to produce it CLOSED-BOOK. Rides the same GRPOLoss with a positive constant - advantage (``--sft-weight``); single-step (old_logps=None) this reduces to - ``-w*logp + beta*KL`` -- a KL-anchored weighted cross-entropy toward the skill.""" - return { - 'data_id': r.get('data_id'), - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'view': 'B', 'diagnosis': '', 'rubric_src': '', 'orig_view': 'A', 'sft': True, - 'response': c['response'], 'skills': c['skills'], - 'skillgen_stop': c.get('skillgen_stop'), - 'advantage': float(args.sft_weight), 'grpo_adv': 0.0, 'kept': True, - 'reward': c['reward'], 'with_pass': c['with_pass']} - - -def _rubric_answer_leaked(r: Dict[str, Any]) -> bool: - """Hard safety gate for external API/rubric diagnosis. Unlike answer-bearing skills - generated by the policy itself, a rubric that contains the target final answer is an - external teacher leak and must not be distilled into view B.""" - return _answer_leaked(r.get('_rubric_diag', ''), r.get('reference_answer', '')) - - -def _viewa_dropped(r: Dict[str, Any], args: argparse.Namespace) -> bool: - """Under --viewa-sft a view-A problem trains ONLY through the SFT path (rubric with - [FAIL], answer-free). No-FAIL and rubric-leaked problems produce no training record - at all (no GRPO backflow: those prompts are query-only and would muddy the pure - view-B GRPO pool), so their skill-gen + executor scoring is skipped entirely.""" - return (bool(args.viewa_sft) and r.get('_view') == 'A' - and (not _rubric_has_fail(r.get('_rubric_diag')) - or _rubric_answer_leaked(r))) - - -def _group_records(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> List[Dict[str, Any]]: - """Training records. Under --viewa-sft, view A is SFT-only: a problem whose rubric - localised a failure (has ``[FAIL]``, answer-free) contributes ONE context-distillation - SFT sample (best parseable open-book skill -- preferring an executor-verified pass, - else any parseable candidate -- rebuilt query-only); no-FAIL / rubric-leaked view-A - problems are dropped (``_viewa_dropped``, generation already skipped). GRPO candidates - come from view B only, keeping that pool purely closed-book. Prompts are rebuilt from - the stored view/diagnosis by ``_skillgen_messages``.""" - out = [] - for r in chunk: - if not r['_hard']: - continue - if args.viewa_sft and r.get('_view') == 'A': - if _viewa_dropped(r, args): - continue - best = _best_sft_candidate(r, args) - if best is not None: - out.append(_sft_record(r, best, args)) - continue - for c in r['_cands']: - if _is_trainable(c, args): - out.append({ - 'data_id': r.get('data_id'), - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'view': r.get('_view', 'A'), 'diagnosis': r.get('_rubric_diag', ''), - 'rubric_src': r.get('_rubric_src', ''), 'sft': False, - 'response': c['response'], 'skills': c['skills'], - 'skillgen_stop': c.get('skillgen_stop'), - 'advantage': c['advantage'], 'grpo_adv': c['grpo_adv'], 'kept': c['kept'], - 'reward': c['reward'], 'with_pass': c['with_pass']}) - return out - - -# =========================================================================== -# Block G -- online GRPO training -# =========================================================================== -def _is_num(v: Any) -> bool: - try: - float(v) - return True - except (TypeError, ValueError): - return False - - -def _train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: - """Training sample = the exact skill-gen prompt (rebuilt by ``_skillgen_messages`` so - train/inference match) + the generated structured guidance response. ``key_rounds`` - selects the final assistant turn; Template masks the prompt and trains the whole - response (the key-round prefix already excludes the prompt-provided ).""" - msgs = _skillgen_messages( - rec['problem'], rec.get('view', 'A'), rec.get('diagnosis', ''), rec.get('rubric_src', '')) - return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], - 'user_data': {'key_rounds': [len(msgs)]}} - - -def _train_chunk(skill_model, ref_model, ckpt: CheckpointEngineManager, samples: List[Dict[str, Any]], - args: argparse.Namespace) -> Dict[str, Any]: - """On-policy GRPO update over one chunk, then sync weights. Micro-batches of - ``sft_batch_size`` accumulate gradients; ONE optimizer step is taken per PPO - mini-batch of ``ppo_mini_batch_size`` samples (0 -> a single step over the whole - chunk, the original behaviour). A frozen reference model provides ref_logps for the - SEAM-style KL penalty. - - Multi-step correctness: with more than one step over the SAME rollout, later - mini-batches see an already-updated policy, so we FREEZE the sampling-policy - ``old_logps`` (recomputed once, before any step) and let GRPOLoss form the PPO ratio - against them. A single step keeps ``old_logps=None`` (ratio==1, pure on-policy). - The batch is padded to a multiple of ``sft_batch_size`` with advantage-0 copies that - contribute no policy gradient. View-A context-distillation samples ride the same loss - with a positive constant advantage (``--sft-weight``); single-step (old_logps=None) - that is ``-w*logp + beta*KL`` (KL-anchored weighted cross-entropy).""" - trajs = [_train_trajectory(rec) for rec in samples] - advs = [float(rec['advantage']) for rec in samples] - rem = (-len(trajs)) % args.sft_batch_size - if rem: - trajs += [trajs[-1]] * rem - advs += [0.0] * rem - - n, sft = len(trajs), args.sft_batch_size - mini = args.ppo_mini_batch_size if args.ppo_mini_batch_size > 0 else n - mini = max(sft, (mini // sft) * sft) # align to a whole number of micro-batches - multi_step = mini < n - - # Freeze ref_logps (frozen model) and -- only when we take multiple steps -- the - # sampling-policy old_logps, BOTH before any optimizer step touches the weights. With - # a single step old_logps stays None so GRPOLoss uses ratio==1 (pure on-policy). - micro_ref, micro_old = [], [] - for i in range(0, n, sft): - mb = trajs[i:i + sft] - micro_ref.append(ref_model.forward_only(inputs=mb).get('logps')) - micro_old.append(skill_model.forward_only(inputs=mb).get('logps') if multi_step else None) - - micro, n_steps = 0, 0 - for ms in range(0, n, mini): - for i in range(ms, min(ms + mini, n), sft): - k = i // sft - skill_model.forward_backward(inputs=trajs[i:i + sft], - advantages=advs[i:i + sft], - old_logps=micro_old[k], - ref_logps=micro_ref[k]) - micro += 1 - skill_model.clip_grad_and_step() - n_steps += 1 - ckpt.sync_weights(merge_and_sync=True) - metric = skill_model.calculate_metric(is_training=True) - n_sft = sum(1 for s in samples if s.get('sft')) - return {'n_samples': len(samples), 'n_sft': n_sft, 'n_grpo': len(samples) - n_sft, - 'n_steps': n_steps, 'n_micro_batches': micro, - 'metric': {k: (float(v) if _is_num(v) else v) for k, v in (metric or {}).items()}} - - -# =========================================================================== -# Block H -- fixed-holdout eval + metric formatting -# =========================================================================== -def run_greedy_eval(base_sampler, skill_sampler, eval_records: List[Dict[str, Any]], - ci: int, rounds: int, base_dp: int, skill_dp: int, args: argparse.Namespace, - base_cache: DiskCache - ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], Dict[str, float]]: - """SEAM ``val-core/math/acc/mean@1`` on the fixed holdout: ONE greedy skill (T=0) per - problem into ONE greedy base solve (T=0). Eval ALWAYS uses view B (query-only, the - deployment form: no rubric, since rubric needs an online teacher unavailable at deploy); - no leak filter (acc scores correctness alone). Baseline reuses the disk cache.""" - baseline_rollout(base_sampler, eval_records, base_dp, args, base_cache) - for r in eval_records: - r['_view'], r['_rubric_diag'] = 'B', '' - sg_out = _run_samples(skill_sampler, [_view_prompt(r, args) for r in eval_records], - 1, args.skill_max_tokens, skill_dp, temperature=0.0) - skills = [] - for seqs in sg_out: - if not seqs: - skills.append(('', '')) - continue - sresp = _clean_text(getattr(seqs[0], 'decoded', '') or '') - skills.append((_extract_skill(sresp) or '', sresp)) - ws_out = _run_samples(base_sampler, - [build_skill_solve_prompt(r['problem'], sk) for r, (sk, _) in zip(eval_records, skills)], - 1, args.max_tokens, base_dp, temperature=0.0) - recs = [] - for r, (sk, sresp), seqs in zip(eval_records, skills, ws_out): - roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - recs.append({ - 'record_type': 'eval_problem', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, - 'data_id': r.get('data_id'), - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'view': r['_view'], 'rubric_diag': r.get('_rubric_diag', ''), - 'baseline_pass': r['_baseline_pass'], 'skill': sk, 'skill_parseable': bool(sk), - 'skill_response': sresp, 'withskill_pred': roll['pred'], - 'withskill_correct': roll['correct'], 'withskill_terminated': roll['terminated'], - 'withskill_stop_reason': roll['stop_reason'], 'withskill_text': roll['text'], - }) - acc = lambda rs: sum(1 for x in rs if x['withskill_correct']) / len(rs) if rs else 0.0 - ws = acc(recs) # all view B (deployment form) - base = sum(x['baseline_pass'] for x in recs) / len(recs) if recs else 0.0 - fmt = (sum(1 for x in recs if x['skill_parseable']) / len(recs)) if recs else 0.0 - term = (sum(1 for x in recs if x['withskill_terminated']) / len(recs)) if recs else 0.0 - summary = {'record_type': 'eval_summary', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, - 'n': len(recs), 'view': 'B', 'acc_mean1': ws, - 'baseline_acc_mean1': base, 'lift_mean1': ws - base, - 'format_mean1': fmt, 'term_mean1': term} - metrics = {'core/math/acc/mean@1': ws, 'core/math/baseline_acc/mean@1': base, - 'core/math/lift/mean@1': ws - base, 'core/math/format/mean@1': fmt, - 'core/math/term/mean@1': term} - return recs, summary, metrics - - -def _trend_line(hist: List[Dict[str, float]], window: int, rounds_done: int) -> Optional[str]: - """Contrast the FIRST vs the most recent ``window`` chunks -- if RFT works, adoption - and lift on recent (fresh) chunks exceed the early baseline.""" - if len(hist) < 2 * window: - return None - base, rec = hist[:window], hist[-window:] - m = lambda xs, k: sum(h[k] for h in xs) / len(xs) - return (f'[trend] first {window} vs last {window} | ' - f'adopt A {m(base,"aA"):.2f}->{m(rec,"aA"):.2f} B {m(base,"aB"):.2f}->{m(rec,"aB"):.2f} | ' - f'lift {m(base,"lift"):+.3f}->{m(rec,"lift"):+.3f} | ' - f'0grad {m(base,"zero_grad"):.2f}->{m(rec,"zero_grad"):.2f} | ' - f'pos/chunk {m(base,"pos"):.1f}->{m(rec,"pos"):.1f} | rounds={rounds_done}') - - -def _swan_metrics(summary: Dict[str, Any], log: Optional[Dict[str, Any]]) -> Dict[str, float]: - """Flat swanlab dict. ``signal/*`` is the primary health group (is GRPO getting a - gradient at all); ``acc/*`` is the effect; the rest diagnose why. acc/adopt/skill are - only logged when the chunk produced a scored group, so idle chunks don't dip charts.""" - sig = summary['signal'] - d: Dict[str, float] = { - # --- signal: the FIRST thing to watch (no variance -> no learning) --- - 'signal/zero_grad_frac': sig['zero_grad_frac'], 'signal/n_groups': sig['n_groups'], - 'signal/reward_mean': sig['reward_mean'], 'signal/reward_std': sig['reward_std'], - 'signal/group_reward_std_mean': sig['group_reward_std_mean'], - 'signal/n_train_samples': summary['n_train_samples'], - 'signal/n_reward_pos': summary['n_reward_pos'], - # --- skill format / leak health --- - 'skill/parse_rate': summary['parse_rate'], 'skill/tokens_mean': summary['skill_tokens_mean'], - 'leak/rate': summary['leak_rate'], 'leak/n': summary['n_leaked'], - # --- base failure taxonomy (loop-out-of-length vs plain wrong) --- - 'fail/loop_frac': summary['fail_loop_frac'], 'fail/wrong_frac': summary['fail_wrong_frac'], - 'fail/frac_first_try': (summary['n_failed_first_try'] / summary['n']) if summary['n'] else 0.0, - # --- view-A routing: dropped = no-FAIL rubric or rubric leak (SFT-only view A) --- - 'viewa/dropped_frac': (summary['n_viewa_dropped'] / summary['view_A']['n'] - if summary['view_A']['n'] else 0.0), - } - bal = summary.get('balance') or {} - if bal.get('enabled'): - d.update({'balance/n_drawn': bal['n_drawn'], 'balance/n_baseline_fresh': bal['n_baseline_fresh'], - 'balance/selected_success_frac': bal['selected_success_frac']}) - xp = summary.get('xproblem') or {} - if xp: - d.update({'xproblem/pair_rate': xp['pair_rate'], 'xproblem/neighbor_sim_mean': xp['neighbor_sim_mean']}) - if sig['n_groups'] > 0: - d.update({'acc/baseline_pass': summary['avg_baseline_pass'], - 'acc/withskill_pass': summary['avg_withskill_pass'], 'acc/lift': summary['avg_lift'], - 'candidate/withskill_pass_parseable': summary['candidate_withskill_pass_parseable'], - 'candidate/withskill_pass_all': summary['candidate_withskill_pass_all'], - 'adopt/A': summary['view_A']['adoption_rate'], - 'adopt/B': summary['view_B']['adoption_rate'], - 'term/withskill': summary['termination_rate_withskill'], - 'term/withskill_trunc_frac': summary['withskill_trunc_frac']}) - if log: - d['train/n_steps'] = log['n_steps'] - d['train/n_micro_batches'] = log['n_micro_batches'] - for k, v in (log.get('metric') or {}).items(): - if not _is_num(v): - continue - if k.startswith('learning rate'): - if 'group 1' in k: - d['train/lr'] = float(v) - else: - d[f'train/{k.replace(" ", "_")}'] = float(v) - return d - - -def _view_a_rubric_leak_metrics(chunk: List[Dict[str, Any]], - pool_answers: Optional[Dict[str, str]] = None) -> Dict[str, float]: - """Swanlab-only audit for answer leakage in view-A rubric text. This never changes - rewards, advantages, filtering, or training records.""" - view_a = [r for r in chunk if r.get('_view') == 'A'] - with_diag = [r for r in view_a if r.get('_rubric_diag')] - target_leaks = sum(1 for r in with_diag - if _answer_leaked(r.get('_rubric_diag', ''), r.get('reference_answer', ''))) - source_leaks = 0 - pool_answers = pool_answers or {} - for r in with_diag: - src = r.get('_rubric_src') - src_ref = pool_answers.get(src, '') if src else r.get('reference_answer', '') - if _answer_leaked(r.get('_rubric_diag', ''), src_ref): - source_leaks += 1 - n = len(with_diag) - return { - 'rubric_leak/n_view_a': float(len(view_a)), - 'rubric_leak/n_checked': float(n), - 'rubric_leak/target_answer_n': float(target_leaks), - 'rubric_leak/target_answer_rate': (target_leaks / n) if n else 0.0, - 'rubric_leak/source_answer_n': float(source_leaks), - 'rubric_leak/source_answer_rate': (source_leaks / n) if n else 0.0, - } - - -# =========================================================================== -# Block F -- components, args, main -# =========================================================================== -def init_components(args: argparse.Namespace): - """Default 8-GPU layout: rank 0 trains the actor, rank 1 hosts a frozen ref model, - 2-3 skill_sampler (synced), 4-7 base_sampler (frozen). Returns - (skill_model, ref_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp).""" - r0, r1 = TRAIN_GPUS, TRAIN_GPUS + REF_GPUS - r2, r3 = r1 + SKILL_SAMPLER_GPUS, NUM_GPUS - twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, lazy_collect=False, groups=[ - DeviceGroup(name='train', ranks=list(range(0, r0)), device_type='GPU'), - DeviceGroup(name='ref', ranks=list(range(r0, r1)), device_type='GPU'), - DeviceGroup(name='skill_sampler', ranks=list(range(r1, r2)), device_type='GPU'), - DeviceGroup(name='base_sampler', ranks=list(range(r2, r3)), device_type='GPU')]) - - train_mesh = DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_DP, fsdp_size=TRAIN_FSDP) - skill_model = TransformersModel(model_id=MODEL_ID, device_mesh=train_mesh, remote_group='train', - ddp_config={'find_unused_parameters': False}) - skill_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) - skill_model.set_template(Template, model_id=MODEL_ID, enable_thinking=False, - max_length=args.max_model_len, truncation_strategy='delete') - skill_model.set_processor(InputProcessor, padding_free=False) - skill_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon, beta=args.kl_beta) - skill_model.set_optimizer('AdamW', lr=args.lr) - skill_model.set_lr_scheduler('CosineWarmupScheduler', num_warmup_steps=10, - num_training_steps=args.max_train_rounds) - - ref_mesh = DeviceMesh.from_sizes(world_size=REF_GPUS, dp_size=REF_DP, fsdp_size=REF_FSDP) - ref_model = TransformersModel(model_id=MODEL_ID, device_mesh=ref_mesh, remote_group='ref', - ddp_config={'find_unused_parameters': False}) - ref_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) - ref_model.set_template(Template, model_id=MODEL_ID, enable_thinking=True, - max_length=args.max_model_len, truncation_strategy='delete') - ref_model.set_processor(InputProcessor, padding_free=False) - ref_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) - - def _sampler(group, world, enable_thinking: bool = True): - s = vLLMSampler(model_id=MODEL_ID, - engine_args={'gpu_memory_utilization': GPU_MEM, - 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, - device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), - remote_group=group) - s.set_template(Template, model_id=MODEL_ID, enable_thinking=enable_thinking, max_length=args.max_model_len) - return s - - skill_sampler = _sampler('skill_sampler', SKILL_SAMPLER_GPUS, enable_thinking=False) - # base_sampler is shared by the main thread and the baseline-prefetch thread -> serialise. - base_sampler = _LockedSampler(_sampler('base_sampler', BASE_SAMPLER_GPUS, enable_thinking=True)) - ckpt = CheckpointEngineManager(model=skill_model, sampler=skill_sampler) - return skill_model, ref_model, skill_sampler, base_sampler, ckpt, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS - - -def _build_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--dataset', choices=('aops', 'math'), default='aops') - p.add_argument('--n', type=int, default=2000, help='Problems loaded into the draw pool.') - p.add_argument('--pool-offset', type=int, default=0, - help='Skip this many shuffled non-eval records before building the train pool; ' - 'useful to avoid cold-start SFT data ranges.') - p.add_argument('--exclude-data-ids', default='', - help='Comma-separated jsonl files whose data_id/problem keys are excluded ' - 'from train/eval selection, e.g. coldstart_sft.jsonl.') - p.add_argument('--seed', type=int, default=42) - p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) - p.add_argument('--eval-size', type=int, default=128, help='Fixed holdout size (0 disables).') - p.add_argument('--eval-every', type=int, default=5, help='Run holdout eval every N chunks.') - p.add_argument('--chunk-size', type=int, default=16) - p.add_argument('--balance', action=argparse.BooleanOptionalAction, default=True) - p.add_argument('--balance-success-frac', type=float, default=0.4, - help='Target fraction of the chunk the base solves (rest are base-fail).') - p.add_argument('--balance-loop-frac', type=float, default=0.5) - p.add_argument('--balance-max-draws-mult', type=int, default=8) - p.add_argument('--n-skills', type=int, default=8) - p.add_argument('--view-b-frac', type=float, default=0.5) - p.add_argument('--viewa-frac-start', type=float, default=None, - help='Enable the view-A curriculum: chunk 0 uses this view-A share ' - '(view_b_frac = 1 - share), decaying linearly to --viewa-frac-end ' - 'over --viewa-decay-chunks chunks, then holding. Overrides ' - '--view-b-frac for every chunk.') - p.add_argument('--viewa-frac-end', type=float, default=0.1) - p.add_argument('--viewa-warmup-chunks', type=int, default=0, - help='Hold the view-A share at --viewa-frac-start for this many chunks ' - 'before the linear decay begins.') - p.add_argument('--viewa-decay-chunks', type=int, default=40) - p.add_argument('--xproblem-rubric', action=argparse.BooleanOptionalAction, default=False, - help='When enabled, view A uses a bag-of-words NEIGHBOUR problem rubric. ' - 'Default is off: each view-A problem uses its own baseline attempt, ' - 'while the API diagnosis prompt is constrained to be answer-free and ' - 'method-level only.') - p.add_argument('--skill-retries', type=int, default=2) - p.add_argument('--skill-gen-temperature', type=float, default=1.0) - p.add_argument('--skill-gen-top-p', type=float, default=1.0) - p.add_argument('--skill-gen-top-k', type=int, default=-1) - p.add_argument('--max-model-len', type=int, default=16384) - p.add_argument('--max-tokens', type=int, default=8192) - p.add_argument('--skill-max-tokens', type=int, default=8192) - p.add_argument('--rubric-workers', type=int, default=16) - p.add_argument('--rubric-retries', type=int, default=2, - help='Retry failed/timeout rubric diagnose calls this many times before ' - 'falling back to an empty diagnosis without caching the failure.') - p.add_argument('--sft-batch-size', type=int, default=8, - help='Driver micro-batch (gradient-accumulation unit); multiple of train dp.') - p.add_argument('--ppo-mini-batch-size', type=int, default=0, - help='Samples per optimizer step (SEAM-style PPO mini-batch). 0 = one step ' - 'over the whole chunk (on-policy, ratio==1). When >0 and smaller than ' - 'the trainable count, multiple steps are taken over the same rollout and ' - 'old_logps are frozen so the PPO ratio/clip stays valid. Rounded down to ' - 'a multiple of --sft-batch-size.') - p.add_argument('--grpo-epsilon', type=float, default=0.2) - p.add_argument('--adv-clip', type=float, default=3.0, - help='Symmetric clip for group-relative advantages; <=0 disables clipping.') - p.add_argument('--kl-beta', type=float, default=0.001, - help='SEAM-style reference KL coefficient for GRPOLoss.') - p.add_argument('--viewa-sft', action=argparse.BooleanOptionalAction, default=True, - help='Route view-A problems to online context distillation (SFT on the best ' - 'passing open-book skill, rebuilt as a query-only prompt) instead of GRPO. ' - 'View B stays GRPO; both share one optimizer step.') - p.add_argument('--sft-weight', type=float, default=0.5, - help='Advantage magnitude (lambda) for view-A SFT distillation samples in the ' - 'shared GRPOLoss. Single-step this is -lambda*logp + beta*KL (weighted CE). ' - 'Keep <~1 so one-signed SFT gradients do not swamp mean-0 GRPO advantages.') - p.add_argument('--sft-target-len', type=int, default=550, - help='Target skill length (chars) for view-A SFT distillation: among passing ' - 'candidates pick the one CLOSEST to this length (not the shortest). Set near ' - 'the empirical pass-rate peak (~500-600 here) so distillation neither shrinks ' - 'rollouts toward zero nor lets them grow unbounded.') - p.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True) - p.add_argument('--lr', type=float, default=6e-6) - p.add_argument('--max-train-rounds', type=int, default=1500) - p.add_argument('--save-rounds', type=int, default=200) - p.add_argument('--trend-every', type=int, default=10) - p.add_argument('--output-dir', default='./output/reflexion_skill') - p.add_argument('--cache-dir', default='', help='Baseline/rubric cache dir (default /cache).') - p.add_argument('--no-cache', action='store_true', help='Disable disk cache read/write.') - p.add_argument('--prefetch-baseline', action=argparse.BooleanOptionalAction, default=True, - help='Prefetch next chunk base baseline on a background thread (overlaps ' - 'with skill-gen; base_sampler is frozen so it never blocks the trainer).') - p.add_argument('--swanlab-project', default='twinkle') - p.add_argument('--swanlab-exp', default='') - args = p.parse_args() - if args.sft_batch_size % TRAIN_DP != 0: - raise ValueError(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple of train dp ({TRAIN_DP})') - if args.chunk_size < 1: - raise ValueError('--chunk-size must be >= 1') - args.balance_success_frac = max(0.0, min(1.0, args.balance_success_frac)) - return args - - -def _write(handle, row: Dict[str, Any]) -> None: - handle.write(json.dumps(row, ensure_ascii=False) + '\n') - - -def main() -> None: - args = _build_args() - records, eval_records, neighbor_map, pool_answers, data_stats = _load_records(args) - if len(records) < args.chunk_size: - raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded ({len(records)}); raise --n') - - os.makedirs(args.output_dir, exist_ok=True) - gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') - eval_path = os.path.join(args.output_dir, 'eval_records.jsonl') - data_path = os.path.join(args.output_dir, 'skill_dataset.jsonl') - train_log_path = os.path.join(args.output_dir, 'train_log.jsonl') - - if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY')): - sys.stderr.write('[rft] WARNING: no LLM backup env; view-A rubric check disabled ' - '(leak filter is deterministic, unaffected)\n') - - use_swan = swanlab is not None and os.environ.get('SWANLAB_MODE') != 'disabled' - if use_swan: - swanlab.init(project=args.swanlab_project, experiment_name=(args.swanlab_exp or None), - config={'model': MODEL_ID, 'dataset': args.dataset, 'n': len(records), - 'eval_n': len(eval_records), 'n_skills': args.n_skills, - 'view_b_frac': args.view_b_frac, 'balance': args.balance, - 'balance_success_frac': args.balance_success_frac, - 'skill_gen_temp': args.skill_gen_temperature, - 'grpo_epsilon': args.grpo_epsilon, 'kl_beta': args.kl_beta, - 'lr': args.lr}) - - skill_model, ref_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp = init_components(args) - checker = build_rubric_checker() - if checker is None: - sys.stderr.write('[rft] no LLM backup env -> view-A rubric process-check DISABLED\n') - - cache_dir = args.cache_dir or os.path.join(args.output_dir, 'cache') - os.makedirs(cache_dir, exist_ok=True) - use_cache = not args.no_cache - base_cache = DiskCache(os.path.join(cache_dir, 'baseline.jsonl'), use_cache) - eval_base_cache = DiskCache(os.path.join(cache_dir, 'eval_baseline.jsonl'), use_cache) - rubric_cache = DiskCache(os.path.join(cache_dir, 'rubric.jsonl'), use_cache) - if args.xproblem_rubric: - sys.stderr.write(f'[rft] xproblem-rubric ON: {len(neighbor_map)} bag-of-words pairs\n') - - cfg = {'record_type': 'config', 'model': MODEL_ID, 'dataset': args.dataset, - 'n': len(records), 'eval_n': len(eval_records), 'seed': args.seed, - 'pool_offset': args.pool_offset, 'exclude_data_ids': args.exclude_data_ids, - 'excluded_records': data_stats.get('excluded_records', 0), - 'numeric_only': args.numeric_only, 'raw_loaded': data_stats['raw_loaded'], - 'numeric_dropped': data_stats['numeric_dropped'], 'eval_every': args.eval_every, - 'n_skills': args.n_skills, 'view_b_frac': args.view_b_frac, - 'viewa_frac_start': args.viewa_frac_start, 'viewa_frac_end': args.viewa_frac_end, - 'viewa_warmup_chunks': args.viewa_warmup_chunks, - 'viewa_decay_chunks': args.viewa_decay_chunks, - 'skill_retries': args.skill_retries, 'balance': args.balance, - 'balance_success_frac': args.balance_success_frac, - 'skill_gen_temp': args.skill_gen_temperature, 'reward': 'greedy_binary(correct)', - 'advantage': 'group_relative', 'format_in_reward': args.format_in_reward, 'cache': use_cache, - 'rubric_check': f'fixed_math_{len(_MATH_RUBRIC)}crit+term(viewA)' if checker else 'disabled', - 'xproblem_rubric': args.xproblem_rubric, - 'viewa_sft': args.viewa_sft, 'sft_weight': args.sft_weight, - 'sft_target_len': args.sft_target_len, - 'adv_clip': args.adv_clip, - 'grpo_epsilon': args.grpo_epsilon, 'kl_beta': args.kl_beta, 'lr': args.lr, - 'train_gpus': TRAIN_GPUS, 'ref_gpus': REF_GPUS, 'ref_fsdp': REF_FSDP, - 'train_fsdp': TRAIN_FSDP, 'train_dp': TRAIN_DP, - 'skill_sampler_gpus': SKILL_SAMPLER_GPUS, 'base_sampler_gpus': BASE_SAMPLER_GPUS, - 'max_train_rounds': args.max_train_rounds, 'started': int(time.time())} - sys.stderr.write(f'[rft] raw={data_stats["raw_loaded"]} numeric_drop={data_stats["numeric_dropped"]} ' - f'train={len(records)} eval={len(eval_records)} {args.dataset}; ' - f'train_gpus={TRAIN_GPUS} ref_gpus={REF_GPUS} train_fsdp={TRAIN_FSDP} ' - f'train_dp={TRAIN_DP} skill_dp={skill_dp} base_dp={base_dp}\n') - - hist: List[Dict[str, float]] = [] - rounds = 0 - pool = ProblemPool(records, args.seed) - with open(gen_path, 'w', encoding='utf-8') as gen_f, \ - open(eval_path, 'w', encoding='utf-8') as eval_f, \ - open(data_path, 'w', encoding='utf-8') as data_f, \ - open(train_log_path, 'w', encoding='utf-8') as tlog: - for f in (gen_f, eval_f, data_f, tlog): - _write(f, cfg) - gstep = 0 - # base baseline (frozen sampler, disk-cached, no weight-sync) is prefetched on a - # background thread while the current chunk generates: the skill-gen phase uses - # skill_sampler, so prefetching the NEXT chunk's baseline on base_sampler overlaps - # it and fills base_cache so the next draw_chunk hits it. base_sampler is wrapped in - # _LockedSampler, so the prefetch and this chunk's with-skill scoring never issue a - # base .sample() concurrently. It never touches the trainer or on-policy generation. - prefetch_pool = ThreadPoolExecutor(max_workers=1) if args.prefetch_baseline else None - pending: Optional[Any] = None - - def _prefetch(peeked: List[Dict[str, Any]]) -> None: - if peeked: - baseline_rollout(base_sampler, peeked, base_dp, args, base_cache) - - # Baseline (round 0) eval BEFORE any training: measures the untrained skill model on - # the fixed holdout so every later eval has a step-0 reference point on the same axis. - if eval_records: - eval_recs, eval_summary, eval_metrics = run_greedy_eval( - base_sampler, skill_sampler, eval_records, -1, rounds, base_dp, skill_dp, - args, eval_base_cache) - for rec in eval_recs: - _write(eval_f, rec) - _write(eval_f, eval_summary) - eval_f.flush() - if use_swan: - swanlab.log({f'eval/{k}': v for k, v in eval_metrics.items()}, step=0) - sys.stderr.write( - f'[eval] g-1 (init): n={eval_summary["n"]} viewB mean@1 ' - f'acc={eval_summary["baseline_acc_mean1"]:.3f}->{eval_summary["acc_mean1"]:.3f} ' - f'lift={eval_summary["lift_mean1"]:+.3f} ' - f'fmt={eval_summary["format_mean1"]:.2f} rounds={rounds}\n') - - # Each chunk is drawn fresh + RE-GENERATED with the current policy (on-policy); - # with --balance, draw_chunk keeps drawing until the base fail:success mix hits target. - while rounds < args.max_train_rounds: - if pending is not None: - pending.result() # finish last round's prefetch before drawing (cache-warm) - pending = None - if args.viewa_frac_start is not None: - args.view_b_frac = _curriculum_view_b_frac(gstep, args) - chunk, balance = draw_chunk(pool, base_sampler, base_dp, args, base_cache) - if prefetch_pool is not None: - peeked = pool.peek(int(args.chunk_size * (args.balance_max_draws_mult if args.balance else 1))) - pending = prefetch_pool.submit(_prefetch, peeked) - full, summary, groups = process_chunk( - base_sampler, skill_sampler, chunk, gstep, base_dp, skill_dp, - args, checker, rubric_cache, base_cache, neighbor_map, pool_answers) - summary['balance'] = balance - summary['view_b_frac'] = round(args.view_b_frac, 4) - - log = None - if groups: - log = _train_chunk(skill_model, ref_model, ckpt, groups, args) - rounds += 1 - log.update({'record_type': 'train_round', 'round': rounds, 'chunk': gstep, - 'epoch': pool.epoch, 'ts': int(time.time())}) - _write(tlog, log) - tlog.flush() - if rounds % args.save_rounds == 0: - skill_model.save(f'skill-rft-{rounds}', output_dir=args.output_dir) - - summary['rounds_done'], summary['epoch'] = rounds, pool.epoch - for rec in full: - _write(gen_f, rec) - _write(gen_f, summary) - gen_f.flush() - for v in groups: - _write(data_f, v) - data_f.flush() - - sa, sb, sig = summary['view_A'], summary['view_B'], summary['signal'] - hist.append({'aA': sa['adoption_rate'], 'aB': sb['adoption_rate'], - 'lift': summary['avg_lift'], 'pos': summary['n_reward_pos'], - 'zero_grad': sig['zero_grad_frac']}) - bal_str = (f'bal {balance["selected_fail"]}f/{balance["selected_success"]}s ' - f'(drew {balance["n_drawn"]}/fresh {balance["n_baseline_fresh"]}' - + ('!' if balance.get('budget_hit') else '') + ') ') if balance.get('enabled') else '' - xp = summary.get('xproblem') - xp_str = f'pair={xp["pair_rate"]:.2f}@{xp["neighbor_sim_mean"]:.2f} ' if xp else '' - tr_str = (f'train={log["n_grpo"]}g/{log["n_sft"]}s ' if log - else f'train={summary["n_train_samples"]} ') - sys.stderr.write( - f'[gen] e{pool.epoch} g{gstep}: {bal_str}n={summary["n"]} ' - f'clean={summary["n_clean"]} leak={summary["leak_rate"]:.2f} {tr_str}' - f'0grad={sig["zero_grad_frac"]:.2f} R={sig["reward_mean"]:.2f}±{sig["reward_std"]:.2f} ' - f'acc={summary["avg_baseline_pass"]:.2f}->{summary["avg_withskill_pass"]:.2f} ' - f'lift={summary["avg_lift"]:+.3f} {xp_str}' - f'A[{sa["n"]} {sa["adoption_rate"]:.2f}] B[{sb["n"]} {sb["adoption_rate"]:.2f}] ' - f'rounds={rounds}\n') - if use_swan: - swan_metrics = _swan_metrics(summary, log) - swan_metrics.update(_view_a_rubric_leak_metrics(chunk, pool_answers)) - swan_metrics['train/view_b_frac'] = float(args.view_b_frac) - swanlab.log(swan_metrics, step=gstep) - - if eval_records and (gstep + 1) % args.eval_every == 0: - eval_recs, eval_summary, eval_metrics = run_greedy_eval( - base_sampler, skill_sampler, eval_records, gstep, rounds, base_dp, skill_dp, - args, eval_base_cache) - for rec in eval_recs: - _write(eval_f, rec) - _write(eval_f, eval_summary) - eval_f.flush() - if use_swan: - swanlab.log({f'eval/{k}': v for k, v in eval_metrics.items()}, step=gstep) - sys.stderr.write( - f'[eval] g{gstep}: n={eval_summary["n"]} viewB mean@1 ' - f'acc={eval_summary["baseline_acc_mean1"]:.3f}->{eval_summary["acc_mean1"]:.3f} ' - f'lift={eval_summary["lift_mean1"]:+.3f} ' - f'fmt={eval_summary["format_mean1"]:.2f} rounds={rounds}\n') - - if (gstep + 1) % args.trend_every == 0: - tl = _trend_line(hist, args.trend_every, rounds) - if tl: - sys.stderr.write(tl + '\n') - gstep += 1 - - if prefetch_pool is not None: - if pending is not None: - pending.result() - prefetch_pool.shutdown(wait=True) - base_cache.close() - eval_base_cache.close() - rubric_cache.close() - skill_model.save('skill-rft-final', output_dir=args.output_dir) - sys.stderr.write(f'[rft] done: {rounds} rounds over {gstep} chunks / {pool.epoch} epochs -> {data_path}\n') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/embedding/train_reflexion_skill_replay.py b/cookbook/exp/embedding/train_reflexion_skill_replay.py deleted file mode 100644 index 0580a6ba4..000000000 --- a/cookbook/exp/embedding/train_reflexion_skill_replay.py +++ /dev/null @@ -1,114 +0,0 @@ -"""Replay-train the reflexion skill model from prebuilt exact RFT data. - -Use ``build_reflexion_skill_data.py`` first to create ``skill_dataset.jsonl``. This -script trains only the skill model from those frozen records; it does not run vLLM -rollouts, leak checks, or rubric diagnosis. - -Launch: - python cookbook/exp/embedding/train_reflexion_skill_replay.py \ - --data ./output/reflexion_skill_data/skill_dataset.jsonl -""" -import argparse -import json -import os -import sys -from collections import defaultdict -from typing import Any, Dict, List - -import train_reflexion_skill_rft as rft - - -def _build_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description=__doc__) - p.add_argument('--data', default='./output/reflexion_skill_data/skill_dataset.jsonl') - p.add_argument('--output-dir', default='./output/reflexion_skill_replay') - p.add_argument('--epochs', type=int, default=1) - p.add_argument('--sft-batch-size', type=int, default=8) - p.add_argument('--max-model-len', type=int, default=16384) - p.add_argument('--grpo-epsilon', type=float, default=0.2) - p.add_argument('--lr', type=float, default=1e-5) - p.add_argument('--save-rounds', type=int, default=50) - return p.parse_args() - - -def _load_chunks(path: str) -> List[List[Dict[str, Any]]]: - chunks: Dict[int, List[Dict[str, Any]]] = defaultdict(list) - fallback_chunk = 0 - with open(path, 'r', encoding='utf-8') as f: - for line_no, line in enumerate(f, 1): - if not line.strip(): - continue - row = json.loads(line) - if row.get('record_type') == 'config': - continue - for key in ('problem', 'response', 'advantage'): - if key not in row: - raise ValueError(f'{path}:{line_no} missing required field {key!r}') - ci = int(row.get('chunk', fallback_chunk)) - chunks[ci].append(row) - if 'chunk' not in row and len(chunks[ci]) >= 64: - fallback_chunk += 1 - return [chunks[k] for k in sorted(chunks) if chunks[k]] - - -def _init_model(args: argparse.Namespace, total_updates: int): - model = 'ms://Qwen/Qwen3-4B' - train_mesh = rft.DeviceMesh.from_sizes( - world_size=rft.TRAIN_GPUS, dp_size=rft.TRAIN_DP, fsdp_size=rft.TRAIN_FSDP) - device_groups = [ - rft.DeviceGroup(name='train', ranks=list(range(rft.TRAIN_GPUS)), device_type='GPU'), - ] - rft.twinkle.initialize(mode='ray', nproc_per_node=rft.TRAIN_GPUS, groups=device_groups, - lazy_collect=False) - model = rft.TransformersModel(model_id=model, device_mesh=train_mesh, - remote_group='train', ddp_config={'find_unused_parameters': False}) - from twinkle.patch.no_split_modules import NoSplitModulesPatch - model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) - model.set_template(rft.Template, model_id=model, - enable_thinking=True, max_length=args.max_model_len, - truncation_strategy='delete') - model.set_processor(rft.InputProcessor, padding_free=False) - model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) - model.set_optimizer('AdamW', lr=args.lr) - model.set_lr_scheduler('CosineWarmupScheduler', num_warmup_steps=10, - num_training_steps=max(1, total_updates)) - return model - - -def main() -> None: - args = _build_args() - if args.sft_batch_size % rft.TRAIN_DP != 0: - raise ValueError(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple ' - f'of the training dp size ({rft.TRAIN_DP})') - chunks = _load_chunks(args.data) - if not chunks: - raise ValueError(f'no train records found in {args.data}') - os.makedirs(args.output_dir, exist_ok=True) - total_updates = len(chunks) * args.epochs - model = _init_model(args, total_updates) - log_path = os.path.join(args.output_dir, 'train_log.jsonl') - cfg = {'record_type': 'config', 'mode': 'offline_replay', 'data': args.data, - 'chunks': len(chunks), 'epochs': args.epochs, 'lr': args.lr, - 'sft_batch_size': args.sft_batch_size} - rounds = 0 - with open(log_path, 'w', encoding='utf-8') as tlog: - tlog.write(json.dumps(cfg, ensure_ascii=False) + '\n') - for epoch in range(args.epochs): - for ci, samples in enumerate(chunks): - log = rft._train_chunk(model, None, samples, args) - rounds += 1 - log.update({'record_type': 'train_round', 'round': rounds, - 'epoch': epoch, 'chunk': ci}) - tlog.write(json.dumps(log, ensure_ascii=False) + '\n') - tlog.flush() - sys.stderr.write( - f'[replay-rft] e{epoch} c{ci}: n={log["n_samples"]} ' - f'micro={log["n_micro_batches"]} metric={log.get("metric")}\n') - if rounds % args.save_rounds == 0: - model.save(f'skill-rft-replay-{rounds}', output_dir=args.output_dir) - model.save('skill-rft-replay-final', output_dir=args.output_dir) - sys.stderr.write(f'[replay-rft] done: {rounds} updates; log -> {log_path}\n') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/embedding/train_reflexion_skill_rft.py b/cookbook/exp/embedding/train_reflexion_skill_rft.py deleted file mode 100644 index b4388b773..000000000 --- a/cookbook/exp/embedding/train_reflexion_skill_rft.py +++ /dev/null @@ -1,1568 +0,0 @@ -"""RFT cold-start for the reflexion skill generator (see reflexion.md §6). - -Trains an INDEPENDENT skill model to write reusable, transferable skills that, -when injected into a FROZEN base solver's system prompt, let the base solve problems -it first got wrong. The base is never trained — it only produces the reward signal. -Scoring is SEAM-style DETERMINISTIC: the base runs each candidate skill once at -temperature 0 (M=1), so the reward ``R in {0,1}`` (answer correct) carries no -sampling noise; the per-candidate advantage is group-relative within a problem -(``A = (R - mean) / (std + eps)``) and the skill model is updated online by GRPO — -problem-groups where every skill scores alike (std=0) contribute no gradient. - -Direction: skill GENERATION + recall. Skill-gen always runs with thinking ON, and -each hard problem is routed to EXACTLY ONE of two views (no reuse — kills memory -leak and holds cost at 1x): view A ``(problem + attempt) -> think + skills`` keeps -the online generator self-bootstrapping; view B ``(problem only) -> think + skills`` -is the deployment form, where the think is grounded on the query alone so it cannot -hallucinate an attempt. Both share the verified skill; the distilled ```` -block is recalled into the base's system prompt at solve time. - -8-GPU layout (three DeviceGroups, one twinkle.initialize): - - ranks 0-3 : ``train`` — skill model, full-param FSDP2, dp=4 - - ranks 4-5 : ``skill_sampler`` — skill model rollouts (vLLM, tp1 dp2) - - ranks 6-7 : ``base_sampler`` — frozen base solver (vLLM, tp1 dp2) -CheckpointEngineManager syncs train -> skill_sampler after every optimizer step; -base_sampler is never synced. - -Leak filtering uses ``LeakVerifier(sampler=None)`` via the backup teacher API -(no local judge, no distillation): set LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / -LLM_BACKUP_MODEL. - -Launch (8 GPUs): - LLM_BACKUP_API_KEY=... LLM_BACKUP_BASE_URL=... LLM_BACKUP_MODEL=... \ - python cookbook/exp/embedding/train_reflexion_skill_rft.py --n 2000 --chunk-size 16 -""" -import argparse -import hashlib -import json -import os -import re -import sys -import time -from typing import Any, Dict, List, Optional, Tuple - -import numpy as np - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.checkpoint_engine import CheckpointEngineManager -from twinkle.model import TransformersModel -from twinkle.processor import InputProcessor -from twinkle.sampler import vLLMSampler -from twinkle.template import Template -from twinkle_agentic.verifier import LeakVerifier, RubricVerifier -from twinkle_agentic.verifier.rubric_verifier import RubricItem - -# Reuse the reference eval's dataset + grading + prompts + sampling config, and the -# phase-0 pipeline's parsing / rollout / injection helpers (Find > Create). -from eval_gpqa_rag import (GEN_GPU_MEM, GEN_MODEL_ID, build_direct_prompt, # noqa: F401 - load_aops, load_math) -from eval_reflexion_skill import (_EX_PROBLEM, _clean_text, # noqa: F401 - _parse_seq, _run_samples, build_skill_solve_prompt) - -logger = get_logger() - -try: - import swanlab -except ImportError: # optional; metric logging degrades to stdout + jsonl only - swanlab = None - - -# -- GPU layout --------------------------------------------------------------- -TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 4)) -SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) -BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 2)) -NUM_GPUS = TRAIN_GPUS + SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS -# FSDP shard group size within a dp replica; TRAIN_DP is the data-parallel axis that -# ``forward_backward`` (slice_dp) splits each mini-batch over. -TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', 2)) -TRAIN_DP = max(1, TRAIN_GPUS // TRAIN_FSDP) - - -# --------------------------------------------------------------------------- -# Skill-generation prompt (DISTILL the useful approach, per the new direction) -# --------------------------------------------------------------------------- -# --- Previous STRICT view-A system prompt (commented out; kept for easy revert). It -# hard-required 3-5 bullets, "output nothing after ", one-imperative-sentence -# items, no-narration, and a strict do-not-reveal block. The soft SEAM-style version -# below drops those four format demands, frames the skills as advisory reminders, and -# explains how they are used. --- -# SKILL_GEN_SYSTEM = ( -# 'You are distilling reusable problem-solving SKILLS from one worked episode. ' -# 'You are shown a competition problem, the guidance the solver was given, and the ' -# "solver's own attempt (its reasoning may be partly right and partly wrong).\n\n" -# 'FIRST, in your private thinking, do ALL of: (a) work out what this TYPE of problem ' -# 'fundamentally requires; (b) pinpoint WHERE THIS attempt actually went wrong ' -# '(when a process-check report is provided below, use its flagged criteria as ' -# 'evidence, but confirm each against the attempt yourself) — ' -# 'the decisive misstep, a missing idea, a wrong turn, or the way it stalled, looped ' -# 'on the same step, or ran the length budget out without ever committing to an ' -# 'answer; and (c) imagine AS MANY DIFFERENT angles as you can — distinct approaches ' -# 'or representations that could crack this problem, alternative solution paths, and ' -# 'the various ways a solver could plausibly go wrong on it (a few words each, do NOT ' -# 'develop them fully). THEN commit to the angle you find most decisive and write a ' -# 'SHORT list of skills that would have PREVENTED that specific ' -# 'failure and would raise the success rate of a SIMILAR solver on SIMILAR problems. ' -# 'Ground each skill in the concrete mistake you found, but state it as a GENERAL, ' -# 'transferable rule — not a patch hard-coded to this problem. Across the 3-5 ' -# 'bullets, prioritise in this order:\n' -# '1. the decisive method or representation this class of problem calls for (what to ' -# 'set up or reach for first);\n' -# '2. the specific mistake that derailed THIS attempt, recast as a general pitfall, ' -# 'plus the quick check that catches it;\n' -# '3. a compact, ordered procedure that reliably drives toward a final answer;\n' -# '4. convergence discipline: once the key quantity is in hand, commit to a single ' -# 'concrete final answer in the required format instead of re-deriving, endless ' -# 'case-splitting, looping on the same check, or overrunning the length budget.\n\n' -# 'OUTPUT FORMAT (strict):\n' -# '- Keep the thinking COMPACT — a quick brainstorm of angles then a decision, not a ' -# 'full solution and not a re-statement of these instructions. AFTER it, ' -# 'output ONLY a markdown bullet list of 3-5 items WRAPPED IN and ' -# 'tags — no preamble, no narration outside the tags. Output nothing after ' -# '.\n' -# '- Each item is ONE short imperative sentence (a method, heuristic, check, or ' -# 'habit).\n' -# '- Inside the tags: no narration, no "The student...", no headings, no restating ' -# 'the problem.\n\n' -# 'CONTENT RULES (strict):\n' -# '- Do NOT reveal the final answer or the multiple-choice option.\n' -# '- Do NOT state the specific numbers, values, or key intermediate results of THIS ' -# 'problem.\n' -# '- Every item must be GENERAL and transferable, not a step-by-step solution to ' -# 'THIS problem.\n\n' -# 'Follow the example below for the exact tags, style, and level of generality.' -# ) -SKILL_GEN_SYSTEM = ( - 'You are a mathematics coach. You are shown a competition problem together with an ' - 'automated process-check of an earlier solver attempt at it -- which solution ' - 'criteria the attempt passed or failed, and suggested fixes for the failures. You ' - 'do NOT see the attempt itself, only this check. Treat the check as privileged ' - 'training scaffolding: study it together with the problem, identify the ' - 'problem-visible features that make each useful flagged failure relevant, then ' - 'rephrase those lessons as self-contained reusable skills. The goal is not to ' - 'continue from the check, cite it, or hide it silently; the goal is to turn it to ' - 'a skill pattern which prevents the model falls into similar pitfalls in the future.\n\n' - 'Good skills name the observable trigger, the method worth reaching for, the ' - 'pitfall to watch, and a quick verification habit. Prefer formulations like ' - '"When a configuration has ...", "Before setting up ...", or "Check whether ..." ' - 'over references to the process-check, failed criteria, or the earlier attempt. ' - "These tips are advisory: they will be placed in a solver's system prompt as gentle " - 'reminders before it works through a SIMILAR problem on its own, without seeing ' - 'this process-check. So keep them general and transferable rather than a worked ' - 'solution to this exact problem, and do not state its specific intermediate values ' - 'or final answer. Think briefly first, then give your tips as a markdown bullet ' - 'list wrapped in and , like the example below.' -) - -# One-shot demo of the recommended mix (method / pitfall+check / procedure / -# convergence), answer-free — anchors both the format and the content priorities. -_EX_SKILLS = ( - '\n' - '- Rewrite each square root by factoring its radicand into a perfect square times ' - 'a remainder, then move the perfect-square factor outside.\n' - '- Avoid the classic trap $\\sqrt{a}+\\sqrt{b}\\ne\\sqrt{a+b}$; only combine terms ' - 'sharing the same simplest radical, and sanity-check by estimating each root.\n' - '- Procedure: simplify every radical, group like radical terms, add their ' - 'coefficients, then reduce to simplest form.\n' - '- Once the expression is in simplest form, commit to that single result as the ' - 'final answer rather than re-checking indefinitely.\n' - '') - - -# View A user template: the problem + the automated rubric process-check of an earlier -# attempt (PASS/FAIL per criterion + suggested fixes). The attempt trajectory is NOT -# shown -- the rubric findings are the evidence the skill model grounds its tips on, -# which avoids feeding the (often long, non-terminating) attempt into the prompt. -SKILL_GEN_USER_RUBRIC = ( - 'Problem:\n{problem}\n\n' - 'Process check of an earlier attempt (automated rubric verifier -- treat as ' - 'evidence, not gospel; PASS/FAIL per criterion with suggested fixes for failures):\n' - '{diagnosis}\n\n' - 'Now output a self-contained skills bullet list. Each bullet should still be useful ' - 'if the process check were removed: connect any useful flagged failure to ' - 'problem-visible features, general methods, and quick checks rather than citing the ' - 'rubric or the earlier attempt. \n\n' - 'Note: **Do not solve the problem, only generate skills**. Now Begin:' -) - - -def build_skillgen_prompt(problem: str, diagnosis: str) -> Dict[str, Any]: - """View A skill-gen prompt: system + one-shot format demo + the real episode - (problem + the rubric process-check of an earlier attempt). The attempt trajectory - is deliberately NOT shown -- the rubric findings localise the failure without the - generator having to re-chew (and often re-solve) a long, possibly non-terminating - attempt. The one-shot demo is query-only; only the real turn carries the diagnosis.""" - return {'messages': [ - {'role': 'system', 'content': SKILL_GEN_SYSTEM}, - # {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=_EX_PROBLEM)}, - # {'role': 'assistant', 'content': _EX_SKILLS}, - {'role': 'user', - 'content': SKILL_GEN_USER_RUBRIC.format(problem=problem, diagnosis=diagnosis)}, - ]} - - -# --------------------------------------------------------------------------- -# View B: query-only skill-gen (deployment form). No attempt is shown — the model -# must reason about the problem TYPE from the query alone, so the think is grounded -# on the query and cannot narrate/fabricate an attempt. Format is deliberately -# distinct from view A so the model learns the two modes as separate contracts. -# --------------------------------------------------------------------------- -# --- Previous STRICT view-B (query-only) system prompt (commented out; kept for revert). -# Same four format demands as the old view A. Soft SEAM-style version below. --- -# SKILL_GEN_SYSTEM_Q = ( -# 'You are distilling reusable problem-solving SKILLS for a CLASS of problems. You ' -# 'are shown ONE competition problem and NOTHING else — no solution, no attempt. ' -# 'FIRST, in your private thinking, imagine AS MANY DIFFERENT angles as you can — ' -# 'distinct approaches or representations that could crack this TYPE of problem, ' -# 'alternative solution paths, and the various ways a solver could plausibly go wrong ' -# 'on it (a few words each, do NOT develop them fully). THEN commit to what you find ' -# 'most decisive and write a SHORT list of skills that would raise a solver\'s success ' -# 'rate on SIMILAR problems. Across the 3-5 bullets, prioritise in this order:\n' -# '1. the decisive method or representation this class of problem calls for (what to ' -# 'set up or reach for first);\n' -# '2. the specific pitfall that derails such problems, plus the quick check that ' -# 'catches it;\n' -# '3. a compact, ordered procedure that reliably drives toward a final answer;\n' -# '4. convergence discipline: once the key quantity is in hand, commit to a single ' -# 'concrete final answer in the required format instead of re-deriving, endless ' -# 'case-splitting, or overrunning the length budget.\n\n' -# 'OUTPUT FORMAT (strict):\n' -# '- Keep the thinking COMPACT — a quick brainstorm of angles then a decision, not a ' -# 'full solution and not a re-statement of these instructions. AFTER ' -# 'it, output ONLY a markdown bullet list of 3-5 items WRAPPED IN and ' -# ' tags — no preamble, no narration outside the tags. Output nothing after ' -# '.\n' -# '- Each item is ONE short imperative sentence (a method, heuristic, check, or ' -# 'habit).\n' -# '- Inside the tags: no narration, no headings, no restating the problem, and no ' -# 'reference to any attempt, student, or solution.\n\n' -# 'CONTENT RULES (strict):\n' -# '- Do NOT solve THIS problem or reveal its final answer or multiple-choice option.\n' -# '- Do NOT state the specific numbers, values, or key intermediate results of THIS ' -# 'problem.\n' -# '- Every item must be GENERAL and transferable to other problems of the same ' -# 'type.\n\n' -# 'Follow the example below for the exact tags, style, and level of generality.' -# ) -SKILL_GEN_SYSTEM_Q = ( - 'You are a mathematics coach. You are shown ONE competition problem and nothing ' - 'else — no solution and no attempt. Think about what approach this KIND of problem ' - 'calls for and where solvers tend to slip, then distil a few reusable tips.\n\n' - "These tips are advisory: they will be placed in a solver's system prompt as gentle " - 'reminders before it works through a SIMILAR problem on its own. So keep them ' - 'general and transferable — the method worth reaching for, the pitfall to watch and ' - 'a quick check, and the discipline to settle on a final answer — rather than a ' - 'worked solution to this exact problem, and without stating its specific ' - 'intermediate values or its final answer. Think briefly first, then give your tips ' - 'as a markdown bullet list wrapped in and , like the example below.' -) - -SKILL_GEN_USER_Q = ( - 'Problem:\n{problem}\n\n' - 'Now reason about this TYPE of problem, then output the skills bullet list.' -) - - -def build_querygen_prompt(problem: str) -> Dict[str, Any]: - """View B skill-gen prompt: system + one-shot demo + the problem ALONE (no - attempt) — matching what is available at deployment (query only).""" - return {'messages': [ - {'role': 'system', 'content': SKILL_GEN_SYSTEM_Q}, - # {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=_EX_PROBLEM)}, - # {'role': 'assistant', 'content': _EX_SKILLS}, - {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=problem)}, - ]} - - -def _assign_view(problem: str, args: argparse.Namespace) -> str: - """Deterministically route a problem to exactly one view (stable across restarts - and across the generation/SFT sides). ``--view-b-frac`` of problems go to view B.""" - h = int(hashlib.md5(f'{args.seed}:{problem}'.encode('utf-8')).hexdigest(), 16) - return 'B' if (h % 100000) / 100000.0 < args.view_b_frac else 'A' - - -def _skillgen_messages(problem: str, view: str, diagnosis: str) -> List[Dict[str, Any]]: - """Single source of truth for the skill-gen prompt, used at BOTH generation and - training time so they can never diverge. View A with a localisable failure uses - problem + rubric findings (NO trajectory); view B -- or a view-A problem whose rubric - flagged NO failure (``[FAIL]`` absent: all-pass or missing diagnosis) -- is query-only. - So view A DEGRADES to view B whenever there is nothing concrete to correct.""" - if view == 'B' or '[FAIL]' not in (diagnosis or ''): - return build_querygen_prompt(problem)['messages'] - return build_skillgen_prompt(problem, diagnosis)['messages'] - - -def _view_prompt(r: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: - """The skill-gen prompt for problem ``r`` under its assigned view (routing in - ``_skillgen_messages``: view A carries the rubric process-check; view B, and any - view-A problem with no rubric failure, is query-only).""" - return {'messages': _skillgen_messages(r['problem'], r['_view'], r.get('_rubric_diag', ''))} - - -_BULLET_RE = re.compile(r'(?m)^\s*(?:[-*]|\d+[.)])\s') -# Trajectory/meta references that betray CoT fragments leaking into the block; any -# hit fails the purity gate (the problem is then re-sampled, per --skill-retries). -_META_RE = re.compile( - r'\b(the student|the solver|the attempt|this attempt|the trace|the response|' - r'in the (?:attempt|trace|response|solution)|as (?:shown|seen|noted) above|' - r'the (?:above|previous|earlier)|my (?:reasoning|analysis)|i (?:think|need|will))\b', - re.IGNORECASE) - - -def _is_clean_block(block: str) -> bool: - """Purity gate for thinking-ON skill-gen: the block must be a pure bullet list - (every non-empty line a bullet — no prose/CoT fragments) with no meta/trajectory - reference. Answer leak is caught separately by the backup-teacher leak stage.""" - lines = [ln for ln in (l.strip() for l in block.splitlines()) if ln] - if not lines or not all(_BULLET_RE.match(ln) for ln in lines): - return False - return _META_RE.search(block) is None - - -def _extract_skills_block(text: str) -> Optional[str]: - """Return the clean ``...`` block, or None if not parseable. - - Skill-gen runs with thinking ON, so the model must end its reasoning with an explicit - ```` before committing an answer (whether the opening ```` is emitted by - the model or pre-injected by the chat template). We therefore REQUIRE ```` and - read only the text after the last one; its absence means the token budget was exhausted - mid-reasoning (nothing committed, per reflexion.md §6.8) — reject so a draft or a - system-prompt demo echo inside the CoT can never be mistaken for the answer. Within - the answer take the ```` block (closing tag optional), strip stray tags, and - require ``_is_clean_block`` — prose-mixed / meta-referencing fragments are rejected - for re-sampling.""" - low = text.lower() - end_think = low.rfind('') - if end_think < 0: - return None # reasoning never closed -> no committed answer - answer = text[end_think + len(''):] - low_a = answer.lower() - s = low_a.find('') - if s < 0: - return None - inner = s + len('') - e = low_a.find('', inner) - block = (answer[inner:e] if e >= 0 else answer[inner:]).strip() - block = re.sub(r'', '', block, flags=re.IGNORECASE).strip() - if not _is_clean_block(block): - return None - return block - - -# --------------------------------------------------------------------------- -# OPTIONAL stricter leak criterion (currently UNUSED -- the run uses answer_only=True, -# which flags ONLY the final answer). This variant ALSO flags concrete intermediate KEY -# results, while still permitting method / plan / pitfalls / checks. To enable, pass -# judge_system=_LEAK_JUDGE_SYSTEM to the LeakVerifier below. -# --------------------------------------------------------------------------- -_LEAK_JUDGE_SYSTEM = """\ -You check whether a HINT that will be shown to someone solving a math TASK gives away -this task's own results. - -The hint may FREELY describe the general method, which approach or technique to use, the -steps or plan to follow, common pitfalls, and sanity checks -- even when that points -strongly at HOW to solve THIS task. Describing the approach is expected of a good hint. - -The hint LEAKS only if, for THIS specific task, it states either: -- the final answer or final result (a value, expression, choice, label, or verbatim - output); or -- a concrete decisive INTERMEDIATE key result -- a specific computed value, quantity, or - fact unique to this task that hands over a key step of the answer. - -If it names only the method / plan / pitfalls / checks WITHOUT stating those concrete -intermediate values or the final result, it does NOT leak. - -Reply with exactly one word: LEAK or CLEAN.""" - - -# --------------------------------------------------------------------------- -# Rubric process-check (view A only): a frozen teacher diagnoses the base's failed -# attempt so the skill model grounds its error analysis on a verified fault -# localisation instead of guessing. Teacher-only (sampler=None -> every diagnose() -# hits llm_backup); mirrors eval_dualline_math's fixed math rubric. -# --------------------------------------------------------------------------- -_RFT_DIAG_SYSTEM = """\ -You are a process error checker for a math solution attempt. You are given a math -problem, a rubric, and one attempted solution segment. Decide PASS or FAIL for each -criterion and explain only the process error type. - -Output STRICT JSON (no prose outside it) with this shape: -{ - "items": [ - {"index": 1, "verdict": "PASS", "reason": "", - "fix": ""}, - {"index": 2, "verdict": "FAIL", "reason": "", - "fix": ""} - ], - "overall": "OK" | "ISSUES", - "summary": "" -} - -Rules: -- Judge every criterion independently and literally; a [Hard Rule] is FAIL unless - unambiguously satisfied. -- Judge ONLY what is observable in THIS segment. -- Content inside ... (or ) is internal reasoning, not - user-facing output; ignore it for "output only X" style criteria. -- For PASS items, leave "fix" as "". -- For FAIL items, "reason", "fix", and "summary" must describe only the flawed - step, theorem, arithmetic operation, case split, or verification habit. -- NEVER try to solve the query or state the correct final answer, corrected final expression, option letter, - graph/choice label, or any exact value that the answer should become. -- NEVER write phrases like "the correct answer is", "which gives", "yielding", - "should be ", "Option ", or "Graph ". -- If a fix would require naming a corrected value, replace it with a method-level - instruction such as "redo that computation carefully" or "apply the theorem with - the correct quantities". -- Keep every "reason" and "fix" clear and concise. -- "overall" is "OK" only if NO criterion is FAIL. -- Output only the JSON object.""" - - -_RFT_DIAG_USER = """\ -## Task / query (context) -{query} - -## Rubric -{rubric} - -## Segment -{segment} - -Now output the diagnostic JSON object.""" - - -class _RftRubricVerifier(RubricVerifier): - def _diagnose_trajectory(self, query: str, rubric_block: str, segment_text: str) -> dict: - user = _RFT_DIAG_USER.format(query=query, rubric=rubric_block, segment=segment_text) - return {'messages': [ - {'role': 'system', 'content': _RFT_DIAG_SYSTEM}, - {'role': 'user', 'content': user}, - ]} - - -_MATH_RUBRIC = [ - ('The reasoning contains no arithmetic or algebraic error', True), - ('Each step follows logically from the previous ones', True), - ('No formula or theorem is misstated or misapplied', True), - ('The approach is on track to answer the actual question asked', False), - ('No step contradicts an earlier established fact', False), -] - - -def _build_rubric_checker() -> Optional['RubricVerifier']: - """Fixed math-process rubric verifier, teacher-served. None if no LLM backup env.""" - if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') - or os.environ.get('OPENAI_API_KEY')): - return None - return _RftRubricVerifier( - fixed_rubric=[RubricItem(t, is_hard=h) for t, h in _MATH_RUBRIC], gate=True) - - -def _format_diagnosis(detail) -> str: - """One line per criterion (PASS/FAIL + reason + fix on FAIL) then a summary — the - compact evidence block appended to the view-A skill-gen prompt.""" - rub = detail.rubric - lines = [] - for it in detail.items: - text = rub[it.index - 1].text if 0 < it.index <= len(rub) else f'criterion {it.index}' - if it.verdict: - lines.append(f'- [PASS] {text}') - else: - tail = f': {it.reason}' if it.reason else '' - tail += f' (fix: {it.fix})' if it.fix else '' - lines.append(f'- [FAIL] {text}{tail}') - if detail.summary: - lines.append(f'Summary: {detail.summary}') - return '\n'.join(lines) - - -def _diagnose_views(checker, hard: List[Dict[str, Any]], args: argparse.Namespace, - diag_cache: Optional[Dict[str, str]] = None) -> None: - """Rubric-check every view-A problem's greedy attempt in parallel, stashing the - formatted findings on ``r['_rubric_diag']`` (view B stays empty). A checker error - or empty result degrades to no diagnosis (the plain view-A prompt).""" - from concurrent.futures import ThreadPoolExecutor - targets = [r for r in hard if r.get('_view') == 'A'] - if not checker or not targets: - return - - def _cache_key(r: Dict[str, Any]) -> str: - init_text = r.get('_init', [{}])[0].get('text', '') - return hashlib.md5(f'{r["problem"]}\n{init_text}'.encode('utf-8')).hexdigest() - - pending = [] - for r in targets: - key = _cache_key(r) - if diag_cache is not None and key in diag_cache: - r['_rubric_diag'] = diag_cache[key] - else: - pending.append((r, key)) - if not pending: - return - - def _run(item: Tuple[Dict[str, Any], str]) -> Tuple[Dict[str, Any], str, str, bool]: - r, key = item - seg = {'messages': [ - {'role': 'user', 'content': r['problem']}, - {'role': 'assistant', 'content': r['_init'][0]['text']}, - ]} - try: - return r, key, _format_diagnosis(checker.diagnose(seg, query=r['problem'])), True - except Exception as exc: # teacher hiccup -> fall back to no-diagnosis prompt - logger.warning(f'[rubric] diagnose error: {exc}') - return r, key, '', False - - workers = max(1, min(args.rubric_workers, len(pending))) - with ThreadPoolExecutor(max_workers=workers) as ex: - for r, key, diag, ok in ex.map(_run, pending): - r['_rubric_diag'] = diag - if ok and diag_cache is not None: - diag_cache[key] = diag - - -# --------------------------------------------------------------------------- -# Online data generation (one chunk; every candidate is recorded, untruncated) -# --------------------------------------------------------------------------- -def _roll(x: Dict[str, Any]) -> Dict[str, Any]: - """Full (untruncated) rollout record for offline analysis.""" - return {k: x[k] for k in ('pred', 'correct', 'terminated', 'passed', - 'stop_reason', 'gen_tokens', 'text')} - - -def _empty_roll() -> Dict[str, Any]: - """Fallback rollout when the sampler returned nothing for a prompt.""" - return {'pred': '', 'correct': False, 'terminated': False, 'passed': False, - 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} - - -def _assign_advantages(hard: List[Dict[str, Any]], args: argparse.Namespace) -> None: - """Group-relative advantage (SEAM-style) over each problem's clean, scored candidates, - using the DETERMINISTIC greedy reward ``R in {0, 1}`` (answer CORRECT only; - termination is NOT part of the reward -- monitored via `terminated`/`passed` only): - - A_j = (R_j - mean_R) / (std_R + eps) - - Groups where every candidate shares the same reward (``std_R == 0``: all solve or all - fail) get advantage 0 and contribute no gradient -- GRPO's own group variance - auto-selects the informative problems, so no explicit difficulty / marginal gate is - needed. Because the reward is deterministic (M=1 greedy, no pass@k sampling), the - std-normalisation no longer amplifies rollout noise (the reason it was dropped for the - old stochastic marginal). ``kept`` marks above-average candidates (for reporting only). - """ - eps = 1e-6 - for r in hard: - for c in r['_cands']: - c['advantage'], c['grpo_adv'], c['kept'] = 0.0, 0.0, False - cs = ([c for c in r['_cands'] if c.get('reward') is not None] if args.format_in_reward - else [c for c in r['_cands'] if c['leaked'] is False and c.get('reward') is not None]) - if len(cs) < 2: - continue - rewards = [c['reward'] for c in cs] - mean_r = sum(rewards) / len(rewards) - std = (sum((x - mean_r) ** 2 for x in rewards) / len(rewards)) ** 0.5 - if std < 1e-9: - continue # all candidates equal (all solve / all fail) -> no learning signal - for c in cs: - adv = (c['reward'] - mean_r) / (std + eps) - c['advantage'], c['grpo_adv'], c['kept'] = adv, adv, c['reward'] > mean_r - - -def _full_record(r: Dict[str, Any], ci: int) -> Dict[str, Any]: - """Complete per-problem record: init attempt, baseline, and ALL candidates - (parseable/leaked/scored alike) with full text — nothing dropped or truncated.""" - init = r['_init'][0] - rec: Dict[str, Any] = { - 'record_type': 'problem', 'chunk': ci, 'problem': r['problem'], - 'reference_answer': r['reference_answer'], 'level': r.get('level', ''), - 'failed_first_try': r['_failed'], - 'init_attempt': {'text': init['text'], 'pred': init['pred'], - 'correct': init['correct'], 'terminated': init['terminated'], - 'stop_reason': init['stop_reason'], 'gen_tokens': init['gen_tokens']}, - } - rec['baseline_pass'] = r['_baseline_pass'] - rec['is_hard'] = r['_hard'] - rec['view'] = r.get('_view', '') - rec['rubric_diag'] = r.get('_rubric_diag', '') - rec['baseline_rolls'] = [_roll(x) for x in r['_baseline_rolls']] - rec['candidates'] = [{ - 'skills': c['skills'], 'response': c['response'], 'parseable': c['parseable'], - 'leaked': c['leaked'], 'leak_reason': c['leak_reason'], 'leak_source': c['leak_source'], - 'with_pass': c['with_pass'], 'reward': c.get('reward'), - 'advantage': c.get('advantage'), 'grpo_adv': c.get('grpo_adv'), 'kept': c.get('kept'), - 'rolls': [_roll(x) for x in c['rolls']], - } for c in r['_cands']] - return rec - - -def _view_stats(hard: List[Dict[str, Any]], view: str) -> Dict[str, Any]: - """Per-view yield: hard problems, clean candidates, and the ADOPTION rate — - the fraction of hard problems that produced at least one clean, non-zero-advantage - candidate (i.e. a record that actually reaches training). Watching A vs B and - early vs late tells whether query-only (B) catches up to trajectory-grounded (A).""" - hv = [r for r in hard if r.get('_view') == view] - cands = [c for r in hv for c in r['_cands'] if c['parseable']] - clean = [c for c in cands if c['leaked'] is False] - adopted = sum(1 for r in hv - if any(c['leaked'] is False and abs(c.get('advantage') or 0.0) > 1e-9 - for c in r['_cands'])) - return { - 'n_hard': len(hv), 'n_candidates_parseable': len(cands), - 'n_clean': len(clean), 'n_adopted_problems': adopted, - 'adoption_rate': (adopted / len(hv)) if hv else 0.0, - } - - -def _is_trainable(c: Dict[str, Any], args: argparse.Namespace) -> bool: - """A candidate reaches the GRPO update iff its advantage is non-zero. With - --format-in-reward every candidate carries a reward (unparseable/leaked score 0), - so non-zero advantage is the only gate; otherwise it must also be clean and scored. - Single source of truth for both the summary counts and ``_group_records``.""" - adv_nz = abs(c.get('advantage') or 0.0) > 1e-9 - if args.format_in_reward: - return adv_nz - return c['leaked'] is False and c.get('with_pass') is not None and adv_nz - - -def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespace) -> Dict[str, Any]: - """Per-chunk aggregates — watch these across chunks to see if the RFT'd skill - model produces better skills over time (yield, leak rate, lift, termination).""" - failed = [r for r in chunk if r['_failed']] - hard = [r for r in chunk if r['_hard']] - all_cands = [c for r in chunk for c in r['_cands']] - cands = [c for c in all_cands if c['parseable']] - scored = [c for c in cands if c['with_pass'] is not None] - ws_rolls = [x for c in scored for x in c['rolls']] - # With --format-in-reward, unparseable/leaked candidates also carry a (0) reward and are - # trained, so count trainables over ALL candidates; else only clean scored ones. - train_cands = [c for c in all_cands if _is_trainable(c, args)] - base_acc = (sum(r['_baseline_pass'] for r in hard) / len(hard)) if hard else 0.0 - ws_acc = (sum(c['with_pass'] for c in scored) / len(scored)) if scored else 0.0 - # -- signal-source monitor: how much of the GRPO signal comes from base-FAIL problems - # (the offensive "rescue a failure" signal we want) vs base-success (defensive "don't - # break an easy one"). abs_adv_from_fail_frac ~0.1 was the diagnosed failure mode. -- - fail_cands = [c for r in chunk if r['_failed'] for c in r['_cands']] - abs_adv = lambda cs: sum(abs(c.get('advantage') or 0.0) for c in cs) - total_abs = abs_adv(all_cands) - return { - 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), - 'n_failed_first_try': len(failed), 'n_hard': len(hard), - 'n_generated': len(all_cands), - 'n_candidates_parseable': len(cands), - 'n_unparseable': len(all_cands) - len(cands), - 'n_leaked': sum(1 for c in cands if c['leaked']), - 'n_clean': sum(1 for c in cands if c['leaked'] is False), - 'n_reward_pos': sum(1 for c in scored if c['reward']), - 'n_train_samples': len(train_cands), - 'n_train_from_fail': sum(1 for c in fail_cands if _is_trainable(c, args)), - 'abs_adv_from_fail_frac': (abs_adv(fail_cands) / total_abs) if total_abs > 0 else 0.0, - 'avg_baseline_pass_on_hard': base_acc, - 'avg_withskill_pass': ws_acc, - 'avg_lift': ws_acc - base_acc, - 'termination_rate_withskill': (sum(1 for x in ws_rolls if x['terminated']) / len(ws_rolls)) if ws_rolls else 0.0, - 'view_A': _view_stats(hard, 'A'), 'view_B': _view_stats(hard, 'B'), - } - - -def _group_records(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> List[Dict[str, Any]]: - """GRPO training records: every clean, scored skill candidate with a NON-ZERO - advantage (positive pushes the skill up, negative down; the group-relative - usefulness-over-base advantage was set in _assign_advantages). Each carries its - ``view`` and the rubric ``diagnosis``; the prompt (identical to generation) is rebuilt - from those by ``_skillgen_messages`` -- no trajectory is stored or replayed.""" - out = [] - for r in chunk: - if not r['_hard']: - continue - view = r.get('_view', 'A') - for c in r['_cands']: - if _is_trainable(c, args): - out.append({ - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'view': view, - 'diagnosis': r.get('_rubric_diag', ''), - 'response': c['response'], 'skills': c['skills'], - 'advantage': c['advantage'], 'grpo_adv': c['grpo_adv'], 'kept': c['kept'], - 'reward': c['reward'], 'with_pass': c['with_pass'], - }) - return out - - -def _apply_baseline(r: Dict[str, Any], roll: Dict[str, Any]) -> None: - """Write a (cached or fresh) greedy baseline roll onto a problem and RESET the per-chunk - working state, so a problem reused in a later chunk never carries prior skill candidates.""" - r['_baseline_rolls'], r['_cands'] = [roll], [] - r['_init'] = [roll] - r['_failed'] = not roll['correct'] - r['_baseline_pass'] = 1.0 if roll['correct'] else 0.0 - r['_hard'] = True # process EVERY selected problem; group variance selects - - -def _baseline_rollout(base_sampler, problems: List[Dict[str, Any]], base_dp: int, - args: argparse.Namespace, cache: Dict[str, Dict[str, Any]]) -> int: - """Phase 1: base solves each problem GREEDILY once (T=0, M=1), keyed-cached by problem - text across chunks. The base sampler is FROZEN and decoding is greedy, so a problem's - baseline never changes over the run -- a cache hit is exact and skips the sampler. - Returns the number of FRESH sampler rollouts (cache misses) for efficiency reporting.""" - todo = [r for r in problems if r['problem'] not in cache] - if todo: - base_out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], - 1, args.max_tokens, base_dp, temperature=0.0) - for r, seqs in zip(todo, base_out): - cache[r['problem']] = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - for r in problems: - _apply_baseline(r, cache[r['problem']]) - return len(todo) - - -def _baseline_class(r: Dict[str, Any]) -> str: - """Bucket a baselined problem by its greedy outcome: ``success`` (base solved it), - ``fail_loop`` (ran the length budget out / never terminated -- the mode skills rescue - best), or ``fail_wrong`` (terminated cleanly but the answer is wrong).""" - roll = r['_init'][0] - if roll['correct']: - return 'success' - if roll['stop_reason'] == 'length' or not roll['terminated']: - return 'fail_loop' - return 'fail_wrong' - - -_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') - - -def _norm_num_text(num: str) -> str: - try: - f = float(num) - return str(int(f)) if f == int(f) else str(f) - except Exception: - return str(num).strip() - - -def _numeric_value(raw: Any) -> Optional[str]: - if raw is None: - return None - s = str(raw).strip().strip('$').strip() - s = s.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac') - s = re.sub(r'\\!|\\,|\\;|\\ |\\left|\\right|\s', '', s) - m = re.fullmatch(r'\\frac\{(-?\d+)\}\{(-?\d+)\}', s) - if m: - a, b = int(m.group(1)), int(m.group(2)) - return _norm_num_text(str(a / b)) if b else None - m = re.fullmatch(r'(-?\d+)/(-?\d+)', s) - if m: - a, b = int(m.group(1)), int(m.group(2)) - return _norm_num_text(str(a / b)) if b else None - if _NUM_RE.fullmatch(s): - return _norm_num_text(s) - return None - - -def _numeric_only_records(records: List[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], int]: - out = [] - dropped = 0 - for r in records: - ref = _numeric_value(r.get('reference_answer')) - if ref is None: - dropped += 1 - continue - rr = dict(r) - rr['reference_answer'] = ref - out.append(rr) - return out, dropped - - -def _load_records(args: argparse.Namespace) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], Dict[str, int]]: - need_split = args.eval_size > 0 - load_n = 0 if (args.numeric_only or need_split) else args.n - records = (load_aops(n=load_n, seed=args.seed) if args.dataset == 'aops' - else load_math(n=load_n, seed=args.seed)) - raw_n = len(records) - dropped = 0 - if args.numeric_only: - records, dropped = _numeric_only_records(records) - rng = np.random.RandomState(args.seed) - rng.shuffle(records) - eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 - eval_records = [dict(r) for r in records[:eval_n]] - train_pool = records[eval_n:] - train_n = args.n if args.n > 0 else len(train_pool) - train_records = [dict(r) for r in train_pool[:train_n]] - overlap = {r['problem'] for r in train_records} & {r['problem'] for r in eval_records} - if overlap: - raise ValueError(f'fixed eval/train overlap detected: {len(overlap)} duplicated problems') - stats = {'raw_loaded': raw_n, 'numeric_dropped': dropped, - 'train_records': len(train_records), 'eval_records': len(eval_records)} - return train_records, eval_records, stats - - -class _ProblemPool: - """Cyclic draw source over the loaded problems. Each full pass reshuffles with - ``seed + epoch`` and bumps ``epoch`` (matching the old per-epoch reshuffle); the - initial pass keeps the loader's shuffled order. Draws never run out.""" - - def __init__(self, records: List[Dict[str, Any]], seed: int): - self._records = list(records) - self._seed = seed - self._cursor = 0 - self.epoch = 0 - self.baseline_cache: Dict[str, Dict[str, Any]] = {} # problem text -> frozen greedy roll - - def draw(self, k: int) -> List[Dict[str, Any]]: - """Return ``k`` DISTINCT problems (unique within this call, so one chunk never - processes the same problem twice even when the cursor wraps mid-draw). ``k`` is - always << pool size, so this terminates.""" - out: List[Dict[str, Any]] = [] - seen: set = set() - while len(out) < k: - if self._cursor >= len(self._records): - self.epoch += 1 - np.random.RandomState(self._seed + self.epoch).shuffle(self._records) - self._cursor = 0 - r = self._records[self._cursor] - self._cursor += 1 - if id(r) not in seen: - seen.add(id(r)) - out.append(r) - return out - - -def _select_balanced(buckets: Dict[str, List[Dict[str, Any]]], n_success: int, - n_fail: int, n_fail_loop: int) -> List[Dict[str, Any]]: - """Pick the chunk from the baselined buckets: ``n_fail`` base-fails (split toward - ``n_fail_loop`` loop-fails, best-effort) + ``n_success`` base-successes. If a bucket - is too thin to hit ``chunk_size`` the shortfall is topped up from leftovers (the - ratio then drifts, which the caller logs).""" - loop, wrong, succ = buckets['fail_loop'], buckets['fail_wrong'], buckets['success'] - take_loop = min(n_fail_loop, len(loop)) - take_wrong = min(n_fail - take_loop, len(wrong)) - take_loop = min(n_fail - take_wrong, len(loop)) # give loop the remainder if wrong is short - sel = loop[:take_loop] + wrong[:take_wrong] + succ[:n_success] - target = n_success + n_fail - if len(sel) < target: - used = {id(x) for x in sel} - leftover = [x for b in (loop, wrong, succ) for x in b if id(x) not in used] - sel += leftover[:target - len(sel)] - return sel - - -def _draw_chunk(pool: _ProblemPool, base_sampler, base_dp: int, args: argparse.Namespace - ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: - """Draw one training chunk, running baseline rollout (Phase 1) on every drawn problem. - - With ``--balance`` off, draw ``chunk_size`` problems and return them. With it on, keep - drawing+baselining in ``chunk_size`` batches, bucketing by ``_baseline_class``, until the - target base fail:success mix is reachable or the draw budget is hit; then select a - balanced subset. Returns ``(chunk, stats)`` where stats records the realised mix.""" - if not args.balance: - chunk = pool.draw(args.chunk_size) - n_fresh = _baseline_rollout(base_sampler, chunk, base_dp, args, pool.baseline_cache) - return chunk, {'enabled': False, 'n_drawn': len(chunk), 'n_baseline_fresh': n_fresh} - - n_success = max(0, min(args.chunk_size, round(args.chunk_size * args.balance_success_frac))) - n_fail = args.chunk_size - n_success - n_fail_loop = round(n_fail * args.balance_loop_frac) - buckets: Dict[str, List[Dict[str, Any]]] = {'success': [], 'fail_loop': [], 'fail_wrong': []} - budget = args.chunk_size * args.balance_max_draws_mult - n_drawn, n_fresh = 0, 0 - seen: set = set() # dedupe across batches: the pool can re-serve a problem after a wrap - while n_drawn < budget: - if (len(buckets['success']) >= n_success - and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail): - break # enough of both classes buffered to satisfy the target split - batch = pool.draw(args.chunk_size) - n_fresh += _baseline_rollout(base_sampler, batch, base_dp, args, pool.baseline_cache) - n_drawn += len(batch) - for r in batch: - if id(r) in seen: - continue - seen.add(id(r)) - buckets[_baseline_class(r)].append(r) - - target_reached = (len(buckets['success']) >= n_success - and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail) - chunk = _select_balanced(buckets, n_success, n_fail, n_fail_loop) - sel_success = sum(1 for r in chunk if not r['_failed']) - stats = { - 'enabled': True, 'n_drawn': n_drawn, 'n_baseline_fresh': n_fresh, 'n_selected': len(chunk), - 'target_success': n_success, 'target_fail': n_fail, 'target_fail_loop': n_fail_loop, - 'selected_success': sel_success, 'selected_fail': len(chunk) - sel_success, - 'selected_fail_loop': sum(1 for r in chunk if _baseline_class(r) == 'fail_loop'), - 'selected_fail_wrong': sum(1 for r in chunk if _baseline_class(r) == 'fail_wrong'), - 'selected_success_frac': (sel_success / len(chunk)) if chunk else 0.0, - 'budget_hit': not target_reached, # stopped short of the target mix, not by choice - } - return chunk, stats - - -def process_chunk(base_sampler, skill_sampler, leak: LeakVerifier, - chunk: List[Dict[str, Any]], ci: int, base_dp: int, skill_dp: int, - args: argparse.Namespace, checker=None, - diag_cache: Optional[Dict[str, str]] = None - ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], List[Dict[str, Any]]]: - """base-solve -> rubric-check (view A) -> skill-gen -> leak-filter -> with-skill pass - -> GRPO advantages, for one chunk. - - Sequential (generate-one-chunk-train-one): generation and the trainer's weight - sync never overlap, so no lock is needed. ``base_sampler`` is frozen (never - synced); ``skill_sampler`` is synced by the trainer between chunks. - - ``chunk`` arrives ALREADY baselined by ``_draw_chunk`` (Phase 1 ran during the - balanced draw), so every problem carries ``_init``/``_failed``/``_baseline_pass``/ - ``_hard``/``_cands`` -- Phase 1 is not repeated here. - """ - # Phase 1 (base greedy solve) ran in _draw_chunk so the balancer could classify by - # outcome; every selected problem is processed (no difficulty gate, SEAM-style): the - # group-relative advantage (Phase 6) gives zero gradient to any problem whose skills - # all score alike, so GRPO's own group variance selects the informative problems. - hard = chunk - - # --- Phase 2: assign each problem's view, then rubric-check the view-A attempts so - # the skill model diagnoses from verified findings instead of guessing. View B is - # query-only and deliberately gets NO rubric (nothing to diagnose without an attempt). --- - for r in hard: - r['_view'] = _assign_view(r['problem'], args) - r['_rubric_diag'] = '' - _diagnose_views(checker, hard, args, diag_cache) - - # --- Phase 3: skill-gen (thinking ON), per-view prompt; re-sample empties. --- - flat: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] - if hard: - pending = list(hard) # problems still without any clean candidate - for _ in range(args.skill_retries + 1): - if not pending: - break - prompts = [_view_prompt(r, args) for r in pending] - sg_out = _run_samples(skill_sampler, prompts, args.n_skills, - args.skill_max_tokens, skill_dp, - temperature=args.skill_gen_temperature, - top_p=args.skill_gen_top_p, - top_k=args.skill_gen_top_k) - still = [] - for r, seqs in zip(pending, sg_out): - got = False - for s in seqs: - resp = _clean_text(getattr(s, 'decoded', '') or '') - block = _extract_skills_block(resp) - cand = {'skills': block or '', 'response': resp, 'parseable': bool(block), - 'view': r['_view'], 'leaked': None, 'leak_reason': '', - 'leak_source': '', 'with_pass': None, - 'reward': None, 'rolls': []} - r['_cands'].append(cand) - if block: - flat.append((r, cand)) - got = True - if not got: - still.append(r) # nothing parseable yet -> retry this problem - pending = still - - # --- Phase 4: leak filter via backup teacher (network only, no lock). VIEW A ONLY -- - # view B is query-only (no trajectory to leak from) and is left exactly like SEAM, which - # runs NO leak filter: its candidates skip the check and are treated as clean. To restore - # leak-checking on view B, drop the ``_view == 'A'`` guard below. --- - for r, c in flat: - if r.get('_view') != 'A': - c['leaked'], c['leak_reason'], c['leak_source'] = False, '', 'skipped_viewB' - flat_a = [(r, c) for r, c in flat if r.get('_view') == 'A'] - if flat_a: - details = leak.leak_batch( - [{'content': c['skills'], 'query': r['problem'], 'reference': r['reference_answer']} - for r, c in flat_a], max_workers=args.leak_workers) - for (r, c), d in zip(flat_a, details): - c['leaked'], c['leak_reason'], c['leak_source'] = bool(d.leaked), d.reason, d.source - - # --- Phase 5: with-skill GREEDY pass (T=0, M=1) on clean candidates. Binary reward - # R = answer CORRECT (deterministic, no pass@k noise), ABSOLUTE -- no baseline - # subtraction; the group mean in Phase 6 is the only baseline. Termination is NOT - # required (monitored only) -- see reflexion.md §7.6. --- - clean = [(r, c) for r, c in flat if c['leaked'] is False] - if clean: - ws_out = _run_samples(base_sampler, - [build_skill_solve_prompt(r['problem'], c['skills']) for r, c in clean], - 1, args.max_tokens, base_dp, temperature=0.0) - for (r, c), seqs in zip(clean, ws_out): - c['rolls'] = [_parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll()] - c['with_pass'] = 1.0 if c['rolls'][0]['correct'] else 0.0 - c['reward'] = c['with_pass'] # valid + clean + correct -> 1 - # Validity-in-reward (SEAM-style, --format-in-reward): every candidate that never reached - # the executor -- unparseable/impure format OR answer-leaked -- scores 0 and STILL joins its - # group, so its whole response (think tokens included) is trained DOWN. Off => those - # candidates are excluded, as before. - if args.format_in_reward: - for r in hard: - for c in r['_cands']: - if c['reward'] is None: - c['reward'] = 0.0 - - # --- Phase 6: group-relative GRPO advantage per problem-group. --- - _assign_advantages(hard, args) - - return ([_full_record(r, ci) for r in chunk], - _chunk_summary(chunk, ci, args), - _group_records(chunk, args)) - - -# --------------------------------------------------------------------------- -# Online RFT training -# --------------------------------------------------------------------------- -def _train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: - """Training sample = the exact skill-gen prompt for this record's view + the - generated (think + skills) response as the target; the GRPO advantage is attached - separately at forward_backward time. - - The prompt is rebuilt by ``_skillgen_messages`` (the same function used at generation), - so train/inference stay identical: view A replays problem + rubric findings, view B (and - no-failure view A) replays the query-only prompt. ``key_rounds`` selects the final - assistant turn (index ``len(msgs)``); the plain ``Template`` then masks the prompt and - trains the whole generated response (reasoning + ```` + skills) -- the key-round - prefix already excludes the prompt-provided ````, so no extra masking is needed.""" - msgs = _skillgen_messages(rec['problem'], rec.get('view', 'A'), rec.get('diagnosis', '')) - full = msgs + [{'role': 'assistant', 'content': rec['response']}] - return {'messages': full, 'user_data': {'key_rounds': [len(msgs)]}} - - -def _train_chunk(skill_model, ckpt: Optional[CheckpointEngineManager], - samples: List[Dict[str, Any]], args: argparse.Namespace) -> Dict[str, Any]: - """One on-policy GRPO optimizer update on THIS chunk's skill candidates, then sync weights. - - Sequential design (generate-one-chunk-train-one): the skills were sampled from the - current policy and trained immediately, so ``old_logps`` is omitted and the GRPO - ratio is ~1. All driver-side mini-batches accumulate into one optimizer step so the - whole rollout chunk stays under the same pre-update policy. The batch is padded to a - multiple of ``sft_batch_size`` with advantage-0 copies that contribute zero gradient. - """ - trajs = [_train_trajectory(rec) for rec in samples] - advs = [float(rec['advantage']) for rec in samples] - rem = (-len(trajs)) % args.sft_batch_size - if rem: - trajs += [trajs[-1]] * rem # zero-advantage pads -> forward only, no gradient - advs += [0.0] * rem - micro_batches = 0 - for i in range(0, len(trajs), args.sft_batch_size): - skill_model.forward_backward(inputs=trajs[i:i + args.sft_batch_size], - advantages=advs[i:i + args.sft_batch_size]) - micro_batches += 1 - skill_model.clip_grad_and_step() - if ckpt is not None: - ckpt.sync_weights(merge_and_sync=True) - metric = skill_model.calculate_metric(is_training=True) - return {'n_samples': len(samples), 'n_steps': 1, 'n_micro_batches': micro_batches, - 'advantages': [float(rec['advantage']) for rec in samples], - 'metric': {k: (float(v) if _is_num(v) else v) for k, v in (metric or {}).items()}} - - -def _is_num(v: Any) -> bool: - try: - float(v) - return True - except (TypeError, ValueError): - return False - - -# --------------------------------------------------------------------------- -# Main -# --------------------------------------------------------------------------- -def _build_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--dataset', choices=('aops', 'math'), default='aops', - help='Problem source. aops (AI-MO competition problems) is much ' - 'harder than MATH, so the base fails more often -> more offensive ' - 'training signal after balanced sampling.') - p.add_argument('--n', type=int, default=2000, - help='Problems to load into the draw pool (cycled/reshuffled across ' - 'epochs; with --balance many more baseline rollouts than this ' - 'may run, but the pool size is fixed here).') - p.add_argument('--seed', type=int, default=42) - p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True, - help='Keep only answers that collapse to one integer/decimal/fraction, ' - 'matching SEAM numeric reward and avoiding non-scalar grading noise.') - p.add_argument('--eval-size', type=int, default=128, - help='Fixed holdout problems, sampled before the train pool after all ' - 'filters; set 0 to disable fixed eval.') - p.add_argument('--eval-every', type=int, default=10, - help='Run fixed holdout eval every N generation chunks when --eval-size > 0.') - p.add_argument('--chunk-size', type=int, default=16, - help='Problems per generation chunk (all sampler calls batched).') - # -- online baseline-balanced sampling (draw+baseline until the chunk hits the - # target base fail:success mix, so the offensive signal is not starved) -- - p.add_argument('--balance', action=argparse.BooleanOptionalAction, default=True, - help='Keep drawing+baselining problems until the chunk matches the ' - 'target base fail:success composition, then select a balanced ' - 'subset. --no-balance draws chunk_size problems directly.') - p.add_argument('--balance-success-frac', type=float, default=0.4, - help='Target fraction of the chunk that the base solves (base-success). ' - '0.4 => 3:2 fail:success; 0.2 => 4:1. The remainder are base-fail.') - p.add_argument('--balance-loop-frac', type=float, default=0.5, - help='Within the base-fail portion, SOFT target fraction of loop-fails ' - '(ran out of length / never terminated) vs non-loop wrong answers. ' - 'Best-effort only: the fail count is filled from whichever bucket ' - 'is available so a thin bucket never starves the chunk.') - p.add_argument('--balance-max-draws-mult', type=int, default=8, - help='Draw budget per chunk as a multiple of chunk_size; once this many ' - 'problems have been baselined the chunk is assembled from whatever ' - 'the buckets hold (ratio may drift; the actual mix is logged).') - p.add_argument('--n-skills', type=int, default=8, - help='Candidate skills generated per hard problem.') - p.add_argument('--view-b-frac', type=float, default=0.5, - help='Fraction of hard problems routed to view B (query-only, ' - 'deployment form); the rest go to view A (problem + attempt). ' - 'Each problem is assigned to EXACTLY ONE view.') - p.add_argument('--skill-retries', type=int, default=2, - help='Extra skill-gen rounds for a hard problem that yielded no ' - 'clean, parseable candidate (thinking-ON purity gate rejects).') - p.add_argument('--skill-gen-temperature', type=float, default=1.0, - help='Sampling temperature for skill-gen (BOTH views). >0 so the ' - 'n_skills candidates per problem are genuinely DIVERSE — a group ' - 'of near-duplicate skills gives GRPO no real good-vs-bad contrast.') - p.add_argument('--skill-gen-top-p', type=float, default=1.0, - help='top_p for skill-gen; 1.0 keeps the full tail for diversity.') - p.add_argument('--skill-gen-top-k', type=int, default=-1, - help='top_k for skill-gen; -1 disables truncation (max diversity). ' - 'A finite value only narrows the candidate pool.') - p.add_argument('--max-model-len', type=int, default=16384) - p.add_argument('--max-tokens', type=int, default=8192, - help='Max generated tokens for solve rollouts.') - p.add_argument('--skill-max-tokens', type=int, default=8192, - help='Max tokens for skill-gen (thinking ON: the model must close ' - ' within this budget or the candidate is dropped, so ' - 'leave ample room).') - p.add_argument('--leak-workers', type=int, default=16, - help='Parallel workers for the LeakVerifier backup judge (capped at 16 ' - 'to avoid the teacher API burst-rate limit; leak and rubric run in ' - 'separate phases so peak teacher concurrency is max(leak,rubric)).') - p.add_argument('--rubric-workers', type=int, default=16, - help='Parallel workers for the view-A rubric diagnose() calls ' - '(teacher-served; requires LLM_BACKUP_* env).') - # -- online GRPO (one on-policy update per generated chunk) -- - p.add_argument('--sft-batch-size', type=int, default=8, - help='Driver-side micro-batch size before the chunk-level optimizer step; ' - 'MUST be a multiple of the training dp size (sliced across dp ranks).') - p.add_argument('--grpo-epsilon', type=float, default=0.2, - help='PPO clip epsilon for GRPOLoss (ratio~1 on-policy, so rarely binds).') - p.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True, - help='Fold output validity into the reward (SEAM-style): unparseable/impure ' - 'or answer-leaked candidates score 0 and join their group to be trained ' - 'DOWN (the whole response, think tokens included). ' - '--no-format-in-reward keeps the reject-and-exclude gate.') - p.add_argument('--lr', type=float, default=1e-5) - p.add_argument('--max-train-rounds', type=int, default=200, - help='Cap on train rounds = trained chunks (also sizes the LR schedule).') - p.add_argument('--save-rounds', type=int, default=50) - p.add_argument('--trend-every', type=int, default=10, - help='Every N chunks, print a [trend] line contrasting the first N ' - 'vs the most recent N chunks (adoption + lift + pos/chunk) ' - 'so the training effect on fresh problems is visible at a glance.') - p.add_argument('--output-dir', default='./output/reflexion_skill_rft') - p.add_argument('--swanlab-project', default='twinkle', - help='swanlab project; logging is skipped when swanlab is not ' - 'installed or SWANLAB_MODE=disabled.') - p.add_argument('--swanlab-exp', default='', - help='swanlab experiment (run) name; empty = auto.') - return p.parse_args() - - -def _trend_line(hist: List[Dict[str, float]], window: int, rounds_done: int) -> Optional[str]: - """Contrast the FIRST ``window`` chunks with the most recent ``window`` chunks so - the online training effect on fresh, never-trained problems is glanceable: if RFT - is working, adoption and lift on recent chunks exceed the early baseline.""" - if len(hist) < 2 * window: - return None # need two non-overlapping windows for a clean before/after - base, rec = hist[:window], hist[-window:] - m = lambda xs, k: sum(h[k] for h in xs) / len(xs) - return (f'[trend] first {window} vs last {window} chunks | ' - f'adopt A {m(base,"aA"):.2f}->{m(rec,"aA"):.2f} ' - f'B {m(base,"aB"):.2f}->{m(rec,"aB"):.2f} | ' - f'lift {m(base,"lift"):+.3f}->{m(rec,"lift"):+.3f} | ' - f'pos/chunk {m(base,"pos"):.1f}->{m(rec,"pos"):.1f} | rounds={rounds_done}') - - -def _query_rows(full: List[Dict[str, Any]]) -> List[Tuple[float, float, float, int, str]]: - """Per hard problem that produced >=1 scored candidate: its no-skill baseline - pass@k, the BEST and MEAN with-skill pass@k over its N skill candidates, the scored - count, and the problem text. Drives both the per-query print and the swanlab passk/* - aggregates.""" - rows = [] - for rec in full: - if rec.get('record_type') != 'problem' or not rec.get('is_hard'): - continue - ps = [c['with_pass'] for c in rec.get('candidates', []) if c.get('with_pass') is not None] - if not ps: - continue - rows.append((rec['baseline_pass'], max(ps), sum(ps) / len(ps), len(ps), rec['problem'])) - return rows - - -def _clean_metric(metric: Optional[Dict[str, Any]]) -> Dict[str, float]: - """Numeric GRPO metrics for swanlab: collapse the duplicate per-group LR to a single - ``lr`` and drop non-numeric fields (e.g. 'total time elapse').""" - out: Dict[str, float] = {} - for k, v in (metric or {}).items(): - if not _is_num(v): - continue - if k.startswith('learning rate'): - if 'group 1' in k: - out['lr'] = float(v) - else: - out[k.replace(' ', '_')] = float(v) - return out - - -def _swan_metrics(summary: Dict[str, Any], log: Optional[Dict[str, Any]], - rows: List[Tuple[float, float, float, int, str]]) -> Dict[str, float]: - """Flat metric dict for swanlab = external reflexion metrics + (when this chunk was - trained) the GRPO built-in metric. acc/adopt/term are only emitted on chunks that had - hard problems, and passk/* only when scored candidates exist, so idle chunks don't dip - the charts to zero.""" - d: Dict[str, float] = { - 'gen/n_hard': summary['n_hard'], 'gen/n_clean': summary['n_clean'], - 'gen/n_leaked': summary['n_leaked'], 'gen/n_train_samples': summary['n_train_samples'], - 'gen/n_reward_pos': summary['n_reward_pos'], - 'gen/n_train_from_fail': summary['n_train_from_fail'], - 'gen/abs_adv_from_fail_frac': summary['abs_adv_from_fail_frac'], - } - bal = summary.get('balance') or {} - if bal.get('enabled'): - d.update({'balance/n_drawn': bal['n_drawn'], - 'balance/n_baseline_fresh': bal['n_baseline_fresh'], - 'balance/selected_success_frac': bal['selected_success_frac'], - 'balance/selected_fail_loop': bal['selected_fail_loop'], - 'balance/selected_fail_wrong': bal['selected_fail_wrong']}) - if summary['n_hard'] > 0: - d.update({ - 'acc/baseline_pass': summary['avg_baseline_pass_on_hard'], - 'acc/withskill_pass': summary['avg_withskill_pass'], - 'acc/lift': summary['avg_lift'], - 'adopt/A': summary['view_A']['adoption_rate'], - 'adopt/B': summary['view_B']['adoption_rate'], - 'term/withskill': summary['termination_rate_withskill'], - }) - if rows: - m = lambda i: sum(r[i] for r in rows) / len(rows) - d.update({'passk/baseline_mean': m(0), 'passk/bestN_mean': m(1), 'passk/avgN_mean': m(2)}) - if log: - d['train/n_steps'] = log['n_steps'] - if 'n_micro_batches' in log: - d['train/n_micro_batches'] = log['n_micro_batches'] - d.update({f'train/{k}': v for k, v in _clean_metric(log.get('metric')).items()}) - return d - - -def _prefix_metrics(metrics: Dict[str, float], prefix: str) -> Dict[str, float]: - return {f'{prefix}/{k}': v for k, v in metrics.items()} - - -def _greedy_eval_metrics(recs: List[Dict[str, Any]], ci: int, rounds: int - ) -> Tuple[Dict[str, Any], Dict[str, float]]: - """Aggregate the greedy holdout into SEAM ``mean@1`` metrics: overall + per-view acc, - the frozen-baseline acc, and their lift -- all single-sample-per-problem means (no - candidate averaging, no pass@k), so acc is directly comparable to SEAM's - ``val-core/math/acc/mean@1`` (correctness only; format/leak not gated).""" - def acc(rs: List[Dict[str, Any]]) -> float: - return sum(1 for x in rs if x['withskill_correct']) / len(rs) if rs else 0.0 - def bacc(rs: List[Dict[str, Any]]) -> float: - return sum(x['baseline_pass'] for x in rs) / len(rs) if rs else 0.0 - A = [x for x in recs if x['view'] == 'A'] - B = [x for x in recs if x['view'] == 'B'] - ws, base = acc(recs), bacc(recs) - summary = { - 'record_type': 'eval_summary', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, - 'n': len(recs), 'n_A': len(A), 'n_B': len(B), - 'acc_mean1': ws, 'baseline_acc_mean1': base, 'lift_mean1': ws - base, - 'acc_A_mean1': acc(A), 'acc_B_mean1': acc(B), - 'format_mean1': (sum(1 for x in recs if x['skill_parseable']) / len(recs)) if recs else 0.0, - 'term_mean1': (sum(1 for x in recs if x['withskill_terminated']) / len(recs)) if recs else 0.0, - } - metrics = { - 'core/math/acc/mean@1': ws, - 'core/math/baseline_acc/mean@1': base, - 'core/math/lift/mean@1': ws - base, - 'core/math/format/mean@1': summary['format_mean1'], - 'core/math/term/mean@1': summary['term_mean1'], - } - if A: - metrics['core/math/acc_A/mean@1'] = summary['acc_A_mean1'] - if B: - metrics['core/math/acc_B/mean@1'] = summary['acc_B_mean1'] - return summary, metrics - - -def _run_greedy_eval(base_sampler, skill_sampler, - eval_records: List[Dict[str, Any]], eval_cache: Dict[str, Dict[str, Any]], - ci: int, rounds: int, base_dp: int, skill_dp: int, - args: argparse.Namespace, checker=None, - diag_cache: Optional[Dict[str, str]] = None - ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], Dict[str, float]]: - """SEAM ``val-core/math/acc/mean@1`` analogue on the fixed holdout: ONE greedy skill per - problem (T=0) injected into ONE greedy base solve (T=0), so acc is a single-sample - pass@1 per problem averaged over problems. Each problem keeps its assigned view; view A - still gets the rubric process-check, view B stays query-only -- the mixed A/B acc is the - deployment number. No leak filter: like SEAM's val, acc scores correctness alone.""" - _baseline_rollout(base_sampler, eval_records, base_dp, args, eval_cache) # frozen greedy baseline - for r in eval_records: - r['_view'] = _assign_view(r['problem'], args) - r['_rubric_diag'] = '' - _diagnose_views(checker, eval_records, args, diag_cache) - sg_out = _run_samples(skill_sampler, [_view_prompt(r, args) for r in eval_records], - 1, args.skill_max_tokens, skill_dp, temperature=0.0) - skills = [] - for seqs in sg_out: - resp = _clean_text(getattr(seqs[0], 'decoded', '') or '') if seqs else '' - skills.append((_extract_skills_block(resp) or '', resp)) - ws_out = _run_samples(base_sampler, - [build_skill_solve_prompt(r['problem'], sk) for r, (sk, _) in zip(eval_records, skills)], - 1, args.max_tokens, base_dp, temperature=0.0) - recs = [] - for r, (sk, sresp), seqs in zip(eval_records, skills, ws_out): - roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - recs.append({ - 'record_type': 'eval_problem', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'view': r['_view'], 'rubric_diag': r.get('_rubric_diag', ''), - 'baseline_pass': r['_baseline_pass'], - 'skill': sk, 'skill_parseable': bool(sk), 'skill_response': sresp, - 'withskill_pred': roll['pred'], 'withskill_correct': roll['correct'], - 'withskill_terminated': roll['terminated'], 'withskill_stop_reason': roll['stop_reason'], - 'withskill_text': roll['text'], - }) - summary, metrics = _greedy_eval_metrics(recs, ci, rounds) - return recs, summary, metrics - - -def _validate_run_config(args: argparse.Namespace, records: List[Dict[str, Any]]) -> None: - """Fail fast on configs that would SILENTLY hang the online sampler: _ProblemPool.draw(k) - dedups within a call, so it never returns unless the pool holds >= chunk_size problems; - a zero draw budget or chunk size yields empty chunks that never advance ``rounds``.""" - if not records: - raise ValueError(f'loaded 0 {args.dataset} problems; check the dataset source') - if args.chunk_size < 1: - raise ValueError(f'--chunk-size must be >= 1 (got {args.chunk_size})') - if args.eval_size < 0: - raise ValueError(f'--eval-size must be >= 0 (got {args.eval_size})') - if args.eval_size > 0 and args.eval_every < 1: - raise ValueError(f'--eval-every must be >= 1 when eval is enabled (got {args.eval_every})') - if len(records) < args.chunk_size: - raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded problems ' - f'({len(records)}); raise --n or lower --chunk-size') - if args.balance_max_draws_mult < 1: - raise ValueError(f'--balance-max-draws-mult must be >= 1 (got {args.balance_max_draws_mult})') - if not 0.0 <= args.balance_success_frac <= 1.0: - raise ValueError(f'--balance-success-frac must be in [0, 1] (got {args.balance_success_frac})') - if not 0.0 <= args.balance_loop_frac <= 1.0: - raise ValueError(f'--balance-loop-frac must be in [0, 1] (got {args.balance_loop_frac})') - - -def main() -> None: - args = _build_args() - if args.sft_batch_size % TRAIN_DP != 0: - raise ValueError(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple ' - f'of the training dp size ({TRAIN_DP})') - # LR schedule now follows chunk-level optimizer updates, not driver micro-batches. - steps_per_round = 1 - records, eval_records, data_stats = _load_records(args) - _validate_run_config(args, records) - os.makedirs(args.output_dir, exist_ok=True) - data_path = os.path.join(args.output_dir, 'skill_dataset.jsonl') - gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') - eval_path = os.path.join(args.output_dir, 'eval_records.jsonl') - train_log_path = os.path.join(args.output_dir, 'train_log.jsonl') - - if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY')): - sys.stderr.write('[rft] WARNING: no LLM_BACKUP_API_KEY/OPENAI_API_KEY — ' - 'LeakVerifier will report no_llm and skip leak filtering\n') - - use_swan = swanlab is not None and os.environ.get('SWANLAB_MODE') != 'disabled' - if use_swan: - swanlab.init(project=args.swanlab_project, - experiment_name=(args.swanlab_exp or None), - config={'model': GEN_MODEL_ID, 'dataset': args.dataset, - 'n': len(records), 'eval_n': len(eval_records), - 'raw_loaded': data_stats['raw_loaded'], - 'numeric_only': args.numeric_only, - 'numeric_dropped': data_stats['numeric_dropped'], - 'n_skills': args.n_skills, 'view_b_frac': args.view_b_frac, - 'balance': args.balance, - 'balance_success_frac': args.balance_success_frac, - 'skill_gen_temp': args.skill_gen_temperature, - 'grpo_epsilon': args.grpo_epsilon, 'lr': args.lr}) - - # -- Device groups: train (FSDP2) + two independent vLLM samplers. -- - r0, r1, r2 = TRAIN_GPUS, TRAIN_GPUS + SKILL_SAMPLER_GPUS, NUM_GPUS - device_groups = [ - DeviceGroup(name='train', ranks=list(range(0, r0)), device_type='GPU'), - DeviceGroup(name='skill_sampler', ranks=list(range(r0, r1)), device_type='GPU'), - DeviceGroup(name='base_sampler', ranks=list(range(r1, r2)), device_type='GPU'), - ] - twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, groups=device_groups, - lazy_collect=False) - - # -- Skill model: full-param FSDP2, GRPO policy update. -- - train_mesh = DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_DP, fsdp_size=TRAIN_FSDP) - skill_model = TransformersModel(model_id=GEN_MODEL_ID, device_mesh=train_mesh, - remote_group='train', - ddp_config={'find_unused_parameters': False}) - from twinkle.patch.no_split_modules import NoSplitModulesPatch - skill_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) - skill_model.set_template(Template, model_id=GEN_MODEL_ID, - enable_thinking=True, max_length=args.max_model_len, - truncation_strategy='delete') - skill_model.set_processor(InputProcessor, padding_free=False) - skill_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) - skill_model.set_optimizer('AdamW', lr=args.lr) - skill_model.set_lr_scheduler('CosineWarmupScheduler', num_warmup_steps=10, - num_training_steps=args.max_train_rounds * steps_per_round) - - # -- Two vLLM samplers: skill (synced) + base (frozen). -- - skill_dp, base_dp = SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS - skill_sampler = vLLMSampler( - model_id=GEN_MODEL_ID, - engine_args={'gpu_memory_utilization': GEN_GPU_MEM, - 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, - device_mesh=DeviceMesh.from_sizes(world_size=SKILL_SAMPLER_GPUS, dp_size=skill_dp), - remote_group='skill_sampler') - skill_sampler.set_template(Template, model_id=GEN_MODEL_ID, - enable_thinking=True, max_length=args.max_model_len) - base_sampler = vLLMSampler( - model_id=GEN_MODEL_ID, - engine_args={'gpu_memory_utilization': GEN_GPU_MEM, - 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, - device_mesh=DeviceMesh.from_sizes(world_size=BASE_SAMPLER_GPUS, dp_size=base_dp), - remote_group='base_sampler') - base_sampler.set_template(Template, model_id=GEN_MODEL_ID, - enable_thinking=True, max_length=args.max_model_len) - - ckpt = CheckpointEngineManager(model=skill_model, sampler=skill_sampler) - leak = LeakVerifier(sampler=None, answer_only=True) # flag ONLY the final answer (view A only; view B skips leak) - # leak = LeakVerifier(sampler=None, judge_system=_LEAK_JUDGE_SYSTEM) # stricter: also flag concrete intermediate key results - checker = _build_rubric_checker() # view-A process-check (teacher-only); None if no LLM backup - if checker is None: - sys.stderr.write('[rft] no LLM backup env -> view-A rubric process-check DISABLED ' - '(skill-gen diagnoses from the attempt alone)\n') - - sys.stderr.write(f'[rft] raw={data_stats["raw_loaded"]} numeric_drop={data_stats["numeric_dropped"]} ' - f'train={len(records)} eval={len(eval_records)} {args.dataset} problems; ' - f'train_gpus={TRAIN_GPUS} skill_dp={skill_dp} base_dp={base_dp}\n') - - # -- Sequential: generate one chunk, train on it, sync -> exact on-policy GRPO. - # Generation dominates wall-clock, so not overlapping training costs little, and - # it removes all producer/consumer concurrency (no thread, no lock). -- - cfg = {'record_type': 'config', 'model': GEN_MODEL_ID, 'dataset': args.dataset, - 'n': len(records), 'eval_n': len(eval_records), 'seed': args.seed, - 'numeric_only': args.numeric_only, - 'raw_loaded': data_stats['raw_loaded'], - 'numeric_dropped': data_stats['numeric_dropped'], - 'eval_every': args.eval_every, - 'n_skills': args.n_skills, - 'view_b_frac': args.view_b_frac, 'skill_retries': args.skill_retries, - 'balance': args.balance, 'balance_success_frac': args.balance_success_frac, - 'balance_loop_frac': args.balance_loop_frac, - 'balance_max_draws_mult': args.balance_max_draws_mult, - 'skill_gen_temp': args.skill_gen_temperature, - 'skill_gen_top_p': args.skill_gen_top_p, 'skill_gen_top_k': args.skill_gen_top_k, - 'reward': 'greedy_binary(correct)', 'advantage': 'group_relative', - 'format_in_reward': args.format_in_reward, - 'rubric_check': 'fixed_math_5crit(viewA)' if checker else 'disabled', - 'grpo_epsilon': args.grpo_epsilon, 'lr': args.lr, - 'max_train_rounds': args.max_train_rounds, 'started': int(time.time())} - hist: List[Dict[str, float]] = [] - rounds = 0 - pool = _ProblemPool(records, args.seed) - eval_cache: Dict[str, Dict[str, Any]] = {} - rubric_cache: Dict[str, str] = {} - eval_rubric_cache: Dict[str, str] = {} - with open(gen_path, 'w', encoding='utf-8') as gen_f, \ - open(eval_path, 'w', encoding='utf-8') as eval_f, \ - open(data_path, 'w', encoding='utf-8') as data_f, \ - open(train_log_path, 'w', encoding='utf-8') as tlog: - for f in (gen_f, eval_f, data_f, tlog): - f.write(json.dumps(cfg, ensure_ascii=False) + '\n') - f.flush() - gstep = 0 - # Each chunk is drawn fresh from the pool (which reshuffles + bumps epoch on every - # full pass) and RE-GENERATED with the current (improved) policy, so every chunk - # stays on-policy (no importance correction) -- the online analogue of SEAM's - # fixed-data epochs. With --balance, _draw_chunk keeps drawing+baselining until the - # base fail:success mix hits the target before this chunk is trained on. - while rounds < args.max_train_rounds: - chunk, balance = _draw_chunk(pool, base_sampler, base_dp, args) - full, summary, groups = process_chunk( - base_sampler, skill_sampler, leak, chunk, gstep, base_dp, skill_dp, - args, checker, rubric_cache) - summary['balance'] = balance - - log = None - if groups: # on-policy GRPO update on this chunk, then weights sync - log = _train_chunk(skill_model, ckpt, groups, args) - rounds += 1 - log.update({'record_type': 'train_round', 'round': rounds, - 'chunk': gstep, 'epoch': pool.epoch, 'ts': int(time.time())}) - tlog.write(json.dumps(log, ensure_ascii=False) + '\n') - tlog.flush() - if rounds % args.save_rounds == 0: - skill_model.save(f'skill-rft-{rounds}', output_dir=args.output_dir) - - summary['rounds_done'], summary['epoch'] = rounds, pool.epoch - for rec in full: - gen_f.write(json.dumps(rec, ensure_ascii=False) + '\n') - gen_f.write(json.dumps(summary, ensure_ascii=False) + '\n') - gen_f.flush() - for v in groups: - data_f.write(json.dumps(v, ensure_ascii=False) + '\n') - data_f.flush() - - sa, sb = summary['view_A'], summary['view_B'] - hist.append({'aA': sa['adoption_rate'], 'aB': sb['adoption_rate'], - 'lift': summary['avg_lift'], 'pos': summary['n_reward_pos']}) - bal_str = (f'bal {balance["selected_fail"]}f/{balance["selected_success"]}s ' - f'(loop {balance["selected_fail_loop"]} drew {balance["n_drawn"]}/' - f'fresh {balance["n_baseline_fresh"]}' - + ('!' if balance.get('budget_hit') else '') + ') ' - ) if balance.get('enabled') else '' - sys.stderr.write( - f'[gen] e{pool.epoch} g{gstep}: {bal_str}hard={summary["n_hard"]} ' - f'clean={summary["n_clean"]} train={summary["n_train_samples"]} ' - f'(fail {summary["n_train_from_fail"]} adv%{summary["abs_adv_from_fail_frac"]:.2f}) ' - f'acc={summary["avg_baseline_pass_on_hard"]:.2f}->{summary["avg_withskill_pass"]:.2f} ' - f'lift={summary["avg_lift"]:+.3f} ' - f'A[{sa["n_hard"]}h {sa["adoption_rate"]:.2f}] ' - f'B[{sb["n_hard"]}h {sb["adoption_rate"]:.2f}] ' - f'rounds={rounds}' - + (f' metric={log.get("metric")}' if log else '') + '\n') - # -- per-query passk (base vs best/avg of N skills) + swanlab metrics -- - rows = _query_rows(full) - for base_p, best_p, avg_p, nsc, prob in rows: - logger.info(f'[q] g{gstep} base={base_p:.2f} bestN={best_p:.2f} avgN={avg_p:.2f} ' - f'n={nsc} | {prob[:70].replace(chr(10), " ")}') - if use_swan: - swanlab.log(_swan_metrics(summary, log, rows), step=gstep) - - if eval_records and (gstep + 1) % args.eval_every == 0: - eval_recs, eval_summary, eval_metrics = _run_greedy_eval( - base_sampler, skill_sampler, eval_records, eval_cache, gstep, - rounds, base_dp, skill_dp, args, checker, eval_rubric_cache) - for rec in eval_recs: - eval_f.write(json.dumps(rec, ensure_ascii=False) + '\n') - eval_f.write(json.dumps(eval_summary, ensure_ascii=False) + '\n') - eval_f.flush() - if use_swan: - swanlab.log(_prefix_metrics(eval_metrics, 'eval'), step=gstep) - sys.stderr.write( - f'[eval] g{gstep}: n={eval_summary["n"]} mean@1 ' - f'acc={eval_summary["baseline_acc_mean1"]:.3f}->{eval_summary["acc_mean1"]:.3f} ' - f'lift={eval_summary["lift_mean1"]:+.3f} ' - f'A[{eval_summary["n_A"]} {eval_summary["acc_A_mean1"]:.3f}] ' - f'B[{eval_summary["n_B"]} {eval_summary["acc_B_mean1"]:.3f}] ' - f'fmt={eval_summary["format_mean1"]:.2f} rounds={rounds}\n') - - if (gstep + 1) % args.trend_every == 0: - tl = _trend_line(hist, args.trend_every, rounds) - if tl: - sys.stderr.write(tl + '\n') - gstep += 1 - - skill_model.save('skill-rft-final', output_dir=args.output_dir) - sys.stderr.write(f'[rft] done: {rounds} train rounds over {gstep} chunks / {pool.epoch} epochs; ' - f'data -> {data_path}\n') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/embedding/train_skill_v2.py b/cookbook/exp/embedding/train_skill_v2.py deleted file mode 100644 index 6e4d7d201..000000000 --- a/cookbook/exp/embedding/train_skill_v2.py +++ /dev/null @@ -1,1450 +0,0 @@ -"""Simplified GRPO + buffer-distill training for the reflexion skill generator (v2). - -Key differences from train_reflexion_skill.py: -- No view A/B split: all skill-gen is query-only (deployment form). -- No baseline rollout in training, no balance selection. -- thinking OFF; skill model outputs optional analysis then block. -- Reward = parseable × (correct AND terminated) × min(1, len_budget/skill_len). -- Buffer A: adv=0 (all-fail) problems accumulate failure trajectories. -- Buffer B: batch rubric → regenerate skill → pass@k validate → SFT injection. -- SFT is event-driven: buffer B reaches threshold → one SFT pass → eval. - -Launch: - LLM_BACKUP_API_KEY=... python cookbook/exp/embedding/train_skill_v2.py \ - --dataset aops --n 5000 --chunk-size 16 --lr 6e-6 -""" -import argparse -import copy -import hashlib -import json -import math -import os -import re -import sys -import threading -import time -from concurrent.futures import ThreadPoolExecutor -from typing import Any, Dict, List, Optional, Set, Tuple - -import numpy as np - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.checkpoint_engine import CheckpointEngineManager -from twinkle.data_format import SamplingParams -from twinkle.dataset import Dataset, DatasetMeta -from twinkle.model import TransformersModel -from twinkle.patch.no_split_modules import NoSplitModulesPatch -from twinkle.processor import InputProcessor -from twinkle.sampler import vLLMSampler -from twinkle.template import Template -from twinkle_agentic.verifier import RubricVerifier -from twinkle_agentic.verifier.rubric_verifier import RubricItem - -logger = get_logger() - -try: - import swanlab -except ImportError: - swanlab = None - -MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') -GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.8)) -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') - -TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 2)) -REF_GPUS = int(os.environ.get('REF_GPUS', 2)) -SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) -BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 2)) -NUM_GPUS = TRAIN_GPUS + REF_GPUS + SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS -TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', min(1, TRAIN_GPUS))) -REF_FSDP = int(os.environ.get('REF_FSDP', min(1, REF_GPUS))) -TRAIN_DP = TRAIN_GPUS // TRAIN_FSDP -REF_DP = REF_GPUS // REF_FSDP - - -# =========================================================================== -# Section A — boxed extraction + answer grading (verbatim from v1) -# =========================================================================== -_BOXED_RE = re.compile(r'\\boxed\s*\{') - - -def extract_boxed(text: str) -> Optional[str]: - if not text: - return None - last = None - for m in _BOXED_RE.finditer(text): - depth, i = 1, m.end() - while i < len(text) and depth > 0: - depth += (text[i] == '{') - (text[i] == '}') - i += 1 - if depth == 0: - last = text[m.end():i - 1].strip() - return last - - -def normalize_answer(ans: str) -> str: - if not ans: - return '' - s = str(ans).strip() - m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) - if m: - return m.group(1) - s = s.strip('$').strip().replace('\u2212', '-') - 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').replace(r'\tfrac', r'\frac') - s = s.replace(r'\,', '').replace(r'\;', '').replace(r'\!', '') - s = re.sub(r'\\(?:quad|qquad|\s)', '', s) - s = re.sub(r'\s+', '', s) - s = re.sub(r'\\sqrt([0-9A-Za-z])', r'\\sqrt{\1}', s) - s = re.sub(r'\\frac\{([^{}]+)\}([^{}\\])', r'\\frac{\1}{\2}', s) - s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) - s = re.sub(r'\{[a-zA-Z]+\}$', '', s) - s = re.sub(r'\^\\circ|\^\{\\circ\}|\u00b0|\\circ', 'deg', s) - s = s.replace(r'\minus{}', '-').replace(r'\minus', '-') - - def _frac_to_slash(mt): - text = mt.group(0) - pos = text.index('{') + 1 - depth, num_start = 1, pos - while depth > 0: - depth += (text[pos] == '{') - (text[pos] == '}') - pos += 1 - numer = text[num_start:pos - 1] - pos += 1 - den_start, depth = pos, 1 - while depth > 0: - depth += (text[pos] == '{') - (text[pos] == '}') - pos += 1 - return f'({numer})/({text[den_start:pos - 1]})' - - s = re.sub(r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', _frac_to_slash, s) - s = re.sub(r'(? bool: - try: - va, vb = float(a.replace('(', '').replace(')', '')), 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) - return va is not None and vb is not None and abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) - - -_MCQ_REF_RE = re.compile( - r'^\\(?:textbf|text|mathrm|mathbf)\{\(?([A-E])\)?\s*\}\s*(.+)$' - r'|^\(?([A-E])\)\s+(.+)$') -_VAR_PREFIX_RE = re.compile(r'^(?:[A-Za-z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)$', re.DOTALL) - - -def _split_mcq(ans): - s = ans.strip() - m = _MCQ_REF_RE.match(s) - if m: - return (m.group(1) or m.group(3)), ((m.group(2) or m.group(4) or '').strip() or None) - bl = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) - return (bl.group(1), None) if bl else (None, s or None) - - -def _strip_var_prefix(ans): - m = _VAR_PREFIX_RE.match((ans or '').strip()) - return m.group(1).strip() if m else (ans or '') - - -def _math_verify_equal(predicted: str, reference: str) -> bool: - try: - from math_verify import parse, verify - gold = parse(r'\boxed{' + reference + '}', parsing_timeout=1) - pred = parse(r'\boxed{' + predicted + '}', parsing_timeout=1) - return bool(gold and pred and verify(gold, pred, timeout_seconds=1)) - except Exception: - return False - - -def answers_match(predicted: str, reference: str) -> bool: - if not predicted or not reference: - return False - norm_p, norm_r = normalize_answer(predicted), normalize_answer(reference) - if norm_p == norm_r or norm_p.lower() == norm_r.lower() or _try_numeric_equal(norm_p, norm_r): - return True - stripped_p = normalize_answer(_strip_var_prefix(predicted)) - stripped_r = normalize_answer(_strip_var_prefix(reference)) - if stripped_p and stripped_r: - if (stripped_p == stripped_r or stripped_p.lower() == stripped_r.lower() - or _try_numeric_equal(stripped_p, stripped_r)): - return True - p_letter, p_value = _split_mcq(predicted) - r_letter, r_value = _split_mcq(reference) - if p_letter and r_letter and p_letter == r_letter: - return True - if (p_letter and p_value is None) and (r_letter and r_value is None): - return False - 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 and (p_val == r_val or p_val.lower() == r_val.lower() or _try_numeric_equal(p_val, r_val)): - return True - for left, right in ((norm_p, norm_r), (stripped_p, stripped_r)): - tl, tr = re.sub(r'[\s()\[\]{}\\]', '', left or ''), re.sub(r'[\s()\[\]{}\\]', '', right or '') - if ',' in tl and tl == tr: - return True - if '=' in norm_r and '=' not in norm_p: - if any(part and (part == norm_p or _try_numeric_equal(part, norm_p)) for part in norm_r.split('=')): - return True - if '=' in norm_p and '=' not in norm_r: - if any(part and (part == norm_r or _try_numeric_equal(part, norm_r)) for part in norm_p.split('=')): - return True - return _math_verify_equal(stripped_p or norm_p, stripped_r or norm_r) - - -_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') - - -def _numeric_value(raw) -> Optional[str]: - if raw is None: - return None - s = str(raw).strip().strip('$').strip() - s = s.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac') - s = re.sub(r'\\!|\\,|\\;|\\ |\\left|\\right|\s', '', s) - for pat in (r'\\frac\{(-?\d+)\}\{(-?\d+)\}', r'(-?\d+)/(-?\d+)'): - m = re.fullmatch(pat, s) - if m: - a, b = int(m.group(1)), int(m.group(2)) - return str(int(a / b)) if (b and a / b == int(a / b)) else (str(a / b) if b else None) - return (str(int(float(s))) if float(s) == int(float(s)) else str(float(s))) if _NUM_RE.fullmatch(s) else None - - -def _answer_leaked(skill: str, reference: str) -> bool: - if not skill: - return False - # Suffix guard: reject only a following DIGIT or a following '.' (decimal point), - # NOT a sentence-ending '.'. Old '(?![\d.])' let leaks like "...= 675." slip through - # because the trailing period satisfied the [\d.] class. 中文注释：尾断言只排除"后接数字" - # 或"后接小数点+数字"，不排除句末句号，堵住 "答案." 这类泄漏漏检。 - for cand in {_numeric_value(reference), (str(reference).strip() or None)}: - if cand and re.search(r'(?') - - -def _clean_text(decoded: Optional[str]) -> str: - return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() - - -def _extract_skill(text: str) -> Optional[str]: - """Parse ... block from skill-gen output.""" - low = text.lower() - end_think = low.rfind('') - answer = text[end_think + len(''):] if end_think >= 0 else text - open_tag, close_tag = '', '' - s = answer.lower().rfind(open_tag) - if s < 0: - return None - inner = s + len(open_tag) - e = answer.lower().find(close_tag, inner) - if e < 0: - return None - block = answer[inner:e].strip() - block = re.sub(r'', '', block, flags=re.IGNORECASE).strip() - return block or None - - -def _parse_seq(seq, gold: str) -> Dict[str, Any]: - text = _clean_text(getattr(seq, 'decoded', '') or '') - pred = extract_boxed(text) - correct = bool(pred) and answers_match(pred, gold) - terminated = getattr(seq, 'stop_reason', None) != 'length' - return {'pred': pred, 'correct': correct, 'terminated': terminated, - 'stop_reason': getattr(seq, 'stop_reason', None), - 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} - - -def _empty_roll(): - return {'pred': '', 'correct': False, 'terminated': False, - 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} - - -def _run_samples(sampler, prompts, num_samples, max_tokens, gen_dp, - temperature=None, top_p=None, top_k=None): - if not prompts: - return [] - params = SamplingParams( - max_tokens=max_tokens, - temperature=GEN_TEMPERATURE if temperature is None else temperature, - top_p=GEN_TOP_P if top_p is None else top_p, - num_samples=num_samples, **({} if top_k is None else {'top_k': top_k})) - padded = prompts - if gen_dp > 1 and 0 < len(prompts) < gen_dp: - padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] - responses = sampler.sample(padded, params)[:len(prompts)] - return [list(r.sequences) if (r and r.sequences) else [] for r in responses] - - -# =========================================================================== -# Section C — data loading (simplified: no balance, no xproblem, no views) -# =========================================================================== -def _boxed_batch(rows, dataset): - sols = rows['solution'] - metas = rows.get('metadata', [None] * len(sols)) - refs = [extract_boxed(s or '') for s in sols] - keep = [bool(ref) and (dataset != 'aops' or bool((meta or {}).get('boxed'))) - for ref, meta in zip(refs, metas)] - return {**rows, 'reference_answer': refs, '_keep': keep} - - -def load_problems(dataset: str, n: int, seed: int) -> List[Dict[str, Any]]: - ds_id = AOPS_DATASET_ID if dataset == 'aops' else os.environ.get('MATH_DATASET_ID', 'modelscope/competition_math') - ds = Dataset(DatasetMeta(dataset_id=f'ms://{ds_id}', split='train')) - nproc = min(32, os.cpu_count() or 1) - ds.map(lambda rows: _boxed_batch(rows, dataset), num_proc=nproc) - ds.filter(lambda row: row['_keep'], num_proc=nproc) - out = [{'data_id': f'{dataset}:{i}', 'problem': row['problem'], - 'reference_answer': row['reference_answer']} - for i, row in enumerate(ds.dataset)] - rng = np.random.RandomState(seed) - rng.shuffle(out) - return out[:n] if (n and n < len(out)) else out - - -def _load_records(args): - records = load_problems(args.dataset, 0, args.seed) - raw_n = len(records) - if args.numeric_only: - records = [{**r, 'reference_answer': v} - for r, v in ((r, _numeric_value(r.get('reference_answer'))) for r in records) - if v is not None] - np.random.RandomState(args.seed).shuffle(records) - # exclude - excl_ids, excl_probs = set(), set() - for path in (args.exclude_data_ids or '').split(','): - path = path.strip() - if not path or not os.path.exists(path): - continue - with open(path) as f: - for line in f: - if not line.strip(): continue - row = json.loads(line) - if row.get('record_type') in {'config', 'summary'}: continue - did = str(row.get('data_id', '')).strip() - if did: excl_ids.add(did) - else: - p = str(row.get('problem', '')).strip() - if p: excl_probs.add(p) - if excl_ids or excl_probs: - records = [r for r in records - if str(r.get('data_id', '')) not in excl_ids - and str(r.get('problem', '')).strip() not in excl_probs] - eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 - eval_records = records[:eval_n] - # Dedup by problem TEXT: index slices are disjoint, but duplicate problem statements - # across the boundary would still leak eval into train. Drop any train record whose - # problem appears in eval, then guard with an explicit overlap assertion. - # 中文注释：train/eval 去重——按题面文本剔除，防止数据集内重复题目跨界泄漏；末尾硬断言无交集。 - eval_probs = {r['problem'] for r in eval_records} - train_records = [r for r in records[eval_n:] if r['problem'] not in eval_probs] - if args.n > 0: - train_records = train_records[:args.n] - if {r['problem'] for r in train_records} & eval_probs: - raise ValueError('eval/train overlap detected after dedup') - logger.info(f'[data] raw={raw_n} train={len(train_records)} eval={len(eval_records)}') - return train_records, eval_records - - -# =========================================================================== -# Section D — DiskCache, ProblemPool, LockedSampler -# =========================================================================== -class DiskCache: - def __init__(self, path: str, enabled: bool = True): - self._mem: Dict[str, Any] = {} - self._fh = None - self._lock = threading.Lock() - if not enabled: - return - if os.path.exists(path): - with open(path, encoding='utf-8') as f: - for line in f: - if line.strip(): - row = json.loads(line) - self._mem[row['key']] = row['value'] - self._fh = open(path, 'a', encoding='utf-8') - - @staticmethod - def key_for(*parts): - return hashlib.md5('\x1f'.join(parts).encode('utf-8')).hexdigest() - - def get(self, key): return self._mem.get(key) - def __contains__(self, key): return key in self._mem - - def put(self, key, value): - with self._lock: - self._mem[key] = value - if self._fh: - self._fh.write(json.dumps({'key': key, 'value': value}, ensure_ascii=False) + '\n') - self._fh.flush() - - def close(self): - if self._fh: self._fh.close() - - -class ProblemPool: - def __init__(self, records, seed): - self._records = list(records) - self._seed, self._cursor, self.epoch = seed, 0, 0 - - def draw(self, k): - out, seen = [], set() - while len(out) < k: - if self._cursor >= len(self._records): - self.epoch += 1 - np.random.RandomState(self._seed + self.epoch).shuffle(self._records) - self._cursor = 0 - r = self._records[self._cursor] - self._cursor += 1 - if id(r) not in seen: - seen.add(id(r)) - out.append(r) - return out - - -class _LockedSampler: - def __init__(self, sampler): - self._sampler = sampler - self._lock = threading.Lock() - - def sample(self, *a, **kw): - with self._lock: - return self._sampler.sample(*a, **kw) - - def __getattr__(self, name): - return getattr(self._sampler, name) - - -# =========================================================================== -# Section E — Rubric (teacher diagnosis, batched at distill time) -# =========================================================================== -_RFT_DIAG_SYSTEM = """\ -You are a strategy-level process checker for a math solution attempt. You are given a -math problem, a rubric, and one attempted solution segment. Decide PASS or FAIL for each -criterion, and write the diagnosis so it can become useful reusable guidance for solving -similar problems without seeing this segment. - -Output STRICT JSON (no prose outside it) with this shape: -{"items": [{"index": 1, "verdict": "PASS"|"FAIL", "reason": "...", "fix": ""}], "overall": "OK"|"ISSUES", "summary": "..."} - -Rules: -- Judge every criterion independently. -- The diagnosis must stay answer-free. -- For FAIL items: describe the process problem at strategy level. -- A fix suggests the LOCAL correction direction without solving. -- Never reveal the final answer or a corrected expression. -- If segment was cut off (no \\boxed{}), mark length-budget as FAIL. -- Keep "reason" and "fix" concise: one short sentence each. -- Output only the JSON object.""" - -_RFT_DIAG_USER = """\ -## Task / query -{query} - -## Rubric -{rubric} - -## Segment -{segment} - -Now output the diagnostic JSON object.""" - -_MATH_RUBRIC = [ - ('The attempt chooses a method suitable for the problem structure', False), - ('The attempt identifies the key constraint, invariant, or quantity before computing', False), - ('Algebraic and logical transformations preserve validity at each step', True), - ('The attempt checks required constraints, domains, boundary cases, or validity conditions', False), - ('The attempt avoids redundant casework, looping, or re-deriving known facts', False), - ('The attempt reaches a final boxed answer within the length budget', False), - ('The approach stays focused on the actual question asked', False), -] -_RUBRIC_VERSION = 'rubric_v5_viewb_strategy' - - -class _RftRubricVerifier(RubricVerifier): - def _diagnose_trajectory(self, query, rubric_block, segment_text): - return {'messages': [ - {'role': 'system', 'content': _RFT_DIAG_SYSTEM}, - {'role': 'user', 'content': _RFT_DIAG_USER.format( - query=query, rubric=rubric_block, segment=segment_text)}]} - - -def build_rubric_checker(): - if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') - or os.environ.get('OPENAI_API_KEY')): - return None - return _RftRubricVerifier( - fixed_rubric=[RubricItem(t, is_hard=h) for t, h in _MATH_RUBRIC], gate=True) - - -def _format_diagnosis(detail) -> str: - rub = detail.rubric - lines = [] - for it in detail.items: - text = rub[it.index - 1].text if 0 < it.index <= len(rub) else f'criterion {it.index}' - if it.verdict: - lines.append(f'- [PASS] {text}') - else: - tail = f': {it.reason}' if it.reason else '' - tail += f' (fix: {it.fix})' if it.fix else '' - lines.append(f'- [FAIL] {text}{tail}') - if detail.summary: - lines.append(f'Summary: {detail.summary}') - return '\n'.join(lines) - - -def _diagnose_entry(checker, entry: Dict[str, Any]) -> Optional[str]: - """Run the teacher rubric on ONE buffer-A failure trajectory → formatted diagnosis text - (or None on API error). Pure network/CPU (no GPU), so it can run on a background thread - while GRPO trains. Shared by the background pre-diagnosis pool and distill_buffer's - fallback for any entry the pool did not reach in time. - 中文注释：单条失败轨迹的 rubric 诊断（纯 API，不吃 GPU）。后台预诊断与 distill 补诊断共用。""" - seg_text = entry['fail_segment'] - if entry.get('fail_stop_reason') == 'length': - seg_text += ('\n\n[Process note: this attempt was cut off at the token budget ' - 'and never produced a final \\boxed{} answer.]') - seg = {'messages': [{'role': 'user', 'content': entry['problem']}, - {'role': 'assistant', 'content': seg_text}]} - try: - return _format_diagnosis(checker.diagnose(seg, query=entry['problem'])) - except Exception as exc: - logger.warning(f'[rubric] diagnose error: {exc}') - return None - - -# =========================================================================== -# Section F — NEW: prompts, reward, buffer logic -# =========================================================================== - -# ---- Skill-gen system prompt (query-only, thinking OFF) ---- -# 中文注释：skillmodel 系统提示词。thinking 关闭；允许在 之前输出简短分析； -# 要求 skill ≤600 字符、切题、不给答案、不啰嗦；方向弱列举（pitfall/技术点/step/overview/确信输出）。 -SKILL_GEN_SYSTEM = """\ -You are a math guidance writer. Write short, reusable guidance for the problem below. - -Rules: -- You may briefly analyze the problem BEFORE the tag. -- Output your guidance inside .... -- Keep the guidance within 600 characters, concise and problem-specific. -- Focus on: pitfalls may be happened to avoid, key techniques, brief step outlines, or how to help to converge to the answer quickly. -- Do NOT calculate the final answer. -- Do NOT be verbose or generic. - -Output format: -[optional brief analysis] - -Your reusable solving guidance here. -""" - -# ---- Executor system prompt (with skill injection) ---- -# 中文注释：executor 系统提示词，将 skill 注入 solver 的 system prompt 前缀。 -DIRECT_SYSTEM = ( - 'You are an expert competition mathematician. Solve the following problem ' - 'step by step. Provide your final answer inside \\boxed{}.') -_SKILL_SOLVE_PREFIX = ( - DIRECT_SYSTEM + '\n\n' - 'Before you start, keep these reminders in mind to avoid common mistakes on this ' - 'type of problem:\n') -_SKILL_SOLVE_SUFFIX = '\nApply them where relevant, but rely on your own reasoning to reach the answer.' - - -def build_direct_prompt(problem): - return {'messages': [{'role': 'system', 'content': DIRECT_SYSTEM}, - {'role': 'user', 'content': problem}]} - - -def build_skill_solve_prompt(problem, skill): - skill = (skill or '').strip() - if not skill: - return build_direct_prompt(problem) - return {'messages': [ - {'role': 'system', 'content': _SKILL_SOLVE_PREFIX + skill + _SKILL_SOLVE_SUFFIX}, - {'role': 'user', 'content': problem}]} - - -# ---- Rubric-guided regeneration prompt (buffer B distillation) ---- -# 中文注释：蒸馏重生成提示词。给旧 skill + rubric 诊断，要求产出改进后的 skill。 -# 只输出 块，≤600字符，不含答案。 -REGEN_SYSTEM = """\ -You are a math guidance writer. You previously wrote guidance for a problem, but the \ -solver still failed. A process-check diagnosed the failure. Revise your guidance to \ -address the diagnosed issues. - -Rules: -- Output ONLY a ... block (no analysis). -- Keep within 600 characters, concise, problem-specific. -- Address the diagnosed failure points, also keep the good parts of the old skills. -- Do NOT include the final answer.""" - -REGEN_USER = """\ -Problem: -{problem} - -Previous guidance (did not help): -{orig_skill} - -Process-check diagnosis: -{rubric_diag} - -Write improved guidance:""" - - -def _skillgen_prompt(problem: str) -> Dict[str, Any]: - """Skill-gen prompt: query-only, no view split.""" - return {'messages': [ - {'role': 'system', 'content': SKILL_GEN_SYSTEM}, - {'role': 'user', 'content': f'Problem:\n{problem}'}]} - - -def _regen_prompt(problem: str, orig_skill: str, rubric_diag: str) -> Dict[str, Any]: - """Regeneration prompt for buffer B distillation.""" - return {'messages': [ - {'role': 'system', 'content': REGEN_SYSTEM}, - {'role': 'user', 'content': REGEN_USER.format( - problem=problem, orig_skill=orig_skill, rubric_diag=rubric_diag)}]} - - -# ---- Reward ---- -# 中文注释：reward = parseable × (correct AND terminated) × min(1, budget/len) -# parseable=0 的候选 reward=0 仍参与 group（格式压力）；截断=失败；超长乘法衰减。 -def _skill_reward(parseable: bool, correct: bool, terminated: bool, - skill_len: int, len_budget: int) -> float: - if not parseable: - return 0.0 - base = 1.0 if (correct and terminated) else 0.0 - len_factor = min(1.0, len_budget / max(skill_len, 1)) - return base * len_factor - - -# ---- Buffer A: collect adv=0 all-fail problems ---- -def _collect_buffer_a(chunk, args) -> List[Dict[str, Any]]: - """Collect problems where all candidates got reward 0 (adv=0, GRPO blind spot). - Store one representative failure trajectory for later rubric diagnosis.""" - entries = [] - for r in chunk: - cs = [c for c in r['_cands'] if c.get('reward') is not None] - if len(cs) < 2: - continue - rewards = [c['reward'] for c in cs] - if max(rewards) > 0: - continue # has signal, not all-fail - # Representative trajectory for rubric + regen seed: prefer a terminated-wrong - # parseable candidate (complete reasoning to diagnose). Its skill becomes the regen - # seed, so pick the MOST SUBSTANTIAL one WITHIN budget (longest ≤ len_budget) — a - # rich-but-not-bloated starting point — rather than an arbitrary [0] or a near-empty - # skill. If all seeds exceed budget, take the one closest to budget (shortest-over). - # 中文注释：代表轨迹既做 rubric 诊断又做 regen 种子——优先"跑完但答错"的候选（完整推理）， - # 其 skill 取预算内最长(最有实质)的作种子；若全超预算则取最接近预算的，避免随机/近空种子。 - budget = args.len_budget - - def _seed_key(c): - L = len(c.get('skills') or '') - return (L <= budget, L if L <= budget else -L) - - parseable = [c for c in cs if c.get('skills')] - term_wrong = [c for c in parseable if c['rolls'] and c['rolls'][0].get('terminated')] - pool_c = term_wrong or parseable or cs - rep = max(pool_c, key=_seed_key) - stop_dist = {} - for c in cs: - sr = c['rolls'][0]['stop_reason'] if c['rolls'] else 'none' - stop_dist[sr] = stop_dist.get(sr, 0) + 1 - entries.append({ - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'data_id': r.get('data_id', ''), - 'orig_skill': rep.get('skills', ''), - 'orig_len': len(rep.get('skills', '')), - 'fail_segment': rep['rolls'][0]['text'] if rep['rolls'] else '', - 'fail_stop_reason': rep['rolls'][0]['stop_reason'] if rep['rolls'] else 'none', - 'stop_reason_dist': stop_dist, - }) - return entries - - -# ---- Buffer B distillation ---- -def distill_buffer(entries: List[Dict[str, Any]], skill_sampler, base_sampler, - checker, skill_dp: int, base_dp: int, - args) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: - """Batch rubric → regenerate K distinct skills → greedy-validate → return SFT records. - 中文注释：蒸馏流程（方案 B，多样性在 skill 侧、executor 用贪心）： - 1. 批量 rubric 诊断失败轨迹；2. 仅 [FAIL] 项用高温重生成 K 个不同候选 skill； - 3. 每个候选过 ≤budget+无leak+去重 过滤；4. 每个存活候选用 executor 贪心(T=0)解 1 次； - 5. gate：≥m 个不同候选达成 terminated-correct → select 长度最接近 budget 的一个入 buffer B。 - 返回 (sft_records, distill_records)：后者逐 entry 记录 rubric_diag/候选 skill/贪心解结果/漏斗 - stage，落盘到 distill_records.jsonl 供复盘（否则 rubric 诊断与候选明细只存在于内存）。""" - if not checker or not entries: - return [], [] - - # Step 1: ensure every entry has a rubric diagnosis. Entries pre-diagnosed in the - # background (see _prediagnose in main) already carry '_rubric_diag'; only the misses - # are diagnosed here (in parallel), so the GPU-idle API wait is normally hidden. - # 中文注释：优先用后台预诊断结果；只对没预诊断到的条目并行补跑，隐藏 API 等待。 - pending = [e for e in entries if not e.get('_rubric_diag')] - if pending: - workers = min(args.rubric_workers, len(pending)) - with ThreadPoolExecutor(max_workers=max(1, workers)) as ex: - diags = list(ex.map(lambda e: _diagnose_entry(checker, e), pending)) - for entry, diag in zip(pending, diags): - entry['_rubric_diag'] = diag or '' - for entry in entries: - entry['rubric_diag'] = entry.get('_rubric_diag') or '' - - # Builder for the structured distill audit records (one per buffer-A entry). Closure over - # `entries`/`args`; takes the per-entry regen skills + greedy solve results (may be empty - # for the early-exit funnel stages). 中文注释：构造逐 entry 的蒸馏审计记录（含漏斗 stage）。 - def _mk_distill(results_by_entry, per_entry_skills, has_fail): - hf_index = {id(e): ei for ei, e in enumerate(has_fail)} - recs = [] - for e in entries: - ei = hf_index.get(id(e)) - cand_results = results_by_entry.get(ei, []) if ei is not None else [] - n_pass = sum(1 for c in cand_results if c['correct'] and c['terminated']) - n_cand = len(per_entry_skills[ei]) if (ei is not None and ei < len(per_entry_skills)) else 0 - if ei is None: - stage = 'no_fail' # rubric 未给出任何 [FAIL] - elif n_cand == 0: - stage = 'no_valid_regen' # 有 [FAIL] 但重生成无一条过 ≤budget/无leak/去重 - elif n_pass >= args.passatk_m: - stage = 'accepted' # ≥m 个候选贪心解对 → 入 buffer B - else: - stage = 'rejected' # 有候选但 args.len_budget: - continue - if _answer_leaked(skill, entry['reference_answer']): - continue - if skill in seen: - continue # 去重：同一 skill 只算一个"不同候选" - seen.add(skill) - skills.append(skill) - per_entry_skills.append(skills) - if skills: - n_entries_with_cands += 1 - - # Flatten to (entry_idx, skill) for ONE batched greedy executor solve per candidate skill. - # 中文注释：每个候选 skill 只用 executor 贪心(T=0)解 1 次——与部署/eval 口径一致； - # k 次 rollout 摊到 k 个不同 skill 上，多样性来自 skill 侧而非 executor 侧。 - flat_idx, flat_prompts = [], [] - for ei, skills in enumerate(per_entry_skills): - for sk in skills: - flat_idx.append((ei, sk)) - flat_prompts.append(build_skill_solve_prompt(has_fail[ei]['problem'], sk)) - if not flat_prompts: - logger.info(f'[distill] {len(has_fail)} [FAIL] entries, 0 valid regen skills') - return [], _mk_distill({}, per_entry_skills, has_fail) - solve_out = _run_samples(base_sampler, flat_prompts, 1, args.max_tokens, base_dp, - temperature=0.0) - - # Gather greedy results back per entry: record EVERY candidate skill's solve outcome - # (not just passers) so distill_records can show why an entry was rejected. - # 中文注释：回收每个候选 skill 的贪心解结果（含未通过的），供审计记录还原被拒原因。 - results_by_entry: Dict[int, List[Dict[str, Any]]] = {} - for (ei, sk), seqs in zip(flat_idx, solve_out): - roll = _parse_seq(seqs[0], has_fail[ei]['reference_answer']) if seqs else _empty_roll() - results_by_entry.setdefault(ei, []).append( - {'skill': sk, 'len': len(sk), 'correct': roll['correct'], 'terminated': roll['terminated']}) - - # Step 4: gate ≥ m distinct greedy-effective skills; select the survivor CLOSEST to the - # length budget (short is the floor, but not so short it degrades to answer-dumping). - # 中文注释：gate——≥m 个不同 skill 在贪心下 terminated-correct；select——在通过的候选里 - # 选长度最接近 budget 的一个入 buffer B（短是地板，但别短到退化成吐答案）。 - sft_records = [] - for ei, cand_results in results_by_entry.items(): - passers = [c['skill'] for c in cand_results if c['correct'] and c['terminated']] - if len(passers) < args.passatk_m: - continue - entry = has_fail[ei] - best = min(passers, key=lambda sk: abs(len(sk) - args.len_budget)) - sft_records.append({ - 'problem': entry['problem'], 'reference_answer': entry['reference_answer'], - 'data_id': entry.get('data_id', ''), - 'response': f'\n{best}\n', - 'skills': best, 'sft': True, - 'n_pass_skills': len(passers), 'n_cand_skills': len(per_entry_skills[ei]), - }) - - logger.info(f'[distill] {len(entries)} A → {len(has_fail)} [FAIL] → ' - f'{n_entries_with_cands} w/cands → {len(sft_records)} validated B ' - f'(gate m={args.passatk_m}/k={k})') - return sft_records, _mk_distill(results_by_entry, per_entry_skills, has_fail) - - - -# =========================================================================== -# Section G — GRPO advantages + training -# =========================================================================== -def _assign_advantages(chunk, args): - """Group-relative advantage: A = (R - mean) / (std + eps). std==0 → adv=0 (skipped).""" - eps = 1e-6 - adv_clip = abs(float(getattr(args, 'adv_clip', 0.0) or 0.0)) - for r in chunk: - for c in r['_cands']: - c['advantage'], c['kept'] = 0.0, False - cs = [c for c in r['_cands'] if c.get('reward') is not None] - if len(cs) < 2: - continue - rewards = [c['reward'] for c in cs] - mean_r = sum(rewards) / len(rewards) - std = (sum((x - mean_r) ** 2 for x in rewards) / len(rewards)) ** 0.5 - if std < 1e-9: - continue - for c in cs: - raw = (c['reward'] - mean_r) / (std + eps) - c['advantage'] = max(-adv_clip, min(adv_clip, raw)) if adv_clip > 0 else raw - c['kept'] = c['reward'] > mean_r - - -def _train_trajectory(rec): - """Rebuild the query-only skill-gen prompt (train/inference match) + response. - GRPO records carry the full generated response; SFT records carry only the - cleaned block. key_rounds selects the final assistant turn.""" - msgs = _skillgen_prompt(rec['problem'])['messages'] - return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], - 'user_data': {'key_rounds': [len(msgs)]}} - - -def _train_step(skill_model, ref_model, ckpt, samples, args): - """On-policy GRPO update over one batch, then sync weights. SFT samples ride the - same GRPOLoss with a positive constant advantage (--sft-weight).""" - trajs = [_train_trajectory(rec) for rec in samples] - advs = [float(rec['advantage']) for rec in samples] - rem = (-len(trajs)) % args.sft_batch_size - if rem: - trajs += [trajs[-1]] * rem - advs += [0.0] * rem - n, sft = len(trajs), args.sft_batch_size - mini = args.ppo_mini_batch_size if args.ppo_mini_batch_size > 0 else n - mini = max(sft, (mini // sft) * sft) - multi_step = mini < n - micro_ref, micro_old = [], [] - for i in range(0, n, sft): - mb = trajs[i:i + sft] - micro_ref.append(ref_model.forward_only(inputs=mb).get('logps')) - micro_old.append(skill_model.forward_only(inputs=mb).get('logps') if multi_step else None) - micro, n_steps = 0, 0 - for ms in range(0, n, mini): - for i in range(ms, min(ms + mini, n), sft): - k = i // sft - skill_model.forward_backward(inputs=trajs[i:i + sft], advantages=advs[i:i + sft], - old_logps=micro_old[k], ref_logps=micro_ref[k]) - micro += 1 - skill_model.clip_grad_and_step() - n_steps += 1 - ckpt.sync_weights(merge_and_sync=True) - metric = skill_model.calculate_metric(is_training=True) - n_sft = sum(1 for s in samples if s.get('sft')) - return {'n_samples': len(samples), 'n_sft': n_sft, 'n_grpo': len(samples) - n_sft, - 'n_steps': n_steps, 'n_micro_batches': micro, - 'metric': {k: (float(v) if _is_num(v) else v) for k, v in (metric or {}).items()}} - - -def _is_num(v): - try: - float(v); return True - except (TypeError, ValueError): - return False - - -# =========================================================================== -# Section H — chunk processing, records, eval (+ hard-slice rescue) -# =========================================================================== -def process_chunk(base_sampler, skill_sampler, chunk, ci, base_dp, skill_dp, args): - """skill-gen (query-only) → leak audit → with-skill greedy pass → reward → advantages. - Returns (full_records, summary, grpo_train_records, buffer_a_entries).""" - for r in chunk: - r['_cands'] = [] - # skill-gen (thinking OFF), re-sample problems with no clean candidate - flat = [] - pending = list(chunk) - for _ in range(args.skill_retries + 1): - if not pending: - break - sg_out = _run_samples(skill_sampler, [_skillgen_prompt(r['problem']) for r in pending], - args.n_skills, args.skill_max_tokens, skill_dp, - temperature=args.skill_gen_temperature, top_p=args.skill_gen_top_p, - top_k=args.skill_gen_top_k) - still = [] - for r, seqs in zip(pending, sg_out): - got = False - for s in seqs: - resp = _clean_text(getattr(s, 'decoded', '') or '') - block = _extract_skill(resp) or '' - cand = {'skills': block, 'response': resp, 'parseable': bool(block), - 'leaked': None, 'with_pass': None, 'reward': None, 'rolls': [], - 'advantage': 0.0, 'kept': False, - 'skillgen_stop': getattr(s, 'stop_reason', None), - 'skillgen_tokens': len(getattr(s, 'tokens', None) or [])} - r['_cands'].append(cand) - if block: - flat.append((r, cand)) - got = True - if not got: - still.append(r) - pending = still - - # leak audit (deterministic, observability only) - for r, c in flat: - c['leaked'] = _answer_leaked(c['skills'], r['reference_answer']) - - # with-skill greedy pass (T=0) - if flat: - ws_out = _run_samples(base_sampler, - [build_skill_solve_prompt(r['problem'], c['skills']) for r, c in flat], - 1, args.max_tokens, base_dp, temperature=0.0) - for (r, c), seqs in zip(flat, ws_out): - roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - c['rolls'] = [roll] - c['with_pass'] = 1.0 if roll['correct'] else 0.0 - c['reward'] = _skill_reward(c['parseable'], roll['correct'], roll['terminated'], - len(c['skills']), args.len_budget) - # unparseable candidates score 0 and still join the group (format pressure) - for r in chunk: - for c in r['_cands']: - if c['reward'] is None: - c['reward'] = 0.0 - - _assign_advantages(chunk, args) - - grpo = [] - for r in chunk: - for c in r['_cands']: - if abs(c.get('advantage') or 0.0) > 1e-9: - grpo.append({'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'data_id': r.get('data_id', ''), 'response': c['response'], - 'skills': c['skills'], 'advantage': c['advantage'], - 'kept': c['kept'], 'reward': c['reward'], 'sft': False}) - buffer_a = _collect_buffer_a(chunk, args) - return _full_records(chunk, ci), _chunk_summary(chunk, ci), grpo, buffer_a - - -def _roll(x): - return {k: x[k] for k in ('pred', 'correct', 'terminated', 'stop_reason', 'gen_tokens', 'text')} - - -def _full_records(chunk, ci): - out = [] - for r in chunk: - out.append({ - 'record_type': 'problem', 'chunk': ci, 'problem': r['problem'], - 'reference_answer': r['reference_answer'], 'data_id': r.get('data_id', ''), - 'candidates': [{ - 'skills': c['skills'], 'response': c['response'], 'parseable': c['parseable'], - 'leaked': c['leaked'], 'with_pass': c['with_pass'], 'reward': c.get('reward'), - 'advantage': c.get('advantage'), 'kept': c.get('kept'), - 'skillgen_stop': c.get('skillgen_stop'), 'skillgen_tokens': c.get('skillgen_tokens'), - 'rolls': [_roll(x) for x in c['rolls']], - } for c in r['_cands']], - }) - return out - - -def _mean(xs): - return sum(xs) / len(xs) if xs else 0.0 - - -def _std(xs): - if len(xs) < 2: - return 0.0 - m = _mean(xs) - return (sum((x - m) ** 2 for x in xs) / len(xs)) ** 0.5 - - -def _chunk_summary(chunk, ci): - all_cands = [c for r in chunk for c in r['_cands']] - cands = [c for c in all_cands if c['parseable']] - scored = [c for c in cands if c['with_pass'] is not None] - ws_rolls = [x for c in scored for x in c['rolls']] - # signal: fraction of groups with zero reward variance (no gradient) - group_vars, all_rewards, zero_grad, groups = [], [], 0, 0 - for r in chunk: - rewards = [c['reward'] for c in r['_cands'] if c.get('reward') is not None] - if len(rewards) < 2: - continue - groups += 1 - all_rewards.extend(rewards) - v = _std(rewards) - group_vars.append(v) - if v < 1e-9: - zero_grad += 1 - n_train = sum(1 for c in all_cands if abs(c.get('advantage') or 0.0) > 1e-9) - trunc = sum(1 for x in ws_rolls if x['stop_reason'] == 'length') - ws_acc = _mean([1.0 if any(c.get('reward') for c in r['_cands']) else 0.0 - for r in chunk if r['_cands']]) - return { - 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), - 'n_generated': len(all_cands), 'n_candidates_parseable': len(cands), - 'parse_rate': (len(cands) / len(all_cands)) if all_cands else 0.0, - 'n_leaked': sum(1 for c in cands if c['leaked']), - 'leak_rate': (sum(1 for c in cands if c['leaked']) / len(cands)) if cands else 0.0, - 'n_train_samples': n_train, 'n_groups': groups, - 'zero_grad_frac': (zero_grad / groups) if groups else 0.0, - 'reward_mean': _mean(all_rewards), 'reward_std': _std(all_rewards), - 'group_reward_std_mean': _mean(group_vars), - 'skill_tokens_mean': _mean([c.get('skillgen_tokens') or 0 for c in cands]), - 'skill_chars_mean': _mean([len(c['skills']) for c in cands]), - 'avg_withskill_pass': ws_acc, - 'candidate_withskill_pass': _mean([c['with_pass'] for c in scored]), - 'withskill_trunc_frac': (trunc / len(ws_rolls)) if ws_rolls else 0.0, - 'termination_rate_withskill': _mean([1.0 if x['terminated'] else 0.0 for x in ws_rolls]), - } - - -def run_greedy_eval(base_sampler, skill_sampler, eval_records, ci, rounds, - base_dp, skill_dp, args, base_cache): - """SEAM mean@1 on the fixed holdout: greedy skill (T=0) → greedy base solve (T=0). - Adds hard-slice (baseline_pass==0) rescue rate as a zero-cost secondary readout.""" - # baseline (frozen, cached) - todo = [r for r in eval_records if DiskCache.key_for(r['problem']) not in base_cache] - if todo: - out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], - 1, args.max_tokens, base_dp, temperature=0.0) - for r, seqs in zip(todo, out): - roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - base_cache.put(DiskCache.key_for(r['problem']), roll) - for r in eval_records: - br = base_cache.get(DiskCache.key_for(r['problem'])) - r['_baseline_pass'] = 1.0 if br['correct'] else 0.0 - # skill-gen (greedy) → with-skill (greedy) - sg_out = _run_samples(skill_sampler, [_skillgen_prompt(r['problem']) for r in eval_records], - 1, args.skill_max_tokens, skill_dp, temperature=0.0) - skills = [] - for seqs in sg_out: - if not seqs: - skills.append(('', '')) - continue - sresp = _clean_text(getattr(seqs[0], 'decoded', '') or '') - skills.append((_extract_skill(sresp) or '', sresp)) - ws_out = _run_samples(base_sampler, - [build_skill_solve_prompt(r['problem'], sk) for r, (sk, _) in zip(eval_records, skills)], - 1, args.max_tokens, base_dp, temperature=0.0) - recs = [] - for r, (sk, sresp), seqs in zip(eval_records, skills, ws_out): - roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - recs.append({ - 'record_type': 'eval_problem', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, - 'data_id': r.get('data_id', ''), 'problem': r['problem'], - 'reference_answer': r['reference_answer'], 'baseline_pass': r['_baseline_pass'], - 'skill': sk, 'skill_parseable': bool(sk), 'skill_chars': len(sk), - 'withskill_pred': roll['pred'], 'withskill_correct': roll['correct'], - 'withskill_terminated': roll['terminated'], 'withskill_stop_reason': roll['stop_reason'], - 'withskill_text': roll['text'], - }) - n = len(recs) - ws = (sum(1 for x in recs if x['withskill_correct']) / n) if n else 0.0 - base = (sum(x['baseline_pass'] for x in recs) / n) if n else 0.0 - fmt = (sum(1 for x in recs if x['skill_parseable']) / n) if n else 0.0 - term = (sum(1 for x in recs if x['withskill_terminated']) / n) if n else 0.0 - # 中文注释：难题子片救活率——baseline_pass==0 的子集里 with-skill 做对的比例。 - # 零成本（复用已算字段），是 buffer B 回路的目标量（见 skill_quality_analysis.md 第 14 节）。 - hard = [x for x in recs if not x['baseline_pass']] - hard_rescued = sum(1 for x in hard if x['withskill_correct']) - hard_rescue_rate = (hard_rescued / len(hard)) if hard else 0.0 - summary = {'record_type': 'eval_summary', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, - 'n': n, 'acc_mean1': ws, 'baseline_acc_mean1': base, 'lift_mean1': ws - base, - 'format_mean1': fmt, 'term_mean1': term, - 'hard_n': len(hard), 'hard_rescued': hard_rescued, 'hard_rescue_rate': hard_rescue_rate} - metrics = {'core/math/acc/mean@1': ws, 'core/math/baseline_acc/mean@1': base, - 'core/math/lift/mean@1': ws - base, 'core/math/format/mean@1': fmt, - 'core/math/term/mean@1': term, 'core/math/hard_rescue/mean@1': hard_rescue_rate} - return recs, summary, metrics - - -# =========================================================================== -# Section I — components, args, main -# =========================================================================== -def init_components(args): - r0, r1 = TRAIN_GPUS, TRAIN_GPUS + REF_GPUS - r2, r3 = r1 + SKILL_SAMPLER_GPUS, NUM_GPUS - twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, lazy_collect=False, groups=[ - DeviceGroup(name='train', ranks=list(range(0, r0)), device_type='GPU'), - DeviceGroup(name='ref', ranks=list(range(r0, r1)), device_type='GPU'), - DeviceGroup(name='skill_sampler', ranks=list(range(r1, r2)), device_type='GPU'), - DeviceGroup(name='base_sampler', ranks=list(range(r2, r3)), device_type='GPU')]) - - train_mesh = DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_DP, fsdp_size=TRAIN_FSDP) - skill_model = TransformersModel(model_id=MODEL_ID, device_mesh=train_mesh, remote_group='train', - ddp_config={'find_unused_parameters': False}) - skill_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) - skill_model.set_template(Template, model_id=MODEL_ID, enable_thinking=False, - max_length=args.max_model_len, truncation_strategy='delete') - skill_model.set_processor(InputProcessor, padding_free=False) - skill_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon, beta=args.kl_beta) - skill_model.set_optimizer('AdamW', lr=args.lr) - skill_model.set_lr_scheduler('CosineWarmupScheduler', num_warmup_steps=10, - num_training_steps=args.max_train_rounds) - - ref_mesh = DeviceMesh.from_sizes(world_size=REF_GPUS, dp_size=REF_DP, fsdp_size=REF_FSDP) - ref_model = TransformersModel(model_id=MODEL_ID, device_mesh=ref_mesh, remote_group='ref', - ddp_config={'find_unused_parameters': False}) - ref_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) - ref_model.set_template(Template, model_id=MODEL_ID, enable_thinking=False, - max_length=args.max_model_len, truncation_strategy='delete') - ref_model.set_processor(InputProcessor, padding_free=False) - ref_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) - - def _sampler(group, world, enable_thinking): - s = vLLMSampler(model_id=MODEL_ID, - engine_args={'gpu_memory_utilization': GPU_MEM, - 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, - device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), - remote_group=group) - s.set_template(Template, model_id=MODEL_ID, enable_thinking=enable_thinking, max_length=args.max_model_len) - return s - - skill_sampler = _sampler('skill_sampler', SKILL_SAMPLER_GPUS, enable_thinking=False) - base_sampler = _LockedSampler(_sampler('base_sampler', BASE_SAMPLER_GPUS, enable_thinking=True)) - ckpt = CheckpointEngineManager(model=skill_model, sampler=skill_sampler) - return skill_model, ref_model, skill_sampler, base_sampler, ckpt, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS - - -def _build_args(): - p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--dataset', choices=('aops', 'math'), default='aops') - p.add_argument('--n', type=int, default=2000, help='Problems loaded into the draw pool (0=all).') - p.add_argument('--exclude-data-ids', default='', - help='Comma-separated jsonl files whose data_id/problem keys are excluded.') - p.add_argument('--seed', type=int, default=42) - p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) - p.add_argument('--eval-size', type=int, default=128, help='Fixed holdout size (0 disables).') - p.add_argument('--eval-every', type=int, default=5, help='Run holdout eval every N chunks.') - p.add_argument('--chunk-size', type=int, default=16) - p.add_argument('--n-skills', type=int, default=8) - p.add_argument('--skill-retries', type=int, default=2) - p.add_argument('--skill-gen-temperature', type=float, default=1.0) - p.add_argument('--skill-gen-top-p', type=float, default=1.0) - p.add_argument('--skill-gen-top-k', type=int, default=-1) - p.add_argument('--max-model-len', type=int, default=16384) - p.add_argument('--max-tokens', type=int, default=8192) - p.add_argument('--skill-max-tokens', type=int, default=4096) - p.add_argument('--len-budget', type=int, default=600, - help='Skill length budget (chars). Reward multiplied by min(1, budget/len).') - # --- buffer / distillation --- - p.add_argument('--distill-trigger', type=int, default=300, - help='Start draining buffer A into distillation once it reaches this many entries.') - p.add_argument('--distill-batch', type=int, default=64, - help='Entries distilled per iteration while buffer A is over --distill-trigger ' - '(incremental drain: bounds per-step latency instead of one big stall).') - p.add_argument('--sft-trigger', type=int, default=100, - help='Run one SFT pass + eval when buffer B reaches this many validated entries. ' - 'Kept low: the distill funnel (has-FAIL × valid-regen × pass@k) yields only ' - '~10-15%% of buffer A, so a high threshold would rarely fire the SFT loop.') - # Plan B validation: diversity lives in the SKILL side, the executor stays at the - # deployment (greedy) decoding口径. For each buffer-A problem we regenerate K distinct - # candidate skills (high temperature), run each through ONE greedy (T=0) executor solve, - # and accept the problem iff >= M distinct skills reach a terminated-correct solve. This - # validates "the problem admits several skills that work under greedy decoding" (matches - # eval口径) rather than "one skill passes m/k times under a high-temperature executor". - p.add_argument('--passatk-k', type=int, default=8, - help='Plan B: number of DISTINCT candidate skills regenerated per problem ' - '(skill-side diversity; executor stays greedy).') - p.add_argument('--passatk-skill-temp', type=float, default=1.0, - help='Skill-model temperature when regenerating the K candidate skills ' - '(needs >0 for diversity across candidates).') - p.add_argument('--passatk-skill-top-p', type=float, default=1.0, - help='Skill-model top-p when regenerating the K candidate skills.') - p.add_argument('--passatk-m', type=int, default=2, - help='Plan B: min number of DISTINCT candidate skills that must reach a ' - 'terminated-correct GREEDY solve to accept the problem into buffer B. ' - 'Lower than pass@k-over-one-skill (default 2): requiring m distinct ' - 'greedy-effective skills is already a strong, low-noise bar.') - p.add_argument('--sft-weight', type=float, default=0.5, - help='Advantage magnitude for SFT distillation samples (-w*logp + beta*KL).') - p.add_argument('--rubric-workers', type=int, default=16) - # --- GRPO --- - p.add_argument('--sft-batch-size', type=int, default=8) - p.add_argument('--ppo-mini-batch-size', type=int, default=0) - p.add_argument('--grpo-epsilon', type=float, default=0.2) - p.add_argument('--adv-clip', type=float, default=3.0) - p.add_argument('--kl-beta', type=float, default=0.001) - p.add_argument('--lr', type=float, default=6e-6) - p.add_argument('--max-train-rounds', type=int, default=1500) - p.add_argument('--save-rounds', type=int, default=200) - p.add_argument('--output-dir', default='./output/skill_v2') - p.add_argument('--cache-dir', default='') - p.add_argument('--no-cache', action='store_true') - p.add_argument('--swanlab-project', default='twinkle') - p.add_argument('--swanlab-exp', default='') - args = p.parse_args() - if args.sft_batch_size % TRAIN_DP != 0: - raise ValueError(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple of train dp ({TRAIN_DP})') - if args.chunk_size < 1: - raise ValueError('--chunk-size must be >= 1') - return args - - -def _write(handle, row): - handle.write(json.dumps(row, ensure_ascii=False) + '\n') - - -def _swan_metrics(summary, log): - # Lean metric set: each carries independent information. Dropped as redundant — - # 中文注释：删除冗余项（换算重复）：n_groups(≈chunk_size)、reward_std(池化,组内方差已够)、 - # skill_tokens_mean(与chars重复)、leak/n(=rate×n)、candidate_withskill(与问题级重复)、 - # term/withskill(=1-trunc)、train/n_steps(恒为1)。 - d = { - 'signal/zero_grad_frac': summary['zero_grad_frac'], - 'signal/reward_mean': summary['reward_mean'], - 'signal/group_reward_std_mean': summary['group_reward_std_mean'], - 'signal/n_train_samples': summary['n_train_samples'], - 'skill/parse_rate': summary['parse_rate'], 'skill/chars_mean': summary['skill_chars_mean'], - 'leak/rate': summary['leak_rate'], - } - if summary['n_groups'] > 0: - d.update({'acc/withskill_pass': summary['avg_withskill_pass'], - 'term/withskill_trunc_frac': summary['withskill_trunc_frac']}) - if log: - d['train/n_grpo'] = log['n_grpo'] - d['train/n_sft'] = log['n_sft'] - for k, v in (log.get('metric') or {}).items(): - if not _is_num(v): - continue - if k.startswith('learning rate'): - if 'group 1' in k: - d['train/lr'] = float(v) - else: - d[f'train/{k.replace(" ", "_")}'] = float(v) - return d - - -def main(): - args = _build_args() - records, eval_records = _load_records(args) - if len(records) < args.chunk_size: - raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded ({len(records)}); raise --n') - - os.makedirs(args.output_dir, exist_ok=True) - gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') - eval_path = os.path.join(args.output_dir, 'eval_records.jsonl') - sft_path = os.path.join(args.output_dir, 'skill_dataset.jsonl') - train_log_path = os.path.join(args.output_dir, 'train_log.jsonl') - buffer_a_path = os.path.join(args.output_dir, 'buffer_a.jsonl') - distill_path = os.path.join(args.output_dir, 'distill_records.jsonl') - - use_swan = swanlab is not None and os.environ.get('SWANLAB_MODE') != 'disabled' - if use_swan: - swanlab.init(project=args.swanlab_project, experiment_name=(args.swanlab_exp or None), - config={'model': MODEL_ID, 'dataset': args.dataset, 'n': len(records), - 'eval_n': len(eval_records), 'n_skills': args.n_skills, - 'len_budget': args.len_budget, 'distill_trigger': args.distill_trigger, - 'sft_trigger': args.sft_trigger, 'passatk_k': args.passatk_k, - 'passatk_m': args.passatk_m, 'passatk_skill_temp': args.passatk_skill_temp, - 'sft_weight': args.sft_weight, 'lr': args.lr}) - - skill_model, ref_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp = init_components(args) - checker = build_rubric_checker() - if checker is None: - sys.stderr.write('[v2] no LLM backup env -> buffer B distillation DISABLED (GRPO only)\n') - - cache_dir = args.cache_dir or os.path.join(args.output_dir, 'cache') - os.makedirs(cache_dir, exist_ok=True) - eval_base_cache = DiskCache(os.path.join(cache_dir, 'eval_baseline.jsonl'), not args.no_cache) - - cfg = {'record_type': 'config', 'model': MODEL_ID, 'dataset': args.dataset, - 'n': len(records), 'eval_n': len(eval_records), 'seed': args.seed, - 'n_skills': args.n_skills, 'len_budget': args.len_budget, - 'distill_trigger': args.distill_trigger, 'sft_trigger': args.sft_trigger, - 'passatk_k': args.passatk_k, 'passatk_m': args.passatk_m, - 'passatk_skill_temp': args.passatk_skill_temp, 'passatk_skill_top_p': args.passatk_skill_top_p, - 'sft_weight': args.sft_weight, - 'grpo_epsilon': args.grpo_epsilon, 'kl_beta': args.kl_beta, 'lr': args.lr, - 'rubric_check': bool(checker), 'max_train_rounds': args.max_train_rounds, - 'started': int(time.time())} - - hist_a: List[Dict[str, Any]] = [] # buffer A accumulator (in-memory + jsonl) - sft_queue: List[Dict[str, Any]] = [] # buffer B: validated SFT records awaiting an SFT pass - rounds = 0 # GRPO rounds only (gates --max-train-rounds + save cadence) - sft_rounds = 0 # SFT passes (separate: must NOT eat the GRPO round budget) - pool = ProblemPool(records, args.seed) - - # Background rubric pre-diagnosis (做法 B): the moment a failure trajectory lands in - # buffer A, fire its teacher-rubric call on a daemon thread pool. The API round-trip - # then overlaps with GRPO GPU work, so by the time --distill-trigger fires the - # diagnoses are usually already cached on each entry ('_rubric_diag'); distill_buffer - # only pays for the stragglers. Entries are dicts held by reference, so the worker - # writes the result straight onto the entry. - # 中文注释：失败轨迹一进 buffer A 就后台异步跑 rubric，API 等待藏进 GPU 训练时间； - # 到蒸馏时诊断多已缓存在条目上，distill_buffer 只补漏。 - prediag_pool = (ThreadPoolExecutor(max_workers=max(1, args.rubric_workers), - thread_name_prefix='rubric-prediag') - if checker else None) - - def _prediagnose(entry: Dict[str, Any]): - entry['_rubric_diag'] = _diagnose_entry(checker, entry) or '' - - with open(gen_path, 'w', encoding='utf-8') as gen_f, \ - open(eval_path, 'w', encoding='utf-8') as eval_f, \ - open(sft_path, 'w', encoding='utf-8') as sft_f, \ - open(train_log_path, 'w', encoding='utf-8') as tlog, \ - open(distill_path, 'w', encoding='utf-8') as distill_f, \ - open(buffer_a_path, 'w', encoding='utf-8') as buf_f: - for f in (gen_f, eval_f, sft_f, tlog, distill_f): - _write(f, cfg) - - def _do_eval(gstep): - recs, summary, metrics = run_greedy_eval( - base_sampler, skill_sampler, eval_records, gstep, rounds, base_dp, skill_dp, - args, eval_base_cache) - for rec in recs: - _write(eval_f, rec) - _write(eval_f, summary) - eval_f.flush() - if use_swan: - swanlab.log({f'eval/{k}': v for k, v in metrics.items()}, step=max(gstep, 0)) - sys.stderr.write( - f'[eval] g{gstep}: n={summary["n"]} acc={summary["baseline_acc_mean1"]:.3f}' - f'->{summary["acc_mean1"]:.3f} lift={summary["lift_mean1"]:+.3f} ' - f'hard_rescue={summary["hard_rescue_rate"]:.3f}({summary["hard_rescued"]}/{summary["hard_n"]}) ' - f'fmt={summary["format_mean1"]:.2f} rounds={rounds}\n') - - if eval_records: - _do_eval(-1) - - gstep = 0 - while rounds < args.max_train_rounds: - chunk = pool.draw(args.chunk_size) - full, summary, grpo, buffer_a = process_chunk( - base_sampler, skill_sampler, chunk, gstep, base_dp, skill_dp, args) - - # accumulate buffer A (only when a rubric checker exists to consume it; - # 中文注释：无 checker 时蒸馏永不触发，不累积以免内存无限增长) - if checker: - for e in buffer_a: - _write(buf_f, e) - prediag_pool.submit(_prediagnose, e) # 后台异步预诊断，不阻塞主循环 - buf_f.flush() - hist_a.extend(buffer_a) - - # GRPO train step (only when there is signal) - log = None - if grpo: - log = _train_step(skill_model, ref_model, ckpt, grpo, args) - rounds += 1 - log.update({'record_type': 'train_round', 'round': rounds, 'chunk': gstep, - 'epoch': pool.epoch, 'kind': 'grpo', 'ts': int(time.time())}) - _write(tlog, log) - tlog.flush() - if rounds % args.save_rounds == 0: - skill_model.save(f'skill-v2-{rounds}', output_dir=args.output_dir) - - summary['rounds_done'], summary['epoch'] = rounds, pool.epoch - summary['buffer_a_size'], summary['sft_queue_size'] = len(hist_a), len(sft_queue) - for rec in full: - _write(gen_f, rec) - _write(gen_f, summary) - gen_f.flush() - - sys.stderr.write( - f'[gen] e{pool.epoch} g{gstep}: n={summary["n"]} ' - f'clean={summary["n_candidates_parseable"]} 0grad={summary["zero_grad_frac"]:.2f} ' - f'R={summary["reward_mean"]:.2f}+-{summary["reward_std"]:.2f} ' - f'ws_acc={summary["avg_withskill_pass"]:.2f} chars={summary["skill_chars_mean"]:.0f} ' - f'bufA={len(hist_a)} bufB={len(sft_queue)} rounds={rounds}\n') - if use_swan: - m = _swan_metrics(summary, log) - m['buffer/a_size'] = float(len(hist_a)) - m['buffer/b_size'] = float(len(sft_queue)) - swanlab.log(m, step=gstep) - - # --- distillation: once buffer A fills, drain it INCREMENTALLY in bounded - # batches (--distill-batch) so a large buffer never stalls the loop for tens - # of minutes; each iteration processes one batch, interleaved with GRPO. - # 中文注释：增量分批蒸馏——buffer A 满后每轮只处理 --distill-batch 条，把一次性 - # 几十分钟阻塞摊成每轮几分钟小停顿；两段验证(见 distill_buffer)再砍验证算力。 - if checker and len(hist_a) >= args.distill_trigger: - batch = hist_a[:args.distill_batch] - hist_a = hist_a[args.distill_batch:] - new_sft, distill_recs = distill_buffer(batch, skill_sampler, base_sampler, checker, - skill_dp, base_dp, args) - for rec in distill_recs: # 逐 entry 审计记录：rubric_diag + 候选 skill + 贪心解 + stage - rec['chunk'] = gstep - _write(distill_f, rec) - distill_f.flush() - for rec in new_sft: - _write(sft_f, rec) - sft_f.flush() - sft_queue.extend(new_sft) - - # --- SFT trigger: buffer B full → one SFT pass + eval --- - did_eval = False - if len(sft_queue) >= args.sft_trigger: - sys.stderr.write(f'[sft] triggered at bufB={len(sft_queue)}\n') - sft_samples = [{**s, 'advantage': float(args.sft_weight)} for s in sft_queue] - sft_log = _train_step(skill_model, ref_model, ckpt, sft_samples, args) - sft_rounds += 1 # 中文注释：SFT 用独立计数，不占用 GRPO 的 rounds 配额/save 节奏 - sft_log.update({'record_type': 'train_round', 'round': rounds, 'sft_round': sft_rounds, - 'chunk': gstep, 'epoch': pool.epoch, 'kind': 'sft', 'ts': int(time.time())}) - _write(tlog, sft_log) - tlog.flush() - sft_queue = [] - skill_model.save(f'skill-v2-sft{sft_rounds}', output_dir=args.output_dir) # 大改动后落盘 - if eval_records: # 中文注释：SFT 后立即 eval，测灾难性遗忘/真提升（第 11.3/13.4 节） - _do_eval(gstep) - did_eval = True - - if eval_records and not did_eval and (gstep + 1) % args.eval_every == 0: - _do_eval(gstep) - gstep += 1 - - if prediag_pool is not None: - prediag_pool.shutdown(wait=False, cancel_futures=True) # 丢弃未完成的后台预诊断 - eval_base_cache.close() - skill_model.save('skill-v2-final', output_dir=args.output_dir) - sys.stderr.write(f'[v2] done: {rounds} rounds over {gstep} chunks / {pool.epoch} epochs\n') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/cold_start/train_cold_start.py b/cookbook/exp/legacy/cold_start/train_cold_start.py similarity index 100% rename from cookbook/exp/cold_start/train_cold_start.py rename to cookbook/exp/legacy/cold_start/train_cold_start.py diff --git a/cookbook/exp/condenser/dataset.py b/cookbook/exp/legacy/condenser/dataset.py similarity index 100% rename from cookbook/exp/condenser/dataset.py rename to cookbook/exp/legacy/condenser/dataset.py diff --git a/cookbook/exp/condenser/make_condenser_dataset.py b/cookbook/exp/legacy/condenser/make_condenser_dataset.py similarity index 100% rename from cookbook/exp/condenser/make_condenser_dataset.py rename to cookbook/exp/legacy/condenser/make_condenser_dataset.py diff --git a/cookbook/exp/condenser/train_condenser_ddp.py b/cookbook/exp/legacy/condenser/train_condenser_ddp.py similarity index 100% rename from cookbook/exp/condenser/train_condenser_ddp.py rename to cookbook/exp/legacy/condenser/train_condenser_ddp.py diff --git a/cookbook/exp/condenser/untested/eval_condensed.py b/cookbook/exp/legacy/condenser/untested/eval_condensed.py similarity index 100% rename from cookbook/exp/condenser/untested/eval_condensed.py rename to cookbook/exp/legacy/condenser/untested/eval_condensed.py diff --git a/cookbook/exp/data_pipeline/audit_rubric.py b/cookbook/exp/legacy/data_pipeline/audit_rubric.py similarity index 100% rename from cookbook/exp/data_pipeline/audit_rubric.py rename to cookbook/exp/legacy/data_pipeline/audit_rubric.py diff --git a/cookbook/exp/data_pipeline/process_and_save.py b/cookbook/exp/legacy/data_pipeline/process_and_save.py similarity index 100% rename from cookbook/exp/data_pipeline/process_and_save.py rename to cookbook/exp/legacy/data_pipeline/process_and_save.py diff --git a/cookbook/exp/rl/grpo.py b/cookbook/exp/legacy/rl/grpo.py similarity index 100% rename from cookbook/exp/rl/grpo.py rename to cookbook/exp/legacy/rl/grpo.py diff --git a/cookbook/exp/rl/rag_hint_grpo.py b/cookbook/exp/legacy/rl/rag_hint_grpo.py similarity index 100% rename from cookbook/exp/rl/rag_hint_grpo.py rename to cookbook/exp/legacy/rl/rag_hint_grpo.py diff --git a/cookbook/exp/skill2lora/good_skill_hard_fail/analyze_3way.py b/cookbook/exp/skill2lora/good_skill_hard_fail/analyze_3way.py new file mode 100644 index 000000000..a80756be7 --- /dev/null +++ b/cookbook/exp/skill2lora/good_skill_hard_fail/analyze_3way.py @@ -0,0 +1,177 @@ +#!/usr/bin/env python3 +# analyze_3way.py — 三路探针对比分析（A/B/C），全部结论可复现、可引用。 +# A = 老环境 nothink : skillcfg_full_off.jsonl + reflexion_full_off.jsonl (根目录) +# B = 新环境 think : skillcfg_full_on.jsonl + reflexion_full_on.jsonl (根目录) +# C = 新环境 nothink : env_runs/vllm_0.23.0/skillcfg_full_off.jsonl + reflexion_full_off.jsonl +# 对比轴: +# vLLM/环境影响 = A vs C (都 nothink) +# think 影响 = C vs B (都新环境) +# 用法: python3 analyze_3way.py > analysis_out/report.txt 2>&1 +import json, os, math, statistics as st, re, random +from collections import defaultdict, Counter + +def _f(x): + try: + return float(x) + except (TypeError, ValueError): + return 0.0 + +HERE = os.path.dirname(os.path.abspath(__file__)) +GROUPS = { + 'A_old_nothink': ('skillcfg_full_off.jsonl', 'reflexion_full_off.jsonl'), + 'B_new_think': ('skillcfg_full_on.jsonl', 'reflexion_full_on.jsonl'), + 'C_new_nothink': ('env_runs/vllm_0.23.0/skillcfg_full_off.jsonl', + 'env_runs/vllm_0.23.0/reflexion_full_off.jsonl'), +} +SKILL_CFGS = ['P1_narrative','P2_combo','P3_toy','P4_card','P5_pitfall','P6_seam','P7_minimal'] +REFL_CFGS = ['R4_blind','D1_needle','D2_narr','D3_toyfix'] + + +def load(path): + """流式读取 -> list[dict]（只保留分析需要的字段，省内存）""" + keep = ('config','data_id','sample_idx','baseline_pass','parseable','skill_chars', + 'leaked','skillgen_stop','skillgen_tokens','withskill_correct','withskill_stop', + 'withskill_tokens','skill') + rows = [] + with open(os.path.join(HERE, path)) as f: + for line in f: + line = line.strip() + if not line: + continue + r = json.loads(line) + rows.append({k: r.get(k) for k in keep}) + return rows + + +def agg_metrics(rows, cfgs): + """返回 {config: metrics}""" + out = {} + for name in cfgs: + sub = [t for t in rows if t['config'] == name] + n = len(sub) + if n == 0: + continue + parse = sum(bool(t['parseable']) for t in sub)/n + leak = sum(bool(t['leaked']) for t in sub)/n + trunc = sum(1 for t in sub if t['withskill_stop']=='length')/n + sg_trunc = sum(1 for t in sub if t['skillgen_stop']=='length')/n + chars = [t['skill_chars'] for t in sub if t['parseable']] + med_chars = int(st.median(chars)) if chars else 0 + acc = sum(bool(t['withskill_correct']) for t in sub)/n + base = sum(_f(t['baseline_pass']) for t in sub)/n + # 题级 pass@k / hard 救活@k + byq = defaultdict(list) + for t in sub: + byq[t['data_id']].append(t) + p_at_k = sum(1 for v in byq.values() if any(x['withskill_correct'] for x in v))/len(byq) + hardq = {d:v for d,v in byq.items() if _f(v[0]['baseline_pass'])==0} + rescue = (sum(1 for v in hardq.values() if any(x['withskill_correct'] for x in v))/len(hardq)) if hardq else 0.0 + # 去混杂子集: leaked=0 且 parseable=1 + clean = [t for t in sub if (not t['leaked']) and t['parseable']] + acc_clean = (sum(bool(t['withskill_correct']) for t in clean)/len(clean)) if clean else float('nan') + out[name] = dict(n=n, parse=parse, leak=leak, trunc=trunc, sg_trunc=sg_trunc, + med_chars=med_chars, acc=acc, base=base, lift=acc-base, + p_at_k=p_at_k, rescue=rescue, n_clean=len(clean), acc_clean=acc_clean, + sg_tokens_med=int(st.median([t['skillgen_tokens'] or 0 for t in sub]))) + return out + + +def diversity(rows, cfgs): + """题内 8 rollout 多样性: 去重率 + 词级 pairwise Jaccard(距离) + 字符长度 CV""" + _word = re.compile(r"[A-Za-z]+|\d+") + out = {} + for name in cfgs: + sub = [t for t in rows if t['config']==name] + byq = defaultdict(list) + for t in sub: + byq[t['data_id']].append(t) + uniq_ratios, jac_dists, all_chars = [], [], [] + for v in byq.values(): + skills = [(t['skill'] or '') for t in v] + all_chars += [len(s) for s in skills] + uniq_ratios.append(len(set(skills))/len(skills)) + sets = [set(_word.findall(s.lower())) for s in skills] + ds = [] + for i in range(len(sets)): + for j in range(i+1, len(sets)): + a,b = sets[i],sets[j] + if not a and not b: + ds.append(0.0); continue + inter=len(a&b); uni=len(a|b) or 1 + ds.append(1 - inter/uni) # 1=完全不同,0=完全相同 + if ds: + jac_dists.append(sum(ds)/len(ds)) + cv = (st.pstdev(all_chars)/ (sum(all_chars)/len(all_chars))) if all_chars and sum(all_chars) else 0.0 + out[name] = dict(uniq=sum(uniq_ratios)/len(uniq_ratios), + jac=sum(jac_dists)/len(jac_dists) if jac_dists else 0.0, + char_cv=cv) + return out + + +def fmt_table(title, gm, cfgs, cols): + print(f"\n### {title}") + head = "%-14s " % "config" + " ".join("%-9s" % c for c,_ in cols) + print(head); print("-"*len(head)) + for name in cfgs: + if name not in gm: + continue + m = gm[name] + row = "%-14s " % name + " ".join(("%-9.3f" if isinstance(m[k],float) else "%-9d") % m[k] for _,k in cols) + print(row) + + +def main(): + os.makedirs(os.path.join(HERE,'analysis_out'), exist_ok=True) + data = {} + for g,(sf,rf) in GROUPS.items(): + data[g] = dict(skill=load(sf), refl=load(rf)) + print(f"[load] {g}: skillcfg={len(data[g]['skill'])} reflexion={len(data[g]['refl'])}") + + # ---------- 对齐校验: 三组是否同题同 idx ---------- + print("\n" + "="*80 + "\n[对齐校验] 三组 (config,data_id,sample_idx) 键集合是否一致") + def keyset(rows): + return set((t['config'],t['data_id'],t['sample_idx']) for t in rows) + ka,kb,kc = keyset(data['A_old_nothink']['skill']),keyset(data['B_new_think']['skill']),keyset(data['C_new_nothink']['skill']) + print(f" A∩C 交集/并集(nothink 对比): {len(ka&kc)}/{len(ka|kc)} A独有={len(ka-kc)} C独有={len(kc-ka)}") + print(f" C∩B 交集/并集(think 对比): {len(kc&kb)}/{len(kc|kb)} C独有={len(kc-kb)} B独有={len(kb-kc)}") + + cols_skill = [('n','n'),('parse','parse'),('leak','leak'),('trunc','trunc'), + ('sgTrunc','sg_trunc'),('chars','med_chars'),('base','base'), + ('acc@1','acc'),('lift','lift'),('pass@k','p_at_k'),('rescue@k','rescue'), + ('accClean','acc_clean')] + # ---------- Q1+Q4: 每组每 config 指标 ---------- + print("\n" + "="*80 + "\n[Q1/Q4] skillcfg 每类准确率与质量指标") + GM = {} + for g in GROUPS: + GM[g] = agg_metrics(data[g]['skill'], SKILL_CFGS) + fmt_table(f"{g} (skillcfg)", GM[g], SKILL_CFGS, cols_skill) + + print("\n" + "="*80 + "\n[Q1/Q4] reflexion 每类救活指标") + RM = {} + for g in GROUPS: + RM[g] = agg_metrics(data[g]['refl'], REFL_CFGS) + fmt_table(f"{g} (reflexion)", RM[g], REFL_CFGS, cols_skill) + + # ---------- Q2: 横向差分 ---------- + print("\n" + "="*80 + "\n[Q2] 环境(vLLM/栈)影响 = C - A (同 nothink) ；think 影响 = B - C (同新环境)") + print("%-14s %-22s %-22s" % ("config","env(C-A) acc/lift/parse","think(B-C) acc/lift/parse")) + for name in SKILL_CFGS: + a,c,b = GM['A_old_nothink'].get(name),GM['C_new_nothink'].get(name),GM['B_new_think'].get(name) + if not(a and c and b): continue + env = f"{c['acc']-a['acc']:+.3f}/{c['lift']-a['lift']:+.3f}/{c['parse']-a['parse']:+.3f}" + thk = f"{b['acc']-c['acc']:+.3f}/{b['lift']-c['lift']:+.3f}/{b['parse']-c['parse']:+.3f}" + print("%-14s %-22s %-22s" % (name, env, thk)) + + # ---------- Q3: 多样性 ---------- + print("\n" + "="*80 + "\n[Q3] skill 生成多样性 (题内 8 rollout；uniq=去重率 jac=词级平均两两距离 char_cv=长度变异)") + for g in GROUPS: + dv = diversity(data[g]['skill'], SKILL_CFGS) + print(f"\n### {g}") + print("%-14s %-8s %-8s %-8s" % ("config","uniq","jac","char_cv")) + for name in SKILL_CFGS: + m=dv[name]; print("%-14s %-8.3f %-8.3f %-8.3f" % (name,m['uniq'],m['jac'],m['char_cv'])) + + print("\n[done] 详见各节表格；抽样交叉验证见 sample_probe.py 输出") + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/skill2lora/good_skill_hard_fail/eval_skill_probe.py b/cookbook/exp/skill2lora/good_skill_hard_fail/eval_skill_probe.py new file mode 100644 index 000000000..0a0619f9d --- /dev/null +++ b/cookbook/exp/skill2lora/good_skill_hard_fail/eval_skill_probe.py @@ -0,0 +1,284 @@ +#!/usr/bin/env python3 +"""eval_skill_probe.py — 自包含的"单次 eval"探针，用于手工迭代 skill。 + +目的：固定一批难题（executor 持续解不出的），保持 executor 侧 prompt/采样/判分与 +train_skill_v2.py 的 eval 完全一致（v2 模式：单 user turn + \\boxed{} 答案格式， +executor 用 base_sampler、enable_thinking=True、greedy 温度 0），但允许自由替换每题 +的 skill(experience) 内容。反复替换 skill 重跑，一旦某题解对，就把该 skill 落盘到 +winning_skills.jsonl，从而观察"能让 executor 生效的 skill 到底长什么样"。 + +不 import train_skill_v2.py（自包含）；executor 侧逻辑逐段复刻自该文件（2026-07）。 + +用法： + # 1) 生成 trials 模板（默认 10 道难题，skill 待填）： + python3 eval_skill_probe.py --init + # 2) 编辑 trials.jsonl，给每题填不同的 skill，然后跑： + python3 eval_skill_probe.py + # 只测某题 / 调大 max_tokens（验证"截断"型失败是否靠加长度能救）： + python3 eval_skill_probe.py --only seam:val:128 --max-tokens 12000 + +trials.jsonl 每行一个 JSON（# 开头的行会被忽略，可当注释）： + {"data_id": "seam:val:128", "skill": "..."} # 用该 skill 解题 + {"data_id": "seam:val:128", "skill": ""} # 空 skill = baseline + {"data_id": "seam:val:128", "tag": "v3", "skill":"..."}# tag 便于区分同题多次试验 +problem / reference_answer 默认按 data_id 从 eval_records.jsonl 解析；也可在行内直接 +提供 "problem" / "reference_answer" 覆盖（用于测 eval 集之外的题）。 + +GPU：仅起一个 executor sampler。默认 EXEC_GPUS=2；若训练在占卡，先用 +CUDA_VISIBLE_DEVICES 指定空闲卡，或设 EXEC_GPUS=1。 +""" +import argparse +import copy +import json +import os +import re +import sys +from typing import Dict, Optional + +import twinkle +from twinkle import DeviceGroup, DeviceMesh +from twinkle.data_format import SamplingParams +from twinkle.sampler import vLLMSampler +from twinkle.template import Template + +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') +GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.8)) +EXEC_GPUS = int(os.environ.get('EXEC_GPUS', 2)) +MAX_MODEL_LEN = int(os.environ.get('MAX_MODEL_LEN', 16384)) +DEFAULT_EVAL_RECORDS = os.path.join(SCRIPT_DIR, '..', 'skill_v2', 'eval_records.jsonl') + +# 与本文件夹 10 个 case 对应的难题（executor 全程持续解不出、baseline 也全 0）。 +DEFAULT_DATA_IDS = ['seam:val:81', 'seam:val:92', 'seam:val:148', 'seam:val:128', 'seam:val:127', + 'seam:val:29', 'seam:val:94', 'seam:val:176', 'seam:val:12', 'seam:val:151'] + +# =========================================================================== +# executor prompt / answer format —— 逐字复刻 train_skill_v2.py 的 v2 分支 +# =========================================================================== +_ANSWER_FORMAT_V2 = ('Present your reasoning, then put ONLY the final numeric result inside ' + '\\boxed{}. For example: \\boxed{42}.') + + +def build_direct_prompt(problem): + content = f'The problem you need to solve:\n{problem}\n\n' + _ANSWER_FORMAT_V2 + return {'messages': [{'role': 'user', 'content': content}]} + + +def build_skill_solve_prompt(problem, skill): + skill = (skill or '').strip() + if not skill: + return build_direct_prompt(problem) + content = (f'The problem you need to solve:\n{problem}\n\n' + 'Skill hint:\nFor this problem, a skill-generation model has analyzed it and ' + 'provided some advisory skills:\n' + f'{skill}\n' + 'Prefer using its techniques when they fit, but if you have a more efficient or ' + 'clearer correct method, you may use it. If you diverge from this advice, briefly ' + 'explain why. Be concise and accurate.\n' + + _ANSWER_FORMAT_V2) + return {'messages': [{'role': 'user', 'content': content}]} + + +# =========================================================================== +# boxed 抽取 + SEAM lpem 风格数值判分 —— 逐字复刻 +# =========================================================================== +_BOXED_RE = re.compile(r'\\boxed\s*\{') + + +def extract_boxed(text: str) -> Optional[str]: + if not text: + return None + last = None + for m in _BOXED_RE.finditer(text): + depth, i = 1, m.end() + while i < len(text) and depth > 0: + depth += (text[i] == '{') - (text[i] == '}') + i += 1 + if depth == 0: + last = text[m.end():i - 1].strip() + return last + + +_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') + + +def _clean_text(decoded: Optional[str]) -> str: + return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() + + +_SEAM_TAG_RE = re.compile(r'<\s*answer\s*>(.*?)<\s*/\s*answer\s*>', re.I | re.S) +_SEAM_BOX_RE = re.compile(r'(?:\\{1,2}\(|)\\{1,2}boxed\s*\{\s*([^}]*)\s*}(?:\)|)', re.S) +_SEAM_INLINE_RE = re.compile(r'\$([^$]+)\$|\\\(([^)]+)\\\)', re.S) +_SEAM_FRAC_RE = re.compile(r'(-?\d+(?:\.\d+)?)/(-?\d+(?:\.\d+)?)') +_SEAM_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') + + +def _seam_norm(num: str) -> str: + try: + f = float(num) + return str(int(f)) if f == int(f) else str(f) + except Exception: + return num.strip() + + +def _seam_sanitize(txt: str) -> str: + txt = (txt or '').strip() + if (m := _SEAM_TAG_RE.search(txt)): + txt = m.group(1).strip() + elif (m := _SEAM_BOX_RE.search(txt)): + txt = m.group(1).strip() + elif (m := _SEAM_INLINE_RE.search(txt)): + txt = (m.group(1) or m.group(2)).strip() + txt = re.sub(r'\\frac\s*\{\s*([^}]+?)\s*}\s*\{\s*([^}]+?)\s*}', r'\1/\2', txt) + if (m := _SEAM_FRAC_RE.search(txt)): + p, q = map(float, m.groups()) + if q: + return _seam_norm(str(p / q)) + if (m := _SEAM_NUM_RE.search(txt)): + return _seam_norm(m.group()) + return txt + + +def _parse_seq(seq, gold: str) -> Dict: + text = _clean_text(getattr(seq, 'decoded', '') or '') + raw = extract_boxed(text) + pred = _seam_sanitize(raw) if raw else None + correct = bool(pred) and (pred == _seam_sanitize(str(gold))) + terminated = getattr(seq, 'stop_reason', None) != 'length' + return {'pred': pred, 'correct': correct, 'terminated': terminated, + 'stop_reason': getattr(seq, 'stop_reason', None), + 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} + + +def _run_samples(sampler, prompts, max_tokens, gen_dp): + """greedy（T=0）单样本，对齐 v2 eval。gen_dp>len 时按最后一条 padding 补齐。""" + if not prompts: + return [] + params = SamplingParams(max_tokens=max_tokens, temperature=0.0, top_p=1.0, num_samples=1) + padded = prompts + if gen_dp > 1 and 0 < len(prompts) < gen_dp: + padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] + responses = sampler.sample(padded, params)[:len(prompts)] + return [list(r.sequences) if (r and r.sequences) else [] for r in responses] + + +# =========================================================================== +# problem 查表 + trials 载入 +# =========================================================================== +def load_problems(eval_records_path) -> Dict[str, Dict]: + probs: Dict[str, Dict] = {} + if not os.path.exists(eval_records_path): + return probs + for line in open(eval_records_path): + line = line.strip() + if not line: + continue + try: + r = json.loads(line) + except Exception: + continue + if r.get('record_type') != 'eval_problem': + continue + did = r.get('data_id') + if did and did not in probs: + probs[did] = {'problem': r['problem'], 'reference_answer': r['reference_answer']} + return probs + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument('--trials', default=os.path.join(SCRIPT_DIR, 'trials.jsonl')) + ap.add_argument('--eval-records', default=DEFAULT_EVAL_RECORDS) + ap.add_argument('--out', default=os.path.join(SCRIPT_DIR, 'winning_skills.jsonl')) + ap.add_argument('--max-tokens', type=int, default=8192, help='对齐 v2 eval 的 --max-tokens') + ap.add_argument('--only', default=None, help='只跑某个 data_id') + ap.add_argument('--init', action='store_true', help='生成 trials.jsonl 模板后退出') + ap.add_argument('--dump-text', action='store_true', help='把每个 trial 的 executor 全文落到 probe_texts/') + args = ap.parse_args() + + problems = load_problems(args.eval_records) + + if args.init: + with open(args.trials, 'w') as f: + f.write('# 每行一个 trial；# 开头的行被忽略。skill="" 即 baseline。改 skill 后重跑本脚本。\n') + for did in DEFAULT_DATA_IDS: + p = problems.get(did, {}) + f.write(json.dumps({'data_id': did, 'tag': 'baseline', 'skill': '', + 'reference_answer': p.get('reference_answer')}, ensure_ascii=False) + '\n') + print(f'已写模板 {args.trials}（{len(DEFAULT_DATA_IDS)} 题，skill 待填）。编辑后去掉 --init 再跑。') + return + + if not os.path.exists(args.trials): + print(f'找不到 {args.trials}，先跑：python3 eval_skill_probe.py --init') + sys.exit(1) + + trials = [] + for line in open(args.trials): + line = line.strip() + if not line or line.startswith('#'): + continue + trials.append(json.loads(line)) + if args.only: + trials = [t for t in trials if t.get('data_id') == args.only] + for t in trials: + p = problems.get(t.get('data_id'), {}) + t.setdefault('problem', p.get('problem')) + t.setdefault('reference_answer', p.get('reference_answer')) + skipped = [t for t in trials if t.get('problem') is None] + for t in skipped: + print(f'[warn] data_id={t.get("data_id")} 无题面（eval_records 找不到且行内未给 problem），跳过') + trials = [t for t in trials if t.get('problem') is not None] + if not trials: + print('没有可跑的 trial。') + return + + # 仅起一个 executor sampler（enable_thinking=True，对齐 v2 eval 的 base_sampler） + twinkle.initialize(mode='ray', nproc_per_node=EXEC_GPUS, lazy_collect=False, + groups=[DeviceGroup(name='exec', ranks=list(range(EXEC_GPUS)), device_type='GPU')]) + sampler = vLLMSampler(model_id=MODEL_ID, + engine_args={'gpu_memory_utilization': GPU_MEM, + 'max_model_len': MAX_MODEL_LEN, 'tensor_parallel_size': 1}, + device_mesh=DeviceMesh.from_sizes(world_size=EXEC_GPUS, dp_size=EXEC_GPUS), + remote_group='exec') + sampler.set_template(Template, model_id=MODEL_ID, enable_thinking=True, max_length=MAX_MODEL_LEN) + + prompts = [build_skill_solve_prompt(t['problem'], t.get('skill', '')) for t in trials] + outs = _run_samples(sampler, prompts, args.max_tokens, EXEC_GPUS) + + if args.dump_text: + os.makedirs(os.path.join(SCRIPT_DIR, 'probe_texts'), exist_ok=True) + + n_ok = 0 + print('\n' + '=' * 94) + print('%-16s %-10s %-4s %-10s %-10s %-7s %s' % ('data_id', 'tag', '对?', 'pred', 'gold', 'tokens', 'note')) + print('-' * 94) + win_f = open(args.out, 'a') + for idx, (t, seqs) in enumerate(zip(trials, outs)): + roll = _parse_seq(seqs[0], t['reference_answer']) if seqs else { + 'pred': None, 'correct': False, 'terminated': False, + 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} + ok = '✓' if roll['correct'] else '✗' + note = '[截断]' if roll['stop_reason'] == 'length' else '' + if roll['correct']: + n_ok += 1 + print('%-16s %-10s %-4s %-10s %-10s %-7d %s' % ( + t.get('data_id', ''), str(t.get('tag', ''))[:10], ok, + str(roll['pred'])[:10], str(t['reference_answer'])[:10], roll['gen_tokens'], note)) + if args.dump_text: + fn = os.path.join(SCRIPT_DIR, 'probe_texts', + 'trial_%02d_%s_%s.txt' % (idx, str(t.get('data_id', '')).replace(':', '_'), + str(t.get('tag', '')))) + with open(fn, 'w') as tf: + tf.write('SKILL:\n' + (t.get('skill') or '') + '\n\n' + '=' * 60 + '\nEXECUTOR OUTPUT:\n' + roll['text']) + if roll['correct']: + win_f.write(json.dumps({'data_id': t.get('data_id'), 'tag': t.get('tag'), + 'reference_answer': t['reference_answer'], 'pred': roll['pred'], + 'gen_tokens': roll['gen_tokens'], 'max_tokens': args.max_tokens, + 'skill': t.get('skill', '')}, ensure_ascii=False) + '\n') + win_f.close() + print('-' * 94) + print(f'共 {len(trials)} 个 trial，解对 {n_ok} 个。成功的 skill 已追加到 {args.out}') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/skill2lora/good_skill_hard_fail/leak_decomp.py b/cookbook/exp/skill2lora/good_skill_hard_fail/leak_decomp.py new file mode 100644 index 000000000..665e24ec2 --- /dev/null +++ b/cookbook/exp/skill2lora/good_skill_hard_fail/leak_decomp.py @@ -0,0 +1,146 @@ +#!/usr/bin/env python3 +# leak_decomp.py — 严格重审 "think 提升是否由答案泄漏驱动"。 +# 上次教训: 宽松判定曾导致严重误判。本脚本做四路独立检验: +# T1 安慰剂测试: 用"别题答案"跑同一 leak 规则 → 估计 leak 标记的偶然假阳性地板 +# (think skill 长且数字密集, 答案数字偶然出现的概率天然更高) +# T2 条件分解: parseable 记录拆 leaked/clean, 分别算 acc + 占比 → 贡献分解 +# T3 反事实: 把 leaked 样本的 acc 替换成同组 clean acc → think 优势还剩多少 +# T4 严格判分: 对 withskill_text 用严格 boxed 精确匹配重新判分, +# 检验现行 grade(_seam_sanitize 取首数字等宽松步骤) 是否给 think 虚增 acc +# 难度控制: T2/T3 同时在 baseline_pass==0 (hard) 子集上重复, 排除"泄漏样本恰好是简单题"。 +import json, os, re +from collections import defaultdict + +HERE = os.path.dirname(os.path.abspath(__file__)) +FILES = {'A_old_off': 'skillcfg_full_off.jsonl', + 'B_new_on': 'skillcfg_full_on.jsonl', + 'C_new_off': 'env_runs/vllm_0.23.0/skillcfg_full_off.jsonl'} +CFGS = ['P1_narrative','P2_combo','P3_toy','P4_card','P5_pitfall','P6_seam','P7_minimal'] + +_NUM = re.compile(r'-?\d+(\.\d+)?') + +def sanitize(x): + s = str(x).strip() + m = _NUM.search(s) + if m and m.group() == s: + try: + f = float(s) + return str(int(f)) if f == int(f) else str(f) + except Exception: + pass + return s + +def leak_rule(skill, ans): + """逐字复刻 skill_config_probe.answer_leaked 的数字边界规则""" + if not skill: + return False + g = sanitize(ans) + if not g or not re.fullmatch(r'-?\d+(\.\d+)?', g): + return None # 不适用(非纯数字答案) + return bool(re.search(r'(? 0: + depth += (text[i] == '{') - (text[i] == '}') + i += 1 + if depth == 0: + last = text[m.end():i-1].strip() + return last + +def strict_correct(text, gold): + """严格口径: 最后一个 boxed 的内容去掉 latex 修饰后必须与 gold 精确相等(串或数值)。 + 不做 '从文本中捞第一个数字' 这类宽松回退。""" + raw = last_boxed(text) + if raw is None: + return False + s = raw.replace('\\!','').replace('\\,','').replace('\\ ',' ') + s = re.sub(r'\\text\s*\{([^}]*)\}', r'\1', s) + s = s.replace('$','').replace('{','').replace('}','').replace('\\','').strip() + g = str(gold).strip() + if s == g: + return True + try: + return float(s) == float(g) + except Exception: + return False + +def main(): + # 先取 200 题的答案表(placebo 用): data_id -> answer + answers = {} + with open(os.path.join(HERE, FILES['A_old_off'])) as f: + for line in f: + r = json.loads(line) + answers.setdefault(r['data_id'], r['reference_answer']) + dids = sorted(answers) + placebo = {} + for i, d in enumerate(dids): + # 找下一个"值不同"的答案做安慰剂 + for j in range(1, len(dids)): + cand = answers[dids[(i+j) % len(dids)]] + if sanitize(cand) != sanitize(answers[d]): + placebo[d] = cand + break + + for grp, path in FILES.items(): + # 聚合器: cfg -> 统计 + S = defaultdict(lambda: defaultdict(float)) + with open(os.path.join(HERE, path)) as f: + for line in f: + r = json.loads(line) + cfg = r['config']; s = S[cfg] + skill = r['skill'] or '' + corr = 1.0 if r['withskill_correct'] else 0.0 + hard = (r.get('baseline_pass') or 0) == 0 + s['n'] += 1 + # --- T4 严格判分 --- + sc = 1.0 if strict_correct(r.get('withskill_text',''), r['reference_answer']) else 0.0 + s['acc_loose'] += corr; s['acc_strict'] += sc + s['loose_only'] += 1.0 if (corr and not sc) else 0.0 + # --- T1 安慰剂 --- + lk = leak_rule(skill, r['reference_answer']) + if lk is not None and r['parseable']: + s['n_lk'] += 1 + s['leak'] += 1.0 if lk else 0.0 + pl = leak_rule(skill, placebo[r['data_id']]) + s['placebo'] += 1.0 if pl else 0.0 + # --- T2 条件分解 --- + if r['parseable']: + key = 'L' if r['leaked'] else 'Cn' + s[f'n_{key}'] += 1; s[f'acc_{key}'] += corr + if hard: + s[f'nh_{key}'] += 1; s[f'acch_{key}'] += corr + else: + s['n_U'] += 1; s['acc_U'] += corr + if hard: + s['nh_U'] += 1; s['acch_U'] += corr + print(f"\n{'='*100}\n### {grp} ({path})") + print("%-14s %6s | %7s %7s %9s | %5s %6s | %5s %6s | %5s %6s | %8s %8s %9s" % ( + 'config','n','leak%','placebo%','净leak%','nL','accL','nCn','accCn','nU','accU','accLoose','accStrict','looseOnly%')) + for cfg in CFGS: + s = S[cfg] + n = s['n'] or 1 + nlk = s['n_lk'] or 1 + lk, pl = s['leak']/nlk, s['placebo']/nlk + aL = s['acc_L']/s['n_L'] if s['n_L'] else float('nan') + aC = s['acc_Cn']/s['n_Cn'] if s['n_Cn'] else float('nan') + aU = s['acc_U']/s['n_U'] if s['n_U'] else float('nan') + print("%-14s %6d | %7.3f %7.3f %9.3f | %5d %6.3f | %5d %6.3f | %5d %6.3f | %8.3f %8.3f %9.3f" % ( + cfg, s['n'], lk, pl, lk-pl, s['n_L'], aL, s['n_Cn'], aC, s['n_U'], aU, + s['acc_loose']/n, s['acc_strict']/n, s['loose_only']/n)) + # hard 子集(排除"泄漏样本挑了简单题") + print(" --- hard(baseline=0) 子集: leaked vs clean 的 acc ---") + for cfg in CFGS: + s = S[cfg] + ahL = s['acch_L']/s['nh_L'] if s['nh_L'] else float('nan') + ahC = s['acch_Cn']/s['nh_Cn'] if s['nh_Cn'] else float('nan') + ahU = s['acch_U']/s['nh_U'] if s['nh_U'] else float('nan') + print(" %-14s hard: leaked %4d/%.3f clean %4d/%.3f unparse %4d/%.3f" % ( + cfg, s['nh_L'], ahL, s['nh_Cn'], ahC, s['nh_U'], ahU)) + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/skill2lora/good_skill_hard_fail/reflexion_probe.py b/cookbook/exp/skill2lora/good_skill_hard_fail/reflexion_probe.py new file mode 100644 index 000000000..899776864 --- /dev/null +++ b/cookbook/exp/skill2lora/good_skill_hard_fail/reflexion_probe.py @@ -0,0 +1,504 @@ +#!/usr/bin/env python3 +"""reflexion_probe.py — reflexion 链路探针（纯推理，不训练）。 + +链路：executor 失败轨迹（复用 v2 eval 的 baseline 缓存原文） + → LLM rubric 诊断（qwen-plus，7 条 _MATH_RUBRIC，[PASS]/[FAIL]+fix 文本，防泄漏硬规则） + → skillmodel(Qwen3-4B, T=0.5) 条件于 (题目 + 失败轨迹 + 诊断) 生成 + → executor(Qwen3-4B, T=0, v2 eval 逐字口径) 带 skill 重试 +目的：在"有一次真实失败 + rubric 诊断"的条件下，比较 skillmodel 的 prompt 写法 × +thinking on/off 哪种救活率最高，为 buffer B 如何用 rubric 经验提供依据。 + +题目：eval 200 题中 baseline=0 的失败题抽 N 道（seed 固定），全部是"裸解必错"题， +因此 executor 重试的 acc 即救活率。诊断按 data_id 缓存（rubric_diag_cache.jsonl）， +on/off 两次运行复用，不重复调 API。 + +prompt 变体（信息量递增，用于分离各级情报的增量价值）： + R4_blind 题目（无失败、无诊断；= 上轮 P5_pitfall 原文，跨轮对照锚点） + R0_trace_only 题目 + 失败轨迹尾部（无诊断；消融 rubric 的增量） + R1_needle 题目 + 失败 + 诊断 → 纠错针（WARNING/INSTEAD + 纪律后缀） + R2_narrative 题目 + 失败 + 诊断 → 训练现版叙述式（对照 v2 regen 路径） + R3_toy_fix 题目 + 失败 + 诊断 → 针对诊断错误点的玩具题示范 + +用法（~/.env 需含 LLM_BACKUP_BASE_URL / LLM_BACKUP_API_KEY，脚本自动加载）： + EXEC_GPUS=2 SKILL_GPUS=2 python3 reflexion_probe.py --skill-thinking off + EXEC_GPUS=2 SKILL_GPUS=2 python3 reflexion_probe.py --skill-thinking on +输出：reflexion_{tag}.jsonl（含 skill-gen 全文与 executor 全文）+ stdout 汇总。 +""" +import argparse +import copy +import hashlib +import json +import os +import random +import re +import statistics as st +import sys +from concurrent.futures import ThreadPoolExecutor +from typing import Dict, List, Optional + +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) + + +def _load_home_env(): + """加载 ~/.env（KEY=VALUE 简单格式；不覆盖已存在的环境变量）。""" + p = os.path.expanduser('~/.env') + if not os.path.exists(p): + return + for line in open(p): + line = line.strip() + if not line or line.startswith('#') or '=' not in line: + continue + k, v = line.split('=', 1) + os.environ.setdefault(k.strip(), v.strip().strip('"').strip("'")) + + +_load_home_env() + +import twinkle # noqa: E402 +from twinkle import DeviceGroup, DeviceMesh # noqa: E402 +from twinkle.data_format import SamplingParams # noqa: E402 +from twinkle.sampler import vLLMSampler # noqa: E402 +from twinkle.template import Template # noqa: E402 + +MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') +GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.8)) +EXEC_GPUS = int(os.environ.get('EXEC_GPUS', 2)) +SKILL_GPUS = int(os.environ.get('SKILL_GPUS', 2)) +MAX_MODEL_LEN = int(os.environ.get('MAX_MODEL_LEN', 16384)) +DIAG_MODEL = os.environ.get('LLM_BACKUP_MODEL', 'qwen-plus') +EVAL_RECORDS = os.path.join(SCRIPT_DIR, '..', 'skill_v2', 'eval_records.jsonl') +BASE_CACHE = os.path.join(SCRIPT_DIR, '..', 'skill_v2', 'cache', 'eval_baseline.jsonl') +DIAG_CACHE = os.path.join(SCRIPT_DIR, 'rubric_diag_cache.jsonl') + +# =========================================================================== +# executor 侧（逐字复刻 v2 eval，与 skill_config_probe.py 相同） +# =========================================================================== +_ANSWER_FORMAT_V2 = ('Present your reasoning, then put ONLY the final numeric result inside ' + '\\boxed{}. For example: \\boxed{42}.') + + +def build_direct_prompt(problem): + content = f'The problem you need to solve:\n{problem}\n\n' + _ANSWER_FORMAT_V2 + return {'messages': [{'role': 'user', 'content': content}]} + + +def build_skill_solve_prompt(problem, skill): + skill = (skill or '').strip() + if not skill: + return build_direct_prompt(problem) + content = (f'The problem you need to solve:\n{problem}\n\n' + 'Skill hint:\nFor this problem, a skill-generation model has analyzed it and ' + 'provided some advisory skills:\n' + f'{skill}\n' + 'Prefer using its techniques when they fit, but if you have a more efficient or ' + 'clearer correct method, you may use it. If you diverge from this advice, briefly ' + 'explain why. Be concise and accurate.\n' + + _ANSWER_FORMAT_V2) + return {'messages': [{'role': 'user', 'content': content}]} + + +_BOXED_RE = re.compile(r'\\boxed\s*\{') + + +def extract_boxed(text): + if not text: + return None + last = None + for m in _BOXED_RE.finditer(text): + depth, i = 1, m.end() + while i < len(text) and depth > 0: + depth += (text[i] == '{') - (text[i] == '}') + i += 1 + if depth == 0: + last = text[m.end():i - 1].strip() + return last + + +_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') + + +def _clean_text(d): + return _SPECIAL_TOKEN_RE.sub('', d or '').rstrip() + + +_SEAM_TAG_RE = re.compile(r'<\s*answer\s*>(.*?)<\s*/\s*answer\s*>', re.I | re.S) +_SEAM_BOX_RE = re.compile(r'(?:\\{1,2}\(|)\\{1,2}boxed\s*\{\s*([^}]*)\s*}(?:\)|)', re.S) +_SEAM_INLINE_RE = re.compile(r'\$([^$]+)\$|\\\(([^)]+)\\\)', re.S) +_SEAM_FRAC_RE = re.compile(r'(-?\d+(?:\.\d+)?)/(-?\d+(?:\.\d+)?)') +_SEAM_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') + + +def _seam_norm(num): + try: + f = float(num) + return str(int(f)) if f == int(f) else str(f) + except Exception: + return num.strip() + + +def _seam_sanitize(txt): + txt = (txt or '').strip() + if (m := _SEAM_TAG_RE.search(txt)): + txt = m.group(1).strip() + elif (m := _SEAM_BOX_RE.search(txt)): + txt = m.group(1).strip() + elif (m := _SEAM_INLINE_RE.search(txt)): + txt = (m.group(1) or m.group(2)).strip() + txt = re.sub(r'\\frac\s*\{\s*([^}]+?)\s*}\s*\{\s*([^}]+?)\s*}', r'\1/\2', txt) + if (m := _SEAM_FRAC_RE.search(txt)): + p, q = map(float, m.groups()) + if q: + return _seam_norm(str(p / q)) + if (m := _SEAM_NUM_RE.search(txt)): + return _seam_norm(m.group()) + return txt + + +def grade(seq, gold): + text = _clean_text(getattr(seq, 'decoded', '') or '') + raw = extract_boxed(text) + pred = _seam_sanitize(raw) if raw else None + return {'pred': pred, 'correct': bool(pred) and (pred == _seam_sanitize(str(gold))), + 'stop_reason': getattr(seq, 'stop_reason', None), + 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} + + +def extract_skill(text): + low = text.lower() + end_think = low.rfind('') + answer = text[end_think + len(''):] if end_think >= 0 else text + s = answer.lower().rfind('') + if s < 0: + return None + inner = s + len('') + e = answer.lower().find('', inner) + if e < 0: + return None + block = answer[inner:e].strip() + return re.sub(r'', '', block, flags=re.I).strip() or None + + +def answer_leaked(text, reference): + """双口径泄漏检测：A=裸数字任意位置；B 口径由分析端按 |gts|>=10 复算。""" + if not text: + return False + g = _seam_sanitize(str(reference)) + if not g or not re.fullmatch(r'-?\d+(\.\d+)?', g): + return False + return bool(re.search(r'(? +or +- [FAIL] : (fix: ) +Then a final line: Summary: <2-3 sentences naming the single most damaging error and the correct turn to take>. + +HARD RULES: never state, compute, or hint at the problem's final numeric answer or any final-stage numeric result; describe errors and directions only. Keep the whole output under 250 words.""" + +DIAG_USER = """## Problem +{problem} + +## Rubric +{rubric} + +## Failed attempt (may be truncated) +{segment} + +Now output the diagnostic lines.""" + + +def diagnose(client, problem, fail_text, gold): + rubric = '\n'.join(f'{i+1}. {t}' for i, t in enumerate(MATH_RUBRIC)) + seg = fail_text[-4000:] + msg = [{'role': 'system', 'content': DIAG_SYSTEM}, + {'role': 'user', 'content': DIAG_USER.format(problem=problem, rubric=rubric, segment=seg)}] + r = client.chat.completions.create(model=DIAG_MODEL, messages=msg, max_tokens=600, + temperature=0.2, timeout=120) + text = (r.choices[0].message.content or '').strip() + # 防泄漏兜底：诊断若带出 gts 数值，重试一次更严的指令；仍泄漏则截去含数值的行 + if answer_leaked(text, gold): + msg.append({'role': 'assistant', 'content': text}) + msg.append({'role': 'user', 'content': 'Your output contained a forbidden final numeric value. ' + 'Rewrite the SAME diagnosis with every final-stage number removed.'}) + r = client.chat.completions.create(model=DIAG_MODEL, messages=msg, max_tokens=600, + temperature=0.2, timeout=120) + text = (r.choices[0].message.content or '').strip() + if answer_leaked(text, gold): + g = _seam_sanitize(str(gold)) + text = '\n'.join(l for l in text.splitlines() + if not re.search(r'(? 输出） +# =========================================================================== +# R4：无情报锚点 = 上轮 P5_pitfall 原文（跨实验可比）。 +R4_BLIND_SYS = """\ +You are a skill-generation model. A separate executor model will solve the problem; it only sees your block. + +First think privately: solve the problem in your head AND identify the single most likely way a solver goes wrong on this type (a tempting but wrong turn, an off-by-one, a wasteful brute-force, a wrong branch). Then, inside , write under 90 words: +- WARNING: name that most likely mistake concretely and say why it is wrong. +- INSTEAD: one or two sentences pointing to the correct turn (technique name + where to apply it), without solving the problem or revealing any numeric result. +- End with: "Avoid re-checking loops; box a bare number as soon as it is computed." +""" + +# R0：只有失败轨迹（无诊断）——消融 rubric 的增量价值。 +R0_TRACE_SYS = """\ +You are a skill-generation model. A separate executor model previously FAILED this problem; you will see the tail of its failed attempt. The executor will retry seeing ONLY your block. + +First think privately: read the failed attempt, find where it went wrong, and decide the correct turn. Then, inside , write under 90 words: +- WARNING: the concrete mistake the previous attempt made (quote its wrong move briefly). +- INSTEAD: one or two sentences pointing to the correct turn, without solving the problem or revealing any numeric result. +- End with: "Avoid re-checking loops; box a bare number as soon as it is computed." +""" + +# R1：query + 诊断（无轨迹）→ 纠错针（把 rubric 的 FAIL/fix 转译成对 executor 的直接行为指令）。 +R1_NEEDLE_SYS = """\ +You are a skill-generation model. A separate executor model previously FAILED this problem. You will see an expert rubric diagnosis of that failure (you will NOT see the failed attempt itself). The executor will retry seeing ONLY your block. + +First think privately: from the diagnosis, pinpoint the decisive error. Then, inside , write under 90 words: +- WARNING: the decisive mistake (grounded in the diagnosis, stated concretely for THIS problem). +- INSTEAD: the corrective route distilled from the diagnosis's fix directions - technique name + where to apply it. Do not solve the problem; never state any numeric result. +- End with: "Avoid re-checking loops; box a bare number as soon as it is computed." +""" + +# R2：query + 诊断（无轨迹）→ 训练现版叙述式（对照 v2 _regen_prompt 的文体路径）。 +R2_NARR_SYS = """\ +You are a skill-generation model. A separate executor model previously FAILED this problem. You will see an expert rubric diagnosis of that failure (you will NOT see the failed attempt itself). The executor will retry seeing ONLY your block. + +First, think privately: work the problem out and understand why the attempt failed. Then write the block as ONE coherent analysis narrative (not a bullet list): name what the problem is essentially asking, walk through the recommended approach, and weave in - informed by the diagnosis - the specific pitfall that sank the previous attempt and how to avoid it. Do NOT solve the problem, do NOT reveal or compute the final answer, and do NOT substitute the problem's specific numbers into the steps. Keep it to roughly one focused paragraph. + +Put ONLY the methodology inside . +""" + +# R3：query + 诊断（无轨迹）→ 玩具题示范（针对诊断指出的错误技巧点造 toy，异数字防泄漏）。 +R3_TOYFIX_SYS = """\ +You are a skill-generation model. A separate executor model previously FAILED this problem. You will see an expert rubric diagnosis of that failure (you will NOT see the failed attempt itself). The executor will retry seeing ONLY your block. + +First think privately: from the diagnosis, identify the ONE technique the executor got wrong. Then, inside , do exactly this (under 110 words): +1. Invent a MINIATURE problem exercising that same technique with DIFFERENT, much smaller numbers, and solve the miniature completely in at most 5 short lines, making the correct move (the one the failed attempt missed) explicit. +2. One transfer sentence: "Your problem needs the same move where the previous attempt went wrong - apply it, then box a bare number." +Hard rules: never use any number from the original problem; never state its answer. +""" + +PROMPTS = { + 'R4_blind': ('none', R4_BLIND_SYS), + 'D1_needle': ('diag', R1_NEEDLE_SYS), + 'D2_narr': ('diag', R2_NARR_SYS), + 'D3_toyfix': ('diag', R3_TOYFIX_SYS), +} + + +def skillgen_prompt(name, problem, fail_tail, diag): + mode, sys_p = PROMPTS[name] + user = f'Problem:\n{problem}' + if mode in ('trace', 'trace+diag'): + user += f'\n\nFailed attempt (tail):\n{fail_tail}' + if mode in ('diag', 'trace+diag'): + user += f'\n\nExpert rubric diagnosis of the failure:\n{diag}' + return {'messages': [{'role': 'system', 'content': sys_p}, {'role': 'user', 'content': user}]} + + +# =========================================================================== +# 数据与主流程 +# =========================================================================== +def md5_key(problem): + return hashlib.md5('\x1f'.join([problem]).encode('utf-8')).hexdigest() + + +def load_fail_problems(n, seed): + probs = {} + for line in open(EVAL_RECORDS): + line = line.strip() + if not line: + continue + r = json.loads(line) + if r.get('record_type') != 'eval_problem' or r.get('chunk') != -1: + continue + if r['data_id'] not in probs and float(r.get('baseline_pass', 1)) == 0: + probs[r['data_id']] = {'data_id': r['data_id'], 'problem': r['problem'], + 'reference_answer': r['reference_answer']} + cache = {} + for line in open(BASE_CACHE): + try: + c = json.loads(line) + cache[c['key']] = c['value'] + except Exception: + continue + items = [] + for p in sorted(probs.values(), key=lambda x: x['data_id']): + v = cache.get(md5_key(p['problem'])) + if v and not v.get('correct') and (v.get('text') or '').strip(): + p['fail_text'] = v['text'] + items.append(p) + rng = random.Random(seed) + sample = rng.sample(items, min(n, len(items))) + sample.sort(key=lambda x: x['data_id']) + print(f'[抽样] baseline 失败且有轨迹全文的题 {len(items)} -> 抽 {len(sample)}（seed={seed}）') + return sample + + +def run_batch(sampler, prompts, max_tokens, temperature, top_p, dp, num_samples=1): + if not prompts: + return [] + params = SamplingParams(max_tokens=max_tokens, temperature=temperature, top_p=top_p, num_samples=num_samples) + padded = prompts + if dp > 1 and 0 < len(prompts) < dp: + padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(dp - len(prompts))] + responses = sampler.sample(padded, params)[:len(prompts)] + return [list(r.sequences) if (r and r.sequences) else [] for r in responses] + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument('--skill-thinking', choices=('on', 'off'), required=True) + ap.add_argument('--prompts', default=','.join(PROMPTS.keys())) + ap.add_argument('--n-problems', type=int, default=130) + ap.add_argument('--n-rollouts', type=int, default=8, help='每题每思路的 skill 采样数（T>0）；executor 对每条 skill 各跑一次 greedy') + ap.add_argument('--seed', type=int, default=42) + ap.add_argument('--skill-temperature', type=float, default=0.5) + ap.add_argument('--skill-max-tokens', type=int, default=8192) + ap.add_argument('--max-tokens', type=int, default=8192) + ap.add_argument('--fail-tail-chars', type=int, default=1800) + ap.add_argument('--out', default=None) + args = ap.parse_args() + + names = [x for x in args.prompts.split(',') if x] + problems = load_fail_problems(args.n_problems, args.seed) + out_path = args.out or os.path.join(SCRIPT_DIR, f'reflexion_{args.skill_thinking}.jsonl') + + # ---- 阶段1：rubric 诊断（带磁盘缓存，8 线程并发）---- + diag_cache = {} + if os.path.exists(DIAG_CACHE): + for line in open(DIAG_CACHE): + try: + c = json.loads(line) + diag_cache[c['data_id']] = c['diag'] + except Exception: + continue + todo = [p for p in problems if p['data_id'] not in diag_cache] + if todo: + from openai import OpenAI + client = OpenAI(api_key=os.environ['LLM_BACKUP_API_KEY'], + base_url=os.environ['LLM_BACKUP_BASE_URL']) + print(f'[诊断] 需调 API {len(todo)} 题（model={DIAG_MODEL}），其余 {len(problems)-len(todo)} 题走缓存') + + def _one(p): + try: + return p['data_id'], diagnose(client, p['problem'], p['fail_text'], p['reference_answer']) + except Exception as e: + return p['data_id'], f'[DIAG_ERROR] {e}' + with ThreadPoolExecutor(max_workers=8) as ex: + with open(DIAG_CACHE, 'a') as f: + for did, diag in ex.map(_one, todo): + diag_cache[did] = diag + f.write(json.dumps({'data_id': did, 'diag': diag}, ensure_ascii=False) + '\n') + n_err = sum(1 for p in problems if str(diag_cache.get(p['data_id'], '')).startswith('[DIAG_ERROR]')) + print(f'[诊断] 完成，失败 {n_err} 题') + + # ---- 阶段2：skill-gen + executor ---- + think = args.skill_thinking == 'on' + twinkle.initialize(mode='ray', nproc_per_node=SKILL_GPUS + EXEC_GPUS, lazy_collect=False, groups=[ + DeviceGroup(name='skill', ranks=list(range(SKILL_GPUS)), device_type='GPU'), + DeviceGroup(name='exec', ranks=list(range(SKILL_GPUS, SKILL_GPUS + EXEC_GPUS)), device_type='GPU')]) + + def make_sampler(group, world, enable_thinking): + s = vLLMSampler(model_id=MODEL_ID, + engine_args={'gpu_memory_utilization': GPU_MEM, + 'max_model_len': MAX_MODEL_LEN, 'tensor_parallel_size': 1}, + device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), + remote_group=group) + s.set_template(Template, model_id=MODEL_ID, enable_thinking=enable_thinking, max_length=MAX_MODEL_LEN) + return s + + skill_sampler = make_sampler('skill', SKILL_GPUS, enable_thinking=think) + exec_sampler = make_sampler('exec', EXEC_GPUS, enable_thinking=True) + + sg_prompts, meta = [], [] + for name in names: + for p in problems: + tail = p['fail_text'][-args.fail_tail_chars:] + sg_prompts.append(skillgen_prompt(name, p['problem'], tail, diag_cache.get(p['data_id'], ''))) + meta.append((name, p)) + print(f'[skill-gen] {len(sg_prompts)} 条 x {args.n_rollouts} rollouts, thinking={args.skill_thinking}, T={args.skill_temperature}') + sg_out = run_batch(skill_sampler, sg_prompts, args.skill_max_tokens, + args.skill_temperature, 0.95, SKILL_GPUS, num_samples=args.n_rollouts) + + trials = [] + for (name, p), seqs in zip(meta, sg_out): + seqs = list(seqs or []) + for si in range(args.n_rollouts): + seq = seqs[si] if si < len(seqs) else None + full = _clean_text(getattr(seq, 'decoded', '') or '') if seq is not None else '' + sk = extract_skill(full) or '' + trials.append({'config': name, 'data_id': p['data_id'], 'sample_idx': si, + 'problem': p['problem'], + 'reference_answer': p['reference_answer'], + 'diag': diag_cache.get(p['data_id'], ''), + 'diag_leaked': answer_leaked(diag_cache.get(p['data_id'], ''), p['reference_answer']), + 'skill': sk, 'parseable': bool(sk), 'skill_chars': len(sk), + 'leaked': answer_leaked(sk, p['reference_answer']), + 'skillgen_full': full, + 'skillgen_stop': getattr(seq, 'stop_reason', None) if seq is not None else 'empty', + 'skillgen_tokens': len(getattr(seq, 'tokens', None) or []) if seq is not None else 0}) + + ex_prompts = [build_skill_solve_prompt(t['problem'], t['skill']) for t in trials] + print(f'[executor] {len(ex_prompts)} 条, T=0') + ex_out = run_batch(exec_sampler, ex_prompts, args.max_tokens, 0.0, 1.0, EXEC_GPUS) + + with open(out_path, 'w') as f: + for t, seqs in zip(trials, ex_out): + roll = grade(seqs[0], t['reference_answer']) if seqs else { + 'pred': None, 'correct': False, 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} + t.update({'withskill_pred': roll['pred'], 'withskill_correct': roll['correct'], + 'withskill_stop': roll['stop_reason'], 'withskill_tokens': roll['gen_tokens'], + 'withskill_text': roll['text'], 'skill_thinking': args.skill_thinking}) + f.write(json.dumps(t, ensure_ascii=False) + '\n') + + # ---- 汇总（全是 baseline=0 的题，acc 即救活率；另报题级 pass@k）---- + print('\n' + '=' * 112) + print('%-11s %-6s %-6s %-6s %-6s %-8s %-8s %-8s %-8s' % ( + 'config', 'n', 'parse', 'leakA', 'trunc', 'skill字符', '救活@1', '救活@k', '救活@1(gts>=10无泄漏)')) + print('-' * 112) + for name in names: + sub = [t for t in trials if t['config'] == name] + n = len(sub) + parse = sum(t['parseable'] for t in sub) / n + leak = sum(t['leaked'] for t in sub) / n + trunc = sum(1 for t in sub if t['withskill_stop'] == 'length') / n + chars = int(st.median([t['skill_chars'] for t in sub if t['parseable']] or [0])) + acc = sum(t['withskill_correct'] for t in sub) / n + byq = {} + for t in sub: + byq.setdefault(t['data_id'], []).append(t) + p_at_k = sum(1 for v in byq.values() if any(x['withskill_correct'] for x in v)) / len(byq) + clean = [t for t in sub if not t['leaked']] + cacc = sum(t['withskill_correct'] for t in clean) / max(1, len(clean)) + print('%-11s %-6d %-6.2f %-6.2f %-6.2f %-8d %-8.3f %-8.3f %-.3f(n=%d)' % ( + name, n, parse, leak, trunc, chars, acc, p_at_k, cacc, len(clean))) + print('-' * 112) + print(f'明细（含 skill-gen/executor 全文）已写入 {out_path}') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/skill2lora/good_skill_hard_fail/sample_probe.py b/cookbook/exp/skill2lora/good_skill_hard_fail/sample_probe.py new file mode 100644 index 000000000..c0892aa90 --- /dev/null +++ b/cookbook/exp/skill2lora/good_skill_hard_fail/sample_probe.py @@ -0,0 +1,136 @@ +#!/usr/bin/env python3 +# sample_probe.py — 抽样交叉验证，为统计结论提供可人工核对的具体证据（含 文件:行号:data_id）。 +# 5 项交叉验证: +# V1 A vs C 样本级逐字一致率 (环境影响的最硬证据; 同 key 对比 skill 与 withskill_correct) +# V2 leak 标记真伪 (随机抽 leaked=True/False 各若干, 核对 reference_answer 是否真出现在 skill) +# V3 think 泄漏机制 (B 组 leaked=True 样本, 展示 think 段算出答案->写进 skill) +# V4 parse 失败成因 (B 组 P7/parse=False 样本, 确认 skillgen 被 think 吃满预算而截断) +# V5 accClean 样本量 (打印各 config n_clean, 防止小样本误读高 accClean) +import json, os, re, random +from collections import defaultdict + +HERE = os.path.dirname(os.path.abspath(__file__)) +FILES = { + 'A': 'skillcfg_full_off.jsonl', + 'B': 'skillcfg_full_on.jsonl', + 'C': 'env_runs/vllm_0.23.0/skillcfg_full_off.jsonl', +} +random.seed(0) + + +def load_indexed(path): + """返回 {(config,data_id,sample_idx): (lineno, record)}""" + d = {} + with open(os.path.join(HERE, path)) as f: + for i, line in enumerate(f, 1): + line = line.strip() + if not line: + continue + r = json.loads(line) + d[(r['config'], r['data_id'], r['sample_idx'])] = (i, r) + return d + + +def leaked_check(skill, ref): + """独立复现 answer_leaked 的判定意图: 整数答案是否以数字边界出现在 skill""" + s = str(ref).strip() + if not re.fullmatch(r'-?\d+(\.\d+)?', s): + return None # 非纯数字答案, 泄漏判定本就不适用 + return re.search(r'(?skill 搬答案）") + shown = 0 + for k in B: + _,r = B[k] + if r['config']=='P1_narrative' and r['leaked'] and str(r['reference_answer']).lstrip('-').isdigit(): + ln,_ = B[k] + full = r['skillgen_full'] or '' + ans = str(r['reference_answer']) + think_end = full.lower().find('') + in_think = ans in full[:think_end] if think_end>0 else False + in_skill = ans in (r['skill'] or '') + print(f" {fb}:{ln} key={k} ref={ans} 答案在think段={in_think} 在skill={in_skill}") + idx = (r['skill'] or '').find(ans) + if idx>=0: + print(f" skill 命中片段: ...{(r['skill'])[max(0,idx-45):idx+len(ans)+25]!r}...") + shown += 1 + if shown>=3: break + + print("\n"+"="*90) + print("[V4] parse 失败成因（B 组 P7_minimal, parseable=False）") + cnt_len=0; cnt_noskill=0; shown=0 + for k in B: + _,r = B[k] + if r['config']!='P7_minimal' or r['parseable']: + continue + full = r['skillgen_full'] or '' + has_close = '' in full.lower() + if r['skillgen_stop']=='length': cnt_len+=1 + if not has_close: cnt_noskill+=1 + if shown<3: + ln,_ = B[k] + print(f" {fb}:{ln} key={k} stop={r['skillgen_stop']} sg_tokens={r['skillgen_tokens']} " + f"含={has_close} full尾50={full[-50:]!r}") + shown+=1 + total_pf = sum(1 for k in B if B[k][1]['config']=='P7_minimal' and not B[k][1]['parseable']) + print(f" P7 parse失败共 {total_pf}: 其中 stop=length {cnt_len} 无 {cnt_noskill}") + + print("\n"+"="*90) + print("[V5] accClean 样本量核对（leaked=0 且 parseable=1 的 n_clean, 防小样本误读）") + for grp,(D,fn) in {'A':(A,fa),'B':(B,fb),'C':(C,fc)}.items(): + by=defaultdict(lambda:[0,0]) + for k in D: + _,r=D[k] + if (not r['leaked']) and r['parseable']: + by[r['config']][0]+=1 + by[r['config']][1]+= 1 if r['withskill_correct'] else 0 + cells=" ".join(f"{c.split('_')[0]}={by[c][0]}({(by[c][1]/by[c][0] if by[c][0] else 0):.2f})" + for c in ['P1_narrative','P5_pitfall','P7_minimal']) + print(f" [{grp}] n_clean(acc): {cells}") + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/skill2lora/good_skill_hard_fail/skill_config_probe.py b/cookbook/exp/skill2lora/good_skill_hard_fail/skill_config_probe.py new file mode 100644 index 000000000..685475ce7 --- /dev/null +++ b/cookbook/exp/skill2lora/good_skill_hard_fail/skill_config_probe.py @@ -0,0 +1,419 @@ +#!/usr/bin/env python3 +"""skill_config_probe.py — Qwen3-4B skill 模型 × executor 的纯推理配置探针。 + +目的：不训练，只推理。用 Qwen3-4B 当 skill 生成模型（T=0.5），产出 喂给 +executor（T=0，与 train_skill_v2.py 的 v2 eval 逐字同口径），比较不同配置的效果： + - skill 模型 enable_thinking：on / off（由 --skill-thinking 指定，跑两次对比） + - skill 模型 system prompt：7 种变体（含训练现版对照、六思路混合模板、toy 类比等） + +题目：从 eval_records.jsonl(chunk=-1) 的 200 题按 baseline_pass 分层抽 50 题， +通过/失败比例与整集一致（难度配比同实际 eval），seed 固定保证跨配置可比。 +baseline 直接复用缓存的 baseline_pass（同模型同 greedy 口径，无需重跑）。 + +用法： + EXEC_GPUS=2 SKILL_GPUS=2 python3 skill_config_probe.py --skill-thinking off + EXEC_GPUS=2 SKILL_GPUS=2 python3 skill_config_probe.py --skill-thinking on + python3 skill_config_probe.py --skill-thinking off --prompts P2_combo,P3_toy # 只跑子集 +输出：skillcfg_{tag}.jsonl（逐 trial）+ stdout 汇总表。 +""" +import argparse +import copy +import json +import os +import random +import re +import statistics as st +import sys +from typing import Dict, List, Optional + +import twinkle +from twinkle import DeviceGroup, DeviceMesh +from twinkle.data_format import SamplingParams +from twinkle.sampler import vLLMSampler +from twinkle.template import Template + +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') +GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.8)) +EXEC_GPUS = int(os.environ.get('EXEC_GPUS', 2)) +SKILL_GPUS = int(os.environ.get('SKILL_GPUS', 2)) +MAX_MODEL_LEN = int(os.environ.get('MAX_MODEL_LEN', 16384)) +DEFAULT_EVAL_RECORDS = os.path.join(SCRIPT_DIR, '..', 'skill_v2', 'eval_records.jsonl') + +# =========================================================================== +# executor 侧 —— 逐字复刻 train_skill_v2.py v2 分支（与 eval 口径一致） +# =========================================================================== +_ANSWER_FORMAT_V2 = ('Present your reasoning, then put ONLY the final numeric result inside ' + '\\boxed{}. For example: \\boxed{42}.') + + +def build_direct_prompt(problem): + content = f'The problem you need to solve:\n{problem}\n\n' + _ANSWER_FORMAT_V2 + return {'messages': [{'role': 'user', 'content': content}]} + + +def build_skill_solve_prompt(problem, skill): + skill = (skill or '').strip() + if not skill: + return build_direct_prompt(problem) + content = (f'The problem you need to solve:\n{problem}\n\n' + 'Skill hint:\nFor this problem, a skill-generation model has analyzed it and ' + 'provided some advisory skills:\n' + f'{skill}\n' + 'Prefer using its techniques when they fit, but if you have a more efficient or ' + 'clearer correct method, you may use it. If you diverge from this advice, briefly ' + 'explain why. Be concise and accurate.\n' + + _ANSWER_FORMAT_V2) + return {'messages': [{'role': 'user', 'content': content}]} + + +_BOXED_RE = re.compile(r'\\boxed\s*\{') + + +def extract_boxed(text: str) -> Optional[str]: + if not text: + return None + last = None + for m in _BOXED_RE.finditer(text): + depth, i = 1, m.end() + while i < len(text) and depth > 0: + depth += (text[i] == '{') - (text[i] == '}') + i += 1 + if depth == 0: + last = text[m.end():i - 1].strip() + return last + + +_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') + + +def _clean_text(decoded: Optional[str]) -> str: + return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() + + +_SEAM_TAG_RE = re.compile(r'<\s*answer\s*>(.*?)<\s*/\s*answer\s*>', re.I | re.S) +_SEAM_BOX_RE = re.compile(r'(?:\\{1,2}\(|)\\{1,2}boxed\s*\{\s*([^}]*)\s*}(?:\)|)', re.S) +_SEAM_INLINE_RE = re.compile(r'\$([^$]+)\$|\\\(([^)]+)\\\)', re.S) +_SEAM_FRAC_RE = re.compile(r'(-?\d+(?:\.\d+)?)/(-?\d+(?:\.\d+)?)') +_SEAM_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') + + +def _seam_norm(num: str) -> str: + try: + f = float(num) + return str(int(f)) if f == int(f) else str(f) + except Exception: + return num.strip() + + +def _seam_sanitize(txt: str) -> str: + txt = (txt or '').strip() + if (m := _SEAM_TAG_RE.search(txt)): + txt = m.group(1).strip() + elif (m := _SEAM_BOX_RE.search(txt)): + txt = m.group(1).strip() + elif (m := _SEAM_INLINE_RE.search(txt)): + txt = (m.group(1) or m.group(2)).strip() + txt = re.sub(r'\\frac\s*\{\s*([^}]+?)\s*}\s*\{\s*([^}]+?)\s*}', r'\1/\2', txt) + if (m := _SEAM_FRAC_RE.search(txt)): + p, q = map(float, m.groups()) + if q: + return _seam_norm(str(p / q)) + if (m := _SEAM_NUM_RE.search(txt)): + return _seam_norm(m.group()) + return txt + + +def grade(seq, gold) -> Dict: + text = _clean_text(getattr(seq, 'decoded', '') or '') + raw = extract_boxed(text) + pred = _seam_sanitize(raw) if raw else None + correct = bool(pred) and (pred == _seam_sanitize(str(gold))) + return {'pred': pred, 'correct': correct, + 'stop_reason': getattr(seq, 'stop_reason', None), + 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} + + +# ---- skill 抽取（v2 泛化版：砍 think 后取最后一个 块）---- +def extract_skill(text: str) -> Optional[str]: + low = text.lower() + end_think = low.rfind('') + answer = text[end_think + len(''):] if end_think >= 0 else text + s = answer.lower().rfind('') + if s < 0: + return None + inner = s + len('') + e = answer.lower().find('', inner) + if e < 0: + return None + block = answer[inner:e].strip() + block = re.sub(r'', '', block, flags=re.IGNORECASE).strip() + return block or None + + +def answer_leaked(skill: str, reference) -> bool: + """诊断用：skill 文本中是否出现 gts 数值（digit-boundary，简版）。""" + if not skill: + return False + g = _seam_sanitize(str(reference)) + if not g or not re.fullmatch(r'-?\d+(\.\d+)?', g): + return False + return bool(re.search(r'(? 块便于解析） +# =========================================================================== +# P1：训练现版（方案1）SKILL_GEN_SYSTEM 逐字对照组。 +P1_NARRATIVE = """\ +You are a skill-generation model. Your block will be fed to a SEPARATE downstream executor model that must solve the problem on its own. The executor will NOT see your private reasoning — it only sees what is inside .... + +First, think privately: actually work the problem out in your head to make sure you understand it, then step back and abstract WHAT MAKES THIS TYPE OF PROBLEM SOLVABLE into transferable methodology. + +Then write the block following these rules: +- Give general, transferable solving techniques for this TYPE of problem: the key concepts/theorems it relies on, the recommended strategy and steps, and the common pitfalls to avoid — plus a brief reason for each piece of advice so the executor understands why. +- Write it as one coherent analysis narrative (not a bullet list): first name what the problem is essentially asking, then walk through how to approach it, blending concepts, steps, pitfalls and reasons into a single connected story. +- CRITICAL: Do NOT solve the problem for the executor. Do NOT reveal or compute the final answer, and do NOT substitute the problem's specific given numbers into the steps or state any intermediate numeric results. Leave ALL concrete numbers for the executor to compute on its own. If you catch yourself writing a specific number from the problem, replace it with a description of the quantity instead. +- Keep it concise: aim for roughly one focused paragraph. + +Put ONLY the methodology inside . +""" + +# P2：六思路混合模板——主体二选一（toy 类比 / 路线卡片）+ 永远加纪律后缀；限长。 +P2_COMBO = """\ +You are a skill-generation model. Your block is the ONLY thing a separate executor model will see; it must help the executor solve the problem quickly within a tight token budget. + +First think privately and solve the problem in your head. Then write a SHORT block (under 120 words) with exactly this structure: +1. MAIN PART - pick ONE of the two forms, whichever fits the problem better: + (a) Toy example: invent a tiny problem of the SAME type but with DIFFERENT, smaller numbers, solve the toy completely in 2-4 lines showing the key trick, then add one sentence: "Your problem has the same shape - apply the same steps to its numbers." + (b) Route card: name the problem type, then give the key formula / recurrence / lemma / reduction that cracks it (no derivation, no solving), and say what single quantity to compute. +2. LAST LINE - always end with exactly this discipline line: "Single pass: no re-deriving, no re-checking; once computed, box a bare number immediately." + +Never use the original problem's own numbers in the main part; never state or imply the final answer. + +Put everything inside . +""" + +# P3：纯 toy 类比——完整解一道异数字同型玩具题，靠示范迁移；天然 answer-free。 +P3_TOY = """\ +You are a skill-generation model. A separate executor model will solve the problem; it only sees your block. + +First think privately and identify the core technique this problem needs. Then, inside , do exactly one thing: invent a MINIATURE problem of the same type with DIFFERENT and much smaller numbers, and solve that miniature completely in at most 5 short lines, making the key trick explicit. Finish with one transfer sentence: "Your problem has the same shape - repeat these steps with its own numbers, then box a bare number." + +Hard rules: never mention or use any number that appears in the original problem; never state the original problem's answer; keep the whole block under 100 words. +""" + +# P4：路线卡片——极简结构化卡片（类型/公式/起点/目标），无叙述。 +P4_CARD = """\ +You are a skill-generation model. A separate executor model will solve the problem; it only sees your block. + +First think privately and find the standard route for this problem. Then output, inside , an ultra-compact ROUTE CARD with at most 4 lines: +TYPE: +KEY: +START: +COMPUTE: + +No derivations, no explanations, no solving, never state the final answer, never copy the problem's numbers into KEY. +""" + +# P5：预判纠错——无失败情报版“纠错针”：预判本题最可能的错误走向并拦截。 +P5_PITFALL = """\ +You are a skill-generation model. A separate executor model will solve the problem; it only sees your block. + +First think privately: solve the problem in your head AND identify the single most likely way a solver goes wrong on this type (a tempting but wrong turn, an off-by-one, a wasteful brute-force, a wrong branch). Then, inside , write under 90 words: +- WARNING: name that most likely mistake concretely and say why it is wrong. +- INSTEAD: one or two sentences pointing to the correct turn (technique name + where to apply it), without solving the problem or revealing any numeric result. +- End with: "Avoid re-checking loops; box a bare number as soon as it is computed." +""" + +# P6：SEAM 经验风格（英文，输出改为 统一解析）——概念/策略/易错三段式对照组。 +P6_SEAM = """\ +You are a problem-solving guidance model. Read the math problem and distill a concise, reusable piece of solving experience that will help a SEPARATE solver model reach the correct answer. +Rules: +- Do NOT solve the problem and do NOT reveal or compute the final answer. +- State the key concepts/theorems, the recommended strategy/steps, and the common pitfalls to avoid. +- Output ONLY the experience, wrapped EXACTLY as ... . +""" + +# P7：一句话下界对照——只准一句话点出关键恒等式/技巧。 +P7_MINIMAL = """\ +You are a skill-generation model. A separate executor model will solve the problem; it only sees your block. +Inside , write EXACTLY ONE sentence (max 30 words) naming the single key identity, theorem, or technique that cracks this problem. Nothing else. Never state the answer. +""" + +PROMPTS = { + 'P1_narrative': P1_NARRATIVE, + 'P2_combo': P2_COMBO, + 'P3_toy': P3_TOY, + 'P4_card': P4_CARD, + 'P5_pitfall': P5_PITFALL, + 'P6_seam': P6_SEAM, + 'P7_minimal': P7_MINIMAL, +} + + +def skillgen_prompt(system: str, problem: str) -> Dict: + # 与训练 _skillgen_prompt 同构：system + user('Problem:\n...') + return {'messages': [{'role': 'system', 'content': system}, + {'role': 'user', 'content': f'Problem:\n{problem}'}]} + + +# =========================================================================== +# 题目分层抽样：与 200 题集 baseline 通过率同配比 +# =========================================================================== +def load_problems(path, n, seed): + probs = {} + for line in open(path): + line = line.strip() + if not line: + continue + try: + r = json.loads(line) + except Exception: + continue + if r.get('record_type') != 'eval_problem' or r.get('chunk') != -1: + continue + did = r.get('data_id') + if did and did not in probs: + probs[did] = {'data_id': did, 'problem': r['problem'], + 'reference_answer': r['reference_answer'], + 'baseline_pass': float(r.get('baseline_pass', 0))} + items = sorted(probs.values(), key=lambda x: x['data_id']) + passed = [x for x in items if x['baseline_pass'] > 0] + failed = [x for x in items if x['baseline_pass'] == 0] + ratio = len(passed) / max(1, len(items)) + n_pass = round(n * ratio) + rng = random.Random(seed) + sample = rng.sample(passed, min(n_pass, len(passed))) + \ + rng.sample(failed, min(n - n_pass, len(failed))) + sample.sort(key=lambda x: x['data_id']) + print(f'[抽样] 全集 {len(items)} 题 baseline率 {ratio:.3f} -> 抽 {len(sample)} 题 ' + f'(pass {sum(1 for x in sample if x["baseline_pass"] > 0)} / fail ' + f'{sum(1 for x in sample if x["baseline_pass"] == 0)})，seed={seed}') + return sample + + +def run_batch(sampler, prompts, max_tokens, temperature, top_p, dp, num_samples=1): + if not prompts: + return [] + params = SamplingParams(max_tokens=max_tokens, temperature=temperature, + top_p=top_p, num_samples=num_samples) + padded = prompts + if dp > 1 and 0 < len(prompts) < dp: + padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(dp - len(prompts))] + responses = sampler.sample(padded, params)[:len(prompts)] + return [list(r.sequences) if (r and r.sequences) else [] for r in responses] + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument('--skill-thinking', choices=('on', 'off'), required=True) + ap.add_argument('--prompts', default=','.join(PROMPTS.keys())) + ap.add_argument('--n-problems', type=int, default=200) + ap.add_argument('--n-rollouts', type=int, default=8, help='每题每思路的 skill 采样数（T>0）；executor 对每条 skill 各跑一次 greedy') + ap.add_argument('--seed', type=int, default=42) + ap.add_argument('--skill-temperature', type=float, default=0.5) + ap.add_argument('--skill-top-p', type=float, default=0.95) + ap.add_argument('--skill-max-tokens', type=int, default=8192) + ap.add_argument('--max-tokens', type=int, default=8192, help='executor，对齐 eval') + ap.add_argument('--eval-records', default=DEFAULT_EVAL_RECORDS) + ap.add_argument('--out', default=None) + args = ap.parse_args() + + names = [x for x in args.prompts.split(',') if x] + for x in names: + if x not in PROMPTS: + print(f'未知 prompt: {x}，可选: {list(PROMPTS)}') + sys.exit(1) + problems = load_problems(args.eval_records, args.n_problems, args.seed) + think = args.skill_thinking == 'on' + out_path = args.out or os.path.join(SCRIPT_DIR, f'skillcfg_{args.skill_thinking}.jsonl') + + # 两个采样器：skill(thinking 可切) + exec(enable_thinking=True, T=0, 对齐 v2 eval 的 base_sampler) + twinkle.initialize(mode='ray', nproc_per_node=SKILL_GPUS + EXEC_GPUS, lazy_collect=False, groups=[ + DeviceGroup(name='skill', ranks=list(range(SKILL_GPUS)), device_type='GPU'), + DeviceGroup(name='exec', ranks=list(range(SKILL_GPUS, SKILL_GPUS + EXEC_GPUS)), device_type='GPU')]) + + def make_sampler(group, world, enable_thinking): + s = vLLMSampler(model_id=MODEL_ID, + engine_args={'gpu_memory_utilization': GPU_MEM, + 'max_model_len': MAX_MODEL_LEN, 'tensor_parallel_size': 1}, + device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), + remote_group=group) + s.set_template(Template, model_id=MODEL_ID, enable_thinking=enable_thinking, max_length=MAX_MODEL_LEN) + return s + + skill_sampler = make_sampler('skill', SKILL_GPUS, enable_thinking=think) + exec_sampler = make_sampler('exec', EXEC_GPUS, enable_thinking=True) + + # ---- 1) 所有配置的 skill-gen 一次性 batch ---- + sg_prompts, meta = [], [] + for name in names: + for p in problems: + sg_prompts.append(skillgen_prompt(PROMPTS[name], p['problem'])) + meta.append((name, p)) + print(f'[skill-gen] {len(sg_prompts)} 条 x {args.n_rollouts} rollouts (prompts={len(names)} x 题={len(problems)}), ' + f'thinking={args.skill_thinking}, T={args.skill_temperature}') + sg_out = run_batch(skill_sampler, sg_prompts, args.skill_max_tokens, + args.skill_temperature, args.skill_top_p, SKILL_GPUS, + num_samples=args.n_rollouts) + + trials = [] + for (name, p), seqs in zip(meta, sg_out): + seqs = list(seqs or []) + for si in range(args.n_rollouts): + seq = seqs[si] if si < len(seqs) else None + full = _clean_text(getattr(seq, 'decoded', '') or '') if seq is not None else '' + sk = extract_skill(full) or '' + trials.append({'config': name, 'data_id': p['data_id'], 'sample_idx': si, + 'problem': p['problem'], + 'reference_answer': p['reference_answer'], 'baseline_pass': p['baseline_pass'], + 'skill': sk, 'parseable': bool(sk), 'skill_chars': len(sk), + 'leaked': answer_leaked(sk, p['reference_answer']), + 'skillgen_full': full, + 'skillgen_stop': getattr(seq, 'stop_reason', None) if seq is not None else 'empty', + 'skillgen_tokens': len(getattr(seq, 'tokens', None) or []) if seq is not None else 0}) + + # ---- 2) 所有配置的 executor 一次性 batch（空 skill 走 direct，等价 baseline 口径）---- + ex_prompts = [build_skill_solve_prompt(t['problem'], t['skill']) for t in trials] + print(f'[executor] {len(ex_prompts)} 条, T=0, max_tokens={args.max_tokens}') + ex_out = run_batch(exec_sampler, ex_prompts, args.max_tokens, 0.0, 1.0, EXEC_GPUS) + + with open(out_path, 'w') as f: + for t, seqs in zip(trials, ex_out): + roll = grade(seqs[0], t['reference_answer']) if seqs else { + 'pred': None, 'correct': False, 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} + t.update({'withskill_pred': roll['pred'], 'withskill_correct': roll['correct'], + 'withskill_stop': roll['stop_reason'], 'withskill_tokens': roll['gen_tokens'], + 'withskill_text': roll['text'], + 'skill_thinking': args.skill_thinking}) + f.write(json.dumps(t, ensure_ascii=False) + '\n') + + # ---- 3) 汇总（含题级 pass@k）---- + print('\n' + '=' * 118) + print('%-14s %-5s %-6s %-6s %-6s %-8s %-7s %-6s %-8s %-8s %-10s' % ( + 'config', 'n', 'parse', 'leak', 'trunc', 'skill字符', 'mean@1', 'base', 'lift', 'pass@k', 'hard救活@k')) + print('-' * 118) + for name in names: + sub = [t for t in trials if t['config'] == name] + n = len(sub) + parse = sum(t['parseable'] for t in sub) / n + leak = sum(t['leaked'] for t in sub) / n + trunc = sum(1 for t in sub if t['withskill_stop'] == 'length') / n + chars = int(st.median([t['skill_chars'] for t in sub if t['parseable']] or [0])) + acc = sum(t['withskill_correct'] for t in sub) / n + base = sum(t['baseline_pass'] for t in sub) / n + byq = {} + for t in sub: + byq.setdefault(t['data_id'], []).append(t) + p_at_k = sum(1 for v in byq.values() if any(x['withskill_correct'] for x in v)) / len(byq) + hardq = {d: v for d, v in byq.items() if v[0]['baseline_pass'] == 0} + rescue_k = (sum(1 for v in hardq.values() if any(x['withskill_correct'] for x in v)) / len(hardq)) if hardq else 0.0 + print('%-14s %-5d %-6.2f %-6.2f %-6.2f %-8d %-7.3f %-6.3f %+-8.3f %-8.3f %-10.3f' % ( + name, n, parse, leak, trunc, chars, acc, base, acc - base, p_at_k, rescue_k)) + print('-' * 118) + print(f'逐 trial 明细（含 skill-gen/executor 全文）已写入 {out_path}') + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/skill2lora/run_ablate12.sh b/cookbook/exp/skill2lora/run_ablate12.sh index 8a72cde22..a90fb9173 100644 --- a/cookbook/exp/skill2lora/run_ablate12.sh +++ b/cookbook/exp/skill2lora/run_ablate12.sh @@ -18,12 +18,25 @@ # Env knobs (all optional): # DEEPMATH_DIR=$HERE/../../../deepmath_103k TRAIN_N=5000 MAX_UPDATES=50 EVAL_EVERY=5 # LR=1e-6 RUN_SFT=1 FORCE=1 ONLY="E5 E6" SLEEP=30 SWANLAB_PROJECT=twinkle +# MIN_LEVEL=6 CHUNK_SIZE=32 (gradient-signal fix: E1/E5 audit — level<=5 all-pass +# dominated, 16-problem chunks leave only ~6 mixed groups per update; eval split unaffected) # ============================================================================== set -euo pipefail +# avoid backward-pass OOM from allocator fragmentation (E6 crash: 15GiB reserved-unallocated); +# inherited by the Ray training actors via twinkle's runtime env passthrough +export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" + HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" cd "$HERE" +# central env file (optional): put all knobs in one place. ENV_FILE=xxx overrides the path. +ENV_FILE="${ENV_FILE:-$HERE/ablate12.env}" +if [ -f "$ENV_FILE" ]; then + echo "[ablate12] loading env from $ENV_FILE" + set -a; . "$ENV_FILE"; set +a +fi + OUT_ROOT="${OUT_ROOT:-$HERE/output.ablate12}" # DeepMath-103K (difficulty-stratified loader in skill_ablate/data.py); replaces the old # SEAM/aops input — see skill_quality_analysis.md 组成漂移修正. @@ -33,6 +46,8 @@ EVAL_SIZE="${EVAL_SIZE:-128}" MAX_UPDATES="${MAX_UPDATES:-50}" EVAL_EVERY="${EVAL_EVERY:-5}" LR="${LR:-1e-6}" +MIN_LEVEL="${MIN_LEVEL:-6}" +CHUNK_SIZE="${CHUNK_SIZE:-32}" SLEEP="${SLEEP:-30}" SWANLAB_PROJECT="${SWANLAB_PROJECT:-twinkle}" RUN_SFT="${RUN_SFT:-0}" @@ -93,6 +108,8 @@ while IFS=$'\t' read -r NAME EXP_DIR THINK SMT OPTIONAL; do --skill-max-tokens "$SMT" \ --max-updates "$MAX_UPDATES" \ --eval-every-updates "$EVAL_EVERY" \ + --min-level "$MIN_LEVEL" \ + --chunk-size "$CHUNK_SIZE" \ --lr "$LR" \ --swanlab-project "$SWANLAB_PROJECT" \ $FORCE_FLAG \ diff --git a/cookbook/exp/skill2lora/skill_ablate/__init__.py b/cookbook/exp/skill2lora/skill_ablate/__init__.py new file mode 100644 index 000000000..374d0aeef --- /dev/null +++ b/cookbook/exp/skill2lora/skill_ablate/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Skill-generation ablation package (view A / view B × think × style × training method). + +Design: reuse train_skill_v2.py primitives verbatim (import, never edit); this package only +adds the experiment matrix, the sample pool, the rubric double-cache, and the pluggable +training methods on top. See cookbook/exp/skill2lora/skill_quality_analysis.md sections +"AI 最终清单" and "AI 接口方案" for the frozen design decisions this code implements. +""" diff --git a/cookbook/exp/skill2lora/skill_ablate/config.py b/cookbook/exp/skill2lora/skill_ablate/config.py new file mode 100644 index 000000000..080697ba7 --- /dev/null +++ b/cookbook/exp/skill2lora/skill_ablate/config.py @@ -0,0 +1,147 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Declarative ablation matrix (E1-E12) + run order. + +Frozen decisions (skill_quality_analysis.md): +- 12 experiments; run nothink before think, the SFT method (E12) LAST and manually gated. +- Unified knobs: executor frozen T=0; skill-model train T=1.0 × 8 rollouts; eval T=0.5 × 4 + rollouts (query-only, no rubric); skill-max-tokens = 8192 (think) / 4096 (nothink). +- view B = query-only; view A = rubric line, eval still query-only (knowledge-transfer probe). + +This module is intentionally dependency-free (pure stdlib) so it can be imported and unit- +smoke-tested without torch / twinkle / a GPU. +""" +from dataclasses import dataclass +from typing import Dict, List + +# --- training methods (internal keys) -------------------------------------------------- +# bnpo view B query-only GRPO/BNPO main loop (reuses v2 process_chunk verbatim). +# rl_ab view A RL, AB split: first bare-problem greedy solve; WRONG problems -> A line +# (skill sampled under query+rubric), RIGHT problems -> B line (query-only); +# both go through executor greedy -> reward -> in-group BNPO, trained together. +# rl_err view A RL, error-only: same as rl_ab but the B line is NOT trained +# (single-variable contrast vs rl_ab on "does training the right-answer B line help"). +# opsd view A On-Policy Self-Distillation: student(query-only) logps pulled toward +# teacher(query+rubric) logps per token (loss='opsd'); error problems only. +# improve_sft view A improve-skill + SFT: first-pass 1 skill; correct -> positive pool +# (no leak, <=4096 chars); wrong -> rubric regen (2-in-8 pick 1) -> negative pool; +# 1:1 accumulate -> SFT. +# sft view A plain SFT: bare-problem wrong -> rubric -> regen (2-in-8) -> accumulate SFT. +METHODS = ('bnpo', 'rl_ab', 'rl_err', 'opsd', 'improve_sft', 'sft') +VIEW_OF_METHOD = {'bnpo': 'B', 'rl_ab': 'A', 'rl_err': 'A', + 'opsd': 'A', 'improve_sft': 'A', 'sft': 'A'} +STYLES = ('narrative', 'pitfall') +THINKINGS = ('on', 'off') + + +@dataclass(frozen=True) +class ExpSpec: + name: str # E1..E13 + method: str # one of METHODS + thinking: str # 'on' | 'off' + style: str # 'narrative' | 'pitfall' (ignored by align='seam': SEAM prompts bypass style) + optional: bool = False # E12(sft): manually gated (RUN_SFT=1), runs last + align: str = 'v2' # 'v2' | 'seam' — sets v2._ALIGN_MODE (prompt/判分/executor 嵌套全开关) + + @property + def view(self) -> str: + return VIEW_OF_METHOD[self.method] + + @property + def needs_rubric(self) -> bool: + return self.view == 'A' + + @property + def skill_max_tokens(self) -> int: + # think must have room for + (4096 truncates to an empty block). + # seam align: 人工拍板用 8192（不复刻 SEAM 原版 4096：think 模式下 4096 会把大量候选截断在 + # 里、压低 parseable，与“think 模式 skill-max-tokens 必须 8192”的矩阵规范保持一致）。 + if self.align == 'seam': + return 8192 + return 8192 if self.thinking == 'on' else 4096 + + @property + def loss(self) -> str: + return 'opsd' if self.method == 'opsd' else 'bnpo' + + @property + def exp_dir(self) -> str: + # output.ablate12/E{n}_{method}_{think}_{style}/ (seam align 加后缀区分) + suffix = '_seam' if self.align == 'seam' else '' + return f'{self.name}_{self.method}_{self.thinking}_{self.style}{suffix}' + + @property + def swanlab_exp(self) -> str: + return f'ablate12_{self.exp_dir}' + + +# --- the 12-experiment matrix (declarative; order field below drives execution) -------- +MATRIX: List[ExpSpec] = [ + # group 1 — view B BNPO: think × style, no-rubric baseline + ExpSpec('E1', 'bnpo', 'off', 'pitfall'), + ExpSpec('E2', 'bnpo', 'off', 'narrative'), + ExpSpec('E3', 'bnpo', 'on', 'pitfall'), + ExpSpec('E4', 'bnpo', 'on', 'narrative'), + # group 2 — view A RL-AB-mix: same grid as E1-E4, isolates "rubric rescues zero-grad groups" + ExpSpec('E5', 'rl_ab', 'off', 'pitfall'), + ExpSpec('E6', 'rl_ab', 'off', 'narrative'), + ExpSpec('E7', 'rl_ab', 'on', 'pitfall'), + ExpSpec('E8', 'rl_ab', 'on', 'narrative'), + # group 3 — view A training-method comparison (fixed think+narrative), sft last & optional + ExpSpec('E9', 'rl_err', 'on', 'narrative'), + ExpSpec('E10', 'opsd', 'on', 'narrative'), + ExpSpec('E11', 'improve_sft', 'on', 'narrative'), + ExpSpec('E12', 'sft', 'on', 'narrative', optional=True), + # group 4 — SEAM-align ablation: same data pipeline as the rest of the matrix, but ALL + # prompt/parsing/executor-nesting rules follow SEAM (align='seam' -> v2._ALIGN_MODE): + # actor uses SEAM EXPERIENCE_PROMPT (), executor sees the nested + # prompt_text+response_text(+think), lpem-parity greedy scoring, actor budget 4096. + # Query-only BNPO main loop (= SEAM's training form); eval stays the matrix-unified + # query-only readout so E13 is directly comparable with E1-E12. + ExpSpec('E13', 'bnpo', 'on', 'narrative', align='seam'), +] + +# execution order: all nothink first, then think; E13 (seam-align baseline) right after E6; +# the data-hungry SFT method dead last. +RUN_ORDER: List[str] = ['E1', 'E2', 'E5', 'E6', 'E13', 'E3', 'E7', 'E8', 'E9', 'E10', 'E11', 'E4', 'E12'] + +BY_NAME: Dict[str, ExpSpec] = {e.name: e for e in MATRIX} + + +def get_spec(name: str) -> ExpSpec: + key = name.strip().upper() + if key not in BY_NAME: + raise KeyError(f'unknown experiment {name!r}; valid: {sorted(BY_NAME)}') + return BY_NAME[key] + + +def ordered_specs(include_optional: bool = True) -> List[ExpSpec]: + specs = [BY_NAME[n] for n in RUN_ORDER] + return specs if include_optional else [s for s in specs if not s.optional] + + +def _self_check() -> None: + """Invariants that guard against typos when editing the matrix.""" + assert set(BY_NAME) == set(RUN_ORDER), 'RUN_ORDER must cover every matrix entry exactly once' + assert len(RUN_ORDER) == len(set(RUN_ORDER)) == len(MATRIX), 'duplicate / missing names' + for e in MATRIX: + assert e.method in METHODS, f'{e.name}: bad method {e.method}' + assert e.thinking in THINKINGS and e.style in STYLES, f'{e.name}: bad think/style' + assert e.align in ('v2', 'seam'), f'{e.name}: bad align {e.align}' + # nothink-before-think ordering within contiguous runs is a soft convention, not asserted. + + +if __name__ == '__main__': + import sys + _self_check() + if '--plan' in sys.argv: + # machine-readable run plan for the launcher: nameexp_dirthinksmtoptional + for e in ordered_specs(): + print(f'{e.name}\t{e.exp_dir}\t{e.thinking}\t{e.skill_max_tokens}\t{int(e.optional)}') + sys.exit(0) + print(f'{len(MATRIX)} experiments; run order: {" -> ".join(RUN_ORDER)}') + hdr = f'{"name":<4} {"view":<4} {"method":<12} {"think":<6} {"style":<10} {"align":<5} {"smt":<5} {"loss":<5} opt' + print(hdr) + print('-' * len(hdr)) + for e in ordered_specs(): + print(f'{e.name:<4} {e.view:<4} {e.method:<12} {e.thinking:<6} {e.style:<10} ' + f'{e.align:<5} {e.skill_max_tokens:<5} {e.loss:<5} {"Y" if e.optional else ""}') diff --git a/cookbook/exp/skill2lora/skill_ablate/data.py b/cookbook/exp/skill2lora/skill_ablate/data.py new file mode 100644 index 000000000..ac84c6d03 --- /dev/null +++ b/cookbook/exp/skill2lora/skill_ablate/data.py @@ -0,0 +1,96 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""DeepMath-103K loader with difficulty-stratified train/eval split. + +Dataset: AI-ModelScope/DeepMath-103K (columns: question / final_answer / difficulty / topic / +r1_solution_1..3). We keep only rows whose final_answer normalizes to a number via v2 +``_numeric_value`` (the \\boxed{} judging pipeline is numeric-exact; answers like ``\\phi^4`` +cannot be scored and are dropped). + +Stratified split (skill_quality_analysis.md 组成漂移修正): difficulty is bucketed to its +rounded integer level; ``eval_size`` problems are sampled with per-bucket quotas proportional +to the pool (largest-remainder rounding), the rest form the train pool — so train and eval +difficulty proportions match by construction. All sampling is seeded and file-order stable: +``data_id = dm::`` is reproducible across runs/experiments. + +Train-only difficulty floor (``--min-level``): E1/E5 gradient audit showed level<=5 groups are +dominated by all-pass (level 3: 63-74% all-pass, corr(level, mixed_rate)=0.92), i.e. mostly +zero-gradient. The floor drops those rows from the *train pool only*; the eval split keeps the +full-level stratification so eval/baseline stay comparable across experiments. +""" +import glob +import os +from collections import defaultdict +from typing import Any, Dict, List, Tuple + +import numpy as np + +import train_skill_v2 as v2 + + +def _read_rows(deepmath_dir: str) -> List[Dict[str, Any]]: + import pyarrow.parquet as pq + paths = sorted(glob.glob(os.path.join(deepmath_dir, '**', '*.parquet'), recursive=True)) + if not paths: + raise FileNotFoundError(f'no parquet files under --deepmath-dir {deepmath_dir}') + rows: List[Dict[str, Any]] = [] + for p in paths: + t = pq.read_table(p, columns=['question', 'final_answer', 'difficulty']) + rows.extend(t.to_pylist()) + return rows + + +def load_deepmath_records(args) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + """-> (train_records, eval_records), each record {'data_id','problem','reference_answer'}.""" + rows = _read_rows(args.deepmath_dir) + pool: List[Dict[str, Any]] = [] + for i, r in enumerate(rows): # global row index over sorted files = stable id + problem = (r.get('question') or '').strip() + num = v2._numeric_value(r.get('final_answer')) + if not problem or num is None: + continue + lvl = int(round(float(r.get('difficulty') or 0))) + pool.append({'data_id': f'dm:{lvl}:{i}', 'problem': problem, + 'reference_answer': num, '_level': lvl}) + + # bucket by level, seeded shuffle inside each bucket + buckets: Dict[int, List[Dict[str, Any]]] = defaultdict(list) + for rec in pool: + buckets[rec['_level']].append(rec) + rng = np.random.RandomState(args.seed) + for lvl in sorted(buckets): + rng.shuffle(buckets[lvl]) + + # eval quota per bucket: proportional, largest-remainder rounding + eval_n = min(args.eval_size, len(pool)) if args.eval_size > 0 else 0 + quota = {lvl: eval_n * len(b) / len(pool) for lvl, b in buckets.items()} + take = {lvl: int(q) for lvl, q in quota.items()} + for lvl in sorted(quota, key=lambda x: quota[x] - int(quota[x]), reverse=True): + if sum(take.values()) >= eval_n: + break + take[lvl] += 1 + + eval_records, train_records = [], [] + for lvl in sorted(buckets): + b = buckets[lvl] + eval_records.extend(b[:take[lvl]]) + train_records.extend(b[take[lvl]:]) + min_level = int(getattr(args, 'min_level', 0) or 0) + if min_level > 0: # train-only floor; eval keeps full-level mix (see module docstring) + n_before = len(train_records) + train_records = [r for r in train_records if r['_level'] >= min_level] + v2.logger.info(f'[data] min_level={min_level}: train pool {n_before} -> {len(train_records)}') + rng.shuffle(train_records) # ProblemPool reshuffles too; this decorrelates level runs + if args.n > 0: # optional stratified-in-expectation downsample (pool already shuffled) + train_records = train_records[:args.n] + + def _lvls(rs): + c = defaultdict(int) + for r in rs: + c[r['_level']] += 1 + return {k: round(v / len(rs), 3) for k, v in sorted(c.items())} + v2.logger.info(f'[data] DeepMath: pool={len(pool)} (numeric-only of {len(rows)}) ' + f'train={len(train_records)} eval={len(eval_records)}') + v2.logger.info(f'[data] level mix train={_lvls(train_records)} eval={_lvls(eval_records)}') + for r in eval_records + train_records: + r.pop('_level', None) + return train_records, eval_records diff --git a/cookbook/exp/skill2lora/skill_ablate/main.py b/cookbook/exp/skill2lora/skill_ablate/main.py new file mode 100644 index 000000000..6035d2846 --- /dev/null +++ b/cookbook/exp/skill2lora/skill_ablate/main.py @@ -0,0 +1,155 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Entry point: run one ablation experiment by name (E1..E12) or explicit knobs. + +Usage: + python -m skill_ablate.main --exp E5 --seam-parquet-dir /root/data/seam \ + --output-dir output.ablate12/E5_rl_ab_off_pitfall + +Defaults mirror train_skill_v2._build_args so reused v2 primitives behave identically; only +the ablation-specific knobs are added (--exp / --max-updates / --eval-every-updates / +--improve-skill-temperature / --skill-char-limit / --pool-max / --rubric-global-dir). +""" +import argparse +import sys + +import train_skill_v2 as v2 + +from .config import METHODS, STYLES, THINKINGS, get_spec, ExpSpec +from .trainer import run_experiment + + +def _build_args(argv=None): + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + # --- experiment selection: either --exp E5, or explicit --method/--thinking/--style --- + p.add_argument('--exp', default='', help='experiment name E1..E12 (fills method/think/style)') + p.add_argument('--method', choices=METHODS, default=None) + p.add_argument('--thinking', choices=THINKINGS, default=None) + p.add_argument('--style', choices=STYLES, default=None) + + # --- data (mirror v2) --- + p.add_argument('--dataset', choices=('aops', 'math'), default='aops') + p.add_argument('--n', type=int, default=0) + p.add_argument('--exclude-data-ids', default='') + p.add_argument('--seed', type=int, default=42) + p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) + p.add_argument('--eval-size', type=int, default=128) + p.add_argument('--seam-parquet-dir', type=str, default='') + p.add_argument('--deepmath-dir', type=str, default='', + help='DeepMath-103K parquet dir; when set, overrides --seam-parquet-dir/--dataset ' + 'and uses the difficulty-stratified split (eval/train same level mix).') + p.add_argument('--min-level', type=int, default=0, + help='train-only difficulty floor for DeepMath (eval keeps full-level mix). ' + '0 = off. E1/E5 audit: level<=5 is all-pass dominated (zero gradient); ' + 'recommended 6.') + + # --- eval口径 (4 rollouts × T=0.5, 与旧臂 E1-E13 同口径, 2026-07-28 拍板回退) --- + # 曾短暂改为 SEAM val 口径(1×greedy)，为保持与已完成 6 臂可横比而回退；需要 SEAM 口径时 + # 显式传 --eval-rollouts 1 --eval-skill-temperature 0.0。 + p.add_argument('--eval-rollouts', type=int, default=4) + p.add_argument('--eval-skill-temperature', type=float, default=0.5) + + # --- skill-gen / rollout (mirror v2) --- + p.add_argument('--chunk-size', type=int, default=16) + p.add_argument('--n-skills', type=int, default=8) + p.add_argument('--skill-gen-temperature', type=float, default=1.0) + p.add_argument('--skill-gen-top-p', type=float, default=1.0) + p.add_argument('--skill-gen-top-k', type=int, default=-1) + p.add_argument('--max-model-len', type=int, default=16384) + p.add_argument('--max-tokens', type=int, default=8192) + p.add_argument('--skill-max-tokens', type=int, default=None, + help='default per-experiment: 8192 (think) / 4096 (nothink); an explicit ' + 'value here wins over the experiment default.') + p.add_argument('--align-mode', choices=('v2', 'seam'), default='v2') + p.add_argument('--len-budget', type=int, default=None, + help='regen skill length target (chars). Default per style (ablation stats ' + '#9): narrative~1100 / pitfall~300. Used to pick the regen survivor.') + + # --- rubric / regen / distill (mirror v2 + new) --- + p.add_argument('--passatk-k', type=int, default=8) + p.add_argument('--passatk-skill-temp', type=float, default=1.0) + p.add_argument('--passatk-skill-top-p', type=float, default=1.0) + p.add_argument('--passatk-m', type=int, default=2) + p.add_argument('--rubric-workers', type=int, default=16) + p.add_argument('--improve-skill-temperature', type=float, default=0.5, + help='temperature for the single first-pass skill in opsd / improve_sft.') + p.add_argument('--skill-char-limit', type=int, default=4096, + help='hard char cap for SFT-seed skills (skill_quality_analysis.md #15-1/#18).') + + # --- GRPO / optim (mirror v2) --- + p.add_argument('--sft-batch-size', type=int, default=16, + help='batch = TRAIN_DP multiple; also the SFT-pool draw size.') + p.add_argument('--train-micro-batch', type=int, default=0, + help='forward_backward micro size; 0 = follow --sft-batch-size. think/8192 ' + 'experiments auto-halve to 8 (fp32 master + 8k-token logits OOM guard); ' + 'gradient is micro-normalized so this is mathematically equivalent.') + p.add_argument('--ppo-mini-batch-size', type=int, default=0) + p.add_argument('--grpo-epsilon', type=float, default=0.2) + p.add_argument('--adv-clip', type=float, default=0.0) + p.add_argument('--kl-beta', type=float, default=0.001) + p.add_argument('--lr', type=float, default=1e-6, + help='stable 1e-6, no warmup / no decay (ablation spec).') + p.add_argument('--sft-weight', type=float, default=1.0, + help='advantage magnitude for SFT samples (ablation spec: 1.0).') + p.add_argument('--drop-zero-adv', action='store_true', + help='reserved single-point ablation (接口方案 #10): drop zero-advantage ' + 'candidates from RL training batches instead of keeping them in the ' + 'token-mean denominator (SEAM口径). Default off; NOT part of E1-E12.') + + # --- run control (new) --- + p.add_argument('--max-updates', type=int, default=50, + help='stop after this many PARAMETER UPDATES (the "step" unit).') + p.add_argument('--eval-every-updates', type=int, default=5) + p.add_argument('--save-every-updates', type=int, default=0, + help='save a weights-only checkpoint (-u) every N updates; 0 = only ' + 'the final save. lr is constant so optimizer state is NOT saved.') + p.add_argument('--resume-from', default='', + help='checkpoint dir name under --output-dir (e.g. E1-final / E1-u50) or an ' + 'absolute path. Loads weights into skill_model, restores updates/chunk ' + 'position from train_state.json (falls back to DONE.json), appends to ' + 'the record files, and bypasses the DONE.json skip guard. Raise ' + '--max-updates beyond the restored count to actually continue.') + p.add_argument('--pool-max', type=int, default=2048, + help='SFT pool per-queue cap (drop oldest; bounds majority backlog).') + p.add_argument('--rubric-global-dir', default='', + help="dir for the cross-experiment global rubric cache " + "(default: parent of --output-dir).") + + # --- output / logging --- + p.add_argument('--output-dir', default='./output.ablate12/exp') + p.add_argument('--no-cache', action='store_true') + p.add_argument('--force', action='store_true', + help='rerun even if /DONE.json marks this experiment complete.') + p.add_argument('--swanlab-project', default='twinkle') + + args = p.parse_args(argv) + + # resolve the experiment spec + if args.exp: + spec = get_spec(args.exp) + else: + if not (args.method and args.thinking and args.style): + p.error('provide --exp E5, or all of --method/--thinking/--style') + spec = ExpSpec(name='Ex', method=args.method, thinking=args.thinking, style=args.style) + + # sanity: Ray dp rule — sft/eval batch must divide TRAIN_DP + if args.sft_batch_size % v2.TRAIN_DP != 0: + p.error(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple of TRAIN_DP ({v2.TRAIN_DP})') + if args.train_micro_batch and args.train_micro_batch % v2.TRAIN_DP != 0: + p.error(f'--train-micro-batch ({args.train_micro_batch}) must be a multiple of TRAIN_DP ({v2.TRAIN_DP})') + if args.chunk_size < 1: + p.error('--chunk-size must be >= 1') + args.rubric_global_dir = args.rubric_global_dir or None + return args, spec + + +def main(argv=None): + args, spec = _build_args(argv) + sys.stderr.write(f'[ablate] running {spec.name}: view={spec.view} method={spec.method} ' + f'thinking={spec.thinking} style={spec.style} loss={spec.loss} ' + f'smt={spec.skill_max_tokens} -> {args.output_dir}\n') + run_experiment(args, spec) + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/skill2lora/skill_ablate/methods.py b/cookbook/exp/skill2lora/skill_ablate/methods.py new file mode 100644 index 000000000..b1a6cd5c4 --- /dev/null +++ b/cookbook/exp/skill2lora/skill_ablate/methods.py @@ -0,0 +1,665 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Pluggable training methods (E1-E12) over a shared context. + +Each method implements TrainMethod: ``step(chunk, ci) -> dict`` where the dict carries at +least ``n_updates`` (parameter-update count, the "step" unit per skill_quality_analysis.md +#2), plus ``metrics`` (scalars to log) and ``gen_records`` (per-problem rollout audit rows +for gen_records.jsonl). New methods / losses are added by writing one class + one +METHOD_REGISTRY entry — the trainer never changes. + +Reuse map (all imported from train_skill_v2 / rollouting, never edited): +- bnpo (view B): v2 ``process_chunk`` + ``_train_step`` verbatim (query-only). +- rl_ab / rl_err (view A): bare greedy solve -> wrong=A(query+rubric) / right=B(query-only), + 8 skills each -> executor greedy reward -> group advantage -> BNPO on the rubric trajectory + (train-with-rubric); rl_err drops the B line from training. Rubric API calls run on threads + WHILE the B line rolls out on GPU (API/GPU overlap). +- opsd (view A): error problems, 1 student skill (query-only, T=0.5); teacher forward + (student prompt + rubric appended to the SYSTEM prompt, same response) -> per-token OPSD + KL (loss='opsd'). Teacher logps are extracted RESPONSE-ONLY via a client-side template + encode (teacher/student prompt lengths differ, so the full-sequence form would misalign). +- improve_sft (view A): first-pass 1 skill (query-only, T=0.5); correct -> positive SFT seed + (no leak, <=4096 chars); parseable-but-wrong -> rubric regen (2-in-8 pick 1) -> negative + SFT seed; unparseable first pass is SKIPPED (no trajectory to diagnose); balanced 1:1 pool + (majority side down-sampled per chunk, never backlogged) -> SFT. +- sft (view A): bare wrong -> rubric -> regen (query+rubric) 2-in-8 -> plain pool -> SFT. + +Training-batch helper ``_train_batch`` mirrors v2 ``_train_step`` exactly (empty-response +filter, drop_last to TRAIN_DP, micro-batch by sft_batch_size, ppo_mini_batch_size multi-step +with pre-computed ref/old logps, clip+step per mini, ckpt sync, calculate_metric) but takes a +swappable trajectory builder (query-only vs query+rubric) and an optional teacher builder for +OPSD. In the OPSD path no ref forward is done at all: OPSDLoss uses only teacher_logps +(kl_beta / ref_logps play no role there). +""" +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, List, Optional, Tuple + +import numpy as np + +import train_skill_v2 as v2 +from train_skill_v2 import ( + TRAIN_DP, + _answer_leaked, + _assign_advantages, + _clean_text, + _empty_roll, + _extract_skill, + _parse_seq, + _regen_prompt, + _run_samples, + _skill_reward, + _skillgen_prompt, + _train_step, + build_direct_prompt, + build_skill_solve_prompt, + process_chunk, +) + +from .pool import NEG, POS, SamplePool +from .rollouting import ( + opsd_teacher_trajectory, + query_only_train_trajectory, + rubric_skillgen_prompt, + rubric_train_trajectory, +) + + +@dataclass +class MethodContext: + """Everything a method needs; assembled once by the trainer.""" + skill_model: Any + ref_model: Any + skill_sampler: Any + base_sampler: Any + ckpt: Any + skill_dp: int + base_dp: int + args: Any + checker: Any = None + rubric_cache: Any = None # GlobalRubricCache (RL/SFT) or LocalRubricCache (improve/opsd) + pool: Optional[SamplePool] = None # SFT-family accumulator (None for RL/OPSD) + encode_template: Any = None # client-side Template clone (OPSD teacher alignment only) + extra: Dict[str, Any] = field(default_factory=dict) + + +# =========================================================================================== +# shared low-level helpers (reuse v2 primitives; only orchestration is new) +# =========================================================================================== +def _bare_solve(ctx: MethodContext, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Bare-problem greedy (T=0) executor solve; returns one roll per record (order-aligned).""" + out = _run_samples(ctx.base_sampler, [build_direct_prompt(r['problem']) for r in records], + 1, ctx.args.max_tokens, ctx.base_dp, temperature=0.0) + return [(_parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll()) + for r, seqs in zip(records, out)] + + +def _rubric_entry(record: Dict[str, Any], roll: Dict[str, Any]) -> Dict[str, Any]: + """Build the entry _diagnose_entry expects from a failure trajectory.""" + return {'problem': record['problem'], 'reference_answer': record['reference_answer'], + 'data_id': record.get('data_id', ''), + 'fail_segment': roll.get('text', ''), + 'fail_stop_reason': roll.get('stop_reason', 'none')} + + +def _diagnose_parallel(ctx: MethodContext, + jobs: List[Tuple[Dict[str, Any], Optional[str]]]) -> List[str]: + """Run rubric diagnoses in parallel threads (pure API, no GPU; DiskCache.put is locked). + + ``jobs`` is a list of (entry, skill_or_None); returns diagnoses aligned to jobs + ('' on cache-off / API error). Parallelism = --rubric-workers (was serial before).""" + if not jobs or ctx.rubric_cache is None: + return [''] * len(jobs) + workers = max(1, min(ctx.args.rubric_workers, len(jobs))) + with ThreadPoolExecutor(max_workers=workers) as ex: + return list(ex.map( + lambda j: ctx.rubric_cache.get_or_diagnose(j[0], ctx.checker, skill=j[1]) or '', + jobs)) + + +def _skillgen_solve(ctx: MethodContext, items: List[Dict[str, Any]], n_skills: int, + temperature: float) -> None: + """For each item {record, prompt}: sample n skills, greedy-solve each, attach a + ``_cands`` list (v2 shape) onto the record so ``_assign_advantages`` can be reused.""" + args = ctx.args + sg_out = _run_samples(ctx.skill_sampler, [it['prompt'] for it in items], n_skills, + args.skill_max_tokens, ctx.skill_dp, + temperature=temperature, top_p=args.skill_gen_top_p, + top_k=args.skill_gen_top_k) + flat = [] + for it, seqs in zip(items, sg_out): + it['record']['_cands'] = [] + for s in seqs: + resp = _clean_text(getattr(s, 'decoded', '') or '') + block = _extract_skill(resp) or '' + cand = {'skills': block, 'response': resp, 'parseable': bool(block), + 'leaked': None, 'with_pass': None, 'reward': None, 'rolls': [], + 'advantage': 0.0, 'kept': False, + 'skillgen_stop': getattr(s, 'stop_reason', None), + 'skillgen_tokens': len(getattr(s, 'tokens', None) or [])} + it['record']['_cands'].append(cand) + if block: + flat.append((it, cand)) + for it, c in flat: + c['leaked'] = _answer_leaked(c['skills'], it['record']['reference_answer']) + if flat: + ws = _run_samples(ctx.base_sampler, + [build_skill_solve_prompt(it['record']['problem'], c['skills']) for it, c in flat], + 1, args.max_tokens, ctx.base_dp, temperature=0.0) + for (it, c), seqs in zip(flat, ws): + roll = _parse_seq(seqs[0], it['record']['reference_answer']) if seqs else _empty_roll() + c['rolls'] = [roll] + c['with_pass'] = 1.0 if roll['correct'] else 0.0 + c['reward'] = _skill_reward(c['parseable'], roll['correct']) + for it in items: + for c in it['record']['_cands']: + if c['reward'] is None: + c['reward'] = 0.0 + + +def _grpo_records(records: List[Dict[str, Any]], with_rubric: bool) -> List[Dict[str, Any]]: + """Flatten per-problem _cands into GRPO train records (v2 shape). + ``with_rubric`` tags each record so the trajectory builder knows which prompt to rebuild.""" + recs = [] + for r in records: + for c in r.get('_cands', []): + if c.get('reward') is None: + continue + recs.append({'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'data_id': r.get('data_id', ''), 'response': c['response'], + 'skills': c['skills'], 'advantage': c['advantage'], + 'kept': c['kept'], 'reward': c['reward'], 'rubric': r.get('_rubric', ''), + 'with_rubric': with_rubric, 'sft': False}) + return recs + + +def _leak_split(pairs: List[Tuple[bool, bool]]) -> Dict[str, float]: + """#10 monitoring curves: leaked&correct vs leaked&wrong rates over parseable skills. + "泄露正确答案可接受"不等于"泄露无害"——有害的是错误数值注入，两条曲线拆开监控。""" + n = len(pairs) + if not n: + return {'leak/correct_rate': 0.0, 'leak/wrong_rate': 0.0} + return {'leak/correct_rate': sum(1 for lk, ok in pairs if lk and ok) / n, + 'leak/wrong_rate': sum(1 for lk, ok in pairs if lk and not ok) / n} + + +def _cand_leak_pairs(records: List[Dict[str, Any]]) -> List[Tuple[bool, bool]]: + return [(bool(c['leaked']), bool(c['with_pass'])) + for r in records for c in r.get('_cands', []) + if c.get('parseable') and c.get('with_pass') is not None] + + +def _cand_pass_metrics(records: List[Dict[str, Any]]) -> Dict[str, float]: + """Mean-family train metrics, same family as eval acc_mean1 (ws_acc is pass@8-inflated): + candidate_pass = P(correct | parseable); clean_pass = P(correct | parseable & terminated). + Sharper channel split by ANSWER AVAILABILITY (truncated rolls still count correct when a + balanced \\boxed{} landed before the budget — 13% of E1 truncations did): + answered_rate = P(pred emitted) and answered_pass = P(correct | answered), the content-only + channel (E1 vs E5 core comparison curve); plus parse/trunc rates for the format channel.""" + cands = [c for r in records for c in r.get('_cands', [])] + if not cands: + return {} + m = {'skill/parse_rate': sum(1 for c in cands if c.get('parseable')) / len(cands)} + scored = [c for c in cands + if c.get('parseable') and c.get('with_pass') is not None and c.get('rolls')] + if scored: + m['acc/candidate_pass'] = sum(c['with_pass'] for c in scored) / len(scored) + clean = [c for c in scored if c['rolls'][0].get('stop_reason') != 'length'] + m['term/withskill_trunc_frac'] = 1.0 - len(clean) / len(scored) + if clean: + m['acc/clean_pass'] = sum(c['with_pass'] for c in clean) / len(clean) + answered = [c for c in scored if c['rolls'][0].get('pred') not in (None, '')] + m['term/answered_rate'] = len(answered) / len(scored) + if answered: + m['acc/answered_pass'] = sum(c['with_pass'] for c in answered) / len(answered) + return m + + +def _train_metrics(metric: Optional[Dict[str, Any]]) -> Dict[str, float]: + """calculate_metric -> swan-ready train/* scalars (same key handling as v2 _swan_metrics).""" + d = {} + for k, val in (metric or {}).items(): + if not v2._is_num(val): + continue + if k.startswith('learning rate'): + if 'group 1' in k: + d['train/lr'] = float(val) + else: + d[f'train/{k.replace(" ", "_")}'] = float(val) + return d + + +# =========================================================================================== +# OPSD teacher alignment: client-side encode -> response-only teacher logps +# =========================================================================================== +def _align_teacher(ctx: MethodContext, samples, traj_fn, teacher_fn): + """Encode student & teacher trajectories with a CLIENT-SIDE clone of the remote template + and keep only samples whose response-token counts match (they always should — same + assistant text — but max-length 'delete' truncation or tokenizer drift must not silently + misalign the distillation). Returns (kept_samples, teacher_label_positions, n_dropped). + + teacher_label_positions[i] are the (rolled-)label indices of sample i inside its OWN + unpadded teacher sequence; right padding keeps these indices valid in the padded batch, + so they can slice the teacher's full-sequence logps down to the response-only form that + OPSDLoss requires (full-sequence form would misalign: prompts differ in length). + """ + tmpl = ctx.encode_template + assert tmpl is not None, 'OPSD needs ctx.encode_template (built by the trainer)' + keep, pos_lists, dropped = [], [], 0 + for s in samples: + st = tmpl.encode(traj_fn(s)) + tt = tmpl.encode(teacher_fn(s)) + if st is None or tt is None: # deleted by max-length truncation + dropped += 1 + continue + spos = np.where(np.asarray(st.get('labels')) != -100)[0] + tpos = np.where(np.asarray(tt.get('labels')) != -100)[0] + if len(spos) != len(tpos) or len(tpos) == 0: + dropped += 1 + continue + keep.append(s) + pos_lists.append(tpos) + return keep, pos_lists, dropped + + +def _gather_response_logps(full_logps, pos_lists) -> List[List[float]]: + """Slice full-sequence [B, S] teacher logps down to per-sample response-only lists.""" + rows = [] + for i, pos in enumerate(pos_lists): + row = full_logps[i] + row = row.tolist() if hasattr(row, 'tolist') else list(row) + assert len(pos) == 0 or int(pos[-1]) < len(row), \ + f'teacher logps row {i} shorter than label positions ({len(row)} <= {int(pos[-1])})' + rows.append([float(row[int(p)]) for p in pos]) + return rows + + +def _train_batch(ctx: MethodContext, samples: List[Dict[str, Any]], + traj_fn: Callable[[Dict[str, Any]], Dict[str, Any]], + teacher_fn: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None, + ) -> Tuple[int, Dict[str, float]]: + """Parameter update(s) over ``samples`` with a swappable trajectory builder. + + Mirrors v2 ``_train_step`` (empty-response filter, drop_last to TRAIN_DP, micro-batch by + sft_batch_size, ppo_mini_batch_size multi-step with ALL ref/old/teacher logps pre-computed + BEFORE the first optimizer step — the teacher is the trainable model itself, so computing + it after a step would go off-policy). Returns (n_updates, train metrics). + """ + args = ctx.args + samples = [s for s in samples if (s.get('response') or '').strip()] + is_opsd = teacher_fn is not None + teacher_pos: Optional[List] = None + n_align_drop = 0 + if is_opsd: + samples, teacher_pos, n_align_drop = _align_teacher(ctx, samples, traj_fn, teacher_fn) + if not samples: + return 0, {} + n_keep = (len(samples) // TRAIN_DP) * TRAIN_DP + if n_keep == 0: + return 0, {} + samples = samples[:n_keep] + if teacher_pos is not None: + teacher_pos = teacher_pos[:n_keep] + trajs = [traj_fn(s) for s in samples] + advs = None if is_opsd else [float(s['advantage']) for s in samples] + # micro 尺寸与“攒批/采样批（sft_batch_size=16，冻结口径）”解耦：think/8192 实验序列长一倍， + # fp32 主权重后 8 条/卡的 backward 会 OOM，用 --train-micro-batch 切细（梯度按 micro 数归一， + # 数学等价）；默认 0 = 跟随 sft_batch_size，nothink 实验行为不变。 + n = len(trajs) + sft = getattr(args, 'train_micro_batch', 0) or args.sft_batch_size + mini = args.ppo_mini_batch_size if args.ppo_mini_batch_size > 0 else n + mini = max(sft, (mini // sft) * sft) + multi_step = mini < n + # pre-compute every micro's ref/old/teacher logps BEFORE any update (v2 pattern) + micro_ref, micro_old, micro_teacher = [], [], [] + for i in range(0, n, sft): + mb = trajs[i:i + sft] + if is_opsd: + # no ref forward at all: OPSDLoss uses only teacher_logps (kl_beta plays no role) + t_mb = [teacher_fn(s) for s in samples[i:i + sft]] + t_full = ctx.skill_model.forward_only(inputs=t_mb).get('logps') + micro_teacher.append(_gather_response_logps(t_full, teacher_pos[i:i + sft])) + micro_ref.append(None) + micro_old.append(None) + else: + micro_ref.append(ctx.ref_model.forward_only(inputs=mb).get('logps')) + micro_old.append(ctx.skill_model.forward_only(inputs=mb).get('logps') if multi_step else None) + micro_teacher.append(None) + n_steps = 0 + for ms in range(0, n, mini): + for i in range(ms, min(ms + mini, n), sft): + k = i // sft + if is_opsd: + ctx.skill_model.forward_backward(inputs=trajs[i:i + sft], + teacher_logps=micro_teacher[k]) + else: + ctx.skill_model.forward_backward(inputs=trajs[i:i + sft], advantages=advs[i:i + sft], + old_logps=micro_old[k], ref_logps=micro_ref[k]) + ctx.skill_model.clip_grad_and_step() + n_steps += 1 + ctx.ckpt.sync_weights(merge_and_sync=True) + metrics = _train_metrics(ctx.skill_model.calculate_metric(is_training=True)) + metrics['train/n_samples'] = float(n) + if is_opsd and n_align_drop: + metrics['train/n_align_dropped'] = float(n_align_drop) + return n_steps, metrics + + +# =========================================================================================== +# method plugins +# =========================================================================================== +class TrainMethod: + needs_rubric: bool = False + + def __init__(self, ctx: MethodContext): + self.ctx = ctx + + def step(self, chunk: List[Dict[str, Any]], ci: int) -> Dict[str, Any]: + raise NotImplementedError + + +class BnpoMethod(TrainMethod): + """view B, query-only GRPO/BNPO — v2 process_chunk + _train_step verbatim.""" + needs_rubric = False + + def step(self, chunk, ci): + ctx = self.ctx + full, summary, grpo, _buf_a = process_chunk( + ctx.base_sampler, ctx.skill_sampler, chunk, ci, ctx.base_dp, ctx.skill_dp, ctx.args) + if grpo and getattr(ctx.args, 'drop_zero_adv', False): + grpo = [g for g in grpo if abs(g['advantage']) > 1e-9] + n_upd, tmetrics = 0, {} + if grpo: + log = _train_step(ctx.skill_model, ctx.ref_model, ctx.ckpt, grpo, ctx.args) + # step = ACTUAL parameter updates: empty-response filter / drop_last may yield 0 + n_upd = int(log.get('n_steps', 0)) + tmetrics = _train_metrics(log.get('metric')) + tmetrics['train/n_samples'] = float(log.get('n_grpo', 0) + log.get('n_sft', 0)) + metrics = {'signal/zero_grad_frac': summary['zero_grad_frac'], + 'acc/withskill_pass': summary['avg_withskill_pass'], + 'leak/rate': summary['leak_rate'], **tmetrics, + **_cand_pass_metrics(chunk), + **_leak_split(_cand_leak_pairs(chunk))} + return {'n_updates': n_upd, 'summary': summary, 'metrics': metrics, + 'gen_records': full} + + +class _RLViewA(TrainMethod): + """Shared view-A RL: bare solve -> A(query+rubric)/B(query-only) skill-gen -> reward -> + group advantage -> BNPO. ``train_b`` toggles whether the right-answer B line is trained. + + API/GPU overlap: rubric diagnoses for the wrong (A) problems run on background threads + WHILE the right (B) problems' query-only skill-gen + greedy validation run on GPU; the + A-line rollout starts as soon as the diagnoses land.""" + needs_rubric = True + train_b = True + + def step(self, chunk, ci): + ctx = self.ctx + args = ctx.args + rolls = _bare_solve(ctx, chunk) + wrong = [(r, roll) for r, roll in zip(chunk, rolls) if not roll['correct']] + right = [r for r, roll in zip(chunk, rolls) if roll['correct']] + # kick off rubric API calls in the background, then run the B line on GPU meanwhile + diag_pool = ThreadPoolExecutor(max_workers=1) + diag_fut = diag_pool.submit( + _diagnose_parallel, ctx, [(_rubric_entry(r, roll), None) for r, roll in wrong]) + try: + b_items = [{'record': r, 'prompt': _skillgen_prompt(r['problem'])} for r in right] + for r in right: + r['_rubric'] = '' + if b_items: + _skillgen_solve(ctx, b_items, args.n_skills, temperature=args.skill_gen_temperature) + diags = diag_fut.result() + finally: + diag_pool.shutdown(wait=False) + a_items = [] + for (r, _roll), diag in zip(wrong, diags): + r['_rubric'] = diag + a_items.append({'record': r, 'prompt': rubric_skillgen_prompt(r['problem'], diag)}) + if a_items: + _skillgen_solve(ctx, a_items, args.n_skills, temperature=args.skill_gen_temperature) + _assign_advantages(chunk, args) + a_recs = _grpo_records([r for r, _ in wrong], with_rubric=True) + b_recs = _grpo_records(right, with_rubric=False) + train_recs = a_recs + (b_recs if self.train_b else []) + has_signal = any(abs(s['advantage']) > 1e-9 for s in train_recs) + if has_signal and getattr(args, 'drop_zero_adv', False): + train_recs = [s for s in train_recs if abs(s['advantage']) > 1e-9] + n_upd, tmetrics = 0, {} + if has_signal: + n_upd, tmetrics = _train_batch( + ctx, train_recs, + traj_fn=lambda s: (rubric_train_trajectory(s) if s['with_rubric'] + else query_only_train_trajectory(s))) + return {'n_updates': n_upd, + 'metrics': {'signal/n_wrong_A': float(len(wrong)), + 'signal/n_right_B': float(len(right)), **tmetrics, + **_cand_pass_metrics(chunk), + **_leak_split(_cand_leak_pairs(chunk))}, + 'gen_records': v2._full_records(chunk, ci)} + + +class RlAbMethod(_RLViewA): + train_b = True + + +class RlErrMethod(_RLViewA): + train_b = False + + +class OpsdMethod(TrainMethod): + """view A OPSD (skill_quality_analysis.md 改进skill+OPSD): first-pass ONE skill + (query-only, T=improve), executor solve WITH skill; for WRONG problems, diagnose the + with-skill failure (local cache, key=data_id+skill), then distill the skill-gen response + from the query-only (student) toward the rubric-in-system-prompt (teacher) distribution + per token. Rubric API calls run threaded right after the failures are known.""" + needs_rubric = True + + def step(self, chunk, ci): + ctx = self.ctx + args = ctx.args + # first-pass ONE skill per problem (query-only, improve temperature) + sg = _run_samples(ctx.skill_sampler, [_skillgen_prompt(r['problem']) for r in chunk], + 1, args.skill_max_tokens, ctx.skill_dp, + temperature=args.improve_skill_temperature) + first = [] + for r, seqs in zip(chunk, sg): + resp = _clean_text(getattr(seqs[0], 'decoded', '') or '') if seqs else '' + first.append((r, resp, _extract_skill(resp) or '')) + # executor solve WITH skill (only parseable skills) + flat = [(r, resp, sk) for r, resp, sk in first if sk] + roll_by = {} + if flat: + solve = _run_samples(ctx.base_sampler, + [build_skill_solve_prompt(r['problem'], sk) for r, _, sk in flat], + 1, args.max_tokens, ctx.base_dp, temperature=0.0) + for (r, _, sk), seqs in zip(flat, solve): + roll_by[id(r)] = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + wrong = [(r, resp, sk, roll_by[id(r)]) for r, resp, sk in first + if sk and id(r) in roll_by and not roll_by[id(r)]['correct']] + leak_pairs = [(bool(_answer_leaked(sk, r['reference_answer'])), roll_by[id(r)]['correct']) + for r, _resp, sk in flat if id(r) in roll_by] + # rubric diagnoses (threaded; nothing left to overlap on GPU this chunk) + diags = _diagnose_parallel( + ctx, [(_rubric_entry(r, roll), sk) for r, _resp, sk, roll in wrong]) + samples = [{'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'data_id': r.get('data_id', ''), 'response': resp, 'rubric': diag} + for (r, resp, sk, _roll), diag in zip(wrong, diags)] + n_upd, tmetrics = 0, {} + if samples: + n_upd, tmetrics = _train_batch(ctx, samples, + traj_fn=query_only_train_trajectory, # student: query-only + teacher_fn=opsd_teacher_trajectory) # teacher: +rubric in system + return {'n_updates': n_upd, + 'metrics': {'signal/n_wrong': float(len(wrong)), **tmetrics, + **_leak_split(leak_pairs)}, + 'gen_records': [ + {'record_type': 'problem', 'chunk': ci, 'data_id': r.get('data_id', ''), + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'skill': sk, 'parseable': bool(sk), + 'withskill_correct': roll_by[id(r)]['correct'] if id(r) in roll_by else None} + for r, _resp, sk in first]} + + +class _SFTFamily(TrainMethod): + """Shared SFT accumulation + fire. Subclasses fill ``collect`` to add pool samples.""" + needs_rubric = True + + def _sft_record(self, problem, ref, data_id, skill): + return {'problem': problem, 'reference_answer': ref, 'data_id': data_id, + 'response': f'\n{skill}\n', 'skills': skill, + 'advantage': float(self.ctx.args.sft_weight), 'sft': True} + + def collect(self, chunk) -> Tuple[List[Tuple[bool, bool]], List[Dict[str, Any]]]: + """Roll out + fill the pool; returns (leak_pairs, gen_records).""" + raise NotImplementedError + + def step(self, chunk, ci): + ctx = self.ctx + leak_pairs, gen_records = self.collect(chunk) + ctx.pool.rebalance() # 1:1 by the minority side, surplus DISCARDED (#18b 不积压) + n_upd, tmetrics = 0, {} + for batch in ctx.pool.draw_all_ready(): + n, m = _train_batch(ctx, batch, traj_fn=query_only_train_trajectory) # SFT: query-only (#6) + n_upd += n + tmetrics.update(m) + return {'n_updates': n_upd, + 'metrics': {**{f'pool/{k}': float(x) for k, x in ctx.pool.sizes().items()}, + **tmetrics, **_leak_split(leak_pairs)}, + 'gen_records': gen_records} + + # -- shared: regenerate skills under rubric, greedy-validate, pick a 2-in-8 passer -- + def _regen_pick(self, record, diag: str, use_orig_skill: bool, orig_skill: str = ''): + ctx = self.ctx + args = ctx.args + if not diag: + return None + prompt = (_regen_prompt(record['problem'], orig_skill, diag) if use_orig_skill + else rubric_skillgen_prompt(record['problem'], diag)) + sg = _run_samples(ctx.skill_sampler, [prompt], args.passatk_k, args.skill_max_tokens, + ctx.skill_dp, temperature=args.passatk_skill_temp, + top_p=args.passatk_skill_top_p) + seqs = sg[0] if sg else [] + cands = [] + for s in seqs: + resp = _clean_text(getattr(s, 'decoded', '') or '') + skill = _extract_skill(resp) or '' + if not skill or len(skill) > args.skill_char_limit: + continue + if _answer_leaked(skill, record['reference_answer']): + continue + cands.append(skill) + if not cands: + return None + solve = _run_samples(ctx.base_sampler, + [build_skill_solve_prompt(record['problem'], sk) for sk in cands], + 1, args.max_tokens, ctx.base_dp, temperature=0.0) + passers = [] + for sk, seqs2 in zip(cands, solve): + roll2 = _parse_seq(seqs2[0], record['reference_answer']) if seqs2 else _empty_roll() + if roll2['correct'] and roll2['terminated']: + passers.append(sk) + if len(passers) < args.passatk_m: + return None + # pick the passer closest to the length budget (short-but-not-empty floor, as in v2) + return min(passers, key=lambda sk: abs(len(sk) - args.len_budget)) + + +class SftMethod(_SFTFamily): + """Plain SFT: bare wrong -> rubric regen (query+rubric, no orig skill) 2-in-8 -> pool.""" + def collect(self, chunk): + ctx = self.ctx + rolls = _bare_solve(ctx, chunk) + wrong = [(r, roll) for r, roll in zip(chunk, rolls) if not roll['correct']] + diags = _diagnose_parallel(ctx, [(_rubric_entry(r, roll), None) for r, roll in wrong]) + gen_records = [] + for (r, _roll), diag in zip(wrong, diags): + skill = self._regen_pick(r, diag, use_orig_skill=False) + if skill: + ctx.pool.add(self._sft_record(r['problem'], r['reference_answer'], + r.get('data_id', ''), skill), NEG) + gen_records.append({'record_type': 'problem', 'data_id': r.get('data_id', ''), + 'problem': r['problem'], 'regen_accepted': bool(skill)}) + return [], gen_records # bare solve has no skill, so no leak pairs here + + +class ImproveSftMethod(_SFTFamily): + """Improve-skill + SFT: first-pass 1 skill (query-only, T=0.5); correct -> positive pool + (no leak, <=char_limit); parseable-but-wrong -> rubric regen (with orig skill) 2-in-8 -> + negative pool; unparseable first pass is skipped (empty trajectory would only feed the + teacher garbage). Balanced 1:1 pool, majority side discarded per chunk (#15b/#18b). + + API/GPU overlap: the wrong problems' rubric diagnoses run on background threads WHILE + the regen sampling for previously-diagnosed problems occupies the GPU (diagnoses land + before the first regen finishes, so the loop below never blocks on the API).""" + def collect(self, chunk): + ctx = self.ctx + args = ctx.args + # first-pass ONE skill per problem (query-only, improve temperature), greedy-solve + sg = _run_samples(ctx.skill_sampler, [_skillgen_prompt(r['problem']) for r in chunk], + 1, args.skill_max_tokens, ctx.skill_dp, + temperature=args.improve_skill_temperature) + first = [] + for r, seqs in zip(chunk, sg): + resp = _clean_text(getattr(seqs[0], 'decoded', '') or '') if seqs else '' + first.append((r, _extract_skill(resp) or '')) + flat = [(r, sk) for r, sk in first if sk] + rolls_by = {} + if flat: + solve = _run_samples(ctx.base_sampler, + [build_skill_solve_prompt(r['problem'], sk) for r, sk in flat], + 1, args.max_tokens, ctx.base_dp, temperature=0.0) + for (r, sk), seqs in zip(flat, solve): + rolls_by[id(r)] = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + leak_pairs = [(bool(_answer_leaked(sk, r['reference_answer'])), rolls_by[id(r)]['correct']) + for r, sk in flat if id(r) in rolls_by] + wrong = [(r, sk, rolls_by[id(r)]) for r, sk in flat + if id(r) in rolls_by and not rolls_by[id(r)]['correct']] + # launch ALL diagnoses on threads first, then regen serially on GPU as they land + diag_pool = ThreadPoolExecutor(max_workers=max(1, min(args.rubric_workers, len(wrong) or 1))) + futs = [diag_pool.submit( + lambda e=_rubric_entry(r, roll), s=sk: + (ctx.rubric_cache.get_or_diagnose(e, ctx.checker, skill=s) or '') + if ctx.rubric_cache else '') for r, sk, roll in wrong] + gen_records = [] + try: + for r, sk in flat: + roll = rolls_by.get(id(r)) + if roll is None: + continue + if roll['correct']: + # positive seed: first-pass skill that worked, no leak, within char limit + if len(sk) <= args.skill_char_limit \ + and not _answer_leaked(sk, r['reference_answer']): + ctx.pool.add(self._sft_record(r['problem'], r['reference_answer'], + r.get('data_id', ''), sk), POS) + gen_records.append({'record_type': 'problem', 'data_id': r.get('data_id', ''), + 'problem': r['problem'], 'first_correct': True, 'skill': sk}) + for (r, sk, roll), fut in zip(wrong, futs): + # negative seed: rubric regen conditioned on the FAILED first-pass skill + skill = self._regen_pick(r, fut.result(), use_orig_skill=True, orig_skill=sk) + if skill: + ctx.pool.add(self._sft_record(r['problem'], r['reference_answer'], + r.get('data_id', ''), skill), NEG) + gen_records.append({'record_type': 'problem', 'data_id': r.get('data_id', ''), + 'problem': r['problem'], 'first_correct': False, + 'skill': sk, 'regen_accepted': bool(skill)}) + finally: + diag_pool.shutdown(wait=False) + return leak_pairs, gen_records + + +METHOD_REGISTRY: Dict[str, Callable[[MethodContext], TrainMethod]] = { + 'bnpo': BnpoMethod, + 'rl_ab': RlAbMethod, + 'rl_err': RlErrMethod, + 'opsd': OpsdMethod, + 'sft': SftMethod, + 'improve_sft': ImproveSftMethod, +} + + +def build_method(method: str, ctx: MethodContext) -> TrainMethod: + if method not in METHOD_REGISTRY: + raise KeyError(f'unknown method {method!r}; valid: {sorted(METHOD_REGISTRY)}') + return METHOD_REGISTRY[method](ctx) diff --git a/cookbook/exp/skill2lora/skill_ablate/pool.py b/cookbook/exp/skill2lora/skill_ablate/pool.py new file mode 100644 index 000000000..346ba7c69 --- /dev/null +++ b/cookbook/exp/skill2lora/skill_ablate/pool.py @@ -0,0 +1,125 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""SamplePool: accumulate training samples across chunks and emit fixed-size batches. + +Ray-infra size rule (skill_quality_analysis.md #11-12): training must drop_last to a +TRAIN_DP multiple and must NEVER pad new sequences (overfitting guard). The actual +drop_last happens inside v2's ``_train_step``; this pool's only job is "accumulate until a +full batch is available, then hand exactly one batch to ``_train_step``". Batch size is a +TRAIN_DP multiple (default 16 = sft_batch_size), so the batch always divides evenly. + +Two modes: +- plain (balanced=False): a single FIFO queue; ready when it holds >= batch_size; a draw + pops the oldest ``batch_size`` samples and keeps the remainder pooled for next time. +- balanced (balanced=True): separate positive/negative queues for the improve-skill+SFT + 1:1 requirement (#15b/#18b). A batch is half positives + half negatives; ready when BOTH + halves are available. After every chunk the caller invokes ``rebalance()``: the majority + side is down-sampled to the minority side and the surplus is DISCARDED immediately + (#18b "以少的一侧为准下采样多的一侧，其余丢弃不积压") — no stale majority backlog can + accumulate, so early-policy easy positives never train dozens of chunks later. + ``max_pool`` remains as a safety cap only (drop oldest). + +This module is dependency-free (pure stdlib) and unit-testable without torch / a GPU. +""" +from collections import deque +from typing import Any, Deque, Dict, List, Optional + +POS, NEG = 'pos', 'neg' + + +class SamplePool: + def __init__(self, batch_size: int = 16, balanced: bool = False, + max_pool: Optional[int] = None): + if batch_size < 1: + raise ValueError('batch_size must be >= 1') + if balanced and batch_size % 2 != 0: + raise ValueError('balanced pool needs an even batch_size (half pos + half neg)') + self.batch_size = batch_size + self.balanced = balanced + self.max_pool = max_pool + self._q: Deque[Dict[str, Any]] = deque() # plain mode + self._pos: Deque[Dict[str, Any]] = deque() # balanced mode + self._neg: Deque[Dict[str, Any]] = deque() + self._added = 0 + self._emitted = 0 + + # -- ingest ------------------------------------------------------------------------- + def add(self, sample: Dict[str, Any], label: str = NEG) -> None: + self._added += 1 + if not self.balanced: + self._q.append(sample) + self._trim(self._q) + return + if label == POS: + self._pos.append(sample) + self._trim(self._pos) + elif label == NEG: + self._neg.append(sample) + self._trim(self._neg) + else: + raise ValueError(f'label must be {POS!r} or {NEG!r}, got {label!r}') + + def add_many(self, samples: List[Dict[str, Any]], label: str = NEG) -> None: + for s in samples: + self.add(s, label) + + def _trim(self, q: Deque[Dict[str, Any]]) -> None: + if self.max_pool is not None: + while len(q) > self.max_pool: + q.popleft() # drop oldest to bound memory / avoid stale majority backlog + + def rebalance(self) -> int: + """Balanced mode: down-sample the majority queue to the minority size, discarding + the NEWEST surplus (this chunk's excess intake — the #18b "其余丢弃不积压" rule). + Called once per chunk; since intake is re-balanced every chunk, both queues stay + equal-length and no side ever backlogs. Returns the number of discarded samples. + No-op in plain mode.""" + if not self.balanced: + return 0 + target = min(len(self._pos), len(self._neg)) + dropped = 0 + for q in (self._pos, self._neg): + while len(q) > target: + q.pop() # newest first: the surplus was added this chunk + dropped += 1 + return dropped + + # -- state -------------------------------------------------------------------------- + def ready(self) -> bool: + if not self.balanced: + return len(self._q) >= self.batch_size + half = self.batch_size // 2 + return len(self._pos) >= half and len(self._neg) >= half + + def sizes(self) -> Dict[str, int]: + if not self.balanced: + return {'pool': len(self._q)} + return {'pos': len(self._pos), 'neg': len(self._neg)} + + @property + def total_added(self) -> int: + return self._added + + @property + def total_emitted(self) -> int: + return self._emitted + + # -- draw --------------------------------------------------------------------------- + def draw(self) -> List[Dict[str, Any]]: + """Pop exactly one batch; raises if not ready. Remainder stays pooled.""" + if not self.ready(): + raise RuntimeError('draw() called while pool not ready; guard with ready()') + if not self.balanced: + batch = [self._q.popleft() for _ in range(self.batch_size)] + else: + half = self.batch_size // 2 + batch = [self._pos.popleft() for _ in range(half)] + batch += [self._neg.popleft() for _ in range(half)] + self._emitted += len(batch) + return batch + + def draw_all_ready(self) -> List[List[Dict[str, Any]]]: + """Pop as many full batches as currently available (0 or more).""" + out = [] + while self.ready(): + out.append(self.draw()) + return out diff --git a/cookbook/exp/skill2lora/skill_ablate/rollouting.py b/cookbook/exp/skill2lora/skill_ablate/rollouting.py new file mode 100644 index 000000000..bd2526c98 --- /dev/null +++ b/cookbook/exp/skill2lora/skill_ablate/rollouting.py @@ -0,0 +1,143 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Rollout / prompt primitives for the ablation package. + +Reuses v2 verbatim (imported, never edited): +- ``_run_samples`` sampler pad / take-first-N (Ray dp size rule), +- ``_skillgen_prompt`` query-only skill-gen (view B + all eval), +- ``_train_trajectory`` query-only train trajectory (view B, SFT samples, OPSD student), +- ``build_skill_solve_prompt`` / ``build_direct_prompt`` executor prompts, +- ``_parse_seq`` / ``_extract_skill`` / ``_clean_text`` / ``_answer_leaked`` / ``_empty_roll``, +- ``_regen_prompt`` improve-skill regeneration (has orig skill; view A improve_sft), +- style/thinking globals ``_SKILL_STYLE`` / ``SKILL_GEN_SYSTEM`` etc. + +Adds ONLY the view-A rubric-conditioned pieces that v2 lacks: +- ``rubric_skillgen_prompt(problem, rubric)``: skill-gen conditioned on query + rubric + diagnosis, NO prior skill (matches the RL flow step 3 "输入 query+rubric ... skillmodel + rollout"). Style-matched (narrative / pitfall) to the main line. +- ``rubric_train_trajectory(rec)``: train trajectory that REBUILDS the query+rubric prompt + + response, for view-A RL training only (train-with-rubric; the trajectory must match the + prompt the skills were SAMPLED under). The SFT side keeps the query-only + ``_train_trajectory`` so the SFT prompt口径 stays query-only (skill_quality_analysis.md #6). +- ``opsd_teacher_trajectory(rec)``: the OPSD teacher — EXACTLY the student's query-only + prompt with the rubric APPENDED TO THE SYSTEM prompt (设计 871 行 "将 rubric 信息额外加入到 + system prompt 中"). Teacher and student therefore differ ONLY by the appended rubric block + and score the SAME response tokens. (``rubric_skillgen_prompt`` is NOT used here: its + system prompt differs wholesale from ``SKILL_GEN_SYSTEM``, which would confound the + distillation signal with a prompt-style shift.) +""" +from typing import Any, Dict + +import train_skill_v2 as v2 +from train_skill_v2 import ( # noqa: F401 (re-exported for methods.py convenience) + _answer_leaked, + _clean_text, + _empty_roll, + _extract_skill, + _parse_seq, + _regen_prompt, + _run_samples, + _skillgen_prompt, + _train_trajectory, + build_direct_prompt, + build_skill_solve_prompt, +) + +# --- view-A rubric-conditioned skill-gen system prompts -------------------------------- +# 中文注释：view-A 的 rubric 条件 skill-gen 提示词。仿照 skill_quality_analysis.md 的 +# "rubric & skill" 模板，但去掉"你已经生成过一个 skill"的指涉（RL 线首步没有旧 skill，只有 +# query + 一个失败尝试的 rubric 诊断）。文体与主链路一致（narrative / pitfall），且强制"自持、 +# 不指向外部上下文"，因为下游 executor 看不到 rubric——指涉会导致幻觉。改进skill+sft 线另有旧 +# skill，走 v2 的 _regen_prompt（含 orig_skill 字段，即 777-797 模板的逐字英文版），不用这里的提示词。 +# narrative 版末尾拼入与 REGEN_SYSTEM 同一个 few-shot 例子（设计 793-796），锁定文体/长度 +# 分布与主链路一致；程序化提取而非复制，保证与冻结的 v2 逐字相同。 +_REGEN_EXAMPLE = v2.REGEN_SYSTEM.split('Example:\n', 1)[1] +assert _REGEN_EXAMPLE.startswith('') and _REGEN_EXAMPLE.rstrip().endswith(''), \ + 'REGEN_SYSTEM example extraction broke; check v2.REGEN_SYSTEM formatting' + +_RUBRIC_SKILLGEN_NARRATIVE = """\ +You are a skill-generation model. Your block will be fed to a SEPARATE downstream executor model that must solve the problem on its own. The executor will NOT see your private reasoning or the analysis below — it only sees what is inside .... + +An expert rubric analysis of a failed attempt on THIS problem is provided to you. Use it to understand where solving this type of problem tends to break down, then think privately and abstract WHAT MAKES THIS TYPE OF PROBLEM SOLVABLE into transferable methodology. + +Then write the block following these rules: +- Give general, transferable solving techniques for this TYPE of problem as one coherent analysis narrative: first name what the problem is essentially asking, then walk through how to approach it, blending the key concepts, the recommended steps, the pitfalls to avoid (informed by the analysis) and a brief reason for each into a single connected story. +- CRITICAL: Do NOT solve the problem, reveal/compute the final answer, or substitute the problem's specific given numbers. Leave ALL concrete numbers for the executor to compute. +- Self-contained: write in the first person (e.g. "I think the step most likely to go wrong is ..."). NEVER reference "the analysis", "the rubric", or "the previous attempt" — the executor cannot see them, such phrasings cause hallucination. +- Keep it concise: aim for roughly one focused paragraph. + +Put ONLY the methodology inside . + +Example: +""" + _REGEN_EXAMPLE + +_RUBRIC_SKILLGEN_PITFALL = """\ +You are a skill-generation model. A separate executor model will solve the problem; it only sees your block, NOT the analysis below. + +An expert rubric analysis of a failed attempt on THIS problem is provided. From it, pinpoint the single decisive way a solver goes wrong on this type of problem. Then, inside , write under 90 words: +- WARNING: name that decisive mistake concretely, in self-contained first person (e.g. "I think the step most likely to go wrong is ..."), and say why it is wrong. +- INSTEAD: one or two sentences pointing to the correct turn (technique name + where to apply it), without solving the problem or revealing any numeric result. +- End with: "Avoid re-checking loops; box a bare number as soon as it is computed." +Hard rules: the block must be self-contained — never reference "the analysis", "the rubric" or "the previous attempt"; the executor cannot see them. +""" + +_RUBRIC_SKILLGEN_USER = """\ +Problem: +{problem} + +Expert rubric analysis of a failed attempt (for your eyes only; do NOT reference it in the skill): +{rubric} + +Now write the improved guidance:""" + + +def rubric_skillgen_prompt(problem: str, rubric: str) -> Dict[str, Any]: + """View-A skill-gen conditioned on (problem + rubric diagnosis), NO prior skill. + + Style-matched to the main line via v2's ``_SKILL_STYLE`` global (set by main() from + ``--skill-style``). narrative -> narrative rubric prompt; pitfall -> pitfall rubric prompt. + """ + sys_p = _RUBRIC_SKILLGEN_PITFALL if v2._SKILL_STYLE == 'pitfall' else _RUBRIC_SKILLGEN_NARRATIVE + return {'messages': [ + {'role': 'system', 'content': sys_p}, + {'role': 'user', 'content': _RUBRIC_SKILLGEN_USER.format(problem=problem, rubric=rubric)}]} + + +def rubric_train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: + """Train trajectory whose PROMPT is the query+rubric skill-gen prompt + the response. + + Used by view-A RL (rl_ab / rl_err) only: train WITH rubric in the prompt (knowledge- + transfer probe; eval is still query-only via v2 ``_skillgen_prompt``). The rebuilt prompt + matches the prompt the skills were SAMPLED under (on-policy consistency). + ``rec`` must carry 'problem', 'rubric' and 'response'. ``key_rounds`` marks the final + assistant turn as the only trainable span (identical convention to v2 ``_train_trajectory``). + """ + msgs = rubric_skillgen_prompt(rec['problem'], rec.get('rubric', ''))['messages'] + return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], + 'user_data': {'key_rounds': [len(msgs)]}} + + +# 中文注释：OPSD teacher 的特权信息块——按设计 871 行要求放进 SYSTEM prompt，且只做“追加”， +# 保证 teacher 与 student 的 prompt 仅差这一段 rubric（最小差异，蒸馏信号不混入提示词风格漂移）。 +_OPSD_TEACHER_SUFFIX = """ + +[Privileged context — an expert rubric analysis of a failed attempt on this problem. \ +It is visible ONLY to you in this forward pass; the downstream executor never sees it. \ +Use it to judge which guidance actually helps, but do NOT reference it explicitly.] +{rubric}""" + + +def opsd_teacher_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: + """OPSD teacher trajectory: student's query-only prompt + rubric appended to the SYSTEM + prompt + the SAME response. With an empty rubric the teacher degenerates to the student + (zero distillation pull), which is the safe behaviour on rubric API failure.""" + msgs = [dict(m) for m in _skillgen_prompt(rec['problem'])['messages']] + rubric = (rec.get('rubric') or '').strip() + if rubric and msgs[0]['role'] == 'system': + msgs[0]['content'] = msgs[0]['content'] + _OPSD_TEACHER_SUFFIX.format(rubric=rubric) + return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], + 'user_data': {'key_rounds': [len(msgs)]}} + + +def query_only_train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: + """Alias for v2's query-only train trajectory (view B, SFT samples, OPSD student).""" + return _train_trajectory(rec) diff --git a/cookbook/exp/skill2lora/skill_ablate/rubric_cache.py b/cookbook/exp/skill2lora/skill_ablate/rubric_cache.py new file mode 100644 index 000000000..205f829ec --- /dev/null +++ b/cookbook/exp/skill2lora/skill_ablate/rubric_cache.py @@ -0,0 +1,91 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Rubric double-cache (skill_quality_analysis.md #4, #17). + +Two cache scopes, both wrapping v2's ``DiskCache`` (append-only jsonl, in-memory index): + +- GlobalRubricCache (key = data_id): the RL / SFT lines diagnose the BARE-PROBLEM greedy + trajectory. Because the executor is frozen at T=0, that trajectory is deterministic and + identical across experiments, so its rubric diagnosis can be shared across ALL runs via one + global file (``rubric_cache_global.jsonl``) — diagnose each problem once, reuse everywhere. + +- LocalRubricCache (key = md5(data_id + skill)): the improve-skill+SFT / OPSD lines diagnose a + WITH-SKILL trajectory whose skill evolves with the policy, so the diagnosis is experiment- + and step-specific and lives only in that experiment's directory. + +Both reuse v2's ``_diagnose_entry`` (pure teacher-API call, no GPU) for the actual diagnosis, +so there is a single source of truth for the rubric prompt / parsing. +""" +import os +from typing import Any, Dict, Optional + +from train_skill_v2 import DiskCache, _diagnose_entry + + +class _BaseRubricCache: + """Shared get-or-diagnose logic over a DiskCache; subclasses define the key.""" + + def __init__(self, path: str, enabled: bool = True): + self._cache = DiskCache(path, enabled) + + def _key(self, entry: Dict[str, Any], skill: Optional[str]) -> str: + raise NotImplementedError + + def get(self, entry: Dict[str, Any], skill: Optional[str] = None) -> Optional[str]: + return self._cache.get(self._key(entry, skill)) + + def get_or_diagnose(self, entry: Dict[str, Any], checker, + skill: Optional[str] = None) -> str: + """Return cached diagnosis, else run the teacher rubric once and cache it. + + ``entry`` must carry the fields ``_diagnose_entry`` needs: ``problem``, + ``fail_segment``, ``fail_stop_reason`` (and ``reference_answer`` is unused by the + diagnosis but kept for auditing). Returns '' when there is no checker or on API error + (never raises), so the caller can treat "no diagnosis" uniformly. + """ + if checker is None: + return '' + key = self._key(entry, skill) + hit = self._cache.get(key) + if hit is not None: + return hit + diag = _diagnose_entry(checker, entry) or '' + self._cache.put(key, diag) + return diag + + def put(self, entry: Dict[str, Any], diag: str, skill: Optional[str] = None) -> None: + self._cache.put(self._key(entry, skill), diag) + + def __contains__(self, key: str) -> bool: + return key in self._cache + + def close(self) -> None: + self._cache.close() + + +class GlobalRubricCache(_BaseRubricCache): + """key = data_id: bare-problem trajectory diagnosis, shareable across experiments.""" + + def _key(self, entry: Dict[str, Any], skill: Optional[str] = None) -> str: + return DiskCache.key_for('rubric_global', str(entry.get('data_id', ''))) + + +class LocalRubricCache(_BaseRubricCache): + """key = md5(data_id + skill): with-skill trajectory diagnosis, per-experiment only.""" + + def _key(self, entry: Dict[str, Any], skill: Optional[str] = None) -> str: + return DiskCache.key_for('rubric_local', str(entry.get('data_id', '')), skill or '') + + +def build_rubric_cache(scope: str, output_dir: str, + global_dir: Optional[str] = None, enabled: bool = True): + """Factory: scope='global' -> shared file under ``global_dir`` (default output_dir/..); + scope='local' -> per-experiment file under ``output_dir/cache``.""" + if scope == 'global': + base = global_dir or os.path.dirname(os.path.abspath(output_dir.rstrip('/'))) + os.makedirs(base, exist_ok=True) + return GlobalRubricCache(os.path.join(base, 'rubric_cache_global.jsonl'), enabled) + if scope == 'local': + cache_dir = os.path.join(output_dir, 'cache') + os.makedirs(cache_dir, exist_ok=True) + return LocalRubricCache(os.path.join(cache_dir, 'rubric_cache_local.jsonl'), enabled) + raise ValueError(f"scope must be 'global' or 'local', got {scope!r}") diff --git a/cookbook/exp/skill2lora/skill_ablate/trainer.py b/cookbook/exp/skill2lora/skill_ablate/trainer.py new file mode 100644 index 000000000..0d313745a --- /dev/null +++ b/cookbook/exp/skill2lora/skill_ablate/trainer.py @@ -0,0 +1,344 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Unified training loop for one ablation experiment. + +step unit = PARAMETER UPDATE (skill_quality_analysis.md #2): a chunk may yield 0/1/more +updates; we stop at ``--max-updates`` and eval every ``--eval-every-updates`` updates. All +eval is query-only via v2 ``run_greedy_eval`` (T=0.5 × 4 rollouts, no rubric) — the +knowledge-transfer probe for view A. Records/log schema mirror v2 (gen/eval/train_log jsonl, +per-problem rollout rows land in gen_records.jsonl). + +swanlab step axis: ONE global axis = chunk index for train AND eval curves (eval also logs +``eval/updates_done`` so the update count is recoverable); mixing chunk/update axes in one +experiment made curves incomparable. +""" +import json +import os +import sys +import time +from typing import Any + +import train_skill_v2 as v2 + +from .config import ExpSpec +from .data import load_deepmath_records +from .methods import MethodContext, build_method +from .pool import SamplePool +from .rubric_cache import build_rubric_cache + +try: + import swanlab +except ImportError: + swanlab = None + + +def _rubric_scope(method: str) -> str: + """Bare-problem lines (rl/sft) share a GLOBAL cache; with-skill lines (opsd/improve) use + a per-experiment LOCAL cache. bnpo needs no rubric.""" + if method in ('rl_ab', 'rl_err', 'sft'): + return 'global' + if method in ('opsd', 'improve_sft'): + return 'local' + return '' + + +def _build_pool(spec: ExpSpec, args) -> Any: + if spec.method == 'improve_sft': + return SamplePool(batch_size=args.sft_batch_size, balanced=True, + max_pool=args.pool_max) + if spec.method == 'sft': + return SamplePool(batch_size=args.sft_batch_size, balanced=False, + max_pool=args.pool_max) + return None # RL / OPSD / bnpo train per-chunk, no accumulation pool + + +def _load_resume_state(args) -> dict: + """Resolve --resume-from into {'ckpt_dir', 'updates', 'chunk_idx'}. + + Weights-only resume (lr is constant; Adam moments restart — accepted trade-off). + Counter source: /train_state.json (written by _save_ckpt); legacy finished runs + (e.g. E1-final) fall back to /DONE.json {'updates','chunks'}. + """ + ck = args.resume_from + if not os.path.isdir(ck): + ck = os.path.join(args.output_dir, args.resume_from) + if not os.path.isdir(ck): + raise FileNotFoundError(f'--resume-from checkpoint dir not found: {args.resume_from}') + state_path = os.path.join(ck, 'train_state.json') + done_path = os.path.join(args.output_dir, 'DONE.json') + if os.path.exists(state_path): + with open(state_path, encoding='utf-8') as f: + st = json.load(f) + for k in ('chunk_size', 'min_level', 'n', 'seed'): + if k in st and getattr(args, k, None) not in (None, '') and st[k] != getattr(args, k): + sys.stderr.write(f'[ablate] WARNING: resume data config mismatch: {k} ' + f'ckpt={st[k]} vs now={getattr(args, k)} — the continued chunk ' + f'sequence will NOT align with the original run.\n') + elif os.path.exists(done_path): + with open(done_path, encoding='utf-8') as f: + d = json.load(f) + st = {'updates': int(d['updates']), 'chunk_idx': int(d['chunks'])} + sys.stderr.write('[ablate] resume: no train_state.json in ckpt, counters restored ' + 'from DONE.json (legacy run — data-config alignment unverified).\n') + else: + raise FileNotFoundError(f'resume: neither {state_path} nor {done_path} exists; ' + 'cannot restore update/chunk counters.') + return {'ckpt_dir': ck, 'updates': int(st['updates']), 'chunk_idx': int(st['chunk_idx'])} + + +def run_experiment(args, spec: ExpSpec) -> None: + # 0) idempotency: DONE.json is written atomically as the very last step of a successful + # run; if present, this experiment is complete -> skip (unless --force). --resume-from + # bypasses the guard by design: its whole point is extending a finished run. + resume = _load_resume_state(args) if getattr(args, 'resume_from', '') else None + done_path = os.path.join(args.output_dir, 'DONE.json') + if os.path.exists(done_path) and not getattr(args, 'force', False) and resume is None: + sys.stderr.write(f'[ablate] {spec.name} already complete ({done_path}); ' + f'use --force to rerun.\n') + return + if resume is not None: + args.skill_init_model_id = resume['ckpt_dir'] # v2.init_components skill_model bypass + sys.stderr.write(f'[ablate] resuming {spec.name} from {resume["ckpt_dir"]} ' + f'(updates={resume["updates"]} chunk={resume["chunk_idx"]}).\n') + + # 1) style / align globals must be set BEFORE any prompt is built. + v2._ALIGN_MODE = spec.align + v2._SKILL_STYLE = spec.style + args.skill_thinking = spec.thinking + # per-style length budget (#9 statistics: narrative≈1100 / pitfall≈300 chars); + # an explicit --len-budget on the CLI wins. + if args.len_budget is None: + args.len_budget = 1100 if spec.style == 'narrative' else 300 + # explicit --skill-max-tokens on the CLI wins over the per-experiment default. + if args.skill_max_tokens is None: + args.skill_max_tokens = spec.skill_max_tokens + elif args.skill_max_tokens != spec.skill_max_tokens: + sys.stderr.write(f'[ablate] WARNING: --skill-max-tokens {args.skill_max_tokens} ' + f'overrides the {spec.name} default {spec.skill_max_tokens}.\n') + # OOM guard: think/8192 实验的训练序列长一倍，fp32 主权重下 16/2卡 的 micro backward + # 会爆显存（E13 实测 Tried to allocate 37.9GiB）；自动把 micro 减半到 8，梯度归一后数学等价， + # 攒批/采样批（sft_batch_size）冻结口径不变。显式 --train-micro-batch 优先。 + if not args.train_micro_batch and args.skill_max_tokens >= 8192: + args.train_micro_batch = max(v2.TRAIN_DP, args.sft_batch_size // 2) + sys.stderr.write(f'[ablate] train_micro_batch auto-set to {args.train_micro_batch} ' + f'(skill_max_tokens={args.skill_max_tokens} OOM guard).\n') + + records, eval_records = (load_deepmath_records(args) if getattr(args, 'deepmath_dir', '') + else v2._load_records(args)) + if len(records) < args.chunk_size: + raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded ({len(records)})') + + os.makedirs(args.output_dir, exist_ok=True) + gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') + eval_path = os.path.join(args.output_dir, 'eval_records.jsonl') + train_log_path = os.path.join(args.output_dir, 'train_log.jsonl') + + # 2) rubric checker BEFORE any GPU allocation: view A without a teacher API cannot run + # (SFT-family would loop forever on an empty pool, OPSD would distill on empty rubrics). + checker = v2.build_rubric_checker() if spec.needs_rubric else None + if spec.needs_rubric and checker is None: + raise RuntimeError( + f'{spec.name} ({spec.method}) is a view-A experiment and REQUIRES the rubric ' + 'teacher API; set LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL (or OPENAI_API_KEY).') + + # 3) components (v2 verbatim); override loss to OPSD when needed. + skill_model, ref_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp = \ + v2.init_components(args) + if spec.loss == 'opsd': + # no beta: OPSDLoss uses only teacher_logps (no ref-KL term, no ref forward at all) + skill_model.set_loss('OPSDLoss', reverse=True) + + scope = _rubric_scope(spec.method) + rubric_cache = build_rubric_cache(scope, args.output_dir, + global_dir=args.rubric_global_dir, + enabled=not args.no_cache) if scope else None + + # client-side Template clone (OPSD only): encodes student/teacher trajectories locally to + # extract the teacher's RESPONSE-ONLY logp positions (prompts differ in length, so the + # remote full-sequence logps must be sliced before OPSDLoss can align them per token). + # Mirrors init_components' set_template call exactly (same tokenizer/thinking/truncation). + encode_template = None + if spec.method == 'opsd': + encode_template = v2.Template(model_id=v2.MODEL_ID, + enable_thinking=(spec.thinking == 'on'), + max_length=args.max_model_len, + truncation_strategy='delete') + + pool = _build_pool(spec, args) + ctx = MethodContext(skill_model=skill_model, ref_model=ref_model, skill_sampler=skill_sampler, + base_sampler=base_sampler, ckpt=ckpt, skill_dp=skill_dp, base_dp=base_dp, + args=args, checker=checker, rubric_cache=rubric_cache, pool=pool, + encode_template=encode_template) + method = build_method(spec.method, ctx) + + def _save_ckpt(name: str, updates: int, chunk_idx: int, epoch: int) -> None: + """Weights-only checkpoint + barrier + resume state. + + skill_model.save dispatches to the train actors; a subsequent cheap blocking call on + the SAME actors (lr_step — a guaranteed no-op here: no scheduler, constant lr) acts as + the barrier: Ray actor tasks run serially per actor, so when it returns the save has + landed on every rank. Without it the driver could exit on a half-written safetensors. + """ + skill_model.save(name, output_dir=args.output_dir) + skill_model.lr_step() # barrier (see docstring) + state = {'updates': updates, 'chunk_idx': chunk_idx, 'epoch': epoch, + 'seed': args.seed, 'chunk_size': args.chunk_size, 'n': args.n, + 'min_level': int(getattr(args, 'min_level', 0) or 0), + 'lr': args.lr, 'exp': spec.name, 'saved': int(time.time())} + with open(os.path.join(args.output_dir, name, 'train_state.json'), 'w', + encoding='utf-8') as f: + json.dump(state, f) + sys.stderr.write(f'[ablate] checkpoint saved: {name} (updates={updates})\n') + + # 4) eval baseline cache (v2 DiskCache) + swanlab. + # 每次启动强制重算 eval baseline：旧缓存可能来自不同环境/代码版本（torch/vllm/dtype 均影响 T=0 输出）， + # 跨 run 复用会造成 with-skill（现算）vs baseline（陈旧）不可比，lift 虚高/虚低。 + cache_dir = os.path.join(args.output_dir, 'cache') + os.makedirs(cache_dir, exist_ok=True) + _base_cache_path = os.path.join(cache_dir, 'eval_baseline.jsonl') + if os.path.exists(_base_cache_path): + os.remove(_base_cache_path) + sys.stderr.write('[ablate] stale eval_baseline cache removed (recomputed this run).\n') + eval_base_cache = v2.DiskCache(_base_cache_path, not args.no_cache) + + use_swan = swanlab is not None and os.environ.get('SWANLAB_MODE') != 'disabled' + if use_swan: + # timestamp suffix so FORCE reruns never collide in swanlab (接口方案 #9) + swan_exp = f'{spec.swanlab_exp}_{time.strftime("%Y%m%d_%H%M%S")}' + swanlab.init(project=args.swanlab_project, experiment_name=swan_exp, + config={'exp': spec.name, 'view': spec.view, 'method': spec.method, + 'thinking': spec.thinking, 'style': spec.style, 'align': spec.align, + 'loss': spec.loss, 'skill_max_tokens': args.skill_max_tokens, + 'max_updates': args.max_updates, 'lr': args.lr, + 'n_skills': args.n_skills, 'sft_batch_size': args.sft_batch_size, + 'len_budget': args.len_budget, + 'drop_zero_adv': args.drop_zero_adv}) + + cfg = {'record_type': 'config', 'exp': spec.name, 'view': spec.view, 'method': spec.method, + 'thinking': spec.thinking, 'style': spec.style, 'align': spec.align, 'loss': spec.loss, + 'skill_max_tokens': args.skill_max_tokens, 'needs_rubric': spec.needs_rubric, + 'rubric_scope': scope, 'rubric_check': bool(checker), + 'n': len(records), 'eval_n': len(eval_records), 'model': v2.MODEL_ID, + 'max_updates': args.max_updates, 'eval_every_updates': args.eval_every_updates, + 'lr': args.lr, 'n_skills': args.n_skills, 'chunk_size': args.chunk_size, + 'sft_batch_size': args.sft_batch_size, 'len_budget': args.len_budget, + 'skill_char_limit': args.skill_char_limit, 'drop_zero_adv': args.drop_zero_adv, + 'improve_skill_temperature': args.improve_skill_temperature, + 'eval_rollouts': args.eval_rollouts, 'eval_skill_temperature': args.eval_skill_temperature, + 'seam_parquet_dir': (getattr(args, 'seam_parquet_dir', '') or ''), + 'deepmath_dir': (getattr(args, 'deepmath_dir', '') or ''), + 'min_level': int(getattr(args, 'min_level', 0) or 0), + 'save_every_updates': int(getattr(args, 'save_every_updates', 0) or 0), + 'resumed_from': (resume['ckpt_dir'] if resume else ''), + 'resumed_updates': (resume['updates'] if resume else 0), + 'started': int(time.time())} + + # resume appends to the record files (the original history stays intact); a fresh run + # truncates as before. + _fmode = 'a' if resume is not None else 'w' + with open(gen_path, _fmode, encoding='utf-8') as gen_f, \ + open(eval_path, _fmode, encoding='utf-8') as eval_f, \ + open(train_log_path, _fmode, encoding='utf-8') as tlog: + for f in (gen_f, eval_f, tlog): + v2._write(f, cfg) + + def _do_eval(updates_done: int, swan_step: int) -> None: + recs, summary, metrics = v2.run_greedy_eval( + base_sampler, skill_sampler, eval_records, updates_done, updates_done, + base_dp, skill_dp, args, eval_base_cache) + for rec in recs: + v2._write(eval_f, rec) + v2._write(eval_f, summary) + eval_f.flush() + if use_swan: + # same chunk-based axis as the train curves; updates recoverable via the + # logged eval/updates_done scalar. + swanlab.log({**{f'eval/{k}': v for k, v in metrics.items()}, + 'eval/updates_done': float(updates_done)}, step=swan_step) + sys.stderr.write( + f'[eval] u{updates_done}: n={summary["n"]} acc={summary["baseline_acc_mean1"]:.3f}' + f'->{summary["acc_mean1"]:.3f} lift={summary["lift_mean1"]:+.3f} ' + f'hard_rescue={summary["hard_rescue_rate"]:.3f} fmt={summary["format_mean1"]:.2f}\n') + + if eval_records and resume is None: + _do_eval(-1, 0) # baseline before any update (chunk axis position 0) + + pool_pp = v2.ProblemPool(records, args.seed) + updates = 0 + last_eval_at = 0 + last_eval_updates = -1 + chunk_idx = 0 + last_save_at = 0 + if resume is not None: + # push the resumed weights into skill_sampler BEFORE the first chunk (otherwise + # skill-gen would sample from base weights until the first post-update sync). + ckpt.sync_weights(merge_and_sync=True) + updates = resume['updates'] + last_eval_at = updates + last_save_at = updates + # fast-forward the pool: draws are deterministic (RandomState(seed+epoch)), so + # replaying chunk_idx draws restores the exact data position — IF chunk_size / + # data config match the original run (warned in _load_resume_state). + for _ in range(resume['chunk_idx']): + pool_pp.draw(args.chunk_size) + chunk_idx = resume['chunk_idx'] + save_every = int(getattr(args, 'save_every_updates', 0) or 0) + while updates < args.max_updates: + chunk = pool_pp.draw(args.chunk_size) + res = method.step(chunk, chunk_idx) + n_upd = int(res.get('n_updates', 0)) + updates += n_upd + + for rec in res.get('gen_records') or []: + v2._write(gen_f, rec) + gen_f.flush() + + log = {'record_type': 'train_round', 'exp': spec.name, 'chunk': chunk_idx, + 'epoch': pool_pp.epoch, 'updates': updates, 'n_updates_step': n_upd, + 'method': spec.method, 'ts': int(time.time()), + 'metrics': res.get('metrics', {})} + if 'summary' in res: # bnpo carries the v2 chunk summary + log['summary'] = res['summary'] + v2._write(tlog, log) + tlog.flush() + + sys.stderr.write(f'[gen] e{pool_pp.epoch} c{chunk_idx}: +{n_upd}upd ' + f'total={updates}/{args.max_updates} ' + + ' '.join(f'{k}={v:.3g}' for k, v in res.get('metrics', {}).items() + if isinstance(v, (int, float))) + '\n') + if use_swan: + m = {f'{k}': float(v) for k, v in res.get('metrics', {}).items() + if isinstance(v, (int, float))} + m['train/updates'] = float(updates) + m['train/n_updates_step'] = float(n_upd) + swanlab.log(m, step=chunk_idx + 1) # +1: step 0 is the eval baseline + + # eval by parameter-update cadence (same chunk-based swan axis) + if eval_records and updates >= last_eval_at + args.eval_every_updates and n_upd > 0: + _do_eval(updates, chunk_idx + 1) + last_eval_at = updates + last_eval_updates = updates + # periodic weights-only checkpoint (same cadence semantics as eval) + if save_every and updates >= last_save_at + save_every and n_upd > 0: + _save_ckpt(f'{spec.name}-u{updates}', updates, chunk_idx + 1, pool_pp.epoch) + last_save_at = updates + chunk_idx += 1 + + # final readout — skip if the periodic eval already covered this exact update count + if eval_records and updates != last_eval_updates: + _do_eval(updates, chunk_idx + 1) + + eval_base_cache.close() + if rubric_cache is not None: + rubric_cache.close() + # final save goes through _save_ckpt: barrier guarantees the safetensors is fully on disk + # BEFORE DONE.json can exist, and train_state.json makes the final model resumable too. + _save_ckpt(f'{spec.name}-final', updates, chunk_idx, pool_pp.epoch) + # completion sentinel LAST (after the final model lands); temp+rename keeps it atomic so a + # crash can never leave a truthy half-written marker. + tmp = done_path + '.tmp' + with open(tmp, 'w', encoding='utf-8') as f: + json.dump({'exp': spec.name, 'updates': updates, 'chunks': chunk_idx, + 'epochs': pool_pp.epoch, 'finished': int(time.time())}, f) + os.replace(tmp, done_path) + sys.stderr.write(f'[ablate] {spec.name} done: {updates} updates over {chunk_idx} chunks / ' + f'{pool_pp.epoch} epochs\n') diff --git a/cookbook/exp/skill2lora/train_skill_v2.py b/cookbook/exp/skill2lora/train_skill_v2.py index 1d189f680..3b9cc8313 100644 --- a/cookbook/exp/skill2lora/train_skill_v2.py +++ b/cookbook/exp/skill2lora/train_skill_v2.py @@ -331,18 +331,12 @@ def _extract_skill(text: str) -> Optional[str]: def _parse_seq(seq, gold: str) -> Dict[str, Any]: text = _clean_text(getattr(seq, 'decoded', '') or '') - if _ALIGN_MODE == 'seam': - # seam：SEAM lpem parity——整段贪婪 sanitize 成单一数值后精确匹配。 - san = _seam_sanitize(text) - pred = san or None - correct = bool(san) and (san == _seam_sanitize(str(gold))) - else: - # v2：只从 \boxed{} 抽取（executor 被要求把最终数值写进 \boxed{}），再走同一套数值归一 - # （frac/inline/number）后精确匹配；不做“整段抓首个数字”的贪婪回退（避免从推理里误抓）。 - # extract_boxed 取最后一个配平的 \boxed{}、截断（未闭合）时不误取；都没有则判错。 - raw = extract_boxed(text) - pred = _seam_sanitize(raw) if raw else None - correct = bool(pred) and (pred == _seam_sanitize(str(gold))) + # 判分口径统一（人工拍板，2026-07-27）：seam/v2 都只从 \boxed{} 抽取，再走同一套数值归一 + # （frac/inline/number）后精确匹配；不做 lpem 式“整段抓数字”贪婪回退，保证 E13 与 E1-E12 + # 的 acc/lift 横向可比。extract_boxed 取最后一个配平的 \boxed{}、截断时不误取；没有则判错。 + raw = extract_boxed(text) + pred = _seam_sanitize(raw) if raw else None + correct = bool(pred) and (pred == _seam_sanitize(str(gold))) terminated = getattr(seq, 'stop_reason', None) != 'length' return {'pred': pred, 'correct': correct, 'terminated': terminated, 'stop_reason': getattr(seq, 'stop_reason', None), @@ -710,18 +704,17 @@ def _diagnose_entry(checker, entry: Dict[str, Any]) -> Optional[str]: _SKILL_STYLE = 'narrative' # 'narrative' | 'toy' | 'pitfall'；由 main() 依据 --skill-style 设置 # ---- Executor prompt (with skill injection) ---- -# 中文注释：executor 提示词。答案格式对齐 SEAM 的 slove_qwen.txt（numeric-only，降截断）。 -# v2 模式：新版采用英文单 user turn——题目 + “技巧提示（skill 作为 advisory）” + 答案格式（见 -# build_skill_solve_prompt）。v2 执行器输出用 \boxed{}（_ANSWER_FORMAT_V2），先不用 ； -# seam 基线 DIRECT_SYSTEM 仍用 （_ANSWER_FORMAT），对齐 SEAM。skill 不再注入 system。 +# 中文注释：executor 提示词。答案格式已统一为 \boxed{}（人工拍板，2026-07-27）：seam/v2 两模式的 +# 格式说明与判分口径完全一致，保证 E13(seam) 与 E1-E12 横向可比；prompt 结构差异（seam 嵌套/ +# system+user vs v2 单 user）作为方案级差异保留。_ANSWER_FORMAT( 版) 已弃用。 _ANSWER_FORMAT = ('Present your reasoning and answer in the following format:\n' ' Content of Thinking[Final numeric result only]') -# v2 执行器答案格式：把最终数值放进 \boxed{}（不再要求 ）；判分对应走 extract_boxed。 +# 统一执行器答案格式：把最终数值放进 \boxed{}；判分对应走 extract_boxed。 _ANSWER_FORMAT_V2 = ('Present your reasoning, then put ONLY the final numeric result inside ' '\\boxed{}. For example: \\boxed{42}.') DIRECT_SYSTEM = ( 'You are an expert competition mathematician. Be concise and accurate. ' - + _ANSWER_FORMAT) + + _ANSWER_FORMAT_V2) _SKILL_SOLVE_PREFIX = ( 'You are an expert competition mathematician. Be concise and accurate.\n\n' 'Before you start, keep these reminders in mind to avoid common mistakes on this ' @@ -758,8 +751,7 @@ def _diagnose_entry(checker, entry: Dict[str, Any]) -> Optional[str]: 'prefer using its techniques when they fit, but you may use alternative correct methods ' 'if they are more efficient or clearer. If you diverge from the advisory context, briefly ' 'explain why. Be concise and accurate.\n' - 'Present your reasoning and answer in the following format:\n' - ' Content of Thinking[Final numeric result only]') + + _ANSWER_FORMAT_V2) def build_skill_solve_prompt_seam(problem, skill, raw_response=None): @@ -1135,7 +1127,10 @@ def _train_step(skill_model, ref_model, ckpt, samples, args): return {'n_samples': n_in, 'n_sft': 0, 'n_grpo': 0, 'n_empty': n_empty, 'n_steps': 0, 'n_micro_batches': 0, 'metric': {}} trajs, advs = trajs[:n_keep], advs[:n_keep] - n, sft = len(trajs), args.sft_batch_size + # micro 尺寸可由 train_micro_batch 覆盖（ablate think/8192 实验防 backward OOM）； + # 默认 0 = 跟随 sft_batch_size，主训练行为不变。梯度按 micro 数归一，切细数学等价。 + n = len(trajs) + sft = getattr(args, 'train_micro_batch', 0) or args.sft_batch_size mini = args.ppo_mini_batch_size if args.ppo_mini_batch_size > 0 else n mini = max(sft, (mini // sft) * sft) multi_step = mini < n @@ -1409,7 +1404,10 @@ def init_components(args): # 主权重必须 fp32（对齐 verl actor：fp32 master + bf16 autocast）。twinkle 默认不传 dtype 时 # transformers 会按 config 加载 bf16 主权重，lr=1e-6 的更新量(~1e-6)远小于 bf16 ulp(~4e-5)， # optimizer.step 的更新几乎全被舍入吞掉——这是 v2 学不动/与 SEAM 对不上的根因（A/B 实测差 10-20 倍）。 - skill_model = TransformersModel(model_id=MODEL_ID, device_mesh=train_mesh, remote_group='train', + # resume 旁路（skill_ablate）：skill_init_model_id 指向已保存的 checkpoint 目录时，仅 skill_model + # 从该目录初始化（TransformersModel.load 无全量模型路径）；ref/samplers/template 仍用 MODEL_ID。 + _skill_init_id = getattr(args, 'skill_init_model_id', '') or MODEL_ID + skill_model = TransformersModel(model_id=_skill_init_id, device_mesh=train_mesh, remote_group='train', torch_dtype='float32', ddp_config={'find_unused_parameters': False}) skill_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) From 067f808d28249ec4c619d5477ccc7c1256f8d396 Mon Sep 17 00:00:00 2001 From: root Date: Sat, 1 Aug 2026 15:44:17 +0800 Subject: [PATCH 27/60] skill2lora: add E13 SEAM-repro arm (executor nothink, seam align) - config.py: E13 executor_thinking='off' to match the SEAM paper run - run_ablate12.sh: dedicated E13 block (min_level=0 full pool, chunk=128, reward-trunc-penalty=0, eval R=1/T=0) reproducing the SEAM run config - include the code-task/reflexion pipeline modules E13 imports at load time (main/trainer top-level import code_task/data_code/eval_reflexion) --- cookbook/exp/skill2lora/code_task.py | 439 +++++++ cookbook/exp/skill2lora/run_ablate12.sh | 232 +++- .../exp/skill2lora/skill_ablate/config.py | 186 ++- cookbook/exp/skill2lora/skill_ablate/data.py | 24 +- .../exp/skill2lora/skill_ablate/data_code.py | 82 ++ .../skill2lora/skill_ablate/eval_reflexion.py | 208 +++ cookbook/exp/skill2lora/skill_ablate/main.py | 177 ++- .../exp/skill2lora/skill_ablate/methods.py | 1120 ++++++++++++++++- .../exp/skill2lora/skill_ablate/rollouting.py | 91 +- .../skill2lora/skill_ablate/rubric_cache.py | 59 +- .../exp/skill2lora/skill_ablate/trainer.py | 125 +- cookbook/exp/skill2lora/train_skill_v2.py | 227 +++- 12 files changed, 2840 insertions(+), 130 deletions(-) create mode 100644 cookbook/exp/skill2lora/code_task.py create mode 100644 cookbook/exp/skill2lora/skill_ablate/data_code.py create mode 100644 cookbook/exp/skill2lora/skill_ablate/eval_reflexion.py diff --git a/cookbook/exp/skill2lora/code_task.py b/cookbook/exp/skill2lora/code_task.py new file mode 100644 index 000000000..44e93351c --- /dev/null +++ b/cookbook/exp/skill2lora/code_task.py @@ -0,0 +1,439 @@ +"""BigCodeBench task adapter: data / executor prompts / unit-test judging / code rubric. + +为什么存在这个模块（承接 deepmath -> BFCL -> BigCodeBench 三轮 eval-0 探针的结论）： + deepmath —— 76.8% 的失败是"没写完"，救活率几乎全由截断率解释，任何 hint（含乱码）只值一个 + wrapper；rubric 只能说"你超预算了/你在兜圈"，命中率≈随机，rubric skill 增量 −0.056。 + BFCL —— 截断混杂消掉了，但 4B 裸解 0.861 只剩 8% headroom，且最大错误类 43% 是 ground truth + 私有口径问题，judge 看不到答案就无从判断 -> rubric skill 增量 +0.002（零）。 + BigCodeBench —— 判分是跑 unittest：对错是客观的（跑过就是对），而且**失败时机器免费给出可定位的 + 证据**（异常类型 / 断言差异 / 失败用例名）。bcb/bcb_eval0_probe.py 实测（n=274， + nothink，截断 0）：F0_none 0.378、query-only skill 0.382（+0.004）、 + rubric skill 0.513（+0.135，p=4e-5）—— 三个数据集里 rubric 第一次真正有增量。 + 结论（已写入长期记忆）：rubric 有用的前提是"诊断有客观可定位的失败证据"，不是"任务是代码"。 + +本模块只做纯任务逻辑（加载 / prompt / 判分 / rubric 素材），**不 import train_skill_v2**， +所以 v2 可以在模块顶层 import 它而不构成循环依赖。判分口径与 bcb_eval0_probe.py 逐字同源 +（extract_code / _RUNNER / run_tests / _trim_err / spec_constraints 直接搬过来），这样探针读数 +与训练读数可以横比。 +""" +import ast as _ast +import importlib.util +import json +import os +import random +import re +import shutil +import subprocess +import sys +import tempfile +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Optional, Tuple + +_HERE = os.path.dirname(os.path.abspath(__file__)) +DEFAULT_PARQUET = os.path.join(_HERE, '..', '..', '..', 'bigcodebench', 'bcb.parquet') + +# 需要外网 / GUI / 子进程的库：沙箱里会挂或超时，判分噪声与 skill 无关 -> 整题排除。 +EXCLUDE_LIBS = {'requests', 'urllib', 'http', 'smtplib', 'socket', 'ssl', 'ftplib', + 'mechanize', 'wikipedia', 'turtle', 'tkinter', 'subprocess', 'sendgrid', + 'python_http_client', 'django', 'flask', 'flask_login', 'flask_mail', + 'flask_restful', 'flask_wtf', 'wtforms', 'multiprocessing'} +LIB_ALIAS = {'sklearn': 'sklearn', 'cv2': 'cv2', 'PIL': 'PIL', 'bs4': 'bs4', 'yaml': 'yaml', + 'dateutil': 'dateutil', 'Crypto': 'Crypto', 'docx': 'docx', 'pytz': 'pytz', + 'psutil': 'psutil', 'texttable': 'texttable', 'wordcloud': 'wordcloud', + 'skimage': 'skimage', 'PyPDF2': 'PyPDF2'} + + +# =========================================================================== +# 数据 +# =========================================================================== +def _libs(rec) -> List[str]: + v = rec.get('libs') + if isinstance(v, str): + try: + return list(_ast.literal_eval(v)) + except Exception: + return [] + return list(v or []) + + +def _importable(lib: str) -> bool: + try: + return importlib.util.find_spec(LIB_ALIAS.get(lib, lib).split('.')[0]) is not None + except Exception: + return False + + +def load_tasks(path: str, seed: int) -> Tuple[List[Dict[str, Any]], Dict[str, int]]: + """-> (tasks, stats)；tasks 已按 seed 洗牌，每条是一个"判分载荷"（见 payload_of）。""" + import pyarrow.parquet as pq + rows = pq.read_table(path).to_pylist() + keep, drop_missing, drop_excl = [], 0, 0 + for r in rows: + libs = _libs(r) + if set(libs) & EXCLUDE_LIBS: + drop_excl += 1 + continue + if any(not _importable(x) for x in libs): + drop_missing += 1 + continue + keep.append({'task_id': r['task_id'], 'instruct_prompt': r['instruct_prompt'], + 'code_prompt': r['code_prompt'], 'test': r['test'], + 'entry_point': r['entry_point'], 'doc_struct': r['doc_struct'], + 'canonical_solution': r['canonical_solution'], 'libs': libs}) + random.Random(seed).shuffle(keep) + return keep, {'raw': len(rows), 'kept': len(keep), + 'drop_missing_lib': drop_missing, 'drop_needs_net_or_gui': drop_excl} + + +def payload_of(task: Dict[str, Any]) -> Dict[str, Any]: + """训练记录里 ``reference_answer`` 的内容 —— 判分需要的一切。 + + ★ 为什么塞进 reference_answer 而不是新开字段：全流水线（v2 / methods / eval_reflexion）判分 + 都走 ``_parse_seq(seq, r['reference_answer'])`` 这一个入口，把载荷放这里就不用改任何签名， + math 分支也完全不受影响。体积：test 平均 ~3KB，每题每 chunk 落盘一次，可接受。 + """ + return {k: task[k] for k in ('task_id', 'entry_point', 'test', 'code_prompt', 'doc_struct', + 'canonical_solution')} + + +# =========================================================================== +# 代码抽取 + 沙箱跑单测（与 bcb_eval0_probe.py 逐字同源） +# =========================================================================== +def after_think(text: str) -> str: + i = (text or '').rfind('') + return text[i + len(''):] if i >= 0 else (text or '') + + +_FENCE_RE = re.compile(r'```(?:python|py)?\s*\n(.*?)```', re.S) + + +def extract_code(text: str) -> str: + """取最后一个能通过 ast.parse 的代码块；没有围栏就退化为整段（切 think 之后）。""" + body = after_think(text or '') + blocks = _FENCE_RE.findall(body) + for b in reversed(blocks): + try: + _ast.parse(b) + return b + except SyntaxError: + continue + if blocks: + return blocks[-1] + try: + _ast.parse(body) + return body + except SyntaxError: + return '' + + +_RUNNER = """ +import unittest, sys +loader = unittest.TestLoader() +suite = loader.loadTestsFromTestCase(TestCases) +res = unittest.TextTestRunner(verbosity=0, stream=sys.stderr).run(suite) +print('__BCB__', res.testsRun, len(res.failures), len(res.errors)) +sys.exit(0 if res.wasSuccessful() and res.testsRun > 0 else 1) +""" + + +def _trim_err(err: str, limit: int = 1600) -> str: + """保留失败测试名与异常行，砍掉中间冗长的 traceback 帧（这是喂给 rubric 的客观证据）。 + + 随机临时目录名换成 ````：traceback 帧里带着 /tmp/bcb_xxxxxxx/ 这种每次都不同的 + 路径，对 judge 是纯噪声，还会让同一个失败在两次运行里看起来不一样（gen_records 里逐字 + 比对失败原因时会误判成"变了"）。 + """ + err = re.sub(r'/tmp/bcb_[A-Za-z0-9_]+', '', err or '') + lines = [ln for ln in err.splitlines() if ln.strip()] + keep = [ln for ln in lines + if ln.startswith(('FAIL:', 'ERROR:', 'AssertionError', 'Traceback')) + or re.match(r'^\w*(Error|Exception|Warning)\b', ln.strip()) + or ', in ' in ln] + text = '\n'.join(keep or lines[-25:]) + return text[-limit:] + + +def run_tests(code: str, payload: Dict[str, Any], timeout: int) -> Dict[str, Any]: + """在独立进程 + 临时目录里跑该题自带的 unittest；返回 pass/fail + 客观报错。 + + 隔离手段只有"子进程 + 临时 cwd + 超时"，没有容器/seccomp：真正危险的题（外网 / GUI / + 子进程 / multiprocessing）在 load_tasks 阶段就按 EXCLUDE_LIBS 整题剔除了。 + env 里必须清掉 CUDA_VISIBLE_DEVICES —— 否则 numpy/torch 系的测试可能去抢训练用的卡。 + """ + if not code.strip(): + return {'passed': False, 'kind': 'no_code', 'error': 'no parseable code block', + 'n_tests': 0} + if payload['entry_point'] not in code: + return {'passed': False, 'kind': 'no_entry', + 'error': f"function {payload['entry_point']} is not defined in the submitted code", + 'n_tests': 0} + tmp = tempfile.mkdtemp(prefix='bcb_') + try: + src = code + '\n\n' + payload['test'] + '\n' + _RUNNER + path = os.path.join(tmp, 'run_case.py') + with open(path, 'w', encoding='utf-8') as f: + f.write(src) + env = dict(os.environ, MPLBACKEND='Agg', PYTHONHASHSEED='0', + OMP_NUM_THREADS='1', MKL_NUM_THREADS='1', TOKENIZERS_PARALLELISM='false') + env.pop('CUDA_VISIBLE_DEVICES', None) + try: + p = subprocess.run([sys.executable, path], cwd=tmp, env=env, timeout=timeout, + capture_output=True, text=True, errors='replace') + except subprocess.TimeoutExpired: + return {'passed': False, 'kind': 'timeout', + 'error': f'the tests did not finish within {timeout}s', 'n_tests': 0} + out, err = p.stdout or '', p.stderr or '' + n_tests = n_fail = n_err = 0 + for line in out.splitlines(): + if line.startswith('__BCB__'): + _, a, b, c = line.split() + n_tests, n_fail, n_err = int(a), int(b), int(c) + if p.returncode == 0 and n_tests > 0: + return {'passed': True, 'kind': 'pass', 'error': '', 'n_tests': n_tests} + kind = 'assertion' if n_fail else ('exception' if n_err else 'import_or_syntax') + return {'passed': False, 'kind': kind, 'error': _trim_err(err.replace(tmp, '')), + 'n_tests': n_tests} + finally: + shutil.rmtree(tmp, ignore_errors=True) + + +def judge_many(items: List[Optional[Tuple[str, Any, int, Dict[str, Any]]]], + workers: int, timeout: int) -> List[Dict[str, Any]]: + """批量判分。``items[i]`` = (text, stop_reason, gen_tokens, payload) 或 None（无输出）。 + + ★ 必须批量：单测是子进程（导入 pandas/sklearn 后典型 1-3s），而一个 E4 chunk 有 ~290 次判分。 + 串行 ≈12 分钟/chunk，远超同 chunk 的 GPU 时间；线程池（--test-workers）把它压到 ~30s。 + 额外去重：同一题的多个 skill 在 T=0 executor 下经常产出逐字相同的代码，去重后实测省下可观的 + 子进程数（同 (task_id, code) 只跑一次）。 + """ + rolls: List[Dict[str, Any]] = [] + keys: List[Optional[Tuple[str, str]]] = [] + jobs: Dict[Tuple[str, str], Dict[str, Any]] = {} # key -> payload（去重后的待跑集合） + for it in items: + if it is None: + rolls.append(empty_roll()) + keys.append(None) + continue + text, stop, ntok, payload = it + code = extract_code(text) + key = (payload['task_id'], code) + rolls.append({'pred': None, 'correct': False, + 'terminated': stop != 'length', 'stop_reason': stop, + 'gen_tokens': int(ntok or 0), 'text': text, 'code': code, + 'kind': None, 'error': '', 'n_tests': 0}) + keys.append(key) + jobs.setdefault(key, payload) + if jobs: + todo = list(jobs) + with ThreadPoolExecutor(max_workers=max(1, min(workers, len(todo)))) as ex: + res = list(ex.map(lambda k: run_tests(k[1], jobs[k], timeout), todo)) + verdicts = dict(zip(todo, res)) + for r, key in zip(rolls, keys): + v = verdicts.get(key) if key is not None else None + if v is None: + continue + r['correct'] = bool(v['passed']) + # pred 在数学分支是"抽出来的答案，抽不到就是 None"，下游 term/answered_rate 与 + # acc/answered_pass 正是按 "pred is not None" 定义"交了可判的答案"这条通道。 + # ⚠️ 所以这里不能无条件写 kind：kind 恒非空会让 answered_rate 恒为 1.000、 + # answered_pass 退化成 candidate_pass，那两条曲线静默失效。口径对齐为： + # 抽到代码块 -> pred = 判分结论（pass/assertion/...，便于审计）；没抽到 -> None。 + r['pred'] = v['kind'] if r['code'] else None + r['kind'], r['error'], r['n_tests'] = v['kind'], v['error'], v['n_tests'] + return rolls + + +def empty_roll() -> Dict[str, Any]: + return {'pred': None, 'correct': False, 'terminated': False, 'stop_reason': 'empty', + 'gen_tokens': 0, 'text': '', 'code': '', 'kind': 'no_code', 'error': '', 'n_tests': 0} + + +def selftest(tasks: List[Dict[str, Any]], workers: int, timeout: int) -> List[str]: + """参考解答必须跑过它自己的单测 —— 跑不过说明沙箱/依赖不可判定，不是模型的错。 + 返回跑不过的 task_id 列表（调用方据此剔题）。""" + codes = [t['code_prompt'] + (t.get('canonical_solution') or '') for t in tasks] + payloads = [payload_of(t) for t in tasks] + with ThreadPoolExecutor(max_workers=max(1, min(workers, len(tasks) or 1))) as ex: + vers = list(ex.map(lambda p: run_tests(p[0], p[1], timeout), zip(codes, payloads))) + return [t['task_id'] for t, v in zip(tasks, vers) if not v['passed']] + + +# =========================================================================== +# executor prompts +# =========================================================================== +# 本数据集的硬性交付要求（BigCodeBench 官方 instruct 模式口径）。所有臂共用。 +EXEC_SYSTEM = """\ +You are an expert Python engineer. You will be given a task description that ends with the exact \ +import lines and function signature your solution must start with. + +Deliver exactly one fenced Python code block and nothing else after it: +- Reproduce the given imports and the given function signature verbatim, including parameter \ +names, order and default values. +- Add any further imports you need inside the same block; the block must run standalone. +- Return exactly the object type the task says to output. If it says the function should output \ +a tuple, return a tuple in that order; if it names a matplotlib Axes, return the Axes object \ +itself, not the Figure and not None. +- Implement the described behaviour for the general case, including the empty / single-element / \ +missing-column edge cases and any exception the description says to raise. +- Do not call the function, do not print demonstrations, do not add tests, do not use \ +`if __name__ == '__main__'`, and do not read from stdin. +- Do not include explanations outside the code block.""" + +# skill hint 包装语：与 v2 的数学版逐字同构（只把 "problem" 换成 "task"），保证 E4/E17 的 +# executor 输入除 skill 文本外没有第二个变量。 +_WRAPPER = ('Skill hint:\nFor this task, a skill-generation model has analyzed it and ' + 'provided some advisory skills:\n{hint}\n' + 'Prefer using its techniques when they fit, but if you have a clearly better ' + 'implementation, you may diverge. Be concise and accurate.\n') + + +def direct_prompt(problem: str) -> Dict[str, Any]: + return {'messages': [{'role': 'system', 'content': EXEC_SYSTEM}, + {'role': 'user', 'content': problem}]} + + +def skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: + skill = (skill or '').strip() + if not skill: + return direct_prompt(problem) + return {'messages': [{'role': 'system', 'content': EXEC_SYSTEM}, + {'role': 'user', + 'content': problem + '\n\n' + _WRAPPER.format(hint=skill)}]} + + +# =========================================================================== +# skill-gen prompts +# =========================================================================== +# E4（view B, query-only）。与数学版同构：先私下把题做一遍，再只写可迁移的方法论； +# 硬禁止写出解法代码与本题字面量（否则 skill 就是抄答案，测不到"方法论有没有用"）。 +SKILLGEN_SYSTEM = """\ +You are a skill-generation model for a Python implementation task. Your block will be fed to a SEPARATE downstream engineer model that must write the function on its own. The engineer sees the same task description and the same required signature, but NOT your private reasoning. + +First think privately: actually work out how you would implement it, including which library calls do the work. Then step back and write, inside , transferable guidance for THIS TYPE of task: which library functions are the right tool and what their relevant arguments and return shapes are, how to get the return value into the exact type the task demands, which edge cases and exceptions this kind of task always has, and the common mistakes to avoid. + +CRITICAL: do NOT write the solution code, and do NOT paste concrete literal values from this task. Name the API and describe the shape of the answer instead of writing it out. +Keep it to roughly one focused paragraph. Put ONLY the guidance inside .""" + + +def skillgen_prompt(problem: str) -> Dict[str, Any]: + return {'messages': [{'role': 'system', 'content': SKILLGEN_SYSTEM}, + {'role': 'user', 'content': f'Task:\n{problem}'}]} + + +# =========================================================================== +# rubric（判据 + judge prompt + 诊断素材） +# =========================================================================== +DIAG_SYSTEM = """\ +You are a strategy-level code reviewer. You are given a Python task description (with the required signature), a rubric, one attempted implementation, and the REAL error that attempt produced when the task's unit tests were run. Decide PASS or FAIL for each criterion, and write the diagnosis so it becomes reusable guidance for similar tasks without seeing this attempt again. + +Output STRICT JSON (no prose outside it) with this shape: +{"items": [{"index": 1, "verdict": "PASS"|"FAIL", "reason": "...", "fix": ""}], "overall": "OK"|"ISSUES", "summary": "..."} + +Rules: +- Ground every FAIL in the task description or the test error you were given; do not speculate. +- The test error is authoritative evidence: if it names an exception, a wrong type or a failed assertion, the criterion it implicates must be FAIL. +- Never write out corrected code; describe the process problem at strategy level. +- A fix suggests the local correction direction without implementing it. +- Keep "reason" and "fix" concise: one short sentence each. +- Output only the JSON object.""" + +DIAG_USER = """\ +## Task +{query} + +## Rubric +{rubric} + +## Attempted implementation and its test error +{segment} + +Now output the diagnostic JSON object.""" + +# PASS = 该类问题不存在（正向陈述，与 v2._format_diagnosis / gate 语义一致）。 +# 判据按 BigCodeBench 的实际失败模式组织：签名/导入、返回类型、API 用法、选库、边界、异常、逻辑。 +CODE_RUBRIC = [ + ('The implementation is runnable as given: it defines the required function with the exact ' + 'signature asked for and imports everything it uses', True), + ('The value returned matches the output type and structure the task states, element for ' + 'element and in the stated order', False), + ('The library functions used exist and are called with arguments and keyword names that ' + 'those functions actually accept', False), + ('The library chosen for each step is the one the task asks for, used for its intended ' + 'purpose rather than reimplemented by hand', False), + ('Edge cases the task implies (empty input, single element, missing key or column, ' + 'duplicate values) are handled instead of crashing', False), + ('Exactly the exceptions the task specifies are raised for invalid input, and no others ' + 'leak out', False), + ('The core computation implements what the description asks, with no step skipped, ' + 'inverted, or replaced by a placeholder', False), +] +# 版本号进 rubric 缓存键（rubric_cache._key）：判据一改旧诊断必须失效。code 与 math 的诊断 +# 还额外分文件存（trainer 按 task 选文件名），双保险。 +RUBRIC_VERSION = 'rubric_code_v1' + + +def spec_constraints(payload: Dict[str, Any]) -> str: + """任务自身声明的硬约定（签名、必需库、返回规格、应抛异常、文档示例）。 + + 只用题面信息、不含参考解答 —— 训练时同样拿得到，所以是"可得且非泄漏"的判据依据。 + ⚠️ BFCL 那边同类做法（schema_constraints）没能提升命中率，因为那边的对错定义在 gt 私有约定 + 里；这里返回类型/异常是题面明写的，所以这次它是真判据。 + """ + try: + doc = payload['doc_struct'] + doc = json.loads(doc) if isinstance(doc, str) else (doc or {}) + except Exception: + doc = {} + lines = [f"- required signature (must be reproduced verbatim):\n{payload['code_prompt'].strip()}"] + for key, label in (('reqs', 'must use these libraries'), ('returns', 'must return'), + ('raises', 'must raise'), ('params', 'parameters')): + vals = [str(x).strip() for x in (doc.get(key) or []) if str(x).strip()] + if vals: + lines.append(f'- {label}: ' + '; '.join(vals)[:400]) + ex = [str(x) for x in (doc.get('examples') or [])][:8] + if ex: + lines.append('- documented example calls:\n ' + '\n '.join(ex)) + return '\n'.join(lines) + + +def diag_query(problem: str, payload: Dict[str, Any]) -> str: + return problem + '\n\nHard requirements declared by the task:\n' + spec_constraints(payload) + + +def diag_segment(roll: Dict[str, Any]) -> str: + """★ rubric 路线唯一真正有效的一处：给 judge 的不是"输出全文"，而是**提交的代码 + 单测真实 + 报错**。 已在 extract_code 里切掉。BFCL 那轮 judge 手上没有任何客观证据，命中率 25% + ≈ 随机；这里报错是客观事实且不含参考解答。""" + return (f"### Submitted code\n```python\n{roll.get('code') or '(no parseable code block)'}\n```" + f"\n\n### Result of running the task's unit tests\n" + f"outcome: {roll.get('kind') or 'unknown'}\n" + f"{roll.get('error') or '(no error output)'}") + + +# =========================================================================== +# leak / skill 文本监控（代码域口径） +# =========================================================================== +def _canon_lines(payload: Dict[str, Any]) -> List[str]: + out = [] + for ln in (payload.get('canonical_solution') or '').splitlines(): + s = ln.strip() + if len(s) >= 20 and not s.startswith(('#', 'import ', 'from ', 'def ', 'return')): + out.append(s) + return out + + +def leaked(skill: str, payload: Any) -> bool: + """代码域 leak = skill 里出现了参考解答的实质代码行（>=20 字符、非 import/def/注释）。 + 与数学域一致：**只做监控，永不进 reward**（项目既定规则）。""" + if not skill or not isinstance(payload, dict): + return False + return any(ln in skill for ln in _canon_lines(payload)) + + +def skill_has_code(skill: str) -> bool: + """skill 退化监控：本该是方法论的块里出现了代码围栏或成段 def/return。 + 数学域的 digit_fraction / no_math_sentence 在代码域无意义，用这条替代。""" + s = skill or '' + if '```' in s: + return True + return bool(re.search(r'^\s*(def |return |import |for .*:|if .*:)', s, re.M)) diff --git a/cookbook/exp/skill2lora/run_ablate12.sh b/cookbook/exp/skill2lora/run_ablate12.sh index a90fb9173..991a09fa4 100644 --- a/cookbook/exp/skill2lora/run_ablate12.sh +++ b/cookbook/exp/skill2lora/run_ablate12.sh @@ -20,6 +20,8 @@ # LR=1e-6 RUN_SFT=1 FORCE=1 ONLY="E5 E6" SLEEP=30 SWANLAB_PROJECT=twinkle # MIN_LEVEL=6 CHUNK_SIZE=32 (gradient-signal fix: E1/E5 audit — level<=5 all-pass # dominated, 16-problem chunks leave only ~6 mixed groups per update; eval split unaffected) +# task=code 的臂（E4/E17，见 config.py）自动改走 BigCodeBench： +# BCB_PARQUET=... CODE_CHUNK_SIZE=48 CODE_EVAL_SIZE=200 CODE_TRAIN_N=0 TEST_WORKERS=24 # ============================================================================== set -euo pipefail @@ -30,6 +32,10 @@ export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:T HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" cd "$HERE" +# twinkle 是 editable 安装（.pth 指向 CPFS 上的 src/）；CPFS 瞬时抖动会让 site 初始化静默 +# 丢弃该路径，实验启动即死在 ModuleNotFoundError: twinkle（实测复现过一次）。PYTHONPATH 兜底。 +export PYTHONPATH="$(cd "$HERE/../../.." && pwd)/src${PYTHONPATH:+:$PYTHONPATH}" + # central env file (optional): put all knobs in one place. ENV_FILE=xxx overrides the path. ENV_FILE="${ENV_FILE:-$HERE/ablate12.env}" if [ -f "$ENV_FILE" ]; then @@ -38,6 +44,9 @@ if [ -f "$ENV_FILE" ]; then fi OUT_ROOT="${OUT_ROOT:-$HERE/output.ablate12}" +# interpreter with the full torch/vllm/twinkle stack; conda base shells shadow `python` with a +# numpy-less interpreter, so default to the absolute path and allow PYBIN=... to override. +PYBIN="${PYBIN:-/usr/local/bin/python3}" # DeepMath-103K (difficulty-stratified loader in skill_ablate/data.py); replaces the old # SEAM/aops input — see skill_quality_analysis.md 组成漂移修正. DEEPMATH_DIR="${DEEPMATH_DIR:-$(cd "$HERE/../../.." && pwd)/deepmath_103k}" @@ -49,24 +58,115 @@ LR="${LR:-1e-6}" MIN_LEVEL="${MIN_LEVEL:-6}" CHUNK_SIZE="${CHUNK_SIZE:-32}" SLEEP="${SLEEP:-30}" -SWANLAB_PROJECT="${SWANLAB_PROJECT:-twinkle}" +# SWAN_PROJ alias: swanlab>=0.8 的 pydantic Settings 会解析进程 env 里的 SWANLAB_PROJECT 并报错， +# 所以外部换项目请用 SWAN_PROJ=xxx，不要 export SWANLAB_PROJECT。 +# bugfix #10：ENV_FILE 的 set -a 会把文件里的 SWANLAB_PROJECT 自动 export（正好踩中上面的坑）； +# 读完值后 unset 掉 export 属性，再以普通 shell 变量重建，保证子进程 env 里没有它。 +_SWAN_PROJ_VAL="${SWAN_PROJ:-${SWANLAB_PROJECT:-twinkle}}" +unset SWANLAB_PROJECT +SWANLAB_PROJECT="$_SWAN_PROJ_VAL" RUN_SFT="${RUN_SFT:-0}" FORCE="${FORCE:-0}" ONLY="${ONLY:-}" +# leak 一律不进 reward（项目既定要求）：留空则不传 flag，由 main.py 默认值 0 生效。 +# 旧版在这里写死 1.0 且第 141 行无条件传入，会静默盖掉 Python 侧默认值。 +LOGP_LEAK_PENALTY="${LOGP_LEAK_PENALTY:-}" +# E16 passrate_hinge：截断铰链惩罚强度与起点、leak gate、base_tok 筛题阈（空=用 main.py 默认） +REWARD_TRUNC_PENALTY="${REWARD_TRUNC_PENALTY:-}" +REWARD_TRUNC_LO="${REWARD_TRUNC_LO:-}" +REWARD_LEAK_GATE="${REWARD_LEAK_GATE:-}" +BASE_TOK_FLOOR="${BASE_TOK_FLOOR:-}" +# kl_beta：对初始策略的锚，是唯一能提供“恢复力”对抗漂移的旋钮（reward 只能在组内排序）。 +# 空=用 main.py 默认 0.01（2026-07-29 从 0.001 上调；E1-E16 既往臂全跑在 0.001）。 +# RUN_TAG 同时隔离输出目录与 swanlab 实验名，用于并列跑同一 ExpSpec 的多个变体而不互相覆盖。 +KL_BETA="${KL_BETA:-}" +RUN_TAG="${RUN_TAG:-}" +# --- E17 reflexion 臂专属规模 ------------------------------------------------------------- +# 本臂只在【裸 executor 做错的题】上训练与评测，与其他臂不可横比（用户 2026-07-29 +# 拍板：只看本实验自身趋势），所以这些规模刻意与全局默认解耦。 +# 取值全部按 E16 落盘数据标定（.tmp_analysis/e17_param_calib.py，同模型/同题池/同难度门）： +# * K=24：E16 每次更新实际只有 23.46 组（总 1173 组）。累积证据 ∝sqrt(N)，而 E16 全程 +# 在文本层只累到 1.9 sigma —— 再砍组数就没任何判别力了。K=24 x 50 = 1200 组， +# 恰好追平 E16，而成本 24x8x8=1536 rollouts/chunk vs E16 实测 1501，几乎相等。 +# * CHUNK 128：实测裸错率 0.329（全量题池 527/1600，level>=6, T=0, floor=0）。 +# ⭐ 不是 0.446 —— 那是 base_tok>5000 筛选后子集的错误率（偏难），本臂 floor=0。 +# 二项精算 P(凑不满 24)：chunk 64 = 74%、96 = 3.6%、112 = 0.27%、128 = 0.012%。 +# 裸解开销 128 道 greedy 相当于 with-skill 的 8%，买对齐很便宜。 +# * TRAIN_N 8000：128 x 50 = 6400 次抽取，超过默认 5000 会进第二个 epoch（重复题）。 +# * EVAL 384 + --eval-min-level=MIN_LEVEL：指标只在错题子集上有信息量。旧口径（128 道 +# 全难度混合）只能给 ~33 道错题，SE 0.087，趋势根本读不出来；384 道 + 难度对齐 +# 训练池给 ~126 道（SE ~0.045）。正确的题跳过全部 GPU 路径，所以总成本几乎不变。 +REFLEXION_K="${REFLEXION_K:-24}" +E17_CHUNK_SIZE="${E17_CHUNK_SIZE:-128}" +E17_EVAL_SIZE="${E17_EVAL_SIZE:-384}" +E17_TRAIN_N="${E17_TRAIN_N:-8000}" +# E17 专属 kl_beta（2026-07-30 拍板 0.001，回到 E1-E16 旧值）。单独立一个变量而不改 +# 全局 KL_BETA，是为了不隐式改动其他臂重跑时的取值。 +E17_KL_BETA="${E17_KL_BETA:-0.001}" +# --- E18 rejection_sft 臂专属 -------------------------------------------------------------- +# 攒批阈值：16（用户 2026-07-30 拍板，从 128 改小）= 一个 sft_batch_size：chunk 32 每轮 +# 收 ~10 条胜者，大约隔 chunk 就 fire 一次，50 次更新 ~80 chunk 可达；128 要 ~15 chunk/次。 +E18_ACCUMULATE="${E18_ACCUMULATE:-16}" +# --- task=code 专属（2026-07-31 E4/E17 换到 BigCodeBench） -------------------------------- +# 规模按 bcb/bcb_eval0_probe.py + 2026-07-31 dry run 的落盘数据标定，与数学默认解耦： +# * 裸错率 **0.5625**（dry run 实测 18/32：训练侧 10/16 + eval 侧 8/16，think=on/8192）。 +# ⚠️ 不是 probe 的 0.715 —— 那是 nothink + 4096 的读数，think 开着以后模型强不少。 +# K=24 时 P(凑不满) : chunk 48 = 15.4%、56 = 1.6%、**64 = 0.1%**。组数恒定是本臂硬要求 +# （凑不满只会打印 k_short 并用更少的组训练，趋势就被抽样噪声污染），所以取 64。 +# 多花的只有裸解那 16 道 greedy（with-skill 部分由 K 固定，不随 chunk 变）。 +# * EVAL 200：题池 908（1140 剔缺库/外网GUI + 沙箱自检不过的 75），错题约 112 道，SE≈0.047。 +# * TRAIN_N=0 = 用掉剩下的全部题（708 道）。chunk64 x 50 ≈ 4.5 个 epoch 重复题（拍板接受）。 +# * TEST_WORKERS：判一条 rollout = 起一个 python 子进程跑 unittest（导入 pandas/sklearn 后 +# 典型 1-3s），一个 chunk 要判几百条，串行判分比同 chunk 的 GPU 时间还长。 +CODE_CHUNK_SIZE="${CODE_CHUNK_SIZE:-64}" +# ★ E4（bnpo）不受 K=24 那条约束 —— 它把 chunk 里每道题都拿来训练，chunk 直接等于每次更新的 +# 组数。跟着 reflexion 用 64 会白白把 rollout 数翻倍（64x8=512/更新），而且与已跑完的数学 +# E4（全局 CHUNK_SIZE=32）不再同规模、不可比。所以 view-B 的 code 臂单独用 32。 +CODE_BNPO_CHUNK_SIZE="${CODE_BNPO_CHUNK_SIZE:-32}" +CODE_EVAL_SIZE="${CODE_EVAL_SIZE:-200}" +CODE_TRAIN_N="${CODE_TRAIN_N:-0}" +BCB_PARQUET="${BCB_PARQUET:-$(cd "$HERE/../../.." && pwd)/bigcodebench/bcb.parquet}" +TEST_WORKERS="${TEST_WORKERS:-24}" +TEST_TIMEOUT="${TEST_TIMEOUT:-60}" +# --- executor nothink 对照（E19 math / E20 code，2026-07-31） ------------------------------- +# chunk / eval 都**不覆盖**：直接沿用同域 think 臂的值（math 128/384、code 64/200）。 +# 理由一，实测：首个 E19 run 的 c0 只从 64 道里拿到 20 道错题 —— math 关 think 的裸错率约 +# 0.31（baseline acc 0.69），与 think 的 0.329 基本相同，我此前估的 0.85 完全错了。 +# p=0.31 时 chunk 64 的期望错题 20±3.7，几乎每个 chunk 都凑不满 K=24，组数恒定失效。 +# code 侧 nothink p=0.622 > think 的 0.5625，chunk 64 本来就够，也无需覆盖。 +# 理由二，可比性：chunk 与 eval 都与同域 think 臂逐项相同，think/nothink 才是唯一自变量。 +# 保留这个 env 只为应急调参，默认空 = 不覆盖。 +EXEC_NOTHINK_CHUNK_SIZE="${EXEC_NOTHINK_CHUNK_SIZE:-}" +# eval 规模**不覆盖**：两个 nothink 臂各自的对照是同域的 think 臂，评测集必须逐题相同 —— +# E19 用 E17_EVAL_SIZE=384（已跑完的 E17 数学臂就是 n=384，baseline acc 0.674 / lift +0.130）， +# E20 用 CODE_EVAL_SIZE=200（与 E17 code 臂一致）。曾想为省 eval 开销把 nothink 统一压到 200， +# 那会让 E19 的评测集变成 E17 的子集、lift 曲线不再可逐题配对，省的钱不值这个代价。 +# --- E4/E17/E19/E20 统一口径（2026-07-31 用户拍板） ---------------------------------------- +# executor 生成预算统一 15000、skill 模型统一 8192（后者由 plan 的 smt 列给出，think=on 即 8192）。 +# 统一的意义：预算不再是四个臂之间的变量，think/nothink 的对比才是单变量的。 +# 连带两处必须跟着改，否则静默失效： +# 1) --max-model-len：默认 16384 装不下 prompt(约 1-2k) + 15000 输出，vLLM 会截 prompt；提到 20480。 +# 2) --reward-trunc-lo：长度惩罚死区按标定比例 5500/8192 缩放到预算上 = 15000*0.671 ≈ 10000。 +# 不改的话死区停在 5500（占预算 37%），会把远未撞墙的正常答案也纳入惩罚区。 +UNIFIED_MAX_TOKENS="${UNIFIED_MAX_TOKENS:-15000}" +UNIFIED_MAX_MODEL_LEN="${UNIFIED_MAX_MODEL_LEN:-20480}" +UNIFIED_TRUNC_LO="${UNIFIED_TRUNC_LO:-10000}" mkdir -p "$OUT_ROOT" # --- pull the run plan (name \t exp_dir \t think \t smt \t optional) ------------------- -PLAN="$(python3 skill_ablate/config.py --plan)" +# bugfix #9：用 $PYBIN（实验同一解释器）而非裸 python3，避免 plan/快照与实验环境不一致 +PLAN="$("$PYBIN" skill_ablate/config.py --plan)" snapshot_env() { # $1 = target file { echo "=== ablate12 env snapshot @ $(date -u +%FT%TZ) ===" echo "host: $(hostname)" - echo "python: $(python3 -c 'import sys;print(sys.version.split()[0])')" - echo "torch: $(python3 -c 'import torch;print(torch.__version__)' 2>/dev/null || echo NA)" - echo "vllm: $(python3 -c 'import vllm;print(vllm.__version__)' 2>/dev/null || echo NA)" - echo "transformers: $(python3 -c 'import transformers;print(transformers.__version__)' 2>/dev/null || echo NA)" + echo "pybin: $PYBIN" + echo "python: $("$PYBIN" -c 'import sys;print(sys.version.split()[0])')" + echo "torch: $("$PYBIN" -c 'import torch;print(torch.__version__)' 2>/dev/null || echo NA)" + echo "vllm: $("$PYBIN" -c 'import vllm;print(vllm.__version__)' 2>/dev/null || echo NA)" + echo "transformers: $("$PYBIN" -c 'import transformers;print(transformers.__version__)' 2>/dev/null || echo NA)" echo "CUDA_VISIBLE_DEVICES: ${CUDA_VISIBLE_DEVICES:-unset}" echo "nvidia-smi:"; nvidia-smi --query-gpu=name,memory.total --format=csv,noheader 2>/dev/null || echo " (nvidia-smi NA)" echo "GPU layout: TRAIN=${TRAIN_GPUS:-2} REF=${REF_GPUS:-2} SKILL_SAMPLER=${SKILL_SAMPLER_GPUS:-2} BASE_SAMPLER=${BASE_SAMPLER_GPUS:-2}" @@ -74,10 +174,12 @@ snapshot_env() { # $1 = target file } > "$1" } -echo "[ablate12] run order:"; echo "$PLAN" | awk -F'\t' '{printf " %s -> %s (think=%s smt=%s opt=%s)\n",$1,$2,$3,$4,$5}' +echo "[ablate12] run order:"; echo "$PLAN" | awk -F'\t' '{printf " %s -> %s (think=%s smt=%s opt=%s task=%s exec_think=%s)\n",$1,$2,$3,$4,$5,$6,$7}' -while IFS=$'\t' read -r NAME EXP_DIR THINK SMT OPTIONAL; do +while IFS=$'\t' read -r NAME EXP_DIR THINK SMT OPTIONAL TASK EXEC_THINK; do [ -z "$NAME" ] && continue + TASK="${TASK:-math}" + EXEC_THINK="${EXEC_THINK:-on}" if [ -n "$ONLY" ] && ! grep -qw "$NAME" <<< "$ONLY"; then echo "[ablate12] $NAME skipped (not in ONLY='$ONLY')"; continue fi @@ -85,7 +187,7 @@ while IFS=$'\t' read -r NAME EXP_DIR THINK SMT OPTIONAL; do echo "[ablate12] $NAME ($EXP_DIR) skipped: optional; set RUN_SFT=1 to run"; continue fi - EXP_OUT="$OUT_ROOT/$EXP_DIR" + EXP_OUT="$OUT_ROOT/$EXP_DIR${RUN_TAG:+.$RUN_TAG}" if [ -f "$EXP_OUT/DONE.json" ] && [ "$FORCE" != "1" ]; then echo "[ablate12] $NAME already done ($EXP_OUT/DONE.json); FORCE=1 to rerun"; continue fi @@ -93,27 +195,125 @@ while IFS=$'\t' read -r NAME EXP_DIR THINK SMT OPTIONAL; do snapshot_env "$EXP_OUT/env_info.txt" echo "======================================================================" - echo "[ablate12] START $NAME -> $EXP_OUT (think=$THINK skill_max_tokens=$SMT)" + echo "[ablate12] START $NAME -> $EXP_OUT (think=$THINK skill_max_tokens=$SMT"\ +"${KL_BETA:+ kl_beta=$KL_BETA}${RUN_TAG:+ tag=$RUN_TAG})" echo "======================================================================" LOG="$EXP_OUT/run.log" FORCE_FLAG="" [ "$FORCE" = "1" ] && FORCE_FLAG="--force" + E16_FLAGS="" + [ -n "$REWARD_TRUNC_PENALTY" ] && E16_FLAGS="$E16_FLAGS --reward-trunc-penalty $REWARD_TRUNC_PENALTY" + [ -n "$REWARD_TRUNC_LO" ] && E16_FLAGS="$E16_FLAGS --reward-trunc-lo $REWARD_TRUNC_LO" + [ -n "$REWARD_LEAK_GATE" ] && E16_FLAGS="$E16_FLAGS --reward-leak-gate $REWARD_LEAK_GATE" + [ -n "$BASE_TOK_FLOOR" ] && E16_FLAGS="$E16_FLAGS --base-tok-floor $BASE_TOK_FLOOR" + [ -n "$LOGP_LEAK_PENALTY" ] && E16_FLAGS="$E16_FLAGS --logp-leak-penalty $LOGP_LEAK_PENALTY" + [ -n "$KL_BETA" ] && E16_FLAGS="$E16_FLAGS --kl-beta $KL_BETA" + [ -n "$RUN_TAG" ] && E16_FLAGS="$E16_FLAGS --run-tag $RUN_TAG" + CHUNK_ARG="$CHUNK_SIZE" + EVAL_ARG="$EVAL_SIZE" + TRAIN_N_ARG="$TRAIN_N" + MIN_LEVEL_ARG="$MIN_LEVEL" + if [ "$TASK" = "code" ]; then + # BigCodeBench：题池/裸错率/判分方式全变，规模走 CODE_* 默认（见文件头注释）。 + # --min-level 一律传 0：BCB 没有 difficulty 字段，传非 0 只会让 data_code 打一行警告。 + CHUNK_ARG="$CODE_CHUNK_SIZE" + [ "$NAME" = "E4" ] && CHUNK_ARG="$CODE_BNPO_CHUNK_SIZE" + EVAL_ARG="$CODE_EVAL_SIZE" + TRAIN_N_ARG="$CODE_TRAIN_N" + MIN_LEVEL_ARG=0 + E16_FLAGS="$E16_FLAGS --task code --bcb-parquet $BCB_PARQUET"\ +" --test-workers $TEST_WORKERS --test-timeout $TEST_TIMEOUT" + echo "[ablate12] $NAME task=code: chunk=$CHUNK_ARG eval=$EVAL_ARG n=${TRAIN_N_ARG}(0=全池)"\ +" parquet=$BCB_PARQUET test_workers=$TEST_WORKERS" + fi + # reflexion 家族（E17 think / E19 math-nothink / E20 code-nothink）共用同一套协议开关。 + if [ "$NAME" = "E17" ] || [ "$NAME" = "E19" ] || [ "$NAME" = "E20" ]; then + # base_tok_floor 不在这里传：ReflexionMethod.__init__ 强制置 0 并告警，且那样 + # config 指纹里落的就是真实生效值（命令行重复传参只会制造歧义）。 + # ★ code 任务下 E17_* 这组数学标定值不适用（裸错率 0.329 -> 0.5625），保持上面 + # CODE_* 已设好的值不动。 + if [ "$TASK" != "code" ]; then + CHUNK_ARG="$E17_CHUNK_SIZE" + EVAL_ARG="$E17_EVAL_SIZE" + TRAIN_N_ARG="$E17_TRAIN_N" + E16_FLAGS="$E16_FLAGS --eval-min-level $MIN_LEVEL" + fi + E16_FLAGS="$E16_FLAGS --reflexion-k $REFLEXION_K" + # eval 口径改为 SEAM 式确定性单次（2026-07-30 拍板）：R=1 + T=0。 + # 代价：失去跨 4 个 skill 平均的降噪，题级读数从 5 档（0/.25/.5/.75/1）退为 0/1， + # 单点 SE 约为原来的 2 倍；换来的是与 SEAM val_kwargs(n=1,do_sample=False) 同口径。 + # 也因此与 E1-E16（R=4/T=0.5）的 eval 读数不同源，不可横比。 + E16_FLAGS="$E16_FLAGS --eval-rollouts 1 --eval-skill-temperature 0.0" + # kl_beta 回 0.001（与 E1-E16 一致）；放在这里覆盖，显式传的全局 KL_BETA 优先。 + [ -z "$KL_BETA" ] && E16_FLAGS="$E16_FLAGS --kl-beta $E17_KL_BETA" + echo "[ablate12] $NAME reflexion: chunk=$CHUNK_ARG k=$REFLEXION_K eval=$EVAL_ARG"\ +" n=$TRAIN_N_ARG task=$TASK eval_min_level=$MIN_LEVEL_ARG kl_beta=${KL_BETA:-$E17_KL_BETA}"\ +" eval_rollouts=1/T=0 (hard-subset protocol; NOT comparable to E1-E16)" + fi + if [ "$EXEC_THINK" = "off" ]; then + # executor 关 thinking：只加一个 flag。chunk/eval 一律沿用同域 think 臂（见文件头注释： + # 实测 math nothink 裸错率 0.31 ≈ think 的 0.329，压小 chunk 会让 K=24 凑不满）。 + [ -n "$EXEC_NOTHINK_CHUNK_SIZE" ] && CHUNK_ARG="$EXEC_NOTHINK_CHUNK_SIZE" + E16_FLAGS="$E16_FLAGS --executor-thinking off" + echo "[ablate12] $NAME executor=nothink: chunk=$CHUNK_ARG eval=$EVAL_ARG"\ +" (探针实测 nothink 截断 0.000、裸解 0.378>0.324、rubric 增量 +0.135 vs +0.080)" + fi + if [ "$NAME" = "E13" ]; then + # SEAM 论文设置复现（2026-08-01 用户拍板）。E13 本身已是 align='seam'（SEAM EXPERIENCE_PROMPT + # + executor 嵌套 prompt + lpem 整段 sanitize 判分）+ executor nothink（config 已改）+ skill 8192， + # executor 预算走 main.py 默认 max-tokens=8192 / max-model-len=16384 / n-skills=8，与 + # .tmp_analysis/SEAM/scripts/train_deepmath_paper.sh 的 executor 5120+8192 / K=8 同口径。 + # 这里再把三个 run 级默认拉到 SEAM run 口径（都可被同名 env 覆盖）： + # * MIN_LEVEL 0：SEAM 随机全池、不挑难题；我们默认 6=只挑最难档，会压平训练集 acc 曲线， + # 正是"我们从没见过 SEAM 那种上升曲线"的主因（见 2026-08-01 SNR 归因）。 + # * chunk 128：= SEAM train_batch_size；bnpo(view B) 无 reflexion 的 K=24 约束，放大安全。 + # * reward-trunc-penalty 0：SEAM reward = correct×format，无长度惩罚（train_skill_v2 头注）。 + # * eval R=1/T=0：对齐 SEAM val_kwargs(n=1,do_sample=False)，与 E17/E19/E20 同口径， + # 与 E1-E16(R=4/T=0.5) 不可横比。 + CHUNK_ARG="${E13_CHUNK_SIZE:-128}" + MIN_LEVEL_ARG="${E13_MIN_LEVEL:-0}" + [ -z "$REWARD_TRUNC_PENALTY" ] && E16_FLAGS="$E16_FLAGS --reward-trunc-penalty 0" + E16_FLAGS="$E16_FLAGS --eval-rollouts 1 --eval-skill-temperature 0.0" + echo "[ablate12] E13 SEAM-repro: chunk=$CHUNK_ARG min_level=$MIN_LEVEL_ARG"\ +" reward_trunc_penalty=0 eval=R1/T0 executor=nothink skill_max_tokens=$SMT (executor 预算 8192/16384, K=n_skills 默认 8)" + fi + case "$NAME" in + E4|E17|E19|E20) + # 四个臂统一 executor 预算 15000 + max_model_len 20480 + 长度惩罚死区 10000， + # 让 think/nothink 与 math/code 的对比都不夹带预算差异（2026-07-31 拍板）。 + # 显式传的 REWARD_TRUNC_LO 优先（上面已拼进 E16_FLAGS 的不会被这里覆盖）。 + E16_FLAGS="$E16_FLAGS --max-tokens $UNIFIED_MAX_TOKENS"\ +" --max-model-len $UNIFIED_MAX_MODEL_LEN" + [ -z "$REWARD_TRUNC_LO" ] && E16_FLAGS="$E16_FLAGS --reward-trunc-lo $UNIFIED_TRUNC_LO" + echo "[ablate12] $NAME 统一口径: max_tokens=$UNIFIED_MAX_TOKENS"\ +" max_model_len=$UNIFIED_MAX_MODEL_LEN skill_max_tokens=$SMT"\ +" reward_trunc_lo=${REWARD_TRUNC_LO:-$UNIFIED_TRUNC_LO}" + ;; + esac + if [ "$NAME" = "E18" ]; then + # eval 是 nothink 确定性单次（trainer 侧临时切 nothink 模板；R=1 + T=0 与 E17 同拍板）。 + # 其余规模全走全局默认（用户要求"配置不变"）；攒批阈值单独可调。 + E16_FLAGS="$E16_FLAGS --e18-accumulate $E18_ACCUMULATE --eval-rollouts 1 --eval-skill-temperature 0.0" + echo "[ablate12] E18 rejection_sft: chunk=$CHUNK_ARG eval=$EVAL_ARG n=$TRAIN_N_ARG"\ +" accumulate=$E18_ACCUMULATE eval=nothink/R=1/T=0" + fi set +e - python -m skill_ablate.main \ + "$PYBIN" -m skill_ablate.main \ --exp "$NAME" \ --deepmath-dir "$DEEPMATH_DIR" \ - --n "$TRAIN_N" \ - --eval-size "$EVAL_SIZE" \ + --n "$TRAIN_N_ARG" \ + --eval-size "$EVAL_ARG" \ --output-dir "$EXP_OUT" \ --skill-max-tokens "$SMT" \ --max-updates "$MAX_UPDATES" \ --eval-every-updates "$EVAL_EVERY" \ - --min-level "$MIN_LEVEL" \ - --chunk-size "$CHUNK_SIZE" \ + --min-level "$MIN_LEVEL_ARG" \ + --chunk-size "$CHUNK_ARG" \ --lr "$LR" \ --swanlab-project "$SWANLAB_PROJECT" \ + $E16_FLAGS \ $FORCE_FLAG \ - 2>&1 | tee "$LOG" + < /dev/null 2>&1 | tee "$LOG" RC=${PIPESTATUS[0]} set -e if [ "$RC" != "0" ]; then diff --git a/cookbook/exp/skill2lora/skill_ablate/config.py b/cookbook/exp/skill2lora/skill_ablate/config.py index 080697ba7..1d1a125fa 100644 --- a/cookbook/exp/skill2lora/skill_ablate/config.py +++ b/cookbook/exp/skill2lora/skill_ablate/config.py @@ -26,21 +26,56 @@ # (no leak, <=4096 chars); wrong -> rubric regen (2-in-8 pick 1) -> negative pool; # 1:1 accumulate -> SFT. # sft view A plain SFT: bare-problem wrong -> rubric -> regen (2-in-8) -> accumulate SFT. -METHODS = ('bnpo', 'rl_ab', 'rl_err', 'opsd', 'improve_sft', 'sft') -VIEW_OF_METHOD = {'bnpo': 'B', 'rl_ab': 'A', 'rl_err': 'A', - 'opsd': 'A', 'improve_sft': 'A', 'sft': 'A'} +# logp_rl E14+: query-only skill-gen, executor T>0 samples a correct pseudo-GT solution S, +# rubric/API audits that executor response, then reward each skill by +# Δ mean logP_executor(S | problem + skill) with answer-leak penalty. +# logp_gt E15: same dense executor-logP reward as logp_rl, but the target S is DeepMath's +# external R1 reference solution (record 'solution'); NO executor rollout and NO +# rubric audit (view B, query-only) -> runs much faster. Validates the logP path +# against a strong external target instead of a self-sampled pseudo-GT. +# passrate_hinge E16: view B query-only. Reward = pass_rate(M rollouts, T>0) minus a hinge +# truncation penalty (only rollouts whose executor output nears the 8192 budget are +# penalized) minus a leak gate; trained problems are pre-filtered to the danger band +# (baseline executor output long / not all-pass). Data-driven closure of the reward +# probe (skill_quality_analysis.md 2026-07-29): pass_rate is the only real signal, +# trunc is the strongest dense side-signal, base_tok is the strongest problem filter. +METHODS = ('bnpo', 'rl_ab', 'rl_err', 'opsd', 'improve_sft', 'sft', 'logp_rl', 'logp_gt', + 'passrate_hinge', 'reflexion', 'rejection_sft') +VIEW_OF_METHOD = {'bnpo': 'B', 'rl_ab': 'A', 'rl_err': 'A', 'reflexion': 'A', + 'opsd': 'A', 'improve_sft': 'A', 'sft': 'A', 'logp_rl': 'A', 'logp_gt': 'B', + 'passrate_hinge': 'B', 'rejection_sft': 'A'} STYLES = ('narrative', 'pitfall') THINKINGS = ('on', 'off') +TASKS = ('math', 'code') @dataclass(frozen=True) class ExpSpec: - name: str # E1..E13 + name: str # E1..E14 method: str # one of METHODS thinking: str # 'on' | 'off' style: str # 'narrative' | 'pitfall' (ignored by align='seam': SEAM prompts bypass style) optional: bool = False # E12(sft): manually gated (RUN_SFT=1), runs last align: str = 'v2' # 'v2' | 'seam' — sets v2._ALIGN_MODE (prompt/判分/executor 嵌套全开关) + # E14+ 稠密 reward：executor 先按 T>0 采样 K 条找正确伪 GT S；训练 reward 不再 rollout， + # 而是算 Δ mean logP_executor(S | problem + skill)。默认字段复用 reward_rollouts/temperature。 + reward_rollouts: int = 1 + reward_temperature: float = 0.0 + smt_override: int = 0 # force skill_max_tokens regardless of the think rule; 0 = default rule + # 任务族（2026-07-31 用户拍板把 E4/E17 换到 BigCodeBench）： + # 'math' = DeepMath-103K + \boxed{} 数值判分（E1-E16、E18 原样） + # 'code' = BigCodeBench + 跑 unittest 判分，executor/skill-gen/rubric 三套 prompt 全换 + # 依据：数学域上 rubric 无增量甚至负向（−0.056），BFCL 上为零（+0.002），只有 BigCodeBench + # 这种"机器给出可定位失败证据（异常/断言/行号）"的任务上 rubric 才有增量（+0.135, p=4e-5， + # 见 bcb/bcb_eval0_probe.py 与 code_task.py 模块注释）。 + task: str = 'math' + # executor(base_sampler) 的 thinking。E1-E18 全是 'on'（历史口径，勿动）。 + # 'off' = 2026-07-31 新增的 nothink 对照（E19 math / E20 code）：bcb 探针实测 think 的 + # executor 有 34-50% rollout 撞满预算且是字面死循环（8-gram 重复率 p50=0.835），加预算 + # 到 20000 无效；关掉后截断归零、裸解 0.378 > 0.324、rubric 增量 +0.135 vs +0.080。 + # 注意这是 executor 侧；skill 模型仍由 thinking 字段控制（两个新臂都保持 skill think=on， + # 因为 v2 实测 skill 侧 nothink 会把完整解答写进 ，等于换标签的泄漏）。 + executor_thinking: str = 'on' @property def view(self) -> str: @@ -52,6 +87,15 @@ def needs_rubric(self) -> bool: @property def skill_max_tokens(self) -> int: + # 显式 per-spec override 优先。 + # ⚠️ 2026-07-29 实测推翻了此前“think 长度与 skill 质量无关，截掉长尾不损失信号”的探针结论： + # E4 在 8192 与 4096 下的受控 A/B（其余配置全同，比共有的 chunk 0-21）显示 4096 是净损失—— + # parse 率 0.939 -> 0.740，无条件 cand_pass 0.689 -> 0.649，撞 think 顶 0.063 -> 0.268。 + # 4096 臂看起来“exec 更短、trunc 更低”纯属幸存者偏差（pass|parse 0.734 -> 0.877）， + # 因为 26% 的候选连 都没写完就掉出了统计口径。 + # 见 .tmp_analysis/think_budget_ab.py。think 模式一律用 8192。 + if self.smt_override: + return self.smt_override # think must have room for + (4096 truncates to an empty block). # seam align: 人工拍板用 8192（不复刻 SEAM 原版 4096：think 模式下 4096 会把大量候选截断在 # 里、压低 parseable，与“think 模式 skill-max-tokens 必须 8192”的矩阵规范保持一致）。 @@ -65,8 +109,18 @@ def loss(self) -> str: @property def exp_dir(self) -> str: - # output.ablate12/E{n}_{method}_{think}_{style}/ (seam align 加后缀区分) + # output.ablate12/E{n}_{method}_{think}_{style}/ (seam align / multi-rollout reward 加后缀区分) suffix = '_seam' if self.align == 'seam' else '' + if self.reward_rollouts > 1: + suffix += f'_r{self.reward_rollouts}' + # 换数据集必须换目录：否则新语义的 run 会撞上已跑完的数学 run（DONE.json 直接跳过、 + # 曲线与 gen_records 混在一起、config 指纹对不上）。 + if self.task != 'math': + suffix += f'_{self.task}' + # executor 口径也必须进目录名：同一个 ExpSpec 换 executor thinking 后 baseline 缓存 + # （key=problem）与 eval 读数全变，撞同一个目录会把两套语义的曲线混在一起。 + if self.executor_thinking != 'on': + suffix += '_execnothink' return f'{self.name}_{self.method}_{self.thinking}_{self.style}{suffix}' @property @@ -80,12 +134,17 @@ def swanlab_exp(self) -> str: ExpSpec('E1', 'bnpo', 'off', 'pitfall'), ExpSpec('E2', 'bnpo', 'off', 'narrative'), ExpSpec('E3', 'bnpo', 'on', 'pitfall'), - ExpSpec('E4', 'bnpo', 'on', 'narrative'), + # E4/E8 的 smt_override=4096 已被上面 skill_max_tokens 里记录的 A/B 推翻（净损失 4 个点 + # 无条件 pass + 20 个点 parse 率）。E8 保留原值只是为了不改动“已跑完的臂”的可复现配置。 + # ★ E4 于 2026-07-31 换到 code 任务并同时删掉 override（用户拍板）：既然要重跑，就按 + # think=on 的规范值 8192 跑，与 E17 同口径可比。旧数学 E4 的产物在 + # output.ablate12/E4_bnpo_on_narrative/（新 run 落在 ..._code/，互不覆盖）。 + ExpSpec('E4', 'bnpo', 'on', 'narrative', task='code'), # group 2 — view A RL-AB-mix: same grid as E1-E4, isolates "rubric rescues zero-grad groups" ExpSpec('E5', 'rl_ab', 'off', 'pitfall'), ExpSpec('E6', 'rl_ab', 'off', 'narrative'), ExpSpec('E7', 'rl_ab', 'on', 'pitfall'), - ExpSpec('E8', 'rl_ab', 'on', 'narrative'), + ExpSpec('E8', 'rl_ab', 'on', 'narrative', smt_override=4096), # 同 E4：重跑前应删掉此 override # group 3 — view A training-method comparison (fixed think+narrative), sft last & optional ExpSpec('E9', 'rl_err', 'on', 'narrative'), ExpSpec('E10', 'opsd', 'on', 'narrative'), @@ -97,12 +156,107 @@ def swanlab_exp(self) -> str: # prompt_text+response_text(+think), lpem-parity greedy scoring, actor budget 4096. # Query-only BNPO main loop (= SEAM's training form); eval stays the matrix-unified # query-only readout so E13 is directly comparable with E1-E12. - ExpSpec('E13', 'bnpo', 'on', 'narrative', align='seam'), + ExpSpec('E13', 'bnpo', 'on', 'narrative', align='seam', executor_thinking='off'), + # group 5 — E14+ 稠密 reward：E4 的部署形态（query-only skill-gen / eval 不变），训练时 + # executor 先 T=0.7×16 采样，选本地判分正确且非截断的伪 GT S，并用 rubric/API 产审计诊断； + # skill reward = Δ mean logP_executor(S | problem + skill)（leak 不进 reward，只做监控）。该臂同时降测量噪声 + # 和抬内容信号，替代旧版 T=0.5×4 多数 rollout 0/1 reward。 + ExpSpec('E14', 'logp_rl', 'on', 'narrative', reward_rollouts=16, reward_temperature=0.7), + # group 6 — E15 稠密 reward 的 GT 版验证：与 E14 同一套 executor logP reward，但 logP 目标 S + # 换成 DeepMath 自带的 R1 参考解（record 'solution'），不再 executor rollout、不再 rubric 审计 + # （view B / query-only）。用于验证“ΔlogP(强外部参考解 | 题+skill)”是否走得通；因省掉 K 次 + # executor 采样 + rubric API，单 chunk 比 E14 快很多。 + ExpSpec('E15', 'logp_gt', 'on', 'narrative'), + # group 7 — E16 数据驱动收敛臂：探针（skill_quality_analysis.md 2026-07-28/29）判死 logP， + # 确认 pass_rate 是唯一真实信号、trunc 是最强稠密辅助、base_tok>5000 是最强筛题维度。 + # reward = mean_i(correct_i·eff_i) - kappa·mean_i(1-eff_i)，护栏 max(reward, pass_rate-1/M)； + # eff = (1-alpha·len_pen)·(1-beta·loop_pen)，len_pen 从 5500 起二次凸爬升（数据标定见 + # .tmp_analysis/reward_shape_calib.py）。leak 不进 reward，只做监控。 + # skill_max_tokens=8192（2026-07-29 拍板）：4096 的 A/B 判为净损失，见 skill_max_tokens 注释。 + # 训练题预筛危险带（baseline 输出长）。view B query-only，不用 rubric。 + ExpSpec('E16', 'passrate_hinge', 'on', 'narrative', + reward_rollouts=8, reward_temperature=0.5, smt_override=8192), + # group 8 — E17 Reflexion 臂：唯一目的是检验"rubric 注入权重外信息"能否让 skill 变得可学。 + # 与 E16 的三处结构差异（2026-07-29 人工拍板）： + # 1) 只在【裸 executor 做错】的题上训练与评测，做对的题完全不碰。E16 的 lift 分解 + # +0.170 = 救回 +0.230 - 破坏 -0.060（skill 挂在裸对的题上会砸掉 10.6% 保持率）， + # 条件化按构造把破坏项归零。见 .tmp_analysis/lift_source_decomp.py。 + # 2) skill-gen 走 view A（query+rubric）。E16 判定 reward 对"该写什么文本"几乎无信息 + # （结果层 SNR 2.25 -> 文本层 0.046，损失 ~50 倍）；rubric 是外部 API 的诊断，是本臂 + # 唯一的新变量，也是它可能不重演 E16 结局的唯一理由。见 .tmp_analysis/batch_size_math.py。 + # 3) 不用 base_tok 危险带筛题（base_tok_floor=0）。 + # ⚠️ 当初的理由已被实测证伪，保留在此以免重蹈：本以为"去掉筛选后错题集就是 + # 推理错 + 没写完的混合，才测得到方法修正"。实测（e17_param_calib.py + E16 全 1600 + # 题）：全量题池的 527 道错题里 96.96% 是 base_tok>=8192（没写完），只有 3.04% + # （16 道）是写完但答错；而 floor=5000 筛选后是 97.71%。也就是说 floor 不是截断主导的 + # 原因，题目本身是（level>=6 配 8192 预算），去掉筛选只把可用的"方法错"样本从 + # 2.3% 提到 3.0%（每 chunk 24 道错题里平均 0.73 道）。保留 floor=0 只是因为它严格 + # 不差于 floor=5000，不要再把它当成本臂能测到 reflexion 的理由。 + # 后果：rubric 在 ~97% 的题上只能说"你超预算了"，signal/wrong_trunc_frac 会直接 + # 开在 0.97。想真正测"方法修正"必须先把 executor 预算提到 16384。 + # 见 .tmp_analysis/e16_redteam5.py、e17_param_calib.py。 + # 批量对齐：chunk_size=128 裸解后取 --reflexion-k=24 道错题，每次更新恒定 24 题 x 8 候选。 + # K=24 不是拍的：E16 每次更新实际只有 23.46 组（全程 1173 组），而累积证据 ∝sqrt(N) + # 且 E16 在文本层只累到 1.9 sigma，再砍组数就没判别力了。chunk=128 是为了把 + # P(凑不满 24) 压到 0.012%（全量题池裸错率实测 0.329，不是筛选后子集的 0.446）。 + # 见 .tmp_analysis/e17_param_calib.py。 + # reward 形状沿用 E16 的 passrate_hinge（死区 5500/二次凸/leak 只监控），但 rollout 数改为 + # M=1（2026-07-30 拍板）：每个 skill 只让 executor 推理一次，reward = 这一次对/错。 + # ⭐两个连带后果（改前必读）： + # 1) pass_rate 从 9 档（0,1/8,...,1）退化为 2 档（0/1），组内无方差的概率大升 + # —— 这正是 E4(M=1) 62% 零梯度组的成因，M=8 当初就是为了绕开它。 + # 2) methods.py 的不可反转护栏 pen_cap = 1/M - 1e-6：M=8 时是 0.125（惩罚只能微调 + # 同档内排序），M=1 时变成 ~1.0，长度/兜圈惩罚能把做对的候选拉到贴近 0。 + # 形式保证（做对过永远排在没做对前）仍成立，但 reward 量级上从“以正确率为主” + # 变成“以长度惩罚为主”。 + # T=0（对齐 SEAM 的 grm.rollout.temperature=0 与 E4/E13）：打分时 executor 贪心确定性解一次， + # 同一个 skill 重跑 reward 不变 —— 消掉 executor 采样噪声，reward 只反映 skill 本身的差异。 + # （skill 模型自己的采样温度是另一个参数 skill_gen_temperature=1.0，不受此影响。） + # ★ 2026-07-31 换数据集（用户拍板）：task='code' —— BigCodeBench + 跑 unittest 判分。 + # 换的理由是上面 (2)(3) 两条在数学域已被实测封死：97% 的裸失败是"没写完"，rubric 只能 + # 说"你超预算了"，三个数据集横比下来 rubric 增量 −0.056(deepmath) / +0.002(BFCL) / + # **+0.135(BigCodeBench, p=4e-5)**，唯一的差别是 judge 手上有没有机器给出的可定位失败 + # 证据（异常类型/断言差异/失败用例名）。code 分支把这份证据喂进 judge（code_task.diag_segment）。 + # 连带变化：裸错率从数学的 0.329 变成 0.5625（2026-07-31 dry run 实测 18/32，think=on/8192； + # ⚠️ 不是 probe 的 0.715，那是 nothink+4096 的读数），所以 K=24 需要 chunk=64 + # （P(凑不满 24)=0.001；chunk 48 会有 15.4% 的更新组数不足）；题池 908 题，chunk64×50 + # ≈4.5 个 epoch 重复题（用户同意）。--min-level/--eval-min-level 在 code 下自动失效。 + ExpSpec('E17', 'reflexion', 'on', 'narrative', task='code', + reward_rollouts=1, reward_temperature=0.0, smt_override=8192), + # group 9 — E18 拒绝采样 SFT（2026-07-30 用户拍板）：采集与 E17 同源（裸解错题 -> rubric -> + # rubric 条件化 skill-gen think 模式，executor greedy T=0 单次判分），但不做 RL：每题在做对的 + # 候选里按 leak 过滤 -> 长度贴近 len_budget -> 与原始 rubric 相似度最高 三道筛取唯一胜者， + # 写本地数据集文件，攒够 16 条 SFT 一次（weight=1，nothink 布局响应；2026-07-30 从 128 改小， + # 使 50 次更新在 ~80 chunk 内可达），训完同步权重到 vLLM。 + # eval：同一个 skill_sampler vLLM 临时切 nothink 模板跑 query-only greedy（trainer 侧实现）。 + # reward_rollouts/temperature 对本臂无效（判分固定 greedy 单次），填 1/0.0 只为指纹如实。 + ExpSpec('E18', 'rejection_sft', 'on', 'narrative', + reward_rollouts=1, reward_temperature=0.0, smt_override=8192), + # group 10 — E19/E20 executor-nothink 对照（2026-07-31 用户拍板）：与 E17 逐字同构的 + # reflexion 臂，唯一变量是 **executor 关 thinking**（skill 侧仍 think=on）。 + # 起因：bcb 探针在 4096/12288/20000 三档 think 预算下测到 rubric 增量 +0.047/+0.058/+0.080， + # 而 nothink 一档是 +0.135（p=1e-4）；增量与"裸失败里 no_code 的占比"严格反向 + # （69%/54%/50%/0%）。截断样本经查是字面死循环（8-gram 重复率 p50=0.835、同一长句重复 + # 92 次），所以加预算无解、只能关 think。两个臂分别回答： + # E19（math）：deepmath + 数学 prompt/rubric。数学域此前的 rubric 增量是 −0.056，而 + # 那次 97% 的裸失败是"没写完"；关掉 executor think 后失败会变成"写完但答错"， + # 这是第一次能在数学域上把"rubric 是否有用"与"截断"分开测。 + # E20（code）：BigCodeBench + code prompt/rubric，直接把探针里 +0.135 那一档搬到训练。 + # ⚠️ 规模不能照抄：裸错率随 executor 口径变（code think 0.5625 -> nothink 约 0.62； + # math nothink 未测，预期远高于 think 的 0.329），chunk 由 run_ablate12.sh 的 + # EXEC_NOTHINK_* 单独给，dry run 后按实测复核。 + ExpSpec('E19', 'reflexion', 'on', 'narrative', executor_thinking='off', + reward_rollouts=1, reward_temperature=0.0, smt_override=8192), + ExpSpec('E20', 'reflexion', 'on', 'narrative', task='code', executor_thinking='off', + reward_rollouts=1, reward_temperature=0.0, smt_override=8192), ] # execution order: all nothink first, then think; E13 (seam-align baseline) right after E6; -# the data-hungry SFT method dead last. -RUN_ORDER: List[str] = ['E1', 'E2', 'E5', 'E6', 'E13', 'E3', 'E7', 'E8', 'E9', 'E10', 'E11', 'E4', 'E12'] +# E14 (rubric pseudo-GT + executor logP dense reward) right after E7 per 2026-07-28 人工拍板; +# E15 (GT-target logP validation) right after E14; the data-hungry SFT method dead last. +# 2026-07-31 用户拍板：E19/E20（executor nothink）排在 E4/E17 之前先跑 —— 探针已判定 +# nothink 是 rubric 增量最大且唯一无截断混杂的口径，先拿这两个臂的结论。 +RUN_ORDER: List[str] = ['E1', 'E2', 'E5', 'E6', 'E13', 'E3', 'E7', 'E14', 'E15', 'E19', 'E20', 'E4', 'E8', 'E16', 'E17', 'E18', 'E9', 'E10', 'E11', 'E12'] BY_NAME: Dict[str, ExpSpec] = {e.name: e for e in MATRIX} @@ -127,6 +281,12 @@ def _self_check() -> None: assert e.method in METHODS, f'{e.name}: bad method {e.method}' assert e.thinking in THINKINGS and e.style in STYLES, f'{e.name}: bad think/style' assert e.align in ('v2', 'seam'), f'{e.name}: bad align {e.align}' + assert e.task in TASKS, f'{e.name}: bad task {e.task}' + assert e.executor_thinking in THINKINGS, \ + f'{e.name}: bad executor_thinking {e.executor_thinking}' + assert not (e.task == 'code' and e.align == 'seam'), \ + f'{e.name}: seam align is math-only (SEAM prompts/parsing are \\boxed 数值口径)' + assert e.reward_rollouts >= 1, f'{e.name}: bad reward_rollouts {e.reward_rollouts}' # nothink-before-think ordering within contiguous runs is a soft convention, not asserted. @@ -134,9 +294,11 @@ def _self_check() -> None: import sys _self_check() if '--plan' in sys.argv: - # machine-readable run plan for the launcher: nameexp_dirthinksmtoptional + # machine-readable run plan for the launcher: + # nameexp_dirthinksmtoptionaltaskexecutor_thinking for e in ordered_specs(): - print(f'{e.name}\t{e.exp_dir}\t{e.thinking}\t{e.skill_max_tokens}\t{int(e.optional)}') + print(f'{e.name}\t{e.exp_dir}\t{e.thinking}\t{e.skill_max_tokens}\t' + f'{int(e.optional)}\t{e.task}\t{e.executor_thinking}') sys.exit(0) print(f'{len(MATRIX)} experiments; run order: {" -> ".join(RUN_ORDER)}') hdr = f'{"name":<4} {"view":<4} {"method":<12} {"think":<6} {"style":<10} {"align":<5} {"smt":<5} {"loss":<5} opt' diff --git a/cookbook/exp/skill2lora/skill_ablate/data.py b/cookbook/exp/skill2lora/skill_ablate/data.py index ac84c6d03..895fc0523 100644 --- a/cookbook/exp/skill2lora/skill_ablate/data.py +++ b/cookbook/exp/skill2lora/skill_ablate/data.py @@ -34,7 +34,8 @@ def _read_rows(deepmath_dir: str) -> List[Dict[str, Any]]: raise FileNotFoundError(f'no parquet files under --deepmath-dir {deepmath_dir}') rows: List[Dict[str, Any]] = [] for p in paths: - t = pq.read_table(p, columns=['question', 'final_answer', 'difficulty']) + # r1_solution_1: E14-ref 的 logP 目标文本（外部 R1 参考解，强于 executor 自采解） + t = pq.read_table(p, columns=['question', 'final_answer', 'difficulty', 'r1_solution_1']) rows.extend(t.to_pylist()) return rows @@ -50,7 +51,8 @@ def load_deepmath_records(args) -> Tuple[List[Dict[str, Any]], List[Dict[str, An continue lvl = int(round(float(r.get('difficulty') or 0))) pool.append({'data_id': f'dm:{lvl}:{i}', 'problem': problem, - 'reference_answer': num, '_level': lvl}) + 'reference_answer': num, '_level': lvl, + 'solution': (r.get('r1_solution_1') or '').strip()}) # bucket by level, seeded shuffle inside each bucket buckets: Dict[int, List[Dict[str, Any]]] = defaultdict(list) @@ -60,9 +62,16 @@ def load_deepmath_records(args) -> Tuple[List[Dict[str, Any]], List[Dict[str, An for lvl in sorted(buckets): rng.shuffle(buckets[lvl]) - # eval quota per bucket: proportional, largest-remainder rounding - eval_n = min(args.eval_size, len(pool)) if args.eval_size > 0 else 0 - quota = {lvl: eval_n * len(b) / len(pool) for lvl, b in buckets.items()} + # eval quota per bucket: proportional, largest-remainder rounding. + # --eval-min-level>0 把配额限在难度不低于它的桶里（默认 0 = 全难度混合，与旧臂逐字一致）。 + # E17 需要它：该臂的指标只在【裸解做错的题】上有信息量，而全难度混合的错题率只有 + # ~26%（E16 实测 128 -> 33 道，SE 0.087，趋势读不出来）；level>=6 上是 ~43%，同样的 + # GPU 成本能换到 1.6 倍的有效样本，同时与训练池（min_level）同分布。 + eval_min_level = int(getattr(args, 'eval_min_level', 0) or 0) + elig = {lvl: b for lvl, b in buckets.items() if lvl >= eval_min_level} + n_elig = sum(len(b) for b in elig.values()) + eval_n = min(args.eval_size, n_elig) if (args.eval_size > 0 and n_elig) else 0 + quota = {lvl: eval_n * len(b) / n_elig for lvl, b in elig.items()} if n_elig else {} take = {lvl: int(q) for lvl, q in quota.items()} for lvl in sorted(quota, key=lambda x: quota[x] - int(quota[x]), reverse=True): if sum(take.values()) >= eval_n: @@ -72,8 +81,9 @@ def load_deepmath_records(args) -> Tuple[List[Dict[str, Any]], List[Dict[str, An eval_records, train_records = [], [] for lvl in sorted(buckets): b = buckets[lvl] - eval_records.extend(b[:take[lvl]]) - train_records.extend(b[take[lvl]:]) + n_ev = take.get(lvl, 0) + eval_records.extend(b[:n_ev]) + train_records.extend(b[n_ev:]) min_level = int(getattr(args, 'min_level', 0) or 0) if min_level > 0: # train-only floor; eval keeps full-level mix (see module docstring) n_before = len(train_records) diff --git a/cookbook/exp/skill2lora/skill_ablate/data_code.py b/cookbook/exp/skill2lora/skill_ablate/data_code.py new file mode 100644 index 000000000..89b78830e --- /dev/null +++ b/cookbook/exp/skill2lora/skill_ablate/data_code.py @@ -0,0 +1,82 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""BigCodeBench loader for the ablation package (code task family). + +Dataset: bigcodebench/bcb.parquet (v0.1.4, 1140 tasks). Record shape matches the math loader +so nothing downstream changes shape: + data_id = task_id (e.g. BigCodeBench/42) + problem = instruct_prompt (ends with the exact imports + signature to reproduce) + reference_answer = code_task.payload_of(task) # 判分载荷，不是数值答案 + +与 DeepMath loader 的三处结构差异： +1. **没有 difficulty**，所以没有分层抽样，也不存在 --min-level / --eval-min-level（在 code + 模式下被显式忽略并告警）。train/eval 就是同一个 seed 洗牌后的前后切分。 +2. **题池极小**：1140 题，剔掉缺库/需外网GUI子进程的题后约 900，再减 eval 后训练池只有几百题。 + E17 用 chunk 48 × 50 updates = 2400 次抽取 ≈ 3 个 epoch 重复题（用户 2026-07-31 拍板接受）。 + 重复题的 rubric 全部缓存命中，所以重复的成本只在 GPU rollout。 +3. **有"沙箱自检"这道闸**：参考解答跑不过它自己的单测 = 环境不可判定（缺库的边角、随机种子、 + matplotlib 后端等），这类题的 0 分与模型能力无关，必须剔掉，否则它会给每个臂加一层同样的 + 噪声底并稀释 lift。probe 实测 7.5% 属于这一类。自检结果按 parquet 落一个 json 缓存， + 之后各臂零成本复用。 +""" +import json +import os +from typing import Any, Dict, List, Tuple + +import code_task +import train_skill_v2 as v2 + + +def _broken_task_ids(args, tasks: List[Dict[str, Any]]) -> set: + """参考解答跑不过自己单测的题（缓存到 --output-dir 的父目录，跨臂复用）。""" + base = (getattr(args, 'rubric_global_dir', None) + or os.path.dirname(os.path.abspath(str(args.output_dir).rstrip('/')))) + os.makedirs(base, exist_ok=True) + path = os.path.join(base, 'bcb_broken_tasks.json') + if os.path.exists(path): + try: + with open(path, encoding='utf-8') as f: + cached = json.load(f) + if int(cached.get('n_tasks', -1)) == len(tasks): + return set(cached.get('broken') or []) + v2.logger.info(f'[data] {path} 的题数 {cached.get("n_tasks")} != 当前 {len(tasks)}，' + f'重跑自检') + except Exception as exc: + v2.logger.warning(f'[data] 读取 {path} 失败（{exc}），重跑自检') + v2.logger.info(f'[data] 沙箱自检：{len(tasks)} 道题跑参考解答（一次性，之后走缓存）…') + broken = code_task.selftest(tasks, args.test_workers, args.test_timeout) + with open(path, 'w', encoding='utf-8') as f: + json.dump({'n_tasks': len(tasks), 'broken': sorted(broken)}, f, indent=1) + return set(broken) + + +def load_code_records(args) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + """-> (train_records, eval_records)，每条 {'data_id','problem','reference_answer'}。""" + tasks, stats = code_task.load_tasks(args.bcb_parquet, args.seed) + if not tasks: + raise FileNotFoundError(f'no usable BigCodeBench task in {args.bcb_parquet}') + v2.logger.info(f"[data] BigCodeBench: 全集 {stats['raw']}，剔除依赖缺失 " + f"{stats['drop_missing_lib']}、需外网/GUI/子进程 {stats['drop_needs_net_or_gui']}" + f" -> {stats['kept']}") + if getattr(args, 'code_selftest', True): + broken = _broken_task_ids(args, tasks) + if broken: + tasks = [t for t in tasks if t['task_id'] not in broken] + v2.logger.info(f'[data] 剔除参考解答自己跑不过单测的题 {len(broken)} 道 ' + f'（沙箱不可判定，非模型能力）-> 可用 {len(tasks)}') + for k in ('min_level', 'eval_min_level'): + if int(getattr(args, k, 0) or 0): + v2.logger.warning(f'[data] --{k.replace("_", "-")} 在 code 任务下无效（' + f'BigCodeBench 没有 difficulty 字段），已忽略') + recs = [{'data_id': t['task_id'], 'problem': t['instruct_prompt'], + 'reference_answer': code_task.payload_of(t)} for t in tasks] + eval_n = min(args.eval_size, len(recs)) if args.eval_size > 0 else 0 + eval_records, train_records = recs[:eval_n], recs[eval_n:] + if args.n > 0: + train_records = train_records[:args.n] + if not train_records: + raise ValueError(f'--eval-size {args.eval_size} 吃掉了整个题池（可用 {len(recs)}）') + epochs = (args.chunk_size * args.max_updates) / max(1, len(train_records)) + v2.logger.info(f'[data] train={len(train_records)} eval={len(eval_records)}；' + f'按 chunk={args.chunk_size} x max_updates={args.max_updates} 估算约 ' + f'{epochs:.1f} 个 epoch（题会重复，rubric 走缓存）') + return train_records, eval_records diff --git a/cookbook/exp/skill2lora/skill_ablate/eval_reflexion.py b/cookbook/exp/skill2lora/skill_ablate/eval_reflexion.py new file mode 100644 index 000000000..fe79faf48 --- /dev/null +++ b/cookbook/exp/skill2lora/skill_ablate/eval_reflexion.py @@ -0,0 +1,208 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""E17 专用 eval：reflexion 协议 —— 只在裸 executor 做错的题上做 rubric 条件化 skill 干预。 + +与 v2 ``run_greedy_eval`` 的区别只有"作用域"和"skill-gen 的输入"两点： + +1. baseline 正确的题**完全不动**（不生成 skill、不重跑 executor），按定义计 1.0。理由：本臂的 + 命题是"reflexion 能不能救回不会的题"，正确题上的 with-skill rollout 既不提供信息、又在 + E16 里贡献了全部 -0.060 的破坏项，把它留在指标里只会用一个已知的、与命题无关的效应稀释 + 趋势。因此当 rubric 无缺失时 acc 恰好等于 ``base + (1-base) * hard_rescue``，两者只差一个 + 线性缩放，**主指标是 hard_rescue_rate**。 +2. hard 子集的 skill-gen 输入 = query + rubric（与训练同分布），rubric 由裸解轨迹诊断得来。 + 这修掉了 E6 的已知缺陷（训练在 query+rubric 分布、eval 却是 query-only）。 + +成本：rubric 条目的缓存键 = data_id + 裸解轨迹，而 eval baseline 是冻结+缓存的，所以同一道 +eval 题在整个 run 里只会调一次 rubric API；skill-gen / executor 也只跑 hard 子集（≈40%）。 + +⚠️ 口径不与 E1-E16 横比（用户 2026-07-29 拍板：只看本实验自身趋势）。 +""" +import sys +from typing import Any, Dict, List, Tuple + +import train_skill_v2 as v2 +from train_skill_v2 import (_clean_text, _extract_skill, _run_samples, build_direct_prompt, + build_skill_solve_prompt) + +from .methods import _rubric_entry +from .rollouting import rubric_skillgen_prompt + + +def _baseline_rolls(base_sampler, eval_records, base_dp, args, base_cache) -> List[Dict[str, Any]]: + """裸 greedy(T=0) 判分，走 v2 DiskCache（与 run_greedy_eval 同一缓存文件、同一键）。""" + todo = [r for r in eval_records if v2.DiskCache.key_for(r['problem']) not in base_cache] + if todo: + out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], + 1, args.max_tokens, base_dp, temperature=0.0) + # 批量判分：code 任务下每条判分是一个跑单测的子进程，200 道题串行要 ~7 分钟。 + rolls = v2._parse_many([(v2._first_seq(seqs), r['reference_answer']) + for r, seqs in zip(todo, out)]) + for r, roll in zip(todo, rolls): + base_cache.put(v2.DiskCache.key_for(r['problem']), roll) + return [base_cache.get(v2.DiskCache.key_for(r['problem'])) for r in eval_records] + + +def _diagnose(rubric_cache, checker, jobs: List[Dict[str, Any]], workers: int) -> List[str]: + from concurrent.futures import ThreadPoolExecutor + if not jobs or rubric_cache is None: + return [''] * len(jobs) + with ThreadPoolExecutor(max_workers=max(1, min(workers, len(jobs)))) as ex: + return list(ex.map(lambda e: rubric_cache.get_or_diagnose(e, checker) or '', jobs)) + + +def _gen_skills(skill_sampler, prompts, R, skill_dp, args) -> List[List[Tuple[str, str]]]: + """每题采 R 个 skill；返回 [(skill_block, raw_response)] * R（缺位补空串，与 v2 同）。""" + sg_out = _run_samples(skill_sampler, prompts, R, args.skill_max_tokens, skill_dp, + temperature=args.eval_skill_temperature) + per = [] + for seqs in sg_out: + seqs = list(seqs or []) + row = [] + for j in range(R): + s = seqs[j] if j < len(seqs) else None + if s is None: + row.append(('', '')) + else: + sresp = _clean_text(getattr(s, 'decoded', '') or '') + row.append((_extract_skill(sresp) or '', sresp)) + per.append(row) + return per + + +def run_reflexion_eval(base_sampler, skill_sampler, eval_records, ci, rounds, + base_dp, skill_dp, args, base_cache, rubric_cache, checker): + """返回 (recs, summary, metrics)，键名与 v2.run_greedy_eval 兼容（trainer 打印共用）。""" + R = max(1, args.eval_rollouts) + base_rolls = _baseline_rolls(base_sampler, eval_records, base_dp, args, base_cache) + + # ---- 1) 切分：baseline 正确的题按协议原样通过，不做任何 GPU 工作 ---- + hard: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] + recs: List[Dict[str, Any]] = [] + head = {'record_type': 'eval_problem', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, + 'protocol': 'reflexion', 'n_rollouts': R, + 'eval_skill_temperature': args.eval_skill_temperature} + for r, br in zip(eval_records, base_rolls): + if br['correct']: + recs.append({**head, 'data_id': r.get('data_id', ''), 'problem': r['problem'], + 'reference_answer': r['reference_answer'], 'baseline_pass': 1.0, + 'intervened': False, 'rubric_ok': None, + # 协议：不干预 -> 保持 baseline 结果 + 'withskill_acc_mean': 1.0, 'withskill_acc_strict_mean': 1.0, + 'withskill_pass_any': 1.0, 'skill_parseable_mean': 1.0, + 'withskill_terminated_mean': 1.0 if br['terminated'] else 0.0}) + else: + hard.append((r, br)) + + # ---- 2) hard 子集：rubric 诊断（纯 API，缓存命中后零成本） ---- + # entry 直接复用训练侧的 _rubric_entry：判据表与 fail_segment 的构成（code 任务下含单测 + # 真实报错）必须与训练逐字同源，否则 eval 的干预分布与训练分布不同。 + entries = [_rubric_entry(r, br) for r, br in hard] + diags = _diagnose(rubric_cache, checker, entries, args.rubric_workers) + todo = [(r, br, d) for (r, br), d in zip(hard, diags) if d] + n_rubric_missing = len(hard) - len(todo) + + # ---- 3) rubric 条件化 skill-gen -> with-skill greedy 重跑 ---- + per_skills = _gen_skills(skill_sampler, [rubric_skillgen_prompt(r['problem'], d) + for r, _br, d in todo], R, skill_dp, args) \ + if todo else [] + flat_prompts, flat_idx = [], [] + for pi, ((r, _br, _d), row) in enumerate(zip(todo, per_skills)): + for j, (sk, sresp) in enumerate(row): + flat_prompts.append(build_skill_solve_prompt(r['problem'], sk, sresp)) + flat_idx.append((pi, j)) + ws_out = _run_samples(base_sampler, flat_prompts, 1, args.max_tokens, base_dp, + temperature=0.0) if flat_prompts else [] + ws_rolls = v2._parse_many([(v2._first_seq(seqs), todo[pi][0]['reference_answer']) + for (pi, _j), seqs in zip(flat_idx, ws_out)]) + roll_by = {idx: roll for idx, roll in zip(flat_idx, ws_rolls)} + + hard_recs: List[Dict[str, Any]] = [] + for pi, ((r, br, d), row) in enumerate(zip(todo, per_skills)): + rolls = [roll_by[(pi, j)] for j in range(len(row))] + corr = [1.0 if x['correct'] else 0.0 for x in rolls] + parses = [1.0 if sk else 0.0 for sk, _ in row] + terms = [1.0 if x['terminated'] else 0.0 for x in rolls] + rec = {**head, 'data_id': r.get('data_id', ''), 'problem': r['problem'], + 'reference_answer': r['reference_answer'], 'baseline_pass': 0.0, + 'intervened': True, 'rubric_ok': True, 'rubric': d, + 'base_stop_reason': br.get('stop_reason', 'none'), + # code 任务：裸失败的种类（assertion / exception / no_code / timeout / ...）。 + # 这是本臂唯一能区分"方法错"与"格式崩塌"的字段，math 下恒为 None。 + 'base_kind': br.get('kind'), + 'withskill_acc_mean': sum(corr) / len(corr) if corr else 0.0, + # strict：unparseable 计 0（回退成 direct 时会被 baseline 掩护，此处 baseline=0 + # 所以两条曲线分叉纯粹反映格式崩塌） + 'withskill_acc_strict_mean': (sum(c * p for c, p in zip(corr, parses)) / len(corr) + if corr else 0.0), + 'withskill_pass_any': 1.0 if any(corr) else 0.0, + 'skill_parseable_mean': sum(parses) / len(parses) if parses else 0.0, + 'withskill_terminated_mean': sum(terms) / len(terms) if terms else 0.0, + 'skill': row[0][0], 'skill_parseable': bool(row[0][0]), 'skill_chars': len(row[0][0]), + 'withskill_pred': rolls[0]['pred'], 'withskill_correct': rolls[0]['correct'], + 'withskill_terminated': rolls[0]['terminated'], + 'withskill_stop_reason': rolls[0]['stop_reason'], 'withskill_text': rolls[0]['text']} + hard_recs.append(rec) + # rubric 缺失的 hard 题：不进 rescue 分母（与训练侧"缺 rubric 一律丢弃"一致），但必须以 + # 显式零进 acc：它们确实没被干预、baseline 也确实错了。分母漂移靠 hard_rubric_missing + # 可审计（rubric 成功后会永久进缓存，所以只会单调收敛到 0）。 + for (r, br), d in zip(hard, diags): + if not d: + recs.append({**head, 'data_id': r.get('data_id', ''), 'problem': r['problem'], + 'reference_answer': r['reference_answer'], 'baseline_pass': 0.0, + 'intervened': False, 'rubric_ok': False, + 'base_stop_reason': br.get('stop_reason', 'none'), + 'withskill_acc_mean': 0.0, 'withskill_acc_strict_mean': 0.0, + 'withskill_pass_any': 0.0, 'skill_parseable_mean': 0.0, + 'withskill_terminated_mean': 0.0}) + recs.extend(hard_recs) + + # ---- 4) 汇总 ---- + # acc 统一取"全部 eval 行的 withskill_acc_mean 均值"，与任何下游按行求均的脚本逐字一致。 + # 不能写成 base + (1-base)*rescue：那个式子隐含"缺 rubric 的题也按 rescue 率被救"， + # 在缺失不为零时会系统高估 acc（缺失=0 时两者相等）。 + n_all = len(eval_records) + n_hard = len(hard_recs) + base = (sum(1.0 for br in base_rolls if br['correct']) / n_all) if n_all else 0.0 + acc = (sum(x['withskill_acc_mean'] for x in recs) / n_all) if n_all else 0.0 + acc_strict = (sum(x['withskill_acc_strict_mean'] for x in recs) / n_all) if n_all else 0.0 + rescue = (sum(x['withskill_acc_mean'] for x in hard_recs) / n_hard) if n_hard else 0.0 + rescue_strict = (sum(x['withskill_acc_strict_mean'] for x in hard_recs) / n_hard) if n_hard else 0.0 + rescue_any = (sum(x['withskill_pass_any'] for x in hard_recs) / n_hard) if n_hard else 0.0 + fmt = (sum(x['skill_parseable_mean'] for x in hard_recs) / n_hard) if n_hard else 0.0 + term = (sum(x['withskill_terminated_mean'] for x in hard_recs) / n_hard) if n_hard else 0.0 + # 错题结构：'length' = 没写完（E16 实测占裸失败的 97.6%），其余 = 写完但答错。 + # 本臂能不能测到"方法修正"完全取决于后者不为零，所以逐次 eval 都上报。 + trunc = (sum(1.0 for x in hard_recs if x.get('base_stop_reason') == 'length') / n_hard + if n_hard else 0.0) + # code 任务：裸失败的种类分布。math 上截断率就够（97.6% 是没写完），代码域必须分开看 + # —— 只有 assertion/exception 这类才是 rubric 有客观证据可诊断的失败，no_code/timeout + # 是格式或环境问题。键名带 kind_ 前缀，随 summary 落盘（不进 swanlab 三条主指标）。 + kind_fracs = {} + if v2._TASK == 'code' and n_hard: + for kind in ('assertion', 'exception', 'import_or_syntax', 'no_code', 'no_entry', + 'timeout'): + kind_fracs[f'hard_base_kind_{kind}'] = ( + sum(1.0 for x in hard_recs if x.get('base_kind') == kind) / n_hard) + + summary = {'record_type': 'eval_summary', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, + 'protocol': 'reflexion', 'n': n_all, 'n_rollouts': R, + 'eval_skill_temperature': args.eval_skill_temperature, + 'baseline_acc_mean1': base, 'acc_mean1': acc, 'lift_mean1': acc - base, + 'acc_strict_mean1': acc_strict, 'lift_strict_mean1': acc_strict - base, + 'format_mean1': fmt, 'term_mean1': term, + # ★ 主指标 + 'hard_n': n_hard, 'hard_rescue_rate': rescue, + 'hard_rescue_strict_rate': rescue_strict, 'hard_rescue_pass_any': rescue_any, + 'hard_rescued': sum(x['withskill_acc_mean'] for x in hard_recs), + 'hard_rubric_missing': n_rubric_missing, + 'hard_base_trunc_frac': trunc, **kind_fracs} + # swanlab 只上报三条（2026-07-30 精简，命名不缩写）。口径：只用错题子集（baseline 做对 + # 的题不干预、不计入），所以 baseline_accuracy 恒为 0、with_skill_accuracy 就是救活率。 + # 其余读数（strict / pass_any / format / term / 混合 acc）仍全量在 summary 里落盘。 + # 不带 'eval/' 前缀：trainer.py 会统一加（f'eval/{k}'），写了会变成 eval/eval/xxx。 + metrics = {'baseline_accuracy': 0.0, + 'with_skill_accuracy': rescue, + 'lift': rescue} + if n_hard and n_hard < 60: + sys.stderr.write(f'[eval] WARNING: reflexion protocol has only {n_hard} hard problems; ' + f'SE(rescue) ~ {(0.25 / n_hard) ** 0.5:.3f} — raise --eval-size.\n') + return recs, summary, metrics diff --git a/cookbook/exp/skill2lora/skill_ablate/main.py b/cookbook/exp/skill2lora/skill_ablate/main.py index 6035d2846..4bd26c3b5 100644 --- a/cookbook/exp/skill2lora/skill_ablate/main.py +++ b/cookbook/exp/skill2lora/skill_ablate/main.py @@ -10,11 +10,13 @@ --improve-skill-temperature / --skill-char-limit / --pool-max / --rubric-global-dir). """ import argparse +import dataclasses import sys +import code_task import train_skill_v2 as v2 -from .config import METHODS, STYLES, THINKINGS, get_spec, ExpSpec +from .config import METHODS, STYLES, TASKS, THINKINGS, get_spec, ExpSpec from .trainer import run_experiment @@ -25,10 +27,32 @@ def _build_args(argv=None): p.add_argument('--exp', default='', help='experiment name E1..E12 (fills method/think/style)') p.add_argument('--method', choices=METHODS, default=None) p.add_argument('--thinking', choices=THINKINGS, default=None) + p.add_argument('--executor-thinking', choices=THINKINGS, default=None, + help="executor(base_sampler) 的 thinking；默认取实验自己的 spec。'off' 是 " + 'E19/E20 的核心变量：think 的 executor 在 BigCodeBench 上 34-50% 的 ' + 'rollout 陷入字面死循环撞满预算，关掉后截断归零、裸解与 rubric 增量都更高。') p.add_argument('--style', choices=STYLES, default=None) # --- data (mirror v2) --- p.add_argument('--dataset', choices=('aops', 'math'), default='aops') + p.add_argument('--task', choices=TASKS, default=None, + help="task family; default = the experiment's own spec.task. 'code' switches " + 'the whole pipeline to BigCodeBench: executor / skill-gen / rubric prompts, ' + 'unit-test judging instead of \\boxed{} matching, and --bcb-parquet as the ' + 'data source (--deepmath-dir / --seam-parquet-dir / --dataset are ignored).') + p.add_argument('--bcb-parquet', default=code_task.DEFAULT_PARQUET, + help='BigCodeBench parquet (task=code only).') + p.add_argument('--test-workers', type=int, default=24, + help='task=code: thread pool for the unit-test subprocesses. Judging one ' + 'rollout starts a python subprocess (1-3s typical), and a chunk needs ' + 'hundreds of them — serial judging costs more wall clock than the GPU ' + 'rollouts themselves.') + p.add_argument('--test-timeout', type=int, default=60, + help='task=code: wall-clock cap per unit-test run (seconds).') + p.add_argument('--code-selftest', action=argparse.BooleanOptionalAction, default=True, + help='task=code: drop tasks whose OWN canonical solution fails their unit ' + 'tests (sandbox-undecidable, not a model failure; ~7.5% measured). ' + 'Result is cached next to the rubric cache and reused across arms.') p.add_argument('--n', type=int, default=0) p.add_argument('--exclude-data-ids', default='') p.add_argument('--seed', type=int, default=42) @@ -42,6 +66,13 @@ def _build_args(argv=None): help='train-only difficulty floor for DeepMath (eval keeps full-level mix). ' '0 = off. E1/E5 audit: level<=5 is all-pass dominated (zero gradient); ' 'recommended 6.') + p.add_argument('--eval-min-level', type=int, default=0, + help='difficulty floor for the EVAL split too (0 = off, full-level mix as in ' + 'E1-E16). Set = --min-level for arms whose readout only carries ' + 'information on problems the bare executor fails (E17): the full mix is ' + 'only ~26%% bare-wrong vs ~43%% at level>=6, so the same eval budget buys ' + '1.6x the effective sample AND matches the train distribution. Changes ' + 'the eval set -> not comparable to arms run with 0.') # --- eval口径 (4 rollouts × T=0.5, 与旧臂 E1-E13 同口径, 2026-07-28 拍板回退) --- # 曾短暂改为 SEAM val 口径(1×greedy)，为保持与已完成 6 臂可横比而回退；需要 SEAM 口径时 @@ -60,7 +91,81 @@ def _build_args(argv=None): p.add_argument('--skill-max-tokens', type=int, default=None, help='default per-experiment: 8192 (think) / 4096 (nothink); an explicit ' 'value here wins over the experiment default.') + # E14+ 稠密 reward：logp_rl 用 reward_rollouts/temperature 做 executor 采样，找正确伪 GT S； + # 之后不再 rollout 判 reward，而是算 Δ mean logP_executor(S | problem + skill)。显式传值覆盖 spec。 + p.add_argument('--reward-rollouts', type=int, default=None) + p.add_argument('--reward-temperature', type=float, default=None) + # leak 一律不进 reward（用户既定要求，反复强调）：该指标假阳性极高——DeepMath 的 gold 约一半是 + # 单字符，raw leak 率 0.53-0.65 而加 >=2 字符门后只剩 0.05-0.12，5-10 倍虚高；把它当 reward 项会 + # 用噪声主导组内排序，并且 -1.0 量级的离群值会抬高组 std、连带压小其余候选的 advantage。 + # leak 只保留监控口径（_leak_split -> leak/correct_rate、leak/wrong_rate）。默认 0 = 关闭。 + p.add_argument('--logp-leak-penalty', type=float, default=0.0, + help='E14/E15 logp reward leak penalty. DEFAULT 0 = OFF: leak is a ' + 'monitoring-only metric by project rule, never a reward term. The ' + 'unparseable-skill floor stays fixed at -1.0 regardless.') + # --- E16 passrate_hinge reward knobs ------------------------------------------------- + # 标定依据：.tmp_analysis/reward_shape_calib.py（E4 11155 条 rollout + DeepMath r1_solution_1）。 + # 关键发现：本数据集上长度本身无害——P(对|token) 在 5500 以下平在 0.96-0.98，5500-7500 微降到 + # 0.90，7500 以后断崖到 0.226（每个 difficulty 层内同形）；且 GT 参考解比模型正确答案更长 + # （p50 4377 vs 3444，GT p90 9816 已超 8192 预算）。所以“越短越好”在这里是错的，必须给死区。 + p.add_argument('--reward-trunc-penalty', type=float, default=0.12, + help='E16 alpha_len: per-rollout length penalty weight. eff *= (1 - this * ' + 'len_pen), len_pen = ((tok - lo) / (budget - lo)) ** pow, 0 below lo. ' + 'Sized so the TOTAL deduction (1+kappa)*(1-eff) stays under one pass_rate ' + 'quantum (1/M): the length signal may break ties but must never override ' + 'pass_rate, which is what eval actually measures. Its absolute size barely ' + 'matters anyway — inside an all-fail group A=(R-mean)/std rescales the ' + 'spread to unit size, so the CURVE SHAPE carries the information.') + p.add_argument('--reward-trunc-lo', type=int, default=5500, + help='E16 length dead zone: rollouts under this many executor tokens are not ' + 'penalized at all. Calibrated: P(correct|tok) is flat 0.96-0.98 below 5500.') + p.add_argument('--reward-len-pow', type=float, default=2.0, + help='E16 length ramp exponent (>1 = convex, marginal penalty grows with ' + 'length, concentrating it in the last ~1300 tokens before the budget).') + p.add_argument('--reward-loop-penalty', type=float, default=0.04, + help='E16 beta_loop: self-revision marker penalty weight. Deliberately small ' + '— inside a fixed token band ~85%% of the raw marker effect is just the ' + 'mechanical "longer output has more markers" correlation, and the same ' + 'one-pass-quantum budget is shared with the length term.') + p.add_argument('--reward-loop-lo', type=float, default=2.0, + help='E16 marker density (per 1k tokens) below which no loop penalty applies.') + p.add_argument('--reward-loop-hi', type=float, default=9.0, + help='E16 marker density at which the loop penalty saturates at 1.0.') + p.add_argument('--reward-ineff-kappa', type=float, default=0.10, + help='E16 kappa: reward -= this * mean(1 - eff). Keeps FAILING candidates ' + 'separable (a pure product collapses them all to 0 = E4 pathology). Small ' + 'on purpose: in an all-fail group A=(R-mean)/std rescales the spread back ' + 'to unit size anyway, and a large value would reverse pass_rate ordering.') + p.add_argument('--reward-leak-gate', type=float, default=0.0, + help='E16 reward leak penalty. DEFAULT 0 = OFF: leak is a monitoring-only ' + 'metric by project rule, never a reward term (see --logp-leak-penalty).') + p.add_argument('--base-tok-floor', type=int, default=5000, + help='E16 problem filter: only train problems whose baseline (no-skill) greedy ' + 'output exceeds this many tokens (probe: base_tok vs skill lift +0.62, the ' + 'strongest problem-side signal). 0 disables the filter.') + # --- E17 reflexion 臂 ----------------------------------------------------------------- + p.add_argument('--reflexion-k', type=int, default=24, + help='E17 batch alignment: train on EXACTLY this many bare-wrong problems per ' + 'chunk. The chunk is drawn at --chunk-size, bare-solved, and wrong ones ' + 'are taken (with rubric backfill) until K is reached; dynamic ' + 'over-drawing is impossible because resume replays fixed-size pool ' + 'draws. Group count per update must be constant or the step size, the ' + 'noise floor and the online SNR readout all move with the draw. ' + 'DEFAULT 24 matches the 23.46 groups/update E16 actually ran (1173 ' + 'total): cumulative evidence scales as sqrt(N) and E16 only reached ' + '1.9 sigma on any text-level direction, so a smaller K has no ' + 'discriminating power left. Needs --chunk-size ~5x (bare error rate ' + '0.329 measured over the FULL level>=6 pool, 527/1600).') p.add_argument('--align-mode', choices=('v2', 'seam'), default='v2') + # --- E18 rejection_sft 臂 --------------------------------------------------------------- + p.add_argument('--e18-accumulate', type=int, default=16, + help='E18: fire one SFT update only after this many accepted (rejection-' + 'sampled) skills have accumulated in the pool. Winners are also ' + 'appended to /e18_sft_dataset.jsonl for offline reuse. ' + 'DEFAULT 16 (2026-07-30 拍板) = one sft_batch_size, so a chunk of 32 ' + '(bare error rate 0.329 -> ~10 accepted) fires roughly every other ' + 'chunk and --max-updates 50 is reachable; 128 would need ~15 chunks ' + 'per update.') p.add_argument('--len-budget', type=int, default=None, help='regen skill length target (chars). Default per style (ablation stats ' '#9): narrative~1100 / pitfall~300. Used to pick the regen survivor.') @@ -86,7 +191,19 @@ def _build_args(argv=None): p.add_argument('--ppo-mini-batch-size', type=int, default=0) p.add_argument('--grpo-epsilon', type=float, default=0.2) p.add_argument('--adv-clip', type=float, default=0.0) - p.add_argument('--kl-beta', type=float, default=0.001) + # 对初始策略的锚。漂移分析（.tmp_analysis/why_no_correction.py 等）显示：“executor 能干净收束” + # 是初始 skill 分布自带的脆弱属性，reward 里能反对它被磨耗的可迁移成分只有 11%，因此 + # 锚本身就是一个直接对症的旋钮（选择压力只能在组内排序，锚才能提供恢复力）。 + # 2026-07-29 拍板：默认 0.001 -> 0.01。⚠️ 注意锚是无方向的刹车，它同等抵制那个方向 + # 正确的 +0.078 收束签名比较信号；取值未经标定（既往臂全部跑在 0.001，无可用对照）。 + p.add_argument('--kl-beta', type=float, default=0.01, + help='KL anchor to the reference (initial, never-synced) policy. Raised from ' + '0.001 to 0.01 on 2026-07-29 to oppose the intrinsic drift that erodes ' + 'executor termination. Recorded in the config fingerprint. Note: the ' + 'anchor is undirected -- it also brakes the (weak) useful gradient.') + p.add_argument('--run-tag', default='', + help='suffix for the swanlab experiment name, to tell apart variants that share ' + 'the same ExpSpec (e.g. RUN_TAG=kl01 for the --kl-beta 0.01 arm).') p.add_argument('--lr', type=float, default=1e-6, help='stable 1e-6, no warmup / no decay (ablation spec).') p.add_argument('--sft-weight', type=float, default=1.0, @@ -131,6 +248,19 @@ def _build_args(argv=None): if not (args.method and args.thinking and args.style): p.error('provide --exp E5, or all of --method/--thinking/--style') spec = ExpSpec(name='Ex', method=args.method, thinking=args.thinking, style=args.style) + # 显式 --task 覆盖 spec（ExpSpec 是 frozen dataclass）；用于临时把某个数学臂放到 code 上跑。 + if args.task and args.task != spec.task: + sys.stderr.write(f'[ablate] WARNING: --task {args.task} overrides {spec.name} ' + f'spec.task={spec.task}\n') + spec = dataclasses.replace(spec, task=args.task) + args.task = spec.task + # executor thinking 同理（frozen dataclass -> replace）。v2.build_* 读 args.executor_thinking。 + if args.executor_thinking and args.executor_thinking != spec.executor_thinking: + sys.stderr.write(f'[ablate] WARNING: --executor-thinking {args.executor_thinking} ' + f'overrides {spec.name} spec.executor_thinking=' + f'{spec.executor_thinking}\n') + spec = dataclasses.replace(spec, executor_thinking=args.executor_thinking) + args.executor_thinking = spec.executor_thinking # sanity: Ray dp rule — sft/eval batch must divide TRAIN_DP if args.sft_batch_size % v2.TRAIN_DP != 0: @@ -139,15 +269,52 @@ def _build_args(argv=None): p.error(f'--train-micro-batch ({args.train_micro_batch}) must be a multiple of TRAIN_DP ({v2.TRAIN_DP})') if args.chunk_size < 1: p.error('--chunk-size must be >= 1') + if spec.method == 'reflexion': + if args.reflexion_k < 1: + p.error('--reflexion-k must be >= 1') + if args.reflexion_k > args.chunk_size: + p.error(f'--reflexion-k ({args.reflexion_k}) > --chunk-size ({args.chunk_size}): the ' + f'aligned batch can never be filled') + # 2-sigma 余量检查：错题数 ~ Binom(chunk, p_wrong)。余量不够时组数会随 chunk 抽样抖， + # 而组数恒定是本臂的硬要求。⭐ math: 0.329 = E16 全量题池实测（527/1600）；切勿用 + # 0.446，那是 base_tok>5000 筛选后子集的错误率（偏难），而本臂 floor=0。 + # ⭐ code: 0.56 —— 不是 probe 的 0.715。probe 的 pass 0.285 是 nothink + max_tokens 4096 + # 的读数；本臂 executor 是 think=on + 8192，同一批题实测裸错率只有 18/32=0.5625 + # （2026-07-31 dry run：训练侧 10/16、eval 侧 8/16）。用 0.715 会把 chunk 需求算小 + # 一半：P(48 道里凑不满 24) 在 0.715 下是 0.000，在 0.5625 下是 0.154。 + # ⭐ executor nothink（E19/E20）另算：code nothink 探针实测 pass 0.378 -> p_wrong 0.622； + # math nothink 实测 0.31（首个 E19 run 的 c0：64 道里 20 道错，baseline acc≈0.69）， + # 与 think 的 0.329 基本相同 —— 曾按 0.85 保守估是错的，会把 chunk 需求算小一半。 + if spec.executor_thinking == 'off': + _p_wrong = 0.622 if spec.task == 'code' else 0.31 + else: + _p_wrong = 0.5625 if spec.task == 'code' else 0.329 + _mu = args.chunk_size * _p_wrong + _sd = (args.chunk_size * _p_wrong * (1 - _p_wrong)) ** 0.5 + if args.reflexion_k > _mu - 2 * _sd: + sys.stderr.write( + f'[ablate] WARNING: --chunk-size {args.chunk_size} gives {_mu:.1f}+-{_sd:.1f} ' + f'wrong problems, less than 2 sigma of headroom over --reflexion-k ' + f'{args.reflexion_k}; expect signal/k_short > 0 and a drifting group count. ' + f'Use --chunk-size >= {int((args.reflexion_k + 2 * _sd) / _p_wrong) + 1}.\n') + if spec.method == 'rejection_sft': + if args.e18_accumulate < 1: + p.error('--e18-accumulate must be >= 1') + # 池 batch 整体交给 _train_batch，drop_last 到 TRAIN_DP 倍数会静默丢尾部真样本； + # 强制倍数关系把丢样本量钉在 0。 + if args.e18_accumulate % v2.TRAIN_DP != 0: + p.error(f'--e18-accumulate ({args.e18_accumulate}) must be a multiple of ' + f'TRAIN_DP ({v2.TRAIN_DP})') args.rubric_global_dir = args.rubric_global_dir or None return args, spec def main(argv=None): args, spec = _build_args(argv) - sys.stderr.write(f'[ablate] running {spec.name}: view={spec.view} method={spec.method} ' - f'thinking={spec.thinking} style={spec.style} loss={spec.loss} ' - f'smt={spec.skill_max_tokens} -> {args.output_dir}\n') + sys.stderr.write(f'[ablate] running {spec.name}: task={spec.task} view={spec.view} ' + f'method={spec.method} thinking={spec.thinking} style={spec.style} ' + f'executor_thinking={spec.executor_thinking} ' + f'loss={spec.loss} smt={spec.skill_max_tokens} -> {args.output_dir}\n') run_experiment(args, spec) diff --git a/cookbook/exp/skill2lora/skill_ablate/methods.py b/cookbook/exp/skill2lora/skill_ablate/methods.py index b1a6cd5c4..ebe8cd468 100644 --- a/cookbook/exp/skill2lora/skill_ablate/methods.py +++ b/cookbook/exp/skill2lora/skill_ablate/methods.py @@ -30,6 +30,11 @@ OPSD. In the OPSD path no ref forward is done at all: OPSDLoss uses only teacher_logps (kl_beta / ref_logps play no role there). """ +import json +import os +import re +import sys +from collections import Counter from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from typing import Any, Callable, Dict, List, Optional, Tuple @@ -89,12 +94,22 @@ def _bare_solve(ctx: MethodContext, records: List[Dict[str, Any]]) -> List[Dict[ """Bare-problem greedy (T=0) executor solve; returns one roll per record (order-aligned).""" out = _run_samples(ctx.base_sampler, [build_direct_prompt(r['problem']) for r in records], 1, ctx.args.max_tokens, ctx.base_dp, temperature=0.0) - return [(_parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll()) - for r, seqs in zip(records, out)] + return v2._parse_many([(v2._first_seq(seqs), r['reference_answer']) + for r, seqs in zip(records, out)]) def _rubric_entry(record: Dict[str, Any], roll: Dict[str, Any]) -> Dict[str, Any]: - """Build the entry _diagnose_entry expects from a failure trajectory.""" + """Build the entry _diagnose_entry expects from a failure trajectory. + + code 任务：fail_segment 不是"输出全文"，而是**提交的代码 + 单测真实报错**（异常类型 / + 断言差异 / 失败用例名）。这是三个数据集横比下 rubric 唯一真正产生增量的原因 —— 数学与 + BFCL 上 judge 手里没有任何客观证据，只能猜，命中率≈随机。 + """ + if v2._TASK == 'code': + return {'problem': record['problem'], 'reference_answer': record['reference_answer'], + 'data_id': record.get('data_id', ''), + 'fail_segment': v2.code_task.diag_segment(roll), + 'fail_stop_reason': roll.get('stop_reason', 'none')} return {'problem': record['problem'], 'reference_answer': record['reference_answer'], 'data_id': record.get('data_id', ''), 'fail_segment': roll.get('text', ''), @@ -142,11 +157,15 @@ def _skillgen_solve(ctx: MethodContext, items: List[Dict[str, Any]], n_skills: i for it, c in flat: c['leaked'] = _answer_leaked(c['skills'], it['record']['reference_answer']) if flat: + # V2 fix: pass the raw skill-gen response like v2 process_chunk does — seam align + # nests the actor's full response_text into the executor prompt (no-op in v2 mode). ws = _run_samples(ctx.base_sampler, - [build_skill_solve_prompt(it['record']['problem'], c['skills']) for it, c in flat], + [build_skill_solve_prompt(it['record']['problem'], c['skills'], c.get('response')) + for it, c in flat], 1, args.max_tokens, ctx.base_dp, temperature=0.0) - for (it, c), seqs in zip(flat, ws): - roll = _parse_seq(seqs[0], it['record']['reference_answer']) if seqs else _empty_roll() + judged = v2._parse_many([(v2._first_seq(seqs), it['record']['reference_answer']) + for (it, _c), seqs in zip(flat, ws)]) + for (it, c), roll in zip(flat, judged): c['rolls'] = [roll] c['with_pass'] = 1.0 if roll['correct'] else 0.0 c['reward'] = _skill_reward(c['parseable'], roll['correct']) @@ -172,6 +191,15 @@ def _grpo_records(records: List[Dict[str, Any]], with_rubric: bool) -> List[Dict return recs +def _leak_blocks(skill: str, reference) -> bool: + """Filtering-grade leak gate (bugfix #4): only answers >=2 chars are informative — the + same gate E14/E16 already apply to their reward penalty. Single-char golds ('2' etc.) + substring-match ordinary math prose (~84% false positives measured), which starved the + SFT-family pools by discarding nearly every regen candidate. Monitoring paths keep the + raw ``_answer_leaked`` so the recorded leak/rate 口径 is unchanged.""" + return len(str(reference).strip()) >= 2 and _answer_leaked(skill, reference) + + def _leak_split(pairs: List[Tuple[bool, bool]]) -> Dict[str, float]: """#10 monitoring curves: leaked&correct vs leaked&wrong rates over parseable skills. "泄露正确答案可接受"不等于"泄露无害"——有害的是错误数值注入，两条曲线拆开监控。""" @@ -183,7 +211,9 @@ def _leak_split(pairs: List[Tuple[bool, bool]]) -> Dict[str, float]: def _cand_leak_pairs(records: List[Dict[str, Any]]) -> List[Tuple[bool, bool]]: - return [(bool(c['leaked']), bool(c['with_pass'])) + # bugfix #7: with_pass is a float pass RATE under M>1 rollouts (E16) — bool(0.25) would + # count a partial pass as "correct"; compare > 0 instead (identical for greedy 0/1 arms). + return [(bool(c['leaked']), (c['with_pass'] or 0) > 0) for r in records for c in r.get('_cands', []) if c.get('parseable') and c.get('with_pass') is not None] @@ -273,6 +303,69 @@ def _gather_response_logps(full_logps, pos_lists) -> List[List[float]]: return rows +# =========================================================================================== +# E14+ helpers: executor pseudo-GT + dense logP reward +# =========================================================================================== +def _executor_answer_trajectory(problem: str, skill: str, answer_text: str, + raw_response: Optional[str] = None) -> Dict[str, Any]: + """Teacher-forcing trajectory for executor logP(S | problem + skill).""" + msgs = [dict(m) for m in build_skill_solve_prompt(problem, skill, raw_response)['messages']] + return {'messages': msgs + [{'role': 'assistant', 'content': answer_text}], + 'user_data': {'key_rounds': [len(msgs)]}} + + +def _set_ref_executor_template(ctx: MethodContext, enable_thinking: bool) -> None: + """Temporarily reuse ref_model as frozen executor scorer, then restore skill/ref layout.""" + ctx.ref_model.set_template(v2.Template, model_id=v2.MODEL_ID, + enable_thinking=enable_thinking, + max_length=ctx.args.max_model_len, + truncation_strategy='delete') + + +def _mean_logp_rows(full_logps, pos_lists) -> List[float]: + rows = _gather_response_logps(full_logps, pos_lists) + return [(sum(x) / len(x)) if x else float('-inf') for x in rows] + + +def _score_executor_mean_logps(ctx: MethodContext, trajs: List[Dict[str, Any]]) -> List[Optional[float]]: + """Return mean response-token logP under the frozen executor template; None means truncated.""" + tmpl = ctx.encode_template + assert tmpl is not None, 'logp_rl needs ctx.encode_template' + out: List[Optional[float]] = [None] * len(trajs) + valid_trajs, pos_lists, valid_idx = [], [], [] + for i, tr in enumerate(trajs): + enc = tmpl.encode(tr) + if enc is None: + continue + pos = np.where(np.asarray(enc.get('labels')) != -100)[0] + if len(pos) == 0: + continue + valid_trajs.append(tr) + pos_lists.append(pos) + valid_idx.append(i) + if not valid_trajs: + return out + sft = getattr(ctx.args, 'train_micro_batch', 0) or ctx.args.sft_batch_size + dp = max(1, v2.REF_DP) + _set_ref_executor_template(ctx, enable_thinking=True) + try: + for st in range(0, len(valid_trajs), sft): + mb = valid_trajs[st:st + sft] + mb_pos = pos_lists[st:st + sft] + n_mb = len(mb) + # forward_only 按 slice_dp 切分，非 dp 整倍数的尾批用末尾样本补齐，输出只取前 n_mb 行 + if n_mb % dp: + pad = dp - (n_mb % dp) + mb = mb + [mb[-1]] * pad + mb_pos = list(mb_pos) + [mb_pos[-1]] * pad + vals = _mean_logp_rows(ctx.ref_model.forward_only(inputs=mb).get('logps'), mb_pos)[:n_mb] + for j, val in enumerate(vals): + out[valid_idx[st + j]] = float(val) + finally: + _set_ref_executor_template(ctx, enable_thinking=(ctx.args.skill_thinking == 'on')) + return out + + def _train_batch(ctx: MethodContext, samples: List[Dict[str, Any]], traj_fn: Callable[[Dict[str, Any]], Dict[str, Any]], teacher_fn: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None, @@ -413,14 +506,20 @@ def step(self, chunk, ci): finally: diag_pool.shutdown(wait=False) a_items = [] + degraded = [] # bugfix #2: rubric 缺失（API 失败/坏缓存）→ 降级 query-only B 线，绝不训练空 rubric prompt for (r, _roll), diag in zip(wrong, diags): r['_rubric'] = diag - a_items.append({'record': r, 'prompt': rubric_skillgen_prompt(r['problem'], diag)}) + if diag: + a_items.append({'record': r, 'prompt': rubric_skillgen_prompt(r['problem'], diag)}) + else: + degraded.append({'record': r, 'prompt': _skillgen_prompt(r['problem'])}) + if degraded: + _skillgen_solve(ctx, degraded, args.n_skills, temperature=args.skill_gen_temperature) if a_items: _skillgen_solve(ctx, a_items, args.n_skills, temperature=args.skill_gen_temperature) _assign_advantages(chunk, args) - a_recs = _grpo_records([r for r, _ in wrong], with_rubric=True) - b_recs = _grpo_records(right, with_rubric=False) + a_recs = _grpo_records([it['record'] for it in a_items], with_rubric=True) + b_recs = _grpo_records(right + [it['record'] for it in degraded], with_rubric=False) train_recs = a_recs + (b_recs if self.train_b else []) has_signal = any(abs(s['advantage']) > 1e-9 for s in train_recs) if has_signal and getattr(args, 'drop_zero_adv', False): @@ -432,8 +531,9 @@ def step(self, chunk, ci): traj_fn=lambda s: (rubric_train_trajectory(s) if s['with_rubric'] else query_only_train_trajectory(s))) return {'n_updates': n_upd, - 'metrics': {'signal/n_wrong_A': float(len(wrong)), - 'signal/n_right_B': float(len(right)), **tmetrics, + 'metrics': {'signal/n_wrong_A': float(len(a_items)), + 'signal/n_right_B': float(len(right)), + 'signal/n_rubric_missing': float(len(degraded)), **tmetrics, **_cand_pass_metrics(chunk), **_leak_split(_cand_leak_pairs(chunk))}, 'gen_records': v2._full_records(chunk, ci)} @@ -501,11 +601,572 @@ def step(self, chunk, ci): for r, _resp, sk in first]} +class LogpRlMethod(TrainMethod): + """E14+: rubric-audited pseudo-GT + executor logP dense reward. + + Per problem, sample K executor attempts without skill at T>0, keep the first locally-correct + and non-truncated response S, audit it through the rubric API/cache, then score each generated + skill by Δ mean logP_executor(S | problem + skill). Problems with no S are skipped this round. + """ + needs_rubric = True + + def step(self, chunk, ci): + usable, sg = self._prepare(chunk, ci) + return self._score_and_train(chunk, ci, usable, sg) + + def _skillgen(self, usable): + ctx = self.ctx + args = ctx.args + return _run_samples(ctx.skill_sampler, [_skillgen_prompt(r['problem']) for r in usable], + args.n_skills, args.skill_max_tokens, ctx.skill_dp, + temperature=args.skill_gen_temperature, top_p=args.skill_gen_top_p, + top_k=args.skill_gen_top_k) + + def _prepare(self, chunk, ci): + """executor T>0 采 K 条 -> 选本地判分正确且非截断的伪 GT S -> rubric API 后台审计 + （与 skill-gen 的 GPU rollout 重叠）。返回 (usable, 每题 skill-gen 序列)。""" + ctx = self.ctx + args = ctx.args + for r in chunk: + r['_cands'] = [] + K = max(1, int(getattr(args, 'reward_rollouts', 1) or 1)) + temp = float(getattr(args, 'reward_temperature', 0.0) or 0.0) + out = _run_samples(ctx.base_sampler, [build_direct_prompt(r['problem']) for r in chunk], + K, args.max_tokens, ctx.base_dp, temperature=temp) + usable, audit_jobs = [], [] + for r, seqs in zip(chunk, out): + rolls = [_parse_seq(s, r['reference_answer']) for s in (seqs or [])] + ok = next((x for x in rolls if x['correct'] and x.get('stop_reason') != 'length'), None) + r['_pseudo_rolls'] = rolls + if ok is None: + continue + r['_pseudo_roll'] = ok + r['_pseudo_solution'] = ok.get('text', '') + usable.append(r) + audit_jobs.append((_rubric_entry(r, ok), ok.get('text', ''))) + # rubric 审计是纯 API：后台线程跑，与 skill-gen 的 GPU rollout 重叠（API/GPU overlap） + diag_pool = ThreadPoolExecutor(max_workers=1) + diag_fut = diag_pool.submit(_diagnose_parallel, ctx, audit_jobs) + try: + sg = self._skillgen(usable) + diags = diag_fut.result() + finally: + diag_pool.shutdown(wait=False) + for r, diag in zip(usable, diags): + r['_rubric'] = diag + return usable, sg + + def _score_and_train(self, chunk, ci, usable, sg): + ctx = self.ctx + args = ctx.args + flat = [] + for r, seqs in zip(usable, sg): + for si, s in enumerate(seqs or []): + resp = _clean_text(getattr(s, 'decoded', '') or '') + block = _extract_skill(resp) or '' + pseudo = dict(r['_pseudo_roll']) + if si > 0: + pseudo['text'] = '' # 磁盘保护：伪 GT 全文每题只在首个候选保留一份 + cand = {'skills': block, 'response': resp, 'parseable': bool(block), + 'leaked': None, 'with_pass': None, 'reward': None, 'rolls': [pseudo], + 'advantage': 0.0, 'kept': False, + 'skillgen_stop': getattr(s, 'stop_reason', None), + 'skillgen_tokens': len(getattr(s, 'tokens', None) or []), + 'logp_base': None, 'logp_skill': None, 'logp_delta': None} + r['_cands'].append(cand) + if block: + cand['leaked'] = _answer_leaked(block, r['reference_answer']) + flat.append((r, cand)) + for c in r['_cands']: + if c['leaked'] is None: + c['leaked'] = False + if flat: + base_trajs = [_executor_answer_trajectory(r['problem'], '', r['_pseudo_solution']) + for r in usable] + base_logps = _score_executor_mean_logps(ctx, base_trajs) + base_by_id = {id(r): lp for r, lp in zip(usable, base_logps)} + cand_trajs = [_executor_answer_trajectory(r['problem'], c['skills'], r['_pseudo_solution'], + c.get('response')) for r, c in flat] + cand_logps = _score_executor_mean_logps(ctx, cand_trajs) + # bugfix #14: logP 目标超长被 truncation='delete' 删掉时 lp=None → reward 地板， + # 整组塔到 -1.0 会零梯度空转；这里显式监控 encode 失败占比（E15 的 R1 参考解尤其长）。 + _n_enc = len(base_logps) + len(cand_logps) + enc_fail_frac = ((sum(1 for v in base_logps if v is None) + + sum(1 for v in cand_logps if v is None)) / _n_enc) if _n_enc else 0.0 + # leak 一律不进 reward（项目既定要求）：--logp-leak-penalty 默认 0，leaked 只做监控。 + # 该指标假阳性极高（单字符 gold 误报 ~84%，c0 实测 raw leak/rate=0.889），-1.0 量级是 + # delta 信号的 50~100 倍，会用噪声主导组内 advantage 并抬高组 std 压小其余候选。 + # format 地板（unparseable / logP 编码失败）固定 -1.0，与 leak 完全解耦、不受影响。 + floor = 1.0 + leak_pen = abs(float(getattr(args, 'logp_leak_penalty', 0.0))) + for (r, c), lp in zip(flat, cand_logps): + base_lp = base_by_id.get(id(r)) + c['logp_base'], c['logp_skill'] = base_lp, lp + if base_lp is None or lp is None: + c['reward'] = -floor + continue + delta = float(lp) - float(base_lp) + c['logp_delta'] = delta + # leak_pen 默认 0（leak 只做监控口径，不进 reward）。若显式开启，仍只对 >=2 字符的 + # 答案生效：usable 子集答案多为 0/1/2/4 等单字符，_answer_leaked 子串匹配在正常数学 + # 叙述里误报率 ~84%。leaked 字段全量记录，监控口径不变。 + informative = len(str(r['reference_answer']).strip()) >= 2 + c['reward'] = delta - (leak_pen if (c.get('leaked') and informative) else 0.0) + for r in chunk: + for c in r.get('_cands', []): + if c['reward'] is None: + c['reward'] = -1.0 + _assign_advantages(chunk, args) + grpo = _grpo_records(usable, with_rubric=False) + has_signal = any(abs(s['advantage']) > 1e-9 for s in grpo) + if has_signal and getattr(args, 'drop_zero_adv', False): + grpo = [s for s in grpo if abs(s['advantage']) > 1e-9] + n_upd, tmetrics = 0, {} + if has_signal: + n_upd, tmetrics = _train_batch(ctx, grpo, traj_fn=query_only_train_trajectory) + summary = v2._chunk_summary(chunk, ci) + # logp_rl 的 reward 不是 executor rollout 通过率；覆盖旧 BNPO summary 中与 pass 绑定的字段， + # 避免 train_log 把“reward 非零”误读成 with-skill pass。 + summary['avg_withskill_pass'] = 0.0 + summary['candidate_withskill_pass'] = 0.0 + summary['withskill_trunc_frac'] = 0.0 + summary['termination_rate_withskill'] = 0.0 + rewards = [c['reward'] for r in usable for c in r.get('_cands', [])] + deltas = [c['logp_delta'] for r in usable for c in r.get('_cands', []) + if c.get('logp_delta') is not None] + metrics = {'signal/pseudo_gt_rate': float(len(usable)) / max(1, len(chunk)), + 'signal/n_pseudo_gt': float(len(usable)), + 'signal/zero_grad_frac': summary['zero_grad_frac'], + 'logp/encode_fail_frac': (enc_fail_frac if flat else 0.0), + 'logp/reward_mean': v2._mean(rewards), + 'logp/reward_std': v2._std(rewards), + 'logp/delta_mean': v2._mean(deltas), + 'logp/delta_std': v2._std(deltas), + 'leak/rate': summary['leak_rate'], **tmetrics, + **_leak_split(_cand_leak_pairs(chunk))} + return {'n_updates': n_upd, 'summary': summary, 'metrics': metrics, + 'gen_records': v2._full_records(chunk, ci)} + + +class LogpGtMethod(LogpRlMethod): + """E15: identical dense executor-logP reward, but the target S is DeepMath's external R1 + reference solution (record 'solution'), NOT an executor-sampled pseudo-GT. No executor + rollout and no rubric audit -> much faster; every problem carrying a solution is usable.""" + needs_rubric = False + + def _prepare(self, chunk, ci): + for r in chunk: + r['_cands'] = [] + usable = [] + for r in chunk: + sol = (r.get('solution') or '').strip() + r['_pseudo_rolls'] = [] + if not sol: + continue + # 合成 roll：logP 目标是外部 R1 参考解，无 executor 采样；stop_reason='gt' 仅作标记， + # 字段与 _parse_seq 输出对齐（_roll 序列化需 pred/correct/terminated/stop_reason/gen_tokens/text）。 + r['_pseudo_roll'] = {'pred': str(r['reference_answer']), 'correct': True, + 'terminated': True, 'stop_reason': 'gt', + 'gen_tokens': 0, 'text': sol} + r['_pseudo_solution'] = sol + r['_rubric'] = '' + usable.append(r) + return usable, self._skillgen(usable) + + +def _hinge_trunc(rolls: List[Dict[str, Any]], lo: int, budget: int = 8192) -> float: + """Mean hinge over rollouts: 0 below ``lo`` tokens, ramps to 1.0 at the ``budget`` line. + Probe (skill_quality_analysis.md 2026-07-29): the length->correctness link exists ONLY near + the truncation budget (zero-truncation groups show +0.03), so penalize the danger zone only, + smoothly, giving gradient BEFORE the hard 'length' cutoff fires.""" + if not rolls: + return 0.0 + span = max(1, budget - lo) + return sum(max(0.0, (int(x.get('gen_tokens') or 0) - lo) / span) for x in rolls) / len(rolls) + + +# 自我推翻标记词。词表由 .tmp_analysis/reward_shape_calib.py 在 E4 11155 条 rollout 上判别力筛出 +# （错误密度/正确密度比值）：confusing 3.06、contradiction 2.49、mistake|error|wrong 1.84、 +# alternatively 1.59。刻意排除的词："let me check" 比值 0.75（**反向**——检查一次是好行为，罚它 +# 有害）、"hmm" 0.86、"recompute|again" 1.00（零信号）、"wait" 仅 1.51 且覆盖率 0.993（几乎人人都写， +# 区分度低）。 +_LOOP_MARKER_RE = re.compile( + r'\b(?:confusing|confused|contradiction|contradicts|contradictory|mistake|error|wrong' + r'|alternatively)\b', re.I) + + +def _loop_density(roll: Dict[str, Any]) -> float: + """Self-revision marker density, occurrences per 1000 generated tokens.""" + tok = int(roll.get('gen_tokens') or 0) + if tok <= 0: + return 0.0 + return len(_LOOP_MARKER_RE.findall(roll.get('text') or '')) / tok * 1000.0 + + +# --- skill 套话度监控 ------------------------------------------------------------------- +# ★ 定位：这两条是退化监控器（看趋势），不是质量预测器（看绝对值）。 +# 2026-07-30 在 E17 的1246 个真实候选上试过 7 种定义（.tmp_analysis/generic_index_*.py）， +# 题内配对对 pass 的预测力全部在 -0.033 到 +0.012 之间，都弱。根因已查清：narrative +# 文体里“元指令”和“具体动作”是同一句话里交织的（“Avoid rechecking the same rounding +# or error estimates”既是元指令又指向具体对象），句子级二分类根本不成立，调词表无法解决。 +# 保留词频法的依据是方向正确：指数低组 leak=0.610 / 高组 leak=0.474，即套话越多越不给 +# 具体内容。但绝对值偏高（`avoid` 在本数据上命中 1223 次，大部分在具体建议里）， +# ★ 因此只能比较同一根曲线的前后变化，不能拿绝对值评判 skill 好坏。 +_GENERIC_ADVICE_RE = re.compile( + r'\b(carefully|careful|make sure|makes sure|ensure|ensuring|be sure|avoid|avoiding|' + r'remember|keep in mind|bear in mind|double[- ]check|double[- ]checking|verify|verifying|' + r'validate|validating|consider|considering|systematically|systematic|efficiently|efficient|' + r'properly|correctly|accurately|appropriately|appropriate|rigorously|rigorous|' + r'manage|managing|track|tracking|monitor|monitoring|streamline|streamlining|' + r'focus on|stay focused|trust|commit to|committing|self[- ]correct\w*|' + r'step by step|methodical\w*|thorough\w*|concise\w*|precise\w*|' + r'token budget|length budget|within the budget|redundan\w*|unnecessar\w*)\b', re.I) +_WORD_RE = re.compile(r"[A-Za-z][A-Za-z'\-]*") +_SENT_RE = re.compile(r'(?<=[.!?;:])\s+|\n+') +# “句中有没有数学内容”的锚点：数字 | LaTeX/算式 | 句中大写专名（定理名）。 +# 不维护数学名词/动词词表：V5 试过，加了两张表反而把具体洞察误判成空话、预测力降到 0.000。 +_MATH_ANCHOR_RE = re.compile(r'\d|\\[A-Za-z]+|\$|\^|_\{|=|≤|≥|≠|∈|∑|∏|√|(? float: + """泛泛建议词占 skill 总词数的比例；空文本返回 0。只看趋势，见上方注释。""" + words = _WORD_RE.findall(text or '') + if not words: + return 0.0 + return len(_GENERIC_ADVICE_RE.findall(text)) / len(words) + + +def _no_math_sentence_fraction(text: str) -> float: + """不含任何数学内容（数字/算式/定理名）的句子占比。 + + 名字只陈述它实际测的东西，不声称它能判“空话”—— 实测题内配对 pass 只有 -0.016， + 但它不经由 leak 中介（词频法那一条经由），两条一起看才能分开“套话变多”与“不给答案”。 + """ + sents = [s.strip() for s in _SENT_RE.split(text or '') if len(s.strip()) >= 15] + if not sents: + return 0.0 + return sum(1 for s in sents if not _MATH_ANCHOR_RE.search(s)) / len(sents) + + +def _digit_fraction(text: str) -> float: + """数字字符占 skill 总字符数的比例（递答案/递具体中间量的代理指标）。""" + t = text or '' + if not t: + return 0.0 + return sum(1 for ch in t if ch.isdigit()) / len(t) + + +def _skill_text_metrics(skills: List[str]) -> Dict[str, float]: + """skill 文本面板（按任务分派）。 + + math: 数字占比 / 泛泛建议词占比 / 无数学内容句占比。 + code: 数字占比与"句中有没有数学锚点"在代码域没有语义（API 名、参数、类型天然带大写与 + 符号），换成 skill_contains_code_fraction —— skill 本该是方法论，写出代码围栏或成段 + def/return 就是退化成抄实现，这是代码域最该盯的那条退化曲线。泛泛建议词那条保留： + 它的词表是任务无关的（carefully / make sure / step by step ...）。 + """ + generic = v2._mean([_generic_advice_fraction(s) for s in skills]) + if v2._TASK == 'code': + return {'train/skill_generic_advice_fraction': generic, + 'train/skill_contains_code_fraction': v2._mean( + [1.0 if v2.code_task.skill_has_code(s) else 0.0 for s in skills])} + return {'train/skill_digit_fraction': v2._mean([_digit_fraction(s) for s in skills]), + 'train/skill_generic_advice_fraction': generic, + 'train/skill_no_math_sentence_fraction': v2._mean( + [_no_math_sentence_fraction(s) for s in skills])} + + +# --- E18 拒绝采样第三道筛：skill 与 rubric 诊断的词频余弦 -------------------------------- +# 定位：在"executor 已做对"的候选里挑与诊断内容对得上的那条，压掉两类假赢家—— +# 与诊断无关的碰巧做对（含泄露式速通的残余）和与谁都不像的空泛套话。 +# 刻意用去停用词的词频余弦而不是 tfidf/语义模型：可迁移性判别器一节实测"仅 tfidf" +# in-sample 0.983 / OOS 0.541 是纯过拟合；词频余弦纯 stdlib、确定性、可离线复算。 +_SIM_STOPWORDS = frozenset( + 'the a an and or of to in is are be for with that this it on as by from at not no was ' + 'were will would can could should may might do does did have has had you your we they ' + 'he she its if then than so but into over under out up down when where which what how ' + 'why all any each more most other some such only own same very'.split()) +_SIM_WORD_RE = re.compile(r"[a-z][a-z'\-]{2,}") + + +def _rubric_similarity(skill: str, rubric: str) -> float: + """内容词词频余弦 ∈ [0,1]；任一侧无内容词返回 0。""" + ca = Counter(w for w in _SIM_WORD_RE.findall((skill or '').lower()) + if w not in _SIM_STOPWORDS) + cb = Counter(w for w in _SIM_WORD_RE.findall((rubric or '').lower()) + if w not in _SIM_STOPWORDS) + if not ca or not cb: + return 0.0 + dot = float(sum(v * cb[k] for k, v in ca.items() if k in cb)) + na = sum(v * v for v in ca.values()) ** 0.5 + nb = sum(v * v for v in cb.values()) ** 0.5 + return dot / (na * nb) if na and nb else 0.0 + + +def _efficiency_terms(rolls: List[Dict[str, Any]], *, budget: int, len_lo: int, len_pow: float, + alpha_len: float, beta_loop: float, loop_lo: float, loop_hi: float + ) -> Tuple[float, float, Dict[str, float]]: + """Per-rollout efficiency factor, aggregated. Returns (mean_score, mean_inefficiency, diag). + + Calibration (.tmp_analysis/reward_shape_calib.py, E4 11155 rollouts + DeepMath r1_solution_1): + * Length is NOT harmful per se on this dataset. P(correct | tokens) is flat at 0.96-0.98 up + to 5500 tokens, dips to 0.932/0.903 at 5500-7500, then collapses to 0.226 past 7500 — and + the same shape holds inside every difficulty stratum. Non-truncated long rollouts still + pass at ~0.93, so the damage comes from hitting the wall, not from being long. + * The GT reference solutions (r1_solution_1) are LONGER than the model's correct answers + (p50 4377 vs 3444 tokens; GT p90 9816 already exceeds the 8192 budget). "Shorter is + better" is empirically false here, so a monotone-from-zero length penalty would tax the + median GOOD answer — hence the dead zone below ``len_lo`` (人工拍板 A, 2026-07-29). + * Convex ramp (``len_pow`` > 1) concentrates the penalty in the last ~1300 tokens before the + budget, matching the flat-then-cliff damage curve while still making the marginal penalty + grow with length. + * Marker density has a real but small INDEPENDENT effect: raw dose-response is pass + 0.908 -> 0.412 across density 0-2 -> 6-9, but inside a fixed token band it shrinks to + 0.976 -> 0.931 and 0.971 -> 0.851, i.e. ~85% of the raw effect is just "longer outputs + mechanically contain more markers". Hence ``beta_loop`` is deliberately small. + + Composition is multiplicative per rollout (人工拍板): eff = (1-a*len_pen)*(1-b*loop_pen), + score = correct * eff. Per-rollout (not per-candidate) so a short correct rollout is never + punished for a sibling rollout that burned the budget. + """ + if not rolls: + return 0.0, 0.0, {'len_pen': 0.0, 'loop_pen': 0.0, 'eff': 1.0, 'loop_density': 0.0} + span = max(1, budget - len_lo) + d_span = max(1e-9, loop_hi - loop_lo) + s_sum = ineff_sum = lp_sum = mp_sum = eff_sum = dens_sum = 0.0 + for x in rolls: + tok = int(x.get('gen_tokens') or 0) + len_pen = min(1.0, (max(0, tok - len_lo) / span) ** len_pow) + dens = _loop_density(x) + loop_pen = min(1.0, max(0.0, (dens - loop_lo) / d_span)) + eff = (1.0 - alpha_len * len_pen) * (1.0 - beta_loop * loop_pen) + s_sum += eff if x.get('correct') else 0.0 + ineff_sum += 1.0 - eff + lp_sum += len_pen + mp_sum += loop_pen + eff_sum += eff + dens_sum += dens + n = float(len(rolls)) + diag = {'len_pen': lp_sum / n, 'loop_pen': mp_sum / n, + 'eff': eff_sum / n, 'loop_density': dens_sum / n} + return s_sum / n, ineff_sum / n, diag + + +class PassrateHingeMethod(TrainMethod): + """E16 (view B, query-only): the data-driven closure of the reward probe. + + Per chunk: (1) baseline greedy solve (T=0) to measure each problem's no-skill executor + output length, keep only the danger band ``base_tok > --base-tok-floor`` (probe: base_tok + vs skill lift +0.62 — the strongest problem filter; a soft floor keeps >=TRAIN_DP problems + so a chunk never fully empties). (2) query-only skill-gen, N candidates. (3) score each + parseable skill over M=``reward_rollouts`` executor rollouts at T=``reward_temperature`` + with a per-rollout multiplicative efficiency factor (see ``_efficiency_terms``): + + eff_i = (1 - alpha_len * len_pen_i) * (1 - beta_loop * loop_pen_i) + reward = mean_i(correct_i * eff_i) - kappa * mean_i(1 - eff_i) + unparseable -> -1.0 floor + + The ``- kappa * mean(1 - eff)`` tail is what keeps FAILING candidates separable: a pure + product would send every wrong candidate back to 0, which is exactly E4's pathology (format + failure / burned budget / wrong method all collapsing onto one reward value). + + Coefficient sizing: the total deduction ``(1 + kappa) * (1 - eff)`` is kept under ONE + pass_rate quantum (1/M), and ``reward = max(reward, pass_rate - 1/M)`` enforces that as a + hard guard. Rationale: pass_rate is what eval measures, so the efficiency signal may break + ties but must never rank "solved it once" below "never solved it". The absolute coefficient + size barely matters — in an all-fail group every reward is -kappa*(1-eff) and A=(R-mean)/std + rescales that spread back to unit magnitude, so the CURVE SHAPE, not the scale, is the signal. + + leak 不参与 reward（项目既定要求；``--reward-leak-gate`` 默认 0），只走监控口径。 + + Group-relative advantage + BNPO on the query-only trajectory (identical to bnpo/logp). + """ + needs_rubric = False + + def __init__(self, ctx: MethodContext): + super().__init__(ctx) + # bugfix #13: 长度死区 len_lo 不随 --max-tokens 联动；lo >= budget 时惩罚恒 0，静默失效。 + # 按标定比例（5500/8192）自动缩放并告警。 + lo = int(getattr(ctx.args, 'reward_trunc_lo', 5500) or 5500) + if lo >= ctx.args.max_tokens: + new_lo = max(1, int(ctx.args.max_tokens * 5500 / 8192)) + sys.stderr.write(f'[ablate] WARNING: --reward-trunc-lo {lo} >= --max-tokens ' + f'{ctx.args.max_tokens} disables the length penalty; ' + f'rescaled to {new_lo}.\n') + ctx.args.reward_trunc_lo = new_lo + + def _score_candidates(self, kept, prompts): + """skill-gen -> M-rollout with-skill solve -> 效率加权 reward（不做 advantage）。 + + 从 step() 里抽出只为了让 E17(ReflexionMethod) 在不复制 reward 代码的前提下换掉 prompt + （query-only -> query+rubric）。行为与抽出前逐字一致；prompts 与 kept 同序同长。 + """ + ctx = self.ctx + args = ctx.args + M = max(1, int(getattr(args, 'reward_rollouts', 8) or 8)) + # `or 0.5` 会把显式的 0.0 当成“未设”静默提到 0.5（SEAM 口径 m=1/T=0 因此根本 + # 设不进来）。改成 None 判定：未设才用默认。已跑完的 E16 显式传 0.5，行为不变。 + _t = getattr(args, 'reward_temperature', None) + temp = 0.5 if _t is None else float(_t) + alpha_len = abs(float(getattr(args, 'reward_trunc_penalty', 0.12))) + len_lo = int(getattr(args, 'reward_trunc_lo', 5500) or 5500) + len_pow = max(1.0, float(getattr(args, 'reward_len_pow', 2.0) or 2.0)) + beta_loop = abs(float(getattr(args, 'reward_loop_penalty', 0.04))) + loop_lo = float(getattr(args, 'reward_loop_lo', 2.0)) + loop_hi = float(getattr(args, 'reward_loop_hi', 9.0)) + kappa = abs(float(getattr(args, 'reward_ineff_kappa', 0.10))) + # 不可反转护栏：总扣分封顶在一个 pass_rate 量子（1/M）以内，保证"做对过一次"永远排在 + # "一次没做对"之前。量级校验：max 扣分 = (1 + kappa) * (1 - eff_min)。 + pen_cap = 1.0 / M - 1e-6 + leak_gate = abs(float(getattr(args, 'reward_leak_gate', 0.0))) # 默认 0：leak 只做监控 + sg = _run_samples(ctx.skill_sampler, list(prompts), + args.n_skills, args.skill_max_tokens, ctx.skill_dp, + temperature=args.skill_gen_temperature, top_p=args.skill_gen_top_p, + top_k=args.skill_gen_top_k) + flat = [] + for r, seqs in zip(kept, sg): + for s in seqs or []: + resp = _clean_text(getattr(s, 'decoded', '') or '') + block = _extract_skill(resp) or '' + cand = {'skills': block, 'response': resp, 'parseable': bool(block), + 'leaked': None, 'with_pass': None, 'reward': None, 'rolls': [], + 'advantage': 0.0, 'kept': False, + 'skillgen_stop': getattr(s, 'stop_reason', None), + 'skillgen_tokens': len(getattr(s, 'tokens', None) or []), + 'trunc_pen': None, 'pass_rate': None, + 'loop_pen': None, 'eff': None, 'loop_density': None} + r['_cands'].append(cand) + if block: + cand['leaked'] = _answer_leaked(block, r['reference_answer']) + flat.append((r, cand)) + # M-rollout with-skill solve (T>0) -> per-rollout efficiency-weighted reward + if flat: + solve = _run_samples( + ctx.base_sampler, + # V2 fix: pass the raw response like v2 process_chunk (seam nesting; v2-mode no-op) + [build_skill_solve_prompt(r['problem'], c['skills'], c.get('response')) for r, c in flat], + M, args.max_tokens, ctx.base_dp, temperature=temp) + # 判分批量化（code 任务：跑单测的子进程必须并行，见 v2._parse_many） + pairs, spans = [], [] + for (r, _c), seqs in zip(flat, solve): + start = len(pairs) + pairs.extend((s, r['reference_answer']) for s in (seqs or [])) + spans.append((start, len(pairs))) + judged = v2._parse_many(pairs) + for (r, c), (a, b) in zip(flat, spans): + rolls = judged[a:b] or [_empty_roll()] + c['rolls'] = rolls + pr = sum(1.0 for x in rolls if x['correct']) / len(rolls) + score, ineff, diag = _efficiency_terms( + rolls, budget=args.max_tokens, len_lo=len_lo, len_pow=len_pow, + alpha_len=alpha_len, beta_loop=beta_loop, loop_lo=loop_lo, loop_hi=loop_hi) + c['pass_rate'], c['with_pass'] = pr, pr + c['trunc_pen'] = diag['len_pen'] # 名字保留，语义为长度惩罚（swanlab 面板连续） + c['loop_pen'], c['eff'] = diag['loop_pen'], diag['eff'] + c['loop_density'] = diag['loop_density'] + reward = score - kappa * ineff + # 护栏：总扣分不得超过一个 pass 量子，否则会反转 pass_rate 排序 + reward = max(reward, pr - pen_cap) + informative = len(str(r['reference_answer']).strip()) >= 2 + if c['leaked'] and informative and leak_gate > 0: + reward -= leak_gate + c['reward'] = reward + for r in kept: + for c in r['_cands']: + if c['reward'] is None: # unparseable / no rolls -> format floor + c['reward'] = -1.0 + + def _reward_panel(self, kept) -> Dict[str, float]: + """reward / 效率面板（E16 与 E17 共用，面板曲线口径保持一致）。""" + def g(k): + return [c[k] for r in kept for c in r['_cands'] if c.get(k) is not None] + rewards = g('reward') + return {'acc/pass_rate_mean': v2._mean(g('pass_rate')), + 'term/trunc_pen_mean': v2._mean(g('trunc_pen')), # = 长度惩罚 len_pen + 'term/loop_pen_mean': v2._mean(g('loop_pen')), + 'term/eff_mean': v2._mean(g('eff')), + 'term/loop_density_mean': v2._mean(g('loop_density')), + 'reward/mean': v2._mean(rewards), 'reward/std': v2._std(rewards)} + + def _gen_records(self, kept, ci) -> List[Dict[str, Any]]: + """v2._full_records 加上 E16/E17 特有的字段（pass_rate / 各惩罚项 / base_* / rubric）。""" + gen_records = v2._full_records(kept, ci) + kept_by_id = {r.get('data_id', ''): r for r in kept} + for gr in gen_records: + r = kept_by_id.get(gr.get('data_id', '')) + if r is not None: + gr['base_tok'] = r['_base_tok'] + gr['base_correct'] = r['_base_correct'] + if r.get('_base_stop') is not None: + gr['base_stop'] = r['_base_stop'] + if r.get('_rubric') is not None: + gr['rubric'] = r['_rubric'] + for gc, c in zip(gr.get('candidates', []), (r['_cands'] if r else [])): + gc['pass_rate'] = c.get('pass_rate') + gc['trunc_pen'] = c.get('trunc_pen') + gc['loop_pen'] = c.get('loop_pen') + gc['eff'] = c.get('eff') + gc['loop_density'] = c.get('loop_density') + return gen_records + + def step(self, chunk, ci): + ctx = self.ctx + args = ctx.args + # 1) baseline (no-skill) greedy solve -> base_tok danger-band filter + base_rolls = _bare_solve(ctx, chunk) + for r, br in zip(chunk, base_rolls): + r['_cands'] = [] + r['_base_tok'] = int(br.get('gen_tokens') or 0) + r['_base_correct'] = bool(br['correct']) + floor_tok = int(getattr(args, 'base_tok_floor', 5000) or 0) + if floor_tok > 0: + kept = [r for r in chunk if r['_base_tok'] > floor_tok] + min_keep = max(TRAIN_DP, 4) + if len(kept) < min_keep: # soft floor: never waste a whole chunk on a thin draw + kept = sorted(chunk, key=lambda r: r['_base_tok'], reverse=True)[:min_keep] + else: + kept = list(chunk) + # 2+3) query-only skill-gen -> M-rollout with-skill solve -> 效率加权 reward + self._score_candidates(kept, [_skillgen_prompt(r['problem']) for r in kept]) + _assign_advantages(kept, args) + grpo = _grpo_records(kept, with_rubric=False) + has_signal = any(abs(s['advantage']) > 1e-9 for s in grpo) + if has_signal and getattr(args, 'drop_zero_adv', False): + grpo = [s for s in grpo if abs(s['advantage']) > 1e-9] + n_upd, tmetrics = 0, {} + if has_signal: + n_upd, tmetrics = _train_batch(ctx, grpo, traj_fn=query_only_train_trajectory) + summary = v2._chunk_summary(kept, ci) + base_toks = [r['_base_tok'] for r in chunk] + metrics = {'signal/zero_grad_frac': summary['zero_grad_frac'], + 'signal/n_kept': float(len(kept)), + 'signal/kept_frac': float(len(kept)) / max(1, len(chunk)), + 'signal/base_tok_mean': v2._mean(base_toks), + 'signal/base_correct_frac': v2._mean([1.0 if r['_base_correct'] else 0.0 for r in chunk]), + **self._reward_panel(kept), + 'leak/rate': summary['leak_rate'], **tmetrics, + **_leak_split(_cand_leak_pairs(kept))} + gen_records = self._gen_records(kept, ci) + # bugfix #15: 被 base_tok 门槛筛掉的题也落盘一行精简记录，筛选器本身可审计 + kept_ids = {id(r) for r in kept} + for r in chunk: + if id(r) not in kept_ids: + gen_records.append({'record_type': 'problem_dropped', 'chunk': ci, + 'data_id': r.get('data_id', ''), + 'base_tok': r['_base_tok'], 'base_correct': r['_base_correct']}) + return {'n_updates': n_upd, 'summary': summary, 'metrics': metrics, + 'gen_records': gen_records} + + class _SFTFamily(TrainMethod): """Shared SFT accumulation + fire. Subclasses fill ``collect`` to add pool samples.""" needs_rubric = True def _sft_record(self, problem, ref, data_id, skill): + # 实测（2026-07-29，.tmp_analysis/verify_v1v2_think.py）：thinking-on 模板对这种裸 + # content 会自动注入空 think 块（\n\n\n\n，且计入 labels）， + # 编码后是合法的 nothink 布局，与带实质 think 的 GRPO 样本混训 token 布局兼容； + # 副作用是把模型推向短/空 think，属设计权衡而非 bug（review #5 定案）。 return {'problem': problem, 'reference_answer': ref, 'data_id': data_id, 'response': f'\n{skill}\n', 'skills': skill, 'advantage': float(self.ctx.args.sft_weight), 'sft': True} @@ -518,11 +1179,19 @@ def step(self, chunk, ci): ctx = self.ctx leak_pairs, gen_records = self.collect(chunk) ctx.pool.rebalance() # 1:1 by the minority side, surplus DISCARDED (#18b 不积压) - n_upd, tmetrics = 0, {} + n_upd = 0 + batch_metrics: List[Dict[str, float]] = [] for batch in ctx.pool.draw_all_ready(): n, m = _train_batch(ctx, batch, traj_fn=query_only_train_trajectory) # SFT: query-only (#6) n_upd += n - tmetrics.update(m) + if m: + batch_metrics.append(m) + # bugfix #16: 多 batch 时按键均值聚合，不再相互覆盖只留最后一批 + tmetrics = {} + if batch_metrics: + keys = set().union(*batch_metrics) + tmetrics = {k: sum(bm[k] for bm in batch_metrics if k in bm) + / sum(1 for bm in batch_metrics if k in bm) for k in keys} return {'n_updates': n_upd, 'metrics': {**{f'pool/{k}': float(x) for k, x in ctx.pool.sizes().items()}, **tmetrics, **_leak_split(leak_pairs)}, @@ -546,7 +1215,7 @@ def _regen_pick(self, record, diag: str, use_orig_skill: bool, orig_skill: str = skill = _extract_skill(resp) or '' if not skill or len(skill) > args.skill_char_limit: continue - if _answer_leaked(skill, record['reference_answer']): + if _leak_blocks(skill, record['reference_answer']): # bugfix #4: informative gate continue cands.append(skill) if not cands: @@ -630,7 +1299,7 @@ def collect(self, chunk): if roll['correct']: # positive seed: first-pass skill that worked, no leak, within char limit if len(sk) <= args.skill_char_limit \ - and not _answer_leaked(sk, r['reference_answer']): + and not _leak_blocks(sk, r['reference_answer']): ctx.pool.add(self._sft_record(r['problem'], r['reference_answer'], r.get('data_id', ''), sk), POS) gen_records.append({'record_type': 'problem', 'data_id': r.get('data_id', ''), @@ -649,6 +1318,420 @@ def collect(self, chunk): return leak_pairs, gen_records +def _answered(roll: Dict[str, Any]) -> float: + """"这条 rollout 到底交没交出一个可判的答案"。 + + math 看 ```` / ``\\boxed``；code 域这两个标记恒不出现，必须换成"抽出了可解析的代码 + 块"（judge_many 把 no_code 记在 kind 里）。不换的话该特征在 code run 上组内方差恒为 0， + observe 会整组跳过 —— 面板上信号最强的那条曲线（E17 实测 cum_sigma 23.70）会静默消失。 + """ + if v2._TASK == 'code': + return 1.0 if (roll.get('code') and roll.get('kind') != 'no_code') else 0.0 + text = roll.get('text') or '' + return 1.0 if ('' in text or '\\boxed' in text) else 0.0 + + +class SnrProbe: + """在线可学信号探针：把组内 advantage 投影到 skill 文本特征上，逐 chunk 上报累积证据。 + + 为什么需要它（E16 事后分析的直接产物，见 .tmp_analysis/batch_size_math.py）： + E16 的 reward 在【结果层】极其确定——"哪个候选让 executor 收住了"每组 SNR 2.25、97% 的组 + 方向一致；但同一个信号投影到任何【skill 文本特征】上，SNR 塌到 0.046、同向组占比 0.481 + （= 抛硬币）。信息是在"从结果归因到文本"这一步丢的，损失约 50 倍。后果是：整个 50 步 run + 在唯一有内容的方向上只累积到 sqrt(1600) x 0.047 ~ 1.9 sigma，连显著性门槛都没到，而策略却 + 以恒定步长（GRPO 组内标准化让 mean|a| 恒为 0.74，信号退化成噪声时步长不会变小）持续扩散 + 离开初始分布——而"冻结 executor 能在预算内收住"恰恰是初始分布自带的脆弱属性。 + + 所以判断一个新臂值不值得跑满 50 步，看的不是 reward/mean（它被 parse 地板的构成变化掩盖， + E16 实测总均值 +0.026 = parse 构成 +0.160 + 层内 -0.134），而是这里的 cum_sigma： + cum_sigma = sqrt(N_groups) * |mean(g)| / std(g) + g 是每组 advantage 与组内标准化目标量的协方差。跑 10-15 个 chunk 就能看出 rubric 条件下的 + 文本层信号是否比 0.046 高一个量级；不高就该停，省 80% 卡时。 + """ + + # (键, 取值函数, 期望方向)；dir=-1 表示"我们希望 reward 压低该量"，上报时已翻正号， + # 所以 mu>0 一律读作"reward 在往我们想要的方向推"。skill_digits 也是 -1：E16 的主导机制 + # 是答案中继（obey|断言对 = 0.99），所以"skill 里的数字变多"是需要报警的方向、不是 + # 鼓励的方向；先前写 +1 会让面板上的正值被误读成好消息。 + # 注：决策变量 cum_sigma 用 |mean|，与方向约定无关，只有 mu / agree 的可读性依赖它。 + TARGETS = ( + ('skill_chars', lambda c: float(len(c.get('skills') or '')), -1), + ('skill_digits', lambda c: float(sum(ch.isdigit() for ch in (c.get('skills') or ''))), -1), + ('think_tokens', lambda c: float(c.get('skillgen_tokens') or 0), -1), + ('exec_trunc', lambda c: _roll_mean(c, lambda x: 1.0 if x.get('stop_reason') == 'length' + else 0.0), -1), + ('exec_tok', lambda c: _roll_mean(c, lambda x: float(int(x.get('gen_tokens') or 0))), -1), + # ---- 2026-07-30 新增 4 条。先用 .tmp_analysis/snr_feature_scan.py 在 E17 的 168 组/ + # 1344 候选上扫过 13 个备选（同一口径），只留下 cum_sigma 过 5 且不与现有项重复的。 + # 被剔掉的（全部 < 3）：exec_tok_spread 2.87、skill_mean_sentence_len 2.67、 + # skill_hedge_frac 2.23、skill_n_sentences 2.04、skill_action_verb_frac 1.72、 + # skill_proper_noun_frac 1.54、skill_step_marker_frac 0.35。 + # ★ exec_answered 实测 23.70、agree 0.966，比旧冠军 exec_trunc(22.21) 还高：把“reward + # 到底在推什么”说得比截断率更直——不是“别写太长”而是“把答案写出来”。dir=+1。 + # （判据按 _TASK 分派，见 _answered） + ('exec_answered', lambda c: _roll_mean(c, _answered), +1), + # loop_pen 就在 reward 公式里（beta=0.04）却一直没进 SNR 面板，那一项到底有没起作用 + # 之前看不到。实测 10.70 / agree 0.789。复用 _loop_density（与 reward 同一个函数）。 + ('exec_loop_density', lambda c: _roll_mean(c, lambda x: _loop_density(x)), -1), + # 唯一直接量化“reward 在多大程度上奖励泄露”的 SNR。实测 8.80 / agree 0.819 / mu=-0.424。 + # 只在组内 leak 有差异的组（实测 83/168）取到值，但那正是要看的那些组。 + # ★ 这是监控形态，不得进 reward（项目既定要求）。dir=-1。 + ('skill_leaked', lambda c: (1.0 if c.get('leaked') else 0.0), -1), + # 比现有的绝对数字数 skill_digits(10.11) 剔掉了 skill 长度混杂；两条都留，差值能 + # 分开“数字变多”与“skill 变长”。实测 6.55 / agree 0.721。 + # 空 skill 返回 0.0（不是 None）：与 skill_chars 口径一致，否则 parse 率随训练上升 + # 会让本特征的取样面系统漂动（正是 observe 里那段注释警告的假趋势源）。 + ('skill_digit_fraction', lambda c: _digit_fraction(c.get('skills') or ''), -1), + ) + + def __init__(self): + # 每个 target 一个累积器：(n, sum(g), sum(g^2), 正号计数) + self._acc: Dict[str, List[float]] = {k: [0.0, 0.0, 0.0, 0.0] for k, _, _ in self.TARGETS} + + def observe(self, records: List[Dict[str, Any]]) -> Dict[str, float]: + """records = 本次更新真正参与训练的 per-problem 记录（带 '_cands'）。返回 swan 标量。""" + out: Dict[str, float] = {} + for key, fn, direction in self.TARGETS: + gs = [] + for r in records: + cs = [c for c in r.get('_cands', []) if c.get('reward') is not None + and c.get('advantage') is not None] + # 只丢"该特征取不到值"的候选，不丢整组。不可解析候选没有 rollout，exec_* 取值 + # 为 None；而 parse 率会随训练上升（E16 实测：全可解析组占比 c0-9 0.757 -> + # c40-49 0.979），所以"任一 None 丢整组"会让取样面随时间系统性扩大 24pp + # —— 那本身就是一个假趋势源，而 cum_sigma 正是要用来判断趋势的。 + # （g = sum(a_i z_i)/n 对 a 加常数不变，z 在存活子集上重新中心化后仍无偏。） + pairs = [(c, fn(c)) for c in cs] + pairs = [(c, v) for c, v in pairs if v is not None] + if len(pairs) < 2: + continue + # advantage 全 0 的组（组内 reward 无差异，_assign_advantages 已置 0）不携带方向 + # 信息；计入只会把均值往 0 拉、同时虚增 n，使 cum_sigma 系统偏低。 + if all(abs(float(c['advantage'])) < 1e-12 for c, _ in pairs): + continue + vals = [v for _, v in pairs] + mu = sum(vals) / len(vals) + var = sum((v - mu) ** 2 for v in vals) / len(vals) + if var < 1e-12: # 组内该特征无差异 -> 这一组对该方向不提供信息 + continue + sd = var ** 0.5 + g = sum(float(c['advantage']) * (v - mu) / sd for c, v in pairs) / len(pairs) + gs.append(direction * g) + if not gs: + continue + a = self._acc[key] + a[0] += len(gs) + a[1] += sum(gs) + a[2] += sum(x * x for x in gs) + a[3] += sum(1.0 for x in gs if x > 0) # 只记正号数，与漂动的运行均值解耦 + n_cum = a[0] + mean_cum = a[1] / n_cum + # 累积 std（总体口径；n 已达数百，与样本口径无实质差别） + var_cum = max(a[2] / n_cum - mean_cum ** 2, 0.0) + sd_cum = var_cum ** 0.5 + out[f'snr/{key}_mu'] = sum(gs) / len(gs) # 本 chunk 的每组均值 + out[f'snr/{key}_mu_cum'] = mean_cum + out[f'snr/{key}_snr_cum'] = (abs(mean_cum) / sd_cum) if sd_cum > 1e-12 else 0.0 + # 同向组占比：取两个符号桶的多数侧，不依赖当时的运行均值（旧实现拿 + # mean_cum 做参系，早期均值不稳时会把同一批组判到不同侧）。 + out[f'snr/{key}_agree_cum'] = max(a[3], n_cum - a[3]) / n_cum + # ★ 决策变量：整个 run 至今在该方向上累积的证据（sigma）。E16 全程只到 1.9。 + out[f'snr/{key}_cum_sigma'] = ((n_cum ** 0.5) * abs(mean_cum) / sd_cum + if sd_cum > 1e-12 else 0.0) + out[f'snr/{key}_n_groups'] = n_cum + return out + + +def _roll_mean(cand: Dict[str, Any], fn: Callable[[Dict[str, Any]], float]) -> Optional[float]: + rolls = cand.get('rolls') or [] + if not rolls: + return None + return sum(fn(x) for x in rolls) / len(rolls) + + +class ReflexionMethod(PassrateHingeMethod): + """E17 —— Reflexion 条件化臂：只在【裸 executor 做错】的题上，用 rubric 生成 skill 并训练。 + + 与父类 PassrateHingeMethod（E16）完全共享 reward 形状（M-rollout pass_rate x 效率加权、 + 死区 5500 起的二次凸长度惩罚、循环惩罚、不可反转护栏、leak 只监控），三处结构差异： + + 1) 选题：父类按 base_tok 危险带筛（floor=5000），本类按【裸解错误】筛，并把批量对齐到 + 恰好 --reflexion-k 道题。对齐的理由：每次更新的组数必须恒定，否则步长与噪声逐 chunk + 变化，SnrProbe 的累积证据和任何趋势读数都会被批量抖动污染。 + 实现上不动态多抽（MethodContext 拿不到 ProblemPool，且 trainer 的断点恢复靠"重放 N 次 + 等长 draw"，变长抽取会破坏恢复），而是把 chunk_size 放大到 ~5K、裸解后逐批取错题 + 并对 rubric 缺失做回填，直到凑满 K；真凑不够时用现有的全部并上报 k_short。 + 标定（.tmp_analysis/e17_param_calib.py，E16 落盘）：全量题池裸错率 0.329（527/1600）， + E16 每次更新实际 23.46 组 / 全程 1173 组，所以 K=24 才能追平 E16 的证据量（累积 + 证据 ∝sqrt(N)，E16 全程在文本层只累到 1.9 sigma，再砍组数就没判别力了）； + chunk=128 时 P(错题<24) = 0.012%。⭐ 不要用 0.446，那是 base_tok>5000 筛选后 + 子集的错误率（偏难），而本臂 floor=0；用它会把 chunk 低估到 96（P=3.6%）。 + 2) skill-gen 走 view A：prompt = rubric_skillgen_prompt(problem, diag)，训练轨迹相应换成 + rubric_train_trajectory（否则会在 query-only 轨迹上训一个 query+rubric 分布下采出的 + response，与 E6 的已知缺陷同型）。rubric API 缺失（失败/坏缓存）的题一律丢弃而不降级 + 成 query-only —— 本臂的唯一自变量就是 rubric，降级样本会把它稀释掉。 + 3) base_tok_floor 强制视为 0。⚠️ 但要知道这几乎不起作用：实测全量题池的 527 道错题里 + 96.96% 是没写完（base_tok>=8192），floor=5000 筛选后也只是 97.71% —— 截断是题目 + （level>=6 配 8192 预算）造成的，不是筛选造成的。所以 rubric 在 ~97% 的题上只能说 + "你超预算了"，signal/wrong_trunc_frac 会直接开在 0.97。 + + API/GPU 重叠：rubric 诊断是纯 API，父类的 GPU 路径在它之后才开始，所以这里先起线程池发 + 诊断、同时不做别的 GPU 工作（本臂没有 B 线可以并行），诊断落地后再进 skill-gen。 + """ + + needs_rubric = True + + def __init__(self, ctx: MethodContext): + super().__init__(ctx) + # 题集由【裸解错】定义，base_tok 危险带筛选强制关闭（用户 2026-07-29 拍板）：错题集 + # 必须保留"推理错 + 没写完"的混合，否则就只剩长尾截断题、测不到方法修正。 + # 在这里而不是在 shell 里置 0，是为了让 config 指纹（trainer 在 build_method 之后才 + # 落盘）记录真实生效值，而不是一个未被读取的默认 5000。 + if int(getattr(ctx.args, 'base_tok_floor', 0) or 0): + sys.stderr.write('[ablate] reflexion: --base-tok-floor forced to 0 (problem set is ' + 'defined by bare-solve failure, not by output length).\n') + ctx.args.base_tok_floor = 0 + self.snr = SnrProbe() + + def step(self, chunk, ci): + ctx = self.ctx + args = ctx.args + K = max(1, int(getattr(args, 'reflexion_k', 16) or 16)) + # 1) 裸解全 chunk，挑错题并对齐到恰好 K 道 + base_rolls = _bare_solve(ctx, chunk) + for r, br in zip(chunk, base_rolls): + r['_cands'] = [] + r['_base_tok'] = int(br.get('gen_tokens') or 0) + r['_base_correct'] = bool(br['correct']) + r['_base_stop'] = br.get('stop_reason') + wrong = [(r, br) for r, br in zip(chunk, base_rolls) if not br['correct']] + # 2) rubric 诊断（纯 API，线程并行；缓存键 = data_id + 裸解轨迹，跳臂全局缓存）。 + # 逐批回填到恰好 K 道：直接 wrong[:K] 会让 rubric 缺失把本 chunk 的组数打到 K 以下， + # 而组数恒定是本臂的硬要求（否则步长、噪声底与 SnrProbe 的累积证据全跟着抖）。 + # 回填只花网络时间不占 GPU；E6/E7 实测缺失率 0，正常路径上循环只进一轮。 + picked, dropped_no_rubric, cursor = [], 0, 0 + while len(picked) < K and cursor < len(wrong): + take = wrong[cursor:cursor + (K - len(picked))] + cursor += len(take) + diags = _diagnose_parallel(ctx, [(_rubric_entry(r, br), None) for r, br in take]) + for (r, br), diag in zip(take, diags): + r['_rubric'] = diag or '' + if diag: + picked.append((r, br, diag)) + else: + dropped_no_rubric += 1 # 不降级成 query-only：rubric 是唯一自变量 + k_short = max(0, K - len(picked)) # 错题不够 + rubric 全打水两种原因合计 + # 组数恒定是本臂的硬要求（步长、噪声底与 SnrProbe 的累积证据都跟着它抖）。原来靠 + # signal/k_short + signal/n_rubric_missing 两条面板指标暴露，面板精简后改走 stderr， + # 否则这个不变量会变成静默失败。 + if k_short or dropped_no_rubric: + sys.stderr.write(f'[E17] c{ci}: WARNING 组数 {len(picked)}/{K}' + f'（缺 {k_short}；其中 rubric 拉不到 {dropped_no_rubric} 道）\n') + kept = [r for r, _br, _d in picked] + # 3) 复用父类的 skill-gen -> M-rollout -> 效率加权 reward -> advantage 流水线 + if kept: + self._score_candidates(kept, [rubric_skillgen_prompt(r['problem'], d) + for r, _br, d in picked]) + _assign_advantages(kept, args) + grpo = _grpo_records(kept, with_rubric=True) + has_signal = any(abs(s['advantage']) > 1e-9 for s in grpo) + if has_signal and getattr(args, 'drop_zero_adv', False): + grpo = [s for s in grpo if abs(s['advantage']) > 1e-9] + n_upd, tmetrics = 0, {} + if has_signal: + n_upd, tmetrics = _train_batch(ctx, grpo, traj_fn=rubric_train_trajectory) + # 4) 指标（2026-07-30 精简）：只保留用户指定的这几条，命名不缩写。 + # 去掉的 leak/* term/* signal/* 仍全量落在 gen_records 里，离线随时可算。 + summary = v2._chunk_summary(kept, ci) if kept else {'zero_grad_frac': 1.0, 'leak_rate': 0.0} + # ★ 两个口径必须分开（否则退化会被隐掉）： + # all = 所有候选，包括解析失败的（skill 为空、reward 拍 -1.0 地板、从未跑 executor） + # scored= 只有真跑了 executor 的（pass_rate 不为 None） + # 只用 scored 算 skill 文本指标会把“退化成空 skill”这个最重要的信号完全遮住 + # （v6 那轮实测空 skill 率 7.4%，全部是 skill 自己写到 8192 撞顶）； + # reward 也必须含地板，否则面板上的 reward 不等于优化器真正看到的那个。 + all_cands = [c for r in kept for c in r['_cands']] + scored = [c for c in all_cands if c.get('pass_rate') is not None] + # roll 里只有 gen_tokens（见 v2._parse_seq），没有 tokens 字段。 + exec_tokens = [float(int(x.get('gen_tokens') or 0)) + for c in scored for x in (c.get('rolls') or [])] + skills = [c.get('skills') or '' for c in all_cands] + # 训练题一律是裸解做错的题，所以 baseline 恒为 0；显式上报是为了让 lift 曲线自解释。 + baseline_accuracy = 0.0 + with_skill_accuracy = v2._mean([c['pass_rate'] for c in scored]) + rewards = [c['reward'] for c in all_cands if c.get('reward') is not None] + metrics = { + 'train/baseline_accuracy': baseline_accuracy, + 'train/with_skill_accuracy': with_skill_accuracy, + 'train/lift': with_skill_accuracy - baseline_accuracy, + 'train/reward_mean': v2._mean(rewards), + 'train/reward_std': v2._std(rewards), + 'train/zero_gradient_fraction': summary['zero_grad_frac'], + # 解析失败率：with_skill_accuracy 只在 scored 上算，这条负责把分母的变化讲出来。 + 'train/skill_parse_failure_rate': (1.0 - len(scored) / len(all_cands)) if all_cands else 0.0, + 'train/skill_length_characters': v2._mean([float(len(s)) for s in skills]), + **_skill_text_metrics(skills), + 'train/executor_length_tokens': v2._mean(exec_tokens), + 'train/executor_loop_rate': v2._mean([c['loop_pen'] for c in scored + if c.get('loop_pen') is not None]), + **tmetrics, # train/loss, train/grad_norm, train/lr, train/iters, train/n_samples + **self.snr.observe(kept), # snr/*：eval 在 R=1 下 MDE 很大，方向判断只能靠这个 + } + gen_records = self._gen_records(kept, ci) + kept_ids = {id(r) for r in kept} + for r in chunk: + if id(r) not in kept_ids: + # drop_reason 让对齐可审计：否则 dump 里分不出"裸解对"、"超过 K 没用上"、 + # "rubric 拉不到"三种丢弃，而只有第三种是需要报警的。 + reason = ('base_correct' if r['_base_correct'] + else 'no_rubric' if r.get('_rubric') == '' else 'beyond_k') + gen_records.append({'record_type': 'problem_dropped', 'chunk': ci, + 'data_id': r.get('data_id', ''), 'drop_reason': reason, + 'base_tok': r['_base_tok'], 'base_correct': r['_base_correct'], + 'base_stop': r['_base_stop']}) + return {'n_updates': n_upd, 'summary': summary, 'metrics': metrics, + 'gen_records': gen_records} + + +class RejectionSftMethod(TrainMethod): + """E18 —— 拒绝采样 SFT（2026-07-30 用户拍板的 9 步方案）： + + 每 chunk：① 全部 query 裸解一次（greedy T=0）判对错 → ② 错题过 rubric 诊断（缺失即丢， + 不降级）→ ③ 按 E17 的 rollout 方式：rubric 条件化 skill-gen（think 模式、T=1.0 × n_skills）， + 每个 skill 让 executor 推理一次（greedy T=0）→ ④ 三道筛选一条： + a. 只留做对的； + b. leak 过滤：含最终答案的丢掉（用 _leak_blocks 的 >=2 字符门：裸 _answer_leaked 对 + 单字符 gold 误报 ~84%，会把池饿死——SFT 家族 bugfix #4 的同一个门；超 + skill_char_limit 的一并丢）； + c. 长度预筛：取离 len_budget 最近的前一半 → 其中与原始 rubric 词频余弦相似度最高的。 + ⑤ 胜者写进本地数据集文件 e18_sft_dataset.jsonl（append-only，含 rubric/相似度/pass 全审计字段） + 并入池 → ⑥ 池满 --e18-accumulate（16，2026-07-30 从 128 改小）条就 SFT 一次（advantage=--sft-weight=1，轨迹用 + query-only + 裸 响应 = nothink 布局：thinking-on 模板会自动注入空 think 块， + 与 Qwen3 enable_thinking=False 的生成布局逐 token 一致，review #5 定案）。 + ⑦ _train_batch 内部 ckpt.sync_weights 把新权重推到 vLLM → ⑧ eval 在 trainer 侧：同一个 + skill_sampler vLLM 临时切 nothink 模板跑 query-only greedy eval（只换客户端编码，引擎不动）。 + + 与 SftMethod(E12) 的本质区别：E12 靠 rubric 重生成 2-in-8 验证入池（无拒绝排序）； + E18 在做对的候选里再按 leak/长度/rubric 对齐度三道筛取唯一胜者，且留下可复现的 + 本地数据集文件。train-with-rubric/train-query-only 的选择沿用 SFT 家族定案 #6： + 用 query-only 轨迹训，避免采集分布（query+rubric）与部署分布（query-only）错配。 + """ + needs_rubric = True + + def step(self, chunk, ci): + ctx = self.ctx + args = ctx.args + # ① 裸解全 chunk（greedy T=0 单次），判对错 + base_rolls = _bare_solve(ctx, chunk) + for r, br in zip(chunk, base_rolls): + r['_cands'] = [] + r['_rubric'] = '' + r['_base_correct'] = bool(br['correct']) + wrong = [(r, br) for r, br in zip(chunk, base_rolls) if not br['correct']] + # ② rubric 诊断（纯 API 线程并行；缺失即丢，不降级 query-only：没有诊断就没有 + # 相似度筛的参照系，与 E17「rubric 是唯一自变量」的丢弃规则同型） + diags = _diagnose_parallel(ctx, [(_rubric_entry(r, br), None) for r, br in wrong]) + todo = [] + for (r, br), diag in zip(wrong, diags): + r['_rubric'] = diag or '' + if diag: + todo.append((r, diag)) + n_rubric_missing = len(wrong) - len(todo) + # ③ rubric 条件化 skill-gen（sampler 模板是 think-on，即用户要求的 think 模式采集）； + # _skillgen_solve 内部对每个可解析 skill 跑 executor greedy T=0 单次，正是本臂口径。 + items = [{'record': r, 'prompt': rubric_skillgen_prompt(r['problem'], d)} + for r, d in todo] + if items: + _skillgen_solve(ctx, items, args.n_skills, temperature=args.skill_gen_temperature) + # ④ 逐题三道筛：做对 -> 不 leak/不超长 -> 长度预筛前半 + rubric 相似度最高 + accepted = [] + sims_pool = [] + n_pass_cands = n_leak_dropped = 0 + for r, d in todo: + passers = [c for c in r['_cands'] + if c.get('parseable') and (c.get('with_pass') or 0) > 0] + n_pass_cands += len(passers) + survivors = [c for c in passers + if len(c['skills']) <= args.skill_char_limit + and not _leak_blocks(c['skills'], r['reference_answer'])] + n_leak_dropped += len(passers) - len(survivors) + if not survivors: + continue + # 长度选择：取离 len_budget 最近的前一半（至少 1 条），再在其中比相似度。 + # 两阶而非加权求和：两个量纲不同（字符距 vs 余弦），权重没法标定。 + by_len = sorted(survivors, key=lambda c: abs(len(c['skills']) - args.len_budget)) + shortlist = by_len[:max(1, (len(by_len) + 1) // 2)] + for c in shortlist: + c['rubric_similarity'] = _rubric_similarity(c['skills'], d) + best = max(shortlist, key=lambda c: c['rubric_similarity']) + best['kept'] = True + sims_pool.append(best['rubric_similarity']) + accepted.append({'problem': r['problem'], + 'reference_answer': r['reference_answer'], + 'data_id': r.get('data_id', ''), + 'response': f"\n{best['skills']}\n", + 'skills': best['skills'], + 'advantage': float(args.sft_weight), 'sft': True, + # 审计字段（只进数据集文件，不进训练轨迹） + 'rubric': d, 'chunk': ci, + 'rubric_similarity': best['rubric_similarity'], + 'skill_chars': len(best['skills']), + 'n_candidates_passed': len(passers)}) + # ⑤ 胜者落盘本地数据集（append-only，逐 chunk 开关避免长持句柄）+ 入池 + if accepted: + with open(os.path.join(args.output_dir, 'e18_sft_dataset.jsonl'), 'a', + encoding='utf-8') as f: + for s in accepted: + f.write(json.dumps(s, ensure_ascii=False) + '\n') + for s in accepted: + # 训练样本只留 _train_trajectory 需要的键（rubric 不进 query-only 轨迹） + ctx.pool.add({k: s[k] for k in ('problem', 'reference_answer', 'data_id', + 'response', 'skills', 'advantage', 'sft')}, NEG) + # ⑥+⑦ 池满 --e18-accumulate 条即 SFT；_train_batch 内部已含 ckpt.sync_weights + n_upd = 0 + batch_metrics: List[Dict[str, float]] = [] + for batch in ctx.pool.draw_all_ready(): + n, m = _train_batch(ctx, batch, traj_fn=query_only_train_trajectory) + n_upd += n + if m: + batch_metrics.append(m) + tmetrics = {} + if batch_metrics: + keys = set().union(*batch_metrics) + tmetrics = {k: sum(bm[k] for bm in batch_metrics if k in bm) + / sum(1 for bm in batch_metrics if k in bm) for k in keys} + # 指标：命名不缩写（沿用 E17 面板约定） + n_wrong = len(wrong) + metrics = { + 'train/pool_size': float(ctx.pool.sizes().get('pool', 0)), + 'train/accept_rate': (len(accepted) / len(todo)) if todo else 0.0, + 'train/candidate_pass_rate': (n_pass_cands / (len(todo) * args.n_skills)) + if todo else 0.0, + 'train/leak_or_overlength_dropped_fraction': (n_leak_dropped / n_pass_cands) + if n_pass_cands else 0.0, + 'train/selected_rubric_similarity': v2._mean(sims_pool), + 'train/selected_skill_length_characters': v2._mean( + [float(s['skill_chars']) for s in accepted]), + 'signal/n_wrong': float(n_wrong), + 'signal/n_rubric_missing': float(n_rubric_missing), + **tmetrics, + } + # gen_records：v2._full_records 不落 rubric/相似度，补上审计字段（筛选器本身可审计） + gen_records = v2._full_records(chunk, ci) + by_id = {r.get('data_id', ''): r for r in chunk} + for gr in gen_records: + r = by_id.get(gr.get('data_id', '')) + if r is None: + continue + gr['base_correct'] = r.get('_base_correct') + if r.get('_rubric'): + gr['rubric'] = r['_rubric'] + for gc, c in zip(gr.get('candidates', []), r.get('_cands', [])): + if c.get('rubric_similarity') is not None: + gc['rubric_similarity'] = c['rubric_similarity'] + return {'n_updates': n_upd, 'metrics': metrics, + 'gen_records': gen_records} + + METHOD_REGISTRY: Dict[str, Callable[[MethodContext], TrainMethod]] = { 'bnpo': BnpoMethod, 'rl_ab': RlAbMethod, @@ -656,6 +1739,11 @@ def collect(self, chunk): 'opsd': OpsdMethod, 'sft': SftMethod, 'improve_sft': ImproveSftMethod, + 'logp_rl': LogpRlMethod, + 'logp_gt': LogpGtMethod, + 'passrate_hinge': PassrateHingeMethod, + 'reflexion': ReflexionMethod, + 'rejection_sft': RejectionSftMethod, } diff --git a/cookbook/exp/skill2lora/skill_ablate/rollouting.py b/cookbook/exp/skill2lora/skill_ablate/rollouting.py index bd2526c98..9b5f0ab25 100644 --- a/cookbook/exp/skill2lora/skill_ablate/rollouting.py +++ b/cookbook/exp/skill2lora/skill_ablate/rollouting.py @@ -50,25 +50,43 @@ # skill，走 v2 的 _regen_prompt（含 orig_skill 字段，即 777-797 模板的逐字英文版），不用这里的提示词。 # narrative 版末尾拼入与 REGEN_SYSTEM 同一个 few-shot 例子（设计 793-796），锁定文体/长度 # 分布与主链路一致；程序化提取而非复制，保证与冻结的 v2 逐字相同。 -_REGEN_EXAMPLE = v2.REGEN_SYSTEM.split('Example:\n', 1)[1] -assert _REGEN_EXAMPLE.startswith('') and _REGEN_EXAMPLE.rstrip().endswith(''), \ - 'REGEN_SYSTEM example extraction broke; check v2.REGEN_SYSTEM formatting' - +# 2026-07-30：narrative 版补上收束纪律句的**指令**（pitfall 版一直有，v2 的 SKILL_GEN_SYSTEM / +# REGEN_SYSTEM 也有，只有这里漏了）。few-shot 例子里本来就带这句，但光靠示范不管用——E17 首个 +# run 实测出现率 0.000-0.006，而同期 E16（有指令）是 1.000。上一轮把这句评为干预优先级 #1， +# 详见 skill_quality_analysis.md「E17 reflexion 臂截断漂移归因」第四节。 +# +# 2026-07-30 第二次改（用户拍板）：从“连贯叙述”改为“问题-根因-规避”结构。 +# ★ 命名提醒：`--skill-style narrative` 这个名字从此名不副实（已不再要求叙事体）， +# 但没改：它进了 exp_dir / swanlab_exp / config 指纹，改名会让旧 run 数据对不上。 +# ★ 只改 _RUBRIC_SKILLGEN_NARRATIVE（E17 专用，训练与 eval 同一个）；v2.REGEN_SYSTEM / +# SKILL_GEN_SYSTEM 未动，所以 E1-E16 的 query-only 链路不受影响。 _RUBRIC_SKILLGEN_NARRATIVE = """\ You are a skill-generation model. Your block will be fed to a SEPARATE downstream executor model that must solve the problem on its own. The executor will NOT see your private reasoning or the analysis below — it only sees what is inside .... -An expert rubric analysis of a failed attempt on THIS problem is provided to you. Use it to understand where solving this type of problem tends to break down, then think privately and abstract WHAT MAKES THIS TYPE OF PROBLEM SOLVABLE into transferable methodology. +An expert rubric analysis of a failed attempt on THIS problem is provided to you. Work from it to derive a COMPLETE, CONCRETE and ACTIONABLE set of instructions for how to avoid going wrong on this problem, and put your full line of thinking inside the block. Then write the block following these rules: -- Give general, transferable solving techniques for this TYPE of problem as one coherent analysis narrative: first name what the problem is essentially asking, then walk through how to approach it, blending the key concepts, the recommended steps, the pitfalls to avoid (informed by the analysis) and a brief reason for each into a single connected story. +- Structure it as problem -> location -> root cause -> countermeasure. Start by naming the failure mode(s) that ACTUALLY occurred on this problem. For each one: state WHERE it happens (which step, which formula or which theorem is being applied), WHAT goes wrong there, WHY it goes wrong (a misunderstanding of a specific concept / a wrong value substituted in / a missing precondition / ...), and WHAT to guarantee in order to avoid it ("to avoid this, make sure that when , you "). +- Cover exactly the failures that are real — no more. If only ONE thing went wrong (e.g. the mathematics was sound and the attempt merely failed to finish), write about that one thing in depth and stop. NEVER invent extra failure modes, and never pad the list to match a pattern: a made-up warning plants a wrong formula in the executor's head. +- If there are several failure modes, walk them in the order the executor will meet them, with ordinals ("First, ...; Second, when moving on to , ..."), so it reads as a checklist; with a single failure mode, no ordinals are needed. +- Be concrete and executable. Every countermeasure must name the object it applies to (the formula, the quantity, the case being split). Do NOT write advice that would read the same on any other problem. - CRITICAL: Do NOT solve the problem, reveal/compute the final answer, or substitute the problem's specific given numbers. Leave ALL concrete numbers for the executor to compute. -- Self-contained: write in the first person (e.g. "I think the step most likely to go wrong is ..."). NEVER reference "the analysis", "the rubric", or "the previous attempt" — the executor cannot see them, such phrasings cause hallucination. -- Keep it concise: aim for roughly one focused paragraph. +- Self-contained: NEVER reference "the analysis", "the rubric", or "the previous attempt" — the executor cannot see them, such phrasings cause hallucination. Address the solver directly. +- Close by telling the solver to commit and emit the answer in one pass without hesitating, naming the failure that hesitating causes. End the block with this exact sentence: "Avoid re-checking loops; box a bare number as soon as it is computed." +- Keep it under about 300 words. + +Put ONLY the guidance inside . -Put ONLY the methodology inside . +Example (several real failure modes): + +The recurring problems on this problem are: miscounting because symmetric configurations are treated as distinct, applying a permutation formula where the objects are actually indistinguishable, and dropping the division that removes duplicates. First, when you set up the count, the failure appears at the moment you choose between a permutation and a combination: the wrong branch is taken because "distinguishable" is read off the surface wording instead of from whether swapping two objects yields a genuinely different configuration. To avoid this, before writing any formula, state explicitly for each set of objects whether swapping two of its members changes the configuration, and only then pick the formula. Second, when moving on to the total count, compute it as if every object were ordered and distinguishable, because that quantity is unambiguous; the error to guard against here is folding the symmetry correction into this step, which makes the correction impossible to audit later. Third, when you apply the symmetry correction, the failure is using the wrong duplication factor — it comes from counting how many objects look alike rather than how many orderings map to the same configuration. To avoid this, derive the factor by asking how many distinct orderings of the interchangeable choices give the identical configuration, and divide the total by exactly that. Finally, commit to the result and emit the answer in one pass without hesitating; hesitating here restarts the case split and burns the budget before any answer is produced. Avoid re-checking loops; box a bare number as soon as it is computed. + -Example: -""" + _REGEN_EXAMPLE +Example (a single real failure mode — the mathematics was sound, so nothing is invented): + +The one recurring problem on this problem is not mathematical: the derivation stays on the right track, but a correct intermediate result gets questioned instead of used, the same quantity is re-derived to double-check it, and the attempt is cut off before any answer is written. The failure appears after the setup is complete, at the moment the first candidate value is in hand; it happens because re-verifying feels safer than committing, yet every re-check repeats the same computation and produces nothing new. To avoid this, once an intermediate quantity is computed, treat it as settled and build the next step directly on it; choose one method at the start, stay on it, and write the final line as soon as the last quantity is evaluated. Avoid re-checking loops; box a bare number as soon as it is computed. + +""" _RUBRIC_SKILLGEN_PITFALL = """\ You are a skill-generation model. A separate executor model will solve the problem; it only sees your block, NOT the analysis below. @@ -89,13 +107,64 @@ Now write the improved guidance:""" +# --- code 任务（BigCodeBench）的 rubric 条件化 skill-gen ----------------------------------- +# 与数学版同一个骨架（问题→位置→根因→规避 / 只写真实发生的失败 / 自持不指涉 rubric）， +# 把"公式、定理、代入数字、boxed 裸数"换成"该调哪个 API、参数与返回形状、边界与异常、 +# 交付一个 code block"。收尾纪律句也换掉：代码域的对应失败不是"兜圈不给答案"，而是 +# "反复重写实现 / 输出解释与演示代码 / 改动给定签名"。 +# ★ 这一版的信息优势来自 rubric 里带**单测真实报错**（code_task.diag_segment）： +# bcb_eval0_probe 实测 rubric skill 0.513 vs query-only 0.382（+0.135, p=4e-5）。 +_CODE_RUBRIC_SKILLGEN = """\ +You are a skill-generation model. Your block will be fed to a SEPARATE downstream engineer model that must implement the function on its own. The engineer sees the same task description and the same required signature, but NOT your private reasoning or the analysis below — it only sees what is inside .... + +An expert review of a failed attempt at THIS task is provided to you, including the real error its unit tests produced. Work from it to derive a COMPLETE, CONCRETE and ACTIONABLE set of instructions for how to avoid going wrong on this task, and put your full line of thinking inside the block. + +Then write the block following these rules: +- Structure it as problem -> location -> root cause -> countermeasure. Start by naming the failure mode(s) that ACTUALLY occurred. For each one: state WHERE it happens (which step of the implementation, which library call, which returned object), WHAT goes wrong there, WHY it goes wrong (a wrong assumption about what an API returns / a keyword the API does not accept / an unhandled empty or missing-column input / the wrong object handed back to the caller / ...), and WHAT to guarantee in order to avoid it ("to avoid this, make sure that when , you "). +- Cover exactly the failures that are real — no more. If only ONE thing went wrong, write about that one thing in depth and stop. NEVER invent extra failure modes: a made-up warning sends the engineer after an API that is not the problem. +- If there are several failure modes, walk them in the order the engineer will meet them, with ordinals ("First, ...; Second, when building the return value, ..."), so it reads as a checklist. +- Be concrete and executable. Every countermeasure must name the object it applies to: the library function, the argument, the returned type, the edge case, the exception. Do NOT write advice that would read the same on any other task. +- CRITICAL: Do NOT write the solution code and do NOT paste concrete literal values from this task. Name the API and describe the shape of the value it returns instead of writing the call out. +- Self-contained: NEVER reference "the analysis", "the review", "the test error" or "the previous attempt" — the engineer cannot see them, such phrasings cause hallucination. Address the engineer directly. +- Close by telling the engineer to deliver one single fenced code block that reproduces the given imports and signature verbatim, with no explanation, no demonstration call and no tests around it. +- Keep it under about 300 words. + +Put ONLY the guidance inside . + +Example (several real failure modes): + +The recurring problems on this type of task are: handing back the wrong object to the caller, assuming a grouping call returns a plain container when it returns an indexed one, and crashing instead of returning a well-defined result when the input is empty. First, when you build the return value, the failure appears at the very last line: a plotting helper is asked for and the figure is returned instead of the axes it drew on, or a tuple is required and only its first element comes back. To avoid this, re-read the sentence in the task that names the output, and make the last line return exactly that many objects in exactly that order, taking the axes object from the plotting call itself rather than from the figure. Second, when you aggregate, the failure is treating the result of the grouping call as a list: it is an indexed object whose labels are the group keys, so positional access silently reads the wrong group. To avoid this, convert it explicitly with the accessor the library provides before you index into it, and sort by the key the task names rather than relying on insertion order. Third, when the input has no rows or the named column is absent, the failure is an exception escaping from the aggregation. To avoid this, decide up front which of the two the task demands — a defined empty result or a specific raised exception — and write that branch before the main computation. Finally, deliver one single fenced code block that reproduces the given imports and signature verbatim, with no explanation, no demonstration call and no tests around it. + + +Example (a single real failure mode — nothing is invented): + +The one recurring problem on this type of task is a keyword that the library function does not accept: the call is the right one for the job, but it is invoked with an argument name borrowed from a similar function in another module, so it raises before any of the logic runs. The failure appears at the single line that does the real work, and it happens because the argument list is recalled from memory instead of from the function being called. To avoid this, when you reach that call, pass only the arguments you are certain that exact function declares, prefer positional arguments for the ones the task names explicitly, and if a behaviour you need is not available as a keyword there, achieve it with a following operation instead of inventing a parameter. Everything else in this task is straightforward once the call succeeds. Deliver one single fenced code block that reproduces the given imports and signature verbatim, with no explanation, no demonstration call and no tests around it. + +""" + +_CODE_RUBRIC_SKILLGEN_USER = """\ +Task: +{problem} + +Expert review of a failed attempt, with the real unit-test error (for your eyes only; do NOT \ +reference it in the skill): +{rubric} + +Now write the guidance:""" + def rubric_skillgen_prompt(problem: str, rubric: str) -> Dict[str, Any]: """View-A skill-gen conditioned on (problem + rubric diagnosis), NO prior skill. Style-matched to the main line via v2's ``_SKILL_STYLE`` global (set by main() from ``--skill-style``). narrative -> narrative rubric prompt; pitfall -> pitfall rubric prompt. + code 任务只有 narrative 一版（E4/E17 都是 narrative；pitfall 未移植，落到同一个 prompt）。 """ + if v2._TASK == 'code': + return {'messages': [ + {'role': 'system', 'content': _CODE_RUBRIC_SKILLGEN}, + {'role': 'user', 'content': _CODE_RUBRIC_SKILLGEN_USER.format( + problem=problem, rubric=rubric)}]} sys_p = _RUBRIC_SKILLGEN_PITFALL if v2._SKILL_STYLE == 'pitfall' else _RUBRIC_SKILLGEN_NARRATIVE return {'messages': [ {'role': 'system', 'content': sys_p}, diff --git a/cookbook/exp/skill2lora/skill_ablate/rubric_cache.py b/cookbook/exp/skill2lora/skill_ablate/rubric_cache.py index 205f829ec..60048a656 100644 --- a/cookbook/exp/skill2lora/skill_ablate/rubric_cache.py +++ b/cookbook/exp/skill2lora/skill_ablate/rubric_cache.py @@ -7,6 +7,11 @@ trajectory. Because the executor is frozen at T=0, that trajectory is deterministic and identical across experiments, so its rubric diagnosis can be shared across ALL runs via one global file (``rubric_cache_global.jsonl``) — diagnose each problem once, reuse everywhere. + ⚠️ 该共享前提只在"executor 口径完全相同"时成立。轨迹本身**不在键里**，所以任何改变裸解 + 轨迹的开关（task 域、executor thinking 开关、以后若改 executor 模型）都必须体现在**文件名** + 上，否则新口径的 run 会命中旧口径的诊断（键只有 data_id，必然命中）。实测这个文件已经攒了 + 2630 条 think 轨迹的数学诊断 —— E19（executor nothink）若共用会几乎全程读到与自己失败无关 + 的诊断，而 rubric 内容正是该臂唯一的自变量。见 build_rubric_cache 的 tag 拼装。 - LocalRubricCache (key = md5(data_id + skill)): the improve-skill+SFT / OPSD lines diagnose a WITH-SKILL trajectory whose skill evolves with the policy, so the diagnosis is experiment- @@ -18,9 +23,16 @@ import os from typing import Any, Dict, Optional +import train_skill_v2 as v2 from train_skill_v2 import DiskCache, _diagnose_entry +def _version() -> str: + """动态读 v2._RUBRIC_VERSION —— 不能 from-import：v2.set_task('code') 会在运行时把它换成 + code 判据的版本号，而 from-import 会把加载那一刻的值钉死，导致代码域诊断用数学域的键。""" + return v2._RUBRIC_VERSION + + class _BaseRubricCache: """Shared get-or-diagnose logic over a DiskCache; subclasses define the key.""" @@ -46,9 +58,14 @@ def get_or_diagnose(self, entry: Dict[str, Any], checker, return '' key = self._key(entry, skill) hit = self._cache.get(key) - if hit is not None: + # bugfix #1: 旧版把 API 失败（_diagnose_entry 返回 None）也以 '' 永久写进缓存， + # 一次瞬时抖动会跨实验毒化全局缓存且永不重试。现在：只缓存真诊断；历史残留的 + # '' 条目视为 miss，下次调用自动重试并用真诊断覆盖。 + if hit: return hit - diag = _diagnose_entry(checker, entry) or '' + diag = _diagnose_entry(checker, entry) + if diag is None: # transient API failure: do NOT cache, retry on the next call + return '' self._cache.put(key, diag) return diag @@ -63,29 +80,49 @@ def close(self) -> None: class GlobalRubricCache(_BaseRubricCache): - """key = data_id: bare-problem trajectory diagnosis, shareable across experiments.""" + """key = (rubric 版本, data_id): bare-problem trajectory diagnosis, shareable across experiments. + + 版本号必须进键：该文件跨实验共享且 append-only，一旦判据表改了而键不变，旧 + taxonomy 的诊断会被静默当成新判据的结果返回（旧版本号仅定义未使用）。 + """ def _key(self, entry: Dict[str, Any], skill: Optional[str] = None) -> str: - return DiskCache.key_for('rubric_global', str(entry.get('data_id', ''))) + return DiskCache.key_for('rubric_global', _version(), + str(entry.get('data_id', ''))) class LocalRubricCache(_BaseRubricCache): - """key = md5(data_id + skill): with-skill trajectory diagnosis, per-experiment only.""" + """key = md5(rubric 版本 + data_id + skill): with-skill trajectory diagnosis, per-experiment only.""" def _key(self, entry: Dict[str, Any], skill: Optional[str] = None) -> str: - return DiskCache.key_for('rubric_local', str(entry.get('data_id', '')), skill or '') + return DiskCache.key_for('rubric_local', _version(), + str(entry.get('data_id', '')), skill or '') -def build_rubric_cache(scope: str, output_dir: str, - global_dir: Optional[str] = None, enabled: bool = True): +def build_rubric_cache(scope: str, output_dir: str, global_dir: Optional[str] = None, + enabled: bool = True, task: str = 'math', executor_thinking: str = 'on'): """Factory: scope='global' -> shared file under ``global_dir`` (default output_dir/..); - scope='local' -> per-experiment file under ``output_dir/cache``.""" + scope='local' -> per-experiment file under ``output_dir/cache``. + + ``task`` 进文件名：代码域与数学域的诊断内容完全不同源（判据表、judge prompt、segment 里 + 有没有单测报错），版本号已经能隔开键，分文件是第二道保险，也让缓存体积可分别管理。 + + ★ ``executor_thinking`` 也必须进文件名（2026-07-31 bugfix）：global 缓存跨实验共享的**唯一 + 依据**是"executor 冻结在 T=0，所以同一道题的裸解轨迹在所有实验里逐字相同"。E19/E20 把 + executor 的 thinking 关掉后这个前提就不成立了 —— 裸解轨迹完全变了（think 那边大量是 + "撞预算没写出代码"，nothink 这边是"写完但答错"），而诊断正是对着这条轨迹做的。共用一个 + 文件会让 nothink 臂直接读到 think 臂的旧诊断（键只有 data_id，必然命中）， + rubric 内容与本臂的真实失败无关 —— 而 rubric 内容恰恰是这两个臂唯一的自变量。 + """ + tag = '' if task == 'math' else f'_{task}' + if executor_thinking != 'on': + tag += '_execnothink' if scope == 'global': base = global_dir or os.path.dirname(os.path.abspath(output_dir.rstrip('/'))) os.makedirs(base, exist_ok=True) - return GlobalRubricCache(os.path.join(base, 'rubric_cache_global.jsonl'), enabled) + return GlobalRubricCache(os.path.join(base, f'rubric_cache_global{tag}.jsonl'), enabled) if scope == 'local': cache_dir = os.path.join(output_dir, 'cache') os.makedirs(cache_dir, exist_ok=True) - return LocalRubricCache(os.path.join(cache_dir, 'rubric_cache_local.jsonl'), enabled) + return LocalRubricCache(os.path.join(cache_dir, f'rubric_cache_local{tag}.jsonl'), enabled) raise ValueError(f"scope must be 'global' or 'local', got {scope!r}") diff --git a/cookbook/exp/skill2lora/skill_ablate/trainer.py b/cookbook/exp/skill2lora/skill_ablate/trainer.py index 0d313745a..d055cb8fe 100644 --- a/cookbook/exp/skill2lora/skill_ablate/trainer.py +++ b/cookbook/exp/skill2lora/skill_ablate/trainer.py @@ -21,6 +21,8 @@ from .config import ExpSpec from .data import load_deepmath_records +from .data_code import load_code_records +from .eval_reflexion import run_reflexion_eval from .methods import MethodContext, build_method from .pool import SamplePool from .rubric_cache import build_rubric_cache @@ -34,9 +36,9 @@ def _rubric_scope(method: str) -> str: """Bare-problem lines (rl/sft) share a GLOBAL cache; with-skill lines (opsd/improve) use a per-experiment LOCAL cache. bnpo needs no rubric.""" - if method in ('rl_ab', 'rl_err', 'sft'): + if method in ('rl_ab', 'rl_err', 'sft', 'reflexion', 'rejection_sft'): return 'global' - if method in ('opsd', 'improve_sft'): + if method in ('opsd', 'improve_sft', 'logp_rl'): return 'local' return '' @@ -48,6 +50,11 @@ def _build_pool(spec: ExpSpec, args) -> Any: if spec.method == 'sft': return SamplePool(batch_size=args.sft_batch_size, balanced=False, max_pool=args.pool_max) + if spec.method == 'rejection_sft': + # E18：攒够 --e18-accumulate（128）条才 fire 一次 SFT；_train_batch 内部仍按 + # sft_batch_size 切 micro，所以这里只需保证是 TRAIN_DP 倍数（main.py 校验）。 + return SamplePool(batch_size=args.e18_accumulate, balanced=False, + max_pool=args.pool_max) return None # RL / OPSD / bnpo train per-chunk, no accumulation pool @@ -68,6 +75,11 @@ def _load_resume_state(args) -> dict: if os.path.exists(state_path): with open(state_path, encoding='utf-8') as f: st = json.load(f) + # bugfix #17: 旧/手写 train_state.json 缺键时给出可读报错而非裸 KeyError + missing = [k for k in ('updates', 'chunk_idx') if k not in st] + if missing: + raise ValueError(f'resume: {state_path} missing keys {missing}; ' + f'present keys: {sorted(st)}') for k in ('chunk_size', 'min_level', 'n', 'seed'): if k in st and getattr(args, k, None) not in (None, '') and st[k] != getattr(args, k): sys.stderr.write(f'[ablate] WARNING: resume data config mismatch: {k} ' @@ -100,7 +112,9 @@ def run_experiment(args, spec: ExpSpec) -> None: sys.stderr.write(f'[ablate] resuming {spec.name} from {resume["ckpt_dir"]} ' f'(updates={resume["updates"]} chunk={resume["chunk_idx"]}).\n') - # 1) style / align globals must be set BEFORE any prompt is built. + # 1) task / style / align globals must be set BEFORE any prompt is built. + # set_task 必须在最前：它同时决定 prompt 分派、判分方式与 rubric 判据版本（缓存键）。 + v2.set_task(spec.task, getattr(args, 'test_workers', 24), getattr(args, 'test_timeout', 60)) v2._ALIGN_MODE = spec.align v2._SKILL_STYLE = spec.style args.skill_thinking = spec.thinking @@ -114,6 +128,11 @@ def run_experiment(args, spec: ExpSpec) -> None: elif args.skill_max_tokens != spec.skill_max_tokens: sys.stderr.write(f'[ablate] WARNING: --skill-max-tokens {args.skill_max_tokens} ' f'overrides the {spec.name} default {spec.skill_max_tokens}.\n') + # E14 信噪比消融：训练判分 rollout 数/温度，CLI 显式值优先，否则取 spec（E1-E13=1×T0）。 + if getattr(args, 'reward_rollouts', None) is None: + args.reward_rollouts = spec.reward_rollouts + if getattr(args, 'reward_temperature', None) is None: + args.reward_temperature = spec.reward_temperature # OOM guard: think/8192 实验的训练序列长一倍，fp32 主权重下 16/2卡 的 micro backward # 会爆显存（E13 实测 Tried to allocate 37.9GiB）；自动把 micro 减半到 8，梯度归一后数学等价， # 攒批/采样批（sft_batch_size）冻结口径不变。显式 --train-micro-batch 优先。 @@ -122,8 +141,11 @@ def run_experiment(args, spec: ExpSpec) -> None: sys.stderr.write(f'[ablate] train_micro_batch auto-set to {args.train_micro_batch} ' f'(skill_max_tokens={args.skill_max_tokens} OOM guard).\n') - records, eval_records = (load_deepmath_records(args) if getattr(args, 'deepmath_dir', '') - else v2._load_records(args)) + if spec.task == 'code': + records, eval_records = load_code_records(args) + else: + records, eval_records = (load_deepmath_records(args) if getattr(args, 'deepmath_dir', '') + else v2._load_records(args)) if len(records) < args.chunk_size: raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded ({len(records)})') @@ -150,16 +172,18 @@ def run_experiment(args, spec: ExpSpec) -> None: scope = _rubric_scope(spec.method) rubric_cache = build_rubric_cache(scope, args.output_dir, global_dir=args.rubric_global_dir, - enabled=not args.no_cache) if scope else None + enabled=not args.no_cache, + task=spec.task, + executor_thinking=spec.executor_thinking) if scope else None - # client-side Template clone (OPSD only): encodes student/teacher trajectories locally to - # extract the teacher's RESPONSE-ONLY logp positions (prompts differ in length, so the - # remote full-sequence logps must be sliced before OPSDLoss can align them per token). - # Mirrors init_components' set_template call exactly (same tokenizer/thinking/truncation). + # client-side Template clone: + # - OPSD: skill-model template, used to align response-token positions for teacher logps. + # - logp_rl / logp_gt: executor template (thinking on), used to slice logP(S | executor prompt). encode_template = None - if spec.method == 'opsd': + if spec.method in ('opsd', 'logp_rl', 'logp_gt'): encode_template = v2.Template(model_id=v2.MODEL_ID, - enable_thinking=(spec.thinking == 'on'), + enable_thinking=(True if spec.method in ('logp_rl', 'logp_gt') + else spec.thinking == 'on'), max_length=args.max_model_len, truncation_strategy='delete') @@ -202,8 +226,11 @@ def _save_ckpt(name: str, updates: int, chunk_idx: int, epoch: int) -> None: use_swan = swanlab is not None and os.environ.get('SWANLAB_MODE') != 'disabled' if use_swan: - # timestamp suffix so FORCE reruns never collide in swanlab (接口方案 #9) - swan_exp = f'{spec.swanlab_exp}_{time.strftime("%Y%m%d_%H%M%S")}' + # timestamp suffix so FORCE reruns never collide in swanlab (接口方案 #9); + # --run-tag 用于区分共用同一 ExpSpec 的变体（如 kl_beta 0.001 vs 0.01）。 + _tag = (getattr(args, 'run_tag', '') or '').strip() + swan_exp = (f'{spec.swanlab_exp}' + (f'_{_tag}' if _tag else '') + + f'_{time.strftime("%Y%m%d_%H%M%S")}') swanlab.init(project=args.swanlab_project, experiment_name=swan_exp, config={'exp': spec.name, 'view': spec.view, 'method': spec.method, 'thinking': spec.thinking, 'style': spec.style, 'align': spec.align, @@ -211,10 +238,23 @@ def _save_ckpt(name: str, updates: int, chunk_idx: int, epoch: int) -> None: 'max_updates': args.max_updates, 'lr': args.lr, 'n_skills': args.n_skills, 'sft_batch_size': args.sft_batch_size, 'len_budget': args.len_budget, + 'run_tag': _tag, + 'reward_rollouts': args.reward_rollouts, + 'reward_temperature': args.reward_temperature, + 'kl_beta': getattr(args, 'kl_beta', 0.01), + 'grpo_epsilon': getattr(args, 'grpo_epsilon', 0.2), + 'adv_clip': getattr(args, 'adv_clip', 0.0), + 'logp_leak_penalty': getattr(args, 'logp_leak_penalty', 0.0), + 'reward_leak_gate': getattr(args, 'reward_leak_gate', 0.0), + 'reward_trunc_penalty': getattr(args, 'reward_trunc_penalty', 0.25), + 'reward_trunc_lo': getattr(args, 'reward_trunc_lo', 6000), + 'base_tok_floor': getattr(args, 'base_tok_floor', 5000), 'drop_zero_adv': args.drop_zero_adv}) - cfg = {'record_type': 'config', 'exp': spec.name, 'view': spec.view, 'method': spec.method, + cfg = {'record_type': 'config', 'exp': spec.name, 'task': spec.task, 'view': spec.view, + 'method': spec.method, 'thinking': spec.thinking, 'style': spec.style, 'align': spec.align, 'loss': spec.loss, + 'executor_thinking': spec.executor_thinking, 'skill_max_tokens': args.skill_max_tokens, 'needs_rubric': spec.needs_rubric, 'rubric_scope': scope, 'rubric_check': bool(checker), 'n': len(records), 'eval_n': len(eval_records), 'model': v2.MODEL_ID, @@ -223,10 +263,28 @@ def _save_ckpt(name: str, updates: int, chunk_idx: int, epoch: int) -> None: 'sft_batch_size': args.sft_batch_size, 'len_budget': args.len_budget, 'skill_char_limit': args.skill_char_limit, 'drop_zero_adv': args.drop_zero_adv, 'improve_skill_temperature': args.improve_skill_temperature, + 'reward_rollouts': args.reward_rollouts, + 'reward_temperature': args.reward_temperature, + 'run_tag': (getattr(args, 'run_tag', '') or ''), + # loss 侧旋钮同样入账：kl_beta 是唯一对抗漂移的恢复力，此前不落盘导致已跑的臂 + # 无法从 gen_records 反推当时的锚强度（全部是旧默认 0.001；现默认 0.01）。 + 'kl_beta': getattr(args, 'kl_beta', 0.01), + 'grpo_epsilon': getattr(args, 'grpo_epsilon', 0.2), + 'adv_clip': getattr(args, 'adv_clip', 0.0), + # reward 公式的所有旋钮全部入账：之前 leak 惩罚默认开着却不落盘，导致已跑的臂无法 + # 从 gen_records 反推当时用的是哪个 reward。leak 一律不进 reward，此处应恒为 0。 + 'logp_leak_penalty': getattr(args, 'logp_leak_penalty', 0.0), + 'reward_leak_gate': getattr(args, 'reward_leak_gate', 0.0), + 'reward_trunc_penalty': getattr(args, 'reward_trunc_penalty', 0.25), + 'reward_trunc_lo': getattr(args, 'reward_trunc_lo', 6000), + 'base_tok_floor': getattr(args, 'base_tok_floor', 5000), + 'reflexion_k': int(getattr(args, 'reflexion_k', 0) or 0), + 'eval_protocol': ('reflexion' if spec.method == 'reflexion' else 'query_only'), 'eval_rollouts': args.eval_rollouts, 'eval_skill_temperature': args.eval_skill_temperature, 'seam_parquet_dir': (getattr(args, 'seam_parquet_dir', '') or ''), 'deepmath_dir': (getattr(args, 'deepmath_dir', '') or ''), 'min_level': int(getattr(args, 'min_level', 0) or 0), + 'eval_min_level': int(getattr(args, 'eval_min_level', 0) or 0), 'save_every_updates': int(getattr(args, 'save_every_updates', 0) or 0), 'resumed_from': (resume['ckpt_dir'] if resume else ''), 'resumed_updates': (resume['updates'] if resume else 0), @@ -242,9 +300,33 @@ def _save_ckpt(name: str, updates: int, chunk_idx: int, epoch: int) -> None: v2._write(f, cfg) def _do_eval(updates_done: int, swan_step: int) -> None: - recs, summary, metrics = v2.run_greedy_eval( - base_sampler, skill_sampler, eval_records, updates_done, updates_done, - base_dp, skill_dp, args, eval_base_cache) + # E17 用 reflexion 协议 eval（只干预裸解做错的题、skill-gen 带 rubric），与训练 + # 同分布；其余臂一律走 v2 的 query-only 全量 eval。两者 summary 键名兼容。 + if spec.method == 'reflexion': + recs, summary, metrics = run_reflexion_eval( + base_sampler, skill_sampler, eval_records, updates_done, updates_done, + base_dp, skill_dp, args, eval_base_cache, rubric_cache, checker) + elif spec.method == 'rejection_sft': + # E18：eval 用 nothink rollout，且与采集共用同一个 skill_sampler vLLM（用户 + # 2026-07-30 拍板）。set_template 只换客户端编码（sampler/base.py:106）， + # 引擎不重建；训练响应是裸 （空 think 布局），与 nothink 生成布局 + # 逐 token 一致，所以这才是本臂的同分布读数。finally 必须切回 think， + # 否则下一个 chunk 的采集会静默变成 nothink。 + skill_sampler.set_template(v2.Template, model_id=v2.MODEL_ID, + enable_thinking=False, + max_length=args.max_model_len) + try: + recs, summary, metrics = v2.run_greedy_eval( + base_sampler, skill_sampler, eval_records, updates_done, updates_done, + base_dp, skill_dp, args, eval_base_cache) + finally: + skill_sampler.set_template(v2.Template, model_id=v2.MODEL_ID, + enable_thinking=(spec.thinking == 'on'), + max_length=args.max_model_len) + else: + recs, summary, metrics = v2.run_greedy_eval( + base_sampler, skill_sampler, eval_records, updates_done, updates_done, + base_dp, skill_dp, args, eval_base_cache) for rec in recs: v2._write(eval_f, rec) v2._write(eval_f, summary) @@ -321,6 +403,13 @@ def _do_eval(updates_done: int, swan_step: int) -> None: if save_every and updates >= last_save_at + save_every and n_upd > 0: _save_ckpt(f'{spec.name}-u{updates}', updates, chunk_idx + 1, pool_pp.epoch) last_save_at = updates + # bugfix #12: 方法把候选/rollout 全文挂在共享 record dict 上（ProblemPool 跨 epoch + # 复用同一批对象），gen_records 已落盘后不清理会让 driver RAM 随触达题数线性增长 + # （5000 题 × 8 候选 × 几十 KB 可达数 GB）。训练/日志都结束后安全清理。 + for r in chunk: + for k in ('_cands', '_pseudo_rolls', '_pseudo_roll', '_pseudo_solution', + '_rubric', '_base_tok', '_base_correct', '_base_stop'): + r.pop(k, None) chunk_idx += 1 # final readout — skip if the periodic eval already covered this exact update count diff --git a/cookbook/exp/skill2lora/train_skill_v2.py b/cookbook/exp/skill2lora/train_skill_v2.py index 3b9cc8313..3e59fe245 100644 --- a/cookbook/exp/skill2lora/train_skill_v2.py +++ b/cookbook/exp/skill2lora/train_skill_v2.py @@ -43,6 +43,9 @@ from twinkle_agentic.verifier import RubricVerifier from twinkle_agentic.verifier.rubric_verifier import RubricItem +# 任务适配器（BigCodeBench）。code_task 刻意不 import 本模块，所以这里可以顶层 import。 +import code_task + logger = get_logger() try: @@ -67,6 +70,32 @@ REF_DP = REF_GPUS // REF_FSDP +# =========================================================================== +# Section A0 — task switch (math = DeepMath \boxed{}, code = BigCodeBench unittest) +# =========================================================================== +# 与 _ALIGN_MODE / _SKILL_STYLE 同型的模块级开关（由 set_task 设置，trainer 在任何 prompt +# 构造之前调用）。**math 分支逐字不变**，所以 E1-E16 的行为、判分与可复现性不受影响。 +# 分派点一共 7 处：build_direct_prompt / build_skill_solve_prompt / _skillgen_prompt / +# _parse_seq(+_parse_many) / _answer_leaked / build_rubric_checker / _diagnose_entry。 +# 换成 code 时 reference_answer 不再是数值，而是 code_task.payload_of() 的判分载荷 +# （task_id / entry_point / test / code_prompt / doc_struct / canonical_solution）。 +_TASK = 'math' # 'math' | 'code' +_CODE_TEST_WORKERS = 24 # 单测线程池（子进程并行度） +_CODE_TEST_TIMEOUT = 60 # 单题单测墙钟上限（秒） + + +def set_task(task: str, test_workers: int = 24, test_timeout: int = 60) -> None: + """Set the task family. MUST run before any prompt is built or any roll is judged.""" + global _TASK, _CODE_TEST_WORKERS, _CODE_TEST_TIMEOUT, _RUBRIC_VERSION + if task not in ('math', 'code'): + raise ValueError(f"task must be 'math' or 'code', got {task!r}") + _TASK = task + _CODE_TEST_WORKERS = max(1, int(test_workers)) + _CODE_TEST_TIMEOUT = max(5, int(test_timeout)) + # 判据表换了，rubric 缓存键必须跟着换（rubric_cache 动态读 v2._RUBRIC_VERSION）。 + _RUBRIC_VERSION = code_task.RUBRIC_VERSION if task == 'code' else _RUBRIC_VERSION_MATH + + # =========================================================================== # Section A — boxed extraction + answer grading (verbatim from v1) # =========================================================================== @@ -249,9 +278,12 @@ def _numeric_value(raw) -> Optional[str]: return (str(int(float(s))) if float(s) == int(float(s)) else str(float(s))) if _NUM_RE.fullmatch(s) else None -def _answer_leaked(skill: str, reference: str) -> bool: +def _answer_leaked(skill: str, reference) -> bool: if not skill: return False + if _TASK == 'code': + # 代码域：leak = skill 里出现参考解答的实质代码行（与数学域一致，只做监控） + return code_task.leaked(skill, reference) # Suffix guard: reject only a following DIGIT or a following '.' (decimal point), # NOT a sentence-ending '.'. Old '(?![\d.])' let leaks like "...= 675." slip through # because the trailing period satisfied the [\d.] class. 中文注释：尾断言只排除"后接数字" @@ -288,6 +320,10 @@ def _seam_sanitize(txt: str) -> str: txt = m.group(1).strip() elif (m := _SEAM_INLINE_RE.search(txt)): txt = (m.group(1) or m.group(2)).strip() + # bugfix 2026-07-29：\dfrac/\tfrac/\cfrac 不被下行字面 \frac 正则匹配，曾致 + # \boxed{\dfrac{1}{2}} 落到 _SEAM_NUM_RE 抓首个数字 → pred='1'（分数答案题全判错； + # 实测被标"错"的 boxed rolls 中 70-85% 实为正确，见 skill_quality_analysis.md 末章）。 + txt = txt.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac').replace(r'\cfrac', r'\frac') txt = re.sub(r'\\frac\s*\{\s*([^}]+?)\s*}\s*\{\s*([^}]+?)\s*}', r'\1/\2', txt) if (m := _SEAM_FRAC_RE.search(txt)): p, q = map(float, m.groups()) @@ -329,7 +365,9 @@ def _extract_skill(text: str) -> Optional[str]: return block or None -def _parse_seq(seq, gold: str) -> Dict[str, Any]: +def _parse_seq(seq, gold) -> Dict[str, Any]: + if _TASK == 'code': + return _parse_many([(seq, gold)])[0] text = _clean_text(getattr(seq, 'decoded', '') or '') # 判分口径统一（人工拍板，2026-07-27）：seam/v2 都只从 \boxed{} 抽取，再走同一套数值归一 # （frac/inline/number）后精确匹配；不做 lpem 式“整段抓数字”贪婪回退，保证 E13 与 E1-E12 @@ -343,7 +381,34 @@ def _parse_seq(seq, gold: str) -> Dict[str, Any]: 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} +def _parse_many(pairs) -> List[Dict[str, Any]]: + """批量判分入口。``pairs`` = [(seq_or_None, gold)]，返回同序 roll 列表。 + + math 分支逐条 _parse_seq（与逐条调用 bit 一致）；code 分支必须批量 —— 判分要起子进程跑 + unittest（典型 1-3s），一个 chunk 有几百次判分，串行会比同 chunk 的 GPU 时间还长一个量级。 + 所有 rollout 汇合点（process_chunk / run_greedy_eval / methods / eval_reflexion）都走这里。 + """ + if _TASK != 'code': + return [(_parse_seq(s, g) if s is not None else _empty_roll()) for s, g in pairs] + items = [] + for s, g in pairs: + if s is None: + items.append(None) + continue + items.append((_clean_text(getattr(s, 'decoded', '') or ''), + getattr(s, 'stop_reason', None), + len(getattr(s, 'tokens', None) or []), g)) + return code_task.judge_many(items, _CODE_TEST_WORKERS, _CODE_TEST_TIMEOUT) + + +def _first_seq(seqs): + """rollout 列表 -> 首个 sequence 或 None（判分批量化后统一用它取 seq）。""" + return seqs[0] if seqs else None + + def _empty_roll(): + if _TASK == 'code': + return code_task.empty_roll() return {'pred': '', 'correct': False, 'terminated': False, 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} @@ -566,7 +631,7 @@ def __getattr__(self, name): - For FAIL items: describe the process problem at strategy level. - A fix suggests the LOCAL correction direction without solving. - Never reveal the final answer or a corrected expression. -- If segment was cut off (no final reached), mark length-budget as FAIL. +- If segment was cut off (no final reached), mark the output-format criterion as FAIL. - Keep "reason" and "fix" concise: one short sentence each. - Output only the JSON object.""" @@ -582,16 +647,28 @@ def __getattr__(self, name): Now output the diagnostic JSON object.""" +# 判据按「错误类型」组织，但文本一律写成正向陈述 —— 全流程的语义是 PASS=没问题 / +# FAIL=该类错误存在（见 _format_diagnosis 与 gate），写成否定句会把 PASS/FAIL 反过来。 _MATH_RUBRIC = [ - ('The attempt chooses a method suitable for the problem structure', False), - ('The attempt identifies the key constraint, invariant, or quantity before computing', False), - ('Algebraic and logical transformations preserve validity at each step', True), - ('The attempt checks required constraints, domains, boundary cases, or validity conditions', False), - ('The attempt avoids redundant casework, looping, or re-deriving known facts', False), - ('The attempt reaches a final within the length budget', False), - ('The approach stays focused on the actual question asked', False), + # 1. 代数计算错误 + ('Arithmetic and algebraic manipulations are carried out correctly', True), + # 2. 公式定理使用错误 + ('Formulas and theorems are invoked correctly and their preconditions hold', False), + # 3. 起始方法论错误 + ('The initial approach is viable for this problem rather than a dead end', False), + # 4. 题目目标分析错误 + ('The attempt correctly identifies what the problem actually asks for', False), + # 5. 输出格式错误 + ('The attempt reaches a final answer in the required output format', False), + # 6. 对计算过程反复犹豫 + ('The attempt commits to its computation instead of repeatedly second-guessing it', False), + # 7. 构成自相矛盾 + ('The attempt stays internally consistent and never contradicts its own results', False), ] -_RUBRIC_VERSION = 'rubric_v5_viewb_strategy' +# 版本号进 rubric 缓存键（GlobalRubricCache._key）：判据一改，旧诊断必须失效， +# 否则 rubric_cache_global.jsonl 里按 data_id 存的旧taxonomy诊断会被当成新判据的结果返回。 +_RUBRIC_VERSION_MATH = 'rubric_v6_error_taxonomy' +_RUBRIC_VERSION = _RUBRIC_VERSION_MATH # set_task('code') 会换成 code_task.RUBRIC_VERSION class _RftRubricVerifier(RubricVerifier): @@ -602,10 +679,23 @@ def _diagnose_trajectory(self, query, rubric_block, segment_text): query=query, rubric=rubric_block, segment=segment_text)}]} +class _CodeRubricVerifier(RubricVerifier): + """代码域 judge：判据 = code_task.CODE_RUBRIC，且 segment 里带**单测真实报错**。""" + + def _diagnose_trajectory(self, query, rubric_block, segment_text): + return {'messages': [ + {'role': 'system', 'content': code_task.DIAG_SYSTEM}, + {'role': 'user', 'content': code_task.DIAG_USER.format( + query=query, rubric=rubric_block, segment=segment_text)}]} + + def build_rubric_checker(): if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') or os.environ.get('OPENAI_API_KEY')): return None + if _TASK == 'code': + return _CodeRubricVerifier( + fixed_rubric=[RubricItem(t, is_hard=h) for t, h in code_task.CODE_RUBRIC], gate=True) return _RftRubricVerifier( fixed_rubric=[RubricItem(t, is_hard=h) for t, h in _MATH_RUBRIC], gate=True) @@ -632,6 +722,17 @@ def _diagnose_entry(checker, entry: Dict[str, Any]) -> Optional[str]: while GRPO trains. Shared by the background pre-diagnosis pool and distill_buffer's fallback for any entry the pool did not reach in time. 中文注释：单条失败轨迹的 rubric 诊断（纯 API，不吃 GPU）。后台预诊断与 distill 补诊断共用。""" + if _TASK == 'code': + # 代码域：query 里补上题面声明的硬约定（签名/返回/异常/示例，不含参考解答）， + # segment 由调用方（_rubric_entry）拼成"提交的代码 + 单测真实报错"。 + query = code_task.diag_query(entry['problem'], entry['reference_answer']) + seg = {'messages': [{'role': 'user', 'content': query}, + {'role': 'assistant', 'content': entry['fail_segment']}]} + try: + return _format_diagnosis(checker.diagnose(seg, query=query)) + except Exception as exc: + logger.warning(f'[rubric] diagnose error: {exc}') + return None seg_text = entry['fail_segment'] if entry.get('fail_stop_reason') == 'length': seg_text += ('\n\n[Process note: this attempt was cut off at the token budget ' @@ -668,12 +769,13 @@ def _diagnose_entry(checker, entry: Dict[str, Any]) -> Optional[str]: - Write it as one coherent analysis narrative (not a bullet list): first name what the problem is essentially asking, then walk through how to approach it, blending concepts, steps, pitfalls and reasons into a single connected story. - CRITICAL: Do NOT solve the problem for the executor. Do NOT reveal or compute the final answer, and do NOT substitute the problem's specific given numbers into the steps or state any intermediate numeric results. Leave ALL concrete numbers for the executor to compute on its own. If you catch yourself writing a specific number from the problem, replace it with a description of the quantity instead. - Keep it concise: aim for roughly one focused paragraph. +- End the block with this exact sentence: "Avoid re-checking loops; box a bare number as soon as it is computed." Put ONLY the methodology inside . Example: -This problem is essentially asking for the units (last) digit of an integer raised to a high power; first get clear on what the problem is asking before deciding where to start. Since only the last digit matters, you should first look only at the units digit of the base, because the units digit of an integer power is determined solely by the units digit of the base and the higher digits do not affect the result — so at this step be careful not to expand or compute the whole large number, which is both unnecessary and error-prone. Next, repeatedly multiply this units digit by itself and record the units digit each time, until it starts to repeat, thereby obtaining its cycle period. The part about "determining the period length" is important here: be careful not to count one term too many or too few, otherwise all the later positioning will be off. Finally, take the given exponent modulo the period length and land on the corresponding term within the period; here pay special attention that when the remainder is 0 it corresponds to the last term of the period rather than the first. Overall, I summarize the approach for this kind of problem as "first recognize that it asks for the units digit of a power, then fix on the units digit to find the cycle period, and finally use the exponent modulo to locate the term", while leaving the concrete numbers for the downstream solver to substitute and compute on its own. +This problem is essentially asking for the units (last) digit of an integer raised to a high power; first get clear on what the problem is asking before deciding where to start. Since only the last digit matters, you should first look only at the units digit of the base, because the units digit of an integer power is determined solely by the units digit of the base and the higher digits do not affect the result — so at this step be careful not to expand or compute the whole large number, which is both unnecessary and error-prone. Next, repeatedly multiply this units digit by itself and record the units digit each time, until it starts to repeat, thereby obtaining its cycle period. The part about "determining the period length" is important here: be careful not to count one term too many or too few, otherwise all the later positioning will be off. Finally, take the given exponent modulo the period length and land on the corresponding term within the period; here pay special attention that when the remainder is 0 it corresponds to the last term of the period rather than the first. Overall, I summarize the approach for this kind of problem as "first recognize that it asks for the units digit of a power, then fix on the units digit to find the cycle period, and finally use the exponent modulo to locate the term", while leaving the concrete numbers for the downstream solver to substitute and compute on its own. Avoid re-checking loops; box a bare number as soon as it is computed. """ @@ -770,6 +872,8 @@ def build_skill_solve_prompt_seam(problem, skill, raw_response=None): def build_direct_prompt(problem): + if _TASK == 'code': + return code_task.direct_prompt(problem) if _ALIGN_MODE == 'seam': # seam 基线保持英文原样，不受 v2 prompt 改动影响 return {'messages': [{'role': 'system', 'content': DIRECT_SYSTEM}, @@ -781,6 +885,9 @@ def build_direct_prompt(problem): def build_skill_solve_prompt(problem, skill, raw_response=None): skill = (skill or '').strip() + if _TASK == 'code': + # 空 skill -> 干净 direct（与数学分支同规则，见下方注释） + return code_task.skill_solve_prompt(problem, skill) if not skill: # 空 skill → 干净 direct。训练侧根本不会用空 skill 走 executor（process_chunk 只对非空 flat 跑， # 空候选直接 reward=0），故此分支仅影响 eval 口径——让空 skill 题 withskill==baseline、对 lift 贡献 0， @@ -813,13 +920,14 @@ def build_skill_solve_prompt(problem, skill, raw_response=None): Your steps: 1. Re-read and understand the original problem. 2. Tell a coherent analysis story for this problem as one flowing narrative: first identify what it is essentially asking, then walk through how to approach it, naturally weaving together the solving points that were already correct last time, the pitfalls that actually tripped up the solving process and how to avoid them, and your reasoning for why you give this advice, blended into a single connected story, and leave the concrete numbers for the downstream solver to compute. -3. Put the above inside . +3. End the block with this exact sentence: "Avoid re-checking loops; box a bare number as soon as it is computed." +4. Put the above inside . Output requirement: Write your judgments and pitfall reminders about this problem directly in the first person (e.g. "I think this step tends to ...", "A common mistake is ..., so you need to ..."), and phrase the issues you find as self-contained, general techniques. Do NOT use phrasings that point to external context such as "according to the given analysis/hints" or "the previous skill" — the downstream executor cannot see that context, and such phrasings will cause hallucination. Example: -This problem is essentially asking "how many arrangements satisfy the given constraints", which is a counting problem; first get clear on "what exactly is being counted" before deciding whether to use permutations or combinations. Since it is counting, you should first clearly define the objects being counted and the constraints, and judge whether the elements are distinguishable and whether order matters, because this directly determines whether you will need to divide out duplicates later. Next, first compute a total as if things were "ordered/distinguishable", then find which seemingly different arrangements actually correspond to the same configuration. The part about "recognizing symmetry and determining the duplication factor" is important here: I think the step most likely to go wrong in this problem is ignoring symmetry and treating essentially identical configurations as different, which makes the result too large; I think it is also easy to directly miss the "divide by the duplication factor" step — as long as the choices can be interchanged, you must divide out duplicates, otherwise you overcount. Finally, divide the total by the duplication factor to get the truly non-duplicated count; here pay special attention not to jump straight to permutation/combination formulas, but first think clearly about whether the elements are distinguishable and then decide whether to divide out duplicates. Overall, I summarize the approach for this kind of problem as "first recognize that it is a counting problem and judge whether the elements are distinguishable, then compute the total, recognize symmetry and remove duplicates", because I judge that the loss points for such problems almost all concentrate on overcounting; while leaving the concrete numbers for the downstream solver to substitute and compute on its own. +This problem is essentially asking "how many arrangements satisfy the given constraints", which is a counting problem; first get clear on "what exactly is being counted" before deciding whether to use permutations or combinations. Since it is counting, you should first clearly define the objects being counted and the constraints, and judge whether the elements are distinguishable and whether order matters, because this directly determines whether you will need to divide out duplicates later. Next, first compute a total as if things were "ordered/distinguishable", then find which seemingly different arrangements actually correspond to the same configuration. The part about "recognizing symmetry and determining the duplication factor" is important here: I think the step most likely to go wrong in this problem is ignoring symmetry and treating essentially identical configurations as different, which makes the result too large; I think it is also easy to directly miss the "divide by the duplication factor" step — as long as the choices can be interchanged, you must divide out duplicates, otherwise you overcount. Finally, divide the total by the duplication factor to get the truly non-duplicated count; here pay special attention not to jump straight to permutation/combination formulas, but first think clearly about whether the elements are distinguishable and then decide whether to divide out duplicates. Overall, I summarize the approach for this kind of problem as "first recognize that it is a counting problem and judge whether the elements are distinguishable, then compute the total, recognize symmetry and remove duplicates", because I judge that the loss points for such problems almost all concentrate on overcounting; while leaving the concrete numbers for the downstream solver to substitute and compute on its own. Avoid re-checking loops; box a bare number as soon as it is computed. """ REGEN_USER = """\ @@ -860,6 +968,9 @@ def build_skill_solve_prompt(problem, skill, raw_response=None): def _skillgen_prompt(problem: str) -> Dict[str, Any]: """Skill-gen prompt: query-only. seam mode uses SEAM EXPERIENCE_PROMPT (single user turn, output); v2 uses SKILL_GEN_SYSTEM ().""" + if _TASK == 'code': + # 代码域只做 narrative 一种文体（E4/E17 都是 narrative；toy/pitfall 未移植） + return code_task.skillgen_prompt(problem) if _ALIGN_MODE == 'seam': return {'messages': [{'role': 'user', 'content': _SEAM_EXPERIENCE_PROMPT.format(problem=problem)}]} # 中文注释：按 --skill-style 选主链路文体（narrative=现版叙述式 / toy / pitfall）。 @@ -880,10 +991,11 @@ def _regen_prompt(problem: str, orig_skill: str, rubric_diag: str) -> Dict[str, # ---- Reward ---- -# 中文注释：reward = parseable × correct（对齐 SEAM lpem：去 terminated、去长度惩罚）。 -# parseable=0 的候选 reward=0 仍参与 group（格式压力）。 -def _skill_reward(parseable: bool, correct: bool) -> float: - return 1.0 if (parseable and correct) else 0.0 +# 中文注释：reward = parseable × 通过率（对齐 SEAM lpem：去 terminated、去长度惩罚）。 +# parseable=0 的候选 reward=0 仍参与 group（格式压力）。correct 兼容 bool（greedy 0/1， +# E1-E13）与 float 通过率（E14 多 rollout 判分，见 process_chunk reward_rollouts）。 +def _skill_reward(parseable: bool, correct) -> float: + return float(correct) if parseable else 0.0 # ---- Buffer A: collect adv=0 all-fail problems ---- @@ -1195,16 +1307,29 @@ def process_chunk(base_sampler, skill_sampler, chunk, ci, base_dp, skill_dp, arg for r, c in flat: c['leaked'] = _answer_leaked(c['skills'], r['reference_answer']) - # with-skill greedy pass (T=0) + # with-skill executor pass：默认 greedy×1（E1-E13，reward 0/1 与旧口径 bit 一致）； + # E14: reward_rollouts>1 时 T=reward_temperature × K 采样，reward = parseable × 通过率， + # 把内容信号从 greedy 0/1 量化里释放出来（提升组内 std>0 比例）。 if flat: + K = max(1, int(getattr(args, 'reward_rollouts', 1) or 1)) + rT = float(getattr(args, 'reward_temperature', 0.0) or 0.0) ws_out = _run_samples(base_sampler, [build_skill_solve_prompt(r['problem'], c['skills'], c.get('response')) for r, c in flat], - 1, args.max_tokens, base_dp, temperature=0.0) + K, args.max_tokens, base_dp, temperature=rT) + # 判分一次性批量化（code 任务要起子进程跑单测，逐条会比 GPU 还慢一个量级） + pairs, spans = [], [] for (r, c), seqs in zip(flat, ws_out): - roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - c['rolls'] = [roll] - c['with_pass'] = 1.0 if roll['correct'] else 0.0 - c['reward'] = _skill_reward(c['parseable'], roll['correct']) + start = len(pairs) + pairs.extend((s, r['reference_answer']) for s in (seqs or [])) + spans.append((start, len(pairs))) + judged = _parse_many(pairs) + for (r, c), (a, b) in zip(flat, spans): + rolls = judged[a:b] or [_empty_roll()] + for x in rolls[1:]: + x['text'] = '' # 磁盘保护：K>1 时只留首 rollout 全文（gen_records 体积控制） + c['rolls'] = rolls + c['with_pass'] = sum(1.0 for x in rolls if x['correct']) / len(rolls) + c['reward'] = _skill_reward(c['parseable'], c['with_pass']) # unparseable candidates score 0 and still join the group (format pressure) for r in chunk: for c in r['_cands']: @@ -1229,7 +1354,12 @@ def process_chunk(base_sampler, skill_sampler, chunk, ci, base_dp, skill_dp, arg def _roll(x): - return {k: x[k] for k in ('pred', 'correct', 'terminated', 'stop_reason', 'gen_tokens', 'text')} + out = {k: x[k] for k in ('pred', 'correct', 'terminated', 'stop_reason', 'gen_tokens', 'text')} + # 代码域审计字段：判分结论 / 单测报错 / 用例数（离线分析错误类型分布靠它，_trim_err 已限长） + for k in ('kind', 'error', 'n_tests'): + if k in x: + out[k] = x[k] + return out def _full_records(chunk, ci): @@ -1244,6 +1374,8 @@ def _full_records(chunk, ci): 'advantage': c.get('advantage'), 'kept': c.get('kept'), 'skillgen_stop': c.get('skillgen_stop'), 'skillgen_tokens': c.get('skillgen_tokens'), 'rolls': [_roll(x) for x in c['rolls']], + 'logp_base': c.get('logp_base'), 'logp_skill': c.get('logp_skill'), + 'logp_delta': c.get('logp_delta'), } for c in r['_cands']], }) return out @@ -1279,7 +1411,9 @@ def _chunk_summary(chunk, ci): zero_grad += 1 n_train = sum(1 for c in all_cands if abs(c.get('advantage') or 0.0) > 1e-9) trunc = sum(1 for x in ws_rolls if x['stop_reason'] == 'length') - ws_acc = _mean([1.0 if any(c.get('reward') for c in r['_cands']) else 0.0 + # bugfix（ablate #6）：旧版 any(c.get('reward')) 用 truthiness，负 reward（E16 hinge/leak_gate、 + # E14 地板 -1.0）也被当“通过”；改用 with_pass>0（真正的 executor 通过率），greedy 0/1 臂语义不变。 + ws_acc = _mean([1.0 if any((c.get('with_pass') or 0) > 0 for c in r['_cands']) else 0.0 for r in chunk if r['_cands']]) return { 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), @@ -1311,8 +1445,9 @@ def run_greedy_eval(base_sampler, skill_sampler, eval_records, ci, rounds, if todo: out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], 1, args.max_tokens, base_dp, temperature=0.0) - for r, seqs in zip(todo, out): - roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() + rolls = _parse_many([(_first_seq(seqs), r['reference_answer']) + for r, seqs in zip(todo, out)]) + for r, roll in zip(todo, rolls): base_cache.put(DiskCache.key_for(r['problem']), roll) for r in eval_records: br = base_cache.get(DiskCache.key_for(r['problem'])) @@ -1341,9 +1476,9 @@ def run_greedy_eval(base_sampler, skill_sampler, eval_records, ci, rounds, flat_prompts.append(build_skill_solve_prompt(r['problem'], sk, sresp)) flat_idx.append((pi, j)) ws_out = _run_samples(base_sampler, flat_prompts, 1, args.max_tokens, base_dp, temperature=0.0) - roll_by = {} - for (pi, j), seqs in zip(flat_idx, ws_out): - roll_by[(pi, j)] = _parse_seq(seqs[0], eval_records[pi]['reference_answer']) if seqs else _empty_roll() + judged = _parse_many([(_first_seq(seqs), eval_records[pi]['reference_answer']) + for (pi, _j), seqs in zip(flat_idx, ws_out)]) + roll_by = {idx: roll for idx, roll in zip(flat_idx, judged)} recs = [] for pi, (r, row) in enumerate(zip(eval_records, per_skills)): rolls = [roll_by[(pi, j)] for j in range(len(row))] @@ -1351,12 +1486,17 @@ def run_greedy_eval(base_sampler, skill_sampler, eval_records, ci, rounds, parses = [1.0 if sk else 0.0 for sk, _ in row] terms = [1.0 if x['terminated'] else 0.0 for x in rolls] acc_mean = sum(corr) / len(corr) if corr else 0.0 + # bugfix（ablate #8）：unparseable skill 的 rollout 实际走了 direct 回退（≈baseline）， + # 主指标里格式崩塌会被 baseline 成绩掩护。strict 通道：unparseable 计 0，格式失败 + # 直接计入代价；主指标口径不变（与历史臂可比），两条曲线分叉即格式崩塌告警。 + acc_strict = (sum(c * p for c, p in zip(corr, parses)) / len(corr)) if corr else 0.0 recs.append({ 'record_type': 'eval_problem', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, 'data_id': r.get('data_id', ''), 'problem': r['problem'], 'reference_answer': r['reference_answer'], 'baseline_pass': r['_baseline_pass'], 'n_rollouts': len(row), 'eval_skill_temperature': args.eval_skill_temperature, 'withskill_acc_mean': acc_mean, # per-problem mean over R rollouts + 'withskill_acc_strict_mean': acc_strict, # unparseable counted wrong 'withskill_pass_any': 1.0 if any(corr) else 0.0, # pass@R (bonus readout) 'skill_parseable_mean': sum(parses) / len(parses) if parses else 0.0, 'withskill_terminated_mean': sum(terms) / len(terms) if terms else 0.0, @@ -1369,6 +1509,7 @@ def run_greedy_eval(base_sampler, skill_sampler, eval_records, ci, rounds, n = len(recs) # acc = 跨题平均的"每题 R 次平均正确率"（mean-over-rollouts） ws = (sum(x['withskill_acc_mean'] for x in recs) / n) if n else 0.0 + ws_strict = (sum(x['withskill_acc_strict_mean'] for x in recs) / n) if n else 0.0 pass_any = (sum(x['withskill_pass_any'] for x in recs) / n) if n else 0.0 base = (sum(x['baseline_pass'] for x in recs) / n) if n else 0.0 fmt = (sum(x['skill_parseable_mean'] for x in recs) / n) if n else 0.0 @@ -1379,11 +1520,13 @@ def run_greedy_eval(base_sampler, skill_sampler, eval_records, ci, rounds, hard_rescued = sum(x['withskill_acc_mean'] for x in hard) # 期望救活数(分数) summary = {'record_type': 'eval_summary', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, 'n': n, 'acc_mean1': ws, 'baseline_acc_mean1': base, 'lift_mean1': ws - base, + 'acc_strict_mean1': ws_strict, 'lift_strict_mean1': ws_strict - base, 'acc_pass_any': pass_any, 'n_rollouts': R, 'eval_skill_temperature': args.eval_skill_temperature, 'format_mean1': fmt, 'term_mean1': term, 'hard_n': len(hard), 'hard_rescued': hard_rescued, 'hard_rescue_rate': hard_rescue_rate} metrics = {'core/math/acc/mean@1': ws, 'core/math/baseline_acc/mean@1': base, 'core/math/lift/mean@1': ws - base, 'core/math/format/mean@1': fmt, + 'core/math/acc_strict/mean@1': ws_strict, 'core/math/lift_strict/mean@1': ws_strict - base, 'core/math/term/mean@1': term, 'core/math/hard_rescue/mean@1': hard_rescue_rate} return recs, summary, metrics @@ -1416,7 +1559,8 @@ def init_components(args): # 因此不构成 SEAM 那种“把 think 喂给 executor”的泄漏。skill_model/ref_model/skill_sampler 三者 # enable_thinking 必须一致，否则训练轨迹 token 布局与采样对不上。 # 中文注释：skill_model/ref_model/skill_sampler 三者 enable_thinking 由 --skill-thinking 统一控制 - # （必须一致，否则训练轨迹 token 布局与采样对不上）；base_sampler（executor）恒 thinking on。 + # （必须一致，否则训练轨迹 token 布局与采样对不上）；base_sampler（executor）走独立开关 + # --executor-thinking（默认 on，E1-E18 全部 on；E19/E20 为 off，见下方 base_sampler 处注释）。 _think = args.skill_thinking == 'on' skill_model.set_template(Template, model_id=MODEL_ID, enable_thinking=_think, max_length=args.max_model_len, truncation_strategy='delete') @@ -1453,7 +1597,16 @@ def _sampler(group, world, enable_thinking): # 方案1：skill 采样器开 thinking，与 skill_model/ref_model 一致（actor 先想再写 ）。 skill_sampler = _sampler('skill_sampler', SKILL_SAMPLER_GPUS, enable_thinking=_think) - base_sampler = _LockedSampler(_sampler('base_sampler', BASE_SAMPLER_GPUS, enable_thinking=True)) + # executor（base_sampler）的 thinking 单独一个开关，默认 on（E1-E18 全部如此）。 + # ⭐ 为什么要能关（2026-07-31 bcb 探针实测，n=275 同题配对）：BigCodeBench 上 think 的 + # executor 有 34-50% 的 rollout 撞满预算、连代码块都没写出来（截断样本 8-gram 重复率 + # p50=0.835、同一长句重复 92 次 = 字面死循环），把预算从 4096 加到 20000 也只把截断 + # 从 0.496 压到 0.338 且 pass 不涨。关掉 thinking 后截断归零、裸解反而更高 + # （0.378 vs 0.324），rubric 增量从 +0.080 抬到 +0.135（p=1e-4）。 + # 见 bcb/bcb_eval0_{nothink,think12k,think20k}.jsonl。 + _exec_think = getattr(args, 'executor_thinking', 'on') == 'on' + base_sampler = _LockedSampler(_sampler('base_sampler', BASE_SAMPLER_GPUS, + enable_thinking=_exec_think)) ckpt = CheckpointEngineManager(model=skill_model, sampler=skill_sampler) return skill_model, ref_model, skill_sampler, base_sampler, ckpt, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS @@ -1494,6 +1647,10 @@ def _build_args(): '主链路与 regen 同文体。') p.add_argument('--skill-thinking', choices=('on', 'off'), default='on', help='skill_model/ref_model/skill_sampler 三者的 enable_thinking（必须一致）') + p.add_argument('--executor-thinking', choices=('on', 'off'), default='on', + help='executor(base_sampler) 的 enable_thinking。off 用于 BigCodeBench 这类' + '"解答短、难点在选 API 而非多步推理"的任务：think 下 34-50% 的 rollout ' + '陷入字面死循环撞满预算，关掉后截断归零且裸解更高（见 build 处注释）。') p.add_argument('--align-mode', choices=('v2', 'seam'), default='v2', help="SEAM-alignment toggle for PROMPT/SKILL FORMAT only. " "'v2'=clean single-user executor prompt + skill-gen. " @@ -1546,7 +1703,9 @@ def _build_args(): p.add_argument('--adv-clip', type=float, default=0.0, help='clip group-relative advantage to [-adv_clip, adv_clip]; ' '0 = no clipping (matches SEAM/verl GRPO which does not clip advantages)') - p.add_argument('--kl-beta', type=float, default=0.001) + # 与 skill_ablate/main.py 的默认值必须一致（两个入口默认值分叉 = 静默不可比）。 + # 2026-07-29 拍板 0.001 -> 0.01：对抗侵蚀 executor 收束能力的自发漂移。 + p.add_argument('--kl-beta', type=float, default=0.01) p.add_argument('--lr', type=float, default=6e-6) p.add_argument('--max-train-rounds', type=int, default=1500) p.add_argument('--save-rounds', type=int, default=200) From 987c2abd86aab1598ca2d4c1ee43600ba134b51b Mon Sep 17 00:00:00 2001 From: root Date: Sat, 1 Aug 2026 16:20:43 +0800 Subject: [PATCH 28/60] skill2lora: add E21 freeform-style arm (model picks any useful skill form) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - train_skill_v2.py: SKILL_GEN_FREEFORM/REGEN_FREEFORM_SYSTEM (a 'menu' prompt letting the skill model choose whatever form helps this problem — analysis, concept, pitfall, tiny example, blunt directive, even 'let's think step by step'), wired into style dispatch + --skill-style choices; the freeform prompt carries wrapper examples so open-form outputs stay parseable - config.py: E21 = bnpo/view-B/freeform (thinking on; see comment for why not off), STYLES + RUN_ORDER updated, self-check passes - trainer.py: freeform shares narrative's 1100-char len budget --- .../exp/skill2lora/skill_ablate/config.py | 13 +++- .../exp/skill2lora/skill_ablate/trainer.py | 4 +- cookbook/exp/skill2lora/train_skill_v2.py | 66 +++++++++++++++++-- 3 files changed, 72 insertions(+), 11 deletions(-) diff --git a/cookbook/exp/skill2lora/skill_ablate/config.py b/cookbook/exp/skill2lora/skill_ablate/config.py index 1d1a125fa..ea84b1d33 100644 --- a/cookbook/exp/skill2lora/skill_ablate/config.py +++ b/cookbook/exp/skill2lora/skill_ablate/config.py @@ -44,7 +44,7 @@ VIEW_OF_METHOD = {'bnpo': 'B', 'rl_ab': 'A', 'rl_err': 'A', 'reflexion': 'A', 'opsd': 'A', 'improve_sft': 'A', 'sft': 'A', 'logp_rl': 'A', 'logp_gt': 'B', 'passrate_hinge': 'B', 'rejection_sft': 'A'} -STYLES = ('narrative', 'pitfall') +STYLES = ('narrative', 'pitfall', 'freeform') THINKINGS = ('on', 'off') TASKS = ('math', 'code') @@ -249,6 +249,15 @@ def swanlab_exp(self) -> str: reward_rollouts=1, reward_temperature=0.0, smt_override=8192), ExpSpec('E20', 'reflexion', 'on', 'narrative', task='code', executor_thinking='off', reward_rollouts=1, reward_temperature=0.0, smt_override=8192), + # group 11 — E21 freeform 文体（2026-08-01 用户拍板）：以 E2 为模板的 query-only BNPO（view B、 + # 无 rubric），唯一变量是 skill-gen system 换成 SKILL_GEN_FREEFORM“招式菜单”：不锁 narrative/ + # pitfall/toy 固定文体，让模型按题自选最有用的形态（分析/概念/预判纠错/迷你示范/直白执行 + # 指令，甚至 “let's think step by step”），T=1.0×8 自然铺开、组内择优。动机：固定 hint + # 消融实测 hint 内容语义贡献≈0、增益几乎全来自“有个 skill 块 + 催答案收尾”（fixed_hint_probe.py： + # A9_wrapperonly/A4_garbage 与有义 hint 打平、A7_budget 最高 +0.16），故放开文体看模型能否自选出更优组合。 + # ★ 刷 thinking='on'（非照 E2 的 off）：freeform prompt 依赖“先私下想再选形态”，nothink 下无处 + # 思考会把推理直接写进 （line 758 记录的泄漏失败模式）。其余同 E2：bnpo/math/narrative-长度预算。 + ExpSpec('E21', 'bnpo', 'on', 'freeform'), ] # execution order: all nothink first, then think; E13 (seam-align baseline) right after E6; @@ -256,7 +265,7 @@ def swanlab_exp(self) -> str: # E15 (GT-target logP validation) right after E14; the data-hungry SFT method dead last. # 2026-07-31 用户拍板：E19/E20（executor nothink）排在 E4/E17 之前先跑 —— 探针已判定 # nothink 是 rubric 增量最大且唯一无截断混杂的口径，先拿这两个臂的结论。 -RUN_ORDER: List[str] = ['E1', 'E2', 'E5', 'E6', 'E13', 'E3', 'E7', 'E14', 'E15', 'E19', 'E20', 'E4', 'E8', 'E16', 'E17', 'E18', 'E9', 'E10', 'E11', 'E12'] +RUN_ORDER: List[str] = ['E1', 'E2', 'E5', 'E6', 'E13', 'E3', 'E7', 'E14', 'E15', 'E19', 'E20', 'E21', 'E4', 'E8', 'E16', 'E17', 'E18', 'E9', 'E10', 'E11', 'E12'] BY_NAME: Dict[str, ExpSpec] = {e.name: e for e in MATRIX} diff --git a/cookbook/exp/skill2lora/skill_ablate/trainer.py b/cookbook/exp/skill2lora/skill_ablate/trainer.py index d055cb8fe..62cc3e8c8 100644 --- a/cookbook/exp/skill2lora/skill_ablate/trainer.py +++ b/cookbook/exp/skill2lora/skill_ablate/trainer.py @@ -119,9 +119,9 @@ def run_experiment(args, spec: ExpSpec) -> None: v2._SKILL_STYLE = spec.style args.skill_thinking = spec.thinking # per-style length budget (#9 statistics: narrative≈1100 / pitfall≈300 chars); - # an explicit --len-budget on the CLI wins. + # freeform 可能产出叙述式长文本，按 narrative 档给 1100 以免误伤；explicit --len-budget on the CLI wins. if args.len_budget is None: - args.len_budget = 1100 if spec.style == 'narrative' else 300 + args.len_budget = 1100 if spec.style in ('narrative', 'freeform') else 300 # explicit --skill-max-tokens on the CLI wins over the per-experiment default. if args.skill_max_tokens is None: args.skill_max_tokens = spec.skill_max_tokens diff --git a/cookbook/exp/skill2lora/train_skill_v2.py b/cookbook/exp/skill2lora/train_skill_v2.py index 3e59fe245..78d951904 100644 --- a/cookbook/exp/skill2lora/train_skill_v2.py +++ b/cookbook/exp/skill2lora/train_skill_v2.py @@ -803,7 +803,44 @@ def _diagnose_entry(checker, entry: Dict[str, Any]) -> Optional[str]: - End with: "Avoid re-checking loops; box a bare number as soon as it is computed." """ -_SKILL_STYLE = 'narrative' # 'narrative' | 'toy' | 'pitfall'；由 main() 依据 --skill-style 设置 +# freeform 主链路（2026-08-01 用户拍板）：不锁死 narrative/pitfall/toy 固定文体，而是给模型一份 +# "招式菜单"，让它按题自选最有用的形态（可组合、可极简）。设计目标是让 T=1.0×8 的候选自然铺开 +# 到不同形态（分析 / 概念 / 预判纠错 / 迷你示范 / 直白执行指令，甚至 "let's think step by step"）， +# 由 GRPO/BNPO 组内择优。依据：固定 hint 消融（good_skill_hard_fail/fixed_hint_probe.py）显示 hint +# 的"内容语义"贡献≈0、增益几乎全来自"存在一个 skill 块 + 催答案收尾"（A7_budget 最高 +0.16、 +# exec_answered 3.95σ），所以放开文体、把"催收尾"作为可选招式之一，看模型能否自选出更优组合。 +# 硬约束仍与 narrative 一致：只输出 块、不解题、不代入本题数值、不给最终答案。 +SKILL_GEN_FREEFORM = """\ +You are a skill-generation model. Your block will be fed to a SEPARATE downstream executor model that must solve the problem on its own. The executor will NOT see your private reasoning — it only sees what is inside .... + +First, think privately: actually work the problem out in your head until you understand what really makes it solvable, then decide what ONE kind of help would most raise a fresh solver's chance on THIS specific problem. + +There is NO fixed format and no required style. Different problems are helped by different things — pick whatever you judge most useful here. Any of the following is allowed (the list is not exhaustive, and you may blend a couple if they genuinely help): +- a short transferable analysis of what this TYPE of problem is really asking and the recommended approach; +- naming the key concept / theorem / trick to reach for; +- a WARNING about the single most likely wrong turn on this type, and the correct move instead; +- a tiny worked example of the SAME type using DIFFERENT, smaller numbers (never the problem's own); +- a blunt execution directive that keeps the solver on track (e.g. "commit to one method and don't keep second-guessing", "let's think step by step", or "box a bare number as soon as it is computed"); +- or plain, nothing-fancy encouragement if that is honestly all this problem needs. +- Any other freeform skill you can imagine to try on this query + +Choose the form that fits THIS problem; do not pad. If one sharp sentence is the best help, give only that sentence; if a short focused paragraph is warranted, keep it tight. Being genuinely useful matters far more than being long or elaborate. + +Hard rules (always apply, whatever form you pick): +- Do NOT solve the problem for the executor. Do NOT reveal or compute the final answer, and do NOT substitute the problem's specific given numbers or state any intermediate numeric result — leave ALL concrete numbers for the executor to compute. +- Put ONLY your chosen help inside , and nothing else. + +Whatever form you choose, it MUST be wrapped in a single block with a proper closing tag. For example, a rich form: + +This is a modular-arithmetic problem: reduce each factor modulo the given modulus before multiplying, and never expand the full product — that is the whole trick. Commit to that reduction and don't second-guess it, then box a bare number as soon as it is computed. + +or, when the problem only needs a nudge, a minimal form is equally valid: + +Let's think step by step, and box a bare number as soon as it is computed. + +""" + +_SKILL_STYLE = 'narrative' # 'narrative' | 'toy' | 'pitfall' | 'freeform'；由 main() 依据 --skill-style 设置 # ---- Executor prompt (with skill injection) ---- # 中文注释：executor 提示词。答案格式已统一为 \boxed{}（人工拍板，2026-07-27）：seam/v2 两模式的 @@ -964,6 +1001,19 @@ def build_skill_solve_prompt(problem, skill, raw_response=None): Hard rules: the block must be self-contained - never reference "the diagnosis" or "the previous skill"; the executor cannot see them. """ +# freeform 的 regen 版（buffer B 蒸馏用）。bnpo/view-B 臂不会走 regen，此处仅为分派完整性与 +# 未来 view-A + freeform 组合预留；同样放开形态、保留"自持、不指涉外部上下文、不泄漏"硬规则。 +REGEN_FREEFORM_SYSTEM = """\ +You are a skill-generation model. A separate executor model previously FAILED this problem even with your earlier skill. You will see that earlier skill and an expert rubric diagnosis of the failure. The executor will retry seeing ONLY your new block. + +First think privately: from the diagnosis, pinpoint the ONE thing that actually went wrong. Then choose whatever form of help would best fix it for THIS problem — there is no fixed format. It may be a short transferable analysis, the key concept to reach for, a WARNING naming the decisive mistake plus the correct move instead, a tiny worked example with DIFFERENT smaller numbers, or a blunt execution directive (e.g. "box a bare number as soon as it is computed"). Blend a couple only if it genuinely helps, and do not pad. + +Hard rules: +- Do NOT solve the problem or reveal/compute the final answer, and do NOT substitute the problem's own numbers. +- The block must be self-contained — never reference "the diagnosis", "the previous skill", or any context the executor cannot see, or it will hallucinate. +- Put ONLY your chosen help inside . +""" + def _skillgen_prompt(problem: str) -> Dict[str, Any]: """Skill-gen prompt: query-only. seam mode uses SEAM EXPERIENCE_PROMPT (single user turn, @@ -974,7 +1024,8 @@ def _skillgen_prompt(problem: str) -> Dict[str, Any]: if _ALIGN_MODE == 'seam': return {'messages': [{'role': 'user', 'content': _SEAM_EXPERIENCE_PROMPT.format(problem=problem)}]} # 中文注释：按 --skill-style 选主链路文体（narrative=现版叙述式 / toy / pitfall）。 - sys_p = {'toy': SKILL_GEN_TOY, 'pitfall': SKILL_GEN_PITFALL}.get(_SKILL_STYLE, SKILL_GEN_SYSTEM) + sys_p = {'toy': SKILL_GEN_TOY, 'pitfall': SKILL_GEN_PITFALL, + 'freeform': SKILL_GEN_FREEFORM}.get(_SKILL_STYLE, SKILL_GEN_SYSTEM) return {'messages': [ {'role': 'system', 'content': sys_p}, {'role': 'user', 'content': f'Problem:\n{problem}'}]} @@ -983,7 +1034,8 @@ def _skillgen_prompt(problem: str) -> Dict[str, Any]: def _regen_prompt(problem: str, orig_skill: str, rubric_diag: str) -> Dict[str, Any]: """Regeneration prompt for buffer B distillation.""" # 中文注释：regen 与主链路同文体（--skill-style），user 模板复用 REGEN_USER 三字段。 - sys_p = {'toy': REGEN_TOY_SYSTEM, 'pitfall': REGEN_PITFALL_SYSTEM}.get(_SKILL_STYLE, REGEN_SYSTEM) + sys_p = {'toy': REGEN_TOY_SYSTEM, 'pitfall': REGEN_PITFALL_SYSTEM, + 'freeform': REGEN_FREEFORM_SYSTEM}.get(_SKILL_STYLE, REGEN_SYSTEM) return {'messages': [ {'role': 'system', 'content': sys_p}, {'role': 'user', 'content': REGEN_USER.format( @@ -1642,9 +1694,9 @@ def _build_args(): # 会截断成空块（_extract_skill 找不到 返回 None）。提到 8192 给两段都留足空间。 p.add_argument('--skill-max-tokens', type=int, default=8192) # 中文注释：文体消融开关——主链路与 buffer B regen 同文体（分布一致才可联合训练）。 - p.add_argument('--skill-style', choices=('narrative', 'toy', 'pitfall'), default='narrative', - help='skill文体: narrative=现版叙述式; toy=异数字玩具题示范; pitfall=预判纠错。' - '主链路与 regen 同文体。') + p.add_argument('--skill-style', choices=('narrative', 'toy', 'pitfall', 'freeform'), default='narrative', + help='skill文体: narrative=现版叙述式; toy=异数字玩具题示范; pitfall=预判纠错; ' + 'freeform=招式菜单/模型按题自选形态。主链路与 regen 同文体。') p.add_argument('--skill-thinking', choices=('on', 'off'), default='on', help='skill_model/ref_model/skill_sampler 三者的 enable_thinking（必须一致）') p.add_argument('--executor-thinking', choices=('on', 'off'), default='on', @@ -1760,7 +1812,7 @@ def main(): args = _build_args() global _ALIGN_MODE, _SKILL_STYLE _ALIGN_MODE = args.align_mode # 'v2' | 'seam' - _SKILL_STYLE = args.skill_style # 'narrative' | 'toy' | 'pitfall' + _SKILL_STYLE = args.skill_style # 'narrative' | 'toy' | 'pitfall' | 'freeform' records, eval_records = _load_records(args) if len(records) < args.chunk_size: raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded ({len(records)}); raise --n') From 2502665ddec7d1c5b92b6a2191090a0c06fb5e65 Mon Sep 17 00:00:00 2001 From: root Date: Sat, 1 Aug 2026 17:01:12 +0800 Subject: [PATCH 29/60] loss: fix BNPO/SEAMBNPO token-mean to be micro/dp-split invariant The BNPO family returned an already-normalized per-group token-mean with num_tokens=0, so the framework's PER-TOKEN-MEAN path equal-weighted micro/dp groups -> a double average (group token-mean, then equal weight over groups) that sits between token-mean and sequence-mean and biases toward short responses (degrades to pure sequence-mean as groups multiply). This diverged from verl/SEAM's true token-mean and was non-orthogonal to skill-length study. - grpo.py: BNPOLoss gains token_mean_scope='global'(default)|'micro'. 'global' returns the token SUM and reports num_tokens=sum(mask), routing into the framework SUM-loss path -> exact global token-mean, invariant to how the batch is split. 'micro' preserves the old behavior to reproduce E1-E20. Added a _loss_num_tokens hook (default 0) so GRPO/DRGRPO/OPSD are untouched. No public interface change; downstream grad + metric already branch on num_tokens. - tests/loss/test_bnpo_token_mean.py: assert global is split-invariant (==true token-mean), micro reproduces the biased double-average, SEAM inherits global. - run_ablate12.sh: E13(140G)/E21(80G) OOM'd in train forward at micro=8; set per-arm train_micro_batch defaults (E13=2xdp, E21=1xdp). The global token-mean fix makes shrinking micro mathematically equivalent, so effective batch and comparability are unchanged. --- cookbook/exp/skill2lora/run_ablate12.sh | 19 +++++- src/twinkle/loss/grpo.py | 46 ++++++++++++-- tests/loss/test_bnpo_token_mean.py | 82 +++++++++++++++++++++++++ 3 files changed, 142 insertions(+), 5 deletions(-) create mode 100644 tests/loss/test_bnpo_token_mean.py diff --git a/cookbook/exp/skill2lora/run_ablate12.sh b/cookbook/exp/skill2lora/run_ablate12.sh index 991a09fa4..4016c970e 100644 --- a/cookbook/exp/skill2lora/run_ablate12.sh +++ b/cookbook/exp/skill2lora/run_ablate12.sh @@ -274,9 +274,26 @@ while IFS=$'\t' read -r NAME EXP_DIR THINK SMT OPTIONAL TASK EXEC_THINK; do MIN_LEVEL_ARG="${E13_MIN_LEVEL:-0}" [ -z "$REWARD_TRUNC_PENALTY" ] && E16_FLAGS="$E16_FLAGS --reward-trunc-penalty 0" E16_FLAGS="$E16_FLAGS --eval-rollouts 1 --eval-skill-temperature 0.0" - echo "[ablate12] E13 SEAM-repro: chunk=$CHUNK_ARG min_level=$MIN_LEVEL_ARG"\ + # 显存：E13 8×140G 实测 micro=8（trainer.py 的 8192 自动档）仍 OOM（132G 已用 + 4.77G）。 + # 降 train_micro_batch 到 2×dp（默认 TRAIN_GPUS=2 → 4）；grpo.py 的 global token-mean 修复后 + # 切 micro 数学等价，不改有效批量(chunk=128=SEAM batch)与可比性，只降训练前向峰值显存。 + # 显式 TRAIN_MICRO_BATCH 优先；E13_TRAIN_MICRO_BATCH 单独可调。值恒为 TRAIN_GPUS 倍数(满足 %TRAIN_DP)。 + _tmb=${E13_TRAIN_MICRO_BATCH:-$((2*${TRAIN_GPUS:-2}))} + [ -n "$TRAIN_MICRO_BATCH" ] && _tmb=$TRAIN_MICRO_BATCH + E16_FLAGS="$E16_FLAGS --train-micro-batch $_tmb" + echo "[ablate12] E13 SEAM-repro: chunk=$CHUNK_ARG min_level=$MIN_LEVEL_ARG train_micro_batch=$_tmb"\ " reward_trunc_penalty=0 eval=R1/T0 executor=nothink skill_max_tokens=$SMT (executor 预算 8192/16384, K=n_skills 默认 8)" fi + if [ "$NAME" = "E21" ]; then + # 显存：E21 每卡 80G，4B 不分片 + fp32 Adam 主权重≈64G、余量仅 ~16G，micro=8（自动档）必 OOM。 + # 默认 train_micro_batch=1×dp（= TRAIN_GPUS，最小且满足 %TRAIN_DP==0）；global token-mean 修复后 + # 切 micro 数学等价。若仍 OOM：改用 TRAIN_GPUS=1（micro→1）或 TRAIN_FSDP=2 分片权重腾显存。 + # 显式 TRAIN_MICRO_BATCH 优先；E21_TRAIN_MICRO_BATCH 单独可调。 + _tmb=${E21_TRAIN_MICRO_BATCH:-${TRAIN_GPUS:-2}} + [ -n "$TRAIN_MICRO_BATCH" ] && _tmb=$TRAIN_MICRO_BATCH + E16_FLAGS="$E16_FLAGS --train-micro-batch $_tmb" + echo "[ablate12] E21 freeform: train_micro_batch=$_tmb (80G OOM guard; global token-mean 使切 micro 等价，不改批量)" + fi case "$NAME" in E4|E17|E19|E20) # 四个臂统一 executor 预算 15000 + max_model_len 20480 + 长度惩罚死区 10000， diff --git a/src/twinkle/loss/grpo.py b/src/twinkle/loss/grpo.py index a01ecda84..12d0416fe 100644 --- a/src/twinkle/loss/grpo.py +++ b/src/twinkle/loss/grpo.py @@ -131,6 +131,17 @@ def _aggregate_loss( # Each sequence contributes equally regardless of length. return ((per_token_loss * loss_mask).sum(-1) / loss_mask.sum(-1).clamp(min=1.0)).mean() + def _loss_num_tokens(self, loss_mask: 'torch.Tensor'): + """Token denominator reported in ``LossOutput.num_tokens``. + + 0 (default) => framework uses the PER-TOKEN-MEAN accumulation path, where each + micro/dp group is equal-weighted. Subclasses that want a strict GLOBAL token-mean + (the SUM-loss path in transformers.py / megatron.py) return ``Σmask`` instead, so + the accumulated gradient is divided by the global token count and the result is + invariant to how the batch is split into micro/dp groups. + """ + return 0 + def _pad_and_align_to_batch( self, data: 'Union[torch.Tensor, List, np.ndarray]', @@ -315,7 +326,7 @@ def __call__( loss = self._aggregate_loss(per_token_loss, loss_mask, **kwargs) - return LossOutput(loss=loss, num_tokens=0) + return LossOutput(loss=loss, num_tokens=self._loss_num_tokens(loss_mask)) class GSPOLoss(GRPOLoss): @@ -410,16 +421,43 @@ class BNPOLoss(GRPOLoss): BNPO (Batch-Normalized Policy Optimization) Loss. Normalizes by total completion tokens across batch. + + ``token_mean_scope``: + 'global' (default, correct): return the UN-normalized token sum and report + ``num_tokens=Σmask``, so the framework's SUM-loss path divides the accumulated + gradient by the GLOBAL token count => exact token-mean, invariant to how the + batch is split into micro/dp groups (matches verl/SEAM BNPO). + 'micro' (legacy): per-(micro÷dp)-group token-mean, equal-weighted across groups + (``num_tokens=0`` => PER-TOKEN-MEAN accumulation). This is the pre-fix behavior; + it double-averages (group token-mean, then equal-weight over groups), biasing + toward short responses and degrading to sequence-mean as groups multiply. Keep it + ONLY to reproduce arms trained before the 2026-08-01 fix (skill2lora E1–E20). """ + def __init__(self, *args, token_mean_scope: str = 'global', **kwargs): + super().__init__(*args, **kwargs) + assert token_mean_scope in ('global', 'micro'), \ + f'token_mean_scope must be global|micro, got {token_mean_scope!r}' + self.token_mean_scope = token_mean_scope + def _aggregate_loss( self, per_token_loss: 'torch.Tensor', loss_mask: 'torch.Tensor', **kwargs, ) -> 'torch.Tensor': - """Sum over all tokens, divide by total token count.""" - return (per_token_loss * loss_mask).sum() / loss_mask.sum().clamp(min=1.0) + """global: return the token SUM (the global division is done downstream via + num_tokens=Σmask). micro (legacy): local token-mean, later equal-weighted across + micro/dp groups.""" + summed = (per_token_loss * loss_mask).sum() + if self.token_mean_scope == 'global': + return summed + return summed / loss_mask.sum().clamp(min=1.0) + + def _loss_num_tokens(self, loss_mask: 'torch.Tensor'): + if self.token_mean_scope == 'global': + return loss_mask.sum().clamp(min=1.0) + return 0 class SEAMBNPOLoss(BNPOLoss): @@ -518,7 +556,7 @@ def __call__( per_token_loss = per_token_loss - self.entropy_coef * entropies.to(per_token_loss.dtype) loss = self._aggregate_loss(per_token_loss, loss_mask, **kwargs) - return LossOutput(loss=loss, num_tokens=0) + return LossOutput(loss=loss, num_tokens=self._loss_num_tokens(loss_mask)) class DRGRPOLoss(GRPOLoss): diff --git a/tests/loss/test_bnpo_token_mean.py b/tests/loss/test_bnpo_token_mean.py new file mode 100644 index 000000000..9e9ddd6ff --- /dev/null +++ b/tests/loss/test_bnpo_token_mean.py @@ -0,0 +1,82 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""BNPO token-mean fix: 'global' scope must be invariant to how a batch is split into +micro/dp groups (strict global token-mean); 'micro' scope reproduces the pre-fix +double-average bias that skews toward short responses. + +The framework combines groups per LossOutput semantics (transformers.py / metric/loss.py): + effective_loss = Σ_g loss_g / Σ_g num_tokens_g +so for scope='global' (loss_g = token sum, num_tokens_g = Σmask) this collapses to the +single-shot token-mean for ANY partition; for scope='micro' (loss_g = token-mean, +num_tokens_g = 0 -> treated as 1) it becomes the equal-weighted mean of per-group means. +""" +import torch + +from twinkle.loss.grpo import BNPOLoss + + +def _combine(loss_fn, ptl, mask, groups): + """Mimic the framework accumulation over `groups` (lists of row indices).""" + tot_loss = 0.0 + tot_tok = 0.0 + for idx in groups: + g_ptl, g_mask = ptl[idx], mask[idx] + loss_g = loss_fn._aggregate_loss(g_ptl, g_mask) + ntok = loss_fn._loss_num_tokens(g_mask) + ntok = float(ntok if not torch.is_tensor(ntok) else ntok.item()) + if ntok <= 0: # micro path: num_tokens=0 -> framework uses 1 per group + ntok = 1.0 + tot_loss = tot_loss + loss_g + tot_tok += ntok + return float(tot_loss) / tot_tok + + +def _fixture(): + # two very different response lengths (short=2 tok, long=6 tok) -> maximally exposes bias + ptl = torch.tensor([ + [1.0, 1.0, 0.0, 0.0, 0.0, 0.0], # short: mean per-token loss 1.0 over 2 tokens + [0.5, 0.5, 0.5, 0.5, 0.5, 0.5], # long : mean per-token loss 0.5 over 6 tokens + ]) + mask = torch.tensor([ + [1., 1., 0., 0., 0., 0.], + [1., 1., 1., 1., 1., 1.], + ]) + return ptl, mask + + +def test_global_is_split_invariant(): + ptl, mask = _fixture() + loss = BNPOLoss(token_mean_scope='global') + whole = _combine(loss, ptl, mask, [[0, 1]]) + split = _combine(loss, ptl, mask, [[0], [1]]) + true_token_mean = float((ptl * mask).sum() / mask.sum()) # (2*1 + 6*0.5)/8 = 0.625 + assert abs(whole - true_token_mean) < 1e-6 + assert abs(split - true_token_mean) < 1e-6 # <-- the fix: split == whole + assert abs(whole - split) < 1e-6 + + +def test_micro_is_biased_and_split_dependent(): + ptl, mask = _fixture() + loss = BNPOLoss(token_mean_scope='micro') + whole = _combine(loss, ptl, mask, [[0, 1]]) # one group -> token-mean 0.625 + split = _combine(loss, ptl, mask, [[0], [1]]) # per-group means (1.0, 0.5) -> 0.75 + assert abs(whole - 0.625) < 1e-6 + assert abs(split - 0.75) < 1e-6 # short response over-weighted + assert split > whole # bias toward short is real + # and it disagrees with the correct global answer + assert abs(split - 0.625) > 1e-3 + + +def test_seam_inherits_global_by_default(): + from twinkle.loss.grpo import SEAMBNPOLoss + seam = SEAMBNPOLoss(epsilon=0.2, beta=0.001) + assert seam.token_mean_scope == 'global' + ptl, mask = _fixture() + split = _combine(seam, ptl, mask, [[0], [1]]) + assert abs(split - 0.625) < 1e-6 + + +if __name__ == '__main__': + test_global_is_split_invariant() + test_micro_is_biased_and_split_dependent() + test_seam_inherits_global_by_default() + print('OK: global split-invariant; micro reproduces the biased double-average') From 1eadc442069b98668f3ac66929a554cb200ed4f1 Mon Sep 17 00:00:00 2001 From: root Date: Sat, 1 Aug 2026 17:04:43 +0800 Subject: [PATCH 30/60] run_ablate12: guard $TRAIN_MICRO_BATCH for set -u (unbound variable) E13/E21 blocks referenced $TRAIN_MICRO_BATCH directly; under set -u an unset env aborts with 'unbound variable'. Use ${TRAIN_MICRO_BATCH:-}. --- cookbook/exp/skill2lora/run_ablate12.sh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cookbook/exp/skill2lora/run_ablate12.sh b/cookbook/exp/skill2lora/run_ablate12.sh index 4016c970e..0c9db64c6 100644 --- a/cookbook/exp/skill2lora/run_ablate12.sh +++ b/cookbook/exp/skill2lora/run_ablate12.sh @@ -279,7 +279,7 @@ while IFS=$'\t' read -r NAME EXP_DIR THINK SMT OPTIONAL TASK EXEC_THINK; do # 切 micro 数学等价，不改有效批量(chunk=128=SEAM batch)与可比性，只降训练前向峰值显存。 # 显式 TRAIN_MICRO_BATCH 优先；E13_TRAIN_MICRO_BATCH 单独可调。值恒为 TRAIN_GPUS 倍数(满足 %TRAIN_DP)。 _tmb=${E13_TRAIN_MICRO_BATCH:-$((2*${TRAIN_GPUS:-2}))} - [ -n "$TRAIN_MICRO_BATCH" ] && _tmb=$TRAIN_MICRO_BATCH + [ -n "${TRAIN_MICRO_BATCH:-}" ] && _tmb=$TRAIN_MICRO_BATCH E16_FLAGS="$E16_FLAGS --train-micro-batch $_tmb" echo "[ablate12] E13 SEAM-repro: chunk=$CHUNK_ARG min_level=$MIN_LEVEL_ARG train_micro_batch=$_tmb"\ " reward_trunc_penalty=0 eval=R1/T0 executor=nothink skill_max_tokens=$SMT (executor 预算 8192/16384, K=n_skills 默认 8)" @@ -290,7 +290,7 @@ while IFS=$'\t' read -r NAME EXP_DIR THINK SMT OPTIONAL TASK EXEC_THINK; do # 切 micro 数学等价。若仍 OOM：改用 TRAIN_GPUS=1（micro→1）或 TRAIN_FSDP=2 分片权重腾显存。 # 显式 TRAIN_MICRO_BATCH 优先；E21_TRAIN_MICRO_BATCH 单独可调。 _tmb=${E21_TRAIN_MICRO_BATCH:-${TRAIN_GPUS:-2}} - [ -n "$TRAIN_MICRO_BATCH" ] && _tmb=$TRAIN_MICRO_BATCH + [ -n "${TRAIN_MICRO_BATCH:-}" ] && _tmb=$TRAIN_MICRO_BATCH E16_FLAGS="$E16_FLAGS --train-micro-batch $_tmb" echo "[ablate12] E21 freeform: train_micro_batch=$_tmb (80G OOM guard; global token-mean 使切 micro 等价，不改批量)" fi From 6de2a971207e7cf737d366fcdccb78570908bc76 Mon Sep 17 00:00:00 2001 From: root Date: Sun, 2 Aug 2026 00:58:05 +0800 Subject: [PATCH 31/60] fix --- cookbook/exp/skill2lora/analyze_more.py | 172 +++++++ cookbook/exp/skill2lora/logp_corr_probe.py | 429 ++++++++++++++++++ cookbook/exp/skill2lora/rubric_effect.py | 138 ++++++ cookbook/exp/skill2lora/run_ablate12.sh | 33 +- cookbook/exp/skill2lora/skill_feature_corr.py | 128 ++++++ cookbook/exp/skill2lora/train_skill_v2.py | 123 +++-- cookbook/exp/skill2lora/watchdog_e14.sh | 40 ++ src/twinkle/loss/grpo.py | 38 +- .../model/transformers/transformers.py | 5 + 9 files changed, 1062 insertions(+), 44 deletions(-) create mode 100644 cookbook/exp/skill2lora/analyze_more.py create mode 100644 cookbook/exp/skill2lora/logp_corr_probe.py create mode 100644 cookbook/exp/skill2lora/rubric_effect.py create mode 100644 cookbook/exp/skill2lora/skill_feature_corr.py create mode 100644 cookbook/exp/skill2lora/watchdog_e14.sh diff --git a/cookbook/exp/skill2lora/analyze_more.py b/cookbook/exp/skill2lora/analyze_more.py new file mode 100644 index 000000000..1543139dc --- /dev/null +++ b/cookbook/exp/skill2lora/analyze_more.py @@ -0,0 +1,172 @@ +#!/usr/bin/env python3 +"""analyze_more.py — logp_corr 数据的补充相关性分析（7 项此前未做的）。纯 CPU。 +用法：/usr/local/bin/python3 analyze_more.py +""" +import json +import math +import os +from collections import defaultdict + +import numpy as np + +D = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'logp_corr') +pairs = [json.loads(l) for l in open(os.path.join(D, 'pairs.jsonl'))] +problems = json.load(open(os.path.join(D, 'problems.json'))) +rolls = [json.loads(l) for l in open(os.path.join(D, 'rollout_results.jsonl'))] +npz = np.load(os.path.join(D, 'token_logps.npz')) + +pas = {r['key']: float(np.mean(r['pass'])) for r in rolls if r['kind'] == 'skill' and r['mode'].startswith('t05')} +tok = {r['key']: float(np.mean(r['tokens'])) for r in rolls if r['kind'] == 'skill' and r['mode'].startswith('t05')} +npass = {r['key']: int(np.sum(r['pass'])) for r in rolls if r['kind'] == 'skill' and r['mode'].startswith('t05')} +gr = {r['key']: float(r['pass'][0]) for r in rolls if r['kind'] == 'skill' and r['mode'] == 'greedy' and r['n']} +bpas = {r['key']: float(np.mean(r['pass'])) for r in rolls if r['kind'] == 'base' and r['mode'].startswith('t05')} +btok = {r['key']: float(np.mean(r['tokens'])) for r in rolls if r['kind'] == 'base' and r['mode'].startswith('t05')} +prob_by = {p['data_id']: p for p in problems} + + +def rank(x): + x = np.asarray(x, float) + o = np.argsort(x, kind='mergesort') + r = np.empty(len(x)) + r[o] = np.arange(len(x)) + for v in np.unique(x): + m = x == v + if m.sum() > 1: + r[m] = r[m].mean() + return r + + +def sp(a, b): + a, b = np.asarray(a, float), np.asarray(b, float) + m = ~(np.isnan(a) | np.isnan(b)) + if m.sum() < 3: + return np.nan + ra, rb = rank(a[m]), rank(b[m]) + if ra.std() == 0 or rb.std() == 0: + return np.nan + return float(((ra - ra.mean()) * (rb - rb.mean())).mean() / (ra.std() * rb.std())) + + +# ============ ① 题目侧特征 -> skill 收益（筛题实证） ============ +print('=' * 70) +print('① 题目侧特征 -> skill 平均收益（n=61 题；spearman 跨题）') +by_p = defaultdict(list) +for pr in pairs: + if pr['pair_id'] in pas: + by_p[pr['data_id']].append(pr) +prows = [] +for did, prs in by_p.items(): + p = prob_by[did] + lv = int(did.split(':')[1]) + lifts = [pas[x['pair_id']] - bpas[did] for x in prs] + passes = [pas[x['pair_id']] for x in prs] + prows.append({'level': lv, 'base_pass': bpas[did], 'base_tok': btok[did], + 'prob_chars': len(p['problem']), 'gt_chars': len(p['gt']), + 'mean_lift': float(np.mean(lifts)), 'grp_std': float(np.std(passes)), + 'frac_helped': float(np.mean([l > 0 for l in lifts]))}) +for fk in ['level', 'base_pass', 'base_tok', 'prob_chars', 'gt_chars']: + v = [r[fk] for r in prows] + print('%-12s vs mean_lift %+0.3f | vs 组可分性(grp_std) %+0.3f | vs frac_helped %+0.3f' % ( + fk, sp(v, [r['mean_lift'] for r in prows]), sp(v, [r['grp_std'] for r in prows]), + sp(v, [r['frac_helped'] for r in prows]))) +bt = np.array([r['base_tok'] for r in prows]) +ml = np.array([r['mean_lift'] for r in prows]) +for lo, hi in [(0, 3000), (3000, 5000), (5000, 9999)]: + m = (bt >= lo) & (bt < hi) + if m.sum(): + print(' base_tok[%d,%d): n=%d mean_lift=%+.3f frac(lift>0)=%.2f' % ( + lo, hi, m.sum(), ml[m].mean(), np.mean([r['frac_helped'] for r, mm in zip(prows, m) if mm]))) + +# ============ ② skill 生成 think 长度 -> 好坏 ============ +print('\n' + '=' * 70) +print('② skillgen_tokens（skill 生成总 token 含 think）组内 vs pass/输出长') +cs1, cs2 = [], [] +for did, prs in by_p.items(): + if len(prs) < 4: + continue + sg = [x.get('skillgen_tokens') or np.nan for x in prs] + c = sp(sg, [pas[x['pair_id']] for x in prs]) + if not np.isnan(c): + cs1.append(c) + c = sp(sg, [tok[x['pair_id']] for x in prs]) + if not np.isnan(c): + cs2.append(c) +print(' vs pass8: mean=%+.3f se=%.3f n=%d' % (np.mean(cs1), np.std(cs1) / np.sqrt(len(cs1)), len(cs1))) +print(' vs exec_tokens: mean=%+.3f se=%.3f n=%d' % (np.mean(cs2), np.std(cs2) / np.sqrt(len(cs2)), len(cs2))) + +# ============ ③ |delta| 当干预强度计：|delta| vs |lift| ============ +print('\n' + '=' * 70) +print('③ |ΔlogP| 是否预测"干预幅度"|lift|（不看方向）') +ad, al, dtk = [], [], [] +for pr in pairs: + k = pr['pair_id'] + if k not in pas or pr.get('logp_delta_train') is None: + continue + ad.append(abs(pr['logp_delta_train'])) + al.append(abs(pas[k] - bpas[pr['data_id']])) + dtk.append(abs(tok[k] - btok[pr['data_id']])) +print(' |delta| vs |lift| 全局 sp=%+.3f (n=%d)' % (sp(ad, al), len(ad))) +print(' |delta| vs |Δexec_tokens| 全局 sp=%+.3f' % sp(ad, dtk)) +cs = [] +for did, prs in by_p.items(): + if len(prs) < 4: + continue + a = [abs(x['logp_delta_train']) for x in prs] + b = [abs(pas[x['pair_id']] - bpas[did]) for x in prs] + c = sp(a, b) + if not np.isnan(c): + cs.append(c) +print(' 组内: mean=%+.3f se=%.3f n=%d' % (np.mean(cs), np.std(cs) / np.sqrt(len(cs)), len(cs))) + +# ============ ④ delta 的位置衰减：skill 影响是否集中在 GT 前段 ============ +print('\n' + '=' * 70) +print('④ per-token delta 的位置分布（四分位段的 mean|delta|，跨 476 对平均）') +qsum = np.zeros(4) +qcnt = 0 +for pr in pairs: + k, did = pr['pair_id'], pr['data_id'] + if f'base|{did}' not in npz.files or f'skill|{k}' not in npz.files: + continue + b, s = npz[f'base|{did}'], npz[f'skill|{k}'] + if len(b) != len(s) or len(b) < 40: + continue + d = np.abs(s - b) + d = d[~np.isnan(d)] + if len(d) < 40: + continue + qs = np.array_split(d, 4) + qsum += np.array([q.mean() for q in qs]) + qcnt += 1 +print(' Q1(前1/4)=%.4f Q2=%.4f Q3=%.4f Q4(末1/4)=%.4f (n=%d)' % (*(qsum / qcnt), qcnt)) + +# ============ ⑤ pass8 分布形态：混沌（U形/过散）还是二项噪声 ============ +print('\n' + '=' * 70) +print('⑤ 混合组 (00.5）案例 vs 其余：输出长度') +dis, rest = [], [] +for k in pas: + if k in gr: + (dis if abs(gr[k] - pas[k]) > 0.5 else rest).append(tok[k]) +print(' 分歧组 n=%d mean_tok=%d p75=%d | 其余 n=%d mean_tok=%d p75=%d' % ( + len(dis), np.mean(dis), np.percentile(dis, 75), len(rest), np.mean(rest), np.percentile(rest, 75))) + +# ============ ⑦ leak clean 分解：leak 到底贡献多少 lift ============ +print('\n' + '=' * 70) +print('⑦ leak 分解（lift = pass8 - base_pass8）') +for name, cond in [('leaked', lambda p: p['leaked']), ('clean', lambda p: not p['leaked'])]: + ls = [pas[p['pair_id']] - bpas[p['data_id']] for p in pairs if p['pair_id'] in pas and cond(p)] + print(' %-7s n=%-4d mean_lift=%+.4f frac(lift>0)=%.2f' % ( + name, len(ls), np.mean(ls), np.mean([x > 0 for x in ls]))) diff --git a/cookbook/exp/skill2lora/logp_corr_probe.py b/cookbook/exp/skill2lora/logp_corr_probe.py new file mode 100644 index 000000000..42ce06690 --- /dev/null +++ b/cookbook/exp/skill2lora/logp_corr_probe.py @@ -0,0 +1,429 @@ +#!/usr/bin/env python3 +"""logp_corr_probe.py — 找"低噪声、强预测 skill 好坏"的指标（E15 数据驱动 reward 选型）。 + +问题：E15 用 mean ΔlogP(GT|题+skill) 当稠密 reward，20 步 delta 不爬。要回答两件事： + 1) ΔlogP（及各种 per-token 聚合变体）到底和 skill 的真实有效性（executor 多 rollout + 通过率）相关吗？相关性多强？ + 2) 各候选指标的噪声多大？（greedy×1 判分 vs 8-rollout 真值 的一致性 = 老 0/1 reward 的噪声） + +数据：E15 gen_records（题/skill/GT 参考解/训练期 fp32 mean logps 都在盘上）。 +三阶段（分开进程跑，互不污染 Ray/vllm）： + --phase rollout twinkle vLLMSampler dp=8：每对 (题,skill) T=0.5×8 rollout + greedy×1， + 外加每题 baseline（无 skill）同口径 → 真值 pass_rate / lift。 + --phase logps 原生 vllm prompt_logprobs=0 + twinkle Template.encode 的 labels 定位 + response 段（与训练 _score_executor_mean_logps 同一模板/同一切位）， + 对 base(题+GT) 与 skill(题+skill+GT) 各算一遍逐 token logp → npz。 + --phase analyze CPU：各指标 vs 真值的 Spearman/AUC（全局 + 组内），噪声对比表。 + +用法（8 卡空闲时）： + cd cookbook/exp/skill2lora + PYTHONPATH=../../src python3 logp_corr_probe.py --phase rollout + PYTHONPATH=../../src python3 logp_corr_probe.py --phase logps + PYTHONPATH=../../src python3 logp_corr_probe.py --phase analyze +""" +import argparse +import copy +import json +import os +import re +import sys +from collections import defaultdict +from typing import Dict, List, Optional + +import numpy as np + +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +E15_DIR = os.path.join(SCRIPT_DIR, 'output.ablate12', 'E15_logp_gt_on_narrative') +OUT_DIR = os.path.join(SCRIPT_DIR, 'logp_corr') +MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') +MAX_MODEL_LEN = int(os.environ.get('MAX_MODEL_LEN', 16384)) +MAX_TOKENS = 8192 # executor 解题预算，对齐 v2 +N_PROBLEMS = int(os.environ.get('PROBE_PROBLEMS', 64)) +N_ROLLOUTS = int(os.environ.get('PROBE_ROLLOUTS', 8)) +SEED = 42 + +# ---- executor prompt / 判分：逐字复刻 train_skill_v2 v2 分支（与 eval_skill_probe 相同） ---- +_ANSWER_FORMAT_V2 = ('Present your reasoning, then put ONLY the final numeric result inside ' + '\\boxed{}. For example: \\boxed{42}.') + + +def build_direct_prompt(problem): + content = f'The problem you need to solve:\n{problem}\n\n' + _ANSWER_FORMAT_V2 + return {'messages': [{'role': 'user', 'content': content}]} + + +def build_skill_solve_prompt(problem, skill): + skill = (skill or '').strip() + if not skill: + return build_direct_prompt(problem) + content = (f'The problem you need to solve:\n{problem}\n\n' + 'Skill hint:\nFor this problem, a skill-generation model has analyzed it and ' + 'provided some advisory skills:\n' + f'{skill}\n' + 'Prefer using its techniques when they fit, but if you have a more efficient or ' + 'clearer correct method, you may use it. If you diverge from this advice, briefly ' + 'explain why. Be concise and accurate.\n' + + _ANSWER_FORMAT_V2) + return {'messages': [{'role': 'user', 'content': content}]} + + +_BOXED_RE = re.compile(r'\\boxed\s*\{') + + +def extract_boxed(text): + if not text: + return None + last = None + for m in _BOXED_RE.finditer(text): + depth, i = 1, m.end() + while i < len(text) and depth > 0: + depth += (text[i] == '{') - (text[i] == '}') + i += 1 + if depth == 0: + last = text[m.end():i - 1].strip() + return last + + +_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') +_SEAM_TAG_RE = re.compile(r'<\s*answer\s*>(.*?)<\s*/\s*answer\s*>', re.I | re.S) +_SEAM_BOX_RE = re.compile(r'(?:\\{1,2}\(|)\\{1,2}boxed\s*\{\s*([^}]*)\s*}(?:\)|)', re.S) +_SEAM_INLINE_RE = re.compile(r'\$([^$]+)\$|\\\(([^)]+)\\\)', re.S) +_SEAM_FRAC_RE = re.compile(r'(-?\d+(?:\.\d+)?)/(-?\d+(?:\.\d+)?)') +_SEAM_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') + + +def _seam_norm(num): + try: + f = float(num) + return str(int(f)) if f == int(f) else str(f) + except Exception: + return num.strip() + + +def _seam_sanitize(txt): + txt = (txt or '').strip() + if (m := _SEAM_TAG_RE.search(txt)): + txt = m.group(1).strip() + elif (m := _SEAM_BOX_RE.search(txt)): + txt = m.group(1).strip() + elif (m := _SEAM_INLINE_RE.search(txt)): + txt = (m.group(1) or m.group(2)).strip() + # bugfix 2026-07-29：\dfrac/\tfrac/\cfrac 归一，同 train_skill_v2._seam_sanitize + txt = txt.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac').replace(r'\cfrac', r'\frac') + txt = re.sub(r'\\frac\s*\{\s*([^}]+?)\s*}\s*\{\s*([^}]+?)\s*}', r'\1/\2', txt) + if (m := _SEAM_FRAC_RE.search(txt)): + p, q = map(float, m.groups()) + if q: + return _seam_norm(str(p / q)) + if (m := _SEAM_NUM_RE.search(txt)): + return _seam_norm(m.group()) + return txt + + +def _judge(decoded, gold): + text = _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() + raw = extract_boxed(text) + pred = _seam_sanitize(raw) if raw else None + return bool(pred) and (pred == _seam_sanitize(str(gold))) + + +# ---- 配对采样：E15 gen_records -> pairs.jsonl ----------------------------------------- +def load_pairs(): + """64 题（seeded）× 组内全部 parseable 且有 delta 的候选（<=8）；GT 取每题首候选 rolls[0].text。""" + by_id = {} + for line in open(os.path.join(E15_DIR, 'gen_records.jsonl')): + r = json.loads(line) + if r.get('record_type') != 'problem' or not r.get('candidates'): + continue + by_id.setdefault(r['data_id'], r) # data_id 在 epoch 内唯一 + ids = sorted(by_id) + rng = np.random.RandomState(SEED) + pick = list(rng.permutation(len(ids))[:N_PROBLEMS]) + pairs, problems = [], [] + for k in pick: + r = by_id[ids[k]] + gt = next((c['rolls'][0]['text'] for c in r['candidates'] + if c.get('rolls') and c['rolls'][0].get('text')), '') + if not gt: + continue + cands = [c for c in r['candidates'] if c['parseable'] and c.get('logp_delta') is not None] + if len(cands) < 4: + continue + problems.append({'data_id': r['data_id'], 'problem': r['problem'], + 'reference_answer': r['reference_answer'], 'gt': gt}) + for j, c in enumerate(cands[:8]): + pairs.append({'pair_id': f'{r["data_id"]}#{j}', 'data_id': r['data_id'], + 'skill': c['skills'], 'leaked': bool(c['leaked']), + 'skill_chars': len(c['skills']), + 'skillgen_tokens': c.get('skillgen_tokens'), + 'logp_base_train': c['logp_base'], 'logp_skill_train': c['logp_skill'], + 'logp_delta_train': c['logp_delta']}) + return problems, pairs + + +# ---- phase: rollout ------------------------------------------------------------------- +def phase_rollout(): + import twinkle + from twinkle import DeviceGroup, DeviceMesh + from twinkle.data_format import SamplingParams + from twinkle.sampler import vLLMSampler + from twinkle.template import Template + + problems, pairs = load_pairs() + os.makedirs(OUT_DIR, exist_ok=True) + json.dump(problems, open(os.path.join(OUT_DIR, 'problems.json'), 'w')) + with open(os.path.join(OUT_DIR, 'pairs.jsonl'), 'w') as f: + for p in pairs: + f.write(json.dumps(p, ensure_ascii=False) + '\n') + print(f'[rollout] problems={len(problems)} pairs={len(pairs)}', flush=True) + + n_gpu = int(os.environ.get('EXEC_GPUS', 8)) + twinkle.initialize(mode='ray', nproc_per_node=n_gpu, lazy_collect=False, + groups=[DeviceGroup(name='exec', ranks=list(range(n_gpu)), device_type='GPU')]) + sampler = vLLMSampler(model_id=MODEL_ID, + engine_args={'gpu_memory_utilization': 0.85, + 'max_model_len': MAX_MODEL_LEN, 'tensor_parallel_size': 1}, + device_mesh=DeviceMesh.from_sizes(world_size=n_gpu, dp_size=n_gpu), + remote_group='exec') + sampler.set_template(Template, model_id=MODEL_ID, enable_thinking=True, max_length=MAX_MODEL_LEN) + + prob_by = {p['data_id']: p for p in problems} + prompts, metas = [], [] + for p in problems: # baseline（无 skill） + prompts.append(build_direct_prompt(p['problem'])) + metas.append(('base', p['data_id'])) + for pr in pairs: # with-skill + prompts.append(build_skill_solve_prompt(prob_by[pr['data_id']]['problem'], pr['skill'])) + metas.append(('skill', pr['pair_id'])) + + def run(params, tag): + padded = prompts if len(prompts) % n_gpu == 0 else \ + prompts + [copy.deepcopy(prompts[-1])] * (n_gpu - len(prompts) % n_gpu) + outs = sampler.sample(padded, params)[:len(prompts)] + rows = [] + for (kind, key), resp in zip(metas, outs): + gold = prob_by[key.split('#')[0]]['reference_answer'] if '#' in key \ + else prob_by[key]['reference_answer'] + seqs = list(resp.sequences) if (resp and resp.sequences) else [] + rows.append({'kind': kind, 'key': key, 'mode': tag, + 'n': len(seqs), + 'pass': [bool(_judge(getattr(s, 'decoded', '') or '', gold)) for s in seqs], + 'trunc': [getattr(s, 'stop_reason', None) == 'length' for s in seqs], + 'tokens': [len(getattr(s, 'tokens', None) or []) for s in seqs]}) + return rows + + rows = run(SamplingParams(max_tokens=MAX_TOKENS, temperature=0.5, top_p=1.0, + num_samples=N_ROLLOUTS), f't05x{N_ROLLOUTS}') + rows += run(SamplingParams(max_tokens=MAX_TOKENS, temperature=0.0, top_p=1.0, + num_samples=1), 'greedy') + with open(os.path.join(OUT_DIR, 'rollout_results.jsonl'), 'w') as f: + for r in rows: + f.write(json.dumps(r) + '\n') + print(f'[rollout] done: {len(rows)} rows -> rollout_results.jsonl', flush=True) + + +# ---- phase: logps --------------------------------------------------------------------- +def phase_logps(): + """原生 vllm prompt_logprobs=0；token 布局与训练一致：twinkle Template.encode 的 labels + != -100 即 response(GT) 段位置。base 与 skill 两条轨迹各存一行 float32。""" + from twinkle.template import Template + from vllm import LLM, SamplingParams as VSP + from vllm.inputs import TokensPrompt + + problems = json.load(open(os.path.join(OUT_DIR, 'problems.json'))) + pairs = [json.loads(l) for l in open(os.path.join(OUT_DIR, 'pairs.jsonl'))] + prob_by = {p['data_id']: p for p in problems} + tmpl = Template(model_id=MODEL_ID, enable_thinking=True, max_length=MAX_MODEL_LEN, + truncation_strategy='delete') + + def encode(problem, skill, gt): + msgs = [dict(m) for m in build_skill_solve_prompt(problem, skill)['messages']] + enc = tmpl.encode({'messages': msgs + [{'role': 'assistant', 'content': gt}], + 'user_data': {'key_rounds': [len(msgs)]}}) + if enc is None: + return None, None + ids = [int(x) for x in enc['input_ids']] # numpy int64 -> int（vllm msgspec 拒收 np 类型） + pos = np.where(np.asarray(enc['labels']) != -100)[0] + return ids, pos + + jobs, keys = [], [] # key: ('base', data_id) / ('skill', pair_id) + for p in problems: + ids, pos = encode(p['problem'], '', p['gt']) + if ids is not None and len(pos): + jobs.append((ids, pos)) + keys.append(('base', p['data_id'])) + for pr in pairs: + p = prob_by[pr['data_id']] + ids, pos = encode(p['problem'], pr['skill'], p['gt']) + if ids is not None and len(pos): + jobs.append((ids, pos)) + keys.append(('skill', pr['pair_id'])) + print(f'[logps] encoded jobs={len(jobs)} (skipped {len(problems)+len(pairs)-len(jobs)})', flush=True) + + llm = LLM(model=_local_model_path(), max_model_len=MAX_MODEL_LEN, + gpu_memory_utilization=0.85, tensor_parallel_size=1) + sp = VSP(max_tokens=1, temperature=0.0, prompt_logprobs=0) + outs = llm.generate([TokensPrompt(prompt_token_ids=ids) for ids, _ in jobs], sp) + + store = {} + for (ids, pos), (kind, key), out in zip(jobs, keys, outs): + plp = out.prompt_logprobs + row = np.full(len(pos), np.nan, dtype=np.float32) + for i, p_ in enumerate(pos): + d = plp[int(p_)] if int(p_) < len(plp) else None + if d: + lp = d.get(ids[int(p_)]) + if lp is not None: + row[i] = lp.logprob + store[f'{kind}|{key}'] = row + np.savez_compressed(os.path.join(OUT_DIR, 'token_logps.npz'), **store) + print(f'[logps] saved {len(store)} rows -> token_logps.npz', flush=True) + + +def _local_model_path(): + from modelscope.hub.snapshot_download import snapshot_download + return snapshot_download(MODEL_ID, local_files_only=True) + + +# ---- phase: analyze ------------------------------------------------------------------- +def _rank(x): + x = np.asarray(x, dtype=np.float64) + order = np.argsort(x, kind='mergesort') + r = np.empty(len(x)) + r[order] = np.arange(len(x)) + for v in np.unique(x): # 平均并列名次 + m = x == v + if m.sum() > 1: + r[m] = r[m].mean() + return r + + +def spearman(a, b): + a, b = np.asarray(a, float), np.asarray(b, float) + m = ~(np.isnan(a) | np.isnan(b)) + if m.sum() < 3: + return np.nan + ra, rb = _rank(a[m]), _rank(b[m]) + sa, sb = ra.std(), rb.std() + return np.nan if sa == 0 or sb == 0 else float(((ra - ra.mean()) * (rb - rb.mean())).mean() / (sa * sb)) + + +def auc(score, label): + score, label = np.asarray(score, float), np.asarray(label, bool) + m = ~np.isnan(score) + score, label = score[m], label[m] + if label.sum() == 0 or (~label).sum() == 0: + return np.nan + r = _rank(score) + return float((r[label].sum() - label.sum() * (label.sum() - 1) / 2) / (label.sum() * (~label).sum())) + + +def phase_analyze(): + problems = json.load(open(os.path.join(OUT_DIR, 'problems.json'))) + pairs = [json.loads(l) for l in open(os.path.join(OUT_DIR, 'pairs.jsonl'))] + rolls = [json.loads(l) for l in open(os.path.join(OUT_DIR, 'rollout_results.jsonl'))] + npz = np.load(os.path.join(OUT_DIR, 'token_logps.npz')) + + base_pass, pair_pass, pair_greedy, base_greedy, pair_trunc = {}, {}, {}, {}, {} + for r in rolls: + rate = float(np.mean(r['pass'])) if r['n'] else np.nan + if r['kind'] == 'base' and r['mode'].startswith('t05'): + base_pass[r['key']] = rate + elif r['kind'] == 'skill' and r['mode'].startswith('t05'): + pair_pass[r['key']] = rate + pair_trunc[r['key']] = float(np.mean(r['trunc'])) if r['n'] else np.nan + elif r['kind'] == 'skill' and r['mode'] == 'greedy': + pair_greedy[r['key']] = float(r['pass'][0]) if r['n'] else np.nan + elif r['kind'] == 'base' and r['mode'] == 'greedy': + base_greedy[r['key']] = float(r['pass'][0]) if r['n'] else np.nan + + rows = [] + for pr in pairs: + key, did = pr['pair_id'], pr['data_id'] + b = npz[f'base|{did}'] if f'base|{did}' in npz.files else None + s = npz[f'skill|{key}'] if f'skill|{key}' in npz.files else None + if key not in pair_pass or did not in base_pass: + continue + row = {'pair_id': key, 'data_id': did, + 'truth_pass8': pair_pass[key], 'truth_lift': pair_pass[key] - base_pass[did], + 'base_pass8': base_pass[did], 'greedy1': pair_greedy.get(key, np.nan), + 'trunc_rate': pair_trunc.get(key, np.nan), + 'delta_train': pr['logp_delta_train'], + 'leaked': float(pr['leaked']), 'skill_chars': float(pr['skill_chars'])} + if b is not None and s is not None and len(b) == len(s): + d = s - b + ok = ~(np.isnan(d)) + d, bb = d[ok], b[ok] + if len(d): + row['delta_mean'] = float(d.mean()) + row['delta_sum'] = float(d.sum()) + k = min(50, len(d)) + row['delta_top50'] = float(d[np.argsort(-np.abs(d))[:k]].mean()) + unc = bb < -1.0 # executor 本来拿不准的 token + row['delta_uncertain'] = float(d[unc].mean()) if unc.sum() >= 5 else np.nan + row['delta_tail100'] = float(d[-min(100, len(d)):].mean()) + row['frac_improved'] = float((d > 0).mean()) + row['base_mean_ck'] = float(bb.mean()) + rows.append(row) + print(f'[analyze] usable pairs={len(rows)}') + json.dump(rows, open(os.path.join(OUT_DIR, 'pair_table.json'), 'w')) + + # 交叉校验：vllm 重算 delta vs 训练 fp32 delta + dm = [r.get('delta_mean', np.nan) for r in rows] + dt = [r['delta_train'] for r in rows] + print(f'\n[校验] corr(delta_vllm, delta_train) spearman={spearman(dm, dt):.3f}') + + metrics = ['delta_train', 'delta_mean', 'delta_sum', 'delta_top50', 'delta_uncertain', + 'delta_tail100', 'frac_improved', 'greedy1', 'skill_chars', 'leaked', 'trunc_rate'] + truth = np.array([r['truth_pass8'] for r in rows]) + lift = np.array([r['truth_lift'] for r in rows]) + helped = lift > 0 + + print('\n=== 全局相关性（n=%d 对）：指标 vs 8-rollout 真值 ===' % len(rows)) + print('%-16s %-14s %-14s %-10s' % ('metric', 'sp(pass8)', 'sp(lift)', 'AUC(lift>0)')) + for m in metrics: + v = np.array([r.get(m, np.nan) for r in rows], float) + print('%-16s %-14s %-14s %-10s' % ( + m, f'{spearman(v, truth):+.3f}', f'{spearman(v, lift):+.3f}', f'{auc(v, helped):.3f}')) + + # 组内（GRPO 真正用的信号）：每题 >=4 候选的组内 spearman 均值 + print('\n=== 组内相关性（每题组内 spearman 的均值±se）===') + by_p = defaultdict(list) + for r in rows: + by_p[r['data_id']].append(r) + for m in metrics: + cs = [] + for did, rs in by_p.items(): + if len(rs) < 4: + continue + v = [r.get(m, np.nan) for r in rs] + t = [r['truth_pass8'] for r in rs] + c = spearman(v, t) + if not np.isnan(c): + cs.append(c) + if cs: + cs = np.array(cs) + print('%-16s mean=%+.3f se=%.3f n_groups=%d' % (m, cs.mean(), cs.std() / np.sqrt(len(cs)), len(cs))) + + # 噪声对比：greedy×1 vs 真值；基线 greedy vs 基线8rollout + g = np.array([r.get('greedy1', np.nan) for r in rows], float) + m = ~np.isnan(g) + hard_wrong = np.abs(g[m] - truth[m]) > 0.5 + print(f'\n[噪声] greedy×1 与 8-rollout 真值强不一致率（|diff|>0.5）: {hard_wrong.mean():.3f} (n={m.sum()})') + bg = np.array([base_greedy.get(p["data_id"], np.nan) for p in problems], float) + bp = np.array([base_pass.get(p["data_id"], np.nan) for p in problems], float) + mm = ~(np.isnan(bg) | np.isnan(bp)) + print(f'[噪声] baseline greedy 与 baseline pass8 强不一致率: {(np.abs(bg[mm]-bp[mm])>0.5).mean():.3f} (n={mm.sum()})') + print(f'[分布] truth_pass8 mean={np.nanmean(truth):.3f} lift>0 比例={np.mean(helped):.3f} ' + f'lift<0 比例={np.mean(lift<0):.3f}') + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument('--phase', choices=('rollout', 'logps', 'analyze'), required=True) + args = ap.parse_args() + {'rollout': phase_rollout, 'logps': phase_logps, 'analyze': phase_analyze}[args.phase]() + + +if __name__ == '__main__': + main() diff --git a/cookbook/exp/skill2lora/rubric_effect.py b/cookbook/exp/skill2lora/rubric_effect.py new file mode 100644 index 000000000..9829dee2c --- /dev/null +++ b/cookbook/exp/skill2lora/rubric_effect.py @@ -0,0 +1,138 @@ +#!/usr/bin/env python3 +"""rubric_effect.py — rubric 作用的三路数据分析（纯 CPU，现有数据）。 +① 臂级：E5/E6/E7(rl_ab, rubric 条件) vs E1/E2/E3(bnpo, query-only) eval lift 对照 +② 题级：同一错题上，rubric 条件生成的 skill vs query-only 生成的 skill 的 executor 通过率 +③ rubric 文本特征 vs A 线拯救率（组内 any-pass） +""" +import json +import hashlib +import os +import re +from collections import defaultdict + +import numpy as np + +BASE = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'output.ablate12') + +# gen_records 未存 rubric 文本；用全局缓存反查（key = md5('rubric_global\x1f'+data_id)）。 +# 能查到 => 该题被诊断过 => 裸题答错、属 A 线（rl_ab 只诊断错题）。 +_RUBRIC = {} +with open(os.path.join(BASE, 'rubric_cache_global.jsonl')) as f: + for l in f: + d = json.loads(l) + _RUBRIC[d['key']] = d.get('value') or '' + + +def rubric_of(data_id): + k = hashlib.md5(('\x1f'.join(['rubric_global', str(data_id)])).encode('utf-8')).hexdigest() + return _RUBRIC.get(k) + + +def load_gen(exp): + rows = [] + with open(os.path.join(BASE, exp, 'gen_records.jsonl')) as f: + for l in f: + r = json.loads(l) + rows.append(r) + return rows + + +def probe_schema(exp): + rows = load_gen(exp) + tps = defaultdict(int) + for r in rows: + tps[r.get('record_type')] += 1 + print(exp, dict(tps)) + for r in rows: + if r.get('record_type') == 'problem': + print(' problem keys:', sorted(r.keys())[:30]) + cands = r.get('cands') or r.get('_cands') or [] + if cands: + print(' cand keys:', sorted(cands[0].keys())) + break + + +if __name__ == '__main__': + import sys + if len(sys.argv) > 1 and sys.argv[1] == 'schema': + probe_schema('E7_rl_ab_on_pitfall') + probe_schema('E3_bnpo_on_pitfall') + sys.exit(0) + + # ---------- ② 题级同题对照 ---------- + # E7 A 线（rubric 非空）的候选 vs E3 同 data_id 的候选（query-only），配对比较组均值 + for pair in [('E7_rl_ab_on_pitfall', 'E3_bnpo_on_pitfall'), + ('E5_rl_ab_off_pitfall', 'E1_bnpo_off_pitfall'), + ('E6_rl_ab_off_narrative', 'E2_bnpo_off_narrative')]: + ea, eb = pair + ga, gb = load_gen(ea), load_gen(eb) + + def group_pass(rows, need_rubric=None): + out = {} + for r in rows: + if r.get('record_type') != 'problem': + continue + did = r.get('data_id', '') + rub = rubric_of(did) + if need_rubric is True and not rub: + continue + cands = [c for c in (r.get('candidates') or []) if c.get('parseable') + and c.get('with_pass') is not None] + if not cands: + continue + # 同一题可能多 chunk 出现，取第一次（早期，policy 漂移最小） + if did not in out: + out[did] = (np.mean([float(c['with_pass']) for c in cands]), rub or '') + return out + + pa = group_pass(ga, need_rubric=True) # A 线（rubric 条件） + pb = group_pass(gb) # query-only + common = sorted(set(pa) & set(pb)) + if not common: + print(f'[②] {ea} vs {eb}: 无同题交集') + continue + da = np.array([pa[d][0] for d in common]) + db = np.array([pb[d][0] for d in common]) + diff = da - db + print(f'[②] {ea.split("_")[0]}(rubric) vs {eb.split("_")[0]}(query-only) 同题 n={len(common)}: ' + f'rubric臂组均pass={da.mean():.3f} qonly臂={db.mean():.3f} ' + f'配对差={diff.mean():+.4f}±{diff.std()/np.sqrt(len(diff)):.4f} ' + f'win/tie/lose={int((diff>0).sum())}/{int((diff==0).sum())}/{int((diff<0).sum())}') + + # ---------- ③ rubric 特征 vs 拯救率 ---------- + ga = load_gen('E7_rl_ab_on_pitfall') + rows = [] + for r in ga: + if r.get('record_type') != 'problem': + continue + rub = (rubric_of(r.get('data_id', '')) or '').strip() + if not rub: + continue + cands = [c for c in (r.get('candidates') or []) if c.get('parseable') + and c.get('with_pass') is not None] + if not cands: + continue + n_fail = len(re.findall(r'\[FAIL\]', rub)) + n_pass = len(re.findall(r'\[PASS\]', rub)) + rows.append({'len': len(rub), 'n_fail': n_fail, 'n_pass': n_pass, + 'n_crit': n_fail + n_pass, + 'has_fix': int('fix:' in rub), + 'rescue': float(np.mean([float(c['with_pass']) for c in cands])), + 'any': float(any(c['with_pass'] for c in cands))}) + if rows: + def sp(a, b): + a, b = np.asarray(a, float), np.asarray(b, float) + ra = np.argsort(np.argsort(a)).astype(float) + rb = np.argsort(np.argsort(b)).astype(float) + if ra.std() == 0 or rb.std() == 0: + return np.nan + return float(((ra - ra.mean()) * (rb - rb.mean())).mean() / (ra.std() * rb.std())) + print(f'\n[③] E7 A线 rubric 特征 vs 拯救率 (n={len(rows)} 题, ' + f'mean rescue={np.mean([r["rescue"] for r in rows]):.3f}, ' + f'any-pass={np.mean([r["any"] for r in rows]):.3f})') + for fk in ['len', 'n_fail', 'n_pass', 'n_crit', 'has_fix']: + v = [r[fk] for r in rows] + print(' %-8s vs rescue %+0.3f | vs any-pass %+0.3f' % ( + fk, sp(v, [r['rescue'] for r in rows]), sp(v, [r['any'] for r in rows]))) + ls = np.array([r['len'] for r in rows]) + print(' rubric len 分布: p25=%d p50=%d p75=%d' % tuple(np.percentile(ls, [25, 50, 75]))) diff --git a/cookbook/exp/skill2lora/run_ablate12.sh b/cookbook/exp/skill2lora/run_ablate12.sh index 0c9db64c6..a19563051 100644 --- a/cookbook/exp/skill2lora/run_ablate12.sh +++ b/cookbook/exp/skill2lora/run_ablate12.sh @@ -274,25 +274,42 @@ while IFS=$'\t' read -r NAME EXP_DIR THINK SMT OPTIONAL TASK EXEC_THINK; do MIN_LEVEL_ARG="${E13_MIN_LEVEL:-0}" [ -z "$REWARD_TRUNC_PENALTY" ] && E16_FLAGS="$E16_FLAGS --reward-trunc-penalty 0" E16_FLAGS="$E16_FLAGS --eval-rollouts 1 --eval-skill-temperature 0.0" - # 显存：E13 8×140G 实测 micro=8（trainer.py 的 8192 自动档）仍 OOM（132G 已用 + 4.77G）。 - # 降 train_micro_batch 到 2×dp（默认 TRAIN_GPUS=2 → 4）；grpo.py 的 global token-mean 修复后 - # 切 micro 数学等价，不改有效批量(chunk=128=SEAM batch)与可比性，只降训练前向峰值显存。 - # 显式 TRAIN_MICRO_BATCH 优先；E13_TRAIN_MICRO_BATCH 单独可调。值恒为 TRAIN_GPUS 倍数(满足 %TRAIN_DP)。 + # ---- 与 SEAM 逐行对齐的优化器参数（2026-08-01）---------------------------------- + # 上一版 E13 只对了 chunk/min_level/惩罚/eval，三个真正控制更新幅度的参数全部跑默认值， + # 与 SEAM 差得很远（实测后果：think 25 步从 3977 塔到 1942 token、reward 0.816→0.734）。 + # SEAM 侧真值来自 scripts/train_deepmath_paper.sh + verl 的 fsdp_workers.py:198-199 归一化： + # ppo_mini_batch_size=20 × rollout.n=8 = 160 全局序列（再 /4gpu = 40/gpu） + # → 每 batch 的 optimizer step 数 = 256/40 = 7（verl data.split 的余数也算一步）， + # twinkle 端 range(0,1024,160) 同样 7 步；且 mini 1: + r[m] = r[m].mean() + return r + + +def sp(a, b): + a, b = np.asarray(a, float), np.asarray(b, float) + m = ~(np.isnan(a) | np.isnan(b)) + if m.sum() < 3: + return np.nan + ra, rb = rank(a[m]), rank(b[m]) + if ra.std() == 0 or rb.std() == 0: + return np.nan + return float(((ra - ra.mean()) * (rb - rb.mean())).mean() / (ra.std() * rb.std())) + + +rows = [] +for pr in pairs: + k = pr['pair_id'] + if k not in pas: + continue + f = feats(pr['skill']) + f.update({'pair_id': k, 'data_id': pr['data_id'], 'pass': pas[k], + 'trunc': trc.get(k, np.nan), 'tok': tok.get(k, np.nan), + 'leaked': float(pr['leaked'])}) + rows.append(f) + +FKEYS = ['chars', 'words', 'sent_len', 'num_density', 'formula_density', 'latex_density', + 'step_markers', 'imper_density', 'hedge_density', 'pitfall_density', 'uniq_ratio'] + +print(f'[skill_feat] n_pairs={len(rows)}') +# 全局 +print('\n=== 全局 spearman：skill 特征 vs 真值 ===') +print('%-16s %-10s %-10s %-10s' % ('feature', 'sp(pass)', 'sp(trunc)', 'sp(tok)')) +for fk in FKEYS: + v = [r[fk] for r in rows] + print('%-16s %+.3f %+.3f %+.3f' % ( + fk, sp(v, [r['pass'] for r in rows]), + sp(v, [r['trunc'] for r in rows]), sp(v, [r['tok'] for r in rows]))) + +# 组内 +by = defaultdict(list) +for r in rows: + by[r['data_id']].append(r) +groups = [g for g in by.values() if len(g) >= 4] +print(f'\n=== 组内 spearman（每题内，n_groups={len(groups)}，mean±se）vs pass_rate ===') +for fk in FKEYS: + cs = [] + for g in groups: + c = sp([r[fk] for r in g], [r['pass'] for r in g]) + if not np.isnan(c): + cs.append(c) + if cs: + cs = np.array(cs) + print('%-16s mean=%+.3f se=%.3f n=%d' % (fk, cs.mean(), cs.std() / np.sqrt(len(cs)), len(cs))) + +print(f'\n=== 组内 spearman vs trunc_rate（截断通道）===') +for fk in FKEYS: + cs = [] + for g in groups: + c = sp([r[fk] for r in g], [r['trunc'] for r in g]) + if not np.isnan(c): + cs.append(c) + if cs: + cs = np.array(cs) + print('%-16s mean=%+.3f se=%.3f n=%d' % (fk, cs.mean(), cs.std() / np.sqrt(len(cs)), len(cs))) diff --git a/cookbook/exp/skill2lora/train_skill_v2.py b/cookbook/exp/skill2lora/train_skill_v2.py index 78d951904..d5af4f03c 100644 --- a/cookbook/exp/skill2lora/train_skill_v2.py +++ b/cookbook/exp/skill2lora/train_skill_v2.py @@ -311,8 +311,12 @@ def _seam_norm(num: str) -> str: return num.strip() -def _seam_sanitize(txt: str) -> str: - """Port of SEAM lpem.sanitize_math_answer + normalize_number_format.""" +def _seam_sanitize(txt: str, dfrac_fix: bool = True) -> str: + """Port of SEAM lpem.sanitize_math_answer + normalize_number_format. + + ``dfrac_fix=False`` gives BIT parity with the upstream function (which only rewrites the + literal ``\\frac``); it is used by the ``align='seam'`` judge so that E13's acc is + reproducible against SEAM's step_summary. See _parse_seq.""" txt = (txt or '').strip() if (m := _SEAM_TAG_RE.search(txt)): txt = m.group(1).strip() @@ -323,7 +327,9 @@ def _seam_sanitize(txt: str) -> str: # bugfix 2026-07-29：\dfrac/\tfrac/\cfrac 不被下行字面 \frac 正则匹配，曾致 # \boxed{\dfrac{1}{2}} 落到 _SEAM_NUM_RE 抓首个数字 → pred='1'（分数答案题全判错； # 实测被标"错"的 boxed rolls 中 70-85% 实为正确，见 skill_quality_analysis.md 末章）。 - txt = txt.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac').replace(r'\cfrac', r'\frac') + # 注：SEAM 上游没有这一步，seam 对齐口径下必须关掉（dfrac_fix=False）。 + if dfrac_fix: + txt = txt.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac').replace(r'\cfrac', r'\frac') txt = re.sub(r'\\frac\s*\{\s*([^}]+?)\s*}\s*\{\s*([^}]+?)\s*}', r'\1/\2', txt) if (m := _SEAM_FRAC_RE.search(txt)): p, q = map(float, m.groups()) @@ -346,12 +352,16 @@ def _clean_text(decoded: Optional[str]) -> str: def _extract_skill(text: str) -> Optional[str]: """Parse the skill block: in seam mode (SEAM format_pass parity), else .""" + if _ALIGN_MODE == 'seam': + # ⭐ 整段搜、取首个匹配，不先剔掉 —— 与 SEAM 两处实现逐字一致： + # lpem.format_pass：MEMORY_RE.search(resp)；fsdp_workers.py:836：_re.search(..., response_text)。 + # 旧实现只搜 之后，会把"只在 think 里写了 memory_item"的候选当成格式失败 + # （reward=0），而 SEAM 那边算格式通过 —— 直接影响 format 率与组内 reward 方差。 + m = re.search(r'(.*?)', text or '', re.DOTALL | re.IGNORECASE) + return (m.group(1).strip() or None) if m else None low = text.lower() end_think = low.rfind('') answer = text[end_think + len(''):] if end_think >= 0 else text - if _ALIGN_MODE == 'seam': - m = re.search(r'(.*?)', answer, re.DOTALL | re.IGNORECASE) - return (m.group(1).strip() or None) if m else None open_tag, close_tag = '', '' s = answer.lower().rfind(open_tag) if s < 0: @@ -369,9 +379,22 @@ def _parse_seq(seq, gold) -> Dict[str, Any]: if _TASK == 'code': return _parse_many([(seq, gold)])[0] text = _clean_text(getattr(seq, 'decoded', '') or '') - # 判分口径统一（人工拍板，2026-07-27）：seam/v2 都只从 \boxed{} 抽取，再走同一套数值归一 - # （frac/inline/number）后精确匹配；不做 lpem 式“整段抓数字”贪婪回退，保证 E13 与 E1-E12 - # 的 acc/lift 横向可比。extract_boxed 取最后一个配平的 \boxed{}、截断时不误取；没有则判错。 + if _ALIGN_MODE == 'seam': + # ⭐ seam 对齐口径（2026-08-02 修正）= SEAM_JUDGE=answer，即 train_deepmath_paper.sh 的默认值： + # 整段文本走 sanitize 级联 → \boxed → $..$/\(..\) → 分数 → **首个数字**。 + # 之前这里走的是 boxed-only（等价 SEAM_JUDGE=boxed），与 executor prompt 要求 + # "......" 直接冲突 —— 换成原版 prompt 后 boxed-only + # 会把几乎所有 rollout 判错。prompt 与判分必须成对切换，见 _SEAM_SOLVE_ADVISORY 注释。 + # dfrac_fix=False 是为了与 lpem.sanitize_math_answer 逐字节一致。 + pred = _seam_sanitize(text, dfrac_fix=False) or None + correct = bool(pred) and (pred == _seam_sanitize(str(gold), dfrac_fix=False)) + terminated = getattr(seq, 'stop_reason', None) != 'length' + return {'pred': pred, 'correct': correct, 'terminated': terminated, + 'stop_reason': getattr(seq, 'stop_reason', None), + 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} + # v2（E1-E12/E14-E21）：只从 \boxed{} 抽取，再走同一套数值归一（frac/inline/number）后精确匹配； + # 不做 lpem 式"整段抓数字"贪婪回退，保证这些臂之间的 acc/lift 横向可比。extract_boxed 取最后一个 + # 配平的 \boxed{}、截断时不误取；没有则判错。 raw = extract_boxed(text) pred = _seam_sanitize(raw) if raw else None correct = bool(pred) and (pred == _seam_sanitize(str(gold))) @@ -884,16 +907,26 @@ def _diagnose_entry(checker, entry: Dict[str, Any]) -> Optional[str]: '- Output ONLY the experience, wrapped EXACTLY as ' ' ... .\n\n' 'Problem:\n{problem}') +# 逐字节复刻 SEAM templates/slove_qwen.txt（注意第一行末尾那个空格，EOF 无换行）。 _SEAM_SOLVE_ADVISORY = ( - 'The above is a Q&A dialogue between a user and a problem-solving guidance model.\n' + 'The above is a Q&A dialogue between a user and a problem-solving guidance model. \n' 'Treat the output of the guidance model as advisory context to solve the math problem: ' 'prefer using its techniques when they fit, but you may use alternative correct methods ' 'if they are more efficient or clearer. If you diverge from the advisory context, briefly ' 'explain why. Be concise and accurate.\n' - + _ANSWER_FORMAT_V2) + + _ANSWER_FORMAT) +# seam 模式的 baseline/空-skill 回退 system。必须用 _ANSWER_FORMAT（/）而不是 +# _ANSWER_FORMAT_V2（\boxed{}）：SEAM 的 fsdp_workers.py:850-856 在 SEAM_JUDGE=answer（默认）下 +# 就是这个串，且 executor 关 thinking 时 Qwen3 会自动补一个空 ，与"请输出 " +# 的要求撞在一起，抽答案更容易失败 —— 这正是 SEAM baseline 只有 0.57 的原因。用 boxed 会把 +# baseline 抬到 0.72、给定可解析 skill 的 acc 抬到 0.923（SEAM 0.865），曲线水平就对不上。 +# 判分侧无需改动：_seam_sanitize 优先匹配 、其次 boxed，两种格式都吃。 +DIRECT_SYSTEM_SEAM = ( + 'You are an expert competition mathematician. Be concise and accurate. ' + + _ANSWER_FORMAT) -def build_skill_solve_prompt_seam(problem, skill, raw_response=None): +def build_skill_solve_prompt_seam(problem, skill, raw_response=None, resp_terminated=True): """SEAM executor prompt. Non-empty skills use the actor's raw response_text, preserving actor exactly as SEAM's reward worker does: prompt_text + response_text + grm. If raw_response is missing, fall back to reconstructing a minimal response.""" @@ -904,6 +937,12 @@ def build_skill_solve_prompt_seam(problem, skill, raw_response=None): + '<|im_end|>\n<|im_start|>assistant\n') if not response_text: response_text = f'{skill}' + elif resp_terminated: + # SEAM 的 response_text 是 skip_special_tokens=False 解码的（fsdp_workers.py:828），正常终止的 + # rollout 末尾带着 EOS，即 "<|im_end|>"；twinkle 的 _clean_text 把 <|...|> 全剔了。 + # 差这一个 token 也会改变 executor 的贪心解码，补回来。截断的 rollout（stop=length） + # SEAM 那边也没有 EOS，所以不补。 + response_text = response_text + '<|im_end|>' content = prompt_text + response_text + '\n' + _SEAM_SOLVE_ADVISORY return {'messages': [{'role': 'user', 'content': content}]} @@ -912,15 +951,15 @@ def build_direct_prompt(problem): if _TASK == 'code': return code_task.direct_prompt(problem) if _ALIGN_MODE == 'seam': - # seam 基线保持英文原样，不受 v2 prompt 改动影响 - return {'messages': [{'role': 'system', 'content': DIRECT_SYSTEM}, + # seam 基线逐字复刻 SEAM fsdp_workers.py:850-858（system + user，/ 格式） + return {'messages': [{'role': 'system', 'content': DIRECT_SYSTEM_SEAM}, {'role': 'user', 'content': problem}]} # v2：英文 executor 基线——与带 skill 版同格式，仅去掉“技巧提示”部分 content = f'The problem you need to solve:\n{problem}\n\n' + _ANSWER_FORMAT_V2 return {'messages': [{'role': 'user', 'content': content}]} -def build_skill_solve_prompt(problem, skill, raw_response=None): +def build_skill_solve_prompt(problem, skill, raw_response=None, resp_terminated=True): skill = (skill or '').strip() if _TASK == 'code': # 空 skill -> 干净 direct（与数学分支同规则，见下方注释） @@ -929,10 +968,12 @@ def build_skill_solve_prompt(problem, skill, raw_response=None): # 空 skill → 干净 direct。训练侧根本不会用空 skill 走 executor（process_chunk 只对非空 flat 跑， # 空候选直接 reward=0），故此分支仅影响 eval 口径——让空 skill 题 withskill==baseline、对 lift 贡献 0， # 去掉空壳嵌套的框架水分，指标更干净。 + # （seam 模式下这正好等于 SEAM fsdp_workers.py:846-858 的 else 分支。） return build_direct_prompt(problem) if _ALIGN_MODE == 'seam': # 非空 skill 走 SEAM 原始 reward worker 路径：executor 可见 actor 完整 response_text（含 ）。 - return build_skill_solve_prompt_seam(problem, skill, raw_response=raw_response) + return build_skill_solve_prompt_seam(problem, skill, raw_response=raw_response, + resp_terminated=resp_terminated) # v2：英文 executor——题目 + 技巧提示(skill 作为 advisory) + 答案格式，单 user turn content = (f'The problem you need to solve:\n{problem}\n\n' 'Skill hint:\nFor this problem, a skill-generation model has analyzed it and ' @@ -1359,23 +1400,34 @@ def process_chunk(base_sampler, skill_sampler, chunk, ci, base_dp, skill_dp, arg for r, c in flat: c['leaked'] = _answer_leaked(c['skills'], r['reference_answer']) + # ⭐ executor 覆盖面（seam 对齐，2026-08-02）：SEAM 对**每一条**候选都跑 executor —— + # fsdp_workers.py:842-858，抽不到 时走 else 分支（direct_system + 裸题目）， + # 它的 acc 照样计入 reward_extra_info["acc"]，也就是 train/with_skill_accuracy 的分子。 + # 旧实现只跑 parseable 的，于是 twinkle 算不出同口径的 withskill 准确率，只能拿 + # P(correct|parseable) 去对 SEAM 的 P(correct)，两条曲线分母不同、根本无法对齐。 + # reward 口径不变：_skill_reward 仍然乘 parseable，空 skill 永远 reward=0。仅 seam 模式开， + # 其余臂（E1-E12/E14-E21）行为与开销不动。 + exec_list = ([(r, c) for r in chunk for c in r['_cands']] if _ALIGN_MODE == 'seam' else flat) + # with-skill executor pass：默认 greedy×1（E1-E13，reward 0/1 与旧口径 bit 一致）； # E14: reward_rollouts>1 时 T=reward_temperature × K 采样，reward = parseable × 通过率， # 把内容信号从 greedy 0/1 量化里释放出来（提升组内 std>0 比例）。 - if flat: + if exec_list: K = max(1, int(getattr(args, 'reward_rollouts', 1) or 1)) rT = float(getattr(args, 'reward_temperature', 0.0) or 0.0) ws_out = _run_samples(base_sampler, - [build_skill_solve_prompt(r['problem'], c['skills'], c.get('response')) for r, c in flat], + [build_skill_solve_prompt(r['problem'], c['skills'], c.get('response'), + resp_terminated=(c.get('skillgen_stop') != 'length')) + for r, c in exec_list], K, args.max_tokens, base_dp, temperature=rT) # 判分一次性批量化（code 任务要起子进程跑单测，逐条会比 GPU 还慢一个量级） pairs, spans = [], [] - for (r, c), seqs in zip(flat, ws_out): + for (r, c), seqs in zip(exec_list, ws_out): start = len(pairs) pairs.extend((s, r['reference_answer']) for s in (seqs or [])) spans.append((start, len(pairs))) judged = _parse_many(pairs) - for (r, c), (a, b) in zip(flat, spans): + for (r, c), (a, b) in zip(exec_list, spans): rolls = judged[a:b] or [_empty_roll()] for x in rolls[1:]: x['text'] = '' # 磁盘保护：K>1 时只留首 rollout 全文（gen_records 体积控制） @@ -1465,8 +1517,16 @@ def _chunk_summary(chunk, ci): trunc = sum(1 for x in ws_rolls if x['stop_reason'] == 'length') # bugfix（ablate #6）：旧版 any(c.get('reward')) 用 truthiness，负 reward（E16 hinge/leak_gate、 # E14 地板 -1.0）也被当“通过”；改用 with_pass>0（真正的 executor 通过率），greedy 0/1 臂语义不变。 - ws_acc = _mean([1.0 if any((c.get('with_pass') or 0) > 0 for c in r['_cands']) else 0.0 + # 同时限 parseable：seam 模式下 unparseable 候选也有 with_pass（走 direct 回退），不限的话 + # 这条 pass@K 会被 baseline 成绩推高、与历史臂不可比。 + ws_acc = _mean([1.0 if any((c.get('with_pass') or 0) > 0 and c.get('parseable') for c in r['_cands']) + else 0.0 for r in chunk if r['_cands']]) + # ⭐ SEAM 同口径的 withskill 准确率：**全部**候选上的 mean(correct)，不看格式、不条件化。 + # = ray_trainer.py:1529 float(np.mean(reward_extra_infos_dict["acc"])) + # = step_summary 的 withskill_pass = swanlab 的 train/with_skill_accuracy。 + # 只有 seam 模式会给 unparseable 候选跑 executor，其余臂这个值等于 candidate_withskill_pass。 + pass_all = [c['with_pass'] for c in all_cands if c['with_pass'] is not None] return { 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), 'n_generated': len(all_cands), 'n_candidates_parseable': len(cands), @@ -1481,6 +1541,8 @@ def _chunk_summary(chunk, ci): 'skill_chars_mean': _mean([len(c['skills']) for c in cands]), 'avg_withskill_pass': ws_acc, 'candidate_withskill_pass': _mean([c['with_pass'] for c in scored]), + 'withskill_pass_all_cands': _mean(pass_all), + 'n_exec_cands': len(pass_all), 'withskill_trunc_frac': (trunc / len(ws_rolls)) if ws_rolls else 0.0, 'termination_rate_withskill': _mean([1.0 if x['terminated'] else 0.0 for x in ws_rolls]), } @@ -1516,16 +1578,17 @@ def run_greedy_eval(base_sampler, skill_sampler, eval_records, ci, rounds, for j in range(R): s = seqs[j] if j < len(seqs) else None if s is None: - row.append(('', '')) + row.append(('', '', 'stop')) else: sresp = _clean_text(getattr(s, 'decoded', '') or '') - row.append((_extract_skill(sresp) or '', sresp)) + row.append((_extract_skill(sresp) or '', sresp, getattr(s, 'stop_reason', None))) per_skills.append(row) # flatten R×N for a single batched greedy executor pass flat_prompts, flat_idx = [], [] for pi, (r, row) in enumerate(zip(eval_records, per_skills)): - for j, (sk, sresp) in enumerate(row): - flat_prompts.append(build_skill_solve_prompt(r['problem'], sk, sresp)) + for j, (sk, sresp, sstop) in enumerate(row): + flat_prompts.append(build_skill_solve_prompt(r['problem'], sk, sresp, + resp_terminated=(sstop != 'length'))) flat_idx.append((pi, j)) ws_out = _run_samples(base_sampler, flat_prompts, 1, args.max_tokens, base_dp, temperature=0.0) judged = _parse_many([(_first_seq(seqs), eval_records[pi]['reference_answer']) @@ -1535,7 +1598,7 @@ def run_greedy_eval(base_sampler, skill_sampler, eval_records, ci, rounds, for pi, (r, row) in enumerate(zip(eval_records, per_skills)): rolls = [roll_by[(pi, j)] for j in range(len(row))] corr = [1.0 if x['correct'] else 0.0 for x in rolls] - parses = [1.0 if sk else 0.0 for sk, _ in row] + parses = [1.0 if sk else 0.0 for sk, _sresp, _sstop in row] terms = [1.0 if x['terminated'] else 0.0 for x in rolls] acc_mean = sum(corr) / len(corr) if corr else 0.0 # bugfix（ablate #8）：unparseable skill 的 rollout 实际走了 direct 回退（≈baseline）， @@ -1794,6 +1857,14 @@ def _swan_metrics(summary, log): if summary['n_groups'] > 0: d.update({'acc/withskill_pass': summary['avg_withskill_pass'], 'term/withskill_trunc_frac': summary['withskill_trunc_frac']}) + # ⭐ 与 SEAM 同名同口径的三条（ray_trainer.py:1569-1583），专为 swanlab 叠图对齐而发： + # train/with_skill_accuracy = 全部候选的 mean(correct)；acc/reward_mean = mean(correct∧format)； + # skill/format_rate = SEAM 的 format_mean。 + # 注：旧的 acc/withskill_pass 是**题级 pass@K**，与 SEAM 同名指标不同口径（它接近 1 + # 且会随训练小幅下行）—— 之前把这两条叠在一张图上看"趋势相反"就是这个原因。 + d['train/with_skill_accuracy'] = summary['withskill_pass_all_cands'] + d['acc/reward_mean'] = summary['reward_mean'] + d['skill/format_rate'] = summary['parse_rate'] if log: d['train/n_grpo'] = log['n_grpo'] d['train/n_sft'] = log['n_sft'] diff --git a/cookbook/exp/skill2lora/watchdog_e14.sh b/cookbook/exp/skill2lora/watchdog_e14.sh new file mode 100644 index 000000000..b9789958f --- /dev/null +++ b/cookbook/exp/skill2lora/watchdog_e14.sh @@ -0,0 +1,40 @@ +#!/bin/bash +# One-shot watchdog (2026-07-28): the running ablate12 launcher holds the OLD plan +# (E7 -> E8). E14 (reward SNR ablation) was inserted after E7 in config.py, so when E7 +# finishes (DONE.json) OR the launcher dies (E7 crash), swap to a fresh launcher that +# reads the new plan: completed arms skip via DONE.json, so it starts E14 directly. +set -u +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +DONE="$HERE/output.ablate12/E7_rl_ab_on_pitfall/DONE.json" +PIDF=/tmp/ablate12.pid +LOG="$HERE/watchdog_e14.log" + +echo "[watchdog $(date +%H:%M:%S)] waiting for E7 DONE.json or launcher exit" >> "$LOG" +while true; do + [ -f "$DONE" ] && { echo "[watchdog $(date +%H:%M:%S)] E7 DONE.json found" >> "$LOG"; break; } + OLD=$(cat "$PIDF" 2>/dev/null || echo "") + if [ -n "$OLD" ] && ! kill -0 "$OLD" 2>/dev/null; then + echo "[watchdog $(date +%H:%M:%S)] launcher $OLD died without E7 DONE (crash?); restarting anyway" >> "$LOG" + break + fi + sleep 60 +done + +sleep 10 +OLD=$(cat "$PIDF" 2>/dev/null || echo "") +[ -n "$OLD" ] && kill "$OLD" 2>/dev/null && echo "[watchdog] killed old launcher $OLD" >> "$LOG" +# kill any experiment python the old launcher may have just started (E8 race window) +pkill -f "skill_ablate.main" 2>/dev/null && echo "[watchdog] killed stray skill_ablate.main" >> "$LOG" + +# wait for the 8 GPUs to drain (engine teardown), max 15 min +for i in $(seq 1 90); do + USED=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | awk '{s+=$1} END {print s}') + [ "${USED:-1}" -lt 1000 ] && break + sleep 10 +done +echo "[watchdog $(date +%H:%M:%S)] GPUs drained (used=${USED:-?}MiB); relaunching" >> "$LOG" + +cd "$HERE" +nohup bash run_ablate12.sh > run_ablate12.nohup.log 2>&1 & +echo $! > "$PIDF" +echo "[watchdog $(date +%H:%M:%S)] new launcher pid=$(cat $PIDF) (plan includes E14 after E7)" >> "$LOG" diff --git a/src/twinkle/loss/grpo.py b/src/twinkle/loss/grpo.py index 12d0416fe..66e4ffcc1 100644 --- a/src/twinkle/loss/grpo.py +++ b/src/twinkle/loss/grpo.py @@ -423,22 +423,40 @@ class BNPOLoss(GRPOLoss): Normalizes by total completion tokens across batch. ``token_mean_scope``: - 'global' (default, correct): return the UN-normalized token sum and report - ``num_tokens=Σmask``, so the framework's SUM-loss path divides the accumulated - gradient by the GLOBAL token count => exact token-mean, invariant to how the - batch is split into micro/dp groups (matches verl/SEAM BNPO). - 'micro' (legacy): per-(micro÷dp)-group token-mean, equal-weighted across groups - (``num_tokens=0`` => PER-TOKEN-MEAN accumulation). This is the pre-fix behavior; - it double-averages (group token-mean, then equal-weight over groups), biasing - toward short responses and degrading to sequence-mean as groups multiply. Keep it - ONLY to reproduce arms trained before the 2026-08-01 fix (skill2lora E1–E20). + 'micro' (default, matches verl/SEAM): per-(micro÷dp)-group token-mean, + equal-weighted across groups (``num_tokens=0`` => PER-TOKEN-MEAN accumulation). + This is what verl actually does -- see verl/workers/actor/dp_actor.py: pg_loss = + agg_loss(..., 'token-mean') is ``masked_mean`` computed WITHIN each micro-batch, + then ``loss = policy_loss * (1/gradient_accumulation)`` before ``backward()``. + So verl's effective gradient is the equal-weighted mean of per-micro token-means, + NOT a global token-mean. + 'global': return the UN-normalized token sum and report ``num_tokens=Σmask``, so the + framework's SUM-loss path divides the accumulated gradient by the GLOBAL token + count => strict token-mean, invariant to micro/dp splitting. + + Why 'global' is NOT the default, despite being the "textbook" token-mean + (measured on skill2lora E13, 2026-08-01): + Group-relative advantages cancel exactly per group (mean A = 0), but the + TOKEN-weighted mean does not: it equals -cov(len, A)/mean(len). With + corr(len, A) = -0.42 (long skill-gen responses hit the 8192 budget, lose their + closing tag, and score 0), 'global' yields a per-token pg_loss of +0.031 versus + verl/SEAM's +3.2e-4 -- a ~100x coherent "emit fewer tokens" gradient. Under + 'global', E13 collapsed its from 3977 to 1942 tokens in 25 updates + (SEAM: -17% in 76 updates) and overshot the optimum: corr(len, correct) flipped + from -0.42 to +0.23 while reward fell 0.816 -> 0.734. 'micro' localizes the + normalization, so the length coupling largely cancels (it degenerates to + sequence-mean as the micro size approaches 1). """ - def __init__(self, *args, token_mean_scope: str = 'global', **kwargs): + def __init__(self, *args, token_mean_scope: str = 'micro', **kwargs): super().__init__(*args, **kwargs) assert token_mean_scope in ('global', 'micro'), \ f'token_mean_scope must be global|micro, got {token_mean_scope!r}' self.token_mean_scope = token_mean_scope + # 'global' 返回的是 token 和（梯度在下游按 num_tokens=Σmask 归一）。必须同步告诉展示层 + # 这是 sum-reduction，否则 LossMetric（metric/loss.py）不会除以 num_tokens，会把每个 micro 的 + # token 和当均值直接平均，展示出一个被 token 数放大的巨大 loss（梯度不受影响，纯展示失真）。 + self.reduction = 'sum' if token_mean_scope == 'global' else 'mean' def _aggregate_loss( self, diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index ee2ffe774..9d0203c1f 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -182,6 +182,11 @@ def __init__( memory_efficient_init: bool = False, **kwargs): os.environ['TOKENIZERS_PARALLELISM'] = 'true' + # Opt-out of the cuDNN SDPA backend (falls back to flash/mem-efficient, numerically + # equivalent): sporadic `mha_graph.execute` RuntimeError on Blackwell + CUDA 13 + # (ablate12 E7 crashed at update 21 mid-forward). Env-gated to keep default behavior. + if os.environ.get('TWINKLE_DISABLE_CUDNN_SDP', '0') == '1': + torch.backends.cuda.enable_cudnn_sdp(False) self._try_init_process_group() super(PreTrainedModel, self).__init__() # The Default tokenizer will be used to save with a model if no template was set. From 3e23a7c454f3a95786d8431dfbf49ebacea55e42 Mon Sep 17 00:00:00 2001 From: root Date: Sun, 2 Aug 2026 00:59:33 +0800 Subject: [PATCH 32/60] wip --- cookbook/exp/skill2lora/skill_ablate/methods.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/cookbook/exp/skill2lora/skill_ablate/methods.py b/cookbook/exp/skill2lora/skill_ablate/methods.py index ebe8cd468..ea323d3b0 100644 --- a/cookbook/exp/skill2lora/skill_ablate/methods.py +++ b/cookbook/exp/skill2lora/skill_ablate/methods.py @@ -468,7 +468,18 @@ def step(self, chunk, ci): tmetrics = _train_metrics(log.get('metric')) tmetrics['train/n_samples'] = float(log.get('n_grpo', 0) + log.get('n_sft', 0)) metrics = {'signal/zero_grad_frac': summary['zero_grad_frac'], + 'signal/group_reward_std_mean': summary['group_reward_std_mean'], + 'signal/n_train_samples': float(summary['n_train_samples']), + 'signal/n_groups': float(summary['n_groups']), + # 旧名（题级 pass@K，历史面板兼容） 'acc/withskill_pass': summary['avg_withskill_pass'], + # ⭐ 与 SEAM ray_trainer.py:1569-1583 同名同口径，用于 swanlab 直接叠图对齐： + # train/with_skill_accuracy = 全部候选的 mean(correct)（SEAM withskill_pass） + # acc/reward_mean = mean(correct∧format)（SEAM reward_mean） + # skill/format_rate = SEAM format_mean + 'train/with_skill_accuracy': summary['withskill_pass_all_cands'], + 'acc/reward_mean': summary['reward_mean'], + 'skill/format_rate': summary['parse_rate'], 'leak/rate': summary['leak_rate'], **tmetrics, **_cand_pass_metrics(chunk), **_leak_split(_cand_leak_pairs(chunk))} From c68395711ab75f572b09b3d95c9ce33c5a74f7cc Mon Sep 17 00:00:00 2001 From: root Date: Sun, 2 Aug 2026 17:29:29 +0800 Subject: [PATCH 33/60] fix(rl): keep sampled tokens intact end-to-end and align skill2lora with SEAM Framework fixes: - pad_and_stack_tensors: do not pad dim 0 when concat=True. dim 0 is the concatenation axis, so padding it fabricated all-pad sample rows; the metric side (align_logps_to_mask) then rejected the row count and silently skipped ratio/kl/clip for the whole step ("old_logps shape (3, N) does not match logps_mb shape (2, N)"). Added regression tests. - data_format: add pack_user_data() and make user_data_get() raise on a raw mapping. user_data must be a list of (key, json) pairs; a dict was silently ignored by readers, dropping key_rounds without any error. Updated all callers under cookbook/exp. - GRPOMetric: add anomaly probes for the trainable token distribution -- logp_min / logp_frac_lt_5 / logp_frac_lt_10, plus sampler_logp_mae, sampler_logp_max_abs and sampler_token_delta for a direct token-level reconciliation against the sampler's own logprobs. Behavioural metrics (format/acc/reward) cannot see an encoding mismatch; these can. skill2lora / SEAM alignment: - Feed the sampler's token ids straight into training (prompt encoded locally, response concatenated as tokens); sampler_token_delta is 0 across chunks. - Request logprobs from the sampler and pass them to GRPOMetric as sampler_logps for the reconciliation above. - Register GRPOMetric on the skill model; it was never added, so none of the RL metrics were on the path. - Do not double-prefix keys that already carry "train/". - Training-side baseline now runs n_skills draws per problem instead of one broadcast draw: SEAM sends the whole 1024-row batch to the baseline pass and its greedy results are not identical within a batch. - Report group_reward_std_mean with ddof=0 to match SEAM's np.std. The advantage-side std stays ddof=1, same as verl. - enable_prefix_caching=True on both samplers; verl hardcodes it in vllm_rollout_spmd, and it changes greedy results for repeated prompts. - --train-order-file pins the training order to SEAM's realized batch sequence (verl shuffles its dataloader, so the same pool yields different batches). - Correct the BNPO token_mean_scope tests: 'micro' is the default and matches verl's equal-weighted micro-means; 'global' is split-invariant and reports a sum for display. Verified against the SEAM paper-repro dump (Qwen3-4B, deepmath, 128x8): problem sets match row-for-row on the first eval and chunks 0-2, sampler_token_delta is 0 on every chunk, and all aggregate deltas sit inside the measured engine noise floor (first eval, identical weights and greedy on both sides: 114/128 problems agree, disagreements symmetric 6 vs 8). --- .gitignore | 2 +- .../legacy/build_reflexion_coldstart_sft.py | 3 +- cookbook/exp/legacy/train_reflexion_skill.py | 4 +- .../exp/legacy/train_reflexion_skill_old.py | 4 +- .../exp/legacy/train_reflexion_skill_rft.py | 3 +- .../exp/legacy/train_reflexion_skill_seam.py | 4 +- .../good_skill_hard_fail/eval_skill_probe.py | 2 + .../good_skill_hard_fail/reflexion_probe.py | 2 + .../skill_config_probe.py | 2 + cookbook/exp/skill2lora/logp_corr_probe.py | 5 +- cookbook/exp/skill2lora/run_ablate12.sh | 33 +- cookbook/exp/skill2lora/skill_ablate/main.py | 5 + .../exp/skill2lora/skill_ablate/methods.py | 123 ++++++-- .../exp/skill2lora/skill_ablate/rollouting.py | 34 ++- .../exp/skill2lora/skill_ablate/trainer.py | 24 +- cookbook/exp/skill2lora/train_skill_v2.py | 284 ++++++++++++++++-- src/twinkle/data_format/__init__.py | 2 +- src/twinkle/data_format/trajectory.py | 20 ++ src/twinkle/metric/grpo.py | 80 ++++- src/twinkle/utils/torch_utils.py | 15 +- tests/loss/test_bnpo_token_mean.py | 42 ++- tests/utils/test_utils.py | 9 +- 22 files changed, 602 insertions(+), 100 deletions(-) diff --git a/.gitignore b/.gitignore index db4719f34..d21498bcb 100644 --- a/.gitignore +++ b/.gitignore @@ -59,7 +59,7 @@ coverage.xml # Translations *.mo *.pot - +deepmath_103k # Django stuff: *.log local_settings.py diff --git a/cookbook/exp/legacy/build_reflexion_coldstart_sft.py b/cookbook/exp/legacy/build_reflexion_coldstart_sft.py index 548dfad58..a7748b859 100644 --- a/cookbook/exp/legacy/build_reflexion_coldstart_sft.py +++ b/cookbook/exp/legacy/build_reflexion_coldstart_sft.py @@ -29,6 +29,7 @@ import twinkle from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import pack_user_data from twinkle.sampler import vLLMSampler from cookbook.exp.embedding.train_reflexion_skill import ( @@ -468,7 +469,7 @@ def gen_one(r: Dict[str, Any]): messages = _sft_messages(r['problem'], response) sft_row = { 'messages': messages, - 'user_data': {'key_rounds': [len(messages) - 1]}, + 'user_data': pack_user_data({'key_rounds': [len(messages) - 1]}), 'data_id': r.get('data_id'), 'problem': r['problem'], 'reference_answer': r['reference_answer'], 'skills': g['skills'], 'response': response, diff --git a/cookbook/exp/legacy/train_reflexion_skill.py b/cookbook/exp/legacy/train_reflexion_skill.py index beccbc7b3..731478c97 100644 --- a/cookbook/exp/legacy/train_reflexion_skill.py +++ b/cookbook/exp/legacy/train_reflexion_skill.py @@ -43,7 +43,7 @@ import twinkle from twinkle import DeviceGroup, DeviceMesh, get_logger from twinkle.checkpoint_engine import CheckpointEngineManager -from twinkle.data_format import SamplingParams +from twinkle.data_format import SamplingParams, pack_user_data from twinkle.dataset import Dataset, DatasetMeta from twinkle.model import TransformersModel from twinkle.patch.no_split_modules import NoSplitModulesPatch @@ -1421,7 +1421,7 @@ def _train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: msgs = _skillgen_messages( rec['problem'], rec.get('view', 'A'), rec.get('diagnosis', ''), rec.get('rubric_src', '')) return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], - 'user_data': {'key_rounds': [len(msgs)]}} + 'user_data': pack_user_data({'key_rounds': [len(msgs)]})} def _train_chunk(skill_model, ref_model, ckpt: CheckpointEngineManager, samples: List[Dict[str, Any]], diff --git a/cookbook/exp/legacy/train_reflexion_skill_old.py b/cookbook/exp/legacy/train_reflexion_skill_old.py index 757d9d915..3bc9d5162 100644 --- a/cookbook/exp/legacy/train_reflexion_skill_old.py +++ b/cookbook/exp/legacy/train_reflexion_skill_old.py @@ -43,7 +43,7 @@ import twinkle from twinkle import DeviceGroup, DeviceMesh, get_logger from twinkle.checkpoint_engine import CheckpointEngineManager -from twinkle.data_format import SamplingParams +from twinkle.data_format import SamplingParams, pack_user_data from twinkle.dataset import Dataset, DatasetMeta from twinkle.model import TransformersModel from twinkle.patch.no_split_modules import NoSplitModulesPatch @@ -1465,7 +1465,7 @@ def _train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: msgs = _skillgen_messages( rec['problem'], rec.get('view', 'A'), rec.get('diagnosis', ''), rec.get('rubric_src', '')) return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], - 'user_data': {'key_rounds': [len(msgs)]}} + 'user_data': pack_user_data({'key_rounds': [len(msgs)]})} def _train_chunk(skill_model, ref_model, ckpt: CheckpointEngineManager, samples: List[Dict[str, Any]], diff --git a/cookbook/exp/legacy/train_reflexion_skill_rft.py b/cookbook/exp/legacy/train_reflexion_skill_rft.py index b4388b773..55197b5d9 100644 --- a/cookbook/exp/legacy/train_reflexion_skill_rft.py +++ b/cookbook/exp/legacy/train_reflexion_skill_rft.py @@ -45,6 +45,7 @@ import twinkle from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import pack_user_data from twinkle.checkpoint_engine import CheckpointEngineManager from twinkle.model import TransformersModel from twinkle.processor import InputProcessor @@ -1017,7 +1018,7 @@ def _train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: prefix already excludes the prompt-provided ````, so no extra masking is needed.""" msgs = _skillgen_messages(rec['problem'], rec.get('view', 'A'), rec.get('diagnosis', '')) full = msgs + [{'role': 'assistant', 'content': rec['response']}] - return {'messages': full, 'user_data': {'key_rounds': [len(msgs)]}} + return {'messages': full, 'user_data': pack_user_data({'key_rounds': [len(msgs)]})} def _train_chunk(skill_model, ckpt: Optional[CheckpointEngineManager], diff --git a/cookbook/exp/legacy/train_reflexion_skill_seam.py b/cookbook/exp/legacy/train_reflexion_skill_seam.py index 757d9d915..3bc9d5162 100644 --- a/cookbook/exp/legacy/train_reflexion_skill_seam.py +++ b/cookbook/exp/legacy/train_reflexion_skill_seam.py @@ -43,7 +43,7 @@ import twinkle from twinkle import DeviceGroup, DeviceMesh, get_logger from twinkle.checkpoint_engine import CheckpointEngineManager -from twinkle.data_format import SamplingParams +from twinkle.data_format import SamplingParams, pack_user_data from twinkle.dataset import Dataset, DatasetMeta from twinkle.model import TransformersModel from twinkle.patch.no_split_modules import NoSplitModulesPatch @@ -1465,7 +1465,7 @@ def _train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: msgs = _skillgen_messages( rec['problem'], rec.get('view', 'A'), rec.get('diagnosis', ''), rec.get('rubric_src', '')) return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], - 'user_data': {'key_rounds': [len(msgs)]}} + 'user_data': pack_user_data({'key_rounds': [len(msgs)]})} def _train_chunk(skill_model, ref_model, ckpt: CheckpointEngineManager, samples: List[Dict[str, Any]], diff --git a/cookbook/exp/skill2lora/good_skill_hard_fail/eval_skill_probe.py b/cookbook/exp/skill2lora/good_skill_hard_fail/eval_skill_probe.py index 0a0619f9d..48c3f6578 100644 --- a/cookbook/exp/skill2lora/good_skill_hard_fail/eval_skill_probe.py +++ b/cookbook/exp/skill2lora/good_skill_hard_fail/eval_skill_probe.py @@ -129,6 +129,8 @@ def _seam_sanitize(txt: str) -> str: txt = m.group(1).strip() elif (m := _SEAM_INLINE_RE.search(txt)): txt = (m.group(1) or m.group(2)).strip() + # bugfix 2026-07-29：\dfrac/\tfrac/\cfrac 归一，同 train_skill_v2._seam_sanitize + txt = txt.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac').replace(r'\cfrac', r'\frac') txt = re.sub(r'\\frac\s*\{\s*([^}]+?)\s*}\s*\{\s*([^}]+?)\s*}', r'\1/\2', txt) if (m := _SEAM_FRAC_RE.search(txt)): p, q = map(float, m.groups()) diff --git a/cookbook/exp/skill2lora/good_skill_hard_fail/reflexion_probe.py b/cookbook/exp/skill2lora/good_skill_hard_fail/reflexion_probe.py index 899776864..90cdd8252 100644 --- a/cookbook/exp/skill2lora/good_skill_hard_fail/reflexion_probe.py +++ b/cookbook/exp/skill2lora/good_skill_hard_fail/reflexion_probe.py @@ -144,6 +144,8 @@ def _seam_sanitize(txt): txt = m.group(1).strip() elif (m := _SEAM_INLINE_RE.search(txt)): txt = (m.group(1) or m.group(2)).strip() + # bugfix 2026-07-29：\dfrac/\tfrac/\cfrac 归一，同 train_skill_v2._seam_sanitize + txt = txt.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac').replace(r'\cfrac', r'\frac') txt = re.sub(r'\\frac\s*\{\s*([^}]+?)\s*}\s*\{\s*([^}]+?)\s*}', r'\1/\2', txt) if (m := _SEAM_FRAC_RE.search(txt)): p, q = map(float, m.groups()) diff --git a/cookbook/exp/skill2lora/good_skill_hard_fail/skill_config_probe.py b/cookbook/exp/skill2lora/good_skill_hard_fail/skill_config_probe.py index 685475ce7..069a74d7c 100644 --- a/cookbook/exp/skill2lora/good_skill_hard_fail/skill_config_probe.py +++ b/cookbook/exp/skill2lora/good_skill_hard_fail/skill_config_probe.py @@ -114,6 +114,8 @@ def _seam_sanitize(txt: str) -> str: txt = m.group(1).strip() elif (m := _SEAM_INLINE_RE.search(txt)): txt = (m.group(1) or m.group(2)).strip() + # bugfix 2026-07-29：\dfrac/\tfrac/\cfrac 归一，同 train_skill_v2._seam_sanitize + txt = txt.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac').replace(r'\cfrac', r'\frac') txt = re.sub(r'\\frac\s*\{\s*([^}]+?)\s*}\s*\{\s*([^}]+?)\s*}', r'\1/\2', txt) if (m := _SEAM_FRAC_RE.search(txt)): p, q = map(float, m.groups()) diff --git a/cookbook/exp/skill2lora/logp_corr_probe.py b/cookbook/exp/skill2lora/logp_corr_probe.py index 42ce06690..86ec3fbb6 100644 --- a/cookbook/exp/skill2lora/logp_corr_probe.py +++ b/cookbook/exp/skill2lora/logp_corr_probe.py @@ -31,6 +31,7 @@ from typing import Dict, List, Optional import numpy as np +from twinkle.data_format import pack_user_data SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) E15_DIR = os.path.join(SCRIPT_DIR, 'output.ablate12', 'E15_logp_gt_on_narrative') @@ -165,7 +166,7 @@ def load_pairs(): def phase_rollout(): import twinkle from twinkle import DeviceGroup, DeviceMesh - from twinkle.data_format import SamplingParams + from twinkle.data_format import SamplingParams, pack_user_data from twinkle.sampler import vLLMSampler from twinkle.template import Template @@ -239,7 +240,7 @@ def phase_logps(): def encode(problem, skill, gt): msgs = [dict(m) for m in build_skill_solve_prompt(problem, skill)['messages']] enc = tmpl.encode({'messages': msgs + [{'role': 'assistant', 'content': gt}], - 'user_data': {'key_rounds': [len(msgs)]}}) + 'user_data': pack_user_data({'key_rounds': [len(msgs)]})}) if enc is None: return None, None ids = [int(x) for x in enc['input_ids']] # numpy int64 -> int（vllm msgspec 拒收 np 类型） diff --git a/cookbook/exp/skill2lora/run_ablate12.sh b/cookbook/exp/skill2lora/run_ablate12.sh index a19563051..a61f1edbc 100644 --- a/cookbook/exp/skill2lora/run_ablate12.sh +++ b/cookbook/exp/skill2lora/run_ablate12.sh @@ -286,13 +286,32 @@ while IFS=$'\t' read -r NAME EXP_DIR THINK SMT OPTIONAL TASK EXEC_THINK; do # 显式同名 env 仍可覆盖。 E16_FLAGS="$E16_FLAGS --ppo-mini-batch-size ${E13_PPO_MINI:-160}" [ -z "$KL_BETA" ] && E16_FLAGS="$E16_FLAGS --kl-beta ${E13_KL_BETA:-0.001}" - # 显存：E13 8×140G 实测 micro=8（trainer.py 的 8192 自动档）OOM（132G 已用 + 4.77G）， - # micro=4（2 序列/卡）已实测跑得通，所以继续用 2×dp。 - # 注：verl 的聚合单元是 5 序列/卡（ppo_micro_batch_size_per_gpu=5），这里是 2；在 - # token_mean_scope='micro' 下这只是聚合粒度差（实测偏离真 token-mean 的倍率 - # 1.415@2 vs 1.35@5，约 5%），而且方向是更靠近 sequence-mean、长度耦合更弱， - # 不会重新引入 'global' 那个 97 倍的“少写 token”梯度。想完全对齐需 TRAIN_FSDP=2 腾显存后设 5。 - _tmb=${E13_TRAIN_MICRO_BATCH:-$((2*${TRAIN_GPUS:-2}))} + # ---- 钉住训练 batch 序列（2026-08-02）------------------------------------------- + # verl 的 dataloader 默认 data.shuffle=True，所以 SEAM 的 step k 不是 train.parquet 的 + # 第 k 个 128 切片。实测（align5 chunk0 vs SEAM step1）：同一个 5000 题池、同 batch + # size，但两边第一步喂的 128 题交集只有 1 题 -> withskill 0.848 vs 0.789、 + # baseline 0.562 vs 0.634、lift +0.285 vs +0.155，全是抽样差。 + # seam_train_order.jsonl 是从 SEAM rollout dump 反推的真实喂题序列（40×128=5120 行， + # 用 .tmp_analysis/mk_seam_train_order.py 生成），挂上后 chunk k 逐题 == SEAM step k+1。 + # 置空 E13_TRAIN_ORDER 即可回到自己的 shuffle（但就不能逐 step 对了）。 + _order=${E13_TRAIN_ORDER-/mnt/data/yzhao/tastelikefeet/twinkle/.tmp_analysis/seam_train_order.jsonl} + if [ -n "$_order" ] && [ -f "$_order" ]; then + E16_FLAGS="$E16_FLAGS --train-order-file $_order" + elif [ -n "$_order" ]; then + echo "[ablate12] WARN: train order file not found: $_order (fallback to shuffle)" + fi + # 显存 / 聚合粒度（2026-08-02 修正）： + # verl 的等权聚合单元是 ppo_micro_batch_size_per_gpu=5 条（dp_actor.py 在每个 micro 内部 + # masked_mean，再 *1/gas 累加），每 step 共 160条/5 = 32 个「5 条组」等权。 + # 旧配置 TRAIN_FSDP=1/DP=2 + micro=4（=2 条/卡）给出 80 个「2 条组」等权 —— micro 越小越 + # 靠 sequence-mean，短序列的每 token 权重越高、压长度的分量越强。align6 实测后果： + # actor 输出 12382 -> 9265 chars 只用 6 个 chunk，format 0.883 -> 0.988； + # SEAM 走完同一段要 40 步（12248 -> 9981，format 0.883 -> 0.943）。 + # 即长度收缩快约 6 倍，是 40 步里唯一超出 SEAM 自身噪声的偏离（format Δ0.056 = 2.4×sd）。 + # TRAIN_FSDP=2 把 fp32 权重+Adam(约 64G) 分到 2 卡，腾出的显存正好够 5 条/卡； + # 此时 TRAIN_DP=1、sft=5 -> 32 个「5 条组」等权，与 verl 逐组一致。 + # REF_FSDP 必须同为 2：否则 REF_DP=2 而 micro=5 不整除，会报 Batch too small。 + _tmb=${E13_TRAIN_MICRO_BATCH:-$([ "${TRAIN_FSDP:-1}" = 2 ] && echo 5 || echo $((2*${TRAIN_GPUS:-2})))} [ -n "${TRAIN_MICRO_BATCH:-}" ] && _tmb=$TRAIN_MICRO_BATCH E16_FLAGS="$E16_FLAGS --train-micro-batch $_tmb" echo "[ablate12] E13 SEAM-repro: chunk=$CHUNK_ARG min_level=$MIN_LEVEL_ARG train_micro_batch=$_tmb"\ diff --git a/cookbook/exp/skill2lora/skill_ablate/main.py b/cookbook/exp/skill2lora/skill_ablate/main.py index 4bd26c3b5..2311059f3 100644 --- a/cookbook/exp/skill2lora/skill_ablate/main.py +++ b/cookbook/exp/skill2lora/skill_ablate/main.py @@ -59,6 +59,11 @@ def _build_args(argv=None): p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) p.add_argument('--eval-size', type=int, default=128) p.add_argument('--seam-parquet-dir', type=str, default='') + p.add_argument('--train-order-file', type=str, default='', + help='jsonl of {data_id,problem,reference_answer} in a FIXED training order; ' + 'replaces the train split and disables ProblemPool shuffling, so chunk k ' + 'is exactly batch k. Used to pin twinkle to SEAM\'s realized batch ' + 'sequence (verl shuffles its dataloader, so same pool != same batches).') p.add_argument('--deepmath-dir', type=str, default='', help='DeepMath-103K parquet dir; when set, overrides --seam-parquet-dir/--dataset ' 'and uses the difficulty-stratified split (eval/train same level mix).') diff --git a/cookbook/exp/skill2lora/skill_ablate/methods.py b/cookbook/exp/skill2lora/skill_ablate/methods.py index ea323d3b0..2e1128f88 100644 --- a/cookbook/exp/skill2lora/skill_ablate/methods.py +++ b/cookbook/exp/skill2lora/skill_ablate/methods.py @@ -40,6 +40,7 @@ from typing import Any, Callable, Dict, List, Optional, Tuple import numpy as np +from twinkle.data_format import pack_user_data import train_skill_v2 as v2 from train_skill_v2 import ( @@ -139,18 +140,22 @@ def _skillgen_solve(ctx: MethodContext, items: List[Dict[str, Any]], n_skills: i sg_out = _run_samples(ctx.skill_sampler, [it['prompt'] for it in items], n_skills, args.skill_max_tokens, ctx.skill_dp, temperature=temperature, top_p=args.skill_gen_top_p, - top_k=args.skill_gen_top_k) + top_k=args.skill_gen_top_k, logprobs=1) flat = [] for it, seqs in zip(items, sg_out): it['record']['_cands'] = [] for s in seqs: resp = _clean_text(getattr(s, 'decoded', '') or '') block = _extract_skill(resp) or '' + # tokens = 采样端真实吐出的 token id；训练样本只能用它拼（见 v2.build_train_feature）。 + # logprobs 同长，只给 GRPOMetric 做采样/训练对账，不进 loss。 + _toks = [int(t) for t in (getattr(s, 'tokens', None) or [])] cand = {'skills': block, 'response': resp, 'parseable': bool(block), 'leaked': None, 'with_pass': None, 'reward': None, 'rolls': [], - 'advantage': 0.0, 'kept': False, + 'advantage': 0.0, 'kept': False, 'tokens': _toks, + 'logprobs': v2.sampler_logprobs(s), 'skillgen_stop': getattr(s, 'stop_reason', None), - 'skillgen_tokens': len(getattr(s, 'tokens', None) or [])} + 'skillgen_tokens': len(_toks)} it['record']['_cands'].append(cand) if block: flat.append((it, cand)) @@ -177,7 +182,11 @@ def _skillgen_solve(ctx: MethodContext, items: List[Dict[str, Any]], n_skills: i def _grpo_records(records: List[Dict[str, Any]], with_rubric: bool) -> List[Dict[str, Any]]: """Flatten per-problem _cands into GRPO train records (v2 shape). - ``with_rubric`` tags each record so the trajectory builder knows which prompt to rebuild.""" + ``with_rubric`` tags each record so the trajectory builder knows which prompt to rebuild. + + ``tokens`` 必须带上：它是训练样本的唯一来源（response 只留给 reward/审计），少了它 + trajectory builder 就会回退到 decode->重编码那条有偏差的路。``logprobs`` 同长，给 + GRPOMetric 做采样/训练对账。""" recs = [] for r in records: for c in r.get('_cands', []): @@ -186,6 +195,8 @@ def _grpo_records(records: List[Dict[str, Any]], with_rubric: bool) -> List[Dict recs.append({'problem': r['problem'], 'reference_answer': r['reference_answer'], 'data_id': r.get('data_id', ''), 'response': c['response'], 'skills': c['skills'], 'advantage': c['advantage'], + 'tokens': c.get('tokens') or [], + 'logprobs': c.get('logprobs') or [], 'kept': c['kept'], 'reward': c['reward'], 'rubric': r.get('_rubric', ''), 'with_rubric': with_rubric, 'sft': False}) return recs @@ -253,6 +264,9 @@ def _train_metrics(metric: Optional[Dict[str, Any]]) -> Dict[str, float]: if k.startswith('learning rate'): if 'group 1' in k: d['train/lr'] = float(val) + elif k.startswith('train/'): + # GRPOMetric 等已经自带 train/ 前缀，不能再套一层。 + d[k.replace(' ', '_')] = float(val) else: d[f'train/{k.replace(" ", "_")}'] = float(val) return d @@ -261,6 +275,12 @@ def _train_metrics(metric: Optional[Dict[str, Any]]) -> Dict[str, float]: # =========================================================================================== # OPSD teacher alignment: client-side encode -> response-only teacher logps # =========================================================================================== +def _encode_for_align(tmpl, built): + """``traj_fn``/``teacher_fn`` 现在可能直接返回成品 InputFeature（token 直通路径）， + 那就无需再编码；只有回退的 messages 形式才走 tmpl.encode。""" + return built if 'input_ids' in built else tmpl.encode(built) + + def _align_teacher(ctx: MethodContext, samples, traj_fn, teacher_fn): """Encode student & teacher trajectories with a CLIENT-SIDE clone of the remote template and keep only samples whose response-token counts match (they always should — same @@ -276,8 +296,8 @@ def _align_teacher(ctx: MethodContext, samples, traj_fn, teacher_fn): assert tmpl is not None, 'OPSD needs ctx.encode_template (built by the trainer)' keep, pos_lists, dropped = [], [], 0 for s in samples: - st = tmpl.encode(traj_fn(s)) - tt = tmpl.encode(teacher_fn(s)) + st = _encode_for_align(tmpl, traj_fn(s)) + tt = _encode_for_align(tmpl, teacher_fn(s)) if st is None or tt is None: # deleted by max-length truncation dropped += 1 continue @@ -307,11 +327,21 @@ def _gather_response_logps(full_logps, pos_lists) -> List[List[float]]: # E14+ helpers: executor pseudo-GT + dense logP reward # =========================================================================================== def _executor_answer_trajectory(problem: str, skill: str, answer_text: str, - raw_response: Optional[str] = None) -> Dict[str, Any]: - """Teacher-forcing trajectory for executor logP(S | problem + skill).""" + raw_response: Optional[str] = None, + answer_tokens: Optional[List[int]] = None, + template=None) -> Dict[str, Any]: + """Teacher-forcing sample for executor logP(S | problem + skill). + + S 是 executor 自己采出来的伪 GT（E14），所以带了 ``answer_tokens`` 时直接拼采样 token： + 这里的 prompt 每次都不同（换 skill）但被打分的 token 必须是同一串，正是 + ``build_train_feature`` 的场景。E15 的 S 是 DeepMath 的外部 R1 参考解，本来就不是模型 + 产出、没有对应 token，只能走 messages 编码。 + """ msgs = [dict(m) for m in build_skill_solve_prompt(problem, skill, raw_response)['messages']] + if answer_tokens: + return v2.build_train_feature(msgs, answer_tokens, template=template) return {'messages': msgs + [{'role': 'assistant', 'content': answer_text}], - 'user_data': {'key_rounds': [len(msgs)]}} + 'user_data': pack_user_data({'key_rounds': [len(msgs)]})} def _set_ref_executor_template(ctx: MethodContext, enable_thinking: bool) -> None: @@ -334,7 +364,7 @@ def _score_executor_mean_logps(ctx: MethodContext, trajs: List[Dict[str, Any]]) out: List[Optional[float]] = [None] * len(trajs) valid_trajs, pos_lists, valid_idx = [], [], [] for i, tr in enumerate(trajs): - enc = tmpl.encode(tr) + enc = _encode_for_align(tmpl, tr) if enc is None: continue pos = np.where(np.asarray(enc.get('labels')) != -100)[0] @@ -378,7 +408,8 @@ def _train_batch(ctx: MethodContext, samples: List[Dict[str, Any]], it after a step would go off-policy). Returns (n_updates, train metrics). """ args = ctx.args - samples = [s for s in samples if (s.get('response') or '').strip()] + samples = [s for s in samples + if (s.get('tokens') if s.get('tokens') is not None else (s.get('response') or '').strip())] is_opsd = teacher_fn is not None teacher_pos: Optional[List] = None n_align_drop = 0 @@ -403,9 +434,15 @@ def _train_batch(ctx: MethodContext, samples: List[Dict[str, Any]], mini = max(sft, (mini // sft) * sft) multi_step = mini < n # pre-compute every micro's ref/old/teacher logps BEFORE any update (v2 pattern) - micro_ref, micro_old, micro_teacher = [], [], [] + micro_ref, micro_old, micro_teacher, micro_smp = [], [], [], [] + # 采样端 logprob，只给 GRPOMetric 做对账（sampler_logp_mae / sampler_token_delta），不进 loss。 + # 整个 micro 都带齐了才传：SFT 样本的 response 是合成文本、没有采样 logprob，混进去会 + # 让 token_delta 无法解释（它的语义是「应恒为 0」）。 + smp_all = [list(s.get('logprobs') or []) for s in samples] for i in range(0, n, sft): mb = trajs[i:i + sft] + smp = smp_all[i:i + sft] + micro_smp.append(smp if all(smp) else None) if is_opsd: # no ref forward at all: OPSDLoss uses only teacher_logps (kl_beta plays no role) t_mb = [teacher_fn(s) for s in samples[i:i + sft]] @@ -423,10 +460,12 @@ def _train_batch(ctx: MethodContext, samples: List[Dict[str, Any]], k = i // sft if is_opsd: ctx.skill_model.forward_backward(inputs=trajs[i:i + sft], - teacher_logps=micro_teacher[k]) + teacher_logps=micro_teacher[k], + sampler_logps=micro_smp[k]) else: ctx.skill_model.forward_backward(inputs=trajs[i:i + sft], advantages=advs[i:i + sft], - old_logps=micro_old[k], ref_logps=micro_ref[k]) + old_logps=micro_old[k], ref_logps=micro_ref[k], + sampler_logps=micro_smp[k]) ctx.skill_model.clip_grad_and_step() n_steps += 1 ctx.ckpt.sync_weights(merge_and_sync=True) @@ -483,6 +522,13 @@ def step(self, chunk, ci): 'leak/rate': summary['leak_rate'], **tmetrics, **_cand_pass_metrics(chunk), **_leak_split(_cand_leak_pairs(chunk))} + if summary.get('baseline_pass_train') is not None: + # 训练侧 no-skill baseline：SEAM 只在 step1 跑（ray_trainer.py:1461），所以 twinkle 也只在 + # chunk 0 有值，两边 lift 就是同一个可比的点。 + metrics['acc/baseline_pass'] = summary['baseline_pass_train'] + metrics['acc/lift'] = summary['lift_train'] + metrics['train/baseline_accuracy'] = summary['baseline_pass_train'] + metrics['train/lift'] = summary['lift_train'] return {'n_updates': n_upd, 'summary': summary, 'metrics': metrics, 'gen_records': full} @@ -572,10 +618,16 @@ def step(self, chunk, ci): # first-pass ONE skill per problem (query-only, improve temperature) sg = _run_samples(ctx.skill_sampler, [_skillgen_prompt(r['problem']) for r in chunk], 1, args.skill_max_tokens, ctx.skill_dp, - temperature=args.improve_skill_temperature) - first = [] + temperature=args.improve_skill_temperature, logprobs=1) + first, toks_by, lps_by = [], {}, {} for r, seqs in zip(chunk, sg): resp = _clean_text(getattr(seqs[0], 'decoded', '') or '') if seqs else '' + # 采样 token 另存一份（每题只一个 skill，用 id(r) 做键就够）：它是 student/teacher 两边 + # 要打分的同一串 token，不能拿 decode 后的文本重编码（见 v2.build_train_feature）。 + # logprobs 同长，只给 GRPOMetric 做采样/训练对账。 + toks_by[id(r)] = ([int(t) for t in (getattr(seqs[0], 'tokens', None) or [])] + if seqs else []) + lps_by[id(r)] = v2.sampler_logprobs(seqs[0]) if seqs else [] first.append((r, resp, _extract_skill(resp) or '')) # executor solve WITH skill (only parseable skills) flat = [(r, resp, sk) for r, resp, sk in first if sk] @@ -594,7 +646,9 @@ def step(self, chunk, ci): diags = _diagnose_parallel( ctx, [(_rubric_entry(r, roll), sk) for r, _resp, sk, roll in wrong]) samples = [{'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'data_id': r.get('data_id', ''), 'response': resp, 'rubric': diag} + 'data_id': r.get('data_id', ''), 'response': resp, 'rubric': diag, + 'tokens': toks_by.get(id(r)) or [], + 'logprobs': lps_by.get(id(r)) or []} for (r, resp, sk, _roll), diag in zip(wrong, diags)] n_upd, tmetrics = 0, {} if samples: @@ -631,7 +685,7 @@ def _skillgen(self, usable): return _run_samples(ctx.skill_sampler, [_skillgen_prompt(r['problem']) for r in usable], args.n_skills, args.skill_max_tokens, ctx.skill_dp, temperature=args.skill_gen_temperature, top_p=args.skill_gen_top_p, - top_k=args.skill_gen_top_k) + top_k=args.skill_gen_top_k, logprobs=1) def _prepare(self, chunk, ci): """executor T>0 采 K 条 -> 选本地判分正确且非截断的伪 GT S -> rubric API 后台审计 @@ -647,12 +701,18 @@ def _prepare(self, chunk, ci): usable, audit_jobs = [], [] for r, seqs in zip(chunk, out): rolls = [_parse_seq(s, r['reference_answer']) for s in (seqs or [])] - ok = next((x for x in rolls if x['correct'] and x.get('stop_reason') != 'length'), None) + ok_i = next((j for j, x in enumerate(rolls) + if x['correct'] and x.get('stop_reason') != 'length'), None) r['_pseudo_rolls'] = rolls - if ok is None: + if ok_i is None: continue + ok = rolls[ok_i] r['_pseudo_roll'] = ok r['_pseudo_solution'] = ok.get('text', '') + # 伪 GT 是 executor 采样产出，打分时直接用它的 token，不拿 _clean_text 后的文本重编码。 + # 从 seq 取而不是从 roll 取：roll 会被每个候选持有，token 存进去会把内存撑爆 + # （E16 那种 M=8 的臂一个 chunk 就是上千条 rollout）。 + r['_pseudo_tokens'] = [int(t) for t in (getattr(seqs[ok_i], 'tokens', None) or [])] usable.append(r) audit_jobs.append((_rubric_entry(r, ok), ok.get('text', ''))) # rubric 审计是纯 API：后台线程跑，与 skill-gen 的 GPU rollout 重叠（API/GPU overlap） @@ -678,11 +738,13 @@ def _score_and_train(self, chunk, ci, usable, sg): pseudo = dict(r['_pseudo_roll']) if si > 0: pseudo['text'] = '' # 磁盘保护：伪 GT 全文每题只在首个候选保留一份 + _toks = [int(t) for t in (getattr(s, 'tokens', None) or [])] cand = {'skills': block, 'response': resp, 'parseable': bool(block), 'leaked': None, 'with_pass': None, 'reward': None, 'rolls': [pseudo], - 'advantage': 0.0, 'kept': False, + 'advantage': 0.0, 'kept': False, 'tokens': _toks, + 'logprobs': v2.sampler_logprobs(s), 'skillgen_stop': getattr(s, 'stop_reason', None), - 'skillgen_tokens': len(getattr(s, 'tokens', None) or []), + 'skillgen_tokens': len(_toks), 'logp_base': None, 'logp_skill': None, 'logp_delta': None} r['_cands'].append(cand) if block: @@ -692,12 +754,17 @@ def _score_and_train(self, chunk, ci, usable, sg): if c['leaked'] is None: c['leaked'] = False if flat: - base_trajs = [_executor_answer_trajectory(r['problem'], '', r['_pseudo_solution']) + base_trajs = [_executor_answer_trajectory(r['problem'], '', r['_pseudo_solution'], + answer_tokens=r.get('_pseudo_tokens'), + template=ctx.encode_template) for r in usable] base_logps = _score_executor_mean_logps(ctx, base_trajs) base_by_id = {id(r): lp for r, lp in zip(usable, base_logps)} cand_trajs = [_executor_answer_trajectory(r['problem'], c['skills'], r['_pseudo_solution'], - c.get('response')) for r, c in flat] + c.get('response'), + answer_tokens=r.get('_pseudo_tokens'), + template=ctx.encode_template) + for r, c in flat] cand_logps = _score_executor_mean_logps(ctx, cand_trajs) # bugfix #14: logP 目标超长被 truncation='delete' 删掉时 lp=None → reward 地板， # 整组塔到 -1.0 会零梯度空转；这里显式监控 encode 失败占比（E15 的 R1 参考解尤其长）。 @@ -1033,17 +1100,19 @@ def _score_candidates(self, kept, prompts): sg = _run_samples(ctx.skill_sampler, list(prompts), args.n_skills, args.skill_max_tokens, ctx.skill_dp, temperature=args.skill_gen_temperature, top_p=args.skill_gen_top_p, - top_k=args.skill_gen_top_k) + top_k=args.skill_gen_top_k, logprobs=1) flat = [] for r, seqs in zip(kept, sg): for s in seqs or []: resp = _clean_text(getattr(s, 'decoded', '') or '') block = _extract_skill(resp) or '' + _toks = [int(t) for t in (getattr(s, 'tokens', None) or [])] cand = {'skills': block, 'response': resp, 'parseable': bool(block), 'leaked': None, 'with_pass': None, 'reward': None, 'rolls': [], - 'advantage': 0.0, 'kept': False, + 'advantage': 0.0, 'kept': False, 'tokens': _toks, + 'logprobs': v2.sampler_logprobs(s), 'skillgen_stop': getattr(s, 'stop_reason', None), - 'skillgen_tokens': len(getattr(s, 'tokens', None) or []), + 'skillgen_tokens': len(_toks), 'trunc_pen': None, 'pass_rate': None, 'loop_pen': None, 'eff': None, 'loop_density': None} r['_cands'].append(cand) diff --git a/cookbook/exp/skill2lora/skill_ablate/rollouting.py b/cookbook/exp/skill2lora/skill_ablate/rollouting.py index 9b5f0ab25..e3e9fff93 100644 --- a/cookbook/exp/skill2lora/skill_ablate/rollouting.py +++ b/cookbook/exp/skill2lora/skill_ablate/rollouting.py @@ -27,6 +27,8 @@ """ from typing import Any, Dict +from twinkle.data_format import pack_user_data + import train_skill_v2 as v2 from train_skill_v2 import ( # noqa: F401 (re-exported for methods.py convenience) _answer_leaked, @@ -172,17 +174,21 @@ def rubric_skillgen_prompt(problem: str, rubric: str) -> Dict[str, Any]: def rubric_train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: - """Train trajectory whose PROMPT is the query+rubric skill-gen prompt + the response. + """Train sample whose PROMPT is the query+rubric skill-gen prompt + the sampled response. + + Used by view-A RL (rl_ab / rl_err / reflexion) only: train WITH rubric in the prompt + (knowledge-transfer probe; eval is still query-only via v2 ``_skillgen_prompt``). The + rebuilt prompt matches the prompt the skills were SAMPLED under (on-policy consistency). - Used by view-A RL (rl_ab / rl_err) only: train WITH rubric in the prompt (knowledge- - transfer probe; eval is still query-only via v2 ``_skillgen_prompt``). The rebuilt prompt - matches the prompt the skills were SAMPLED under (on-policy consistency). - ``rec`` must carry 'problem', 'rubric' and 'response'. ``key_rounds`` marks the final - assistant turn as the only trainable span (identical convention to v2 ``_train_trajectory``). + response 段直接拼采样返回的 token（``rec['tokens']``），绝不 decode 后重新过模板 —— + 重渲染会把模板自己补的换行/EOS/空思考块训进去，详见 v2.build_train_feature 的注释。 + 只有合成文本（没有 tokens 的 SFT 记录）才回退到 messages 编码。 """ msgs = rubric_skillgen_prompt(rec['problem'], rec.get('rubric', ''))['messages'] + if rec.get('tokens'): + return v2.build_train_feature(msgs, rec['tokens']) return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], - 'user_data': {'key_rounds': [len(msgs)]}} + 'user_data': pack_user_data({'key_rounds': [len(msgs)]})} # 中文注释：OPSD teacher 的特权信息块——按设计 871 行要求放进 SYSTEM prompt，且只做“追加”， @@ -196,15 +202,21 @@ def rubric_train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: def opsd_teacher_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: - """OPSD teacher trajectory: student's query-only prompt + rubric appended to the SYSTEM - prompt + the SAME response. With an empty rubric the teacher degenerates to the student - (zero distillation pull), which is the safe behaviour on rubric API failure.""" + """OPSD teacher: student's query-only prompt + rubric appended to the SYSTEM prompt + the + SAME sampled response tokens. With an empty rubric the teacher degenerates to the student + (zero distillation pull), which is the safe behaviour on rubric API failure. + + teacher 与 student 必须打分**同一串 token**，所以两边的 response 段都直接拼 ``rec['tokens']``； + 只有 prompt 段不同（多一段 rubric）。 + """ msgs = [dict(m) for m in _skillgen_prompt(rec['problem'])['messages']] rubric = (rec.get('rubric') or '').strip() if rubric and msgs[0]['role'] == 'system': msgs[0]['content'] = msgs[0]['content'] + _OPSD_TEACHER_SUFFIX.format(rubric=rubric) + if rec.get('tokens'): + return v2.build_train_feature(msgs, rec['tokens']) return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], - 'user_data': {'key_rounds': [len(msgs)]}} + 'user_data': pack_user_data({'key_rounds': [len(msgs)]})} def query_only_train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: diff --git a/cookbook/exp/skill2lora/skill_ablate/trainer.py b/cookbook/exp/skill2lora/skill_ablate/trainer.py index 62cc3e8c8..506398de9 100644 --- a/cookbook/exp/skill2lora/skill_ablate/trainer.py +++ b/cookbook/exp/skill2lora/skill_ablate/trainer.py @@ -146,6 +146,20 @@ def run_experiment(args, spec: ExpSpec) -> None: else: records, eval_records = (load_deepmath_records(args) if getattr(args, 'deepmath_dir', '') else v2._load_records(args)) + # ⭐ --train-order-file：用 SEAM 已实现的 batch 序列覆盖 train 分支（eval 划分不动 —— twinkle 的 + # eval 128 题已逐题验证与 SEAM val.parquet 一致）。verl 的 dataloader 会 shuffle，所以两边 + # 即使同题池、同 batch size，step k 实际喂的 128 题也几乎不重叠（实测交集 1/128）， + # 单此一项就能把 acc/baseline/lift 拉开 6-7 个点。给了 order 文件后 chunk k 逐题 == SEAM step k+1。 + order_file = (getattr(args, 'train_order_file', '') or '').strip() + if order_file: + records = v2.load_train_order_file(order_file) + ev = {r['problem'] for r in eval_records} + overlap = sum(1 for r in records if r['problem'] in ev) + if overlap: + raise ValueError(f'--train-order-file overlaps eval on {overlap} rows: {order_file}') + sys.stderr.write(f'[ablate] train order pinned to {order_file}: {len(records)} rows ' + f'({len(records) // max(1, args.chunk_size)} chunks of {args.chunk_size}), ' + f'ProblemPool shuffle DISABLED.\n') if len(records) < args.chunk_size: raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded ({len(records)})') @@ -186,6 +200,13 @@ def run_experiment(args, spec: ExpSpec) -> None: else spec.thinking == 'on'), max_length=args.max_model_len, truncation_strategy='delete') + # 所有臂都需要的 skill 侧模板副本：训练样本的 prompt 段在客户端编码，response 段直接拼 + # 采样返回的 token（见 v2.build_train_feature）。必须与 skill_model/skill_sampler 同配置， + # 否则 prompt 段的 token 会与采样时对不上。 + v2.set_encode_template(v2.Template(model_id=v2.MODEL_ID, + enable_thinking=(spec.thinking == 'on'), + max_length=args.max_model_len, + truncation_strategy='delete')) pool = _build_pool(spec, args) ctx = MethodContext(skill_model=skill_model, ref_model=ref_model, skill_sampler=skill_sampler, @@ -344,7 +365,8 @@ def _do_eval(updates_done: int, swan_step: int) -> None: if eval_records and resume is None: _do_eval(-1, 0) # baseline before any update (chunk axis position 0) - pool_pp = v2.ProblemPool(records, args.seed) + pool_pp = v2.ProblemPool(records, args.seed, + fixed_order=bool((getattr(args, 'train_order_file', '') or '').strip())) updates = 0 last_eval_at = 0 last_eval_updates = -1 diff --git a/cookbook/exp/skill2lora/train_skill_v2.py b/cookbook/exp/skill2lora/train_skill_v2.py index d5af4f03c..40aec9094 100644 --- a/cookbook/exp/skill2lora/train_skill_v2.py +++ b/cookbook/exp/skill2lora/train_skill_v2.py @@ -33,7 +33,7 @@ import twinkle from twinkle import DeviceGroup, DeviceMesh, get_logger from twinkle.checkpoint_engine import CheckpointEngineManager -from twinkle.data_format import SamplingParams +from twinkle.data_format import SamplingParams, pack_user_data from twinkle.dataset import Dataset, DatasetMeta from twinkle.model import TransformersModel from twinkle.patch.no_split_modules import NoSplitModulesPatch @@ -437,14 +437,15 @@ def _empty_roll(): def _run_samples(sampler, prompts, num_samples, max_tokens, gen_dp, - temperature=None, top_p=None, top_k=None): + temperature=None, top_p=None, top_k=None, logprobs=None): if not prompts: return [] params = SamplingParams( max_tokens=max_tokens, temperature=GEN_TEMPERATURE if temperature is None else temperature, top_p=GEN_TOP_P if top_p is None else top_p, - num_samples=num_samples, **({} if top_k is None else {'top_k': top_k})) + num_samples=num_samples, **({} if top_k is None else {'top_k': top_k}), + **({} if logprobs is None else {'logprobs': logprobs})) padded = prompts if gen_dp > 1 and 0 < len(prompts) < gen_dp: padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] @@ -452,6 +453,23 @@ def _run_samples(sampler, prompts, num_samples, max_tokens, gen_dp, return [list(r.sequences) if (r and r.sequences) else [] for r in responses] +def sampler_logprobs(seq): + """从 ``SampledSequence.logprobs`` 抽出采样 token 自己的 logprob（逐 token 一个 float）。 + + 采样器返回的形状是 ``[[(token_id, logprob), ...], ...]``（logprobs=1 时每位只有一项， + 就是被采中的那个 token）。抽出来喂给 GRPOMetric 的 ``sampler_logps``，它会把这串值与 + 训练 forward 算出的 logp 逐 token 对账 —— 序列一致时两边只差引擎精度。 + 采样时 T=1/top_p=1/top_k=-1（skill-gen 的固定参数），logits 没经过任何处理，所以 + vLLM 的 processed_logprobs 就等于原始 logprob，与 trainer 的取值口径一致。 + """ + out = [] + for item in (getattr(seq, 'logprobs', None) or []): + if not item: + return [] + out.append(float(item[0][1])) + return out + + # =========================================================================== # Section C — data loading (simplified: no balance, no xproblem, no views) # =========================================================================== @@ -499,6 +517,38 @@ def _load_seam_parquet(path: str) -> List[Dict[str, Any]]: return out +def load_train_order_file(path: str) -> List[Dict[str, Any]]: + """Read a fixed training ORDER file (jsonl) -> records in file order, duplicates kept. + + Why this exists: verl's dataloader shuffles (data.shuffle defaults True), so SEAM's step-k + batch is NOT train.parquet[k*128:(k+1)*128]. Measured 2026-08-02: twinkle chunk 0 and SEAM + step 1 drew from the SAME 5000-problem pool but shared only 1 of 128 problems, which alone + put ~6-7 accuracy points between the two curves. This file is SEAM's REALIZED batch + sequence, reverse-engineered from its rollout dump (40 steps x 128 problems, in order), so + feeding it with ProblemPool(fixed_order=True) makes chunk k == SEAM step k+1 problem by + problem. Generated by .tmp_analysis/mk_seam_train_order.py. + Each line: {'data_id','problem','reference_answer'[,'level','seam_step']}. + """ + out: List[Dict[str, Any]] = [] + with open(path, encoding='utf-8') as f: + for line in f: + line = line.strip() + if not line: + continue + d = json.loads(line) + problem, ref = (d.get('problem') or '').strip(), d.get('reference_answer') + if not problem or ref is None: + continue + rec = {'data_id': str(d.get('data_id', '')), 'problem': problem, + 'reference_answer': str(ref)} + if d.get('level') is not None: + rec['level'] = d['level'] + out.append(rec) + if not out: + raise ValueError(f'--train-order-file {path} produced 0 records') + return out + + def _load_records(args): seam_dir = (getattr(args, 'seam_parquet_dir', '') or '').strip() if seam_dir: # 直读 SEAM parquet：按文件顺序取 train，val 整份当 eval，跳过 load/numeric/shuffle/split @@ -600,16 +650,24 @@ class ProblemPool: SEAM's verl dataloader uses a sampler and drop_last=True; this mirrors that behavior more closely than the old cursor loop that carried a short epoch tail into the next batch. + + ``fixed_order=True`` disables the permutation and walks ``records`` in file order. That is + what --train-order-file needs: the order file already IS verl's realized batch sequence + (reverse-engineered from SEAM's rollout dump), so any reshuffle here would destroy it. """ - def __init__(self, records, seed): + def __init__(self, records, seed, fixed_order=False): self._records = list(records) self._seed, self._cursor, self.epoch = seed, 0, 0 + self._fixed_order = bool(fixed_order) self._order: List[int] = [] self._reset_epoch() def _reset_epoch(self): - rng = np.random.RandomState(self._seed + self.epoch) - self._order = list(rng.permutation(len(self._records))) + if self._fixed_order: + self._order = list(range(len(self._records))) + else: + rng = np.random.RandomState(self._seed + self.epoch) + self._order = list(rng.permutation(len(self._records))) self._cursor = 0 def draw(self, k): @@ -896,6 +954,68 @@ def _diagnose_entry(checker, entry: Dict[str, Any]) -> Optional[str]: # enable_thinking=True 时 Qwen3 生成起点默认进入 thinking；这里不改共享 sampler 的 assistant 前缀注入。 _ALIGN_MODE = 'v2' # 'v2' | 'seam'；由 main() 依据 --align-mode 设置 +# =========================================================================================== +# SEAM 对齐定案（2026-08-02，E13 = --align-mode seam） +# 目的：把追查过程固化下来，避免以后重复审查同样的地方。 +# +# 【已逐条/逐字节校验对齐的 12 项】 +# 1. 训练喂题序列 —— verl dataloader data.shuffle 默认 True，SEAM 的 step k 不是 +# train.parquet[k*128:(k+1)*128]。实测 twinkle chunk0 与 SEAM step1 的 128 题交集只有 +# 1 题，单这一项就拉开 6-7 个 acc 点。已从 SEAM rollout dump 反推真实序列，用 +# --train-order-file + ProblemPool(fixed_order=True) 钉死（5120 行，missing=0）。 +# 2. actor prompt / executor prompt（with-skill 与 baseline 两路）—— 含 chat template 与 +# Qwen3 的空 \n\n\n\n，逐字节 True。 +# 3. templates/slove_qwen.txt advisory —— 逐字节 True（含第一行末尾那个空格、EOF 无换行）。 +# 4. 判分口径 —— SEAM_JUDGE=answer 的整段级联（ -> \boxed -> $..$ -> 分数 -> 首个 +# 数字）。用 SEAM dump replay 20480/20480 与其 acc 逐条一致；旧的 boxed-only 判分对 +# SEAM 自己的输出只给 acc≈0.05。 +# 5. format 抽取 —— 整段搜首个 ，20480/20480 与 SEAM format 一致。 +# 6. loss 聚合 —— BNPO token_mean_scope='micro'。verl 侧证据：dp_actor.py 在每个 micro 内部 +# masked_mean，再乘 loss_scale_factor = n_micro / ppo_mini_batch_size(=40/卡)，满 micro 时 +# 恰好等权；即 verl 也是「每 5 条组等权」而不是全局 token-mean。 +# 7. 聚合粒度 —— verl 的等权单元是 ppo_micro_batch_size_per_gpu=5 条。TRAIN_FSDP=2 腾出显存后 +# train_micro_batch=5、TRAIN_DP=1，每 optimizer step 32 个「5 条组」等权，与 verl 逐组一致。 +# 8. 每 step 的 optimizer step 数 —— verl 的 ppo_mini_batch_size 经 fsdp_workers.py:198-199 +# 归一化为 per-GPU 40，每卡 256 条 => 7 个 optimizer step；twinkle mini=160、n=1024 也是 7。 +# 9. PPO 语义 —— clip_ratio 0.2 / clip_ratio_c 3.0 / ppo_epochs 1 / entropy_coeff 0 / +# use_kl_loss+low_var_kl+coef 0.001 / clip_grad 1.0 / AdamW wd 0.01 全部一致；old_logps 在 +# 任何更新前统一预计算（multi_step=True），所以 7 个 mini-step 的 PPO clip 真实有约束力。 +# 10. advantage —— A=(r-mean)/std（unbiased std，norm_adv_by_std_in_grpo 语义）；drop_zero_adv=False， +# 零 adv 样本照样占分母。实测 |adv| 均值与 SEAM adv_absmax_mean 逐 step 吻合（0.20-0.28）。 +# 11. 数值精度 —— fp32 master + bf16 计算（torch_dtype='float32' + mixed_precision='bf16'）， +# 与 verl FSDP 的 fp32 master + param_dtype=bf16 等价；lr 恒定 1e-6 无 warmup/decay。 +# 12. 采样超参 —— actor temperature 1.0 / top_p 1.0 / top_k 禁用；executor greedy + nothink； +# 8192/8192 预算，executor prompt 实测 ≤8522 < 10752 不触发左截断。 +# +# 【残余偏离已定案并修复：截断 rollout 的思考段被 chat template 补了一个空思考块】 +# 现象：两边起点几乎同一点（actor 输出 3851 vs 3842 tokens、format 0.874 vs 0.883），但 +# twinkle 6 个 chunk 就走完 SEAM 40 步的长度降幅（3851->3177 vs 3842->3239，均 -16~17%）， +# 于是 format 0.874->0.957 而 SEAM 40 步才 0.883->0.943。acc/reward/lift/zero_grad/grp_std/ +# n_train 均已在 SEAM 自身 step 间噪声(sd≈0.03)内逐 step 对应。 +# 定位方法（.tmp_analysis/probe_grad.py + probe_twinkle.py）：第 1 个 optimizer step 时 ratio≡1、 +# KL≡0，梯度完全由 (advantage, logits) 决定，所以拿 SEAM step1 dump 的同一批样本做一次 +# forward+backward 就能逐层对。结果： +# * 环境差异排除 —— 同一份参考实现在 torch 2.11/tf 5.12 与 SEAM 的 torch 2.7/tf 4.53 下 +# tokenization md5 相同、loss 差 0.24%、grad_norm 差 0.9%。 +# * 逐组定位 —— 只有「非零 advantage 的样本恰好是截断样本」的 micro 组对不上（g4 +# 0.5927 -> 2.0552，3.47 倍），不含截断样本的组差 <1%。 +# * 真因 —— 撞 8192 上限的 rollout 思考段没有 （160 条里 18 条截断，其中 16 条 +# 无 ，且无 的样本 100% 是截断样本），Template 的 pre-pipeline +# _to_standard_reasoning_content 拆不出思考段，只能置 reasoning_content=''，Qwen3 模板于是 +# 渲染成空思考块 + 原文，原文自带的 变成紧跟在闭合标签之后的 token —— 那个 +# 位置模型输出 的概率≈0，logp 极低、梯度极大。 +# * 因果验证 —— 在参考实现里复刻这一编码（probe_grad.py --emulate-think-bug）后，逐组 +# grad_norm 从 [0.9999, 0.5733, 0.5927, 0.5307] 变成 [1.2520, 0.5672, 2.0539, 0.5310]， +# 与 twinkle 实测 [1.2573, 0.5787, 2.0552, 0.5338] 逐组重合。 +# * 修复验证 —— Template._fix_unfinished_last_round 上线后 g4 2.0552 -> 0.5922 +# （ref 0.5927），整批 160 条 32 组 0.3505 -> 0.1151（ref 0.1149，差 0.2%）。 +# 为何只影响长度/format、不影响 acc/reward：截断样本 reward=0（没闭合 ）、 +# advantage 为负，那个巨大梯度全压在“长输出”这一模式上；而截断率会随训练自我消解 +# （chunk0 11.6% -> chunk6 1.2%），所以偏差在前几步最猛、之后自行消失 —— 正好解释了 +# 「twinkle 6 步冲完然后平稳 vs SEAM 40 步缓慢上升」。 +# 另注：grad_norm 不可直接比（twinkle 只记 7 个 optimizer step 中的最后一个，verl 记均值）。 +# =========================================================================================== + _SEAM_EXPERIENCE_PROMPT = ( 'You are a problem-solving guidance model. Read the math problem below and ' 'distill a concise, reusable piece of solving experience that will help a ' @@ -1306,13 +1426,52 @@ def _assign_advantages(chunk, args): c['kept'] = c['reward'] > mean_r +# =========================================================================================== +# 训练样本构造：response 段一律直接用采样返回的 token，严禁 decode 后重新过模板 +# =========================================================================================== +# 为什么（2026-08-02 定案，.tmp_analysis/probe_token_direct.py 在 160 条真实样本上实测）： +# 把采样产出 decode 成文本、再塞回 messages 让 chat template 重新渲染一遍，可训练区就不再 +# 等于模型真实生成的 token —— 模板会给 assistant 角色行带一个换行、给每条 message 补结尾 +# EOS，思考段没闭合时还会先渲染一个空思考块。实测 160/160 条的可训练区都比采样 token 多； +# 截断样本更被塞进一个紧跟 之后的 （p≈0、logp≈-20），单个 micro 的 +# grad_norm 从 0.59 抬到 2.05，整批放大 3 倍，直接把 format_rate 曲线提前顶到 0.99。 +# +# 正确形状：prompt 段照常编码（它本来就是模板产物，重编码无风险），response 段原样拼采样 +# token。同一条实测里这条路径 160/160 逐 token 相等。采样器已经把 token 备好了 +# （SampledSequence.tokens），所以候选记录只需把它带下来。 +_ENCODE_TEMPLATE = None + + +def set_encode_template(template) -> None: + """注入 skill 模型的客户端 Template 副本，供 build_train_feature 编码 prompt 段。""" + global _ENCODE_TEMPLATE + _ENCODE_TEMPLATE = template + + +def build_train_feature(prompt_messages, tokens, template=None): + """prompt 段编码 + 原样拼采样 token，返回可直接喂模型的 InputFeature。 + + labels 只盖住 ``tokens`` 那一段（prompt 段全 -100），与采样端真实生成的 token 逐一对应。 + ``template`` 默认用 skill 模型那份；executor 侧打分（E14 的 logP reward）要传自己的。 + """ + tmpl = template or _ENCODE_TEMPLATE + assert tmpl is not None, 'call set_encode_template() before building train features' + feat = tmpl.encode({'messages': [dict(m) for m in prompt_messages]}, add_generation_prompt=True) + return tmpl.concat_input_feature(feat, [int(t) for t in tokens]) + + def _train_trajectory(rec): - """Rebuild the query-only skill-gen prompt (train/inference match) + response. - GRPO records carry the full generated response; SFT records carry only the - cleaned block. key_rounds selects the final assistant turn.""" + """query-only skill-gen prompt + 采样产出。 + + GRPO 记录带 ``tokens``（采样端 vLLM 真实吐出的 token id），直接拼成 InputFeature； + SFT 记录的 response 是程序合成的 ```` 文本（本来就不是采样产出、没有对应 token）， + 只能走 messages 编码，``key_rounds`` 标出最后一轮为唯一可训练区。 + """ msgs = _skillgen_prompt(rec['problem'])['messages'] + if rec.get('tokens'): + return build_train_feature(msgs, rec['tokens']) return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], - 'user_data': {'key_rounds': [len(msgs)]}} + 'user_data': pack_user_data({'key_rounds': [len(msgs)]})} def _train_step(skill_model, ref_model, ckpt, samples, args): @@ -1321,10 +1480,12 @@ def _train_step(skill_model, ref_model, ckpt, samples, args): # 过滤空/纯空白 response：其可训练 token 为 0，会让持有它的 DP rank 跳过 backward， # 与对端 all-reduce 失步 → NCCL 死锁（find_unused_parameters=False）。高熵采样偶发首 token 即 EOS。 n_in = len(samples) - samples = [rec for rec in samples if (rec.get('response') or '').strip()] + samples = [rec for rec in samples + if (rec.get('tokens') if rec.get('tokens') is not None else (rec.get('response') or '').strip())] n_empty = n_in - len(samples) trajs = [_train_trajectory(rec) for rec in samples] advs = [float(rec['advantage']) for rec in samples] + smp_all = [list(rec.get('logprobs') or []) for rec in samples] # drop_last 到 TRAIN_DP 整数倍：每个 micro（末尾那个可短于 sft）仍能被 dp 均分，零 padding 假样本。 # 只丢尾部 ≤ dp-1 条真样本；若整批不足 dp（n 0 else n mini = max(sft, (mini // sft) * sft) multi_step = mini < n - micro_ref, micro_old = [], [] + micro_ref, micro_old, micro_smp = [], [], [] for i in range(0, n, sft): mb = trajs[i:i + sft] micro_ref.append(ref_model.forward_only(inputs=mb).get('logps')) micro_old.append(skill_model.forward_only(inputs=mb).get('logps') if multi_step else None) + # 采样端 logprob，只给 GRPOMetric 做对账（sampler_logp_mae / sampler_token_delta）。 + # 整个 micro 都带齐了才传：SFT 样本的 response 是合成文本、根本没有采样 logprob， + # 混进去只会让 token_delta 无法解释（它的语义是「应恒为 0」）。 + smp = [smp_all[j] for j in range(i, min(i + sft, n))] + micro_smp.append(smp if all(smp) else None) micro, n_steps = 0, 0 for ms in range(0, n, mini): for i in range(ms, min(ms + mini, n), sft): k = i // sft skill_model.forward_backward(inputs=trajs[i:i + sft], advantages=advs[i:i + sft], - old_logps=micro_old[k], ref_logps=micro_ref[k]) + old_logps=micro_old[k], ref_logps=micro_ref[k], + sampler_logps=micro_smp[k]) micro += 1 skill_model.clip_grad_and_step() n_steps += 1 @@ -1382,16 +1549,20 @@ def process_chunk(base_sampler, skill_sampler, chunk, ci, base_dp, skill_dp, arg sg_out = _run_samples(skill_sampler, [_skillgen_prompt(r['problem']) for r in chunk], args.n_skills, args.skill_max_tokens, skill_dp, temperature=args.skill_gen_temperature, top_p=args.skill_gen_top_p, - top_k=args.skill_gen_top_k) + top_k=args.skill_gen_top_k, logprobs=1) for r, seqs in zip(chunk, sg_out): for s in seqs: resp = _clean_text(getattr(s, 'decoded', '') or '') block = _extract_skill(resp) or '' + # tokens = 采样端真实吐出的 token id，训练样本只能用它拼（见 build_train_feature）。 + # logprobs 同长，只给 GRPOMetric 做采样/训练对账（见 sampler_logprobs），不进 loss。 + _toks = [int(t) for t in (getattr(s, 'tokens', None) or [])] cand = {'skills': block, 'response': resp, 'parseable': bool(block), 'leaked': None, 'with_pass': None, 'reward': None, 'rolls': [], - 'advantage': 0.0, 'kept': False, + 'advantage': 0.0, 'kept': False, 'tokens': _toks, + 'logprobs': sampler_logprobs(s), 'skillgen_stop': getattr(s, 'stop_reason', None), - 'skillgen_tokens': len(getattr(s, 'tokens', None) or [])} + 'skillgen_tokens': len(_toks)} r['_cands'].append(cand) if block: flat.append((r, cand)) @@ -1440,6 +1611,29 @@ def process_chunk(base_sampler, skill_sampler, chunk, ci, base_dp, skill_dp, arg if c['reward'] is None: c['reward'] = 0.0 + # ⭐ 训练侧 no-skill baseline（seam 对齐，2026-08-02）。SEAM 只在第 1 个训练步跑一次 + # （ray_trainer.py:1461 `if self.global_steps == 1`，注释写明是提速、reward/梯度不受影响， + # 后续 step 的 step_summary.lift 为 None），所以 SEAM 的 train lift 只有 step1 一个点 + # （0.7891-0.6338=+0.1553）。这里逐条照抄：仅 ci==0、direct prompt、T=0 + # （= SEAM use_experience=False 分支）。 + # + # 每题必须跑 n_skills 次，不能跑 1 次再广播：SEAM 把**整个 1024 行 batch**送进 + # generate_sequences_as_grm_baseline，同一道题的 8 行是 8 个独立请求。executor 虽然是 + # greedy，vLLM 的批内非确定性（左 padding 长度随 batch 变、chunked prefill 的规约顺序） + # 让同一 prompt 的 8 次结果并不总相同 —— step1 dump 实测 128 题里有 10 题的 8 次 + # baseline_acc 不全同。所以"同题 8 条结果相同、取 1 次即可"这个旧假设是错的，按题 + # 展开取均值才与 SEAM 的 np.mean(baseline_acc) 同口径。 + if _ALIGN_MODE == 'seam' and ci == 0 and chunk: + K_b = max(1, int(args.n_skills)) + b_pairs = [r for r in chunk for _ in range(K_b)] + b_out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in b_pairs], + 1, args.max_tokens, base_dp, temperature=0.0) + b_rolls = _parse_many([(_first_seq(seqs), r['reference_answer']) + for r, seqs in zip(b_pairs, b_out)]) + for i, r in enumerate(chunk): + hits = [1.0 if x['correct'] else 0.0 for x in b_rolls[i * K_b:(i + 1) * K_b]] + r['_train_baseline_pass'] = _mean(hits) if hits else None + _assign_advantages(chunk, args) # SEAM/verl 对齐：每个 dataloader batch 都进入 actor update。零 adv 候选 PG 贡献 0， @@ -1452,6 +1646,9 @@ def process_chunk(base_sampler, skill_sampler, chunk, ci, base_dp, skill_dp, arg grpo.append({'problem': r['problem'], 'reference_answer': r['reference_answer'], 'data_id': r.get('data_id', ''), 'response': c['response'], 'skills': c['skills'], 'advantage': c['advantage'], + # tokens 是训练样本的唯一来源；response 只留给 reward/审计（见 _train_trajectory） + 'tokens': c.get('tokens') or [], + 'logprobs': c.get('logprobs') or [], 'kept': c['kept'], 'reward': c['reward'], 'sft': False}) buffer_a = _collect_buffer_a(chunk, args) return _full_records(chunk, ci), _chunk_summary(chunk, ci), grpo, buffer_a @@ -1496,6 +1693,22 @@ def _std(xs): return float(torch.std(torch.tensor(xs, dtype=torch.float32)).item()) +def _pstd(xs): + """总体标准差（ddof=0）—— 报表用。 + + SEAM 的 step_summary 用 ``np.std``（ddof=0）算 reward_std / group_reward_std_mean + （ray_trainer.py:654,671），而 :func:`_std` 是 ``torch.std``（ddof=1）。组内只有 8 条时 + 两者差 sqrt(8/7)=1.069，足以把 group_reward_std_mean 抬高 ~0.01（c0 实测 + 0.14993 vs SEAM 0.14074；改 ddof=0 后为 0.14025）。注意分开：**advantage 的组内 + std 必须继续用 ddof=1**，verl 的 compute_grpo_outcome_advantage 用的也是 + ``torch.std`` 默认 unbiased，两边本来就一致，改了反而会把梯度弄不一致。 + """ + if len(xs) < 2: + return 0.0 + m = sum(xs) / len(xs) + return (sum((x - m) ** 2 for x in xs) / len(xs)) ** 0.5 + + def _chunk_summary(chunk, ci): all_cands = [c for r in chunk for c in r['_cands']] cands = [c for c in all_cands if c['parseable']] @@ -1509,7 +1722,8 @@ def _chunk_summary(chunk, ci): continue groups += 1 all_rewards.extend(rewards) - v = _std(rewards) + # 报表口径对齐 SEAM 的 np.std（ddof=0），见 _pstd。advantage 那边仍用 ddof=1。 + v = _pstd(rewards) group_vars.append(v) if v < 1e-9: zero_grad += 1 @@ -1527,6 +1741,10 @@ def _chunk_summary(chunk, ci): # = step_summary 的 withskill_pass = swanlab 的 train/with_skill_accuracy。 # 只有 seam 模式会给 unparseable 候选跑 executor，其余臂这个值等于 candidate_withskill_pass。 pass_all = [c['with_pass'] for c in all_cands if c['with_pass'] is not None] + # 训练侧 baseline（仅 seam 模式的 chunk 0 有），按候选展开与 SEAM 的 np.mean(baseline_acc) 同口径。 + b_all = [r['_train_baseline_pass'] for r in chunk if r.get('_train_baseline_pass') is not None + for _c in r['_cands']] + base_pass = _mean(b_all) if b_all else None return { 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), 'n_generated': len(all_cands), 'n_candidates_parseable': len(cands), @@ -1535,7 +1753,7 @@ def _chunk_summary(chunk, ci): 'leak_rate': (sum(1 for c in cands if c['leaked']) / len(cands)) if cands else 0.0, 'n_train_samples': n_train, 'n_groups': groups, 'zero_grad_frac': (zero_grad / groups) if groups else 0.0, - 'reward_mean': _mean(all_rewards), 'reward_std': _std(all_rewards), + 'reward_mean': _mean(all_rewards), 'reward_std': _pstd(all_rewards), 'group_reward_std_mean': _mean(group_vars), 'skill_tokens_mean': _mean([c.get('skillgen_tokens') or 0 for c in cands]), 'skill_chars_mean': _mean([len(c['skills']) for c in cands]), @@ -1543,6 +1761,8 @@ def _chunk_summary(chunk, ci): 'candidate_withskill_pass': _mean([c['with_pass'] for c in scored]), 'withskill_pass_all_cands': _mean(pass_all), 'n_exec_cands': len(pass_all), + 'baseline_pass_train': base_pass, + 'lift_train': (_mean(pass_all) - base_pass) if base_pass is not None else None, 'withskill_trunc_frac': (trunc / len(ws_rolls)) if ws_rolls else 0.0, 'termination_rate_withskill': _mean([1.0 if x['terminated'] else 0.0 for x in ws_rolls]), } @@ -1680,10 +1900,20 @@ def init_components(args): skill_model.set_template(Template, model_id=MODEL_ID, enable_thinking=_think, max_length=args.max_model_len, truncation_strategy='delete') skill_model.set_processor(InputProcessor, padding_free=False) + # 客户端同配置副本：训练样本的 prompt 段在这里编码，response 段直接拼采样返回的 + # token（见 build_train_feature）。skill_ablate/trainer.py 走自己那一句，两边不互干扰。 + set_encode_template(Template(model_id=MODEL_ID, enable_thinking=_think, + max_length=args.max_model_len, truncation_strategy='delete')) # loss 统一用 SEAM 对齐的 SEAMBNPOLoss（verl PPO clip + low_var_kl + token-mean）； # v2/seam 两模式一致，不再随 align-mode 变。 _loss_cls = 'SEAMBNPOLoss' skill_model.set_loss(_loss_cls, epsilon=args.grpo_epsilon, beta=args.kl_beta) + # RL 异常 token 监控：除了 ratio/kl/entropy/clip 那一套，GRPOMetric 还会吐 + # train/logp_min、train/logp_frac_lt_10、train/sampler_logp_mae、train/sampler_token_delta + # —— 训练序列一旦混进模型没生成过的 token（编码错位、模板凭空补的 EOS/空思考块）， + # 前两条会直接炸、后两条直接违反断言；而行为层指标（format/acc/reward）完全看不出来。 + # temperature 不传：skill-gen 就是 T=1，与采样端 logprob 取值口径天然一致。 + skill_model.add_metric('GRPOMetric', is_training=True, epsilon=args.grpo_epsilon) skill_model.set_optimizer('AdamW', lr=args.lr) # 对齐 SEAM：恒定 lr（无 warmup、无 decay）。SEAM 用 get_constant_schedule_with_warmup( # num_warmup_steps=0)+warmup_style=constant，全程恒定 1e-6。这里直接不设 scheduler， @@ -1702,9 +1932,16 @@ def init_components(args): ref_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) def _sampler(group, world, enable_thinking): + # enable_prefix_caching 必须开：verl 的 vllm_rollout_spmd 把它硬编码成 True + # （vllm_rollout_spmd.py:182），actor rollout 与 grm/executor rollout 共用这个引擎， + # 所以 SEAM 两条采样通路全程带前缀缓存。twinkle 默认是 False，实测差别不只是速度： + # 同一批里 8 个完全相同的 baseline prompt，缓存关掉时逐条重算、结果 128/128 题全同 + # （608=76x8 精确整除），缓存打开后首条算新 KV、后 7 条复用缓存 KV，数值路径不同， + # SEAM 那边就有 10/128 题的 8 次贪心结果不全同。要和 SEAM 同口径就得同样开着。 s = vLLMSampler(model_id=MODEL_ID, engine_args={'gpu_memory_utilization': GPU_MEM, - 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, + 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1, + 'enable_prefix_caching': True}, device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), remote_group=group) s.set_template(Template, model_id=MODEL_ID, enable_thinking=enable_thinking, max_length=args.max_model_len) @@ -1865,6 +2102,12 @@ def _swan_metrics(summary, log): d['train/with_skill_accuracy'] = summary['withskill_pass_all_cands'] d['acc/reward_mean'] = summary['reward_mean'] d['skill/format_rate'] = summary['parse_rate'] + if summary.get('baseline_pass_train') is not None: + # SEAM 只在 step1 跑训练侧 baseline，所以这两条也只在 chunk 0 有值（与 SEAM 同步）。 + d['acc/baseline_pass'] = summary['baseline_pass_train'] + d['acc/lift'] = summary['lift_train'] + d['train/baseline_accuracy'] = summary['baseline_pass_train'] + d['train/lift'] = summary['lift_train'] if log: d['train/n_grpo'] = log['n_grpo'] d['train/n_sft'] = log['n_sft'] @@ -1874,6 +2117,9 @@ def _swan_metrics(summary, log): if k.startswith('learning rate'): if 'group 1' in k: d['train/lr'] = float(v) + elif k.startswith('train/'): + # GRPOMetric 等已经自带 train/ 前缀，不能再套一层。 + d[k.replace(' ', '_')] = float(v) else: d[f'train/{k.replace(" ", "_")}'] = float(v) return d diff --git a/src/twinkle/data_format/__init__.py b/src/twinkle/data_format/__init__.py index c93bebd2d..d51f09dfa 100644 --- a/src/twinkle/data_format/__init__.py +++ b/src/twinkle/data_format/__init__.py @@ -3,4 +3,4 @@ from .message import Message, Tool, ToolCall from .output import LossOutput, ModelOutput from .sampling import SampledSequence, SampleResponse, SamplingParams -from .trajectory import Trajectory, pack_value, user_data_get +from .trajectory import Trajectory, pack_user_data, pack_value, user_data_get diff --git a/src/twinkle/data_format/trajectory.py b/src/twinkle/data_format/trajectory.py index 992df28d1..5044bcd02 100644 --- a/src/twinkle/data_format/trajectory.py +++ b/src/twinkle/data_format/trajectory.py @@ -1,6 +1,7 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import json import sys +from collections.abc import Mapping from typing import Any, List, Optional, Tuple, Union from .message import Message, Tool @@ -28,8 +29,27 @@ def pack_value(value: Any) -> str: return json.dumps(value, ensure_ascii=False, default=str) +def pack_user_data(values: Any) -> List[Tuple[str, str]]: + """Build a canonical ``user_data`` payload from a plain mapping. + + ``user_data`` must be a list of ``(key, json_string)`` pairs: a dict cannot be written to + PyArrow (its struct schema would differ from shard to shard), and readers go through + :func:`user_data_get`, which only understands the packed form — a dict is silently ignored + there, so e.g. ``key_rounds`` would be dropped without any error. Always pack before + attaching to a trajectory. + """ + if values is None: + return [] + if isinstance(values, Mapping): + return [(k, v if isinstance(v, str) else pack_value(v)) for k, v in values.items()] + return [(k, v if isinstance(v, str) else pack_value(v)) for k, v in values] + + def user_data_get(items: Any, key: str, default: Any = None) -> Any: """Look up the first value matching ``key`` in packed user_data, decoded.""" + if isinstance(items, Mapping): + raise TypeError('user_data must be a list of (key, json_string) pairs, got a mapping. ' + 'Wrap it with twinkle.data_format.pack_user_data().') if not isinstance(items, list): return default for entry in items: diff --git a/src/twinkle/metric/grpo.py b/src/twinkle/metric/grpo.py index bd85aab67..ea22d005d 100644 --- a/src/twinkle/metric/grpo.py +++ b/src/twinkle/metric/grpo.py @@ -53,6 +53,14 @@ def reset(self): self.clip_n_total: float = 0.0 self.high_kl_records: list = [] self._gsi_cursor: int = 0 + # 异常 token 探针：logp 尾部统计 + 与采样端的对账。 + self.min_new_logp: float = 0.0 + self.n_logp_lt5: int = 0 + self.n_logp_lt10: int = 0 + self.sum_sampler_abs: float = 0.0 + self.max_sampler_abs: float = 0.0 + self.n_sampler_matched: int = 0 + self.n_sampler_given: int = 0 @staticmethod def _as_mb_list(logps_val) -> Optional[List]: @@ -111,6 +119,7 @@ def _accumulate_mb( entropies: Optional['torch.Tensor'] = None, adv_slice: Any = None, gsi_base: int = 0, + sampler_slice: Any = None, ) -> int: """Reduce one microbatch into ``self.sum_*`` counters. @@ -150,14 +159,22 @@ def _accumulate_mb( # Rescaling keeps ``logp_diff`` / ``approx_kl`` unchanged because # both new and old logps receive the same multiplier. scale = self.temperature - logps_f = logps.float() - if scale > 0.0 and scale != 1.0: - logps_f = logps_f * scale + logps_raw = logps.float() + logps_f = logps_raw * scale if (scale > 0.0 and scale != 1.0) else logps_raw mask_f = mask.float() self.n_tokens += n_tok self.sum_new += float((logps_f * mask_f).sum().item()) + cur_min = float(logps_raw.masked_fill(~mask, 0.0).min().item()) + if cur_min < self.min_new_logp: + self.min_new_logp = cur_min + self.n_logp_lt5 += int(((logps_raw < -5.0) & mask).sum().item()) + self.n_logp_lt10 += int(((logps_raw < -10.0) & mask).sum().item()) + + if sampler_slice is not None: + self._accumulate_sampler(logps_raw, sampler_slice, mask, mask_f) + # Entropy is loss-type-agnostic; aligned to logps shape by the model forward. if entropies is not None and torch.is_tensor(entropies) and entropies.numel() > 0: ent_f = entropies.float() @@ -231,6 +248,36 @@ def _accumulate_clip( self.sum_clip_high += float((is_high.float() * mask_f).sum().item()) self.clip_n_total += float(mask_f.sum().item()) + def _accumulate_sampler( + self, + logps_raw: 'torch.Tensor', + sampler_slice: Any, + mask: 'torch.Tensor', + mask_f: 'torch.Tensor', + ) -> None: + smp = align_logps_to_mask(sampler_slice, mask, logps_raw.dtype) + if smp is None: + return + rows = sampler_slice if isinstance(sampler_slice, (list, tuple)) else [sampler_slice] + lens = [] + for row in rows: + try: + lens.append(int(len(row))) + except TypeError: + lens.append(1) + self.n_sampler_given += sum(lens) + cov = align_logps_to_mask([[1.0] * n for n in lens], mask, logps_raw.dtype) + if cov is None: + return + valid = cov * mask_f + diff_abs = (logps_raw - smp).abs() * valid + self.sum_sampler_abs += float(diff_abs.sum().item()) + self.n_sampler_matched += int(valid.sum().item()) + if diff_abs.numel() > 0: + cur = float(diff_abs.max().item()) + if cur > self.max_sampler_abs: + self.max_sampler_abs = cur + def accumulate( self, inputs: Union[InputFeature, List[InputFeature]], @@ -238,6 +285,7 @@ def accumulate( *, old_logps: Any = None, advantages: Any = None, + sampler_logps: Any = None, **kwargs, ): import torch @@ -268,6 +316,9 @@ def accumulate( flat_adv: Optional[List] = None if advantages is not None and isinstance(advantages, (list, tuple)): flat_adv = list(advantages) + flat_sampler: Optional[List] = None + if sampler_logps is not None and isinstance(sampler_logps, (list, tuple)): + flat_sampler = list(sampler_logps) cursor = 0 n_mb = min(len(inputs_list), len(logps_list)) @@ -305,8 +356,10 @@ def accumulate( old_slice = None adv_mb = flat_adv[cursor:cursor + num_seq_est] if flat_adv is not None else None + smp_mb = flat_sampler[cursor:cursor + num_seq_est] if flat_sampler is not None else None gsi_base = self._gsi_cursor - advanced = self._accumulate_mb(labels, logps_mb, old_slice, ent_mb, adv_mb, gsi_base=gsi_base) + advanced = self._accumulate_mb(labels, logps_mb, old_slice, ent_mb, adv_mb, + gsi_base=gsi_base, sampler_slice=smp_mb) self._gsi_cursor += advanced cursor += advanced @@ -326,6 +379,13 @@ def calculate(self) -> Dict[str, Any]: 'sum_clip_low': self.sum_clip_low, 'sum_clip_high': self.sum_clip_high, 'clip_n_total': self.clip_n_total, + 'min_new_logp': self.min_new_logp, + 'n_logp_lt5': self.n_logp_lt5, + 'n_logp_lt10': self.n_logp_lt10, + 'sum_sampler_abs': self.sum_sampler_abs, + 'max_sampler_abs': self.max_sampler_abs, + 'n_sampler_matched': self.n_sampler_matched, + 'n_sampler_given': self.n_sampler_given, }] all_results = self.gather_results(local) @@ -340,6 +400,10 @@ def calculate(self) -> Dict[str, Any]: results: Dict[str, Any] = { 'train/policy_confidence': math.exp(mean_new), 'train/mean_new_logp': mean_new, + 'train/n_trainable_tokens': n_total, + 'train/logp_min': min(r.get('min_new_logp', 0.0) for r in all_results), + 'train/logp_frac_lt_5': sum(r.get('n_logp_lt5', 0) for r in all_results) / n_total, + 'train/logp_frac_lt_10': sum(r.get('n_logp_lt10', 0) for r in all_results) / n_total, } if any(r['has_old'] for r in all_results): mean_old = sum(r['sum_old'] for r in all_results) / n_total @@ -365,6 +429,14 @@ def calculate(self) -> Dict[str, Any]: results['train/clip_ratio_high'] = sum_high / clip_n results['train/clip_ratio'] = (sum_low + sum_high) / clip_n + # 采样端对账（只在调用方传了 sampler_logps 时出现）。两条都是断言型指标： + # sampler_logp_mae 应在引擎精度量级（bf16 约 1e-2），sampler_token_delta 应恒为 0。 + n_smp = sum(r.get('n_sampler_matched', 0) for r in all_results) + if n_smp > 0: + results['train/sampler_logp_mae'] = sum(r.get('sum_sampler_abs', 0.0) for r in all_results) / n_smp + results['train/sampler_logp_max_abs'] = max(r.get('max_sampler_abs', 0.0) for r in all_results) + results['train/sampler_token_delta'] = n_total - sum(r.get('n_sampler_given', 0) for r in all_results) + # Underscore-prefixed key bypasses swanlab numeric coercion; script can pop and consume. if self.high_kl_records: results['_high_kl_records'] = list(self.high_kl_records) diff --git a/src/twinkle/utils/torch_utils.py b/src/twinkle/utils/torch_utils.py index 84a335852..7c376e9c6 100644 --- a/src/twinkle/utils/torch_utils.py +++ b/src/twinkle/utils/torch_utils.py @@ -247,17 +247,22 @@ def pad_and_stack_tensors(tensors: List['torch.Tensor'], pad_value: float = -200 t = t.unsqueeze(0) expanded_tensors.append(t) - max_shape = [] - for dim in range(max_ndim): - max_shape.append(max(t.shape[dim] for t in expanded_tensors)) + # dim 0 是 concat 的拼接维，torch.cat 本来就不要求它对齐 —— 把它也 pad 到最大值会**凭空造出 + # 不存在的样本行**（例如 dp rank0 收 3 行、rank1 收 2 行时，结果是 3+3=6 行而不是 5 行，多出来 + # 的那行全是 pad_value）。这些假行流进下游后：损失侧 GRPOLoss._pad_and_align_to_batch 靠 + # `data[i] for i in range(batch_size)` 把它们丢掉所以侥幸无害，但指标侧 align_logps_to_mask 是 + # 严格判等，行数一多就整步跳过 ratio/kl/clip（日志里的 `old_logps shape (3, N) does not match + # logps_mb shape (2, N)` 就是它）。所以 concat 时只对齐 dim>=1，stack 时才需要全维对齐。 + pad_from = 1 if concat else 0 + max_shape = [max(t.shape[dim] for t in expanded_tensors) for dim in range(max_ndim)] padded_tensors = [] for t in expanded_tensors: - if list(t.shape) == max_shape: + if all(t.shape[dim] == max_shape[dim] for dim in range(pad_from, max_ndim)): padded_tensors.append(t) else: pad_params = [] - for dim in range(max_ndim - 1, -1, -1): + for dim in range(max_ndim - 1, pad_from - 1, -1): pad_params.extend([0, max_shape[dim] - t.shape[dim]]) padded = torch.nn.functional.pad(t, pad_params, value=pad_value) padded_tensors.append(padded) diff --git a/tests/loss/test_bnpo_token_mean.py b/tests/loss/test_bnpo_token_mean.py index 9e9ddd6ff..3514498ed 100644 --- a/tests/loss/test_bnpo_token_mean.py +++ b/tests/loss/test_bnpo_token_mean.py @@ -1,7 +1,16 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""BNPO token-mean fix: 'global' scope must be invariant to how a batch is split into -micro/dp groups (strict global token-mean); 'micro' scope reproduces the pre-fix -double-average bias that skews toward short responses. +"""BNPO ``token_mean_scope`` semantics. + +'micro' (the DEFAULT) reproduces verl/SEAM: ``masked_mean`` inside each micro-batch, then +an equal-weighted average across micro/dp groups (see verl/workers/actor/dp_actor.py -- +``pg_loss = agg_loss(..., 'token-mean')`` per micro, then ``* 1/gradient_accumulation`` +before ``backward()``). It is deliberately NOT split-invariant. + +'global' is the strict, split-invariant token-mean. It is available but NOT the default: +with group-relative advantages the token-weighted mean does not cancel (it equals +-cov(len, A)/mean(len)), which on skill2lora E13 produced a ~100x stronger coherent +"emit fewer tokens" gradient than verl and collapsed the response length. See BNPOLoss's +docstring for the measurements. The framework combines groups per LossOutput semantics (transformers.py / metric/loss.py): effective_loss = Σ_g loss_g / Σ_g num_tokens_g @@ -54,29 +63,36 @@ def test_global_is_split_invariant(): assert abs(whole - split) < 1e-6 -def test_micro_is_biased_and_split_dependent(): +def test_micro_matches_verl_equal_weighted_micro_means(): ptl, mask = _fixture() loss = BNPOLoss(token_mean_scope='micro') whole = _combine(loss, ptl, mask, [[0, 1]]) # one group -> token-mean 0.625 split = _combine(loss, ptl, mask, [[0], [1]]) # per-group means (1.0, 0.5) -> 0.75 assert abs(whole - 0.625) < 1e-6 - assert abs(split - 0.75) < 1e-6 # short response over-weighted - assert split > whole # bias toward short is real - # and it disagrees with the correct global answer + assert abs(split - 0.75) < 1e-6 # equal weight per micro, as verl does + # Not split-invariant, by design: this is exactly verl's behaviour. assert abs(split - 0.625) > 1e-3 -def test_seam_inherits_global_by_default(): +def test_seam_defaults_to_micro_like_verl(): from twinkle.loss.grpo import SEAMBNPOLoss seam = SEAMBNPOLoss(epsilon=0.2, beta=0.001) - assert seam.token_mean_scope == 'global' + assert seam.token_mean_scope == 'micro' + assert seam.reduction == 'mean' # display layer must not divide again ptl, mask = _fixture() split = _combine(seam, ptl, mask, [[0], [1]]) - assert abs(split - 0.625) < 1e-6 + assert abs(split - 0.75) < 1e-6 + + +def test_global_reports_sum_reduction_for_display(): + """'global' returns a token SUM, so LossMetric must be told reduction='sum' or the + logged loss is inflated by the token count.""" + assert BNPOLoss(token_mean_scope='global').reduction == 'sum' if __name__ == '__main__': test_global_is_split_invariant() - test_micro_is_biased_and_split_dependent() - test_seam_inherits_global_by_default() - print('OK: global split-invariant; micro reproduces the biased double-average') + test_micro_matches_verl_equal_weighted_micro_means() + test_seam_defaults_to_micro_like_verl() + test_global_reports_sum_reduction_for_display() + print('OK: micro (default) == verl equal-weighted micro-means; global is split-invariant') diff --git a/tests/utils/test_utils.py b/tests/utils/test_utils.py index 641da8ad1..b9b1aa97a 100644 --- a/tests/utils/test_utils.py +++ b/tests/utils/test_utils.py @@ -151,7 +151,14 @@ def test_same_shape(self): def test_different_length(self): tensors = [torch.randn(3), torch.randn(5)] result = pad_and_stack_tensors(tensors, pad_value=0) - assert result.shape == (10, ) # padded to max length then concat + # concat 沿 dim 0，而 dim 0 就是拼接维：不能 pad，否则会插入不存在的元素（旧行为给 (10,)） + assert result.shape == (8, ) + + def test_concat_does_not_pad_batch_dim(self): + tensors = [torch.randn(3, 4), torch.randn(2, 6)] + result = pad_and_stack_tensors(tensors, pad_value=0) + # 只对齐 seq 维；行数必须是 3+2，不能被拉成 3+3 + assert result.shape == (5, 6) def test_different_length_stack(self): tensors = [torch.randn(3), torch.randn(5)] From 3f3e76341c97bf53d5eaff17f308ef55ee68aa17 Mon Sep 17 00:00:00 2001 From: root Date: Sun, 2 Aug 2026 18:40:46 +0800 Subject: [PATCH 34/60] fix(metric): align old_logps to the mask when it is wider than the local logps GRPOMetric dropped ratio/kl for most steps with a warning like old_logps shape (2, 7754) does not match logps_mb shape (2, 6447) Root cause: old_logps comes from a ref/old model forward, whose sequence dim is padded to the max over the WHOLE micro batch (before the dp split), while the training forward pads only to the local rank's own max. Whenever the longest sample of the micro batch lives on another rank, old is strictly wider -- so the exact-shape gate in GRPOMetric threw the step away. Row counts always matched, which is why this is a different bug from the batch-dim one fixed earlier. Gradients were never affected: GRPOLoss._pad_and_align_to_batch already had the full-sequence branch, so the loss always indexed the right tokens. Only the panel went blank. - align_logps_to_mask: add the same full-sequence branch as the loss (slice to seq_len, then index by mask) instead of taking the first n_pos values, which would have read prompt positions and misaligned every ratio. Response-only rows (sampler_logps / advantages) keep their current behaviour. - GRPOMetric: accept any old_logps tensor whose row count matches and whose sequence width is >=; keep warning + skipping on row mismatch, since that one really is a bug that must stay visible. - tests: 4 regressions, including a cross-check that the metric and the loss align to byte-identical tensors. --- src/twinkle/metric/grpo.py | 23 ++++++++++++++----- src/twinkle/utils/transformers_utils.py | 22 ++++++++++++++++-- tests/metric/test_metrics.py | 29 ++++++++++++++++++++++++ tests/utils/test_utils.py | 30 +++++++++++++++++++++++++ 4 files changed, 96 insertions(+), 8 deletions(-) diff --git a/src/twinkle/metric/grpo.py b/src/twinkle/metric/grpo.py index ea22d005d..7c2e98fd9 100644 --- a/src/twinkle/metric/grpo.py +++ b/src/twinkle/metric/grpo.py @@ -339,18 +339,29 @@ def accumulate( if flat_old is not None: old_slice = flat_old[cursor:cursor + num_seq_est] elif old_logps is not None and hasattr(old_logps, 'shape'): - # Uncommon: aligned global tensor. Only honour when it - # exactly matches the single-mb shape; otherwise drop. + # Aligned tensor from a ref/old model forward. Its seq width is the max over + # the WHOLE micro batch (padded before the dp split), while ``logps_mb`` is + # padded only to this rank's own max — so old is routinely LONGER, and + # requiring exact equality here threw away ratio/kl on most steps whenever the + # longest sample of the micro batch lived on another rank. The loss never had + # this problem (GRPOLoss._pad_and_align_to_batch has the full-sequence branch), + # so the gradients were right all along and only the panel went blank. + # align_logps_to_mask now shares that branch; accept anything it can align. import torch as _torch # noqa: F811 - if _torch.is_tensor(old_logps) and old_logps.shape == logps_mb.shape: + usable = (_torch.is_tensor(old_logps) + and old_logps.dim() == logps_mb.dim() + and old_logps.shape[0] == logps_mb.shape[0] + and old_logps.shape[-1] >= logps_mb.shape[-1]) + if usable: old_slice = old_logps else: if mb_idx == 0: # Warn once per accumulate call (not per mb) to avoid log spam. old_shape = tuple(old_logps.shape) if _torch.is_tensor(old_logps) else 'unknown' - logger.warning(f'GRPOMetric: old_logps shape {old_shape} does not match ' - f'logps_mb shape {tuple(logps_mb.shape)}; ratio/kl metrics will ' - f'be skipped for this step.') + logger.warning(f'GRPOMetric: old_logps shape {old_shape} cannot be aligned to ' + f'logps_mb shape {tuple(logps_mb.shape)} (row count must match and ' + f'seq width must be >=); ratio/kl metrics will be skipped for ' + f'this step.') old_slice = None else: old_slice = None diff --git a/src/twinkle/utils/transformers_utils.py b/src/twinkle/utils/transformers_utils.py index 12963c13f..73a71752f 100644 --- a/src/twinkle/utils/transformers_utils.py +++ b/src/twinkle/utils/transformers_utils.py @@ -14,6 +14,20 @@ def align_logps_to_mask( mask: 'torch.Tensor', dtype: 'torch.dtype', ) -> Optional['torch.Tensor']: + """Scatter ragged per-sample values onto the trainable positions of ``mask``. + + Two per-sample forms are supported, disambiguated by length exactly like + ``GRPOLoss._pad_and_align_to_batch`` (the two MUST agree, otherwise the metric + reports ratios computed on different tokens than the loss optimises): + * Response-only form (``len == mask[i].sum()``): scattered directly. + * Full-sequence form (``len >= mask.shape[1]``, right-padded): sliced to + ``seq_len`` and indexed by ``mask[i]`` first. This is what a ref/old model + forward returns; its padding width is the max over the WHOLE micro batch + before the dp split, so it is routinely LONGER than the local ``logps`` + (which is padded only to the local rank's max). Taking ``vals[:n_pos]`` + instead would read prompt positions and silently misalign every ratio. + Anything shorter than both is unusable and returns None rather than guessing. + """ import torch device = mask.device @@ -40,8 +54,12 @@ def align_logps_to_mask( result[i, pos] = float(sample) continue vals = torch.as_tensor(sample, dtype=dtype, device=device).flatten() - n = min(len(pos), int(vals.numel())) - if n > 0: + n = int(vals.numel()) + if n == len(pos): + result[i, pos] = vals + elif n >= seq_len: + result[i, pos] = vals[:seq_len][mask[i]] + elif n > 0: result[i, pos[:n]] = vals[:n] return result diff --git a/tests/metric/test_metrics.py b/tests/metric/test_metrics.py index e4651ee86..e98846a53 100644 --- a/tests/metric/test_metrics.py +++ b/tests/metric/test_metrics.py @@ -334,6 +334,35 @@ def test_grpo_metric_entropy(self): result = m.calculate() assert 'train/entropy' in result + def test_grpo_metric_old_logps_wider_than_logps(self): + """old_logps 来自 forward_only，序列维 pad 到整个 micro batch 的最大长度（dp split 之前）， + 而 logps 只 pad 到本 rank 的最大长度 —— old 比 new 宽是常态，不能因此丢掉 ratio/kl。 + 取值必须落在 mask 位上（不是行首 N 个），所以 mean_old_logp 是错位的判据。 + """ + m = _no_dist_metric(GRPOMetric) + labels = torch.tensor([[-100, -100, 10, 11, 12, 13], + [-100, -100, -100, -100, 20, 21]]) + logps = torch.zeros(2, 6) + logps[0, 2:6] = torch.tensor([-1.1, -2.1, -3.1, -4.1]) + logps[1, 4:6] = torch.tensor([-5.1, -6.1]) + old_logps = torch.zeros(2, 9) # 右 pad 到 9 > 6 + old_logps[0, 2:6] = torch.tensor([-1.0, -2.0, -3.0, -4.0]) + old_logps[1, 4:6] = torch.tensor([-5.0, -6.0]) + m.accumulate({'labels': labels}, {'logps': logps}, old_logps=old_logps) + result = m.calculate() + assert 'train/approx_kl' in result + assert abs(result['train/mean_old_logp'] - (-3.5)) < 1e-6 + assert abs(result['train/logp_diff_mean'] - (-0.1)) < 1e-6 + + def test_grpo_metric_old_logps_row_mismatch_skipped(self): + """行数不匹配是另一类真 bug（凭空 pad 出的假样本行），必须继续被丢弃而不是硬对齐。""" + m = _no_dist_metric(GRPOMetric) + labels = torch.tensor([[-100, 10, 11, 12]]) + logps = torch.randn(1, 4) + m.accumulate({'labels': labels}, {'logps': logps}, old_logps=torch.zeros(3, 4)) + result = m.calculate() + assert 'train/approx_kl' not in result + def test_grpo_metric_reset(self): m = _no_dist_metric(GRPOMetric) labels = torch.tensor([[1, 2, -100, -100]]) diff --git a/tests/utils/test_utils.py b/tests/utils/test_utils.py index b9b1aa97a..448d59909 100644 --- a/tests/utils/test_utils.py +++ b/tests/utils/test_utils.py @@ -365,6 +365,36 @@ def test_returns_none_for_unsupported(self): result = align_logps_to_mask(42, mask, torch.float32) assert result is None + def test_full_sequence_form_indexes_by_mask(self): + """全序列形式（len >= seq_len，右 pad）必须先按 mask 取位置再 scatter。 + + 这是 ref/old 模型 forward 返回的形式：它的 pad 宽度是 dp split 前整个 micro batch + 的最大长度，所以常常比本 rank 的 logps 更宽。若退化成取行首 n_pos 个，读到的 + 就是 prompt 位置，每一个 IS ratio 都会错位。 + """ + mask = torch.tensor([[False, False, True, True], + [False, False, False, True]]) + full = torch.zeros(2, 7) # 7 > seq_len=4 + full[0, 2:4] = torch.tensor([-1.0, -2.0]) + full[1, 3] = -3.0 + result = align_logps_to_mask(full, mask, torch.float32) + assert result.shape == (2, 4) + assert result[0, 2].item() == pytest.approx(-1.0) + assert result[0, 3].item() == pytest.approx(-2.0) + assert result[1, 3].item() == pytest.approx(-3.0) + assert result[0, :2].abs().sum().item() == 0.0 + assert result[1, :3].abs().sum().item() == 0.0 + + def test_full_sequence_matches_grpo_loss_alignment(self): + """指标侧与损失侧必须对齐到**同一批 token**，否则面板上的 ratio 不是优化器看到的。""" + from twinkle.loss.grpo import GRPOLoss + mask = torch.tensor([[False, True, True, True], + [False, False, True, True]]) + full = torch.randn(2, 9) + got = align_logps_to_mask(full, mask, torch.float32) + want = GRPOLoss()._pad_and_align_to_batch(full, mask, mask.device, torch.float32) + assert torch.equal(got, want) + class TestFilterFromConfigKwargs: From 06e5f55bfa1ecee82b0b616ef5a3ce24f4e25459 Mon Sep 17 00:00:00 2001 From: root Date: Sun, 9 Aug 2026 00:19:20 +0800 Subject: [PATCH 35/60] =?UTF-8?q?feat(human=5Fe18):=20KodCode=20=E9=87=87?= =?UTF-8?q?=E9=9B=86=E6=94=AF=E6=8C=81=E5=A4=9A=E6=9C=BA=E5=88=86=E7=89=87?= =?UTF-8?q?=E4=B8=8E=E4=BA=A7=E7=89=A9=E5=90=88=E5=B9=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 无共享存储时用 crc32(data_id) % SHARD_N 划分题池，两机无需通信即可 保证互不重叠：crc32 是纯函数，跨进程/跨机/跨重启恒定（hash() 受 PYTHONHASHSEED 影响，会造成重叠+遗漏）。分片在 resume 过滤之前执行。 - SHARD_N / SHARD_ID 环境变量，默认 1/0 即原单机路径 - RUN_ID 在 SHARD_N>1 时加 .sN 后缀，避免两机同秒启动撞同一个 run - shard_tool.py: seed 导出已跑 data_id（B 机需要，否则会重跑 A 机做过的题） merge 合并多机产物（sft 按 data_id 去重 / candidates 按 (id,run,idx) / collect_log 加 src 标来源，因两机 chunk 都从 0 编号） - 可控的失败轨迹输入机制 KOD_USE_TRAJ，默认关闭 验证：真实题池 SHARD_N=2/3/4 均无重叠无遗漏、难度无偏；seed+merge 往返 27096 条逐条等价；3 个 PYTHONHASHSEED 结果一致。 --- cookbook/human_e18/README.md | 77 +++ cookbook/human_e18/e18_collect_kod.py | 738 ++++++++++++++++++++++ cookbook/human_e18/e18_kodcode.py | 404 ++++++++++++ cookbook/human_e18/e18_multidiag.py | 245 ++++++++ cookbook/human_e18/e18_prompts.py | 264 ++++++++ cookbook/human_e18/e18_rejection_sft.py | 775 ++++++++++++++++++++++++ cookbook/human_e18/e18_select.py | 142 +++++ cookbook/human_e18/e18_sft_kod.py | 543 +++++++++++++++++ cookbook/human_e18/e19_logp_select.py | 193 ++++++ cookbook/human_e18/e20_success_skill.py | 311 ++++++++++ cookbook/human_e18/e21_paired_rubric.py | 214 +++++++ cookbook/human_e18/run_collect_kod.sh | 8 + cookbook/human_e18/run_sft_kod.sh | 18 + cookbook/human_e18/shard_tool.py | 134 ++++ 14 files changed, 4066 insertions(+) create mode 100644 cookbook/human_e18/README.md create mode 100644 cookbook/human_e18/e18_collect_kod.py create mode 100644 cookbook/human_e18/e18_kodcode.py create mode 100644 cookbook/human_e18/e18_multidiag.py create mode 100644 cookbook/human_e18/e18_prompts.py create mode 100644 cookbook/human_e18/e18_rejection_sft.py create mode 100644 cookbook/human_e18/e18_select.py create mode 100644 cookbook/human_e18/e18_sft_kod.py create mode 100644 cookbook/human_e18/e19_logp_select.py create mode 100644 cookbook/human_e18/e20_success_skill.py create mode 100644 cookbook/human_e18/e21_paired_rubric.py create mode 100755 cookbook/human_e18/run_collect_kod.sh create mode 100755 cookbook/human_e18/run_sft_kod.sh create mode 100644 cookbook/human_e18/shard_tool.py diff --git a/cookbook/human_e18/README.md b/cookbook/human_e18/README.md new file mode 100644 index 000000000..62150ae5a --- /dev/null +++ b/cookbook/human_e18/README.md @@ -0,0 +1,77 @@ +# E18 — BigCodeBench 上的拒绝采样 SFT + +从 `cookbook/exp/skill2lora/skill_ablate/config.py` L244-252 的 `ExpSpec('E18', 'rejection_sft', ...)` +单独脱离出来的自包含实现，目录布局仿照 `cookbook/human`（E23）。 + +## 与 E23 的关系 + +两臂**共用**环境层与教师 judge（直接 import `../human/e23_bcb.py`、`../human/e23_rubric.py`， +不拷贝），所以数据过滤、沙箱单测判分、rubric 诊断三处逐字同源、结果可直接比。差别只在训练方法： + +| | E23 | E18（本目录） | +|---|---|---| +| 方法 | GRPO（组内归一化 advantage） | 拒绝采样 SFT（只用正样本） | +| loss | `SEAMBNPOLoss`（PPO clip + KL） | **`CrossEntropyLoss`**（纯交叉熵，无 ratio/clip/KL/advantage） | +| 梯度 | 正负都有；60% 组零方差→零梯度 | 只有正梯度，无零梯度浪费 | +| 每题产出 | 8 个候选全部进 batch | 两道筛后**唯一胜者**进池 | +| 训练轨迹 | 采样 token 直通（query+rubric） | messages 编码（**query-only**） | +| eval | 带 rubric | **query-only**（部署口径） | +| 卡数 | 8（含 ref） | 6（SFT 无需 ref 模型） | + +「选择」全部发生在两道筛（只有胜者入池），到 loss 这一层就是普通的「拟合目标文本」， +所以不传 `advantages` —— `CrossEntropyLoss` 不读该参数，传了是静默无效。 + +## 两道筛（本臂唯一自变量，见 `e18_select.py`） + +1. **增量达阈**：`with_pass >= base_pass_rate + MIN_PASS_GAIN`（默认 +2/8）。仅仅「没弄坏」 + （8/8 -> 8/8）不够格 —— 那种样本对「学会写有效 skill」没有监督信号； +2. **不超长**：超 `SKILL_CHAR_LIMIT` 直接丢； +3. **pass_rate 最大 -> 并列内 rubric 相似度**：先按客观效果取最大档，并列内取词频余弦相似度 + 最高的。长度只在相似度也并列时做确定性拆平（取较短者）。 + +⚠️ 原先第 3 道筛是「先按离 `LEN_BUDGET`（400）的距离取前一半」，已删：实测 **66% 的题 8 个 +候选全部 `with_pass=1.0`**（天花板打平），此时长度成了事实上的唯一决策依据，而它有系统性 +偏差——中文表达同样内容字符数天然更少（357 vs 705），永远更贴近 400，于是「离预算最近」被 +翻译成「选中文模板」，把信息量大的长英文候选全部淘汰。 + +⚠️ 原先有一道泄漏门（`leak_blocks`），已删：BCB 的 `reference_answer` 是 ~2500 字符的 dict， +子串匹配要求 skill 逐字包含整个 dict 的 repr，而 skill 上限 1500 字符 —— 触发概率恒为 0。 +真正的泄漏通道在 `test` 的断言期望值与 `canonical_solution`，需另写检测。 + +⚠️ 存活候选 ≤2 条时相似度那阶形同虚设，胜者由长度决定 —— 详见 `select_winner` 的 docstring。 + +## 文件 + +| 文件 | 内容 | +|---|---| +| `e18_rejection_sft.py` | 采集 + SFT 主循环、eval、swanlab | +| `e18_prompts.py` | executor / skill-gen / 训练轨迹三处 prompt | +| `e18_select.py` | 三道拒绝筛（泄漏门 + 词频余弦 + 两阶选择） | + +## 跑法 + +```bash +cd cookbook/human_e18 +# 教师 API 必需（rubric 是选择的参照系，不可降级） +export LLM_BACKUP_API_KEY=... LLM_BACKUP_BASE_URL=... +nohup python -u e18_rejection_sft.py > nohup.e18.$(date +%m%d-%H%M).log 2>&1 & +``` + +主要环境变量（默认值见文件头）： + +| 变量 | 默认 | 说明 | +|---|---|---| +| `ACCUMULATE` | 16 | 攒够多少条胜者 SFT 一次（须为 `TRAIN_DP` 整数倍） | +| `N_SKILLS` | 8 | 每题候选数 = 拒绝采样的池大小 | +| `MAX_UPDATES` | 50 | 总更新数 | +| `SKILL_CHAR_LIMIT` | 1500 | 超过直接丢 | +| `EVAL_SIZE` | 200 | holdout 题数；0 = 不 eval | +| `SWAN_PROJ` | twinkle | swanlab 项目（**勿** export `SWANLAB_PROJECT`） | + +## 产物 + +| 文件 | 内容 | +|---|---| +| `output.e18/e18_sft_dataset.jsonl` | **主产物**：每条胜者 + rubric/相似度/pass 全审计字段，可离线复算、换超参重训而不必重跑 GPU | +| `output.e18/train_log.jsonl` | 逐 chunk 指标（accept_rate / candidate_pass_rate / eval/*） | +| `output.e18/E18-final/` | 最终 skill 模型权重 | diff --git a/cookbook/human_e18/e18_collect_kod.py b/cookbook/human_e18/e18_collect_kod.py new file mode 100644 index 000000000..55efca255 --- /dev/null +++ b/cookbook/human_e18/e18_collect_kod.py @@ -0,0 +1,738 @@ +# -*- coding: utf-8 -*- +"""E18 冷启动数据采集（KodCode 域，**无 SFT**）。 + +与 `e18_rejection_sft.py` 的关系：把它的采集半部分原样搬过来，删掉训练/eval/权重同步。 +所以 `collect_chunk` 的逻辑、三道筛、指标口径、落盘字段全部逐字一致 —— 这样采出来的 +冷启动数据集与在线 run 的样本同分布，后续接 SFT 时不需要再对齐一次。 + +**删掉了什么，以及为什么** +* `TransformersModel` / `set_loss` / `set_optimizer` / `train_batch`：不训练。 +* `CheckpointEngineManager` / `_sync_trained_to_sampler` / `_restore_base_weights`： + 没有训练权重要推给 sampler，skill_sampler 全程是初始模型。 +* `run_eval` / `EVAL_SIZE`：eval 衡量的是「训练后的 skill 模型在部署口径下的能力」， + 不训练时它恒等于 baseline，跑它纯浪费 GPU。故 `load_records(eval_size=0)`，题全进采集池。 + +**8 卡全部给 rollout**：原来 train 占 2 张、skill_sampler 2 张、base_sampler 4 张。 +现在 train 那 2 张转给两个 sampler。base_sampler 仍拿大头（默认 6）——它是唯一瓶颈： +每 chunk 它要跑裸解 CHUNK_SIZE*BARE_ROLLOUTS + 重解 CHUNK_SIZE*N_SKILLS*EXEC_ROLLOUTS， +序列数比 skill_sampler 多一个量级，且 EXEC_MAX_TOKENS 远大于 SKILL_MAX_TOKENS。 + +产物（都在 OUTPUT_DIR 下，append-only）： +* `e18_sft_dataset.jsonl`：胜者，字段与在线 run 完全一致，直接可喂 SFT。 +* `e18_candidates.jsonl`：**全部** skill 候选（含落选与解析失败的），靠 `kept` 区分选与未选。 + 用于离线重算阀值、判断胜者是真更好还是拆平局选出来的。 +* `collect_log.jsonl`：逐 chunk 指标，用于监控 accept_rate / degrade_rate 是否异常。 +""" +import json +import os +import shutil +import sys +import time +import zlib +from dataclasses import dataclass +from typing import Any, Dict, List, Tuple + +import torch +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import SamplingParams +from twinkle.dataloader import DataLoader +from twinkle.sampler import vLLMSampler +from twinkle.template import Template + +_HERE = os.path.dirname(os.path.abspath(__file__)) +_COOKBOOK = os.path.abspath(os.path.join(_HERE, '..')) +for _p in (_HERE, os.path.join(_COOKBOOK, 'human')): + if _p not in sys.path: + sys.path.insert(0, _p) + +from e23_rubric import build_checker, class_metrics # noqa: E402 + +from e18_kodcode import (clean_text, empty_roll, extract_skill, # noqa: E402 + judge_seqs, load_records) +from e18_multidiag import MultiDiagCache, multidiag_metrics # noqa: E402 +from e18_prompts import (direct_prompt, format_trajectory, # noqa: E402 + skill_solve_prompt, skillgen_prompt) +from e18_select import gain_stats, select_winner # noqa: E402 + +logger = get_logger() + +# ========== Configuration ========== +MODEL_ID = os.environ.get('MODEL_ID', 'Qwen/Qwen3-4B') +OUTPUT_DIR = os.environ.get('OUTPUT_DIR', os.path.join(_HERE, 'output.e18.kod')) + +# ⭐ 8 卡全给 rollout：不训练，所以没有 train 组。base_sampler 拿大头（瓶颈见文件头注释）。 +SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) +BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 6)) +NUM_GPUS = SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS +GPU_MEM = float(os.environ.get('GPU_MEM', 0.8)) + +SEED = int(os.environ.get('SEED', 42)) +# ⭐ 续跑开关。为何默认关：开着会把旧 run 的样本接着往同一批数据里添，而跨 run 的 +# 模型/prompt 可能已经变了 —— 默认开启等于默认允许污染，违反归档机制的初衷。 +# KOD_RESUME=1 时做三件事（见 archive_output_dir / resume_done_ids）： +# 1. 归档**之前**先从 e18_candidates.jsonl 读出已跑过的 data_id； +# 2. 三个 jsonl 复制回新目录（而不是只搬 broken_tasks 缓存），让计数接着累积； +# 3. 把已采题从题池里 filter 掉。 +# 为何必须有它：DataLoader 用固定 SEED+shuffle，重启后取题顺序逐字相同，不过滤就会 +# 把前 N 个 chunk 的题原样重跑一遂（实测 65 chunk ≈ 11 小时 8 卡），纯浪费。 +KOD_RESUME = os.environ.get('KOD_RESUME', '0') == '1' +CHUNK_SIZE = int(os.environ.get('CHUNK_SIZE', 64)) +# ⭐ 4 而不是 8：skill 候选池只用来「邀出」候选，最终只有 1 条胜者入池。实测 66% 的题 +# 8 个候选的 with_pass 全部并列，多出来的 4 条几乎不改变胜者，却要占掉一半 executor 序列。 +# ⭐ N_SKILLS 现在是**两阶段的总上限**：先生成 N_SKILLS_STAGE1 个，全部不达标才补到 N_SKILLS。 +# 实测（4224 题）：39.7% 的题前 2 个候选就能出胜者 -> 平均只花 3.21 个候选， +# 产出与固定 4 个**完全相同**（1134 题），executor 序列从 32 降到 25.6。 +N_SKILLS = int(os.environ.get('N_SKILLS', 4)) +N_SKILLS_STAGE1 = int(os.environ.get('N_SKILLS_STAGE1', 2)) +# ⭐ 天花板短路：base_pass_rate 已打满的题直接跳过，不生成任何 skill 候选。 +# 依据：实测 1750 道 base=1.0 的题，入池 **0 条** —— 因为门槛是 +# pass_gain >= MIN_PASS_GAIN(0.25)，而 base=1.0 时 with_pass 最大也是 1.0，gain 恒为 0。 +# 所以这 41% 的题在**本离线采集脚本**里是纯浪费（占 37% 算力、零产出）。 +# ⚠️ 与 collect_chunk 原注释「全量 rollout」的设计意图相反：那条理由（避免 skill 模型 +# 只学会救难题）适用于**在线 RL**（每步需要 reward 信号，包括平局组）； +# 本脚本是纯离线采集，产物只有 e18_sft_dataset.jsonl，天花板题从不进该文件。 +# 若日后要拿这份代码回到在线 RL，必须把本开关置 0。 +SKIP_CEILING = int(os.environ.get('SKIP_CEILING', 1)) +# ⭐ 粗筛 rollout 次数：先给每个候选跑 PROBE_ROLLOUTS 次，只对**并列最高**者补到 +# EXEC_ROLLOUTS 次。实测（1515 题）：M=2 的 top-1 与 8 次一致率 68.9%、M=4 为 77.1%； +# 而 51% 的题 4 候选 with_pass 全同分，前 2 次就能看出打平并提前停手。 +# 0 = 关闭粗筛（所有候选直接跑足 EXEC_ROLLOUTS 次，与历史 run 逐字一致）。 +PROBE_ROLLOUTS = int(os.environ.get('PROBE_ROLLOUTS', 2)) +# 采够多少条胜者就停。0 = 把题池跑完。 +TARGET_SAMPLES = int(os.environ.get('TARGET_SAMPLES', 20000)) +MAX_CHUNKS = int(os.environ.get('MAX_CHUNKS', 0)) # 0 = 不限 + +# ⭐ 多机并行分片。SHARD_ID 取值 0..SHARD_N-1，每台机器只跑 +# `crc32(data_id) % SHARD_N == SHARD_ID` 的题。 +# 为何用哈希而不是「切片取前后一半」： +# 1. 无状态 —— 两台机器不需要任何通信/共享盘就能保证不重叠（NAS 不通用时的必要条件） +# 2. 对题池变化稳健 —— 万一两边题池大小不一致（版本/缓存差异），按下标切会错位重叠， +# 而哈希绑定在 data_id 上，永远不会 +# 3. 难度分布无偏 —— crc32 与 gpt_pass_percentage 无相关，两片难度同分布 +# ❗ 两台机器必须用**相同的 SHARD_N**，否则分片不构成划分（会既重叠又遗漏）。 +# ❗ TARGET_SAMPLES 是**本分片的**目标：想总共 20000 就两边各填 10000。 +SHARD_N = int(os.environ.get('SHARD_N', 1)) +SHARD_ID = int(os.environ.get('SHARD_ID', 0)) + +SKILL_MAX_TOKENS = int(os.environ.get('SKILL_MAX_TOKENS', 8192)) +# ⭐ 把失败轨迹（裸解时做错的那份代码）一并给 skillmodel。默认**关**： +# 开了就换了 prompt 口径，与已采的 1750 条不同源，不能默认静默切换。 +# 长度安全性已实测（.tmp_analysis/len_budget.py，400 条真实数据）： +# prompt 预算 = MAX_MODEL_LEN(16000) - SKILL_MAX_TOKENS(8192) = 7808 token +# 现状不带轨迹 中位 1802 / 最大 3192 +# 带 1 条(3x 最坏) 中位 2710 / 最大 6162 -> 仍在预算内，**无需改 MAX_MODEL_LEN** +# 带 2 条(3x 最坏) 最大 9132 -> 1.5% 超预算，所以默认只给 1 条 +# 超预算的后果不是报错而是 vLLM 返回空序列 -> parseable=False -> 该候选白跑， +# 难以从日志发现，所以宁可保守。 +USE_TRAJ = int(os.environ.get('KOD_USE_TRAJ', 0)) +TRAJ_N = int(os.environ.get('KOD_TRAJ_N', 1)) +# 单条轨迹的字符上限。4000 字符 ≈ 1540 token（代码 2.6 chars/token）， +# 加上现状最大 3192 仍不到 7808。超长者由 format_trajectory 头尾各留一半。 +TRAJ_MAX_CHARS = int(os.environ.get('KOD_TRAJ_MAX_CHARS', 4000)) +SKILL_GEN_TEMPERATURE = float(os.environ.get('SKILL_GEN_TEMPERATURE', 1.0)) +SKILL_GEN_TOP_P = float(os.environ.get('SKILL_GEN_TOP_P', 1.0)) +SKILL_GEN_TOP_K = int(os.environ.get('SKILL_GEN_TOP_K', -1)) +EXEC_MAX_TOKENS = int(os.environ.get('EXEC_MAX_TOKENS', 15000)) +MAX_MODEL_LEN = int(os.environ.get('MAX_MODEL_LEN', 16000)) + +EXEC_ROLLOUTS = int(os.environ.get('EXEC_ROLLOUTS', 8)) +# ⭐ 裸解单独用 4 次，而带 skill 重解仍是 EXEC_ROLLOUTS(8)。两侧刻意**不对称**： +# 裸解只承担「这题有没有提升空间」的粗判（rate<1 即进诊断），4 次足够；而 with_pass 要 +# 在候选之间排序、还要减去 base 算增量，精度需求高得多，降它会直接动摇入池门槛的语义。 +# ⚠️ 代价（已知并接受）：base_pass_rate 只有 5 档（0,.25,.5,.75,1），与 with_pass 的 9 档 +# 不同分母。于是 pass_gain = with_pass - base_pass_rate 的零点变粗，base 的采样标准误从 +# 0.177 升到 0.25 —— 判「有没有空间」够用，但别再把 pass_gain 的绝对值当精密量看。 +# 另外 n_ceiling（base 打满 4/4）会比 8/8 更容易达成，天花板题占比会上升。 +BARE_ROLLOUTS = int(os.environ.get('BARE_ROLLOUTS', 4)) +EXEC_TEMPERATURE = float(os.environ.get('EXEC_TEMPERATURE', 0.6)) +EXEC_TOP_P = float(os.environ.get('EXEC_TOP_P', 0.95)) + +SKILL_CHAR_LIMIT = int(os.environ.get('SKILL_CHAR_LIMIT', 1500)) +# 门槛仍按 EXEC_ROLLOUTS(8) 算 = +2/8 = +0.25，与历史 run 逐字可比。 +# 不能拿 BARE_ROLLOUTS 做分母：with_pass 是 8 次采的，两边分母必须是同一个。 +MIN_GAIN_ROLLOUTS = int(os.environ.get('MIN_GAIN_ROLLOUTS', 2)) +MIN_PASS_GAIN = MIN_GAIN_ROLLOUTS / max(1, EXEC_ROLLOUTS) + +# ⭐ 多机并行时 RUN_ID 必须带机器标识：原来只有时间戳，两台机器同一秒启动会撞成 +# 同一个 run，合并后就再也分不出某条数据是哪台机器产的（排查单机异常时必须能分开）。 +# 分片号是天然的机器标识，比 hostname 稳（容器重建后 hostname 会变），所以 +# SHARD_N>1 时后缀 `.sN` —— 单机跑时 RUN_ID 保持原格式，历史 run 的比对不受影响。 +RUN_ID = time.strftime('%m%d-%H%M%S') + (f'.s{SHARD_ID}' if SHARD_N > 1 else '') + + +@dataclass +class Runtime: + skill_sampler: Any + base_sampler: Any + checker: Any + rubric_cache: MultiDiagCache + + +# =========================================================================== +# 采样工具（与 e18_rejection_sft 逐字一致） +# =========================================================================== +def run_samples(sampler, prompts, num_samples, max_tokens, gen_dp, + temperature=None, top_p=None, top_k=None, logprobs=None): + """采样。prompts 少于 dp 时补齐再截回 —— Ray 的 dp 切分要求每个 rank 至少一条。 + + ⭐ 与 e18_rejection_sft.run_samples 逐字一致。三处踩过的坑： + 1. 字段名是 `num_samples` 不是 `n`（SamplingParams 没有 `n`，传了直接 TypeError）。 + 2. 走 `sampler.sample(prompts, params)`，不是 pack_user_data + generate_sequences。 + 3. dp 补齐不能省：最后一个 chunk 不满、或 flat 很少时，条数 < dp 会直接报错。 + """ + if not prompts: + return [] + import copy + params = SamplingParams( + max_tokens=max_tokens, + temperature=0.6 if temperature is None else temperature, + top_p=0.95 if top_p is None else top_p, + num_samples=num_samples, + **({} if top_k is None else {'top_k': top_k}), + **({} if logprobs is None else {'logprobs': logprobs})) + padded = prompts + if gen_dp > 1 and 0 < len(prompts) < gen_dp: + padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] + responses = sampler.sample(padded, params)[:len(prompts)] + return [list(r.sequences) if (r and r.sequences) else [] for r in responses] + + +def first_seq(seqs): + return seqs[0] if seqs else None + + +def seq_text(seq) -> str: + return clean_text(getattr(seq, 'decoded', '') or '') if seq is not None else '' + + +def _mean(xs) -> float: + xs = [float(x) for x in xs if x is not None] + return sum(xs) / len(xs) if xs else 0.0 + + +# =========================================================================== +# 采集：裸解 -> 诊断 -> skill-gen -> executor 重解 -> 三道筛 +# =========================================================================== +def _pass_rate(rolls) -> float: + return (sum(1.0 for x in rolls if x['correct']) / len(rolls)) if rolls else 0.0 + + +def bare_solve(rt: Runtime, records, rollouts: int = None) -> List[List[Dict[str, Any]]]: + """裸题重解 `rollouts` 次，每条记录返回一个 roll 列表（长度 = 实际采到的序列数）。 + + 返回**嵌套**列表而不是单个 roll：调用方靠 _pass_rate() 取连续值。 + 判分全部汇到一次 judge_seqs（它内部按 (task_id, code) 去重，相同代码只跑一次单测）。 + + ⭐ M==1 时必须降到 temperature=0.0（与原版一致）：单次采样就不需要多样性， + 带温度只会引入无意义的方差；max(1, ...) 防 rollouts 传 0 时采不到任何序列。 + """ + M = max(1, rollouts if rollouts is not None else EXEC_ROLLOUTS) + out = run_samples(rt.base_sampler, [direct_prompt(r['problem']) for r in records], + M, EXEC_MAX_TOKENS, BASE_SAMPLER_GPUS, + temperature=(0.0 if M == 1 else EXEC_TEMPERATURE), + top_p=(None if M == 1 else EXEC_TOP_P)) + pairs, spans = [], [] + for r, seqs in zip(records, out): + seqs = list(seqs or []) + spans.append(len(seqs)) + pairs.extend((s, r['reference_answer']) for s in seqs) + judged = judge_seqs(pairs) + rolls, i = [], 0 + for n in spans: + rolls.append(judged[i:i + n] if n else [empty_roll()]) + i += n + return rolls + + +# =========================================================================== +# 采集（与 e18_rejection_sft.collect_chunk 逐字一致） +# =========================================================================== +def _pick_trajectory(rolls: List[Dict[str, Any]]) -> str: + """从裸解的 rolls 里挑出最值得给 skill-gen 看的失败代码。 + + ⭐ 挑选而不是全给：BARE_ROLLOUTS=4 条里常有 2-3 条是**同一个错** + （judge_seqs 内部按 (task_id, code) 去重就是因为重复普遍），全给只会重复占预算。 + + 优先级：有报错信息 > 代码短。 + 为何偏好**短**代码：长代码往往是思维链泄到正文里的 no_code / import_or_syntax + 废文，信息密度低；短而完整的错解才能看出逻辑问题。同时也直接压低了长度风险。 + """ + bad = [x for x in (rolls or []) if not x.get('correct') and (x.get('code') or '').strip()] + if not bad: + return '' + bad.sort(key=lambda x: (0 if x.get('error') else 1, len(x.get('code') or ''))) + blocks = [format_trajectory(x.get('code'), x.get('error'), x.get('kind'), + max_chars=TRAJ_MAX_CHARS) + for x in bad[:max(1, TRAJ_N)]] + return '\n\n'.join(blocks) + + +def _gen_candidates(rt: Runtime, todo, n_skills: int) -> List[List[Dict[str, Any]]]: + """给 todo 里每道题生成 n_skills 个 skill 候选（只生成，不判分）。 + + todo 元素 = (record, rubric, base_rate, trajectory)；trajectory 在 USE_TRAJ=0 时恒为 ''。 + """ + sg = run_samples(rt.skill_sampler, + [skillgen_prompt(r['problem'], d, eval=False, trajectory=tj) + for r, d, _br, tj in todo], + n_skills, SKILL_MAX_TOKENS, SKILL_SAMPLER_GPUS, + temperature=SKILL_GEN_TEMPERATURE, top_p=SKILL_GEN_TOP_P, + top_k=SKILL_GEN_TOP_K) + out = [] + for seqs in sg: + cands = [] + for s in seqs or []: + resp = seq_text(s) + block = extract_skill(resp) + cands.append({'skills': block, 'response': resp, 'parseable': bool(block), + 'with_pass': None, 'kept': False, + 'skillgen_stop': getattr(s, 'stop_reason', None)}) + out.append(cands) + return out + + +def _judge_candidates(rt: Runtime, flat, rollouts: int) -> None: + """对 flat=[(record, cand), ...] 跑 `rollouts` 次重解并原地回写 with_pass。 + + ⭐ 原地累加而非覆盖：两阶段粗筛里同一个候选会被判两次（先 PROBE 后补足）， + 第二次必须把两次的样本**合并**算 pass_rate，否则前 PROBE_ROLLOUTS 次白扔。 + 用 _n_correct/_n_total 累计，with_pass 每次由累计值重算。 + """ + if not flat: + return + ws = run_samples(rt.base_sampler, + [skill_solve_prompt(r['problem'], c['skills']) for r, c in flat], + rollouts, EXEC_MAX_TOKENS, BASE_SAMPLER_GPUS, + temperature=EXEC_TEMPERATURE, top_p=EXEC_TOP_P) + pairs, spans = [], [] + for (r, _c), seqs in zip(flat, ws): + seqs = list(seqs or []) + spans.append(len(seqs)) + pairs.extend((s, r['reference_answer']) for s in seqs) + judged = judge_seqs(pairs) + i = 0 + for (_r, c), n in zip(flat, spans): + rr = judged[i:i + n] + i += n + c['_n_correct'] = c.get('_n_correct', 0) + sum(1 for x in rr if x['correct']) + c['_n_total'] = c.get('_n_total', 0) + n + c['with_pass'] = (c['_n_correct'] / c['_n_total']) if c['_n_total'] else 0.0 + c['n_rollouts'] = c['_n_total'] + if c.get('roll_kind') in (None, 'pass'): + c['roll_kind'] = (next((x['kind'] for x in rr if not x['correct']), + c.get('roll_kind') or 'pass') if rr else 'empty') + + +def _resolve_with_probe(rt: Runtime, todo, per_task, base_rates) -> None: + """两阶段 rollout：先粗筛 PROBE_ROLLOUTS 次，只把**可能是胜者**的候选补到 EXEC_ROLLOUTS。 + + 补足的判据是「粗筛并列最高」而不是「粗筛最高」：粗筛只有 2 次采样， + 并列极其常见（实测 51% 的题全同分），只补单个最高者会把真胜者漏掉。 + ⚠️ 只有 parseable 的候选参与；不可解析的候选 with_pass 保持 None，与历史行为一致。 + """ + alive = [(r, c) for (r, _d, _br, _tj), cands in zip(todo, per_task) + for c in cands if c.get('parseable')] + if not alive: + return + if not PROBE_ROLLOUTS or PROBE_ROLLOUTS >= EXEC_ROLLOUTS: + _judge_candidates(rt, alive, EXEC_ROLLOUTS) + return + _judge_candidates(rt, alive, PROBE_ROLLOUTS) + # 逐题挑「粗筛并列最高」者补足到 EXEC_ROLLOUTS + need = [] + for (r, _d, base_rate, _tj), cands in zip(todo, per_task): + ok = [c for c in cands if c.get('parseable')] + if not ok: + continue + top = max(c['with_pass'] for c in ok) + # 粗筛就已经够不到门槛的题：补足也不可能入池，直接省掉。 + if top + 1e-9 < base_rate + MIN_PASS_GAIN - (1.0 / max(1, EXEC_ROLLOUTS)): + continue + need.extend((r, c) for c in ok if c['with_pass'] >= top - 1e-9) + _judge_candidates(rt, need, EXEC_ROLLOUTS - PROBE_ROLLOUTS) + + +def collect_chunk(rt: Runtime, chunk, ci: int) -> Tuple[List[Dict[str, Any]], Dict[str, float]]: + """一个 chunk 的采集，返回 (胜者列表, 指标)。 + + 三层省算力（每层只放行需要下一层的题），全部可用环境变量关掉回到历史行为： + 1. SKIP_CEILING : base 打满的题不生成候选（实测入池 0 条，纯浪费） + 2. N_SKILLS_STAGE1: 先 2 个候选，全不达标才补到 N_SKILLS + 3. PROBE_ROLLOUTS: 每候选先 2 次，只对并列最高者补到 EXEC_ROLLOUTS + + 诊断（rubric）只对**做错的题**拉，且**每一种失败模式各诊一次后合并**（见 e18_multidiag）。 + 无诊断的题不丢弃：skillgen_prompt 内部会填兜底文案。 + """ + base_rolls = bare_solve(rt, chunk, rollouts=BARE_ROLLOUTS) + base_rates = [_pass_rate(rr) for rr in base_rolls] + base_acc = _mean(base_rates) + wrong = [(r, rr) for r, rr, rate in zip(chunk, base_rolls, base_rates) if rate < 1.0] + + before = rt.rubric_cache.stats.copy() + diags = rt.rubric_cache.diagnose_many(rt.checker, wrong) + rmetrics = multidiag_metrics(rt.rubric_cache.stats - before) + diag_by_id = {id(r): d for (r, _rr), d in zip(wrong, diags) if d} + n_rubric_missing = len(wrong) - len(diag_by_id) + n_multi = sum(1 for d in diag_by_id.values() if d.count('FAILURE ') > 1) + + # ⭐ todo 元素是四元组 (record, rubric, base_rate, trajectory)。 + # trajectory 只在 USE_TRAJ=1 时非空；关闭时恒为 '' -> skillgen_prompt 走原模板， + # 与历史 run 逐字一致。轨迹取自**裸解**的 rolls（就是当初拿去要诊断的那批）， + # 所以与 rubric 同源、描述的是同一次失败 —— 这正是 E22 当时做不到的（那次是重采的）。 + traj_by_id = ({id(r): _pick_trajectory(rr) for r, rr in zip(chunk, base_rolls)} + if USE_TRAJ else {}) + all_tasks = [(r, diag_by_id.get(id(r), ''), rate, traj_by_id.get(id(r), '')) + for r, rate in zip(chunk, base_rates)] + # ⭐ 天花板短路。n_skipped 单独记账，指标分母用 todo（非天花板题）—— + # 于是 baseline_accuracy / candidate_pass_rate / lift 三个指标的口径变了， + # **与 SKIP_CEILING=0 的历史 run 不可直接比较**，看趋势时注意这一点。 + if SKIP_CEILING: + todo = [t for t in all_tasks if t[2] + MIN_PASS_GAIN <= 1.0 + 1e-9] + else: + todo = all_tasks + n_skipped = len(all_tasks) - len(todo) + + per_task = _gen_candidates(rt, todo, min(N_SKILLS_STAGE1, N_SKILLS)) if todo else [] + if per_task: + _resolve_with_probe(rt, todo, per_task, base_rates) + + # ⭐ 阶段2：阶段1 全部不达标的题，再生成剩余候选。 + # 实测 39.7% 的题阶段1 就出胜者 -> 这批题省掉一半候选；6.1% 的题靠阶段2 救回。 + n_stage2_tasks = n_stage2_saved = 0 + remain = N_SKILLS - min(N_SKILLS_STAGE1, N_SKILLS) + if per_task and remain > 0: + idx2 = [i for i, ((_r, _d, br, _tj), cands) in enumerate(zip(todo, per_task)) + if not any(c.get('parseable') + and (c.get('with_pass') or 0.0) >= br + MIN_PASS_GAIN - 1e-9 + for c in cands)] + if idx2: + n_stage2_tasks = len(idx2) + todo2 = [todo[i] for i in idx2] + extra = _gen_candidates(rt, todo2, remain) + per_task2 = [[] for _ in todo] + for i, cands in zip(idx2, extra): + per_task2[i] = cands + _resolve_with_probe(rt, todo, per_task2, base_rates) + for i, cands in zip(idx2, extra): + per_task[i].extend(cands) + n_stage2_saved = sum( + 1 for i, cands in zip(idx2, extra) + if any(c.get('parseable') + and (c.get('with_pass') or 0.0) >= todo[i][2] + MIN_PASS_GAIN - 1e-9 + for c in cands)) + + accepted, sims = [], [] + n_pass_cands = n_survivors = 0 + n_acc_hard = n_acc_easy = 0 + gtot = {'improved': 0, 'tied': 0, 'degraded': 0} + gains = [] + for (r, d, base_rate, _tj), cands in zip(todo, per_task): + passers = [c for c in cands + if c.get('parseable') and (c.get('with_pass') or 0) >= base_rate] + n_pass_cands += len(passers) + for k, v in gain_stats(cands, base_rate).items(): + gtot[k] += v + best = select_winner(cands, d, r['reference_answer'], + skill_char_limit=SKILL_CHAR_LIMIT, + base_pass_rate=base_rate, min_pass_gain=MIN_PASS_GAIN) + if best is None: + continue + n_survivors += 1 + sims.append(best['rubric_similarity']) + gains.append(best['pass_gain']) + if base_rate >= 1.0: + n_acc_easy += 1 + else: + n_acc_hard += 1 + accepted.append({ + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'data_id': r.get('data_id', ''), 'skills': best['skills'], + 'response': f"\n{best['skills']}\n", + 'base_pass_rate': base_rate, 'with_pass_rate': best['with_pass'], + 'pass_gain': best['pass_gain'], 'gain_kind': best['gain_kind'], + 'rubric': d, 'chunk': ci, 'run': RUN_ID, + 'rubric_similarity': best['rubric_similarity'], + 'skill_chars': len(best['skills']), + 'n_candidates_passed': len(passers)}) + + # ⭐ 分母改成**实际生成的候选总数**：两阶段下每道题的候选数不再是定值 N_SKILLS + # （阶段1 就达标的题只有 N_SKILLS_STAGE1 个），写死 N_SKILLS 会把分母虚抬、 + # 让 candidate_pass_rate 和 lift 系统性偏低。 + n_cands_total = sum(len(cs) for cs in per_task) + _cand_pass = (n_pass_cands / n_cands_total) if n_cands_total else 0.0 + # 必须在选择循环之后：kept / with_pass / rubric_similarity 都是 select_winner 原地回写的。 + dump_candidates(todo, per_task, ci) + metrics = { + 'train/baseline_accuracy': base_acc, + 'train/accept_rate': (len(accepted) / len(todo)) if todo else 0.0, + 'train/candidate_pass_rate': _cand_pass, + 'train/lift': _cand_pass - base_acc, + 'train/selected_rubric_similarity': _mean(sims), + 'train/selected_skill_length_characters': _mean( + [float(s['skill_chars']) for s in accepted]), + 'signal/n_wrong': float(len(wrong)), + 'signal/n_rubric_missing': float(n_rubric_missing), + 'signal/n_multi_cause': float(n_multi), + 'signal/n_accepted': float(len(accepted)), + 'signal/n_accepted_hard': float(n_acc_hard), + 'signal/n_accepted_easy': float(n_acc_easy), + 'signal/n_ceiling': float(sum(1 for _r, _d, br, _tj in todo + if br + MIN_PASS_GAIN > 1.0 + 1e-9)), + 'signal/min_pass_gain': MIN_PASS_GAIN, + # 三层省算力的实测记账（用于事后核对真省了多少，而不是只信离线模拟） + 'saving/n_ceiling_skipped': float(n_skipped), + 'saving/n_stage2_tasks': float(n_stage2_tasks), + 'saving/n_stage2_rescued': float(n_stage2_saved), + 'saving/candidates_per_task': (n_cands_total / len(todo)) if todo else 0.0, + 'saving/exec_sequences': float(sum( + c.get('_n_total', 0) for cs in per_task for c in cs)), + # ⭐ 轨迹注入的实测记账。traj/rate 远低于 1 就说明很多题拿不到可用失败代码 + # （全对、或全是 no_code 空代码），此时该开关的实际覆盖面比以为的小。 + # traj/chars 盯长度风险：除以 2.6 就是多吃的 token 数。 + 'traj/enabled': float(USE_TRAJ), + 'traj/rate': (sum(1 for _r, _d, _br, tj in todo if tj) / len(todo)) if todo else 0.0, + 'traj/chars': _mean([float(len(tj)) for _r, _d, _br, tj in todo if tj]), + 'gain/improved_candidates': float(gtot['improved']), + 'gain/tied_candidates': float(gtot['tied']), + 'gain/degraded_candidates': float(gtot['degraded']), + 'gain/degrade_rate': (gtot['degraded'] / max(1, sum(gtot.values()))), + 'gain/improve_rate': (gtot['improved'] / max(1, sum(gtot.values()))), + 'gain/selected_pass_gain': _mean(gains), + } | rmetrics | class_metrics([d for _r, d, _br, _tj in todo if d]) + return accepted, metrics + + +def dump_dataset(accepted) -> None: + """胜者落盘 append-only 的 SFT 数据集（字段与在线 run 完全一致）。""" + if not accepted: + return + path = os.path.join(OUTPUT_DIR, 'e18_sft_dataset.jsonl') + with open(path, 'a', encoding='utf-8') as f: + for s in accepted: + f.write(json.dumps(s, ensure_ascii=False) + '\n') + + +def dump_candidates(todo, per_task, ci: int) -> None: + """**全部** skill 候选落盘（选上的、没选上的、甚至解不出 的），append-only。 + + 为何必需：`e18_sft_dataset.jsonl` 只存胜者，而拒绝采样的全部信息量在「同一题的 N 条 + 候选之间的差异」里 —— 丢掉落选者就无法回答：胜者是真的更好，还是只是采样噪声 + （候选全部并列时靠拆平局选出来的）？也无法事后重算阀值：改 MIN_PASS_GAIN / + SKILL_CHAR_LIMIT 后想知道会多收多少条，必须有落选者的 with_pass 才能离线重放。 + + `kept` 区分选与未选：`select_winner` 对胜者**原地** 置 True（e18_select.py:107）， + 本函数必须在 select_winner 之后调用，否则全部候选都是 False。 + 同理 `with_pass` / `pass_gain` / `rubric_similarity` 也是选择阶段回写的。 + + 不存 `response` 全文（包含 think 链，体量是 skills 的十倍量级），只存抽取后的 + skills 与长度/截断信息；解析失败（parseable=False）时 skills 为空串，靠 skillgen_stop + 判断是撞预算截断还是真的不守格式。 + """ + path = os.path.join(OUTPUT_DIR, 'e18_candidates.jsonl') + with open(path, 'a', encoding='utf-8') as f: + for (r, d, base_rate, _tj), cands in zip(todo, per_task): + for j, c in enumerate(cands): + wp = c.get('with_pass') + f.write(json.dumps({ + 'chunk': ci, 'run': RUN_ID, + 'data_id': r.get('data_id', ''), + 'task_id': r['reference_answer'].get('task_id', ''), + 'cand_idx': j, + 'kept': bool(c.get('kept')), + 'parseable': bool(c.get('parseable')), + 'skills': c.get('skills', ''), + 'skill_chars': len(c.get('skills') or ''), + 'skillgen_stop': c.get('skillgen_stop'), + 'base_pass_rate': base_rate, + 'with_pass_rate': wp, + # 未参与重解（解析失败）时 with_pass 是 None，pass_gain 也留 None， + # 不能当 0 存 —— 否则离线统计会把「没跑」混成「跑了但零增益」。 + 'pass_gain': (None if wp is None else round(wp - base_rate, 6)), + 'gain_kind': c.get('gain_kind'), + 'n_rollouts': c.get('n_rollouts'), + 'roll_kind': c.get('roll_kind'), + 'rubric_similarity': c.get('rubric_similarity'), + 'rubric': d, + }, ensure_ascii=False) + '\n') + + +# =========================================================================== +# main +# =========================================================================== +def build_runtime(checker, rubric_cache) -> Runtime: + """两组卡：skill_sampler / base_sampler(executor)。不训练，所以没有 train 组。""" + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, lazy_collect=False, groups=[ + DeviceGroup(name='skill_sampler', ranks=list(range(0, SKILL_SAMPLER_GPUS)), + device_type='GPU'), + DeviceGroup(name='base_sampler', ranks=list(range(SKILL_SAMPLER_GPUS, NUM_GPUS)), + device_type='GPU')]) + + def _sampler(group, world, enable_thinking): + s = vLLMSampler(model_id=MODEL_ID, + engine_args={'gpu_memory_utilization': GPU_MEM, + 'max_model_len': MAX_MODEL_LEN, + 'tensor_parallel_size': 1}, + device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), + remote_group=group) + s.set_template(Template, model_id=MODEL_ID, enable_thinking=enable_thinking, + max_length=MAX_MODEL_LEN) + return s + + # 与在线 run 同口径：skill_sampler 开 think（采集要多样性），executor 也开 think。 + skill_sampler = _sampler('skill_sampler', SKILL_SAMPLER_GPUS, enable_thinking=True) + base_sampler = _sampler('base_sampler', BASE_SAMPLER_GPUS, enable_thinking=True) + return Runtime(skill_sampler=skill_sampler, base_sampler=base_sampler, + checker=checker, rubric_cache=rubric_cache) + + +def count_existing_samples() -> int: + """续跑：已有胜者条数（= e18_sft_dataset.jsonl 的有效行数）。同样要在归档前调。 + + 这里才该用 sft_dataset 而不是 candidates：它要回答的是「已经凑了多少条胜者」， + 用来接着比 TARGET_SAMPLES；而 resume_done_ids() 要回答的是「哪些题不用再跑」。 + 两个问题不同源，切勿合并。 + """ + path = os.path.join(OUTPUT_DIR, 'e18_sft_dataset.jsonl') + if not os.path.exists(path): + return 0 + n = 0 + with open(path, 'r', encoding='utf-8') as f: + for line in f: + if line.strip(): + n += 1 + return n + + +def resume_done_ids() -> set: + """续跑：读出已经跑过的 data_id。必须在 archive_output_dir() **之前**调用。 + + ⭐ 读 `e18_candidates.jsonl` 而不是 `e18_sft_dataset.jsonl`：后者只有**胜者**（实测 + accept_rate 约 22%），拿它去重会把剩下 78%「跑过但没能入池」的题当成未跑、 + 下次重新烧一遍 GPU。候选文件是全量落盘的，覆盖面才完整。 + + 容错：进程被 kill 时最后一行可能是写一半的残行，json.loads 会抛异常 -> + 逐行 try 跳过（丢掉一道题的去重信息只是多跑一题，而整个函数抛异常会直接弄挂启动）。 + """ + path = os.path.join(OUTPUT_DIR, 'e18_candidates.jsonl') + if not os.path.exists(path): + return set() + done = set() + bad = 0 + with open(path, 'r', encoding='utf-8') as f: + for line in f: + line = line.strip() + if not line: + continue + try: + d = json.loads(line)['data_id'] + except Exception: + bad += 1 + continue + if d: + done.add(str(d)) + logger.info(f'[resume] 已跑过 {len(done)} 题' + + (f'（跳过 {bad} 行残缺/半行）' if bad else '')) + return done + + +def archive_output_dir(carry_data: bool = False) -> None: + """启动时把已存在的 OUTPUT_DIR 整个 mv 走，保证本 run 写入空目录。 + + 为何必需：`e18_sft_dataset.jsonl` / `collect_log.jsonl` 都是 `open(..., 'a')` 追写。 + 没有这一步时，重启一次就把新旧 run 的样本焊在同一个文件里，而且不报错。 + 用 mv 而不是删：旧 run 的样本是可复用的分析素材。 + 沙箱自检缓存（kod_broken_tasks.json）会搬回新目录 —— 那是纯函数结果，重跑一次很贵。 + + carry_data=True（续跑）时额外把三个产物 jsonl 也复制回新目录，于是新 run 接着往后面 + 追写、总数连续。用 copy2 而不是 move：归档副本保留完整快照，万一续跑又崩了还能回溯。 + 注意：续跑后同一份数据里会存在多个 `run` 值，离线分析要按 run 字段分组看。 + """ + if not os.path.isdir(OUTPUT_DIR) or not os.listdir(OUTPUT_DIR): + os.makedirs(OUTPUT_DIR, exist_ok=True) + return + stamp = time.strftime('%m%d-%H%M%S', time.localtime(os.path.getmtime(OUTPUT_DIR))) + dst = f'{OUTPUT_DIR}.bak-{stamp}' + i = 1 + while os.path.exists(dst): + dst = f'{OUTPUT_DIR}.bak-{stamp}-{i}' + i += 1 + shutil.move(OUTPUT_DIR, dst) + os.makedirs(OUTPUT_DIR, exist_ok=True) + logger.info(f'[init] 旧输出目录已归档 -> {dst}') + cache = os.path.join(dst, 'kod_broken_tasks.json') + if os.path.exists(cache): + shutil.copy2(cache, os.path.join(OUTPUT_DIR, 'kod_broken_tasks.json')) + logger.info('[init] 沙箱自检缓存已搬回新目录（避免重跑）') + if carry_data: + for name in ('e18_sft_dataset.jsonl', 'e18_candidates.jsonl', 'collect_log.jsonl'): + src = os.path.join(dst, name) + if os.path.exists(src): + shutil.copy2(src, os.path.join(OUTPUT_DIR, name)) + logger.info('[resume] 旧产物已复制回新目录，本 run 接着追写') + + +def main(): + t_start = time.time() + # ⭐ 顺序强制：读已采 id 必须在归档**之前**，否则 OUTPUT_DIR 已经被 mv 走、读到空集。 + done_ids = resume_done_ids() if KOD_RESUME else set() + n_done = count_existing_samples() if KOD_RESUME else 0 + archive_output_dir(carry_data=KOD_RESUME) + checker = build_checker() + # eval_size=0：不训练就没有 eval 的意义，题全进采集池。 + train_dataset, _ = load_records(SEED, 0, OUTPUT_DIR) + # ⭐ 分片必须在 resume 过滤**之前**做，且用 crc32(data_id) 而不是下标： + # 两台机器无共享存储，只能靠纯函数保证不重叠。zlib.crc32 跨进程/跨平台稳定 + # （而 hash() 受 PYTHONHASHSEED 影响，每次启动都不同 —— 用它会造成重叠+遗漏）。 + if SHARD_N > 1: + before = len(train_dataset) + train_dataset.filter( + lambda r: zlib.crc32(str(r['data_id']).encode()) % SHARD_N == SHARD_ID) + logger.info(f'[shard] {SHARD_ID}/{SHARD_N}: 题池 {before} -> {len(train_dataset)}') + if done_ids: + before = len(train_dataset) + train_dataset.filter(lambda r: r['data_id'] not in done_ids) + logger.info(f'[resume] 题池剔除已采 {before - len(train_dataset)} 题 -> 剩 {len(train_dataset)}') + logger.info(f'[data] 采集池 ={len(train_dataset)} 题') + rt = build_runtime(checker, MultiDiagCache()) + logger.info(f'E18-collect start: chunk={CHUNK_SIZE} n_skills={N_SKILLS} ' + f'bare_rollouts={BARE_ROLLOUTS} exec_rollouts={EXEC_ROLLOUTS} ' + f'min_pass_gain={MIN_PASS_GAIN:.3g} ' + f'use_traj={USE_TRAJ}(n={TRAJ_N},max_chars={TRAJ_MAX_CHARS}) ' + f'shard={SHARD_ID}/{SHARD_N} ' + f'target={TARGET_SAMPLES} gpus={SKILL_SAMPLER_GPUS}+{BASE_SAMPLER_GPUS} ' + f'resume={int(KOD_RESUME)} n_done={n_done} ' + f'output={OUTPUT_DIR}') + + # ⭐ n_total 从已有数量起算，不是 0：它是 TARGET_SAMPLES 的比较对象，从 0 起算会变成 + # 「再采 TARGET_SAMPLES 条」而不是「凑到 TARGET_SAMPLES 条」。 + n_total, ci = n_done, 0 + log_path = os.path.join(OUTPUT_DIR, 'collect_log.jsonl') + with open(log_path, 'a', encoding='utf-8') as log_fh: + loader = DataLoader(dataset=train_dataset, batch_size=CHUNK_SIZE, num_workers=0, + shuffle=True, drop_last=False, + generator=torch.Generator().manual_seed(SEED)) + for chunk in loader: + if TARGET_SAMPLES and n_total >= TARGET_SAMPLES: + break + if MAX_CHUNKS and ci >= MAX_CHUNKS: + break + t0 = time.time() + accepted, metrics = collect_chunk(rt, chunk, ci) + dump_dataset(accepted) + n_total += len(accepted) + row = {'chunk': ci, 'run': RUN_ID, 'seconds': round(time.time() - t0, 1), + 'n_collected': float(n_total), **metrics} + log_fh.write(json.dumps(row, ensure_ascii=False) + '\n') + log_fh.flush() + logger.info(f'[c{ci}] collected={n_total}' + + (f'/{TARGET_SAMPLES}' if TARGET_SAMPLES else '') + + ' ' + ' '.join(f'{k}={v:.4g}' for k, v in row.items() + if isinstance(v, float))) + ci += 1 + logger.info(f'[完成] 共 {n_total} 条，{ci} 个 chunk，' + f'耗时 {(time.time() - t_start) / 60:.1f} 分钟 -> {OUTPUT_DIR}') + + +if __name__ == '__main__': + main() diff --git a/cookbook/human_e18/e18_kodcode.py b/cookbook/human_e18/e18_kodcode.py new file mode 100644 index 000000000..bc753e500 --- /dev/null +++ b/cookbook/human_e18/e18_kodcode.py @@ -0,0 +1,404 @@ +# -*- coding: utf-8 -*- +"""KodCode-V1 数据域适配：与 `e23_bcb.py` **同契约**的加载 + 沙箱判分。 + +为什么另起一个文件而不改 e23_bcb：BCB 与 KodCode 的单测框架不同（unittest vs pytest）、 +题面/入口点的来源字段也不同，但**对上层的接口必须逐字一致** —— `e18_rejection_sft.py` +只认 `load_records / judge_seqs / empty_roll / run_tests` 这组签名和 `{'data_id', +'problem', 'reference_answer'}` 这个记录形状。保持接口一致，换域时上层零改动。 + +与 e23_bcb 的对齐点（改任何一处都会让两域的 pass_rate 不可比）： +* `run_tests` 返回 `{'passed', 'kind', 'error'}`，`kind` 取值集合完全相同： + `pass / no_code / no_entry / timeout / assertion / exception / import_or_syntax`。 + `e18_multidiag._signature` 拿 kind 做失败签名，取值不一致会让诊断缓存串味。 +* `judge_seqs` 同 `(task_id, code)` 只判一次、线程池并发、返回 roll 的字段集相同。 +* `_trim_err` 只保留失败测试名与异常行，并把随机临时目录名归一化成 `` —— + 否则同一个失败在两次运行里字符串不同，`_signature` 会算出两个签名、缓存永远不命中。 +* 参考解答跑不过自己单测的题一律剔除（BCB 实测 ~7.5%，KodCode 实测 ~5%）， + 自检结果落盘缓存。这类题不是模型的错，留着会把 base_pass_rate 永久压低。 + +KodCode 特有的两点： +1. 单测是 pytest 风格且 **181/200 靠 `from solution import X`** 取被测函数，所以沙箱里必须 + 把提交代码写成 `solution.py`（而不是 BCB 那样把代码和 test 拼进同一个文件）。 +2. 自带 `gpt_pass_percentage`（教师多次尝试的通过率），是**免费的先验难度**。E18 原先要靠 + 跑 8 次 bare rollout 才能找出难题，这里可以直接按阈值筛，省掉这部分 GPU。 +""" +import ast as _ast +import json +import os +import re +import shutil +import subprocess +import sys +import tempfile +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Optional, Tuple + +from twinkle import get_logger +from twinkle.dataset import Dataset, DatasetMeta + +logger = get_logger() + +# ========== 配置 ========== +KOD_DATASET = os.environ.get('KOD_DATASET', 'ms://AI-ModelScope/KodCode-V1') +KOD_SUBSET = os.environ.get('KOD_SUBSET', 'default') +KOD_SPLIT = os.environ.get('KOD_SPLIT', 'train') +# 难度窗口：只留「教师也常做错、但并非无解」的题。 +# 上界 0.3 -> executor 大概率失败（有诊断可采）；下界 >0 -> 排除疑似不可解。 +KOD_MAX_PASS_PCT = float(os.environ.get('KOD_MAX_PASS_PCT', 0.3)) +KOD_MIN_PASS_PCT = float(os.environ.get('KOD_MIN_PASS_PCT', 0.0)) +KOD_MAX_TASKS = int(os.environ.get('KOD_MAX_TASKS', 0)) # 0 = 不截断 +# ⭐ 默认关沙箱自检：全量 73747 题跑一遍参考解答要 ~30 小时（子进程）。 +# 关掉的代价：参考解答自己都跑不过单测的坏题（实测 ~12.5%）会留在题池里， +# 但它们在采集时会自然显形 —— base_pass_rate 恒为 0、且任何 skill 都拿不到 +# +MIN_PASS_GAIN，于是 select_winner 返回 None、不入池。只浪费 rollout，不污染数据集。 +KOD_SELFCHECK = os.environ.get('KOD_SELFCHECK', '0') == '1' +# ⭐ TEST_WORKERS 默认 96，而不是继承 BCB 的 24：判分是子进程，与 GPU 采样**串行**， +# 每 chunk 要跑 CHUNK_SIZE*(BARE_ROLLOUTS + N_SKILLS*EXEC_ROLLOUTS) ≈ 2300 次，并发不够就直接拆 GPU 空转。 +# BCB 用 24 是因为它的单测要 import pandas/sklearn/matplotlib（单次 1-3s、内存大）； +# KodCode 是纯算法题，单测 0.05-0.3s、几乎不导包，可以开得高得多。 +# 上限卡在 min(96, 核数一半)：留余量给 vLLM 的调度/集合线程，别把宿主打满反而拖慢采样。 +TEST_WORKERS = int(os.environ.get( + 'TEST_WORKERS', max(24, min(96, (os.cpu_count() or 24) // 2)))) +TEST_TIMEOUT = int(os.environ.get('TEST_TIMEOUT', 60)) + +_FENCE_RE = re.compile(r'```(?:python|py)?\s*\n(.*?)```', re.S) +_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') + + +# ========== 文本处理（与 e23_bcb 逐字一致） ========== +def after_think(text: str) -> str: + """只取 之后的正文；没有闭合标签就原样返回。""" + idx = text.rfind('') + return text[idx + len(''):] if idx >= 0 else text + + +def clean_text(decoded: Optional[str]) -> str: + return _SPECIAL_TOKEN_RE.sub('', decoded or '').strip() + + +def extract_code(text: str) -> str: + """取最后一个完整的 ``` 代码块；没有围栏时退化成整段正文。 + + 取**最后一个**而不是第一个：模型常先给一版草稿再给最终版，最后一个才是它的结论。 + """ + body = after_think(text) + blocks = _FENCE_RE.findall(body) + if blocks: + return blocks[-1].strip() + # 没有围栏：可能是 nothink 直出代码。剔掉明显的自然语言行后返回。 + return body.strip() + + +def extract_skill(text: str) -> str: + """取 ... 里的内容；没有标签时返回空串（视为格式失败）。""" + body = after_think(text) + m = re.search(r'(.*?)', body, re.S | re.IGNORECASE) + return m.group(1).strip() if m else '' + + +# ========== 沙箱判分 ========== +# ⭐ 纯 pytest 驱动，但用插件把 (n_tests, n_fail, n_err) 拿回来 —— 与 e23_bcb 的 +# `__BCB__ n f e` 输出契约同形，好让 kind 的判定规则完全一致。 +# 只统计 when=='call'：setup/teardown 阶段的失败算 error（多半是 import 不了 solution）。 +# +# ⭐ 断言 vs 异常的区分必须走 `report.longrepr.reprcrash.message`，**不能**用 +# `'AssertionError' in str(longrepr)`：pytest 默认开断言重写（assertion rewriting）， +# 失败摘要长这样 `E assert -1 == 3`，整段里根本没有 "AssertionError" 这个词， +# 于是所有断言失败都会被误判成 exception。实测踩过：断言不符返回了 kind='exception'。 +# kind 错了会让 `e18_multidiag._signature` 的失败签名串味、rubric 缓存失效。 +_RUNNER = r""" +import sys, pytest + + +class _Collect: + def __init__(self): + self.n_tests = self.n_fail = self.n_err = 0 + + @staticmethod + def _is_assertion(report): + crash = getattr(getattr(report, 'longrepr', None), 'reprcrash', None) + msg = getattr(crash, 'message', '') or '' + # 断言重写后首行是 "assert ..."；未重写时是 "AssertionError: ..."。 + return msg.startswith('assert') or msg.startswith('AssertionError') + + def pytest_runtest_logreport(self, report): + if report.when == 'call': + self.n_tests += 1 + if report.failed: + if self._is_assertion(report): + self.n_fail += 1 + else: + self.n_err += 1 + elif report.failed: + self.n_err += 1 + + +c = _Collect() +rc = pytest.main(['-q', '--no-header', '-p', 'no:cacheprovider', + '--tb=short', 'test_solution.py'], plugins=[c]) +print('__KOD__', c.n_tests, c.n_fail, c.n_err) +sys.exit(0 if int(rc) == 0 else 1) +""" + + +def _trim_err(err: str, limit: int = 1600) -> str: + """保留失败测试名与异常行，砍掉冗长 traceback 帧 —— 这是喂给 rubric 的客观证据。 + 随机临时目录名换成 ，否则同一个失败在两次运行里看起来不一样 + （`e18_multidiag._signature` 会因此算出不同签名，诊断缓存永久不命中）。""" + err = re.sub(r'/tmp/kod_[A-Za-z0-9_]+', '', err or '') + lines = [ln for ln in err.splitlines() if ln.strip()] + keep = [ln for ln in lines + if ln.startswith(('FAILED', 'FAIL:', 'ERROR:', 'AssertionError', 'Traceback', 'E ')) + or re.match(r'^\w*(Error|Exception|Warning)\b', ln.strip()) + or ', in ' in ln] + return '\n'.join(keep or lines[-25:])[-limit:] + + +def run_tests(code: str, payload: Dict[str, Any], timeout: int = TEST_TIMEOUT) -> Dict[str, Any]: + """子进程里跑「提交代码(solution.py) + 官方 test(test_solution.py)」。 + + -> {'passed', 'kind', 'error'}，kind 取值与 e23_bcb.run_tests 完全一致。 + + 与 BCB 的唯一实质差异：代码单独落成 `solution.py`，因为 KodCode 的单测靠 + `from solution import X` 取被测函数，拼进同一个文件会 ImportError。 + """ + if not code.strip(): + return {'passed': False, 'kind': 'no_code', 'error': 'no parseable code block'} + entry = payload.get('entry_point') or '' + if entry and entry not in code: + return {'passed': False, 'kind': 'no_entry', + 'error': f'function {entry} is not defined in the submitted code'} + tmp = tempfile.mkdtemp(prefix='kod_') + try: + with open(os.path.join(tmp, 'solution.py'), 'w', encoding='utf-8') as f: + f.write(code) + with open(os.path.join(tmp, 'test_solution.py'), 'w', encoding='utf-8') as f: + f.write(payload['test']) + with open(os.path.join(tmp, '_run.py'), 'w', encoding='utf-8') as f: + f.write(_RUNNER) + env = dict(os.environ, MPLBACKEND='Agg', PYTHONHASHSEED='0', OMP_NUM_THREADS='1', + MKL_NUM_THREADS='1', TOKENIZERS_PARALLELISM='false') + env.pop('CUDA_VISIBLE_DEVICES', None) + # cwd=tmp 让 `from solution import X` 能找到同目录的 solution.py。 + try: + p = subprocess.run([sys.executable, '_run.py'], cwd=tmp, env=env, timeout=timeout, + capture_output=True, text=True, errors='replace') + except subprocess.TimeoutExpired: + return {'passed': False, 'kind': 'timeout', + 'error': f'the tests did not finish within {timeout}s'} + n_tests = n_fail = n_err = 0 + for line in (p.stdout or '').splitlines(): + if line.startswith('__KOD__'): + _, a, b, c = line.split() + n_tests, n_fail, n_err = int(a), int(b), int(c) + if p.returncode == 0 and n_tests > 0: + return {'passed': True, 'kind': 'pass', 'error': ''} + kind = 'assertion' if n_fail else ('exception' if n_err else 'import_or_syntax') + merged = ((p.stdout or '') + '\n' + (p.stderr or '')).replace(tmp, '') + return {'passed': False, 'kind': kind, 'error': _trim_err(merged)} + finally: + shutil.rmtree(tmp, ignore_errors=True) + + +def empty_roll() -> Dict[str, Any]: + return {'correct': False, 'stop_reason': 'empty', 'gen_tokens': 0, 'text': '', 'code': '', + 'kind': 'no_code', 'error': ''} + + +def judge_seqs(pairs: List[Tuple[Any, Dict[str, Any]]]) -> List[Dict[str, Any]]: + """[(采样 sequence 或 None, payload)] -> rolls。所有判分都汇合到这里。 + + 必须批量：单测是子进程，一个 chunk 几百次判分串行会比同 chunk 的 GPU 时间还长一个量级。 + 同 (task_id, code) 只跑一次 —— T=0 的 executor 经常对同一题产出逐字相同的代码。 + """ + rolls: List[Dict[str, Any]] = [] + keys: List[Optional[Tuple[str, str]]] = [] + jobs: Dict[Tuple[str, str], Dict[str, Any]] = {} + for seq, payload in pairs: + if seq is None: + rolls.append(empty_roll()) + keys.append(None) + continue + text = clean_text(getattr(seq, 'decoded', '') or '') + code = extract_code(text) + key = (payload['task_id'], code) + rolls.append({'correct': False, 'stop_reason': getattr(seq, 'stop_reason', None), + 'gen_tokens': len(getattr(seq, 'tokens', None) or []), + 'text': text, 'code': code, 'kind': None, 'error': ''}) + keys.append(key) + jobs.setdefault(key, payload) + if jobs: + todo = list(jobs) + with ThreadPoolExecutor(max_workers=max(1, min(TEST_WORKERS, len(todo)))) as ex: + verdicts = dict(zip(todo, ex.map(lambda k: run_tests(k[1], jobs[k]), todo))) + for roll, key in zip(rolls, keys): + v = verdicts.get(key) if key is not None else None + if v is not None: + roll['correct'] = bool(v['passed']) + roll['kind'], roll['error'] = v['kind'], v['error'] + return rolls + + +# ========== 数据 ========== +# reference_answer 的字段集：判分需要的一切。与 e23_bcb 的 _PAYLOAD_KEYS 同名对齐， +# 上层（e18_select / e18_multidiag / dump_dataset）拿到的 key 才一致。 +_PAYLOAD_KEYS = ('task_id', 'entry_point', 'test', 'code_prompt', 'doc_struct', + 'canonical_solution') + + +def _entry_point(row: Dict[str, Any]) -> str: + """从 test_info 拿被测函数名。拿不到就回落到 test 里的 `from solution import X`。""" + ti = row.get('test_info') + if ti is not None: + try: + items = list(ti) if not isinstance(ti, str) else _ast.literal_eval(ti) + for it in items: + name = (it or {}).get('function_name') + if name: + return str(name) + except Exception: + pass + m = re.search(r'from\s+solution\s+import\s+([A-Za-z_]\w*)', row.get('test') or '') + return m.group(1) if m else '' + + +def _code_prompt(row: Dict[str, Any]) -> str: + """函数签名，用于给 executor 固定入口点（对齐 BCB 的 code_prompt 语义）。""" + ti = row.get('test_info') + if ti is not None: + try: + items = list(ti) if not isinstance(ti, str) else _ast.literal_eval(ti) + for it in items: + decl = (it or {}).get('function_declaration') + if decl: + return str(decl) + except Exception: + pass + return '' + + +def _usable(row: Dict[str, Any]) -> bool: + """能进题池的最低门槛。 + + 要求 test 通过 `from solution import` 取函数：11.7% 的题直接裸调函数名， + 在「代码写进 solution.py」的沙箱布局下必然 NameError —— 那是 harness 不兼容， + 不是模型的错，留着会把 base_pass_rate 永久压低。 + """ + test = row.get('test') or '' + if 'def test_' not in test: + return False + if not re.search(r'from\s+solution\s+import|import\s+solution\b', test): + return False + return bool((row.get('solution') or '').strip()) and bool(_entry_point(row)) + + +# ⭐ 题面尾部必须追加函数签名：实测只有 **8%** 的 KodCode question 提到了被测函数名， +# 而单测靠 `from solution import ` 取函数。不补签名的后果：executor 把函数叫成 +# 任何名字都算错，92% 的题无论 skill 好坏都是 0 分 —— pass_rate 全城 0、整个采集废掉。 +# BCB 不需要这一步是因为它的 instruct_prompt 自带 `def task_func(...)` 骨架。 +_SIG_HINT = ('\n\nYou should write self-contained code starting with:\n```\n{decl}\n```') + + +def _problem_text(row: Dict[str, Any]) -> str: + """题面 = question + 函数签名（签名已在题面里就不重复追加）。""" + q = row.get('question') or '' + decl = _code_prompt(row) + if not decl: + return q + if decl.strip() in q: + return q + return q + _SIG_HINT.format(decl=decl.strip()) + + +def _to_record(batch: Dict[str, List]) -> Dict[str, List]: + """原始 KodCode 列 -> {'data_id', 'problem', 'reference_answer'}。 + + Dataset.map 强制 batched=True，所以这里收发的都是列式 batch。 + """ + n = len(batch['question_id']) + rows = [{k: batch[k][i] for k in batch} for i in range(n)] + return { + 'data_id': [str(r['question_id']) for r in rows], + 'problem': [_problem_text(r) for r in rows], + 'reference_answer': [{ + 'task_id': str(r['question_id']), + 'entry_point': _entry_point(r), + 'test': r['test'], + 'code_prompt': _code_prompt(r), + 'doc_struct': '', + 'canonical_solution': r['solution'], + # 教师先验难度：保留下来供离线分析（不参与判分）。 + 'gpt_pass_percentage': float(r.get('gpt_pass_percentage') or 0.0), + 'gpt_difficulty': r.get('gpt_difficulty') or '', + } for r in rows], + } + + +def _broken_tasks(ds: Dataset, output_dir: str) -> set: + """参考解答跑不过自己的单测 = 数据缺陷或沙箱不可判定，不是模型的错（实测约 5%）。 + 自检一次后落盘缓存，题数不变则复用。必须在 map 之后调用（读的是 reference_answer）。""" + path = os.path.join(output_dir, 'kod_broken_tasks.json') + if os.path.exists(path): + try: + with open(path, encoding='utf-8') as f: + c = json.load(f) + if int(c.get('n_tasks', -1)) == len(ds): + logger.info(f'[data] 复用沙箱自检缓存：剔除 {len(c["broken"])} 道') + return set(c['broken']) + except Exception as exc: + logger.warning(f'[data] 读取 {path} 失败（{exc}），重跑自检') + logger.info(f'[data] 沙箱自检：{len(ds)} 道题跑参考解答（一次性，之后走缓存）…') + rows = [ds[i] for i in range(len(ds))] + jobs = [(r['reference_answer']['canonical_solution'], r['reference_answer']) for r in rows] + with ThreadPoolExecutor(max_workers=max(1, min(TEST_WORKERS, len(jobs)))) as ex: + vers = list(ex.map(lambda p: run_tests(p[0], p[1]), jobs)) + broken = {r['data_id'] for r, v in zip(rows, vers) if not v['passed']} + with open(path, 'w', encoding='utf-8') as f: + json.dump({'n_tasks': len(ds), 'broken': sorted(broken)}, f, indent=1) + logger.info(f'[data] 自检完成：剔除 {len(broken)}/{len(rows)} ' + f'({100.0 * len(broken) / max(1, len(rows)):.1f}%)') + return broken + + +def load_records(seed: int, eval_size: int, + output_dir: str) -> Tuple[Dataset, List[Dict[str, Any]]]: + """-> (train_dataset, eval_records)，每条记录是 {'data_id', 'problem', 'reference_answer'}。 + + 签名与 `e23_bcb.load_records` 逐字一致，上层可直接换 import。 + + 与 BCB 的差异：KodCode 自带 `gpt_pass_percentage`，所以**先按难度窗口过滤**再自检 —— + 自检要跑一遍全部参考解答（子进程，很贵），先筛掉容易题能省掉大部分开销。 + """ + ds = Dataset(DatasetMeta(KOD_DATASET, subset_name=KOD_SUBSET, split=KOD_SPLIT)) + n_raw = len(ds) + + ds.filter(lambda r: KOD_MIN_PASS_PCT < float(r.get('gpt_pass_percentage') or 0.0) + <= KOD_MAX_PASS_PCT) + n_hard = len(ds) + ds.filter(_usable) + logger.info(f'[data] KodCode: 全集 {n_raw}，难度窗口 ' + f'({KOD_MIN_PASS_PCT}, {KOD_MAX_PASS_PCT}] 保留 {n_hard}、' + f'harness 不兼容剔除 {n_hard - len(ds)} -> {len(ds)}') + if KOD_MAX_TASKS and len(ds) > KOD_MAX_TASKS: + # ⭐ 必须用 filter 而不是 `ds.dataset = ds.dataset.select(...)`：Dataset.map 内部读的是 + # `self.datasets`（未截断的副本）并回写 `self.dataset`，直接赋值 self.dataset 会在 + # 下一句 map 里被静默覆盖 —— 实测踩过：截断到 40 题后自检仍在跑 73747 题。 + keep = set(ds.dataset.shuffle(seed=seed)['question_id'][:KOD_MAX_TASKS]) + ds.filter(lambda r: r['question_id'] in keep) + logger.info(f'[data] KOD_MAX_TASKS 截断 -> {len(ds)}') + + ds.map(_to_record, remove_columns=ds.dataset.column_names) + + broken = _broken_tasks(ds, output_dir) if KOD_SELFCHECK else set() + if broken: + ds.filter(lambda r: r['data_id'] not in broken) + logger.info(f'[data] 剔除参考解答自己跑不过单测的题 {len(broken)} 道 -> 可用 {len(ds)}') + elif not KOD_SELFCHECK: + logger.info(f'[data] 跳过沙箱自检（KOD_SELFCHECK=0），题池 {len(ds)}；' + f'坏题会在采集时因 base_pass_rate=0 自然不入池') + + shuffled = ds.dataset.shuffle(seed=seed) + n_eval = min(eval_size, len(shuffled)) if eval_size > 0 else 0 + eval_records = list(shuffled.select(range(n_eval))) + train_dataset = Dataset(DatasetMeta(data=shuffled.select(range(n_eval, len(shuffled))))) + return train_dataset, eval_records diff --git a/cookbook/human_e18/e18_multidiag.py b/cookbook/human_e18/e18_multidiag.py new file mode 100644 index 000000000..d1166a705 --- /dev/null +++ b/cookbook/human_e18/e18_multidiag.py @@ -0,0 +1,245 @@ +"""E18 的多轨迹诊断：给**每一条失败的 rollout** 各诊一次，再合并成一份 rubric。 + +为什么需要它（E23 的 `RubricCache.get_or_diagnose` 不能直接复用）： + 1. 它的缓存键是 `RUBRIC_VERSION + _PROMPT_KEY + data_id`，**不含失败轨迹内容**。同一题调 + N 次会全部命中第一次的结果 —— 想「每条 rollout 判一次」在那个键下是做不到的。 + 2. 它一题只产出一个决定性根因，且 `_format` 把 secondary 渲染成 + `ALSO OFF (do not write about these)`，**主动禁止** skill-gen 覆盖第二处。 + 实测 E23 有 73% 的零 reward 组正是「修对第一处、挂在第二处」。 + +本模块只做加法，不改 `e23_rubric.py`（它是 E18/E23 共用、必须逐字同源的判分/诊断层）： +复用其 `diag_query` / `diag_segment` / `_validate` / `_CLASS_SHORT`，但换一个**按失败内容** +分桶的缓存键，并自己渲染合并文本。 + +⭐ 去重按 `(kind, 报错签名)` 而不是按 rollout 逐条：8 次采样里常有 5 次是同一个 assertion， +逐条诊断纯属浪费 API。同一签名只诊一次，但记下它出现了几次（`n_seen`），合并时按频次排序 —— +出现 5 次的根因显然比只出现 1 次的更该先修。 +""" +import hashlib +import json +import os +import re +import threading +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Optional, Tuple + +from twinkle import get_logger + +import e23_rubric as R + +logger = get_logger() + +# 每题最多诊断几种**不同**的失败签名。3 与 e23_rubric 的 MAX_INDEPENDENT_CAUSES 一致： +# 超过 3 个独立根因的题，教师侧本来就判为不可救（诊断会互相矛盾，skill 也写不下）。 +MAX_DIAG_PER_TASK = int(os.environ.get('MAX_DIAG_PER_TASK', 3)) +# 一个签名至少要出现几次才值得诊断。默认 1 = 全诊；设 2 可以过滤掉只出现一次的偶发错误。 +MIN_SIGNATURE_COUNT = int(os.environ.get('MIN_SIGNATURE_COUNT', 1)) +# 并行度按**题**算：24 个线程各认领一道题，题内的多个签名仍串行，所以同时在飞的 HTTP +# request 就是 24。默认从 8 提到 24：纯 API 等待不占 GPU也不占 CPU，8 并行下一个 +# 64 题 chunk 的诊断阶段实测要 ~26 分钟，而这段时间 8 张卡全部空转。 +# 上限受教师侧限流约束，碰到 429 就把这个值调回。 +DIAG_WORKERS = int(os.environ.get('RUBRIC_WORKERS', 24)) + +# 失败签名：只取**结构化字段**（哪个测试挂了 + 什么异常类），不用报错正文。 +# +# ⭐ 为何不对报错正文做归一化（前一版的做法）：那靠的是「把数字/路径/引号内容逐个替成 +# 占位符」，而异常消息里的变量形式永远枚不完（裸数值、列宽对齐的空白、repr 片段……）； +# 漏一个，同一个 bug 就被当成 N 个不同根因，白花 N 倍教师 API 且 N 条诊断在说同一件事。 +# 现在只依赖两个结构化信号，行为可预测，不随报错排版变化。 +# +# 两个信号都来自 unittest 的固定输出格式，且 e23_bcb._trim_err 保证保留（它只留 +# FAIL:/ERROR:/Traceback/异常类名/带 ', in ' 的帧行）： +# * 失败的测试方法名 —— 区分「同一异常但挂在不同测试上」（那是不同根因）； +# * 异常类名 —— 区分 KeyError / AttributeError / AssertionError。 +# 只用 kind + 异常类会把前者错误合并，所以两个都要。 +_TEST_RE = re.compile(r'^(?:FAIL|ERROR):\s*(\w+)', re.M) +_EXC_RE = re.compile(r'^([A-Za-z_][\w.]*(?:Error|Exception|Warning))\b', re.M) + + +def _signature(roll: Dict[str, Any]) -> str: + """失败轨迹 -> 稳定签名 = kind + 失败测试名集合 + 异常类集合。 + + 集合都排序去重，所以「两个测试挂了」不会因报错顺序不同而分成两个签名。 + + ⭐ 拿不到任何结构化信号时（如 kind='timeout' / 'no_code'，根本没跑到 unittest）， + 两个集合都为空，签名退化成单独的 kind —— 这正是想要的：同一题的 8 次超时是 + 同一件事，只该诊一次。 + """ + kind = str(roll.get('kind') or 'unknown') + err = str(roll.get('error') or '') + tests = ','.join(sorted(set(_TEST_RE.findall(err)))) + excs = ','.join(sorted(set(_EXC_RE.findall(err)))) + return f'{kind}\x00{tests}\x00{excs}' + + +def bucket_failures(rolls: List[Dict[str, Any]]) -> List[Tuple[Dict[str, Any], int]]: + """把一题的 rolls 按失败签名分桶，返回 [(代表 roll, 出现次数), ...]，按次数降序。 + + 只取失败的 roll。代表 roll 用该桶里第一条 —— 同签名意味着同一组 (测试, 异常类)， + 但报错正文仍可能略有差异（具体数值）；教师看到的是这条代表的完整报错，信息不丢。 + """ + buckets: Dict[str, Dict[str, Any]] = {} + for r in rolls: + if r.get('correct'): + continue + sig = _signature(r) + b = buckets.get(sig) + if b is None: + buckets[sig] = {'roll': r, 'n': 1} + else: + b['n'] += 1 + out = [(b['roll'], b['n']) for b in buckets.values()] + out.sort(key=lambda t: -t[1]) # 高频根因优先 + return out + + +class MultiDiagCache: + """按 (data_id, 失败签名) 缓存单条诊断；磁盘格式与 e23_rubric 的缓存文件同构但独立成文件。 + + 独立文件的理由：键的语义不同（这里含失败签名），混进同一个文件会让 E23 的缓存读取拿到 + 对不上的条目。E23 那份缓存**完全不动**，两个实验各自可复现。 + """ + + def __init__(self, path: str = None): + self.path = path or os.path.join( + os.path.dirname(os.path.abspath(__file__)), + 'multidiag_cache_code_v1.jsonl') + self._idx: Dict[str, Any] = {} + import collections + self.stats: collections.Counter = collections.Counter() + if os.path.exists(self.path): + with open(self.path, encoding='utf-8') as f: + for line in f: + try: + rec = json.loads(line) + self._idx[rec['key']] = rec['value'] + except Exception: + continue + logger.info(f'[multidiag] 缓存载入 {len(self._idx)} 条：{self.path}') + self._fh = open(self.path, 'a', encoding='utf-8') + # ⭐ _put 会被 DIAG_WORKERS 个线程并发调用，write + flush 两步不是原子的：无锁时 + # 两条记录会交错成半行，下次启动 json.loads 解不开就默默丢掉（except: continue）， + # 表现是「明明诊过却反复花 API 钱」且无任何报错。并行度提到 24 后这个概率不再可忽。 + self._lock = threading.Lock() + + def _put(self, key: str, value: Any) -> None: + line = json.dumps({'key': key, 'value': value}, ensure_ascii=False) + '\n' + with self._lock: + self._idx[key] = value + self._fh.write(line) + self._fh.flush() + + def _one(self, checker, record: Dict[str, Any], roll: Dict[str, Any]) -> Optional[Dict]: + """诊断单条失败轨迹，返回 validate 过的 diag dict（不可救/失败返回 None）。 + + 与 e23_rubric.get_or_diagnose 的差别只有缓存键：这里把失败签名放进键，所以同一题的 + 不同失败模式各占一个槽位。校验逻辑直接复用 R._validate，保持判据完全一致。 + """ + sig = _signature(roll) + key = hashlib.md5( + f"{R.RUBRIC_VERSION}\x00{R._PROMPT_KEY}\x00" + f"{record.get('data_id', '')}\x00{sig}".encode('utf-8')).hexdigest() + if key in self._idx: + cached = self._idx[key] + if not cached: + self.stats['hit_dropped'] += 1 + return None + self.stats['hit'] += 1 + return cached + query = R.diag_query(record['problem'], record['reference_answer']) + try: + obj = checker.classify(query, R.diag_segment(roll)) + except Exception as exc: + logger.warning(f'[multidiag] classify error: {exc}') + obj = None + if obj is None: + # 同 e23_rubric：API 故障**绝不缓存**，否则一次抖动会永久丢掉这个失败模式。 + self.stats['api_fail'] += 1 + return None + diag = R._validate(obj, query) + if diag is None: + self._put(key, None) # 稳定判决：这个失败模式教师给不出可用分类 + self.stats['dropped_unaddressable'] += 1 + return None + self.stats['ok'] += 1 + self.stats[f"class_{diag['class']}"] += 1 + self._put(key, diag) + return diag + + def diagnose_task(self, checker, record: Dict[str, Any], + rolls: List[Dict[str, Any]]) -> str: + """一题的多失败模式诊断 -> 合并后的 rubric 文本（无可用诊断返回 ''）。""" + buckets = [(r, n) for r, n in bucket_failures(rolls) if n >= MIN_SIGNATURE_COUNT] + if not buckets: + return '' + buckets = buckets[:MAX_DIAG_PER_TASK] + diags = [] + for roll, n in buckets: + d = self._one(checker, record, roll) + if d: + diags.append((d, n)) + return merge_diags(diags, n_rollouts=len(rolls)) + + def diagnose_many(self, checker, + jobs: List[Tuple[Dict[str, Any], List[Dict[str, Any]]]]) -> List[str]: + """并行版。jobs = [(record, rolls), ...]，返回对齐的 rubric 文本列表。 + + 并行度按**题**而不是按签名：一个线程认领一道题，题内的最多 + MAX_DIAG_PER_TASK 次 classify 仍串行，所以同时在飞的 request 数 = min(DIAG_WORKERS, 题数)。 + 题数通常远大于 DIAG_WORKERS，所以实际并行就是 DIAG_WORKERS。 + + 不把签名也展平成任务（那会再快 ~3 倍）的原因不是缓存安全 —— _put 已加锁； + 而是签名数预先不知道，展平后难以把结果按题对齐回去，且同题的多条诊断本身 + 就要合并。纯 API 等待不占 GPU。 + """ + if not jobs: + return [] + workers = max(1, min(DIAG_WORKERS, len(jobs))) + with ThreadPoolExecutor(max_workers=workers) as ex: + return list(ex.map(lambda j: self.diagnose_task(checker, j[0], j[1]), jobs)) + + def close(self): + self._fh.close() + + +def merge_diags(diags: List[Tuple[Dict[str, Any], int]], n_rollouts: int = 0) -> str: + """把多条诊断拼成一份给 skill-gen 的文本。 + + ⭐ 与 e23_rubric._format 的两处关键差别： + 1. **不再输出** `ALSO OFF (do not write about these)`。那条禁令是 E23 单根因口径的产物， + 而本模块的全部目的就是让 skill 覆盖多处根因 —— 留着它会自相矛盾。 + 2. 带上 `seen k/M times` 频次。skill-gen 据此知道哪个根因更普遍、该先写哪个； + 只翻车 1/8 次的偶发问题不该和翻车 5/8 次的主因同等对待。 + + evidence 仍然**不进**文本（与 _format 一致）：它是单测报错原文，断言 diff 里带期望值， + 是最强的答案泄漏通道。 + """ + if not diags: + return '' + if len(diags) == 1: + d, n = diags[0] + head = [f"DECISIVE FAILURE: {d['class']} — {R._CLASS_SHORT[d['class']]}", + f"WHAT WENT WRONG: {d['reason']}", + f"PRIOR THAT WOULD HAVE PREVENTED IT: {d['prior']}"] + if n_rollouts and n: + head.insert(1, f'OBSERVED: this failure appeared in {n}/{n_rollouts} attempts.') + return '\n'.join(head) + lines = [f'The attempt failed in {len(diags)} distinct ways across {n_rollouts} attempts. ' + f'Address ALL of them — fixing only the first will still fail the tests.'] + for i, (d, n) in enumerate(diags, 1): + freq = f' (seen {n}/{n_rollouts})' if n_rollouts else '' + lines.append( + f"\nFAILURE {i}: {d['class']} — {R._CLASS_SHORT[d['class']]}{freq}" + f"\n WHAT WENT WRONG: {d['reason']}" + f"\n PRIOR THAT WOULD HAVE PREVENTED IT: {d['prior']}") + return '\n'.join(lines) + + +def multidiag_metrics(stats) -> Dict[str, float]: + """缓存/诊断计数 -> train_log 指标（与 e23_rubric.cache_metrics 同风格）。""" + tot = max(1, stats['hit'] + stats['ok'] + stats['hit_dropped'] + + stats['dropped_unaddressable'] + stats['api_fail']) + return {'rubric/hit_rate': (stats['hit'] + stats['hit_dropped']) / tot, + 'rubric/ok': float(stats['ok']), + 'rubric/dropped_unaddressable': float(stats['dropped_unaddressable']), + 'rubric/api_fail': float(stats['api_fail'])} diff --git a/cookbook/human_e18/e18_prompts.py b/cookbook/human_e18/e18_prompts.py new file mode 100644 index 000000000..28a6e0c12 --- /dev/null +++ b/cookbook/human_e18/e18_prompts.py @@ -0,0 +1,264 @@ +"""E18 的全部 prompt 文本与拼装函数：executor / 教师 judge / skill-gen 三处。 + +与 cookbook/human/e23_prompts.py 同源（同一套 BigCodeBench 交付要求与失败分类表），差别只在 +skill-gen 的系统提示：E18 是**拒绝采样 SFT**，采集时要 think 模式、训练时用 nothink 布局， +所以这里额外提供 query-only 的训练轨迹拼装（train_prompt）。 +""" +# flake8: noqa: E501 +# prompt 正文按「一段一行」书写，折行会改变真正发给模型的文本，故整文件豁免行长检查。 +from typing import Any, Dict + +# =========================================================================== +# executor +# =========================================================================== +# BigCodeBench 官方 instruct 模式的硬性交付要求（与 e23 逐字相同，保证跨实验可比）。 +EXEC_SYSTEM = """\ +You are an expert Python engineer. You will be given a task description that ends with the exact \ +import lines and function signature your solution must start with. + +Deliver exactly one fenced Python code block and nothing else after it: +- Reproduce the given imports and the given function signature verbatim, including parameter \ +names, order and default values. +- Add any further imports you need inside the same block; the block must run standalone. +- Return exactly the object type the task says to output. If it says the function should output \ +a tuple, return a tuple in that order; if it names a matplotlib Axes, return the Axes object \ +itself, not the Figure and not None. +- Implement the described behaviour for the general case, including the empty / single-element / \ +missing-column edge cases and any exception the description says to raise. +- Do not call the function, do not print demonstrations, do not add tests, do not use \ +`if __name__ == '__main__'`, and do not read from stdin. +- Do not include explanations outside the code block.""" + +# E18 的 executor 只看抽出来的 块，**不看** actor 的 ：本臂的产物是要写进 +# SFT 数据集、将来 query-only 部署的 skill 文本，采集期就必须按「部署时 executor 能看到什么」 +# 来判分，否则筛出来的胜者依赖一段部署时不存在的思考过程。（E23 是相反的口径，故意保留差异。） +_WRAPPER_SKILL_ONLY = ( + 'Hint:\n{hint}\n') + + +def direct_prompt(problem: str) -> Dict[str, Any]: + return {'messages': [{'role': 'system', 'content': EXEC_SYSTEM}, + {'role': 'user', 'content': problem}]} + + +def skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: + """题面 + 抽出的 块。skill 为空 -> 干净 direct（等价裸解，不塞空指导）。""" + skill = (skill or '').strip() + if not skill: + return direct_prompt(problem) + return {'messages': [{'role': 'system', 'content': EXEC_SYSTEM}, + {'role': 'user', + 'content': problem + '\n\n' + _WRAPPER_SKILL_ONLY.format(hint=skill)}]} + + +# =========================================================================== +# skill-gen（采集用：rubric 条件化、think 模式） +# =========================================================================== +# ⭐ 全英文、且**不给回复格式模板**。 +# 旧版是中文，并在类型 3/4 里给了带 `...` 占位符的回复示范（“根据你曾经犯过的错误……”、 +# “该问题属于...问题，因此可以拆解为...步骤”）。实测后果：模型把骨架连 `...` 一起抄下来， +# 胜者中 79/148 含空模板、中文占比从 65% 升到 100%，唯一词数 376→123，candidate_pass_rate +# 从 0.856 塌到 0.701。语言还与选择器共谋：中文字符数天然更少，在旧的 LEN_BUDGET 口径下永远 +# 更贴近预算而胜出。故：只说**要写什么**，不给句式；措辞由模型自己生成。 +# executor 与 BCB 题面均为英文，skill 也必须是英文才能与之对齐。 +# ⭐ 类型列表的**描述粒度必须齐平**，否则列表本身就是偏置：旧版 3/4 带了展开要求 +# （“Explain what the weak point is and why…”、“the concrete steps needed…”），1/2/5 却只有光板 +# 一句 —— 实测胜者里 t4 占 77-100%、t3 占 62-80%，而 t2 恒为 0%、t1 不过 4%。 +# 同理不写“vary across attempts”：单次采样看不到其他 rollout，该句对单条生成无法执行； +# 多样性靠 N_SKILLS 个独立 rollout 的采样噪声，以及“只选一类 + 五类等价”的显式声明。 +# ⭐ 2026-08-06：改为 **narrative 文体**（移植 skill2lora 的 SKILL_GEN_SYSTEM，见 +# cookbook/exp/skill2lora/train_skill_v2.py:843-861）。原版（下方 _SKILLGEN_SYSTEM_TYPED， +# 已注释停用）是「五类里挑一类」的列表式，实测问题： +# 1. 46.6% 的胜者是「用 A 不要用 B」的 API 纠正，靠的是 rubric 里的库行为知识； +# 2. 12.3% 直接把只存在于隐藏单测里的列名（'closing_price' 之类）写进 skill —— query 里 +# 根本没有，eval 时 query-only 无从得知，训练等于教模型凭空猜列名； +# 3. 5.9% 用「The critical error was...」这种事后复盘句式指代一次 eval 时不存在的失败。 +# narrative 的三条硬约束正好对上后两条：强制第一人称自持句式、**明令禁止**指向外部上下文 +# （skill2lora 的原注释写明这类措辞「会导致幻觉」）、不许代入本题具体数值。 +# ⭐ rubric 的用法：单独一段说明「有诊断时当作证据用」，与下方禁指代外部上下文那条并不矛盾 —— +# 两者分属不同层：**任务指令层**要求模型靠诊断定位软肋（WHERE），**输出层**要求把它转写成 +# executor 能直接执行的前瞻告诫（不得提及诊断本身）。这正是 skill2lora 的 REGEN_SYSTEM +# （train_skill_v2.py:1130-1147）的做法：步骤 2 要求「weaving together ... the pitfalls that +# actually tripped up the solving process」，而 Output requirement 同时禁止 "according to the given +# analysis/hints"。差异在于：REGEN 是「旧 skill + 诊断 -> 重写」的蒸馏场景，E18 是首次生成， +# 所以此处写成条件句（"If a grader's diagnosis ... is supplied"），无诊断时自动退化成纯预判。 +# 为何不能直接写「根据你之前犯过的错误……」：executor 看不到任何「之前」，而且训练目标是 +# query-only 的 —— 实测 5.9% 的胜者写成「The critical error was...」，eval 时模型无错可指只能编。 +# 所以保留“定位到具体步骤/API”这个内核，只把时态从「已发生」换成「容易在此处发生」。 +# ⚠️ 代价（skill2lora 已记录）：few-shot 例子占整条 prompt 的 65%，每次采样必然命中，会锁死 +# 文体与长度 —— 多样性下降是预期内的，换来的是「删掉本题依然成立」的可迁移性。 +# ⚠️ 本文件的 few-shot 例子必须是**代码域**的（原版是数学域的「末位数字/计数问题」，直接搬过来 +# 会把 executor 往数学叙述上带）；收尾纪律句同理换成 BCB 的交付要求，不用 boxed 那句。 +SKILLGEN_SYSTEM = """\ +You are a skill-writing expert. Your block will be fed to a SEPARATE downstream executor model that must solve the Python task on its own. The executor will NOT see your private reasoning — it only sees what is inside .... + +First, think privately: work out where the executor is most likely to get stuck, then step back and abstract WHAT MAKES THIS TYPE OF TASK GO WRONG into transferable guidance. + +Second, if a grader's diagnosis of a failed attempt is supplied, use it as your evidence for WHERE the weak point is: read what actually went wrong, then decide which part of the approach needs the executor's attention. Fold that insight into the narrative as guidance the executor can act on before it starts — name the step or the API where the trouble lives and say what to do there instead, e.g. "the place this tends to go wrong is when you ..., so at that point you should ...". Do not report the diagnosis; convert it into advice. + +Then write the block following these rules: +- Give general, transferable techniques for this TYPE of task: the library behaviour it relies on, the recommended approach, and the common pitfalls to avoid — plus a brief reason for each piece of advice so the executor understands why. +- Write it as one coherent analysis narrative (not a bullet list): first name what the task is essentially asking, then walk through how to approach it, blending the API contracts, steps, pitfalls and reasons into a single connected story. +- Write your judgements directly in the first person (e.g. "I think this step tends to ...", "A common mistake is ..., so you need to ..."), and phrase every issue as a self-contained, general technique. +- CRITICAL: Do NOT use phrasings that point to external context, such as "according to the given diagnosis", "the failed attempt", or "the previous error". The executor cannot see that context, and such phrasings will cause hallucination. State the pitfall as something that tends to happen at a particular step, not as something that already happened. +- CRITICAL: Do NOT name a column, key, or literal value that the task description does not itself state. If the task never names its columns, say how to discover them from the input instead of guessing names. +- Name the concrete API, argument, or keyword involved whenever the task description supports it. +- Keep it concise: aim for roughly one focused paragraph. + +Put ONLY the methodology inside . + +Example: + +This task is essentially asking you to reshape tabular input and hand back a plot object, so the delivery contract matters here as much as the computation; I would pin down exactly what type the function must return before writing any logic, because returning a Figure where an Axes was requested fails even when every number is right. The first place this tends to go wrong is the input itself: it arrives as a plain container, and I find the single most common break in this type of task is assuming it is already a DataFrame — dictionaries and lists of tuples carry none of the frame methods, so reaching for column-based access on them raises immediately, and at that point you should check what the object actually is and build the frame from it explicitly. The next place to slow down is naming: let the task description dictate the column names and read them off the signature or the docstring rather than inventing plausible-sounding ones, and when the description never states them, derive them from the input's own keys instead of hard-coding a guess, because a name that merely sounds right will pass your own reading and still miss. A common mistake is treating an empty or single-element input as impossible, so decide up front whether it should yield an empty result or raise, and write that branch before the main path. Finally, when plotting, create the Axes explicitly and return that same object, since helper calls that draw on the current figure make it easy to hand back something you never configured. Overall I summarise this type of task as "fix the return contract, verify the input's real type, take names from the description, then handle the empty case before the happy path", because that is where the failures concentrate. + +""" + +# =========================================================================== +# 【已停用】原「五类挑一类」列表式 skill-gen 提示（2026-08-06 换成上方 narrative） +# =========================================================================== +# 保留全文仅为记录历史口径与可回退：把下面的字符串改名回 SKILLGEN_SYSTEM 即可复原。 +# 停用原因见上方 narrative 块的注释（隐藏契约泄漏 12.3% / 复盘句式 5.9%）。 +# 注意它自身也修过两轮：类型描述粒度齐平（t1-t5 各 9-15 词）、以及那条前瞻视角规则 —— +# 这两笔修改都已被 narrative 的硬约束覆盖，回退时才需要重新评估。 +_SKILLGEN_SYSTEM_TYPED = """\ +You are a skill-writing expert. Your job is to write an advisory note that makes a downstream executor model solve the given Python task more accurately. + +First decide where the executor is most likely to get stuck, then write the advice you believe helps most. Pick the ONE kind below that fits this task best. All five are equally worth choosing, and a single sharp sentence often beats a long note: +1. A plain instruction about how to approach the work, such as what to be careful about. +2. A calibration cue about how much to deliberate, or about trusting its own judgement. +3. A generalized lesson drawn from the grader's diagnosis of a previous failed attempt, if one is supplied. Explain what the weak point is and why the lesson prevents it. +4. A decomposition of the task into the concrete steps needed to solve it. +5. Any other kind of skill you judge useful, including an angle you would not normally try. + +Rules: +- Do NOT solve the task and do NOT write code. You only write advice. +- Name the concrete API, argument, key, or value involved whenever you can. +- Be direct and specific. No filler, no restating the task, no placeholder text. +- Write forward-looking advice to someone who has not attempted the task yet. Do not refer to an error, mistake, or attempt as something that already happened. +- Choose your own wording and structure; there is no fixed format to follow. + +Wrap your skills in ... . +""" + +SKILLGEN_USER = """\ +TASK +{problem} + +GRADER'S DIAGNOSIS OF THE FAILED ATTEMPT +{rubric} + +Write the advisory note now, wrapped in .""" + +# ⭐ 带失败代码的变体。与 SKILLGEN_USER 的差别只有多出的 FAILED ATTEMPT 段， +# 段序是「题面 -> 失败代码 -> 诊断」：诊断紧贴写作指令，因为它才是要被消化的主结论； +# 把代码放中间让模型先看到证据再看结论，而不是反过来。 +SKILLGEN_USER_TRAJ = """\ +TASK +{problem} + +CODE FROM A FAILED ATTEMPT (for your analysis only — the executor will never see it) +{trajectory} + +GRADER'S DIAGNOSIS OF THE FAILED ATTEMPT +{rubric} + +Write the advisory note now, wrapped in .""" + + +def format_trajectory(code: str, error: str = '', kind: str = '', + max_chars: int = 4000) -> str: + """把一条失败 rollout 整理成给 skill-gen 看的文本块。 + + ⭐ 头尾各留一半而不是直接截前 max_chars：Python 失败代码的关键信息经常在**末尾** + （未闭合的分支、漏掉的 return、被截断的行），只留开头会把根因裁掉。 + + ⚠️ error 只取前 400 字符：pytest 的 longrepr 能有几千字符且大量重复的堆栈帧， + 全带上会把预算吃光，而判别失败模式只需要头部的异常类型与消息。 + """ + code = (code or '').strip() + if not code: + return '(the attempt produced no extractable code)' + if len(code) > max_chars: + half = max_chars // 2 + code = (code[:half] + '\n\n... [%d characters omitted] ...\n\n' % (len(code) - max_chars) + + code[-half:]) + out = ['```python', code, '```'] + if error: + out.append('OBSERVED ERROR: ' + ' '.join(str(error).split())[:400]) + if kind: + out.append('FAILURE CATEGORY: %s' % kind) + return '\n'.join(out) + + +# ⭐ 训推一致：这份同时做**训练 prompt** 与 **eval prompt**，与 SKILLGEN_SYSTEM 一起改成英文； +# 两边语言不一致会让模型在采集与部署时面对不同分布。 +SKILLGEN_SYSTEM_EVAL = """\ +You are a skill-writing expert. Your job is to write an advisory note that makes a downstream executor model solve the given Python task more accurately. + +First decide where the executor is most likely to get stuck, then write the advice you believe helps most. + +1. You have been trained on many kinds of skills, from a one-line caution to a full step decomposition. +2. Your memory already holds what works best for different kinds of problems. +3. Analyse the task and choose the skills you judge most useful. A single sharp sentence often beats a long note. + +Rules: +- Do NOT solve the task and do NOT write code. You only write advice. +- Name the concrete API, argument, key, or value involved whenever you can. +- Be direct and specific. No filler, no restating the task, no placeholder text. +- Write forward-looking advice to someone who has not attempted the task yet. Do not refer to an error, mistake, or attempt as something that already happened. + +Wrap your skills in ... . +""" + +SKILLGEN_USER_EVAL = """\ +TASK +{problem} + +Write the advisory note now, wrapped in .""" + + +def skillgen_prompt(problem: str, rubric: str, eval: bool, + trajectory: str = '') -> Dict[str, Any]: + """trajectory 非空时切到带失败代码的 user 模板（由 KOD_USE_TRAJ 控制，默认关）。 + + ⚠️ 只换 user 模板、**不换 system**：SKILLGEN_SYSTEM 里那条 + "Do NOT use phrasings that point to external context ... 'the failed attempt'" + 的禁令对带轨迹的情形更重要（模型看到真实代码后更容易写成事后复盘）， + 换掉 system 会同时丢掉这条约束。 + """ + if not eval: + if not rubric: + rubric = ('No diagnosis is available for this task. Consider the other kinds of ' + 'skill instead.') + user = (SKILLGEN_USER_TRAJ.format(problem=problem, rubric=rubric, + trajectory=trajectory) + if trajectory else + SKILLGEN_USER.format(problem=problem, rubric=rubric)) + return {'messages': [{'role': 'system', 'content': SKILLGEN_SYSTEM}, + {'role': 'user', 'content': user}]} + else: + return {'messages': [{'role': 'system', 'content': SKILLGEN_SYSTEM_EVAL}, + {'role': 'user', 'content': SKILLGEN_USER_EVAL.format(problem=problem)}]} + + +# =========================================================================== +# 已废弃：训练/eval 统一走 skillgen_prompt(..., eval=True) +# =========================================================================== +# ⭐ 不要再用 TRAIN_SYSTEM / train_prompt。 +# 训推一致要求「训练 prompt 与 eval prompt 逐字相同」，而 eval 用的是 SKILLGEN_SYSTEM_EVAL； +# 再并行维护一份英文 TRAIN_SYSTEM 只会让两边默默分叉。保留它仅为记录历史口径。 +TRAIN_SYSTEM = """\ +You are a problem-solving coach for a Python engineer. Given a task, write a short advisory note that anticipates the most likely decisive mistake and prevents it. + +Requirements: +- Wrap the note in and tags. +- State what to do, in the imperative. Name the concrete API, argument, key, or value involved. +- Do NOT solve the task, do NOT write code, and do NOT state the expected output value. +- Keep it under 90 words.""" + + +def train_prompt(problem: str) -> Dict[str, Any]: + """已废弃。训练与 eval 统一用 `skillgen_prompt(problem, '', eval=True)`。""" + raise NotImplementedError( + 'train_prompt 已废弃：训练/eval 请用 skillgen_prompt(problem, \'\', eval=True)，' + '以保证两边的 system/user 逐字一致（训推一致）。') diff --git a/cookbook/human_e18/e18_rejection_sft.py b/cookbook/human_e18/e18_rejection_sft.py new file mode 100644 index 000000000..01229d832 --- /dev/null +++ b/cookbook/human_e18/e18_rejection_sft.py @@ -0,0 +1,775 @@ +#!/usr/bin/env python3 +import json +import os +import shutil +import sys +import time +from dataclasses import dataclass +from typing import Any, Dict, List, Tuple + +import torch +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.data_format import SamplingParams, pack_user_data +from twinkle.dataloader import DataLoader +from twinkle.model import TransformersModel +from twinkle.patch.no_split_modules import NoSplitModulesPatch +from twinkle.processor import InputProcessor +from twinkle.sampler import vLLMSampler +from twinkle.template import Template + +# 环境层与教师 judge 复用 human/ 下的 e23 模块（不拷贝，保证判分/诊断逐字同源）。 +_HERE = os.path.dirname(os.path.abspath(__file__)) +_HUMAN = os.path.abspath(os.path.join(_HERE, '..', 'human')) +if _HUMAN not in sys.path: + sys.path.insert(0, _HUMAN) + +from e23_bcb import clean_text, empty_roll, extract_skill, judge_seqs, load_records # noqa: E402 +from e23_rubric import build_checker, class_metrics # noqa: E402 +# 多轨迹诊断：每种失败模式各诊一次再合并。不用 e23_rubric.RubricCache 是因为它的缓存键 +# 只含 data_id，同一题调 N 次会全部命中第一次的结果 —— 「每条 rollout 各诊一次」在那个键下做不到。 +from e18_multidiag import MultiDiagCache, multidiag_metrics # noqa: E402 + +from e18_prompts import direct_prompt, skill_solve_prompt, skillgen_prompt # noqa: E402 +from e18_select import gain_stats, select_winner # noqa: E402 + +try: + import swanlab +except ImportError: + swanlab = None + +logger = get_logger() + +# ========== Configuration ========== +MODEL_ID = os.environ.get('MODEL_ID', 'Qwen/Qwen3-4B') +OUTPUT_DIR = os.environ.get('OUTPUT_DIR', os.path.join(_HERE, 'output.e18')) + +TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 2)) +SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) +# ⭐ base_sampler 拿 4 张（而不是和其他两组一样的 2）：它是唯一瓶颈。每 chunk 的序列数 +# 相差一个量级 —— skill_sampler 只跑 CHUNK_SIZE*N_SKILLS（512），而 base_sampler 要跑 +# 裸解 CHUNK_SIZE*EXEC_ROLLOUTS（512）+ 重解 CHUNK_SIZE*N_SKILLS*EXEC_ROLLOUTS（4096）= 4608， +# 且 EXEC_MAX_TOKENS（15000）远大于 SKILL_MAX_TOKENS（8192），token 预算相差约 16 倍。 +# 给 skill_sampler 加卡几乎无收益，加在这里才能缩短 wall clock。 +BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 4)) +# ⭐ 没有 ref 模型：SFT 是纯交叉熵，不需要 KL 参考。E23 的 REF_GPUS 那两张转给了 base_sampler， +# 而不是空着 —— 8 卡机器上「省卡」没有意义，只会让瓶颈环节白白排队。 +NUM_GPUS = TRAIN_GPUS + SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS +TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', 1)) +TRAIN_DP = TRAIN_GPUS // TRAIN_FSDP +GPU_MEM = float(os.environ.get('GPU_MEM', 0.8)) + +SEED = int(os.environ.get('SEED', 42)) +CHUNK_SIZE = int(os.environ.get('CHUNK_SIZE', 64)) # 每轮裸解多少题去筛错题 +N_SKILLS = int(os.environ.get('N_SKILLS', 8)) # 每题采多少 skill 候选（拒绝采样的池） +# 攒够多少条胜者才 SFT 一次。必须是 TRAIN_DP 的整数倍（dp 切分要求），否则末尾会被丢。 +ACCUMULATE = int(os.environ.get('ACCUMULATE', 16)) +MAX_UPDATES = int(os.environ.get('MAX_UPDATES', 200)) +EVAL_SIZE = int(os.environ.get('EVAL_SIZE', 100)) +EVAL_EVERY_UPDATES = int(os.environ.get('EVAL_EVERY_UPDATES', 10)) +SAVE_EVERY_UPDATES = int(os.environ.get('SAVE_EVERY_UPDATES', 50)) # 0 = 只在结束时存 + +# skill-gen 采集：think 开、T=1（要多样性才有拒绝采样的意义） +SKILL_MAX_TOKENS = int(os.environ.get('SKILL_MAX_TOKENS', 8192)) +SKILL_GEN_TEMPERATURE = float(os.environ.get('SKILL_GEN_TEMPERATURE', 1.0)) +SKILL_GEN_TOP_P = float(os.environ.get('SKILL_GEN_TOP_P', 1.0)) +SKILL_GEN_TOP_K = int(os.environ.get('SKILL_GEN_TOP_K', -1)) +EXEC_MAX_TOKENS = int(os.environ.get('EXEC_MAX_TOKENS', 15000)) +MAX_MODEL_LEN = int(os.environ.get('MAX_MODEL_LEN', 16000)) + +# ⭐ pass@k 评分：executor 不再用 greedy 单次，而是每个 prompt 重解 EXEC_ROLLOUTS 次取通过率。 +# 为何必须：M=1/T=0 时 pass_rate 只有 0/1，易题上「加任何 skill 都对」，正例标签与 skill +# 质量无关 —— 等于往数据集里灌随机 skill。跑 8 次取连续 pass_rate 后，同一题的不同 skill 之间 +# 才有方差（如 8/8 vs 5/8），能真正排序。代价：executor GPU 时间乘 ~8 倍。 +# 温度必须 >0：T=0 下 8 次采样会逐字相同（judge_seqs 还会按 code 去重），pass_rate 退回 0/1。 +EXEC_ROLLOUTS = int(os.environ.get('EXEC_ROLLOUTS', 8)) +EXEC_TEMPERATURE = float(os.environ.get('EXEC_TEMPERATURE', 0.6)) +EXEC_TOP_P = float(os.environ.get('EXEC_TOP_P', 0.95)) + +LR = float(os.environ.get('LR', 1e-5)) # 恒定 lr，无 warmup 无 decay +TRAIN_MICRO_BATCH = int(os.environ.get('TRAIN_MICRO_BATCH', max(TRAIN_DP, ACCUMULATE // 2))) + +# 三道筛的两个阈值 +SKILL_CHAR_LIMIT = int(os.environ.get('SKILL_CHAR_LIMIT', 1500)) # 超过直接丢 + +# ⭐ 入池门槛：skill 至少要多做对 MIN_GAIN_ROLLOUTS 次（默认 2，即 +2/8 = +0.25）。 +# 为何不收 tie（8/8 -> 8/8）：那种样本只能证明 skill 无害，对「学会写有效 skill」没有任何 +# 监督信号 —— 易题上加任何 skill 都是 8/8，收它等于往数据集里灌随机文本。 +# 为何阈值是 2 而不是 1：8 次采样下 +1/8 在采样噪声量级内（二项分布标准误约 0.17）， +# 分不清是真提升还是波动；要求 +2/8 才能把噪声挤出去。用 rollout 数而不是写死 0.25， +# 是为了改 EXEC_ROLLOUTS 时该语义（「多做对几次」）保持不变。 +MIN_GAIN_ROLLOUTS = int(os.environ.get('MIN_GAIN_ROLLOUTS', 2)) +MIN_PASS_GAIN = MIN_GAIN_ROLLOUTS / max(1, EXEC_ROLLOUTS) + +SWAN_PROJ = os.environ.get('SWAN_PROJ', 'twinkle') +RUN_TAG = os.environ.get('RUN_TAG', '').strip() +RUN_ID = time.strftime('%m%d-%H%M%S') + + +@dataclass +class Runtime: + skill_model: Any + skill_sampler: Any + base_sampler: Any + ckpt: Any + checker: Any + rubric_cache: MultiDiagCache + + +# =========================================================================== +# 采样 / 小工具 +# =========================================================================== +def run_samples(sampler, prompts, num_samples, max_tokens, gen_dp, + temperature=None, top_p=None, top_k=None, logprobs=None): + """采样。prompts 少于 dp 时补齐再截回 —— Ray 的 dp 切分要求每个 rank 至少一条。""" + if not prompts: + return [] + import copy + params = SamplingParams( + max_tokens=max_tokens, + temperature=0.6 if temperature is None else temperature, + top_p=0.95 if top_p is None else top_p, + num_samples=num_samples, + **({} if top_k is None else {'top_k': top_k}), + **({} if logprobs is None else {'logprobs': logprobs})) + padded = prompts + if gen_dp > 1 and 0 < len(prompts) < gen_dp: + padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] + responses = sampler.sample(padded, params)[:len(prompts)] + return [list(r.sequences) if (r and r.sequences) else [] for r in responses] + + +def first_seq(seqs): + return seqs[0] if seqs else None + + +def seq_text(seq) -> str: + return clean_text(getattr(seq, 'decoded', '') or '') if seq is not None else '' + + +def _mean(xs) -> float: + xs = [float(x) for x in xs if x is not None] + return sum(xs) / len(xs) if xs else 0.0 + + +# =========================================================================== +# 采集：裸解 -> 诊断 -> skill-gen -> executor 重解 -> 三道筛 +# =========================================================================== +def _pass_rate(rolls) -> float: + return (sum(1.0 for x in rolls if x['correct']) / len(rolls)) if rolls else 0.0 + + +def bare_solve(rt: Runtime, records, rollouts: int = None) -> List[List[Dict[str, Any]]]: + """裸题重解 `rollouts` 次，每条记录返回一个 roll 列表（长度 = 实际采到的序列数）。 + + 返回**嵌套**列表而不是单个 roll：调用方靠 _pass_rate() 取连续值。 + 判分全部汇到一次 judge_seqs（它内部按 (task_id, code) 去重，相同代码只跑一次单测）。 + """ + M = max(1, rollouts if rollouts is not None else EXEC_ROLLOUTS) + out = run_samples(rt.base_sampler, [direct_prompt(r['problem']) for r in records], + M, EXEC_MAX_TOKENS, BASE_SAMPLER_GPUS, + temperature=(0.0 if M == 1 else EXEC_TEMPERATURE), + top_p=(None if M == 1 else EXEC_TOP_P)) + pairs, spans = [], [] + for r, seqs in zip(records, out): + seqs = list(seqs or []) + spans.append(len(seqs)) + pairs.extend((s, r['reference_answer']) for s in seqs) + judged = judge_seqs(pairs) + rolls, i = [], 0 + for n in spans: + rolls.append(judged[i:i + n] if n else [empty_roll()]) + i += n + return rolls + + +def collect_chunk(rt: Runtime, chunk, ci: int) -> Tuple[List[Dict[str, Any]], Dict[str, float]]: + """一个 chunk 的采集，返回 (胜者列表, 指标)。 + + ⭐ **全量 rollout**：chunk 里每一道题都要采 skill 候选，包括裸解已经做对的。 + 理由：部署时 skill 模型面对的是任意题，不知道 executor 会不会做对；只在错题上训会让 + 它只学会「救难题」、在简单题上也写一大堆纠错式提示。 + + 诊断（rubric）只对**做错的题**拉，且**每一种失败模式各诊一次后合并**（见 e18_multidiag）： + 8 次 rollout 往往挂在不同地方，只拿「第一条失败轨迹」去诊等于抛硬币选根因，而且旧缓存键 + 只含 data_id，会把那次随机结果**永久写进缓存**。现在按失败签名分桶、按频次排序， + 合并文本里明确要求「Address ALL of them」。 + 无诊断（做对的、或 API 失败的）的题**不丢弃**：skillgen_prompt 内部会填兜底文案。 + 难易判定改用**连续 pass_rate**：裸解跑 EXEC_ROLLOUTS 次，base_pass_rate < 1 即视为「有提升 + 空间」。这比单次贪心稳得多 —— 单次 T=0 的对/错在临界题上换个种子就翻转。 + """ + base_rolls = bare_solve(rt, chunk) + base_rates = [_pass_rate(rr) for rr in base_rolls] + base_acc = _mean(base_rates) + # 诊断目标：没能每次都对的题。整组 rolls 都传进去 —— 由 multidiag 自己按失败签名分桶， + # 每种模式诊一次（同签名只花一次 API），再合并成一份 rubric。 + wrong = [(r, rr) for r, rr, rate in zip(chunk, base_rolls, base_rates) if rate < 1.0] + + before = rt.rubric_cache.stats.copy() + diags = rt.rubric_cache.diagnose_many(rt.checker, wrong) + rmetrics = multidiag_metrics(rt.rubric_cache.stats - before) + diag_by_id = {id(r): d for (r, _rr), d in zip(wrong, diags) if d} + n_rubric_missing = len(wrong) - len(diag_by_id) + # 多根因覆盖率：合并文本里有几段 FAILURE。持续=1 说明多轨迹诊断没带来新信息。 + n_multi = sum(1 for d in diag_by_id.values() if d.count('FAILURE ') > 1) + + # todo 现在是**全量** chunk：(record, rubric_or_empty, base_pass_rate) + todo = [(r, diag_by_id.get(id(r), ''), rate) + for r, rate in zip(chunk, base_rates)] + + # skill-gen：think 模式、T=1、每题 N 个候选（这就是拒绝采样的候选池）。 + # eval=False -> 用 SKILLGEN_SYSTEM（thinking 采集口径，允许多类型 skill）。 + sg = run_samples(rt.skill_sampler, + [skillgen_prompt(r['problem'], d, eval=False) for r, d, _br in todo], + N_SKILLS, SKILL_MAX_TOKENS, SKILL_SAMPLER_GPUS, + temperature=SKILL_GEN_TEMPERATURE, top_p=SKILL_GEN_TOP_P, + top_k=SKILL_GEN_TOP_K) + per_task: List[List[Dict[str, Any]]] = [] + flat = [] + for (r, _d, _br), seqs in zip(todo, sg): + cands = [] + for s in seqs or []: + resp = seq_text(s) + block = extract_skill(resp) + c = {'skills': block, 'response': resp, 'parseable': bool(block), + 'with_pass': None, 'kept': False, + 'skillgen_stop': getattr(s, 'stop_reason', None)} + cands.append(c) + if block: + flat.append((r, c)) + per_task.append(cands) + + # executor 带 skill 重解：**每个 skill 跑 EXEC_ROLLOUTS 次**，取连续 pass_rate。 + # 这是本次改造的核心：只有连续值才能在「全部都能做对」的易题上区分 skill 好坏。 + if flat: + ws = run_samples(rt.base_sampler, + [skill_solve_prompt(r['problem'], c['skills']) for r, c in flat], + EXEC_ROLLOUTS, EXEC_MAX_TOKENS, BASE_SAMPLER_GPUS, + temperature=EXEC_TEMPERATURE, top_p=EXEC_TOP_P) + pairs, spans = [], [] + for (r, _c), seqs in zip(flat, ws): + seqs = list(seqs or []) + spans.append(len(seqs)) + pairs.extend((s, r['reference_answer']) for s in seqs) + judged = judge_seqs(pairs) + i = 0 + for (_r, c), n in zip(flat, spans): + rr = judged[i:i + n] + i += n + c['with_pass'] = _pass_rate(rr) + c['n_rollouts'] = n + c['roll_kind'] = (next((x['kind'] for x in rr if not x['correct']), 'pass') + if rr else 'empty') + + # 三道筛：第一道改成**pass_rate 严格不降 + 取最大**（详见 select_winner） + accepted, sims = [], [] + n_pass_cands = n_survivors = 0 + n_acc_hard = n_acc_easy = 0 + gtot = {'improved': 0, 'tied': 0, 'degraded': 0} + gains = [] + for (r, d, base_rate), cands in zip(todo, per_task): + passers = [c for c in cands + if c.get('parseable') and (c.get('with_pass') or 0) >= base_rate] + n_pass_cands += len(passers) + for k, v in gain_stats(cands, base_rate).items(): + gtot[k] += v + best = select_winner(cands, d, r['reference_answer'], + skill_char_limit=SKILL_CHAR_LIMIT, + base_pass_rate=base_rate, min_pass_gain=MIN_PASS_GAIN) + if best is None: + continue + n_survivors += 1 + sims.append(best['rubric_similarity']) + gains.append(best['pass_gain']) + if base_rate >= 1.0: + n_acc_easy += 1 + else: + n_acc_hard += 1 + accepted.append({ + 'problem': r['problem'], 'reference_answer': r['reference_answer'], + 'data_id': r.get('data_id', ''), 'skills': best['skills'], + 'response': f"\n{best['skills']}\n", + # 审计字段（只进数据集文件，不进训练轨迹） + # base_pass_rate / with_pass / pass_gain：连续口径下判断“skill 到底有没有用”的依据。 + 'base_pass_rate': base_rate, 'with_pass_rate': best['with_pass'], + 'pass_gain': best['pass_gain'], 'gain_kind': best['gain_kind'], + 'rubric': d, 'chunk': ci, 'run': RUN_ID, + 'rubric_similarity': best['rubric_similarity'], + 'skill_chars': len(best['skills']), + 'n_candidates_passed': len(passers)}) + + _cand_pass = (n_pass_cands / max(1, len(todo) * N_SKILLS)) if todo else 0.0 + metrics = { + 'train/baseline_accuracy': base_acc, + 'train/accept_rate': (len(accepted) / len(todo)) if todo else 0.0, + 'train/candidate_pass_rate': _cand_pass, + # ⭐ 采集侧 lift：**全部**候选的平均增量，无选择偏差。不要用 + # `gain/selected_pass_gain` 代替它：后者只统计已通过 +MIN_PASS_GAIN 门槛的胜者， + # 按定义恒为正，衡量的是「被选中那条有多好」而非「模型平均能写多好」。 + # 与 `eval/lift` 也不可直接比：这里的 prompt 带 rubric（教师诊断），eval 是 query-only， + # 所以 train/lift 包含了“教师诊断的价值”，两者的差距正是本实验要缩小的东西。 + 'train/lift': _cand_pass - base_acc, + 'train/selected_rubric_similarity': _mean(sims), + 'train/selected_skill_length_characters': _mean( + [float(s['skill_chars']) for s in accepted]), + 'signal/n_wrong': float(len(wrong)), + 'signal/n_rubric_missing': float(n_rubric_missing), + 'signal/n_multi_cause': float(n_multi), + 'signal/n_accepted': float(len(accepted)), + # 分层接受数：易题 = base_pass_rate 已经 1.0（跑 8 次全对）。 + 'signal/n_accepted_hard': float(n_acc_hard), + 'signal/n_accepted_easy': float(n_acc_easy), + # 天花板题数：base_pass_rate 高到拿不到 +MIN_PASS_GAIN（如 8/8），结构性无法入池。 + # 它与 n_accepted_easy 合看：前者持续很大就说明大量 GPU 花在了注定不入池的题上。 + 'signal/n_ceiling': float(sum(1 for _r, _d, br in todo + if br + MIN_PASS_GAIN > 1.0 + 1e-9)), + 'signal/min_pass_gain': MIN_PASS_GAIN, + # 候选级增量分解（相对裸解 pass_rate）。degraded 最重要：skill 把通过率拉低了， + # 它不会反映在 accept_rate 上，持续偏高就说明 skill-gen 在写有害提示。 + 'gain/improved_candidates': float(gtot['improved']), + 'gain/tied_candidates': float(gtot['tied']), + 'gain/degraded_candidates': float(gtot['degraded']), + 'gain/degrade_rate': (gtot['degraded'] / max(1, sum(gtot.values()))), + 'gain/improve_rate': (gtot['improved'] / max(1, sum(gtot.values()))), + # 胜者的平均 pass_rate 增量：这才是「入池样本到底有多有用」的直接度量。 + # 持续趋近 0 就说明池子里全是「写了也白写」的 skill。 + 'gain/selected_pass_gain': _mean(gains), + } | rmetrics | class_metrics([d for _r, d, _br in todo if d]) + return accepted, metrics + + +def dump_dataset(accepted) -> None: + """胜者落盘 append-only 的 SFT 数据集（含 rubric/相似度/pass 全审计字段）。 + + 这份文件是 E18 的主产物：它让「筛选器选了什么」可离线复算、可跨 run 复用（不必重跑 GPU + 就能换 SFT 超参再训一遍）。 + """ + if not accepted: + return + path = os.path.join(OUTPUT_DIR, 'e18_sft_dataset.jsonl') + with open(path, 'a', encoding='utf-8') as f: + for s in accepted: + f.write(json.dumps(s, ensure_ascii=False) + '\n') + + +# =========================================================================== +# 训练：纯 SFT（query-only 轨迹） +# =========================================================================== +def train_batch(rt: Runtime, samples) -> Tuple[int, Dict[str, float]]: + """在攒够的胜者上做 1 个 optimizer step。 + + ⭐ 训推一致的关键：训练 prompt 段用 **`skillgen_prompt(..., eval=True)`**，与 run_eval / + 部署时用的系统提示逐字相同（SKILLGEN_SYSTEM_EVAL，不带 rubric）。 + 采集时用的是带诊断的 thinking 口径（SKILLGEN_SYSTEM）—— 那只是为了**邀出**好 skill， + 不能拿去当训练分布：线上没有诊断可用，拿带诊断的 prompt 去训会学成「看着诊断改写」。 + + 响应段走 messages 编码而不是拼采样 token：胜者的 `` 文本是程序合成的（采集时原 + 响应带 、且包装不同），本来就没有对应的采样 token 序列。`key_rounds` 标出最后一轮 + （assistant）为唯一可训区，prompt 段全 -100。 + + loss 是纯 `CrossEntropyLoss`（不走 GRPO）：不传 advantages。拒绝采样的“选择”已经完成于三道筛 + （只有胜者入池），此处只需拟合目标文本。若要给样本加权，得在 loss 内部乘，而不是传 + advantages —— CrossEntropyLoss 不读这个参数，传了也是静默无效（所以此处根本不传）。 + """ + samples = [s for s in samples if (s.get('response') or '').strip()] + n = (len(samples) // TRAIN_DP) * TRAIN_DP # dp 切分要求整倍数 + if n == 0: + return 0, {} + samples = samples[:n] + trajs = [] + for s in samples: + msgs = skillgen_prompt(s['problem'], '', eval=True)['messages'] + trajs.append({'messages': msgs + [{'role': 'assistant', 'content': s['response']}], + 'user_data': pack_user_data({'key_rounds': [len(msgs)]})}) + micro = max(TRAIN_DP, min(TRAIN_MICRO_BATCH, n)) + for i in range(0, n, micro): + rt.skill_model.forward_backward(inputs=trajs[i:i + micro]) + rt.skill_model.clip_grad_and_step() + # ⭐ 这里**不**同步权重。同步只在 run_eval 前后发生（见 sync_for_eval），以保证 + # skill_sampler 在**采集**阶段永远是初始权重。 + # 为何：采集用 thinking 口径邀出候选，而训练目标是 nothink 的 纯文本。 + # 每步同步会把“别推理、直接吐 skills”回灌采集端，而下一轮采集又要求它 thinking + # —— 两个分布互相拉扯，训练每次都赢。实测（output.e18.en）：4 步内 skill 长度 + # 281->573、出现 `Motor virtue` / `spectral misfire` 这类退化文本，candidate_pass_rate + # 从 0.801 跌到 0.404、train/lift 转负（-0.055）。采集固定用初始权重能切断这个回路。 + # 代价：训练对采集零反馈，本质上退化成“离线数据生成 + 独立 SFT”。eval 仍用最新 + # 权重，所以 eval/lift 依旧反映 skill_model 的真实进步。 + + metrics = {'train/n_samples': float(n)} + # ⭐ 必须用 float() 尝试转换而不是 isinstance 判数值型：twinkle 的 LossMetric.calculate() + # 把 loss / grad_norm 格式化成**字符串**后才返回（`f'{avg_loss:.4f}'`），用 + # isinstance(val, (int, float)) 会把这两个最关键的优化指标静默丢弃 —— 且丢在写文件之前， + # 所以 train_log / swanlab / 日志里全部看不到，事后也无法找回。 + # 转不成的（'total time elapse'='12.3 minutes'、'speed'='1.2 iters/s'）才跳过。 + for k, val in (rt.skill_model.calculate_metric(is_training=True) or {}).items(): + if isinstance(val, bool): + continue + try: + fval = float(val) + except (TypeError, ValueError): + continue + if k.startswith('learning rate'): + if 'group 1' in k: + metrics['train/lr'] = fval + elif k.startswith('train/'): + metrics[k.replace(' ', '_')] = fval + else: + metrics[f'train/{k.replace(" ", "_")}'] = fval + return 1, metrics + + +# =========================================================================== +# eval：query-only（部署口径） +# =========================================================================== +def _sync_trained_to_sampler(rt: Runtime) -> None: + """把当前（已训练）权重临时推给 skill_sampler，供 eval 使用。 + + 只是「临时」：eval 一结束就由 _restore_base_weights 把初始权重灌回去。skill_model 侧 + 不落盘、不 load，训练权重与优化器状态全程不受影响。 + """ + rt.ckpt.sync_weights(merge_and_sync=True) + rt.skill_sampler.reset_prefix_cache() + + +def _restore_base_weights(rt: Runtime) -> None: + """把 skill_sampler 恢复到**初始**权重，供下一轮采集使用。 + + ⭐ 走 `skill_sampler.load_weights_from_path()`（不传参 = sampler 自己的 model_id，即原始 + 预训练权重），从磁盘直接流进 vLLM。关键是它**完全不碰 skill_model**： + 不 save、不 load，训练权重和 AdamW 动量都不受影响。 + + 对比曾经考虑过的「save 训练权重 -> skill_model.load(初始) -> sync -> load 回训练权重」： + 那条路要在训练模型上来回 load 两次，一旦中途失败（OOM / 磁盘满 / 进程被杀），训练端就 + 停在初始权重上却带着原来的优化器状态继续跑，训练成果被静默清零且日志上看不出来。 + 现在最坏情况只是采集端权重不对（下一次 eval 前的 sync 会覆盖掉），训练端不可能被破坏。 + + reset_prefix_cache 必须跟着走：prefix cache 里缓存的是旧权重算出的 KV，换完权重不清就会 + 拿旧 KV 拼新权重的输出。 + """ + rt.skill_sampler.load_weights_from_path() + rt.skill_sampler.reset_prefix_cache() + + +def run_eval(rt: Runtime, eval_records, base_cache: Dict[str, Dict[str, Any]], + updates: int = 0) -> Dict[str, float]: + """部署口径 eval：**nothink + SKILLGEN_SYSTEM_EVAL + 不给 rubric**，与训练分布逐字一致。 + + ⭐ skill_sampler 建立时是 enable_thinking=True（采集需要 thinking 多样性），而训练/部署是 + nothink。所以 eval 必须先把同一个引擎的**客户端模板**临时切成 nothink，跑完再切回； + 只换编码模板，引擎本身不动（与 skill_ablate/trainer.py 的做法同源）。 + 不切的后果：eval 会带着 thinking 布局生成，与训练的 nothink 布局错配 —— 测出来的数不是 + 部署时的真实能力。 + + baseline（executor 冻结）整个 run 只算一次，之后走 base_cache（缓存的是 **pass_rate**）。 + eval 也跑 EXEC_ROLLOUTS 次取通过率：与采集口径一致，否则 lift 不可比。 + + ⭐ 权重由调用方负责：进来之前已 sync 成最新训练权重，出去之后立刻恢复初始权重 + （见 main 里的 _sync_trained_to_sampler / _restore_base_weights）。本函数只管跑分。 + """ + todo = [r for r in eval_records if r['data_id'] not in base_cache] + for r, rr in zip(todo, bare_solve(rt, todo) if todo else []): + base_cache[r['data_id']] = _pass_rate(rr) + base_rates = [base_cache[r['data_id']] for r in eval_records] + base_acc = _mean(base_rates) + + # skill 模型：nothink + eval 系统提示、greedy 出一条 skill + rt.skill_sampler.set_template(Template, model_id=MODEL_ID, enable_thinking=False, + max_length=MAX_MODEL_LEN) + try: + sg = run_samples(rt.skill_sampler, + [skillgen_prompt(r['problem'], '', eval=True) for r in eval_records], + 1, SKILL_MAX_TOKENS, SKILL_SAMPLER_GPUS, temperature=0.0) + finally: + # 必须切回：下一个 chunk 的采集要 thinking。放 finally 里是为了 eval 中途报错也不会 + # 把采集口径永久卡在 nothink 上。 + rt.skill_sampler.set_template(Template, model_id=MODEL_ID, enable_thinking=True, + max_length=MAX_MODEL_LEN) + skills = [extract_skill(seq_text(first_seq(s))) for s in sg] + n_parsed = sum(1 for s in skills if s) + # executor 也跑 EXEC_ROLLOUTS 次，与 baseline / 采集同口径 + ws = run_samples(rt.base_sampler, + [skill_solve_prompt(r['problem'], s) for r, s in zip(eval_records, skills)], + EXEC_ROLLOUTS, EXEC_MAX_TOKENS, BASE_SAMPLER_GPUS, + temperature=EXEC_TEMPERATURE, top_p=EXEC_TOP_P) + pairs, spans = [], [] + for r, seqs in zip(eval_records, ws): + seqs = list(seqs or []) + spans.append(len(seqs)) + pairs.extend((s, r['reference_answer']) for s in seqs) + judged = judge_seqs(pairs) + rates, i = [], 0 + for n in spans: + rates.append(_pass_rate(judged[i:i + n])) + i += n + acc = _mean(rates) + # ⭐ skill 原文落盘：聚合指标看不出「写成了什么样」，lift 为负时必须能回到原文归因。 + dump_eval_skills(eval_records, skills, base_rates, rates, updates) + # 难题 = baseline 没能每次都对（pass_rate < 1）；rescue 改成平均 pass_rate 增量。 + hard = [(b, w) for b, w in zip(base_rates, rates) if b < 1.0] + return { + 'eval/accuracy': acc, + 'eval/baseline_accuracy': base_acc, + 'eval/lift': acc - base_acc, + 'eval/format_rate': n_parsed / max(1, len(eval_records)), + 'eval/hard_rescue_rate': (_mean([w - b for b, w in hard]) if hard else 0.0), + 'eval/n_hard': float(len(hard)), + 'eval/skill_length_characters': _mean([float(len(s)) for s in skills]), + } + + +def dump_eval_skills(eval_records, skills, base_rates, rates, updates) -> None: + """eval 产出的 skill 原文落盘（append-only）。 + + 为何必需：`run_eval` 以前只回传聚合指标，`skills` 是局部变量、用完即弃，而 + SAVE_EVERY_UPDATES 默认 50，所以 lift 为负时既看不到 skill 原文、也没有 ckpt 可以重跑 + —— 无法区分「内容写得差」和「注入方式不对」。这份产物让部署口径可离线归因。 + + 逐题存 base/with pass_rate，所以可以直接筛出被 skill 带坏的题（with < base）。 + """ + path = os.path.join(OUTPUT_DIR, 'eval_skills.jsonl') + with open(path, 'a', encoding='utf-8') as f: + for r, s, b, w in zip(eval_records, skills, base_rates, rates): + f.write(json.dumps({ + 'updates': updates, + 'data_id': r.get('data_id'), + 'task_id': r['reference_answer'].get('task_id'), + 'base_pass_rate': b, + 'with_pass_rate': w, + 'pass_gain': w - b, + 'skill_chars': len(s), + 'problem': r['problem'], + 'skills': s, + }, ensure_ascii=False) + '\n') + + +# =========================================================================== +# main +# =========================================================================== +def build_runtime(checker, rubric_cache) -> Runtime: + """三组卡：train / skill_sampler / base_sampler(executor)。SFT 不需要 ref 模型。""" + r0 = TRAIN_GPUS + r1 = r0 + SKILL_SAMPLER_GPUS + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, lazy_collect=False, groups=[ + DeviceGroup(name='train', ranks=list(range(0, r0)), device_type='GPU'), + DeviceGroup(name='skill_sampler', ranks=list(range(r0, r1)), device_type='GPU'), + DeviceGroup(name='base_sampler', ranks=list(range(r1, NUM_GPUS)), device_type='GPU')]) + + skill_model = TransformersModel( + model_id=MODEL_ID, remote_group='train', + device_mesh=DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_DP, + fsdp_size=TRAIN_FSDP), + ddp_config={'find_unused_parameters': False}, torch_dtype='float32') + skill_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) + # ⭐ 训练模板 enable_thinking=False：训练/部署都是 nothink，必须一致。 + # 采集（skill_sampler）才是 thinking，那只用来邀出候选，不是训练分布。 + # 反例：改成 True 会把 `\n\n\n` 那 4 个固定 token 也纳入可训区，而 run_eval + # 是 nothink 生成的 —— 两边可训/生成区对不上，就不再是训推一致。 + skill_model.set_template(Template, model_id=MODEL_ID, enable_thinking=False, + max_length=MAX_MODEL_LEN, truncation_strategy='delete') + skill_model.set_processor(InputProcessor, padding_free=False) + # ⭐ 纯 SFT：交叉熵，不走 GRPO 那一套。 + # CrossEntropyLoss 只看 inputs['labels']（即 key_rounds 圈出的 assistant 段），没有 ratio / + # clip / KL / advantage，也不需要 old_logps、ref 模型。拒绝采样的“选择”已经全部发生在 + # 三道筛里（胜者才入池），到了 loss 这一层就是普通的“拟合这条目标文本”，不应再有 RL 项。 + # reduction='mean'（默认）-> num_tokens=0，每个 micro 自己 token-mean，梯度按 micro 数归一。 + skill_model.set_loss('CrossEntropyLoss') + skill_model.set_optimizer('AdamW', lr=LR) + + def _sampler(group, world, enable_thinking): + s = vLLMSampler(model_id=MODEL_ID, + engine_args={'gpu_memory_utilization': GPU_MEM, + 'max_model_len': MAX_MODEL_LEN, + 'tensor_parallel_size': 1}, + device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), + remote_group=group) + s.set_template(Template, model_id=MODEL_ID, enable_thinking=enable_thinking, + max_length=MAX_MODEL_LEN) + return s + + # skill_sampler 开 think：采集要多样性（拒绝采样的前提）。eval 时同一引擎跑 query-only。 + skill_sampler = _sampler('skill_sampler', SKILL_SAMPLER_GPUS, enable_thinking=True) + # executor 关 think（与 E23 一致）：开 think 有大量 rollout 撞满预算连代码块都写不出来。 + base_sampler = _sampler('base_sampler', BASE_SAMPLER_GPUS, enable_thinking=True) + return Runtime( + skill_model=skill_model, skill_sampler=skill_sampler, base_sampler=base_sampler, + ckpt=CheckpointEngineManager(model=skill_model, sampler=skill_sampler), + checker=checker, rubric_cache=rubric_cache) + + +def swan_init(): + if swanlab is None or os.environ.get('SWANLAB_MODE') == 'disabled': + logger.info('[swanlab] 未启用，只写 train_log.jsonl') + return None + name = 'E18_rejection_sft_code' + (f'_{RUN_TAG}' if RUN_TAG else '') + f'_{RUN_ID}' + swanlab.init(project=SWAN_PROJ, experiment_name=name, config={ + 'model_id': MODEL_ID, 'seed': SEED, 'chunk_size': CHUNK_SIZE, 'n_skills': N_SKILLS, + 'accumulate': ACCUMULATE, 'max_updates': MAX_UPDATES, 'lr': LR, + 'loss': 'CrossEntropyLoss', + 'exec_rollouts': EXEC_ROLLOUTS, 'exec_temperature': EXEC_TEMPERATURE, + 'exec_top_p': EXEC_TOP_P, 'min_gain_rollouts': MIN_GAIN_ROLLOUTS, + 'min_pass_gain': MIN_PASS_GAIN, + 'skill_char_limit': SKILL_CHAR_LIMIT, 'skill_max_tokens': SKILL_MAX_TOKENS, + 'exec_max_tokens': EXEC_MAX_TOKENS, 'eval_size': EVAL_SIZE, + 'skill_gen_temperature': SKILL_GEN_TEMPERATURE, + 'run_tag': RUN_TAG, 'run_id': RUN_ID, 'output_dir': OUTPUT_DIR}) + logger.info(f'[swanlab] project={SWAN_PROJ} experiment={name}') + return swanlab + + +def swan_log(swan, row: Dict[str, Any], step: int) -> None: + if swan is None: + return + m = {k: float(v) for k, v in row.items() + if isinstance(v, (int, float)) and not isinstance(v, bool)} + try: + swan.log(m, step=step) + except Exception as e: + logger.warning(f'[swanlab] log 失败（已忽略）：{e}') + + +# 不参与污染判定的文件：跑一次要花大量 CPU/沙箱时间（~900 道题跑参考解答）， +# 且内容只依赖题池、与哪个 run 无关，所以要从旧目录携带到新目录。 +_CARRY_OVER = ('bcb_broken_tasks.json', ) + + +def archive_output_dir() -> None: + """启动时把已存在的 OUTPUT_DIR 整个 mv 走，保证本 run 写入空目录。 + + 为何必需：`e18_sft_dataset.jsonl` / `train_log.jsonl` 都是 `open(..., 'a')` 追写。 + 没有这一步时，重启一次就把新旧 run 的样本焊在同一个文件里，而且不报错。 + 实测后果（output.e18.en）：崩溃 run 的 602 条（硬缺陷 41.4%：ttr<0.45 有 170 条、 + 超长 125 条、重复 118 条）与新 run 的 274 条混在一起，旧数据占 69%。 + 本 run 的训练不受影响（训练吃的是内存里的 pool，这个文件只写不读），但它的存在 + 意义就是“不重跑 GPU 就能换 SFT 超参再训一遍”（见 dump_dataset）—— 那个场景下 + 污染会直接进训练，且无任何报错，只是模型更差。 + + 用 mv 而不是删：旧 run 的诊断数据（eval 曲线、崩溃样本）是可复用的分析素材， + 丢了就得重跑 GPU 才能拿回。 + """ + if not os.path.isdir(OUTPUT_DIR): + return + if not os.listdir(OUTPUT_DIR): # 空目录直接用，不制造无意义的归档 + return + # 归档名带旧目录的 mtime（而不是当前时间）：同一批旧数据无论何时重启都归到 + # 同一个名字上，看名字就知道里面是哪段时间的 run。 + stamp = time.strftime('%m%d-%H%M%S', time.localtime(os.path.getmtime(OUTPUT_DIR))) + dst = f'{OUTPUT_DIR}.bak-{stamp}' + n = 1 + while os.path.exists(dst): # 同一秒重启两次也不能覆盖已有归档 + dst = f'{OUTPUT_DIR}.bak-{stamp}.{n}' + n += 1 + shutil.move(OUTPUT_DIR, dst) + os.makedirs(OUTPUT_DIR, exist_ok=True) + carried = [] + for name in _CARRY_OVER: + src = os.path.join(dst, name) + if os.path.exists(src): + shutil.copy2(src, os.path.join(OUTPUT_DIR, name)) + carried.append(name) + logger.warning(f'[output] 已存在的 {OUTPUT_DIR} 已归档到 {dst}' + + (f'（携带缓存：{", ".join(carried)}）' if carried else '')) + + +def main(): + archive_output_dir() + os.makedirs(OUTPUT_DIR, exist_ok=True) + if ACCUMULATE % TRAIN_DP != 0: + raise ValueError(f'ACCUMULATE({ACCUMULATE}) 必须是 TRAIN_DP({TRAIN_DP}) 的整数倍，' + f'否则 dp 切分会丢掉尾部样本') + checker = build_checker() + if checker is None: + raise RuntimeError('没有教师 API（LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / ' + 'OPENAI_API_KEY 都没设）；rubric 是三道筛的参照系，无法降级运行') + train_dataset, eval_records = load_records(SEED, EVAL_SIZE, OUTPUT_DIR) + if len(train_dataset) < CHUNK_SIZE: + raise ValueError(f'训练池 {len(train_dataset)} 小于 CHUNK_SIZE {CHUNK_SIZE}') + logger.info(f'[data] train={len(train_dataset)} eval={len(eval_records)}') + + rt = build_runtime(checker, MultiDiagCache()) + swan = swan_init() + base_cache: Dict[str, Dict[str, Any]] = {} + pool: List[Dict[str, Any]] = [] # 攒够 ACCUMULATE 条就 SFT 一次 + updates, ci, si, epoch, last_eval = 0, 0, 0, 0, 0 + logger.info(f'E18 start: lr={LR} chunk={CHUNK_SIZE} n_skills={N_SKILLS} ' + f'accumulate={ACCUMULATE} max_updates={MAX_UPDATES} output={OUTPUT_DIR}') + + with open(os.path.join(OUTPUT_DIR, 'train_log.jsonl'), 'a', encoding='utf-8') as log_fh: + # ⭐ updates=0 的 baseline eval：没有这一点，`eval/*` 曲线的第一个数据点已经是训了 + # EVAL_EVERY_UPDATES 步之后的值，无法区分「训练带来的提升」和「初始就有的能力」。 + # 注意 base_cache 缓存的是 executor 裸解 pass_rate（executor 的基线），不是 skill 模型 + # 的基线 —— 两回事，不能互代。 + # 副作用：这次 eval 会把裸解结果写进 base_cache，所以后续 eval 能直接复用， + # 整体 GPU 成本并非净增一整次 eval。 + if eval_records: + row0 = {'step': 0, 'chunk': 0, 'updates': 0, 'epoch': 0, + 'signal/pool_size': 0.0} + row0.update(run_eval(rt, eval_records, base_cache, updates=0)) + row0['eval/updates_done'] = 0 + log_fh.write(json.dumps(row0, ensure_ascii=False) + '\n') + log_fh.flush() + swan_log(swan, row0, si) + logger.info('[baseline u0] ' + + ' '.join(f'{k}={v:.4g}' for k, v in row0.items() + if isinstance(v, float))) + si += 1 + + while updates < MAX_UPDATES: + loader = DataLoader(dataset=train_dataset, batch_size=CHUNK_SIZE, num_workers=0, + shuffle=True, drop_last=True, + generator=torch.Generator().manual_seed(SEED + epoch)) + for chunk in loader: + if updates >= MAX_UPDATES: + break + t0 = time.time() + accepted, metrics = collect_chunk(rt, chunk, ci) + dump_dataset(accepted) + pool.extend(accepted) + + n_upd, tmetrics = 0, {} + while len(pool) >= ACCUMULATE and updates < MAX_UPDATES: + batch, pool = pool[:ACCUMULATE], pool[ACCUMULATE:] + k, tm = train_batch(rt, batch) + n_upd += k + updates += k + tmetrics = tm or tmetrics + + row = {'step': si, 'chunk': ci, 'updates': updates, 'epoch': epoch, + 'seconds': round(time.time() - t0, 1), + 'signal/pool_size': float(len(pool)), + **metrics, **tmetrics} + if eval_records and (updates - last_eval >= EVAL_EVERY_UPDATES + or updates >= MAX_UPDATES) and updates > 0: + # ⭐ eval 前把训练权重临时推给 sampler，跑完立刻把初始权重灌回去。 + # train_batch 已不再逐步同步，所以采集阶段永远是初始模型； + # 只有这一段区间内 sampler 带的是训练权重。finally 保证 eval 中途 + # 报错也不会把退化权重永久留在采集端。 + _sync_trained_to_sampler(rt) + try: + row.update(run_eval(rt, eval_records, base_cache, + updates=updates)) + finally: + _restore_base_weights(rt) + row['eval/updates_done'] = updates + last_eval = updates + log_fh.write(json.dumps(row, ensure_ascii=False) + '\n') + log_fh.flush() + swan_log(swan, row, si) + logger.info(f'[s{si} c{ci} u{updates}/{MAX_UPDATES}] ' + + ' '.join(f'{k}={v:.4g}' for k, v in row.items() + if isinstance(v, float))) + if SAVE_EVERY_UPDATES and updates and updates % SAVE_EVERY_UPDATES == 0: + rt.skill_model.save(f'E18-u{updates}', output_dir=OUTPUT_DIR) + si += 1 + ci += 1 + epoch += 1 + + rt.skill_model.save('E18-final', output_dir=OUTPUT_DIR) + rt.rubric_cache.close() + if swan is not None: + swan.finish() + logger.info(f'done: updates={updates} chunks={ci} -> {OUTPUT_DIR}/E18-final') + + +if __name__ == '__main__': + main() diff --git a/cookbook/human_e18/e18_select.py b/cookbook/human_e18/e18_select.py new file mode 100644 index 000000000..b2d3a28d5 --- /dev/null +++ b/cookbook/human_e18/e18_select.py @@ -0,0 +1,142 @@ +"""E18 的拒绝筛：增量达阈 -> 不超长 -> pass_rate 最大（并列内长度/rubric 相似度拆平局）。 + +这是本臂**唯一的自变量**。E12 靠「rubric 重生成 + 2-in-8 验证」入池，没有排序；E18 在做对的 +候选里再排一次序，只留唯一胜者，并把打分全过程写进数据集文件供事后审计。 + +`_rubric_similarity` 与 skill_ablate/methods.py 逐字同源（含那条踩坑注释），搬过来是为了让 +本目录自包含、不再 import 那棵 2 万行的 methods.py。同源搬过来的泄漏门已删，理由见下。 +""" +import re +from collections import Counter +from typing import Any, Dict, List, Optional + +# --- 第二道筛已删：泄漏检测对 coding 任务恒为 False ------------------------------------ +# ⭐ 原 `answer_leaked` / `leak_blocks` 从 skill_ablate/methods.py 搬来，但那边是**数学**任务： +# `reference_answer` 是 `'42'` 这样的答案字符串，`str(ref) in skill` 的子串匹配有意义。 +# BCB 的 `reference_answer` 是个 ~2500 字符的 dict（task_id / entry_point / test / code_prompt / +# doc_struct / canonical_solution），该匹配要求 skill 逐字包含整个 dict 的 Python repr +# （含 `{'task_id': 'BigCodeBench/598', ...}`）—— 而 skill 上限 SKILL_CHAR_LIMIT=1500 字符， +# 物理上装不进。实测：只有把整个 dict 原样粘进去才判 True，截断 50 字符加前缀就已为 False， +# 触发概率恒为 0。 +# 真正的泄漏通道在 `test`（隐藏单测的断言期望值）与 `canonical_solution`，需要另写检测； +# 留着一个恒 False 的门只会给人「已经防住了」的假安全感，所以整体删除。 + + +# --- 第二道筛：skill 与 rubric 诊断的词频余弦 ---------------------------------------------- +# 定位：在「executor 已做对」的候选里挑与诊断内容对得上的那条，压掉两类假赢家 —— +# 与诊断无关的碰巧做对（含泄漏式速通的残余），以及与谁都不像的空泛套话。 +# 刻意用去停用词的词频余弦而不是 tfidf/语义模型：可迁移性判别器一节实测「仅 tfidf」 +# in-sample 0.983 / OOS 0.541 是纯过拟合；词频余弦纯 stdlib、确定性、可离线复算。 +_SIM_STOPWORDS = frozenset( + 'the a an and or of to in is are be for with that this it on as by from at not no was ' + 'were will would can could should may might do does did have has had you your we they ' + 'he she its if then than so but into over under out up down when where which what how ' + 'why all any each more most other some such only own same very'.split()) +_SIM_WORD_RE = re.compile(r"[a-z][a-z'\-]{2,}") + + +def rubric_similarity(skill: str, rubric: str) -> float: + """内容词词频余弦 ∈ [0,1]；任一侧无内容词返回 0。""" + ca = Counter(w for w in _SIM_WORD_RE.findall((skill or '').lower()) + if w not in _SIM_STOPWORDS) + cb = Counter(w for w in _SIM_WORD_RE.findall((rubric or '').lower()) + if w not in _SIM_STOPWORDS) + if not ca or not cb: + return 0.0 + dot = float(sum(v * cb[k] for k, v in ca.items() if k in cb)) + na = sum(v * v for v in ca.values()) ** 0.5 + nb = sum(v * v for v in cb.values()) ** 0.5 + return dot / (na * nb) if na and nb else 0.0 + + +# --- 两道筛合体 ----------------------------------------------------------------------------- +def select_winner(cands: List[Dict[str, Any]], rubric: str, reference: Any = None, *, + len_budget: int = 0, skill_char_limit: int, + base_pass_rate: float = 0.0, + min_pass_gain: float = 0.0) -> Optional[Dict[str, Any]]: + """按 **pass_rate 增量**选胜者；增量不达 min_pass_gain 就返回 None。 + + ⭐ 为何不能用旧的「做对就入池」（M=1/T=0）：那时 with_pass 只有 0/1，易题上「加任何 + skill 都对」，8 个候选全部 with_pass=1 —— 此时「挑哪条」完全由长度决定（易题 rubric 为空， + 相似度恒 0），等于往数据集里灌随机 skill。现在 executor 每个 skill 跑 M=8 次， + with_pass 是连续通过率，同一题的不同 skill 之间才有方差（如 8/8 vs 5/8）。 + + 筛选顺序： + a. **增量达阈**：`with_pass >= base_pass_rate + min_pass_gain`。 + min_pass_gain>0 时，仅仅「没弄坏」（tie，如 8/8 -> 8/8）**不够格**：那种样本对 + 「学会写有效 skill」没有监督信号。拉低通过率的更是直接淘汰。 + b. 超 skill_char_limit 过滤。（原本还有一道泄漏门，已删 —— 它在 coding 任务上恒为 False， + 详见文件头部注释。） + c. 取 **pass_rate 最大**的一档（允许并列）；并列内部按 **rubric 相似度**取高。 + + ⭐ c 的层次顺序很关键：pass_rate 是**客观效果**，rubric 相似度是**内容对齐**， + 长度只是**形式偏好**。长度已彻底移出择优路径（只在相似度也并列时做确定性拆平， + 取较短者）：旧版先按 `abs(len - len_budget)` 砍掉一半候选，而实测 66% 的题 8 个候选全部 + with_pass=1.0，于是长度成了事实上的唯一决策依据 —— 它把信息量大的长候选系统性淘汰。 + + ⭐ base_pass_rate 接近 1 时 a 几乎不可满足（天花板效应）：8/8 的题永远拿不到 +2/8， + 因此会被成建制排除在训练集外 —— 这是调用方想要的行为（只训真正有提升空间的题）， + 但意味着池子会偏向中等难度题，看 `signal/n_accepted_easy` 确认。 + + 胜者会被标上 `pass_gain`（= with_pass - base_pass_rate）与 `gain_kind`： + * 'improve'：pass_gain > 0；'tie'：== 0（仅当 min_pass_gain=0 时才可能返回）。 + + rubric 为空（易题无诊断）时相似度恒 0，并列内退化成取较短者 —— 此时已经是 + 「效果完全相同」的候选，拿什么拆平局都不影响效果，只是个确定性要求。 + + `reference` 已不再使用（泄漏门删除后唯一的消费方消失），`len_budget` 同样已废弃； + 两个形参保留只为不打断现有调用方的写法。 + """ + need = base_pass_rate + min_pass_gain + ok = [c for c in cands + if c.get('parseable') and (c.get('with_pass') or 0.0) >= need - 1e-9] + survivors = [c for c in ok if len(c['skills']) <= skill_char_limit] + if not survivors: + return None + # a/c：先按客观效果取最大档，再在并列内按 rubric 相似度拆平局。 + top = max((c.get('with_pass') or 0.0) for c in survivors) + tied = [c for c in survivors if (c.get('with_pass') or 0.0) >= top - 1e-9] + for c in tied: + c['rubric_similarity'] = rubric_similarity(c['skills'], rubric) + # 长度只作**最后的确定性拆平**（相似度也并列时），不再参与择优： + # 实测 66% 的题 8 个候选全部 with_pass=1.0，此时 `abs(len - LEN_BUDGET)` 事实上成了唯一 + # 决策依据，而它有系统性偏差 —— 中文表达同样内容的字符数天然更少（357 vs 705），永远 + # 更贴近预算，于是「离 400 最近」被翻译成了「选中文模板」，把信息量大的长英文候选全部 + # 淘汰。长度是形式偏好，不该越过效果与内容对齐，故彻底移出择优路径。 + best = max(tied, key=lambda c: (c['rubric_similarity'], -len(c['skills']))) + best['kept'] = True + best['pass_gain'] = round((best.get('with_pass') or 0.0) - base_pass_rate, 6) + best['gain_kind'] = 'improve' if best['pass_gain'] > 1e-9 else 'tie' + return best + + +def gain_stats(cands: List[Dict[str, Any]], base_pass_rate: float) -> Dict[str, int]: + """候选级增量计数（相对裸解 pass_rate），给 train_log 做监控。 + + `degraded` 是关键项：skill 把通过率拉低了。它完全不会反映在 accept_rate 上 + （那些候选只是默默被汰掉），持续偏高就是 skill-gen 在写有害提示的直接证据。 + """ + out = {'improved': 0, 'tied': 0, 'degraded': 0} + for c in cands: + if not c.get('parseable'): + continue + wp = c.get('with_pass') + if wp is None: + continue + if wp > base_pass_rate + 1e-9: + out['improved'] += 1 + elif wp < base_pass_rate - 1e-9: + out['degraded'] += 1 + else: + out['tied'] += 1 + return out + + +def filter_stats(passers: int, survivors: int) -> Dict[str, float]: + """第一道筛后的超长丢弃率，进 train_log。 + + 指标名保留 `leak_or_overlength_dropped_fraction` 不改：泄漏门删除前它的触发率恒为 0， + 所以新旧 run 的这个数值本来就完全可比（一直只在统计超长），改名反而会断掉曲线。 + """ + return {'train/leak_or_overlength_dropped_fraction': + ((passers - survivors) / passers) if passers else 0.0} diff --git a/cookbook/human_e18/e18_sft_kod.py b/cookbook/human_e18/e18_sft_kod.py new file mode 100644 index 000000000..6f0c4d1ea --- /dev/null +++ b/cookbook/human_e18/e18_sft_kod.py @@ -0,0 +1,543 @@ +# -*- coding: utf-8 -*- +"""E18-KOD 离线 SFT：把 `e18_collect_kod.py` 采到的胜者用 **nothink 口径**训一遍，并在 +首尾各跑一次 eval，验证训练是否真的带来提升。 + +与在线版 `e18_rejection_sft.py` 的关系：**保留训练 + eval，去掉采集**。rubric 诊断、 +拒绝采样、chunk 循环全部移除 —— 数据已经在 `e18_sft_dataset.jsonl` 里落好了。 + +卡位：4 训练 + 2 skill_sampler + 2 base_sampler(executor)。为何不是 8 卡全训练： +eval 要 skillmodel 生成 skill、executor 跑代码，两者都需要 vLLM 引擎常驻。 + +训推一致的三个不可动点（与在线版逐字对齐，改任何一处就不再是同一个实验）： +1. **prompt 段用 `skillgen_prompt(..., eval=True)`** —— 即 SKILLGEN_SYSTEM_EVAL、不带 rubric。 + 采集时用的是带诊断的 thinking 口径（SKILLGEN_SYSTEM），那只是为了「邀出」好 skill； + 线上没有诊断可用，拿带诊断的 prompt 去训会学成「看着诊断改写」。 +2. **模板 enable_thinking=False** —— 训练/部署都是 nothink，必须一致。改成 True 会把 + think 那几个固定 token 也纳入可训区，与部署时的生成区对不上。 +3. **`key_rounds=[len(msgs)]`** 标出最后一轮（assistant）为唯一可训区，prompt 段全 -100。 + +loss 是纯 `CrossEntropyLoss`（twinkle 没有 SFTLoss 这个类）：拒绝采样的「选择」已经 +发生在采集侧的三道筛里（只有胜者入池），到 loss 这层就是普通的拟合目标文本， +不该再有 ratio / clip / KL / advantage。 + +⭐ eval 口径与在线版的**唯一差异**：executor 只跑 1 次且 temperature=0（在线版是 8 次 +T=0.6 取通过率）。这是按需求指定的 —— 省 GPU、且 greedy 单次可复现。代价写在 +run_eval 的注释里：pass_rate 退化成 0/1 二值，单题不可比，只能看 100 题的均值。 + +产物（OUTPUT_DIR 下）： +* `sft_log.jsonl`：逐 step 的 loss / grad_norm / lr，用来看收敛。 +* `eval_log.jsonl`：首尾两次 eval 的聚合指标。 +* `eval_skills.jsonl`：eval 生成的 skill 原文 + 逐题 base/with，用来归因。 +* `KODSFT-final/`：权重（SAVE_EVERY_STEPS>0 时还有 `KODSFT-s/`）。 +""" +import json +import os +import random +import shutil +import sys +import time +from typing import Any, Dict, List, Tuple + +import torch +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.data_format import SamplingParams, pack_user_data +from twinkle.model import TransformersModel +from twinkle.patch.no_split_modules import NoSplitModulesPatch +from twinkle.processor import InputProcessor +from twinkle.sampler import vLLMSampler +from twinkle.template import Template + +_HERE = os.path.dirname(os.path.abspath(__file__)) +_COOKBOOK = os.path.abspath(os.path.join(_HERE, '..')) +for _p in (_HERE, os.path.join(_COOKBOOK, 'human')): + if _p not in sys.path: + sys.path.insert(0, _p) + +from e18_kodcode import (clean_text, empty_roll, extract_skill, # noqa: E402 + judge_seqs, load_records) +from e18_prompts import direct_prompt, skill_solve_prompt, skillgen_prompt # noqa: E402 + +logger = get_logger() + +# ========== 配置 ========== +MODEL_ID = os.environ.get('MODEL_ID', 'Qwen/Qwen3-4B') +DATA_PATH = os.environ.get( + 'DATA_PATH', os.path.join(_HERE, 'output.e18.kod', 'e18_sft_dataset.jsonl')) +OUTPUT_DIR = os.environ.get('OUTPUT_DIR', os.path.join(_HERE, 'output.e18.kod.sft')) + +# 8 卡分三组：4 训练 + 2 skillmodel 推理 + 2 executor。 +# ⭐ 为何 executor 只需 2 张（采集时是 6 张）：eval 只跑 100 题 × 1 次 = 100 个序列， +# 而采集是每 chunk 2304 个。序列数少两个量级，再加卡只会让训练侧变慢。 +TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 4)) +SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) +BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 2)) +NUM_GPUS = TRAIN_GPUS + SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS +TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', 1)) +TRAIN_DP = TRAIN_GPUS // TRAIN_FSDP +GPU_MEM = float(os.environ.get('GPU_MEM', 0.8)) + +SEED = int(os.environ.get('SEED', 42)) +EPOCHS = float(os.environ.get('EPOCHS', 3)) +# ⭐ 必须是 TRAIN_DP 的整倍数：dp 切分会把不足一轮的尾部丢掉。 +BATCH_SIZE = int(os.environ.get('BATCH_SIZE', 16)) +MICRO_BATCH = int(os.environ.get('MICRO_BATCH', 8)) +LR = float(os.environ.get('LR', 1e-5)) +MAX_MODEL_LEN = int(os.environ.get('MAX_MODEL_LEN', 16000)) + +# ---- eval 口径 ---- +EVAL_SIZE = int(os.environ.get('EVAL_SIZE', 100)) +SKILL_MAX_TOKENS = int(os.environ.get('SKILL_MAX_TOKENS', 8192)) +EXEC_MAX_TOKENS = int(os.environ.get('EXEC_MAX_TOKENS', 15000)) +# ⭐ executor 跑 1 次、temperature=0（按需求指定，与在线版的 8×T=0.6 不同）。 +# 后果：单题 pass_rate 只能是 0 或 1，所以**不要看单题差异**，只看 100 题均值； +# 也因此 hard_rescue 这类分层指标失去意义（不再计算）。greedy 的好处是可复现。 +EXEC_ROLLOUTS = int(os.environ.get('EXEC_ROLLOUTS', 1)) +EXEC_TEMPERATURE = float(os.environ.get('EXEC_TEMPERATURE', 0.0)) + +# 数据清洗门槛（见 load_samples 的注释，都是实测抽检出来的缺陷） +MIN_CHARS = int(os.environ.get('MIN_CHARS', 200)) +MAX_CHARS = int(os.environ.get('MAX_CHARS', 1500)) +DROP_CJK = os.environ.get('DROP_CJK', '1') == '1' + +SAVE_EVERY_STEPS = int(os.environ.get('SAVE_EVERY_STEPS', 0)) # 0 = 只在结束时存 +LOG_EVERY_STEPS = int(os.environ.get('LOG_EVERY_STEPS', 1)) +RUN_ID = time.strftime('%m%d-%H%M%S') + + +# =========================================================================== +# 数据 +# =========================================================================== +def load_samples() -> List[Dict[str, Any]]: + """读胜者并做格式清洗。返回 [{'problem', 'response'}]。 + + ⭐ 为何要在 SFT 侧再清一遍（采集侧已有 SKILL_CHAR_LIMIT）：抽检 1104 条实测出三类 + 残留缺陷，占比虽小但都会被模型逐字学走： + * CJK 混入 3/1104（'动态规划' 这种孤立中文词）—— 部署口径是纯英文，学走就成了双语输出； + * 残留 `` 标签 2/1104 —— response 外层已经由采集侧包了一层，内层再出现就是嵌套； + * 截断（无终止标点）7/1104 —— 学截断等于学「说半句就停」。 + 合计约 1.1%，宁可丢掉也不喂进去。 + + `response` 用采集时存的原字段（已是 `\\n...\\n` 包装），不重新拼： + 与在线版 `train_batch` 消费的是同一个字段，保持逐字一致。 + """ + if not os.path.exists(DATA_PATH): + raise FileNotFoundError(f'找不到数据集：{DATA_PATH}') + raw, bad_json = [], 0 + with open(DATA_PATH, 'r', encoding='utf-8') as f: + for line in f: + line = line.strip() + if not line: + continue + try: + raw.append(json.loads(line)) + except Exception: + # 采集进程被 kill 时最后一行可能是半行，跳过而不是让整个训练起不来。 + bad_json += 1 + drop = {'no_response': 0, 'too_short': 0, 'too_long': 0, 'cjk': 0, + 'nested_tag': 0, 'truncated': 0} + out = [] + for r in raw: + resp = (r.get('response') or '').strip() + skills = (r.get('skills') or '').strip() + if not resp or not skills: + drop['no_response'] += 1 + continue + if len(skills) < MIN_CHARS: + drop['too_short'] += 1 + continue + if len(skills) > MAX_CHARS: + drop['too_long'] += 1 + continue + if DROP_CJK and any('\u4e00' <= ch <= '\u9fff' for ch in skills): + drop['cjk'] += 1 + continue + if ' 1: + # 续跑（KOD_RESUME=1）后同一份文件里会有多个 run，训练是全量混合 —— 这里只提示， + # 不自动过滤：要不要分开训是实验设计问题，不该由脚本替你决定。 + logger.warning(f'[data] 数据里含 {len(runs)} 个 run：{runs}（全部混合训练）') + return out + + +def load_eval_records(train_data_ids: set): + """eval 集：从 KodCode 题池里挑 EVAL_SIZE 道**没有生成过 skill** 的题。 + + ⭐ 必须排掉采集跑过的题，而且排的是 **candidates 里的全量 id**、不是仅胜者： + 采集时一道题跑了 4 个候选但只有 ~22% 能入池，剩下的题虽然不在训练集里，却已经 + 被用来挑过 skill —— 拿它当 eval 会高估（数据选择偏差：那些题本身就是「skill 救不了」 + 或「本来就全对」的）。所以传进来的 train_data_ids 应该来自 `e18_candidates.jsonl`。 + + load_records 的 seed 与采集侧一致，所以题池顺序可复现；取前 EVAL_SIZE 条未采过的。 + """ + ds, _ = load_records(SEED, 0, OUTPUT_DIR) + out = [] + for r in ds.dataset: + if r['data_id'] in train_data_ids: + continue + out.append(r) + if len(out) >= EVAL_SIZE: + break + logger.info(f'[eval] 选了 {len(out)} 道未采集过的题（排除已跑 {len(train_data_ids)} 题）') + return out + + +def collected_data_ids() -> set: + """采集阶段跑过的全部 data_id（含未入池的），用来从 eval 集里排除。 + + 优先读 `e18_candidates.jsonl`（全量）；没有它才退回 sft_dataset（仅胜者，覆盖不全， + 会让 eval 集混进采集过的题）。 + """ + ids = set() + cand = os.path.join(os.path.dirname(DATA_PATH), 'e18_candidates.jsonl') + src = cand if os.path.exists(cand) else DATA_PATH + with open(src, 'r', encoding='utf-8') as f: + for line in f: + line = line.strip() + if not line: + continue + try: + ids.add(str(json.loads(line)['data_id'])) + except Exception: + continue + logger.info(f'[eval] 排除源 {os.path.basename(src)}：{len(ids)} 道已跑题') + return ids + + +# =========================================================================== +# 训练 +# =========================================================================== +def build_model(): + """三组卡：train(4) / skill_sampler(2) / base_sampler(2)。 + + 返回 (model, skill_sampler, base_sampler, ckpt)。训练组配置与在线版 build_runtime 一致。 + + ⭐ 两个 sampler 都直接建成 **enable_thinking=False**：本脚本没有采集阶段，不需要 + thinking 多样性，而 eval 就是部署口径（nothink）。因此不需要像在线版 run_eval 那样 + 每次 eval 前后临时切模板再切回 —— 少一个可能切错的状态。 + + ⭐ ckpt 只绑 skill_sampler：训练的是 skillmodel，executor 必须全程冻结， + 否则首尾两次 eval 的 baseline 不可比（分母都变了就无法归因给 skill）。 + """ + r0, r1 = TRAIN_GPUS, TRAIN_GPUS + SKILL_SAMPLER_GPUS + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, lazy_collect=False, groups=[ + DeviceGroup(name='train', ranks=list(range(0, r0)), device_type='GPU'), + DeviceGroup(name='skill_sampler', ranks=list(range(r0, r1)), device_type='GPU'), + DeviceGroup(name='base_sampler', ranks=list(range(r1, NUM_GPUS)), device_type='GPU')]) + model = TransformersModel( + model_id=MODEL_ID, remote_group='train', + device_mesh=DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_DP, + fsdp_size=TRAIN_FSDP), + ddp_config={'find_unused_parameters': False}, torch_dtype='float32') + model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) + # ⭐ enable_thinking=False：训练/部署都是 nothink，必须一致（见文件头注释第 2 点）。 + model.set_template(Template, model_id=MODEL_ID, enable_thinking=False, + max_length=MAX_MODEL_LEN, truncation_strategy='delete') + model.set_processor(InputProcessor, padding_free=False) + model.set_loss('CrossEntropyLoss') + model.set_optimizer('AdamW', lr=LR) + + def _mk_sampler(group: str, world: int): + s = vLLMSampler( + model_id=MODEL_ID, remote_group=group, + device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), + engine_args={'gpu_memory_utilization': GPU_MEM, + 'max_model_len': MAX_MODEL_LEN, 'tensor_parallel_size': 1}) + s.set_template(Template, model_id=MODEL_ID, enable_thinking=False, + max_length=MAX_MODEL_LEN) + return s + + skill_sampler = _mk_sampler('skill_sampler', SKILL_SAMPLER_GPUS) + base_sampler = _mk_sampler('base_sampler', BASE_SAMPLER_GPUS) + ckpt = CheckpointEngineManager(model=model, sampler=skill_sampler) + return model, skill_sampler, base_sampler, ckpt + + +# =========================================================================== +# eval：部署口径（nothink + SKILLGEN_SYSTEM_EVAL + 不给 rubric） +# =========================================================================== +def run_samples(sampler, prompts, num_samples, max_tokens, gen_dp, + temperature=None, top_p=None): + """采样。与 e18_collect_kod.run_samples 逐字一致（去掉本脚本用不到的 top_k/logprobs）。 + + 三处踩过的坑： + 1. 字段名是 `num_samples` 不是 `n`（SamplingParams 没有 `n`，传了直接 TypeError）。 + 2. 走 `sampler.sample(prompts, params)`，不是 pack_user_data + generate_sequences。 + 3. dp 补齐不能省：条数 < dp 会直接报错（eval 只 100 题、dp=2 虽然安全， + 但 EVAL_SIZE 调小到 1 时就会触发）。 + """ + if not prompts: + return [] + import copy + params = SamplingParams( + max_tokens=max_tokens, + temperature=0.6 if temperature is None else temperature, + top_p=0.95 if top_p is None else top_p, + num_samples=num_samples) + padded = prompts + if gen_dp > 1 and 0 < len(prompts) < gen_dp: + padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] + responses = sampler.sample(padded, params)[:len(prompts)] + return [list(r.sequences) if (r and r.sequences) else [] for r in responses] + + +def first_seq(seqs): + return seqs[0] if seqs else None + + +def seq_text(seq) -> str: + return clean_text(getattr(seq, 'decoded', '') or '') if seq is not None else '' + + +def _mean(xs) -> float: + xs = [float(x) for x in xs if x is not None] + return sum(xs) / len(xs) if xs else 0.0 + + +def _pass_rate(rolls) -> float: + return (sum(1.0 for x in rolls if x['correct']) / len(rolls)) if rolls else 0.0 + + +def _judge_batch(records, prompts, sampler, gen_dp) -> List[float]: + """跑 executor + 判分，返回逐题 pass_rate。 + + spans 记每题实际拿到几条序列（可能为 0），不能直接按 EXEC_ROLLOUTS 切 judged： + 采样失败的题会返回空列表，按固定步长切会整体错位。 + """ + ws = run_samples(sampler, prompts, EXEC_ROLLOUTS, EXEC_MAX_TOKENS, gen_dp, + temperature=EXEC_TEMPERATURE) + pairs, spans = [], [] + for r, seqs in zip(records, ws): + seqs = list(seqs or []) + spans.append(len(seqs)) + pairs.extend((s, r['reference_answer']) for s in seqs) + judged = judge_seqs(pairs) if pairs else [] + rates, i = [], 0 + for n in spans: + rates.append(_pass_rate(judged[i:i + n]) if n else 0.0) + i += n + return rates + + +def run_eval(skill_sampler, base_sampler, eval_records, + base_cache: Dict[str, float], tag: str, step: int) -> Dict[str, float]: + """部署口径 eval：skillmodel 生成 1 个 skill -> executor 跑 1 次 -> 看提升。 + + 与训练分布逐字一致：nothink + SKILLGEN_SYSTEM_EVAL + 不给 rubric。 + + ⭐ baseline 整个 run 只算一次，之后走 base_cache：executor 权重全程冻结，同一道题的 + 裸解结果不会变（T=0 更是确定性的）。重算不仅浪费 GPU，还会因为 vLLM 的 + 非确定性引入假的 baseline 漂移，把本来该归因给 skill 的差异污染掉。 + + ⭐ 只看 `lift`（= accuracy - baseline）的**首尾差异**。因为 EXEC_ROLLOUTS=1 + T=0， + 单题 pass_rate 只能是 0/1，100 题的均值标准误差约 0.05 —— 所以 lift 变化小于 + 约 0.07 时不要当成真实效果（双样本差值的噪声更大）。这是 1 次 rollout 的固有代价。 + """ + t0 = time.time() + todo = [r for r in eval_records if r['data_id'] not in base_cache] + if todo: + rates = _judge_batch(todo, [direct_prompt(r['problem']) for r in todo], + base_sampler, BASE_SAMPLER_GPUS) + for r, rate in zip(todo, rates): + base_cache[r['data_id']] = rate + base_rates = [base_cache[r['data_id']] for r in eval_records] + + sg = run_samples(skill_sampler, + [skillgen_prompt(r['problem'], '', eval=True) for r in eval_records], + 1, SKILL_MAX_TOKENS, SKILL_SAMPLER_GPUS, temperature=0.0) + skills = [extract_skill(seq_text(first_seq(s))) for s in sg] + n_parsed = sum(1 for s in skills if s) + with_rates = _judge_batch( + eval_records, + [skill_solve_prompt(r['problem'], s) for r, s in zip(eval_records, skills)], + base_sampler, BASE_SAMPLER_GPUS) + + acc, base_acc = _mean(with_rates), _mean(base_rates) + m = {'eval/accuracy': acc, 'eval/baseline_accuracy': base_acc, + 'eval/lift': acc - base_acc, + 'eval/format_rate': n_parsed / max(1, len(eval_records)), + 'eval/skill_length_characters': _mean([float(len(s)) for s in skills]), + 'eval/n_improved': float(sum(1 for b, w in zip(base_rates, with_rates) if w > b)), + 'eval/n_hurt': float(sum(1 for b, w in zip(base_rates, with_rates) if w < b)), + 'eval/seconds': round(time.time() - t0, 1)} + # ⭐ skill 原文落盘：聚合指标看不出「写成了什么样」，lift 为负时必须能回到原文归因。 + with open(os.path.join(OUTPUT_DIR, 'eval_skills.jsonl'), 'a', encoding='utf-8') as f: + for r, s, b, w in zip(eval_records, skills, base_rates, with_rates): + f.write(json.dumps({'tag': tag, 'step': step, 'run': RUN_ID, + 'data_id': r['data_id'], 'base': b, 'with': w, + 'skill_chars': len(s), 'skill': s}, + ensure_ascii=False) + '\n') + return m + + +# =========================================================================== +# 训练 +# =========================================================================== + + +def make_trajs(batch) -> List[Dict[str, Any]]: + """样本 -> twinkle 轨迹。与在线版 `train_batch` 的构造逐字相同。""" + trajs = [] + for s in batch: + msgs = skillgen_prompt(s['problem'], '', eval=True)['messages'] + trajs.append({'messages': msgs + [{'role': 'assistant', 'content': s['response']}], + 'user_data': pack_user_data({'key_rounds': [len(msgs)]})}) + return trajs + + +def step_metrics(model) -> Dict[str, float]: + """取本 step 的优化指标。 + + ⭐ 必须用 float() 试转而不是 isinstance 判数值型：twinkle 的 LossMetric.calculate() + 把 loss / grad_norm 格式化成**字符串**后才返回（`f'{avg_loss:.4f}'`），用 + isinstance(val, (int, float)) 会把这两个最关键的指标静默丢弃 —— 而这个脚本唯一的目的 + 就是看收敛，丢了 loss 就什么都看不到了。 + 转不成的（'total time elapse'='12.3 minutes'）才跳过。 + """ + out: Dict[str, float] = {} + for k, val in (model.calculate_metric(is_training=True) or {}).items(): + if isinstance(val, bool): + continue + try: + fval = float(val) + except (TypeError, ValueError): + continue + if k.startswith('learning rate'): + if 'group 1' in k: + out['lr'] = fval + else: + out[k.replace(' ', '_')] = fval + return out + + +def archive_output_dir() -> None: + """启动时把已存在的 OUTPUT_DIR 整个 mv 走（`sft_log.jsonl` 是追写的）。 + + 与 e18_collect_kod 同一套机制：不这么做，重跑一次就把两条 loss 曲线焊在一个文件里， + 而且不报错 —— 看收敛时会看到一条莫名其妙回弹的曲线。 + """ + if not os.path.isdir(OUTPUT_DIR) or not os.listdir(OUTPUT_DIR): + os.makedirs(OUTPUT_DIR, exist_ok=True) + return + stamp = time.strftime('%m%d-%H%M%S', time.localtime(os.path.getmtime(OUTPUT_DIR))) + dst = f'{OUTPUT_DIR}.bak-{stamp}' + i = 1 + while os.path.exists(dst): + dst = f'{OUTPUT_DIR}.bak-{stamp}-{i}' + i += 1 + shutil.move(OUTPUT_DIR, dst) + os.makedirs(OUTPUT_DIR, exist_ok=True) + logger.info(f'[init] 旧输出目录已归档 -> {dst}') + + +def log_eval(m: Dict[str, float], tag: str, step: int) -> None: + with open(os.path.join(OUTPUT_DIR, 'eval_log.jsonl'), 'a', encoding='utf-8') as f: + f.write(json.dumps({'tag': tag, 'step': step, 'run': RUN_ID, **m}, + ensure_ascii=False) + '\n') + logger.info(f'[eval:{tag}] ' + ' '.join(f'{k}={v:.4g}' for k, v in m.items())) + + +def main(): + t0 = time.time() + archive_output_dir() + samples = load_samples() + if len(samples) < BATCH_SIZE: + raise RuntimeError(f'可用样本 {len(samples)} 条 < BATCH_SIZE {BATCH_SIZE}') + if BATCH_SIZE % TRAIN_DP: + raise RuntimeError(f'BATCH_SIZE({BATCH_SIZE}) 必须是 TRAIN_DP({TRAIN_DP}) 的整倍数') + # eval 题池先选好（纯 CPU），再开 GPU：选错了就不用白等 vLLM 启动的几分钟。 + eval_records = load_eval_records(collected_data_ids()) + if not eval_records: + raise RuntimeError('eval 集为空（题池里的题已全部被采集过？）') + + model, skill_sampler, base_sampler, ckpt = build_model() + steps_per_epoch = len(samples) // BATCH_SIZE + total_steps = int(steps_per_epoch * EPOCHS) + logger.info(f'E18-KOD-SFT start: n={len(samples)} bs={BATCH_SIZE} micro={MICRO_BATCH} ' + f'lr={LR} epochs={EPOCHS} steps/epoch={steps_per_epoch} ' + f'total_steps={total_steps} eval_n={len(eval_records)} ' + f'exec_rollouts={EXEC_ROLLOUTS}@T{EXEC_TEMPERATURE} ' + f'gpus={TRAIN_GPUS}+{SKILL_SAMPLER_GPUS}+{BASE_SAMPLER_GPUS} out={OUTPUT_DIR}') + + # ---- 首次 eval（step 0，未训练的初始权重）---- + # ⭐ 不用 sync：skill_sampler 刚建立，拿的就是 MODEL_ID 的原始权重，与训练端同源。 + base_cache: Dict[str, float] = {} + m_before = run_eval(skill_sampler, base_sampler, eval_records, base_cache, 'before', 0) + log_eval(m_before, 'before', 0) + + log_path = os.path.join(OUTPUT_DIR, 'sft_log.jsonl') + rng = random.Random(SEED) + step = 0 + with open(log_path, 'a', encoding='utf-8') as log_fh: + epoch = 0 + while step < total_steps: + order = list(range(len(samples))) + rng.shuffle(order) # 每个 epoch 重洗，种子固定所以可复现 + for bi in range(steps_per_epoch): + if step >= total_steps: + break + batch = [samples[j] for j in order[bi * BATCH_SIZE:(bi + 1) * BATCH_SIZE]] + trajs = make_trajs(batch) + micro = max(TRAIN_DP, min(MICRO_BATCH, len(trajs))) + t_step = time.time() + for i in range(0, len(trajs), micro): + model.forward_backward(inputs=trajs[i:i + micro]) + model.clip_grad_and_step() + step += 1 + row = {'step': step, 'epoch': epoch, 'run': RUN_ID, + 'n_samples': len(batch), 'seconds': round(time.time() - t_step, 2)} + row.update(step_metrics(model)) + log_fh.write(json.dumps(row, ensure_ascii=False) + '\n') + log_fh.flush() + if step % LOG_EVERY_STEPS == 0: + logger.info('[s%d/%d ep%d] ' % (step, total_steps, epoch) + + ' '.join(f'{k}={v:.4g}' for k, v in row.items() + if isinstance(v, float))) + if SAVE_EVERY_STEPS and step % SAVE_EVERY_STEPS == 0: + # ⭐ API 是 `model.save(tag, output_dir=)`，不是 save_checkpoint(path) + # —— 与在线版 e18_rejection_sft.py:762 的用法一致。 + model.save(f'KODSFT-s{step}', output_dir=OUTPUT_DIR) + logger.info(f'[save] KODSFT-s{step}') + epoch += 1 + + model.save('KODSFT-final', output_dir=OUTPUT_DIR) + + # ---- 尾次 eval（训练后权重）---- + # ⭐ 必须先 sync_weights 把训好的权重推给 skill_sampler，否则这次 eval 跑的还是初始 + # 权重 —— 两次结果几乎相同，看起来像「训了没效果」，而真因是权重根本没上去。 + # merge_and_sync=True：全参数训练需要先在 dp 间 merge 再推。 + # reset_prefix_cache 必须跟着走：prefix cache 里是旧权重算出的 KV，不清就会拿旧 KV + # 拼新权重的输出，得到一个既不是训前也不是训后的嵌合态。 + ckpt.sync_weights(merge_and_sync=True) + skill_sampler.reset_prefix_cache() + # base_cache 沿用：executor 全程未动，baseline 不需重算（也不应重算，见 run_eval）。 + m_after = run_eval(skill_sampler, base_sampler, eval_records, base_cache, 'after', step) + log_eval(m_after, 'after', step) + + d_lift = m_after['eval/lift'] - m_before['eval/lift'] + logger.info('[result] lift %.4f -> %.4f (Δ %+.4f) accuracy %.4f -> %.4f baseline %.4f' + % (m_before['eval/lift'], m_after['eval/lift'], d_lift, + m_before['eval/accuracy'], m_after['eval/accuracy'], + m_after['eval/baseline_accuracy'])) + # ⭐ 噪声底提醒：EXEC_ROLLOUTS=1 + T=0 下单题 pass_rate 是 0/1，n=100 的均值标准误约 + # 0.05，首尾差值的噪声更大。不把这句写进日志，很容易把 ±0.05 的漂动当成结论。 + if abs(d_lift) < 0.07: + logger.warning(f'[result] Δlift {d_lift:+.4f} 在噪声量级内（n={len(eval_records)}、' + f'rollout=1、T=0 时约 ±0.07），不足以断定有/无效果。') + logger.info(f'[done] steps={step} 用时 {(time.time() - t0) / 60:.1f} 分钟 -> {OUTPUT_DIR}') + + +if __name__ == '__main__': + main() diff --git a/cookbook/human_e18/e19_logp_select.py b/cookbook/human_e18/e19_logp_select.py new file mode 100644 index 000000000..08d311af4 --- /dev/null +++ b/cookbook/human_e18/e19_logp_select.py @@ -0,0 +1,193 @@ +# -*- coding: utf-8 -*- +"""E19：能否不靠 8 次 executor rollout 就选出最好的 skill？（RLT-style logp/熵 打分） + +映射自 skill_quality_analysis.md 第 18 节的「用法2（只搬 reward）」。核心问题： +现在选 skill 要给每个候选跑 8 次 executor + 判分（8 候选 = 64 次解码），太贵。 +能不能用 **frozen executor 上的一次 teacher-forcing forward**（无解码）替代？ + +流程（与 e18 采集逐字同源，**不用 GT 造 prompt**）： + 1. 裸解 BARE_ROLLOUTS 次 -> 取**失败的** trajectory（错误代码 + 报错） + 2. 失败 traj + rubric 一起喂给 skillmodel -> 生成 N_SKILLS 个 narrative skill + 3. 每个 skill 让 executor rollout EXEC_ROLLOUTS 次 -> with_pass_rate（**这是 ground truth， + 只用于离线评估选择器，不参与任何打分特征**） + 4. 同时对每个 skill 算 cheap 特征（一次 forward，见下）-> 落盘 + 5. 离线比：cheap selector 的命中率 vs 8-rollout oracle + +⭐ 为什么 GT 不进 prompt：用户明确要求。RLT 的 teacher 是**开卷**的（输入含标准答案）， +因此它必须靠 r_KL 压泄漏，且 teacher 不能直接部署。本实验的 skillmodel 保持**闭卷** +（输入只有题面 + 自己的失败 traj + rubric），GT 只在两个地方出现： + (a) 判分（pytest 跑测试）—— 本来就在判分侧，不进生成器； + (b) r_SS 的打分目标 —— 只流经 reward 计算的 forward，不进生成器上下文。 +所以本实验**结构上无泄漏**，不需要 RLT 的 r_KL 来压——但仍然实现了 leak 监控项， +因为 rubric 里可能夹带答案片段（见 leak_frac）。 + +⭐ 打分目标选 `canonical_solution` 而不是「修对后的代码」：RLT 的 r_SS 是 +logp(标准答案 | 讲解, 题)。我们没有「修对后的代码」这种东西（那要先跑通才知道）， +所以用数据集自带的参考解。代价：参考解的写法风格与 executor 的自然写法不同， +logp 会偏低且带常数偏移 —— 但我们只在**同题的候选之间**比较排序，常数偏移会被抵消。 + +cheap 特征（全部来自一次 prompt_logprobs forward，零解码）： + * r_ss : mean logp(参考解 token | 题 + skill) <- RLT 主项，越高越好 + * r_ss_min : min-k 平均（最难的 10% token） <- RLT 的 α·min 项 + * ppl : exp(-r_ss)，困惑度 + * ent_* : 参考解位置上的预测熵（需要 topk） + * skill 自身的 logp/熵（生成时顺带拿到） +""" +import json +import os +import sys +import time +from collections import defaultdict +from typing import Any, Dict, List, Optional, Tuple + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import SamplingParams +from twinkle.sampler import vLLMSampler +from twinkle.template import Template + +_HERE = os.path.dirname(os.path.abspath(__file__)) +for _p in (_HERE, os.path.abspath(os.path.join(_HERE, '..', 'human'))): + if _p not in sys.path: + sys.path.insert(0, _p) + +from e18_kodcode import (clean_text, extract_code, judge_seqs, # noqa: E402 + load_records) +from e18_prompts import direct_prompt, skill_solve_prompt # noqa: E402 +from e18_multidiag import MultiDiagCache # noqa: E402 +from e23_rubric import build_checker # noqa: E402 + +logger = get_logger() + +MODEL_ID = os.environ.get('MODEL_ID', 'Qwen/Qwen3-4B') +OUTPUT_DIR = os.environ.get('OUTPUT_DIR', os.path.join(_HERE, 'output.e19.logp')) +SEED = int(os.environ.get('SEED', 42)) + +SKILL_GPUS = int(os.environ.get('SKILL_GPUS', 4)) +EXEC_GPUS = int(os.environ.get('EXEC_GPUS', 4)) +NUM_GPUS = SKILL_GPUS + EXEC_GPUS +GPU_MEM = float(os.environ.get('GPU_MEM', 0.85)) +MAX_MODEL_LEN = int(os.environ.get('MAX_MODEL_LEN', 24000)) + +N_TASKS = int(os.environ.get('N_TASKS', 24)) # 「几组错误的」——先小规模看效果 +N_SKILLS = int(os.environ.get('N_SKILLS', 8)) # 每题 8 个 skill 候选 +BARE_ROLLOUTS = int(os.environ.get('BARE_ROLLOUTS', 4)) +EXEC_ROLLOUTS = int(os.environ.get('EXEC_ROLLOUTS', 8)) # ground truth 用 +SKILL_MAX_TOKENS = int(os.environ.get('SKILL_MAX_TOKENS', 8192)) +EXEC_MAX_TOKENS = int(os.environ.get('EXEC_MAX_TOKENS', 15000)) +EXEC_TEMPERATURE = float(os.environ.get('EXEC_TEMPERATURE', 0.6)) +SKILL_TEMPERATURE = float(os.environ.get('SKILL_TEMPERATURE', 1.0)) +TOPK = int(os.environ.get('TOPK', 20)) # 算熵用的 top-k +RUN_ID = time.strftime('%m%d-%H%M%S') + + +# =========================================================================== +# prompt：失败 traj + rubric -> narrative skill（闭卷，无 GT） +# =========================================================================== +SKILLGEN_SYSTEM = ( + 'You are helping a Python programmer who is about to attempt a coding task. ' + 'You have seen this programmer fail this exact task before, and you have a ' + 'diagnosis of what went wrong.\n\n' + 'Write a short piece of guidance (a "skill") that would have prevented that ' + 'failure. Requirements:\n' + '- Write flowing prose, not bullet points or headings.\n' + '- Name the concrete API, argument, keyword, or edge case involved.\n' + '- Refer to the past failure as something that already happened, and say what ' + 'to do instead.\n' + '- Do NOT include any code block, and do NOT write a full solution.\n' + '- Keep it under 200 words.\n' + 'Wrap your guidance in and tags.') + + +def skillgen_prompt(problem: str, failed_code: str, error: str, rubric: str) -> Dict[str, Any]: + """闭卷 skill 生成 prompt：题面 + 自己的失败代码 + 报错 + rubric 诊断。 + + ⭐ 这里**没有** canonical_solution / reference answer。这是与 RLT teacher 的关键区别： + RLT 开卷（输入含答案）所以必须用 r_KL 压泄漏；我们闭卷，泄漏在结构上不可能发生。 + """ + user = (f'Task the programmer was given:\n{problem}\n\n' + f'The code they wrote (it failed):\n```python\n{failed_code}\n```\n\n' + f'How it failed:\n{error}\n\n' + f'Diagnosis:\n{rubric}\n\n' + 'Write the guidance that would have prevented this failure.') + return {'messages': [{'role': 'system', 'content': SKILLGEN_SYSTEM}, + {'role': 'user', 'content': user}]} + + +# =========================================================================== +# 采样 / 判分 工具（与 e18 同源） +# =========================================================================== +def run_samples(sampler, prompts, num_samples, max_tokens, gen_dp, + temperature=None, top_p=None, logprobs=None): + if not prompts: + return [] + import copy + params = SamplingParams( + max_tokens=max_tokens, + temperature=0.6 if temperature is None else temperature, + top_p=0.95 if top_p is None else top_p, + num_samples=num_samples, + **({} if logprobs is None else {'logprobs': logprobs})) + padded = prompts + if gen_dp > 1 and 0 < len(prompts) < gen_dp: + padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] + resp = sampler.sample(padded, params)[:len(prompts)] + return [list(r.sequences) if (r and r.sequences) else [] for r in resp] + + +def seq_text(seq) -> str: + return clean_text(getattr(seq, 'decoded', '') or '') if seq is not None else '' + + +def _mean(xs) -> float: + xs = [float(x) for x in xs if x is not None] + return sum(xs) / len(xs) if xs else 0.0 + + +def extract_skill(text: str) -> str: + if '' not in text: + return '' + body = text.split('', 1)[1] + return body.split('', 1)[0].strip() if '' in body else '' + + +# =========================================================================== +# ⭐ 核心：teacher-forcing 打分（一次 forward，零解码） +# =========================================================================== +def score_teacher_forcing(exec_sampler, problem: str, skill: str, target_code: str, + gen_dp: int) -> Dict[str, float]: + """算 logp(target_code | 题 + skill) —— RLT 的 r_SS，用 prompt_logprobs 实现。 + + ⭐ 机制：把「题+skill」当 prompt、把 target_code **拼在 prompt 末尾**，然后 + `max_tokens=1` + `prompt_logprobs=TOPK` 采样。vLLM 会回传**每个 prompt token 的 + logprob**（vllm_engine.py:324-343 取的是实际 token 的 logprob），于是我们免费拿到 + 了 target_code 每个 token 在该上下文下的条件概率 —— 这就是 teacher forcing， + 且**不解码任何 token**，比 8 次 rollout 便宜两个量级。 + + ⭐ 必须用 continue_final_message 让模板把 target_code 编成 **assistant 内容**而不是 + 新一轮 user：编错角色会导致 logp 分布完全不同（模型在算「用户会说这段代码的概率」）。 + 这里靠传入 assistant 角色的消息 + 模板拼接实现。 + + ⭐ 只取 target 段的 logprob，必须切掉 prompt 段。切点靠**两次编码求长度差**确定： + 先编「题+skill」得 n_ctx，再编「题+skill+target」得 n_all，则 target 占 + [n_ctx, n_all)。不能用固定偏移或字符串查找 —— tokenizer 会在边界合并 token。 + """ + return {} + + +def _entropy_from_topk(topk: List[Optional[List[Tuple[int, float]]]]) -> List[float]: + """由 top-k logprob 估计每位置的熵。 + + ⚠️ 这是**截断熵**（只看 top-k），不是真熵：尾部质量被忽略，所以系统性偏低。 + 但我们只做同题候选间的排序比较，偏差方向一致，可用。k=TOPK=20 时通常覆盖 >90% 概率质量。 + """ + import math + out = [] + for lps in topk or []: + if not lps: + out.append(0.0) + continue + ps = [math.exp(lp) for _, lp in lps] + z = sum(ps) or 1.0 + out.append(-sum((p / z) * math.log(max(p / z, 1e-12)) for p in ps)) + return out diff --git a/cookbook/human_e18/e20_success_skill.py b/cookbook/human_e18/e20_success_skill.py new file mode 100644 index 000000000..cc035a17d --- /dev/null +++ b/cookbook/human_e18/e20_success_skill.py @@ -0,0 +1,311 @@ +# -*- coding: utf-8 -*- +"""E20：**成功** trajectory -> narrative skill（无 rubric），什么情况下该保留？ + +与 E18 的差别只有一处：题源从「裸解失败」换成「裸解**第一次就成功**」， +于是 skillmodel 拿到的是一条**成功的** trajectory，而且**没有 rubric** +（rubric 是失败诊断，成功的题没有可诊断的失败）。 + +⭐ 本实验的核心难点：这些题裸解已经通过，**pass_rate 没有提升空间**。 +所以 E18 那套「+0.25 增益」门槛在这里恒不成立，直接套用会得出「一条都不该留」 +的空洞结论。必须换保留判据。本实验同时量四条候选判据： + + J1 **不倒退 (do-no-harm)**：加了 skill 后 pass_rate 不下降。 + —— 这是**必要条件**，不是充分条件（什么都不说的废话 skill 也满足）。 + J2 **稳健性提升**：裸解 M 次里**并非全对**（0 narrative skill（无 rubric） +# =========================================================================== +# ⭐ 与 E18 的 SKILLGEN_SYSTEM 保持同一 narrative 家族（散文体、第一人称、不提外部上下文、 +# 不写代码块），只把「诊断失败」换成「复盘一次成功」。刻意**不引入** rubric 位 —— +# 本实验的自变量就是「没有 rubric」。 +SKILLGEN_SYSTEM_SUCCESS = ( + 'You are a Python programmer writing a note to your future self.\n\n' + 'You just solved a coding task on the first attempt. Write down the one ' + 'insight that made it work, so that next time you meet a task of this shape ' + 'you get it right immediately again.\n\n' + 'Requirements:\n' + '- Write one flowing narrative in the first person, not bullet points or headings.\n' + '- Name the concrete API, argument, keyword, data shape, or edge case that mattered.\n' + '- Write it as guidance that transfers to other tasks of the same shape, not a ' + 'description of this one task. Do not mention the specific function name you wrote.\n' + '- Do NOT include any code block, and do NOT restate the solution.\n' + '- If nothing non-obvious was involved, say so briefly instead of inventing a lesson.\n' + '- Keep it under 150 words.\n' + 'Wrap the note in and tags.') + + +def skillgen_success_prompt(problem: str, code: str) -> Dict[str, Any]: + """成功复盘 prompt。**无 rubric、无 GT** —— 只有题面和自己刚写对的代码。 + + ⭐ 「If nothing non-obvious was involved, say so briefly」这一句是刻意加的逃生口： + 首次成功的题很多是**平凡题**，逼模型硬编一条"经验"只会得到套话。给它说"没什么特别" + 的许可，才能让 J1（不倒退）这个判据真正区分出「有料」和「没料」。 + """ + user = (f'The task:\n{problem}\n\n' + f'The solution you wrote, which passed on the first attempt:\n' + f'```python\n{code}\n```\n\n' + 'Write the note to your future self.') + return {'messages': [{'role': 'system', 'content': SKILLGEN_SYSTEM_SUCCESS}, + {'role': 'user', 'content': user}]} + + +# =========================================================================== +# 工具（与 e18_collect_kod 同源） +# =========================================================================== +def run_samples(sampler, prompts, num_samples, max_tokens, gen_dp, + temperature=None, top_p=None): + if not prompts: + return [] + import copy + params = SamplingParams( + max_tokens=max_tokens, + temperature=0.6 if temperature is None else temperature, + top_p=0.95 if top_p is None else top_p, + num_samples=num_samples) + padded = prompts + if gen_dp > 1 and 0 < len(prompts) < gen_dp: + padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] + resp = sampler.sample(padded, params)[:len(prompts)] + return [list(r.sequences) if (r and r.sequences) else [] for r in resp] + + +def seq_text(seq) -> str: + return clean_text(getattr(seq, 'decoded', '') or '') if seq is not None else '' + + +def _mean(xs) -> float: + xs = [float(x) for x in xs if x is not None] + return sum(xs) / len(xs) if xs else 0.0 + + +def _pass_rate(rolls) -> float: + return (sum(1.0 for x in rolls if x['correct']) / len(rolls)) if rolls else 0.0 + + +def extract_skill(text: str) -> str: + if '' not in text: + return '' + body = text.split('', 1)[1] + return body.split('', 1)[0].strip() if '' in body else '' + + +# =========================================================================== +# 保留判据 +# =========================================================================== +IDENT = re.compile(r'\bdef\s+(\w+)') +# 「无信息」自述：模型用了 prompt 给的逃生口，说明它自己认为这题没什么可学的 +NOINFO = ('nothing non-obvious', 'nothing particularly', 'nothing special', + 'straightforward', 'no special', 'nothing unusual', 'not much to') + + +def decide_keep(base_rate: float, with_rate: float, skill: str, + gt_code: str, base_tokens: float, with_tokens: float) -> Dict[str, Any]: + """四条判据各自独立判定，不合成单一分数。 + + ⭐ 为何不合成一个总分：这四条问的是**不同的问题**，权重取决于 skill 池的用途 + （做 SFT 目标 vs 做检索库），此处只如实报出各判据的通过情况，把权衡留给决策。 + + J1 do-no-harm：with >= base。必要不充分 —— 一句废话也满足，所以**不能单独用它保留**。 + J2 稳健性：仅对 0=10%），说明少走弯路。 + 10% 门槛是任意的，但比"变短一点"要求高，避免采样噪声。 + J4 迁移性：skill 不含 GT 里的函数名 + 没使用「没什么特别」的自述。 + 含函数名 -> 只对本题有效；自述无信息 -> 模型自己承认没料。 + """ + harmless = with_rate >= base_rate - 1e-9 + robust = (base_rate < 1.0) and (with_rate >= 1.0 - 1e-9) + shorter = harmless and with_tokens > 0 and base_tokens > 0 and \ + (with_tokens <= 0.90 * base_tokens) + names = set(IDENT.findall(gt_code or '')) + low = (skill or '').lower() + has_name = any(n and n.lower() in low for n in names) + self_noinfo = any(p in low for p in NOINFO) + transfer = (not has_name) and (not self_noinfo) + return {'J1_harmless': harmless, 'J2_robust': robust, 'J3_shorter': shorter, + 'J4_transfer': transfer, 'has_own_fn_name': has_name, + 'self_says_noinfo': self_noinfo} + + +# =========================================================================== +# 主流程 +# =========================================================================== +def build_runtime(): + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, lazy_collect=False, groups=[ + DeviceGroup(name='skill', ranks=list(range(SKILL_GPUS)), device_type='GPU'), + DeviceGroup(name='exec', ranks=list(range(SKILL_GPUS, NUM_GPUS)), + device_type='GPU')]) + + def mk(group, world, thinking): + s = vLLMSampler(model_id=MODEL_ID, remote_group=group, + device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), + engine_args={'gpu_memory_utilization': GPU_MEM, + 'max_model_len': MAX_MODEL_LEN, + 'tensor_parallel_size': 1}) + s.set_template(Template, model_id=MODEL_ID, enable_thinking=thinking, + max_length=MAX_MODEL_LEN) + return s + + # skill 侧开 thinking（与 E18 采集口径一致，靠 thinking 拿多样性）；executor 关。 + return mk('skill', SKILL_GPUS, True), mk('exec', EXEC_GPUS, False) + + +def main(): + os.makedirs(OUTPUT_DIR, exist_ok=True) + t0 = time.time() + ds, _ = load_records(SEED, 0, OUTPUT_DIR) + pool = [r for i, r in enumerate(ds.dataset) if i < N_TASKS * POOL_MULT] + logger.info(f'E20 start: 题池 {len(pool)}（目标首次成功 {N_TASKS} 题）' + f' n_skills={N_SKILLS} bare={BARE_ROLLOUTS} exec={EXEC_ROLLOUTS}') + skill_sampler, exec_sampler = build_runtime() + + # ---- 1. 裸解，挑「第一次就通过」的题 ---- + bare = run_samples(exec_sampler, [direct_prompt(r['problem']) for r in pool], + BARE_ROLLOUTS, EXEC_MAX_TOKENS, EXEC_GPUS, + temperature=EXEC_TEMPERATURE, top_p=EXEC_TOP_P) + pairs, spans = [], [] + for r, seqs in zip(pool, bare): + seqs = list(seqs or []) + spans.append(len(seqs)) + pairs.extend((s, r['reference_answer']) for s in seqs) + judged = judge_seqs(pairs) if pairs else [] + picked, i = [], 0 + for r, n in zip(pool, spans): + rolls = judged[i:i + n] + i += n + if not rolls or not rolls[0]['correct']: + continue # ⭐ 只要**第 1 次**就通过的题（按用户要求） + picked.append({'rec': r, 'base_rate': _pass_rate(rolls), + 'code': rolls[0].get('code') or '', + 'base_tokens': _mean([x.get('gen_tokens') for x in rolls])}) + if len(picked) >= N_TASKS: + break + logger.info(f'[bare] 首次成功 {len(picked)}/{len(pool)} 题' + f'（其中 base<1 的不稳定题 {sum(1 for p in picked if p["base_rate"] < 1.0)}）') + if not picked: + raise RuntimeError('没有首次成功的题') + + # ---- 2. 成功 traj -> narrative skill（无 rubric）---- + sg = run_samples(skill_sampler, + [skillgen_success_prompt(p['rec']['problem'], p['code']) + for p in picked], + N_SKILLS, SKILL_MAX_TOKENS, SKILL_GPUS, + temperature=SKILL_TEMPERATURE) + flat = [] + for p, seqs in zip(picked, sg): + for ci in range(N_SKILLS): + seq = seqs[ci] if seqs and ci < len(seqs) else None + sk = extract_skill(seq_text(seq)) + flat.append({'p': p, 'cand_idx': ci, 'skill': sk, + 'stop': getattr(seq, 'stop_reason', None) if seq else None}) + n_ok = sum(1 for f in flat if f['skill']) + logger.info(f'[skillgen] {len(flat)} 候选，可解析 {n_ok} ({100*n_ok/len(flat):.0f}%)') + + # ---- 3. 带 skill 重解 ---- + todo = [f for f in flat if f['skill']] + ws = run_samples(exec_sampler, + [skill_solve_prompt(f['p']['rec']['problem'], f['skill']) for f in todo], + EXEC_ROLLOUTS, EXEC_MAX_TOKENS, EXEC_GPUS, + temperature=EXEC_TEMPERATURE, top_p=EXEC_TOP_P) + pairs, spans = [], [] + for f, seqs in zip(todo, ws): + seqs = list(seqs or []) + spans.append(len(seqs)) + pairs.extend((s, f['p']['rec']['reference_answer']) for s in seqs) + judged = judge_seqs(pairs) if pairs else [] + i = 0 + for f, n in zip(todo, spans): + rolls = judged[i:i + n] + i += n + f['with_rate'] = _pass_rate(rolls) + f['with_tokens'] = _mean([x.get('gen_tokens') for x in rolls]) + + # ---- 4. 判据 ---- + out = os.path.join(OUTPUT_DIR, 'e20_candidates.jsonl') + with open(out, 'w', encoding='utf-8') as fh: + for f in flat: + p = f['p'] + if not f['skill']: + row = {'data_id': p['rec']['data_id'], 'cand_idx': f['cand_idx'], + 'parseable': False, 'base_rate': p['base_rate'], + 'with_rate': None, 'skill': '', 'skill_chars': 0} + else: + d = decide_keep(p['base_rate'], f['with_rate'], f['skill'], + p['rec']['reference_answer'].get('canonical_solution', ''), + p['base_tokens'], f['with_tokens']) + row = {'data_id': p['rec']['data_id'], 'cand_idx': f['cand_idx'], + 'parseable': True, 'base_rate': p['base_rate'], + 'with_rate': f['with_rate'], + 'delta': round(f['with_rate'] - p['base_rate'], 6), + 'base_tokens': round(p['base_tokens'], 1), + 'with_tokens': round(f['with_tokens'], 1), + 'skill_chars': len(f['skill']), 'stop': f['stop'], + 'skill': f['skill'], **d} + fh.write(json.dumps(row, ensure_ascii=False) + '\n') + logger.info(f'[done] 落盘 {out}，用时 {(time.time()-t0)/60:.1f} 分钟') + + +if __name__ == '__main__': + main() diff --git a/cookbook/human_e18/e21_paired_rubric.py b/cookbook/human_e18/e21_paired_rubric.py new file mode 100644 index 000000000..b56659d4b --- /dev/null +++ b/cookbook/human_e18/e21_paired_rubric.py @@ -0,0 +1,214 @@ +# -*- coding: utf-8 -*- +"""E21：rubric 值多少钱？在**同一批题**上做成功复盘 vs 失败诊断的配对对照。 + +目标题型 = 「首次成功但不稳定」（0 < base_pass_rate < 1）。 +⭐ 为什么只能用这类题：它们**同时拥有**成功轨迹和失败轨迹，所以同一道题可以同时喂给 +两个 arm，构成**配对设计**（paired design）。base=1.0 的题没有失败轨迹（无法出 rubric）， +base=0 的题没有成功轨迹（无法做成功复盘）—— 只有这个交集能做干净对照。 +配对比独立分组强得多：题目难度是最大的方差来源，配对把它消掉了。 + +两个 arm，除了输入信号完全同构（同题、同 N_SKILLS、同 executor、同温度、同 rollout 数）： + + arm SUCCESS : 成功代码 -> narrative skill（无 rubric） [E20 的 prompt] + arm RUBRIC : 失败代码 + 报错 -> 教师诊断出 rubric -> narrative skill [E18 的 prompt] + +判据统一为 J2（升到全对）：base<1 的题加 skill 后 with_pass_rate 是否达到 1.0。 +这是唯一在两个 arm 上都可测、且不受选择偏差污染的口径。 + +⚠️ 已知的不对称（诚实记录，不是 bug）： + 1. RUBRIC arm 多消耗一次教师 API 调用（不占 GPU，但不是零成本）。 + 2. 两个 arm 的 prompt 家族不同（SKILLGEN_SYSTEM vs SKILLGEN_SYSTEM_SUCCESS）， + 所以测的是「成功复盘管线」vs「失败诊断管线」的**整体**差异， + 不是「rubric 这一个字段」的净效应。要拆到字段级需要第三个 arm + （失败代码但不给 rubric），本脚本用 FAILONLY arm 补上。 + 3. FAILONLY arm 复用 SKILLGEN_SYSTEM 但把 rubric 位填成「无诊断」的官方 fallback + （skillgen_prompt 内建该分支），所以 arm 间 system prompt 一致， + RUBRIC vs FAILONLY 的差值才是 rubric 字段的净贡献。 +""" +import json +import os +import sys +import time +from collections import defaultdict +from typing import Any, Dict, List + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import SamplingParams +from twinkle.sampler import vLLMSampler +from twinkle.template import Template + +_HERE = os.path.dirname(os.path.abspath(__file__)) +for _p in (_HERE, os.path.abspath(os.path.join(_HERE, '..', 'human'))): + if _p not in sys.path: + sys.path.insert(0, _p) + +from e18_kodcode import clean_text, judge_seqs, load_records # noqa: E402 +from e18_prompts import direct_prompt, skill_solve_prompt, skillgen_prompt # noqa: E402 +from e20_success_skill import (extract_skill, run_samples, seq_text, # noqa: E402 + skillgen_success_prompt, _mean, _pass_rate) +from e23_rubric import RubricCache, build_checker # noqa: E402 + +logger = get_logger() + +MODEL_ID = os.environ.get('MODEL_ID', 'Qwen/Qwen3-4B') +OUTPUT_DIR = os.environ.get('OUTPUT_DIR', os.path.join(_HERE, 'output.e21.paired')) +SEED = int(os.environ.get('SEED', 42)) +SKILL_GPUS = int(os.environ.get('SKILL_GPUS', 4)) +EXEC_GPUS = int(os.environ.get('EXEC_GPUS', 4)) +NUM_GPUS = SKILL_GPUS + EXEC_GPUS +GPU_MEM = float(os.environ.get('GPU_MEM', 0.85)) +MAX_MODEL_LEN = int(os.environ.get('MAX_MODEL_LEN', 24000)) + +POOL = int(os.environ.get('POOL', 400)) # 题池；不稳定题约占 4-5% +N_SKILLS = int(os.environ.get('N_SKILLS', 4)) +BARE_ROLLOUTS = int(os.environ.get('BARE_ROLLOUTS', 4)) +EXEC_ROLLOUTS = int(os.environ.get('EXEC_ROLLOUTS', 8)) +SKILL_MAX_TOKENS = int(os.environ.get('SKILL_MAX_TOKENS', 8192)) +EXEC_MAX_TOKENS = int(os.environ.get('EXEC_MAX_TOKENS', 15000)) +EXEC_TEMPERATURE = float(os.environ.get('EXEC_TEMPERATURE', 0.6)) +EXEC_TOP_P = float(os.environ.get('EXEC_TOP_P', 0.95)) +SKILL_TEMPERATURE = float(os.environ.get('SKILL_TEMPERATURE', 1.0)) + + +def build_runtime(): + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, lazy_collect=False, groups=[ + DeviceGroup(name='skill', ranks=list(range(SKILL_GPUS)), device_type='GPU'), + DeviceGroup(name='exec', ranks=list(range(SKILL_GPUS, NUM_GPUS)), + device_type='GPU')]) + + def mk(group, world, thinking): + s = vLLMSampler(model_id=MODEL_ID, remote_group=group, + device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), + engine_args={'gpu_memory_utilization': GPU_MEM, + 'max_model_len': MAX_MODEL_LEN, + 'tensor_parallel_size': 1}) + s.set_template(Template, model_id=MODEL_ID, enable_thinking=thinking, + max_length=MAX_MODEL_LEN) + return s + + return mk('skill', SKILL_GPUS, True), mk('exec', EXEC_GPUS, False) + + +def main(): + os.makedirs(OUTPUT_DIR, exist_ok=True) + t0 = time.time() + ds, _ = load_records(SEED, 0, OUTPUT_DIR) + pool = [r for i, r in enumerate(ds.dataset) if i < POOL] + logger.info(f'E21 start: 题池 {len(pool)} n_skills={N_SKILLS} ' + f'bare={BARE_ROLLOUTS} exec={EXEC_ROLLOUTS}') + skill_sampler, exec_sampler = build_runtime() + + # ---- 1. 裸解，挑「首次成功但不稳定」的题 ---- + bare = run_samples(exec_sampler, [direct_prompt(r['problem']) for r in pool], + BARE_ROLLOUTS, EXEC_MAX_TOKENS, EXEC_GPUS, + temperature=EXEC_TEMPERATURE, top_p=EXEC_TOP_P) + pairs, spans = [], [] + for r, seqs in zip(pool, bare): + seqs = list(seqs or []) + spans.append(len(seqs)) + pairs.extend((s, r['reference_answer']) for s in seqs) + judged = judge_seqs(pairs) if pairs else [] + picked, i, n_first_ok = [], 0, 0 + for r, n in zip(pool, spans): + rolls = judged[i:i + n] + i += n + if not rolls or not rolls[0]['correct']: + continue + n_first_ok += 1 + rate = _pass_rate(rolls) + if rate >= 1.0: + continue # 稳定成功 -> 无失败轨迹，做不了配对 + bad = next((x for x in rolls[1:] if not x['correct']), None) + if bad is None: + continue + picked.append({'rec': r, 'base_rate': rate, + 'good_code': rolls[0].get('code') or '', + 'bad_roll': bad, + 'base_tokens': _mean([x.get('gen_tokens') for x in rolls])}) + logger.info(f'[bare] 首次成功 {n_first_ok}/{len(pool)}；' + f'其中**不稳定**（可配对）{len(picked)} 题') + if not picked: + raise RuntimeError('没有可配对的不稳定题') + + # ---- 2. 教师诊断出 rubric（纯 API，不占 GPU）---- + # ⭐ 缓存必须用**本 run 私有**的路径，不能复用全局 RUBRIC_CACHE_PATH： + # RubricCache 的键里**不含轨迹**，其跨 run 复用的前提是「executor 冻结在 T=0， + # 同一题裸解逐字相同」（见 e23_rubric.py:227 的说明）。本实验裸解用 T=0.6， + # 轨迹每次不同，共用全局缓存会拿到**别的轨迹**的诊断，静默污染 RUBRIC arm。 + cache = RubricCache(os.path.join(OUTPUT_DIR, 'diag_cache.jsonl')) + checker = build_checker() + rubrics = cache.diagnose_many(checker, [(p['rec'], p['bad_roll']) for p in picked]) + n_rub = sum(1 for x in rubrics for _ in [x] if x) + logger.info(f'[rubric] {n_rub}/{len(picked)} 题拿到诊断') + for p, rb in zip(picked, rubrics): + p['rubric'] = rb or '' + + # ---- 3. 三个 arm 生成 skill ---- + # ⭐ 三个 arm 一次性拼进同一个 sample 调用，保证同一批权重、同一批 KV cache 状态， + # 避免"先跑完 A 再跑 B"引入的引擎状态差异。 + arms: List[str] = ['SUCCESS', 'RUBRIC', 'FAILONLY'] + prompts, meta = [], [] + for p in picked: + prompts.append(skillgen_success_prompt(p['rec']['problem'], p['good_code'])) + meta.append((p, 'SUCCESS')) + prompts.append(skillgen_prompt(p['rec']['problem'], p['rubric'], False)) + meta.append((p, 'RUBRIC')) + # FAILONLY：同 system prompt，rubric 位走内建的「无诊断」fallback + prompts.append(skillgen_prompt(p['rec']['problem'], '', False)) + meta.append((p, 'FAILONLY')) + sg = run_samples(skill_sampler, prompts, N_SKILLS, SKILL_MAX_TOKENS, + SKILL_GPUS, temperature=SKILL_TEMPERATURE) + flat = [] + for (p, arm), seqs in zip(meta, sg): + for ci in range(N_SKILLS): + seq = seqs[ci] if seqs and ci < len(seqs) else None + flat.append({'p': p, 'arm': arm, 'cand_idx': ci, + 'skill': extract_skill(seq_text(seq))}) + for arm in arms: + g = [f for f in flat if f['arm'] == arm] + n = sum(1 for f in g if f['skill']) + logger.info(f'[skillgen] {arm}: {n}/{len(g)} 可解析 ({100*n/max(1,len(g)):.0f}%)') + + # ---- 4. 带 skill 重解 ---- + todo = [f for f in flat if f['skill']] + ws = run_samples(exec_sampler, + [skill_solve_prompt(f['p']['rec']['problem'], f['skill']) for f in todo], + EXEC_ROLLOUTS, EXEC_MAX_TOKENS, EXEC_GPUS, + temperature=EXEC_TEMPERATURE, top_p=EXEC_TOP_P) + pairs, spans = [], [] + for f, seqs in zip(todo, ws): + seqs = list(seqs or []) + spans.append(len(seqs)) + pairs.extend((s, f['p']['rec']['reference_answer']) for s in seqs) + judged = judge_seqs(pairs) if pairs else [] + i = 0 + for f, n in zip(todo, spans): + rolls = judged[i:i + n] + i += n + f['with_rate'] = _pass_rate(rolls) + f['with_tokens'] = _mean([x.get('gen_tokens') for x in rolls]) + + # ---- 5. 落盘 ---- + out = os.path.join(OUTPUT_DIR, 'e21_candidates.jsonl') + with open(out, 'w', encoding='utf-8') as fh: + for f in flat: + p = f['p'] + row = {'data_id': p['rec']['data_id'], 'arm': f['arm'], + 'cand_idx': f['cand_idx'], 'base_rate': p['base_rate'], + 'parseable': bool(f['skill']), + 'with_rate': f.get('with_rate'), + 'base_tokens': round(p['base_tokens'], 1), + 'with_tokens': round(f.get('with_tokens') or 0, 1), + 'has_rubric': bool(p['rubric']), + 'skill_chars': len(f['skill']), 'skill': f['skill']} + if f.get('with_rate') is not None: + row['delta'] = round(f['with_rate'] - p['base_rate'], 6) + row['J2_robust'] = f['with_rate'] >= 1.0 - 1e-9 + fh.write(json.dumps(row, ensure_ascii=False) + '\n') + cache.close() + logger.info(f'[done] 落盘 {out}，用时 {(time.time()-t0)/60:.1f} 分钟') + + +if __name__ == '__main__': + main() diff --git a/cookbook/human_e18/run_collect_kod.sh b/cookbook/human_e18/run_collect_kod.sh new file mode 100755 index 000000000..55464e691 --- /dev/null +++ b/cookbook/human_e18/run_collect_kod.sh @@ -0,0 +1,8 @@ +#!/bin/bash +# KodCode 冷启动数据采集：8 卡全 rollout，无 SFT，narrative prompt 不变。 +cd /mnt/data/yzhao/tastelikefeet/twinkle/cookbook/human_e18 +export PYTHONPATH=/mnt/data/yzhao/tastelikefeet/twinkle/src:${PYTHONPATH} +export SKILL_SAMPLER_GPUS=2 BASE_SAMPLER_GPUS=6 # token 预算比 16.5:1，2+6 实测理论最优 +export KOD_SELFCHECK=0 # 不自检（坏题采集时自然不入池） +export TARGET_SAMPLES=5000 +exec /usr/local/bin/python -u e18_collect_kod.py diff --git a/cookbook/human_e18/run_sft_kod.sh b/cookbook/human_e18/run_sft_kod.sh new file mode 100755 index 000000000..3686b3400 --- /dev/null +++ b/cookbook/human_e18/run_sft_kod.sh @@ -0,0 +1,18 @@ +#!/bin/bash +# E18-KOD 离线 SFT + 首尾 eval：训一遍，看 loss 收敛 + 训练前后 lift 差异。 +# ⚠️ 会占满 8 卡（4 训练 + 2 skillmodel + 2 executor）—— 采集进程若还在跑，先确认它已停。 +# ⚠️ eval 要跑 judge（本地 pytest 沙箱），不需要教师 API key。 +cd /mnt/data/yzhao/tastelikefeet/twinkle/cookbook/human_e18 +export PYTHONPATH=/mnt/data/yzhao/tastelikefeet/twinkle/src:${PYTHONPATH} +export TRAIN_GPUS=4 # 4 训练 +export SKILL_SAMPLER_GPUS=2 # 2 张出 skill +export BASE_SAMPLER_GPUS=2 # 2 张跑 executor +export BATCH_SIZE=16 # 必须是 TRAIN_DP(=4) 的整倍数 +export MICRO_BATCH=8 +export LR=1e-5 # 与在线版一致：恒定 lr、无 warmup/decay +export EPOCHS=3 +export EVAL_SIZE=100 # 首尾各 100 题，选自 KodCode 中未生成过 skill 的题 +export EXEC_ROLLOUTS=1 # executor 单次 +export EXEC_TEMPERATURE=0.0 # greedy，可复现 +export SAVE_EVERY_STEPS=0 # 0 = 只在结束时存权重 +exec /usr/local/bin/python -u e18_sft_kod.py diff --git a/cookbook/human_e18/shard_tool.py b/cookbook/human_e18/shard_tool.py new file mode 100644 index 000000000..5615edcb4 --- /dev/null +++ b/cookbook/human_e18/shard_tool.py @@ -0,0 +1,134 @@ +# -*- coding: utf-8 -*- +"""多机分片采集的两个配套工具：导出「已跑过的题」种子 + 合并多机产物。 + + python3 shard_tool.py seed # 在 A 机跑，产出给 B 机的种子 + python3 shard_tool.py merge [dir2 ...] # 合并任意台机器的产物 + +为什么需要 seed： + crc32 分片只保证「以后」两台机器不撞，但 A 机已经跑掉的题会均匀落在所有分片上 + （实测 7374 题在 SHARD_N=2 下是 3667/3707）。B 机目录是空的，resume_done_ids() + 读不到任何东西，就会把属于自己分片的那 3707 题重跑一遍 —— 白烧一半算力， + 而且合并时同一 data_id 出现两份。种子文件就是把 A 机的 done id 搬给 B 机。 + ❗ 种子只需要 e18_candidates.jsonl 的 data_id 字段（resume_done_ids 只读这个）， + 所以导出的是**精简行**，不是整个 27000 行的候选文件 —— B 机不需要 A 机的 skill 正文。 +""" +import collections +import json +import os +import sys + +FILES = ('e18_sft_dataset.jsonl', 'e18_candidates.jsonl', 'collect_log.jsonl') + + +def _rows(path): + if not os.path.exists(path): + return + with open(path, encoding='utf-8') as f: + for ln in f: + ln = ln.strip() + if not ln: + continue + try: + yield json.loads(ln) + except Exception: + continue + + +def cmd_seed(src, out): + """把 src 里所有跑过的 data_id 导成最小的 candidates 行。""" + ids = {r['data_id'] for r in _rows(os.path.join(src, 'e18_candidates.jsonl')) + if r.get('data_id')} + with open(out, 'w', encoding='utf-8') as f: + for i in sorted(ids): + # resume_done_ids() 只取 data_id；其余字段给最小合法值， + # 保证这些行即使被别的分析脚本读到也不会伪装成真实候选： + # kept=False + parseable=False + seed=True 一眼可滤。 + f.write(json.dumps({'data_id': i, 'kept': False, 'parseable': False, + 'seed': True, 'run': 'SEED', 'chunk': -1}, + ensure_ascii=False) + '\n') + print('导出 %d 个已跑 data_id -> %s' % (len(ids), out)) + print('用法：拷到 B 机的 OUTPUT_DIR/e18_candidates.jsonl，再用 KOD_RESUME=1 启动') + + +def cmd_merge(out_dir, dirs): + os.makedirs(out_dir, exist_ok=True) + report = {} + # ---- 1. sft_dataset：按 data_id 去重（同题多机重复时保留 pass_gain 更高者）---- + best, dup = {}, 0 + for d in dirs: + for r in _rows(os.path.join(d, 'e18_sft_dataset.jsonl')): + k = r.get('data_id') + if not k: + continue + if k in best: + dup += 1 + # 保留 gain 高的：重复只可能来自"种子没同步"的意外， + # 此时保留更优样本比保留先到者更合理 + if (r.get('pass_gain') or -9) <= (best[k].get('pass_gain') or -9): + continue + best[k] = r + with open(os.path.join(out_dir, 'e18_sft_dataset.jsonl'), 'w', encoding='utf-8') as f: + for r in best.values(): + f.write(json.dumps(r, ensure_ascii=False) + '\n') + report['sft'] = (len(best), dup) + + # ---- 2. candidates：按 (data_id, run, cand_idx) 去重，丢掉 seed 占位行 ---- + seen, rows, nseed = set(), [], 0 + for d in dirs: + for r in _rows(os.path.join(d, 'e18_candidates.jsonl')): + if r.get('seed'): + nseed += 1 + continue + k = (r.get('data_id'), r.get('run'), r.get('cand_idx')) + if k in seen: + continue + seen.add(k) + rows.append(r) + with open(os.path.join(out_dir, 'e18_candidates.jsonl'), 'w', encoding='utf-8') as f: + for r in rows: + f.write(json.dumps(r, ensure_ascii=False) + '\n') + report['cand'] = (len(rows), nseed) + + # ---- 3. collect_log：chunk 编号两机都从 0 起，直接 cat 会产生歧义 ---- + # 不重编号（会破坏与 run.log 的对照），改为加 src 字段标机器来源， + # 并按 (run, chunk) 唯一化。分析脚本本来就该按 run 分组看 chunk。 + lg, seen2 = [], set() + for d in dirs: + tag = os.path.basename(os.path.normpath(d)) + for r in _rows(os.path.join(d, 'collect_log.jsonl')): + k = (r.get('run'), r.get('chunk')) + if k in seen2: + continue + seen2.add(k) + r['src'] = tag + lg.append(r) + lg.sort(key=lambda r: (str(r.get('run')), r.get('chunk') or 0)) + with open(os.path.join(out_dir, 'collect_log.jsonl'), 'w', encoding='utf-8') as f: + for r in lg: + f.write(json.dumps(r, ensure_ascii=False) + '\n') + report['log'] = (len(lg), 0) + + print('=== 合并完成 -> %s ===' % out_dir) + print(' e18_sft_dataset.jsonl %6d 条 (跨机重复丢弃 %d)' % report['sft']) + print(' e18_candidates.jsonl %6d 条 (滤掉种子占位 %d)' % report['cand']) + print(' collect_log.jsonl %6d 条' % report['log'][0]) + # 交叉校验：胜者的 data_id 必须都能在候选里找到 + cid = {r.get('data_id') for r in rows} + miss = [k for k in best if k not in cid] + print(' 一致性：胜者 data_id 在候选中缺失 %d 个 %s' + % (len(miss), '(OK)' if not miss else '<- 异常')) + per = collections.Counter(str(r.get('run')).split('.')[-1] for r in rows) + print(' 按分片计候选数：%s' % dict(per)) + + +if __name__ == '__main__': + if len(sys.argv) < 3: + print(__doc__) + sys.exit(2) + if sys.argv[1] == 'seed': + cmd_seed(sys.argv[2], sys.argv[3]) + elif sys.argv[1] == 'merge': + cmd_merge(sys.argv[2], sys.argv[3:]) + else: + print(__doc__) + sys.exit(2) From 19e1ecf7c26d068d66dba14e588781a83cba0cc8 Mon Sep 17 00:00:00 2001 From: root Date: Sun, 9 Aug 2026 00:21:50 +0800 Subject: [PATCH 36/60] fix --- .../sampler/vllm_sampler/vllm_sampler.py | 54 +++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py index 0433ef5a8..70369bf6d 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py @@ -480,6 +480,60 @@ async def _receive_and_load(): self._run_in_loop(_receive_and_load()) + @remote_function(dispatch='all', collect='first', lazy_collect=False) + def load_weights_from_path(self, path: Optional[str] = None): + """Reload base weights into the running vLLM engine directly from a checkpoint. + + Unlike :meth:`receive_weights`, this does **not** involve the training model: + weights are read from disk and streamed straight into vLLM. This is what makes + it possible to restore the sampler to a known checkpoint without a trainer + round-trip -- no ``save``/``load`` on the training model, so training weights + and optimizer state are never touched. + + Weights are yielded **lazily** one tensor at a time (never materialising a full + state dict) because ``VLLMEngine.update_weights`` accepts a generator and packs + tensors into fixed-size transfer buckets itself. Tensors stay on CPU, so the + engine takes its shared-memory path rather than CUDA IPC. + + Names are passed through untouched: safetensors files already store canonical + HF names, which is exactly what the worker's ``model.load_weights()`` expects + (it does the q/k/v -> qkv and gate/up -> gate_up stacking internally). + + Args: + path: Local checkpoint dir or a hub model id. Defaults to the ``model_id`` + the sampler was constructed with, i.e. the original pretrained weights. + """ + import glob + import json + from safetensors import safe_open + + path = path or self.model_id + checkpoint_dir = path if os.path.exists(path) else HubOperation.download_model(path) + + index_path = os.path.join(checkpoint_dir, 'model.safetensors.index.json') + if os.path.exists(index_path): + with open(index_path, encoding='utf-8') as f: + weight_map = json.load(f)['weight_map'] + shards = [os.path.join(checkpoint_dir, s) for s in sorted(set(weight_map.values()))] + else: + shards = sorted(glob.glob(os.path.join(checkpoint_dir, '*.safetensors'))) + if not shards: + raise FileNotFoundError(f'No .safetensors weights found under {checkpoint_dir}') + + def _iter_weights(): + # safe_open + get_tensor reads one tensor at a time (mmap-backed), so peak + # host memory is a single tensor rather than the whole shard. + for shard in shards: + with safe_open(shard, framework='pt', device='cpu') as f: + for name in f.keys(): + yield name, f.get_tensor(name) + + self._run_in_loop(self.engine.update_weights(_iter_weights(), base_sync_done=False)) + # A base-model load invalidates any previously synced LoRA adapter, mirroring + # the `not base_sync_done` branch of receive_weights(). + self.engine.invalidate_synced_lora() + logger.info(f'Reloaded base weights from {checkpoint_dir} ({len(shards)} shard(s))') + @remote_function(dispatch='all', collect='first', lazy_collect=False) def shutdown(self): """Gracefully shutdown the vLLM engine and background event loop. From 7eea3930103b7c8cd0628484012b7f7f702ef5b4 Mon Sep 17 00:00:00 2001 From: root Date: Sun, 9 Aug 2026 00:33:27 +0800 Subject: [PATCH 37/60] =?UTF-8?q?fix(cookbook):=20=E8=A1=A5=E6=8F=90?= =?UTF-8?q?=E4=BA=A4=20human/=20=E4=B8=8B=20e18=20=E4=BE=9D=E8=B5=96?= =?UTF-8?q?=E7=9A=84=20e23=20=E6=A8=A1=E5=9D=97?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit e18_collect_kod.py 把 cookbook/human 加进 sys.path 后 import e23_rubric， 但该目录此前 0 文件入库，导致新机器 clone 后启动即 ModuleNotFoundError: No module named 'e23_rubric'。 依赖链已静态遍历确认闭合： e18_collect_kod -> e18_{kodcode,multidiag,prompts,select} -> e23_rubric -> e23_prompts e23_bcb 由 e18 其他脚本引用，一并提交。 --- cookbook/human/e23_bcb.py | 278 +++++++++++++++++++++++++ cookbook/human/e23_prompts.py | 225 ++++++++++++++++++++ cookbook/human/e23_rubric.py | 309 ++++++++++++++++++++++++++++ cookbook/human/skill_drift_stats.py | 168 +++++++++++++++ 4 files changed, 980 insertions(+) create mode 100644 cookbook/human/e23_bcb.py create mode 100644 cookbook/human/e23_prompts.py create mode 100644 cookbook/human/e23_rubric.py create mode 100644 cookbook/human/skill_drift_stats.py diff --git a/cookbook/human/e23_bcb.py b/cookbook/human/e23_bcb.py new file mode 100644 index 000000000..4ff11acaa --- /dev/null +++ b/cookbook/human/e23_bcb.py @@ -0,0 +1,278 @@ +"""BigCodeBench 环境层：数据加载、沙箱单测判分、模型输出解析。 + +与训练完全解耦 —— 这里只回答「一段模型文本能不能跑过官方单测」，不认识 skill / rubric / GRPO。 +""" +import ast as _ast +import importlib.util +import json +import os +import re +import shutil +import subprocess +import sys +import tempfile +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Optional, Tuple + +from twinkle import get_logger +from twinkle.dataset import Dataset, DatasetMeta + +logger = get_logger() +_HERE = os.path.dirname(os.path.abspath(__file__)) + +# ModelScope 数据集 id 或本地 parquet 路径都行：DatasetMeta 按 os.path.exists 自行分流， +# 走本地文件时 subset / split 会被忽略。 +BCB_DATASET = os.environ.get('BCB_DATASET', 'ms://bigcode/bigcodebench') +BCB_SUBSET = os.environ.get('BCB_SUBSET', 'default') +BCB_SPLIT = os.environ.get('BCB_SPLIT', 'v0.1.0_hf') +TEST_WORKERS = int(os.environ.get('TEST_WORKERS', 24)) # 跑单测的线程池（每线程一个子进程） +TEST_TIMEOUT = int(os.environ.get('TEST_TIMEOUT', 60)) # 单题单测墙钟上限（秒） + +# 需要外网 / GUI / 子进程的库：沙箱里会挂或超时，判分噪声与 skill 无关 -> 整题排除。 +EXCLUDE_LIBS = {'requests', 'urllib', 'http', 'smtplib', 'socket', 'ssl', 'ftplib', + 'mechanize', 'wikipedia', 'turtle', 'tkinter', 'subprocess', 'sendgrid', + 'python_http_client', 'django', 'flask', 'flask_login', 'flask_mail', + 'flask_restful', 'flask_wtf', 'wtforms', 'multiprocessing'} +LIB_ALIAS = {'cv2': 'cv2', 'PIL': 'PIL', 'bs4': 'bs4', 'yaml': 'yaml', 'dateutil': 'dateutil', + 'Crypto': 'Crypto', 'docx': 'docx', 'pytz': 'pytz', 'psutil': 'psutil', + 'texttable': 'texttable', 'wordcloud': 'wordcloud', 'skimage': 'skimage', + 'PyPDF2': 'PyPDF2', 'sklearn': 'sklearn'} + +_FENCE_RE = re.compile(r'```(?:python|py)?\s*\n(.*?)```', re.S) +_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') + + +# ========== 输出解析 ========== +def after_think(text: str) -> str: + i = (text or '').rfind('') + return text[i + len(''):] if i >= 0 else (text or '') + + +def clean_text(decoded: Optional[str]) -> str: + """只剔 <|...|> 这类特殊 token 的**字面量**。 必须保留 —— E23 要把它递给 executor。""" + return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() + + +def extract_code(text: str) -> str: + """取最后一个能通过 ast.parse 的代码块；没有围栏就退化为整段（切 think 之后）。""" + body = after_think(text or '') + blocks = _FENCE_RE.findall(body) + for b in reversed(blocks): + try: + _ast.parse(b) + return b + except SyntaxError: + continue + if blocks: + return blocks[-1] + try: + _ast.parse(body) + return body + except SyntaxError: + return '' + + +def extract_skill(text: str) -> str: + """抽 块；任何畸形（未闭合 / 只写在 think 里 / 空块）一律返回 ''。 + + 只在 **之后**找：写在思考过程里的 不算产出。返回 '' 时 executor 走干净 + direct（见 skill_solve_prompt），不会拿到半截内容。 + """ + answer = after_think(text) + s = answer.lower().rfind('') + if s < 0: + return '' + inner = s + len('') + e = answer.lower().find('', inner) + if e < 0: + return '' + return re.sub(r'', '', + answer[inner:e].strip(), flags=re.IGNORECASE).strip() + + +# ========== 沙箱判分 ========== +_RUNNER = """ +import unittest, sys +loader = unittest.TestLoader() +suite = loader.loadTestsFromTestCase(TestCases) +res = unittest.TextTestRunner(verbosity=0, stream=sys.stderr).run(suite) +print('__BCB__', res.testsRun, len(res.failures), len(res.errors)) +sys.exit(0 if res.wasSuccessful() and res.testsRun > 0 else 1) +""" + + +def _trim_err(err: str, limit: int = 1600) -> str: + """保留失败测试名与异常行，砍掉冗长 traceback 帧 —— 这是喂给 rubric 的客观证据。 + 随机临时目录名换成 ，否则同一个失败在两次运行里看起来不一样。""" + err = re.sub(r'/tmp/bcb_[A-Za-z0-9_]+', '', err or '') + lines = [ln for ln in err.splitlines() if ln.strip()] + keep = [ln for ln in lines + if ln.startswith(('FAIL:', 'ERROR:', 'AssertionError', 'Traceback')) + or re.match(r'^\w*(Error|Exception|Warning)\b', ln.strip()) + or ', in ' in ln] + return '\n'.join(keep or lines[-25:])[-limit:] + + +def run_tests(code: str, payload: Dict[str, Any], timeout: int = TEST_TIMEOUT) -> Dict[str, Any]: + """子进程里跑「提交代码 + 官方 test + _RUNNER」。-> {'passed', 'kind', 'error'}。""" + if not code.strip(): + return {'passed': False, 'kind': 'no_code', 'error': 'no parseable code block'} + if payload['entry_point'] not in code: + return {'passed': False, 'kind': 'no_entry', + 'error': f"function {payload['entry_point']} is not defined in the submitted code"} + tmp = tempfile.mkdtemp(prefix='bcb_') + try: + path = os.path.join(tmp, 'run_case.py') + with open(path, 'w', encoding='utf-8') as f: + f.write(code + '\n\n' + payload['test'] + '\n' + _RUNNER) + env = dict(os.environ, MPLBACKEND='Agg', PYTHONHASHSEED='0', OMP_NUM_THREADS='1', + MKL_NUM_THREADS='1', TOKENIZERS_PARALLELISM='false') + env.pop('CUDA_VISIBLE_DEVICES', None) + try: + p = subprocess.run([sys.executable, path], cwd=tmp, env=env, timeout=timeout, + capture_output=True, text=True, errors='replace') + except subprocess.TimeoutExpired: + return {'passed': False, 'kind': 'timeout', + 'error': f'the tests did not finish within {timeout}s'} + n_tests = n_fail = n_err = 0 + for line in (p.stdout or '').splitlines(): + if line.startswith('__BCB__'): + _, a, b, c = line.split() + n_tests, n_fail, n_err = int(a), int(b), int(c) + if p.returncode == 0 and n_tests > 0: + return {'passed': True, 'kind': 'pass', 'error': ''} + kind = 'assertion' if n_fail else ('exception' if n_err else 'import_or_syntax') + return {'passed': False, 'kind': kind, + 'error': _trim_err((p.stderr or '').replace(tmp, ''))} + finally: + shutil.rmtree(tmp, ignore_errors=True) + + +def empty_roll() -> Dict[str, Any]: + return {'correct': False, 'stop_reason': 'empty', 'gen_tokens': 0, 'text': '', 'code': '', + 'kind': 'no_code', 'error': ''} + + +def judge_seqs(pairs: List[Tuple[Any, Dict[str, Any]]]) -> List[Dict[str, Any]]: + """[(采样 sequence 或 None, payload)] -> rolls。所有判分都汇合到这里。 + + 必须批量：单测是子进程（导入 pandas/sklearn 后典型 1-3s），一个 chunk 几百次判分串行会比同 + chunk 的 GPU 时间还长一个量级。同 (task_id, code) 只跑一次 —— T=0 的 executor 经常对同一题 + 产出逐字相同的代码。 + """ + rolls: List[Dict[str, Any]] = [] + keys: List[Optional[Tuple[str, str]]] = [] + jobs: Dict[Tuple[str, str], Dict[str, Any]] = {} + for seq, payload in pairs: + if seq is None: + rolls.append(empty_roll()) + keys.append(None) + continue + text = clean_text(getattr(seq, 'decoded', '') or '') + code = extract_code(text) + key = (payload['task_id'], code) + rolls.append({'correct': False, 'stop_reason': getattr(seq, 'stop_reason', None), + 'gen_tokens': len(getattr(seq, 'tokens', None) or []), + 'text': text, 'code': code, 'kind': None, 'error': ''}) + keys.append(key) + jobs.setdefault(key, payload) + if jobs: + todo = list(jobs) + with ThreadPoolExecutor(max_workers=max(1, min(TEST_WORKERS, len(todo)))) as ex: + verdicts = dict(zip(todo, ex.map(lambda k: run_tests(k[1], jobs[k]), todo))) + for roll, key in zip(rolls, keys): + v = verdicts.get(key) if key is not None else None + if v is not None: + roll['correct'] = bool(v['passed']) + roll['kind'], roll['error'] = v['kind'], v['error'] + return rolls + + +# ========== 数据 ========== +# reference_answer 的字段集：判分需要的一切（code 域不是数值答案）。 +_PAYLOAD_KEYS = ('task_id', 'entry_point', 'test', 'code_prompt', 'doc_struct', + 'canonical_solution') + + +def _importable(lib: str) -> bool: + try: + return importlib.util.find_spec(LIB_ALIAS.get(lib, lib).split('.')[0]) is not None + except Exception: + return False + + +def _row_libs(row: Dict[str, Any]) -> List[str]: + v = row.get('libs') + if isinstance(v, str): # 数据集里存的是 list 的字符串形式 + try: + return list(_ast.literal_eval(v)) + except Exception: + return [] + return list(v or []) + + +def _to_record(batch: Dict[str, List]) -> Dict[str, List]: + """原始 BCB 列 -> {'data_id', 'problem', 'reference_answer'}。 + + Dataset.map 强制 batched=True，所以这里收发的都是列式 batch。 + """ + return {'data_id': list(batch['task_id']), + 'problem': list(batch['instruct_prompt']), + 'reference_answer': [{k: batch[k][i] for k in _PAYLOAD_KEYS} + for i in range(len(batch['task_id']))]} + + +def _broken_tasks(ds: Dataset, output_dir: str) -> set: + """参考解答跑不过自己的单测 = 沙箱/依赖不可判定，不是模型的错（实测约 7.5%）。 + 自检一次后落盘缓存，题数不变则复用。必须在 map 之前调用（要读原始列）。""" + path = os.path.join(output_dir, 'bcb_broken_tasks.json') + if os.path.exists(path): + try: + with open(path, encoding='utf-8') as f: + c = json.load(f) + if int(c.get('n_tasks', -1)) == len(ds): + return set(c['broken']) + except Exception as exc: + logger.warning(f'[data] 读取 {path} 失败（{exc}），重跑自检') + logger.info(f'[data] 沙箱自检：{len(ds)} 道题跑参考解答（一次性，之后走缓存）…') + rows = [ds[i] for i in range(len(ds))] + jobs = [(r['code_prompt'] + (r['canonical_solution'] or ''), {k: r[k] for k in _PAYLOAD_KEYS}) + for r in rows] + with ThreadPoolExecutor(max_workers=max(1, min(TEST_WORKERS, len(jobs)))) as ex: + vers = list(ex.map(lambda p: run_tests(p[0], p[1]), jobs)) + broken = {r['task_id'] for r, v in zip(rows, vers) if not v['passed']} + with open(path, 'w', encoding='utf-8') as f: + json.dump({'n_tasks': len(ds), 'broken': sorted(broken)}, f, indent=1) + return broken + + +def load_records(seed: int, eval_size: int, + output_dir: str) -> Tuple[Dataset, List[Dict[str, Any]]]: + """-> (train_dataset, eval_records)，每条记录是 {'data_id', 'problem', 'reference_answer'}。 + + 训练侧返回 Dataset 交给调用方喂 DataLoader；holdout 是固定的一小批、每轮 eval 整体遍历， + 没有分批的意义，直接物化成 list。 + + BigCodeBench 没有 difficulty 字段，所以不做分层：过滤完按 seed 洗牌，前 eval_size 道作 + holdout。题池很小（1140 -> 剔除后约 900），重复抽到的题 rubric 全部缓存命中，成本只在 + GPU rollout。 + """ + ds = Dataset(DatasetMeta(BCB_DATASET, subset_name=BCB_SUBSET, split=BCB_SPLIT)) + n_raw = len(ds) + ds.filter(lambda r: not (set(_row_libs(r)) & EXCLUDE_LIBS)) + n_kept_libs = len(ds) + ds.filter(lambda r: all(_importable(x) for x in _row_libs(r))) + logger.info(f'[data] BigCodeBench: 全集 {n_raw}，剔除需外网/GUI/子进程 {n_raw - n_kept_libs}、' + f'依赖缺失 {n_kept_libs - len(ds)} -> {len(ds)}') + + broken = _broken_tasks(ds, output_dir) + if broken: + ds.filter(lambda r: r['task_id'] not in broken) + logger.info(f'[data] 剔除参考解答自己跑不过单测的题 {len(broken)} 道 -> 可用 {len(ds)}') + ds.map(_to_record, remove_columns=ds.dataset.column_names) + + shuffled = ds.dataset.shuffle(seed=seed) + n_eval = min(eval_size, len(shuffled)) if eval_size > 0 else 0 + eval_records = list(shuffled.select(range(n_eval))) + train_dataset = Dataset(DatasetMeta(data=shuffled.select(range(n_eval, len(shuffled))))) + return train_dataset, eval_records diff --git a/cookbook/human/e23_prompts.py b/cookbook/human/e23_prompts.py new file mode 100644 index 000000000..5e73f2f40 --- /dev/null +++ b/cookbook/human/e23_prompts.py @@ -0,0 +1,225 @@ +"""E23 的全部 prompt 文本与拼装函数：executor / 教师 judge / skill-gen 三处。""" +# flake8: noqa: E501 +# prompt 正文按「一段一行」书写，折行会改变真正发给模型的文本，故整文件豁免行长检查。 +import json +from typing import Any, Dict + +# =========================================================================== +# executor +# =========================================================================== +# BigCodeBench 官方 instruct 模式的硬性交付要求。 +EXEC_SYSTEM = """\ +You are an expert Python engineer. You will be given a task description that ends with the exact \ +import lines and function signature your solution must start with. + +Deliver exactly one fenced Python code block and nothing else after it: +- Reproduce the given imports and the given function signature verbatim, including parameter \ +names, order and default values. +- Add any further imports you need inside the same block; the block must run standalone. +- Return exactly the object type the task says to output. If it says the function should output \ +a tuple, return a tuple in that order; if it names a matplotlib Axes, return the Axes object \ +itself, not the Figure and not None. +- Implement the described behaviour for the general case, including the empty / single-element / \ +missing-column edge cases and any exception the description says to raise. +- Do not call the function, do not print demonstrations, do not add tests, do not use \ +`if __name__ == '__main__'`, and do not read from stdin. +- Do not include explanations outside the code block.""" + +# ⭐ E23 的定义性特征：executor 看到 actor 的**完整产出**（ + ），不是只看抽出来的 +# 块。第一句必须说明「接下来这坨是什么」，否则 executor 会把 当成任务描述的一 +# 部分。第二句与「只给 skills」的对照臂逐字相同，使两臂只差「能否看到推理过程」这一个变量。 +_WRAPPER_WITH_THINK = ( + 'Guidance model transcript:\nFor this task, a separate problem-solving guidance model was ' + 'asked to analyse it and produce advisory skills. Its full output follows, including its ' + 'private reasoning:\n{hint}\n' + 'Prefer using its techniques when they fit, but if you have a clearly better ' + 'implementation, you may diverge. Be concise and accurate.\n') + + +def direct_prompt(problem: str) -> Dict[str, Any]: + return {'messages': [{'role': 'system', 'content': EXEC_SYSTEM}, + {'role': 'user', 'content': problem}]} + + +def skill_solve_prompt(problem: str, skill: str, raw_response: str = '') -> Dict[str, Any]: + """题面 + actor 全文（含 ）。 + + skill 为空（actor 只写了 ，或写到上限没闭合）-> 干净 direct，而不是把一坨无结论的思考 + 过程当指导塞进去。故意**不补 <|im_end|>**：那会让 advisory 落在非法 ChatML 位置并改变紧邻 + token 的 BPE 切分，只为「与 SEAM 逐 token 对齐」才值得，这里的代价是多一个变量。 + """ + skill = (skill or '').strip() + if not skill: + return direct_prompt(problem) + hint = (raw_response or '').strip() or skill + return {'messages': [{'role': 'system', 'content': EXEC_SYSTEM}, + {'role': 'user', + 'content': problem + '\n\n' + _WRAPPER_WITH_THINK.format(hint=hint)}]} + + +# =========================================================================== +# 教师 judge（rubric 诊断） +# =========================================================================== +DIAG_SYSTEM = """\ +You are a code-failure classifier. You are given a Python task description (with the required signature), a list of FAILURE CLASSES, one attempted implementation, and the REAL error that attempt produced when the task's unit tests were run. + +Your job is NOT to review the code. It is to (a) name the SINGLE decisive failure class, (b) quote the evidence for it out of the test error, and (c) state the prior knowledge that would have PREVENTED it. + +Output STRICT JSON (no prose outside it), either + +{"addressable": true, "class": "", "evidence": "", "required_value": "", "required_value_source": "", "reason": "", "prior": "", "secondary": [""], "independent_causes": 1} + +or, when you cannot ground a decisive class in the test error: + +{"addressable": false, "why": ""} + +Rules: +- "evidence" MUST be copied verbatim out of the test error you were given. If the error does not let you single out ONE class, output addressable=false instead of guessing. Never invent evidence. +- The error tells you WHICH assertion fired first, not WHAT caused it. A failing shape, type or value assertion is normally the last link of the chain, not the class. Ask "which wrong belief produced this?" and classify THAT. Classify by the assertion itself only when nothing upstream is wrong: when the computation is right and merely the kind of object handed back is not what the caller reads. +- "class" is the one class the evidence implicates. List other classes that are also off in "secondary"; leave it empty if there are none. +- "independent_causes" counts mutually independent root causes, not symptoms of one. Three or more means no single short warning could have saved this attempt. +- "prior" is the hard part. It MUST STAY TRUE AND USEFUL IF THIS TASK IS DELETED: a fact about a library, an API default, a format directive, or a testing convention. It must NOT contain any identifier, literal, column name, file name or number taken from this task, and must NOT describe the steps of this task's solution. Write "pandas writes JSON Lines rather than one JSON document when the lines flag is set", NOT "pass lines=False here". +- "prior" must say WHY the wrong result arose -- the rule the attempt had backwards -- and NOT what the correct result should look like. "The result must carry one row per input key" merely restates the assertion that failed and teaches nothing, because the reader already has the task description; "a merge defaults to an inner join and silently drops keys missing on either side" is the belief that was actually absent. +- For class TESTCONTRACT and class EXCEPTION, "required_value" MUST NOT be null. Those two classes are BY DEFINITION about failing to match something the caller demands -- a title, a label, a key, an exception -- so name that thing. If you find yourself unable to name it, you have the wrong class. +- "required_value" and "required_value_source" are a CITATION, not an opinion. Whenever the failure comes down to matching something the tests demand, put that demanded thing in "required_value", then go back to the TASK section and copy out the fragment that states it into "required_value_source". You may only fill "required_value_source" with text you can actually see in the TASK section -- it is checked against it verbatim. If the task never states it, write null, and the diagnosis will be discarded, which is the correct outcome: an engineer reading only the task could not have known it either. +- That check is where a plausible-looking diagnosis does the most damage. "Tests assert exact string equality on titles" is a true and transferable sentence, and it is still worthless when the title itself appears nowhere in the task -- the engineer learns that the string matters but not what it is. Do not let a good-sounding prior talk you out of the citation. +- Set addressable=false rather than writing a vacuous prior such as "implement the description carefully". A sentence that merely says validation must be written, or that a parameter must satisfy the tests, is vacuous however factual it sounds; a usable prior names a library, function, format or convention that the reader could look up. +- Output only the JSON object.""" + +DIAG_USER = """\ +## Task +{query} + +## Failure classes +{rubric} + +## Attempted implementation and its test error +{segment} + +Now output the classification JSON object.""" + +# ⭐ 判据按「什么先验知识能预测这个失败」切分，而不是按「代码哪里错了」。 +# 依据是 E23.t1 的 480 组实测：旧表第 3 条只管**签名合法性**（调用能否被接受），于是所有「调用被 +# 接受、但行为/默认值不符合假设」的失败（to_json(lines=True)、mkdir 不带 parents、glob('*') 含 +# 目录、strftime('%Z') 多后缀）全掉进兜底项「核心计算错」——它在 76% 的题上 FAIL，且 leak 率 +# 0.106 是全体 0.021 的 5 倍：它救得回题，靠的正是让 skill 把解法写出来。BEHAVIOUR 就是补这个洞， +# 占旧兜底 FAIL 的 53%。每条都按**报错里的可观测特征**定义，这才能既判得准又分得开。 +# +# ⚠️ 「按可观测特征定义」的代价是教师容易**照最先响的那条断言分类**，于是 SHAPE 会接手本属 +# BEHAVIOUR 的题（BCB/441：einsum 输入下标 ikl 应为 jkl，症状是 shape 断言先挂 -> 判 SHAPE -> +# prior 只讲输出下标决定形状 -> 8 个候选一起只改输出下标，形状对了数值仍错，整组 reward 0）。 +# SHAPE 的定义因此显式排除「算错导致形状/数值不对」，DIAG_SYSTEM 里也有一条反症状规则兜着。 +# (code, 给 skill-gen 的短标签, 给教师判定用的完整定义) +FAILURE_CLASSES = [ + ('BEHAVIOUR', 'a library call behaves differently from what was assumed', + 'The call is accepted, but the function\'s real behaviour, default value or precondition ' + 'differs from what the code assumed: a flag that changes the output format, a default that ' + 'does not do what its name suggests, a required preparation step, a half-open range.'), + ('SIGNATURE', 'the call itself is not accepted', + 'The function, attribute or module does not exist, or it does not take the argument names, ' + 'positions or count that were passed.'), + ('SHAPE', 'the object handed back is not the one the task asks for', + 'The KIND of object is wrong irrespective of the values inside it: wrong container or element ' + 'type, wrong nesting, an unconsumed lazy object where a value was expected, a figure handed ' + 'back where the caller reads an axes. NOT for a result whose shape or values came out wrong ' + 'because the computation was wrong -- that belongs to whichever class names the wrong ' + 'assumption in the computation, usually BEHAVIOUR.'), + ('TESTCONTRACT', 'what the caller inspects was never set, or a required side effect never ran', + 'The returned object exists but does not carry what the tests read off it, or a side effect ' + 'the task requires was never performed: labels and titles left unset, a resource not cleaned ' + 'up, a call the task says to make never made.'), + ('DETERMINISM', 'the result is not reproducible or not exactly comparable', + 'Unseeded randomness, reliance on iteration order, missing or wrong sorting, rounding or ' + 'precision other than stated.'), + ('NORMALISATION', 'a text or boundary convention is wrong', + 'Case sensitivity, surrounding whitespace, regex anchoring, separator handling, inclusive ' + 'versus exclusive bounds, off-by-one.'), + ('DEGENERATE', 'a degenerate input crashes instead of being handled', + 'Empty, single-element, all-equal or missing-key / missing-column input.'), + ('EXCEPTION', 'the exception contract is not met', + 'The exception the task specifies is not raised, is raised as a different type, or an ' + 'unrelated exception escapes.'), +] + + +def render_classes() -> str: + """完整定义只给教师；skill-gen 那侧只看短标签，避免它照着举例去写不相干的失败模式。""" + return '\n'.join(f'- {code}: {desc}' for code, _short, desc in FAILURE_CLASSES) + + +def diag_query(problem: str, payload: Dict[str, Any]) -> str: + """题面 + 任务自身声明的硬约定（签名、必需库、返回规格、应抛异常、文档示例）。 + 只用题面信息、不含参考解答 —— 训练时同样拿得到，所以是「可得且非泄漏」的判据依据。""" + try: + doc = payload['doc_struct'] + doc = json.loads(doc) if isinstance(doc, str) else (doc or {}) + except Exception: + doc = {} + lines = [f"- required signature (must be reproduced verbatim):\n" + f"{payload['code_prompt'].strip()}"] + for key, label in (('reqs', 'must use these libraries'), ('returns', 'must return'), + ('raises', 'must raise'), ('params', 'parameters')): + vals = [str(x).strip() for x in (doc.get(key) or []) if str(x).strip()] + if vals: + lines.append(f'- {label}: ' + '; '.join(vals)) + ex = [str(x) for x in (doc.get('examples') or [])] + if ex: + lines.append('- documented example calls:\n ' + '\n '.join(ex)) + return problem + '\n\nHard requirements declared by the task:\n' + '\n'.join(lines) + + +def diag_segment(roll: Dict[str, Any]) -> str: + """★ rubric 路线唯一真正有效的一处：给 judge 的不是「输出全文」，而是**提交的代码 + 单测真实 + 报错**（ 已被 extract_code 切掉）。报错是客观事实且不含参考解答 —— 这正是 code 域 + rubric 有增量（+0.135, p=4e-5）而数学 / BFCL 域没有的原因。""" + return (f"### Submitted code\n```python\n{roll.get('code') or '(no parseable code block)'}\n```" + f"\n\n### Result of running the task's unit tests\n" + f"outcome: {roll.get('kind') or 'unknown'}\n" + f"{roll.get('error') or '(no error output)'}") + + +# =========================================================================== +# skill-gen（pitfall 文体 + rubric 条件化） +# =========================================================================== +# 与 narrative 文体的对照变量是**广度 vs 聚焦**：narrative 穷尽所有失败模式（~300 词），pitfall +# 只挑决定性的那一个（<90 词）。E23 选 pitfall 是因为 executor 已经能看到 actor 的 ， +# narrative 会与思考过程大面积重复，pitfall 让两者的分工是「过程 vs 结论」。 +SKILLGEN_SYSTEM = """\ +You are a skill-generation model. Your block will be fed to a SEPARATE downstream engineer model that must implement the function on its own. The engineer sees the same task description and the same required signature, but NOT your private reasoning or the analysis below — it only sees what is inside .... + +A diagnosis of a failed attempt at THIS task is provided to you: the class of the decisive failure, and the prior knowledge that would have prevented it. That prior is your material — it is a fact that stays true for other tasks too. Your job is to land it at the exact point in THIS task where it bites. + +Then, inside , write under 90 words: +- WARNING: name where the failure strikes — the operation or the hand-off point it happens at — and the wrong assumption behind it, in the terms of the failure class you were given. +- INSTEAD: one or two sentences naming what to guarantee at that exact point, phrased as the general rule rather than as this task's answer. +- End by telling the engineer to deliver one single fenced code block that reproduces the given imports and signature verbatim, with no explanation, no demonstration call and no tests around it. + +Hard rules: +- Write about the failure class you were given. Do NOT substitute a different one and do NOT add a second: a made-up warning sends the engineer after something that is not the problem. +- Do NOT write the solution, and do NOT copy identifiers, column names, file names or literal values out of this task. Name the operation and describe the shape of what it returns instead of writing the call out. +- Self-contained: NEVER reference "the diagnosis", "the analysis", "the review" or "the previous attempt" — the engineer cannot see them, such phrasings cause hallucination. Address the engineer directly. + +Put ONLY the guidance inside . + +Example: + +WARNING: the call doing the real work is accepted but does not behave as its name suggests: in the mode this task nudges you towards, it emits one record per row instead of one whole document, so the reader rejects it. +INSTEAD: confirm what that call emits in the mode you pick, and choose the mode whose output the consumer on the other side expects. +Deliver one fenced code block reproducing the given imports and signature verbatim, with no explanation, demonstration call or tests. + +""" + +SKILLGEN_USER = """\ +Task: +{problem} + +Diagnosis of a failed attempt (for your eyes only; do NOT reference it in the skill): +{rubric} + +Now write the guidance:""" + + +def skillgen_prompt(problem: str, rubric: str) -> Dict[str, Any]: + return {'messages': [{'role': 'system', 'content': SKILLGEN_SYSTEM}, + {'role': 'user', 'content': SKILLGEN_USER.format(problem=problem, + rubric=rubric)}]} diff --git a/cookbook/human/e23_rubric.py b/cookbook/human/e23_rubric.py new file mode 100644 index 000000000..185e408fc --- /dev/null +++ b/cookbook/human/e23_rubric.py @@ -0,0 +1,309 @@ +"""教师 judge（失败分类 + 可迁移先验）与它的磁盘缓存。 + +诊断是 E23 的唯一自变量：分类不出决定性类的题一律丢弃，绝不降级成 query-only，否则自变量被稀释。 +与 v1 的区别是判据从「逐条 PASS/FAIL 的代码 review」换成「单一决定性类 + 报错证据 + 可迁移先验」， +理由见 e23_prompts.FAILURE_CLASSES 上方。 +""" +import collections +import hashlib +import json +import os +import re +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Optional, Tuple + +from twinkle import get_logger +from twinkle_agentic.utils.llm_backup import llm_backup +from twinkle_agentic.verifier import RubricVerifier +from twinkle_agentic.verifier.rubric_verifier import RubricItem, _extract_json_obj, _short_hash + +from e23_prompts import (DIAG_SYSTEM, DIAG_USER, FAILURE_CLASSES, diag_query, diag_segment, + render_classes) + +logger = get_logger() +_HERE = os.path.dirname(os.path.abspath(__file__)) + +RUBRIC_WORKERS = int(os.environ.get('RUBRIC_WORKERS', 16)) +# **轨迹不在缓存键里** —— 任何改变裸解轨迹的开关（executor 的 thinking / executor 模型 / 任务域） +# 都必须体现在文件名上，否则会命中旧口径的诊断。 +RUBRIC_CACHE_PATH = os.environ.get( + 'RUBRIC_CACHE_PATH', os.path.join(_HERE, 'rubric_cache_global_code_execnothink_v3.jsonl')) +RUBRIC_VERSION = 'rubric_code_v3_mechanism' +# 独立根因数 >= 此值就丢题：实测（E23.t1，480 组）只坏 1 处的题救回率 0.504，坏 >=3 处的只有 +# 0.244 —— 一条 90 词内的警告救不回同时坏三处的尝试，训它只是稀释 batch。 +MAX_INDEPENDENT_CAUSES = 3 +# 判定温度。>0 换来的是标签多样性，代价是同一道题重判可能换类 —— 正好可以拿来测标签自一致率 +# （没有真值时这是唯一的准确性代理）；要逐字复现的诊断就设 0。 +DIAG_TEMPERATURE = float(os.environ.get('DIAG_TEMPERATURE', 0.3)) +# 教师 API 配错时每题都会失败 -> 每题都丢 -> 主循环一晚上零更新还不报错。一条都没成功过就早失败。 +API_FAIL_ABORT = int(os.environ.get('API_FAIL_ABORT', 20)) +_CLASS_SHORT = {code: short for code, short, _ in FAILURE_CLASSES} +# 分类表进 llm_backup 的置信度键：判据表一改，学生/教师一致性统计必须从头算。 +_CLASSES_KEY = _short_hash(render_classes()) +# ⭐ 提示词哈希进缓存键。之前只靠 RUBRIC_VERSION 这个手写标签，改了 DIAG_SYSTEM 却忘了改标签就会 +# 全量命中旧诊断、新规则一次都不执行，而且**毫无迹象**（日志里全是缓存命中，看起来一切正常）。 +# 实测踩过：加完「题面没给就弃权」那条规则后重跑，24 道题全命中 20 分钟前的旧判决。哈希兜住这个。 +_PROMPT_KEY = _short_hash(DIAG_SYSTEM + DIAG_USER + render_classes()) +# 先验里出现引号字面量 = 抄了本题的列名/文件名/期望值。默认只统计不丢弃：先看住这个比例，确认 +# 教师是否真的守住了「删掉本题这句话依然成立」，再决定要不要收紧成硬丢。 +PRIOR_REJECT_LITERALS = os.environ.get('PRIOR_REJECT_LITERALS', '0') == '1' +_LITERAL_RE = re.compile(r"'[^']{2,}'|\"[^\"]{2,}\"") + + +def _same_class(a: str, b: str) -> bool: + """student/teacher 是否算一致：只比决定性类。reason/prior 是自由文本，逐字比毫无意义。""" + ca = (_extract_json_obj(a) or {}).get('class') + return bool(ca) and ca == (_extract_json_obj(b) or {}).get('class') + + +class CodeRubricVerifier(RubricVerifier): + """代码域 judge：不走父类的逐条 PASS/FAIL，改成单一决定性类 + 证据 + 先验。 + + 复用父类的调用层而**不动共享代码**：父类的 _parse_diagnosis 只认 index/verdict/reason/fix， + evidence / prior / class 会被静默丢掉。 + + ⭐ 必须经 @llm_backup，不能直接调 _sample_text：本臂没有学生 sampler（build_checker 不传）， + _sample_text 此时按设计恒返回 ''，教师 API 完全由这个装饰器提供。直接调的话每道题都拿到空 + 响应、被判成「不可救」写进缓存，跑一遍就把缓存永久毒化。 + """ + + @llm_backup(key_params=['query', 'rubric_key'], comparator=_same_class) + def _classify_once(self, trajectory, sampling_params, query: str = None, + rubric_key: str = '') -> str: + return self._sample_text(trajectory, sampling_params, self.score_lora_path) + + def classify(self, query: str, segment_text: str) -> Optional[Dict[str, Any]]: + """返回解析出的 JSON；**None 专指「没拿到可解析响应」**（调用层故障，调用方不得缓存）。""" + traj = {'messages': [ + {'role': 'system', 'content': DIAG_SYSTEM}, + {'role': 'user', 'content': DIAG_USER.format( + query=query, rubric=render_classes(), segment=segment_text)}]} + raw = self._classify_once( + trajectory=traj, + sampling_params=self._diagnose_sampling_params(None, temperature=DIAG_TEMPERATURE), + query=query, rubric_key=_CLASSES_KEY) + return _extract_json_obj(raw) + + +def build_checker(): + """没有教师 API 就返回 None（调用方据此报错退出）。 + + fixed_rubric / gate 只为让父类构造合法：本臂从不读 detail.scalar，gate 与 is_hard 都不生效。 + """ + if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') + or os.environ.get('OPENAI_API_KEY')): + return None + return CodeRubricVerifier( + fixed_rubric=[RubricItem(f'{c}: {d}', is_hard=False) for c, _s, d in FAILURE_CLASSES], + gate=False) + + +# 这两类**按定义**就是「对不上调用方要求的某个东西」（前者：调用方读的东西没被设上；后者：该抛的 +# 异常没抛），required_value 为 null 在这里不是合法答案，只能是教师在走捷径。 +_CITATION_REQUIRED = {'TESTCONTRACT', 'EXCEPTION'} + + +def _cites_task(obj: Dict[str, Any], query: str) -> bool: + """教师声称「单测要的这个值题面里给了」时，核验它引的原文**真的**在题面里。 + + ⭐ 这一关不能只靠 DIAG_SYSTEM 的软性要求，两轮实测都被绕过： + 第一轮（只加「题面没给就弃权」的指令）—— 109/222/409 这类「标题/键名只存在于隐藏单测里」的题 + 照收不误，教师写一条 "tests assert exact string equality on titles" 的先验，这句话本身正确且 + 可迁移，于是自己说服自己题可救；可工程师读完仍不知道那字符串是什么，8 个候选必然一起挂。 + 第二轮（加 required_value / required_value_source 引用字段）—— 同样三道题又溜过去了，因为 + 教师把 required_value 填成 null，声称「本题失败与对不上指定值无关」，整关就不适用了。 + 所以对上面那两类**强制**要求引用：拿不出题面原文 = 题面确实没有 = 丢题。 + """ + val = str(obj.get('required_value') or '').strip() + cls = str(obj.get('class') or '').strip().upper() + if not val or val.lower() == 'null': + # 只有「这类失败本来就跟指定值无关」时才放行；两个强制类填 null 一律当没引用处理。 + return cls not in _CITATION_REQUIRED + src = str(obj.get('required_value_source') or '').strip() + if not src or src.lower() == 'null': + return False # 教师自认题面没给 -> 谁也做不出来 + return _squash(src) in _squash(query) # 引文对不上题面 = 编的 + + +def _validate(obj: Any, query: str) -> Optional[Dict[str, Any]]: + """把教师输出收敛成可用诊断；不合格返回 None -> 该题丢出训练集。 + + 弃权（addressable=false）、分类不在表内、拿不出报错证据、先验为空、独立根因 >= 3、单测要的值 + 题面里查无出处 —— 全部按「这题没法用一条不含解法的短警告救」处理。宁可丢题也不喂空洞诊断： + 每 chunk 产出约 40 道错题只需 24 道，丢得起（实测 t1 的 beyond_k 就丢了 328 道）。 + """ + if not isinstance(obj, dict) or not obj.get('addressable'): + return None + if not _cites_task(obj, query): + return None + cls = str(obj.get('class') or '').strip().upper() + prior = str(obj.get('prior') or '').strip() + evidence = str(obj.get('evidence') or '').strip() + if cls not in _CLASS_SHORT or not prior or not evidence: + return None + try: + n_causes = int(obj.get('independent_causes') or 1) + except (TypeError, ValueError): + n_causes = 1 + if n_causes >= MAX_INDEPENDENT_CAUSES: + return None + if PRIOR_REJECT_LITERALS and _LITERAL_RE.search(prior): + return None + secondary = [str(x).strip().upper() for x in (obj.get('secondary') or [])] + # 引用字段一并落缓存（**不进 _format**，不给 skill-gen 看）：上一轮排查时缓存里没有它们， + # 只能靠反推才确认教师是把 required_value 填了 null 溜过去的，白绕一圈。 + return {'class': cls, 'reason': str(obj.get('reason') or '').strip(), 'prior': prior, + 'evidence': evidence, 'n_causes': n_causes, + 'required_value': str(obj.get('required_value') or '').strip(), + 'required_value_source': str(obj.get('required_value_source') or '').strip(), + 'secondary': [s for s in secondary if s in _CLASS_SHORT and s != cls]} + + +# 单测报错里「期望值」的两种常见形态。⭐ 只用来**统计**教师该弃权时有没有弃权，绝不拦截：实测在 +# 352 道题上约两成假阳 —— 命中的字面量可能只是大小写/标点与题面不同（'Performance'、'Random +# Walk'），也可能抓到的是单测**构造输入**用的键而不是要求返回的键（BCB/524 的 'bird'/'fish'）。 +# 归一化比较能消掉前一类，后一类消不掉，所以这个信号只配当监控，判决交给看得见题面的教师。 +_EXPECT_RE = re.compile(r"""!=\s*(['"])(.{2,60}?)\1|KeyError:\s*(['"])(.{1,40}?)\3""") + + +def _squash(s: str) -> str: + return re.sub(r'[^a-z0-9]', '', (s or '').lower()) + + +def _unseen_expected_literal(evidence: str, query: str) -> bool: + """报错要求的字面量在题面里查无此物 = 这题从题面根本做不出来，教师本该弃权。""" + hay = _squash(query) + for m in _EXPECT_RE.finditer(evidence or ''): + lit = m.group(2) or m.group(4) + if lit and _squash(lit) not in hay: + return True + return False + + +def _format(d: Dict[str, Any]) -> str: + """诊断 -> 喂给 skill-gen 的纯文本。 + + ⭐ evidence 故意**不进** skill-gen：它是单测报错原文，断言 diff 里带期望值，是最强的泄漏通道。 + 它留在缓存里只为两件事 —— 逼教师把分类落到客观事实上，以及事后审计。 + """ + lines = [f"DECISIVE FAILURE: {d['class']} — {_CLASS_SHORT[d['class']]}", + f"WHAT WENT WRONG: {d['reason']}", + f"PRIOR THAT WOULD HAVE PREVENTED IT: {d['prior']}"] + if d['secondary']: + lines.append('ALSO OFF (do not write about these): ' + ', '.join(d['secondary'])) + return '\n'.join(lines) + + +def class_metrics(diag_texts: List[str]) -> Dict[str, float]: + """标签退化监控：一组题用到几个决定性类、最大那类占多少。 + + 判据表的全部价值在于分得开：一旦某一类吃掉大半（旧表的兜底项吃了 76% 的题），所有题的诊断 + 就长得一样，skill-gen 失去逐题条件化，等于退回没有 rubric 的对照臂。 + """ + codes = [m.group(1) for m in + (re.match(r'DECISIVE FAILURE:\s*([A-Z]+)', (t or '').split('\n', 1)[0]) + for t in diag_texts) if m] + c = collections.Counter(codes) + return {'signal/rubric_n_classes': float(len(c)), + 'signal/rubric_top_share': max(c.values()) / len(codes) if codes else 0.0} + + +def cache_metrics(delta: collections.Counter) -> Dict[str, float]: + """一步之内 rubric 侧的计数，进 train_log。 + + unseen_literal 是弃权规则的自查通道：教师本该判「题面没给」却收下的题数。要和 dropped 一起 + 读 —— dropped 不涨而 unseen_literal 在涨，就说明 DIAG_SYSTEM 那条规则没被遵守。 + """ + return {'signal/rubric_dropped': float(delta['dropped_unaddressable'] + delta['hit_dropped']), + 'signal/rubric_dropped_not_stated': float(delta['dropped_not_stated']), + 'signal/rubric_api_fail': float(delta['api_fail']), + 'signal/rubric_unseen_literal': float(delta['unseen_literal']), + 'signal/rubric_prior_has_literal': float(delta['prior_has_literal'])} + + +class RubricCache: + """append-only jsonl + 内存索引，键 = md5(RUBRIC_VERSION, 提示词哈希, data_id)，值 = 诊断 JSON。 + + 能跨 run 复用的**唯一依据**是「executor 冻结在 T=0，所以同一道题的裸解轨迹在所有 run 里逐字 + 相同」。轨迹不在键里，见 RUBRIC_CACHE_PATH 上方的换名要求。 + 存 JSON 而不是成品文本：改 _format 的排版不必重拉 API，evidence 也留得住可审计。 + """ + + def __init__(self, path: str = RUBRIC_CACHE_PATH): + self.path = path + self._idx: Dict[str, Any] = {} + self.stats: collections.Counter = collections.Counter() + if os.path.exists(path): + with open(path, encoding='utf-8') as f: + for line in f: + try: + rec = json.loads(line) + self._idx[rec['key']] = rec['value'] + except Exception: + continue + logger.info(f'[rubric] 缓存载入 {len(self._idx)} 条：{path}') + self._fh = open(path, 'a', encoding='utf-8') + + def _put(self, key: str, value: Any) -> None: + self._idx[key] = value + self._fh.write(json.dumps({'key': key, 'value': value}, ensure_ascii=False) + '\n') + self._fh.flush() + + def get_or_diagnose(self, checker, record: Dict[str, Any], roll: Dict[str, Any]) -> str: + """返回诊断文本；教师弃权/不可救返回 '' 并**缓存该判决**，API 失败返回 '' 但**不缓存**。 + + 两者都让调用方丢掉这道题，但只有前者是稳定判决 —— 把一次瞬时抖动写进缓存会永久毒化它。 + """ + key = hashlib.md5(f"{RUBRIC_VERSION}\x00{_PROMPT_KEY}\x00" + f"{record.get('data_id', '')}".encode('utf-8')).hexdigest() + if key in self._idx: + cached = self._idx[key] + if not cached: + self.stats['hit_dropped'] += 1 + return '' + self.stats['hit'] += 1 + return _format(cached) + query = diag_query(record['problem'], record['reference_answer']) + try: + obj = checker.classify(query, diag_segment(roll)) + except Exception as exc: + logger.warning(f'[rubric] classify error: {exc}') + obj = None + if obj is None: + # 拿不到可解析响应 = 调用层故障，**绝不能**当成「教师判不可救」写进缓存：那会把一次 + # 抖动变成永久丢题。教师彻底不通时一条都不会成功，与其空转一夜不如早失败。 + self.stats['api_fail'] += 1 + if self.stats['api_fail'] >= API_FAIL_ABORT and not self.stats['ok']: + raise RuntimeError( + f'[rubric] 连续 {self.stats["api_fail"]} 次拿不到教师响应且无一成功；' + f'检查 LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / LLM_BACKUP_MODEL') + return '' + diag = _validate(obj, query) + if diag is None: + # 缓存空值 = 「这题教师给不出可用分类」，是稳定判决，下次直接跳过不再花 API 钱。 + if isinstance(obj, dict) and obj.get('addressable') and not _cites_task(obj, query): + # 单独计数：这是「题面根本没给」，与「教师主动弃权」是两种不同的丢题原因。 + self.stats['dropped_not_stated'] += 1 + self._put(key, None) + self.stats['dropped_unaddressable'] += 1 + return '' + if _LITERAL_RE.search(diag['prior']): + self.stats['prior_has_literal'] += 1 + if _unseen_expected_literal(diag['evidence'], query): + # 教师收下了一道「答案只存在于隐藏单测里」的题。不改判决（机械检测精度不够），但这个 + # 计数持续偏高就说明 DIAG_SYSTEM 的弃权规则没被遵守。 + self.stats['unseen_literal'] += 1 + self.stats['ok'] += 1 + self.stats[f"class_{diag['class']}"] += 1 + self._put(key, diag) + return _format(diag) + + def diagnose_many(self, checker, pairs: List[Tuple[Dict, Dict]]) -> List[str]: + """并行拉诊断（纯 API 调用，不占 GPU）。pairs = [(record, roll), ...]。""" + if not pairs: + return [] + with ThreadPoolExecutor(max_workers=max(1, min(RUBRIC_WORKERS, len(pairs)))) as ex: + return list(ex.map(lambda rb: self.get_or_diagnose(checker, rb[0], rb[1]), pairs)) + + def close(self): + self._fh.close() diff --git a/cookbook/human/skill_drift_stats.py b/cookbook/human/skill_drift_stats.py new file mode 100644 index 000000000..405ae95af --- /dev/null +++ b/cookbook/human/skill_drift_stats.py @@ -0,0 +1,168 @@ +#!/usr/bin/env python3 +# Copyright (c) ModelScope Contributors. All rights reserved. +"""统计 E23 skill 随训练 step 的「长度」与「词频」漂移，量化「skill 是否越来越空泛」。 + +⚠️ 数据来源与偏差 +------------------ +本脚本读 output.e23/zero_reward_groups.jsonl —— 它**只落整组 reward=0 的题**（8 个候选全错）。 +这不是全部 skill，是一个偏难的子集。所以： + * 「长度随 step 涨」的结论对这个子集成立，但推广到全体 skill 时要记住这一点； + * 词频漂移（guarantee->confirm 等）同理。 +若要无偏统计需改 e23 落全量 skill；当前只有这一份带 step 标签的 skill 文本。 + +口径 +---- +* skill = 候选的 块（extract_skill 已抽好，存在 candidates[].skills）。 +* INSTEAD 段单独抽出来看动词，因为「空泛化」主要发生在这一段（WARNING 段是描述过去的错误）。 +* 词频对比早期（step<=2）vs 晚期（step>=13）两窗，报 log2 比值最大的上升/下降词。 +""" +import collections +import json +import math +import os +import re +import statistics as st +import sys + +HERE = os.path.dirname(os.path.abspath(__file__)) +DEFAULT = os.path.join(HERE, 'output.e23', 'zero_reward_groups.jsonl') + +# 英文停用词（够用即可，不引第三方）。 +STOP = set('the a an of to in on for and or is are be it its this that with as by from at ' + 'you your not no if then when will would should must can may a an s t re ve'.split()) +INSTEAD_RE = re.compile(r'INSTEAD:\s*(.*?)(?:\n(?:Deliver|WARNING)|\Z)', re.S) +WARNING_RE = re.compile(r'WARNING:\s*(.*?)(?:\n(?:INSTEAD|Deliver)|\Z)', re.S) +# 空泛/认知性措辞：不要求 executor 改任何代码，只要求「认同一条命题」。 +VAGUE_RE = re.compile( + r'\b(confirm that|in any|for any|general principle|it may|might|unpredictab\w+|' + r'violat\w+ the principle|be aware|keep in mind|note that|understand that|' + r'is a general|in general|conceptually|principle)\b', re.I) +# 祈使/可执行动词：直接命令 executor 怎么写。 +IMPER_RE = re.compile(r'^(use|set|derive|write|apply|replace|compute|return|add|remove|' + r'ensure|guarantee|call|pass|cast|convert|assign|initialize|import|' + r'check|handle|raise|match)\b', re.I) + + +def words(text): + return [w for w in re.findall(r"[a-zA-Z_][a-zA-Z_']+", (text or '').lower()) + if w not in STOP and len(w) > 2] + + +def load(): + path = sys.argv[1] if len(sys.argv) > 1 else DEFAULT + rows = [json.loads(l) for l in open(path, encoding='utf-8') if l.strip()] + # 每个候选一条记录：(step, skill_text) + out = [] + for r in rows: + s = r['step'] + for c in r['candidates']: + out.append((s, c.get('skills') or '')) + return out, path + + +def instead_seg(sk): + m = INSTEAD_RE.search(sk or '') + return (m.group(1).strip() if m else '') + + +def warning_seg(sk): + m = WARNING_RE.search(sk or '') + return (m.group(1).strip() if m else '') + + +def per_step_table(data): + by = collections.defaultdict(list) + for s, sk in data: + by[s].append(sk) + steps = sorted(by) + print('=' * 88) + print('一、skill 长度与结构 随 step（每行 = 该 step 全部候选的均值；n=候选数）') + print('=' * 88) + print(f"{'step':>4} {'n':>4} {'skill字符':>9} {'skill词数':>9} {'INSTEAD词数':>11} " + f"{'空泛词/条':>9} {'祈使开头%':>9}") + rows_for_trend = [] + for s in steps: + sks = by[s] + chars = st.mean(len(x) for x in sks) + nwords = st.mean(len(words(x)) for x in sks) + ins = [instead_seg(x) for x in sks] + inw = st.mean(len(i.split()) for i in ins) if ins else 0 + vague = st.mean(len(VAGUE_RE.findall(x)) for x in sks) + imper = sum(1 for i in ins if IMPER_RE.match(i)) / len(ins) if ins else 0 + print(f'{s:>4} {len(sks):>4} {chars:>9.1f} {nwords:>9.1f} {inw:>11.1f} ' + f'{vague:>9.2f} {imper:>8.0%}') + rows_for_trend.append((s, chars, vague, imper)) + return by, steps, rows_for_trend + + +def trend(rows_for_trend): + """对 (step, y) 做最小二乘斜率 + t，判断长度/空泛度/祈使率是否真在漂移。""" + print('\n' + '=' * 88) + print('二、趋势显著性（OLS 斜率 / step，|t|>2 才算真漂移）') + print('=' * 88) + xs = [r[0] for r in rows_for_trend] + n = len(xs) + mx = st.mean(xs) + sxx = sum((x - mx) ** 2 for x in xs) + for idx, name in ((1, 'skill 字符数'), (2, '空泛词/条'), (3, 'INSTEAD 祈使开头率')): + ys = [r[idx] for r in rows_for_trend] + my = st.mean(ys) + slope = sum((x - mx) * (y - my) for x, y in zip(xs, ys)) / sxx + resid = [y - (my + slope * (x - mx)) for x, y in zip(xs, ys)] + s2 = sum(e * e for e in resid) / (n - 2) + se = math.sqrt(s2 / sxx) if sxx else 0.0 + t = slope / se if se else 0.0 + tag = '显著' if abs(t) > 2 else '不显著' + print(f' {name:20s} 斜率={slope:+.4f}/step t={t:+.2f} [{tag}] ' + f'首={ys[0]:.3f} 尾={ys[-1]:.3f}') + + +def word_freq_shift(by, steps): + """早期 vs 晚期两窗的词频比。""" + print('\n' + '=' * 88) + print('三、词频漂移：早期(step<=2) vs 晚期(step>=13)') + print('=' * 88) + early_steps = [s for s in steps if s <= 2] + late_steps = [s for s in steps if s >= 13] + + def counts(win): + c = collections.Counter() + ntok = 0 + for s in win: + for sk in by[s]: + w = words(instead_seg(sk)) # 只看 INSTEAD 段 + c.update(w) + ntok += len(w) + return c, ntok + ce, ne = counts(early_steps) + cl, nl = counts(late_steps) + print(f'早期窗 step={early_steps} INSTEAD总词={ne}；晚期窗 step={late_steps} 总词={nl}') + # 频率（每千词），加平滑 + vocab = set(ce) | set(cl) + rows = [] + for w in vocab: + fe = (ce[w] + 0.5) / (ne + 1) * 1000 + fl = (cl[w] + 0.5) / (nl + 1) * 1000 + if ce[w] + cl[w] < 4: # 太稀疏的词不看 + continue + rows.append((math.log2(fl / fe), w, ce[w], cl[w], fe, fl)) + rows.sort(reverse=True) + print(f'\n{"↑晚期变多的词":22s}{"早/千":>8}{"晚/千":>8}{"log2比":>8}') + for lr, w, e, l, fe, fl in rows[:15]: + print(f' {w:20s}{fe:8.1f}{fl:8.1f}{lr:+8.2f}') + print(f'\n{"↓晚期变少的词":22s}{"早/千":>8}{"晚/千":>8}{"log2比":>8}') + for lr, w, e, l, fe, fl in rows[-15:][::-1]: + print(f' {w:20s}{fe:8.1f}{fl:8.1f}{lr:+8.2f}') + + +def main(): + data, path = load() + print(f'[数据] {path}') + print(f'[数据] {len(data)} 个候选 skill（来自整组 reward=0 的题，是偏难子集，非全量）\n') + by, steps, rows_for_trend = per_step_table(data) + trend(rows_for_trend) + word_freq_shift(by, steps) + + +if __name__ == '__main__': + main() From f86f1c1301137f734bf70ae34e2cfe1cfd354a97 Mon Sep 17 00:00:00 2001 From: root Date: Sun, 9 Aug 2026 10:27:49 +0800 Subject: [PATCH 38/60] =?UTF-8?q?feat(human=5Fe18):=20=E6=96=B0=E5=A2=9E?= =?UTF-8?q?=20preflight.py=20=E5=BC=80=E8=B7=91=E5=89=8D=E8=87=AA=E6=A3=80?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit judge 的失败是静默的：pytest 缺失不会让进程崩，只会让每题判 incorrect， 表现为 baseline_accuracy=0 / n_wrong=64/64 / collected 恒 0。B 机因此白跑 10 小时 65 个 chunk。 自检照抄 e18_kodcode.run_tests 的真实结构（solution.py + test_solution.py + _run.py + subprocess/sys.executable），只有走同一条路径，'通过'才等价于 judge 会判通过。同时检查依赖版本、教师 API key、分片参数、resume 种子。 --- cookbook/human_e18/preflight.py | 147 ++++++++++++++++++++++++++++++++ 1 file changed, 147 insertions(+) create mode 100644 cookbook/human_e18/preflight.py diff --git a/cookbook/human_e18/preflight.py b/cookbook/human_e18/preflight.py new file mode 100644 index 000000000..03ea0606d --- /dev/null +++ b/cookbook/human_e18/preflight.py @@ -0,0 +1,147 @@ +# -*- coding: utf-8 -*- +"""开跑前自检：把「跑 10 小时才发现环境不对」提前到 30 秒内暴露。 + +⭐ 为何必须存在：judge 的失败是**静默**的 —— pytest 缺失不会让进程崩， +只会让每道题都判 incorrect，日志上表现为 baseline_accuracy=0 / n_wrong=64/64， +而 collected 一直是 0。B 机曾因此白跑 10 小时 65 个 chunk。 + +⚠️ 关键点：judge 用 `subprocess.run([sys.executable, '_run.py'])` 起子进程， +所以 pytest 必须装在**启动脚本所用的那个解释器**里，不是 `which pytest` 指的那个。 +本脚本用 sys.executable 自查，跟 judge 走同一条路径。 + +用法：python3 preflight.py （用你打算启动采集的同一个 python3） +""" +import os +import subprocess +import sys +import tempfile + +FAIL, WARN = [], [] + + +def ck(name, cond, note='', hard=True): + if not cond: + (FAIL if hard else WARN).append(name) + tag = 'OK ' if cond else ('BAD' if hard else 'WARN') + print(' [%s] %-44s %s' % (tag, name, note)) + + +print('=== 0. 解释器 ===') +print(' sys.executable = %s' % sys.executable) +print(' version = %s' % sys.version.split()[0]) + +print() +print('=== 1. judge 沙箱（最关键，静默失败源）===') +# 完整复刻 judge 的执行路径：子进程 + pytest 跑一个必过的单测 +# ⭐ 完整照抄 e18_kodcode.run_tests 的真实结构，而不是自己另写一个 pytest 调用： +# - solution.py 被测代码（单测靠 `from solution import X` 取） +# - test_solution.py 单测文件（pytest 只收集 test_*.py，名字不能改） +# - _run.py runner，pytest.main 显式指定 test_solution.py +# - subprocess + cwd=tmp + sys.executable +# 只有走同一条路径，「通过」才真的等价于 judge 会判通过。 +_RUNNER = '''import sys, pytest +rc = pytest.main(['-q', '--no-header', '-p', 'no:cacheprovider', + '--tb=short', 'test_solution.py']) +print('__KOD__', rc) +sys.exit(0 if int(rc) == 0 else 1) +''' +_SOLUTION = 'def add(a, b):\n return a + b\n' +_TEST = 'from solution import add\n\n\ndef test_add():\n assert add(1, 2) == 3\n' + +with tempfile.TemporaryDirectory() as td: + for name, body in (('solution.py', _SOLUTION), + ('test_solution.py', _TEST), + ('_run.py', _RUNNER)): + with open(os.path.join(td, name), 'w') as f: + f.write(body) + try: + env = dict(os.environ, PYTHONHASHSEED='0') + env.pop('CUDA_VISIBLE_DEVICES', None) + r = subprocess.run([sys.executable, '_run.py'], cwd=td, env=env, + capture_output=True, text=True, timeout=120) + ok = (r.returncode == 0) + last = (r.stdout or r.stderr).strip().split('\n')[-1][:60] + ck('judge 沙箱可判对一份正确解', ok, 'rc=%d %s' % (r.returncode, last)) + if not ok: + print(' -> 修复：%s -m pip install pytest' % sys.executable) + print(' -> 完整输出：') + for ln in (r.stdout + r.stderr).strip().split('\n')[-6:]: + print(' %s' % ln[:100]) + except Exception as e: + ck('judge 沙箱可判对一份正确解', False, str(e)[:70]) + +print() +print('=== 2. 依赖包 ===') +import importlib.metadata as _md +# A 机实测版本，作为对照基线 +BASE = {'pytest': '9.1.1', 'vllm': '0.23.0', 'torch': '2.11.0+cu130', + 'transformers': '5.14.1', 'modelscope': '1.38.1', 'datasets': '4.8.4', + 'openai': '2.45.0', 'numpy': '2.5.1'} +for pkg, want in BASE.items(): + try: + got = _md.version(pkg) + # 版本不一致只告警：主版本差异才真会出问题，补丁号差异通常无害 + same_major = got.split('.')[0] == want.split('.')[0] + ck(pkg, True, '%s (A机 %s)%s' % (got, want, '' if same_major else ' <- 主版本不同')) + if not same_major: + WARN.append(pkg + '-major') + except Exception: + ck(pkg, False, '未安装 (A机 %s)' % want) + +print() +print('=== 3. 教师 API（rubric 的唯一来源）===') +# 不发真请求，只查变量在不在：缺 key 会让 build_checker 返回 None, +# 后果是 rubric 全空 -> n_rubric_missing == 题数 -> 一条都收不到 +for v in ('LLM_BACKUP_API_KEY', 'LLM_BACKUP_BASE_URL'): + val = os.environ.get(v, '') + ck(v, bool(val), ('已设置(%d字符)' % len(val)) if val else '缺失 -> rubric 会全空') + +print() +print('=== 4. 分片参数 ===') +sn = int(os.environ.get('SHARD_N', 1)) +si = int(os.environ.get('SHARD_ID', 0)) +ck('SHARD_N/SHARD_ID 合法', sn >= 1 and 0 <= si < sn, 'SHARD_N=%d SHARD_ID=%d' % (sn, si)) +ck('多机模式已开启', sn > 1, '单机模式' if sn == 1 else '分片 %d/%d' % (si, sn), hard=False) + +print() +print('=== 5. resume 种子 ===') +_HERE = os.path.dirname(os.path.abspath(__file__)) +od = os.environ.get('OUTPUT_DIR', os.path.join(_HERE, 'output.e18.kod')) +cand = os.path.join(od, 'e18_candidates.jsonl') +if os.path.exists(cand): + import json + import zlib + n = mine = 0 + with open(cand, encoding='utf-8') as f: + for ln in f: + ln = ln.strip() + if not ln: + continue + try: + d = json.loads(ln).get('data_id') + except Exception: + continue + if d: + n += 1 + if zlib.crc32(str(d).encode()) % sn == si: + mine += 1 + ck('种子文件存在', True, '%d 个 id，其中属于本分片 %d 个' % (n, mine)) + ck('种子含本分片的题', mine > 0 or sn == 1, + '本分片会跳过 %d 题' % mine, hard=False) +else: + ck('种子文件存在', False, + '缺 %s -> 会重跑别的机器已做过的题' % cand, hard=False) + +print() +print('=== 6. 数据集缓存 ===') +ck('未设 HF_DATASETS_OFFLINE', os.environ.get('HF_DATASETS_OFFLINE', '') not in ('1', 'true'), + '设了会因本机缓存 config 名带 hash 后缀而加载失败') + +print() +if FAIL: + print('结论: 不可启动 —— 必须先修: %s' % FAIL) +elif WARN: + print('结论: 可启动，但注意: %s' % WARN) +else: + print('结论: 全部通过，可以启动') +sys.exit(1 if FAIL else 0) From f5169ba5fc9fb4c650c152f103298ca6e1e1d07a Mon Sep 17 00:00:00 2001 From: root Date: Sun, 16 Aug 2026 18:31:26 +0800 Subject: [PATCH 39/60] wip --- .../untested/eval_condensed_compressed.sh | 29 - .../untested/eval_condensed_native.sh | 25 - 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 - cookbook/exp/embedding/eval_math_by_level.sh | 59 - .../exp/embedding/train_reflexion_skill.sh | 64 - .../embedding/train_reflexion_skill_rft.sh | 35 - cookbook/exp/embedding/train_skill_v2.sh | 59 - cookbook/exp/legacy/ablation_all.sh | 84 - cookbook/exp/legacy/ablation_direct.sh | 12 - .../exp/legacy/ablation_rag_api_condenser.sh | 17 - .../legacy/ablation_rag_local_condenser.sh | 18 - cookbook/exp/legacy/ablation_rag_raw.sh | 15 - .../legacy/build_reflexion_coldstart_sft.py | 520 ---- .../exp/legacy/build_reflexion_skill_data.py | 1123 -------- .../exp/legacy/build_thinking_rag_index.py | 1159 -------- .../exp/legacy/cold_start/train_cold_start.py | 333 --- cookbook/exp/legacy/compare_math_levels.py | 91 - cookbook/exp/legacy/condenser/dataset.py | 459 ---- .../condenser/make_condenser_dataset.py | 737 ------ .../legacy/condenser/train_condenser_ddp.py | 100 - .../condenser/untested/eval_condensed.py | 382 --- .../exp/legacy/data_pipeline/audit_rubric.py | 167 -- .../legacy/data_pipeline/process_and_save.py | 371 --- cookbook/exp/legacy/dataset_hard.py | 202 -- cookbook/exp/legacy/dataset_index.py | 718 ----- cookbook/exp/legacy/dataset_think.py | 456 ---- cookbook/exp/legacy/eval_dualline_math.py | 689 ----- cookbook/exp/legacy/eval_gpqa_rag.py | 1547 ----------- cookbook/exp/legacy/eval_math_by_level.sh | 59 - cookbook/exp/legacy/eval_rag_recall.py | 187 -- cookbook/exp/legacy/eval_reflexion_skill.py | 762 ------ cookbook/exp/legacy/grpo_baseline.py | 593 ----- cookbook/exp/legacy/grpo_condensed.py | 955 ------- cookbook/exp/legacy/make_condensed_sft.py | 945 ------- cookbook/exp/legacy/make_embedding_dataset.py | 758 ------ cookbook/exp/legacy/reannotate_groundtruth.py | 389 --- cookbook/exp/legacy/rl/grpo.py | 787 ------ cookbook/exp/legacy/rl/rag_hint_grpo.py | 1480 ----------- .../exp/legacy/train_embedding_full_ddp.py | 270 -- cookbook/exp/legacy/train_extract_ddp.py | 119 - cookbook/exp/legacy/train_reflexion_skill.py | 1990 -------------- cookbook/exp/legacy/train_reflexion_skill.sh | 64 - .../exp/legacy/train_reflexion_skill_old.py | 2022 -------------- .../exp/legacy/train_reflexion_skill_old.sh | 57 - .../legacy/train_reflexion_skill_replay.py | 114 - .../exp/legacy/train_reflexion_skill_rft.py | 1569 ----------- .../exp/legacy/train_reflexion_skill_rft.sh | 35 - .../exp/legacy/train_reflexion_skill_seam.py | 2022 -------------- cookbook/exp/legacy/train_skill_v2_ablate.sh | 35 - cookbook/exp/legacy/train_skill_v2_ablate3.sh | 62 - cookbook/exp/skill2lora/analyze_more.py | 172 -- cookbook/exp/skill2lora/code_task.py | 439 ---- .../good_skill_hard_fail/analyze_3way.py | 177 -- .../good_skill_hard_fail/eval_skill_probe.py | 286 -- .../good_skill_hard_fail/leak_decomp.py | 146 -- .../good_skill_hard_fail/reflexion_probe.py | 506 ---- .../good_skill_hard_fail/sample_probe.py | 136 - .../skill_config_probe.py | 421 --- cookbook/exp/skill2lora/logp_corr_probe.py | 430 --- cookbook/exp/skill2lora/rubric_effect.py | 138 - cookbook/exp/skill2lora/run_ablate12.sh | 379 --- .../exp/skill2lora/skill_ablate/__init__.py | 8 - .../exp/skill2lora/skill_ablate/config.py | 318 --- cookbook/exp/skill2lora/skill_ablate/data.py | 106 - .../exp/skill2lora/skill_ablate/data_code.py | 82 - .../skill2lora/skill_ablate/eval_reflexion.py | 208 -- cookbook/exp/skill2lora/skill_ablate/main.py | 327 --- .../exp/skill2lora/skill_ablate/methods.py | 1833 ------------- cookbook/exp/skill2lora/skill_ablate/pool.py | 125 - .../exp/skill2lora/skill_ablate/rollouting.py | 224 -- .../skill2lora/skill_ablate/rubric_cache.py | 128 - .../exp/skill2lora/skill_ablate/trainer.py | 455 ---- cookbook/exp/skill2lora/skill_feature_corr.py | 128 - cookbook/exp/skill2lora/train_skill_v2.py | 2326 ----------------- cookbook/exp/skill2lora/train_skill_v2.sh | 69 - cookbook/exp/skill2lora/watchdog_e14.sh | 40 - cookbook/human/e23_bcb.py | 278 -- cookbook/human/e23_prompts.py | 225 -- cookbook/human/e23_rubric.py | 309 --- cookbook/human/skill_drift_stats.py | 168 -- cookbook/human_e18/README.md | 77 - cookbook/human_e18/e18_collect_kod.py | 738 ------ cookbook/human_e18/e18_kodcode.py | 404 --- cookbook/human_e18/e18_multidiag.py | 245 -- cookbook/human_e18/e18_prompts.py | 264 -- cookbook/human_e18/e18_rejection_sft.py | 775 ------ cookbook/human_e18/e18_select.py | 142 - cookbook/human_e18/e18_sft_kod.py | 543 ---- cookbook/human_e18/e19_logp_select.py | 193 -- cookbook/human_e18/e20_success_skill.py | 311 --- cookbook/human_e18/e21_paired_rubric.py | 214 -- cookbook/human_e18/preflight.py | 147 -- cookbook/human_e18/run_collect_kod.sh | 8 - cookbook/human_e18/run_sft_kod.sh | 18 - cookbook/human_e18/shard_tool.py | 134 - 99 files changed, 39691 deletions(-) delete mode 100755 cookbook/exp/condenser/untested/eval_condensed_compressed.sh delete mode 100755 cookbook/exp/condenser/untested/eval_condensed_native.sh delete mode 100755 cookbook/exp/embedding/ablation_all.sh delete mode 100755 cookbook/exp/embedding/ablation_direct.sh delete mode 100755 cookbook/exp/embedding/ablation_rag_api_condenser.sh delete mode 100755 cookbook/exp/embedding/ablation_rag_local_condenser.sh delete mode 100755 cookbook/exp/embedding/ablation_rag_raw.sh delete mode 100755 cookbook/exp/embedding/eval_math_by_level.sh delete mode 100755 cookbook/exp/embedding/train_reflexion_skill.sh delete mode 100644 cookbook/exp/embedding/train_reflexion_skill_rft.sh delete mode 100644 cookbook/exp/embedding/train_skill_v2.sh delete mode 100755 cookbook/exp/legacy/ablation_all.sh delete mode 100755 cookbook/exp/legacy/ablation_direct.sh delete mode 100755 cookbook/exp/legacy/ablation_rag_api_condenser.sh delete mode 100755 cookbook/exp/legacy/ablation_rag_local_condenser.sh delete mode 100755 cookbook/exp/legacy/ablation_rag_raw.sh delete mode 100644 cookbook/exp/legacy/build_reflexion_coldstart_sft.py delete mode 100644 cookbook/exp/legacy/build_reflexion_skill_data.py delete mode 100644 cookbook/exp/legacy/build_thinking_rag_index.py delete mode 100644 cookbook/exp/legacy/cold_start/train_cold_start.py delete mode 100644 cookbook/exp/legacy/compare_math_levels.py delete mode 100644 cookbook/exp/legacy/condenser/dataset.py delete mode 100644 cookbook/exp/legacy/condenser/make_condenser_dataset.py delete mode 100644 cookbook/exp/legacy/condenser/train_condenser_ddp.py delete mode 100644 cookbook/exp/legacy/condenser/untested/eval_condensed.py delete mode 100644 cookbook/exp/legacy/data_pipeline/audit_rubric.py delete mode 100644 cookbook/exp/legacy/data_pipeline/process_and_save.py delete mode 100644 cookbook/exp/legacy/dataset_hard.py delete mode 100644 cookbook/exp/legacy/dataset_index.py delete mode 100644 cookbook/exp/legacy/dataset_think.py delete mode 100644 cookbook/exp/legacy/eval_dualline_math.py delete mode 100644 cookbook/exp/legacy/eval_gpqa_rag.py delete mode 100755 cookbook/exp/legacy/eval_math_by_level.sh delete mode 100644 cookbook/exp/legacy/eval_rag_recall.py delete mode 100644 cookbook/exp/legacy/eval_reflexion_skill.py delete mode 100644 cookbook/exp/legacy/grpo_baseline.py delete mode 100644 cookbook/exp/legacy/grpo_condensed.py delete mode 100644 cookbook/exp/legacy/make_condensed_sft.py delete mode 100644 cookbook/exp/legacy/make_embedding_dataset.py delete mode 100644 cookbook/exp/legacy/reannotate_groundtruth.py delete mode 100644 cookbook/exp/legacy/rl/grpo.py delete mode 100644 cookbook/exp/legacy/rl/rag_hint_grpo.py delete mode 100644 cookbook/exp/legacy/train_embedding_full_ddp.py delete mode 100644 cookbook/exp/legacy/train_extract_ddp.py delete mode 100644 cookbook/exp/legacy/train_reflexion_skill.py delete mode 100755 cookbook/exp/legacy/train_reflexion_skill.sh delete mode 100644 cookbook/exp/legacy/train_reflexion_skill_old.py delete mode 100755 cookbook/exp/legacy/train_reflexion_skill_old.sh delete mode 100644 cookbook/exp/legacy/train_reflexion_skill_replay.py delete mode 100644 cookbook/exp/legacy/train_reflexion_skill_rft.py delete mode 100644 cookbook/exp/legacy/train_reflexion_skill_rft.sh delete mode 100644 cookbook/exp/legacy/train_reflexion_skill_seam.py delete mode 100644 cookbook/exp/legacy/train_skill_v2_ablate.sh delete mode 100755 cookbook/exp/legacy/train_skill_v2_ablate3.sh delete mode 100644 cookbook/exp/skill2lora/analyze_more.py delete mode 100644 cookbook/exp/skill2lora/code_task.py delete mode 100644 cookbook/exp/skill2lora/good_skill_hard_fail/analyze_3way.py delete mode 100644 cookbook/exp/skill2lora/good_skill_hard_fail/eval_skill_probe.py delete mode 100644 cookbook/exp/skill2lora/good_skill_hard_fail/leak_decomp.py delete mode 100644 cookbook/exp/skill2lora/good_skill_hard_fail/reflexion_probe.py delete mode 100644 cookbook/exp/skill2lora/good_skill_hard_fail/sample_probe.py delete mode 100644 cookbook/exp/skill2lora/good_skill_hard_fail/skill_config_probe.py delete mode 100644 cookbook/exp/skill2lora/logp_corr_probe.py delete mode 100644 cookbook/exp/skill2lora/rubric_effect.py delete mode 100644 cookbook/exp/skill2lora/run_ablate12.sh delete mode 100644 cookbook/exp/skill2lora/skill_ablate/__init__.py delete mode 100644 cookbook/exp/skill2lora/skill_ablate/config.py delete mode 100644 cookbook/exp/skill2lora/skill_ablate/data.py delete mode 100644 cookbook/exp/skill2lora/skill_ablate/data_code.py delete mode 100644 cookbook/exp/skill2lora/skill_ablate/eval_reflexion.py delete mode 100644 cookbook/exp/skill2lora/skill_ablate/main.py delete mode 100644 cookbook/exp/skill2lora/skill_ablate/methods.py delete mode 100644 cookbook/exp/skill2lora/skill_ablate/pool.py delete mode 100644 cookbook/exp/skill2lora/skill_ablate/rollouting.py delete mode 100644 cookbook/exp/skill2lora/skill_ablate/rubric_cache.py delete mode 100644 cookbook/exp/skill2lora/skill_ablate/trainer.py delete mode 100644 cookbook/exp/skill2lora/skill_feature_corr.py delete mode 100644 cookbook/exp/skill2lora/train_skill_v2.py delete mode 100644 cookbook/exp/skill2lora/train_skill_v2.sh delete mode 100644 cookbook/exp/skill2lora/watchdog_e14.sh delete mode 100644 cookbook/human/e23_bcb.py delete mode 100644 cookbook/human/e23_prompts.py delete mode 100644 cookbook/human/e23_rubric.py delete mode 100644 cookbook/human/skill_drift_stats.py delete mode 100644 cookbook/human_e18/README.md delete mode 100644 cookbook/human_e18/e18_collect_kod.py delete mode 100644 cookbook/human_e18/e18_kodcode.py delete mode 100644 cookbook/human_e18/e18_multidiag.py delete mode 100644 cookbook/human_e18/e18_prompts.py delete mode 100644 cookbook/human_e18/e18_rejection_sft.py delete mode 100644 cookbook/human_e18/e18_select.py delete mode 100644 cookbook/human_e18/e18_sft_kod.py delete mode 100644 cookbook/human_e18/e19_logp_select.py delete mode 100644 cookbook/human_e18/e20_success_skill.py delete mode 100644 cookbook/human_e18/e21_paired_rubric.py delete mode 100644 cookbook/human_e18/preflight.py delete mode 100755 cookbook/human_e18/run_collect_kod.sh delete mode 100755 cookbook/human_e18/run_sft_kod.sh delete mode 100644 cookbook/human_e18/shard_tool.py diff --git a/cookbook/exp/condenser/untested/eval_condensed_compressed.sh b/cookbook/exp/condenser/untested/eval_condensed_compressed.sh deleted file mode 100755 index ce814ae14..000000000 --- a/cookbook/exp/condenser/untested/eval_condensed_compressed.sh +++ /dev/null @@ -1,29 +0,0 @@ -#!/bin/sh -# Compressed run: chunk → condense via Qwen3.5-4B-Condenser LoRA → extract_condensed tool loop. -# Identical --dataset / --limit / --model_id as eval_condensed_native.sh for an A/B comparison. -set -eu - -DATASET="/mnt/data/yzhao/datasets/musique_ans_v1.0_dev.jsonl" -MODEL_ID="ms://Qwen/Qwen3.5-4B" -CONDENSER_LORA="ms://twinkle-kit/Qwen3.5-4B-Condenser" -LIMIT="500" -NUM_GPUS="4" -OUT_DIR="eval_out" - -CUDA_VISIBLE_DEVICES=0,1,2,3 \ -python cookbook/exp/eval_condensed.py \ - --mode condensed \ - --dataset_format musique \ - --dataset "${DATASET}" \ - --model_id "${MODEL_ID}" \ - --condenser_lora "${CONDENSER_LORA}" \ - --limit "${LIMIT}" \ - --num_gpus "${NUM_GPUS}" \ - --batch_size 8 \ - --max_model_len 32768 \ - --max_new_tokens 2048 \ - --max_turns 4 \ - --max_trajectory_tokens 8192 \ - --chunk_size 1024 \ - --temperature 0.0 \ - --out_dir "${OUT_DIR}" diff --git a/cookbook/exp/condenser/untested/eval_condensed_native.sh b/cookbook/exp/condenser/untested/eval_condensed_native.sh deleted file mode 100755 index 3a84cff26..000000000 --- a/cookbook/exp/condenser/untested/eval_condensed_native.sh +++ /dev/null @@ -1,25 +0,0 @@ -#!/bin/sh -# Native baseline: full original context, single-turn QA, no compression, no tools. -# Compare against eval_condensed_compressed.sh on identical --dataset / --limit / --model_id. -set -eu - -DATASET="/mnt/data/yzhao/datasets/musique_ans_v1.0_dev.jsonl" -MODEL_ID="ms://Qwen/Qwen3.5-4B" -LIMIT="500" -NUM_GPUS="4" -OUT_DIR="eval_out" - -CUDA_VISIBLE_DEVICES=0,1,2,3 \ -python cookbook/exp/eval_condensed.py \ - --mode native \ - --dataset_format musique \ - --dataset "${DATASET}" \ - --model_id "${MODEL_ID}" \ - --limit "${LIMIT}" \ - --num_gpus "${NUM_GPUS}" \ - --batch_size 8 \ - --max_model_len 32768 \ - --max_new_tokens 2048 \ - --max_trajectory_tokens 8192 \ - --temperature 0.0 \ - --out_dir "${OUT_DIR}" diff --git a/cookbook/exp/embedding/ablation_all.sh b/cookbook/exp/embedding/ablation_all.sh deleted file mode 100755 index 2be4c7fc7..000000000 --- a/cookbook/exp/embedding/ablation_all.sh +++ /dev/null @@ -1,84 +0,0 @@ -#!/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 deleted file mode 100755 index f5e5d1456..000000000 --- a/cookbook/exp/embedding/ablation_direct.sh +++ /dev/null @@ -1,12 +0,0 @@ -#!/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 deleted file mode 100755 index fb47718e6..000000000 --- a/cookbook/exp/embedding/ablation_rag_api_condenser.sh +++ /dev/null @@ -1,17 +0,0 @@ -#!/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 deleted file mode 100755 index defa863cd..000000000 --- a/cookbook/exp/embedding/ablation_rag_local_condenser.sh +++ /dev/null @@ -1,18 +0,0 @@ -#!/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 deleted file mode 100755 index 86e35bd1e..000000000 --- a/cookbook/exp/embedding/ablation_rag_raw.sh +++ /dev/null @@ -1,15 +0,0 @@ -#!/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/eval_math_by_level.sh b/cookbook/exp/embedding/eval_math_by_level.sh deleted file mode 100755 index d4f8bbdb3..000000000 --- a/cookbook/exp/embedding/eval_math_by_level.sh +++ /dev/null @@ -1,59 +0,0 @@ -#!/bin/bash -# MATH (Hendrycks) difficulty-stratified evaluation. -# -# Goal: measure how the (raw) RAG gain over direct varies with problem -# difficulty (Level 1-5). Runs raw RAG first (retrieve -> qwen3.7-max condense -# -> inject, no hint filtering; it writes the problem-id file), then direct on -# the *same* problems for a paired comparison. -# -# Usage: -# COMPRESS_API_KEY=sk-xxx bash cookbook/exp/embedding/eval_math_by_level.sh -# -# Env knobs: -# PER_LEVEL problems per difficulty level (default 100 -> 500 total) -# SEED stratified-sampling seed (default 100; must match across runs) -# DB_PATH LanceDB retrieval index -# SIM / TOPK retrieval threshold / top-k - -set -euo pipefail - -export COMPRESS_API_KEY="${COMPRESS_API_KEY:?Set COMPRESS_API_KEY}" - -SCRIPT="cookbook/exp/embedding/eval_gpqa_rag.py" -PER_LEVEL="${PER_LEVEL:-100}" -SEED="${SEED:-100}" -SIM="${SIM:-0.75}" -TOPK="${TOPK:-1}" -OUTDIR="./output/thinking_rag" -DB_PATH="${DB_PATH:-./output.oldemb/thinking_rag/lance.db}" - -mkdir -p "$OUTDIR" - -echo "============================================================" -echo " MATH by level: raw RAG (qwen3.7-max condenser, no hint)" -echo " per_level=$PER_LEVEL seed=$SEED" -echo "============================================================" -python "$SCRIPT" \ - --dataset math --math-split test \ - --mode rag \ - --per-level "$PER_LEVEL" --seed "$SEED" \ - --db-path "$DB_PATH" \ - --sim-threshold "$SIM" --top-k "$TOPK" \ - --condense \ - --output "$OUTDIR/math_rag_results.jsonl" - -echo "" -echo "============================================================" -echo " MATH by level: Direct (same problems as raw RAG)" -echo "============================================================" -# Direct reads math_rag_problem_ids.json (written above) to match the subset. -python "$SCRIPT" \ - --dataset math --math-split test \ - --mode direct \ - --per-level "$PER_LEVEL" --seed "$SEED" \ - --output "$OUTDIR/math_direct_results.jsonl" - -echo "" -echo "============================================================" -echo " Done. Compare with: python cookbook/exp/embedding/compare_math_levels.py" -echo "============================================================" diff --git a/cookbook/exp/embedding/train_reflexion_skill.sh b/cookbook/exp/embedding/train_reflexion_skill.sh deleted file mode 100755 index e732a6a3f..000000000 --- a/cookbook/exp/embedding/train_reflexion_skill.sh +++ /dev/null @@ -1,64 +0,0 @@ -#!/bin/bash -# Online GRPO RFT for the reflexion skill generator (unified, self-contained, cached). -# GPUs: 8 — default high-memory layout uses rank 0 for actor training, rank 1 for -# the frozen ref model, ranks 2-3 for skill sampler (synced), and ranks 4-7 for -# base sampler (frozen). Override TRAIN_GPUS / REF_GPUS / SKILL_SAMPLER_GPUS / -# BASE_SAMPLER_GPUS for other layouts. Per chunk: base greedy -# solve -> rubric process-check (view A) -> -# skill-gen (thinking ON, N candidates) -> deterministic leak filter -> with-skill greedy -# pass -> group-relative advantage -> ONE on-policy GRPO step -> sync weights. -# -# Baseline rollouts + rubric diagnoses are disk-cached (output-dir/cache/*.jsonl), so a -# restart skips re-sampling them; skill-gen is on-policy and never cached. The next chunk's -# baseline is prefetched on a background thread (overlaps skill-gen; base sampler is frozen). -# -# The view-A rubric process-check uses the backup teacher API (set LLM_BACKUP_*). Without -# it the run still works: view A degrades to query-only and the leak filter stays -# deterministic (no teacher needed). - -set -euo pipefail - -export GEN_MODEL_ID=${GEN_MODEL_ID:-Qwen/Qwen3-4B} -export GEN_GPU_MEM=${GEN_GPU_MEM:-0.8} -# Datasets are pulled from ModelScope via twinkle.Dataset (ms://AI-MO/aops or -# ms://modelscope/competition_math); override AOPS_DATASET_ID / MATH_DATASET_ID to change. -# Teacher API for the view-A rubric process-check (optional; leak filter is deterministic). -export LLM_BACKUP_API_KEY=${LLM_BACKUP_API_KEY:-} -export LLM_BACKUP_BASE_URL=${LLM_BACKUP_BASE_URL:-https://dashscope.aliyuncs.com/compatible-mode/v1} -export LLM_BACKUP_MODEL=${LLM_BACKUP_MODEL:-qwen3.7-max} -export EXCLUDE_DATA_IDS=${EXCLUDE_DATA_IDS:-./output/reflexion_coldstart_sft/coldstart_sft.jsonl} - -python cookbook/exp/embedding/train_reflexion_skill.py \ - --dataset aops \ - --n 10000 \ - --numeric-only \ - --chunk-size 64 \ - --n-skills 8 \ - --viewa-frac-start 1.0 \ - --viewa-frac-end 0.1 \ - --viewa-warmup-chunks 20 \ - --viewa-decay-chunks 40 \ - --skill-retries 2 \ - --balance \ - --balance-success-frac 0.2 \ - --balance-loop-frac 0.5 \ - --balance-max-draws-mult 8 \ - --max-tokens 8192 \ - --skill-max-tokens 4096 \ - --max-model-len 16384 \ - --eval-size 128 \ - --exclude-data-ids "${EXCLUDE_DATA_IDS}" \ - --eval-every 5 \ - --sft-batch-size 8 \ - --ppo-mini-batch-size 0 \ - --grpo-epsilon 0.2 \ - --kl-beta 0.001 \ - --format-in-reward \ - --lr 1e-6 \ - --max-train-rounds 1500 \ - --save-rounds 200 \ - --trend-every 10 \ - --prefetch-baseline \ - --output-dir ./output/reflexion_skill_curriculum \ - --swanlab-project twinkle \ - --swanlab-exp reflexion_skill_curriculum diff --git a/cookbook/exp/embedding/train_reflexion_skill_rft.sh b/cookbook/exp/embedding/train_reflexion_skill_rft.sh deleted file mode 100644 index b5e0f1d82..000000000 --- a/cookbook/exp/embedding/train_reflexion_skill_rft.sh +++ /dev/null @@ -1,35 +0,0 @@ -#!/bin/bash -# RFT cold-start for the reflexion skill generator (see reflexion.md §6). -# GPUs: 8 — ranks 0-3 train (skill model, FSDP2), 4-5 skill sampler, 6-7 base sampler. -# Leak filtering uses the backup teacher API (no local judge): set LLM_BACKUP_*. - -set -euo pipefail - -export GEN_MODEL_ID=${GEN_MODEL_ID:-Qwen/Qwen3-4B} -# Local MATH copy (modelscope download cache). Override MATH_DATA_DIR if the -# cache hash dir changes or the data lives elsewhere. -export MATH_DATA_DIR=${MATH_DATA_DIR:-/mnt/workspace/.cache/modelscope/hub/datasets/downloads/extracted/0744cd2d347a7e8f85f7087d950b2ed38b626a5c808c5399e2d8a0923d42d013/MATH} -export LLM_BACKUP_API_KEY=${LLM_BACKUP_API_KEY:?set LLM_BACKUP_API_KEY for the leak judge} -export LLM_BACKUP_BASE_URL=${LLM_BACKUP_BASE_URL:-https://dashscope.aliyuncs.com/compatible-mode/v1} -export LLM_BACKUP_MODEL=${LLM_BACKUP_MODEL:-qwen3.7-max} - -python cookbook/exp/embedding/train_reflexion_skill_rft.py \ - --dataset aops \ - --n 5000 \ - --chunk-size 16 \ - --n-skills 8 \ - --view-b-frac 0.5 \ - --skill-retries 2 \ - --balance \ - --balance-success-frac 0.4 \ - --balance-loop-frac 0.5 \ - --balance-max-draws-mult 8 \ - --max-tokens 25000 \ - --max-model-len 30000 \ - --sft-batch-size 8 \ - --grpo-epsilon 0.2 \ - --lr 6e-6 \ - --max-train-rounds 1500 \ - --save-rounds 25 \ - --trend-every 10 \ - --output-dir ./output/reflexion_skill_rft diff --git a/cookbook/exp/embedding/train_skill_v2.sh b/cookbook/exp/embedding/train_skill_v2.sh deleted file mode 100644 index f6c2d91ea..000000000 --- a/cookbook/exp/embedding/train_skill_v2.sh +++ /dev/null @@ -1,59 +0,0 @@ -#!/bin/bash -# train_skill_v2.sh — 简化 GRPO + buffer distill 训练启动脚本 -# 用法: bash cookbook/exp/embedding/train_skill_v2.sh -# -# 环境变量: -# LLM_BACKUP_API_KEY - rubric 诊断用的教师 API key（必须，否则 buffer B 蒸馏不可用） -# LLM_BACKUP_BASE_URL - 教师 API base URL -# LLM_BACKUP_MODEL - 教师模型 ID -# GEN_MODEL_ID - 训练 skill 模型 ID（默认 Qwen/Qwen3-4B） -# TRAIN_GPUS / REF_GPUS / SKILL_SAMPLER_GPUS / BASE_SAMPLER_GPUS — GPU 分配 - -set -euo pipefail -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" - -# 默认输出目录 -OUTPUT_DIR="${OUTPUT_DIR:-./output/skill_v2}" - -# 去重/排斥数据（冷启动 SFT 数据避免重叠） -EXCLUDE="${EXCLUDE_DATA_IDS:-}" - -# 提前建目录：tee 需在 python 建目录前就能打开日志文件 -mkdir -p "${OUTPUT_DIR}" - -python3 "${SCRIPT_DIR}/train_skill_v2.py" \ - --dataset aops \ - --n 20000 \ - --numeric-only \ - --eval-size 200 \ - --eval-every 10 \ - --chunk-size 32 \ - --n-skills 8 \ - --skill-retries 2 \ - --skill-gen-temperature 1.0 \ - --skill-gen-top-p 1.0 \ - --skill-gen-top-k -1 \ - --max-model-len 16384 \ - --max-tokens 8192 \ - --skill-max-tokens 4096 \ - --len-budget 600 \ - --distill-trigger 150 \ - --distill-batch 64 \ - --sft-trigger 100 \ - --passatk-k 8 \ - --passatk-m 2 \ - --sft-weight 1.0 \ - --rubric-workers 16 \ - --sft-batch-size 8 \ - --ppo-mini-batch-size 0 \ - --grpo-epsilon 0.2 \ - --adv-clip 3.0 \ - --kl-beta 0.001 \ - --lr 1e-6 \ - --max-train-rounds 1500 \ - --save-rounds 200 \ - --output-dir "${OUTPUT_DIR}" \ - --swanlab-project twinkle \ - --swanlab-exp "skill_v2_$(date +%Y%m%d_%H%M%S)" \ - ${EXCLUDE:+--exclude-data-ids "${EXCLUDE}"} \ - "$@" 2>&1 | tee "${OUTPUT_DIR}/run.log" diff --git a/cookbook/exp/legacy/ablation_all.sh b/cookbook/exp/legacy/ablation_all.sh deleted file mode 100755 index 2be4c7fc7..000000000 --- a/cookbook/exp/legacy/ablation_all.sh +++ /dev/null @@ -1,84 +0,0 @@ -#!/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/legacy/ablation_direct.sh b/cookbook/exp/legacy/ablation_direct.sh deleted file mode 100755 index f5e5d1456..000000000 --- a/cookbook/exp/legacy/ablation_direct.sh +++ /dev/null @@ -1,12 +0,0 @@ -#!/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/legacy/ablation_rag_api_condenser.sh b/cookbook/exp/legacy/ablation_rag_api_condenser.sh deleted file mode 100755 index fb47718e6..000000000 --- a/cookbook/exp/legacy/ablation_rag_api_condenser.sh +++ /dev/null @@ -1,17 +0,0 @@ -#!/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/legacy/ablation_rag_local_condenser.sh b/cookbook/exp/legacy/ablation_rag_local_condenser.sh deleted file mode 100755 index defa863cd..000000000 --- a/cookbook/exp/legacy/ablation_rag_local_condenser.sh +++ /dev/null @@ -1,18 +0,0 @@ -#!/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/legacy/ablation_rag_raw.sh b/cookbook/exp/legacy/ablation_rag_raw.sh deleted file mode 100755 index 86e35bd1e..000000000 --- a/cookbook/exp/legacy/ablation_rag_raw.sh +++ /dev/null @@ -1,15 +0,0 @@ -#!/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/legacy/build_reflexion_coldstart_sft.py b/cookbook/exp/legacy/build_reflexion_coldstart_sft.py deleted file mode 100644 index a7748b859..000000000 --- a/cookbook/exp/legacy/build_reflexion_coldstart_sft.py +++ /dev/null @@ -1,520 +0,0 @@ -"""Build a cold-start SFT corpus for reflexion skill generation on AOPS. - -Pipeline: - AOPS problems -> frozen base greedy attempt -> strategy-level rubric API diagnosis - -> answer-free API skill target -> query-only SFT examples. - -This is intentionally offline: no GRPO, no actor training, and no skill-model rollout. -The API is treated as an external teacher, so both diagnosis and generated skill targets -are filtered if they reveal the target final answer. - -Example: - LLM_BACKUP_API_KEY=... LLM_BACKUP_BASE_URL=... LLM_BACKUP_MODEL=... \ - python cookbook/exp/embedding/build_reflexion_coldstart_sft.py \ - --dataset aops --n 10000 --output-dir ./output/reflexion_coldstart_sft --overwrite -""" -import argparse -import json -import os -import sys -import time -import urllib.error -import urllib.request -from concurrent.futures import ThreadPoolExecutor -from typing import Any, Dict, List, Optional, Tuple - -_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../..')) -if _REPO_ROOT not in sys.path: - sys.path.insert(0, _REPO_ROOT) - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.data_format import pack_user_data -from twinkle.sampler import vLLMSampler - -from cookbook.exp.embedding.train_reflexion_skill import ( - MODEL_ID, - GPU_MEM, - SamplingParams, - DiskCache, - _MATH_RUBRIC, - _RUBRIC_VERSION, - _answer_leaked, - _clean_text, - _empty_roll, - _format_diagnosis, - _numeric_value, - _parse_seq, - _run_samples, - _skillgen_messages, - _load_excluded_records, - build_direct_prompt, - build_skill_solve_prompt, - build_rubric_checker, - extract_boxed, - load_problems, -) - -logger = get_logger() - -COLDSTART_SYSTEM = """\ -You are writing cold-start training targets for a math skill generator. You are given a -competition problem and an answer-free process diagnosis of a previous attempt. - -Write concise, reusable guidance that a query-only solver could use before solving this -problem or similar problems. Focus on route choice, structural observations, constraints, -validity checks, and length-control habits. - -Output exactly one XML-style block: - -Your reusable guidance here. - - -Rules: -- Do not mention the diagnosis, rubric, previous attempt, or API. -- Do not reveal the final answer, a corrected value/expression, an option label, or a - step-by-step solution. -- It is okay to name methods, checks, pitfalls, and local strategy directions. -- Keep it short and useful: 3-6 compact sentences or bullets. -""" - -COLDSTART_USER = """\ -Problem: -{problem} - -Answer-free process diagnosis: -{diagnosis} - -Now write the reusable skill guidance. -""" - -_SPECIAL_TOKEN_NOTE = 'process diagnosis leaked target answer' - - -def _api_config() -> Tuple[str, str, str]: - api_key = os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY') - base_url = os.environ.get('LLM_BACKUP_BASE_URL') or os.environ.get('OPENAI_BASE_URL') or 'https://api.openai.com/v1' - model = os.environ.get('LLM_BACKUP_MODEL') or os.environ.get('OPENAI_MODEL') or 'gpt-4o-mini' - if not api_key: - raise RuntimeError('Set LLM_BACKUP_API_KEY or OPENAI_API_KEY for cold-start API generation.') - return api_key, base_url.rstrip('/'), model - - -def _chat_complete(messages: List[Dict[str, str]], max_tokens: int, temperature: float, - retries: int = 3, timeout: int = 120) -> str: - api_key, base_url, model = _api_config() - url = f'{base_url}/chat/completions' - payload = { - 'model': model, - 'messages': messages, - 'temperature': temperature, - 'max_tokens': max_tokens, - } - data = json.dumps(payload).encode('utf-8') - headers = {'Content-Type': 'application/json', 'Authorization': f'Bearer {api_key}'} - last_err = None - for attempt in range(max(1, retries)): - req = urllib.request.Request(url, data=data, headers=headers, method='POST') - try: - with urllib.request.urlopen(req, timeout=timeout) as resp: - obj = json.loads(resp.read().decode('utf-8')) - return obj['choices'][0]['message']['content'] - except (urllib.error.HTTPError, urllib.error.URLError, TimeoutError, KeyError, json.JSONDecodeError) as exc: - last_err = exc - if attempt + 1 < max(1, retries): - time.sleep(min(8.0, 1.0 * (2 ** attempt))) - continue - raise RuntimeError(f'chat completion failed after {retries} attempts: {last_err}') - - -def _extract_skill_block(text: str) -> Optional[str]: - low = (text or '').lower() - end_think = low.rfind('') - answer = text[end_think + len(''):] if end_think >= 0 else (text or '') - low = answer.lower() - s = low.rfind('') - if s < 0: - return None - inner = s + len('') - e = low.find('', inner) - if e < 0: - return None - block = answer[inner:e].strip() - return block or None - - -def _skill_response(block: str) -> str: - return f'\n{block.strip()}\n' - - -def _baseline_rollout(base_sampler, problems: List[Dict[str, Any]], base_dp: int, - args: argparse.Namespace, cache: DiskCache) -> int: - todo = [r for r in problems if DiskCache.key_for(r['problem']) not in cache] - if todo: - outs = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], - 1, args.max_tokens, base_dp, temperature=0.0) - for r, seqs in zip(todo, outs): - roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - cache.put(DiskCache.key_for(r['problem']), roll) - for r in problems: - roll = cache.get(DiskCache.key_for(r['problem'])) - r['_init'] = [roll] - r['_baseline_pass'] = 1.0 if roll.get('correct') else 0.0 - r['_failed'] = not roll.get('correct') - return len(todo) - - -def _diagnose_one(checker, r: Dict[str, Any], args: argparse.Namespace) -> str: - init = r['_init'][0] - seg_text = init.get('text', '') - if init.get('stop_reason') == 'length' or not init.get('terminated'): - seg_text += ('\n\n[Process note: this attempt was cut off at the token budget ' - 'and never produced a final \\boxed{} answer.]') - seg = {'messages': [{'role': 'user', 'content': r['problem']}, - {'role': 'assistant', 'content': seg_text}]} - attempts = max(1, args.rubric_retries + 1) - for attempt in range(attempts): - try: - return _format_diagnosis(checker.diagnose(seg, query=r['problem'])) - except Exception as exc: - if attempt + 1 < attempts: - logger.warning(f'[rubric] diagnose error: {exc}; retry {attempt + 1}/{args.rubric_retries}') - time.sleep(min(4.0, 0.5 * (2 ** attempt))) - else: - logger.warning(f'[rubric] diagnose failed: {exc}') - return '' - - -def _diagnose_batch(checker, rows: List[Dict[str, Any]], args: argparse.Namespace, - cache: DiskCache) -> int: - pending = [] - for r in rows: - init = r['_init'][0] - term = 'L' if (init.get('stop_reason') == 'length' or not init.get('terminated')) else 'T' - key = DiskCache.key_for(r['problem'], init.get('text', ''), _RUBRIC_VERSION, term) - if key in cache: - r['_rubric_diag'] = cache.get(key) - else: - pending.append((r, key)) - if not pending: - return 0 - - def run(item): - r, key = item - diag = _diagnose_one(checker, r, args) - return r, key, diag - - workers = max(1, min(args.rubric_workers, len(pending))) - fresh = 0 - with ThreadPoolExecutor(max_workers=workers) as ex: - for r, key, diag in ex.map(run, pending): - r['_rubric_diag'] = diag or '' - if diag: - cache.put(key, diag) - fresh += 1 - return fresh - - -def _target_key(problem: str, diagnosis: str, sample_idx: int) -> str: - return DiskCache.key_for('coldstart_skill_v2', str(sample_idx), problem, diagnosis) - - -def _generate_skill_targets(r: Dict[str, Any], args: argparse.Namespace, - cache: DiskCache) -> List[Dict[str, Any]]: - out = [] - messages = [ - {'role': 'system', 'content': COLDSTART_SYSTEM}, - {'role': 'user', 'content': COLDSTART_USER.format(problem=r['problem'], diagnosis=r.get('_rubric_diag', ''))}, - ] - for sample_idx in range(max(1, int(args.api_samples))): - key = _target_key(r['problem'], r.get('_rubric_diag', ''), sample_idx) - if key in cache: - resp = cache.get(key) - else: - resp = _chat_complete(messages, max_tokens=args.api_max_tokens, - temperature=args.api_temperature, retries=args.api_retries, - timeout=args.api_timeout) - cache.put(key, resp) - block = _extract_skill_block(resp) or '' - leaked = _answer_leaked(resp + '\n' + block, r['reference_answer']) - out.append({'sample_idx': sample_idx, 'raw_response': resp, - 'skills': block, 'skill_leak': leaked}) - return out - - -def _sft_messages(problem: str, response: str) -> List[Dict[str, str]]: - msgs = _skillgen_messages(problem, 'B', '') - return msgs + [{'role': 'assistant', 'content': response}] - - -def _init_base_sampler(args: argparse.Namespace): - twinkle.initialize(mode='ray', nproc_per_node=args.base_gpus, lazy_collect=False, - groups=[DeviceGroup(name='base_sampler', ranks=list(range(args.base_gpus)), device_type='GPU')]) - sampler = vLLMSampler( - model_id=MODEL_ID, - engine_args={'gpu_memory_utilization': GPU_MEM, - 'max_model_len': args.max_model_len, - 'tensor_parallel_size': 1}, - device_mesh=DeviceMesh.from_sizes(world_size=args.base_gpus, dp_size=args.base_gpus), - remote_group='base_sampler') - sampler.set_template('Template', model_id=MODEL_ID, enable_thinking=True, - max_length=args.max_model_len) - return sampler, args.base_gpus - - -def _select_records(args: argparse.Namespace) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: - load_n = 0 if args.numeric_only or args.eval_size > 0 else max(args.n, args.target_size + args.eval_size) - records = load_problems(args.dataset, load_n, args.seed) - raw_n = len(records) - if args.numeric_only: - records = [{**r, 'reference_answer': v} - for r, v in ((r, _numeric_value(r.get('reference_answer'))) for r in records) - if v is not None] - import numpy as np - np.random.RandomState(args.seed).shuffle(records) - exclude_ids, exclude_problems = _load_excluded_records(getattr(args, 'exclude_data_ids', '')) - excluded = 0 - if exclude_ids or exclude_problems: - before = len(records) - records = [r for r in records - if str(r.get('data_id', '')) not in exclude_ids - and str(r.get('problem', '')).strip() not in exclude_problems] - excluded = before - len(records) - eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 - pool = [dict(r) for r in records[eval_n:]] - pool_offset = max(0, int(getattr(args, 'pool_offset', 0) or 0)) - if pool_offset: - if pool_offset >= len(pool): - raise ValueError(f'--pool-offset ({pool_offset}) leaves no cold-start records from pool size {len(pool)}') - pool = pool[pool_offset:] - n = min(args.n, len(pool)) if args.n > 0 else min(len(pool), max(args.target_size * 2, args.target_size + 512)) - stats = {'raw_loaded': raw_n, 'numeric_dropped': raw_n - len(records) - excluded, - 'excluded_records': excluded, 'eval_size': eval_n, - 'pool_offset': pool_offset, 'pool_selected': n} - return pool[:n], stats - - -def _build_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--dataset', choices=('aops', 'math'), default='aops') - p.add_argument('--target-size', type=int, default=10000, help='Number of accepted SFT examples to write.') - p.add_argument('--n', type=int, default=0, help='Raw train-pool size after eval split; 0 auto-selects.') - p.add_argument('--pool-offset', type=int, default=0, - help='Skip this many shuffled non-eval records before building the cold-start pool.') - p.add_argument('--exclude-data-ids', default='', - help='Comma-separated jsonl files whose data_id/problem keys are excluded, ' - 'useful for building non-overlapping shards.') - p.add_argument('--eval-size', type=int, default=128) - p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) - p.add_argument('--seed', type=int, default=42) - p.add_argument('--output-dir', default='./output/reflexion_coldstart_sft') - p.add_argument('--cache-dir', default='') - p.add_argument('--overwrite', action='store_true') - p.add_argument('--no-cache', action='store_true') - p.add_argument('--chunk-size', type=int, default=64) - p.add_argument('--base-gpus', type=int, default=int(os.environ.get('BASE_GPUS', 4))) - p.add_argument('--max-model-len', type=int, default=16384) - p.add_argument('--max-tokens', type=int, default=8192) - p.add_argument('--rubric-workers', type=int, default=16) - p.add_argument('--rubric-retries', type=int, default=2) - p.add_argument('--api-workers', type=int, default=16) - p.add_argument('--api-samples', type=int, default=4, - help='API skill targets sampled per problem before executor verification.') - p.add_argument('--verify-targets', action=argparse.BooleanOptionalAction, default=True, - help='Run frozen base executor with each API skill target and keep a successful one.') - p.add_argument('--keep-unverified-targets', action='store_true', - help='If all executor checks fail, keep the first clean target anyway. Default skips it.') - p.add_argument('--api-retries', type=int, default=3) - p.add_argument('--api-timeout', type=int, default=120) - p.add_argument('--api-max-tokens', type=int, default=768) - p.add_argument('--api-temperature', type=float, default=0.2) - p.add_argument('--require-fail', action=argparse.BooleanOptionalAction, default=True, - help='Only keep API diagnoses containing [FAIL]. Use --no-require-fail to keep OK diagnoses too.') - return p.parse_args() - - -def _write(f, row: Dict[str, Any]) -> None: - f.write(json.dumps(row, ensure_ascii=False) + '\n') - - -def main() -> None: - args = _build_args() - if args.target_size <= 0: - raise ValueError('--target-size must be positive') - records, data_stats = _select_records(args) - if not records: - raise ValueError('no records selected') - - os.makedirs(args.output_dir, exist_ok=True) - sft_path = os.path.join(args.output_dir, 'coldstart_sft.jsonl') - rec_path = os.path.join(args.output_dir, 'coldstart_records.jsonl') - for path in (sft_path, rec_path): - if os.path.exists(path) and not args.overwrite: - raise FileExistsError(f'{path} exists; pass --overwrite') - - checker = build_rubric_checker() - if checker is None: - raise RuntimeError('No rubric checker available; set LLM_BACKUP_API_KEY/BASE_URL or OPENAI_API_KEY.') - _api_config() - base_sampler, base_dp = _init_base_sampler(args) - - cache_dir = args.cache_dir or os.path.join(args.output_dir, 'cache') - os.makedirs(cache_dir, exist_ok=True) - use_cache = not args.no_cache - base_cache = DiskCache(os.path.join(cache_dir, 'baseline.jsonl'), use_cache) - rubric_cache = DiskCache(os.path.join(cache_dir, 'rubric.jsonl'), use_cache) - skill_cache = DiskCache(os.path.join(cache_dir, 'api_skill.jsonl'), use_cache) - - cfg = { - 'record_type': 'config', 'mode': 'coldstart_sft_build', 'dataset': args.dataset, - 'target_size': args.target_size, 'selected_records': len(records), 'seed': args.seed, - 'pool_offset': args.pool_offset, 'exclude_data_ids': args.exclude_data_ids, - 'numeric_only': args.numeric_only, **data_stats, - 'rubric_version': _RUBRIC_VERSION, 'rubric_check': f'fixed_math_{len(_MATH_RUBRIC)}crit', - 'api_model': os.environ.get('LLM_BACKUP_MODEL') or os.environ.get('OPENAI_MODEL') or 'gpt-4o-mini', - 'api_samples': args.api_samples, 'verify_targets': args.verify_targets, - 'keep_unverified_targets': args.keep_unverified_targets, - 'require_fail': args.require_fail, 'started': int(time.time()), - } - - accepted = 0 - skipped_no_diag = skipped_no_fail = skipped_api_leak = 0 - skipped_no_skill = skipped_skill_leak = skipped_executor_fail = 0 - processed = 0 - with open(sft_path, 'w', encoding='utf-8') as sft_f, open(rec_path, 'w', encoding='utf-8') as rec_f: - _write(rec_f, cfg) - for start in range(0, len(records), args.chunk_size): - if accepted >= args.target_size: - break - chunk = [dict(r) for r in records[start:start + args.chunk_size]] - _baseline_rollout(base_sampler, chunk, base_dp, args, base_cache) - _diagnose_batch(checker, chunk, args, rubric_cache) - - def gen_one(r: Dict[str, Any]): - return r, _generate_skill_targets(r, args, skill_cache) - - candidates = [] - for r in chunk: - processed += 1 - diag = r.get('_rubric_diag', '') or '' - if not diag: - skipped_no_diag += 1 - continue - if args.require_fail and '[FAIL]' not in diag: - skipped_no_fail += 1 - continue - if _answer_leaked(diag, r['reference_answer']): - skipped_api_leak += 1 - continue - candidates.append(r) - - generated = [] - workers = max(1, min(args.api_workers, len(candidates))) - if candidates: - with ThreadPoolExecutor(max_workers=workers) as ex: - for r, targets in ex.map(gen_one, candidates): - for target in targets: - skills = target.get('skills', '') - if not skills: - skipped_no_skill += 1 - continue - if target.get('skill_leak'): - skipped_skill_leak += 1 - continue - target['r'] = r - target['response'] = _skill_response(skills) - generated.append(target) - - selected = [] - selected_keys = set() - if generated and args.verify_targets: - verify_prompts = [build_skill_solve_prompt(g['r']['problem'], g['skills']) for g in generated] - verify_outs = _run_samples(base_sampler, verify_prompts, 1, args.max_tokens, - base_dp, temperature=0.0) - attempted_keys = set() - for g, seqs in zip(generated, verify_outs): - r = g['r'] - key = r.get('data_id') or r['problem'] - attempted_keys.add(key) - roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - g['target_roll'] = roll - if key not in selected_keys and roll.get('correct') and roll.get('terminated'): - g['executor_verified'] = True - selected.append(g) - selected_keys.add(key) - if args.keep_unverified_targets: - for g in generated: - r = g['r'] - key = r.get('data_id') or r['problem'] - if key not in selected_keys: - g['executor_verified'] = False - g.setdefault('target_roll', {}) - selected.append(g) - selected_keys.add(key) - skipped_executor_fail += len(attempted_keys - selected_keys) - elif generated: - for g in generated: - r = g['r'] - key = r.get('data_id') or r['problem'] - if key not in selected_keys: - g['executor_verified'] = False - selected.append(g) - selected_keys.add(key) - - for g in selected: - if accepted >= args.target_size: - break - r = g['r'] - response = g['response'] - messages = _sft_messages(r['problem'], response) - sft_row = { - 'messages': messages, - 'user_data': pack_user_data({'key_rounds': [len(messages) - 1]}), - 'data_id': r.get('data_id'), - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'skills': g['skills'], 'response': response, - 'view': 'B', 'sft': True, 'source': 'api_coldstart', - 'api_sample_idx': g.get('sample_idx'), - 'executor_verified': g.get('executor_verified', False), - 'baseline_correct': r['_init'][0].get('correct'), - 'baseline_terminated': r['_init'][0].get('terminated'), - 'baseline_stop_reason': r['_init'][0].get('stop_reason'), - 'target_correct': (g.get('target_roll') or {}).get('correct'), - 'target_terminated': (g.get('target_roll') or {}).get('terminated'), - 'target_stop_reason': (g.get('target_roll') or {}).get('stop_reason'), - 'diagnosis': r.get('_rubric_diag', ''), - } - audit = { - 'record_type': 'coldstart_problem', 'accepted': True, - 'data_id': r.get('data_id'), - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'baseline': r['_init'][0], 'diagnosis': r.get('_rubric_diag', ''), - 'raw_skill_response': g.get('raw_response'), 'skills': g['skills'], - 'api_sample_idx': g.get('sample_idx'), - 'executor_verified': g.get('executor_verified', False), - 'target_roll': g.get('target_roll'), - } - _write(sft_f, sft_row) - _write(rec_f, audit) - accepted += 1 - sys.stderr.write( - f'[coldstart] processed={processed} accepted={accepted}/{args.target_size} ' - f'skip(no_diag={skipped_no_diag}, no_fail={skipped_no_fail}, api_leak={skipped_api_leak}, ' - f'no_skill={skipped_no_skill}, skill_leak={skipped_skill_leak}, ' - f'executor_fail={skipped_executor_fail})\n') - sft_f.flush(); rec_f.flush() - - summary = { - 'record_type': 'summary', 'processed': processed, 'accepted': accepted, - 'skipped_no_diag': skipped_no_diag, 'skipped_no_fail': skipped_no_fail, - 'skipped_api_leak': skipped_api_leak, 'skipped_no_skill': skipped_no_skill, - 'skipped_skill_leak': skipped_skill_leak, - 'skipped_executor_fail': skipped_executor_fail, 'finished': int(time.time()), - } - with open(rec_path, 'a', encoding='utf-8') as rec_f: - _write(rec_f, summary) - sys.stderr.write(f'[coldstart] wrote {accepted} SFT rows to {sft_path}\n') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/legacy/build_reflexion_skill_data.py b/cookbook/exp/legacy/build_reflexion_skill_data.py deleted file mode 100644 index 59fe2382c..000000000 --- a/cookbook/exp/legacy/build_reflexion_skill_data.py +++ /dev/null @@ -1,1123 +0,0 @@ -"""Offline builder for reflexion skill RFT data (self-contained, cached). - -Runs the SAME pipeline as the online trainer -- base greedy solve -> rubric -process-check (view A) -> skill-gen -> leak filter -> with-skill greedy pass -> -group-relative GRPO advantage -- but never updates the skill model. It emits -``skill_dataset.jsonl`` (trainer-schema training records), ``gen_records.jsonl`` -(full per-problem traces) and ``eval_holdout.jsonl`` (the fixed holdout). - -The expensive base rollouts and rubric diagnoses are cached to disk (one jsonl -each, keyed by an md5 of their inputs) so a re-run skips them entirely. - -8 GPUs: ranks 0-3 skill_sampler (vLLM tp1 dp4), ranks 4-7 base_sampler. Leak / -rubric use the backup teacher API: set LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / -LLM_BACKUP_MODEL. - -Launch: - LLM_BACKUP_API_KEY=... python cookbook/exp/embedding/build_reflexion_skill_data.py \ - --total-problems 3200 --base-success-frac 0.3 -""" -import argparse -import copy -import hashlib -import json -import math -import os -import re -import sys -import time -from concurrent.futures import ThreadPoolExecutor -from typing import Any, Dict, List, Optional, Tuple - -import numpy as np - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.data_format import SamplingParams -from twinkle.dataset import Dataset, DatasetMeta -from twinkle.sampler import vLLMSampler -from twinkle_agentic.verifier import LeakVerifier, RubricVerifier -from twinkle_agentic.verifier.rubric_verifier import RubricItem - -logger = get_logger() - -MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') -GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.8)) -GEN_TEMPERATURE = float(os.environ.get('GEN_TEMPERATURE', 0.6)) -GEN_TOP_P = float(os.environ.get('GEN_TOP_P', 0.95)) -AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') -MATH_DATASET_ID = os.environ.get('MATH_DATASET_ID', 'modelscope/competition_math') - - -# =========================================================================== -# Block A -- boxed extraction + answer grading -# =========================================================================== -_BOXED_RE = re.compile(r'\\boxed\s*\{') - - -def extract_boxed(text: str) -> Optional[str]: - """Last ``\\boxed{...}`` content, brace-balanced.""" - if not text: - return None - last = None - for m in _BOXED_RE.finditer(text): - depth, i = 1, m.end() - while i < len(text) and depth > 0: - depth += (text[i] == '{') - (text[i] == '}') - i += 1 - if depth == 0: - last = text[m.end():i - 1].strip() - return last - - -def normalize_answer(ans: str) -> str: - if not ans: - return '' - s = str(ans).strip() - m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) - if m: - return m.group(1) - s = s.strip('$').strip().replace('−', '-') - 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').replace(r'\tfrac', r'\frac') - s = s.replace(r'\,', '').replace(r'\;', '').replace(r'\!', '') - s = re.sub(r'\\(?:quad|qquad|\s)', '', s) - s = re.sub(r'\s+', '', s) - s = re.sub(r'\\sqrt([0-9A-Za-z])', r'\\sqrt{\1}', s) - s = re.sub(r'\\frac\{([^{}]+)\}([^{}\\])', r'\\frac{\1}{\2}', s) - s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) - s = re.sub(r'\{[a-zA-Z]+\}$', '', s) - s = re.sub(r'\^\\circ|\^\{\\circ\}|°|\\circ', 'deg', s) - s = s.replace(r'\minus{}', '-').replace(r'\minus', '-') - - def _frac_to_slash(mt): - text = mt.group(0) - pos = text.index('{') + 1 - depth, num_start = 1, pos - while depth > 0: - depth += (text[pos] == '{') - (text[pos] == '}') - pos += 1 - numer = text[num_start:pos - 1] - pos += 1 - den_start, depth = pos, 1 - while depth > 0: - depth += (text[pos] == '{') - (text[pos] == '}') - pos += 1 - return f'({numer})/({text[den_start:pos - 1]})' - - s = re.sub(r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', _frac_to_slash, s) - s = re.sub(r'(? bool: - try: - va, vb = float(a.replace('(', '').replace(')', '')), 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) - return va is not None and vb is not None and abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) - - -_MCQ_REF_RE = re.compile( - r'^\\(?:textbf|text|mathrm|mathbf)\{\(?([A-E])\)?\s*\}\s*(.+)$' - r'|^\(?([A-E])\)\s+(.+)$') -_VAR_PREFIX_RE = re.compile(r'^(?:[A-Za-z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)$', re.DOTALL) - - -def _split_mcq(ans: str): - s = ans.strip() - m = _MCQ_REF_RE.match(s) - if m: - return (m.group(1) or m.group(3)), ((m.group(2) or m.group(4) or '').strip() or None) - bl = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) - return (bl.group(1), None) if bl else (None, s or None) - - -def _strip_var_prefix(ans: str) -> str: - m = _VAR_PREFIX_RE.match((ans or '').strip()) - return m.group(1).strip() if m else (ans or '') - - -def _normalize_tuple(ans: str) -> str: - return re.sub(r'[\s()\[\]{}\\]', '', ans or '') - - -def _math_verify_equal(predicted: str, reference: str) -> bool: - try: - from math_verify import parse, verify - gold = parse(r'\boxed{' + reference + '}', parsing_timeout=1) - pred = parse(r'\boxed{' + predicted + '}', parsing_timeout=1) - return bool(gold and pred and verify(gold, pred, timeout_seconds=1)) - except Exception: - return False - - -def answers_match(predicted: str, reference: str) -> bool: - if not predicted or not reference: - return False - norm_p, norm_r = normalize_answer(predicted), normalize_answer(reference) - if norm_p == norm_r or norm_p.lower() == norm_r.lower() or _try_numeric_equal(norm_p, norm_r): - return True - stripped_p = normalize_answer(_strip_var_prefix(predicted)) - stripped_r = normalize_answer(_strip_var_prefix(reference)) - if stripped_p and stripped_r: - if (stripped_p == stripped_r or stripped_p.lower() == stripped_r.lower() - or _try_numeric_equal(stripped_p, stripped_r)): - return True - p_letter, p_value = _split_mcq(predicted) - r_letter, r_value = _split_mcq(reference) - if p_letter and r_letter and p_letter == r_letter: - return True - if (p_letter and p_value is None) and (r_letter and r_value is None): - return False - 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 and (p_val == r_val or p_val.lower() == r_val.lower() or _try_numeric_equal(p_val, r_val)): - return True - for left, right in ((norm_p, norm_r), (stripped_p, stripped_r)): - tl, tr = _normalize_tuple(left), _normalize_tuple(right) - if ',' in tl and tl == tr: - return True - if '=' in norm_r and '=' not in norm_p: - if any(part and (part == norm_p or _try_numeric_equal(part, norm_p)) for part in norm_r.split('=')): - return True - if '=' in norm_p and '=' not in norm_r: - if any(part and (part == norm_r or _try_numeric_equal(part, norm_r)) for part in norm_p.split('=')): - return True - return _math_verify_equal(stripped_p or norm_p, stripped_r or norm_r) - - -# =========================================================================== -# Block B -- prompts, skill parsing, batched sampling -# =========================================================================== -DIRECT_SYSTEM = ( - 'You are an expert competition mathematician. Solve the following problem ' - 'step by step. Provide your final answer inside \\boxed{}.') - -# Appended to solve turns: box BOTH the letter and value of an MCQ so the model -# never loops deciding which form to box. -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.') - -_SKILL_SOLVE_PREFIX = ( - DIRECT_SYSTEM + '\n\n' - 'Before you start, keep these reminders in mind to avoid common mistakes on this ' - 'type of problem:\n') -_SKILL_SOLVE_SUFFIX = '\nApply them where relevant, but rely on your own reasoning to reach the answer.' - - -def build_direct_prompt(problem: str) -> Dict[str, Any]: - return {'messages': [{'role': 'system', 'content': DIRECT_SYSTEM}, - {'role': 'user', 'content': problem}]} - - -def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: - # Concatenation (not .format): DIRECT_SYSTEM/skill contain literal braces. - return {'messages': [ - {'role': 'system', 'content': _SKILL_SOLVE_PREFIX + skill + _SKILL_SOLVE_SUFFIX}, - {'role': 'user', 'content': problem + MCQ_INSTRUCTION}]} - - -# -- skill-gen prompts (view A: problem + rubric findings; view B: query only) -- -SKILL_GEN_SYSTEM = ( - 'You are a mathematics coach. You are shown a competition problem together with an ' - 'automated process-check of an earlier solver attempt at it -- which solution ' - 'criteria the attempt passed or failed, and suggested fixes for the failures. You ' - 'do NOT see the attempt itself, only this check. Treat the check as privileged ' - 'training scaffolding: study it together with the problem, identify the ' - 'problem-visible features that make each useful flagged failure relevant, then ' - 'rephrase those lessons as self-contained reusable skills. The goal is not to ' - 'continue from the check, cite it, or hide it silently; the goal is to turn it into ' - 'a problem-triggered reasoning pattern a query-only solver could reproduce later.\n\n' - 'Good skills name the observable trigger, the method worth reaching for, the ' - 'pitfall to watch, and a quick verification habit. Prefer formulations like ' - '"When a configuration has ...", "Before setting up ...", or "Check whether ..." ' - 'over references to the process-check, failed criteria, or the earlier attempt. ' - "These tips are advisory: they will be placed in a solver's system prompt as gentle " - 'reminders before it works through a SIMILAR problem on its own, without seeing ' - 'this process-check. So keep them general and transferable rather than a worked ' - 'solution to this exact problem, and do not state its specific intermediate values ' - 'or final answer. Think briefly first, then give your tips as a markdown bullet ' - 'list wrapped in and , like the example below.') - -SKILL_GEN_SYSTEM_Q = ( - 'You are a mathematics coach. You are shown ONE competition problem and nothing ' - 'else — no solution and no attempt. Think about what approach this KIND of problem ' - 'calls for and where solvers tend to slip, then distil a few reusable tips.\n\n' - "These tips are advisory: they will be placed in a solver's system prompt as gentle " - 'reminders before it works through a SIMILAR problem on its own. So keep them ' - 'general and transferable — the method worth reaching for, the pitfall to watch and ' - 'a quick check, and the discipline to settle on a final answer — rather than a ' - 'worked solution to this exact problem, and without stating its specific ' - 'intermediate values or its final answer. Think briefly first, then give your tips ' - 'as a markdown bullet list wrapped in and , like the example below.') - -SKILL_GEN_USER_Q = ( - 'Problem:\n{problem}\n\n' - 'Now reason about this TYPE of problem, then output the skills bullet list.') - -SKILL_GEN_USER_RUBRIC = ( - 'Problem:\n{problem}\n\n' - 'Process check of an earlier attempt (automated rubric verifier -- treat as ' - 'evidence, not gospel; PASS/FAIL per criterion with suggested fixes for failures):\n' - '{diagnosis}\n\n' - 'Now output a self-contained skills bullet list. Each bullet should still be useful ' - 'if the process check were removed: connect any useful flagged failure to ' - 'problem-visible features, general methods, and quick checks rather than citing the ' - 'rubric or the earlier attempt.') - -_EX_PROBLEM = 'Simplify $\\sqrt{72} + \\sqrt{18}$ and give the result.' -_EX_SKILLS = ( - '\n' - '- Rewrite each square root by factoring its radicand into a perfect square times ' - 'a remainder, then move the perfect-square factor outside.\n' - '- Avoid the classic trap $\\sqrt{a}+\\sqrt{b}\\ne\\sqrt{a+b}$; only combine terms ' - 'sharing the same simplest radical, and sanity-check by estimating each root.\n' - '- Procedure: simplify every radical, group like radical terms, add their ' - 'coefficients, then reduce to simplest form.\n' - '- Once the expression is in simplest form, commit to that single result as the ' - 'final answer rather than re-checking indefinitely.\n' - '') - - -def _skillgen_messages(problem: str, view: str, diagnosis: str) -> List[Dict[str, Any]]: - """Single source of truth for the skill-gen prompt. View A with a localisable - failure uses problem + rubric findings; view B -- or a view-A problem whose rubric - flagged nothing (no ``[FAIL]``) -- degrades to the query-only prompt.""" - if view == 'B' or '[FAIL]' not in (diagnosis or ''): - return [{'role': 'system', 'content': SKILL_GEN_SYSTEM_Q}, - {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=_EX_PROBLEM)}, - {'role': 'assistant', 'content': _EX_SKILLS}, - {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=problem)}] - return [{'role': 'system', 'content': SKILL_GEN_SYSTEM}, - {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=_EX_PROBLEM)}, - {'role': 'assistant', 'content': _EX_SKILLS}, - {'role': 'user', 'content': SKILL_GEN_USER_RUBRIC.format(problem=problem, diagnosis=diagnosis)}] - - -def _assign_view(problem: str, args: argparse.Namespace) -> str: - h = int(hashlib.md5(f'{args.seed}:{problem}'.encode('utf-8')).hexdigest(), 16) - return 'B' if (h % 100000) / 100000.0 < args.view_b_frac else 'A' - - -def _view_prompt(r: Dict[str, Any]) -> Dict[str, Any]: - return {'messages': _skillgen_messages(r['problem'], r['_view'], r.get('_rubric_diag', ''))} - - -_BULLET_RE = re.compile(r'(?m)^\s*(?:[-*]|\d+[.)])\s') -_META_RE = re.compile( - r'\b(the student|the solver|the attempt|this attempt|the trace|the response|' - r'in the (?:attempt|trace|response|solution)|as (?:shown|seen|noted) above|' - r'the (?:above|previous|earlier)|my (?:reasoning|analysis)|i (?:think|need|will))\b', - re.IGNORECASE) -_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') - - -def _clean_text(decoded: Optional[str]) -> str: - return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() - - -def _is_clean_block(block: str) -> bool: - """Pure bullet list (every non-empty line a bullet) with no meta/trajectory ref.""" - lines = [ln for ln in (l.strip() for l in block.splitlines()) if ln] - if not lines or not all(_BULLET_RE.match(ln) for ln in lines): - return False - return _META_RE.search(block) is None - - -def _extract_skills_block(text: str) -> Optional[str]: - """Clean ``...`` block, or None. Requires ```` (skill-gen - runs thinking ON); reads only the answer after the last one, so a mid-reasoning draft - or a demo echo can never be mistaken for the answer.""" - low = text.lower() - end_think = low.rfind('') - if end_think < 0: - return None - answer = text[end_think + len(''):] - low_a = answer.lower() - s = low_a.find('') - if s < 0: - return None - inner = s + len('') - e = low_a.find('', inner) - block = (answer[inner:e] if e >= 0 else answer[inner:]).strip() - block = re.sub(r'', '', block, flags=re.IGNORECASE).strip() - return block if _is_clean_block(block) else None - - -def _parse_seq(seq, gold: str) -> Dict[str, Any]: - """Grade one sampled sequence into a rollout record.""" - text = _clean_text(getattr(seq, 'decoded', '') or '') - pred = extract_boxed(text) - correct = bool(pred) and answers_match(pred, gold) - terminated = getattr(seq, 'stop_reason', None) != 'length' - return {'pred': pred, 'correct': correct, 'terminated': terminated, - 'passed': bool(correct and terminated), - 'stop_reason': getattr(seq, 'stop_reason', None), - 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} - - -def _run_samples(sampler, prompts: List[Any], num_samples: int, max_tokens: int, - gen_dp: int, temperature: Optional[float] = None, - top_p: Optional[float] = None, top_k: Optional[int] = None) -> List[List[Any]]: - """One batched sampler call -> per-prompt list of raw sequences. vLLM dp needs - batch len >= dp, so pad the tail and slice back.""" - if not prompts: - return [] - params = SamplingParams( - max_tokens=max_tokens, - temperature=GEN_TEMPERATURE if temperature is None else temperature, - top_p=GEN_TOP_P if top_p is None else top_p, - num_samples=num_samples, **({} if top_k is None else {'top_k': top_k})) - padded = prompts - if gen_dp > 1 and 0 < len(prompts) < gen_dp: - padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] - responses = sampler.sample(padded, params)[:len(prompts)] - return [list(r.sequences) if (r and r.sequences) else [] for r in responses] - - -# =========================================================================== -# Block C -- data loading via twinkle.Dataset + numeric filtering -# =========================================================================== -def load_problems(dataset: str, n: int, seed: int) -> List[Dict[str, Any]]: - """Load boxed-answer problems as ``{problem, reference_answer, level?}`` via - twinkle.Dataset (ModelScope hub), sampled to ``n`` (0 = all).""" - ds_id = AOPS_DATASET_ID if dataset == 'aops' else MATH_DATASET_ID - rows = Dataset(DatasetMeta(dataset_id=f'ms://{ds_id}', split='train')).dataset - out: List[Dict[str, Any]] = [] - for row in rows: - if dataset == 'aops' and not (row.get('metadata') or {}).get('boxed'): - continue - ref = extract_boxed(row.get('solution', '')) - if not ref: - continue - rec = {'problem': row['problem'], 'reference_answer': ref} - if row.get('level'): - rec['level'] = row['level'] - out.append(rec) - logger.info(f'[data] {dataset}: {len(out)} boxed problems') - rng = np.random.RandomState(seed) - rng.shuffle(out) - return out[:n] if (n and n < len(out)) else out - - -_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') - - -def _norm_num_text(num: str) -> str: - try: - f = float(num) - return str(int(f)) if f == int(f) else str(f) - except Exception: - return str(num).strip() - - -def _numeric_value(raw: Any) -> Optional[str]: - """Collapse an answer to a single int/decimal/fraction, or None.""" - if raw is None: - return None - s = str(raw).strip().strip('$').strip() - s = s.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac') - s = re.sub(r'\\!|\\,|\\;|\\ |\\left|\\right|\s', '', s) - for pat in (r'\\frac\{(-?\d+)\}\{(-?\d+)\}', r'(-?\d+)/(-?\d+)'): - m = re.fullmatch(pat, s) - if m: - a, b = int(m.group(1)), int(m.group(2)) - return _norm_num_text(str(a / b)) if b else None - return _norm_num_text(s) if _NUM_RE.fullmatch(s) else None - - -def _load_records(args: argparse.Namespace) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], Dict[str, int]]: - """Load, numeric-filter, shuffle, then split a fixed eval holdout off the front.""" - # Load all when filtering or splitting (else the eval holdout could starve train). - load_n = 0 if (args.numeric_only or args.eval_size > 0) else args.n - records = load_problems(args.dataset, load_n, args.seed) - raw_n, dropped = len(records), 0 - if args.numeric_only: - kept = [] - for r in records: - ref = _numeric_value(r.get('reference_answer')) - if ref is None: - dropped += 1 - continue - kept.append({**r, 'reference_answer': ref}) - records = kept - np.random.RandomState(args.seed).shuffle(records) - eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 - eval_records = [dict(r) for r in records[:eval_n]] - pool = records[eval_n:] - train_n = args.n if args.n > 0 else len(pool) - train_records = [dict(r) for r in pool[:train_n]] - overlap = {r['problem'] for r in train_records} & {r['problem'] for r in eval_records} - if overlap: - raise ValueError(f'eval/train overlap: {len(overlap)} problems') - stats = {'raw_loaded': raw_n, 'numeric_dropped': dropped, - 'train_records': len(train_records), 'eval_records': len(eval_records)} - return train_records, eval_records, stats - - -# =========================================================================== -# Block D -- disk cache, problem pool, baseline rollout, rubric check -# =========================================================================== -class DiskCache: - """Append-only jsonl kv cache (md5 key -> value). Loads on init, appends on put. - Disabled instances (``enabled=False``) always miss and never write.""" - - def __init__(self, path: str, enabled: bool = True): - self.path, self.enabled = path, enabled - self._mem: Dict[str, Any] = {} - self._fh = None - if not enabled: - return - if os.path.exists(path): - with open(path, encoding='utf-8') as f: - for line in f: - if line.strip(): - row = json.loads(line) - self._mem[row['key']] = row['value'] - self._fh = open(path, 'a', encoding='utf-8') - - @staticmethod - def key_for(*parts: str) -> str: - return hashlib.md5('\x1f'.join(parts).encode('utf-8')).hexdigest() - - def __contains__(self, key: str) -> bool: - return key in self._mem - - def get(self, key: str) -> Any: - return self._mem.get(key) - - def put(self, key: str, value: Any) -> None: - self._mem[key] = value - if self._fh is not None: - self._fh.write(json.dumps({'key': key, 'value': value}, ensure_ascii=False) + '\n') - self._fh.flush() - - def close(self) -> None: - if self._fh is not None: - self._fh.close() - - -class ProblemPool: - """Cyclic draw source. Each full pass reshuffles with ``seed + epoch``; the initial - pass keeps the loader order. ``draw(k)`` returns k distinct problems (k << pool).""" - - def __init__(self, records: List[Dict[str, Any]], seed: int): - self._records = list(records) - self._seed, self._cursor, self.epoch = seed, 0, 0 - - def draw(self, k: int) -> List[Dict[str, Any]]: - out, seen = [], set() - while len(out) < k: - if self._cursor >= len(self._records): - self.epoch += 1 - np.random.RandomState(self._seed + self.epoch).shuffle(self._records) - self._cursor = 0 - r = self._records[self._cursor] - self._cursor += 1 - if id(r) not in seen: - seen.add(id(r)) - out.append(r) - return out - - -def _empty_roll() -> Dict[str, Any]: - return {'pred': '', 'correct': False, 'terminated': False, 'passed': False, - 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} - - -def _apply_baseline(r: Dict[str, Any], roll: Dict[str, Any]) -> None: - """Attach a greedy baseline roll and reset per-chunk working state.""" - r['_baseline_rolls'], r['_cands'], r['_init'] = [roll], [], [roll] - r['_failed'] = not roll['correct'] - r['_baseline_pass'] = 1.0 if roll['correct'] else 0.0 - r['_hard'] = True # process every problem; group variance selects (SEAM-style) - - -def baseline_rollout(base_sampler, problems: List[Dict[str, Any]], base_dp: int, - args: argparse.Namespace, cache: DiskCache) -> int: - """Phase 1: base solves each problem greedily once (T=0, M=1), disk-cached by - problem text. The base is frozen + greedy so the cache is exact. Returns the number - of fresh (cache-miss) rollouts.""" - todo = [r for r in problems if DiskCache.key_for(r['problem']) not in cache] - if todo: - out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], - 1, args.max_tokens, base_dp, temperature=0.0) - for r, seqs in zip(todo, out): - roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - cache.put(DiskCache.key_for(r['problem']), roll) - for r in problems: - _apply_baseline(r, cache.get(DiskCache.key_for(r['problem']))) - return len(todo) - - -# -- rubric process-check (view A): teacher diagnoses the base's attempt -- -_RFT_DIAG_SYSTEM = """\ -You are a process error checker for a math solution attempt. You are given a math -problem, a rubric, and one attempted solution segment. Decide PASS or FAIL for each -criterion and explain only the process error type. - -Output STRICT JSON (no prose outside it) with this shape: -{ - "items": [ - {"index": 1, "verdict": "PASS", "reason": "", - "fix": ""}, - {"index": 2, "verdict": "FAIL", "reason": "", - "fix": ""} - ], - "overall": "OK" | "ISSUES", - "summary": "" -} - -Rules: -- Judge every criterion independently and literally; a [Hard Rule] is FAIL unless - unambiguously satisfied. -- Judge ONLY what is observable in THIS segment. -- Content inside ... (or ) is internal reasoning, not - user-facing output; ignore it for "output only X" style criteria. -- For PASS items, leave "fix" as "". -- For FAIL items, "reason", "fix", and "summary" must describe only the flawed - step, theorem, arithmetic operation, case split, or verification habit. -- NEVER state the correct final answer, corrected final expression, option letter, - graph/choice label, or any exact value that the answer should become. -- NEVER write phrases like "the correct answer is", "which gives", "yielding", - "should be ", "Option ", or "Graph ". -- If a fix would require naming a corrected value, replace it with a method-level - instruction such as "redo that computation carefully" or "apply the theorem with - the correct quantities". -- Keep every "reason" and "fix" clear and concise — one short sentence each. -- "overall" is "OK" only if NO criterion is FAIL. -- Output only the JSON object.""" - -_RFT_DIAG_USER = """\ -## Task / query (context) -{query} - -## Rubric -{rubric} - -## Segment -{segment} - -Now output the diagnostic JSON object.""" - -_MATH_RUBRIC = [ - ('The reasoning contains no arithmetic or algebraic error', True), - ('Each step follows logically from the previous ones', True), - ('No formula or theorem is misstated or misapplied', True), - ('The approach is on track to answer the actual question asked', False), - ('No step contradicts an earlier established fact', False), -] - - -class _RftRubricVerifier(RubricVerifier): - def _diagnose_trajectory(self, query: str, rubric_block: str, segment_text: str) -> dict: - return {'messages': [ - {'role': 'system', 'content': _RFT_DIAG_SYSTEM}, - {'role': 'user', 'content': _RFT_DIAG_USER.format( - query=query, rubric=rubric_block, segment=segment_text)}]} - - -def build_rubric_checker() -> Optional[RubricVerifier]: - """Fixed math-process rubric verifier (teacher-served). None if no LLM backup env.""" - if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') - or os.environ.get('OPENAI_API_KEY')): - return None - return _RftRubricVerifier( - fixed_rubric=[RubricItem(t, is_hard=h) for t, h in _MATH_RUBRIC], gate=True) - - -def _format_diagnosis(detail) -> str: - """One line per criterion (PASS/FAIL + reason + fix) then a summary.""" - rub = detail.rubric - lines = [] - for it in detail.items: - text = rub[it.index - 1].text if 0 < it.index <= len(rub) else f'criterion {it.index}' - if it.verdict: - lines.append(f'- [PASS] {text}') - else: - tail = f': {it.reason}' if it.reason else '' - tail += f' (fix: {it.fix})' if it.fix else '' - lines.append(f'- [FAIL] {text}{tail}') - if detail.summary: - lines.append(f'Summary: {detail.summary}') - return '\n'.join(lines) - - -def diagnose_views(checker, hard: List[Dict[str, Any]], args: argparse.Namespace, - cache: DiskCache) -> None: - """Rubric-check every view-A problem's greedy attempt in parallel (disk-cached by - problem + attempt), stashing the formatted findings on ``r['_rubric_diag']``.""" - targets = [r for r in hard if r.get('_view') == 'A'] - if not checker or not targets: - return - - def _key(r: Dict[str, Any]) -> str: - return DiskCache.key_for(r['problem'], r.get('_init', [{}])[0].get('text', '')) - - pending = [] - for r in targets: - key = _key(r) - if key in cache: - r['_rubric_diag'] = cache.get(key) - else: - pending.append((r, key)) - if not pending: - return - - def _run(item): - r, key = item - seg = {'messages': [{'role': 'user', 'content': r['problem']}, - {'role': 'assistant', 'content': r['_init'][0]['text']}]} - try: - return r, key, _format_diagnosis(checker.diagnose(seg, query=r['problem'])) - except Exception as exc: # teacher hiccup -> no-diagnosis prompt (not cached) - logger.warning(f'[rubric] diagnose error: {exc}') - return r, key, None - - workers = max(1, min(args.rubric_workers, len(pending))) - with ThreadPoolExecutor(max_workers=workers) as ex: - for r, key, diag in ex.map(_run, pending): - r['_rubric_diag'] = diag or '' - if diag is not None: - cache.put(key, diag) - - -# =========================================================================== -# Block E -- chunk draw, pipeline, record building -# =========================================================================== -def _baseline_class(r: Dict[str, Any]) -> str: - """success | fail_loop (out of length / never terminated) | fail_wrong.""" - roll = r['_init'][0] - if roll['correct']: - return 'success' - return 'fail_loop' if (roll['stop_reason'] == 'length' or not roll['terminated']) else 'fail_wrong' - - -def _select_balanced(buckets: Dict[str, List[Dict[str, Any]]], n_success: int, - n_fail: int, n_fail_loop: int) -> List[Dict[str, Any]]: - """Pick n_fail base-fails (toward n_fail_loop loop-fails, best-effort) + n_success - base-successes; top up any shortfall from leftovers.""" - loop, wrong, succ = buckets['fail_loop'], buckets['fail_wrong'], buckets['success'] - take_loop = min(n_fail_loop, len(loop)) - take_wrong = min(n_fail - take_loop, len(wrong)) - take_loop = min(n_fail - take_wrong, len(loop)) - sel = loop[:take_loop] + wrong[:take_wrong] + succ[:n_success] - target = n_success + n_fail - if len(sel) < target: - used = {id(x) for x in sel} - sel += [x for b in (loop, wrong, succ) for x in b if id(x) not in used][:target - len(sel)] - return sel - - -def draw_chunk(pool: ProblemPool, base_sampler, base_dp: int, args: argparse.Namespace, - cache: DiskCache) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: - """Draw one chunk, running baseline rollout (Phase 1) on every drawn problem. With - ``--balance``, keep drawing+baselining until the target base fail:success mix is - reachable (or the budget is hit), then select a balanced subset.""" - if not args.balance: - chunk = pool.draw(args.chunk_size) - fresh = baseline_rollout(base_sampler, chunk, base_dp, args, cache) - return chunk, {'enabled': False, 'n_drawn': len(chunk), 'n_baseline_fresh': fresh} - - n_success = max(0, min(args.chunk_size, round(args.chunk_size * args.balance_success_frac))) - n_fail = args.chunk_size - n_success - n_fail_loop = round(n_fail * args.balance_loop_frac) - buckets: Dict[str, List[Dict[str, Any]]] = {'success': [], 'fail_loop': [], 'fail_wrong': []} - budget, n_drawn, n_fresh, seen = args.chunk_size * args.balance_max_draws_mult, 0, 0, set() - while n_drawn < budget: - if (len(buckets['success']) >= n_success - and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail): - break - batch = pool.draw(args.chunk_size) - n_fresh += baseline_rollout(base_sampler, batch, base_dp, args, cache) - n_drawn += len(batch) - for r in batch: - if id(r) not in seen: - seen.add(id(r)) - buckets[_baseline_class(r)].append(r) - - reached = (len(buckets['success']) >= n_success - and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail) - chunk = _select_balanced(buckets, n_success, n_fail, n_fail_loop) - sel_success = sum(1 for r in chunk if not r['_failed']) - stats = { - 'enabled': True, 'n_drawn': n_drawn, 'n_baseline_fresh': n_fresh, 'n_selected': len(chunk), - 'target_success': n_success, 'target_fail': n_fail, 'target_fail_loop': n_fail_loop, - 'selected_success': sel_success, 'selected_fail': len(chunk) - sel_success, - 'selected_fail_loop': sum(1 for r in chunk if _baseline_class(r) == 'fail_loop'), - 'selected_fail_wrong': sum(1 for r in chunk if _baseline_class(r) == 'fail_wrong'), - 'selected_success_frac': (sel_success / len(chunk)) if chunk else 0.0, - 'budget_hit': not reached, - } - return chunk, stats - - -def _assign_advantages(hard: List[Dict[str, Any]], args: argparse.Namespace) -> None: - """Group-relative advantage over each problem's scored candidates using the greedy - binary reward R in {0,1}: ``A = (R - mean) / (std + eps)``. std==0 groups (all solve - / all fail) get advantage 0 and no gradient -- GRPO's variance selects informative - problems, so no explicit difficulty gate is needed.""" - eps = 1e-6 - for r in hard: - for c in r['_cands']: - c['advantage'], c['grpo_adv'], c['kept'] = 0.0, 0.0, False - cs = ([c for c in r['_cands'] if c.get('reward') is not None] if args.format_in_reward - else [c for c in r['_cands'] if c['leaked'] is False and c.get('reward') is not None]) - if len(cs) < 2: - continue - rewards = [c['reward'] for c in cs] - mean_r = sum(rewards) / len(rewards) - std = (sum((x - mean_r) ** 2 for x in rewards) / len(rewards)) ** 0.5 - if std < 1e-9: - continue - for c in cs: - adv = (c['reward'] - mean_r) / (std + eps) - c['advantage'], c['grpo_adv'], c['kept'] = adv, adv, c['reward'] > mean_r - - -def process_chunk(base_sampler, skill_sampler, leak: LeakVerifier, chunk: List[Dict[str, Any]], - ci: int, base_dp: int, skill_dp: int, args: argparse.Namespace, - checker, rubric_cache: DiskCache - ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], List[Dict[str, Any]]]: - """view assign -> rubric-check (view A) -> skill-gen -> leak-filter -> with-skill - greedy pass -> GRPO advantages. ``chunk`` arrives already baselined by draw_chunk.""" - hard = chunk - - # Phase 2: view routing + view-A rubric check (view B is query-only, no rubric). - for r in hard: - r['_view'], r['_rubric_diag'] = _assign_view(r['problem'], args), '' - diagnose_views(checker, hard, args, rubric_cache) - - # Phase 3: skill-gen (thinking ON), per-view prompt; re-sample empties. - flat: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] - pending = list(hard) - for _ in range(args.skill_retries + 1): - if not pending: - break - sg_out = _run_samples(skill_sampler, [_view_prompt(r) for r in pending], - args.n_skills, args.skill_max_tokens, skill_dp, - temperature=args.skill_gen_temperature, - top_p=args.skill_gen_top_p, top_k=args.skill_gen_top_k) - still = [] - for r, seqs in zip(pending, sg_out): - got = False - for s in seqs: - resp = _clean_text(getattr(s, 'decoded', '') or '') - block = _extract_skills_block(resp) - cand = {'skills': block or '', 'response': resp, 'parseable': bool(block), - 'view': r['_view'], 'leaked': None, 'leak_reason': '', - 'leak_source': '', 'with_pass': None, 'reward': None, 'rolls': []} - r['_cands'].append(cand) - if block: - flat.append((r, cand)) - got = True - if not got: - still.append(r) - pending = still - - # Phase 4: leak filter (view A only; view B is query-only -> treated clean, SEAM-like). - for r, c in flat: - if r.get('_view') != 'A': - c['leaked'], c['leak_reason'], c['leak_source'] = False, '', 'skipped_viewB' - flat_a = [(r, c) for r, c in flat if r.get('_view') == 'A'] - if flat_a: - details = leak.leak_batch( - [{'content': c['skills'], 'query': r['problem'], 'reference': r['reference_answer']} - for r, c in flat_a], max_workers=args.leak_workers) - for (r, c), d in zip(flat_a, details): - c['leaked'], c['leak_reason'], c['leak_source'] = bool(d.leaked), d.reason, d.source - - # Phase 5: with-skill greedy pass (T=0, M=1) on clean candidates. Reward = correct, - # absolute (no baseline subtraction); the group mean in Phase 6 is the only baseline. - clean = [(r, c) for r, c in flat if c['leaked'] is False] - if clean: - ws_out = _run_samples(base_sampler, - [build_skill_solve_prompt(r['problem'], c['skills']) for r, c in clean], - 1, args.max_tokens, base_dp, temperature=0.0) - for (r, c), seqs in zip(clean, ws_out): - c['rolls'] = [_parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll()] - c['with_pass'] = 1.0 if c['rolls'][0]['correct'] else 0.0 - c['reward'] = c['with_pass'] - if args.format_in_reward: # unparseable/leaked candidates score 0 and still join the group - for r in hard: - for c in r['_cands']: - if c['reward'] is None: - c['reward'] = 0.0 - - # Phase 6: group-relative GRPO advantage. - _assign_advantages(hard, args) - return ([_full_record(r, ci) for r in chunk], _chunk_summary(chunk, ci, args), - _group_records(chunk, args)) - - -def _roll(x: Dict[str, Any]) -> Dict[str, Any]: - return {k: x[k] for k in ('pred', 'correct', 'terminated', 'passed', - 'stop_reason', 'gen_tokens', 'text')} - - -def _is_trainable(c: Dict[str, Any], args: argparse.Namespace) -> bool: - """A candidate reaches the GRPO update iff its advantage is non-zero (and, without - --format-in-reward, is also clean and scored).""" - adv_nz = abs(c.get('advantage') or 0.0) > 1e-9 - if args.format_in_reward: - return adv_nz - return c['leaked'] is False and c.get('with_pass') is not None and adv_nz - - -def _full_record(r: Dict[str, Any], ci: int) -> Dict[str, Any]: - """Complete per-problem trace: init attempt, baseline, and all candidates.""" - init = r['_init'][0] - return { - 'record_type': 'problem', 'chunk': ci, 'problem': r['problem'], - 'reference_answer': r['reference_answer'], 'level': r.get('level', ''), - 'failed_first_try': r['_failed'], - 'init_attempt': {'text': init['text'], 'pred': init['pred'], 'correct': init['correct'], - 'terminated': init['terminated'], 'stop_reason': init['stop_reason'], - 'gen_tokens': init['gen_tokens']}, - 'baseline_pass': r['_baseline_pass'], 'is_hard': r['_hard'], - 'view': r.get('_view', ''), 'rubric_diag': r.get('_rubric_diag', ''), - 'baseline_rolls': [_roll(x) for x in r['_baseline_rolls']], - 'candidates': [{ - 'skills': c['skills'], 'response': c['response'], 'parseable': c['parseable'], - 'leaked': c['leaked'], 'leak_reason': c['leak_reason'], 'leak_source': c['leak_source'], - 'with_pass': c['with_pass'], 'reward': c.get('reward'), 'advantage': c.get('advantage'), - 'grpo_adv': c.get('grpo_adv'), 'kept': c.get('kept'), - 'rolls': [_roll(x) for x in c['rolls']], - } for c in r['_cands']], - } - - -def _view_stats(hard: List[Dict[str, Any]], view: str) -> Dict[str, Any]: - hv = [r for r in hard if r.get('_view') == view] - cands = [c for r in hv for c in r['_cands'] if c['parseable']] - clean = [c for c in cands if c['leaked'] is False] - adopted = sum(1 for r in hv - if any(c['leaked'] is False and abs(c.get('advantage') or 0.0) > 1e-9 - for c in r['_cands'])) - return {'n_hard': len(hv), 'n_candidates_parseable': len(cands), 'n_clean': len(clean), - 'n_adopted_problems': adopted, 'adoption_rate': (adopted / len(hv)) if hv else 0.0} - - -def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespace) -> Dict[str, Any]: - hard = [r for r in chunk if r['_hard']] - all_cands = [c for r in chunk for c in r['_cands']] - cands = [c for c in all_cands if c['parseable']] - scored = [c for c in cands if c['with_pass'] is not None] - ws_rolls = [x for c in scored for x in c['rolls']] - train_cands = [c for c in all_cands if _is_trainable(c, args)] - fail_cands = [c for r in chunk if r['_failed'] for c in r['_cands']] - base_acc = (sum(r['_baseline_pass'] for r in hard) / len(hard)) if hard else 0.0 - ws_acc = (sum(c['with_pass'] for c in scored) / len(scored)) if scored else 0.0 - abs_adv = lambda cs: sum(abs(c.get('advantage') or 0.0) for c in cs) - total_abs = abs_adv(all_cands) - return { - 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), - 'n_failed_first_try': sum(1 for r in chunk if r['_failed']), 'n_hard': len(hard), - 'n_generated': len(all_cands), 'n_candidates_parseable': len(cands), - 'n_unparseable': len(all_cands) - len(cands), - 'n_leaked': sum(1 for c in cands if c['leaked']), - 'n_clean': sum(1 for c in cands if c['leaked'] is False), - 'n_reward_pos': sum(1 for c in scored if c['reward']), 'n_train_samples': len(train_cands), - 'n_train_from_fail': sum(1 for c in fail_cands if _is_trainable(c, args)), - 'abs_adv_from_fail_frac': (abs_adv(fail_cands) / total_abs) if total_abs > 0 else 0.0, - 'avg_baseline_pass_on_hard': base_acc, 'avg_withskill_pass': ws_acc, - 'avg_lift': ws_acc - base_acc, - 'termination_rate_withskill': (sum(1 for x in ws_rolls if x['terminated']) / len(ws_rolls)) if ws_rolls else 0.0, - 'view_A': _view_stats(hard, 'A'), 'view_B': _view_stats(hard, 'B'), - } - - -def _group_records(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> List[Dict[str, Any]]: - """GRPO training records: every trainable candidate with its view + rubric diagnosis - (the prompt is rebuilt from those by ``_skillgen_messages``, no trajectory stored).""" - out = [] - for r in chunk: - if not r['_hard']: - continue - for c in r['_cands']: - if _is_trainable(c, args): - out.append({ - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'view': r.get('_view', 'A'), 'diagnosis': r.get('_rubric_diag', ''), - 'response': c['response'], 'skills': c['skills'], - 'advantage': c['advantage'], 'grpo_adv': c['grpo_adv'], 'kept': c['kept'], - 'reward': c['reward'], 'with_pass': c['with_pass']}) - return out - - -# =========================================================================== -# Block F -- samplers, args, main -# =========================================================================== -def init_samplers(args: argparse.Namespace): - """8 GPUs: ranks 0-3 skill_sampler, ranks 4-7 base_sampler (both vLLM tp1 dp4).""" - twinkle.initialize(mode='ray', nproc_per_node=8, lazy_collect=False, groups=[ - DeviceGroup(name='skill_sampler', ranks=list(range(0, 4)), device_type='GPU'), - DeviceGroup(name='base_sampler', ranks=list(range(4, 8)), device_type='GPU')]) - samplers = [] - for group in ('skill_sampler', 'base_sampler'): - s = vLLMSampler( - model_id=MODEL_ID, - engine_args={'gpu_memory_utilization': GPU_MEM, - 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, - device_mesh=DeviceMesh.from_sizes(world_size=4, dp_size=4), remote_group=group) - s.set_template('Template', model_id=MODEL_ID, enable_thinking=True, - max_length=args.max_model_len) - samplers.append(s) - return samplers[1], samplers[0], 4, 4 # base_sampler, skill_sampler, base_dp, skill_dp - - -def _build_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--total-problems', type=int, default=3200, - help='Final number of problems selected into generated chunks.') - p.add_argument('--base-success-frac', type=float, default=0.3, - help='Target fraction of selected problems the frozen base solves.') - p.add_argument('--output-dir', default='./output/reflexion_skill_data') - p.add_argument('--cache-dir', default='', - help='Baseline/rubric cache dir (default /cache).') - p.add_argument('--no-cache', action='store_true', help='Disable disk cache read/write.') - p.add_argument('--overwrite', action='store_true', help='Replace existing output jsonl.') - p.add_argument('--seed', type=int, default=42) - p.add_argument('--dataset', choices=('aops', 'math'), default='aops') - p.add_argument('--n', type=int, default=0, - help='Raw train-pool size; 0 derives it from --total-problems.') - p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) - p.add_argument('--eval-size', type=int, default=128) - p.add_argument('--chunk-size', type=int, default=16) - p.add_argument('--balance', action=argparse.BooleanOptionalAction, default=True) - p.add_argument('--balance-loop-frac', type=float, default=0.5) - p.add_argument('--balance-max-draws-mult', type=int, default=8) - p.add_argument('--n-skills', type=int, default=8) - p.add_argument('--view-b-frac', type=float, default=0.5) - p.add_argument('--skill-retries', type=int, default=2) - p.add_argument('--skill-gen-temperature', type=float, default=1.0) - p.add_argument('--skill-gen-top-p', type=float, default=1.0) - p.add_argument('--skill-gen-top-k', type=int, default=-1) - p.add_argument('--max-model-len', type=int, default=16384) - p.add_argument('--max-tokens', type=int, default=8192) - p.add_argument('--skill-max-tokens', type=int, default=8192) - p.add_argument('--leak-workers', type=int, default=16) - p.add_argument('--rubric-workers', type=int, default=16) - p.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True) - args = p.parse_args() - if args.total_problems <= 0 or args.chunk_size <= 0: - raise ValueError('--total-problems and --chunk-size must be positive') - if not 0.0 <= args.base_success_frac <= 1.0: - raise ValueError('--base-success-frac must be in [0, 1]') - args.chunks = math.ceil(args.total_problems / args.chunk_size) - args.balance_success_frac = args.base_success_frac - if args.n <= 0: - args.n = max(args.total_problems + args.eval_size, math.ceil(args.total_problems * 1.5)) - return args - - -def _write(handle, row: Dict[str, Any]) -> None: - handle.write(json.dumps(row, ensure_ascii=False) + '\n') - - -def main() -> None: - args = _build_args() - records, eval_records, data_stats = _load_records(args) - if not records: - raise ValueError(f'loaded 0 {args.dataset} problems') - if len(records) < args.chunk_size: - raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded ({len(records)})') - - os.makedirs(args.output_dir, exist_ok=True) - data_path = os.path.join(args.output_dir, 'skill_dataset.jsonl') - gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') - eval_path = os.path.join(args.output_dir, 'eval_holdout.jsonl') - for path in (data_path, gen_path, eval_path): - if os.path.exists(path) and not args.overwrite: - raise FileExistsError(f'{path} exists; pass --overwrite to replace it') - - if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY')): - sys.stderr.write('[build] WARNING: no LLM backup env; leak/rubric checks degrade\n') - - base_sampler, skill_sampler, base_dp, skill_dp = init_samplers(args) - leak = LeakVerifier(sampler=None, answer_only=True) - checker = build_rubric_checker() - pool = ProblemPool(records, args.seed) - - cache_dir = args.cache_dir or os.path.join(args.output_dir, 'cache') - os.makedirs(cache_dir, exist_ok=True) - use_cache = not args.no_cache - baseline_cache = DiskCache(os.path.join(cache_dir, 'baseline.jsonl'), use_cache) - rubric_cache = DiskCache(os.path.join(cache_dir, 'rubric.jsonl'), use_cache) - - cfg = { - 'record_type': 'config', 'mode': 'offline_data_build', 'model': MODEL_ID, - 'dataset': args.dataset, 'n': len(records), 'eval_n': len(eval_records), - 'total_problems': args.total_problems, 'seed': args.seed, 'numeric_only': args.numeric_only, - 'raw_loaded': data_stats['raw_loaded'], 'numeric_dropped': data_stats['numeric_dropped'], - 'chunks': args.chunks, 'chunk_size': args.chunk_size, 'n_skills': args.n_skills, - 'view_b_frac': args.view_b_frac, 'balance': args.balance, - 'base_success_frac': args.base_success_frac, 'balance_success_frac': args.balance_success_frac, - 'reward': 'greedy_binary(correct)', 'advantage': 'group_relative', - 'format_in_reward': args.format_in_reward, 'cache': use_cache, - 'rubric_check': 'fixed_math_5crit(viewA)' if checker else 'disabled', - 'started': int(time.time()), - } - total_groups, selected = 0, 0 - with open(gen_path, 'w', encoding='utf-8') as gen_f, \ - open(data_path, 'w', encoding='utf-8') as data_f, \ - open(eval_path, 'w', encoding='utf-8') as eval_f: - for handle in (gen_f, data_f, eval_f): - _write(handle, cfg) - for rec in eval_records: - _write(eval_f, {'record_type': 'eval_holdout', **rec}) - eval_f.flush() - - full_chunk_size = args.chunk_size - for ci in range(args.chunks): - remaining = args.total_problems - selected - if remaining <= 0: - break - args.chunk_size = min(full_chunk_size, remaining) # last chunk may be short - chunk, balance = draw_chunk(pool, base_sampler, base_dp, args, baseline_cache) - full, summary, groups = process_chunk( - base_sampler, skill_sampler, leak, chunk, ci, base_dp, skill_dp, - args, checker, rubric_cache) - summary['balance'] = balance - for rec in full: - _write(gen_f, rec) - _write(gen_f, summary) - gen_f.flush() - for row in groups: - _write(data_f, {'chunk': ci, **row}) - data_f.flush() - total_groups += len(groups) - selected += len(chunk) - sys.stderr.write( - f'[build] g{ci}: problems={selected}/{args.total_problems} ' - f'train={len(groups)} total={total_groups} ' - f'acc={summary["avg_baseline_pass_on_hard"]:.2f}->{summary["avg_withskill_pass"]:.2f} ' - f'lift={summary["avg_lift"]:+.3f}\n') - - baseline_cache.close() - rubric_cache.close() - sys.stderr.write(f'[build] done: {total_groups} train records -> {data_path}; trace -> {gen_path}\n') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/legacy/build_thinking_rag_index.py b/cookbook/exp/legacy/build_thinking_rag_index.py deleted file mode 100644 index a71bae060..000000000 --- a/cookbook/exp/legacy/build_thinking_rag_index.py +++ /dev/null @@ -1,1159 +0,0 @@ -"""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 ``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 - ``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 -import threading -import time -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 - -# --------------------------------------------------------------------------- -# Compress prompts — MUST match train_embedding_full_ddp.py exactly. -# --------------------------------------------------------------------------- -_HERE = Path(__file__).resolve().parent -sys.path.insert(0, str(_HERE)) - -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 - -_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', - 'output/embedding_full_transformers/last-checkpoint', -) -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)) - -# Dataset mix caps (only used in 'both' mode). None = no cap. -THINK_CAP: Optional[int] = int(os.environ.get('THINK_CAP', 400_000)) or None -INDEX_CAP: Optional[int] = int(os.environ.get('INDEX_CAP', 400_000)) or None -MIX_SHUFFLE_SEED = 100 - -# Concurrency knobs for API fallback and prefetch pipeline. -API_CONCURRENCY = int(os.environ.get('API_CONCURRENCY', 8)) -API_MIN_INTERVAL = float(os.environ.get('API_MIN_INTERVAL', 0.1)) -PREFETCH_WORKERS = int(os.environ.get('PREFETCH_WORKERS', 2)) - -# Hard-templated hints: the condenser SFT prior maps `Skill` to the legacy -# `Use when: / numbered steps / Output:` skeleton on long inputs; embedding the -# exact 4-line body template + explicit negative constraints is the only way to -# override it deterministically across query and cot sides. -RAG_QUERY_HINT = ( - '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.') -RAG_THINKING_HINT = ( - '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.') - -# 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 -# =========================================================================== - -_LEGACY_USE_WHEN_RE = re.compile(r'(?im)^\s*Use when\s*:') -_SCHEMA_MARKERS = ('Problem:', 'Skill:', 'Knowledge:') - - -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 - 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 - 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 _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) - - -_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. - - API fallback runs concurrently (up to API_CONCURRENCY workers) for speed. - """ - pairs = _vllm_compress(sampler, texts, query_hint) - 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[i] = text - else: - fallback_indices.append(i) - - if fallback_indices and api is not None: - from concurrent.futures import ThreadPoolExecutor, as_completed - with ThreadPoolExecutor(max_workers=API_CONCURRENCY) as pool: - futures = {} - for idx in fallback_indices: - msgs = _build_compress_messages(texts[idx], query_hint) - futures[pool.submit(_api_compress_throttled, api, msgs)] = idx - for fut in as_completed(futures): - idx = futures[fut] - api_text = fut.result() - if api_text and not _is_truncated_compression(api_text): - results[idx] = api_text - - return results - - -def _resolve_compressed_multi(sampler: vLLMSampler, api: Optional[OpenAIClient], - texts: List[str], hints: List[str]) -> List[Optional[str]]: - """Like _resolve_compressed but each text has its own per-item hint. - - Merges all texts into a SINGLE vLLM batch call (instead of one per hint), - dramatically reducing round-trip overhead when processing interleaved - query+cot pairs with different hint strings. - - Args: - sampler: vLLM condenser sampler. - api: Optional OpenAI-compatible API client for fallback. - texts: List of raw texts to compress (may contain empty strings to skip). - hints: Per-text hint strings (same length as texts). - - Returns: - List of compressed texts (None where compression failed entirely). - """ - assert len(texts) == len(hints), f'texts({len(texts)}) != hints({len(hints)})' - if not texts: - return [] - - # Skip texts that would exceed the condenser's context window. - _max_input_chars = (CONDENSE_MAX_MODEL_LEN - CONDENSE_MAX_TOKENS) * 3 - skip_mask = [len(t) > _max_input_chars for t in texts] - - # Build prompts per-item (each text gets its own hint as the query parameter). - prompts = [{'messages': _build_compress_messages(t, h)} - for t, h, skip in zip(texts, hints, skip_mask) if not skip] - active_indices = [i for i, skip in enumerate(skip_mask) if not skip] - params = TwinkleSamplingParams( - max_tokens=CONDENSE_MAX_TOKENS, - temperature=COMPRESS_TEMPERATURE, - top_p=COMPRESS_TOP_P, - num_samples=1, - ) - - # Single vLLM batch call — the key throughput win. - responses = sampler.sample(prompts, params) if prompts else [] - - results: List[Optional[str]] = [None] * len(texts) - fallback_indices: List[int] = [] - for resp_idx, orig_idx in enumerate(active_indices): - resp = responses[resp_idx] - seq = resp.sequences[0] if resp and resp.sequences else None - if seq is None: - fallback_indices.append(orig_idx) - continue - text = seq.decoded or '' - text = re.sub(r'<\|[^|]+\|>', '', text).rstrip() - text = _strip_outer_codefence(text) - if seq.stop_reason != 'length' and not _is_truncated_compression(text): - results[orig_idx] = text - else: - fallback_indices.append(orig_idx) - - # Concurrent API fallback for failed items. - if fallback_indices and api is not None: - from concurrent.futures import ThreadPoolExecutor, as_completed - with ThreadPoolExecutor(max_workers=API_CONCURRENCY) as pool: - futures = {} - for idx in fallback_indices: - msgs = _build_compress_messages(texts[idx], hints[idx]) - futures[pool.submit(_api_compress_throttled, api, msgs)] = idx - for fut in as_completed(futures): - idx = futures[fut] - api_text = fut.result() - if api_text and not _is_truncated_compression(api_text): - results[idx] = api_text - - return results - - -# =========================================================================== -# 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 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' - + (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 - 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 _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 [] - # 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 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: - 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), - 'cot_raw_head': _short(r['cot_raw'], 200), - }) - continue - 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 - 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) - 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: - 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), - '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) - 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 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: - 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, - 'source': row.get('source') or 'unknown', - 'query_raw': user_query, - 'cot_raw': cot, - }) - if len(batch) >= args.batch_size: - pending_futures.append(prefetch_pool.submit(_process_batch, list(batch))) - batch.clear() - - # Flush remainder - if 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] 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: - 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], RAG_QUERY_HINT) - 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='both', - choices=['dataset_index', 'dataset_think', 'both'], - help='Which loader to use: dataset_index (RAG profile), ' - 'dataset_think (training mix), or both (50/50 mix).') - p.add_argument('--limit', type=int, default=0, - help='Stop building once this many rows are kept (0 = no cap).') - p.add_argument('--max-rows', type=int, default=0, - help='Truncate corpus to this many rows AFTER get_dataset (0 = no cap). ' - 'Use this instead of --total to avoid invalidating the dataset cache.') - p.add_argument('--batch-size', type=int, default=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', - 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 - 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. - 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/legacy/cold_start/train_cold_start.py b/cookbook/exp/legacy/cold_start/train_cold_start.py deleted file mode 100644 index 9c662e89d..000000000 --- a/cookbook/exp/legacy/cold_start/train_cold_start.py +++ /dev/null @@ -1,333 +0,0 @@ -import json -import os -from functools import partial -from pathlib import Path -from typing import Any, Dict, Iterator, List - -from peft import LoraConfig - -import twinkle -from twinkle import DeviceMesh, DeviceGroup, get_device_placement, get_logger -from twinkle.dataloader import DataLoader -from twinkle.dataset import Dataset, PackingDataset -from twinkle.dataset.base import DatasetMeta -from twinkle.model import MegatronModel -from twinkle_agentic.preprocessor import ( - QualityPreprocessor, - IntentClassifier, HardFilter, RefuseFilter, DeadLoopFilter, TokenSoupFilter, MessageSanityFilter, - SpecialCharsFilter, ModelFilter, DedupFilter, - MessageNormalizer, -) -from twinkle_agentic.preprocessor.experimental import SamplerBackend # noqa: F401 - -logger = get_logger() - -# ── Model ──────────────────────────────────────────────────────────────────── -MODEL_ID = 'ms://Qwen/Qwen3-4B' -TEMPLATE_NAME = 'Template' -MAX_LENGTH = 80000 - -# ── GPU allocation ─────────────────────────────────────────────────────────── -MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 8)) -SAMPLER_GPUS = int(os.environ.get('SAMPLER_GPUS', 0)) -NUM_GPUS = MODEL_GPUS + SAMPLER_GPUS - -# ── Training ───────────────────────────────────────────────────────────────── -BATCH_SIZE = int(os.environ.get('BATCH_SIZE', 1)) -LEARNING_RATE = float(os.environ.get('LR', 1e-5)) -GRADIENT_ACCUMULATION_STEPS = int(os.environ.get('GRAD_ACCUM', 4)) -LOG_INTERVAL = 1 -SAVE_INTERVAL = 500 -NUM_STEPS = int(os.environ.get('NUM_STEPS', 5000)) - -# ── Output ─────────────────────────────────────────────────────────────────── -OUTPUT_DIR = './output/streaming_sft' -TRAINED_DATA_PATH = os.path.join(OUTPUT_DIR, 'trained_data.jsonl') -DROPPED_DATA_PATH = os.path.join(OUTPUT_DIR, 'dropped_data.jsonl') -ADAPTER_NAME = 'default' - -# ── Data source ────────────────────────────────────────────────────────────── -CSV_PATH = os.environ.get('CSV_PATH') -DATASET_TOTAL = int(os.environ.get('DATASET_TOTAL', 10000)) # 0 = full materialized dataset -# Worker count for HF Dataset.map(num_proc=N); spawn start method is forced in twinkle.dataset.base. -MAP_NUM_PROC = int(os.environ.get('MAP_NUM_PROC', 16)) - - -def _canonicalize_tool_call(tc: Any) -> Dict[str, Any]: - """Coerce ``tool_calls[i]`` to a fixed-schema dict for stable Arrow inference. - - Keeps ``function.arguments`` as the OpenAI-native JSON string so every row - sees a uniform ``string`` field; any string→dict decoding is the - chat_template's concern (see ``Template._apply_chat_template``). - - The decoded form is enforced to be a JSON object so the chat_template's - ``|items`` filter never receives list/scalar/null — those originate from - dirty CSV rows and are coerced to ``{}`` here, the ingestion boundary. - """ - tc = tc if isinstance(tc, dict) else {} - fn = tc.get('function') if isinstance(tc.get('function'), dict) else {} - args = fn.get('arguments') - if isinstance(args, dict): - args_str = json.dumps(args, ensure_ascii=False) - elif isinstance(args, str) and args.strip(): - try: - decoded = json.loads(args) - except json.JSONDecodeError: - decoded = {} - if not isinstance(decoded, dict): - decoded = {} - args_str = json.dumps(decoded, ensure_ascii=False) - else: - args_str = '{}' - return { - 'id': str(tc.get('id') or ''), - 'type': str(tc.get('type') or 'function'), - 'function': { - 'name': str(fn.get('name') or ''), - 'arguments': args_str, - }, - } - - -def _stream_csv_rows(csv_path: str, max_rows: int = 0) -> Iterator[Dict[str, Any]]: - """Stream the custom CSV: each line is `ts,model,req_id,messages_json` (no quoting). - - The first 3 fields are scalar; the remainder of the line is a JSON array of - chat messages, possibly containing commas — so we split on the first 3 commas only. - ``max_rows`` caps the yielded rows at ingestion time so Arrow never materializes - the unused tail. - """ - emitted = 0 - with open(csv_path, 'rb') as f: - bad_bytes = 0 - for raw in f: - try: - line = raw.decode('utf-8').rstrip('\n').rstrip('\r') - except UnicodeDecodeError: - bad_bytes += 1 - continue - if not line: - continue - parts = line.split(',', 3) - if len(parts) < 4: - continue - ts, _model, req_id, msgs_raw = parts - try: - raw_msgs = json.loads(msgs_raw) - except json.JSONDecodeError: - continue - messages: List[Dict[str, Any]] = [] - for m in raw_msgs: - role = m.get('role', '') - content = m.get('content') - # User content arrives as [{'type':'text','text':...}, ...]; flatten to plain string. - if isinstance(content, list): - content = ''.join( - p.get('text', '') for p in content - if isinstance(p, dict) and p.get('type') == 'text') - if content is None: - content = '' - if not isinstance(content, str): - continue - raw_tcs = m.get('tool_calls') if role == 'assistant' else None - tc_list = [_canonicalize_tool_call(tc) for tc in raw_tcs] if raw_tcs else [] - if role == 'assistant': - if not content and not tc_list: - continue - if m.get('reasoning_content'): - content = f"{m['reasoning_content']}{content}" - elif role == 'tool': - pass - elif not content: - continue - # tool_calls stored as JSON string (empty -> ''): keeps Arrow schema as a - # stable Value(string) regardless of empty-list / heterogeneous-struct shards. - # Template._apply_chat_template decodes it back to list before jinja render. - messages.append({ - 'role': role, - 'content': content, - 'tool_calls': json.dumps(tc_list, ensure_ascii=False) if tc_list else '', - 'tool_call_id': str(m.get('tool_call_id') or '') if role == 'tool' else '', - }) - if not messages: - continue - yield { - 'id': f'csv__{ts}__{req_id}', - 'source': Path(csv_path).stem, - 'model_id': _model, - 'messages': messages, - 'user_data': [], - } - emitted += 1 - if max_rows and emitted >= max_rows: - break - - -# ── QualityPreprocessor config ─────────────────────────────────────────────── -SENSITIVE_WORDS_FILE = str( - Path(__file__).resolve().parent.parent.parent / 'sensitive_words.txt') -# chr_min cutoff: keep round if chr_min < threshold (low chr_min = hard). -CHR_MIN_THRESHOLD = float(os.environ.get('CHR_MIN_THRESHOLD', 0.5)) -REFINE_TEMPERATURE = float(os.environ.get('REFINE_TEMPERATURE', 0.6)) -REFINE_MAX_TOKENS = int(os.environ.get('REFINE_MAX_TOKENS', 4096)) - -# ── Pass@4 LLM-as-judge (grades each diagnostic rollout vs GT) ─────────────── -# Set JUDGE_MODEL='' to disable; otherwise judge runs over every diagnostic round. -JUDGE_MODEL = os.environ.get('JUDGE_MODEL', 'qwen3.7-max') -JUDGE_BASE_URL = os.environ.get('JUDGE_BASE_URL', 'https://dashscope.aliyuncs.com/compatible-mode/v1') -JUDGE_API_KEY = os.environ.get('JUDGE_API_KEY', 'EMPTY') -JUDGE_TEMPERATURE = float(os.environ.get('JUDGE_TEMPERATURE', 0.3)) -JUDGE_MAX_TOKENS = int(os.environ.get('JUDGE_MAX_TOKENS', 32000)) -JUDGE_MAX_WORKERS = int(os.environ.get('JUDGE_MAX_WORKERS', 16)) - - -def build_dataset(backend: SamplerBackend) -> Dataset: - """Materialize the local CSV, convert to SFT messages format, run QualityPreprocessor. - - Switched from streaming IterableDataset to in-memory Dataset so HF - `Dataset.map(num_proc=N)` can parallelize the QualityPreprocessor pipeline. - """ - os.makedirs(OUTPUT_DIR, exist_ok=True) - - # Custom CSV format (commas inside JSON) — feed framework via callable, not csv loader. - meta = DatasetMeta( - dataset_id=Path(CSV_PATH).stem, - data=partial(_stream_csv_rows, csv_path=CSV_PATH, max_rows=DATASET_TOTAL), - ) - dataset = PackingDataset(meta) - - qp = QualityPreprocessor( - pipeline=[ - ModelFilter(), - MessageNormalizer(), - HardFilter( - min_user_chars_cjk=14, min_user_chars=24, - system_deny_keywords=[ - '角色扮演', '扮演', '人设', 'roleplay', 'role play', 'cosplay', - '群聊模拟', '虚拟角色', '二次元', 'OC设定', - ], - max_rounds=30, - ), - RefuseFilter(), - DeadLoopFilter(), - MessageSanityFilter(sensitive_words_file='.temp/sensitive_words.txt'), - SpecialCharsFilter(max_ratio=0.6), - TokenSoupFilter(max_chars=8000), - IntentClassifier(), - # ScoreFilter( - # template=template, - # backend=backend, - # scorers=[ - # ChrMinScorer(), - # ], - # ), - # PIIPresidioFilter(languages=('en', 'zh')), - ], - dropped_log_path=DROPPED_DATA_PATH, - ) - dataset.map(qp, num_proc=8, load_from_cache_file=True) - dataset.map( - QualityPreprocessor(pipeline=[DedupFilter()]), - num_proc=1, - batch_size=len(dataset.dataset), - load_from_cache_file=True, - ) - - print(len(dataset.dataset)) - dataset.set_template( - TEMPLATE_NAME, - model_id=MODEL_ID, - max_length=MAX_LENGTH, - truncation_strategy='delete', - enable_thinking=False, - ) - dataset.encode(num_proc=16, load_from_cache_file=True) - dataset.pack_dataset() - return dataset - - -def save_checkpoint(model: MegatronModel, checkpoint_name: str, dataloader: DataLoader): - model.save( - checkpoint_name, - output_dir=OUTPUT_DIR, - adapter_name=ADAPTER_NAME, - save_optimizer=True, - consumed_train_samples=dataloader.get_state()['consumed_train_samples'], - ) - - -def train(): - # ── Ray mode: GPUs 0-3 for training, GPUs 4-7 for vLLMSampler ──────────── - device_groups = [ - DeviceGroup(name='model', ranks=list(range(MODEL_GPUS)), device_type='GPU'), - # DeviceGroup(name='sampler', ranks=list(range(MODEL_GPUS, NUM_GPUS)), device_type='GPU', gpus_per_worker=2), - ] - model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, dp_size=1, cp_size=8) - # sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, dp_size=SAMPLER_GPUS // 2, tp_size=2) - twinkle.initialize(mode='local', nproc_per_node=NUM_GPUS, groups=device_groups, - global_device_mesh=model_mesh, lazy_collect=False) - - # ── vLLMSampler on GPUs 4-7 (Ray actor, no HTTP overhead) ──────────────── - # sampler = vLLMSampler( - # model_id=MODEL_ID, - # engine_args={ - # 'gpu_memory_utilization': 0.6, - # 'max_model_len': MAX_LENGTH, - # }, - # device_mesh=sampler_mesh, - # remote_group='sampler', - # ) - # sampler.set_template(TEMPLATE_NAME, model_id=MODEL_ID) - # backend = SamplerBackend(sampler) - # logger.info(f'vLLMSampler ready on GPUs {MODEL_GPUS}-{NUM_GPUS - 1}') - - # ── Dataset with full QualityPreprocessor (uses SamplerBackend) ─────────── - dataset = build_dataset(None) - dataloader = DataLoader( - dataset=dataset, - batch_size=BATCH_SIZE, - ) - - # ── Model (LoRA on 4 GPUs) ──────────────────────────────────────────────── - model = MegatronModel( - model_id=MODEL_ID, - device_mesh=model_mesh, - # remote_group='model', - # attn_implementation='flash_attention_2', - ) - - lora_config = LoraConfig(r=16, lora_alpha=32, target_modules='all-linear') - model.add_adapter_to_model( - ADAPTER_NAME, lora_config, - gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) - model.set_optimizer(optimizer_cls='default', lr=LEARNING_RATE) - model.set_lr_scheduler( - scheduler_cls='default', - lr_warmup_steps=2, - lr_decay_steps=len(dataloader)) - - logger.info(get_device_placement()) - logger.info(model.get_train_configs()) - logger.info(f'Total steps: {NUM_STEPS}, model GPUs: {MODEL_GPUS}, sampler GPUs: {SAMPLER_GPUS}') - - for cur_step, batch in enumerate(dataloader): - model.forward_backward(inputs=batch) - model.clip_grad_and_step() - - if cur_step % LOG_INTERVAL == 0: - metric = model.calculate_metric(is_training=True) - logger.info(f'Step {cur_step}/{NUM_STEPS}, metric: {metric}') - - if cur_step % SAVE_INTERVAL == 0: - save_checkpoint(model, f'step-{cur_step}', dataloader) - - if cur_step >= NUM_STEPS: - break - - save_checkpoint(model, 'last-checkpoint', dataloader) - logger.info(f'Training complete. Trained data saved to: {TRAINED_DATA_PATH}') - logger.info(f'Dropped data saved to: {DROPPED_DATA_PATH}') - - -if __name__ == '__main__': - train() diff --git a/cookbook/exp/legacy/compare_math_levels.py b/cookbook/exp/legacy/compare_math_levels.py deleted file mode 100644 index d488909b9..000000000 --- a/cookbook/exp/legacy/compare_math_levels.py +++ /dev/null @@ -1,91 +0,0 @@ -"""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/legacy/condenser/dataset.py b/cookbook/exp/legacy/condenser/dataset.py deleted file mode 100644 index 32c30de4b..000000000 --- a/cookbook/exp/legacy/condenser/dataset.py +++ /dev/null @@ -1,459 +0,0 @@ -import hashlib -import json -import os -import re -from pathlib import Path -from typing import Any, Dict, List, Optional -from datasets import Features, Value -from modelscope import dataset_snapshot_download - -from twinkle.dataset import Dataset, DatasetMeta -from twinkle.preprocessor import Preprocessor - -_TARGET_FEATURES = Features({ - 'id': Value('string'), - 'source': Value('string'), - 'messages': [{'role': Value('string'), 'content': Value('string')}], -}) - - -def _hash_id(prefix: str, content: str) -> str: - """Stable id from MD5 of content; collision-free for textual datasets.""" - return f'{prefix}__{hashlib.md5(content.encode("utf-8")).hexdigest()[:16]}' - - -def _register(dataset, processor_cls, meta: DatasetMeta, init_args: Optional[Dict[str, Any]] = None, - load_from_cache_file: bool = True) -> None: - """Add dataset and run preprocessor; auto-strip every input column to enforce - the universal ``{id, source, messages}`` output schema.""" - dataset.add_dataset(meta) - cols = list(dataset.datasets[meta.get_id()].column_names) - dataset.map( - processor_cls, - dataset_meta=meta, - init_args=init_args or {}, - remove_columns=cols, - load_from_cache_file=load_from_cache_file, - features=_TARGET_FEATURES, - ) - - -# ===== MuSiQue ===== -MUSIQUE_REPO = 'voidful/MuSiQue' - - -class MusiqueProcessor(Preprocessor): - """MuSiQue raw row → multiple ``{id, source, messages}`` rows, one per paragraph.""" - - 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: - if row.get('answerable') is False: - continue - parent = str(row.get('id', '')) - for idx, p in enumerate(row.get('paragraphs') or []): - text = (p.get('paragraph_text') or '').strip() - if not text: - continue - out.append({ - 'id': f'musique__{parent}__{idx}', - 'source': 'musique', - 'messages': [{'role': 'assistant', 'content': text}], - }) - return self.map_row_to_col(out, keys=['id', 'source', 'messages']) - - -# Repo 仅含原始 JSONL 无 HF 元数据，必须先快照下载再以文件路径注册。 -_musique_jsonl = Path(dataset_snapshot_download(MUSIQUE_REPO)) / 'musique_ans_v1.0_train.jsonl' -if not _musique_jsonl.is_file(): - raise FileNotFoundError(f'MuSiQue raw file not found: {_musique_jsonl}') - - -# ===== swift/github-code ===== -GITHUB_CODE_REPO = 'ms://swift/github-code' - - -class GithubCodeProcessor(Preprocessor): - """github-code row → ``{id, source, messages}``；按代码长度均匀采样。 - - 把 ``[length_min, length_max)`` 切 ``n_buckets`` 桶，每桶配额 ``target/n_buckets``， - 桶满或超界即丢；近似得到 ``target`` 条且长度均匀分布的样本。 - 依赖 batched map 单进程下实例状态跨 batch 共享（``num_proc>1`` 会失效）。 - """ - - def __init__(self, target: int = 30000, length_min: int = 500, - length_max: int = 40000, n_buckets: int = 30): - self.length_min = length_min - self.length_max = length_max - self.n_buckets = n_buckets - self.bucket_quota = max(1, target // n_buckets) - self.bucket_count = [0] * n_buckets - - def _bucket(self, n: int) -> int: - if n < self.length_min or n >= self.length_max: - return -1 - idx = int((n - self.length_min) / (self.length_max - self.length_min) * self.n_buckets) - return min(idx, self.n_buckets - 1) - - 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: - code = row.get('code') or '' - if not isinstance(code, str): - continue - b = self._bucket(len(code)) - if b < 0 or self.bucket_count[b] >= self.bucket_quota: - continue - self.bucket_count[b] += 1 - lang = row.get('language') or 'unknown' - out.append({ - 'id': _hash_id(f'github_code__{lang}', code), - 'source': 'github-code', - 'messages': [{'role': 'assistant', 'content': code}], - }) - return self.map_row_to_col(out, keys=['id', 'source', 'messages']) - - -# ===== modelscope/competition_math ===== -COMPETITION_MATH_REPO = 'ms://modelscope/competition_math' - - -class MathProcessor(Preprocessor): - """competition_math row → ``{id, source, messages}`` (user/assistant pair).""" - - 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: - problem = (row.get('problem') or '').strip() - solution = (row.get('solution') or '').strip() - if not problem or not solution: - continue - out.append({ - 'id': _hash_id('math', f'{problem}\n{solution}'), - 'source': 'competition_math', - 'messages': [ - {'role': 'assistant', 'content': solution}, - ], - }) - return self.map_row_to_col(out, keys=['id', 'source', 'messages']) - - -# ===== nampdn-ai/tiny-textbooks ===== -TINY_TEXTBOOKS_REPO = 'ms://AI-ModelScope/tiny-textbooks' - - -class TinyTextbooksProcessor(Preprocessor): - """tiny-textbooks row → ``{id, source, messages}`` (user/assistant pair).""" - - 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 '').strip() - textbook = (row.get('textbook') or '').strip() - if not text or not textbook: - continue - out.append({ - 'id': _hash_id('tinytb', f'{text}\n{textbook}'), - 'source': 'tiny-textbooks', - 'messages': [ - {'role': 'assistant', 'content': textbook}, - ], - }) - return self.map_row_to_col(out, keys=['id', 'source', 'messages']) - - -# ===== Passage Explosion for Compression Distillation ===== -# Each message content >= threshold becomes a standalone row: messages=[{role:user, content:X}] - -_MIN_PASSAGE_LEN = 500 # CJK-equivalent units - - -def _effective_len(text: str) -> int: - """CJK chars count double; threshold 500 ≈ 500 Chinese chars ≈ 1000 Latin chars.""" - cjk = sum(1 for c in text if '\u4e00' <= c <= '\u9fff' or '\u3000' <= c <= '\u303f') - return cjk * 2 + (len(text) - cjk) - - -def _extract_content(msg: dict) -> str: - """Extract text content from a message dict, handling multimodal list-content.""" - content = msg.get('content') - if isinstance(content, list): - content = '\n'.join( - p.get('text', '') if isinstance(p, dict) else str(p) for p in content) - if not isinstance(content, str): - return '' - return content.strip() - - -class PassageExplodeProcessor(Preprocessor): - """Explode multi-turn messages into individual long passages for compression distillation.""" - - def __init__(self, source: str): - self.source = source - - 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: - messages = row.get('messages') - if isinstance(messages, str): - try: - messages = json.loads(messages) - except (ValueError, TypeError): - continue - if not isinstance(messages, list): - continue - for msg in messages: - if not isinstance(msg, dict): - continue - role = msg.get('role') or '' - if role == 'system': - continue - content = _extract_content(msg) - if not content or _effective_len(content) < _MIN_PASSAGE_LEN: - continue - out.append({ - 'id': _hash_id(self.source, content), - 'source': self.source, - 'messages': [{'role': 'assistant', 'content': content}], - }) - return self.map_row_to_col(out, keys=['id', 'source', 'messages']) - - -# ===== Reasoning / CoT datasets — explode query and assistant separately ===== -_THINK_RE = re.compile(r'(.*?)', re.DOTALL) - - -class CotExplodeProcessor(Preprocessor): - """Base for CoT datasets: explode query and full assistant content as separate passages.""" - - def _extract_rows(self, rows: List[Dict[str, Any]]) -> List[tuple]: - """Subclass returns list of (query, cot, response) tuples.""" - raise NotImplementedError - - def __call__(self, rows: Dict[str, List[Any]]) -> Dict[str, List[Any]]: - rows_list = self.map_col_to_row(rows) - out: List[Dict[str, Any]] = [] - for query, cot, response, source in self._extract_rows(rows_list): - if cot: - response = _THINK_RE.sub('', response).strip() - assistant_content = f'{cot}{response}' if cot else response - for text in (query, assistant_content): - if not text or _effective_len(text) < _MIN_PASSAGE_LEN: - continue - out.append({ - 'id': _hash_id(source, text), - 'source': source, - 'messages': [{'role': 'assistant', 'content': text}], - }) - return self.map_row_to_col(out, keys=['id', 'source', 'messages']) - - -# -- Chinese-DeepSeek-R1-Distill-data-110k -- -CN_R1_DISTILL_REPO = 'ms://AI-ModelScope/Chinese-DeepSeek-R1-Distill-data-110k' - - -class ChineseR1DistillProcessor(CotExplodeProcessor): - """input → query, reasoning_content → cot, content → response.""" - - def _extract_rows(self, rows): - for row in rows: - query = (row.get('input') or '').strip() - cot = (row.get('reasoning_content') or '').strip() - response = (row.get('content') or '').strip() - if not query or not response: - continue - yield query, cot, response, 'Chinese-DeepSeek-R1-Distill-data-110k' - - -# -- Opus-4.6-Reasoning-3000x-filtered -- -OPUS_REASONING_REPO = 'ms://nohurry/Opus-4.6-Reasoning-3000x-filtered' - - -class OpusReasoningProcessor(CotExplodeProcessor): - """problem → query, thinking → cot, solution → response.""" - - def _extract_rows(self, rows): - for row in rows: - query = (row.get('problem') or '').strip() - cot = (row.get('thinking') or '').strip() - response = (row.get('solution') or '').strip() - if not query or not response: - continue - yield query, cot, response, 'Opus-4.6-Reasoning-3000x-filtered' - - -# -- claude-opus-4.6-10000x -- -CLAUDE_OPUS_REPO = 'ms://Roman1111111/claude-opus-4.6-10000x' - - -class ClaudeOpusProcessor(CotExplodeProcessor): - """messages (OpenAI format) → extract user/assistant, split or reasoning field.""" - - def _extract_rows(self, rows): - for row in rows: - messages = row.get('messages') - if not isinstance(messages, list): - continue - query = '' - assistant_text = '' - reasoning = '' - for msg in messages: - if not isinstance(msg, dict): - continue - role = msg.get('role') or '' - content = msg.get('content') or '' - if not isinstance(content, str): - continue - if role == 'user' and not query: - query = content.strip() - elif role == 'assistant' and not assistant_text: - assistant_text = content.strip() - reasoning = (msg.get('reasoning') or '').strip() - break - if not query or not assistant_text: - continue - cot = reasoning - if not cot: - m = _THINK_RE.search(assistant_text) - if m: - cot = m.group(1).strip() - assistant_text = assistant_text[m.end():].strip() - response = assistant_text if not reasoning else _THINK_RE.sub('', assistant_text).strip() - if not response: - continue - yield query, cot, response, 'claude-opus-4.6-10000x' - - -# -- angrygiraffe-claude-opus-4.6-4.7-reasoning-8.7k -- -ANGRYGIRAFFE_REPO = 'ms://hf/angrygiraffe-claude-opus-4.6-4.7-reasoning-8.7k' - - -class AngrygiraffeOpusReasoningProcessor(CotExplodeProcessor): - """messages (OpenAI format) → extract first user/assistant, split tag.""" - - def _extract_rows(self, rows): - for row in rows: - messages = row.get('messages') - if not isinstance(messages, list): - continue - query = '' - assistant_text = '' - for msg in messages: - if not isinstance(msg, dict): - continue - role = msg.get('role') or '' - content = msg.get('content') or '' - if not isinstance(content, str): - continue - if role == 'user' and not query: - query = content.strip() - elif role == 'assistant' and not assistant_text: - assistant_text = content.strip() - break - if not query or not assistant_text: - continue - m = _THINK_RE.search(assistant_text) - if m: - cot = m.group(1).strip() - response = assistant_text[m.end():].strip() - else: - cot = '' - response = assistant_text - if not response: - continue - yield query, cot, response, 'angrygiraffe-claude-opus-4.6-4.7-reasoning-8.7k' - - -_BASE_SIZES = { - 'tiny_textbooks': 10000, - 'musique': 1000, - 'github_code': 30000, - 'competition_math': 7500, - 'toucan': 10000, - 'swe_smith': 1000, - 'cn_r1_distill': 10000, - 'opus_reasoning': 3000, - 'claude_opus': 10000, - 'angrygiraffe': 20000, -} - - -def _scaled_sizes(total: Optional[int]) -> Dict[str, int]: - if total is None: - 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 get_dataset(total: Optional[int] = None, load_from_cache_file: bool = True) -> Dataset: - """Build the unified compression-distillation dataset. - - If ``total`` is given, every per-source row count in ``_BASE_SIZES`` is - scaled proportionally so the input-row sum approximates ``total``. - """ - sizes = _scaled_sizes(total) - dataset = Dataset() - - _register(dataset, TinyTextbooksProcessor, - DatasetMeta(dataset_id=TINY_TEXTBOOKS_REPO, split='train', - data_slice=range(sizes['tiny_textbooks'])), - load_from_cache_file=load_from_cache_file) - - _register(dataset, MusiqueProcessor, - DatasetMeta(str(_musique_jsonl), data_slice=range(sizes['musique'])), - load_from_cache_file=load_from_cache_file) - - _register(dataset, GithubCodeProcessor, - DatasetMeta(dataset_id=GITHUB_CODE_REPO, subset_name='all-apache-2.0', split='train'), - init_args={'target': sizes['github_code']}, - load_from_cache_file=load_from_cache_file) - - _register(dataset, MathProcessor, - DatasetMeta(dataset_id=COMPETITION_MATH_REPO, subset_name='default', split='train', - data_slice=range(sizes['competition_math'])), - load_from_cache_file=load_from_cache_file) - - _register(dataset, PassageExplodeProcessor, - DatasetMeta(dataset_id='ms://Agent-Ark/Toucan-1.5M', subset_name='Kimi-K2', split='train', - data_slice=range(sizes['toucan'])), - init_args={'source': 'toucan'}, - load_from_cache_file=load_from_cache_file) - - _register(dataset, PassageExplodeProcessor, - DatasetMeta(dataset_id='ms://SWE-bench/SWE-smith-trajectories', split='tool', - data_slice=range(sizes['swe_smith'])), - init_args={'source': 'swe-smith'}, - load_from_cache_file=load_from_cache_file) - - _register(dataset, ChineseR1DistillProcessor, - DatasetMeta(dataset_id=CN_R1_DISTILL_REPO, split='train', - data_slice=range(sizes['cn_r1_distill'])), - load_from_cache_file=load_from_cache_file) - - _register(dataset, OpusReasoningProcessor, - DatasetMeta(dataset_id=OPUS_REASONING_REPO, split='train', - data_slice=range(sizes['opus_reasoning'])), - load_from_cache_file=load_from_cache_file) - - _register(dataset, ClaudeOpusProcessor, - DatasetMeta(dataset_id=CLAUDE_OPUS_REPO, split='train', - data_slice=range(sizes['claude_opus'])), - load_from_cache_file=load_from_cache_file) - - _register(dataset, AngrygiraffeOpusReasoningProcessor, - DatasetMeta(dataset_id=ANGRYGIRAFFE_REPO, split='train', - data_slice=range(sizes['angrygiraffe'])), - load_from_cache_file=load_from_cache_file) - - dataset.mix_dataset(False) - return dataset - - -if __name__ == '__main__': - dataset = get_dataset(load_from_cache_file=True) - print(len(dataset)) diff --git a/cookbook/exp/legacy/condenser/make_condenser_dataset.py b/cookbook/exp/legacy/condenser/make_condenser_dataset.py deleted file mode 100644 index cf56a44e3..000000000 --- a/cookbook/exp/legacy/condenser/make_condenser_dataset.py +++ /dev/null @@ -1,737 +0,0 @@ -import argparse -import hashlib -import json -import os -import random -import re -import sys -import threading -from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait -from typing import Any, Dict, Iterator, List, Optional, Set - -from tqdm import tqdm - -from twinkle.data_format.sampling import SamplingParams -from twinkle_agentic.protocol.openai import OpenAI - - -# ═══════════════════════════════════════════════════════════════════════════════ -# Prompts -# ═══════════════════════════════════════════════════════════════════════════════ - -QUERY_GEN_SYSTEM = """\ -You are a query designer. Given a source passage, enumerate distinct information \ -queries a reader might ask of it. Each query must steer toward a meaningfully \ -DIFFERENT compression of the same source — different facets, not rephrasings of \ -the same need. - -Category hints (not exhaustive — combine or invent as fits the source): -- Interface extraction (code): class / method signatures, parameter and return types -- Functional summary: what the passage accomplishes at a high level -- Error & pitfall analysis: bugs, anti-patterns, failure modes, edge cases -- Experience distillation: lessons learned, best practices, do's and don'ts -- Skill extraction (knowledge-as-skill): WHAT this passage lets you do, HOW to \ -apply it as reusable steps, WHEN to invoke it (trigger conditions / use cases) -- Abstract analysis: design patterns, architectural decisions, trade-offs -- Information summary: key facts, entities, numbers, relationships -- Dependency & context: prerequisites, imports, environment, related modules - -Rules: -1. SHAPE — each query is one short imperative or interrogative sentence (e.g. \ -"List all public method signatures with parameter and return types", "What race \ -conditions does this code contain?"). -2. DISTINCT — reject any pair whose answers would substantially overlap; \ -rephrasings of the same information need do NOT count as separate queries. -3. SKILL FOR KNOWLEDGE — when the source reads as tutorial / experience / \ -how-to / domain knowledge, ALWAYS include exactly one skill-style query asking \ -what the reader can accomplish with it and how to apply it (phrased in the \ -source language). -4. ANSWERABLE — skip queries the source cannot actually answer, and skip \ -trivial queries that would just reproduce the source verbatim. -5. SCALE — short / single-purpose → 1; medium → 2; rich / multi-topic → 3–4. \ -Do not pad. -6. LANGUAGE — query language MUST match the source language. -7. OUTPUT — a single JSON array of strings; no preamble, no code fences, \ -nothing else.\ -""" - -QUERY_GEN_USER = 'Analyze the following text and return a JSON array of queries.\n\n{text}' - -COMPRESS_SYSTEM = """\ -You are a compression 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)`. - -Examples: - -Query: List all public method signatures with parameter and return types -Source: (a Python HTTP client class with retry decorator, structured logging, \ -and request helpers) -## Summary -Topic: Python HTTP client class — public surface of retried request helpers. -retry_request(url:str, max_retries:int=3, timeout:float=10.0) -> Response -fetch_json(endpoint:str, params:dict|None=None) -> dict -post_data(endpoint:str, payload:dict, headers:dict|None=None) -> Response - -## More -- decorators: @retry config — exponential backoff (base=2.0, max=60s) -- logging: structured per-request logs with request_id and latency_ms -- private helpers: _build_headers, _parse_error — not in public surface -─── -Query: What can this passage help you accomplish, and how to use it? -Source: (a tutorial on configuring Linux cgroups v2 caps for a systemd service) -## Summary -Topic: Linux cgroups v2 — per-service CPU / memory caps via systemd slice units. -Use when: needing per-service CPU/memory caps on systemd hosts. -1.create slice unit /etc/systemd/system/.slice with CPUQuota=, MemoryMax= -2.attach service via Slice=.slice in [Service] -3.systemctl daemon-reload + restart service -4.verify: systemctl status shows Tasks/CPU/Memory inside slice -Output: hard caps enforced by kernel cgroup v2. - -## More -- pitfalls: cgroup v1/v2 mode detection, MemorySwapMax behavior on OOM -- delegation: Delegate=yes for nested controllers in container managers -- examples: nginx and postgres slice templates with concrete numeric caps -- diagnostics: systemd-cgls / systemd-cgtop walkthrough -─── -Query: 总结这段代码的错误和改进经验 -Source: (一段有 race condition 和未关闭资源的 Go 代码) -## Summary -Topic: Go HTTP fetch 循环 — 并发写共享 map + 未关闭响应体导致的稳定性缺陷。 -1.race: 并发写 map 未锁 → sync.RWMutex 或 sync.Map -2.泄漏: resp.Body 未 Close → 请求后立即 defer resp.Body.Close() -3.吞错: err 未检查 → 每处 err!=nil 必处理或上抛 - -## More -- (none) - -Now begin.\ -""" - -COMPRESS_USER = '## Query\n{query}\n\n## Source\n{text}' - -# Short system prompt embedded in emitted SFT samples — the long COMPRESS_SYSTEM -# is for data generation only; training samples carry only the binding contract. -COMPRESS_SYSTEM_TRAIN = """\ -You are a compression 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. - -Output skeleton: - -## Summary -Topic: - - -## More -- : -- ... - -Rules: -1. Line 1 of `## Summary` is ALWAYS `Topic: ...`. -2. Body is maximally dense; every token carries query-relevant signal. -3. Never silently drop a fact — anything cut for length MUST appear as a key \ -under `## More` (do not duplicate inline material here). -4. No fabrication, no extrapolation, no misleading partial truths. -5. Match the source language. No outer code fences, no meta-commentary.\ -""" - -# Fixed queries — used directly (no Phase-1 LLM generation) for a proportion of items. -FIXED_QUERY_NEED = ( - '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.') -FIXED_QUERY_SKILL = ( - '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.') -FIXED_QUERIES = [FIXED_QUERY_NEED, FIXED_QUERY_SKILL] -FIXED_QUERY_RATIO = 0.3 - - -# ═══════════════════════════════════════════════════════════════════════════════ -# Core logic -# ═══════════════════════════════════════════════════════════════════════════════ - -def _extract_json_array(text: str) -> Optional[List[str]]: - """Best-effort extraction of a JSON string array from LLM output.""" - text = text.strip() - # Try direct parse first - if text.startswith('['): - try: - arr = json.loads(text) - if isinstance(arr, list) and all(isinstance(x, str) for x in arr): - return arr - except json.JSONDecodeError: - pass - # Fallback: find first [...] block - m = re.search(r'\[.*\]', text, re.DOTALL) - if m: - try: - arr = json.loads(m.group()) - if isinstance(arr, list) and all(isinstance(x, str) for x in arr): - return arr - except json.JSONDecodeError: - pass - return None - - -def generate_queries(api: OpenAI, text: str) -> List[str]: - """Phase 1: ask the LLM what queries can be asked about ``text``.""" - trajectory = { - 'messages': [ - {'role': 'system', 'content': QUERY_GEN_SYSTEM}, - {'role': 'user', 'content': QUERY_GEN_USER.format(text=text)}, - ] - } - sp = SamplingParams(temperature=0.7, max_tokens=1024) - for attempt in range(2): - try: - reply = api(trajectory, sp, extra_body={'enable_thinking': True}) - except Exception as exc: - sys.stderr.write(f'[query_gen] error: {exc}\n') - return [] - content = reply.get('content') or '' - queries = _extract_json_array(content) - if queries: - return queries - if attempt == 0: - sys.stderr.write('[query_gen] retry: failed to parse JSON array\n') - return [] - - -def compress_for_query(api: OpenAI, text: str, query: str, - thinking_budget: int = 1024) -> Optional[str]: - """Phase 2: compress ``text`` w.r.t. ``query``. Returns compressed content or None.""" - trajectory = { - 'messages': [ - {'role': 'system', 'content': COMPRESS_SYSTEM}, - {'role': 'user', 'content': COMPRESS_USER.format(query=query, text=text)}, - ] - } - sp = SamplingParams(temperature=0.3, max_tokens=16384) - for attempt in range(2): - try: - reply = api(trajectory, sp, extra_body={ - 'enable_thinking': False, - 'thinking_budget': thinking_budget, - }) - except Exception as exc: - sys.stderr.write(f'[compress] error: {exc}\n') - return None - content = (reply.get('content') or '').strip() - if not content: - if attempt == 0: - sys.stderr.write('[compress] retry: empty response\n') - continue - # Strip whole-answer code fence if present. - m = re.match(r'^```[a-zA-Z]*\n(.*?)\n```\s*$', content, re.DOTALL) - if m: - content = m.group(1).strip() - if not (re.search(r'(?im)^##\s*Summary\b', content) - and re.search(r'(?im)^##\s*More\b', content)): - if attempt == 0: - sys.stderr.write('[compress] retry: missing required sections\n') - continue - return content - return None - - -def _query_hash(query: str) -> str: - """Stable short hash of a query string — embedded in sample id for resume.""" - return hashlib.md5(query.strip().encode('utf-8')).hexdigest()[:8] - - -def process_item( - api: OpenAI, - item: Dict[str, Any], - done_sample_ids: Optional[Set[str]] = None, - thinking_budget: int = 1024, - fixed_query_ratio: float = FIXED_QUERY_RATIO, -) -> List[Dict[str, Any]]: - """Run both phases on one dataset item. Returns list of SFT samples. - - Input rows come from ``dataset.py`` (single assistant message) or - ``dataset_think.py`` (user query + assistant with reasoning_content). - For thinking-data rows, ``FIXED_QUERY_NEED`` is applied to the query - and ``FIXED_QUERY_SKILL`` to the CoT, skipping Phase-1 generation. - - ``done_sample_ids`` (full sample ids already on disk for this item) - lets resume skip queries that were already emitted, keyed by query - content hash so a phase-1 reorder still resolves correctly. - """ - done = done_sample_ids or set() - messages = item.get('messages') or [] - - # Detect thinking-data: user message + assistant with reasoning_content - user_query = '' - cot_text = '' - assistant_text = '' - for m in messages: - if not isinstance(m, dict): - continue - role = m.get('role', '') - if role == 'user' and not user_query: - user_query = (m.get('content') or '').strip() - elif role == 'assistant': - cot_text = (m.get('reasoning_content') or '').strip() - assistant_text = (m.get('content') or '').strip() - break - - item_id = item.get('id') - if not item_id: - return [] - source = item.get('source', 'unknown') - - # Thinking-data path: compress query and CoT separately with fixed queries - if user_query and cot_text: - pairs = [(user_query, FIXED_QUERY_NEED), (cot_text, FIXED_QUERY_SKILL)] - samples: List[Dict[str, Any]] = [] - for text, query in pairs: - if len(text) < 100: - continue - sample_id = f'{item_id}__{_query_hash(query)}' - if sample_id in done: - continue - compressed = compress_for_query(api, text, query, thinking_budget=thinking_budget) - if not compressed: - continue - sft_messages = [ - {'role': 'system', 'content': COMPRESS_SYSTEM_TRAIN}, - {'role': 'user', 'content': COMPRESS_USER.format(query=query, text=text)}, - {'role': 'assistant', 'content': compressed}, - ] - samples.append({ - 'id': sample_id, - 'source': source, - 'query': query, - 'original_len': len(text), - 'compressed_len': len(compressed), - 'original_tokens': 0, - 'compressed_tokens': 0, - 'messages': sft_messages, - '__src': text, - '__cmp': compressed, - }) - return samples - - # Plain-data path: single assistant message - text = assistant_text - if not text or len(text) < 100: - return [] - - queries = generate_queries(api, text) - if not queries: - return [] - queries = queries[:2] - - # Mix in fixed queries for a proportion of items - if random.random() < fixed_query_ratio: - queries = list(FIXED_QUERIES) - - samples: List[Dict[str, Any]] = [] - for query in queries: - sample_id = f'{item_id}__{_query_hash(query)}' - if sample_id in done: - continue - compressed = compress_for_query(api, text, query, thinking_budget=thinking_budget) - if not compressed: - continue - sft_messages = [ - {'role': 'system', 'content': COMPRESS_SYSTEM_TRAIN}, - {'role': 'user', 'content': COMPRESS_USER.format(query=query, text=text)}, - {'role': 'assistant', 'content': compressed}, - ] - samples.append({ - 'id': sample_id, - 'source': source, - 'query': query, - 'original_len': len(text), - 'compressed_len': len(compressed), - 'original_tokens': 0, - 'compressed_tokens': 0, - 'messages': sft_messages, - # Stashed for sparse tokenization on main thread; popped before write. - '__src': text, - '__cmp': compressed, - }) - return samples - - -def process_failure( - api: OpenAI, - item: Dict[str, Any], - thinking_budget: int = 1024, -) -> List[Dict[str, Any]]: - """Re-compress a single failure record (id, query, text already pinned). - - Used by ``--failures`` mode: query and source passage are taken verbatim - from the original failure entry, so Phase-1 generation is skipped and the - output id matches the original sample id. - """ - sid = item.get('id') or '' - query = (item.get('query') or '').strip() - text = (item.get('text') or '').strip() - if not sid or not query or not text: - return [] - compressed = compress_for_query(api, text, query, thinking_budget=thinking_budget) - if not compressed: - return [] - sft_messages = [ - {'role': 'system', 'content': COMPRESS_SYSTEM_TRAIN}, - {'role': 'user', 'content': COMPRESS_USER.format(query=query, text=text)}, - {'role': 'assistant', 'content': compressed}, - ] - return [{ - 'id': sid, - 'source': item.get('source', 'failure_regen'), - 'query': query, - 'original_len': len(text), - 'compressed_len': len(compressed), - 'original_tokens': 0, - 'compressed_tokens': 0, - 'messages': sft_messages, - '__src': text, - '__cmp': compressed, - }] - - -# ═══════════════════════════════════════════════════════════════════════════════ -# I/O helpers -# ═══════════════════════════════════════════════════════════════════════════════ - -def iter_input(path: str) -> Iterator[Dict[str, Any]]: - """Stream JSONL dataset row-by-row (no full-file load).""" - with open(path, 'r', encoding='utf-8') as fh: - for line in fh: - line = line.strip() - if not line: - continue - try: - yield json.loads(line) - except json.JSONDecodeError: - continue - - -def iter_dataset_py(total: Optional[int], load_from_cache_file: bool) -> Iterator[Dict[str, Any]]: - """Stream rows directly from ``dataset.py::get_dataset`` without any JSONL hop.""" - # Lazy import: dataset.py triggers HF / ModelScope downloads at module load. - sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - from cookbook.exp.condenser.dataset import get_dataset - hf = get_dataset(total=total, load_from_cache_file=load_from_cache_file) - sys.stderr.write(f'Loaded dataset.py::get_dataset: {len(hf)} rows\n') - for row in hf: - yield row - - -def iter_dataset_think_py(total: Optional[int], load_from_cache_file: bool) -> Iterator[Dict[str, Any]]: - """Stream rows from ``dataset_think.py::get_dataset`` (query + CoT data).""" - sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - from dataset_think import get_dataset - hf = get_dataset(total=total, load_from_cache_file=load_from_cache_file) - sys.stderr.write(f'Loaded dataset_think.py::get_dataset: {len(hf)} rows\n') - for row in hf: - yield row - - -def iter_failures(path: str, skip_ids: Optional[Set[str]] = None) -> Iterator[Dict[str, Any]]: - """Stream records from a ``failures.jsonl`` for re-compression. - - Each input record carries a full sample id, the original query, and a - user message whose body embeds the source passage after a ``## Passage`` - or ``## Source`` header. The yielded item is shaped for ``process_failure`` - (id, source, query, text). Items whose id is in ``skip_ids`` are skipped. - """ - skip = skip_ids or set() - n_total = n_skipped = n_yielded = n_bad = 0 - with open(path, 'r', encoding='utf-8') as fh: - for line in fh: - line = line.strip() - if not line: - continue - n_total += 1 - try: - obj = json.loads(line) - except json.JSONDecodeError: - n_bad += 1 - continue - sid = obj.get('id') or '' - if not sid: - n_bad += 1 - continue - if sid in skip: - n_skipped += 1 - continue - query = (obj.get('query') or '').strip() - user_content = '' - for m in obj.get('messages') or []: - if isinstance(m, dict) and m.get('role') == 'user': - user_content = m.get('content') or '' - break - text = '' - for sep in ('## Passage\n', '## Source\n'): - if sep in user_content: - text = user_content.split(sep, 1)[1].strip() - break - if not query or not text: - sys.stderr.write(f'[failures] skip {sid}: missing query/passage\n') - n_bad += 1 - continue - n_yielded += 1 - yield { - 'id': sid, - 'source': obj.get('source', 'failure_regen'), - 'query': query, - 'text': text, - } - sys.stderr.write( - f'[failures] total={n_total} yielded={n_yielded} ' - f'resume_skipped={n_skipped} malformed={n_bad}\n') - - -def load_done_sample_ids(path: str) -> Set[str]: - """Collect already-written full sample ids (``base__hash``) for resume.""" - if not os.path.exists(path): - return set() - done: Set[str] = set() - with open(path, 'r', encoding='utf-8') as fh: - for line in fh: - try: - obj = json.loads(line) - except json.JSONDecodeError: - continue - sid = obj.get('id', '') - if sid: - done.add(sid) - return done - - -# ═══════════════════════════════════════════════════════════════════════════════ -# Main -# ═══════════════════════════════════════════════════════════════════════════════ - -def main() -> None: - parser = argparse.ArgumentParser( - description='Two-phase query-diverse condenser dataset builder.') - parser.add_argument('--input', default=None, - help='Optional JSONL override; default uses dataset.py::get_dataset') - parser.add_argument('--output', required=True, - help='Output JSONL file for SFT samples') - parser.add_argument('--total', type=int, default=0, - help='Total input rows for proportional scaling in dataset.py (0 = base sizes)') - parser.add_argument('--no-cache', action='store_true', - help='Disable load_from_cache_file when calling dataset.py::get_dataset') - parser.add_argument('--model', required=True, - help='API model name') - parser.add_argument('--api-key', default=os.environ.get('OPENAI_API_KEY')) - parser.add_argument('--base-url', default=os.environ.get('OPENAI_BASE_URL')) - parser.add_argument('--concurrency', type=int, default=32, - help='Number of parallel workers') - parser.add_argument('--limit', type=int, default=0, - help='Max items to process (0 = all)') - parser.add_argument('--thinking-budget', type=int, default=1024, - help='Max thinking tokens for phase-2 compress (shorter = faster, cheaper)') - parser.add_argument('--tokenizer', default='Qwen/Qwen3.5-4B', - help='HF/ModelScope tokenizer id for sparse token-ratio probe') - parser.add_argument('--tokenize-every', type=int, default=1000, - help='Tokenize one sample every N writes; others get tokens=0') - parser.add_argument('--fixed-query-ratio', type=float, default=FIXED_QUERY_RATIO, - help='Proportion of plain-data items using fixed queries instead of LLM-generated ones') - parser.add_argument('--source', choices=['think', 'plain', 'both'], default='think', - help='Data source: think=dataset_think.py (query+CoT), plain=dataset.py, both=chain both') - parser.add_argument('--failures', default=None, - help='Path to a failures.jsonl; when set, re-generate compressions for every record ' - 'using its original (query, passage) pair and ignore --input/--source.') - args = parser.parse_args() - - out_dir = os.path.dirname(args.output) - if out_dir: - os.makedirs(out_dir, exist_ok=True) - - done_sample_ids = load_done_sample_ids(args.output) - # Group done sample ids by base item id so each worker only sees its slice. - done_per_item: Dict[str, Set[str]] = {} - for sid in done_sample_ids: - if '__' in sid: - base = sid.rsplit('__', 1)[0] - done_per_item.setdefault(base, set()).add(sid) - sys.stderr.write( - f'Resume: {len(done_sample_ids)} samples on disk across ' - f'{len(done_per_item)} items.\n') - - api = OpenAI(model=args.model, api_key=args.api_key, base_url=args.base_url) - - from modelscope import AutoTokenizer - tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True) - - def iter_pending() -> Iterator[Dict[str, Any]]: - if args.failures: - source_iter = iter_failures(args.failures, done_sample_ids) - elif args.input: - source_iter = iter_input(args.input) - else: - import itertools - sources = [] - if args.source in ('plain', 'both'): - sources.append(iter_dataset_py( - total=args.total or None, - load_from_cache_file=not args.no_cache, - )) - if args.source in ('think', 'both'): - sources.append(iter_dataset_think_py( - total=args.total or None, - load_from_cache_file=not args.no_cache, - )) - source_iter = itertools.chain(*sources) - emitted = 0 - for it in source_iter: - iid = it.get('id') - if not iid: - sys.stderr.write('[skip] row missing "id" field\n') - continue - if args.limit > 0 and emitted >= args.limit: - return - yield it - emitted += 1 - - write_lock = threading.Lock() - out_fh = open(args.output, 'a', encoding='utf-8') - items_done = 0 - items_failed = 0 - samples_emitted = 0 - pbar = tqdm(desc='condense', unit='item', dynamic_ncols=True) - - items_iter = iter_pending() - in_flight: Dict[Any, str] = {} - # Sliding window: keep ~2x concurrency tasks queued so the pool never starves. - window = max(args.concurrency * 2, args.concurrency + 4) - - try: - with ThreadPoolExecutor(max_workers=args.concurrency) as ex: - exhausted = False - while True: - while not exhausted and len(in_flight) < window: - try: - it = next(items_iter) - except StopIteration: - exhausted = True - break - iid = it['id'] - if args.failures: - fut = ex.submit( - process_failure, api, it, args.thinking_budget, - ) - else: - fut = ex.submit( - process_item, api, it, done_per_item.get(iid), - args.thinking_budget, args.fixed_query_ratio, - ) - in_flight[fut] = iid - if not in_flight: - break - done, _ = wait(list(in_flight.keys()), return_when=FIRST_COMPLETED) - for fut in done: - iid = in_flight.pop(fut) - try: - samples = fut.result() - except Exception as exc: - sys.stderr.write(f'[item {iid}] crashed: {exc}\n') - items_failed += 1 - pbar.update(1) - continue - if not samples: - items_failed += 1 - pbar.update(1) - continue - with write_lock: - for s in samples: - src = s.pop('__src', '') - cmp = s.pop('__cmp', '') - samples_emitted += 1 - if (samples_emitted - 1) % args.tokenize_every == 0: - s['original_tokens'] = len(tokenizer(src).input_ids) - s['compressed_tokens'] = len(tokenizer(cmp).input_ids) - out_fh.write(json.dumps(s, ensure_ascii=False) + '\n') - out_fh.flush() - items_done += 1 - pbar.set_postfix( - done=items_done, failed=items_failed, - samples=samples_emitted, refresh=False, - ) - pbar.update(1) - finally: - out_fh.close() - pbar.close() - - sys.stderr.write( - f'Done. items_done={items_done}, samples={samples_emitted}, ' - f'failed={items_failed}\n') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/legacy/condenser/train_condenser_ddp.py b/cookbook/exp/legacy/condenser/train_condenser_ddp.py deleted file mode 100644 index 997235781..000000000 --- a/cookbook/exp/legacy/condenser/train_condenser_ddp.py +++ /dev/null @@ -1,100 +0,0 @@ -"""Ray LoRA SFT for the condenser model on condense_300K. - -Launch: - python cookbook/exp/train_condenser_ddp.py -""" -from pathlib import Path - -from peft import LoraConfig -from tqdm import tqdm - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_device_placement, get_logger -from twinkle.dataloader import DataLoader -from twinkle.dataset import Dataset, DatasetMeta -from twinkle.model import TransformersModel -from twinkle.preprocessor import Preprocessor - -logger = get_logger() - -MODEL_ID = 'ms://Qwen/Qwen3.5-4B' -DATASET_ID = 'ms://twinkle-kit/condense_300K' -TEMPLATE_NAME = 'Qwen3_5Template' - -DP_SIZE = 8 -BATCH_SIZE = 8 -LEARNING_RATE = 1e-5 -GRADIENT_ACCUMULATION_STEPS = 8 -LOG_INTERVAL = 20 -EVAL_INTERVAL = 200 -EVAL_SAMPLES = 100 -NUM_EPOCHS = 1 - -OUTPUT_DIR = './output/condenser_ddp' -RESUME_FROM_CHECKPOINT = None -RESUME_ONLY_MODEL = False -IGNORE_DATA_SKIP = False -ADAPTER_NAME = 'default' - -class LegacySectionRenameProcessor(Preprocessor): - """Rewrite legacy `## Read inline` / `## Call extract_compressed for` headers to `## Summary` / `## More`.""" - - _REPLACEMENTS = ( - ('## Read inline', '## Summary'), - ('## Call extract_compressed for', '## More'), - ) - - def __call__(self, batch): - new_messages = [] - for msgs in batch['messages']: - patched = [] - for m in msgs: - content = m.get('content', '') or '' - for old, new in self._REPLACEMENTS: - content = content.replace(old, new) - patched.append({**m, 'content': content}) - new_messages.append(patched) - return {'messages': new_messages} - - -def build_dataset() -> Dataset: - dataset = Dataset(dataset_meta=DatasetMeta('/mnt/workspace/yzhao/tastelikefeet/condense_300K/train.jsonl')) - dataset.map(LegacySectionRenameProcessor(), remove_columns=[], num_proc=16) - dataset.set_template(TEMPLATE_NAME, model_id=MODEL_ID, max_length=40000, enable_thinking=False, truncation_strategy='delete') - dataset.encode(load_from_cache_file=True, num_proc=64) - return dataset - - -def train(): - device_groups = [DeviceGroup(name='model', ranks=DP_SIZE, device_type='GPU')] - model_mesh = DeviceMesh.from_sizes(world_size=DP_SIZE, dp_size=4, fsdp_size=2) - twinkle.initialize(mode='ray', nproc_per_node=DP_SIZE, groups=device_groups, global_device_mesh=model_mesh) - - dataset = build_dataset() - dataloader = DataLoader(dataset=dataset, batch_size=BATCH_SIZE, shuffle=True) - - model = TransformersModel(model_id=MODEL_ID, device_mesh=model_mesh, remote_group='model') - - model.set_optimizer(optimizer_cls='AdamW', lr=LEARNING_RATE) - total_optim_steps = (len(dataloader) * NUM_EPOCHS) // GRADIENT_ACCUMULATION_STEPS - model.set_lr_scheduler( - scheduler_cls='CosineWarmupScheduler', num_warmup_steps=50, num_training_steps=total_optim_steps) - - logger.info(get_device_placement()) - logger.info(model.get_train_configs()) - logger.info(f'Total micro-steps: {len(dataloader) * NUM_EPOCHS}, optim steps: {total_optim_steps}') - - for i in range(NUM_EPOCHS): - for cur_step, batch in enumerate(dataloader): - model.forward_backward(inputs=batch) - model.clip_grad_and_step(gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) - if cur_step % LOG_INTERVAL == 0: - metric = model.calculate_metric(is_training=True) - logger.info(f'Step {cur_step}/{len(dataloader) * NUM_EPOCHS}, metric: {metric}') - if cur_step % 4000 == 0: - model.save(f'step_{cur_step}', output_dir=OUTPUT_DIR) - model.save('last_checkpoint', output_dir=OUTPUT_DIR) - - -if __name__ == '__main__': - train() diff --git a/cookbook/exp/legacy/condenser/untested/eval_condensed.py b/cookbook/exp/legacy/condenser/untested/eval_condensed.py deleted file mode 100644 index 730aaf3a8..000000000 --- a/cookbook/exp/legacy/condenser/untested/eval_condensed.py +++ /dev/null @@ -1,382 +0,0 @@ -"""Evaluation: native (full ctx) vs condensed (chunk → condense → extract_condensed tool). - -Reuses the training-time data shape and prompt so the comparison is apples-to-apples. - -Launch: - # native baseline (full HotpotQA context, no compression, no tool) - python cookbook/exp/eval_condensed.py --mode native \\ - --dataset /path/to/hotpot_dev_fullwiki.jsonl - - # condensed (chunk → condense via Qwen3.5-4B-Condenser → extract_condensed tool) - python cookbook/exp/eval_condensed.py --mode condensed \\ - --dataset /path/to/hotpot_dev_fullwiki.jsonl - -Outputs (under --out_dir / _/): - predictions.jsonl one row per sample with pred / gold / f1 / em / token-counts / tool-calls - summary.json aggregate metrics -""" -import argparse -import json -import os -import re -import time -import uuid -from collections import Counter -from typing import Any, Dict, List, Optional - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.data_format import Message, SamplingParams, Trajectory -from twinkle.dataloader import DataLoader -from twinkle.dataset import Dataset, DatasetMeta -from twinkle.sampler import vLLMSampler -from twinkle.template import Qwen3_5Template -from twinkle_agentic.chunker.native import NativeChunker -from twinkle_agentic.condenser import ModelCondenser -from twinkle_agentic.reward.f1 import _f1_score -from twinkle_agentic.rollout.multi_turn import MultiTurnRollout -from twinkle_agentic.rollout.multi_turn_condense import MultiTurnCondenseRollout -from twinkle_agentic.tools.tool_manager import ToolManager -from twinkle.preprocessor.base import Preprocessor - -# Reuse training assets so eval and train share data shape + condensed prompt. -from cookbook.exp.legacy.grpo_condensed import ( - SYSTEM_PROMPT as CONDENSED_SYSTEM_PROMPT, - HotpotQAProcessor, - _BOXED_RE, - _last_assistant_text, -) - - -class MuSiQueProcessor(Preprocessor): - """MuSiQue-Ans → Trajectory adapter. - - MuSiQue native schema (per row): - id, question, paragraphs=[{idx, title, paragraph_text, is_supporting}], answer, - answer_aliases=[...], answerable, question_decomposition=[...] - - Maps to the same Trajectory(messages, user_data) shape that - :class:`HotpotQAProcessor` produces, so downstream rollout code is - schema-agnostic. ``ground_truth`` carries answer + answer_aliases. - """ - - def __init__(self, system: str): - self.system = system - - def __call__(self, rows: Dict[str, List[Any]]) -> Dict[str, List[Any]]: - rows = self.map_col_to_row(rows) - out = [self.preprocess(r) for r in rows] - out = [r for r in out if r is not None] - return self.map_row_to_col(out) - - @staticmethod - def _format_context(paragraphs: List[Dict[str, Any]]) -> str: - lines = [] - for p in paragraphs or []: - title = (p.get('title') or '').strip() - body = (p.get('paragraph_text') or '').strip() - if not body: - continue - lines.append(f'{title}: {body}' if title else body) - return '\n\n'.join(lines) - - def preprocess(self, row: Dict[str, Any]) -> Optional[Trajectory]: - if row.get('answerable') is False: - return None - question = (row.get('question') or '').strip() - if not question: - return None - gold_main = (row.get('answer') or '').strip() - aliases = row.get('answer_aliases') or [] - gold = [g for g in dict.fromkeys([gold_main] + list(aliases)) if g] - if not gold: - return None - paragraphs = row.get('paragraphs') or [] - context_block = self._format_context(paragraphs) - user_msg = f'Question: {question}\n\nContext:\n\n{context_block}' - messages = [ - Message(role='system', content=self.system), - Message(role='user', content=user_msg), - ] - sf_titles = list(dict.fromkeys( - (p.get('title') or '').strip() - for p in paragraphs - if p.get('is_supporting') and (p.get('title') or '').strip())) - user_data = [('ground_truth', g) for g in gold] + [('sf_title', t) for t in sf_titles] - return Trajectory(messages=messages, user_data=user_data) - -logger = get_logger() - -NATIVE_SYSTEM_PROMPT = """You are a careful multi-hop QA assistant. - -The user message contains a Question and a Context. Read both, reason step by step, -then commit to a final answer. - -## Output Format -End your final response with \\boxed{answer}. -Keep the boxed text short: a name, entity, date, or "yes"/"no". -Answers not inside \\boxed{} will not be scored.""" - - -def parse_args(): - p = argparse.ArgumentParser() - p.add_argument('--mode', choices=['native', 'condensed'], required=True) - p.add_argument('--dataset', required=True, - help='Eval set jsonl. HotpotQA or MuSiQue-Ans schema (see --dataset_format).') - p.add_argument('--dataset_format', choices=['hotpotqa', 'musique'], default='musique', - help='Schema of --dataset. MuSiQue-Ans (default) is harder multi-hop and OOD vs training.') - p.add_argument('--model_id', default='ms://Qwen/Qwen3.5-4B') - p.add_argument('--lora_path', default=None, - help='Optional LoRA adapter on top of model_id (e.g. trained QA LoRA).') - p.add_argument('--condenser_lora', default='ms://twinkle-kit/Qwen3.5-4B-Condenser') - p.add_argument('--limit', type=int, default=500) - p.add_argument('--num_gpus', type=int, default=4) - p.add_argument('--batch_size', type=int, default=8) - p.add_argument('--max_model_len', type=int, default=32768) - p.add_argument('--max_new_tokens', type=int, default=2048) - p.add_argument('--max_turns', type=int, default=4) - p.add_argument('--max_trajectory_tokens', type=int, default=8192) - p.add_argument('--chunk_size', type=int, default=1024) - p.add_argument('--temperature', type=float, default=0.0) - p.add_argument('--out_dir', default='eval_out') - p.add_argument('--seed', type=int, default=42) - return p.parse_args() - - -def build_dataset(path: str, dataset_format: str, model_id: str, - max_length: int, limit: int, system: str) -> Dataset: - """Load eval JSONL and produce Trajectory rows tagged with ground_truth user_data.""" - ds = Dataset() - ds.add_dataset(DatasetMeta(path)) - if limit > 0 and len(ds) > limit: - ds = ds.select(range(limit)) - ds.set_template( - 'Qwen3_5Template', model_id=model_id, max_length=max_length, - truncation_strategy='delete', enable_thinking=False) - if dataset_format == 'musique': - # MuSiQue-Ans cols (drop everything; we keep only the produced messages/user_data) - cols = ['id', 'question', 'paragraphs', 'answer', 'answer_aliases', - 'answerable', 'question_decomposition'] - ds.map(MuSiQueProcessor(system=system), remove_columns=cols) - else: - cols = ['id', 'question', 'question_fixed', 'answers', 'original_answer', - 'type', 'level', 'verdict', 'reasoning', 'supporting_facts', 'context'] - ds.map(HotpotQAProcessor(system=system), remove_columns=cols) - return ds - - -def extract_boxed(text: str) -> Optional[str]: - """Pull the inner text of the LAST `\\boxed{...}` marker, brace-balanced enough for short answers.""" - if not text: - return None - matches = _BOXED_RE.findall(text) - if not matches: - return None - last = matches[-1] - return last[len(r'\boxed{'):-1].strip() - - -def best_f1_em(pred: str, golds: List[str]) -> Dict[str, float]: - """Max-over-references SQuAD-style F1 / EM, reusing the training reward's normalizer.""" - if not golds: - return {'f1': 0.0, 'em': 0.0} - if not pred: - return {'f1': 0.0, 'em': 0.0} - best_f1, best_em = 0.0, 0.0 - for g in golds: - f1, em = _f1_score(pred, g) - if f1 > best_f1: - best_f1 = f1 - if em > best_em: - best_em = em - return {'f1': best_f1, 'em': best_em} - - -def _user_text(traj_or_msg) -> str: - """Concat all text parts of the first user message — used to count original context tokens.""" - msgs = traj_or_msg if isinstance(traj_or_msg, list) else (traj_or_msg.get('messages') or []) - for m in msgs: - role = m.get('role') if isinstance(m, dict) else getattr(m, 'role', None) - if role != 'user': - continue - content = m.get('content') if isinstance(m, dict) else getattr(m, 'content', None) - if isinstance(content, str): - return content - if isinstance(content, list): - return ''.join(p.get('text') or '' for p in content if isinstance(p, dict) and p.get('type') == 'text') - return '' - return '' - - -def _count_tool_calls(traj: Dict[str, Any]) -> int: - return sum(len(m.get('tool_calls') or []) - for m in (traj.get('messages') or []) if m.get('role') == 'assistant') - - -def main(): - args = parse_args() - run_id = time.strftime('%Y%m%d_%H%M%S') + '_' + uuid.uuid4().hex[:6] - out_dir = os.path.join(args.out_dir, f'{args.mode}_{run_id}') - os.makedirs(out_dir, exist_ok=True) - - device_groups = [DeviceGroup(name='sampler', ranks=list(range(args.num_gpus)), device_type='GPU')] - sampler_mesh = DeviceMesh.from_sizes(world_size=args.num_gpus, dp_size=args.num_gpus) - twinkle.initialize(mode='ray', nproc_per_node=args.num_gpus, - groups=device_groups, lazy_collect=False) - - system = CONDENSED_SYSTEM_PROMPT if args.mode == 'condensed' else NATIVE_SYSTEM_PROMPT - ds = build_dataset(args.dataset, args.dataset_format, args.model_id, - args.max_model_len, args.limit, system) - logger.info('Eval dataset: %d rows from %s (mode=%s, format=%s)', - len(ds), args.dataset, args.mode, args.dataset_format) - - sampler = vLLMSampler( - model_id=args.model_id, - engine_args={ - 'gpu_memory_utilization': 0.85, 'max_model_len': args.max_model_len, - 'max_lora_rank': 32, 'enable_lora': True, - 'enable_tower_connector_lora': True, 'max_loras': 5, - 'seed': args.seed, - }, - device_mesh=sampler_mesh, remote_group='sampler') - sampler.set_template('Qwen3_5Template', model_id=args.model_id, - enable_thinking=False, max_length=args.max_model_len) - template = Qwen3_5Template(args.model_id, max_length=args.max_model_len, enable_thinking=False) - - # stop=[''] only matters for condensed mode where the model issues tool calls - sampling_params = SamplingParams( - max_tokens=args.max_new_tokens, num_samples=1, - temperature=args.temperature, top_p=0.95, - stop=[''] if args.mode == 'condensed' else None, - ) - - if args.mode == 'condensed': - chunker = NativeChunker(chunk_size=args.chunk_size, passage_boundary_re=r'(?<=\n\n)') - # Chunk-level extraction of the question line; \A anchor avoids matching "Question:" inside passages. - _q_re = re.compile(r'\AQuestion:\s*(.+)') - - def _q_from_chunk(chunk): - c = chunk.get('content') - if chunk.get('type') != 'text' or not isinstance(c, str): - return None - m = _q_re.search(c) - return m.group(1).strip() if m else None - - condenser = ModelCondenser( - sampler=sampler, compression_ratio=2.0, - sampling_params=SamplingParams(max_tokens=1024, num_samples=1, - temperature=0.4, top_p=0.9), - min_chars=200, template=template, - lora_path=args.condenser_lora, skip_pattern=r'^Question:', - related_query=_q_from_chunk, - ) - rollout = MultiTurnCondenseRollout( - sampler=sampler, template=template, tool_manager=ToolManager(), - chunker=chunker, condenser=condenser, - sampling_params=sampling_params, - max_turns=args.max_turns, max_trajectory_tokens=args.max_trajectory_tokens, - ) - else: - # max_turns=1, no tools: reduces to single-turn QA over the full original context - rollout = MultiTurnRollout( - sampler=sampler, template=template, tool_manager=ToolManager(), - sampling_params=sampling_params, - max_turns=1, max_trajectory_tokens=args.max_trajectory_tokens, - ) - - dataloader = DataLoader(dataset=ds, batch_size=args.batch_size, - min_batch_size=1, shuffle=False) - - pred_path = os.path.join(out_dir, 'predictions.jsonl') - pf = open(pred_path, 'w', encoding='utf-8') - - agg = Counter() - sums = {'f1': 0.0, 'em': 0.0, - 'prompt_tok': 0, 'comp_tok': 0, 'orig_ctx_tok': 0, - 'turns': 0, 'tool_calls': 0} - t0 = time.time() - - for batch in dataloader: - trajs = rollout(batch) - - for src, traj in zip(batch, trajs): - text = _last_assistant_text(traj) or '' - pred = extract_boxed(text) or '' - golds = [v for k, v in (src.user_data or []) if k == 'ground_truth' and v] - - scores = best_f1_em(pred, golds) - ids = traj.get('input_ids') or [] - comp_tok = sum(1 for l in (traj.get('labels') or []) if l != -100) - prompt_tok = max(0, len(ids) - comp_tok) - tool_calls = _count_tool_calls(traj) - - # Original (uncondensed) context size — feed only the user msg, not the system prompt, - # so the compression ratio stays comparable across modes. - orig_user = _user_text(src.messages) - orig_ctx_tok = len(template.tokenizer.encode(orig_user)) if orig_user else 0 - - agg['n'] += 1 - agg['no_box'] += int(_BOXED_RE.search(text) is None) - agg['tool_use'] += int(tool_calls > 0) - sums['f1'] += scores['f1'] - sums['em'] += scores['em'] - sums['prompt_tok'] += prompt_tok - sums['comp_tok'] += comp_tok - sums['orig_ctx_tok'] += orig_ctx_tok - sums['turns'] += int(traj.get('turns') or 1) - sums['tool_calls'] += tool_calls - - pf.write(json.dumps({ - 'pred': pred, - 'gold': golds, - 'f1': scores['f1'], - 'em': scores['em'], - 'prompt_tok': prompt_tok, - 'comp_tok': comp_tok, - 'orig_ctx_tok': orig_ctx_tok, - 'tool_calls': tool_calls, - 'turns': int(traj.get('turns') or 1), - 'no_boxed': _BOXED_RE.search(text) is None, - 'response': text, - }, ensure_ascii=False) + '\n') - - logger.info('[eval] %d / %d processed', agg['n'], len(ds)) - - pf.close() - wall = time.time() - t0 - n = max(1, agg['n']) - summary = { - 'mode': args.mode, - 'dataset_format': args.dataset_format, - 'model_id': args.model_id, - 'lora_path': args.lora_path, - 'condenser_lora': args.condenser_lora if args.mode == 'condensed' else None, - 'dataset': args.dataset, - 'n_samples': agg['n'], - # quality - 'f1': sums['f1'] / n, - 'em': sums['em'] / n, - 'no_boxed_rate': agg['no_box'] / n, - # cost - 'avg_prompt_tokens': sums['prompt_tok'] / n, - 'avg_completion_tokens': sums['comp_tok'] / n, - 'avg_orig_context_tokens': sums['orig_ctx_tok'] / n, - 'compression_ratio': (sums['prompt_tok'] / sums['orig_ctx_tok'] - if sums['orig_ctx_tok'] else None), - # tool / multi-turn behavior - 'avg_turns': sums['turns'] / n, - 'avg_tool_calls': sums['tool_calls'] / n, - 'tool_use_rate': agg['tool_use'] / n, - # wall - 'wall_time_sec': wall, - 'samples_per_sec': agg['n'] / wall if wall > 0 else 0.0, - } - with open(os.path.join(out_dir, 'summary.json'), 'w', encoding='utf-8') as f: - json.dump(summary, f, indent=2, ensure_ascii=False) - - logger.info('Done. Output: %s', out_dir) - logger.info('Summary: %s', json.dumps(summary, indent=2, ensure_ascii=False)) - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/legacy/data_pipeline/audit_rubric.py b/cookbook/exp/legacy/data_pipeline/audit_rubric.py deleted file mode 100644 index 2d3107373..000000000 --- a/cookbook/exp/legacy/data_pipeline/audit_rubric.py +++ /dev/null @@ -1,167 +0,0 @@ -"""Audit rubric scoring accuracy: re-score a spread of kept agent trajectories -with the SAME teacher RubricVerifier used in the pipeline, and print, per chosen -trajectory, the generated rubric + per-criterion pass rate + a readable segment -summary so a human can judge whether the score is *right* (good and bad alike). - -Run (same env as the pipeline): - LLM_BACKUP_MODEL=qwen3.7-max \ - LLM_BACKUP_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1 \ - LLM_BACKUP_API_KEY=sk-... \ - python cookbook/exp/data_pipeline/audit_rubric.py -""" -import json -import os -from typing import Any, Dict, List, Optional - -PROCESSED = os.environ.get( - 'PROCESSED_PATH', './output/data_pipeline/processed_20260531_200.jsonl') -N_HIGH = int(os.environ.get('AUDIT_N_HIGH', 3)) -N_LOW = int(os.environ.get('AUDIT_N_LOW', 3)) -N_MID = int(os.environ.get('AUDIT_N_MID', 2)) -# Same stabilization policy as the pipeline (skeleton|fixed|off). -RUBRIC_MODE = os.environ.get('TRAJ_RUBRIC_MODE', 'skeleton').strip().lower() -# Repeat each re-score REPEAT times to measure score variance (jitter). >1 to -# check whether the stabilization actually lowered the spread. -REPEAT = int(os.environ.get('AUDIT_REPEAT', 1)) - - -def _lab(row: Dict[str, Any], key: str) -> Optional[Any]: - for kv in (row.get('user_data') or []): - if isinstance(kv, list) and len(kv) == 2 and kv[0] == key: - try: - return json.loads(kv[1]) - except Exception: - return kv[1] - return None - - -def _seg_summary(messages: List[dict], max_chars: int = 900) -> str: - parts = [] - for m in messages: - role = m.get('role', '?') - content = m.get('content') or '' - tc = m.get('tool_calls') - if tc: - try: - calls = json.loads(tc) if isinstance(tc, str) else tc - names = ','.join(c.get('function', {}).get('name', '?') for c in calls) - content = (content + f' [tool_calls: {names}]').strip() - except Exception: - pass - content = content.replace('\n', ' ') - if len(content) > 200: - content = content[:200] + '…' - parts.append(f' {role}: {content}') - text = '\n'.join(parts) - return text if len(text) <= max_chars else text[:max_chars] + '\n …(truncated)' - - -def _infer_intent(messages: List[dict]) -> Optional[str]: - """Structural intent for the whole trajectory (same detectors as the scorer).""" - from twinkle_agentic.preprocessor.intent_classifier import ( - CodeDetector, MathDetector, ToolCallDetector) - # tool_calls arrive JSON-encoded in the processed jsonl; decode so the - # ToolCallDetector (which reads normalized tool_calls) can see them. - norm = [] - for m in messages: - m = dict(m) - tc = m.get('tool_calls') - if isinstance(tc, str) and tc.strip(): - try: - m['tool_calls'] = json.loads(tc) - except Exception: - m['tool_calls'] = [] - norm.append(m) - for det in (ToolCallDetector(), CodeDetector(), MathDetector()): - try: - if det(norm): - return det.intent - except Exception: - continue - return None - - -def main() -> None: - from twinkle_agentic.verifier import (RubricVerifier, - default_intent_base_rubrics, - default_intent_fixed_rubrics) - - rows = [json.loads(l) for l in open(PROCESSED)] - scored = [r for r in rows if _lab(r, 'traj_score') is not None] - scored.sort(key=lambda r: _lab(r, 'traj_score')) - if not scored: - print('no scored rows found') - return - - picks: List[Dict[str, Any]] = [] - picks += scored[:N_LOW] # lowest - mid = len(scored) // 2 - picks += scored[mid:mid + N_MID] # middle - picks += scored[-N_HIGH:] # highest - # de-dup by id, preserve order - seen = set() - uniq = [] - for r in picks: - if r.get('id') not in seen: - seen.add(r.get('id')) - uniq.append(r) - - intent_base = intent_fixed = None - if RUBRIC_MODE == 'skeleton': - intent_base = default_intent_base_rubrics() - elif RUBRIC_MODE == 'fixed': - intent_fixed = default_intent_fixed_rubrics() - rv = RubricVerifier( - max_votes=5, max_votes_long=3, min_votes_long=2, long_margin_threshold=0.18, - min_votes_high=3, high_score_threshold=0.85, - intent_base_rubrics=intent_base, intent_rubrics=intent_fixed) - print(f'[audit] rubric_mode={RUBRIC_MODE} repeat={REPEAT}') - - for r in uniq: - stored = _lab(r, 'traj_score') - seg_scores = _lab(r, 'segment_scores') - messages = r.get('messages') or [] - intent = _infer_intent(messages) - print('=' * 100) - print(f"id={r.get('id')} model={r.get('model_id')} n_msgs={len(messages)} intent={intent}") - print(f"stored traj_score={stored} level={_lab(r,'traj_level')} " - f"segment_scores={seg_scores} safety={_lab(r,'safety_score')}") - print('-- trajectory summary --') - print(_seg_summary(messages)) - traj = {'messages': messages} - if r.get('tools'): - traj['tools'] = r['tools'] - - scalars: List[float] = [] - det = None - for _ in range(max(1, REPEAT)): - try: - det = rv.score_detail(traj, intent=intent) - except Exception as e: - print(f'!! re-score failed: {e}') - det = None - break - scalars.append(det.scalar) - if det is None: - continue - print('-- teacher re-score --') - print(f" llm_scalar={det.llm_scalar:.3f} hard_pass_rate={det.hard_pass_rate:.3f} " - f"scalar={det.scalar:.3f} gated={det.gated} n_votes={det.n_votes}") - if REPEAT > 1: - lo, hi = min(scalars), max(scalars) - mean = sum(scalars) / len(scalars) - var = sum((s - mean) ** 2 for s in scalars) / len(scalars) - print(f" [variance over {REPEAT}] mean={mean:.3f} spread={hi - lo:.3f} " - f"std={var ** 0.5:.3f} scalars={[round(s, 3) for s in scalars]}") - rubric = det.rubric or [] - rates = det.per_item_pass_rate or [] - for i, it in enumerate(rubric): - rate = rates[i] if i < len(rates) else None - kind = 'HARD' if getattr(it, 'is_hard', False) else 'prin' - rate_s = 'n/a' if rate is None else f'{rate:.2f}' - print(f' [{kind}] pass={rate_s} {it.text}') - print('=' * 100) - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/legacy/data_pipeline/process_and_save.py b/cookbook/exp/legacy/data_pipeline/process_and_save.py deleted file mode 100644 index 3a19b1569..000000000 --- a/cookbook/exp/legacy/data_pipeline/process_and_save.py +++ /dev/null @@ -1,371 +0,0 @@ -"""Standalone dataset-processing demo: run the full agentic QualityPreprocessor -over a slice of the raw OpenClaw CSV, then persist the cleaned + scored + tagged -trajectories with ``Dataset.save_as`` for inspection. - -This is NOT a training script — no template/encode/pack. The point is to see how -the preprocessor behaves end-to-end and to keep ALL enrichment: - -- **tags / scores in ``user_data``** (AUDIT A5 envelope): per-round & trajectory - scores (``TrajectoryScorer``), safety (``SafetyScorer``), intent key-rounds - (``IntentClassifier``), structural-noise ratio (``StructuralNoiseTagger``), - and provenance/lineage (``ProvenanceStamp``). -- **important fields preserved**: ``id``/``source``/``model_id``/``messages`` - are never dropped; ``MessageNormalizer`` now passes through - ``reasoning_content``/``thinking`` (AUDIT P3). - -Pipeline follows the tag-then-filter architecture: mappers annotate, a final -read-only ``TrajectoryOutcomeFilter`` drops on the tags (no DAG, linear order). - -Run: - CSV_PATH=/mnt/data/yzhao/tastelikefeet/bc/20260531.csv \ - USE_RUBRIC=1 SELECT_FRAC=0.1 DATASET_TOTAL=2000 \ - python cookbook/exp/data_pipeline/process_and_save.py - -With gating, rubric LLM cost ~ ``SELECT_FRAC * N_kept`` (not ``N_kept``). Ingest -more rows in pass 1; only the global top fraction gets rubric in pass 2. -""" -import json -import os -from functools import partial -from pathlib import Path -from typing import Any, Dict, Iterator, List - -from twinkle.dataset import Dataset -from twinkle.dataset.base import DatasetMeta -from twinkle.utils import get_logger -from twinkle_agentic.preprocessor import (DeadLoopFilter, HardFilter, - IntentClassifier, LanguageFilter, - MessageNormalizer, - MessageSanityFilter, ModelFilter, - ProvenanceStamp, QualityPreprocessor, - RefuseFilter, SafetyScorer, - SpecialCharsFilter, - StructuralNoiseTagger, - TokenSoupFilter, - TrajectoryOutcomeFilter, - TrajectoryScorer, - ValueSelector, - merge_dropped_shards, - run_quality_pipeline, - select_top_for_rubric, - truncate_dropped_logs) -from twinkle_agentic.preprocessor import label_schema as L - -logger = get_logger() - -# ── Config ──────────────────────────────────────────────────────────────────── -CSV_PATH = os.environ.get('CSV_PATH', '/mnt/data/yzhao/tastelikefeet/bc/20260531.csv') -# Default 2000: pass-1 (filter + value_score) scales with N; with USE_RUBRIC=1 and -# SELECT_FRAC=0.1, rubric cost stays ~10% of survivors (similar to old 200×full rubric). -DATASET_TOTAL = int(os.environ.get('DATASET_TOTAL', 2000)) -MAP_NUM_PROC = int(os.environ.get('MAP_NUM_PROC', 8)) -OUTPUT_DIR = os.environ.get('OUTPUT_DIR', './output/data_pipeline') -_OUTPUT_BASENAME = os.environ.get( - 'OUTPUT_BASENAME', - f'processed_{Path(CSV_PATH).stem}_{DATASET_TOTAL}.jsonl', -) -OUTPUT_PATH = os.path.join(OUTPUT_DIR, _OUTPUT_BASENAME) -DROPPED_PATH = os.path.join(OUTPUT_DIR, 'dropped.jsonl') -PIPELINE_VERSION = os.environ.get('PIPELINE_VERSION', 'audit-v1') -# Set to keep only trajectories above this fused score. None -> keep all (inspect scores only). -MIN_TRAJ_SCORE = os.environ.get('MIN_TRAJ_SCORE') -MIN_TRAJ_SCORE = float(MIN_TRAJ_SCORE) if MIN_TRAJ_SCORE else None -# TrajectoryScorer runs hard-only by default (deterministic, fast). Set USE_RUBRIC=1 -# to attach a RubricVerifier so per-segment scores get real LLM semantic signal -# (needs LLM_BACKUP_* / OPENAI_API_KEY; much slower). Without it every clean -# trajectory tends to collapse to level 4 because only hard checks discriminate. -USE_RUBRIC = os.environ.get('USE_RUBRIC', '') not in ('', '0', 'false', 'False') -TRAJ_FUSION = os.environ.get('TRAJ_FUSION', 'hard_soft_blend') -TRAJ_HARD_CEIL_SKIP = os.environ.get('TRAJ_HARD_CEIL_SKIP') -TRAJ_HARD_CEIL_SKIP = float(TRAJ_HARD_CEIL_SKIP) if TRAJ_HARD_CEIL_SKIP else (0.92 if USE_RUBRIC else None) -TRAJ_SCORER_WORKERS = int(os.environ.get('TRAJ_SCORER_WORKERS', '2' if USE_RUBRIC else '1')) -# Rubric stabilization policy (reduces per-call score jitter for template-like -# intents). 'skeleton' = half-fixed core + generated tail (flexible, DEFAULT), -# 'fixed' = fully fixed per intent (max stability), 'off' = pure generation. -TRAJ_RUBRIC_MODE = os.environ.get('TRAJ_RUBRIC_MODE', 'skeleton').strip().lower() -# Active-learning pre-selection: only the global top SELECT_FRAC by value_score -# gets an (expensive) rubric pass; the rest are hard-only. 1.0 = label everyone -# (disables gating). Only takes effect with USE_RUBRIC=1. -SELECT_FRAC = float(os.environ.get('SELECT_FRAC', '0.1')) -SELECT_MIN = int(os.environ.get('SELECT_MIN', '0')) -SELECT_MAX = os.environ.get('SELECT_MAX') -SELECT_MAX = int(SELECT_MAX) if SELECT_MAX else None -# Persist the full per-segment rubric diagnosis (per-criterion verdict + reason + -# fix + raw teacher output) for rubric-scored rows — the SFT corpus to distill a -# PRM / error-checker LoRA later. Defaults ON when rubric labeling is enabled -# (one extra teacher call per scored segment). Set PERSIST_DIAGNOSIS=0 to skip. -PERSIST_DIAGNOSIS = os.environ.get( - 'PERSIST_DIAGNOSIS', '1' if USE_RUBRIC else '0') not in ('', '0', 'false', 'False') - - -# ── CSV ingestion (custom format: `ts,model,req_id,messages_json`) ───────────── -def _canonicalize_tool_call(tc: Any) -> Dict[str, Any]: - """Coerce a raw tool_call into a fixed-schema dict for stable Arrow inference.""" - tc = tc if isinstance(tc, dict) else {} - fn = tc.get('function') if isinstance(tc.get('function'), dict) else {} - args = fn.get('arguments') - if isinstance(args, dict): - args_str = json.dumps(args, ensure_ascii=False) - elif isinstance(args, str) and args.strip(): - try: - decoded = json.loads(args) - except json.JSONDecodeError: - decoded = {} - args_str = json.dumps(decoded if isinstance(decoded, dict) else {}, ensure_ascii=False) - else: - args_str = '{}' - return { - 'id': str(tc.get('id') or ''), - 'type': str(tc.get('type') or 'function'), - 'function': {'name': str(fn.get('name') or ''), 'arguments': args_str}, - } - - -def _stream_csv_rows(csv_path: str, max_rows: int = 0) -> Iterator[Dict[str, Any]]: - """Stream the custom CSV. First 3 fields are scalar; the rest of the line is a - JSON array of chat messages (may contain commas) — split on the first 3 commas. - - ``reasoning_content`` is folded into a ``...`` prefix so it - survives as visible content and is later re-exposed by MessageNormalizer (P3). - """ - emitted = 0 - with open(csv_path, 'rb') as f: - for raw in f: - try: - line = raw.decode('utf-8').rstrip('\n').rstrip('\r') - except UnicodeDecodeError: - continue - if not line: - continue - parts = line.split(',', 3) - if len(parts) < 4: - continue - ts, model, req_id, msgs_raw = parts - try: - raw_msgs = json.loads(msgs_raw) - except json.JSONDecodeError: - continue - messages: List[Dict[str, Any]] = [] - for m in raw_msgs: - role = m.get('role', '') - content = m.get('content') - if isinstance(content, list): - content = ''.join(p.get('text', '') for p in content - if isinstance(p, dict) and p.get('type') == 'text') - if content is None: - content = '' - if not isinstance(content, str): - continue - raw_tcs = m.get('tool_calls') if role == 'assistant' else None - tc_list = [_canonicalize_tool_call(tc) for tc in raw_tcs] if raw_tcs else [] - if role == 'assistant': - if not content and not tc_list: - continue - if m.get('reasoning_content'): - content = f"{m['reasoning_content']}{content}" - elif role != 'tool' and not content: - continue - messages.append({ - 'role': role, - 'content': content, - 'tool_calls': json.dumps(tc_list, ensure_ascii=False) if tc_list else '', - 'tool_call_id': str(m.get('tool_call_id') or '') if role == 'tool' else '', - }) - if not messages: - continue - yield { - 'id': f'csv__{ts}__{req_id}', - 'source': Path(csv_path).stem, - 'model_id': model, - 'messages': messages, - 'user_data': [], - # Pre-declare the only top-level column a downstream step adds - # (IntentClassifier). Without it the ingest schema lacks `intent`, - # and HF datasets.map(num_proc>1) infers per-shard features from - # the FIRST finished writer; shards that added `intent` later get - # that column dropped on concat -> rows with intent/value/prov all - # None (observed as 140/327 at 500 rows x num_proc=32). Declaring - # it up front keeps the Arrow schema identical across shards. - 'intent': None, - } - emitted += 1 - if max_rows and emitted >= max_rows: - break - - -def _build_trajectory_scorer() -> TrajectoryScorer: - """Hard-only by default; attach an LLM RubricVerifier when USE_RUBRIC is set.""" - rubric_verifier = None - if USE_RUBRIC: - from twinkle_agentic.verifier import (RubricVerifier, - default_intent_base_rubrics, - default_intent_fixed_rubrics) - # Intent-aware rubric stabilization: half-fixed skeleton (default) keeps - # the generator flexible while anchoring a shared core; 'fixed' drops - # generation entirely for tool_call/code/math; 'off' = pure generation. - intent_base = intent_fixed = None - if TRAJ_RUBRIC_MODE == 'skeleton': - intent_base = default_intent_base_rubrics() - elif TRAJ_RUBRIC_MODE == 'fixed': - intent_fixed = default_intent_fixed_rubrics() - rubric_verifier = RubricVerifier( - max_votes=5, - max_votes_long=3, - min_votes_long=2, - long_margin_threshold=0.18, - # Re-sample top-band segments so "looks perfect" isn't a lucky draw. - min_votes_high=3, - high_score_threshold=0.85, - intent_base_rubrics=intent_base, - intent_rubrics=intent_fixed, - ) - return TrajectoryScorer( - rubric_verifier=rubric_verifier, - fusion=TRAJ_FUSION, - hard_ceil_skip=TRAJ_HARD_CEIL_SKIP, - scorer_workers=TRAJ_SCORER_WORKERS, - reconcile_max_messages=int(os.environ.get('TRAJ_RECONCILE_MAX_MSGS', '80')), - # Persist the full rubric diagnosis (verdict+reason+fix+raw) for scored - # segments — the SFT corpus for a distilled PRM/checker LoRA. - persist_diagnosis=PERSIST_DIAGNOSIS, - ) - - -def build_pipeline_pass1() -> QualityPreprocessor: - """Pass 1: clean + tag + cheap value scoring (NO LLM rubric). - - Everything here is deterministic/parallel-safe. It ends by stamping a - ``value_score`` on every surviving row so the driver can then pick the global - top fraction for the (expensive) rubric pass. When rubric labeling is off, - the whole pipeline is a single pass and ValueSelector is skipped. - """ - steps = [ - # 0) lineage first, so even dropped rows carry provenance in the log. - ProvenanceStamp(source=Path(CSV_PATH).stem, pipeline_version=PIPELINE_VERSION), - # 1) canonicalize message schema (heartbeat strip, tool-call normalize, - # reasoning passthrough — P3), then structural / content filters. - MessageNormalizer(), - ModelFilter(), - LanguageFilter(allowed=('en', 'zh')), - # Shallow-chat round cap is 40; agent traces capped at 20 logical rounds - # (min user/assistant counts) for pipeline experiments — raise for prod. - HardFilter(min_user_chars_cjk=14, min_user_chars=24, max_rounds=40, - agent_max_rounds=20), - RefuseFilter(), - DeadLoopFilter(), - MessageSanityFilter(), - SpecialCharsFilter(max_ratio=0.6), - TokenSoupFilter(max_chars=8000), - # 2) taggers (never drop): intent key-rounds, structural noise ratio. - IntentClassifier(), - StructuralNoiseTagger(), - ] - if USE_RUBRIC and SELECT_FRAC < 1.0: - # Active-learning pre-selection: cheap value_score for top-fraction gating. - steps.append(ValueSelector()) - if not USE_RUBRIC: - # Single-pass mode: fold scoring + safety + outcome filter in here. - # No selection happened, so safety scores every row (gated=None). - steps += _pass2_tail(gated=False) - # drop_mode='mark': map returns equal-length columns (dropped rows flagged), - # and run_quality_pipeline materializes the removal via Dataset.filter — so a - # partially filtered batch can never leave ghost rows (no remove_columns hack). - return QualityPreprocessor(pipeline=steps, dropped_log_path=DROPPED_PATH, - drop_mode='mark') - - -def _pass2_tail(gated: bool) -> list: - """Scoring + safety + outcome filter (the LLM-touching tail). - - When ``gated`` is True (two-pass active-learning mode) both the rubric scorer - and the safety scorer only spend an LLM call on rows pre-selected by - ValueSelector (``selected_for_rubric``); everyone else is hard/neutral only. - """ - gate = L.KEY_SELECTED_FOR_RUBRIC if gated else None - return [ - _build_trajectory_scorer(), # hard-only, or LLM rubric when USE_RUBRIC=1 - # Fixed safety rubric; gated post-selection so the LLM safety pass runs - # only on selected rows (neutral-safe otherwise). - SafetyScorer(gate_label=gate), - # read-only outcome filter: drops on the tags above (D6). Enabled only - # when MIN_TRAJ_SCORE is set, else we keep everything to inspect. - TrajectoryOutcomeFilter( - min_traj_score=MIN_TRAJ_SCORE if MIN_TRAJ_SCORE is not None else 0.0, - min_safety_score=None, - drop_unsafe_flag=False, - ), - ] - - -def build_pipeline_pass2(gated: bool = True) -> QualityPreprocessor: - """Pass 2: rubric + safety (both gated on ``selected_for_rubric``) + filter.""" - return QualityPreprocessor(pipeline=_pass2_tail(gated=gated), - dropped_log_path=DROPPED_PATH, drop_mode='mark') - - -def _print_sample(dataset: Dataset, n: int = 3) -> None: - """Show the enrichment kept on a few rows so you can eyeball tags/scores.""" - hf = dataset.dataset - show = min(n, len(hf)) - logger.info(f'── sample of {show} processed rows (tags/scores in user_data) ──') - for i in range(show): - row = hf[i] - logger.info( - f"[{row.get('id')}] model={row.get('model_id')} n_msgs={len(row.get('messages') or [])}\n" - f" intent = {row.get('intent')}\n" - f" traj_score = {L.get_label(row, L.KEY_TRAJ_SCORE)} " - f"level={L.get_label(row, L.KEY_TRAJ_LEVEL)} " - f"conf={L.get_label(row, L.KEY_TRAJ_CONFIDENCE)}\n" - f" round_scores= {L.get_label(row, L.KEY_ROUND_SCORES)}\n" - f" safety = {L.get_label(row, L.KEY_SAFETY_SCORE)} " - f"(unsafe={L.get_label(row, L.KEY_SAFETY_UNSAFE)})\n" - f" noise_ratio = {L.get_label(row, 'structural_noise_ratio')}\n" - f" provenance = {L.get_label(row, L.KEY_PROVENANCE)}") - - -def main() -> None: - os.makedirs(OUTPUT_DIR, exist_ok=True) - logger.info(f'Loading up to {DATASET_TOTAL} rows from {CSV_PATH}') - - meta = DatasetMeta( - dataset_id=Path(CSV_PATH).stem, - data=partial(_stream_csv_rows, csv_path=CSV_PATH, max_rows=DATASET_TOTAL), - ) - dataset = Dataset(meta) - logger.info(f'Ingested {len(dataset.dataset)} rows.') - - truncate_dropped_logs(DROPPED_PATH) - # Pass 1: clean + tag (+ value_score when gating). run_quality_pipeline runs - # the pipeline as map(equal-length columns, dropped rows flagged) + a single - # Dataset.filter — map never changes row count, so no ghost rows can appear - # (unlike a filtering map, which needs remove_columns and still risks - # partial-batch ghosting). num_proc parallelizes across shards. - run_quality_pipeline(dataset, build_pipeline_pass1(), - num_proc=MAP_NUM_PROC, load_from_cache_file=False) - merge_dropped_shards(DROPPED_PATH) - logger.info(f'After pass 1 (clean+tag): {len(dataset.dataset)} rows kept.') - - if USE_RUBRIC: - gating = SELECT_FRAC < 1.0 - if gating: - # Global top-fraction: pick the most valuable rows for the LLM pass. - _, n_sel = select_top_for_rubric( - dataset, select_frac=SELECT_FRAC, - min_select=SELECT_MIN, max_select=SELECT_MAX) - logger.info(f'Value-gated rubric: {n_sel} rows selected for LLM labeling ' - f'(frac={SELECT_FRAC}).') - # Pass 2: rubric + safety (both gated when a selection ran) + outcome filter. - run_quality_pipeline(dataset, build_pipeline_pass2(gated=gating), - num_proc=MAP_NUM_PROC, load_from_cache_file=False) - merge_dropped_shards(DROPPED_PATH) - - logger.info(f'After pipeline: {len(dataset.dataset)} rows kept.') - _print_sample(dataset) - - dataset.save_as(OUTPUT_PATH, format='jsonl') - logger.info(f'Saved processed dataset -> {OUTPUT_PATH}') - logger.info(f'Dropped rows log -> {DROPPED_PATH}') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/legacy/dataset_hard.py b/cookbook/exp/legacy/dataset_hard.py deleted file mode 100644 index 9fa059b95..000000000 --- a/cookbook/exp/legacy/dataset_hard.py +++ /dev/null @@ -1,202 +0,0 @@ -"""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/legacy/dataset_index.py b/cookbook/exp/legacy/dataset_index.py deleted file mode 100644 index 7d2905a59..000000000 --- a/cookbook/exp/legacy/dataset_index.py +++ /dev/null @@ -1,718 +0,0 @@ -"""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/legacy/dataset_think.py b/cookbook/exp/legacy/dataset_think.py deleted file mode 100644 index 38618ced1..000000000 --- a/cookbook/exp/legacy/dataset_think.py +++ /dev/null @@ -1,456 +0,0 @@ -import hashlib -import re -from typing import Any, Dict, List, Optional - -from twinkle.dataset import Dataset, DatasetMeta -from twinkle.preprocessor import Preprocessor - -_THINK_RE = re.compile(r'(.*?)', re.DOTALL) - - -def _hash_id(prefix: str, content: str) -> str: - return f'{prefix}__{hashlib.md5(content.encode("utf-8")).hexdigest()[:16]}' - - -def _register(dataset, processor_cls, meta: DatasetMeta, init_args: Optional[Dict[str, Any]] = None, - load_from_cache_file: bool = True) -> None: - """Add dataset and run preprocessor; auto-strip every input column to enforce - the universal ``{id, source, query, cot, response}`` output schema.""" - dataset.add_dataset(meta) - cols = list(dataset.datasets[meta.get_id()].column_names) - dataset.map( - processor_cls, - dataset_meta=meta, - init_args=init_args or {}, - remove_columns=cols, - load_from_cache_file=load_from_cache_file, - ) - - -# ===== Modotte/CodeX-2M-Thinking ===== -CODEX_THINKING_REPO = 'ms://Modotte/CodeX-2M-Thinking' - - -class CodeXThinkingProcessor(Preprocessor): - """CodeX-2M-Thinking row → ``{id, source, query, cot, response}``。 - - 输入 schema: ``input``（问题）、``output``（含 ``...`` + 答案）。 - 拆分 output 为 cot（think 标签内容）和 response（标签之后的正文）。 - 丢弃缺失 input/output 或无法解析 think 标签的行。 - """ - - 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 '').strip() - output = (row.get('output') or '').strip() - if not query or not output: - continue - m = _THINK_RE.search(output) - if not m: - continue - cot = m.group(1).strip() - response = output[m.end():].strip() - if not cot or not response: - continue - out.append({ - 'id': _hash_id('codex_think', f'{query}\n{response}'), - 'source': 'CodeX-2M-Thinking', - 'query': query, - 'cot': cot, - 'response': response, - }) - return self.map_row_to_col(out) - - -# ===== open-thoughts/OpenThoughts3-1.2M ===== -OPEN_THOUGHTS_REPO = 'ms://open-thoughts/OpenThoughts3-1.2M' - - -class OpenThoughtsProcessor(Preprocessor): - """OpenThoughts3 row → ``{id, source, query, cot, response}``。 - - 输入 schema: ``conversations`` (messages 格式 list[{from/value}])。 - 取第一个 human 作 query，第一个 gpt 的 value 按 ``...`` 拆 cot/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: - convs = row.get('conversations') - if not isinstance(convs, list): - continue - query = '' - assistant_text = '' - for msg in convs: - if not isinstance(msg, dict): - continue - role = msg.get('from') or msg.get('role') or '' - value = msg.get('value') or msg.get('content') or '' - if role in ('human', 'user') and not query: - query = value.strip() - elif role in ('gpt', 'assistant') and not assistant_text: - assistant_text = value.strip() - break - if not query or not assistant_text: - continue - m = _THINK_RE.search(assistant_text) - if not m: - continue - cot = m.group(1).strip() - response = assistant_text[m.end():].strip() - if not cot or not response: - continue - out.append({ - 'id': _hash_id('openthoughts', f'{query}\n{response}'), - 'source': 'OpenThoughts3-1.2M', - 'query': query, - 'cot': cot, - 'response': response, - }) - return self.map_row_to_col(out) - - -# ===== GAIR/LIMO-v2 ===== -LIMO_REPO = 'ms://GAIR/LIMO-v2' - - -class LIMOProcessor(Preprocessor): - """LIMO-v2 row → ``{id, source, query, cot, response}``。 - - 输入 schema: ``question``、``solution``（含 ``...`` + 答案）。 - 拆分 solution 为 cot 和 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('question') or '').strip() - solution = (row.get('solution') or '').strip() - if not query or not solution: - continue - m = _THINK_RE.search(solution) - if m: - cot = m.group(1).strip() - response = solution[m.end():].strip() - else: - # 无 think 标签时，solution 整体作为 response，cot 留空 - cot = '' - response = solution - if not response: - continue - out.append({ - 'id': _hash_id('limo', f'{query}\n{response}'), - 'source': 'LIMO-v2', - 'query': query, - 'cot': cot, - 'response': response, - }) - return self.map_row_to_col(out) - - -# ===== AI-ModelScope/Chinese-DeepSeek-R1-Distill-data-110k ===== -CN_R1_DISTILL_REPO = 'ms://AI-ModelScope/Chinese-DeepSeek-R1-Distill-data-110k' - - -class ChineseR1DistillProcessor(Preprocessor): - """Chinese-DeepSeek-R1-Distill row → ``{id, source, query, cot, response}``。 - - 输入已有三列: ``input`` → query, ``reasoning_content`` → cot, ``content`` → 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('input') or '').strip() - cot = (row.get('reasoning_content') or '').strip() - response = (row.get('content') or '').strip() - if not query or not response: - continue - if cot: - response = _THINK_RE.sub('', response).strip() - if not response: - continue - out.append({ - 'id': _hash_id('cn_r1_distill', f'{query}\n{response}'), - 'source': 'Chinese-DeepSeek-R1-Distill-data-110k', - 'query': query, - 'cot': cot, - 'response': response, - }) - return self.map_row_to_col(out) - - -# ===== nohurry/Opus-4.6-Reasoning-3000x-filtered ===== -OPUS_REASONING_REPO = 'ms://nohurry/Opus-4.6-Reasoning-3000x-filtered' - - -class OpusReasoningProcessor(Preprocessor): - """Opus-4.6-Reasoning-3000x-filtered row → ``{id, source, query, cot, response}``。 - - 输入已有三列: ``problem`` → query, ``thinking`` → cot, ``solution`` → 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 '').strip() - cot = (row.get('thinking') or '').strip() - response = (row.get('solution') or '').strip() - if not query or not response: - continue - if cot: - response = _THINK_RE.sub('', response).strip() - if not response: - continue - out.append({ - 'id': _hash_id('opus_reasoning', f'{query}\n{response}'), - 'source': 'Opus-4.6-Reasoning-3000x-filtered', - 'query': query, - 'cot': cot, - 'response': response, - }) - return self.map_row_to_col(out) - - -# ===== Roman1111111/claude-opus-4.6-10000x ===== -CLAUDE_OPUS_REPO = 'ms://Roman1111111/claude-opus-4.6-10000x' - - -class ClaudeOpusProcessor(Preprocessor): - """claude-opus-4.6-10000x row → ``{id, source, query, cot, response}``。 - - 输入 schema: ``messages`` (OpenAI 格式 list[{role, content}])。 - 取首个 user 作 query，首个 assistant 按 ``...`` 拆 cot/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: - messages = row.get('messages') - if not isinstance(messages, list): - continue - query = '' - assistant_text = '' - for msg in messages: - if not isinstance(msg, dict): - continue - role = msg.get('role') or '' - content = msg.get('content') or '' - if not isinstance(content, str): - continue - if role == 'user' and not query: - query = content.strip() - elif role == 'assistant' and not assistant_text: - assistant_text = content.strip() - break - if not query or not assistant_text: - continue - m = _THINK_RE.search(assistant_text) - if m: - cot = m.group(1).strip() - response = assistant_text[m.end():].strip() - else: - cot = '' - response = assistant_text - if not response: - continue - out.append({ - 'id': _hash_id('claude_opus', f'{query}\n{response}'), - 'source': 'claude-opus-4.6-10000x', - 'query': query, - 'cot': cot, - 'response': response, - }) - return self.map_row_to_col(out) - - -ANGRYGIRAFFE_REPO = 'ms://hf/angrygiraffe-claude-opus-4.6-4.7-reasoning-8.7k' - - -class AngrygiraffeOpusReasoningProcessor(Preprocessor): - """angrygiraffe/claude-opus-4.6-4.7-reasoning-8.7k row → ``{id, source, query, cot, response}``。 - - 输入 schema: ``messages`` (OpenAI 格式 list[{role, content}])。 - 取首个 user 作 query，首个 assistant 按 ``...`` 拆 cot/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: - messages = row.get('messages') - if not isinstance(messages, list): - continue - query = '' - assistant_text = '' - for msg in messages: - if not isinstance(msg, dict): - continue - role = msg.get('role') or '' - content = msg.get('content') or '' - if not isinstance(content, str): - continue - if role == 'user' and not query: - query = content.strip() - elif role == 'assistant' and not assistant_text: - assistant_text = content.strip() - break - if not query or not assistant_text: - continue - m = _THINK_RE.search(assistant_text) - if m: - cot = m.group(1).strip() - response = assistant_text[m.end():].strip() - else: - cot = '' - response = assistant_text - if not response: - continue - out.append({ - 'id': _hash_id('angrygiraffe_opus', f'{query}\n{response}'), - 'source': 'angrygiraffe-claude-opus-4.6-4.7-reasoning-8.7k', - 'query': query, - 'cot': cot, - 'response': response, - }) - return self.map_row_to_col(out) - - -_BASE_SIZES = { - 'codex_think': 100000, - 'open_thoughts': 400000, - 'cn_r1_distill': 100000, - 'opus_reasoning': 3000, - 'claude_opus': 10000, - 'angrygiraffe': 38000, -} - - -def _scaled_sizes(total: Optional[int]) -> Dict[str, int]: - if total is None: - 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, CodeXThinkingProcessor, - DatasetMeta(dataset_id=CODEX_THINKING_REPO, split='train', - data_slice=range(sizes['codex_think'])), - load_from_cache_file=load_from_cache_file) - - _register(dataset, OpenThoughtsProcessor, - DatasetMeta(dataset_id=OPEN_THOUGHTS_REPO, split='train', - data_slice=range(sizes['open_thoughts'])), - load_from_cache_file=load_from_cache_file) - - _register(dataset, LIMOProcessor, - DatasetMeta(dataset_id=LIMO_REPO, split='train'), - load_from_cache_file=load_from_cache_file) - - _register(dataset, ChineseR1DistillProcessor, - DatasetMeta(dataset_id=CN_R1_DISTILL_REPO, split='train', - data_slice=range(sizes['cn_r1_distill'])), - load_from_cache_file=load_from_cache_file) - - _register(dataset, OpusReasoningProcessor, - DatasetMeta(dataset_id=OPUS_REASONING_REPO, split='train', - data_slice=range(sizes['opus_reasoning'])), - load_from_cache_file=load_from_cache_file) - - _register(dataset, ClaudeOpusProcessor, - DatasetMeta(dataset_id=CLAUDE_OPUS_REPO, split='train', - data_slice=range(sizes['claude_opus'])), - load_from_cache_file=load_from_cache_file) - - _register(dataset, AngrygiraffeOpusReasoningProcessor, - DatasetMeta(dataset_id=ANGRYGIRAFFE_REPO, split='train', - data_slice=range(sizes['angrygiraffe'])), - load_from_cache_file=load_from_cache_file) - - dataset.mix_dataset(False) - return dataset - - -class ToMessagesProcessor(Preprocessor): - """Convert {query, cot, response} → {id, source, messages}.""" - - 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('query') or '' - cot = row.get('cot') or '' - response = row.get('response') or '' - if not cot: - continue - assistant_content = f'{cot}' - out.append({ - 'id': row.get('id', ''), - 'source': row.get('source', ''), - 'messages': [ - {'role': 'user', 'content': query}, - {'role': 'assistant', 'content': assistant_content, - 'reasoning_content': cot}, - ], - }) - return self.map_row_to_col(out, keys=['id', 'source', 'messages']) - - -def get_dataset(total: Optional[int] = None, dropped_log: Optional[str] = None, - load_from_cache_file: bool = True) -> Dataset: - """Build, convert to messages format, and quality-filter the CoT dataset. - - If ``total`` is given, every per-source row count in ``_BASE_SIZES`` is - scaled proportionally so the input-row sum approximates ``total``. - """ - from twinkle_agentic.preprocessor import ( - DeadLoopFilter, - FixUnicodeFilter, - HardFilter, - IntentClassifier, - MessageSanityFilter, - QualityPreprocessor, - RefuseFilter, - RemoveRepeatSentencesFilter, - TokenNumFilter, - TokenSoupFilter, - ) - - dataset = _build_dataset(total=total, 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.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/legacy/eval_dualline_math.py b/cookbook/exp/legacy/eval_dualline_math.py deleted file mode 100644 index bd1c183c7..000000000 --- a/cookbook/exp/legacy/eval_dualline_math.py +++ /dev/null @@ -1,689 +0,0 @@ -"""Dual-line math evaluation: baseline vs online process-checking + rubric injection. - -This is **Phase 0 of DESIGN §11.6** ("参数化 memory: 查错 LoRA"): before training any -LoRA, test the *upper bound* of the mechanism "pause every N tokens, let a strong -teacher check the partial reasoning for rubric errors, inject the found issue back -into the context, then resume". If even the strongest teacher checking online cannot -lift math accuracy, distilling that ability into a LoRA is pointless — so we gate on -this first. - -It deliberately reuses the SAME dataset loader, sampling params and answer grader as -``eval_gpqa_rag.py`` so the two lines are directly comparable: - - - **Line A — baseline** (``--mode baseline``): the student model solves each problem - in a single pass (identical to ``eval_gpqa_rag.py --mode direct``). - - **Line B — dualline** (``--mode dualline``, default): the student generates in - ``--chunk-tokens`` slices; between slices a teacher ``RubricVerifier.diagnose()`` - inspects the full reasoning so far (query + all prior response). When it reports - process issues, the finding is injected back as a first-person self-correction - (in the student's own voice) and generation resumes. - -The teacher checker is the ``llm_backup`` teacher API (no student sampler is given to -the verifier, so every check is served by the teacher — exactly the Phase-0 setup). -Configure it via the ``LLM_BACKUP_*`` env vars (see ``utils/llm_backup.py``). - -Continuation is done at the token level (crude on purpose — §11.6 says experiment -performance is not a concern): each slice re-feeds the prior ``new_input_feature`` and, -on injection, splices the tokenized note in before resuming. - -The dataset defaults to AoPS (``--dataset aops``), which auto-downloads from -ModelScope so no local data path is needed; pass ``--dataset math`` to use the -local Hendrycks MATH set instead. Both lines MUST share ``--dataset``, ``--n``, -``--target-eval`` and ``--seed`` to stay a paired comparison. - -Launch examples: - # Dual-line on 200 AoPS problems (needs LLM_BACKUP_* for the teacher checker) - LLM_BACKUP_API_KEY=sk-... LLM_BACKUP_BASE_URL=... \\ - python cookbook/exp/embedding/eval_dualline_math.py \\ - --n 200 --target-eval 200 --seed 42 - - # Paired baseline on the same subset (no checker calls) - python cookbook/exp/embedding/eval_dualline_math.py --mode baseline \\ - --n 200 --target-eval 200 --seed 42 -""" -import argparse -import copy -import json -import os -import sys -import time -from collections import defaultdict -from typing import Any, Dict, List, Optional - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.data_format import SamplingParams as TwinkleSamplingParams -from twinkle.sampler import vLLMSampler - -# Reuse the reference eval's dataset + grading + prompts verbatim so the two -# lines are measured on identical footing. -from eval_gpqa_rag import (GEN_MODEL_ID, GEN_GPU_MEM, GEN_GPUS, GEN_TEMPERATURE, - GEN_TOP_P, answers_match, build_direct_prompt, - extract_boxed, load_aops, load_math) - -# Dualline eval defaults (override via --max-model-len or DUALLINE_MAX_MODEL_LEN). -DUALLINE_DEFAULT_MAX_MODEL_LEN = int(os.environ.get('DUALLINE_MAX_MODEL_LEN', 32000)) -DUALLINE_DEFAULT_MAX_GEN_TOKENS = int( - os.environ.get('DUALLINE_MAX_GEN_TOKENS', DUALLINE_DEFAULT_MAX_MODEL_LEN)) - -# vLLM parallel: default tp=1, dp=GEN_GPUS (override with GEN_TP / keep GEN_GPUS=8). -GEN_TP = int(os.environ.get('GEN_TP', 1)) - -logger = get_logger() - -# --------------------------------------------------------------------------- -# Dual-line config -# --------------------------------------------------------------------------- -CHUNK_TOKENS = int(os.environ.get('DUALLINE_CHUNK_TOKENS', 512)) -MAX_CHECKS = int(os.environ.get('DUALLINE_MAX_CHECKS', 8)) -MAX_INJECTIONS = int(os.environ.get('DUALLINE_MAX_INJECTIONS', 3)) -# Only inject when the checker is confident enough that something is wrong. -CHECK_SCORE_FLOOR = float(os.environ.get('DUALLINE_CHECK_FLOOR', 0.6)) -# The note is written in the student's own first-person voice so, when spliced -# back in, the running model treats it as its own mid-thought self-correction -# rather than an external interruption (which tended to derail generation toward -# max-length). Kept short to limit disruption. -INJECT_TEMPLATE = ( - '\n\nWait — reviewing my reasoning above, I realize there is a problem: {issue}\n' - 'Let me correct this and continue.\n\n') - -# When context hits max_model_len (or sample fails), dump query + generation here. -OVERFLOW_DUMP_DIR = os.environ.get( - 'DUALLINE_OVERFLOW_DUMP_DIR', './output/dualline/overflow_dumps') - - -def _decode(tokenizer, ids: List[int]) -> str: - return tokenizer.decode(ids, skip_special_tokens=True) - - -def _input_ids_len(cur_inputs: Any) -> Optional[int]: - """Length of the tokenized prompt fed to vLLM on this step, if known.""" - if not cur_inputs: - return None - item = cur_inputs[0] - if isinstance(item, dict) and 'input_ids' in item: - ids = item['input_ids'] - return len(ids) if ids is not None else None - return None - - -def _dump_dualline_state( - *, - reason: str, - problem: str, - debug_idx: Optional[int], - chunk_tokens: int, - cur_inputs: Any, - gen_ids: List[int], - injected_ids: List[int], - tokenizer, - n_checks: int, - n_injections: int, - findings: List[Dict[str, Any]], - total_new: int, - finished: bool, - max_model_len: int, - error: Optional[str] = None, -) -> str: - """Persist state for post-mortem (student CoT vs checker injection). Returns path.""" - os.makedirs(OVERFLOW_DUMP_DIR, exist_ok=True) - tag = f'idx{debug_idx}' if debug_idx is not None else 'idx_unknown' - path = os.path.join( - OVERFLOW_DUMP_DIR, f'{tag}_{reason}_{int(time.time())}.json') - - partial_cot = _decode(tokenizer, gen_ids) if tokenizer and gen_ids else '' - injected_text = (_decode(tokenizer, injected_ids) - if tokenizer and injected_ids else '') - ctx_len = _input_ids_len(cur_inputs) - - payload: Dict[str, Any] = { - 'reason': reason, - 'error': error, - 'query': problem, - 'debug_idx': debug_idx, - 'gen_token_count': len(gen_ids), - 'injected_token_count': len(injected_ids), - 'context_input_ids_len': ctx_len, - 'max_model_len': max_model_len, - 'chunk_tokens': chunk_tokens, - 'total_new': total_new, - 'n_checks': n_checks, - 'n_injections': n_injections, - 'findings': findings, - 'finished': finished, - 'partial_cot': partial_cot, - 'injected_text': injected_text, - 'partial_cot_chars': len(partial_cot), - 'context_is_message_prompt': ctx_len is None, - } - with open(path, 'w', encoding='utf-8') as f: - json.dump(payload, f, ensure_ascii=False, indent=2) - cot_path = path.replace('.json', '_partial_cot.txt') - with open(cot_path, 'w', encoding='utf-8') as f: - f.write(partial_cot) - sys.stderr.write(f'[dualline] overflow dump -> {path}\n') - return path - -# --------------------------------------------------------------------------- -# Teacher checker (Phase-0: pure teacher via llm_backup) -# --------------------------------------------------------------------------- -def _build_checker(): - """RubricVerifier with no student sampler -> every diagnose() hits the teacher. - - Uses a fixed, math-oriented process rubric so we do not spend a rubric- - generation call per slice (the segment here is a partial CoT, not a finished - trajectory). Falls back to auto-generated rubrics if fixed_rubric is cleared. - """ - from twinkle_agentic.verifier import RubricVerifier - from twinkle_agentic.verifier.rubric_verifier import RubricItem - - fixed = [ - RubricItem('The reasoning contains no arithmetic or algebraic error so far', - is_hard=True), - RubricItem('Each step follows logically from the previous ones', is_hard=True), - RubricItem('No formula or theorem is misstated or misapplied', is_hard=True), - RubricItem('The approach is on track to answer the actual question asked', - is_hard=False), - RubricItem('No step contradicts an earlier established fact', is_hard=False), - ] - return RubricVerifier(fixed_rubric=fixed, gate=True) - - -def _checker_available() -> bool: - return bool(os.environ.get('LLM_BACKUP_API_KEY') - or os.environ.get('LLM_BACKUP_BASE_URL') - or os.environ.get('OPENAI_API_KEY')) - - -def _diagnose_partial(checker, problem: str, partial_cot: str): - """Run the teacher checker on the reasoning so far; return (issue_or_None, detail). - - ``partial_cot`` is the FULL reasoning generated so far (all prior chunks plus - any self-corrections already spliced in), not just the latest slice, so the - teacher judges the whole derivation in context. We label it as in-progress so - it grades correctness of the steps rather than penalizing the absence of a - final answer. - """ - seg_content = ( - '[The following is the full reasoning so far, still in progress and not ' - 'yet complete. Judge only whether the reasoning up to this point is ' - 'mathematically correct; do not expect a final answer here.]\n\n' + partial_cot) - seg = {'messages': [ - {'role': 'user', 'content': problem}, - {'role': 'assistant', 'content': seg_content}, - ]} - try: - detail = checker.diagnose(seg, query=problem) - except Exception as exc: - logger.warning(f'[dualline] checker error: {exc}') - return None, None - if detail.overall_ok: - return None, detail - if detail.scalar >= CHECK_SCORE_FLOOR: - # Checker leans "mostly fine"; don't disrupt on a marginal signal. - return None, detail - fails = [it for it in detail.items if not it.verdict] - if not fails: - return None, detail - # Prefer a fix if the checker gave one; else the reason. - parts = [] - for it in fails[:2]: - msg = it.fix or it.reason - if msg: - parts.append(msg) - issue = ' '.join(parts).strip() or detail.summary - return (issue or None), detail - - -def _pad_batch_for_dp(items: List[Any], gen_dp: int) -> List[Any]: - """``slice_dp`` needs batch len >= DP world size (every rank gets work). - - Only kicks in on the tail rounds when fewer than ``gen_dp`` problems are - still active; the padded replicas are dropped by the caller. - """ - if gen_dp <= 1 or not items or len(items) >= gen_dp: - return items - pad = [copy.deepcopy(items[-1]) for _ in range(gen_dp - len(items))] - return items + pad - - -class _DualState: - """Per-problem generation state for the batched dualline loop. - - All problems advance together, one ``chunk_tokens`` slice per round. A - problem stays *active* until it emits EOS, hits ``max_gen_tokens``, would - overflow ``max_model_len``, or a sample call fails. Because the problems - share every round's ``sampler.sample`` call, the vLLM engine batches them - (and, with dp>1, spreads them across ranks) instead of running one at a - time. - """ - - __slots__ = ('idx', 'problem', 'cur_input', 'gen_ids', 'injected_ids', - 'n_checks', 'n_injections', 'findings', 'total_new', - 'finished', 'stopped_reason', 'context_input_ids_len', - 'pending_partial_cot', 'prompt_len') - - def __init__(self, idx: int, problem: str, prompt: Any): - self.idx = idx - self.problem = problem - self.cur_input: Any = prompt # str prompt (round 0) or input_feature - self.gen_ids: List[int] = [] # student-generated token ids only - self.injected_ids: List[int] = [] # spliced-in ids (excluded from answer) - self.n_checks = 0 - self.n_injections = 0 - self.findings: List[Dict[str, Any]] = [] - self.total_new = 0 - self.finished = False - self.stopped_reason: Optional[str] = None - self.context_input_ids_len: Optional[int] = None - self.pending_partial_cot: Optional[str] = None - self.prompt_len: Optional[int] = None # token len of the fixed prompt prefix - - def cur_input_len(self) -> Optional[int]: - item = self.cur_input - if isinstance(item, dict) and 'input_ids' in item: - ids = item['input_ids'] - return len(ids) if ids is not None else None - return None - - def result(self, tokenizer) -> Dict[str, Any]: - if self.context_input_ids_len is None: - self.context_input_ids_len = self.cur_input_len() - return { - 'text': _decode(tokenizer, self.gen_ids), - 'finished': self.finished, - 'stopped_reason': self.stopped_reason, - 'context_input_ids_len': self.context_input_ids_len, - 'n_checks': self.n_checks, - 'n_injections': self.n_injections, - 'findings': self.findings, - 'gen_tokens': len(self.gen_ids), - } - - -def _dump_state_obj(st: '_DualState', tokenizer, chunk_tokens: int, - max_model_len: int, reason: str, error: str) -> None: - _dump_dualline_state( - reason=reason, - problem=st.problem, - debug_idx=st.idx, - chunk_tokens=chunk_tokens, - cur_inputs=[st.cur_input], - gen_ids=st.gen_ids, - injected_ids=st.injected_ids, - tokenizer=tokenizer, - n_checks=st.n_checks, - n_injections=st.n_injections, - findings=st.findings, - total_new=st.total_new, - finished=st.finished, - max_model_len=max_model_len, - error=error, - ) - - -# --------------------------------------------------------------------------- -# Batched token-level segmented generation with mid-stream injection -# --------------------------------------------------------------------------- -def run_dualline_batch(sampler, tokenizer, problems: List[str], checker, - base_params: TwinkleSamplingParams, - chunk_tokens: int, - max_model_len: int, - max_gen_tokens: int, - gen_dp: int = 1, - diagnose_workers: int = 8) -> List[Dict[str, Any]]: - """Advance every problem in lock-step slices, sharing one sampler call/round. - - Each round: (1) preflight-drop any problem that would overflow the context, - (2) one ``sampler.sample`` over all still-active problems (vLLM batches + - spreads over dp ranks), (3) for the length-capped ones, run the teacher - diagnoses concurrently and splice injections, then loop. - - Returns per-problem result dicts in the original ``problems`` order. - """ - from concurrent.futures import ThreadPoolExecutor - - chunk_params = TwinkleSamplingParams( - max_tokens=chunk_tokens, temperature=base_params.temperature, - top_p=base_params.top_p, num_samples=1) - - states = [_DualState(i, p, build_direct_prompt(p)) - for i, p in enumerate(problems)] - active = list(states) - round_no = 0 - - while active: - round_no += 1 - - # (1) Preflight: drop problems that would overflow the context window, - # and those that already reached the generation-token cap. - survivors: List[_DualState] = [] - for st in active: - if st.total_new >= max_gen_tokens: - st.stopped_reason = st.stopped_reason or 'max_gen_tokens' - continue - ctx_len = st.cur_input_len() - if ctx_len is not None and ctx_len + chunk_tokens >= max_model_len: - st.context_input_ids_len = ctx_len - st.stopped_reason = 'context_full' - _dump_state_obj( - st, tokenizer, chunk_tokens, max_model_len, - reason='preflight_context_full', - error=(f'context len {ctx_len} + chunk {chunk_tokens} ' - f'>= max_model_len {max_model_len}')) - continue - survivors.append(st) - active = survivors - if not active: - break - - # (2) One shared sampler call over all active problems. On tail rounds - # with fewer active problems than dp ranks, pad to keep slice_dp happy - # and drop the padded responses. The context-overflow preflight above - # guarantees every input still fits, so a length-capped slice should - # never raise here; let any genuine engine error propagate instead of - # masking it as a whole-round failure. - batch_inputs = [st.cur_input for st in active] - padded = _pad_batch_for_dp(batch_inputs, gen_dp) - responses = sampler.sample(padded, chunk_params) - responses = responses[:len(active)] - - # (3) Consume each problem's slice; queue the ones needing a check. - to_diagnose: List[_DualState] = [] - next_active: List[_DualState] = [] - for st, resp in zip(active, responses): - seq = resp.sequences[0] if resp and resp.sequences else None - if seq is None: - st.stopped_reason = st.stopped_reason or 'empty_response' - continue - st.gen_ids.extend(seq.tokens) - st.total_new += len(seq.tokens) - st.cur_input = seq.new_input_feature - if st.prompt_len is None: - # Fixed prompt prefix = everything before this round's generation. - st.prompt_len = len(st.cur_input['input_ids']) - len(seq.tokens) - - if seq.stop_reason != 'length': - st.finished = True # EOS / stop -> done - continue - if st.n_checks >= MAX_CHECKS or not checker: - next_active.append(st) # keep generating, no more checks - continue - # Diagnose the FULL reasoning generated so far (all prior chunks plus - # any self-corrections already spliced in), so the teacher judges the - # whole derivation in context rather than an isolated tail slice. - st.pending_partial_cot = _decode( - tokenizer, st.cur_input['input_ids'][st.prompt_len:]) - st.n_checks += 1 - to_diagnose.append(st) - - # Concurrent teacher diagnoses for this round's length-capped problems. - if to_diagnose: - def _run(st: _DualState): - return st, _diagnose_partial( - checker, st.problem, st.pending_partial_cot) - workers = max(1, min(diagnose_workers, len(to_diagnose))) - with ThreadPoolExecutor(max_workers=workers) as ex: - for st, (issue, _detail) in ex.map(_run, to_diagnose): - st.pending_partial_cot = None - if issue and st.n_injections < MAX_INJECTIONS: - note = INJECT_TEMPLATE.format(issue=issue) - note_ids = tokenizer.encode(note, add_special_tokens=False) - feat = dict(st.cur_input) - feat['input_ids'] = list(feat['input_ids']) + note_ids - if 'labels' in feat: - feat['labels'] = list(feat['labels']) + note_ids - st.cur_input = feat - st.injected_ids.extend(note_ids) - st.n_injections += 1 - st.findings.append( - {'at_token': st.total_new, 'issue': issue}) - next_active.append(st) - - active = next_active - n_done = sum(1 for s in states if s.finished or s.stopped_reason) - sys.stderr.write( - f'[dualline] round {round_no}: active={len(active)} ' - f'done={n_done}/{len(states)}\n') - - return [st.result(tokenizer) for st in states] - - -def _load_tokenizer(model_id: str): - """Load the tokenizer from ModelScope (matches the vLLM sampler source). - - The box runs offline, so ``transformers.AutoTokenizer`` (which resolves via - the HF hub) fails with ``Network is unreachable``. ModelScope's AutoTokenizer - downloads/reads from the ModelScope cache instead — the same place the vLLM - sampler already pulled the model from. Falls back to transformers only if the - ModelScope path is unavailable. - """ - try: - from modelscope import AutoTokenizer as MSAutoTokenizer - return MSAutoTokenizer.from_pretrained(model_id, trust_remote_code=True) - except Exception as exc: - sys.stderr.write(f'[dualline] modelscope tokenizer load failed ({exc}); ' - f'falling back to transformers\n') - from transformers import AutoTokenizer - return AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) - - -# --------------------------------------------------------------------------- -# Main -# --------------------------------------------------------------------------- -def main(): - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--mode', choices=['baseline', 'dualline'], default='dualline') - p.add_argument('--dataset', choices=['aops', 'math'], default='aops', - help='Evaluation dataset. "aops" (default) auto-downloads from ' - 'ModelScope (no local path needed); "math" reads local ' - 'MATH_DATA_DIR, stratified by difficulty level.') - p.add_argument('--math-split', default='test') - p.add_argument('--per-level', type=int, default=0, - help='MATH only: problems per level. 0 => --n split across levels.') - p.add_argument('--n', type=int, default=32, - help='Pool size sampled from the dataset (MATH is stratified ' - 'by level; AoPS is a flat shuffle).') - p.add_argument('--target-eval', type=int, default=32, - help='Stop after this many problems are evaluated (0 = all sampled).') - p.add_argument('--max-model-len', type=int, default=DUALLINE_DEFAULT_MAX_MODEL_LEN, - help='vLLM max_model_len / template max_length (default 32000).') - p.add_argument('--max-gen-tokens', type=int, default=DUALLINE_DEFAULT_MAX_GEN_TOKENS, - help='Cap total generated tokens per problem (default: same as ' - 'max-model-len / DUALLINE_MAX_GEN_TOKENS).') - p.add_argument('--chunk-tokens', type=int, default=CHUNK_TOKENS, - help='Generate this many tokens between checker pauses.') - p.add_argument('--batch-size', type=int, default=16, - help='Baseline mode batch size. Dualline runs all problems ' - 'concurrently (one shared sampler call per slice-round).') - p.add_argument('--diagnose-workers', type=int, - default=int(os.environ.get('DUALLINE_DIAGNOSE_WORKERS', 8)), - help='Concurrency for teacher diagnose() calls within a round.') - 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/dualline/{args.dataset}_{args.mode}_results.jsonl' - - is_dual = (args.mode == 'dualline') - if is_dual and not _checker_available(): - sys.stderr.write( - '[dualline] ERROR: --mode dualline needs a teacher checker but no ' - 'LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / OPENAI_API_KEY is set.\n' - ' Set them, or run --mode baseline for the paired baseline.\n') - sys.exit(1) - - 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) - if args.target_eval > 0: - records = records[:args.target_eval] - max_model_len = args.max_model_len - max_gen_tokens = args.max_gen_tokens - sys.stderr.write( - f'[dualline] evaluating {len(records)} problems ' - f'(mode={args.mode}, dataset={args.dataset}, ' - f'max_model_len={max_model_len}, max_gen_tokens={max_gen_tokens})\n') - - device_groups = [ - DeviceGroup(name='sampler', ranks=list(range(GEN_GPUS)), - device_type='GPU', gpus_per_worker=GEN_TP), - ] - if GEN_GPUS % GEN_TP != 0: - raise ValueError(f'GEN_GPUS ({GEN_GPUS}) must be divisible by GEN_TP ({GEN_TP})') - gen_dp = GEN_GPUS // GEN_TP - gen_mesh = DeviceMesh.from_sizes(world_size=GEN_GPUS, dp_size=gen_dp, tp_size=GEN_TP) - 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': max_model_len, - 'tensor_parallel_size': GEN_TP, - }, - device_mesh=gen_mesh, - remote_group='sampler', - ) - sampler.set_template('Qwen3_5Template', model_id=GEN_MODEL_ID, - enable_thinking=True, max_length=max_model_len) - sys.stderr.write( - f'[dualline] vLLM sampler ready (model={GEN_MODEL_ID}, ' - f'tp={GEN_TP}, dp={gen_dp})\n') - - gen_params = TwinkleSamplingParams( - max_tokens=max_gen_tokens, temperature=GEN_TEMPERATURE, - top_p=GEN_TOP_P, num_samples=1) - - checker = None - tokenizer = None - if is_dual: - checker = _build_checker() - tokenizer = _load_tokenizer(GEN_MODEL_ID) - sys.stderr.write('[dualline] teacher checker ready (llm_backup teacher)\n') - - correct = 0 - total = 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') - - def _grade_and_log(rec, idx, raw_output, extra=None): - nonlocal correct, total - predicted = extract_boxed(raw_output) - is_correct = answers_match(predicted, rec['reference_answer']) - if is_correct: - correct += 1 - total += 1 - debug_rec = { - 'idx': idx, - 'reference_answer': rec['reference_answer'], - 'predicted': predicted, - 'is_correct': is_correct, - 'problem': rec['problem'], - 'model_output': raw_output, - } - if rec.get('level'): - debug_rec['level'] = rec['level'] - if rec.get('type'): - debug_rec['type'] = rec['type'] - if extra: - debug_rec.update(extra) - debug_records.append(debug_rec) - out_f.write(json.dumps(debug_rec, ensure_ascii=False) + '\n') - out_f.flush() - - if is_dual: - problems = [rec['problem'] for rec in records] - results = run_dualline_batch( - sampler, tokenizer, problems, checker, gen_params, - args.chunk_tokens, max_model_len, max_gen_tokens, - gen_dp=gen_dp, diagnose_workers=args.diagnose_workers) - for idx, (rec, result) in enumerate(zip(records, results)): - _grade_and_log(rec, idx, result['text'], extra={ - 'n_checks': result['n_checks'], - 'n_injections': result['n_injections'], - 'findings': result['findings'], - 'finished': result['finished'], - 'stopped_reason': result.get('stopped_reason'), - 'context_input_ids_len': result.get('context_input_ids_len'), - 'gen_tokens': result['gen_tokens'], - }) - stop_tag = (f' stop={result["stopped_reason"]}' - if result.get('stopped_reason') else '') - sys.stderr.write( - f' [idx {idx}] correct={debug_records[-1]["is_correct"]} ' - f'gen={result["gen_tokens"]} checks={result["n_checks"]} ' - f'inj={result["n_injections"]}{stop_tag}\n') - acc = correct / total if total else 0 - sys.stderr.write( - f'[dualline] batched eval done: acc={acc:.4f} ({correct}/{total})\n') - else: - import re - for batch_start in range(0, len(records), args.batch_size): - batch = records[batch_start:batch_start + args.batch_size] - prompts = [build_direct_prompt(r['problem']) for r in batch] - if gen_dp > 1 and len(prompts) < gen_dp: - prompts = _pad_batch_for_dp(prompts, gen_dp) - pad_n = len(prompts) - len(batch) - else: - pad_n = 0 - responses = sampler.sample(prompts, gen_params) - if pad_n: - responses = responses[:len(batch)] - 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 = re.sub(r'<\|[^|]+\|>', '', seq.decoded or '').rstrip() - _grade_and_log(rec, batch_start + i, raw_output) - acc = correct / total if total else 0 - sys.stderr.write(f' [{total}/{len(records)}] acc={acc:.4f} ' - f'({correct}/{total})\n') - - overall = correct / total if total else 0 - print(f'\n{"=" * 60}') - print(f'MATH dual-line — mode={args.mode}, model={GEN_MODEL_ID}') - print(f' n={total}, seed={args.seed}, chunk_tokens={args.chunk_tokens}, ' - f'max_model_len={max_model_len}') - print(f'{"=" * 60}') - print(f'Overall accuracy: {overall:.4f} ({correct}/{total})') - - if is_dual: - tot_checks = sum(r.get('n_checks', 0) for r in debug_records) - tot_inj = sum(r.get('n_injections', 0) for r in debug_records) - n_with_inj = sum(1 for r in debug_records if r.get('n_injections', 0) > 0) - print(f' checker: {tot_checks} checks, {tot_inj} injections across ' - f'{n_with_inj}/{total} problems') - n_ctx_full = sum( - 1 for r in debug_records if r.get('stopped_reason') == 'context_full') - n_sample_fail = sum( - 1 for r in debug_records if r.get('stopped_reason') == 'sample_failed') - n_unfinished = sum(1 for r in debug_records if not r.get('finished')) - print(f' length: context_full={n_ctx_full}/{total}, ' - f'sample_failed={n_sample_fail}/{total}, ' - f'unfinished(no EOS)={n_unfinished}/{total}') - - if any(r.get('level') for r in debug_records): - per = defaultdict(lambda: [0, 0]) - for r in debug_records: - lv = r.get('level', 'Unknown') - per[lv][1] += 1 - if r['is_correct']: - per[lv][0] += 1 - print('\nPer-level accuracy:') - for lv in sorted(per.keys()): - c, t = per[lv] - print(f' {lv:>10}: {c/t:.4f} ({c}/{t})') - - out_f.close() - print(f'\n[output] {len(debug_records)} records saved to {args.output}') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/legacy/eval_gpqa_rag.py b/cookbook/exp/legacy/eval_gpqa_rag.py deleted file mode 100644 index 7954e5fd2..000000000 --- a/cookbook/exp/legacy/eval_gpqa_rag.py +++ /dev/null @@ -1,1547 +0,0 @@ -"""Math evaluation: direct vs RAG-augmented with Qwen3.5-4B. - -Datasets (``--dataset``): - - ``math`` (default): MATH (Hendrycks), stratified by difficulty (Level 1-5) - so RAG gain can be plotted against difficulty. - - ``aops``: AoPS competition problems (metadata.boxed only). - -Modes (``--mode``): - - ``direct``: The model solves problems directly (4 GPUs, TP=4). - - ``rag`` (default): Retrieve top-k thinking traces from LanceDB, condense - them (API qwen3.7-max), inject as 1-shot examples, then solve - (8 GPUs: DP=4 embedding + TP=4 vLLM). - -Defaults implement **raw RAG on MATH**: ``--dataset math --mode rag --condense`` -with hint filtering OFF. The API condenser needs ``COMPRESS_API_KEY`` (or a -local condenser via ``EVAL_CONDENSER_GPUS``); otherwise pass ``--no-condense``. - -Optional ``--hint`` flag (rag mode only): - After retrieval + condensing, call an API model to filter and refine the - traces — keeping only applicable methods — then inject the refined trace. - -Reference answers are the ``\\boxed{...}`` content of each solution. - -Launch examples: - # Default: raw RAG on MATH, stratified 100/level (needs COMPRESS_API_KEY) - COMPRESS_API_KEY=sk-... python cookbook/exp/embedding/eval_gpqa_rag.py - - # Paired direct baseline on the same MATH subset - python cookbook/exp/embedding/eval_gpqa_rag.py --mode direct - - # Raw RAG without condenser (inject raw retrieved traces) - python cookbook/exp/embedding/eval_gpqa_rag.py --no-condense - - # Add hint filtering back on top of condensing - COMPRESS_API_KEY=sk-... python cookbook/exp/embedding/eval_gpqa_rag.py --hint - - # Fall back to the old AoPS dataset - python cookbook/exp/embedding/eval_gpqa_rag.py --dataset aops -""" -import argparse -import json -import os -import random -import re -import sys -import threading -import time -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Any, Dict, List, Optional - -import numpy as np -import torch - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.data_format import SamplingParams as TwinkleSamplingParams -from twinkle.loss import InfonceLoss -from twinkle.model import TransformersModel -from twinkle.processor import InputProcessor -from twinkle.sampler import vLLMSampler -from twinkle.template import Qwen3_5Template -from twinkle_agentic.protocol.openai import OpenAI as OpenAIClient - -logger = get_logger() - -# -- Condenser config ---------------------------------------------------------- -CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') -CONDENSE_API_KEY = os.environ.get('COMPRESS_API_KEY', '') -CONDENSE_BASE_URL = os.environ.get('COMPRESS_BASE_URL', 'https://dashscope.aliyuncs.com/compatible-mode/v1') -CONDENSE_API_MODEL = os.environ.get('COMPRESS_MODEL', 'qwen3.7-max') -CONDENSE_API_CONCURRENCY = int(os.environ.get('API_CONCURRENCY', 32)) -CONDENSE_API_MIN_INTERVAL = float(os.environ.get('API_MIN_INTERVAL', 0.1)) -CONDENSE_TEMPERATURE = 0.2 -CONDENSE_MAX_TOKENS = 8192 - -# -- Hint analysis config ------------------------------------------------------ -HINT_ANALYSIS_MAX_TOKENS = int(os.environ.get('HINT_ANALYSIS_MAX_TOKENS', 2000)) -HINT_ANALYSIS_TEMPERATURE = 0.2 - -HINT_ANALYSIS_SYSTEM = ( - 'You are a mathematical reasoning trace filter. ' - 'Given a target problem and reasoning traces retrieved from SIMILAR (but different) problems, ' - 'your task is to FILTER and REFINE the traces into a clean reference.\n\n' - 'Rules:\n' - '1. KEEP: solution steps, methods, formulas, techniques, and key insights ' - 'that are directly applicable to solving the target problem.\n' - '2. REMOVE: problem-specific numeric calculations that do not transfer, ' - 'dead-end explorations, irrelevant approaches, verbose restatements, ' - 'and any content that would mislead the solver on the target problem.\n' - '3. Output the refined trace directly as actionable solution steps. ' - 'Preserve the original mathematical expressions and step structure.\n' - '4. Do NOT solve the target problem. Do NOT add your own solutions or commentary.\n' - '5. Do NOT output the answer to either problem.\n' - '6. If the traces are entirely irrelevant, output exactly: "No applicable methods."' -) - -HINT_ANALYSIS_USER = ( - '## Target Problem\n{query}\n\n' - '## Retrieved Reasoning Traces\n{thinking}' -) - -# --------------------------------------------------------------------------- -# Config -# --------------------------------------------------------------------------- -# -- Gen/Embed config --------------------------------------------------------- -GEN_MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3.5-4B') -EMBED_MODEL_ID = os.environ.get( - 'EMBED_MODEL_ID', 'output.oldemb/embedding_full_transformers/last-checkpoint') - -GEN_GPUS = int(os.environ.get('GEN_GPUS', 8)) -EMB_GPUS = int(os.environ.get('EMB_GPUS', 2)) -EMBED_MAX_LENGTH = int(os.environ.get('EMBED_MAX_LENGTH', 20000)) - -GEN_GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.85)) -GEN_MAX_MODEL_LEN = int(os.environ.get('GEN_MAX_MODEL_LEN', 65536)) -GEN_MAX_TOKENS = int(os.environ.get('GEN_MAX_TOKENS', 65536)) -GEN_TEMPERATURE = float(os.environ.get('GEN_TEMPERATURE', 0.6)) -GEN_TOP_P = float(os.environ.get('GEN_TOP_P', 0.95)) - -AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') -MATH_DATA_DIR = os.environ.get('MATH_DATA_DIR', './output/math_data/MATH') - - -# --------------------------------------------------------------------------- -# Condenser prompts & validation -# --------------------------------------------------------------------------- - -COMPRESS_SYSTEM = """\ -You are a reasoning-trace condenser. Given a verbose reasoning trace, \ -extract the TRANSFERABLE KNOWLEDGE as an EXECUTABLE SOLUTION SKELETON \ -that would help a reader solve SIMILAR problems in the same domain. - -Your output is the ENTIRE useful content — there is no expansion tool, no second pass. \ -The reader will apply this knowledge to a DIFFERENT problem, so focus on what transfers. - -Principles: -1. OUTPUT AN EXECUTABLE STEP CHAIN: numbered steps that a solver can directly follow. \ -Each step should state WHAT to do and HOW (with the formula/technique), not just \ -name the concept. -2. INCLUDE FULL FORMULAS: theorems, identities, inequalities — state each \ -with its COMPLETE MATHEMATICAL EXPRESSION, not just the name. -3. STATE APPLICABILITY: what structural features of a problem signal that this \ -approach works (e.g. "when the constraint is a sum of squares"). -4. PRESERVE KEY INSIGHTS: the non-obvious ideas or tricks that make the approach \ -work — the things a solver would NOT think of without guidance. -5. REMOVE: problem-specific numeric calculations, dead-end explorations, \ -hesitations, verbose restatements, and trivial arithmetic. -6. FORMAT: Start with a one-line "Applicability" statement, then numbered steps, \ -then key formulas. Keep it concise and actionable. -7. NO meta-commentary about the compression process. NO preamble. -""" - -COMPRESS_USER = ( - '## Reader Problem (context only — do NOT solve it)\n{query}\n\n' - '## Reasoning Trace to Condense\n{text}') - - -def _is_truncated_compression(text: str) -> bool: - if not text or not text.strip(): - return True - lines = [l.strip() for l in text.strip().splitlines() if l.strip()] - if len(lines) < 3: - return True - last_line = lines[-1] - # Truncated if last line looks incomplete (no terminal punctuation/formula) - if last_line and last_line[-1] not in '.。!！)）]】}\\$': - # Allow lines ending with numbers, boxed answers, etc. - if not re.search(r'\d$|\\boxed|\$|\)$', last_line): - return True - return False - - -# -- API rate limiter ---------------------------------------------------------- -_api_semaphore = threading.Semaphore(CONDENSE_API_CONCURRENCY) -_api_bucket_lock = threading.Lock() -_api_tokens = [float(CONDENSE_API_CONCURRENCY)] -_api_last_refill = [time.monotonic()] - - -def _api_throttle(): - _api_semaphore.acquire() - wait = 0.0 - try: - with _api_bucket_lock: - now = time.monotonic() - elapsed = now - _api_last_refill[0] - refill = elapsed / CONDENSE_API_MIN_INTERVAL - _api_tokens[0] = min(float(CONDENSE_API_CONCURRENCY), _api_tokens[0] + refill) - _api_last_refill[0] = now - if _api_tokens[0] >= 1.0: - _api_tokens[0] -= 1.0 - else: - wait = (1.0 - _api_tokens[0]) * CONDENSE_API_MIN_INTERVAL - _api_tokens[0] = 0.0 - finally: - _api_semaphore.release() - if wait > 0: - time.sleep(wait) - - -def _api_condense_single(api_client: OpenAIClient, messages: List[Dict]) -> Optional[str]: - _api_throttle() - trajectory = {'messages': messages} - sp = TwinkleSamplingParams(temperature=CONDENSE_TEMPERATURE, max_tokens=CONDENSE_MAX_TOKENS) - try: - reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) - except Exception as exc: - logger.warning(f'[condense-api] error: {exc}') - return None - content = (reply.get('content') or '').strip() - if not content: - return None - m = re.match(r'^```[a-zA-Z]*\n(.*?)\n```\s*$', content, re.DOTALL) - if m: - content = m.group(1).strip() - return content - - -def _api_hint_analysis_batch( - api_client: OpenAIClient, - problems: List[str], - condensed_examples: List[List[Dict[str, str]]], -) -> List[Optional[str]]: - """Call API to pre-analyze RAG relevance for each problem.""" - _MAX_HINT_INPUT = 8000 - results: List[Optional[str]] = [None] * len(problems) - tasks = [] - for i, prob in enumerate(problems): - if not condensed_examples[i]: - continue - traces = [ex.get('thinking', '') for ex in condensed_examples[i]] - merged_thinking = '\n---\n'.join(traces) - if len(merged_thinking) > _MAX_HINT_INPUT: - merged_thinking = merged_thinking[:_MAX_HINT_INPUT] + '\n[...truncated]' - user_msg = HINT_ANALYSIS_USER.format(query=prob, thinking=merged_thinking) - msgs = [ - {'role': 'system', 'content': HINT_ANALYSIS_SYSTEM}, - {'role': 'user', 'content': user_msg}, - ] - tasks.append((i, msgs)) - - if not tasks: - return results - - def _call_one(idx, msgs): - _api_throttle() - try: - trajectory = {'messages': msgs} - sp = TwinkleSamplingParams( - temperature=HINT_ANALYSIS_TEMPERATURE, - max_tokens=HINT_ANALYSIS_MAX_TOKENS) - reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) - content = (reply.get('content') or '').strip() - # Treat "No applicable methods." as empty (will trigger fallback) - if not content or content == 'No applicable methods.': - return idx, None - return idx, content - except Exception as exc: - logger.warning(f'[hint-analysis] error for idx={idx}: {exc}') - return idx, None - - with ThreadPoolExecutor(max_workers=min(len(tasks), CONDENSE_API_CONCURRENCY)) as pool: - futs = [pool.submit(_call_one, idx, msgs) for idx, msgs in tasks] - for fut in as_completed(futs): - idx, analysis = fut.result() - results[idx] = analysis - - n_success = sum(1 for r in results if r) - logger.info(f'[hint-analysis] completed {n_success}/{len(tasks)} analyses') - return results - - -# --------------------------------------------------------------------------- -# LLM-based decontamination -# --------------------------------------------------------------------------- - -_DECONTAM_JUDGE_PROMPT = ( - 'We are building a RAG-augmented math training system. Problem A is the test ' - 'question; Problem B was retrieved from a knowledge base.\n' - 'Answer YES only if A and B are essentially the SAME specific problem — ' - 'i.e. solving B directly gives you A\'s answer (just different wording/notation/' - 'format/negation).\n' - 'Answer NO if they merely share the same method/topic but have different ' - 'specific values, equations, or geometric configurations — learning B\'s ' - 'approach still requires independent work to solve A.\n' - 'Problem A: {prob_a}\n' - 'Problem B: {prob_b}\n' - 'Answer only YES or NO.' -) - - -def _llm_judge_same_problem( - api_client: OpenAIClient, pairs: List[tuple], -) -> List[bool]: - """Batch LLM judge: are (problem_a, problem_b) the same problem? - - Returns list of bools (True = same problem = should filter). - """ - if not pairs or not api_client: - return [False] * len(pairs) - - results = [False] * len(pairs) - - def _judge_one(idx, pa, pb): - prompt = _DECONTAM_JUDGE_PROMPT.format(prob_a=pa, prob_b=pb) - msgs = [{'role': 'user', 'content': prompt}] - _api_throttle() - try: - trajectory = {'messages': msgs} - sp = TwinkleSamplingParams(temperature=0.1, max_tokens=8) - reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) - answer = (reply.get('content') or '').strip().upper() - return idx, 'YES' in answer - except Exception: - return idx, False - - with ThreadPoolExecutor(max_workers=min(len(pairs), CONDENSE_API_CONCURRENCY)) as pool: - futs = [pool.submit(_judge_one, i, pa, pb) for i, (pa, pb) in enumerate(pairs)] - for fut in as_completed(futs): - idx, is_same = fut.result() - results[idx] = is_same - return results - - -def _llm_decontaminate( - api_client: OpenAIClient, - problems: List[str], - all_examples: List[List[Dict[str, str]]], -) -> List[List[Dict[str, str]]]: - """Apply LLM-based decontamination: remove retrievals judged as same problem.""" - judge_pairs = [] # (qi, ret_idx, prob_a, prob_b) - for qi, exs in enumerate(all_examples): - for ri, ex in enumerate(exs): - judge_pairs.append((qi, ri, problems[qi], ex.get('query', ''))) - - if not judge_pairs: - return all_examples - - pairs_input = [(pa, pb) for _, _, pa, pb in judge_pairs] - verdicts = _llm_judge_same_problem(api_client, pairs_input) - to_remove = set() - for vi, (qi, ri, _, _) in enumerate(judge_pairs): - if verdicts[vi]: - to_remove.add((qi, ri)) - - if to_remove: - logger.info(f'[decontam-llm] filtered {len(to_remove)} same-problem retrievals') - for qi in range(len(all_examples)): - all_examples[qi] = [ - ex for ri, ex in enumerate(all_examples[qi]) - if (qi, ri) not in to_remove - ] - return all_examples - - -def condense_traces( - examples_batch: List[List[Dict[str, str]]], - problems: List[str], - api_client: OpenAIClient, - condenser_sampler=None, - compress_params=None, - special_tokens: set = None, - max_output_len: int = 2000, - dp_size: int = 1, -) -> List[List[Dict[str, str]]]: - """Compress retrieved thinking traces with query-aware condenser. - - Primary: local vLLM condenser (if provided). - Fallback: API condenser. - Final fallback: raw trace truncated to max_output_len. - """ - result: List[List[Dict[str, str]]] = [] - # Flatten all (batch_idx, ex_idx, problem, example) for batch processing - tasks = [] - for bi, (exs, prob) in enumerate(zip(examples_batch, problems)): - for ei, ex in enumerate(exs): - tasks.append((bi, ei, prob, ex)) - - if not tasks: - return [[] for _ in examples_batch] - - # Build condense prompts (aligned with make_embedding_dataset.py hard path) - prompts = [] - for _, _, prob, ex in tasks: - user_msg = COMPRESS_USER.format(query=prob, text=ex['thinking']) - prompts.append([{'role': 'system', 'content': COMPRESS_SYSTEM}, - {'role': 'user', 'content': user_msg}]) - - # Phase 1: local vLLM condenser - condensed = [None] * len(tasks) - condense_sources = ['raw'] * len(tasks) - fallback_indices = [] - - if condenser_sampler is not None and compress_params is not None: - sampler_inputs = [{'messages': p} for p in prompts] - # The local vLLM sampler runs data-parallel across ``dp_size`` workers - # and requires at least one item per worker (it errors with - # "Batch too small for N workers" otherwise). Pad the batch up to a - # multiple of dp_size by repeating the last item, run, then keep only - # the first ``n_real`` responses and drop the padding. - n_real = len(sampler_inputs) - pad_size = 0 - if dp_size > 1 and n_real > 0 and n_real % dp_size != 0: - pad_size = dp_size - (n_real % dp_size) - sampler_inputs = sampler_inputs + [sampler_inputs[-1]] * pad_size - try: - responses = condenser_sampler.sample(sampler_inputs, compress_params) - except Exception as exc: - logger.warning(f'[condense] sampler error: {exc}') - responses = [None] * len(sampler_inputs) - if pad_size: - responses = responses[:n_real] - for ri, resp in enumerate(responses): - seq = resp.sequences[0] if resp and resp.sequences else None - text = '' - if seq and seq.stop_reason != 'length' and seq.decoded: - text = seq.decoded - if special_tokens: - for tok in special_tokens: - text = text.replace(tok, '') - text = text.rstrip() - if text and not _is_truncated_compression(text): - condensed[ri] = text - condense_sources[ri] = 'local' - else: - fallback_indices.append(ri) - else: - fallback_indices = list(range(len(tasks))) - - # Phase 2: API fallback - if fallback_indices and api_client: - with ThreadPoolExecutor(max_workers=CONDENSE_API_CONCURRENCY) as pool: - futures = {} - for ri in fallback_indices: - futures[pool.submit(_api_condense_single, api_client, prompts[ri])] = ri - for fut in as_completed(futures): - ri = futures[fut] - api_result = fut.result() - if api_result and not _is_truncated_compression(api_result): - condensed[ri] = api_result - condense_sources[ri] = 'api' - - # Phase 3: assemble results (fallback to raw truncation) - result = [[] for _ in examples_batch] - for ti, (bi, ei, prob, ex) in enumerate(tasks): - compressed = condensed[ti] - raw_len = len(ex['thinking']) - sim_val = ex.get('_sim', 0.0) - if compressed: - result[bi].append({'query': ex['query'], - 'thinking': _strip_condenser_markers(compressed), - '_condense_source': condense_sources[ti], - '_raw_trace_len': raw_len, '_sim': sim_val}) - else: - result[bi].append({'query': ex['query'], - 'thinking': ex['thinking'][:max_output_len], - '_condense_source': 'raw', - '_raw_trace_len': raw_len, '_sim': sim_val}) - - n_ok = sum(1 for c in condensed if c) - logger.info(f'[condense] {n_ok}/{len(tasks)} compressed ok, ' - f'{len(tasks) - n_ok} fell back to raw truncation') - return result - - -def _strip_condenser_markers(text: str) -> str: - """Light cleanup of condenser output. - - Removes any residual markdown headers or meta-lines that don't carry - solution content. Keeps numbered steps and equations intact. - """ - # Remove legacy ## headers if condenser still emits them - if '## More' in text: - text = text.split('## More', 1)[0] - text = re.sub(r'^##\s*Summary\s*\n?', '', text, flags=re.MULTILINE) - text = re.sub(r'^Topic:\s*.*\n?', '', text, flags=re.MULTILINE) - # Remove meta-commentary lines - text = re.sub(r'^\s*\(Note:.*\)\s*$', '', text, flags=re.MULTILINE) - return text.strip() - - -# --------------------------------------------------------------------------- -# Boxed answer extraction -# --------------------------------------------------------------------------- -_BOXED_RE = re.compile(r'\\boxed\s*\{') - - -def extract_boxed(text: str) -> Optional[str]: - """Extract the last \\boxed{...} content, handling nested braces.""" - if not text: - return None - last_match = None - for m in _BOXED_RE.finditer(text): - start = m.end() - depth = 1 - i = start - while i < len(text) and depth > 0: - if text[i] == '{': - depth += 1 - elif text[i] == '}': - depth -= 1 - i += 1 - if depth == 0: - last_match = text[start:i - 1].strip() - return last_match - - -def normalize_answer(ans: str) -> str: - """Normalize a math answer string for comparison.""" - if not ans: - return '' - s = str(ans).strip() - m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) - if m: - return m.group(1) - s = s.strip('$').strip() - s = s.replace('−', '-') - 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.replace(r'\,', '').replace(r'\;', '').replace(r'\!', '') - s = re.sub(r'\\(?:quad|qquad|\s)', '', s) - s = re.sub(r'\s+', '', s) - s = re.sub(r'\\sqrt([0-9A-Za-z])', r'\\sqrt{\1}', s) - s = re.sub(r'\\frac\{([^{}]+)\}([^{}\\])', r'\\frac{\1}{\2}', s) - s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) - s = re.sub(r'\{[a-zA-Z]+\}$', '', s) - s = re.sub(r'\^\\circ|\^\{\\circ\}|°|\\circ', 'deg', 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'(? 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 - frac_re = re.compile(r'^\(([^)]+)\)/\(([^)]+)\)$') - - def _eval_frac(s): - m = frac_re.match(s) - if m: - try: - return float(m.group(1)) / float(m.group(2)) - except (ValueError, ZeroDivisionError): - pass - return None - - va, vb = _eval_frac(a), _eval_frac(b) - if va is not None and vb is not None: - return abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) - return False - - -_MCQ_REF_RE = re.compile( - r'^\\(?:textbf|text|mathrm|mathbf)\{\(?([A-E])\)?\s*\}\s*(.+)$' - r'|^\(?([A-E])\)\s+(.+)$' -) -_VAR_PREFIX_RE = re.compile(r'^(?:[A-Za-z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)$', re.DOTALL) - - -def _split_mcq(ans: str): - """Split an MCQ answer into (letter, value) components.""" - 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) - 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 _strip_var_prefix(ans: str) -> str: - m = _VAR_PREFIX_RE.match((ans or '').strip()) - return m.group(1).strip() if m else (ans or '') - - -def _normalize_tuple(ans: str) -> str: - return re.sub(r'[\s()\[\]{}\\]', '', ans or '') - - -def _math_verify_equal(predicted: str, reference: str) -> bool: - try: - from math_verify import parse, verify - gold = parse(r'\boxed{' + reference + '}', parsing_timeout=1) - pred = parse(r'\boxed{' + predicted + '}', parsing_timeout=1) - return bool(gold and pred and verify(gold, pred, timeout_seconds=1)) - except Exception: - 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 or norm_p.lower() == norm_r.lower(): - return True - if _try_numeric_equal(norm_p, norm_r): - return True - - stripped_p = normalize_answer(_strip_var_prefix(predicted)) - stripped_r = normalize_answer(_strip_var_prefix(reference)) - if stripped_p and stripped_r: - if stripped_p == stripped_r or stripped_p.lower() == stripped_r.lower(): - return True - if _try_numeric_equal(stripped_p, stripped_r): - return True - - p_letter, p_value = _split_mcq(predicted) - r_letter, r_value = _split_mcq(reference) - if p_letter and r_letter and p_letter == r_letter: - return True - if (p_letter and p_value is None) and (r_letter and r_value is None): - return False - - 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 p_val.lower() == r_val.lower() or _try_numeric_equal(p_val, r_val): - return True - - for left, right in ((norm_p, norm_r), (stripped_p, stripped_r)): - tuple_l, tuple_r = _normalize_tuple(left), _normalize_tuple(right) - if ',' in tuple_l and tuple_l == tuple_r: - return True - - if '=' in norm_r and '=' not in norm_p: - if any(part and (part == norm_p or _try_numeric_equal(part, norm_p)) for part in norm_r.split('=')): - return True - if '=' in norm_p and '=' not in norm_r: - if any(part and (part == norm_r or _try_numeric_equal(part, norm_r)) for part in norm_p.split('=')): - return True - - return _math_verify_equal(stripped_p or norm_p, stripped_r or norm_r) - - -# --------------------------------------------------------------------------- -# Dataset loading -# --------------------------------------------------------------------------- - -def _load_aops_from_modelscope(): - """Download the AoPS repo natively from ModelScope and read its parquet. - - Primary loader: ``dataset_snapshot_download`` pulls the dataset repo files - (parquet) straight from the ModelScope hub WITHOUT going through the - ``datasets``/HF-filesystem path used by ``MsDataset.load`` — that path is - broken on this modelscope build (``HfFileSystem.find() got multiple values - for 'maxdepth'``). We then read the local parquet with the ``datasets`` - backend (reading local files does not trigger the HfFileSystem bug). - """ - import glob - - from datasets import Dataset as HFDataset - from modelscope.hub.snapshot_download import dataset_snapshot_download - - local = dataset_snapshot_download(AOPS_DATASET_ID) - files = sorted(glob.glob(os.path.join(local, '**', '*.parquet'), - recursive=True)) - if not files: - # Older snapshots may materialize an arrow file instead of parquet. - files = sorted(glob.glob(os.path.join(local, '**', '*train*.arrow'), - recursive=True)) - if files: - sys.stderr.write(f'[aops] modelscope snapshot arrow: {files[0]}\n') - return HFDataset.from_file(files[0]) - return None - sys.stderr.write(f'[aops] modelscope snapshot parquet: {files[0]}\n') - return HFDataset.from_parquet(files if len(files) > 1 else files[0]) - - -def load_aops(n: int, seed: int = 42) -> List[Dict[str, Any]]: - """Load AoPS boxed problems, sample n, extract reference answers. - - Uses ModelScope as the data source (native repo snapshot download). - """ - ds = None - try: - ds = _load_aops_from_modelscope() - except Exception as exc: - sys.stderr.write(f'[aops] modelscope snapshot download failed ({exc}); ' - f'trying MsDataset.load\n') - if ds is None: - from modelscope import MsDataset - ds = MsDataset.load(AOPS_DATASET_ID, split='train', - download_mode='reuse_dataset_if_exists') - boxed = [] - for row in ds: - if not row['metadata'].get('boxed'): - continue - ref = extract_boxed(row['solution']) - if not ref: - continue - boxed.append({ - 'problem': row['problem'], - 'solution': row['solution'], - 'reference_answer': ref, - 'tags': row.get('tags', []), - }) - sys.stderr.write(f'[aops] {len(boxed)} boxed problems with extractable answers\n') - rng = random.Random(seed) - rng.shuffle(boxed) - if n > 0 and n < len(boxed): - boxed = boxed[:n] - sys.stderr.write(f'[aops] sampled {n} problems\n') - return boxed - - -def load_math(n: int, seed: int = 42, split: str = 'test', - per_level: int = 0) -> List[Dict[str, Any]]: - """Load the MATH (Hendrycks) dataset from local extracted JSON files. - - Each problem's reference answer is the ``\\boxed{}`` content of its - ``solution`` (MATH solutions always end in a boxed answer). - - Sampling is *stratified by level* so every difficulty (Level 1-5) is - represented equally — required to measure how RAG gain varies with - difficulty. ``per_level`` (if >0) fixes the count per level; otherwise - ``n`` is split evenly across the 5 levels. When both are 0, all problems - are returned. The final list is shuffled with ``seed`` so index order is - stable/comparable across direct vs rag runs. - """ - import glob - root = os.path.join(MATH_DATA_DIR, split) - files = glob.glob(os.path.join(root, '*', '*.json')) - if not files: - raise FileNotFoundError( - f'[math] no problems found under {root!r}; set MATH_DATA_DIR or ' - f'extract MATH.zip there') - - by_level: Dict[str, List[Dict[str, Any]]] = {} - n_no_box = 0 - for fp in files: - try: - with open(fp, 'r', encoding='utf-8') as fin: - row = json.load(fin) - except Exception: - continue - ref = extract_boxed(row.get('solution', '')) - if not ref: - n_no_box += 1 - continue - level = row.get('level', 'Unknown') - by_level.setdefault(level, []).append({ - 'problem': row['problem'], - 'solution': row['solution'], - 'reference_answer': ref, - 'level': level, - 'type': row.get('type', ''), - }) - - total = sum(len(v) for v in by_level.values()) - sys.stderr.write( - f'[math] {total} problems with boxed answers across ' - f'{len(by_level)} levels (skipped {n_no_box} without boxed)\n') - - levels = sorted(by_level.keys()) - rng = random.Random(seed) - - # Decide how many per level. - if per_level <= 0 and n > 0: - per_level = max(1, n // max(1, len(levels))) - - sampled: List[Dict[str, Any]] = [] - for lv in levels: - pool = by_level[lv] - rng.shuffle(pool) - take = pool if per_level <= 0 else pool[:per_level] - sampled.extend(take) - sys.stderr.write(f'[math] {lv}: took {len(take)}/{len(pool)}\n') - - rng.shuffle(sampled) - sys.stderr.write(f'[math] total sampled: {len(sampled)}\n') - return sampled - - -# --------------------------------------------------------------------------- -# Prompt building -# --------------------------------------------------------------------------- - -DIRECT_SYSTEM = ( - 'You are an expert competition mathematician. Solve the following problem ' - 'step by step. Provide your final answer inside \\boxed{}.' -) - -RAG_SYSTEM = ( - 'You are an expert competition mathematician. Solve the following problem ' - 'step by step. Provide your final answer inside \\boxed{}.\n\n' - 'You will first see example problem-solving traces or skills. ' - 'Learn from the reasoning methodology demonstrated in these examples, ' - 'then thinking to solve the actual problem.' -) - -RAG_FOLLOWUP = ( - 'The above is a reference solution to a similar problem. ' - 'You may use any applicable techniques from it, or ignore it ' - 'if you find a better approach. ' - 'Solve the problem step by step and put your final answer in \\boxed{}.' -) - -HINT_FOLLOWUP = ( - 'The above are applicable solution approaches extracted from similar problems. ' - 'You may use any applicable techniques from them, or ignore them ' - 'if you find a better approach. ' - 'Solve the problem step by step and put your final answer in \\boxed{}.' -) - -# Reminder appended to the final user turn. Without this, the reasoning model -# can loop indefinitely on multiple-choice problems, oscillating between boxing -# the option letter and boxing the value (e.g. "I'll box B. I'll box 21. ...") -# and never terminating. Boxing BOTH the letter and value removes the ambiguity -# (the grader accepts either), so the model has no format decision to agonize over. -MCQ_INSTRUCTION = ( - '\n\nNote: If the problem is multiple-choice (it lists options such as ' - '(A), (B), (C), ...), put BOTH the option letter and its value in the box, ' - 'e.g. \\boxed{(B) 21}. Otherwise, box the value directly. Decide the answer ' - 'format once and do not deliberate over which form to box.' -) - - -def build_direct_prompt(problem: str) -> Dict[str, Any]: - return { - 'messages': [ - {'role': 'system', 'content': DIRECT_SYSTEM}, - {'role': 'user', 'content': problem}, - ] - } - - -def build_hint_prompt(problem: str, hint_analysis: str) -> Dict[str, Any]: - """Build prompt with pre-analyzed hint in a multi-turn conversation. - - Mirrors ``build_rag_prompt``: the hint is presented as an assistant - "extracted approaches" turn (instead of being buried in the system - prompt), followed by a user instruction that provides a clear closing - directive to solve the problem and box the answer. Keeping the final - solve/box instruction in a dedicated user turn (rather than in the - system prompt) helps the reasoning model terminate cleanly. - """ - messages: List[Dict[str, str]] = [ - {'role': 'system', 'content': DIRECT_SYSTEM}, - {'role': 'user', 'content': problem}, - {'role': 'assistant', - 'content': ('Here are applicable solution approaches extracted from ' - f'similar problems:\n\n{hint_analysis}')}, - {'role': 'user', 'content': HINT_FOLLOWUP + MCQ_INSTRUCTION}, - ] - return {'messages': messages} - - -def build_rag_prompt(problem: str, - examples: List[Dict[str, str]]) -> Dict[str, Any]: - """Approach B: multi-turn assistant format. - - The trace is presented as an assistant "retrieval" turn, followed by - a user instruction that constrains the model to use methodology only. - """ - messages: List[Dict[str, str]] = [{'role': 'system', 'content': DIRECT_SYSTEM}] - messages.append({'role': 'user', 'content': problem}) - # Build trace content from retrieved examples - trace_parts = [] - for i, ex in enumerate(examples, 1): - trace_parts.append(f'[Retrieved Example {i}]\nProblem: {ex["query"]}\n' - f'Reasoning:\n{ex["thinking"]}') - trace_text = '\n\n'.join(trace_parts) - messages.append({'role': 'assistant', - 'content': f'I found relevant reasoning traces from the knowledge base!\n\n{trace_text}'}) - messages.append({'role': 'user', 'content': RAG_FOLLOWUP + MCQ_INSTRUCTION}) - return {'messages': messages} - - -# --------------------------------------------------------------------------- -# 13-gram Jaccard decontamination -# --------------------------------------------------------------------------- - -def _normalize_for_ngram(text: str) -> str: - """Normalize text for n-gram comparison: strip LaTeX markup, lowercase.""" - text = text.lower() - text = re.sub(r'\$+', '', text) - text = re.sub(r'\\[a-z]+\{([^}]*)\}', r'\1', text) - text = re.sub(r'\\[a-z]+', ' ', text) - text = re.sub(r'[{}\\^_$]', '', text) - text = re.sub(r'\s+', ' ', text).strip() - return text - - -def _ngram_jaccard(text_a: str, text_b: str, n: int = 13) -> float: - """13-gram character-level Jaccard similarity.""" - a = _normalize_for_ngram(text_a) - b = _normalize_for_ngram(text_b) - if len(a) < n or len(b) < n: - return 0.0 - grams_a = set(a[i:i + n] for i in range(len(a) - n + 1)) - grams_b = set(b[i:i + n] for i in range(len(b) - n + 1)) - if not grams_a or not grams_b: - return 0.0 - return len(grams_a & grams_b) / len(grams_a | grams_b) - - -# --------------------------------------------------------------------------- -# Embedding / RAG helpers -# --------------------------------------------------------------------------- - -def _wrap_anchor(text: str) -> List[Dict[str, str]]: - return [ - {'role': 'user', 'content': text}, - {'role': 'assistant', 'content': 'Match the correct response here.'}, - ] - - -def get_embeddings(model: TransformersModel, template: Qwen3_5Template, - texts: List[str], dp_size: int) -> np.ndarray: - if not texts: - return np.zeros((0,), dtype=np.float32) - n = len(texts) - pad_n = (-n) % dp_size - padded = list(texts) + [' '] * pad_n if pad_n else list(texts) - features = [] - for t in padded: - feat = template.encode({'messages': _wrap_anchor(t or ' ')}) - feat['labels'] = [1] - features.append(feat) - out = model.forward_only(inputs=features, task='embedding', return_logits=True) - emb = out['embeddings'] - if isinstance(emb, torch.Tensor): - emb = emb.detach().to(torch.float32).cpu().numpy() - emb = np.asarray(emb, dtype=np.float32) - return emb[:n] if pad_n else emb - - -def retrieve_examples(tbl, query_vecs: np.ndarray, top_k: int, - use_thinking_raw: bool, sim_threshold: float = 0.0, - problems: List[str] = None, - decontam_threshold: float = 0.0, - ) -> List[List[Dict[str, str]]]: - thinking_field = 'thinking_raw' if use_thinking_raw else 'cot_compressed' - fetch_limit = top_k + 50 if decontam_threshold > 0 else top_k - n_queries = len(query_vecs) - all_examples: List[List[Dict[str, str]]] = [None] * n_queries - decontam_skipped = 0 - _decontam_lock = threading.Lock() - - def _search_one(qi: int): - nonlocal decontam_skipped - vec = query_vecs[qi] - results = ( - tbl.search(vec.astype(np.float32).tolist()) - .metric('dot') - .limit(fetch_limit) - .select(['query_raw', thinking_field, '_distance']) - .to_list() - ) - problem_text = problems[qi] if problems else '' - examples = [] - local_skipped = 0 - for r in results: - if len(examples) >= top_k: - break - sim = 1.0 - r.get('_distance', 0.0) - if sim < sim_threshold: - continue - q = r.get('query_raw', '') - t = r.get(thinking_field, '') - if not t: - continue - if decontam_threshold > 0 and problem_text and q: - ng_sim = _ngram_jaccard(problem_text, q) - if ng_sim > decontam_threshold: - local_skipped += 1 - continue - examples.append({'query': q, 'thinking': t, '_sim': round(sim, 4), - '_raw_trace_len': len(t)}) - all_examples[qi] = examples - if local_skipped: - with _decontam_lock: - decontam_skipped += local_skipped - - with ThreadPoolExecutor(max_workers=min(n_queries, 16)) as pool: - list(pool.map(_search_one, range(n_queries))) - - if decontam_skipped > 0: - logger.info(f'[decontam] skipped {decontam_skipped} leaked retrievals ' - f'(13-gram Jaccard > {decontam_threshold})') - return all_examples - - -# --------------------------------------------------------------------------- -# Main -# --------------------------------------------------------------------------- - -def main(): - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--mode', choices=['direct', 'rag'], default='rag') - p.add_argument('--dataset', choices=['aops', 'math'], default='math', - help='Evaluation dataset. "math" = MATH (Hendrycks), ' - 'stratified by level for a difficulty-vs-gain curve.') - p.add_argument('--math-split', default='test', - help='MATH split to load (test/train).') - p.add_argument('--per-level', type=int, default=100, - help='MATH only: problems per difficulty level (default 100 ' - '-> 500 total across Level 1-5). If 0, --n is split ' - 'evenly across the 5 levels.') - p.add_argument('--n', type=int, default=0, - help='Pool size: sample this many problems (0 = all boxed). ' - 'In RAG mode with --target-eval, set this to 0 for max coverage.') - p.add_argument('--target-eval', type=int, default=0, - help='Stop after this many problems are successfully evaluated ' - '(0 = no limit, evaluate the entire sampled set — the ' - 'default, so all 500 stratified MATH problems are run). ' - 'RAG mode: counts problems with valid traces after ' - 'decontam; direct mode: ignored, evaluates all filtered.') - p.add_argument('--db-path', default='./output.oldemb/thinking_rag/lance.db') - p.add_argument('--table', default='thinking_traces') - p.add_argument('--top-k', type=int, default=1) - p.add_argument('--use-cot-compressed', action='store_true', - help='Use pre-compressed cot_compressed field instead of thinking_raw.') - p.add_argument('--sim-threshold', type=float, default=0.75, - help='Minimum cosine similarity for retrieved traces. ' - 'Traces below this are discarded at retrieval time.') - p.add_argument('--decontam-threshold', type=float, default=0.20, - help='13-gram Jaccard threshold for leak detection. ' - 'Retrieved traces above this are skipped (0=disabled).') - p.add_argument('--llm-decontam', action='store_true', default=True, - help='LLM-based decontamination (default ON): API judges whether ' - 'retrieved problem is the same as the test problem. ' - 'Applied after 13-gram decontam, before condensing. ' - 'Use --no-llm-decontam to disable.') - p.add_argument('--no-llm-decontam', dest='llm_decontam', action='store_false', - help='Disable LLM-based decontamination.') - p.add_argument('--max-trace-len', type=int, default=12000) - p.add_argument('--condense', action='store_true', default=True, - help='Enable condenser re-compression on retrieved traces ' - '(default ON). Use --no-condense to inject raw traces.') - p.add_argument('--no-condense', dest='condense', action='store_false', - help='Disable condenser; inject raw retrieved traces.') - p.add_argument('--condense-max-len', type=int, default=2000, - help='Max chars of condensed trace (fallback truncation).') - p.add_argument('--batch-size', type=int, default=16) - p.add_argument('--seed', type=int, default=42) - p.add_argument('--hint', action='store_true', default=False, - help='Enable API hint filtering on retrieved traces (default OFF; ' - 'raw RAG injects the condensed trace directly). ' - 'In rag mode: retrieve → condense → API filters trace → refined system prompt. ' - 'In direct mode: ignored (no traces to filter).') - p.add_argument('--no-hint', dest='hint', action='store_false', - help='Disable API hint filtering; inject condensed trace directly.') - p.add_argument('--problem-ids-file', default=None, - help='File listing problem indices evaluated by RAG mode. ' - 'RAG mode writes this file; direct mode reads it to ' - 'evaluate the same subset (use --no-filter to disable). ' - 'Defaults to a dataset-specific path.') - p.add_argument('--no-filter', action='store_true', - help='In direct mode, evaluate ALL sampled problems ' - 'instead of filtering to RAG subset.') - p.add_argument('--output', default=None) - args = p.parse_args() - - # Dataset-specific default paths (keeps aops and math runs from colliding). - if args.problem_ids_file is None: - args.problem_ids_file = ( - f'./output/thinking_rag/{args.dataset}_rag_problem_ids.json') - - if args.output is None: - suffix = f'{args.mode}_hint' if (args.hint and args.mode == 'rag') else args.mode - args.output = ( - f'./output/thinking_rag/{args.dataset}_{suffix}_results.jsonl') - - if args.condense and args.use_cot_compressed: - logger.warning('--condense requires thinking_raw, ignoring --use-cot-compressed') - args.use_cot_compressed = False - - if args.dataset == 'math': - records = load_math(n=args.n, seed=args.seed, split=args.math_split, - per_level=args.per_level) - else: - records = load_aops(n=args.n, seed=args.seed) - - is_rag = (args.mode == 'rag') - - # Direct mode: filter to same problems RAG evaluated (controlled comparison) - original_indices = list(range(len(records))) # track original indices - if not is_rag and not args.no_filter: - if os.path.exists(args.problem_ids_file): - with open(args.problem_ids_file) as f: - content = f.read().strip() - if content.startswith('['): - valid_indices = set(json.loads(content)) - else: - valid_indices = set(int(line) for line in content.splitlines() if line.strip()) - filtered = [(i, r) for i, r in enumerate(records) if i in valid_indices] - original_indices = [i for i, _ in filtered] - records = [r for _, r in filtered] - sys.stderr.write( - f'[direct] filtered to {len(records)} problems ' - f'from {args.problem_ids_file}\n') - else: - sys.stderr.write( - f'[direct] WARNING: {args.problem_ids_file} not found, ' - f'running all {len(records)} problems\n') - - condenser_gpus = int(os.environ.get('EVAL_CONDENSER_GPUS', 0)) if args.condense else 0 - - # Raw RAG relies on the API condenser (qwen3.7-max). Fail fast with a clear - # message if it's enabled without an API key and without a local condenser. - if is_rag and args.condense and not CONDENSE_API_KEY and condenser_gpus == 0: - sys.stderr.write( - '[condense] ERROR: --condense is ON but COMPRESS_API_KEY is unset ' - 'and no local condenser (EVAL_CONDENSER_GPUS=0).\n' - ' Fix one of:\n' - ' - export COMPRESS_API_KEY=sk-... (use API condenser)\n' - ' - EVAL_CONDENSER_GPUS=2 python ... (use local vLLM condenser)\n' - ' - pass --no-condense (inject raw traces)\n') - sys.exit(1) - - if is_rag: - num_gpus = EMB_GPUS + GEN_GPUS + condenser_gpus - device_groups = [ - DeviceGroup(name='emb_model', ranks=list(range(EMB_GPUS)), - device_type='GPU'), - DeviceGroup(name='sampler', - ranks=list(range(EMB_GPUS, EMB_GPUS + GEN_GPUS)), - device_type='GPU', gpus_per_worker=GEN_GPUS), - ] - if condenser_gpus > 0: - cond_start = EMB_GPUS + GEN_GPUS - device_groups.append( - DeviceGroup(name='condenser', - ranks=list(range(cond_start, cond_start + condenser_gpus)), - device_type='GPU')) - emb_mesh = DeviceMesh.from_sizes(world_size=EMB_GPUS, dp_size=EMB_GPUS) - gen_mesh = DeviceMesh.from_sizes(world_size=GEN_GPUS, tp_size=GEN_GPUS) - twinkle.initialize(mode='ray', nproc_per_node=num_gpus, - groups=device_groups, lazy_collect=False) - else: - device_groups = [ - DeviceGroup(name='sampler', ranks=list(range(GEN_GPUS)), - device_type='GPU', gpus_per_worker=GEN_GPUS), - ] - gen_mesh = DeviceMesh.from_sizes(world_size=GEN_GPUS, tp_size=GEN_GPUS) - twinkle.initialize(mode='ray', nproc_per_node=GEN_GPUS, - groups=device_groups, lazy_collect=False) - - sampler = vLLMSampler( - model_id=GEN_MODEL_ID, - engine_args={ - 'gpu_memory_utilization': GEN_GPU_MEM, - 'max_model_len': GEN_MAX_MODEL_LEN, - }, - device_mesh=gen_mesh, - remote_group='sampler', - ) - sampler.set_template('Qwen3_5Template', model_id=GEN_MODEL_ID, - enable_thinking=True, max_length=GEN_MAX_MODEL_LEN) - sys.stderr.write(f'[aops] vLLM sampler ready (model={GEN_MODEL_ID})\n') - - gen_params = TwinkleSamplingParams( - max_tokens=GEN_MAX_TOKENS, - temperature=GEN_TEMPERATURE, - top_p=GEN_TOP_P, - num_samples=1, - ) - - emb_model = emb_template = tbl = None - if is_rag: - import lancedb - db = lancedb.connect(args.db_path) - if args.table not in db.table_names(): - raise SystemExit(f'Table "{args.table}" not found in {args.db_path}') - tbl = db.open_table(args.table) - sys.stderr.write(f'[aops] LanceDB rows={tbl.count_rows()}\n') - - emb_model = TransformersModel( - model_id=EMBED_MODEL_ID, device_mesh=emb_mesh, - remote_group='emb_model') - emb_model.set_processor(InputProcessor) - emb_model.set_loss(InfonceLoss, temperature=0.03, use_batch=True) - emb_template = Qwen3_5Template( - model_id=EMBED_MODEL_ID, max_length=EMBED_MAX_LENGTH, - truncation_strategy='delete', enable_thinking=False) - sys.stderr.write('[aops] embedding model ready\n') - - # -- Condenser setup (API primary + optional local vLLM) ------------------- - condenser_api_client = None - condenser_sampler_obj = None - condenser_params = None - condenser_special_tokens = None - - if args.condense and is_rag: - condenser_api_client = OpenAIClient( - model=CONDENSE_API_MODEL, api_key=CONDENSE_API_KEY, - base_url=CONDENSE_BASE_URL) - sys.stderr.write(f'[condense] API client ready (model={CONDENSE_API_MODEL})\n') - - if condenser_gpus > 0: - condenser_mesh = DeviceMesh.from_sizes( - world_size=condenser_gpus, dp_size=condenser_gpus) - condenser_sampler_obj = vLLMSampler( - model_id=CONDENSE_MODEL_ID, - engine_args={'gpu_memory_utilization': 0.8, 'max_model_len': 32768}, - device_mesh=condenser_mesh, - remote_group='condenser', - ) - condenser_sampler_obj.set_template( - 'Qwen3_5Template', model_id=CONDENSE_MODEL_ID, - enable_thinking=False, truncation_strategy='delete', - max_length=32768) - condenser_template = Qwen3_5Template( - model_id=CONDENSE_MODEL_ID, max_length=32768, - enable_thinking=False, truncation_strategy='delete') - condenser_special_tokens = set(condenser_template.tokenizer.all_special_tokens) - condenser_params = TwinkleSamplingParams( - max_tokens=CONDENSE_MAX_TOKENS, - temperature=CONDENSE_TEMPERATURE, - top_p=0.5, num_samples=1) - sys.stderr.write(f'[condense] local vLLM ready (model={CONDENSE_MODEL_ID})\n') - - # -- Hint analysis API client (reuses condenser API config) ----------------- - hint_api_client = None - if args.hint and is_rag: - if condenser_api_client is not None: - hint_api_client = condenser_api_client - else: - hint_api_client = OpenAIClient( - model=CONDENSE_API_MODEL, api_key=CONDENSE_API_KEY, - base_url=CONDENSE_BASE_URL) - sys.stderr.write(f'[hint] API hint analysis enabled (model={CONDENSE_API_MODEL})\n') - - # -- LLM decontam API client --------------------------------------------------- - decontam_api_client = None - if args.llm_decontam and is_rag: - if hint_api_client is not None: - decontam_api_client = hint_api_client - elif condenser_api_client is not None: - decontam_api_client = condenser_api_client - else: - decontam_api_client = OpenAIClient( - model=CONDENSE_API_MODEL, api_key=CONDENSE_API_KEY, - base_url=CONDENSE_BASE_URL) - sys.stderr.write(f'[decontam-llm] LLM decontamination enabled (model={CONDENSE_API_MODEL})\n') - - correct_count = 0 - total_count = 0 - skipped_indices: List[int] = [] # problems skipped by RAG (no valid trace) - evaluated_indices: List[int] = [] # problems actually evaluated - debug_records: List[Dict[str, Any]] = [] - - os.makedirs(os.path.dirname(args.output) or '.', exist_ok=True) - out_f = open(args.output, 'w', encoding='utf-8') - - # Open problem-ids files for incremental writing (RAG mode only) - ids_f = None - skip_f = None - if is_rag: - os.makedirs(os.path.dirname(args.problem_ids_file) or '.', exist_ok=True) - ids_f = open(args.problem_ids_file, 'w', encoding='utf-8') - skip_path = args.problem_ids_file.replace('.json', '_skipped.json') - skip_f = open(skip_path, 'w', encoding='utf-8') - - # -- RAG batch preparation (embed + retrieve + decontam + condense + hint) -- - def _prepare_rag_batch(batch_start: int): - """Prepare a RAG batch: returns (prompts, batch, all_examples, - hint_analyses, kept_global_indices, batch_skipped_indices) or None.""" - batch_end = min(batch_start + args.batch_size, len(records)) - batch = records[batch_start:batch_end] - problems = [r['problem'] for r in batch] - - query_vecs = get_embeddings(emb_model, emb_template, problems, EMB_GPUS) - use_raw = not args.use_cot_compressed - all_examples = retrieve_examples(tbl, query_vecs, args.top_k, - use_raw, args.sim_threshold, - problems=problems, - decontam_threshold=args.decontam_threshold) - if args.use_cot_compressed: - for exs in all_examples: - for ex in exs: - ex['thinking'] = _strip_condenser_markers(ex['thinking']) - - if args.llm_decontam and decontam_api_client: - all_examples = _llm_decontaminate( - decontam_api_client, problems, all_examples) - - if args.condense and condenser_api_client: - all_examples = condense_traces( - all_examples, problems, condenser_api_client, - condenser_sampler=condenser_sampler_obj, - compress_params=condenser_params, - special_tokens=condenser_special_tokens, - max_output_len=args.condense_max_len, - dp_size=condenser_gpus) - - hint_analyses = None - if args.hint and hint_api_client: - hint_analyses = _api_hint_analysis_batch( - hint_api_client, problems, all_examples) - - keep_mask = [] - for pi, (r, examples) in enumerate(zip(batch, all_examples)): - if not examples: - keep_mask.append(False) - elif hint_analyses and hint_analyses[pi]: - keep_mask.append(True) - else: - usable = [ex for ex in examples - if len(ex['thinking']) <= args.max_trace_len] - keep_mask.append(bool(usable)) - - batch_skipped = [] - for pi, keep in enumerate(keep_mask): - if not keep: - batch_skipped.append(batch_start + pi) - - kept_batch = [] - kept_examples = [] - kept_hints = [] - kept_global_indices = [] - for pi, keep in enumerate(keep_mask): - if keep: - kept_batch.append(batch[pi]) - kept_examples.append(all_examples[pi]) - kept_hints.append(hint_analyses[pi] if hint_analyses else None) - kept_global_indices.append(batch_start + pi) - - if not kept_batch: - return None, None, None, None, None, batch_skipped - - prompts = [] - for pi, (r, examples) in enumerate(zip(kept_batch, kept_examples)): - if kept_hints[pi]: - prompts.append(build_hint_prompt(r['problem'], kept_hints[pi])) - else: - filtered = [{'query': ex['query'], 'thinking': ex['thinking']} - for ex in examples - if len(ex['thinking']) <= args.max_trace_len] - prompts.append(build_rag_prompt(r['problem'], filtered)) - - return prompts, kept_batch, kept_examples, kept_hints, kept_global_indices, batch_skipped - - target_reached = False - batch_starts = list(range(0, len(records), args.batch_size)) - - if is_rag: - # Pipeline: prefetch next batch while current batch generates - from concurrent.futures import Future - prefetch_pool = ThreadPoolExecutor(max_workers=1) - # Prepare first batch synchronously - cur_result = _prepare_rag_batch(batch_starts[0]) - - for bi, batch_start in enumerate(batch_starts): - if target_reached: - break - prompts, batch, all_examples, hint_analyses, kept_global_indices, batch_skipped = cur_result - skipped_indices.extend(batch_skipped or []) - if skip_f and batch_skipped: - for sid in batch_skipped: - skip_f.write(f'{sid}\n') - skip_f.flush() - - # Submit next batch preparation in background - next_future: Optional[Future] = None - if bi + 1 < len(batch_starts) and not target_reached: - next_future = prefetch_pool.submit(_prepare_rag_batch, batch_starts[bi + 1]) - - if prompts is None: - # Entire batch skipped - cur_result = next_future.result() if next_future else None - continue - - # Generate (runs on gen GPU while next batch prepares on emb GPU + API) - responses = sampler.sample(prompts, gen_params) - - for i, (rec, resp) in enumerate(zip(batch, responses)): - seq = resp.sequences[0] if resp and resp.sequences else None - raw_output = '' - if seq is not None: - raw_output = seq.decoded or '' - raw_output = re.sub(r'<\|[^|]+\|>', '', raw_output).rstrip() - - predicted = extract_boxed(raw_output) - is_correct = answers_match(predicted, rec['reference_answer']) - if is_correct: - correct_count += 1 - total_count += 1 - - global_idx = kept_global_indices[i] - evaluated_indices.append(global_idx) - if ids_f: - ids_f.write(f'{global_idx}\n') - ids_f.flush() - - debug_rec = { - 'idx': global_idx, - 'reference_answer': rec['reference_answer'], - 'predicted': predicted, - 'is_correct': is_correct, - 'problem': rec['problem'], - 'model_output': raw_output, - } - if rec.get('level'): - debug_rec['level'] = rec['level'] - if rec.get('type'): - debug_rec['type'] = rec['type'] - debug_rec['num_traces'] = len(all_examples[i]) - if all_examples[i]: - ex0 = all_examples[i][0] - debug_rec['similarity'] = ex0.get('_sim', 0.0) - debug_rec['retrieved_query'] = ex0.get('query', '') - debug_rec['raw_trace_len'] = ex0.get('_raw_trace_len', 0) - debug_rec['condensed_trace'] = ex0['thinking'] - debug_rec['condensed_trace_len'] = len(ex0['thinking']) - debug_rec['condense_source'] = ex0.get('_condense_source', '') - if hint_analyses and hint_analyses[i]: - debug_rec['hint_analysis'] = hint_analyses[i] - debug_records.append(debug_rec) - out_f.write(json.dumps(debug_rec, ensure_ascii=False) + '\n') - out_f.flush() - - acc = correct_count / total_count if total_count else 0 - sys.stderr.write( - f' [{total_count}/{args.target_eval}] ' - f'acc={acc:.4f} ({correct_count}/{total_count})\n') - - if args.target_eval > 0 and total_count >= args.target_eval: - target_reached = True - - # Collect prefetched result for next iteration (skip if done) - if not target_reached and next_future: - cur_result = next_future.result() - else: - cur_result = None - - prefetch_pool.shutdown(wait=True) - else: - # Direct mode: no pipeline needed, just batch generate - for batch_start in batch_starts: - batch_end = min(batch_start + args.batch_size, len(records)) - batch = records[batch_start:batch_end] - prompts = [build_direct_prompt(r['problem']) for r in batch] - - responses = sampler.sample(prompts, gen_params) - - for i, (rec, resp) in enumerate(zip(batch, responses)): - seq = resp.sequences[0] if resp and resp.sequences else None - raw_output = '' - if seq is not None: - raw_output = seq.decoded or '' - raw_output = re.sub(r'<\|[^|]+\|>', '', raw_output).rstrip() - - predicted = extract_boxed(raw_output) - is_correct = answers_match(predicted, rec['reference_answer']) - if is_correct: - correct_count += 1 - total_count += 1 - - global_idx = original_indices[batch_start + i] - evaluated_indices.append(global_idx) - - debug_rec = { - 'idx': global_idx, - 'reference_answer': rec['reference_answer'], - 'predicted': predicted, - 'is_correct': is_correct, - 'problem': rec['problem'], - 'model_output': raw_output, - } - if rec.get('level'): - debug_rec['level'] = rec['level'] - if rec.get('type'): - debug_rec['type'] = rec['type'] - debug_records.append(debug_rec) - out_f.write(json.dumps(debug_rec, ensure_ascii=False) + '\n') - out_f.flush() - - acc = correct_count / total_count if total_count else 0 - sys.stderr.write( - f' [{total_count}/{len(records)}] ' - f'acc={acc:.4f} ({correct_count}/{total_count})\n') - - overall_acc = correct_count / total_count if total_count else 0 - print(f'\n{"=" * 60}') - print(f'{args.dataset.upper()} — mode={args.mode}, model={GEN_MODEL_ID}') - print(f' n={total_count}, seed={args.seed}') - if is_rag: - print(f' evaluated={len(evaluated_indices)}, skipped={len(skipped_indices)}') - print(f'{"=" * 60}') - print(f'Overall accuracy: {overall_acc:.4f} ({correct_count}/{total_count})') - - # Per-level breakdown (MATH: the difficulty-vs-gain curve we care about). - if any(r.get('level') for r in debug_records): - from collections import defaultdict - per = defaultdict(lambda: [0, 0]) # level -> [correct, total] - for r in debug_records: - lv = r.get('level', 'Unknown') - per[lv][1] += 1 - if r['is_correct']: - per[lv][0] += 1 - print(f'\nPer-level accuracy:') - for lv in sorted(per.keys()): - c, t = per[lv] - print(f' {lv:>10}: {c/t:.4f} ({c}/{t})') - - out_f.close() - print(f'\n[output] {len(debug_records)} records saved to {args.output}') - - if ids_f: - ids_f.close() - print(f'[output] problem IDs ({len(evaluated_indices)}) saved to {args.problem_ids_file}') - if skip_f: - skip_f.close() - if skipped_indices: - print(f'[output] skipped IDs ({len(skipped_indices)}) saved to ' - f'{args.problem_ids_file.replace(".json", "_skipped.json")}') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/legacy/eval_math_by_level.sh b/cookbook/exp/legacy/eval_math_by_level.sh deleted file mode 100755 index d4f8bbdb3..000000000 --- a/cookbook/exp/legacy/eval_math_by_level.sh +++ /dev/null @@ -1,59 +0,0 @@ -#!/bin/bash -# MATH (Hendrycks) difficulty-stratified evaluation. -# -# Goal: measure how the (raw) RAG gain over direct varies with problem -# difficulty (Level 1-5). Runs raw RAG first (retrieve -> qwen3.7-max condense -# -> inject, no hint filtering; it writes the problem-id file), then direct on -# the *same* problems for a paired comparison. -# -# Usage: -# COMPRESS_API_KEY=sk-xxx bash cookbook/exp/embedding/eval_math_by_level.sh -# -# Env knobs: -# PER_LEVEL problems per difficulty level (default 100 -> 500 total) -# SEED stratified-sampling seed (default 100; must match across runs) -# DB_PATH LanceDB retrieval index -# SIM / TOPK retrieval threshold / top-k - -set -euo pipefail - -export COMPRESS_API_KEY="${COMPRESS_API_KEY:?Set COMPRESS_API_KEY}" - -SCRIPT="cookbook/exp/embedding/eval_gpqa_rag.py" -PER_LEVEL="${PER_LEVEL:-100}" -SEED="${SEED:-100}" -SIM="${SIM:-0.75}" -TOPK="${TOPK:-1}" -OUTDIR="./output/thinking_rag" -DB_PATH="${DB_PATH:-./output.oldemb/thinking_rag/lance.db}" - -mkdir -p "$OUTDIR" - -echo "============================================================" -echo " MATH by level: raw RAG (qwen3.7-max condenser, no hint)" -echo " per_level=$PER_LEVEL seed=$SEED" -echo "============================================================" -python "$SCRIPT" \ - --dataset math --math-split test \ - --mode rag \ - --per-level "$PER_LEVEL" --seed "$SEED" \ - --db-path "$DB_PATH" \ - --sim-threshold "$SIM" --top-k "$TOPK" \ - --condense \ - --output "$OUTDIR/math_rag_results.jsonl" - -echo "" -echo "============================================================" -echo " MATH by level: Direct (same problems as raw RAG)" -echo "============================================================" -# Direct reads math_rag_problem_ids.json (written above) to match the subset. -python "$SCRIPT" \ - --dataset math --math-split test \ - --mode direct \ - --per-level "$PER_LEVEL" --seed "$SEED" \ - --output "$OUTDIR/math_direct_results.jsonl" - -echo "" -echo "============================================================" -echo " Done. Compare with: python cookbook/exp/embedding/compare_math_levels.py" -echo "============================================================" diff --git a/cookbook/exp/legacy/eval_rag_recall.py b/cookbook/exp/legacy/eval_rag_recall.py deleted file mode 100644 index 19bc5ad6d..000000000 --- a/cookbook/exp/legacy/eval_rag_recall.py +++ /dev/null @@ -1,187 +0,0 @@ -"""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/legacy/eval_reflexion_skill.py b/cookbook/exp/legacy/eval_reflexion_skill.py deleted file mode 100644 index 3f7ddcc0d..000000000 --- a/cookbook/exp/legacy/eval_reflexion_skill.py +++ /dev/null @@ -1,762 +0,0 @@ -"""Phase-0 measurement for the reflexion self-skill scheme (see reflexion.md). - -Question this script answers: **on problems the base model first gets wrong, does -letting the SAME base model reflect on its failed attempt, distill a general -"skill", and re-solve WITH that skill in the system prompt, actually raise its -pass@k?** No LoRA is trained here — this is the upper-bound / go-no-go gate before -investing in a Skill-LoRA. If the base model's own skills don't help, training a -LoRA to produce them is pointless. - -It deliberately reuses ``eval_gpqa_rag`` verbatim (dataset, grader, prompts, sampling -config) so numbers are comparable with the other AoPS lines. Only the base model + -one vLLM sampler are used; the dataset is AoPS; validation is on the SAME problem -(no similar-problem retrieval). - -Per chunk of problems (all sampler calls are BATCHED across the whole chunk — never -one problem at a time): - 1. Initial solve — 1 rollout each; keep only problems the model got wrong. - 2. Skill generation — for each failed problem, the base model reads its own failed - attempt and produces N candidate skills (general reminders, no answer/solution). - 3. Leak filter — drop skills that leak the gold answer or a full solution. - 4. Baseline pass@k — K rollouts of the plain problem (the "no-skill" control). - 5. With-skill pass@k — K rollouts of the problem with each surviving skill in the - system prompt. - 6. Score — marginal = with-skill pass@k − baseline pass@k; keep the best skill. -A "pass" = answer correct AND generation terminated (no length cutoff). - -Everything useful (failed attempt, every candidate skill + leak flag, baseline and -per-skill rollout stats, marginals, best skill) is written to a JSONL **incrementally -after each chunk**, so partial runs are fully analysable. - -Launch (8 GPUs, tp=1 dp=8 by default): - python cookbook/exp/embedding/eval_reflexion_skill.py --n 64 --chunk-size 16 -""" -import argparse -import copy -import json -import os -import re -import sys -import threading -import time -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Any, Dict, List, Optional, Tuple - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.data_format import SamplingParams as TwinkleSamplingParams -from twinkle.sampler import vLLMSampler -from twinkle_agentic.protocol.openai import OpenAI as OpenAIClient - -# Reuse the reference eval's dataset + grading + prompts + sampling config so this -# line is directly comparable with eval_gpqa_rag / eval_dualline_math. -from eval_gpqa_rag import (DIRECT_SYSTEM, GEN_GPU_MEM, GEN_GPUS, GEN_MODEL_ID, - GEN_TEMPERATURE, GEN_TOP_P, MCQ_INSTRUCTION, answers_match, - build_direct_prompt, extract_boxed, load_aops) - -logger = get_logger() - -# vLLM parallel: tp=1, dp=GEN_GPUS by default (override GEN_TP; keep GEN_GPUS=8). -GEN_TP = int(os.environ.get('GEN_TP', 1)) - -# Leak-detector API (reuses eval_gpqa_rag's env names). A strong external model -# judges whether a candidate skill leaks THIS problem's answer/solution — catching -# what the string filter cannot (multiple-choice letters, derived-result leakage). -LEAK_API_KEY = os.environ.get('COMPRESS_API_KEY', '') -LEAK_BASE_URL = os.environ.get('COMPRESS_BASE_URL', - 'https://dashscope.aliyuncs.com/compatible-mode/v1') -LEAK_API_MODEL = os.environ.get('COMPRESS_MODEL', 'qwen3.7-max') - -# Global call spacer so bursts of leak-judge calls stay under the QPS limit. -_api_lock = threading.Lock() -_api_next = [0.0] - - -def _api_throttle(min_interval: float) -> None: - with _api_lock: - now = time.monotonic() - wait = max(0.0, _api_next[0] - now) - _api_next[0] = max(now, _api_next[0]) + min_interval - if wait > 0: - time.sleep(wait) - - -# --------------------------------------------------------------------------- -# Prompts (self-reflection skill generation + skill-conditioned solving) -# --------------------------------------------------------------------------- -SKILL_GEN_SYSTEM = ( - "You are a meticulous mathematics coach. You are shown a competition problem and a " - "student's FAILED attempt. Produce a SHORT list of general, reusable skills that " - 'would prevent this class of mistake on SIMILAR problems.\n\n' - 'OUTPUT FORMAT (strict):\n' - '- You may reason briefly first, but the final answer MUST be a markdown bullet ' - 'list of 3-5 items WRAPPED IN and tags. Output nothing after ' - '.\n' - '- Each item is ONE short imperative sentence (a rule, check, or habit).\n' - '- Inside the tags: no diagnosis narration, no "The student...", no headings, no ' - 'restating the problem or the examples.\n\n' - 'CONTENT RULES (strict):\n' - '- Do NOT reveal the final answer or the multiple-choice option.\n' - '- Do NOT state the specific numbers, values, or key intermediate results of THIS ' - 'problem.\n' - '- Do NOT give a step-by-step solution to THIS problem. Every item must be GENERAL ' - 'and transferable to other problems of the same type.\n\n' - 'Follow the example below for the exact tags, style, and level of generality.' -) - -SKILL_GEN_USER = ( - 'Problem:\n{problem}\n\n' - "The student's failed attempt (it may be long or may fail to terminate):\n" - '{attempt}\n\n' - 'Now output the skills bullet list.' -) - -# One-shot demonstration of the required format and generality (answer-free). -_EX_PROBLEM = 'Simplify $\\sqrt{72} + \\sqrt{18}$ and give the result.' -_EX_ATTEMPT = ( - 'The student added the radicands directly to get $\\sqrt{90}$ and concluded it ' - 'could not be simplified, never factoring out the perfect squares first.') -_EX_SKILLS = ( - '\n' - '- Before adding square roots, factor each radicand into a perfect square times a ' - 'remainder and move the perfect-square root outside.\n' - '- Never add radicands directly: $\\sqrt{a}+\\sqrt{b}\\ne\\sqrt{a+b}$.\n' - '- Only combine radical terms after reducing them to the same simplest radical ' - 'form.\n' - '- Sanity-check the simplified result by estimating each root numerically.\n' - '') - - -def build_skillgen_prompt(problem: str, attempt: str) -> Dict[str, Any]: - return {'messages': [ - {'role': 'system', 'content': SKILL_GEN_SYSTEM}, - {'role': 'user', - 'content': SKILL_GEN_USER.format(problem=_EX_PROBLEM, attempt=_EX_ATTEMPT)}, - {'role': 'assistant', 'content': _EX_SKILLS}, - {'role': 'user', - 'content': SKILL_GEN_USER.format(problem=problem, attempt=attempt)}, - ]} - - -# The skill is injected into the SYSTEM prompt (per reflexion.md), on top of the -# exact DIRECT_SYSTEM used by the baseline so the only difference is the reminders. -# Built by concatenation (NOT str.format): DIRECT_SYSTEM and the skill may contain -# literal braces (e.g. ``\boxed{}``, LaTeX), which would break ``.format``. -_SKILL_SOLVE_PREFIX = ( - DIRECT_SYSTEM + '\n\n' - 'Before you start, keep these reminders in mind to avoid common mistakes on this ' - 'type of problem:\n') -_SKILL_SOLVE_SUFFIX = ( - '\nApply them where relevant, but rely on your own reasoning to reach the answer.') - - -def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: - return {'messages': [ - {'role': 'system', 'content': _SKILL_SOLVE_PREFIX + skill + _SKILL_SOLVE_SUFFIX}, - {'role': 'user', 'content': problem + MCQ_INSTRUCTION}, - ]} - - -# --------------------------------------------------------------------------- -# Parsing / grading / leak filtering -# --------------------------------------------------------------------------- -_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') - - -def _clean_text(decoded: Optional[str]) -> str: - return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() - - -_BULLET_RE = re.compile(r'^\s*(?:[-*]|\d+[.)])\s') - - -def _extract_skill_list(text: str) -> str: - """Pull just the clean skill list out of a (possibly thinking-laden) output. - - The model is instructed to wrap the final list in ``...``, so - prefer that (robust to any preceding reasoning, closed or unterminated). Fall - back to dropping a ```` block and keeping from the first bullet onward. - """ - low = text.lower() - if '' in low: - start = low.index('') + len('') - end = low.index('') if '' in low else len(text) - return text[start:end].strip() - if '' in text: - text = text.rsplit('', 1)[-1] - text = re.sub(r'.*?', '', text, flags=re.DOTALL).strip() - lines = text.splitlines() - for i, line in enumerate(lines): - if _BULLET_RE.match(line): - return '\n'.join(lines[i:]).strip() - return text.strip() - - -def _bound_attempt(text: str, gen_tokens: int, budget_tokens: int) -> str: - """Keep a failed attempt within the skill-gen context budget. - - Round-2 (skill-gen) input contains the FULL round-1 attempt, and failed - attempts are often the ones that ran to the token cap (repetition loops), so - feeding them verbatim overflows max_model_len. Keep the head (real reasoning + - where it went wrong) plus a short tail (the final wrong answer); drop the - redundant middle. Token->char conversion uses THIS attempt's observed - chars-per-token so the cut fits precisely. - """ - if not text or gen_tokens <= budget_tokens: - return text - cpt = len(text) / max(1, gen_tokens) - head_tok = int(budget_tokens * 0.7) - tail_tok = budget_tokens - head_tok - head = text[:int(head_tok * cpt)] - tail = text[-int(tail_tok * cpt):] if tail_tok > 0 else '' - return f'{head}\n\n[... attempt truncated for length ...]\n\n{tail}' - - -def _parse_seq(seq, gold: str) -> Dict[str, Any]: - """Turn one sampled sequence into a graded rollout record. - - ``pass`` requires BOTH a correct boxed answer AND clean termination (a length - cutoff means the model never actually committed to the answer). - """ - text = _clean_text(getattr(seq, 'decoded', '') or '') - pred = extract_boxed(text) - correct = bool(pred) and answers_match(pred, gold) - terminated = getattr(seq, 'stop_reason', None) != 'length' - return { - 'pred': pred, - 'correct': correct, - 'terminated': terminated, - 'passed': bool(correct and terminated), - 'stop_reason': getattr(seq, 'stop_reason', None), - 'gen_tokens': len(getattr(seq, 'tokens', None) or []), - 'text': text, - } - - -def _pass_rate(rolls: List[Dict[str, Any]]) -> float: - return sum(1 for r in rolls if r['passed']) / len(rolls) if rolls else 0.0 - - -def _skill_leaks(skill: str, gold: str) -> Tuple[bool, str]: - """Reject a skill that leaks the answer, or is a degenerate / non-list output.""" - if not skill.strip(): - return True, 'empty' - bullets = [ln for ln in skill.splitlines() if _BULLET_RE.match(ln)] - if len(bullets) < 2: - return True, 'too_short' - low = skill.lower() - if 'item 1' in low and 'item 2' in low: # model echoed the format placeholder - return True, 'placeholder' - if '\\boxed' in skill: - return True, 'contains_boxed' - g = (gold or '').strip() - # Raw substring match is only trustworthy when the answer is specific enough that - # an incidental hit is unlikely. Short answers ('D', 'E', '1') would match almost - # any text, so leave those to the API judge instead of false-flagging every skill. - if len(g) >= 4 and g.lower() in skill.lower(): - return True, 'contains_gold_answer' - # Standalone multi-digit numbers from the gold answer leaking into the skill. - for num in re.findall(r'-?\d{2,}', g): - if re.search(r'(? Optional[bool]: - """Return True (leak) / False (clean) / None (unparseable or API error after retries). - - Only transient API errors are retried (with exponential backoff); an unparseable - verdict is deterministic at temperature 0, so retrying it is pointless. - """ - msgs = [ - {'role': 'system', 'content': _LEAK_JUDGE_SYSTEM}, - {'role': 'user', 'content': _LEAK_JUDGE_USER.format( - problem=problem[:4000], gold=gold, skill=skill[:4000])}, - ] - for attempt in range(retries + 1): - _api_throttle(min_interval) - try: - reply = api({'messages': msgs}, - TwinkleSamplingParams(temperature=0.0, max_tokens=16), - extra_body={'enable_thinking': False}) - except Exception as exc: # noqa: BLE001 — broad catch is intentional - logger.warning(f'[leak-judge] error (attempt {attempt + 1}/{retries + 1}): {exc}') - if attempt < retries: - time.sleep(min(4.0, 0.5 * 2 ** attempt)) # exponential backoff - continue - return None - verdict = (reply.get('content') or '').strip().upper() - if 'CLEAN' in verdict: - return False - if 'LEAK' in verdict: - return True - return None # unparseable — deterministic at temp 0, no point retrying - - -def _api_leak_batch(api: OpenAIClient, items: List[Tuple[int, str, str, str]], - concurrency: int, min_interval: float, - retries: int) -> Dict[int, Optional[bool]]: - """Judge many (key, problem, gold, skill) tuples in parallel; key -> verdict.""" - verdicts: Dict[int, Optional[bool]] = {} - if not items: - return verdicts - with ThreadPoolExecutor(max_workers=min(len(items), concurrency)) as pool: - futs = {pool.submit(_api_leak_judge_one, api, p, g, s, min_interval, retries): k - for (k, p, g, s) in items} - for fut in as_completed(futs): - verdicts[futs[fut]] = fut.result() - return verdicts - - -# --------------------------------------------------------------------------- -# Batched sampling (one shared sampler.sample per phase — never per problem) -# --------------------------------------------------------------------------- -def _pad_for_dp(prompts: List[Any], gen_dp: int) -> List[Any]: - """vLLM dp needs batch len >= dp; pad tail rounds and let the caller slice back.""" - if gen_dp <= 1 or not prompts or len(prompts) >= gen_dp: - return prompts - pad = [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] - return prompts + pad - - -def _run_samples(sampler, prompts: List[Any], num_samples: int, max_tokens: int, - gen_dp: int, temperature: Optional[float] = None, - top_p: Optional[float] = None, - top_k: Optional[int] = None) -> List[List[Any]]: - """One batched sampler call; return per-prompt list of raw sampled sequences. - - ``temperature``/``top_p``/``top_k`` default to the module's sampling config; pass - ``temperature=0.0`` for deterministic greedy decoding (SEAM-style executor scoring), - or a high ``temperature`` with ``top_k=-1`` for diverse multi-candidate sampling.""" - if not prompts: - return [] - params = TwinkleSamplingParams( - max_tokens=max_tokens, - temperature=GEN_TEMPERATURE if temperature is None else temperature, - top_p=GEN_TOP_P if top_p is None else top_p, - num_samples=num_samples, - **({} if top_k is None else {'top_k': top_k})) - padded = _pad_for_dp(prompts, gen_dp) - responses = sampler.sample(padded, params)[:len(prompts)] - return [list(r.sequences) if (r and r.sequences) else [] for r in responses] - - -def _set_thinking(sampler, args: argparse.Namespace, enabled: bool) -> None: - """Toggle the remote template's thinking mode. - - Skill generation wants thinking OFF so the model emits the short ```` - list directly (with thinking ON it burns the token budget reasoning and often - never reaches the list); solving wants it ON. ``set_template`` is a - remote_function, so this propagates to every sampler worker. - """ - sampler.set_template('Qwen3_5Template', model_id=GEN_MODEL_ID, - enable_thinking=enabled, max_length=args.max_model_len) - - -def _bounded_attempt_for(r: Dict[str, Any], args: argparse.Namespace) -> str: - """Per-problem bound so problem + attempt + skill output fits the context window.""" - prob_est = len(r['problem']) // 2 # conservative problem token estimate - budget = max(1024, args.max_model_len - args.skill_max_tokens - - args.attempt_reserve_tokens - prob_est) - return _bound_attempt(r['_init'][0]['text'], r['_init'][0]['gen_tokens'], budget) - - -def _filter_candidates(api: Optional[OpenAIClient], - cands: List[Tuple[Dict[str, Any], str]], - args: argparse.Namespace) -> None: - """Apply string + API leak filters to (problem, skill) candidates; append results. - - The cheap string filter runs first; the API judge only sees skills that pass it, - which is what catches MCQ-letter / derived-result / full-solution leakage. - """ - prepared = [] # [r, text, leaked(bool|None), reason] - for r, text in cands: - leaked, reason = _skill_leaks(text, r['reference_answer']) - prepared.append([r, text, True if leaked else None, reason]) - if api is not None: - items = [(i, prepared[i][0]['problem'], prepared[i][0]['reference_answer'], - prepared[i][1]) for i in range(len(prepared)) if prepared[i][2] is None] - verdicts = _api_leak_batch(api, items, args.api_concurrency, args.api_min_interval, - args.api_retries) - for key, _p, _g, _s in items: - v = verdicts.get(key) - if v is True: - prepared[key][2], prepared[key][3] = True, 'api_leak' - elif v is False: - prepared[key][2], prepared[key][3] = False, '' - else: - prepared[key][2], prepared[key][3] = False, 'api_uncertain' - for r, text, leaked, reason in prepared: - r['_skills'].append({'skill': text, 'leaked': bool(leaked), 'leak_reason': reason}) - - -def _build_skills(sampler, api: Optional[OpenAIClient], failed: List[Dict[str, Any]], - gen_dp: int, args: argparse.Namespace) -> None: - """Generate + extract + leak-filter skills, re-rolling problems short on clean ones. - - Clean skills accumulate across rounds; only problems still below ``min_survivors`` - clean skills are re-rolled, up to ``skill_retries`` extra rounds. - """ - for r in failed: - r['_skills'] = [] - todo = list(failed) - _set_thinking(sampler, args, False) # skill-gen: emit the list directly, no CoT - try: - for _ in range(args.skill_retries + 1): - if not todo: - break - sg_out = _run_samples( - sampler, - [build_skillgen_prompt(r['problem'], _bounded_attempt_for(r, args)) for r in todo], - args.n_skills, args.skill_max_tokens, gen_dp) - cands = [(r, _extract_skill_list(_clean_text(getattr(s, 'decoded', '') or ''))) - for r, seqs in zip(todo, sg_out) for s in seqs] - _filter_candidates(api, cands, args) - todo = [r for r in failed - if sum(1 for sk in r['_skills'] if not sk['leaked']) < args.min_survivors] - finally: - _set_thinking(sampler, args, True) # restore for solving phases - tot = sum(len(r['_skills']) for r in failed) - leaked = sum(1 for r in failed for sk in r['_skills'] if sk['leaked']) - sys.stderr.write(f' phase2: skills={tot} leaked={leaked} ' - f'({leaked / max(1, tot):.0%}); {len(todo)} still short of ' - f'{args.min_survivors} clean\n') - - -# --------------------------------------------------------------------------- -# Per-chunk pipeline -# --------------------------------------------------------------------------- -def process_chunk(sampler, api: Optional[OpenAIClient], chunk: List[Dict[str, Any]], - gen_dp: int, args: argparse.Namespace) -> List[Dict[str, Any]]: - """Run all 6 phases for one chunk (batched) and return per-problem records.""" - # --- Phase 1: initial solve, keep only the ones the model got wrong. --- - init_out = _run_samples( - sampler, [build_direct_prompt(r['problem']) for r in chunk], - args.init_samples, args.max_tokens, gen_dp) - for r, seqs in zip(chunk, init_out): - r['_init'] = [_parse_seq(s, r['reference_answer']) for s in seqs] - r['_init_pass'] = _pass_rate(r['_init']) - r['_failed'] = r['_init_pass'] == 0.0 - failed = [r for r in chunk if r['_failed']] - sys.stderr.write(f' phase1: {len(chunk)-len(failed)}/{len(chunk)} solved on ' - f'first try, {len(failed)} failed -> reflect\n') - - if failed: - # --- Baseline pass@k FIRST: defines which failures are genuinely hard. --- - # (A single initial rollout is noisy; an easy problem can fail phase 1 yet - # have a high pass@k, so measure the marginal only on truly hard problems.) - base_out = _run_samples( - sampler, [build_direct_prompt(r['problem']) for r in failed], - args.pass_k, args.max_tokens, gen_dp) - for r, seqs in zip(failed, base_out): - r['_baseline'] = [_parse_seq(s, r['reference_answer']) for s in seqs] - r['_baseline_pass'] = _pass_rate(r['_baseline']) - r['_hard'] = r['_baseline_pass'] <= args.hard_baseline_max - r['_skills'] = [] - r['_best'] = None - hard = [r for r in failed if r['_hard']] - sys.stderr.write(f' baseline: {len(hard)}/{len(failed)} failures are hard ' - f'(pass@{args.pass_k} <= {args.hard_baseline_max})\n') - - if hard: - # --- Skills (generate + leak filter + re-rollout) for HARD problems only. --- - _build_skills(sampler, api, hard, gen_dp, args) - - # --- With-skill pass@k (flatten hard-problem x surviving skill). --- - flat: List[Tuple[int, int]] = [] - ws_prompts: List[Any] = [] - for ri, r in enumerate(hard): - for si, sk in enumerate(r['_skills']): - if sk['leaked'] or not sk['skill'].strip(): - continue - flat.append((ri, si)) - ws_prompts.append(build_skill_solve_prompt(r['problem'], sk['skill'])) - ws_out = _run_samples(sampler, ws_prompts, args.pass_k, args.max_tokens, gen_dp) - for (ri, si), seqs in zip(flat, ws_out): - r = hard[ri] - sk = r['_skills'][si] - sk['rolls'] = [_parse_seq(s, r['reference_answer']) for s in seqs] - sk['with_pass'] = _pass_rate(sk['rolls']) - sk['marginal'] = sk['with_pass'] - r['_baseline_pass'] - - # --- Pick the best (highest marginal) surviving skill per hard problem. --- - for r in hard: - scored = [sk for sk in r['_skills'] if 'marginal' in sk] - r['_best'] = max(scored, key=lambda s: s['marginal']) if scored else None - - return [_make_record(r, args) for r in chunk] - - -def _roll_summary(roll: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: - out = {k: roll[k] for k in ('pred', 'correct', 'terminated', 'passed', - 'stop_reason', 'gen_tokens')} - if args.store_rollout_text: - out['text'] = roll['text'][:args.store_rollout_chars] - return out - - -def _make_record(r: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: - """Assemble the incremental JSONL record for one problem (solved or failed).""" - rec: Dict[str, Any] = { - 'problem': r['problem'], - 'reference_answer': r['reference_answer'], - 'tags': r.get('tags', []), - 'failed_first_try': r['_failed'], - 'init_pass_rate': r['_init_pass'], - 'init_attempt': { - 'text': r['_init'][0]['text'][:args.store_init_chars], - 'pred': r['_init'][0]['pred'], - 'stop_reason': r['_init'][0]['stop_reason'], - 'gen_tokens': r['_init'][0]['gen_tokens'], - }, - } - if not r['_failed']: - return rec - - best = r.get('_best') - rec['baseline_pass'] = r['_baseline_pass'] - # Genuinely hard = low baseline pass@k; only these count in the marginal stats. - rec['is_hard'] = bool(r.get('_hard')) - rec['baseline_rolls'] = [_roll_summary(x, args) for x in r['_baseline']] - rec['skills'] = [{ - 'skill': sk['skill'], - 'leaked': sk['leaked'], - 'leak_reason': sk['leak_reason'], - 'with_pass': sk.get('with_pass'), - 'marginal': sk.get('marginal'), - 'rolls': [_roll_summary(x, args) for x in sk.get('rolls', [])], - } for sk in r.get('_skills', [])] - rec['best_skill'] = best['skill'] if best else None - rec['best_marginal'] = best['marginal'] if best else None - rec['best_with_pass'] = best['with_pass'] if best else None - # "rescued" = a leak-free skill turned a fully-failing problem into some passes. - rec['rescued'] = bool(best and r['_baseline_pass'] == 0.0 and best['with_pass'] > 0.0) - rec['helped'] = bool(best and best['marginal'] > 0.0) - return rec - - -# --------------------------------------------------------------------------- -# Running summary -# --------------------------------------------------------------------------- -def _update_summary(summ: Dict[str, Any], recs: List[Dict[str, Any]]) -> None: - for rec in recs: - summ['n_total'] += 1 - if not rec['failed_first_try']: - summ['n_solved_first'] += 1 - continue - summ['n_failed'] += 1 - if not rec.get('is_hard'): - summ['n_failed_easy'] += 1 # failed phase 1 but easy on pass@k — excluded - continue - summ['n_hard'] += 1 - base = rec.get('baseline_pass', 0.0) - summ['sum_baseline_pass'] += base - if rec.get('best_marginal') is not None: - summ['n_with_skill'] += 1 - summ['sum_best_with_pass'] += rec.get('best_with_pass', 0.0) - summ['sum_best_marginal'] += rec.get('best_marginal', 0.0) - else: - # No clean skill produced for this hard problem -> skill adds no gain - # (count it honestly as marginal 0 rather than dropping it from the average). - summ['sum_best_with_pass'] += base - summ['n_helped'] += int(rec.get('helped', False)) - summ['n_rescued'] += int(rec.get('rescued', False)) - - -def _summary_report(summ: Dict[str, Any]) -> Dict[str, Any]: - nh = max(1, summ['n_hard']) - return { - 'record_type': 'summary', - 'n_total': summ['n_total'], - 'n_solved_first_try': summ['n_solved_first'], - 'n_failed_first_try': summ['n_failed'], - 'n_failed_but_easy': summ['n_failed_easy'], - 'n_hard': summ['n_hard'], - 'n_hard_with_skill': summ['n_with_skill'], - # Averages are over ALL hard problems; a hard problem with no clean skill - # counts as zero gain (with_pass == baseline), so with - base == marginal. - 'avg_baseline_pass_on_hard': summ['sum_baseline_pass'] / nh, - 'avg_best_with_skill_pass_on_hard': summ['sum_best_with_pass'] / nh, - 'avg_best_marginal_on_hard': summ['sum_best_marginal'] / nh, - 'n_helped_by_skill': summ['n_helped'], - 'n_rescued_from_zero': summ['n_rescued'], - 'frac_hard_helped': summ['n_helped'] / nh, - } - - -# --------------------------------------------------------------------------- -# Main -# --------------------------------------------------------------------------- -def main() -> None: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--n', type=int, default=64, help='AoPS problems to sample.') - p.add_argument('--seed', type=int, default=42) - p.add_argument('--chunk-size', type=int, default=16, - help='Problems per chunk. All sampler calls within a chunk are ' - 'batched; results are flushed to disk after each chunk.') - p.add_argument('--init-samples', type=int, default=1, - help='Rollouts for the initial solve. A problem is "failed" (and ' - 'sent to reflection) only if all initial rollouts are wrong.') - p.add_argument('--n-skills', type=int, default=8, - help='Candidate skills generated per failed problem.') - p.add_argument('--pass-k', type=int, default=8, - help='Rollouts per (baseline / with-skill) pass@k estimate.') - p.add_argument('--hard-baseline-max', type=float, default=0.25, - help='A failed problem counts as "hard" (included in the marginal ' - 'stats) only if its baseline pass@k <= this. Filters out easy ' - 'problems that merely failed the single initial rollout.') - p.add_argument('--max-model-len', type=int, default=30000, - help='Context window (engine + template). MUST exceed --max-tokens: ' - 'the round-2 skill-gen input holds the full round-1 attempt ' - 'plus the problem.') - p.add_argument('--max-tokens', type=int, default=20000, - help='Max generated tokens for solving rollouts (round-1 output cap).') - p.add_argument('--skill-max-tokens', type=int, default=2048, - help='Max tokens for skill generation. Enough to finish any thinking ' - 'and emit the short bullet list (which is then extracted).') - p.add_argument('--attempt-reserve-tokens', type=int, default=2048, - help='Tokens reserved for system prompt + wrappers when bounding the ' - 'failed attempt fed into skill generation (the problem length ' - 'is accounted for separately, per-problem).') - p.add_argument('--min-survivors', type=int, default=2, - help='Re-roll a problem\'s skills if fewer than this many survive the ' - 'leak filters.') - p.add_argument('--skill-retries', type=int, default=1, - help='Max extra skill-generation rounds for problems short on clean ' - 'skills (0 = no retry).') - p.add_argument('--api-concurrency', type=int, default=32, - help='Parallel workers for the API leak judge (max 32 recommended).') - p.add_argument('--api-min-interval', type=float, default=0.1, - help='Minimum seconds between API leak-judge calls (QPS guard).') - p.add_argument('--api-retries', type=int, default=3, - help='Retries on transient API errors per leak-judge call (exponential ' - 'backoff); only after these are exhausted is a skill kept as ' - 'api_uncertain.') - p.add_argument('--disable-api-leak', action='store_true', - help='Skip the API leak judge even if COMPRESS_API_KEY is set ' - '(string filter only).') - p.add_argument('--output', default='./output/reflexion_phase0/aops_results.jsonl') - p.add_argument('--store-init-chars', type=int, default=8000, - help='Truncate the stored failed-attempt text to this many chars.') - p.add_argument('--store-rollout-text', action='store_true', - help='Also store (truncated) text of every rollout, not just stats.') - p.add_argument('--store-rollout-chars', type=int, default=2000) - args = p.parse_args() - - records = load_aops(n=args.n, seed=args.seed) - sys.stderr.write(f'[reflexion] {len(records)} AoPS problems, chunk={args.chunk_size}, ' - f'init_samples={args.init_samples}, n_skills={args.n_skills}, ' - f'pass_k={args.pass_k}, max_tokens={args.max_tokens}\n') - - # --- 8-GPU vLLM sampler (tp=GEN_TP, dp=GEN_GPUS/GEN_TP). --- - if GEN_GPUS % GEN_TP != 0: - raise ValueError(f'GEN_GPUS ({GEN_GPUS}) must be divisible by GEN_TP ({GEN_TP})') - gen_dp = GEN_GPUS // GEN_TP - gen_mesh = DeviceMesh.from_sizes(world_size=GEN_GPUS, dp_size=gen_dp, tp_size=GEN_TP) - twinkle.initialize( - mode='ray', nproc_per_node=GEN_GPUS, - groups=[DeviceGroup(name='sampler', ranks=list(range(GEN_GPUS)), - device_type='GPU', gpus_per_worker=GEN_TP)], - lazy_collect=False) - sampler = vLLMSampler( - model_id=GEN_MODEL_ID, - engine_args={'gpu_memory_utilization': GEN_GPU_MEM, - 'max_model_len': args.max_model_len, - 'tensor_parallel_size': GEN_TP}, - device_mesh=gen_mesh, remote_group='sampler') - sampler.set_template('Qwen3_5Template', model_id=GEN_MODEL_ID, - enable_thinking=True, max_length=args.max_model_len) - sys.stderr.write(f'[reflexion] sampler ready (model={GEN_MODEL_ID}, tp={GEN_TP}, ' - f'dp={gen_dp})\n') - - # --- API leak judge (optional): reuses eval_gpqa_rag's COMPRESS_* env. --- - api: Optional[OpenAIClient] = None - if LEAK_API_KEY and not args.disable_api_leak: - api = OpenAIClient(model=LEAK_API_MODEL, api_key=LEAK_API_KEY, - base_url=LEAK_BASE_URL) - sys.stderr.write(f'[reflexion] leak judge ON via API model={LEAK_API_MODEL} ' - f'(concurrency={args.api_concurrency})\n') - else: - sys.stderr.write('[reflexion] leak judge OFF (string filter only) — set ' - 'COMPRESS_API_KEY to enable the API judge\n') - - os.makedirs(os.path.dirname(args.output) or '.', exist_ok=True) - summ = {k: 0 for k in ('n_total', 'n_solved_first', 'n_failed', 'n_failed_easy', - 'n_hard', 'n_with_skill', 'n_helped', 'n_rescued')} - summ.update({'sum_baseline_pass': 0.0, 'sum_best_with_pass': 0.0, - 'sum_best_marginal': 0.0}) - - with open(args.output, 'w', encoding='utf-8') as out_f: - # Line 1: run config, for reproducibility / later analysis. - out_f.write(json.dumps({ - 'record_type': 'config', 'model': GEN_MODEL_ID, 'dataset': 'aops', - 'n': len(records), 'seed': args.seed, 'init_samples': args.init_samples, - 'n_skills': args.n_skills, 'pass_k': args.pass_k, - 'hard_baseline_max': args.hard_baseline_max, - 'max_model_len': args.max_model_len, 'max_tokens': args.max_tokens, - 'skill_max_tokens': args.skill_max_tokens, - 'api_leak_judge': api is not None, - 'api_leak_model': LEAK_API_MODEL if api is not None else None, - 'min_survivors': args.min_survivors, 'skill_retries': args.skill_retries, - 'gpus': GEN_GPUS, 'tp': GEN_TP, 'started': int(time.time()), - }, ensure_ascii=False) + '\n') - out_f.flush() - - n_chunks = (len(records) + args.chunk_size - 1) // args.chunk_size - for ci in range(n_chunks): - chunk = records[ci * args.chunk_size:(ci + 1) * args.chunk_size] - sys.stderr.write(f'[reflexion] chunk {ci+1}/{n_chunks} ({len(chunk)} problems)\n') - recs = process_chunk(sampler, api, chunk, gen_dp, args) - for rec in recs: # incremental write per problem - out_f.write(json.dumps(rec, ensure_ascii=False) + '\n') - out_f.flush() - _update_summary(summ, recs) - rep = _summary_report(summ) - sys.stderr.write( - f' running: failed={rep["n_failed_first_try"]} ' - f'(easy={rep["n_failed_but_easy"]}) hard={rep["n_hard"]} ' - f'base_pass={rep["avg_baseline_pass_on_hard"]:.3f} ' - f'skill_pass={rep["avg_best_with_skill_pass_on_hard"]:.3f} ' - f'helped={rep["n_helped_by_skill"]} rescued={rep["n_rescued_from_zero"]}\n') - - report = _summary_report(summ) - out_f.write(json.dumps(report, ensure_ascii=False) + '\n') - out_f.flush() - - print('\n' + '=' * 60) - print(f'Reflexion Phase-0 — model={GEN_MODEL_ID}, dataset=aops, n={report["n_total"]}') - print('=' * 60) - print(f'solved on first try : {report["n_solved_first_try"]}/{report["n_total"]}') - print(f'failed first try : {report["n_failed_first_try"]} ' - f'(easy, excluded: {report["n_failed_but_easy"]})') - print(f'hard (baseline pass@{args.pass_k}<= {args.hard_baseline_max}) : {report["n_hard"]}') - print(f' avg baseline pass@{args.pass_k:<2} : {report["avg_baseline_pass_on_hard"]:.4f}') - print(f' avg best-skill pass@{args.pass_k:<2} : {report["avg_best_with_skill_pass_on_hard"]:.4f}') - print(f' avg best marginal : {report["avg_best_marginal_on_hard"]:+.4f}') - print(f' helped by a skill : {report["n_helped_by_skill"]}/{report["n_hard"]}') - print(f' rescued from 0 pass : {report["n_rescued_from_zero"]}/{report["n_hard"]}') - print(f'\n[output] {args.output}') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/legacy/grpo_baseline.py b/cookbook/exp/legacy/grpo_baseline.py deleted file mode 100644 index 237f9b065..000000000 --- a/cookbook/exp/legacy/grpo_baseline.py +++ /dev/null @@ -1,593 +0,0 @@ -"""HotpotQA GRPO baseline — full context, no chunking, no compression, no tools. - -This is the **control group** for ``grpo_condensed.py``. Both scripts share: - * dataset (HotpotQA fullwiki, hard split) - * preprocessing (``HotpotQAProcessor`` with ``[K] Title: ...`` passages) - * GRPO infra (model / sampler / device mesh / hyperparams) - * rollout class (``MultiTurnRollout`` from ``multi_turn.py``) - -The only differences are intentional: - * no ``NativeChunker`` / ``ModelCondenser`` (full passages go in verbatim) - * no tools registered (``ToolManager()`` is empty) - * ``max_turns=1`` so the rollout is effectively single-turn - * simplified system prompt (no ```` / ``extract_condensed`` syntax) - * ``F1Reward + CoTReward`` only (no ``ToolExploreReward``) - * traces → ``rollout_trace_baseline.jsonl`` - * checkpoints prefixed ``hotpotqa-grpo-baseline-*`` - -Keeping the same ``MultiTurnRollout`` code path on both sides means any -training-loop-level discrepancy between the two runs is attributable to -the chunk+condense pipeline, not to differences in rollout plumbing. -""" - -import math -import os -import re -from typing import Any, Dict, List, Optional - -import swanlab -from peft import LoraConfig - -import twinkle -from twinkle import DeviceMesh, DeviceGroup, get_logger -from twinkle.advantage import GRPOAdvantage -from twinkle.checkpoint_engine import CheckpointEngineManager -from twinkle.data_format import Message, SamplingParams, Trajectory -from twinkle.dataloader import DataLoader -from twinkle.dataset import Dataset, DatasetMeta -from twinkle.metric import CompletionRewardMetric -from twinkle.model import TransformersModel -from twinkle.preprocessor.base import Preprocessor -from twinkle.processor import InputProcessor -from twinkle.sampler import vLLMSampler -from twinkle.template import Qwen3_5Template -from twinkle_agentic.reward import F1Reward, CoTReward -from twinkle_agentic.rollout.multi_turn import MultiTurnRollout -from twinkle_agentic.tools.tool_manager import ToolManager - -logger = get_logger() - -MODEL_ID = os.environ.get('MODEL_ID', 'ms://Qwen/Qwen3.5-4B') -USE_MEGATRON = bool(int(os.environ.get('USE_MEGATRON', '1'))) - -MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 4)) -SAMPLER_GPUS = int(os.environ.get('SAMPLER_GPUS', 4)) -NUM_GPUS = MODEL_GPUS + SAMPLER_GPUS - -NUM_GENERATIONS = int(os.environ.get('NUM_GENERATIONS', 8)) -MAX_NEW_TOKENS = int(os.environ.get('MAX_NEW_TOKENS', 4096)) -LEARNING_RATE = float(os.environ.get('LR', 1e-5)) -NUM_EPOCHS = int(os.environ.get('NUM_EPOCHS', 1)) -MAX_STEPS = int(os.environ.get('MAX_STEPS', 0)) -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)) -GRADIENT_ACCUMULATION_STEPS = int(os.environ.get('GRADIENT_ACCUMULATION_STEPS', 1)) -ADAPTER_NAME = 'default' -SAVE_STEPS = int(os.environ.get('SAVE_STEPS', 1000)) -LORA_RANK = int(os.environ.get('LORA_RANK', 16)) - -# Single-turn baseline; tools are not registered, but we keep MultiTurnRollout -# to share the rollout code path with the condensed variant. ``max_turns=1`` -# guarantees the loop runs exactly one sampling pass per trajectory. -MAX_TURNS = int(os.environ.get('MAX_TURNS', 1)) - -HOTPOTQA_NUM_PROC = int(os.environ.get('HOTPOTQA_NUM_PROC', 16)) -HOTPOTQA_MAX_LENGTH = int(os.environ.get('HOTPOTQA_MAX_LENGTH', 64000)) - -F1_REWARD_WEIGHT = float(os.environ.get('F1_REWARD_WEIGHT', 1.0)) -COT_REWARD_WEIGHT = float(os.environ.get('COT_REWARD_WEIGHT', 0.2)) - -# KL penalty coefficient; 0 disables KL (and skips the ref forward pass entirely). -KL_BETA = float(os.environ.get('KL_BETA', 0.02)) - -# Entropy bonus coefficient; 0 disables entropy compute path. -ENTROPY_COEF = float(os.environ.get('ENTROPY_COEF', 0.0)) - -# CISPO token-level IS clamp thresholds (asymmetric: 0.2 / 0.28). -CISPO_EPS_LOW = float(os.environ.get('CISPO_EPS_LOW', 0.2)) -CISPO_EPS_HIGH = float(os.environ.get('CISPO_EPS_HIGH', 0.2)) - -# High-KL token capture: top-K per microbatch dumped into log_dict['_high_kl_records']. 0 = disabled. -HIGH_KL_TOPK = int(os.environ.get('HIGH_KL_TOPK', 0)) - -DATASET_PATH = os.environ.get( - 'DATASET_PATH', - os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), - 'hotpotqa_fullwiki_reannotated_12k.jsonl')) -F1_BINARY_THRESHOLD = float(os.environ.get('F1_BINARY_THRESHOLD', 0.5)) - -_ROLLOUT_TRACE_DIR = os.environ.get( - 'ROLLOUT_TRACE_BASELINE_DIR', 'rollout_trace_baseline') - -SYSTEM_PROMPT = """You are a careful multi-hop QA assistant. - -You will receive a question and a set of supporting passages. Each passage \ -is shown inline as plain text in the form `[K] Title: ...`, where `K` is the \ -passage index. All passages are already complete — there is no extraction \ -or expansion step. - -## Workflow - -Step 1: Read every passage and identify which ones are relevant to the question. -Step 2: Reason step by step, citing the passage indices you used. - Step N: From passage [K], I learn that [fact A]. - Step N+1: From passage [M], I learn that [fact B]. - Step N+2: Combining these, the answer is ... -Step 3: Emit the final answer in `\\boxed{...}`. - -Only answer when you are confident in the supporting facts. - -## Output Format -End your final response with \\boxed{answer}, e.g. \\boxed{Delhi}. -Keep the boxed text short: a name, entity, date, or "yes"/"no". -Answers not inside \\boxed{} will not be scored.""" - - -_F1_REWARD: Optional[F1Reward] = F1Reward() -_COT_REWARD: Optional[CoTReward] = CoTReward() - - -def compute_rewards(trajectories: List[Dict[str, Any]]): - f1_raw = _F1_REWARD(trajectories) - f1 = [1.0 if v >= F1_BINARY_THRESHOLD else 0.0 for v in f1_raw] if F1_BINARY_THRESHOLD > 0 else f1_raw - cot = _COT_REWARD(trajectories) - total = [ - F1_REWARD_WEIGHT * a + COT_REWARD_WEIGHT * c - for a, c in zip(f1, cot) - ] - return total, f1, cot - - -class HotpotQAProcessor(Preprocessor): - """Preprocessor for the reannotated HotpotQA JSONL. Passages are emitted - as ``[K] Title: ...`` lines. Rows with ``verdict='drop'`` are excluded; - ``question_fixed`` is used in place of ``question`` when present.""" - - def __init__(self, system: str = SYSTEM_PROMPT): - self.system = system - - def __call__(self, rows: Dict[str, List[Any]]) -> Dict[str, List[Any]]: - rows = self.map_col_to_row(rows) - rows = [self.preprocess(row) for row in rows] - rows = [r for r in rows if r is not None] - rows = self.map_row_to_col(rows) - return rows - - @staticmethod - def _format_context(context: Dict[str, Any]) -> str: - titles = context.get('title', []) or [] - sentences = context.get('sentences', []) or [] - lines = [] - for i, (title, sents) in enumerate(zip(titles, sentences), start=1): - if isinstance(sents, list): - body = ' '.join(s.strip() for s in sents if s and s.strip()) - else: - body = str(sents).strip() - lines.append(f'[{i}] {title}: {body}') - return '\n\n'.join(lines) - - def preprocess(self, row: Dict[str, Any]) -> Optional[Trajectory]: - if (row.get('verdict') or '').strip().lower() == 'drop': - return None - question = row.get('question_fixed') or row['question'] - answers = row.get('answers') - if isinstance(answers, list) and answers: - golds = [str(a).strip() for a in answers if str(a).strip()] - else: - golds = [s for s in [(row.get('answer', '') or '').strip()] if s] - context_block = self._format_context(row.get('context', {}) or {}) - user_msg = f'Question: {question}\n\nContext:\n\n{context_block}' - messages = [ - Message(role='system', content=self.system), - Message(role='user', content=user_msg), - ] - return Trajectory(messages=messages, user_data=[('ground_truth', g) for g in golds]) - - -def create_hotpotqa_dataset() -> Dataset: - dataset = Dataset() - dataset.add_dataset(DatasetMeta(DATASET_PATH)) - logger.info('[dataset] loaded %s: %d rows', DATASET_PATH, len(dataset)) - - dataset.set_template( - 'Qwen3_5Template', model_id=MODEL_ID, max_length=HOTPOTQA_MAX_LENGTH, - truncation_strategy='delete', enable_thinking=False) - _HOTPOTQA_COLS = ['id', 'question', 'question_fixed', 'answers', - 'original_answer', 'type', 'level', 'verdict', - 'reasoning', 'supporting_facts', 'context'] - dataset.map(HotpotQAProcessor(system=SYSTEM_PROMPT), - remove_columns=_HOTPOTQA_COLS) - return dataset - - -# Matches a LaTeX ``\boxed{...}`` final-answer marker — used to flag -# rollouts that never committed an answer. Brace-balanced is overkill for -# a logging heuristic; a non-greedy ``[^}]*`` is good enough. -_BOXED_RE = re.compile(r'\\boxed\{[^}]*\}') - -# Pulls the leading number out of pre-formatted metric strings such as -# ``'0.03 iters/s'`` / ``'1.000000e-05'`` / ``'30 seconds'`` emitted by -# ``TrainMetric`` and ``GRPOMetric``. We use this in ``_coerce_for_swanlab`` -# so swanlab can build line charts instead of dropping those keys with a -# ``failed to create chart for key '...': invalid value type`` warning. -_LEADING_NUMBER_RE = re.compile(r'[-+]?\d*\.?\d+(?:[eE][-+]?\d+)?') - - -def _coerce_for_swanlab(log_dict: Dict[str, Any]) -> Dict[str, Any]: - """Cast string-valued metrics to float for swanlab line charts. - - ``TrainMetric.calculate()`` and ``GRPOMetric.calculate()`` return - pre-formatted strings (``'0.03 iters/s'``, ``'1.000000e-05'``, - ``'30 seconds'``, ``'0.8321'``). swanlab cannot build a line chart - from a string value and emits one warning per key per step. We extract - the leading number where possible; keys whose value can't be parsed - as a scalar are left as-is so they still show up in the text log. - """ - coerced: Dict[str, Any] = {} - for k, v in log_dict.items(): - if isinstance(v, bool) or isinstance(v, (int, float)): - coerced[k] = v - continue - if isinstance(v, str): - m = _LEADING_NUMBER_RE.search(v) - if m: - try: - coerced[k] = float(m.group()) - continue - except ValueError: - pass - coerced[k] = v - return coerced - - -def _last_assistant_text(trajectory: Dict[str, Any]) -> Optional[str]: - """Return the text of the last ``assistant`` message, or ``None``. - - ``content`` can be ``str`` | ``None`` | ``dict`` (single multimodal - part) | ``list[dict]`` (multiple parts). The downstream caller feeds - this into ``_BOXED_RE.search(...)``, so we collapse the visible text - into a single string and ignore non-text parts (images etc.). - """ - for m in reversed(trajectory.get('messages', [])): - if m.get('role') != 'assistant': - continue - c = m.get('content') - if c is None: - return None - if isinstance(c, str): - return c - if isinstance(c, dict): - return c.get('text') if c.get('type') == 'text' else None - if isinstance(c, list): - parts = [p.get('text') or '' for p in c - if isinstance(p, dict) and p.get('type') == 'text'] - return '\n'.join(parts) if parts else None - return str(c) - return None - - -def _compute_rollout_diagnostics( - trajectories: List[Dict[str, Any]], - n_turns_per_rollout: List[int], - per_rollout_completion_length: List[int], - f1_rewards: Optional[List[float]] = None, - old_logps: Optional[List[List[float]]] = None, -) -> Dict[str, float]: - """Aggregate rollout diagnostics for swanlab logging. - - Stripped-down version of the condensed variant's diagnostics — without - chunking we only care about (a) the longest non-trainable prefix - (system prompt + full passages), and (b) whether the rollout produced - a `\\boxed{}` final answer at all. ``avg_turns`` is logged for symmetry - even though it should be exactly 1.0 with ``MAX_TURNS=1``. - """ - out: Dict[str, float] = {} - if n_turns_per_rollout: - out['avg_turns'] = sum(n_turns_per_rollout) / len(n_turns_per_rollout) - - _max_non_trainable = 0 - for t, comp_len in zip(trajectories, per_rollout_completion_length): - ids = t.get('input_ids') or [] - non_trainable = max(0, len(ids) - int(comp_len or 0)) - if non_trainable > _max_non_trainable: - _max_non_trainable = non_trainable - out['non_trainable_tokens'] = _max_non_trainable - - if trajectories: - n_no_boxed = sum( - 0 if _BOXED_RE.search(_last_assistant_text(t) or '') else 1 - for t in trajectories) - out['no_boxed_rate'] = n_no_boxed / len(trajectories) - - def _content_chars(c: Any) -> int: - if not c: - return 0 - if isinstance(c, str): - return len(c) - if isinstance(c, dict): - if c.get('type') == 'text': - return len(c.get('text') or '') - return 0 - if isinstance(c, list): - total = 0 - for part in c: - if isinstance(part, dict) and part.get('type') == 'text': - total += len(part.get('text') or '') - elif isinstance(part, str): - total += len(part) - return total - # Unknown shape -- fall back to ``str()`` length rather than - # crashing, so a template quirk never breaks metric logging. - return len(str(c)) - - msg_chars_total, prompt_chars, asst_chars = [], [], [] - for t in trajectories: - total_i = prompt_i = asst_i = 0 - for m in (t.get('messages') or []): - role = m.get('role') - if role == 'system': - continue - n = _content_chars(m.get('content')) - total_i += n - if role in ('user', 'tool'): - prompt_i += n - elif role == 'assistant': - asst_i += n - msg_chars_total.append(total_i) - prompt_chars.append(prompt_i) - asst_chars.append(asst_i) - out['avg_chars_total_no_sys'] = sum(msg_chars_total) / len(msg_chars_total) - out['avg_chars_prompt_no_sys'] = sum(prompt_chars) / len(prompt_chars) - out['avg_chars_assistant'] = sum(asst_chars) / len(asst_chars) - - if f1_rewards is not None and old_logps is not None and f1_rewards: - per_traj_mean = [(sum(lp) / len(lp)) if lp else 0.0 for lp in old_logps] - pos_logp = [m for m, f1 in zip(per_traj_mean, f1_rewards) if f1 > 0] - zero_logp = [m for m, f1 in zip(per_traj_mean, f1_rewards) if f1 <= 0] - out['f1_correct_rate'] = len(pos_logp) / len(f1_rewards) - out['f1_zero_rate'] = len(zero_logp) / len(f1_rewards) - out['mean_old_logp_f1_pos'] = (sum(pos_logp) / len(pos_logp)) if pos_logp else 0.0 - out['mean_old_logp_f1_zero'] = (sum(zero_logp) / len(zero_logp)) if zero_logp else 0.0 - out['policy_confidence_f1_pos'] = math.exp(out['mean_old_logp_f1_pos']) - out['policy_confidence_f1_zero'] = math.exp(out['mean_old_logp_f1_zero']) - return out - - -def main(): - swanlab.init(project='twinkle') - - device_groups = [ - DeviceGroup(name='model', ranks=list(range(MODEL_GPUS)), device_type='GPU'), - DeviceGroup(name='sampler', ranks=list(range(MODEL_GPUS, 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, dp_size=SAMPLER_GPUS) - twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, - groups=device_groups, lazy_collect=False) - - logger.info('Building HotpotQA dataset (baseline, full context)') - _prebuilt_dataset = create_hotpotqa_dataset() - logger.info('Dataset ready: %d rows', len(_prebuilt_dataset)) - - GLOBAL_BATCH_SIZE = BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS - batches_per_epoch = max(1, len(_prebuilt_dataset) // GLOBAL_BATCH_SIZE) - # Single-turn baseline: every rollout produces exactly one assistant - # turn, so the per-batch optim-step count equals - # ceil(GLOBAL_BATCH_SIZE * NUM_GENERATIONS / MINI_BATCH_SIZE). - optim_steps_per_batch = max(1, (GLOBAL_BATCH_SIZE * NUM_GENERATIONS - + MINI_BATCH_SIZE - 1) // MINI_BATCH_SIZE) - steps_per_epoch = batches_per_epoch * optim_steps_per_batch - derived_total_steps = NUM_EPOCHS * steps_per_epoch - total_steps = min(MAX_STEPS, derived_total_steps) if MAX_STEPS > 0 else derived_total_steps - logger.info('Training horizon: %d steps (%d epochs × %d batches × %d steps/batch)', - total_steps, NUM_EPOCHS, batches_per_epoch, optim_steps_per_batch) - - lora_config = LoraConfig( - target_modules='all-linear', r=LORA_RANK, - lora_alpha=LORA_RANK * 2, lora_dropout=0.05) - - if USE_MEGATRON: - from twinkle.model.megatron import MegatronModel - model = MegatronModel( - model_id=MODEL_ID, device_mesh=model_mesh, remote_group='model', - mixed_precision='bf16', variable_seq_lengths=True) - else: - model = TransformersModel( - model_id=MODEL_ID, device_mesh=model_mesh, remote_group='model') - - model.add_adapter_to_model(ADAPTER_NAME, lora_config, - gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) - if USE_MEGATRON: - model.set_optimizer('default', lr=LEARNING_RATE) - model.set_lr_scheduler('default', lr_decay_steps=total_steps, max_lr=LEARNING_RATE) - else: - model.set_optimizer('AdamW', lr=LEARNING_RATE) - model.set_lr_scheduler('CosineAnnealingLR', T_max=total_steps, eta_min=0) - - model.set_loss('GRPOLoss', epsilon=CISPO_EPS_LOW, epsilon_high=CISPO_EPS_HIGH, - beta=KL_BETA, entropy_coef=ENTROPY_COEF) - model.set_processor(InputProcessor, padding_free=True) - model.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=False, max_length=HOTPOTQA_MAX_LENGTH) - - model.add_metric('GRPOMetric', is_training=True, - epsilon=CISPO_EPS_LOW, epsilon_high=CISPO_EPS_HIGH, - top_k_kl=HIGH_KL_TOPK) - - sampler = vLLMSampler( - model_id=MODEL_ID, - engine_args={ - 'gpu_memory_utilization': 0.8, 'max_model_len': 32768, - 'max_lora_rank': 32, 'enable_lora': True, - 'enable_tower_connector_lora': True, - }, - device_mesh=sampler_mesh, remote_group='sampler') - sampler.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=False, max_length=HOTPOTQA_MAX_LENGTH) - rollout_template = Qwen3_5Template( - MODEL_ID, max_length=HOTPOTQA_MAX_LENGTH, enable_thinking=False) - - ckpt_manager = CheckpointEngineManager(model=model, sampler=sampler) - - dataloader = DataLoader( - dataset=lambda: _prebuilt_dataset, - batch_size=GLOBAL_BATCH_SIZE, min_batch_size=GLOBAL_BATCH_SIZE) - - 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) - - def _trace_should_store(traj): - return True - - def _trace_is_success(traj): - return _F1_REWARD([traj])[0] > 0.0 - - rollout = MultiTurnRollout( - sampler=sampler, - template=rollout_template, - tool_manager=ToolManager(), - sampling_params=sampling_params, - max_turns=MAX_TURNS, - trace_dir=_ROLLOUT_TRACE_DIR or None, - trace_callback=_trace_should_store, - success_callback=_trace_is_success, - ) - - optim_step = 0 - logger.info('Starting HotpotQA GRPO baseline (no chunk / no condense / no tools)') - - def _epoch_cycle(dl, n_epochs): - for ep in range(1, n_epochs + 1): - logger.info(f'=== Epoch {ep}/{n_epochs} (step={optim_step}/{total_steps}) ===') - for batch in dl: - yield batch - - for batch in _epoch_cycle(dataloader, NUM_EPOCHS): - if optim_step >= total_steps: - break - - # Single source of truth for the step shown in swanlab / logger / rollout-trace filename. - batch_step = optim_step - - metrics.reset() - expand_prompts = [p for prompt in batch for p in [prompt] * NUM_GENERATIONS] - - ckpt_manager.sync_weights(merge_and_sync=False) - sampler.reset_prefix_cache() - - # Single batched rollout: each trajectory produces exactly one - # assistant turn (tools are unregistered, ``max_turns=1``). - all_trajectories: List[Dict[str, Any]] = rollout(expand_prompts) - n_turns_per_rollout = [int(t.get('turns') or 0) for t in all_trajectories] - per_rollout_completion_length = [ - sum(1 for l in (t.get('labels') or []) if l != -100) - for t in all_trajectories] - - total_rewards, f1_rewards, cot_rewards = compute_rewards(all_trajectories) - - rollout_advantages = advantage_fn( - total_rewards, num_generations=NUM_GENERATIONS, scale='group').tolist() - - all_f1_labels: List[bool] = [f > 0 for f in f1_rewards] - n_pos = sum(1 for p in all_f1_labels if p) - n_neg = sum(1 for p in all_f1_labels if not p) - pos_with_neg_adv = sum(1 for p, a in zip(all_f1_labels, rollout_advantages) if p and a < 0) - neg_with_pos_adv = sum(1 for p, a in zip(all_f1_labels, rollout_advantages) if not p and a > 0) - - all_old_logps: List[List[float]] = [ - [lp[0][1] for lp in (t.get('logprobs') or [])] for t in all_trajectories] - - # Skip homogeneous groups where gradient signal is meaningless - f1_pos_rate = n_pos / len(f1_rewards) if f1_rewards else 0.5 - if f1_pos_rate > 0.9 or f1_pos_rate < 0.1: - logger.info('[skip-homogeneous] f1_pos_rate=%.3f, skipping training update', f1_pos_rate) - metrics.accumulate( - completion_lengths=per_rollout_completion_length, - rewards={'total': total_rewards, 'f1': f1_rewards, 'cot': cot_rewards}) - log_dict = metrics.calculate() - log_dict.update(_compute_rollout_diagnostics( - all_trajectories, n_turns_per_rollout, per_rollout_completion_length, - f1_rewards=f1_rewards, old_logps=all_old_logps)) - log_dict['skipped'] = True - log_dict['pos_neg_adv_rate'] = pos_with_neg_adv / n_pos if n_pos else 0.0 - log_dict['neg_pos_adv_rate'] = neg_with_pos_adv / n_neg if n_neg else 0.0 - log_dict['adv_max'] = max(rollout_advantages) if rollout_advantages else 0.0 - log_dict['adv_min'] = min(rollout_advantages) if rollout_advantages else 0.0 - swanlab.log(_coerce_for_swanlab(log_dict), step=batch_step) - metrics.reset() - logger.info(f'[Step {batch_step}/{total_steps}] [SKIPPED] {log_dict}') - optim_step += optim_steps_per_batch - continue - - metrics.accumulate( - completion_lengths=per_rollout_completion_length, - rewards={'total': total_rewards, 'f1': f1_rewards, 'cot': cot_rewards}) - - all_input_data: List[Any] = list(all_trajectories) - advantages: List[float] = list(rollout_advantages) - - total_completions = len(all_input_data) - aligned_completions = (total_completions // MODEL_GPUS) * MODEL_GPUS - if aligned_completions < total_completions: - logger.info( - '[dp-align] dropping %d tail sample(s): total=%d -> aligned=%d (dp=%d)', - total_completions - aligned_completions, - total_completions, aligned_completions, MODEL_GPUS) - for mb_start in range(0, aligned_completions, MINI_BATCH_SIZE): - mb_end = min(mb_start + MINI_BATCH_SIZE, aligned_completions) - mb_inputs = all_input_data[mb_start:mb_end] - # Reference log-probs for KL: same policy with LoRA disabled (= base model). - ref_logps = None - if KL_BETA > 0.0: - ref_outputs = model.forward_only(inputs=mb_inputs, disable_lora=True) - ref_logps = ref_outputs.get('logps') if isinstance(ref_outputs, dict) else getattr(ref_outputs, 'logps', None) - model.forward_backward( - inputs=mb_inputs, - old_logps=all_old_logps[mb_start:mb_end], - advantages=advantages[mb_start:mb_end], - ref_logps=ref_logps, - positive_mask=all_f1_labels[mb_start:mb_end], - micro_batch_size=MICRO_BATCH_SIZE) - model.clip_grad_and_step() - optim_step += 1 - if optim_step >= total_steps: - break - if optim_step % SAVE_STEPS == 0: - model.save(f'hotpotqa-grpo-baseline-checkpoint-{optim_step}') - - log_dict = metrics.calculate() - log_dict.update(model.calculate_metric(is_training=True)) - log_dict.update(_compute_rollout_diagnostics( - all_trajectories, n_turns_per_rollout, per_rollout_completion_length, - f1_rewards=f1_rewards, old_logps=all_old_logps)) - log_dict['pos_neg_adv_rate'] = pos_with_neg_adv / n_pos if n_pos else 0.0 - log_dict['neg_pos_adv_rate'] = neg_with_pos_adv / n_neg if n_neg else 0.0 - log_dict['adv_max'] = max(rollout_advantages) if rollout_advantages else 0.0 - log_dict['adv_min'] = min(rollout_advantages) if rollout_advantages else 0.0 - # Pop high-KL token records before swanlab.log: list-of-dict won't render as a chart. - _hk = log_dict.pop('_high_kl_records', None) - if _hk: - _tok = rollout_template.tokenizer - for r in _hk: - gsi = r.get('gsi') - tid = all_trajectories[gsi].get('id') if gsi is not None and 0 <= gsi < len(all_trajectories) else None - try: - tok_text = _tok.decode([r['token_id']]) - except Exception: - tok_text = None - logger.info( - '[high-kl] step=%d gsi=%s tid=%s pos=%s tok=%r kl=%.4f r=%.4f lp_new=%.4f lp_old=%.4f', - batch_step, gsi, tid, r.get('pos'), tok_text, - r.get('kl'), r.get('ratio'), r.get('logp_new'), r.get('logp_old')) - swanlab.log(_coerce_for_swanlab(log_dict), step=batch_step) - metrics.reset() - logger.info(f'[Step {batch_step}/{total_steps}] {log_dict}') - - logger.info(f'Training completed. optim_steps={optim_step}') - model.save('hotpotqa-grpo-baseline-final') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/legacy/grpo_condensed.py b/cookbook/exp/legacy/grpo_condensed.py deleted file mode 100644 index 83eb49ac7..000000000 --- a/cookbook/exp/legacy/grpo_condensed.py +++ /dev/null @@ -1,955 +0,0 @@ -import copy -import math -import os -import re -from typing import Any, Dict, List, Optional - -import torch -import swanlab -from peft import LoraConfig - -import twinkle -from twinkle import DeviceMesh, DeviceGroup, get_logger -from twinkle.advantage import GRPOAdvantage -from twinkle.checkpoint_engine import CheckpointEngineManager -from twinkle.data_format import Message, SamplingParams, Trajectory -from twinkle.dataloader import DataLoader -from twinkle.dataset import Dataset, DatasetMeta -from twinkle.metric import CompletionRewardMetric -from twinkle.model import TransformersModel -from twinkle.preprocessor.base import Preprocessor -from twinkle.processor import InputProcessor -from twinkle.sampler import vLLMSampler -from twinkle.template import Qwen3_5Template -from twinkle_agentic.chunker.native import NativeChunker -from twinkle_agentic.condenser import ModelCondenser -from twinkle_agentic.reward import F1Reward, CoTReward, ToolExploreReward -from twinkle_agentic.rollout.multi_turn_condense import MultiTurnCondenseRollout -from twinkle_agentic.tools.tool_manager import ToolManager - -logger = get_logger() - -MODEL_ID = os.environ.get('MODEL_ID', 'ms://Qwen/Qwen3.5-4B') -USE_MEGATRON = bool(int(os.environ.get('USE_MEGATRON', '0'))) - -MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 4)) -SAMPLER_GPUS = int(os.environ.get('SAMPLER_GPUS', 4)) -NUM_GPUS = MODEL_GPUS + SAMPLER_GPUS - -NUM_GENERATIONS = int(os.environ.get('NUM_GENERATIONS', 8)) -MAX_NEW_TOKENS = int(os.environ.get('MAX_NEW_TOKENS', 4096)) -LEARNING_RATE = float(os.environ.get('LR', 1e-5)) -NUM_EPOCHS = int(os.environ.get('NUM_EPOCHS', 1)) -MAX_STEPS = int(os.environ.get('MAX_STEPS', 0)) -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)) -GRADIENT_ACCUMULATION_STEPS = int(os.environ.get('GRADIENT_ACCUMULATION_STEPS', 1)) -ADAPTER_NAME = 'default' -SAVE_STEPS = int(os.environ.get('SAVE_STEPS', 1000)) -LORA_RANK = int(os.environ.get('LORA_RANK', 16)) - -MAX_TURNS = int(os.environ.get('MAX_TURNS', 4)) -MAX_TRAJECTORY_TOKENS = int(os.environ.get('MAX_TRAJECTORY_TOKENS', 8192)) -CHUNK_SIZE = int(os.environ.get('CHUNK_SIZE', 1024)) - -HOTPOTQA_NUM_PROC = int(os.environ.get('HOTPOTQA_NUM_PROC', 16)) -HOTPOTQA_MAX_LENGTH = int(os.environ.get('HOTPOTQA_MAX_LENGTH', 64000)) - -F1_REWARD_WEIGHT = float(os.environ.get('F1_REWARD_WEIGHT', 1.0)) -COT_REWARD_WEIGHT = float(os.environ.get('COT_REWARD_WEIGHT', 0)) -TOOL_BONUS_WEIGHT = float(os.environ.get('TOOL_BONUS_WEIGHT', 0.0)) -TOOL_BONUS_F1_THRESHOLD = float( - os.environ.get('TOOL_BONUS_F1_THRESHOLD', 0.5)) - -# KL penalty coefficient; 0 disables KL (and skips the ref forward pass entirely). -# CISPO is token-level and DOES support per-token KL — small positive value (e.g. 0.005) recommended as anchor. -KL_BETA = float(os.environ.get('KL_BETA', 0.01)) - -# Entropy bonus coefficient; 0 disables the entropy compute path entirely. -# Typical GRPO values: 0.001–0.01. Loss is: L = L_PPO + beta*KL - entropy_coef*H. -ENTROPY_COEF = float(os.environ.get('ENTROPY_COEF', 0.0)) - -# Per-token oracle bonus coefficient; 0 disables. Typical: 0.05–0.2. -# Loss becomes: L = L_PPO + beta*KL - entropy_coef*H - token_bonus_coef*(oracle_logps - rollout_logps) -ORACLE_BONUS_COEF = float(os.environ.get('ORACLE_BONUS_COEF', 0.0)) - -# CISPO token-level IS clamp thresholds (MiniMax CISPO defaults: 0.2 / 0.28 asymmetric). -CISPO_EPS_LOW = float(os.environ.get('CISPO_EPS_LOW', 0.2)) -CISPO_EPS_HIGH = float(os.environ.get('CISPO_EPS_HIGH', 0.2)) - -# High-KL token capture: top-K per microbatch dumped into log_dict['_high_kl_records']. 0 = disabled. -HIGH_KL_TOPK = int(os.environ.get('HIGH_KL_TOPK', 0)) - -INIT_LORA_PATH = os.environ.get('INIT_LORA_PATH', 'output/condensed_sft_ddp/last-checkpoint') -DATASET_PATH = os.environ.get( - 'DATASET_PATH', - os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), - 'hotpotqa_fullwiki_reannotated_12k.jsonl')) -F1_BINARY_THRESHOLD = float(os.environ.get('F1_BINARY_THRESHOLD', 0.5)) - -_ROLLOUT_TRACE_DIR = os.environ.get('ROLLOUT_TRACE_DIR', 'rollout_trace') -ORACLE_HINT = bool(int(os.environ.get('ORACLE_HINT', '0'))) - - -# [EXP-ORACLE] staged hint injection — appended to the Question line so skip_pattern keeps it uncompressed. -def _oracle_hint_stage(step: int, total_steps: int) -> int: - """0 = explicit titles, 1 = vague count, 2 = no hint.""" - return 0 - # if total_steps <= 0: - # return 0 - # third = max(1, total_steps // 3) - # if step < third: - # return 0 - # if step < 2 * third: - # return 1 - # return 2 - - - -def _make_oracle_hint_callback(total_steps: int): - """Return a post_compress_callback that injects oracle hints with actual block IDs. - - Called by MultiTurnCondenseRollout after compression + metadata merge, so - ``compressed['user_data']`` carries sf_titles and ``chunks`` carries the - condensed/raw status of each passage. - - Stages (determined by global_step / total_steps): - 0 — explicit block IDs for supporting-fact passages - 1 — block count only (no IDs) - 2 — no hint - """ - _q_split = re.compile(r'(Question:\s*.+?)(\n\nContext:)', re.DOTALL) - - def _callback(compressed, chunks, **kwargs): - step = kwargs.get('global_step', 0) - stage = _oracle_hint_stage(step, total_steps) - if stage == 2: - return compressed - - user_data = compressed.get('user_data') or [] - sf_titles = [v for k, v in user_data if k == 'sf_title' and v] - if not sf_titles: - return compressed - sf_set = set(sf_titles) - - # Map sf_titles → block IDs by walking condensed chunks - block_id = 0 - sf_block_ids = [] - for c in chunks.chunks: - if c.get('type') != 'text': - continue - content = c.get('content') - if not isinstance(content, str) or not content: - continue - if c.get('role') == 'tool': - continue - raw = c.get('raw') - if not (isinstance(raw, dict) and raw.get('condensed')): - continue - block_id += 1 - original = raw.get('original', '') - if isinstance(original, str): - for title in sf_set: - if original.startswith(f'{title}: ') or original.startswith(f'{title}:'): - sf_block_ids.append(block_id) - break - - if stage == 0: - if sf_block_ids: - ids_str = ', '.join(str(b) for b in sf_block_ids) - hint = (f'\n[Oracle Hint] Block {ids_str} contain(s) the supporting facts. ' - 'Call `extract_condensed` to expand them if you need more detail information.') - else: - n = len(sf_set) - word = {1: 'One', 2: 'Two', 3: 'Three'}.get(n, str(n)) - hint = (f'\n[Oracle Hint] {word} short passage(s) contain the supporting facts; ' - 'they are uncompressed — read them directly.') - else: - hint = (f'\n[Oracle Hint] Some compressed block(s) contain the supporting facts; ' - 'call `extract_condensed` to expand them if you need more detail information.') - - for m in (compressed.get('messages') or []): - if m.get('role') != 'user': - continue - c = m.get('content') - if isinstance(c, str): - m['content'] = _q_split.sub( - lambda g: g.group(1) + hint + g.group(2), c, count=1) - elif isinstance(c, list): - for part in c: - if isinstance(part, dict) and part.get('type') == 'text': - part['text'] = _q_split.sub( - lambda g: g.group(1) + hint + g.group(2), - part.get('text') or '', count=1) - break - break - return compressed - - return _callback - -SYSTEM_PROMPT = """You are a careful multi-hop QA assistant. - -## Context Format (Mixed) -The context you receive is a **mix of two forms**: - -1. **Compressed blocks** — long passages wrapped in `...`, \ - displayed as a Markdown digest in **telegraphic style** (no \ - articles / "is" / "are"; colons and commas mean "is" / "has") \ - with two sections: - - **Summary**: overview plus facts strongly related to the question, stated explicitly. - - **More**: a collapsed INDEX of category keywords hinting at extra details hidden in the full text (call `extract_condensed` to see them). - Reading example: `India: 7th largest by area. Borders: Pakistan, \ - China.` means "India is the 7th largest country by area and \ - shares borders with Pakistan and China." -2. **Raw passages** — short passages shown inline as plain text (`Title: \ - body`) **without** any `` wrapping. These are already the full \ - text; nothing is hidden. - -Only the ``-wrapped blocks are compressed and can be expanded. \ -Block ids `N` are 1-based and assigned in the order compressed blocks \ -appear in the context, so they are always contiguous (``, \ -``, ``, ...). Raw passages have no block id and cannot \ -be extracted — they are already complete. - -## Workflow - -### Phase 1 — Scan and Decide -Step 1: Read each compressed block's Summary, and read raw \ -passages directly, to get an overview. -Step 2: For compressed blocks, check the More keywords to judge whether \ -hidden details are needed. -Step 3: Decide which compressed blocks to expand, then call \ -`extract_condensed` with their block ids. Raw passages need no extraction. - -### Phase 2 — Reason and Answer -After the tool returns the full text, continue stepping through the evidence: -Step N: From block X (or the raw passage titled "..."), I learn that [fact A]. -Step N+1: From block Y, I need to call `extract_condensed` to get more information, because this block is related to... -Step N+2: Combining these, the answer is ... -\\boxed{answer} - -You may call `extract_condensed` several times to expand more blocks if the information is not enough, only answer the question if you are sure about the facts. -The `blocks` parameter accepts **exactly one integer** per call (e.g. `3`); lists are rejected. Expand additional blocks by issuing separate `extract_condensed` calls, one per block. Only pass ids that actually appear as `` in the context, and do **not** request the same block twice — its text is already in the conversation after the first expansion. - -## Tool Call Format - - - -3 - - - - -## Output Format -End your final response with \\boxed{answer}, e.g. \\boxed{Delhi}. -Keep the boxed text short: a name, entity, date, or "yes"/"no". -Answers not inside \\boxed{} will not be scored.""" - - -_F1_REWARD: Optional[F1Reward] = F1Reward() -_COT_REWARD: Optional[CoTReward] = CoTReward() -_TOOL_EXPLORE_REWARD: Optional[ToolExploreReward] = ToolExploreReward( - f1_threshold=TOOL_BONUS_F1_THRESHOLD) - - -def compute_rewards(trajectories: List[Dict[str, Any]]): - f1_raw = _F1_REWARD(trajectories) - f1 = [1.0 if v >= F1_BINARY_THRESHOLD else 0.0 for v in f1_raw] if F1_BINARY_THRESHOLD > 0 else f1_raw - cot = _COT_REWARD(trajectories) - tool_explore = _TOOL_EXPLORE_REWARD(trajectories) - total = [ - F1_REWARD_WEIGHT * a + COT_REWARD_WEIGHT * c + TOOL_BONUS_WEIGHT * te - for a, c, te in zip(f1, cot, tool_explore) - ] - return total, f1, cot, tool_explore - - -class HotpotQAProcessor(Preprocessor): - def __init__(self, system: str = SYSTEM_PROMPT): - self.system = system - - def __call__(self, rows: Dict[str, List[Any]]) -> Dict[str, List[Any]]: - rows = self.map_col_to_row(rows) - rows = [self.preprocess(row) for row in rows] - rows = [r for r in rows if r is not None] - rows = self.map_row_to_col(rows) - return rows - - @staticmethod - def _format_context(context: Dict[str, Any]) -> str: - titles = context.get('title', []) or [] - sentences = context.get('sentences', []) or [] - lines = [] - for title, sents in zip(titles, sentences): - if isinstance(sents, list): - body = ' '.join(s.strip() for s in sents if s and s.strip()) - else: - body = str(sents).strip() - lines.append(f'{title}: {body}') - return '\n\n'.join(lines) - - def preprocess(self, row: Dict[str, Any]) -> Optional[Trajectory]: - if (row.get('verdict') or '').strip().lower() == 'drop': - return None - question = row.get('question_fixed') or row['question'] - answers = row.get('answers') - if isinstance(answers, list) and answers: - gold = [str(a).strip() for a in answers if str(a).strip()] - else: - gold = [s for s in [(row.get('answer', '') or '').strip()] if s] - context_block = self._format_context(row.get('context', {}) or {}) - user_msg = f'Question: {question}\n\nContext:\n\n{context_block}' - messages = [ - Message(role='system', content=self.system), - Message(role='user', content=user_msg), - ] - # [EXP-ORACLE] carry supporting_facts titles via user_data; rollout injects post-compression block hint - sf = row.get('supporting_facts') or {} - sf_titles = sf.get('title') or [] - sf_unique = list(dict.fromkeys(t for t in sf_titles if t)) - user_data = [('ground_truth', g) for g in gold] + [('sf_title', t) for t in sf_unique] - return Trajectory(messages=messages, user_data=user_data) - - -def create_hotpotqa_dataset() -> Dataset: - dataset = Dataset() - dataset.add_dataset(DatasetMeta(DATASET_PATH)) - logger.info('[dataset] loaded %s: %d rows', DATASET_PATH, len(dataset)) - - dataset.set_template( - 'Qwen3_5Template', model_id=MODEL_ID, max_length=HOTPOTQA_MAX_LENGTH, - truncation_strategy='delete', enable_thinking=False) - _HOTPOTQA_COLS = ['id', 'question', 'question_fixed', 'answers', - 'original_answer', 'type', 'level', 'verdict', - 'reasoning', 'supporting_facts', 'context'] - dataset.map(HotpotQAProcessor(system=SYSTEM_PROMPT), remove_columns=_HOTPOTQA_COLS) - return dataset - - -# Matches a LaTeX ``\boxed{...}`` final-answer marker — used to flag -# rollouts that never committed an answer. Brace-balanced is overkill for -# a logging heuristic; a non-greedy ``[^}]*`` is good enough. -_BOXED_RE = re.compile(r'\\boxed\{[^}]*\}') - -# Pulls the leading number out of pre-formatted metric strings such as -# ``'0.03 iters/s'`` / ``'1.000000e-05'`` / ``'30 seconds'`` emitted by -# ``TrainMetric`` and ``GRPOMetric``. We use this in ``_coerce_for_swanlab`` -# so swanlab can build line charts instead of dropping those keys with a -# ``failed to create chart for key '...': invalid value type`` warning. -_LEADING_NUMBER_RE = re.compile(r'[-+]?\d*\.?\d+(?:[eE][-+]?\d+)?') - - -def _coerce_for_swanlab(log_dict: Dict[str, Any]) -> Dict[str, Any]: - """Cast string-valued metrics to float for swanlab line charts. - - ``TrainMetric.calculate()`` and ``GRPOMetric.calculate()`` return - pre-formatted strings (``'0.03 iters/s'``, ``'1.000000e-05'``, - ``'30 seconds'``, ``'0.8321'``). swanlab cannot build a line chart - from a string value and emits one warning per key per step. We extract - the leading number where possible; keys whose value can't be parsed - as a scalar are left as-is so they still show up in the text log. - """ - coerced: Dict[str, Any] = {} - for k, v in log_dict.items(): - if isinstance(v, bool) or isinstance(v, (int, float)): - coerced[k] = v - continue - if isinstance(v, str): - m = _LEADING_NUMBER_RE.search(v) - if m: - try: - coerced[k] = float(m.group()) - continue - except ValueError: - pass - coerced[k] = v - return coerced - - -def _last_assistant_text(trajectory: Dict[str, Any]) -> Optional[str]: - """Return the text of the last ``assistant`` message, or ``None``. - - ``content`` can be ``str`` | ``None`` | ``dict`` (single multimodal - part) | ``list[dict]`` (multiple parts). The downstream caller feeds - this into ``_BOXED_RE.search(...)``, so we collapse the visible text - into a single string and ignore non-text parts (images etc.). - """ - for m in reversed(trajectory.get('messages', [])): - if m.get('role') != 'assistant': - continue - c = m.get('content') - if c is None: - return None - if isinstance(c, str): - return c - if isinstance(c, dict): - return c.get('text') if c.get('type') == 'text' else None - if isinstance(c, list): - parts = [p.get('text') or '' for p in c - if isinstance(p, dict) and p.get('type') == 'text'] - return '\n'.join(parts) if parts else None - return str(c) - return None - - -def _compute_rollout_diagnostics( - trajectories: List[Dict[str, Any]], - n_turns_per_rollout: List[int], - per_rollout_completion_length: List[int], - f1_rewards: Optional[List[float]] = None, - old_logps: Optional[List[List[float]]] = None, -) -> Dict[str, float]: - """Aggregate rollout diagnostics for swanlab logging. - - All inputs are already flat: - * ``trajectories[i]`` is the merged trajectory dict returned by - :class:`MultiTurnCondenseRollout` (contains ``messages``, - ``input_ids``, ``labels``, ``turns`` at top level). - * ``n_turns_per_rollout[i] == trajectories[i]['turns']``. - * ``per_rollout_completion_length[i]`` == number of trainable - tokens in the trajectory (labels != -100). - """ - out: Dict[str, float] = {} - if n_turns_per_rollout: - out['avg_turns'] = sum(n_turns_per_rollout) / len(n_turns_per_rollout) - - # ``non_trainable_tokens`` is the longest non-trainable prefix across - # the batch: ``len(input_ids) - sum(1 for l in labels if l != -100)``. - # Tracks how much the condensed context + system prompt is eating the - # context budget (it does NOT equal the first-turn prompt length - # because multi-turn runs also contribute non-trainable tokens from - # the ``tool`` observations between assistant turns). - _max_non_trainable = 0 - for t, comp_len in zip(trajectories, per_rollout_completion_length): - ids = t.get('input_ids') or [] - non_trainable = max(0, len(ids) - int(comp_len or 0)) - if non_trainable > _max_non_trainable: - _max_non_trainable = non_trainable - out['non_trainable_tokens'] = _max_non_trainable - - if trajectories: - tool_counts = [ - sum(len(m.get('tool_calls') or []) - for m in t.get('messages', []) if m.get('role') == 'assistant') - for t in trajectories] - out['avg_tool_calls'] = sum(tool_counts) / len(tool_counts) - out['tool_use_rate'] = sum(1 for c in tool_counts if c > 0) / len(tool_counts) - n_no_boxed = sum( - 0 if _BOXED_RE.search(_last_assistant_text(t) or '') else 1 - for t in trajectories) - out['no_boxed_rate'] = n_no_boxed / len(trajectories) - def _content_chars(c: Any) -> int: - if not c: - return 0 - if isinstance(c, str): - return len(c) - if isinstance(c, dict): - if c.get('type') == 'text': - return len(c.get('text') or '') - return 0 - if isinstance(c, list): - total = 0 - for part in c: - if isinstance(part, dict) and part.get('type') == 'text': - total += len(part.get('text') or '') - elif isinstance(part, str): - total += len(part) - return total - # Unknown shape -- fall back to ``str()`` length rather than - # crashing, so a template quirk never breaks metric logging. - return len(str(c)) - - msg_chars_total, prompt_chars, asst_chars = [], [], [] - for t in trajectories: - total_i = prompt_i = asst_i = 0 - for m in (t.get('messages') or []): - role = m.get('role') - if role == 'system': - continue - n = _content_chars(m.get('content')) - total_i += n - if role in ('user', 'tool'): - prompt_i += n - elif role == 'assistant': - asst_i += n - msg_chars_total.append(total_i) - prompt_chars.append(prompt_i) - asst_chars.append(asst_i) - out['avg_chars_total_no_sys'] = sum(msg_chars_total) / len(msg_chars_total) - out['avg_chars_prompt_no_sys'] = sum(prompt_chars) / len(prompt_chars) - out['avg_chars_assistant'] = sum(asst_chars) / len(asst_chars) - - if f1_rewards is not None and old_logps is not None and f1_rewards: - per_traj_mean = [ - (sum(lp) / len(lp)) if lp else 0.0 for lp in old_logps] - pos_logp = [m for m, f1 in zip(per_traj_mean, f1_rewards) if f1 > 0] - zero_logp = [m for m, f1 in zip(per_traj_mean, f1_rewards) if f1 <= 0] - out['f1_correct_rate'] = len(pos_logp) / len(f1_rewards) - out['f1_zero_rate'] = len(zero_logp) / len(f1_rewards) - out['mean_old_logp_f1_pos'] = (sum(pos_logp) / len(pos_logp)) if pos_logp else 0.0 - out['mean_old_logp_f1_zero'] = (sum(zero_logp) / len(zero_logp)) if zero_logp else 0.0 - out['policy_confidence_f1_pos'] = math.exp(out['mean_old_logp_f1_pos']) - out['policy_confidence_f1_zero'] = math.exp(out['mean_old_logp_f1_zero']) - return out - - -def _build_oracle_inputs( - mb_inputs: List[Dict[str, Any]], - f1_labels: List[bool], - template, -) -> Optional[List[Dict[str, Any]]]: - """Build oracle-context inputs at the TOKEN level for per-token bonus computation. - - The approach: - 1. Find ``first_trainable`` from labels (first position != -100). - Due to NTP shift, input_ids[first_trainable] is the last prefix token (e.g. \\n - after ``assistant``) and labels[first_trainable] is the first response token target. - 2. Construct oracle messages: [system, user_with_oracle_suffix]. - 3. Encode with template (add_generation_prompt=True) → oracle_prefix_ids ending with - the same assistant header token. - 4. Concatenate: oracle_prefix_ids + input_ids[first_trainable+1:] (response tokens). - 5. Labels: [-100]*(len(oracle_prefix)-1) + labels[first_trainable:] so the last prefix - position predicts the first response token. - - For F1=0 samples: copied unchanged (bonus zeroed by _compute_token_bonus). - """ - _q_line_re = re.compile(r'Question:\s*(.+?)(?:\n|$)', re.DOTALL) - oracle_inputs = [] - any_modified = False - - for inp, is_pos in zip(mb_inputs, f1_labels): - if not is_pos: - oracle_inputs.append(inp) - continue - - user_data = inp.get('user_data') or [] - sf_titles = [v for k, v in user_data if k == 'sf_title' and v] - gts = [v for k, v in user_data if k == 'ground_truth' and v] - if not sf_titles and not gts: - oracle_inputs.append(inp) - continue - - labels = inp.get('labels') or [] - input_ids = inp.get('input_ids') or [] - if not labels or not input_ids: - oracle_inputs.append(inp) - continue - - # 1. Find first trainable position - first_trainable = None - for i, l in enumerate(labels): - if l != -100: - first_trainable = i - break - - assert first_trainable is not None - - # 2. Extract question from first user message - question = None - msgs = inp.get('messages') or [] - for m in msgs: - if m.get('role') != 'user': - continue - c = m.get('content') - text = c if isinstance(c, str) else ( - next((p.get('text') for p in c if isinstance(p, dict) and p.get('type') == 'text'), '') - if isinstance(c, list) else '') - q_match = _q_line_re.match(text or '') - if q_match: - question = q_match.group(1).strip() - break - - if not question: - oracle_inputs.append(inp) - continue - - # 3. Build oracle user message (concise: question + oracle hints only) - hint_parts = [] - if sf_titles: - hint_parts.append('Supporting passages: ' + ', '.join(f'"{t}"' for t in sf_titles)) - if gts: - hint_parts.append('Answer: ' + '; '.join(gts)) - hint_parts.append('You must call `extract_condensed` to read the right original passage from the condensed block with thinking steps, and give the final correct answer') - oracle_suffix = '\n[Oracle Context] ' + '. '.join(hint_parts) + '.' - oracle_user_content = f'Question: {question}{oracle_suffix}' - - oracle_msgs = [ - Message(role='system', content=SYSTEM_PROMPT), - Message(role='user', content=oracle_user_content), - ] - - # 4. Encode oracle prefix (ends with <|im_start|>assistant\n) - oracle_feature = template.encode( - Trajectory(messages=oracle_msgs), add_generation_prompt=True) - oracle_prefix_ids = list(oracle_feature['input_ids']) - - # 5. Splice: oracle_prefix + response_tokens - response_tokens = list(input_ids[first_trainable + 1:]) - response_labels = list(labels[first_trainable:]) - - oracle_input_ids = oracle_prefix_ids + response_tokens - # Last position of oracle prefix predicts first response token - oracle_labels = [-100] * (len(oracle_prefix_ids) - 1) + response_labels - - assert len(oracle_input_ids) == len(oracle_labels) - seq_len = len(oracle_input_ids) - # Start from original keys to keep collator-compatible shape - oi = dict(inp) - oi['input_ids'] = oracle_input_ids - oi['labels'] = oracle_labels - oi['attention_mask'] = [1] * seq_len - oi['messages'] = None - oi['length'] = seq_len - # Replicate mrope position_ids shape from original input - orig_pos = inp.get('position_ids') - if isinstance(orig_pos, torch.Tensor) and orig_pos.dim() == 3: - n_dims = orig_pos.shape[0] - pos_range = torch.arange(seq_len).unsqueeze(0).unsqueeze(0) - oi['position_ids'] = pos_range.expand(n_dims, 1, seq_len) - else: - oi['position_ids'] = list(range(seq_len)) - if 'mm_token_type_ids' in inp: - oi['mm_token_type_ids'] = torch.zeros(1, seq_len) - oracle_inputs.append(oi) - any_modified = True - - return oracle_inputs if any_modified else None - - -def _compute_token_bonus( - oracle_logps: Any, - old_logps: List[List[float]], - f1_labels: List[bool], - oracle_inputs: List[Dict[str, Any]], -) -> List[List[float]]: - """Compute per-token bonus = oracle_logps - rollout_logps, zeroed for F1=0 samples. - - oracle_logps is full-sequence form [batch, padded_seq] from forward_only + collector. - We extract valid positions using oracle_inputs[i]['labels'] mask to get response-only - logps aligned 1:1 with old_logps. - """ - import torch - - if isinstance(oracle_logps, torch.Tensor): - oracle_logps = oracle_logps.float().cpu() - - bonus = [] - for i, (is_pos, old_lp) in enumerate(zip(f1_labels, old_logps)): - if not is_pos or not old_lp: - bonus.append([0.0] * len(old_lp) if old_lp else []) - continue - - n = len(old_lp) - oracle_labels = oracle_inputs[i].get('labels') or [] - - # Build mask from oracle labels to extract valid (trainable) positions - if isinstance(oracle_logps, torch.Tensor): - orc_row = oracle_logps[i] - mask = torch.tensor([l != -100 for l in oracle_labels], dtype=torch.bool) - seq_len = min(len(mask), orc_row.numel()) - orc_valid = orc_row[:seq_len][mask[:seq_len]].tolist() - else: - orc_row = oracle_logps[i] if i < len(oracle_logps) else [] - if isinstance(orc_row, torch.Tensor): - orc_row = orc_row.float().cpu().tolist() - elif not isinstance(orc_row, (list, tuple)): - orc_row = [] - orc_valid = [v for v, l in zip(orc_row, oracle_labels) if l != -100] - - assert len(orc_valid) == n - bonus.append([o - r for o, r in zip(orc_valid, old_lp)]) - return bonus - - -def main(): - swanlab.init(project='twinkle') - - device_groups = [ - DeviceGroup(name='model', ranks=list(range(MODEL_GPUS)), device_type='GPU'), - DeviceGroup(name='sampler', ranks=list(range(MODEL_GPUS, 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, dp_size=SAMPLER_GPUS) - twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, - groups=device_groups, lazy_collect=False) - - logger.info('Building HotpotQA dataset') - _prebuilt_dataset = create_hotpotqa_dataset() - logger.info('Dataset ready: %d rows', len(_prebuilt_dataset)) - - GLOBAL_BATCH_SIZE = BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS - batches_per_epoch = max(1, len(_prebuilt_dataset) // GLOBAL_BATCH_SIZE) - optim_steps_per_batch = max(1, (GLOBAL_BATCH_SIZE * NUM_GENERATIONS - + MINI_BATCH_SIZE - 1) // MINI_BATCH_SIZE) - steps_per_epoch = batches_per_epoch * optim_steps_per_batch - derived_total_steps = NUM_EPOCHS * steps_per_epoch - total_steps = min(MAX_STEPS, derived_total_steps) if MAX_STEPS > 0 else derived_total_steps - logger.info('Training horizon: %d steps (%d epochs × %d batches × %d steps/batch)', - total_steps, NUM_EPOCHS, batches_per_epoch, optim_steps_per_batch) - - lora_config = LoraConfig( - target_modules='all-linear', r=LORA_RANK, - lora_alpha=LORA_RANK * 2, lora_dropout=0.05) - - if USE_MEGATRON: - from twinkle.model.megatron import MegatronModel - model = MegatronModel( - model_id=MODEL_ID, device_mesh=model_mesh, remote_group='model', - mixed_precision='bf16', variable_seq_lengths=True) - else: - model = TransformersModel( - model_id=MODEL_ID, device_mesh=model_mesh, remote_group='model') - - model.add_adapter_to_model(ADAPTER_NAME, lora_config, - gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) - if INIT_LORA_PATH: - model.load(INIT_LORA_PATH, adapter_name=ADAPTER_NAME) - logger.info('Loaded cold-start LoRA from %s', INIT_LORA_PATH) - if USE_MEGATRON: - model.set_optimizer('default', lr=LEARNING_RATE) - model.set_lr_scheduler('default', lr_decay_steps=total_steps, max_lr=LEARNING_RATE) - else: - model.set_optimizer('AdamW', lr=LEARNING_RATE) - model.set_lr_scheduler('CosineAnnealingLR', T_max=total_steps, eta_min=0) - - model.set_loss('GRPOLoss', epsilon=CISPO_EPS_LOW, epsilon_high=CISPO_EPS_HIGH, - beta=KL_BETA, entropy_coef=ENTROPY_COEF, token_bonus_coef=ORACLE_BONUS_COEF) - model.set_processor(InputProcessor, padding_free=True) - model.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=False, max_length=HOTPOTQA_MAX_LENGTH) - - model.add_metric('GRPOMetric', is_training=True, - epsilon=CISPO_EPS_LOW, epsilon_high=CISPO_EPS_HIGH, - top_k_kl=HIGH_KL_TOPK) - - sampler = vLLMSampler( - model_id=MODEL_ID, - engine_args={ - 'gpu_memory_utilization': 0.8, 'max_model_len': 32768, - 'max_lora_rank': 32, 'enable_lora': True, - 'enable_tower_connector_lora': True, - 'max_loras': 5 - }, - device_mesh=sampler_mesh, remote_group='sampler') - sampler.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=False, max_length=HOTPOTQA_MAX_LENGTH) - rollout_template = Qwen3_5Template( - MODEL_ID, max_length=HOTPOTQA_MAX_LENGTH, enable_thinking=False) - - ckpt_manager = CheckpointEngineManager(model=model, sampler=sampler) - chunker = NativeChunker( - chunk_size=CHUNK_SIZE, - passage_boundary_re=r'(?<=\n\n)', - ) - # ``\A`` anchor: prevents a ``Question:`` line inside a passage from being misread as the query. - _question_re = re.compile(r'\AQuestion:\s*(.+)') - - def _extract_question(chunk): - content = chunk.get('content') - if chunk.get('type') != 'text' or not isinstance(content, str): - return None - m = _question_re.search(content) - return m.group(1).strip() if m else None - - condenser = ModelCondenser( - sampler=sampler, - compression_ratio=2.0, - sampling_params=SamplingParams( - max_tokens=1024, num_samples=1, temperature=0.4, top_p=0.9), - min_chars=200, - template=rollout_template, - lora_path='ms://twinkle-kit/Qwen3.5-4B-Condenser', - skip_pattern=r'^Question:', - related_query=_extract_question, - ) - - dataloader = DataLoader( - dataset=lambda: _prebuilt_dataset, - batch_size=GLOBAL_BATCH_SIZE, min_batch_size=GLOBAL_BATCH_SIZE) - - 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, - stop=['']) - - def _trace_should_store(traj): - return _F1_REWARD([traj])[0] == 0.0 - - def _trace_is_success(traj): - return _F1_REWARD([traj])[0] > 0.0 - - rollout = MultiTurnCondenseRollout( - sampler=sampler, - template=rollout_template, - tool_manager=ToolManager(), - chunker=chunker, - condenser=condenser, - sampling_params=sampling_params, - max_turns=MAX_TURNS, - max_trajectory_tokens=MAX_TRAJECTORY_TOKENS, - trace_dir=_ROLLOUT_TRACE_DIR or None, - trace_callback=_trace_should_store, - success_callback=_trace_is_success, - post_compress_callback=( - _make_oracle_hint_callback(total_steps) if ORACLE_HINT else None), - ) - - optim_step = 0 - logger.info('Starting HotpotQA GRPO training (LLM condenser variant)') - - def _epoch_cycle(dl, n_epochs): - for ep in range(1, n_epochs + 1): - logger.info(f'=== Epoch {ep}/{n_epochs} (step={optim_step}/{total_steps}) ===') - for batch in dl: - yield batch - - for batch in _epoch_cycle(dataloader, NUM_EPOCHS): - if optim_step >= total_steps: - break - - # Single source of truth for the step shown in swanlab / logger / rollout-trace filename. - # Equals the number of optimizer updates already completed when this rollout was sampled. - batch_step = optim_step - - metrics.reset() - expand_prompts = [p for prompt in batch for p in [prompt] * NUM_GENERATIONS] - - ckpt_manager.sync_weights(merge_and_sync=False) - sampler.reset_prefix_cache() - - # Batched multi-turn rollout with chunk+condense pre-processing. - # Each returned trajectory is a flat dict containing ``messages``, - # ``input_ids``, ``labels``, ``attention_mask``, ``position_ids``, - # ``turns``, ``logprobs``, ``stop_reason``, ``truncated``. - all_trajectories: List[Dict[str, Any]] = rollout(expand_prompts, global_step=batch_step) - n_turns_per_rollout = [int(t.get('turns') or 0) for t in all_trajectories] - per_rollout_completion_length = [ - sum(1 for l in (t.get('labels') or []) if l != -100) - for t in all_trajectories] - - total_rewards, f1_rewards, cot_rewards, tool_explore_rewards = \ - compute_rewards(all_trajectories) - - rollout_advantages = advantage_fn( - total_rewards, num_generations=NUM_GENERATIONS, scale='group').tolist() - - all_f1_labels: List[bool] = [f > 0 for f in f1_rewards] - n_pos = sum(1 for p in all_f1_labels if p) - n_neg = sum(1 for p in all_f1_labels if not p) - pos_with_neg_adv = sum(1 for p, a in zip(all_f1_labels, rollout_advantages) if p and a < 0) - neg_with_pos_adv = sum(1 for p, a in zip(all_f1_labels, rollout_advantages) if not p and a > 0) - - # Skip homogeneous groups where gradient signal is meaningless - f1_pos_rate = n_pos / len(f1_rewards) if f1_rewards else 0.5 - if f1_pos_rate > 0.9 or f1_pos_rate < 0.1: - logger.info('[skip-homogeneous] f1_pos_rate=%.3f, skipping training update', f1_pos_rate) - metrics.accumulate( - completion_lengths=per_rollout_completion_length, - rewards={'total': total_rewards, 'f1': f1_rewards, - 'cot': cot_rewards, 'tool_explore': tool_explore_rewards}) - log_dict = metrics.calculate() - log_dict.update(_compute_rollout_diagnostics( - all_trajectories, n_turns_per_rollout, per_rollout_completion_length, - f1_rewards=f1_rewards, old_logps=[[lp[0][1] for lp in (t.get('logprobs') or [])] for t in all_trajectories])) - log_dict['skipped'] = True - log_dict['pos_neg_adv_rate'] = pos_with_neg_adv / n_pos if n_pos else 0.0 - log_dict['neg_pos_adv_rate'] = neg_with_pos_adv / n_neg if n_neg else 0.0 - log_dict['adv_max'] = max(rollout_advantages) if rollout_advantages else 0.0 - log_dict['adv_min'] = min(rollout_advantages) if rollout_advantages else 0.0 - swanlab.log(_coerce_for_swanlab(log_dict), step=batch_step) - metrics.reset() - logger.info(f'[Step {batch_step}/{total_steps}] [SKIPPED] {log_dict}') - optim_step += optim_steps_per_batch - continue - - metrics.accumulate( - completion_lengths=per_rollout_completion_length, - rewards={'total': total_rewards, 'f1': f1_rewards, - 'cot': cot_rewards, 'tool_explore': tool_explore_rewards}) - - all_input_data: List[Any] = [] - all_old_logps: List[List[float]] = [] - advantages: List[float] = [] - for t, adv in zip(all_trajectories, rollout_advantages): - all_input_data.append(t) - all_old_logps.append([lp[0][1] for lp in (t.get('logprobs') or [])]) - advantages.append(adv) - - total_completions = len(all_input_data) - aligned_completions = (total_completions // MODEL_GPUS) * MODEL_GPUS - if aligned_completions < total_completions: - logger.info( - '[dp-align] dropping %d tail sample(s): total=%d -> aligned=%d (dp=%d)', - total_completions - aligned_completions, - total_completions, aligned_completions, MODEL_GPUS) - for mb_start in range(0, aligned_completions, MINI_BATCH_SIZE): - mb_end = min(mb_start + MINI_BATCH_SIZE, aligned_completions) - mb_inputs = all_input_data[mb_start:mb_end] - # Reference log-probs for KL: same policy model with LoRA adapter disabled (= base model). - # Skipped when KL_BETA == 0 to save one extra forward per mini-batch. - ref_logps = None - if KL_BETA > 0.0: - ref_outputs = model.forward_only(inputs=mb_inputs, disable_lora=True) - ref_logps = ref_outputs.get('logps') if isinstance(ref_outputs, dict) else getattr(ref_outputs, 'logps', None) - # [EXP-ORACLE] per-token bonus: forward with oracle context, diff against rollout logps - mb_token_bonus = None - if ORACLE_BONUS_COEF > 0.0: - mb_oracle_inputs = _build_oracle_inputs( - mb_inputs, all_f1_labels[mb_start:mb_end], rollout_template) - if mb_oracle_inputs is not None: - oracle_outputs = model.forward_only(inputs=mb_oracle_inputs) - oracle_logps = oracle_outputs.get('logps') if isinstance(oracle_outputs, dict) else getattr(oracle_outputs, 'logps', None) - if oracle_logps is not None: - mb_token_bonus = _compute_token_bonus( - oracle_logps, all_old_logps[mb_start:mb_end], - all_f1_labels[mb_start:mb_end], mb_oracle_inputs) - model.forward_backward( - inputs=mb_inputs, - old_logps=all_old_logps[mb_start:mb_end], - advantages=advantages[mb_start:mb_end], - ref_logps=ref_logps, - token_bonus=mb_token_bonus, - positive_mask=all_f1_labels[mb_start:mb_end], - micro_batch_size=MICRO_BATCH_SIZE) - model.clip_grad_and_step() - optim_step += 1 - if optim_step >= total_steps: - break - if optim_step % SAVE_STEPS == 0: - model.save(f'hotpotqa-grpo-tools-llmcondense-checkpoint-{optim_step}') - - log_dict = metrics.calculate() - log_dict.update(model.calculate_metric(is_training=True)) - log_dict.update(_compute_rollout_diagnostics( - all_trajectories, n_turns_per_rollout, per_rollout_completion_length, - f1_rewards=f1_rewards, old_logps=all_old_logps)) - log_dict['pos_neg_adv_rate'] = pos_with_neg_adv / n_pos if n_pos else 0.0 - log_dict['neg_pos_adv_rate'] = neg_with_pos_adv / n_neg if n_neg else 0.0 - log_dict['adv_max'] = max(rollout_advantages) if rollout_advantages else 0.0 - log_dict['adv_min'] = min(rollout_advantages) if rollout_advantages else 0.0 - # Pop high-KL token records before swanlab.log: list-of-dict won't render as a chart. - _hk = log_dict.pop('_high_kl_records', None) - if _hk: - _tok = rollout_template.tokenizer - for r in _hk: - gsi = r.get('gsi') - tid = all_trajectories[gsi].get('id') if gsi is not None and 0 <= gsi < len(all_trajectories) else None - try: - tok_text = _tok.decode([r['token_id']]) - except Exception: - tok_text = None - logger.info( - '[high-kl] step=%d gsi=%s tid=%s pos=%s tok=%r kl=%.4f r=%.4f lp_new=%.4f lp_old=%.4f', - batch_step, gsi, tid, r.get('pos'), tok_text, - r.get('kl'), r.get('ratio'), r.get('logp_new'), r.get('logp_old')) - swanlab.log(_coerce_for_swanlab(log_dict), step=batch_step) - metrics.reset() - logger.info(f'[Step {batch_step}/{total_steps}] {log_dict}') - - logger.info(f'Training completed. optim_steps={optim_step}') - model.save('hotpotqa-grpo-tools-llmcondense-final') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/legacy/make_condensed_sft.py b/cookbook/exp/legacy/make_condensed_sft.py deleted file mode 100644 index 3b9855ac2..000000000 --- a/cookbook/exp/legacy/make_condensed_sft.py +++ /dev/null @@ -1,945 +0,0 @@ -"""Cold-start SFT dataset builder for the condensed multi-hop QA task. - -Pipeline per HotpotQA distractor row: - 1. Build the standard system + user-with-context trajectory using the - production ``SYSTEM_PROMPT`` and ``_format_context`` from - ``cookbook/rl/grpo_condensed.py`` so the offline data matches what - the policy sees at training/inference time. - 2. Run the production ``NativeChunker`` + ``ModelCondenser`` on the - row to produce ``...`` compressed text. - 3. **Validation pass** (super-LLM, ``enable_thinking=True``, no oracle, - no tools): judge whether the question / supporting_facts / GT are - well-formed against the raw passages; return strict JSON - ``{"verdict": "ok"|"fix"|"drop", ...}`` with fixed SF + GT when - applicable. ``drop`` skips the row. - 4. **Oracle rollout pass** via :class:`APIMultiTurnRollout` with a - trajectory-bound :class:`ExtractCondensed` tool. The oracle hint - (SF titles + GT) is injected into the system prompt **only for - the API call**; it is stripped before saving. The model emits - OpenAI-shape ``tool_calls`` for ``extract_condensed``, the rollout - dispatches them through :class:`ToolManager` and feeds back the - pre-compression passage text as a ``tool`` message, looping until - the model finalises with ``\\boxed{...}`` or hits ``MAX_TURNS``. - 5. Accept iff F1(boxed, used_gt) >= ``F1_ACCEPT_THRESHOLD``. On miss, - retry once with a higher temperature. - 6. Convert OpenAI-shape ``tool_calls`` into the textual - ``N`` - format consumed by the training chat template (mirrors - ``grpo_condensed.SYSTEM_PROMPT`` L232-239), restore the clean - system prompt, and emit one JSONL line. - -Run:: - - python cookbook/rl/make_condensed_sft.py \\ - --output hotpotqa_sft_coldstart.jsonl \\ - --model --api-key $KEY --base-url $URL \\ - --total 9000 --easy 1500 --medium 3000 --hard 4500 \\ - --concurrency 16 --seed 42 \\ - --condenser-model-id ms://Qwen/Qwen3.5-4B \\ - --condenser-lora ms://twinkle-kit/Qwen3.5-4B-Condenser -""" -from __future__ import annotations - -import argparse -import json -import os -import random -import re -import sys -import threading -from concurrent.futures import ThreadPoolExecutor -from typing import Any, Dict, List, Optional, Tuple - -from datasets import load_dataset - -from twinkle.data_format.sampling import SamplingParams -from twinkle.sampler import vLLMSampler -from twinkle.template import Qwen3_5Template -from twinkle_agentic.chunker.native import NativeChunker -from twinkle_agentic.condenser import ModelCondenser -from twinkle_agentic.data_format import Chunks -from twinkle_agentic.protocol.openai import OpenAI -from twinkle_agentic.reward.f1 import _extract_final_answer, _f1_score -from twinkle_agentic.rollout import APIMultiTurnRollout -from twinkle_agentic.tools.extract_condensed import ExtractCondensed -from twinkle_agentic.tools.tool_manager import ToolManager - - -# -------------------------------------------------------------------------- -# Constants mirrored from grpo_condensed.py so the SFT data matches the -# runtime contract byte-for-byte. Re-import would pull the whole training -# module; copying these few strings keeps the builder standalone. -# -------------------------------------------------------------------------- -SYSTEM_PROMPT = """You are a careful multi-hop QA assistant. - -## Context Format (Mixed) -The context you receive is a **mix of two forms**: - -1. **Compressed blocks** — long passages wrapped in `...`, \ -displayed as a Markdown digest in **telegraphic style** (no \ -articles / "is" / "are"; colons and commas mean "is" / "has") \ -with two sections: - - **Summary**: overview plus facts strongly related to the question, stated explicitly. - - **More**: a collapsed INDEX of category keywords hinting at extra details hidden in the full text (call `extract_condensed` to see them). - Reading example: `India: 7th largest by area. Borders: Pakistan, \ -China.` means "India is the 7th largest country by area and \ -shares borders with Pakistan and China." -2. **Raw passages** — short passages shown inline as plain text (`Title: \ -body`) **without** any `` wrapping. These are already the full \ -text; nothing is hidden. - -Only the ``-wrapped blocks are compressed and can be expanded. \ -Block ids `N` are 1-based and assigned in the order compressed blocks \ -appear in the context, so they are always contiguous (``, \ -``, ``, ...). Raw passages have no block id and cannot \ -be extracted — they are already complete. - -## Workflow - -### Phase 1 — Scan and Decide -Step 1: Read each compressed block's Summary, and read raw \ -passages directly, to get an overview. -Step 2: For compressed blocks, check the More keywords to judge whether \ -hidden details are needed. -Step 3: Decide which compressed blocks to expand, then call \ -`extract_condensed` with their block ids. Raw passages need no extraction. - -### Phase 2 — Reason and Answer -After the tool returns the full text, continue stepping through the evidence: -Step N: From block X (or the raw passage titled "..."), I learn that [fact A]. -Step N+1: From block Y, I need to call `extract_condensed` to get more information, because this block is related to... -Step N+2: Combining these, the answer is ... -\\boxed{answer} - -You may call `extract_condensed` several times to expand more blocks if the information is not enough, only answer the question if you are sure about the facts. -The `blocks` parameter accepts **exactly one integer** per call (e.g. `3`); lists are rejected. Expand additional blocks by issuing separate `extract_condensed` calls, one per block. Only pass ids that actually appear as `` in the context, and do **not** request the same block twice — its text is already in the conversation after the first expansion. - -## Tool Call Format - - - -3 - - - - -## Output Format -End your final response with \\boxed{answer}, e.g. \\boxed{Delhi}. -Keep the boxed text short: a name, entity, date, or "yes"/"no". -Answers not inside \\boxed{} will not be scored.""" - - -# Oracle suffix appended ONLY for API generation; stripped before save. -_ORACLE_HINT_TEMPLATE = ( - '\n\n## Oracle hint (PRIVATE — do NOT quote verbatim)\n' - 'The following supporting-fact titles and ground-truth answer are ' - 'provided to make your final answer reliable. Use them as a signpost ' - 'while you reason from the context; your final `\\boxed{{...}}` MUST ' - 'paraphrase the ground truth using evidence from the blocks (after ' - 'expanding compressed blocks when needed), not just echo it.\n' - 'Supporting facts (titles): {sf}\n' - 'Ground truth: {gt}\n' - 'You MUST still call `extract_condensed` on EVERY compressed block ' - 'whose Summary or More keywords touch any supporting-fact title, even ' - 'if the Summary already seems to state the answer — the compressed ' - 'Summary occasionally loses pronoun referents or attribution and the ' - 'raw passage is the authoritative source.' -) - - -VALIDATION_SYSTEM = ( - 'You are a HotpotQA annotation auditor. Read the raw passages, the ' - 'question, the supplied supporting-fact titles and the supplied ' - 'ground-truth answer. Decide whether this row is usable for training ' - 'a multi-hop QA model.\n\n' - 'Pathologies to catch (drop or fix):\n' - ' - question template leakage: the question literally contains the ' - 'answer, references a passage id, or is malformed;\n' - ' - subject/answer mismatch: the GT does not actually answer the ' - 'question given the passages (e.g. the question asks about an event ' - 'X but GT is from a sibling event Y);\n' - ' - GT entity not present in any passage AND not directly inferable ' - 'by a 2-hop bridge from the passages;\n' - ' - supporting-fact titles obviously incomplete for a 2-hop question.\n' - '\n' - 'Return STRICT JSON ONLY (no markdown fence, no preamble) with this ' - 'exact shape:\n' - ' {"verdict": "ok"|"fix"|"drop", "reason": "", ' - '"fixed_supporting_facts": ["", ...], ' - '"fixed_ground_truth": "<short answer>"}\n' - 'Use verdict "ok" when the supplied SF + GT are correct (then ' - '"fixed_supporting_facts" and "fixed_ground_truth" MAY be empty). ' - 'Use verdict "fix" when the question is answerable but SF or GT are ' - 'wrong/incomplete -- fill the fixed fields with the corrected values, ' - 'titles drawn verbatim from the passage titles below. Use verdict ' - '"drop" when the question itself is invalid or unanswerable from the ' - 'given passages.' -) - - -VALIDATION_USER_TEMPLATE = ( - 'Question: {question}\n' - '\n' - 'Supplied supporting-fact titles: {sf}\n' - 'Supplied ground truth: {gt}\n' - '\n' - 'Passage titles (verbatim):\n{titles}\n' - '\n' - 'Passages (raw, uncompressed):\n\n{passages}' -) - - -# JSON Schema for the OpenAI API; the in-process ExtractCondensed tool's -# tool_info() emits a free-form description that the OpenAI SDK rejects. -EXTRACT_CONDENSED_TOOL: Dict[str, Any] = { - 'type': 'function', - 'function': { - 'name': 'extract_condensed', - 'description': ( - 'Recover the full, uncompressed text of ONE previously ' - 'condensed passage, identified by its <block_N> tag. Use ' - 'this tool whenever you need to re-read the original detail ' - 'of a compressed block. Each call expands exactly one block; ' - 'issue separate calls for additional blocks, and do not ' - 'request the same block twice.'), - 'parameters': { - 'type': 'object', - 'properties': { - 'blocks': { - 'type': 'integer', - 'description': ( - 'The 1-indexed block number N appearing inside ' - '<block_N>...</block_N>. Exactly one block per ' - 'call (e.g. 3); lists are rejected.'), - }, - }, - 'required': ['blocks'], - }, - }, -} - - -F1_ACCEPT_THRESHOLD: float = 0.5 -ROLLOUT_MAX_TURNS: int = 8 -ROLLOUT_MAX_TOKENS: int = 2048 -VALIDATION_MAX_TOKENS: int = 1024 -ROLLOUT_TEMPERATURE_LADDER: Tuple[float, ...] = (0.4, 0.7) - - -# -------------------------------------------------------------------------- -# Trajectory + chunk helpers (mirror HotpotQAProcessor + production prompt). -# -------------------------------------------------------------------------- -def _format_passage(title: str, sentences: Any) -> str: - if isinstance(sentences, list): - body = ' '.join(s.strip() for s in sentences if s and s.strip()) - else: - body = str(sentences).strip() - return f'{title}: {body}' - - -def _format_context(titles: List[str], sentences_list: List[Any]) -> str: - return '\n\n'.join( - _format_passage(t, s) for t, s in zip(titles, sentences_list)) - - -def _build_initial_trajectory(row: Dict[str, Any]) -> Dict[str, Any]: - """Build the pre-compression trajectory dict the chunker expects.""" - ctx = row.get('context') or {} - titles = list(ctx.get('title') or []) - sentences_list = list(ctx.get('sentences') or []) - user_msg = ( - f"Question: {row['question']}\n\n" - f'Context:\n\n{_format_context(titles, sentences_list)}') - return { - 'messages': [ - {'role': 'system', 'content': SYSTEM_PROMPT}, - {'role': 'user', 'content': user_msg}, - ], - } - - -def _extract_question_from_chunk(chunk): - content = chunk.get('content') - if chunk.get('type') != 'text' or not isinstance(content, str): - return None - m = re.search(r'\AQuestion:\s*(.+)', content) - return m.group(1).strip() if m else None - - -# -------------------------------------------------------------------------- -# Per-batch compression (re-use MultiTurnCondenseRollout's batching trick: -# merge all per-row chunks into ONE Chunks so the sampler sees a packed batch). -# -------------------------------------------------------------------------- -def compress_rows( - rows: List[Dict[str, Any]], - chunker: NativeChunker, - condenser: ModelCondenser, -) -> List[Tuple[Dict[str, Any], Chunks]]: - """Return ``[(compressed_trajectory_dict, per_row_Chunks), ...]``. - - ``compressed_trajectory_dict`` already has ``<block_N>...</block_N>`` - wrapping in its user message (see :meth:`Chunks.to_trajectory`). - ``per_row_Chunks`` carries ``raw.original`` snapshots so - :class:`ExtractCondensed` can return the pre-compression text. - """ - if not rows: - return [] - initial = [_build_initial_trajectory(r) for r in rows] - per_row_chunks = [chunker(t) for t in initial] - merged_list: List[Any] = [] - boundaries: List[int] = [] - for ck in per_row_chunks: - merged_list.extend(ck.chunks) - boundaries.append(len(merged_list)) - merged = condenser(Chunks(chunks=merged_list)) - out: List[Tuple[Dict[str, Any], Chunks]] = [] - start = 0 - for end in boundaries: - slc = Chunks(chunks=list(merged.chunks[start:end])) - out.append((slc.to_trajectory(), slc)) - start = end - return out - - -# -------------------------------------------------------------------------- -# Stage 1: validation pass. -# -------------------------------------------------------------------------- -_JSON_FENCE_RE = re.compile(r'```(?:json)?\s*\n(.*?)\n```', re.DOTALL) - - -def _extract_json_object(text: str) -> Optional[Dict[str, Any]]: - """Best-effort JSON parse: strip fence, then locate first ``{...}`` block.""" - if not text: - return None - candidate = text.strip() - m = _JSON_FENCE_RE.search(candidate) - if m: - candidate = m.group(1).strip() - depth = 0 - start = -1 - for i, ch in enumerate(candidate): - if ch == '{': - if depth == 0: - start = i - depth += 1 - elif ch == '}': - depth -= 1 - if depth == 0 and start != -1: - blob = candidate[start:i + 1] - try: - return json.loads(blob) - except json.JSONDecodeError: - start = -1 - continue - return None - - -def validate_row( - api: OpenAI, row: Dict[str, Any], original_gt: List[str], sf_titles: List[str], -) -> Optional[Dict[str, Any]]: - """Return parsed JSON verdict, or ``None`` on unrecoverable parse failure.""" - ctx = row.get('context') or {} - titles = list(ctx.get('title') or []) - sentences_list = list(ctx.get('sentences') or []) - passages = _format_context(titles, sentences_list) - user = VALIDATION_USER_TEMPLATE.format( - question=row['question'], - sf=json.dumps(sf_titles, ensure_ascii=False), - gt=json.dumps(original_gt, ensure_ascii=False), - titles='\n'.join(f'- {t}' for t in titles), - passages=passages, - ) - trajectory = { - 'messages': [ - {'role': 'system', 'content': VALIDATION_SYSTEM}, - {'role': 'user', 'content': user}, - ], - } - sp = SamplingParams( - temperature=0.0, max_tokens=VALIDATION_MAX_TOKENS, num_samples=1) - for attempt in range(2): - try: - reply = api( - trajectory, sp, extra_body={'enable_thinking': True}) - except Exception as exc: - sys.stderr.write(f'[validate] row={row.get("id")} attempt={attempt} api error: {exc}\n') - return None - content = reply.get('content') or '' - parsed = _extract_json_object(content) - if parsed and parsed.get('verdict') in ('ok', 'fix', 'drop'): - return parsed - return None - - -def resolve_validation( - verdict: Dict[str, Any], original_gt: List[str], sf_titles: List[str], -) -> Tuple[List[str], List[str]]: - """Pick the SF + GT list to use downstream based on verdict.""" - v = verdict.get('verdict') - if v == 'fix': - fixed_gt = verdict.get('fixed_ground_truth') or '' - fixed_sf = verdict.get('fixed_supporting_facts') or [] - gt_list: List[str] = [] - if isinstance(fixed_gt, list): - gt_list = [str(x).strip() for x in fixed_gt if str(x).strip()] - elif isinstance(fixed_gt, str) and fixed_gt.strip(): - gt_list = [fixed_gt.strip()] - if not gt_list: - gt_list = original_gt - sf_list = ( - [str(x).strip() for x in fixed_sf if str(x).strip()] - if isinstance(fixed_sf, list) else sf_titles) - if not sf_list: - sf_list = sf_titles - return gt_list, sf_list - return original_gt, sf_titles - - -# -------------------------------------------------------------------------- -# Stage 2 prep: build oracle trajectory + per-trajectory ToolManager. -# -------------------------------------------------------------------------- -def _oracle_system_prompt(sf_titles: List[str], gt_list: List[str]) -> str: - sf_render = ', '.join(repr(t) for t in sf_titles) if sf_titles else '(none)' - gt_render = ' | '.join(gt_list) if gt_list else '(unknown)' - return SYSTEM_PROMPT + _ORACLE_HINT_TEMPLATE.format( - sf=sf_render, gt=gt_render) - - -def _build_oracle_trajectory( - compressed_traj: Dict[str, Any], - sf_titles: List[str], - gt_list: List[str], -) -> Dict[str, Any]: - """Replace the system message with the oracle-suffixed variant and - attach the JSON-schema tools field consumed by the OpenAI API.""" - oracle_sp = _oracle_system_prompt(sf_titles, gt_list) - out_messages: List[Dict[str, Any]] = [] - sys_inserted = False - for m in compressed_traj.get('messages') or []: - if m.get('role') == 'system' and not sys_inserted: - out_messages.append({'role': 'system', 'content': oracle_sp}) - sys_inserted = True - else: - out_messages.append(dict(m)) - if not sys_inserted: - out_messages.insert(0, {'role': 'system', 'content': oracle_sp}) - return { - 'messages': out_messages, - 'tools': [EXTRACT_CONDENSED_TOOL], - } - - -def _make_tool_manager(chunks: Chunks) -> ToolManager: - """One ToolManager + ExtractCondensed per trajectory; the tool keeps - a ``_already_expanded`` set, so reusing across trials would lie to - the model on retry.""" - tm = ToolManager() - tm.register(ExtractCondensed(chunks)) - return tm - - -# -------------------------------------------------------------------------- -# Stage 3 + 4: F1 acceptance + conversion to training-runtime format. -# -------------------------------------------------------------------------- -def boxed_f1(boxed: str, gt_list: List[str]) -> float: - if not boxed or not gt_list: - return 0.0 - return max(_f1_score(boxed, g)[0] for g in gt_list) - - -def _last_assistant_text(messages: List[Dict[str, Any]]) -> str: - for m in reversed(messages): - if m.get('role') == 'assistant' and isinstance(m.get('content'), str): - return m['content'] - return '' - - -def _format_tool_call_text(blocks: int) -> str: - return ( - '<tool_call>\n' - '<function=extract_condensed>\n' - '<parameter=blocks>\n' - f'{blocks}\n' - '</parameter>\n' - '</function>\n' - '</tool_call>' - ) - - -def convert_to_runtime_messages( - api_messages: List[Dict[str, Any]], -) -> List[Dict[str, Any]]: - """OpenAI tool_calls -> textual <tool_call> format consumed by the - training chat template. The first system message has its oracle - suffix stripped (we just replace it with the clean SYSTEM_PROMPT). - """ - out: List[Dict[str, Any]] = [] - sys_done = False - for m in api_messages: - role = m.get('role') - if role == 'system' and not sys_done: - out.append({'role': 'system', 'content': SYSTEM_PROMPT}) - sys_done = True - continue - if role == 'assistant': - content = m.get('content') or '' - tool_calls = m.get('tool_calls') or [] - if tool_calls: - pieces = [content.rstrip()] if content else [] - for tc in tool_calls: - fn = tc.get('function') or {} - args_raw = fn.get('arguments') - try: - args = ( - json.loads(args_raw) if isinstance(args_raw, str) - else (args_raw or {})) - except json.JSONDecodeError: - args = {} - blocks_val = args.get('blocks', args.get('block')) - try: - n = int(blocks_val) - except (TypeError, ValueError): - continue - pieces.append(_format_tool_call_text(n)) - text = '\n\n'.join(p for p in pieces if p) - out.append({'role': 'assistant', 'content': text}) - else: - out.append({'role': 'assistant', 'content': content}) - continue - if role == 'tool': - out.append({'role': 'tool', 'content': m.get('content') or ''}) - continue - out.append({k: v for k, v in m.items() if k in ('role', 'content')}) - return out - - -def trajectory_achieved_ratio(chunks: Chunks) -> float: - total_src = 0 - total_cmp = 0 - for c in chunks.chunks: - if c.get('type') != 'text': - continue - raw = c.get('raw') - if not (isinstance(raw, dict) and raw.get('condensed')): - continue - original = raw.get('original') - compressed = c.get('content') - if isinstance(original, str) and isinstance(compressed, str): - total_src += len(original) - total_cmp += len(compressed) - return round(total_cmp / total_src, 4) if total_src else 0.0 - - -def build_record( - row: Dict[str, Any], - runtime_messages: List[Dict[str, Any]], - chunks: Chunks, - verdict: Dict[str, Any], - original_gt: List[str], - used_gt: List[str], - used_sf: List[str], - boxed: str, - f1: float, - num_tool_calls: int, -) -> Dict[str, Any]: - ctx = row.get('context') or {} - titles = list(ctx.get('title') or []) - sentences_list = list(ctx.get('sentences') or []) - raw_passages = [ - { - 'title': t, - 'sentences': list(s) if isinstance(s, list) else [str(s)], - } - for t, s in zip(titles, sentences_list) - ] - sf_full = row.get('supporting_facts') or {} - return { - 'id': row['id'], - 'level': row.get('level'), - 'type': row.get('type'), - 'messages': runtime_messages, - 'tools': [EXTRACT_CONDENSED_TOOL], - 'meta': { - 'num_tool_calls': num_tool_calls, - 'achieved_ratio': trajectory_achieved_ratio(chunks), - 'validation_verdict': verdict.get('verdict'), - 'validation_reason': verdict.get('reason'), - 'original_question': row.get('question'), - 'original_answer': row.get('answer'), - 'original_gt': original_gt, - 'used_gt': used_gt, - 'used_supporting_facts': used_sf, - 'original_supporting_facts': { - 'title': list(sf_full.get('title') or []), - 'sent_id': list(sf_full.get('sent_id') or []), - }, - 'original_passages': raw_passages, - 'f1': round(f1, 4), - 'boxed': boxed, - }, - } - - -# -------------------------------------------------------------------------- -# Per-batch pipeline orchestration. -# -------------------------------------------------------------------------- -def _extract_original_gt_sf(row: Dict[str, Any]) -> Tuple[List[str], List[str]]: - answers = row.get('answers') - if isinstance(answers, list) and answers: - original_gt = [str(a).strip() for a in answers if str(a).strip()] - else: - original_gt = [(row.get('answer', '') or '').strip()] - original_gt = [g for g in original_gt if g] - sf = row.get('supporting_facts') or {} - sf_titles = list(dict.fromkeys(t for t in (sf.get('title') or []) if t)) - return original_gt, sf_titles - - -def _validate_in_parallel( - api: OpenAI, batch: List[Dict[str, Any]], pool: ThreadPoolExecutor, -) -> Tuple[List[Optional[Dict[str, Any]]], List[Tuple[List[str], List[str]]]]: - """Run ``validate_row`` for every row in parallel (one OpenAI call each).""" - futures = [] - payloads: List[Tuple[List[str], List[str]]] = [] - for row in batch: - original_gt, sf_titles = _extract_original_gt_sf(row) - payloads.append((original_gt, sf_titles)) - futures.append(pool.submit( - validate_row, api, row, original_gt, sf_titles)) - verdicts: List[Optional[Dict[str, Any]]] = [f.result() for f in futures] - return verdicts, payloads - - -def _num_tool_calls(messages: List[Dict[str, Any]]) -> int: - return sum( - len(m.get('tool_calls') or []) - for m in messages if m.get('role') == 'assistant') - - -def process_batch( - api: OpenAI, - rollout: APIMultiTurnRollout, - batch: List[Dict[str, Any]], - chunker: NativeChunker, - condenser: ModelCondenser, - validation_pool: ThreadPoolExecutor, -) -> List[Dict[str, Any]]: - """Validate -> compress -> rollout (T-ladder) -> accept. Returns the - list of accepted JSONL records for the batch.""" - if not batch: - return [] - # 1. Validation in parallel. - verdicts, payloads = _validate_in_parallel(api, batch, validation_pool) - - survivors_meta: List[Dict[str, Any]] = [] - for row, verdict, (original_gt, sf_titles) in zip(batch, verdicts, payloads): - if verdict is None or verdict.get('verdict') == 'drop': - continue - if not original_gt: - continue - used_gt, used_sf = resolve_validation(verdict, original_gt, sf_titles) - if not used_gt: - continue - survivors_meta.append({ - 'row': row, 'verdict': verdict, - 'original_gt': original_gt, - 'used_gt': used_gt, 'used_sf': used_sf, - }) - if not survivors_meta: - return [] - - # 2. Compress survivors (one packed batch through ModelCondenser). - survivor_rows = [m['row'] for m in survivors_meta] - try: - compressed = compress_rows(survivor_rows, chunker, condenser) - except Exception as exc: - sys.stderr.write(f'[compress] batch crashed: {exc}\n') - return [] - - # 3. Build oracle trajectories + per-trajectory ToolManagers. - trajs: List[Dict[str, Any]] = [] - chunks_list: List[Chunks] = [] - for meta, (compressed_traj, chunks) in zip(survivors_meta, compressed): - trajs.append(_build_oracle_trajectory( - compressed_traj, meta['used_sf'], meta['used_gt'])) - chunks_list.append(chunks) - - # 4. Temperature ladder. Each rung gets fresh ExtractCondensed tools so - # a retry does not see the previous attempt's already-expanded set. - accepted: List[Dict[str, Any]] = [] - pending_idx = list(range(len(trajs))) - for temperature in ROLLOUT_TEMPERATURE_LADDER: - if not pending_idx: - break - sp = SamplingParams( - temperature=temperature, max_tokens=ROLLOUT_MAX_TOKENS, num_samples=1) - run_trajs = [trajs[i] for i in pending_idx] - run_tms = [_make_tool_manager(chunks_list[i]) for i in pending_idx] - try: - outs = rollout( - run_trajs, tool_manager=run_tms, sampling_params=sp) - except Exception as exc: - sys.stderr.write(f'[rollout] batch crashed at T={temperature}: {exc}\n') - return accepted - next_pending: List[int] = [] - for local_pos, traj_idx in enumerate(pending_idx): - out_traj = outs[local_pos] - if out_traj.get('stop_reason') == 'api_error': - continue # hard-drop API failures, do not retry - messages = out_traj.get('messages') or [] - boxed = _extract_final_answer(_last_assistant_text(messages)) - meta = survivors_meta[traj_idx] - f1 = boxed_f1(boxed, meta['used_gt']) - if f1 >= F1_ACCEPT_THRESHOLD: - runtime_messages = convert_to_runtime_messages(messages) - accepted.append(build_record( - row=meta['row'], - runtime_messages=runtime_messages, - chunks=chunks_list[traj_idx], - verdict=meta['verdict'], - original_gt=meta['original_gt'], - used_gt=meta['used_gt'], - used_sf=meta['used_sf'], - boxed=boxed, f1=f1, - num_tool_calls=_num_tool_calls(messages))) - else: - next_pending.append(traj_idx) - pending_idx = next_pending - return accepted - - -# -------------------------------------------------------------------------- -# Stratified sampling + resume. -# -------------------------------------------------------------------------- -LEVELS: Tuple[str, str, str] = ('easy', 'medium', 'hard') - - -def stratified_sample( - ds, per_level: Dict[str, int], seed: int, -) -> List[Dict[str, Any]]: - rng = random.Random(seed) - buckets: Dict[str, List[int]] = {lv: [] for lv in LEVELS} - for i, lv in enumerate(ds['level']): - if lv in buckets: - buckets[lv].append(i) - picked: List[int] = [] - for lv in LEVELS: - need = per_level[lv] - pool = buckets[lv] - if len(pool) < need: - raise RuntimeError( - f'level={lv} has only {len(pool)} rows, need {need}') - picked.extend(rng.sample(pool, need)) - rng.shuffle(picked) - return [ds[int(i)] for i in picked] - - -def load_done_ids(path: str) -> set: - if not os.path.exists(path): - return set() - done = set() - with open(path, 'r', encoding='utf-8') as fh: - for line in fh: - try: - obj = json.loads(line) - except json.JSONDecodeError: - continue - rid = obj.get('id') - if rid: - done.add(rid) - return done - - -def apply_reannotation_overlay( - rows: List[Dict[str, Any]], path: str, -) -> List[Dict[str, Any]]: - """Drop verdict=drop ids; overlay ``question_fixed`` and multi-form ``answers``. - - The validation stage in ``process_batch`` still runs on every survivor - because the audit ran on a different HF subset (fullwiki) than this - builder's default (distractor) and passage contexts differ. - """ - overrides: Dict[str, Dict[str, Any]] = {} - drop_ids: set = set() - with open(path, 'r', encoding='utf-8') as fh: - for line in fh: - line = line.strip() - if not line: - continue - try: - obj = json.loads(line) - except json.JSONDecodeError: - continue - rid = obj.get('id') - if not rid: - continue - if obj.get('verdict') == 'drop': - drop_ids.add(rid) - else: - overrides[rid] = obj - out: List[Dict[str, Any]] = [] - overridden = 0 - for row in rows: - rid = row.get('id') - if rid in drop_ids: - continue - ov = overrides.get(rid) - if ov is not None: - row = dict(row) - qfix = (ov.get('question_fixed') or '').strip() - if qfix: - row['question'] = qfix - ans = [str(a).strip() for a in (ov.get('answers') or []) if str(a).strip()] - if ans: - row['answers'] = ans - overridden += 1 - out.append(row) - sys.stderr.write( - f'[REANNOTATED] {path}: {len(rows)} -> {len(out)} rows ' - f'(dropped={len(drop_ids)}, overridden={overridden})\n') - return out - - -# -------------------------------------------------------------------------- -# CLI + main loop. -# -------------------------------------------------------------------------- -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser() - parser.add_argument('--output', required=True) - parser.add_argument('--model', required=True, - help='Super-LLM model name (OpenAI-protocol).') - parser.add_argument('--api-key', default=os.environ.get('OPENAI_API_KEY')) - parser.add_argument('--base-url', default=os.environ.get('OPENAI_BASE_URL')) - parser.add_argument('--total', type=int, default=12000) - parser.add_argument('--easy', type=int, default=2000) - parser.add_argument('--medium', type=int, default=4000) - parser.add_argument('--hard', type=int, default=6000) - parser.add_argument('--concurrency', type=int, default=16) - parser.add_argument('--seed', type=int, default=42) - parser.add_argument('--reannotated', default=os.environ.get('REANNOTATED_FILE', ''), - help='Path to wrong_ids_reannotated.jsonl. Drops verdict=drop ids and overlays question_fixed + multi-form answers. Validation stage still runs because the audit was on a different HF subset.') - parser.add_argument('--hf-subset', default='distractor') - parser.add_argument('--hf-split', default='train') - parser.add_argument('--condenser-model-id', - default=os.environ.get('MODEL_ID', 'ms://Qwen/Qwen3.5-4B')) - parser.add_argument('--condenser-lora', - default='ms://twinkle-kit/Qwen3.5-4B-Condenser') - parser.add_argument('--chunk-size', type=int, default=1024) - parser.add_argument('--hotpotqa-max-length', type=int, default=64000) - parser.add_argument('--compress-batch-size', type=int, default=32, - help='How many rows to feed to ModelCondenser at once.') - parser.add_argument('--gpu-memory-utilization', type=float, default=0.8) - return parser.parse_args() - - -def build_condenser(args: argparse.Namespace) -> Tuple[NativeChunker, ModelCondenser]: - sampler = vLLMSampler( - model_id=args.condenser_model_id, - engine_args={ - 'gpu_memory_utilization': args.gpu_memory_utilization, - 'max_model_len': max(8192, args.hotpotqa_max_length), - 'max_lora_rank': 32, - 'enable_lora': True, - 'max_loras': 2, - }, - ) - sampler.set_template( - 'Qwen3_5Template', model_id=args.condenser_model_id, - enable_thinking=False, max_length=args.hotpotqa_max_length) - rollout_template = Qwen3_5Template( - args.condenser_model_id, max_length=args.hotpotqa_max_length, - enable_thinking=False) - chunker = NativeChunker( - chunk_size=args.chunk_size, - passage_boundary_re=r'(?<=\n\n)', - ) - condenser = ModelCondenser( - sampler=sampler, - compression_ratio=2.0, - sampling_params=SamplingParams( - max_tokens=1024, num_samples=1, temperature=0.4, top_p=0.9), - min_chars=200, - template=rollout_template, - lora_path=args.condenser_lora or None, - skip_pattern=r'^Question:', - related_query=_extract_question_from_chunk, - ) - return chunker, condenser - - -def main() -> None: - args = parse_args() - if args.easy + args.medium + args.hard != args.total: - raise ValueError( - f'--easy + --medium + --hard ({args.easy + args.medium + args.hard}) ' - f'must equal --total ({args.total})') - per_level = {'easy': args.easy, 'medium': args.medium, 'hard': args.hard} - - sys.stderr.write( - f'Loading hotpotqa/hotpot_qa:{args.hf_subset}:{args.hf_split}...\n') - ds = load_dataset( - 'hotpotqa/hotpot_qa', args.hf_subset, split=args.hf_split) - - rows = stratified_sample(ds, per_level=per_level, seed=args.seed) - if args.reannotated.strip(): - rows = apply_reannotation_overlay(rows, args.reannotated.strip()) - done = load_done_ids(args.output) - sys.stderr.write(f'Resume: {len(done)} rows already emitted.\n') - pending = [r for r in rows if r['id'] not in done] - sys.stderr.write(f'Pending: {len(pending)} / {len(rows)}\n') - - chunker, condenser = build_condenser(args) - api = OpenAI( - model=args.model, api_key=args.api_key, base_url=args.base_url) - - # APIMultiTurnRollout itself owns the per-trajectory thread pool. The - # validation phase runs on a separate pool of equal size; both phases - # are network-bound so we never need more threads than ``concurrency``. - rollout = APIMultiTurnRollout( - api=api, - tool_manager=ToolManager(), # placeholder; per-call list overrides - sampling_params=SamplingParams( - temperature=ROLLOUT_TEMPERATURE_LADDER[0], - max_tokens=ROLLOUT_MAX_TOKENS, num_samples=1), - max_turns=ROLLOUT_MAX_TURNS, - concurrency=args.concurrency, - extra_body={'enable_thinking': False}, - ) - - write_lock = threading.Lock() - out_fh = open(args.output, 'a', encoding='utf-8') - accepted_total = 0 - seen_total = 0 - - with ThreadPoolExecutor(max_workers=args.concurrency) as validation_pool: - try: - for start in range(0, len(pending), args.compress_batch_size): - batch = pending[start:start + args.compress_batch_size] - seen_total += len(batch) - try: - records = process_batch( - api, rollout, batch, chunker, condenser, - validation_pool) - except Exception as exc: - sys.stderr.write( - f'[batch {start}-{start + len(batch)}] crashed: {exc}\n') - continue - with write_lock: - for record in records: - out_fh.write( - json.dumps(record, ensure_ascii=False) + '\n') - out_fh.flush() - accepted_total += len(records) - sys.stderr.write( - f'[progress] seen={seen_total}/{len(pending)} ' - f'accepted={accepted_total} ' - f'(+{len(records)} from this batch)\n') - finally: - out_fh.close() - - sys.stderr.write( - f'Done. accepted={accepted_total} total_pending={len(pending)}\n') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/legacy/make_embedding_dataset.py b/cookbook/exp/legacy/make_embedding_dataset.py deleted file mode 100644 index 847f222fc..000000000 --- a/cookbook/exp/legacy/make_embedding_dataset.py +++ /dev/null @@ -1,758 +0,0 @@ -"""Offline compression pipeline: raw datasets → condenser → pre-compressed embedding dataset. - -Loads think/index/hard datasets, compresses query/cot/negatives via vLLM condenser -with API fallback, saves a single HF Dataset ready for embedding training. - -Output schema: {anchor_text, positive_text, negative_texts, source} - -Launch (8 GPUs — 4 for vLLM condenser): - python cookbook/exp/embedding/make_embedding_dataset.py -""" -import json -import os -import re -import sys -import threading -import time -from concurrent.futures import ThreadPoolExecutor, as_completed -from pathlib import Path -from typing import Any, Dict, List, Optional - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.data_format import SamplingParams -from twinkle.sampler import vLLMSampler -from twinkle.template import Qwen3_5Template -from twinkle.utils.parallel import PosixFileLock -from twinkle_agentic.protocol.openai import OpenAI as OpenAIClient - -sys.path.insert(0, str(Path(__file__).resolve().parent)) -from dataset_think import get_dataset as get_dataset_think # noqa: E402 -from dataset_index import get_dataset as get_dataset_index # noqa: E402 -from dataset_hard import get_dataset as get_dataset_hard # noqa: E402 - -logger = get_logger() - -# -- Model config ------------------------------------------------------------- -CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') -TEMPLATE_NAME = 'Qwen3_5Template' - -# -- GPU placement (condenser only) ------------------------------------------- -CONDENSER_GPUS = int(os.environ.get('CONDENSER_GPUS', 8)) - -# -- Dataset caps ------------------------------------------------------------- -TOTAL_SAMPLES: Optional[int] = None -THINK_CAP: Optional[int] = int(os.environ.get('THINK_CAP', 100_000)) -INDEX_CAP: Optional[int] = int(os.environ.get('INDEX_CAP', 100_000)) -HARD_CAP: Optional[int] = int(os.environ.get('HARD_CAP', 0)) or None -HARD_MAX_NEGATIVES = int(os.environ.get('HARD_MAX_NEGATIVES', 8)) - -# -- Compression params ------------------------------------------------------- -MIN_TEXT_CHARS = 256 -DATASET_MAX_TOKENS = 32768 -COMPRESS_TEMPERATURE = 0.2 -COMPRESS_TOP_P = 0.5 -COMPRESS_MAX_MODEL_LEN = 32768 -BATCH_SIZE = int(os.environ.get('COMPRESS_BATCH_SIZE', 128)) - -# -- API fallback ------------------------------------------------------------- -COMPRESS_API_KEY = os.environ.get('COMPRESS_API_KEY', '') -COMPRESS_BASE_URL = os.environ.get('COMPRESS_BASE_URL', 'https://dashscope.aliyuncs.com/compatible-mode/v1') -COMPRESS_MODEL = os.environ.get('COMPRESS_MODEL', 'qwen3.7-max') -API_MIN_INTERVAL = float(os.environ.get('API_MIN_INTERVAL', 0.1)) -API_CONCURRENCY = int(os.environ.get('API_CONCURRENCY', 24)) -SAMPLER_TIMEOUT = float(os.environ.get('SAMPLER_TIMEOUT', 300)) - -# -- Output ------------------------------------------------------------------- -OUTPUT_DIR = os.environ.get('EMB_DATASET_OUTPUT', './output/embedding_dataset') -RESULTS_JSONL = f'{OUTPUT_DIR}/results.jsonl' -PROGRESS_FILE = f'{OUTPUT_DIR}/progress.json' - -# ============================================================================= -# Prompts -# ============================================================================= - -COMPRESS_SYSTEM = """\ -You are a compression and summary assistant. For the (query, source) pair, emit a Markdown \ -answer with TWO sections, designed to pair with the `extract_compressed` tool: \ -the reader absorbs `## Summary` directly, then calls `extract_compressed` \ -on any topic-key listed under `## More` to recover its \ -fuller content. - - `## Summary` — extreme-density text the reader reads directly. - `## More` — a topic index whose keys are valid arguments \ -to `extract_compressed` for recovering material not captured inline. - -Together the two sections must form a COMPLETE, NON-DISTORTING inventory of the \ -source for the query — nothing essential lost, nothing implied that the source \ -does not support. NO preamble, NO meta-commentary, NO code fences wrapping the \ -whole output. - -Output skeleton: - -## Summary -Topic: <what the source is about + scope, one line> -<dense body answering the query> - -## More -- <topic-key>: <one-line hint of what is revealed when expanded> -- ... - -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: <trigger>`; numbered telegraphic \ -steps `1.do X 2.then Y`; close with `Output: <result>` when relevant -- Procedural → numbered short steps -- Analytical / design → hierarchical bullets with abbreviations - -`## Summary` rules: -1. TOPIC LINE — line 1 is ALWAYS `Topic: <subject — scope>`, 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>: <one-line hint>`: - • 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: <specific pattern name — scope>\n' - 'Problem: <what concrete problem is being asked>\n' - 'Skill: <which specific method/technique/pattern is required to solve it>\n' - 'Knowledge: <which domains/concepts/facts must be invoked>\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: <specific pattern name — scope>\n' - 'Problem: <what concrete problem this trace tackled>\n' - 'Skill: <which specific method/technique/pattern was applied>\n' - 'Knowledge: <which domains/concepts/facts were used>\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: <problem class — mathematical/logical domain>\n' - 'Problem: <what abstract TYPE of problem needs solving, no specific values>\n' - 'Skill: <which general method/technique is required>\n' - 'Knowledge: <which theoretical concepts/formulas must be invoked>\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: <method/technique name — scope>\n' - 'Problem: <what abstract type of problem this method solves>\n' - 'Skill: <key steps of the methodology in abstract terms>\n' - 'Knowledge: <theoretical basis and prerequisites>\n' - 'Then emit the mandatory ## More section as usual. ' - 'Topic must name the method class, never mention specific numbers.') - - -# ============================================================================= -# Validation & API fallback -# ============================================================================= - -_LEGACY_USE_WHEN_RE = re.compile(r'(?im)^\s*Use when\s*:') -_SCHEMA_MARKERS = ('Problem:', 'Skill:', 'Knowledge:') - - -def _is_truncated_compression(text: str, schema: str = 'new') -> bool: - if not text or not text.strip(): - return True - if '## More' not in text or '## Summary' not in text: - return True - after_more = text.split('## More', 1)[1].strip() - if not after_more: - return True - last_line = after_more.splitlines()[-1].strip() - if not (last_line.startswith('-') or last_line.endswith(')')): - return True - if schema == 'new': - summary_body = text.split('## Summary', 1)[1].split('## More', 1)[0] - if _LEGACY_USE_WHEN_RE.search(summary_body): - return True - if not all(marker in summary_body for marker in _SCHEMA_MARKERS): - return True - return False - - -_api_semaphore = threading.Semaphore(API_CONCURRENCY) -_api_bucket_lock = threading.Lock() -_api_tokens = [float(API_CONCURRENCY)] -_api_last_refill = [time.monotonic()] - - -def _api_throttle(): - """Token-bucket rate limiter: API_CONCURRENCY requests per API_MIN_INTERVAL*API_CONCURRENCY window.""" - _api_semaphore.acquire() - try: - with _api_bucket_lock: - now = time.monotonic() - elapsed = now - _api_last_refill[0] - refill = elapsed / API_MIN_INTERVAL - _api_tokens[0] = min(float(API_CONCURRENCY), _api_tokens[0] + refill) - _api_last_refill[0] = now - if _api_tokens[0] >= 1.0: - _api_tokens[0] -= 1.0 - else: - wait = (1.0 - _api_tokens[0]) * API_MIN_INTERVAL - _api_tokens[0] = 0.0 - time.sleep(wait) - finally: - _api_semaphore.release() - - -def _api_compress(api_client: OpenAIClient, prompt: Dict[str, Any]) -> Optional[str]: - _api_throttle() - trajectory = {'messages': prompt['messages']} - sp = SamplingParams(temperature=0.2, max_tokens=8192) - try: - reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) - except Exception as exc: - logger.warning(f'[api_fallback] error: {exc}') - return None - content = (reply.get('content') or '').strip() - if not content: - return None - m = re.match(r'^```[a-zA-Z]*\n(.*?)\n```\s*$', content, re.DOTALL) - if m: - content = m.group(1).strip() - return content - - -# ============================================================================= -# Core compression logic -# ============================================================================= - -def _extract_query_cot(row: Dict[str, Any]): - messages = row.get('messages') or [] - query, cot = '', '' - for m in messages: - if not isinstance(m, dict): - continue - role = m.get('role') or '' - if role == 'user' and not query: - query = (m.get('content') or '').strip() - elif role == 'assistant': - cot = (m.get('reasoning_content') or '').strip() - break - return query, cot - - -def _compress_batch_phase1( - rows: List[Dict[str, Any]], - condenser_sampler, - compress_params: SamplingParams, - special_tokens: set, - source_type: str, -) -> Optional[Dict[str, Any]]: - """Phase 1 (GPU): build prompts → vLLM sample → validate. Returns state for phase 2.""" - _MAX_COT_CHARS = 30_000 - - if source_type == 'hard': - return _compress_hard_phase1(rows, condenser_sampler, compress_params, - special_tokens, source_type) - - prompts: List[Optional[Dict[str, Any]]] = [] - meta: List[Dict[str, Any]] = [] - for i, row in enumerate(rows): - query, cot = _extract_query_cot(row) - if not query or len(cot) < MIN_TEXT_CHARS or len(cot) > _MAX_COT_CHARS: - continue - schema = 'legacy' if (i % 2 == 0) else 'new' - q_hint = EMBED_QUERY_Q_LEGACY if schema == 'legacy' else EMBED_QUERY_Q - c_hint = EMBED_QUERY_COT_LEGACY if schema == 'legacy' else EMBED_QUERY_COT - - if len(query) < MIN_TEXT_CHARS: - prompts.append(None) - else: - user = COMPRESS_USER.format(query=q_hint, text=query) - prompts.append({'messages': [ - {'role': 'system', 'content': COMPRESS_SYSTEM}, - {'role': 'user', 'content': user}, - ]}) - user_c = COMPRESS_USER.format(query=c_hint, text=cot) - prompts.append({'messages': [ - {'role': 'system', 'content': COMPRESS_SYSTEM}, - {'role': 'user', 'content': user_c}, - ]}) - meta.append({'query_raw': query, 'cot_raw': cot, 'schema': schema, - 'q_hint': q_hint, 'source': source_type, - 'row_id': row.get('id', str(i))}) - - if not prompts: - return {'final': []} - - sampler_input = [p for p in prompts if p is not None] - sampler_pos = [ri for ri, p in enumerate(prompts) if p is not None] - try: - sampler_responses = condenser_sampler.sample(sampler_input, compress_params) - except Exception as exc: - logger.warning(f'[compress] sampler error: {exc}') - sampler_responses = [None] * len(sampler_input) - - responses = [None] * len(prompts) - for resp, pos in zip(sampler_responses, sampler_pos): - responses[pos] = resp - - decoded: List[str] = [] - fallback_indices: List[int] = [] - for ri in range(len(prompts)): - pair_idx = ri // 2 - schema = meta[pair_idx]['schema'] - if prompts[ri] is None: - decoded.append(meta[pair_idx]['query_raw']) - continue - resp = responses[ri] - seq = resp.sequences[0] if resp and resp.sequences else None - text = '' - if seq and seq.stop_reason != 'length' and seq.decoded: - text = seq.decoded - for tok in special_tokens: - text = text.replace(tok, '') - text = text.rstrip() - if not _is_truncated_compression(text, schema): - decoded.append(text) - else: - decoded.append('') - fallback_indices.append(ri) - - return {'prompts': prompts, 'meta': meta, 'decoded': decoded, - 'fallback_indices': fallback_indices} - - -def _compress_batch_phase2( - state: Dict[str, Any], - api_client: OpenAIClient, -) -> List[Dict[str, Any]]: - """Phase 2 (no GPU): API fallback → build results.""" - if 'final' in state: - return state['final'] - - prompts = state['prompts'] - decoded = state['decoded'] - fallback_indices = state['fallback_indices'] - is_hard = state.get('hard', False) - meta = state.get('meta') # None for hard - - # Track which prompts used API fallback - api_set: set = set() - if fallback_indices: - api_futures = {} - with ThreadPoolExecutor(max_workers=API_CONCURRENCY) as pool: - for ri in fallback_indices: - api_futures[pool.submit(_api_compress, api_client, prompts[ri])] = ri - for fut in as_completed(api_futures): - ri = api_futures[fut] - api_result = fut.result() - schema = 'new' if is_hard else meta[ri // 2]['schema'] - if api_result and not _is_truncated_compression(api_result, schema): - decoded[ri] = api_result - api_set.add(ri) - - state['api_set'] = api_set - if is_hard: - return _build_hard_results(state) - return _build_think_index_results(state) - - -def _build_think_index_results(state: Dict[str, Any]) -> List[Dict[str, Any]]: - meta = state['meta'] - decoded = state['decoded'] - api_set = state.get('api_set', set()) - results = [] - for pair_idx in range(len(meta)): - q_text = decoded[pair_idx * 2] - c_text = decoded[pair_idx * 2 + 1] - if not q_text or not c_text: - continue - q_method = 'api' if (pair_idx * 2) in api_set else 'vllm' - c_method = 'api' if (pair_idx * 2 + 1) in api_set else 'vllm' - results.append({ - 'anchor_text': q_text, - 'positive_text': c_text, - 'negative_texts': [], - 'source': meta[pair_idx]['source'], - 'query_raw': meta[pair_idx]['query_raw'], - 'cot_raw': meta[pair_idx]['cot_raw'], - 'anchor_method': q_method, - 'positive_method': c_method, - }) - return results - - -def _build_hard_results(state: Dict[str, Any]) -> List[Dict[str, Any]]: - group_sizes = state['group_sizes'] - decoded = state['decoded'] - source_type = state['source_type'] - raw_groups = state['raw_groups'] - api_set = state.get('api_set', set()) - results = [] - offset = 0 - for gi, gs in enumerate(group_sizes): - q_text = decoded[offset] - c_text = decoded[offset + 1] - if not q_text or not c_text: - offset += gs - continue - neg_texts = [] - neg_raws = [] - neg_methods = [] - for ni in range(2, gs): - nt = decoded[offset + ni] - if nt: - neg_texts.append(nt) - neg_raws.append(raw_groups[gi]['negs_raw'][ni - 2]) - neg_methods.append('api' if (offset + ni) in api_set else 'vllm') - q_method = 'api' if offset in api_set else 'vllm' - c_method = 'api' if (offset + 1) in api_set else 'vllm' - results.append({ - 'anchor_text': q_text, - 'positive_text': c_text, - 'negative_texts': neg_texts, - 'source': source_type, - 'query_raw': raw_groups[gi]['query_raw'], - 'cot_raw': raw_groups[gi]['cot_raw'], - 'negs_raw': neg_raws, - 'anchor_method': q_method, - 'positive_method': c_method, - 'neg_methods': neg_methods, - }) - offset += gs - return results - - -def _compress_hard_phase1( - rows: List[Dict[str, Any]], - condenser_sampler, - compress_params: SamplingParams, - special_tokens: set, - source_type: str, -) -> Dict[str, Any]: - """Phase 1 for hard rows: vLLM sample + validate. Returns state for phase 2.""" - _MAX_COT_CHARS = 30_000 - - prompts: List[Dict[str, Any]] = [] - group_sizes: List[int] = [] - row_ids: List[str] = [] - raw_groups: List[Dict[str, Any]] = [] - - for row in rows: - query, cot = _extract_query_cot(row) - if not query or not cot or len(cot) > _MAX_COT_CHARS: - continue - negatives = row.get('negatives') or [] - valid_negs = [n for n in negatives - if n and len(n) <= _MAX_COT_CHARS] - - user_q = COMPRESS_USER.format(query=EMBED_QUERY_REASONIR_Q, text=query) - prompts.append({'messages': [ - {'role': 'system', 'content': COMPRESS_SYSTEM}, - {'role': 'user', 'content': user_q}, - ]}) - user_c = COMPRESS_USER.format(query=EMBED_QUERY_REASONIR_COT, text=cot) - prompts.append({'messages': [ - {'role': 'system', 'content': COMPRESS_SYSTEM}, - {'role': 'user', 'content': user_c}, - ]}) - for neg in valid_negs: - user_n = COMPRESS_USER.format(query=EMBED_QUERY_REASONIR_COT, text=neg) - prompts.append({'messages': [ - {'role': 'system', 'content': COMPRESS_SYSTEM}, - {'role': 'user', 'content': user_n}, - ]}) - group_sizes.append(2 + len(valid_negs)) - row_ids.append(row.get('id', '')) - raw_groups.append({'query_raw': query, 'cot_raw': cot, 'negs_raw': valid_negs}) - - if not prompts: - return {'hard': True, 'prompts': [], 'group_sizes': [], 'row_ids': [], - 'decoded': [], 'fallback_indices': [], 'source_type': source_type, - 'raw_groups': []} - - try: - responses = condenser_sampler.sample(prompts, compress_params) - except Exception as exc: - logger.warning(f'[compress-hard] sampler error: {exc}') - responses = [None] * len(prompts) - - decoded: List[str] = [] - fallback_indices: List[int] = [] - for ri, resp in enumerate(responses): - seq = resp.sequences[0] if resp and resp.sequences else None - text = '' - if seq and seq.stop_reason != 'length' and seq.decoded: - text = seq.decoded - for tok in special_tokens: - text = text.replace(tok, '') - text = text.rstrip() - if text and not _is_truncated_compression(text, 'new'): - decoded.append(text) - else: - decoded.append('') - fallback_indices.append(ri) - - return {'hard': True, 'prompts': prompts, 'group_sizes': group_sizes, - 'row_ids': row_ids, 'decoded': decoded, - 'fallback_indices': fallback_indices, 'source_type': source_type, - 'raw_groups': raw_groups} - - -# ============================================================================= -# Main pipeline -# ============================================================================= - -def main(): - device_groups = [ - DeviceGroup(name='condenser_sampler', - ranks=list(range(CONDENSER_GPUS)), - device_type='GPU'), - ] - condenser_mesh = DeviceMesh.from_sizes( - world_size=CONDENSER_GPUS, dp_size=CONDENSER_GPUS) - twinkle.initialize(mode='ray', nproc_per_node=CONDENSER_GPUS, groups=device_groups) - - # -- Load raw datasets ---------------------------------------------------- - from datasets import Dataset as HFDataset - - dataset_think = get_dataset_think(total=TOTAL_SAMPLES, load_from_cache_file=True) - if THINK_CAP and len(dataset_think.dataset) > THINK_CAP: - dataset_think.dataset = dataset_think.dataset.select(range(THINK_CAP)) - ds_think = dataset_think.dataset - logger.info(f'[load] think={len(ds_think)}') - - ds_index_obj = get_dataset_index(total=None, load_from_cache_file=True) - ds_index = ds_index_obj.dataset - if INDEX_CAP and len(ds_index) > INDEX_CAP: - ds_index = ds_index.select(range(INDEX_CAP)) - logger.info(f'[load] index={len(ds_index)}') - - ds_hard_raw = get_dataset_hard(max_negatives=HARD_MAX_NEGATIVES, load_from_cache_file=True) - if HARD_CAP and len(ds_hard_raw) > HARD_CAP: - ds_hard_raw = ds_hard_raw.select(range(HARD_CAP)) - n_hard = len(ds_hard_raw) - logger.info(f'[load] hard={n_hard}') - - # Convert hard to messages schema - hard_rows_list = [] - if n_hard > 0: - h_ids = ds_hard_raw['id'] - h_queries = ds_hard_raw['query'] - h_cots = ds_hard_raw['cot'] - h_responses = ds_hard_raw['response'] if 'response' in ds_hard_raw.column_names else [''] * n_hard - h_negatives = ds_hard_raw['negatives'] - for i in range(n_hard): - hard_rows_list.append({ - 'id': h_ids[i], - 'messages': [ - {'role': 'user', 'content': h_queries[i]}, - {'role': 'assistant', 'reasoning_content': h_cots[i], - 'content': h_responses[i] or ''}, - ], - 'negatives': h_negatives[i], - }) - - # Batch-convert HF Datasets to list-of-dicts - def _ds_to_rows(ds): - return [dict(zip(ds.column_names, vals)) for vals in zip(*(ds[c] for c in ds.column_names))] - - think_rows = _ds_to_rows(ds_think) - index_rows = _ds_to_rows(ds_index) - - # -- Setup condenser ------------------------------------------------------ - condenser_template = Qwen3_5Template( - model_id=CONDENSE_MODEL_ID, max_length=DATASET_MAX_TOKENS, - enable_thinking=False, truncation_strategy='delete') - special_tokens = set(condenser_template.tokenizer.all_special_tokens) - - condenser_sampler = vLLMSampler( - model_id=CONDENSE_MODEL_ID, - engine_args={'gpu_memory_utilization': 0.8, 'max_model_len': COMPRESS_MAX_MODEL_LEN}, - device_mesh=condenser_mesh, - remote_group='condenser_sampler', - ) - condenser_sampler.set_template( - TEMPLATE_NAME, model_id=CONDENSE_MODEL_ID, enable_thinking=False, - truncation_strategy='delete', max_length=DATASET_MAX_TOKENS) - condenser_sampler._ray_get_timeout = SAMPLER_TIMEOUT - compress_params = SamplingParams( - max_tokens=8192, temperature=COMPRESS_TEMPERATURE, - top_p=COMPRESS_TOP_P, num_samples=1) - - api_client = OpenAIClient( - model=COMPRESS_MODEL, api_key=COMPRESS_API_KEY, base_url=COMPRESS_BASE_URL) - - # -- Resume support ---------------------------------------------------------- - os.makedirs(OUTPUT_DIR, exist_ok=True) - progress = {'think': 0, 'index': 0, 'hard': 0} - if os.path.exists(PROGRESS_FILE): - with open(PROGRESS_FILE, 'r') as f: - progress = json.load(f) - logger.info(f'[resume] loaded progress: {progress}') - - _results_lock = PosixFileLock(RESULTS_JSONL + '.lock') - - def _flush_results(records: List[Dict[str, Any]]): - if not records: - return - lines = [json.dumps(r, ensure_ascii=False) + '\n' for r in records] - with _results_lock: - with open(RESULTS_JSONL, 'a', encoding='utf-8') as f: - f.writelines(lines) - - def _save_progress(): - tmp = PROGRESS_FILE + '.tmp' - with open(tmp, 'w') as f: - json.dump(progress, f) - os.replace(tmp, PROGRESS_FILE) - - # -- Process in batches (pipelined: vLLM batch N+1 overlaps API fallback N) - - total_flushed = 0 - if os.path.exists(RESULTS_JSONL): - with open(RESULTS_JSONL, 'r', encoding='utf-8') as f: - total_flushed = sum(1 for l in f if l.strip()) - if total_flushed: - logger.info(f'[resume] {total_flushed} records already in results.jsonl') - - def _process_source(rows, source_type, label): - nonlocal total_flushed - n_total = len(rows) - skip = progress.get(source_type, 0) - if skip >= n_total: - logger.info(f'[{label}] skipped (already done {skip}/{n_total})') - return - if skip > 0: - logger.info(f'[{label}] resuming from row {skip}/{n_total}') - - bg_pool = ThreadPoolExecutor(max_workers=1) - pending = None # (future, batch_start, batch_len) - - def _drain_pending(): - nonlocal total_flushed, pending - if pending is None: - return - fut, p_start, p_len = pending - batch_results = fut.result() - _flush_results(batch_results) - total_flushed += len(batch_results) - progress[source_type] = p_start + p_len - _save_progress() - pending = None - - for start in range(skip, n_total, BATCH_SIZE): - batch = rows[start:start + BATCH_SIZE] - state = _compress_batch_phase1( - batch, condenser_sampler, compress_params, - special_tokens, source_type) - _drain_pending() - pending = ( - bg_pool.submit(_compress_batch_phase2, state, api_client), - start, len(batch)) - n_done = start + len(batch) - if n_done % (BATCH_SIZE * 10) == 0 or n_done >= n_total: - logger.info(f'[{label}] {n_done}/{n_total} vLLM done, ' - f'{total_flushed} records flushed (last batch pending)') - - _drain_pending() - bg_pool.shutdown(wait=False) - logger.info(f'[{label}] complete, {total_flushed} total records flushed') - - _process_source(hard_rows_list, 'hard', 'hard') - _process_source(think_rows, 'think', 'think') - _process_source(index_rows, 'index', 'index') - - # -- Convert JSONL → HF Dataset ------------------------------------------- - logger.info(f'[save] converting results.jsonl to HF Dataset...') - all_results = [] - with open(RESULTS_JSONL, 'r', encoding='utf-8') as f: - for line_no, line in enumerate(f, 1): - if not line.strip(): - continue - try: - all_results.append(json.loads(line)) - except json.JSONDecodeError: - logger.warning(f'[save] skipping malformed line {line_no} (truncated resume?)') - logger.info(f'[save] total records: {len(all_results)}') - out_ds = HFDataset.from_dict({ - 'anchor_text': [r['anchor_text'] for r in all_results], - 'positive_text': [r['positive_text'] for r in all_results], - 'negative_texts': [r['negative_texts'] for r in all_results], - 'source': [r['source'] for r in all_results], - 'query_raw': [r.get('query_raw', '') for r in all_results], - 'cot_raw': [r.get('cot_raw', '') for r in all_results], - 'negs_raw': [r.get('negs_raw', []) for r in all_results], - }) - out_ds.save_to_disk(OUTPUT_DIR + '/dataset') - logger.info(f'[save] dataset saved to {OUTPUT_DIR}/dataset') - logger.info(f'[stats] think={sum(1 for r in all_results if r["source"]=="think")} ' - f'index={sum(1 for r in all_results if r["source"]=="index")} ' - f'hard={sum(1 for r in all_results if r["source"]=="hard")}') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/legacy/reannotate_groundtruth.py b/cookbook/exp/legacy/reannotate_groundtruth.py deleted file mode 100644 index 137ebb4b9..000000000 --- a/cookbook/exp/legacy/reannotate_groundtruth.py +++ /dev/null @@ -1,389 +0,0 @@ -"""Re-annotate HotpotQA ground truth using a super-LLM to ensure correctness. - -The original HotpotQA dataset has annotation issues: - - GT doesn't match the question type (asks "where", GT gives a name) - - Partial/incomplete answers for multi-hop questions - - Single form when multiple valid forms exist (e.g. "2" vs "two") - - Question itself malformed (wrong question word, truncation, presupposition - mismatch with the answer type) - -This script: - 1. Loads HotpotQA fullwiki train split. - 2. By default (--only-forced), re-annotates ONLY the IDs listed in - wrong_ids.txt (the 340 known-bad cases). - Pass --no-only-forced to fall back to stratified 3000-per-level sampling - with wrong_ids force-included. - 3. For each row, sends question + full context + original GT to a super-LLM. - 4. The LLM emits one of four verdicts and (when applicable) a multi-form - answer list and/or a repaired question: - - keep: original Q + A are both correct - - fix_answer: Q is fine; A is wrong/incomplete - - fix_question: Q is malformed but repairable into a well-formed Q - that the same passages answer with the same gold facts - - drop: Q cannot be repaired without changing the fact, OR - passages do not support any answer - 5. Outputs ONE JSONL file containing all rows (including drop). Each row has - verdict, question, question_fixed, answers, reasoning. Downstream filters - by verdict. - -Run (re-clean wrong_ids.txt only, default): - python reannotate_groundtruth.py \ - --model qwen-max --api-key $OPENAI_API_KEY \ - --base-url https://dashscope.aliyuncs.com/compatible-mode/v1 \ - --output hotpotqa_reannotated_wrong.jsonl --concurrency 16 -""" -import argparse -import json -import os -import random -import re -import sys -import threading -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Any, Dict, List, Optional, Tuple - -from datasets import load_dataset - -from twinkle.data_format.sampling import SamplingParams -from twinkle_agentic.protocol.openai import OpenAI - - -VERIFY_SYSTEM = """You are a dataset quality auditor for a multi-hop QA benchmark (HotpotQA). - -Given a Question, supporting Context passages, and the dataset's Original Answer, output ONE of four verdicts and a multi-form answer list grounded in the passages. - -VERDICTS -- "keep": original question + original answer are both correct. -- "fix_answer": question is fine; original answer is wrong/incomplete. -- "fix_question": question is malformed (wrong question word, broken grammar, truncated, or presupposition mismatch with the answer type) but can be REPAIRED into a well-formed question that the SAME passages answer with the SAME gold facts. -- "drop": question cannot be repaired without changing the underlying fact, OR the passages do not support any answer. - -MULTI-FORM ANSWER RULES (apply to keep / fix_answer / fix_question) -1. Output ALL acceptable surface forms whenever applicable: - - Number variants: arabic + english word + hyphen-prefix form (e.g. "3", "three", "three-door", "3-door") - - Range variants: start, end, and full range string (e.g. "1901", "1902", "1901-1902", "1901-2") - - Location variants: city / state-or-province / country (e.g. "Everett", "Washington", "WA", "United States") - - Person variants: legal name / nickname / full name (e.g. "Allan", "Heywood", "Allan Stewart Konigsberg") - - Entity-role pairs for role-of-X questions: BOTH the role AND the entity (e.g. "chauffeur", "Hitler's chauffeur") - - Show-vs-character pairs for best-known-for questions: BOTH the show AND the character (e.g. "M*A*S*H", "Major Frank Burns") - - Common abbreviations (e.g. "NYC", "New York City", "New York") - - With/without titles (e.g. "Dr. Smith", "Smith") - - Different date formats if applicable (e.g. "July 4, 1776", "4 July 1776") -2. Each answer is SHORT (a name, entity, number, date, or yes/no). -3. yes/no answers MUST be lowercase ["yes"] or ["no"]. -4. Do NOT hallucinate. Every answer must be grounded in the provided passages. - -QUESTION REWRITE RULES (verdict = fix_question) -1. question_fixed MUST be answerable by the SAME passages and yield the SAME factual answer as the original gold facts. -2. Allowed edits: swap question word (Where -> Did / Who / What), repair grammar, complete truncation, align question word with the answer type. -3. FORBIDDEN: changing intent, injecting the answer into the question, adding facts not in the passages. -4. If you cannot satisfy these constraints, downgrade to "drop". - -DROP RULES (verdict = drop) -- answers MUST be [] and question_fixed MUST be null. - -OUTPUT FORMAT (JSON only, no markdown fence, no explanation) -{"verdict": "keep|fix_answer|fix_question|drop", "question_fixed": "..." | null, "answers": ["..."], "reasoning": "one sentence"}""" - -VERIFY_USER = """## Question -{question} - -## Original Answer (may be wrong) -{original_answer} - -## Supporting Passages -{context} - -## Task -Audit the row per the system rules. Pick exactly one verdict (keep / fix_answer / fix_question / drop), produce the multi-form answers list (or [] for drop), and write a one-sentence reasoning. If verdict=fix_question, also produce question_fixed; otherwise set it to null. -Return a single JSON object only.""" - - -LEVELS: Tuple[str, str, str] = ('easy', 'medium', 'hard') - - -def _format_context(context: Dict[str, Any]) -> str: - titles = context.get('title', []) or [] - sentences = context.get('sentences', []) or [] - lines = [] - for i, (title, sents) in enumerate(zip(titles, sentences), start=1): - if isinstance(sents, list): - body = ' '.join(s.strip() for s in sents if s and s.strip()) - else: - body = str(sents).strip() - lines.append(f'[{i}] {title}: {body}') - return '\n\n'.join(lines) - - -_JSON_RE = re.compile(r'\{[^{}]*"verdict"\s*:\s*"[^"]+"[^{}]*"answers"\s*:\s*\[.*?\][^{}]*\}', re.DOTALL) - -_VALID_VERDICTS = ('keep', 'fix_answer', 'fix_question', 'drop') - - -def _parse_response(text: str) -> Optional[Dict[str, Any]]: - text = text.strip() - if text.startswith('```'): - first_nl = text.find('\n') - last_fence = text.rfind('```') - if first_nl != -1 and last_fence > first_nl: - text = text[first_nl + 1:last_fence].strip() - try: - obj = json.loads(text) - if isinstance(obj, dict) and 'answers' in obj: - return obj - except json.JSONDecodeError: - pass - m = _JSON_RE.search(text) - if m: - try: - return json.loads(m.group(0)) - except json.JSONDecodeError: - pass - return None - - -def _validate_verdict( - verdict: Optional[str], answers: List[str], - qfix: Optional[str], original_question: str, -) -> bool: - if verdict not in _VALID_VERDICTS: - return False - if verdict == 'drop': - return not answers and qfix is None - if not answers: - return False - if verdict == 'fix_question': - return bool(qfix) and qfix.strip() != original_question.strip() - return qfix is None - - -def verify_answer( - api: OpenAI, model: str, row: Dict[str, Any], -) -> Optional[Dict[str, Any]]: - question = row['question'] - original_answer = row.get('answer', '') or '' - context_str = _format_context(row.get('context', {}) or {}) - - user_content = VERIFY_USER.format( - question=question, - original_answer=original_answer, - context=context_str) - - trajectory = { - 'messages': [ - {'role': 'system', 'content': VERIFY_SYSTEM}, - {'role': 'user', 'content': user_content}, - ] - } - sp = SamplingParams(temperature=0.1, max_tokens=512) - - for attempt in range(3): - try: - reply = api(trajectory, sp, extra_body={'enable_thinking': True}) - except Exception as exc: - sys.stderr.write(f'[verify] {row["id"]}: API error: {exc}\n') - if attempt < 2: - continue - return None - - content = reply.get('content') or '' - parsed = _parse_response(content) - if parsed: - verdict = parsed.get('verdict') - answers_raw = parsed.get('answers') - answers = ( - [str(a).strip() for a in answers_raw if str(a).strip()] - if isinstance(answers_raw, list) else []) - qfix_raw = parsed.get('question_fixed') - qfix = (qfix_raw.strip() or None) if isinstance(qfix_raw, str) else None - if _validate_verdict(verdict, answers, qfix, question): - return { - 'id': row['id'], - 'verdict': verdict, - 'question': question, - 'question_fixed': qfix, - 'original_answer': original_answer, - 'answers': answers, - 'reasoning': parsed.get('reasoning', ''), - 'level': row.get('level', ''), - 'type': row.get('type', ''), - 'context': row.get('context', {}), - 'supporting_facts': row.get('supporting_facts', {}), - } - sys.stderr.write( - f'[verify retry {attempt+1}] {row["id"]}: ' - f'parse failed, content={content[:200]!r}\n') - - sys.stderr.write(f'[verify drop] {row["id"]}: all attempts failed\n') - return None - - -def stratified_sample_with_forced( - ds, per_level: Dict[str, int], forced_ids: frozenset, seed: int, -) -> List[Dict[str, Any]]: - rng = random.Random(seed) - buckets: Dict[str, List[int]] = {lv: [] for lv in LEVELS} - forced_indices: List[int] = [] - forced_levels: Dict[str, int] = {lv: 0 for lv in LEVELS} - - for i in range(len(ds)): - row_id = ds[i]['id'] - level = (ds[i].get('level') or '').strip().lower() - if row_id in forced_ids: - forced_indices.append(i) - if level in forced_levels: - forced_levels[level] += 1 - elif level in buckets: - buckets[level].append(i) - - picked_set = set(forced_indices) - for lv in LEVELS: - need = max(0, per_level[lv] - forced_levels[lv]) - pool = [idx for idx in buckets[lv] if idx not in picked_set] - if len(pool) < need: - sys.stderr.write( - f'Warning: level={lv} has {len(pool)} available, need {need}\n') - need = len(pool) - sampled = rng.sample(pool, need) - picked_set.update(sampled) - - picked = sorted(picked_set) - rng.shuffle(picked) - return [ds[int(i)] for i in picked] - - -def select_forced_only(ds, forced_ids: frozenset, seed: int) -> List[Dict[str, Any]]: - """Pick exactly the rows whose id is in forced_ids; warn on missing.""" - indices: List[int] = [] - found: set = set() - for i in range(len(ds)): - rid = ds[i]['id'] - if rid in forced_ids: - indices.append(i) - found.add(rid) - missing = forced_ids - found - if missing: - sys.stderr.write( - f'Warning: {len(missing)} forced ids not found in dataset, ' - f'e.g. {sorted(missing)[:5]}\n') - rng = random.Random(seed) - rng.shuffle(indices) - return [ds[int(i)] for i in indices] - - -def load_done_ids(path: str) -> set: - if not os.path.exists(path): - return set() - done = set() - with open(path, 'r', encoding='utf-8') as fh: - for line in fh: - try: - obj = json.loads(line) - except json.JSONDecodeError: - continue - rid = obj.get('id') - if rid: - done.add(rid) - return done - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument('--output', required=True) - parser.add_argument('--model', required=True) - parser.add_argument('--api-key', default=os.environ.get('OPENAI_API_KEY')) - parser.add_argument('--base-url', default=os.environ.get('OPENAI_BASE_URL')) - parser.add_argument('--total', type=int, default=12000) - parser.add_argument('--easy', type=int, default=2000) - parser.add_argument('--medium', type=int, default=4000) - parser.add_argument('--hard', type=int, default=6000) - parser.add_argument('--concurrency', type=int, default=16) - parser.add_argument('--seed', type=int, default=42) - parser.add_argument('--wrong-ids', default='cookbook/rl/wrong_ids.txt') - parser.add_argument('--hf-subset', default='fullwiki') - parser.add_argument('--hf-split', default='train') - parser.add_argument( - '--only-forced', action=argparse.BooleanOptionalAction, default=False, - help='If set, re-annotate ONLY IDs in --wrong-ids; default is stratified sampling with wrong_ids force-included.') - args = parser.parse_args() - - forced_ids: frozenset = frozenset() - if args.wrong_ids and os.path.exists(args.wrong_ids): - with open(args.wrong_ids, 'r', encoding='utf-8') as fh: - forced_ids = frozenset(ln.strip() for ln in fh if ln.strip()) - sys.stderr.write(f'Forced IDs loaded: {len(forced_ids)}\n') - - if args.only_forced and not forced_ids: - raise ValueError( - f'--only-forced is set but no IDs loaded from {args.wrong_ids!r}') - - sys.stderr.write( - f'Loading hotpotqa/hotpot_qa:{args.hf_subset}:{args.hf_split}...\n') - ds = load_dataset( - 'hotpotqa/hotpot_qa', args.hf_subset, split=args.hf_split) - - if args.only_forced: - rows = select_forced_only(ds, forced_ids=forced_ids, seed=args.seed) - sys.stderr.write( - f'Selected {len(rows)} rows (only-forced mode, ' - f'requested={len(forced_ids)})\n') - else: - if args.easy + args.medium + args.hard != args.total: - raise ValueError( - f'--easy + --medium + --hard ({args.easy + args.medium + args.hard}) ' - f'must equal --total ({args.total})') - per_level = {'easy': args.easy, 'medium': args.medium, 'hard': args.hard} - rows = stratified_sample_with_forced( - ds, per_level=per_level, forced_ids=forced_ids, seed=args.seed) - sys.stderr.write( - f'Selected {len(rows)} rows (stratified per_level={per_level}, ' - f'forced={len(forced_ids)})\n') - - done = load_done_ids(args.output) - sys.stderr.write(f'Resume: {len(done)} rows already done, skipping.\n') - pending = [row for row in rows if row['id'] not in done] - sys.stderr.write(f'Pending: {len(pending)} / {len(rows)}\n') - - api = OpenAI( - model=args.model, api_key=args.api_key, base_url=args.base_url) - - write_lock = threading.Lock() - out_fh = open(args.output, 'a', encoding='utf-8') - rows_done = 0 - rows_failed = 0 - try: - with ThreadPoolExecutor(max_workers=args.concurrency) as ex: - futures = { - ex.submit(verify_answer, api, args.model, row): row['id'] - for row in pending - } - for fut in as_completed(futures): - rid = futures[fut] - try: - result = fut.result() - except Exception as exc: - sys.stderr.write(f'[row {rid}] crashed: {exc}\n') - rows_failed += 1 - continue - if result is None: - rows_failed += 1 - continue - with write_lock: - out_fh.write( - json.dumps(result, ensure_ascii=False) + '\n') - out_fh.flush() - rows_done += 1 - if rows_done % 100 == 0: - sys.stderr.write( - f'[progress] done={rows_done} ' - f'failed={rows_failed}\n') - finally: - out_fh.close() - - sys.stderr.write( - f'Done. rows_done={rows_done}, failed={rows_failed}, ' - f'total_pending={len(pending)}\n') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/legacy/rl/grpo.py b/cookbook/exp/legacy/rl/grpo.py deleted file mode 100644 index 7f43d93f4..000000000 --- a/cookbook/exp/legacy/rl/grpo.py +++ /dev/null @@ -1,787 +0,0 @@ -"""Pure GRPO training on AoPS dataset (no RAG, ablation baseline). - -Architecture (8 GPUs): - - 4 GPUs: sampler/rollout (vLLM TP=4) - - 4 GPUs: training model (FSDP) - -Pipeline per step: - 1. DataLoader yields a batch of math problems - 2. Sampler generates rollouts - 3. Reward (accuracy + format + gibberish) → GRPO advantage → model update - -Launch: - python cookbook/exp/rl/grpo.py -""" -import json -import os -import re -import random -from typing import Any, Dict, List, Tuple - -import numpy as np -import torch - -import twinkle -from twinkle import DeviceMesh, DeviceGroup, get_device_placement, get_logger -from twinkle.advantage import GRPOAdvantage -from twinkle.checkpoint_engine import CheckpointEngineManager -from twinkle.data_format import SamplingParams -from twinkle.dataloader import DataLoader -from twinkle.dataset import Dataset, DatasetMeta -from twinkle.metric import CompletionRewardMetric -from twinkle.model import TransformersModel -from twinkle.processor import InputProcessor -from twinkle.reward.base import Reward -from twinkle.sampler import vLLMSampler -from twinkle.template import Qwen3_5Template - -logger = get_logger() - -# ============================================================================ -# Configuration -# ============================================================================ -MODEL_ID = os.environ.get('MODEL_ID', 'ms://Qwen/Qwen3.5-4B') - -# GPU layout: 4 rollout + 4 train = 8 -SAMPLER_GPUS = int(os.environ.get('SAMPLER_GPUS', 4)) -MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 4)) -NUM_GPUS = SAMPLER_GPUS + MODEL_GPUS - -# Training hyperparams -NUM_GENERATIONS = int(os.environ.get('NUM_GENERATIONS', 8)) -MAX_NEW_TOKENS = int(os.environ.get('MAX_NEW_TOKENS', 32768)) -LEARNING_RATE = float(os.environ.get('LR', 1e-5)) -MAX_STEPS = int(os.environ.get('MAX_STEPS', 5000)) -BATCH_SIZE = int(os.environ.get('BATCH_SIZE', 8)) -MINI_BATCH_SIZE = int(os.environ.get('MINI_BATCH_SIZE', 8)) -MICRO_BATCH_SIZE = int(os.environ.get('MICRO_BATCH_SIZE', 8)) -GRADIENT_ACCUMULATION_STEPS = int(os.environ.get('GRADIENT_ACCUMULATION_STEPS', 1)) -SAVE_STEPS = int(os.environ.get('SAVE_STEPS', 100)) -ADV_CLIP = float(os.environ.get('ADV_CLIP', 1.0)) -LOSS_SPIKE_THRESHOLD = float(os.environ.get('LOSS_SPIKE_THRESHOLD', 10.0)) - -# Dataset -AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') -AOPS_SEED = int(os.environ.get('AOPS_SEED', 100)) - -# Output / diagnostics -OUTPUT_DIR = os.environ.get('OUTPUT_DIR', './outputs/grpo') - -# System prompt -SYSTEM_PROMPT = ( - 'You are an expert competition mathematician. ' - 'Solve the problem step by step. Put your final answer inside \\boxed{}. ' - 'For multiple-choice questions, put the option LETTER (A/B/C/D/E) inside \\boxed{}.' -) - - -# ============================================================================ -# Reward -# ============================================================================ -class AoPSAccuracyReward(Reward): - """Accuracy reward via boxed answer extraction + robust equivalence matching.""" - - @staticmethod - def extract_boxed(text: str) -> str: - idx = text.rfind('\\boxed{') - if idx == -1: - return '' - start = idx + len('\\boxed{') - depth = 1 - j = start - while j < len(text) and depth > 0: - if text[j] == '{': - depth += 1 - elif text[j] == '}': - depth -= 1 - j += 1 - if depth == 0: - return text[start:j - 1].strip() - return '' - - _MCQ_GT_RE = re.compile( - r'^\\?(?:textbf|mathbf|text|mathrm)\{?\(?([A-E])[)}\s\\]*(.*)', - re.DOTALL) - _MCQ_PAREN_RE = re.compile(r'^\(?([A-E])\)?[.:\s\\]+(.*)', re.DOTALL) - _MCQ_SINGLE_LETTER_RE = re.compile(r'^[A-E]$') - _VAR_PREFIX_RE = re.compile( - r'^(?:[a-zA-Z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)', re.DOTALL) - _EQ_RHS_RE = re.compile(r'^.+=\s*(.+)$') - - @staticmethod - def normalize_answer(ans: str) -> str: - if not ans: - return '' - s = ans.strip() - m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf)?\{?\(?([A-E])\)?\}?$', s) - if m: - return m.group(1) - s = s.replace(' ', '') - s = s.replace(r'\,', '') - s = s.replace(r'\;', '') - s = s.replace(r'\!', '') - s = re.sub(r'\\(?:text|mathrm|mathbf|textbf|operatorname)\{([^}]*)\}', r'\1', s) - s = s.replace(r'\displaystyle', '') - s = re.sub(r'\\(?:left|right)[.()\[\]|]', '', s) - s = s.replace(r'\dfrac', r'\frac') - s = s.replace(r'\tfrac', r'\frac') - # \frac shorthand without braces: \frac ab → \frac{a}{b} - s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) - s = s.strip('$').strip() - s = re.sub(r'\{[a-zA-Z]+\}$', '', s) - s = re.sub(r'\^\{\\circ\}|\^\\circ|°|\\circ', '', s) - s = re.sub(r'\\(?:quad|qquad|\s)', '', s) - s = s.replace(r'\minus{}', '-').replace(r'\minus', '-') - - def _frac_to_slash(m): - text = m.group(0) - pos = text.index('{') + 1 - depth, num_start = 1, pos - while depth > 0: - if text[pos] == '{': depth += 1 - elif text[pos] == '}': depth -= 1 - pos += 1 - numer = text[num_start:pos - 1] - pos += 1 - den_start = pos - depth = 1 - while depth > 0: - if text[pos] == '{': depth += 1 - elif text[pos] == '}': depth -= 1 - pos += 1 - denom = text[den_start:pos - 1] - return f'({numer})/({denom})' - - s = re.sub( - r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', - _frac_to_slash, s) - s = re.sub(r'(?<!\w)(\d+)/(\d+)(?!\w)', r'(\1)/(\2)', s) - return s - - @classmethod - def _strip_var_prefix(cls, s: str) -> str: - m = cls._VAR_PREFIX_RE.match(s) - return m.group(1).strip() if m else s - - @classmethod - def _extract_mcq_parts(cls, s: str): - m = cls._MCQ_GT_RE.match(s) - if m: - return m.group(1), m.group(2).strip() - m = cls._MCQ_PAREN_RE.match(s) - if m: - return m.group(1), m.group(2).strip() - m2 = re.search(r'\(?([A-E])\)?\s*$', s) - if m2 and len(s) > 3: - return m2.group(1), s[:m2.start()].strip() - return None, None - - @staticmethod - def _try_numeric_equal(a: str, b: str) -> bool: - import math - - def _try_eval(s: str): - try: - return float(s.replace('(', '').replace(')', '')) - except (ValueError, ZeroDivisionError): - pass - s_stripped = re.sub(r'[a-zA-Z]+$', '', s.replace('(', '').replace(')', '')).strip() - if s_stripped and s_stripped != s: - try: - return float(s_stripped) - except (ValueError, ZeroDivisionError): - pass - frac_re = re.compile(r'^\(([^)]+)\)/\(([^)]+)\)$') - m = frac_re.match(s) - if m: - try: - return float(m.group(1)) / float(m.group(2)) - except (ValueError, ZeroDivisionError): - pass - expr = s - expr = expr.replace('\\pi', str(math.pi)) - expr = expr.replace('\\e', str(math.e)) - expr = re.sub(r'\\sqrt\{([^}]+)\}', r'(\1)**0.5', expr) - expr = re.sub(r'\\sqrt\[3\]\{([^}]+)\}', r'(\1)**(1/3)', expr) - expr = re.sub(r'\\sqrt\[([^]]+)\]\{([^}]+)\}', r'(\2)**(1/\1)', expr) - expr = expr.replace('{', '(').replace('}', ')') - expr = expr.replace('\\cdot', '*').replace('\\times', '*') - expr = re.sub(r'(\d)\(', r'\1*(', expr) - try: - val = eval(expr, {"__builtins__": {}, "math": math, "pi": math.pi, "e": math.e}, {}) - return float(val) - except Exception: - pass - return None - - va, vb = _try_eval(a), _try_eval(b) - if va is not None and vb is not None: - return abs(va - vb) < 1e-6 * max(1, abs(va), abs(vb)) - return False - - @classmethod - def _strip_quantifiers(cls, s: str) -> str: - """Strip universal/existential quantifier wrappers.""" - s = re.sub(r'^\\forall\s*\w+\s*\\in\s*\\mathbb\s*\{?[A-Z]\}?\s*[:,]\s*', '', s) - s = re.sub(r'\s*\(\\forall[^)]*\)\s*$', '', s) - s = re.sub(r'\s*\(for\s+all[^)]*\)\s*$', '', s, flags=re.IGNORECASE) - return s.strip() - - @classmethod - def _try_param_rename(cls, a: str, b: str) -> bool: - """Check if a == b up to consistent single free-parameter rename (ax vs cx).""" - if not a or not b or len(a) != len(b) or len(a) > 80: - return False - diffs = [(i, a[i], b[i]) for i in range(len(a)) if a[i] != b[i]] - if not diffs: - return False - src_chars = set(d[1] for d in diffs) - dst_chars = set(d[2] for d in diffs) - if len(src_chars) == 1 and len(dst_chars) == 1: - src, dst = src_chars.pop(), dst_chars.pop() - if src.isalpha() and dst.isalpha(): - return a.replace(src, dst) == b - return False - - @classmethod - def _normalize_tuple(cls, s: str) -> str: - # Strip set-builder conditions: \mid ... or | ... - s = re.sub(r'\\mid.*$', '', s) - s = re.sub(r'\|[^,]*$', '', s) - return re.sub(r'[\s()\[\]{}\\]', '', s) - - @classmethod - def _try_sympy_equal(cls, a: str, b: str) -> bool: - try: - from sympy.parsing.latex import parse_latex - from sympy import simplify, nsimplify - expr_a = parse_latex(a) - expr_b = parse_latex(b) - diff = simplify(nsimplify(expr_a - expr_b)) - return diff == 0 - except Exception: - return False - - @classmethod - def answers_match(cls, predicted: str, reference: str) -> bool: - if not predicted or not reference: - return False - - norm_p = cls.normalize_answer(predicted) - norm_r = cls.normalize_answer(reference) - - if norm_p == norm_r: - return True - if norm_p.lower() == norm_r.lower(): - return True - if cls._try_numeric_equal(norm_p, norm_r): - return True - - stripped_p = cls.normalize_answer(cls._strip_var_prefix(predicted)) - stripped_r = cls.normalize_answer(cls._strip_var_prefix(reference)) - if stripped_p and stripped_r and stripped_p == stripped_r: - return True - if stripped_p and stripped_r and cls._try_numeric_equal(stripped_p, stripped_r): - return True - - ref_letter, ref_value = cls._extract_mcq_parts(reference) - if ref_letter: - if norm_p == ref_letter or predicted.strip().upper() == ref_letter: - return True - if ref_value: - norm_ref_val = cls.normalize_answer(ref_value) - if norm_p == norm_ref_val or cls._try_numeric_equal(norm_p, norm_ref_val): - return True - pred_letter, pred_value = cls._extract_mcq_parts(predicted) - if pred_letter: - if norm_r == pred_letter or reference.strip().upper() == pred_letter: - return True - if pred_value: - norm_pred_val = cls.normalize_answer(pred_value) - if norm_r == norm_pred_val or cls._try_numeric_equal(norm_r, norm_pred_val): - return True - if cls._MCQ_SINGLE_LETTER_RE.match(reference.strip()): - if cls._MCQ_SINGLE_LETTER_RE.match(predicted.strip().upper()): - return predicted.strip().upper() == reference.strip().upper() - - tuple_p = cls._normalize_tuple(norm_p) - tuple_r = cls._normalize_tuple(norm_r) - if ',' in tuple_p and tuple_p == tuple_r: - return True - if stripped_p and stripped_r: - tuple_sp = cls._normalize_tuple(stripped_p) - tuple_sr = cls._normalize_tuple(stripped_r) - if ',' in tuple_sp and tuple_sp == tuple_sr: - return True - - if '=' in norm_r and '=' not in norm_p: - parts = norm_r.split('=') - for part in parts: - part = part.strip() - if part == norm_p or cls._try_numeric_equal(part, norm_p): - return True - if '=' in norm_p and '=' not in norm_r: - parts = norm_p.split('=') - for part in parts: - part = part.strip() - if part == norm_r or cls._try_numeric_equal(part, norm_r): - return True - if '=' in norm_p and '=' in norm_r: - pp = [x.strip() for x in norm_p.split('=')] - rp = [x.strip() for x in norm_r.split('=')] - if set(pp) == set(rp): - return True - - def _sort_factors(s): - tokens = re.findall(r'\\?[a-zA-Z]+\{[^}]*\}|\\?[a-zA-Z]+|\d+|[^a-zA-Z\d\\{}]', s) - return ''.join(sorted(tokens)) - if _sort_factors(norm_p) == _sort_factors(norm_r): - return True - - if cls._try_sympy_equal(predicted, reference): - return True - - # --- Strategy 9b: quantifier stripping + param rename --- - q_stripped_r = cls._strip_quantifiers(reference) - q_stripped_p = cls._strip_quantifiers(predicted) - if q_stripped_r != reference or q_stripped_p != predicted: - norm_qr = cls.normalize_answer(cls._strip_var_prefix(q_stripped_r)) - norm_qp = cls.normalize_answer(cls._strip_var_prefix(q_stripped_p)) - if norm_qr and norm_qp: - if norm_qr == norm_qp: - return True - if cls._try_param_rename(norm_qr, norm_qp): - return True - - # --- Strategy 10: param rename on var-prefix-stripped forms --- - if stripped_p and stripped_r and cls._try_param_rename(stripped_p, stripped_r): - return True - - return False - - def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: - rewards = [] - for traj in trajectories: - messages = traj.get('messages', []) - completion = '' - for msg in reversed(messages): - if msg.get('role') == 'assistant': - completion = msg.get('content', '') - break - user_data = traj.get('user_data') or [] - gt = '' - for item in user_data: - if item[0] == 'ground_truth': - gt = item[1] - break - predicted = self.extract_boxed(completion) - correct = self.answers_match(predicted, gt) - rewards.append(1.0 if correct else 0.0) - return rewards - - -class FormatReward(Reward): - """Reward for having \\boxed{} in the output.""" - - def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: - rewards = [] - for traj in trajectories: - messages = traj.get('messages', []) - completion = '' - for msg in reversed(messages): - if msg.get('role') == 'assistant': - completion = msg.get('content', '') - break - has_boxed = '\\boxed{' in completion - rewards.append(0.5 if has_boxed else 0.0) - return rewards - - -class GibberishPenalty(Reward): - """Negative reward for degenerate outputs (gibberish/random unicode tail).""" - - TAIL_CHARS = 400 - GIBBERISH_THRESHOLD = 0.20 - - @classmethod - def is_gibberish(cls, text: str) -> bool: - if not text: - return False - tail = text[-cls.TAIL_CHARS:] if len(text) > cls.TAIL_CHARS else text - non_math_non_ascii = 0 - for c in tail: - code = ord(c) - if code > 127 and not (0x4e00 <= code <= 0x9fff): - non_math_non_ascii += 1 - return non_math_non_ascii > len(tail) * cls.GIBBERISH_THRESHOLD - - def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: - rewards = [] - for traj in trajectories: - messages = traj.get('messages', []) - completion = '' - for msg in reversed(messages): - if msg.get('role') == 'assistant': - completion = msg.get('content', '') - break - rewards.append(-0.5 if self.is_gibberish(completion) else 0.0) - return rewards - - -def compute_rewards(trajectories: List[Dict[str, Any]] - ) -> Tuple[List[float], List[float], List[float]]: - acc_fn = AoPSAccuracyReward() - fmt_fn = FormatReward() - gib_fn = GibberishPenalty() - acc = acc_fn(trajectories) - fmt = fmt_fn(trajectories) - gib = gib_fn(trajectories) - total = [a + f + g for a, f, g in zip(acc, fmt, gib)] - return total, fmt, acc - - -# ============================================================================ -# Dataset: AoPS boxed problems -# ============================================================================ -def create_aops_dataset(): - """Load AoPS and create GRPO-style dataset (prompt only, with ground_truth in user_data).""" - from modelscope import MsDataset - from twinkle.data_format import Message, Trajectory - - ds = MsDataset.load(AOPS_DATASET_ID, split='train', - download_mode='reuse_dataset_if_exists') - rows = [] - for row in ds: - if not row['metadata'].get('boxed'): - continue - ref = AoPSAccuracyReward.extract_boxed(row['solution']) - if not ref: - continue - rows.append({'problem': row['problem'], 'ground_truth': ref}) - - logger.info(f'[aops] loaded {len(rows)} boxed problems') - rng = random.Random(AOPS_SEED) - rng.shuffle(rows) - - trajectories = [] - for r in rows: - traj = Trajectory( - messages=[ - Message(role='system', content=SYSTEM_PROMPT), - Message(role='user', content=r['problem']), - ], - user_data=[('ground_truth', r['ground_truth'])], - ) - trajectories.append(traj) - - data_meta = DatasetMeta(data=trajectories) - dataset = Dataset(data_meta) - dataset.set_template('Qwen3_5Template', model_id=MODEL_ID, - max_length=16384, truncation_strategy='delete', - enable_thinking=True) - dataset.encode(add_generation_prompt=True) - return dataset - - -# ============================================================================ -# Main -# ============================================================================ -def main(): - sampler_start = 0 - model_start = sampler_start + SAMPLER_GPUS - - device_groups = [ - DeviceGroup(name='sampler', ranks=list(range(sampler_start, model_start)), - device_type='GPU'), - DeviceGroup(name='model', ranks=list(range(model_start, NUM_GPUS)), - device_type='GPU'), - ] - - model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, fsdp_size=MODEL_GPUS, ulysses_size=2) - sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, dp_size=SAMPLER_GPUS) - - twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, - groups=device_groups, lazy_collect=False) - - # -- Training model (full-parameter) -- - model = TransformersModel( - model_id=MODEL_ID, device_mesh=model_mesh, remote_group='model') - model.set_optimizer('AdamW', lr=LEARNING_RATE) - model.set_lr_scheduler('CosineAnnealingLR', T_max=MAX_STEPS, eta_min=0) - model.set_loss('GSPOLoss', epsilon=0.2, epsilon_high=0.28, beta=0.04) - model.set_processor(InputProcessor) - model.set_template('Qwen3_5Template', model_id=MODEL_ID, - enable_thinking=True, max_length=32768) - - # -- Rollout sampler -- - sampler = vLLMSampler( - model_id=MODEL_ID, - engine_args={ - 'gpu_memory_utilization': 0.8, - 'max_model_len': 32768, - }, - device_mesh=sampler_mesh, - remote_group='sampler', - ) - sampler.set_template('Qwen3_5Template', model_id=MODEL_ID, - enable_thinking=True, max_length=32768) - - # -- Checkpoint & DataLoader -- - ckpt_manager = CheckpointEngineManager(model=model, sampler=sampler) - - GLOBAL_BATCH_SIZE = BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS - dataloader = DataLoader( - dataset=create_aops_dataset, - batch_size=GLOBAL_BATCH_SIZE, - min_batch_size=GLOBAL_BATCH_SIZE, - device_mesh=model_mesh, - remote_group='model', - ) - - advantage_fn = GRPOAdvantage() - metrics = CompletionRewardMetric() - sampling_params = SamplingParams( - max_tokens=MAX_NEW_TOKENS, num_samples=1, logprobs=1, - temperature=1.0, top_p=0.95) - - optim_step = 0 - logger.info('Starting pure GRPO training (no RAG)') - logger.info(get_device_placement()) - - # -- Diagnostics -- - os.makedirs(OUTPUT_DIR, exist_ok=True) - diag_path = os.path.join(OUTPUT_DIR, 'diagnostics.jsonl') - diag_f = open(diag_path, 'w', encoding='utf-8') - logger.info(f'[diag] diagnostics → {diag_path}') - - def _content_to_str(content): - if isinstance(content, str): - return content - if isinstance(content, list): - return ''.join( - b.get('text', '') if isinstance(b, dict) else str(b) - for b in content) - return str(content) - - for batch in dataloader: - if optim_step >= MAX_STEPS: - break - - metrics.reset() - - # Build prompts (direct, no RAG) - prompts = [] - for item in batch: - msgs = item.get('messages', []) - prob = '' - for m in msgs: - if m.get('role') == 'user': - prob = m.get('content', '') - if isinstance(prob, list): - prob = ''.join(p.get('text', '') for p in prob if isinstance(p, dict)) - break - ud = item.get('user_data', []) - gt = '' - for pair in ud: - if pair[0] == 'ground_truth': - gt = pair[1] - break - prompts.append({ - 'messages': [ - {'role': 'system', 'content': SYSTEM_PROMPT}, - {'role': 'user', 'content': prob}, - ], - 'user_data': [('ground_truth', gt)], - }) - - # Expand for NUM_GENERATIONS and sample - expand_prompts = [] - for prompt in prompts: - expand_prompts.extend([prompt] * NUM_GENERATIONS) - - ckpt_manager.sync_weights(merge_and_sync=False) - sampler.reset_prefix_cache() - - sample_responses = sampler.sample(expand_prompts, sampling_params) - - # Collect rollouts - all_input_data: List[Dict[str, Any]] = [] - all_old_logps: List[List[float]] = [] - all_completion_lengths: List[int] = [] - - for sample_response in sample_responses: - for sequence in sample_response.sequences: - all_input_data.append(sequence.new_input_feature) - all_old_logps.append([logprob[0][1] for logprob in sequence.logprobs]) - all_completion_lengths.append(len(sequence.tokens)) - - # Rewards - total_rewards, format_rewards, accuracy_rewards = compute_rewards(all_input_data) - - # Zero out rewards for rollouts that hit the max_tokens ceiling - max_len_threshold = int(MAX_NEW_TOKENS * 0.95) - for i in range(len(all_input_data)): - if all_completion_lengths[i] >= max_len_threshold: - total_rewards[i] = 0.0 - accuracy_rewards[i] = 0.0 - format_rewards[i] = 0.0 - - # Per-step reward summary - n_correct = sum(1 for a in accuracy_rewards if a > 0) - diag_f.write(json.dumps({ - 'step': optim_step, 'type': 'reward_summary', - 'n_samples': len(accuracy_rewards), - 'accuracy': n_correct / len(accuracy_rewards) if accuracy_rewards else 0, - 'mean_reward': sum(total_rewards) / len(total_rewards) if total_rewards else 0, - }, ensure_ascii=False) + '\n') - - metrics.accumulate( - completion_lengths=all_completion_lengths, - rewards={ - 'total': total_rewards, - 'format': format_rewards, - 'accuracy': accuracy_rewards, - }, - ) - - # GRPO advantage - advantages = advantage_fn( - total_rewards, num_generations=NUM_GENERATIONS, scale='group').tolist() - if ADV_CLIP > 0: - advantages = [max(-ADV_CLIP, min(ADV_CLIP, a)) for a in advantages] - - # Log rollout responses - _extract_boxed = AoPSAccuracyReward.extract_boxed - for ridx, traj in enumerate(all_input_data): - msgs = traj.get('messages', []) - assistant_text = _content_to_str(next( - (m['content'] for m in reversed(msgs) if m.get('role') == 'assistant'), '')) - user_text = _content_to_str(next( - (m['content'] for m in msgs if m.get('role') == 'user'), '')) - user_data = traj.get('user_data') or [] - gt = next((v for k, v in user_data if k == 'ground_truth'), '') - problem_idx = ridx // NUM_GENERATIONS - grp_start = problem_idx * NUM_GENERATIONS - grp_end = grp_start + NUM_GENERATIONS - grp_acc = sum(accuracy_rewards[grp_start:grp_end]) / NUM_GENERATIONS - - diag_f.write(json.dumps({ - 'step': optim_step, 'type': 'rollout', - 'idx': ridx, - 'problem_idx': problem_idx, - 'problem': user_text, - 'response': assistant_text, - 'ground_truth': gt, - 'predicted': _extract_boxed(assistant_text), - 'reward': total_rewards[ridx], - 'accuracy_reward': accuracy_rewards[ridx], - 'format_reward': format_rewards[ridx], - 'advantage': advantages[ridx], - 'completion_length': all_completion_lengths[ridx], - 'group_accuracy': grp_acc, - }, ensure_ascii=False) + '\n') - - diag_f.flush() - - # Filter out low-signal problem groups (DAPO-style dynamic sampling) - # Skip groups where accuracy is too low (<0.1) or too high (>0.9) - # to avoid gradient dominated by gibberish/format noise or no learning signal. - filtered_inputs, filtered_old_logps, filtered_advantages = [], [], [] - for g in range(BATCH_SIZE): - g_start = g * NUM_GENERATIONS - g_end = g_start + NUM_GENERATIONS - grp_adv = advantages[g_start:g_end] - if all(abs(a) < 1e-8 for a in grp_adv): - continue - grp_acc_rate = sum(accuracy_rewards[g_start:g_end]) / NUM_GENERATIONS - if grp_acc_rate < 0.2 or grp_acc_rate > 0.8: - continue - filtered_inputs.extend(all_input_data[g_start:g_end]) - filtered_old_logps.extend(all_old_logps[g_start:g_end]) - filtered_advantages.extend(grp_adv) - - # Mini-batch training with gradient accumulation - # Process MICRO_BATCH_SIZE samples per forward, accumulate grad_accum_steps - # times before one optimizer step. clip_grad_norm normalizes by accumulated - # num_tokens, ensuring mathematical equivalence with larger batch forward. - total_completions = len(filtered_inputs) - if total_completions == 0: - logger.info(f'[Step {optim_step}] all groups filtered (uniform rewards), skip training') - continue - - grad_accum_steps = MINI_BATCH_SIZE // MICRO_BATCH_SIZE - accum_count = 0 - for mb_start in range(0, total_completions, MICRO_BATCH_SIZE): - mb_end = min(mb_start + MICRO_BATCH_SIZE, total_completions) - mb_inputs = filtered_inputs[mb_start:mb_end] - mb_old_logps = filtered_old_logps[mb_start:mb_end] - mb_advantages = filtered_advantages[mb_start:mb_end] - - outputs = model.forward_backward( - inputs=mb_inputs, - old_logps=mb_old_logps, - ref_logps=mb_old_logps, - advantages=mb_advantages, - ) - accum_count += 1 - - if accum_count % grad_accum_steps == 0: - skip_step = False - try: - loss_val = outputs.get('loss', None) - if loss_val is not None: - if hasattr(loss_val, 'item'): - loss_val = loss_val.item() - if loss_val > LOSS_SPIKE_THRESHOLD: - skip_step = True - logger.warning( - f'[Step {optim_step}] Loss spike: {loss_val:.4f} > ' - f'{LOSS_SPIKE_THRESHOLD}, skipping update') - except Exception: - pass - - if skip_step: - model.zero_grad() - else: - model.clip_grad_and_step() - optim_step += 1 - - if optim_step >= MAX_STEPS: - break - if optim_step % SAVE_STEPS == 0: - model.save(f'grpo-checkpoint-{optim_step}') - - # Flush remaining accumulated gradients (incomplete window at tail) - if accum_count % grad_accum_steps != 0: - skip_step = False - try: - loss_val = outputs.get('loss', None) - if loss_val is not None: - if hasattr(loss_val, 'item'): - loss_val = loss_val.item() - if loss_val > LOSS_SPIKE_THRESHOLD: - skip_step = True - logger.warning( - f'[Step {optim_step}] Loss spike (tail): {loss_val:.4f} > ' - f'{LOSS_SPIKE_THRESHOLD}, skipping update') - except Exception: - pass - - if skip_step: - model.zero_grad() - else: - model.clip_grad_and_step() - optim_step += 1 - - log_dict = metrics.calculate() - log_dict.update(model.calculate_metric(is_training=True)) - metrics.reset() - logger.info(f'[Step {optim_step}/{MAX_STEPS}] {log_dict}') - - diag_f.close() - logger.info(f'Training completed. optim_steps={optim_step}') - model.save('grpo-final') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/legacy/rl/rag_hint_grpo.py b/cookbook/exp/legacy/rl/rag_hint_grpo.py deleted file mode 100644 index b35b3706b..000000000 --- a/cookbook/exp/legacy/rl/rag_hint_grpo.py +++ /dev/null @@ -1,1480 +0,0 @@ -"""RAG-hint GRPO training: retrieve thinking traces and condense as hints for RL. - -Architecture (8 GPUs): - - 1 GPU: condenser (vLLM, compress retrieved traces) - - 1 GPU: embedding model (encode queries for retrieval) - - 4 GPUs: sampler/rollout (vLLM TP=4) - - 2 GPUs: training model (FSDP/DP) - -Pipeline per step: - 1. DataLoader yields a batch of math problems - 2. [Async] Embedding model encodes problems → retrieve from LanceDB → condenser compresses - 3. Build RAG-hint prompts (one-shot in system, with analysis prefix) - 4. Sampler generates rollouts (response starts with forced analysis prefix) - 5. Reward (accuracy + format) → GRPO advantage → model update - -Launch: - python cookbook/exp/rl/rag_hint_grpo.py -""" -import json -import os -import re -import random -import threading -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Any, Dict, List, Optional, Tuple - -import numpy as np -import torch - -import twinkle -from twinkle import DeviceMesh, DeviceGroup, get_device_placement, get_logger -from twinkle.advantage import GRPOAdvantage -from twinkle.checkpoint_engine import CheckpointEngineManager -from twinkle.data_format import SamplingParams -from twinkle.dataloader import DataLoader -from twinkle.dataset import Dataset, DatasetMeta -from twinkle.loss import InfonceLoss -from twinkle.metric import CompletionRewardMetric -from twinkle.model import TransformersModel -from twinkle.processor import InputProcessor -from twinkle.reward.base import Reward -from twinkle.sampler import vLLMSampler -from twinkle.template import Qwen3_5Template -from twinkle_agentic.protocol.openai import OpenAI as OpenAIClient - -logger = get_logger() - -# ============================================================================ -# Configuration -# ============================================================================ -MODEL_ID = os.environ.get('MODEL_ID', 'ms://Qwen/Qwen3.5-4B') - -# GPU layout: 1 condenser + 1 embedding + 4 rollout + 2 train = 8 -CONDENSER_GPUS = int(os.environ.get('CONDENSER_GPUS', 1)) -EMB_GPUS = int(os.environ.get('EMB_GPUS', 1)) -SAMPLER_GPUS = int(os.environ.get('SAMPLER_GPUS', 2)) -MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 4)) -NUM_GPUS = CONDENSER_GPUS + EMB_GPUS + SAMPLER_GPUS + MODEL_GPUS - -# Training hyperparams -NUM_GENERATIONS = int(os.environ.get('NUM_GENERATIONS', 8)) -MAX_NEW_TOKENS = int(os.environ.get('MAX_NEW_TOKENS', 32768)) -LEARNING_RATE = float(os.environ.get('LR', 1e-5)) -MAX_STEPS = int(os.environ.get('MAX_STEPS', 5000)) -BATCH_SIZE = int(os.environ.get('BATCH_SIZE', 8)) -MINI_BATCH_SIZE = int(os.environ.get('MINI_BATCH_SIZE', 8)) -MICRO_BATCH_SIZE = int(os.environ.get('MICRO_BATCH_SIZE', 8)) -GRADIENT_ACCUMULATION_STEPS = int(os.environ.get('GRADIENT_ACCUMULATION_STEPS', 1)) -SAVE_STEPS = int(os.environ.get('SAVE_STEPS', 100)) -ADV_CLIP = float(os.environ.get('ADV_CLIP', 1.0)) -LOSS_SPIKE_THRESHOLD = float(os.environ.get('LOSS_SPIKE_THRESHOLD', 10.0)) - -# RAG config -DB_PATH = os.environ.get('DB_PATH', './output.oldemb/thinking_rag/lance.db') -DB_TABLE = os.environ.get('DB_TABLE', 'thinking_traces') -TOP_K = int(os.environ.get('TOP_K', 2)) -SIM_THRESHOLD = float(os.environ.get('SIM_THRESHOLD', 0.75)) -MAX_TRACE_LEN = int(os.environ.get('MAX_TRACE_LEN', 8192)) -EMBED_MODEL_ID = os.environ.get( - 'EMBED_MODEL_ID', 'output.oldemb/embedding_full_transformers/last-checkpoint') -EMBED_MAX_LENGTH = int(os.environ.get('EMBED_MAX_LENGTH', 32000)) - -# Condenser config -CONDENSE_MODEL_ID = os.environ.get('CONDENSE_MODEL_ID', 'ms://twinkle-kit/Qwen3.5-4B-CM-v2') -CONDENSE_API_KEY = os.environ.get('COMPRESS_API_KEY', '') -CONDENSE_BASE_URL = os.environ.get('COMPRESS_BASE_URL', 'https://dashscope.aliyuncs.com/compatible-mode/v1') -CONDENSE_API_MODEL = os.environ.get('COMPRESS_MODEL', 'qwen3.7-max') -CONDENSE_API_CONCURRENCY = int(os.environ.get('API_CONCURRENCY', 16)) -CONDENSE_TEMPERATURE = 0.2 -CONDENSE_MAX_TOKENS = 8192 - -# Dataset -AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') -AOPS_SEED = int(os.environ.get('AOPS_SEED', 100)) - -# Decontamination & RAG fallback -DECONTAM_THRESHOLD = float(os.environ.get('DECONTAM_THRESHOLD', 0.20)) -RAG_FALLBACK_SIM = float(os.environ.get('RAG_FALLBACK_SIM', 0.60)) - -# Output / diagnostics -OUTPUT_DIR = os.environ.get('OUTPUT_DIR', './outputs/rag_hint_grpo') - -# Forced analysis prefix appended at the start of assistant response -ANALYSIS_PREFIX = '' - -# Fixed opening inside <hint> block — model must produce this EXACT prefix -HINT_REQUIRED_PREFIX = "Let's analyze the RAG example step by step." - -# Hint analysis config (API pre-analysis) -HINT_ANALYSIS_MAX_TOKENS = int(os.environ.get('HINT_ANALYSIS_MAX_TOKENS', 400)) -HINT_ANALYSIS_TEMPERATURE = 0.3 - -# ============================================================================ -# Condenser prompt (strategy-level extraction) -# ============================================================================ -COMPRESS_SYSTEM = """\ -You are a reasoning-trace condenser. Given a verbose reasoning trace, \ -extract the TRANSFERABLE KNOWLEDGE as an EXECUTABLE SOLUTION SKELETON \ -that would help a reader solve SIMILAR problems in the same domain. - -The reader will apply this knowledge to a DIFFERENT problem, so focus on what transfers. \ -NEVER output the final answer or conclusion of the original problem. \ -NEVER include problem-specific numeric results. - -Principles: -1. OUTPUT AN EXECUTABLE STEP CHAIN: numbered steps that a solver can directly follow. \ -Each step should state WHAT/WHY/HOW (with the formula/technique), not just name the concept. -2. INCLUDE FULL FORMULAS: theorems, identities — state each with COMPLETE MATHEMATICAL EXPRESSION. -3. STATE APPLICABILITY: what structural features signal that this approach works. -4. PRESERVE KEY INSIGHTS: non-obvious ideas that make the approach work. -5. REMOVE: problem-specific numeric calculations, final answers, dead-end explorations, hesitations. -6. FORMAT: Start with "Applicability:" one-line, then numbered steps. Keep concise. -7. NO meta-commentary. NO preamble. NO final answer. -""" - -COMPRESS_USER = ( - '## Reader Problem (context only — do NOT solve it)\n{query}\n\n' - '## Reasoning Trace to Condense\n{text}') - -# ============================================================================ -# RAG system prompt template (few-shot in system) -# ============================================================================ -SYSTEM_WITH_RAG_HEADER = ( - 'You are an expert competition mathematician. ' - 'Below are condensed reasoning examples from similar problems.\n\n' - '## Output Format (STRICT)\n' - 'Your response MUST begin with a <hint> block as the VERY FIRST content. ' - 'Do NOT output any text before <hint>.\n\n' - 'The <hint> block MUST start with EXACTLY this sentence (copy verbatim):\n' - '"Let\'s analyze the RAG example step by step."\n\n' - 'Then continue your analysis:\n' - '- Walk through each example\'s methodology and identify which steps, ' - 'formulas, and concepts are CORRECT and APPLICABLE to the current problem.\n' - '- Identify which parts are WRONG, IRRELEVANT, MISLEADING, or based on ' - 'assumptions that do NOT hold for this problem.\n' - '- End with a one-line verdict: "Useful: ..." and "Discard: ..."\n\n' - 'Example format:\n' - '<hint>\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' - '</hint>\n\n' - 'After the </hint> block, solve the actual problem step by step using ONLY ' - 'the validated useful parts. Put your final answer inside \\boxed{}. ' - 'For multiple-choice questions, put the option LETTER (A/B/C/D/E) inside \\boxed{}.\n\n' -) - -# System prompt for pre-analyzed RAG (hint analysis done by API, model just solves) -_PREANALYSIS_BEFORE = ( - 'You are an expert competition mathematician.\n\n' - '## RAG Analysis (pre-computed)\n' -) -_PREANALYSIS_AFTER = ( - '\n\n## Instructions\n' - 'Use the useful methods/formulas identified above to solve the problem. ' - 'Ignore anything marked as irrelevant. ' - 'Solve step by step and put your final answer inside \\boxed{}. ' - 'For multiple-choice questions, put the option LETTER (A/B/C/D/E) inside \\boxed{}.' -) - - -def build_preanalysis_system(hint_analysis: str) -> str: - """Build system prompt with pre-analyzed hint. Uses concatenation to avoid .format() issues with math braces.""" - return _PREANALYSIS_BEFORE + hint_analysis + _PREANALYSIS_AFTER - -# API prompt for hint analysis generation -HINT_ANALYSIS_SYSTEM = ( - 'You are a mathematical methodology analyst. ' - 'Given a target problem and a condensed reasoning trace from a SIMILAR (but different) problem, ' - 'analyze which methods, formulas, and techniques from the trace are APPLICABLE to the target problem ' - 'and which are IRRELEVANT or MISLEADING.\n\n' - 'Output format (strict):\n' - '- Useful: [list specific methods/formulas/techniques that transfer to the target]\n' - '- Discard: [list parts that are irrelevant or would mislead]\n' - '- Key insight: [one sentence on how to apply the useful parts]\n\n' - 'Rules:\n' - '1. Be concise — at most 200 words total.\n' - '2. Focus ONLY on transferable methodology, never solve the target problem.\n' - '3. Never output the answer to either problem.\n' - '4. If the trace is entirely irrelevant, say "Useful: None. Discard: All."' -) - -HINT_ANALYSIS_USER = ( - '## Target Problem\n{query}\n\n' - '## Condensed Trace (from similar problem)\n{thinking}' -) - -EXAMPLE_TEMPLATE = ( - '--- Example {idx} ---\n' - 'Problem: {example_query}\n' - 'Methodology:\n{example_thinking}\n' - '--- End Example {idx} ---\n' -) - -SYSTEM_DIRECT = ( - 'You are an expert competition mathematician. ' - 'Solve the problem step by step. Put your final answer inside \\boxed{}. ' - 'For multiple-choice questions, put the option LETTER (A/B/C/D/E) inside \\boxed{}.' -) - - -# ============================================================================ -# Condenser utilities -# ============================================================================ -_api_semaphore = threading.Semaphore(CONDENSE_API_CONCURRENCY) - - -def _api_condense_single(api_client: OpenAIClient, messages: List[Dict]) -> Optional[str]: - _api_semaphore.acquire() - try: - trajectory = {'messages': messages} - sp = SamplingParams(temperature=CONDENSE_TEMPERATURE, max_tokens=CONDENSE_MAX_TOKENS) - reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) - content = (reply.get('content') or '').strip() - if not content: - return None - return content - except Exception as exc: - logger.warning(f'[condense-api] error: {exc}') - return None - finally: - _api_semaphore.release() - - -def _api_hint_analysis_batch( - api_client: OpenAIClient, - problems: List[str], - condensed_examples: List[List[Dict[str, str]]], -) -> List[Optional[str]]: - """Call API to pre-analyze RAG relevance for each problem. ~300 tokens per call.""" - results: List[Optional[str]] = [None] * len(problems) - tasks = [] # (idx, messages) - for i, prob in enumerate(problems): - if not condensed_examples[i]: - continue - # Merge all condensed traces into one block - traces = [] - for ex in condensed_examples[i]: - traces.append(ex.get('thinking', '')) - merged_thinking = '\n---\n'.join(traces) - user_msg = HINT_ANALYSIS_USER.format(query=prob, thinking=merged_thinking) - msgs = [ - {'role': 'system', 'content': HINT_ANALYSIS_SYSTEM}, - {'role': 'user', 'content': user_msg}, - ] - tasks.append((i, msgs)) - - if not tasks: - return results - - def _call_one(idx, msgs): - _api_semaphore.acquire() - try: - trajectory = {'messages': msgs} - sp = SamplingParams( - temperature=HINT_ANALYSIS_TEMPERATURE, - max_tokens=HINT_ANALYSIS_MAX_TOKENS) - reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) - content = (reply.get('content') or '').strip() - return idx, content if content else None - except Exception as exc: - logger.warning(f'[hint-analysis] error for idx={idx}: {exc}') - return idx, None - finally: - _api_semaphore.release() - - with ThreadPoolExecutor(max_workers=min(len(tasks), CONDENSE_API_CONCURRENCY)) as pool: - futs = [pool.submit(_call_one, idx, msgs) for idx, msgs in tasks] - for fut in as_completed(futs): - idx, analysis = fut.result() - results[idx] = analysis - - n_success = sum(1 for r in results if r) - logger.info(f'[hint-analysis] completed {n_success}/{len(tasks)} analyses') - return results - - -# ============================================================================ -# Embedding & Retrieval -# ============================================================================ -def _normalize_for_ngram(text: str) -> str: - """Normalize text for n-gram comparison: strip LaTeX markup, lowercase.""" - text = text.lower() - text = re.sub(r'\$+', '', text) - text = re.sub(r'\\[a-z]+\{([^}]*)\}', r'\1', text) - text = re.sub(r'\\[a-z]+', ' ', text) - text = re.sub(r'[{}\\^_$]', '', text) - text = re.sub(r'\s+', ' ', text).strip() - return text - - -def _ngram_jaccard(text_a: str, text_b: str, n: int = 13) -> float: - """13-gram character-level Jaccard similarity for decontamination.""" - a = _normalize_for_ngram(text_a) - b = _normalize_for_ngram(text_b) - if len(a) < n or len(b) < n: - return 0.0 - grams_a = set(a[i:i + n] for i in range(len(a) - n + 1)) - grams_b = set(b[i:i + n] for i in range(len(b) - n + 1)) - if not grams_a or not grams_b: - return 0.0 - return len(grams_a & grams_b) / len(grams_a | grams_b) - - -def _wrap_anchor(text: str) -> List[Dict[str, str]]: - return [ - {'role': 'user', 'content': text}, - {'role': 'assistant', 'content': 'Match the correct response here.'}, - ] - - -_DECONTAM_JUDGE_PROMPT = ( - 'We are building a RAG-augmented math training system. Problem A is the test ' - 'question; Problem B was retrieved from a knowledge base.\n' - 'Answer YES only if A and B are essentially the SAME specific problem — ' - 'i.e. solving B directly gives you A\'s answer (just different wording/notation/' - 'format/negation).\n' - 'Answer NO if they merely share the same method/topic but have different ' - 'specific values, equations, or geometric configurations — learning B\'s ' - 'approach still requires independent work to solve A.\n' - 'Problem A: {prob_a}\n' - 'Problem B: {prob_b}\n' - 'Answer only YES or NO.' -) - - -def _llm_judge_same_problem( - api_client, pairs: List[tuple], -) -> List[bool]: - """Batch LLM judge: are (problem_a, problem_b) the same problem? - - Each pair text is truncated to 200 chars to keep latency low. - Returns list of bools (True = same problem = should filter). - """ - if not pairs or not api_client: - return [False] * len(pairs) - - results = [False] * len(pairs) - - def _judge_one(idx, pa, pb): - prompt = _DECONTAM_JUDGE_PROMPT.format(prob_a=pa, prob_b=pb) - msgs = [{'role': 'user', 'content': prompt}] - try: - trajectory = {'messages': msgs} - sp = SamplingParams(temperature=0.1, max_tokens=8) - reply = api_client(trajectory, sp, extra_body={'enable_thinking': False}) - answer = (reply.get('content') or '').strip().upper() - return idx, 'YES' in answer - except Exception: - return idx, False - - with ThreadPoolExecutor(max_workers=min(len(pairs), CONDENSE_API_CONCURRENCY)) as pool: - futs = [pool.submit(_judge_one, i, pa, pb) for i, (pa, pb) in enumerate(pairs)] - for fut in as_completed(futs): - idx, is_same = fut.result() - results[idx] = is_same - return results - - -def get_embeddings(model: TransformersModel, template: Qwen3_5Template, - texts: List[str], dp_size: int) -> np.ndarray: - if not texts: - return np.zeros((0,), dtype=np.float32) - n = len(texts) - pad_n = (-n) % dp_size - padded = list(texts) + [' '] * pad_n if pad_n else list(texts) - features = [] - for t in padded: - feat = template.encode({'messages': _wrap_anchor(t or ' ')}) - feat['labels'] = [1] - features.append(feat) - out = model.forward_only(inputs=features, task='embedding', return_logits=True) - emb = out['embeddings'] - if isinstance(emb, torch.Tensor): - emb = emb.detach().to(torch.float32).cpu().numpy() - emb = np.asarray(emb, dtype=np.float32) - return emb[:n] if pad_n else emb - - -def retrieve_topk(tbl, query_vecs: np.ndarray, problems: List[str], - sim_threshold: float - ) -> List[List[Dict[str, Any]]]: - """Retrieve top-K thinking_raw per query with decontamination and length filter. - - Returns per-query list of dicts with keys: query, thinking, sim. - """ - results = [] - decontam_skipped = 0 - for qi, vec in enumerate(query_vecs): - hits = ( - tbl.search(vec.astype(np.float32).tolist()) - .metric('dot') - .limit(TOP_K + 50) - .select(['query_raw', 'thinking_raw', '_distance']) - .to_list() - ) - matched = [] - problem_text = problems[qi] if problems else '' - for h in hits: - if len(matched) >= TOP_K: - break - sim = 1.0 - h.get('_distance', 0.0) - if sim < sim_threshold: - continue - q = h.get('query_raw', '') - t = h.get('thinking_raw', '') - if not t: - continue - # Decontamination: skip if retrieved problem is too similar to current - if DECONTAM_THRESHOLD > 0 and problem_text and q: - if _ngram_jaccard(problem_text, q) > DECONTAM_THRESHOLD: - decontam_skipped += 1 - continue - # Drop traces exceeding max length (don't truncate — they'll be condensed poorly) - if len(t) > MAX_TRACE_LEN * 4: - continue - matched.append({'query': q, 'thinking': t, 'sim': sim}) - results.append(matched) - if decontam_skipped > 0: - logger.info(f'[decontam] skipped {decontam_skipped} leaked retrievals') - return results - - -# ============================================================================ -# Reward -# ============================================================================ -class AoPSAccuracyReward(Reward): - """Accuracy reward via boxed answer extraction + robust equivalence matching.""" - - @staticmethod - def extract_boxed(text: str) -> str: - idx = text.rfind('\\boxed{') - if idx == -1: - return '' - start = idx + len('\\boxed{') - depth = 1 - j = start - while j < len(text) and depth > 0: - if text[j] == '{': - depth += 1 - elif text[j] == '}': - depth -= 1 - j += 1 - if depth == 0: - return text[start:j - 1].strip() - return '' - - # --- MCQ letter regex (matches \textbf{(C) }value, (C) value, etc.) --- - _MCQ_GT_RE = re.compile( - r'^\\?(?:textbf|mathbf|text|mathrm)\{?\(?([A-E])[)}\s\\]*(.*)', - re.DOTALL) - _MCQ_PAREN_RE = re.compile(r'^\(?([A-E])\)?[.:\s\\]+(.*)', re.DOTALL) - _MCQ_SINGLE_LETTER_RE = re.compile(r'^[A-E]$') - # variable prefix: f(x)=..., m=..., N=..., P(n+1)=..., (x,y)=... - _VAR_PREFIX_RE = re.compile( - r'^(?:[a-zA-Z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)', re.DOTALL) - # GT with derivation: "18×1+999×2=2016" → extract RHS - _EQ_RHS_RE = re.compile(r'^.+=\s*(.+)$') - - @staticmethod - def normalize_answer(ans: str) -> str: - if not ans: - return '' - s = ans.strip() - # Pure MCQ letter - m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf)?\{?\(?([A-E])\)?\}?$', s) - if m: - return m.group(1) - s = s.replace(' ', '') - s = s.replace(r'\,', '') - s = s.replace(r'\;', '') - s = s.replace(r'\!', '') - # Remove text-mode wrappers but keep content (unwrap braces) - s = re.sub(r'\\(?:text|mathrm|mathbf|textbf|operatorname)\{([^}]*)\}', r'\1', s) - s = s.replace(r'\displaystyle', '') - s = re.sub(r'\\(?:left|right)[.()\[\]|]', '', s) - s = s.replace(r'\dfrac', r'\frac') - s = s.replace(r'\tfrac', r'\frac') - # \frac shorthand without braces: \frac ab → \frac{a}{b} - s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) - s = s.strip('$').strip() - # Remove trailing unit braces: {cm}, {kg}, etc. - s = re.sub(r'\{[a-zA-Z]+\}$', '', s) - # Degree normalization — strip entirely (degrees are contextual) - s = re.sub(r'\^\{\\circ\}|\^\\circ|°|\\circ', '', s) - # Remove \quad, \qquad, \ etc spacing - s = re.sub(r'\\(?:quad|qquad|\s)', '', s) - # Normalize minus: \minus{} → - - s = s.replace(r'\minus{}', '-').replace(r'\minus', '-') - - def _frac_to_slash(m): - text = m.group(0) - pos = text.index('{') + 1 - depth, num_start = 1, pos - while depth > 0: - if text[pos] == '{': depth += 1 - elif text[pos] == '}': depth -= 1 - pos += 1 - numer = text[num_start:pos - 1] - pos += 1 - den_start = pos - depth = 1 - while depth > 0: - if text[pos] == '{': depth += 1 - elif text[pos] == '}': depth -= 1 - pos += 1 - denom = text[den_start:pos - 1] - return f'({numer})/({denom})' - - s = re.sub( - r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', - _frac_to_slash, s) - s = re.sub(r'(?<!\w)(\d+)/(\d+)(?!\w)', r'(\1)/(\2)', s) - return s - - @classmethod - def _strip_var_prefix(cls, s: str) -> str: - """Strip variable assignment prefix: 'f(x)=x+1' → 'x+1', 'N=1006' → '1006'.""" - m = cls._VAR_PREFIX_RE.match(s) - return m.group(1).strip() if m else s - - @classmethod - def _extract_mcq_parts(cls, s: str): - """Extract (letter, value) from MCQ-formatted string. Returns (None, None) if not MCQ.""" - m = cls._MCQ_GT_RE.match(s) - if m: - return m.group(1), m.group(2).strip() - m = cls._MCQ_PAREN_RE.match(s) - if m: - return m.group(1), m.group(2).strip() - # GT ends with " (A)" pattern: "2+2\sqrt{7} (A)" - m2 = re.search(r'\(?([A-E])\)?\s*$', s) - if m2 and len(s) > 3: - return m2.group(1), s[:m2.start()].strip() - return None, None - - @staticmethod - def _try_numeric_equal(a: str, b: str) -> bool: - """Try numeric equality after normalization. Handles fracs and simple expressions.""" - import math - - def _try_eval(s: str): - # Direct float - try: - return float(s.replace('(', '').replace(')', '')) - except (ValueError, ZeroDivisionError): - pass - # Strip trailing unit-like suffix and retry - s_stripped = re.sub(r'[a-zA-Z]+$', '', s.replace('(', '').replace(')', '')).strip() - if s_stripped and s_stripped != s: - try: - return float(s_stripped) - except (ValueError, ZeroDivisionError): - pass - # Fraction pattern (a)/(b) - frac_re = re.compile(r'^\(([^)]+)\)/\(([^)]+)\)$') - m = frac_re.match(s) - if m: - try: - return float(m.group(1)) / float(m.group(2)) - except (ValueError, ZeroDivisionError): - pass - # Try evaluating simple math expressions (pi, sqrt, etc.) - expr = s - expr = expr.replace('\\pi', str(math.pi)) - expr = expr.replace('\\e', str(math.e)) - expr = re.sub(r'\\sqrt\{([^}]+)\}', r'(\1)**0.5', expr) - expr = re.sub(r'\\sqrt\[3\]\{([^}]+)\}', r'(\1)**(1/3)', expr) - expr = re.sub(r'\\sqrt\[([^]]+)\]\{([^}]+)\}', r'(\2)**(1/\1)', expr) - expr = expr.replace('{', '(').replace('}', ')') - expr = expr.replace('\\cdot', '*').replace('\\times', '*') - expr = re.sub(r'(\d)\(', r'\1*(', expr) - try: - val = eval(expr, {"__builtins__": {}, "math": math, "pi": math.pi, "e": math.e}, {}) - return float(val) - except Exception: - pass - return None - - va, vb = _try_eval(a), _try_eval(b) - if va is not None and vb is not None: - return abs(va - vb) < 1e-6 * max(1, abs(va), abs(vb)) - return False - - @classmethod - def _strip_quantifiers(cls, s: str) -> str: - """Strip universal/existential quantifier wrappers.""" - s = re.sub(r'^\\forall\s*\w+\s*\\in\s*\\mathbb\s*\{?[A-Z]\}?\s*[:,]\s*', '', s) - s = re.sub(r'\s*\(\\forall[^)]*\)\s*$', '', s) - s = re.sub(r'\s*\(for\s+all[^)]*\)\s*$', '', s, flags=re.IGNORECASE) - return s.strip() - - @classmethod - def _try_param_rename(cls, a: str, b: str) -> bool: - """Check if a == b up to consistent single free-parameter rename (ax vs cx).""" - if not a or not b or len(a) != len(b) or len(a) > 80: - return False - diffs = [(i, a[i], b[i]) for i in range(len(a)) if a[i] != b[i]] - if not diffs: - return False - src_chars = set(d[1] for d in diffs) - dst_chars = set(d[2] for d in diffs) - if len(src_chars) == 1 and len(dst_chars) == 1: - src, dst = src_chars.pop(), dst_chars.pop() - if src.isalpha() and dst.isalpha(): - return a.replace(src, dst) == b - return False - - @classmethod - def _normalize_tuple(cls, s: str) -> str: - """Normalize tuple formatting: (2, 5, 609) → 2,5,609.""" - # Strip set-builder conditions: \mid ... or | ... - s = re.sub(r'\\mid.*$', '', s) - s = re.sub(r'\|[^,]*$', '', s) - return re.sub(r'[\s()\[\]{}\\]', '', s) - - @classmethod - def _try_sympy_equal(cls, a: str, b: str) -> bool: - """Optional sympy-based algebraic equivalence (graceful fallback if unavailable).""" - try: - from sympy.parsing.latex import parse_latex - from sympy import simplify, nsimplify - expr_a = parse_latex(a) - expr_b = parse_latex(b) - diff = simplify(nsimplify(expr_a - expr_b)) - return diff == 0 - except Exception: - return False - - @classmethod - def answers_match(cls, predicted: str, reference: str) -> bool: - if not predicted or not reference: - return False - - norm_p = cls.normalize_answer(predicted) - norm_r = cls.normalize_answer(reference) - - # --- Strategy 1: direct string equality --- - if norm_p == norm_r: - return True - - # --- Strategy 2: case-insensitive --- - if norm_p.lower() == norm_r.lower(): - return True - - # --- Strategy 3: numeric equality --- - if cls._try_numeric_equal(norm_p, norm_r): - return True - - # --- Strategy 4: variable-prefix stripping (both sides) --- - stripped_p = cls.normalize_answer(cls._strip_var_prefix(predicted)) - stripped_r = cls.normalize_answer(cls._strip_var_prefix(reference)) - if stripped_p and stripped_r and stripped_p == stripped_r: - return True - if stripped_p and stripped_r and cls._try_numeric_equal(stripped_p, stripped_r): - return True - - # --- Strategy 5: MCQ double matching --- - # Extract letter+value from reference - ref_letter, ref_value = cls._extract_mcq_parts(reference) - if ref_letter: - # pred matches the letter? - if norm_p == ref_letter or predicted.strip().upper() == ref_letter: - return True - # pred matches the value? - if ref_value: - norm_ref_val = cls.normalize_answer(ref_value) - if norm_p == norm_ref_val or cls._try_numeric_equal(norm_p, norm_ref_val): - return True - # Extract from predicted side too (pred="B", ref has value) - pred_letter, pred_value = cls._extract_mcq_parts(predicted) - if pred_letter: - if norm_r == pred_letter or reference.strip().upper() == pred_letter: - return True - if pred_value: - norm_pred_val = cls.normalize_answer(pred_value) - if norm_r == norm_pred_val or cls._try_numeric_equal(norm_r, norm_pred_val): - return True - # MCQ: GT is single letter, pred is numeric/expression → match if pred chose option - if cls._MCQ_SINGLE_LETTER_RE.match(reference.strip()): - if cls._MCQ_SINGLE_LETTER_RE.match(predicted.strip().upper()): - return predicted.strip().upper() == reference.strip().upper() - # pred is a value, GT is just a letter: we accept pred=letter match only - # (can't verify value without options text) - - # --- Strategy 6: tuple/set normalization --- - tuple_p = cls._normalize_tuple(norm_p) - tuple_r = cls._normalize_tuple(norm_r) - if ',' in tuple_p and tuple_p == tuple_r: - return True - if stripped_p and stripped_r: - tuple_sp = cls._normalize_tuple(stripped_p) - tuple_sr = cls._normalize_tuple(stripped_r) - if ',' in tuple_sp and tuple_sp == tuple_sr: - return True - - # --- Strategy 7: equation reorder (a+b=c vs c=a+b, or lhs=rhs swapped) --- - if '=' in norm_r and '=' not in norm_p: - # GT has derivation like "18*1+999*2=2016", pred is "2016" - parts = norm_r.split('=') - for part in parts: - part = part.strip() - if part == norm_p or cls._try_numeric_equal(part, norm_p): - return True - if '=' in norm_p and '=' not in norm_r: - parts = norm_p.split('=') - for part in parts: - part = part.strip() - if part == norm_r or cls._try_numeric_equal(part, norm_r): - return True - if '=' in norm_p and '=' in norm_r: - # Both have =: try matching LHS=RHS in any order - pp = [x.strip() for x in norm_p.split('=')] - rp = [x.strip() for x in norm_r.split('=')] - if set(pp) == set(rp): - return True - - # --- Strategy 8: multiplicative reorder (27\pi\sqrt{6} vs 27\sqrt{6}\pi) --- - def _sort_factors(s): - tokens = re.findall(r'\\?[a-zA-Z]+\{[^}]*\}|\\?[a-zA-Z]+|\d+|[^a-zA-Z\d\\{}]', s) - return ''.join(sorted(tokens)) - if _sort_factors(norm_p) == _sort_factors(norm_r): - return True - - # --- Strategy 9: sympy algebraic equivalence (optional, slow) --- - if cls._try_sympy_equal(predicted, reference): - return True - - # --- Strategy 9b: quantifier stripping + param rename --- - q_stripped_r = cls._strip_quantifiers(reference) - q_stripped_p = cls._strip_quantifiers(predicted) - if q_stripped_r != reference or q_stripped_p != predicted: - norm_qr = cls.normalize_answer(cls._strip_var_prefix(q_stripped_r)) - norm_qp = cls.normalize_answer(cls._strip_var_prefix(q_stripped_p)) - if norm_qr and norm_qp: - if norm_qr == norm_qp: - return True - if cls._try_param_rename(norm_qr, norm_qp): - return True - - # --- Strategy 10: param rename on var-prefix-stripped forms --- - if stripped_p and stripped_r and cls._try_param_rename(stripped_p, stripped_r): - return True - - return False - - def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: - rewards = [] - for traj in trajectories: - messages = traj.get('messages', []) - completion = '' - for msg in reversed(messages): - if msg.get('role') == 'assistant': - completion = msg.get('content', '') - break - user_data = traj.get('user_data') or [] - gt = '' - for item in user_data: - if item[0] == 'ground_truth': - gt = item[1] - break - predicted = self.extract_boxed(completion) - correct = self.answers_match(predicted, gt) - rewards.append(1.0 if correct else 0.0) - return rewards - - -class FormatReward(Reward): - """Reward for having \\boxed{} and <hint>...</hint> analysis in the output.""" - - _HINT_RE = re.compile(r'<hint>(.*?)</hint>', re.DOTALL) - _THINK_RE = re.compile(r'^.*?</think>', 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 <think>...</think> block to get the visible response.""" - think_end = text.find('</think>') - if think_end >= 0: - return text[think_end + len('</think>'):] - return text - - def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: - rewards = [] - for traj in trajectories: - messages = traj.get('messages', []) - completion = '' - sys_content = '' - for msg in messages: - if msg.get('role') == 'system': - sys_content = self._to_text(msg.get('content', '')) - for msg in reversed(messages): - if msg.get('role') == 'assistant': - completion = self._to_text(msg.get('content', '')) - break - has_boxed = '\\boxed{' in completion - # Only check hint tags for RAG prompts (system contains examples) - is_rag = 'condensed reasoning examples from similar problems' in sys_content - if is_rag: - # Check hint in VISIBLE response only (after </think>) - visible = self._visible_response(completion) - hint_match = self._HINT_RE.search(visible) - has_good_hint = False - if hint_match: - hint_text = hint_match.group(1).strip() - hint_pos = hint_match.start() - # Hint must be near the start of visible output - at_beginning = hint_pos < max(len(visible) * 0.05, 200) - is_substantive = len(hint_text) >= self._MIN_HINT_LEN - # Reject hints that are too long (model dumping thinking) - not_dump = len(hint_text) <= self._MAX_HINT_LEN - # Must start with the required prefix - has_prefix = hint_text.startswith(HINT_REQUIRED_PREFIX) - has_good_hint = (at_beginning and is_substantive - and not_dump and has_prefix) - # 0.3 for boxed + 0.2 for good hint = 0.5 max - reward = (0.3 if has_boxed else 0.0) + (0.2 if has_good_hint else 0.0) - else: - reward = 0.5 if has_boxed else 0.0 - rewards.append(reward) - return rewards - - -class GibberishPenalty(Reward): - """Negative reward for degenerate outputs (gibberish/random unicode tail).""" - - TAIL_CHARS = 400 - GIBBERISH_THRESHOLD = 0.20 # >20% non-math non-ascii in tail - - @classmethod - def is_gibberish(cls, text: str) -> bool: - if not text: - return False - tail = text[-cls.TAIL_CHARS:] if len(text) > cls.TAIL_CHARS else text - non_math_non_ascii = 0 - for c in tail: - code = ord(c) - # Allow: ASCII, common CJK (for Chinese math), LaTeX symbols - if code > 127 and not (0x4e00 <= code <= 0x9fff): - non_math_non_ascii += 1 - return non_math_non_ascii > len(tail) * cls.GIBBERISH_THRESHOLD - - def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: - rewards = [] - for traj in trajectories: - messages = traj.get('messages', []) - completion = '' - for msg in reversed(messages): - if msg.get('role') == 'assistant': - completion = msg.get('content', '') - break - rewards.append(-0.5 if self.is_gibberish(completion) else 0.0) - return rewards - - -def compute_rewards(trajectories: List[Dict[str, Any]] - ) -> Tuple[List[float], List[float], List[float]]: - acc_fn = AoPSAccuracyReward() - fmt_fn = FormatReward() - gib_fn = GibberishPenalty() - acc = acc_fn(trajectories) - fmt = fmt_fn(trajectories) - gib = gib_fn(trajectories) - total = [a + f + g for a, f, g in zip(acc, fmt, gib)] - return total, fmt, acc - - -# ============================================================================ -# Dataset: AoPS boxed problems -# ============================================================================ -def create_aops_dataset(): - """Load AoPS and create GRPO-style dataset (prompt only, with ground_truth in user_data).""" - from modelscope import MsDataset - from twinkle.data_format import Message, Trajectory - - ds = MsDataset.load(AOPS_DATASET_ID, split='train', - download_mode='reuse_dataset_if_exists') - rows = [] - for row in ds: - if not row['metadata'].get('boxed'): - continue - ref = AoPSAccuracyReward.extract_boxed(row['solution']) - if not ref: - continue - rows.append({'problem': row['problem'], 'ground_truth': ref}) - - logger.info(f'[aops] loaded {len(rows)} boxed problems') - rng = random.Random(AOPS_SEED) - rng.shuffle(rows) - - # Build Trajectory list (prompt-only for GRPO) - trajectories = [] - for r in rows: - # Use direct system prompt as placeholder — will be replaced by RAG pipeline - traj = Trajectory( - messages=[ - Message(role='system', content=SYSTEM_DIRECT), - Message(role='user', content=r['problem']), - ], - user_data=[('ground_truth', r['ground_truth'])], - ) - trajectories.append(traj) - - data_meta = DatasetMeta(data=trajectories) - dataset = Dataset(data_meta) - dataset.set_template('Qwen3_5Template', model_id=MODEL_ID, - max_length=16384, truncation_strategy='delete', - enable_thinking=True) - dataset.encode(add_generation_prompt=True) - return dataset - - -# ============================================================================ -# Main -# ============================================================================ -def main(): - # GPU rank allocation - cond_start = 0 - emb_start = cond_start + CONDENSER_GPUS - sampler_start = emb_start + EMB_GPUS - model_start = sampler_start + SAMPLER_GPUS - - device_groups = [ - DeviceGroup(name='condenser', ranks=list(range(cond_start, emb_start)), - device_type='GPU'), - DeviceGroup(name='emb_model', ranks=list(range(emb_start, sampler_start)), - device_type='GPU'), - DeviceGroup(name='sampler', ranks=list(range(sampler_start, model_start)), - device_type='GPU'), - DeviceGroup(name='model', ranks=list(range(model_start, NUM_GPUS)), - device_type='GPU'), - ] - - model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, fsdp_size=MODEL_GPUS, ulysses_size=2) - sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, dp_size=SAMPLER_GPUS) - emb_mesh = DeviceMesh.from_sizes(world_size=EMB_GPUS, dp_size=EMB_GPUS) - condenser_mesh = DeviceMesh.from_sizes(world_size=CONDENSER_GPUS, dp_size=CONDENSER_GPUS) - - twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, - groups=device_groups, lazy_collect=False) - - # -- Training model (full-parameter) -- - model = TransformersModel( - model_id=MODEL_ID, device_mesh=model_mesh, remote_group='model') - model.set_optimizer('AdamW', lr=LEARNING_RATE) - model.set_lr_scheduler('CosineAnnealingLR', T_max=MAX_STEPS, eta_min=0) - model.set_loss('GSPOLoss', epsilon=0.2, epsilon_high=0.28, beta=0.04) - model.set_processor(InputProcessor) - model.set_template('Qwen3_5Template', model_id=MODEL_ID, - enable_thinking=True, max_length=32768) - - # -- Rollout sampler -- - sampler = vLLMSampler( - model_id=MODEL_ID, - engine_args={ - 'gpu_memory_utilization': 0.8, - 'max_model_len': 32768, - }, - device_mesh=sampler_mesh, - remote_group='sampler', - ) - sampler.set_template('Qwen3_5Template', model_id=MODEL_ID, - enable_thinking=True, max_length=32768) - - # -- Embedding model -- - emb_model = TransformersModel( - model_id=EMBED_MODEL_ID, device_mesh=emb_mesh, remote_group='emb_model') - emb_model.set_processor(InputProcessor) - emb_model.set_loss(InfonceLoss, temperature=0.03, use_batch=True) - emb_template = Qwen3_5Template( - model_id=EMBED_MODEL_ID, max_length=EMBED_MAX_LENGTH, - truncation_strategy='delete', enable_thinking=False) - - # -- Condenser sampler -- - condenser_sampler = vLLMSampler( - model_id=CONDENSE_MODEL_ID, - engine_args={'gpu_memory_utilization': 0.85, 'max_model_len': 32768}, - device_mesh=condenser_mesh, - remote_group='condenser', - ) - condenser_sampler.set_template( - 'Qwen3_5Template', model_id=CONDENSE_MODEL_ID, - enable_thinking=False, truncation_strategy='delete', max_length=32768) - condenser_template = Qwen3_5Template( - model_id=CONDENSE_MODEL_ID, max_length=32768, - enable_thinking=False, truncation_strategy='delete') - condenser_special_tokens = set(condenser_template.tokenizer.all_special_tokens) - compress_params = SamplingParams( - max_tokens=CONDENSE_MAX_TOKENS, temperature=CONDENSE_TEMPERATURE, - top_p=0.5, num_samples=1) - - # -- API client (condenser fallback) -- - api_client = None - if CONDENSE_API_KEY: - api_client = OpenAIClient( - model=CONDENSE_API_MODEL, api_key=CONDENSE_API_KEY, - base_url=CONDENSE_BASE_URL) - - # -- LanceDB -- - import lancedb - db = lancedb.connect(DB_PATH) - tbl = db.open_table(DB_TABLE) - logger.info(f'[rag] LanceDB ready, rows={tbl.count_rows()}') - - # -- Checkpoint & DataLoader -- - ckpt_manager = CheckpointEngineManager(model=model, sampler=sampler) - - GLOBAL_BATCH_SIZE = BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS - dataloader = DataLoader( - dataset=create_aops_dataset, - batch_size=GLOBAL_BATCH_SIZE, - min_batch_size=GLOBAL_BATCH_SIZE, - device_mesh=model_mesh, - remote_group='model', - ) - - advantage_fn = GRPOAdvantage() - metrics = CompletionRewardMetric() - sampling_params = SamplingParams( - max_tokens=MAX_NEW_TOKENS, num_samples=1, logprobs=1, - temperature=1.0, top_p=0.95) - - optim_step = 0 - logger.info('Starting RAG-hint GRPO training') - logger.info(get_device_placement()) - - # -- Prefetch: overlap RAG data preparation with training -- - prefetch_pool = ThreadPoolExecutor(max_workers=1) - - def _extract_text(content) -> str: - """Extract plain text from content (str or list-of-parts format).""" - if isinstance(content, str): - return content - if isinstance(content, list): - return ''.join( - p.get('text', '') for p in content if isinstance(p, dict) and p.get('type') == 'text') - return str(content) if content else '' - - def prepare_rag_batch(batch): - """Embed → retrieve → condense → build prompts. Runs in background thread.""" - problems = [] - ground_truths = [] - for item in batch: - msgs = item.get('messages', []) - prob = '' - for m in msgs: - if m.get('role') == 'user': - prob = _extract_text(m.get('content', '')) - break - problems.append(prob) - ud = item.get('user_data', []) - gt = '' - for pair in ud: - if pair[0] == 'ground_truth': - gt = pair[1] - break - ground_truths.append(gt) - - # Embed & retrieve - query_vecs = get_embeddings(emb_model, emb_template, problems, EMB_GPUS) - retrieved = retrieve_topk(tbl, query_vecs, problems, SIM_THRESHOLD) - raw_retrieved_counts = [len(r) for r in retrieved] - - # LLM-based decontamination: judge ALL retrievals via API - if api_client: - judge_pairs = [] # (qi, ret_idx, prob_a, prob_b) - for qi, rets in enumerate(retrieved): - for ri, ret in enumerate(rets): - judge_pairs.append((qi, ri, problems[qi], ret['query'])) - - if judge_pairs: - pairs_input = [(pa, pb) for _, _, pa, pb in judge_pairs] - verdicts = _llm_judge_same_problem(api_client, pairs_input) - to_remove = set() - for vi, (qi, ri, _, _) in enumerate(judge_pairs): - if verdicts[vi]: - to_remove.add((qi, ri)) - if to_remove: - logger.info(f'[decontam-llm] filtered {len(to_remove)} same-problem retrievals') - for qi in range(len(retrieved)): - retrieved[qi] = [ - ret for ri, ret in enumerate(retrieved[qi]) - if (qi, ri) not in to_remove - ] - - # Condense (batch local vLLM + API fallback) - condensed_examples: List[List[Dict[str, str]]] = [[] for _ in range(len(problems))] - tasks_to_condense = [] - for i, rets in enumerate(retrieved): - for j, ret in enumerate(rets): - tasks_to_condense.append((i, j, problems[i], ret)) - - if tasks_to_condense: - condense_prompts = [] - for idx, _j, prob, ret in tasks_to_condense: - user_msg = COMPRESS_USER.format(query=prob, text=ret['thinking']) - condense_prompts.append({'messages': [ - {'role': 'system', 'content': COMPRESS_SYSTEM}, - {'role': 'user', 'content': user_msg}]}) - - try: - condense_responses = condenser_sampler.sample(condense_prompts, compress_params) - except Exception as exc: - logger.warning(f'[condense] local batch error: {exc}') - condense_responses = [None] * len(condense_prompts) - - api_fallback_indices = [] - for ci, (idx, _j, prob, ret) in enumerate(tasks_to_condense): - resp = condense_responses[ci] if condense_responses else None - seq = resp.sequences[0] if resp and resp.sequences else None - text = '' - if seq and seq.stop_reason != 'length' and seq.decoded: - text = seq.decoded - for tok in condenser_special_tokens: - text = text.replace(tok, '') - text = text.strip() - if text: - condensed_examples[idx].append({'query': ret['query'], 'thinking': text}) - else: - api_fallback_indices.append(ci) - - if api_fallback_indices and api_client: - def _fallback(ci): - return ci, _api_condense_single(api_client, condense_prompts[ci]['messages']) - with ThreadPoolExecutor(max_workers=CONDENSE_API_CONCURRENCY) as pool: - futs = [pool.submit(_fallback, ci) for ci in api_fallback_indices] - for fut in as_completed(futs): - ci, result = fut.result() - idx, _j, prob, ret = tasks_to_condense[ci] - text = result if result else ret['thinking'][:MAX_TRACE_LEN] - condensed_examples[idx].append({'query': ret['query'], 'thinking': text}) - elif api_fallback_indices: - for ci in api_fallback_indices: - idx, _j, prob, ret = tasks_to_condense[ci] - condensed_examples[idx].append( - {'query': ret['query'], 'thinking': ret['thinking'][:MAX_TRACE_LEN]}) - - # API hint analysis: pre-compute RAG relevance verdict - hint_analyses = [None] * len(problems) - if api_client: - hint_analyses = _api_hint_analysis_batch(api_client, problems, condensed_examples) - - # Build prompts with rag_fallback_sim check - rag_prompts = [] - rag_debug_records = [] - for i, prob in enumerate(problems): - examples = condensed_examples[i] - rets = retrieved[i] - best_sim = max((r['sim'] for r in rets), default=0.0) - use_rag = bool(examples) and best_sim >= RAG_FALLBACK_SIM - - if use_rag: - # If API hint analysis succeeded, use pre-analyzed prompt (no <hint> needed) - if hint_analyses[i]: - rag_sys_content = build_preanalysis_system(hint_analyses[i]) - else: - # Fallback: old-style prompt with self-analysis requirement - parts = [SYSTEM_WITH_RAG_HEADER] - for eidx, ex in enumerate(examples, 1): - parts.append(EXAMPLE_TEMPLATE.format( - idx=eidx, - example_query=ex['query'], - example_thinking=ex['thinking'])) - rag_sys_content = ''.join(parts) - - # RAG group (only RAG, no paired NoRAG) - rag_prompts.append({ - 'messages': [ - {'role': 'system', 'content': rag_sys_content}, - {'role': 'user', 'content': prob}, - ], - 'user_data': [('ground_truth', ground_truths[i])], - 'assistant_prefix': ANALYSIS_PREFIX, - }) - rag_debug_records.append({ - 'problem': prob[:200], - 'ground_truth': ground_truths[i], - 'best_sim': round(best_sim, 4), - 'num_raw_retrieved': raw_retrieved_counts[i], - 'num_retrieved': len(rets), - 'num_condensed': len(examples), - 'use_rag': True, - 'has_preanalysis': hint_analyses[i] is not None, - 'preanalysis_len': len(hint_analyses[i]) if hint_analyses[i] else 0, - 'top_retrieved_query': rets[0]['query'][:200] if rets else '', - 'condensed_len': len(examples[0].get('thinking', '')) if examples else 0, - }) - else: - # No hint found — skip this query entirely - continue - - return rag_prompts, rag_debug_records - - # Submit first batch prefetch - os.makedirs(OUTPUT_DIR, exist_ok=True) - rag_log_path = os.path.join(OUTPUT_DIR, 'rag_diagnostics.jsonl') - rag_log_f = open(rag_log_path, 'w', encoding='utf-8') - logger.info(f'[rag] diagnostics → {rag_log_path}') - - batch_iter = iter(dataloader) - pending_future = None - try: - first_batch = next(batch_iter) - pending_future = prefetch_pool.submit(prepare_rag_batch, first_batch) - except StopIteration: - pass - - try: - while pending_future is not None: - if optim_step >= MAX_STEPS: - break - - metrics.reset() - rag_prompts, rag_debug_records = pending_future.result() - - # Write RAG diagnostics - for rec in rag_debug_records: - rec['step'] = optim_step - rag_log_f.write(json.dumps(rec, ensure_ascii=False) + '\n') - rag_log_f.flush() - - # Submit next batch prefetch (overlaps with rollout + training) - pending_future = None - try: - next_batch = next(batch_iter) - pending_future = prefetch_pool.submit(prepare_rag_batch, next_batch) - except StopIteration: - pass - - # ---- Expand for NUM_GENERATIONS and sample ---- - expand_prompts = [] - for prompt in rag_prompts: - expand_prompts.extend([prompt] * NUM_GENERATIONS) - - if not expand_prompts: - logger.warning(f'[Step {optim_step}] empty prompt list after RAG processing, skip') - continue - - ckpt_manager.sync_weights(merge_and_sync=False) - sampler.reset_prefix_cache() - - sample_responses = sampler.sample(expand_prompts, sampling_params) - - # ---- Collect rollouts ---- - all_input_data: List[Dict[str, Any]] = [] - all_old_logps: List[List[float]] = [] - all_completion_lengths: List[int] = [] - - for sample_response in sample_responses: - for sequence in sample_response.sequences: - all_input_data.append(sequence.new_input_feature) - all_old_logps.append([logprob[0][1] for logprob in sequence.logprobs]) - all_completion_lengths.append(len(sequence.tokens)) - - # ---- Rewards ---- - total_rewards, format_rewards, accuracy_rewards = compute_rewards(all_input_data) - - # Zero out rewards for rollouts that hit the max_tokens ceiling - max_len_threshold = int(MAX_NEW_TOKENS * 0.95) - for i in range(len(all_input_data)): - if all_completion_lengths[i] >= max_len_threshold: - total_rewards[i] = 0.0 - accuracy_rewards[i] = 0.0 - format_rewards[i] = 0.0 - - # Per-step reward summary to diagnostics - n_correct = sum(1 for a in accuracy_rewards if a > 0) - rag_log_f.write(json.dumps({ - 'step': optim_step, 'type': 'reward_summary', - 'n_samples': len(accuracy_rewards), - 'accuracy': n_correct / len(accuracy_rewards) if accuracy_rewards else 0, - 'mean_reward': sum(total_rewards) / len(total_rewards) if total_rewards else 0, - }, ensure_ascii=False) + '\n') - - metrics.accumulate( - completion_lengths=all_completion_lengths, - rewards={ - 'total': total_rewards, - 'format': format_rewards, - 'accuracy': accuracy_rewards, - }, - ) - - # ---- GRPO advantage ---- - advantages = advantage_fn( - total_rewards, num_generations=NUM_GENERATIONS, scale='group').tolist() - if ADV_CLIP > 0: - advantages = [max(-ADV_CLIP, min(ADV_CLIP, a)) for a in advantages] - - # Log all rollout responses (after advantage computation) - _extract_boxed = AoPSAccuracyReward.extract_boxed - def _content_to_str(content): - """Convert message content (str or list of blocks) to plain text.""" - if isinstance(content, str): - return content - if isinstance(content, list): - return ''.join( - b.get('text', '') if isinstance(b, dict) else str(b) - for b in content) - return str(content) - - for ridx, traj in enumerate(all_input_data): - msgs = traj.get('messages', []) - assistant_text = _content_to_str(next( - (m['content'] for m in reversed(msgs) if m.get('role') == 'assistant'), '')) - user_text = _content_to_str(next( - (m['content'] for m in msgs if m.get('role') == 'user'), '')) - sys_text = _content_to_str(next( - (m['content'] for m in msgs if m.get('role') == 'system'), '')) - user_data = traj.get('user_data') or [] - gt = next((v for k, v in user_data if k == 'ground_truth'), '') - problem_idx = ridx // NUM_GENERATIONS - use_rag = ('condensed reasoning examples from similar problems' in sys_text - or 'RAG Analysis (pre-computed)' in sys_text) - # Per-problem group accuracy (all generations for same problem) - grp_start = problem_idx * NUM_GENERATIONS - grp_end = grp_start + NUM_GENERATIONS - grp_acc = sum(accuracy_rewards[grp_start:grp_end]) / NUM_GENERATIONS - - rag_log_f.write(json.dumps({ - 'step': optim_step, 'type': 'rollout', - 'idx': ridx, - 'problem_idx': problem_idx, - 'problem': user_text, - 'system': sys_text, - 'response': assistant_text, - 'ground_truth': gt, - 'predicted': _extract_boxed(assistant_text), - 'use_rag': use_rag, - 'best_sim': rag_debug_records[problem_idx].get('best_sim', 0.0) if problem_idx < len(rag_debug_records) else 0.0, - 'reward': total_rewards[ridx], - 'accuracy_reward': accuracy_rewards[ridx], - 'format_reward': format_rewards[ridx], - 'advantage': advantages[ridx], - 'completion_length': all_completion_lengths[ridx], - 'group_accuracy': grp_acc, - }, ensure_ascii=False) + '\n') - - rag_log_f.flush() - - # ---- Filter out low-signal problem groups (DAPO-style dynamic sampling) ---- - # Skip groups where accuracy is too low (<0.1) or too high (>0.9) - # to avoid gradient dominated by gibberish/format noise or no learning signal. - filtered_inputs, filtered_old_logps, filtered_advantages = [], [], [] - actual_num_groups = len(all_input_data) // NUM_GENERATIONS - for g in range(actual_num_groups): - g_start = g * NUM_GENERATIONS - g_end = g_start + NUM_GENERATIONS - grp_adv = advantages[g_start:g_end] - if all(abs(a) < 1e-8 for a in grp_adv): - continue - grp_acc_rate = sum(accuracy_rewards[g_start:g_end]) / NUM_GENERATIONS - if grp_acc_rate < 0.2 or grp_acc_rate > 0.8: - continue - filtered_inputs.extend(all_input_data[g_start:g_end]) - filtered_old_logps.extend(all_old_logps[g_start:g_end]) - filtered_advantages.extend(grp_adv) - - # ---- Mini-batch training with gradient accumulation ---- - # Process MICRO_BATCH_SIZE samples per forward, accumulate grad_accum_steps - # times before one optimizer step. clip_grad_norm normalizes by accumulated - # num_tokens, ensuring mathematical equivalence with larger batch forward. - total_completions = len(filtered_inputs) - if total_completions == 0: - logger.info(f'[Step {optim_step}] all groups filtered (uniform rewards), skip training') - continue - - grad_accum_steps = MINI_BATCH_SIZE // MICRO_BATCH_SIZE - accum_count = 0 - for mb_start in range(0, total_completions, MICRO_BATCH_SIZE): - mb_end = min(mb_start + MICRO_BATCH_SIZE, total_completions) - mb_inputs = filtered_inputs[mb_start:mb_end] - mb_old_logps = filtered_old_logps[mb_start:mb_end] - mb_advantages = filtered_advantages[mb_start:mb_end] - - outputs = model.forward_backward( - inputs=mb_inputs, - old_logps=mb_old_logps, - ref_logps=mb_old_logps, - advantages=mb_advantages, - ) - accum_count += 1 - - if accum_count % grad_accum_steps == 0: - # Loss spike skip: discard explosive gradients - skip_step = False - try: - loss_val = outputs.get('loss', None) - if loss_val is not None: - if hasattr(loss_val, 'item'): - loss_val = loss_val.item() - if loss_val > LOSS_SPIKE_THRESHOLD: - skip_step = True - logger.warning( - f'[Step {optim_step}] Loss spike: {loss_val:.4f} > ' - f'{LOSS_SPIKE_THRESHOLD}, skipping update') - except Exception: - pass - - if skip_step: - model.zero_grad() - else: - model.clip_grad_and_step() - optim_step += 1 - - if optim_step >= MAX_STEPS: - break - if optim_step % SAVE_STEPS == 0: - model.save(f'rag-hint-grpo-checkpoint-{optim_step}') - - # Flush remaining accumulated gradients (incomplete window at tail) - if accum_count % grad_accum_steps != 0: - skip_step = False - try: - loss_val = outputs.get('loss', None) - if loss_val is not None: - if hasattr(loss_val, 'item'): - loss_val = loss_val.item() - if loss_val > LOSS_SPIKE_THRESHOLD: - skip_step = True - logger.warning( - f'[Step {optim_step}] Loss spike (tail): {loss_val:.4f} > ' - f'{LOSS_SPIKE_THRESHOLD}, skipping update') - except Exception: - pass - - if skip_step: - model.zero_grad() - else: - model.clip_grad_and_step() - optim_step += 1 - - log_dict = metrics.calculate() - log_dict.update(model.calculate_metric(is_training=True)) - metrics.reset() - logger.info(f'[Step {optim_step}/{MAX_STEPS}] {log_dict}') - finally: - prefetch_pool.shutdown(wait=False) - rag_log_f.close() - - logger.info(f'Training completed. optim_steps={optim_step}') - model.save('rag-hint-grpo-final') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/legacy/train_embedding_full_ddp.py b/cookbook/exp/legacy/train_embedding_full_ddp.py deleted file mode 100644 index 97ab3b128..000000000 --- a/cookbook/exp/legacy/train_embedding_full_ddp.py +++ /dev/null @@ -1,270 +0,0 @@ -"""Full-parameter embedding training on pre-compressed dataset. - -Reads the pre-compressed HF Dataset produced by make_embedding_dataset.py, -encodes features, trains with InfoNCE loss. - -Architecture (4 GPUs): - - Ranks 0-3: Trainable embedding model, InfoNCE loss. - -Launch: - python cookbook/exp/embedding/train_embedding_full_ddp.py -""" -import os -import time -from typing import Any, Dict, List, Literal, Optional - -import swanlab - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_device_placement, get_logger -from twinkle.loss import InfonceLoss -from twinkle.metric import EmbeddingMetric -from twinkle.model import TransformersModel -from twinkle.processor import InputProcessor -from twinkle.template import Qwen3_5Template, Template - -logger = get_logger() - -# -- Backend selection -------------------------------------------------------- -BACKEND: Literal['transformers', 'megatron'] = 'transformers' - -MODEL_ID = os.environ.get('MODEL_ID', 'ms://Qwen/Qwen3.5-4B') - -# -- GPU placement ------------------------------------------------------------ -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', 64)) -LEARNING_RATE = 1e-5 -GRADIENT_ACCUMULATION_STEPS = 1 -LOG_INTERVAL = 2 -SAVE_INTERVAL = 2000 -NUM_EPOCHS = 1 - -# -- Dataset path (output of make_embedding_dataset.py) ----------------------- -DATASET_PATH = os.environ.get('EMB_DATASET_PATH', 'ms://twinkle-kit/qth-embedding') -MIX_SHUFFLE_SEED = 42 - -# -- Resume from checkpoint --------------------------------------------------- -RESUME_CHECKPOINT = os.environ.get('RESUME_CHECKPOINT', '') -RESUME_STEP = int(os.environ.get('RESUME_STEP', 0)) - -# -- Output ------------------------------------------------------------------- -OUTPUT_DIR = f'./output/embedding_full_{BACKEND}' - - -# ============================================================================= -# Model builders -# ============================================================================= - -def build_model(device_mesh: DeviceMesh): - model_id = RESUME_CHECKPOINT if RESUME_CHECKPOINT else MODEL_ID - if BACKEND == 'transformers': - model = TransformersModel( - model_id=model_id, - device_mesh=device_mesh, - remote_group='model', - ddp_config={'find_unused_parameters': True}, - ) - from twinkle.patch.no_split_modules import NoSplitModulesPatch - model.apply_patch(NoSplitModulesPatch({'Qwen3_5DecoderLayer'})) - return model - if BACKEND == 'megatron': - from twinkle.model import MegatronModel - return MegatronModel( - model_id=MODEL_ID, - device_mesh=device_mesh, - remote_group='model', - mixed_precision='bf16', - variable_seq_lengths=True, - ) - raise ValueError(f'Unknown BACKEND={BACKEND!r}') - - -def setup_optimizer(model, total_steps: int): - if BACKEND == 'transformers': - model.set_optimizer(optimizer_cls='AdamW', lr=LEARNING_RATE) - model.set_lr_scheduler( - scheduler_cls='CosineWarmupScheduler', - num_warmup_steps=200, - num_training_steps=total_steps, - ) - return - if BACKEND == 'megatron': - model.set_optimizer(optimizer_cls='default', lr=LEARNING_RATE) - model.set_lr_scheduler( - scheduler_cls='default', - lr_warmup_steps=50, - lr_decay_steps=total_steps, - ) - return - raise ValueError(f'Unknown BACKEND={BACKEND!r}') - - -def save_checkpoint(model, name: str): - model.save(name, output_dir=OUTPUT_DIR) - - -# ============================================================================= -# Feature encoding -# ============================================================================= - -def _get_first_feature(decoded_text: str, template: Template, role: str) -> Optional[Dict[str, Any]]: - if not decoded_text: - return None - if role == 'anchor': - feat = template.encode({'messages': [ - {'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 - - -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 - - -# ============================================================================= -# Main training -# ============================================================================= - -def train(): - device_groups = [ - DeviceGroup(name='model', - ranks=list(range(MODEL_GPUS)), - device_type='GPU'), - ] - model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, dp_size=MODEL_GPUS) - twinkle.initialize(mode='ray', nproc_per_node=MODEL_GPUS, groups=device_groups) - - # -- Load pre-compressed dataset ------------------------------------------ - from twinkle.dataset import Dataset as TwinkleDataset, DatasetMeta - logger.info(f'[data] loading pre-compressed dataset from {DATASET_PATH}') - dataset = TwinkleDataset(DatasetMeta(dataset_id=DATASET_PATH), download_mode='force_redownload') - dataset = dataset.dataset.shuffle(seed=MIX_SHUFFLE_SEED) - logger.info(f'[data] {len(dataset)} rows loaded') - - # -- Compute steps -------------------------------------------------------- - rows_per_step = BATCH_SIZE - total_steps = (len(dataset) // rows_per_step) * NUM_EPOCHS - optimizer_steps = total_steps // GRADIENT_ACCUMULATION_STEPS - - # -- Model ---------------------------------------------------------------- - model = build_model(model_mesh) - model.set_processor(InputProcessor) - model.set_loss(InfonceLoss, temperature=TEMPERATURE, use_batch=True, - hard_negatives=HARD_NEGATIVES) - setup_optimizer(model, optimizer_steps) - model.add_metric(EmbeddingMetric, is_training=True) - - 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 steps: {total_steps}, optimizer steps: {optimizer_steps}') - - swanlab.init(project='twinkle', config={ - 'backend': BACKEND, - 'model_id': MODEL_ID, - 'batch_size': BATCH_SIZE, - 'lr': LEARNING_RATE, - 'temperature': TEMPERATURE, - 'emb_max_length': EMB_MAX_LENGTH, - 'dataset_path': DATASET_PATH, - }) - - # -- Train loop ----------------------------------------------------------- - cur_step = 0 - _skip_rows = RESUME_STEP * rows_per_step # approximate rows to skip - - for epoch in range(NUM_EPOCHS): - for start in range(0, len(dataset), rows_per_step): - if start < _skip_rows: - continue - - 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() - features = _encode_batch(rows_list, emb_template) - t_encode = time.monotonic() - t0 - - if len(features) < 4: - continue - - t1 = time.monotonic() - model.forward_backward(inputs=features, task='embedding') - model.clip_grad_and_step( - gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) - t_train = time.monotonic() - t1 - cur_step += 1 - - if cur_step % LOG_INTERVAL == 0: - metric = model.calculate_metric(is_training=True) - logger.info( - f'Epoch {epoch} Step {cur_step}/{total_steps}, ' - f'metric: {metric} | ' - f'encode={t_encode:.2f}s train={t_train:.2f}s') - log_dict = {} - for k, v in metric.items(): - if not v: - continue - try: - log_dict[k] = float(v) - except (ValueError, TypeError): - pass - log_dict['epoch'] = epoch - log_dict['encode_sec'] = round(t_encode, 3) - log_dict['train_sec'] = round(t_train, 3) - swanlab.log(log_dict, step=cur_step) - if cur_step % SAVE_INTERVAL == 0: - save_checkpoint(model, f'step_{cur_step}') - - save_checkpoint(model, 'last-checkpoint') - # Force sync: resolve any pending lazy remote calls (save) before exit - model.calculate_metric(is_training=True) - logger.info(f'Training complete. Final step: {cur_step}') - - -if __name__ == '__main__': - train() diff --git a/cookbook/exp/legacy/train_extract_ddp.py b/cookbook/exp/legacy/train_extract_ddp.py deleted file mode 100644 index 38d3c1f5f..000000000 --- a/cookbook/exp/legacy/train_extract_ddp.py +++ /dev/null @@ -1,119 +0,0 @@ -"""DDP LoRA SFT for the policy on hotpotqa_distractor_reannotated_sft_12k.jsonl. - -The JSONL is the output of ``cookbook/rl/make_condensed_sft.py``: each row -already carries ``messages`` (system / user / assistant with textual -``<tool_call>`` blocks / tool) plus an OpenAI-shape ``tools`` schema, ready -for ``Qwen3_5Template`` to render. ``enable_thinking=False`` matches the -RL runtime contract. - -Launch: - torchrun --nproc_per_node=8 cookbook/rl/train_condensed_sft_ddp.py -""" -from pathlib import Path - -from peft import LoraConfig - -import twinkle -from twinkle import DeviceMesh, get_device_placement, get_logger -from twinkle.dataloader import DataLoader -from twinkle.dataset import Dataset, DatasetMeta -from twinkle.model import TransformersModel - -logger = get_logger() - -MODEL_ID = 'ms://Qwen/Qwen3.5-4B' -DATASET_PATH = str( - Path(__file__).resolve().parent.parent.parent - / 'hotpotqa_distractor_reannotated_sft_12k.jsonl') -TEMPLATE_NAME = 'Qwen3_5Template' -# Multi-hop with compressed context + multi-turn extract_condensed CoT; -# raw audit: most samples land well under 16k after condensation. -MAX_LENGTH = 32000 - -DP_SIZE = 8 -BATCH_SIZE = 16 -LEARNING_RATE = 1e-4 -GRADIENT_ACCUMULATION_STEPS = 2 -LOG_INTERVAL = 20 -NUM_EPOCHS = 2 - -OUTPUT_DIR = './output/condensed_sft_ddp' -RESUME_FROM_CHECKPOINT = None -RESUME_ONLY_MODEL = False -IGNORE_DATA_SKIP = False -ADAPTER_NAME = 'default' - -device_mesh = DeviceMesh.from_sizes(dp_size=DP_SIZE) -twinkle.initialize(mode='local', global_device_mesh=device_mesh) - - -def build_dataset(num_samples: int = None) -> Dataset: - meta_kwargs = {} - if num_samples is not None: - meta_kwargs['data_slice'] = range(num_samples) - dataset = Dataset(dataset_meta=DatasetMeta(DATASET_PATH, **meta_kwargs)) - # ``truncation_strategy='delete'`` drops overlong rows instead of slicing — - # a sliced multi-turn trajectory would lose `\boxed{}` and break SFT signal. - dataset.set_template( - TEMPLATE_NAME, - model_id=MODEL_ID, - max_length=MAX_LENGTH, - truncation_strategy='delete', - enable_thinking=False) - dataset.encode(load_from_cache_file=True, num_proc=16) - return dataset - - -def save_checkpoint(model: TransformersModel, checkpoint_name: str, dataloader: DataLoader): - model.save( - checkpoint_name, - output_dir=OUTPUT_DIR, - adapter_name=ADAPTER_NAME, - save_optimizer=True, - consumed_train_samples=dataloader.get_state()['consumed_train_samples'], - ) - - -def train(): - dataset = build_dataset() - dataloader = DataLoader(dataset=dataset, batch_size=BATCH_SIZE) - - model = TransformersModel(model_id=MODEL_ID, ddp_config={'find_unused_parameters': True}) - model.model._no_split_modules = {'Qwen3_5DecoderLayer'} - - lora_config = LoraConfig(r=16, lora_alpha=32, target_modules='all-linear') - model.add_adapter_to_model(ADAPTER_NAME, lora_config, gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) - model.set_optimizer(optimizer_cls='AdamW', lr=LEARNING_RATE) - model.set_lr_scheduler( - scheduler_cls='CosineWarmupScheduler', - num_warmup_steps=50, - num_training_steps=len(dataloader) * NUM_EPOCHS // GRADIENT_ACCUMULATION_STEPS) - - if RESUME_FROM_CHECKPOINT: - checkpoint_path = Path(RESUME_FROM_CHECKPOINT).expanduser().resolve() - kwargs = {'adapter_name': ADAPTER_NAME} if ADAPTER_NAME else {} - progress = model.resume_from_checkpoint( - str(checkpoint_path), resume_only_model=RESUME_ONLY_MODEL, **kwargs) - if not IGNORE_DATA_SKIP: - dataloader.resume_from_checkpoint(progress['consumed_train_samples']) - - logger.info(get_device_placement()) - logger.info(model.get_train_configs()) - logger.info(f'Total steps: {len(dataloader) * NUM_EPOCHS}') - - optimizer_group = model.optimizer_group[ADAPTER_NAME] - - for epoch in range(NUM_EPOCHS): - for batch in dataloader: - model.forward_backward(inputs=batch) - model.clip_grad_and_step() - cur_step = optimizer_group.cur_step - if cur_step % LOG_INTERVAL == 0: - metric = model.calculate_metric(is_training=True) - logger.info(f'Epoch {epoch} Step {cur_step}/{len(dataloader) * NUM_EPOCHS}, metric: {metric}') - save_checkpoint(model, f'epoch-{epoch}', dataloader) - save_checkpoint(model, 'last-checkpoint', dataloader) - - -if __name__ == '__main__': - train() diff --git a/cookbook/exp/legacy/train_reflexion_skill.py b/cookbook/exp/legacy/train_reflexion_skill.py deleted file mode 100644 index 731478c97..000000000 --- a/cookbook/exp/legacy/train_reflexion_skill.py +++ /dev/null @@ -1,1990 +0,0 @@ -"""Online GRPO RFT for the reflexion skill generator (self-contained, cached). - -Trains an INDEPENDENT skill model to write reusable skills that, injected into a -FROZEN base solver's system prompt, raise its accuracy. The base is never trained; -it only produces the reward. Per chunk: base greedy solve -> rubric process-check -(view A) -> skill-gen (thinking ON, N candidates) -> leak filter -> with-skill -greedy pass -> group-relative advantage -> ONE on-policy GRPO step -> sync weights. -Reward is deterministic (T=0, M=1) binary correctness; A = (R - mean) / (std + eps) -within each problem-group, so std=0 groups give no gradient (GRPO variance selects). - -Each problem is routed to EXACTLY ONE view: A = problem + rubric findings, B = -query only (deployment form). Skill-gen trains only the final structured guidance turn. - -Base greedy rollouts and rubric diagnoses are cached to disk (jsonl, md5-keyed) so -restarts skip them; skill-gen is on-policy and never cached. - -8 GPUs default for high-memory cards: rank 0 trains the actor, rank 1 hosts a -frozen reference model, ranks 2-3 skill_sampler (synced), ranks 4-7 base_sampler -(frozen). Override TRAIN_GPUS / REF_GPUS / SKILL_SAMPLER_GPUS / BASE_SAMPLER_GPUS -for other layouts. -Leak / rubric use the backup teacher API: set LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / -LLM_BACKUP_MODEL. - -Launch: - LLM_BACKUP_API_KEY=... python cookbook/exp/embedding/train_reflexion_skill.py \ - --dataset aops --n 5000 --chunk-size 16 --lr 6e-6 -""" -import argparse -import copy -import hashlib -import json -import math -import os -import re -import sys -import threading -import time -from concurrent.futures import ThreadPoolExecutor -from typing import Any, Dict, List, Optional, Set, Tuple - -import numpy as np - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.checkpoint_engine import CheckpointEngineManager -from twinkle.data_format import SamplingParams, pack_user_data -from twinkle.dataset import Dataset, DatasetMeta -from twinkle.model import TransformersModel -from twinkle.patch.no_split_modules import NoSplitModulesPatch -from twinkle.processor import InputProcessor -from twinkle.sampler import vLLMSampler -from twinkle.template import Template -from twinkle_agentic.verifier import RubricVerifier -from twinkle_agentic.verifier.rubric_verifier import RubricItem - -logger = get_logger() - -try: - import swanlab -except ImportError: - swanlab = None - -MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') -GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.8)) -GEN_TEMPERATURE = float(os.environ.get('GEN_TEMPERATURE', 0.6)) -GEN_TOP_P = float(os.environ.get('GEN_TOP_P', 0.95)) -AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') -MATH_DATASET_ID = os.environ.get('MATH_DATASET_ID', 'modelscope/competition_math') - -# GPU layout: train actor (synced to skill_sampler) + frozen ref model + skill_sampler + base_sampler. -# High-memory cards can keep Qwen3-4B actor/ref as one GPU each and spend more GPUs -# on vLLM data-parallel sampling. The base side is heavier here because every -# clean skill candidate is re-solved by the frozen base model, plus balance/eval baselines. -TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 2)) -REF_GPUS = int(os.environ.get('REF_GPUS', 2)) -SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) -BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 2)) -NUM_GPUS = TRAIN_GPUS + REF_GPUS + SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS -TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', min(1, TRAIN_GPUS))) -REF_FSDP = int(os.environ.get('REF_FSDP', min(1, REF_GPUS))) -if min(TRAIN_GPUS, REF_GPUS, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS, TRAIN_FSDP, REF_FSDP) < 1: - raise ValueError('TRAIN_GPUS, REF_GPUS, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS, TRAIN_FSDP and REF_FSDP must all be >= 1') -if TRAIN_GPUS % TRAIN_FSDP != 0: - raise ValueError(f'TRAIN_GPUS ({TRAIN_GPUS}) must be divisible by TRAIN_FSDP ({TRAIN_FSDP})') -if REF_GPUS % REF_FSDP != 0: - raise ValueError(f'REF_GPUS ({REF_GPUS}) must be divisible by REF_FSDP ({REF_FSDP})') -TRAIN_DP = TRAIN_GPUS // TRAIN_FSDP -REF_DP = REF_GPUS // REF_FSDP - - -# =========================================================================== -# Block A -- boxed extraction + answer grading -# =========================================================================== -_BOXED_RE = re.compile(r'\\boxed\s*\{') - - -def extract_boxed(text: str) -> Optional[str]: - """Last ``\\boxed{...}`` content, brace-balanced.""" - if not text: - return None - last = None - for m in _BOXED_RE.finditer(text): - depth, i = 1, m.end() - while i < len(text) and depth > 0: - depth += (text[i] == '{') - (text[i] == '}') - i += 1 - if depth == 0: - last = text[m.end():i - 1].strip() - return last - - -def normalize_answer(ans: str) -> str: - if not ans: - return '' - s = str(ans).strip() - m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) - if m: - return m.group(1) - s = s.strip('$').strip().replace('−', '-') - 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').replace(r'\tfrac', r'\frac') - s = s.replace(r'\,', '').replace(r'\;', '').replace(r'\!', '') - s = re.sub(r'\\(?:quad|qquad|\s)', '', s) - s = re.sub(r'\s+', '', s) - s = re.sub(r'\\sqrt([0-9A-Za-z])', r'\\sqrt{\1}', s) - s = re.sub(r'\\frac\{([^{}]+)\}([^{}\\])', r'\\frac{\1}{\2}', s) - s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) - s = re.sub(r'\{[a-zA-Z]+\}$', '', s) - s = re.sub(r'\^\\circ|\^\{\\circ\}|°|\\circ', 'deg', s) - s = s.replace(r'\minus{}', '-').replace(r'\minus', '-') - - def _frac_to_slash(mt): - text = mt.group(0) - pos = text.index('{') + 1 - depth, num_start = 1, pos - while depth > 0: - depth += (text[pos] == '{') - (text[pos] == '}') - pos += 1 - numer = text[num_start:pos - 1] - pos += 1 - den_start, depth = pos, 1 - while depth > 0: - depth += (text[pos] == '{') - (text[pos] == '}') - pos += 1 - return f'({numer})/({text[den_start:pos - 1]})' - - s = re.sub(r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', _frac_to_slash, s) - s = re.sub(r'(?<!\w)(\d+)/(\d+)(?!\w)', r'(\1)/(\2)', s) - return s - - -def _try_numeric_equal(a: str, b: str) -> bool: - try: - va, vb = float(a.replace('(', '').replace(')', '')), 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) - return va is not None and vb is not None and abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) - - -_MCQ_REF_RE = re.compile( - r'^\\(?:textbf|text|mathrm|mathbf)\{\(?([A-E])\)?\s*\}\s*(.+)$' - r'|^\(?([A-E])\)\s+(.+)$') -_VAR_PREFIX_RE = re.compile(r'^(?:[A-Za-z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)$', re.DOTALL) - - -def _split_mcq(ans: str): - s = ans.strip() - m = _MCQ_REF_RE.match(s) - if m: - return (m.group(1) or m.group(3)), ((m.group(2) or m.group(4) or '').strip() or None) - bl = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) - return (bl.group(1), None) if bl else (None, s or None) - - -def _strip_var_prefix(ans: str) -> str: - m = _VAR_PREFIX_RE.match((ans or '').strip()) - return m.group(1).strip() if m else (ans or '') - - -def _normalize_tuple(ans: str) -> str: - return re.sub(r'[\s()\[\]{}\\]', '', ans or '') - - -def _math_verify_equal(predicted: str, reference: str) -> bool: - try: - from math_verify import parse, verify - gold = parse(r'\boxed{' + reference + '}', parsing_timeout=1) - pred = parse(r'\boxed{' + predicted + '}', parsing_timeout=1) - return bool(gold and pred and verify(gold, pred, timeout_seconds=1)) - except Exception: - return False - - -def answers_match(predicted: str, reference: str) -> bool: - if not predicted or not reference: - return False - norm_p, norm_r = normalize_answer(predicted), normalize_answer(reference) - if norm_p == norm_r or norm_p.lower() == norm_r.lower() or _try_numeric_equal(norm_p, norm_r): - return True - stripped_p = normalize_answer(_strip_var_prefix(predicted)) - stripped_r = normalize_answer(_strip_var_prefix(reference)) - if stripped_p and stripped_r: - if (stripped_p == stripped_r or stripped_p.lower() == stripped_r.lower() - or _try_numeric_equal(stripped_p, stripped_r)): - return True - p_letter, p_value = _split_mcq(predicted) - r_letter, r_value = _split_mcq(reference) - if p_letter and r_letter and p_letter == r_letter: - return True - if (p_letter and p_value is None) and (r_letter and r_value is None): - return False - 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 and (p_val == r_val or p_val.lower() == r_val.lower() or _try_numeric_equal(p_val, r_val)): - return True - for left, right in ((norm_p, norm_r), (stripped_p, stripped_r)): - tl, tr = _normalize_tuple(left), _normalize_tuple(right) - if ',' in tl and tl == tr: - return True - if '=' in norm_r and '=' not in norm_p: - if any(part and (part == norm_p or _try_numeric_equal(part, norm_p)) for part in norm_r.split('=')): - return True - if '=' in norm_p and '=' not in norm_r: - if any(part and (part == norm_r or _try_numeric_equal(part, norm_r)) for part in norm_p.split('=')): - return True - return _math_verify_equal(stripped_p or norm_p, stripped_r or norm_r) - - -# =========================================================================== -# Block B -- prompts, skill parsing, batched sampling -# =========================================================================== -DIRECT_SYSTEM = ( - 'You are an expert competition mathematician. Solve the following problem ' - 'step by step. Provide your final answer inside \\boxed{}.') - -_SKILL_SOLVE_PREFIX = ( - DIRECT_SYSTEM + '\n\n' - 'Before you start, keep these reminders in mind to avoid common mistakes on this ' - 'type of problem:\n') -_SKILL_SOLVE_SUFFIX = '\nApply them where relevant, but rely on your own reasoning to reach the answer.' - - -def build_direct_prompt(problem: str) -> Dict[str, Any]: - return {'messages': [{'role': 'system', 'content': DIRECT_SYSTEM}, - {'role': 'user', 'content': problem}]} - - -def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: - skill = (skill or '').strip() - if not skill: - return build_direct_prompt(problem) - # Concatenation (not .format): DIRECT_SYSTEM/skill contain literal braces. - return {'messages': [ - {'role': 'system', 'content': _SKILL_SOLVE_PREFIX + skill + _SKILL_SOLVE_SUFFIX}, - {'role': 'user', 'content': problem}]} - - -# -- skill-gen prompts (view A: problem + rubric source; view B: query only) -- -# Kept deliberately short: this is the RL policy's system prompt, so over-specifying -# the output hurts convergence. The concrete output format is appended separately by -# ``_SKILL_OUTPUT`` so the task text stays format-agnostic. -SKILL_GEN_SYSTEM = ( - 'You are a math guidance writer. A process-check on a related problem hints at ' - 'likely mistakes. Write short reusable guidance for this and similar problems, ' - 'and note what to watch out for.\n') - -SKILL_GEN_SYSTEM_Q = ( - 'You are a math guidance writer. Write short reusable guidance for this and ' - 'similar problems.\n') - -_SKILL_OUTPUT = ( - 'Output only:\n<skills>\nYour reusable solving guidance here.\n</skills>') - -SKILL_GEN_USER_Q = ( - 'Problem:\n{problem}\n\n') - -SKILL_GEN_USER_RUBRIC = ( - 'Target problem:\n{problem}\n\n' - 'Problem used for the process check:\n{rubric_problem}\n\n' - 'Process check:\n' - '{diagnosis}\n\n') - - -def _rubric_has_fail(diagnosis: str) -> bool: - """View A uses the open-book rubric prompt (and, under --viewa-sft, SFT distillation) - IFF its rubric localised a failure. No ``[FAIL]`` => nothing to inject: generation - degrades to query-only and the problem is trained by GRPO exactly like view B. Single - source of truth so ``_skillgen_messages`` and ``_group_records`` never diverge.""" - return '[FAIL]' in (diagnosis or '') - - -def _skillgen_messages(problem: str, view: str, diagnosis: str, - rubric_problem: str = '') -> List[Dict[str, Any]]: - """Single source of truth for the skill-gen prompt (used at BOTH generation and - training so they never diverge). View A with a localisable failure uses the target - problem plus the rubric source problem and findings; view B -- or a view-A problem - whose rubric flagged nothing (no ``[FAIL]``) -- degrades to the query-only prompt.""" - if view == 'B' or not _rubric_has_fail(diagnosis): - return [{'role': 'system', 'content': SKILL_GEN_SYSTEM_Q + _SKILL_OUTPUT}, - {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=problem)}] - rubric_problem = rubric_problem or problem - return [{'role': 'system', 'content': SKILL_GEN_SYSTEM + _SKILL_OUTPUT}, - {'role': 'user', 'content': SKILL_GEN_USER_RUBRIC.format( - problem=problem, rubric_problem=rubric_problem, diagnosis=diagnosis)}] - - -def _assign_view(problem: str, args: argparse.Namespace) -> str: - h = int(hashlib.md5(f'{args.seed}:{problem}'.encode('utf-8')).hexdigest(), 16) - return 'B' if (h % 100000) / 100000.0 < args.view_b_frac else 'A' - - -def _curriculum_view_b_frac(gstep: int, args: argparse.Namespace) -> float: - """View-A anneal (--viewa-frac-start): the view-A share holds at ``viewa_frac_start`` - for the first ``viewa_warmup_chunks`` chunks (pure-SFT warmup when start==1.0), then - decays linearly to ``viewa_frac_end`` over ``viewa_decay_chunks`` chunks and holds. - Because _assign_view is a fixed hash against a moving threshold, the B set grows - MONOTONICALLY: a problem trained open-book (A) early can only reappear closed-book - (B) later, never the reverse.""" - t = min(max(gstep - args.viewa_warmup_chunks, 0) / max(args.viewa_decay_chunks, 1), 1.0) - share = args.viewa_frac_start + (args.viewa_frac_end - args.viewa_frac_start) * t - return 1.0 - share - - -def _view_prompt(r: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: - return {'messages': _skillgen_messages( - r['problem'], r['_view'], r.get('_rubric_diag', ''), r.get('_rubric_src', ''))} - - -_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') - - -def _clean_text(decoded: Optional[str]) -> str: - return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() - - -def _extract_tag_block(answer: str, tag: str, allow_empty: bool = False) -> Optional[str]: - low = answer.lower() - open_tag, close_tag = f'<{tag}>', f'</{tag}>' - s = low.rfind(open_tag) - if s < 0: - return None - inner = s + len(open_tag) - e = low.find(close_tag, inner) - if e < 0: - return None - block = answer[inner:e].strip() - block = re.sub(r'</?(?:skills|skill|diagnose|pitfall|strategy|think)>', '', block, flags=re.IGNORECASE).strip() - return block if (block or allow_empty) else None - - -def _extract_skill(text: str) -> Optional[str]: - """Parse skill-generation output: return the inner text of a non-empty ``<skills>`` - block, or None. If a ``</think>`` marker is present, parse only the text after the - last one; otherwise parse the full response.""" - low = text.lower() - end_think = low.rfind('</think>') - answer = text[end_think + len('</think>'):] if end_think >= 0 else text - return _extract_tag_block(answer, 'skills') - - -def _parse_seq(seq, gold: str) -> Dict[str, Any]: - """Grade one sampled sequence into a rollout record.""" - text = _clean_text(getattr(seq, 'decoded', '') or '') - pred = extract_boxed(text) - correct = bool(pred) and answers_match(pred, gold) - terminated = getattr(seq, 'stop_reason', None) != 'length' - return {'pred': pred, 'correct': correct, 'terminated': terminated, - 'passed': bool(correct and terminated), - 'stop_reason': getattr(seq, 'stop_reason', None), - 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} - - -def _run_samples(sampler, prompts: List[Any], num_samples: int, max_tokens: int, - gen_dp: int, temperature: Optional[float] = None, - top_p: Optional[float] = None, top_k: Optional[int] = None) -> List[List[Any]]: - """One batched sampler call -> per-prompt list of raw sequences. vLLM dp needs - batch len >= dp, so pad the tail and slice back.""" - if not prompts: - return [] - params = SamplingParams( - max_tokens=max_tokens, - temperature=GEN_TEMPERATURE if temperature is None else temperature, - top_p=GEN_TOP_P if top_p is None else top_p, - num_samples=num_samples, **({} if top_k is None else {'top_k': top_k})) - padded = prompts - if gen_dp > 1 and 0 < len(prompts) < gen_dp: - padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] - responses = sampler.sample(padded, params)[:len(prompts)] - return [list(r.sequences) if (r and r.sequences) else [] for r in responses] - - -# =========================================================================== -# Block C -- data loading via twinkle.Dataset + numeric filtering -# =========================================================================== -def _boxed_batch(rows: Dict[str, List[Any]], dataset: str) -> Dict[str, List[Any]]: - """Batched (columnar) mapper for ``Dataset.map``: annotate each row with its boxed - ``reference_answer`` and a ``_keep`` flag (aops also needs the ``boxed`` metadata).""" - sols = rows['solution'] - metas = rows.get('metadata', [None] * len(sols)) - refs = [extract_boxed(s or '') for s in sols] - keep = [bool(ref) and (dataset != 'aops' or bool((meta or {}).get('boxed'))) - for ref, meta in zip(refs, metas)] - return {**rows, 'reference_answer': refs, '_keep': keep} - - -def load_problems(dataset: str, n: int, seed: int, num_proc: int = 0) -> List[Dict[str, Any]]: - """Load boxed-answer problems as ``{problem, reference_answer, level?}`` via - twinkle.Dataset (ModelScope hub), sampled to ``n`` (0 = all). Boxed extraction (regex - + brace scan over every solution) is parallelised by ``Dataset.map``/``.filter``; - ``num_proc`` defaults to all cores (set 1 to force serial).""" - ds_id = AOPS_DATASET_ID if dataset == 'aops' else MATH_DATASET_ID - ds = Dataset(DatasetMeta(dataset_id=f'ms://{ds_id}', split='train')) - nproc = num_proc if num_proc > 0 else min(32, os.cpu_count() or 1) - ds.map(lambda rows: _boxed_batch(rows, dataset), num_proc=nproc) - ds.filter(lambda row: row['_keep'], num_proc=nproc) - has_level = 'level' in ds.dataset.column_names - out = [{'data_id': f'{dataset}:{i}', 'problem': row['problem'], - 'reference_answer': row['reference_answer'], - **({'level': row['level']} if has_level and row.get('level') else {})} - for i, row in enumerate(ds.dataset)] - logger.info(f'[data] {dataset}: {len(out)} boxed problems ({nproc} procs)') - rng = np.random.RandomState(seed) - rng.shuffle(out) - return out[:n] if (n and n < len(out)) else out - - -# --------------------------------------------------------------------------- -# Bag-of-words neighbour pairing (cross-problem rubric transfer, --xproblem-rubric) -# --------------------------------------------------------------------------- -# Common English + math-scaffolding words that carry no problem-type signal. Kept small -# and deterministic on purpose (no external stopword list): what survives is the domain -# vocabulary ('triangle', 'prime', 'polynomial', ...) that actually defines the type. -_BOW_STOP = frozenset(""" -a an the of to in on at for and or but if is are be was were been being this that these those -with without into onto from by as it its their his her our your my we you they he she them -find compute determine calculate evaluate solve show prove given let suppose consider assume -what which when where how many much value values number numbers expression form terms term -such that then than so if only when each every all any some both one two three four five six -seven eight nine ten first second third last non over under about above below between -problem answer result equal equals sum difference product total following there here have has -had do does did can could will would should may might must not no yes if then else -""".split()) - -_WORD_RE = re.compile(r'[a-z]+') - - -def _stem(w: str) -> str: - """Crude suffix stripper so 'prime'/'primes', 'triangle'/'triangles' collapse to one - type token. Not linguistically correct -- just enough to merge the common plural/gerund - variants that otherwise split a type's vocabulary and starve the df filter.""" - if len(w) > 4 and w.endswith('ies'): - return w[:-3] + 'y' # properties -> property - if len(w) > 4 and w.endswith('es') and w[-3] in 'sxzh': - return w[:-2] # boxes -> box (keep primes -> prime below) - for suf in ('ing', 'ed', 's'): - if len(w) > len(suf) + 2 and w.endswith(suf): - return w[:-len(suf)] - return w - - -def _tokenize(problem: str) -> List[str]: - """Deterministic bag-of-words tokens for type matching: lowercase alphabetic words - (numbers dropped -- they are instance detail, not type), minus generic stopwords, then - stemmed, so only domain terms remain. Latex control words (frac, sqrt, ...) survive.""" - return [_stem(w) for w in _WORD_RE.findall((problem or '').lower()) - if len(w) > 2 and w not in _BOW_STOP] - - -class BagOfWordsIndex: - """TF-IDF cosine nearest-neighbour over problem statements. Sparse dict vectors + - an inverted index, so ``nearest`` scores only problems sharing a term (near-linear in - practice, no dense NxN). Deterministic; ties break on lower index for reproducibility. - - Answer-aware: ``nearest`` skips the same index, near-verbatim duplicates (cosine - >= ``sim_max``) and -- crucially for anti-leak -- any candidate whose answer equals the - query's, so a neighbour rubric can never hand over the query's own answer.""" - - def __init__(self, problems: List[str], answers: Optional[List[str]] = None, - min_df: int = 2, max_df_frac: float = 0.5): - self._toks = [_tokenize(p) for p in problems] - self._ans = [(_numeric_value(a) or (str(a).strip() if a else '')) for a in answers] \ - if answers is not None else [''] * len(problems) - n = len(self._toks) - df: Dict[str, int] = {} - for toks in self._toks: - for w in set(toks): - df[w] = df.get(w, 0) + 1 - max_df = max(min_df, int(max_df_frac * n)) - self._idf = {w: math.log((n + 1) / (c + 1)) + 1.0 - for w, c in df.items() if min_df <= c <= max_df} - self._vecs: List[Dict[str, float]] = [self._vectorize(t) for t in self._toks] - self._inverted: Dict[str, List[int]] = {} - for i, v in enumerate(self._vecs): - for w in v: - self._inverted.setdefault(w, []).append(i) - - def _vectorize(self, toks: List[str]) -> Dict[str, float]: - tf: Dict[str, float] = {} - for w in toks: - if w in self._idf: - tf[w] = tf.get(w, 0.0) + 1.0 - vec = {w: c * self._idf[w] for w, c in tf.items()} - norm = math.sqrt(sum(x * x for x in vec.values())) - return {w: x / norm for w, x in vec.items()} if norm > 0 else {} - - def nearest(self, i: int, sim_max: float = 0.98) -> Tuple[int, float]: - """(index, cosine) of the most similar *distinct* problem: not i, not a duplicate - (cosine < ``sim_max``) and not the same answer. (-1, 0.0) if none qualifies.""" - vi = self._vecs[i] - if not vi: - return -1, 0.0 - ai = self._ans[i] - scores: Dict[int, float] = {} - for w, xi in vi.items(): - for j in self._inverted.get(w, ()): - if j != i: - scores[j] = scores.get(j, 0.0) + xi * self._vecs[j].get(w, 0.0) - best_j, best_s = -1, 0.0 - for j, s in scores.items(): - if s >= sim_max or (ai and self._ans[j] == ai): - continue - if s > best_s or (s == best_s and (best_j < 0 or j < best_j)): - best_j, best_s = j, s - return best_j, best_s - - -def build_pairs(records: List[Dict[str, Any]], n: int, seed: int, sim_max: float = 0.98 - ) -> Tuple[List[Dict[str, Any]], Dict[str, Tuple[str, float]]]: - """Single-pass cross-problem pairing over the whole pool (one index build). - - Returns ``(subset, neighbour_map)`` where ``subset`` is the ``n`` problems with the - strongest qualifying neighbour (dense same-type pairs; n<=0/n>=len keeps all, shuffled) - and ``neighbour_map[q] = (p, cosine)`` gives each kept problem its analogue P. P is drawn - from the *full* pool (richer analogues) but never a duplicate or same-answer problem, so - P's rubric can transfer method without ever leaking Q's answer.""" - index = BagOfWordsIndex([r['problem'] for r in records], - [str(r.get('reference_answer', '')) for r in records]) - nbr = [index.nearest(i, sim_max) for i in range(len(records))] - order = sorted(range(len(records)), key=lambda i: (-nbr[i][1], i)) - keep = order[:n] if (0 < n < len(records)) else list(range(len(records))) - rng = np.random.RandomState(seed) - rng.shuffle(keep) - subset = [records[i] for i in keep] - neighbour_map = {records[i]['problem']: (records[nbr[i][0]]['problem'], nbr[i][1]) - for i in keep if nbr[i][0] >= 0} - return subset, neighbour_map - - -_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') - - -def _norm_num_text(num: str) -> str: - try: - f = float(num) - return str(int(f)) if f == int(f) else str(f) - except Exception: - return str(num).strip() - - -def _numeric_value(raw: Any) -> Optional[str]: - """Collapse an answer to a single int/decimal/fraction, or None.""" - if raw is None: - return None - s = str(raw).strip().strip('$').strip() - s = s.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac') - s = re.sub(r'\\!|\\,|\\;|\\ |\\left|\\right|\s', '', s) - for pat in (r'\\frac\{(-?\d+)\}\{(-?\d+)\}', r'(-?\d+)/(-?\d+)'): - m = re.fullmatch(pat, s) - if m: - a, b = int(m.group(1)), int(m.group(2)) - return _norm_num_text(str(a / b)) if b else None - return _norm_num_text(s) if _NUM_RE.fullmatch(s) else None - - -def _answer_leaked(skill: str, reference: str) -> bool: - """Audit whether a generated skill contains the final answer verbatim. This is NOT - a training filter: if the skill model derives an answer from the problem, that is a - legitimate answer-bearing skill under this experiment. The real leakage boundary is the - external API/rubric diagnosis, which is constrained by prompt to stay answer-free.""" - if not skill: - return False - for cand in {_numeric_value(reference), (str(reference).strip() or None)}: - if cand and re.search(r'(?<![\d.])' + re.escape(cand) + r'(?![\d.])', skill): - return True - return False - - -def _load_excluded_records(paths_arg: str) -> Tuple[Set[str], Set[str]]: - """Read jsonl files and collect data_id/problem keys that must be excluded. - - The cold-start SFT builder writes stable ``data_id`` values. ``problem`` is kept as a - backward-compatible fallback for older jsonl files produced before data_id existed.""" - ids: Set[str] = set() - problems: Set[str] = set() - for raw_path in (paths_arg or '').split(','): - path = raw_path.strip() - if not path or not os.path.exists(path): - continue - with open(path, encoding='utf-8') as f: - for line in f: - if not line.strip(): - continue - row = json.loads(line) - if row.get('record_type') in {'config', 'summary'}: - continue - data_id = str(row.get('data_id') or '').strip() - problem = str(row.get('problem') or '').strip() - if data_id: - ids.add(data_id) - elif problem: - problems.add(problem) - return ids, problems - - -def _load_records(args: argparse.Namespace - ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], - Dict[str, Tuple[str, float]], Dict[str, str], Dict[str, int]]: - """Load, numeric-filter, shuffle, split a fixed eval holdout, and (when --xproblem-rubric) - select a same-type-dense train subset with its neighbour map -- all in one pass. - Also returns ``pool_answers`` (every candidate neighbour P's true answer) so P's baseline - can be graded/cached correctly even when P is not itself a training problem.""" - # Load all when filtering or splitting (else the eval holdout could starve train). - load_n = 0 if (args.numeric_only or args.eval_size > 0) else args.n - records = load_problems(args.dataset, load_n, args.seed) - raw_n, dropped = len(records), 0 - if args.numeric_only: - records = [{**r, 'reference_answer': v} - for r, v in ((r, _numeric_value(r.get('reference_answer'))) for r in records) - if v is not None] - dropped = raw_n - len(records) - np.random.RandomState(args.seed).shuffle(records) - exclude_ids, exclude_problems = _load_excluded_records(getattr(args, 'exclude_data_ids', '')) - excluded = 0 - if exclude_ids or exclude_problems: - before = len(records) - records = [r for r in records - if str(r.get('data_id', '')) not in exclude_ids - and str(r.get('problem', '')).strip() not in exclude_problems] - excluded = before - len(records) - eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 - eval_records = [dict(r) for r in records[:eval_n]] - pool = records[eval_n:] - pool_offset = max(0, int(getattr(args, 'pool_offset', 0) or 0)) - if pool_offset: - if pool_offset >= len(pool): - raise ValueError(f'--pool-offset ({pool_offset}) leaves no train records from pool size {len(pool)}') - pool = pool[pool_offset:] - train_n = args.n if args.n > 0 else len(pool) - # Cross-problem rubric: pick the train_n problems with the strongest qualifying neighbour - # (dense same-type pairs) and build {Q -> (P, sim)} in one index pass. Otherwise keep the - # first train_n (already shuffled) with no neighbours. - if args.xproblem_rubric: - subset, neighbor_map = build_pairs(pool, train_n, args.seed) - train_records = [dict(r) for r in subset] - pool_answers = {r['problem']: str(r.get('reference_answer', '')) for r in pool} - else: - train_records, neighbor_map, pool_answers = [dict(r) for r in pool[:train_n]], {}, {} - if {r['problem'] for r in train_records} & {r['problem'] for r in eval_records}: - raise ValueError('eval/train overlap detected') - stats = {'raw_loaded': raw_n, 'numeric_dropped': dropped, - 'excluded_records': excluded, 'pool_offset': pool_offset, - 'train_records': len(train_records), 'eval_records': len(eval_records)} - return train_records, eval_records, neighbor_map, pool_answers, stats - - -# =========================================================================== -# Block D -- disk cache, problem pool, baseline rollout, rubric check -# =========================================================================== -class DiskCache: - """Append-only jsonl kv cache (md5 key -> value). Loads on init, appends on put. - Disabled instances always miss and never write.""" - - def __init__(self, path: str, enabled: bool = True): - self._mem: Dict[str, Any] = {} - self._fh = None - self._lock = threading.Lock() # base baseline is prefetched on a background thread - if not enabled: - return - if os.path.exists(path): - with open(path, encoding='utf-8') as f: - for line in f: - if line.strip(): - row = json.loads(line) - self._mem[row['key']] = row['value'] - self._fh = open(path, 'a', encoding='utf-8') - - @staticmethod - def key_for(*parts: str) -> str: - return hashlib.md5('\x1f'.join(parts).encode('utf-8')).hexdigest() - - def __contains__(self, key: str) -> bool: - with self._lock: - return key in self._mem - - def get(self, key: str) -> Any: - with self._lock: - return self._mem.get(key) - - def put(self, key: str, value: Any) -> None: - with self._lock: - self._mem[key] = value - if self._fh is not None: - self._fh.write(json.dumps({'key': key, 'value': value}, ensure_ascii=False) + '\n') - self._fh.flush() - - def close(self) -> None: - if self._fh is not None: - self._fh.close() - - -class _LockedSampler: - """Serialises every ``.sample()`` on a sampler behind one lock. The base sampler is - shared by the main thread (with-skill scoring, eval) and the baseline-prefetch thread; - ``sample`` is a slice_dp/flatten remote call whose collect is NOT safe to interleave - across two callers, so concurrent calls could mis-join sequences. The lock keeps base - calls serial (prefetch still overlaps the skill-gen phase, which uses skill_sampler).""" - - def __init__(self, sampler): - self._sampler = sampler - self._lock = threading.Lock() - - def sample(self, *args, **kwargs): - with self._lock: - return self._sampler.sample(*args, **kwargs) - - def __getattr__(self, name): - return getattr(self._sampler, name) - - -class ProblemPool: - """Cyclic draw source. Each full pass reshuffles with ``seed + epoch``; the initial - pass keeps the loader order. ``draw(k)`` returns k distinct problems (k << pool).""" - - def __init__(self, records: List[Dict[str, Any]], seed: int): - self._records = list(records) - self._seed, self._cursor, self.epoch = seed, 0, 0 - - def draw(self, k: int) -> List[Dict[str, Any]]: - out, seen = [], set() - while len(out) < k: - if self._cursor >= len(self._records): - self.epoch += 1 - np.random.RandomState(self._seed + self.epoch).shuffle(self._records) - self._cursor = 0 - r = self._records[self._cursor] - self._cursor += 1 - if id(r) not in seen: - seen.add(id(r)) - out.append(r) - return out - - def peek(self, k: int) -> List[Dict[str, Any]]: - """The next k distinct problems draw() would return, WITHOUT advancing state - (no cursor move, no reshuffle). Used to prefetch their baseline into base_cache - while the current chunk trains -- a wrong guess (cross-epoch reshuffle) only - misses the cache, never corrupts the draw.""" - out, seen, cur = [], set(), self._cursor - recs = self._records - while len(out) < k and cur < len(recs): # stop at epoch edge; don't simulate reshuffle - r = recs[cur] - cur += 1 - if id(r) not in seen: - seen.add(id(r)) - out.append(r) - return out - - -def _empty_roll() -> Dict[str, Any]: - return {'pred': '', 'correct': False, 'terminated': False, 'passed': False, - 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} - - -def _apply_baseline(r: Dict[str, Any], roll: Dict[str, Any]) -> None: - """Attach a greedy baseline roll and reset per-chunk working state.""" - r['_baseline_rolls'], r['_cands'], r['_init'] = [roll], [], [roll] - r['_failed'] = not roll['correct'] - r['_baseline_pass'] = 1.0 if roll['correct'] else 0.0 - r['_hard'] = True # process every problem; group variance selects (SEAM-style) - - -def baseline_rollout(base_sampler, problems: List[Dict[str, Any]], base_dp: int, - args: argparse.Namespace, cache: DiskCache) -> int: - """Base solves each problem greedily once (T=0, M=1), disk-cached by problem text. - The base is frozen + greedy so a cache hit is exact. Returns fresh (miss) count.""" - todo = [r for r in problems if DiskCache.key_for(r['problem']) not in cache] - if todo: - out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], - 1, args.max_tokens, base_dp, temperature=0.0) - for r, seqs in zip(todo, out): - roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - cache.put(DiskCache.key_for(r['problem']), roll) - for r in problems: - _apply_baseline(r, cache.get(DiskCache.key_for(r['problem']))) - return len(todo) - - -# -- rubric process-check (view A): teacher diagnoses the base's attempt -- -_RFT_DIAG_SYSTEM = """\ -You are a strategy-level process checker for a math solution attempt. You are given a -math problem, a rubric, and one attempted solution segment. Decide PASS or FAIL for each -criterion, and write the diagnosis so it can become useful reusable guidance for solving -similar problems without seeing this segment. - -Output STRICT JSON (no prose outside it) with this shape: -{ - "items": [ - {"index": 1, "verdict": "PASS", "reason": "<why the process satisfies it>", - "fix": ""}, - {"index": 2, "verdict": "FAIL", "reason": "<what reusable process issue is present>", - "fix": "<local strategy correction, without solving the problem>"} - ], - "overall": "OK" | "ISSUES", - "summary": "<one sentence naming the reusable process issue, not the answer>" -} - -Rules: -- Judge every criterion independently; a [Hard Rule] is FAIL unless clearly satisfied. -- Judge ONLY what is observable in THIS segment. Ignore hidden <think> or <thinking> - content for output-format criteria. -- The API diagnosis is an external teacher signal, so it must stay answer-free. -- Prefer diagnosis that transfers to view-B skill generation: name the route choice, - structural observation, missing check, or length-control habit that a solver should - remember before solving a similar problem. -- For PASS items, leave "fix" as "". -- For FAIL items, describe the process problem at strategy level: unsuitable method, - missed structure, invalid transformation, missing constraint check, redundant cases, - off-track approach, contradiction, or inefficient/unfinished reasoning. -- A fix may suggest the LOCAL correction direction, such as identify the key structure, - verify constraints, preserve equivalence, reduce redundant cases, or choose a more - direct route. Do not carry out the correction. -- Never reveal the final answer, a corrected value/expression, an option label, or a - step-by-step solution that would let another model copy the solve. -- If the segment contains a process note saying it was cut off before a final boxed - answer, mark the length-budget criterion as FAIL and suggest a method-level way to - finish faster. -- Keep every "reason" and "fix" concise: one short sentence each. -- "overall" is "OK" only if NO criterion is FAIL. -- Output only the JSON object.""" - -_RFT_DIAG_USER = """\ -## Task / query (context) -{query} - -## Rubric -{rubric} - -## Segment -{segment} - -Now output the diagnostic JSON object.""" - -_MATH_RUBRIC = [ - ('The attempt chooses a method suitable for the problem structure', False), - ('The attempt identifies the key constraint, invariant, or quantity before computing', False), - ('Algebraic and logical transformations preserve validity at each step', True), - ('The attempt checks required constraints, domains, boundary cases, or validity conditions', False), - ('The attempt avoids redundant casework, looping, or re-deriving known facts', False), - ('The attempt reaches a final boxed answer within the length budget', False), - ('The approach stays focused on the actual question asked', False), -] - -# Bump when the rubric criteria or the diagnosis prompt change so the disk-cached -# diagnoses written under an older rubric are not silently reused. -_RUBRIC_VERSION = 'rubric_v5_viewb_strategy' - - -class _RftRubricVerifier(RubricVerifier): - def _diagnose_trajectory(self, query: str, rubric_block: str, segment_text: str) -> dict: - return {'messages': [ - {'role': 'system', 'content': _RFT_DIAG_SYSTEM}, - {'role': 'user', 'content': _RFT_DIAG_USER.format( - query=query, rubric=rubric_block, segment=segment_text)}]} - - -def build_rubric_checker() -> Optional[RubricVerifier]: - """Fixed math-process rubric verifier (teacher-served). None if no LLM backup env.""" - if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') - or os.environ.get('OPENAI_API_KEY')): - return None - return _RftRubricVerifier( - fixed_rubric=[RubricItem(t, is_hard=h) for t, h in _MATH_RUBRIC], gate=True) - - -def _format_diagnosis(detail) -> str: - """One line per criterion (PASS/FAIL + reason + fix) then a summary.""" - rub = detail.rubric - lines = [] - for it in detail.items: - text = rub[it.index - 1].text if 0 < it.index <= len(rub) else f'criterion {it.index}' - if it.verdict: - lines.append(f'- [PASS] {text}') - else: - tail = f': {it.reason}' if it.reason else '' - tail += f' (fix: {it.fix})' if it.fix else '' - lines.append(f'- [FAIL] {text}{tail}') - if detail.summary: - lines.append(f'Summary: {detail.summary}') - return '\n'.join(lines) - - -def diagnose_views(checker, problems: List[Dict[str, Any]], args: argparse.Namespace, - cache: DiskCache) -> None: - """Rubric-check every view-A problem's greedy attempt in parallel (disk-cached by - problem + attempt), stashing the formatted findings on ``r['_rubric_diag']``.""" - targets = [r for r in problems if r.get('_view') == 'A'] - if not checker or not targets: - return - - def _key(r: Dict[str, Any]) -> str: - init = r.get('_init', [{}])[0] - term = 'L' if (init.get('stop_reason') == 'length' or not init.get('terminated')) else 'T' - return DiskCache.key_for(r['problem'], init.get('text', ''), _RUBRIC_VERSION, term) - - pending = [] - for r in targets: - key = _key(r) - if key in cache: - r['_rubric_diag'] = cache.get(key) - else: - pending.append((r, key)) - if not pending: - return - - def _run(item): - r, key = item - init = r['_init'][0] - seg_text = init['text'] - if init.get('stop_reason') == 'length' or not init.get('terminated'): - seg_text += ('\n\n[Process note: this attempt was cut off at the token budget ' - 'and never produced a final \\boxed{} answer.]') - seg = {'messages': [{'role': 'user', 'content': r['problem']}, - {'role': 'assistant', 'content': seg_text}]} - attempts = max(1, args.rubric_retries + 1) - for attempt in range(attempts): - try: - return r, key, _format_diagnosis(checker.diagnose(seg, query=r['problem'])) - except Exception as exc: # teacher hiccup -> retry, then no-diagnosis prompt (not cached) - if attempt + 1 < attempts: - logger.warning(f'[rubric] diagnose error: {exc}; retry {attempt + 1}/{args.rubric_retries}') - time.sleep(min(2.0, 0.5 * (2 ** attempt))) - continue - logger.warning(f'[rubric] diagnose error: {exc}; giving up after {attempts} attempts') - return r, key, None - - workers = max(1, min(args.rubric_workers, len(pending))) - with ThreadPoolExecutor(max_workers=workers) as ex: - for r, key, diag in ex.map(_run, pending): - r['_rubric_diag'] = diag or '' - if diag is not None: - cache.put(key, diag) - - -# =========================================================================== -# Block E -- chunk draw, generation pipeline, record building -# =========================================================================== -def _baseline_class(r: Dict[str, Any]) -> str: - """success | fail_loop (out of length / never terminated) | fail_wrong.""" - roll = r['_init'][0] - if roll['correct']: - return 'success' - return 'fail_loop' if (roll['stop_reason'] == 'length' or not roll['terminated']) else 'fail_wrong' - - -def _select_balanced(buckets: Dict[str, List[Dict[str, Any]]], n_success: int, - n_fail: int, n_fail_loop: int) -> List[Dict[str, Any]]: - """Pick n_fail base-fails (toward n_fail_loop loop-fails, best-effort) + n_success - base-successes; top up any shortfall from leftovers.""" - loop, wrong, succ = buckets['fail_loop'], buckets['fail_wrong'], buckets['success'] - take_loop = min(n_fail_loop, len(loop)) - take_wrong = min(n_fail - take_loop, len(wrong)) - take_loop = min(n_fail - take_wrong, len(loop)) - sel = loop[:take_loop] + wrong[:take_wrong] + succ[:n_success] - target = n_success + n_fail - if len(sel) < target: - used = {id(x) for x in sel} - sel += [x for b in (loop, wrong, succ) for x in b if id(x) not in used][:target - len(sel)] - return sel - - -def draw_chunk(pool: ProblemPool, base_sampler, base_dp: int, args: argparse.Namespace, - cache: DiskCache) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: - """Draw one chunk, baselining every drawn problem. With ``--balance``, keep - drawing+baselining until the target base fail:success mix is reachable (or the budget - is hit), then select a balanced subset.""" - if not args.balance: - chunk = pool.draw(args.chunk_size) - fresh = baseline_rollout(base_sampler, chunk, base_dp, args, cache) - return chunk, {'enabled': False, 'n_drawn': len(chunk), 'n_baseline_fresh': fresh} - - n_success = max(0, min(args.chunk_size, round(args.chunk_size * args.balance_success_frac))) - n_fail = args.chunk_size - n_success - n_fail_loop = round(n_fail * args.balance_loop_frac) - buckets: Dict[str, List[Dict[str, Any]]] = {'success': [], 'fail_loop': [], 'fail_wrong': []} - budget, n_drawn, n_fresh, seen = args.chunk_size * args.balance_max_draws_mult, 0, 0, set() - while n_drawn < budget: - if (len(buckets['success']) >= n_success - and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail): - break - batch = pool.draw(args.chunk_size) - n_fresh += baseline_rollout(base_sampler, batch, base_dp, args, cache) - n_drawn += len(batch) - for r in batch: - if id(r) not in seen: - seen.add(id(r)) - buckets[_baseline_class(r)].append(r) - - reached = (len(buckets['success']) >= n_success - and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail) - chunk = _select_balanced(buckets, n_success, n_fail, n_fail_loop) - sel_success = sum(1 for r in chunk if not r['_failed']) - stats = { - 'enabled': True, 'n_drawn': n_drawn, 'n_baseline_fresh': n_fresh, 'n_selected': len(chunk), - 'target_success': n_success, 'target_fail': n_fail, 'target_fail_loop': n_fail_loop, - 'selected_success': sel_success, 'selected_fail': len(chunk) - sel_success, - 'selected_fail_loop': sum(1 for r in chunk if _baseline_class(r) == 'fail_loop'), - 'selected_fail_wrong': sum(1 for r in chunk if _baseline_class(r) == 'fail_wrong'), - 'selected_success_frac': (sel_success / len(chunk)) if chunk else 0.0, - 'budget_hit': not reached, - } - return chunk, stats - - -def _assign_advantages(hard: List[Dict[str, Any]], args: argparse.Namespace) -> None: - """Group-relative advantage over each problem's scored candidates using the greedy - binary reward R in {0,1}: ``A = (R - mean) / (std + eps)``. std==0 groups get no - gradient -- GRPO's variance selects informative problems (no explicit difficulty gate). - A symmetric clip limits both positive and negative outliers from sparse 1-vs-many groups.""" - eps = 1e-6 - adv_clip = abs(float(getattr(args, 'adv_clip', 0.0) or 0.0)) - for r in hard: - for c in r['_cands']: - c['advantage'], c['grpo_adv'], c['kept'] = 0.0, 0.0, False - cs = [c for c in r['_cands'] if c.get('reward') is not None] - if len(cs) < 2: - continue - rewards = [c['reward'] for c in cs] - mean_r = sum(rewards) / len(rewards) - std = (sum((x - mean_r) ** 2 for x in rewards) / len(rewards)) ** 0.5 - if std < 1e-9: - continue - for c in cs: - raw_adv = (c['reward'] - mean_r) / (std + eps) - adv = max(-adv_clip, min(adv_clip, raw_adv)) if adv_clip > 0 else raw_adv - c['advantage'], c['grpo_adv'], c['kept'] = adv, adv, c['reward'] > mean_r - - -def _best_sft_candidate(r: Dict[str, Any], args: argparse.Namespace) -> Optional[Dict[str, Any]]: - """Pick ONE view-A candidate to distill (online context distillation). ONLY - executor-verified PASSING skills (reward==1) are distilled: the earlier fallback to - unverified skills meant ~60% of SFT targets had failed their own executor pass - (measured on the sft35 run) and the model was imitating plausible-but-wrong skills. - Problems with no passing candidate now yield NO SFT record. Answer-bearing skills - produced by the skill model itself are allowed here; only the external API/rubric - diagnosis must be answer-free. Among the passing candidates, take the one whose skill - length is CLOSEST to ``--sft-target-len`` -- an empirically high-pass-rate length - (~500-600 chars in this run) -- breaking ties by the fewest executor solve tokens. - Targeting a length (rather than the minimum) avoids a distillation feedback loop that - would otherwise drive rollouts ever shorter.""" - eligible = [c for c in r['_cands'] if c.get('parseable') and c.get('skills')] - passing = [c for c in eligible if c.get('reward') == 1.0] - if not passing: - return None - target = int(getattr(args, 'sft_target_len', 550) or 550) - - def _solve_tokens(c: Dict[str, Any]) -> int: - rolls = c.get('rolls') or [] - return rolls[0].get('gen_tokens', 1 << 30) if rolls else (1 << 30) - - return min(passing, key=lambda c: (abs(len(c['skills']) - target), _solve_tokens(c))) - - -def apply_neighbor_rubric(base_sampler, chunk: List[Dict[str, Any]], - neighbor_map: Dict[str, Tuple[str, float]], - pool_answers: Dict[str, str], base_dp: int, - args: argparse.Namespace, checker, - base_cache: DiskCache, rubric_cache: DiskCache) -> None: - """Cross-problem rubric (--xproblem-rubric): for every view-A problem Q, replace its own - rubric with the rubric of its bag-of-words neighbour P (Q keeps being the solved/scored - problem). P is baselined + diagnosed here (both disk-cached) as a stub carrying P's REAL - answer, so P's baseline grades correctly and legitimately shares the baseline cache with - P-as-training-problem. P's findings are copied onto Q with the neighbour text + similarity - for audit. A P that can't be diagnosed -> Q degrades to view B. Since P's answer differs - from Q's (guaranteed at pairing time), any answer P's rubric leaks is useless for Q.""" - targets = [r for r in chunk if r.get('_view') == 'A' and neighbor_map.get(r['problem'])] - if not targets: - return - stubs, by_problem = [], {} - for r in targets: - p, _ = neighbor_map[r['problem']] - if p not in by_problem: - stub = {'problem': p, 'reference_answer': pool_answers.get(p, ''), '_view': 'A'} - by_problem[p] = stub - stubs.append(stub) - baseline_rollout(base_sampler, stubs, base_dp, args, base_cache) - diagnose_views(checker, stubs, args, rubric_cache) - for r in targets: - p, sim = neighbor_map[r['problem']] - r['_rubric_diag'] = by_problem[p].get('_rubric_diag', '') - r['_rubric_src'], r['_neighbor_sim'] = p, sim - - -def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], - ci: int, base_dp: int, skill_dp: int, args: argparse.Namespace, - checker, rubric_cache: DiskCache, base_cache: DiskCache = None, - neighbor_map: Optional[Dict[str, Tuple[str, float]]] = None, - pool_answers: Optional[Dict[str, str]] = None - ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], List[Dict[str, Any]]]: - """view assign -> rubric-check (view A) -> skill-gen -> answer audit -> with-skill - greedy pass -> GRPO advantages. ``chunk`` arrives already baselined by draw_chunk. - With --xproblem-rubric, view A gets its NEIGHBOUR's rubric (see apply_neighbor_rubric).""" - hard = chunk - for r in hard: - r['_view'], r['_rubric_diag'] = _assign_view(r['problem'], args), '' - if args.xproblem_rubric and neighbor_map: - apply_neighbor_rubric(base_sampler, hard, neighbor_map, pool_answers or {}, base_dp, - args, checker, base_cache, rubric_cache) - else: - diagnose_views(checker, hard, args, rubric_cache) - - # skill-gen (thinking OFF), per-view prompt; re-sample problems with no clean candidate. - # View-A problems that cannot yield a training record (no [FAIL] rubric, or rubric - # leaked the answer) are dropped from training entirely -- skip their generation. - flat: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] - pending = [r for r in hard if not _viewa_dropped(r, args)] - for _ in range(args.skill_retries + 1): - if not pending: - break - sg_out = _run_samples(skill_sampler, [_view_prompt(r, args) for r in pending], - args.n_skills, args.skill_max_tokens, skill_dp, - temperature=args.skill_gen_temperature, - top_p=args.skill_gen_top_p, top_k=args.skill_gen_top_k) - still = [] - for r, seqs in zip(pending, sg_out): - got = False - for s in seqs: - resp = _clean_text(getattr(s, 'decoded', '') or '') - block = _extract_skill(resp) or '' - cand = {'skills': block, 'response': resp, 'parseable': bool(block), - 'view': r['_view'], 'leaked': None, 'leak_reason': '', - 'leak_source': '', 'with_pass': None, 'reward': None, 'rolls': [], - 'skillgen_stop': getattr(s, 'stop_reason', None), - 'skillgen_tokens': len(getattr(s, 'tokens', None) or [])} - r['_cands'].append(cand) - if block: - flat.append((r, cand)) - got = True - if not got: - still.append(r) - pending = still - - # Answer-bearing skill audit: deterministic verbatim-answer check only (no LLM). This - # is observability only; it records metrics for swanlab/jsonl, but does not block - # scoring, reward, advantage assignment, SFT candidate selection, or GRPO training. - for r, c in flat: - leaked = _answer_leaked(c['skills'], r['reference_answer']) - c['leaked'] = leaked - c['leak_reason'] = 'answer_verbatim' if leaked else '' - c['leak_source'] = 'deterministic' - - # with-skill greedy pass (T=0, M=1); reward = correct, absolute (group mean is baseline). - scored_inputs = flat - if scored_inputs: - ws_out = _run_samples(base_sampler, - [build_skill_solve_prompt(r['problem'], c['skills']) for r, c in scored_inputs], - 1, args.max_tokens, base_dp, temperature=0.0) - for (r, c), seqs in zip(scored_inputs, ws_out): - c['rolls'] = [_parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll()] - c['with_pass'] = 1.0 if c['rolls'][0]['correct'] else 0.0 - c['reward'] = c['with_pass'] - if args.format_in_reward: # unparseable candidates score 0 and still join the group - for r in hard: - for c in r['_cands']: - if c['reward'] is None: - c['reward'] = 0.0 - - _assign_advantages(hard, args) - return ([_full_record(r, ci) for r in chunk], _chunk_summary(chunk, ci, args), - _group_records(chunk, args)) - - -def _roll(x: Dict[str, Any]) -> Dict[str, Any]: - return {k: x[k] for k in ('pred', 'correct', 'terminated', 'passed', - 'stop_reason', 'gen_tokens', 'text')} - - -def _is_trainable(c: Dict[str, Any], args: argparse.Namespace) -> bool: - """Reaches the GRPO update iff advantage is non-zero. Leak flags are audit-only.""" - adv_nz = abs(c.get('advantage') or 0.0) > 1e-9 - if args.format_in_reward: - return adv_nz - return c.get('with_pass') is not None and adv_nz - - -def _full_record(r: Dict[str, Any], ci: int) -> Dict[str, Any]: - """Complete per-problem trace: init attempt, baseline, and all candidates.""" - init = r['_init'][0] - return { - 'record_type': 'problem', 'chunk': ci, 'problem': r['problem'], - 'reference_answer': r['reference_answer'], 'level': r.get('level', ''), - 'failed_first_try': r['_failed'], - 'init_attempt': {'text': init['text'], 'pred': init['pred'], 'correct': init['correct'], - 'terminated': init['terminated'], 'stop_reason': init['stop_reason'], - 'gen_tokens': init['gen_tokens']}, - 'baseline_pass': r['_baseline_pass'], 'is_hard': r['_hard'], - 'view': r.get('_view', ''), 'rubric_diag': r.get('_rubric_diag', ''), - # xproblem audit: which neighbour's rubric this problem borrowed, and how similar. - 'rubric_src': r.get('_rubric_src', ''), 'neighbor_sim': r.get('_neighbor_sim'), - 'baseline_rolls': [_roll(x) for x in r['_baseline_rolls']], - 'candidates': [{ - 'skills': c['skills'], 'response': c['response'], 'parseable': c['parseable'], - 'skillgen_stop': c.get('skillgen_stop'), 'skillgen_tokens': c.get('skillgen_tokens'), - 'leaked': c['leaked'], 'leak_reason': c['leak_reason'], 'leak_source': c['leak_source'], - 'with_pass': c['with_pass'], 'reward': c.get('reward'), 'advantage': c.get('advantage'), - 'grpo_adv': c.get('grpo_adv'), 'kept': c.get('kept'), - 'rolls': [_roll(x) for x in c['rolls']], - } for c in r['_cands']], - } - - -def _view_stats(problems: List[Dict[str, Any]], view: str) -> Dict[str, Any]: - pv = [r for r in problems if r.get('_view') == view] - cands = [c for r in pv for c in r['_cands'] if c['parseable']] - clean = [c for c in cands if c['leaked'] is False] - adopted = sum(1 for r in pv - if any(abs(c.get('advantage') or 0.0) > 1e-9 for c in r['_cands'])) - return {'n': len(pv), 'n_candidates_parseable': len(cands), 'n_clean': len(clean), - 'n_adopted_problems': adopted, 'adoption_rate': (adopted / len(pv)) if pv else 0.0} - - -def _mean(xs: List[float]) -> float: - return sum(xs) / len(xs) if xs else 0.0 - - -def _std(xs: List[float]) -> float: - if len(xs) < 2: - return 0.0 - m = _mean(xs) - return (sum((x - m) ** 2 for x in xs) / len(xs)) ** 0.5 - - -def _signal_stats(problems: List[Dict[str, Any]]) -> Dict[str, float]: - """The heart of 'is there a learning signal': per problem, the scored candidates form a - GRPO group. A group with zero reward variance (all skills solve, or none do -- the - hard-problem dead zone) gives NO gradient. Tracks that fraction plus reward level and - within-group variance so a collapse (all-0 or all-1) is visible immediately.""" - group_vars, all_rewards, zero_grad, groups = [], [], 0, 0 - for r in problems: - rewards = [c['reward'] for c in r['_cands'] if c.get('reward') is not None] - if len(rewards) < 2: - continue - groups += 1 - all_rewards.extend(rewards) - v = _std(rewards) - group_vars.append(v) - if v < 1e-9: # every skill got the same reward -> GRPO skips this problem - zero_grad += 1 - return {'n_groups': groups, 'zero_grad_frac': (zero_grad / groups) if groups else 0.0, - 'reward_mean': _mean(all_rewards), 'reward_std': _std(all_rewards), - 'group_reward_std_mean': _mean(group_vars)} - - -def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespace) -> Dict[str, Any]: - all_cands = [c for r in chunk for c in r['_cands']] - cands = [c for c in all_cands if c['parseable']] - scored = [c for c in cands if c['with_pass'] is not None] - clean = [c for c in cands if c['leaked'] is False] - ws_rolls = [x for c in scored for x in c['rolls']] - # viewa-dropped problems generate no candidates; keep acc/* on the generated subset - # so the with-skill/lift trend stays comparable across view_b_frac settings. - gen_probs = [r for r in chunk if r['_cands']] - base_acc = _mean([1.0 if r['_baseline_pass'] else 0.0 for r in gen_probs]) - ws_acc = _mean([1.0 if any(c.get('reward') for c in r['_cands']) else 0.0 for r in gen_probs]) - cand_pass_parseable = _mean([c['with_pass'] for c in scored]) - cand_pass_all = _mean([1.0 if c.get('reward') else 0.0 for c in all_cands]) - # base failure taxonomy (you asked whether skills fail because the base loops out of length) - classes = [_baseline_class(r) for r in chunk] - n_fail = sum(1 for c in classes if c != 'success') - skill_tokens = [c.get('skillgen_tokens') or 0 for c in cands] # skill-gen response length - trunc = sum(1 for r in chunk for c in r['_cands'] - for x in c['rolls'] if x['stop_reason'] == 'length') - rubric_answer_leaks = sum( - 1 for r in chunk - if r.get('_view') == 'A' and _rubric_has_fail(r.get('_rubric_diag')) and _rubric_answer_leaked(r)) - return { - 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), - 'n_failed_first_try': sum(1 for r in chunk if r['_failed']), - 'n_generated': len(all_cands), 'n_candidates_parseable': len(cands), - 'n_unparseable': len(all_cands) - len(cands), - 'parse_rate': (len(cands) / len(all_cands)) if all_cands else 0.0, - 'n_leaked': sum(1 for c in cands if c['leaked']), 'n_clean': len(clean), - 'leak_rate': (sum(1 for c in cands if c['leaked']) / len(cands)) if cands else 0.0, - 'n_reward_pos': sum(1 for c in scored if c['reward']), - 'n_rubric_answer_leaked': rubric_answer_leaks, - 'n_viewa_dropped': sum(1 for r in chunk if _viewa_dropped(r, args)), - 'n_train_samples': sum(1 for c in all_cands if _is_trainable(c, args)), - 'signal': _signal_stats(chunk), - 'fail_loop_frac': (sum(1 for c in classes if c == 'fail_loop') / n_fail) if n_fail else 0.0, - 'fail_wrong_frac': (sum(1 for c in classes if c == 'fail_wrong') / n_fail) if n_fail else 0.0, - 'skill_tokens_mean': _mean(skill_tokens), - 'withskill_trunc_frac': (trunc / len(ws_rolls)) if ws_rolls else 0.0, - 'avg_baseline_pass': base_acc, 'avg_withskill_pass': ws_acc, - 'avg_lift': ws_acc - base_acc, - 'candidate_withskill_pass_parseable': cand_pass_parseable, - 'candidate_withskill_pass_all': cand_pass_all, - 'termination_rate_withskill': _mean([1.0 if x['terminated'] else 0.0 for x in ws_rolls]), - 'view_A': _view_stats(chunk, 'A'), 'view_B': _view_stats(chunk, 'B'), - **_xproblem_stats(chunk, args), - } - - -def _xproblem_stats(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> Dict[str, Any]: - """Cross-problem pairing health: of the view-A problems, how many actually got a - neighbour's rubric (pair_rate) and how similar those neighbours were. Empty when off.""" - if not args.xproblem_rubric: - return {} - view_a = [r for r in chunk if r.get('_view') == 'A'] - paired = [r for r in view_a if r.get('_rubric_src')] - return {'xproblem': { - 'n_view_a': len(view_a), 'n_paired': len(paired), - 'pair_rate': (len(paired) / len(view_a)) if view_a else 0.0, - 'neighbor_sim_mean': _mean([r.get('_neighbor_sim', 0.0) for r in paired])}} - - -def _sft_record(r: Dict[str, Any], c: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: - """A view-A context-distillation sample: the OPEN-BOOK (rubric-conditioned) skill ``c`` - is the target, but the prompt is rebuilt as query-only (view B, no rubric) so the model - learns to produce it CLOSED-BOOK. Rides the same GRPOLoss with a positive constant - advantage (``--sft-weight``); single-step (old_logps=None) this reduces to - ``-w*logp + beta*KL`` -- a KL-anchored weighted cross-entropy toward the skill.""" - return { - 'data_id': r.get('data_id'), - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'view': 'B', 'diagnosis': '', 'rubric_src': '', 'orig_view': 'A', 'sft': True, - 'response': c['response'], 'skills': c['skills'], - 'skillgen_stop': c.get('skillgen_stop'), - 'advantage': float(args.sft_weight), 'grpo_adv': 0.0, 'kept': True, - 'reward': c['reward'], 'with_pass': c['with_pass']} - - -def _rubric_answer_leaked(r: Dict[str, Any]) -> bool: - """Hard safety gate for external API/rubric diagnosis. Unlike answer-bearing skills - generated by the policy itself, a rubric that contains the target final answer is an - external teacher leak and must not be distilled into view B.""" - return _answer_leaked(r.get('_rubric_diag', ''), r.get('reference_answer', '')) - - -def _viewa_dropped(r: Dict[str, Any], args: argparse.Namespace) -> bool: - """Under --viewa-sft a view-A problem trains ONLY through the SFT path (rubric with - [FAIL], answer-free). No-FAIL and rubric-leaked problems produce no training record - at all (no GRPO backflow: those prompts are query-only and would muddy the pure - view-B GRPO pool), so their skill-gen + executor scoring is skipped entirely.""" - return (bool(args.viewa_sft) and r.get('_view') == 'A' - and (not _rubric_has_fail(r.get('_rubric_diag')) - or _rubric_answer_leaked(r))) - - -def _group_records(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> List[Dict[str, Any]]: - """Training records. Under --viewa-sft, view A is SFT-only: a problem whose rubric - localised a failure (has ``[FAIL]``, answer-free) contributes ONE context-distillation - SFT sample (best parseable open-book skill -- preferring an executor-verified pass, - else any parseable candidate -- rebuilt query-only); no-FAIL / rubric-leaked view-A - problems are dropped (``_viewa_dropped``, generation already skipped). GRPO candidates - come from view B only, keeping that pool purely closed-book. Prompts are rebuilt from - the stored view/diagnosis by ``_skillgen_messages``.""" - out = [] - for r in chunk: - if not r['_hard']: - continue - if args.viewa_sft and r.get('_view') == 'A': - if _viewa_dropped(r, args): - continue - best = _best_sft_candidate(r, args) - if best is not None: - out.append(_sft_record(r, best, args)) - continue - for c in r['_cands']: - if _is_trainable(c, args): - out.append({ - 'data_id': r.get('data_id'), - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'view': r.get('_view', 'A'), 'diagnosis': r.get('_rubric_diag', ''), - 'rubric_src': r.get('_rubric_src', ''), 'sft': False, - 'response': c['response'], 'skills': c['skills'], - 'skillgen_stop': c.get('skillgen_stop'), - 'advantage': c['advantage'], 'grpo_adv': c['grpo_adv'], 'kept': c['kept'], - 'reward': c['reward'], 'with_pass': c['with_pass']}) - return out - - -# =========================================================================== -# Block G -- online GRPO training -# =========================================================================== -def _is_num(v: Any) -> bool: - try: - float(v) - return True - except (TypeError, ValueError): - return False - - -def _train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: - """Training sample = the exact skill-gen prompt (rebuilt by ``_skillgen_messages`` so - train/inference match) + the generated structured guidance response. ``key_rounds`` - selects the final assistant turn; Template masks the prompt and trains the whole - response (the key-round prefix already excludes the prompt-provided <think>).""" - msgs = _skillgen_messages( - rec['problem'], rec.get('view', 'A'), rec.get('diagnosis', ''), rec.get('rubric_src', '')) - return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], - 'user_data': pack_user_data({'key_rounds': [len(msgs)]})} - - -def _train_chunk(skill_model, ref_model, ckpt: CheckpointEngineManager, samples: List[Dict[str, Any]], - args: argparse.Namespace) -> Dict[str, Any]: - """On-policy GRPO update over one chunk, then sync weights. Micro-batches of - ``sft_batch_size`` accumulate gradients; ONE optimizer step is taken per PPO - mini-batch of ``ppo_mini_batch_size`` samples (0 -> a single step over the whole - chunk, the original behaviour). A frozen reference model provides ref_logps for the - SEAM-style KL penalty. - - Multi-step correctness: with more than one step over the SAME rollout, later - mini-batches see an already-updated policy, so we FREEZE the sampling-policy - ``old_logps`` (recomputed once, before any step) and let GRPOLoss form the PPO ratio - against them. A single step keeps ``old_logps=None`` (ratio==1, pure on-policy). - The batch is padded to a multiple of ``sft_batch_size`` with advantage-0 copies that - contribute no policy gradient. View-A context-distillation samples ride the same loss - with a positive constant advantage (``--sft-weight``); single-step (old_logps=None) - that is ``-w*logp + beta*KL`` (KL-anchored weighted cross-entropy).""" - trajs = [_train_trajectory(rec) for rec in samples] - advs = [float(rec['advantage']) for rec in samples] - rem = (-len(trajs)) % args.sft_batch_size - if rem: - trajs += [trajs[-1]] * rem - advs += [0.0] * rem - - n, sft = len(trajs), args.sft_batch_size - mini = args.ppo_mini_batch_size if args.ppo_mini_batch_size > 0 else n - mini = max(sft, (mini // sft) * sft) # align to a whole number of micro-batches - multi_step = mini < n - - # Freeze ref_logps (frozen model) and -- only when we take multiple steps -- the - # sampling-policy old_logps, BOTH before any optimizer step touches the weights. With - # a single step old_logps stays None so GRPOLoss uses ratio==1 (pure on-policy). - micro_ref, micro_old = [], [] - for i in range(0, n, sft): - mb = trajs[i:i + sft] - micro_ref.append(ref_model.forward_only(inputs=mb).get('logps')) - micro_old.append(skill_model.forward_only(inputs=mb).get('logps') if multi_step else None) - - micro, n_steps = 0, 0 - for ms in range(0, n, mini): - for i in range(ms, min(ms + mini, n), sft): - k = i // sft - skill_model.forward_backward(inputs=trajs[i:i + sft], - advantages=advs[i:i + sft], - old_logps=micro_old[k], - ref_logps=micro_ref[k]) - micro += 1 - skill_model.clip_grad_and_step() - n_steps += 1 - ckpt.sync_weights(merge_and_sync=True) - metric = skill_model.calculate_metric(is_training=True) - n_sft = sum(1 for s in samples if s.get('sft')) - return {'n_samples': len(samples), 'n_sft': n_sft, 'n_grpo': len(samples) - n_sft, - 'n_steps': n_steps, 'n_micro_batches': micro, - 'metric': {k: (float(v) if _is_num(v) else v) for k, v in (metric or {}).items()}} - - -# =========================================================================== -# Block H -- fixed-holdout eval + metric formatting -# =========================================================================== -def run_greedy_eval(base_sampler, skill_sampler, eval_records: List[Dict[str, Any]], - ci: int, rounds: int, base_dp: int, skill_dp: int, args: argparse.Namespace, - base_cache: DiskCache - ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], Dict[str, float]]: - """SEAM ``val-core/math/acc/mean@1`` on the fixed holdout: ONE greedy skill (T=0) per - problem into ONE greedy base solve (T=0). Eval ALWAYS uses view B (query-only, the - deployment form: no rubric, since rubric needs an online teacher unavailable at deploy); - no leak filter (acc scores correctness alone). Baseline reuses the disk cache.""" - baseline_rollout(base_sampler, eval_records, base_dp, args, base_cache) - for r in eval_records: - r['_view'], r['_rubric_diag'] = 'B', '' - sg_out = _run_samples(skill_sampler, [_view_prompt(r, args) for r in eval_records], - 1, args.skill_max_tokens, skill_dp, temperature=0.0) - skills = [] - for seqs in sg_out: - if not seqs: - skills.append(('', '')) - continue - sresp = _clean_text(getattr(seqs[0], 'decoded', '') or '') - skills.append((_extract_skill(sresp) or '', sresp)) - ws_out = _run_samples(base_sampler, - [build_skill_solve_prompt(r['problem'], sk) for r, (sk, _) in zip(eval_records, skills)], - 1, args.max_tokens, base_dp, temperature=0.0) - recs = [] - for r, (sk, sresp), seqs in zip(eval_records, skills, ws_out): - roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - recs.append({ - 'record_type': 'eval_problem', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, - 'data_id': r.get('data_id'), - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'view': r['_view'], 'rubric_diag': r.get('_rubric_diag', ''), - 'baseline_pass': r['_baseline_pass'], 'skill': sk, 'skill_parseable': bool(sk), - 'skill_response': sresp, 'withskill_pred': roll['pred'], - 'withskill_correct': roll['correct'], 'withskill_terminated': roll['terminated'], - 'withskill_stop_reason': roll['stop_reason'], 'withskill_text': roll['text'], - }) - acc = lambda rs: sum(1 for x in rs if x['withskill_correct']) / len(rs) if rs else 0.0 - ws = acc(recs) # all view B (deployment form) - base = sum(x['baseline_pass'] for x in recs) / len(recs) if recs else 0.0 - fmt = (sum(1 for x in recs if x['skill_parseable']) / len(recs)) if recs else 0.0 - term = (sum(1 for x in recs if x['withskill_terminated']) / len(recs)) if recs else 0.0 - summary = {'record_type': 'eval_summary', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, - 'n': len(recs), 'view': 'B', 'acc_mean1': ws, - 'baseline_acc_mean1': base, 'lift_mean1': ws - base, - 'format_mean1': fmt, 'term_mean1': term} - metrics = {'core/math/acc/mean@1': ws, 'core/math/baseline_acc/mean@1': base, - 'core/math/lift/mean@1': ws - base, 'core/math/format/mean@1': fmt, - 'core/math/term/mean@1': term} - return recs, summary, metrics - - -def _trend_line(hist: List[Dict[str, float]], window: int, rounds_done: int) -> Optional[str]: - """Contrast the FIRST vs the most recent ``window`` chunks -- if RFT works, adoption - and lift on recent (fresh) chunks exceed the early baseline.""" - if len(hist) < 2 * window: - return None - base, rec = hist[:window], hist[-window:] - m = lambda xs, k: sum(h[k] for h in xs) / len(xs) - return (f'[trend] first {window} vs last {window} | ' - f'adopt A {m(base,"aA"):.2f}->{m(rec,"aA"):.2f} B {m(base,"aB"):.2f}->{m(rec,"aB"):.2f} | ' - f'lift {m(base,"lift"):+.3f}->{m(rec,"lift"):+.3f} | ' - f'0grad {m(base,"zero_grad"):.2f}->{m(rec,"zero_grad"):.2f} | ' - f'pos/chunk {m(base,"pos"):.1f}->{m(rec,"pos"):.1f} | rounds={rounds_done}') - - -def _swan_metrics(summary: Dict[str, Any], log: Optional[Dict[str, Any]]) -> Dict[str, float]: - """Flat swanlab dict. ``signal/*`` is the primary health group (is GRPO getting a - gradient at all); ``acc/*`` is the effect; the rest diagnose why. acc/adopt/skill are - only logged when the chunk produced a scored group, so idle chunks don't dip charts.""" - sig = summary['signal'] - d: Dict[str, float] = { - # --- signal: the FIRST thing to watch (no variance -> no learning) --- - 'signal/zero_grad_frac': sig['zero_grad_frac'], 'signal/n_groups': sig['n_groups'], - 'signal/reward_mean': sig['reward_mean'], 'signal/reward_std': sig['reward_std'], - 'signal/group_reward_std_mean': sig['group_reward_std_mean'], - 'signal/n_train_samples': summary['n_train_samples'], - 'signal/n_reward_pos': summary['n_reward_pos'], - # --- skill format / leak health --- - 'skill/parse_rate': summary['parse_rate'], 'skill/tokens_mean': summary['skill_tokens_mean'], - 'leak/rate': summary['leak_rate'], 'leak/n': summary['n_leaked'], - # --- base failure taxonomy (loop-out-of-length vs plain wrong) --- - 'fail/loop_frac': summary['fail_loop_frac'], 'fail/wrong_frac': summary['fail_wrong_frac'], - 'fail/frac_first_try': (summary['n_failed_first_try'] / summary['n']) if summary['n'] else 0.0, - # --- view-A routing: dropped = no-FAIL rubric or rubric leak (SFT-only view A) --- - 'viewa/dropped_frac': (summary['n_viewa_dropped'] / summary['view_A']['n'] - if summary['view_A']['n'] else 0.0), - } - bal = summary.get('balance') or {} - if bal.get('enabled'): - d.update({'balance/n_drawn': bal['n_drawn'], 'balance/n_baseline_fresh': bal['n_baseline_fresh'], - 'balance/selected_success_frac': bal['selected_success_frac']}) - xp = summary.get('xproblem') or {} - if xp: - d.update({'xproblem/pair_rate': xp['pair_rate'], 'xproblem/neighbor_sim_mean': xp['neighbor_sim_mean']}) - if sig['n_groups'] > 0: - d.update({'acc/baseline_pass': summary['avg_baseline_pass'], - 'acc/withskill_pass': summary['avg_withskill_pass'], 'acc/lift': summary['avg_lift'], - 'candidate/withskill_pass_parseable': summary['candidate_withskill_pass_parseable'], - 'candidate/withskill_pass_all': summary['candidate_withskill_pass_all'], - 'adopt/A': summary['view_A']['adoption_rate'], - 'adopt/B': summary['view_B']['adoption_rate'], - 'term/withskill': summary['termination_rate_withskill'], - 'term/withskill_trunc_frac': summary['withskill_trunc_frac']}) - if log: - d['train/n_steps'] = log['n_steps'] - d['train/n_micro_batches'] = log['n_micro_batches'] - for k, v in (log.get('metric') or {}).items(): - if not _is_num(v): - continue - if k.startswith('learning rate'): - if 'group 1' in k: - d['train/lr'] = float(v) - else: - d[f'train/{k.replace(" ", "_")}'] = float(v) - return d - - -def _view_a_rubric_leak_metrics(chunk: List[Dict[str, Any]], - pool_answers: Optional[Dict[str, str]] = None) -> Dict[str, float]: - """Swanlab-only audit for answer leakage in view-A rubric text. This never changes - rewards, advantages, filtering, or training records.""" - view_a = [r for r in chunk if r.get('_view') == 'A'] - with_diag = [r for r in view_a if r.get('_rubric_diag')] - target_leaks = sum(1 for r in with_diag - if _answer_leaked(r.get('_rubric_diag', ''), r.get('reference_answer', ''))) - source_leaks = 0 - pool_answers = pool_answers or {} - for r in with_diag: - src = r.get('_rubric_src') - src_ref = pool_answers.get(src, '') if src else r.get('reference_answer', '') - if _answer_leaked(r.get('_rubric_diag', ''), src_ref): - source_leaks += 1 - n = len(with_diag) - return { - 'rubric_leak/n_view_a': float(len(view_a)), - 'rubric_leak/n_checked': float(n), - 'rubric_leak/target_answer_n': float(target_leaks), - 'rubric_leak/target_answer_rate': (target_leaks / n) if n else 0.0, - 'rubric_leak/source_answer_n': float(source_leaks), - 'rubric_leak/source_answer_rate': (source_leaks / n) if n else 0.0, - } - - -# =========================================================================== -# Block F -- components, args, main -# =========================================================================== -def init_components(args: argparse.Namespace): - """Default 8-GPU layout: rank 0 trains the actor, rank 1 hosts a frozen ref model, - 2-3 skill_sampler (synced), 4-7 base_sampler (frozen). Returns - (skill_model, ref_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp).""" - r0, r1 = TRAIN_GPUS, TRAIN_GPUS + REF_GPUS - r2, r3 = r1 + SKILL_SAMPLER_GPUS, NUM_GPUS - twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, lazy_collect=False, groups=[ - DeviceGroup(name='train', ranks=list(range(0, r0)), device_type='GPU'), - DeviceGroup(name='ref', ranks=list(range(r0, r1)), device_type='GPU'), - DeviceGroup(name='skill_sampler', ranks=list(range(r1, r2)), device_type='GPU'), - DeviceGroup(name='base_sampler', ranks=list(range(r2, r3)), device_type='GPU')]) - - train_mesh = DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_DP, fsdp_size=TRAIN_FSDP) - skill_model = TransformersModel(model_id=MODEL_ID, device_mesh=train_mesh, remote_group='train', - ddp_config={'find_unused_parameters': False}) - skill_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) - skill_model.set_template(Template, model_id=MODEL_ID, enable_thinking=False, - max_length=args.max_model_len, truncation_strategy='delete') - skill_model.set_processor(InputProcessor, padding_free=False) - skill_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon, beta=args.kl_beta) - skill_model.set_optimizer('AdamW', lr=args.lr) - skill_model.set_lr_scheduler('CosineWarmupScheduler', num_warmup_steps=10, - num_training_steps=args.max_train_rounds) - - ref_mesh = DeviceMesh.from_sizes(world_size=REF_GPUS, dp_size=REF_DP, fsdp_size=REF_FSDP) - ref_model = TransformersModel(model_id=MODEL_ID, device_mesh=ref_mesh, remote_group='ref', - ddp_config={'find_unused_parameters': False}) - ref_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) - ref_model.set_template(Template, model_id=MODEL_ID, enable_thinking=True, - max_length=args.max_model_len, truncation_strategy='delete') - ref_model.set_processor(InputProcessor, padding_free=False) - ref_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) - - def _sampler(group, world, enable_thinking: bool = True): - s = vLLMSampler(model_id=MODEL_ID, - engine_args={'gpu_memory_utilization': GPU_MEM, - 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, - device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), - remote_group=group) - s.set_template(Template, model_id=MODEL_ID, enable_thinking=enable_thinking, max_length=args.max_model_len) - return s - - skill_sampler = _sampler('skill_sampler', SKILL_SAMPLER_GPUS, enable_thinking=False) - # base_sampler is shared by the main thread and the baseline-prefetch thread -> serialise. - base_sampler = _LockedSampler(_sampler('base_sampler', BASE_SAMPLER_GPUS, enable_thinking=True)) - ckpt = CheckpointEngineManager(model=skill_model, sampler=skill_sampler) - return skill_model, ref_model, skill_sampler, base_sampler, ckpt, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS - - -def _build_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--dataset', choices=('aops', 'math'), default='aops') - p.add_argument('--n', type=int, default=2000, help='Problems loaded into the draw pool.') - p.add_argument('--pool-offset', type=int, default=0, - help='Skip this many shuffled non-eval records before building the train pool; ' - 'useful to avoid cold-start SFT data ranges.') - p.add_argument('--exclude-data-ids', default='', - help='Comma-separated jsonl files whose data_id/problem keys are excluded ' - 'from train/eval selection, e.g. coldstart_sft.jsonl.') - p.add_argument('--seed', type=int, default=42) - p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) - p.add_argument('--eval-size', type=int, default=128, help='Fixed holdout size (0 disables).') - p.add_argument('--eval-every', type=int, default=5, help='Run holdout eval every N chunks.') - p.add_argument('--chunk-size', type=int, default=16) - p.add_argument('--balance', action=argparse.BooleanOptionalAction, default=True) - p.add_argument('--balance-success-frac', type=float, default=0.4, - help='Target fraction of the chunk the base solves (rest are base-fail).') - p.add_argument('--balance-loop-frac', type=float, default=0.5) - p.add_argument('--balance-max-draws-mult', type=int, default=8) - p.add_argument('--n-skills', type=int, default=8) - p.add_argument('--view-b-frac', type=float, default=0.5) - p.add_argument('--viewa-frac-start', type=float, default=None, - help='Enable the view-A curriculum: chunk 0 uses this view-A share ' - '(view_b_frac = 1 - share), decaying linearly to --viewa-frac-end ' - 'over --viewa-decay-chunks chunks, then holding. Overrides ' - '--view-b-frac for every chunk.') - p.add_argument('--viewa-frac-end', type=float, default=0.1) - p.add_argument('--viewa-warmup-chunks', type=int, default=0, - help='Hold the view-A share at --viewa-frac-start for this many chunks ' - 'before the linear decay begins.') - p.add_argument('--viewa-decay-chunks', type=int, default=40) - p.add_argument('--xproblem-rubric', action=argparse.BooleanOptionalAction, default=False, - help='When enabled, view A uses a bag-of-words NEIGHBOUR problem rubric. ' - 'Default is off: each view-A problem uses its own baseline attempt, ' - 'while the API diagnosis prompt is constrained to be answer-free and ' - 'method-level only.') - p.add_argument('--skill-retries', type=int, default=2) - p.add_argument('--skill-gen-temperature', type=float, default=1.0) - p.add_argument('--skill-gen-top-p', type=float, default=1.0) - p.add_argument('--skill-gen-top-k', type=int, default=-1) - p.add_argument('--max-model-len', type=int, default=16384) - p.add_argument('--max-tokens', type=int, default=8192) - p.add_argument('--skill-max-tokens', type=int, default=8192) - p.add_argument('--rubric-workers', type=int, default=16) - p.add_argument('--rubric-retries', type=int, default=2, - help='Retry failed/timeout rubric diagnose calls this many times before ' - 'falling back to an empty diagnosis without caching the failure.') - p.add_argument('--sft-batch-size', type=int, default=8, - help='Driver micro-batch (gradient-accumulation unit); multiple of train dp.') - p.add_argument('--ppo-mini-batch-size', type=int, default=0, - help='Samples per optimizer step (SEAM-style PPO mini-batch). 0 = one step ' - 'over the whole chunk (on-policy, ratio==1). When >0 and smaller than ' - 'the trainable count, multiple steps are taken over the same rollout and ' - 'old_logps are frozen so the PPO ratio/clip stays valid. Rounded down to ' - 'a multiple of --sft-batch-size.') - p.add_argument('--grpo-epsilon', type=float, default=0.2) - p.add_argument('--adv-clip', type=float, default=3.0, - help='Symmetric clip for group-relative advantages; <=0 disables clipping.') - p.add_argument('--kl-beta', type=float, default=0.001, - help='SEAM-style reference KL coefficient for GRPOLoss.') - p.add_argument('--viewa-sft', action=argparse.BooleanOptionalAction, default=True, - help='Route view-A problems to online context distillation (SFT on the best ' - 'passing open-book skill, rebuilt as a query-only prompt) instead of GRPO. ' - 'View B stays GRPO; both share one optimizer step.') - p.add_argument('--sft-weight', type=float, default=0.5, - help='Advantage magnitude (lambda) for view-A SFT distillation samples in the ' - 'shared GRPOLoss. Single-step this is -lambda*logp + beta*KL (weighted CE). ' - 'Keep <~1 so one-signed SFT gradients do not swamp mean-0 GRPO advantages.') - p.add_argument('--sft-target-len', type=int, default=550, - help='Target skill length (chars) for view-A SFT distillation: among passing ' - 'candidates pick the one CLOSEST to this length (not the shortest). Set near ' - 'the empirical pass-rate peak (~500-600 here) so distillation neither shrinks ' - 'rollouts toward zero nor lets them grow unbounded.') - p.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True) - p.add_argument('--lr', type=float, default=6e-6) - p.add_argument('--max-train-rounds', type=int, default=1500) - p.add_argument('--save-rounds', type=int, default=200) - p.add_argument('--trend-every', type=int, default=10) - p.add_argument('--output-dir', default='./output/reflexion_skill') - p.add_argument('--cache-dir', default='', help='Baseline/rubric cache dir (default <output-dir>/cache).') - p.add_argument('--no-cache', action='store_true', help='Disable disk cache read/write.') - p.add_argument('--prefetch-baseline', action=argparse.BooleanOptionalAction, default=True, - help='Prefetch next chunk base baseline on a background thread (overlaps ' - 'with skill-gen; base_sampler is frozen so it never blocks the trainer).') - p.add_argument('--swanlab-project', default='twinkle') - p.add_argument('--swanlab-exp', default='') - args = p.parse_args() - if args.sft_batch_size % TRAIN_DP != 0: - raise ValueError(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple of train dp ({TRAIN_DP})') - if args.chunk_size < 1: - raise ValueError('--chunk-size must be >= 1') - args.balance_success_frac = max(0.0, min(1.0, args.balance_success_frac)) - return args - - -def _write(handle, row: Dict[str, Any]) -> None: - handle.write(json.dumps(row, ensure_ascii=False) + '\n') - - -def main() -> None: - args = _build_args() - records, eval_records, neighbor_map, pool_answers, data_stats = _load_records(args) - if len(records) < args.chunk_size: - raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded ({len(records)}); raise --n') - - os.makedirs(args.output_dir, exist_ok=True) - gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') - eval_path = os.path.join(args.output_dir, 'eval_records.jsonl') - data_path = os.path.join(args.output_dir, 'skill_dataset.jsonl') - train_log_path = os.path.join(args.output_dir, 'train_log.jsonl') - - if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY')): - sys.stderr.write('[rft] WARNING: no LLM backup env; view-A rubric check disabled ' - '(leak filter is deterministic, unaffected)\n') - - use_swan = swanlab is not None and os.environ.get('SWANLAB_MODE') != 'disabled' - if use_swan: - swanlab.init(project=args.swanlab_project, experiment_name=(args.swanlab_exp or None), - config={'model': MODEL_ID, 'dataset': args.dataset, 'n': len(records), - 'eval_n': len(eval_records), 'n_skills': args.n_skills, - 'view_b_frac': args.view_b_frac, 'balance': args.balance, - 'balance_success_frac': args.balance_success_frac, - 'skill_gen_temp': args.skill_gen_temperature, - 'grpo_epsilon': args.grpo_epsilon, 'kl_beta': args.kl_beta, - 'lr': args.lr}) - - skill_model, ref_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp = init_components(args) - checker = build_rubric_checker() - if checker is None: - sys.stderr.write('[rft] no LLM backup env -> view-A rubric process-check DISABLED\n') - - cache_dir = args.cache_dir or os.path.join(args.output_dir, 'cache') - os.makedirs(cache_dir, exist_ok=True) - use_cache = not args.no_cache - base_cache = DiskCache(os.path.join(cache_dir, 'baseline.jsonl'), use_cache) - eval_base_cache = DiskCache(os.path.join(cache_dir, 'eval_baseline.jsonl'), use_cache) - rubric_cache = DiskCache(os.path.join(cache_dir, 'rubric.jsonl'), use_cache) - if args.xproblem_rubric: - sys.stderr.write(f'[rft] xproblem-rubric ON: {len(neighbor_map)} bag-of-words pairs\n') - - cfg = {'record_type': 'config', 'model': MODEL_ID, 'dataset': args.dataset, - 'n': len(records), 'eval_n': len(eval_records), 'seed': args.seed, - 'pool_offset': args.pool_offset, 'exclude_data_ids': args.exclude_data_ids, - 'excluded_records': data_stats.get('excluded_records', 0), - 'numeric_only': args.numeric_only, 'raw_loaded': data_stats['raw_loaded'], - 'numeric_dropped': data_stats['numeric_dropped'], 'eval_every': args.eval_every, - 'n_skills': args.n_skills, 'view_b_frac': args.view_b_frac, - 'viewa_frac_start': args.viewa_frac_start, 'viewa_frac_end': args.viewa_frac_end, - 'viewa_warmup_chunks': args.viewa_warmup_chunks, - 'viewa_decay_chunks': args.viewa_decay_chunks, - 'skill_retries': args.skill_retries, 'balance': args.balance, - 'balance_success_frac': args.balance_success_frac, - 'skill_gen_temp': args.skill_gen_temperature, 'reward': 'greedy_binary(correct)', - 'advantage': 'group_relative', 'format_in_reward': args.format_in_reward, 'cache': use_cache, - 'rubric_check': f'fixed_math_{len(_MATH_RUBRIC)}crit+term(viewA)' if checker else 'disabled', - 'xproblem_rubric': args.xproblem_rubric, - 'viewa_sft': args.viewa_sft, 'sft_weight': args.sft_weight, - 'sft_target_len': args.sft_target_len, - 'adv_clip': args.adv_clip, - 'grpo_epsilon': args.grpo_epsilon, 'kl_beta': args.kl_beta, 'lr': args.lr, - 'train_gpus': TRAIN_GPUS, 'ref_gpus': REF_GPUS, 'ref_fsdp': REF_FSDP, - 'train_fsdp': TRAIN_FSDP, 'train_dp': TRAIN_DP, - 'skill_sampler_gpus': SKILL_SAMPLER_GPUS, 'base_sampler_gpus': BASE_SAMPLER_GPUS, - 'max_train_rounds': args.max_train_rounds, 'started': int(time.time())} - sys.stderr.write(f'[rft] raw={data_stats["raw_loaded"]} numeric_drop={data_stats["numeric_dropped"]} ' - f'train={len(records)} eval={len(eval_records)} {args.dataset}; ' - f'train_gpus={TRAIN_GPUS} ref_gpus={REF_GPUS} train_fsdp={TRAIN_FSDP} ' - f'train_dp={TRAIN_DP} skill_dp={skill_dp} base_dp={base_dp}\n') - - hist: List[Dict[str, float]] = [] - rounds = 0 - pool = ProblemPool(records, args.seed) - with open(gen_path, 'w', encoding='utf-8') as gen_f, \ - open(eval_path, 'w', encoding='utf-8') as eval_f, \ - open(data_path, 'w', encoding='utf-8') as data_f, \ - open(train_log_path, 'w', encoding='utf-8') as tlog: - for f in (gen_f, eval_f, data_f, tlog): - _write(f, cfg) - gstep = 0 - # base baseline (frozen sampler, disk-cached, no weight-sync) is prefetched on a - # background thread while the current chunk generates: the skill-gen phase uses - # skill_sampler, so prefetching the NEXT chunk's baseline on base_sampler overlaps - # it and fills base_cache so the next draw_chunk hits it. base_sampler is wrapped in - # _LockedSampler, so the prefetch and this chunk's with-skill scoring never issue a - # base .sample() concurrently. It never touches the trainer or on-policy generation. - prefetch_pool = ThreadPoolExecutor(max_workers=1) if args.prefetch_baseline else None - pending: Optional[Any] = None - - def _prefetch(peeked: List[Dict[str, Any]]) -> None: - if peeked: - baseline_rollout(base_sampler, peeked, base_dp, args, base_cache) - - # Baseline (round 0) eval BEFORE any training: measures the untrained skill model on - # the fixed holdout so every later eval has a step-0 reference point on the same axis. - if eval_records: - eval_recs, eval_summary, eval_metrics = run_greedy_eval( - base_sampler, skill_sampler, eval_records, -1, rounds, base_dp, skill_dp, - args, eval_base_cache) - for rec in eval_recs: - _write(eval_f, rec) - _write(eval_f, eval_summary) - eval_f.flush() - if use_swan: - swanlab.log({f'eval/{k}': v for k, v in eval_metrics.items()}, step=0) - sys.stderr.write( - f'[eval] g-1 (init): n={eval_summary["n"]} viewB mean@1 ' - f'acc={eval_summary["baseline_acc_mean1"]:.3f}->{eval_summary["acc_mean1"]:.3f} ' - f'lift={eval_summary["lift_mean1"]:+.3f} ' - f'fmt={eval_summary["format_mean1"]:.2f} rounds={rounds}\n') - - # Each chunk is drawn fresh + RE-GENERATED with the current policy (on-policy); - # with --balance, draw_chunk keeps drawing until the base fail:success mix hits target. - while rounds < args.max_train_rounds: - if pending is not None: - pending.result() # finish last round's prefetch before drawing (cache-warm) - pending = None - if args.viewa_frac_start is not None: - args.view_b_frac = _curriculum_view_b_frac(gstep, args) - chunk, balance = draw_chunk(pool, base_sampler, base_dp, args, base_cache) - if prefetch_pool is not None: - peeked = pool.peek(int(args.chunk_size * (args.balance_max_draws_mult if args.balance else 1))) - pending = prefetch_pool.submit(_prefetch, peeked) - full, summary, groups = process_chunk( - base_sampler, skill_sampler, chunk, gstep, base_dp, skill_dp, - args, checker, rubric_cache, base_cache, neighbor_map, pool_answers) - summary['balance'] = balance - summary['view_b_frac'] = round(args.view_b_frac, 4) - - log = None - if groups: - log = _train_chunk(skill_model, ref_model, ckpt, groups, args) - rounds += 1 - log.update({'record_type': 'train_round', 'round': rounds, 'chunk': gstep, - 'epoch': pool.epoch, 'ts': int(time.time())}) - _write(tlog, log) - tlog.flush() - if rounds % args.save_rounds == 0: - skill_model.save(f'skill-rft-{rounds}', output_dir=args.output_dir) - - summary['rounds_done'], summary['epoch'] = rounds, pool.epoch - for rec in full: - _write(gen_f, rec) - _write(gen_f, summary) - gen_f.flush() - for v in groups: - _write(data_f, v) - data_f.flush() - - sa, sb, sig = summary['view_A'], summary['view_B'], summary['signal'] - hist.append({'aA': sa['adoption_rate'], 'aB': sb['adoption_rate'], - 'lift': summary['avg_lift'], 'pos': summary['n_reward_pos'], - 'zero_grad': sig['zero_grad_frac']}) - bal_str = (f'bal {balance["selected_fail"]}f/{balance["selected_success"]}s ' - f'(drew {balance["n_drawn"]}/fresh {balance["n_baseline_fresh"]}' - + ('!' if balance.get('budget_hit') else '') + ') ') if balance.get('enabled') else '' - xp = summary.get('xproblem') - xp_str = f'pair={xp["pair_rate"]:.2f}@{xp["neighbor_sim_mean"]:.2f} ' if xp else '' - tr_str = (f'train={log["n_grpo"]}g/{log["n_sft"]}s ' if log - else f'train={summary["n_train_samples"]} ') - sys.stderr.write( - f'[gen] e{pool.epoch} g{gstep}: {bal_str}n={summary["n"]} ' - f'clean={summary["n_clean"]} leak={summary["leak_rate"]:.2f} {tr_str}' - f'0grad={sig["zero_grad_frac"]:.2f} R={sig["reward_mean"]:.2f}±{sig["reward_std"]:.2f} ' - f'acc={summary["avg_baseline_pass"]:.2f}->{summary["avg_withskill_pass"]:.2f} ' - f'lift={summary["avg_lift"]:+.3f} {xp_str}' - f'A[{sa["n"]} {sa["adoption_rate"]:.2f}] B[{sb["n"]} {sb["adoption_rate"]:.2f}] ' - f'rounds={rounds}\n') - if use_swan: - swan_metrics = _swan_metrics(summary, log) - swan_metrics.update(_view_a_rubric_leak_metrics(chunk, pool_answers)) - swan_metrics['train/view_b_frac'] = float(args.view_b_frac) - swanlab.log(swan_metrics, step=gstep) - - if eval_records and (gstep + 1) % args.eval_every == 0: - eval_recs, eval_summary, eval_metrics = run_greedy_eval( - base_sampler, skill_sampler, eval_records, gstep, rounds, base_dp, skill_dp, - args, eval_base_cache) - for rec in eval_recs: - _write(eval_f, rec) - _write(eval_f, eval_summary) - eval_f.flush() - if use_swan: - swanlab.log({f'eval/{k}': v for k, v in eval_metrics.items()}, step=gstep) - sys.stderr.write( - f'[eval] g{gstep}: n={eval_summary["n"]} viewB mean@1 ' - f'acc={eval_summary["baseline_acc_mean1"]:.3f}->{eval_summary["acc_mean1"]:.3f} ' - f'lift={eval_summary["lift_mean1"]:+.3f} ' - f'fmt={eval_summary["format_mean1"]:.2f} rounds={rounds}\n') - - if (gstep + 1) % args.trend_every == 0: - tl = _trend_line(hist, args.trend_every, rounds) - if tl: - sys.stderr.write(tl + '\n') - gstep += 1 - - if prefetch_pool is not None: - if pending is not None: - pending.result() - prefetch_pool.shutdown(wait=True) - base_cache.close() - eval_base_cache.close() - rubric_cache.close() - skill_model.save('skill-rft-final', output_dir=args.output_dir) - sys.stderr.write(f'[rft] done: {rounds} rounds over {gstep} chunks / {pool.epoch} epochs -> {data_path}\n') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/legacy/train_reflexion_skill.sh b/cookbook/exp/legacy/train_reflexion_skill.sh deleted file mode 100755 index e732a6a3f..000000000 --- a/cookbook/exp/legacy/train_reflexion_skill.sh +++ /dev/null @@ -1,64 +0,0 @@ -#!/bin/bash -# Online GRPO RFT for the reflexion skill generator (unified, self-contained, cached). -# GPUs: 8 — default high-memory layout uses rank 0 for actor training, rank 1 for -# the frozen ref model, ranks 2-3 for skill sampler (synced), and ranks 4-7 for -# base sampler (frozen). Override TRAIN_GPUS / REF_GPUS / SKILL_SAMPLER_GPUS / -# BASE_SAMPLER_GPUS for other layouts. Per chunk: base greedy -# solve -> rubric process-check (view A) -> -# skill-gen (thinking ON, N candidates) -> deterministic leak filter -> with-skill greedy -# pass -> group-relative advantage -> ONE on-policy GRPO step -> sync weights. -# -# Baseline rollouts + rubric diagnoses are disk-cached (output-dir/cache/*.jsonl), so a -# restart skips re-sampling them; skill-gen is on-policy and never cached. The next chunk's -# baseline is prefetched on a background thread (overlaps skill-gen; base sampler is frozen). -# -# The view-A rubric process-check uses the backup teacher API (set LLM_BACKUP_*). Without -# it the run still works: view A degrades to query-only and the leak filter stays -# deterministic (no teacher needed). - -set -euo pipefail - -export GEN_MODEL_ID=${GEN_MODEL_ID:-Qwen/Qwen3-4B} -export GEN_GPU_MEM=${GEN_GPU_MEM:-0.8} -# Datasets are pulled from ModelScope via twinkle.Dataset (ms://AI-MO/aops or -# ms://modelscope/competition_math); override AOPS_DATASET_ID / MATH_DATASET_ID to change. -# Teacher API for the view-A rubric process-check (optional; leak filter is deterministic). -export LLM_BACKUP_API_KEY=${LLM_BACKUP_API_KEY:-} -export LLM_BACKUP_BASE_URL=${LLM_BACKUP_BASE_URL:-https://dashscope.aliyuncs.com/compatible-mode/v1} -export LLM_BACKUP_MODEL=${LLM_BACKUP_MODEL:-qwen3.7-max} -export EXCLUDE_DATA_IDS=${EXCLUDE_DATA_IDS:-./output/reflexion_coldstart_sft/coldstart_sft.jsonl} - -python cookbook/exp/embedding/train_reflexion_skill.py \ - --dataset aops \ - --n 10000 \ - --numeric-only \ - --chunk-size 64 \ - --n-skills 8 \ - --viewa-frac-start 1.0 \ - --viewa-frac-end 0.1 \ - --viewa-warmup-chunks 20 \ - --viewa-decay-chunks 40 \ - --skill-retries 2 \ - --balance \ - --balance-success-frac 0.2 \ - --balance-loop-frac 0.5 \ - --balance-max-draws-mult 8 \ - --max-tokens 8192 \ - --skill-max-tokens 4096 \ - --max-model-len 16384 \ - --eval-size 128 \ - --exclude-data-ids "${EXCLUDE_DATA_IDS}" \ - --eval-every 5 \ - --sft-batch-size 8 \ - --ppo-mini-batch-size 0 \ - --grpo-epsilon 0.2 \ - --kl-beta 0.001 \ - --format-in-reward \ - --lr 1e-6 \ - --max-train-rounds 1500 \ - --save-rounds 200 \ - --trend-every 10 \ - --prefetch-baseline \ - --output-dir ./output/reflexion_skill_curriculum \ - --swanlab-project twinkle \ - --swanlab-exp reflexion_skill_curriculum diff --git a/cookbook/exp/legacy/train_reflexion_skill_old.py b/cookbook/exp/legacy/train_reflexion_skill_old.py deleted file mode 100644 index 3bc9d5162..000000000 --- a/cookbook/exp/legacy/train_reflexion_skill_old.py +++ /dev/null @@ -1,2022 +0,0 @@ -"""Online GRPO RFT for the reflexion skill generator (self-contained, cached). - -Trains an INDEPENDENT skill model to write reusable skills that, injected into a -FROZEN base solver's system prompt, raise its accuracy. The base is never trained; -it only produces the reward. Per chunk: base greedy solve -> rubric process-check -(view A) -> skill-gen (thinking ON, N candidates) -> leak filter -> with-skill -greedy pass -> group-relative advantage -> ONE on-policy GRPO step -> sync weights. -Reward is deterministic (T=0, M=1) binary correctness; A = (R - mean) / (std + eps) -within each problem-group, so std=0 groups give no gradient (GRPO variance selects). - -Each problem is routed to EXACTLY ONE view: A = problem + rubric findings, B = -query only (deployment form). Skill-gen trains only the final structured guidance turn. - -Base greedy rollouts and rubric diagnoses are cached to disk (jsonl, md5-keyed) so -restarts skip them; skill-gen is on-policy and never cached. - -8 GPUs default for high-memory cards: rank 0 trains the actor, rank 1 hosts a -frozen reference model, ranks 2-3 skill_sampler (synced), ranks 4-7 base_sampler -(frozen). Override TRAIN_GPUS / REF_GPUS / SKILL_SAMPLER_GPUS / BASE_SAMPLER_GPUS -for other layouts. -Leak / rubric use the backup teacher API: set LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / -LLM_BACKUP_MODEL. - -Launch: - LLM_BACKUP_API_KEY=... python cookbook/exp/embedding/train_reflexion_skill.py \ - --dataset aops --n 5000 --chunk-size 16 --lr 6e-6 -""" -import argparse -import copy -import hashlib -import json -import math -import os -import re -import sys -import threading -import time -from concurrent.futures import ThreadPoolExecutor -from typing import Any, Dict, List, Optional, Set, Tuple - -import numpy as np - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.checkpoint_engine import CheckpointEngineManager -from twinkle.data_format import SamplingParams, pack_user_data -from twinkle.dataset import Dataset, DatasetMeta -from twinkle.model import TransformersModel -from twinkle.patch.no_split_modules import NoSplitModulesPatch -from twinkle.processor import InputProcessor -from twinkle.sampler import vLLMSampler -from twinkle.template import Template -from twinkle_agentic.verifier import RubricVerifier -from twinkle_agentic.verifier.rubric_verifier import RubricItem - -logger = get_logger() - -try: - import swanlab -except ImportError: - swanlab = None - -MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') -GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.8)) -GEN_TEMPERATURE = float(os.environ.get('GEN_TEMPERATURE', 0.6)) -GEN_TOP_P = float(os.environ.get('GEN_TOP_P', 0.95)) -AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') -MATH_DATASET_ID = os.environ.get('MATH_DATASET_ID', 'modelscope/competition_math') - -# GPU layout: train actor (synced to skill_sampler) + frozen ref model + skill_sampler + base_sampler. -# High-memory cards can keep Qwen3-4B actor/ref as one GPU each and spend more GPUs -# on vLLM data-parallel sampling. The base side is heavier here because every -# clean skill candidate is re-solved by the frozen base model, plus balance/eval baselines. -TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 2)) -REF_GPUS = int(os.environ.get('REF_GPUS', 2)) -SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) -BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 2)) -NUM_GPUS = TRAIN_GPUS + REF_GPUS + SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS -TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', min(1, TRAIN_GPUS))) -REF_FSDP = int(os.environ.get('REF_FSDP', min(1, REF_GPUS))) -if min(TRAIN_GPUS, REF_GPUS, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS, TRAIN_FSDP, REF_FSDP) < 1: - raise ValueError('TRAIN_GPUS, REF_GPUS, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS, TRAIN_FSDP and REF_FSDP must all be >= 1') -if TRAIN_GPUS % TRAIN_FSDP != 0: - raise ValueError(f'TRAIN_GPUS ({TRAIN_GPUS}) must be divisible by TRAIN_FSDP ({TRAIN_FSDP})') -if REF_GPUS % REF_FSDP != 0: - raise ValueError(f'REF_GPUS ({REF_GPUS}) must be divisible by REF_FSDP ({REF_FSDP})') -TRAIN_DP = TRAIN_GPUS // TRAIN_FSDP -REF_DP = REF_GPUS // REF_FSDP - - -# =========================================================================== -# Block A -- boxed extraction + answer grading -# =========================================================================== -_BOXED_RE = re.compile(r'\\boxed\s*\{') - - -def extract_boxed(text: str) -> Optional[str]: - """Last ``\\boxed{...}`` content, brace-balanced.""" - if not text: - return None - last = None - for m in _BOXED_RE.finditer(text): - depth, i = 1, m.end() - while i < len(text) and depth > 0: - depth += (text[i] == '{') - (text[i] == '}') - i += 1 - if depth == 0: - last = text[m.end():i - 1].strip() - return last - - -def normalize_answer(ans: str) -> str: - if not ans: - return '' - s = str(ans).strip() - m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) - if m: - return m.group(1) - s = s.strip('$').strip().replace('−', '-') - 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').replace(r'\tfrac', r'\frac') - s = s.replace(r'\,', '').replace(r'\;', '').replace(r'\!', '') - s = re.sub(r'\\(?:quad|qquad|\s)', '', s) - s = re.sub(r'\s+', '', s) - s = re.sub(r'\\sqrt([0-9A-Za-z])', r'\\sqrt{\1}', s) - s = re.sub(r'\\frac\{([^{}]+)\}([^{}\\])', r'\\frac{\1}{\2}', s) - s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) - s = re.sub(r'\{[a-zA-Z]+\}$', '', s) - s = re.sub(r'\^\\circ|\^\{\\circ\}|°|\\circ', 'deg', s) - s = s.replace(r'\minus{}', '-').replace(r'\minus', '-') - - def _frac_to_slash(mt): - text = mt.group(0) - pos = text.index('{') + 1 - depth, num_start = 1, pos - while depth > 0: - depth += (text[pos] == '{') - (text[pos] == '}') - pos += 1 - numer = text[num_start:pos - 1] - pos += 1 - den_start, depth = pos, 1 - while depth > 0: - depth += (text[pos] == '{') - (text[pos] == '}') - pos += 1 - return f'({numer})/({text[den_start:pos - 1]})' - - s = re.sub(r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', _frac_to_slash, s) - s = re.sub(r'(?<!\w)(\d+)/(\d+)(?!\w)', r'(\1)/(\2)', s) - return s - - -def _try_numeric_equal(a: str, b: str) -> bool: - try: - va, vb = float(a.replace('(', '').replace(')', '')), 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) - return va is not None and vb is not None and abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) - - -_MCQ_REF_RE = re.compile( - r'^\\(?:textbf|text|mathrm|mathbf)\{\(?([A-E])\)?\s*\}\s*(.+)$' - r'|^\(?([A-E])\)\s+(.+)$') -_VAR_PREFIX_RE = re.compile(r'^(?:[A-Za-z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)$', re.DOTALL) - - -def _split_mcq(ans: str): - s = ans.strip() - m = _MCQ_REF_RE.match(s) - if m: - return (m.group(1) or m.group(3)), ((m.group(2) or m.group(4) or '').strip() or None) - bl = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) - return (bl.group(1), None) if bl else (None, s or None) - - -def _strip_var_prefix(ans: str) -> str: - m = _VAR_PREFIX_RE.match((ans or '').strip()) - return m.group(1).strip() if m else (ans or '') - - -def _normalize_tuple(ans: str) -> str: - return re.sub(r'[\s()\[\]{}\\]', '', ans or '') - - -def _math_verify_equal(predicted: str, reference: str) -> bool: - try: - from math_verify import parse, verify - gold = parse(r'\boxed{' + reference + '}', parsing_timeout=1) - pred = parse(r'\boxed{' + predicted + '}', parsing_timeout=1) - return bool(gold and pred and verify(gold, pred, timeout_seconds=1)) - except Exception: - return False - - -def answers_match(predicted: str, reference: str) -> bool: - if not predicted or not reference: - return False - norm_p, norm_r = normalize_answer(predicted), normalize_answer(reference) - if norm_p == norm_r or norm_p.lower() == norm_r.lower() or _try_numeric_equal(norm_p, norm_r): - return True - stripped_p = normalize_answer(_strip_var_prefix(predicted)) - stripped_r = normalize_answer(_strip_var_prefix(reference)) - if stripped_p and stripped_r: - if (stripped_p == stripped_r or stripped_p.lower() == stripped_r.lower() - or _try_numeric_equal(stripped_p, stripped_r)): - return True - p_letter, p_value = _split_mcq(predicted) - r_letter, r_value = _split_mcq(reference) - if p_letter and r_letter and p_letter == r_letter: - return True - if (p_letter and p_value is None) and (r_letter and r_value is None): - return False - 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 and (p_val == r_val or p_val.lower() == r_val.lower() or _try_numeric_equal(p_val, r_val)): - return True - for left, right in ((norm_p, norm_r), (stripped_p, stripped_r)): - tl, tr = _normalize_tuple(left), _normalize_tuple(right) - if ',' in tl and tl == tr: - return True - if '=' in norm_r and '=' not in norm_p: - if any(part and (part == norm_p or _try_numeric_equal(part, norm_p)) for part in norm_r.split('=')): - return True - if '=' in norm_p and '=' not in norm_r: - if any(part and (part == norm_r or _try_numeric_equal(part, norm_r)) for part in norm_p.split('=')): - return True - return _math_verify_equal(stripped_p or norm_p, stripped_r or norm_r) - - -# =========================================================================== -# Block B -- prompts, skill parsing, batched sampling -# =========================================================================== -DIRECT_SYSTEM = ( - 'You are an expert competition mathematician. Solve the following problem ' - 'step by step. Provide your final answer inside \\boxed{}.') - -_SKILL_SOLVE_PREFIX = ( - DIRECT_SYSTEM + '\n\n' - 'Before you start, keep these reminders in mind to avoid common mistakes on this ' - 'type of problem:\n') -_SKILL_SOLVE_SUFFIX = '\nApply them where relevant, but rely on your own reasoning to reach the answer.' - - -def build_direct_prompt(problem: str) -> Dict[str, Any]: - return {'messages': [{'role': 'system', 'content': DIRECT_SYSTEM}, - {'role': 'user', 'content': problem}]} - - -def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: - skill = (skill or '').strip() - if not skill: - return build_direct_prompt(problem) - # Concatenation (not .format): DIRECT_SYSTEM/skill contain literal braces. - return {'messages': [ - {'role': 'system', 'content': _SKILL_SOLVE_PREFIX + skill + _SKILL_SOLVE_SUFFIX}, - {'role': 'user', 'content': problem}]} - - -# -- skill-gen prompts (view A: problem + rubric source; view B: query only) -- -# Kept deliberately short: this is the RL policy's system prompt, so over-specifying -# the output hurts convergence. The concrete output format is appended separately by -# ``_SKILL_OUTPUT`` so the task text stays format-agnostic. -SKILL_GEN_SYSTEM = ( - 'You are a math guidance writer. A process-check on a related problem hints at ' - 'likely mistakes. Write short reusable guidance for this and similar problems, ' - 'and note what to watch out for.\n') - -SKILL_GEN_SYSTEM_Q = ( - 'You are a math guidance writer. Write short reusable guidance for this and ' - 'similar problems.\n') - -_SKILL_OUTPUT = ( - 'Output only:\n<skills>\nYour reusable solving guidance here.\n</skills>') - -SKILL_GEN_USER_Q = ( - 'Problem:\n{problem}\n\n') - -SKILL_GEN_USER_RUBRIC = ( - 'Target problem:\n{problem}\n\n' - 'Problem used for the process check:\n{rubric_problem}\n\n' - 'Process check:\n' - '{diagnosis}\n\n') - - -def _rubric_has_fail(diagnosis: str) -> bool: - """View A uses the open-book rubric prompt (and, under --viewa-sft, SFT distillation) - IFF its rubric localised a failure. No ``[FAIL]`` => nothing to inject: generation - degrades to query-only and the problem is trained by GRPO exactly like view B. Single - source of truth so ``_skillgen_messages`` and ``_group_records`` never diverge.""" - return '[FAIL]' in (diagnosis or '') - - -def _skillgen_messages(problem: str, view: str, diagnosis: str, - rubric_problem: str = '') -> List[Dict[str, Any]]: - """Single source of truth for the skill-gen prompt (used at BOTH generation and - training so they never diverge). View A with a localisable failure uses the target - problem plus the rubric source problem and findings; view B -- or a view-A problem - whose rubric flagged nothing (no ``[FAIL]``) -- degrades to the query-only prompt.""" - if view == 'B' or not _rubric_has_fail(diagnosis): - return [{'role': 'system', 'content': SKILL_GEN_SYSTEM_Q + _SKILL_OUTPUT}, - {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=problem)}] - rubric_problem = rubric_problem or problem - return [{'role': 'system', 'content': SKILL_GEN_SYSTEM + _SKILL_OUTPUT}, - {'role': 'user', 'content': SKILL_GEN_USER_RUBRIC.format( - problem=problem, rubric_problem=rubric_problem, diagnosis=diagnosis)}] - - -def _assign_view(problem: str, args: argparse.Namespace) -> str: - h = int(hashlib.md5(f'{args.seed}:{problem}'.encode('utf-8')).hexdigest(), 16) - return 'B' if (h % 100000) / 100000.0 < args.view_b_frac else 'A' - - -def _view_prompt(r: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: - return {'messages': _skillgen_messages( - r['problem'], r['_view'], r.get('_rubric_diag', ''), r.get('_rubric_src', ''))} - - -_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') - - -def _clean_text(decoded: Optional[str]) -> str: - return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() - - -def _extract_tag_block(answer: str, tag: str, allow_empty: bool = False) -> Optional[str]: - low = answer.lower() - open_tag, close_tag = f'<{tag}>', f'</{tag}>' - s = low.rfind(open_tag) - if s < 0: - return None - inner = s + len(open_tag) - e = low.find(close_tag, inner) - if e < 0: - return None - block = answer[inner:e].strip() - block = re.sub(r'</?(?:skills|skill|diagnose|pitfall|strategy|think)>', '', block, flags=re.IGNORECASE).strip() - return block if (block or allow_empty) else None - - -def _extract_skill(text: str) -> Optional[str]: - """Parse skill-generation output: return the inner text of a non-empty ``<skills>`` - block, or None. If a ``</think>`` marker is present, parse only the text after the - last one; otherwise parse the full response.""" - low = text.lower() - end_think = low.rfind('</think>') - answer = text[end_think + len('</think>'):] if end_think >= 0 else text - return _extract_tag_block(answer, 'skills') - - -def _parse_seq(seq, gold: str) -> Dict[str, Any]: - """Grade one sampled sequence into a rollout record.""" - text = _clean_text(getattr(seq, 'decoded', '') or '') - pred = extract_boxed(text) - correct = bool(pred) and answers_match(pred, gold) - terminated = getattr(seq, 'stop_reason', None) != 'length' - return {'pred': pred, 'correct': correct, 'terminated': terminated, - 'passed': bool(correct and terminated), - 'stop_reason': getattr(seq, 'stop_reason', None), - 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} - - -def _run_samples(sampler, prompts: List[Any], num_samples: int, max_tokens: int, - gen_dp: int, temperature: Optional[float] = None, - top_p: Optional[float] = None, top_k: Optional[int] = None) -> List[List[Any]]: - """One batched sampler call -> per-prompt list of raw sequences. vLLM dp needs - batch len >= dp, so pad the tail and slice back.""" - if not prompts: - return [] - params = SamplingParams( - max_tokens=max_tokens, - temperature=GEN_TEMPERATURE if temperature is None else temperature, - top_p=GEN_TOP_P if top_p is None else top_p, - num_samples=num_samples, **({} if top_k is None else {'top_k': top_k})) - padded = prompts - if gen_dp > 1 and 0 < len(prompts) < gen_dp: - padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] - responses = sampler.sample(padded, params)[:len(prompts)] - return [list(r.sequences) if (r and r.sequences) else [] for r in responses] - - -# =========================================================================== -# Block C -- data loading via twinkle.Dataset + numeric filtering -# =========================================================================== -def _boxed_batch(rows: Dict[str, List[Any]], dataset: str) -> Dict[str, List[Any]]: - """Batched (columnar) mapper for ``Dataset.map``: annotate each row with its boxed - ``reference_answer`` and a ``_keep`` flag (aops also needs the ``boxed`` metadata).""" - sols = rows['solution'] - metas = rows.get('metadata', [None] * len(sols)) - refs = [extract_boxed(s or '') for s in sols] - keep = [bool(ref) and (dataset != 'aops' or bool((meta or {}).get('boxed'))) - for ref, meta in zip(refs, metas)] - return {**rows, 'reference_answer': refs, '_keep': keep} - - -def load_problems(dataset: str, n: int, seed: int, num_proc: int = 0) -> List[Dict[str, Any]]: - """Load boxed-answer problems as ``{problem, reference_answer, level?}`` via - twinkle.Dataset (ModelScope hub), sampled to ``n`` (0 = all). Boxed extraction (regex - + brace scan over every solution) is parallelised by ``Dataset.map``/``.filter``; - ``num_proc`` defaults to all cores (set 1 to force serial).""" - ds_id = AOPS_DATASET_ID if dataset == 'aops' else MATH_DATASET_ID - ds = Dataset(DatasetMeta(dataset_id=f'ms://{ds_id}', split='train')) - nproc = num_proc if num_proc > 0 else min(32, os.cpu_count() or 1) - ds.map(lambda rows: _boxed_batch(rows, dataset), num_proc=nproc) - ds.filter(lambda row: row['_keep'], num_proc=nproc) - has_level = 'level' in ds.dataset.column_names - out = [{'data_id': f'{dataset}:{i}', 'problem': row['problem'], - 'reference_answer': row['reference_answer'], - **({'level': row['level']} if has_level and row.get('level') else {})} - for i, row in enumerate(ds.dataset)] - logger.info(f'[data] {dataset}: {len(out)} boxed problems ({nproc} procs)') - rng = np.random.RandomState(seed) - rng.shuffle(out) - return out[:n] if (n and n < len(out)) else out - - -# --------------------------------------------------------------------------- -# Bag-of-words neighbour pairing (cross-problem rubric transfer, --xproblem-rubric) -# --------------------------------------------------------------------------- -# Common English + math-scaffolding words that carry no problem-type signal. Kept small -# and deterministic on purpose (no external stopword list): what survives is the domain -# vocabulary ('triangle', 'prime', 'polynomial', ...) that actually defines the type. -_BOW_STOP = frozenset(""" -a an the of to in on at for and or but if is are be was were been being this that these those -with without into onto from by as it its their his her our your my we you they he she them -find compute determine calculate evaluate solve show prove given let suppose consider assume -what which when where how many much value values number numbers expression form terms term -such that then than so if only when each every all any some both one two three four five six -seven eight nine ten first second third last non over under about above below between -problem answer result equal equals sum difference product total following there here have has -had do does did can could will would should may might must not no yes if then else -""".split()) - -_WORD_RE = re.compile(r'[a-z]+') - - -def _stem(w: str) -> str: - """Crude suffix stripper so 'prime'/'primes', 'triangle'/'triangles' collapse to one - type token. Not linguistically correct -- just enough to merge the common plural/gerund - variants that otherwise split a type's vocabulary and starve the df filter.""" - if len(w) > 4 and w.endswith('ies'): - return w[:-3] + 'y' # properties -> property - if len(w) > 4 and w.endswith('es') and w[-3] in 'sxzh': - return w[:-2] # boxes -> box (keep primes -> prime below) - for suf in ('ing', 'ed', 's'): - if len(w) > len(suf) + 2 and w.endswith(suf): - return w[:-len(suf)] - return w - - -def _tokenize(problem: str) -> List[str]: - """Deterministic bag-of-words tokens for type matching: lowercase alphabetic words - (numbers dropped -- they are instance detail, not type), minus generic stopwords, then - stemmed, so only domain terms remain. Latex control words (frac, sqrt, ...) survive.""" - return [_stem(w) for w in _WORD_RE.findall((problem or '').lower()) - if len(w) > 2 and w not in _BOW_STOP] - - -class BagOfWordsIndex: - """TF-IDF cosine nearest-neighbour over problem statements. Sparse dict vectors + - an inverted index, so ``nearest`` scores only problems sharing a term (near-linear in - practice, no dense NxN). Deterministic; ties break on lower index for reproducibility. - - Answer-aware: ``nearest`` skips the same index, near-verbatim duplicates (cosine - >= ``sim_max``) and -- crucially for anti-leak -- any candidate whose answer equals the - query's, so a neighbour rubric can never hand over the query's own answer.""" - - def __init__(self, problems: List[str], answers: Optional[List[str]] = None, - min_df: int = 2, max_df_frac: float = 0.5): - self._toks = [_tokenize(p) for p in problems] - self._ans = [(_numeric_value(a) or (str(a).strip() if a else '')) for a in answers] \ - if answers is not None else [''] * len(problems) - n = len(self._toks) - df: Dict[str, int] = {} - for toks in self._toks: - for w in set(toks): - df[w] = df.get(w, 0) + 1 - max_df = max(min_df, int(max_df_frac * n)) - self._idf = {w: math.log((n + 1) / (c + 1)) + 1.0 - for w, c in df.items() if min_df <= c <= max_df} - self._vecs: List[Dict[str, float]] = [self._vectorize(t) for t in self._toks] - self._inverted: Dict[str, List[int]] = {} - for i, v in enumerate(self._vecs): - for w in v: - self._inverted.setdefault(w, []).append(i) - - def _vectorize(self, toks: List[str]) -> Dict[str, float]: - tf: Dict[str, float] = {} - for w in toks: - if w in self._idf: - tf[w] = tf.get(w, 0.0) + 1.0 - vec = {w: c * self._idf[w] for w, c in tf.items()} - norm = math.sqrt(sum(x * x for x in vec.values())) - return {w: x / norm for w, x in vec.items()} if norm > 0 else {} - - def nearest(self, i: int, sim_max: float = 0.98) -> Tuple[int, float]: - """(index, cosine) of the most similar *distinct* problem: not i, not a duplicate - (cosine < ``sim_max``) and not the same answer. (-1, 0.0) if none qualifies.""" - vi = self._vecs[i] - if not vi: - return -1, 0.0 - ai = self._ans[i] - scores: Dict[int, float] = {} - for w, xi in vi.items(): - for j in self._inverted.get(w, ()): - if j != i: - scores[j] = scores.get(j, 0.0) + xi * self._vecs[j].get(w, 0.0) - best_j, best_s = -1, 0.0 - for j, s in scores.items(): - if s >= sim_max or (ai and self._ans[j] == ai): - continue - if s > best_s or (s == best_s and (best_j < 0 or j < best_j)): - best_j, best_s = j, s - return best_j, best_s - - -def build_pairs(records: List[Dict[str, Any]], n: int, seed: int, sim_max: float = 0.98 - ) -> Tuple[List[Dict[str, Any]], Dict[str, Tuple[str, float]]]: - """Single-pass cross-problem pairing over the whole pool (one index build). - - Returns ``(subset, neighbour_map)`` where ``subset`` is the ``n`` problems with the - strongest qualifying neighbour (dense same-type pairs; n<=0/n>=len keeps all, shuffled) - and ``neighbour_map[q] = (p, cosine)`` gives each kept problem its analogue P. P is drawn - from the *full* pool (richer analogues) but never a duplicate or same-answer problem, so - P's rubric can transfer method without ever leaking Q's answer.""" - index = BagOfWordsIndex([r['problem'] for r in records], - [str(r.get('reference_answer', '')) for r in records]) - nbr = [index.nearest(i, sim_max) for i in range(len(records))] - order = sorted(range(len(records)), key=lambda i: (-nbr[i][1], i)) - keep = order[:n] if (0 < n < len(records)) else list(range(len(records))) - rng = np.random.RandomState(seed) - rng.shuffle(keep) - subset = [records[i] for i in keep] - neighbour_map = {records[i]['problem']: (records[nbr[i][0]]['problem'], nbr[i][1]) - for i in keep if nbr[i][0] >= 0} - return subset, neighbour_map - - -_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') - - -def _norm_num_text(num: str) -> str: - try: - f = float(num) - return str(int(f)) if f == int(f) else str(f) - except Exception: - return str(num).strip() - - -def _numeric_value(raw: Any) -> Optional[str]: - """Collapse an answer to a single int/decimal/fraction, or None.""" - if raw is None: - return None - s = str(raw).strip().strip('$').strip() - s = s.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac') - s = re.sub(r'\\!|\\,|\\;|\\ |\\left|\\right|\s', '', s) - for pat in (r'\\frac\{(-?\d+)\}\{(-?\d+)\}', r'(-?\d+)/(-?\d+)'): - m = re.fullmatch(pat, s) - if m: - a, b = int(m.group(1)), int(m.group(2)) - return _norm_num_text(str(a / b)) if b else None - return _norm_num_text(s) if _NUM_RE.fullmatch(s) else None - - -def _answer_leaked(skill: str, reference: str) -> bool: - """Audit whether a generated skill contains the final answer verbatim. This is NOT - a training filter: if the skill model derives an answer from the problem, that is a - legitimate answer-bearing skill under this experiment. The real leakage boundary is the - external API/rubric diagnosis, which is constrained by prompt to stay answer-free.""" - if not skill: - return False - for cand in {_numeric_value(reference), (str(reference).strip() or None)}: - if cand and re.search(r'(?<![\d.])' + re.escape(cand) + r'(?![\d.])', skill): - return True - return False - - -def _load_excluded_records(paths_arg: str) -> Tuple[Set[str], Set[str]]: - """Read jsonl files and collect data_id/problem keys that must be excluded. - - The cold-start SFT builder writes stable ``data_id`` values. ``problem`` is kept as a - backward-compatible fallback for older jsonl files produced before data_id existed.""" - ids: Set[str] = set() - problems: Set[str] = set() - for raw_path in (paths_arg or '').split(','): - path = raw_path.strip() - if not path or not os.path.exists(path): - continue - with open(path, encoding='utf-8') as f: - for line in f: - if not line.strip(): - continue - row = json.loads(line) - if row.get('record_type') in {'config', 'summary'}: - continue - data_id = str(row.get('data_id') or '').strip() - problem = str(row.get('problem') or '').strip() - if data_id: - ids.add(data_id) - elif problem: - problems.add(problem) - return ids, problems - - -def _load_seam_parquet(path: str) -> List[Dict[str, Any]]: - """Read a SEAM ``build_aops_dataset.py`` parquet (VERL RLHF schema) into twinkle records, - PRESERVING file row order. ``problem <- extra_info.problem`` and - ``reference_answer <- reward_model.ground_truth``. No shuffle/filter: - the parquet is already SEAM's numeric-filtered, seed-42-shuffled, truncated split.""" - import pyarrow.parquet as pq - rows = pq.read_table(path).to_pylist() - out: List[Dict[str, Any]] = [] - for i, r in enumerate(rows): - ei = r.get('extra_info') or {} - rm = r.get('reward_model') or {} - problem = (ei.get('problem') or '').strip() - ref = rm.get('ground_truth') - if not problem or ref is None: - continue - out.append({'problem': problem, 'reference_answer': str(ref), - 'data_id': f"seam:{ei.get('split', '')}:{ei.get('index', i)}"}) - return out - - -def _load_records_from_seam(args: argparse.Namespace, seam_dir: str - ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], - Dict[str, Tuple[str, float]], Dict[str, str], Dict[str, int]]: - """Data entry that mirrors a SEAM run EXACTLY: read ``train.parquet``/``val.parquet`` from - ``seam_dir`` in file order, use ``val`` as the eval holdout, take the first ``--n`` train rows - (post ``--pool-offset``) with NO shuffle. ``--numeric-only``/``--eval-size``/internal shuffle - are bypassed (the parquet is already the authoritative split).""" - tp, vp = os.path.join(seam_dir, 'train.parquet'), os.path.join(seam_dir, 'val.parquet') - if not (os.path.exists(tp) and os.path.exists(vp)): - raise FileNotFoundError( - f'--seam-parquet-dir needs both train.parquet and val.parquet in {seam_dir}') - if args.xproblem_rubric: - raise ValueError('--xproblem-rubric is unsupported with --seam-parquet-dir ' - '(SEAM parquet carries no neighbour structure).') - pool = _load_seam_parquet(tp) # already SEAM-shuffled + truncated, in file order - eval_records = [dict(r) for r in _load_seam_parquet(vp)] # SEAM's exact val holdout - pool_offset = max(0, int(getattr(args, 'pool_offset', 0) or 0)) - if pool_offset: - if pool_offset >= len(pool): - raise ValueError(f'--pool-offset ({pool_offset}) leaves no train records ' - f'from SEAM train pool size {len(pool)}') - pool = pool[pool_offset:] - train_n = args.n if args.n > 0 else len(pool) - train_records = [dict(r) for r in pool[:train_n]] - if {r['problem'] for r in train_records} & {r['problem'] for r in eval_records}: - raise ValueError('eval/train overlap detected in SEAM parquet') - stats = {'raw_loaded': len(pool) + len(eval_records), 'numeric_dropped': 0, - 'excluded_records': 0, 'pool_offset': pool_offset, - 'train_records': len(train_records), 'eval_records': len(eval_records), - 'source': 'seam_parquet', 'seam_parquet_dir': seam_dir} - return train_records, eval_records, {}, {}, stats - - -def _load_records(args: argparse.Namespace - ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], - Dict[str, Tuple[str, float]], Dict[str, str], Dict[str, int]]: - """Load, numeric-filter, shuffle, split a fixed eval holdout, and (when --xproblem-rubric) - select a same-type-dense train subset with its neighbour map -- all in one pass. - Also returns ``pool_answers`` (every candidate neighbour P's true answer) so P's baseline - can be graded/cached correctly even when P is not itself a training problem.""" - seam_dir = (getattr(args, 'seam_parquet_dir', '') or '').strip() - if seam_dir: # read SEAM parquet in file order, bypassing load/filter/shuffle/split - return _load_records_from_seam(args, seam_dir) - # Load all when filtering or splitting (else the eval holdout could starve train). - load_n = 0 if (args.numeric_only or args.eval_size > 0) else args.n - records = load_problems(args.dataset, load_n, args.seed) - raw_n, dropped = len(records), 0 - if args.numeric_only: - records = [{**r, 'reference_answer': v} - for r, v in ((r, _numeric_value(r.get('reference_answer'))) for r in records) - if v is not None] - dropped = raw_n - len(records) - np.random.RandomState(args.seed).shuffle(records) - exclude_ids, exclude_problems = _load_excluded_records(getattr(args, 'exclude_data_ids', '')) - excluded = 0 - if exclude_ids or exclude_problems: - before = len(records) - records = [r for r in records - if str(r.get('data_id', '')) not in exclude_ids - and str(r.get('problem', '')).strip() not in exclude_problems] - excluded = before - len(records) - eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 - eval_records = [dict(r) for r in records[:eval_n]] - pool = records[eval_n:] - pool_offset = max(0, int(getattr(args, 'pool_offset', 0) or 0)) - if pool_offset: - if pool_offset >= len(pool): - raise ValueError(f'--pool-offset ({pool_offset}) leaves no train records from pool size {len(pool)}') - pool = pool[pool_offset:] - train_n = args.n if args.n > 0 else len(pool) - # Cross-problem rubric: pick the train_n problems with the strongest qualifying neighbour - # (dense same-type pairs) and build {Q -> (P, sim)} in one index pass. Otherwise keep the - # first train_n (already shuffled) with no neighbours. - if args.xproblem_rubric: - subset, neighbor_map = build_pairs(pool, train_n, args.seed) - train_records = [dict(r) for r in subset] - pool_answers = {r['problem']: str(r.get('reference_answer', '')) for r in pool} - else: - train_records, neighbor_map, pool_answers = [dict(r) for r in pool[:train_n]], {}, {} - if {r['problem'] for r in train_records} & {r['problem'] for r in eval_records}: - raise ValueError('eval/train overlap detected') - stats = {'raw_loaded': raw_n, 'numeric_dropped': dropped, - 'excluded_records': excluded, 'pool_offset': pool_offset, - 'train_records': len(train_records), 'eval_records': len(eval_records)} - return train_records, eval_records, neighbor_map, pool_answers, stats - - -# =========================================================================== -# Block D -- disk cache, problem pool, baseline rollout, rubric check -# =========================================================================== -class DiskCache: - """Append-only jsonl kv cache (md5 key -> value). Loads on init, appends on put. - Disabled instances always miss and never write.""" - - def __init__(self, path: str, enabled: bool = True): - self._mem: Dict[str, Any] = {} - self._fh = None - self._lock = threading.Lock() # base baseline is prefetched on a background thread - if not enabled: - return - if os.path.exists(path): - with open(path, encoding='utf-8') as f: - for line in f: - if line.strip(): - row = json.loads(line) - self._mem[row['key']] = row['value'] - self._fh = open(path, 'a', encoding='utf-8') - - @staticmethod - def key_for(*parts: str) -> str: - return hashlib.md5('\x1f'.join(parts).encode('utf-8')).hexdigest() - - def __contains__(self, key: str) -> bool: - with self._lock: - return key in self._mem - - def get(self, key: str) -> Any: - with self._lock: - return self._mem.get(key) - - def put(self, key: str, value: Any) -> None: - with self._lock: - self._mem[key] = value - if self._fh is not None: - self._fh.write(json.dumps({'key': key, 'value': value}, ensure_ascii=False) + '\n') - self._fh.flush() - - def close(self) -> None: - if self._fh is not None: - self._fh.close() - - -class _LockedSampler: - """Serialises every ``.sample()`` on a sampler behind one lock. The base sampler is - shared by the main thread (with-skill scoring, eval) and the baseline-prefetch thread; - ``sample`` is a slice_dp/flatten remote call whose collect is NOT safe to interleave - across two callers, so concurrent calls could mis-join sequences. The lock keeps base - calls serial (prefetch still overlaps the skill-gen phase, which uses skill_sampler).""" - - def __init__(self, sampler): - self._sampler = sampler - self._lock = threading.Lock() - - def sample(self, *args, **kwargs): - with self._lock: - return self._sampler.sample(*args, **kwargs) - - def __getattr__(self, name): - return getattr(self._sampler, name) - - -class ProblemPool: - """Cyclic draw source. Each full pass reshuffles with ``seed + epoch``; the initial - pass keeps the loader order. ``draw(k)`` returns k distinct problems (k << pool).""" - - def __init__(self, records: List[Dict[str, Any]], seed: int): - self._records = list(records) - self._seed, self._cursor, self.epoch = seed, 0, 0 - - def draw(self, k: int) -> List[Dict[str, Any]]: - out, seen = [], set() - while len(out) < k: - if self._cursor >= len(self._records): - self.epoch += 1 - np.random.RandomState(self._seed + self.epoch).shuffle(self._records) - self._cursor = 0 - r = self._records[self._cursor] - self._cursor += 1 - if id(r) not in seen: - seen.add(id(r)) - out.append(r) - return out - - def peek(self, k: int) -> List[Dict[str, Any]]: - """The next k distinct problems draw() would return, WITHOUT advancing state - (no cursor move, no reshuffle). Used to prefetch their baseline into base_cache - while the current chunk trains -- a wrong guess (cross-epoch reshuffle) only - misses the cache, never corrupts the draw.""" - out, seen, cur = [], set(), self._cursor - recs = self._records - while len(out) < k and cur < len(recs): # stop at epoch edge; don't simulate reshuffle - r = recs[cur] - cur += 1 - if id(r) not in seen: - seen.add(id(r)) - out.append(r) - return out - - -def _empty_roll() -> Dict[str, Any]: - return {'pred': '', 'correct': False, 'terminated': False, 'passed': False, - 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} - - -def _apply_baseline(r: Dict[str, Any], roll: Dict[str, Any]) -> None: - """Attach a greedy baseline roll and reset per-chunk working state.""" - r['_baseline_rolls'], r['_cands'], r['_init'] = [roll], [], [roll] - r['_failed'] = not roll['correct'] - r['_baseline_pass'] = 1.0 if roll['correct'] else 0.0 - r['_hard'] = True # process every problem; group variance selects (SEAM-style) - - -def baseline_rollout(base_sampler, problems: List[Dict[str, Any]], base_dp: int, - args: argparse.Namespace, cache: DiskCache) -> int: - """Base solves each problem greedily once (T=0, M=1), disk-cached by problem text. - The base is frozen + greedy so a cache hit is exact. Returns fresh (miss) count.""" - todo = [r for r in problems if DiskCache.key_for(r['problem']) not in cache] - if todo: - out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], - 1, args.max_tokens, base_dp, temperature=0.0) - for r, seqs in zip(todo, out): - roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - cache.put(DiskCache.key_for(r['problem']), roll) - for r in problems: - _apply_baseline(r, cache.get(DiskCache.key_for(r['problem']))) - return len(todo) - - -# -- rubric process-check (view A): teacher diagnoses the base's attempt -- -_RFT_DIAG_SYSTEM = """\ -You are a strategy-level process checker for a math solution attempt. You are given a -math problem, a rubric, and one attempted solution segment. Decide PASS or FAIL for each -criterion, and write the diagnosis so it can become useful reusable guidance for solving -similar problems without seeing this segment. - -Output STRICT JSON (no prose outside it) with this shape: -{ - "items": [ - {"index": 1, "verdict": "PASS", "reason": "<why the process satisfies it>", - "fix": ""}, - {"index": 2, "verdict": "FAIL", "reason": "<what reusable process issue is present>", - "fix": "<local strategy correction, without solving the problem>"} - ], - "overall": "OK" | "ISSUES", - "summary": "<one sentence naming the reusable process issue, not the answer>" -} - -Rules: -- Judge every criterion independently; a [Hard Rule] is FAIL unless clearly satisfied. -- Judge ONLY what is observable in THIS segment. Ignore hidden <think> or <thinking> - content for output-format criteria. -- The API diagnosis is an external teacher signal, so it must stay answer-free. -- Prefer diagnosis that transfers to view-B skill generation: name the route choice, - structural observation, missing check, or length-control habit that a solver should - remember before solving a similar problem. -- For PASS items, leave "fix" as "". -- For FAIL items, describe the process problem at strategy level: unsuitable method, - missed structure, invalid transformation, missing constraint check, redundant cases, - off-track approach, contradiction, or inefficient/unfinished reasoning. -- A fix may suggest the LOCAL correction direction, such as identify the key structure, - verify constraints, preserve equivalence, reduce redundant cases, or choose a more - direct route. Do not carry out the correction. -- Never reveal the final answer, a corrected value/expression, an option label, or a - step-by-step solution that would let another model copy the solve. -- If the segment contains a process note saying it was cut off before a final boxed - answer, mark the length-budget criterion as FAIL and suggest a method-level way to - finish faster. -- Keep every "reason" and "fix" concise: one short sentence each. -- "overall" is "OK" only if NO criterion is FAIL. -- Output only the JSON object.""" - -_RFT_DIAG_USER = """\ -## Task / query (context) -{query} - -## Rubric -{rubric} - -## Segment -{segment} - -Now output the diagnostic JSON object.""" - -_MATH_RUBRIC = [ - ('The attempt chooses a method suitable for the problem structure', False), - ('The attempt identifies the key constraint, invariant, or quantity before computing', False), - ('Algebraic and logical transformations preserve validity at each step', True), - ('The attempt checks required constraints, domains, boundary cases, or validity conditions', False), - ('The attempt avoids redundant casework, looping, or re-deriving known facts', False), - ('The attempt reaches a final boxed answer within the length budget', False), - ('The approach stays focused on the actual question asked', False), -] - -# Bump when the rubric criteria or the diagnosis prompt change so the disk-cached -# diagnoses written under an older rubric are not silently reused. -_RUBRIC_VERSION = 'rubric_v5_viewb_strategy' - - -class _RftRubricVerifier(RubricVerifier): - def _diagnose_trajectory(self, query: str, rubric_block: str, segment_text: str) -> dict: - return {'messages': [ - {'role': 'system', 'content': _RFT_DIAG_SYSTEM}, - {'role': 'user', 'content': _RFT_DIAG_USER.format( - query=query, rubric=rubric_block, segment=segment_text)}]} - - -def build_rubric_checker() -> Optional[RubricVerifier]: - """Fixed math-process rubric verifier (teacher-served). None if no LLM backup env.""" - if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') - or os.environ.get('OPENAI_API_KEY')): - return None - return _RftRubricVerifier( - fixed_rubric=[RubricItem(t, is_hard=h) for t, h in _MATH_RUBRIC], gate=True) - - -def _format_diagnosis(detail) -> str: - """One line per criterion (PASS/FAIL + reason + fix) then a summary.""" - rub = detail.rubric - lines = [] - for it in detail.items: - text = rub[it.index - 1].text if 0 < it.index <= len(rub) else f'criterion {it.index}' - if it.verdict: - lines.append(f'- [PASS] {text}') - else: - tail = f': {it.reason}' if it.reason else '' - tail += f' (fix: {it.fix})' if it.fix else '' - lines.append(f'- [FAIL] {text}{tail}') - if detail.summary: - lines.append(f'Summary: {detail.summary}') - return '\n'.join(lines) - - -def diagnose_views(checker, problems: List[Dict[str, Any]], args: argparse.Namespace, - cache: DiskCache) -> None: - """Rubric-check every view-A problem's greedy attempt in parallel (disk-cached by - problem + attempt), stashing the formatted findings on ``r['_rubric_diag']``.""" - targets = [r for r in problems if r.get('_view') == 'A'] - if not checker or not targets: - return - - def _key(r: Dict[str, Any]) -> str: - init = r.get('_init', [{}])[0] - term = 'L' if (init.get('stop_reason') == 'length' or not init.get('terminated')) else 'T' - return DiskCache.key_for(r['problem'], init.get('text', ''), _RUBRIC_VERSION, term) - - pending = [] - for r in targets: - key = _key(r) - if key in cache: - r['_rubric_diag'] = cache.get(key) - else: - pending.append((r, key)) - if not pending: - return - - def _run(item): - r, key = item - init = r['_init'][0] - seg_text = init['text'] - if init.get('stop_reason') == 'length' or not init.get('terminated'): - seg_text += ('\n\n[Process note: this attempt was cut off at the token budget ' - 'and never produced a final \\boxed{} answer.]') - seg = {'messages': [{'role': 'user', 'content': r['problem']}, - {'role': 'assistant', 'content': seg_text}]} - attempts = max(1, args.rubric_retries + 1) - for attempt in range(attempts): - try: - return r, key, _format_diagnosis(checker.diagnose(seg, query=r['problem'])) - except Exception as exc: # teacher hiccup -> retry, then no-diagnosis prompt (not cached) - if attempt + 1 < attempts: - logger.warning(f'[rubric] diagnose error: {exc}; retry {attempt + 1}/{args.rubric_retries}') - time.sleep(min(2.0, 0.5 * (2 ** attempt))) - continue - logger.warning(f'[rubric] diagnose error: {exc}; giving up after {attempts} attempts') - return r, key, None - - workers = max(1, min(args.rubric_workers, len(pending))) - with ThreadPoolExecutor(max_workers=workers) as ex: - for r, key, diag in ex.map(_run, pending): - r['_rubric_diag'] = diag or '' - if diag is not None: - cache.put(key, diag) - - -# =========================================================================== -# Block E -- chunk draw, generation pipeline, record building -# =========================================================================== -def _baseline_class(r: Dict[str, Any]) -> str: - """success | fail_loop (out of length / never terminated) | fail_wrong.""" - roll = r['_init'][0] - if roll['correct']: - return 'success' - return 'fail_loop' if (roll['stop_reason'] == 'length' or not roll['terminated']) else 'fail_wrong' - - -def _select_balanced(buckets: Dict[str, List[Dict[str, Any]]], n_success: int, - n_fail: int, n_fail_loop: int) -> List[Dict[str, Any]]: - """Pick n_fail base-fails (toward n_fail_loop loop-fails, best-effort) + n_success - base-successes; top up any shortfall from leftovers.""" - loop, wrong, succ = buckets['fail_loop'], buckets['fail_wrong'], buckets['success'] - take_loop = min(n_fail_loop, len(loop)) - take_wrong = min(n_fail - take_loop, len(wrong)) - take_loop = min(n_fail - take_wrong, len(loop)) - sel = loop[:take_loop] + wrong[:take_wrong] + succ[:n_success] - target = n_success + n_fail - if len(sel) < target: - used = {id(x) for x in sel} - sel += [x for b in (loop, wrong, succ) for x in b if id(x) not in used][:target - len(sel)] - return sel - - -def draw_chunk(pool: ProblemPool, base_sampler, base_dp: int, args: argparse.Namespace, - cache: DiskCache) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: - """Draw one chunk, baselining every drawn problem. With ``--balance``, keep - drawing+baselining until the target base fail:success mix is reachable (or the budget - is hit), then select a balanced subset.""" - if not args.balance: - chunk = pool.draw(args.chunk_size) - fresh = baseline_rollout(base_sampler, chunk, base_dp, args, cache) - return chunk, {'enabled': False, 'n_drawn': len(chunk), 'n_baseline_fresh': fresh} - - n_success = max(0, min(args.chunk_size, round(args.chunk_size * args.balance_success_frac))) - n_fail = args.chunk_size - n_success - n_fail_loop = round(n_fail * args.balance_loop_frac) - buckets: Dict[str, List[Dict[str, Any]]] = {'success': [], 'fail_loop': [], 'fail_wrong': []} - budget, n_drawn, n_fresh, seen = args.chunk_size * args.balance_max_draws_mult, 0, 0, set() - while n_drawn < budget: - if (len(buckets['success']) >= n_success - and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail): - break - batch = pool.draw(args.chunk_size) - n_fresh += baseline_rollout(base_sampler, batch, base_dp, args, cache) - n_drawn += len(batch) - for r in batch: - if id(r) not in seen: - seen.add(id(r)) - buckets[_baseline_class(r)].append(r) - - reached = (len(buckets['success']) >= n_success - and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail) - chunk = _select_balanced(buckets, n_success, n_fail, n_fail_loop) - sel_success = sum(1 for r in chunk if not r['_failed']) - stats = { - 'enabled': True, 'n_drawn': n_drawn, 'n_baseline_fresh': n_fresh, 'n_selected': len(chunk), - 'target_success': n_success, 'target_fail': n_fail, 'target_fail_loop': n_fail_loop, - 'selected_success': sel_success, 'selected_fail': len(chunk) - sel_success, - 'selected_fail_loop': sum(1 for r in chunk if _baseline_class(r) == 'fail_loop'), - 'selected_fail_wrong': sum(1 for r in chunk if _baseline_class(r) == 'fail_wrong'), - 'selected_success_frac': (sel_success / len(chunk)) if chunk else 0.0, - 'budget_hit': not reached, - } - return chunk, stats - - -def _assign_advantages(hard: List[Dict[str, Any]], args: argparse.Namespace) -> None: - """Group-relative advantage over each problem's scored candidates using the greedy - binary reward R in {0,1}: ``A = (R - mean) / (std + eps)``. std==0 groups get no - gradient -- GRPO's variance selects informative problems (no explicit difficulty gate). - A symmetric clip limits both positive and negative outliers from sparse 1-vs-many groups.""" - eps = 1e-6 - adv_clip = abs(float(getattr(args, 'adv_clip', 0.0) or 0.0)) - for r in hard: - for c in r['_cands']: - c['advantage'], c['grpo_adv'], c['kept'] = 0.0, 0.0, False - cs = [c for c in r['_cands'] if c.get('reward') is not None] - if len(cs) < 2: - continue - rewards = [c['reward'] for c in cs] - mean_r = sum(rewards) / len(rewards) - std = (sum((x - mean_r) ** 2 for x in rewards) / len(rewards)) ** 0.5 - if std < 1e-9: - continue - for c in cs: - raw_adv = (c['reward'] - mean_r) / (std + eps) - adv = max(-adv_clip, min(adv_clip, raw_adv)) if adv_clip > 0 else raw_adv - c['advantage'], c['grpo_adv'], c['kept'] = adv, adv, c['reward'] > mean_r - - -def _best_sft_candidate(r: Dict[str, Any], args: argparse.Namespace) -> Optional[Dict[str, Any]]: - """Pick ONE view-A candidate to distill (online context distillation). PREFER the - executor-verified PASSING skills (reward==1); if NONE passed -- common on the hard - problems that are exactly the cases worth distilling -- FALL BACK to any parseable - open-book skill regardless of the executor outcome. Answer-bearing skills produced by - the skill model itself are allowed here; only the external API/rubric diagnosis must be - answer-free. Within the chosen tier, take the one whose skill length is CLOSEST to - ``--sft-target-len`` -- an empirically high-pass-rate length (~500-600 chars in this - run) -- breaking ties by the fewest executor solve tokens. Targeting a length (rather - than the minimum) avoids a distillation feedback loop that would otherwise drive - rollouts ever shorter. None only when no parseable candidate exists at all.""" - eligible = [c for c in r['_cands'] if c.get('parseable') and c.get('skills')] - if not eligible: - return None - passing = [c for c in eligible if c.get('reward') == 1.0] - cs = passing or eligible - target = int(getattr(args, 'sft_target_len', 550) or 550) - - def _solve_tokens(c: Dict[str, Any]) -> int: - rolls = c.get('rolls') or [] - return rolls[0].get('gen_tokens', 1 << 30) if rolls else (1 << 30) - - return min(cs, key=lambda c: (abs(len(c['skills']) - target), _solve_tokens(c))) - - -def apply_neighbor_rubric(base_sampler, chunk: List[Dict[str, Any]], - neighbor_map: Dict[str, Tuple[str, float]], - pool_answers: Dict[str, str], base_dp: int, - args: argparse.Namespace, checker, - base_cache: DiskCache, rubric_cache: DiskCache) -> None: - """Cross-problem rubric (--xproblem-rubric): for every view-A problem Q, replace its own - rubric with the rubric of its bag-of-words neighbour P (Q keeps being the solved/scored - problem). P is baselined + diagnosed here (both disk-cached) as a stub carrying P's REAL - answer, so P's baseline grades correctly and legitimately shares the baseline cache with - P-as-training-problem. P's findings are copied onto Q with the neighbour text + similarity - for audit. A P that can't be diagnosed -> Q degrades to view B. Since P's answer differs - from Q's (guaranteed at pairing time), any answer P's rubric leaks is useless for Q.""" - targets = [r for r in chunk if r.get('_view') == 'A' and neighbor_map.get(r['problem'])] - if not targets: - return - stubs, by_problem = [], {} - for r in targets: - p, _ = neighbor_map[r['problem']] - if p not in by_problem: - stub = {'problem': p, 'reference_answer': pool_answers.get(p, ''), '_view': 'A'} - by_problem[p] = stub - stubs.append(stub) - baseline_rollout(base_sampler, stubs, base_dp, args, base_cache) - diagnose_views(checker, stubs, args, rubric_cache) - for r in targets: - p, sim = neighbor_map[r['problem']] - r['_rubric_diag'] = by_problem[p].get('_rubric_diag', '') - r['_rubric_src'], r['_neighbor_sim'] = p, sim - - -def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], - ci: int, base_dp: int, skill_dp: int, args: argparse.Namespace, - checker, rubric_cache: DiskCache, base_cache: DiskCache = None, - neighbor_map: Optional[Dict[str, Tuple[str, float]]] = None, - pool_answers: Optional[Dict[str, str]] = None - ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], List[Dict[str, Any]]]: - """view assign -> rubric-check (view A) -> skill-gen -> answer audit -> with-skill - greedy pass -> GRPO advantages. ``chunk`` arrives already baselined by draw_chunk. - With --xproblem-rubric, view A gets its NEIGHBOUR's rubric (see apply_neighbor_rubric).""" - hard = chunk - for r in hard: - r['_view'], r['_rubric_diag'] = _assign_view(r['problem'], args), '' - if args.xproblem_rubric and neighbor_map: - apply_neighbor_rubric(base_sampler, hard, neighbor_map, pool_answers or {}, base_dp, - args, checker, base_cache, rubric_cache) - else: - diagnose_views(checker, hard, args, rubric_cache) - - # skill-gen (thinking OFF), per-view prompt; re-sample problems with no clean candidate. - # View-A problems that cannot yield a training record (no [FAIL] rubric, or rubric - # leaked the answer) are dropped from training entirely -- skip their generation. - flat: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] - pending = [r for r in hard if not _viewa_dropped(r, args)] - for _ in range(args.skill_retries + 1): - if not pending: - break - sg_out = _run_samples(skill_sampler, [_view_prompt(r, args) for r in pending], - args.n_skills, args.skill_max_tokens, skill_dp, - temperature=args.skill_gen_temperature, - top_p=args.skill_gen_top_p, top_k=args.skill_gen_top_k) - still = [] - for r, seqs in zip(pending, sg_out): - got = False - for s in seqs: - resp = _clean_text(getattr(s, 'decoded', '') or '') - block = _extract_skill(resp) or '' - cand = {'skills': block, 'response': resp, 'parseable': bool(block), - 'view': r['_view'], 'leaked': None, 'leak_reason': '', - 'leak_source': '', 'with_pass': None, 'reward': None, 'rolls': [], - 'skillgen_stop': getattr(s, 'stop_reason', None), - 'skillgen_tokens': len(getattr(s, 'tokens', None) or [])} - r['_cands'].append(cand) - if block: - flat.append((r, cand)) - got = True - if not got: - still.append(r) - pending = still - - # Answer-bearing skill audit: deterministic verbatim-answer check only (no LLM). This - # is observability only; it records metrics for swanlab/jsonl, but does not block - # scoring, reward, advantage assignment, SFT candidate selection, or GRPO training. - for r, c in flat: - leaked = _answer_leaked(c['skills'], r['reference_answer']) - c['leaked'] = leaked - c['leak_reason'] = 'answer_verbatim' if leaked else '' - c['leak_source'] = 'deterministic' - - # with-skill greedy pass (T=0, M=1); reward = correct, absolute (group mean is baseline). - scored_inputs = flat - if scored_inputs: - ws_out = _run_samples(base_sampler, - [build_skill_solve_prompt(r['problem'], c['skills']) for r, c in scored_inputs], - 1, args.max_tokens, base_dp, temperature=0.0) - for (r, c), seqs in zip(scored_inputs, ws_out): - c['rolls'] = [_parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll()] - c['with_pass'] = 1.0 if c['rolls'][0]['correct'] else 0.0 - c['reward'] = c['with_pass'] - if args.format_in_reward: # unparseable candidates score 0 and still join the group - for r in hard: - for c in r['_cands']: - if c['reward'] is None: - c['reward'] = 0.0 - - _assign_advantages(hard, args) - return ([_full_record(r, ci) for r in chunk], _chunk_summary(chunk, ci, args), - _group_records(chunk, args)) - - -def _roll(x: Dict[str, Any]) -> Dict[str, Any]: - return {k: x[k] for k in ('pred', 'correct', 'terminated', 'passed', - 'stop_reason', 'gen_tokens', 'text')} - - -def _is_trainable(c: Dict[str, Any], args: argparse.Namespace) -> bool: - """Reaches the GRPO update iff advantage is non-zero. Leak flags are audit-only.""" - adv_nz = abs(c.get('advantage') or 0.0) > 1e-9 - if args.format_in_reward: - return adv_nz - return c.get('with_pass') is not None and adv_nz - - -def _full_record(r: Dict[str, Any], ci: int) -> Dict[str, Any]: - """Complete per-problem trace: init attempt, baseline, and all candidates.""" - init = r['_init'][0] - return { - 'record_type': 'problem', 'chunk': ci, 'problem': r['problem'], - 'reference_answer': r['reference_answer'], 'level': r.get('level', ''), - 'failed_first_try': r['_failed'], - 'init_attempt': {'text': init['text'], 'pred': init['pred'], 'correct': init['correct'], - 'terminated': init['terminated'], 'stop_reason': init['stop_reason'], - 'gen_tokens': init['gen_tokens']}, - 'baseline_pass': r['_baseline_pass'], 'is_hard': r['_hard'], - 'view': r.get('_view', ''), 'rubric_diag': r.get('_rubric_diag', ''), - # xproblem audit: which neighbour's rubric this problem borrowed, and how similar. - 'rubric_src': r.get('_rubric_src', ''), 'neighbor_sim': r.get('_neighbor_sim'), - 'baseline_rolls': [_roll(x) for x in r['_baseline_rolls']], - 'candidates': [{ - 'skills': c['skills'], 'response': c['response'], 'parseable': c['parseable'], - 'skillgen_stop': c.get('skillgen_stop'), 'skillgen_tokens': c.get('skillgen_tokens'), - 'leaked': c['leaked'], 'leak_reason': c['leak_reason'], 'leak_source': c['leak_source'], - 'with_pass': c['with_pass'], 'reward': c.get('reward'), 'advantage': c.get('advantage'), - 'grpo_adv': c.get('grpo_adv'), 'kept': c.get('kept'), - 'rolls': [_roll(x) for x in c['rolls']], - } for c in r['_cands']], - } - - -def _view_stats(problems: List[Dict[str, Any]], view: str) -> Dict[str, Any]: - pv = [r for r in problems if r.get('_view') == view] - cands = [c for r in pv for c in r['_cands'] if c['parseable']] - clean = [c for c in cands if c['leaked'] is False] - adopted = sum(1 for r in pv - if any(abs(c.get('advantage') or 0.0) > 1e-9 for c in r['_cands'])) - return {'n': len(pv), 'n_candidates_parseable': len(cands), 'n_clean': len(clean), - 'n_adopted_problems': adopted, 'adoption_rate': (adopted / len(pv)) if pv else 0.0} - - -def _mean(xs: List[float]) -> float: - return sum(xs) / len(xs) if xs else 0.0 - - -def _std(xs: List[float]) -> float: - if len(xs) < 2: - return 0.0 - m = _mean(xs) - return (sum((x - m) ** 2 for x in xs) / len(xs)) ** 0.5 - - -def _signal_stats(problems: List[Dict[str, Any]]) -> Dict[str, float]: - """The heart of 'is there a learning signal': per problem, the scored candidates form a - GRPO group. A group with zero reward variance (all skills solve, or none do -- the - hard-problem dead zone) gives NO gradient. Tracks that fraction plus reward level and - within-group variance so a collapse (all-0 or all-1) is visible immediately.""" - group_vars, all_rewards, zero_grad, groups = [], [], 0, 0 - for r in problems: - rewards = [c['reward'] for c in r['_cands'] if c.get('reward') is not None] - if len(rewards) < 2: - continue - groups += 1 - all_rewards.extend(rewards) - v = _std(rewards) - group_vars.append(v) - if v < 1e-9: # every skill got the same reward -> GRPO skips this problem - zero_grad += 1 - return {'n_groups': groups, 'zero_grad_frac': (zero_grad / groups) if groups else 0.0, - 'reward_mean': _mean(all_rewards), 'reward_std': _std(all_rewards), - 'group_reward_std_mean': _mean(group_vars)} - - -def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespace) -> Dict[str, Any]: - all_cands = [c for r in chunk for c in r['_cands']] - cands = [c for c in all_cands if c['parseable']] - scored = [c for c in cands if c['with_pass'] is not None] - clean = [c for c in cands if c['leaked'] is False] - ws_rolls = [x for c in scored for x in c['rolls']] - # viewa-dropped problems generate no candidates; keep acc/* on the generated subset - # so the with-skill/lift trend stays comparable across view_b_frac settings. - gen_probs = [r for r in chunk if r['_cands']] - base_acc = _mean([1.0 if r['_baseline_pass'] else 0.0 for r in gen_probs]) - ws_acc = _mean([1.0 if any(c.get('reward') for c in r['_cands']) else 0.0 for r in gen_probs]) - cand_pass_parseable = _mean([c['with_pass'] for c in scored]) - cand_pass_all = _mean([1.0 if c.get('reward') else 0.0 for c in all_cands]) - # base failure taxonomy (you asked whether skills fail because the base loops out of length) - classes = [_baseline_class(r) for r in chunk] - n_fail = sum(1 for c in classes if c != 'success') - skill_tokens = [c.get('skillgen_tokens') or 0 for c in cands] # skill-gen response length - trunc = sum(1 for r in chunk for c in r['_cands'] - for x in c['rolls'] if x['stop_reason'] == 'length') - rubric_answer_leaks = sum( - 1 for r in chunk - if r.get('_view') == 'A' and _rubric_has_fail(r.get('_rubric_diag')) and _rubric_answer_leaked(r)) - return { - 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), - 'n_failed_first_try': sum(1 for r in chunk if r['_failed']), - 'n_generated': len(all_cands), 'n_candidates_parseable': len(cands), - 'n_unparseable': len(all_cands) - len(cands), - 'parse_rate': (len(cands) / len(all_cands)) if all_cands else 0.0, - 'n_leaked': sum(1 for c in cands if c['leaked']), 'n_clean': len(clean), - 'leak_rate': (sum(1 for c in cands if c['leaked']) / len(cands)) if cands else 0.0, - 'n_reward_pos': sum(1 for c in scored if c['reward']), - 'n_rubric_answer_leaked': rubric_answer_leaks, - 'n_viewa_dropped': sum(1 for r in chunk if _viewa_dropped(r, args)), - 'n_train_samples': sum(1 for c in all_cands if _is_trainable(c, args)), - 'signal': _signal_stats(chunk), - 'fail_loop_frac': (sum(1 for c in classes if c == 'fail_loop') / n_fail) if n_fail else 0.0, - 'fail_wrong_frac': (sum(1 for c in classes if c == 'fail_wrong') / n_fail) if n_fail else 0.0, - 'skill_tokens_mean': _mean(skill_tokens), - 'withskill_trunc_frac': (trunc / len(ws_rolls)) if ws_rolls else 0.0, - 'avg_baseline_pass': base_acc, 'avg_withskill_pass': ws_acc, - 'avg_lift': ws_acc - base_acc, - 'candidate_withskill_pass_parseable': cand_pass_parseable, - 'candidate_withskill_pass_all': cand_pass_all, - 'termination_rate_withskill': _mean([1.0 if x['terminated'] else 0.0 for x in ws_rolls]), - 'view_A': _view_stats(chunk, 'A'), 'view_B': _view_stats(chunk, 'B'), - **_xproblem_stats(chunk, args), - } - - -def _xproblem_stats(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> Dict[str, Any]: - """Cross-problem pairing health: of the view-A problems, how many actually got a - neighbour's rubric (pair_rate) and how similar those neighbours were. Empty when off.""" - if not args.xproblem_rubric: - return {} - view_a = [r for r in chunk if r.get('_view') == 'A'] - paired = [r for r in view_a if r.get('_rubric_src')] - return {'xproblem': { - 'n_view_a': len(view_a), 'n_paired': len(paired), - 'pair_rate': (len(paired) / len(view_a)) if view_a else 0.0, - 'neighbor_sim_mean': _mean([r.get('_neighbor_sim', 0.0) for r in paired])}} - - -def _sft_record(r: Dict[str, Any], c: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: - """A view-A context-distillation sample: the OPEN-BOOK (rubric-conditioned) skill ``c`` - is the target, but the prompt is rebuilt as query-only (view B, no rubric) so the model - learns to produce it CLOSED-BOOK. Rides the same GRPOLoss with a positive constant - advantage (``--sft-weight``); single-step (old_logps=None) this reduces to - ``-w*logp + beta*KL`` -- a KL-anchored weighted cross-entropy toward the skill.""" - return { - 'data_id': r.get('data_id'), - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'view': 'B', 'diagnosis': '', 'rubric_src': '', 'orig_view': 'A', 'sft': True, - 'response': c['response'], 'skills': c['skills'], - 'skillgen_stop': c.get('skillgen_stop'), - 'advantage': float(args.sft_weight), 'grpo_adv': 0.0, 'kept': True, - 'reward': c['reward'], 'with_pass': c['with_pass']} - - -def _rubric_answer_leaked(r: Dict[str, Any]) -> bool: - """Hard safety gate for external API/rubric diagnosis. Unlike answer-bearing skills - generated by the policy itself, a rubric that contains the target final answer is an - external teacher leak and must not be distilled into view B.""" - return _answer_leaked(r.get('_rubric_diag', ''), r.get('reference_answer', '')) - - -def _viewa_dropped(r: Dict[str, Any], args: argparse.Namespace) -> bool: - """Under --viewa-sft a view-A problem trains ONLY through the SFT path (rubric with - [FAIL], answer-free). No-FAIL and rubric-leaked problems produce no training record - at all (no GRPO backflow: those prompts are query-only and would muddy the pure - view-B GRPO pool), so their skill-gen + executor scoring is skipped entirely.""" - return (bool(args.viewa_sft) and r.get('_view') == 'A' - and (not _rubric_has_fail(r.get('_rubric_diag')) - or _rubric_answer_leaked(r))) - - -def _group_records(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> List[Dict[str, Any]]: - """Training records. Under --viewa-sft, view A is SFT-only: a problem whose rubric - localised a failure (has ``[FAIL]``, answer-free) contributes ONE context-distillation - SFT sample (best parseable open-book skill -- preferring an executor-verified pass, - else any parseable candidate -- rebuilt query-only); no-FAIL / rubric-leaked view-A - problems are dropped (``_viewa_dropped``, generation already skipped). GRPO candidates - come from view B only, keeping that pool purely closed-book. Prompts are rebuilt from - the stored view/diagnosis by ``_skillgen_messages``.""" - out = [] - for r in chunk: - if not r['_hard']: - continue - if args.viewa_sft and r.get('_view') == 'A': - if _viewa_dropped(r, args): - continue - best = _best_sft_candidate(r, args) - if best is not None: - out.append(_sft_record(r, best, args)) - continue - for c in r['_cands']: - if _is_trainable(c, args): - out.append({ - 'data_id': r.get('data_id'), - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'view': r.get('_view', 'A'), 'diagnosis': r.get('_rubric_diag', ''), - 'rubric_src': r.get('_rubric_src', ''), 'sft': False, - 'response': c['response'], 'skills': c['skills'], - 'skillgen_stop': c.get('skillgen_stop'), - 'advantage': c['advantage'], 'grpo_adv': c['grpo_adv'], 'kept': c['kept'], - 'reward': c['reward'], 'with_pass': c['with_pass']}) - return out - - -# =========================================================================== -# Block G -- online GRPO training -# =========================================================================== -def _is_num(v: Any) -> bool: - try: - float(v) - return True - except (TypeError, ValueError): - return False - - -def _train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: - """Training sample = the exact skill-gen prompt (rebuilt by ``_skillgen_messages`` so - train/inference match) + the generated structured guidance response. ``key_rounds`` - selects the final assistant turn; Template masks the prompt and trains the whole - response (the key-round prefix already excludes the prompt-provided <think>).""" - msgs = _skillgen_messages( - rec['problem'], rec.get('view', 'A'), rec.get('diagnosis', ''), rec.get('rubric_src', '')) - return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], - 'user_data': pack_user_data({'key_rounds': [len(msgs)]})} - - -def _train_chunk(skill_model, ref_model, ckpt: CheckpointEngineManager, samples: List[Dict[str, Any]], - args: argparse.Namespace) -> Dict[str, Any]: - """On-policy GRPO update over one chunk, then sync weights. Micro-batches of - ``sft_batch_size`` accumulate gradients; ONE optimizer step is taken per PPO - mini-batch of ``ppo_mini_batch_size`` samples (0 -> a single step over the whole - chunk, the original behaviour). A frozen reference model provides ref_logps for the - SEAM-style KL penalty. - - Multi-step correctness: with more than one step over the SAME rollout, later - mini-batches see an already-updated policy, so we FREEZE the sampling-policy - ``old_logps`` (recomputed once, before any step) and let GRPOLoss form the PPO ratio - against them. A single step keeps ``old_logps=None`` (ratio==1, pure on-policy). - The batch is padded to a multiple of ``sft_batch_size`` with advantage-0 copies that - contribute no policy gradient. View-A context-distillation samples ride the same loss - with a positive constant advantage (``--sft-weight``); single-step (old_logps=None) - that is ``-w*logp + beta*KL`` (KL-anchored weighted cross-entropy).""" - trajs = [_train_trajectory(rec) for rec in samples] - advs = [float(rec['advantage']) for rec in samples] - rem = (-len(trajs)) % args.sft_batch_size - if rem: - trajs += [trajs[-1]] * rem - advs += [0.0] * rem - - n, sft = len(trajs), args.sft_batch_size - mini = args.ppo_mini_batch_size if args.ppo_mini_batch_size > 0 else n - mini = max(sft, (mini // sft) * sft) # align to a whole number of micro-batches - multi_step = mini < n - - # Freeze ref_logps (frozen model) and -- only when we take multiple steps -- the - # sampling-policy old_logps, BOTH before any optimizer step touches the weights. With - # a single step old_logps stays None so GRPOLoss uses ratio==1 (pure on-policy). - micro_ref, micro_old = [], [] - for i in range(0, n, sft): - mb = trajs[i:i + sft] - micro_ref.append(ref_model.forward_only(inputs=mb).get('logps')) - micro_old.append(skill_model.forward_only(inputs=mb).get('logps') if multi_step else None) - - micro, n_steps = 0, 0 - for ms in range(0, n, mini): - for i in range(ms, min(ms + mini, n), sft): - k = i // sft - skill_model.forward_backward(inputs=trajs[i:i + sft], - advantages=advs[i:i + sft], - old_logps=micro_old[k], - ref_logps=micro_ref[k]) - micro += 1 - skill_model.clip_grad_and_step() - n_steps += 1 - ckpt.sync_weights(merge_and_sync=True) - metric = skill_model.calculate_metric(is_training=True) - n_sft = sum(1 for s in samples if s.get('sft')) - return {'n_samples': len(samples), 'n_sft': n_sft, 'n_grpo': len(samples) - n_sft, - 'n_steps': n_steps, 'n_micro_batches': micro, - 'metric': {k: (float(v) if _is_num(v) else v) for k, v in (metric or {}).items()}} - - -# =========================================================================== -# Block H -- fixed-holdout eval + metric formatting -# =========================================================================== -def run_greedy_eval(base_sampler, skill_sampler, eval_records: List[Dict[str, Any]], - ci: int, rounds: int, base_dp: int, skill_dp: int, args: argparse.Namespace, - base_cache: DiskCache - ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], Dict[str, float]]: - """SEAM ``val-core/math/acc/mean@1`` on the fixed holdout: ONE greedy skill (T=0) per - problem into ONE greedy base solve (T=0). Eval ALWAYS uses view B (query-only, the - deployment form: no rubric, since rubric needs an online teacher unavailable at deploy); - no leak filter (acc scores correctness alone). Baseline reuses the disk cache.""" - baseline_rollout(base_sampler, eval_records, base_dp, args, base_cache) - for r in eval_records: - r['_view'], r['_rubric_diag'] = 'B', '' - sg_out = _run_samples(skill_sampler, [_view_prompt(r, args) for r in eval_records], - 1, args.skill_max_tokens, skill_dp, temperature=0.0) - skills = [] - for seqs in sg_out: - if not seqs: - skills.append(('', '')) - continue - sresp = _clean_text(getattr(seqs[0], 'decoded', '') or '') - skills.append((_extract_skill(sresp) or '', sresp)) - ws_out = _run_samples(base_sampler, - [build_skill_solve_prompt(r['problem'], sk) for r, (sk, _) in zip(eval_records, skills)], - 1, args.max_tokens, base_dp, temperature=0.0) - recs = [] - for r, (sk, sresp), seqs in zip(eval_records, skills, ws_out): - roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - recs.append({ - 'record_type': 'eval_problem', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, - 'data_id': r.get('data_id'), - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'view': r['_view'], 'rubric_diag': r.get('_rubric_diag', ''), - 'baseline_pass': r['_baseline_pass'], 'skill': sk, 'skill_parseable': bool(sk), - 'skill_response': sresp, 'withskill_pred': roll['pred'], - 'withskill_correct': roll['correct'], 'withskill_terminated': roll['terminated'], - 'withskill_stop_reason': roll['stop_reason'], 'withskill_text': roll['text'], - }) - acc = lambda rs: sum(1 for x in rs if x['withskill_correct']) / len(rs) if rs else 0.0 - ws = acc(recs) # all view B (deployment form) - base = sum(x['baseline_pass'] for x in recs) / len(recs) if recs else 0.0 - fmt = (sum(1 for x in recs if x['skill_parseable']) / len(recs)) if recs else 0.0 - term = (sum(1 for x in recs if x['withskill_terminated']) / len(recs)) if recs else 0.0 - summary = {'record_type': 'eval_summary', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, - 'n': len(recs), 'view': 'B', 'acc_mean1': ws, - 'baseline_acc_mean1': base, 'lift_mean1': ws - base, - 'format_mean1': fmt, 'term_mean1': term} - metrics = {'core/math/acc/mean@1': ws, 'core/math/baseline_acc/mean@1': base, - 'core/math/lift/mean@1': ws - base, 'core/math/format/mean@1': fmt, - 'core/math/term/mean@1': term} - return recs, summary, metrics - - -def _trend_line(hist: List[Dict[str, float]], window: int, rounds_done: int) -> Optional[str]: - """Contrast the FIRST vs the most recent ``window`` chunks -- if RFT works, adoption - and lift on recent (fresh) chunks exceed the early baseline.""" - if len(hist) < 2 * window: - return None - base, rec = hist[:window], hist[-window:] - m = lambda xs, k: sum(h[k] for h in xs) / len(xs) - return (f'[trend] first {window} vs last {window} | ' - f'adopt A {m(base,"aA"):.2f}->{m(rec,"aA"):.2f} B {m(base,"aB"):.2f}->{m(rec,"aB"):.2f} | ' - f'lift {m(base,"lift"):+.3f}->{m(rec,"lift"):+.3f} | ' - f'0grad {m(base,"zero_grad"):.2f}->{m(rec,"zero_grad"):.2f} | ' - f'pos/chunk {m(base,"pos"):.1f}->{m(rec,"pos"):.1f} | rounds={rounds_done}') - - -def _swan_metrics(summary: Dict[str, Any], log: Optional[Dict[str, Any]]) -> Dict[str, float]: - """Flat swanlab dict. ``signal/*`` is the primary health group (is GRPO getting a - gradient at all); ``acc/*`` is the effect; the rest diagnose why. acc/adopt/skill are - only logged when the chunk produced a scored group, so idle chunks don't dip charts.""" - sig = summary['signal'] - d: Dict[str, float] = { - # --- signal: the FIRST thing to watch (no variance -> no learning) --- - 'signal/zero_grad_frac': sig['zero_grad_frac'], 'signal/n_groups': sig['n_groups'], - 'signal/reward_mean': sig['reward_mean'], 'signal/reward_std': sig['reward_std'], - 'signal/group_reward_std_mean': sig['group_reward_std_mean'], - 'signal/n_train_samples': summary['n_train_samples'], - 'signal/n_reward_pos': summary['n_reward_pos'], - # --- skill format / leak health --- - 'skill/parse_rate': summary['parse_rate'], 'skill/tokens_mean': summary['skill_tokens_mean'], - 'leak/rate': summary['leak_rate'], 'leak/n': summary['n_leaked'], - # --- base failure taxonomy (loop-out-of-length vs plain wrong) --- - 'fail/loop_frac': summary['fail_loop_frac'], 'fail/wrong_frac': summary['fail_wrong_frac'], - 'fail/frac_first_try': (summary['n_failed_first_try'] / summary['n']) if summary['n'] else 0.0, - # --- view-A routing: dropped = no-FAIL rubric or rubric leak (SFT-only view A) --- - 'viewa/dropped_frac': (summary['n_viewa_dropped'] / summary['view_A']['n'] - if summary['view_A']['n'] else 0.0), - } - bal = summary.get('balance') or {} - if bal.get('enabled'): - d.update({'balance/n_drawn': bal['n_drawn'], 'balance/n_baseline_fresh': bal['n_baseline_fresh'], - 'balance/selected_success_frac': bal['selected_success_frac']}) - xp = summary.get('xproblem') or {} - if xp: - d.update({'xproblem/pair_rate': xp['pair_rate'], 'xproblem/neighbor_sim_mean': xp['neighbor_sim_mean']}) - if sig['n_groups'] > 0: - d.update({'acc/baseline_pass': summary['avg_baseline_pass'], - 'acc/withskill_pass': summary['avg_withskill_pass'], 'acc/lift': summary['avg_lift'], - 'candidate/withskill_pass_parseable': summary['candidate_withskill_pass_parseable'], - 'candidate/withskill_pass_all': summary['candidate_withskill_pass_all'], - 'adopt/A': summary['view_A']['adoption_rate'], - 'adopt/B': summary['view_B']['adoption_rate'], - 'term/withskill': summary['termination_rate_withskill'], - 'term/withskill_trunc_frac': summary['withskill_trunc_frac']}) - if log: - d['train/n_steps'] = log['n_steps'] - d['train/n_micro_batches'] = log['n_micro_batches'] - for k, v in (log.get('metric') or {}).items(): - if not _is_num(v): - continue - if k.startswith('learning rate'): - if 'group 1' in k: - d['train/lr'] = float(v) - else: - d[f'train/{k.replace(" ", "_")}'] = float(v) - return d - - -def _view_a_rubric_leak_metrics(chunk: List[Dict[str, Any]], - pool_answers: Optional[Dict[str, str]] = None) -> Dict[str, float]: - """Swanlab-only audit for answer leakage in view-A rubric text. This never changes - rewards, advantages, filtering, or training records.""" - view_a = [r for r in chunk if r.get('_view') == 'A'] - with_diag = [r for r in view_a if r.get('_rubric_diag')] - target_leaks = sum(1 for r in with_diag - if _answer_leaked(r.get('_rubric_diag', ''), r.get('reference_answer', ''))) - source_leaks = 0 - pool_answers = pool_answers or {} - for r in with_diag: - src = r.get('_rubric_src') - src_ref = pool_answers.get(src, '') if src else r.get('reference_answer', '') - if _answer_leaked(r.get('_rubric_diag', ''), src_ref): - source_leaks += 1 - n = len(with_diag) - return { - 'rubric_leak/n_view_a': float(len(view_a)), - 'rubric_leak/n_checked': float(n), - 'rubric_leak/target_answer_n': float(target_leaks), - 'rubric_leak/target_answer_rate': (target_leaks / n) if n else 0.0, - 'rubric_leak/source_answer_n': float(source_leaks), - 'rubric_leak/source_answer_rate': (source_leaks / n) if n else 0.0, - } - - -# =========================================================================== -# Block F -- components, args, main -# =========================================================================== -def init_components(args: argparse.Namespace): - """Default 8-GPU layout: rank 0 trains the actor, rank 1 hosts a frozen ref model, - 2-3 skill_sampler (synced), 4-7 base_sampler (frozen). Returns - (skill_model, ref_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp).""" - r0, r1 = TRAIN_GPUS, TRAIN_GPUS + REF_GPUS - r2, r3 = r1 + SKILL_SAMPLER_GPUS, NUM_GPUS - twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, lazy_collect=False, groups=[ - DeviceGroup(name='train', ranks=list(range(0, r0)), device_type='GPU'), - DeviceGroup(name='ref', ranks=list(range(r0, r1)), device_type='GPU'), - DeviceGroup(name='skill_sampler', ranks=list(range(r1, r2)), device_type='GPU'), - DeviceGroup(name='base_sampler', ranks=list(range(r2, r3)), device_type='GPU')]) - - train_mesh = DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_DP, fsdp_size=TRAIN_FSDP) - skill_model = TransformersModel(model_id=MODEL_ID, device_mesh=train_mesh, remote_group='train', - ddp_config={'find_unused_parameters': False}) - skill_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) - skill_model.set_template(Template, model_id=MODEL_ID, enable_thinking=False, - max_length=args.max_model_len, truncation_strategy='delete') - skill_model.set_processor(InputProcessor, padding_free=False) - skill_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon, beta=args.kl_beta) - skill_model.set_optimizer('AdamW', lr=args.lr) - skill_model.set_lr_scheduler('CosineWarmupScheduler', num_warmup_steps=10, - num_training_steps=args.max_train_rounds) - - ref_mesh = DeviceMesh.from_sizes(world_size=REF_GPUS, dp_size=REF_DP, fsdp_size=REF_FSDP) - ref_model = TransformersModel(model_id=MODEL_ID, device_mesh=ref_mesh, remote_group='ref', - ddp_config={'find_unused_parameters': False}) - ref_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) - ref_model.set_template(Template, model_id=MODEL_ID, enable_thinking=True, - max_length=args.max_model_len, truncation_strategy='delete') - ref_model.set_processor(InputProcessor, padding_free=False) - ref_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) - - def _sampler(group, world, enable_thinking: bool = True): - s = vLLMSampler(model_id=MODEL_ID, - engine_args={'gpu_memory_utilization': GPU_MEM, - 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, - device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), - remote_group=group) - s.set_template(Template, model_id=MODEL_ID, enable_thinking=enable_thinking, max_length=args.max_model_len) - return s - - skill_sampler = _sampler('skill_sampler', SKILL_SAMPLER_GPUS, enable_thinking=False) - # base_sampler is shared by the main thread and the baseline-prefetch thread -> serialise. - base_sampler = _LockedSampler(_sampler('base_sampler', BASE_SAMPLER_GPUS, enable_thinking=True)) - ckpt = CheckpointEngineManager(model=skill_model, sampler=skill_sampler) - return skill_model, ref_model, skill_sampler, base_sampler, ckpt, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS - - -def _build_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--dataset', choices=('aops', 'math'), default='aops') - p.add_argument('--n', type=int, default=2000, help='Problems loaded into the draw pool.') - p.add_argument('--pool-offset', type=int, default=0, - help='Skip this many shuffled non-eval records before building the train pool; ' - 'useful to avoid cold-start SFT data ranges.') - p.add_argument('--exclude-data-ids', default='', - help='Comma-separated jsonl files whose data_id/problem keys are excluded ' - 'from train/eval selection, e.g. coldstart_sft.jsonl.') - p.add_argument('--seed', type=int, default=42) - p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) - p.add_argument('--eval-size', type=int, default=128, help='Fixed holdout size (0 disables).') - p.add_argument('--eval-every', type=int, default=5, help='Run holdout eval every N chunks.') - p.add_argument('--chunk-size', type=int, default=16) - p.add_argument('--balance', action=argparse.BooleanOptionalAction, default=True) - p.add_argument('--balance-success-frac', type=float, default=0.4, - help='Target fraction of the chunk the base solves (rest are base-fail).') - p.add_argument('--balance-loop-frac', type=float, default=0.5) - p.add_argument('--balance-max-draws-mult', type=int, default=8) - p.add_argument('--seam-parquet-dir', type=str, default='', - help='Read SEAM build_aops_dataset.py train.parquet/val.parquet directly, in ' - 'file order (problem<-extra_info.problem, answer<-reward_model.ground_truth). ' - 'val.parquet becomes the eval holdout. Bypasses load/--numeric-only/' - '--eval-size/internal shuffle so the input data matches a SEAM run exactly.') - p.add_argument('--n-skills', type=int, default=8) - p.add_argument('--view-b-frac', type=float, default=0.5) - p.add_argument('--xproblem-rubric', action=argparse.BooleanOptionalAction, default=False, - help='When enabled, view A uses a bag-of-words NEIGHBOUR problem rubric. ' - 'Default is off: each view-A problem uses its own baseline attempt, ' - 'while the API diagnosis prompt is constrained to be answer-free and ' - 'method-level only.') - p.add_argument('--skill-retries', type=int, default=2) - p.add_argument('--skill-gen-temperature', type=float, default=1.0) - p.add_argument('--skill-gen-top-p', type=float, default=1.0) - p.add_argument('--skill-gen-top-k', type=int, default=-1) - p.add_argument('--max-model-len', type=int, default=16384) - p.add_argument('--max-tokens', type=int, default=8192) - p.add_argument('--skill-max-tokens', type=int, default=8192) - p.add_argument('--rubric-workers', type=int, default=16) - p.add_argument('--rubric-retries', type=int, default=2, - help='Retry failed/timeout rubric diagnose calls this many times before ' - 'falling back to an empty diagnosis without caching the failure.') - p.add_argument('--sft-batch-size', type=int, default=8, - help='Driver micro-batch (gradient-accumulation unit); multiple of train dp.') - p.add_argument('--ppo-mini-batch-size', type=int, default=0, - help='Samples per optimizer step (SEAM-style PPO mini-batch). 0 = one step ' - 'over the whole chunk (on-policy, ratio==1). When >0 and smaller than ' - 'the trainable count, multiple steps are taken over the same rollout and ' - 'old_logps are frozen so the PPO ratio/clip stays valid. Rounded down to ' - 'a multiple of --sft-batch-size.') - p.add_argument('--grpo-epsilon', type=float, default=0.2) - p.add_argument('--adv-clip', type=float, default=3.0, - help='Symmetric clip for group-relative advantages; <=0 disables clipping.') - p.add_argument('--kl-beta', type=float, default=0.001, - help='SEAM-style reference KL coefficient for GRPOLoss.') - p.add_argument('--viewa-sft', action=argparse.BooleanOptionalAction, default=True, - help='Route view-A problems to online context distillation (SFT on the best ' - 'passing open-book skill, rebuilt as a query-only prompt) instead of GRPO. ' - 'View B stays GRPO; both share one optimizer step.') - p.add_argument('--sft-weight', type=float, default=0.5, - help='Advantage magnitude (lambda) for view-A SFT distillation samples in the ' - 'shared GRPOLoss. Single-step this is -lambda*logp + beta*KL (weighted CE). ' - 'Keep <~1 so one-signed SFT gradients do not swamp mean-0 GRPO advantages.') - p.add_argument('--sft-target-len', type=int, default=550, - help='Target skill length (chars) for view-A SFT distillation: among passing ' - 'candidates pick the one CLOSEST to this length (not the shortest). Set near ' - 'the empirical pass-rate peak (~500-600 here) so distillation neither shrinks ' - 'rollouts toward zero nor lets them grow unbounded.') - p.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True) - p.add_argument('--lr', type=float, default=6e-6) - p.add_argument('--max-train-rounds', type=int, default=1500) - p.add_argument('--save-rounds', type=int, default=200) - p.add_argument('--trend-every', type=int, default=10) - p.add_argument('--output-dir', default='./output/reflexion_skill') - p.add_argument('--cache-dir', default='', help='Baseline/rubric cache dir (default <output-dir>/cache).') - p.add_argument('--no-cache', action='store_true', help='Disable disk cache read/write.') - p.add_argument('--prefetch-baseline', action=argparse.BooleanOptionalAction, default=True, - help='Prefetch next chunk base baseline on a background thread (overlaps ' - 'with skill-gen; base_sampler is frozen so it never blocks the trainer).') - p.add_argument('--swanlab-project', default='twinkle') - p.add_argument('--swanlab-exp', default='') - args = p.parse_args() - if args.sft_batch_size % TRAIN_DP != 0: - raise ValueError(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple of train dp ({TRAIN_DP})') - if args.chunk_size < 1: - raise ValueError('--chunk-size must be >= 1') - args.balance_success_frac = max(0.0, min(1.0, args.balance_success_frac)) - return args - - -def _write(handle, row: Dict[str, Any]) -> None: - handle.write(json.dumps(row, ensure_ascii=False) + '\n') - - -def main() -> None: - args = _build_args() - records, eval_records, neighbor_map, pool_answers, data_stats = _load_records(args) - if len(records) < args.chunk_size: - raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded ({len(records)}); raise --n') - - os.makedirs(args.output_dir, exist_ok=True) - gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') - eval_path = os.path.join(args.output_dir, 'eval_records.jsonl') - data_path = os.path.join(args.output_dir, 'skill_dataset.jsonl') - train_log_path = os.path.join(args.output_dir, 'train_log.jsonl') - - if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY')): - sys.stderr.write('[rft] WARNING: no LLM backup env; view-A rubric check disabled ' - '(leak filter is deterministic, unaffected)\n') - - use_swan = swanlab is not None and os.environ.get('SWANLAB_MODE') != 'disabled' - if use_swan: - swanlab.init(project=args.swanlab_project, experiment_name=(args.swanlab_exp or None), - config={'model': MODEL_ID, 'dataset': args.dataset, 'n': len(records), - 'eval_n': len(eval_records), 'n_skills': args.n_skills, - 'view_b_frac': args.view_b_frac, 'balance': args.balance, - 'balance_success_frac': args.balance_success_frac, - 'skill_gen_temp': args.skill_gen_temperature, - 'grpo_epsilon': args.grpo_epsilon, 'kl_beta': args.kl_beta, - 'lr': args.lr}) - - skill_model, ref_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp = init_components(args) - checker = build_rubric_checker() - if checker is None: - sys.stderr.write('[rft] no LLM backup env -> view-A rubric process-check DISABLED\n') - - cache_dir = args.cache_dir or os.path.join(args.output_dir, 'cache') - os.makedirs(cache_dir, exist_ok=True) - use_cache = not args.no_cache - base_cache = DiskCache(os.path.join(cache_dir, 'baseline.jsonl'), use_cache) - eval_base_cache = DiskCache(os.path.join(cache_dir, 'eval_baseline.jsonl'), use_cache) - rubric_cache = DiskCache(os.path.join(cache_dir, 'rubric.jsonl'), use_cache) - if args.xproblem_rubric: - sys.stderr.write(f'[rft] xproblem-rubric ON: {len(neighbor_map)} bag-of-words pairs\n') - - cfg = {'record_type': 'config', 'model': MODEL_ID, 'dataset': args.dataset, - 'n': len(records), 'eval_n': len(eval_records), 'seed': args.seed, - 'pool_offset': args.pool_offset, 'exclude_data_ids': args.exclude_data_ids, - 'excluded_records': data_stats.get('excluded_records', 0), - 'numeric_only': args.numeric_only, 'raw_loaded': data_stats['raw_loaded'], - 'numeric_dropped': data_stats['numeric_dropped'], 'eval_every': args.eval_every, - 'n_skills': args.n_skills, 'view_b_frac': args.view_b_frac, - 'skill_retries': args.skill_retries, 'balance': args.balance, - 'balance_success_frac': args.balance_success_frac, - 'skill_gen_temp': args.skill_gen_temperature, 'reward': 'greedy_binary(correct)', - 'advantage': 'group_relative', 'format_in_reward': args.format_in_reward, 'cache': use_cache, - 'rubric_check': f'fixed_math_{len(_MATH_RUBRIC)}crit+term(viewA)' if checker else 'disabled', - 'xproblem_rubric': args.xproblem_rubric, - 'viewa_sft': args.viewa_sft, 'sft_weight': args.sft_weight, - 'sft_target_len': args.sft_target_len, - 'adv_clip': args.adv_clip, - 'grpo_epsilon': args.grpo_epsilon, 'kl_beta': args.kl_beta, 'lr': args.lr, - 'train_gpus': TRAIN_GPUS, 'ref_gpus': REF_GPUS, 'ref_fsdp': REF_FSDP, - 'train_fsdp': TRAIN_FSDP, 'train_dp': TRAIN_DP, - 'skill_sampler_gpus': SKILL_SAMPLER_GPUS, 'base_sampler_gpus': BASE_SAMPLER_GPUS, - 'max_train_rounds': args.max_train_rounds, 'started': int(time.time())} - sys.stderr.write(f'[rft] raw={data_stats["raw_loaded"]} numeric_drop={data_stats["numeric_dropped"]} ' - f'train={len(records)} eval={len(eval_records)} {args.dataset}; ' - f'train_gpus={TRAIN_GPUS} ref_gpus={REF_GPUS} train_fsdp={TRAIN_FSDP} ' - f'train_dp={TRAIN_DP} skill_dp={skill_dp} base_dp={base_dp}\n') - - hist: List[Dict[str, float]] = [] - rounds = 0 - pool = ProblemPool(records, args.seed) - with open(gen_path, 'w', encoding='utf-8') as gen_f, \ - open(eval_path, 'w', encoding='utf-8') as eval_f, \ - open(data_path, 'w', encoding='utf-8') as data_f, \ - open(train_log_path, 'w', encoding='utf-8') as tlog: - for f in (gen_f, eval_f, data_f, tlog): - _write(f, cfg) - gstep = 0 - # base baseline (frozen sampler, disk-cached, no weight-sync) is prefetched on a - # background thread while the current chunk generates: the skill-gen phase uses - # skill_sampler, so prefetching the NEXT chunk's baseline on base_sampler overlaps - # it and fills base_cache so the next draw_chunk hits it. base_sampler is wrapped in - # _LockedSampler, so the prefetch and this chunk's with-skill scoring never issue a - # base .sample() concurrently. It never touches the trainer or on-policy generation. - prefetch_pool = ThreadPoolExecutor(max_workers=1) if args.prefetch_baseline else None - pending: Optional[Any] = None - - def _prefetch(peeked: List[Dict[str, Any]]) -> None: - if peeked: - baseline_rollout(base_sampler, peeked, base_dp, args, base_cache) - - # Baseline (round 0) eval BEFORE any training: measures the untrained skill model on - # the fixed holdout so every later eval has a step-0 reference point on the same axis. - if eval_records: - eval_recs, eval_summary, eval_metrics = run_greedy_eval( - base_sampler, skill_sampler, eval_records, -1, rounds, base_dp, skill_dp, - args, eval_base_cache) - for rec in eval_recs: - _write(eval_f, rec) - _write(eval_f, eval_summary) - eval_f.flush() - if use_swan: - swanlab.log({f'eval/{k}': v for k, v in eval_metrics.items()}, step=0) - sys.stderr.write( - f'[eval] g-1 (init): n={eval_summary["n"]} viewB mean@1 ' - f'acc={eval_summary["baseline_acc_mean1"]:.3f}->{eval_summary["acc_mean1"]:.3f} ' - f'lift={eval_summary["lift_mean1"]:+.3f} ' - f'fmt={eval_summary["format_mean1"]:.2f} rounds={rounds}\n') - - # Each chunk is drawn fresh + RE-GENERATED with the current policy (on-policy); - # with --balance, draw_chunk keeps drawing until the base fail:success mix hits target. - while rounds < args.max_train_rounds: - if pending is not None: - pending.result() # finish last round's prefetch before drawing (cache-warm) - pending = None - chunk, balance = draw_chunk(pool, base_sampler, base_dp, args, base_cache) - if prefetch_pool is not None: - peeked = pool.peek(int(args.chunk_size * (args.balance_max_draws_mult if args.balance else 1))) - pending = prefetch_pool.submit(_prefetch, peeked) - full, summary, groups = process_chunk( - base_sampler, skill_sampler, chunk, gstep, base_dp, skill_dp, - args, checker, rubric_cache, base_cache, neighbor_map, pool_answers) - summary['balance'] = balance - - log = None - if groups: - log = _train_chunk(skill_model, ref_model, ckpt, groups, args) - rounds += 1 - log.update({'record_type': 'train_round', 'round': rounds, 'chunk': gstep, - 'epoch': pool.epoch, 'ts': int(time.time())}) - _write(tlog, log) - tlog.flush() - if rounds % args.save_rounds == 0: - skill_model.save(f'skill-rft-{rounds}', output_dir=args.output_dir) - - summary['rounds_done'], summary['epoch'] = rounds, pool.epoch - for rec in full: - _write(gen_f, rec) - _write(gen_f, summary) - gen_f.flush() - for v in groups: - _write(data_f, v) - data_f.flush() - - sa, sb, sig = summary['view_A'], summary['view_B'], summary['signal'] - hist.append({'aA': sa['adoption_rate'], 'aB': sb['adoption_rate'], - 'lift': summary['avg_lift'], 'pos': summary['n_reward_pos'], - 'zero_grad': sig['zero_grad_frac']}) - bal_str = (f'bal {balance["selected_fail"]}f/{balance["selected_success"]}s ' - f'(drew {balance["n_drawn"]}/fresh {balance["n_baseline_fresh"]}' - + ('!' if balance.get('budget_hit') else '') + ') ') if balance.get('enabled') else '' - xp = summary.get('xproblem') - xp_str = f'pair={xp["pair_rate"]:.2f}@{xp["neighbor_sim_mean"]:.2f} ' if xp else '' - tr_str = (f'train={log["n_grpo"]}g/{log["n_sft"]}s ' if log - else f'train={summary["n_train_samples"]} ') - sys.stderr.write( - f'[gen] e{pool.epoch} g{gstep}: {bal_str}n={summary["n"]} ' - f'clean={summary["n_clean"]} leak={summary["leak_rate"]:.2f} {tr_str}' - f'0grad={sig["zero_grad_frac"]:.2f} R={sig["reward_mean"]:.2f}±{sig["reward_std"]:.2f} ' - f'acc={summary["avg_baseline_pass"]:.2f}->{summary["avg_withskill_pass"]:.2f} ' - f'lift={summary["avg_lift"]:+.3f} {xp_str}' - f'A[{sa["n"]} {sa["adoption_rate"]:.2f}] B[{sb["n"]} {sb["adoption_rate"]:.2f}] ' - f'rounds={rounds}\n') - if use_swan: - swan_metrics = _swan_metrics(summary, log) - swan_metrics.update(_view_a_rubric_leak_metrics(chunk, pool_answers)) - swanlab.log(swan_metrics, step=gstep) - - if eval_records and (gstep + 1) % args.eval_every == 0: - eval_recs, eval_summary, eval_metrics = run_greedy_eval( - base_sampler, skill_sampler, eval_records, gstep, rounds, base_dp, skill_dp, - args, eval_base_cache) - for rec in eval_recs: - _write(eval_f, rec) - _write(eval_f, eval_summary) - eval_f.flush() - if use_swan: - swanlab.log({f'eval/{k}': v for k, v in eval_metrics.items()}, step=gstep) - sys.stderr.write( - f'[eval] g{gstep}: n={eval_summary["n"]} viewB mean@1 ' - f'acc={eval_summary["baseline_acc_mean1"]:.3f}->{eval_summary["acc_mean1"]:.3f} ' - f'lift={eval_summary["lift_mean1"]:+.3f} ' - f'fmt={eval_summary["format_mean1"]:.2f} rounds={rounds}\n') - - if (gstep + 1) % args.trend_every == 0: - tl = _trend_line(hist, args.trend_every, rounds) - if tl: - sys.stderr.write(tl + '\n') - gstep += 1 - - if prefetch_pool is not None: - if pending is not None: - pending.result() - prefetch_pool.shutdown(wait=True) - base_cache.close() - eval_base_cache.close() - rubric_cache.close() - skill_model.save('skill-rft-final', output_dir=args.output_dir) - sys.stderr.write(f'[rft] done: {rounds} rounds over {gstep} chunks / {pool.epoch} epochs -> {data_path}\n') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/legacy/train_reflexion_skill_old.sh b/cookbook/exp/legacy/train_reflexion_skill_old.sh deleted file mode 100755 index 093bf3521..000000000 --- a/cookbook/exp/legacy/train_reflexion_skill_old.sh +++ /dev/null @@ -1,57 +0,0 @@ -#!/bin/bash -# Online GRPO RFT for the reflexion skill generator (unified, self-contained, cached). -# GPUs: 8 — default high-memory layout uses rank 0 for actor training, rank 1 for -# the frozen ref model, ranks 2-3 for skill sampler (synced), and ranks 4-7 for -# base sampler (frozen). Override TRAIN_GPUS / REF_GPUS / SKILL_SAMPLER_GPUS / -# BASE_SAMPLER_GPUS for other layouts. Per chunk: base greedy -# solve -> rubric process-check (view A) -> -# skill-gen (thinking ON, N candidates) -> deterministic leak filter -> with-skill greedy -# pass -> group-relative advantage -> ONE on-policy GRPO step -> sync weights. -# -# Baseline rollouts + rubric diagnoses are disk-cached (output-dir/cache/*.jsonl), so a -# restart skips re-sampling them; skill-gen is on-policy and never cached. The next chunk's -# baseline is prefetched on a background thread (overlaps skill-gen; base sampler is frozen). -# -# The view-A rubric process-check uses the backup teacher API (set LLM_BACKUP_*). Without -# it the run still works: view A degrades to query-only and the leak filter stays -# deterministic (no teacher needed). - -set -euo pipefail - -export GEN_MODEL_ID=${GEN_MODEL_ID:-Qwen/Qwen3-4B} -export GEN_GPU_MEM=${GEN_GPU_MEM:-0.8} -# Datasets are pulled from ModelScope via twinkle.Dataset (ms://AI-MO/aops or -# ms://modelscope/competition_math); override AOPS_DATASET_ID / MATH_DATASET_ID to change. -# Teacher API for the view-A rubric process-check (optional; leak filter is deterministic). -export LLM_BACKUP_API_KEY=${LLM_BACKUP_API_KEY:-} -export LLM_BACKUP_BASE_URL=${LLM_BACKUP_BASE_URL:-https://dashscope.aliyuncs.com/compatible-mode/v1} -export LLM_BACKUP_MODEL=${LLM_BACKUP_MODEL:-qwen3.7-max} - -python cookbook/exp/embedding/train_reflexion_skill.py \ - --dataset aops \ - --n 5000 \ - --seam-parquet-dir /root/data/seam \ - --numeric-only \ - --chunk-size 32 \ - --n-skills 16 \ - --view-b-frac 0.5 \ - --skill-retries 2 \ - --no-balance \ - --max-tokens 8192 \ - --skill-max-tokens 4096 \ - --max-model-len 16384 \ - --eval-size 128 \ - --eval-every 5 \ - --sft-batch-size 8 \ - --ppo-mini-batch-size 0 \ - --grpo-epsilon 0.2 \ - --kl-beta 0.001 \ - --format-in-reward \ - --lr 1e-6 \ - --max-train-rounds 1500 \ - --save-rounds 200 \ - --trend-every 10 \ - --prefetch-baseline \ - --output-dir ./output/reflexion_skill \ - --swanlab-project twinkle \ - --swanlab-exp reflexion_skill_rft diff --git a/cookbook/exp/legacy/train_reflexion_skill_replay.py b/cookbook/exp/legacy/train_reflexion_skill_replay.py deleted file mode 100644 index 0580a6ba4..000000000 --- a/cookbook/exp/legacy/train_reflexion_skill_replay.py +++ /dev/null @@ -1,114 +0,0 @@ -"""Replay-train the reflexion skill model from prebuilt exact RFT data. - -Use ``build_reflexion_skill_data.py`` first to create ``skill_dataset.jsonl``. This -script trains only the skill model from those frozen records; it does not run vLLM -rollouts, leak checks, or rubric diagnosis. - -Launch: - python cookbook/exp/embedding/train_reflexion_skill_replay.py \ - --data ./output/reflexion_skill_data/skill_dataset.jsonl -""" -import argparse -import json -import os -import sys -from collections import defaultdict -from typing import Any, Dict, List - -import train_reflexion_skill_rft as rft - - -def _build_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description=__doc__) - p.add_argument('--data', default='./output/reflexion_skill_data/skill_dataset.jsonl') - p.add_argument('--output-dir', default='./output/reflexion_skill_replay') - p.add_argument('--epochs', type=int, default=1) - p.add_argument('--sft-batch-size', type=int, default=8) - p.add_argument('--max-model-len', type=int, default=16384) - p.add_argument('--grpo-epsilon', type=float, default=0.2) - p.add_argument('--lr', type=float, default=1e-5) - p.add_argument('--save-rounds', type=int, default=50) - return p.parse_args() - - -def _load_chunks(path: str) -> List[List[Dict[str, Any]]]: - chunks: Dict[int, List[Dict[str, Any]]] = defaultdict(list) - fallback_chunk = 0 - with open(path, 'r', encoding='utf-8') as f: - for line_no, line in enumerate(f, 1): - if not line.strip(): - continue - row = json.loads(line) - if row.get('record_type') == 'config': - continue - for key in ('problem', 'response', 'advantage'): - if key not in row: - raise ValueError(f'{path}:{line_no} missing required field {key!r}') - ci = int(row.get('chunk', fallback_chunk)) - chunks[ci].append(row) - if 'chunk' not in row and len(chunks[ci]) >= 64: - fallback_chunk += 1 - return [chunks[k] for k in sorted(chunks) if chunks[k]] - - -def _init_model(args: argparse.Namespace, total_updates: int): - model = 'ms://Qwen/Qwen3-4B' - train_mesh = rft.DeviceMesh.from_sizes( - world_size=rft.TRAIN_GPUS, dp_size=rft.TRAIN_DP, fsdp_size=rft.TRAIN_FSDP) - device_groups = [ - rft.DeviceGroup(name='train', ranks=list(range(rft.TRAIN_GPUS)), device_type='GPU'), - ] - rft.twinkle.initialize(mode='ray', nproc_per_node=rft.TRAIN_GPUS, groups=device_groups, - lazy_collect=False) - model = rft.TransformersModel(model_id=model, device_mesh=train_mesh, - remote_group='train', ddp_config={'find_unused_parameters': False}) - from twinkle.patch.no_split_modules import NoSplitModulesPatch - model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) - model.set_template(rft.Template, model_id=model, - enable_thinking=True, max_length=args.max_model_len, - truncation_strategy='delete') - model.set_processor(rft.InputProcessor, padding_free=False) - model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) - model.set_optimizer('AdamW', lr=args.lr) - model.set_lr_scheduler('CosineWarmupScheduler', num_warmup_steps=10, - num_training_steps=max(1, total_updates)) - return model - - -def main() -> None: - args = _build_args() - if args.sft_batch_size % rft.TRAIN_DP != 0: - raise ValueError(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple ' - f'of the training dp size ({rft.TRAIN_DP})') - chunks = _load_chunks(args.data) - if not chunks: - raise ValueError(f'no train records found in {args.data}') - os.makedirs(args.output_dir, exist_ok=True) - total_updates = len(chunks) * args.epochs - model = _init_model(args, total_updates) - log_path = os.path.join(args.output_dir, 'train_log.jsonl') - cfg = {'record_type': 'config', 'mode': 'offline_replay', 'data': args.data, - 'chunks': len(chunks), 'epochs': args.epochs, 'lr': args.lr, - 'sft_batch_size': args.sft_batch_size} - rounds = 0 - with open(log_path, 'w', encoding='utf-8') as tlog: - tlog.write(json.dumps(cfg, ensure_ascii=False) + '\n') - for epoch in range(args.epochs): - for ci, samples in enumerate(chunks): - log = rft._train_chunk(model, None, samples, args) - rounds += 1 - log.update({'record_type': 'train_round', 'round': rounds, - 'epoch': epoch, 'chunk': ci}) - tlog.write(json.dumps(log, ensure_ascii=False) + '\n') - tlog.flush() - sys.stderr.write( - f'[replay-rft] e{epoch} c{ci}: n={log["n_samples"]} ' - f'micro={log["n_micro_batches"]} metric={log.get("metric")}\n') - if rounds % args.save_rounds == 0: - model.save(f'skill-rft-replay-{rounds}', output_dir=args.output_dir) - model.save('skill-rft-replay-final', output_dir=args.output_dir) - sys.stderr.write(f'[replay-rft] done: {rounds} updates; log -> {log_path}\n') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/legacy/train_reflexion_skill_rft.py b/cookbook/exp/legacy/train_reflexion_skill_rft.py deleted file mode 100644 index 55197b5d9..000000000 --- a/cookbook/exp/legacy/train_reflexion_skill_rft.py +++ /dev/null @@ -1,1569 +0,0 @@ -"""RFT cold-start for the reflexion skill generator (see reflexion.md §6). - -Trains an INDEPENDENT skill model to write reusable, transferable skills that, -when injected into a FROZEN base solver's system prompt, let the base solve problems -it first got wrong. The base is never trained — it only produces the reward signal. -Scoring is SEAM-style DETERMINISTIC: the base runs each candidate skill once at -temperature 0 (M=1), so the reward ``R in {0,1}`` (answer correct) carries no -sampling noise; the per-candidate advantage is group-relative within a problem -(``A = (R - mean) / (std + eps)``) and the skill model is updated online by GRPO — -problem-groups where every skill scores alike (std=0) contribute no gradient. - -Direction: skill GENERATION + recall. Skill-gen always runs with thinking ON, and -each hard problem is routed to EXACTLY ONE of two views (no reuse — kills memory -leak and holds cost at 1x): view A ``(problem + attempt) -> think + skills`` keeps -the online generator self-bootstrapping; view B ``(problem only) -> think + skills`` -is the deployment form, where the think is grounded on the query alone so it cannot -hallucinate an attempt. Both share the verified skill; the distilled ``<skills>`` -block is recalled into the base's system prompt at solve time. - -8-GPU layout (three DeviceGroups, one twinkle.initialize): - - ranks 0-3 : ``train`` — skill model, full-param FSDP2, dp=4 - - ranks 4-5 : ``skill_sampler`` — skill model rollouts (vLLM, tp1 dp2) - - ranks 6-7 : ``base_sampler`` — frozen base solver (vLLM, tp1 dp2) -CheckpointEngineManager syncs train -> skill_sampler after every optimizer step; -base_sampler is never synced. - -Leak filtering uses ``LeakVerifier(sampler=None)`` via the backup teacher API -(no local judge, no distillation): set LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / -LLM_BACKUP_MODEL. - -Launch (8 GPUs): - LLM_BACKUP_API_KEY=... LLM_BACKUP_BASE_URL=... LLM_BACKUP_MODEL=... \ - python cookbook/exp/embedding/train_reflexion_skill_rft.py --n 2000 --chunk-size 16 -""" -import argparse -import hashlib -import json -import os -import re -import sys -import time -from typing import Any, Dict, List, Optional, Tuple - -import numpy as np - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.data_format import pack_user_data -from twinkle.checkpoint_engine import CheckpointEngineManager -from twinkle.model import TransformersModel -from twinkle.processor import InputProcessor -from twinkle.sampler import vLLMSampler -from twinkle.template import Template -from twinkle_agentic.verifier import LeakVerifier, RubricVerifier -from twinkle_agentic.verifier.rubric_verifier import RubricItem - -# Reuse the reference eval's dataset + grading + prompts + sampling config, and the -# phase-0 pipeline's parsing / rollout / injection helpers (Find > Create). -from eval_gpqa_rag import (GEN_GPU_MEM, GEN_MODEL_ID, build_direct_prompt, # noqa: F401 - load_aops, load_math) -from eval_reflexion_skill import (_EX_PROBLEM, _clean_text, # noqa: F401 - _parse_seq, _run_samples, build_skill_solve_prompt) - -logger = get_logger() - -try: - import swanlab -except ImportError: # optional; metric logging degrades to stdout + jsonl only - swanlab = None - - -# -- GPU layout --------------------------------------------------------------- -TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 4)) -SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) -BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 2)) -NUM_GPUS = TRAIN_GPUS + SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS -# FSDP shard group size within a dp replica; TRAIN_DP is the data-parallel axis that -# ``forward_backward`` (slice_dp) splits each mini-batch over. -TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', 2)) -TRAIN_DP = max(1, TRAIN_GPUS // TRAIN_FSDP) - - -# --------------------------------------------------------------------------- -# Skill-generation prompt (DISTILL the useful approach, per the new direction) -# --------------------------------------------------------------------------- -# --- Previous STRICT view-A system prompt (commented out; kept for easy revert). It -# hard-required 3-5 bullets, "output nothing after </skills>", one-imperative-sentence -# items, no-narration, and a strict do-not-reveal block. The soft SEAM-style version -# below drops those four format demands, frames the skills as advisory reminders, and -# explains how they are used. --- -# SKILL_GEN_SYSTEM = ( -# 'You are distilling reusable problem-solving SKILLS from one worked episode. ' -# 'You are shown a competition problem, the guidance the solver was given, and the ' -# "solver's own attempt (its reasoning may be partly right and partly wrong).\n\n" -# 'FIRST, in your private thinking, do ALL of: (a) work out what this TYPE of problem ' -# 'fundamentally requires; (b) pinpoint WHERE THIS attempt actually went wrong ' -# '(when a process-check report is provided below, use its flagged criteria as ' -# 'evidence, but confirm each against the attempt yourself) — ' -# 'the decisive misstep, a missing idea, a wrong turn, or the way it stalled, looped ' -# 'on the same step, or ran the length budget out without ever committing to an ' -# 'answer; and (c) imagine AS MANY DIFFERENT angles as you can — distinct approaches ' -# 'or representations that could crack this problem, alternative solution paths, and ' -# 'the various ways a solver could plausibly go wrong on it (a few words each, do NOT ' -# 'develop them fully). THEN commit to the angle you find most decisive and write a ' -# 'SHORT list of skills that would have PREVENTED that specific ' -# 'failure and would raise the success rate of a SIMILAR solver on SIMILAR problems. ' -# 'Ground each skill in the concrete mistake you found, but state it as a GENERAL, ' -# 'transferable rule — not a patch hard-coded to this problem. Across the 3-5 ' -# 'bullets, prioritise in this order:\n' -# '1. the decisive method or representation this class of problem calls for (what to ' -# 'set up or reach for first);\n' -# '2. the specific mistake that derailed THIS attempt, recast as a general pitfall, ' -# 'plus the quick check that catches it;\n' -# '3. a compact, ordered procedure that reliably drives toward a final answer;\n' -# '4. convergence discipline: once the key quantity is in hand, commit to a single ' -# 'concrete final answer in the required format instead of re-deriving, endless ' -# 'case-splitting, looping on the same check, or overrunning the length budget.\n\n' -# 'OUTPUT FORMAT (strict):\n' -# '- Keep the thinking COMPACT — a quick brainstorm of angles then a decision, not a ' -# 'full solution and not a re-statement of these instructions. AFTER it, ' -# 'output ONLY a markdown bullet list of 3-5 items WRAPPED IN <skills> and </skills> ' -# 'tags — no preamble, no narration outside the tags. Output nothing after ' -# '</skills>.\n' -# '- Each item is ONE short imperative sentence (a method, heuristic, check, or ' -# 'habit).\n' -# '- Inside the tags: no narration, no "The student...", no headings, no restating ' -# 'the problem.\n\n' -# 'CONTENT RULES (strict):\n' -# '- Do NOT reveal the final answer or the multiple-choice option.\n' -# '- Do NOT state the specific numbers, values, or key intermediate results of THIS ' -# 'problem.\n' -# '- Every item must be GENERAL and transferable, not a step-by-step solution to ' -# 'THIS problem.\n\n' -# 'Follow the example below for the exact tags, style, and level of generality.' -# ) -SKILL_GEN_SYSTEM = ( - 'You are a mathematics coach. You are shown a competition problem together with an ' - 'automated process-check of an earlier solver attempt at it -- which solution ' - 'criteria the attempt passed or failed, and suggested fixes for the failures. You ' - 'do NOT see the attempt itself, only this check. Treat the check as privileged ' - 'training scaffolding: study it together with the problem, identify the ' - 'problem-visible features that make each useful flagged failure relevant, then ' - 'rephrase those lessons as self-contained reusable skills. The goal is not to ' - 'continue from the check, cite it, or hide it silently; the goal is to turn it to ' - 'a skill pattern which prevents the model falls into similar pitfalls in the future.\n\n' - 'Good skills name the observable trigger, the method worth reaching for, the ' - 'pitfall to watch, and a quick verification habit. Prefer formulations like ' - '"When a configuration has ...", "Before setting up ...", or "Check whether ..." ' - 'over references to the process-check, failed criteria, or the earlier attempt. ' - "These tips are advisory: they will be placed in a solver's system prompt as gentle " - 'reminders before it works through a SIMILAR problem on its own, without seeing ' - 'this process-check. So keep them general and transferable rather than a worked ' - 'solution to this exact problem, and do not state its specific intermediate values ' - 'or final answer. Think briefly first, then give your tips as a markdown bullet ' - 'list wrapped in <skills> and </skills>, like the example below.' -) - -# One-shot demo of the recommended mix (method / pitfall+check / procedure / -# convergence), answer-free — anchors both the format and the content priorities. -_EX_SKILLS = ( - '<skills>\n' - '- Rewrite each square root by factoring its radicand into a perfect square times ' - 'a remainder, then move the perfect-square factor outside.\n' - '- Avoid the classic trap $\\sqrt{a}+\\sqrt{b}\\ne\\sqrt{a+b}$; only combine terms ' - 'sharing the same simplest radical, and sanity-check by estimating each root.\n' - '- Procedure: simplify every radical, group like radical terms, add their ' - 'coefficients, then reduce to simplest form.\n' - '- Once the expression is in simplest form, commit to that single result as the ' - 'final answer rather than re-checking indefinitely.\n' - '</skills>') - - -# View A user template: the problem + the automated rubric process-check of an earlier -# attempt (PASS/FAIL per criterion + suggested fixes). The attempt trajectory is NOT -# shown -- the rubric findings are the evidence the skill model grounds its tips on, -# which avoids feeding the (often long, non-terminating) attempt into the prompt. -SKILL_GEN_USER_RUBRIC = ( - 'Problem:\n{problem}\n\n' - 'Process check of an earlier attempt (automated rubric verifier -- treat as ' - 'evidence, not gospel; PASS/FAIL per criterion with suggested fixes for failures):\n' - '{diagnosis}\n\n' - 'Now output a self-contained skills bullet list. Each bullet should still be useful ' - 'if the process check were removed: connect any useful flagged failure to ' - 'problem-visible features, general methods, and quick checks rather than citing the ' - 'rubric or the earlier attempt. \n\n' - 'Note: **Do not solve the problem, only generate skills**. Now Begin:' -) - - -def build_skillgen_prompt(problem: str, diagnosis: str) -> Dict[str, Any]: - """View A skill-gen prompt: system + one-shot format demo + the real episode - (problem + the rubric process-check of an earlier attempt). The attempt trajectory - is deliberately NOT shown -- the rubric findings localise the failure without the - generator having to re-chew (and often re-solve) a long, possibly non-terminating - attempt. The one-shot demo is query-only; only the real turn carries the diagnosis.""" - return {'messages': [ - {'role': 'system', 'content': SKILL_GEN_SYSTEM}, - # {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=_EX_PROBLEM)}, - # {'role': 'assistant', 'content': _EX_SKILLS}, - {'role': 'user', - 'content': SKILL_GEN_USER_RUBRIC.format(problem=problem, diagnosis=diagnosis)}, - ]} - - -# --------------------------------------------------------------------------- -# View B: query-only skill-gen (deployment form). No attempt is shown — the model -# must reason about the problem TYPE from the query alone, so the think is grounded -# on the query and cannot narrate/fabricate an attempt. Format is deliberately -# distinct from view A so the model learns the two modes as separate contracts. -# --------------------------------------------------------------------------- -# --- Previous STRICT view-B (query-only) system prompt (commented out; kept for revert). -# Same four format demands as the old view A. Soft SEAM-style version below. --- -# SKILL_GEN_SYSTEM_Q = ( -# 'You are distilling reusable problem-solving SKILLS for a CLASS of problems. You ' -# 'are shown ONE competition problem and NOTHING else — no solution, no attempt. ' -# 'FIRST, in your private thinking, imagine AS MANY DIFFERENT angles as you can — ' -# 'distinct approaches or representations that could crack this TYPE of problem, ' -# 'alternative solution paths, and the various ways a solver could plausibly go wrong ' -# 'on it (a few words each, do NOT develop them fully). THEN commit to what you find ' -# 'most decisive and write a SHORT list of skills that would raise a solver\'s success ' -# 'rate on SIMILAR problems. Across the 3-5 bullets, prioritise in this order:\n' -# '1. the decisive method or representation this class of problem calls for (what to ' -# 'set up or reach for first);\n' -# '2. the specific pitfall that derails such problems, plus the quick check that ' -# 'catches it;\n' -# '3. a compact, ordered procedure that reliably drives toward a final answer;\n' -# '4. convergence discipline: once the key quantity is in hand, commit to a single ' -# 'concrete final answer in the required format instead of re-deriving, endless ' -# 'case-splitting, or overrunning the length budget.\n\n' -# 'OUTPUT FORMAT (strict):\n' -# '- Keep the thinking COMPACT — a quick brainstorm of angles then a decision, not a ' -# 'full solution and not a re-statement of these instructions. AFTER ' -# 'it, output ONLY a markdown bullet list of 3-5 items WRAPPED IN <skills> and ' -# '</skills> tags — no preamble, no narration outside the tags. Output nothing after ' -# '</skills>.\n' -# '- Each item is ONE short imperative sentence (a method, heuristic, check, or ' -# 'habit).\n' -# '- Inside the tags: no narration, no headings, no restating the problem, and no ' -# 'reference to any attempt, student, or solution.\n\n' -# 'CONTENT RULES (strict):\n' -# '- Do NOT solve THIS problem or reveal its final answer or multiple-choice option.\n' -# '- Do NOT state the specific numbers, values, or key intermediate results of THIS ' -# 'problem.\n' -# '- Every item must be GENERAL and transferable to other problems of the same ' -# 'type.\n\n' -# 'Follow the example below for the exact tags, style, and level of generality.' -# ) -SKILL_GEN_SYSTEM_Q = ( - 'You are a mathematics coach. You are shown ONE competition problem and nothing ' - 'else — no solution and no attempt. Think about what approach this KIND of problem ' - 'calls for and where solvers tend to slip, then distil a few reusable tips.\n\n' - "These tips are advisory: they will be placed in a solver's system prompt as gentle " - 'reminders before it works through a SIMILAR problem on its own. So keep them ' - 'general and transferable — the method worth reaching for, the pitfall to watch and ' - 'a quick check, and the discipline to settle on a final answer — rather than a ' - 'worked solution to this exact problem, and without stating its specific ' - 'intermediate values or its final answer. Think briefly first, then give your tips ' - 'as a markdown bullet list wrapped in <skills> and </skills>, like the example below.' -) - -SKILL_GEN_USER_Q = ( - 'Problem:\n{problem}\n\n' - 'Now reason about this TYPE of problem, then output the skills bullet list.' -) - - -def build_querygen_prompt(problem: str) -> Dict[str, Any]: - """View B skill-gen prompt: system + one-shot demo + the problem ALONE (no - attempt) — matching what is available at deployment (query only).""" - return {'messages': [ - {'role': 'system', 'content': SKILL_GEN_SYSTEM_Q}, - # {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=_EX_PROBLEM)}, - # {'role': 'assistant', 'content': _EX_SKILLS}, - {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=problem)}, - ]} - - -def _assign_view(problem: str, args: argparse.Namespace) -> str: - """Deterministically route a problem to exactly one view (stable across restarts - and across the generation/SFT sides). ``--view-b-frac`` of problems go to view B.""" - h = int(hashlib.md5(f'{args.seed}:{problem}'.encode('utf-8')).hexdigest(), 16) - return 'B' if (h % 100000) / 100000.0 < args.view_b_frac else 'A' - - -def _skillgen_messages(problem: str, view: str, diagnosis: str) -> List[Dict[str, Any]]: - """Single source of truth for the skill-gen prompt, used at BOTH generation and - training time so they can never diverge. View A with a localisable failure uses - problem + rubric findings (NO trajectory); view B -- or a view-A problem whose rubric - flagged NO failure (``[FAIL]`` absent: all-pass or missing diagnosis) -- is query-only. - So view A DEGRADES to view B whenever there is nothing concrete to correct.""" - if view == 'B' or '[FAIL]' not in (diagnosis or ''): - return build_querygen_prompt(problem)['messages'] - return build_skillgen_prompt(problem, diagnosis)['messages'] - - -def _view_prompt(r: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: - """The skill-gen prompt for problem ``r`` under its assigned view (routing in - ``_skillgen_messages``: view A carries the rubric process-check; view B, and any - view-A problem with no rubric failure, is query-only).""" - return {'messages': _skillgen_messages(r['problem'], r['_view'], r.get('_rubric_diag', ''))} - - -_BULLET_RE = re.compile(r'(?m)^\s*(?:[-*]|\d+[.)])\s') -# Trajectory/meta references that betray CoT fragments leaking into the block; any -# hit fails the purity gate (the problem is then re-sampled, per --skill-retries). -_META_RE = re.compile( - r'\b(the student|the solver|the attempt|this attempt|the trace|the response|' - r'in the (?:attempt|trace|response|solution)|as (?:shown|seen|noted) above|' - r'the (?:above|previous|earlier)|my (?:reasoning|analysis)|i (?:think|need|will))\b', - re.IGNORECASE) - - -def _is_clean_block(block: str) -> bool: - """Purity gate for thinking-ON skill-gen: the block must be a pure bullet list - (every non-empty line a bullet — no prose/CoT fragments) with no meta/trajectory - reference. Answer leak is caught separately by the backup-teacher leak stage.""" - lines = [ln for ln in (l.strip() for l in block.splitlines()) if ln] - if not lines or not all(_BULLET_RE.match(ln) for ln in lines): - return False - return _META_RE.search(block) is None - - -def _extract_skills_block(text: str) -> Optional[str]: - """Return the clean ``<skills>...</skills>`` block, or None if not parseable. - - Skill-gen runs with thinking ON, so the model must end its reasoning with an explicit - ``</think>`` before committing an answer (whether the opening ``<think>`` is emitted by - the model or pre-injected by the chat template). We therefore REQUIRE ``</think>`` and - read only the text after the last one; its absence means the token budget was exhausted - mid-reasoning (nothing committed, per reflexion.md §6.8) — reject so a draft or a - system-prompt demo echo inside the CoT can never be mistaken for the answer. Within - the answer take the ``<skills>`` block (closing tag optional), strip stray tags, and - require ``_is_clean_block`` — prose-mixed / meta-referencing fragments are rejected - for re-sampling.""" - low = text.lower() - end_think = low.rfind('</think>') - if end_think < 0: - return None # reasoning never closed -> no committed answer - answer = text[end_think + len('</think>'):] - low_a = answer.lower() - s = low_a.find('<skills>') - if s < 0: - return None - inner = s + len('<skills>') - e = low_a.find('</skills>', inner) - block = (answer[inner:e] if e >= 0 else answer[inner:]).strip() - block = re.sub(r'</?(?:skills|think)>', '', block, flags=re.IGNORECASE).strip() - if not _is_clean_block(block): - return None - return block - - -# --------------------------------------------------------------------------- -# OPTIONAL stricter leak criterion (currently UNUSED -- the run uses answer_only=True, -# which flags ONLY the final answer). This variant ALSO flags concrete intermediate KEY -# results, while still permitting method / plan / pitfalls / checks. To enable, pass -# judge_system=_LEAK_JUDGE_SYSTEM to the LeakVerifier below. -# --------------------------------------------------------------------------- -_LEAK_JUDGE_SYSTEM = """\ -You check whether a HINT that will be shown to someone solving a math TASK gives away -this task's own results. - -The hint may FREELY describe the general method, which approach or technique to use, the -steps or plan to follow, common pitfalls, and sanity checks -- even when that points -strongly at HOW to solve THIS task. Describing the approach is expected of a good hint. - -The hint LEAKS only if, for THIS specific task, it states either: -- the final answer or final result (a value, expression, choice, label, or verbatim - output); or -- a concrete decisive INTERMEDIATE key result -- a specific computed value, quantity, or - fact unique to this task that hands over a key step of the answer. - -If it names only the method / plan / pitfalls / checks WITHOUT stating those concrete -intermediate values or the final result, it does NOT leak. - -Reply with exactly one word: LEAK or CLEAN.""" - - -# --------------------------------------------------------------------------- -# Rubric process-check (view A only): a frozen teacher diagnoses the base's failed -# attempt so the skill model grounds its error analysis on a verified fault -# localisation instead of guessing. Teacher-only (sampler=None -> every diagnose() -# hits llm_backup); mirrors eval_dualline_math's fixed math rubric. -# --------------------------------------------------------------------------- -_RFT_DIAG_SYSTEM = """\ -You are a process error checker for a math solution attempt. You are given a math -problem, a rubric, and one attempted solution segment. Decide PASS or FAIL for each -criterion and explain only the process error type. - -Output STRICT JSON (no prose outside it) with this shape: -{ - "items": [ - {"index": 1, "verdict": "PASS", "reason": "<why the process satisfies it>", - "fix": ""}, - {"index": 2, "verdict": "FAIL", "reason": "<what process step is wrong>", - "fix": "<method-level correction, without computing the corrected result>"} - ], - "overall": "OK" | "ISSUES", - "summary": "<one sentence naming the process issue, not the answer>" -} - -Rules: -- Judge every criterion independently and literally; a [Hard Rule] is FAIL unless - unambiguously satisfied. -- Judge ONLY what is observable in THIS segment. -- Content inside <think>...</think> (or <thinking>) is internal reasoning, not - user-facing output; ignore it for "output only X" style criteria. -- For PASS items, leave "fix" as "". -- For FAIL items, "reason", "fix", and "summary" must describe only the flawed - step, theorem, arithmetic operation, case split, or verification habit. -- NEVER try to solve the query or state the correct final answer, corrected final expression, option letter, - graph/choice label, or any exact value that the answer should become. -- NEVER write phrases like "the correct answer is", "which gives", "yielding", - "should be <value>", "Option <letter>", or "Graph <letter>". -- If a fix would require naming a corrected value, replace it with a method-level - instruction such as "redo that computation carefully" or "apply the theorem with - the correct quantities". -- Keep every "reason" and "fix" clear and concise. -- "overall" is "OK" only if NO criterion is FAIL. -- Output only the JSON object.""" - - -_RFT_DIAG_USER = """\ -## Task / query (context) -{query} - -## Rubric -{rubric} - -## Segment -{segment} - -Now output the diagnostic JSON object.""" - - -class _RftRubricVerifier(RubricVerifier): - def _diagnose_trajectory(self, query: str, rubric_block: str, segment_text: str) -> dict: - user = _RFT_DIAG_USER.format(query=query, rubric=rubric_block, segment=segment_text) - return {'messages': [ - {'role': 'system', 'content': _RFT_DIAG_SYSTEM}, - {'role': 'user', 'content': user}, - ]} - - -_MATH_RUBRIC = [ - ('The reasoning contains no arithmetic or algebraic error', True), - ('Each step follows logically from the previous ones', True), - ('No formula or theorem is misstated or misapplied', True), - ('The approach is on track to answer the actual question asked', False), - ('No step contradicts an earlier established fact', False), -] - - -def _build_rubric_checker() -> Optional['RubricVerifier']: - """Fixed math-process rubric verifier, teacher-served. None if no LLM backup env.""" - if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') - or os.environ.get('OPENAI_API_KEY')): - return None - return _RftRubricVerifier( - fixed_rubric=[RubricItem(t, is_hard=h) for t, h in _MATH_RUBRIC], gate=True) - - -def _format_diagnosis(detail) -> str: - """One line per criterion (PASS/FAIL + reason + fix on FAIL) then a summary — the - compact evidence block appended to the view-A skill-gen prompt.""" - rub = detail.rubric - lines = [] - for it in detail.items: - text = rub[it.index - 1].text if 0 < it.index <= len(rub) else f'criterion {it.index}' - if it.verdict: - lines.append(f'- [PASS] {text}') - else: - tail = f': {it.reason}' if it.reason else '' - tail += f' (fix: {it.fix})' if it.fix else '' - lines.append(f'- [FAIL] {text}{tail}') - if detail.summary: - lines.append(f'Summary: {detail.summary}') - return '\n'.join(lines) - - -def _diagnose_views(checker, hard: List[Dict[str, Any]], args: argparse.Namespace, - diag_cache: Optional[Dict[str, str]] = None) -> None: - """Rubric-check every view-A problem's greedy attempt in parallel, stashing the - formatted findings on ``r['_rubric_diag']`` (view B stays empty). A checker error - or empty result degrades to no diagnosis (the plain view-A prompt).""" - from concurrent.futures import ThreadPoolExecutor - targets = [r for r in hard if r.get('_view') == 'A'] - if not checker or not targets: - return - - def _cache_key(r: Dict[str, Any]) -> str: - init_text = r.get('_init', [{}])[0].get('text', '') - return hashlib.md5(f'{r["problem"]}\n{init_text}'.encode('utf-8')).hexdigest() - - pending = [] - for r in targets: - key = _cache_key(r) - if diag_cache is not None and key in diag_cache: - r['_rubric_diag'] = diag_cache[key] - else: - pending.append((r, key)) - if not pending: - return - - def _run(item: Tuple[Dict[str, Any], str]) -> Tuple[Dict[str, Any], str, str, bool]: - r, key = item - seg = {'messages': [ - {'role': 'user', 'content': r['problem']}, - {'role': 'assistant', 'content': r['_init'][0]['text']}, - ]} - try: - return r, key, _format_diagnosis(checker.diagnose(seg, query=r['problem'])), True - except Exception as exc: # teacher hiccup -> fall back to no-diagnosis prompt - logger.warning(f'[rubric] diagnose error: {exc}') - return r, key, '', False - - workers = max(1, min(args.rubric_workers, len(pending))) - with ThreadPoolExecutor(max_workers=workers) as ex: - for r, key, diag, ok in ex.map(_run, pending): - r['_rubric_diag'] = diag - if ok and diag_cache is not None: - diag_cache[key] = diag - - -# --------------------------------------------------------------------------- -# Online data generation (one chunk; every candidate is recorded, untruncated) -# --------------------------------------------------------------------------- -def _roll(x: Dict[str, Any]) -> Dict[str, Any]: - """Full (untruncated) rollout record for offline analysis.""" - return {k: x[k] for k in ('pred', 'correct', 'terminated', 'passed', - 'stop_reason', 'gen_tokens', 'text')} - - -def _empty_roll() -> Dict[str, Any]: - """Fallback rollout when the sampler returned nothing for a prompt.""" - return {'pred': '', 'correct': False, 'terminated': False, 'passed': False, - 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} - - -def _assign_advantages(hard: List[Dict[str, Any]], args: argparse.Namespace) -> None: - """Group-relative advantage (SEAM-style) over each problem's clean, scored candidates, - using the DETERMINISTIC greedy reward ``R in {0, 1}`` (answer CORRECT only; - termination is NOT part of the reward -- monitored via `terminated`/`passed` only): - - A_j = (R_j - mean_R) / (std_R + eps) - - Groups where every candidate shares the same reward (``std_R == 0``: all solve or all - fail) get advantage 0 and contribute no gradient -- GRPO's own group variance - auto-selects the informative problems, so no explicit difficulty / marginal gate is - needed. Because the reward is deterministic (M=1 greedy, no pass@k sampling), the - std-normalisation no longer amplifies rollout noise (the reason it was dropped for the - old stochastic marginal). ``kept`` marks above-average candidates (for reporting only). - """ - eps = 1e-6 - for r in hard: - for c in r['_cands']: - c['advantage'], c['grpo_adv'], c['kept'] = 0.0, 0.0, False - cs = ([c for c in r['_cands'] if c.get('reward') is not None] if args.format_in_reward - else [c for c in r['_cands'] if c['leaked'] is False and c.get('reward') is not None]) - if len(cs) < 2: - continue - rewards = [c['reward'] for c in cs] - mean_r = sum(rewards) / len(rewards) - std = (sum((x - mean_r) ** 2 for x in rewards) / len(rewards)) ** 0.5 - if std < 1e-9: - continue # all candidates equal (all solve / all fail) -> no learning signal - for c in cs: - adv = (c['reward'] - mean_r) / (std + eps) - c['advantage'], c['grpo_adv'], c['kept'] = adv, adv, c['reward'] > mean_r - - -def _full_record(r: Dict[str, Any], ci: int) -> Dict[str, Any]: - """Complete per-problem record: init attempt, baseline, and ALL candidates - (parseable/leaked/scored alike) with full text — nothing dropped or truncated.""" - init = r['_init'][0] - rec: Dict[str, Any] = { - 'record_type': 'problem', 'chunk': ci, 'problem': r['problem'], - 'reference_answer': r['reference_answer'], 'level': r.get('level', ''), - 'failed_first_try': r['_failed'], - 'init_attempt': {'text': init['text'], 'pred': init['pred'], - 'correct': init['correct'], 'terminated': init['terminated'], - 'stop_reason': init['stop_reason'], 'gen_tokens': init['gen_tokens']}, - } - rec['baseline_pass'] = r['_baseline_pass'] - rec['is_hard'] = r['_hard'] - rec['view'] = r.get('_view', '') - rec['rubric_diag'] = r.get('_rubric_diag', '') - rec['baseline_rolls'] = [_roll(x) for x in r['_baseline_rolls']] - rec['candidates'] = [{ - 'skills': c['skills'], 'response': c['response'], 'parseable': c['parseable'], - 'leaked': c['leaked'], 'leak_reason': c['leak_reason'], 'leak_source': c['leak_source'], - 'with_pass': c['with_pass'], 'reward': c.get('reward'), - 'advantage': c.get('advantage'), 'grpo_adv': c.get('grpo_adv'), 'kept': c.get('kept'), - 'rolls': [_roll(x) for x in c['rolls']], - } for c in r['_cands']] - return rec - - -def _view_stats(hard: List[Dict[str, Any]], view: str) -> Dict[str, Any]: - """Per-view yield: hard problems, clean candidates, and the ADOPTION rate — - the fraction of hard problems that produced at least one clean, non-zero-advantage - candidate (i.e. a record that actually reaches training). Watching A vs B and - early vs late tells whether query-only (B) catches up to trajectory-grounded (A).""" - hv = [r for r in hard if r.get('_view') == view] - cands = [c for r in hv for c in r['_cands'] if c['parseable']] - clean = [c for c in cands if c['leaked'] is False] - adopted = sum(1 for r in hv - if any(c['leaked'] is False and abs(c.get('advantage') or 0.0) > 1e-9 - for c in r['_cands'])) - return { - 'n_hard': len(hv), 'n_candidates_parseable': len(cands), - 'n_clean': len(clean), 'n_adopted_problems': adopted, - 'adoption_rate': (adopted / len(hv)) if hv else 0.0, - } - - -def _is_trainable(c: Dict[str, Any], args: argparse.Namespace) -> bool: - """A candidate reaches the GRPO update iff its advantage is non-zero. With - --format-in-reward every candidate carries a reward (unparseable/leaked score 0), - so non-zero advantage is the only gate; otherwise it must also be clean and scored. - Single source of truth for both the summary counts and ``_group_records``.""" - adv_nz = abs(c.get('advantage') or 0.0) > 1e-9 - if args.format_in_reward: - return adv_nz - return c['leaked'] is False and c.get('with_pass') is not None and adv_nz - - -def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespace) -> Dict[str, Any]: - """Per-chunk aggregates — watch these across chunks to see if the RFT'd skill - model produces better skills over time (yield, leak rate, lift, termination).""" - failed = [r for r in chunk if r['_failed']] - hard = [r for r in chunk if r['_hard']] - all_cands = [c for r in chunk for c in r['_cands']] - cands = [c for c in all_cands if c['parseable']] - scored = [c for c in cands if c['with_pass'] is not None] - ws_rolls = [x for c in scored for x in c['rolls']] - # With --format-in-reward, unparseable/leaked candidates also carry a (0) reward and are - # trained, so count trainables over ALL candidates; else only clean scored ones. - train_cands = [c for c in all_cands if _is_trainable(c, args)] - base_acc = (sum(r['_baseline_pass'] for r in hard) / len(hard)) if hard else 0.0 - ws_acc = (sum(c['with_pass'] for c in scored) / len(scored)) if scored else 0.0 - # -- signal-source monitor: how much of the GRPO signal comes from base-FAIL problems - # (the offensive "rescue a failure" signal we want) vs base-success (defensive "don't - # break an easy one"). abs_adv_from_fail_frac ~0.1 was the diagnosed failure mode. -- - fail_cands = [c for r in chunk if r['_failed'] for c in r['_cands']] - abs_adv = lambda cs: sum(abs(c.get('advantage') or 0.0) for c in cs) - total_abs = abs_adv(all_cands) - return { - 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), - 'n_failed_first_try': len(failed), 'n_hard': len(hard), - 'n_generated': len(all_cands), - 'n_candidates_parseable': len(cands), - 'n_unparseable': len(all_cands) - len(cands), - 'n_leaked': sum(1 for c in cands if c['leaked']), - 'n_clean': sum(1 for c in cands if c['leaked'] is False), - 'n_reward_pos': sum(1 for c in scored if c['reward']), - 'n_train_samples': len(train_cands), - 'n_train_from_fail': sum(1 for c in fail_cands if _is_trainable(c, args)), - 'abs_adv_from_fail_frac': (abs_adv(fail_cands) / total_abs) if total_abs > 0 else 0.0, - 'avg_baseline_pass_on_hard': base_acc, - 'avg_withskill_pass': ws_acc, - 'avg_lift': ws_acc - base_acc, - 'termination_rate_withskill': (sum(1 for x in ws_rolls if x['terminated']) / len(ws_rolls)) if ws_rolls else 0.0, - 'view_A': _view_stats(hard, 'A'), 'view_B': _view_stats(hard, 'B'), - } - - -def _group_records(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> List[Dict[str, Any]]: - """GRPO training records: every clean, scored skill candidate with a NON-ZERO - advantage (positive pushes the skill up, negative down; the group-relative - usefulness-over-base advantage was set in _assign_advantages). Each carries its - ``view`` and the rubric ``diagnosis``; the prompt (identical to generation) is rebuilt - from those by ``_skillgen_messages`` -- no trajectory is stored or replayed.""" - out = [] - for r in chunk: - if not r['_hard']: - continue - view = r.get('_view', 'A') - for c in r['_cands']: - if _is_trainable(c, args): - out.append({ - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'view': view, - 'diagnosis': r.get('_rubric_diag', ''), - 'response': c['response'], 'skills': c['skills'], - 'advantage': c['advantage'], 'grpo_adv': c['grpo_adv'], 'kept': c['kept'], - 'reward': c['reward'], 'with_pass': c['with_pass'], - }) - return out - - -def _apply_baseline(r: Dict[str, Any], roll: Dict[str, Any]) -> None: - """Write a (cached or fresh) greedy baseline roll onto a problem and RESET the per-chunk - working state, so a problem reused in a later chunk never carries prior skill candidates.""" - r['_baseline_rolls'], r['_cands'] = [roll], [] - r['_init'] = [roll] - r['_failed'] = not roll['correct'] - r['_baseline_pass'] = 1.0 if roll['correct'] else 0.0 - r['_hard'] = True # process EVERY selected problem; group variance selects - - -def _baseline_rollout(base_sampler, problems: List[Dict[str, Any]], base_dp: int, - args: argparse.Namespace, cache: Dict[str, Dict[str, Any]]) -> int: - """Phase 1: base solves each problem GREEDILY once (T=0, M=1), keyed-cached by problem - text across chunks. The base sampler is FROZEN and decoding is greedy, so a problem's - baseline never changes over the run -- a cache hit is exact and skips the sampler. - Returns the number of FRESH sampler rollouts (cache misses) for efficiency reporting.""" - todo = [r for r in problems if r['problem'] not in cache] - if todo: - base_out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], - 1, args.max_tokens, base_dp, temperature=0.0) - for r, seqs in zip(todo, base_out): - cache[r['problem']] = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - for r in problems: - _apply_baseline(r, cache[r['problem']]) - return len(todo) - - -def _baseline_class(r: Dict[str, Any]) -> str: - """Bucket a baselined problem by its greedy outcome: ``success`` (base solved it), - ``fail_loop`` (ran the length budget out / never terminated -- the mode skills rescue - best), or ``fail_wrong`` (terminated cleanly but the answer is wrong).""" - roll = r['_init'][0] - if roll['correct']: - return 'success' - if roll['stop_reason'] == 'length' or not roll['terminated']: - return 'fail_loop' - return 'fail_wrong' - - -_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') - - -def _norm_num_text(num: str) -> str: - try: - f = float(num) - return str(int(f)) if f == int(f) else str(f) - except Exception: - return str(num).strip() - - -def _numeric_value(raw: Any) -> Optional[str]: - if raw is None: - return None - s = str(raw).strip().strip('$').strip() - s = s.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac') - s = re.sub(r'\\!|\\,|\\;|\\ |\\left|\\right|\s', '', s) - m = re.fullmatch(r'\\frac\{(-?\d+)\}\{(-?\d+)\}', s) - if m: - a, b = int(m.group(1)), int(m.group(2)) - return _norm_num_text(str(a / b)) if b else None - m = re.fullmatch(r'(-?\d+)/(-?\d+)', s) - if m: - a, b = int(m.group(1)), int(m.group(2)) - return _norm_num_text(str(a / b)) if b else None - if _NUM_RE.fullmatch(s): - return _norm_num_text(s) - return None - - -def _numeric_only_records(records: List[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], int]: - out = [] - dropped = 0 - for r in records: - ref = _numeric_value(r.get('reference_answer')) - if ref is None: - dropped += 1 - continue - rr = dict(r) - rr['reference_answer'] = ref - out.append(rr) - return out, dropped - - -def _load_records(args: argparse.Namespace) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], Dict[str, int]]: - need_split = args.eval_size > 0 - load_n = 0 if (args.numeric_only or need_split) else args.n - records = (load_aops(n=load_n, seed=args.seed) if args.dataset == 'aops' - else load_math(n=load_n, seed=args.seed)) - raw_n = len(records) - dropped = 0 - if args.numeric_only: - records, dropped = _numeric_only_records(records) - rng = np.random.RandomState(args.seed) - rng.shuffle(records) - eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 - eval_records = [dict(r) for r in records[:eval_n]] - train_pool = records[eval_n:] - train_n = args.n if args.n > 0 else len(train_pool) - train_records = [dict(r) for r in train_pool[:train_n]] - overlap = {r['problem'] for r in train_records} & {r['problem'] for r in eval_records} - if overlap: - raise ValueError(f'fixed eval/train overlap detected: {len(overlap)} duplicated problems') - stats = {'raw_loaded': raw_n, 'numeric_dropped': dropped, - 'train_records': len(train_records), 'eval_records': len(eval_records)} - return train_records, eval_records, stats - - -class _ProblemPool: - """Cyclic draw source over the loaded problems. Each full pass reshuffles with - ``seed + epoch`` and bumps ``epoch`` (matching the old per-epoch reshuffle); the - initial pass keeps the loader's shuffled order. Draws never run out.""" - - def __init__(self, records: List[Dict[str, Any]], seed: int): - self._records = list(records) - self._seed = seed - self._cursor = 0 - self.epoch = 0 - self.baseline_cache: Dict[str, Dict[str, Any]] = {} # problem text -> frozen greedy roll - - def draw(self, k: int) -> List[Dict[str, Any]]: - """Return ``k`` DISTINCT problems (unique within this call, so one chunk never - processes the same problem twice even when the cursor wraps mid-draw). ``k`` is - always << pool size, so this terminates.""" - out: List[Dict[str, Any]] = [] - seen: set = set() - while len(out) < k: - if self._cursor >= len(self._records): - self.epoch += 1 - np.random.RandomState(self._seed + self.epoch).shuffle(self._records) - self._cursor = 0 - r = self._records[self._cursor] - self._cursor += 1 - if id(r) not in seen: - seen.add(id(r)) - out.append(r) - return out - - -def _select_balanced(buckets: Dict[str, List[Dict[str, Any]]], n_success: int, - n_fail: int, n_fail_loop: int) -> List[Dict[str, Any]]: - """Pick the chunk from the baselined buckets: ``n_fail`` base-fails (split toward - ``n_fail_loop`` loop-fails, best-effort) + ``n_success`` base-successes. If a bucket - is too thin to hit ``chunk_size`` the shortfall is topped up from leftovers (the - ratio then drifts, which the caller logs).""" - loop, wrong, succ = buckets['fail_loop'], buckets['fail_wrong'], buckets['success'] - take_loop = min(n_fail_loop, len(loop)) - take_wrong = min(n_fail - take_loop, len(wrong)) - take_loop = min(n_fail - take_wrong, len(loop)) # give loop the remainder if wrong is short - sel = loop[:take_loop] + wrong[:take_wrong] + succ[:n_success] - target = n_success + n_fail - if len(sel) < target: - used = {id(x) for x in sel} - leftover = [x for b in (loop, wrong, succ) for x in b if id(x) not in used] - sel += leftover[:target - len(sel)] - return sel - - -def _draw_chunk(pool: _ProblemPool, base_sampler, base_dp: int, args: argparse.Namespace - ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: - """Draw one training chunk, running baseline rollout (Phase 1) on every drawn problem. - - With ``--balance`` off, draw ``chunk_size`` problems and return them. With it on, keep - drawing+baselining in ``chunk_size`` batches, bucketing by ``_baseline_class``, until the - target base fail:success mix is reachable or the draw budget is hit; then select a - balanced subset. Returns ``(chunk, stats)`` where stats records the realised mix.""" - if not args.balance: - chunk = pool.draw(args.chunk_size) - n_fresh = _baseline_rollout(base_sampler, chunk, base_dp, args, pool.baseline_cache) - return chunk, {'enabled': False, 'n_drawn': len(chunk), 'n_baseline_fresh': n_fresh} - - n_success = max(0, min(args.chunk_size, round(args.chunk_size * args.balance_success_frac))) - n_fail = args.chunk_size - n_success - n_fail_loop = round(n_fail * args.balance_loop_frac) - buckets: Dict[str, List[Dict[str, Any]]] = {'success': [], 'fail_loop': [], 'fail_wrong': []} - budget = args.chunk_size * args.balance_max_draws_mult - n_drawn, n_fresh = 0, 0 - seen: set = set() # dedupe across batches: the pool can re-serve a problem after a wrap - while n_drawn < budget: - if (len(buckets['success']) >= n_success - and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail): - break # enough of both classes buffered to satisfy the target split - batch = pool.draw(args.chunk_size) - n_fresh += _baseline_rollout(base_sampler, batch, base_dp, args, pool.baseline_cache) - n_drawn += len(batch) - for r in batch: - if id(r) in seen: - continue - seen.add(id(r)) - buckets[_baseline_class(r)].append(r) - - target_reached = (len(buckets['success']) >= n_success - and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail) - chunk = _select_balanced(buckets, n_success, n_fail, n_fail_loop) - sel_success = sum(1 for r in chunk if not r['_failed']) - stats = { - 'enabled': True, 'n_drawn': n_drawn, 'n_baseline_fresh': n_fresh, 'n_selected': len(chunk), - 'target_success': n_success, 'target_fail': n_fail, 'target_fail_loop': n_fail_loop, - 'selected_success': sel_success, 'selected_fail': len(chunk) - sel_success, - 'selected_fail_loop': sum(1 for r in chunk if _baseline_class(r) == 'fail_loop'), - 'selected_fail_wrong': sum(1 for r in chunk if _baseline_class(r) == 'fail_wrong'), - 'selected_success_frac': (sel_success / len(chunk)) if chunk else 0.0, - 'budget_hit': not target_reached, # stopped short of the target mix, not by choice - } - return chunk, stats - - -def process_chunk(base_sampler, skill_sampler, leak: LeakVerifier, - chunk: List[Dict[str, Any]], ci: int, base_dp: int, skill_dp: int, - args: argparse.Namespace, checker=None, - diag_cache: Optional[Dict[str, str]] = None - ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], List[Dict[str, Any]]]: - """base-solve -> rubric-check (view A) -> skill-gen -> leak-filter -> with-skill pass - -> GRPO advantages, for one chunk. - - Sequential (generate-one-chunk-train-one): generation and the trainer's weight - sync never overlap, so no lock is needed. ``base_sampler`` is frozen (never - synced); ``skill_sampler`` is synced by the trainer between chunks. - - ``chunk`` arrives ALREADY baselined by ``_draw_chunk`` (Phase 1 ran during the - balanced draw), so every problem carries ``_init``/``_failed``/``_baseline_pass``/ - ``_hard``/``_cands`` -- Phase 1 is not repeated here. - """ - # Phase 1 (base greedy solve) ran in _draw_chunk so the balancer could classify by - # outcome; every selected problem is processed (no difficulty gate, SEAM-style): the - # group-relative advantage (Phase 6) gives zero gradient to any problem whose skills - # all score alike, so GRPO's own group variance selects the informative problems. - hard = chunk - - # --- Phase 2: assign each problem's view, then rubric-check the view-A attempts so - # the skill model diagnoses from verified findings instead of guessing. View B is - # query-only and deliberately gets NO rubric (nothing to diagnose without an attempt). --- - for r in hard: - r['_view'] = _assign_view(r['problem'], args) - r['_rubric_diag'] = '' - _diagnose_views(checker, hard, args, diag_cache) - - # --- Phase 3: skill-gen (thinking ON), per-view prompt; re-sample empties. --- - flat: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] - if hard: - pending = list(hard) # problems still without any clean candidate - for _ in range(args.skill_retries + 1): - if not pending: - break - prompts = [_view_prompt(r, args) for r in pending] - sg_out = _run_samples(skill_sampler, prompts, args.n_skills, - args.skill_max_tokens, skill_dp, - temperature=args.skill_gen_temperature, - top_p=args.skill_gen_top_p, - top_k=args.skill_gen_top_k) - still = [] - for r, seqs in zip(pending, sg_out): - got = False - for s in seqs: - resp = _clean_text(getattr(s, 'decoded', '') or '') - block = _extract_skills_block(resp) - cand = {'skills': block or '', 'response': resp, 'parseable': bool(block), - 'view': r['_view'], 'leaked': None, 'leak_reason': '', - 'leak_source': '', 'with_pass': None, - 'reward': None, 'rolls': []} - r['_cands'].append(cand) - if block: - flat.append((r, cand)) - got = True - if not got: - still.append(r) # nothing parseable yet -> retry this problem - pending = still - - # --- Phase 4: leak filter via backup teacher (network only, no lock). VIEW A ONLY -- - # view B is query-only (no trajectory to leak from) and is left exactly like SEAM, which - # runs NO leak filter: its candidates skip the check and are treated as clean. To restore - # leak-checking on view B, drop the ``_view == 'A'`` guard below. --- - for r, c in flat: - if r.get('_view') != 'A': - c['leaked'], c['leak_reason'], c['leak_source'] = False, '', 'skipped_viewB' - flat_a = [(r, c) for r, c in flat if r.get('_view') == 'A'] - if flat_a: - details = leak.leak_batch( - [{'content': c['skills'], 'query': r['problem'], 'reference': r['reference_answer']} - for r, c in flat_a], max_workers=args.leak_workers) - for (r, c), d in zip(flat_a, details): - c['leaked'], c['leak_reason'], c['leak_source'] = bool(d.leaked), d.reason, d.source - - # --- Phase 5: with-skill GREEDY pass (T=0, M=1) on clean candidates. Binary reward - # R = answer CORRECT (deterministic, no pass@k noise), ABSOLUTE -- no baseline - # subtraction; the group mean in Phase 6 is the only baseline. Termination is NOT - # required (monitored only) -- see reflexion.md §7.6. --- - clean = [(r, c) for r, c in flat if c['leaked'] is False] - if clean: - ws_out = _run_samples(base_sampler, - [build_skill_solve_prompt(r['problem'], c['skills']) for r, c in clean], - 1, args.max_tokens, base_dp, temperature=0.0) - for (r, c), seqs in zip(clean, ws_out): - c['rolls'] = [_parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll()] - c['with_pass'] = 1.0 if c['rolls'][0]['correct'] else 0.0 - c['reward'] = c['with_pass'] # valid + clean + correct -> 1 - # Validity-in-reward (SEAM-style, --format-in-reward): every candidate that never reached - # the executor -- unparseable/impure format OR answer-leaked -- scores 0 and STILL joins its - # group, so its whole response (think tokens included) is trained DOWN. Off => those - # candidates are excluded, as before. - if args.format_in_reward: - for r in hard: - for c in r['_cands']: - if c['reward'] is None: - c['reward'] = 0.0 - - # --- Phase 6: group-relative GRPO advantage per problem-group. --- - _assign_advantages(hard, args) - - return ([_full_record(r, ci) for r in chunk], - _chunk_summary(chunk, ci, args), - _group_records(chunk, args)) - - -# --------------------------------------------------------------------------- -# Online RFT training -# --------------------------------------------------------------------------- -def _train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: - """Training sample = the exact skill-gen prompt for this record's view + the - generated (think + skills) response as the target; the GRPO advantage is attached - separately at forward_backward time. - - The prompt is rebuilt by ``_skillgen_messages`` (the same function used at generation), - so train/inference stay identical: view A replays problem + rubric findings, view B (and - no-failure view A) replays the query-only prompt. ``key_rounds`` selects the final - assistant turn (index ``len(msgs)``); the plain ``Template`` then masks the prompt and - trains the whole generated response (reasoning + ``</think>`` + skills) -- the key-round - prefix already excludes the prompt-provided ``<think>``, so no extra masking is needed.""" - msgs = _skillgen_messages(rec['problem'], rec.get('view', 'A'), rec.get('diagnosis', '')) - full = msgs + [{'role': 'assistant', 'content': rec['response']}] - return {'messages': full, 'user_data': pack_user_data({'key_rounds': [len(msgs)]})} - - -def _train_chunk(skill_model, ckpt: Optional[CheckpointEngineManager], - samples: List[Dict[str, Any]], args: argparse.Namespace) -> Dict[str, Any]: - """One on-policy GRPO optimizer update on THIS chunk's skill candidates, then sync weights. - - Sequential design (generate-one-chunk-train-one): the skills were sampled from the - current policy and trained immediately, so ``old_logps`` is omitted and the GRPO - ratio is ~1. All driver-side mini-batches accumulate into one optimizer step so the - whole rollout chunk stays under the same pre-update policy. The batch is padded to a - multiple of ``sft_batch_size`` with advantage-0 copies that contribute zero gradient. - """ - trajs = [_train_trajectory(rec) for rec in samples] - advs = [float(rec['advantage']) for rec in samples] - rem = (-len(trajs)) % args.sft_batch_size - if rem: - trajs += [trajs[-1]] * rem # zero-advantage pads -> forward only, no gradient - advs += [0.0] * rem - micro_batches = 0 - for i in range(0, len(trajs), args.sft_batch_size): - skill_model.forward_backward(inputs=trajs[i:i + args.sft_batch_size], - advantages=advs[i:i + args.sft_batch_size]) - micro_batches += 1 - skill_model.clip_grad_and_step() - if ckpt is not None: - ckpt.sync_weights(merge_and_sync=True) - metric = skill_model.calculate_metric(is_training=True) - return {'n_samples': len(samples), 'n_steps': 1, 'n_micro_batches': micro_batches, - 'advantages': [float(rec['advantage']) for rec in samples], - 'metric': {k: (float(v) if _is_num(v) else v) for k, v in (metric or {}).items()}} - - -def _is_num(v: Any) -> bool: - try: - float(v) - return True - except (TypeError, ValueError): - return False - - -# --------------------------------------------------------------------------- -# Main -# --------------------------------------------------------------------------- -def _build_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--dataset', choices=('aops', 'math'), default='aops', - help='Problem source. aops (AI-MO competition problems) is much ' - 'harder than MATH, so the base fails more often -> more offensive ' - 'training signal after balanced sampling.') - p.add_argument('--n', type=int, default=2000, - help='Problems to load into the draw pool (cycled/reshuffled across ' - 'epochs; with --balance many more baseline rollouts than this ' - 'may run, but the pool size is fixed here).') - p.add_argument('--seed', type=int, default=42) - p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True, - help='Keep only answers that collapse to one integer/decimal/fraction, ' - 'matching SEAM numeric reward and avoiding non-scalar grading noise.') - p.add_argument('--eval-size', type=int, default=128, - help='Fixed holdout problems, sampled before the train pool after all ' - 'filters; set 0 to disable fixed eval.') - p.add_argument('--eval-every', type=int, default=10, - help='Run fixed holdout eval every N generation chunks when --eval-size > 0.') - p.add_argument('--chunk-size', type=int, default=16, - help='Problems per generation chunk (all sampler calls batched).') - # -- online baseline-balanced sampling (draw+baseline until the chunk hits the - # target base fail:success mix, so the offensive signal is not starved) -- - p.add_argument('--balance', action=argparse.BooleanOptionalAction, default=True, - help='Keep drawing+baselining problems until the chunk matches the ' - 'target base fail:success composition, then select a balanced ' - 'subset. --no-balance draws chunk_size problems directly.') - p.add_argument('--balance-success-frac', type=float, default=0.4, - help='Target fraction of the chunk that the base solves (base-success). ' - '0.4 => 3:2 fail:success; 0.2 => 4:1. The remainder are base-fail.') - p.add_argument('--balance-loop-frac', type=float, default=0.5, - help='Within the base-fail portion, SOFT target fraction of loop-fails ' - '(ran out of length / never terminated) vs non-loop wrong answers. ' - 'Best-effort only: the fail count is filled from whichever bucket ' - 'is available so a thin bucket never starves the chunk.') - p.add_argument('--balance-max-draws-mult', type=int, default=8, - help='Draw budget per chunk as a multiple of chunk_size; once this many ' - 'problems have been baselined the chunk is assembled from whatever ' - 'the buckets hold (ratio may drift; the actual mix is logged).') - p.add_argument('--n-skills', type=int, default=8, - help='Candidate skills generated per hard problem.') - p.add_argument('--view-b-frac', type=float, default=0.5, - help='Fraction of hard problems routed to view B (query-only, ' - 'deployment form); the rest go to view A (problem + attempt). ' - 'Each problem is assigned to EXACTLY ONE view.') - p.add_argument('--skill-retries', type=int, default=2, - help='Extra skill-gen rounds for a hard problem that yielded no ' - 'clean, parseable candidate (thinking-ON purity gate rejects).') - p.add_argument('--skill-gen-temperature', type=float, default=1.0, - help='Sampling temperature for skill-gen (BOTH views). >0 so the ' - 'n_skills candidates per problem are genuinely DIVERSE — a group ' - 'of near-duplicate skills gives GRPO no real good-vs-bad contrast.') - p.add_argument('--skill-gen-top-p', type=float, default=1.0, - help='top_p for skill-gen; 1.0 keeps the full tail for diversity.') - p.add_argument('--skill-gen-top-k', type=int, default=-1, - help='top_k for skill-gen; -1 disables truncation (max diversity). ' - 'A finite value only narrows the candidate pool.') - p.add_argument('--max-model-len', type=int, default=16384) - p.add_argument('--max-tokens', type=int, default=8192, - help='Max generated tokens for solve rollouts.') - p.add_argument('--skill-max-tokens', type=int, default=8192, - help='Max tokens for skill-gen (thinking ON: the model must close ' - '</think> within this budget or the candidate is dropped, so ' - 'leave ample room).') - p.add_argument('--leak-workers', type=int, default=16, - help='Parallel workers for the LeakVerifier backup judge (capped at 16 ' - 'to avoid the teacher API burst-rate limit; leak and rubric run in ' - 'separate phases so peak teacher concurrency is max(leak,rubric)).') - p.add_argument('--rubric-workers', type=int, default=16, - help='Parallel workers for the view-A rubric diagnose() calls ' - '(teacher-served; requires LLM_BACKUP_* env).') - # -- online GRPO (one on-policy update per generated chunk) -- - p.add_argument('--sft-batch-size', type=int, default=8, - help='Driver-side micro-batch size before the chunk-level optimizer step; ' - 'MUST be a multiple of the training dp size (sliced across dp ranks).') - p.add_argument('--grpo-epsilon', type=float, default=0.2, - help='PPO clip epsilon for GRPOLoss (ratio~1 on-policy, so rarely binds).') - p.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True, - help='Fold output validity into the reward (SEAM-style): unparseable/impure ' - 'or answer-leaked candidates score 0 and join their group to be trained ' - 'DOWN (the whole response, think tokens included). ' - '--no-format-in-reward keeps the reject-and-exclude gate.') - p.add_argument('--lr', type=float, default=1e-5) - p.add_argument('--max-train-rounds', type=int, default=200, - help='Cap on train rounds = trained chunks (also sizes the LR schedule).') - p.add_argument('--save-rounds', type=int, default=50) - p.add_argument('--trend-every', type=int, default=10, - help='Every N chunks, print a [trend] line contrasting the first N ' - 'vs the most recent N chunks (adoption + lift + pos/chunk) ' - 'so the training effect on fresh problems is visible at a glance.') - p.add_argument('--output-dir', default='./output/reflexion_skill_rft') - p.add_argument('--swanlab-project', default='twinkle', - help='swanlab project; logging is skipped when swanlab is not ' - 'installed or SWANLAB_MODE=disabled.') - p.add_argument('--swanlab-exp', default='', - help='swanlab experiment (run) name; empty = auto.') - return p.parse_args() - - -def _trend_line(hist: List[Dict[str, float]], window: int, rounds_done: int) -> Optional[str]: - """Contrast the FIRST ``window`` chunks with the most recent ``window`` chunks so - the online training effect on fresh, never-trained problems is glanceable: if RFT - is working, adoption and lift on recent chunks exceed the early baseline.""" - if len(hist) < 2 * window: - return None # need two non-overlapping windows for a clean before/after - base, rec = hist[:window], hist[-window:] - m = lambda xs, k: sum(h[k] for h in xs) / len(xs) - return (f'[trend] first {window} vs last {window} chunks | ' - f'adopt A {m(base,"aA"):.2f}->{m(rec,"aA"):.2f} ' - f'B {m(base,"aB"):.2f}->{m(rec,"aB"):.2f} | ' - f'lift {m(base,"lift"):+.3f}->{m(rec,"lift"):+.3f} | ' - f'pos/chunk {m(base,"pos"):.1f}->{m(rec,"pos"):.1f} | rounds={rounds_done}') - - -def _query_rows(full: List[Dict[str, Any]]) -> List[Tuple[float, float, float, int, str]]: - """Per hard problem that produced >=1 scored candidate: its no-skill baseline - pass@k, the BEST and MEAN with-skill pass@k over its N skill candidates, the scored - count, and the problem text. Drives both the per-query print and the swanlab passk/* - aggregates.""" - rows = [] - for rec in full: - if rec.get('record_type') != 'problem' or not rec.get('is_hard'): - continue - ps = [c['with_pass'] for c in rec.get('candidates', []) if c.get('with_pass') is not None] - if not ps: - continue - rows.append((rec['baseline_pass'], max(ps), sum(ps) / len(ps), len(ps), rec['problem'])) - return rows - - -def _clean_metric(metric: Optional[Dict[str, Any]]) -> Dict[str, float]: - """Numeric GRPO metrics for swanlab: collapse the duplicate per-group LR to a single - ``lr`` and drop non-numeric fields (e.g. 'total time elapse').""" - out: Dict[str, float] = {} - for k, v in (metric or {}).items(): - if not _is_num(v): - continue - if k.startswith('learning rate'): - if 'group 1' in k: - out['lr'] = float(v) - else: - out[k.replace(' ', '_')] = float(v) - return out - - -def _swan_metrics(summary: Dict[str, Any], log: Optional[Dict[str, Any]], - rows: List[Tuple[float, float, float, int, str]]) -> Dict[str, float]: - """Flat metric dict for swanlab = external reflexion metrics + (when this chunk was - trained) the GRPO built-in metric. acc/adopt/term are only emitted on chunks that had - hard problems, and passk/* only when scored candidates exist, so idle chunks don't dip - the charts to zero.""" - d: Dict[str, float] = { - 'gen/n_hard': summary['n_hard'], 'gen/n_clean': summary['n_clean'], - 'gen/n_leaked': summary['n_leaked'], 'gen/n_train_samples': summary['n_train_samples'], - 'gen/n_reward_pos': summary['n_reward_pos'], - 'gen/n_train_from_fail': summary['n_train_from_fail'], - 'gen/abs_adv_from_fail_frac': summary['abs_adv_from_fail_frac'], - } - bal = summary.get('balance') or {} - if bal.get('enabled'): - d.update({'balance/n_drawn': bal['n_drawn'], - 'balance/n_baseline_fresh': bal['n_baseline_fresh'], - 'balance/selected_success_frac': bal['selected_success_frac'], - 'balance/selected_fail_loop': bal['selected_fail_loop'], - 'balance/selected_fail_wrong': bal['selected_fail_wrong']}) - if summary['n_hard'] > 0: - d.update({ - 'acc/baseline_pass': summary['avg_baseline_pass_on_hard'], - 'acc/withskill_pass': summary['avg_withskill_pass'], - 'acc/lift': summary['avg_lift'], - 'adopt/A': summary['view_A']['adoption_rate'], - 'adopt/B': summary['view_B']['adoption_rate'], - 'term/withskill': summary['termination_rate_withskill'], - }) - if rows: - m = lambda i: sum(r[i] for r in rows) / len(rows) - d.update({'passk/baseline_mean': m(0), 'passk/bestN_mean': m(1), 'passk/avgN_mean': m(2)}) - if log: - d['train/n_steps'] = log['n_steps'] - if 'n_micro_batches' in log: - d['train/n_micro_batches'] = log['n_micro_batches'] - d.update({f'train/{k}': v for k, v in _clean_metric(log.get('metric')).items()}) - return d - - -def _prefix_metrics(metrics: Dict[str, float], prefix: str) -> Dict[str, float]: - return {f'{prefix}/{k}': v for k, v in metrics.items()} - - -def _greedy_eval_metrics(recs: List[Dict[str, Any]], ci: int, rounds: int - ) -> Tuple[Dict[str, Any], Dict[str, float]]: - """Aggregate the greedy holdout into SEAM ``mean@1`` metrics: overall + per-view acc, - the frozen-baseline acc, and their lift -- all single-sample-per-problem means (no - candidate averaging, no pass@k), so acc is directly comparable to SEAM's - ``val-core/math/acc/mean@1`` (correctness only; format/leak not gated).""" - def acc(rs: List[Dict[str, Any]]) -> float: - return sum(1 for x in rs if x['withskill_correct']) / len(rs) if rs else 0.0 - def bacc(rs: List[Dict[str, Any]]) -> float: - return sum(x['baseline_pass'] for x in rs) / len(rs) if rs else 0.0 - A = [x for x in recs if x['view'] == 'A'] - B = [x for x in recs if x['view'] == 'B'] - ws, base = acc(recs), bacc(recs) - summary = { - 'record_type': 'eval_summary', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, - 'n': len(recs), 'n_A': len(A), 'n_B': len(B), - 'acc_mean1': ws, 'baseline_acc_mean1': base, 'lift_mean1': ws - base, - 'acc_A_mean1': acc(A), 'acc_B_mean1': acc(B), - 'format_mean1': (sum(1 for x in recs if x['skill_parseable']) / len(recs)) if recs else 0.0, - 'term_mean1': (sum(1 for x in recs if x['withskill_terminated']) / len(recs)) if recs else 0.0, - } - metrics = { - 'core/math/acc/mean@1': ws, - 'core/math/baseline_acc/mean@1': base, - 'core/math/lift/mean@1': ws - base, - 'core/math/format/mean@1': summary['format_mean1'], - 'core/math/term/mean@1': summary['term_mean1'], - } - if A: - metrics['core/math/acc_A/mean@1'] = summary['acc_A_mean1'] - if B: - metrics['core/math/acc_B/mean@1'] = summary['acc_B_mean1'] - return summary, metrics - - -def _run_greedy_eval(base_sampler, skill_sampler, - eval_records: List[Dict[str, Any]], eval_cache: Dict[str, Dict[str, Any]], - ci: int, rounds: int, base_dp: int, skill_dp: int, - args: argparse.Namespace, checker=None, - diag_cache: Optional[Dict[str, str]] = None - ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], Dict[str, float]]: - """SEAM ``val-core/math/acc/mean@1`` analogue on the fixed holdout: ONE greedy skill per - problem (T=0) injected into ONE greedy base solve (T=0), so acc is a single-sample - pass@1 per problem averaged over problems. Each problem keeps its assigned view; view A - still gets the rubric process-check, view B stays query-only -- the mixed A/B acc is the - deployment number. No leak filter: like SEAM's val, acc scores correctness alone.""" - _baseline_rollout(base_sampler, eval_records, base_dp, args, eval_cache) # frozen greedy baseline - for r in eval_records: - r['_view'] = _assign_view(r['problem'], args) - r['_rubric_diag'] = '' - _diagnose_views(checker, eval_records, args, diag_cache) - sg_out = _run_samples(skill_sampler, [_view_prompt(r, args) for r in eval_records], - 1, args.skill_max_tokens, skill_dp, temperature=0.0) - skills = [] - for seqs in sg_out: - resp = _clean_text(getattr(seqs[0], 'decoded', '') or '') if seqs else '' - skills.append((_extract_skills_block(resp) or '', resp)) - ws_out = _run_samples(base_sampler, - [build_skill_solve_prompt(r['problem'], sk) for r, (sk, _) in zip(eval_records, skills)], - 1, args.max_tokens, base_dp, temperature=0.0) - recs = [] - for r, (sk, sresp), seqs in zip(eval_records, skills, ws_out): - roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - recs.append({ - 'record_type': 'eval_problem', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'view': r['_view'], 'rubric_diag': r.get('_rubric_diag', ''), - 'baseline_pass': r['_baseline_pass'], - 'skill': sk, 'skill_parseable': bool(sk), 'skill_response': sresp, - 'withskill_pred': roll['pred'], 'withskill_correct': roll['correct'], - 'withskill_terminated': roll['terminated'], 'withskill_stop_reason': roll['stop_reason'], - 'withskill_text': roll['text'], - }) - summary, metrics = _greedy_eval_metrics(recs, ci, rounds) - return recs, summary, metrics - - -def _validate_run_config(args: argparse.Namespace, records: List[Dict[str, Any]]) -> None: - """Fail fast on configs that would SILENTLY hang the online sampler: _ProblemPool.draw(k) - dedups within a call, so it never returns unless the pool holds >= chunk_size problems; - a zero draw budget or chunk size yields empty chunks that never advance ``rounds``.""" - if not records: - raise ValueError(f'loaded 0 {args.dataset} problems; check the dataset source') - if args.chunk_size < 1: - raise ValueError(f'--chunk-size must be >= 1 (got {args.chunk_size})') - if args.eval_size < 0: - raise ValueError(f'--eval-size must be >= 0 (got {args.eval_size})') - if args.eval_size > 0 and args.eval_every < 1: - raise ValueError(f'--eval-every must be >= 1 when eval is enabled (got {args.eval_every})') - if len(records) < args.chunk_size: - raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded problems ' - f'({len(records)}); raise --n or lower --chunk-size') - if args.balance_max_draws_mult < 1: - raise ValueError(f'--balance-max-draws-mult must be >= 1 (got {args.balance_max_draws_mult})') - if not 0.0 <= args.balance_success_frac <= 1.0: - raise ValueError(f'--balance-success-frac must be in [0, 1] (got {args.balance_success_frac})') - if not 0.0 <= args.balance_loop_frac <= 1.0: - raise ValueError(f'--balance-loop-frac must be in [0, 1] (got {args.balance_loop_frac})') - - -def main() -> None: - args = _build_args() - if args.sft_batch_size % TRAIN_DP != 0: - raise ValueError(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple ' - f'of the training dp size ({TRAIN_DP})') - # LR schedule now follows chunk-level optimizer updates, not driver micro-batches. - steps_per_round = 1 - records, eval_records, data_stats = _load_records(args) - _validate_run_config(args, records) - os.makedirs(args.output_dir, exist_ok=True) - data_path = os.path.join(args.output_dir, 'skill_dataset.jsonl') - gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') - eval_path = os.path.join(args.output_dir, 'eval_records.jsonl') - train_log_path = os.path.join(args.output_dir, 'train_log.jsonl') - - if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY')): - sys.stderr.write('[rft] WARNING: no LLM_BACKUP_API_KEY/OPENAI_API_KEY — ' - 'LeakVerifier will report no_llm and skip leak filtering\n') - - use_swan = swanlab is not None and os.environ.get('SWANLAB_MODE') != 'disabled' - if use_swan: - swanlab.init(project=args.swanlab_project, - experiment_name=(args.swanlab_exp or None), - config={'model': GEN_MODEL_ID, 'dataset': args.dataset, - 'n': len(records), 'eval_n': len(eval_records), - 'raw_loaded': data_stats['raw_loaded'], - 'numeric_only': args.numeric_only, - 'numeric_dropped': data_stats['numeric_dropped'], - 'n_skills': args.n_skills, 'view_b_frac': args.view_b_frac, - 'balance': args.balance, - 'balance_success_frac': args.balance_success_frac, - 'skill_gen_temp': args.skill_gen_temperature, - 'grpo_epsilon': args.grpo_epsilon, 'lr': args.lr}) - - # -- Device groups: train (FSDP2) + two independent vLLM samplers. -- - r0, r1, r2 = TRAIN_GPUS, TRAIN_GPUS + SKILL_SAMPLER_GPUS, NUM_GPUS - device_groups = [ - DeviceGroup(name='train', ranks=list(range(0, r0)), device_type='GPU'), - DeviceGroup(name='skill_sampler', ranks=list(range(r0, r1)), device_type='GPU'), - DeviceGroup(name='base_sampler', ranks=list(range(r1, r2)), device_type='GPU'), - ] - twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, groups=device_groups, - lazy_collect=False) - - # -- Skill model: full-param FSDP2, GRPO policy update. -- - train_mesh = DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_DP, fsdp_size=TRAIN_FSDP) - skill_model = TransformersModel(model_id=GEN_MODEL_ID, device_mesh=train_mesh, - remote_group='train', - ddp_config={'find_unused_parameters': False}) - from twinkle.patch.no_split_modules import NoSplitModulesPatch - skill_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) - skill_model.set_template(Template, model_id=GEN_MODEL_ID, - enable_thinking=True, max_length=args.max_model_len, - truncation_strategy='delete') - skill_model.set_processor(InputProcessor, padding_free=False) - skill_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) - skill_model.set_optimizer('AdamW', lr=args.lr) - skill_model.set_lr_scheduler('CosineWarmupScheduler', num_warmup_steps=10, - num_training_steps=args.max_train_rounds * steps_per_round) - - # -- Two vLLM samplers: skill (synced) + base (frozen). -- - skill_dp, base_dp = SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS - skill_sampler = vLLMSampler( - model_id=GEN_MODEL_ID, - engine_args={'gpu_memory_utilization': GEN_GPU_MEM, - 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, - device_mesh=DeviceMesh.from_sizes(world_size=SKILL_SAMPLER_GPUS, dp_size=skill_dp), - remote_group='skill_sampler') - skill_sampler.set_template(Template, model_id=GEN_MODEL_ID, - enable_thinking=True, max_length=args.max_model_len) - base_sampler = vLLMSampler( - model_id=GEN_MODEL_ID, - engine_args={'gpu_memory_utilization': GEN_GPU_MEM, - 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, - device_mesh=DeviceMesh.from_sizes(world_size=BASE_SAMPLER_GPUS, dp_size=base_dp), - remote_group='base_sampler') - base_sampler.set_template(Template, model_id=GEN_MODEL_ID, - enable_thinking=True, max_length=args.max_model_len) - - ckpt = CheckpointEngineManager(model=skill_model, sampler=skill_sampler) - leak = LeakVerifier(sampler=None, answer_only=True) # flag ONLY the final answer (view A only; view B skips leak) - # leak = LeakVerifier(sampler=None, judge_system=_LEAK_JUDGE_SYSTEM) # stricter: also flag concrete intermediate key results - checker = _build_rubric_checker() # view-A process-check (teacher-only); None if no LLM backup - if checker is None: - sys.stderr.write('[rft] no LLM backup env -> view-A rubric process-check DISABLED ' - '(skill-gen diagnoses from the attempt alone)\n') - - sys.stderr.write(f'[rft] raw={data_stats["raw_loaded"]} numeric_drop={data_stats["numeric_dropped"]} ' - f'train={len(records)} eval={len(eval_records)} {args.dataset} problems; ' - f'train_gpus={TRAIN_GPUS} skill_dp={skill_dp} base_dp={base_dp}\n') - - # -- Sequential: generate one chunk, train on it, sync -> exact on-policy GRPO. - # Generation dominates wall-clock, so not overlapping training costs little, and - # it removes all producer/consumer concurrency (no thread, no lock). -- - cfg = {'record_type': 'config', 'model': GEN_MODEL_ID, 'dataset': args.dataset, - 'n': len(records), 'eval_n': len(eval_records), 'seed': args.seed, - 'numeric_only': args.numeric_only, - 'raw_loaded': data_stats['raw_loaded'], - 'numeric_dropped': data_stats['numeric_dropped'], - 'eval_every': args.eval_every, - 'n_skills': args.n_skills, - 'view_b_frac': args.view_b_frac, 'skill_retries': args.skill_retries, - 'balance': args.balance, 'balance_success_frac': args.balance_success_frac, - 'balance_loop_frac': args.balance_loop_frac, - 'balance_max_draws_mult': args.balance_max_draws_mult, - 'skill_gen_temp': args.skill_gen_temperature, - 'skill_gen_top_p': args.skill_gen_top_p, 'skill_gen_top_k': args.skill_gen_top_k, - 'reward': 'greedy_binary(correct)', 'advantage': 'group_relative', - 'format_in_reward': args.format_in_reward, - 'rubric_check': 'fixed_math_5crit(viewA)' if checker else 'disabled', - 'grpo_epsilon': args.grpo_epsilon, 'lr': args.lr, - 'max_train_rounds': args.max_train_rounds, 'started': int(time.time())} - hist: List[Dict[str, float]] = [] - rounds = 0 - pool = _ProblemPool(records, args.seed) - eval_cache: Dict[str, Dict[str, Any]] = {} - rubric_cache: Dict[str, str] = {} - eval_rubric_cache: Dict[str, str] = {} - with open(gen_path, 'w', encoding='utf-8') as gen_f, \ - open(eval_path, 'w', encoding='utf-8') as eval_f, \ - open(data_path, 'w', encoding='utf-8') as data_f, \ - open(train_log_path, 'w', encoding='utf-8') as tlog: - for f in (gen_f, eval_f, data_f, tlog): - f.write(json.dumps(cfg, ensure_ascii=False) + '\n') - f.flush() - gstep = 0 - # Each chunk is drawn fresh from the pool (which reshuffles + bumps epoch on every - # full pass) and RE-GENERATED with the current (improved) policy, so every chunk - # stays on-policy (no importance correction) -- the online analogue of SEAM's - # fixed-data epochs. With --balance, _draw_chunk keeps drawing+baselining until the - # base fail:success mix hits the target before this chunk is trained on. - while rounds < args.max_train_rounds: - chunk, balance = _draw_chunk(pool, base_sampler, base_dp, args) - full, summary, groups = process_chunk( - base_sampler, skill_sampler, leak, chunk, gstep, base_dp, skill_dp, - args, checker, rubric_cache) - summary['balance'] = balance - - log = None - if groups: # on-policy GRPO update on this chunk, then weights sync - log = _train_chunk(skill_model, ckpt, groups, args) - rounds += 1 - log.update({'record_type': 'train_round', 'round': rounds, - 'chunk': gstep, 'epoch': pool.epoch, 'ts': int(time.time())}) - tlog.write(json.dumps(log, ensure_ascii=False) + '\n') - tlog.flush() - if rounds % args.save_rounds == 0: - skill_model.save(f'skill-rft-{rounds}', output_dir=args.output_dir) - - summary['rounds_done'], summary['epoch'] = rounds, pool.epoch - for rec in full: - gen_f.write(json.dumps(rec, ensure_ascii=False) + '\n') - gen_f.write(json.dumps(summary, ensure_ascii=False) + '\n') - gen_f.flush() - for v in groups: - data_f.write(json.dumps(v, ensure_ascii=False) + '\n') - data_f.flush() - - sa, sb = summary['view_A'], summary['view_B'] - hist.append({'aA': sa['adoption_rate'], 'aB': sb['adoption_rate'], - 'lift': summary['avg_lift'], 'pos': summary['n_reward_pos']}) - bal_str = (f'bal {balance["selected_fail"]}f/{balance["selected_success"]}s ' - f'(loop {balance["selected_fail_loop"]} drew {balance["n_drawn"]}/' - f'fresh {balance["n_baseline_fresh"]}' - + ('!' if balance.get('budget_hit') else '') + ') ' - ) if balance.get('enabled') else '' - sys.stderr.write( - f'[gen] e{pool.epoch} g{gstep}: {bal_str}hard={summary["n_hard"]} ' - f'clean={summary["n_clean"]} train={summary["n_train_samples"]} ' - f'(fail {summary["n_train_from_fail"]} adv%{summary["abs_adv_from_fail_frac"]:.2f}) ' - f'acc={summary["avg_baseline_pass_on_hard"]:.2f}->{summary["avg_withskill_pass"]:.2f} ' - f'lift={summary["avg_lift"]:+.3f} ' - f'A[{sa["n_hard"]}h {sa["adoption_rate"]:.2f}] ' - f'B[{sb["n_hard"]}h {sb["adoption_rate"]:.2f}] ' - f'rounds={rounds}' - + (f' metric={log.get("metric")}' if log else '') + '\n') - # -- per-query passk (base vs best/avg of N skills) + swanlab metrics -- - rows = _query_rows(full) - for base_p, best_p, avg_p, nsc, prob in rows: - logger.info(f'[q] g{gstep} base={base_p:.2f} bestN={best_p:.2f} avgN={avg_p:.2f} ' - f'n={nsc} | {prob[:70].replace(chr(10), " ")}') - if use_swan: - swanlab.log(_swan_metrics(summary, log, rows), step=gstep) - - if eval_records and (gstep + 1) % args.eval_every == 0: - eval_recs, eval_summary, eval_metrics = _run_greedy_eval( - base_sampler, skill_sampler, eval_records, eval_cache, gstep, - rounds, base_dp, skill_dp, args, checker, eval_rubric_cache) - for rec in eval_recs: - eval_f.write(json.dumps(rec, ensure_ascii=False) + '\n') - eval_f.write(json.dumps(eval_summary, ensure_ascii=False) + '\n') - eval_f.flush() - if use_swan: - swanlab.log(_prefix_metrics(eval_metrics, 'eval'), step=gstep) - sys.stderr.write( - f'[eval] g{gstep}: n={eval_summary["n"]} mean@1 ' - f'acc={eval_summary["baseline_acc_mean1"]:.3f}->{eval_summary["acc_mean1"]:.3f} ' - f'lift={eval_summary["lift_mean1"]:+.3f} ' - f'A[{eval_summary["n_A"]} {eval_summary["acc_A_mean1"]:.3f}] ' - f'B[{eval_summary["n_B"]} {eval_summary["acc_B_mean1"]:.3f}] ' - f'fmt={eval_summary["format_mean1"]:.2f} rounds={rounds}\n') - - if (gstep + 1) % args.trend_every == 0: - tl = _trend_line(hist, args.trend_every, rounds) - if tl: - sys.stderr.write(tl + '\n') - gstep += 1 - - skill_model.save('skill-rft-final', output_dir=args.output_dir) - sys.stderr.write(f'[rft] done: {rounds} train rounds over {gstep} chunks / {pool.epoch} epochs; ' - f'data -> {data_path}\n') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/legacy/train_reflexion_skill_rft.sh b/cookbook/exp/legacy/train_reflexion_skill_rft.sh deleted file mode 100644 index b5e0f1d82..000000000 --- a/cookbook/exp/legacy/train_reflexion_skill_rft.sh +++ /dev/null @@ -1,35 +0,0 @@ -#!/bin/bash -# RFT cold-start for the reflexion skill generator (see reflexion.md §6). -# GPUs: 8 — ranks 0-3 train (skill model, FSDP2), 4-5 skill sampler, 6-7 base sampler. -# Leak filtering uses the backup teacher API (no local judge): set LLM_BACKUP_*. - -set -euo pipefail - -export GEN_MODEL_ID=${GEN_MODEL_ID:-Qwen/Qwen3-4B} -# Local MATH copy (modelscope download cache). Override MATH_DATA_DIR if the -# cache hash dir changes or the data lives elsewhere. -export MATH_DATA_DIR=${MATH_DATA_DIR:-/mnt/workspace/.cache/modelscope/hub/datasets/downloads/extracted/0744cd2d347a7e8f85f7087d950b2ed38b626a5c808c5399e2d8a0923d42d013/MATH} -export LLM_BACKUP_API_KEY=${LLM_BACKUP_API_KEY:?set LLM_BACKUP_API_KEY for the leak judge} -export LLM_BACKUP_BASE_URL=${LLM_BACKUP_BASE_URL:-https://dashscope.aliyuncs.com/compatible-mode/v1} -export LLM_BACKUP_MODEL=${LLM_BACKUP_MODEL:-qwen3.7-max} - -python cookbook/exp/embedding/train_reflexion_skill_rft.py \ - --dataset aops \ - --n 5000 \ - --chunk-size 16 \ - --n-skills 8 \ - --view-b-frac 0.5 \ - --skill-retries 2 \ - --balance \ - --balance-success-frac 0.4 \ - --balance-loop-frac 0.5 \ - --balance-max-draws-mult 8 \ - --max-tokens 25000 \ - --max-model-len 30000 \ - --sft-batch-size 8 \ - --grpo-epsilon 0.2 \ - --lr 6e-6 \ - --max-train-rounds 1500 \ - --save-rounds 25 \ - --trend-every 10 \ - --output-dir ./output/reflexion_skill_rft diff --git a/cookbook/exp/legacy/train_reflexion_skill_seam.py b/cookbook/exp/legacy/train_reflexion_skill_seam.py deleted file mode 100644 index 3bc9d5162..000000000 --- a/cookbook/exp/legacy/train_reflexion_skill_seam.py +++ /dev/null @@ -1,2022 +0,0 @@ -"""Online GRPO RFT for the reflexion skill generator (self-contained, cached). - -Trains an INDEPENDENT skill model to write reusable skills that, injected into a -FROZEN base solver's system prompt, raise its accuracy. The base is never trained; -it only produces the reward. Per chunk: base greedy solve -> rubric process-check -(view A) -> skill-gen (thinking ON, N candidates) -> leak filter -> with-skill -greedy pass -> group-relative advantage -> ONE on-policy GRPO step -> sync weights. -Reward is deterministic (T=0, M=1) binary correctness; A = (R - mean) / (std + eps) -within each problem-group, so std=0 groups give no gradient (GRPO variance selects). - -Each problem is routed to EXACTLY ONE view: A = problem + rubric findings, B = -query only (deployment form). Skill-gen trains only the final structured guidance turn. - -Base greedy rollouts and rubric diagnoses are cached to disk (jsonl, md5-keyed) so -restarts skip them; skill-gen is on-policy and never cached. - -8 GPUs default for high-memory cards: rank 0 trains the actor, rank 1 hosts a -frozen reference model, ranks 2-3 skill_sampler (synced), ranks 4-7 base_sampler -(frozen). Override TRAIN_GPUS / REF_GPUS / SKILL_SAMPLER_GPUS / BASE_SAMPLER_GPUS -for other layouts. -Leak / rubric use the backup teacher API: set LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / -LLM_BACKUP_MODEL. - -Launch: - LLM_BACKUP_API_KEY=... python cookbook/exp/embedding/train_reflexion_skill.py \ - --dataset aops --n 5000 --chunk-size 16 --lr 6e-6 -""" -import argparse -import copy -import hashlib -import json -import math -import os -import re -import sys -import threading -import time -from concurrent.futures import ThreadPoolExecutor -from typing import Any, Dict, List, Optional, Set, Tuple - -import numpy as np - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.checkpoint_engine import CheckpointEngineManager -from twinkle.data_format import SamplingParams, pack_user_data -from twinkle.dataset import Dataset, DatasetMeta -from twinkle.model import TransformersModel -from twinkle.patch.no_split_modules import NoSplitModulesPatch -from twinkle.processor import InputProcessor -from twinkle.sampler import vLLMSampler -from twinkle.template import Template -from twinkle_agentic.verifier import RubricVerifier -from twinkle_agentic.verifier.rubric_verifier import RubricItem - -logger = get_logger() - -try: - import swanlab -except ImportError: - swanlab = None - -MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') -GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.8)) -GEN_TEMPERATURE = float(os.environ.get('GEN_TEMPERATURE', 0.6)) -GEN_TOP_P = float(os.environ.get('GEN_TOP_P', 0.95)) -AOPS_DATASET_ID = os.environ.get('AOPS_DATASET_ID', 'AI-MO/aops') -MATH_DATASET_ID = os.environ.get('MATH_DATASET_ID', 'modelscope/competition_math') - -# GPU layout: train actor (synced to skill_sampler) + frozen ref model + skill_sampler + base_sampler. -# High-memory cards can keep Qwen3-4B actor/ref as one GPU each and spend more GPUs -# on vLLM data-parallel sampling. The base side is heavier here because every -# clean skill candidate is re-solved by the frozen base model, plus balance/eval baselines. -TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 2)) -REF_GPUS = int(os.environ.get('REF_GPUS', 2)) -SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) -BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 2)) -NUM_GPUS = TRAIN_GPUS + REF_GPUS + SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS -TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', min(1, TRAIN_GPUS))) -REF_FSDP = int(os.environ.get('REF_FSDP', min(1, REF_GPUS))) -if min(TRAIN_GPUS, REF_GPUS, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS, TRAIN_FSDP, REF_FSDP) < 1: - raise ValueError('TRAIN_GPUS, REF_GPUS, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS, TRAIN_FSDP and REF_FSDP must all be >= 1') -if TRAIN_GPUS % TRAIN_FSDP != 0: - raise ValueError(f'TRAIN_GPUS ({TRAIN_GPUS}) must be divisible by TRAIN_FSDP ({TRAIN_FSDP})') -if REF_GPUS % REF_FSDP != 0: - raise ValueError(f'REF_GPUS ({REF_GPUS}) must be divisible by REF_FSDP ({REF_FSDP})') -TRAIN_DP = TRAIN_GPUS // TRAIN_FSDP -REF_DP = REF_GPUS // REF_FSDP - - -# =========================================================================== -# Block A -- boxed extraction + answer grading -# =========================================================================== -_BOXED_RE = re.compile(r'\\boxed\s*\{') - - -def extract_boxed(text: str) -> Optional[str]: - """Last ``\\boxed{...}`` content, brace-balanced.""" - if not text: - return None - last = None - for m in _BOXED_RE.finditer(text): - depth, i = 1, m.end() - while i < len(text) and depth > 0: - depth += (text[i] == '{') - (text[i] == '}') - i += 1 - if depth == 0: - last = text[m.end():i - 1].strip() - return last - - -def normalize_answer(ans: str) -> str: - if not ans: - return '' - s = str(ans).strip() - m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) - if m: - return m.group(1) - s = s.strip('$').strip().replace('−', '-') - 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').replace(r'\tfrac', r'\frac') - s = s.replace(r'\,', '').replace(r'\;', '').replace(r'\!', '') - s = re.sub(r'\\(?:quad|qquad|\s)', '', s) - s = re.sub(r'\s+', '', s) - s = re.sub(r'\\sqrt([0-9A-Za-z])', r'\\sqrt{\1}', s) - s = re.sub(r'\\frac\{([^{}]+)\}([^{}\\])', r'\\frac{\1}{\2}', s) - s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) - s = re.sub(r'\{[a-zA-Z]+\}$', '', s) - s = re.sub(r'\^\\circ|\^\{\\circ\}|°|\\circ', 'deg', s) - s = s.replace(r'\minus{}', '-').replace(r'\minus', '-') - - def _frac_to_slash(mt): - text = mt.group(0) - pos = text.index('{') + 1 - depth, num_start = 1, pos - while depth > 0: - depth += (text[pos] == '{') - (text[pos] == '}') - pos += 1 - numer = text[num_start:pos - 1] - pos += 1 - den_start, depth = pos, 1 - while depth > 0: - depth += (text[pos] == '{') - (text[pos] == '}') - pos += 1 - return f'({numer})/({text[den_start:pos - 1]})' - - s = re.sub(r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', _frac_to_slash, s) - s = re.sub(r'(?<!\w)(\d+)/(\d+)(?!\w)', r'(\1)/(\2)', s) - return s - - -def _try_numeric_equal(a: str, b: str) -> bool: - try: - va, vb = float(a.replace('(', '').replace(')', '')), 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) - return va is not None and vb is not None and abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) - - -_MCQ_REF_RE = re.compile( - r'^\\(?:textbf|text|mathrm|mathbf)\{\(?([A-E])\)?\s*\}\s*(.+)$' - r'|^\(?([A-E])\)\s+(.+)$') -_VAR_PREFIX_RE = re.compile(r'^(?:[A-Za-z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)$', re.DOTALL) - - -def _split_mcq(ans: str): - s = ans.strip() - m = _MCQ_REF_RE.match(s) - if m: - return (m.group(1) or m.group(3)), ((m.group(2) or m.group(4) or '').strip() or None) - bl = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) - return (bl.group(1), None) if bl else (None, s or None) - - -def _strip_var_prefix(ans: str) -> str: - m = _VAR_PREFIX_RE.match((ans or '').strip()) - return m.group(1).strip() if m else (ans or '') - - -def _normalize_tuple(ans: str) -> str: - return re.sub(r'[\s()\[\]{}\\]', '', ans or '') - - -def _math_verify_equal(predicted: str, reference: str) -> bool: - try: - from math_verify import parse, verify - gold = parse(r'\boxed{' + reference + '}', parsing_timeout=1) - pred = parse(r'\boxed{' + predicted + '}', parsing_timeout=1) - return bool(gold and pred and verify(gold, pred, timeout_seconds=1)) - except Exception: - return False - - -def answers_match(predicted: str, reference: str) -> bool: - if not predicted or not reference: - return False - norm_p, norm_r = normalize_answer(predicted), normalize_answer(reference) - if norm_p == norm_r or norm_p.lower() == norm_r.lower() or _try_numeric_equal(norm_p, norm_r): - return True - stripped_p = normalize_answer(_strip_var_prefix(predicted)) - stripped_r = normalize_answer(_strip_var_prefix(reference)) - if stripped_p and stripped_r: - if (stripped_p == stripped_r or stripped_p.lower() == stripped_r.lower() - or _try_numeric_equal(stripped_p, stripped_r)): - return True - p_letter, p_value = _split_mcq(predicted) - r_letter, r_value = _split_mcq(reference) - if p_letter and r_letter and p_letter == r_letter: - return True - if (p_letter and p_value is None) and (r_letter and r_value is None): - return False - 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 and (p_val == r_val or p_val.lower() == r_val.lower() or _try_numeric_equal(p_val, r_val)): - return True - for left, right in ((norm_p, norm_r), (stripped_p, stripped_r)): - tl, tr = _normalize_tuple(left), _normalize_tuple(right) - if ',' in tl and tl == tr: - return True - if '=' in norm_r and '=' not in norm_p: - if any(part and (part == norm_p or _try_numeric_equal(part, norm_p)) for part in norm_r.split('=')): - return True - if '=' in norm_p and '=' not in norm_r: - if any(part and (part == norm_r or _try_numeric_equal(part, norm_r)) for part in norm_p.split('=')): - return True - return _math_verify_equal(stripped_p or norm_p, stripped_r or norm_r) - - -# =========================================================================== -# Block B -- prompts, skill parsing, batched sampling -# =========================================================================== -DIRECT_SYSTEM = ( - 'You are an expert competition mathematician. Solve the following problem ' - 'step by step. Provide your final answer inside \\boxed{}.') - -_SKILL_SOLVE_PREFIX = ( - DIRECT_SYSTEM + '\n\n' - 'Before you start, keep these reminders in mind to avoid common mistakes on this ' - 'type of problem:\n') -_SKILL_SOLVE_SUFFIX = '\nApply them where relevant, but rely on your own reasoning to reach the answer.' - - -def build_direct_prompt(problem: str) -> Dict[str, Any]: - return {'messages': [{'role': 'system', 'content': DIRECT_SYSTEM}, - {'role': 'user', 'content': problem}]} - - -def build_skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: - skill = (skill or '').strip() - if not skill: - return build_direct_prompt(problem) - # Concatenation (not .format): DIRECT_SYSTEM/skill contain literal braces. - return {'messages': [ - {'role': 'system', 'content': _SKILL_SOLVE_PREFIX + skill + _SKILL_SOLVE_SUFFIX}, - {'role': 'user', 'content': problem}]} - - -# -- skill-gen prompts (view A: problem + rubric source; view B: query only) -- -# Kept deliberately short: this is the RL policy's system prompt, so over-specifying -# the output hurts convergence. The concrete output format is appended separately by -# ``_SKILL_OUTPUT`` so the task text stays format-agnostic. -SKILL_GEN_SYSTEM = ( - 'You are a math guidance writer. A process-check on a related problem hints at ' - 'likely mistakes. Write short reusable guidance for this and similar problems, ' - 'and note what to watch out for.\n') - -SKILL_GEN_SYSTEM_Q = ( - 'You are a math guidance writer. Write short reusable guidance for this and ' - 'similar problems.\n') - -_SKILL_OUTPUT = ( - 'Output only:\n<skills>\nYour reusable solving guidance here.\n</skills>') - -SKILL_GEN_USER_Q = ( - 'Problem:\n{problem}\n\n') - -SKILL_GEN_USER_RUBRIC = ( - 'Target problem:\n{problem}\n\n' - 'Problem used for the process check:\n{rubric_problem}\n\n' - 'Process check:\n' - '{diagnosis}\n\n') - - -def _rubric_has_fail(diagnosis: str) -> bool: - """View A uses the open-book rubric prompt (and, under --viewa-sft, SFT distillation) - IFF its rubric localised a failure. No ``[FAIL]`` => nothing to inject: generation - degrades to query-only and the problem is trained by GRPO exactly like view B. Single - source of truth so ``_skillgen_messages`` and ``_group_records`` never diverge.""" - return '[FAIL]' in (diagnosis or '') - - -def _skillgen_messages(problem: str, view: str, diagnosis: str, - rubric_problem: str = '') -> List[Dict[str, Any]]: - """Single source of truth for the skill-gen prompt (used at BOTH generation and - training so they never diverge). View A with a localisable failure uses the target - problem plus the rubric source problem and findings; view B -- or a view-A problem - whose rubric flagged nothing (no ``[FAIL]``) -- degrades to the query-only prompt.""" - if view == 'B' or not _rubric_has_fail(diagnosis): - return [{'role': 'system', 'content': SKILL_GEN_SYSTEM_Q + _SKILL_OUTPUT}, - {'role': 'user', 'content': SKILL_GEN_USER_Q.format(problem=problem)}] - rubric_problem = rubric_problem or problem - return [{'role': 'system', 'content': SKILL_GEN_SYSTEM + _SKILL_OUTPUT}, - {'role': 'user', 'content': SKILL_GEN_USER_RUBRIC.format( - problem=problem, rubric_problem=rubric_problem, diagnosis=diagnosis)}] - - -def _assign_view(problem: str, args: argparse.Namespace) -> str: - h = int(hashlib.md5(f'{args.seed}:{problem}'.encode('utf-8')).hexdigest(), 16) - return 'B' if (h % 100000) / 100000.0 < args.view_b_frac else 'A' - - -def _view_prompt(r: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: - return {'messages': _skillgen_messages( - r['problem'], r['_view'], r.get('_rubric_diag', ''), r.get('_rubric_src', ''))} - - -_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') - - -def _clean_text(decoded: Optional[str]) -> str: - return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() - - -def _extract_tag_block(answer: str, tag: str, allow_empty: bool = False) -> Optional[str]: - low = answer.lower() - open_tag, close_tag = f'<{tag}>', f'</{tag}>' - s = low.rfind(open_tag) - if s < 0: - return None - inner = s + len(open_tag) - e = low.find(close_tag, inner) - if e < 0: - return None - block = answer[inner:e].strip() - block = re.sub(r'</?(?:skills|skill|diagnose|pitfall|strategy|think)>', '', block, flags=re.IGNORECASE).strip() - return block if (block or allow_empty) else None - - -def _extract_skill(text: str) -> Optional[str]: - """Parse skill-generation output: return the inner text of a non-empty ``<skills>`` - block, or None. If a ``</think>`` marker is present, parse only the text after the - last one; otherwise parse the full response.""" - low = text.lower() - end_think = low.rfind('</think>') - answer = text[end_think + len('</think>'):] if end_think >= 0 else text - return _extract_tag_block(answer, 'skills') - - -def _parse_seq(seq, gold: str) -> Dict[str, Any]: - """Grade one sampled sequence into a rollout record.""" - text = _clean_text(getattr(seq, 'decoded', '') or '') - pred = extract_boxed(text) - correct = bool(pred) and answers_match(pred, gold) - terminated = getattr(seq, 'stop_reason', None) != 'length' - return {'pred': pred, 'correct': correct, 'terminated': terminated, - 'passed': bool(correct and terminated), - 'stop_reason': getattr(seq, 'stop_reason', None), - 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} - - -def _run_samples(sampler, prompts: List[Any], num_samples: int, max_tokens: int, - gen_dp: int, temperature: Optional[float] = None, - top_p: Optional[float] = None, top_k: Optional[int] = None) -> List[List[Any]]: - """One batched sampler call -> per-prompt list of raw sequences. vLLM dp needs - batch len >= dp, so pad the tail and slice back.""" - if not prompts: - return [] - params = SamplingParams( - max_tokens=max_tokens, - temperature=GEN_TEMPERATURE if temperature is None else temperature, - top_p=GEN_TOP_P if top_p is None else top_p, - num_samples=num_samples, **({} if top_k is None else {'top_k': top_k})) - padded = prompts - if gen_dp > 1 and 0 < len(prompts) < gen_dp: - padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] - responses = sampler.sample(padded, params)[:len(prompts)] - return [list(r.sequences) if (r and r.sequences) else [] for r in responses] - - -# =========================================================================== -# Block C -- data loading via twinkle.Dataset + numeric filtering -# =========================================================================== -def _boxed_batch(rows: Dict[str, List[Any]], dataset: str) -> Dict[str, List[Any]]: - """Batched (columnar) mapper for ``Dataset.map``: annotate each row with its boxed - ``reference_answer`` and a ``_keep`` flag (aops also needs the ``boxed`` metadata).""" - sols = rows['solution'] - metas = rows.get('metadata', [None] * len(sols)) - refs = [extract_boxed(s or '') for s in sols] - keep = [bool(ref) and (dataset != 'aops' or bool((meta or {}).get('boxed'))) - for ref, meta in zip(refs, metas)] - return {**rows, 'reference_answer': refs, '_keep': keep} - - -def load_problems(dataset: str, n: int, seed: int, num_proc: int = 0) -> List[Dict[str, Any]]: - """Load boxed-answer problems as ``{problem, reference_answer, level?}`` via - twinkle.Dataset (ModelScope hub), sampled to ``n`` (0 = all). Boxed extraction (regex - + brace scan over every solution) is parallelised by ``Dataset.map``/``.filter``; - ``num_proc`` defaults to all cores (set 1 to force serial).""" - ds_id = AOPS_DATASET_ID if dataset == 'aops' else MATH_DATASET_ID - ds = Dataset(DatasetMeta(dataset_id=f'ms://{ds_id}', split='train')) - nproc = num_proc if num_proc > 0 else min(32, os.cpu_count() or 1) - ds.map(lambda rows: _boxed_batch(rows, dataset), num_proc=nproc) - ds.filter(lambda row: row['_keep'], num_proc=nproc) - has_level = 'level' in ds.dataset.column_names - out = [{'data_id': f'{dataset}:{i}', 'problem': row['problem'], - 'reference_answer': row['reference_answer'], - **({'level': row['level']} if has_level and row.get('level') else {})} - for i, row in enumerate(ds.dataset)] - logger.info(f'[data] {dataset}: {len(out)} boxed problems ({nproc} procs)') - rng = np.random.RandomState(seed) - rng.shuffle(out) - return out[:n] if (n and n < len(out)) else out - - -# --------------------------------------------------------------------------- -# Bag-of-words neighbour pairing (cross-problem rubric transfer, --xproblem-rubric) -# --------------------------------------------------------------------------- -# Common English + math-scaffolding words that carry no problem-type signal. Kept small -# and deterministic on purpose (no external stopword list): what survives is the domain -# vocabulary ('triangle', 'prime', 'polynomial', ...) that actually defines the type. -_BOW_STOP = frozenset(""" -a an the of to in on at for and or but if is are be was were been being this that these those -with without into onto from by as it its their his her our your my we you they he she them -find compute determine calculate evaluate solve show prove given let suppose consider assume -what which when where how many much value values number numbers expression form terms term -such that then than so if only when each every all any some both one two three four five six -seven eight nine ten first second third last non over under about above below between -problem answer result equal equals sum difference product total following there here have has -had do does did can could will would should may might must not no yes if then else -""".split()) - -_WORD_RE = re.compile(r'[a-z]+') - - -def _stem(w: str) -> str: - """Crude suffix stripper so 'prime'/'primes', 'triangle'/'triangles' collapse to one - type token. Not linguistically correct -- just enough to merge the common plural/gerund - variants that otherwise split a type's vocabulary and starve the df filter.""" - if len(w) > 4 and w.endswith('ies'): - return w[:-3] + 'y' # properties -> property - if len(w) > 4 and w.endswith('es') and w[-3] in 'sxzh': - return w[:-2] # boxes -> box (keep primes -> prime below) - for suf in ('ing', 'ed', 's'): - if len(w) > len(suf) + 2 and w.endswith(suf): - return w[:-len(suf)] - return w - - -def _tokenize(problem: str) -> List[str]: - """Deterministic bag-of-words tokens for type matching: lowercase alphabetic words - (numbers dropped -- they are instance detail, not type), minus generic stopwords, then - stemmed, so only domain terms remain. Latex control words (frac, sqrt, ...) survive.""" - return [_stem(w) for w in _WORD_RE.findall((problem or '').lower()) - if len(w) > 2 and w not in _BOW_STOP] - - -class BagOfWordsIndex: - """TF-IDF cosine nearest-neighbour over problem statements. Sparse dict vectors + - an inverted index, so ``nearest`` scores only problems sharing a term (near-linear in - practice, no dense NxN). Deterministic; ties break on lower index for reproducibility. - - Answer-aware: ``nearest`` skips the same index, near-verbatim duplicates (cosine - >= ``sim_max``) and -- crucially for anti-leak -- any candidate whose answer equals the - query's, so a neighbour rubric can never hand over the query's own answer.""" - - def __init__(self, problems: List[str], answers: Optional[List[str]] = None, - min_df: int = 2, max_df_frac: float = 0.5): - self._toks = [_tokenize(p) for p in problems] - self._ans = [(_numeric_value(a) or (str(a).strip() if a else '')) for a in answers] \ - if answers is not None else [''] * len(problems) - n = len(self._toks) - df: Dict[str, int] = {} - for toks in self._toks: - for w in set(toks): - df[w] = df.get(w, 0) + 1 - max_df = max(min_df, int(max_df_frac * n)) - self._idf = {w: math.log((n + 1) / (c + 1)) + 1.0 - for w, c in df.items() if min_df <= c <= max_df} - self._vecs: List[Dict[str, float]] = [self._vectorize(t) for t in self._toks] - self._inverted: Dict[str, List[int]] = {} - for i, v in enumerate(self._vecs): - for w in v: - self._inverted.setdefault(w, []).append(i) - - def _vectorize(self, toks: List[str]) -> Dict[str, float]: - tf: Dict[str, float] = {} - for w in toks: - if w in self._idf: - tf[w] = tf.get(w, 0.0) + 1.0 - vec = {w: c * self._idf[w] for w, c in tf.items()} - norm = math.sqrt(sum(x * x for x in vec.values())) - return {w: x / norm for w, x in vec.items()} if norm > 0 else {} - - def nearest(self, i: int, sim_max: float = 0.98) -> Tuple[int, float]: - """(index, cosine) of the most similar *distinct* problem: not i, not a duplicate - (cosine < ``sim_max``) and not the same answer. (-1, 0.0) if none qualifies.""" - vi = self._vecs[i] - if not vi: - return -1, 0.0 - ai = self._ans[i] - scores: Dict[int, float] = {} - for w, xi in vi.items(): - for j in self._inverted.get(w, ()): - if j != i: - scores[j] = scores.get(j, 0.0) + xi * self._vecs[j].get(w, 0.0) - best_j, best_s = -1, 0.0 - for j, s in scores.items(): - if s >= sim_max or (ai and self._ans[j] == ai): - continue - if s > best_s or (s == best_s and (best_j < 0 or j < best_j)): - best_j, best_s = j, s - return best_j, best_s - - -def build_pairs(records: List[Dict[str, Any]], n: int, seed: int, sim_max: float = 0.98 - ) -> Tuple[List[Dict[str, Any]], Dict[str, Tuple[str, float]]]: - """Single-pass cross-problem pairing over the whole pool (one index build). - - Returns ``(subset, neighbour_map)`` where ``subset`` is the ``n`` problems with the - strongest qualifying neighbour (dense same-type pairs; n<=0/n>=len keeps all, shuffled) - and ``neighbour_map[q] = (p, cosine)`` gives each kept problem its analogue P. P is drawn - from the *full* pool (richer analogues) but never a duplicate or same-answer problem, so - P's rubric can transfer method without ever leaking Q's answer.""" - index = BagOfWordsIndex([r['problem'] for r in records], - [str(r.get('reference_answer', '')) for r in records]) - nbr = [index.nearest(i, sim_max) for i in range(len(records))] - order = sorted(range(len(records)), key=lambda i: (-nbr[i][1], i)) - keep = order[:n] if (0 < n < len(records)) else list(range(len(records))) - rng = np.random.RandomState(seed) - rng.shuffle(keep) - subset = [records[i] for i in keep] - neighbour_map = {records[i]['problem']: (records[nbr[i][0]]['problem'], nbr[i][1]) - for i in keep if nbr[i][0] >= 0} - return subset, neighbour_map - - -_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') - - -def _norm_num_text(num: str) -> str: - try: - f = float(num) - return str(int(f)) if f == int(f) else str(f) - except Exception: - return str(num).strip() - - -def _numeric_value(raw: Any) -> Optional[str]: - """Collapse an answer to a single int/decimal/fraction, or None.""" - if raw is None: - return None - s = str(raw).strip().strip('$').strip() - s = s.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac') - s = re.sub(r'\\!|\\,|\\;|\\ |\\left|\\right|\s', '', s) - for pat in (r'\\frac\{(-?\d+)\}\{(-?\d+)\}', r'(-?\d+)/(-?\d+)'): - m = re.fullmatch(pat, s) - if m: - a, b = int(m.group(1)), int(m.group(2)) - return _norm_num_text(str(a / b)) if b else None - return _norm_num_text(s) if _NUM_RE.fullmatch(s) else None - - -def _answer_leaked(skill: str, reference: str) -> bool: - """Audit whether a generated skill contains the final answer verbatim. This is NOT - a training filter: if the skill model derives an answer from the problem, that is a - legitimate answer-bearing skill under this experiment. The real leakage boundary is the - external API/rubric diagnosis, which is constrained by prompt to stay answer-free.""" - if not skill: - return False - for cand in {_numeric_value(reference), (str(reference).strip() or None)}: - if cand and re.search(r'(?<![\d.])' + re.escape(cand) + r'(?![\d.])', skill): - return True - return False - - -def _load_excluded_records(paths_arg: str) -> Tuple[Set[str], Set[str]]: - """Read jsonl files and collect data_id/problem keys that must be excluded. - - The cold-start SFT builder writes stable ``data_id`` values. ``problem`` is kept as a - backward-compatible fallback for older jsonl files produced before data_id existed.""" - ids: Set[str] = set() - problems: Set[str] = set() - for raw_path in (paths_arg or '').split(','): - path = raw_path.strip() - if not path or not os.path.exists(path): - continue - with open(path, encoding='utf-8') as f: - for line in f: - if not line.strip(): - continue - row = json.loads(line) - if row.get('record_type') in {'config', 'summary'}: - continue - data_id = str(row.get('data_id') or '').strip() - problem = str(row.get('problem') or '').strip() - if data_id: - ids.add(data_id) - elif problem: - problems.add(problem) - return ids, problems - - -def _load_seam_parquet(path: str) -> List[Dict[str, Any]]: - """Read a SEAM ``build_aops_dataset.py`` parquet (VERL RLHF schema) into twinkle records, - PRESERVING file row order. ``problem <- extra_info.problem`` and - ``reference_answer <- reward_model.ground_truth``. No shuffle/filter: - the parquet is already SEAM's numeric-filtered, seed-42-shuffled, truncated split.""" - import pyarrow.parquet as pq - rows = pq.read_table(path).to_pylist() - out: List[Dict[str, Any]] = [] - for i, r in enumerate(rows): - ei = r.get('extra_info') or {} - rm = r.get('reward_model') or {} - problem = (ei.get('problem') or '').strip() - ref = rm.get('ground_truth') - if not problem or ref is None: - continue - out.append({'problem': problem, 'reference_answer': str(ref), - 'data_id': f"seam:{ei.get('split', '')}:{ei.get('index', i)}"}) - return out - - -def _load_records_from_seam(args: argparse.Namespace, seam_dir: str - ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], - Dict[str, Tuple[str, float]], Dict[str, str], Dict[str, int]]: - """Data entry that mirrors a SEAM run EXACTLY: read ``train.parquet``/``val.parquet`` from - ``seam_dir`` in file order, use ``val`` as the eval holdout, take the first ``--n`` train rows - (post ``--pool-offset``) with NO shuffle. ``--numeric-only``/``--eval-size``/internal shuffle - are bypassed (the parquet is already the authoritative split).""" - tp, vp = os.path.join(seam_dir, 'train.parquet'), os.path.join(seam_dir, 'val.parquet') - if not (os.path.exists(tp) and os.path.exists(vp)): - raise FileNotFoundError( - f'--seam-parquet-dir needs both train.parquet and val.parquet in {seam_dir}') - if args.xproblem_rubric: - raise ValueError('--xproblem-rubric is unsupported with --seam-parquet-dir ' - '(SEAM parquet carries no neighbour structure).') - pool = _load_seam_parquet(tp) # already SEAM-shuffled + truncated, in file order - eval_records = [dict(r) for r in _load_seam_parquet(vp)] # SEAM's exact val holdout - pool_offset = max(0, int(getattr(args, 'pool_offset', 0) or 0)) - if pool_offset: - if pool_offset >= len(pool): - raise ValueError(f'--pool-offset ({pool_offset}) leaves no train records ' - f'from SEAM train pool size {len(pool)}') - pool = pool[pool_offset:] - train_n = args.n if args.n > 0 else len(pool) - train_records = [dict(r) for r in pool[:train_n]] - if {r['problem'] for r in train_records} & {r['problem'] for r in eval_records}: - raise ValueError('eval/train overlap detected in SEAM parquet') - stats = {'raw_loaded': len(pool) + len(eval_records), 'numeric_dropped': 0, - 'excluded_records': 0, 'pool_offset': pool_offset, - 'train_records': len(train_records), 'eval_records': len(eval_records), - 'source': 'seam_parquet', 'seam_parquet_dir': seam_dir} - return train_records, eval_records, {}, {}, stats - - -def _load_records(args: argparse.Namespace - ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], - Dict[str, Tuple[str, float]], Dict[str, str], Dict[str, int]]: - """Load, numeric-filter, shuffle, split a fixed eval holdout, and (when --xproblem-rubric) - select a same-type-dense train subset with its neighbour map -- all in one pass. - Also returns ``pool_answers`` (every candidate neighbour P's true answer) so P's baseline - can be graded/cached correctly even when P is not itself a training problem.""" - seam_dir = (getattr(args, 'seam_parquet_dir', '') or '').strip() - if seam_dir: # read SEAM parquet in file order, bypassing load/filter/shuffle/split - return _load_records_from_seam(args, seam_dir) - # Load all when filtering or splitting (else the eval holdout could starve train). - load_n = 0 if (args.numeric_only or args.eval_size > 0) else args.n - records = load_problems(args.dataset, load_n, args.seed) - raw_n, dropped = len(records), 0 - if args.numeric_only: - records = [{**r, 'reference_answer': v} - for r, v in ((r, _numeric_value(r.get('reference_answer'))) for r in records) - if v is not None] - dropped = raw_n - len(records) - np.random.RandomState(args.seed).shuffle(records) - exclude_ids, exclude_problems = _load_excluded_records(getattr(args, 'exclude_data_ids', '')) - excluded = 0 - if exclude_ids or exclude_problems: - before = len(records) - records = [r for r in records - if str(r.get('data_id', '')) not in exclude_ids - and str(r.get('problem', '')).strip() not in exclude_problems] - excluded = before - len(records) - eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 - eval_records = [dict(r) for r in records[:eval_n]] - pool = records[eval_n:] - pool_offset = max(0, int(getattr(args, 'pool_offset', 0) or 0)) - if pool_offset: - if pool_offset >= len(pool): - raise ValueError(f'--pool-offset ({pool_offset}) leaves no train records from pool size {len(pool)}') - pool = pool[pool_offset:] - train_n = args.n if args.n > 0 else len(pool) - # Cross-problem rubric: pick the train_n problems with the strongest qualifying neighbour - # (dense same-type pairs) and build {Q -> (P, sim)} in one index pass. Otherwise keep the - # first train_n (already shuffled) with no neighbours. - if args.xproblem_rubric: - subset, neighbor_map = build_pairs(pool, train_n, args.seed) - train_records = [dict(r) for r in subset] - pool_answers = {r['problem']: str(r.get('reference_answer', '')) for r in pool} - else: - train_records, neighbor_map, pool_answers = [dict(r) for r in pool[:train_n]], {}, {} - if {r['problem'] for r in train_records} & {r['problem'] for r in eval_records}: - raise ValueError('eval/train overlap detected') - stats = {'raw_loaded': raw_n, 'numeric_dropped': dropped, - 'excluded_records': excluded, 'pool_offset': pool_offset, - 'train_records': len(train_records), 'eval_records': len(eval_records)} - return train_records, eval_records, neighbor_map, pool_answers, stats - - -# =========================================================================== -# Block D -- disk cache, problem pool, baseline rollout, rubric check -# =========================================================================== -class DiskCache: - """Append-only jsonl kv cache (md5 key -> value). Loads on init, appends on put. - Disabled instances always miss and never write.""" - - def __init__(self, path: str, enabled: bool = True): - self._mem: Dict[str, Any] = {} - self._fh = None - self._lock = threading.Lock() # base baseline is prefetched on a background thread - if not enabled: - return - if os.path.exists(path): - with open(path, encoding='utf-8') as f: - for line in f: - if line.strip(): - row = json.loads(line) - self._mem[row['key']] = row['value'] - self._fh = open(path, 'a', encoding='utf-8') - - @staticmethod - def key_for(*parts: str) -> str: - return hashlib.md5('\x1f'.join(parts).encode('utf-8')).hexdigest() - - def __contains__(self, key: str) -> bool: - with self._lock: - return key in self._mem - - def get(self, key: str) -> Any: - with self._lock: - return self._mem.get(key) - - def put(self, key: str, value: Any) -> None: - with self._lock: - self._mem[key] = value - if self._fh is not None: - self._fh.write(json.dumps({'key': key, 'value': value}, ensure_ascii=False) + '\n') - self._fh.flush() - - def close(self) -> None: - if self._fh is not None: - self._fh.close() - - -class _LockedSampler: - """Serialises every ``.sample()`` on a sampler behind one lock. The base sampler is - shared by the main thread (with-skill scoring, eval) and the baseline-prefetch thread; - ``sample`` is a slice_dp/flatten remote call whose collect is NOT safe to interleave - across two callers, so concurrent calls could mis-join sequences. The lock keeps base - calls serial (prefetch still overlaps the skill-gen phase, which uses skill_sampler).""" - - def __init__(self, sampler): - self._sampler = sampler - self._lock = threading.Lock() - - def sample(self, *args, **kwargs): - with self._lock: - return self._sampler.sample(*args, **kwargs) - - def __getattr__(self, name): - return getattr(self._sampler, name) - - -class ProblemPool: - """Cyclic draw source. Each full pass reshuffles with ``seed + epoch``; the initial - pass keeps the loader order. ``draw(k)`` returns k distinct problems (k << pool).""" - - def __init__(self, records: List[Dict[str, Any]], seed: int): - self._records = list(records) - self._seed, self._cursor, self.epoch = seed, 0, 0 - - def draw(self, k: int) -> List[Dict[str, Any]]: - out, seen = [], set() - while len(out) < k: - if self._cursor >= len(self._records): - self.epoch += 1 - np.random.RandomState(self._seed + self.epoch).shuffle(self._records) - self._cursor = 0 - r = self._records[self._cursor] - self._cursor += 1 - if id(r) not in seen: - seen.add(id(r)) - out.append(r) - return out - - def peek(self, k: int) -> List[Dict[str, Any]]: - """The next k distinct problems draw() would return, WITHOUT advancing state - (no cursor move, no reshuffle). Used to prefetch their baseline into base_cache - while the current chunk trains -- a wrong guess (cross-epoch reshuffle) only - misses the cache, never corrupts the draw.""" - out, seen, cur = [], set(), self._cursor - recs = self._records - while len(out) < k and cur < len(recs): # stop at epoch edge; don't simulate reshuffle - r = recs[cur] - cur += 1 - if id(r) not in seen: - seen.add(id(r)) - out.append(r) - return out - - -def _empty_roll() -> Dict[str, Any]: - return {'pred': '', 'correct': False, 'terminated': False, 'passed': False, - 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} - - -def _apply_baseline(r: Dict[str, Any], roll: Dict[str, Any]) -> None: - """Attach a greedy baseline roll and reset per-chunk working state.""" - r['_baseline_rolls'], r['_cands'], r['_init'] = [roll], [], [roll] - r['_failed'] = not roll['correct'] - r['_baseline_pass'] = 1.0 if roll['correct'] else 0.0 - r['_hard'] = True # process every problem; group variance selects (SEAM-style) - - -def baseline_rollout(base_sampler, problems: List[Dict[str, Any]], base_dp: int, - args: argparse.Namespace, cache: DiskCache) -> int: - """Base solves each problem greedily once (T=0, M=1), disk-cached by problem text. - The base is frozen + greedy so a cache hit is exact. Returns fresh (miss) count.""" - todo = [r for r in problems if DiskCache.key_for(r['problem']) not in cache] - if todo: - out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], - 1, args.max_tokens, base_dp, temperature=0.0) - for r, seqs in zip(todo, out): - roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - cache.put(DiskCache.key_for(r['problem']), roll) - for r in problems: - _apply_baseline(r, cache.get(DiskCache.key_for(r['problem']))) - return len(todo) - - -# -- rubric process-check (view A): teacher diagnoses the base's attempt -- -_RFT_DIAG_SYSTEM = """\ -You are a strategy-level process checker for a math solution attempt. You are given a -math problem, a rubric, and one attempted solution segment. Decide PASS or FAIL for each -criterion, and write the diagnosis so it can become useful reusable guidance for solving -similar problems without seeing this segment. - -Output STRICT JSON (no prose outside it) with this shape: -{ - "items": [ - {"index": 1, "verdict": "PASS", "reason": "<why the process satisfies it>", - "fix": ""}, - {"index": 2, "verdict": "FAIL", "reason": "<what reusable process issue is present>", - "fix": "<local strategy correction, without solving the problem>"} - ], - "overall": "OK" | "ISSUES", - "summary": "<one sentence naming the reusable process issue, not the answer>" -} - -Rules: -- Judge every criterion independently; a [Hard Rule] is FAIL unless clearly satisfied. -- Judge ONLY what is observable in THIS segment. Ignore hidden <think> or <thinking> - content for output-format criteria. -- The API diagnosis is an external teacher signal, so it must stay answer-free. -- Prefer diagnosis that transfers to view-B skill generation: name the route choice, - structural observation, missing check, or length-control habit that a solver should - remember before solving a similar problem. -- For PASS items, leave "fix" as "". -- For FAIL items, describe the process problem at strategy level: unsuitable method, - missed structure, invalid transformation, missing constraint check, redundant cases, - off-track approach, contradiction, or inefficient/unfinished reasoning. -- A fix may suggest the LOCAL correction direction, such as identify the key structure, - verify constraints, preserve equivalence, reduce redundant cases, or choose a more - direct route. Do not carry out the correction. -- Never reveal the final answer, a corrected value/expression, an option label, or a - step-by-step solution that would let another model copy the solve. -- If the segment contains a process note saying it was cut off before a final boxed - answer, mark the length-budget criterion as FAIL and suggest a method-level way to - finish faster. -- Keep every "reason" and "fix" concise: one short sentence each. -- "overall" is "OK" only if NO criterion is FAIL. -- Output only the JSON object.""" - -_RFT_DIAG_USER = """\ -## Task / query (context) -{query} - -## Rubric -{rubric} - -## Segment -{segment} - -Now output the diagnostic JSON object.""" - -_MATH_RUBRIC = [ - ('The attempt chooses a method suitable for the problem structure', False), - ('The attempt identifies the key constraint, invariant, or quantity before computing', False), - ('Algebraic and logical transformations preserve validity at each step', True), - ('The attempt checks required constraints, domains, boundary cases, or validity conditions', False), - ('The attempt avoids redundant casework, looping, or re-deriving known facts', False), - ('The attempt reaches a final boxed answer within the length budget', False), - ('The approach stays focused on the actual question asked', False), -] - -# Bump when the rubric criteria or the diagnosis prompt change so the disk-cached -# diagnoses written under an older rubric are not silently reused. -_RUBRIC_VERSION = 'rubric_v5_viewb_strategy' - - -class _RftRubricVerifier(RubricVerifier): - def _diagnose_trajectory(self, query: str, rubric_block: str, segment_text: str) -> dict: - return {'messages': [ - {'role': 'system', 'content': _RFT_DIAG_SYSTEM}, - {'role': 'user', 'content': _RFT_DIAG_USER.format( - query=query, rubric=rubric_block, segment=segment_text)}]} - - -def build_rubric_checker() -> Optional[RubricVerifier]: - """Fixed math-process rubric verifier (teacher-served). None if no LLM backup env.""" - if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') - or os.environ.get('OPENAI_API_KEY')): - return None - return _RftRubricVerifier( - fixed_rubric=[RubricItem(t, is_hard=h) for t, h in _MATH_RUBRIC], gate=True) - - -def _format_diagnosis(detail) -> str: - """One line per criterion (PASS/FAIL + reason + fix) then a summary.""" - rub = detail.rubric - lines = [] - for it in detail.items: - text = rub[it.index - 1].text if 0 < it.index <= len(rub) else f'criterion {it.index}' - if it.verdict: - lines.append(f'- [PASS] {text}') - else: - tail = f': {it.reason}' if it.reason else '' - tail += f' (fix: {it.fix})' if it.fix else '' - lines.append(f'- [FAIL] {text}{tail}') - if detail.summary: - lines.append(f'Summary: {detail.summary}') - return '\n'.join(lines) - - -def diagnose_views(checker, problems: List[Dict[str, Any]], args: argparse.Namespace, - cache: DiskCache) -> None: - """Rubric-check every view-A problem's greedy attempt in parallel (disk-cached by - problem + attempt), stashing the formatted findings on ``r['_rubric_diag']``.""" - targets = [r for r in problems if r.get('_view') == 'A'] - if not checker or not targets: - return - - def _key(r: Dict[str, Any]) -> str: - init = r.get('_init', [{}])[0] - term = 'L' if (init.get('stop_reason') == 'length' or not init.get('terminated')) else 'T' - return DiskCache.key_for(r['problem'], init.get('text', ''), _RUBRIC_VERSION, term) - - pending = [] - for r in targets: - key = _key(r) - if key in cache: - r['_rubric_diag'] = cache.get(key) - else: - pending.append((r, key)) - if not pending: - return - - def _run(item): - r, key = item - init = r['_init'][0] - seg_text = init['text'] - if init.get('stop_reason') == 'length' or not init.get('terminated'): - seg_text += ('\n\n[Process note: this attempt was cut off at the token budget ' - 'and never produced a final \\boxed{} answer.]') - seg = {'messages': [{'role': 'user', 'content': r['problem']}, - {'role': 'assistant', 'content': seg_text}]} - attempts = max(1, args.rubric_retries + 1) - for attempt in range(attempts): - try: - return r, key, _format_diagnosis(checker.diagnose(seg, query=r['problem'])) - except Exception as exc: # teacher hiccup -> retry, then no-diagnosis prompt (not cached) - if attempt + 1 < attempts: - logger.warning(f'[rubric] diagnose error: {exc}; retry {attempt + 1}/{args.rubric_retries}') - time.sleep(min(2.0, 0.5 * (2 ** attempt))) - continue - logger.warning(f'[rubric] diagnose error: {exc}; giving up after {attempts} attempts') - return r, key, None - - workers = max(1, min(args.rubric_workers, len(pending))) - with ThreadPoolExecutor(max_workers=workers) as ex: - for r, key, diag in ex.map(_run, pending): - r['_rubric_diag'] = diag or '' - if diag is not None: - cache.put(key, diag) - - -# =========================================================================== -# Block E -- chunk draw, generation pipeline, record building -# =========================================================================== -def _baseline_class(r: Dict[str, Any]) -> str: - """success | fail_loop (out of length / never terminated) | fail_wrong.""" - roll = r['_init'][0] - if roll['correct']: - return 'success' - return 'fail_loop' if (roll['stop_reason'] == 'length' or not roll['terminated']) else 'fail_wrong' - - -def _select_balanced(buckets: Dict[str, List[Dict[str, Any]]], n_success: int, - n_fail: int, n_fail_loop: int) -> List[Dict[str, Any]]: - """Pick n_fail base-fails (toward n_fail_loop loop-fails, best-effort) + n_success - base-successes; top up any shortfall from leftovers.""" - loop, wrong, succ = buckets['fail_loop'], buckets['fail_wrong'], buckets['success'] - take_loop = min(n_fail_loop, len(loop)) - take_wrong = min(n_fail - take_loop, len(wrong)) - take_loop = min(n_fail - take_wrong, len(loop)) - sel = loop[:take_loop] + wrong[:take_wrong] + succ[:n_success] - target = n_success + n_fail - if len(sel) < target: - used = {id(x) for x in sel} - sel += [x for b in (loop, wrong, succ) for x in b if id(x) not in used][:target - len(sel)] - return sel - - -def draw_chunk(pool: ProblemPool, base_sampler, base_dp: int, args: argparse.Namespace, - cache: DiskCache) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: - """Draw one chunk, baselining every drawn problem. With ``--balance``, keep - drawing+baselining until the target base fail:success mix is reachable (or the budget - is hit), then select a balanced subset.""" - if not args.balance: - chunk = pool.draw(args.chunk_size) - fresh = baseline_rollout(base_sampler, chunk, base_dp, args, cache) - return chunk, {'enabled': False, 'n_drawn': len(chunk), 'n_baseline_fresh': fresh} - - n_success = max(0, min(args.chunk_size, round(args.chunk_size * args.balance_success_frac))) - n_fail = args.chunk_size - n_success - n_fail_loop = round(n_fail * args.balance_loop_frac) - buckets: Dict[str, List[Dict[str, Any]]] = {'success': [], 'fail_loop': [], 'fail_wrong': []} - budget, n_drawn, n_fresh, seen = args.chunk_size * args.balance_max_draws_mult, 0, 0, set() - while n_drawn < budget: - if (len(buckets['success']) >= n_success - and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail): - break - batch = pool.draw(args.chunk_size) - n_fresh += baseline_rollout(base_sampler, batch, base_dp, args, cache) - n_drawn += len(batch) - for r in batch: - if id(r) not in seen: - seen.add(id(r)) - buckets[_baseline_class(r)].append(r) - - reached = (len(buckets['success']) >= n_success - and len(buckets['fail_loop']) + len(buckets['fail_wrong']) >= n_fail) - chunk = _select_balanced(buckets, n_success, n_fail, n_fail_loop) - sel_success = sum(1 for r in chunk if not r['_failed']) - stats = { - 'enabled': True, 'n_drawn': n_drawn, 'n_baseline_fresh': n_fresh, 'n_selected': len(chunk), - 'target_success': n_success, 'target_fail': n_fail, 'target_fail_loop': n_fail_loop, - 'selected_success': sel_success, 'selected_fail': len(chunk) - sel_success, - 'selected_fail_loop': sum(1 for r in chunk if _baseline_class(r) == 'fail_loop'), - 'selected_fail_wrong': sum(1 for r in chunk if _baseline_class(r) == 'fail_wrong'), - 'selected_success_frac': (sel_success / len(chunk)) if chunk else 0.0, - 'budget_hit': not reached, - } - return chunk, stats - - -def _assign_advantages(hard: List[Dict[str, Any]], args: argparse.Namespace) -> None: - """Group-relative advantage over each problem's scored candidates using the greedy - binary reward R in {0,1}: ``A = (R - mean) / (std + eps)``. std==0 groups get no - gradient -- GRPO's variance selects informative problems (no explicit difficulty gate). - A symmetric clip limits both positive and negative outliers from sparse 1-vs-many groups.""" - eps = 1e-6 - adv_clip = abs(float(getattr(args, 'adv_clip', 0.0) or 0.0)) - for r in hard: - for c in r['_cands']: - c['advantage'], c['grpo_adv'], c['kept'] = 0.0, 0.0, False - cs = [c for c in r['_cands'] if c.get('reward') is not None] - if len(cs) < 2: - continue - rewards = [c['reward'] for c in cs] - mean_r = sum(rewards) / len(rewards) - std = (sum((x - mean_r) ** 2 for x in rewards) / len(rewards)) ** 0.5 - if std < 1e-9: - continue - for c in cs: - raw_adv = (c['reward'] - mean_r) / (std + eps) - adv = max(-adv_clip, min(adv_clip, raw_adv)) if adv_clip > 0 else raw_adv - c['advantage'], c['grpo_adv'], c['kept'] = adv, adv, c['reward'] > mean_r - - -def _best_sft_candidate(r: Dict[str, Any], args: argparse.Namespace) -> Optional[Dict[str, Any]]: - """Pick ONE view-A candidate to distill (online context distillation). PREFER the - executor-verified PASSING skills (reward==1); if NONE passed -- common on the hard - problems that are exactly the cases worth distilling -- FALL BACK to any parseable - open-book skill regardless of the executor outcome. Answer-bearing skills produced by - the skill model itself are allowed here; only the external API/rubric diagnosis must be - answer-free. Within the chosen tier, take the one whose skill length is CLOSEST to - ``--sft-target-len`` -- an empirically high-pass-rate length (~500-600 chars in this - run) -- breaking ties by the fewest executor solve tokens. Targeting a length (rather - than the minimum) avoids a distillation feedback loop that would otherwise drive - rollouts ever shorter. None only when no parseable candidate exists at all.""" - eligible = [c for c in r['_cands'] if c.get('parseable') and c.get('skills')] - if not eligible: - return None - passing = [c for c in eligible if c.get('reward') == 1.0] - cs = passing or eligible - target = int(getattr(args, 'sft_target_len', 550) or 550) - - def _solve_tokens(c: Dict[str, Any]) -> int: - rolls = c.get('rolls') or [] - return rolls[0].get('gen_tokens', 1 << 30) if rolls else (1 << 30) - - return min(cs, key=lambda c: (abs(len(c['skills']) - target), _solve_tokens(c))) - - -def apply_neighbor_rubric(base_sampler, chunk: List[Dict[str, Any]], - neighbor_map: Dict[str, Tuple[str, float]], - pool_answers: Dict[str, str], base_dp: int, - args: argparse.Namespace, checker, - base_cache: DiskCache, rubric_cache: DiskCache) -> None: - """Cross-problem rubric (--xproblem-rubric): for every view-A problem Q, replace its own - rubric with the rubric of its bag-of-words neighbour P (Q keeps being the solved/scored - problem). P is baselined + diagnosed here (both disk-cached) as a stub carrying P's REAL - answer, so P's baseline grades correctly and legitimately shares the baseline cache with - P-as-training-problem. P's findings are copied onto Q with the neighbour text + similarity - for audit. A P that can't be diagnosed -> Q degrades to view B. Since P's answer differs - from Q's (guaranteed at pairing time), any answer P's rubric leaks is useless for Q.""" - targets = [r for r in chunk if r.get('_view') == 'A' and neighbor_map.get(r['problem'])] - if not targets: - return - stubs, by_problem = [], {} - for r in targets: - p, _ = neighbor_map[r['problem']] - if p not in by_problem: - stub = {'problem': p, 'reference_answer': pool_answers.get(p, ''), '_view': 'A'} - by_problem[p] = stub - stubs.append(stub) - baseline_rollout(base_sampler, stubs, base_dp, args, base_cache) - diagnose_views(checker, stubs, args, rubric_cache) - for r in targets: - p, sim = neighbor_map[r['problem']] - r['_rubric_diag'] = by_problem[p].get('_rubric_diag', '') - r['_rubric_src'], r['_neighbor_sim'] = p, sim - - -def process_chunk(base_sampler, skill_sampler, chunk: List[Dict[str, Any]], - ci: int, base_dp: int, skill_dp: int, args: argparse.Namespace, - checker, rubric_cache: DiskCache, base_cache: DiskCache = None, - neighbor_map: Optional[Dict[str, Tuple[str, float]]] = None, - pool_answers: Optional[Dict[str, str]] = None - ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], List[Dict[str, Any]]]: - """view assign -> rubric-check (view A) -> skill-gen -> answer audit -> with-skill - greedy pass -> GRPO advantages. ``chunk`` arrives already baselined by draw_chunk. - With --xproblem-rubric, view A gets its NEIGHBOUR's rubric (see apply_neighbor_rubric).""" - hard = chunk - for r in hard: - r['_view'], r['_rubric_diag'] = _assign_view(r['problem'], args), '' - if args.xproblem_rubric and neighbor_map: - apply_neighbor_rubric(base_sampler, hard, neighbor_map, pool_answers or {}, base_dp, - args, checker, base_cache, rubric_cache) - else: - diagnose_views(checker, hard, args, rubric_cache) - - # skill-gen (thinking OFF), per-view prompt; re-sample problems with no clean candidate. - # View-A problems that cannot yield a training record (no [FAIL] rubric, or rubric - # leaked the answer) are dropped from training entirely -- skip their generation. - flat: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] - pending = [r for r in hard if not _viewa_dropped(r, args)] - for _ in range(args.skill_retries + 1): - if not pending: - break - sg_out = _run_samples(skill_sampler, [_view_prompt(r, args) for r in pending], - args.n_skills, args.skill_max_tokens, skill_dp, - temperature=args.skill_gen_temperature, - top_p=args.skill_gen_top_p, top_k=args.skill_gen_top_k) - still = [] - for r, seqs in zip(pending, sg_out): - got = False - for s in seqs: - resp = _clean_text(getattr(s, 'decoded', '') or '') - block = _extract_skill(resp) or '' - cand = {'skills': block, 'response': resp, 'parseable': bool(block), - 'view': r['_view'], 'leaked': None, 'leak_reason': '', - 'leak_source': '', 'with_pass': None, 'reward': None, 'rolls': [], - 'skillgen_stop': getattr(s, 'stop_reason', None), - 'skillgen_tokens': len(getattr(s, 'tokens', None) or [])} - r['_cands'].append(cand) - if block: - flat.append((r, cand)) - got = True - if not got: - still.append(r) - pending = still - - # Answer-bearing skill audit: deterministic verbatim-answer check only (no LLM). This - # is observability only; it records metrics for swanlab/jsonl, but does not block - # scoring, reward, advantage assignment, SFT candidate selection, or GRPO training. - for r, c in flat: - leaked = _answer_leaked(c['skills'], r['reference_answer']) - c['leaked'] = leaked - c['leak_reason'] = 'answer_verbatim' if leaked else '' - c['leak_source'] = 'deterministic' - - # with-skill greedy pass (T=0, M=1); reward = correct, absolute (group mean is baseline). - scored_inputs = flat - if scored_inputs: - ws_out = _run_samples(base_sampler, - [build_skill_solve_prompt(r['problem'], c['skills']) for r, c in scored_inputs], - 1, args.max_tokens, base_dp, temperature=0.0) - for (r, c), seqs in zip(scored_inputs, ws_out): - c['rolls'] = [_parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll()] - c['with_pass'] = 1.0 if c['rolls'][0]['correct'] else 0.0 - c['reward'] = c['with_pass'] - if args.format_in_reward: # unparseable candidates score 0 and still join the group - for r in hard: - for c in r['_cands']: - if c['reward'] is None: - c['reward'] = 0.0 - - _assign_advantages(hard, args) - return ([_full_record(r, ci) for r in chunk], _chunk_summary(chunk, ci, args), - _group_records(chunk, args)) - - -def _roll(x: Dict[str, Any]) -> Dict[str, Any]: - return {k: x[k] for k in ('pred', 'correct', 'terminated', 'passed', - 'stop_reason', 'gen_tokens', 'text')} - - -def _is_trainable(c: Dict[str, Any], args: argparse.Namespace) -> bool: - """Reaches the GRPO update iff advantage is non-zero. Leak flags are audit-only.""" - adv_nz = abs(c.get('advantage') or 0.0) > 1e-9 - if args.format_in_reward: - return adv_nz - return c.get('with_pass') is not None and adv_nz - - -def _full_record(r: Dict[str, Any], ci: int) -> Dict[str, Any]: - """Complete per-problem trace: init attempt, baseline, and all candidates.""" - init = r['_init'][0] - return { - 'record_type': 'problem', 'chunk': ci, 'problem': r['problem'], - 'reference_answer': r['reference_answer'], 'level': r.get('level', ''), - 'failed_first_try': r['_failed'], - 'init_attempt': {'text': init['text'], 'pred': init['pred'], 'correct': init['correct'], - 'terminated': init['terminated'], 'stop_reason': init['stop_reason'], - 'gen_tokens': init['gen_tokens']}, - 'baseline_pass': r['_baseline_pass'], 'is_hard': r['_hard'], - 'view': r.get('_view', ''), 'rubric_diag': r.get('_rubric_diag', ''), - # xproblem audit: which neighbour's rubric this problem borrowed, and how similar. - 'rubric_src': r.get('_rubric_src', ''), 'neighbor_sim': r.get('_neighbor_sim'), - 'baseline_rolls': [_roll(x) for x in r['_baseline_rolls']], - 'candidates': [{ - 'skills': c['skills'], 'response': c['response'], 'parseable': c['parseable'], - 'skillgen_stop': c.get('skillgen_stop'), 'skillgen_tokens': c.get('skillgen_tokens'), - 'leaked': c['leaked'], 'leak_reason': c['leak_reason'], 'leak_source': c['leak_source'], - 'with_pass': c['with_pass'], 'reward': c.get('reward'), 'advantage': c.get('advantage'), - 'grpo_adv': c.get('grpo_adv'), 'kept': c.get('kept'), - 'rolls': [_roll(x) for x in c['rolls']], - } for c in r['_cands']], - } - - -def _view_stats(problems: List[Dict[str, Any]], view: str) -> Dict[str, Any]: - pv = [r for r in problems if r.get('_view') == view] - cands = [c for r in pv for c in r['_cands'] if c['parseable']] - clean = [c for c in cands if c['leaked'] is False] - adopted = sum(1 for r in pv - if any(abs(c.get('advantage') or 0.0) > 1e-9 for c in r['_cands'])) - return {'n': len(pv), 'n_candidates_parseable': len(cands), 'n_clean': len(clean), - 'n_adopted_problems': adopted, 'adoption_rate': (adopted / len(pv)) if pv else 0.0} - - -def _mean(xs: List[float]) -> float: - return sum(xs) / len(xs) if xs else 0.0 - - -def _std(xs: List[float]) -> float: - if len(xs) < 2: - return 0.0 - m = _mean(xs) - return (sum((x - m) ** 2 for x in xs) / len(xs)) ** 0.5 - - -def _signal_stats(problems: List[Dict[str, Any]]) -> Dict[str, float]: - """The heart of 'is there a learning signal': per problem, the scored candidates form a - GRPO group. A group with zero reward variance (all skills solve, or none do -- the - hard-problem dead zone) gives NO gradient. Tracks that fraction plus reward level and - within-group variance so a collapse (all-0 or all-1) is visible immediately.""" - group_vars, all_rewards, zero_grad, groups = [], [], 0, 0 - for r in problems: - rewards = [c['reward'] for c in r['_cands'] if c.get('reward') is not None] - if len(rewards) < 2: - continue - groups += 1 - all_rewards.extend(rewards) - v = _std(rewards) - group_vars.append(v) - if v < 1e-9: # every skill got the same reward -> GRPO skips this problem - zero_grad += 1 - return {'n_groups': groups, 'zero_grad_frac': (zero_grad / groups) if groups else 0.0, - 'reward_mean': _mean(all_rewards), 'reward_std': _std(all_rewards), - 'group_reward_std_mean': _mean(group_vars)} - - -def _chunk_summary(chunk: List[Dict[str, Any]], ci: int, args: argparse.Namespace) -> Dict[str, Any]: - all_cands = [c for r in chunk for c in r['_cands']] - cands = [c for c in all_cands if c['parseable']] - scored = [c for c in cands if c['with_pass'] is not None] - clean = [c for c in cands if c['leaked'] is False] - ws_rolls = [x for c in scored for x in c['rolls']] - # viewa-dropped problems generate no candidates; keep acc/* on the generated subset - # so the with-skill/lift trend stays comparable across view_b_frac settings. - gen_probs = [r for r in chunk if r['_cands']] - base_acc = _mean([1.0 if r['_baseline_pass'] else 0.0 for r in gen_probs]) - ws_acc = _mean([1.0 if any(c.get('reward') for c in r['_cands']) else 0.0 for r in gen_probs]) - cand_pass_parseable = _mean([c['with_pass'] for c in scored]) - cand_pass_all = _mean([1.0 if c.get('reward') else 0.0 for c in all_cands]) - # base failure taxonomy (you asked whether skills fail because the base loops out of length) - classes = [_baseline_class(r) for r in chunk] - n_fail = sum(1 for c in classes if c != 'success') - skill_tokens = [c.get('skillgen_tokens') or 0 for c in cands] # skill-gen response length - trunc = sum(1 for r in chunk for c in r['_cands'] - for x in c['rolls'] if x['stop_reason'] == 'length') - rubric_answer_leaks = sum( - 1 for r in chunk - if r.get('_view') == 'A' and _rubric_has_fail(r.get('_rubric_diag')) and _rubric_answer_leaked(r)) - return { - 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), - 'n_failed_first_try': sum(1 for r in chunk if r['_failed']), - 'n_generated': len(all_cands), 'n_candidates_parseable': len(cands), - 'n_unparseable': len(all_cands) - len(cands), - 'parse_rate': (len(cands) / len(all_cands)) if all_cands else 0.0, - 'n_leaked': sum(1 for c in cands if c['leaked']), 'n_clean': len(clean), - 'leak_rate': (sum(1 for c in cands if c['leaked']) / len(cands)) if cands else 0.0, - 'n_reward_pos': sum(1 for c in scored if c['reward']), - 'n_rubric_answer_leaked': rubric_answer_leaks, - 'n_viewa_dropped': sum(1 for r in chunk if _viewa_dropped(r, args)), - 'n_train_samples': sum(1 for c in all_cands if _is_trainable(c, args)), - 'signal': _signal_stats(chunk), - 'fail_loop_frac': (sum(1 for c in classes if c == 'fail_loop') / n_fail) if n_fail else 0.0, - 'fail_wrong_frac': (sum(1 for c in classes if c == 'fail_wrong') / n_fail) if n_fail else 0.0, - 'skill_tokens_mean': _mean(skill_tokens), - 'withskill_trunc_frac': (trunc / len(ws_rolls)) if ws_rolls else 0.0, - 'avg_baseline_pass': base_acc, 'avg_withskill_pass': ws_acc, - 'avg_lift': ws_acc - base_acc, - 'candidate_withskill_pass_parseable': cand_pass_parseable, - 'candidate_withskill_pass_all': cand_pass_all, - 'termination_rate_withskill': _mean([1.0 if x['terminated'] else 0.0 for x in ws_rolls]), - 'view_A': _view_stats(chunk, 'A'), 'view_B': _view_stats(chunk, 'B'), - **_xproblem_stats(chunk, args), - } - - -def _xproblem_stats(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> Dict[str, Any]: - """Cross-problem pairing health: of the view-A problems, how many actually got a - neighbour's rubric (pair_rate) and how similar those neighbours were. Empty when off.""" - if not args.xproblem_rubric: - return {} - view_a = [r for r in chunk if r.get('_view') == 'A'] - paired = [r for r in view_a if r.get('_rubric_src')] - return {'xproblem': { - 'n_view_a': len(view_a), 'n_paired': len(paired), - 'pair_rate': (len(paired) / len(view_a)) if view_a else 0.0, - 'neighbor_sim_mean': _mean([r.get('_neighbor_sim', 0.0) for r in paired])}} - - -def _sft_record(r: Dict[str, Any], c: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]: - """A view-A context-distillation sample: the OPEN-BOOK (rubric-conditioned) skill ``c`` - is the target, but the prompt is rebuilt as query-only (view B, no rubric) so the model - learns to produce it CLOSED-BOOK. Rides the same GRPOLoss with a positive constant - advantage (``--sft-weight``); single-step (old_logps=None) this reduces to - ``-w*logp + beta*KL`` -- a KL-anchored weighted cross-entropy toward the skill.""" - return { - 'data_id': r.get('data_id'), - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'view': 'B', 'diagnosis': '', 'rubric_src': '', 'orig_view': 'A', 'sft': True, - 'response': c['response'], 'skills': c['skills'], - 'skillgen_stop': c.get('skillgen_stop'), - 'advantage': float(args.sft_weight), 'grpo_adv': 0.0, 'kept': True, - 'reward': c['reward'], 'with_pass': c['with_pass']} - - -def _rubric_answer_leaked(r: Dict[str, Any]) -> bool: - """Hard safety gate for external API/rubric diagnosis. Unlike answer-bearing skills - generated by the policy itself, a rubric that contains the target final answer is an - external teacher leak and must not be distilled into view B.""" - return _answer_leaked(r.get('_rubric_diag', ''), r.get('reference_answer', '')) - - -def _viewa_dropped(r: Dict[str, Any], args: argparse.Namespace) -> bool: - """Under --viewa-sft a view-A problem trains ONLY through the SFT path (rubric with - [FAIL], answer-free). No-FAIL and rubric-leaked problems produce no training record - at all (no GRPO backflow: those prompts are query-only and would muddy the pure - view-B GRPO pool), so their skill-gen + executor scoring is skipped entirely.""" - return (bool(args.viewa_sft) and r.get('_view') == 'A' - and (not _rubric_has_fail(r.get('_rubric_diag')) - or _rubric_answer_leaked(r))) - - -def _group_records(chunk: List[Dict[str, Any]], args: argparse.Namespace) -> List[Dict[str, Any]]: - """Training records. Under --viewa-sft, view A is SFT-only: a problem whose rubric - localised a failure (has ``[FAIL]``, answer-free) contributes ONE context-distillation - SFT sample (best parseable open-book skill -- preferring an executor-verified pass, - else any parseable candidate -- rebuilt query-only); no-FAIL / rubric-leaked view-A - problems are dropped (``_viewa_dropped``, generation already skipped). GRPO candidates - come from view B only, keeping that pool purely closed-book. Prompts are rebuilt from - the stored view/diagnosis by ``_skillgen_messages``.""" - out = [] - for r in chunk: - if not r['_hard']: - continue - if args.viewa_sft and r.get('_view') == 'A': - if _viewa_dropped(r, args): - continue - best = _best_sft_candidate(r, args) - if best is not None: - out.append(_sft_record(r, best, args)) - continue - for c in r['_cands']: - if _is_trainable(c, args): - out.append({ - 'data_id': r.get('data_id'), - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'view': r.get('_view', 'A'), 'diagnosis': r.get('_rubric_diag', ''), - 'rubric_src': r.get('_rubric_src', ''), 'sft': False, - 'response': c['response'], 'skills': c['skills'], - 'skillgen_stop': c.get('skillgen_stop'), - 'advantage': c['advantage'], 'grpo_adv': c['grpo_adv'], 'kept': c['kept'], - 'reward': c['reward'], 'with_pass': c['with_pass']}) - return out - - -# =========================================================================== -# Block G -- online GRPO training -# =========================================================================== -def _is_num(v: Any) -> bool: - try: - float(v) - return True - except (TypeError, ValueError): - return False - - -def _train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: - """Training sample = the exact skill-gen prompt (rebuilt by ``_skillgen_messages`` so - train/inference match) + the generated structured guidance response. ``key_rounds`` - selects the final assistant turn; Template masks the prompt and trains the whole - response (the key-round prefix already excludes the prompt-provided <think>).""" - msgs = _skillgen_messages( - rec['problem'], rec.get('view', 'A'), rec.get('diagnosis', ''), rec.get('rubric_src', '')) - return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], - 'user_data': pack_user_data({'key_rounds': [len(msgs)]})} - - -def _train_chunk(skill_model, ref_model, ckpt: CheckpointEngineManager, samples: List[Dict[str, Any]], - args: argparse.Namespace) -> Dict[str, Any]: - """On-policy GRPO update over one chunk, then sync weights. Micro-batches of - ``sft_batch_size`` accumulate gradients; ONE optimizer step is taken per PPO - mini-batch of ``ppo_mini_batch_size`` samples (0 -> a single step over the whole - chunk, the original behaviour). A frozen reference model provides ref_logps for the - SEAM-style KL penalty. - - Multi-step correctness: with more than one step over the SAME rollout, later - mini-batches see an already-updated policy, so we FREEZE the sampling-policy - ``old_logps`` (recomputed once, before any step) and let GRPOLoss form the PPO ratio - against them. A single step keeps ``old_logps=None`` (ratio==1, pure on-policy). - The batch is padded to a multiple of ``sft_batch_size`` with advantage-0 copies that - contribute no policy gradient. View-A context-distillation samples ride the same loss - with a positive constant advantage (``--sft-weight``); single-step (old_logps=None) - that is ``-w*logp + beta*KL`` (KL-anchored weighted cross-entropy).""" - trajs = [_train_trajectory(rec) for rec in samples] - advs = [float(rec['advantage']) for rec in samples] - rem = (-len(trajs)) % args.sft_batch_size - if rem: - trajs += [trajs[-1]] * rem - advs += [0.0] * rem - - n, sft = len(trajs), args.sft_batch_size - mini = args.ppo_mini_batch_size if args.ppo_mini_batch_size > 0 else n - mini = max(sft, (mini // sft) * sft) # align to a whole number of micro-batches - multi_step = mini < n - - # Freeze ref_logps (frozen model) and -- only when we take multiple steps -- the - # sampling-policy old_logps, BOTH before any optimizer step touches the weights. With - # a single step old_logps stays None so GRPOLoss uses ratio==1 (pure on-policy). - micro_ref, micro_old = [], [] - for i in range(0, n, sft): - mb = trajs[i:i + sft] - micro_ref.append(ref_model.forward_only(inputs=mb).get('logps')) - micro_old.append(skill_model.forward_only(inputs=mb).get('logps') if multi_step else None) - - micro, n_steps = 0, 0 - for ms in range(0, n, mini): - for i in range(ms, min(ms + mini, n), sft): - k = i // sft - skill_model.forward_backward(inputs=trajs[i:i + sft], - advantages=advs[i:i + sft], - old_logps=micro_old[k], - ref_logps=micro_ref[k]) - micro += 1 - skill_model.clip_grad_and_step() - n_steps += 1 - ckpt.sync_weights(merge_and_sync=True) - metric = skill_model.calculate_metric(is_training=True) - n_sft = sum(1 for s in samples if s.get('sft')) - return {'n_samples': len(samples), 'n_sft': n_sft, 'n_grpo': len(samples) - n_sft, - 'n_steps': n_steps, 'n_micro_batches': micro, - 'metric': {k: (float(v) if _is_num(v) else v) for k, v in (metric or {}).items()}} - - -# =========================================================================== -# Block H -- fixed-holdout eval + metric formatting -# =========================================================================== -def run_greedy_eval(base_sampler, skill_sampler, eval_records: List[Dict[str, Any]], - ci: int, rounds: int, base_dp: int, skill_dp: int, args: argparse.Namespace, - base_cache: DiskCache - ) -> Tuple[List[Dict[str, Any]], Dict[str, Any], Dict[str, float]]: - """SEAM ``val-core/math/acc/mean@1`` on the fixed holdout: ONE greedy skill (T=0) per - problem into ONE greedy base solve (T=0). Eval ALWAYS uses view B (query-only, the - deployment form: no rubric, since rubric needs an online teacher unavailable at deploy); - no leak filter (acc scores correctness alone). Baseline reuses the disk cache.""" - baseline_rollout(base_sampler, eval_records, base_dp, args, base_cache) - for r in eval_records: - r['_view'], r['_rubric_diag'] = 'B', '' - sg_out = _run_samples(skill_sampler, [_view_prompt(r, args) for r in eval_records], - 1, args.skill_max_tokens, skill_dp, temperature=0.0) - skills = [] - for seqs in sg_out: - if not seqs: - skills.append(('', '')) - continue - sresp = _clean_text(getattr(seqs[0], 'decoded', '') or '') - skills.append((_extract_skill(sresp) or '', sresp)) - ws_out = _run_samples(base_sampler, - [build_skill_solve_prompt(r['problem'], sk) for r, (sk, _) in zip(eval_records, skills)], - 1, args.max_tokens, base_dp, temperature=0.0) - recs = [] - for r, (sk, sresp), seqs in zip(eval_records, skills, ws_out): - roll = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - recs.append({ - 'record_type': 'eval_problem', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, - 'data_id': r.get('data_id'), - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'view': r['_view'], 'rubric_diag': r.get('_rubric_diag', ''), - 'baseline_pass': r['_baseline_pass'], 'skill': sk, 'skill_parseable': bool(sk), - 'skill_response': sresp, 'withskill_pred': roll['pred'], - 'withskill_correct': roll['correct'], 'withskill_terminated': roll['terminated'], - 'withskill_stop_reason': roll['stop_reason'], 'withskill_text': roll['text'], - }) - acc = lambda rs: sum(1 for x in rs if x['withskill_correct']) / len(rs) if rs else 0.0 - ws = acc(recs) # all view B (deployment form) - base = sum(x['baseline_pass'] for x in recs) / len(recs) if recs else 0.0 - fmt = (sum(1 for x in recs if x['skill_parseable']) / len(recs)) if recs else 0.0 - term = (sum(1 for x in recs if x['withskill_terminated']) / len(recs)) if recs else 0.0 - summary = {'record_type': 'eval_summary', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, - 'n': len(recs), 'view': 'B', 'acc_mean1': ws, - 'baseline_acc_mean1': base, 'lift_mean1': ws - base, - 'format_mean1': fmt, 'term_mean1': term} - metrics = {'core/math/acc/mean@1': ws, 'core/math/baseline_acc/mean@1': base, - 'core/math/lift/mean@1': ws - base, 'core/math/format/mean@1': fmt, - 'core/math/term/mean@1': term} - return recs, summary, metrics - - -def _trend_line(hist: List[Dict[str, float]], window: int, rounds_done: int) -> Optional[str]: - """Contrast the FIRST vs the most recent ``window`` chunks -- if RFT works, adoption - and lift on recent (fresh) chunks exceed the early baseline.""" - if len(hist) < 2 * window: - return None - base, rec = hist[:window], hist[-window:] - m = lambda xs, k: sum(h[k] for h in xs) / len(xs) - return (f'[trend] first {window} vs last {window} | ' - f'adopt A {m(base,"aA"):.2f}->{m(rec,"aA"):.2f} B {m(base,"aB"):.2f}->{m(rec,"aB"):.2f} | ' - f'lift {m(base,"lift"):+.3f}->{m(rec,"lift"):+.3f} | ' - f'0grad {m(base,"zero_grad"):.2f}->{m(rec,"zero_grad"):.2f} | ' - f'pos/chunk {m(base,"pos"):.1f}->{m(rec,"pos"):.1f} | rounds={rounds_done}') - - -def _swan_metrics(summary: Dict[str, Any], log: Optional[Dict[str, Any]]) -> Dict[str, float]: - """Flat swanlab dict. ``signal/*`` is the primary health group (is GRPO getting a - gradient at all); ``acc/*`` is the effect; the rest diagnose why. acc/adopt/skill are - only logged when the chunk produced a scored group, so idle chunks don't dip charts.""" - sig = summary['signal'] - d: Dict[str, float] = { - # --- signal: the FIRST thing to watch (no variance -> no learning) --- - 'signal/zero_grad_frac': sig['zero_grad_frac'], 'signal/n_groups': sig['n_groups'], - 'signal/reward_mean': sig['reward_mean'], 'signal/reward_std': sig['reward_std'], - 'signal/group_reward_std_mean': sig['group_reward_std_mean'], - 'signal/n_train_samples': summary['n_train_samples'], - 'signal/n_reward_pos': summary['n_reward_pos'], - # --- skill format / leak health --- - 'skill/parse_rate': summary['parse_rate'], 'skill/tokens_mean': summary['skill_tokens_mean'], - 'leak/rate': summary['leak_rate'], 'leak/n': summary['n_leaked'], - # --- base failure taxonomy (loop-out-of-length vs plain wrong) --- - 'fail/loop_frac': summary['fail_loop_frac'], 'fail/wrong_frac': summary['fail_wrong_frac'], - 'fail/frac_first_try': (summary['n_failed_first_try'] / summary['n']) if summary['n'] else 0.0, - # --- view-A routing: dropped = no-FAIL rubric or rubric leak (SFT-only view A) --- - 'viewa/dropped_frac': (summary['n_viewa_dropped'] / summary['view_A']['n'] - if summary['view_A']['n'] else 0.0), - } - bal = summary.get('balance') or {} - if bal.get('enabled'): - d.update({'balance/n_drawn': bal['n_drawn'], 'balance/n_baseline_fresh': bal['n_baseline_fresh'], - 'balance/selected_success_frac': bal['selected_success_frac']}) - xp = summary.get('xproblem') or {} - if xp: - d.update({'xproblem/pair_rate': xp['pair_rate'], 'xproblem/neighbor_sim_mean': xp['neighbor_sim_mean']}) - if sig['n_groups'] > 0: - d.update({'acc/baseline_pass': summary['avg_baseline_pass'], - 'acc/withskill_pass': summary['avg_withskill_pass'], 'acc/lift': summary['avg_lift'], - 'candidate/withskill_pass_parseable': summary['candidate_withskill_pass_parseable'], - 'candidate/withskill_pass_all': summary['candidate_withskill_pass_all'], - 'adopt/A': summary['view_A']['adoption_rate'], - 'adopt/B': summary['view_B']['adoption_rate'], - 'term/withskill': summary['termination_rate_withskill'], - 'term/withskill_trunc_frac': summary['withskill_trunc_frac']}) - if log: - d['train/n_steps'] = log['n_steps'] - d['train/n_micro_batches'] = log['n_micro_batches'] - for k, v in (log.get('metric') or {}).items(): - if not _is_num(v): - continue - if k.startswith('learning rate'): - if 'group 1' in k: - d['train/lr'] = float(v) - else: - d[f'train/{k.replace(" ", "_")}'] = float(v) - return d - - -def _view_a_rubric_leak_metrics(chunk: List[Dict[str, Any]], - pool_answers: Optional[Dict[str, str]] = None) -> Dict[str, float]: - """Swanlab-only audit for answer leakage in view-A rubric text. This never changes - rewards, advantages, filtering, or training records.""" - view_a = [r for r in chunk if r.get('_view') == 'A'] - with_diag = [r for r in view_a if r.get('_rubric_diag')] - target_leaks = sum(1 for r in with_diag - if _answer_leaked(r.get('_rubric_diag', ''), r.get('reference_answer', ''))) - source_leaks = 0 - pool_answers = pool_answers or {} - for r in with_diag: - src = r.get('_rubric_src') - src_ref = pool_answers.get(src, '') if src else r.get('reference_answer', '') - if _answer_leaked(r.get('_rubric_diag', ''), src_ref): - source_leaks += 1 - n = len(with_diag) - return { - 'rubric_leak/n_view_a': float(len(view_a)), - 'rubric_leak/n_checked': float(n), - 'rubric_leak/target_answer_n': float(target_leaks), - 'rubric_leak/target_answer_rate': (target_leaks / n) if n else 0.0, - 'rubric_leak/source_answer_n': float(source_leaks), - 'rubric_leak/source_answer_rate': (source_leaks / n) if n else 0.0, - } - - -# =========================================================================== -# Block F -- components, args, main -# =========================================================================== -def init_components(args: argparse.Namespace): - """Default 8-GPU layout: rank 0 trains the actor, rank 1 hosts a frozen ref model, - 2-3 skill_sampler (synced), 4-7 base_sampler (frozen). Returns - (skill_model, ref_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp).""" - r0, r1 = TRAIN_GPUS, TRAIN_GPUS + REF_GPUS - r2, r3 = r1 + SKILL_SAMPLER_GPUS, NUM_GPUS - twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, lazy_collect=False, groups=[ - DeviceGroup(name='train', ranks=list(range(0, r0)), device_type='GPU'), - DeviceGroup(name='ref', ranks=list(range(r0, r1)), device_type='GPU'), - DeviceGroup(name='skill_sampler', ranks=list(range(r1, r2)), device_type='GPU'), - DeviceGroup(name='base_sampler', ranks=list(range(r2, r3)), device_type='GPU')]) - - train_mesh = DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_DP, fsdp_size=TRAIN_FSDP) - skill_model = TransformersModel(model_id=MODEL_ID, device_mesh=train_mesh, remote_group='train', - ddp_config={'find_unused_parameters': False}) - skill_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) - skill_model.set_template(Template, model_id=MODEL_ID, enable_thinking=False, - max_length=args.max_model_len, truncation_strategy='delete') - skill_model.set_processor(InputProcessor, padding_free=False) - skill_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon, beta=args.kl_beta) - skill_model.set_optimizer('AdamW', lr=args.lr) - skill_model.set_lr_scheduler('CosineWarmupScheduler', num_warmup_steps=10, - num_training_steps=args.max_train_rounds) - - ref_mesh = DeviceMesh.from_sizes(world_size=REF_GPUS, dp_size=REF_DP, fsdp_size=REF_FSDP) - ref_model = TransformersModel(model_id=MODEL_ID, device_mesh=ref_mesh, remote_group='ref', - ddp_config={'find_unused_parameters': False}) - ref_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) - ref_model.set_template(Template, model_id=MODEL_ID, enable_thinking=True, - max_length=args.max_model_len, truncation_strategy='delete') - ref_model.set_processor(InputProcessor, padding_free=False) - ref_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) - - def _sampler(group, world, enable_thinking: bool = True): - s = vLLMSampler(model_id=MODEL_ID, - engine_args={'gpu_memory_utilization': GPU_MEM, - 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1}, - device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), - remote_group=group) - s.set_template(Template, model_id=MODEL_ID, enable_thinking=enable_thinking, max_length=args.max_model_len) - return s - - skill_sampler = _sampler('skill_sampler', SKILL_SAMPLER_GPUS, enable_thinking=False) - # base_sampler is shared by the main thread and the baseline-prefetch thread -> serialise. - base_sampler = _LockedSampler(_sampler('base_sampler', BASE_SAMPLER_GPUS, enable_thinking=True)) - ckpt = CheckpointEngineManager(model=skill_model, sampler=skill_sampler) - return skill_model, ref_model, skill_sampler, base_sampler, ckpt, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS - - -def _build_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--dataset', choices=('aops', 'math'), default='aops') - p.add_argument('--n', type=int, default=2000, help='Problems loaded into the draw pool.') - p.add_argument('--pool-offset', type=int, default=0, - help='Skip this many shuffled non-eval records before building the train pool; ' - 'useful to avoid cold-start SFT data ranges.') - p.add_argument('--exclude-data-ids', default='', - help='Comma-separated jsonl files whose data_id/problem keys are excluded ' - 'from train/eval selection, e.g. coldstart_sft.jsonl.') - p.add_argument('--seed', type=int, default=42) - p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) - p.add_argument('--eval-size', type=int, default=128, help='Fixed holdout size (0 disables).') - p.add_argument('--eval-every', type=int, default=5, help='Run holdout eval every N chunks.') - p.add_argument('--chunk-size', type=int, default=16) - p.add_argument('--balance', action=argparse.BooleanOptionalAction, default=True) - p.add_argument('--balance-success-frac', type=float, default=0.4, - help='Target fraction of the chunk the base solves (rest are base-fail).') - p.add_argument('--balance-loop-frac', type=float, default=0.5) - p.add_argument('--balance-max-draws-mult', type=int, default=8) - p.add_argument('--seam-parquet-dir', type=str, default='', - help='Read SEAM build_aops_dataset.py train.parquet/val.parquet directly, in ' - 'file order (problem<-extra_info.problem, answer<-reward_model.ground_truth). ' - 'val.parquet becomes the eval holdout. Bypasses load/--numeric-only/' - '--eval-size/internal shuffle so the input data matches a SEAM run exactly.') - p.add_argument('--n-skills', type=int, default=8) - p.add_argument('--view-b-frac', type=float, default=0.5) - p.add_argument('--xproblem-rubric', action=argparse.BooleanOptionalAction, default=False, - help='When enabled, view A uses a bag-of-words NEIGHBOUR problem rubric. ' - 'Default is off: each view-A problem uses its own baseline attempt, ' - 'while the API diagnosis prompt is constrained to be answer-free and ' - 'method-level only.') - p.add_argument('--skill-retries', type=int, default=2) - p.add_argument('--skill-gen-temperature', type=float, default=1.0) - p.add_argument('--skill-gen-top-p', type=float, default=1.0) - p.add_argument('--skill-gen-top-k', type=int, default=-1) - p.add_argument('--max-model-len', type=int, default=16384) - p.add_argument('--max-tokens', type=int, default=8192) - p.add_argument('--skill-max-tokens', type=int, default=8192) - p.add_argument('--rubric-workers', type=int, default=16) - p.add_argument('--rubric-retries', type=int, default=2, - help='Retry failed/timeout rubric diagnose calls this many times before ' - 'falling back to an empty diagnosis without caching the failure.') - p.add_argument('--sft-batch-size', type=int, default=8, - help='Driver micro-batch (gradient-accumulation unit); multiple of train dp.') - p.add_argument('--ppo-mini-batch-size', type=int, default=0, - help='Samples per optimizer step (SEAM-style PPO mini-batch). 0 = one step ' - 'over the whole chunk (on-policy, ratio==1). When >0 and smaller than ' - 'the trainable count, multiple steps are taken over the same rollout and ' - 'old_logps are frozen so the PPO ratio/clip stays valid. Rounded down to ' - 'a multiple of --sft-batch-size.') - p.add_argument('--grpo-epsilon', type=float, default=0.2) - p.add_argument('--adv-clip', type=float, default=3.0, - help='Symmetric clip for group-relative advantages; <=0 disables clipping.') - p.add_argument('--kl-beta', type=float, default=0.001, - help='SEAM-style reference KL coefficient for GRPOLoss.') - p.add_argument('--viewa-sft', action=argparse.BooleanOptionalAction, default=True, - help='Route view-A problems to online context distillation (SFT on the best ' - 'passing open-book skill, rebuilt as a query-only prompt) instead of GRPO. ' - 'View B stays GRPO; both share one optimizer step.') - p.add_argument('--sft-weight', type=float, default=0.5, - help='Advantage magnitude (lambda) for view-A SFT distillation samples in the ' - 'shared GRPOLoss. Single-step this is -lambda*logp + beta*KL (weighted CE). ' - 'Keep <~1 so one-signed SFT gradients do not swamp mean-0 GRPO advantages.') - p.add_argument('--sft-target-len', type=int, default=550, - help='Target skill length (chars) for view-A SFT distillation: among passing ' - 'candidates pick the one CLOSEST to this length (not the shortest). Set near ' - 'the empirical pass-rate peak (~500-600 here) so distillation neither shrinks ' - 'rollouts toward zero nor lets them grow unbounded.') - p.add_argument('--format-in-reward', action=argparse.BooleanOptionalAction, default=True) - p.add_argument('--lr', type=float, default=6e-6) - p.add_argument('--max-train-rounds', type=int, default=1500) - p.add_argument('--save-rounds', type=int, default=200) - p.add_argument('--trend-every', type=int, default=10) - p.add_argument('--output-dir', default='./output/reflexion_skill') - p.add_argument('--cache-dir', default='', help='Baseline/rubric cache dir (default <output-dir>/cache).') - p.add_argument('--no-cache', action='store_true', help='Disable disk cache read/write.') - p.add_argument('--prefetch-baseline', action=argparse.BooleanOptionalAction, default=True, - help='Prefetch next chunk base baseline on a background thread (overlaps ' - 'with skill-gen; base_sampler is frozen so it never blocks the trainer).') - p.add_argument('--swanlab-project', default='twinkle') - p.add_argument('--swanlab-exp', default='') - args = p.parse_args() - if args.sft_batch_size % TRAIN_DP != 0: - raise ValueError(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple of train dp ({TRAIN_DP})') - if args.chunk_size < 1: - raise ValueError('--chunk-size must be >= 1') - args.balance_success_frac = max(0.0, min(1.0, args.balance_success_frac)) - return args - - -def _write(handle, row: Dict[str, Any]) -> None: - handle.write(json.dumps(row, ensure_ascii=False) + '\n') - - -def main() -> None: - args = _build_args() - records, eval_records, neighbor_map, pool_answers, data_stats = _load_records(args) - if len(records) < args.chunk_size: - raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded ({len(records)}); raise --n') - - os.makedirs(args.output_dir, exist_ok=True) - gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') - eval_path = os.path.join(args.output_dir, 'eval_records.jsonl') - data_path = os.path.join(args.output_dir, 'skill_dataset.jsonl') - train_log_path = os.path.join(args.output_dir, 'train_log.jsonl') - - if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('OPENAI_API_KEY')): - sys.stderr.write('[rft] WARNING: no LLM backup env; view-A rubric check disabled ' - '(leak filter is deterministic, unaffected)\n') - - use_swan = swanlab is not None and os.environ.get('SWANLAB_MODE') != 'disabled' - if use_swan: - swanlab.init(project=args.swanlab_project, experiment_name=(args.swanlab_exp or None), - config={'model': MODEL_ID, 'dataset': args.dataset, 'n': len(records), - 'eval_n': len(eval_records), 'n_skills': args.n_skills, - 'view_b_frac': args.view_b_frac, 'balance': args.balance, - 'balance_success_frac': args.balance_success_frac, - 'skill_gen_temp': args.skill_gen_temperature, - 'grpo_epsilon': args.grpo_epsilon, 'kl_beta': args.kl_beta, - 'lr': args.lr}) - - skill_model, ref_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp = init_components(args) - checker = build_rubric_checker() - if checker is None: - sys.stderr.write('[rft] no LLM backup env -> view-A rubric process-check DISABLED\n') - - cache_dir = args.cache_dir or os.path.join(args.output_dir, 'cache') - os.makedirs(cache_dir, exist_ok=True) - use_cache = not args.no_cache - base_cache = DiskCache(os.path.join(cache_dir, 'baseline.jsonl'), use_cache) - eval_base_cache = DiskCache(os.path.join(cache_dir, 'eval_baseline.jsonl'), use_cache) - rubric_cache = DiskCache(os.path.join(cache_dir, 'rubric.jsonl'), use_cache) - if args.xproblem_rubric: - sys.stderr.write(f'[rft] xproblem-rubric ON: {len(neighbor_map)} bag-of-words pairs\n') - - cfg = {'record_type': 'config', 'model': MODEL_ID, 'dataset': args.dataset, - 'n': len(records), 'eval_n': len(eval_records), 'seed': args.seed, - 'pool_offset': args.pool_offset, 'exclude_data_ids': args.exclude_data_ids, - 'excluded_records': data_stats.get('excluded_records', 0), - 'numeric_only': args.numeric_only, 'raw_loaded': data_stats['raw_loaded'], - 'numeric_dropped': data_stats['numeric_dropped'], 'eval_every': args.eval_every, - 'n_skills': args.n_skills, 'view_b_frac': args.view_b_frac, - 'skill_retries': args.skill_retries, 'balance': args.balance, - 'balance_success_frac': args.balance_success_frac, - 'skill_gen_temp': args.skill_gen_temperature, 'reward': 'greedy_binary(correct)', - 'advantage': 'group_relative', 'format_in_reward': args.format_in_reward, 'cache': use_cache, - 'rubric_check': f'fixed_math_{len(_MATH_RUBRIC)}crit+term(viewA)' if checker else 'disabled', - 'xproblem_rubric': args.xproblem_rubric, - 'viewa_sft': args.viewa_sft, 'sft_weight': args.sft_weight, - 'sft_target_len': args.sft_target_len, - 'adv_clip': args.adv_clip, - 'grpo_epsilon': args.grpo_epsilon, 'kl_beta': args.kl_beta, 'lr': args.lr, - 'train_gpus': TRAIN_GPUS, 'ref_gpus': REF_GPUS, 'ref_fsdp': REF_FSDP, - 'train_fsdp': TRAIN_FSDP, 'train_dp': TRAIN_DP, - 'skill_sampler_gpus': SKILL_SAMPLER_GPUS, 'base_sampler_gpus': BASE_SAMPLER_GPUS, - 'max_train_rounds': args.max_train_rounds, 'started': int(time.time())} - sys.stderr.write(f'[rft] raw={data_stats["raw_loaded"]} numeric_drop={data_stats["numeric_dropped"]} ' - f'train={len(records)} eval={len(eval_records)} {args.dataset}; ' - f'train_gpus={TRAIN_GPUS} ref_gpus={REF_GPUS} train_fsdp={TRAIN_FSDP} ' - f'train_dp={TRAIN_DP} skill_dp={skill_dp} base_dp={base_dp}\n') - - hist: List[Dict[str, float]] = [] - rounds = 0 - pool = ProblemPool(records, args.seed) - with open(gen_path, 'w', encoding='utf-8') as gen_f, \ - open(eval_path, 'w', encoding='utf-8') as eval_f, \ - open(data_path, 'w', encoding='utf-8') as data_f, \ - open(train_log_path, 'w', encoding='utf-8') as tlog: - for f in (gen_f, eval_f, data_f, tlog): - _write(f, cfg) - gstep = 0 - # base baseline (frozen sampler, disk-cached, no weight-sync) is prefetched on a - # background thread while the current chunk generates: the skill-gen phase uses - # skill_sampler, so prefetching the NEXT chunk's baseline on base_sampler overlaps - # it and fills base_cache so the next draw_chunk hits it. base_sampler is wrapped in - # _LockedSampler, so the prefetch and this chunk's with-skill scoring never issue a - # base .sample() concurrently. It never touches the trainer or on-policy generation. - prefetch_pool = ThreadPoolExecutor(max_workers=1) if args.prefetch_baseline else None - pending: Optional[Any] = None - - def _prefetch(peeked: List[Dict[str, Any]]) -> None: - if peeked: - baseline_rollout(base_sampler, peeked, base_dp, args, base_cache) - - # Baseline (round 0) eval BEFORE any training: measures the untrained skill model on - # the fixed holdout so every later eval has a step-0 reference point on the same axis. - if eval_records: - eval_recs, eval_summary, eval_metrics = run_greedy_eval( - base_sampler, skill_sampler, eval_records, -1, rounds, base_dp, skill_dp, - args, eval_base_cache) - for rec in eval_recs: - _write(eval_f, rec) - _write(eval_f, eval_summary) - eval_f.flush() - if use_swan: - swanlab.log({f'eval/{k}': v for k, v in eval_metrics.items()}, step=0) - sys.stderr.write( - f'[eval] g-1 (init): n={eval_summary["n"]} viewB mean@1 ' - f'acc={eval_summary["baseline_acc_mean1"]:.3f}->{eval_summary["acc_mean1"]:.3f} ' - f'lift={eval_summary["lift_mean1"]:+.3f} ' - f'fmt={eval_summary["format_mean1"]:.2f} rounds={rounds}\n') - - # Each chunk is drawn fresh + RE-GENERATED with the current policy (on-policy); - # with --balance, draw_chunk keeps drawing until the base fail:success mix hits target. - while rounds < args.max_train_rounds: - if pending is not None: - pending.result() # finish last round's prefetch before drawing (cache-warm) - pending = None - chunk, balance = draw_chunk(pool, base_sampler, base_dp, args, base_cache) - if prefetch_pool is not None: - peeked = pool.peek(int(args.chunk_size * (args.balance_max_draws_mult if args.balance else 1))) - pending = prefetch_pool.submit(_prefetch, peeked) - full, summary, groups = process_chunk( - base_sampler, skill_sampler, chunk, gstep, base_dp, skill_dp, - args, checker, rubric_cache, base_cache, neighbor_map, pool_answers) - summary['balance'] = balance - - log = None - if groups: - log = _train_chunk(skill_model, ref_model, ckpt, groups, args) - rounds += 1 - log.update({'record_type': 'train_round', 'round': rounds, 'chunk': gstep, - 'epoch': pool.epoch, 'ts': int(time.time())}) - _write(tlog, log) - tlog.flush() - if rounds % args.save_rounds == 0: - skill_model.save(f'skill-rft-{rounds}', output_dir=args.output_dir) - - summary['rounds_done'], summary['epoch'] = rounds, pool.epoch - for rec in full: - _write(gen_f, rec) - _write(gen_f, summary) - gen_f.flush() - for v in groups: - _write(data_f, v) - data_f.flush() - - sa, sb, sig = summary['view_A'], summary['view_B'], summary['signal'] - hist.append({'aA': sa['adoption_rate'], 'aB': sb['adoption_rate'], - 'lift': summary['avg_lift'], 'pos': summary['n_reward_pos'], - 'zero_grad': sig['zero_grad_frac']}) - bal_str = (f'bal {balance["selected_fail"]}f/{balance["selected_success"]}s ' - f'(drew {balance["n_drawn"]}/fresh {balance["n_baseline_fresh"]}' - + ('!' if balance.get('budget_hit') else '') + ') ') if balance.get('enabled') else '' - xp = summary.get('xproblem') - xp_str = f'pair={xp["pair_rate"]:.2f}@{xp["neighbor_sim_mean"]:.2f} ' if xp else '' - tr_str = (f'train={log["n_grpo"]}g/{log["n_sft"]}s ' if log - else f'train={summary["n_train_samples"]} ') - sys.stderr.write( - f'[gen] e{pool.epoch} g{gstep}: {bal_str}n={summary["n"]} ' - f'clean={summary["n_clean"]} leak={summary["leak_rate"]:.2f} {tr_str}' - f'0grad={sig["zero_grad_frac"]:.2f} R={sig["reward_mean"]:.2f}±{sig["reward_std"]:.2f} ' - f'acc={summary["avg_baseline_pass"]:.2f}->{summary["avg_withskill_pass"]:.2f} ' - f'lift={summary["avg_lift"]:+.3f} {xp_str}' - f'A[{sa["n"]} {sa["adoption_rate"]:.2f}] B[{sb["n"]} {sb["adoption_rate"]:.2f}] ' - f'rounds={rounds}\n') - if use_swan: - swan_metrics = _swan_metrics(summary, log) - swan_metrics.update(_view_a_rubric_leak_metrics(chunk, pool_answers)) - swanlab.log(swan_metrics, step=gstep) - - if eval_records and (gstep + 1) % args.eval_every == 0: - eval_recs, eval_summary, eval_metrics = run_greedy_eval( - base_sampler, skill_sampler, eval_records, gstep, rounds, base_dp, skill_dp, - args, eval_base_cache) - for rec in eval_recs: - _write(eval_f, rec) - _write(eval_f, eval_summary) - eval_f.flush() - if use_swan: - swanlab.log({f'eval/{k}': v for k, v in eval_metrics.items()}, step=gstep) - sys.stderr.write( - f'[eval] g{gstep}: n={eval_summary["n"]} viewB mean@1 ' - f'acc={eval_summary["baseline_acc_mean1"]:.3f}->{eval_summary["acc_mean1"]:.3f} ' - f'lift={eval_summary["lift_mean1"]:+.3f} ' - f'fmt={eval_summary["format_mean1"]:.2f} rounds={rounds}\n') - - if (gstep + 1) % args.trend_every == 0: - tl = _trend_line(hist, args.trend_every, rounds) - if tl: - sys.stderr.write(tl + '\n') - gstep += 1 - - if prefetch_pool is not None: - if pending is not None: - pending.result() - prefetch_pool.shutdown(wait=True) - base_cache.close() - eval_base_cache.close() - rubric_cache.close() - skill_model.save('skill-rft-final', output_dir=args.output_dir) - sys.stderr.write(f'[rft] done: {rounds} rounds over {gstep} chunks / {pool.epoch} epochs -> {data_path}\n') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/legacy/train_skill_v2_ablate.sh b/cookbook/exp/legacy/train_skill_v2_ablate.sh deleted file mode 100644 index 0b859e759..000000000 --- a/cookbook/exp/legacy/train_skill_v2_ablate.sh +++ /dev/null @@ -1,35 +0,0 @@ -#!/bin/bash -# train_skill_v2_ablate.sh — skill 文体消融：toy vs pitfall，各 30 rounds，顺序执行 -# -# 设计（对应探针实验结论 CONCLUSIONS_config.md / CONCLUSIONS_reflexion.md）： -# 1. thinking 可控（--skill-thinking），本轮两个实验均为 off； -# 2. 两个文体方向：toy（异数字玩具题示范）与 pitfall（预判纠错）； -# 同一文体在主链路（query-only）与 buffer B regen（rubric 诊断）下输出格式一致， -# 保证 GRPO 与 SFT 样本分布一致可联合训练； -# 3. 每个实验 --max-train-rounds 30 结束，输出分目录： -# output.ablate_toy/skill_v2 与 output.ablate_pitfall/skill_v2 -# -# 用法: bash cookbook/exp/embedding/train_skill_v2_ablate.sh -# 注：train_skill_v2.sh 末尾以 "$@" 透传附加参数，argparse 同名参数后者覆盖前者， -# 故此处的 --max-train-rounds 30 会覆盖基础脚本里的 1500。 - -set -euo pipefail -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" - -for STYLE in toy pitfall; do - echo "==============================================================" - echo "[ablate] 开始 skill-style=${STYLE} (thinking=off, 30 rounds)" - echo "==============================================================" - OUTPUT_DIR="./output.ablate_${STYLE}/skill_v2" \ - bash "${SCRIPT_DIR}/train_skill_v2.sh" \ - --skill-style "${STYLE}" \ - --skill-thinking off \ - --max-train-rounds 30 - echo "[ablate] skill-style=${STYLE} 完成" - # 中文注释：两次连跑之间等 Ray/vLLM 完全退出，避免引擎初始化竞态（曾复现过一次） - sleep 30 -done - -echo "[ablate] 两个消融实验全部完成：" -echo " toy -> ./output.ablate_toy/skill_v2" -echo " pitfall -> ./output.ablate_pitfall/skill_v2" diff --git a/cookbook/exp/legacy/train_skill_v2_ablate3.sh b/cookbook/exp/legacy/train_skill_v2_ablate3.sh deleted file mode 100755 index e8e5f3285..000000000 --- a/cookbook/exp/legacy/train_skill_v2_ablate3.sh +++ /dev/null @@ -1,62 +0,0 @@ -#!/bin/bash -# train_skill_v2_ablate3.sh — 三路 prompt×thinking 消融，各 40 GRPO rounds，顺序执行。 -# -# 三组（单变量：skill 文体 + thinking）： -# 1) narrative + think -> output.ablate_narrative_think/skill_v2 -# 2) pitfall + think -> output.ablate_pitfall_think/skill_v2 -# 3) pitfall + nothink -> output.ablate_pitfall_nothink/skill_v2 -# -# 关键设计（对应探针结论 + 训练脚本代码事实）： -# A. 离线 SFT/buffer B 暂时关闭：本脚本 unset LLM_BACKUP_*/OPENAI 环境，使 -# build_rubric_checker() 返回 None -> 训练走 "GRPO only"（train_skill_v2.py:1596-1597,1673）。 -# 理由：40 rounds×chunk16≈640 题，远够不到 distill-trigger(150) 与 sft-trigger(100)， -# SFT 本就不会触发；关掉还能省去每条失败轨迹的 qwen-plus 预诊断 API 开销与后台线程。 -# B. think 组把 --skill-max-tokens 覆盖回 8192：基础 .sh 写死 4096，装不下 think 段+完整 -# <skills>，会截断成空块导致 parse 崩（train_skill_v2.py:1455-1457）。nothink 组维持 4096。 -# C. eval 保留（--eval-every/--eval-size 沿用基础 .sh 的 5/200），这是三组对比的产出信号，不能关。 -# -# ⚠ 重要提醒（务必在 swanlab 盯 leak/rate 曲线）： -# 主 GRPO 路径 reward = parseable AND correct，leak 仅作 observability 审计、不进 reward -# （train_skill_v2.py:1195-1197,1208）。buffer B 关闭后训练回路里【没有任何泄漏防御】。 -# 探针实测 narrative+think 净泄漏≈0.46、pitfall+think 也偏高 —— 两个 think 组极可能 -# reward-hacking：靠把答案写进 skill 拿高 reward，reward_mean/eval-lift 会虚高。 -# 解读 think 组时必须同时看 leak/rate；pitfall+nothink 泄漏≈0，是干净对照基线。 - -set -euo pipefail -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" - -# A. 关闭 buffer B / 离线 SFT / rubric 预诊断（GRPO only） -unset LLM_BACKUP_API_KEY LLM_BACKUP_BASE_URL LLM_BACKUP_MODEL OPENAI_API_KEY || true - -ROUNDS="${ROUNDS:-40}" - -# 每行: TAG STYLE THINKING(on/off) SKILL_MAX_TOKENS -RUNS=( - "narrative_think narrative on 8192" - "pitfall_think pitfall on 8192" - "pitfall_nothink pitfall off 4096" -) - -for spec in "${RUNS[@]}"; do - read -r TAG STYLE THINKING SMT <<< "${spec}" - OUT="./output.ablate_${TAG}/skill_v2" - echo "==============================================================" - echo "[ablate3] TAG=${TAG} style=${STYLE} thinking=${THINKING} skill_max_tokens=${SMT} rounds=${ROUNDS}" - echo " 输出目录: ${OUT} (GRPO only, buffer B 已关)" - echo "==============================================================" - OUTPUT_DIR="${OUT}" \ - bash "${SCRIPT_DIR}/train_skill_v2.sh" \ - --skill-style "${STYLE}" \ - --skill-thinking "${THINKING}" \ - --skill-max-tokens "${SMT}" \ - --max-train-rounds "${ROUNDS}" \ - --swanlab-exp "ablate3_${TAG}_$(date +%Y%m%d_%H%M%S)" - echo "[ablate3] ${TAG} 完成 -> ${OUT}" - # 两次连跑之间等 Ray/vLLM 完全退出，避免引擎初始化竞态（曾复现过） - sleep 30 -done - -echo "[ablate3] 三组全部完成：" -echo " narrative+think -> ./output.ablate_narrative_think/skill_v2" -echo " pitfall+think -> ./output.ablate_pitfall_think/skill_v2" -echo " pitfall+nothink -> ./output.ablate_pitfall_nothink/skill_v2" diff --git a/cookbook/exp/skill2lora/analyze_more.py b/cookbook/exp/skill2lora/analyze_more.py deleted file mode 100644 index 1543139dc..000000000 --- a/cookbook/exp/skill2lora/analyze_more.py +++ /dev/null @@ -1,172 +0,0 @@ -#!/usr/bin/env python3 -"""analyze_more.py — logp_corr 数据的补充相关性分析（7 项此前未做的）。纯 CPU。 -用法：/usr/local/bin/python3 analyze_more.py -""" -import json -import math -import os -from collections import defaultdict - -import numpy as np - -D = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'logp_corr') -pairs = [json.loads(l) for l in open(os.path.join(D, 'pairs.jsonl'))] -problems = json.load(open(os.path.join(D, 'problems.json'))) -rolls = [json.loads(l) for l in open(os.path.join(D, 'rollout_results.jsonl'))] -npz = np.load(os.path.join(D, 'token_logps.npz')) - -pas = {r['key']: float(np.mean(r['pass'])) for r in rolls if r['kind'] == 'skill' and r['mode'].startswith('t05')} -tok = {r['key']: float(np.mean(r['tokens'])) for r in rolls if r['kind'] == 'skill' and r['mode'].startswith('t05')} -npass = {r['key']: int(np.sum(r['pass'])) for r in rolls if r['kind'] == 'skill' and r['mode'].startswith('t05')} -gr = {r['key']: float(r['pass'][0]) for r in rolls if r['kind'] == 'skill' and r['mode'] == 'greedy' and r['n']} -bpas = {r['key']: float(np.mean(r['pass'])) for r in rolls if r['kind'] == 'base' and r['mode'].startswith('t05')} -btok = {r['key']: float(np.mean(r['tokens'])) for r in rolls if r['kind'] == 'base' and r['mode'].startswith('t05')} -prob_by = {p['data_id']: p for p in problems} - - -def rank(x): - x = np.asarray(x, float) - o = np.argsort(x, kind='mergesort') - r = np.empty(len(x)) - r[o] = np.arange(len(x)) - for v in np.unique(x): - m = x == v - if m.sum() > 1: - r[m] = r[m].mean() - return r - - -def sp(a, b): - a, b = np.asarray(a, float), np.asarray(b, float) - m = ~(np.isnan(a) | np.isnan(b)) - if m.sum() < 3: - return np.nan - ra, rb = rank(a[m]), rank(b[m]) - if ra.std() == 0 or rb.std() == 0: - return np.nan - return float(((ra - ra.mean()) * (rb - rb.mean())).mean() / (ra.std() * rb.std())) - - -# ============ ① 题目侧特征 -> skill 收益（筛题实证） ============ -print('=' * 70) -print('① 题目侧特征 -> skill 平均收益（n=61 题；spearman 跨题）') -by_p = defaultdict(list) -for pr in pairs: - if pr['pair_id'] in pas: - by_p[pr['data_id']].append(pr) -prows = [] -for did, prs in by_p.items(): - p = prob_by[did] - lv = int(did.split(':')[1]) - lifts = [pas[x['pair_id']] - bpas[did] for x in prs] - passes = [pas[x['pair_id']] for x in prs] - prows.append({'level': lv, 'base_pass': bpas[did], 'base_tok': btok[did], - 'prob_chars': len(p['problem']), 'gt_chars': len(p['gt']), - 'mean_lift': float(np.mean(lifts)), 'grp_std': float(np.std(passes)), - 'frac_helped': float(np.mean([l > 0 for l in lifts]))}) -for fk in ['level', 'base_pass', 'base_tok', 'prob_chars', 'gt_chars']: - v = [r[fk] for r in prows] - print('%-12s vs mean_lift %+0.3f | vs 组可分性(grp_std) %+0.3f | vs frac_helped %+0.3f' % ( - fk, sp(v, [r['mean_lift'] for r in prows]), sp(v, [r['grp_std'] for r in prows]), - sp(v, [r['frac_helped'] for r in prows]))) -bt = np.array([r['base_tok'] for r in prows]) -ml = np.array([r['mean_lift'] for r in prows]) -for lo, hi in [(0, 3000), (3000, 5000), (5000, 9999)]: - m = (bt >= lo) & (bt < hi) - if m.sum(): - print(' base_tok[%d,%d): n=%d mean_lift=%+.3f frac(lift>0)=%.2f' % ( - lo, hi, m.sum(), ml[m].mean(), np.mean([r['frac_helped'] for r, mm in zip(prows, m) if mm]))) - -# ============ ② skill 生成 think 长度 -> 好坏 ============ -print('\n' + '=' * 70) -print('② skillgen_tokens（skill 生成总 token 含 think）组内 vs pass/输出长') -cs1, cs2 = [], [] -for did, prs in by_p.items(): - if len(prs) < 4: - continue - sg = [x.get('skillgen_tokens') or np.nan for x in prs] - c = sp(sg, [pas[x['pair_id']] for x in prs]) - if not np.isnan(c): - cs1.append(c) - c = sp(sg, [tok[x['pair_id']] for x in prs]) - if not np.isnan(c): - cs2.append(c) -print(' vs pass8: mean=%+.3f se=%.3f n=%d' % (np.mean(cs1), np.std(cs1) / np.sqrt(len(cs1)), len(cs1))) -print(' vs exec_tokens: mean=%+.3f se=%.3f n=%d' % (np.mean(cs2), np.std(cs2) / np.sqrt(len(cs2)), len(cs2))) - -# ============ ③ |delta| 当干预强度计：|delta| vs |lift| ============ -print('\n' + '=' * 70) -print('③ |ΔlogP| 是否预测"干预幅度"|lift|（不看方向）') -ad, al, dtk = [], [], [] -for pr in pairs: - k = pr['pair_id'] - if k not in pas or pr.get('logp_delta_train') is None: - continue - ad.append(abs(pr['logp_delta_train'])) - al.append(abs(pas[k] - bpas[pr['data_id']])) - dtk.append(abs(tok[k] - btok[pr['data_id']])) -print(' |delta| vs |lift| 全局 sp=%+.3f (n=%d)' % (sp(ad, al), len(ad))) -print(' |delta| vs |Δexec_tokens| 全局 sp=%+.3f' % sp(ad, dtk)) -cs = [] -for did, prs in by_p.items(): - if len(prs) < 4: - continue - a = [abs(x['logp_delta_train']) for x in prs] - b = [abs(pas[x['pair_id']] - bpas[did]) for x in prs] - c = sp(a, b) - if not np.isnan(c): - cs.append(c) -print(' 组内: mean=%+.3f se=%.3f n=%d' % (np.mean(cs), np.std(cs) / np.sqrt(len(cs)), len(cs))) - -# ============ ④ delta 的位置衰减：skill 影响是否集中在 GT 前段 ============ -print('\n' + '=' * 70) -print('④ per-token delta 的位置分布（四分位段的 mean|delta|，跨 476 对平均）') -qsum = np.zeros(4) -qcnt = 0 -for pr in pairs: - k, did = pr['pair_id'], pr['data_id'] - if f'base|{did}' not in npz.files or f'skill|{k}' not in npz.files: - continue - b, s = npz[f'base|{did}'], npz[f'skill|{k}'] - if len(b) != len(s) or len(b) < 40: - continue - d = np.abs(s - b) - d = d[~np.isnan(d)] - if len(d) < 40: - continue - qs = np.array_split(d, 4) - qsum += np.array([q.mean() for q in qs]) - qcnt += 1 -print(' Q1(前1/4)=%.4f Q2=%.4f Q3=%.4f Q4(末1/4)=%.4f (n=%d)' % (*(qsum / qcnt), qcnt)) - -# ============ ⑤ pass8 分布形态：混沌（U形/过散）还是二项噪声 ============ -print('\n' + '=' * 70) -print('⑤ 混合组 (0<pass8<1) 的 k/8 直方图（U 形=题级混沌，钟形=独立二项）') -ks = [npass[k] for k in npass if 0 < npass[k] < 8] -hist = np.bincount(ks, minlength=9)[1:8] -print(' k=1..7: %s (n=%d)' % (hist.tolist(), len(ks))) -# 对照：以每对自身 p=k/8 的独立二项，条件在 0<k<8 上的期望形状 -exp = np.zeros(7) -for k in ks: - p = k / 8 - probs = np.array([math.comb(8, i) * p**i * (1 - p)**(8 - i) for i in range(1, 8)]) - exp += probs / probs.sum() -print(' 二项参照: %s' % np.round(exp, 1).tolist()) - -# ============ ⑥ greedy vs pass8 强分歧案例的特征（危险区验证） ============ -print('\n' + '=' * 70) -print('⑥ greedy 与真值强分歧（|greedy-pass8|>0.5）案例 vs 其余：输出长度') -dis, rest = [], [] -for k in pas: - if k in gr: - (dis if abs(gr[k] - pas[k]) > 0.5 else rest).append(tok[k]) -print(' 分歧组 n=%d mean_tok=%d p75=%d | 其余 n=%d mean_tok=%d p75=%d' % ( - len(dis), np.mean(dis), np.percentile(dis, 75), len(rest), np.mean(rest), np.percentile(rest, 75))) - -# ============ ⑦ leak clean 分解：leak 到底贡献多少 lift ============ -print('\n' + '=' * 70) -print('⑦ leak 分解（lift = pass8 - base_pass8）') -for name, cond in [('leaked', lambda p: p['leaked']), ('clean', lambda p: not p['leaked'])]: - ls = [pas[p['pair_id']] - bpas[p['data_id']] for p in pairs if p['pair_id'] in pas and cond(p)] - print(' %-7s n=%-4d mean_lift=%+.4f frac(lift>0)=%.2f' % ( - name, len(ls), np.mean(ls), np.mean([x > 0 for x in ls]))) diff --git a/cookbook/exp/skill2lora/code_task.py b/cookbook/exp/skill2lora/code_task.py deleted file mode 100644 index 44e93351c..000000000 --- a/cookbook/exp/skill2lora/code_task.py +++ /dev/null @@ -1,439 +0,0 @@ -"""BigCodeBench task adapter: data / executor prompts / unit-test judging / code rubric. - -为什么存在这个模块（承接 deepmath -> BFCL -> BigCodeBench 三轮 eval-0 探针的结论）： - deepmath —— 76.8% 的失败是"没写完"，救活率几乎全由截断率解释，任何 hint（含乱码）只值一个 - wrapper；rubric 只能说"你超预算了/你在兜圈"，命中率≈随机，rubric skill 增量 −0.056。 - BFCL —— 截断混杂消掉了，但 4B 裸解 0.861 只剩 8% headroom，且最大错误类 43% 是 ground truth - 私有口径问题，judge 看不到答案就无从判断 -> rubric skill 增量 +0.002（零）。 - BigCodeBench —— 判分是跑 unittest：对错是客观的（跑过就是对），而且**失败时机器免费给出可定位的 - 证据**（异常类型 / 断言差异 / 失败用例名）。bcb/bcb_eval0_probe.py 实测（n=274， - nothink，截断 0）：F0_none 0.378、query-only skill 0.382（+0.004）、 - rubric skill 0.513（+0.135，p=4e-5）—— 三个数据集里 rubric 第一次真正有增量。 - 结论（已写入长期记忆）：rubric 有用的前提是"诊断有客观可定位的失败证据"，不是"任务是代码"。 - -本模块只做纯任务逻辑（加载 / prompt / 判分 / rubric 素材），**不 import train_skill_v2**， -所以 v2 可以在模块顶层 import 它而不构成循环依赖。判分口径与 bcb_eval0_probe.py 逐字同源 -（extract_code / _RUNNER / run_tests / _trim_err / spec_constraints 直接搬过来），这样探针读数 -与训练读数可以横比。 -""" -import ast as _ast -import importlib.util -import json -import os -import random -import re -import shutil -import subprocess -import sys -import tempfile -from concurrent.futures import ThreadPoolExecutor -from typing import Any, Dict, List, Optional, Tuple - -_HERE = os.path.dirname(os.path.abspath(__file__)) -DEFAULT_PARQUET = os.path.join(_HERE, '..', '..', '..', 'bigcodebench', 'bcb.parquet') - -# 需要外网 / GUI / 子进程的库：沙箱里会挂或超时，判分噪声与 skill 无关 -> 整题排除。 -EXCLUDE_LIBS = {'requests', 'urllib', 'http', 'smtplib', 'socket', 'ssl', 'ftplib', - 'mechanize', 'wikipedia', 'turtle', 'tkinter', 'subprocess', 'sendgrid', - 'python_http_client', 'django', 'flask', 'flask_login', 'flask_mail', - 'flask_restful', 'flask_wtf', 'wtforms', 'multiprocessing'} -LIB_ALIAS = {'sklearn': 'sklearn', 'cv2': 'cv2', 'PIL': 'PIL', 'bs4': 'bs4', 'yaml': 'yaml', - 'dateutil': 'dateutil', 'Crypto': 'Crypto', 'docx': 'docx', 'pytz': 'pytz', - 'psutil': 'psutil', 'texttable': 'texttable', 'wordcloud': 'wordcloud', - 'skimage': 'skimage', 'PyPDF2': 'PyPDF2'} - - -# =========================================================================== -# 数据 -# =========================================================================== -def _libs(rec) -> List[str]: - v = rec.get('libs') - if isinstance(v, str): - try: - return list(_ast.literal_eval(v)) - except Exception: - return [] - return list(v or []) - - -def _importable(lib: str) -> bool: - try: - return importlib.util.find_spec(LIB_ALIAS.get(lib, lib).split('.')[0]) is not None - except Exception: - return False - - -def load_tasks(path: str, seed: int) -> Tuple[List[Dict[str, Any]], Dict[str, int]]: - """-> (tasks, stats)；tasks 已按 seed 洗牌，每条是一个"判分载荷"（见 payload_of）。""" - import pyarrow.parquet as pq - rows = pq.read_table(path).to_pylist() - keep, drop_missing, drop_excl = [], 0, 0 - for r in rows: - libs = _libs(r) - if set(libs) & EXCLUDE_LIBS: - drop_excl += 1 - continue - if any(not _importable(x) for x in libs): - drop_missing += 1 - continue - keep.append({'task_id': r['task_id'], 'instruct_prompt': r['instruct_prompt'], - 'code_prompt': r['code_prompt'], 'test': r['test'], - 'entry_point': r['entry_point'], 'doc_struct': r['doc_struct'], - 'canonical_solution': r['canonical_solution'], 'libs': libs}) - random.Random(seed).shuffle(keep) - return keep, {'raw': len(rows), 'kept': len(keep), - 'drop_missing_lib': drop_missing, 'drop_needs_net_or_gui': drop_excl} - - -def payload_of(task: Dict[str, Any]) -> Dict[str, Any]: - """训练记录里 ``reference_answer`` 的内容 —— 判分需要的一切。 - - ★ 为什么塞进 reference_answer 而不是新开字段：全流水线（v2 / methods / eval_reflexion）判分 - 都走 ``_parse_seq(seq, r['reference_answer'])`` 这一个入口，把载荷放这里就不用改任何签名， - math 分支也完全不受影响。体积：test 平均 ~3KB，每题每 chunk 落盘一次，可接受。 - """ - return {k: task[k] for k in ('task_id', 'entry_point', 'test', 'code_prompt', 'doc_struct', - 'canonical_solution')} - - -# =========================================================================== -# 代码抽取 + 沙箱跑单测（与 bcb_eval0_probe.py 逐字同源） -# =========================================================================== -def after_think(text: str) -> str: - i = (text or '').rfind('</think>') - return text[i + len('</think>'):] if i >= 0 else (text or '') - - -_FENCE_RE = re.compile(r'```(?:python|py)?\s*\n(.*?)```', re.S) - - -def extract_code(text: str) -> str: - """取最后一个能通过 ast.parse 的代码块；没有围栏就退化为整段（切 think 之后）。""" - body = after_think(text or '') - blocks = _FENCE_RE.findall(body) - for b in reversed(blocks): - try: - _ast.parse(b) - return b - except SyntaxError: - continue - if blocks: - return blocks[-1] - try: - _ast.parse(body) - return body - except SyntaxError: - return '' - - -_RUNNER = """ -import unittest, sys -loader = unittest.TestLoader() -suite = loader.loadTestsFromTestCase(TestCases) -res = unittest.TextTestRunner(verbosity=0, stream=sys.stderr).run(suite) -print('__BCB__', res.testsRun, len(res.failures), len(res.errors)) -sys.exit(0 if res.wasSuccessful() and res.testsRun > 0 else 1) -""" - - -def _trim_err(err: str, limit: int = 1600) -> str: - """保留失败测试名与异常行，砍掉中间冗长的 traceback 帧（这是喂给 rubric 的客观证据）。 - - 随机临时目录名换成 ``<sandbox>``：traceback 帧里带着 /tmp/bcb_xxxxxxx/ 这种每次都不同的 - 路径，对 judge 是纯噪声，还会让同一个失败在两次运行里看起来不一样（gen_records 里逐字 - 比对失败原因时会误判成"变了"）。 - """ - err = re.sub(r'/tmp/bcb_[A-Za-z0-9_]+', '<sandbox>', err or '') - lines = [ln for ln in err.splitlines() if ln.strip()] - keep = [ln for ln in lines - if ln.startswith(('FAIL:', 'ERROR:', 'AssertionError', 'Traceback')) - or re.match(r'^\w*(Error|Exception|Warning)\b', ln.strip()) - or ', in ' in ln] - text = '\n'.join(keep or lines[-25:]) - return text[-limit:] - - -def run_tests(code: str, payload: Dict[str, Any], timeout: int) -> Dict[str, Any]: - """在独立进程 + 临时目录里跑该题自带的 unittest；返回 pass/fail + 客观报错。 - - 隔离手段只有"子进程 + 临时 cwd + 超时"，没有容器/seccomp：真正危险的题（外网 / GUI / - 子进程 / multiprocessing）在 load_tasks 阶段就按 EXCLUDE_LIBS 整题剔除了。 - env 里必须清掉 CUDA_VISIBLE_DEVICES —— 否则 numpy/torch 系的测试可能去抢训练用的卡。 - """ - if not code.strip(): - return {'passed': False, 'kind': 'no_code', 'error': 'no parseable code block', - 'n_tests': 0} - if payload['entry_point'] not in code: - return {'passed': False, 'kind': 'no_entry', - 'error': f"function {payload['entry_point']} is not defined in the submitted code", - 'n_tests': 0} - tmp = tempfile.mkdtemp(prefix='bcb_') - try: - src = code + '\n\n' + payload['test'] + '\n' + _RUNNER - path = os.path.join(tmp, 'run_case.py') - with open(path, 'w', encoding='utf-8') as f: - f.write(src) - env = dict(os.environ, MPLBACKEND='Agg', PYTHONHASHSEED='0', - OMP_NUM_THREADS='1', MKL_NUM_THREADS='1', TOKENIZERS_PARALLELISM='false') - env.pop('CUDA_VISIBLE_DEVICES', None) - try: - p = subprocess.run([sys.executable, path], cwd=tmp, env=env, timeout=timeout, - capture_output=True, text=True, errors='replace') - except subprocess.TimeoutExpired: - return {'passed': False, 'kind': 'timeout', - 'error': f'the tests did not finish within {timeout}s', 'n_tests': 0} - out, err = p.stdout or '', p.stderr or '' - n_tests = n_fail = n_err = 0 - for line in out.splitlines(): - if line.startswith('__BCB__'): - _, a, b, c = line.split() - n_tests, n_fail, n_err = int(a), int(b), int(c) - if p.returncode == 0 and n_tests > 0: - return {'passed': True, 'kind': 'pass', 'error': '', 'n_tests': n_tests} - kind = 'assertion' if n_fail else ('exception' if n_err else 'import_or_syntax') - return {'passed': False, 'kind': kind, 'error': _trim_err(err.replace(tmp, '<sandbox>')), - 'n_tests': n_tests} - finally: - shutil.rmtree(tmp, ignore_errors=True) - - -def judge_many(items: List[Optional[Tuple[str, Any, int, Dict[str, Any]]]], - workers: int, timeout: int) -> List[Dict[str, Any]]: - """批量判分。``items[i]`` = (text, stop_reason, gen_tokens, payload) 或 None（无输出）。 - - ★ 必须批量：单测是子进程（导入 pandas/sklearn 后典型 1-3s），而一个 E4 chunk 有 ~290 次判分。 - 串行 ≈12 分钟/chunk，远超同 chunk 的 GPU 时间；线程池（--test-workers）把它压到 ~30s。 - 额外去重：同一题的多个 skill 在 T=0 executor 下经常产出逐字相同的代码，去重后实测省下可观的 - 子进程数（同 (task_id, code) 只跑一次）。 - """ - rolls: List[Dict[str, Any]] = [] - keys: List[Optional[Tuple[str, str]]] = [] - jobs: Dict[Tuple[str, str], Dict[str, Any]] = {} # key -> payload（去重后的待跑集合） - for it in items: - if it is None: - rolls.append(empty_roll()) - keys.append(None) - continue - text, stop, ntok, payload = it - code = extract_code(text) - key = (payload['task_id'], code) - rolls.append({'pred': None, 'correct': False, - 'terminated': stop != 'length', 'stop_reason': stop, - 'gen_tokens': int(ntok or 0), 'text': text, 'code': code, - 'kind': None, 'error': '', 'n_tests': 0}) - keys.append(key) - jobs.setdefault(key, payload) - if jobs: - todo = list(jobs) - with ThreadPoolExecutor(max_workers=max(1, min(workers, len(todo)))) as ex: - res = list(ex.map(lambda k: run_tests(k[1], jobs[k], timeout), todo)) - verdicts = dict(zip(todo, res)) - for r, key in zip(rolls, keys): - v = verdicts.get(key) if key is not None else None - if v is None: - continue - r['correct'] = bool(v['passed']) - # pred 在数学分支是"抽出来的答案，抽不到就是 None"，下游 term/answered_rate 与 - # acc/answered_pass 正是按 "pred is not None" 定义"交了可判的答案"这条通道。 - # ⚠️ 所以这里不能无条件写 kind：kind 恒非空会让 answered_rate 恒为 1.000、 - # answered_pass 退化成 candidate_pass，那两条曲线静默失效。口径对齐为： - # 抽到代码块 -> pred = 判分结论（pass/assertion/...，便于审计）；没抽到 -> None。 - r['pred'] = v['kind'] if r['code'] else None - r['kind'], r['error'], r['n_tests'] = v['kind'], v['error'], v['n_tests'] - return rolls - - -def empty_roll() -> Dict[str, Any]: - return {'pred': None, 'correct': False, 'terminated': False, 'stop_reason': 'empty', - 'gen_tokens': 0, 'text': '', 'code': '', 'kind': 'no_code', 'error': '', 'n_tests': 0} - - -def selftest(tasks: List[Dict[str, Any]], workers: int, timeout: int) -> List[str]: - """参考解答必须跑过它自己的单测 —— 跑不过说明沙箱/依赖不可判定，不是模型的错。 - 返回跑不过的 task_id 列表（调用方据此剔题）。""" - codes = [t['code_prompt'] + (t.get('canonical_solution') or '') for t in tasks] - payloads = [payload_of(t) for t in tasks] - with ThreadPoolExecutor(max_workers=max(1, min(workers, len(tasks) or 1))) as ex: - vers = list(ex.map(lambda p: run_tests(p[0], p[1], timeout), zip(codes, payloads))) - return [t['task_id'] for t, v in zip(tasks, vers) if not v['passed']] - - -# =========================================================================== -# executor prompts -# =========================================================================== -# 本数据集的硬性交付要求（BigCodeBench 官方 instruct 模式口径）。所有臂共用。 -EXEC_SYSTEM = """\ -You are an expert Python engineer. You will be given a task description that ends with the exact \ -import lines and function signature your solution must start with. - -Deliver exactly one fenced Python code block and nothing else after it: -- Reproduce the given imports and the given function signature verbatim, including parameter \ -names, order and default values. -- Add any further imports you need inside the same block; the block must run standalone. -- Return exactly the object type the task says to output. If it says the function should output \ -a tuple, return a tuple in that order; if it names a matplotlib Axes, return the Axes object \ -itself, not the Figure and not None. -- Implement the described behaviour for the general case, including the empty / single-element / \ -missing-column edge cases and any exception the description says to raise. -- Do not call the function, do not print demonstrations, do not add tests, do not use \ -`if __name__ == '__main__'`, and do not read from stdin. -- Do not include explanations outside the code block.""" - -# skill hint 包装语：与 v2 的数学版逐字同构（只把 "problem" 换成 "task"），保证 E4/E17 的 -# executor 输入除 skill 文本外没有第二个变量。 -_WRAPPER = ('Skill hint:\nFor this task, a skill-generation model has analyzed it and ' - 'provided some advisory skills:\n{hint}\n' - 'Prefer using its techniques when they fit, but if you have a clearly better ' - 'implementation, you may diverge. Be concise and accurate.\n') - - -def direct_prompt(problem: str) -> Dict[str, Any]: - return {'messages': [{'role': 'system', 'content': EXEC_SYSTEM}, - {'role': 'user', 'content': problem}]} - - -def skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: - skill = (skill or '').strip() - if not skill: - return direct_prompt(problem) - return {'messages': [{'role': 'system', 'content': EXEC_SYSTEM}, - {'role': 'user', - 'content': problem + '\n\n' + _WRAPPER.format(hint=skill)}]} - - -# =========================================================================== -# skill-gen prompts -# =========================================================================== -# E4（view B, query-only）。与数学版同构：先私下把题做一遍，再只写可迁移的方法论； -# 硬禁止写出解法代码与本题字面量（否则 skill 就是抄答案，测不到"方法论有没有用"）。 -SKILLGEN_SYSTEM = """\ -You are a skill-generation model for a Python implementation task. Your <skills> block will be fed to a SEPARATE downstream engineer model that must write the function on its own. The engineer sees the same task description and the same required signature, but NOT your private reasoning. - -First think privately: actually work out how you would implement it, including which library calls do the work. Then step back and write, inside <skills></skills>, transferable guidance for THIS TYPE of task: which library functions are the right tool and what their relevant arguments and return shapes are, how to get the return value into the exact type the task demands, which edge cases and exceptions this kind of task always has, and the common mistakes to avoid. - -CRITICAL: do NOT write the solution code, and do NOT paste concrete literal values from this task. Name the API and describe the shape of the answer instead of writing it out. -Keep it to roughly one focused paragraph. Put ONLY the guidance inside <skills></skills>.""" - - -def skillgen_prompt(problem: str) -> Dict[str, Any]: - return {'messages': [{'role': 'system', 'content': SKILLGEN_SYSTEM}, - {'role': 'user', 'content': f'Task:\n{problem}'}]} - - -# =========================================================================== -# rubric（判据 + judge prompt + 诊断素材） -# =========================================================================== -DIAG_SYSTEM = """\ -You are a strategy-level code reviewer. You are given a Python task description (with the required signature), a rubric, one attempted implementation, and the REAL error that attempt produced when the task's unit tests were run. Decide PASS or FAIL for each criterion, and write the diagnosis so it becomes reusable guidance for similar tasks without seeing this attempt again. - -Output STRICT JSON (no prose outside it) with this shape: -{"items": [{"index": 1, "verdict": "PASS"|"FAIL", "reason": "...", "fix": ""}], "overall": "OK"|"ISSUES", "summary": "..."} - -Rules: -- Ground every FAIL in the task description or the test error you were given; do not speculate. -- The test error is authoritative evidence: if it names an exception, a wrong type or a failed assertion, the criterion it implicates must be FAIL. -- Never write out corrected code; describe the process problem at strategy level. -- A fix suggests the local correction direction without implementing it. -- Keep "reason" and "fix" concise: one short sentence each. -- Output only the JSON object.""" - -DIAG_USER = """\ -## Task -{query} - -## Rubric -{rubric} - -## Attempted implementation and its test error -{segment} - -Now output the diagnostic JSON object.""" - -# PASS = 该类问题不存在（正向陈述，与 v2._format_diagnosis / gate 语义一致）。 -# 判据按 BigCodeBench 的实际失败模式组织：签名/导入、返回类型、API 用法、选库、边界、异常、逻辑。 -CODE_RUBRIC = [ - ('The implementation is runnable as given: it defines the required function with the exact ' - 'signature asked for and imports everything it uses', True), - ('The value returned matches the output type and structure the task states, element for ' - 'element and in the stated order', False), - ('The library functions used exist and are called with arguments and keyword names that ' - 'those functions actually accept', False), - ('The library chosen for each step is the one the task asks for, used for its intended ' - 'purpose rather than reimplemented by hand', False), - ('Edge cases the task implies (empty input, single element, missing key or column, ' - 'duplicate values) are handled instead of crashing', False), - ('Exactly the exceptions the task specifies are raised for invalid input, and no others ' - 'leak out', False), - ('The core computation implements what the description asks, with no step skipped, ' - 'inverted, or replaced by a placeholder', False), -] -# 版本号进 rubric 缓存键（rubric_cache._key）：判据一改旧诊断必须失效。code 与 math 的诊断 -# 还额外分文件存（trainer 按 task 选文件名），双保险。 -RUBRIC_VERSION = 'rubric_code_v1' - - -def spec_constraints(payload: Dict[str, Any]) -> str: - """任务自身声明的硬约定（签名、必需库、返回规格、应抛异常、文档示例）。 - - 只用题面信息、不含参考解答 —— 训练时同样拿得到，所以是"可得且非泄漏"的判据依据。 - ⚠️ BFCL 那边同类做法（schema_constraints）没能提升命中率，因为那边的对错定义在 gt 私有约定 - 里；这里返回类型/异常是题面明写的，所以这次它是真判据。 - """ - try: - doc = payload['doc_struct'] - doc = json.loads(doc) if isinstance(doc, str) else (doc or {}) - except Exception: - doc = {} - lines = [f"- required signature (must be reproduced verbatim):\n{payload['code_prompt'].strip()}"] - for key, label in (('reqs', 'must use these libraries'), ('returns', 'must return'), - ('raises', 'must raise'), ('params', 'parameters')): - vals = [str(x).strip() for x in (doc.get(key) or []) if str(x).strip()] - if vals: - lines.append(f'- {label}: ' + '; '.join(vals)[:400]) - ex = [str(x) for x in (doc.get('examples') or [])][:8] - if ex: - lines.append('- documented example calls:\n ' + '\n '.join(ex)) - return '\n'.join(lines) - - -def diag_query(problem: str, payload: Dict[str, Any]) -> str: - return problem + '\n\nHard requirements declared by the task:\n' + spec_constraints(payload) - - -def diag_segment(roll: Dict[str, Any]) -> str: - """★ rubric 路线唯一真正有效的一处：给 judge 的不是"输出全文"，而是**提交的代码 + 单测真实 - 报错**。<think> 已在 extract_code 里切掉。BFCL 那轮 judge 手上没有任何客观证据，命中率 25% - ≈ 随机；这里报错是客观事实且不含参考解答。""" - return (f"### Submitted code\n```python\n{roll.get('code') or '(no parseable code block)'}\n```" - f"\n\n### Result of running the task's unit tests\n" - f"outcome: {roll.get('kind') or 'unknown'}\n" - f"{roll.get('error') or '(no error output)'}") - - -# =========================================================================== -# leak / skill 文本监控（代码域口径） -# =========================================================================== -def _canon_lines(payload: Dict[str, Any]) -> List[str]: - out = [] - for ln in (payload.get('canonical_solution') or '').splitlines(): - s = ln.strip() - if len(s) >= 20 and not s.startswith(('#', 'import ', 'from ', 'def ', 'return')): - out.append(s) - return out - - -def leaked(skill: str, payload: Any) -> bool: - """代码域 leak = skill 里出现了参考解答的实质代码行（>=20 字符、非 import/def/注释）。 - 与数学域一致：**只做监控，永不进 reward**（项目既定规则）。""" - if not skill or not isinstance(payload, dict): - return False - return any(ln in skill for ln in _canon_lines(payload)) - - -def skill_has_code(skill: str) -> bool: - """skill 退化监控：本该是方法论的块里出现了代码围栏或成段 def/return。 - 数学域的 digit_fraction / no_math_sentence 在代码域无意义，用这条替代。""" - s = skill or '' - if '```' in s: - return True - return bool(re.search(r'^\s*(def |return |import |for .*:|if .*:)', s, re.M)) diff --git a/cookbook/exp/skill2lora/good_skill_hard_fail/analyze_3way.py b/cookbook/exp/skill2lora/good_skill_hard_fail/analyze_3way.py deleted file mode 100644 index a80756be7..000000000 --- a/cookbook/exp/skill2lora/good_skill_hard_fail/analyze_3way.py +++ /dev/null @@ -1,177 +0,0 @@ -#!/usr/bin/env python3 -# analyze_3way.py — 三路探针对比分析（A/B/C），全部结论可复现、可引用。 -# A = 老环境 nothink : skillcfg_full_off.jsonl + reflexion_full_off.jsonl (根目录) -# B = 新环境 think : skillcfg_full_on.jsonl + reflexion_full_on.jsonl (根目录) -# C = 新环境 nothink : env_runs/vllm_0.23.0/skillcfg_full_off.jsonl + reflexion_full_off.jsonl -# 对比轴: -# vLLM/环境影响 = A vs C (都 nothink) -# think 影响 = C vs B (都新环境) -# 用法: python3 analyze_3way.py > analysis_out/report.txt 2>&1 -import json, os, math, statistics as st, re, random -from collections import defaultdict, Counter - -def _f(x): - try: - return float(x) - except (TypeError, ValueError): - return 0.0 - -HERE = os.path.dirname(os.path.abspath(__file__)) -GROUPS = { - 'A_old_nothink': ('skillcfg_full_off.jsonl', 'reflexion_full_off.jsonl'), - 'B_new_think': ('skillcfg_full_on.jsonl', 'reflexion_full_on.jsonl'), - 'C_new_nothink': ('env_runs/vllm_0.23.0/skillcfg_full_off.jsonl', - 'env_runs/vllm_0.23.0/reflexion_full_off.jsonl'), -} -SKILL_CFGS = ['P1_narrative','P2_combo','P3_toy','P4_card','P5_pitfall','P6_seam','P7_minimal'] -REFL_CFGS = ['R4_blind','D1_needle','D2_narr','D3_toyfix'] - - -def load(path): - """流式读取 -> list[dict]（只保留分析需要的字段，省内存）""" - keep = ('config','data_id','sample_idx','baseline_pass','parseable','skill_chars', - 'leaked','skillgen_stop','skillgen_tokens','withskill_correct','withskill_stop', - 'withskill_tokens','skill') - rows = [] - with open(os.path.join(HERE, path)) as f: - for line in f: - line = line.strip() - if not line: - continue - r = json.loads(line) - rows.append({k: r.get(k) for k in keep}) - return rows - - -def agg_metrics(rows, cfgs): - """返回 {config: metrics}""" - out = {} - for name in cfgs: - sub = [t for t in rows if t['config'] == name] - n = len(sub) - if n == 0: - continue - parse = sum(bool(t['parseable']) for t in sub)/n - leak = sum(bool(t['leaked']) for t in sub)/n - trunc = sum(1 for t in sub if t['withskill_stop']=='length')/n - sg_trunc = sum(1 for t in sub if t['skillgen_stop']=='length')/n - chars = [t['skill_chars'] for t in sub if t['parseable']] - med_chars = int(st.median(chars)) if chars else 0 - acc = sum(bool(t['withskill_correct']) for t in sub)/n - base = sum(_f(t['baseline_pass']) for t in sub)/n - # 题级 pass@k / hard 救活@k - byq = defaultdict(list) - for t in sub: - byq[t['data_id']].append(t) - p_at_k = sum(1 for v in byq.values() if any(x['withskill_correct'] for x in v))/len(byq) - hardq = {d:v for d,v in byq.items() if _f(v[0]['baseline_pass'])==0} - rescue = (sum(1 for v in hardq.values() if any(x['withskill_correct'] for x in v))/len(hardq)) if hardq else 0.0 - # 去混杂子集: leaked=0 且 parseable=1 - clean = [t for t in sub if (not t['leaked']) and t['parseable']] - acc_clean = (sum(bool(t['withskill_correct']) for t in clean)/len(clean)) if clean else float('nan') - out[name] = dict(n=n, parse=parse, leak=leak, trunc=trunc, sg_trunc=sg_trunc, - med_chars=med_chars, acc=acc, base=base, lift=acc-base, - p_at_k=p_at_k, rescue=rescue, n_clean=len(clean), acc_clean=acc_clean, - sg_tokens_med=int(st.median([t['skillgen_tokens'] or 0 for t in sub]))) - return out - - -def diversity(rows, cfgs): - """题内 8 rollout 多样性: 去重率 + 词级 pairwise Jaccard(距离) + 字符长度 CV""" - _word = re.compile(r"[A-Za-z]+|\d+") - out = {} - for name in cfgs: - sub = [t for t in rows if t['config']==name] - byq = defaultdict(list) - for t in sub: - byq[t['data_id']].append(t) - uniq_ratios, jac_dists, all_chars = [], [], [] - for v in byq.values(): - skills = [(t['skill'] or '') for t in v] - all_chars += [len(s) for s in skills] - uniq_ratios.append(len(set(skills))/len(skills)) - sets = [set(_word.findall(s.lower())) for s in skills] - ds = [] - for i in range(len(sets)): - for j in range(i+1, len(sets)): - a,b = sets[i],sets[j] - if not a and not b: - ds.append(0.0); continue - inter=len(a&b); uni=len(a|b) or 1 - ds.append(1 - inter/uni) # 1=完全不同,0=完全相同 - if ds: - jac_dists.append(sum(ds)/len(ds)) - cv = (st.pstdev(all_chars)/ (sum(all_chars)/len(all_chars))) if all_chars and sum(all_chars) else 0.0 - out[name] = dict(uniq=sum(uniq_ratios)/len(uniq_ratios), - jac=sum(jac_dists)/len(jac_dists) if jac_dists else 0.0, - char_cv=cv) - return out - - -def fmt_table(title, gm, cfgs, cols): - print(f"\n### {title}") - head = "%-14s " % "config" + " ".join("%-9s" % c for c,_ in cols) - print(head); print("-"*len(head)) - for name in cfgs: - if name not in gm: - continue - m = gm[name] - row = "%-14s " % name + " ".join(("%-9.3f" if isinstance(m[k],float) else "%-9d") % m[k] for _,k in cols) - print(row) - - -def main(): - os.makedirs(os.path.join(HERE,'analysis_out'), exist_ok=True) - data = {} - for g,(sf,rf) in GROUPS.items(): - data[g] = dict(skill=load(sf), refl=load(rf)) - print(f"[load] {g}: skillcfg={len(data[g]['skill'])} reflexion={len(data[g]['refl'])}") - - # ---------- 对齐校验: 三组是否同题同 idx ---------- - print("\n" + "="*80 + "\n[对齐校验] 三组 (config,data_id,sample_idx) 键集合是否一致") - def keyset(rows): - return set((t['config'],t['data_id'],t['sample_idx']) for t in rows) - ka,kb,kc = keyset(data['A_old_nothink']['skill']),keyset(data['B_new_think']['skill']),keyset(data['C_new_nothink']['skill']) - print(f" A∩C 交集/并集(nothink 对比): {len(ka&kc)}/{len(ka|kc)} A独有={len(ka-kc)} C独有={len(kc-ka)}") - print(f" C∩B 交集/并集(think 对比): {len(kc&kb)}/{len(kc|kb)} C独有={len(kc-kb)} B独有={len(kb-kc)}") - - cols_skill = [('n','n'),('parse','parse'),('leak','leak'),('trunc','trunc'), - ('sgTrunc','sg_trunc'),('chars','med_chars'),('base','base'), - ('acc@1','acc'),('lift','lift'),('pass@k','p_at_k'),('rescue@k','rescue'), - ('accClean','acc_clean')] - # ---------- Q1+Q4: 每组每 config 指标 ---------- - print("\n" + "="*80 + "\n[Q1/Q4] skillcfg 每类准确率与质量指标") - GM = {} - for g in GROUPS: - GM[g] = agg_metrics(data[g]['skill'], SKILL_CFGS) - fmt_table(f"{g} (skillcfg)", GM[g], SKILL_CFGS, cols_skill) - - print("\n" + "="*80 + "\n[Q1/Q4] reflexion 每类救活指标") - RM = {} - for g in GROUPS: - RM[g] = agg_metrics(data[g]['refl'], REFL_CFGS) - fmt_table(f"{g} (reflexion)", RM[g], REFL_CFGS, cols_skill) - - # ---------- Q2: 横向差分 ---------- - print("\n" + "="*80 + "\n[Q2] 环境(vLLM/栈)影响 = C - A (同 nothink) ；think 影响 = B - C (同新环境)") - print("%-14s %-22s %-22s" % ("config","env(C-A) acc/lift/parse","think(B-C) acc/lift/parse")) - for name in SKILL_CFGS: - a,c,b = GM['A_old_nothink'].get(name),GM['C_new_nothink'].get(name),GM['B_new_think'].get(name) - if not(a and c and b): continue - env = f"{c['acc']-a['acc']:+.3f}/{c['lift']-a['lift']:+.3f}/{c['parse']-a['parse']:+.3f}" - thk = f"{b['acc']-c['acc']:+.3f}/{b['lift']-c['lift']:+.3f}/{b['parse']-c['parse']:+.3f}" - print("%-14s %-22s %-22s" % (name, env, thk)) - - # ---------- Q3: 多样性 ---------- - print("\n" + "="*80 + "\n[Q3] skill 生成多样性 (题内 8 rollout；uniq=去重率 jac=词级平均两两距离 char_cv=长度变异)") - for g in GROUPS: - dv = diversity(data[g]['skill'], SKILL_CFGS) - print(f"\n### {g}") - print("%-14s %-8s %-8s %-8s" % ("config","uniq","jac","char_cv")) - for name in SKILL_CFGS: - m=dv[name]; print("%-14s %-8.3f %-8.3f %-8.3f" % (name,m['uniq'],m['jac'],m['char_cv'])) - - print("\n[done] 详见各节表格；抽样交叉验证见 sample_probe.py 输出") - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/skill2lora/good_skill_hard_fail/eval_skill_probe.py b/cookbook/exp/skill2lora/good_skill_hard_fail/eval_skill_probe.py deleted file mode 100644 index 48c3f6578..000000000 --- a/cookbook/exp/skill2lora/good_skill_hard_fail/eval_skill_probe.py +++ /dev/null @@ -1,286 +0,0 @@ -#!/usr/bin/env python3 -"""eval_skill_probe.py — 自包含的"单次 eval"探针，用于手工迭代 skill。 - -目的：固定一批难题（executor 持续解不出的），保持 executor 侧 prompt/采样/判分与 -train_skill_v2.py 的 eval 完全一致（v2 模式：单 user turn + \\boxed{} 答案格式， -executor 用 base_sampler、enable_thinking=True、greedy 温度 0），但允许自由替换每题 -的 skill(experience) 内容。反复替换 skill 重跑，一旦某题解对，就把该 skill 落盘到 -winning_skills.jsonl，从而观察"能让 executor 生效的 skill 到底长什么样"。 - -不 import train_skill_v2.py（自包含）；executor 侧逻辑逐段复刻自该文件（2026-07）。 - -用法： - # 1) 生成 trials 模板（默认 10 道难题，skill 待填）： - python3 eval_skill_probe.py --init - # 2) 编辑 trials.jsonl，给每题填不同的 skill，然后跑： - python3 eval_skill_probe.py - # 只测某题 / 调大 max_tokens（验证"截断"型失败是否靠加长度能救）： - python3 eval_skill_probe.py --only seam:val:128 --max-tokens 12000 - -trials.jsonl 每行一个 JSON（# 开头的行会被忽略，可当注释）： - {"data_id": "seam:val:128", "skill": "..."} # 用该 skill 解题 - {"data_id": "seam:val:128", "skill": ""} # 空 skill = baseline - {"data_id": "seam:val:128", "tag": "v3", "skill":"..."}# tag 便于区分同题多次试验 -problem / reference_answer 默认按 data_id 从 eval_records.jsonl 解析；也可在行内直接 -提供 "problem" / "reference_answer" 覆盖（用于测 eval 集之外的题）。 - -GPU：仅起一个 executor sampler。默认 EXEC_GPUS=2；若训练在占卡，先用 -CUDA_VISIBLE_DEVICES 指定空闲卡，或设 EXEC_GPUS=1。 -""" -import argparse -import copy -import json -import os -import re -import sys -from typing import Dict, Optional - -import twinkle -from twinkle import DeviceGroup, DeviceMesh -from twinkle.data_format import SamplingParams -from twinkle.sampler import vLLMSampler -from twinkle.template import Template - -SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) -MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') -GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.8)) -EXEC_GPUS = int(os.environ.get('EXEC_GPUS', 2)) -MAX_MODEL_LEN = int(os.environ.get('MAX_MODEL_LEN', 16384)) -DEFAULT_EVAL_RECORDS = os.path.join(SCRIPT_DIR, '..', 'skill_v2', 'eval_records.jsonl') - -# 与本文件夹 10 个 case 对应的难题（executor 全程持续解不出、baseline 也全 0）。 -DEFAULT_DATA_IDS = ['seam:val:81', 'seam:val:92', 'seam:val:148', 'seam:val:128', 'seam:val:127', - 'seam:val:29', 'seam:val:94', 'seam:val:176', 'seam:val:12', 'seam:val:151'] - -# =========================================================================== -# executor prompt / answer format —— 逐字复刻 train_skill_v2.py 的 v2 分支 -# =========================================================================== -_ANSWER_FORMAT_V2 = ('Present your reasoning, then put ONLY the final numeric result inside ' - '\\boxed{}. For example: \\boxed{42}.') - - -def build_direct_prompt(problem): - content = f'The problem you need to solve:\n{problem}\n\n' + _ANSWER_FORMAT_V2 - return {'messages': [{'role': 'user', 'content': content}]} - - -def build_skill_solve_prompt(problem, skill): - skill = (skill or '').strip() - if not skill: - return build_direct_prompt(problem) - content = (f'The problem you need to solve:\n{problem}\n\n' - 'Skill hint:\nFor this problem, a skill-generation model has analyzed it and ' - 'provided some advisory skills:\n' - f'{skill}\n' - 'Prefer using its techniques when they fit, but if you have a more efficient or ' - 'clearer correct method, you may use it. If you diverge from this advice, briefly ' - 'explain why. Be concise and accurate.\n' - + _ANSWER_FORMAT_V2) - return {'messages': [{'role': 'user', 'content': content}]} - - -# =========================================================================== -# boxed 抽取 + SEAM lpem 风格数值判分 —— 逐字复刻 -# =========================================================================== -_BOXED_RE = re.compile(r'\\boxed\s*\{') - - -def extract_boxed(text: str) -> Optional[str]: - if not text: - return None - last = None - for m in _BOXED_RE.finditer(text): - depth, i = 1, m.end() - while i < len(text) and depth > 0: - depth += (text[i] == '{') - (text[i] == '}') - i += 1 - if depth == 0: - last = text[m.end():i - 1].strip() - return last - - -_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') - - -def _clean_text(decoded: Optional[str]) -> str: - return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() - - -_SEAM_TAG_RE = re.compile(r'<\s*answer\s*>(.*?)<\s*/\s*answer\s*>', re.I | re.S) -_SEAM_BOX_RE = re.compile(r'(?:\\{1,2}\(|)\\{1,2}boxed\s*\{\s*([^}]*)\s*}(?:\)|)', re.S) -_SEAM_INLINE_RE = re.compile(r'\$([^$]+)\$|\\\(([^)]+)\\\)', re.S) -_SEAM_FRAC_RE = re.compile(r'(-?\d+(?:\.\d+)?)/(-?\d+(?:\.\d+)?)') -_SEAM_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') - - -def _seam_norm(num: str) -> str: - try: - f = float(num) - return str(int(f)) if f == int(f) else str(f) - except Exception: - return num.strip() - - -def _seam_sanitize(txt: str) -> str: - txt = (txt or '').strip() - if (m := _SEAM_TAG_RE.search(txt)): - txt = m.group(1).strip() - elif (m := _SEAM_BOX_RE.search(txt)): - txt = m.group(1).strip() - elif (m := _SEAM_INLINE_RE.search(txt)): - txt = (m.group(1) or m.group(2)).strip() - # bugfix 2026-07-29：\dfrac/\tfrac/\cfrac 归一，同 train_skill_v2._seam_sanitize - txt = txt.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac').replace(r'\cfrac', r'\frac') - txt = re.sub(r'\\frac\s*\{\s*([^}]+?)\s*}\s*\{\s*([^}]+?)\s*}', r'\1/\2', txt) - if (m := _SEAM_FRAC_RE.search(txt)): - p, q = map(float, m.groups()) - if q: - return _seam_norm(str(p / q)) - if (m := _SEAM_NUM_RE.search(txt)): - return _seam_norm(m.group()) - return txt - - -def _parse_seq(seq, gold: str) -> Dict: - text = _clean_text(getattr(seq, 'decoded', '') or '') - raw = extract_boxed(text) - pred = _seam_sanitize(raw) if raw else None - correct = bool(pred) and (pred == _seam_sanitize(str(gold))) - terminated = getattr(seq, 'stop_reason', None) != 'length' - return {'pred': pred, 'correct': correct, 'terminated': terminated, - 'stop_reason': getattr(seq, 'stop_reason', None), - 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} - - -def _run_samples(sampler, prompts, max_tokens, gen_dp): - """greedy（T=0）单样本，对齐 v2 eval。gen_dp>len 时按最后一条 padding 补齐。""" - if not prompts: - return [] - params = SamplingParams(max_tokens=max_tokens, temperature=0.0, top_p=1.0, num_samples=1) - padded = prompts - if gen_dp > 1 and 0 < len(prompts) < gen_dp: - padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] - responses = sampler.sample(padded, params)[:len(prompts)] - return [list(r.sequences) if (r and r.sequences) else [] for r in responses] - - -# =========================================================================== -# problem 查表 + trials 载入 -# =========================================================================== -def load_problems(eval_records_path) -> Dict[str, Dict]: - probs: Dict[str, Dict] = {} - if not os.path.exists(eval_records_path): - return probs - for line in open(eval_records_path): - line = line.strip() - if not line: - continue - try: - r = json.loads(line) - except Exception: - continue - if r.get('record_type') != 'eval_problem': - continue - did = r.get('data_id') - if did and did not in probs: - probs[did] = {'problem': r['problem'], 'reference_answer': r['reference_answer']} - return probs - - -def main(): - ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - ap.add_argument('--trials', default=os.path.join(SCRIPT_DIR, 'trials.jsonl')) - ap.add_argument('--eval-records', default=DEFAULT_EVAL_RECORDS) - ap.add_argument('--out', default=os.path.join(SCRIPT_DIR, 'winning_skills.jsonl')) - ap.add_argument('--max-tokens', type=int, default=8192, help='对齐 v2 eval 的 --max-tokens') - ap.add_argument('--only', default=None, help='只跑某个 data_id') - ap.add_argument('--init', action='store_true', help='生成 trials.jsonl 模板后退出') - ap.add_argument('--dump-text', action='store_true', help='把每个 trial 的 executor 全文落到 probe_texts/') - args = ap.parse_args() - - problems = load_problems(args.eval_records) - - if args.init: - with open(args.trials, 'w') as f: - f.write('# 每行一个 trial；# 开头的行被忽略。skill="" 即 baseline。改 skill 后重跑本脚本。\n') - for did in DEFAULT_DATA_IDS: - p = problems.get(did, {}) - f.write(json.dumps({'data_id': did, 'tag': 'baseline', 'skill': '', - 'reference_answer': p.get('reference_answer')}, ensure_ascii=False) + '\n') - print(f'已写模板 {args.trials}（{len(DEFAULT_DATA_IDS)} 题，skill 待填）。编辑后去掉 --init 再跑。') - return - - if not os.path.exists(args.trials): - print(f'找不到 {args.trials}，先跑：python3 eval_skill_probe.py --init') - sys.exit(1) - - trials = [] - for line in open(args.trials): - line = line.strip() - if not line or line.startswith('#'): - continue - trials.append(json.loads(line)) - if args.only: - trials = [t for t in trials if t.get('data_id') == args.only] - for t in trials: - p = problems.get(t.get('data_id'), {}) - t.setdefault('problem', p.get('problem')) - t.setdefault('reference_answer', p.get('reference_answer')) - skipped = [t for t in trials if t.get('problem') is None] - for t in skipped: - print(f'[warn] data_id={t.get("data_id")} 无题面（eval_records 找不到且行内未给 problem），跳过') - trials = [t for t in trials if t.get('problem') is not None] - if not trials: - print('没有可跑的 trial。') - return - - # 仅起一个 executor sampler（enable_thinking=True，对齐 v2 eval 的 base_sampler） - twinkle.initialize(mode='ray', nproc_per_node=EXEC_GPUS, lazy_collect=False, - groups=[DeviceGroup(name='exec', ranks=list(range(EXEC_GPUS)), device_type='GPU')]) - sampler = vLLMSampler(model_id=MODEL_ID, - engine_args={'gpu_memory_utilization': GPU_MEM, - 'max_model_len': MAX_MODEL_LEN, 'tensor_parallel_size': 1}, - device_mesh=DeviceMesh.from_sizes(world_size=EXEC_GPUS, dp_size=EXEC_GPUS), - remote_group='exec') - sampler.set_template(Template, model_id=MODEL_ID, enable_thinking=True, max_length=MAX_MODEL_LEN) - - prompts = [build_skill_solve_prompt(t['problem'], t.get('skill', '')) for t in trials] - outs = _run_samples(sampler, prompts, args.max_tokens, EXEC_GPUS) - - if args.dump_text: - os.makedirs(os.path.join(SCRIPT_DIR, 'probe_texts'), exist_ok=True) - - n_ok = 0 - print('\n' + '=' * 94) - print('%-16s %-10s %-4s %-10s %-10s %-7s %s' % ('data_id', 'tag', '对?', 'pred', 'gold', 'tokens', 'note')) - print('-' * 94) - win_f = open(args.out, 'a') - for idx, (t, seqs) in enumerate(zip(trials, outs)): - roll = _parse_seq(seqs[0], t['reference_answer']) if seqs else { - 'pred': None, 'correct': False, 'terminated': False, - 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} - ok = '✓' if roll['correct'] else '✗' - note = '[截断]' if roll['stop_reason'] == 'length' else '' - if roll['correct']: - n_ok += 1 - print('%-16s %-10s %-4s %-10s %-10s %-7d %s' % ( - t.get('data_id', ''), str(t.get('tag', ''))[:10], ok, - str(roll['pred'])[:10], str(t['reference_answer'])[:10], roll['gen_tokens'], note)) - if args.dump_text: - fn = os.path.join(SCRIPT_DIR, 'probe_texts', - 'trial_%02d_%s_%s.txt' % (idx, str(t.get('data_id', '')).replace(':', '_'), - str(t.get('tag', '')))) - with open(fn, 'w') as tf: - tf.write('SKILL:\n' + (t.get('skill') or '') + '\n\n' + '=' * 60 + '\nEXECUTOR OUTPUT:\n' + roll['text']) - if roll['correct']: - win_f.write(json.dumps({'data_id': t.get('data_id'), 'tag': t.get('tag'), - 'reference_answer': t['reference_answer'], 'pred': roll['pred'], - 'gen_tokens': roll['gen_tokens'], 'max_tokens': args.max_tokens, - 'skill': t.get('skill', '')}, ensure_ascii=False) + '\n') - win_f.close() - print('-' * 94) - print(f'共 {len(trials)} 个 trial，解对 {n_ok} 个。成功的 skill 已追加到 {args.out}') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/skill2lora/good_skill_hard_fail/leak_decomp.py b/cookbook/exp/skill2lora/good_skill_hard_fail/leak_decomp.py deleted file mode 100644 index 665e24ec2..000000000 --- a/cookbook/exp/skill2lora/good_skill_hard_fail/leak_decomp.py +++ /dev/null @@ -1,146 +0,0 @@ -#!/usr/bin/env python3 -# leak_decomp.py — 严格重审 "think 提升是否由答案泄漏驱动"。 -# 上次教训: 宽松判定曾导致严重误判。本脚本做四路独立检验: -# T1 安慰剂测试: 用"别题答案"跑同一 leak 规则 → 估计 leak 标记的偶然假阳性地板 -# (think skill 长且数字密集, 答案数字偶然出现的概率天然更高) -# T2 条件分解: parseable 记录拆 leaked/clean, 分别算 acc + 占比 → 贡献分解 -# T3 反事实: 把 leaked 样本的 acc 替换成同组 clean acc → think 优势还剩多少 -# T4 严格判分: 对 withskill_text 用严格 boxed 精确匹配重新判分, -# 检验现行 grade(_seam_sanitize 取首数字等宽松步骤) 是否给 think 虚增 acc -# 难度控制: T2/T3 同时在 baseline_pass==0 (hard) 子集上重复, 排除"泄漏样本恰好是简单题"。 -import json, os, re -from collections import defaultdict - -HERE = os.path.dirname(os.path.abspath(__file__)) -FILES = {'A_old_off': 'skillcfg_full_off.jsonl', - 'B_new_on': 'skillcfg_full_on.jsonl', - 'C_new_off': 'env_runs/vllm_0.23.0/skillcfg_full_off.jsonl'} -CFGS = ['P1_narrative','P2_combo','P3_toy','P4_card','P5_pitfall','P6_seam','P7_minimal'] - -_NUM = re.compile(r'-?\d+(\.\d+)?') - -def sanitize(x): - s = str(x).strip() - m = _NUM.search(s) - if m and m.group() == s: - try: - f = float(s) - return str(int(f)) if f == int(f) else str(f) - except Exception: - pass - return s - -def leak_rule(skill, ans): - """逐字复刻 skill_config_probe.answer_leaked 的数字边界规则""" - if not skill: - return False - g = sanitize(ans) - if not g or not re.fullmatch(r'-?\d+(\.\d+)?', g): - return None # 不适用(非纯数字答案) - return bool(re.search(r'(?<![\d.])' + re.escape(g) + r'(?!\d)(?!\.\d)', skill)) - -_BOXED = re.compile(r'\\boxed\s*\{') - -def last_boxed(text): - last = None - for m in _BOXED.finditer(text or ''): - depth, i = 1, m.end() - while i < len(text) and depth > 0: - depth += (text[i] == '{') - (text[i] == '}') - i += 1 - if depth == 0: - last = text[m.end():i-1].strip() - return last - -def strict_correct(text, gold): - """严格口径: 最后一个 boxed 的内容去掉 latex 修饰后必须与 gold 精确相等(串或数值)。 - 不做 '从文本中捞第一个数字' 这类宽松回退。""" - raw = last_boxed(text) - if raw is None: - return False - s = raw.replace('\\!','').replace('\\,','').replace('\\ ',' ') - s = re.sub(r'\\text\s*\{([^}]*)\}', r'\1', s) - s = s.replace('$','').replace('{','').replace('}','').replace('\\','').strip() - g = str(gold).strip() - if s == g: - return True - try: - return float(s) == float(g) - except Exception: - return False - -def main(): - # 先取 200 题的答案表(placebo 用): data_id -> answer - answers = {} - with open(os.path.join(HERE, FILES['A_old_off'])) as f: - for line in f: - r = json.loads(line) - answers.setdefault(r['data_id'], r['reference_answer']) - dids = sorted(answers) - placebo = {} - for i, d in enumerate(dids): - # 找下一个"值不同"的答案做安慰剂 - for j in range(1, len(dids)): - cand = answers[dids[(i+j) % len(dids)]] - if sanitize(cand) != sanitize(answers[d]): - placebo[d] = cand - break - - for grp, path in FILES.items(): - # 聚合器: cfg -> 统计 - S = defaultdict(lambda: defaultdict(float)) - with open(os.path.join(HERE, path)) as f: - for line in f: - r = json.loads(line) - cfg = r['config']; s = S[cfg] - skill = r['skill'] or '' - corr = 1.0 if r['withskill_correct'] else 0.0 - hard = (r.get('baseline_pass') or 0) == 0 - s['n'] += 1 - # --- T4 严格判分 --- - sc = 1.0 if strict_correct(r.get('withskill_text',''), r['reference_answer']) else 0.0 - s['acc_loose'] += corr; s['acc_strict'] += sc - s['loose_only'] += 1.0 if (corr and not sc) else 0.0 - # --- T1 安慰剂 --- - lk = leak_rule(skill, r['reference_answer']) - if lk is not None and r['parseable']: - s['n_lk'] += 1 - s['leak'] += 1.0 if lk else 0.0 - pl = leak_rule(skill, placebo[r['data_id']]) - s['placebo'] += 1.0 if pl else 0.0 - # --- T2 条件分解 --- - if r['parseable']: - key = 'L' if r['leaked'] else 'Cn' - s[f'n_{key}'] += 1; s[f'acc_{key}'] += corr - if hard: - s[f'nh_{key}'] += 1; s[f'acch_{key}'] += corr - else: - s['n_U'] += 1; s['acc_U'] += corr - if hard: - s['nh_U'] += 1; s['acch_U'] += corr - print(f"\n{'='*100}\n### {grp} ({path})") - print("%-14s %6s | %7s %7s %9s | %5s %6s | %5s %6s | %5s %6s | %8s %8s %9s" % ( - 'config','n','leak%','placebo%','净leak%','nL','accL','nCn','accCn','nU','accU','accLoose','accStrict','looseOnly%')) - for cfg in CFGS: - s = S[cfg] - n = s['n'] or 1 - nlk = s['n_lk'] or 1 - lk, pl = s['leak']/nlk, s['placebo']/nlk - aL = s['acc_L']/s['n_L'] if s['n_L'] else float('nan') - aC = s['acc_Cn']/s['n_Cn'] if s['n_Cn'] else float('nan') - aU = s['acc_U']/s['n_U'] if s['n_U'] else float('nan') - print("%-14s %6d | %7.3f %7.3f %9.3f | %5d %6.3f | %5d %6.3f | %5d %6.3f | %8.3f %8.3f %9.3f" % ( - cfg, s['n'], lk, pl, lk-pl, s['n_L'], aL, s['n_Cn'], aC, s['n_U'], aU, - s['acc_loose']/n, s['acc_strict']/n, s['loose_only']/n)) - # hard 子集(排除"泄漏样本挑了简单题") - print(" --- hard(baseline=0) 子集: leaked vs clean 的 acc ---") - for cfg in CFGS: - s = S[cfg] - ahL = s['acch_L']/s['nh_L'] if s['nh_L'] else float('nan') - ahC = s['acch_Cn']/s['nh_Cn'] if s['nh_Cn'] else float('nan') - ahU = s['acch_U']/s['nh_U'] if s['nh_U'] else float('nan') - print(" %-14s hard: leaked %4d/%.3f clean %4d/%.3f unparse %4d/%.3f" % ( - cfg, s['nh_L'], ahL, s['nh_Cn'], ahC, s['nh_U'], ahU)) - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/skill2lora/good_skill_hard_fail/reflexion_probe.py b/cookbook/exp/skill2lora/good_skill_hard_fail/reflexion_probe.py deleted file mode 100644 index 90cdd8252..000000000 --- a/cookbook/exp/skill2lora/good_skill_hard_fail/reflexion_probe.py +++ /dev/null @@ -1,506 +0,0 @@ -#!/usr/bin/env python3 -"""reflexion_probe.py — reflexion 链路探针（纯推理，不训练）。 - -链路：executor 失败轨迹（复用 v2 eval 的 baseline 缓存原文） - → LLM rubric 诊断（qwen-plus，7 条 _MATH_RUBRIC，[PASS]/[FAIL]+fix 文本，防泄漏硬规则） - → skillmodel(Qwen3-4B, T=0.5) 条件于 (题目 + 失败轨迹 + 诊断) 生成 <skills> - → executor(Qwen3-4B, T=0, v2 eval 逐字口径) 带 skill 重试 -目的：在"有一次真实失败 + rubric 诊断"的条件下，比较 skillmodel 的 prompt 写法 × -thinking on/off 哪种救活率最高，为 buffer B 如何用 rubric 经验提供依据。 - -题目：eval 200 题中 baseline=0 的失败题抽 N 道（seed 固定），全部是"裸解必错"题， -因此 executor 重试的 acc 即救活率。诊断按 data_id 缓存（rubric_diag_cache.jsonl）， -on/off 两次运行复用，不重复调 API。 - -prompt 变体（信息量递增，用于分离各级情报的增量价值）： - R4_blind 题目（无失败、无诊断；= 上轮 P5_pitfall 原文，跨轮对照锚点） - R0_trace_only 题目 + 失败轨迹尾部（无诊断；消融 rubric 的增量） - R1_needle 题目 + 失败 + 诊断 → 纠错针（WARNING/INSTEAD + 纪律后缀） - R2_narrative 题目 + 失败 + 诊断 → 训练现版叙述式（对照 v2 regen 路径） - R3_toy_fix 题目 + 失败 + 诊断 → 针对诊断错误点的玩具题示范 - -用法（~/.env 需含 LLM_BACKUP_BASE_URL / LLM_BACKUP_API_KEY，脚本自动加载）： - EXEC_GPUS=2 SKILL_GPUS=2 python3 reflexion_probe.py --skill-thinking off - EXEC_GPUS=2 SKILL_GPUS=2 python3 reflexion_probe.py --skill-thinking on -输出：reflexion_{tag}.jsonl（含 skill-gen 全文与 executor 全文）+ stdout 汇总。 -""" -import argparse -import copy -import hashlib -import json -import os -import random -import re -import statistics as st -import sys -from concurrent.futures import ThreadPoolExecutor -from typing import Dict, List, Optional - -SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) - - -def _load_home_env(): - """加载 ~/.env（KEY=VALUE 简单格式；不覆盖已存在的环境变量）。""" - p = os.path.expanduser('~/.env') - if not os.path.exists(p): - return - for line in open(p): - line = line.strip() - if not line or line.startswith('#') or '=' not in line: - continue - k, v = line.split('=', 1) - os.environ.setdefault(k.strip(), v.strip().strip('"').strip("'")) - - -_load_home_env() - -import twinkle # noqa: E402 -from twinkle import DeviceGroup, DeviceMesh # noqa: E402 -from twinkle.data_format import SamplingParams # noqa: E402 -from twinkle.sampler import vLLMSampler # noqa: E402 -from twinkle.template import Template # noqa: E402 - -MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') -GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.8)) -EXEC_GPUS = int(os.environ.get('EXEC_GPUS', 2)) -SKILL_GPUS = int(os.environ.get('SKILL_GPUS', 2)) -MAX_MODEL_LEN = int(os.environ.get('MAX_MODEL_LEN', 16384)) -DIAG_MODEL = os.environ.get('LLM_BACKUP_MODEL', 'qwen-plus') -EVAL_RECORDS = os.path.join(SCRIPT_DIR, '..', 'skill_v2', 'eval_records.jsonl') -BASE_CACHE = os.path.join(SCRIPT_DIR, '..', 'skill_v2', 'cache', 'eval_baseline.jsonl') -DIAG_CACHE = os.path.join(SCRIPT_DIR, 'rubric_diag_cache.jsonl') - -# =========================================================================== -# executor 侧（逐字复刻 v2 eval，与 skill_config_probe.py 相同） -# =========================================================================== -_ANSWER_FORMAT_V2 = ('Present your reasoning, then put ONLY the final numeric result inside ' - '\\boxed{}. For example: \\boxed{42}.') - - -def build_direct_prompt(problem): - content = f'The problem you need to solve:\n{problem}\n\n' + _ANSWER_FORMAT_V2 - return {'messages': [{'role': 'user', 'content': content}]} - - -def build_skill_solve_prompt(problem, skill): - skill = (skill or '').strip() - if not skill: - return build_direct_prompt(problem) - content = (f'The problem you need to solve:\n{problem}\n\n' - 'Skill hint:\nFor this problem, a skill-generation model has analyzed it and ' - 'provided some advisory skills:\n' - f'{skill}\n' - 'Prefer using its techniques when they fit, but if you have a more efficient or ' - 'clearer correct method, you may use it. If you diverge from this advice, briefly ' - 'explain why. Be concise and accurate.\n' - + _ANSWER_FORMAT_V2) - return {'messages': [{'role': 'user', 'content': content}]} - - -_BOXED_RE = re.compile(r'\\boxed\s*\{') - - -def extract_boxed(text): - if not text: - return None - last = None - for m in _BOXED_RE.finditer(text): - depth, i = 1, m.end() - while i < len(text) and depth > 0: - depth += (text[i] == '{') - (text[i] == '}') - i += 1 - if depth == 0: - last = text[m.end():i - 1].strip() - return last - - -_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') - - -def _clean_text(d): - return _SPECIAL_TOKEN_RE.sub('', d or '').rstrip() - - -_SEAM_TAG_RE = re.compile(r'<\s*answer\s*>(.*?)<\s*/\s*answer\s*>', re.I | re.S) -_SEAM_BOX_RE = re.compile(r'(?:\\{1,2}\(|)\\{1,2}boxed\s*\{\s*([^}]*)\s*}(?:\)|)', re.S) -_SEAM_INLINE_RE = re.compile(r'\$([^$]+)\$|\\\(([^)]+)\\\)', re.S) -_SEAM_FRAC_RE = re.compile(r'(-?\d+(?:\.\d+)?)/(-?\d+(?:\.\d+)?)') -_SEAM_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') - - -def _seam_norm(num): - try: - f = float(num) - return str(int(f)) if f == int(f) else str(f) - except Exception: - return num.strip() - - -def _seam_sanitize(txt): - txt = (txt or '').strip() - if (m := _SEAM_TAG_RE.search(txt)): - txt = m.group(1).strip() - elif (m := _SEAM_BOX_RE.search(txt)): - txt = m.group(1).strip() - elif (m := _SEAM_INLINE_RE.search(txt)): - txt = (m.group(1) or m.group(2)).strip() - # bugfix 2026-07-29：\dfrac/\tfrac/\cfrac 归一，同 train_skill_v2._seam_sanitize - txt = txt.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac').replace(r'\cfrac', r'\frac') - txt = re.sub(r'\\frac\s*\{\s*([^}]+?)\s*}\s*\{\s*([^}]+?)\s*}', r'\1/\2', txt) - if (m := _SEAM_FRAC_RE.search(txt)): - p, q = map(float, m.groups()) - if q: - return _seam_norm(str(p / q)) - if (m := _SEAM_NUM_RE.search(txt)): - return _seam_norm(m.group()) - return txt - - -def grade(seq, gold): - text = _clean_text(getattr(seq, 'decoded', '') or '') - raw = extract_boxed(text) - pred = _seam_sanitize(raw) if raw else None - return {'pred': pred, 'correct': bool(pred) and (pred == _seam_sanitize(str(gold))), - 'stop_reason': getattr(seq, 'stop_reason', None), - 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} - - -def extract_skill(text): - low = text.lower() - end_think = low.rfind('</think>') - answer = text[end_think + len('</think>'):] if end_think >= 0 else text - s = answer.lower().rfind('<skills>') - if s < 0: - return None - inner = s + len('<skills>') - e = answer.lower().find('</skills>', inner) - if e < 0: - return None - block = answer[inner:e].strip() - return re.sub(r'</?(?:skills|skill|think)>', '', block, flags=re.I).strip() or None - - -def answer_leaked(text, reference): - """双口径泄漏检测：A=裸数字任意位置；B 口径由分析端按 |gts|>=10 复算。""" - if not text: - return False - g = _seam_sanitize(str(reference)) - if not g or not re.fullmatch(r'-?\d+(\.\d+)?', g): - return False - return bool(re.search(r'(?<![\d.])' + re.escape(g) + r'(?!\d)(?!\.\d)', text)) - - -# =========================================================================== -# rubric 诊断（对齐 v2 的 _MATH_RUBRIC；qwen-plus；防泄漏硬规则；文本格式 [PASS]/[FAIL]+fix） -# =========================================================================== -MATH_RUBRIC = [ - 'The attempt chooses a method suitable for the problem structure', - 'The attempt identifies the key constraint, invariant, or quantity before computing', - 'Algebraic and logical transformations preserve validity at each step', - 'The attempt checks required constraints, domains, boundary cases, or validity conditions', - 'The attempt avoids redundant casework, looping, or re-deriving known facts', - 'The attempt reaches a final answer within the length budget', - 'The approach stays focused on the actual question asked', -] - -# 中文注释：诊断 system prompt——硬规则禁止给出最终答案/最终数值，输出 [PASS]/[FAIL]+reason+fix -# 的紧凑文本（与 v2 _format_diagnosis 的落盘形态一致），供 skillmodel 直接消费。 -DIAG_SYSTEM = """\ -You are a rigorous math-competition grader. You will see a problem and a FAILED solution attempt (possibly truncated). Evaluate the attempt against each rubric criterion. - -Output format: one line per criterion, exactly: -- [PASS] <criterion text> -or -- [FAIL] <criterion text>: <one-sentence reason> (fix: <one-sentence concrete correction direction>) -Then a final line: Summary: <2-3 sentences naming the single most damaging error and the correct turn to take>. - -HARD RULES: never state, compute, or hint at the problem's final numeric answer or any final-stage numeric result; describe errors and directions only. Keep the whole output under 250 words.""" - -DIAG_USER = """## Problem -{problem} - -## Rubric -{rubric} - -## Failed attempt (may be truncated) -{segment} - -Now output the diagnostic lines.""" - - -def diagnose(client, problem, fail_text, gold): - rubric = '\n'.join(f'{i+1}. {t}' for i, t in enumerate(MATH_RUBRIC)) - seg = fail_text[-4000:] - msg = [{'role': 'system', 'content': DIAG_SYSTEM}, - {'role': 'user', 'content': DIAG_USER.format(problem=problem, rubric=rubric, segment=seg)}] - r = client.chat.completions.create(model=DIAG_MODEL, messages=msg, max_tokens=600, - temperature=0.2, timeout=120) - text = (r.choices[0].message.content or '').strip() - # 防泄漏兜底：诊断若带出 gts 数值，重试一次更严的指令；仍泄漏则截去含数值的行 - if answer_leaked(text, gold): - msg.append({'role': 'assistant', 'content': text}) - msg.append({'role': 'user', 'content': 'Your output contained a forbidden final numeric value. ' - 'Rewrite the SAME diagnosis with every final-stage number removed.'}) - r = client.chat.completions.create(model=DIAG_MODEL, messages=msg, max_tokens=600, - temperature=0.2, timeout=120) - text = (r.choices[0].message.content or '').strip() - if answer_leaked(text, gold): - g = _seam_sanitize(str(gold)) - text = '\n'.join(l for l in text.splitlines() - if not re.search(r'(?<![\d.])' + re.escape(g) + r'(?!\d)(?!\.\d)', l)) - return text - - -# =========================================================================== -# skillmodel prompt 变体（英文 prompt + 中文注释；统一 <skills> 输出） -# =========================================================================== -# R4：无情报锚点 = 上轮 P5_pitfall 原文（跨实验可比）。 -R4_BLIND_SYS = """\ -You are a skill-generation model. A separate executor model will solve the problem; it only sees your <skills> block. - -First think privately: solve the problem in your head AND identify the single most likely way a solver goes wrong on this type (a tempting but wrong turn, an off-by-one, a wasteful brute-force, a wrong branch). Then, inside <skills></skills>, write under 90 words: -- WARNING: name that most likely mistake concretely and say why it is wrong. -- INSTEAD: one or two sentences pointing to the correct turn (technique name + where to apply it), without solving the problem or revealing any numeric result. -- End with: "Avoid re-checking loops; box a bare number as soon as it is computed." -""" - -# R0：只有失败轨迹（无诊断）——消融 rubric 的增量价值。 -R0_TRACE_SYS = """\ -You are a skill-generation model. A separate executor model previously FAILED this problem; you will see the tail of its failed attempt. The executor will retry seeing ONLY your <skills> block. - -First think privately: read the failed attempt, find where it went wrong, and decide the correct turn. Then, inside <skills></skills>, write under 90 words: -- WARNING: the concrete mistake the previous attempt made (quote its wrong move briefly). -- INSTEAD: one or two sentences pointing to the correct turn, without solving the problem or revealing any numeric result. -- End with: "Avoid re-checking loops; box a bare number as soon as it is computed." -""" - -# R1：query + 诊断（无轨迹）→ 纠错针（把 rubric 的 FAIL/fix 转译成对 executor 的直接行为指令）。 -R1_NEEDLE_SYS = """\ -You are a skill-generation model. A separate executor model previously FAILED this problem. You will see an expert rubric diagnosis of that failure (you will NOT see the failed attempt itself). The executor will retry seeing ONLY your <skills> block. - -First think privately: from the diagnosis, pinpoint the decisive error. Then, inside <skills></skills>, write under 90 words: -- WARNING: the decisive mistake (grounded in the diagnosis, stated concretely for THIS problem). -- INSTEAD: the corrective route distilled from the diagnosis's fix directions - technique name + where to apply it. Do not solve the problem; never state any numeric result. -- End with: "Avoid re-checking loops; box a bare number as soon as it is computed." -""" - -# R2：query + 诊断（无轨迹）→ 训练现版叙述式（对照 v2 _regen_prompt 的文体路径）。 -R2_NARR_SYS = """\ -You are a skill-generation model. A separate executor model previously FAILED this problem. You will see an expert rubric diagnosis of that failure (you will NOT see the failed attempt itself). The executor will retry seeing ONLY your <skills> block. - -First, think privately: work the problem out and understand why the attempt failed. Then write the <skills> block as ONE coherent analysis narrative (not a bullet list): name what the problem is essentially asking, walk through the recommended approach, and weave in - informed by the diagnosis - the specific pitfall that sank the previous attempt and how to avoid it. Do NOT solve the problem, do NOT reveal or compute the final answer, and do NOT substitute the problem's specific numbers into the steps. Keep it to roughly one focused paragraph. - -Put ONLY the methodology inside <skills></skills>. -""" - -# R3：query + 诊断（无轨迹）→ 玩具题示范（针对诊断指出的错误技巧点造 toy，异数字防泄漏）。 -R3_TOYFIX_SYS = """\ -You are a skill-generation model. A separate executor model previously FAILED this problem. You will see an expert rubric diagnosis of that failure (you will NOT see the failed attempt itself). The executor will retry seeing ONLY your <skills> block. - -First think privately: from the diagnosis, identify the ONE technique the executor got wrong. Then, inside <skills></skills>, do exactly this (under 110 words): -1. Invent a MINIATURE problem exercising that same technique with DIFFERENT, much smaller numbers, and solve the miniature completely in at most 5 short lines, making the correct move (the one the failed attempt missed) explicit. -2. One transfer sentence: "Your problem needs the same move where the previous attempt went wrong - apply it, then box a bare number." -Hard rules: never use any number from the original problem; never state its answer. -""" - -PROMPTS = { - 'R4_blind': ('none', R4_BLIND_SYS), - 'D1_needle': ('diag', R1_NEEDLE_SYS), - 'D2_narr': ('diag', R2_NARR_SYS), - 'D3_toyfix': ('diag', R3_TOYFIX_SYS), -} - - -def skillgen_prompt(name, problem, fail_tail, diag): - mode, sys_p = PROMPTS[name] - user = f'Problem:\n{problem}' - if mode in ('trace', 'trace+diag'): - user += f'\n\nFailed attempt (tail):\n{fail_tail}' - if mode in ('diag', 'trace+diag'): - user += f'\n\nExpert rubric diagnosis of the failure:\n{diag}' - return {'messages': [{'role': 'system', 'content': sys_p}, {'role': 'user', 'content': user}]} - - -# =========================================================================== -# 数据与主流程 -# =========================================================================== -def md5_key(problem): - return hashlib.md5('\x1f'.join([problem]).encode('utf-8')).hexdigest() - - -def load_fail_problems(n, seed): - probs = {} - for line in open(EVAL_RECORDS): - line = line.strip() - if not line: - continue - r = json.loads(line) - if r.get('record_type') != 'eval_problem' or r.get('chunk') != -1: - continue - if r['data_id'] not in probs and float(r.get('baseline_pass', 1)) == 0: - probs[r['data_id']] = {'data_id': r['data_id'], 'problem': r['problem'], - 'reference_answer': r['reference_answer']} - cache = {} - for line in open(BASE_CACHE): - try: - c = json.loads(line) - cache[c['key']] = c['value'] - except Exception: - continue - items = [] - for p in sorted(probs.values(), key=lambda x: x['data_id']): - v = cache.get(md5_key(p['problem'])) - if v and not v.get('correct') and (v.get('text') or '').strip(): - p['fail_text'] = v['text'] - items.append(p) - rng = random.Random(seed) - sample = rng.sample(items, min(n, len(items))) - sample.sort(key=lambda x: x['data_id']) - print(f'[抽样] baseline 失败且有轨迹全文的题 {len(items)} -> 抽 {len(sample)}（seed={seed}）') - return sample - - -def run_batch(sampler, prompts, max_tokens, temperature, top_p, dp, num_samples=1): - if not prompts: - return [] - params = SamplingParams(max_tokens=max_tokens, temperature=temperature, top_p=top_p, num_samples=num_samples) - padded = prompts - if dp > 1 and 0 < len(prompts) < dp: - padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(dp - len(prompts))] - responses = sampler.sample(padded, params)[:len(prompts)] - return [list(r.sequences) if (r and r.sequences) else [] for r in responses] - - -def main(): - ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - ap.add_argument('--skill-thinking', choices=('on', 'off'), required=True) - ap.add_argument('--prompts', default=','.join(PROMPTS.keys())) - ap.add_argument('--n-problems', type=int, default=130) - ap.add_argument('--n-rollouts', type=int, default=8, help='每题每思路的 skill 采样数（T>0）；executor 对每条 skill 各跑一次 greedy') - ap.add_argument('--seed', type=int, default=42) - ap.add_argument('--skill-temperature', type=float, default=0.5) - ap.add_argument('--skill-max-tokens', type=int, default=8192) - ap.add_argument('--max-tokens', type=int, default=8192) - ap.add_argument('--fail-tail-chars', type=int, default=1800) - ap.add_argument('--out', default=None) - args = ap.parse_args() - - names = [x for x in args.prompts.split(',') if x] - problems = load_fail_problems(args.n_problems, args.seed) - out_path = args.out or os.path.join(SCRIPT_DIR, f'reflexion_{args.skill_thinking}.jsonl') - - # ---- 阶段1：rubric 诊断（带磁盘缓存，8 线程并发）---- - diag_cache = {} - if os.path.exists(DIAG_CACHE): - for line in open(DIAG_CACHE): - try: - c = json.loads(line) - diag_cache[c['data_id']] = c['diag'] - except Exception: - continue - todo = [p for p in problems if p['data_id'] not in diag_cache] - if todo: - from openai import OpenAI - client = OpenAI(api_key=os.environ['LLM_BACKUP_API_KEY'], - base_url=os.environ['LLM_BACKUP_BASE_URL']) - print(f'[诊断] 需调 API {len(todo)} 题（model={DIAG_MODEL}），其余 {len(problems)-len(todo)} 题走缓存') - - def _one(p): - try: - return p['data_id'], diagnose(client, p['problem'], p['fail_text'], p['reference_answer']) - except Exception as e: - return p['data_id'], f'[DIAG_ERROR] {e}' - with ThreadPoolExecutor(max_workers=8) as ex: - with open(DIAG_CACHE, 'a') as f: - for did, diag in ex.map(_one, todo): - diag_cache[did] = diag - f.write(json.dumps({'data_id': did, 'diag': diag}, ensure_ascii=False) + '\n') - n_err = sum(1 for p in problems if str(diag_cache.get(p['data_id'], '')).startswith('[DIAG_ERROR]')) - print(f'[诊断] 完成，失败 {n_err} 题') - - # ---- 阶段2：skill-gen + executor ---- - think = args.skill_thinking == 'on' - twinkle.initialize(mode='ray', nproc_per_node=SKILL_GPUS + EXEC_GPUS, lazy_collect=False, groups=[ - DeviceGroup(name='skill', ranks=list(range(SKILL_GPUS)), device_type='GPU'), - DeviceGroup(name='exec', ranks=list(range(SKILL_GPUS, SKILL_GPUS + EXEC_GPUS)), device_type='GPU')]) - - def make_sampler(group, world, enable_thinking): - s = vLLMSampler(model_id=MODEL_ID, - engine_args={'gpu_memory_utilization': GPU_MEM, - 'max_model_len': MAX_MODEL_LEN, 'tensor_parallel_size': 1}, - device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), - remote_group=group) - s.set_template(Template, model_id=MODEL_ID, enable_thinking=enable_thinking, max_length=MAX_MODEL_LEN) - return s - - skill_sampler = make_sampler('skill', SKILL_GPUS, enable_thinking=think) - exec_sampler = make_sampler('exec', EXEC_GPUS, enable_thinking=True) - - sg_prompts, meta = [], [] - for name in names: - for p in problems: - tail = p['fail_text'][-args.fail_tail_chars:] - sg_prompts.append(skillgen_prompt(name, p['problem'], tail, diag_cache.get(p['data_id'], ''))) - meta.append((name, p)) - print(f'[skill-gen] {len(sg_prompts)} 条 x {args.n_rollouts} rollouts, thinking={args.skill_thinking}, T={args.skill_temperature}') - sg_out = run_batch(skill_sampler, sg_prompts, args.skill_max_tokens, - args.skill_temperature, 0.95, SKILL_GPUS, num_samples=args.n_rollouts) - - trials = [] - for (name, p), seqs in zip(meta, sg_out): - seqs = list(seqs or []) - for si in range(args.n_rollouts): - seq = seqs[si] if si < len(seqs) else None - full = _clean_text(getattr(seq, 'decoded', '') or '') if seq is not None else '' - sk = extract_skill(full) or '' - trials.append({'config': name, 'data_id': p['data_id'], 'sample_idx': si, - 'problem': p['problem'], - 'reference_answer': p['reference_answer'], - 'diag': diag_cache.get(p['data_id'], ''), - 'diag_leaked': answer_leaked(diag_cache.get(p['data_id'], ''), p['reference_answer']), - 'skill': sk, 'parseable': bool(sk), 'skill_chars': len(sk), - 'leaked': answer_leaked(sk, p['reference_answer']), - 'skillgen_full': full, - 'skillgen_stop': getattr(seq, 'stop_reason', None) if seq is not None else 'empty', - 'skillgen_tokens': len(getattr(seq, 'tokens', None) or []) if seq is not None else 0}) - - ex_prompts = [build_skill_solve_prompt(t['problem'], t['skill']) for t in trials] - print(f'[executor] {len(ex_prompts)} 条, T=0') - ex_out = run_batch(exec_sampler, ex_prompts, args.max_tokens, 0.0, 1.0, EXEC_GPUS) - - with open(out_path, 'w') as f: - for t, seqs in zip(trials, ex_out): - roll = grade(seqs[0], t['reference_answer']) if seqs else { - 'pred': None, 'correct': False, 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} - t.update({'withskill_pred': roll['pred'], 'withskill_correct': roll['correct'], - 'withskill_stop': roll['stop_reason'], 'withskill_tokens': roll['gen_tokens'], - 'withskill_text': roll['text'], 'skill_thinking': args.skill_thinking}) - f.write(json.dumps(t, ensure_ascii=False) + '\n') - - # ---- 汇总（全是 baseline=0 的题，acc 即救活率；另报题级 pass@k）---- - print('\n' + '=' * 112) - print('%-11s %-6s %-6s %-6s %-6s %-8s %-8s %-8s %-8s' % ( - 'config', 'n', 'parse', 'leakA', 'trunc', 'skill字符', '救活@1', '救活@k', '救活@1(gts>=10无泄漏)')) - print('-' * 112) - for name in names: - sub = [t for t in trials if t['config'] == name] - n = len(sub) - parse = sum(t['parseable'] for t in sub) / n - leak = sum(t['leaked'] for t in sub) / n - trunc = sum(1 for t in sub if t['withskill_stop'] == 'length') / n - chars = int(st.median([t['skill_chars'] for t in sub if t['parseable']] or [0])) - acc = sum(t['withskill_correct'] for t in sub) / n - byq = {} - for t in sub: - byq.setdefault(t['data_id'], []).append(t) - p_at_k = sum(1 for v in byq.values() if any(x['withskill_correct'] for x in v)) / len(byq) - clean = [t for t in sub if not t['leaked']] - cacc = sum(t['withskill_correct'] for t in clean) / max(1, len(clean)) - print('%-11s %-6d %-6.2f %-6.2f %-6.2f %-8d %-8.3f %-8.3f %-.3f(n=%d)' % ( - name, n, parse, leak, trunc, chars, acc, p_at_k, cacc, len(clean))) - print('-' * 112) - print(f'明细（含 skill-gen/executor 全文）已写入 {out_path}') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/skill2lora/good_skill_hard_fail/sample_probe.py b/cookbook/exp/skill2lora/good_skill_hard_fail/sample_probe.py deleted file mode 100644 index c0892aa90..000000000 --- a/cookbook/exp/skill2lora/good_skill_hard_fail/sample_probe.py +++ /dev/null @@ -1,136 +0,0 @@ -#!/usr/bin/env python3 -# sample_probe.py — 抽样交叉验证，为统计结论提供可人工核对的具体证据（含 文件:行号:data_id）。 -# 5 项交叉验证: -# V1 A vs C 样本级逐字一致率 (环境影响的最硬证据; 同 key 对比 skill 与 withskill_correct) -# V2 leak 标记真伪 (随机抽 leaked=True/False 各若干, 核对 reference_answer 是否真出现在 skill) -# V3 think 泄漏机制 (B 组 leaked=True 样本, 展示 think 段算出答案->写进 skill) -# V4 parse 失败成因 (B 组 P7/parse=False 样本, 确认 skillgen 被 think 吃满预算而截断) -# V5 accClean 样本量 (打印各 config n_clean, 防止小样本误读高 accClean) -import json, os, re, random -from collections import defaultdict - -HERE = os.path.dirname(os.path.abspath(__file__)) -FILES = { - 'A': 'skillcfg_full_off.jsonl', - 'B': 'skillcfg_full_on.jsonl', - 'C': 'env_runs/vllm_0.23.0/skillcfg_full_off.jsonl', -} -random.seed(0) - - -def load_indexed(path): - """返回 {(config,data_id,sample_idx): (lineno, record)}""" - d = {} - with open(os.path.join(HERE, path)) as f: - for i, line in enumerate(f, 1): - line = line.strip() - if not line: - continue - r = json.loads(line) - d[(r['config'], r['data_id'], r['sample_idx'])] = (i, r) - return d - - -def leaked_check(skill, ref): - """独立复现 answer_leaked 的判定意图: 整数答案是否以数字边界出现在 skill""" - s = str(ref).strip() - if not re.fullmatch(r'-?\d+(\.\d+)?', s): - return None # 非纯数字答案, 泄漏判定本就不适用 - return re.search(r'(?<!\d)' + re.escape(s) + r'(?!\d)', skill or '') is not None - - -def main(): - A = load_indexed(FILES['A']) - B = load_indexed(FILES['B']) - C = load_indexed(FILES['C']) - fa, fb, fc = FILES['A'], FILES['B'], FILES['C'] - - print("="*90) - print("[V1] A vs C 样本级逐字一致率（同 config/data_id/sample_idx；环境=torch/cuda/transformers 变更的净效应）") - keys = [k for k in A if k in C] - same_skill = sum(1 for k in keys if (A[k][1]['skill'] or '') == (C[k][1]['skill'] or '')) - same_exec = sum(1 for k in keys if (A[k][1]['withskill_text'] or '') == (C[k][1]['withskill_text'] or '')) - same_corr = sum(1 for k in keys if bool(A[k][1]['withskill_correct']) == bool(C[k][1]['withskill_correct'])) - print(f" N={len(keys)} skill 逐字一致={same_skill/len(keys):.3f} " - f"executor 全文逐字一致={same_exec/len(keys):.3f} correct 一致={same_corr/len(keys):.3f}") - # 展示一个 skill 不同但 correct 相同 / 一个逐字相同 的实例 - diff_ex = next((k for k in keys if (A[k][1]['skill'] or '')!=(C[k][1]['skill'] or '')), None) - if diff_ex: - la,_ = A[diff_ex]; lc,_ = C[diff_ex] - print(f" 例(skill 不同): key={diff_ex} A={fa}:{la} C={fc}:{lc}") - print(f" A.skill[:90]={ (A[diff_ex][1]['skill'] or '')[:90]!r}") - print(f" C.skill[:90]={ (C[diff_ex][1]['skill'] or '')[:90]!r}") - - print("\n"+"="*90) - print("[V2] leak 标记真伪核对（随机抽样, 独立重算数字边界匹配）") - for grp,(D,fn) in {'A':(A,fa),'B':(B,fb),'C':(C,fc)}.items(): - pos = [k for k in D if D[k][1]['leaked']] - neg = [k for k in D if not D[k][1]['leaked']] - random.shuffle(pos); random.shuffle(neg) - tp = 0; checked_pos = pos[:60] - for k in checked_pos: - _,r = D[k] - v = leaked_check(r['skill'], r['reference_answer']) - if v: tp += 1 - fp_free = 0; checked_neg = neg[:60] - for k in checked_neg: - _,r = D[k] - v = leaked_check(r['skill'], r['reference_answer']) - if v is False or v is None: - fp_free += 1 - print(f" [{grp}] leaked=True 抽{len(checked_pos)} 复算确含答案={tp}/{len(checked_pos)} " - f"leaked=False 抽{len(checked_neg)} 复算确不含={fp_free}/{len(checked_neg)}") - - print("\n"+"="*90) - print("[V3] think 泄漏机制实例（B 组 leaked=True, 展示 think 段->skill 搬答案）") - shown = 0 - for k in B: - _,r = B[k] - if r['config']=='P1_narrative' and r['leaked'] and str(r['reference_answer']).lstrip('-').isdigit(): - ln,_ = B[k] - full = r['skillgen_full'] or '' - ans = str(r['reference_answer']) - think_end = full.lower().find('</think>') - in_think = ans in full[:think_end] if think_end>0 else False - in_skill = ans in (r['skill'] or '') - print(f" {fb}:{ln} key={k} ref={ans} 答案在think段={in_think} 在skill={in_skill}") - idx = (r['skill'] or '').find(ans) - if idx>=0: - print(f" skill 命中片段: ...{(r['skill'])[max(0,idx-45):idx+len(ans)+25]!r}...") - shown += 1 - if shown>=3: break - - print("\n"+"="*90) - print("[V4] parse 失败成因（B 组 P7_minimal, parseable=False）") - cnt_len=0; cnt_noskill=0; shown=0 - for k in B: - _,r = B[k] - if r['config']!='P7_minimal' or r['parseable']: - continue - full = r['skillgen_full'] or '' - has_close = '</skills>' in full.lower() - if r['skillgen_stop']=='length': cnt_len+=1 - if not has_close: cnt_noskill+=1 - if shown<3: - ln,_ = B[k] - print(f" {fb}:{ln} key={k} stop={r['skillgen_stop']} sg_tokens={r['skillgen_tokens']} " - f"含</skills>={has_close} full尾50={full[-50:]!r}") - shown+=1 - total_pf = sum(1 for k in B if B[k][1]['config']=='P7_minimal' and not B[k][1]['parseable']) - print(f" P7 parse失败共 {total_pf}: 其中 stop=length {cnt_len} 无</skills> {cnt_noskill}") - - print("\n"+"="*90) - print("[V5] accClean 样本量核对（leaked=0 且 parseable=1 的 n_clean, 防小样本误读）") - for grp,(D,fn) in {'A':(A,fa),'B':(B,fb),'C':(C,fc)}.items(): - by=defaultdict(lambda:[0,0]) - for k in D: - _,r=D[k] - if (not r['leaked']) and r['parseable']: - by[r['config']][0]+=1 - by[r['config']][1]+= 1 if r['withskill_correct'] else 0 - cells=" ".join(f"{c.split('_')[0]}={by[c][0]}({(by[c][1]/by[c][0] if by[c][0] else 0):.2f})" - for c in ['P1_narrative','P5_pitfall','P7_minimal']) - print(f" [{grp}] n_clean(acc): {cells}") - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/skill2lora/good_skill_hard_fail/skill_config_probe.py b/cookbook/exp/skill2lora/good_skill_hard_fail/skill_config_probe.py deleted file mode 100644 index 069a74d7c..000000000 --- a/cookbook/exp/skill2lora/good_skill_hard_fail/skill_config_probe.py +++ /dev/null @@ -1,421 +0,0 @@ -#!/usr/bin/env python3 -"""skill_config_probe.py — Qwen3-4B skill 模型 × executor 的纯推理配置探针。 - -目的：不训练，只推理。用 Qwen3-4B 当 skill 生成模型（T=0.5），产出 <skills> 喂给 -executor（T=0，与 train_skill_v2.py 的 v2 eval 逐字同口径），比较不同配置的效果： - - skill 模型 enable_thinking：on / off（由 --skill-thinking 指定，跑两次对比） - - skill 模型 system prompt：7 种变体（含训练现版对照、六思路混合模板、toy 类比等） - -题目：从 eval_records.jsonl(chunk=-1) 的 200 题按 baseline_pass 分层抽 50 题， -通过/失败比例与整集一致（难度配比同实际 eval），seed 固定保证跨配置可比。 -baseline 直接复用缓存的 baseline_pass（同模型同 greedy 口径，无需重跑）。 - -用法： - EXEC_GPUS=2 SKILL_GPUS=2 python3 skill_config_probe.py --skill-thinking off - EXEC_GPUS=2 SKILL_GPUS=2 python3 skill_config_probe.py --skill-thinking on - python3 skill_config_probe.py --skill-thinking off --prompts P2_combo,P3_toy # 只跑子集 -输出：skillcfg_{tag}.jsonl（逐 trial）+ stdout 汇总表。 -""" -import argparse -import copy -import json -import os -import random -import re -import statistics as st -import sys -from typing import Dict, List, Optional - -import twinkle -from twinkle import DeviceGroup, DeviceMesh -from twinkle.data_format import SamplingParams -from twinkle.sampler import vLLMSampler -from twinkle.template import Template - -SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) -MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') -GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.8)) -EXEC_GPUS = int(os.environ.get('EXEC_GPUS', 2)) -SKILL_GPUS = int(os.environ.get('SKILL_GPUS', 2)) -MAX_MODEL_LEN = int(os.environ.get('MAX_MODEL_LEN', 16384)) -DEFAULT_EVAL_RECORDS = os.path.join(SCRIPT_DIR, '..', 'skill_v2', 'eval_records.jsonl') - -# =========================================================================== -# executor 侧 —— 逐字复刻 train_skill_v2.py v2 分支（与 eval 口径一致） -# =========================================================================== -_ANSWER_FORMAT_V2 = ('Present your reasoning, then put ONLY the final numeric result inside ' - '\\boxed{}. For example: \\boxed{42}.') - - -def build_direct_prompt(problem): - content = f'The problem you need to solve:\n{problem}\n\n' + _ANSWER_FORMAT_V2 - return {'messages': [{'role': 'user', 'content': content}]} - - -def build_skill_solve_prompt(problem, skill): - skill = (skill or '').strip() - if not skill: - return build_direct_prompt(problem) - content = (f'The problem you need to solve:\n{problem}\n\n' - 'Skill hint:\nFor this problem, a skill-generation model has analyzed it and ' - 'provided some advisory skills:\n' - f'{skill}\n' - 'Prefer using its techniques when they fit, but if you have a more efficient or ' - 'clearer correct method, you may use it. If you diverge from this advice, briefly ' - 'explain why. Be concise and accurate.\n' - + _ANSWER_FORMAT_V2) - return {'messages': [{'role': 'user', 'content': content}]} - - -_BOXED_RE = re.compile(r'\\boxed\s*\{') - - -def extract_boxed(text: str) -> Optional[str]: - if not text: - return None - last = None - for m in _BOXED_RE.finditer(text): - depth, i = 1, m.end() - while i < len(text) and depth > 0: - depth += (text[i] == '{') - (text[i] == '}') - i += 1 - if depth == 0: - last = text[m.end():i - 1].strip() - return last - - -_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') - - -def _clean_text(decoded: Optional[str]) -> str: - return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() - - -_SEAM_TAG_RE = re.compile(r'<\s*answer\s*>(.*?)<\s*/\s*answer\s*>', re.I | re.S) -_SEAM_BOX_RE = re.compile(r'(?:\\{1,2}\(|)\\{1,2}boxed\s*\{\s*([^}]*)\s*}(?:\)|)', re.S) -_SEAM_INLINE_RE = re.compile(r'\$([^$]+)\$|\\\(([^)]+)\\\)', re.S) -_SEAM_FRAC_RE = re.compile(r'(-?\d+(?:\.\d+)?)/(-?\d+(?:\.\d+)?)') -_SEAM_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') - - -def _seam_norm(num: str) -> str: - try: - f = float(num) - return str(int(f)) if f == int(f) else str(f) - except Exception: - return num.strip() - - -def _seam_sanitize(txt: str) -> str: - txt = (txt or '').strip() - if (m := _SEAM_TAG_RE.search(txt)): - txt = m.group(1).strip() - elif (m := _SEAM_BOX_RE.search(txt)): - txt = m.group(1).strip() - elif (m := _SEAM_INLINE_RE.search(txt)): - txt = (m.group(1) or m.group(2)).strip() - # bugfix 2026-07-29：\dfrac/\tfrac/\cfrac 归一，同 train_skill_v2._seam_sanitize - txt = txt.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac').replace(r'\cfrac', r'\frac') - txt = re.sub(r'\\frac\s*\{\s*([^}]+?)\s*}\s*\{\s*([^}]+?)\s*}', r'\1/\2', txt) - if (m := _SEAM_FRAC_RE.search(txt)): - p, q = map(float, m.groups()) - if q: - return _seam_norm(str(p / q)) - if (m := _SEAM_NUM_RE.search(txt)): - return _seam_norm(m.group()) - return txt - - -def grade(seq, gold) -> Dict: - text = _clean_text(getattr(seq, 'decoded', '') or '') - raw = extract_boxed(text) - pred = _seam_sanitize(raw) if raw else None - correct = bool(pred) and (pred == _seam_sanitize(str(gold))) - return {'pred': pred, 'correct': correct, - 'stop_reason': getattr(seq, 'stop_reason', None), - 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} - - -# ---- skill 抽取（v2 泛化版：砍 think 后取最后一个 <skills> 块）---- -def extract_skill(text: str) -> Optional[str]: - low = text.lower() - end_think = low.rfind('</think>') - answer = text[end_think + len('</think>'):] if end_think >= 0 else text - s = answer.lower().rfind('<skills>') - if s < 0: - return None - inner = s + len('<skills>') - e = answer.lower().find('</skills>', inner) - if e < 0: - return None - block = answer[inner:e].strip() - block = re.sub(r'</?(?:skills|skill|think)>', '', block, flags=re.IGNORECASE).strip() - return block or None - - -def answer_leaked(skill: str, reference) -> bool: - """诊断用：skill 文本中是否出现 gts 数值（digit-boundary，简版）。""" - if not skill: - return False - g = _seam_sanitize(str(reference)) - if not g or not re.fullmatch(r'-?\d+(\.\d+)?', g): - return False - return bool(re.search(r'(?<![\d.])' + re.escape(g) + r'(?!\d)(?!\.\d)', skill)) - - -# =========================================================================== -# skill 模型 system prompt 变体（英文 prompt + 中文注释；输出统一 <skills> 块便于解析） -# =========================================================================== -# P1：训练现版（方案1）SKILL_GEN_SYSTEM 逐字对照组。 -P1_NARRATIVE = """\ -You are a skill-generation model. Your <skills> block will be fed to a SEPARATE downstream executor model that must solve the problem on its own. The executor will NOT see your private reasoning — it only sees what is inside <skills>...</skills>. - -First, think privately: actually work the problem out in your head to make sure you understand it, then step back and abstract WHAT MAKES THIS TYPE OF PROBLEM SOLVABLE into transferable methodology. - -Then write the <skills> block following these rules: -- Give general, transferable solving techniques for this TYPE of problem: the key concepts/theorems it relies on, the recommended strategy and steps, and the common pitfalls to avoid — plus a brief reason for each piece of advice so the executor understands why. -- Write it as one coherent analysis narrative (not a bullet list): first name what the problem is essentially asking, then walk through how to approach it, blending concepts, steps, pitfalls and reasons into a single connected story. -- CRITICAL: Do NOT solve the problem for the executor. Do NOT reveal or compute the final answer, and do NOT substitute the problem's specific given numbers into the steps or state any intermediate numeric results. Leave ALL concrete numbers for the executor to compute on its own. If you catch yourself writing a specific number from the problem, replace it with a description of the quantity instead. -- Keep it concise: aim for roughly one focused paragraph. - -Put ONLY the methodology inside <skills></skills>. -""" - -# P2：六思路混合模板——主体二选一（toy 类比 / 路线卡片）+ 永远加纪律后缀；限长。 -P2_COMBO = """\ -You are a skill-generation model. Your <skills> block is the ONLY thing a separate executor model will see; it must help the executor solve the problem quickly within a tight token budget. - -First think privately and solve the problem in your head. Then write a SHORT <skills> block (under 120 words) with exactly this structure: -1. MAIN PART - pick ONE of the two forms, whichever fits the problem better: - (a) Toy example: invent a tiny problem of the SAME type but with DIFFERENT, smaller numbers, solve the toy completely in 2-4 lines showing the key trick, then add one sentence: "Your problem has the same shape - apply the same steps to its numbers." - (b) Route card: name the problem type, then give the key formula / recurrence / lemma / reduction that cracks it (no derivation, no solving), and say what single quantity to compute. -2. LAST LINE - always end with exactly this discipline line: "Single pass: no re-deriving, no re-checking; once computed, box a bare number immediately." - -Never use the original problem's own numbers in the main part; never state or imply the final answer. - -Put everything inside <skills></skills>. -""" - -# P3：纯 toy 类比——完整解一道异数字同型玩具题，靠示范迁移；天然 answer-free。 -P3_TOY = """\ -You are a skill-generation model. A separate executor model will solve the problem; it only sees your <skills> block. - -First think privately and identify the core technique this problem needs. Then, inside <skills></skills>, do exactly one thing: invent a MINIATURE problem of the same type with DIFFERENT and much smaller numbers, and solve that miniature completely in at most 5 short lines, making the key trick explicit. Finish with one transfer sentence: "Your problem has the same shape - repeat these steps with its own numbers, then box a bare number." - -Hard rules: never mention or use any number that appears in the original problem; never state the original problem's answer; keep the whole block under 100 words. -""" - -# P4：路线卡片——极简结构化卡片（类型/公式/起点/目标），无叙述。 -P4_CARD = """\ -You are a skill-generation model. A separate executor model will solve the problem; it only sees your <skills> block. - -First think privately and find the standard route for this problem. Then output, inside <skills></skills>, an ultra-compact ROUTE CARD with at most 4 lines: -TYPE: <the problem type in a few words> -KEY: <the one formula / recurrence / lemma / substitution that cracks it> -START: <what to set up first> -COMPUTE: <the single final quantity the executor must evaluate and box as a bare number> - -No derivations, no explanations, no solving, never state the final answer, never copy the problem's numbers into KEY. -""" - -# P5：预判纠错——无失败情报版“纠错针”：预判本题最可能的错误走向并拦截。 -P5_PITFALL = """\ -You are a skill-generation model. A separate executor model will solve the problem; it only sees your <skills> block. - -First think privately: solve the problem in your head AND identify the single most likely way a solver goes wrong on this type (a tempting but wrong turn, an off-by-one, a wasteful brute-force, a wrong branch). Then, inside <skills></skills>, write under 90 words: -- WARNING: name that most likely mistake concretely and say why it is wrong. -- INSTEAD: one or two sentences pointing to the correct turn (technique name + where to apply it), without solving the problem or revealing any numeric result. -- End with: "Avoid re-checking loops; box a bare number as soon as it is computed." -""" - -# P6：SEAM 经验风格（英文，输出改为 <skills> 统一解析）——概念/策略/易错三段式对照组。 -P6_SEAM = """\ -You are a problem-solving guidance model. Read the math problem and distill a concise, reusable piece of solving experience that will help a SEPARATE solver model reach the correct answer. -Rules: -- Do NOT solve the problem and do NOT reveal or compute the final answer. -- State the key concepts/theorems, the recommended strategy/steps, and the common pitfalls to avoid. -- Output ONLY the experience, wrapped EXACTLY as <skills> ... </skills>. -""" - -# P7：一句话下界对照——只准一句话点出关键恒等式/技巧。 -P7_MINIMAL = """\ -You are a skill-generation model. A separate executor model will solve the problem; it only sees your <skills> block. -Inside <skills></skills>, write EXACTLY ONE sentence (max 30 words) naming the single key identity, theorem, or technique that cracks this problem. Nothing else. Never state the answer. -""" - -PROMPTS = { - 'P1_narrative': P1_NARRATIVE, - 'P2_combo': P2_COMBO, - 'P3_toy': P3_TOY, - 'P4_card': P4_CARD, - 'P5_pitfall': P5_PITFALL, - 'P6_seam': P6_SEAM, - 'P7_minimal': P7_MINIMAL, -} - - -def skillgen_prompt(system: str, problem: str) -> Dict: - # 与训练 _skillgen_prompt 同构：system + user('Problem:\n...') - return {'messages': [{'role': 'system', 'content': system}, - {'role': 'user', 'content': f'Problem:\n{problem}'}]} - - -# =========================================================================== -# 题目分层抽样：与 200 题集 baseline 通过率同配比 -# =========================================================================== -def load_problems(path, n, seed): - probs = {} - for line in open(path): - line = line.strip() - if not line: - continue - try: - r = json.loads(line) - except Exception: - continue - if r.get('record_type') != 'eval_problem' or r.get('chunk') != -1: - continue - did = r.get('data_id') - if did and did not in probs: - probs[did] = {'data_id': did, 'problem': r['problem'], - 'reference_answer': r['reference_answer'], - 'baseline_pass': float(r.get('baseline_pass', 0))} - items = sorted(probs.values(), key=lambda x: x['data_id']) - passed = [x for x in items if x['baseline_pass'] > 0] - failed = [x for x in items if x['baseline_pass'] == 0] - ratio = len(passed) / max(1, len(items)) - n_pass = round(n * ratio) - rng = random.Random(seed) - sample = rng.sample(passed, min(n_pass, len(passed))) + \ - rng.sample(failed, min(n - n_pass, len(failed))) - sample.sort(key=lambda x: x['data_id']) - print(f'[抽样] 全集 {len(items)} 题 baseline率 {ratio:.3f} -> 抽 {len(sample)} 题 ' - f'(pass {sum(1 for x in sample if x["baseline_pass"] > 0)} / fail ' - f'{sum(1 for x in sample if x["baseline_pass"] == 0)})，seed={seed}') - return sample - - -def run_batch(sampler, prompts, max_tokens, temperature, top_p, dp, num_samples=1): - if not prompts: - return [] - params = SamplingParams(max_tokens=max_tokens, temperature=temperature, - top_p=top_p, num_samples=num_samples) - padded = prompts - if dp > 1 and 0 < len(prompts) < dp: - padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(dp - len(prompts))] - responses = sampler.sample(padded, params)[:len(prompts)] - return [list(r.sequences) if (r and r.sequences) else [] for r in responses] - - -def main(): - ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - ap.add_argument('--skill-thinking', choices=('on', 'off'), required=True) - ap.add_argument('--prompts', default=','.join(PROMPTS.keys())) - ap.add_argument('--n-problems', type=int, default=200) - ap.add_argument('--n-rollouts', type=int, default=8, help='每题每思路的 skill 采样数（T>0）；executor 对每条 skill 各跑一次 greedy') - ap.add_argument('--seed', type=int, default=42) - ap.add_argument('--skill-temperature', type=float, default=0.5) - ap.add_argument('--skill-top-p', type=float, default=0.95) - ap.add_argument('--skill-max-tokens', type=int, default=8192) - ap.add_argument('--max-tokens', type=int, default=8192, help='executor，对齐 eval') - ap.add_argument('--eval-records', default=DEFAULT_EVAL_RECORDS) - ap.add_argument('--out', default=None) - args = ap.parse_args() - - names = [x for x in args.prompts.split(',') if x] - for x in names: - if x not in PROMPTS: - print(f'未知 prompt: {x}，可选: {list(PROMPTS)}') - sys.exit(1) - problems = load_problems(args.eval_records, args.n_problems, args.seed) - think = args.skill_thinking == 'on' - out_path = args.out or os.path.join(SCRIPT_DIR, f'skillcfg_{args.skill_thinking}.jsonl') - - # 两个采样器：skill(thinking 可切) + exec(enable_thinking=True, T=0, 对齐 v2 eval 的 base_sampler) - twinkle.initialize(mode='ray', nproc_per_node=SKILL_GPUS + EXEC_GPUS, lazy_collect=False, groups=[ - DeviceGroup(name='skill', ranks=list(range(SKILL_GPUS)), device_type='GPU'), - DeviceGroup(name='exec', ranks=list(range(SKILL_GPUS, SKILL_GPUS + EXEC_GPUS)), device_type='GPU')]) - - def make_sampler(group, world, enable_thinking): - s = vLLMSampler(model_id=MODEL_ID, - engine_args={'gpu_memory_utilization': GPU_MEM, - 'max_model_len': MAX_MODEL_LEN, 'tensor_parallel_size': 1}, - device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), - remote_group=group) - s.set_template(Template, model_id=MODEL_ID, enable_thinking=enable_thinking, max_length=MAX_MODEL_LEN) - return s - - skill_sampler = make_sampler('skill', SKILL_GPUS, enable_thinking=think) - exec_sampler = make_sampler('exec', EXEC_GPUS, enable_thinking=True) - - # ---- 1) 所有配置的 skill-gen 一次性 batch ---- - sg_prompts, meta = [], [] - for name in names: - for p in problems: - sg_prompts.append(skillgen_prompt(PROMPTS[name], p['problem'])) - meta.append((name, p)) - print(f'[skill-gen] {len(sg_prompts)} 条 x {args.n_rollouts} rollouts (prompts={len(names)} x 题={len(problems)}), ' - f'thinking={args.skill_thinking}, T={args.skill_temperature}') - sg_out = run_batch(skill_sampler, sg_prompts, args.skill_max_tokens, - args.skill_temperature, args.skill_top_p, SKILL_GPUS, - num_samples=args.n_rollouts) - - trials = [] - for (name, p), seqs in zip(meta, sg_out): - seqs = list(seqs or []) - for si in range(args.n_rollouts): - seq = seqs[si] if si < len(seqs) else None - full = _clean_text(getattr(seq, 'decoded', '') or '') if seq is not None else '' - sk = extract_skill(full) or '' - trials.append({'config': name, 'data_id': p['data_id'], 'sample_idx': si, - 'problem': p['problem'], - 'reference_answer': p['reference_answer'], 'baseline_pass': p['baseline_pass'], - 'skill': sk, 'parseable': bool(sk), 'skill_chars': len(sk), - 'leaked': answer_leaked(sk, p['reference_answer']), - 'skillgen_full': full, - 'skillgen_stop': getattr(seq, 'stop_reason', None) if seq is not None else 'empty', - 'skillgen_tokens': len(getattr(seq, 'tokens', None) or []) if seq is not None else 0}) - - # ---- 2) 所有配置的 executor 一次性 batch（空 skill 走 direct，等价 baseline 口径）---- - ex_prompts = [build_skill_solve_prompt(t['problem'], t['skill']) for t in trials] - print(f'[executor] {len(ex_prompts)} 条, T=0, max_tokens={args.max_tokens}') - ex_out = run_batch(exec_sampler, ex_prompts, args.max_tokens, 0.0, 1.0, EXEC_GPUS) - - with open(out_path, 'w') as f: - for t, seqs in zip(trials, ex_out): - roll = grade(seqs[0], t['reference_answer']) if seqs else { - 'pred': None, 'correct': False, 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} - t.update({'withskill_pred': roll['pred'], 'withskill_correct': roll['correct'], - 'withskill_stop': roll['stop_reason'], 'withskill_tokens': roll['gen_tokens'], - 'withskill_text': roll['text'], - 'skill_thinking': args.skill_thinking}) - f.write(json.dumps(t, ensure_ascii=False) + '\n') - - # ---- 3) 汇总（含题级 pass@k）---- - print('\n' + '=' * 118) - print('%-14s %-5s %-6s %-6s %-6s %-8s %-7s %-6s %-8s %-8s %-10s' % ( - 'config', 'n', 'parse', 'leak', 'trunc', 'skill字符', 'mean@1', 'base', 'lift', 'pass@k', 'hard救活@k')) - print('-' * 118) - for name in names: - sub = [t for t in trials if t['config'] == name] - n = len(sub) - parse = sum(t['parseable'] for t in sub) / n - leak = sum(t['leaked'] for t in sub) / n - trunc = sum(1 for t in sub if t['withskill_stop'] == 'length') / n - chars = int(st.median([t['skill_chars'] for t in sub if t['parseable']] or [0])) - acc = sum(t['withskill_correct'] for t in sub) / n - base = sum(t['baseline_pass'] for t in sub) / n - byq = {} - for t in sub: - byq.setdefault(t['data_id'], []).append(t) - p_at_k = sum(1 for v in byq.values() if any(x['withskill_correct'] for x in v)) / len(byq) - hardq = {d: v for d, v in byq.items() if v[0]['baseline_pass'] == 0} - rescue_k = (sum(1 for v in hardq.values() if any(x['withskill_correct'] for x in v)) / len(hardq)) if hardq else 0.0 - print('%-14s %-5d %-6.2f %-6.2f %-6.2f %-8d %-7.3f %-6.3f %+-8.3f %-8.3f %-10.3f' % ( - name, n, parse, leak, trunc, chars, acc, base, acc - base, p_at_k, rescue_k)) - print('-' * 118) - print(f'逐 trial 明细（含 skill-gen/executor 全文）已写入 {out_path}') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/skill2lora/logp_corr_probe.py b/cookbook/exp/skill2lora/logp_corr_probe.py deleted file mode 100644 index 86ec3fbb6..000000000 --- a/cookbook/exp/skill2lora/logp_corr_probe.py +++ /dev/null @@ -1,430 +0,0 @@ -#!/usr/bin/env python3 -"""logp_corr_probe.py — 找"低噪声、强预测 skill 好坏"的指标（E15 数据驱动 reward 选型）。 - -问题：E15 用 mean ΔlogP(GT|题+skill) 当稠密 reward，20 步 delta 不爬。要回答两件事： - 1) ΔlogP（及各种 per-token 聚合变体）到底和 skill 的真实有效性（executor 多 rollout - 通过率）相关吗？相关性多强？ - 2) 各候选指标的噪声多大？（greedy×1 判分 vs 8-rollout 真值 的一致性 = 老 0/1 reward 的噪声） - -数据：E15 gen_records（题/skill/GT 参考解/训练期 fp32 mean logps 都在盘上）。 -三阶段（分开进程跑，互不污染 Ray/vllm）： - --phase rollout twinkle vLLMSampler dp=8：每对 (题,skill) T=0.5×8 rollout + greedy×1， - 外加每题 baseline（无 skill）同口径 → 真值 pass_rate / lift。 - --phase logps 原生 vllm prompt_logprobs=0 + twinkle Template.encode 的 labels 定位 - response 段（与训练 _score_executor_mean_logps 同一模板/同一切位）， - 对 base(题+GT) 与 skill(题+skill+GT) 各算一遍逐 token logp → npz。 - --phase analyze CPU：各指标 vs 真值的 Spearman/AUC（全局 + 组内），噪声对比表。 - -用法（8 卡空闲时）： - cd cookbook/exp/skill2lora - PYTHONPATH=../../src python3 logp_corr_probe.py --phase rollout - PYTHONPATH=../../src python3 logp_corr_probe.py --phase logps - PYTHONPATH=../../src python3 logp_corr_probe.py --phase analyze -""" -import argparse -import copy -import json -import os -import re -import sys -from collections import defaultdict -from typing import Dict, List, Optional - -import numpy as np -from twinkle.data_format import pack_user_data - -SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) -E15_DIR = os.path.join(SCRIPT_DIR, 'output.ablate12', 'E15_logp_gt_on_narrative') -OUT_DIR = os.path.join(SCRIPT_DIR, 'logp_corr') -MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') -MAX_MODEL_LEN = int(os.environ.get('MAX_MODEL_LEN', 16384)) -MAX_TOKENS = 8192 # executor 解题预算，对齐 v2 -N_PROBLEMS = int(os.environ.get('PROBE_PROBLEMS', 64)) -N_ROLLOUTS = int(os.environ.get('PROBE_ROLLOUTS', 8)) -SEED = 42 - -# ---- executor prompt / 判分：逐字复刻 train_skill_v2 v2 分支（与 eval_skill_probe 相同） ---- -_ANSWER_FORMAT_V2 = ('Present your reasoning, then put ONLY the final numeric result inside ' - '\\boxed{}. For example: \\boxed{42}.') - - -def build_direct_prompt(problem): - content = f'The problem you need to solve:\n{problem}\n\n' + _ANSWER_FORMAT_V2 - return {'messages': [{'role': 'user', 'content': content}]} - - -def build_skill_solve_prompt(problem, skill): - skill = (skill or '').strip() - if not skill: - return build_direct_prompt(problem) - content = (f'The problem you need to solve:\n{problem}\n\n' - 'Skill hint:\nFor this problem, a skill-generation model has analyzed it and ' - 'provided some advisory skills:\n' - f'{skill}\n' - 'Prefer using its techniques when they fit, but if you have a more efficient or ' - 'clearer correct method, you may use it. If you diverge from this advice, briefly ' - 'explain why. Be concise and accurate.\n' - + _ANSWER_FORMAT_V2) - return {'messages': [{'role': 'user', 'content': content}]} - - -_BOXED_RE = re.compile(r'\\boxed\s*\{') - - -def extract_boxed(text): - if not text: - return None - last = None - for m in _BOXED_RE.finditer(text): - depth, i = 1, m.end() - while i < len(text) and depth > 0: - depth += (text[i] == '{') - (text[i] == '}') - i += 1 - if depth == 0: - last = text[m.end():i - 1].strip() - return last - - -_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') -_SEAM_TAG_RE = re.compile(r'<\s*answer\s*>(.*?)<\s*/\s*answer\s*>', re.I | re.S) -_SEAM_BOX_RE = re.compile(r'(?:\\{1,2}\(|)\\{1,2}boxed\s*\{\s*([^}]*)\s*}(?:\)|)', re.S) -_SEAM_INLINE_RE = re.compile(r'\$([^$]+)\$|\\\(([^)]+)\\\)', re.S) -_SEAM_FRAC_RE = re.compile(r'(-?\d+(?:\.\d+)?)/(-?\d+(?:\.\d+)?)') -_SEAM_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') - - -def _seam_norm(num): - try: - f = float(num) - return str(int(f)) if f == int(f) else str(f) - except Exception: - return num.strip() - - -def _seam_sanitize(txt): - txt = (txt or '').strip() - if (m := _SEAM_TAG_RE.search(txt)): - txt = m.group(1).strip() - elif (m := _SEAM_BOX_RE.search(txt)): - txt = m.group(1).strip() - elif (m := _SEAM_INLINE_RE.search(txt)): - txt = (m.group(1) or m.group(2)).strip() - # bugfix 2026-07-29：\dfrac/\tfrac/\cfrac 归一，同 train_skill_v2._seam_sanitize - txt = txt.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac').replace(r'\cfrac', r'\frac') - txt = re.sub(r'\\frac\s*\{\s*([^}]+?)\s*}\s*\{\s*([^}]+?)\s*}', r'\1/\2', txt) - if (m := _SEAM_FRAC_RE.search(txt)): - p, q = map(float, m.groups()) - if q: - return _seam_norm(str(p / q)) - if (m := _SEAM_NUM_RE.search(txt)): - return _seam_norm(m.group()) - return txt - - -def _judge(decoded, gold): - text = _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() - raw = extract_boxed(text) - pred = _seam_sanitize(raw) if raw else None - return bool(pred) and (pred == _seam_sanitize(str(gold))) - - -# ---- 配对采样：E15 gen_records -> pairs.jsonl ----------------------------------------- -def load_pairs(): - """64 题（seeded）× 组内全部 parseable 且有 delta 的候选（<=8）；GT 取每题首候选 rolls[0].text。""" - by_id = {} - for line in open(os.path.join(E15_DIR, 'gen_records.jsonl')): - r = json.loads(line) - if r.get('record_type') != 'problem' or not r.get('candidates'): - continue - by_id.setdefault(r['data_id'], r) # data_id 在 epoch 内唯一 - ids = sorted(by_id) - rng = np.random.RandomState(SEED) - pick = list(rng.permutation(len(ids))[:N_PROBLEMS]) - pairs, problems = [], [] - for k in pick: - r = by_id[ids[k]] - gt = next((c['rolls'][0]['text'] for c in r['candidates'] - if c.get('rolls') and c['rolls'][0].get('text')), '') - if not gt: - continue - cands = [c for c in r['candidates'] if c['parseable'] and c.get('logp_delta') is not None] - if len(cands) < 4: - continue - problems.append({'data_id': r['data_id'], 'problem': r['problem'], - 'reference_answer': r['reference_answer'], 'gt': gt}) - for j, c in enumerate(cands[:8]): - pairs.append({'pair_id': f'{r["data_id"]}#{j}', 'data_id': r['data_id'], - 'skill': c['skills'], 'leaked': bool(c['leaked']), - 'skill_chars': len(c['skills']), - 'skillgen_tokens': c.get('skillgen_tokens'), - 'logp_base_train': c['logp_base'], 'logp_skill_train': c['logp_skill'], - 'logp_delta_train': c['logp_delta']}) - return problems, pairs - - -# ---- phase: rollout ------------------------------------------------------------------- -def phase_rollout(): - import twinkle - from twinkle import DeviceGroup, DeviceMesh - from twinkle.data_format import SamplingParams, pack_user_data - from twinkle.sampler import vLLMSampler - from twinkle.template import Template - - problems, pairs = load_pairs() - os.makedirs(OUT_DIR, exist_ok=True) - json.dump(problems, open(os.path.join(OUT_DIR, 'problems.json'), 'w')) - with open(os.path.join(OUT_DIR, 'pairs.jsonl'), 'w') as f: - for p in pairs: - f.write(json.dumps(p, ensure_ascii=False) + '\n') - print(f'[rollout] problems={len(problems)} pairs={len(pairs)}', flush=True) - - n_gpu = int(os.environ.get('EXEC_GPUS', 8)) - twinkle.initialize(mode='ray', nproc_per_node=n_gpu, lazy_collect=False, - groups=[DeviceGroup(name='exec', ranks=list(range(n_gpu)), device_type='GPU')]) - sampler = vLLMSampler(model_id=MODEL_ID, - engine_args={'gpu_memory_utilization': 0.85, - 'max_model_len': MAX_MODEL_LEN, 'tensor_parallel_size': 1}, - device_mesh=DeviceMesh.from_sizes(world_size=n_gpu, dp_size=n_gpu), - remote_group='exec') - sampler.set_template(Template, model_id=MODEL_ID, enable_thinking=True, max_length=MAX_MODEL_LEN) - - prob_by = {p['data_id']: p for p in problems} - prompts, metas = [], [] - for p in problems: # baseline（无 skill） - prompts.append(build_direct_prompt(p['problem'])) - metas.append(('base', p['data_id'])) - for pr in pairs: # with-skill - prompts.append(build_skill_solve_prompt(prob_by[pr['data_id']]['problem'], pr['skill'])) - metas.append(('skill', pr['pair_id'])) - - def run(params, tag): - padded = prompts if len(prompts) % n_gpu == 0 else \ - prompts + [copy.deepcopy(prompts[-1])] * (n_gpu - len(prompts) % n_gpu) - outs = sampler.sample(padded, params)[:len(prompts)] - rows = [] - for (kind, key), resp in zip(metas, outs): - gold = prob_by[key.split('#')[0]]['reference_answer'] if '#' in key \ - else prob_by[key]['reference_answer'] - seqs = list(resp.sequences) if (resp and resp.sequences) else [] - rows.append({'kind': kind, 'key': key, 'mode': tag, - 'n': len(seqs), - 'pass': [bool(_judge(getattr(s, 'decoded', '') or '', gold)) for s in seqs], - 'trunc': [getattr(s, 'stop_reason', None) == 'length' for s in seqs], - 'tokens': [len(getattr(s, 'tokens', None) or []) for s in seqs]}) - return rows - - rows = run(SamplingParams(max_tokens=MAX_TOKENS, temperature=0.5, top_p=1.0, - num_samples=N_ROLLOUTS), f't05x{N_ROLLOUTS}') - rows += run(SamplingParams(max_tokens=MAX_TOKENS, temperature=0.0, top_p=1.0, - num_samples=1), 'greedy') - with open(os.path.join(OUT_DIR, 'rollout_results.jsonl'), 'w') as f: - for r in rows: - f.write(json.dumps(r) + '\n') - print(f'[rollout] done: {len(rows)} rows -> rollout_results.jsonl', flush=True) - - -# ---- phase: logps --------------------------------------------------------------------- -def phase_logps(): - """原生 vllm prompt_logprobs=0；token 布局与训练一致：twinkle Template.encode 的 labels - != -100 即 response(GT) 段位置。base 与 skill 两条轨迹各存一行 float32。""" - from twinkle.template import Template - from vllm import LLM, SamplingParams as VSP - from vllm.inputs import TokensPrompt - - problems = json.load(open(os.path.join(OUT_DIR, 'problems.json'))) - pairs = [json.loads(l) for l in open(os.path.join(OUT_DIR, 'pairs.jsonl'))] - prob_by = {p['data_id']: p for p in problems} - tmpl = Template(model_id=MODEL_ID, enable_thinking=True, max_length=MAX_MODEL_LEN, - truncation_strategy='delete') - - def encode(problem, skill, gt): - msgs = [dict(m) for m in build_skill_solve_prompt(problem, skill)['messages']] - enc = tmpl.encode({'messages': msgs + [{'role': 'assistant', 'content': gt}], - 'user_data': pack_user_data({'key_rounds': [len(msgs)]})}) - if enc is None: - return None, None - ids = [int(x) for x in enc['input_ids']] # numpy int64 -> int（vllm msgspec 拒收 np 类型） - pos = np.where(np.asarray(enc['labels']) != -100)[0] - return ids, pos - - jobs, keys = [], [] # key: ('base', data_id) / ('skill', pair_id) - for p in problems: - ids, pos = encode(p['problem'], '', p['gt']) - if ids is not None and len(pos): - jobs.append((ids, pos)) - keys.append(('base', p['data_id'])) - for pr in pairs: - p = prob_by[pr['data_id']] - ids, pos = encode(p['problem'], pr['skill'], p['gt']) - if ids is not None and len(pos): - jobs.append((ids, pos)) - keys.append(('skill', pr['pair_id'])) - print(f'[logps] encoded jobs={len(jobs)} (skipped {len(problems)+len(pairs)-len(jobs)})', flush=True) - - llm = LLM(model=_local_model_path(), max_model_len=MAX_MODEL_LEN, - gpu_memory_utilization=0.85, tensor_parallel_size=1) - sp = VSP(max_tokens=1, temperature=0.0, prompt_logprobs=0) - outs = llm.generate([TokensPrompt(prompt_token_ids=ids) for ids, _ in jobs], sp) - - store = {} - for (ids, pos), (kind, key), out in zip(jobs, keys, outs): - plp = out.prompt_logprobs - row = np.full(len(pos), np.nan, dtype=np.float32) - for i, p_ in enumerate(pos): - d = plp[int(p_)] if int(p_) < len(plp) else None - if d: - lp = d.get(ids[int(p_)]) - if lp is not None: - row[i] = lp.logprob - store[f'{kind}|{key}'] = row - np.savez_compressed(os.path.join(OUT_DIR, 'token_logps.npz'), **store) - print(f'[logps] saved {len(store)} rows -> token_logps.npz', flush=True) - - -def _local_model_path(): - from modelscope.hub.snapshot_download import snapshot_download - return snapshot_download(MODEL_ID, local_files_only=True) - - -# ---- phase: analyze ------------------------------------------------------------------- -def _rank(x): - x = np.asarray(x, dtype=np.float64) - order = np.argsort(x, kind='mergesort') - r = np.empty(len(x)) - r[order] = np.arange(len(x)) - for v in np.unique(x): # 平均并列名次 - m = x == v - if m.sum() > 1: - r[m] = r[m].mean() - return r - - -def spearman(a, b): - a, b = np.asarray(a, float), np.asarray(b, float) - m = ~(np.isnan(a) | np.isnan(b)) - if m.sum() < 3: - return np.nan - ra, rb = _rank(a[m]), _rank(b[m]) - sa, sb = ra.std(), rb.std() - return np.nan if sa == 0 or sb == 0 else float(((ra - ra.mean()) * (rb - rb.mean())).mean() / (sa * sb)) - - -def auc(score, label): - score, label = np.asarray(score, float), np.asarray(label, bool) - m = ~np.isnan(score) - score, label = score[m], label[m] - if label.sum() == 0 or (~label).sum() == 0: - return np.nan - r = _rank(score) - return float((r[label].sum() - label.sum() * (label.sum() - 1) / 2) / (label.sum() * (~label).sum())) - - -def phase_analyze(): - problems = json.load(open(os.path.join(OUT_DIR, 'problems.json'))) - pairs = [json.loads(l) for l in open(os.path.join(OUT_DIR, 'pairs.jsonl'))] - rolls = [json.loads(l) for l in open(os.path.join(OUT_DIR, 'rollout_results.jsonl'))] - npz = np.load(os.path.join(OUT_DIR, 'token_logps.npz')) - - base_pass, pair_pass, pair_greedy, base_greedy, pair_trunc = {}, {}, {}, {}, {} - for r in rolls: - rate = float(np.mean(r['pass'])) if r['n'] else np.nan - if r['kind'] == 'base' and r['mode'].startswith('t05'): - base_pass[r['key']] = rate - elif r['kind'] == 'skill' and r['mode'].startswith('t05'): - pair_pass[r['key']] = rate - pair_trunc[r['key']] = float(np.mean(r['trunc'])) if r['n'] else np.nan - elif r['kind'] == 'skill' and r['mode'] == 'greedy': - pair_greedy[r['key']] = float(r['pass'][0]) if r['n'] else np.nan - elif r['kind'] == 'base' and r['mode'] == 'greedy': - base_greedy[r['key']] = float(r['pass'][0]) if r['n'] else np.nan - - rows = [] - for pr in pairs: - key, did = pr['pair_id'], pr['data_id'] - b = npz[f'base|{did}'] if f'base|{did}' in npz.files else None - s = npz[f'skill|{key}'] if f'skill|{key}' in npz.files else None - if key not in pair_pass or did not in base_pass: - continue - row = {'pair_id': key, 'data_id': did, - 'truth_pass8': pair_pass[key], 'truth_lift': pair_pass[key] - base_pass[did], - 'base_pass8': base_pass[did], 'greedy1': pair_greedy.get(key, np.nan), - 'trunc_rate': pair_trunc.get(key, np.nan), - 'delta_train': pr['logp_delta_train'], - 'leaked': float(pr['leaked']), 'skill_chars': float(pr['skill_chars'])} - if b is not None and s is not None and len(b) == len(s): - d = s - b - ok = ~(np.isnan(d)) - d, bb = d[ok], b[ok] - if len(d): - row['delta_mean'] = float(d.mean()) - row['delta_sum'] = float(d.sum()) - k = min(50, len(d)) - row['delta_top50'] = float(d[np.argsort(-np.abs(d))[:k]].mean()) - unc = bb < -1.0 # executor 本来拿不准的 token - row['delta_uncertain'] = float(d[unc].mean()) if unc.sum() >= 5 else np.nan - row['delta_tail100'] = float(d[-min(100, len(d)):].mean()) - row['frac_improved'] = float((d > 0).mean()) - row['base_mean_ck'] = float(bb.mean()) - rows.append(row) - print(f'[analyze] usable pairs={len(rows)}') - json.dump(rows, open(os.path.join(OUT_DIR, 'pair_table.json'), 'w')) - - # 交叉校验：vllm 重算 delta vs 训练 fp32 delta - dm = [r.get('delta_mean', np.nan) for r in rows] - dt = [r['delta_train'] for r in rows] - print(f'\n[校验] corr(delta_vllm, delta_train) spearman={spearman(dm, dt):.3f}') - - metrics = ['delta_train', 'delta_mean', 'delta_sum', 'delta_top50', 'delta_uncertain', - 'delta_tail100', 'frac_improved', 'greedy1', 'skill_chars', 'leaked', 'trunc_rate'] - truth = np.array([r['truth_pass8'] for r in rows]) - lift = np.array([r['truth_lift'] for r in rows]) - helped = lift > 0 - - print('\n=== 全局相关性（n=%d 对）：指标 vs 8-rollout 真值 ===' % len(rows)) - print('%-16s %-14s %-14s %-10s' % ('metric', 'sp(pass8)', 'sp(lift)', 'AUC(lift>0)')) - for m in metrics: - v = np.array([r.get(m, np.nan) for r in rows], float) - print('%-16s %-14s %-14s %-10s' % ( - m, f'{spearman(v, truth):+.3f}', f'{spearman(v, lift):+.3f}', f'{auc(v, helped):.3f}')) - - # 组内（GRPO 真正用的信号）：每题 >=4 候选的组内 spearman 均值 - print('\n=== 组内相关性（每题组内 spearman 的均值±se）===') - by_p = defaultdict(list) - for r in rows: - by_p[r['data_id']].append(r) - for m in metrics: - cs = [] - for did, rs in by_p.items(): - if len(rs) < 4: - continue - v = [r.get(m, np.nan) for r in rs] - t = [r['truth_pass8'] for r in rs] - c = spearman(v, t) - if not np.isnan(c): - cs.append(c) - if cs: - cs = np.array(cs) - print('%-16s mean=%+.3f se=%.3f n_groups=%d' % (m, cs.mean(), cs.std() / np.sqrt(len(cs)), len(cs))) - - # 噪声对比：greedy×1 vs 真值；基线 greedy vs 基线8rollout - g = np.array([r.get('greedy1', np.nan) for r in rows], float) - m = ~np.isnan(g) - hard_wrong = np.abs(g[m] - truth[m]) > 0.5 - print(f'\n[噪声] greedy×1 与 8-rollout 真值强不一致率（|diff|>0.5）: {hard_wrong.mean():.3f} (n={m.sum()})') - bg = np.array([base_greedy.get(p["data_id"], np.nan) for p in problems], float) - bp = np.array([base_pass.get(p["data_id"], np.nan) for p in problems], float) - mm = ~(np.isnan(bg) | np.isnan(bp)) - print(f'[噪声] baseline greedy 与 baseline pass8 强不一致率: {(np.abs(bg[mm]-bp[mm])>0.5).mean():.3f} (n={mm.sum()})') - print(f'[分布] truth_pass8 mean={np.nanmean(truth):.3f} lift>0 比例={np.mean(helped):.3f} ' - f'lift<0 比例={np.mean(lift<0):.3f}') - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument('--phase', choices=('rollout', 'logps', 'analyze'), required=True) - args = ap.parse_args() - {'rollout': phase_rollout, 'logps': phase_logps, 'analyze': phase_analyze}[args.phase]() - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/skill2lora/rubric_effect.py b/cookbook/exp/skill2lora/rubric_effect.py deleted file mode 100644 index 9829dee2c..000000000 --- a/cookbook/exp/skill2lora/rubric_effect.py +++ /dev/null @@ -1,138 +0,0 @@ -#!/usr/bin/env python3 -"""rubric_effect.py — rubric 作用的三路数据分析（纯 CPU，现有数据）。 -① 臂级：E5/E6/E7(rl_ab, rubric 条件) vs E1/E2/E3(bnpo, query-only) eval lift 对照 -② 题级：同一错题上，rubric 条件生成的 skill vs query-only 生成的 skill 的 executor 通过率 -③ rubric 文本特征 vs A 线拯救率（组内 any-pass） -""" -import json -import hashlib -import os -import re -from collections import defaultdict - -import numpy as np - -BASE = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'output.ablate12') - -# gen_records 未存 rubric 文本；用全局缓存反查（key = md5('rubric_global\x1f'+data_id)）。 -# 能查到 => 该题被诊断过 => 裸题答错、属 A 线（rl_ab 只诊断错题）。 -_RUBRIC = {} -with open(os.path.join(BASE, 'rubric_cache_global.jsonl')) as f: - for l in f: - d = json.loads(l) - _RUBRIC[d['key']] = d.get('value') or '' - - -def rubric_of(data_id): - k = hashlib.md5(('\x1f'.join(['rubric_global', str(data_id)])).encode('utf-8')).hexdigest() - return _RUBRIC.get(k) - - -def load_gen(exp): - rows = [] - with open(os.path.join(BASE, exp, 'gen_records.jsonl')) as f: - for l in f: - r = json.loads(l) - rows.append(r) - return rows - - -def probe_schema(exp): - rows = load_gen(exp) - tps = defaultdict(int) - for r in rows: - tps[r.get('record_type')] += 1 - print(exp, dict(tps)) - for r in rows: - if r.get('record_type') == 'problem': - print(' problem keys:', sorted(r.keys())[:30]) - cands = r.get('cands') or r.get('_cands') or [] - if cands: - print(' cand keys:', sorted(cands[0].keys())) - break - - -if __name__ == '__main__': - import sys - if len(sys.argv) > 1 and sys.argv[1] == 'schema': - probe_schema('E7_rl_ab_on_pitfall') - probe_schema('E3_bnpo_on_pitfall') - sys.exit(0) - - # ---------- ② 题级同题对照 ---------- - # E7 A 线（rubric 非空）的候选 vs E3 同 data_id 的候选（query-only），配对比较组均值 - for pair in [('E7_rl_ab_on_pitfall', 'E3_bnpo_on_pitfall'), - ('E5_rl_ab_off_pitfall', 'E1_bnpo_off_pitfall'), - ('E6_rl_ab_off_narrative', 'E2_bnpo_off_narrative')]: - ea, eb = pair - ga, gb = load_gen(ea), load_gen(eb) - - def group_pass(rows, need_rubric=None): - out = {} - for r in rows: - if r.get('record_type') != 'problem': - continue - did = r.get('data_id', '') - rub = rubric_of(did) - if need_rubric is True and not rub: - continue - cands = [c for c in (r.get('candidates') or []) if c.get('parseable') - and c.get('with_pass') is not None] - if not cands: - continue - # 同一题可能多 chunk 出现，取第一次（早期，policy 漂移最小） - if did not in out: - out[did] = (np.mean([float(c['with_pass']) for c in cands]), rub or '') - return out - - pa = group_pass(ga, need_rubric=True) # A 线（rubric 条件） - pb = group_pass(gb) # query-only - common = sorted(set(pa) & set(pb)) - if not common: - print(f'[②] {ea} vs {eb}: 无同题交集') - continue - da = np.array([pa[d][0] for d in common]) - db = np.array([pb[d][0] for d in common]) - diff = da - db - print(f'[②] {ea.split("_")[0]}(rubric) vs {eb.split("_")[0]}(query-only) 同题 n={len(common)}: ' - f'rubric臂组均pass={da.mean():.3f} qonly臂={db.mean():.3f} ' - f'配对差={diff.mean():+.4f}±{diff.std()/np.sqrt(len(diff)):.4f} ' - f'win/tie/lose={int((diff>0).sum())}/{int((diff==0).sum())}/{int((diff<0).sum())}') - - # ---------- ③ rubric 特征 vs 拯救率 ---------- - ga = load_gen('E7_rl_ab_on_pitfall') - rows = [] - for r in ga: - if r.get('record_type') != 'problem': - continue - rub = (rubric_of(r.get('data_id', '')) or '').strip() - if not rub: - continue - cands = [c for c in (r.get('candidates') or []) if c.get('parseable') - and c.get('with_pass') is not None] - if not cands: - continue - n_fail = len(re.findall(r'\[FAIL\]', rub)) - n_pass = len(re.findall(r'\[PASS\]', rub)) - rows.append({'len': len(rub), 'n_fail': n_fail, 'n_pass': n_pass, - 'n_crit': n_fail + n_pass, - 'has_fix': int('fix:' in rub), - 'rescue': float(np.mean([float(c['with_pass']) for c in cands])), - 'any': float(any(c['with_pass'] for c in cands))}) - if rows: - def sp(a, b): - a, b = np.asarray(a, float), np.asarray(b, float) - ra = np.argsort(np.argsort(a)).astype(float) - rb = np.argsort(np.argsort(b)).astype(float) - if ra.std() == 0 or rb.std() == 0: - return np.nan - return float(((ra - ra.mean()) * (rb - rb.mean())).mean() / (ra.std() * rb.std())) - print(f'\n[③] E7 A线 rubric 特征 vs 拯救率 (n={len(rows)} 题, ' - f'mean rescue={np.mean([r["rescue"] for r in rows]):.3f}, ' - f'any-pass={np.mean([r["any"] for r in rows]):.3f})') - for fk in ['len', 'n_fail', 'n_pass', 'n_crit', 'has_fix']: - v = [r[fk] for r in rows] - print(' %-8s vs rescue %+0.3f | vs any-pass %+0.3f' % ( - fk, sp(v, [r['rescue'] for r in rows]), sp(v, [r['any'] for r in rows]))) - ls = np.array([r['len'] for r in rows]) - print(' rubric len 分布: p25=%d p50=%d p75=%d' % tuple(np.percentile(ls, [25, 50, 75]))) diff --git a/cookbook/exp/skill2lora/run_ablate12.sh b/cookbook/exp/skill2lora/run_ablate12.sh deleted file mode 100644 index a61f1edbc..000000000 --- a/cookbook/exp/skill2lora/run_ablate12.sh +++ /dev/null @@ -1,379 +0,0 @@ -#!/usr/bin/env bash -# ============================================================================== -# run_ablate12.sh — sequential launcher for the 12-experiment skill ablation. -# -# Reads the run plan from skill_ablate/config.py (single source of truth for order / -# dir names / think / skill-max-tokens / optional gate), then runs each experiment via -# `python -m skill_ablate.main --exp E{n}` in RUN_ORDER. -# -# Per experiment: -# - isolated product dir output.ablate12/<exp_dir>/ -# - idempotent: a successful run writes <exp_dir>/DONE.json (atomic, last step); both this -# script and skill_ablate.main skip completed experiments unless FORCE=1 -# - env snapshot output.ablate12/<exp_dir>/env_info.txt -# - skill-max-tokens 8192 (think) / 4096 (nothink) [from the plan] -# - E12 (sft, optional) SKIPPED unless RUN_SFT=1 -# - sleep between runs to let the previous Ray/vLLM engine tear down (avoid contention) -# -# Env knobs (all optional): -# DEEPMATH_DIR=$HERE/../../../deepmath_103k TRAIN_N=5000 MAX_UPDATES=50 EVAL_EVERY=5 -# LR=1e-6 RUN_SFT=1 FORCE=1 ONLY="E5 E6" SLEEP=30 SWANLAB_PROJECT=twinkle -# MIN_LEVEL=6 CHUNK_SIZE=32 (gradient-signal fix: E1/E5 audit — level<=5 all-pass -# dominated, 16-problem chunks leave only ~6 mixed groups per update; eval split unaffected) -# task=code 的臂（E4/E17，见 config.py）自动改走 BigCodeBench： -# BCB_PARQUET=... CODE_CHUNK_SIZE=48 CODE_EVAL_SIZE=200 CODE_TRAIN_N=0 TEST_WORKERS=24 -# ============================================================================== -set -euo pipefail - -# avoid backward-pass OOM from allocator fragmentation (E6 crash: 15GiB reserved-unallocated); -# inherited by the Ray training actors via twinkle's runtime env passthrough -export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" - -HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -cd "$HERE" - -# twinkle 是 editable 安装（.pth 指向 CPFS 上的 src/）；CPFS 瞬时抖动会让 site 初始化静默 -# 丢弃该路径，实验启动即死在 ModuleNotFoundError: twinkle（实测复现过一次）。PYTHONPATH 兜底。 -export PYTHONPATH="$(cd "$HERE/../../.." && pwd)/src${PYTHONPATH:+:$PYTHONPATH}" - -# central env file (optional): put all knobs in one place. ENV_FILE=xxx overrides the path. -ENV_FILE="${ENV_FILE:-$HERE/ablate12.env}" -if [ -f "$ENV_FILE" ]; then - echo "[ablate12] loading env from $ENV_FILE" - set -a; . "$ENV_FILE"; set +a -fi - -OUT_ROOT="${OUT_ROOT:-$HERE/output.ablate12}" -# interpreter with the full torch/vllm/twinkle stack; conda base shells shadow `python` with a -# numpy-less interpreter, so default to the absolute path and allow PYBIN=... to override. -PYBIN="${PYBIN:-/usr/local/bin/python3}" -# DeepMath-103K (difficulty-stratified loader in skill_ablate/data.py); replaces the old -# SEAM/aops input — see skill_quality_analysis.md 组成漂移修正. -DEEPMATH_DIR="${DEEPMATH_DIR:-$(cd "$HERE/../../.." && pwd)/deepmath_103k}" -TRAIN_N="${TRAIN_N:-5000}" -EVAL_SIZE="${EVAL_SIZE:-128}" -MAX_UPDATES="${MAX_UPDATES:-50}" -EVAL_EVERY="${EVAL_EVERY:-5}" -LR="${LR:-1e-6}" -MIN_LEVEL="${MIN_LEVEL:-6}" -CHUNK_SIZE="${CHUNK_SIZE:-32}" -SLEEP="${SLEEP:-30}" -# SWAN_PROJ alias: swanlab>=0.8 的 pydantic Settings 会解析进程 env 里的 SWANLAB_PROJECT 并报错， -# 所以外部换项目请用 SWAN_PROJ=xxx，不要 export SWANLAB_PROJECT。 -# bugfix #10：ENV_FILE 的 set -a 会把文件里的 SWANLAB_PROJECT 自动 export（正好踩中上面的坑）； -# 读完值后 unset 掉 export 属性，再以普通 shell 变量重建，保证子进程 env 里没有它。 -_SWAN_PROJ_VAL="${SWAN_PROJ:-${SWANLAB_PROJECT:-twinkle}}" -unset SWANLAB_PROJECT -SWANLAB_PROJECT="$_SWAN_PROJ_VAL" -RUN_SFT="${RUN_SFT:-0}" -FORCE="${FORCE:-0}" -ONLY="${ONLY:-}" -# leak 一律不进 reward（项目既定要求）：留空则不传 flag，由 main.py 默认值 0 生效。 -# 旧版在这里写死 1.0 且第 141 行无条件传入，会静默盖掉 Python 侧默认值。 -LOGP_LEAK_PENALTY="${LOGP_LEAK_PENALTY:-}" -# E16 passrate_hinge：截断铰链惩罚强度与起点、leak gate、base_tok 筛题阈（空=用 main.py 默认） -REWARD_TRUNC_PENALTY="${REWARD_TRUNC_PENALTY:-}" -REWARD_TRUNC_LO="${REWARD_TRUNC_LO:-}" -REWARD_LEAK_GATE="${REWARD_LEAK_GATE:-}" -BASE_TOK_FLOOR="${BASE_TOK_FLOOR:-}" -# kl_beta：对初始策略的锚，是唯一能提供“恢复力”对抗漂移的旋钮（reward 只能在组内排序）。 -# 空=用 main.py 默认 0.01（2026-07-29 从 0.001 上调；E1-E16 既往臂全跑在 0.001）。 -# RUN_TAG 同时隔离输出目录与 swanlab 实验名，用于并列跑同一 ExpSpec 的多个变体而不互相覆盖。 -KL_BETA="${KL_BETA:-}" -RUN_TAG="${RUN_TAG:-}" -# --- E17 reflexion 臂专属规模 ------------------------------------------------------------- -# 本臂只在【裸 executor 做错的题】上训练与评测，与其他臂不可横比（用户 2026-07-29 -# 拍板：只看本实验自身趋势），所以这些规模刻意与全局默认解耦。 -# 取值全部按 E16 落盘数据标定（.tmp_analysis/e17_param_calib.py，同模型/同题池/同难度门）： -# * K=24：E16 每次更新实际只有 23.46 组（总 1173 组）。累积证据 ∝sqrt(N)，而 E16 全程 -# 在文本层只累到 1.9 sigma —— 再砍组数就没任何判别力了。K=24 x 50 = 1200 组， -# 恰好追平 E16，而成本 24x8x8=1536 rollouts/chunk vs E16 实测 1501，几乎相等。 -# * CHUNK 128：实测裸错率 0.329（全量题池 527/1600，level>=6, T=0, floor=0）。 -# ⭐ 不是 0.446 —— 那是 base_tok>5000 筛选后子集的错误率（偏难），本臂 floor=0。 -# 二项精算 P(凑不满 24)：chunk 64 = 74%、96 = 3.6%、112 = 0.27%、128 = 0.012%。 -# 裸解开销 128 道 greedy 相当于 with-skill 的 8%，买对齐很便宜。 -# * TRAIN_N 8000：128 x 50 = 6400 次抽取，超过默认 5000 会进第二个 epoch（重复题）。 -# * EVAL 384 + --eval-min-level=MIN_LEVEL：指标只在错题子集上有信息量。旧口径（128 道 -# 全难度混合）只能给 ~33 道错题，SE 0.087，趋势根本读不出来；384 道 + 难度对齐 -# 训练池给 ~126 道（SE ~0.045）。正确的题跳过全部 GPU 路径，所以总成本几乎不变。 -REFLEXION_K="${REFLEXION_K:-24}" -E17_CHUNK_SIZE="${E17_CHUNK_SIZE:-128}" -E17_EVAL_SIZE="${E17_EVAL_SIZE:-384}" -E17_TRAIN_N="${E17_TRAIN_N:-8000}" -# E17 专属 kl_beta（2026-07-30 拍板 0.001，回到 E1-E16 旧值）。单独立一个变量而不改 -# 全局 KL_BETA，是为了不隐式改动其他臂重跑时的取值。 -E17_KL_BETA="${E17_KL_BETA:-0.001}" -# --- E18 rejection_sft 臂专属 -------------------------------------------------------------- -# 攒批阈值：16（用户 2026-07-30 拍板，从 128 改小）= 一个 sft_batch_size：chunk 32 每轮 -# 收 ~10 条胜者，大约隔 chunk 就 fire 一次，50 次更新 ~80 chunk 可达；128 要 ~15 chunk/次。 -E18_ACCUMULATE="${E18_ACCUMULATE:-16}" -# --- task=code 专属（2026-07-31 E4/E17 换到 BigCodeBench） -------------------------------- -# 规模按 bcb/bcb_eval0_probe.py + 2026-07-31 dry run 的落盘数据标定，与数学默认解耦： -# * 裸错率 **0.5625**（dry run 实测 18/32：训练侧 10/16 + eval 侧 8/16，think=on/8192）。 -# ⚠️ 不是 probe 的 0.715 —— 那是 nothink + 4096 的读数，think 开着以后模型强不少。 -# K=24 时 P(凑不满) : chunk 48 = 15.4%、56 = 1.6%、**64 = 0.1%**。组数恒定是本臂硬要求 -# （凑不满只会打印 k_short 并用更少的组训练，趋势就被抽样噪声污染），所以取 64。 -# 多花的只有裸解那 16 道 greedy（with-skill 部分由 K 固定，不随 chunk 变）。 -# * EVAL 200：题池 908（1140 剔缺库/外网GUI + 沙箱自检不过的 75），错题约 112 道，SE≈0.047。 -# * TRAIN_N=0 = 用掉剩下的全部题（708 道）。chunk64 x 50 ≈ 4.5 个 epoch 重复题（拍板接受）。 -# * TEST_WORKERS：判一条 rollout = 起一个 python 子进程跑 unittest（导入 pandas/sklearn 后 -# 典型 1-3s），一个 chunk 要判几百条，串行判分比同 chunk 的 GPU 时间还长。 -CODE_CHUNK_SIZE="${CODE_CHUNK_SIZE:-64}" -# ★ E4（bnpo）不受 K=24 那条约束 —— 它把 chunk 里每道题都拿来训练，chunk 直接等于每次更新的 -# 组数。跟着 reflexion 用 64 会白白把 rollout 数翻倍（64x8=512/更新），而且与已跑完的数学 -# E4（全局 CHUNK_SIZE=32）不再同规模、不可比。所以 view-B 的 code 臂单独用 32。 -CODE_BNPO_CHUNK_SIZE="${CODE_BNPO_CHUNK_SIZE:-32}" -CODE_EVAL_SIZE="${CODE_EVAL_SIZE:-200}" -CODE_TRAIN_N="${CODE_TRAIN_N:-0}" -BCB_PARQUET="${BCB_PARQUET:-$(cd "$HERE/../../.." && pwd)/bigcodebench/bcb.parquet}" -TEST_WORKERS="${TEST_WORKERS:-24}" -TEST_TIMEOUT="${TEST_TIMEOUT:-60}" -# --- executor nothink 对照（E19 math / E20 code，2026-07-31） ------------------------------- -# chunk / eval 都**不覆盖**：直接沿用同域 think 臂的值（math 128/384、code 64/200）。 -# 理由一，实测：首个 E19 run 的 c0 只从 64 道里拿到 20 道错题 —— math 关 think 的裸错率约 -# 0.31（baseline acc 0.69），与 think 的 0.329 基本相同，我此前估的 0.85 完全错了。 -# p=0.31 时 chunk 64 的期望错题 20±3.7，几乎每个 chunk 都凑不满 K=24，组数恒定失效。 -# code 侧 nothink p=0.622 > think 的 0.5625，chunk 64 本来就够，也无需覆盖。 -# 理由二，可比性：chunk 与 eval 都与同域 think 臂逐项相同，think/nothink 才是唯一自变量。 -# 保留这个 env 只为应急调参，默认空 = 不覆盖。 -EXEC_NOTHINK_CHUNK_SIZE="${EXEC_NOTHINK_CHUNK_SIZE:-}" -# eval 规模**不覆盖**：两个 nothink 臂各自的对照是同域的 think 臂，评测集必须逐题相同 —— -# E19 用 E17_EVAL_SIZE=384（已跑完的 E17 数学臂就是 n=384，baseline acc 0.674 / lift +0.130）， -# E20 用 CODE_EVAL_SIZE=200（与 E17 code 臂一致）。曾想为省 eval 开销把 nothink 统一压到 200， -# 那会让 E19 的评测集变成 E17 的子集、lift 曲线不再可逐题配对，省的钱不值这个代价。 -# --- E4/E17/E19/E20 统一口径（2026-07-31 用户拍板） ---------------------------------------- -# executor 生成预算统一 15000、skill 模型统一 8192（后者由 plan 的 smt 列给出，think=on 即 8192）。 -# 统一的意义：预算不再是四个臂之间的变量，think/nothink 的对比才是单变量的。 -# 连带两处必须跟着改，否则静默失效： -# 1) --max-model-len：默认 16384 装不下 prompt(约 1-2k) + 15000 输出，vLLM 会截 prompt；提到 20480。 -# 2) --reward-trunc-lo：长度惩罚死区按标定比例 5500/8192 缩放到预算上 = 15000*0.671 ≈ 10000。 -# 不改的话死区停在 5500（占预算 37%），会把远未撞墙的正常答案也纳入惩罚区。 -UNIFIED_MAX_TOKENS="${UNIFIED_MAX_TOKENS:-15000}" -UNIFIED_MAX_MODEL_LEN="${UNIFIED_MAX_MODEL_LEN:-20480}" -UNIFIED_TRUNC_LO="${UNIFIED_TRUNC_LO:-10000}" - -mkdir -p "$OUT_ROOT" - -# --- pull the run plan (name \t exp_dir \t think \t smt \t optional) ------------------- -# bugfix #9：用 $PYBIN（实验同一解释器）而非裸 python3，避免 plan/快照与实验环境不一致 -PLAN="$("$PYBIN" skill_ablate/config.py --plan)" - -snapshot_env() { # $1 = target file - { - echo "=== ablate12 env snapshot @ $(date -u +%FT%TZ) ===" - echo "host: $(hostname)" - echo "pybin: $PYBIN" - echo "python: $("$PYBIN" -c 'import sys;print(sys.version.split()[0])')" - echo "torch: $("$PYBIN" -c 'import torch;print(torch.__version__)' 2>/dev/null || echo NA)" - echo "vllm: $("$PYBIN" -c 'import vllm;print(vllm.__version__)' 2>/dev/null || echo NA)" - echo "transformers: $("$PYBIN" -c 'import transformers;print(transformers.__version__)' 2>/dev/null || echo NA)" - echo "CUDA_VISIBLE_DEVICES: ${CUDA_VISIBLE_DEVICES:-unset}" - echo "nvidia-smi:"; nvidia-smi --query-gpu=name,memory.total --format=csv,noheader 2>/dev/null || echo " (nvidia-smi NA)" - echo "GPU layout: TRAIN=${TRAIN_GPUS:-2} REF=${REF_GPUS:-2} SKILL_SAMPLER=${SKILL_SAMPLER_GPUS:-2} BASE_SAMPLER=${BASE_SAMPLER_GPUS:-2}" - echo "LLM_BACKUP set: $([ -n "${LLM_BACKUP_API_KEY:-}${LLM_BACKUP_BASE_URL:-}${OPENAI_API_KEY:-}" ] && echo yes || echo no)" - } > "$1" -} - -echo "[ablate12] run order:"; echo "$PLAN" | awk -F'\t' '{printf " %s -> %s (think=%s smt=%s opt=%s task=%s exec_think=%s)\n",$1,$2,$3,$4,$5,$6,$7}' - -while IFS=$'\t' read -r NAME EXP_DIR THINK SMT OPTIONAL TASK EXEC_THINK; do - [ -z "$NAME" ] && continue - TASK="${TASK:-math}" - EXEC_THINK="${EXEC_THINK:-on}" - if [ -n "$ONLY" ] && ! grep -qw "$NAME" <<< "$ONLY"; then - echo "[ablate12] $NAME skipped (not in ONLY='$ONLY')"; continue - fi - if [ "$OPTIONAL" = "1" ] && [ "$RUN_SFT" != "1" ]; then - echo "[ablate12] $NAME ($EXP_DIR) skipped: optional; set RUN_SFT=1 to run"; continue - fi - - EXP_OUT="$OUT_ROOT/$EXP_DIR${RUN_TAG:+.$RUN_TAG}" - if [ -f "$EXP_OUT/DONE.json" ] && [ "$FORCE" != "1" ]; then - echo "[ablate12] $NAME already done ($EXP_OUT/DONE.json); FORCE=1 to rerun"; continue - fi - mkdir -p "$EXP_OUT" - snapshot_env "$EXP_OUT/env_info.txt" - - echo "======================================================================" - echo "[ablate12] START $NAME -> $EXP_OUT (think=$THINK skill_max_tokens=$SMT"\ -"${KL_BETA:+ kl_beta=$KL_BETA}${RUN_TAG:+ tag=$RUN_TAG})" - echo "======================================================================" - LOG="$EXP_OUT/run.log" - FORCE_FLAG="" - [ "$FORCE" = "1" ] && FORCE_FLAG="--force" - E16_FLAGS="" - [ -n "$REWARD_TRUNC_PENALTY" ] && E16_FLAGS="$E16_FLAGS --reward-trunc-penalty $REWARD_TRUNC_PENALTY" - [ -n "$REWARD_TRUNC_LO" ] && E16_FLAGS="$E16_FLAGS --reward-trunc-lo $REWARD_TRUNC_LO" - [ -n "$REWARD_LEAK_GATE" ] && E16_FLAGS="$E16_FLAGS --reward-leak-gate $REWARD_LEAK_GATE" - [ -n "$BASE_TOK_FLOOR" ] && E16_FLAGS="$E16_FLAGS --base-tok-floor $BASE_TOK_FLOOR" - [ -n "$LOGP_LEAK_PENALTY" ] && E16_FLAGS="$E16_FLAGS --logp-leak-penalty $LOGP_LEAK_PENALTY" - [ -n "$KL_BETA" ] && E16_FLAGS="$E16_FLAGS --kl-beta $KL_BETA" - [ -n "$RUN_TAG" ] && E16_FLAGS="$E16_FLAGS --run-tag $RUN_TAG" - CHUNK_ARG="$CHUNK_SIZE" - EVAL_ARG="$EVAL_SIZE" - TRAIN_N_ARG="$TRAIN_N" - MIN_LEVEL_ARG="$MIN_LEVEL" - if [ "$TASK" = "code" ]; then - # BigCodeBench：题池/裸错率/判分方式全变，规模走 CODE_* 默认（见文件头注释）。 - # --min-level 一律传 0：BCB 没有 difficulty 字段，传非 0 只会让 data_code 打一行警告。 - CHUNK_ARG="$CODE_CHUNK_SIZE" - [ "$NAME" = "E4" ] && CHUNK_ARG="$CODE_BNPO_CHUNK_SIZE" - EVAL_ARG="$CODE_EVAL_SIZE" - TRAIN_N_ARG="$CODE_TRAIN_N" - MIN_LEVEL_ARG=0 - E16_FLAGS="$E16_FLAGS --task code --bcb-parquet $BCB_PARQUET"\ -" --test-workers $TEST_WORKERS --test-timeout $TEST_TIMEOUT" - echo "[ablate12] $NAME task=code: chunk=$CHUNK_ARG eval=$EVAL_ARG n=${TRAIN_N_ARG}(0=全池)"\ -" parquet=$BCB_PARQUET test_workers=$TEST_WORKERS" - fi - # reflexion 家族（E17 think / E19 math-nothink / E20 code-nothink）共用同一套协议开关。 - if [ "$NAME" = "E17" ] || [ "$NAME" = "E19" ] || [ "$NAME" = "E20" ]; then - # base_tok_floor 不在这里传：ReflexionMethod.__init__ 强制置 0 并告警，且那样 - # config 指纹里落的就是真实生效值（命令行重复传参只会制造歧义）。 - # ★ code 任务下 E17_* 这组数学标定值不适用（裸错率 0.329 -> 0.5625），保持上面 - # CODE_* 已设好的值不动。 - if [ "$TASK" != "code" ]; then - CHUNK_ARG="$E17_CHUNK_SIZE" - EVAL_ARG="$E17_EVAL_SIZE" - TRAIN_N_ARG="$E17_TRAIN_N" - E16_FLAGS="$E16_FLAGS --eval-min-level $MIN_LEVEL" - fi - E16_FLAGS="$E16_FLAGS --reflexion-k $REFLEXION_K" - # eval 口径改为 SEAM 式确定性单次（2026-07-30 拍板）：R=1 + T=0。 - # 代价：失去跨 4 个 skill 平均的降噪，题级读数从 5 档（0/.25/.5/.75/1）退为 0/1， - # 单点 SE 约为原来的 2 倍；换来的是与 SEAM val_kwargs(n=1,do_sample=False) 同口径。 - # 也因此与 E1-E16（R=4/T=0.5）的 eval 读数不同源，不可横比。 - E16_FLAGS="$E16_FLAGS --eval-rollouts 1 --eval-skill-temperature 0.0" - # kl_beta 回 0.001（与 E1-E16 一致）；放在这里覆盖，显式传的全局 KL_BETA 优先。 - [ -z "$KL_BETA" ] && E16_FLAGS="$E16_FLAGS --kl-beta $E17_KL_BETA" - echo "[ablate12] $NAME reflexion: chunk=$CHUNK_ARG k=$REFLEXION_K eval=$EVAL_ARG"\ -" n=$TRAIN_N_ARG task=$TASK eval_min_level=$MIN_LEVEL_ARG kl_beta=${KL_BETA:-$E17_KL_BETA}"\ -" eval_rollouts=1/T=0 (hard-subset protocol; NOT comparable to E1-E16)" - fi - if [ "$EXEC_THINK" = "off" ]; then - # executor 关 thinking：只加一个 flag。chunk/eval 一律沿用同域 think 臂（见文件头注释： - # 实测 math nothink 裸错率 0.31 ≈ think 的 0.329，压小 chunk 会让 K=24 凑不满）。 - [ -n "$EXEC_NOTHINK_CHUNK_SIZE" ] && CHUNK_ARG="$EXEC_NOTHINK_CHUNK_SIZE" - E16_FLAGS="$E16_FLAGS --executor-thinking off" - echo "[ablate12] $NAME executor=nothink: chunk=$CHUNK_ARG eval=$EVAL_ARG"\ -" (探针实测 nothink 截断 0.000、裸解 0.378>0.324、rubric 增量 +0.135 vs +0.080)" - fi - if [ "$NAME" = "E13" ]; then - # SEAM 论文设置复现（2026-08-01 用户拍板）。E13 本身已是 align='seam'（SEAM EXPERIENCE_PROMPT - # + executor 嵌套 prompt + lpem 整段 sanitize 判分）+ executor nothink（config 已改）+ skill 8192， - # executor 预算走 main.py 默认 max-tokens=8192 / max-model-len=16384 / n-skills=8，与 - # .tmp_analysis/SEAM/scripts/train_deepmath_paper.sh 的 executor 5120+8192 / K=8 同口径。 - # 这里再把三个 run 级默认拉到 SEAM run 口径（都可被同名 env 覆盖）： - # * MIN_LEVEL 0：SEAM 随机全池、不挑难题；我们默认 6=只挑最难档，会压平训练集 acc 曲线， - # 正是"我们从没见过 SEAM 那种上升曲线"的主因（见 2026-08-01 SNR 归因）。 - # * chunk 128：= SEAM train_batch_size；bnpo(view B) 无 reflexion 的 K=24 约束，放大安全。 - # * reward-trunc-penalty 0：SEAM reward = correct×format，无长度惩罚（train_skill_v2 头注）。 - # * eval R=1/T=0：对齐 SEAM val_kwargs(n=1,do_sample=False)，与 E17/E19/E20 同口径， - # 与 E1-E16(R=4/T=0.5) 不可横比。 - CHUNK_ARG="${E13_CHUNK_SIZE:-128}" - MIN_LEVEL_ARG="${E13_MIN_LEVEL:-0}" - [ -z "$REWARD_TRUNC_PENALTY" ] && E16_FLAGS="$E16_FLAGS --reward-trunc-penalty 0" - E16_FLAGS="$E16_FLAGS --eval-rollouts 1 --eval-skill-temperature 0.0" - # ---- 与 SEAM 逐行对齐的优化器参数（2026-08-01）---------------------------------- - # 上一版 E13 只对了 chunk/min_level/惩罚/eval，三个真正控制更新幅度的参数全部跑默认值， - # 与 SEAM 差得很远（实测后果：think 25 步从 3977 塔到 1942 token、reward 0.816→0.734）。 - # SEAM 侧真值来自 scripts/train_deepmath_paper.sh + verl 的 fsdp_workers.py:198-199 归一化： - # ppo_mini_batch_size=20 × rollout.n=8 = 160 全局序列（再 /4gpu = 40/gpu） - # → 每 batch 的 optimizer step 数 = 256/40 = 7（verl data.split 的余数也算一步）， - # twinkle 端 range(0,1024,160) 同样 7 步；且 mini<n 使 multi_step=True、启用 PPO ratio， - # 与 verl 用 rollout 时 old_log_prob 的行为一致（旧值 mini=n → 1 步/batch、ratio 恒=1）。 - # kl_loss_coef=0.001（main.py 默认 0.01，旧 E13 就是 0.01，整整差 10 倍）。 - # 显式同名 env 仍可覆盖。 - E16_FLAGS="$E16_FLAGS --ppo-mini-batch-size ${E13_PPO_MINI:-160}" - [ -z "$KL_BETA" ] && E16_FLAGS="$E16_FLAGS --kl-beta ${E13_KL_BETA:-0.001}" - # ---- 钉住训练 batch 序列（2026-08-02）------------------------------------------- - # verl 的 dataloader 默认 data.shuffle=True，所以 SEAM 的 step k 不是 train.parquet 的 - # 第 k 个 128 切片。实测（align5 chunk0 vs SEAM step1）：同一个 5000 题池、同 batch - # size，但两边第一步喂的 128 题交集只有 1 题 -> withskill 0.848 vs 0.789、 - # baseline 0.562 vs 0.634、lift +0.285 vs +0.155，全是抽样差。 - # seam_train_order.jsonl 是从 SEAM rollout dump 反推的真实喂题序列（40×128=5120 行， - # 用 .tmp_analysis/mk_seam_train_order.py 生成），挂上后 chunk k 逐题 == SEAM step k+1。 - # 置空 E13_TRAIN_ORDER 即可回到自己的 shuffle（但就不能逐 step 对了）。 - _order=${E13_TRAIN_ORDER-/mnt/data/yzhao/tastelikefeet/twinkle/.tmp_analysis/seam_train_order.jsonl} - if [ -n "$_order" ] && [ -f "$_order" ]; then - E16_FLAGS="$E16_FLAGS --train-order-file $_order" - elif [ -n "$_order" ]; then - echo "[ablate12] WARN: train order file not found: $_order (fallback to shuffle)" - fi - # 显存 / 聚合粒度（2026-08-02 修正）： - # verl 的等权聚合单元是 ppo_micro_batch_size_per_gpu=5 条（dp_actor.py 在每个 micro 内部 - # masked_mean，再 *1/gas 累加），每 step 共 160条/5 = 32 个「5 条组」等权。 - # 旧配置 TRAIN_FSDP=1/DP=2 + micro=4（=2 条/卡）给出 80 个「2 条组」等权 —— micro 越小越 - # 靠 sequence-mean，短序列的每 token 权重越高、压长度的分量越强。align6 实测后果： - # actor 输出 12382 -> 9265 chars 只用 6 个 chunk，format 0.883 -> 0.988； - # SEAM 走完同一段要 40 步（12248 -> 9981，format 0.883 -> 0.943）。 - # 即长度收缩快约 6 倍，是 40 步里唯一超出 SEAM 自身噪声的偏离（format Δ0.056 = 2.4×sd）。 - # TRAIN_FSDP=2 把 fp32 权重+Adam(约 64G) 分到 2 卡，腾出的显存正好够 5 条/卡； - # 此时 TRAIN_DP=1、sft=5 -> 32 个「5 条组」等权，与 verl 逐组一致。 - # REF_FSDP 必须同为 2：否则 REF_DP=2 而 micro=5 不整除，会报 Batch too small。 - _tmb=${E13_TRAIN_MICRO_BATCH:-$([ "${TRAIN_FSDP:-1}" = 2 ] && echo 5 || echo $((2*${TRAIN_GPUS:-2})))} - [ -n "${TRAIN_MICRO_BATCH:-}" ] && _tmb=$TRAIN_MICRO_BATCH - E16_FLAGS="$E16_FLAGS --train-micro-batch $_tmb" - echo "[ablate12] E13 SEAM-repro: chunk=$CHUNK_ARG min_level=$MIN_LEVEL_ARG train_micro_batch=$_tmb"\ -" ppo_mini=${E13_PPO_MINI:-160} kl_beta=${KL_BETA:-${E13_KL_BETA:-0.001}} reward_trunc_penalty=0 eval=R1/T0"\ -" executor=nothink skill_max_tokens=$SMT (executor 预算 8192/16384, K=n_skills 默认 8)" - fi - if [ "$NAME" = "E21" ]; then - # 显存：E21 每卡 80G，4B 不分片 + fp32 Adam 主权重≈64G、余量仅 ~16G，micro=8（自动档）必 OOM。 - # 默认 train_micro_batch=1×dp（= TRAIN_GPUS，最小且满足 %TRAIN_DP==0）。 - # 注：token_mean_scope 已改回 'micro'（= verl 口径），所以切 micro 不再是数学等价， - # 而是会改变聚合粒度（micro 越小越接近 sequence-mean）。E21 与其它臂横比时需记一笔。 - # 若仍 OOM：TRAIN_FSDP=2 REF_FSDP=2（分片权重、且 dp→1 允许 micro=1）。 - # 显式 TRAIN_MICRO_BATCH 优先；E21_TRAIN_MICRO_BATCH 单独可调。 - _tmb=${E21_TRAIN_MICRO_BATCH:-${TRAIN_GPUS:-2}} - [ -n "${TRAIN_MICRO_BATCH:-}" ] && _tmb=$TRAIN_MICRO_BATCH - E16_FLAGS="$E16_FLAGS --train-micro-batch $_tmb" - echo "[ablate12] E21 freeform: train_micro_batch=$_tmb (80G OOM guard; micro 口径下会改变聚合粒度)" - fi - case "$NAME" in - E4|E17|E19|E20) - # 四个臂统一 executor 预算 15000 + max_model_len 20480 + 长度惩罚死区 10000， - # 让 think/nothink 与 math/code 的对比都不夹带预算差异（2026-07-31 拍板）。 - # 显式传的 REWARD_TRUNC_LO 优先（上面已拼进 E16_FLAGS 的不会被这里覆盖）。 - E16_FLAGS="$E16_FLAGS --max-tokens $UNIFIED_MAX_TOKENS"\ -" --max-model-len $UNIFIED_MAX_MODEL_LEN" - [ -z "$REWARD_TRUNC_LO" ] && E16_FLAGS="$E16_FLAGS --reward-trunc-lo $UNIFIED_TRUNC_LO" - echo "[ablate12] $NAME 统一口径: max_tokens=$UNIFIED_MAX_TOKENS"\ -" max_model_len=$UNIFIED_MAX_MODEL_LEN skill_max_tokens=$SMT"\ -" reward_trunc_lo=${REWARD_TRUNC_LO:-$UNIFIED_TRUNC_LO}" - ;; - esac - if [ "$NAME" = "E18" ]; then - # eval 是 nothink 确定性单次（trainer 侧临时切 nothink 模板；R=1 + T=0 与 E17 同拍板）。 - # 其余规模全走全局默认（用户要求"配置不变"）；攒批阈值单独可调。 - E16_FLAGS="$E16_FLAGS --e18-accumulate $E18_ACCUMULATE --eval-rollouts 1 --eval-skill-temperature 0.0" - echo "[ablate12] E18 rejection_sft: chunk=$CHUNK_ARG eval=$EVAL_ARG n=$TRAIN_N_ARG"\ -" accumulate=$E18_ACCUMULATE eval=nothink/R=1/T=0" - fi - set +e - "$PYBIN" -m skill_ablate.main \ - --exp "$NAME" \ - --deepmath-dir "$DEEPMATH_DIR" \ - --n "$TRAIN_N_ARG" \ - --eval-size "$EVAL_ARG" \ - --output-dir "$EXP_OUT" \ - --skill-max-tokens "$SMT" \ - --max-updates "$MAX_UPDATES" \ - --eval-every-updates "$EVAL_EVERY" \ - --min-level "$MIN_LEVEL_ARG" \ - --chunk-size "$CHUNK_ARG" \ - --lr "$LR" \ - --swanlab-project "$SWANLAB_PROJECT" \ - $E16_FLAGS \ - $FORCE_FLAG \ - < /dev/null 2>&1 | tee "$LOG" - RC=${PIPESTATUS[0]} - set -e - if [ "$RC" != "0" ]; then - echo "[ablate12] $NAME FAILED (rc=$RC); see $LOG. Stopping."; exit "$RC" - fi - echo "[ablate12] $NAME done. Sleeping ${SLEEP}s for engine teardown..." - sleep "$SLEEP" -done <<< "$PLAN" - -echo "[ablate12] all requested experiments finished." diff --git a/cookbook/exp/skill2lora/skill_ablate/__init__.py b/cookbook/exp/skill2lora/skill_ablate/__init__.py deleted file mode 100644 index 374d0aeef..000000000 --- a/cookbook/exp/skill2lora/skill_ablate/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Skill-generation ablation package (view A / view B × think × style × training method). - -Design: reuse train_skill_v2.py primitives verbatim (import, never edit); this package only -adds the experiment matrix, the sample pool, the rubric double-cache, and the pluggable -training methods on top. See cookbook/exp/skill2lora/skill_quality_analysis.md sections -"AI 最终清单" and "AI 接口方案" for the frozen design decisions this code implements. -""" diff --git a/cookbook/exp/skill2lora/skill_ablate/config.py b/cookbook/exp/skill2lora/skill_ablate/config.py deleted file mode 100644 index ea84b1d33..000000000 --- a/cookbook/exp/skill2lora/skill_ablate/config.py +++ /dev/null @@ -1,318 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Declarative ablation matrix (E1-E12) + run order. - -Frozen decisions (skill_quality_analysis.md): -- 12 experiments; run nothink before think, the SFT method (E12) LAST and manually gated. -- Unified knobs: executor frozen T=0; skill-model train T=1.0 × 8 rollouts; eval T=0.5 × 4 - rollouts (query-only, no rubric); skill-max-tokens = 8192 (think) / 4096 (nothink). -- view B = query-only; view A = rubric line, eval still query-only (knowledge-transfer probe). - -This module is intentionally dependency-free (pure stdlib) so it can be imported and unit- -smoke-tested without torch / twinkle / a GPU. -""" -from dataclasses import dataclass -from typing import Dict, List - -# --- training methods (internal keys) -------------------------------------------------- -# bnpo view B query-only GRPO/BNPO main loop (reuses v2 process_chunk verbatim). -# rl_ab view A RL, AB split: first bare-problem greedy solve; WRONG problems -> A line -# (skill sampled under query+rubric), RIGHT problems -> B line (query-only); -# both go through executor greedy -> reward -> in-group BNPO, trained together. -# rl_err view A RL, error-only: same as rl_ab but the B line is NOT trained -# (single-variable contrast vs rl_ab on "does training the right-answer B line help"). -# opsd view A On-Policy Self-Distillation: student(query-only) logps pulled toward -# teacher(query+rubric) logps per token (loss='opsd'); error problems only. -# improve_sft view A improve-skill + SFT: first-pass 1 skill; correct -> positive pool -# (no leak, <=4096 chars); wrong -> rubric regen (2-in-8 pick 1) -> negative pool; -# 1:1 accumulate -> SFT. -# sft view A plain SFT: bare-problem wrong -> rubric -> regen (2-in-8) -> accumulate SFT. -# logp_rl E14+: query-only skill-gen, executor T>0 samples a correct pseudo-GT solution S, -# rubric/API audits that executor response, then reward each skill by -# Δ mean logP_executor(S | problem + skill) with answer-leak penalty. -# logp_gt E15: same dense executor-logP reward as logp_rl, but the target S is DeepMath's -# external R1 reference solution (record 'solution'); NO executor rollout and NO -# rubric audit (view B, query-only) -> runs much faster. Validates the logP path -# against a strong external target instead of a self-sampled pseudo-GT. -# passrate_hinge E16: view B query-only. Reward = pass_rate(M rollouts, T>0) minus a hinge -# truncation penalty (only rollouts whose executor output nears the 8192 budget are -# penalized) minus a leak gate; trained problems are pre-filtered to the danger band -# (baseline executor output long / not all-pass). Data-driven closure of the reward -# probe (skill_quality_analysis.md 2026-07-29): pass_rate is the only real signal, -# trunc is the strongest dense side-signal, base_tok is the strongest problem filter. -METHODS = ('bnpo', 'rl_ab', 'rl_err', 'opsd', 'improve_sft', 'sft', 'logp_rl', 'logp_gt', - 'passrate_hinge', 'reflexion', 'rejection_sft') -VIEW_OF_METHOD = {'bnpo': 'B', 'rl_ab': 'A', 'rl_err': 'A', 'reflexion': 'A', - 'opsd': 'A', 'improve_sft': 'A', 'sft': 'A', 'logp_rl': 'A', 'logp_gt': 'B', - 'passrate_hinge': 'B', 'rejection_sft': 'A'} -STYLES = ('narrative', 'pitfall', 'freeform') -THINKINGS = ('on', 'off') -TASKS = ('math', 'code') - - -@dataclass(frozen=True) -class ExpSpec: - name: str # E1..E14 - method: str # one of METHODS - thinking: str # 'on' | 'off' - style: str # 'narrative' | 'pitfall' (ignored by align='seam': SEAM prompts bypass style) - optional: bool = False # E12(sft): manually gated (RUN_SFT=1), runs last - align: str = 'v2' # 'v2' | 'seam' — sets v2._ALIGN_MODE (prompt/判分/executor 嵌套全开关) - # E14+ 稠密 reward：executor 先按 T>0 采样 K 条找正确伪 GT S；训练 reward 不再 rollout， - # 而是算 Δ mean logP_executor(S | problem + skill)。默认字段复用 reward_rollouts/temperature。 - reward_rollouts: int = 1 - reward_temperature: float = 0.0 - smt_override: int = 0 # force skill_max_tokens regardless of the think rule; 0 = default rule - # 任务族（2026-07-31 用户拍板把 E4/E17 换到 BigCodeBench）： - # 'math' = DeepMath-103K + \boxed{} 数值判分（E1-E16、E18 原样） - # 'code' = BigCodeBench + 跑 unittest 判分，executor/skill-gen/rubric 三套 prompt 全换 - # 依据：数学域上 rubric 无增量甚至负向（−0.056），BFCL 上为零（+0.002），只有 BigCodeBench - # 这种"机器给出可定位失败证据（异常/断言/行号）"的任务上 rubric 才有增量（+0.135, p=4e-5， - # 见 bcb/bcb_eval0_probe.py 与 code_task.py 模块注释）。 - task: str = 'math' - # executor(base_sampler) 的 thinking。E1-E18 全是 'on'（历史口径，勿动）。 - # 'off' = 2026-07-31 新增的 nothink 对照（E19 math / E20 code）：bcb 探针实测 think 的 - # executor 有 34-50% rollout 撞满预算且是字面死循环（8-gram 重复率 p50=0.835），加预算 - # 到 20000 无效；关掉后截断归零、裸解 0.378 > 0.324、rubric 增量 +0.135 vs +0.080。 - # 注意这是 executor 侧；skill 模型仍由 thinking 字段控制（两个新臂都保持 skill think=on， - # 因为 v2 实测 skill 侧 nothink 会把完整解答写进 <skills>，等于换标签的泄漏）。 - executor_thinking: str = 'on' - - @property - def view(self) -> str: - return VIEW_OF_METHOD[self.method] - - @property - def needs_rubric(self) -> bool: - return self.view == 'A' - - @property - def skill_max_tokens(self) -> int: - # 显式 per-spec override 优先。 - # ⚠️ 2026-07-29 实测推翻了此前“think 长度与 skill 质量无关，截掉长尾不损失信号”的探针结论： - # E4 在 8192 与 4096 下的受控 A/B（其余配置全同，比共有的 chunk 0-21）显示 4096 是净损失—— - # parse 率 0.939 -> 0.740，无条件 cand_pass 0.689 -> 0.649，撞 think 顶 0.063 -> 0.268。 - # 4096 臂看起来“exec 更短、trunc 更低”纯属幸存者偏差（pass|parse 0.734 -> 0.877）， - # 因为 26% 的候选连 <skills> 都没写完就掉出了统计口径。 - # 见 .tmp_analysis/think_budget_ab.py。think 模式一律用 8192。 - if self.smt_override: - return self.smt_override - # think must have room for <think> + <skills> (4096 truncates to an empty block). - # seam align: 人工拍板用 8192（不复刻 SEAM 原版 4096：think 模式下 4096 会把大量候选截断在 - # <think> 里、压低 parseable，与“think 模式 skill-max-tokens 必须 8192”的矩阵规范保持一致）。 - if self.align == 'seam': - return 8192 - return 8192 if self.thinking == 'on' else 4096 - - @property - def loss(self) -> str: - return 'opsd' if self.method == 'opsd' else 'bnpo' - - @property - def exp_dir(self) -> str: - # output.ablate12/E{n}_{method}_{think}_{style}/ (seam align / multi-rollout reward 加后缀区分) - suffix = '_seam' if self.align == 'seam' else '' - if self.reward_rollouts > 1: - suffix += f'_r{self.reward_rollouts}' - # 换数据集必须换目录：否则新语义的 run 会撞上已跑完的数学 run（DONE.json 直接跳过、 - # 曲线与 gen_records 混在一起、config 指纹对不上）。 - if self.task != 'math': - suffix += f'_{self.task}' - # executor 口径也必须进目录名：同一个 ExpSpec 换 executor thinking 后 baseline 缓存 - # （key=problem）与 eval 读数全变，撞同一个目录会把两套语义的曲线混在一起。 - if self.executor_thinking != 'on': - suffix += '_execnothink' - return f'{self.name}_{self.method}_{self.thinking}_{self.style}{suffix}' - - @property - def swanlab_exp(self) -> str: - return f'ablate12_{self.exp_dir}' - - -# --- the 12-experiment matrix (declarative; order field below drives execution) -------- -MATRIX: List[ExpSpec] = [ - # group 1 — view B BNPO: think × style, no-rubric baseline - ExpSpec('E1', 'bnpo', 'off', 'pitfall'), - ExpSpec('E2', 'bnpo', 'off', 'narrative'), - ExpSpec('E3', 'bnpo', 'on', 'pitfall'), - # E4/E8 的 smt_override=4096 已被上面 skill_max_tokens 里记录的 A/B 推翻（净损失 4 个点 - # 无条件 pass + 20 个点 parse 率）。E8 保留原值只是为了不改动“已跑完的臂”的可复现配置。 - # ★ E4 于 2026-07-31 换到 code 任务并同时删掉 override（用户拍板）：既然要重跑，就按 - # think=on 的规范值 8192 跑，与 E17 同口径可比。旧数学 E4 的产物在 - # output.ablate12/E4_bnpo_on_narrative/（新 run 落在 ..._code/，互不覆盖）。 - ExpSpec('E4', 'bnpo', 'on', 'narrative', task='code'), - # group 2 — view A RL-AB-mix: same grid as E1-E4, isolates "rubric rescues zero-grad groups" - ExpSpec('E5', 'rl_ab', 'off', 'pitfall'), - ExpSpec('E6', 'rl_ab', 'off', 'narrative'), - ExpSpec('E7', 'rl_ab', 'on', 'pitfall'), - ExpSpec('E8', 'rl_ab', 'on', 'narrative', smt_override=4096), # 同 E4：重跑前应删掉此 override - # group 3 — view A training-method comparison (fixed think+narrative), sft last & optional - ExpSpec('E9', 'rl_err', 'on', 'narrative'), - ExpSpec('E10', 'opsd', 'on', 'narrative'), - ExpSpec('E11', 'improve_sft', 'on', 'narrative'), - ExpSpec('E12', 'sft', 'on', 'narrative', optional=True), - # group 4 — SEAM-align ablation: same data pipeline as the rest of the matrix, but ALL - # prompt/parsing/executor-nesting rules follow SEAM (align='seam' -> v2._ALIGN_MODE): - # actor uses SEAM EXPERIENCE_PROMPT (<memory_item>), executor sees the nested - # prompt_text+response_text(+think), lpem-parity greedy scoring, actor budget 4096. - # Query-only BNPO main loop (= SEAM's training form); eval stays the matrix-unified - # query-only readout so E13 is directly comparable with E1-E12. - ExpSpec('E13', 'bnpo', 'on', 'narrative', align='seam', executor_thinking='off'), - # group 5 — E14+ 稠密 reward：E4 的部署形态（query-only skill-gen / eval 不变），训练时 - # executor 先 T=0.7×16 采样，选本地判分正确且非截断的伪 GT S，并用 rubric/API 产审计诊断； - # skill reward = Δ mean logP_executor(S | problem + skill)（leak 不进 reward，只做监控）。该臂同时降测量噪声 - # 和抬内容信号，替代旧版 T=0.5×4 多数 rollout 0/1 reward。 - ExpSpec('E14', 'logp_rl', 'on', 'narrative', reward_rollouts=16, reward_temperature=0.7), - # group 6 — E15 稠密 reward 的 GT 版验证：与 E14 同一套 executor logP reward，但 logP 目标 S - # 换成 DeepMath 自带的 R1 参考解（record 'solution'），不再 executor rollout、不再 rubric 审计 - # （view B / query-only）。用于验证“ΔlogP(强外部参考解 | 题+skill)”是否走得通；因省掉 K 次 - # executor 采样 + rubric API，单 chunk 比 E14 快很多。 - ExpSpec('E15', 'logp_gt', 'on', 'narrative'), - # group 7 — E16 数据驱动收敛臂：探针（skill_quality_analysis.md 2026-07-28/29）判死 logP， - # 确认 pass_rate 是唯一真实信号、trunc 是最强稠密辅助、base_tok>5000 是最强筛题维度。 - # reward = mean_i(correct_i·eff_i) - kappa·mean_i(1-eff_i)，护栏 max(reward, pass_rate-1/M)； - # eff = (1-alpha·len_pen)·(1-beta·loop_pen)，len_pen 从 5500 起二次凸爬升（数据标定见 - # .tmp_analysis/reward_shape_calib.py）。leak 不进 reward，只做监控。 - # skill_max_tokens=8192（2026-07-29 拍板）：4096 的 A/B 判为净损失，见 skill_max_tokens 注释。 - # 训练题预筛危险带（baseline 输出长）。view B query-only，不用 rubric。 - ExpSpec('E16', 'passrate_hinge', 'on', 'narrative', - reward_rollouts=8, reward_temperature=0.5, smt_override=8192), - # group 8 — E17 Reflexion 臂：唯一目的是检验"rubric 注入权重外信息"能否让 skill 变得可学。 - # 与 E16 的三处结构差异（2026-07-29 人工拍板）： - # 1) 只在【裸 executor 做错】的题上训练与评测，做对的题完全不碰。E16 的 lift 分解 - # +0.170 = 救回 +0.230 - 破坏 -0.060（skill 挂在裸对的题上会砸掉 10.6% 保持率）， - # 条件化按构造把破坏项归零。见 .tmp_analysis/lift_source_decomp.py。 - # 2) skill-gen 走 view A（query+rubric）。E16 判定 reward 对"该写什么文本"几乎无信息 - # （结果层 SNR 2.25 -> 文本层 0.046，损失 ~50 倍）；rubric 是外部 API 的诊断，是本臂 - # 唯一的新变量，也是它可能不重演 E16 结局的唯一理由。见 .tmp_analysis/batch_size_math.py。 - # 3) 不用 base_tok 危险带筛题（base_tok_floor=0）。 - # ⚠️ 当初的理由已被实测证伪，保留在此以免重蹈：本以为"去掉筛选后错题集就是 - # 推理错 + 没写完的混合，才测得到方法修正"。实测（e17_param_calib.py + E16 全 1600 - # 题）：全量题池的 527 道错题里 96.96% 是 base_tok>=8192（没写完），只有 3.04% - # （16 道）是写完但答错；而 floor=5000 筛选后是 97.71%。也就是说 floor 不是截断主导的 - # 原因，题目本身是（level>=6 配 8192 预算），去掉筛选只把可用的"方法错"样本从 - # 2.3% 提到 3.0%（每 chunk 24 道错题里平均 0.73 道）。保留 floor=0 只是因为它严格 - # 不差于 floor=5000，不要再把它当成本臂能测到 reflexion 的理由。 - # 后果：rubric 在 ~97% 的题上只能说"你超预算了"，signal/wrong_trunc_frac 会直接 - # 开在 0.97。想真正测"方法修正"必须先把 executor 预算提到 16384。 - # 见 .tmp_analysis/e16_redteam5.py、e17_param_calib.py。 - # 批量对齐：chunk_size=128 裸解后取 --reflexion-k=24 道错题，每次更新恒定 24 题 x 8 候选。 - # K=24 不是拍的：E16 每次更新实际只有 23.46 组（全程 1173 组），而累积证据 ∝sqrt(N) - # 且 E16 在文本层只累到 1.9 sigma，再砍组数就没判别力了。chunk=128 是为了把 - # P(凑不满 24) 压到 0.012%（全量题池裸错率实测 0.329，不是筛选后子集的 0.446）。 - # 见 .tmp_analysis/e17_param_calib.py。 - # reward 形状沿用 E16 的 passrate_hinge（死区 5500/二次凸/leak 只监控），但 rollout 数改为 - # M=1（2026-07-30 拍板）：每个 skill 只让 executor 推理一次，reward = 这一次对/错。 - # ⭐两个连带后果（改前必读）： - # 1) pass_rate 从 9 档（0,1/8,...,1）退化为 2 档（0/1），组内无方差的概率大升 - # —— 这正是 E4(M=1) 62% 零梯度组的成因，M=8 当初就是为了绕开它。 - # 2) methods.py 的不可反转护栏 pen_cap = 1/M - 1e-6：M=8 时是 0.125（惩罚只能微调 - # 同档内排序），M=1 时变成 ~1.0，长度/兜圈惩罚能把做对的候选拉到贴近 0。 - # 形式保证（做对过永远排在没做对前）仍成立，但 reward 量级上从“以正确率为主” - # 变成“以长度惩罚为主”。 - # T=0（对齐 SEAM 的 grm.rollout.temperature=0 与 E4/E13）：打分时 executor 贪心确定性解一次， - # 同一个 skill 重跑 reward 不变 —— 消掉 executor 采样噪声，reward 只反映 skill 本身的差异。 - # （skill 模型自己的采样温度是另一个参数 skill_gen_temperature=1.0，不受此影响。） - # ★ 2026-07-31 换数据集（用户拍板）：task='code' —— BigCodeBench + 跑 unittest 判分。 - # 换的理由是上面 (2)(3) 两条在数学域已被实测封死：97% 的裸失败是"没写完"，rubric 只能 - # 说"你超预算了"，三个数据集横比下来 rubric 增量 −0.056(deepmath) / +0.002(BFCL) / - # **+0.135(BigCodeBench, p=4e-5)**，唯一的差别是 judge 手上有没有机器给出的可定位失败 - # 证据（异常类型/断言差异/失败用例名）。code 分支把这份证据喂进 judge（code_task.diag_segment）。 - # 连带变化：裸错率从数学的 0.329 变成 0.5625（2026-07-31 dry run 实测 18/32，think=on/8192； - # ⚠️ 不是 probe 的 0.715，那是 nothink+4096 的读数），所以 K=24 需要 chunk=64 - # （P(凑不满 24)=0.001；chunk 48 会有 15.4% 的更新组数不足）；题池 908 题，chunk64×50 - # ≈4.5 个 epoch 重复题（用户同意）。--min-level/--eval-min-level 在 code 下自动失效。 - ExpSpec('E17', 'reflexion', 'on', 'narrative', task='code', - reward_rollouts=1, reward_temperature=0.0, smt_override=8192), - # group 9 — E18 拒绝采样 SFT（2026-07-30 用户拍板）：采集与 E17 同源（裸解错题 -> rubric -> - # rubric 条件化 skill-gen think 模式，executor greedy T=0 单次判分），但不做 RL：每题在做对的 - # 候选里按 leak 过滤 -> 长度贴近 len_budget -> 与原始 rubric 相似度最高 三道筛取唯一胜者， - # 写本地数据集文件，攒够 16 条 SFT 一次（weight=1，nothink 布局响应；2026-07-30 从 128 改小， - # 使 50 次更新在 ~80 chunk 内可达），训完同步权重到 vLLM。 - # eval：同一个 skill_sampler vLLM 临时切 nothink 模板跑 query-only greedy（trainer 侧实现）。 - # reward_rollouts/temperature 对本臂无效（判分固定 greedy 单次），填 1/0.0 只为指纹如实。 - ExpSpec('E18', 'rejection_sft', 'on', 'narrative', - reward_rollouts=1, reward_temperature=0.0, smt_override=8192), - # group 10 — E19/E20 executor-nothink 对照（2026-07-31 用户拍板）：与 E17 逐字同构的 - # reflexion 臂，唯一变量是 **executor 关 thinking**（skill 侧仍 think=on）。 - # 起因：bcb 探针在 4096/12288/20000 三档 think 预算下测到 rubric 增量 +0.047/+0.058/+0.080， - # 而 nothink 一档是 +0.135（p=1e-4）；增量与"裸失败里 no_code 的占比"严格反向 - # （69%/54%/50%/0%）。截断样本经查是字面死循环（8-gram 重复率 p50=0.835、同一长句重复 - # 92 次），所以加预算无解、只能关 think。两个臂分别回答： - # E19（math）：deepmath + 数学 prompt/rubric。数学域此前的 rubric 增量是 −0.056，而 - # 那次 97% 的裸失败是"没写完"；关掉 executor think 后失败会变成"写完但答错"， - # 这是第一次能在数学域上把"rubric 是否有用"与"截断"分开测。 - # E20（code）：BigCodeBench + code prompt/rubric，直接把探针里 +0.135 那一档搬到训练。 - # ⚠️ 规模不能照抄：裸错率随 executor 口径变（code think 0.5625 -> nothink 约 0.62； - # math nothink 未测，预期远高于 think 的 0.329），chunk 由 run_ablate12.sh 的 - # EXEC_NOTHINK_* 单独给，dry run 后按实测复核。 - ExpSpec('E19', 'reflexion', 'on', 'narrative', executor_thinking='off', - reward_rollouts=1, reward_temperature=0.0, smt_override=8192), - ExpSpec('E20', 'reflexion', 'on', 'narrative', task='code', executor_thinking='off', - reward_rollouts=1, reward_temperature=0.0, smt_override=8192), - # group 11 — E21 freeform 文体（2026-08-01 用户拍板）：以 E2 为模板的 query-only BNPO（view B、 - # 无 rubric），唯一变量是 skill-gen system 换成 SKILL_GEN_FREEFORM“招式菜单”：不锁 narrative/ - # pitfall/toy 固定文体，让模型按题自选最有用的形态（分析/概念/预判纠错/迷你示范/直白执行 - # 指令，甚至 “let's think step by step”），T=1.0×8 自然铺开、组内择优。动机：固定 hint - # 消融实测 hint 内容语义贡献≈0、增益几乎全来自“有个 skill 块 + 催答案收尾”（fixed_hint_probe.py： - # A9_wrapperonly/A4_garbage 与有义 hint 打平、A7_budget 最高 +0.16），故放开文体看模型能否自选出更优组合。 - # ★ 刷 thinking='on'（非照 E2 的 off）：freeform prompt 依赖“先私下想再选形态”，nothink 下无处 - # 思考会把推理直接写进 <skills>（line 758 记录的泄漏失败模式）。其余同 E2：bnpo/math/narrative-长度预算。 - ExpSpec('E21', 'bnpo', 'on', 'freeform'), -] - -# execution order: all nothink first, then think; E13 (seam-align baseline) right after E6; -# E14 (rubric pseudo-GT + executor logP dense reward) right after E7 per 2026-07-28 人工拍板; -# E15 (GT-target logP validation) right after E14; the data-hungry SFT method dead last. -# 2026-07-31 用户拍板：E19/E20（executor nothink）排在 E4/E17 之前先跑 —— 探针已判定 -# nothink 是 rubric 增量最大且唯一无截断混杂的口径，先拿这两个臂的结论。 -RUN_ORDER: List[str] = ['E1', 'E2', 'E5', 'E6', 'E13', 'E3', 'E7', 'E14', 'E15', 'E19', 'E20', 'E21', 'E4', 'E8', 'E16', 'E17', 'E18', 'E9', 'E10', 'E11', 'E12'] - -BY_NAME: Dict[str, ExpSpec] = {e.name: e for e in MATRIX} - - -def get_spec(name: str) -> ExpSpec: - key = name.strip().upper() - if key not in BY_NAME: - raise KeyError(f'unknown experiment {name!r}; valid: {sorted(BY_NAME)}') - return BY_NAME[key] - - -def ordered_specs(include_optional: bool = True) -> List[ExpSpec]: - specs = [BY_NAME[n] for n in RUN_ORDER] - return specs if include_optional else [s for s in specs if not s.optional] - - -def _self_check() -> None: - """Invariants that guard against typos when editing the matrix.""" - assert set(BY_NAME) == set(RUN_ORDER), 'RUN_ORDER must cover every matrix entry exactly once' - assert len(RUN_ORDER) == len(set(RUN_ORDER)) == len(MATRIX), 'duplicate / missing names' - for e in MATRIX: - assert e.method in METHODS, f'{e.name}: bad method {e.method}' - assert e.thinking in THINKINGS and e.style in STYLES, f'{e.name}: bad think/style' - assert e.align in ('v2', 'seam'), f'{e.name}: bad align {e.align}' - assert e.task in TASKS, f'{e.name}: bad task {e.task}' - assert e.executor_thinking in THINKINGS, \ - f'{e.name}: bad executor_thinking {e.executor_thinking}' - assert not (e.task == 'code' and e.align == 'seam'), \ - f'{e.name}: seam align is math-only (SEAM prompts/parsing are \\boxed 数值口径)' - assert e.reward_rollouts >= 1, f'{e.name}: bad reward_rollouts {e.reward_rollouts}' - # nothink-before-think ordering within contiguous runs is a soft convention, not asserted. - - -if __name__ == '__main__': - import sys - _self_check() - if '--plan' in sys.argv: - # machine-readable run plan for the launcher: - # name<TAB>exp_dir<TAB>think<TAB>smt<TAB>optional<TAB>task<TAB>executor_thinking - for e in ordered_specs(): - print(f'{e.name}\t{e.exp_dir}\t{e.thinking}\t{e.skill_max_tokens}\t' - f'{int(e.optional)}\t{e.task}\t{e.executor_thinking}') - sys.exit(0) - print(f'{len(MATRIX)} experiments; run order: {" -> ".join(RUN_ORDER)}') - hdr = f'{"name":<4} {"view":<4} {"method":<12} {"think":<6} {"style":<10} {"align":<5} {"smt":<5} {"loss":<5} opt' - print(hdr) - print('-' * len(hdr)) - for e in ordered_specs(): - print(f'{e.name:<4} {e.view:<4} {e.method:<12} {e.thinking:<6} {e.style:<10} ' - f'{e.align:<5} {e.skill_max_tokens:<5} {e.loss:<5} {"Y" if e.optional else ""}') diff --git a/cookbook/exp/skill2lora/skill_ablate/data.py b/cookbook/exp/skill2lora/skill_ablate/data.py deleted file mode 100644 index 895fc0523..000000000 --- a/cookbook/exp/skill2lora/skill_ablate/data.py +++ /dev/null @@ -1,106 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""DeepMath-103K loader with difficulty-stratified train/eval split. - -Dataset: AI-ModelScope/DeepMath-103K (columns: question / final_answer / difficulty / topic / -r1_solution_1..3). We keep only rows whose final_answer normalizes to a number via v2 -``_numeric_value`` (the \\boxed{} judging pipeline is numeric-exact; answers like ``\\phi^4`` -cannot be scored and are dropped). - -Stratified split (skill_quality_analysis.md 组成漂移修正): difficulty is bucketed to its -rounded integer level; ``eval_size`` problems are sampled with per-bucket quotas proportional -to the pool (largest-remainder rounding), the rest form the train pool — so train and eval -difficulty proportions match by construction. All sampling is seeded and file-order stable: -``data_id = dm:<level>:<global_row_index>`` is reproducible across runs/experiments. - -Train-only difficulty floor (``--min-level``): E1/E5 gradient audit showed level<=5 groups are -dominated by all-pass (level 3: 63-74% all-pass, corr(level, mixed_rate)=0.92), i.e. mostly -zero-gradient. The floor drops those rows from the *train pool only*; the eval split keeps the -full-level stratification so eval/baseline stay comparable across experiments. -""" -import glob -import os -from collections import defaultdict -from typing import Any, Dict, List, Tuple - -import numpy as np - -import train_skill_v2 as v2 - - -def _read_rows(deepmath_dir: str) -> List[Dict[str, Any]]: - import pyarrow.parquet as pq - paths = sorted(glob.glob(os.path.join(deepmath_dir, '**', '*.parquet'), recursive=True)) - if not paths: - raise FileNotFoundError(f'no parquet files under --deepmath-dir {deepmath_dir}') - rows: List[Dict[str, Any]] = [] - for p in paths: - # r1_solution_1: E14-ref 的 logP 目标文本（外部 R1 参考解，强于 executor 自采解） - t = pq.read_table(p, columns=['question', 'final_answer', 'difficulty', 'r1_solution_1']) - rows.extend(t.to_pylist()) - return rows - - -def load_deepmath_records(args) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: - """-> (train_records, eval_records), each record {'data_id','problem','reference_answer'}.""" - rows = _read_rows(args.deepmath_dir) - pool: List[Dict[str, Any]] = [] - for i, r in enumerate(rows): # global row index over sorted files = stable id - problem = (r.get('question') or '').strip() - num = v2._numeric_value(r.get('final_answer')) - if not problem or num is None: - continue - lvl = int(round(float(r.get('difficulty') or 0))) - pool.append({'data_id': f'dm:{lvl}:{i}', 'problem': problem, - 'reference_answer': num, '_level': lvl, - 'solution': (r.get('r1_solution_1') or '').strip()}) - - # bucket by level, seeded shuffle inside each bucket - buckets: Dict[int, List[Dict[str, Any]]] = defaultdict(list) - for rec in pool: - buckets[rec['_level']].append(rec) - rng = np.random.RandomState(args.seed) - for lvl in sorted(buckets): - rng.shuffle(buckets[lvl]) - - # eval quota per bucket: proportional, largest-remainder rounding. - # --eval-min-level>0 把配额限在难度不低于它的桶里（默认 0 = 全难度混合，与旧臂逐字一致）。 - # E17 需要它：该臂的指标只在【裸解做错的题】上有信息量，而全难度混合的错题率只有 - # ~26%（E16 实测 128 -> 33 道，SE 0.087，趋势读不出来）；level>=6 上是 ~43%，同样的 - # GPU 成本能换到 1.6 倍的有效样本，同时与训练池（min_level）同分布。 - eval_min_level = int(getattr(args, 'eval_min_level', 0) or 0) - elig = {lvl: b for lvl, b in buckets.items() if lvl >= eval_min_level} - n_elig = sum(len(b) for b in elig.values()) - eval_n = min(args.eval_size, n_elig) if (args.eval_size > 0 and n_elig) else 0 - quota = {lvl: eval_n * len(b) / n_elig for lvl, b in elig.items()} if n_elig else {} - take = {lvl: int(q) for lvl, q in quota.items()} - for lvl in sorted(quota, key=lambda x: quota[x] - int(quota[x]), reverse=True): - if sum(take.values()) >= eval_n: - break - take[lvl] += 1 - - eval_records, train_records = [], [] - for lvl in sorted(buckets): - b = buckets[lvl] - n_ev = take.get(lvl, 0) - eval_records.extend(b[:n_ev]) - train_records.extend(b[n_ev:]) - min_level = int(getattr(args, 'min_level', 0) or 0) - if min_level > 0: # train-only floor; eval keeps full-level mix (see module docstring) - n_before = len(train_records) - train_records = [r for r in train_records if r['_level'] >= min_level] - v2.logger.info(f'[data] min_level={min_level}: train pool {n_before} -> {len(train_records)}') - rng.shuffle(train_records) # ProblemPool reshuffles too; this decorrelates level runs - if args.n > 0: # optional stratified-in-expectation downsample (pool already shuffled) - train_records = train_records[:args.n] - - def _lvls(rs): - c = defaultdict(int) - for r in rs: - c[r['_level']] += 1 - return {k: round(v / len(rs), 3) for k, v in sorted(c.items())} - v2.logger.info(f'[data] DeepMath: pool={len(pool)} (numeric-only of {len(rows)}) ' - f'train={len(train_records)} eval={len(eval_records)}') - v2.logger.info(f'[data] level mix train={_lvls(train_records)} eval={_lvls(eval_records)}') - for r in eval_records + train_records: - r.pop('_level', None) - return train_records, eval_records diff --git a/cookbook/exp/skill2lora/skill_ablate/data_code.py b/cookbook/exp/skill2lora/skill_ablate/data_code.py deleted file mode 100644 index 89b78830e..000000000 --- a/cookbook/exp/skill2lora/skill_ablate/data_code.py +++ /dev/null @@ -1,82 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""BigCodeBench loader for the ablation package (code task family). - -Dataset: bigcodebench/bcb.parquet (v0.1.4, 1140 tasks). Record shape matches the math loader -so nothing downstream changes shape: - data_id = task_id (e.g. BigCodeBench/42) - problem = instruct_prompt (ends with the exact imports + signature to reproduce) - reference_answer = code_task.payload_of(task) # 判分载荷，不是数值答案 - -与 DeepMath loader 的三处结构差异： -1. **没有 difficulty**，所以没有分层抽样，也不存在 --min-level / --eval-min-level（在 code - 模式下被显式忽略并告警）。train/eval 就是同一个 seed 洗牌后的前后切分。 -2. **题池极小**：1140 题，剔掉缺库/需外网GUI子进程的题后约 900，再减 eval 后训练池只有几百题。 - E17 用 chunk 48 × 50 updates = 2400 次抽取 ≈ 3 个 epoch 重复题（用户 2026-07-31 拍板接受）。 - 重复题的 rubric 全部缓存命中，所以重复的成本只在 GPU rollout。 -3. **有"沙箱自检"这道闸**：参考解答跑不过它自己的单测 = 环境不可判定（缺库的边角、随机种子、 - matplotlib 后端等），这类题的 0 分与模型能力无关，必须剔掉，否则它会给每个臂加一层同样的 - 噪声底并稀释 lift。probe 实测 7.5% 属于这一类。自检结果按 parquet 落一个 json 缓存， - 之后各臂零成本复用。 -""" -import json -import os -from typing import Any, Dict, List, Tuple - -import code_task -import train_skill_v2 as v2 - - -def _broken_task_ids(args, tasks: List[Dict[str, Any]]) -> set: - """参考解答跑不过自己单测的题（缓存到 --output-dir 的父目录，跨臂复用）。""" - base = (getattr(args, 'rubric_global_dir', None) - or os.path.dirname(os.path.abspath(str(args.output_dir).rstrip('/')))) - os.makedirs(base, exist_ok=True) - path = os.path.join(base, 'bcb_broken_tasks.json') - if os.path.exists(path): - try: - with open(path, encoding='utf-8') as f: - cached = json.load(f) - if int(cached.get('n_tasks', -1)) == len(tasks): - return set(cached.get('broken') or []) - v2.logger.info(f'[data] {path} 的题数 {cached.get("n_tasks")} != 当前 {len(tasks)}，' - f'重跑自检') - except Exception as exc: - v2.logger.warning(f'[data] 读取 {path} 失败（{exc}），重跑自检') - v2.logger.info(f'[data] 沙箱自检：{len(tasks)} 道题跑参考解答（一次性，之后走缓存）…') - broken = code_task.selftest(tasks, args.test_workers, args.test_timeout) - with open(path, 'w', encoding='utf-8') as f: - json.dump({'n_tasks': len(tasks), 'broken': sorted(broken)}, f, indent=1) - return set(broken) - - -def load_code_records(args) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: - """-> (train_records, eval_records)，每条 {'data_id','problem','reference_answer'}。""" - tasks, stats = code_task.load_tasks(args.bcb_parquet, args.seed) - if not tasks: - raise FileNotFoundError(f'no usable BigCodeBench task in {args.bcb_parquet}') - v2.logger.info(f"[data] BigCodeBench: 全集 {stats['raw']}，剔除依赖缺失 " - f"{stats['drop_missing_lib']}、需外网/GUI/子进程 {stats['drop_needs_net_or_gui']}" - f" -> {stats['kept']}") - if getattr(args, 'code_selftest', True): - broken = _broken_task_ids(args, tasks) - if broken: - tasks = [t for t in tasks if t['task_id'] not in broken] - v2.logger.info(f'[data] 剔除参考解答自己跑不过单测的题 {len(broken)} 道 ' - f'（沙箱不可判定，非模型能力）-> 可用 {len(tasks)}') - for k in ('min_level', 'eval_min_level'): - if int(getattr(args, k, 0) or 0): - v2.logger.warning(f'[data] --{k.replace("_", "-")} 在 code 任务下无效（' - f'BigCodeBench 没有 difficulty 字段），已忽略') - recs = [{'data_id': t['task_id'], 'problem': t['instruct_prompt'], - 'reference_answer': code_task.payload_of(t)} for t in tasks] - eval_n = min(args.eval_size, len(recs)) if args.eval_size > 0 else 0 - eval_records, train_records = recs[:eval_n], recs[eval_n:] - if args.n > 0: - train_records = train_records[:args.n] - if not train_records: - raise ValueError(f'--eval-size {args.eval_size} 吃掉了整个题池（可用 {len(recs)}）') - epochs = (args.chunk_size * args.max_updates) / max(1, len(train_records)) - v2.logger.info(f'[data] train={len(train_records)} eval={len(eval_records)}；' - f'按 chunk={args.chunk_size} x max_updates={args.max_updates} 估算约 ' - f'{epochs:.1f} 个 epoch（题会重复，rubric 走缓存）') - return train_records, eval_records diff --git a/cookbook/exp/skill2lora/skill_ablate/eval_reflexion.py b/cookbook/exp/skill2lora/skill_ablate/eval_reflexion.py deleted file mode 100644 index fe79faf48..000000000 --- a/cookbook/exp/skill2lora/skill_ablate/eval_reflexion.py +++ /dev/null @@ -1,208 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""E17 专用 eval：reflexion 协议 —— 只在裸 executor 做错的题上做 rubric 条件化 skill 干预。 - -与 v2 ``run_greedy_eval`` 的区别只有"作用域"和"skill-gen 的输入"两点： - -1. baseline 正确的题**完全不动**（不生成 skill、不重跑 executor），按定义计 1.0。理由：本臂的 - 命题是"reflexion 能不能救回不会的题"，正确题上的 with-skill rollout 既不提供信息、又在 - E16 里贡献了全部 -0.060 的破坏项，把它留在指标里只会用一个已知的、与命题无关的效应稀释 - 趋势。因此当 rubric 无缺失时 acc 恰好等于 ``base + (1-base) * hard_rescue``，两者只差一个 - 线性缩放，**主指标是 hard_rescue_rate**。 -2. hard 子集的 skill-gen 输入 = query + rubric（与训练同分布），rubric 由裸解轨迹诊断得来。 - 这修掉了 E6 的已知缺陷（训练在 query+rubric 分布、eval 却是 query-only）。 - -成本：rubric 条目的缓存键 = data_id + 裸解轨迹，而 eval baseline 是冻结+缓存的，所以同一道 -eval 题在整个 run 里只会调一次 rubric API；skill-gen / executor 也只跑 hard 子集（≈40%）。 - -⚠️ 口径不与 E1-E16 横比（用户 2026-07-29 拍板：只看本实验自身趋势）。 -""" -import sys -from typing import Any, Dict, List, Tuple - -import train_skill_v2 as v2 -from train_skill_v2 import (_clean_text, _extract_skill, _run_samples, build_direct_prompt, - build_skill_solve_prompt) - -from .methods import _rubric_entry -from .rollouting import rubric_skillgen_prompt - - -def _baseline_rolls(base_sampler, eval_records, base_dp, args, base_cache) -> List[Dict[str, Any]]: - """裸 greedy(T=0) 判分，走 v2 DiskCache（与 run_greedy_eval 同一缓存文件、同一键）。""" - todo = [r for r in eval_records if v2.DiskCache.key_for(r['problem']) not in base_cache] - if todo: - out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], - 1, args.max_tokens, base_dp, temperature=0.0) - # 批量判分：code 任务下每条判分是一个跑单测的子进程，200 道题串行要 ~7 分钟。 - rolls = v2._parse_many([(v2._first_seq(seqs), r['reference_answer']) - for r, seqs in zip(todo, out)]) - for r, roll in zip(todo, rolls): - base_cache.put(v2.DiskCache.key_for(r['problem']), roll) - return [base_cache.get(v2.DiskCache.key_for(r['problem'])) for r in eval_records] - - -def _diagnose(rubric_cache, checker, jobs: List[Dict[str, Any]], workers: int) -> List[str]: - from concurrent.futures import ThreadPoolExecutor - if not jobs or rubric_cache is None: - return [''] * len(jobs) - with ThreadPoolExecutor(max_workers=max(1, min(workers, len(jobs)))) as ex: - return list(ex.map(lambda e: rubric_cache.get_or_diagnose(e, checker) or '', jobs)) - - -def _gen_skills(skill_sampler, prompts, R, skill_dp, args) -> List[List[Tuple[str, str]]]: - """每题采 R 个 skill；返回 [(skill_block, raw_response)] * R（缺位补空串，与 v2 同）。""" - sg_out = _run_samples(skill_sampler, prompts, R, args.skill_max_tokens, skill_dp, - temperature=args.eval_skill_temperature) - per = [] - for seqs in sg_out: - seqs = list(seqs or []) - row = [] - for j in range(R): - s = seqs[j] if j < len(seqs) else None - if s is None: - row.append(('', '')) - else: - sresp = _clean_text(getattr(s, 'decoded', '') or '') - row.append((_extract_skill(sresp) or '', sresp)) - per.append(row) - return per - - -def run_reflexion_eval(base_sampler, skill_sampler, eval_records, ci, rounds, - base_dp, skill_dp, args, base_cache, rubric_cache, checker): - """返回 (recs, summary, metrics)，键名与 v2.run_greedy_eval 兼容（trainer 打印共用）。""" - R = max(1, args.eval_rollouts) - base_rolls = _baseline_rolls(base_sampler, eval_records, base_dp, args, base_cache) - - # ---- 1) 切分：baseline 正确的题按协议原样通过，不做任何 GPU 工作 ---- - hard: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [] - recs: List[Dict[str, Any]] = [] - head = {'record_type': 'eval_problem', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, - 'protocol': 'reflexion', 'n_rollouts': R, - 'eval_skill_temperature': args.eval_skill_temperature} - for r, br in zip(eval_records, base_rolls): - if br['correct']: - recs.append({**head, 'data_id': r.get('data_id', ''), 'problem': r['problem'], - 'reference_answer': r['reference_answer'], 'baseline_pass': 1.0, - 'intervened': False, 'rubric_ok': None, - # 协议：不干预 -> 保持 baseline 结果 - 'withskill_acc_mean': 1.0, 'withskill_acc_strict_mean': 1.0, - 'withskill_pass_any': 1.0, 'skill_parseable_mean': 1.0, - 'withskill_terminated_mean': 1.0 if br['terminated'] else 0.0}) - else: - hard.append((r, br)) - - # ---- 2) hard 子集：rubric 诊断（纯 API，缓存命中后零成本） ---- - # entry 直接复用训练侧的 _rubric_entry：判据表与 fail_segment 的构成（code 任务下含单测 - # 真实报错）必须与训练逐字同源，否则 eval 的干预分布与训练分布不同。 - entries = [_rubric_entry(r, br) for r, br in hard] - diags = _diagnose(rubric_cache, checker, entries, args.rubric_workers) - todo = [(r, br, d) for (r, br), d in zip(hard, diags) if d] - n_rubric_missing = len(hard) - len(todo) - - # ---- 3) rubric 条件化 skill-gen -> with-skill greedy 重跑 ---- - per_skills = _gen_skills(skill_sampler, [rubric_skillgen_prompt(r['problem'], d) - for r, _br, d in todo], R, skill_dp, args) \ - if todo else [] - flat_prompts, flat_idx = [], [] - for pi, ((r, _br, _d), row) in enumerate(zip(todo, per_skills)): - for j, (sk, sresp) in enumerate(row): - flat_prompts.append(build_skill_solve_prompt(r['problem'], sk, sresp)) - flat_idx.append((pi, j)) - ws_out = _run_samples(base_sampler, flat_prompts, 1, args.max_tokens, base_dp, - temperature=0.0) if flat_prompts else [] - ws_rolls = v2._parse_many([(v2._first_seq(seqs), todo[pi][0]['reference_answer']) - for (pi, _j), seqs in zip(flat_idx, ws_out)]) - roll_by = {idx: roll for idx, roll in zip(flat_idx, ws_rolls)} - - hard_recs: List[Dict[str, Any]] = [] - for pi, ((r, br, d), row) in enumerate(zip(todo, per_skills)): - rolls = [roll_by[(pi, j)] for j in range(len(row))] - corr = [1.0 if x['correct'] else 0.0 for x in rolls] - parses = [1.0 if sk else 0.0 for sk, _ in row] - terms = [1.0 if x['terminated'] else 0.0 for x in rolls] - rec = {**head, 'data_id': r.get('data_id', ''), 'problem': r['problem'], - 'reference_answer': r['reference_answer'], 'baseline_pass': 0.0, - 'intervened': True, 'rubric_ok': True, 'rubric': d, - 'base_stop_reason': br.get('stop_reason', 'none'), - # code 任务：裸失败的种类（assertion / exception / no_code / timeout / ...）。 - # 这是本臂唯一能区分"方法错"与"格式崩塌"的字段，math 下恒为 None。 - 'base_kind': br.get('kind'), - 'withskill_acc_mean': sum(corr) / len(corr) if corr else 0.0, - # strict：unparseable 计 0（回退成 direct 时会被 baseline 掩护，此处 baseline=0 - # 所以两条曲线分叉纯粹反映格式崩塌） - 'withskill_acc_strict_mean': (sum(c * p for c, p in zip(corr, parses)) / len(corr) - if corr else 0.0), - 'withskill_pass_any': 1.0 if any(corr) else 0.0, - 'skill_parseable_mean': sum(parses) / len(parses) if parses else 0.0, - 'withskill_terminated_mean': sum(terms) / len(terms) if terms else 0.0, - 'skill': row[0][0], 'skill_parseable': bool(row[0][0]), 'skill_chars': len(row[0][0]), - 'withskill_pred': rolls[0]['pred'], 'withskill_correct': rolls[0]['correct'], - 'withskill_terminated': rolls[0]['terminated'], - 'withskill_stop_reason': rolls[0]['stop_reason'], 'withskill_text': rolls[0]['text']} - hard_recs.append(rec) - # rubric 缺失的 hard 题：不进 rescue 分母（与训练侧"缺 rubric 一律丢弃"一致），但必须以 - # 显式零进 acc：它们确实没被干预、baseline 也确实错了。分母漂移靠 hard_rubric_missing - # 可审计（rubric 成功后会永久进缓存，所以只会单调收敛到 0）。 - for (r, br), d in zip(hard, diags): - if not d: - recs.append({**head, 'data_id': r.get('data_id', ''), 'problem': r['problem'], - 'reference_answer': r['reference_answer'], 'baseline_pass': 0.0, - 'intervened': False, 'rubric_ok': False, - 'base_stop_reason': br.get('stop_reason', 'none'), - 'withskill_acc_mean': 0.0, 'withskill_acc_strict_mean': 0.0, - 'withskill_pass_any': 0.0, 'skill_parseable_mean': 0.0, - 'withskill_terminated_mean': 0.0}) - recs.extend(hard_recs) - - # ---- 4) 汇总 ---- - # acc 统一取"全部 eval 行的 withskill_acc_mean 均值"，与任何下游按行求均的脚本逐字一致。 - # 不能写成 base + (1-base)*rescue：那个式子隐含"缺 rubric 的题也按 rescue 率被救"， - # 在缺失不为零时会系统高估 acc（缺失=0 时两者相等）。 - n_all = len(eval_records) - n_hard = len(hard_recs) - base = (sum(1.0 for br in base_rolls if br['correct']) / n_all) if n_all else 0.0 - acc = (sum(x['withskill_acc_mean'] for x in recs) / n_all) if n_all else 0.0 - acc_strict = (sum(x['withskill_acc_strict_mean'] for x in recs) / n_all) if n_all else 0.0 - rescue = (sum(x['withskill_acc_mean'] for x in hard_recs) / n_hard) if n_hard else 0.0 - rescue_strict = (sum(x['withskill_acc_strict_mean'] for x in hard_recs) / n_hard) if n_hard else 0.0 - rescue_any = (sum(x['withskill_pass_any'] for x in hard_recs) / n_hard) if n_hard else 0.0 - fmt = (sum(x['skill_parseable_mean'] for x in hard_recs) / n_hard) if n_hard else 0.0 - term = (sum(x['withskill_terminated_mean'] for x in hard_recs) / n_hard) if n_hard else 0.0 - # 错题结构：'length' = 没写完（E16 实测占裸失败的 97.6%），其余 = 写完但答错。 - # 本臂能不能测到"方法修正"完全取决于后者不为零，所以逐次 eval 都上报。 - trunc = (sum(1.0 for x in hard_recs if x.get('base_stop_reason') == 'length') / n_hard - if n_hard else 0.0) - # code 任务：裸失败的种类分布。math 上截断率就够（97.6% 是没写完），代码域必须分开看 - # —— 只有 assertion/exception 这类才是 rubric 有客观证据可诊断的失败，no_code/timeout - # 是格式或环境问题。键名带 kind_ 前缀，随 summary 落盘（不进 swanlab 三条主指标）。 - kind_fracs = {} - if v2._TASK == 'code' and n_hard: - for kind in ('assertion', 'exception', 'import_or_syntax', 'no_code', 'no_entry', - 'timeout'): - kind_fracs[f'hard_base_kind_{kind}'] = ( - sum(1.0 for x in hard_recs if x.get('base_kind') == kind) / n_hard) - - summary = {'record_type': 'eval_summary', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, - 'protocol': 'reflexion', 'n': n_all, 'n_rollouts': R, - 'eval_skill_temperature': args.eval_skill_temperature, - 'baseline_acc_mean1': base, 'acc_mean1': acc, 'lift_mean1': acc - base, - 'acc_strict_mean1': acc_strict, 'lift_strict_mean1': acc_strict - base, - 'format_mean1': fmt, 'term_mean1': term, - # ★ 主指标 - 'hard_n': n_hard, 'hard_rescue_rate': rescue, - 'hard_rescue_strict_rate': rescue_strict, 'hard_rescue_pass_any': rescue_any, - 'hard_rescued': sum(x['withskill_acc_mean'] for x in hard_recs), - 'hard_rubric_missing': n_rubric_missing, - 'hard_base_trunc_frac': trunc, **kind_fracs} - # swanlab 只上报三条（2026-07-30 精简，命名不缩写）。口径：只用错题子集（baseline 做对 - # 的题不干预、不计入），所以 baseline_accuracy 恒为 0、with_skill_accuracy 就是救活率。 - # 其余读数（strict / pass_any / format / term / 混合 acc）仍全量在 summary 里落盘。 - # 不带 'eval/' 前缀：trainer.py 会统一加（f'eval/{k}'），写了会变成 eval/eval/xxx。 - metrics = {'baseline_accuracy': 0.0, - 'with_skill_accuracy': rescue, - 'lift': rescue} - if n_hard and n_hard < 60: - sys.stderr.write(f'[eval] WARNING: reflexion protocol has only {n_hard} hard problems; ' - f'SE(rescue) ~ {(0.25 / n_hard) ** 0.5:.3f} — raise --eval-size.\n') - return recs, summary, metrics diff --git a/cookbook/exp/skill2lora/skill_ablate/main.py b/cookbook/exp/skill2lora/skill_ablate/main.py deleted file mode 100644 index 2311059f3..000000000 --- a/cookbook/exp/skill2lora/skill_ablate/main.py +++ /dev/null @@ -1,327 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Entry point: run one ablation experiment by name (E1..E12) or explicit knobs. - -Usage: - python -m skill_ablate.main --exp E5 --seam-parquet-dir /root/data/seam \ - --output-dir output.ablate12/E5_rl_ab_off_pitfall - -Defaults mirror train_skill_v2._build_args so reused v2 primitives behave identically; only -the ablation-specific knobs are added (--exp / --max-updates / --eval-every-updates / ---improve-skill-temperature / --skill-char-limit / --pool-max / --rubric-global-dir). -""" -import argparse -import dataclasses -import sys - -import code_task -import train_skill_v2 as v2 - -from .config import METHODS, STYLES, TASKS, THINKINGS, get_spec, ExpSpec -from .trainer import run_experiment - - -def _build_args(argv=None): - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - # --- experiment selection: either --exp E5, or explicit --method/--thinking/--style --- - p.add_argument('--exp', default='', help='experiment name E1..E12 (fills method/think/style)') - p.add_argument('--method', choices=METHODS, default=None) - p.add_argument('--thinking', choices=THINKINGS, default=None) - p.add_argument('--executor-thinking', choices=THINKINGS, default=None, - help="executor(base_sampler) 的 thinking；默认取实验自己的 spec。'off' 是 " - 'E19/E20 的核心变量：think 的 executor 在 BigCodeBench 上 34-50% 的 ' - 'rollout 陷入字面死循环撞满预算，关掉后截断归零、裸解与 rubric 增量都更高。') - p.add_argument('--style', choices=STYLES, default=None) - - # --- data (mirror v2) --- - p.add_argument('--dataset', choices=('aops', 'math'), default='aops') - p.add_argument('--task', choices=TASKS, default=None, - help="task family; default = the experiment's own spec.task. 'code' switches " - 'the whole pipeline to BigCodeBench: executor / skill-gen / rubric prompts, ' - 'unit-test judging instead of \\boxed{} matching, and --bcb-parquet as the ' - 'data source (--deepmath-dir / --seam-parquet-dir / --dataset are ignored).') - p.add_argument('--bcb-parquet', default=code_task.DEFAULT_PARQUET, - help='BigCodeBench parquet (task=code only).') - p.add_argument('--test-workers', type=int, default=24, - help='task=code: thread pool for the unit-test subprocesses. Judging one ' - 'rollout starts a python subprocess (1-3s typical), and a chunk needs ' - 'hundreds of them — serial judging costs more wall clock than the GPU ' - 'rollouts themselves.') - p.add_argument('--test-timeout', type=int, default=60, - help='task=code: wall-clock cap per unit-test run (seconds).') - p.add_argument('--code-selftest', action=argparse.BooleanOptionalAction, default=True, - help='task=code: drop tasks whose OWN canonical solution fails their unit ' - 'tests (sandbox-undecidable, not a model failure; ~7.5% measured). ' - 'Result is cached next to the rubric cache and reused across arms.') - p.add_argument('--n', type=int, default=0) - p.add_argument('--exclude-data-ids', default='') - p.add_argument('--seed', type=int, default=42) - p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) - p.add_argument('--eval-size', type=int, default=128) - p.add_argument('--seam-parquet-dir', type=str, default='') - p.add_argument('--train-order-file', type=str, default='', - help='jsonl of {data_id,problem,reference_answer} in a FIXED training order; ' - 'replaces the train split and disables ProblemPool shuffling, so chunk k ' - 'is exactly batch k. Used to pin twinkle to SEAM\'s realized batch ' - 'sequence (verl shuffles its dataloader, so same pool != same batches).') - p.add_argument('--deepmath-dir', type=str, default='', - help='DeepMath-103K parquet dir; when set, overrides --seam-parquet-dir/--dataset ' - 'and uses the difficulty-stratified split (eval/train same level mix).') - p.add_argument('--min-level', type=int, default=0, - help='train-only difficulty floor for DeepMath (eval keeps full-level mix). ' - '0 = off. E1/E5 audit: level<=5 is all-pass dominated (zero gradient); ' - 'recommended 6.') - p.add_argument('--eval-min-level', type=int, default=0, - help='difficulty floor for the EVAL split too (0 = off, full-level mix as in ' - 'E1-E16). Set = --min-level for arms whose readout only carries ' - 'information on problems the bare executor fails (E17): the full mix is ' - 'only ~26%% bare-wrong vs ~43%% at level>=6, so the same eval budget buys ' - '1.6x the effective sample AND matches the train distribution. Changes ' - 'the eval set -> not comparable to arms run with 0.') - - # --- eval口径 (4 rollouts × T=0.5, 与旧臂 E1-E13 同口径, 2026-07-28 拍板回退) --- - # 曾短暂改为 SEAM val 口径(1×greedy)，为保持与已完成 6 臂可横比而回退；需要 SEAM 口径时 - # 显式传 --eval-rollouts 1 --eval-skill-temperature 0.0。 - p.add_argument('--eval-rollouts', type=int, default=4) - p.add_argument('--eval-skill-temperature', type=float, default=0.5) - - # --- skill-gen / rollout (mirror v2) --- - p.add_argument('--chunk-size', type=int, default=16) - p.add_argument('--n-skills', type=int, default=8) - p.add_argument('--skill-gen-temperature', type=float, default=1.0) - p.add_argument('--skill-gen-top-p', type=float, default=1.0) - p.add_argument('--skill-gen-top-k', type=int, default=-1) - p.add_argument('--max-model-len', type=int, default=16384) - p.add_argument('--max-tokens', type=int, default=8192) - p.add_argument('--skill-max-tokens', type=int, default=None, - help='default per-experiment: 8192 (think) / 4096 (nothink); an explicit ' - 'value here wins over the experiment default.') - # E14+ 稠密 reward：logp_rl 用 reward_rollouts/temperature 做 executor 采样，找正确伪 GT S； - # 之后不再 rollout 判 reward，而是算 Δ mean logP_executor(S | problem + skill)。显式传值覆盖 spec。 - p.add_argument('--reward-rollouts', type=int, default=None) - p.add_argument('--reward-temperature', type=float, default=None) - # leak 一律不进 reward（用户既定要求，反复强调）：该指标假阳性极高——DeepMath 的 gold 约一半是 - # 单字符，raw leak 率 0.53-0.65 而加 >=2 字符门后只剩 0.05-0.12，5-10 倍虚高；把它当 reward 项会 - # 用噪声主导组内排序，并且 -1.0 量级的离群值会抬高组 std、连带压小其余候选的 advantage。 - # leak 只保留监控口径（_leak_split -> leak/correct_rate、leak/wrong_rate）。默认 0 = 关闭。 - p.add_argument('--logp-leak-penalty', type=float, default=0.0, - help='E14/E15 logp reward leak penalty. DEFAULT 0 = OFF: leak is a ' - 'monitoring-only metric by project rule, never a reward term. The ' - 'unparseable-skill floor stays fixed at -1.0 regardless.') - # --- E16 passrate_hinge reward knobs ------------------------------------------------- - # 标定依据：.tmp_analysis/reward_shape_calib.py（E4 11155 条 rollout + DeepMath r1_solution_1）。 - # 关键发现：本数据集上长度本身无害——P(对|token) 在 5500 以下平在 0.96-0.98，5500-7500 微降到 - # 0.90，7500 以后断崖到 0.226（每个 difficulty 层内同形）；且 GT 参考解比模型正确答案更长 - # （p50 4377 vs 3444，GT p90 9816 已超 8192 预算）。所以“越短越好”在这里是错的，必须给死区。 - p.add_argument('--reward-trunc-penalty', type=float, default=0.12, - help='E16 alpha_len: per-rollout length penalty weight. eff *= (1 - this * ' - 'len_pen), len_pen = ((tok - lo) / (budget - lo)) ** pow, 0 below lo. ' - 'Sized so the TOTAL deduction (1+kappa)*(1-eff) stays under one pass_rate ' - 'quantum (1/M): the length signal may break ties but must never override ' - 'pass_rate, which is what eval actually measures. Its absolute size barely ' - 'matters anyway — inside an all-fail group A=(R-mean)/std rescales the ' - 'spread to unit size, so the CURVE SHAPE carries the information.') - p.add_argument('--reward-trunc-lo', type=int, default=5500, - help='E16 length dead zone: rollouts under this many executor tokens are not ' - 'penalized at all. Calibrated: P(correct|tok) is flat 0.96-0.98 below 5500.') - p.add_argument('--reward-len-pow', type=float, default=2.0, - help='E16 length ramp exponent (>1 = convex, marginal penalty grows with ' - 'length, concentrating it in the last ~1300 tokens before the budget).') - p.add_argument('--reward-loop-penalty', type=float, default=0.04, - help='E16 beta_loop: self-revision marker penalty weight. Deliberately small ' - '— inside a fixed token band ~85%% of the raw marker effect is just the ' - 'mechanical "longer output has more markers" correlation, and the same ' - 'one-pass-quantum budget is shared with the length term.') - p.add_argument('--reward-loop-lo', type=float, default=2.0, - help='E16 marker density (per 1k tokens) below which no loop penalty applies.') - p.add_argument('--reward-loop-hi', type=float, default=9.0, - help='E16 marker density at which the loop penalty saturates at 1.0.') - p.add_argument('--reward-ineff-kappa', type=float, default=0.10, - help='E16 kappa: reward -= this * mean(1 - eff). Keeps FAILING candidates ' - 'separable (a pure product collapses them all to 0 = E4 pathology). Small ' - 'on purpose: in an all-fail group A=(R-mean)/std rescales the spread back ' - 'to unit size anyway, and a large value would reverse pass_rate ordering.') - p.add_argument('--reward-leak-gate', type=float, default=0.0, - help='E16 reward leak penalty. DEFAULT 0 = OFF: leak is a monitoring-only ' - 'metric by project rule, never a reward term (see --logp-leak-penalty).') - p.add_argument('--base-tok-floor', type=int, default=5000, - help='E16 problem filter: only train problems whose baseline (no-skill) greedy ' - 'output exceeds this many tokens (probe: base_tok vs skill lift +0.62, the ' - 'strongest problem-side signal). 0 disables the filter.') - # --- E17 reflexion 臂 ----------------------------------------------------------------- - p.add_argument('--reflexion-k', type=int, default=24, - help='E17 batch alignment: train on EXACTLY this many bare-wrong problems per ' - 'chunk. The chunk is drawn at --chunk-size, bare-solved, and wrong ones ' - 'are taken (with rubric backfill) until K is reached; dynamic ' - 'over-drawing is impossible because resume replays fixed-size pool ' - 'draws. Group count per update must be constant or the step size, the ' - 'noise floor and the online SNR readout all move with the draw. ' - 'DEFAULT 24 matches the 23.46 groups/update E16 actually ran (1173 ' - 'total): cumulative evidence scales as sqrt(N) and E16 only reached ' - '1.9 sigma on any text-level direction, so a smaller K has no ' - 'discriminating power left. Needs --chunk-size ~5x (bare error rate ' - '0.329 measured over the FULL level>=6 pool, 527/1600).') - p.add_argument('--align-mode', choices=('v2', 'seam'), default='v2') - # --- E18 rejection_sft 臂 --------------------------------------------------------------- - p.add_argument('--e18-accumulate', type=int, default=16, - help='E18: fire one SFT update only after this many accepted (rejection-' - 'sampled) skills have accumulated in the pool. Winners are also ' - 'appended to <output-dir>/e18_sft_dataset.jsonl for offline reuse. ' - 'DEFAULT 16 (2026-07-30 拍板) = one sft_batch_size, so a chunk of 32 ' - '(bare error rate 0.329 -> ~10 accepted) fires roughly every other ' - 'chunk and --max-updates 50 is reachable; 128 would need ~15 chunks ' - 'per update.') - p.add_argument('--len-budget', type=int, default=None, - help='regen skill length target (chars). Default per style (ablation stats ' - '#9): narrative~1100 / pitfall~300. Used to pick the regen survivor.') - - # --- rubric / regen / distill (mirror v2 + new) --- - p.add_argument('--passatk-k', type=int, default=8) - p.add_argument('--passatk-skill-temp', type=float, default=1.0) - p.add_argument('--passatk-skill-top-p', type=float, default=1.0) - p.add_argument('--passatk-m', type=int, default=2) - p.add_argument('--rubric-workers', type=int, default=16) - p.add_argument('--improve-skill-temperature', type=float, default=0.5, - help='temperature for the single first-pass skill in opsd / improve_sft.') - p.add_argument('--skill-char-limit', type=int, default=4096, - help='hard char cap for SFT-seed skills (skill_quality_analysis.md #15-1/#18).') - - # --- GRPO / optim (mirror v2) --- - p.add_argument('--sft-batch-size', type=int, default=16, - help='batch = TRAIN_DP multiple; also the SFT-pool draw size.') - p.add_argument('--train-micro-batch', type=int, default=0, - help='forward_backward micro size; 0 = follow --sft-batch-size. think/8192 ' - 'experiments auto-halve to 8 (fp32 master + 8k-token logits OOM guard); ' - 'gradient is micro-normalized so this is mathematically equivalent.') - p.add_argument('--ppo-mini-batch-size', type=int, default=0) - p.add_argument('--grpo-epsilon', type=float, default=0.2) - p.add_argument('--adv-clip', type=float, default=0.0) - # 对初始策略的锚。漂移分析（.tmp_analysis/why_no_correction.py 等）显示：“executor 能干净收束” - # 是初始 skill 分布自带的脆弱属性，reward 里能反对它被磨耗的可迁移成分只有 11%，因此 - # 锚本身就是一个直接对症的旋钮（选择压力只能在组内排序，锚才能提供恢复力）。 - # 2026-07-29 拍板：默认 0.001 -> 0.01。⚠️ 注意锚是无方向的刹车，它同等抵制那个方向 - # 正确的 +0.078 收束签名比较信号；取值未经标定（既往臂全部跑在 0.001，无可用对照）。 - p.add_argument('--kl-beta', type=float, default=0.01, - help='KL anchor to the reference (initial, never-synced) policy. Raised from ' - '0.001 to 0.01 on 2026-07-29 to oppose the intrinsic drift that erodes ' - 'executor termination. Recorded in the config fingerprint. Note: the ' - 'anchor is undirected -- it also brakes the (weak) useful gradient.') - p.add_argument('--run-tag', default='', - help='suffix for the swanlab experiment name, to tell apart variants that share ' - 'the same ExpSpec (e.g. RUN_TAG=kl01 for the --kl-beta 0.01 arm).') - p.add_argument('--lr', type=float, default=1e-6, - help='stable 1e-6, no warmup / no decay (ablation spec).') - p.add_argument('--sft-weight', type=float, default=1.0, - help='advantage magnitude for SFT samples (ablation spec: 1.0).') - p.add_argument('--drop-zero-adv', action='store_true', - help='reserved single-point ablation (接口方案 #10): drop zero-advantage ' - 'candidates from RL training batches instead of keeping them in the ' - 'token-mean denominator (SEAM口径). Default off; NOT part of E1-E12.') - - # --- run control (new) --- - p.add_argument('--max-updates', type=int, default=50, - help='stop after this many PARAMETER UPDATES (the "step" unit).') - p.add_argument('--eval-every-updates', type=int, default=5) - p.add_argument('--save-every-updates', type=int, default=0, - help='save a weights-only checkpoint (<exp>-u<N>) every N updates; 0 = only ' - 'the final save. lr is constant so optimizer state is NOT saved.') - p.add_argument('--resume-from', default='', - help='checkpoint dir name under --output-dir (e.g. E1-final / E1-u50) or an ' - 'absolute path. Loads weights into skill_model, restores updates/chunk ' - 'position from train_state.json (falls back to DONE.json), appends to ' - 'the record files, and bypasses the DONE.json skip guard. Raise ' - '--max-updates beyond the restored count to actually continue.') - p.add_argument('--pool-max', type=int, default=2048, - help='SFT pool per-queue cap (drop oldest; bounds majority backlog).') - p.add_argument('--rubric-global-dir', default='', - help="dir for the cross-experiment global rubric cache " - "(default: parent of --output-dir).") - - # --- output / logging --- - p.add_argument('--output-dir', default='./output.ablate12/exp') - p.add_argument('--no-cache', action='store_true') - p.add_argument('--force', action='store_true', - help='rerun even if <output-dir>/DONE.json marks this experiment complete.') - p.add_argument('--swanlab-project', default='twinkle') - - args = p.parse_args(argv) - - # resolve the experiment spec - if args.exp: - spec = get_spec(args.exp) - else: - if not (args.method and args.thinking and args.style): - p.error('provide --exp E5, or all of --method/--thinking/--style') - spec = ExpSpec(name='Ex', method=args.method, thinking=args.thinking, style=args.style) - # 显式 --task 覆盖 spec（ExpSpec 是 frozen dataclass）；用于临时把某个数学臂放到 code 上跑。 - if args.task and args.task != spec.task: - sys.stderr.write(f'[ablate] WARNING: --task {args.task} overrides {spec.name} ' - f'spec.task={spec.task}\n') - spec = dataclasses.replace(spec, task=args.task) - args.task = spec.task - # executor thinking 同理（frozen dataclass -> replace）。v2.build_* 读 args.executor_thinking。 - if args.executor_thinking and args.executor_thinking != spec.executor_thinking: - sys.stderr.write(f'[ablate] WARNING: --executor-thinking {args.executor_thinking} ' - f'overrides {spec.name} spec.executor_thinking=' - f'{spec.executor_thinking}\n') - spec = dataclasses.replace(spec, executor_thinking=args.executor_thinking) - args.executor_thinking = spec.executor_thinking - - # sanity: Ray dp rule — sft/eval batch must divide TRAIN_DP - if args.sft_batch_size % v2.TRAIN_DP != 0: - p.error(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple of TRAIN_DP ({v2.TRAIN_DP})') - if args.train_micro_batch and args.train_micro_batch % v2.TRAIN_DP != 0: - p.error(f'--train-micro-batch ({args.train_micro_batch}) must be a multiple of TRAIN_DP ({v2.TRAIN_DP})') - if args.chunk_size < 1: - p.error('--chunk-size must be >= 1') - if spec.method == 'reflexion': - if args.reflexion_k < 1: - p.error('--reflexion-k must be >= 1') - if args.reflexion_k > args.chunk_size: - p.error(f'--reflexion-k ({args.reflexion_k}) > --chunk-size ({args.chunk_size}): the ' - f'aligned batch can never be filled') - # 2-sigma 余量检查：错题数 ~ Binom(chunk, p_wrong)。余量不够时组数会随 chunk 抽样抖， - # 而组数恒定是本臂的硬要求。⭐ math: 0.329 = E16 全量题池实测（527/1600）；切勿用 - # 0.446，那是 base_tok>5000 筛选后子集的错误率（偏难），而本臂 floor=0。 - # ⭐ code: 0.56 —— 不是 probe 的 0.715。probe 的 pass 0.285 是 nothink + max_tokens 4096 - # 的读数；本臂 executor 是 think=on + 8192，同一批题实测裸错率只有 18/32=0.5625 - # （2026-07-31 dry run：训练侧 10/16、eval 侧 8/16）。用 0.715 会把 chunk 需求算小 - # 一半：P(48 道里凑不满 24) 在 0.715 下是 0.000，在 0.5625 下是 0.154。 - # ⭐ executor nothink（E19/E20）另算：code nothink 探针实测 pass 0.378 -> p_wrong 0.622； - # math nothink 实测 0.31（首个 E19 run 的 c0：64 道里 20 道错，baseline acc≈0.69）， - # 与 think 的 0.329 基本相同 —— 曾按 0.85 保守估是错的，会把 chunk 需求算小一半。 - if spec.executor_thinking == 'off': - _p_wrong = 0.622 if spec.task == 'code' else 0.31 - else: - _p_wrong = 0.5625 if spec.task == 'code' else 0.329 - _mu = args.chunk_size * _p_wrong - _sd = (args.chunk_size * _p_wrong * (1 - _p_wrong)) ** 0.5 - if args.reflexion_k > _mu - 2 * _sd: - sys.stderr.write( - f'[ablate] WARNING: --chunk-size {args.chunk_size} gives {_mu:.1f}+-{_sd:.1f} ' - f'wrong problems, less than 2 sigma of headroom over --reflexion-k ' - f'{args.reflexion_k}; expect signal/k_short > 0 and a drifting group count. ' - f'Use --chunk-size >= {int((args.reflexion_k + 2 * _sd) / _p_wrong) + 1}.\n') - if spec.method == 'rejection_sft': - if args.e18_accumulate < 1: - p.error('--e18-accumulate must be >= 1') - # 池 batch 整体交给 _train_batch，drop_last 到 TRAIN_DP 倍数会静默丢尾部真样本； - # 强制倍数关系把丢样本量钉在 0。 - if args.e18_accumulate % v2.TRAIN_DP != 0: - p.error(f'--e18-accumulate ({args.e18_accumulate}) must be a multiple of ' - f'TRAIN_DP ({v2.TRAIN_DP})') - args.rubric_global_dir = args.rubric_global_dir or None - return args, spec - - -def main(argv=None): - args, spec = _build_args(argv) - sys.stderr.write(f'[ablate] running {spec.name}: task={spec.task} view={spec.view} ' - f'method={spec.method} thinking={spec.thinking} style={spec.style} ' - f'executor_thinking={spec.executor_thinking} ' - f'loss={spec.loss} smt={spec.skill_max_tokens} -> {args.output_dir}\n') - run_experiment(args, spec) - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/skill2lora/skill_ablate/methods.py b/cookbook/exp/skill2lora/skill_ablate/methods.py deleted file mode 100644 index 2e1128f88..000000000 --- a/cookbook/exp/skill2lora/skill_ablate/methods.py +++ /dev/null @@ -1,1833 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Pluggable training methods (E1-E12) over a shared context. - -Each method implements TrainMethod: ``step(chunk, ci) -> dict`` where the dict carries at -least ``n_updates`` (parameter-update count, the "step" unit per skill_quality_analysis.md -#2), plus ``metrics`` (scalars to log) and ``gen_records`` (per-problem rollout audit rows -for gen_records.jsonl). New methods / losses are added by writing one class + one -METHOD_REGISTRY entry — the trainer never changes. - -Reuse map (all imported from train_skill_v2 / rollouting, never edited): -- bnpo (view B): v2 ``process_chunk`` + ``_train_step`` verbatim (query-only). -- rl_ab / rl_err (view A): bare greedy solve -> wrong=A(query+rubric) / right=B(query-only), - 8 skills each -> executor greedy reward -> group advantage -> BNPO on the rubric trajectory - (train-with-rubric); rl_err drops the B line from training. Rubric API calls run on threads - WHILE the B line rolls out on GPU (API/GPU overlap). -- opsd (view A): error problems, 1 student skill (query-only, T=0.5); teacher forward - (student prompt + rubric appended to the SYSTEM prompt, same response) -> per-token OPSD - KL (loss='opsd'). Teacher logps are extracted RESPONSE-ONLY via a client-side template - encode (teacher/student prompt lengths differ, so the full-sequence form would misalign). -- improve_sft (view A): first-pass 1 skill (query-only, T=0.5); correct -> positive SFT seed - (no leak, <=4096 chars); parseable-but-wrong -> rubric regen (2-in-8 pick 1) -> negative - SFT seed; unparseable first pass is SKIPPED (no trajectory to diagnose); balanced 1:1 pool - (majority side down-sampled per chunk, never backlogged) -> SFT. -- sft (view A): bare wrong -> rubric -> regen (query+rubric) 2-in-8 -> plain pool -> SFT. - -Training-batch helper ``_train_batch`` mirrors v2 ``_train_step`` exactly (empty-response -filter, drop_last to TRAIN_DP, micro-batch by sft_batch_size, ppo_mini_batch_size multi-step -with pre-computed ref/old logps, clip+step per mini, ckpt sync, calculate_metric) but takes a -swappable trajectory builder (query-only vs query+rubric) and an optional teacher builder for -OPSD. In the OPSD path no ref forward is done at all: OPSDLoss uses only teacher_logps -(kl_beta / ref_logps play no role there). -""" -import json -import os -import re -import sys -from collections import Counter -from concurrent.futures import ThreadPoolExecutor -from dataclasses import dataclass, field -from typing import Any, Callable, Dict, List, Optional, Tuple - -import numpy as np -from twinkle.data_format import pack_user_data - -import train_skill_v2 as v2 -from train_skill_v2 import ( - TRAIN_DP, - _answer_leaked, - _assign_advantages, - _clean_text, - _empty_roll, - _extract_skill, - _parse_seq, - _regen_prompt, - _run_samples, - _skill_reward, - _skillgen_prompt, - _train_step, - build_direct_prompt, - build_skill_solve_prompt, - process_chunk, -) - -from .pool import NEG, POS, SamplePool -from .rollouting import ( - opsd_teacher_trajectory, - query_only_train_trajectory, - rubric_skillgen_prompt, - rubric_train_trajectory, -) - - -@dataclass -class MethodContext: - """Everything a method needs; assembled once by the trainer.""" - skill_model: Any - ref_model: Any - skill_sampler: Any - base_sampler: Any - ckpt: Any - skill_dp: int - base_dp: int - args: Any - checker: Any = None - rubric_cache: Any = None # GlobalRubricCache (RL/SFT) or LocalRubricCache (improve/opsd) - pool: Optional[SamplePool] = None # SFT-family accumulator (None for RL/OPSD) - encode_template: Any = None # client-side Template clone (OPSD teacher alignment only) - extra: Dict[str, Any] = field(default_factory=dict) - - -# =========================================================================================== -# shared low-level helpers (reuse v2 primitives; only orchestration is new) -# =========================================================================================== -def _bare_solve(ctx: MethodContext, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: - """Bare-problem greedy (T=0) executor solve; returns one roll per record (order-aligned).""" - out = _run_samples(ctx.base_sampler, [build_direct_prompt(r['problem']) for r in records], - 1, ctx.args.max_tokens, ctx.base_dp, temperature=0.0) - return v2._parse_many([(v2._first_seq(seqs), r['reference_answer']) - for r, seqs in zip(records, out)]) - - -def _rubric_entry(record: Dict[str, Any], roll: Dict[str, Any]) -> Dict[str, Any]: - """Build the entry _diagnose_entry expects from a failure trajectory. - - code 任务：fail_segment 不是"输出全文"，而是**提交的代码 + 单测真实报错**（异常类型 / - 断言差异 / 失败用例名）。这是三个数据集横比下 rubric 唯一真正产生增量的原因 —— 数学与 - BFCL 上 judge 手里没有任何客观证据，只能猜，命中率≈随机。 - """ - if v2._TASK == 'code': - return {'problem': record['problem'], 'reference_answer': record['reference_answer'], - 'data_id': record.get('data_id', ''), - 'fail_segment': v2.code_task.diag_segment(roll), - 'fail_stop_reason': roll.get('stop_reason', 'none')} - return {'problem': record['problem'], 'reference_answer': record['reference_answer'], - 'data_id': record.get('data_id', ''), - 'fail_segment': roll.get('text', ''), - 'fail_stop_reason': roll.get('stop_reason', 'none')} - - -def _diagnose_parallel(ctx: MethodContext, - jobs: List[Tuple[Dict[str, Any], Optional[str]]]) -> List[str]: - """Run rubric diagnoses in parallel threads (pure API, no GPU; DiskCache.put is locked). - - ``jobs`` is a list of (entry, skill_or_None); returns diagnoses aligned to jobs - ('' on cache-off / API error). Parallelism = --rubric-workers (was serial before).""" - if not jobs or ctx.rubric_cache is None: - return [''] * len(jobs) - workers = max(1, min(ctx.args.rubric_workers, len(jobs))) - with ThreadPoolExecutor(max_workers=workers) as ex: - return list(ex.map( - lambda j: ctx.rubric_cache.get_or_diagnose(j[0], ctx.checker, skill=j[1]) or '', - jobs)) - - -def _skillgen_solve(ctx: MethodContext, items: List[Dict[str, Any]], n_skills: int, - temperature: float) -> None: - """For each item {record, prompt}: sample n skills, greedy-solve each, attach a - ``_cands`` list (v2 shape) onto the record so ``_assign_advantages`` can be reused.""" - args = ctx.args - sg_out = _run_samples(ctx.skill_sampler, [it['prompt'] for it in items], n_skills, - args.skill_max_tokens, ctx.skill_dp, - temperature=temperature, top_p=args.skill_gen_top_p, - top_k=args.skill_gen_top_k, logprobs=1) - flat = [] - for it, seqs in zip(items, sg_out): - it['record']['_cands'] = [] - for s in seqs: - resp = _clean_text(getattr(s, 'decoded', '') or '') - block = _extract_skill(resp) or '' - # tokens = 采样端真实吐出的 token id；训练样本只能用它拼（见 v2.build_train_feature）。 - # logprobs 同长，只给 GRPOMetric 做采样/训练对账，不进 loss。 - _toks = [int(t) for t in (getattr(s, 'tokens', None) or [])] - cand = {'skills': block, 'response': resp, 'parseable': bool(block), - 'leaked': None, 'with_pass': None, 'reward': None, 'rolls': [], - 'advantage': 0.0, 'kept': False, 'tokens': _toks, - 'logprobs': v2.sampler_logprobs(s), - 'skillgen_stop': getattr(s, 'stop_reason', None), - 'skillgen_tokens': len(_toks)} - it['record']['_cands'].append(cand) - if block: - flat.append((it, cand)) - for it, c in flat: - c['leaked'] = _answer_leaked(c['skills'], it['record']['reference_answer']) - if flat: - # V2 fix: pass the raw skill-gen response like v2 process_chunk does — seam align - # nests the actor's full response_text into the executor prompt (no-op in v2 mode). - ws = _run_samples(ctx.base_sampler, - [build_skill_solve_prompt(it['record']['problem'], c['skills'], c.get('response')) - for it, c in flat], - 1, args.max_tokens, ctx.base_dp, temperature=0.0) - judged = v2._parse_many([(v2._first_seq(seqs), it['record']['reference_answer']) - for (it, _c), seqs in zip(flat, ws)]) - for (it, c), roll in zip(flat, judged): - c['rolls'] = [roll] - c['with_pass'] = 1.0 if roll['correct'] else 0.0 - c['reward'] = _skill_reward(c['parseable'], roll['correct']) - for it in items: - for c in it['record']['_cands']: - if c['reward'] is None: - c['reward'] = 0.0 - - -def _grpo_records(records: List[Dict[str, Any]], with_rubric: bool) -> List[Dict[str, Any]]: - """Flatten per-problem _cands into GRPO train records (v2 shape). - ``with_rubric`` tags each record so the trajectory builder knows which prompt to rebuild. - - ``tokens`` 必须带上：它是训练样本的唯一来源（response 只留给 reward/审计），少了它 - trajectory builder 就会回退到 decode->重编码那条有偏差的路。``logprobs`` 同长，给 - GRPOMetric 做采样/训练对账。""" - recs = [] - for r in records: - for c in r.get('_cands', []): - if c.get('reward') is None: - continue - recs.append({'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'data_id': r.get('data_id', ''), 'response': c['response'], - 'skills': c['skills'], 'advantage': c['advantage'], - 'tokens': c.get('tokens') or [], - 'logprobs': c.get('logprobs') or [], - 'kept': c['kept'], 'reward': c['reward'], 'rubric': r.get('_rubric', ''), - 'with_rubric': with_rubric, 'sft': False}) - return recs - - -def _leak_blocks(skill: str, reference) -> bool: - """Filtering-grade leak gate (bugfix #4): only answers >=2 chars are informative — the - same gate E14/E16 already apply to their reward penalty. Single-char golds ('2' etc.) - substring-match ordinary math prose (~84% false positives measured), which starved the - SFT-family pools by discarding nearly every regen candidate. Monitoring paths keep the - raw ``_answer_leaked`` so the recorded leak/rate 口径 is unchanged.""" - return len(str(reference).strip()) >= 2 and _answer_leaked(skill, reference) - - -def _leak_split(pairs: List[Tuple[bool, bool]]) -> Dict[str, float]: - """#10 monitoring curves: leaked&correct vs leaked&wrong rates over parseable skills. - "泄露正确答案可接受"不等于"泄露无害"——有害的是错误数值注入，两条曲线拆开监控。""" - n = len(pairs) - if not n: - return {'leak/correct_rate': 0.0, 'leak/wrong_rate': 0.0} - return {'leak/correct_rate': sum(1 for lk, ok in pairs if lk and ok) / n, - 'leak/wrong_rate': sum(1 for lk, ok in pairs if lk and not ok) / n} - - -def _cand_leak_pairs(records: List[Dict[str, Any]]) -> List[Tuple[bool, bool]]: - # bugfix #7: with_pass is a float pass RATE under M>1 rollouts (E16) — bool(0.25) would - # count a partial pass as "correct"; compare > 0 instead (identical for greedy 0/1 arms). - return [(bool(c['leaked']), (c['with_pass'] or 0) > 0) - for r in records for c in r.get('_cands', []) - if c.get('parseable') and c.get('with_pass') is not None] - - -def _cand_pass_metrics(records: List[Dict[str, Any]]) -> Dict[str, float]: - """Mean-family train metrics, same family as eval acc_mean1 (ws_acc is pass@8-inflated): - candidate_pass = P(correct | parseable); clean_pass = P(correct | parseable & terminated). - Sharper channel split by ANSWER AVAILABILITY (truncated rolls still count correct when a - balanced \\boxed{} landed before the budget — 13% of E1 truncations did): - answered_rate = P(pred emitted) and answered_pass = P(correct | answered), the content-only - channel (E1 vs E5 core comparison curve); plus parse/trunc rates for the format channel.""" - cands = [c for r in records for c in r.get('_cands', [])] - if not cands: - return {} - m = {'skill/parse_rate': sum(1 for c in cands if c.get('parseable')) / len(cands)} - scored = [c for c in cands - if c.get('parseable') and c.get('with_pass') is not None and c.get('rolls')] - if scored: - m['acc/candidate_pass'] = sum(c['with_pass'] for c in scored) / len(scored) - clean = [c for c in scored if c['rolls'][0].get('stop_reason') != 'length'] - m['term/withskill_trunc_frac'] = 1.0 - len(clean) / len(scored) - if clean: - m['acc/clean_pass'] = sum(c['with_pass'] for c in clean) / len(clean) - answered = [c for c in scored if c['rolls'][0].get('pred') not in (None, '')] - m['term/answered_rate'] = len(answered) / len(scored) - if answered: - m['acc/answered_pass'] = sum(c['with_pass'] for c in answered) / len(answered) - return m - - -def _train_metrics(metric: Optional[Dict[str, Any]]) -> Dict[str, float]: - """calculate_metric -> swan-ready train/* scalars (same key handling as v2 _swan_metrics).""" - d = {} - for k, val in (metric or {}).items(): - if not v2._is_num(val): - continue - if k.startswith('learning rate'): - if 'group 1' in k: - d['train/lr'] = float(val) - elif k.startswith('train/'): - # GRPOMetric 等已经自带 train/ 前缀，不能再套一层。 - d[k.replace(' ', '_')] = float(val) - else: - d[f'train/{k.replace(" ", "_")}'] = float(val) - return d - - -# =========================================================================================== -# OPSD teacher alignment: client-side encode -> response-only teacher logps -# =========================================================================================== -def _encode_for_align(tmpl, built): - """``traj_fn``/``teacher_fn`` 现在可能直接返回成品 InputFeature（token 直通路径）， - 那就无需再编码；只有回退的 messages 形式才走 tmpl.encode。""" - return built if 'input_ids' in built else tmpl.encode(built) - - -def _align_teacher(ctx: MethodContext, samples, traj_fn, teacher_fn): - """Encode student & teacher trajectories with a CLIENT-SIDE clone of the remote template - and keep only samples whose response-token counts match (they always should — same - assistant text — but max-length 'delete' truncation or tokenizer drift must not silently - misalign the distillation). Returns (kept_samples, teacher_label_positions, n_dropped). - - teacher_label_positions[i] are the (rolled-)label indices of sample i inside its OWN - unpadded teacher sequence; right padding keeps these indices valid in the padded batch, - so they can slice the teacher's full-sequence logps down to the response-only form that - OPSDLoss requires (full-sequence form would misalign: prompts differ in length). - """ - tmpl = ctx.encode_template - assert tmpl is not None, 'OPSD needs ctx.encode_template (built by the trainer)' - keep, pos_lists, dropped = [], [], 0 - for s in samples: - st = _encode_for_align(tmpl, traj_fn(s)) - tt = _encode_for_align(tmpl, teacher_fn(s)) - if st is None or tt is None: # deleted by max-length truncation - dropped += 1 - continue - spos = np.where(np.asarray(st.get('labels')) != -100)[0] - tpos = np.where(np.asarray(tt.get('labels')) != -100)[0] - if len(spos) != len(tpos) or len(tpos) == 0: - dropped += 1 - continue - keep.append(s) - pos_lists.append(tpos) - return keep, pos_lists, dropped - - -def _gather_response_logps(full_logps, pos_lists) -> List[List[float]]: - """Slice full-sequence [B, S] teacher logps down to per-sample response-only lists.""" - rows = [] - for i, pos in enumerate(pos_lists): - row = full_logps[i] - row = row.tolist() if hasattr(row, 'tolist') else list(row) - assert len(pos) == 0 or int(pos[-1]) < len(row), \ - f'teacher logps row {i} shorter than label positions ({len(row)} <= {int(pos[-1])})' - rows.append([float(row[int(p)]) for p in pos]) - return rows - - -# =========================================================================================== -# E14+ helpers: executor pseudo-GT + dense logP reward -# =========================================================================================== -def _executor_answer_trajectory(problem: str, skill: str, answer_text: str, - raw_response: Optional[str] = None, - answer_tokens: Optional[List[int]] = None, - template=None) -> Dict[str, Any]: - """Teacher-forcing sample for executor logP(S | problem + skill). - - S 是 executor 自己采出来的伪 GT（E14），所以带了 ``answer_tokens`` 时直接拼采样 token： - 这里的 prompt 每次都不同（换 skill）但被打分的 token 必须是同一串，正是 - ``build_train_feature`` 的场景。E15 的 S 是 DeepMath 的外部 R1 参考解，本来就不是模型 - 产出、没有对应 token，只能走 messages 编码。 - """ - msgs = [dict(m) for m in build_skill_solve_prompt(problem, skill, raw_response)['messages']] - if answer_tokens: - return v2.build_train_feature(msgs, answer_tokens, template=template) - return {'messages': msgs + [{'role': 'assistant', 'content': answer_text}], - 'user_data': pack_user_data({'key_rounds': [len(msgs)]})} - - -def _set_ref_executor_template(ctx: MethodContext, enable_thinking: bool) -> None: - """Temporarily reuse ref_model as frozen executor scorer, then restore skill/ref layout.""" - ctx.ref_model.set_template(v2.Template, model_id=v2.MODEL_ID, - enable_thinking=enable_thinking, - max_length=ctx.args.max_model_len, - truncation_strategy='delete') - - -def _mean_logp_rows(full_logps, pos_lists) -> List[float]: - rows = _gather_response_logps(full_logps, pos_lists) - return [(sum(x) / len(x)) if x else float('-inf') for x in rows] - - -def _score_executor_mean_logps(ctx: MethodContext, trajs: List[Dict[str, Any]]) -> List[Optional[float]]: - """Return mean response-token logP under the frozen executor template; None means truncated.""" - tmpl = ctx.encode_template - assert tmpl is not None, 'logp_rl needs ctx.encode_template' - out: List[Optional[float]] = [None] * len(trajs) - valid_trajs, pos_lists, valid_idx = [], [], [] - for i, tr in enumerate(trajs): - enc = _encode_for_align(tmpl, tr) - if enc is None: - continue - pos = np.where(np.asarray(enc.get('labels')) != -100)[0] - if len(pos) == 0: - continue - valid_trajs.append(tr) - pos_lists.append(pos) - valid_idx.append(i) - if not valid_trajs: - return out - sft = getattr(ctx.args, 'train_micro_batch', 0) or ctx.args.sft_batch_size - dp = max(1, v2.REF_DP) - _set_ref_executor_template(ctx, enable_thinking=True) - try: - for st in range(0, len(valid_trajs), sft): - mb = valid_trajs[st:st + sft] - mb_pos = pos_lists[st:st + sft] - n_mb = len(mb) - # forward_only 按 slice_dp 切分，非 dp 整倍数的尾批用末尾样本补齐，输出只取前 n_mb 行 - if n_mb % dp: - pad = dp - (n_mb % dp) - mb = mb + [mb[-1]] * pad - mb_pos = list(mb_pos) + [mb_pos[-1]] * pad - vals = _mean_logp_rows(ctx.ref_model.forward_only(inputs=mb).get('logps'), mb_pos)[:n_mb] - for j, val in enumerate(vals): - out[valid_idx[st + j]] = float(val) - finally: - _set_ref_executor_template(ctx, enable_thinking=(ctx.args.skill_thinking == 'on')) - return out - - -def _train_batch(ctx: MethodContext, samples: List[Dict[str, Any]], - traj_fn: Callable[[Dict[str, Any]], Dict[str, Any]], - teacher_fn: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None, - ) -> Tuple[int, Dict[str, float]]: - """Parameter update(s) over ``samples`` with a swappable trajectory builder. - - Mirrors v2 ``_train_step`` (empty-response filter, drop_last to TRAIN_DP, micro-batch by - sft_batch_size, ppo_mini_batch_size multi-step with ALL ref/old/teacher logps pre-computed - BEFORE the first optimizer step — the teacher is the trainable model itself, so computing - it after a step would go off-policy). Returns (n_updates, train metrics). - """ - args = ctx.args - samples = [s for s in samples - if (s.get('tokens') if s.get('tokens') is not None else (s.get('response') or '').strip())] - is_opsd = teacher_fn is not None - teacher_pos: Optional[List] = None - n_align_drop = 0 - if is_opsd: - samples, teacher_pos, n_align_drop = _align_teacher(ctx, samples, traj_fn, teacher_fn) - if not samples: - return 0, {} - n_keep = (len(samples) // TRAIN_DP) * TRAIN_DP - if n_keep == 0: - return 0, {} - samples = samples[:n_keep] - if teacher_pos is not None: - teacher_pos = teacher_pos[:n_keep] - trajs = [traj_fn(s) for s in samples] - advs = None if is_opsd else [float(s['advantage']) for s in samples] - # micro 尺寸与“攒批/采样批（sft_batch_size=16，冻结口径）”解耦：think/8192 实验序列长一倍， - # fp32 主权重后 8 条/卡的 backward 会 OOM，用 --train-micro-batch 切细（梯度按 micro 数归一， - # 数学等价）；默认 0 = 跟随 sft_batch_size，nothink 实验行为不变。 - n = len(trajs) - sft = getattr(args, 'train_micro_batch', 0) or args.sft_batch_size - mini = args.ppo_mini_batch_size if args.ppo_mini_batch_size > 0 else n - mini = max(sft, (mini // sft) * sft) - multi_step = mini < n - # pre-compute every micro's ref/old/teacher logps BEFORE any update (v2 pattern) - micro_ref, micro_old, micro_teacher, micro_smp = [], [], [], [] - # 采样端 logprob，只给 GRPOMetric 做对账（sampler_logp_mae / sampler_token_delta），不进 loss。 - # 整个 micro 都带齐了才传：SFT 样本的 response 是合成文本、没有采样 logprob，混进去会 - # 让 token_delta 无法解释（它的语义是「应恒为 0」）。 - smp_all = [list(s.get('logprobs') or []) for s in samples] - for i in range(0, n, sft): - mb = trajs[i:i + sft] - smp = smp_all[i:i + sft] - micro_smp.append(smp if all(smp) else None) - if is_opsd: - # no ref forward at all: OPSDLoss uses only teacher_logps (kl_beta plays no role) - t_mb = [teacher_fn(s) for s in samples[i:i + sft]] - t_full = ctx.skill_model.forward_only(inputs=t_mb).get('logps') - micro_teacher.append(_gather_response_logps(t_full, teacher_pos[i:i + sft])) - micro_ref.append(None) - micro_old.append(None) - else: - micro_ref.append(ctx.ref_model.forward_only(inputs=mb).get('logps')) - micro_old.append(ctx.skill_model.forward_only(inputs=mb).get('logps') if multi_step else None) - micro_teacher.append(None) - n_steps = 0 - for ms in range(0, n, mini): - for i in range(ms, min(ms + mini, n), sft): - k = i // sft - if is_opsd: - ctx.skill_model.forward_backward(inputs=trajs[i:i + sft], - teacher_logps=micro_teacher[k], - sampler_logps=micro_smp[k]) - else: - ctx.skill_model.forward_backward(inputs=trajs[i:i + sft], advantages=advs[i:i + sft], - old_logps=micro_old[k], ref_logps=micro_ref[k], - sampler_logps=micro_smp[k]) - ctx.skill_model.clip_grad_and_step() - n_steps += 1 - ctx.ckpt.sync_weights(merge_and_sync=True) - metrics = _train_metrics(ctx.skill_model.calculate_metric(is_training=True)) - metrics['train/n_samples'] = float(n) - if is_opsd and n_align_drop: - metrics['train/n_align_dropped'] = float(n_align_drop) - return n_steps, metrics - - -# =========================================================================================== -# method plugins -# =========================================================================================== -class TrainMethod: - needs_rubric: bool = False - - def __init__(self, ctx: MethodContext): - self.ctx = ctx - - def step(self, chunk: List[Dict[str, Any]], ci: int) -> Dict[str, Any]: - raise NotImplementedError - - -class BnpoMethod(TrainMethod): - """view B, query-only GRPO/BNPO — v2 process_chunk + _train_step verbatim.""" - needs_rubric = False - - def step(self, chunk, ci): - ctx = self.ctx - full, summary, grpo, _buf_a = process_chunk( - ctx.base_sampler, ctx.skill_sampler, chunk, ci, ctx.base_dp, ctx.skill_dp, ctx.args) - if grpo and getattr(ctx.args, 'drop_zero_adv', False): - grpo = [g for g in grpo if abs(g['advantage']) > 1e-9] - n_upd, tmetrics = 0, {} - if grpo: - log = _train_step(ctx.skill_model, ctx.ref_model, ctx.ckpt, grpo, ctx.args) - # step = ACTUAL parameter updates: empty-response filter / drop_last may yield 0 - n_upd = int(log.get('n_steps', 0)) - tmetrics = _train_metrics(log.get('metric')) - tmetrics['train/n_samples'] = float(log.get('n_grpo', 0) + log.get('n_sft', 0)) - metrics = {'signal/zero_grad_frac': summary['zero_grad_frac'], - 'signal/group_reward_std_mean': summary['group_reward_std_mean'], - 'signal/n_train_samples': float(summary['n_train_samples']), - 'signal/n_groups': float(summary['n_groups']), - # 旧名（题级 pass@K，历史面板兼容） - 'acc/withskill_pass': summary['avg_withskill_pass'], - # ⭐ 与 SEAM ray_trainer.py:1569-1583 同名同口径，用于 swanlab 直接叠图对齐： - # train/with_skill_accuracy = 全部候选的 mean(correct)（SEAM withskill_pass） - # acc/reward_mean = mean(correct∧format)（SEAM reward_mean） - # skill/format_rate = SEAM format_mean - 'train/with_skill_accuracy': summary['withskill_pass_all_cands'], - 'acc/reward_mean': summary['reward_mean'], - 'skill/format_rate': summary['parse_rate'], - 'leak/rate': summary['leak_rate'], **tmetrics, - **_cand_pass_metrics(chunk), - **_leak_split(_cand_leak_pairs(chunk))} - if summary.get('baseline_pass_train') is not None: - # 训练侧 no-skill baseline：SEAM 只在 step1 跑（ray_trainer.py:1461），所以 twinkle 也只在 - # chunk 0 有值，两边 lift 就是同一个可比的点。 - metrics['acc/baseline_pass'] = summary['baseline_pass_train'] - metrics['acc/lift'] = summary['lift_train'] - metrics['train/baseline_accuracy'] = summary['baseline_pass_train'] - metrics['train/lift'] = summary['lift_train'] - return {'n_updates': n_upd, 'summary': summary, 'metrics': metrics, - 'gen_records': full} - - -class _RLViewA(TrainMethod): - """Shared view-A RL: bare solve -> A(query+rubric)/B(query-only) skill-gen -> reward -> - group advantage -> BNPO. ``train_b`` toggles whether the right-answer B line is trained. - - API/GPU overlap: rubric diagnoses for the wrong (A) problems run on background threads - WHILE the right (B) problems' query-only skill-gen + greedy validation run on GPU; the - A-line rollout starts as soon as the diagnoses land.""" - needs_rubric = True - train_b = True - - def step(self, chunk, ci): - ctx = self.ctx - args = ctx.args - rolls = _bare_solve(ctx, chunk) - wrong = [(r, roll) for r, roll in zip(chunk, rolls) if not roll['correct']] - right = [r for r, roll in zip(chunk, rolls) if roll['correct']] - # kick off rubric API calls in the background, then run the B line on GPU meanwhile - diag_pool = ThreadPoolExecutor(max_workers=1) - diag_fut = diag_pool.submit( - _diagnose_parallel, ctx, [(_rubric_entry(r, roll), None) for r, roll in wrong]) - try: - b_items = [{'record': r, 'prompt': _skillgen_prompt(r['problem'])} for r in right] - for r in right: - r['_rubric'] = '' - if b_items: - _skillgen_solve(ctx, b_items, args.n_skills, temperature=args.skill_gen_temperature) - diags = diag_fut.result() - finally: - diag_pool.shutdown(wait=False) - a_items = [] - degraded = [] # bugfix #2: rubric 缺失（API 失败/坏缓存）→ 降级 query-only B 线，绝不训练空 rubric prompt - for (r, _roll), diag in zip(wrong, diags): - r['_rubric'] = diag - if diag: - a_items.append({'record': r, 'prompt': rubric_skillgen_prompt(r['problem'], diag)}) - else: - degraded.append({'record': r, 'prompt': _skillgen_prompt(r['problem'])}) - if degraded: - _skillgen_solve(ctx, degraded, args.n_skills, temperature=args.skill_gen_temperature) - if a_items: - _skillgen_solve(ctx, a_items, args.n_skills, temperature=args.skill_gen_temperature) - _assign_advantages(chunk, args) - a_recs = _grpo_records([it['record'] for it in a_items], with_rubric=True) - b_recs = _grpo_records(right + [it['record'] for it in degraded], with_rubric=False) - train_recs = a_recs + (b_recs if self.train_b else []) - has_signal = any(abs(s['advantage']) > 1e-9 for s in train_recs) - if has_signal and getattr(args, 'drop_zero_adv', False): - train_recs = [s for s in train_recs if abs(s['advantage']) > 1e-9] - n_upd, tmetrics = 0, {} - if has_signal: - n_upd, tmetrics = _train_batch( - ctx, train_recs, - traj_fn=lambda s: (rubric_train_trajectory(s) if s['with_rubric'] - else query_only_train_trajectory(s))) - return {'n_updates': n_upd, - 'metrics': {'signal/n_wrong_A': float(len(a_items)), - 'signal/n_right_B': float(len(right)), - 'signal/n_rubric_missing': float(len(degraded)), **tmetrics, - **_cand_pass_metrics(chunk), - **_leak_split(_cand_leak_pairs(chunk))}, - 'gen_records': v2._full_records(chunk, ci)} - - -class RlAbMethod(_RLViewA): - train_b = True - - -class RlErrMethod(_RLViewA): - train_b = False - - -class OpsdMethod(TrainMethod): - """view A OPSD (skill_quality_analysis.md 改进skill+OPSD): first-pass ONE skill - (query-only, T=improve), executor solve WITH skill; for WRONG problems, diagnose the - with-skill failure (local cache, key=data_id+skill), then distill the skill-gen response - from the query-only (student) toward the rubric-in-system-prompt (teacher) distribution - per token. Rubric API calls run threaded right after the failures are known.""" - needs_rubric = True - - def step(self, chunk, ci): - ctx = self.ctx - args = ctx.args - # first-pass ONE skill per problem (query-only, improve temperature) - sg = _run_samples(ctx.skill_sampler, [_skillgen_prompt(r['problem']) for r in chunk], - 1, args.skill_max_tokens, ctx.skill_dp, - temperature=args.improve_skill_temperature, logprobs=1) - first, toks_by, lps_by = [], {}, {} - for r, seqs in zip(chunk, sg): - resp = _clean_text(getattr(seqs[0], 'decoded', '') or '') if seqs else '' - # 采样 token 另存一份（每题只一个 skill，用 id(r) 做键就够）：它是 student/teacher 两边 - # 要打分的同一串 token，不能拿 decode 后的文本重编码（见 v2.build_train_feature）。 - # logprobs 同长，只给 GRPOMetric 做采样/训练对账。 - toks_by[id(r)] = ([int(t) for t in (getattr(seqs[0], 'tokens', None) or [])] - if seqs else []) - lps_by[id(r)] = v2.sampler_logprobs(seqs[0]) if seqs else [] - first.append((r, resp, _extract_skill(resp) or '')) - # executor solve WITH skill (only parseable skills) - flat = [(r, resp, sk) for r, resp, sk in first if sk] - roll_by = {} - if flat: - solve = _run_samples(ctx.base_sampler, - [build_skill_solve_prompt(r['problem'], sk) for r, _, sk in flat], - 1, args.max_tokens, ctx.base_dp, temperature=0.0) - for (r, _, sk), seqs in zip(flat, solve): - roll_by[id(r)] = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - wrong = [(r, resp, sk, roll_by[id(r)]) for r, resp, sk in first - if sk and id(r) in roll_by and not roll_by[id(r)]['correct']] - leak_pairs = [(bool(_answer_leaked(sk, r['reference_answer'])), roll_by[id(r)]['correct']) - for r, _resp, sk in flat if id(r) in roll_by] - # rubric diagnoses (threaded; nothing left to overlap on GPU this chunk) - diags = _diagnose_parallel( - ctx, [(_rubric_entry(r, roll), sk) for r, _resp, sk, roll in wrong]) - samples = [{'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'data_id': r.get('data_id', ''), 'response': resp, 'rubric': diag, - 'tokens': toks_by.get(id(r)) or [], - 'logprobs': lps_by.get(id(r)) or []} - for (r, resp, sk, _roll), diag in zip(wrong, diags)] - n_upd, tmetrics = 0, {} - if samples: - n_upd, tmetrics = _train_batch(ctx, samples, - traj_fn=query_only_train_trajectory, # student: query-only - teacher_fn=opsd_teacher_trajectory) # teacher: +rubric in system - return {'n_updates': n_upd, - 'metrics': {'signal/n_wrong': float(len(wrong)), **tmetrics, - **_leak_split(leak_pairs)}, - 'gen_records': [ - {'record_type': 'problem', 'chunk': ci, 'data_id': r.get('data_id', ''), - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'skill': sk, 'parseable': bool(sk), - 'withskill_correct': roll_by[id(r)]['correct'] if id(r) in roll_by else None} - for r, _resp, sk in first]} - - -class LogpRlMethod(TrainMethod): - """E14+: rubric-audited pseudo-GT + executor logP dense reward. - - Per problem, sample K executor attempts without skill at T>0, keep the first locally-correct - and non-truncated response S, audit it through the rubric API/cache, then score each generated - skill by Δ mean logP_executor(S | problem + skill). Problems with no S are skipped this round. - """ - needs_rubric = True - - def step(self, chunk, ci): - usable, sg = self._prepare(chunk, ci) - return self._score_and_train(chunk, ci, usable, sg) - - def _skillgen(self, usable): - ctx = self.ctx - args = ctx.args - return _run_samples(ctx.skill_sampler, [_skillgen_prompt(r['problem']) for r in usable], - args.n_skills, args.skill_max_tokens, ctx.skill_dp, - temperature=args.skill_gen_temperature, top_p=args.skill_gen_top_p, - top_k=args.skill_gen_top_k, logprobs=1) - - def _prepare(self, chunk, ci): - """executor T>0 采 K 条 -> 选本地判分正确且非截断的伪 GT S -> rubric API 后台审计 - （与 skill-gen 的 GPU rollout 重叠）。返回 (usable, 每题 skill-gen 序列)。""" - ctx = self.ctx - args = ctx.args - for r in chunk: - r['_cands'] = [] - K = max(1, int(getattr(args, 'reward_rollouts', 1) or 1)) - temp = float(getattr(args, 'reward_temperature', 0.0) or 0.0) - out = _run_samples(ctx.base_sampler, [build_direct_prompt(r['problem']) for r in chunk], - K, args.max_tokens, ctx.base_dp, temperature=temp) - usable, audit_jobs = [], [] - for r, seqs in zip(chunk, out): - rolls = [_parse_seq(s, r['reference_answer']) for s in (seqs or [])] - ok_i = next((j for j, x in enumerate(rolls) - if x['correct'] and x.get('stop_reason') != 'length'), None) - r['_pseudo_rolls'] = rolls - if ok_i is None: - continue - ok = rolls[ok_i] - r['_pseudo_roll'] = ok - r['_pseudo_solution'] = ok.get('text', '') - # 伪 GT 是 executor 采样产出，打分时直接用它的 token，不拿 _clean_text 后的文本重编码。 - # 从 seq 取而不是从 roll 取：roll 会被每个候选持有，token 存进去会把内存撑爆 - # （E16 那种 M=8 的臂一个 chunk 就是上千条 rollout）。 - r['_pseudo_tokens'] = [int(t) for t in (getattr(seqs[ok_i], 'tokens', None) or [])] - usable.append(r) - audit_jobs.append((_rubric_entry(r, ok), ok.get('text', ''))) - # rubric 审计是纯 API：后台线程跑，与 skill-gen 的 GPU rollout 重叠（API/GPU overlap） - diag_pool = ThreadPoolExecutor(max_workers=1) - diag_fut = diag_pool.submit(_diagnose_parallel, ctx, audit_jobs) - try: - sg = self._skillgen(usable) - diags = diag_fut.result() - finally: - diag_pool.shutdown(wait=False) - for r, diag in zip(usable, diags): - r['_rubric'] = diag - return usable, sg - - def _score_and_train(self, chunk, ci, usable, sg): - ctx = self.ctx - args = ctx.args - flat = [] - for r, seqs in zip(usable, sg): - for si, s in enumerate(seqs or []): - resp = _clean_text(getattr(s, 'decoded', '') or '') - block = _extract_skill(resp) or '' - pseudo = dict(r['_pseudo_roll']) - if si > 0: - pseudo['text'] = '' # 磁盘保护：伪 GT 全文每题只在首个候选保留一份 - _toks = [int(t) for t in (getattr(s, 'tokens', None) or [])] - cand = {'skills': block, 'response': resp, 'parseable': bool(block), - 'leaked': None, 'with_pass': None, 'reward': None, 'rolls': [pseudo], - 'advantage': 0.0, 'kept': False, 'tokens': _toks, - 'logprobs': v2.sampler_logprobs(s), - 'skillgen_stop': getattr(s, 'stop_reason', None), - 'skillgen_tokens': len(_toks), - 'logp_base': None, 'logp_skill': None, 'logp_delta': None} - r['_cands'].append(cand) - if block: - cand['leaked'] = _answer_leaked(block, r['reference_answer']) - flat.append((r, cand)) - for c in r['_cands']: - if c['leaked'] is None: - c['leaked'] = False - if flat: - base_trajs = [_executor_answer_trajectory(r['problem'], '', r['_pseudo_solution'], - answer_tokens=r.get('_pseudo_tokens'), - template=ctx.encode_template) - for r in usable] - base_logps = _score_executor_mean_logps(ctx, base_trajs) - base_by_id = {id(r): lp for r, lp in zip(usable, base_logps)} - cand_trajs = [_executor_answer_trajectory(r['problem'], c['skills'], r['_pseudo_solution'], - c.get('response'), - answer_tokens=r.get('_pseudo_tokens'), - template=ctx.encode_template) - for r, c in flat] - cand_logps = _score_executor_mean_logps(ctx, cand_trajs) - # bugfix #14: logP 目标超长被 truncation='delete' 删掉时 lp=None → reward 地板， - # 整组塔到 -1.0 会零梯度空转；这里显式监控 encode 失败占比（E15 的 R1 参考解尤其长）。 - _n_enc = len(base_logps) + len(cand_logps) - enc_fail_frac = ((sum(1 for v in base_logps if v is None) - + sum(1 for v in cand_logps if v is None)) / _n_enc) if _n_enc else 0.0 - # leak 一律不进 reward（项目既定要求）：--logp-leak-penalty 默认 0，leaked 只做监控。 - # 该指标假阳性极高（单字符 gold 误报 ~84%，c0 实测 raw leak/rate=0.889），-1.0 量级是 - # delta 信号的 50~100 倍，会用噪声主导组内 advantage 并抬高组 std 压小其余候选。 - # format 地板（unparseable / logP 编码失败）固定 -1.0，与 leak 完全解耦、不受影响。 - floor = 1.0 - leak_pen = abs(float(getattr(args, 'logp_leak_penalty', 0.0))) - for (r, c), lp in zip(flat, cand_logps): - base_lp = base_by_id.get(id(r)) - c['logp_base'], c['logp_skill'] = base_lp, lp - if base_lp is None or lp is None: - c['reward'] = -floor - continue - delta = float(lp) - float(base_lp) - c['logp_delta'] = delta - # leak_pen 默认 0（leak 只做监控口径，不进 reward）。若显式开启，仍只对 >=2 字符的 - # 答案生效：usable 子集答案多为 0/1/2/4 等单字符，_answer_leaked 子串匹配在正常数学 - # 叙述里误报率 ~84%。leaked 字段全量记录，监控口径不变。 - informative = len(str(r['reference_answer']).strip()) >= 2 - c['reward'] = delta - (leak_pen if (c.get('leaked') and informative) else 0.0) - for r in chunk: - for c in r.get('_cands', []): - if c['reward'] is None: - c['reward'] = -1.0 - _assign_advantages(chunk, args) - grpo = _grpo_records(usable, with_rubric=False) - has_signal = any(abs(s['advantage']) > 1e-9 for s in grpo) - if has_signal and getattr(args, 'drop_zero_adv', False): - grpo = [s for s in grpo if abs(s['advantage']) > 1e-9] - n_upd, tmetrics = 0, {} - if has_signal: - n_upd, tmetrics = _train_batch(ctx, grpo, traj_fn=query_only_train_trajectory) - summary = v2._chunk_summary(chunk, ci) - # logp_rl 的 reward 不是 executor rollout 通过率；覆盖旧 BNPO summary 中与 pass 绑定的字段， - # 避免 train_log 把“reward 非零”误读成 with-skill pass。 - summary['avg_withskill_pass'] = 0.0 - summary['candidate_withskill_pass'] = 0.0 - summary['withskill_trunc_frac'] = 0.0 - summary['termination_rate_withskill'] = 0.0 - rewards = [c['reward'] for r in usable for c in r.get('_cands', [])] - deltas = [c['logp_delta'] for r in usable for c in r.get('_cands', []) - if c.get('logp_delta') is not None] - metrics = {'signal/pseudo_gt_rate': float(len(usable)) / max(1, len(chunk)), - 'signal/n_pseudo_gt': float(len(usable)), - 'signal/zero_grad_frac': summary['zero_grad_frac'], - 'logp/encode_fail_frac': (enc_fail_frac if flat else 0.0), - 'logp/reward_mean': v2._mean(rewards), - 'logp/reward_std': v2._std(rewards), - 'logp/delta_mean': v2._mean(deltas), - 'logp/delta_std': v2._std(deltas), - 'leak/rate': summary['leak_rate'], **tmetrics, - **_leak_split(_cand_leak_pairs(chunk))} - return {'n_updates': n_upd, 'summary': summary, 'metrics': metrics, - 'gen_records': v2._full_records(chunk, ci)} - - -class LogpGtMethod(LogpRlMethod): - """E15: identical dense executor-logP reward, but the target S is DeepMath's external R1 - reference solution (record 'solution'), NOT an executor-sampled pseudo-GT. No executor - rollout and no rubric audit -> much faster; every problem carrying a solution is usable.""" - needs_rubric = False - - def _prepare(self, chunk, ci): - for r in chunk: - r['_cands'] = [] - usable = [] - for r in chunk: - sol = (r.get('solution') or '').strip() - r['_pseudo_rolls'] = [] - if not sol: - continue - # 合成 roll：logP 目标是外部 R1 参考解，无 executor 采样；stop_reason='gt' 仅作标记， - # 字段与 _parse_seq 输出对齐（_roll 序列化需 pred/correct/terminated/stop_reason/gen_tokens/text）。 - r['_pseudo_roll'] = {'pred': str(r['reference_answer']), 'correct': True, - 'terminated': True, 'stop_reason': 'gt', - 'gen_tokens': 0, 'text': sol} - r['_pseudo_solution'] = sol - r['_rubric'] = '' - usable.append(r) - return usable, self._skillgen(usable) - - -def _hinge_trunc(rolls: List[Dict[str, Any]], lo: int, budget: int = 8192) -> float: - """Mean hinge over rollouts: 0 below ``lo`` tokens, ramps to 1.0 at the ``budget`` line. - Probe (skill_quality_analysis.md 2026-07-29): the length->correctness link exists ONLY near - the truncation budget (zero-truncation groups show +0.03), so penalize the danger zone only, - smoothly, giving gradient BEFORE the hard 'length' cutoff fires.""" - if not rolls: - return 0.0 - span = max(1, budget - lo) - return sum(max(0.0, (int(x.get('gen_tokens') or 0) - lo) / span) for x in rolls) / len(rolls) - - -# 自我推翻标记词。词表由 .tmp_analysis/reward_shape_calib.py 在 E4 11155 条 rollout 上判别力筛出 -# （错误密度/正确密度比值）：confusing 3.06、contradiction 2.49、mistake|error|wrong 1.84、 -# alternatively 1.59。刻意排除的词："let me check" 比值 0.75（**反向**——检查一次是好行为，罚它 -# 有害）、"hmm" 0.86、"recompute|again" 1.00（零信号）、"wait" 仅 1.51 且覆盖率 0.993（几乎人人都写， -# 区分度低）。 -_LOOP_MARKER_RE = re.compile( - r'\b(?:confusing|confused|contradiction|contradicts|contradictory|mistake|error|wrong' - r'|alternatively)\b', re.I) - - -def _loop_density(roll: Dict[str, Any]) -> float: - """Self-revision marker density, occurrences per 1000 generated tokens.""" - tok = int(roll.get('gen_tokens') or 0) - if tok <= 0: - return 0.0 - return len(_LOOP_MARKER_RE.findall(roll.get('text') or '')) / tok * 1000.0 - - -# --- skill 套话度监控 ------------------------------------------------------------------- -# ★ 定位：这两条是退化监控器（看趋势），不是质量预测器（看绝对值）。 -# 2026-07-30 在 E17 的1246 个真实候选上试过 7 种定义（.tmp_analysis/generic_index_*.py）， -# 题内配对对 pass 的预测力全部在 -0.033 到 +0.012 之间，都弱。根因已查清：narrative -# 文体里“元指令”和“具体动作”是同一句话里交织的（“Avoid rechecking the same rounding -# or error estimates”既是元指令又指向具体对象），句子级二分类根本不成立，调词表无法解决。 -# 保留词频法的依据是方向正确：指数低组 leak=0.610 / 高组 leak=0.474，即套话越多越不给 -# 具体内容。但绝对值偏高（`avoid` 在本数据上命中 1223 次，大部分在具体建议里）， -# ★ 因此只能比较同一根曲线的前后变化，不能拿绝对值评判 skill 好坏。 -_GENERIC_ADVICE_RE = re.compile( - r'\b(carefully|careful|make sure|makes sure|ensure|ensuring|be sure|avoid|avoiding|' - r'remember|keep in mind|bear in mind|double[- ]check|double[- ]checking|verify|verifying|' - r'validate|validating|consider|considering|systematically|systematic|efficiently|efficient|' - r'properly|correctly|accurately|appropriately|appropriate|rigorously|rigorous|' - r'manage|managing|track|tracking|monitor|monitoring|streamline|streamlining|' - r'focus on|stay focused|trust|commit to|committing|self[- ]correct\w*|' - r'step by step|methodical\w*|thorough\w*|concise\w*|precise\w*|' - r'token budget|length budget|within the budget|redundan\w*|unnecessar\w*)\b', re.I) -_WORD_RE = re.compile(r"[A-Za-z][A-Za-z'\-]*") -_SENT_RE = re.compile(r'(?<=[.!?;:])\s+|\n+') -# “句中有没有数学内容”的锚点：数字 | LaTeX/算式 | 句中大写专名（定理名）。 -# 不维护数学名词/动词词表：V5 试过，加了两张表反而把具体洞察误判成空话、预测力降到 0.000。 -_MATH_ANCHOR_RE = re.compile(r'\d|\\[A-Za-z]+|\$|\^|_\{|=|≤|≥|≠|∈|∑|∏|√|(?<!^)(?<![.!?]\s)\b[A-Z][a-z]{2,}') - - -def _generic_advice_fraction(text: str) -> float: - """泛泛建议词占 skill 总词数的比例；空文本返回 0。只看趋势，见上方注释。""" - words = _WORD_RE.findall(text or '') - if not words: - return 0.0 - return len(_GENERIC_ADVICE_RE.findall(text)) / len(words) - - -def _no_math_sentence_fraction(text: str) -> float: - """不含任何数学内容（数字/算式/定理名）的句子占比。 - - 名字只陈述它实际测的东西，不声称它能判“空话”—— 实测题内配对 pass 只有 -0.016， - 但它不经由 leak 中介（词频法那一条经由），两条一起看才能分开“套话变多”与“不给答案”。 - """ - sents = [s.strip() for s in _SENT_RE.split(text or '') if len(s.strip()) >= 15] - if not sents: - return 0.0 - return sum(1 for s in sents if not _MATH_ANCHOR_RE.search(s)) / len(sents) - - -def _digit_fraction(text: str) -> float: - """数字字符占 skill 总字符数的比例（递答案/递具体中间量的代理指标）。""" - t = text or '' - if not t: - return 0.0 - return sum(1 for ch in t if ch.isdigit()) / len(t) - - -def _skill_text_metrics(skills: List[str]) -> Dict[str, float]: - """skill 文本面板（按任务分派）。 - - math: 数字占比 / 泛泛建议词占比 / 无数学内容句占比。 - code: 数字占比与"句中有没有数学锚点"在代码域没有语义（API 名、参数、类型天然带大写与 - 符号），换成 skill_contains_code_fraction —— skill 本该是方法论，写出代码围栏或成段 - def/return 就是退化成抄实现，这是代码域最该盯的那条退化曲线。泛泛建议词那条保留： - 它的词表是任务无关的（carefully / make sure / step by step ...）。 - """ - generic = v2._mean([_generic_advice_fraction(s) for s in skills]) - if v2._TASK == 'code': - return {'train/skill_generic_advice_fraction': generic, - 'train/skill_contains_code_fraction': v2._mean( - [1.0 if v2.code_task.skill_has_code(s) else 0.0 for s in skills])} - return {'train/skill_digit_fraction': v2._mean([_digit_fraction(s) for s in skills]), - 'train/skill_generic_advice_fraction': generic, - 'train/skill_no_math_sentence_fraction': v2._mean( - [_no_math_sentence_fraction(s) for s in skills])} - - -# --- E18 拒绝采样第三道筛：skill 与 rubric 诊断的词频余弦 -------------------------------- -# 定位：在"executor 已做对"的候选里挑与诊断内容对得上的那条，压掉两类假赢家—— -# 与诊断无关的碰巧做对（含泄露式速通的残余）和与谁都不像的空泛套话。 -# 刻意用去停用词的词频余弦而不是 tfidf/语义模型：可迁移性判别器一节实测"仅 tfidf" -# in-sample 0.983 / OOS 0.541 是纯过拟合；词频余弦纯 stdlib、确定性、可离线复算。 -_SIM_STOPWORDS = frozenset( - 'the a an and or of to in is are be for with that this it on as by from at not no was ' - 'were will would can could should may might do does did have has had you your we they ' - 'he she its if then than so but into over under out up down when where which what how ' - 'why all any each more most other some such only own same very'.split()) -_SIM_WORD_RE = re.compile(r"[a-z][a-z'\-]{2,}") - - -def _rubric_similarity(skill: str, rubric: str) -> float: - """内容词词频余弦 ∈ [0,1]；任一侧无内容词返回 0。""" - ca = Counter(w for w in _SIM_WORD_RE.findall((skill or '').lower()) - if w not in _SIM_STOPWORDS) - cb = Counter(w for w in _SIM_WORD_RE.findall((rubric or '').lower()) - if w not in _SIM_STOPWORDS) - if not ca or not cb: - return 0.0 - dot = float(sum(v * cb[k] for k, v in ca.items() if k in cb)) - na = sum(v * v for v in ca.values()) ** 0.5 - nb = sum(v * v for v in cb.values()) ** 0.5 - return dot / (na * nb) if na and nb else 0.0 - - -def _efficiency_terms(rolls: List[Dict[str, Any]], *, budget: int, len_lo: int, len_pow: float, - alpha_len: float, beta_loop: float, loop_lo: float, loop_hi: float - ) -> Tuple[float, float, Dict[str, float]]: - """Per-rollout efficiency factor, aggregated. Returns (mean_score, mean_inefficiency, diag). - - Calibration (.tmp_analysis/reward_shape_calib.py, E4 11155 rollouts + DeepMath r1_solution_1): - * Length is NOT harmful per se on this dataset. P(correct | tokens) is flat at 0.96-0.98 up - to 5500 tokens, dips to 0.932/0.903 at 5500-7500, then collapses to 0.226 past 7500 — and - the same shape holds inside every difficulty stratum. Non-truncated long rollouts still - pass at ~0.93, so the damage comes from hitting the wall, not from being long. - * The GT reference solutions (r1_solution_1) are LONGER than the model's correct answers - (p50 4377 vs 3444 tokens; GT p90 9816 already exceeds the 8192 budget). "Shorter is - better" is empirically false here, so a monotone-from-zero length penalty would tax the - median GOOD answer — hence the dead zone below ``len_lo`` (人工拍板 A, 2026-07-29). - * Convex ramp (``len_pow`` > 1) concentrates the penalty in the last ~1300 tokens before the - budget, matching the flat-then-cliff damage curve while still making the marginal penalty - grow with length. - * Marker density has a real but small INDEPENDENT effect: raw dose-response is pass - 0.908 -> 0.412 across density 0-2 -> 6-9, but inside a fixed token band it shrinks to - 0.976 -> 0.931 and 0.971 -> 0.851, i.e. ~85% of the raw effect is just "longer outputs - mechanically contain more markers". Hence ``beta_loop`` is deliberately small. - - Composition is multiplicative per rollout (人工拍板): eff = (1-a*len_pen)*(1-b*loop_pen), - score = correct * eff. Per-rollout (not per-candidate) so a short correct rollout is never - punished for a sibling rollout that burned the budget. - """ - if not rolls: - return 0.0, 0.0, {'len_pen': 0.0, 'loop_pen': 0.0, 'eff': 1.0, 'loop_density': 0.0} - span = max(1, budget - len_lo) - d_span = max(1e-9, loop_hi - loop_lo) - s_sum = ineff_sum = lp_sum = mp_sum = eff_sum = dens_sum = 0.0 - for x in rolls: - tok = int(x.get('gen_tokens') or 0) - len_pen = min(1.0, (max(0, tok - len_lo) / span) ** len_pow) - dens = _loop_density(x) - loop_pen = min(1.0, max(0.0, (dens - loop_lo) / d_span)) - eff = (1.0 - alpha_len * len_pen) * (1.0 - beta_loop * loop_pen) - s_sum += eff if x.get('correct') else 0.0 - ineff_sum += 1.0 - eff - lp_sum += len_pen - mp_sum += loop_pen - eff_sum += eff - dens_sum += dens - n = float(len(rolls)) - diag = {'len_pen': lp_sum / n, 'loop_pen': mp_sum / n, - 'eff': eff_sum / n, 'loop_density': dens_sum / n} - return s_sum / n, ineff_sum / n, diag - - -class PassrateHingeMethod(TrainMethod): - """E16 (view B, query-only): the data-driven closure of the reward probe. - - Per chunk: (1) baseline greedy solve (T=0) to measure each problem's no-skill executor - output length, keep only the danger band ``base_tok > --base-tok-floor`` (probe: base_tok - vs skill lift +0.62 — the strongest problem filter; a soft floor keeps >=TRAIN_DP problems - so a chunk never fully empties). (2) query-only skill-gen, N candidates. (3) score each - parseable skill over M=``reward_rollouts`` executor rollouts at T=``reward_temperature`` - with a per-rollout multiplicative efficiency factor (see ``_efficiency_terms``): - - eff_i = (1 - alpha_len * len_pen_i) * (1 - beta_loop * loop_pen_i) - reward = mean_i(correct_i * eff_i) - kappa * mean_i(1 - eff_i) - unparseable -> -1.0 floor - - The ``- kappa * mean(1 - eff)`` tail is what keeps FAILING candidates separable: a pure - product would send every wrong candidate back to 0, which is exactly E4's pathology (format - failure / burned budget / wrong method all collapsing onto one reward value). - - Coefficient sizing: the total deduction ``(1 + kappa) * (1 - eff)`` is kept under ONE - pass_rate quantum (1/M), and ``reward = max(reward, pass_rate - 1/M)`` enforces that as a - hard guard. Rationale: pass_rate is what eval measures, so the efficiency signal may break - ties but must never rank "solved it once" below "never solved it". The absolute coefficient - size barely matters — in an all-fail group every reward is -kappa*(1-eff) and A=(R-mean)/std - rescales that spread back to unit magnitude, so the CURVE SHAPE, not the scale, is the signal. - - leak 不参与 reward（项目既定要求；``--reward-leak-gate`` 默认 0），只走监控口径。 - - Group-relative advantage + BNPO on the query-only trajectory (identical to bnpo/logp). - """ - needs_rubric = False - - def __init__(self, ctx: MethodContext): - super().__init__(ctx) - # bugfix #13: 长度死区 len_lo 不随 --max-tokens 联动；lo >= budget 时惩罚恒 0，静默失效。 - # 按标定比例（5500/8192）自动缩放并告警。 - lo = int(getattr(ctx.args, 'reward_trunc_lo', 5500) or 5500) - if lo >= ctx.args.max_tokens: - new_lo = max(1, int(ctx.args.max_tokens * 5500 / 8192)) - sys.stderr.write(f'[ablate] WARNING: --reward-trunc-lo {lo} >= --max-tokens ' - f'{ctx.args.max_tokens} disables the length penalty; ' - f'rescaled to {new_lo}.\n') - ctx.args.reward_trunc_lo = new_lo - - def _score_candidates(self, kept, prompts): - """skill-gen -> M-rollout with-skill solve -> 效率加权 reward（不做 advantage）。 - - 从 step() 里抽出只为了让 E17(ReflexionMethod) 在不复制 reward 代码的前提下换掉 prompt - （query-only -> query+rubric）。行为与抽出前逐字一致；prompts 与 kept 同序同长。 - """ - ctx = self.ctx - args = ctx.args - M = max(1, int(getattr(args, 'reward_rollouts', 8) or 8)) - # `or 0.5` 会把显式的 0.0 当成“未设”静默提到 0.5（SEAM 口径 m=1/T=0 因此根本 - # 设不进来）。改成 None 判定：未设才用默认。已跑完的 E16 显式传 0.5，行为不变。 - _t = getattr(args, 'reward_temperature', None) - temp = 0.5 if _t is None else float(_t) - alpha_len = abs(float(getattr(args, 'reward_trunc_penalty', 0.12))) - len_lo = int(getattr(args, 'reward_trunc_lo', 5500) or 5500) - len_pow = max(1.0, float(getattr(args, 'reward_len_pow', 2.0) or 2.0)) - beta_loop = abs(float(getattr(args, 'reward_loop_penalty', 0.04))) - loop_lo = float(getattr(args, 'reward_loop_lo', 2.0)) - loop_hi = float(getattr(args, 'reward_loop_hi', 9.0)) - kappa = abs(float(getattr(args, 'reward_ineff_kappa', 0.10))) - # 不可反转护栏：总扣分封顶在一个 pass_rate 量子（1/M）以内，保证"做对过一次"永远排在 - # "一次没做对"之前。量级校验：max 扣分 = (1 + kappa) * (1 - eff_min)。 - pen_cap = 1.0 / M - 1e-6 - leak_gate = abs(float(getattr(args, 'reward_leak_gate', 0.0))) # 默认 0：leak 只做监控 - sg = _run_samples(ctx.skill_sampler, list(prompts), - args.n_skills, args.skill_max_tokens, ctx.skill_dp, - temperature=args.skill_gen_temperature, top_p=args.skill_gen_top_p, - top_k=args.skill_gen_top_k, logprobs=1) - flat = [] - for r, seqs in zip(kept, sg): - for s in seqs or []: - resp = _clean_text(getattr(s, 'decoded', '') or '') - block = _extract_skill(resp) or '' - _toks = [int(t) for t in (getattr(s, 'tokens', None) or [])] - cand = {'skills': block, 'response': resp, 'parseable': bool(block), - 'leaked': None, 'with_pass': None, 'reward': None, 'rolls': [], - 'advantage': 0.0, 'kept': False, 'tokens': _toks, - 'logprobs': v2.sampler_logprobs(s), - 'skillgen_stop': getattr(s, 'stop_reason', None), - 'skillgen_tokens': len(_toks), - 'trunc_pen': None, 'pass_rate': None, - 'loop_pen': None, 'eff': None, 'loop_density': None} - r['_cands'].append(cand) - if block: - cand['leaked'] = _answer_leaked(block, r['reference_answer']) - flat.append((r, cand)) - # M-rollout with-skill solve (T>0) -> per-rollout efficiency-weighted reward - if flat: - solve = _run_samples( - ctx.base_sampler, - # V2 fix: pass the raw response like v2 process_chunk (seam nesting; v2-mode no-op) - [build_skill_solve_prompt(r['problem'], c['skills'], c.get('response')) for r, c in flat], - M, args.max_tokens, ctx.base_dp, temperature=temp) - # 判分批量化（code 任务：跑单测的子进程必须并行，见 v2._parse_many） - pairs, spans = [], [] - for (r, _c), seqs in zip(flat, solve): - start = len(pairs) - pairs.extend((s, r['reference_answer']) for s in (seqs or [])) - spans.append((start, len(pairs))) - judged = v2._parse_many(pairs) - for (r, c), (a, b) in zip(flat, spans): - rolls = judged[a:b] or [_empty_roll()] - c['rolls'] = rolls - pr = sum(1.0 for x in rolls if x['correct']) / len(rolls) - score, ineff, diag = _efficiency_terms( - rolls, budget=args.max_tokens, len_lo=len_lo, len_pow=len_pow, - alpha_len=alpha_len, beta_loop=beta_loop, loop_lo=loop_lo, loop_hi=loop_hi) - c['pass_rate'], c['with_pass'] = pr, pr - c['trunc_pen'] = diag['len_pen'] # 名字保留，语义为长度惩罚（swanlab 面板连续） - c['loop_pen'], c['eff'] = diag['loop_pen'], diag['eff'] - c['loop_density'] = diag['loop_density'] - reward = score - kappa * ineff - # 护栏：总扣分不得超过一个 pass 量子，否则会反转 pass_rate 排序 - reward = max(reward, pr - pen_cap) - informative = len(str(r['reference_answer']).strip()) >= 2 - if c['leaked'] and informative and leak_gate > 0: - reward -= leak_gate - c['reward'] = reward - for r in kept: - for c in r['_cands']: - if c['reward'] is None: # unparseable / no rolls -> format floor - c['reward'] = -1.0 - - def _reward_panel(self, kept) -> Dict[str, float]: - """reward / 效率面板（E16 与 E17 共用，面板曲线口径保持一致）。""" - def g(k): - return [c[k] for r in kept for c in r['_cands'] if c.get(k) is not None] - rewards = g('reward') - return {'acc/pass_rate_mean': v2._mean(g('pass_rate')), - 'term/trunc_pen_mean': v2._mean(g('trunc_pen')), # = 长度惩罚 len_pen - 'term/loop_pen_mean': v2._mean(g('loop_pen')), - 'term/eff_mean': v2._mean(g('eff')), - 'term/loop_density_mean': v2._mean(g('loop_density')), - 'reward/mean': v2._mean(rewards), 'reward/std': v2._std(rewards)} - - def _gen_records(self, kept, ci) -> List[Dict[str, Any]]: - """v2._full_records 加上 E16/E17 特有的字段（pass_rate / 各惩罚项 / base_* / rubric）。""" - gen_records = v2._full_records(kept, ci) - kept_by_id = {r.get('data_id', ''): r for r in kept} - for gr in gen_records: - r = kept_by_id.get(gr.get('data_id', '')) - if r is not None: - gr['base_tok'] = r['_base_tok'] - gr['base_correct'] = r['_base_correct'] - if r.get('_base_stop') is not None: - gr['base_stop'] = r['_base_stop'] - if r.get('_rubric') is not None: - gr['rubric'] = r['_rubric'] - for gc, c in zip(gr.get('candidates', []), (r['_cands'] if r else [])): - gc['pass_rate'] = c.get('pass_rate') - gc['trunc_pen'] = c.get('trunc_pen') - gc['loop_pen'] = c.get('loop_pen') - gc['eff'] = c.get('eff') - gc['loop_density'] = c.get('loop_density') - return gen_records - - def step(self, chunk, ci): - ctx = self.ctx - args = ctx.args - # 1) baseline (no-skill) greedy solve -> base_tok danger-band filter - base_rolls = _bare_solve(ctx, chunk) - for r, br in zip(chunk, base_rolls): - r['_cands'] = [] - r['_base_tok'] = int(br.get('gen_tokens') or 0) - r['_base_correct'] = bool(br['correct']) - floor_tok = int(getattr(args, 'base_tok_floor', 5000) or 0) - if floor_tok > 0: - kept = [r for r in chunk if r['_base_tok'] > floor_tok] - min_keep = max(TRAIN_DP, 4) - if len(kept) < min_keep: # soft floor: never waste a whole chunk on a thin draw - kept = sorted(chunk, key=lambda r: r['_base_tok'], reverse=True)[:min_keep] - else: - kept = list(chunk) - # 2+3) query-only skill-gen -> M-rollout with-skill solve -> 效率加权 reward - self._score_candidates(kept, [_skillgen_prompt(r['problem']) for r in kept]) - _assign_advantages(kept, args) - grpo = _grpo_records(kept, with_rubric=False) - has_signal = any(abs(s['advantage']) > 1e-9 for s in grpo) - if has_signal and getattr(args, 'drop_zero_adv', False): - grpo = [s for s in grpo if abs(s['advantage']) > 1e-9] - n_upd, tmetrics = 0, {} - if has_signal: - n_upd, tmetrics = _train_batch(ctx, grpo, traj_fn=query_only_train_trajectory) - summary = v2._chunk_summary(kept, ci) - base_toks = [r['_base_tok'] for r in chunk] - metrics = {'signal/zero_grad_frac': summary['zero_grad_frac'], - 'signal/n_kept': float(len(kept)), - 'signal/kept_frac': float(len(kept)) / max(1, len(chunk)), - 'signal/base_tok_mean': v2._mean(base_toks), - 'signal/base_correct_frac': v2._mean([1.0 if r['_base_correct'] else 0.0 for r in chunk]), - **self._reward_panel(kept), - 'leak/rate': summary['leak_rate'], **tmetrics, - **_leak_split(_cand_leak_pairs(kept))} - gen_records = self._gen_records(kept, ci) - # bugfix #15: 被 base_tok 门槛筛掉的题也落盘一行精简记录，筛选器本身可审计 - kept_ids = {id(r) for r in kept} - for r in chunk: - if id(r) not in kept_ids: - gen_records.append({'record_type': 'problem_dropped', 'chunk': ci, - 'data_id': r.get('data_id', ''), - 'base_tok': r['_base_tok'], 'base_correct': r['_base_correct']}) - return {'n_updates': n_upd, 'summary': summary, 'metrics': metrics, - 'gen_records': gen_records} - - -class _SFTFamily(TrainMethod): - """Shared SFT accumulation + fire. Subclasses fill ``collect`` to add pool samples.""" - needs_rubric = True - - def _sft_record(self, problem, ref, data_id, skill): - # 实测（2026-07-29，.tmp_analysis/verify_v1v2_think.py）：thinking-on 模板对这种裸 - # <skills> content 会自动注入空 think 块（<think>\n\n</think>\n\n，且计入 labels）， - # 编码后是合法的 nothink 布局，与带实质 think 的 GRPO 样本混训 token 布局兼容； - # 副作用是把模型推向短/空 think，属设计权衡而非 bug（review #5 定案）。 - return {'problem': problem, 'reference_answer': ref, 'data_id': data_id, - 'response': f'<skills>\n{skill}\n</skills>', 'skills': skill, - 'advantage': float(self.ctx.args.sft_weight), 'sft': True} - - def collect(self, chunk) -> Tuple[List[Tuple[bool, bool]], List[Dict[str, Any]]]: - """Roll out + fill the pool; returns (leak_pairs, gen_records).""" - raise NotImplementedError - - def step(self, chunk, ci): - ctx = self.ctx - leak_pairs, gen_records = self.collect(chunk) - ctx.pool.rebalance() # 1:1 by the minority side, surplus DISCARDED (#18b 不积压) - n_upd = 0 - batch_metrics: List[Dict[str, float]] = [] - for batch in ctx.pool.draw_all_ready(): - n, m = _train_batch(ctx, batch, traj_fn=query_only_train_trajectory) # SFT: query-only (#6) - n_upd += n - if m: - batch_metrics.append(m) - # bugfix #16: 多 batch 时按键均值聚合，不再相互覆盖只留最后一批 - tmetrics = {} - if batch_metrics: - keys = set().union(*batch_metrics) - tmetrics = {k: sum(bm[k] for bm in batch_metrics if k in bm) - / sum(1 for bm in batch_metrics if k in bm) for k in keys} - return {'n_updates': n_upd, - 'metrics': {**{f'pool/{k}': float(x) for k, x in ctx.pool.sizes().items()}, - **tmetrics, **_leak_split(leak_pairs)}, - 'gen_records': gen_records} - - # -- shared: regenerate skills under rubric, greedy-validate, pick a 2-in-8 passer -- - def _regen_pick(self, record, diag: str, use_orig_skill: bool, orig_skill: str = ''): - ctx = self.ctx - args = ctx.args - if not diag: - return None - prompt = (_regen_prompt(record['problem'], orig_skill, diag) if use_orig_skill - else rubric_skillgen_prompt(record['problem'], diag)) - sg = _run_samples(ctx.skill_sampler, [prompt], args.passatk_k, args.skill_max_tokens, - ctx.skill_dp, temperature=args.passatk_skill_temp, - top_p=args.passatk_skill_top_p) - seqs = sg[0] if sg else [] - cands = [] - for s in seqs: - resp = _clean_text(getattr(s, 'decoded', '') or '') - skill = _extract_skill(resp) or '' - if not skill or len(skill) > args.skill_char_limit: - continue - if _leak_blocks(skill, record['reference_answer']): # bugfix #4: informative gate - continue - cands.append(skill) - if not cands: - return None - solve = _run_samples(ctx.base_sampler, - [build_skill_solve_prompt(record['problem'], sk) for sk in cands], - 1, args.max_tokens, ctx.base_dp, temperature=0.0) - passers = [] - for sk, seqs2 in zip(cands, solve): - roll2 = _parse_seq(seqs2[0], record['reference_answer']) if seqs2 else _empty_roll() - if roll2['correct'] and roll2['terminated']: - passers.append(sk) - if len(passers) < args.passatk_m: - return None - # pick the passer closest to the length budget (short-but-not-empty floor, as in v2) - return min(passers, key=lambda sk: abs(len(sk) - args.len_budget)) - - -class SftMethod(_SFTFamily): - """Plain SFT: bare wrong -> rubric regen (query+rubric, no orig skill) 2-in-8 -> pool.""" - def collect(self, chunk): - ctx = self.ctx - rolls = _bare_solve(ctx, chunk) - wrong = [(r, roll) for r, roll in zip(chunk, rolls) if not roll['correct']] - diags = _diagnose_parallel(ctx, [(_rubric_entry(r, roll), None) for r, roll in wrong]) - gen_records = [] - for (r, _roll), diag in zip(wrong, diags): - skill = self._regen_pick(r, diag, use_orig_skill=False) - if skill: - ctx.pool.add(self._sft_record(r['problem'], r['reference_answer'], - r.get('data_id', ''), skill), NEG) - gen_records.append({'record_type': 'problem', 'data_id': r.get('data_id', ''), - 'problem': r['problem'], 'regen_accepted': bool(skill)}) - return [], gen_records # bare solve has no skill, so no leak pairs here - - -class ImproveSftMethod(_SFTFamily): - """Improve-skill + SFT: first-pass 1 skill (query-only, T=0.5); correct -> positive pool - (no leak, <=char_limit); parseable-but-wrong -> rubric regen (with orig skill) 2-in-8 -> - negative pool; unparseable first pass is skipped (empty trajectory would only feed the - teacher garbage). Balanced 1:1 pool, majority side discarded per chunk (#15b/#18b). - - API/GPU overlap: the wrong problems' rubric diagnoses run on background threads WHILE - the regen sampling for previously-diagnosed problems occupies the GPU (diagnoses land - before the first regen finishes, so the loop below never blocks on the API).""" - def collect(self, chunk): - ctx = self.ctx - args = ctx.args - # first-pass ONE skill per problem (query-only, improve temperature), greedy-solve - sg = _run_samples(ctx.skill_sampler, [_skillgen_prompt(r['problem']) for r in chunk], - 1, args.skill_max_tokens, ctx.skill_dp, - temperature=args.improve_skill_temperature) - first = [] - for r, seqs in zip(chunk, sg): - resp = _clean_text(getattr(seqs[0], 'decoded', '') or '') if seqs else '' - first.append((r, _extract_skill(resp) or '')) - flat = [(r, sk) for r, sk in first if sk] - rolls_by = {} - if flat: - solve = _run_samples(ctx.base_sampler, - [build_skill_solve_prompt(r['problem'], sk) for r, sk in flat], - 1, args.max_tokens, ctx.base_dp, temperature=0.0) - for (r, sk), seqs in zip(flat, solve): - rolls_by[id(r)] = _parse_seq(seqs[0], r['reference_answer']) if seqs else _empty_roll() - leak_pairs = [(bool(_answer_leaked(sk, r['reference_answer'])), rolls_by[id(r)]['correct']) - for r, sk in flat if id(r) in rolls_by] - wrong = [(r, sk, rolls_by[id(r)]) for r, sk in flat - if id(r) in rolls_by and not rolls_by[id(r)]['correct']] - # launch ALL diagnoses on threads first, then regen serially on GPU as they land - diag_pool = ThreadPoolExecutor(max_workers=max(1, min(args.rubric_workers, len(wrong) or 1))) - futs = [diag_pool.submit( - lambda e=_rubric_entry(r, roll), s=sk: - (ctx.rubric_cache.get_or_diagnose(e, ctx.checker, skill=s) or '') - if ctx.rubric_cache else '') for r, sk, roll in wrong] - gen_records = [] - try: - for r, sk in flat: - roll = rolls_by.get(id(r)) - if roll is None: - continue - if roll['correct']: - # positive seed: first-pass skill that worked, no leak, within char limit - if len(sk) <= args.skill_char_limit \ - and not _leak_blocks(sk, r['reference_answer']): - ctx.pool.add(self._sft_record(r['problem'], r['reference_answer'], - r.get('data_id', ''), sk), POS) - gen_records.append({'record_type': 'problem', 'data_id': r.get('data_id', ''), - 'problem': r['problem'], 'first_correct': True, 'skill': sk}) - for (r, sk, roll), fut in zip(wrong, futs): - # negative seed: rubric regen conditioned on the FAILED first-pass skill - skill = self._regen_pick(r, fut.result(), use_orig_skill=True, orig_skill=sk) - if skill: - ctx.pool.add(self._sft_record(r['problem'], r['reference_answer'], - r.get('data_id', ''), skill), NEG) - gen_records.append({'record_type': 'problem', 'data_id': r.get('data_id', ''), - 'problem': r['problem'], 'first_correct': False, - 'skill': sk, 'regen_accepted': bool(skill)}) - finally: - diag_pool.shutdown(wait=False) - return leak_pairs, gen_records - - -def _answered(roll: Dict[str, Any]) -> float: - """"这条 rollout 到底交没交出一个可判的答案"。 - - math 看 ``<answer>`` / ``\\boxed``；code 域这两个标记恒不出现，必须换成"抽出了可解析的代码 - 块"（judge_many 把 no_code 记在 kind 里）。不换的话该特征在 code run 上组内方差恒为 0， - observe 会整组跳过 —— 面板上信号最强的那条曲线（E17 实测 cum_sigma 23.70）会静默消失。 - """ - if v2._TASK == 'code': - return 1.0 if (roll.get('code') and roll.get('kind') != 'no_code') else 0.0 - text = roll.get('text') or '' - return 1.0 if ('<answer>' in text or '\\boxed' in text) else 0.0 - - -class SnrProbe: - """在线可学信号探针：把组内 advantage 投影到 skill 文本特征上，逐 chunk 上报累积证据。 - - 为什么需要它（E16 事后分析的直接产物，见 .tmp_analysis/batch_size_math.py）： - E16 的 reward 在【结果层】极其确定——"哪个候选让 executor 收住了"每组 SNR 2.25、97% 的组 - 方向一致；但同一个信号投影到任何【skill 文本特征】上，SNR 塌到 0.046、同向组占比 0.481 - （= 抛硬币）。信息是在"从结果归因到文本"这一步丢的，损失约 50 倍。后果是：整个 50 步 run - 在唯一有内容的方向上只累积到 sqrt(1600) x 0.047 ~ 1.9 sigma，连显著性门槛都没到，而策略却 - 以恒定步长（GRPO 组内标准化让 mean|a| 恒为 0.74，信号退化成噪声时步长不会变小）持续扩散 - 离开初始分布——而"冻结 executor 能在预算内收住"恰恰是初始分布自带的脆弱属性。 - - 所以判断一个新臂值不值得跑满 50 步，看的不是 reward/mean（它被 parse 地板的构成变化掩盖， - E16 实测总均值 +0.026 = parse 构成 +0.160 + 层内 -0.134），而是这里的 cum_sigma： - cum_sigma = sqrt(N_groups) * |mean(g)| / std(g) - g 是每组 advantage 与组内标准化目标量的协方差。跑 10-15 个 chunk 就能看出 rubric 条件下的 - 文本层信号是否比 0.046 高一个量级；不高就该停，省 80% 卡时。 - """ - - # (键, 取值函数, 期望方向)；dir=-1 表示"我们希望 reward 压低该量"，上报时已翻正号， - # 所以 mu>0 一律读作"reward 在往我们想要的方向推"。skill_digits 也是 -1：E16 的主导机制 - # 是答案中继（obey|断言对 = 0.99），所以"skill 里的数字变多"是需要报警的方向、不是 - # 鼓励的方向；先前写 +1 会让面板上的正值被误读成好消息。 - # 注：决策变量 cum_sigma 用 |mean|，与方向约定无关，只有 mu / agree 的可读性依赖它。 - TARGETS = ( - ('skill_chars', lambda c: float(len(c.get('skills') or '')), -1), - ('skill_digits', lambda c: float(sum(ch.isdigit() for ch in (c.get('skills') or ''))), -1), - ('think_tokens', lambda c: float(c.get('skillgen_tokens') or 0), -1), - ('exec_trunc', lambda c: _roll_mean(c, lambda x: 1.0 if x.get('stop_reason') == 'length' - else 0.0), -1), - ('exec_tok', lambda c: _roll_mean(c, lambda x: float(int(x.get('gen_tokens') or 0))), -1), - # ---- 2026-07-30 新增 4 条。先用 .tmp_analysis/snr_feature_scan.py 在 E17 的 168 组/ - # 1344 候选上扫过 13 个备选（同一口径），只留下 cum_sigma 过 5 且不与现有项重复的。 - # 被剔掉的（全部 < 3）：exec_tok_spread 2.87、skill_mean_sentence_len 2.67、 - # skill_hedge_frac 2.23、skill_n_sentences 2.04、skill_action_verb_frac 1.72、 - # skill_proper_noun_frac 1.54、skill_step_marker_frac 0.35。 - # ★ exec_answered 实测 23.70、agree 0.966，比旧冠军 exec_trunc(22.21) 还高：把“reward - # 到底在推什么”说得比截断率更直——不是“别写太长”而是“把答案写出来”。dir=+1。 - # （判据按 _TASK 分派，见 _answered） - ('exec_answered', lambda c: _roll_mean(c, _answered), +1), - # loop_pen 就在 reward 公式里（beta=0.04）却一直没进 SNR 面板，那一项到底有没起作用 - # 之前看不到。实测 10.70 / agree 0.789。复用 _loop_density（与 reward 同一个函数）。 - ('exec_loop_density', lambda c: _roll_mean(c, lambda x: _loop_density(x)), -1), - # 唯一直接量化“reward 在多大程度上奖励泄露”的 SNR。实测 8.80 / agree 0.819 / mu=-0.424。 - # 只在组内 leak 有差异的组（实测 83/168）取到值，但那正是要看的那些组。 - # ★ 这是监控形态，不得进 reward（项目既定要求）。dir=-1。 - ('skill_leaked', lambda c: (1.0 if c.get('leaked') else 0.0), -1), - # 比现有的绝对数字数 skill_digits(10.11) 剔掉了 skill 长度混杂；两条都留，差值能 - # 分开“数字变多”与“skill 变长”。实测 6.55 / agree 0.721。 - # 空 skill 返回 0.0（不是 None）：与 skill_chars 口径一致，否则 parse 率随训练上升 - # 会让本特征的取样面系统漂动（正是 observe 里那段注释警告的假趋势源）。 - ('skill_digit_fraction', lambda c: _digit_fraction(c.get('skills') or ''), -1), - ) - - def __init__(self): - # 每个 target 一个累积器：(n, sum(g), sum(g^2), 正号计数) - self._acc: Dict[str, List[float]] = {k: [0.0, 0.0, 0.0, 0.0] for k, _, _ in self.TARGETS} - - def observe(self, records: List[Dict[str, Any]]) -> Dict[str, float]: - """records = 本次更新真正参与训练的 per-problem 记录（带 '_cands'）。返回 swan 标量。""" - out: Dict[str, float] = {} - for key, fn, direction in self.TARGETS: - gs = [] - for r in records: - cs = [c for c in r.get('_cands', []) if c.get('reward') is not None - and c.get('advantage') is not None] - # 只丢"该特征取不到值"的候选，不丢整组。不可解析候选没有 rollout，exec_* 取值 - # 为 None；而 parse 率会随训练上升（E16 实测：全可解析组占比 c0-9 0.757 -> - # c40-49 0.979），所以"任一 None 丢整组"会让取样面随时间系统性扩大 24pp - # —— 那本身就是一个假趋势源，而 cum_sigma 正是要用来判断趋势的。 - # （g = sum(a_i z_i)/n 对 a 加常数不变，z 在存活子集上重新中心化后仍无偏。） - pairs = [(c, fn(c)) for c in cs] - pairs = [(c, v) for c, v in pairs if v is not None] - if len(pairs) < 2: - continue - # advantage 全 0 的组（组内 reward 无差异，_assign_advantages 已置 0）不携带方向 - # 信息；计入只会把均值往 0 拉、同时虚增 n，使 cum_sigma 系统偏低。 - if all(abs(float(c['advantage'])) < 1e-12 for c, _ in pairs): - continue - vals = [v for _, v in pairs] - mu = sum(vals) / len(vals) - var = sum((v - mu) ** 2 for v in vals) / len(vals) - if var < 1e-12: # 组内该特征无差异 -> 这一组对该方向不提供信息 - continue - sd = var ** 0.5 - g = sum(float(c['advantage']) * (v - mu) / sd for c, v in pairs) / len(pairs) - gs.append(direction * g) - if not gs: - continue - a = self._acc[key] - a[0] += len(gs) - a[1] += sum(gs) - a[2] += sum(x * x for x in gs) - a[3] += sum(1.0 for x in gs if x > 0) # 只记正号数，与漂动的运行均值解耦 - n_cum = a[0] - mean_cum = a[1] / n_cum - # 累积 std（总体口径；n 已达数百，与样本口径无实质差别） - var_cum = max(a[2] / n_cum - mean_cum ** 2, 0.0) - sd_cum = var_cum ** 0.5 - out[f'snr/{key}_mu'] = sum(gs) / len(gs) # 本 chunk 的每组均值 - out[f'snr/{key}_mu_cum'] = mean_cum - out[f'snr/{key}_snr_cum'] = (abs(mean_cum) / sd_cum) if sd_cum > 1e-12 else 0.0 - # 同向组占比：取两个符号桶的多数侧，不依赖当时的运行均值（旧实现拿 - # mean_cum 做参系，早期均值不稳时会把同一批组判到不同侧）。 - out[f'snr/{key}_agree_cum'] = max(a[3], n_cum - a[3]) / n_cum - # ★ 决策变量：整个 run 至今在该方向上累积的证据（sigma）。E16 全程只到 1.9。 - out[f'snr/{key}_cum_sigma'] = ((n_cum ** 0.5) * abs(mean_cum) / sd_cum - if sd_cum > 1e-12 else 0.0) - out[f'snr/{key}_n_groups'] = n_cum - return out - - -def _roll_mean(cand: Dict[str, Any], fn: Callable[[Dict[str, Any]], float]) -> Optional[float]: - rolls = cand.get('rolls') or [] - if not rolls: - return None - return sum(fn(x) for x in rolls) / len(rolls) - - -class ReflexionMethod(PassrateHingeMethod): - """E17 —— Reflexion 条件化臂：只在【裸 executor 做错】的题上，用 rubric 生成 skill 并训练。 - - 与父类 PassrateHingeMethod（E16）完全共享 reward 形状（M-rollout pass_rate x 效率加权、 - 死区 5500 起的二次凸长度惩罚、循环惩罚、不可反转护栏、leak 只监控），三处结构差异： - - 1) 选题：父类按 base_tok 危险带筛（floor=5000），本类按【裸解错误】筛，并把批量对齐到 - 恰好 --reflexion-k 道题。对齐的理由：每次更新的组数必须恒定，否则步长与噪声逐 chunk - 变化，SnrProbe 的累积证据和任何趋势读数都会被批量抖动污染。 - 实现上不动态多抽（MethodContext 拿不到 ProblemPool，且 trainer 的断点恢复靠"重放 N 次 - 等长 draw"，变长抽取会破坏恢复），而是把 chunk_size 放大到 ~5K、裸解后逐批取错题 - 并对 rubric 缺失做回填，直到凑满 K；真凑不够时用现有的全部并上报 k_short。 - 标定（.tmp_analysis/e17_param_calib.py，E16 落盘）：全量题池裸错率 0.329（527/1600）， - E16 每次更新实际 23.46 组 / 全程 1173 组，所以 K=24 才能追平 E16 的证据量（累积 - 证据 ∝sqrt(N)，E16 全程在文本层只累到 1.9 sigma，再砍组数就没判别力了）； - chunk=128 时 P(错题<24) = 0.012%。⭐ 不要用 0.446，那是 base_tok>5000 筛选后 - 子集的错误率（偏难），而本臂 floor=0；用它会把 chunk 低估到 96（P=3.6%）。 - 2) skill-gen 走 view A：prompt = rubric_skillgen_prompt(problem, diag)，训练轨迹相应换成 - rubric_train_trajectory（否则会在 query-only 轨迹上训一个 query+rubric 分布下采出的 - response，与 E6 的已知缺陷同型）。rubric API 缺失（失败/坏缓存）的题一律丢弃而不降级 - 成 query-only —— 本臂的唯一自变量就是 rubric，降级样本会把它稀释掉。 - 3) base_tok_floor 强制视为 0。⚠️ 但要知道这几乎不起作用：实测全量题池的 527 道错题里 - 96.96% 是没写完（base_tok>=8192），floor=5000 筛选后也只是 97.71% —— 截断是题目 - （level>=6 配 8192 预算）造成的，不是筛选造成的。所以 rubric 在 ~97% 的题上只能说 - "你超预算了"，signal/wrong_trunc_frac 会直接开在 0.97。 - - API/GPU 重叠：rubric 诊断是纯 API，父类的 GPU 路径在它之后才开始，所以这里先起线程池发 - 诊断、同时不做别的 GPU 工作（本臂没有 B 线可以并行），诊断落地后再进 skill-gen。 - """ - - needs_rubric = True - - def __init__(self, ctx: MethodContext): - super().__init__(ctx) - # 题集由【裸解错】定义，base_tok 危险带筛选强制关闭（用户 2026-07-29 拍板）：错题集 - # 必须保留"推理错 + 没写完"的混合，否则就只剩长尾截断题、测不到方法修正。 - # 在这里而不是在 shell 里置 0，是为了让 config 指纹（trainer 在 build_method 之后才 - # 落盘）记录真实生效值，而不是一个未被读取的默认 5000。 - if int(getattr(ctx.args, 'base_tok_floor', 0) or 0): - sys.stderr.write('[ablate] reflexion: --base-tok-floor forced to 0 (problem set is ' - 'defined by bare-solve failure, not by output length).\n') - ctx.args.base_tok_floor = 0 - self.snr = SnrProbe() - - def step(self, chunk, ci): - ctx = self.ctx - args = ctx.args - K = max(1, int(getattr(args, 'reflexion_k', 16) or 16)) - # 1) 裸解全 chunk，挑错题并对齐到恰好 K 道 - base_rolls = _bare_solve(ctx, chunk) - for r, br in zip(chunk, base_rolls): - r['_cands'] = [] - r['_base_tok'] = int(br.get('gen_tokens') or 0) - r['_base_correct'] = bool(br['correct']) - r['_base_stop'] = br.get('stop_reason') - wrong = [(r, br) for r, br in zip(chunk, base_rolls) if not br['correct']] - # 2) rubric 诊断（纯 API，线程并行；缓存键 = data_id + 裸解轨迹，跳臂全局缓存）。 - # 逐批回填到恰好 K 道：直接 wrong[:K] 会让 rubric 缺失把本 chunk 的组数打到 K 以下， - # 而组数恒定是本臂的硬要求（否则步长、噪声底与 SnrProbe 的累积证据全跟着抖）。 - # 回填只花网络时间不占 GPU；E6/E7 实测缺失率 0，正常路径上循环只进一轮。 - picked, dropped_no_rubric, cursor = [], 0, 0 - while len(picked) < K and cursor < len(wrong): - take = wrong[cursor:cursor + (K - len(picked))] - cursor += len(take) - diags = _diagnose_parallel(ctx, [(_rubric_entry(r, br), None) for r, br in take]) - for (r, br), diag in zip(take, diags): - r['_rubric'] = diag or '' - if diag: - picked.append((r, br, diag)) - else: - dropped_no_rubric += 1 # 不降级成 query-only：rubric 是唯一自变量 - k_short = max(0, K - len(picked)) # 错题不够 + rubric 全打水两种原因合计 - # 组数恒定是本臂的硬要求（步长、噪声底与 SnrProbe 的累积证据都跟着它抖）。原来靠 - # signal/k_short + signal/n_rubric_missing 两条面板指标暴露，面板精简后改走 stderr， - # 否则这个不变量会变成静默失败。 - if k_short or dropped_no_rubric: - sys.stderr.write(f'[E17] c{ci}: WARNING 组数 {len(picked)}/{K}' - f'（缺 {k_short}；其中 rubric 拉不到 {dropped_no_rubric} 道）\n') - kept = [r for r, _br, _d in picked] - # 3) 复用父类的 skill-gen -> M-rollout -> 效率加权 reward -> advantage 流水线 - if kept: - self._score_candidates(kept, [rubric_skillgen_prompt(r['problem'], d) - for r, _br, d in picked]) - _assign_advantages(kept, args) - grpo = _grpo_records(kept, with_rubric=True) - has_signal = any(abs(s['advantage']) > 1e-9 for s in grpo) - if has_signal and getattr(args, 'drop_zero_adv', False): - grpo = [s for s in grpo if abs(s['advantage']) > 1e-9] - n_upd, tmetrics = 0, {} - if has_signal: - n_upd, tmetrics = _train_batch(ctx, grpo, traj_fn=rubric_train_trajectory) - # 4) 指标（2026-07-30 精简）：只保留用户指定的这几条，命名不缩写。 - # 去掉的 leak/* term/* signal/* 仍全量落在 gen_records 里，离线随时可算。 - summary = v2._chunk_summary(kept, ci) if kept else {'zero_grad_frac': 1.0, 'leak_rate': 0.0} - # ★ 两个口径必须分开（否则退化会被隐掉）： - # all = 所有候选，包括解析失败的（skill 为空、reward 拍 -1.0 地板、从未跑 executor） - # scored= 只有真跑了 executor 的（pass_rate 不为 None） - # 只用 scored 算 skill 文本指标会把“退化成空 skill”这个最重要的信号完全遮住 - # （v6 那轮实测空 skill 率 7.4%，全部是 skill 自己写到 8192 撞顶）； - # reward 也必须含地板，否则面板上的 reward 不等于优化器真正看到的那个。 - all_cands = [c for r in kept for c in r['_cands']] - scored = [c for c in all_cands if c.get('pass_rate') is not None] - # roll 里只有 gen_tokens（见 v2._parse_seq），没有 tokens 字段。 - exec_tokens = [float(int(x.get('gen_tokens') or 0)) - for c in scored for x in (c.get('rolls') or [])] - skills = [c.get('skills') or '' for c in all_cands] - # 训练题一律是裸解做错的题，所以 baseline 恒为 0；显式上报是为了让 lift 曲线自解释。 - baseline_accuracy = 0.0 - with_skill_accuracy = v2._mean([c['pass_rate'] for c in scored]) - rewards = [c['reward'] for c in all_cands if c.get('reward') is not None] - metrics = { - 'train/baseline_accuracy': baseline_accuracy, - 'train/with_skill_accuracy': with_skill_accuracy, - 'train/lift': with_skill_accuracy - baseline_accuracy, - 'train/reward_mean': v2._mean(rewards), - 'train/reward_std': v2._std(rewards), - 'train/zero_gradient_fraction': summary['zero_grad_frac'], - # 解析失败率：with_skill_accuracy 只在 scored 上算，这条负责把分母的变化讲出来。 - 'train/skill_parse_failure_rate': (1.0 - len(scored) / len(all_cands)) if all_cands else 0.0, - 'train/skill_length_characters': v2._mean([float(len(s)) for s in skills]), - **_skill_text_metrics(skills), - 'train/executor_length_tokens': v2._mean(exec_tokens), - 'train/executor_loop_rate': v2._mean([c['loop_pen'] for c in scored - if c.get('loop_pen') is not None]), - **tmetrics, # train/loss, train/grad_norm, train/lr, train/iters, train/n_samples - **self.snr.observe(kept), # snr/*：eval 在 R=1 下 MDE 很大，方向判断只能靠这个 - } - gen_records = self._gen_records(kept, ci) - kept_ids = {id(r) for r in kept} - for r in chunk: - if id(r) not in kept_ids: - # drop_reason 让对齐可审计：否则 dump 里分不出"裸解对"、"超过 K 没用上"、 - # "rubric 拉不到"三种丢弃，而只有第三种是需要报警的。 - reason = ('base_correct' if r['_base_correct'] - else 'no_rubric' if r.get('_rubric') == '' else 'beyond_k') - gen_records.append({'record_type': 'problem_dropped', 'chunk': ci, - 'data_id': r.get('data_id', ''), 'drop_reason': reason, - 'base_tok': r['_base_tok'], 'base_correct': r['_base_correct'], - 'base_stop': r['_base_stop']}) - return {'n_updates': n_upd, 'summary': summary, 'metrics': metrics, - 'gen_records': gen_records} - - -class RejectionSftMethod(TrainMethod): - """E18 —— 拒绝采样 SFT（2026-07-30 用户拍板的 9 步方案）： - - 每 chunk：① 全部 query 裸解一次（greedy T=0）判对错 → ② 错题过 rubric 诊断（缺失即丢， - 不降级）→ ③ 按 E17 的 rollout 方式：rubric 条件化 skill-gen（think 模式、T=1.0 × n_skills）， - 每个 skill 让 executor 推理一次（greedy T=0）→ ④ 三道筛选一条： - a. 只留做对的； - b. leak 过滤：含最终答案的丢掉（用 _leak_blocks 的 >=2 字符门：裸 _answer_leaked 对 - 单字符 gold 误报 ~84%，会把池饿死——SFT 家族 bugfix #4 的同一个门；超 - skill_char_limit 的一并丢）； - c. 长度预筛：取离 len_budget 最近的前一半 → 其中与原始 rubric 词频余弦相似度最高的。 - ⑤ 胜者写进本地数据集文件 e18_sft_dataset.jsonl（append-only，含 rubric/相似度/pass 全审计字段） - 并入池 → ⑥ 池满 --e18-accumulate（16，2026-07-30 从 128 改小）条就 SFT 一次（advantage=--sft-weight=1，轨迹用 - query-only + 裸 <skills> 响应 = nothink 布局：thinking-on 模板会自动注入空 think 块， - 与 Qwen3 enable_thinking=False 的生成布局逐 token 一致，review #5 定案）。 - ⑦ _train_batch 内部 ckpt.sync_weights 把新权重推到 vLLM → ⑧ eval 在 trainer 侧：同一个 - skill_sampler vLLM 临时切 nothink 模板跑 query-only greedy eval（只换客户端编码，引擎不动）。 - - 与 SftMethod(E12) 的本质区别：E12 靠 rubric 重生成 2-in-8 验证入池（无拒绝排序）； - E18 在做对的候选里再按 leak/长度/rubric 对齐度三道筛取唯一胜者，且留下可复现的 - 本地数据集文件。train-with-rubric/train-query-only 的选择沿用 SFT 家族定案 #6： - 用 query-only 轨迹训，避免采集分布（query+rubric）与部署分布（query-only）错配。 - """ - needs_rubric = True - - def step(self, chunk, ci): - ctx = self.ctx - args = ctx.args - # ① 裸解全 chunk（greedy T=0 单次），判对错 - base_rolls = _bare_solve(ctx, chunk) - for r, br in zip(chunk, base_rolls): - r['_cands'] = [] - r['_rubric'] = '' - r['_base_correct'] = bool(br['correct']) - wrong = [(r, br) for r, br in zip(chunk, base_rolls) if not br['correct']] - # ② rubric 诊断（纯 API 线程并行；缺失即丢，不降级 query-only：没有诊断就没有 - # 相似度筛的参照系，与 E17「rubric 是唯一自变量」的丢弃规则同型） - diags = _diagnose_parallel(ctx, [(_rubric_entry(r, br), None) for r, br in wrong]) - todo = [] - for (r, br), diag in zip(wrong, diags): - r['_rubric'] = diag or '' - if diag: - todo.append((r, diag)) - n_rubric_missing = len(wrong) - len(todo) - # ③ rubric 条件化 skill-gen（sampler 模板是 think-on，即用户要求的 think 模式采集）； - # _skillgen_solve 内部对每个可解析 skill 跑 executor greedy T=0 单次，正是本臂口径。 - items = [{'record': r, 'prompt': rubric_skillgen_prompt(r['problem'], d)} - for r, d in todo] - if items: - _skillgen_solve(ctx, items, args.n_skills, temperature=args.skill_gen_temperature) - # ④ 逐题三道筛：做对 -> 不 leak/不超长 -> 长度预筛前半 + rubric 相似度最高 - accepted = [] - sims_pool = [] - n_pass_cands = n_leak_dropped = 0 - for r, d in todo: - passers = [c for c in r['_cands'] - if c.get('parseable') and (c.get('with_pass') or 0) > 0] - n_pass_cands += len(passers) - survivors = [c for c in passers - if len(c['skills']) <= args.skill_char_limit - and not _leak_blocks(c['skills'], r['reference_answer'])] - n_leak_dropped += len(passers) - len(survivors) - if not survivors: - continue - # 长度选择：取离 len_budget 最近的前一半（至少 1 条），再在其中比相似度。 - # 两阶而非加权求和：两个量纲不同（字符距 vs 余弦），权重没法标定。 - by_len = sorted(survivors, key=lambda c: abs(len(c['skills']) - args.len_budget)) - shortlist = by_len[:max(1, (len(by_len) + 1) // 2)] - for c in shortlist: - c['rubric_similarity'] = _rubric_similarity(c['skills'], d) - best = max(shortlist, key=lambda c: c['rubric_similarity']) - best['kept'] = True - sims_pool.append(best['rubric_similarity']) - accepted.append({'problem': r['problem'], - 'reference_answer': r['reference_answer'], - 'data_id': r.get('data_id', ''), - 'response': f"<skills>\n{best['skills']}\n</skills>", - 'skills': best['skills'], - 'advantage': float(args.sft_weight), 'sft': True, - # 审计字段（只进数据集文件，不进训练轨迹） - 'rubric': d, 'chunk': ci, - 'rubric_similarity': best['rubric_similarity'], - 'skill_chars': len(best['skills']), - 'n_candidates_passed': len(passers)}) - # ⑤ 胜者落盘本地数据集（append-only，逐 chunk 开关避免长持句柄）+ 入池 - if accepted: - with open(os.path.join(args.output_dir, 'e18_sft_dataset.jsonl'), 'a', - encoding='utf-8') as f: - for s in accepted: - f.write(json.dumps(s, ensure_ascii=False) + '\n') - for s in accepted: - # 训练样本只留 _train_trajectory 需要的键（rubric 不进 query-only 轨迹） - ctx.pool.add({k: s[k] for k in ('problem', 'reference_answer', 'data_id', - 'response', 'skills', 'advantage', 'sft')}, NEG) - # ⑥+⑦ 池满 --e18-accumulate 条即 SFT；_train_batch 内部已含 ckpt.sync_weights - n_upd = 0 - batch_metrics: List[Dict[str, float]] = [] - for batch in ctx.pool.draw_all_ready(): - n, m = _train_batch(ctx, batch, traj_fn=query_only_train_trajectory) - n_upd += n - if m: - batch_metrics.append(m) - tmetrics = {} - if batch_metrics: - keys = set().union(*batch_metrics) - tmetrics = {k: sum(bm[k] for bm in batch_metrics if k in bm) - / sum(1 for bm in batch_metrics if k in bm) for k in keys} - # 指标：命名不缩写（沿用 E17 面板约定） - n_wrong = len(wrong) - metrics = { - 'train/pool_size': float(ctx.pool.sizes().get('pool', 0)), - 'train/accept_rate': (len(accepted) / len(todo)) if todo else 0.0, - 'train/candidate_pass_rate': (n_pass_cands / (len(todo) * args.n_skills)) - if todo else 0.0, - 'train/leak_or_overlength_dropped_fraction': (n_leak_dropped / n_pass_cands) - if n_pass_cands else 0.0, - 'train/selected_rubric_similarity': v2._mean(sims_pool), - 'train/selected_skill_length_characters': v2._mean( - [float(s['skill_chars']) for s in accepted]), - 'signal/n_wrong': float(n_wrong), - 'signal/n_rubric_missing': float(n_rubric_missing), - **tmetrics, - } - # gen_records：v2._full_records 不落 rubric/相似度，补上审计字段（筛选器本身可审计） - gen_records = v2._full_records(chunk, ci) - by_id = {r.get('data_id', ''): r for r in chunk} - for gr in gen_records: - r = by_id.get(gr.get('data_id', '')) - if r is None: - continue - gr['base_correct'] = r.get('_base_correct') - if r.get('_rubric'): - gr['rubric'] = r['_rubric'] - for gc, c in zip(gr.get('candidates', []), r.get('_cands', [])): - if c.get('rubric_similarity') is not None: - gc['rubric_similarity'] = c['rubric_similarity'] - return {'n_updates': n_upd, 'metrics': metrics, - 'gen_records': gen_records} - - -METHOD_REGISTRY: Dict[str, Callable[[MethodContext], TrainMethod]] = { - 'bnpo': BnpoMethod, - 'rl_ab': RlAbMethod, - 'rl_err': RlErrMethod, - 'opsd': OpsdMethod, - 'sft': SftMethod, - 'improve_sft': ImproveSftMethod, - 'logp_rl': LogpRlMethod, - 'logp_gt': LogpGtMethod, - 'passrate_hinge': PassrateHingeMethod, - 'reflexion': ReflexionMethod, - 'rejection_sft': RejectionSftMethod, -} - - -def build_method(method: str, ctx: MethodContext) -> TrainMethod: - if method not in METHOD_REGISTRY: - raise KeyError(f'unknown method {method!r}; valid: {sorted(METHOD_REGISTRY)}') - return METHOD_REGISTRY[method](ctx) diff --git a/cookbook/exp/skill2lora/skill_ablate/pool.py b/cookbook/exp/skill2lora/skill_ablate/pool.py deleted file mode 100644 index 346ba7c69..000000000 --- a/cookbook/exp/skill2lora/skill_ablate/pool.py +++ /dev/null @@ -1,125 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""SamplePool: accumulate training samples across chunks and emit fixed-size batches. - -Ray-infra size rule (skill_quality_analysis.md #11-12): training must drop_last to a -TRAIN_DP multiple and must NEVER pad new sequences (overfitting guard). The actual -drop_last happens inside v2's ``_train_step``; this pool's only job is "accumulate until a -full batch is available, then hand exactly one batch to ``_train_step``". Batch size is a -TRAIN_DP multiple (default 16 = sft_batch_size), so the batch always divides evenly. - -Two modes: -- plain (balanced=False): a single FIFO queue; ready when it holds >= batch_size; a draw - pops the oldest ``batch_size`` samples and keeps the remainder pooled for next time. -- balanced (balanced=True): separate positive/negative queues for the improve-skill+SFT - 1:1 requirement (#15b/#18b). A batch is half positives + half negatives; ready when BOTH - halves are available. After every chunk the caller invokes ``rebalance()``: the majority - side is down-sampled to the minority side and the surplus is DISCARDED immediately - (#18b "以少的一侧为准下采样多的一侧，其余丢弃不积压") — no stale majority backlog can - accumulate, so early-policy easy positives never train dozens of chunks later. - ``max_pool`` remains as a safety cap only (drop oldest). - -This module is dependency-free (pure stdlib) and unit-testable without torch / a GPU. -""" -from collections import deque -from typing import Any, Deque, Dict, List, Optional - -POS, NEG = 'pos', 'neg' - - -class SamplePool: - def __init__(self, batch_size: int = 16, balanced: bool = False, - max_pool: Optional[int] = None): - if batch_size < 1: - raise ValueError('batch_size must be >= 1') - if balanced and batch_size % 2 != 0: - raise ValueError('balanced pool needs an even batch_size (half pos + half neg)') - self.batch_size = batch_size - self.balanced = balanced - self.max_pool = max_pool - self._q: Deque[Dict[str, Any]] = deque() # plain mode - self._pos: Deque[Dict[str, Any]] = deque() # balanced mode - self._neg: Deque[Dict[str, Any]] = deque() - self._added = 0 - self._emitted = 0 - - # -- ingest ------------------------------------------------------------------------- - def add(self, sample: Dict[str, Any], label: str = NEG) -> None: - self._added += 1 - if not self.balanced: - self._q.append(sample) - self._trim(self._q) - return - if label == POS: - self._pos.append(sample) - self._trim(self._pos) - elif label == NEG: - self._neg.append(sample) - self._trim(self._neg) - else: - raise ValueError(f'label must be {POS!r} or {NEG!r}, got {label!r}') - - def add_many(self, samples: List[Dict[str, Any]], label: str = NEG) -> None: - for s in samples: - self.add(s, label) - - def _trim(self, q: Deque[Dict[str, Any]]) -> None: - if self.max_pool is not None: - while len(q) > self.max_pool: - q.popleft() # drop oldest to bound memory / avoid stale majority backlog - - def rebalance(self) -> int: - """Balanced mode: down-sample the majority queue to the minority size, discarding - the NEWEST surplus (this chunk's excess intake — the #18b "其余丢弃不积压" rule). - Called once per chunk; since intake is re-balanced every chunk, both queues stay - equal-length and no side ever backlogs. Returns the number of discarded samples. - No-op in plain mode.""" - if not self.balanced: - return 0 - target = min(len(self._pos), len(self._neg)) - dropped = 0 - for q in (self._pos, self._neg): - while len(q) > target: - q.pop() # newest first: the surplus was added this chunk - dropped += 1 - return dropped - - # -- state -------------------------------------------------------------------------- - def ready(self) -> bool: - if not self.balanced: - return len(self._q) >= self.batch_size - half = self.batch_size // 2 - return len(self._pos) >= half and len(self._neg) >= half - - def sizes(self) -> Dict[str, int]: - if not self.balanced: - return {'pool': len(self._q)} - return {'pos': len(self._pos), 'neg': len(self._neg)} - - @property - def total_added(self) -> int: - return self._added - - @property - def total_emitted(self) -> int: - return self._emitted - - # -- draw --------------------------------------------------------------------------- - def draw(self) -> List[Dict[str, Any]]: - """Pop exactly one batch; raises if not ready. Remainder stays pooled.""" - if not self.ready(): - raise RuntimeError('draw() called while pool not ready; guard with ready()') - if not self.balanced: - batch = [self._q.popleft() for _ in range(self.batch_size)] - else: - half = self.batch_size // 2 - batch = [self._pos.popleft() for _ in range(half)] - batch += [self._neg.popleft() for _ in range(half)] - self._emitted += len(batch) - return batch - - def draw_all_ready(self) -> List[List[Dict[str, Any]]]: - """Pop as many full batches as currently available (0 or more).""" - out = [] - while self.ready(): - out.append(self.draw()) - return out diff --git a/cookbook/exp/skill2lora/skill_ablate/rollouting.py b/cookbook/exp/skill2lora/skill_ablate/rollouting.py deleted file mode 100644 index e3e9fff93..000000000 --- a/cookbook/exp/skill2lora/skill_ablate/rollouting.py +++ /dev/null @@ -1,224 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Rollout / prompt primitives for the ablation package. - -Reuses v2 verbatim (imported, never edited): -- ``_run_samples`` sampler pad / take-first-N (Ray dp size rule), -- ``_skillgen_prompt`` query-only skill-gen (view B + all eval), -- ``_train_trajectory`` query-only train trajectory (view B, SFT samples, OPSD student), -- ``build_skill_solve_prompt`` / ``build_direct_prompt`` executor prompts, -- ``_parse_seq`` / ``_extract_skill`` / ``_clean_text`` / ``_answer_leaked`` / ``_empty_roll``, -- ``_regen_prompt`` improve-skill regeneration (has orig skill; view A improve_sft), -- style/thinking globals ``_SKILL_STYLE`` / ``SKILL_GEN_SYSTEM`` etc. - -Adds ONLY the view-A rubric-conditioned pieces that v2 lacks: -- ``rubric_skillgen_prompt(problem, rubric)``: skill-gen conditioned on query + rubric - diagnosis, NO prior skill (matches the RL flow step 3 "输入 query+rubric ... skillmodel - rollout"). Style-matched (narrative / pitfall) to the main line. -- ``rubric_train_trajectory(rec)``: train trajectory that REBUILDS the query+rubric prompt + - response, for view-A RL training only (train-with-rubric; the trajectory must match the - prompt the skills were SAMPLED under). The SFT side keeps the query-only - ``_train_trajectory`` so the SFT prompt口径 stays query-only (skill_quality_analysis.md #6). -- ``opsd_teacher_trajectory(rec)``: the OPSD teacher — EXACTLY the student's query-only - prompt with the rubric APPENDED TO THE SYSTEM prompt (设计 871 行 "将 rubric 信息额外加入到 - system prompt 中"). Teacher and student therefore differ ONLY by the appended rubric block - and score the SAME response tokens. (``rubric_skillgen_prompt`` is NOT used here: its - system prompt differs wholesale from ``SKILL_GEN_SYSTEM``, which would confound the - distillation signal with a prompt-style shift.) -""" -from typing import Any, Dict - -from twinkle.data_format import pack_user_data - -import train_skill_v2 as v2 -from train_skill_v2 import ( # noqa: F401 (re-exported for methods.py convenience) - _answer_leaked, - _clean_text, - _empty_roll, - _extract_skill, - _parse_seq, - _regen_prompt, - _run_samples, - _skillgen_prompt, - _train_trajectory, - build_direct_prompt, - build_skill_solve_prompt, -) - -# --- view-A rubric-conditioned skill-gen system prompts -------------------------------- -# 中文注释：view-A 的 rubric 条件 skill-gen 提示词。仿照 skill_quality_analysis.md 的 -# "rubric & skill" 模板，但去掉"你已经生成过一个 skill"的指涉（RL 线首步没有旧 skill，只有 -# query + 一个失败尝试的 rubric 诊断）。文体与主链路一致（narrative / pitfall），且强制"自持、 -# 不指向外部上下文"，因为下游 executor 看不到 rubric——指涉会导致幻觉。改进skill+sft 线另有旧 -# skill，走 v2 的 _regen_prompt（含 orig_skill 字段，即 777-797 模板的逐字英文版），不用这里的提示词。 -# narrative 版末尾拼入与 REGEN_SYSTEM 同一个 <skills> few-shot 例子（设计 793-796），锁定文体/长度 -# 分布与主链路一致；程序化提取而非复制，保证与冻结的 v2 逐字相同。 -# 2026-07-30：narrative 版补上收束纪律句的**指令**（pitfall 版一直有，v2 的 SKILL_GEN_SYSTEM / -# REGEN_SYSTEM 也有，只有这里漏了）。few-shot 例子里本来就带这句，但光靠示范不管用——E17 首个 -# run 实测出现率 0.000-0.006，而同期 E16（有指令）是 1.000。上一轮把这句评为干预优先级 #1， -# 详见 skill_quality_analysis.md「E17 reflexion 臂截断漂移归因」第四节。 -# -# 2026-07-30 第二次改（用户拍板）：从“连贯叙述”改为“问题-根因-规避”结构。 -# ★ 命名提醒：`--skill-style narrative` 这个名字从此名不副实（已不再要求叙事体）， -# 但没改：它进了 exp_dir / swanlab_exp / config 指纹，改名会让旧 run 数据对不上。 -# ★ 只改 _RUBRIC_SKILLGEN_NARRATIVE（E17 专用，训练与 eval 同一个）；v2.REGEN_SYSTEM / -# SKILL_GEN_SYSTEM 未动，所以 E1-E16 的 query-only 链路不受影响。 -_RUBRIC_SKILLGEN_NARRATIVE = """\ -You are a skill-generation model. Your <skills> block will be fed to a SEPARATE downstream executor model that must solve the problem on its own. The executor will NOT see your private reasoning or the analysis below — it only sees what is inside <skills>...</skills>. - -An expert rubric analysis of a failed attempt on THIS problem is provided to you. Work from it to derive a COMPLETE, CONCRETE and ACTIONABLE set of instructions for how to avoid going wrong on this problem, and put your full line of thinking inside the <skills> block. - -Then write the <skills> block following these rules: -- Structure it as problem -> location -> root cause -> countermeasure. Start by naming the failure mode(s) that ACTUALLY occurred on this problem. For each one: state WHERE it happens (which step, which formula or which theorem is being applied), WHAT goes wrong there, WHY it goes wrong (a misunderstanding of a specific concept / a wrong value substituted in / a missing precondition / ...), and WHAT to guarantee in order to avoid it ("to avoid this, make sure that when <condition>, you <action>"). -- Cover exactly the failures that are real — no more. If only ONE thing went wrong (e.g. the mathematics was sound and the attempt merely failed to finish), write about that one thing in depth and stop. NEVER invent extra failure modes, and never pad the list to match a pattern: a made-up warning plants a wrong formula in the executor's head. -- If there are several failure modes, walk them in the order the executor will meet them, with ordinals ("First, ...; Second, when moving on to <step>, ..."), so it reads as a checklist; with a single failure mode, no ordinals are needed. -- Be concrete and executable. Every countermeasure must name the object it applies to (the formula, the quantity, the case being split). Do NOT write advice that would read the same on any other problem. -- CRITICAL: Do NOT solve the problem, reveal/compute the final answer, or substitute the problem's specific given numbers. Leave ALL concrete numbers for the executor to compute. -- Self-contained: NEVER reference "the analysis", "the rubric", or "the previous attempt" — the executor cannot see them, such phrasings cause hallucination. Address the solver directly. -- Close by telling the solver to commit and emit the answer in one pass without hesitating, naming the failure that hesitating causes. End the block with this exact sentence: "Avoid re-checking loops; box a bare number as soon as it is computed." -- Keep it under about 300 words. - -Put ONLY the guidance inside <skills></skills>. - -Example (several real failure modes): -<skills> -The recurring problems on this problem are: miscounting because symmetric configurations are treated as distinct, applying a permutation formula where the objects are actually indistinguishable, and dropping the division that removes duplicates. First, when you set up the count, the failure appears at the moment you choose between a permutation and a combination: the wrong branch is taken because "distinguishable" is read off the surface wording instead of from whether swapping two objects yields a genuinely different configuration. To avoid this, before writing any formula, state explicitly for each set of objects whether swapping two of its members changes the configuration, and only then pick the formula. Second, when moving on to the total count, compute it as if every object were ordered and distinguishable, because that quantity is unambiguous; the error to guard against here is folding the symmetry correction into this step, which makes the correction impossible to audit later. Third, when you apply the symmetry correction, the failure is using the wrong duplication factor — it comes from counting how many objects look alike rather than how many orderings map to the same configuration. To avoid this, derive the factor by asking how many distinct orderings of the interchangeable choices give the identical configuration, and divide the total by exactly that. Finally, commit to the result and emit the answer in one pass without hesitating; hesitating here restarts the case split and burns the budget before any answer is produced. Avoid re-checking loops; box a bare number as soon as it is computed. -</skills> - -Example (a single real failure mode — the mathematics was sound, so nothing is invented): -<skills> -The one recurring problem on this problem is not mathematical: the derivation stays on the right track, but a correct intermediate result gets questioned instead of used, the same quantity is re-derived to double-check it, and the attempt is cut off before any answer is written. The failure appears after the setup is complete, at the moment the first candidate value is in hand; it happens because re-verifying feels safer than committing, yet every re-check repeats the same computation and produces nothing new. To avoid this, once an intermediate quantity is computed, treat it as settled and build the next step directly on it; choose one method at the start, stay on it, and write the final line as soon as the last quantity is evaluated. Avoid re-checking loops; box a bare number as soon as it is computed. -</skills> -""" - -_RUBRIC_SKILLGEN_PITFALL = """\ -You are a skill-generation model. A separate executor model will solve the problem; it only sees your <skills> block, NOT the analysis below. - -An expert rubric analysis of a failed attempt on THIS problem is provided. From it, pinpoint the single decisive way a solver goes wrong on this type of problem. Then, inside <skills></skills>, write under 90 words: -- WARNING: name that decisive mistake concretely, in self-contained first person (e.g. "I think the step most likely to go wrong is ..."), and say why it is wrong. -- INSTEAD: one or two sentences pointing to the correct turn (technique name + where to apply it), without solving the problem or revealing any numeric result. -- End with: "Avoid re-checking loops; box a bare number as soon as it is computed." -Hard rules: the block must be self-contained — never reference "the analysis", "the rubric" or "the previous attempt"; the executor cannot see them. -""" - -_RUBRIC_SKILLGEN_USER = """\ -Problem: -{problem} - -Expert rubric analysis of a failed attempt (for your eyes only; do NOT reference it in the skill): -{rubric} - -Now write the improved <skills> guidance:""" - -# --- code 任务（BigCodeBench）的 rubric 条件化 skill-gen ----------------------------------- -# 与数学版同一个骨架（问题→位置→根因→规避 / 只写真实发生的失败 / 自持不指涉 rubric）， -# 把"公式、定理、代入数字、boxed 裸数"换成"该调哪个 API、参数与返回形状、边界与异常、 -# 交付一个 code block"。收尾纪律句也换掉：代码域的对应失败不是"兜圈不给答案"，而是 -# "反复重写实现 / 输出解释与演示代码 / 改动给定签名"。 -# ★ 这一版的信息优势来自 rubric 里带**单测真实报错**（code_task.diag_segment）： -# bcb_eval0_probe 实测 rubric skill 0.513 vs query-only 0.382（+0.135, p=4e-5）。 -_CODE_RUBRIC_SKILLGEN = """\ -You are a skill-generation model. Your <skills> block will be fed to a SEPARATE downstream engineer model that must implement the function on its own. The engineer sees the same task description and the same required signature, but NOT your private reasoning or the analysis below — it only sees what is inside <skills>...</skills>. - -An expert review of a failed attempt at THIS task is provided to you, including the real error its unit tests produced. Work from it to derive a COMPLETE, CONCRETE and ACTIONABLE set of instructions for how to avoid going wrong on this task, and put your full line of thinking inside the <skills> block. - -Then write the <skills> block following these rules: -- Structure it as problem -> location -> root cause -> countermeasure. Start by naming the failure mode(s) that ACTUALLY occurred. For each one: state WHERE it happens (which step of the implementation, which library call, which returned object), WHAT goes wrong there, WHY it goes wrong (a wrong assumption about what an API returns / a keyword the API does not accept / an unhandled empty or missing-column input / the wrong object handed back to the caller / ...), and WHAT to guarantee in order to avoid it ("to avoid this, make sure that when <condition>, you <action>"). -- Cover exactly the failures that are real — no more. If only ONE thing went wrong, write about that one thing in depth and stop. NEVER invent extra failure modes: a made-up warning sends the engineer after an API that is not the problem. -- If there are several failure modes, walk them in the order the engineer will meet them, with ordinals ("First, ...; Second, when building the return value, ..."), so it reads as a checklist. -- Be concrete and executable. Every countermeasure must name the object it applies to: the library function, the argument, the returned type, the edge case, the exception. Do NOT write advice that would read the same on any other task. -- CRITICAL: Do NOT write the solution code and do NOT paste concrete literal values from this task. Name the API and describe the shape of the value it returns instead of writing the call out. -- Self-contained: NEVER reference "the analysis", "the review", "the test error" or "the previous attempt" — the engineer cannot see them, such phrasings cause hallucination. Address the engineer directly. -- Close by telling the engineer to deliver one single fenced code block that reproduces the given imports and signature verbatim, with no explanation, no demonstration call and no tests around it. -- Keep it under about 300 words. - -Put ONLY the guidance inside <skills></skills>. - -Example (several real failure modes): -<skills> -The recurring problems on this type of task are: handing back the wrong object to the caller, assuming a grouping call returns a plain container when it returns an indexed one, and crashing instead of returning a well-defined result when the input is empty. First, when you build the return value, the failure appears at the very last line: a plotting helper is asked for and the figure is returned instead of the axes it drew on, or a tuple is required and only its first element comes back. To avoid this, re-read the sentence in the task that names the output, and make the last line return exactly that many objects in exactly that order, taking the axes object from the plotting call itself rather than from the figure. Second, when you aggregate, the failure is treating the result of the grouping call as a list: it is an indexed object whose labels are the group keys, so positional access silently reads the wrong group. To avoid this, convert it explicitly with the accessor the library provides before you index into it, and sort by the key the task names rather than relying on insertion order. Third, when the input has no rows or the named column is absent, the failure is an exception escaping from the aggregation. To avoid this, decide up front which of the two the task demands — a defined empty result or a specific raised exception — and write that branch before the main computation. Finally, deliver one single fenced code block that reproduces the given imports and signature verbatim, with no explanation, no demonstration call and no tests around it. -</skills> - -Example (a single real failure mode — nothing is invented): -<skills> -The one recurring problem on this type of task is a keyword that the library function does not accept: the call is the right one for the job, but it is invoked with an argument name borrowed from a similar function in another module, so it raises before any of the logic runs. The failure appears at the single line that does the real work, and it happens because the argument list is recalled from memory instead of from the function being called. To avoid this, when you reach that call, pass only the arguments you are certain that exact function declares, prefer positional arguments for the ones the task names explicitly, and if a behaviour you need is not available as a keyword there, achieve it with a following operation instead of inventing a parameter. Everything else in this task is straightforward once the call succeeds. Deliver one single fenced code block that reproduces the given imports and signature verbatim, with no explanation, no demonstration call and no tests around it. -</skills> -""" - -_CODE_RUBRIC_SKILLGEN_USER = """\ -Task: -{problem} - -Expert review of a failed attempt, with the real unit-test error (for your eyes only; do NOT \ -reference it in the skill): -{rubric} - -Now write the <skills> guidance:""" - - -def rubric_skillgen_prompt(problem: str, rubric: str) -> Dict[str, Any]: - """View-A skill-gen conditioned on (problem + rubric diagnosis), NO prior skill. - - Style-matched to the main line via v2's ``_SKILL_STYLE`` global (set by main() from - ``--skill-style``). narrative -> narrative rubric prompt; pitfall -> pitfall rubric prompt. - code 任务只有 narrative 一版（E4/E17 都是 narrative；pitfall 未移植，落到同一个 prompt）。 - """ - if v2._TASK == 'code': - return {'messages': [ - {'role': 'system', 'content': _CODE_RUBRIC_SKILLGEN}, - {'role': 'user', 'content': _CODE_RUBRIC_SKILLGEN_USER.format( - problem=problem, rubric=rubric)}]} - sys_p = _RUBRIC_SKILLGEN_PITFALL if v2._SKILL_STYLE == 'pitfall' else _RUBRIC_SKILLGEN_NARRATIVE - return {'messages': [ - {'role': 'system', 'content': sys_p}, - {'role': 'user', 'content': _RUBRIC_SKILLGEN_USER.format(problem=problem, rubric=rubric)}]} - - -def rubric_train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: - """Train sample whose PROMPT is the query+rubric skill-gen prompt + the sampled response. - - Used by view-A RL (rl_ab / rl_err / reflexion) only: train WITH rubric in the prompt - (knowledge-transfer probe; eval is still query-only via v2 ``_skillgen_prompt``). The - rebuilt prompt matches the prompt the skills were SAMPLED under (on-policy consistency). - - response 段直接拼采样返回的 token（``rec['tokens']``），绝不 decode 后重新过模板 —— - 重渲染会把模板自己补的换行/EOS/空思考块训进去，详见 v2.build_train_feature 的注释。 - 只有合成文本（没有 tokens 的 SFT 记录）才回退到 messages 编码。 - """ - msgs = rubric_skillgen_prompt(rec['problem'], rec.get('rubric', ''))['messages'] - if rec.get('tokens'): - return v2.build_train_feature(msgs, rec['tokens']) - return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], - 'user_data': pack_user_data({'key_rounds': [len(msgs)]})} - - -# 中文注释：OPSD teacher 的特权信息块——按设计 871 行要求放进 SYSTEM prompt，且只做“追加”， -# 保证 teacher 与 student 的 prompt 仅差这一段 rubric（最小差异，蒸馏信号不混入提示词风格漂移）。 -_OPSD_TEACHER_SUFFIX = """ - -[Privileged context — an expert rubric analysis of a failed attempt on this problem. \ -It is visible ONLY to you in this forward pass; the downstream executor never sees it. \ -Use it to judge which guidance actually helps, but do NOT reference it explicitly.] -{rubric}""" - - -def opsd_teacher_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: - """OPSD teacher: student's query-only prompt + rubric appended to the SYSTEM prompt + the - SAME sampled response tokens. With an empty rubric the teacher degenerates to the student - (zero distillation pull), which is the safe behaviour on rubric API failure. - - teacher 与 student 必须打分**同一串 token**，所以两边的 response 段都直接拼 ``rec['tokens']``； - 只有 prompt 段不同（多一段 rubric）。 - """ - msgs = [dict(m) for m in _skillgen_prompt(rec['problem'])['messages']] - rubric = (rec.get('rubric') or '').strip() - if rubric and msgs[0]['role'] == 'system': - msgs[0]['content'] = msgs[0]['content'] + _OPSD_TEACHER_SUFFIX.format(rubric=rubric) - if rec.get('tokens'): - return v2.build_train_feature(msgs, rec['tokens']) - return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], - 'user_data': pack_user_data({'key_rounds': [len(msgs)]})} - - -def query_only_train_trajectory(rec: Dict[str, Any]) -> Dict[str, Any]: - """Alias for v2's query-only train trajectory (view B, SFT samples, OPSD student).""" - return _train_trajectory(rec) diff --git a/cookbook/exp/skill2lora/skill_ablate/rubric_cache.py b/cookbook/exp/skill2lora/skill_ablate/rubric_cache.py deleted file mode 100644 index 60048a656..000000000 --- a/cookbook/exp/skill2lora/skill_ablate/rubric_cache.py +++ /dev/null @@ -1,128 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Rubric double-cache (skill_quality_analysis.md #4, #17). - -Two cache scopes, both wrapping v2's ``DiskCache`` (append-only jsonl, in-memory index): - -- GlobalRubricCache (key = data_id): the RL / SFT lines diagnose the BARE-PROBLEM greedy - trajectory. Because the executor is frozen at T=0, that trajectory is deterministic and - identical across experiments, so its rubric diagnosis can be shared across ALL runs via one - global file (``rubric_cache_global.jsonl``) — diagnose each problem once, reuse everywhere. - ⚠️ 该共享前提只在"executor 口径完全相同"时成立。轨迹本身**不在键里**，所以任何改变裸解 - 轨迹的开关（task 域、executor thinking 开关、以后若改 executor 模型）都必须体现在**文件名** - 上，否则新口径的 run 会命中旧口径的诊断（键只有 data_id，必然命中）。实测这个文件已经攒了 - 2630 条 think 轨迹的数学诊断 —— E19（executor nothink）若共用会几乎全程读到与自己失败无关 - 的诊断，而 rubric 内容正是该臂唯一的自变量。见 build_rubric_cache 的 tag 拼装。 - -- LocalRubricCache (key = md5(data_id + skill)): the improve-skill+SFT / OPSD lines diagnose a - WITH-SKILL trajectory whose skill evolves with the policy, so the diagnosis is experiment- - and step-specific and lives only in that experiment's directory. - -Both reuse v2's ``_diagnose_entry`` (pure teacher-API call, no GPU) for the actual diagnosis, -so there is a single source of truth for the rubric prompt / parsing. -""" -import os -from typing import Any, Dict, Optional - -import train_skill_v2 as v2 -from train_skill_v2 import DiskCache, _diagnose_entry - - -def _version() -> str: - """动态读 v2._RUBRIC_VERSION —— 不能 from-import：v2.set_task('code') 会在运行时把它换成 - code 判据的版本号，而 from-import 会把加载那一刻的值钉死，导致代码域诊断用数学域的键。""" - return v2._RUBRIC_VERSION - - -class _BaseRubricCache: - """Shared get-or-diagnose logic over a DiskCache; subclasses define the key.""" - - def __init__(self, path: str, enabled: bool = True): - self._cache = DiskCache(path, enabled) - - def _key(self, entry: Dict[str, Any], skill: Optional[str]) -> str: - raise NotImplementedError - - def get(self, entry: Dict[str, Any], skill: Optional[str] = None) -> Optional[str]: - return self._cache.get(self._key(entry, skill)) - - def get_or_diagnose(self, entry: Dict[str, Any], checker, - skill: Optional[str] = None) -> str: - """Return cached diagnosis, else run the teacher rubric once and cache it. - - ``entry`` must carry the fields ``_diagnose_entry`` needs: ``problem``, - ``fail_segment``, ``fail_stop_reason`` (and ``reference_answer`` is unused by the - diagnosis but kept for auditing). Returns '' when there is no checker or on API error - (never raises), so the caller can treat "no diagnosis" uniformly. - """ - if checker is None: - return '' - key = self._key(entry, skill) - hit = self._cache.get(key) - # bugfix #1: 旧版把 API 失败（_diagnose_entry 返回 None）也以 '' 永久写进缓存， - # 一次瞬时抖动会跨实验毒化全局缓存且永不重试。现在：只缓存真诊断；历史残留的 - # '' 条目视为 miss，下次调用自动重试并用真诊断覆盖。 - if hit: - return hit - diag = _diagnose_entry(checker, entry) - if diag is None: # transient API failure: do NOT cache, retry on the next call - return '' - self._cache.put(key, diag) - return diag - - def put(self, entry: Dict[str, Any], diag: str, skill: Optional[str] = None) -> None: - self._cache.put(self._key(entry, skill), diag) - - def __contains__(self, key: str) -> bool: - return key in self._cache - - def close(self) -> None: - self._cache.close() - - -class GlobalRubricCache(_BaseRubricCache): - """key = (rubric 版本, data_id): bare-problem trajectory diagnosis, shareable across experiments. - - 版本号必须进键：该文件跨实验共享且 append-only，一旦判据表改了而键不变，旧 - taxonomy 的诊断会被静默当成新判据的结果返回（旧版本号仅定义未使用）。 - """ - - def _key(self, entry: Dict[str, Any], skill: Optional[str] = None) -> str: - return DiskCache.key_for('rubric_global', _version(), - str(entry.get('data_id', ''))) - - -class LocalRubricCache(_BaseRubricCache): - """key = md5(rubric 版本 + data_id + skill): with-skill trajectory diagnosis, per-experiment only.""" - - def _key(self, entry: Dict[str, Any], skill: Optional[str] = None) -> str: - return DiskCache.key_for('rubric_local', _version(), - str(entry.get('data_id', '')), skill or '') - - -def build_rubric_cache(scope: str, output_dir: str, global_dir: Optional[str] = None, - enabled: bool = True, task: str = 'math', executor_thinking: str = 'on'): - """Factory: scope='global' -> shared file under ``global_dir`` (default output_dir/..); - scope='local' -> per-experiment file under ``output_dir/cache``. - - ``task`` 进文件名：代码域与数学域的诊断内容完全不同源（判据表、judge prompt、segment 里 - 有没有单测报错），版本号已经能隔开键，分文件是第二道保险，也让缓存体积可分别管理。 - - ★ ``executor_thinking`` 也必须进文件名（2026-07-31 bugfix）：global 缓存跨实验共享的**唯一 - 依据**是"executor 冻结在 T=0，所以同一道题的裸解轨迹在所有实验里逐字相同"。E19/E20 把 - executor 的 thinking 关掉后这个前提就不成立了 —— 裸解轨迹完全变了（think 那边大量是 - "撞预算没写出代码"，nothink 这边是"写完但答错"），而诊断正是对着这条轨迹做的。共用一个 - 文件会让 nothink 臂直接读到 think 臂的旧诊断（键只有 data_id，必然命中）， - rubric 内容与本臂的真实失败无关 —— 而 rubric 内容恰恰是这两个臂唯一的自变量。 - """ - tag = '' if task == 'math' else f'_{task}' - if executor_thinking != 'on': - tag += '_execnothink' - if scope == 'global': - base = global_dir or os.path.dirname(os.path.abspath(output_dir.rstrip('/'))) - os.makedirs(base, exist_ok=True) - return GlobalRubricCache(os.path.join(base, f'rubric_cache_global{tag}.jsonl'), enabled) - if scope == 'local': - cache_dir = os.path.join(output_dir, 'cache') - os.makedirs(cache_dir, exist_ok=True) - return LocalRubricCache(os.path.join(cache_dir, f'rubric_cache_local{tag}.jsonl'), enabled) - raise ValueError(f"scope must be 'global' or 'local', got {scope!r}") diff --git a/cookbook/exp/skill2lora/skill_ablate/trainer.py b/cookbook/exp/skill2lora/skill_ablate/trainer.py deleted file mode 100644 index 506398de9..000000000 --- a/cookbook/exp/skill2lora/skill_ablate/trainer.py +++ /dev/null @@ -1,455 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Unified training loop for one ablation experiment. - -step unit = PARAMETER UPDATE (skill_quality_analysis.md #2): a chunk may yield 0/1/more -updates; we stop at ``--max-updates`` and eval every ``--eval-every-updates`` updates. All -eval is query-only via v2 ``run_greedy_eval`` (T=0.5 × 4 rollouts, no rubric) — the -knowledge-transfer probe for view A. Records/log schema mirror v2 (gen/eval/train_log jsonl, -per-problem rollout rows land in gen_records.jsonl). - -swanlab step axis: ONE global axis = chunk index for train AND eval curves (eval also logs -``eval/updates_done`` so the update count is recoverable); mixing chunk/update axes in one -experiment made curves incomparable. -""" -import json -import os -import sys -import time -from typing import Any - -import train_skill_v2 as v2 - -from .config import ExpSpec -from .data import load_deepmath_records -from .data_code import load_code_records -from .eval_reflexion import run_reflexion_eval -from .methods import MethodContext, build_method -from .pool import SamplePool -from .rubric_cache import build_rubric_cache - -try: - import swanlab -except ImportError: - swanlab = None - - -def _rubric_scope(method: str) -> str: - """Bare-problem lines (rl/sft) share a GLOBAL cache; with-skill lines (opsd/improve) use - a per-experiment LOCAL cache. bnpo needs no rubric.""" - if method in ('rl_ab', 'rl_err', 'sft', 'reflexion', 'rejection_sft'): - return 'global' - if method in ('opsd', 'improve_sft', 'logp_rl'): - return 'local' - return '' - - -def _build_pool(spec: ExpSpec, args) -> Any: - if spec.method == 'improve_sft': - return SamplePool(batch_size=args.sft_batch_size, balanced=True, - max_pool=args.pool_max) - if spec.method == 'sft': - return SamplePool(batch_size=args.sft_batch_size, balanced=False, - max_pool=args.pool_max) - if spec.method == 'rejection_sft': - # E18：攒够 --e18-accumulate（128）条才 fire 一次 SFT；_train_batch 内部仍按 - # sft_batch_size 切 micro，所以这里只需保证是 TRAIN_DP 倍数（main.py 校验）。 - return SamplePool(batch_size=args.e18_accumulate, balanced=False, - max_pool=args.pool_max) - return None # RL / OPSD / bnpo train per-chunk, no accumulation pool - - -def _load_resume_state(args) -> dict: - """Resolve --resume-from into {'ckpt_dir', 'updates', 'chunk_idx'}. - - Weights-only resume (lr is constant; Adam moments restart — accepted trade-off). - Counter source: <ckpt>/train_state.json (written by _save_ckpt); legacy finished runs - (e.g. E1-final) fall back to <output_dir>/DONE.json {'updates','chunks'}. - """ - ck = args.resume_from - if not os.path.isdir(ck): - ck = os.path.join(args.output_dir, args.resume_from) - if not os.path.isdir(ck): - raise FileNotFoundError(f'--resume-from checkpoint dir not found: {args.resume_from}') - state_path = os.path.join(ck, 'train_state.json') - done_path = os.path.join(args.output_dir, 'DONE.json') - if os.path.exists(state_path): - with open(state_path, encoding='utf-8') as f: - st = json.load(f) - # bugfix #17: 旧/手写 train_state.json 缺键时给出可读报错而非裸 KeyError - missing = [k for k in ('updates', 'chunk_idx') if k not in st] - if missing: - raise ValueError(f'resume: {state_path} missing keys {missing}; ' - f'present keys: {sorted(st)}') - for k in ('chunk_size', 'min_level', 'n', 'seed'): - if k in st and getattr(args, k, None) not in (None, '') and st[k] != getattr(args, k): - sys.stderr.write(f'[ablate] WARNING: resume data config mismatch: {k} ' - f'ckpt={st[k]} vs now={getattr(args, k)} — the continued chunk ' - f'sequence will NOT align with the original run.\n') - elif os.path.exists(done_path): - with open(done_path, encoding='utf-8') as f: - d = json.load(f) - st = {'updates': int(d['updates']), 'chunk_idx': int(d['chunks'])} - sys.stderr.write('[ablate] resume: no train_state.json in ckpt, counters restored ' - 'from DONE.json (legacy run — data-config alignment unverified).\n') - else: - raise FileNotFoundError(f'resume: neither {state_path} nor {done_path} exists; ' - 'cannot restore update/chunk counters.') - return {'ckpt_dir': ck, 'updates': int(st['updates']), 'chunk_idx': int(st['chunk_idx'])} - - -def run_experiment(args, spec: ExpSpec) -> None: - # 0) idempotency: DONE.json is written atomically as the very last step of a successful - # run; if present, this experiment is complete -> skip (unless --force). --resume-from - # bypasses the guard by design: its whole point is extending a finished run. - resume = _load_resume_state(args) if getattr(args, 'resume_from', '') else None - done_path = os.path.join(args.output_dir, 'DONE.json') - if os.path.exists(done_path) and not getattr(args, 'force', False) and resume is None: - sys.stderr.write(f'[ablate] {spec.name} already complete ({done_path}); ' - f'use --force to rerun.\n') - return - if resume is not None: - args.skill_init_model_id = resume['ckpt_dir'] # v2.init_components skill_model bypass - sys.stderr.write(f'[ablate] resuming {spec.name} from {resume["ckpt_dir"]} ' - f'(updates={resume["updates"]} chunk={resume["chunk_idx"]}).\n') - - # 1) task / style / align globals must be set BEFORE any prompt is built. - # set_task 必须在最前：它同时决定 prompt 分派、判分方式与 rubric 判据版本（缓存键）。 - v2.set_task(spec.task, getattr(args, 'test_workers', 24), getattr(args, 'test_timeout', 60)) - v2._ALIGN_MODE = spec.align - v2._SKILL_STYLE = spec.style - args.skill_thinking = spec.thinking - # per-style length budget (#9 statistics: narrative≈1100 / pitfall≈300 chars); - # freeform 可能产出叙述式长文本，按 narrative 档给 1100 以免误伤；explicit --len-budget on the CLI wins. - if args.len_budget is None: - args.len_budget = 1100 if spec.style in ('narrative', 'freeform') else 300 - # explicit --skill-max-tokens on the CLI wins over the per-experiment default. - if args.skill_max_tokens is None: - args.skill_max_tokens = spec.skill_max_tokens - elif args.skill_max_tokens != spec.skill_max_tokens: - sys.stderr.write(f'[ablate] WARNING: --skill-max-tokens {args.skill_max_tokens} ' - f'overrides the {spec.name} default {spec.skill_max_tokens}.\n') - # E14 信噪比消融：训练判分 rollout 数/温度，CLI 显式值优先，否则取 spec（E1-E13=1×T0）。 - if getattr(args, 'reward_rollouts', None) is None: - args.reward_rollouts = spec.reward_rollouts - if getattr(args, 'reward_temperature', None) is None: - args.reward_temperature = spec.reward_temperature - # OOM guard: think/8192 实验的训练序列长一倍，fp32 主权重下 16/2卡 的 micro backward - # 会爆显存（E13 实测 Tried to allocate 37.9GiB）；自动把 micro 减半到 8，梯度归一后数学等价， - # 攒批/采样批（sft_batch_size）冻结口径不变。显式 --train-micro-batch 优先。 - if not args.train_micro_batch and args.skill_max_tokens >= 8192: - args.train_micro_batch = max(v2.TRAIN_DP, args.sft_batch_size // 2) - sys.stderr.write(f'[ablate] train_micro_batch auto-set to {args.train_micro_batch} ' - f'(skill_max_tokens={args.skill_max_tokens} OOM guard).\n') - - if spec.task == 'code': - records, eval_records = load_code_records(args) - else: - records, eval_records = (load_deepmath_records(args) if getattr(args, 'deepmath_dir', '') - else v2._load_records(args)) - # ⭐ --train-order-file：用 SEAM 已实现的 batch 序列覆盖 train 分支（eval 划分不动 —— twinkle 的 - # eval 128 题已逐题验证与 SEAM val.parquet 一致）。verl 的 dataloader 会 shuffle，所以两边 - # 即使同题池、同 batch size，step k 实际喂的 128 题也几乎不重叠（实测交集 1/128）， - # 单此一项就能把 acc/baseline/lift 拉开 6-7 个点。给了 order 文件后 chunk k 逐题 == SEAM step k+1。 - order_file = (getattr(args, 'train_order_file', '') or '').strip() - if order_file: - records = v2.load_train_order_file(order_file) - ev = {r['problem'] for r in eval_records} - overlap = sum(1 for r in records if r['problem'] in ev) - if overlap: - raise ValueError(f'--train-order-file overlaps eval on {overlap} rows: {order_file}') - sys.stderr.write(f'[ablate] train order pinned to {order_file}: {len(records)} rows ' - f'({len(records) // max(1, args.chunk_size)} chunks of {args.chunk_size}), ' - f'ProblemPool shuffle DISABLED.\n') - if len(records) < args.chunk_size: - raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded ({len(records)})') - - os.makedirs(args.output_dir, exist_ok=True) - gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') - eval_path = os.path.join(args.output_dir, 'eval_records.jsonl') - train_log_path = os.path.join(args.output_dir, 'train_log.jsonl') - - # 2) rubric checker BEFORE any GPU allocation: view A without a teacher API cannot run - # (SFT-family would loop forever on an empty pool, OPSD would distill on empty rubrics). - checker = v2.build_rubric_checker() if spec.needs_rubric else None - if spec.needs_rubric and checker is None: - raise RuntimeError( - f'{spec.name} ({spec.method}) is a view-A experiment and REQUIRES the rubric ' - 'teacher API; set LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL (or OPENAI_API_KEY).') - - # 3) components (v2 verbatim); override loss to OPSD when needed. - skill_model, ref_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp = \ - v2.init_components(args) - if spec.loss == 'opsd': - # no beta: OPSDLoss uses only teacher_logps (no ref-KL term, no ref forward at all) - skill_model.set_loss('OPSDLoss', reverse=True) - - scope = _rubric_scope(spec.method) - rubric_cache = build_rubric_cache(scope, args.output_dir, - global_dir=args.rubric_global_dir, - enabled=not args.no_cache, - task=spec.task, - executor_thinking=spec.executor_thinking) if scope else None - - # client-side Template clone: - # - OPSD: skill-model template, used to align response-token positions for teacher logps. - # - logp_rl / logp_gt: executor template (thinking on), used to slice logP(S | executor prompt). - encode_template = None - if spec.method in ('opsd', 'logp_rl', 'logp_gt'): - encode_template = v2.Template(model_id=v2.MODEL_ID, - enable_thinking=(True if spec.method in ('logp_rl', 'logp_gt') - else spec.thinking == 'on'), - max_length=args.max_model_len, - truncation_strategy='delete') - # 所有臂都需要的 skill 侧模板副本：训练样本的 prompt 段在客户端编码，response 段直接拼 - # 采样返回的 token（见 v2.build_train_feature）。必须与 skill_model/skill_sampler 同配置， - # 否则 prompt 段的 token 会与采样时对不上。 - v2.set_encode_template(v2.Template(model_id=v2.MODEL_ID, - enable_thinking=(spec.thinking == 'on'), - max_length=args.max_model_len, - truncation_strategy='delete')) - - pool = _build_pool(spec, args) - ctx = MethodContext(skill_model=skill_model, ref_model=ref_model, skill_sampler=skill_sampler, - base_sampler=base_sampler, ckpt=ckpt, skill_dp=skill_dp, base_dp=base_dp, - args=args, checker=checker, rubric_cache=rubric_cache, pool=pool, - encode_template=encode_template) - method = build_method(spec.method, ctx) - - def _save_ckpt(name: str, updates: int, chunk_idx: int, epoch: int) -> None: - """Weights-only checkpoint + barrier + resume state. - - skill_model.save dispatches to the train actors; a subsequent cheap blocking call on - the SAME actors (lr_step — a guaranteed no-op here: no scheduler, constant lr) acts as - the barrier: Ray actor tasks run serially per actor, so when it returns the save has - landed on every rank. Without it the driver could exit on a half-written safetensors. - """ - skill_model.save(name, output_dir=args.output_dir) - skill_model.lr_step() # barrier (see docstring) - state = {'updates': updates, 'chunk_idx': chunk_idx, 'epoch': epoch, - 'seed': args.seed, 'chunk_size': args.chunk_size, 'n': args.n, - 'min_level': int(getattr(args, 'min_level', 0) or 0), - 'lr': args.lr, 'exp': spec.name, 'saved': int(time.time())} - with open(os.path.join(args.output_dir, name, 'train_state.json'), 'w', - encoding='utf-8') as f: - json.dump(state, f) - sys.stderr.write(f'[ablate] checkpoint saved: {name} (updates={updates})\n') - - # 4) eval baseline cache (v2 DiskCache) + swanlab. - # 每次启动强制重算 eval baseline：旧缓存可能来自不同环境/代码版本（torch/vllm/dtype 均影响 T=0 输出）， - # 跨 run 复用会造成 with-skill（现算）vs baseline（陈旧）不可比，lift 虚高/虚低。 - cache_dir = os.path.join(args.output_dir, 'cache') - os.makedirs(cache_dir, exist_ok=True) - _base_cache_path = os.path.join(cache_dir, 'eval_baseline.jsonl') - if os.path.exists(_base_cache_path): - os.remove(_base_cache_path) - sys.stderr.write('[ablate] stale eval_baseline cache removed (recomputed this run).\n') - eval_base_cache = v2.DiskCache(_base_cache_path, not args.no_cache) - - use_swan = swanlab is not None and os.environ.get('SWANLAB_MODE') != 'disabled' - if use_swan: - # timestamp suffix so FORCE reruns never collide in swanlab (接口方案 #9); - # --run-tag 用于区分共用同一 ExpSpec 的变体（如 kl_beta 0.001 vs 0.01）。 - _tag = (getattr(args, 'run_tag', '') or '').strip() - swan_exp = (f'{spec.swanlab_exp}' + (f'_{_tag}' if _tag else '') - + f'_{time.strftime("%Y%m%d_%H%M%S")}') - swanlab.init(project=args.swanlab_project, experiment_name=swan_exp, - config={'exp': spec.name, 'view': spec.view, 'method': spec.method, - 'thinking': spec.thinking, 'style': spec.style, 'align': spec.align, - 'loss': spec.loss, 'skill_max_tokens': args.skill_max_tokens, - 'max_updates': args.max_updates, 'lr': args.lr, - 'n_skills': args.n_skills, 'sft_batch_size': args.sft_batch_size, - 'len_budget': args.len_budget, - 'run_tag': _tag, - 'reward_rollouts': args.reward_rollouts, - 'reward_temperature': args.reward_temperature, - 'kl_beta': getattr(args, 'kl_beta', 0.01), - 'grpo_epsilon': getattr(args, 'grpo_epsilon', 0.2), - 'adv_clip': getattr(args, 'adv_clip', 0.0), - 'logp_leak_penalty': getattr(args, 'logp_leak_penalty', 0.0), - 'reward_leak_gate': getattr(args, 'reward_leak_gate', 0.0), - 'reward_trunc_penalty': getattr(args, 'reward_trunc_penalty', 0.25), - 'reward_trunc_lo': getattr(args, 'reward_trunc_lo', 6000), - 'base_tok_floor': getattr(args, 'base_tok_floor', 5000), - 'drop_zero_adv': args.drop_zero_adv}) - - cfg = {'record_type': 'config', 'exp': spec.name, 'task': spec.task, 'view': spec.view, - 'method': spec.method, - 'thinking': spec.thinking, 'style': spec.style, 'align': spec.align, 'loss': spec.loss, - 'executor_thinking': spec.executor_thinking, - 'skill_max_tokens': args.skill_max_tokens, 'needs_rubric': spec.needs_rubric, - 'rubric_scope': scope, 'rubric_check': bool(checker), - 'n': len(records), 'eval_n': len(eval_records), 'model': v2.MODEL_ID, - 'max_updates': args.max_updates, 'eval_every_updates': args.eval_every_updates, - 'lr': args.lr, 'n_skills': args.n_skills, 'chunk_size': args.chunk_size, - 'sft_batch_size': args.sft_batch_size, 'len_budget': args.len_budget, - 'skill_char_limit': args.skill_char_limit, 'drop_zero_adv': args.drop_zero_adv, - 'improve_skill_temperature': args.improve_skill_temperature, - 'reward_rollouts': args.reward_rollouts, - 'reward_temperature': args.reward_temperature, - 'run_tag': (getattr(args, 'run_tag', '') or ''), - # loss 侧旋钮同样入账：kl_beta 是唯一对抗漂移的恢复力，此前不落盘导致已跑的臂 - # 无法从 gen_records 反推当时的锚强度（全部是旧默认 0.001；现默认 0.01）。 - 'kl_beta': getattr(args, 'kl_beta', 0.01), - 'grpo_epsilon': getattr(args, 'grpo_epsilon', 0.2), - 'adv_clip': getattr(args, 'adv_clip', 0.0), - # reward 公式的所有旋钮全部入账：之前 leak 惩罚默认开着却不落盘，导致已跑的臂无法 - # 从 gen_records 反推当时用的是哪个 reward。leak 一律不进 reward，此处应恒为 0。 - 'logp_leak_penalty': getattr(args, 'logp_leak_penalty', 0.0), - 'reward_leak_gate': getattr(args, 'reward_leak_gate', 0.0), - 'reward_trunc_penalty': getattr(args, 'reward_trunc_penalty', 0.25), - 'reward_trunc_lo': getattr(args, 'reward_trunc_lo', 6000), - 'base_tok_floor': getattr(args, 'base_tok_floor', 5000), - 'reflexion_k': int(getattr(args, 'reflexion_k', 0) or 0), - 'eval_protocol': ('reflexion' if spec.method == 'reflexion' else 'query_only'), - 'eval_rollouts': args.eval_rollouts, 'eval_skill_temperature': args.eval_skill_temperature, - 'seam_parquet_dir': (getattr(args, 'seam_parquet_dir', '') or ''), - 'deepmath_dir': (getattr(args, 'deepmath_dir', '') or ''), - 'min_level': int(getattr(args, 'min_level', 0) or 0), - 'eval_min_level': int(getattr(args, 'eval_min_level', 0) or 0), - 'save_every_updates': int(getattr(args, 'save_every_updates', 0) or 0), - 'resumed_from': (resume['ckpt_dir'] if resume else ''), - 'resumed_updates': (resume['updates'] if resume else 0), - 'started': int(time.time())} - - # resume appends to the record files (the original history stays intact); a fresh run - # truncates as before. - _fmode = 'a' if resume is not None else 'w' - with open(gen_path, _fmode, encoding='utf-8') as gen_f, \ - open(eval_path, _fmode, encoding='utf-8') as eval_f, \ - open(train_log_path, _fmode, encoding='utf-8') as tlog: - for f in (gen_f, eval_f, tlog): - v2._write(f, cfg) - - def _do_eval(updates_done: int, swan_step: int) -> None: - # E17 用 reflexion 协议 eval（只干预裸解做错的题、skill-gen 带 rubric），与训练 - # 同分布；其余臂一律走 v2 的 query-only 全量 eval。两者 summary 键名兼容。 - if spec.method == 'reflexion': - recs, summary, metrics = run_reflexion_eval( - base_sampler, skill_sampler, eval_records, updates_done, updates_done, - base_dp, skill_dp, args, eval_base_cache, rubric_cache, checker) - elif spec.method == 'rejection_sft': - # E18：eval 用 nothink rollout，且与采集共用同一个 skill_sampler vLLM（用户 - # 2026-07-30 拍板）。set_template 只换客户端编码（sampler/base.py:106）， - # 引擎不重建；训练响应是裸 <skills>（空 think 布局），与 nothink 生成布局 - # 逐 token 一致，所以这才是本臂的同分布读数。finally 必须切回 think， - # 否则下一个 chunk 的采集会静默变成 nothink。 - skill_sampler.set_template(v2.Template, model_id=v2.MODEL_ID, - enable_thinking=False, - max_length=args.max_model_len) - try: - recs, summary, metrics = v2.run_greedy_eval( - base_sampler, skill_sampler, eval_records, updates_done, updates_done, - base_dp, skill_dp, args, eval_base_cache) - finally: - skill_sampler.set_template(v2.Template, model_id=v2.MODEL_ID, - enable_thinking=(spec.thinking == 'on'), - max_length=args.max_model_len) - else: - recs, summary, metrics = v2.run_greedy_eval( - base_sampler, skill_sampler, eval_records, updates_done, updates_done, - base_dp, skill_dp, args, eval_base_cache) - for rec in recs: - v2._write(eval_f, rec) - v2._write(eval_f, summary) - eval_f.flush() - if use_swan: - # same chunk-based axis as the train curves; updates recoverable via the - # logged eval/updates_done scalar. - swanlab.log({**{f'eval/{k}': v for k, v in metrics.items()}, - 'eval/updates_done': float(updates_done)}, step=swan_step) - sys.stderr.write( - f'[eval] u{updates_done}: n={summary["n"]} acc={summary["baseline_acc_mean1"]:.3f}' - f'->{summary["acc_mean1"]:.3f} lift={summary["lift_mean1"]:+.3f} ' - f'hard_rescue={summary["hard_rescue_rate"]:.3f} fmt={summary["format_mean1"]:.2f}\n') - - if eval_records and resume is None: - _do_eval(-1, 0) # baseline before any update (chunk axis position 0) - - pool_pp = v2.ProblemPool(records, args.seed, - fixed_order=bool((getattr(args, 'train_order_file', '') or '').strip())) - updates = 0 - last_eval_at = 0 - last_eval_updates = -1 - chunk_idx = 0 - last_save_at = 0 - if resume is not None: - # push the resumed weights into skill_sampler BEFORE the first chunk (otherwise - # skill-gen would sample from base weights until the first post-update sync). - ckpt.sync_weights(merge_and_sync=True) - updates = resume['updates'] - last_eval_at = updates - last_save_at = updates - # fast-forward the pool: draws are deterministic (RandomState(seed+epoch)), so - # replaying chunk_idx draws restores the exact data position — IF chunk_size / - # data config match the original run (warned in _load_resume_state). - for _ in range(resume['chunk_idx']): - pool_pp.draw(args.chunk_size) - chunk_idx = resume['chunk_idx'] - save_every = int(getattr(args, 'save_every_updates', 0) or 0) - while updates < args.max_updates: - chunk = pool_pp.draw(args.chunk_size) - res = method.step(chunk, chunk_idx) - n_upd = int(res.get('n_updates', 0)) - updates += n_upd - - for rec in res.get('gen_records') or []: - v2._write(gen_f, rec) - gen_f.flush() - - log = {'record_type': 'train_round', 'exp': spec.name, 'chunk': chunk_idx, - 'epoch': pool_pp.epoch, 'updates': updates, 'n_updates_step': n_upd, - 'method': spec.method, 'ts': int(time.time()), - 'metrics': res.get('metrics', {})} - if 'summary' in res: # bnpo carries the v2 chunk summary - log['summary'] = res['summary'] - v2._write(tlog, log) - tlog.flush() - - sys.stderr.write(f'[gen] e{pool_pp.epoch} c{chunk_idx}: +{n_upd}upd ' - f'total={updates}/{args.max_updates} ' - + ' '.join(f'{k}={v:.3g}' for k, v in res.get('metrics', {}).items() - if isinstance(v, (int, float))) + '\n') - if use_swan: - m = {f'{k}': float(v) for k, v in res.get('metrics', {}).items() - if isinstance(v, (int, float))} - m['train/updates'] = float(updates) - m['train/n_updates_step'] = float(n_upd) - swanlab.log(m, step=chunk_idx + 1) # +1: step 0 is the eval baseline - - # eval by parameter-update cadence (same chunk-based swan axis) - if eval_records and updates >= last_eval_at + args.eval_every_updates and n_upd > 0: - _do_eval(updates, chunk_idx + 1) - last_eval_at = updates - last_eval_updates = updates - # periodic weights-only checkpoint (same cadence semantics as eval) - if save_every and updates >= last_save_at + save_every and n_upd > 0: - _save_ckpt(f'{spec.name}-u{updates}', updates, chunk_idx + 1, pool_pp.epoch) - last_save_at = updates - # bugfix #12: 方法把候选/rollout 全文挂在共享 record dict 上（ProblemPool 跨 epoch - # 复用同一批对象），gen_records 已落盘后不清理会让 driver RAM 随触达题数线性增长 - # （5000 题 × 8 候选 × 几十 KB 可达数 GB）。训练/日志都结束后安全清理。 - for r in chunk: - for k in ('_cands', '_pseudo_rolls', '_pseudo_roll', '_pseudo_solution', - '_rubric', '_base_tok', '_base_correct', '_base_stop'): - r.pop(k, None) - chunk_idx += 1 - - # final readout — skip if the periodic eval already covered this exact update count - if eval_records and updates != last_eval_updates: - _do_eval(updates, chunk_idx + 1) - - eval_base_cache.close() - if rubric_cache is not None: - rubric_cache.close() - # final save goes through _save_ckpt: barrier guarantees the safetensors is fully on disk - # BEFORE DONE.json can exist, and train_state.json makes the final model resumable too. - _save_ckpt(f'{spec.name}-final', updates, chunk_idx, pool_pp.epoch) - # completion sentinel LAST (after the final model lands); temp+rename keeps it atomic so a - # crash can never leave a truthy half-written marker. - tmp = done_path + '.tmp' - with open(tmp, 'w', encoding='utf-8') as f: - json.dump({'exp': spec.name, 'updates': updates, 'chunks': chunk_idx, - 'epochs': pool_pp.epoch, 'finished': int(time.time())}, f) - os.replace(tmp, done_path) - sys.stderr.write(f'[ablate] {spec.name} done: {updates} updates over {chunk_idx} chunks / ' - f'{pool_pp.epoch} epochs\n') diff --git a/cookbook/exp/skill2lora/skill_feature_corr.py b/cookbook/exp/skill2lora/skill_feature_corr.py deleted file mode 100644 index 46ec5e814..000000000 --- a/cookbook/exp/skill2lora/skill_feature_corr.py +++ /dev/null @@ -1,128 +0,0 @@ -#!/usr/bin/env python3 -"""skill_feature_corr.py — 探"什么样的 skill 文本特征决定 executor 好坏"。 - -对象：logp_corr/ 已有的 476 对 (题,skill) + 8-rollout 真值（pass_rate / trunc_rate / mean_tokens）。 -在 skill 文本上抽一批可解释特征，与真值做组内 spearman（GRPO 实际用的口径）+ 全局 spearman。 -纯 CPU，无 GPU。用法：/usr/local/bin/python3 skill_feature_corr.py -""" -import json -import os -import re -from collections import defaultdict - -import numpy as np - -SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) -OUT = os.path.join(SCRIPT_DIR, 'logp_corr') - -pairs = [json.loads(l) for l in open(os.path.join(OUT, 'pairs.jsonl'))] -rolls = [json.loads(l) for l in open(os.path.join(OUT, 'rollout_results.jsonl'))] -pas = {r['key']: float(np.mean(r['pass'])) for r in rolls if r['kind'] == 'skill' and r['mode'].startswith('t05')} -trc = {r['key']: float(np.mean(r['trunc'])) for r in rolls if r['kind'] == 'skill' and r['mode'].startswith('t05')} -tok = {r['key']: float(np.mean(r['tokens'])) for r in rolls if r['kind'] == 'skill' and r['mode'].startswith('t05')} - -# ---- skill 文本特征 ------------------------------------------------------------------- -_NUM = re.compile(r'\d') -_FORMULA = re.compile(r'[=+\-*/^]|\\frac|\\sqrt|\\sum|\\int|\$') -_STEP = re.compile(r'(?im)^\s*(step\s*\d|[0-9]+[.)]|first|second|third|next|then|finally)\b') -_IMPER = re.compile(r'(?i)\b(use|apply|consider|note|remember|compute|calculate|check|verify|' - r'identify|recall|avoid|ensure|find|start|begin|rewrite|simplify|substitute)\b') -_HEDGE = re.compile(r'(?i)\b(might|maybe|perhaps|possibly|could|may|try)\b') -_LATEX = re.compile(r'\\[a-zA-Z]+') -_PITFALL = re.compile(r'(?i)\b(mistake|error|pitfall|careful|caution|wrong|avoid|common|trap|' - r'incorrect|forget|overlook)\b') - - -def feats(skill): - s = skill or '' - words = s.split() - nw = max(1, len(words)) - sents = [x for x in re.split(r'[.!?\n]', s) if x.strip()] - return { - 'chars': len(s), - 'words': nw, - 'sent_len': nw / max(1, len(sents)), # 平均句长（可读性） - 'num_density': len(_NUM.findall(s)) / nw, # 数字密度（具体计算 vs 抽象） - 'formula_density': len(_FORMULA.findall(s)) / nw, - 'latex_density': len(_LATEX.findall(s)) / nw, - 'step_markers': len(_STEP.findall(s)), # 步骤/编号结构 - 'imper_density': len(_IMPER.findall(s)) / nw, # 祈使动词（指令性） - 'hedge_density': len(_HEDGE.findall(s)) / nw, # 模糊限定词（不确定） - 'pitfall_density': len(_PITFALL.findall(s)) / nw, - 'uniq_ratio': len(set(words)) / nw, # 词汇多样性（低=啰嗦重复） - } - - -def rank(x): - x = np.asarray(x, float) - o = np.argsort(x, kind='mergesort') - r = np.empty(len(x)) - r[o] = np.arange(len(x)) - for v in np.unique(x): - m = x == v - if m.sum() > 1: - r[m] = r[m].mean() - return r - - -def sp(a, b): - a, b = np.asarray(a, float), np.asarray(b, float) - m = ~(np.isnan(a) | np.isnan(b)) - if m.sum() < 3: - return np.nan - ra, rb = rank(a[m]), rank(b[m]) - if ra.std() == 0 or rb.std() == 0: - return np.nan - return float(((ra - ra.mean()) * (rb - rb.mean())).mean() / (ra.std() * rb.std())) - - -rows = [] -for pr in pairs: - k = pr['pair_id'] - if k not in pas: - continue - f = feats(pr['skill']) - f.update({'pair_id': k, 'data_id': pr['data_id'], 'pass': pas[k], - 'trunc': trc.get(k, np.nan), 'tok': tok.get(k, np.nan), - 'leaked': float(pr['leaked'])}) - rows.append(f) - -FKEYS = ['chars', 'words', 'sent_len', 'num_density', 'formula_density', 'latex_density', - 'step_markers', 'imper_density', 'hedge_density', 'pitfall_density', 'uniq_ratio'] - -print(f'[skill_feat] n_pairs={len(rows)}') -# 全局 -print('\n=== 全局 spearman：skill 特征 vs 真值 ===') -print('%-16s %-10s %-10s %-10s' % ('feature', 'sp(pass)', 'sp(trunc)', 'sp(tok)')) -for fk in FKEYS: - v = [r[fk] for r in rows] - print('%-16s %+.3f %+.3f %+.3f' % ( - fk, sp(v, [r['pass'] for r in rows]), - sp(v, [r['trunc'] for r in rows]), sp(v, [r['tok'] for r in rows]))) - -# 组内 -by = defaultdict(list) -for r in rows: - by[r['data_id']].append(r) -groups = [g for g in by.values() if len(g) >= 4] -print(f'\n=== 组内 spearman（每题内，n_groups={len(groups)}，mean±se）vs pass_rate ===') -for fk in FKEYS: - cs = [] - for g in groups: - c = sp([r[fk] for r in g], [r['pass'] for r in g]) - if not np.isnan(c): - cs.append(c) - if cs: - cs = np.array(cs) - print('%-16s mean=%+.3f se=%.3f n=%d' % (fk, cs.mean(), cs.std() / np.sqrt(len(cs)), len(cs))) - -print(f'\n=== 组内 spearman vs trunc_rate（截断通道）===') -for fk in FKEYS: - cs = [] - for g in groups: - c = sp([r[fk] for r in g], [r['trunc'] for r in g]) - if not np.isnan(c): - cs.append(c) - if cs: - cs = np.array(cs) - print('%-16s mean=%+.3f se=%.3f n=%d' % (fk, cs.mean(), cs.std() / np.sqrt(len(cs)), len(cs))) diff --git a/cookbook/exp/skill2lora/train_skill_v2.py b/cookbook/exp/skill2lora/train_skill_v2.py deleted file mode 100644 index 40aec9094..000000000 --- a/cookbook/exp/skill2lora/train_skill_v2.py +++ /dev/null @@ -1,2326 +0,0 @@ -"""Simplified GRPO + buffer-distill training for the reflexion skill generator (v2). - -Key differences from train_reflexion_skill.py: -- No view A/B split: all skill-gen is query-only (deployment form). -- No baseline rollout in training, no balance selection. -- thinking ON for the skill model: actor reasons in <think> then emits a distilled - <skills> block; the <think> is stripped by _extract_skill and NEVER reaches the executor - (executor only consumes <skills>), so it is not SEAM-style think leakage. -- Reward = parseable × correct (aligned with SEAM lpem: no terminated, no length penalty). -- Buffer A: adv=0 (all-fail) problems accumulate failure trajectories. -- Buffer B: batch rubric → regenerate skill → pass@k validate → SFT injection. -- SFT is event-driven: buffer B reaches threshold → one SFT pass → eval. - -Launch: - LLM_BACKUP_API_KEY=... python cookbook/exp/skill2lora/train_skill_v2.py \ - --dataset aops --n 5000 --chunk-size 16 --lr 6e-6 -""" -import argparse -import copy -import hashlib -import json -import math -import os -import re -import sys -import threading -import time -from concurrent.futures import ThreadPoolExecutor -from typing import Any, Dict, List, Optional, Set, Tuple - -import numpy as np - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.checkpoint_engine import CheckpointEngineManager -from twinkle.data_format import SamplingParams, pack_user_data -from twinkle.dataset import Dataset, DatasetMeta -from twinkle.model import TransformersModel -from twinkle.patch.no_split_modules import NoSplitModulesPatch -from twinkle.processor import InputProcessor -from twinkle.sampler import vLLMSampler -from twinkle.template import Template -from twinkle_agentic.verifier import RubricVerifier -from twinkle_agentic.verifier.rubric_verifier import RubricItem - -# 任务适配器（BigCodeBench）。code_task 刻意不 import 本模块，所以这里可以顶层 import。 -import code_task - -logger = get_logger() - -try: - import swanlab -except ImportError: - swanlab = None - -MODEL_ID = os.environ.get('GEN_MODEL_ID', 'Qwen/Qwen3-4B') -GPU_MEM = float(os.environ.get('GEN_GPU_MEM', 0.8)) -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') - -TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 2)) -REF_GPUS = int(os.environ.get('REF_GPUS', 2)) -SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) -BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 2)) -NUM_GPUS = TRAIN_GPUS + REF_GPUS + SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS -TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', min(1, TRAIN_GPUS))) -REF_FSDP = int(os.environ.get('REF_FSDP', min(1, REF_GPUS))) -TRAIN_DP = TRAIN_GPUS // TRAIN_FSDP -REF_DP = REF_GPUS // REF_FSDP - - -# =========================================================================== -# Section A0 — task switch (math = DeepMath \boxed{}, code = BigCodeBench unittest) -# =========================================================================== -# 与 _ALIGN_MODE / _SKILL_STYLE 同型的模块级开关（由 set_task 设置，trainer 在任何 prompt -# 构造之前调用）。**math 分支逐字不变**，所以 E1-E16 的行为、判分与可复现性不受影响。 -# 分派点一共 7 处：build_direct_prompt / build_skill_solve_prompt / _skillgen_prompt / -# _parse_seq(+_parse_many) / _answer_leaked / build_rubric_checker / _diagnose_entry。 -# 换成 code 时 reference_answer 不再是数值，而是 code_task.payload_of() 的判分载荷 -# （task_id / entry_point / test / code_prompt / doc_struct / canonical_solution）。 -_TASK = 'math' # 'math' | 'code' -_CODE_TEST_WORKERS = 24 # 单测线程池（子进程并行度） -_CODE_TEST_TIMEOUT = 60 # 单题单测墙钟上限（秒） - - -def set_task(task: str, test_workers: int = 24, test_timeout: int = 60) -> None: - """Set the task family. MUST run before any prompt is built or any roll is judged.""" - global _TASK, _CODE_TEST_WORKERS, _CODE_TEST_TIMEOUT, _RUBRIC_VERSION - if task not in ('math', 'code'): - raise ValueError(f"task must be 'math' or 'code', got {task!r}") - _TASK = task - _CODE_TEST_WORKERS = max(1, int(test_workers)) - _CODE_TEST_TIMEOUT = max(5, int(test_timeout)) - # 判据表换了，rubric 缓存键必须跟着换（rubric_cache 动态读 v2._RUBRIC_VERSION）。 - _RUBRIC_VERSION = code_task.RUBRIC_VERSION if task == 'code' else _RUBRIC_VERSION_MATH - - -# =========================================================================== -# Section A — boxed extraction + answer grading (verbatim from v1) -# =========================================================================== -_BOXED_RE = re.compile(r'\\boxed\s*\{') - - -def extract_boxed(text: str) -> Optional[str]: - if not text: - return None - last = None - for m in _BOXED_RE.finditer(text): - depth, i = 1, m.end() - while i < len(text) and depth > 0: - depth += (text[i] == '{') - (text[i] == '}') - i += 1 - if depth == 0: - last = text[m.end():i - 1].strip() - return last - - -# SEAM-style answer format: executor emits <think>...</think><answer>[numeric only]</answer>. -# 中文注释：executor 答案格式对齐 SEAM——优先解析 <answer>…</answer>，回退到 \boxed{} -# （兼容旧轨迹/rubric 提示）。取最后一个 <answer>，容忍缺失闭合标签（截断时取到 EOS）。 -_ANSWER_RE = re.compile(r'<answer>(.*?)</answer>', re.DOTALL | re.IGNORECASE) -_ANSWER_OPEN_RE = re.compile(r'<answer>(.*)', re.DOTALL | re.IGNORECASE) - - -def extract_answer(text: str) -> Optional[str]: - """Extract the final answer, preferring SEAM's <answer>…</answer>, falling back to \\boxed{}.""" - if not text: - return None - matches = _ANSWER_RE.findall(text) - if matches: - return matches[-1].strip() or None - # tolerate a truncated / unclosed <answer> tag (e.g. cut at token budget) - m = _ANSWER_OPEN_RE.search(text) - if m: - return m.group(1).strip() or None - return extract_boxed(text) - - -def normalize_answer(ans: str) -> str: - if not ans: - return '' - s = str(ans).strip() - m = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) - if m: - return m.group(1) - s = s.strip('$').strip().replace('\u2212', '-') - 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').replace(r'\tfrac', r'\frac') - s = s.replace(r'\,', '').replace(r'\;', '').replace(r'\!', '') - s = re.sub(r'\\(?:quad|qquad|\s)', '', s) - s = re.sub(r'\s+', '', s) - s = re.sub(r'\\sqrt([0-9A-Za-z])', r'\\sqrt{\1}', s) - s = re.sub(r'\\frac\{([^{}]+)\}([^{}\\])', r'\\frac{\1}{\2}', s) - s = re.sub(r'\\frac([^{\\])([^{\\])', r'\\frac{\1}{\2}', s) - s = re.sub(r'\{[a-zA-Z]+\}$', '', s) - s = re.sub(r'\^\\circ|\^\{\\circ\}|\u00b0|\\circ', 'deg', s) - s = s.replace(r'\minus{}', '-').replace(r'\minus', '-') - - def _frac_to_slash(mt): - text = mt.group(0) - pos = text.index('{') + 1 - depth, num_start = 1, pos - while depth > 0: - depth += (text[pos] == '{') - (text[pos] == '}') - pos += 1 - numer = text[num_start:pos - 1] - pos += 1 - den_start, depth = pos, 1 - while depth > 0: - depth += (text[pos] == '{') - (text[pos] == '}') - pos += 1 - return f'({numer})/({text[den_start:pos - 1]})' - - s = re.sub(r'\\frac\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', _frac_to_slash, s) - s = re.sub(r'(?<!\w)(\d+)/(\d+)(?!\w)', r'(\1)/(\2)', s) - return s - - -def _try_numeric_equal(a: str, b: str) -> bool: - try: - va, vb = float(a.replace('(', '').replace(')', '')), 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) - return va is not None and vb is not None and abs(va - vb) < 1e-9 * max(1, abs(va), abs(vb)) - - -_MCQ_REF_RE = re.compile( - r'^\\(?:textbf|text|mathrm|mathbf)\{\(?([A-E])\)?\s*\}\s*(.+)$' - r'|^\(?([A-E])\)\s+(.+)$') -_VAR_PREFIX_RE = re.compile(r'^(?:[A-Za-z](?:\([^)]*\))?|\([^)]*\))\s*=\s*(.+)$', re.DOTALL) - - -def _split_mcq(ans): - s = ans.strip() - m = _MCQ_REF_RE.match(s) - if m: - return (m.group(1) or m.group(3)), ((m.group(2) or m.group(4) or '').strip() or None) - bl = re.match(r'^\\?(?:textbf|text|mathrm|mathbf|mathbb)?\{?\(?([A-E])\)?\}?$', s) - return (bl.group(1), None) if bl else (None, s or None) - - -def _strip_var_prefix(ans): - m = _VAR_PREFIX_RE.match((ans or '').strip()) - return m.group(1).strip() if m else (ans or '') - - -def _math_verify_equal(predicted: str, reference: str) -> bool: - try: - from math_verify import parse, verify - gold = parse(r'\boxed{' + reference + '}', parsing_timeout=1) - pred = parse(r'\boxed{' + predicted + '}', parsing_timeout=1) - return bool(gold and pred and verify(gold, pred, timeout_seconds=1)) - except Exception: - return False - - -def answers_match(predicted: str, reference: str) -> bool: - if not predicted or not reference: - return False - norm_p, norm_r = normalize_answer(predicted), normalize_answer(reference) - if norm_p == norm_r or norm_p.lower() == norm_r.lower() or _try_numeric_equal(norm_p, norm_r): - return True - stripped_p = normalize_answer(_strip_var_prefix(predicted)) - stripped_r = normalize_answer(_strip_var_prefix(reference)) - if stripped_p and stripped_r: - if (stripped_p == stripped_r or stripped_p.lower() == stripped_r.lower() - or _try_numeric_equal(stripped_p, stripped_r)): - return True - p_letter, p_value = _split_mcq(predicted) - r_letter, r_value = _split_mcq(reference) - if p_letter and r_letter and p_letter == r_letter: - return True - if (p_letter and p_value is None) and (r_letter and r_value is None): - return False - 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 and (p_val == r_val or p_val.lower() == r_val.lower() or _try_numeric_equal(p_val, r_val)): - return True - for left, right in ((norm_p, norm_r), (stripped_p, stripped_r)): - tl, tr = re.sub(r'[\s()\[\]{}\\]', '', left or ''), re.sub(r'[\s()\[\]{}\\]', '', right or '') - if ',' in tl and tl == tr: - return True - if '=' in norm_r and '=' not in norm_p: - if any(part and (part == norm_p or _try_numeric_equal(part, norm_p)) for part in norm_r.split('=')): - return True - if '=' in norm_p and '=' not in norm_r: - if any(part and (part == norm_r or _try_numeric_equal(part, norm_r)) for part in norm_p.split('=')): - return True - return _math_verify_equal(stripped_p or norm_p, stripped_r or norm_r) - - -_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') - - -def _numeric_value(raw) -> Optional[str]: - if raw is None: - return None - s = str(raw).strip().strip('$').strip() - s = s.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac') - s = re.sub(r'\\!|\\,|\\;|\\ |\\left|\\right|\s', '', s) - for pat in (r'\\frac\{(-?\d+)\}\{(-?\d+)\}', r'(-?\d+)/(-?\d+)'): - m = re.fullmatch(pat, s) - if m: - a, b = int(m.group(1)), int(m.group(2)) - return str(int(a / b)) if (b and a / b == int(a / b)) else (str(a / b) if b else None) - return (str(int(float(s))) if float(s) == int(float(s)) else str(float(s))) if _NUM_RE.fullmatch(s) else None - - -def _answer_leaked(skill: str, reference) -> bool: - if not skill: - return False - if _TASK == 'code': - # 代码域：leak = skill 里出现参考解答的实质代码行（与数学域一致，只做监控） - return code_task.leaked(skill, reference) - # Suffix guard: reject only a following DIGIT or a following '.<digit>' (decimal point), - # NOT a sentence-ending '.'. Old '(?![\d.])' let leaks like "...= 675." slip through - # because the trailing period satisfied the [\d.] class. 中文注释：尾断言只排除"后接数字" - # 或"后接小数点+数字"，不排除句末句号，堵住 "答案." 这类泄漏漏检。 - for cand in {_numeric_value(reference), (str(reference).strip() or None)}: - if cand and re.search(r'(?<![\d.])' + re.escape(cand) + r'(?!\d)(?!\.\d)', skill): - return True - return False - - -# ---- SEAM lpem-style numeric correctness (seam=整段 sanitize；v2=\boxed{} 锚定) ---- -# 中文注释：复刻 SEAM lpem 判分——抽 <answer>/boxed/$...$/分数/首个数字，float 归一后纯数值精确匹配。 -_SEAM_TAG_RE = re.compile(r'<\s*answer\s*>(.*?)<\s*/\s*answer\s*>', re.I | re.S) -_SEAM_BOX_RE = re.compile(r'(?:\\{1,2}\(|)\\{1,2}boxed\s*\{\s*([^}]*)\s*}(?:\)|)', re.S) -_SEAM_INLINE_RE = re.compile(r'\$([^$]+)\$|\\\(([^)]+)\\\)', re.S) -_SEAM_FRAC_RE = re.compile(r'(-?\d+(?:\.\d+)?)/(-?\d+(?:\.\d+)?)') -_SEAM_NUM_RE = re.compile(r'-?\d+(?:\.\d+)?') - - -def _seam_norm(num: str) -> str: - try: - f = float(num) - return str(int(f)) if f == int(f) else str(f) - except Exception: - return num.strip() - - -def _seam_sanitize(txt: str, dfrac_fix: bool = True) -> str: - """Port of SEAM lpem.sanitize_math_answer + normalize_number_format. - - ``dfrac_fix=False`` gives BIT parity with the upstream function (which only rewrites the - literal ``\\frac``); it is used by the ``align='seam'`` judge so that E13's acc is - reproducible against SEAM's step_summary. See _parse_seq.""" - txt = (txt or '').strip() - if (m := _SEAM_TAG_RE.search(txt)): - txt = m.group(1).strip() - elif (m := _SEAM_BOX_RE.search(txt)): - txt = m.group(1).strip() - elif (m := _SEAM_INLINE_RE.search(txt)): - txt = (m.group(1) or m.group(2)).strip() - # bugfix 2026-07-29：\dfrac/\tfrac/\cfrac 不被下行字面 \frac 正则匹配，曾致 - # \boxed{\dfrac{1}{2}} 落到 _SEAM_NUM_RE 抓首个数字 → pred='1'（分数答案题全判错； - # 实测被标"错"的 boxed rolls 中 70-85% 实为正确，见 skill_quality_analysis.md 末章）。 - # 注：SEAM 上游没有这一步，seam 对齐口径下必须关掉（dfrac_fix=False）。 - if dfrac_fix: - txt = txt.replace(r'\dfrac', r'\frac').replace(r'\tfrac', r'\frac').replace(r'\cfrac', r'\frac') - txt = re.sub(r'\\frac\s*\{\s*([^}]+?)\s*}\s*\{\s*([^}]+?)\s*}', r'\1/\2', txt) - if (m := _SEAM_FRAC_RE.search(txt)): - p, q = map(float, m.groups()) - if q: - return _seam_norm(str(p / q)) - if (m := _SEAM_NUM_RE.search(txt)): - return _seam_norm(m.group()) - return txt - - -# =========================================================================== -# Section B — sampling / parsing utilities -# =========================================================================== -_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') - - -def _clean_text(decoded: Optional[str]) -> str: - return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() - - -def _extract_skill(text: str) -> Optional[str]: - """Parse the skill block: <memory_item> in seam mode (SEAM format_pass parity), else <skills>.""" - if _ALIGN_MODE == 'seam': - # ⭐ 整段搜、取首个匹配，不先剔掉 <think> —— 与 SEAM 两处实现逐字一致： - # lpem.format_pass：MEMORY_RE.search(resp)；fsdp_workers.py:836：_re.search(..., response_text)。 - # 旧实现只搜 </think> 之后，会把"只在 think 里写了 memory_item"的候选当成格式失败 - # （reward=0），而 SEAM 那边算格式通过 —— 直接影响 format 率与组内 reward 方差。 - m = re.search(r'<memory_item>(.*?)</memory_item>', text or '', re.DOTALL | re.IGNORECASE) - return (m.group(1).strip() or None) if m else None - low = text.lower() - end_think = low.rfind('</think>') - answer = text[end_think + len('</think>'):] if end_think >= 0 else text - open_tag, close_tag = '<skills>', '</skills>' - s = answer.lower().rfind(open_tag) - if s < 0: - return None - inner = s + len(open_tag) - e = answer.lower().find(close_tag, inner) - if e < 0: - return None - block = answer[inner:e].strip() - block = re.sub(r'</?(?:skills|skill|diagnose|pitfall|strategy|think)>', '', block, flags=re.IGNORECASE).strip() - return block or None - - -def _parse_seq(seq, gold) -> Dict[str, Any]: - if _TASK == 'code': - return _parse_many([(seq, gold)])[0] - text = _clean_text(getattr(seq, 'decoded', '') or '') - if _ALIGN_MODE == 'seam': - # ⭐ seam 对齐口径（2026-08-02 修正）= SEAM_JUDGE=answer，即 train_deepmath_paper.sh 的默认值： - # 整段文本走 sanitize 级联 <answer> → \boxed → $..$/\(..\) → 分数 → **首个数字**。 - # 之前这里走的是 boxed-only（等价 SEAM_JUDGE=boxed），与 executor prompt 要求 - # "<think>...</think><answer>...</answer>" 直接冲突 —— 换成原版 prompt 后 boxed-only - # 会把几乎所有 rollout 判错。prompt 与判分必须成对切换，见 _SEAM_SOLVE_ADVISORY 注释。 - # dfrac_fix=False 是为了与 lpem.sanitize_math_answer 逐字节一致。 - pred = _seam_sanitize(text, dfrac_fix=False) or None - correct = bool(pred) and (pred == _seam_sanitize(str(gold), dfrac_fix=False)) - terminated = getattr(seq, 'stop_reason', None) != 'length' - return {'pred': pred, 'correct': correct, 'terminated': terminated, - 'stop_reason': getattr(seq, 'stop_reason', None), - 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} - # v2（E1-E12/E14-E21）：只从 \boxed{} 抽取，再走同一套数值归一（frac/inline/number）后精确匹配； - # 不做 lpem 式"整段抓数字"贪婪回退，保证这些臂之间的 acc/lift 横向可比。extract_boxed 取最后一个 - # 配平的 \boxed{}、截断时不误取；没有则判错。 - raw = extract_boxed(text) - pred = _seam_sanitize(raw) if raw else None - correct = bool(pred) and (pred == _seam_sanitize(str(gold))) - terminated = getattr(seq, 'stop_reason', None) != 'length' - return {'pred': pred, 'correct': correct, 'terminated': terminated, - 'stop_reason': getattr(seq, 'stop_reason', None), - 'gen_tokens': len(getattr(seq, 'tokens', None) or []), 'text': text} - - -def _parse_many(pairs) -> List[Dict[str, Any]]: - """批量判分入口。``pairs`` = [(seq_or_None, gold)]，返回同序 roll 列表。 - - math 分支逐条 _parse_seq（与逐条调用 bit 一致）；code 分支必须批量 —— 判分要起子进程跑 - unittest（典型 1-3s），一个 chunk 有几百次判分，串行会比同 chunk 的 GPU 时间还长一个量级。 - 所有 rollout 汇合点（process_chunk / run_greedy_eval / methods / eval_reflexion）都走这里。 - """ - if _TASK != 'code': - return [(_parse_seq(s, g) if s is not None else _empty_roll()) for s, g in pairs] - items = [] - for s, g in pairs: - if s is None: - items.append(None) - continue - items.append((_clean_text(getattr(s, 'decoded', '') or ''), - getattr(s, 'stop_reason', None), - len(getattr(s, 'tokens', None) or []), g)) - return code_task.judge_many(items, _CODE_TEST_WORKERS, _CODE_TEST_TIMEOUT) - - -def _first_seq(seqs): - """rollout 列表 -> 首个 sequence 或 None（判分批量化后统一用它取 seq）。""" - return seqs[0] if seqs else None - - -def _empty_roll(): - if _TASK == 'code': - return code_task.empty_roll() - return {'pred': '', 'correct': False, 'terminated': False, - 'stop_reason': 'empty', 'gen_tokens': 0, 'text': ''} - - -def _run_samples(sampler, prompts, num_samples, max_tokens, gen_dp, - temperature=None, top_p=None, top_k=None, logprobs=None): - if not prompts: - return [] - params = SamplingParams( - max_tokens=max_tokens, - temperature=GEN_TEMPERATURE if temperature is None else temperature, - top_p=GEN_TOP_P if top_p is None else top_p, - num_samples=num_samples, **({} if top_k is None else {'top_k': top_k}), - **({} if logprobs is None else {'logprobs': logprobs})) - padded = prompts - if gen_dp > 1 and 0 < len(prompts) < gen_dp: - padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] - responses = sampler.sample(padded, params)[:len(prompts)] - return [list(r.sequences) if (r and r.sequences) else [] for r in responses] - - -def sampler_logprobs(seq): - """从 ``SampledSequence.logprobs`` 抽出采样 token 自己的 logprob（逐 token 一个 float）。 - - 采样器返回的形状是 ``[[(token_id, logprob), ...], ...]``（logprobs=1 时每位只有一项， - 就是被采中的那个 token）。抽出来喂给 GRPOMetric 的 ``sampler_logps``，它会把这串值与 - 训练 forward 算出的 logp 逐 token 对账 —— 序列一致时两边只差引擎精度。 - 采样时 T=1/top_p=1/top_k=-1（skill-gen 的固定参数），logits 没经过任何处理，所以 - vLLM 的 processed_logprobs 就等于原始 logprob，与 trainer 的取值口径一致。 - """ - out = [] - for item in (getattr(seq, 'logprobs', None) or []): - if not item: - return [] - out.append(float(item[0][1])) - return out - - -# =========================================================================== -# Section C — data loading (simplified: no balance, no xproblem, no views) -# =========================================================================== -def _boxed_batch(rows, dataset): - sols = rows['solution'] - metas = rows.get('metadata', [None] * len(sols)) - refs = [extract_boxed(s or '') for s in sols] - keep = [bool(ref) and (dataset != 'aops' or bool((meta or {}).get('boxed'))) - for ref, meta in zip(refs, metas)] - return {**rows, 'reference_answer': refs, '_keep': keep} - - -def load_problems(dataset: str, n: int, seed: int) -> List[Dict[str, Any]]: - ds_id = AOPS_DATASET_ID if dataset == 'aops' else os.environ.get('MATH_DATASET_ID', 'modelscope/competition_math') - ds = Dataset(DatasetMeta(dataset_id=f'ms://{ds_id}', split='train')) - nproc = min(32, os.cpu_count() or 1) - ds.map(lambda rows: _boxed_batch(rows, dataset), num_proc=nproc) - ds.filter(lambda row: row['_keep'], num_proc=nproc) - out = [{'data_id': f'{dataset}:{i}', 'problem': row['problem'], - 'reference_answer': row['reference_answer']} - for i, row in enumerate(ds.dataset)] - rng = np.random.RandomState(seed) - rng.shuffle(out) - return out[:n] if (n and n < len(out)) else out - - -def _load_seam_parquet(path: str) -> List[Dict[str, Any]]: - """Read a SEAM ``build_aops_dataset.py`` parquet (VERL RLHF schema) into twinkle records, - PRESERVING file row order. ``problem <- extra_info.problem`` and - ``reference_answer <- reward_model.ground_truth``. No shuffle/filter: the parquet is already - SEAM's numeric-filtered, seed-42-shuffled, truncated split. - 中文注释：直读 SEAM parquet、保持文件顺序，用于让 twinkle 与 SEAM 输入同一批数据。""" - import pyarrow.parquet as pq - rows = pq.read_table(path).to_pylist() - out: List[Dict[str, Any]] = [] - for i, r in enumerate(rows): - ei = r.get('extra_info') or {} - rm = r.get('reward_model') or {} - problem = (ei.get('problem') or '').strip() - ref = rm.get('ground_truth') - if not problem or ref is None: - continue - out.append({'data_id': f"seam:{ei.get('split', '')}:{ei.get('index', i)}", - 'problem': problem, 'reference_answer': str(ref)}) - return out - - -def load_train_order_file(path: str) -> List[Dict[str, Any]]: - """Read a fixed training ORDER file (jsonl) -> records in file order, duplicates kept. - - Why this exists: verl's dataloader shuffles (data.shuffle defaults True), so SEAM's step-k - batch is NOT train.parquet[k*128:(k+1)*128]. Measured 2026-08-02: twinkle chunk 0 and SEAM - step 1 drew from the SAME 5000-problem pool but shared only 1 of 128 problems, which alone - put ~6-7 accuracy points between the two curves. This file is SEAM's REALIZED batch - sequence, reverse-engineered from its rollout dump (40 steps x 128 problems, in order), so - feeding it with ProblemPool(fixed_order=True) makes chunk k == SEAM step k+1 problem by - problem. Generated by .tmp_analysis/mk_seam_train_order.py. - Each line: {'data_id','problem','reference_answer'[,'level','seam_step']}. - """ - out: List[Dict[str, Any]] = [] - with open(path, encoding='utf-8') as f: - for line in f: - line = line.strip() - if not line: - continue - d = json.loads(line) - problem, ref = (d.get('problem') or '').strip(), d.get('reference_answer') - if not problem or ref is None: - continue - rec = {'data_id': str(d.get('data_id', '')), 'problem': problem, - 'reference_answer': str(ref)} - if d.get('level') is not None: - rec['level'] = d['level'] - out.append(rec) - if not out: - raise ValueError(f'--train-order-file {path} produced 0 records') - return out - - -def _load_records(args): - seam_dir = (getattr(args, 'seam_parquet_dir', '') or '').strip() - if seam_dir: # 直读 SEAM parquet：按文件顺序取 train，val 整份当 eval，跳过 load/numeric/shuffle/split - tp, vp = os.path.join(seam_dir, 'train.parquet'), os.path.join(seam_dir, 'val.parquet') - if not (os.path.exists(tp) and os.path.exists(vp)): - raise FileNotFoundError( - f'--seam-parquet-dir needs both train.parquet and val.parquet in {seam_dir}') - pool = _load_seam_parquet(tp) # already SEAM-shuffled + truncated, file order - eval_records = _load_seam_parquet(vp) # SEAM's exact val holdout - eval_probs = {r['problem'] for r in eval_records} - train_records = [r for r in pool if r['problem'] not in eval_probs] - if args.n > 0: - train_records = train_records[:args.n] - if {r['problem'] for r in train_records} & eval_probs: - raise ValueError('eval/train overlap detected in SEAM parquet') - logger.info(f'[data] SEAM parquet: train={len(train_records)} eval={len(eval_records)} dir={seam_dir}') - return train_records, eval_records - records = load_problems(args.dataset, 0, args.seed) - raw_n = len(records) - if args.numeric_only: - records = [{**r, 'reference_answer': v} - for r, v in ((r, _numeric_value(r.get('reference_answer'))) for r in records) - if v is not None] - np.random.RandomState(args.seed).shuffle(records) - # exclude - excl_ids, excl_probs = set(), set() - for path in (args.exclude_data_ids or '').split(','): - path = path.strip() - if not path or not os.path.exists(path): - continue - with open(path) as f: - for line in f: - if not line.strip(): continue - row = json.loads(line) - if row.get('record_type') in {'config', 'summary'}: continue - did = str(row.get('data_id', '')).strip() - if did: excl_ids.add(did) - else: - p = str(row.get('problem', '')).strip() - if p: excl_probs.add(p) - if excl_ids or excl_probs: - records = [r for r in records - if str(r.get('data_id', '')) not in excl_ids - and str(r.get('problem', '')).strip() not in excl_probs] - eval_n = min(args.eval_size, len(records)) if args.eval_size > 0 else 0 - eval_records = records[:eval_n] - # Dedup by problem TEXT: index slices are disjoint, but duplicate problem statements - # across the boundary would still leak eval into train. Drop any train record whose - # problem appears in eval, then guard with an explicit overlap assertion. - # 中文注释：train/eval 去重——按题面文本剔除，防止数据集内重复题目跨界泄漏；末尾硬断言无交集。 - eval_probs = {r['problem'] for r in eval_records} - train_records = [r for r in records[eval_n:] if r['problem'] not in eval_probs] - if args.n > 0: - train_records = train_records[:args.n] - if {r['problem'] for r in train_records} & eval_probs: - raise ValueError('eval/train overlap detected after dedup') - logger.info(f'[data] raw={raw_n} train={len(train_records)} eval={len(eval_records)}') - return train_records, eval_records - - -# =========================================================================== -# Section D — DiskCache, ProblemPool, LockedSampler -# =========================================================================== -class DiskCache: - def __init__(self, path: str, enabled: bool = True): - self._mem: Dict[str, Any] = {} - self._fh = None - self._lock = threading.Lock() - if not enabled: - return - if os.path.exists(path): - with open(path, encoding='utf-8') as f: - for line in f: - if line.strip(): - row = json.loads(line) - self._mem[row['key']] = row['value'] - self._fh = open(path, 'a', encoding='utf-8') - - @staticmethod - def key_for(*parts): - return hashlib.md5('\x1f'.join(parts).encode('utf-8')).hexdigest() - - def get(self, key): return self._mem.get(key) - def __contains__(self, key): return key in self._mem - - def put(self, key, value): - with self._lock: - self._mem[key] = value - if self._fh: - self._fh.write(json.dumps({'key': key, 'value': value}, ensure_ascii=False) + '\n') - self._fh.flush() - - def close(self): - if self._fh: self._fh.close() - - -class ProblemPool: - """Dataloader-like batch sampler: epoch-wise seeded RandomSampler + drop_last. - - SEAM's verl dataloader uses a sampler and drop_last=True; this mirrors that behavior more - closely than the old cursor loop that carried a short epoch tail into the next batch. - - ``fixed_order=True`` disables the permutation and walks ``records`` in file order. That is - what --train-order-file needs: the order file already IS verl's realized batch sequence - (reverse-engineered from SEAM's rollout dump), so any reshuffle here would destroy it. - """ - def __init__(self, records, seed, fixed_order=False): - self._records = list(records) - self._seed, self._cursor, self.epoch = seed, 0, 0 - self._fixed_order = bool(fixed_order) - self._order: List[int] = [] - self._reset_epoch() - - def _reset_epoch(self): - if self._fixed_order: - self._order = list(range(len(self._records))) - else: - rng = np.random.RandomState(self._seed + self.epoch) - self._order = list(rng.permutation(len(self._records))) - self._cursor = 0 - - def draw(self, k): - if k > len(self._records): - raise ValueError(f'batch size {k} exceeds dataset size {len(self._records)}') - if self._cursor + k > len(self._order): - self.epoch += 1 - self._reset_epoch() - idx = self._order[self._cursor:self._cursor + k] - self._cursor += k - return [self._records[i] for i in idx] - - -class _LockedSampler: - def __init__(self, sampler): - self._sampler = sampler - self._lock = threading.Lock() - - def sample(self, *a, **kw): - with self._lock: - return self._sampler.sample(*a, **kw) - - def __getattr__(self, name): - return getattr(self._sampler, name) - - -# =========================================================================== -# Section E — Rubric (teacher diagnosis, batched at distill time) -# =========================================================================== -_RFT_DIAG_SYSTEM = """\ -You are a strategy-level process checker for a math solution attempt. You are given a -math problem, a rubric, and one attempted solution segment. Decide PASS or FAIL for each -criterion, and write the diagnosis so it can become useful reusable guidance for solving -similar problems without seeing this segment. - -Output STRICT JSON (no prose outside it) with this shape: -{"items": [{"index": 1, "verdict": "PASS"|"FAIL", "reason": "...", "fix": ""}], "overall": "OK"|"ISSUES", "summary": "..."} - -Rules: -- Judge every criterion independently. -- The diagnosis must stay answer-free. -- For FAIL items: describe the process problem at strategy level. -- A fix suggests the LOCAL correction direction without solving. -- Never reveal the final answer or a corrected expression. -- If segment was cut off (no final <answer> reached), mark the output-format criterion as FAIL. -- Keep "reason" and "fix" concise: one short sentence each. -- Output only the JSON object.""" - -_RFT_DIAG_USER = """\ -## Task / query -{query} - -## Rubric -{rubric} - -## Segment -{segment} - -Now output the diagnostic JSON object.""" - -# 判据按「错误类型」组织，但文本一律写成正向陈述 —— 全流程的语义是 PASS=没问题 / -# FAIL=该类错误存在（见 _format_diagnosis 与 gate），写成否定句会把 PASS/FAIL 反过来。 -_MATH_RUBRIC = [ - # 1. 代数计算错误 - ('Arithmetic and algebraic manipulations are carried out correctly', True), - # 2. 公式定理使用错误 - ('Formulas and theorems are invoked correctly and their preconditions hold', False), - # 3. 起始方法论错误 - ('The initial approach is viable for this problem rather than a dead end', False), - # 4. 题目目标分析错误 - ('The attempt correctly identifies what the problem actually asks for', False), - # 5. 输出格式错误 - ('The attempt reaches a final answer in the required output format', False), - # 6. 对计算过程反复犹豫 - ('The attempt commits to its computation instead of repeatedly second-guessing it', False), - # 7. 构成自相矛盾 - ('The attempt stays internally consistent and never contradicts its own results', False), -] -# 版本号进 rubric 缓存键（GlobalRubricCache._key）：判据一改，旧诊断必须失效， -# 否则 rubric_cache_global.jsonl 里按 data_id 存的旧taxonomy诊断会被当成新判据的结果返回。 -_RUBRIC_VERSION_MATH = 'rubric_v6_error_taxonomy' -_RUBRIC_VERSION = _RUBRIC_VERSION_MATH # set_task('code') 会换成 code_task.RUBRIC_VERSION - - -class _RftRubricVerifier(RubricVerifier): - def _diagnose_trajectory(self, query, rubric_block, segment_text): - return {'messages': [ - {'role': 'system', 'content': _RFT_DIAG_SYSTEM}, - {'role': 'user', 'content': _RFT_DIAG_USER.format( - query=query, rubric=rubric_block, segment=segment_text)}]} - - -class _CodeRubricVerifier(RubricVerifier): - """代码域 judge：判据 = code_task.CODE_RUBRIC，且 segment 里带**单测真实报错**。""" - - def _diagnose_trajectory(self, query, rubric_block, segment_text): - return {'messages': [ - {'role': 'system', 'content': code_task.DIAG_SYSTEM}, - {'role': 'user', 'content': code_task.DIAG_USER.format( - query=query, rubric=rubric_block, segment=segment_text)}]} - - -def build_rubric_checker(): - if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') - or os.environ.get('OPENAI_API_KEY')): - return None - if _TASK == 'code': - return _CodeRubricVerifier( - fixed_rubric=[RubricItem(t, is_hard=h) for t, h in code_task.CODE_RUBRIC], gate=True) - return _RftRubricVerifier( - fixed_rubric=[RubricItem(t, is_hard=h) for t, h in _MATH_RUBRIC], gate=True) - - -def _format_diagnosis(detail) -> str: - rub = detail.rubric - lines = [] - for it in detail.items: - text = rub[it.index - 1].text if 0 < it.index <= len(rub) else f'criterion {it.index}' - if it.verdict: - lines.append(f'- [PASS] {text}') - else: - tail = f': {it.reason}' if it.reason else '' - tail += f' (fix: {it.fix})' if it.fix else '' - lines.append(f'- [FAIL] {text}{tail}') - if detail.summary: - lines.append(f'Summary: {detail.summary}') - return '\n'.join(lines) - - -def _diagnose_entry(checker, entry: Dict[str, Any]) -> Optional[str]: - """Run the teacher rubric on ONE buffer-A failure trajectory → formatted diagnosis text - (or None on API error). Pure network/CPU (no GPU), so it can run on a background thread - while GRPO trains. Shared by the background pre-diagnosis pool and distill_buffer's - fallback for any entry the pool did not reach in time. - 中文注释：单条失败轨迹的 rubric 诊断（纯 API，不吃 GPU）。后台预诊断与 distill 补诊断共用。""" - if _TASK == 'code': - # 代码域：query 里补上题面声明的硬约定（签名/返回/异常/示例，不含参考解答）， - # segment 由调用方（_rubric_entry）拼成"提交的代码 + 单测真实报错"。 - query = code_task.diag_query(entry['problem'], entry['reference_answer']) - seg = {'messages': [{'role': 'user', 'content': query}, - {'role': 'assistant', 'content': entry['fail_segment']}]} - try: - return _format_diagnosis(checker.diagnose(seg, query=query)) - except Exception as exc: - logger.warning(f'[rubric] diagnose error: {exc}') - return None - seg_text = entry['fail_segment'] - if entry.get('fail_stop_reason') == 'length': - seg_text += ('\n\n[Process note: this attempt was cut off at the token budget ' - 'and never produced a final <answer>.]') - seg = {'messages': [{'role': 'user', 'content': entry['problem']}, - {'role': 'assistant', 'content': seg_text}]} - try: - return _format_diagnosis(checker.diagnose(seg, query=entry['problem'])) - except Exception as exc: - logger.warning(f'[rubric] diagnose error: {exc}') - return None - - -# =========================================================================== -# Section F — NEW: prompts, reward, buffer logic -# =========================================================================== - -# ---- Skill-gen system prompt (query-only; used in v2 mode) ---- -# 中文注释：skillmodel 系统提示词（仅 v2 模式；seam 模式改用 _SEAM_EXPERIENCE_PROMPT）。 -# 方案1：thinking 开启。让 skill 模型在 <think> 里“先把本题实际解一遍、想清楚”，再在 <skills> 里 -# 只写抽象出来的“通用方法论”（不含本题任何具体数字/中间结果/答案）。<think> 会被 _extract_skill -# 用 rfind('</think>') 砍掉、绝不流给 executor（避免 SEAM 那种 think 泄漏），executor 只吃 <skills>。 -# 之所以要开 thinking：nothinking 下模型无处安放解题过程，只能把“完整解答+答案”直接写进 <skills> -# （实测 <skills> 前置分析长度=0、且常把答案算出来写进块内），等于换标签的泄漏且质量差。开 thinking 后 -# “先解题、再提炼”两步显式分离，<skills> 才可能是真正可迁移、不代入本题数值的方法论。 -# skill_model/ref_model/skill_sampler 三者 enable_thinking 必须一致，否则训练轨迹 token 布局与采样对不上。 -SKILL_GEN_SYSTEM = """\ -You are a skill-generation model. Your <skills> block will be fed to a SEPARATE downstream executor model that must solve the problem on its own. The executor will NOT see your private reasoning — it only sees what is inside <skills>...</skills>. - -First, think privately: actually work the problem out in your head to make sure you understand it, then step back and abstract WHAT MAKES THIS TYPE OF PROBLEM SOLVABLE into transferable methodology. - -Then write the <skills> block following these rules: -- Give general, transferable solving techniques for this TYPE of problem: the key concepts/theorems it relies on, the recommended strategy and steps, and the common pitfalls to avoid — plus a brief reason for each piece of advice so the executor understands why. -- Write it as one coherent analysis narrative (not a bullet list): first name what the problem is essentially asking, then walk through how to approach it, blending concepts, steps, pitfalls and reasons into a single connected story. -- CRITICAL: Do NOT solve the problem for the executor. Do NOT reveal or compute the final answer, and do NOT substitute the problem's specific given numbers into the steps or state any intermediate numeric results. Leave ALL concrete numbers for the executor to compute on its own. If you catch yourself writing a specific number from the problem, replace it with a description of the quantity instead. -- Keep it concise: aim for roughly one focused paragraph. -- End the block with this exact sentence: "Avoid re-checking loops; box a bare number as soon as it is computed." - -Put ONLY the methodology inside <skills></skills>. - -Example: -<skills> -This problem is essentially asking for the units (last) digit of an integer raised to a high power; first get clear on what the problem is asking before deciding where to start. Since only the last digit matters, you should first look only at the units digit of the base, because the units digit of an integer power is determined solely by the units digit of the base and the higher digits do not affect the result — so at this step be careful not to expand or compute the whole large number, which is both unnecessary and error-prone. Next, repeatedly multiply this units digit by itself and record the units digit each time, until it starts to repeat, thereby obtaining its cycle period. The part about "determining the period length" is important here: be careful not to count one term too many or too few, otherwise all the later positioning will be off. Finally, take the given exponent modulo the period length and land on the corresponding term within the period; here pay special attention that when the remainder is 0 it corresponds to the last term of the period rather than the first. Overall, I summarize the approach for this kind of problem as "first recognize that it asks for the units digit of a power, then fix on the units digit to find the cycle period, and finally use the exponent modulo to locate the term", while leaving the concrete numbers for the downstream solver to substitute and compute on its own. Avoid re-checking loops; box a bare number as soon as it is computed. -</skills> -""" - -# ---- Skill 文体消融（--skill-style）---- -# 中文注释：探针实验（CONCLUSIONS_config/reflexion.md）验证的两种高性价比文体。 -# 关键约束：同一文体在主链路（query-only 预判）与 buffer B regen（rubric 诊断条件）下 -# 输出格式必须一致（toy=迷你题示范+迁移句；pitfall=WARNING/INSTEAD/纪律句）， -# 否则 GRPO 与 SFT 样本分布不一致无法联合训练。 -# toy 主链路：探针 P3_toy 原文（异数字玩具题，天然 answer-free）。 -SKILL_GEN_TOY = """\ -You are a skill-generation model. A separate executor model will solve the problem; it only sees your <skills> block. - -First think privately and identify the core technique this problem needs. Then, inside <skills></skills>, do exactly one thing: invent a MINIATURE problem of the same type with DIFFERENT and much smaller numbers, and solve that miniature completely in at most 5 short lines, making the key trick explicit. Finish with one transfer sentence: "Your problem has the same shape - repeat these steps with its own numbers, then box a bare number." - -Hard rules: never mention or use any number that appears in the original problem; never state the original problem's answer; keep the whole block under 100 words. -""" - -# pitfall 主链路：探针 P5_pitfall 原文（预判最可能错误走向并拦截）。 -SKILL_GEN_PITFALL = """\ -You are a skill-generation model. A separate executor model will solve the problem; it only sees your <skills> block. - -First think privately: solve the problem in your head AND identify the single most likely way a solver goes wrong on this type (a tempting but wrong turn, an off-by-one, a wasteful brute-force, a wrong branch). Then, inside <skills></skills>, write under 90 words: -- WARNING: name that most likely mistake concretely and say why it is wrong. -- INSTEAD: one or two sentences pointing to the correct turn (technique name + where to apply it), without solving the problem or revealing any numeric result. -- End with: "Avoid re-checking loops; box a bare number as soon as it is computed." -""" - -# freeform 主链路（2026-08-01 用户拍板）：不锁死 narrative/pitfall/toy 固定文体，而是给模型一份 -# "招式菜单"，让它按题自选最有用的形态（可组合、可极简）。设计目标是让 T=1.0×8 的候选自然铺开 -# 到不同形态（分析 / 概念 / 预判纠错 / 迷你示范 / 直白执行指令，甚至 "let's think step by step"）， -# 由 GRPO/BNPO 组内择优。依据：固定 hint 消融（good_skill_hard_fail/fixed_hint_probe.py）显示 hint -# 的"内容语义"贡献≈0、增益几乎全来自"存在一个 skill 块 + 催答案收尾"（A7_budget 最高 +0.16、 -# exec_answered 3.95σ），所以放开文体、把"催收尾"作为可选招式之一，看模型能否自选出更优组合。 -# 硬约束仍与 narrative 一致：只输出 <skills> 块、不解题、不代入本题数值、不给最终答案。 -SKILL_GEN_FREEFORM = """\ -You are a skill-generation model. Your <skills> block will be fed to a SEPARATE downstream executor model that must solve the problem on its own. The executor will NOT see your private reasoning — it only sees what is inside <skills>...</skills>. - -First, think privately: actually work the problem out in your head until you understand what really makes it solvable, then decide what ONE kind of help would most raise a fresh solver's chance on THIS specific problem. - -There is NO fixed format and no required style. Different problems are helped by different things — pick whatever you judge most useful here. Any of the following is allowed (the list is not exhaustive, and you may blend a couple if they genuinely help): -- a short transferable analysis of what this TYPE of problem is really asking and the recommended approach; -- naming the key concept / theorem / trick to reach for; -- a WARNING about the single most likely wrong turn on this type, and the correct move instead; -- a tiny worked example of the SAME type using DIFFERENT, smaller numbers (never the problem's own); -- a blunt execution directive that keeps the solver on track (e.g. "commit to one method and don't keep second-guessing", "let's think step by step", or "box a bare number as soon as it is computed"); -- or plain, nothing-fancy encouragement if that is honestly all this problem needs. -- Any other freeform skill you can imagine to try on this query - -Choose the form that fits THIS problem; do not pad. If one sharp sentence is the best help, give only that sentence; if a short focused paragraph is warranted, keep it tight. Being genuinely useful matters far more than being long or elaborate. - -Hard rules (always apply, whatever form you pick): -- Do NOT solve the problem for the executor. Do NOT reveal or compute the final answer, and do NOT substitute the problem's specific given numbers or state any intermediate numeric result — leave ALL concrete numbers for the executor to compute. -- Put ONLY your chosen help inside <skills></skills>, and nothing else. - -Whatever form you choose, it MUST be wrapped in a single <skills></skills> block with a proper closing tag. For example, a rich form: -<skills> -This is a modular-arithmetic problem: reduce each factor modulo the given modulus before multiplying, and never expand the full product — that is the whole trick. Commit to that reduction and don't second-guess it, then box a bare number as soon as it is computed. -</skills> -or, when the problem only needs a nudge, a minimal form is equally valid: -<skills> -Let's think step by step, and box a bare number as soon as it is computed. -</skills> -""" - -_SKILL_STYLE = 'narrative' # 'narrative' | 'toy' | 'pitfall' | 'freeform'；由 main() 依据 --skill-style 设置 - -# ---- Executor prompt (with skill injection) ---- -# 中文注释：executor 提示词。答案格式已统一为 \boxed{}（人工拍板，2026-07-27）：seam/v2 两模式的 -# 格式说明与判分口径完全一致，保证 E13(seam) 与 E1-E12 横向可比；prompt 结构差异（seam 嵌套/ -# system+user vs v2 单 user）作为方案级差异保留。_ANSWER_FORMAT(<answer> 版) 已弃用。 -_ANSWER_FORMAT = ('Present your reasoning and answer in the following format:\n' - '<think> Content of Thinking</think><answer>[Final numeric result only]</answer>') -# 统一执行器答案格式：把最终数值放进 \boxed{}；判分对应走 extract_boxed。 -_ANSWER_FORMAT_V2 = ('Present your reasoning, then put ONLY the final numeric result inside ' - '\\boxed{}. For example: \\boxed{42}.') -DIRECT_SYSTEM = ( - 'You are an expert competition mathematician. Be concise and accurate. ' - + _ANSWER_FORMAT_V2) -_SKILL_SOLVE_PREFIX = ( - 'You are an expert competition mathematician. Be concise and accurate.\n\n' - 'Before you start, keep these reminders in mind to avoid common mistakes on this ' - 'type of problem:\n') -_SKILL_SOLVE_SUFFIX = ('\nApply them where relevant, but rely on your own reasoning to reach the answer.\n\n' - + _ANSWER_FORMAT) - - -# ---- Unified SEAM-alignment mode (toggle: prompt + skill-gen format only) ---- -# 中文注释：SEAM 对齐开关（由 --align-mode 控制）。 -# 1) executor 输入：seam=复刻 SEAM reward worker 的非空 experience 路径： -# prompt_text + actor 原始 response_text（保留 <think> 与 <memory_item>）+ 解题 advisory。 -# v2=干净单 user turn（题目 + “Skill hint”advisory + 答案格式；空 skill 回退 direct）。 -# 2) skill-gen prompt：seam=SEAM EXPERIENCE_PROMPT（单 user turn，输出 <memory_item>）；v2=SKILL_GEN_SYSTEM（<skills>）。 -# 注：reward 判分（_parse_seq）与 loss 聚合（set_loss）已统一为 SEAM 口径（lpem 纯数值匹配 + token-mean）。 -# 关于强制 <think>：SEAM 原版在 executor 末尾裸拼 "<think>\n"。v2 走 messages+模板路径，base_sampler -# enable_thinking=True 时 Qwen3 生成起点默认进入 thinking；这里不改共享 sampler 的 assistant 前缀注入。 -_ALIGN_MODE = 'v2' # 'v2' | 'seam'；由 main() 依据 --align-mode 设置 - -# =========================================================================================== -# SEAM 对齐定案（2026-08-02，E13 = --align-mode seam） -# 目的：把追查过程固化下来，避免以后重复审查同样的地方。 -# -# 【已逐条/逐字节校验对齐的 12 项】 -# 1. 训练喂题序列 —— verl dataloader data.shuffle 默认 True，SEAM 的 step k 不是 -# train.parquet[k*128:(k+1)*128]。实测 twinkle chunk0 与 SEAM step1 的 128 题交集只有 -# 1 题，单这一项就拉开 6-7 个 acc 点。已从 SEAM rollout dump 反推真实序列，用 -# --train-order-file + ProblemPool(fixed_order=True) 钉死（5120 行，missing=0）。 -# 2. actor prompt / executor prompt（with-skill 与 baseline 两路）—— 含 chat template 与 -# Qwen3 的空 <think>\n\n</think>\n\n，逐字节 True。 -# 3. templates/slove_qwen.txt advisory —— 逐字节 True（含第一行末尾那个空格、EOF 无换行）。 -# 4. 判分口径 —— SEAM_JUDGE=answer 的整段级联（<answer> -> \boxed -> $..$ -> 分数 -> 首个 -# 数字）。用 SEAM dump replay 20480/20480 与其 acc 逐条一致；旧的 boxed-only 判分对 -# SEAM 自己的输出只给 acc≈0.05。 -# 5. format 抽取 —— 整段搜首个 <memory_item>，20480/20480 与 SEAM format 一致。 -# 6. loss 聚合 —— BNPO token_mean_scope='micro'。verl 侧证据：dp_actor.py 在每个 micro 内部 -# masked_mean，再乘 loss_scale_factor = n_micro / ppo_mini_batch_size(=40/卡)，满 micro 时 -# 恰好等权；即 verl 也是「每 5 条组等权」而不是全局 token-mean。 -# 7. 聚合粒度 —— verl 的等权单元是 ppo_micro_batch_size_per_gpu=5 条。TRAIN_FSDP=2 腾出显存后 -# train_micro_batch=5、TRAIN_DP=1，每 optimizer step 32 个「5 条组」等权，与 verl 逐组一致。 -# 8. 每 step 的 optimizer step 数 —— verl 的 ppo_mini_batch_size 经 fsdp_workers.py:198-199 -# 归一化为 per-GPU 40，每卡 256 条 => 7 个 optimizer step；twinkle mini=160、n=1024 也是 7。 -# 9. PPO 语义 —— clip_ratio 0.2 / clip_ratio_c 3.0 / ppo_epochs 1 / entropy_coeff 0 / -# use_kl_loss+low_var_kl+coef 0.001 / clip_grad 1.0 / AdamW wd 0.01 全部一致；old_logps 在 -# 任何更新前统一预计算（multi_step=True），所以 7 个 mini-step 的 PPO clip 真实有约束力。 -# 10. advantage —— A=(r-mean)/std（unbiased std，norm_adv_by_std_in_grpo 语义）；drop_zero_adv=False， -# 零 adv 样本照样占分母。实测 |adv| 均值与 SEAM adv_absmax_mean 逐 step 吻合（0.20-0.28）。 -# 11. 数值精度 —— fp32 master + bf16 计算（torch_dtype='float32' + mixed_precision='bf16'）， -# 与 verl FSDP 的 fp32 master + param_dtype=bf16 等价；lr 恒定 1e-6 无 warmup/decay。 -# 12. 采样超参 —— actor temperature 1.0 / top_p 1.0 / top_k 禁用；executor greedy + nothink； -# 8192/8192 预算，executor prompt 实测 ≤8522 < 10752 不触发左截断。 -# -# 【残余偏离已定案并修复：截断 rollout 的思考段被 chat template 补了一个空思考块】 -# 现象：两边起点几乎同一点（actor 输出 3851 vs 3842 tokens、format 0.874 vs 0.883），但 -# twinkle 6 个 chunk 就走完 SEAM 40 步的长度降幅（3851->3177 vs 3842->3239，均 -16~17%）， -# 于是 format 0.874->0.957 而 SEAM 40 步才 0.883->0.943。acc/reward/lift/zero_grad/grp_std/ -# n_train 均已在 SEAM 自身 step 间噪声(sd≈0.03)内逐 step 对应。 -# 定位方法（.tmp_analysis/probe_grad.py + probe_twinkle.py）：第 1 个 optimizer step 时 ratio≡1、 -# KL≡0，梯度完全由 (advantage, logits) 决定，所以拿 SEAM step1 dump 的同一批样本做一次 -# forward+backward 就能逐层对。结果： -# * 环境差异排除 —— 同一份参考实现在 torch 2.11/tf 5.12 与 SEAM 的 torch 2.7/tf 4.53 下 -# tokenization md5 相同、loss 差 0.24%、grad_norm 差 0.9%。 -# * 逐组定位 —— 只有「非零 advantage 的样本恰好是截断样本」的 micro 组对不上（g4 -# 0.5927 -> 2.0552，3.47 倍），不含截断样本的组差 <1%。 -# * 真因 —— 撞 8192 上限的 rollout 思考段没有 </think>（160 条里 18 条截断，其中 16 条 -# 无 </think>，且无 </think> 的样本 100% 是截断样本），Template 的 pre-pipeline -# _to_standard_reasoning_content 拆不出思考段，只能置 reasoning_content=''，Qwen3 模板于是 -# 渲染成空思考块 + 原文，原文自带的 <think> 变成紧跟在闭合标签之后的 token —— 那个 -# 位置模型输出 <think> 的概率≈0，logp 极低、梯度极大。 -# * 因果验证 —— 在参考实现里复刻这一编码（probe_grad.py --emulate-think-bug）后，逐组 -# grad_norm 从 [0.9999, 0.5733, 0.5927, 0.5307] 变成 [1.2520, 0.5672, 2.0539, 0.5310]， -# 与 twinkle 实测 [1.2573, 0.5787, 2.0552, 0.5338] 逐组重合。 -# * 修复验证 —— Template._fix_unfinished_last_round 上线后 g4 2.0552 -> 0.5922 -# （ref 0.5927），整批 160 条 32 组 0.3505 -> 0.1151（ref 0.1149，差 0.2%）。 -# 为何只影响长度/format、不影响 acc/reward：截断样本 reward=0（没闭合 <memory_item>）、 -# advantage 为负，那个巨大梯度全压在“长输出”这一模式上；而截断率会随训练自我消解 -# （chunk0 11.6% -> chunk6 1.2%），所以偏差在前几步最猛、之后自行消失 —— 正好解释了 -# 「twinkle 6 步冲完然后平稳 vs SEAM 40 步缓慢上升」。 -# 另注：grad_norm 不可直接比（twinkle 只记 7 个 optimizer step 中的最后一个，verl 记均值）。 -# =========================================================================================== - -_SEAM_EXPERIENCE_PROMPT = ( - 'You are a problem-solving guidance model. Read the math problem below and ' - 'distill a concise, reusable piece of solving experience that will help a ' - 'SEPARATE solver model reach the correct answer.\n' - 'Rules:\n' - '- Do NOT solve the problem and do NOT reveal or compute the final answer.\n' - '- State the key concepts/theorems, the recommended strategy/steps, and the ' - 'common pitfalls to avoid.\n' - '- Output ONLY the experience, wrapped EXACTLY as ' - '<memory_item> ... </memory_item>.\n\n' - 'Problem:\n{problem}') -# 逐字节复刻 SEAM templates/slove_qwen.txt（注意第一行末尾那个空格，EOF 无换行）。 -_SEAM_SOLVE_ADVISORY = ( - 'The above is a Q&A dialogue between a user and a problem-solving guidance model. \n' - 'Treat the output of the guidance model as advisory context to solve the math problem: ' - 'prefer using its techniques when they fit, but you may use alternative correct methods ' - 'if they are more efficient or clearer. If you diverge from the advisory context, briefly ' - 'explain why. Be concise and accurate.\n' - + _ANSWER_FORMAT) -# seam 模式的 baseline/空-skill 回退 system。必须用 _ANSWER_FORMAT（<think>/<answer>）而不是 -# _ANSWER_FORMAT_V2（\boxed{}）：SEAM 的 fsdp_workers.py:850-856 在 SEAM_JUDGE=answer（默认）下 -# 就是这个串，且 executor 关 thinking 时 Qwen3 会自动补一个空 <think></think>，与"请输出 <think>" -# 的要求撞在一起，抽答案更容易失败 —— 这正是 SEAM baseline 只有 0.57 的原因。用 boxed 会把 -# baseline 抬到 0.72、给定可解析 skill 的 acc 抬到 0.923（SEAM 0.865），曲线水平就对不上。 -# 判分侧无需改动：_seam_sanitize 优先匹配 <answer>、其次 boxed，两种格式都吃。 -DIRECT_SYSTEM_SEAM = ( - 'You are an expert competition mathematician. Be concise and accurate. ' - + _ANSWER_FORMAT) - - -def build_skill_solve_prompt_seam(problem, skill, raw_response=None, resp_terminated=True): - """SEAM executor prompt. Non-empty skills use the actor's raw response_text, preserving - actor <think> exactly as SEAM's reward worker does: prompt_text + response_text + grm. - If raw_response is missing, fall back to reconstructing a minimal <memory_item> response.""" - skill = (skill or '').strip() - response_text = (raw_response or '').strip() - prompt_text = ('<|im_start|>user\n' - + _SEAM_EXPERIENCE_PROMPT.format(problem=problem) - + '<|im_end|>\n<|im_start|>assistant\n') - if not response_text: - response_text = f'<memory_item>{skill}</memory_item>' - elif resp_terminated: - # SEAM 的 response_text 是 skip_special_tokens=False 解码的（fsdp_workers.py:828），正常终止的 - # rollout 末尾带着 EOS，即 "</memory_item><|im_end|>"；twinkle 的 _clean_text 把 <|...|> 全剔了。 - # 差这一个 token 也会改变 executor 的贪心解码，补回来。截断的 rollout（stop=length） - # SEAM 那边也没有 EOS，所以不补。 - response_text = response_text + '<|im_end|>' - content = prompt_text + response_text + '\n' + _SEAM_SOLVE_ADVISORY - return {'messages': [{'role': 'user', 'content': content}]} - - -def build_direct_prompt(problem): - if _TASK == 'code': - return code_task.direct_prompt(problem) - if _ALIGN_MODE == 'seam': - # seam 基线逐字复刻 SEAM fsdp_workers.py:850-858（system + user，<think>/<answer> 格式） - return {'messages': [{'role': 'system', 'content': DIRECT_SYSTEM_SEAM}, - {'role': 'user', 'content': problem}]} - # v2：英文 executor 基线——与带 skill 版同格式，仅去掉“技巧提示”部分 - content = f'The problem you need to solve:\n{problem}\n\n' + _ANSWER_FORMAT_V2 - return {'messages': [{'role': 'user', 'content': content}]} - - -def build_skill_solve_prompt(problem, skill, raw_response=None, resp_terminated=True): - skill = (skill or '').strip() - if _TASK == 'code': - # 空 skill -> 干净 direct（与数学分支同规则，见下方注释） - return code_task.skill_solve_prompt(problem, skill) - if not skill: - # 空 skill → 干净 direct。训练侧根本不会用空 skill 走 executor（process_chunk 只对非空 flat 跑， - # 空候选直接 reward=0），故此分支仅影响 eval 口径——让空 skill 题 withskill==baseline、对 lift 贡献 0， - # 去掉空壳嵌套的框架水分，指标更干净。 - # （seam 模式下这正好等于 SEAM fsdp_workers.py:846-858 的 else 分支。） - return build_direct_prompt(problem) - if _ALIGN_MODE == 'seam': - # 非空 skill 走 SEAM 原始 reward worker 路径：executor 可见 actor 完整 response_text（含 <think>）。 - return build_skill_solve_prompt_seam(problem, skill, raw_response=raw_response, - resp_terminated=resp_terminated) - # v2：英文 executor——题目 + 技巧提示(skill 作为 advisory) + 答案格式，单 user turn - content = (f'The problem you need to solve:\n{problem}\n\n' - 'Skill hint:\nFor this problem, a skill-generation model has analyzed it and ' - 'provided some advisory skills:\n' - f'{skill}\n' - 'Prefer using its techniques when they fit, but if you have a more efficient or ' - 'clearer correct method, you may use it. If you diverge from this advice, briefly ' - 'explain why. Be concise and accurate.\n' - + _ANSWER_FORMAT_V2) - return {'messages': [{'role': 'user', 'content': content}]} - - -# ---- Rubric-guided regeneration prompt (buffer B distillation) ---- -# 中文注释：蒸馏重生成提示词。给旧 skill + rubric 诊断，要求产出改进后的 skill（<skills>）。 -# 输出要求第一人称自持句式、不指向外部上下文（防幻觉），含一个连贯叙述式示例。 -REGEN_SYSTEM = """\ -You are a skill-generation model. Your skill will be fed to a downstream executor model to help it solve the problem better. -You may give general, transferable solving techniques, together with why you give this advice, so the downstream model can follow it. Do NOT reveal or compute the final answer. - -You previously generated a skill, but that skill did not help the model. The executor's actual solving process has now been analyzed. You need to regenerate the skill based on your previous skill and the mistakes the model actually made, so as to help the model solve this problem. - -Your steps: -1. Re-read and understand the original problem. -2. Tell a coherent analysis story for this problem as one flowing narrative: first identify what it is essentially asking, then walk through how to approach it, naturally weaving together the solving points that were already correct last time, the pitfalls that actually tripped up the solving process and how to avoid them, and your reasoning for why you give this advice, blended into a single connected story, and leave the concrete numbers for the downstream solver to compute. -3. End the block with this exact sentence: "Avoid re-checking loops; box a bare number as soon as it is computed." -4. Put the above inside <skills></skills>. - -Output requirement: Write your judgments and pitfall reminders about this problem directly in the first person (e.g. "I think this step tends to ...", "A common mistake is ..., so you need to ..."), and phrase the issues you find as self-contained, general techniques. Do NOT use phrasings that point to external context such as "according to the given analysis/hints" or "the previous skill" — the downstream executor cannot see that context, and such phrasings will cause hallucination. - -Example: -<skills> -This problem is essentially asking "how many arrangements satisfy the given constraints", which is a counting problem; first get clear on "what exactly is being counted" before deciding whether to use permutations or combinations. Since it is counting, you should first clearly define the objects being counted and the constraints, and judge whether the elements are distinguishable and whether order matters, because this directly determines whether you will need to divide out duplicates later. Next, first compute a total as if things were "ordered/distinguishable", then find which seemingly different arrangements actually correspond to the same configuration. The part about "recognizing symmetry and determining the duplication factor" is important here: I think the step most likely to go wrong in this problem is ignoring symmetry and treating essentially identical configurations as different, which makes the result too large; I think it is also easy to directly miss the "divide by the duplication factor" step — as long as the choices can be interchanged, you must divide out duplicates, otherwise you overcount. Finally, divide the total by the duplication factor to get the truly non-duplicated count; here pay special attention not to jump straight to permutation/combination formulas, but first think clearly about whether the elements are distinguishable and then decide whether to divide out duplicates. Overall, I summarize the approach for this kind of problem as "first recognize that it is a counting problem and judge whether the elements are distinguishable, then compute the total, recognize symmetry and remove duplicates", because I judge that the loss points for such problems almost all concentrate on overcounting; while leaving the concrete numbers for the downstream solver to substitute and compute on its own. Avoid re-checking loops; box a bare number as soon as it is computed. -</skills>""" - -REGEN_USER = """\ -Original problem: -{problem} - -Previously generated skill (did not help the executor): -{orig_skill} - -Analysis of the executor's actual solving process: -{rubric_diag} - -Now rewrite the improved <skills> guidance:""" - -# 中文注释：buffer B regen 的 toy/pitfall 文体版（与主链路同文体，保证训练分布一致）。 -# 源自 reflexion 探针 D3_toyfix / D1_needle（diag-only 口径），另加防指涉硬规则 -# （不许写 "according to the diagnosis / the previous skill"，executor 看不到这些上下文）。 -REGEN_TOY_SYSTEM = """\ -You are a skill-generation model. A separate executor model previously FAILED this problem even with your earlier skill. You will see that earlier skill and an expert rubric diagnosis of the failure. The executor will retry seeing ONLY your new <skills> block. - -First think privately: from the diagnosis, identify the ONE technique the executor got wrong. Then, inside <skills></skills>, do exactly this (under 110 words): -1. Invent a MINIATURE problem exercising that same technique with DIFFERENT, much smaller numbers, and solve the miniature completely in at most 5 short lines, making the correct move (the one the failed attempt missed) explicit. -2. One transfer sentence: "Your problem has the same shape - repeat these steps with its own numbers, then box a bare number." -Hard rules: never use any number from the original problem; never state its answer; the block must be self-contained - never reference "the diagnosis", "the previous skill" or any context the executor cannot see. -""" - -REGEN_PITFALL_SYSTEM = """\ -You are a skill-generation model. A separate executor model previously FAILED this problem even with your earlier skill. You will see that earlier skill and an expert rubric diagnosis of the failure. The executor will retry seeing ONLY your new <skills> block. - -First think privately: from the diagnosis, pinpoint the decisive error. Then, inside <skills></skills>, write under 90 words: -- WARNING: the decisive mistake, stated concretely for THIS problem in self-contained first person (e.g. "I think the step most likely to go wrong is ..."). -- INSTEAD: one or two sentences pointing to the correct turn (technique name + where to apply it), without solving the problem or revealing any numeric result. -- End with: "Avoid re-checking loops; box a bare number as soon as it is computed." -Hard rules: the block must be self-contained - never reference "the diagnosis" or "the previous skill"; the executor cannot see them. -""" - -# freeform 的 regen 版（buffer B 蒸馏用）。bnpo/view-B 臂不会走 regen，此处仅为分派完整性与 -# 未来 view-A + freeform 组合预留；同样放开形态、保留"自持、不指涉外部上下文、不泄漏"硬规则。 -REGEN_FREEFORM_SYSTEM = """\ -You are a skill-generation model. A separate executor model previously FAILED this problem even with your earlier skill. You will see that earlier skill and an expert rubric diagnosis of the failure. The executor will retry seeing ONLY your new <skills> block. - -First think privately: from the diagnosis, pinpoint the ONE thing that actually went wrong. Then choose whatever form of help would best fix it for THIS problem — there is no fixed format. It may be a short transferable analysis, the key concept to reach for, a WARNING naming the decisive mistake plus the correct move instead, a tiny worked example with DIFFERENT smaller numbers, or a blunt execution directive (e.g. "box a bare number as soon as it is computed"). Blend a couple only if it genuinely helps, and do not pad. - -Hard rules: -- Do NOT solve the problem or reveal/compute the final answer, and do NOT substitute the problem's own numbers. -- The block must be self-contained — never reference "the diagnosis", "the previous skill", or any context the executor cannot see, or it will hallucinate. -- Put ONLY your chosen help inside <skills></skills>. -""" - - -def _skillgen_prompt(problem: str) -> Dict[str, Any]: - """Skill-gen prompt: query-only. seam mode uses SEAM EXPERIENCE_PROMPT (single user turn, - <memory_item> output); v2 uses SKILL_GEN_SYSTEM (<skills>).""" - if _TASK == 'code': - # 代码域只做 narrative 一种文体（E4/E17 都是 narrative；toy/pitfall 未移植） - return code_task.skillgen_prompt(problem) - if _ALIGN_MODE == 'seam': - return {'messages': [{'role': 'user', 'content': _SEAM_EXPERIENCE_PROMPT.format(problem=problem)}]} - # 中文注释：按 --skill-style 选主链路文体（narrative=现版叙述式 / toy / pitfall）。 - sys_p = {'toy': SKILL_GEN_TOY, 'pitfall': SKILL_GEN_PITFALL, - 'freeform': SKILL_GEN_FREEFORM}.get(_SKILL_STYLE, SKILL_GEN_SYSTEM) - return {'messages': [ - {'role': 'system', 'content': sys_p}, - {'role': 'user', 'content': f'Problem:\n{problem}'}]} - - -def _regen_prompt(problem: str, orig_skill: str, rubric_diag: str) -> Dict[str, Any]: - """Regeneration prompt for buffer B distillation.""" - # 中文注释：regen 与主链路同文体（--skill-style），user 模板复用 REGEN_USER 三字段。 - sys_p = {'toy': REGEN_TOY_SYSTEM, 'pitfall': REGEN_PITFALL_SYSTEM, - 'freeform': REGEN_FREEFORM_SYSTEM}.get(_SKILL_STYLE, REGEN_SYSTEM) - return {'messages': [ - {'role': 'system', 'content': sys_p}, - {'role': 'user', 'content': REGEN_USER.format( - problem=problem, orig_skill=orig_skill, rubric_diag=rubric_diag)}]} - - -# ---- Reward ---- -# 中文注释：reward = parseable × 通过率（对齐 SEAM lpem：去 terminated、去长度惩罚）。 -# parseable=0 的候选 reward=0 仍参与 group（格式压力）。correct 兼容 bool（greedy 0/1， -# E1-E13）与 float 通过率（E14 多 rollout 判分，见 process_chunk reward_rollouts）。 -def _skill_reward(parseable: bool, correct) -> float: - return float(correct) if parseable else 0.0 - - -# ---- Buffer A: collect adv=0 all-fail problems ---- -def _collect_buffer_a(chunk, args) -> List[Dict[str, Any]]: - """Collect problems where all candidates got reward 0 (adv=0, GRPO blind spot). - Store one representative failure trajectory for later rubric diagnosis.""" - entries = [] - for r in chunk: - cs = [c for c in r['_cands'] if c.get('reward') is not None] - if len(cs) < 2: - continue - rewards = [c['reward'] for c in cs] - if max(rewards) > 0: - continue # has signal, not all-fail - # Representative trajectory for rubric + regen seed: prefer a terminated-wrong - # parseable candidate (complete reasoning to diagnose). Its skill becomes the regen - # seed, so pick the MOST SUBSTANTIAL one WITHIN budget (longest ≤ len_budget) — a - # rich-but-not-bloated starting point — rather than an arbitrary [0] or a near-empty - # skill. If all seeds exceed budget, take the one closest to budget (shortest-over). - # 中文注释：代表轨迹既做 rubric 诊断又做 regen 种子——优先"跑完但答错"的候选（完整推理）， - # 其 skill 取预算内最长(最有实质)的作种子；若全超预算则取最接近预算的，避免随机/近空种子。 - budget = args.len_budget - - def _seed_key(c): - L = len(c.get('skills') or '') - return (L <= budget, L if L <= budget else -L) - - parseable = [c for c in cs if c.get('skills')] - term_wrong = [c for c in parseable if c['rolls'] and c['rolls'][0].get('terminated')] - pool_c = term_wrong or parseable or cs - rep = max(pool_c, key=_seed_key) - stop_dist = {} - for c in cs: - sr = c['rolls'][0]['stop_reason'] if c['rolls'] else 'none' - stop_dist[sr] = stop_dist.get(sr, 0) + 1 - entries.append({ - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'data_id': r.get('data_id', ''), - 'orig_skill': rep.get('skills', ''), - 'orig_len': len(rep.get('skills', '')), - 'fail_segment': rep['rolls'][0]['text'] if rep['rolls'] else '', - 'fail_stop_reason': rep['rolls'][0]['stop_reason'] if rep['rolls'] else 'none', - 'stop_reason_dist': stop_dist, - }) - return entries - - -# ---- Buffer B distillation ---- -def distill_buffer(entries: List[Dict[str, Any]], skill_sampler, base_sampler, - checker, skill_dp: int, base_dp: int, - args) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: - """Batch rubric → regenerate K distinct skills → greedy-validate → return SFT records. - 中文注释：蒸馏流程（方案 B，多样性在 skill 侧、executor 用贪心）： - 1. 批量 rubric 诊断失败轨迹；2. 仅 [FAIL] 项用高温重生成 K 个不同候选 skill； - 3. 每个候选过 ≤budget+无leak+去重 过滤；4. 每个存活候选用 executor 贪心(T=0)解 1 次； - 5. gate：≥m 个不同候选达成 terminated-correct → select 长度最接近 budget 的一个入 buffer B。 - 返回 (sft_records, distill_records)：后者逐 entry 记录 rubric_diag/候选 skill/贪心解结果/漏斗 - stage，落盘到 distill_records.jsonl 供复盘（否则 rubric 诊断与候选明细只存在于内存）。""" - if not checker or not entries: - return [], [] - - # Step 1: ensure every entry has a rubric diagnosis. Entries pre-diagnosed in the - # background (see _prediagnose in main) already carry '_rubric_diag'; only the misses - # are diagnosed here (in parallel), so the GPU-idle API wait is normally hidden. - # 中文注释：优先用后台预诊断结果；只对没预诊断到的条目并行补跑，隐藏 API 等待。 - pending = [e for e in entries if not e.get('_rubric_diag')] - if pending: - workers = min(args.rubric_workers, len(pending)) - with ThreadPoolExecutor(max_workers=max(1, workers)) as ex: - diags = list(ex.map(lambda e: _diagnose_entry(checker, e), pending)) - for entry, diag in zip(pending, diags): - entry['_rubric_diag'] = diag or '' - for entry in entries: - entry['rubric_diag'] = entry.get('_rubric_diag') or '' - - # Builder for the structured distill audit records (one per buffer-A entry). Closure over - # `entries`/`args`; takes the per-entry regen skills + greedy solve results (may be empty - # for the early-exit funnel stages). 中文注释：构造逐 entry 的蒸馏审计记录（含漏斗 stage）。 - def _mk_distill(results_by_entry, per_entry_skills, has_fail): - hf_index = {id(e): ei for ei, e in enumerate(has_fail)} - recs = [] - for e in entries: - ei = hf_index.get(id(e)) - cand_results = results_by_entry.get(ei, []) if ei is not None else [] - n_pass = sum(1 for c in cand_results if c['correct'] and c['terminated']) - n_cand = len(per_entry_skills[ei]) if (ei is not None and ei < len(per_entry_skills)) else 0 - if ei is None: - stage = 'no_fail' # rubric 未给出任何 [FAIL] - elif n_cand == 0: - stage = 'no_valid_regen' # 有 [FAIL] 但重生成无一条过 ≤budget/无leak/去重 - elif n_pass >= args.passatk_m: - stage = 'accepted' # ≥m 个候选贪心解对 → 入 buffer B - else: - stage = 'rejected' # 有候选但 <m 个解对 - recs.append({ - 'record_type': 'distill', 'stage': stage, - 'data_id': e.get('data_id', ''), 'problem': e['problem'], - 'reference_answer': e['reference_answer'], - 'orig_skill': e.get('orig_skill', ''), 'orig_len': e.get('orig_len', 0), - 'fail_stop_reason': e.get('fail_stop_reason', ''), - 'rubric_diag': e.get('rubric_diag', ''), - 'n_cand_skills': n_cand, 'n_pass_skills': n_pass, - 'candidates': cand_results, # each: {skill, len, correct, terminated} - }) - return recs - - # Step 2: filter to entries with [FAIL] diagnosis - has_fail = [e for e in entries if '[FAIL]' in e.get('rubric_diag', '')] - if not has_fail: - logger.info('[distill] no [FAIL] diagnoses, skipping regen') - return [], _mk_distill({}, [], []) - - # Step 3: regenerate K DISTINCT candidate skills per [FAIL] entry (skill-model, high T) and - # greedy-validate each. Entries that do NOT yet have >= m distinct greedy-passing skills are - # RETRIED for up to --distill-retries extra rounds: regenerate more skills, dedup against those - # already seen for that entry, greedy-solve, and accumulate — rescuing more problems into buffer B. - # 中文注释：对每条 [FAIL] 高温采 K 个不同候选 skill 并贪心验证；还没凑够 m 个“贪心解对”的条目， - # 再重生成 --distill-retries 轮（新候选去重、贪心解、累计），把更多题救进 buffer B。 - k = args.passatk_k - per_entry_skills: List[List[str]] = [[] for _ in has_fail] # 累计去重候选(跨轮) - results_by_entry: Dict[int, List[Dict[str, Any]]] = {} # 累计每候选贪心结果(跨轮) - passers_by_entry: Dict[int, Set[str]] = {ei: set() for ei in range(len(has_fail))} - pending_ei = list(range(len(has_fail))) # 还没凑够 m 个 passer 的条目 - for _ in range(args.distill_retries + 1): - if not pending_ei: - break - regen_prompts = [_regen_prompt(has_fail[ei]['problem'], has_fail[ei]['orig_skill'], - has_fail[ei]['rubric_diag']) for ei in pending_ei] - regen_out = _run_samples(skill_sampler, regen_prompts, k, args.skill_max_tokens, skill_dp, - temperature=args.passatk_skill_temp, top_p=args.passatk_skill_top_p) - # 过滤(可解析/≤budget/无leak/对本条目去重) → 收集本轮新候选 - new_flat_idx, new_flat_prompts = [], [] - for ei, seqs in zip(pending_ei, regen_out): - seen = set(per_entry_skills[ei]) - for s in (seqs or []): - resp = _clean_text(getattr(s, 'decoded', '') or '') - skill = _extract_skill(resp) - if not skill or len(skill) > args.len_budget: - continue - if _answer_leaked(skill, has_fail[ei]['reference_answer']): - continue - if skill in seen: - continue - seen.add(skill) - per_entry_skills[ei].append(skill) - new_flat_idx.append((ei, skill)) - new_flat_prompts.append(build_skill_solve_prompt(has_fail[ei]['problem'], skill)) - # 本轮新候选各用 executor 贪心(T=0)解 1 次，累计结果与 distinct passer - if new_flat_prompts: - solve_out = _run_samples(base_sampler, new_flat_prompts, 1, args.max_tokens, base_dp, - temperature=0.0) - for (ei, sk), seqs in zip(new_flat_idx, solve_out): - roll = _parse_seq(seqs[0], has_fail[ei]['reference_answer']) if seqs else _empty_roll() - results_by_entry.setdefault(ei, []).append( - {'skill': sk, 'len': len(sk), 'correct': roll['correct'], 'terminated': roll['terminated']}) - if roll['correct'] and roll['terminated']: - passers_by_entry[ei].add(sk) - # 仍不足 m 个 distinct passer 的条目进入下一轮重试 - pending_ei = [ei for ei in pending_ei if len(passers_by_entry[ei]) < args.passatk_m] - - n_entries_with_cands = sum(1 for sk in per_entry_skills if sk) - if not results_by_entry: - logger.info(f'[distill] {len(has_fail)} [FAIL] entries, 0 valid regen skills') - return [], _mk_distill({}, per_entry_skills, has_fail) - - # Step 4: gate ≥ m distinct greedy-effective skills; select the survivor CLOSEST to the - # length budget (short is the floor, but not so short it degrades to answer-dumping). - # 中文注释：gate——≥m 个不同 skill 在贪心下 terminated-correct；select——在通过的候选里 - # 选长度最接近 budget 的一个入 buffer B（短是地板，但别短到退化成吐答案）。 - sft_records = [] - for ei, cand_results in results_by_entry.items(): - passers = [c['skill'] for c in cand_results if c['correct'] and c['terminated']] - if len(passers) < args.passatk_m: - continue - entry = has_fail[ei] - best = min(passers, key=lambda sk: abs(len(sk) - args.len_budget)) - sft_records.append({ - 'problem': entry['problem'], 'reference_answer': entry['reference_answer'], - 'data_id': entry.get('data_id', ''), - 'response': f'<skills>\n{best}\n</skills>', - 'skills': best, 'sft': True, - 'n_pass_skills': len(passers), 'n_cand_skills': len(per_entry_skills[ei]), - }) - - logger.info(f'[distill] {len(entries)} A → {len(has_fail)} [FAIL] → ' - f'{n_entries_with_cands} w/cands → {len(sft_records)} validated B ' - f'(gate m={args.passatk_m}/k={k})') - return sft_records, _mk_distill(results_by_entry, per_entry_skills, has_fail) - - - -# =========================================================================== -# Section G — GRPO advantages + training -# =========================================================================== -def _assign_advantages(chunk, args): - """Group-relative advantage: A = (R - mean) / (std + eps). std==0 → adv=0 (skipped).""" - eps = 1e-6 - adv_clip = abs(float(getattr(args, 'adv_clip', 0.0) or 0.0)) - for r in chunk: - for c in r['_cands']: - c['advantage'], c['kept'] = 0.0, False - cs = [c for c in r['_cands'] if c.get('reward') is not None] - if len(cs) < 2: - continue - rewards = [c['reward'] for c in cs] - mean_r = sum(rewards) / len(rewards) - # SEAM/verl uses torch.std's default unbiased=True for GRPO group std. - import torch - std = float(torch.std(torch.tensor(rewards, dtype=torch.float32)).item()) - if std < 1e-9: - continue - for c in cs: - raw = (c['reward'] - mean_r) / (std + eps) - c['advantage'] = max(-adv_clip, min(adv_clip, raw)) if adv_clip > 0 else raw - c['kept'] = c['reward'] > mean_r - - -# =========================================================================================== -# 训练样本构造：response 段一律直接用采样返回的 token，严禁 decode 后重新过模板 -# =========================================================================================== -# 为什么（2026-08-02 定案，.tmp_analysis/probe_token_direct.py 在 160 条真实样本上实测）： -# 把采样产出 decode 成文本、再塞回 messages 让 chat template 重新渲染一遍，可训练区就不再 -# 等于模型真实生成的 token —— 模板会给 assistant 角色行带一个换行、给每条 message 补结尾 -# EOS，思考段没闭合时还会先渲染一个空思考块。实测 160/160 条的可训练区都比采样 token 多； -# 截断样本更被塞进一个紧跟 </think> 之后的 <think>（p≈0、logp≈-20），单个 micro 的 -# grad_norm 从 0.59 抬到 2.05，整批放大 3 倍，直接把 format_rate 曲线提前顶到 0.99。 -# -# 正确形状：prompt 段照常编码（它本来就是模板产物，重编码无风险），response 段原样拼采样 -# token。同一条实测里这条路径 160/160 逐 token 相等。采样器已经把 token 备好了 -# （SampledSequence.tokens），所以候选记录只需把它带下来。 -_ENCODE_TEMPLATE = None - - -def set_encode_template(template) -> None: - """注入 skill 模型的客户端 Template 副本，供 build_train_feature 编码 prompt 段。""" - global _ENCODE_TEMPLATE - _ENCODE_TEMPLATE = template - - -def build_train_feature(prompt_messages, tokens, template=None): - """prompt 段编码 + 原样拼采样 token，返回可直接喂模型的 InputFeature。 - - labels 只盖住 ``tokens`` 那一段（prompt 段全 -100），与采样端真实生成的 token 逐一对应。 - ``template`` 默认用 skill 模型那份；executor 侧打分（E14 的 logP reward）要传自己的。 - """ - tmpl = template or _ENCODE_TEMPLATE - assert tmpl is not None, 'call set_encode_template() before building train features' - feat = tmpl.encode({'messages': [dict(m) for m in prompt_messages]}, add_generation_prompt=True) - return tmpl.concat_input_feature(feat, [int(t) for t in tokens]) - - -def _train_trajectory(rec): - """query-only skill-gen prompt + 采样产出。 - - GRPO 记录带 ``tokens``（采样端 vLLM 真实吐出的 token id），直接拼成 InputFeature； - SFT 记录的 response 是程序合成的 ``<skills>`` 文本（本来就不是采样产出、没有对应 token）， - 只能走 messages 编码，``key_rounds`` 标出最后一轮为唯一可训练区。 - """ - msgs = _skillgen_prompt(rec['problem'])['messages'] - if rec.get('tokens'): - return build_train_feature(msgs, rec['tokens']) - return {'messages': msgs + [{'role': 'assistant', 'content': rec['response']}], - 'user_data': pack_user_data({'key_rounds': [len(msgs)]})} - - -def _train_step(skill_model, ref_model, ckpt, samples, args): - """On-policy GRPO update over one batch, then sync weights. SFT samples ride the - same BNPOLoss with a positive constant advantage (--sft-weight).""" - # 过滤空/纯空白 response：其可训练 token 为 0，会让持有它的 DP rank 跳过 backward， - # 与对端 all-reduce 失步 → NCCL 死锁（find_unused_parameters=False）。高熵采样偶发首 token 即 EOS。 - n_in = len(samples) - samples = [rec for rec in samples - if (rec.get('tokens') if rec.get('tokens') is not None else (rec.get('response') or '').strip())] - n_empty = n_in - len(samples) - trajs = [_train_trajectory(rec) for rec in samples] - advs = [float(rec['advantage']) for rec in samples] - smp_all = [list(rec.get('logprobs') or []) for rec in samples] - # drop_last 到 TRAIN_DP 整数倍：每个 micro（末尾那个可短于 sft）仍能被 dp 均分，零 padding 假样本。 - # 只丢尾部 ≤ dp-1 条真样本；若整批不足 dp（n<dp）则直接跳过该步。 - n_keep = (len(trajs) // TRAIN_DP) * TRAIN_DP - if n_keep == 0: - return {'n_samples': n_in, 'n_sft': 0, 'n_grpo': 0, 'n_empty': n_empty, - 'n_steps': 0, 'n_micro_batches': 0, 'metric': {}} - trajs, advs = trajs[:n_keep], advs[:n_keep] - # micro 尺寸可由 train_micro_batch 覆盖（ablate think/8192 实验防 backward OOM）； - # 默认 0 = 跟随 sft_batch_size，主训练行为不变。梯度按 micro 数归一，切细数学等价。 - n = len(trajs) - sft = getattr(args, 'train_micro_batch', 0) or args.sft_batch_size - mini = args.ppo_mini_batch_size if args.ppo_mini_batch_size > 0 else n - mini = max(sft, (mini // sft) * sft) - multi_step = mini < n - micro_ref, micro_old, micro_smp = [], [], [] - for i in range(0, n, sft): - mb = trajs[i:i + sft] - micro_ref.append(ref_model.forward_only(inputs=mb).get('logps')) - micro_old.append(skill_model.forward_only(inputs=mb).get('logps') if multi_step else None) - # 采样端 logprob，只给 GRPOMetric 做对账（sampler_logp_mae / sampler_token_delta）。 - # 整个 micro 都带齐了才传：SFT 样本的 response 是合成文本、根本没有采样 logprob， - # 混进去只会让 token_delta 无法解释（它的语义是「应恒为 0」）。 - smp = [smp_all[j] for j in range(i, min(i + sft, n))] - micro_smp.append(smp if all(smp) else None) - micro, n_steps = 0, 0 - for ms in range(0, n, mini): - for i in range(ms, min(ms + mini, n), sft): - k = i // sft - skill_model.forward_backward(inputs=trajs[i:i + sft], advantages=advs[i:i + sft], - old_logps=micro_old[k], ref_logps=micro_ref[k], - sampler_logps=micro_smp[k]) - micro += 1 - skill_model.clip_grad_and_step() - n_steps += 1 - ckpt.sync_weights(merge_and_sync=True) - metric = skill_model.calculate_metric(is_training=True) - n_sft = sum(1 for s in samples if s.get('sft')) - return {'n_samples': n_in, 'n_sft': n_sft, 'n_grpo': len(samples) - n_sft, 'n_empty': n_empty, - 'n_steps': n_steps, 'n_micro_batches': micro, - 'metric': {k: (float(v) if _is_num(v) else v) for k, v in (metric or {}).items()}} - - -def _is_num(v): - try: - float(v); return True - except (TypeError, ValueError): - return False - - -# =========================================================================== -# Section H — chunk processing, records, eval (+ hard-slice rescue) -# =========================================================================== -def process_chunk(base_sampler, skill_sampler, chunk, ci, base_dp, skill_dp, args): - """skill-gen (query-only) → leak audit → with-skill greedy pass → reward → advantages. - Returns (full_records, summary, grpo_train_records, buffer_a_entries).""" - for r in chunk: - r['_cands'] = [] - # skill-gen (single pass, no retry)：每题恒 n_skills 个候选、组大小固定，对齐 SEAM rollout.n。 - # （全 0 的难题不在这里重采，而是进 buffer A → rubric 重生成，见 distill_buffer。） - flat = [] - sg_out = _run_samples(skill_sampler, [_skillgen_prompt(r['problem']) for r in chunk], - args.n_skills, args.skill_max_tokens, skill_dp, - temperature=args.skill_gen_temperature, top_p=args.skill_gen_top_p, - top_k=args.skill_gen_top_k, logprobs=1) - for r, seqs in zip(chunk, sg_out): - for s in seqs: - resp = _clean_text(getattr(s, 'decoded', '') or '') - block = _extract_skill(resp) or '' - # tokens = 采样端真实吐出的 token id，训练样本只能用它拼（见 build_train_feature）。 - # logprobs 同长，只给 GRPOMetric 做采样/训练对账（见 sampler_logprobs），不进 loss。 - _toks = [int(t) for t in (getattr(s, 'tokens', None) or [])] - cand = {'skills': block, 'response': resp, 'parseable': bool(block), - 'leaked': None, 'with_pass': None, 'reward': None, 'rolls': [], - 'advantage': 0.0, 'kept': False, 'tokens': _toks, - 'logprobs': sampler_logprobs(s), - 'skillgen_stop': getattr(s, 'stop_reason', None), - 'skillgen_tokens': len(_toks)} - r['_cands'].append(cand) - if block: - flat.append((r, cand)) - - # leak audit (deterministic, observability only) - for r, c in flat: - c['leaked'] = _answer_leaked(c['skills'], r['reference_answer']) - - # ⭐ executor 覆盖面（seam 对齐，2026-08-02）：SEAM 对**每一条**候选都跑 executor —— - # fsdp_workers.py:842-858，抽不到 <memory_item> 时走 else 分支（direct_system + 裸题目）， - # 它的 acc 照样计入 reward_extra_info["acc"]，也就是 train/with_skill_accuracy 的分子。 - # 旧实现只跑 parseable 的，于是 twinkle 算不出同口径的 withskill 准确率，只能拿 - # P(correct|parseable) 去对 SEAM 的 P(correct)，两条曲线分母不同、根本无法对齐。 - # reward 口径不变：_skill_reward 仍然乘 parseable，空 skill 永远 reward=0。仅 seam 模式开， - # 其余臂（E1-E12/E14-E21）行为与开销不动。 - exec_list = ([(r, c) for r in chunk for c in r['_cands']] if _ALIGN_MODE == 'seam' else flat) - - # with-skill executor pass：默认 greedy×1（E1-E13，reward 0/1 与旧口径 bit 一致）； - # E14: reward_rollouts>1 时 T=reward_temperature × K 采样，reward = parseable × 通过率， - # 把内容信号从 greedy 0/1 量化里释放出来（提升组内 std>0 比例）。 - if exec_list: - K = max(1, int(getattr(args, 'reward_rollouts', 1) or 1)) - rT = float(getattr(args, 'reward_temperature', 0.0) or 0.0) - ws_out = _run_samples(base_sampler, - [build_skill_solve_prompt(r['problem'], c['skills'], c.get('response'), - resp_terminated=(c.get('skillgen_stop') != 'length')) - for r, c in exec_list], - K, args.max_tokens, base_dp, temperature=rT) - # 判分一次性批量化（code 任务要起子进程跑单测，逐条会比 GPU 还慢一个量级） - pairs, spans = [], [] - for (r, c), seqs in zip(exec_list, ws_out): - start = len(pairs) - pairs.extend((s, r['reference_answer']) for s in (seqs or [])) - spans.append((start, len(pairs))) - judged = _parse_many(pairs) - for (r, c), (a, b) in zip(exec_list, spans): - rolls = judged[a:b] or [_empty_roll()] - for x in rolls[1:]: - x['text'] = '' # 磁盘保护：K>1 时只留首 rollout 全文（gen_records 体积控制） - c['rolls'] = rolls - c['with_pass'] = sum(1.0 for x in rolls if x['correct']) / len(rolls) - c['reward'] = _skill_reward(c['parseable'], c['with_pass']) - # unparseable candidates score 0 and still join the group (format pressure) - for r in chunk: - for c in r['_cands']: - if c['reward'] is None: - c['reward'] = 0.0 - - # ⭐ 训练侧 no-skill baseline（seam 对齐，2026-08-02）。SEAM 只在第 1 个训练步跑一次 - # （ray_trainer.py:1461 `if self.global_steps == 1`，注释写明是提速、reward/梯度不受影响， - # 后续 step 的 step_summary.lift 为 None），所以 SEAM 的 train lift 只有 step1 一个点 - # （0.7891-0.6338=+0.1553）。这里逐条照抄：仅 ci==0、direct prompt、T=0 - # （= SEAM use_experience=False 分支）。 - # - # 每题必须跑 n_skills 次，不能跑 1 次再广播：SEAM 把**整个 1024 行 batch**送进 - # generate_sequences_as_grm_baseline，同一道题的 8 行是 8 个独立请求。executor 虽然是 - # greedy，vLLM 的批内非确定性（左 padding 长度随 batch 变、chunked prefill 的规约顺序） - # 让同一 prompt 的 8 次结果并不总相同 —— step1 dump 实测 128 题里有 10 题的 8 次 - # baseline_acc 不全同。所以"同题 8 条结果相同、取 1 次即可"这个旧假设是错的，按题 - # 展开取均值才与 SEAM 的 np.mean(baseline_acc) 同口径。 - if _ALIGN_MODE == 'seam' and ci == 0 and chunk: - K_b = max(1, int(args.n_skills)) - b_pairs = [r for r in chunk for _ in range(K_b)] - b_out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in b_pairs], - 1, args.max_tokens, base_dp, temperature=0.0) - b_rolls = _parse_many([(_first_seq(seqs), r['reference_answer']) - for r, seqs in zip(b_pairs, b_out)]) - for i, r in enumerate(chunk): - hits = [1.0 if x['correct'] else 0.0 for x in b_rolls[i * K_b:(i + 1) * K_b]] - r['_train_baseline_pass'] = _mean(hits) if hits else None - - _assign_advantages(chunk, args) - - # SEAM/verl 对齐：每个 dataloader batch 都进入 actor update。零 adv 候选 PG 贡献 0， - # 但仍计入 token-mean 分母，并在 beta>0 时贡献 KL 锚定；不再因整 chunk 无信号而跳过 step。 - grpo = [] - for r in chunk: - for c in r['_cands']: - if c.get('reward') is None: - continue - grpo.append({'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'data_id': r.get('data_id', ''), 'response': c['response'], - 'skills': c['skills'], 'advantage': c['advantage'], - # tokens 是训练样本的唯一来源；response 只留给 reward/审计（见 _train_trajectory） - 'tokens': c.get('tokens') or [], - 'logprobs': c.get('logprobs') or [], - 'kept': c['kept'], 'reward': c['reward'], 'sft': False}) - buffer_a = _collect_buffer_a(chunk, args) - return _full_records(chunk, ci), _chunk_summary(chunk, ci), grpo, buffer_a - - -def _roll(x): - out = {k: x[k] for k in ('pred', 'correct', 'terminated', 'stop_reason', 'gen_tokens', 'text')} - # 代码域审计字段：判分结论 / 单测报错 / 用例数（离线分析错误类型分布靠它，_trim_err 已限长） - for k in ('kind', 'error', 'n_tests'): - if k in x: - out[k] = x[k] - return out - - -def _full_records(chunk, ci): - out = [] - for r in chunk: - out.append({ - 'record_type': 'problem', 'chunk': ci, 'problem': r['problem'], - 'reference_answer': r['reference_answer'], 'data_id': r.get('data_id', ''), - 'candidates': [{ - 'skills': c['skills'], 'response': c['response'], 'parseable': c['parseable'], - 'leaked': c['leaked'], 'with_pass': c['with_pass'], 'reward': c.get('reward'), - 'advantage': c.get('advantage'), 'kept': c.get('kept'), - 'skillgen_stop': c.get('skillgen_stop'), 'skillgen_tokens': c.get('skillgen_tokens'), - 'rolls': [_roll(x) for x in c['rolls']], - 'logp_base': c.get('logp_base'), 'logp_skill': c.get('logp_skill'), - 'logp_delta': c.get('logp_delta'), - } for c in r['_cands']], - }) - return out - - -def _mean(xs): - return sum(xs) / len(xs) if xs else 0.0 - - -def _std(xs): - if len(xs) < 2: - return 0.0 - import torch - return float(torch.std(torch.tensor(xs, dtype=torch.float32)).item()) - - -def _pstd(xs): - """总体标准差（ddof=0）—— 报表用。 - - SEAM 的 step_summary 用 ``np.std``（ddof=0）算 reward_std / group_reward_std_mean - （ray_trainer.py:654,671），而 :func:`_std` 是 ``torch.std``（ddof=1）。组内只有 8 条时 - 两者差 sqrt(8/7)=1.069，足以把 group_reward_std_mean 抬高 ~0.01（c0 实测 - 0.14993 vs SEAM 0.14074；改 ddof=0 后为 0.14025）。注意分开：**advantage 的组内 - std 必须继续用 ddof=1**，verl 的 compute_grpo_outcome_advantage 用的也是 - ``torch.std`` 默认 unbiased，两边本来就一致，改了反而会把梯度弄不一致。 - """ - if len(xs) < 2: - return 0.0 - m = sum(xs) / len(xs) - return (sum((x - m) ** 2 for x in xs) / len(xs)) ** 0.5 - - -def _chunk_summary(chunk, ci): - all_cands = [c for r in chunk for c in r['_cands']] - cands = [c for c in all_cands if c['parseable']] - scored = [c for c in cands if c['with_pass'] is not None] - ws_rolls = [x for c in scored for x in c['rolls']] - # signal: fraction of groups with zero reward variance (no gradient) - group_vars, all_rewards, zero_grad, groups = [], [], 0, 0 - for r in chunk: - rewards = [c['reward'] for c in r['_cands'] if c.get('reward') is not None] - if len(rewards) < 2: - continue - groups += 1 - all_rewards.extend(rewards) - # 报表口径对齐 SEAM 的 np.std（ddof=0），见 _pstd。advantage 那边仍用 ddof=1。 - v = _pstd(rewards) - group_vars.append(v) - if v < 1e-9: - zero_grad += 1 - n_train = sum(1 for c in all_cands if abs(c.get('advantage') or 0.0) > 1e-9) - trunc = sum(1 for x in ws_rolls if x['stop_reason'] == 'length') - # bugfix（ablate #6）：旧版 any(c.get('reward')) 用 truthiness，负 reward（E16 hinge/leak_gate、 - # E14 地板 -1.0）也被当“通过”；改用 with_pass>0（真正的 executor 通过率），greedy 0/1 臂语义不变。 - # 同时限 parseable：seam 模式下 unparseable 候选也有 with_pass（走 direct 回退），不限的话 - # 这条 pass@K 会被 baseline 成绩推高、与历史臂不可比。 - ws_acc = _mean([1.0 if any((c.get('with_pass') or 0) > 0 and c.get('parseable') for c in r['_cands']) - else 0.0 - for r in chunk if r['_cands']]) - # ⭐ SEAM 同口径的 withskill 准确率：**全部**候选上的 mean(correct)，不看格式、不条件化。 - # = ray_trainer.py:1529 float(np.mean(reward_extra_infos_dict["acc"])) - # = step_summary 的 withskill_pass = swanlab 的 train/with_skill_accuracy。 - # 只有 seam 模式会给 unparseable 候选跑 executor，其余臂这个值等于 candidate_withskill_pass。 - pass_all = [c['with_pass'] for c in all_cands if c['with_pass'] is not None] - # 训练侧 baseline（仅 seam 模式的 chunk 0 有），按候选展开与 SEAM 的 np.mean(baseline_acc) 同口径。 - b_all = [r['_train_baseline_pass'] for r in chunk if r.get('_train_baseline_pass') is not None - for _c in r['_cands']] - base_pass = _mean(b_all) if b_all else None - return { - 'record_type': 'summary', 'chunk': ci, 'n': len(chunk), - 'n_generated': len(all_cands), 'n_candidates_parseable': len(cands), - 'parse_rate': (len(cands) / len(all_cands)) if all_cands else 0.0, - 'n_leaked': sum(1 for c in cands if c['leaked']), - 'leak_rate': (sum(1 for c in cands if c['leaked']) / len(cands)) if cands else 0.0, - 'n_train_samples': n_train, 'n_groups': groups, - 'zero_grad_frac': (zero_grad / groups) if groups else 0.0, - 'reward_mean': _mean(all_rewards), 'reward_std': _pstd(all_rewards), - 'group_reward_std_mean': _mean(group_vars), - 'skill_tokens_mean': _mean([c.get('skillgen_tokens') or 0 for c in cands]), - 'skill_chars_mean': _mean([len(c['skills']) for c in cands]), - 'avg_withskill_pass': ws_acc, - 'candidate_withskill_pass': _mean([c['with_pass'] for c in scored]), - 'withskill_pass_all_cands': _mean(pass_all), - 'n_exec_cands': len(pass_all), - 'baseline_pass_train': base_pass, - 'lift_train': (_mean(pass_all) - base_pass) if base_pass is not None else None, - 'withskill_trunc_frac': (trunc / len(ws_rolls)) if ws_rolls else 0.0, - 'termination_rate_withskill': _mean([1.0 if x['terminated'] else 0.0 for x in ws_rolls]), - } - - -def run_greedy_eval(base_sampler, skill_sampler, eval_records, ci, rounds, - base_dp, skill_dp, args, base_cache): - """Holdout readout: skill-gen (T=args.eval_skill_temperature, args.eval_rollouts rollouts) - -> greedy base solve (T=0). acc = per-problem mean correctness over the rollouts, averaged over - problems (falls back to single greedy when eval_rollouts=1 & temp=0). Adds hard-slice - (baseline_pass==0) rescue rate as a zero-cost secondary readout.""" - # baseline (frozen, cached) - todo = [r for r in eval_records if DiskCache.key_for(r['problem']) not in base_cache] - if todo: - out = _run_samples(base_sampler, [build_direct_prompt(r['problem']) for r in todo], - 1, args.max_tokens, base_dp, temperature=0.0) - rolls = _parse_many([(_first_seq(seqs), r['reference_answer']) - for r, seqs in zip(todo, out)]) - for r, roll in zip(todo, rolls): - base_cache.put(DiskCache.key_for(r['problem']), roll) - for r in eval_records: - br = base_cache.get(DiskCache.key_for(r['problem'])) - r['_baseline_pass'] = 1.0 if br['correct'] else 0.0 - # skill-gen (T=eval_skill_temperature, R rollouts) → with-skill greedy → mean acc over rollouts - R = max(1, args.eval_rollouts) - sg_out = _run_samples(skill_sampler, [_skillgen_prompt(r['problem']) for r in eval_records], - R, args.skill_max_tokens, skill_dp, temperature=args.eval_skill_temperature) - # per problem -> list of R (skill, sresp) - per_skills = [] - for seqs in sg_out: - seqs = list(seqs or []) - row = [] - for j in range(R): - s = seqs[j] if j < len(seqs) else None - if s is None: - row.append(('', '', 'stop')) - else: - sresp = _clean_text(getattr(s, 'decoded', '') or '') - row.append((_extract_skill(sresp) or '', sresp, getattr(s, 'stop_reason', None))) - per_skills.append(row) - # flatten R×N for a single batched greedy executor pass - flat_prompts, flat_idx = [], [] - for pi, (r, row) in enumerate(zip(eval_records, per_skills)): - for j, (sk, sresp, sstop) in enumerate(row): - flat_prompts.append(build_skill_solve_prompt(r['problem'], sk, sresp, - resp_terminated=(sstop != 'length'))) - flat_idx.append((pi, j)) - ws_out = _run_samples(base_sampler, flat_prompts, 1, args.max_tokens, base_dp, temperature=0.0) - judged = _parse_many([(_first_seq(seqs), eval_records[pi]['reference_answer']) - for (pi, _j), seqs in zip(flat_idx, ws_out)]) - roll_by = {idx: roll for idx, roll in zip(flat_idx, judged)} - recs = [] - for pi, (r, row) in enumerate(zip(eval_records, per_skills)): - rolls = [roll_by[(pi, j)] for j in range(len(row))] - corr = [1.0 if x['correct'] else 0.0 for x in rolls] - parses = [1.0 if sk else 0.0 for sk, _sresp, _sstop in row] - terms = [1.0 if x['terminated'] else 0.0 for x in rolls] - acc_mean = sum(corr) / len(corr) if corr else 0.0 - # bugfix（ablate #8）：unparseable skill 的 rollout 实际走了 direct 回退（≈baseline）， - # 主指标里格式崩塌会被 baseline 成绩掩护。strict 通道：unparseable 计 0，格式失败 - # 直接计入代价；主指标口径不变（与历史臂可比），两条曲线分叉即格式崩塌告警。 - acc_strict = (sum(c * p for c, p in zip(corr, parses)) / len(corr)) if corr else 0.0 - recs.append({ - 'record_type': 'eval_problem', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, - 'data_id': r.get('data_id', ''), 'problem': r['problem'], - 'reference_answer': r['reference_answer'], 'baseline_pass': r['_baseline_pass'], - 'n_rollouts': len(row), 'eval_skill_temperature': args.eval_skill_temperature, - 'withskill_acc_mean': acc_mean, # per-problem mean over R rollouts - 'withskill_acc_strict_mean': acc_strict, # unparseable counted wrong - 'withskill_pass_any': 1.0 if any(corr) else 0.0, # pass@R (bonus readout) - 'skill_parseable_mean': sum(parses) / len(parses) if parses else 0.0, - 'withskill_terminated_mean': sum(terms) / len(terms) if terms else 0.0, - # 首个 rollout 的明细留作肉眼抽查 - 'skill': row[0][0], 'skill_parseable': bool(row[0][0]), 'skill_chars': len(row[0][0]), - 'withskill_pred': rolls[0]['pred'], 'withskill_correct': rolls[0]['correct'], - 'withskill_terminated': rolls[0]['terminated'], 'withskill_stop_reason': rolls[0]['stop_reason'], - 'withskill_text': rolls[0]['text'], - }) - n = len(recs) - # acc = 跨题平均的"每题 R 次平均正确率"（mean-over-rollouts） - ws = (sum(x['withskill_acc_mean'] for x in recs) / n) if n else 0.0 - ws_strict = (sum(x['withskill_acc_strict_mean'] for x in recs) / n) if n else 0.0 - pass_any = (sum(x['withskill_pass_any'] for x in recs) / n) if n else 0.0 - base = (sum(x['baseline_pass'] for x in recs) / n) if n else 0.0 - fmt = (sum(x['skill_parseable_mean'] for x in recs) / n) if n else 0.0 - term = (sum(x['withskill_terminated_mean'] for x in recs) / n) if n else 0.0 - # 中文注释：难题子片救活率——baseline_pass==0 子集里 with-skill 的平均正确率（同 mean-over-rollouts 口径）。 - hard = [x for x in recs if not x['baseline_pass']] - hard_rescue_rate = (sum(x['withskill_acc_mean'] for x in hard) / len(hard)) if hard else 0.0 - hard_rescued = sum(x['withskill_acc_mean'] for x in hard) # 期望救活数(分数) - summary = {'record_type': 'eval_summary', 'split': 'eval', 'chunk': ci, 'rounds_done': rounds, - 'n': n, 'acc_mean1': ws, 'baseline_acc_mean1': base, 'lift_mean1': ws - base, - 'acc_strict_mean1': ws_strict, 'lift_strict_mean1': ws_strict - base, - 'acc_pass_any': pass_any, 'n_rollouts': R, 'eval_skill_temperature': args.eval_skill_temperature, - 'format_mean1': fmt, 'term_mean1': term, - 'hard_n': len(hard), 'hard_rescued': hard_rescued, 'hard_rescue_rate': hard_rescue_rate} - metrics = {'core/math/acc/mean@1': ws, 'core/math/baseline_acc/mean@1': base, - 'core/math/lift/mean@1': ws - base, 'core/math/format/mean@1': fmt, - 'core/math/acc_strict/mean@1': ws_strict, 'core/math/lift_strict/mean@1': ws_strict - base, - 'core/math/term/mean@1': term, 'core/math/hard_rescue/mean@1': hard_rescue_rate} - return recs, summary, metrics - - -# =========================================================================== -# Section I — components, args, main -# =========================================================================== -def init_components(args): - r0, r1 = TRAIN_GPUS, TRAIN_GPUS + REF_GPUS - r2, r3 = r1 + SKILL_SAMPLER_GPUS, NUM_GPUS - twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, lazy_collect=False, groups=[ - DeviceGroup(name='train', ranks=list(range(0, r0)), device_type='GPU'), - DeviceGroup(name='ref', ranks=list(range(r0, r1)), device_type='GPU'), - DeviceGroup(name='skill_sampler', ranks=list(range(r1, r2)), device_type='GPU'), - DeviceGroup(name='base_sampler', ranks=list(range(r2, r3)), device_type='GPU')]) - - train_mesh = DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_DP, fsdp_size=TRAIN_FSDP) - # 主权重必须 fp32（对齐 verl actor：fp32 master + bf16 autocast）。twinkle 默认不传 dtype 时 - # transformers 会按 config 加载 bf16 主权重，lr=1e-6 的更新量(~1e-6)远小于 bf16 ulp(~4e-5)， - # optimizer.step 的更新几乎全被舍入吞掉——这是 v2 学不动/与 SEAM 对不上的根因（A/B 实测差 10-20 倍）。 - # resume 旁路（skill_ablate）：skill_init_model_id 指向已保存的 checkpoint 目录时，仅 skill_model - # 从该目录初始化（TransformersModel.load 无全量模型路径）；ref/samplers/template 仍用 MODEL_ID。 - _skill_init_id = getattr(args, 'skill_init_model_id', '') or MODEL_ID - skill_model = TransformersModel(model_id=_skill_init_id, device_mesh=train_mesh, remote_group='train', - torch_dtype='float32', - ddp_config={'find_unused_parameters': False}) - skill_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) - # 方案1：skill 模型开 thinking——让 actor 先在 <think> 里解题+提炼，<skills> 只放通用方法论。 - # <think> 由 _extract_skill 的 rfind('</think>') 砍掉，绝不流给 executor（executor 只吃 <skills>）， - # 因此不构成 SEAM 那种“把 think 喂给 executor”的泄漏。skill_model/ref_model/skill_sampler 三者 - # enable_thinking 必须一致，否则训练轨迹 token 布局与采样对不上。 - # 中文注释：skill_model/ref_model/skill_sampler 三者 enable_thinking 由 --skill-thinking 统一控制 - # （必须一致，否则训练轨迹 token 布局与采样对不上）；base_sampler（executor）走独立开关 - # --executor-thinking（默认 on，E1-E18 全部 on；E19/E20 为 off，见下方 base_sampler 处注释）。 - _think = args.skill_thinking == 'on' - skill_model.set_template(Template, model_id=MODEL_ID, enable_thinking=_think, - max_length=args.max_model_len, truncation_strategy='delete') - skill_model.set_processor(InputProcessor, padding_free=False) - # 客户端同配置副本：训练样本的 prompt 段在这里编码，response 段直接拼采样返回的 - # token（见 build_train_feature）。skill_ablate/trainer.py 走自己那一句，两边不互干扰。 - set_encode_template(Template(model_id=MODEL_ID, enable_thinking=_think, - max_length=args.max_model_len, truncation_strategy='delete')) - # loss 统一用 SEAM 对齐的 SEAMBNPOLoss（verl PPO clip + low_var_kl + token-mean）； - # v2/seam 两模式一致，不再随 align-mode 变。 - _loss_cls = 'SEAMBNPOLoss' - skill_model.set_loss(_loss_cls, epsilon=args.grpo_epsilon, beta=args.kl_beta) - # RL 异常 token 监控：除了 ratio/kl/entropy/clip 那一套，GRPOMetric 还会吐 - # train/logp_min、train/logp_frac_lt_10、train/sampler_logp_mae、train/sampler_token_delta - # —— 训练序列一旦混进模型没生成过的 token（编码错位、模板凭空补的 EOS/空思考块）， - # 前两条会直接炸、后两条直接违反断言；而行为层指标（format/acc/reward）完全看不出来。 - # temperature 不传：skill-gen 就是 T=1，与采样端 logprob 取值口径天然一致。 - skill_model.add_metric('GRPOMetric', is_training=True, epsilon=args.grpo_epsilon) - skill_model.set_optimizer('AdamW', lr=args.lr) - # 对齐 SEAM：恒定 lr（无 warmup、无 decay）。SEAM 用 get_constant_schedule_with_warmup( - # num_warmup_steps=0)+warmup_style=constant，全程恒定 1e-6。这里直接不设 scheduler， - # skill_model.step() 对 lr_scheduler is None 有保护（transformers.py:822-824），lr 恒为 args.lr。 - # 之前的 CosineWarmupScheduler(warmup=10, cosine decay→0) 会让训练中后期有效 lr 持续衰减、 - # 更新幅度变小，与 SEAM 不一致，故移除。 - - ref_mesh = DeviceMesh.from_sizes(world_size=REF_GPUS, dp_size=REF_DP, fsdp_size=REF_FSDP) - ref_model = TransformersModel(model_id=MODEL_ID, device_mesh=ref_mesh, remote_group='ref', - ddp_config={'find_unused_parameters': False}) - ref_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) - # 方案1：与 skill_model 保持一致开 thinking（三者 enable_thinking 必须一致）。 - ref_model.set_template(Template, model_id=MODEL_ID, enable_thinking=_think, - max_length=args.max_model_len, truncation_strategy='delete') - ref_model.set_processor(InputProcessor, padding_free=False) - ref_model.set_loss('GRPOLoss', epsilon=args.grpo_epsilon) - - def _sampler(group, world, enable_thinking): - # enable_prefix_caching 必须开：verl 的 vllm_rollout_spmd 把它硬编码成 True - # （vllm_rollout_spmd.py:182），actor rollout 与 grm/executor rollout 共用这个引擎， - # 所以 SEAM 两条采样通路全程带前缀缓存。twinkle 默认是 False，实测差别不只是速度： - # 同一批里 8 个完全相同的 baseline prompt，缓存关掉时逐条重算、结果 128/128 题全同 - # （608=76x8 精确整除），缓存打开后首条算新 KV、后 7 条复用缓存 KV，数值路径不同， - # SEAM 那边就有 10/128 题的 8 次贪心结果不全同。要和 SEAM 同口径就得同样开着。 - s = vLLMSampler(model_id=MODEL_ID, - engine_args={'gpu_memory_utilization': GPU_MEM, - 'max_model_len': args.max_model_len, 'tensor_parallel_size': 1, - 'enable_prefix_caching': True}, - device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), - remote_group=group) - s.set_template(Template, model_id=MODEL_ID, enable_thinking=enable_thinking, max_length=args.max_model_len) - return s - - # 方案1：skill 采样器开 thinking，与 skill_model/ref_model 一致（actor 先想再写 <skills>）。 - skill_sampler = _sampler('skill_sampler', SKILL_SAMPLER_GPUS, enable_thinking=_think) - # executor（base_sampler）的 thinking 单独一个开关，默认 on（E1-E18 全部如此）。 - # ⭐ 为什么要能关（2026-07-31 bcb 探针实测，n=275 同题配对）：BigCodeBench 上 think 的 - # executor 有 34-50% 的 rollout 撞满预算、连代码块都没写出来（截断样本 8-gram 重复率 - # p50=0.835、同一长句重复 92 次 = 字面死循环），把预算从 4096 加到 20000 也只把截断 - # 从 0.496 压到 0.338 且 pass 不涨。关掉 thinking 后截断归零、裸解反而更高 - # （0.378 vs 0.324），rubric 增量从 +0.080 抬到 +0.135（p=1e-4）。 - # 见 bcb/bcb_eval0_{nothink,think12k,think20k}.jsonl。 - _exec_think = getattr(args, 'executor_thinking', 'on') == 'on' - base_sampler = _LockedSampler(_sampler('base_sampler', BASE_SAMPLER_GPUS, - enable_thinking=_exec_think)) - ckpt = CheckpointEngineManager(model=skill_model, sampler=skill_sampler) - return skill_model, ref_model, skill_sampler, base_sampler, ckpt, SKILL_SAMPLER_GPUS, BASE_SAMPLER_GPUS - - -def _build_args(): - p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--dataset', choices=('aops', 'math'), default='aops') - p.add_argument('--n', type=int, default=0, - help='Problems loaded into the draw pool (0=all; keep 0 to match a SEAM run).') - p.add_argument('--exclude-data-ids', default='', - help='Comma-separated jsonl files whose data_id/problem keys are excluded.') - p.add_argument('--seed', type=int, default=42) - p.add_argument('--numeric-only', action=argparse.BooleanOptionalAction, default=True) - p.add_argument('--eval-size', type=int, default=128, help='Fixed holdout size (0 disables).') - p.add_argument('--seam-parquet-dir', type=str, default='', - help='Read SEAM build_aops_dataset.py train.parquet/val.parquet directly, in ' - 'file order (problem<-extra_info.problem, answer<-reward_model.ground_truth). ' - 'val.parquet becomes the eval holdout. Bypasses load/--numeric-only/' - '--eval-size/internal shuffle so the input data matches a SEAM run exactly.') - p.add_argument('--eval-every', type=int, default=5, help='Run holdout eval every N chunks.') - p.add_argument('--eval-rollouts', type=int, default=1, - help='Eval: skill rollouts per holdout problem. SEAM validation uses one greedy rollout.') - p.add_argument('--eval-skill-temperature', type=float, default=0.0, - help='Eval: skill-model sampling temperature. SEAM validation uses greedy T=0.') - p.add_argument('--chunk-size', type=int, default=16) - p.add_argument('--n-skills', type=int, default=8) - p.add_argument('--skill-gen-temperature', type=float, default=1.0) - p.add_argument('--skill-gen-top-p', type=float, default=1.0) - p.add_argument('--skill-gen-top-k', type=int, default=-1) - p.add_argument('--max-model-len', type=int, default=16384) - p.add_argument('--max-tokens', type=int, default=8192) - # 方案1：开 thinking 后 skill 模型要先写 <think> 分析再写 <skills>，4096 装不下 think+完整 skills - # 会截断成空块（_extract_skill 找不到 </skills> 返回 None）。提到 8192 给两段都留足空间。 - p.add_argument('--skill-max-tokens', type=int, default=8192) - # 中文注释：文体消融开关——主链路与 buffer B regen 同文体（分布一致才可联合训练）。 - p.add_argument('--skill-style', choices=('narrative', 'toy', 'pitfall', 'freeform'), default='narrative', - help='skill文体: narrative=现版叙述式; toy=异数字玩具题示范; pitfall=预判纠错; ' - 'freeform=招式菜单/模型按题自选形态。主链路与 regen 同文体。') - p.add_argument('--skill-thinking', choices=('on', 'off'), default='on', - help='skill_model/ref_model/skill_sampler 三者的 enable_thinking（必须一致）') - p.add_argument('--executor-thinking', choices=('on', 'off'), default='on', - help='executor(base_sampler) 的 enable_thinking。off 用于 BigCodeBench 这类' - '"解答短、难点在选 API 而非多步推理"的任务：think 下 34-50% 的 rollout ' - '陷入字面死循环撞满预算，关掉后截断归零且裸解更高（见 build 处注释）。') - p.add_argument('--align-mode', choices=('v2', 'seam'), default='v2', - help="SEAM-alignment toggle for PROMPT/SKILL FORMAT only. " - "'v2'=clean single-user executor prompt + <skills> skill-gen. " - "'seam'=nested single-user executor prompt + EXPERIENCE_PROMPT/<memory_item> skill-gen. " - 'Reward (lpem numeric-only) and loss (BNPOLoss token-mean) are SEAM-style in BOTH modes.') - p.add_argument('--len-budget', type=int, default=1200, - help='Skill length budget (chars). ONLY used in distillation: drop regen skills ' - 'longer than this, pick the buffer-A seed / buffer-B survivor closest to it. ' - 'Does NOT affect GRPO reward or eval (reward = parseable AND correct).') - # --- buffer / distillation --- - p.add_argument('--distill-trigger', type=int, default=300, - help='Start draining buffer A into distillation once it reaches this many entries.') - p.add_argument('--distill-batch', type=int, default=64, - help='Entries distilled per iteration while buffer A is over --distill-trigger ' - '(incremental drain: bounds per-step latency instead of one big stall).') - p.add_argument('--sft-trigger', type=int, default=100, - help='Run one SFT pass + eval when buffer B reaches this many validated entries. ' - 'Kept low: the distill funnel (has-FAIL × valid-regen × pass@k) yields only ' - '~10-15%% of buffer A, so a high threshold would rarely fire the SFT loop.') - # Plan B validation: diversity lives in the SKILL side, the executor stays at the - # deployment (greedy) decoding口径. For each buffer-A problem we regenerate K distinct - # candidate skills (high temperature), run each through ONE greedy (T=0) executor solve, - # and accept the problem iff >= M distinct skills reach a terminated-correct solve. This - # validates "the problem admits several skills that work under greedy decoding" (matches - # eval口径) rather than "one skill passes m/k times under a high-temperature executor". - p.add_argument('--passatk-k', type=int, default=8, - help='Plan B: number of DISTINCT candidate skills regenerated per problem ' - '(skill-side diversity; executor stays greedy).') - p.add_argument('--passatk-skill-temp', type=float, default=1.0, - help='Skill-model temperature when regenerating the K candidate skills ' - '(needs >0 for diversity across candidates).') - p.add_argument('--passatk-skill-top-p', type=float, default=1.0, - help='Skill-model top-p when regenerating the K candidate skills.') - p.add_argument('--passatk-m', type=int, default=2, - help='Plan B: min number of DISTINCT candidate skills that must reach a ' - 'terminated-correct GREEDY solve to accept the problem into buffer B. ' - 'Lower than pass@k-over-one-skill (default 2): requiring m distinct ' - 'greedy-effective skills is already a strong, low-noise bar.') - p.add_argument('--distill-retries', type=int, default=1, - help='Extra regeneration rounds in distillation for [FAIL] entries that have ' - 'not yet reached m distinct greedy-passing skills (0=single pass). Each ' - 'extra round regenerates more skills (deduped) to rescue more into buffer B.') - p.add_argument('--sft-weight', type=float, default=0.5, - help='Advantage magnitude for SFT distillation samples (-w*logp + beta*KL).') - p.add_argument('--rubric-workers', type=int, default=16) - # --- GRPO --- - p.add_argument('--sft-batch-size', type=int, default=8) - p.add_argument('--ppo-mini-batch-size', type=int, default=0) - p.add_argument('--grpo-epsilon', type=float, default=0.2) - p.add_argument('--adv-clip', type=float, default=0.0, - help='clip group-relative advantage to [-adv_clip, adv_clip]; ' - '0 = no clipping (matches SEAM/verl GRPO which does not clip advantages)') - # 与 skill_ablate/main.py 的默认值必须一致（两个入口默认值分叉 = 静默不可比）。 - # 2026-07-29 拍板 0.001 -> 0.01：对抗侵蚀 executor 收束能力的自发漂移。 - p.add_argument('--kl-beta', type=float, default=0.01) - p.add_argument('--lr', type=float, default=6e-6) - p.add_argument('--max-train-rounds', type=int, default=1500) - p.add_argument('--save-rounds', type=int, default=200) - p.add_argument('--output-dir', default='./output/skill_v2') - p.add_argument('--cache-dir', default='') - p.add_argument('--no-cache', action='store_true') - p.add_argument('--swanlab-project', default='twinkle') - p.add_argument('--swanlab-exp', default='') - args = p.parse_args() - if args.sft_batch_size % TRAIN_DP != 0: - raise ValueError(f'--sft-batch-size ({args.sft_batch_size}) must be a multiple of train dp ({TRAIN_DP})') - if args.chunk_size < 1: - raise ValueError('--chunk-size must be >= 1') - return args - - -def _write(handle, row): - handle.write(json.dumps(row, ensure_ascii=False) + '\n') - - -def _swan_metrics(summary, log): - # Lean metric set: each carries independent information. Dropped as redundant — - # 中文注释：删除冗余项（换算重复）：n_groups(≈chunk_size)、reward_std(池化,组内方差已够)、 - # skill_tokens_mean(与chars重复)、leak/n(=rate×n)、candidate_withskill(与问题级重复)、 - # term/withskill(=1-trunc)、train/n_steps(恒为1)。 - d = { - 'signal/zero_grad_frac': summary['zero_grad_frac'], - 'signal/reward_mean': summary['reward_mean'], - 'signal/group_reward_std_mean': summary['group_reward_std_mean'], - 'signal/n_train_samples': summary['n_train_samples'], - 'skill/parse_rate': summary['parse_rate'], 'skill/chars_mean': summary['skill_chars_mean'], - 'leak/rate': summary['leak_rate'], - } - if summary['n_groups'] > 0: - d.update({'acc/withskill_pass': summary['avg_withskill_pass'], - 'term/withskill_trunc_frac': summary['withskill_trunc_frac']}) - # ⭐ 与 SEAM 同名同口径的三条（ray_trainer.py:1569-1583），专为 swanlab 叠图对齐而发： - # train/with_skill_accuracy = 全部候选的 mean(correct)；acc/reward_mean = mean(correct∧format)； - # skill/format_rate = SEAM 的 format_mean。 - # 注：旧的 acc/withskill_pass 是**题级 pass@K**，与 SEAM 同名指标不同口径（它接近 1 - # 且会随训练小幅下行）—— 之前把这两条叠在一张图上看"趋势相反"就是这个原因。 - d['train/with_skill_accuracy'] = summary['withskill_pass_all_cands'] - d['acc/reward_mean'] = summary['reward_mean'] - d['skill/format_rate'] = summary['parse_rate'] - if summary.get('baseline_pass_train') is not None: - # SEAM 只在 step1 跑训练侧 baseline，所以这两条也只在 chunk 0 有值（与 SEAM 同步）。 - d['acc/baseline_pass'] = summary['baseline_pass_train'] - d['acc/lift'] = summary['lift_train'] - d['train/baseline_accuracy'] = summary['baseline_pass_train'] - d['train/lift'] = summary['lift_train'] - if log: - d['train/n_grpo'] = log['n_grpo'] - d['train/n_sft'] = log['n_sft'] - for k, v in (log.get('metric') or {}).items(): - if not _is_num(v): - continue - if k.startswith('learning rate'): - if 'group 1' in k: - d['train/lr'] = float(v) - elif k.startswith('train/'): - # GRPOMetric 等已经自带 train/ 前缀，不能再套一层。 - d[k.replace(' ', '_')] = float(v) - else: - d[f'train/{k.replace(" ", "_")}'] = float(v) - return d - - -def main(): - args = _build_args() - global _ALIGN_MODE, _SKILL_STYLE - _ALIGN_MODE = args.align_mode # 'v2' | 'seam' - _SKILL_STYLE = args.skill_style # 'narrative' | 'toy' | 'pitfall' | 'freeform' - records, eval_records = _load_records(args) - if len(records) < args.chunk_size: - raise ValueError(f'--chunk-size ({args.chunk_size}) exceeds loaded ({len(records)}); raise --n') - - os.makedirs(args.output_dir, exist_ok=True) - gen_path = os.path.join(args.output_dir, 'gen_records.jsonl') - eval_path = os.path.join(args.output_dir, 'eval_records.jsonl') - sft_path = os.path.join(args.output_dir, 'skill_dataset.jsonl') - train_log_path = os.path.join(args.output_dir, 'train_log.jsonl') - buffer_a_path = os.path.join(args.output_dir, 'buffer_a.jsonl') - distill_path = os.path.join(args.output_dir, 'distill_records.jsonl') - - use_swan = swanlab is not None and os.environ.get('SWANLAB_MODE') != 'disabled' - if use_swan: - swanlab.init(project=args.swanlab_project, experiment_name=(args.swanlab_exp or None), - config={'model': MODEL_ID, 'dataset': args.dataset, 'n': len(records), - 'eval_n': len(eval_records), 'n_skills': args.n_skills, - 'len_budget': args.len_budget, 'distill_trigger': args.distill_trigger, - 'sft_trigger': args.sft_trigger, 'passatk_k': args.passatk_k, - 'passatk_m': args.passatk_m, 'passatk_skill_temp': args.passatk_skill_temp, - 'sft_weight': args.sft_weight, 'lr': args.lr, 'align_mode': args.align_mode}) - - skill_model, ref_model, skill_sampler, base_sampler, ckpt, skill_dp, base_dp = init_components(args) - checker = build_rubric_checker() - if checker is None: - sys.stderr.write('[v2] no LLM backup env -> buffer B distillation DISABLED (GRPO only)\n') - - cache_dir = args.cache_dir or os.path.join(args.output_dir, 'cache') - os.makedirs(cache_dir, exist_ok=True) - # 每次启动强制重算 eval baseline：旧缓存可能来自不同环境/代码版本（torch/vllm/dtype 均影响 T=0 输出）， - # 跨 run 复用会造成 with-skill（现算）vs baseline（陈旧）不可比，lift 虚高/虚低（已实锤过一次）。 - _base_cache_path = os.path.join(cache_dir, 'eval_baseline.jsonl') - if os.path.exists(_base_cache_path): - os.remove(_base_cache_path) - logger.info('stale eval_baseline cache removed (recomputed this run)') - eval_base_cache = DiskCache(_base_cache_path, not args.no_cache) - - cfg = {'record_type': 'config', 'model': MODEL_ID, 'dataset': args.dataset, - 'n': len(records), 'eval_n': len(eval_records), 'seed': args.seed, - 'n_skills': args.n_skills, 'len_budget': args.len_budget, - 'distill_trigger': args.distill_trigger, 'sft_trigger': args.sft_trigger, - 'passatk_k': args.passatk_k, 'passatk_m': args.passatk_m, - 'passatk_skill_temp': args.passatk_skill_temp, 'passatk_skill_top_p': args.passatk_skill_top_p, - 'sft_weight': args.sft_weight, - 'grpo_epsilon': args.grpo_epsilon, 'kl_beta': args.kl_beta, 'lr': args.lr, - 'align_mode': args.align_mode, - 'rubric_check': bool(checker), 'max_train_rounds': args.max_train_rounds, - 'seam_parquet_dir': (getattr(args, 'seam_parquet_dir', '') or ''), - 'started': int(time.time())} - - hist_a: List[Dict[str, Any]] = [] # buffer A accumulator (in-memory + jsonl) - sft_queue: List[Dict[str, Any]] = [] # buffer B: validated SFT records awaiting an SFT pass - rounds = 0 # GRPO rounds only (gates --max-train-rounds + save cadence) - sft_rounds = 0 # SFT passes (separate: must NOT eat the GRPO round budget) - pool = ProblemPool(records, args.seed) - - # Background rubric pre-diagnosis (做法 B): the moment a failure trajectory lands in - # buffer A, fire its teacher-rubric call on a daemon thread pool. The API round-trip - # then overlaps with GRPO GPU work, so by the time --distill-trigger fires the - # diagnoses are usually already cached on each entry ('_rubric_diag'); distill_buffer - # only pays for the stragglers. Entries are dicts held by reference, so the worker - # writes the result straight onto the entry. - # 中文注释：失败轨迹一进 buffer A 就后台异步跑 rubric，API 等待藏进 GPU 训练时间； - # 到蒸馏时诊断多已缓存在条目上，distill_buffer 只补漏。 - prediag_pool = (ThreadPoolExecutor(max_workers=max(1, args.rubric_workers), - thread_name_prefix='rubric-prediag') - if checker else None) - - def _prediagnose(entry: Dict[str, Any]): - entry['_rubric_diag'] = _diagnose_entry(checker, entry) or '' - - with open(gen_path, 'w', encoding='utf-8') as gen_f, \ - open(eval_path, 'w', encoding='utf-8') as eval_f, \ - open(sft_path, 'w', encoding='utf-8') as sft_f, \ - open(train_log_path, 'w', encoding='utf-8') as tlog, \ - open(distill_path, 'w', encoding='utf-8') as distill_f, \ - open(buffer_a_path, 'w', encoding='utf-8') as buf_f: - for f in (gen_f, eval_f, sft_f, tlog, distill_f): - _write(f, cfg) - - def _do_eval(gstep): - recs, summary, metrics = run_greedy_eval( - base_sampler, skill_sampler, eval_records, gstep, rounds, base_dp, skill_dp, - args, eval_base_cache) - for rec in recs: - _write(eval_f, rec) - _write(eval_f, summary) - eval_f.flush() - if use_swan: - swanlab.log({f'eval/{k}': v for k, v in metrics.items()}, step=max(gstep, 0)) - sys.stderr.write( - f'[eval] g{gstep}: n={summary["n"]} acc={summary["baseline_acc_mean1"]:.3f}' - f'->{summary["acc_mean1"]:.3f} lift={summary["lift_mean1"]:+.3f} ' - f'hard_rescue={summary["hard_rescue_rate"]:.3f}({summary["hard_rescued"]}/{summary["hard_n"]}) ' - f'fmt={summary["format_mean1"]:.2f} rounds={rounds}\n') - - if eval_records: - _do_eval(-1) - - gstep = 0 - while rounds < args.max_train_rounds: - chunk = pool.draw(args.chunk_size) - full, summary, grpo, buffer_a = process_chunk( - base_sampler, skill_sampler, chunk, gstep, base_dp, skill_dp, args) - - # accumulate buffer A (only when a rubric checker exists to consume it; - # 中文注释：无 checker 时蒸馏永不触发，不累积以免内存无限增长) - if checker: - for e in buffer_a: - _write(buf_f, e) - prediag_pool.submit(_prediagnose, e) # 后台异步预诊断，不阻塞主循环 - buf_f.flush() - hist_a.extend(buffer_a) - - # GRPO train step (only when there is signal) - log = None - if grpo: - log = _train_step(skill_model, ref_model, ckpt, grpo, args) - rounds += 1 - log.update({'record_type': 'train_round', 'round': rounds, 'chunk': gstep, - 'epoch': pool.epoch, 'kind': 'grpo', 'ts': int(time.time())}) - _write(tlog, log) - tlog.flush() - if rounds % args.save_rounds == 0: - skill_model.save(f'skill-v2-{rounds}', output_dir=args.output_dir) - - summary['rounds_done'], summary['epoch'] = rounds, pool.epoch - summary['buffer_a_size'], summary['sft_queue_size'] = len(hist_a), len(sft_queue) - for rec in full: - _write(gen_f, rec) - _write(gen_f, summary) - gen_f.flush() - - sys.stderr.write( - f'[gen] e{pool.epoch} g{gstep}: n={summary["n"]} ' - f'clean={summary["n_candidates_parseable"]} 0grad={summary["zero_grad_frac"]:.2f} ' - f'R={summary["reward_mean"]:.2f}+-{summary["reward_std"]:.2f} ' - f'ws_acc={summary["avg_withskill_pass"]:.2f} chars={summary["skill_chars_mean"]:.0f} ' - f'bufA={len(hist_a)} bufB={len(sft_queue)} rounds={rounds}\n') - if use_swan: - m = _swan_metrics(summary, log) - m['buffer/a_size'] = float(len(hist_a)) - m['buffer/b_size'] = float(len(sft_queue)) - swanlab.log(m, step=gstep) - - # --- distillation: once buffer A fills, drain it INCREMENTALLY in bounded - # batches (--distill-batch) so a large buffer never stalls the loop for tens - # of minutes; each iteration processes one batch, interleaved with GRPO. - # 中文注释：增量分批蒸馏——buffer A 满后每轮只处理 --distill-batch 条，把一次性 - # 几十分钟阻塞摊成每轮几分钟小停顿；两段验证(见 distill_buffer)再砍验证算力。 - if checker and len(hist_a) >= args.distill_trigger: - batch = hist_a[:args.distill_batch] - hist_a = hist_a[args.distill_batch:] - new_sft, distill_recs = distill_buffer(batch, skill_sampler, base_sampler, checker, - skill_dp, base_dp, args) - for rec in distill_recs: # 逐 entry 审计记录：rubric_diag + 候选 skill + 贪心解 + stage - rec['chunk'] = gstep - _write(distill_f, rec) - distill_f.flush() - for rec in new_sft: - _write(sft_f, rec) - sft_f.flush() - sft_queue.extend(new_sft) - - # --- SFT trigger: buffer B full → one SFT pass + eval --- - did_eval = False - if len(sft_queue) >= args.sft_trigger: - sys.stderr.write(f'[sft] triggered at bufB={len(sft_queue)}\n') - sft_samples = [{**s, 'advantage': float(args.sft_weight)} for s in sft_queue] - sft_log = _train_step(skill_model, ref_model, ckpt, sft_samples, args) - sft_rounds += 1 # 中文注释：SFT 用独立计数，不占用 GRPO 的 rounds 配额/save 节奏 - sft_log.update({'record_type': 'train_round', 'round': rounds, 'sft_round': sft_rounds, - 'chunk': gstep, 'epoch': pool.epoch, 'kind': 'sft', 'ts': int(time.time())}) - _write(tlog, sft_log) - tlog.flush() - sft_queue = [] - skill_model.save(f'skill-v2-sft{sft_rounds}', output_dir=args.output_dir) # 大改动后落盘 - if eval_records: # 中文注释：SFT 后立即 eval，测灾难性遗忘/真提升（第 11.3/13.4 节） - _do_eval(gstep) - did_eval = True - - if eval_records and not did_eval and (gstep + 1) % args.eval_every == 0: - _do_eval(gstep) - gstep += 1 - - if prediag_pool is not None: - prediag_pool.shutdown(wait=False, cancel_futures=True) # 丢弃未完成的后台预诊断 - eval_base_cache.close() - skill_model.save('skill-v2-final', output_dir=args.output_dir) - sys.stderr.write(f'[v2] done: {rounds} rounds over {gstep} chunks / {pool.epoch} epochs\n') - - -if __name__ == '__main__': - main() diff --git a/cookbook/exp/skill2lora/train_skill_v2.sh b/cookbook/exp/skill2lora/train_skill_v2.sh deleted file mode 100644 index 2f7e4573a..000000000 --- a/cookbook/exp/skill2lora/train_skill_v2.sh +++ /dev/null @@ -1,69 +0,0 @@ -#!/bin/bash -# train_skill_v2.sh — 简化 GRPO + buffer distill 训练启动脚本 -# 用法: bash cookbook/exp/skill2lora/train_skill_v2.sh -# -# 环境变量: -# LLM_BACKUP_API_KEY - rubric 诊断用的教师 API key（必须，否则 buffer B 蒸馏不可用） -# LLM_BACKUP_BASE_URL - 教师 API base URL -# LLM_BACKUP_MODEL - 教师模型 ID -# GEN_MODEL_ID - 训练 skill 模型 ID（默认 Qwen/Qwen3-4B） -# TRAIN_GPUS / REF_GPUS / SKILL_SAMPLER_GPUS / BASE_SAMPLER_GPUS — GPU 分配 - -set -euo pipefail -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" - -# 缓解显存碎片（reserved-but-unallocated），降低 forward_backward 阶段 OOM 概率 -export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" - -# 默认输出目录 -OUTPUT_DIR="${OUTPUT_DIR:-./output/skill_v2}" - -# 去重/排斥数据（冷启动 SFT 数据避免重叠） -EXCLUDE="${EXCLUDE_DATA_IDS:-}" - -# 与 SEAM 输入数据对齐：直读 SEAM build_aops_dataset.py 产出的 parquet（同题池 + 同 val） -# 置空则回退到 twinkle 自己的 load+shuffle+split。 -SEAM_PARQUET_DIR="${SEAM_PARQUET_DIR:-/root/data/seam}" - -# 提前建目录：tee 需在 python 建目录前就能打开日志文件 -mkdir -p "${OUTPUT_DIR}" - -python3 "${SCRIPT_DIR}/train_skill_v2.py" \ - --dataset aops \ - --numeric-only \ - --eval-size 200 \ - --eval-every 5 \ - --eval-rollouts 1 \ - --eval-skill-temperature 0.0 \ - --chunk-size 16 \ - --n-skills 8 \ - --distill-retries 1 \ - --skill-gen-temperature 1.0 \ - --skill-gen-top-p 1.0 \ - --skill-gen-top-k -1 \ - --max-model-len 16384 \ - --max-tokens 8192 \ - --skill-max-tokens 4096 \ - --len-budget 600 \ - --distill-trigger 150 \ - --distill-batch 64 \ - --sft-trigger 100 \ - --passatk-k 8 \ - --passatk-m 2 \ - --align-mode seam \ - --sft-weight 1.0 \ - --rubric-workers 16 \ - --sft-batch-size 4 \ - --ppo-mini-batch-size 0 \ - --grpo-epsilon 0.2 \ - --adv-clip 0 \ - --kl-beta 0.001 \ - --lr 1e-6 \ - --max-train-rounds 1500 \ - --save-rounds 200 \ - --output-dir "${OUTPUT_DIR}" \ - --swanlab-project twinkle \ - --swanlab-exp "skill_v2_$(date +%Y%m%d_%H%M%S)" \ - ${EXCLUDE:+--exclude-data-ids "${EXCLUDE}"} \ - ${SEAM_PARQUET_DIR:+--seam-parquet-dir "${SEAM_PARQUET_DIR}"} \ - "$@" 2>&1 | tee "${OUTPUT_DIR}/run.log" diff --git a/cookbook/exp/skill2lora/watchdog_e14.sh b/cookbook/exp/skill2lora/watchdog_e14.sh deleted file mode 100644 index b9789958f..000000000 --- a/cookbook/exp/skill2lora/watchdog_e14.sh +++ /dev/null @@ -1,40 +0,0 @@ -#!/bin/bash -# One-shot watchdog (2026-07-28): the running ablate12 launcher holds the OLD plan -# (E7 -> E8). E14 (reward SNR ablation) was inserted after E7 in config.py, so when E7 -# finishes (DONE.json) OR the launcher dies (E7 crash), swap to a fresh launcher that -# reads the new plan: completed arms skip via DONE.json, so it starts E14 directly. -set -u -HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -DONE="$HERE/output.ablate12/E7_rl_ab_on_pitfall/DONE.json" -PIDF=/tmp/ablate12.pid -LOG="$HERE/watchdog_e14.log" - -echo "[watchdog $(date +%H:%M:%S)] waiting for E7 DONE.json or launcher exit" >> "$LOG" -while true; do - [ -f "$DONE" ] && { echo "[watchdog $(date +%H:%M:%S)] E7 DONE.json found" >> "$LOG"; break; } - OLD=$(cat "$PIDF" 2>/dev/null || echo "") - if [ -n "$OLD" ] && ! kill -0 "$OLD" 2>/dev/null; then - echo "[watchdog $(date +%H:%M:%S)] launcher $OLD died without E7 DONE (crash?); restarting anyway" >> "$LOG" - break - fi - sleep 60 -done - -sleep 10 -OLD=$(cat "$PIDF" 2>/dev/null || echo "") -[ -n "$OLD" ] && kill "$OLD" 2>/dev/null && echo "[watchdog] killed old launcher $OLD" >> "$LOG" -# kill any experiment python the old launcher may have just started (E8 race window) -pkill -f "skill_ablate.main" 2>/dev/null && echo "[watchdog] killed stray skill_ablate.main" >> "$LOG" - -# wait for the 8 GPUs to drain (engine teardown), max 15 min -for i in $(seq 1 90); do - USED=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | awk '{s+=$1} END {print s}') - [ "${USED:-1}" -lt 1000 ] && break - sleep 10 -done -echo "[watchdog $(date +%H:%M:%S)] GPUs drained (used=${USED:-?}MiB); relaunching" >> "$LOG" - -cd "$HERE" -nohup bash run_ablate12.sh > run_ablate12.nohup.log 2>&1 & -echo $! > "$PIDF" -echo "[watchdog $(date +%H:%M:%S)] new launcher pid=$(cat $PIDF) (plan includes E14 after E7)" >> "$LOG" diff --git a/cookbook/human/e23_bcb.py b/cookbook/human/e23_bcb.py deleted file mode 100644 index 4ff11acaa..000000000 --- a/cookbook/human/e23_bcb.py +++ /dev/null @@ -1,278 +0,0 @@ -"""BigCodeBench 环境层：数据加载、沙箱单测判分、模型输出解析。 - -与训练完全解耦 —— 这里只回答「一段模型文本能不能跑过官方单测」，不认识 skill / rubric / GRPO。 -""" -import ast as _ast -import importlib.util -import json -import os -import re -import shutil -import subprocess -import sys -import tempfile -from concurrent.futures import ThreadPoolExecutor -from typing import Any, Dict, List, Optional, Tuple - -from twinkle import get_logger -from twinkle.dataset import Dataset, DatasetMeta - -logger = get_logger() -_HERE = os.path.dirname(os.path.abspath(__file__)) - -# ModelScope 数据集 id 或本地 parquet 路径都行：DatasetMeta 按 os.path.exists 自行分流， -# 走本地文件时 subset / split 会被忽略。 -BCB_DATASET = os.environ.get('BCB_DATASET', 'ms://bigcode/bigcodebench') -BCB_SUBSET = os.environ.get('BCB_SUBSET', 'default') -BCB_SPLIT = os.environ.get('BCB_SPLIT', 'v0.1.0_hf') -TEST_WORKERS = int(os.environ.get('TEST_WORKERS', 24)) # 跑单测的线程池（每线程一个子进程） -TEST_TIMEOUT = int(os.environ.get('TEST_TIMEOUT', 60)) # 单题单测墙钟上限（秒） - -# 需要外网 / GUI / 子进程的库：沙箱里会挂或超时，判分噪声与 skill 无关 -> 整题排除。 -EXCLUDE_LIBS = {'requests', 'urllib', 'http', 'smtplib', 'socket', 'ssl', 'ftplib', - 'mechanize', 'wikipedia', 'turtle', 'tkinter', 'subprocess', 'sendgrid', - 'python_http_client', 'django', 'flask', 'flask_login', 'flask_mail', - 'flask_restful', 'flask_wtf', 'wtforms', 'multiprocessing'} -LIB_ALIAS = {'cv2': 'cv2', 'PIL': 'PIL', 'bs4': 'bs4', 'yaml': 'yaml', 'dateutil': 'dateutil', - 'Crypto': 'Crypto', 'docx': 'docx', 'pytz': 'pytz', 'psutil': 'psutil', - 'texttable': 'texttable', 'wordcloud': 'wordcloud', 'skimage': 'skimage', - 'PyPDF2': 'PyPDF2', 'sklearn': 'sklearn'} - -_FENCE_RE = re.compile(r'```(?:python|py)?\s*\n(.*?)```', re.S) -_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') - - -# ========== 输出解析 ========== -def after_think(text: str) -> str: - i = (text or '').rfind('</think>') - return text[i + len('</think>'):] if i >= 0 else (text or '') - - -def clean_text(decoded: Optional[str]) -> str: - """只剔 <|...|> 这类特殊 token 的**字面量**。<think> 必须保留 —— E23 要把它递给 executor。""" - return _SPECIAL_TOKEN_RE.sub('', decoded or '').rstrip() - - -def extract_code(text: str) -> str: - """取最后一个能通过 ast.parse 的代码块；没有围栏就退化为整段（切 think 之后）。""" - body = after_think(text or '') - blocks = _FENCE_RE.findall(body) - for b in reversed(blocks): - try: - _ast.parse(b) - return b - except SyntaxError: - continue - if blocks: - return blocks[-1] - try: - _ast.parse(body) - return body - except SyntaxError: - return '' - - -def extract_skill(text: str) -> str: - """抽 <skills> 块；任何畸形（未闭合 / 只写在 think 里 / 空块）一律返回 ''。 - - 只在 </think> **之后**找：写在思考过程里的 <skills> 不算产出。返回 '' 时 executor 走干净 - direct（见 skill_solve_prompt），不会拿到半截内容。 - """ - answer = after_think(text) - s = answer.lower().rfind('<skills>') - if s < 0: - return '' - inner = s + len('<skills>') - e = answer.lower().find('</skills>', inner) - if e < 0: - return '' - return re.sub(r'</?(?:skills|skill|diagnose|pitfall|strategy|think)>', '', - answer[inner:e].strip(), flags=re.IGNORECASE).strip() - - -# ========== 沙箱判分 ========== -_RUNNER = """ -import unittest, sys -loader = unittest.TestLoader() -suite = loader.loadTestsFromTestCase(TestCases) -res = unittest.TextTestRunner(verbosity=0, stream=sys.stderr).run(suite) -print('__BCB__', res.testsRun, len(res.failures), len(res.errors)) -sys.exit(0 if res.wasSuccessful() and res.testsRun > 0 else 1) -""" - - -def _trim_err(err: str, limit: int = 1600) -> str: - """保留失败测试名与异常行，砍掉冗长 traceback 帧 —— 这是喂给 rubric 的客观证据。 - 随机临时目录名换成 <sandbox>，否则同一个失败在两次运行里看起来不一样。""" - err = re.sub(r'/tmp/bcb_[A-Za-z0-9_]+', '<sandbox>', err or '') - lines = [ln for ln in err.splitlines() if ln.strip()] - keep = [ln for ln in lines - if ln.startswith(('FAIL:', 'ERROR:', 'AssertionError', 'Traceback')) - or re.match(r'^\w*(Error|Exception|Warning)\b', ln.strip()) - or ', in ' in ln] - return '\n'.join(keep or lines[-25:])[-limit:] - - -def run_tests(code: str, payload: Dict[str, Any], timeout: int = TEST_TIMEOUT) -> Dict[str, Any]: - """子进程里跑「提交代码 + 官方 test + _RUNNER」。-> {'passed', 'kind', 'error'}。""" - if not code.strip(): - return {'passed': False, 'kind': 'no_code', 'error': 'no parseable code block'} - if payload['entry_point'] not in code: - return {'passed': False, 'kind': 'no_entry', - 'error': f"function {payload['entry_point']} is not defined in the submitted code"} - tmp = tempfile.mkdtemp(prefix='bcb_') - try: - path = os.path.join(tmp, 'run_case.py') - with open(path, 'w', encoding='utf-8') as f: - f.write(code + '\n\n' + payload['test'] + '\n' + _RUNNER) - env = dict(os.environ, MPLBACKEND='Agg', PYTHONHASHSEED='0', OMP_NUM_THREADS='1', - MKL_NUM_THREADS='1', TOKENIZERS_PARALLELISM='false') - env.pop('CUDA_VISIBLE_DEVICES', None) - try: - p = subprocess.run([sys.executable, path], cwd=tmp, env=env, timeout=timeout, - capture_output=True, text=True, errors='replace') - except subprocess.TimeoutExpired: - return {'passed': False, 'kind': 'timeout', - 'error': f'the tests did not finish within {timeout}s'} - n_tests = n_fail = n_err = 0 - for line in (p.stdout or '').splitlines(): - if line.startswith('__BCB__'): - _, a, b, c = line.split() - n_tests, n_fail, n_err = int(a), int(b), int(c) - if p.returncode == 0 and n_tests > 0: - return {'passed': True, 'kind': 'pass', 'error': ''} - kind = 'assertion' if n_fail else ('exception' if n_err else 'import_or_syntax') - return {'passed': False, 'kind': kind, - 'error': _trim_err((p.stderr or '').replace(tmp, '<sandbox>'))} - finally: - shutil.rmtree(tmp, ignore_errors=True) - - -def empty_roll() -> Dict[str, Any]: - return {'correct': False, 'stop_reason': 'empty', 'gen_tokens': 0, 'text': '', 'code': '', - 'kind': 'no_code', 'error': ''} - - -def judge_seqs(pairs: List[Tuple[Any, Dict[str, Any]]]) -> List[Dict[str, Any]]: - """[(采样 sequence 或 None, payload)] -> rolls。所有判分都汇合到这里。 - - 必须批量：单测是子进程（导入 pandas/sklearn 后典型 1-3s），一个 chunk 几百次判分串行会比同 - chunk 的 GPU 时间还长一个量级。同 (task_id, code) 只跑一次 —— T=0 的 executor 经常对同一题 - 产出逐字相同的代码。 - """ - rolls: List[Dict[str, Any]] = [] - keys: List[Optional[Tuple[str, str]]] = [] - jobs: Dict[Tuple[str, str], Dict[str, Any]] = {} - for seq, payload in pairs: - if seq is None: - rolls.append(empty_roll()) - keys.append(None) - continue - text = clean_text(getattr(seq, 'decoded', '') or '') - code = extract_code(text) - key = (payload['task_id'], code) - rolls.append({'correct': False, 'stop_reason': getattr(seq, 'stop_reason', None), - 'gen_tokens': len(getattr(seq, 'tokens', None) or []), - 'text': text, 'code': code, 'kind': None, 'error': ''}) - keys.append(key) - jobs.setdefault(key, payload) - if jobs: - todo = list(jobs) - with ThreadPoolExecutor(max_workers=max(1, min(TEST_WORKERS, len(todo)))) as ex: - verdicts = dict(zip(todo, ex.map(lambda k: run_tests(k[1], jobs[k]), todo))) - for roll, key in zip(rolls, keys): - v = verdicts.get(key) if key is not None else None - if v is not None: - roll['correct'] = bool(v['passed']) - roll['kind'], roll['error'] = v['kind'], v['error'] - return rolls - - -# ========== 数据 ========== -# reference_answer 的字段集：判分需要的一切（code 域不是数值答案）。 -_PAYLOAD_KEYS = ('task_id', 'entry_point', 'test', 'code_prompt', 'doc_struct', - 'canonical_solution') - - -def _importable(lib: str) -> bool: - try: - return importlib.util.find_spec(LIB_ALIAS.get(lib, lib).split('.')[0]) is not None - except Exception: - return False - - -def _row_libs(row: Dict[str, Any]) -> List[str]: - v = row.get('libs') - if isinstance(v, str): # 数据集里存的是 list 的字符串形式 - try: - return list(_ast.literal_eval(v)) - except Exception: - return [] - return list(v or []) - - -def _to_record(batch: Dict[str, List]) -> Dict[str, List]: - """原始 BCB 列 -> {'data_id', 'problem', 'reference_answer'}。 - - Dataset.map 强制 batched=True，所以这里收发的都是列式 batch。 - """ - return {'data_id': list(batch['task_id']), - 'problem': list(batch['instruct_prompt']), - 'reference_answer': [{k: batch[k][i] for k in _PAYLOAD_KEYS} - for i in range(len(batch['task_id']))]} - - -def _broken_tasks(ds: Dataset, output_dir: str) -> set: - """参考解答跑不过自己的单测 = 沙箱/依赖不可判定，不是模型的错（实测约 7.5%）。 - 自检一次后落盘缓存，题数不变则复用。必须在 map 之前调用（要读原始列）。""" - path = os.path.join(output_dir, 'bcb_broken_tasks.json') - if os.path.exists(path): - try: - with open(path, encoding='utf-8') as f: - c = json.load(f) - if int(c.get('n_tasks', -1)) == len(ds): - return set(c['broken']) - except Exception as exc: - logger.warning(f'[data] 读取 {path} 失败（{exc}），重跑自检') - logger.info(f'[data] 沙箱自检：{len(ds)} 道题跑参考解答（一次性，之后走缓存）…') - rows = [ds[i] for i in range(len(ds))] - jobs = [(r['code_prompt'] + (r['canonical_solution'] or ''), {k: r[k] for k in _PAYLOAD_KEYS}) - for r in rows] - with ThreadPoolExecutor(max_workers=max(1, min(TEST_WORKERS, len(jobs)))) as ex: - vers = list(ex.map(lambda p: run_tests(p[0], p[1]), jobs)) - broken = {r['task_id'] for r, v in zip(rows, vers) if not v['passed']} - with open(path, 'w', encoding='utf-8') as f: - json.dump({'n_tasks': len(ds), 'broken': sorted(broken)}, f, indent=1) - return broken - - -def load_records(seed: int, eval_size: int, - output_dir: str) -> Tuple[Dataset, List[Dict[str, Any]]]: - """-> (train_dataset, eval_records)，每条记录是 {'data_id', 'problem', 'reference_answer'}。 - - 训练侧返回 Dataset 交给调用方喂 DataLoader；holdout 是固定的一小批、每轮 eval 整体遍历， - 没有分批的意义，直接物化成 list。 - - BigCodeBench 没有 difficulty 字段，所以不做分层：过滤完按 seed 洗牌，前 eval_size 道作 - holdout。题池很小（1140 -> 剔除后约 900），重复抽到的题 rubric 全部缓存命中，成本只在 - GPU rollout。 - """ - ds = Dataset(DatasetMeta(BCB_DATASET, subset_name=BCB_SUBSET, split=BCB_SPLIT)) - n_raw = len(ds) - ds.filter(lambda r: not (set(_row_libs(r)) & EXCLUDE_LIBS)) - n_kept_libs = len(ds) - ds.filter(lambda r: all(_importable(x) for x in _row_libs(r))) - logger.info(f'[data] BigCodeBench: 全集 {n_raw}，剔除需外网/GUI/子进程 {n_raw - n_kept_libs}、' - f'依赖缺失 {n_kept_libs - len(ds)} -> {len(ds)}') - - broken = _broken_tasks(ds, output_dir) - if broken: - ds.filter(lambda r: r['task_id'] not in broken) - logger.info(f'[data] 剔除参考解答自己跑不过单测的题 {len(broken)} 道 -> 可用 {len(ds)}') - ds.map(_to_record, remove_columns=ds.dataset.column_names) - - shuffled = ds.dataset.shuffle(seed=seed) - n_eval = min(eval_size, len(shuffled)) if eval_size > 0 else 0 - eval_records = list(shuffled.select(range(n_eval))) - train_dataset = Dataset(DatasetMeta(data=shuffled.select(range(n_eval, len(shuffled))))) - return train_dataset, eval_records diff --git a/cookbook/human/e23_prompts.py b/cookbook/human/e23_prompts.py deleted file mode 100644 index 5e73f2f40..000000000 --- a/cookbook/human/e23_prompts.py +++ /dev/null @@ -1,225 +0,0 @@ -"""E23 的全部 prompt 文本与拼装函数：executor / 教师 judge / skill-gen 三处。""" -# flake8: noqa: E501 -# prompt 正文按「一段一行」书写，折行会改变真正发给模型的文本，故整文件豁免行长检查。 -import json -from typing import Any, Dict - -# =========================================================================== -# executor -# =========================================================================== -# BigCodeBench 官方 instruct 模式的硬性交付要求。 -EXEC_SYSTEM = """\ -You are an expert Python engineer. You will be given a task description that ends with the exact \ -import lines and function signature your solution must start with. - -Deliver exactly one fenced Python code block and nothing else after it: -- Reproduce the given imports and the given function signature verbatim, including parameter \ -names, order and default values. -- Add any further imports you need inside the same block; the block must run standalone. -- Return exactly the object type the task says to output. If it says the function should output \ -a tuple, return a tuple in that order; if it names a matplotlib Axes, return the Axes object \ -itself, not the Figure and not None. -- Implement the described behaviour for the general case, including the empty / single-element / \ -missing-column edge cases and any exception the description says to raise. -- Do not call the function, do not print demonstrations, do not add tests, do not use \ -`if __name__ == '__main__'`, and do not read from stdin. -- Do not include explanations outside the code block.""" - -# ⭐ E23 的定义性特征：executor 看到 actor 的**完整产出**（<think> + <skills>），不是只看抽出来的 -# <skills> 块。第一句必须说明「接下来这坨是什么」，否则 executor 会把 <think> 当成任务描述的一 -# 部分。第二句与「只给 skills」的对照臂逐字相同，使两臂只差「能否看到推理过程」这一个变量。 -_WRAPPER_WITH_THINK = ( - 'Guidance model transcript:\nFor this task, a separate problem-solving guidance model was ' - 'asked to analyse it and produce advisory skills. Its full output follows, including its ' - 'private reasoning:\n{hint}\n' - 'Prefer using its techniques when they fit, but if you have a clearly better ' - 'implementation, you may diverge. Be concise and accurate.\n') - - -def direct_prompt(problem: str) -> Dict[str, Any]: - return {'messages': [{'role': 'system', 'content': EXEC_SYSTEM}, - {'role': 'user', 'content': problem}]} - - -def skill_solve_prompt(problem: str, skill: str, raw_response: str = '') -> Dict[str, Any]: - """题面 + actor 全文（含 <think>）。 - - skill 为空（actor 只写了 <think>，或写到上限没闭合）-> 干净 direct，而不是把一坨无结论的思考 - 过程当指导塞进去。故意**不补 <|im_end|>**：那会让 advisory 落在非法 ChatML 位置并改变紧邻 - token 的 BPE 切分，只为「与 SEAM 逐 token 对齐」才值得，这里的代价是多一个变量。 - """ - skill = (skill or '').strip() - if not skill: - return direct_prompt(problem) - hint = (raw_response or '').strip() or skill - return {'messages': [{'role': 'system', 'content': EXEC_SYSTEM}, - {'role': 'user', - 'content': problem + '\n\n' + _WRAPPER_WITH_THINK.format(hint=hint)}]} - - -# =========================================================================== -# 教师 judge（rubric 诊断） -# =========================================================================== -DIAG_SYSTEM = """\ -You are a code-failure classifier. You are given a Python task description (with the required signature), a list of FAILURE CLASSES, one attempted implementation, and the REAL error that attempt produced when the task's unit tests were run. - -Your job is NOT to review the code. It is to (a) name the SINGLE decisive failure class, (b) quote the evidence for it out of the test error, and (c) state the prior knowledge that would have PREVENTED it. - -Output STRICT JSON (no prose outside it), either - -{"addressable": true, "class": "<CLASS CODE>", "evidence": "<verbatim fragment of the test error that proves this class>", "required_value": "<the exact value, name or behaviour the tests demand -- a title, a label, a dictionary key, a numeric bound, an exception that must be raised -- or null if the failure is not about matching a demanded value>", "required_value_source": "<verbatim fragment of the TASK section that states that value or demands that behaviour, or null if the task never states it>", "reason": "<one short sentence: which assumption was wrong, at which step>", "prior": "<one sentence of transferable knowledge>", "secondary": ["<CLASS CODE>"], "independent_causes": 1} - -or, when you cannot ground a decisive class in the test error: - -{"addressable": false, "why": "<one short sentence>"} - -Rules: -- "evidence" MUST be copied verbatim out of the test error you were given. If the error does not let you single out ONE class, output addressable=false instead of guessing. Never invent evidence. -- The error tells you WHICH assertion fired first, not WHAT caused it. A failing shape, type or value assertion is normally the last link of the chain, not the class. Ask "which wrong belief produced this?" and classify THAT. Classify by the assertion itself only when nothing upstream is wrong: when the computation is right and merely the kind of object handed back is not what the caller reads. -- "class" is the one class the evidence implicates. List other classes that are also off in "secondary"; leave it empty if there are none. -- "independent_causes" counts mutually independent root causes, not symptoms of one. Three or more means no single short warning could have saved this attempt. -- "prior" is the hard part. It MUST STAY TRUE AND USEFUL IF THIS TASK IS DELETED: a fact about a library, an API default, a format directive, or a testing convention. It must NOT contain any identifier, literal, column name, file name or number taken from this task, and must NOT describe the steps of this task's solution. Write "pandas writes JSON Lines rather than one JSON document when the lines flag is set", NOT "pass lines=False here". -- "prior" must say WHY the wrong result arose -- the rule the attempt had backwards -- and NOT what the correct result should look like. "The result must carry one row per input key" merely restates the assertion that failed and teaches nothing, because the reader already has the task description; "a merge defaults to an inner join and silently drops keys missing on either side" is the belief that was actually absent. -- For class TESTCONTRACT and class EXCEPTION, "required_value" MUST NOT be null. Those two classes are BY DEFINITION about failing to match something the caller demands -- a title, a label, a key, an exception -- so name that thing. If you find yourself unable to name it, you have the wrong class. -- "required_value" and "required_value_source" are a CITATION, not an opinion. Whenever the failure comes down to matching something the tests demand, put that demanded thing in "required_value", then go back to the TASK section and copy out the fragment that states it into "required_value_source". You may only fill "required_value_source" with text you can actually see in the TASK section -- it is checked against it verbatim. If the task never states it, write null, and the diagnosis will be discarded, which is the correct outcome: an engineer reading only the task could not have known it either. -- That check is where a plausible-looking diagnosis does the most damage. "Tests assert exact string equality on titles" is a true and transferable sentence, and it is still worthless when the title itself appears nowhere in the task -- the engineer learns that the string matters but not what it is. Do not let a good-sounding prior talk you out of the citation. -- Set addressable=false rather than writing a vacuous prior such as "implement the description carefully". A sentence that merely says validation must be written, or that a parameter must satisfy the tests, is vacuous however factual it sounds; a usable prior names a library, function, format or convention that the reader could look up. -- Output only the JSON object.""" - -DIAG_USER = """\ -## Task -{query} - -## Failure classes -{rubric} - -## Attempted implementation and its test error -{segment} - -Now output the classification JSON object.""" - -# ⭐ 判据按「什么先验知识能预测这个失败」切分，而不是按「代码哪里错了」。 -# 依据是 E23.t1 的 480 组实测：旧表第 3 条只管**签名合法性**（调用能否被接受），于是所有「调用被 -# 接受、但行为/默认值不符合假设」的失败（to_json(lines=True)、mkdir 不带 parents、glob('*') 含 -# 目录、strftime('%Z') 多后缀）全掉进兜底项「核心计算错」——它在 76% 的题上 FAIL，且 leak 率 -# 0.106 是全体 0.021 的 5 倍：它救得回题，靠的正是让 skill 把解法写出来。BEHAVIOUR 就是补这个洞， -# 占旧兜底 FAIL 的 53%。每条都按**报错里的可观测特征**定义，这才能既判得准又分得开。 -# -# ⚠️ 「按可观测特征定义」的代价是教师容易**照最先响的那条断言分类**，于是 SHAPE 会接手本属 -# BEHAVIOUR 的题（BCB/441：einsum 输入下标 ikl 应为 jkl，症状是 shape 断言先挂 -> 判 SHAPE -> -# prior 只讲输出下标决定形状 -> 8 个候选一起只改输出下标，形状对了数值仍错，整组 reward 0）。 -# SHAPE 的定义因此显式排除「算错导致形状/数值不对」，DIAG_SYSTEM 里也有一条反症状规则兜着。 -# (code, 给 skill-gen 的短标签, 给教师判定用的完整定义) -FAILURE_CLASSES = [ - ('BEHAVIOUR', 'a library call behaves differently from what was assumed', - 'The call is accepted, but the function\'s real behaviour, default value or precondition ' - 'differs from what the code assumed: a flag that changes the output format, a default that ' - 'does not do what its name suggests, a required preparation step, a half-open range.'), - ('SIGNATURE', 'the call itself is not accepted', - 'The function, attribute or module does not exist, or it does not take the argument names, ' - 'positions or count that were passed.'), - ('SHAPE', 'the object handed back is not the one the task asks for', - 'The KIND of object is wrong irrespective of the values inside it: wrong container or element ' - 'type, wrong nesting, an unconsumed lazy object where a value was expected, a figure handed ' - 'back where the caller reads an axes. NOT for a result whose shape or values came out wrong ' - 'because the computation was wrong -- that belongs to whichever class names the wrong ' - 'assumption in the computation, usually BEHAVIOUR.'), - ('TESTCONTRACT', 'what the caller inspects was never set, or a required side effect never ran', - 'The returned object exists but does not carry what the tests read off it, or a side effect ' - 'the task requires was never performed: labels and titles left unset, a resource not cleaned ' - 'up, a call the task says to make never made.'), - ('DETERMINISM', 'the result is not reproducible or not exactly comparable', - 'Unseeded randomness, reliance on iteration order, missing or wrong sorting, rounding or ' - 'precision other than stated.'), - ('NORMALISATION', 'a text or boundary convention is wrong', - 'Case sensitivity, surrounding whitespace, regex anchoring, separator handling, inclusive ' - 'versus exclusive bounds, off-by-one.'), - ('DEGENERATE', 'a degenerate input crashes instead of being handled', - 'Empty, single-element, all-equal or missing-key / missing-column input.'), - ('EXCEPTION', 'the exception contract is not met', - 'The exception the task specifies is not raised, is raised as a different type, or an ' - 'unrelated exception escapes.'), -] - - -def render_classes() -> str: - """完整定义只给教师；skill-gen 那侧只看短标签，避免它照着举例去写不相干的失败模式。""" - return '\n'.join(f'- {code}: {desc}' for code, _short, desc in FAILURE_CLASSES) - - -def diag_query(problem: str, payload: Dict[str, Any]) -> str: - """题面 + 任务自身声明的硬约定（签名、必需库、返回规格、应抛异常、文档示例）。 - 只用题面信息、不含参考解答 —— 训练时同样拿得到，所以是「可得且非泄漏」的判据依据。""" - try: - doc = payload['doc_struct'] - doc = json.loads(doc) if isinstance(doc, str) else (doc or {}) - except Exception: - doc = {} - lines = [f"- required signature (must be reproduced verbatim):\n" - f"{payload['code_prompt'].strip()}"] - for key, label in (('reqs', 'must use these libraries'), ('returns', 'must return'), - ('raises', 'must raise'), ('params', 'parameters')): - vals = [str(x).strip() for x in (doc.get(key) or []) if str(x).strip()] - if vals: - lines.append(f'- {label}: ' + '; '.join(vals)) - ex = [str(x) for x in (doc.get('examples') or [])] - if ex: - lines.append('- documented example calls:\n ' + '\n '.join(ex)) - return problem + '\n\nHard requirements declared by the task:\n' + '\n'.join(lines) - - -def diag_segment(roll: Dict[str, Any]) -> str: - """★ rubric 路线唯一真正有效的一处：给 judge 的不是「输出全文」，而是**提交的代码 + 单测真实 - 报错**（<think> 已被 extract_code 切掉）。报错是客观事实且不含参考解答 —— 这正是 code 域 - rubric 有增量（+0.135, p=4e-5）而数学 / BFCL 域没有的原因。""" - return (f"### Submitted code\n```python\n{roll.get('code') or '(no parseable code block)'}\n```" - f"\n\n### Result of running the task's unit tests\n" - f"outcome: {roll.get('kind') or 'unknown'}\n" - f"{roll.get('error') or '(no error output)'}") - - -# =========================================================================== -# skill-gen（pitfall 文体 + rubric 条件化） -# =========================================================================== -# 与 narrative 文体的对照变量是**广度 vs 聚焦**：narrative 穷尽所有失败模式（~300 词），pitfall -# 只挑决定性的那一个（<90 词）。E23 选 pitfall 是因为 executor 已经能看到 actor 的 <think>， -# narrative 会与思考过程大面积重复，pitfall 让两者的分工是「过程 vs 结论」。 -SKILLGEN_SYSTEM = """\ -You are a skill-generation model. Your <skills> block will be fed to a SEPARATE downstream engineer model that must implement the function on its own. The engineer sees the same task description and the same required signature, but NOT your private reasoning or the analysis below — it only sees what is inside <skills>...</skills>. - -A diagnosis of a failed attempt at THIS task is provided to you: the class of the decisive failure, and the prior knowledge that would have prevented it. That prior is your material — it is a fact that stays true for other tasks too. Your job is to land it at the exact point in THIS task where it bites. - -Then, inside <skills></skills>, write under 90 words: -- WARNING: name where the failure strikes — the operation or the hand-off point it happens at — and the wrong assumption behind it, in the terms of the failure class you were given. -- INSTEAD: one or two sentences naming what to guarantee at that exact point, phrased as the general rule rather than as this task's answer. -- End by telling the engineer to deliver one single fenced code block that reproduces the given imports and signature verbatim, with no explanation, no demonstration call and no tests around it. - -Hard rules: -- Write about the failure class you were given. Do NOT substitute a different one and do NOT add a second: a made-up warning sends the engineer after something that is not the problem. -- Do NOT write the solution, and do NOT copy identifiers, column names, file names or literal values out of this task. Name the operation and describe the shape of what it returns instead of writing the call out. -- Self-contained: NEVER reference "the diagnosis", "the analysis", "the review" or "the previous attempt" — the engineer cannot see them, such phrasings cause hallucination. Address the engineer directly. - -Put ONLY the guidance inside <skills></skills>. - -Example: -<skills> -WARNING: the call doing the real work is accepted but does not behave as its name suggests: in the mode this task nudges you towards, it emits one record per row instead of one whole document, so the reader rejects it. -INSTEAD: confirm what that call emits in the mode you pick, and choose the mode whose output the consumer on the other side expects. -Deliver one fenced code block reproducing the given imports and signature verbatim, with no explanation, demonstration call or tests. -</skills> -""" - -SKILLGEN_USER = """\ -Task: -{problem} - -Diagnosis of a failed attempt (for your eyes only; do NOT reference it in the skill): -{rubric} - -Now write the <skills> guidance:""" - - -def skillgen_prompt(problem: str, rubric: str) -> Dict[str, Any]: - return {'messages': [{'role': 'system', 'content': SKILLGEN_SYSTEM}, - {'role': 'user', 'content': SKILLGEN_USER.format(problem=problem, - rubric=rubric)}]} diff --git a/cookbook/human/e23_rubric.py b/cookbook/human/e23_rubric.py deleted file mode 100644 index 185e408fc..000000000 --- a/cookbook/human/e23_rubric.py +++ /dev/null @@ -1,309 +0,0 @@ -"""教师 judge（失败分类 + 可迁移先验）与它的磁盘缓存。 - -诊断是 E23 的唯一自变量：分类不出决定性类的题一律丢弃，绝不降级成 query-only，否则自变量被稀释。 -与 v1 的区别是判据从「逐条 PASS/FAIL 的代码 review」换成「单一决定性类 + 报错证据 + 可迁移先验」， -理由见 e23_prompts.FAILURE_CLASSES 上方。 -""" -import collections -import hashlib -import json -import os -import re -from concurrent.futures import ThreadPoolExecutor -from typing import Any, Dict, List, Optional, Tuple - -from twinkle import get_logger -from twinkle_agentic.utils.llm_backup import llm_backup -from twinkle_agentic.verifier import RubricVerifier -from twinkle_agentic.verifier.rubric_verifier import RubricItem, _extract_json_obj, _short_hash - -from e23_prompts import (DIAG_SYSTEM, DIAG_USER, FAILURE_CLASSES, diag_query, diag_segment, - render_classes) - -logger = get_logger() -_HERE = os.path.dirname(os.path.abspath(__file__)) - -RUBRIC_WORKERS = int(os.environ.get('RUBRIC_WORKERS', 16)) -# **轨迹不在缓存键里** —— 任何改变裸解轨迹的开关（executor 的 thinking / executor 模型 / 任务域） -# 都必须体现在文件名上，否则会命中旧口径的诊断。 -RUBRIC_CACHE_PATH = os.environ.get( - 'RUBRIC_CACHE_PATH', os.path.join(_HERE, 'rubric_cache_global_code_execnothink_v3.jsonl')) -RUBRIC_VERSION = 'rubric_code_v3_mechanism' -# 独立根因数 >= 此值就丢题：实测（E23.t1，480 组）只坏 1 处的题救回率 0.504，坏 >=3 处的只有 -# 0.244 —— 一条 90 词内的警告救不回同时坏三处的尝试，训它只是稀释 batch。 -MAX_INDEPENDENT_CAUSES = 3 -# 判定温度。>0 换来的是标签多样性，代价是同一道题重判可能换类 —— 正好可以拿来测标签自一致率 -# （没有真值时这是唯一的准确性代理）；要逐字复现的诊断就设 0。 -DIAG_TEMPERATURE = float(os.environ.get('DIAG_TEMPERATURE', 0.3)) -# 教师 API 配错时每题都会失败 -> 每题都丢 -> 主循环一晚上零更新还不报错。一条都没成功过就早失败。 -API_FAIL_ABORT = int(os.environ.get('API_FAIL_ABORT', 20)) -_CLASS_SHORT = {code: short for code, short, _ in FAILURE_CLASSES} -# 分类表进 llm_backup 的置信度键：判据表一改，学生/教师一致性统计必须从头算。 -_CLASSES_KEY = _short_hash(render_classes()) -# ⭐ 提示词哈希进缓存键。之前只靠 RUBRIC_VERSION 这个手写标签，改了 DIAG_SYSTEM 却忘了改标签就会 -# 全量命中旧诊断、新规则一次都不执行，而且**毫无迹象**（日志里全是缓存命中，看起来一切正常）。 -# 实测踩过：加完「题面没给就弃权」那条规则后重跑，24 道题全命中 20 分钟前的旧判决。哈希兜住这个。 -_PROMPT_KEY = _short_hash(DIAG_SYSTEM + DIAG_USER + render_classes()) -# 先验里出现引号字面量 = 抄了本题的列名/文件名/期望值。默认只统计不丢弃：先看住这个比例，确认 -# 教师是否真的守住了「删掉本题这句话依然成立」，再决定要不要收紧成硬丢。 -PRIOR_REJECT_LITERALS = os.environ.get('PRIOR_REJECT_LITERALS', '0') == '1' -_LITERAL_RE = re.compile(r"'[^']{2,}'|\"[^\"]{2,}\"") - - -def _same_class(a: str, b: str) -> bool: - """student/teacher 是否算一致：只比决定性类。reason/prior 是自由文本，逐字比毫无意义。""" - ca = (_extract_json_obj(a) or {}).get('class') - return bool(ca) and ca == (_extract_json_obj(b) or {}).get('class') - - -class CodeRubricVerifier(RubricVerifier): - """代码域 judge：不走父类的逐条 PASS/FAIL，改成单一决定性类 + 证据 + 先验。 - - 复用父类的调用层而**不动共享代码**：父类的 _parse_diagnosis 只认 index/verdict/reason/fix， - evidence / prior / class 会被静默丢掉。 - - ⭐ 必须经 @llm_backup，不能直接调 _sample_text：本臂没有学生 sampler（build_checker 不传）， - _sample_text 此时按设计恒返回 ''，教师 API 完全由这个装饰器提供。直接调的话每道题都拿到空 - 响应、被判成「不可救」写进缓存，跑一遍就把缓存永久毒化。 - """ - - @llm_backup(key_params=['query', 'rubric_key'], comparator=_same_class) - def _classify_once(self, trajectory, sampling_params, query: str = None, - rubric_key: str = '') -> str: - return self._sample_text(trajectory, sampling_params, self.score_lora_path) - - def classify(self, query: str, segment_text: str) -> Optional[Dict[str, Any]]: - """返回解析出的 JSON；**None 专指「没拿到可解析响应」**（调用层故障，调用方不得缓存）。""" - traj = {'messages': [ - {'role': 'system', 'content': DIAG_SYSTEM}, - {'role': 'user', 'content': DIAG_USER.format( - query=query, rubric=render_classes(), segment=segment_text)}]} - raw = self._classify_once( - trajectory=traj, - sampling_params=self._diagnose_sampling_params(None, temperature=DIAG_TEMPERATURE), - query=query, rubric_key=_CLASSES_KEY) - return _extract_json_obj(raw) - - -def build_checker(): - """没有教师 API 就返回 None（调用方据此报错退出）。 - - fixed_rubric / gate 只为让父类构造合法：本臂从不读 detail.scalar，gate 与 is_hard 都不生效。 - """ - if not (os.environ.get('LLM_BACKUP_API_KEY') or os.environ.get('LLM_BACKUP_BASE_URL') - or os.environ.get('OPENAI_API_KEY')): - return None - return CodeRubricVerifier( - fixed_rubric=[RubricItem(f'{c}: {d}', is_hard=False) for c, _s, d in FAILURE_CLASSES], - gate=False) - - -# 这两类**按定义**就是「对不上调用方要求的某个东西」（前者：调用方读的东西没被设上；后者：该抛的 -# 异常没抛），required_value 为 null 在这里不是合法答案，只能是教师在走捷径。 -_CITATION_REQUIRED = {'TESTCONTRACT', 'EXCEPTION'} - - -def _cites_task(obj: Dict[str, Any], query: str) -> bool: - """教师声称「单测要的这个值题面里给了」时，核验它引的原文**真的**在题面里。 - - ⭐ 这一关不能只靠 DIAG_SYSTEM 的软性要求，两轮实测都被绕过： - 第一轮（只加「题面没给就弃权」的指令）—— 109/222/409 这类「标题/键名只存在于隐藏单测里」的题 - 照收不误，教师写一条 "tests assert exact string equality on titles" 的先验，这句话本身正确且 - 可迁移，于是自己说服自己题可救；可工程师读完仍不知道那字符串是什么，8 个候选必然一起挂。 - 第二轮（加 required_value / required_value_source 引用字段）—— 同样三道题又溜过去了，因为 - 教师把 required_value 填成 null，声称「本题失败与对不上指定值无关」，整关就不适用了。 - 所以对上面那两类**强制**要求引用：拿不出题面原文 = 题面确实没有 = 丢题。 - """ - val = str(obj.get('required_value') or '').strip() - cls = str(obj.get('class') or '').strip().upper() - if not val or val.lower() == 'null': - # 只有「这类失败本来就跟指定值无关」时才放行；两个强制类填 null 一律当没引用处理。 - return cls not in _CITATION_REQUIRED - src = str(obj.get('required_value_source') or '').strip() - if not src or src.lower() == 'null': - return False # 教师自认题面没给 -> 谁也做不出来 - return _squash(src) in _squash(query) # 引文对不上题面 = 编的 - - -def _validate(obj: Any, query: str) -> Optional[Dict[str, Any]]: - """把教师输出收敛成可用诊断；不合格返回 None -> 该题丢出训练集。 - - 弃权（addressable=false）、分类不在表内、拿不出报错证据、先验为空、独立根因 >= 3、单测要的值 - 题面里查无出处 —— 全部按「这题没法用一条不含解法的短警告救」处理。宁可丢题也不喂空洞诊断： - 每 chunk 产出约 40 道错题只需 24 道，丢得起（实测 t1 的 beyond_k 就丢了 328 道）。 - """ - if not isinstance(obj, dict) or not obj.get('addressable'): - return None - if not _cites_task(obj, query): - return None - cls = str(obj.get('class') or '').strip().upper() - prior = str(obj.get('prior') or '').strip() - evidence = str(obj.get('evidence') or '').strip() - if cls not in _CLASS_SHORT or not prior or not evidence: - return None - try: - n_causes = int(obj.get('independent_causes') or 1) - except (TypeError, ValueError): - n_causes = 1 - if n_causes >= MAX_INDEPENDENT_CAUSES: - return None - if PRIOR_REJECT_LITERALS and _LITERAL_RE.search(prior): - return None - secondary = [str(x).strip().upper() for x in (obj.get('secondary') or [])] - # 引用字段一并落缓存（**不进 _format**，不给 skill-gen 看）：上一轮排查时缓存里没有它们， - # 只能靠反推才确认教师是把 required_value 填了 null 溜过去的，白绕一圈。 - return {'class': cls, 'reason': str(obj.get('reason') or '').strip(), 'prior': prior, - 'evidence': evidence, 'n_causes': n_causes, - 'required_value': str(obj.get('required_value') or '').strip(), - 'required_value_source': str(obj.get('required_value_source') or '').strip(), - 'secondary': [s for s in secondary if s in _CLASS_SHORT and s != cls]} - - -# 单测报错里「期望值」的两种常见形态。⭐ 只用来**统计**教师该弃权时有没有弃权，绝不拦截：实测在 -# 352 道题上约两成假阳 —— 命中的字面量可能只是大小写/标点与题面不同（'Performance'、'Random -# Walk'），也可能抓到的是单测**构造输入**用的键而不是要求返回的键（BCB/524 的 'bird'/'fish'）。 -# 归一化比较能消掉前一类，后一类消不掉，所以这个信号只配当监控，判决交给看得见题面的教师。 -_EXPECT_RE = re.compile(r"""!=\s*(['"])(.{2,60}?)\1|KeyError:\s*(['"])(.{1,40}?)\3""") - - -def _squash(s: str) -> str: - return re.sub(r'[^a-z0-9]', '', (s or '').lower()) - - -def _unseen_expected_literal(evidence: str, query: str) -> bool: - """报错要求的字面量在题面里查无此物 = 这题从题面根本做不出来，教师本该弃权。""" - hay = _squash(query) - for m in _EXPECT_RE.finditer(evidence or ''): - lit = m.group(2) or m.group(4) - if lit and _squash(lit) not in hay: - return True - return False - - -def _format(d: Dict[str, Any]) -> str: - """诊断 -> 喂给 skill-gen 的纯文本。 - - ⭐ evidence 故意**不进** skill-gen：它是单测报错原文，断言 diff 里带期望值，是最强的泄漏通道。 - 它留在缓存里只为两件事 —— 逼教师把分类落到客观事实上，以及事后审计。 - """ - lines = [f"DECISIVE FAILURE: {d['class']} — {_CLASS_SHORT[d['class']]}", - f"WHAT WENT WRONG: {d['reason']}", - f"PRIOR THAT WOULD HAVE PREVENTED IT: {d['prior']}"] - if d['secondary']: - lines.append('ALSO OFF (do not write about these): ' + ', '.join(d['secondary'])) - return '\n'.join(lines) - - -def class_metrics(diag_texts: List[str]) -> Dict[str, float]: - """标签退化监控：一组题用到几个决定性类、最大那类占多少。 - - 判据表的全部价值在于分得开：一旦某一类吃掉大半（旧表的兜底项吃了 76% 的题），所有题的诊断 - 就长得一样，skill-gen 失去逐题条件化，等于退回没有 rubric 的对照臂。 - """ - codes = [m.group(1) for m in - (re.match(r'DECISIVE FAILURE:\s*([A-Z]+)', (t or '').split('\n', 1)[0]) - for t in diag_texts) if m] - c = collections.Counter(codes) - return {'signal/rubric_n_classes': float(len(c)), - 'signal/rubric_top_share': max(c.values()) / len(codes) if codes else 0.0} - - -def cache_metrics(delta: collections.Counter) -> Dict[str, float]: - """一步之内 rubric 侧的计数，进 train_log。 - - unseen_literal 是弃权规则的自查通道：教师本该判「题面没给」却收下的题数。要和 dropped 一起 - 读 —— dropped 不涨而 unseen_literal 在涨，就说明 DIAG_SYSTEM 那条规则没被遵守。 - """ - return {'signal/rubric_dropped': float(delta['dropped_unaddressable'] + delta['hit_dropped']), - 'signal/rubric_dropped_not_stated': float(delta['dropped_not_stated']), - 'signal/rubric_api_fail': float(delta['api_fail']), - 'signal/rubric_unseen_literal': float(delta['unseen_literal']), - 'signal/rubric_prior_has_literal': float(delta['prior_has_literal'])} - - -class RubricCache: - """append-only jsonl + 内存索引，键 = md5(RUBRIC_VERSION, 提示词哈希, data_id)，值 = 诊断 JSON。 - - 能跨 run 复用的**唯一依据**是「executor 冻结在 T=0，所以同一道题的裸解轨迹在所有 run 里逐字 - 相同」。轨迹不在键里，见 RUBRIC_CACHE_PATH 上方的换名要求。 - 存 JSON 而不是成品文本：改 _format 的排版不必重拉 API，evidence 也留得住可审计。 - """ - - def __init__(self, path: str = RUBRIC_CACHE_PATH): - self.path = path - self._idx: Dict[str, Any] = {} - self.stats: collections.Counter = collections.Counter() - if os.path.exists(path): - with open(path, encoding='utf-8') as f: - for line in f: - try: - rec = json.loads(line) - self._idx[rec['key']] = rec['value'] - except Exception: - continue - logger.info(f'[rubric] 缓存载入 {len(self._idx)} 条：{path}') - self._fh = open(path, 'a', encoding='utf-8') - - def _put(self, key: str, value: Any) -> None: - self._idx[key] = value - self._fh.write(json.dumps({'key': key, 'value': value}, ensure_ascii=False) + '\n') - self._fh.flush() - - def get_or_diagnose(self, checker, record: Dict[str, Any], roll: Dict[str, Any]) -> str: - """返回诊断文本；教师弃权/不可救返回 '' 并**缓存该判决**，API 失败返回 '' 但**不缓存**。 - - 两者都让调用方丢掉这道题，但只有前者是稳定判决 —— 把一次瞬时抖动写进缓存会永久毒化它。 - """ - key = hashlib.md5(f"{RUBRIC_VERSION}\x00{_PROMPT_KEY}\x00" - f"{record.get('data_id', '')}".encode('utf-8')).hexdigest() - if key in self._idx: - cached = self._idx[key] - if not cached: - self.stats['hit_dropped'] += 1 - return '' - self.stats['hit'] += 1 - return _format(cached) - query = diag_query(record['problem'], record['reference_answer']) - try: - obj = checker.classify(query, diag_segment(roll)) - except Exception as exc: - logger.warning(f'[rubric] classify error: {exc}') - obj = None - if obj is None: - # 拿不到可解析响应 = 调用层故障，**绝不能**当成「教师判不可救」写进缓存：那会把一次 - # 抖动变成永久丢题。教师彻底不通时一条都不会成功，与其空转一夜不如早失败。 - self.stats['api_fail'] += 1 - if self.stats['api_fail'] >= API_FAIL_ABORT and not self.stats['ok']: - raise RuntimeError( - f'[rubric] 连续 {self.stats["api_fail"]} 次拿不到教师响应且无一成功；' - f'检查 LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / LLM_BACKUP_MODEL') - return '' - diag = _validate(obj, query) - if diag is None: - # 缓存空值 = 「这题教师给不出可用分类」，是稳定判决，下次直接跳过不再花 API 钱。 - if isinstance(obj, dict) and obj.get('addressable') and not _cites_task(obj, query): - # 单独计数：这是「题面根本没给」，与「教师主动弃权」是两种不同的丢题原因。 - self.stats['dropped_not_stated'] += 1 - self._put(key, None) - self.stats['dropped_unaddressable'] += 1 - return '' - if _LITERAL_RE.search(diag['prior']): - self.stats['prior_has_literal'] += 1 - if _unseen_expected_literal(diag['evidence'], query): - # 教师收下了一道「答案只存在于隐藏单测里」的题。不改判决（机械检测精度不够），但这个 - # 计数持续偏高就说明 DIAG_SYSTEM 的弃权规则没被遵守。 - self.stats['unseen_literal'] += 1 - self.stats['ok'] += 1 - self.stats[f"class_{diag['class']}"] += 1 - self._put(key, diag) - return _format(diag) - - def diagnose_many(self, checker, pairs: List[Tuple[Dict, Dict]]) -> List[str]: - """并行拉诊断（纯 API 调用，不占 GPU）。pairs = [(record, roll), ...]。""" - if not pairs: - return [] - with ThreadPoolExecutor(max_workers=max(1, min(RUBRIC_WORKERS, len(pairs)))) as ex: - return list(ex.map(lambda rb: self.get_or_diagnose(checker, rb[0], rb[1]), pairs)) - - def close(self): - self._fh.close() diff --git a/cookbook/human/skill_drift_stats.py b/cookbook/human/skill_drift_stats.py deleted file mode 100644 index 405ae95af..000000000 --- a/cookbook/human/skill_drift_stats.py +++ /dev/null @@ -1,168 +0,0 @@ -#!/usr/bin/env python3 -# Copyright (c) ModelScope Contributors. All rights reserved. -"""统计 E23 skill 随训练 step 的「长度」与「词频」漂移，量化「skill 是否越来越空泛」。 - -⚠️ 数据来源与偏差 ------------------- -本脚本读 output.e23/zero_reward_groups.jsonl —— 它**只落整组 reward=0 的题**（8 个候选全错）。 -这不是全部 skill，是一个偏难的子集。所以： - * 「长度随 step 涨」的结论对这个子集成立，但推广到全体 skill 时要记住这一点； - * 词频漂移（guarantee->confirm 等）同理。 -若要无偏统计需改 e23 落全量 skill；当前只有这一份带 step 标签的 skill 文本。 - -口径 ----- -* skill = 候选的 <skills> 块（extract_skill 已抽好，存在 candidates[].skills）。 -* INSTEAD 段单独抽出来看动词，因为「空泛化」主要发生在这一段（WARNING 段是描述过去的错误）。 -* 词频对比早期（step<=2）vs 晚期（step>=13）两窗，报 log2 比值最大的上升/下降词。 -""" -import collections -import json -import math -import os -import re -import statistics as st -import sys - -HERE = os.path.dirname(os.path.abspath(__file__)) -DEFAULT = os.path.join(HERE, 'output.e23', 'zero_reward_groups.jsonl') - -# 英文停用词（够用即可，不引第三方）。 -STOP = set('the a an of to in on for and or is are be it its this that with as by from at ' - 'you your not no if then when will would should must can may a an s t re ve'.split()) -INSTEAD_RE = re.compile(r'INSTEAD:\s*(.*?)(?:\n(?:Deliver|WARNING)|\Z)', re.S) -WARNING_RE = re.compile(r'WARNING:\s*(.*?)(?:\n(?:INSTEAD|Deliver)|\Z)', re.S) -# 空泛/认知性措辞：不要求 executor 改任何代码，只要求「认同一条命题」。 -VAGUE_RE = re.compile( - r'\b(confirm that|in any|for any|general principle|it may|might|unpredictab\w+|' - r'violat\w+ the principle|be aware|keep in mind|note that|understand that|' - r'is a general|in general|conceptually|principle)\b', re.I) -# 祈使/可执行动词：直接命令 executor 怎么写。 -IMPER_RE = re.compile(r'^(use|set|derive|write|apply|replace|compute|return|add|remove|' - r'ensure|guarantee|call|pass|cast|convert|assign|initialize|import|' - r'check|handle|raise|match)\b', re.I) - - -def words(text): - return [w for w in re.findall(r"[a-zA-Z_][a-zA-Z_']+", (text or '').lower()) - if w not in STOP and len(w) > 2] - - -def load(): - path = sys.argv[1] if len(sys.argv) > 1 else DEFAULT - rows = [json.loads(l) for l in open(path, encoding='utf-8') if l.strip()] - # 每个候选一条记录：(step, skill_text) - out = [] - for r in rows: - s = r['step'] - for c in r['candidates']: - out.append((s, c.get('skills') or '')) - return out, path - - -def instead_seg(sk): - m = INSTEAD_RE.search(sk or '') - return (m.group(1).strip() if m else '') - - -def warning_seg(sk): - m = WARNING_RE.search(sk or '') - return (m.group(1).strip() if m else '') - - -def per_step_table(data): - by = collections.defaultdict(list) - for s, sk in data: - by[s].append(sk) - steps = sorted(by) - print('=' * 88) - print('一、skill 长度与结构 随 step（每行 = 该 step 全部候选的均值；n=候选数）') - print('=' * 88) - print(f"{'step':>4} {'n':>4} {'skill字符':>9} {'skill词数':>9} {'INSTEAD词数':>11} " - f"{'空泛词/条':>9} {'祈使开头%':>9}") - rows_for_trend = [] - for s in steps: - sks = by[s] - chars = st.mean(len(x) for x in sks) - nwords = st.mean(len(words(x)) for x in sks) - ins = [instead_seg(x) for x in sks] - inw = st.mean(len(i.split()) for i in ins) if ins else 0 - vague = st.mean(len(VAGUE_RE.findall(x)) for x in sks) - imper = sum(1 for i in ins if IMPER_RE.match(i)) / len(ins) if ins else 0 - print(f'{s:>4} {len(sks):>4} {chars:>9.1f} {nwords:>9.1f} {inw:>11.1f} ' - f'{vague:>9.2f} {imper:>8.0%}') - rows_for_trend.append((s, chars, vague, imper)) - return by, steps, rows_for_trend - - -def trend(rows_for_trend): - """对 (step, y) 做最小二乘斜率 + t，判断长度/空泛度/祈使率是否真在漂移。""" - print('\n' + '=' * 88) - print('二、趋势显著性（OLS 斜率 / step，|t|>2 才算真漂移）') - print('=' * 88) - xs = [r[0] for r in rows_for_trend] - n = len(xs) - mx = st.mean(xs) - sxx = sum((x - mx) ** 2 for x in xs) - for idx, name in ((1, 'skill 字符数'), (2, '空泛词/条'), (3, 'INSTEAD 祈使开头率')): - ys = [r[idx] for r in rows_for_trend] - my = st.mean(ys) - slope = sum((x - mx) * (y - my) for x, y in zip(xs, ys)) / sxx - resid = [y - (my + slope * (x - mx)) for x, y in zip(xs, ys)] - s2 = sum(e * e for e in resid) / (n - 2) - se = math.sqrt(s2 / sxx) if sxx else 0.0 - t = slope / se if se else 0.0 - tag = '显著' if abs(t) > 2 else '不显著' - print(f' {name:20s} 斜率={slope:+.4f}/step t={t:+.2f} [{tag}] ' - f'首={ys[0]:.3f} 尾={ys[-1]:.3f}') - - -def word_freq_shift(by, steps): - """早期 vs 晚期两窗的词频比。""" - print('\n' + '=' * 88) - print('三、词频漂移：早期(step<=2) vs 晚期(step>=13)') - print('=' * 88) - early_steps = [s for s in steps if s <= 2] - late_steps = [s for s in steps if s >= 13] - - def counts(win): - c = collections.Counter() - ntok = 0 - for s in win: - for sk in by[s]: - w = words(instead_seg(sk)) # 只看 INSTEAD 段 - c.update(w) - ntok += len(w) - return c, ntok - ce, ne = counts(early_steps) - cl, nl = counts(late_steps) - print(f'早期窗 step={early_steps} INSTEAD总词={ne}；晚期窗 step={late_steps} 总词={nl}') - # 频率（每千词），加平滑 - vocab = set(ce) | set(cl) - rows = [] - for w in vocab: - fe = (ce[w] + 0.5) / (ne + 1) * 1000 - fl = (cl[w] + 0.5) / (nl + 1) * 1000 - if ce[w] + cl[w] < 4: # 太稀疏的词不看 - continue - rows.append((math.log2(fl / fe), w, ce[w], cl[w], fe, fl)) - rows.sort(reverse=True) - print(f'\n{"↑晚期变多的词":22s}{"早/千":>8}{"晚/千":>8}{"log2比":>8}') - for lr, w, e, l, fe, fl in rows[:15]: - print(f' {w:20s}{fe:8.1f}{fl:8.1f}{lr:+8.2f}') - print(f'\n{"↓晚期变少的词":22s}{"早/千":>8}{"晚/千":>8}{"log2比":>8}') - for lr, w, e, l, fe, fl in rows[-15:][::-1]: - print(f' {w:20s}{fe:8.1f}{fl:8.1f}{lr:+8.2f}') - - -def main(): - data, path = load() - print(f'[数据] {path}') - print(f'[数据] {len(data)} 个候选 skill（来自整组 reward=0 的题，是偏难子集，非全量）\n') - by, steps, rows_for_trend = per_step_table(data) - trend(rows_for_trend) - word_freq_shift(by, steps) - - -if __name__ == '__main__': - main() diff --git a/cookbook/human_e18/README.md b/cookbook/human_e18/README.md deleted file mode 100644 index 62150ae5a..000000000 --- a/cookbook/human_e18/README.md +++ /dev/null @@ -1,77 +0,0 @@ -# E18 — BigCodeBench 上的拒绝采样 SFT - -从 `cookbook/exp/skill2lora/skill_ablate/config.py` L244-252 的 `ExpSpec('E18', 'rejection_sft', ...)` -单独脱离出来的自包含实现，目录布局仿照 `cookbook/human`（E23）。 - -## 与 E23 的关系 - -两臂**共用**环境层与教师 judge（直接 import `../human/e23_bcb.py`、`../human/e23_rubric.py`， -不拷贝），所以数据过滤、沙箱单测判分、rubric 诊断三处逐字同源、结果可直接比。差别只在训练方法： - -| | E23 | E18（本目录） | -|---|---|---| -| 方法 | GRPO（组内归一化 advantage） | 拒绝采样 SFT（只用正样本） | -| loss | `SEAMBNPOLoss`（PPO clip + KL） | **`CrossEntropyLoss`**（纯交叉熵，无 ratio/clip/KL/advantage） | -| 梯度 | 正负都有；60% 组零方差→零梯度 | 只有正梯度，无零梯度浪费 | -| 每题产出 | 8 个候选全部进 batch | 两道筛后**唯一胜者**进池 | -| 训练轨迹 | 采样 token 直通（query+rubric） | messages 编码（**query-only**） | -| eval | 带 rubric | **query-only**（部署口径） | -| 卡数 | 8（含 ref） | 6（SFT 无需 ref 模型） | - -「选择」全部发生在两道筛（只有胜者入池），到 loss 这一层就是普通的「拟合目标文本」， -所以不传 `advantages` —— `CrossEntropyLoss` 不读该参数，传了是静默无效。 - -## 两道筛（本臂唯一自变量，见 `e18_select.py`） - -1. **增量达阈**：`with_pass >= base_pass_rate + MIN_PASS_GAIN`（默认 +2/8）。仅仅「没弄坏」 - （8/8 -> 8/8）不够格 —— 那种样本对「学会写有效 skill」没有监督信号； -2. **不超长**：超 `SKILL_CHAR_LIMIT` 直接丢； -3. **pass_rate 最大 -> 并列内 rubric 相似度**：先按客观效果取最大档，并列内取词频余弦相似度 - 最高的。长度只在相似度也并列时做确定性拆平（取较短者）。 - -⚠️ 原先第 3 道筛是「先按离 `LEN_BUDGET`（400）的距离取前一半」，已删：实测 **66% 的题 8 个 -候选全部 `with_pass=1.0`**（天花板打平），此时长度成了事实上的唯一决策依据，而它有系统性 -偏差——中文表达同样内容字符数天然更少（357 vs 705），永远更贴近 400，于是「离预算最近」被 -翻译成「选中文模板」，把信息量大的长英文候选全部淘汰。 - -⚠️ 原先有一道泄漏门（`leak_blocks`），已删：BCB 的 `reference_answer` 是 ~2500 字符的 dict， -子串匹配要求 skill 逐字包含整个 dict 的 repr，而 skill 上限 1500 字符 —— 触发概率恒为 0。 -真正的泄漏通道在 `test` 的断言期望值与 `canonical_solution`，需另写检测。 - -⚠️ 存活候选 ≤2 条时相似度那阶形同虚设，胜者由长度决定 —— 详见 `select_winner` 的 docstring。 - -## 文件 - -| 文件 | 内容 | -|---|---| -| `e18_rejection_sft.py` | 采集 + SFT 主循环、eval、swanlab | -| `e18_prompts.py` | executor / skill-gen / 训练轨迹三处 prompt | -| `e18_select.py` | 三道拒绝筛（泄漏门 + 词频余弦 + 两阶选择） | - -## 跑法 - -```bash -cd cookbook/human_e18 -# 教师 API 必需（rubric 是选择的参照系，不可降级） -export LLM_BACKUP_API_KEY=... LLM_BACKUP_BASE_URL=... -nohup python -u e18_rejection_sft.py > nohup.e18.$(date +%m%d-%H%M).log 2>&1 & -``` - -主要环境变量（默认值见文件头）： - -| 变量 | 默认 | 说明 | -|---|---|---| -| `ACCUMULATE` | 16 | 攒够多少条胜者 SFT 一次（须为 `TRAIN_DP` 整数倍） | -| `N_SKILLS` | 8 | 每题候选数 = 拒绝采样的池大小 | -| `MAX_UPDATES` | 50 | 总更新数 | -| `SKILL_CHAR_LIMIT` | 1500 | 超过直接丢 | -| `EVAL_SIZE` | 200 | holdout 题数；0 = 不 eval | -| `SWAN_PROJ` | twinkle | swanlab 项目（**勿** export `SWANLAB_PROJECT`） | - -## 产物 - -| 文件 | 内容 | -|---|---| -| `output.e18/e18_sft_dataset.jsonl` | **主产物**：每条胜者 + rubric/相似度/pass 全审计字段，可离线复算、换超参重训而不必重跑 GPU | -| `output.e18/train_log.jsonl` | 逐 chunk 指标（accept_rate / candidate_pass_rate / eval/*） | -| `output.e18/E18-final/` | 最终 skill 模型权重 | diff --git a/cookbook/human_e18/e18_collect_kod.py b/cookbook/human_e18/e18_collect_kod.py deleted file mode 100644 index 55efca255..000000000 --- a/cookbook/human_e18/e18_collect_kod.py +++ /dev/null @@ -1,738 +0,0 @@ -# -*- coding: utf-8 -*- -"""E18 冷启动数据采集（KodCode 域，**无 SFT**）。 - -与 `e18_rejection_sft.py` 的关系：把它的采集半部分原样搬过来，删掉训练/eval/权重同步。 -所以 `collect_chunk` 的逻辑、三道筛、指标口径、落盘字段全部逐字一致 —— 这样采出来的 -冷启动数据集与在线 run 的样本同分布，后续接 SFT 时不需要再对齐一次。 - -**删掉了什么，以及为什么** -* `TransformersModel` / `set_loss` / `set_optimizer` / `train_batch`：不训练。 -* `CheckpointEngineManager` / `_sync_trained_to_sampler` / `_restore_base_weights`： - 没有训练权重要推给 sampler，skill_sampler 全程是初始模型。 -* `run_eval` / `EVAL_SIZE`：eval 衡量的是「训练后的 skill 模型在部署口径下的能力」， - 不训练时它恒等于 baseline，跑它纯浪费 GPU。故 `load_records(eval_size=0)`，题全进采集池。 - -**8 卡全部给 rollout**：原来 train 占 2 张、skill_sampler 2 张、base_sampler 4 张。 -现在 train 那 2 张转给两个 sampler。base_sampler 仍拿大头（默认 6）——它是唯一瓶颈： -每 chunk 它要跑裸解 CHUNK_SIZE*BARE_ROLLOUTS + 重解 CHUNK_SIZE*N_SKILLS*EXEC_ROLLOUTS， -序列数比 skill_sampler 多一个量级，且 EXEC_MAX_TOKENS 远大于 SKILL_MAX_TOKENS。 - -产物（都在 OUTPUT_DIR 下，append-only）： -* `e18_sft_dataset.jsonl`：胜者，字段与在线 run 完全一致，直接可喂 SFT。 -* `e18_candidates.jsonl`：**全部** skill 候选（含落选与解析失败的），靠 `kept` 区分选与未选。 - 用于离线重算阀值、判断胜者是真更好还是拆平局选出来的。 -* `collect_log.jsonl`：逐 chunk 指标，用于监控 accept_rate / degrade_rate 是否异常。 -""" -import json -import os -import shutil -import sys -import time -import zlib -from dataclasses import dataclass -from typing import Any, Dict, List, Tuple - -import torch -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.data_format import SamplingParams -from twinkle.dataloader import DataLoader -from twinkle.sampler import vLLMSampler -from twinkle.template import Template - -_HERE = os.path.dirname(os.path.abspath(__file__)) -_COOKBOOK = os.path.abspath(os.path.join(_HERE, '..')) -for _p in (_HERE, os.path.join(_COOKBOOK, 'human')): - if _p not in sys.path: - sys.path.insert(0, _p) - -from e23_rubric import build_checker, class_metrics # noqa: E402 - -from e18_kodcode import (clean_text, empty_roll, extract_skill, # noqa: E402 - judge_seqs, load_records) -from e18_multidiag import MultiDiagCache, multidiag_metrics # noqa: E402 -from e18_prompts import (direct_prompt, format_trajectory, # noqa: E402 - skill_solve_prompt, skillgen_prompt) -from e18_select import gain_stats, select_winner # noqa: E402 - -logger = get_logger() - -# ========== Configuration ========== -MODEL_ID = os.environ.get('MODEL_ID', 'Qwen/Qwen3-4B') -OUTPUT_DIR = os.environ.get('OUTPUT_DIR', os.path.join(_HERE, 'output.e18.kod')) - -# ⭐ 8 卡全给 rollout：不训练，所以没有 train 组。base_sampler 拿大头（瓶颈见文件头注释）。 -SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) -BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 6)) -NUM_GPUS = SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS -GPU_MEM = float(os.environ.get('GPU_MEM', 0.8)) - -SEED = int(os.environ.get('SEED', 42)) -# ⭐ 续跑开关。为何默认关：开着会把旧 run 的样本接着往同一批数据里添，而跨 run 的 -# 模型/prompt 可能已经变了 —— 默认开启等于默认允许污染，违反归档机制的初衷。 -# KOD_RESUME=1 时做三件事（见 archive_output_dir / resume_done_ids）： -# 1. 归档**之前**先从 e18_candidates.jsonl 读出已跑过的 data_id； -# 2. 三个 jsonl 复制回新目录（而不是只搬 broken_tasks 缓存），让计数接着累积； -# 3. 把已采题从题池里 filter 掉。 -# 为何必须有它：DataLoader 用固定 SEED+shuffle，重启后取题顺序逐字相同，不过滤就会 -# 把前 N 个 chunk 的题原样重跑一遂（实测 65 chunk ≈ 11 小时 8 卡），纯浪费。 -KOD_RESUME = os.environ.get('KOD_RESUME', '0') == '1' -CHUNK_SIZE = int(os.environ.get('CHUNK_SIZE', 64)) -# ⭐ 4 而不是 8：skill 候选池只用来「邀出」候选，最终只有 1 条胜者入池。实测 66% 的题 -# 8 个候选的 with_pass 全部并列，多出来的 4 条几乎不改变胜者，却要占掉一半 executor 序列。 -# ⭐ N_SKILLS 现在是**两阶段的总上限**：先生成 N_SKILLS_STAGE1 个，全部不达标才补到 N_SKILLS。 -# 实测（4224 题）：39.7% 的题前 2 个候选就能出胜者 -> 平均只花 3.21 个候选， -# 产出与固定 4 个**完全相同**（1134 题），executor 序列从 32 降到 25.6。 -N_SKILLS = int(os.environ.get('N_SKILLS', 4)) -N_SKILLS_STAGE1 = int(os.environ.get('N_SKILLS_STAGE1', 2)) -# ⭐ 天花板短路：base_pass_rate 已打满的题直接跳过，不生成任何 skill 候选。 -# 依据：实测 1750 道 base=1.0 的题，入池 **0 条** —— 因为门槛是 -# pass_gain >= MIN_PASS_GAIN(0.25)，而 base=1.0 时 with_pass 最大也是 1.0，gain 恒为 0。 -# 所以这 41% 的题在**本离线采集脚本**里是纯浪费（占 37% 算力、零产出）。 -# ⚠️ 与 collect_chunk 原注释「全量 rollout」的设计意图相反：那条理由（避免 skill 模型 -# 只学会救难题）适用于**在线 RL**（每步需要 reward 信号，包括平局组）； -# 本脚本是纯离线采集，产物只有 e18_sft_dataset.jsonl，天花板题从不进该文件。 -# 若日后要拿这份代码回到在线 RL，必须把本开关置 0。 -SKIP_CEILING = int(os.environ.get('SKIP_CEILING', 1)) -# ⭐ 粗筛 rollout 次数：先给每个候选跑 PROBE_ROLLOUTS 次，只对**并列最高**者补到 -# EXEC_ROLLOUTS 次。实测（1515 题）：M=2 的 top-1 与 8 次一致率 68.9%、M=4 为 77.1%； -# 而 51% 的题 4 候选 with_pass 全同分，前 2 次就能看出打平并提前停手。 -# 0 = 关闭粗筛（所有候选直接跑足 EXEC_ROLLOUTS 次，与历史 run 逐字一致）。 -PROBE_ROLLOUTS = int(os.environ.get('PROBE_ROLLOUTS', 2)) -# 采够多少条胜者就停。0 = 把题池跑完。 -TARGET_SAMPLES = int(os.environ.get('TARGET_SAMPLES', 20000)) -MAX_CHUNKS = int(os.environ.get('MAX_CHUNKS', 0)) # 0 = 不限 - -# ⭐ 多机并行分片。SHARD_ID 取值 0..SHARD_N-1，每台机器只跑 -# `crc32(data_id) % SHARD_N == SHARD_ID` 的题。 -# 为何用哈希而不是「切片取前后一半」： -# 1. 无状态 —— 两台机器不需要任何通信/共享盘就能保证不重叠（NAS 不通用时的必要条件） -# 2. 对题池变化稳健 —— 万一两边题池大小不一致（版本/缓存差异），按下标切会错位重叠， -# 而哈希绑定在 data_id 上，永远不会 -# 3. 难度分布无偏 —— crc32 与 gpt_pass_percentage 无相关，两片难度同分布 -# ❗ 两台机器必须用**相同的 SHARD_N**，否则分片不构成划分（会既重叠又遗漏）。 -# ❗ TARGET_SAMPLES 是**本分片的**目标：想总共 20000 就两边各填 10000。 -SHARD_N = int(os.environ.get('SHARD_N', 1)) -SHARD_ID = int(os.environ.get('SHARD_ID', 0)) - -SKILL_MAX_TOKENS = int(os.environ.get('SKILL_MAX_TOKENS', 8192)) -# ⭐ 把失败轨迹（裸解时做错的那份代码）一并给 skillmodel。默认**关**： -# 开了就换了 prompt 口径，与已采的 1750 条不同源，不能默认静默切换。 -# 长度安全性已实测（.tmp_analysis/len_budget.py，400 条真实数据）： -# prompt 预算 = MAX_MODEL_LEN(16000) - SKILL_MAX_TOKENS(8192) = 7808 token -# 现状不带轨迹 中位 1802 / 最大 3192 -# 带 1 条(3x 最坏) 中位 2710 / 最大 6162 -> 仍在预算内，**无需改 MAX_MODEL_LEN** -# 带 2 条(3x 最坏) 最大 9132 -> 1.5% 超预算，所以默认只给 1 条 -# 超预算的后果不是报错而是 vLLM 返回空序列 -> parseable=False -> 该候选白跑， -# 难以从日志发现，所以宁可保守。 -USE_TRAJ = int(os.environ.get('KOD_USE_TRAJ', 0)) -TRAJ_N = int(os.environ.get('KOD_TRAJ_N', 1)) -# 单条轨迹的字符上限。4000 字符 ≈ 1540 token（代码 2.6 chars/token）， -# 加上现状最大 3192 仍不到 7808。超长者由 format_trajectory 头尾各留一半。 -TRAJ_MAX_CHARS = int(os.environ.get('KOD_TRAJ_MAX_CHARS', 4000)) -SKILL_GEN_TEMPERATURE = float(os.environ.get('SKILL_GEN_TEMPERATURE', 1.0)) -SKILL_GEN_TOP_P = float(os.environ.get('SKILL_GEN_TOP_P', 1.0)) -SKILL_GEN_TOP_K = int(os.environ.get('SKILL_GEN_TOP_K', -1)) -EXEC_MAX_TOKENS = int(os.environ.get('EXEC_MAX_TOKENS', 15000)) -MAX_MODEL_LEN = int(os.environ.get('MAX_MODEL_LEN', 16000)) - -EXEC_ROLLOUTS = int(os.environ.get('EXEC_ROLLOUTS', 8)) -# ⭐ 裸解单独用 4 次，而带 skill 重解仍是 EXEC_ROLLOUTS(8)。两侧刻意**不对称**： -# 裸解只承担「这题有没有提升空间」的粗判（rate<1 即进诊断），4 次足够；而 with_pass 要 -# 在候选之间排序、还要减去 base 算增量，精度需求高得多，降它会直接动摇入池门槛的语义。 -# ⚠️ 代价（已知并接受）：base_pass_rate 只有 5 档（0,.25,.5,.75,1），与 with_pass 的 9 档 -# 不同分母。于是 pass_gain = with_pass - base_pass_rate 的零点变粗，base 的采样标准误从 -# 0.177 升到 0.25 —— 判「有没有空间」够用，但别再把 pass_gain 的绝对值当精密量看。 -# 另外 n_ceiling（base 打满 4/4）会比 8/8 更容易达成，天花板题占比会上升。 -BARE_ROLLOUTS = int(os.environ.get('BARE_ROLLOUTS', 4)) -EXEC_TEMPERATURE = float(os.environ.get('EXEC_TEMPERATURE', 0.6)) -EXEC_TOP_P = float(os.environ.get('EXEC_TOP_P', 0.95)) - -SKILL_CHAR_LIMIT = int(os.environ.get('SKILL_CHAR_LIMIT', 1500)) -# 门槛仍按 EXEC_ROLLOUTS(8) 算 = +2/8 = +0.25，与历史 run 逐字可比。 -# 不能拿 BARE_ROLLOUTS 做分母：with_pass 是 8 次采的，两边分母必须是同一个。 -MIN_GAIN_ROLLOUTS = int(os.environ.get('MIN_GAIN_ROLLOUTS', 2)) -MIN_PASS_GAIN = MIN_GAIN_ROLLOUTS / max(1, EXEC_ROLLOUTS) - -# ⭐ 多机并行时 RUN_ID 必须带机器标识：原来只有时间戳，两台机器同一秒启动会撞成 -# 同一个 run，合并后就再也分不出某条数据是哪台机器产的（排查单机异常时必须能分开）。 -# 分片号是天然的机器标识，比 hostname 稳（容器重建后 hostname 会变），所以 -# SHARD_N>1 时后缀 `.sN` —— 单机跑时 RUN_ID 保持原格式，历史 run 的比对不受影响。 -RUN_ID = time.strftime('%m%d-%H%M%S') + (f'.s{SHARD_ID}' if SHARD_N > 1 else '') - - -@dataclass -class Runtime: - skill_sampler: Any - base_sampler: Any - checker: Any - rubric_cache: MultiDiagCache - - -# =========================================================================== -# 采样工具（与 e18_rejection_sft 逐字一致） -# =========================================================================== -def run_samples(sampler, prompts, num_samples, max_tokens, gen_dp, - temperature=None, top_p=None, top_k=None, logprobs=None): - """采样。prompts 少于 dp 时补齐再截回 —— Ray 的 dp 切分要求每个 rank 至少一条。 - - ⭐ 与 e18_rejection_sft.run_samples 逐字一致。三处踩过的坑： - 1. 字段名是 `num_samples` 不是 `n`（SamplingParams 没有 `n`，传了直接 TypeError）。 - 2. 走 `sampler.sample(prompts, params)`，不是 pack_user_data + generate_sequences。 - 3. dp 补齐不能省：最后一个 chunk 不满、或 flat 很少时，条数 < dp 会直接报错。 - """ - if not prompts: - return [] - import copy - params = SamplingParams( - max_tokens=max_tokens, - temperature=0.6 if temperature is None else temperature, - top_p=0.95 if top_p is None else top_p, - num_samples=num_samples, - **({} if top_k is None else {'top_k': top_k}), - **({} if logprobs is None else {'logprobs': logprobs})) - padded = prompts - if gen_dp > 1 and 0 < len(prompts) < gen_dp: - padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] - responses = sampler.sample(padded, params)[:len(prompts)] - return [list(r.sequences) if (r and r.sequences) else [] for r in responses] - - -def first_seq(seqs): - return seqs[0] if seqs else None - - -def seq_text(seq) -> str: - return clean_text(getattr(seq, 'decoded', '') or '') if seq is not None else '' - - -def _mean(xs) -> float: - xs = [float(x) for x in xs if x is not None] - return sum(xs) / len(xs) if xs else 0.0 - - -# =========================================================================== -# 采集：裸解 -> 诊断 -> skill-gen -> executor 重解 -> 三道筛 -# =========================================================================== -def _pass_rate(rolls) -> float: - return (sum(1.0 for x in rolls if x['correct']) / len(rolls)) if rolls else 0.0 - - -def bare_solve(rt: Runtime, records, rollouts: int = None) -> List[List[Dict[str, Any]]]: - """裸题重解 `rollouts` 次，每条记录返回一个 roll 列表（长度 = 实际采到的序列数）。 - - 返回**嵌套**列表而不是单个 roll：调用方靠 _pass_rate() 取连续值。 - 判分全部汇到一次 judge_seqs（它内部按 (task_id, code) 去重，相同代码只跑一次单测）。 - - ⭐ M==1 时必须降到 temperature=0.0（与原版一致）：单次采样就不需要多样性， - 带温度只会引入无意义的方差；max(1, ...) 防 rollouts 传 0 时采不到任何序列。 - """ - M = max(1, rollouts if rollouts is not None else EXEC_ROLLOUTS) - out = run_samples(rt.base_sampler, [direct_prompt(r['problem']) for r in records], - M, EXEC_MAX_TOKENS, BASE_SAMPLER_GPUS, - temperature=(0.0 if M == 1 else EXEC_TEMPERATURE), - top_p=(None if M == 1 else EXEC_TOP_P)) - pairs, spans = [], [] - for r, seqs in zip(records, out): - seqs = list(seqs or []) - spans.append(len(seqs)) - pairs.extend((s, r['reference_answer']) for s in seqs) - judged = judge_seqs(pairs) - rolls, i = [], 0 - for n in spans: - rolls.append(judged[i:i + n] if n else [empty_roll()]) - i += n - return rolls - - -# =========================================================================== -# 采集（与 e18_rejection_sft.collect_chunk 逐字一致） -# =========================================================================== -def _pick_trajectory(rolls: List[Dict[str, Any]]) -> str: - """从裸解的 rolls 里挑出最值得给 skill-gen 看的失败代码。 - - ⭐ 挑选而不是全给：BARE_ROLLOUTS=4 条里常有 2-3 条是**同一个错** - （judge_seqs 内部按 (task_id, code) 去重就是因为重复普遍），全给只会重复占预算。 - - 优先级：有报错信息 > 代码短。 - 为何偏好**短**代码：长代码往往是思维链泄到正文里的 no_code / import_or_syntax - 废文，信息密度低；短而完整的错解才能看出逻辑问题。同时也直接压低了长度风险。 - """ - bad = [x for x in (rolls or []) if not x.get('correct') and (x.get('code') or '').strip()] - if not bad: - return '' - bad.sort(key=lambda x: (0 if x.get('error') else 1, len(x.get('code') or ''))) - blocks = [format_trajectory(x.get('code'), x.get('error'), x.get('kind'), - max_chars=TRAJ_MAX_CHARS) - for x in bad[:max(1, TRAJ_N)]] - return '\n\n'.join(blocks) - - -def _gen_candidates(rt: Runtime, todo, n_skills: int) -> List[List[Dict[str, Any]]]: - """给 todo 里每道题生成 n_skills 个 skill 候选（只生成，不判分）。 - - todo 元素 = (record, rubric, base_rate, trajectory)；trajectory 在 USE_TRAJ=0 时恒为 ''。 - """ - sg = run_samples(rt.skill_sampler, - [skillgen_prompt(r['problem'], d, eval=False, trajectory=tj) - for r, d, _br, tj in todo], - n_skills, SKILL_MAX_TOKENS, SKILL_SAMPLER_GPUS, - temperature=SKILL_GEN_TEMPERATURE, top_p=SKILL_GEN_TOP_P, - top_k=SKILL_GEN_TOP_K) - out = [] - for seqs in sg: - cands = [] - for s in seqs or []: - resp = seq_text(s) - block = extract_skill(resp) - cands.append({'skills': block, 'response': resp, 'parseable': bool(block), - 'with_pass': None, 'kept': False, - 'skillgen_stop': getattr(s, 'stop_reason', None)}) - out.append(cands) - return out - - -def _judge_candidates(rt: Runtime, flat, rollouts: int) -> None: - """对 flat=[(record, cand), ...] 跑 `rollouts` 次重解并原地回写 with_pass。 - - ⭐ 原地累加而非覆盖：两阶段粗筛里同一个候选会被判两次（先 PROBE 后补足）， - 第二次必须把两次的样本**合并**算 pass_rate，否则前 PROBE_ROLLOUTS 次白扔。 - 用 _n_correct/_n_total 累计，with_pass 每次由累计值重算。 - """ - if not flat: - return - ws = run_samples(rt.base_sampler, - [skill_solve_prompt(r['problem'], c['skills']) for r, c in flat], - rollouts, EXEC_MAX_TOKENS, BASE_SAMPLER_GPUS, - temperature=EXEC_TEMPERATURE, top_p=EXEC_TOP_P) - pairs, spans = [], [] - for (r, _c), seqs in zip(flat, ws): - seqs = list(seqs or []) - spans.append(len(seqs)) - pairs.extend((s, r['reference_answer']) for s in seqs) - judged = judge_seqs(pairs) - i = 0 - for (_r, c), n in zip(flat, spans): - rr = judged[i:i + n] - i += n - c['_n_correct'] = c.get('_n_correct', 0) + sum(1 for x in rr if x['correct']) - c['_n_total'] = c.get('_n_total', 0) + n - c['with_pass'] = (c['_n_correct'] / c['_n_total']) if c['_n_total'] else 0.0 - c['n_rollouts'] = c['_n_total'] - if c.get('roll_kind') in (None, 'pass'): - c['roll_kind'] = (next((x['kind'] for x in rr if not x['correct']), - c.get('roll_kind') or 'pass') if rr else 'empty') - - -def _resolve_with_probe(rt: Runtime, todo, per_task, base_rates) -> None: - """两阶段 rollout：先粗筛 PROBE_ROLLOUTS 次，只把**可能是胜者**的候选补到 EXEC_ROLLOUTS。 - - 补足的判据是「粗筛并列最高」而不是「粗筛最高」：粗筛只有 2 次采样， - 并列极其常见（实测 51% 的题全同分），只补单个最高者会把真胜者漏掉。 - ⚠️ 只有 parseable 的候选参与；不可解析的候选 with_pass 保持 None，与历史行为一致。 - """ - alive = [(r, c) for (r, _d, _br, _tj), cands in zip(todo, per_task) - for c in cands if c.get('parseable')] - if not alive: - return - if not PROBE_ROLLOUTS or PROBE_ROLLOUTS >= EXEC_ROLLOUTS: - _judge_candidates(rt, alive, EXEC_ROLLOUTS) - return - _judge_candidates(rt, alive, PROBE_ROLLOUTS) - # 逐题挑「粗筛并列最高」者补足到 EXEC_ROLLOUTS - need = [] - for (r, _d, base_rate, _tj), cands in zip(todo, per_task): - ok = [c for c in cands if c.get('parseable')] - if not ok: - continue - top = max(c['with_pass'] for c in ok) - # 粗筛就已经够不到门槛的题：补足也不可能入池，直接省掉。 - if top + 1e-9 < base_rate + MIN_PASS_GAIN - (1.0 / max(1, EXEC_ROLLOUTS)): - continue - need.extend((r, c) for c in ok if c['with_pass'] >= top - 1e-9) - _judge_candidates(rt, need, EXEC_ROLLOUTS - PROBE_ROLLOUTS) - - -def collect_chunk(rt: Runtime, chunk, ci: int) -> Tuple[List[Dict[str, Any]], Dict[str, float]]: - """一个 chunk 的采集，返回 (胜者列表, 指标)。 - - 三层省算力（每层只放行需要下一层的题），全部可用环境变量关掉回到历史行为： - 1. SKIP_CEILING : base 打满的题不生成候选（实测入池 0 条，纯浪费） - 2. N_SKILLS_STAGE1: 先 2 个候选，全不达标才补到 N_SKILLS - 3. PROBE_ROLLOUTS: 每候选先 2 次，只对并列最高者补到 EXEC_ROLLOUTS - - 诊断（rubric）只对**做错的题**拉，且**每一种失败模式各诊一次后合并**（见 e18_multidiag）。 - 无诊断的题不丢弃：skillgen_prompt 内部会填兜底文案。 - """ - base_rolls = bare_solve(rt, chunk, rollouts=BARE_ROLLOUTS) - base_rates = [_pass_rate(rr) for rr in base_rolls] - base_acc = _mean(base_rates) - wrong = [(r, rr) for r, rr, rate in zip(chunk, base_rolls, base_rates) if rate < 1.0] - - before = rt.rubric_cache.stats.copy() - diags = rt.rubric_cache.diagnose_many(rt.checker, wrong) - rmetrics = multidiag_metrics(rt.rubric_cache.stats - before) - diag_by_id = {id(r): d for (r, _rr), d in zip(wrong, diags) if d} - n_rubric_missing = len(wrong) - len(diag_by_id) - n_multi = sum(1 for d in diag_by_id.values() if d.count('FAILURE ') > 1) - - # ⭐ todo 元素是四元组 (record, rubric, base_rate, trajectory)。 - # trajectory 只在 USE_TRAJ=1 时非空；关闭时恒为 '' -> skillgen_prompt 走原模板， - # 与历史 run 逐字一致。轨迹取自**裸解**的 rolls（就是当初拿去要诊断的那批）， - # 所以与 rubric 同源、描述的是同一次失败 —— 这正是 E22 当时做不到的（那次是重采的）。 - traj_by_id = ({id(r): _pick_trajectory(rr) for r, rr in zip(chunk, base_rolls)} - if USE_TRAJ else {}) - all_tasks = [(r, diag_by_id.get(id(r), ''), rate, traj_by_id.get(id(r), '')) - for r, rate in zip(chunk, base_rates)] - # ⭐ 天花板短路。n_skipped 单独记账，指标分母用 todo（非天花板题）—— - # 于是 baseline_accuracy / candidate_pass_rate / lift 三个指标的口径变了， - # **与 SKIP_CEILING=0 的历史 run 不可直接比较**，看趋势时注意这一点。 - if SKIP_CEILING: - todo = [t for t in all_tasks if t[2] + MIN_PASS_GAIN <= 1.0 + 1e-9] - else: - todo = all_tasks - n_skipped = len(all_tasks) - len(todo) - - per_task = _gen_candidates(rt, todo, min(N_SKILLS_STAGE1, N_SKILLS)) if todo else [] - if per_task: - _resolve_with_probe(rt, todo, per_task, base_rates) - - # ⭐ 阶段2：阶段1 全部不达标的题，再生成剩余候选。 - # 实测 39.7% 的题阶段1 就出胜者 -> 这批题省掉一半候选；6.1% 的题靠阶段2 救回。 - n_stage2_tasks = n_stage2_saved = 0 - remain = N_SKILLS - min(N_SKILLS_STAGE1, N_SKILLS) - if per_task and remain > 0: - idx2 = [i for i, ((_r, _d, br, _tj), cands) in enumerate(zip(todo, per_task)) - if not any(c.get('parseable') - and (c.get('with_pass') or 0.0) >= br + MIN_PASS_GAIN - 1e-9 - for c in cands)] - if idx2: - n_stage2_tasks = len(idx2) - todo2 = [todo[i] for i in idx2] - extra = _gen_candidates(rt, todo2, remain) - per_task2 = [[] for _ in todo] - for i, cands in zip(idx2, extra): - per_task2[i] = cands - _resolve_with_probe(rt, todo, per_task2, base_rates) - for i, cands in zip(idx2, extra): - per_task[i].extend(cands) - n_stage2_saved = sum( - 1 for i, cands in zip(idx2, extra) - if any(c.get('parseable') - and (c.get('with_pass') or 0.0) >= todo[i][2] + MIN_PASS_GAIN - 1e-9 - for c in cands)) - - accepted, sims = [], [] - n_pass_cands = n_survivors = 0 - n_acc_hard = n_acc_easy = 0 - gtot = {'improved': 0, 'tied': 0, 'degraded': 0} - gains = [] - for (r, d, base_rate, _tj), cands in zip(todo, per_task): - passers = [c for c in cands - if c.get('parseable') and (c.get('with_pass') or 0) >= base_rate] - n_pass_cands += len(passers) - for k, v in gain_stats(cands, base_rate).items(): - gtot[k] += v - best = select_winner(cands, d, r['reference_answer'], - skill_char_limit=SKILL_CHAR_LIMIT, - base_pass_rate=base_rate, min_pass_gain=MIN_PASS_GAIN) - if best is None: - continue - n_survivors += 1 - sims.append(best['rubric_similarity']) - gains.append(best['pass_gain']) - if base_rate >= 1.0: - n_acc_easy += 1 - else: - n_acc_hard += 1 - accepted.append({ - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'data_id': r.get('data_id', ''), 'skills': best['skills'], - 'response': f"<skills>\n{best['skills']}\n</skills>", - 'base_pass_rate': base_rate, 'with_pass_rate': best['with_pass'], - 'pass_gain': best['pass_gain'], 'gain_kind': best['gain_kind'], - 'rubric': d, 'chunk': ci, 'run': RUN_ID, - 'rubric_similarity': best['rubric_similarity'], - 'skill_chars': len(best['skills']), - 'n_candidates_passed': len(passers)}) - - # ⭐ 分母改成**实际生成的候选总数**：两阶段下每道题的候选数不再是定值 N_SKILLS - # （阶段1 就达标的题只有 N_SKILLS_STAGE1 个），写死 N_SKILLS 会把分母虚抬、 - # 让 candidate_pass_rate 和 lift 系统性偏低。 - n_cands_total = sum(len(cs) for cs in per_task) - _cand_pass = (n_pass_cands / n_cands_total) if n_cands_total else 0.0 - # 必须在选择循环之后：kept / with_pass / rubric_similarity 都是 select_winner 原地回写的。 - dump_candidates(todo, per_task, ci) - metrics = { - 'train/baseline_accuracy': base_acc, - 'train/accept_rate': (len(accepted) / len(todo)) if todo else 0.0, - 'train/candidate_pass_rate': _cand_pass, - 'train/lift': _cand_pass - base_acc, - 'train/selected_rubric_similarity': _mean(sims), - 'train/selected_skill_length_characters': _mean( - [float(s['skill_chars']) for s in accepted]), - 'signal/n_wrong': float(len(wrong)), - 'signal/n_rubric_missing': float(n_rubric_missing), - 'signal/n_multi_cause': float(n_multi), - 'signal/n_accepted': float(len(accepted)), - 'signal/n_accepted_hard': float(n_acc_hard), - 'signal/n_accepted_easy': float(n_acc_easy), - 'signal/n_ceiling': float(sum(1 for _r, _d, br, _tj in todo - if br + MIN_PASS_GAIN > 1.0 + 1e-9)), - 'signal/min_pass_gain': MIN_PASS_GAIN, - # 三层省算力的实测记账（用于事后核对真省了多少，而不是只信离线模拟） - 'saving/n_ceiling_skipped': float(n_skipped), - 'saving/n_stage2_tasks': float(n_stage2_tasks), - 'saving/n_stage2_rescued': float(n_stage2_saved), - 'saving/candidates_per_task': (n_cands_total / len(todo)) if todo else 0.0, - 'saving/exec_sequences': float(sum( - c.get('_n_total', 0) for cs in per_task for c in cs)), - # ⭐ 轨迹注入的实测记账。traj/rate 远低于 1 就说明很多题拿不到可用失败代码 - # （全对、或全是 no_code 空代码），此时该开关的实际覆盖面比以为的小。 - # traj/chars 盯长度风险：除以 2.6 就是多吃的 token 数。 - 'traj/enabled': float(USE_TRAJ), - 'traj/rate': (sum(1 for _r, _d, _br, tj in todo if tj) / len(todo)) if todo else 0.0, - 'traj/chars': _mean([float(len(tj)) for _r, _d, _br, tj in todo if tj]), - 'gain/improved_candidates': float(gtot['improved']), - 'gain/tied_candidates': float(gtot['tied']), - 'gain/degraded_candidates': float(gtot['degraded']), - 'gain/degrade_rate': (gtot['degraded'] / max(1, sum(gtot.values()))), - 'gain/improve_rate': (gtot['improved'] / max(1, sum(gtot.values()))), - 'gain/selected_pass_gain': _mean(gains), - } | rmetrics | class_metrics([d for _r, d, _br, _tj in todo if d]) - return accepted, metrics - - -def dump_dataset(accepted) -> None: - """胜者落盘 append-only 的 SFT 数据集（字段与在线 run 完全一致）。""" - if not accepted: - return - path = os.path.join(OUTPUT_DIR, 'e18_sft_dataset.jsonl') - with open(path, 'a', encoding='utf-8') as f: - for s in accepted: - f.write(json.dumps(s, ensure_ascii=False) + '\n') - - -def dump_candidates(todo, per_task, ci: int) -> None: - """**全部** skill 候选落盘（选上的、没选上的、甚至解不出 <skills> 的），append-only。 - - 为何必需：`e18_sft_dataset.jsonl` 只存胜者，而拒绝采样的全部信息量在「同一题的 N 条 - 候选之间的差异」里 —— 丢掉落选者就无法回答：胜者是真的更好，还是只是采样噪声 - （候选全部并列时靠拆平局选出来的）？也无法事后重算阀值：改 MIN_PASS_GAIN / - SKILL_CHAR_LIMIT 后想知道会多收多少条，必须有落选者的 with_pass 才能离线重放。 - - `kept` 区分选与未选：`select_winner` 对胜者**原地** 置 True（e18_select.py:107）， - 本函数必须在 select_winner 之后调用，否则全部候选都是 False。 - 同理 `with_pass` / `pass_gain` / `rubric_similarity` 也是选择阶段回写的。 - - 不存 `response` 全文（包含 think 链，体量是 skills 的十倍量级），只存抽取后的 - skills 与长度/截断信息；解析失败（parseable=False）时 skills 为空串，靠 skillgen_stop - 判断是撞预算截断还是真的不守格式。 - """ - path = os.path.join(OUTPUT_DIR, 'e18_candidates.jsonl') - with open(path, 'a', encoding='utf-8') as f: - for (r, d, base_rate, _tj), cands in zip(todo, per_task): - for j, c in enumerate(cands): - wp = c.get('with_pass') - f.write(json.dumps({ - 'chunk': ci, 'run': RUN_ID, - 'data_id': r.get('data_id', ''), - 'task_id': r['reference_answer'].get('task_id', ''), - 'cand_idx': j, - 'kept': bool(c.get('kept')), - 'parseable': bool(c.get('parseable')), - 'skills': c.get('skills', ''), - 'skill_chars': len(c.get('skills') or ''), - 'skillgen_stop': c.get('skillgen_stop'), - 'base_pass_rate': base_rate, - 'with_pass_rate': wp, - # 未参与重解（解析失败）时 with_pass 是 None，pass_gain 也留 None， - # 不能当 0 存 —— 否则离线统计会把「没跑」混成「跑了但零增益」。 - 'pass_gain': (None if wp is None else round(wp - base_rate, 6)), - 'gain_kind': c.get('gain_kind'), - 'n_rollouts': c.get('n_rollouts'), - 'roll_kind': c.get('roll_kind'), - 'rubric_similarity': c.get('rubric_similarity'), - 'rubric': d, - }, ensure_ascii=False) + '\n') - - -# =========================================================================== -# main -# =========================================================================== -def build_runtime(checker, rubric_cache) -> Runtime: - """两组卡：skill_sampler / base_sampler(executor)。不训练，所以没有 train 组。""" - twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, lazy_collect=False, groups=[ - DeviceGroup(name='skill_sampler', ranks=list(range(0, SKILL_SAMPLER_GPUS)), - device_type='GPU'), - DeviceGroup(name='base_sampler', ranks=list(range(SKILL_SAMPLER_GPUS, NUM_GPUS)), - device_type='GPU')]) - - def _sampler(group, world, enable_thinking): - s = vLLMSampler(model_id=MODEL_ID, - engine_args={'gpu_memory_utilization': GPU_MEM, - 'max_model_len': MAX_MODEL_LEN, - 'tensor_parallel_size': 1}, - device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), - remote_group=group) - s.set_template(Template, model_id=MODEL_ID, enable_thinking=enable_thinking, - max_length=MAX_MODEL_LEN) - return s - - # 与在线 run 同口径：skill_sampler 开 think（采集要多样性），executor 也开 think。 - skill_sampler = _sampler('skill_sampler', SKILL_SAMPLER_GPUS, enable_thinking=True) - base_sampler = _sampler('base_sampler', BASE_SAMPLER_GPUS, enable_thinking=True) - return Runtime(skill_sampler=skill_sampler, base_sampler=base_sampler, - checker=checker, rubric_cache=rubric_cache) - - -def count_existing_samples() -> int: - """续跑：已有胜者条数（= e18_sft_dataset.jsonl 的有效行数）。同样要在归档前调。 - - 这里才该用 sft_dataset 而不是 candidates：它要回答的是「已经凑了多少条胜者」， - 用来接着比 TARGET_SAMPLES；而 resume_done_ids() 要回答的是「哪些题不用再跑」。 - 两个问题不同源，切勿合并。 - """ - path = os.path.join(OUTPUT_DIR, 'e18_sft_dataset.jsonl') - if not os.path.exists(path): - return 0 - n = 0 - with open(path, 'r', encoding='utf-8') as f: - for line in f: - if line.strip(): - n += 1 - return n - - -def resume_done_ids() -> set: - """续跑：读出已经跑过的 data_id。必须在 archive_output_dir() **之前**调用。 - - ⭐ 读 `e18_candidates.jsonl` 而不是 `e18_sft_dataset.jsonl`：后者只有**胜者**（实测 - accept_rate 约 22%），拿它去重会把剩下 78%「跑过但没能入池」的题当成未跑、 - 下次重新烧一遍 GPU。候选文件是全量落盘的，覆盖面才完整。 - - 容错：进程被 kill 时最后一行可能是写一半的残行，json.loads 会抛异常 -> - 逐行 try 跳过（丢掉一道题的去重信息只是多跑一题，而整个函数抛异常会直接弄挂启动）。 - """ - path = os.path.join(OUTPUT_DIR, 'e18_candidates.jsonl') - if not os.path.exists(path): - return set() - done = set() - bad = 0 - with open(path, 'r', encoding='utf-8') as f: - for line in f: - line = line.strip() - if not line: - continue - try: - d = json.loads(line)['data_id'] - except Exception: - bad += 1 - continue - if d: - done.add(str(d)) - logger.info(f'[resume] 已跑过 {len(done)} 题' - + (f'（跳过 {bad} 行残缺/半行）' if bad else '')) - return done - - -def archive_output_dir(carry_data: bool = False) -> None: - """启动时把已存在的 OUTPUT_DIR 整个 mv 走，保证本 run 写入空目录。 - - 为何必需：`e18_sft_dataset.jsonl` / `collect_log.jsonl` 都是 `open(..., 'a')` 追写。 - 没有这一步时，重启一次就把新旧 run 的样本焊在同一个文件里，而且不报错。 - 用 mv 而不是删：旧 run 的样本是可复用的分析素材。 - 沙箱自检缓存（kod_broken_tasks.json）会搬回新目录 —— 那是纯函数结果，重跑一次很贵。 - - carry_data=True（续跑）时额外把三个产物 jsonl 也复制回新目录，于是新 run 接着往后面 - 追写、总数连续。用 copy2 而不是 move：归档副本保留完整快照，万一续跑又崩了还能回溯。 - 注意：续跑后同一份数据里会存在多个 `run` 值，离线分析要按 run 字段分组看。 - """ - if not os.path.isdir(OUTPUT_DIR) or not os.listdir(OUTPUT_DIR): - os.makedirs(OUTPUT_DIR, exist_ok=True) - return - stamp = time.strftime('%m%d-%H%M%S', time.localtime(os.path.getmtime(OUTPUT_DIR))) - dst = f'{OUTPUT_DIR}.bak-{stamp}' - i = 1 - while os.path.exists(dst): - dst = f'{OUTPUT_DIR}.bak-{stamp}-{i}' - i += 1 - shutil.move(OUTPUT_DIR, dst) - os.makedirs(OUTPUT_DIR, exist_ok=True) - logger.info(f'[init] 旧输出目录已归档 -> {dst}') - cache = os.path.join(dst, 'kod_broken_tasks.json') - if os.path.exists(cache): - shutil.copy2(cache, os.path.join(OUTPUT_DIR, 'kod_broken_tasks.json')) - logger.info('[init] 沙箱自检缓存已搬回新目录（避免重跑）') - if carry_data: - for name in ('e18_sft_dataset.jsonl', 'e18_candidates.jsonl', 'collect_log.jsonl'): - src = os.path.join(dst, name) - if os.path.exists(src): - shutil.copy2(src, os.path.join(OUTPUT_DIR, name)) - logger.info('[resume] 旧产物已复制回新目录，本 run 接着追写') - - -def main(): - t_start = time.time() - # ⭐ 顺序强制：读已采 id 必须在归档**之前**，否则 OUTPUT_DIR 已经被 mv 走、读到空集。 - done_ids = resume_done_ids() if KOD_RESUME else set() - n_done = count_existing_samples() if KOD_RESUME else 0 - archive_output_dir(carry_data=KOD_RESUME) - checker = build_checker() - # eval_size=0：不训练就没有 eval 的意义，题全进采集池。 - train_dataset, _ = load_records(SEED, 0, OUTPUT_DIR) - # ⭐ 分片必须在 resume 过滤**之前**做，且用 crc32(data_id) 而不是下标： - # 两台机器无共享存储，只能靠纯函数保证不重叠。zlib.crc32 跨进程/跨平台稳定 - # （而 hash() 受 PYTHONHASHSEED 影响，每次启动都不同 —— 用它会造成重叠+遗漏）。 - if SHARD_N > 1: - before = len(train_dataset) - train_dataset.filter( - lambda r: zlib.crc32(str(r['data_id']).encode()) % SHARD_N == SHARD_ID) - logger.info(f'[shard] {SHARD_ID}/{SHARD_N}: 题池 {before} -> {len(train_dataset)}') - if done_ids: - before = len(train_dataset) - train_dataset.filter(lambda r: r['data_id'] not in done_ids) - logger.info(f'[resume] 题池剔除已采 {before - len(train_dataset)} 题 -> 剩 {len(train_dataset)}') - logger.info(f'[data] 采集池 ={len(train_dataset)} 题') - rt = build_runtime(checker, MultiDiagCache()) - logger.info(f'E18-collect start: chunk={CHUNK_SIZE} n_skills={N_SKILLS} ' - f'bare_rollouts={BARE_ROLLOUTS} exec_rollouts={EXEC_ROLLOUTS} ' - f'min_pass_gain={MIN_PASS_GAIN:.3g} ' - f'use_traj={USE_TRAJ}(n={TRAJ_N},max_chars={TRAJ_MAX_CHARS}) ' - f'shard={SHARD_ID}/{SHARD_N} ' - f'target={TARGET_SAMPLES} gpus={SKILL_SAMPLER_GPUS}+{BASE_SAMPLER_GPUS} ' - f'resume={int(KOD_RESUME)} n_done={n_done} ' - f'output={OUTPUT_DIR}') - - # ⭐ n_total 从已有数量起算，不是 0：它是 TARGET_SAMPLES 的比较对象，从 0 起算会变成 - # 「再采 TARGET_SAMPLES 条」而不是「凑到 TARGET_SAMPLES 条」。 - n_total, ci = n_done, 0 - log_path = os.path.join(OUTPUT_DIR, 'collect_log.jsonl') - with open(log_path, 'a', encoding='utf-8') as log_fh: - loader = DataLoader(dataset=train_dataset, batch_size=CHUNK_SIZE, num_workers=0, - shuffle=True, drop_last=False, - generator=torch.Generator().manual_seed(SEED)) - for chunk in loader: - if TARGET_SAMPLES and n_total >= TARGET_SAMPLES: - break - if MAX_CHUNKS and ci >= MAX_CHUNKS: - break - t0 = time.time() - accepted, metrics = collect_chunk(rt, chunk, ci) - dump_dataset(accepted) - n_total += len(accepted) - row = {'chunk': ci, 'run': RUN_ID, 'seconds': round(time.time() - t0, 1), - 'n_collected': float(n_total), **metrics} - log_fh.write(json.dumps(row, ensure_ascii=False) + '\n') - log_fh.flush() - logger.info(f'[c{ci}] collected={n_total}' - + (f'/{TARGET_SAMPLES}' if TARGET_SAMPLES else '') - + ' ' + ' '.join(f'{k}={v:.4g}' for k, v in row.items() - if isinstance(v, float))) - ci += 1 - logger.info(f'[完成] 共 {n_total} 条，{ci} 个 chunk，' - f'耗时 {(time.time() - t_start) / 60:.1f} 分钟 -> {OUTPUT_DIR}') - - -if __name__ == '__main__': - main() diff --git a/cookbook/human_e18/e18_kodcode.py b/cookbook/human_e18/e18_kodcode.py deleted file mode 100644 index bc753e500..000000000 --- a/cookbook/human_e18/e18_kodcode.py +++ /dev/null @@ -1,404 +0,0 @@ -# -*- coding: utf-8 -*- -"""KodCode-V1 数据域适配：与 `e23_bcb.py` **同契约**的加载 + 沙箱判分。 - -为什么另起一个文件而不改 e23_bcb：BCB 与 KodCode 的单测框架不同（unittest vs pytest）、 -题面/入口点的来源字段也不同，但**对上层的接口必须逐字一致** —— `e18_rejection_sft.py` -只认 `load_records / judge_seqs / empty_roll / run_tests` 这组签名和 `{'data_id', -'problem', 'reference_answer'}` 这个记录形状。保持接口一致，换域时上层零改动。 - -与 e23_bcb 的对齐点（改任何一处都会让两域的 pass_rate 不可比）： -* `run_tests` 返回 `{'passed', 'kind', 'error'}`，`kind` 取值集合完全相同： - `pass / no_code / no_entry / timeout / assertion / exception / import_or_syntax`。 - `e18_multidiag._signature` 拿 kind 做失败签名，取值不一致会让诊断缓存串味。 -* `judge_seqs` 同 `(task_id, code)` 只判一次、线程池并发、返回 roll 的字段集相同。 -* `_trim_err` 只保留失败测试名与异常行，并把随机临时目录名归一化成 `<sandbox>` —— - 否则同一个失败在两次运行里字符串不同，`_signature` 会算出两个签名、缓存永远不命中。 -* 参考解答跑不过自己单测的题一律剔除（BCB 实测 ~7.5%，KodCode 实测 ~5%）， - 自检结果落盘缓存。这类题不是模型的错，留着会把 base_pass_rate 永久压低。 - -KodCode 特有的两点： -1. 单测是 pytest 风格且 **181/200 靠 `from solution import X`** 取被测函数，所以沙箱里必须 - 把提交代码写成 `solution.py`（而不是 BCB 那样把代码和 test 拼进同一个文件）。 -2. 自带 `gpt_pass_percentage`（教师多次尝试的通过率），是**免费的先验难度**。E18 原先要靠 - 跑 8 次 bare rollout 才能找出难题，这里可以直接按阈值筛，省掉这部分 GPU。 -""" -import ast as _ast -import json -import os -import re -import shutil -import subprocess -import sys -import tempfile -from concurrent.futures import ThreadPoolExecutor -from typing import Any, Dict, List, Optional, Tuple - -from twinkle import get_logger -from twinkle.dataset import Dataset, DatasetMeta - -logger = get_logger() - -# ========== 配置 ========== -KOD_DATASET = os.environ.get('KOD_DATASET', 'ms://AI-ModelScope/KodCode-V1') -KOD_SUBSET = os.environ.get('KOD_SUBSET', 'default') -KOD_SPLIT = os.environ.get('KOD_SPLIT', 'train') -# 难度窗口：只留「教师也常做错、但并非无解」的题。 -# 上界 0.3 -> executor 大概率失败（有诊断可采）；下界 >0 -> 排除疑似不可解。 -KOD_MAX_PASS_PCT = float(os.environ.get('KOD_MAX_PASS_PCT', 0.3)) -KOD_MIN_PASS_PCT = float(os.environ.get('KOD_MIN_PASS_PCT', 0.0)) -KOD_MAX_TASKS = int(os.environ.get('KOD_MAX_TASKS', 0)) # 0 = 不截断 -# ⭐ 默认关沙箱自检：全量 73747 题跑一遍参考解答要 ~30 小时（子进程）。 -# 关掉的代价：参考解答自己都跑不过单测的坏题（实测 ~12.5%）会留在题池里， -# 但它们在采集时会自然显形 —— base_pass_rate 恒为 0、且任何 skill 都拿不到 -# +MIN_PASS_GAIN，于是 select_winner 返回 None、不入池。只浪费 rollout，不污染数据集。 -KOD_SELFCHECK = os.environ.get('KOD_SELFCHECK', '0') == '1' -# ⭐ TEST_WORKERS 默认 96，而不是继承 BCB 的 24：判分是子进程，与 GPU 采样**串行**， -# 每 chunk 要跑 CHUNK_SIZE*(BARE_ROLLOUTS + N_SKILLS*EXEC_ROLLOUTS) ≈ 2300 次，并发不够就直接拆 GPU 空转。 -# BCB 用 24 是因为它的单测要 import pandas/sklearn/matplotlib（单次 1-3s、内存大）； -# KodCode 是纯算法题，单测 0.05-0.3s、几乎不导包，可以开得高得多。 -# 上限卡在 min(96, 核数一半)：留余量给 vLLM 的调度/集合线程，别把宿主打满反而拖慢采样。 -TEST_WORKERS = int(os.environ.get( - 'TEST_WORKERS', max(24, min(96, (os.cpu_count() or 24) // 2)))) -TEST_TIMEOUT = int(os.environ.get('TEST_TIMEOUT', 60)) - -_FENCE_RE = re.compile(r'```(?:python|py)?\s*\n(.*?)```', re.S) -_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') - - -# ========== 文本处理（与 e23_bcb 逐字一致） ========== -def after_think(text: str) -> str: - """只取 </think> 之后的正文；没有闭合标签就原样返回。""" - idx = text.rfind('</think>') - return text[idx + len('</think>'):] if idx >= 0 else text - - -def clean_text(decoded: Optional[str]) -> str: - return _SPECIAL_TOKEN_RE.sub('', decoded or '').strip() - - -def extract_code(text: str) -> str: - """取最后一个完整的 ``` 代码块；没有围栏时退化成整段正文。 - - 取**最后一个**而不是第一个：模型常先给一版草稿再给最终版，最后一个才是它的结论。 - """ - body = after_think(text) - blocks = _FENCE_RE.findall(body) - if blocks: - return blocks[-1].strip() - # 没有围栏：可能是 nothink 直出代码。剔掉明显的自然语言行后返回。 - return body.strip() - - -def extract_skill(text: str) -> str: - """取 <skills>...</skills> 里的内容；没有标签时返回空串（视为格式失败）。""" - body = after_think(text) - m = re.search(r'<skills>(.*?)</skills>', body, re.S | re.IGNORECASE) - return m.group(1).strip() if m else '' - - -# ========== 沙箱判分 ========== -# ⭐ 纯 pytest 驱动，但用插件把 (n_tests, n_fail, n_err) 拿回来 —— 与 e23_bcb 的 -# `__BCB__ n f e` 输出契约同形，好让 kind 的判定规则完全一致。 -# 只统计 when=='call'：setup/teardown 阶段的失败算 error（多半是 import 不了 solution）。 -# -# ⭐ 断言 vs 异常的区分必须走 `report.longrepr.reprcrash.message`，**不能**用 -# `'AssertionError' in str(longrepr)`：pytest 默认开断言重写（assertion rewriting）， -# 失败摘要长这样 `E assert -1 == 3`，整段里根本没有 "AssertionError" 这个词， -# 于是所有断言失败都会被误判成 exception。实测踩过：断言不符返回了 kind='exception'。 -# kind 错了会让 `e18_multidiag._signature` 的失败签名串味、rubric 缓存失效。 -_RUNNER = r""" -import sys, pytest - - -class _Collect: - def __init__(self): - self.n_tests = self.n_fail = self.n_err = 0 - - @staticmethod - def _is_assertion(report): - crash = getattr(getattr(report, 'longrepr', None), 'reprcrash', None) - msg = getattr(crash, 'message', '') or '' - # 断言重写后首行是 "assert ..."；未重写时是 "AssertionError: ..."。 - return msg.startswith('assert') or msg.startswith('AssertionError') - - def pytest_runtest_logreport(self, report): - if report.when == 'call': - self.n_tests += 1 - if report.failed: - if self._is_assertion(report): - self.n_fail += 1 - else: - self.n_err += 1 - elif report.failed: - self.n_err += 1 - - -c = _Collect() -rc = pytest.main(['-q', '--no-header', '-p', 'no:cacheprovider', - '--tb=short', 'test_solution.py'], plugins=[c]) -print('__KOD__', c.n_tests, c.n_fail, c.n_err) -sys.exit(0 if int(rc) == 0 else 1) -""" - - -def _trim_err(err: str, limit: int = 1600) -> str: - """保留失败测试名与异常行，砍掉冗长 traceback 帧 —— 这是喂给 rubric 的客观证据。 - 随机临时目录名换成 <sandbox>，否则同一个失败在两次运行里看起来不一样 - （`e18_multidiag._signature` 会因此算出不同签名，诊断缓存永久不命中）。""" - err = re.sub(r'/tmp/kod_[A-Za-z0-9_]+', '<sandbox>', err or '') - lines = [ln for ln in err.splitlines() if ln.strip()] - keep = [ln for ln in lines - if ln.startswith(('FAILED', 'FAIL:', 'ERROR:', 'AssertionError', 'Traceback', 'E ')) - or re.match(r'^\w*(Error|Exception|Warning)\b', ln.strip()) - or ', in ' in ln] - return '\n'.join(keep or lines[-25:])[-limit:] - - -def run_tests(code: str, payload: Dict[str, Any], timeout: int = TEST_TIMEOUT) -> Dict[str, Any]: - """子进程里跑「提交代码(solution.py) + 官方 test(test_solution.py)」。 - - -> {'passed', 'kind', 'error'}，kind 取值与 e23_bcb.run_tests 完全一致。 - - 与 BCB 的唯一实质差异：代码单独落成 `solution.py`，因为 KodCode 的单测靠 - `from solution import X` 取被测函数，拼进同一个文件会 ImportError。 - """ - if not code.strip(): - return {'passed': False, 'kind': 'no_code', 'error': 'no parseable code block'} - entry = payload.get('entry_point') or '' - if entry and entry not in code: - return {'passed': False, 'kind': 'no_entry', - 'error': f'function {entry} is not defined in the submitted code'} - tmp = tempfile.mkdtemp(prefix='kod_') - try: - with open(os.path.join(tmp, 'solution.py'), 'w', encoding='utf-8') as f: - f.write(code) - with open(os.path.join(tmp, 'test_solution.py'), 'w', encoding='utf-8') as f: - f.write(payload['test']) - with open(os.path.join(tmp, '_run.py'), 'w', encoding='utf-8') as f: - f.write(_RUNNER) - env = dict(os.environ, MPLBACKEND='Agg', PYTHONHASHSEED='0', OMP_NUM_THREADS='1', - MKL_NUM_THREADS='1', TOKENIZERS_PARALLELISM='false') - env.pop('CUDA_VISIBLE_DEVICES', None) - # cwd=tmp 让 `from solution import X` 能找到同目录的 solution.py。 - try: - p = subprocess.run([sys.executable, '_run.py'], cwd=tmp, env=env, timeout=timeout, - capture_output=True, text=True, errors='replace') - except subprocess.TimeoutExpired: - return {'passed': False, 'kind': 'timeout', - 'error': f'the tests did not finish within {timeout}s'} - n_tests = n_fail = n_err = 0 - for line in (p.stdout or '').splitlines(): - if line.startswith('__KOD__'): - _, a, b, c = line.split() - n_tests, n_fail, n_err = int(a), int(b), int(c) - if p.returncode == 0 and n_tests > 0: - return {'passed': True, 'kind': 'pass', 'error': ''} - kind = 'assertion' if n_fail else ('exception' if n_err else 'import_or_syntax') - merged = ((p.stdout or '') + '\n' + (p.stderr or '')).replace(tmp, '<sandbox>') - return {'passed': False, 'kind': kind, 'error': _trim_err(merged)} - finally: - shutil.rmtree(tmp, ignore_errors=True) - - -def empty_roll() -> Dict[str, Any]: - return {'correct': False, 'stop_reason': 'empty', 'gen_tokens': 0, 'text': '', 'code': '', - 'kind': 'no_code', 'error': ''} - - -def judge_seqs(pairs: List[Tuple[Any, Dict[str, Any]]]) -> List[Dict[str, Any]]: - """[(采样 sequence 或 None, payload)] -> rolls。所有判分都汇合到这里。 - - 必须批量：单测是子进程，一个 chunk 几百次判分串行会比同 chunk 的 GPU 时间还长一个量级。 - 同 (task_id, code) 只跑一次 —— T=0 的 executor 经常对同一题产出逐字相同的代码。 - """ - rolls: List[Dict[str, Any]] = [] - keys: List[Optional[Tuple[str, str]]] = [] - jobs: Dict[Tuple[str, str], Dict[str, Any]] = {} - for seq, payload in pairs: - if seq is None: - rolls.append(empty_roll()) - keys.append(None) - continue - text = clean_text(getattr(seq, 'decoded', '') or '') - code = extract_code(text) - key = (payload['task_id'], code) - rolls.append({'correct': False, 'stop_reason': getattr(seq, 'stop_reason', None), - 'gen_tokens': len(getattr(seq, 'tokens', None) or []), - 'text': text, 'code': code, 'kind': None, 'error': ''}) - keys.append(key) - jobs.setdefault(key, payload) - if jobs: - todo = list(jobs) - with ThreadPoolExecutor(max_workers=max(1, min(TEST_WORKERS, len(todo)))) as ex: - verdicts = dict(zip(todo, ex.map(lambda k: run_tests(k[1], jobs[k]), todo))) - for roll, key in zip(rolls, keys): - v = verdicts.get(key) if key is not None else None - if v is not None: - roll['correct'] = bool(v['passed']) - roll['kind'], roll['error'] = v['kind'], v['error'] - return rolls - - -# ========== 数据 ========== -# reference_answer 的字段集：判分需要的一切。与 e23_bcb 的 _PAYLOAD_KEYS 同名对齐， -# 上层（e18_select / e18_multidiag / dump_dataset）拿到的 key 才一致。 -_PAYLOAD_KEYS = ('task_id', 'entry_point', 'test', 'code_prompt', 'doc_struct', - 'canonical_solution') - - -def _entry_point(row: Dict[str, Any]) -> str: - """从 test_info 拿被测函数名。拿不到就回落到 test 里的 `from solution import X`。""" - ti = row.get('test_info') - if ti is not None: - try: - items = list(ti) if not isinstance(ti, str) else _ast.literal_eval(ti) - for it in items: - name = (it or {}).get('function_name') - if name: - return str(name) - except Exception: - pass - m = re.search(r'from\s+solution\s+import\s+([A-Za-z_]\w*)', row.get('test') or '') - return m.group(1) if m else '' - - -def _code_prompt(row: Dict[str, Any]) -> str: - """函数签名，用于给 executor 固定入口点（对齐 BCB 的 code_prompt 语义）。""" - ti = row.get('test_info') - if ti is not None: - try: - items = list(ti) if not isinstance(ti, str) else _ast.literal_eval(ti) - for it in items: - decl = (it or {}).get('function_declaration') - if decl: - return str(decl) - except Exception: - pass - return '' - - -def _usable(row: Dict[str, Any]) -> bool: - """能进题池的最低门槛。 - - 要求 test 通过 `from solution import` 取函数：11.7% 的题直接裸调函数名， - 在「代码写进 solution.py」的沙箱布局下必然 NameError —— 那是 harness 不兼容， - 不是模型的错，留着会把 base_pass_rate 永久压低。 - """ - test = row.get('test') or '' - if 'def test_' not in test: - return False - if not re.search(r'from\s+solution\s+import|import\s+solution\b', test): - return False - return bool((row.get('solution') or '').strip()) and bool(_entry_point(row)) - - -# ⭐ 题面尾部必须追加函数签名：实测只有 **8%** 的 KodCode question 提到了被测函数名， -# 而单测靠 `from solution import <name>` 取函数。不补签名的后果：executor 把函数叫成 -# 任何名字都算错，92% 的题无论 skill 好坏都是 0 分 —— pass_rate 全城 0、整个采集废掉。 -# BCB 不需要这一步是因为它的 instruct_prompt 自带 `def task_func(...)` 骨架。 -_SIG_HINT = ('\n\nYou should write self-contained code starting with:\n```\n{decl}\n```') - - -def _problem_text(row: Dict[str, Any]) -> str: - """题面 = question + 函数签名（签名已在题面里就不重复追加）。""" - q = row.get('question') or '' - decl = _code_prompt(row) - if not decl: - return q - if decl.strip() in q: - return q - return q + _SIG_HINT.format(decl=decl.strip()) - - -def _to_record(batch: Dict[str, List]) -> Dict[str, List]: - """原始 KodCode 列 -> {'data_id', 'problem', 'reference_answer'}。 - - Dataset.map 强制 batched=True，所以这里收发的都是列式 batch。 - """ - n = len(batch['question_id']) - rows = [{k: batch[k][i] for k in batch} for i in range(n)] - return { - 'data_id': [str(r['question_id']) for r in rows], - 'problem': [_problem_text(r) for r in rows], - 'reference_answer': [{ - 'task_id': str(r['question_id']), - 'entry_point': _entry_point(r), - 'test': r['test'], - 'code_prompt': _code_prompt(r), - 'doc_struct': '', - 'canonical_solution': r['solution'], - # 教师先验难度：保留下来供离线分析（不参与判分）。 - 'gpt_pass_percentage': float(r.get('gpt_pass_percentage') or 0.0), - 'gpt_difficulty': r.get('gpt_difficulty') or '', - } for r in rows], - } - - -def _broken_tasks(ds: Dataset, output_dir: str) -> set: - """参考解答跑不过自己的单测 = 数据缺陷或沙箱不可判定，不是模型的错（实测约 5%）。 - 自检一次后落盘缓存，题数不变则复用。必须在 map 之后调用（读的是 reference_answer）。""" - path = os.path.join(output_dir, 'kod_broken_tasks.json') - if os.path.exists(path): - try: - with open(path, encoding='utf-8') as f: - c = json.load(f) - if int(c.get('n_tasks', -1)) == len(ds): - logger.info(f'[data] 复用沙箱自检缓存：剔除 {len(c["broken"])} 道') - return set(c['broken']) - except Exception as exc: - logger.warning(f'[data] 读取 {path} 失败（{exc}），重跑自检') - logger.info(f'[data] 沙箱自检：{len(ds)} 道题跑参考解答（一次性，之后走缓存）…') - rows = [ds[i] for i in range(len(ds))] - jobs = [(r['reference_answer']['canonical_solution'], r['reference_answer']) for r in rows] - with ThreadPoolExecutor(max_workers=max(1, min(TEST_WORKERS, len(jobs)))) as ex: - vers = list(ex.map(lambda p: run_tests(p[0], p[1]), jobs)) - broken = {r['data_id'] for r, v in zip(rows, vers) if not v['passed']} - with open(path, 'w', encoding='utf-8') as f: - json.dump({'n_tasks': len(ds), 'broken': sorted(broken)}, f, indent=1) - logger.info(f'[data] 自检完成：剔除 {len(broken)}/{len(rows)} ' - f'({100.0 * len(broken) / max(1, len(rows)):.1f}%)') - return broken - - -def load_records(seed: int, eval_size: int, - output_dir: str) -> Tuple[Dataset, List[Dict[str, Any]]]: - """-> (train_dataset, eval_records)，每条记录是 {'data_id', 'problem', 'reference_answer'}。 - - 签名与 `e23_bcb.load_records` 逐字一致，上层可直接换 import。 - - 与 BCB 的差异：KodCode 自带 `gpt_pass_percentage`，所以**先按难度窗口过滤**再自检 —— - 自检要跑一遍全部参考解答（子进程，很贵），先筛掉容易题能省掉大部分开销。 - """ - ds = Dataset(DatasetMeta(KOD_DATASET, subset_name=KOD_SUBSET, split=KOD_SPLIT)) - n_raw = len(ds) - - ds.filter(lambda r: KOD_MIN_PASS_PCT < float(r.get('gpt_pass_percentage') or 0.0) - <= KOD_MAX_PASS_PCT) - n_hard = len(ds) - ds.filter(_usable) - logger.info(f'[data] KodCode: 全集 {n_raw}，难度窗口 ' - f'({KOD_MIN_PASS_PCT}, {KOD_MAX_PASS_PCT}] 保留 {n_hard}、' - f'harness 不兼容剔除 {n_hard - len(ds)} -> {len(ds)}') - if KOD_MAX_TASKS and len(ds) > KOD_MAX_TASKS: - # ⭐ 必须用 filter 而不是 `ds.dataset = ds.dataset.select(...)`：Dataset.map 内部读的是 - # `self.datasets`（未截断的副本）并回写 `self.dataset`，直接赋值 self.dataset 会在 - # 下一句 map 里被静默覆盖 —— 实测踩过：截断到 40 题后自检仍在跑 73747 题。 - keep = set(ds.dataset.shuffle(seed=seed)['question_id'][:KOD_MAX_TASKS]) - ds.filter(lambda r: r['question_id'] in keep) - logger.info(f'[data] KOD_MAX_TASKS 截断 -> {len(ds)}') - - ds.map(_to_record, remove_columns=ds.dataset.column_names) - - broken = _broken_tasks(ds, output_dir) if KOD_SELFCHECK else set() - if broken: - ds.filter(lambda r: r['data_id'] not in broken) - logger.info(f'[data] 剔除参考解答自己跑不过单测的题 {len(broken)} 道 -> 可用 {len(ds)}') - elif not KOD_SELFCHECK: - logger.info(f'[data] 跳过沙箱自检（KOD_SELFCHECK=0），题池 {len(ds)}；' - f'坏题会在采集时因 base_pass_rate=0 自然不入池') - - shuffled = ds.dataset.shuffle(seed=seed) - n_eval = min(eval_size, len(shuffled)) if eval_size > 0 else 0 - eval_records = list(shuffled.select(range(n_eval))) - train_dataset = Dataset(DatasetMeta(data=shuffled.select(range(n_eval, len(shuffled))))) - return train_dataset, eval_records diff --git a/cookbook/human_e18/e18_multidiag.py b/cookbook/human_e18/e18_multidiag.py deleted file mode 100644 index d1166a705..000000000 --- a/cookbook/human_e18/e18_multidiag.py +++ /dev/null @@ -1,245 +0,0 @@ -"""E18 的多轨迹诊断：给**每一条失败的 rollout** 各诊一次，再合并成一份 rubric。 - -为什么需要它（E23 的 `RubricCache.get_or_diagnose` 不能直接复用）： - 1. 它的缓存键是 `RUBRIC_VERSION + _PROMPT_KEY + data_id`，**不含失败轨迹内容**。同一题调 - N 次会全部命中第一次的结果 —— 想「每条 rollout 判一次」在那个键下是做不到的。 - 2. 它一题只产出一个决定性根因，且 `_format` 把 secondary 渲染成 - `ALSO OFF (do not write about these)`，**主动禁止** skill-gen 覆盖第二处。 - 实测 E23 有 73% 的零 reward 组正是「修对第一处、挂在第二处」。 - -本模块只做加法，不改 `e23_rubric.py`（它是 E18/E23 共用、必须逐字同源的判分/诊断层）： -复用其 `diag_query` / `diag_segment` / `_validate` / `_CLASS_SHORT`，但换一个**按失败内容** -分桶的缓存键，并自己渲染合并文本。 - -⭐ 去重按 `(kind, 报错签名)` 而不是按 rollout 逐条：8 次采样里常有 5 次是同一个 assertion， -逐条诊断纯属浪费 API。同一签名只诊一次，但记下它出现了几次（`n_seen`），合并时按频次排序 —— -出现 5 次的根因显然比只出现 1 次的更该先修。 -""" -import hashlib -import json -import os -import re -import threading -from concurrent.futures import ThreadPoolExecutor -from typing import Any, Dict, List, Optional, Tuple - -from twinkle import get_logger - -import e23_rubric as R - -logger = get_logger() - -# 每题最多诊断几种**不同**的失败签名。3 与 e23_rubric 的 MAX_INDEPENDENT_CAUSES 一致： -# 超过 3 个独立根因的题，教师侧本来就判为不可救（诊断会互相矛盾，skill 也写不下）。 -MAX_DIAG_PER_TASK = int(os.environ.get('MAX_DIAG_PER_TASK', 3)) -# 一个签名至少要出现几次才值得诊断。默认 1 = 全诊；设 2 可以过滤掉只出现一次的偶发错误。 -MIN_SIGNATURE_COUNT = int(os.environ.get('MIN_SIGNATURE_COUNT', 1)) -# 并行度按**题**算：24 个线程各认领一道题，题内的多个签名仍串行，所以同时在飞的 HTTP -# request 就是 24。默认从 8 提到 24：纯 API 等待不占 GPU也不占 CPU，8 并行下一个 -# 64 题 chunk 的诊断阶段实测要 ~26 分钟，而这段时间 8 张卡全部空转。 -# 上限受教师侧限流约束，碰到 429 就把这个值调回。 -DIAG_WORKERS = int(os.environ.get('RUBRIC_WORKERS', 24)) - -# 失败签名：只取**结构化字段**（哪个测试挂了 + 什么异常类），不用报错正文。 -# -# ⭐ 为何不对报错正文做归一化（前一版的做法）：那靠的是「把数字/路径/引号内容逐个替成 -# 占位符」，而异常消息里的变量形式永远枚不完（裸数值、列宽对齐的空白、repr 片段……）； -# 漏一个，同一个 bug 就被当成 N 个不同根因，白花 N 倍教师 API 且 N 条诊断在说同一件事。 -# 现在只依赖两个结构化信号，行为可预测，不随报错排版变化。 -# -# 两个信号都来自 unittest 的固定输出格式，且 e23_bcb._trim_err 保证保留（它只留 -# FAIL:/ERROR:/Traceback/异常类名/带 ', in ' 的帧行）： -# * 失败的测试方法名 —— 区分「同一异常但挂在不同测试上」（那是不同根因）； -# * 异常类名 —— 区分 KeyError / AttributeError / AssertionError。 -# 只用 kind + 异常类会把前者错误合并，所以两个都要。 -_TEST_RE = re.compile(r'^(?:FAIL|ERROR):\s*(\w+)', re.M) -_EXC_RE = re.compile(r'^([A-Za-z_][\w.]*(?:Error|Exception|Warning))\b', re.M) - - -def _signature(roll: Dict[str, Any]) -> str: - """失败轨迹 -> 稳定签名 = kind + 失败测试名集合 + 异常类集合。 - - 集合都排序去重，所以「两个测试挂了」不会因报错顺序不同而分成两个签名。 - - ⭐ 拿不到任何结构化信号时（如 kind='timeout' / 'no_code'，根本没跑到 unittest）， - 两个集合都为空，签名退化成单独的 kind —— 这正是想要的：同一题的 8 次超时是 - 同一件事，只该诊一次。 - """ - kind = str(roll.get('kind') or 'unknown') - err = str(roll.get('error') or '') - tests = ','.join(sorted(set(_TEST_RE.findall(err)))) - excs = ','.join(sorted(set(_EXC_RE.findall(err)))) - return f'{kind}\x00{tests}\x00{excs}' - - -def bucket_failures(rolls: List[Dict[str, Any]]) -> List[Tuple[Dict[str, Any], int]]: - """把一题的 rolls 按失败签名分桶，返回 [(代表 roll, 出现次数), ...]，按次数降序。 - - 只取失败的 roll。代表 roll 用该桶里第一条 —— 同签名意味着同一组 (测试, 异常类)， - 但报错正文仍可能略有差异（具体数值）；教师看到的是这条代表的完整报错，信息不丢。 - """ - buckets: Dict[str, Dict[str, Any]] = {} - for r in rolls: - if r.get('correct'): - continue - sig = _signature(r) - b = buckets.get(sig) - if b is None: - buckets[sig] = {'roll': r, 'n': 1} - else: - b['n'] += 1 - out = [(b['roll'], b['n']) for b in buckets.values()] - out.sort(key=lambda t: -t[1]) # 高频根因优先 - return out - - -class MultiDiagCache: - """按 (data_id, 失败签名) 缓存单条诊断；磁盘格式与 e23_rubric 的缓存文件同构但独立成文件。 - - 独立文件的理由：键的语义不同（这里含失败签名），混进同一个文件会让 E23 的缓存读取拿到 - 对不上的条目。E23 那份缓存**完全不动**，两个实验各自可复现。 - """ - - def __init__(self, path: str = None): - self.path = path or os.path.join( - os.path.dirname(os.path.abspath(__file__)), - 'multidiag_cache_code_v1.jsonl') - self._idx: Dict[str, Any] = {} - import collections - self.stats: collections.Counter = collections.Counter() - if os.path.exists(self.path): - with open(self.path, encoding='utf-8') as f: - for line in f: - try: - rec = json.loads(line) - self._idx[rec['key']] = rec['value'] - except Exception: - continue - logger.info(f'[multidiag] 缓存载入 {len(self._idx)} 条：{self.path}') - self._fh = open(self.path, 'a', encoding='utf-8') - # ⭐ _put 会被 DIAG_WORKERS 个线程并发调用，write + flush 两步不是原子的：无锁时 - # 两条记录会交错成半行，下次启动 json.loads 解不开就默默丢掉（except: continue）， - # 表现是「明明诊过却反复花 API 钱」且无任何报错。并行度提到 24 后这个概率不再可忽。 - self._lock = threading.Lock() - - def _put(self, key: str, value: Any) -> None: - line = json.dumps({'key': key, 'value': value}, ensure_ascii=False) + '\n' - with self._lock: - self._idx[key] = value - self._fh.write(line) - self._fh.flush() - - def _one(self, checker, record: Dict[str, Any], roll: Dict[str, Any]) -> Optional[Dict]: - """诊断单条失败轨迹，返回 validate 过的 diag dict（不可救/失败返回 None）。 - - 与 e23_rubric.get_or_diagnose 的差别只有缓存键：这里把失败签名放进键，所以同一题的 - 不同失败模式各占一个槽位。校验逻辑直接复用 R._validate，保持判据完全一致。 - """ - sig = _signature(roll) - key = hashlib.md5( - f"{R.RUBRIC_VERSION}\x00{R._PROMPT_KEY}\x00" - f"{record.get('data_id', '')}\x00{sig}".encode('utf-8')).hexdigest() - if key in self._idx: - cached = self._idx[key] - if not cached: - self.stats['hit_dropped'] += 1 - return None - self.stats['hit'] += 1 - return cached - query = R.diag_query(record['problem'], record['reference_answer']) - try: - obj = checker.classify(query, R.diag_segment(roll)) - except Exception as exc: - logger.warning(f'[multidiag] classify error: {exc}') - obj = None - if obj is None: - # 同 e23_rubric：API 故障**绝不缓存**，否则一次抖动会永久丢掉这个失败模式。 - self.stats['api_fail'] += 1 - return None - diag = R._validate(obj, query) - if diag is None: - self._put(key, None) # 稳定判决：这个失败模式教师给不出可用分类 - self.stats['dropped_unaddressable'] += 1 - return None - self.stats['ok'] += 1 - self.stats[f"class_{diag['class']}"] += 1 - self._put(key, diag) - return diag - - def diagnose_task(self, checker, record: Dict[str, Any], - rolls: List[Dict[str, Any]]) -> str: - """一题的多失败模式诊断 -> 合并后的 rubric 文本（无可用诊断返回 ''）。""" - buckets = [(r, n) for r, n in bucket_failures(rolls) if n >= MIN_SIGNATURE_COUNT] - if not buckets: - return '' - buckets = buckets[:MAX_DIAG_PER_TASK] - diags = [] - for roll, n in buckets: - d = self._one(checker, record, roll) - if d: - diags.append((d, n)) - return merge_diags(diags, n_rollouts=len(rolls)) - - def diagnose_many(self, checker, - jobs: List[Tuple[Dict[str, Any], List[Dict[str, Any]]]]) -> List[str]: - """并行版。jobs = [(record, rolls), ...]，返回对齐的 rubric 文本列表。 - - 并行度按**题**而不是按签名：一个线程认领一道题，题内的最多 - MAX_DIAG_PER_TASK 次 classify 仍串行，所以同时在飞的 request 数 = min(DIAG_WORKERS, 题数)。 - 题数通常远大于 DIAG_WORKERS，所以实际并行就是 DIAG_WORKERS。 - - 不把签名也展平成任务（那会再快 ~3 倍）的原因不是缓存安全 —— _put 已加锁； - 而是签名数预先不知道，展平后难以把结果按题对齐回去，且同题的多条诊断本身 - 就要合并。纯 API 等待不占 GPU。 - """ - if not jobs: - return [] - workers = max(1, min(DIAG_WORKERS, len(jobs))) - with ThreadPoolExecutor(max_workers=workers) as ex: - return list(ex.map(lambda j: self.diagnose_task(checker, j[0], j[1]), jobs)) - - def close(self): - self._fh.close() - - -def merge_diags(diags: List[Tuple[Dict[str, Any], int]], n_rollouts: int = 0) -> str: - """把多条诊断拼成一份给 skill-gen 的文本。 - - ⭐ 与 e23_rubric._format 的两处关键差别： - 1. **不再输出** `ALSO OFF (do not write about these)`。那条禁令是 E23 单根因口径的产物， - 而本模块的全部目的就是让 skill 覆盖多处根因 —— 留着它会自相矛盾。 - 2. 带上 `seen k/M times` 频次。skill-gen 据此知道哪个根因更普遍、该先写哪个； - 只翻车 1/8 次的偶发问题不该和翻车 5/8 次的主因同等对待。 - - evidence 仍然**不进**文本（与 _format 一致）：它是单测报错原文，断言 diff 里带期望值， - 是最强的答案泄漏通道。 - """ - if not diags: - return '' - if len(diags) == 1: - d, n = diags[0] - head = [f"DECISIVE FAILURE: {d['class']} — {R._CLASS_SHORT[d['class']]}", - f"WHAT WENT WRONG: {d['reason']}", - f"PRIOR THAT WOULD HAVE PREVENTED IT: {d['prior']}"] - if n_rollouts and n: - head.insert(1, f'OBSERVED: this failure appeared in {n}/{n_rollouts} attempts.') - return '\n'.join(head) - lines = [f'The attempt failed in {len(diags)} distinct ways across {n_rollouts} attempts. ' - f'Address ALL of them — fixing only the first will still fail the tests.'] - for i, (d, n) in enumerate(diags, 1): - freq = f' (seen {n}/{n_rollouts})' if n_rollouts else '' - lines.append( - f"\nFAILURE {i}: {d['class']} — {R._CLASS_SHORT[d['class']]}{freq}" - f"\n WHAT WENT WRONG: {d['reason']}" - f"\n PRIOR THAT WOULD HAVE PREVENTED IT: {d['prior']}") - return '\n'.join(lines) - - -def multidiag_metrics(stats) -> Dict[str, float]: - """缓存/诊断计数 -> train_log 指标（与 e23_rubric.cache_metrics 同风格）。""" - tot = max(1, stats['hit'] + stats['ok'] + stats['hit_dropped'] - + stats['dropped_unaddressable'] + stats['api_fail']) - return {'rubric/hit_rate': (stats['hit'] + stats['hit_dropped']) / tot, - 'rubric/ok': float(stats['ok']), - 'rubric/dropped_unaddressable': float(stats['dropped_unaddressable']), - 'rubric/api_fail': float(stats['api_fail'])} diff --git a/cookbook/human_e18/e18_prompts.py b/cookbook/human_e18/e18_prompts.py deleted file mode 100644 index 28a6e0c12..000000000 --- a/cookbook/human_e18/e18_prompts.py +++ /dev/null @@ -1,264 +0,0 @@ -"""E18 的全部 prompt 文本与拼装函数：executor / 教师 judge / skill-gen 三处。 - -与 cookbook/human/e23_prompts.py 同源（同一套 BigCodeBench 交付要求与失败分类表），差别只在 -skill-gen 的系统提示：E18 是**拒绝采样 SFT**，采集时要 think 模式、训练时用 nothink 布局， -所以这里额外提供 query-only 的训练轨迹拼装（train_prompt）。 -""" -# flake8: noqa: E501 -# prompt 正文按「一段一行」书写，折行会改变真正发给模型的文本，故整文件豁免行长检查。 -from typing import Any, Dict - -# =========================================================================== -# executor -# =========================================================================== -# BigCodeBench 官方 instruct 模式的硬性交付要求（与 e23 逐字相同，保证跨实验可比）。 -EXEC_SYSTEM = """\ -You are an expert Python engineer. You will be given a task description that ends with the exact \ -import lines and function signature your solution must start with. - -Deliver exactly one fenced Python code block and nothing else after it: -- Reproduce the given imports and the given function signature verbatim, including parameter \ -names, order and default values. -- Add any further imports you need inside the same block; the block must run standalone. -- Return exactly the object type the task says to output. If it says the function should output \ -a tuple, return a tuple in that order; if it names a matplotlib Axes, return the Axes object \ -itself, not the Figure and not None. -- Implement the described behaviour for the general case, including the empty / single-element / \ -missing-column edge cases and any exception the description says to raise. -- Do not call the function, do not print demonstrations, do not add tests, do not use \ -`if __name__ == '__main__'`, and do not read from stdin. -- Do not include explanations outside the code block.""" - -# E18 的 executor 只看抽出来的 <skills> 块，**不看** actor 的 <think>：本臂的产物是要写进 -# SFT 数据集、将来 query-only 部署的 skill 文本，采集期就必须按「部署时 executor 能看到什么」 -# 来判分，否则筛出来的胜者依赖一段部署时不存在的思考过程。（E23 是相反的口径，故意保留差异。） -_WRAPPER_SKILL_ONLY = ( - 'Hint:\n{hint}\n') - - -def direct_prompt(problem: str) -> Dict[str, Any]: - return {'messages': [{'role': 'system', 'content': EXEC_SYSTEM}, - {'role': 'user', 'content': problem}]} - - -def skill_solve_prompt(problem: str, skill: str) -> Dict[str, Any]: - """题面 + 抽出的 <skills> 块。skill 为空 -> 干净 direct（等价裸解，不塞空指导）。""" - skill = (skill or '').strip() - if not skill: - return direct_prompt(problem) - return {'messages': [{'role': 'system', 'content': EXEC_SYSTEM}, - {'role': 'user', - 'content': problem + '\n\n' + _WRAPPER_SKILL_ONLY.format(hint=skill)}]} - - -# =========================================================================== -# skill-gen（采集用：rubric 条件化、think 模式） -# =========================================================================== -# ⭐ 全英文、且**不给回复格式模板**。 -# 旧版是中文，并在类型 3/4 里给了带 `...` 占位符的回复示范（“根据你曾经犯过的错误……”、 -# “该问题属于...问题，因此可以拆解为...步骤”）。实测后果：模型把骨架连 `...` 一起抄下来， -# 胜者中 79/148 含空模板、中文占比从 65% 升到 100%，唯一词数 376→123，candidate_pass_rate -# 从 0.856 塌到 0.701。语言还与选择器共谋：中文字符数天然更少，在旧的 LEN_BUDGET 口径下永远 -# 更贴近预算而胜出。故：只说**要写什么**，不给句式；措辞由模型自己生成。 -# executor 与 BCB 题面均为英文，skill 也必须是英文才能与之对齐。 -# ⭐ 类型列表的**描述粒度必须齐平**，否则列表本身就是偏置：旧版 3/4 带了展开要求 -# （“Explain what the weak point is and why…”、“the concrete steps needed…”），1/2/5 却只有光板 -# 一句 —— 实测胜者里 t4 占 77-100%、t3 占 62-80%，而 t2 恒为 0%、t1 不过 4%。 -# 同理不写“vary across attempts”：单次采样看不到其他 rollout，该句对单条生成无法执行； -# 多样性靠 N_SKILLS 个独立 rollout 的采样噪声，以及“只选一类 + 五类等价”的显式声明。 -# ⭐ 2026-08-06：改为 **narrative 文体**（移植 skill2lora 的 SKILL_GEN_SYSTEM，见 -# cookbook/exp/skill2lora/train_skill_v2.py:843-861）。原版（下方 _SKILLGEN_SYSTEM_TYPED， -# 已注释停用）是「五类里挑一类」的列表式，实测问题： -# 1. 46.6% 的胜者是「用 A 不要用 B」的 API 纠正，靠的是 rubric 里的库行为知识； -# 2. 12.3% 直接把只存在于隐藏单测里的列名（'closing_price' 之类）写进 skill —— query 里 -# 根本没有，eval 时 query-only 无从得知，训练等于教模型凭空猜列名； -# 3. 5.9% 用「The critical error was...」这种事后复盘句式指代一次 eval 时不存在的失败。 -# narrative 的三条硬约束正好对上后两条：强制第一人称自持句式、**明令禁止**指向外部上下文 -# （skill2lora 的原注释写明这类措辞「会导致幻觉」）、不许代入本题具体数值。 -# ⭐ rubric 的用法：单独一段说明「有诊断时当作证据用」，与下方禁指代外部上下文那条并不矛盾 —— -# 两者分属不同层：**任务指令层**要求模型靠诊断定位软肋（WHERE），**输出层**要求把它转写成 -# executor 能直接执行的前瞻告诫（不得提及诊断本身）。这正是 skill2lora 的 REGEN_SYSTEM -# （train_skill_v2.py:1130-1147）的做法：步骤 2 要求「weaving together ... the pitfalls that -# actually tripped up the solving process」，而 Output requirement 同时禁止 "according to the given -# analysis/hints"。差异在于：REGEN 是「旧 skill + 诊断 -> 重写」的蒸馏场景，E18 是首次生成， -# 所以此处写成条件句（"If a grader's diagnosis ... is supplied"），无诊断时自动退化成纯预判。 -# 为何不能直接写「根据你之前犯过的错误……」：executor 看不到任何「之前」，而且训练目标是 -# query-only 的 —— 实测 5.9% 的胜者写成「The critical error was...」，eval 时模型无错可指只能编。 -# 所以保留“定位到具体步骤/API”这个内核，只把时态从「已发生」换成「容易在此处发生」。 -# ⚠️ 代价（skill2lora 已记录）：few-shot 例子占整条 prompt 的 65%，每次采样必然命中，会锁死 -# 文体与长度 —— 多样性下降是预期内的，换来的是「删掉本题依然成立」的可迁移性。 -# ⚠️ 本文件的 few-shot 例子必须是**代码域**的（原版是数学域的「末位数字/计数问题」，直接搬过来 -# 会把 executor 往数学叙述上带）；收尾纪律句同理换成 BCB 的交付要求，不用 boxed 那句。 -SKILLGEN_SYSTEM = """\ -You are a skill-writing expert. Your <skills> block will be fed to a SEPARATE downstream executor model that must solve the Python task on its own. The executor will NOT see your private reasoning — it only sees what is inside <skills>...</skills>. - -First, think privately: work out where the executor is most likely to get stuck, then step back and abstract WHAT MAKES THIS TYPE OF TASK GO WRONG into transferable guidance. - -Second, if a grader's diagnosis of a failed attempt is supplied, use it as your evidence for WHERE the weak point is: read what actually went wrong, then decide which part of the approach needs the executor's attention. Fold that insight into the narrative as guidance the executor can act on before it starts — name the step or the API where the trouble lives and say what to do there instead, e.g. "the place this tends to go wrong is when you ..., so at that point you should ...". Do not report the diagnosis; convert it into advice. - -Then write the <skills> block following these rules: -- Give general, transferable techniques for this TYPE of task: the library behaviour it relies on, the recommended approach, and the common pitfalls to avoid — plus a brief reason for each piece of advice so the executor understands why. -- Write it as one coherent analysis narrative (not a bullet list): first name what the task is essentially asking, then walk through how to approach it, blending the API contracts, steps, pitfalls and reasons into a single connected story. -- Write your judgements directly in the first person (e.g. "I think this step tends to ...", "A common mistake is ..., so you need to ..."), and phrase every issue as a self-contained, general technique. -- CRITICAL: Do NOT use phrasings that point to external context, such as "according to the given diagnosis", "the failed attempt", or "the previous error". The executor cannot see that context, and such phrasings will cause hallucination. State the pitfall as something that tends to happen at a particular step, not as something that already happened. -- CRITICAL: Do NOT name a column, key, or literal value that the task description does not itself state. If the task never names its columns, say how to discover them from the input instead of guessing names. -- Name the concrete API, argument, or keyword involved whenever the task description supports it. -- Keep it concise: aim for roughly one focused paragraph. - -Put ONLY the methodology inside <skills></skills>. - -Example: -<skills> -This task is essentially asking you to reshape tabular input and hand back a plot object, so the delivery contract matters here as much as the computation; I would pin down exactly what type the function must return before writing any logic, because returning a Figure where an Axes was requested fails even when every number is right. The first place this tends to go wrong is the input itself: it arrives as a plain container, and I find the single most common break in this type of task is assuming it is already a DataFrame — dictionaries and lists of tuples carry none of the frame methods, so reaching for column-based access on them raises immediately, and at that point you should check what the object actually is and build the frame from it explicitly. The next place to slow down is naming: let the task description dictate the column names and read them off the signature or the docstring rather than inventing plausible-sounding ones, and when the description never states them, derive them from the input's own keys instead of hard-coding a guess, because a name that merely sounds right will pass your own reading and still miss. A common mistake is treating an empty or single-element input as impossible, so decide up front whether it should yield an empty result or raise, and write that branch before the main path. Finally, when plotting, create the Axes explicitly and return that same object, since helper calls that draw on the current figure make it easy to hand back something you never configured. Overall I summarise this type of task as "fix the return contract, verify the input's real type, take names from the description, then handle the empty case before the happy path", because that is where the failures concentrate. -</skills> -""" - -# =========================================================================== -# 【已停用】原「五类挑一类」列表式 skill-gen 提示（2026-08-06 换成上方 narrative） -# =========================================================================== -# 保留全文仅为记录历史口径与可回退：把下面的字符串改名回 SKILLGEN_SYSTEM 即可复原。 -# 停用原因见上方 narrative 块的注释（隐藏契约泄漏 12.3% / 复盘句式 5.9%）。 -# 注意它自身也修过两轮：类型描述粒度齐平（t1-t5 各 9-15 词）、以及那条前瞻视角规则 —— -# 这两笔修改都已被 narrative 的硬约束覆盖，回退时才需要重新评估。 -_SKILLGEN_SYSTEM_TYPED = """\ -You are a skill-writing expert. Your job is to write an advisory note that makes a downstream executor model solve the given Python task more accurately. - -First decide where the executor is most likely to get stuck, then write the advice you believe helps most. Pick the ONE kind below that fits this task best. All five are equally worth choosing, and a single sharp sentence often beats a long note: -1. A plain instruction about how to approach the work, such as what to be careful about. -2. A calibration cue about how much to deliberate, or about trusting its own judgement. -3. A generalized lesson drawn from the grader's diagnosis of a previous failed attempt, if one is supplied. Explain what the weak point is and why the lesson prevents it. -4. A decomposition of the task into the concrete steps needed to solve it. -5. Any other kind of skill you judge useful, including an angle you would not normally try. - -Rules: -- Do NOT solve the task and do NOT write code. You only write advice. -- Name the concrete API, argument, key, or value involved whenever you can. -- Be direct and specific. No filler, no restating the task, no placeholder text. -- Write forward-looking advice to someone who has not attempted the task yet. Do not refer to an error, mistake, or attempt as something that already happened. -- Choose your own wording and structure; there is no fixed format to follow. - -Wrap your skills in <skills> ... </skills>. -""" - -SKILLGEN_USER = """\ -TASK -{problem} - -GRADER'S DIAGNOSIS OF THE FAILED ATTEMPT -{rubric} - -Write the advisory note now, wrapped in <skills></skills>.""" - -# ⭐ 带失败代码的变体。与 SKILLGEN_USER 的差别只有多出的 FAILED ATTEMPT 段， -# 段序是「题面 -> 失败代码 -> 诊断」：诊断紧贴写作指令，因为它才是要被消化的主结论； -# 把代码放中间让模型先看到证据再看结论，而不是反过来。 -SKILLGEN_USER_TRAJ = """\ -TASK -{problem} - -CODE FROM A FAILED ATTEMPT (for your analysis only — the executor will never see it) -{trajectory} - -GRADER'S DIAGNOSIS OF THE FAILED ATTEMPT -{rubric} - -Write the advisory note now, wrapped in <skills></skills>.""" - - -def format_trajectory(code: str, error: str = '', kind: str = '', - max_chars: int = 4000) -> str: - """把一条失败 rollout 整理成给 skill-gen 看的文本块。 - - ⭐ 头尾各留一半而不是直接截前 max_chars：Python 失败代码的关键信息经常在**末尾** - （未闭合的分支、漏掉的 return、被截断的行），只留开头会把根因裁掉。 - - ⚠️ error 只取前 400 字符：pytest 的 longrepr 能有几千字符且大量重复的堆栈帧， - 全带上会把预算吃光，而判别失败模式只需要头部的异常类型与消息。 - """ - code = (code or '').strip() - if not code: - return '(the attempt produced no extractable code)' - if len(code) > max_chars: - half = max_chars // 2 - code = (code[:half] + '\n\n... [%d characters omitted] ...\n\n' % (len(code) - max_chars) - + code[-half:]) - out = ['```python', code, '```'] - if error: - out.append('OBSERVED ERROR: ' + ' '.join(str(error).split())[:400]) - if kind: - out.append('FAILURE CATEGORY: %s' % kind) - return '\n'.join(out) - - -# ⭐ 训推一致：这份同时做**训练 prompt** 与 **eval prompt**，与 SKILLGEN_SYSTEM 一起改成英文； -# 两边语言不一致会让模型在采集与部署时面对不同分布。 -SKILLGEN_SYSTEM_EVAL = """\ -You are a skill-writing expert. Your job is to write an advisory note that makes a downstream executor model solve the given Python task more accurately. - -First decide where the executor is most likely to get stuck, then write the advice you believe helps most. - -1. You have been trained on many kinds of skills, from a one-line caution to a full step decomposition. -2. Your memory already holds what works best for different kinds of problems. -3. Analyse the task and choose the skills you judge most useful. A single sharp sentence often beats a long note. - -Rules: -- Do NOT solve the task and do NOT write code. You only write advice. -- Name the concrete API, argument, key, or value involved whenever you can. -- Be direct and specific. No filler, no restating the task, no placeholder text. -- Write forward-looking advice to someone who has not attempted the task yet. Do not refer to an error, mistake, or attempt as something that already happened. - -Wrap your skills in <skills> ... </skills>. -""" - -SKILLGEN_USER_EVAL = """\ -TASK -{problem} - -Write the advisory note now, wrapped in <skills></skills>.""" - - -def skillgen_prompt(problem: str, rubric: str, eval: bool, - trajectory: str = '') -> Dict[str, Any]: - """trajectory 非空时切到带失败代码的 user 模板（由 KOD_USE_TRAJ 控制，默认关）。 - - ⚠️ 只换 user 模板、**不换 system**：SKILLGEN_SYSTEM 里那条 - "Do NOT use phrasings that point to external context ... 'the failed attempt'" - 的禁令对带轨迹的情形更重要（模型看到真实代码后更容易写成事后复盘）， - 换掉 system 会同时丢掉这条约束。 - """ - if not eval: - if not rubric: - rubric = ('No diagnosis is available for this task. Consider the other kinds of ' - 'skill instead.') - user = (SKILLGEN_USER_TRAJ.format(problem=problem, rubric=rubric, - trajectory=trajectory) - if trajectory else - SKILLGEN_USER.format(problem=problem, rubric=rubric)) - return {'messages': [{'role': 'system', 'content': SKILLGEN_SYSTEM}, - {'role': 'user', 'content': user}]} - else: - return {'messages': [{'role': 'system', 'content': SKILLGEN_SYSTEM_EVAL}, - {'role': 'user', 'content': SKILLGEN_USER_EVAL.format(problem=problem)}]} - - -# =========================================================================== -# 已废弃：训练/eval 统一走 skillgen_prompt(..., eval=True) -# =========================================================================== -# ⭐ 不要再用 TRAIN_SYSTEM / train_prompt。 -# 训推一致要求「训练 prompt 与 eval prompt 逐字相同」，而 eval 用的是 SKILLGEN_SYSTEM_EVAL； -# 再并行维护一份英文 TRAIN_SYSTEM 只会让两边默默分叉。保留它仅为记录历史口径。 -TRAIN_SYSTEM = """\ -You are a problem-solving coach for a Python engineer. Given a task, write a short advisory note that anticipates the most likely decisive mistake and prevents it. - -Requirements: -- Wrap the note in <skills> and </skills> tags. -- State what to do, in the imperative. Name the concrete API, argument, key, or value involved. -- Do NOT solve the task, do NOT write code, and do NOT state the expected output value. -- Keep it under 90 words.""" - - -def train_prompt(problem: str) -> Dict[str, Any]: - """已废弃。训练与 eval 统一用 `skillgen_prompt(problem, '', eval=True)`。""" - raise NotImplementedError( - 'train_prompt 已废弃：训练/eval 请用 skillgen_prompt(problem, \'\', eval=True)，' - '以保证两边的 system/user 逐字一致（训推一致）。') diff --git a/cookbook/human_e18/e18_rejection_sft.py b/cookbook/human_e18/e18_rejection_sft.py deleted file mode 100644 index 01229d832..000000000 --- a/cookbook/human_e18/e18_rejection_sft.py +++ /dev/null @@ -1,775 +0,0 @@ -#!/usr/bin/env python3 -import json -import os -import shutil -import sys -import time -from dataclasses import dataclass -from typing import Any, Dict, List, Tuple - -import torch -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.checkpoint_engine import CheckpointEngineManager -from twinkle.data_format import SamplingParams, pack_user_data -from twinkle.dataloader import DataLoader -from twinkle.model import TransformersModel -from twinkle.patch.no_split_modules import NoSplitModulesPatch -from twinkle.processor import InputProcessor -from twinkle.sampler import vLLMSampler -from twinkle.template import Template - -# 环境层与教师 judge 复用 human/ 下的 e23 模块（不拷贝，保证判分/诊断逐字同源）。 -_HERE = os.path.dirname(os.path.abspath(__file__)) -_HUMAN = os.path.abspath(os.path.join(_HERE, '..', 'human')) -if _HUMAN not in sys.path: - sys.path.insert(0, _HUMAN) - -from e23_bcb import clean_text, empty_roll, extract_skill, judge_seqs, load_records # noqa: E402 -from e23_rubric import build_checker, class_metrics # noqa: E402 -# 多轨迹诊断：每种失败模式各诊一次再合并。不用 e23_rubric.RubricCache 是因为它的缓存键 -# 只含 data_id，同一题调 N 次会全部命中第一次的结果 —— 「每条 rollout 各诊一次」在那个键下做不到。 -from e18_multidiag import MultiDiagCache, multidiag_metrics # noqa: E402 - -from e18_prompts import direct_prompt, skill_solve_prompt, skillgen_prompt # noqa: E402 -from e18_select import gain_stats, select_winner # noqa: E402 - -try: - import swanlab -except ImportError: - swanlab = None - -logger = get_logger() - -# ========== Configuration ========== -MODEL_ID = os.environ.get('MODEL_ID', 'Qwen/Qwen3-4B') -OUTPUT_DIR = os.environ.get('OUTPUT_DIR', os.path.join(_HERE, 'output.e18')) - -TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 2)) -SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) -# ⭐ base_sampler 拿 4 张（而不是和其他两组一样的 2）：它是唯一瓶颈。每 chunk 的序列数 -# 相差一个量级 —— skill_sampler 只跑 CHUNK_SIZE*N_SKILLS（512），而 base_sampler 要跑 -# 裸解 CHUNK_SIZE*EXEC_ROLLOUTS（512）+ 重解 CHUNK_SIZE*N_SKILLS*EXEC_ROLLOUTS（4096）= 4608， -# 且 EXEC_MAX_TOKENS（15000）远大于 SKILL_MAX_TOKENS（8192），token 预算相差约 16 倍。 -# 给 skill_sampler 加卡几乎无收益，加在这里才能缩短 wall clock。 -BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 4)) -# ⭐ 没有 ref 模型：SFT 是纯交叉熵，不需要 KL 参考。E23 的 REF_GPUS 那两张转给了 base_sampler， -# 而不是空着 —— 8 卡机器上「省卡」没有意义，只会让瓶颈环节白白排队。 -NUM_GPUS = TRAIN_GPUS + SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS -TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', 1)) -TRAIN_DP = TRAIN_GPUS // TRAIN_FSDP -GPU_MEM = float(os.environ.get('GPU_MEM', 0.8)) - -SEED = int(os.environ.get('SEED', 42)) -CHUNK_SIZE = int(os.environ.get('CHUNK_SIZE', 64)) # 每轮裸解多少题去筛错题 -N_SKILLS = int(os.environ.get('N_SKILLS', 8)) # 每题采多少 skill 候选（拒绝采样的池） -# 攒够多少条胜者才 SFT 一次。必须是 TRAIN_DP 的整数倍（dp 切分要求），否则末尾会被丢。 -ACCUMULATE = int(os.environ.get('ACCUMULATE', 16)) -MAX_UPDATES = int(os.environ.get('MAX_UPDATES', 200)) -EVAL_SIZE = int(os.environ.get('EVAL_SIZE', 100)) -EVAL_EVERY_UPDATES = int(os.environ.get('EVAL_EVERY_UPDATES', 10)) -SAVE_EVERY_UPDATES = int(os.environ.get('SAVE_EVERY_UPDATES', 50)) # 0 = 只在结束时存 - -# skill-gen 采集：think 开、T=1（要多样性才有拒绝采样的意义） -SKILL_MAX_TOKENS = int(os.environ.get('SKILL_MAX_TOKENS', 8192)) -SKILL_GEN_TEMPERATURE = float(os.environ.get('SKILL_GEN_TEMPERATURE', 1.0)) -SKILL_GEN_TOP_P = float(os.environ.get('SKILL_GEN_TOP_P', 1.0)) -SKILL_GEN_TOP_K = int(os.environ.get('SKILL_GEN_TOP_K', -1)) -EXEC_MAX_TOKENS = int(os.environ.get('EXEC_MAX_TOKENS', 15000)) -MAX_MODEL_LEN = int(os.environ.get('MAX_MODEL_LEN', 16000)) - -# ⭐ pass@k 评分：executor 不再用 greedy 单次，而是每个 prompt 重解 EXEC_ROLLOUTS 次取通过率。 -# 为何必须：M=1/T=0 时 pass_rate 只有 0/1，易题上「加任何 skill 都对」，正例标签与 skill -# 质量无关 —— 等于往数据集里灌随机 skill。跑 8 次取连续 pass_rate 后，同一题的不同 skill 之间 -# 才有方差（如 8/8 vs 5/8），能真正排序。代价：executor GPU 时间乘 ~8 倍。 -# 温度必须 >0：T=0 下 8 次采样会逐字相同（judge_seqs 还会按 code 去重），pass_rate 退回 0/1。 -EXEC_ROLLOUTS = int(os.environ.get('EXEC_ROLLOUTS', 8)) -EXEC_TEMPERATURE = float(os.environ.get('EXEC_TEMPERATURE', 0.6)) -EXEC_TOP_P = float(os.environ.get('EXEC_TOP_P', 0.95)) - -LR = float(os.environ.get('LR', 1e-5)) # 恒定 lr，无 warmup 无 decay -TRAIN_MICRO_BATCH = int(os.environ.get('TRAIN_MICRO_BATCH', max(TRAIN_DP, ACCUMULATE // 2))) - -# 三道筛的两个阈值 -SKILL_CHAR_LIMIT = int(os.environ.get('SKILL_CHAR_LIMIT', 1500)) # 超过直接丢 - -# ⭐ 入池门槛：skill 至少要多做对 MIN_GAIN_ROLLOUTS 次（默认 2，即 +2/8 = +0.25）。 -# 为何不收 tie（8/8 -> 8/8）：那种样本只能证明 skill 无害，对「学会写有效 skill」没有任何 -# 监督信号 —— 易题上加任何 skill 都是 8/8，收它等于往数据集里灌随机文本。 -# 为何阈值是 2 而不是 1：8 次采样下 +1/8 在采样噪声量级内（二项分布标准误约 0.17）， -# 分不清是真提升还是波动；要求 +2/8 才能把噪声挤出去。用 rollout 数而不是写死 0.25， -# 是为了改 EXEC_ROLLOUTS 时该语义（「多做对几次」）保持不变。 -MIN_GAIN_ROLLOUTS = int(os.environ.get('MIN_GAIN_ROLLOUTS', 2)) -MIN_PASS_GAIN = MIN_GAIN_ROLLOUTS / max(1, EXEC_ROLLOUTS) - -SWAN_PROJ = os.environ.get('SWAN_PROJ', 'twinkle') -RUN_TAG = os.environ.get('RUN_TAG', '').strip() -RUN_ID = time.strftime('%m%d-%H%M%S') - - -@dataclass -class Runtime: - skill_model: Any - skill_sampler: Any - base_sampler: Any - ckpt: Any - checker: Any - rubric_cache: MultiDiagCache - - -# =========================================================================== -# 采样 / 小工具 -# =========================================================================== -def run_samples(sampler, prompts, num_samples, max_tokens, gen_dp, - temperature=None, top_p=None, top_k=None, logprobs=None): - """采样。prompts 少于 dp 时补齐再截回 —— Ray 的 dp 切分要求每个 rank 至少一条。""" - if not prompts: - return [] - import copy - params = SamplingParams( - max_tokens=max_tokens, - temperature=0.6 if temperature is None else temperature, - top_p=0.95 if top_p is None else top_p, - num_samples=num_samples, - **({} if top_k is None else {'top_k': top_k}), - **({} if logprobs is None else {'logprobs': logprobs})) - padded = prompts - if gen_dp > 1 and 0 < len(prompts) < gen_dp: - padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] - responses = sampler.sample(padded, params)[:len(prompts)] - return [list(r.sequences) if (r and r.sequences) else [] for r in responses] - - -def first_seq(seqs): - return seqs[0] if seqs else None - - -def seq_text(seq) -> str: - return clean_text(getattr(seq, 'decoded', '') or '') if seq is not None else '' - - -def _mean(xs) -> float: - xs = [float(x) for x in xs if x is not None] - return sum(xs) / len(xs) if xs else 0.0 - - -# =========================================================================== -# 采集：裸解 -> 诊断 -> skill-gen -> executor 重解 -> 三道筛 -# =========================================================================== -def _pass_rate(rolls) -> float: - return (sum(1.0 for x in rolls if x['correct']) / len(rolls)) if rolls else 0.0 - - -def bare_solve(rt: Runtime, records, rollouts: int = None) -> List[List[Dict[str, Any]]]: - """裸题重解 `rollouts` 次，每条记录返回一个 roll 列表（长度 = 实际采到的序列数）。 - - 返回**嵌套**列表而不是单个 roll：调用方靠 _pass_rate() 取连续值。 - 判分全部汇到一次 judge_seqs（它内部按 (task_id, code) 去重，相同代码只跑一次单测）。 - """ - M = max(1, rollouts if rollouts is not None else EXEC_ROLLOUTS) - out = run_samples(rt.base_sampler, [direct_prompt(r['problem']) for r in records], - M, EXEC_MAX_TOKENS, BASE_SAMPLER_GPUS, - temperature=(0.0 if M == 1 else EXEC_TEMPERATURE), - top_p=(None if M == 1 else EXEC_TOP_P)) - pairs, spans = [], [] - for r, seqs in zip(records, out): - seqs = list(seqs or []) - spans.append(len(seqs)) - pairs.extend((s, r['reference_answer']) for s in seqs) - judged = judge_seqs(pairs) - rolls, i = [], 0 - for n in spans: - rolls.append(judged[i:i + n] if n else [empty_roll()]) - i += n - return rolls - - -def collect_chunk(rt: Runtime, chunk, ci: int) -> Tuple[List[Dict[str, Any]], Dict[str, float]]: - """一个 chunk 的采集，返回 (胜者列表, 指标)。 - - ⭐ **全量 rollout**：chunk 里每一道题都要采 skill 候选，包括裸解已经做对的。 - 理由：部署时 skill 模型面对的是任意题，不知道 executor 会不会做对；只在错题上训会让 - 它只学会「救难题」、在简单题上也写一大堆纠错式提示。 - - 诊断（rubric）只对**做错的题**拉，且**每一种失败模式各诊一次后合并**（见 e18_multidiag）： - 8 次 rollout 往往挂在不同地方，只拿「第一条失败轨迹」去诊等于抛硬币选根因，而且旧缓存键 - 只含 data_id，会把那次随机结果**永久写进缓存**。现在按失败签名分桶、按频次排序， - 合并文本里明确要求「Address ALL of them」。 - 无诊断（做对的、或 API 失败的）的题**不丢弃**：skillgen_prompt 内部会填兜底文案。 - 难易判定改用**连续 pass_rate**：裸解跑 EXEC_ROLLOUTS 次，base_pass_rate < 1 即视为「有提升 - 空间」。这比单次贪心稳得多 —— 单次 T=0 的对/错在临界题上换个种子就翻转。 - """ - base_rolls = bare_solve(rt, chunk) - base_rates = [_pass_rate(rr) for rr in base_rolls] - base_acc = _mean(base_rates) - # 诊断目标：没能每次都对的题。整组 rolls 都传进去 —— 由 multidiag 自己按失败签名分桶， - # 每种模式诊一次（同签名只花一次 API），再合并成一份 rubric。 - wrong = [(r, rr) for r, rr, rate in zip(chunk, base_rolls, base_rates) if rate < 1.0] - - before = rt.rubric_cache.stats.copy() - diags = rt.rubric_cache.diagnose_many(rt.checker, wrong) - rmetrics = multidiag_metrics(rt.rubric_cache.stats - before) - diag_by_id = {id(r): d for (r, _rr), d in zip(wrong, diags) if d} - n_rubric_missing = len(wrong) - len(diag_by_id) - # 多根因覆盖率：合并文本里有几段 FAILURE。持续=1 说明多轨迹诊断没带来新信息。 - n_multi = sum(1 for d in diag_by_id.values() if d.count('FAILURE ') > 1) - - # todo 现在是**全量** chunk：(record, rubric_or_empty, base_pass_rate) - todo = [(r, diag_by_id.get(id(r), ''), rate) - for r, rate in zip(chunk, base_rates)] - - # skill-gen：think 模式、T=1、每题 N 个候选（这就是拒绝采样的候选池）。 - # eval=False -> 用 SKILLGEN_SYSTEM（thinking 采集口径，允许多类型 skill）。 - sg = run_samples(rt.skill_sampler, - [skillgen_prompt(r['problem'], d, eval=False) for r, d, _br in todo], - N_SKILLS, SKILL_MAX_TOKENS, SKILL_SAMPLER_GPUS, - temperature=SKILL_GEN_TEMPERATURE, top_p=SKILL_GEN_TOP_P, - top_k=SKILL_GEN_TOP_K) - per_task: List[List[Dict[str, Any]]] = [] - flat = [] - for (r, _d, _br), seqs in zip(todo, sg): - cands = [] - for s in seqs or []: - resp = seq_text(s) - block = extract_skill(resp) - c = {'skills': block, 'response': resp, 'parseable': bool(block), - 'with_pass': None, 'kept': False, - 'skillgen_stop': getattr(s, 'stop_reason', None)} - cands.append(c) - if block: - flat.append((r, c)) - per_task.append(cands) - - # executor 带 skill 重解：**每个 skill 跑 EXEC_ROLLOUTS 次**，取连续 pass_rate。 - # 这是本次改造的核心：只有连续值才能在「全部都能做对」的易题上区分 skill 好坏。 - if flat: - ws = run_samples(rt.base_sampler, - [skill_solve_prompt(r['problem'], c['skills']) for r, c in flat], - EXEC_ROLLOUTS, EXEC_MAX_TOKENS, BASE_SAMPLER_GPUS, - temperature=EXEC_TEMPERATURE, top_p=EXEC_TOP_P) - pairs, spans = [], [] - for (r, _c), seqs in zip(flat, ws): - seqs = list(seqs or []) - spans.append(len(seqs)) - pairs.extend((s, r['reference_answer']) for s in seqs) - judged = judge_seqs(pairs) - i = 0 - for (_r, c), n in zip(flat, spans): - rr = judged[i:i + n] - i += n - c['with_pass'] = _pass_rate(rr) - c['n_rollouts'] = n - c['roll_kind'] = (next((x['kind'] for x in rr if not x['correct']), 'pass') - if rr else 'empty') - - # 三道筛：第一道改成**pass_rate 严格不降 + 取最大**（详见 select_winner） - accepted, sims = [], [] - n_pass_cands = n_survivors = 0 - n_acc_hard = n_acc_easy = 0 - gtot = {'improved': 0, 'tied': 0, 'degraded': 0} - gains = [] - for (r, d, base_rate), cands in zip(todo, per_task): - passers = [c for c in cands - if c.get('parseable') and (c.get('with_pass') or 0) >= base_rate] - n_pass_cands += len(passers) - for k, v in gain_stats(cands, base_rate).items(): - gtot[k] += v - best = select_winner(cands, d, r['reference_answer'], - skill_char_limit=SKILL_CHAR_LIMIT, - base_pass_rate=base_rate, min_pass_gain=MIN_PASS_GAIN) - if best is None: - continue - n_survivors += 1 - sims.append(best['rubric_similarity']) - gains.append(best['pass_gain']) - if base_rate >= 1.0: - n_acc_easy += 1 - else: - n_acc_hard += 1 - accepted.append({ - 'problem': r['problem'], 'reference_answer': r['reference_answer'], - 'data_id': r.get('data_id', ''), 'skills': best['skills'], - 'response': f"<skills>\n{best['skills']}\n</skills>", - # 审计字段（只进数据集文件，不进训练轨迹） - # base_pass_rate / with_pass / pass_gain：连续口径下判断“skill 到底有没有用”的依据。 - 'base_pass_rate': base_rate, 'with_pass_rate': best['with_pass'], - 'pass_gain': best['pass_gain'], 'gain_kind': best['gain_kind'], - 'rubric': d, 'chunk': ci, 'run': RUN_ID, - 'rubric_similarity': best['rubric_similarity'], - 'skill_chars': len(best['skills']), - 'n_candidates_passed': len(passers)}) - - _cand_pass = (n_pass_cands / max(1, len(todo) * N_SKILLS)) if todo else 0.0 - metrics = { - 'train/baseline_accuracy': base_acc, - 'train/accept_rate': (len(accepted) / len(todo)) if todo else 0.0, - 'train/candidate_pass_rate': _cand_pass, - # ⭐ 采集侧 lift：**全部**候选的平均增量，无选择偏差。不要用 - # `gain/selected_pass_gain` 代替它：后者只统计已通过 +MIN_PASS_GAIN 门槛的胜者， - # 按定义恒为正，衡量的是「被选中那条有多好」而非「模型平均能写多好」。 - # 与 `eval/lift` 也不可直接比：这里的 prompt 带 rubric（教师诊断），eval 是 query-only， - # 所以 train/lift 包含了“教师诊断的价值”，两者的差距正是本实验要缩小的东西。 - 'train/lift': _cand_pass - base_acc, - 'train/selected_rubric_similarity': _mean(sims), - 'train/selected_skill_length_characters': _mean( - [float(s['skill_chars']) for s in accepted]), - 'signal/n_wrong': float(len(wrong)), - 'signal/n_rubric_missing': float(n_rubric_missing), - 'signal/n_multi_cause': float(n_multi), - 'signal/n_accepted': float(len(accepted)), - # 分层接受数：易题 = base_pass_rate 已经 1.0（跑 8 次全对）。 - 'signal/n_accepted_hard': float(n_acc_hard), - 'signal/n_accepted_easy': float(n_acc_easy), - # 天花板题数：base_pass_rate 高到拿不到 +MIN_PASS_GAIN（如 8/8），结构性无法入池。 - # 它与 n_accepted_easy 合看：前者持续很大就说明大量 GPU 花在了注定不入池的题上。 - 'signal/n_ceiling': float(sum(1 for _r, _d, br in todo - if br + MIN_PASS_GAIN > 1.0 + 1e-9)), - 'signal/min_pass_gain': MIN_PASS_GAIN, - # 候选级增量分解（相对裸解 pass_rate）。degraded 最重要：skill 把通过率拉低了， - # 它不会反映在 accept_rate 上，持续偏高就说明 skill-gen 在写有害提示。 - 'gain/improved_candidates': float(gtot['improved']), - 'gain/tied_candidates': float(gtot['tied']), - 'gain/degraded_candidates': float(gtot['degraded']), - 'gain/degrade_rate': (gtot['degraded'] / max(1, sum(gtot.values()))), - 'gain/improve_rate': (gtot['improved'] / max(1, sum(gtot.values()))), - # 胜者的平均 pass_rate 增量：这才是「入池样本到底有多有用」的直接度量。 - # 持续趋近 0 就说明池子里全是「写了也白写」的 skill。 - 'gain/selected_pass_gain': _mean(gains), - } | rmetrics | class_metrics([d for _r, d, _br in todo if d]) - return accepted, metrics - - -def dump_dataset(accepted) -> None: - """胜者落盘 append-only 的 SFT 数据集（含 rubric/相似度/pass 全审计字段）。 - - 这份文件是 E18 的主产物：它让「筛选器选了什么」可离线复算、可跨 run 复用（不必重跑 GPU - 就能换 SFT 超参再训一遍）。 - """ - if not accepted: - return - path = os.path.join(OUTPUT_DIR, 'e18_sft_dataset.jsonl') - with open(path, 'a', encoding='utf-8') as f: - for s in accepted: - f.write(json.dumps(s, ensure_ascii=False) + '\n') - - -# =========================================================================== -# 训练：纯 SFT（query-only 轨迹） -# =========================================================================== -def train_batch(rt: Runtime, samples) -> Tuple[int, Dict[str, float]]: - """在攒够的胜者上做 1 个 optimizer step。 - - ⭐ 训推一致的关键：训练 prompt 段用 **`skillgen_prompt(..., eval=True)`**，与 run_eval / - 部署时用的系统提示逐字相同（SKILLGEN_SYSTEM_EVAL，不带 rubric）。 - 采集时用的是带诊断的 thinking 口径（SKILLGEN_SYSTEM）—— 那只是为了**邀出**好 skill， - 不能拿去当训练分布：线上没有诊断可用，拿带诊断的 prompt 去训会学成「看着诊断改写」。 - - 响应段走 messages 编码而不是拼采样 token：胜者的 `<skills>` 文本是程序合成的（采集时原 - 响应带 <think>、且包装不同），本来就没有对应的采样 token 序列。`key_rounds` 标出最后一轮 - （assistant）为唯一可训区，prompt 段全 -100。 - - loss 是纯 `CrossEntropyLoss`（不走 GRPO）：不传 advantages。拒绝采样的“选择”已经完成于三道筛 - （只有胜者入池），此处只需拟合目标文本。若要给样本加权，得在 loss 内部乘，而不是传 - advantages —— CrossEntropyLoss 不读这个参数，传了也是静默无效（所以此处根本不传）。 - """ - samples = [s for s in samples if (s.get('response') or '').strip()] - n = (len(samples) // TRAIN_DP) * TRAIN_DP # dp 切分要求整倍数 - if n == 0: - return 0, {} - samples = samples[:n] - trajs = [] - for s in samples: - msgs = skillgen_prompt(s['problem'], '', eval=True)['messages'] - trajs.append({'messages': msgs + [{'role': 'assistant', 'content': s['response']}], - 'user_data': pack_user_data({'key_rounds': [len(msgs)]})}) - micro = max(TRAIN_DP, min(TRAIN_MICRO_BATCH, n)) - for i in range(0, n, micro): - rt.skill_model.forward_backward(inputs=trajs[i:i + micro]) - rt.skill_model.clip_grad_and_step() - # ⭐ 这里**不**同步权重。同步只在 run_eval 前后发生（见 sync_for_eval），以保证 - # skill_sampler 在**采集**阶段永远是初始权重。 - # 为何：采集用 thinking 口径邀出候选，而训练目标是 nothink 的 <skills> 纯文本。 - # 每步同步会把“别推理、直接吐 skills”回灌采集端，而下一轮采集又要求它 thinking - # —— 两个分布互相拉扯，训练每次都赢。实测（output.e18.en）：4 步内 skill 长度 - # 281->573、出现 `Motor virtue` / `spectral misfire` 这类退化文本，candidate_pass_rate - # 从 0.801 跌到 0.404、train/lift 转负（-0.055）。采集固定用初始权重能切断这个回路。 - # 代价：训练对采集零反馈，本质上退化成“离线数据生成 + 独立 SFT”。eval 仍用最新 - # 权重，所以 eval/lift 依旧反映 skill_model 的真实进步。 - - metrics = {'train/n_samples': float(n)} - # ⭐ 必须用 float() 尝试转换而不是 isinstance 判数值型：twinkle 的 LossMetric.calculate() - # 把 loss / grad_norm 格式化成**字符串**后才返回（`f'{avg_loss:.4f}'`），用 - # isinstance(val, (int, float)) 会把这两个最关键的优化指标静默丢弃 —— 且丢在写文件之前， - # 所以 train_log / swanlab / 日志里全部看不到，事后也无法找回。 - # 转不成的（'total time elapse'='12.3 minutes'、'speed'='1.2 iters/s'）才跳过。 - for k, val in (rt.skill_model.calculate_metric(is_training=True) or {}).items(): - if isinstance(val, bool): - continue - try: - fval = float(val) - except (TypeError, ValueError): - continue - if k.startswith('learning rate'): - if 'group 1' in k: - metrics['train/lr'] = fval - elif k.startswith('train/'): - metrics[k.replace(' ', '_')] = fval - else: - metrics[f'train/{k.replace(" ", "_")}'] = fval - return 1, metrics - - -# =========================================================================== -# eval：query-only（部署口径） -# =========================================================================== -def _sync_trained_to_sampler(rt: Runtime) -> None: - """把当前（已训练）权重临时推给 skill_sampler，供 eval 使用。 - - 只是「临时」：eval 一结束就由 _restore_base_weights 把初始权重灌回去。skill_model 侧 - 不落盘、不 load，训练权重与优化器状态全程不受影响。 - """ - rt.ckpt.sync_weights(merge_and_sync=True) - rt.skill_sampler.reset_prefix_cache() - - -def _restore_base_weights(rt: Runtime) -> None: - """把 skill_sampler 恢复到**初始**权重，供下一轮采集使用。 - - ⭐ 走 `skill_sampler.load_weights_from_path()`（不传参 = sampler 自己的 model_id，即原始 - 预训练权重），从磁盘直接流进 vLLM。关键是它**完全不碰 skill_model**： - 不 save、不 load，训练权重和 AdamW 动量都不受影响。 - - 对比曾经考虑过的「save 训练权重 -> skill_model.load(初始) -> sync -> load 回训练权重」： - 那条路要在训练模型上来回 load 两次，一旦中途失败（OOM / 磁盘满 / 进程被杀），训练端就 - 停在初始权重上却带着原来的优化器状态继续跑，训练成果被静默清零且日志上看不出来。 - 现在最坏情况只是采集端权重不对（下一次 eval 前的 sync 会覆盖掉），训练端不可能被破坏。 - - reset_prefix_cache 必须跟着走：prefix cache 里缓存的是旧权重算出的 KV，换完权重不清就会 - 拿旧 KV 拼新权重的输出。 - """ - rt.skill_sampler.load_weights_from_path() - rt.skill_sampler.reset_prefix_cache() - - -def run_eval(rt: Runtime, eval_records, base_cache: Dict[str, Dict[str, Any]], - updates: int = 0) -> Dict[str, float]: - """部署口径 eval：**nothink + SKILLGEN_SYSTEM_EVAL + 不给 rubric**，与训练分布逐字一致。 - - ⭐ skill_sampler 建立时是 enable_thinking=True（采集需要 thinking 多样性），而训练/部署是 - nothink。所以 eval 必须先把同一个引擎的**客户端模板**临时切成 nothink，跑完再切回； - 只换编码模板，引擎本身不动（与 skill_ablate/trainer.py 的做法同源）。 - 不切的后果：eval 会带着 thinking 布局生成，与训练的 nothink 布局错配 —— 测出来的数不是 - 部署时的真实能力。 - - baseline（executor 冻结）整个 run 只算一次，之后走 base_cache（缓存的是 **pass_rate**）。 - eval 也跑 EXEC_ROLLOUTS 次取通过率：与采集口径一致，否则 lift 不可比。 - - ⭐ 权重由调用方负责：进来之前已 sync 成最新训练权重，出去之后立刻恢复初始权重 - （见 main 里的 _sync_trained_to_sampler / _restore_base_weights）。本函数只管跑分。 - """ - todo = [r for r in eval_records if r['data_id'] not in base_cache] - for r, rr in zip(todo, bare_solve(rt, todo) if todo else []): - base_cache[r['data_id']] = _pass_rate(rr) - base_rates = [base_cache[r['data_id']] for r in eval_records] - base_acc = _mean(base_rates) - - # skill 模型：nothink + eval 系统提示、greedy 出一条 skill - rt.skill_sampler.set_template(Template, model_id=MODEL_ID, enable_thinking=False, - max_length=MAX_MODEL_LEN) - try: - sg = run_samples(rt.skill_sampler, - [skillgen_prompt(r['problem'], '', eval=True) for r in eval_records], - 1, SKILL_MAX_TOKENS, SKILL_SAMPLER_GPUS, temperature=0.0) - finally: - # 必须切回：下一个 chunk 的采集要 thinking。放 finally 里是为了 eval 中途报错也不会 - # 把采集口径永久卡在 nothink 上。 - rt.skill_sampler.set_template(Template, model_id=MODEL_ID, enable_thinking=True, - max_length=MAX_MODEL_LEN) - skills = [extract_skill(seq_text(first_seq(s))) for s in sg] - n_parsed = sum(1 for s in skills if s) - # executor 也跑 EXEC_ROLLOUTS 次，与 baseline / 采集同口径 - ws = run_samples(rt.base_sampler, - [skill_solve_prompt(r['problem'], s) for r, s in zip(eval_records, skills)], - EXEC_ROLLOUTS, EXEC_MAX_TOKENS, BASE_SAMPLER_GPUS, - temperature=EXEC_TEMPERATURE, top_p=EXEC_TOP_P) - pairs, spans = [], [] - for r, seqs in zip(eval_records, ws): - seqs = list(seqs or []) - spans.append(len(seqs)) - pairs.extend((s, r['reference_answer']) for s in seqs) - judged = judge_seqs(pairs) - rates, i = [], 0 - for n in spans: - rates.append(_pass_rate(judged[i:i + n])) - i += n - acc = _mean(rates) - # ⭐ skill 原文落盘：聚合指标看不出「写成了什么样」，lift 为负时必须能回到原文归因。 - dump_eval_skills(eval_records, skills, base_rates, rates, updates) - # 难题 = baseline 没能每次都对（pass_rate < 1）；rescue 改成平均 pass_rate 增量。 - hard = [(b, w) for b, w in zip(base_rates, rates) if b < 1.0] - return { - 'eval/accuracy': acc, - 'eval/baseline_accuracy': base_acc, - 'eval/lift': acc - base_acc, - 'eval/format_rate': n_parsed / max(1, len(eval_records)), - 'eval/hard_rescue_rate': (_mean([w - b for b, w in hard]) if hard else 0.0), - 'eval/n_hard': float(len(hard)), - 'eval/skill_length_characters': _mean([float(len(s)) for s in skills]), - } - - -def dump_eval_skills(eval_records, skills, base_rates, rates, updates) -> None: - """eval 产出的 skill 原文落盘（append-only）。 - - 为何必需：`run_eval` 以前只回传聚合指标，`skills` 是局部变量、用完即弃，而 - SAVE_EVERY_UPDATES 默认 50，所以 lift 为负时既看不到 skill 原文、也没有 ckpt 可以重跑 - —— 无法区分「内容写得差」和「注入方式不对」。这份产物让部署口径可离线归因。 - - 逐题存 base/with pass_rate，所以可以直接筛出被 skill 带坏的题（with < base）。 - """ - path = os.path.join(OUTPUT_DIR, 'eval_skills.jsonl') - with open(path, 'a', encoding='utf-8') as f: - for r, s, b, w in zip(eval_records, skills, base_rates, rates): - f.write(json.dumps({ - 'updates': updates, - 'data_id': r.get('data_id'), - 'task_id': r['reference_answer'].get('task_id'), - 'base_pass_rate': b, - 'with_pass_rate': w, - 'pass_gain': w - b, - 'skill_chars': len(s), - 'problem': r['problem'], - 'skills': s, - }, ensure_ascii=False) + '\n') - - -# =========================================================================== -# main -# =========================================================================== -def build_runtime(checker, rubric_cache) -> Runtime: - """三组卡：train / skill_sampler / base_sampler(executor)。SFT 不需要 ref 模型。""" - r0 = TRAIN_GPUS - r1 = r0 + SKILL_SAMPLER_GPUS - twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, lazy_collect=False, groups=[ - DeviceGroup(name='train', ranks=list(range(0, r0)), device_type='GPU'), - DeviceGroup(name='skill_sampler', ranks=list(range(r0, r1)), device_type='GPU'), - DeviceGroup(name='base_sampler', ranks=list(range(r1, NUM_GPUS)), device_type='GPU')]) - - skill_model = TransformersModel( - model_id=MODEL_ID, remote_group='train', - device_mesh=DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_DP, - fsdp_size=TRAIN_FSDP), - ddp_config={'find_unused_parameters': False}, torch_dtype='float32') - skill_model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) - # ⭐ 训练模板 enable_thinking=False：训练/部署都是 nothink，必须一致。 - # 采集（skill_sampler）才是 thinking，那只用来邀出候选，不是训练分布。 - # 反例：改成 True 会把 `\n<think>\n\n</think>` 那 4 个固定 token 也纳入可训区，而 run_eval - # 是 nothink 生成的 —— 两边可训/生成区对不上，就不再是训推一致。 - skill_model.set_template(Template, model_id=MODEL_ID, enable_thinking=False, - max_length=MAX_MODEL_LEN, truncation_strategy='delete') - skill_model.set_processor(InputProcessor, padding_free=False) - # ⭐ 纯 SFT：交叉熵，不走 GRPO 那一套。 - # CrossEntropyLoss 只看 inputs['labels']（即 key_rounds 圈出的 assistant 段），没有 ratio / - # clip / KL / advantage，也不需要 old_logps、ref 模型。拒绝采样的“选择”已经全部发生在 - # 三道筛里（胜者才入池），到了 loss 这一层就是普通的“拟合这条目标文本”，不应再有 RL 项。 - # reduction='mean'（默认）-> num_tokens=0，每个 micro 自己 token-mean，梯度按 micro 数归一。 - skill_model.set_loss('CrossEntropyLoss') - skill_model.set_optimizer('AdamW', lr=LR) - - def _sampler(group, world, enable_thinking): - s = vLLMSampler(model_id=MODEL_ID, - engine_args={'gpu_memory_utilization': GPU_MEM, - 'max_model_len': MAX_MODEL_LEN, - 'tensor_parallel_size': 1}, - device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), - remote_group=group) - s.set_template(Template, model_id=MODEL_ID, enable_thinking=enable_thinking, - max_length=MAX_MODEL_LEN) - return s - - # skill_sampler 开 think：采集要多样性（拒绝采样的前提）。eval 时同一引擎跑 query-only。 - skill_sampler = _sampler('skill_sampler', SKILL_SAMPLER_GPUS, enable_thinking=True) - # executor 关 think（与 E23 一致）：开 think 有大量 rollout 撞满预算连代码块都写不出来。 - base_sampler = _sampler('base_sampler', BASE_SAMPLER_GPUS, enable_thinking=True) - return Runtime( - skill_model=skill_model, skill_sampler=skill_sampler, base_sampler=base_sampler, - ckpt=CheckpointEngineManager(model=skill_model, sampler=skill_sampler), - checker=checker, rubric_cache=rubric_cache) - - -def swan_init(): - if swanlab is None or os.environ.get('SWANLAB_MODE') == 'disabled': - logger.info('[swanlab] 未启用，只写 train_log.jsonl') - return None - name = 'E18_rejection_sft_code' + (f'_{RUN_TAG}' if RUN_TAG else '') + f'_{RUN_ID}' - swanlab.init(project=SWAN_PROJ, experiment_name=name, config={ - 'model_id': MODEL_ID, 'seed': SEED, 'chunk_size': CHUNK_SIZE, 'n_skills': N_SKILLS, - 'accumulate': ACCUMULATE, 'max_updates': MAX_UPDATES, 'lr': LR, - 'loss': 'CrossEntropyLoss', - 'exec_rollouts': EXEC_ROLLOUTS, 'exec_temperature': EXEC_TEMPERATURE, - 'exec_top_p': EXEC_TOP_P, 'min_gain_rollouts': MIN_GAIN_ROLLOUTS, - 'min_pass_gain': MIN_PASS_GAIN, - 'skill_char_limit': SKILL_CHAR_LIMIT, 'skill_max_tokens': SKILL_MAX_TOKENS, - 'exec_max_tokens': EXEC_MAX_TOKENS, 'eval_size': EVAL_SIZE, - 'skill_gen_temperature': SKILL_GEN_TEMPERATURE, - 'run_tag': RUN_TAG, 'run_id': RUN_ID, 'output_dir': OUTPUT_DIR}) - logger.info(f'[swanlab] project={SWAN_PROJ} experiment={name}') - return swanlab - - -def swan_log(swan, row: Dict[str, Any], step: int) -> None: - if swan is None: - return - m = {k: float(v) for k, v in row.items() - if isinstance(v, (int, float)) and not isinstance(v, bool)} - try: - swan.log(m, step=step) - except Exception as e: - logger.warning(f'[swanlab] log 失败（已忽略）：{e}') - - -# 不参与污染判定的文件：跑一次要花大量 CPU/沙箱时间（~900 道题跑参考解答）， -# 且内容只依赖题池、与哪个 run 无关，所以要从旧目录携带到新目录。 -_CARRY_OVER = ('bcb_broken_tasks.json', ) - - -def archive_output_dir() -> None: - """启动时把已存在的 OUTPUT_DIR 整个 mv 走，保证本 run 写入空目录。 - - 为何必需：`e18_sft_dataset.jsonl` / `train_log.jsonl` 都是 `open(..., 'a')` 追写。 - 没有这一步时，重启一次就把新旧 run 的样本焊在同一个文件里，而且不报错。 - 实测后果（output.e18.en）：崩溃 run 的 602 条（硬缺陷 41.4%：ttr<0.45 有 170 条、 - 超长 125 条、重复 118 条）与新 run 的 274 条混在一起，旧数据占 69%。 - 本 run 的训练不受影响（训练吃的是内存里的 pool，这个文件只写不读），但它的存在 - 意义就是“不重跑 GPU 就能换 SFT 超参再训一遍”（见 dump_dataset）—— 那个场景下 - 污染会直接进训练，且无任何报错，只是模型更差。 - - 用 mv 而不是删：旧 run 的诊断数据（eval 曲线、崩溃样本）是可复用的分析素材， - 丢了就得重跑 GPU 才能拿回。 - """ - if not os.path.isdir(OUTPUT_DIR): - return - if not os.listdir(OUTPUT_DIR): # 空目录直接用，不制造无意义的归档 - return - # 归档名带旧目录的 mtime（而不是当前时间）：同一批旧数据无论何时重启都归到 - # 同一个名字上，看名字就知道里面是哪段时间的 run。 - stamp = time.strftime('%m%d-%H%M%S', time.localtime(os.path.getmtime(OUTPUT_DIR))) - dst = f'{OUTPUT_DIR}.bak-{stamp}' - n = 1 - while os.path.exists(dst): # 同一秒重启两次也不能覆盖已有归档 - dst = f'{OUTPUT_DIR}.bak-{stamp}.{n}' - n += 1 - shutil.move(OUTPUT_DIR, dst) - os.makedirs(OUTPUT_DIR, exist_ok=True) - carried = [] - for name in _CARRY_OVER: - src = os.path.join(dst, name) - if os.path.exists(src): - shutil.copy2(src, os.path.join(OUTPUT_DIR, name)) - carried.append(name) - logger.warning(f'[output] 已存在的 {OUTPUT_DIR} 已归档到 {dst}' - + (f'（携带缓存：{", ".join(carried)}）' if carried else '')) - - -def main(): - archive_output_dir() - os.makedirs(OUTPUT_DIR, exist_ok=True) - if ACCUMULATE % TRAIN_DP != 0: - raise ValueError(f'ACCUMULATE({ACCUMULATE}) 必须是 TRAIN_DP({TRAIN_DP}) 的整数倍，' - f'否则 dp 切分会丢掉尾部样本') - checker = build_checker() - if checker is None: - raise RuntimeError('没有教师 API（LLM_BACKUP_API_KEY / LLM_BACKUP_BASE_URL / ' - 'OPENAI_API_KEY 都没设）；rubric 是三道筛的参照系，无法降级运行') - train_dataset, eval_records = load_records(SEED, EVAL_SIZE, OUTPUT_DIR) - if len(train_dataset) < CHUNK_SIZE: - raise ValueError(f'训练池 {len(train_dataset)} 小于 CHUNK_SIZE {CHUNK_SIZE}') - logger.info(f'[data] train={len(train_dataset)} eval={len(eval_records)}') - - rt = build_runtime(checker, MultiDiagCache()) - swan = swan_init() - base_cache: Dict[str, Dict[str, Any]] = {} - pool: List[Dict[str, Any]] = [] # 攒够 ACCUMULATE 条就 SFT 一次 - updates, ci, si, epoch, last_eval = 0, 0, 0, 0, 0 - logger.info(f'E18 start: lr={LR} chunk={CHUNK_SIZE} n_skills={N_SKILLS} ' - f'accumulate={ACCUMULATE} max_updates={MAX_UPDATES} output={OUTPUT_DIR}') - - with open(os.path.join(OUTPUT_DIR, 'train_log.jsonl'), 'a', encoding='utf-8') as log_fh: - # ⭐ updates=0 的 baseline eval：没有这一点，`eval/*` 曲线的第一个数据点已经是训了 - # EVAL_EVERY_UPDATES 步之后的值，无法区分「训练带来的提升」和「初始就有的能力」。 - # 注意 base_cache 缓存的是 executor 裸解 pass_rate（executor 的基线），不是 skill 模型 - # 的基线 —— 两回事，不能互代。 - # 副作用：这次 eval 会把裸解结果写进 base_cache，所以后续 eval 能直接复用， - # 整体 GPU 成本并非净增一整次 eval。 - if eval_records: - row0 = {'step': 0, 'chunk': 0, 'updates': 0, 'epoch': 0, - 'signal/pool_size': 0.0} - row0.update(run_eval(rt, eval_records, base_cache, updates=0)) - row0['eval/updates_done'] = 0 - log_fh.write(json.dumps(row0, ensure_ascii=False) + '\n') - log_fh.flush() - swan_log(swan, row0, si) - logger.info('[baseline u0] ' - + ' '.join(f'{k}={v:.4g}' for k, v in row0.items() - if isinstance(v, float))) - si += 1 - - while updates < MAX_UPDATES: - loader = DataLoader(dataset=train_dataset, batch_size=CHUNK_SIZE, num_workers=0, - shuffle=True, drop_last=True, - generator=torch.Generator().manual_seed(SEED + epoch)) - for chunk in loader: - if updates >= MAX_UPDATES: - break - t0 = time.time() - accepted, metrics = collect_chunk(rt, chunk, ci) - dump_dataset(accepted) - pool.extend(accepted) - - n_upd, tmetrics = 0, {} - while len(pool) >= ACCUMULATE and updates < MAX_UPDATES: - batch, pool = pool[:ACCUMULATE], pool[ACCUMULATE:] - k, tm = train_batch(rt, batch) - n_upd += k - updates += k - tmetrics = tm or tmetrics - - row = {'step': si, 'chunk': ci, 'updates': updates, 'epoch': epoch, - 'seconds': round(time.time() - t0, 1), - 'signal/pool_size': float(len(pool)), - **metrics, **tmetrics} - if eval_records and (updates - last_eval >= EVAL_EVERY_UPDATES - or updates >= MAX_UPDATES) and updates > 0: - # ⭐ eval 前把训练权重临时推给 sampler，跑完立刻把初始权重灌回去。 - # train_batch 已不再逐步同步，所以采集阶段永远是初始模型； - # 只有这一段区间内 sampler 带的是训练权重。finally 保证 eval 中途 - # 报错也不会把退化权重永久留在采集端。 - _sync_trained_to_sampler(rt) - try: - row.update(run_eval(rt, eval_records, base_cache, - updates=updates)) - finally: - _restore_base_weights(rt) - row['eval/updates_done'] = updates - last_eval = updates - log_fh.write(json.dumps(row, ensure_ascii=False) + '\n') - log_fh.flush() - swan_log(swan, row, si) - logger.info(f'[s{si} c{ci} u{updates}/{MAX_UPDATES}] ' - + ' '.join(f'{k}={v:.4g}' for k, v in row.items() - if isinstance(v, float))) - if SAVE_EVERY_UPDATES and updates and updates % SAVE_EVERY_UPDATES == 0: - rt.skill_model.save(f'E18-u{updates}', output_dir=OUTPUT_DIR) - si += 1 - ci += 1 - epoch += 1 - - rt.skill_model.save('E18-final', output_dir=OUTPUT_DIR) - rt.rubric_cache.close() - if swan is not None: - swan.finish() - logger.info(f'done: updates={updates} chunks={ci} -> {OUTPUT_DIR}/E18-final') - - -if __name__ == '__main__': - main() diff --git a/cookbook/human_e18/e18_select.py b/cookbook/human_e18/e18_select.py deleted file mode 100644 index b2d3a28d5..000000000 --- a/cookbook/human_e18/e18_select.py +++ /dev/null @@ -1,142 +0,0 @@ -"""E18 的拒绝筛：增量达阈 -> 不超长 -> pass_rate 最大（并列内长度/rubric 相似度拆平局）。 - -这是本臂**唯一的自变量**。E12 靠「rubric 重生成 + 2-in-8 验证」入池，没有排序；E18 在做对的 -候选里再排一次序，只留唯一胜者，并把打分全过程写进数据集文件供事后审计。 - -`_rubric_similarity` 与 skill_ablate/methods.py 逐字同源（含那条踩坑注释），搬过来是为了让 -本目录自包含、不再 import 那棵 2 万行的 methods.py。同源搬过来的泄漏门已删，理由见下。 -""" -import re -from collections import Counter -from typing import Any, Dict, List, Optional - -# --- 第二道筛已删：泄漏检测对 coding 任务恒为 False ------------------------------------ -# ⭐ 原 `answer_leaked` / `leak_blocks` 从 skill_ablate/methods.py 搬来，但那边是**数学**任务： -# `reference_answer` 是 `'42'` 这样的答案字符串，`str(ref) in skill` 的子串匹配有意义。 -# BCB 的 `reference_answer` 是个 ~2500 字符的 dict（task_id / entry_point / test / code_prompt / -# doc_struct / canonical_solution），该匹配要求 skill 逐字包含整个 dict 的 Python repr -# （含 `{'task_id': 'BigCodeBench/598', ...}`）—— 而 skill 上限 SKILL_CHAR_LIMIT=1500 字符， -# 物理上装不进。实测：只有把整个 dict 原样粘进去才判 True，截断 50 字符加前缀就已为 False， -# 触发概率恒为 0。 -# 真正的泄漏通道在 `test`（隐藏单测的断言期望值）与 `canonical_solution`，需要另写检测； -# 留着一个恒 False 的门只会给人「已经防住了」的假安全感，所以整体删除。 - - -# --- 第二道筛：skill 与 rubric 诊断的词频余弦 ---------------------------------------------- -# 定位：在「executor 已做对」的候选里挑与诊断内容对得上的那条，压掉两类假赢家 —— -# 与诊断无关的碰巧做对（含泄漏式速通的残余），以及与谁都不像的空泛套话。 -# 刻意用去停用词的词频余弦而不是 tfidf/语义模型：可迁移性判别器一节实测「仅 tfidf」 -# in-sample 0.983 / OOS 0.541 是纯过拟合；词频余弦纯 stdlib、确定性、可离线复算。 -_SIM_STOPWORDS = frozenset( - 'the a an and or of to in is are be for with that this it on as by from at not no was ' - 'were will would can could should may might do does did have has had you your we they ' - 'he she its if then than so but into over under out up down when where which what how ' - 'why all any each more most other some such only own same very'.split()) -_SIM_WORD_RE = re.compile(r"[a-z][a-z'\-]{2,}") - - -def rubric_similarity(skill: str, rubric: str) -> float: - """内容词词频余弦 ∈ [0,1]；任一侧无内容词返回 0。""" - ca = Counter(w for w in _SIM_WORD_RE.findall((skill or '').lower()) - if w not in _SIM_STOPWORDS) - cb = Counter(w for w in _SIM_WORD_RE.findall((rubric or '').lower()) - if w not in _SIM_STOPWORDS) - if not ca or not cb: - return 0.0 - dot = float(sum(v * cb[k] for k, v in ca.items() if k in cb)) - na = sum(v * v for v in ca.values()) ** 0.5 - nb = sum(v * v for v in cb.values()) ** 0.5 - return dot / (na * nb) if na and nb else 0.0 - - -# --- 两道筛合体 ----------------------------------------------------------------------------- -def select_winner(cands: List[Dict[str, Any]], rubric: str, reference: Any = None, *, - len_budget: int = 0, skill_char_limit: int, - base_pass_rate: float = 0.0, - min_pass_gain: float = 0.0) -> Optional[Dict[str, Any]]: - """按 **pass_rate 增量**选胜者；增量不达 min_pass_gain 就返回 None。 - - ⭐ 为何不能用旧的「做对就入池」（M=1/T=0）：那时 with_pass 只有 0/1，易题上「加任何 - skill 都对」，8 个候选全部 with_pass=1 —— 此时「挑哪条」完全由长度决定（易题 rubric 为空， - 相似度恒 0），等于往数据集里灌随机 skill。现在 executor 每个 skill 跑 M=8 次， - with_pass 是连续通过率，同一题的不同 skill 之间才有方差（如 8/8 vs 5/8）。 - - 筛选顺序： - a. **增量达阈**：`with_pass >= base_pass_rate + min_pass_gain`。 - min_pass_gain>0 时，仅仅「没弄坏」（tie，如 8/8 -> 8/8）**不够格**：那种样本对 - 「学会写有效 skill」没有监督信号。拉低通过率的更是直接淘汰。 - b. 超 skill_char_limit 过滤。（原本还有一道泄漏门，已删 —— 它在 coding 任务上恒为 False， - 详见文件头部注释。） - c. 取 **pass_rate 最大**的一档（允许并列）；并列内部按 **rubric 相似度**取高。 - - ⭐ c 的层次顺序很关键：pass_rate 是**客观效果**，rubric 相似度是**内容对齐**， - 长度只是**形式偏好**。长度已彻底移出择优路径（只在相似度也并列时做确定性拆平， - 取较短者）：旧版先按 `abs(len - len_budget)` 砍掉一半候选，而实测 66% 的题 8 个候选全部 - with_pass=1.0，于是长度成了事实上的唯一决策依据 —— 它把信息量大的长候选系统性淘汰。 - - ⭐ base_pass_rate 接近 1 时 a 几乎不可满足（天花板效应）：8/8 的题永远拿不到 +2/8， - 因此会被成建制排除在训练集外 —— 这是调用方想要的行为（只训真正有提升空间的题）， - 但意味着池子会偏向中等难度题，看 `signal/n_accepted_easy` 确认。 - - 胜者会被标上 `pass_gain`（= with_pass - base_pass_rate）与 `gain_kind`： - * 'improve'：pass_gain > 0；'tie'：== 0（仅当 min_pass_gain=0 时才可能返回）。 - - rubric 为空（易题无诊断）时相似度恒 0，并列内退化成取较短者 —— 此时已经是 - 「效果完全相同」的候选，拿什么拆平局都不影响效果，只是个确定性要求。 - - `reference` 已不再使用（泄漏门删除后唯一的消费方消失），`len_budget` 同样已废弃； - 两个形参保留只为不打断现有调用方的写法。 - """ - need = base_pass_rate + min_pass_gain - ok = [c for c in cands - if c.get('parseable') and (c.get('with_pass') or 0.0) >= need - 1e-9] - survivors = [c for c in ok if len(c['skills']) <= skill_char_limit] - if not survivors: - return None - # a/c：先按客观效果取最大档，再在并列内按 rubric 相似度拆平局。 - top = max((c.get('with_pass') or 0.0) for c in survivors) - tied = [c for c in survivors if (c.get('with_pass') or 0.0) >= top - 1e-9] - for c in tied: - c['rubric_similarity'] = rubric_similarity(c['skills'], rubric) - # 长度只作**最后的确定性拆平**（相似度也并列时），不再参与择优： - # 实测 66% 的题 8 个候选全部 with_pass=1.0，此时 `abs(len - LEN_BUDGET)` 事实上成了唯一 - # 决策依据，而它有系统性偏差 —— 中文表达同样内容的字符数天然更少（357 vs 705），永远 - # 更贴近预算，于是「离 400 最近」被翻译成了「选中文模板」，把信息量大的长英文候选全部 - # 淘汰。长度是形式偏好，不该越过效果与内容对齐，故彻底移出择优路径。 - best = max(tied, key=lambda c: (c['rubric_similarity'], -len(c['skills']))) - best['kept'] = True - best['pass_gain'] = round((best.get('with_pass') or 0.0) - base_pass_rate, 6) - best['gain_kind'] = 'improve' if best['pass_gain'] > 1e-9 else 'tie' - return best - - -def gain_stats(cands: List[Dict[str, Any]], base_pass_rate: float) -> Dict[str, int]: - """候选级增量计数（相对裸解 pass_rate），给 train_log 做监控。 - - `degraded` 是关键项：skill 把通过率拉低了。它完全不会反映在 accept_rate 上 - （那些候选只是默默被汰掉），持续偏高就是 skill-gen 在写有害提示的直接证据。 - """ - out = {'improved': 0, 'tied': 0, 'degraded': 0} - for c in cands: - if not c.get('parseable'): - continue - wp = c.get('with_pass') - if wp is None: - continue - if wp > base_pass_rate + 1e-9: - out['improved'] += 1 - elif wp < base_pass_rate - 1e-9: - out['degraded'] += 1 - else: - out['tied'] += 1 - return out - - -def filter_stats(passers: int, survivors: int) -> Dict[str, float]: - """第一道筛后的超长丢弃率，进 train_log。 - - 指标名保留 `leak_or_overlength_dropped_fraction` 不改：泄漏门删除前它的触发率恒为 0， - 所以新旧 run 的这个数值本来就完全可比（一直只在统计超长），改名反而会断掉曲线。 - """ - return {'train/leak_or_overlength_dropped_fraction': - ((passers - survivors) / passers) if passers else 0.0} diff --git a/cookbook/human_e18/e18_sft_kod.py b/cookbook/human_e18/e18_sft_kod.py deleted file mode 100644 index 6f0c4d1ea..000000000 --- a/cookbook/human_e18/e18_sft_kod.py +++ /dev/null @@ -1,543 +0,0 @@ -# -*- coding: utf-8 -*- -"""E18-KOD 离线 SFT：把 `e18_collect_kod.py` 采到的胜者用 **nothink 口径**训一遍，并在 -首尾各跑一次 eval，验证训练是否真的带来提升。 - -与在线版 `e18_rejection_sft.py` 的关系：**保留训练 + eval，去掉采集**。rubric 诊断、 -拒绝采样、chunk 循环全部移除 —— 数据已经在 `e18_sft_dataset.jsonl` 里落好了。 - -卡位：4 训练 + 2 skill_sampler + 2 base_sampler(executor)。为何不是 8 卡全训练： -eval 要 skillmodel 生成 skill、executor 跑代码，两者都需要 vLLM 引擎常驻。 - -训推一致的三个不可动点（与在线版逐字对齐，改任何一处就不再是同一个实验）： -1. **prompt 段用 `skillgen_prompt(..., eval=True)`** —— 即 SKILLGEN_SYSTEM_EVAL、不带 rubric。 - 采集时用的是带诊断的 thinking 口径（SKILLGEN_SYSTEM），那只是为了「邀出」好 skill； - 线上没有诊断可用，拿带诊断的 prompt 去训会学成「看着诊断改写」。 -2. **模板 enable_thinking=False** —— 训练/部署都是 nothink，必须一致。改成 True 会把 - think 那几个固定 token 也纳入可训区，与部署时的生成区对不上。 -3. **`key_rounds=[len(msgs)]`** 标出最后一轮（assistant）为唯一可训区，prompt 段全 -100。 - -loss 是纯 `CrossEntropyLoss`（twinkle 没有 SFTLoss 这个类）：拒绝采样的「选择」已经 -发生在采集侧的三道筛里（只有胜者入池），到 loss 这层就是普通的拟合目标文本， -不该再有 ratio / clip / KL / advantage。 - -⭐ eval 口径与在线版的**唯一差异**：executor 只跑 1 次且 temperature=0（在线版是 8 次 -T=0.6 取通过率）。这是按需求指定的 —— 省 GPU、且 greedy 单次可复现。代价写在 -run_eval 的注释里：pass_rate 退化成 0/1 二值，单题不可比，只能看 100 题的均值。 - -产物（OUTPUT_DIR 下）： -* `sft_log.jsonl`：逐 step 的 loss / grad_norm / lr，用来看收敛。 -* `eval_log.jsonl`：首尾两次 eval 的聚合指标。 -* `eval_skills.jsonl`：eval 生成的 skill 原文 + 逐题 base/with，用来归因。 -* `KODSFT-final/`：权重（SAVE_EVERY_STEPS>0 时还有 `KODSFT-s<step>/`）。 -""" -import json -import os -import random -import shutil -import sys -import time -from typing import Any, Dict, List, Tuple - -import torch -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.checkpoint_engine import CheckpointEngineManager -from twinkle.data_format import SamplingParams, pack_user_data -from twinkle.model import TransformersModel -from twinkle.patch.no_split_modules import NoSplitModulesPatch -from twinkle.processor import InputProcessor -from twinkle.sampler import vLLMSampler -from twinkle.template import Template - -_HERE = os.path.dirname(os.path.abspath(__file__)) -_COOKBOOK = os.path.abspath(os.path.join(_HERE, '..')) -for _p in (_HERE, os.path.join(_COOKBOOK, 'human')): - if _p not in sys.path: - sys.path.insert(0, _p) - -from e18_kodcode import (clean_text, empty_roll, extract_skill, # noqa: E402 - judge_seqs, load_records) -from e18_prompts import direct_prompt, skill_solve_prompt, skillgen_prompt # noqa: E402 - -logger = get_logger() - -# ========== 配置 ========== -MODEL_ID = os.environ.get('MODEL_ID', 'Qwen/Qwen3-4B') -DATA_PATH = os.environ.get( - 'DATA_PATH', os.path.join(_HERE, 'output.e18.kod', 'e18_sft_dataset.jsonl')) -OUTPUT_DIR = os.environ.get('OUTPUT_DIR', os.path.join(_HERE, 'output.e18.kod.sft')) - -# 8 卡分三组：4 训练 + 2 skillmodel 推理 + 2 executor。 -# ⭐ 为何 executor 只需 2 张（采集时是 6 张）：eval 只跑 100 题 × 1 次 = 100 个序列， -# 而采集是每 chunk 2304 个。序列数少两个量级，再加卡只会让训练侧变慢。 -TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 4)) -SKILL_SAMPLER_GPUS = int(os.environ.get('SKILL_SAMPLER_GPUS', 2)) -BASE_SAMPLER_GPUS = int(os.environ.get('BASE_SAMPLER_GPUS', 2)) -NUM_GPUS = TRAIN_GPUS + SKILL_SAMPLER_GPUS + BASE_SAMPLER_GPUS -TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', 1)) -TRAIN_DP = TRAIN_GPUS // TRAIN_FSDP -GPU_MEM = float(os.environ.get('GPU_MEM', 0.8)) - -SEED = int(os.environ.get('SEED', 42)) -EPOCHS = float(os.environ.get('EPOCHS', 3)) -# ⭐ 必须是 TRAIN_DP 的整倍数：dp 切分会把不足一轮的尾部丢掉。 -BATCH_SIZE = int(os.environ.get('BATCH_SIZE', 16)) -MICRO_BATCH = int(os.environ.get('MICRO_BATCH', 8)) -LR = float(os.environ.get('LR', 1e-5)) -MAX_MODEL_LEN = int(os.environ.get('MAX_MODEL_LEN', 16000)) - -# ---- eval 口径 ---- -EVAL_SIZE = int(os.environ.get('EVAL_SIZE', 100)) -SKILL_MAX_TOKENS = int(os.environ.get('SKILL_MAX_TOKENS', 8192)) -EXEC_MAX_TOKENS = int(os.environ.get('EXEC_MAX_TOKENS', 15000)) -# ⭐ executor 跑 1 次、temperature=0（按需求指定，与在线版的 8×T=0.6 不同）。 -# 后果：单题 pass_rate 只能是 0 或 1，所以**不要看单题差异**，只看 100 题均值； -# 也因此 hard_rescue 这类分层指标失去意义（不再计算）。greedy 的好处是可复现。 -EXEC_ROLLOUTS = int(os.environ.get('EXEC_ROLLOUTS', 1)) -EXEC_TEMPERATURE = float(os.environ.get('EXEC_TEMPERATURE', 0.0)) - -# 数据清洗门槛（见 load_samples 的注释，都是实测抽检出来的缺陷） -MIN_CHARS = int(os.environ.get('MIN_CHARS', 200)) -MAX_CHARS = int(os.environ.get('MAX_CHARS', 1500)) -DROP_CJK = os.environ.get('DROP_CJK', '1') == '1' - -SAVE_EVERY_STEPS = int(os.environ.get('SAVE_EVERY_STEPS', 0)) # 0 = 只在结束时存 -LOG_EVERY_STEPS = int(os.environ.get('LOG_EVERY_STEPS', 1)) -RUN_ID = time.strftime('%m%d-%H%M%S') - - -# =========================================================================== -# 数据 -# =========================================================================== -def load_samples() -> List[Dict[str, Any]]: - """读胜者并做格式清洗。返回 [{'problem', 'response'}]。 - - ⭐ 为何要在 SFT 侧再清一遍（采集侧已有 SKILL_CHAR_LIMIT）：抽检 1104 条实测出三类 - 残留缺陷，占比虽小但都会被模型逐字学走： - * CJK 混入 3/1104（'动态规划' 这种孤立中文词）—— 部署口径是纯英文，学走就成了双语输出； - * 残留 `<skills>` 标签 2/1104 —— response 外层已经由采集侧包了一层，内层再出现就是嵌套； - * 截断（无终止标点）7/1104 —— 学截断等于学「说半句就停」。 - 合计约 1.1%，宁可丢掉也不喂进去。 - - `response` 用采集时存的原字段（已是 `<skills>\\n...\\n</skills>` 包装），不重新拼： - 与在线版 `train_batch` 消费的是同一个字段，保持逐字一致。 - """ - if not os.path.exists(DATA_PATH): - raise FileNotFoundError(f'找不到数据集：{DATA_PATH}') - raw, bad_json = [], 0 - with open(DATA_PATH, 'r', encoding='utf-8') as f: - for line in f: - line = line.strip() - if not line: - continue - try: - raw.append(json.loads(line)) - except Exception: - # 采集进程被 kill 时最后一行可能是半行，跳过而不是让整个训练起不来。 - bad_json += 1 - drop = {'no_response': 0, 'too_short': 0, 'too_long': 0, 'cjk': 0, - 'nested_tag': 0, 'truncated': 0} - out = [] - for r in raw: - resp = (r.get('response') or '').strip() - skills = (r.get('skills') or '').strip() - if not resp or not skills: - drop['no_response'] += 1 - continue - if len(skills) < MIN_CHARS: - drop['too_short'] += 1 - continue - if len(skills) > MAX_CHARS: - drop['too_long'] += 1 - continue - if DROP_CJK and any('\u4e00' <= ch <= '\u9fff' for ch in skills): - drop['cjk'] += 1 - continue - if '<skills' in skills.lower() or '</skills' in skills.lower(): - drop['nested_tag'] += 1 - continue - if not skills.rstrip().endswith(('.', '!', '?', ':', ')')): - drop['truncated'] += 1 - continue - out.append({'problem': r['problem'], 'response': resp, - 'data_id': r.get('data_id', ''), 'run': r.get('run', '')}) - logger.info(f'[data] 读入 {len(raw)} 条' - + (f'（跳过 {bad_json} 行半行）' if bad_json else '') - + f'，清洗后 {len(out)} 条，丢弃明细 {drop}') - runs = sorted({s['run'] for s in out if s['run']}) - if len(runs) > 1: - # 续跑（KOD_RESUME=1）后同一份文件里会有多个 run，训练是全量混合 —— 这里只提示， - # 不自动过滤：要不要分开训是实验设计问题，不该由脚本替你决定。 - logger.warning(f'[data] 数据里含 {len(runs)} 个 run：{runs}（全部混合训练）') - return out - - -def load_eval_records(train_data_ids: set): - """eval 集：从 KodCode 题池里挑 EVAL_SIZE 道**没有生成过 skill** 的题。 - - ⭐ 必须排掉采集跑过的题，而且排的是 **candidates 里的全量 id**、不是仅胜者： - 采集时一道题跑了 4 个候选但只有 ~22% 能入池，剩下的题虽然不在训练集里，却已经 - 被用来挑过 skill —— 拿它当 eval 会高估（数据选择偏差：那些题本身就是「skill 救不了」 - 或「本来就全对」的）。所以传进来的 train_data_ids 应该来自 `e18_candidates.jsonl`。 - - load_records 的 seed 与采集侧一致，所以题池顺序可复现；取前 EVAL_SIZE 条未采过的。 - """ - ds, _ = load_records(SEED, 0, OUTPUT_DIR) - out = [] - for r in ds.dataset: - if r['data_id'] in train_data_ids: - continue - out.append(r) - if len(out) >= EVAL_SIZE: - break - logger.info(f'[eval] 选了 {len(out)} 道未采集过的题（排除已跑 {len(train_data_ids)} 题）') - return out - - -def collected_data_ids() -> set: - """采集阶段跑过的全部 data_id（含未入池的），用来从 eval 集里排除。 - - 优先读 `e18_candidates.jsonl`（全量）；没有它才退回 sft_dataset（仅胜者，覆盖不全， - 会让 eval 集混进采集过的题）。 - """ - ids = set() - cand = os.path.join(os.path.dirname(DATA_PATH), 'e18_candidates.jsonl') - src = cand if os.path.exists(cand) else DATA_PATH - with open(src, 'r', encoding='utf-8') as f: - for line in f: - line = line.strip() - if not line: - continue - try: - ids.add(str(json.loads(line)['data_id'])) - except Exception: - continue - logger.info(f'[eval] 排除源 {os.path.basename(src)}：{len(ids)} 道已跑题') - return ids - - -# =========================================================================== -# 训练 -# =========================================================================== -def build_model(): - """三组卡：train(4) / skill_sampler(2) / base_sampler(2)。 - - 返回 (model, skill_sampler, base_sampler, ckpt)。训练组配置与在线版 build_runtime 一致。 - - ⭐ 两个 sampler 都直接建成 **enable_thinking=False**：本脚本没有采集阶段，不需要 - thinking 多样性，而 eval 就是部署口径（nothink）。因此不需要像在线版 run_eval 那样 - 每次 eval 前后临时切模板再切回 —— 少一个可能切错的状态。 - - ⭐ ckpt 只绑 skill_sampler：训练的是 skillmodel，executor 必须全程冻结， - 否则首尾两次 eval 的 baseline 不可比（分母都变了就无法归因给 skill）。 - """ - r0, r1 = TRAIN_GPUS, TRAIN_GPUS + SKILL_SAMPLER_GPUS - twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, lazy_collect=False, groups=[ - DeviceGroup(name='train', ranks=list(range(0, r0)), device_type='GPU'), - DeviceGroup(name='skill_sampler', ranks=list(range(r0, r1)), device_type='GPU'), - DeviceGroup(name='base_sampler', ranks=list(range(r1, NUM_GPUS)), device_type='GPU')]) - model = TransformersModel( - model_id=MODEL_ID, remote_group='train', - device_mesh=DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_DP, - fsdp_size=TRAIN_FSDP), - ddp_config={'find_unused_parameters': False}, torch_dtype='float32') - model.apply_patch(NoSplitModulesPatch({'Qwen3DecoderLayer'})) - # ⭐ enable_thinking=False：训练/部署都是 nothink，必须一致（见文件头注释第 2 点）。 - model.set_template(Template, model_id=MODEL_ID, enable_thinking=False, - max_length=MAX_MODEL_LEN, truncation_strategy='delete') - model.set_processor(InputProcessor, padding_free=False) - model.set_loss('CrossEntropyLoss') - model.set_optimizer('AdamW', lr=LR) - - def _mk_sampler(group: str, world: int): - s = vLLMSampler( - model_id=MODEL_ID, remote_group=group, - device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), - engine_args={'gpu_memory_utilization': GPU_MEM, - 'max_model_len': MAX_MODEL_LEN, 'tensor_parallel_size': 1}) - s.set_template(Template, model_id=MODEL_ID, enable_thinking=False, - max_length=MAX_MODEL_LEN) - return s - - skill_sampler = _mk_sampler('skill_sampler', SKILL_SAMPLER_GPUS) - base_sampler = _mk_sampler('base_sampler', BASE_SAMPLER_GPUS) - ckpt = CheckpointEngineManager(model=model, sampler=skill_sampler) - return model, skill_sampler, base_sampler, ckpt - - -# =========================================================================== -# eval：部署口径（nothink + SKILLGEN_SYSTEM_EVAL + 不给 rubric） -# =========================================================================== -def run_samples(sampler, prompts, num_samples, max_tokens, gen_dp, - temperature=None, top_p=None): - """采样。与 e18_collect_kod.run_samples 逐字一致（去掉本脚本用不到的 top_k/logprobs）。 - - 三处踩过的坑： - 1. 字段名是 `num_samples` 不是 `n`（SamplingParams 没有 `n`，传了直接 TypeError）。 - 2. 走 `sampler.sample(prompts, params)`，不是 pack_user_data + generate_sequences。 - 3. dp 补齐不能省：条数 < dp 会直接报错（eval 只 100 题、dp=2 虽然安全， - 但 EVAL_SIZE 调小到 1 时就会触发）。 - """ - if not prompts: - return [] - import copy - params = SamplingParams( - max_tokens=max_tokens, - temperature=0.6 if temperature is None else temperature, - top_p=0.95 if top_p is None else top_p, - num_samples=num_samples) - padded = prompts - if gen_dp > 1 and 0 < len(prompts) < gen_dp: - padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] - responses = sampler.sample(padded, params)[:len(prompts)] - return [list(r.sequences) if (r and r.sequences) else [] for r in responses] - - -def first_seq(seqs): - return seqs[0] if seqs else None - - -def seq_text(seq) -> str: - return clean_text(getattr(seq, 'decoded', '') or '') if seq is not None else '' - - -def _mean(xs) -> float: - xs = [float(x) for x in xs if x is not None] - return sum(xs) / len(xs) if xs else 0.0 - - -def _pass_rate(rolls) -> float: - return (sum(1.0 for x in rolls if x['correct']) / len(rolls)) if rolls else 0.0 - - -def _judge_batch(records, prompts, sampler, gen_dp) -> List[float]: - """跑 executor + 判分，返回逐题 pass_rate。 - - spans 记每题实际拿到几条序列（可能为 0），不能直接按 EXEC_ROLLOUTS 切 judged： - 采样失败的题会返回空列表，按固定步长切会整体错位。 - """ - ws = run_samples(sampler, prompts, EXEC_ROLLOUTS, EXEC_MAX_TOKENS, gen_dp, - temperature=EXEC_TEMPERATURE) - pairs, spans = [], [] - for r, seqs in zip(records, ws): - seqs = list(seqs or []) - spans.append(len(seqs)) - pairs.extend((s, r['reference_answer']) for s in seqs) - judged = judge_seqs(pairs) if pairs else [] - rates, i = [], 0 - for n in spans: - rates.append(_pass_rate(judged[i:i + n]) if n else 0.0) - i += n - return rates - - -def run_eval(skill_sampler, base_sampler, eval_records, - base_cache: Dict[str, float], tag: str, step: int) -> Dict[str, float]: - """部署口径 eval：skillmodel 生成 1 个 skill -> executor 跑 1 次 -> 看提升。 - - 与训练分布逐字一致：nothink + SKILLGEN_SYSTEM_EVAL + 不给 rubric。 - - ⭐ baseline 整个 run 只算一次，之后走 base_cache：executor 权重全程冻结，同一道题的 - 裸解结果不会变（T=0 更是确定性的）。重算不仅浪费 GPU，还会因为 vLLM 的 - 非确定性引入假的 baseline 漂移，把本来该归因给 skill 的差异污染掉。 - - ⭐ 只看 `lift`（= accuracy - baseline）的**首尾差异**。因为 EXEC_ROLLOUTS=1 + T=0， - 单题 pass_rate 只能是 0/1，100 题的均值标准误差约 0.05 —— 所以 lift 变化小于 - 约 0.07 时不要当成真实效果（双样本差值的噪声更大）。这是 1 次 rollout 的固有代价。 - """ - t0 = time.time() - todo = [r for r in eval_records if r['data_id'] not in base_cache] - if todo: - rates = _judge_batch(todo, [direct_prompt(r['problem']) for r in todo], - base_sampler, BASE_SAMPLER_GPUS) - for r, rate in zip(todo, rates): - base_cache[r['data_id']] = rate - base_rates = [base_cache[r['data_id']] for r in eval_records] - - sg = run_samples(skill_sampler, - [skillgen_prompt(r['problem'], '', eval=True) for r in eval_records], - 1, SKILL_MAX_TOKENS, SKILL_SAMPLER_GPUS, temperature=0.0) - skills = [extract_skill(seq_text(first_seq(s))) for s in sg] - n_parsed = sum(1 for s in skills if s) - with_rates = _judge_batch( - eval_records, - [skill_solve_prompt(r['problem'], s) for r, s in zip(eval_records, skills)], - base_sampler, BASE_SAMPLER_GPUS) - - acc, base_acc = _mean(with_rates), _mean(base_rates) - m = {'eval/accuracy': acc, 'eval/baseline_accuracy': base_acc, - 'eval/lift': acc - base_acc, - 'eval/format_rate': n_parsed / max(1, len(eval_records)), - 'eval/skill_length_characters': _mean([float(len(s)) for s in skills]), - 'eval/n_improved': float(sum(1 for b, w in zip(base_rates, with_rates) if w > b)), - 'eval/n_hurt': float(sum(1 for b, w in zip(base_rates, with_rates) if w < b)), - 'eval/seconds': round(time.time() - t0, 1)} - # ⭐ skill 原文落盘：聚合指标看不出「写成了什么样」，lift 为负时必须能回到原文归因。 - with open(os.path.join(OUTPUT_DIR, 'eval_skills.jsonl'), 'a', encoding='utf-8') as f: - for r, s, b, w in zip(eval_records, skills, base_rates, with_rates): - f.write(json.dumps({'tag': tag, 'step': step, 'run': RUN_ID, - 'data_id': r['data_id'], 'base': b, 'with': w, - 'skill_chars': len(s), 'skill': s}, - ensure_ascii=False) + '\n') - return m - - -# =========================================================================== -# 训练 -# =========================================================================== - - -def make_trajs(batch) -> List[Dict[str, Any]]: - """样本 -> twinkle 轨迹。与在线版 `train_batch` 的构造逐字相同。""" - trajs = [] - for s in batch: - msgs = skillgen_prompt(s['problem'], '', eval=True)['messages'] - trajs.append({'messages': msgs + [{'role': 'assistant', 'content': s['response']}], - 'user_data': pack_user_data({'key_rounds': [len(msgs)]})}) - return trajs - - -def step_metrics(model) -> Dict[str, float]: - """取本 step 的优化指标。 - - ⭐ 必须用 float() 试转而不是 isinstance 判数值型：twinkle 的 LossMetric.calculate() - 把 loss / grad_norm 格式化成**字符串**后才返回（`f'{avg_loss:.4f}'`），用 - isinstance(val, (int, float)) 会把这两个最关键的指标静默丢弃 —— 而这个脚本唯一的目的 - 就是看收敛，丢了 loss 就什么都看不到了。 - 转不成的（'total time elapse'='12.3 minutes'）才跳过。 - """ - out: Dict[str, float] = {} - for k, val in (model.calculate_metric(is_training=True) or {}).items(): - if isinstance(val, bool): - continue - try: - fval = float(val) - except (TypeError, ValueError): - continue - if k.startswith('learning rate'): - if 'group 1' in k: - out['lr'] = fval - else: - out[k.replace(' ', '_')] = fval - return out - - -def archive_output_dir() -> None: - """启动时把已存在的 OUTPUT_DIR 整个 mv 走（`sft_log.jsonl` 是追写的）。 - - 与 e18_collect_kod 同一套机制：不这么做，重跑一次就把两条 loss 曲线焊在一个文件里， - 而且不报错 —— 看收敛时会看到一条莫名其妙回弹的曲线。 - """ - if not os.path.isdir(OUTPUT_DIR) or not os.listdir(OUTPUT_DIR): - os.makedirs(OUTPUT_DIR, exist_ok=True) - return - stamp = time.strftime('%m%d-%H%M%S', time.localtime(os.path.getmtime(OUTPUT_DIR))) - dst = f'{OUTPUT_DIR}.bak-{stamp}' - i = 1 - while os.path.exists(dst): - dst = f'{OUTPUT_DIR}.bak-{stamp}-{i}' - i += 1 - shutil.move(OUTPUT_DIR, dst) - os.makedirs(OUTPUT_DIR, exist_ok=True) - logger.info(f'[init] 旧输出目录已归档 -> {dst}') - - -def log_eval(m: Dict[str, float], tag: str, step: int) -> None: - with open(os.path.join(OUTPUT_DIR, 'eval_log.jsonl'), 'a', encoding='utf-8') as f: - f.write(json.dumps({'tag': tag, 'step': step, 'run': RUN_ID, **m}, - ensure_ascii=False) + '\n') - logger.info(f'[eval:{tag}] ' + ' '.join(f'{k}={v:.4g}' for k, v in m.items())) - - -def main(): - t0 = time.time() - archive_output_dir() - samples = load_samples() - if len(samples) < BATCH_SIZE: - raise RuntimeError(f'可用样本 {len(samples)} 条 < BATCH_SIZE {BATCH_SIZE}') - if BATCH_SIZE % TRAIN_DP: - raise RuntimeError(f'BATCH_SIZE({BATCH_SIZE}) 必须是 TRAIN_DP({TRAIN_DP}) 的整倍数') - # eval 题池先选好（纯 CPU），再开 GPU：选错了就不用白等 vLLM 启动的几分钟。 - eval_records = load_eval_records(collected_data_ids()) - if not eval_records: - raise RuntimeError('eval 集为空（题池里的题已全部被采集过？）') - - model, skill_sampler, base_sampler, ckpt = build_model() - steps_per_epoch = len(samples) // BATCH_SIZE - total_steps = int(steps_per_epoch * EPOCHS) - logger.info(f'E18-KOD-SFT start: n={len(samples)} bs={BATCH_SIZE} micro={MICRO_BATCH} ' - f'lr={LR} epochs={EPOCHS} steps/epoch={steps_per_epoch} ' - f'total_steps={total_steps} eval_n={len(eval_records)} ' - f'exec_rollouts={EXEC_ROLLOUTS}@T{EXEC_TEMPERATURE} ' - f'gpus={TRAIN_GPUS}+{SKILL_SAMPLER_GPUS}+{BASE_SAMPLER_GPUS} out={OUTPUT_DIR}') - - # ---- 首次 eval（step 0，未训练的初始权重）---- - # ⭐ 不用 sync：skill_sampler 刚建立，拿的就是 MODEL_ID 的原始权重，与训练端同源。 - base_cache: Dict[str, float] = {} - m_before = run_eval(skill_sampler, base_sampler, eval_records, base_cache, 'before', 0) - log_eval(m_before, 'before', 0) - - log_path = os.path.join(OUTPUT_DIR, 'sft_log.jsonl') - rng = random.Random(SEED) - step = 0 - with open(log_path, 'a', encoding='utf-8') as log_fh: - epoch = 0 - while step < total_steps: - order = list(range(len(samples))) - rng.shuffle(order) # 每个 epoch 重洗，种子固定所以可复现 - for bi in range(steps_per_epoch): - if step >= total_steps: - break - batch = [samples[j] for j in order[bi * BATCH_SIZE:(bi + 1) * BATCH_SIZE]] - trajs = make_trajs(batch) - micro = max(TRAIN_DP, min(MICRO_BATCH, len(trajs))) - t_step = time.time() - for i in range(0, len(trajs), micro): - model.forward_backward(inputs=trajs[i:i + micro]) - model.clip_grad_and_step() - step += 1 - row = {'step': step, 'epoch': epoch, 'run': RUN_ID, - 'n_samples': len(batch), 'seconds': round(time.time() - t_step, 2)} - row.update(step_metrics(model)) - log_fh.write(json.dumps(row, ensure_ascii=False) + '\n') - log_fh.flush() - if step % LOG_EVERY_STEPS == 0: - logger.info('[s%d/%d ep%d] ' % (step, total_steps, epoch) - + ' '.join(f'{k}={v:.4g}' for k, v in row.items() - if isinstance(v, float))) - if SAVE_EVERY_STEPS and step % SAVE_EVERY_STEPS == 0: - # ⭐ API 是 `model.save(tag, output_dir=)`，不是 save_checkpoint(path) - # —— 与在线版 e18_rejection_sft.py:762 的用法一致。 - model.save(f'KODSFT-s{step}', output_dir=OUTPUT_DIR) - logger.info(f'[save] KODSFT-s{step}') - epoch += 1 - - model.save('KODSFT-final', output_dir=OUTPUT_DIR) - - # ---- 尾次 eval（训练后权重）---- - # ⭐ 必须先 sync_weights 把训好的权重推给 skill_sampler，否则这次 eval 跑的还是初始 - # 权重 —— 两次结果几乎相同，看起来像「训了没效果」，而真因是权重根本没上去。 - # merge_and_sync=True：全参数训练需要先在 dp 间 merge 再推。 - # reset_prefix_cache 必须跟着走：prefix cache 里是旧权重算出的 KV，不清就会拿旧 KV - # 拼新权重的输出，得到一个既不是训前也不是训后的嵌合态。 - ckpt.sync_weights(merge_and_sync=True) - skill_sampler.reset_prefix_cache() - # base_cache 沿用：executor 全程未动，baseline 不需重算（也不应重算，见 run_eval）。 - m_after = run_eval(skill_sampler, base_sampler, eval_records, base_cache, 'after', step) - log_eval(m_after, 'after', step) - - d_lift = m_after['eval/lift'] - m_before['eval/lift'] - logger.info('[result] lift %.4f -> %.4f (Δ %+.4f) accuracy %.4f -> %.4f baseline %.4f' - % (m_before['eval/lift'], m_after['eval/lift'], d_lift, - m_before['eval/accuracy'], m_after['eval/accuracy'], - m_after['eval/baseline_accuracy'])) - # ⭐ 噪声底提醒：EXEC_ROLLOUTS=1 + T=0 下单题 pass_rate 是 0/1，n=100 的均值标准误约 - # 0.05，首尾差值的噪声更大。不把这句写进日志，很容易把 ±0.05 的漂动当成结论。 - if abs(d_lift) < 0.07: - logger.warning(f'[result] Δlift {d_lift:+.4f} 在噪声量级内（n={len(eval_records)}、' - f'rollout=1、T=0 时约 ±0.07），不足以断定有/无效果。') - logger.info(f'[done] steps={step} 用时 {(time.time() - t0) / 60:.1f} 分钟 -> {OUTPUT_DIR}') - - -if __name__ == '__main__': - main() diff --git a/cookbook/human_e18/e19_logp_select.py b/cookbook/human_e18/e19_logp_select.py deleted file mode 100644 index 08d311af4..000000000 --- a/cookbook/human_e18/e19_logp_select.py +++ /dev/null @@ -1,193 +0,0 @@ -# -*- coding: utf-8 -*- -"""E19：能否不靠 8 次 executor rollout 就选出最好的 skill？（RLT-style logp/熵 打分） - -映射自 skill_quality_analysis.md 第 18 节的「用法2（只搬 reward）」。核心问题： -现在选 skill 要给每个候选跑 8 次 executor + 判分（8 候选 = 64 次解码），太贵。 -能不能用 **frozen executor 上的一次 teacher-forcing forward**（无解码）替代？ - -流程（与 e18 采集逐字同源，**不用 GT 造 prompt**）： - 1. 裸解 BARE_ROLLOUTS 次 -> 取**失败的** trajectory（错误代码 + 报错） - 2. 失败 traj + rubric 一起喂给 skillmodel -> 生成 N_SKILLS 个 narrative skill - 3. 每个 skill 让 executor rollout EXEC_ROLLOUTS 次 -> with_pass_rate（**这是 ground truth， - 只用于离线评估选择器，不参与任何打分特征**） - 4. 同时对每个 skill 算 cheap 特征（一次 forward，见下）-> 落盘 - 5. 离线比：cheap selector 的命中率 vs 8-rollout oracle - -⭐ 为什么 GT 不进 prompt：用户明确要求。RLT 的 teacher 是**开卷**的（输入含标准答案）， -因此它必须靠 r_KL 压泄漏，且 teacher 不能直接部署。本实验的 skillmodel 保持**闭卷** -（输入只有题面 + 自己的失败 traj + rubric），GT 只在两个地方出现： - (a) 判分（pytest 跑测试）—— 本来就在判分侧，不进生成器； - (b) r_SS 的打分目标 —— 只流经 reward 计算的 forward，不进生成器上下文。 -所以本实验**结构上无泄漏**，不需要 RLT 的 r_KL 来压——但仍然实现了 leak 监控项， -因为 rubric 里可能夹带答案片段（见 leak_frac）。 - -⭐ 打分目标选 `canonical_solution` 而不是「修对后的代码」：RLT 的 r_SS 是 -logp(标准答案 | 讲解, 题)。我们没有「修对后的代码」这种东西（那要先跑通才知道）， -所以用数据集自带的参考解。代价：参考解的写法风格与 executor 的自然写法不同， -logp 会偏低且带常数偏移 —— 但我们只在**同题的候选之间**比较排序，常数偏移会被抵消。 - -cheap 特征（全部来自一次 prompt_logprobs forward，零解码）： - * r_ss : mean logp(参考解 token | 题 + skill) <- RLT 主项，越高越好 - * r_ss_min : min-k 平均（最难的 10% token） <- RLT 的 α·min 项 - * ppl : exp(-r_ss)，困惑度 - * ent_* : 参考解位置上的预测熵（需要 topk） - * skill 自身的 logp/熵（生成时顺带拿到） -""" -import json -import os -import sys -import time -from collections import defaultdict -from typing import Any, Dict, List, Optional, Tuple - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.data_format import SamplingParams -from twinkle.sampler import vLLMSampler -from twinkle.template import Template - -_HERE = os.path.dirname(os.path.abspath(__file__)) -for _p in (_HERE, os.path.abspath(os.path.join(_HERE, '..', 'human'))): - if _p not in sys.path: - sys.path.insert(0, _p) - -from e18_kodcode import (clean_text, extract_code, judge_seqs, # noqa: E402 - load_records) -from e18_prompts import direct_prompt, skill_solve_prompt # noqa: E402 -from e18_multidiag import MultiDiagCache # noqa: E402 -from e23_rubric import build_checker # noqa: E402 - -logger = get_logger() - -MODEL_ID = os.environ.get('MODEL_ID', 'Qwen/Qwen3-4B') -OUTPUT_DIR = os.environ.get('OUTPUT_DIR', os.path.join(_HERE, 'output.e19.logp')) -SEED = int(os.environ.get('SEED', 42)) - -SKILL_GPUS = int(os.environ.get('SKILL_GPUS', 4)) -EXEC_GPUS = int(os.environ.get('EXEC_GPUS', 4)) -NUM_GPUS = SKILL_GPUS + EXEC_GPUS -GPU_MEM = float(os.environ.get('GPU_MEM', 0.85)) -MAX_MODEL_LEN = int(os.environ.get('MAX_MODEL_LEN', 24000)) - -N_TASKS = int(os.environ.get('N_TASKS', 24)) # 「几组错误的」——先小规模看效果 -N_SKILLS = int(os.environ.get('N_SKILLS', 8)) # 每题 8 个 skill 候选 -BARE_ROLLOUTS = int(os.environ.get('BARE_ROLLOUTS', 4)) -EXEC_ROLLOUTS = int(os.environ.get('EXEC_ROLLOUTS', 8)) # ground truth 用 -SKILL_MAX_TOKENS = int(os.environ.get('SKILL_MAX_TOKENS', 8192)) -EXEC_MAX_TOKENS = int(os.environ.get('EXEC_MAX_TOKENS', 15000)) -EXEC_TEMPERATURE = float(os.environ.get('EXEC_TEMPERATURE', 0.6)) -SKILL_TEMPERATURE = float(os.environ.get('SKILL_TEMPERATURE', 1.0)) -TOPK = int(os.environ.get('TOPK', 20)) # 算熵用的 top-k -RUN_ID = time.strftime('%m%d-%H%M%S') - - -# =========================================================================== -# prompt：失败 traj + rubric -> narrative skill（闭卷，无 GT） -# =========================================================================== -SKILLGEN_SYSTEM = ( - 'You are helping a Python programmer who is about to attempt a coding task. ' - 'You have seen this programmer fail this exact task before, and you have a ' - 'diagnosis of what went wrong.\n\n' - 'Write a short piece of guidance (a "skill") that would have prevented that ' - 'failure. Requirements:\n' - '- Write flowing prose, not bullet points or headings.\n' - '- Name the concrete API, argument, keyword, or edge case involved.\n' - '- Refer to the past failure as something that already happened, and say what ' - 'to do instead.\n' - '- Do NOT include any code block, and do NOT write a full solution.\n' - '- Keep it under 200 words.\n' - 'Wrap your guidance in <skills> and </skills> tags.') - - -def skillgen_prompt(problem: str, failed_code: str, error: str, rubric: str) -> Dict[str, Any]: - """闭卷 skill 生成 prompt：题面 + 自己的失败代码 + 报错 + rubric 诊断。 - - ⭐ 这里**没有** canonical_solution / reference answer。这是与 RLT teacher 的关键区别： - RLT 开卷（输入含答案）所以必须用 r_KL 压泄漏；我们闭卷，泄漏在结构上不可能发生。 - """ - user = (f'Task the programmer was given:\n{problem}\n\n' - f'The code they wrote (it failed):\n```python\n{failed_code}\n```\n\n' - f'How it failed:\n{error}\n\n' - f'Diagnosis:\n{rubric}\n\n' - 'Write the guidance that would have prevented this failure.') - return {'messages': [{'role': 'system', 'content': SKILLGEN_SYSTEM}, - {'role': 'user', 'content': user}]} - - -# =========================================================================== -# 采样 / 判分 工具（与 e18 同源） -# =========================================================================== -def run_samples(sampler, prompts, num_samples, max_tokens, gen_dp, - temperature=None, top_p=None, logprobs=None): - if not prompts: - return [] - import copy - params = SamplingParams( - max_tokens=max_tokens, - temperature=0.6 if temperature is None else temperature, - top_p=0.95 if top_p is None else top_p, - num_samples=num_samples, - **({} if logprobs is None else {'logprobs': logprobs})) - padded = prompts - if gen_dp > 1 and 0 < len(prompts) < gen_dp: - padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] - resp = sampler.sample(padded, params)[:len(prompts)] - return [list(r.sequences) if (r and r.sequences) else [] for r in resp] - - -def seq_text(seq) -> str: - return clean_text(getattr(seq, 'decoded', '') or '') if seq is not None else '' - - -def _mean(xs) -> float: - xs = [float(x) for x in xs if x is not None] - return sum(xs) / len(xs) if xs else 0.0 - - -def extract_skill(text: str) -> str: - if '<skills>' not in text: - return '' - body = text.split('<skills>', 1)[1] - return body.split('</skills>', 1)[0].strip() if '</skills>' in body else '' - - -# =========================================================================== -# ⭐ 核心：teacher-forcing 打分（一次 forward，零解码） -# =========================================================================== -def score_teacher_forcing(exec_sampler, problem: str, skill: str, target_code: str, - gen_dp: int) -> Dict[str, float]: - """算 logp(target_code | 题 + skill) —— RLT 的 r_SS，用 prompt_logprobs 实现。 - - ⭐ 机制：把「题+skill」当 prompt、把 target_code **拼在 prompt 末尾**，然后 - `max_tokens=1` + `prompt_logprobs=TOPK` 采样。vLLM 会回传**每个 prompt token 的 - logprob**（vllm_engine.py:324-343 取的是实际 token 的 logprob），于是我们免费拿到 - 了 target_code 每个 token 在该上下文下的条件概率 —— 这就是 teacher forcing， - 且**不解码任何 token**，比 8 次 rollout 便宜两个量级。 - - ⭐ 必须用 continue_final_message 让模板把 target_code 编成 **assistant 内容**而不是 - 新一轮 user：编错角色会导致 logp 分布完全不同（模型在算「用户会说这段代码的概率」）。 - 这里靠传入 assistant 角色的消息 + 模板拼接实现。 - - ⭐ 只取 target 段的 logprob，必须切掉 prompt 段。切点靠**两次编码求长度差**确定： - 先编「题+skill」得 n_ctx，再编「题+skill+target」得 n_all，则 target 占 - [n_ctx, n_all)。不能用固定偏移或字符串查找 —— tokenizer 会在边界合并 token。 - """ - return {} - - -def _entropy_from_topk(topk: List[Optional[List[Tuple[int, float]]]]) -> List[float]: - """由 top-k logprob 估计每位置的熵。 - - ⚠️ 这是**截断熵**（只看 top-k），不是真熵：尾部质量被忽略，所以系统性偏低。 - 但我们只做同题候选间的排序比较，偏差方向一致，可用。k=TOPK=20 时通常覆盖 >90% 概率质量。 - """ - import math - out = [] - for lps in topk or []: - if not lps: - out.append(0.0) - continue - ps = [math.exp(lp) for _, lp in lps] - z = sum(ps) or 1.0 - out.append(-sum((p / z) * math.log(max(p / z, 1e-12)) for p in ps)) - return out diff --git a/cookbook/human_e18/e20_success_skill.py b/cookbook/human_e18/e20_success_skill.py deleted file mode 100644 index cc035a17d..000000000 --- a/cookbook/human_e18/e20_success_skill.py +++ /dev/null @@ -1,311 +0,0 @@ -# -*- coding: utf-8 -*- -"""E20：**成功** trajectory -> narrative skill（无 rubric），什么情况下该保留？ - -与 E18 的差别只有一处：题源从「裸解失败」换成「裸解**第一次就成功**」， -于是 skillmodel 拿到的是一条**成功的** trajectory，而且**没有 rubric** -（rubric 是失败诊断，成功的题没有可诊断的失败）。 - -⭐ 本实验的核心难点：这些题裸解已经通过，**pass_rate 没有提升空间**。 -所以 E18 那套「+0.25 增益」门槛在这里恒不成立，直接套用会得出「一条都不该留」 -的空洞结论。必须换保留判据。本实验同时量四条候选判据： - - J1 **不倒退 (do-no-harm)**：加了 skill 后 pass_rate 不下降。 - —— 这是**必要条件**，不是充分条件（什么都不说的废话 skill 也满足）。 - J2 **稳健性提升**：裸解 M 次里**并非全对**（0<base<1）的题，加 skill 后升到 1.0。 - —— 「第一次成功」不等于「稳定成功」，这类题才有真实增量。 - J3 **token 效率**：pass 不变但生成长度显著变短（少走弯路）。 - J4 **迁移性**：skill 不含本题专属标识（函数名/变量名），才可能对别的题有用。 - -判据的取舍理由写在 decide_keep() 里。 - -流程（8 卡：4 skill + 4 executor）： - 1. 裸解 BARE_ROLLOUTS 次，挑**第 1 次就通过**的题 - 2. 把该次成功 trajectory（代码 + 通过信息）喂给 skillmodel，narrative、**无 rubric** - 3. 生成 N_SKILLS 个候选 - 4. 每候选 executor 重解 EXEC_ROLLOUTS 次 - 5. 按 J1-J4 分类，报「各判据下可保留的比例」 -""" -import json -import os -import re -import statistics as st -import sys -import time -from collections import Counter, defaultdict -from typing import Any, Dict, List - -import torch -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.data_format import SamplingParams -from twinkle.sampler import vLLMSampler -from twinkle.template import Template - -_HERE = os.path.dirname(os.path.abspath(__file__)) -for _p in (_HERE, os.path.abspath(os.path.join(_HERE, '..', 'human'))): - if _p not in sys.path: - sys.path.insert(0, _p) - -from e18_kodcode import clean_text, judge_seqs, load_records # noqa: E402 -from e18_prompts import direct_prompt, skill_solve_prompt # noqa: E402 - -logger = get_logger() - -MODEL_ID = os.environ.get('MODEL_ID', 'Qwen/Qwen3-4B') -OUTPUT_DIR = os.environ.get('OUTPUT_DIR', os.path.join(_HERE, 'output.e20.success')) -SEED = int(os.environ.get('SEED', 42)) -SKILL_GPUS = int(os.environ.get('SKILL_GPUS', 4)) -EXEC_GPUS = int(os.environ.get('EXEC_GPUS', 4)) -NUM_GPUS = SKILL_GPUS + EXEC_GPUS -GPU_MEM = float(os.environ.get('GPU_MEM', 0.85)) -MAX_MODEL_LEN = int(os.environ.get('MAX_MODEL_LEN', 24000)) - -N_TASKS = int(os.environ.get('N_TASKS', 64)) # 需要凑够的「首次成功」题数 -POOL_MULT = int(os.environ.get('POOL_MULT', 4)) # 题池放大倍数（首次成功率约 40%） -N_SKILLS = int(os.environ.get('N_SKILLS', 4)) -BARE_ROLLOUTS = int(os.environ.get('BARE_ROLLOUTS', 4)) -EXEC_ROLLOUTS = int(os.environ.get('EXEC_ROLLOUTS', 8)) -SKILL_MAX_TOKENS = int(os.environ.get('SKILL_MAX_TOKENS', 8192)) -EXEC_MAX_TOKENS = int(os.environ.get('EXEC_MAX_TOKENS', 15000)) -EXEC_TEMPERATURE = float(os.environ.get('EXEC_TEMPERATURE', 0.6)) -EXEC_TOP_P = float(os.environ.get('EXEC_TOP_P', 0.95)) -SKILL_TEMPERATURE = float(os.environ.get('SKILL_TEMPERATURE', 1.0)) -SKILL_CHAR_LIMIT = int(os.environ.get('SKILL_CHAR_LIMIT', 1500)) -RUN_ID = time.strftime('%m%d-%H%M%S') - - -# =========================================================================== -# prompt：成功 trajectory -> narrative skill（无 rubric） -# =========================================================================== -# ⭐ 与 E18 的 SKILLGEN_SYSTEM 保持同一 narrative 家族（散文体、第一人称、不提外部上下文、 -# 不写代码块），只把「诊断失败」换成「复盘一次成功」。刻意**不引入** rubric 位 —— -# 本实验的自变量就是「没有 rubric」。 -SKILLGEN_SYSTEM_SUCCESS = ( - 'You are a Python programmer writing a note to your future self.\n\n' - 'You just solved a coding task on the first attempt. Write down the one ' - 'insight that made it work, so that next time you meet a task of this shape ' - 'you get it right immediately again.\n\n' - 'Requirements:\n' - '- Write one flowing narrative in the first person, not bullet points or headings.\n' - '- Name the concrete API, argument, keyword, data shape, or edge case that mattered.\n' - '- Write it as guidance that transfers to other tasks of the same shape, not a ' - 'description of this one task. Do not mention the specific function name you wrote.\n' - '- Do NOT include any code block, and do NOT restate the solution.\n' - '- If nothing non-obvious was involved, say so briefly instead of inventing a lesson.\n' - '- Keep it under 150 words.\n' - 'Wrap the note in <skills> and </skills> tags.') - - -def skillgen_success_prompt(problem: str, code: str) -> Dict[str, Any]: - """成功复盘 prompt。**无 rubric、无 GT** —— 只有题面和自己刚写对的代码。 - - ⭐ 「If nothing non-obvious was involved, say so briefly」这一句是刻意加的逃生口： - 首次成功的题很多是**平凡题**，逼模型硬编一条"经验"只会得到套话。给它说"没什么特别" - 的许可，才能让 J1（不倒退）这个判据真正区分出「有料」和「没料」。 - """ - user = (f'The task:\n{problem}\n\n' - f'The solution you wrote, which passed on the first attempt:\n' - f'```python\n{code}\n```\n\n' - 'Write the note to your future self.') - return {'messages': [{'role': 'system', 'content': SKILLGEN_SYSTEM_SUCCESS}, - {'role': 'user', 'content': user}]} - - -# =========================================================================== -# 工具（与 e18_collect_kod 同源） -# =========================================================================== -def run_samples(sampler, prompts, num_samples, max_tokens, gen_dp, - temperature=None, top_p=None): - if not prompts: - return [] - import copy - params = SamplingParams( - max_tokens=max_tokens, - temperature=0.6 if temperature is None else temperature, - top_p=0.95 if top_p is None else top_p, - num_samples=num_samples) - padded = prompts - if gen_dp > 1 and 0 < len(prompts) < gen_dp: - padded = prompts + [copy.deepcopy(prompts[-1]) for _ in range(gen_dp - len(prompts))] - resp = sampler.sample(padded, params)[:len(prompts)] - return [list(r.sequences) if (r and r.sequences) else [] for r in resp] - - -def seq_text(seq) -> str: - return clean_text(getattr(seq, 'decoded', '') or '') if seq is not None else '' - - -def _mean(xs) -> float: - xs = [float(x) for x in xs if x is not None] - return sum(xs) / len(xs) if xs else 0.0 - - -def _pass_rate(rolls) -> float: - return (sum(1.0 for x in rolls if x['correct']) / len(rolls)) if rolls else 0.0 - - -def extract_skill(text: str) -> str: - if '<skills>' not in text: - return '' - body = text.split('<skills>', 1)[1] - return body.split('</skills>', 1)[0].strip() if '</skills>' in body else '' - - -# =========================================================================== -# 保留判据 -# =========================================================================== -IDENT = re.compile(r'\bdef\s+(\w+)') -# 「无信息」自述：模型用了 prompt 给的逃生口，说明它自己认为这题没什么可学的 -NOINFO = ('nothing non-obvious', 'nothing particularly', 'nothing special', - 'straightforward', 'no special', 'nothing unusual', 'not much to') - - -def decide_keep(base_rate: float, with_rate: float, skill: str, - gt_code: str, base_tokens: float, with_tokens: float) -> Dict[str, Any]: - """四条判据各自独立判定，不合成单一分数。 - - ⭐ 为何不合成一个总分：这四条问的是**不同的问题**，权重取决于 skill 池的用途 - （做 SFT 目标 vs 做检索库），此处只如实报出各判据的通过情况，把权衡留给决策。 - - J1 do-no-harm：with >= base。必要不充分 —— 一句废话也满足，所以**不能单独用它保留**。 - J2 稳健性：仅对 0<base<1（首次成功但不稳定）的题有意义，升到 1.0 才算。 - base 已经=1.0 的题在这条判据下恒为 False（没有可升空间），这是**设计如此**不是 bug。 - J3 token 效率：pass 不降且生成长度显著变短（>=10%），说明少走弯路。 - 10% 门槛是任意的，但比"变短一点"要求高，避免采样噪声。 - J4 迁移性：skill 不含 GT 里的函数名 + 没使用「没什么特别」的自述。 - 含函数名 -> 只对本题有效；自述无信息 -> 模型自己承认没料。 - """ - harmless = with_rate >= base_rate - 1e-9 - robust = (base_rate < 1.0) and (with_rate >= 1.0 - 1e-9) - shorter = harmless and with_tokens > 0 and base_tokens > 0 and \ - (with_tokens <= 0.90 * base_tokens) - names = set(IDENT.findall(gt_code or '')) - low = (skill or '').lower() - has_name = any(n and n.lower() in low for n in names) - self_noinfo = any(p in low for p in NOINFO) - transfer = (not has_name) and (not self_noinfo) - return {'J1_harmless': harmless, 'J2_robust': robust, 'J3_shorter': shorter, - 'J4_transfer': transfer, 'has_own_fn_name': has_name, - 'self_says_noinfo': self_noinfo} - - -# =========================================================================== -# 主流程 -# =========================================================================== -def build_runtime(): - twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, lazy_collect=False, groups=[ - DeviceGroup(name='skill', ranks=list(range(SKILL_GPUS)), device_type='GPU'), - DeviceGroup(name='exec', ranks=list(range(SKILL_GPUS, NUM_GPUS)), - device_type='GPU')]) - - def mk(group, world, thinking): - s = vLLMSampler(model_id=MODEL_ID, remote_group=group, - device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), - engine_args={'gpu_memory_utilization': GPU_MEM, - 'max_model_len': MAX_MODEL_LEN, - 'tensor_parallel_size': 1}) - s.set_template(Template, model_id=MODEL_ID, enable_thinking=thinking, - max_length=MAX_MODEL_LEN) - return s - - # skill 侧开 thinking（与 E18 采集口径一致，靠 thinking 拿多样性）；executor 关。 - return mk('skill', SKILL_GPUS, True), mk('exec', EXEC_GPUS, False) - - -def main(): - os.makedirs(OUTPUT_DIR, exist_ok=True) - t0 = time.time() - ds, _ = load_records(SEED, 0, OUTPUT_DIR) - pool = [r for i, r in enumerate(ds.dataset) if i < N_TASKS * POOL_MULT] - logger.info(f'E20 start: 题池 {len(pool)}（目标首次成功 {N_TASKS} 题）' - f' n_skills={N_SKILLS} bare={BARE_ROLLOUTS} exec={EXEC_ROLLOUTS}') - skill_sampler, exec_sampler = build_runtime() - - # ---- 1. 裸解，挑「第一次就通过」的题 ---- - bare = run_samples(exec_sampler, [direct_prompt(r['problem']) for r in pool], - BARE_ROLLOUTS, EXEC_MAX_TOKENS, EXEC_GPUS, - temperature=EXEC_TEMPERATURE, top_p=EXEC_TOP_P) - pairs, spans = [], [] - for r, seqs in zip(pool, bare): - seqs = list(seqs or []) - spans.append(len(seqs)) - pairs.extend((s, r['reference_answer']) for s in seqs) - judged = judge_seqs(pairs) if pairs else [] - picked, i = [], 0 - for r, n in zip(pool, spans): - rolls = judged[i:i + n] - i += n - if not rolls or not rolls[0]['correct']: - continue # ⭐ 只要**第 1 次**就通过的题（按用户要求） - picked.append({'rec': r, 'base_rate': _pass_rate(rolls), - 'code': rolls[0].get('code') or '', - 'base_tokens': _mean([x.get('gen_tokens') for x in rolls])}) - if len(picked) >= N_TASKS: - break - logger.info(f'[bare] 首次成功 {len(picked)}/{len(pool)} 题' - f'（其中 base<1 的不稳定题 {sum(1 for p in picked if p["base_rate"] < 1.0)}）') - if not picked: - raise RuntimeError('没有首次成功的题') - - # ---- 2. 成功 traj -> narrative skill（无 rubric）---- - sg = run_samples(skill_sampler, - [skillgen_success_prompt(p['rec']['problem'], p['code']) - for p in picked], - N_SKILLS, SKILL_MAX_TOKENS, SKILL_GPUS, - temperature=SKILL_TEMPERATURE) - flat = [] - for p, seqs in zip(picked, sg): - for ci in range(N_SKILLS): - seq = seqs[ci] if seqs and ci < len(seqs) else None - sk = extract_skill(seq_text(seq)) - flat.append({'p': p, 'cand_idx': ci, 'skill': sk, - 'stop': getattr(seq, 'stop_reason', None) if seq else None}) - n_ok = sum(1 for f in flat if f['skill']) - logger.info(f'[skillgen] {len(flat)} 候选，可解析 {n_ok} ({100*n_ok/len(flat):.0f}%)') - - # ---- 3. 带 skill 重解 ---- - todo = [f for f in flat if f['skill']] - ws = run_samples(exec_sampler, - [skill_solve_prompt(f['p']['rec']['problem'], f['skill']) for f in todo], - EXEC_ROLLOUTS, EXEC_MAX_TOKENS, EXEC_GPUS, - temperature=EXEC_TEMPERATURE, top_p=EXEC_TOP_P) - pairs, spans = [], [] - for f, seqs in zip(todo, ws): - seqs = list(seqs or []) - spans.append(len(seqs)) - pairs.extend((s, f['p']['rec']['reference_answer']) for s in seqs) - judged = judge_seqs(pairs) if pairs else [] - i = 0 - for f, n in zip(todo, spans): - rolls = judged[i:i + n] - i += n - f['with_rate'] = _pass_rate(rolls) - f['with_tokens'] = _mean([x.get('gen_tokens') for x in rolls]) - - # ---- 4. 判据 ---- - out = os.path.join(OUTPUT_DIR, 'e20_candidates.jsonl') - with open(out, 'w', encoding='utf-8') as fh: - for f in flat: - p = f['p'] - if not f['skill']: - row = {'data_id': p['rec']['data_id'], 'cand_idx': f['cand_idx'], - 'parseable': False, 'base_rate': p['base_rate'], - 'with_rate': None, 'skill': '', 'skill_chars': 0} - else: - d = decide_keep(p['base_rate'], f['with_rate'], f['skill'], - p['rec']['reference_answer'].get('canonical_solution', ''), - p['base_tokens'], f['with_tokens']) - row = {'data_id': p['rec']['data_id'], 'cand_idx': f['cand_idx'], - 'parseable': True, 'base_rate': p['base_rate'], - 'with_rate': f['with_rate'], - 'delta': round(f['with_rate'] - p['base_rate'], 6), - 'base_tokens': round(p['base_tokens'], 1), - 'with_tokens': round(f['with_tokens'], 1), - 'skill_chars': len(f['skill']), 'stop': f['stop'], - 'skill': f['skill'], **d} - fh.write(json.dumps(row, ensure_ascii=False) + '\n') - logger.info(f'[done] 落盘 {out}，用时 {(time.time()-t0)/60:.1f} 分钟') - - -if __name__ == '__main__': - main() diff --git a/cookbook/human_e18/e21_paired_rubric.py b/cookbook/human_e18/e21_paired_rubric.py deleted file mode 100644 index b56659d4b..000000000 --- a/cookbook/human_e18/e21_paired_rubric.py +++ /dev/null @@ -1,214 +0,0 @@ -# -*- coding: utf-8 -*- -"""E21：rubric 值多少钱？在**同一批题**上做成功复盘 vs 失败诊断的配对对照。 - -目标题型 = 「首次成功但不稳定」（0 < base_pass_rate < 1）。 -⭐ 为什么只能用这类题：它们**同时拥有**成功轨迹和失败轨迹，所以同一道题可以同时喂给 -两个 arm，构成**配对设计**（paired design）。base=1.0 的题没有失败轨迹（无法出 rubric）， -base=0 的题没有成功轨迹（无法做成功复盘）—— 只有这个交集能做干净对照。 -配对比独立分组强得多：题目难度是最大的方差来源，配对把它消掉了。 - -两个 arm，除了输入信号完全同构（同题、同 N_SKILLS、同 executor、同温度、同 rollout 数）： - - arm SUCCESS : 成功代码 -> narrative skill（无 rubric） [E20 的 prompt] - arm RUBRIC : 失败代码 + 报错 -> 教师诊断出 rubric -> narrative skill [E18 的 prompt] - -判据统一为 J2（升到全对）：base<1 的题加 skill 后 with_pass_rate 是否达到 1.0。 -这是唯一在两个 arm 上都可测、且不受选择偏差污染的口径。 - -⚠️ 已知的不对称（诚实记录，不是 bug）： - 1. RUBRIC arm 多消耗一次教师 API 调用（不占 GPU，但不是零成本）。 - 2. 两个 arm 的 prompt 家族不同（SKILLGEN_SYSTEM vs SKILLGEN_SYSTEM_SUCCESS）， - 所以测的是「成功复盘管线」vs「失败诊断管线」的**整体**差异， - 不是「rubric 这一个字段」的净效应。要拆到字段级需要第三个 arm - （失败代码但不给 rubric），本脚本用 FAILONLY arm 补上。 - 3. FAILONLY arm 复用 SKILLGEN_SYSTEM 但把 rubric 位填成「无诊断」的官方 fallback - （skillgen_prompt 内建该分支），所以 arm 间 system prompt 一致， - RUBRIC vs FAILONLY 的差值才是 rubric 字段的净贡献。 -""" -import json -import os -import sys -import time -from collections import defaultdict -from typing import Any, Dict, List - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.data_format import SamplingParams -from twinkle.sampler import vLLMSampler -from twinkle.template import Template - -_HERE = os.path.dirname(os.path.abspath(__file__)) -for _p in (_HERE, os.path.abspath(os.path.join(_HERE, '..', 'human'))): - if _p not in sys.path: - sys.path.insert(0, _p) - -from e18_kodcode import clean_text, judge_seqs, load_records # noqa: E402 -from e18_prompts import direct_prompt, skill_solve_prompt, skillgen_prompt # noqa: E402 -from e20_success_skill import (extract_skill, run_samples, seq_text, # noqa: E402 - skillgen_success_prompt, _mean, _pass_rate) -from e23_rubric import RubricCache, build_checker # noqa: E402 - -logger = get_logger() - -MODEL_ID = os.environ.get('MODEL_ID', 'Qwen/Qwen3-4B') -OUTPUT_DIR = os.environ.get('OUTPUT_DIR', os.path.join(_HERE, 'output.e21.paired')) -SEED = int(os.environ.get('SEED', 42)) -SKILL_GPUS = int(os.environ.get('SKILL_GPUS', 4)) -EXEC_GPUS = int(os.environ.get('EXEC_GPUS', 4)) -NUM_GPUS = SKILL_GPUS + EXEC_GPUS -GPU_MEM = float(os.environ.get('GPU_MEM', 0.85)) -MAX_MODEL_LEN = int(os.environ.get('MAX_MODEL_LEN', 24000)) - -POOL = int(os.environ.get('POOL', 400)) # 题池；不稳定题约占 4-5% -N_SKILLS = int(os.environ.get('N_SKILLS', 4)) -BARE_ROLLOUTS = int(os.environ.get('BARE_ROLLOUTS', 4)) -EXEC_ROLLOUTS = int(os.environ.get('EXEC_ROLLOUTS', 8)) -SKILL_MAX_TOKENS = int(os.environ.get('SKILL_MAX_TOKENS', 8192)) -EXEC_MAX_TOKENS = int(os.environ.get('EXEC_MAX_TOKENS', 15000)) -EXEC_TEMPERATURE = float(os.environ.get('EXEC_TEMPERATURE', 0.6)) -EXEC_TOP_P = float(os.environ.get('EXEC_TOP_P', 0.95)) -SKILL_TEMPERATURE = float(os.environ.get('SKILL_TEMPERATURE', 1.0)) - - -def build_runtime(): - twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, lazy_collect=False, groups=[ - DeviceGroup(name='skill', ranks=list(range(SKILL_GPUS)), device_type='GPU'), - DeviceGroup(name='exec', ranks=list(range(SKILL_GPUS, NUM_GPUS)), - device_type='GPU')]) - - def mk(group, world, thinking): - s = vLLMSampler(model_id=MODEL_ID, remote_group=group, - device_mesh=DeviceMesh.from_sizes(world_size=world, dp_size=world), - engine_args={'gpu_memory_utilization': GPU_MEM, - 'max_model_len': MAX_MODEL_LEN, - 'tensor_parallel_size': 1}) - s.set_template(Template, model_id=MODEL_ID, enable_thinking=thinking, - max_length=MAX_MODEL_LEN) - return s - - return mk('skill', SKILL_GPUS, True), mk('exec', EXEC_GPUS, False) - - -def main(): - os.makedirs(OUTPUT_DIR, exist_ok=True) - t0 = time.time() - ds, _ = load_records(SEED, 0, OUTPUT_DIR) - pool = [r for i, r in enumerate(ds.dataset) if i < POOL] - logger.info(f'E21 start: 题池 {len(pool)} n_skills={N_SKILLS} ' - f'bare={BARE_ROLLOUTS} exec={EXEC_ROLLOUTS}') - skill_sampler, exec_sampler = build_runtime() - - # ---- 1. 裸解，挑「首次成功但不稳定」的题 ---- - bare = run_samples(exec_sampler, [direct_prompt(r['problem']) for r in pool], - BARE_ROLLOUTS, EXEC_MAX_TOKENS, EXEC_GPUS, - temperature=EXEC_TEMPERATURE, top_p=EXEC_TOP_P) - pairs, spans = [], [] - for r, seqs in zip(pool, bare): - seqs = list(seqs or []) - spans.append(len(seqs)) - pairs.extend((s, r['reference_answer']) for s in seqs) - judged = judge_seqs(pairs) if pairs else [] - picked, i, n_first_ok = [], 0, 0 - for r, n in zip(pool, spans): - rolls = judged[i:i + n] - i += n - if not rolls or not rolls[0]['correct']: - continue - n_first_ok += 1 - rate = _pass_rate(rolls) - if rate >= 1.0: - continue # 稳定成功 -> 无失败轨迹，做不了配对 - bad = next((x for x in rolls[1:] if not x['correct']), None) - if bad is None: - continue - picked.append({'rec': r, 'base_rate': rate, - 'good_code': rolls[0].get('code') or '', - 'bad_roll': bad, - 'base_tokens': _mean([x.get('gen_tokens') for x in rolls])}) - logger.info(f'[bare] 首次成功 {n_first_ok}/{len(pool)}；' - f'其中**不稳定**（可配对）{len(picked)} 题') - if not picked: - raise RuntimeError('没有可配对的不稳定题') - - # ---- 2. 教师诊断出 rubric（纯 API，不占 GPU）---- - # ⭐ 缓存必须用**本 run 私有**的路径，不能复用全局 RUBRIC_CACHE_PATH： - # RubricCache 的键里**不含轨迹**，其跨 run 复用的前提是「executor 冻结在 T=0， - # 同一题裸解逐字相同」（见 e23_rubric.py:227 的说明）。本实验裸解用 T=0.6， - # 轨迹每次不同，共用全局缓存会拿到**别的轨迹**的诊断，静默污染 RUBRIC arm。 - cache = RubricCache(os.path.join(OUTPUT_DIR, 'diag_cache.jsonl')) - checker = build_checker() - rubrics = cache.diagnose_many(checker, [(p['rec'], p['bad_roll']) for p in picked]) - n_rub = sum(1 for x in rubrics for _ in [x] if x) - logger.info(f'[rubric] {n_rub}/{len(picked)} 题拿到诊断') - for p, rb in zip(picked, rubrics): - p['rubric'] = rb or '' - - # ---- 3. 三个 arm 生成 skill ---- - # ⭐ 三个 arm 一次性拼进同一个 sample 调用，保证同一批权重、同一批 KV cache 状态， - # 避免"先跑完 A 再跑 B"引入的引擎状态差异。 - arms: List[str] = ['SUCCESS', 'RUBRIC', 'FAILONLY'] - prompts, meta = [], [] - for p in picked: - prompts.append(skillgen_success_prompt(p['rec']['problem'], p['good_code'])) - meta.append((p, 'SUCCESS')) - prompts.append(skillgen_prompt(p['rec']['problem'], p['rubric'], False)) - meta.append((p, 'RUBRIC')) - # FAILONLY：同 system prompt，rubric 位走内建的「无诊断」fallback - prompts.append(skillgen_prompt(p['rec']['problem'], '', False)) - meta.append((p, 'FAILONLY')) - sg = run_samples(skill_sampler, prompts, N_SKILLS, SKILL_MAX_TOKENS, - SKILL_GPUS, temperature=SKILL_TEMPERATURE) - flat = [] - for (p, arm), seqs in zip(meta, sg): - for ci in range(N_SKILLS): - seq = seqs[ci] if seqs and ci < len(seqs) else None - flat.append({'p': p, 'arm': arm, 'cand_idx': ci, - 'skill': extract_skill(seq_text(seq))}) - for arm in arms: - g = [f for f in flat if f['arm'] == arm] - n = sum(1 for f in g if f['skill']) - logger.info(f'[skillgen] {arm}: {n}/{len(g)} 可解析 ({100*n/max(1,len(g)):.0f}%)') - - # ---- 4. 带 skill 重解 ---- - todo = [f for f in flat if f['skill']] - ws = run_samples(exec_sampler, - [skill_solve_prompt(f['p']['rec']['problem'], f['skill']) for f in todo], - EXEC_ROLLOUTS, EXEC_MAX_TOKENS, EXEC_GPUS, - temperature=EXEC_TEMPERATURE, top_p=EXEC_TOP_P) - pairs, spans = [], [] - for f, seqs in zip(todo, ws): - seqs = list(seqs or []) - spans.append(len(seqs)) - pairs.extend((s, f['p']['rec']['reference_answer']) for s in seqs) - judged = judge_seqs(pairs) if pairs else [] - i = 0 - for f, n in zip(todo, spans): - rolls = judged[i:i + n] - i += n - f['with_rate'] = _pass_rate(rolls) - f['with_tokens'] = _mean([x.get('gen_tokens') for x in rolls]) - - # ---- 5. 落盘 ---- - out = os.path.join(OUTPUT_DIR, 'e21_candidates.jsonl') - with open(out, 'w', encoding='utf-8') as fh: - for f in flat: - p = f['p'] - row = {'data_id': p['rec']['data_id'], 'arm': f['arm'], - 'cand_idx': f['cand_idx'], 'base_rate': p['base_rate'], - 'parseable': bool(f['skill']), - 'with_rate': f.get('with_rate'), - 'base_tokens': round(p['base_tokens'], 1), - 'with_tokens': round(f.get('with_tokens') or 0, 1), - 'has_rubric': bool(p['rubric']), - 'skill_chars': len(f['skill']), 'skill': f['skill']} - if f.get('with_rate') is not None: - row['delta'] = round(f['with_rate'] - p['base_rate'], 6) - row['J2_robust'] = f['with_rate'] >= 1.0 - 1e-9 - fh.write(json.dumps(row, ensure_ascii=False) + '\n') - cache.close() - logger.info(f'[done] 落盘 {out}，用时 {(time.time()-t0)/60:.1f} 分钟') - - -if __name__ == '__main__': - main() diff --git a/cookbook/human_e18/preflight.py b/cookbook/human_e18/preflight.py deleted file mode 100644 index 03ea0606d..000000000 --- a/cookbook/human_e18/preflight.py +++ /dev/null @@ -1,147 +0,0 @@ -# -*- coding: utf-8 -*- -"""开跑前自检：把「跑 10 小时才发现环境不对」提前到 30 秒内暴露。 - -⭐ 为何必须存在：judge 的失败是**静默**的 —— pytest 缺失不会让进程崩， -只会让每道题都判 incorrect，日志上表现为 baseline_accuracy=0 / n_wrong=64/64， -而 collected 一直是 0。B 机曾因此白跑 10 小时 65 个 chunk。 - -⚠️ 关键点：judge 用 `subprocess.run([sys.executable, '_run.py'])` 起子进程， -所以 pytest 必须装在**启动脚本所用的那个解释器**里，不是 `which pytest` 指的那个。 -本脚本用 sys.executable 自查，跟 judge 走同一条路径。 - -用法：python3 preflight.py （用你打算启动采集的同一个 python3） -""" -import os -import subprocess -import sys -import tempfile - -FAIL, WARN = [], [] - - -def ck(name, cond, note='', hard=True): - if not cond: - (FAIL if hard else WARN).append(name) - tag = 'OK ' if cond else ('BAD' if hard else 'WARN') - print(' [%s] %-44s %s' % (tag, name, note)) - - -print('=== 0. 解释器 ===') -print(' sys.executable = %s' % sys.executable) -print(' version = %s' % sys.version.split()[0]) - -print() -print('=== 1. judge 沙箱（最关键，静默失败源）===') -# 完整复刻 judge 的执行路径：子进程 + pytest 跑一个必过的单测 -# ⭐ 完整照抄 e18_kodcode.run_tests 的真实结构，而不是自己另写一个 pytest 调用： -# - solution.py 被测代码（单测靠 `from solution import X` 取） -# - test_solution.py 单测文件（pytest 只收集 test_*.py，名字不能改） -# - _run.py runner，pytest.main 显式指定 test_solution.py -# - subprocess + cwd=tmp + sys.executable -# 只有走同一条路径，「通过」才真的等价于 judge 会判通过。 -_RUNNER = '''import sys, pytest -rc = pytest.main(['-q', '--no-header', '-p', 'no:cacheprovider', - '--tb=short', 'test_solution.py']) -print('__KOD__', rc) -sys.exit(0 if int(rc) == 0 else 1) -''' -_SOLUTION = 'def add(a, b):\n return a + b\n' -_TEST = 'from solution import add\n\n\ndef test_add():\n assert add(1, 2) == 3\n' - -with tempfile.TemporaryDirectory() as td: - for name, body in (('solution.py', _SOLUTION), - ('test_solution.py', _TEST), - ('_run.py', _RUNNER)): - with open(os.path.join(td, name), 'w') as f: - f.write(body) - try: - env = dict(os.environ, PYTHONHASHSEED='0') - env.pop('CUDA_VISIBLE_DEVICES', None) - r = subprocess.run([sys.executable, '_run.py'], cwd=td, env=env, - capture_output=True, text=True, timeout=120) - ok = (r.returncode == 0) - last = (r.stdout or r.stderr).strip().split('\n')[-1][:60] - ck('judge 沙箱可判对一份正确解', ok, 'rc=%d %s' % (r.returncode, last)) - if not ok: - print(' -> 修复：%s -m pip install pytest' % sys.executable) - print(' -> 完整输出：') - for ln in (r.stdout + r.stderr).strip().split('\n')[-6:]: - print(' %s' % ln[:100]) - except Exception as e: - ck('judge 沙箱可判对一份正确解', False, str(e)[:70]) - -print() -print('=== 2. 依赖包 ===') -import importlib.metadata as _md -# A 机实测版本，作为对照基线 -BASE = {'pytest': '9.1.1', 'vllm': '0.23.0', 'torch': '2.11.0+cu130', - 'transformers': '5.14.1', 'modelscope': '1.38.1', 'datasets': '4.8.4', - 'openai': '2.45.0', 'numpy': '2.5.1'} -for pkg, want in BASE.items(): - try: - got = _md.version(pkg) - # 版本不一致只告警：主版本差异才真会出问题，补丁号差异通常无害 - same_major = got.split('.')[0] == want.split('.')[0] - ck(pkg, True, '%s (A机 %s)%s' % (got, want, '' if same_major else ' <- 主版本不同')) - if not same_major: - WARN.append(pkg + '-major') - except Exception: - ck(pkg, False, '未安装 (A机 %s)' % want) - -print() -print('=== 3. 教师 API（rubric 的唯一来源）===') -# 不发真请求，只查变量在不在：缺 key 会让 build_checker 返回 None, -# 后果是 rubric 全空 -> n_rubric_missing == 题数 -> 一条都收不到 -for v in ('LLM_BACKUP_API_KEY', 'LLM_BACKUP_BASE_URL'): - val = os.environ.get(v, '') - ck(v, bool(val), ('已设置(%d字符)' % len(val)) if val else '缺失 -> rubric 会全空') - -print() -print('=== 4. 分片参数 ===') -sn = int(os.environ.get('SHARD_N', 1)) -si = int(os.environ.get('SHARD_ID', 0)) -ck('SHARD_N/SHARD_ID 合法', sn >= 1 and 0 <= si < sn, 'SHARD_N=%d SHARD_ID=%d' % (sn, si)) -ck('多机模式已开启', sn > 1, '单机模式' if sn == 1 else '分片 %d/%d' % (si, sn), hard=False) - -print() -print('=== 5. resume 种子 ===') -_HERE = os.path.dirname(os.path.abspath(__file__)) -od = os.environ.get('OUTPUT_DIR', os.path.join(_HERE, 'output.e18.kod')) -cand = os.path.join(od, 'e18_candidates.jsonl') -if os.path.exists(cand): - import json - import zlib - n = mine = 0 - with open(cand, encoding='utf-8') as f: - for ln in f: - ln = ln.strip() - if not ln: - continue - try: - d = json.loads(ln).get('data_id') - except Exception: - continue - if d: - n += 1 - if zlib.crc32(str(d).encode()) % sn == si: - mine += 1 - ck('种子文件存在', True, '%d 个 id，其中属于本分片 %d 个' % (n, mine)) - ck('种子含本分片的题', mine > 0 or sn == 1, - '本分片会跳过 %d 题' % mine, hard=False) -else: - ck('种子文件存在', False, - '缺 %s -> 会重跑别的机器已做过的题' % cand, hard=False) - -print() -print('=== 6. 数据集缓存 ===') -ck('未设 HF_DATASETS_OFFLINE', os.environ.get('HF_DATASETS_OFFLINE', '') not in ('1', 'true'), - '设了会因本机缓存 config 名带 hash 后缀而加载失败') - -print() -if FAIL: - print('结论: 不可启动 —— 必须先修: %s' % FAIL) -elif WARN: - print('结论: 可启动，但注意: %s' % WARN) -else: - print('结论: 全部通过，可以启动') -sys.exit(1 if FAIL else 0) diff --git a/cookbook/human_e18/run_collect_kod.sh b/cookbook/human_e18/run_collect_kod.sh deleted file mode 100755 index 55464e691..000000000 --- a/cookbook/human_e18/run_collect_kod.sh +++ /dev/null @@ -1,8 +0,0 @@ -#!/bin/bash -# KodCode 冷启动数据采集：8 卡全 rollout，无 SFT，narrative prompt 不变。 -cd /mnt/data/yzhao/tastelikefeet/twinkle/cookbook/human_e18 -export PYTHONPATH=/mnt/data/yzhao/tastelikefeet/twinkle/src:${PYTHONPATH} -export SKILL_SAMPLER_GPUS=2 BASE_SAMPLER_GPUS=6 # token 预算比 16.5:1，2+6 实测理论最优 -export KOD_SELFCHECK=0 # 不自检（坏题采集时自然不入池） -export TARGET_SAMPLES=5000 -exec /usr/local/bin/python -u e18_collect_kod.py diff --git a/cookbook/human_e18/run_sft_kod.sh b/cookbook/human_e18/run_sft_kod.sh deleted file mode 100755 index 3686b3400..000000000 --- a/cookbook/human_e18/run_sft_kod.sh +++ /dev/null @@ -1,18 +0,0 @@ -#!/bin/bash -# E18-KOD 离线 SFT + 首尾 eval：训一遍，看 loss 收敛 + 训练前后 lift 差异。 -# ⚠️ 会占满 8 卡（4 训练 + 2 skillmodel + 2 executor）—— 采集进程若还在跑，先确认它已停。 -# ⚠️ eval 要跑 judge（本地 pytest 沙箱），不需要教师 API key。 -cd /mnt/data/yzhao/tastelikefeet/twinkle/cookbook/human_e18 -export PYTHONPATH=/mnt/data/yzhao/tastelikefeet/twinkle/src:${PYTHONPATH} -export TRAIN_GPUS=4 # 4 训练 -export SKILL_SAMPLER_GPUS=2 # 2 张出 skill -export BASE_SAMPLER_GPUS=2 # 2 张跑 executor -export BATCH_SIZE=16 # 必须是 TRAIN_DP(=4) 的整倍数 -export MICRO_BATCH=8 -export LR=1e-5 # 与在线版一致：恒定 lr、无 warmup/decay -export EPOCHS=3 -export EVAL_SIZE=100 # 首尾各 100 题，选自 KodCode 中未生成过 skill 的题 -export EXEC_ROLLOUTS=1 # executor 单次 -export EXEC_TEMPERATURE=0.0 # greedy，可复现 -export SAVE_EVERY_STEPS=0 # 0 = 只在结束时存权重 -exec /usr/local/bin/python -u e18_sft_kod.py diff --git a/cookbook/human_e18/shard_tool.py b/cookbook/human_e18/shard_tool.py deleted file mode 100644 index 5615edcb4..000000000 --- a/cookbook/human_e18/shard_tool.py +++ /dev/null @@ -1,134 +0,0 @@ -# -*- coding: utf-8 -*- -"""多机分片采集的两个配套工具：导出「已跑过的题」种子 + 合并多机产物。 - - python3 shard_tool.py seed <src_dir> <out.jsonl> # 在 A 机跑，产出给 B 机的种子 - python3 shard_tool.py merge <out_dir> <dir1> [dir2 ...] # 合并任意台机器的产物 - -为什么需要 seed： - crc32 分片只保证「以后」两台机器不撞，但 A 机已经跑掉的题会均匀落在所有分片上 - （实测 7374 题在 SHARD_N=2 下是 3667/3707）。B 机目录是空的，resume_done_ids() - 读不到任何东西，就会把属于自己分片的那 3707 题重跑一遍 —— 白烧一半算力， - 而且合并时同一 data_id 出现两份。种子文件就是把 A 机的 done id 搬给 B 机。 - ❗ 种子只需要 e18_candidates.jsonl 的 data_id 字段（resume_done_ids 只读这个）， - 所以导出的是**精简行**，不是整个 27000 行的候选文件 —— B 机不需要 A 机的 skill 正文。 -""" -import collections -import json -import os -import sys - -FILES = ('e18_sft_dataset.jsonl', 'e18_candidates.jsonl', 'collect_log.jsonl') - - -def _rows(path): - if not os.path.exists(path): - return - with open(path, encoding='utf-8') as f: - for ln in f: - ln = ln.strip() - if not ln: - continue - try: - yield json.loads(ln) - except Exception: - continue - - -def cmd_seed(src, out): - """把 src 里所有跑过的 data_id 导成最小的 candidates 行。""" - ids = {r['data_id'] for r in _rows(os.path.join(src, 'e18_candidates.jsonl')) - if r.get('data_id')} - with open(out, 'w', encoding='utf-8') as f: - for i in sorted(ids): - # resume_done_ids() 只取 data_id；其余字段给最小合法值， - # 保证这些行即使被别的分析脚本读到也不会伪装成真实候选： - # kept=False + parseable=False + seed=True 一眼可滤。 - f.write(json.dumps({'data_id': i, 'kept': False, 'parseable': False, - 'seed': True, 'run': 'SEED', 'chunk': -1}, - ensure_ascii=False) + '\n') - print('导出 %d 个已跑 data_id -> %s' % (len(ids), out)) - print('用法：拷到 B 机的 OUTPUT_DIR/e18_candidates.jsonl，再用 KOD_RESUME=1 启动') - - -def cmd_merge(out_dir, dirs): - os.makedirs(out_dir, exist_ok=True) - report = {} - # ---- 1. sft_dataset：按 data_id 去重（同题多机重复时保留 pass_gain 更高者）---- - best, dup = {}, 0 - for d in dirs: - for r in _rows(os.path.join(d, 'e18_sft_dataset.jsonl')): - k = r.get('data_id') - if not k: - continue - if k in best: - dup += 1 - # 保留 gain 高的：重复只可能来自"种子没同步"的意外， - # 此时保留更优样本比保留先到者更合理 - if (r.get('pass_gain') or -9) <= (best[k].get('pass_gain') or -9): - continue - best[k] = r - with open(os.path.join(out_dir, 'e18_sft_dataset.jsonl'), 'w', encoding='utf-8') as f: - for r in best.values(): - f.write(json.dumps(r, ensure_ascii=False) + '\n') - report['sft'] = (len(best), dup) - - # ---- 2. candidates：按 (data_id, run, cand_idx) 去重，丢掉 seed 占位行 ---- - seen, rows, nseed = set(), [], 0 - for d in dirs: - for r in _rows(os.path.join(d, 'e18_candidates.jsonl')): - if r.get('seed'): - nseed += 1 - continue - k = (r.get('data_id'), r.get('run'), r.get('cand_idx')) - if k in seen: - continue - seen.add(k) - rows.append(r) - with open(os.path.join(out_dir, 'e18_candidates.jsonl'), 'w', encoding='utf-8') as f: - for r in rows: - f.write(json.dumps(r, ensure_ascii=False) + '\n') - report['cand'] = (len(rows), nseed) - - # ---- 3. collect_log：chunk 编号两机都从 0 起，直接 cat 会产生歧义 ---- - # 不重编号（会破坏与 run.log 的对照），改为加 src 字段标机器来源， - # 并按 (run, chunk) 唯一化。分析脚本本来就该按 run 分组看 chunk。 - lg, seen2 = [], set() - for d in dirs: - tag = os.path.basename(os.path.normpath(d)) - for r in _rows(os.path.join(d, 'collect_log.jsonl')): - k = (r.get('run'), r.get('chunk')) - if k in seen2: - continue - seen2.add(k) - r['src'] = tag - lg.append(r) - lg.sort(key=lambda r: (str(r.get('run')), r.get('chunk') or 0)) - with open(os.path.join(out_dir, 'collect_log.jsonl'), 'w', encoding='utf-8') as f: - for r in lg: - f.write(json.dumps(r, ensure_ascii=False) + '\n') - report['log'] = (len(lg), 0) - - print('=== 合并完成 -> %s ===' % out_dir) - print(' e18_sft_dataset.jsonl %6d 条 (跨机重复丢弃 %d)' % report['sft']) - print(' e18_candidates.jsonl %6d 条 (滤掉种子占位 %d)' % report['cand']) - print(' collect_log.jsonl %6d 条' % report['log'][0]) - # 交叉校验：胜者的 data_id 必须都能在候选里找到 - cid = {r.get('data_id') for r in rows} - miss = [k for k in best if k not in cid] - print(' 一致性：胜者 data_id 在候选中缺失 %d 个 %s' - % (len(miss), '(OK)' if not miss else '<- 异常')) - per = collections.Counter(str(r.get('run')).split('.')[-1] for r in rows) - print(' 按分片计候选数：%s' % dict(per)) - - -if __name__ == '__main__': - if len(sys.argv) < 3: - print(__doc__) - sys.exit(2) - if sys.argv[1] == 'seed': - cmd_seed(sys.argv[2], sys.argv[3]) - elif sys.argv[1] == 'merge': - cmd_merge(sys.argv[2], sys.argv[3:]) - else: - print(__doc__) - sys.exit(2) From ce2dbe09964d4999395be44100570aaf77971d97 Mon Sep 17 00:00:00 2001 From: root <yuze.zyz@alibaba-inc.com> Date: Sun, 16 Aug 2026 22:53:58 +0800 Subject: [PATCH 40/60] wip --- cookbook/rsi/run_rsi.py | 131 +++++++++++ src/twinkle_agentic/rsi/__init__.py | 9 + src/twinkle_agentic/rsi/rsi_distill.py | 285 ++++++++++++++++++++++ src/twinkle_agentic/rsi/rsi_prepare.py | 175 ++++++++++++++ src/twinkle_agentic/rsi/rsi_refine.py | 277 ++++++++++++++++++++++ src/twinkle_agentic/rsi/rsi_rl.py | 299 ++++++++++++++++++++++++ src/twinkle_agentic/utils/llm_backup.py | 47 +++- 7 files changed, 1218 insertions(+), 5 deletions(-) create mode 100644 cookbook/rsi/run_rsi.py create mode 100644 src/twinkle_agentic/rsi/__init__.py create mode 100644 src/twinkle_agentic/rsi/rsi_distill.py create mode 100644 src/twinkle_agentic/rsi/rsi_prepare.py create mode 100644 src/twinkle_agentic/rsi/rsi_refine.py create mode 100644 src/twinkle_agentic/rsi/rsi_rl.py diff --git a/cookbook/rsi/run_rsi.py b/cookbook/rsi/run_rsi.py new file mode 100644 index 000000000..8bddfea56 --- /dev/null +++ b/cookbook/rsi/run_rsi.py @@ -0,0 +1,131 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""RSI pipeline entry point — run any single stage (or the whole chain) so each +step can be validated in isolation. + +The four stages live in ``twinkle_agentic.rsi`` and each already has its own CLI: + + 1 prepare twinkle_agentic.rsi.rsi_prepare (CPU) raw -> subset + 2 refine twinkle_agentic.rsi.rsi_refine (API) subset-> flows + 3 rl twinkle_agentic.rsi.rsi_rl (ray+GPU) flows -> executor LoRA + 4 distill twinkle_agentic.rsi.rsi_distill (ray+GPU) dump -> role LoRA + +Why this launches SUBPROCESSES instead of importing and calling: + * ``rsi_rl`` runs ``CLI.from_args()`` and ``swanlab.init()`` at IMPORT time, so + merely importing it would parse this launcher's argv and start a run. + * ``rl`` and ``distill`` need DIFFERENT ray topologies (MultiLora+sampler vs a + single TransformersModel group); they cannot share one ray init in-process. +Running each stage as its own ``python -m ...`` process side-steps both — and is +exactly what "run each step separately to validate" needs. + +This launcher invents no parameters: it only wires the default output of one +stage into the input of the next (reusing each script's own default paths) and +forwards any extra flags straight through to the selected stage. + +Examples +-------- +Validate one stage at a time (extra flags after the known ones are forwarded): + + python cookbook/rsi/run_rsi.py --step prepare --raw data/raw.jsonl + python cookbook/rsi/run_rsi.py --step refine --teacher-model qwen3-235b-a22b-instruct-2507 + python cookbook/rsi/run_rsi.py --step rl --model.model_id ms://Qwen/Qwen3-4B --infra.model_gpus 4 + python cookbook/rsi/run_rsi.py --step distill --dump output/rsi/dump/refine.jsonl --adapter refine + +Run the whole chain with default paths (each stage still a fresh process): + + python cookbook/rsi/run_rsi.py --step all --raw data/raw.jsonl +""" +import argparse +import os +import subprocess +import sys + +# Default paths chain one stage into the next. These mirror the defaults baked +# into each stage's own CLI, kept here so --step all wires up with no flags. +DEFAULT_SUBSET = 'output/rsi/subset.jsonl' # rsi_prepare --output +DEFAULT_FLOWS = 'output/rsi/standard_flows.jsonl' # rsi_refine --output / rsi_rl RSI_STD_FLOWS +DEFAULT_DUMP = 'output/rsi/dump/refine.jsonl' # llm_backup LLM_BACKUP_DUMP_PATH / rsi_distill --input + +MODULES = { + 'prepare': 'twinkle_agentic.rsi.rsi_prepare', + 'refine': 'twinkle_agentic.rsi.rsi_refine', + 'rl': 'twinkle_agentic.rsi.rsi_rl', + 'distill': 'twinkle_agentic.rsi.rsi_distill', +} +ORDER = ['prepare', 'refine', 'rl', 'distill'] + + +def _run(module: str, argv: list, env: dict) -> None: + """Run ``python -m module argv...`` as a child process, streaming its output. + + Raises on non-zero exit so --step all stops at the first failing stage + instead of silently feeding a broken artifact into the next stage. + """ + cmd = [sys.executable, '-m', module] + argv + print(f'\n[run_rsi] $ {" ".join(cmd)}', flush=True) + subprocess.run(cmd, env=env, check=True) + + +def _argv_for(step: str, a: argparse.Namespace, extra: list) -> tuple: + """Build (argv, env) for one stage. ``extra`` is forwarded verbatim so each + stage's own flags (teacher creds, twinkle CLI knobs, ...) still work.""" + env = dict(os.environ) + if step == 'prepare': + if not a.raw: + raise SystemExit('[run_rsi] --step prepare 需要 --raw 指向原始数据源') + argv = ['--input', a.raw, '--output', a.subset, '--num-proc', str(a.num_proc)] + if a.dropped_log: + argv += ['--dropped-log', a.dropped_log] + return argv + extra, env + if step == 'refine': + return ['--input', a.subset, '--output', a.flows] + extra, env + if step == 'rl': + # rsi_rl reads the standard-flow path from an env var, not a flag; + # model/infra/rl knobs arrive through `extra` (twinkle CLI). + env['RSI_STD_FLOWS'] = a.flows + return list(extra), env + if step == 'distill': + # rsi_distill accepts --input/--adapter and also honours these env vars. + env['RSI_DUMP_PATH'] = a.dump + argv = ['--input', a.dump] + if a.adapter: + argv += ['--adapter', a.adapter] + return argv + extra, env + raise SystemExit(f'[run_rsi] 未知 step: {step}') + + +def main(): + parser = argparse.ArgumentParser( + description='RSI pipeline launcher — run one stage (validate) or the whole chain.', + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument('--step', required=True, choices=ORDER + ['all'], + help='Which stage to run (or "all" for prepare->refine->rl->distill).') + parser.add_argument('--raw', default='', help='Raw data source for prepare (local path or ms:// id).') + parser.add_argument('--subset', default=DEFAULT_SUBSET, help='prepare output / refine input.') + parser.add_argument('--flows', default=DEFAULT_FLOWS, help='refine output / rl standard-flow input.') + parser.add_argument('--dump', default=os.environ.get('LLM_BACKUP_DUMP_PATH', DEFAULT_DUMP), + help='llm_backup dump JSONL / distill input.') + parser.add_argument('--adapter', default='', help='Adapter name for distill (default: derived from dump name).') + parser.add_argument('--num-proc', type=int, default=int(os.environ.get('RSI_NUM_PROC', '4')), + help='Parallel workers for prepare.') + parser.add_argument('--dropped-log', default='', help='Optional dropped-row log for prepare.') + a, extra = parser.parse_known_args() + + if a.step == 'all': + if extra: + # For 'all' the extras are ambiguous (which stage?); refuse rather than + # forward a flag to a stage that does not accept it. + raise SystemExit(f'[run_rsi] --step all 不接受透传参数 {extra}；请逐个 --step 跑并各自带参数') + for step in ORDER: + if step == 'prepare' and not a.raw: + raise SystemExit('[run_rsi] --step all 需要 --raw 指向原始数据源') + argv, env = _argv_for(step, a, []) + _run(MODULES[step], argv, env) + print('\n[run_rsi] all stages done.', flush=True) + return + + argv, env = _argv_for(a.step, a, extra) + _run(MODULES[a.step], argv, env) + + +if __name__ == '__main__': + main() diff --git a/src/twinkle_agentic/rsi/__init__.py b/src/twinkle_agentic/rsi/__init__.py new file mode 100644 index 000000000..eb7bb499b --- /dev/null +++ b/src/twinkle_agentic/rsi/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""RSI (recursive self-improvement) pipeline scripts. + +Stages (each a standalone entry script): + rsi_prepare.py - step 1: read a raw source, parallel-preprocess, dump a subset. + rsi_refine.py - step 2: re-analyze/strengthen trajectories into a standard flow. + rsi_rl.py - step 3: multi-LoRA RL, one training query per round. + rsi_distill.py - step 4: dump llm_backup data, SFT the auxiliary-role LoRAs. +""" diff --git a/src/twinkle_agentic/rsi/rsi_distill.py b/src/twinkle_agentic/rsi/rsi_distill.py new file mode 100644 index 000000000..bdedc8c00 --- /dev/null +++ b/src/twinkle_agentic/rsi/rsi_distill.py @@ -0,0 +1,285 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""RSI step 4 — turn the data collected by ``llm_backup`` into a per-role LoRA +via plain SFT. + +Where the data comes from +------------------------- +``twinkle_agentic.utils.llm_backup`` now optionally appends one raw record per +teacher call to the JSONL at ``$LLM_BACKUP_DUMP_PATH`` (off unless that env var +is set). Each line is:: + + {"key": <hash>, "trajectory": <the exact model input>, + "student": <student output str>, "teacher": <teacher output str>, + "match": <bool>} + +``trajectory`` is whatever the decorated role passed as its ``trajectory`` arg — +in twinkle_agentic that is an ``{"messages": [...], "tools": [...]}`` dict (see +protocol/openai.py). It is stored verbatim so it can be reshaped here into an +SFT pair without re-running anything. + +What this script trains +------------------------ +SFT target = EVERY teacher output (the ``match`` flag is ignored; decided by the +user). One training sample is:: + + messages = <trajectory messages> + [{"role": "assistant", "content": teacher}] + +and only that final assistant turn is trainable (``key_rounds=[len(msgs)-1]``), +so the LoRA learns to reproduce the teacher's output for that role. Run the +script once per auxiliary role, each with its own dump file and adapter name — +that is the composable unit; there is no multi-role loop here on purpose. + +Status: this is the PLUMBING. It only runs distillation when invoked explicitly +(``python -m twinkle_agentic.rsi.rsi_distill --input ...``); importing it does +nothing. + +Numbers are INHERITED, not invented: + * LR / BATCH_SIZE / MICRO_BATCH / EPOCHS / MAX_MODEL_LEN <- e18_sft_kod.py + * LORA_RANK / alpha=rank*2 / dropout=0.05 <- rsi_rl.py +All are overridable via the env vars below. + +Env vars +-------- + RSI_DUMP_PATH dump JSONL to read (default $LLM_BACKUP_DUMP_PATH) + RSI_ADAPTER adapter name to train + save (default: derived from dump name) + OUTPUT_DIR where the adapter is written (default output/rsi/distill) + MODEL_ID base model (default Qwen/Qwen3-4B) + plus TRAIN_GPUS / TRAIN_FSDP / GPU_MEM / SEED / EPOCHS / BATCH_SIZE / + MICRO_BATCH / LR / MAX_MODEL_LEN / LORA_RANK (all inherited defaults). +""" +import argparse +import json +import os +import random +import shutil +import time +from typing import Any, Dict, List + +from peft import LoraConfig + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import pack_user_data +from twinkle.model import TransformersModel +from twinkle.processor import InputProcessor +from twinkle.template import Template + +logger = get_logger() + +# ── config (env, all defaults inherited from e18_sft_kod.py / rsi_rl.py) ──── +MODEL_ID = os.environ.get('MODEL_ID', 'Qwen/Qwen3-4B') +OUTPUT_DIR = os.environ.get('OUTPUT_DIR', os.path.join('output', 'rsi', 'distill')) + +TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 4)) +TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', 1)) +TRAIN_DP = TRAIN_GPUS // TRAIN_FSDP +GPU_MEM = float(os.environ.get('GPU_MEM', 0.8)) + +SEED = int(os.environ.get('SEED', 42)) +EPOCHS = float(os.environ.get('EPOCHS', 1)) +BATCH_SIZE = int(os.environ.get('BATCH_SIZE', 16)) +MICRO_BATCH = int(os.environ.get('MICRO_BATCH', 8)) +LR = float(os.environ.get('LR', 1e-5)) +MAX_MODEL_LEN = int(os.environ.get('MAX_MODEL_LEN', 16000)) +LORA_RANK = int(os.environ.get('LORA_RANK', 16)) + +LOG_EVERY_STEPS = int(os.environ.get('LOG_EVERY_STEPS', 1)) +RUN_ID = time.strftime('%m%d-%H%M%S') + + +# =========================================================================== +# data: llm_backup dump -> SFT samples (target = every teacher output) +# =========================================================================== +def _trajectory_messages(traj: Any) -> List[Dict[str, Any]]: + """Pull the message list out of a dumped ``trajectory``. + + twinkle_agentic passes trajectory as ``{"messages": [...], "tools": ...}``; + tolerate a bare list of messages too so the loader does not depend on one + role's exact calling convention. + """ + if isinstance(traj, dict): + msgs = traj.get('messages') + elif isinstance(traj, list): + msgs = traj + else: + msgs = None + return msgs if isinstance(msgs, list) else [] + + +def load_samples(dump_path: str) -> List[Dict[str, Any]]: + """Read the llm_backup dump and build one SFT sample per usable record. + + A record is usable when it has a non-empty trajectory message list AND a + non-empty teacher string. The teacher output becomes the sole trainable + assistant turn appended to the trajectory messages. ``match`` is ignored on + purpose (target = every teacher output). + """ + if not os.path.exists(dump_path): + raise FileNotFoundError(f'找不到 llm_backup dump：{dump_path}') + raw, bad_json = [], 0 + with open(dump_path, 'r', encoding='utf-8') as f: + for line in f: + line = line.strip() + if not line: + continue + try: + raw.append(json.loads(line)) + except Exception: + bad_json += 1 + drop = {'no_messages': 0, 'no_teacher': 0} + out = [] + for r in raw: + msgs = _trajectory_messages(r.get('trajectory')) + if not msgs: + drop['no_messages'] += 1 + continue + teacher = r.get('teacher') + if not isinstance(teacher, str) or not teacher.strip(): + drop['no_teacher'] += 1 + continue + sample_msgs = list(msgs) + [{'role': 'assistant', 'content': teacher}] + out.append({'messages': sample_msgs, 'key': r.get('key', '')}) + logger.info(f'[data] 读入 {len(raw)} 条' + + (f'（跳过 {bad_json} 行半行）' if bad_json else '') + + f'，可训 {len(out)} 条，丢弃明细 {drop}') + return out + + +def make_trajs(batch: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """样本 -> twinkle 轨迹。只把最后一轮（teacher 输出）标为可训区。""" + trajs = [] + for s in batch: + msgs = s['messages'] + trajs.append({'messages': msgs, + 'user_data': pack_user_data({'key_rounds': [len(msgs) - 1]})}) + return trajs + + +# =========================================================================== +# model: base + one LoRA adapter (rsi_rl.py's config, TransformersModel SFT) +# =========================================================================== +def build_model(adapter_name: str): + twinkle.initialize(mode='ray', nproc_per_node=TRAIN_GPUS, lazy_collect=False, groups=[ + DeviceGroup(name='train', ranks=list(range(TRAIN_GPUS)), device_type='GPU')]) + model = TransformersModel( + model_id=MODEL_ID, remote_group='train', + device_mesh=DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_DP, + fsdp_size=TRAIN_FSDP), + ddp_config={'find_unused_parameters': False}) + # enable_thinking=True: auxiliary roles produce reasoning before their answer, + # same deployment mode as the teacher that generated the targets. + model.set_template(Template, model_id=MODEL_ID, enable_thinking=True, + max_length=MAX_MODEL_LEN, truncation_strategy='delete') + model.set_processor(InputProcessor, padding_free=False) + model.set_loss('CrossEntropyLoss') + lora_cfg = LoraConfig(target_modules='all-linear', r=LORA_RANK, + lora_alpha=LORA_RANK * 2, lora_dropout=0.05) + model.add_adapter_to_model(adapter_name, lora_cfg, + gradient_accumulation_steps=1) + model.set_optimizer('AdamW', lr=LR) + return model + + +def step_metrics(model) -> Dict[str, float]: + """取本 step 的优化指标（loss/grad_norm/lr）。twinkle 把 loss 格式化成字符串， + 所以用 float() 试转而不是 isinstance 判数值型，否则会静默丢掉 loss。""" + out: Dict[str, float] = {} + for k, val in (model.calculate_metric(is_training=True) or {}).items(): + if isinstance(val, bool): + continue + try: + fval = float(val) + except (TypeError, ValueError): + continue + if k.startswith('learning rate'): + if 'group 1' in k: + out['lr'] = fval + else: + out[k.replace(' ', '_')] = fval + return out + + +def archive_output_dir() -> None: + """启动时把已存在的非空 OUTPUT_DIR 整个 mv 走（sft_log.jsonl 是追写的）， + 否则重跑会把两条 loss 曲线焊进一个文件。与 e18_sft_kod 同一套机制。""" + if not os.path.isdir(OUTPUT_DIR) or not os.listdir(OUTPUT_DIR): + os.makedirs(OUTPUT_DIR, exist_ok=True) + return + stamp = time.strftime('%m%d-%H%M%S', time.localtime(os.path.getmtime(OUTPUT_DIR))) + dst = f'{OUTPUT_DIR}.bak-{stamp}' + i = 1 + while os.path.exists(dst): + dst = f'{OUTPUT_DIR}.bak-{stamp}-{i}' + i += 1 + shutil.move(OUTPUT_DIR, dst) + os.makedirs(OUTPUT_DIR, exist_ok=True) + logger.info(f'[init] 旧输出目录已归档 -> {dst}') + + +# =========================================================================== +# train +# =========================================================================== +def main(): + parser = argparse.ArgumentParser(description='RSI step 4: distill one auxiliary-role LoRA from an llm_backup dump.') + parser.add_argument('--input', default=os.environ.get('RSI_DUMP_PATH', os.environ.get('LLM_BACKUP_DUMP_PATH')), + help='llm_backup dump JSONL (default $RSI_DUMP_PATH / $LLM_BACKUP_DUMP_PATH)') + parser.add_argument('--adapter', default=os.environ.get('RSI_ADAPTER'), + help='adapter name to train and save (default: derived from dump filename)') + args = parser.parse_args() + + dump_path = args.input + if not dump_path: + raise SystemExit('必须给 --input（或设 $RSI_DUMP_PATH / $LLM_BACKUP_DUMP_PATH）指向 llm_backup dump') + adapter_name = args.adapter or os.path.splitext(os.path.basename(dump_path))[0] + + t0 = time.time() + archive_output_dir() + samples = load_samples(dump_path) + if len(samples) < BATCH_SIZE: + raise RuntimeError(f'可用样本 {len(samples)} 条 < BATCH_SIZE {BATCH_SIZE}') + if BATCH_SIZE % TRAIN_DP: + raise RuntimeError(f'BATCH_SIZE({BATCH_SIZE}) 必须是 TRAIN_DP({TRAIN_DP}) 的整倍数') + + model = build_model(adapter_name) + steps_per_epoch = len(samples) // BATCH_SIZE + total_steps = int(steps_per_epoch * EPOCHS) + logger.info(f'RSI-DISTILL start: adapter={adapter_name} n={len(samples)} bs={BATCH_SIZE} ' + f'micro={MICRO_BATCH} lr={LR} epochs={EPOCHS} steps/epoch={steps_per_epoch} ' + f'total_steps={total_steps} rank={LORA_RANK} gpus={TRAIN_GPUS} out={OUTPUT_DIR}') + + log_path = os.path.join(OUTPUT_DIR, 'sft_log.jsonl') + rng = random.Random(SEED) + step = 0 + with open(log_path, 'a', encoding='utf-8') as log_fh: + epoch = 0 + while step < total_steps: + order = list(range(len(samples))) + rng.shuffle(order) # 每 epoch 重洗，种子固定所以可复现 + for bi in range(steps_per_epoch): + if step >= total_steps: + break + batch = [samples[j] for j in order[bi * BATCH_SIZE:(bi + 1) * BATCH_SIZE]] + trajs = make_trajs(batch) + micro = max(TRAIN_DP, min(MICRO_BATCH, len(trajs))) + t_step = time.time() + for i in range(0, len(trajs), micro): + model.forward_backward(inputs=trajs[i:i + micro]) + model.clip_grad_and_step() + step += 1 + row = {'step': step, 'epoch': epoch, 'run': RUN_ID, 'adapter': adapter_name, + 'n_samples': len(batch), 'seconds': round(time.time() - t_step, 2)} + row.update(step_metrics(model)) + log_fh.write(json.dumps(row, ensure_ascii=False) + '\n') + log_fh.flush() + if step % LOG_EVERY_STEPS == 0: + logger.info('[s%d/%d ep%d] ' % (step, total_steps, epoch) + + ' '.join(f'{k}={v:.4g}' for k, v in row.items() + if isinstance(v, float))) + epoch += 1 + + ckpt = model.save(f'{adapter_name}-final', output_dir=OUTPUT_DIR, adapter_name=adapter_name) + logger.info(f'[done] steps={step} 用时 {(time.time() - t0) / 60:.1f} 分钟 -> {ckpt}') + + +if __name__ == '__main__': + main() diff --git a/src/twinkle_agentic/rsi/rsi_prepare.py b/src/twinkle_agentic/rsi/rsi_prepare.py new file mode 100644 index 000000000..dd4fefdef --- /dev/null +++ b/src/twinkle_agentic/rsi/rsi_prepare.py @@ -0,0 +1,175 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""RSI step 1 — read a raw data source, parallel-preprocess it with the +twinkle_agentic preprocessor, and write the surviving subset to disk. + +Usage +----- + python -m twinkle_agentic.rsi.rsi_prepare \ + --input /path/to/raw.jsonl \ + --output output/rsi/subset.jsonl \ + --num-proc 4 + +``--input`` accepts a local ``.jsonl``/``.parquet`` path or an ``ms://`` dataset +id; the raw schema is intentionally not pinned here (decided per data source at +test time). Every row must expose a ``messages`` list — the preprocessor keys +off it. Adapt other schemas in :func:`load_source` before the pipeline runs. + +Pipeline +-------- +Core steps (no external deps, always on), each using the filter's OWN default +thresholds (no thresholds invented here): + + MessageNormalizer -> MessageSanityFilter -> RefuseFilter -> DeadLoopFilter + -> TokenSoupFilter -> HardFilter + +Optional steps, off by default (enabling needs extra packages): + RSI_USE_LANG=1 LanguageFilter (langid, degrades to heuristic) + RSI_USE_DATAJUICER=1 FixUnicode/RemoveRepeat/SpecialChars/TokenNum (data_juicer[, modelscope]) + RSI_USE_PII=1 PIIPresidioFilter (presidio-analyzer/anonymizer) + +``DedupFilter`` is NOT part of the parallel pipeline: its docstring requires it +to see the whole dataset in one call, so it runs once after the parallel pass. +""" +import argparse +import os + +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.utils import get_logger +from twinkle_agentic.preprocessor import (DeadLoopFilter, DedupFilter, HardFilter, MessageNormalizer, + MessageSanityFilter, QualityPreprocessor, RefuseFilter, TokenSoupFilter, + merge_dropped_shards, run_quality_pipeline, truncate_dropped_logs) + +logger = get_logger() + + +def _env_flag(name: str, default: str = '0') -> bool: + return os.environ.get(name, default).strip().lower() in ('1', 'true', 'yes', 'on') + + +def build_pipeline(): + """Return the ordered list of preprocessor steps for the parallel pass. + + DedupFilter is deliberately excluded (see module docstring); it is applied + separately on the full materialized dataset. + """ + steps = [ + MessageNormalizer(), # strip heartbeats, rewrite tool calls, merge consecutive roles + MessageSanityFilter(), # role order / tool-id matching / content integrity / sensitive words + RefuseFilter(), # drop assistant self-referential refusals + DeadLoopFilter(), # drop degenerate / stuck (hesitation, cascade, ngram repeat) + TokenSoupFilter(), # drop garbled text (replacement/control/private-use chars, script chaos) + # min_assistant_chars_2turn=0: a single-turn valid tool call (e.g. `[Func(x=1)]`) + # is only tens of chars; HardFilter's default 80-char floor wrongly drops it + # as a "shallow_reply". Zeroing the floor keeps these tool-call rows (Rule 3 + # still removes genuinely empty assistants). Overridable via env. + HardFilter(min_assistant_chars_2turn=int(os.environ.get('RSI_MIN_ASST_CHARS_2TURN', 0))), + ] + if _env_flag('RSI_USE_LANG'): + from twinkle_agentic.preprocessor import LanguageFilter + steps.append(LanguageFilter()) + if _env_flag('RSI_USE_DATAJUICER'): + from twinkle_agentic.preprocessor import (FixUnicodeFilter, RemoveRepeatSentencesFilter, SpecialCharsFilter, + TokenNumFilter) + steps += [FixUnicodeFilter(), RemoveRepeatSentencesFilter(), SpecialCharsFilter(), TokenNumFilter()] + if _env_flag('RSI_USE_PII'): + from twinkle_agentic.preprocessor import PIIPresidioFilter + steps.append(PIIPresidioFilter()) + return steps + + +# ShareGPT `from` value -> standard message role. ToolACE uses +# system/user/assistant/tool; other ShareGPT variants use human/gpt/observation. +_ROLE_MAP = { + 'system': 'system', + 'user': 'user', 'human': 'user', + 'assistant': 'assistant', 'gpt': 'assistant', 'bot': 'assistant', + 'tool': 'tool', 'observation': 'tool', 'function': 'tool', + 'function_call': 'assistant', 'function_response': 'tool', 'tool_response': 'tool', +} + + +def _row_to_messages(row: dict) -> dict: + """Map one ShareGPT ``conversations`` row to a ``messages`` row. + + Only ``from``->``role`` and ``value``->``content`` are rewritten; the tool + call embedded in an assistant turn is left as-is in ``content`` (ToolACE keeps + it as a bracket-DSL string) and parsed later in rsi_refine/rsi_rl. Turns whose + ``from`` is unknown are dropped so no invalid role reaches the pipeline. + """ + messages = [] + for turn in (row.get('conversations') or []): + if not isinstance(turn, dict): + continue + role = _ROLE_MAP.get(str(turn.get('from', '')).lower()) + if role is None: + continue + messages.append({'role': role, 'content': turn.get('value', '') or ''}) + return {'messages': messages, 'id': row.get('id', '')} + + +def load_source(input_path: str, num_proc: int = 4) -> Dataset: + """Load the raw source into a twinkle Dataset. + + A local path is loaded by extension (jsonl->json, parquet, csv...); anything + else is treated as a hub id (e.g. ``ms://org/name``). Rows are passed through + unchanged except for one adaptation: ShareGPT-style rows (a ``conversations`` + list of ``{"from", "value"}`` turns, e.g. ToolACE) are mapped to a standard + ``messages`` list, because the whole preprocessor keys off ``messages``. Rows + that already carry ``messages`` are left untouched. + """ + ds = Dataset(DatasetMeta(dataset_id=input_path)) + cols = ds.dataset.column_names + if 'messages' not in cols and 'conversations' in cols: + logger.info('[rsi_prepare] ShareGPT `conversations` detected -> mapping to `messages`') + # Materialize + convert in Python then rebuild: twinkle's Dataset.map forces + # batched=True and wraps the fn as a Preprocessor, which does not fit a plain + # per-row schema rewrite. The source is small enough to hold in memory. + rows = [_row_to_messages(r) for r in ds.dataset.to_list()] + ds = Dataset(DatasetMeta(data=rows)) + return ds + + +def main(): + parser = argparse.ArgumentParser(description='RSI step 1: preprocess a raw source into a clean subset.') + parser.add_argument('--input', required=True, help='Local .jsonl/.parquet path or an ms:// dataset id.') + parser.add_argument('--output', default='output/rsi/subset.jsonl', help='Where to write the surviving subset.') + parser.add_argument('--num-proc', type=int, default=int(os.environ.get('RSI_NUM_PROC', '4')), + help='Parallel workers for the preprocessor map pass.') + parser.add_argument('--dropped-log', default='', help='Optional JSONL of dropped-row metadata (empty=off).') + args = parser.parse_args() + + os.makedirs(os.path.dirname(os.path.abspath(args.output)) or '.', exist_ok=True) + + pipeline = build_pipeline() + step_names = [type(s).__name__ for s in pipeline] + logger.info(f'[rsi_prepare] pipeline: {" -> ".join(step_names)} + DedupFilter(global)') + + dataset = load_source(args.input, num_proc=args.num_proc) + n_in = len(dataset.dataset) + logger.info(f'[rsi_prepare] loaded {n_in} rows from {args.input}') + + # 'mark' mode + run_quality_pipeline is the ghost-proof parallel path: + # map returns equal-length columns flagged _keep, then a single filter removes. + if args.dropped_log: + truncate_dropped_logs(args.dropped_log) + qp = QualityPreprocessor(pipeline, dropped_log_path=args.dropped_log, drop_mode='mark') + run_quality_pipeline(dataset, qp, num_proc=args.num_proc) + if args.dropped_log: + merge_dropped_shards(args.dropped_log) + + n_after_pipeline = len(dataset.dataset) + logger.info(f'[rsi_prepare] after parallel pipeline: {n_in} -> {n_after_pipeline}') + + # Global longest-wins dedup — must see the whole dataset at once. + rows = dataset.dataset.to_list() + kept, dropped = DedupFilter()(rows) + logger.info(f'[rsi_prepare] after global dedup: {n_after_pipeline} -> {len(kept)} ' + f'(dropped {len(dropped)} duplicates)') + + out = Dataset(DatasetMeta(data=kept)) + out.save_as(args.output) + logger.info(f'[rsi_prepare] wrote {len(kept)} rows -> {args.output}') + + +if __name__ == '__main__': + main() diff --git a/src/twinkle_agentic/rsi/rsi_refine.py b/src/twinkle_agentic/rsi/rsi_refine.py new file mode 100644 index 000000000..fbfd90724 --- /dev/null +++ b/src/twinkle_agentic/rsi/rsi_refine.py @@ -0,0 +1,277 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""RSI step 2 — re-analyze each preprocessed trajectory into a STANDARD solving +flow via an injectable teacher API, mark the key rounds, and attach a per-round +reward method. + +Input : the subset produced by rsi_prepare.py (rows with a ``messages`` list). +Output: one refined record per trajectory that could be organized: + { + "id": <passthrough id if present>, + "system": <original system message, kept verbatim (holds tool defs)>, + "query": <original first user message, kept verbatim>, + "tools": <original tools, kept verbatim>, + "rounds": [ {intent, type, tool_call, result, code, reward_method} ], + } +Trajectories the teacher marks unorganizable (e.g. missing tools) are written to +a separate ``*.unorganizable.jsonl`` log and excluded from the standard set. + +Design +------ +- The teacher is INJECTABLE: any OpenAI-compatible endpoint, chosen at runtime + via --teacher-model / --teacher-base-url / --teacher-api-key (env fallbacks + RSI_TEACHER_* then LLM_BACKUP_*). Nothing about the model is hardcoded. +- Key round = a round that carries a tool call OR a code block. The reward + method is attached automatically by round content: + tool call present -> 'tool_result' (executable, verifiable) + code only -> 'rubric' (no executable signal here) + Concrete reward thresholds are intentionally left unset (decided in step 3). +- Heartbeat rounds are already stripped upstream by MessageNormalizer (step 1), + so this stage does not re-handle them. +""" +import argparse +import json +import os +import re +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Any, Dict, List, Optional, Tuple + +from twinkle.data_format import SamplingParams +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.utils import get_logger +from twinkle_agentic.protocol.openai import OpenAI + +logger = get_logger() + +REWARD_TOOL_RESULT = 'tool_result' +REWARD_RUBRIC = 'rubric' + +# Runtime prompt is English on purpose: trajectories are predominantly English, +# and mixing languages degrades the teacher. A Chinese rendering was reviewed +# and approved separately. +REORG_SYSTEM = """\ +You are given ONE agent trajectory that solves a single task through multiple \ +rounds of tool calls. Produce the STANDARD solving flow for this task as \ +STRUCTURED JSON, so it can be parsed. + +The original system message (which holds the tool definitions) and the original \ +user query are kept separately — do NOT rewrite them. Your job is to output only \ +the cleaned, correctly-ordered sequence of KEY rounds: the tool calls and their \ +results that materially lead to the solution. + +Rules: +- Work ONLY from what actually happens in the trajectory and its real tool \ +results. Do NOT invent tools, arguments, or results that are not present. +- Remove redundant, failed-and-abandoned, or out-of-order rounds; reorder the \ +remaining rounds into the logical order that reaches the solution. +- Preserve each kept round's tool call (name + arguments) and the tool's \ +returned result verbatim. +- In each tool_call's "arguments", include ONLY the parameters that were \ +actually passed in that call. Do NOT enumerate the tool's full parameter \ +schema, and do NOT add keys whose value is null/empty for unused parameters. +- Write it as a FRESH, clean standard procedure. Do NOT say the original was \ +wrong, and do NOT reference "previous attempts", "the original code", or "the \ +error above". +- If the trajectory cannot be organized into a standard flow (e.g. required \ +tools are missing, the task never reaches a solution), output exactly \ +{"unorganizable": true, "reason": "<short reason>"} and nothing else. +- Otherwise output ONLY valid JSON in this schema (no prose outside the JSON): +{ + "rounds": [ + {"intent": "<one line>", + "type": "tool" or "code", + "tool_call": {"name": "...", "arguments": {...}} or null, + "result": "<verbatim tool/exec result>", + "code": "<code text if type==code, else null>"} + ] +} +""" + +_CODE_FENCE_RE = re.compile(r'```') +_JSON_FENCE_RE = re.compile(r'^\s*```(?:json)?\s*|\s*```\s*$', re.IGNORECASE) + + +def _first_role(messages: List[Dict[str, Any]], role: str) -> Optional[Dict[str, Any]]: + for m in messages: + if isinstance(m, dict) and m.get('role') == role: + return m + return None + + +def _strip_json_fence(text: str) -> str: + """Remove a leading ```json / trailing ``` wrapper if the model added one.""" + text = text.strip() + text = _JSON_FENCE_RE.sub('', text) + return text.strip() + + +def _strip_null_args(tool_call: Any) -> Any: + """Drop arguments whose value is null/empty from a tool_call. + + Backstop for teachers that echo the tool's full parameter schema and pad + unused params with null: the original calls only pass real args, so a null + here is invented noise that would break step-3 argument matching. Keys with + value None or '' are removed; the rest are kept verbatim. + """ + if not isinstance(tool_call, dict): + return tool_call + args = tool_call.get('arguments') + if isinstance(args, dict): + tool_call['arguments'] = {k: v for k, v in args.items() if v is not None and v != ''} + return tool_call + + +def attach_reward_method(round_obj: Dict[str, Any]) -> str: + """Decide the reward method from the round's actual content (not the label). + + A round with a tool call is verifiable by its tool result; a code-only round + has no executable signal here, so it is scored by rubric. + """ + if round_obj.get('tool_call'): + return REWARD_TOOL_RESULT + if round_obj.get('code') or (isinstance(round_obj.get('type'), str) and round_obj['type'] == 'code'): + return REWARD_RUBRIC + # Fallback: treat as rubric (no tool call, no code detected). + return REWARD_RUBRIC + + +def build_teacher(args) -> OpenAI: + """Construct the injectable teacher client from CLI/env (nothing hardcoded).""" + model = args.teacher_model or os.environ.get('RSI_TEACHER_MODEL') or os.environ.get('LLM_BACKUP_MODEL') + base_url = args.teacher_base_url or os.environ.get('RSI_TEACHER_BASE_URL') or os.environ.get('LLM_BACKUP_BASE_URL') + api_key = args.teacher_api_key or os.environ.get('RSI_TEACHER_API_KEY') or os.environ.get('LLM_BACKUP_API_KEY') + if not model: + raise ValueError('No teacher model given. Pass --teacher-model or set RSI_TEACHER_MODEL/LLM_BACKUP_MODEL.') + timeout = float(os.environ.get('RSI_TEACHER_TIMEOUT', '120')) + max_retries = int(os.environ.get('RSI_TEACHER_MAX_RETRIES', '2')) + return OpenAI(model=model, api_key=api_key, base_url=base_url, + client_kwargs={'timeout': timeout, 'max_retries': max_retries}) + + +def reorder_workflow(traj: Dict[str, Any], teacher: OpenAI, + sampling_params: SamplingParams) -> Tuple[Optional[Dict[str, Any]], Optional[str]]: + """Ask the teacher to reorganize one trajectory into the standard-flow JSON. + + Returns (parsed_json, None) on success, or (None, reason) when the teacher + declares it unorganizable or the response is not parseable. + """ + payload = json.dumps({'messages': traj.get('messages', []), 'tools': traj.get('tools', [])}, + ensure_ascii=False) + request = {'messages': [{'role': 'system', 'content': REORG_SYSTEM}, + {'role': 'user', 'content': payload}]} + message = teacher(request, sampling_params) + if isinstance(message, list): + message = message[0] if message else {} + content = message.get('content', '') if isinstance(message, dict) else '' + if not content.strip(): + return None, 'empty_teacher_response' + try: + parsed = json.loads(_strip_json_fence(content)) + except (ValueError, TypeError): + return None, 'unparseable_json' + if parsed.get('unorganizable'): + return None, f"unorganizable:{parsed.get('reason', '')}" + if not isinstance(parsed.get('rounds'), list) or not parsed['rounds']: + return None, 'no_rounds' + return parsed, None + + +def refine_one(traj: Dict[str, Any], teacher: OpenAI, + sampling_params: SamplingParams) -> Tuple[Optional[Dict[str, Any]], Optional[str]]: + """Turn one raw trajectory into a standard-flow record (or a drop reason).""" + messages = traj.get('messages') or [] + system_msg = _first_role(messages, 'system') + query_msg = _first_role(messages, 'user') + if query_msg is None: + return None, 'no_user_query' + + parsed, reason = reorder_workflow(traj, teacher, sampling_params) + if parsed is None: + return None, reason + + rounds = [] + for r in parsed['rounds']: + if not isinstance(r, dict): + continue + r = dict(r) + if r.get('tool_call'): + r['tool_call'] = _strip_null_args(r['tool_call']) + r['reward_method'] = attach_reward_method(r) + rounds.append(r) + if not rounds: + return None, 'no_valid_rounds' + + record = { + 'id': traj.get('id'), + 'system': system_msg, + 'query': query_msg, + 'tools': traj.get('tools', []), + 'rounds': rounds, + } + return record, None + + +def main(): + parser = argparse.ArgumentParser(description='RSI step 2: refine trajectories into standard solving flows.') + parser.add_argument('--input', default='output/rsi/subset.jsonl', help='Subset from rsi_prepare.py.') + parser.add_argument('--output', default='output/rsi/standard_flows.jsonl', help='Refined standard-flow records.') + parser.add_argument('--teacher-model', default='', help='Teacher model id (or RSI_TEACHER_MODEL/LLM_BACKUP_MODEL).') + parser.add_argument('--teacher-base-url', default='', help='Teacher endpoint base url.') + parser.add_argument('--teacher-api-key', default='', help='Teacher API key.') + parser.add_argument('--max-workers', type=int, default=int(os.environ.get('RSI_MAX_WORKERS', '8')), + help='Concurrent teacher API calls.') + # Generation knobs (defaults shown; override to taste). Low temperature keeps + # the reformat deterministic; max-tokens bounds the JSON output size. + parser.add_argument('--temperature', type=float, default=float(os.environ.get('RSI_TEACHER_TEMPERATURE', '0.0'))) + parser.add_argument('--max-tokens', type=int, default=int(os.environ.get('RSI_TEACHER_MAX_TOKENS', '8192'))) + args = parser.parse_args() + + os.makedirs(os.path.dirname(os.path.abspath(args.output)) or '.', exist_ok=True) + unorg_path = os.path.splitext(args.output)[0] + '.unorganizable.jsonl' + + teacher = build_teacher(args) + sampling_params = SamplingParams(max_tokens=args.max_tokens, num_samples=1, + temperature=args.temperature, top_p=1.0) + + rows = Dataset(DatasetMeta(dataset_id=args.input)).dataset.to_list() + logger.info(f'[rsi_refine] loaded {len(rows)} trajectories from {args.input}') + + kept: List[Dict[str, Any]] = [] + dropped: List[Dict[str, Any]] = [] + total = len(rows) + # Log progress every ~5% (at least every 1) so a long run is not a black box. + step = max(1, total // 20) + done = 0 + with ThreadPoolExecutor(max_workers=args.max_workers) as ex: + futures = {ex.submit(refine_one, row, teacher, sampling_params): row for row in rows} + for fut in as_completed(futures): + row = futures[fut] + try: + record, reason = fut.result() + except Exception as e: # noqa: BLE001 + record, reason = None, f'exception:{type(e).__name__}:{e}' + if record is not None: + kept.append(record) + else: + dropped.append({'id': row.get('id'), 'reason': reason}) + done += 1 + if done % step == 0 or done == total: + logger.info(f'[rsi_refine] progress {done}/{total} ({100 * done // total}%) ' + f'kept={len(kept)} dropped={len(dropped)}') + + # Write JSONL directly (NOT via Dataset.save_as): the Arrow table used by + # save_as unifies the per-round ``arguments`` struct across all rows, padding + # every call with the global union of arg names as null (and coercing ints to + # floats). That corrupts the tool calls, so we serialize each record verbatim. + with open(args.output, 'w', encoding='utf-8') as f: + for rec in kept: + f.write(json.dumps(rec, ensure_ascii=False) + '\n') + if dropped: + with open(unorg_path, 'w', encoding='utf-8') as f: + for d in dropped: + f.write(json.dumps(d, ensure_ascii=False) + '\n') + logger.info(f'[rsi_refine] standard flows: {len(kept)}; dropped/unorganizable: {len(dropped)} ' + f'-> {args.output}' + (f' (+ {unorg_path})' if dropped else '')) + + +if __name__ == '__main__': + main() diff --git a/src/twinkle_agentic/rsi/rsi_rl.py b/src/twinkle_agentic/rsi/rsi_rl.py new file mode 100644 index 000000000..fcd947871 --- /dev/null +++ b/src/twinkle_agentic/rsi/rsi_rl.py @@ -0,0 +1,299 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""RSI step 3 — multi-LoRA GRPO where each round of a standard flow becomes its +own training query. + +Idea (confirmed): a multi-turn standard flow is decomposed into one training +query PER key round. For a tool round i the model is shown the fixed prior key +nodes (their tool calls + results) and must roll out {reasoning + the tool call} +for round i. The reward is whether the GENERATED tool call matches the recorded +standard call — name exact + every standard-call argument key/value present in +the generated call (extra args / order ignored). No sandbox is needed because +the standard call is the reference answer. Only the reasoning ("思路") varies +across rollouts; the key node is the target. + +v1 scope: only TOOL rounds are trained. Code rounds (no tool call) still appear +in the prior context but are not turned into training queries yet (their reward +is rubric-based, wired in a later iteration). + +RL data-flow discipline (verified): train ONLY on ``sequence.new_input_feature`` +and use ``sequence.logprobs`` as old_logps — never decode-then-re-encode. The +generated tool call is already parsed into ``new_input_feature['messages'][-1]`` +by the template, and the reference call rides along in ``user_data``. + +Structure mirrors cookbook/rl/grpo/short_math_grpo_multi_lora.py (MultiLoRA +Megatron + filesystem LoRA sync to vLLM). RSI-specific paths come from env vars +so the standard CLI (model/infra/rl knobs) stays identical to the reference: + + RSI_STD_FLOWS standard_flows.jsonl from rsi_refine.py (default output/rsi/standard_flows.jsonl) + RSI_TEMPLATE template name, must match the model (default Qwen3_5Template) + RSI_LORA_SYNC dir for filesystem LoRA sync (default output/rsi/lora_sync) + RSI_ADAPTER executor adapter name (default executor) +""" +import json +import os +from typing import Any, Dict, List, Optional, Tuple + +from peft import LoraConfig + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_device_placement, get_logger +from twinkle.advantage import GRPOAdvantage +from twinkle.cli import CLI +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 MultiLoraMegatronModel +from twinkle.processor import InputProcessor +from twinkle.reward.base import Reward +from twinkle.sampler import vLLMSampler + +logger = get_logger() +args = CLI.from_args() + +# ── RSI-specific paths (env) ─────────────────────────────────────────────── +STD_FLOWS = os.environ.get('RSI_STD_FLOWS', 'output/rsi/standard_flows.jsonl') +TEMPLATE = os.environ.get('RSI_TEMPLATE', 'Qwen3_5Template') +LORA_SYNC_DIR = os.environ.get('RSI_LORA_SYNC', 'output/rsi/lora_sync') +ADAPTER_NAME = os.environ.get('RSI_ADAPTER', 'executor') +REWARD_TOOL_RESULT = 'tool_result' # matches rsi_refine.attach_reward_method + +# ── standard CLI knobs (same shape as the reference script) ──────────────── +MODEL_ID = args.model.model_id or 'ms://Qwen/Qwen3.6-35B-A3B' +MODEL_GPUS = args.infra.model_gpus or 4 +SAMPLER_GPUS = args.infra.sampler_gpus or 2 +SAMPLER_TP = args.sampler.tensor_parallel_size or 2 +NUM_GPUS = MODEL_GPUS + SAMPLER_GPUS +NUM_GENERATIONS = args.rl.num_generations or 8 +MAX_NEW_TOKENS = args.sampling.max_tokens or 4096 +LEARNING_RATE = args.optimizer.learning_rate or 5e-5 +MAX_STEPS = args.training.max_steps or 1000 +BATCH_SIZE = args.training.batch_size or 4 +MINI_BATCH_SIZE = args.training.mini_batch_size or 4 +MICRO_BATCH_SIZE = args.training.micro_batch_size or 1 +GRADIENT_ACCUMULATION_STEPS = args.training.gradient_accumulation_steps or 1 +SAVE_STEPS = args.training.save_steps or 1000 +LORA_RANK = args.lora.lora_r or 16 + +import swanlab +swanlab.init(project='twinkle-rsi') + + +# ── tool-call matching (name exact + standard-call arg subset) ───────────── +def _as_args(a: Any) -> Dict[str, Any]: + if isinstance(a, str): + try: + return json.loads(a) + except (ValueError, TypeError): + return {} + return a or {} + + +def tool_call_matches(gen_call: Optional[Dict[str, Any]], ref_call: Dict[str, Any]) -> bool: + """True iff name matches and every reference arg (key+value) is present.""" + if not gen_call or gen_call.get('name') != ref_call.get('name'): + return False + gen_args = _as_args(gen_call.get('arguments')) + ref_args = _as_args(ref_call.get('arguments')) + for k, v in ref_args.items(): + if k not in gen_args or gen_args[k] != v: + return False + return True + + +class ToolMatchReward(Reward): + """1.0 when the generated tool call matches the recorded standard call.""" + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + rewards = [] + for traj in trajectories: + gen_call = None + for m in reversed(traj.get('messages', []) or []): + if m.get('role') == 'assistant': + tcs = m.get('tool_calls') or [] + if tcs: + gen_call = tcs[0].get('function') + break + ref_call = None + for item in (traj.get('user_data') or []): + if item[0] == 'ref_tool_call': + try: + ref_call = json.loads(item[1]) + except (ValueError, TypeError): + ref_call = None + break + rewards.append(1.0 if (ref_call and tool_call_matches(gen_call, ref_call)) else 0.0) + return rewards + + +# ── decompose standard flows into per-round training trajectories ────────── +def _openai_tool_call(call: Dict[str, Any], idx: int) -> Dict[str, Any]: + args_ = call.get('arguments', {}) + return { + 'id': f'call_{idx}', + 'type': 'function', + 'function': { + 'name': call.get('name', ''), + 'arguments': json.dumps(args_, ensure_ascii=False) if isinstance(args_, dict) else str(args_), + }, + } + + +def _render_prior_round(r: Dict[str, Any], idx: int) -> List[Dict[str, Any]]: + """Render a completed prior round as fixed context messages.""" + result = r.get('result', '') + if r.get('tool_call'): + tc = _openai_tool_call(r['tool_call'], idx) + return [ + {'role': 'assistant', 'content': '', 'tool_calls': [tc]}, + {'role': 'tool', 'content': str(result), 'tool_call_id': tc['id']}, + ] + # Code round: no tool_call_id exists, so keep it template-agnostic. + return [ + {'role': 'assistant', 'content': r.get('code', '') or ''}, + {'role': 'user', 'content': f'[execution result]\n{result}'}, + ] + + +def build_round_trajectories(records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """One training trajectory per TOOL round; prior rounds become fixed context.""" + trajs: List[Dict[str, Any]] = [] + for rec in records: + prefix: List[Dict[str, Any]] = [] + if rec.get('system'): + prefix.append(rec['system']) + if rec.get('query'): + prefix.append(rec['query']) + tools = rec.get('tools') or [] + rounds = rec.get('rounds') or [] + for i, r in enumerate(rounds): + if r.get('reward_method') != REWARD_TOOL_RESULT or not r.get('tool_call'): + continue # v1: train tool rounds only + messages = list(prefix) + for j in range(i): + messages.extend(_render_prior_round(rounds[j], j)) + trajs.append({ + 'messages': messages, + 'tools': tools, + 'user_data': [('ref_tool_call', json.dumps(r['tool_call'], ensure_ascii=False))], + }) + return trajs + + +def create_rsi_dataset(): + records = Dataset(DatasetMeta(dataset_id=STD_FLOWS)).dataset.to_list() + trajs = build_round_trajectories(records) + logger.info(f'[rsi_rl] {len(records)} standard flows -> {len(trajs)} per-round tool queries') + dataset = Dataset(DatasetMeta(data=trajs)) + # enable_thinking=True: we train the reasoning that precedes the tool call. + dataset.set_template(TEMPLATE, model_id=MODEL_ID, max_length=8192, + truncation_strategy='delete', enable_thinking=True) + dataset.encode(add_generation_prompt=True) + return dataset + + +def main(): + device_groups = [ + DeviceGroup(name='model', ranks=list(range(MODEL_GPUS)), device_type='GPU'), + DeviceGroup(name='sampler', ranks=list(range(MODEL_GPUS, NUM_GPUS)), device_type='GPU', + gpus_per_worker=SAMPLER_TP), + ] + model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, tp_size=2, ep_size=2, pp_size=2, sequence_parallel=True) + sampler_dp_size = SAMPLER_GPUS // SAMPLER_TP + sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, dp_size=sampler_dp_size, tp_size=SAMPLER_TP) + + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, groups=device_groups, lazy_collect=False) + + lora_config = LoraConfig(target_modules='all-linear', r=LORA_RANK, lora_alpha=LORA_RANK * 2, lora_dropout=0.05) + + model = MultiLoraMegatronModel(model_id=MODEL_ID, device_mesh=model_mesh, remote_group='model', + mixed_precision='bf16') + model.add_adapter_to_model(ADAPTER_NAME, lora_config, gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) + model.set_optimizer('default', lr=LEARNING_RATE, adapter_name=ADAPTER_NAME) + model.set_lr_scheduler('default', lr_decay_steps=MAX_STEPS, max_lr=LEARNING_RATE, adapter_name=ADAPTER_NAME) + model.set_loss('GRPOLoss', epsilon=0.2, adapter_name=ADAPTER_NAME) + model.set_processor(InputProcessor, adapter_name=ADAPTER_NAME) + model.set_template(TEMPLATE, model_id=MODEL_ID, enable_thinking=True, adapter_name=ADAPTER_NAME) + + sampler = vLLMSampler( + model_id=MODEL_ID, + engine_args={ + 'tensor_parallel_size': SAMPLER_TP, + 'gpu_memory_utilization': 0.8, + 'max_model_len': 8192, + 'max_lora_rank': LORA_RANK, + 'enable_lora': True, + 'enable_tower_connector_lora': True, + }, + device_mesh=sampler_mesh, + remote_group='sampler', + ) + sampler.set_template(TEMPLATE, model_id=MODEL_ID, enable_thinking=True) + + GLOBAL_BATCH_SIZE = BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS + dataloader = DataLoader(dataset=create_rsi_dataset, batch_size=GLOBAL_BATCH_SIZE, + min_batch_size=GLOBAL_BATCH_SIZE, device_mesh=model_mesh, remote_group='model') + + advantage_fn = GRPOAdvantage() + metrics = CompletionRewardMetric() + reward_fn = ToolMatchReward() + 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 RSI per-round GRPO (MultiLoraMegatron, filesystem LoRA sync)') + logger.info(get_device_placement()) + + for batch in dataloader: + if optim_step >= MAX_STEPS: + break + metrics.reset() + expand_prompts = [] + for prompt in batch: + expand_prompts.extend([prompt] * NUM_GENERATIONS) + + lora_sync_path = model.save(f'lora-sync-step-{optim_step}', output_dir=LORA_SYNC_DIR, adapter_name=ADAPTER_NAME) + sampler.reset_prefix_cache() + sample_responses = sampler.sample(expand_prompts, sampling_params, adapter_path=lora_sync_path) + + 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 = reward_fn(all_input_data) + metrics.accumulate(completion_lengths=all_completion_lengths, rewards={'tool_match': rewards}) + advantages = advantage_fn(rewards, num_generations=NUM_GENERATIONS, scale='group').tolist() + + total = len(all_input_data) + for mb_start in range(0, total, MINI_BATCH_SIZE): + mb_end = min(mb_start + MINI_BATCH_SIZE, total) + model.forward_backward( + inputs=all_input_data[mb_start:mb_end], + old_logps=all_old_logps[mb_start:mb_end], + advantages=advantages[mb_start:mb_end], + micro_batch_size=MICRO_BATCH_SIZE, + adapter_name=ADAPTER_NAME, + ) + model.clip_grad_and_step(adapter_name=ADAPTER_NAME) + optim_step += 1 + if optim_step >= MAX_STEPS: + break + if optim_step % SAVE_STEPS == 0: + model.save(f'rsi-executor-checkpoint-{optim_step}', adapter_name=ADAPTER_NAME) + + log_dict = metrics.calculate() + log_dict.update(model.calculate_metric(is_training=True, adapter_name=ADAPTER_NAME)) + swanlab.log(log_dict) + metrics.reset() + logger.info(f'[Step {optim_step}/{MAX_STEPS}] {log_dict}') + + logger.info(f'Training completed. optim_steps={optim_step}') + model.save('rsi-executor-final', adapter_name=ADAPTER_NAME) + + +if __name__ == '__main__': + main() diff --git a/src/twinkle_agentic/utils/llm_backup.py b/src/twinkle_agentic/utils/llm_backup.py index 3ce0c4dd8..a7d804101 100644 --- a/src/twinkle_agentic/utils/llm_backup.py +++ b/src/twinkle_agentic/utils/llm_backup.py @@ -16,6 +16,7 @@ class EvalRecord: student_result: Any teacher_result: Any match: bool + trajectory: Any = None @dataclass @@ -50,14 +51,19 @@ def increment_call(self, key: str) -> int: state.call_count += 1 return state.call_count - def add_record(self, key: str, student_result: Any, teacher_result: Any, match: bool): + def add_record(self, key: str, student_result: Any, teacher_result: Any, match: bool, + trajectory: Any = None): with self._lock: state = self._states[key] state.dataset.append(EvalRecord( student_result=student_result, teacher_result=teacher_result, match=match, + trajectory=trajectory, )) + # File IO deliberately outside the registry lock so a slow disk never + # stalls confidence bookkeeping for other keys. + _maybe_dump(key, trajectory, student_result, teacher_result, match) def refresh_confidence(self, key: str) -> float: with self._lock: @@ -82,6 +88,37 @@ def _compute_confidence(dataset: List[EvalRecord]) -> float: _registry = DistillationRegistry() _teacher_api = None _teacher_lock = threading.Lock() +_dump_lock = threading.Lock() + + +def _maybe_dump(key: str, trajectory: Any, student_result: Any, + teacher_result: Any, match: bool) -> None: + """Append one raw (input -> teacher output) record as JSONL when the env var + ``LLM_BACKUP_DUMP_PATH`` is set. Off by default: no path -> nothing written, + behaviour is identical to before. + + The ``trajectory`` (the exact model input) is stored verbatim so the dump is + directly reshapeable into SFT pairs downstream; ``student``/``teacher``/ + ``match`` are kept too so nothing is thrown away (target selection is decided + by the consumer, not here). + """ + path = os.environ.get('LLM_BACKUP_DUMP_PATH') + if not path: + return + rec = { + 'key': key, + 'trajectory': trajectory, + 'student': student_result, + 'teacher': teacher_result, + 'match': match, + } + try: + line = json.dumps(rec, ensure_ascii=False, default=str) + except (TypeError, ValueError): + return + with _dump_lock: + with open(path, 'a', encoding='utf-8') as f: + f.write(line + '\n') def _get_teacher_api(): @@ -239,7 +276,7 @@ def wrapper(*args, **kwargs): if random.random() < sample_rate: teacher_result = _call_teacher(trajectory, sampling_params) match = comparator(result, teacher_result) - _registry.add_record(key, result, teacher_result, match) + _registry.add_record(key, result, teacher_result, match, trajectory=trajectory) if not match: result = teacher_result else: @@ -247,7 +284,7 @@ def wrapper(*args, **kwargs): teacher_result = _call_teacher(trajectory, sampling_params) student_result = fn(*args, **kwargs) match = comparator(student_result, teacher_result) - _registry.add_record(key, student_result, teacher_result, match) + _registry.add_record(key, student_result, teacher_result, match, trajectory=trajectory) result = teacher_result call_count = _registry.increment_call(key) @@ -298,14 +335,14 @@ async def wrapper(*args, **kwargs): if random.random() < sample_rate: teacher_result = _call_teacher(trajectory, sampling_params) match = comparator(result, teacher_result) - _registry.add_record(key, result, teacher_result, match) + _registry.add_record(key, result, teacher_result, match, trajectory=trajectory) if not match: result = teacher_result else: teacher_result = _call_teacher(trajectory, sampling_params) student_result = await fn(*args, **kwargs) match = comparator(student_result, teacher_result) - _registry.add_record(key, student_result, teacher_result, match) + _registry.add_record(key, student_result, teacher_result, match, trajectory=trajectory) result = teacher_result call_count = _registry.increment_call(key) From 5175833159a4355ae9801b16f8d8beb2026d18a6 Mon Sep 17 00:00:00 2001 From: root <yuze.zyz@alibaba-inc.com> Date: Fri, 21 Aug 2026 17:22:50 +0800 Subject: [PATCH 41/60] wip --- cookbook/rl/grpo/kodcode_grpo.py | 524 +++++++ cookbook/rl/grpo/mbpp_grpo.py | 407 ++++++ cookbook/rl/grpo/short_math_grpo.py | 6 +- cookbook/rl/rsi_agentic/rsi_agent.yaml | 64 + cookbook/rl/rsi_agentic/rsi_agentic_grpo.py | 325 +++++ cookbook/rl/rsi_agentic/tasks.example.jsonl | 4 + cookbook/rsi/run_rsi_selfplay.py | 125 ++ src/twinkle/template/qwen3_5_vl.py | 2 +- src/twinkle/template/tools/__init__.py | 5 + src/twinkle/template/tools/bracket_dsl.py | 180 +++ src/twinkle/utils/__init__.py | 2 +- src/twinkle/utils/utils.py | 21 + src/twinkle_agentic/classifier/__init__.py | 0 src/twinkle_agentic/classifier/base.py | 11 - src/twinkle_agentic/data_format/__init__.py | 1 - src/twinkle_agentic/data_format/chunks.py | 104 -- src/twinkle_agentic/envs/__init__.py | 1 + src/twinkle_agentic/envs/base.py | 18 +- src/twinkle_agentic/envs/env_tool.py | 41 +- src/twinkle_agentic/envs/ms_agent_tool_env.py | 212 +++ src/twinkle_agentic/harness/__init__.py | 14 + src/twinkle_agentic/harness/base.py | 103 ++ src/twinkle_agentic/harness/ms_agent.py | 422 ++++++ src/twinkle_agentic/memory/DESIGN.md | 745 ---------- src/twinkle_agentic/preprocessor/__init__.py | 4 - .../preprocessor/message_normalizer.py | 15 +- .../preprocessor/outcome_filter.py | 80 -- .../preprocessor/safety_scorer.py | 100 -- .../preprocessor/trajectory_scorer.py | 365 ----- .../preprocessor/value_selector.py | 288 ---- src/twinkle_agentic/rollout/__init__.py | 10 +- src/twinkle_agentic/rollout/multi_turn.py | 294 +++- .../rollout/multi_turn_condense.py | 284 ---- src/twinkle_agentic/rsi/rsi_challenge.py | 1044 ++++++++++++++ src/twinkle_agentic/rsi/rsi_prepare.py | 6 +- src/twinkle_agentic/rsi/rsi_rl.py | 1208 +++++++++++++++-- src/twinkle_agentic/segment/__init__.py | 4 - src/twinkle_agentic/segment/base.py | 237 ---- src/twinkle_agentic/segment/llm_segmenter.py | 308 ----- src/twinkle_agentic/summarizer/__init__.py | 4 +- .../summarizer/action_summarizer.py | 86 -- .../summarizer/error_summarizer.py | 0 .../summarizer/fact_summarizer.py | 83 -- .../summarizer/pattern_summarizer.py | 0 .../tools/extract_condensed.py | 150 -- src/twinkle_agentic/tools/tool_manager.py | 113 +- src/twinkle_agentic/train/__init__.py | 0 src/twinkle_agentic/train/base.py | 6 - src/twinkle_agentic/train/cron.py | 5 - src/twinkle_agentic/verifier/__init__.py | 33 +- src/twinkle_agentic/verifier/aggregation.py | 274 ---- src/twinkle_agentic/verifier/base.py | 11 - src/twinkle_agentic/verifier/domain_checks.py | 436 ------ src/twinkle_agentic/verifier/hard_scorer.py | 444 ------ src/twinkle_agentic/verifier/leak_verifier.py | 363 ----- src/twinkle_agentic/verifier/result_check.py | 365 +++++ .../verifier/rubric_library.py | 108 -- .../verifier/rubric_verifier.py | 1077 --------------- tests/preprocessor/test_refuse_filter.py | 58 +- tests/preprocessor/test_value_selector.py | 234 ---- tests/twinkle_agentic/test_agentic_rsi.py | 260 ++++ .../test_aggregation_fusion.py | 12 - .../twinkle_agentic/test_diagnosis_salvage.py | 62 - .../twinkle_agentic/test_extract_condensed.py | 422 ------ tests/twinkle_agentic/test_harness.py | 175 +++ .../twinkle_agentic/test_keyword_condenser.py | 486 ------- tests/twinkle_agentic/test_model_condenser.py | 515 ------- .../test_multi_turn_condense_trace.py | 133 -- tests/twinkle_agentic/test_native_chunker.py | 555 -------- .../test_repeated_calls_spin.py | 47 - .../test_rubric_stabilization.py | 92 -- 71 files changed, 5836 insertions(+), 8357 deletions(-) create mode 100644 cookbook/rl/grpo/kodcode_grpo.py create mode 100644 cookbook/rl/grpo/mbpp_grpo.py create mode 100644 cookbook/rl/rsi_agentic/rsi_agent.yaml create mode 100644 cookbook/rl/rsi_agentic/rsi_agentic_grpo.py create mode 100644 cookbook/rl/rsi_agentic/tasks.example.jsonl create mode 100644 cookbook/rsi/run_rsi_selfplay.py create mode 100644 src/twinkle/template/tools/bracket_dsl.py delete mode 100644 src/twinkle_agentic/classifier/__init__.py delete mode 100644 src/twinkle_agentic/classifier/base.py delete mode 100644 src/twinkle_agentic/data_format/__init__.py delete mode 100644 src/twinkle_agentic/data_format/chunks.py create mode 100644 src/twinkle_agentic/envs/ms_agent_tool_env.py create mode 100644 src/twinkle_agentic/harness/__init__.py create mode 100644 src/twinkle_agentic/harness/base.py create mode 100644 src/twinkle_agentic/harness/ms_agent.py delete mode 100644 src/twinkle_agentic/memory/DESIGN.md delete mode 100644 src/twinkle_agentic/preprocessor/outcome_filter.py delete mode 100644 src/twinkle_agentic/preprocessor/safety_scorer.py delete mode 100644 src/twinkle_agentic/preprocessor/trajectory_scorer.py delete mode 100644 src/twinkle_agentic/preprocessor/value_selector.py delete mode 100644 src/twinkle_agentic/rollout/multi_turn_condense.py create mode 100644 src/twinkle_agentic/rsi/rsi_challenge.py delete mode 100644 src/twinkle_agentic/segment/__init__.py delete mode 100644 src/twinkle_agentic/segment/base.py delete mode 100644 src/twinkle_agentic/segment/llm_segmenter.py delete mode 100644 src/twinkle_agentic/summarizer/action_summarizer.py delete mode 100644 src/twinkle_agentic/summarizer/error_summarizer.py delete mode 100644 src/twinkle_agentic/summarizer/fact_summarizer.py delete mode 100644 src/twinkle_agentic/summarizer/pattern_summarizer.py delete mode 100644 src/twinkle_agentic/tools/extract_condensed.py delete mode 100644 src/twinkle_agentic/train/__init__.py delete mode 100644 src/twinkle_agentic/train/base.py delete mode 100644 src/twinkle_agentic/train/cron.py delete mode 100644 src/twinkle_agentic/verifier/aggregation.py delete mode 100644 src/twinkle_agentic/verifier/base.py delete mode 100644 src/twinkle_agentic/verifier/domain_checks.py delete mode 100644 src/twinkle_agentic/verifier/hard_scorer.py delete mode 100644 src/twinkle_agentic/verifier/leak_verifier.py create mode 100644 src/twinkle_agentic/verifier/result_check.py delete mode 100644 src/twinkle_agentic/verifier/rubric_library.py delete mode 100644 src/twinkle_agentic/verifier/rubric_verifier.py delete mode 100644 tests/preprocessor/test_value_selector.py create mode 100644 tests/twinkle_agentic/test_agentic_rsi.py delete mode 100644 tests/twinkle_agentic/test_aggregation_fusion.py delete mode 100644 tests/twinkle_agentic/test_diagnosis_salvage.py delete mode 100644 tests/twinkle_agentic/test_extract_condensed.py create mode 100644 tests/twinkle_agentic/test_harness.py delete mode 100644 tests/twinkle_agentic/test_keyword_condenser.py delete mode 100644 tests/twinkle_agentic/test_model_condenser.py delete mode 100644 tests/twinkle_agentic/test_multi_turn_condense_trace.py delete mode 100644 tests/twinkle_agentic/test_native_chunker.py delete mode 100644 tests/twinkle_agentic/test_repeated_calls_spin.py delete mode 100644 tests/twinkle_agentic/test_rubric_stabilization.py diff --git a/cookbook/rl/grpo/kodcode_grpo.py b/cookbook/rl/grpo/kodcode_grpo.py new file mode 100644 index 000000000..09195f41f --- /dev/null +++ b/cookbook/rl/grpo/kodcode_grpo.py @@ -0,0 +1,524 @@ +"""GRPO training script for KodCode-V1 (code generation with pytest-verified reward). + +Same structure as short_math_grpo.py, but the reward runs the dataset's own +pytest suite against the generated code instead of comparing a final number. + +Difficulty is filtered by KodCode's own ``gpt_pass_percentage`` so that the +sampled group is unlikely to collapse (all-correct or all-wrong within a group +gives a zero GRPO advantage and therefore no gradient). + +Sandbox judging follows .temp/human_e18/e18_kodcode.py (``run_tests``): the +submitted code is written to ``solution.py`` and the official test to +``test_solution.py``, then pytest runs in a subprocess with a timeout and a 2GB +address-space limit. That logic is inlined here rather than imported, because +Ray deserializes the dataset builder and the reward inside worker processes that +do not share this driver's ``sys.path``. +""" +import ast as _ast +import os +import re +import resource +import shutil +import signal +import subprocess +import sys +import tempfile +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Optional, Tuple + +from peft import LoraConfig + +import swanlab +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.cli import CLI +from twinkle.data_format import Message, SamplingParams, Trajectory +from twinkle.dataloader import DataLoader +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.metric import CompletionRewardMetric +from twinkle.model import TransformersModel +from twinkle.preprocessor import Preprocessor +from twinkle.processor import InputProcessor +from twinkle.reward.base import Reward +from twinkle.sampler import vLLMSampler + +logger = get_logger() +args = CLI.from_args() + +swanlab.init(project='twinkle') + +# ========== Configuration ========== +MODEL_ID = args.model.model_id or 'ms://Qwen/Qwen3-4B' +USE_MEGATRON = args.model.strategy != 'native_fsdp' + +MODEL_GPUS = args.infra.model_gpus or 4 +SAMPLER_GPUS = args.infra.sampler_gpus or 4 +NUM_GPUS = MODEL_GPUS + SAMPLER_GPUS + +NUM_GENERATIONS = args.rl.num_generations or 8 +MAX_NEW_TOKENS = args.sampling.max_tokens or 4096 +LEARNING_RATE = args.optimizer.learning_rate or 1e-5 +MAX_STEPS = args.training.max_steps or 1000 +BATCH_SIZE = args.training.batch_size or 8 +MINI_BATCH_SIZE = args.training.mini_batch_size or 8 +MICRO_BATCH_SIZE = args.training.micro_batch_size or 2 +GRADIENT_ACCUMULATION_STEPS = args.training.gradient_accumulation_steps or 1 +SAVE_STEPS = args.training.save_steps or 1000 +LORA_RANK = args.lora.lora_r or 16 + +# Keep only problems the teacher solved sometimes but not always: a group whose +# 8 samples are all right or all wrong contributes no advantage. +KOD_MIN_PASS_PCT = float(os.environ.get('KOD_MIN_PASS_PCT', 0.2)) +KOD_MAX_PASS_PCT = float(os.environ.get('KOD_MAX_PASS_PCT', 0.8)) +# Judging is a subprocess and runs while the GPUs idle, so keep it wide. +JUDGE_WORKERS = int(os.environ.get('JUDGE_WORKERS', max(24, min(96, (os.cpu_count() or 24) // 2)))) + +SYSTEM_PROMPT = ('You are an expert Python programmer. Write a complete, self-contained ' + 'solution in a single ```python code block. Do not include tests.') + +TEST_TIMEOUT = int(os.environ.get('TEST_TIMEOUT', 60)) + +_FENCE_RE = re.compile(r'```(?:python|py)?\s*\n(.*?)```', re.S) +_SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') + + +# ========== Text handling (same as e18_kodcode) ========== +def after_think(text: str) -> str: + """Keep only what follows </think>; return the text unchanged if unclosed.""" + idx = text.rfind('</think>') + return text[idx + len('</think>'):] if idx >= 0 else text + + +def clean_text(decoded: Optional[str]) -> str: + return _SPECIAL_TOKEN_RE.sub('', decoded or '').strip() + + +def extract_code(text: str) -> str: + """Take the last fenced block; fall back to the whole body when unfenced. + + The last one, not the first: models often draft a version before the final + one, and the last block is their conclusion. + """ + body = after_think(text) + blocks = _FENCE_RE.findall(body) + if blocks: + return blocks[-1].strip() + return body.strip() + + +# ========== Sandbox (same contract as e18_kodcode.run_tests) ========== +# Assertion vs exception must be told apart via ``reprcrash.message``: pytest +# rewrites assertions, so the summary reads "E assert -1 == 3" and the string +# "AssertionError" never appears -- matching on it misclassifies every failed +# assertion as an exception. +_RUNNER = r""" +import sys, pytest + + +class _Collect: + def __init__(self): + self.n_tests = self.n_fail = self.n_err = 0 + + @staticmethod + def _is_assertion(report): + crash = getattr(getattr(report, 'longrepr', None), 'reprcrash', None) + msg = getattr(crash, 'message', '') or '' + return msg.startswith('assert') or msg.startswith('AssertionError') + + def pytest_runtest_logreport(self, report): + if report.when == 'call': + self.n_tests += 1 + if report.failed: + if self._is_assertion(report): + self.n_fail += 1 + else: + self.n_err += 1 + elif report.failed: + self.n_err += 1 + + +c = _Collect() +rc = pytest.main(['-q', '--no-header', '-p', 'no:cacheprovider', + '--tb=short', 'test_solution.py'], plugins=[c]) +print('__KOD__', c.n_tests, c.n_fail, c.n_err) +sys.exit(0 if int(rc) == 0 else 1) +""" + + +def run_tests(code: str, payload: Dict[str, Any], timeout: int = TEST_TIMEOUT) -> Dict[str, Any]: + """Run the submitted code (solution.py) against the official test in a subprocess. + + The code goes into its own ``solution.py`` because KodCode tests grab the + function under test via ``from solution import X``. + """ + if not code.strip(): + return {'passed': False, 'kind': 'no_code', 'error': 'no parseable code block'} + entry = payload.get('entry_point') or '' + if entry and entry not in code: + return {'passed': False, 'kind': 'no_entry', + 'error': f'function {entry} is not defined in the submitted code'} + tmp = tempfile.mkdtemp(prefix='kod_') + try: + with open(os.path.join(tmp, 'solution.py'), 'w', encoding='utf-8') as f: + f.write(code) + with open(os.path.join(tmp, 'test_solution.py'), 'w', encoding='utf-8') as f: + f.write(payload['test']) + with open(os.path.join(tmp, '_run.py'), 'w', encoding='utf-8') as f: + f.write(_RUNNER) + env = dict(os.environ, MPLBACKEND='Agg', PYTHONHASHSEED='0', OMP_NUM_THREADS='1', + MKL_NUM_THREADS='1', TOKENIZERS_PARALLELISM='false') + env.pop('CUDA_VISIBLE_DEVICES', None) + + # start_new_session + killpg on timeout: pytest can fork, and a bare + # kill would leave grandchildren running. RLIMIT_AS caps the child at + # 2GB so a runaway solution cannot take the host down. + def _limit(): + resource.setrlimit(resource.RLIMIT_AS, (2 * 1024**3, 2 * 1024**3)) + + proc = subprocess.Popen([sys.executable, '_run.py'], cwd=tmp, env=env, + stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, + errors='replace', start_new_session=True, preexec_fn=_limit) + try: + stdout, stderr = proc.communicate(timeout=timeout) + returncode = proc.returncode + except subprocess.TimeoutExpired: + try: + os.killpg(proc.pid, signal.SIGKILL) + except ProcessLookupError: + pass + try: + proc.communicate(timeout=5) + except Exception: + pass + return {'passed': False, 'kind': 'timeout', + 'error': f'the tests did not finish within {timeout}s'} + n_tests = n_fail = n_err = 0 + for line in (stdout or '').splitlines(): + if line.startswith('__KOD__'): + _, a, b, c = line.split() + n_tests, n_fail, n_err = int(a), int(b), int(c) + if returncode == 0 and n_tests > 0: + return {'passed': True, 'kind': 'pass', 'error': ''} + kind = 'assertion' if n_fail else ('exception' if n_err else 'import_or_syntax') + return {'passed': False, 'kind': kind, 'error': ''} + finally: + shutil.rmtree(tmp, ignore_errors=True) + + +# ========== Row helpers (same as e18_kodcode) ========== +def _entry_point(row: Dict[str, Any]) -> str: + """Function under test, from test_info; else from ``from solution import X``.""" + ti = row.get('test_info') + if ti is not None: + try: + items = list(ti) if not isinstance(ti, str) else _ast.literal_eval(ti) + for it in items: + name = (it or {}).get('function_name') + if name: + return str(name) + except Exception: + pass + m = re.search(r'from\s+solution\s+import\s+([A-Za-z_]\w*)', row.get('test') or '') + return m.group(1) if m else '' + + +def _code_prompt(row: Dict[str, Any]) -> str: + """The function signature, used to pin the entry point for the model.""" + ti = row.get('test_info') + if ti is not None: + try: + items = list(ti) if not isinstance(ti, str) else _ast.literal_eval(ti) + for it in items: + decl = (it or {}).get('function_declaration') + if decl: + return str(decl) + except Exception: + pass + return '' + + +def _usable(row: Dict[str, Any]) -> bool: + """Minimum bar to enter the pool. + + The test must import from ``solution``: 11.7% of rows call bare function + names, which can never resolve under this sandbox layout, so keeping them + would permanently depress the reward for reasons unrelated to the model. + """ + test = row.get('test') or '' + if 'def test_' not in test: + return False + if not re.search(r'from\s+solution\s+import|import\s+solution\b', test): + return False + return bool((row.get('solution') or '').strip()) and bool(_entry_point(row)) + + +# ========== Reward ========== +class KodCodePytestReward(Reward): + """1.0 when the generated code passes the problem's own pytest suite. + + The suite is carried per-sample through ``user_data`` (``kod_payload``), so + each trajectory is judged against its own tests. Judging runs in a thread + pool because every verdict is a separate subprocess. + """ + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + jobs: List[Tuple[int, str, Dict[str, Any]]] = [] + rewards = [0.0] * len(trajectories) + for i, traj in enumerate(trajectories): + payload = None + for item in traj.get('user_data') or []: + if item[0] == 'kod_payload': + payload = item[1] + break + if payload is None: + continue + completion = '' + for msg in reversed(traj.get('messages', []) or []): + if msg.get('role') == 'assistant': + completion = msg.get('content', '') or '' + break + jobs.append((i, extract_code(completion), payload)) + + if not jobs: + return rewards + # Same (code, test) pair judged once: identical completions are common. + uniq: Dict[Tuple[str, str], Dict[str, Any]] = {} + for _, code, payload in jobs: + uniq.setdefault((payload['task_id'], code), payload) + todo = list(uniq) + with ThreadPoolExecutor(max_workers=max(1, min(JUDGE_WORKERS, len(todo)))) as ex: + verdicts = dict(zip(todo, ex.map(lambda k: run_tests(k[1], uniq[k]), todo))) + for i, code, payload in jobs: + v = verdicts.get((payload['task_id'], code)) + rewards[i] = 1.0 if (v and v['passed']) else 0.0 + return rewards + + +class KodCodeFormatReward(Reward): + """1.0 when the completion contains a parseable python code block.""" + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + rewards = [] + for traj in trajectories: + completion = '' + for msg in reversed(traj.get('messages', []) or []): + if msg.get('role') == 'assistant': + completion = msg.get('content', '') or '' + break + rewards.append(1.0 if extract_code(completion).strip() else 0.0) + return rewards + + +# ========== Dataset ========== +# Only 8% of KodCode questions name the function under test, but the tests grab +# it via ``from solution import <name>``. Append the signature or nearly every +# sample scores 0 regardless of how good the answer is. +_SIG_HINT = '\n\nYou should write self-contained code starting with:\n```\n{decl}\n```' + + +class KodCodeProcessor(Preprocessor): + """KodCode row -> prompt-only Trajectory carrying its pytest suite.""" + + def __init__(self, system=SYSTEM_PROMPT): + self.system = system + + def __call__(self, rows: Dict[str, List[Any]]) -> Dict[str, List[Any]]: + rows = self.map_col_to_row(rows) + rows = [self.preprocess(row) for row in rows] + return self.map_row_to_col(rows) + + def preprocess(self, row) -> Trajectory: + question = row.get('question') or '' + decl = _code_prompt(row) + if decl and decl.strip() not in question: + question = question + _SIG_HINT.format(decl=decl.strip()) + payload = { + 'task_id': str(row.get('question_id') or ''), + 'entry_point': _entry_point(row), + 'test': row.get('test') or '', + } + return Trajectory( + messages=[ + Message(role='system', content=self.system), + Message(role='user', content=question), + ], + user_data=[('kod_payload', payload)], + ) + + +def create_kodcode_dataset(): + dataset = Dataset() + dataset.add_dataset(DatasetMeta('ms://AI-ModelScope/KodCode-V1', split='train')) + # Filter before templating: the full set is 73747 rows. + dataset.filter(lambda r: KOD_MIN_PASS_PCT <= float(r.get('gpt_pass_percentage') or 0.0) + <= KOD_MAX_PASS_PCT) + # Tests must import from ``solution``; the 11.7% that call bare names can + # never pass in this sandbox layout and would only drag the reward down. + dataset.filter(_usable) + dataset.set_template('Template', model_id=MODEL_ID, max_length=4096, + truncation_strategy='delete', enable_thinking=True) + dataset.map(KodCodeProcessor()) + dataset.encode(add_generation_prompt=True) + return dataset + + +def compute_rewards( + trajectories: List[Dict[str, Any]], +) -> Tuple[List[float], List[float], List[float]]: + pass_rewards = KodCodePytestReward()(trajectories) + format_rewards = KodCodeFormatReward()(trajectories) + total_rewards = [p + f for p, f in zip(pass_rewards, format_rewards)] + return total_rewards, format_rewards, pass_rewards + + +# ========== Main ========== +def main(): + device_groups = [ + DeviceGroup(name='model', ranks=list(range(MODEL_GPUS)), device_type='GPU'), + DeviceGroup(name='sampler', ranks=list(range(MODEL_GPUS, 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, dp_size=SAMPLER_GPUS) + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, groups=device_groups, lazy_collect=False) + + lora_config = LoraConfig( + target_modules='all-linear', + r=LORA_RANK, + lora_alpha=LORA_RANK * 2, + lora_dropout=0.05, + ) + + if USE_MEGATRON: + from twinkle.model.megatron import MegatronModel + model = MegatronModel( + model_id=MODEL_ID, + device_mesh=model_mesh, + remote_group='model', + mixed_precision='bf16', + variable_seq_lengths=True, + ) + else: + model = TransformersModel( + model_id=MODEL_ID, + device_mesh=model_mesh, + remote_group='model', + ) + + model.add_adapter_to_model('default', lora_config, + gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) + if USE_MEGATRON: + model.set_optimizer('default', lr=LEARNING_RATE) + model.set_lr_scheduler('default', lr_decay_steps=MAX_STEPS, max_lr=LEARNING_RATE) + else: + 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('Template', model_id=MODEL_ID, enable_thinking=True) + + sampler = vLLMSampler( + model_id=MODEL_ID, + engine_args={ + 'gpu_memory_utilization': 0.8, + 'max_model_len': 8192, + 'max_lora_rank': 32, + 'enable_lora': True, + 'enable_tower_connector_lora': True, + }, + device_mesh=sampler_mesh, + remote_group='sampler', + ) + sampler.set_template('Template', model_id=MODEL_ID, enable_thinking=True) + + ckpt_manager = CheckpointEngineManager(model=model, sampler=sampler) + + GLOBAL_BATCH_SIZE = BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS + dataloader = DataLoader( + dataset=create_kodcode_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(f'Starting KodCode GRPO (pass_pct window ' + f'[{KOD_MIN_PASS_PCT}, {KOD_MAX_PASS_PCT}], judge workers {JUDGE_WORKERS})') + logger.info(get_device_placement()) + + for batch in dataloader: + if optim_step >= MAX_STEPS: + break + + metrics.reset() + expand_prompts = [] + for prompt in batch: + 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) + + 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)) + + total_rewards, format_rewards, pass_rewards = compute_rewards(all_input_data) + + metrics.accumulate( + completion_lengths=all_completion_lengths, + rewards={ + 'total': total_rewards, + 'format': format_rewards, + 'pass': pass_rewards, + }, + ) + + advantages = advantage_fn(total_rewards, num_generations=NUM_GENERATIONS, + scale='group').tolist() + + 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) + model.forward_backward( + inputs=all_input_data[mb_start:mb_end], + old_logps=all_old_logps[mb_start:mb_end], + advantages=advantages[mb_start:mb_end], + 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'kodcode-grpo-checkpoint-{optim_step}') + + log_dict = metrics.calculate() + log_dict.update(model.calculate_metric(is_training=True)) + swanlab.log(log_dict) + metrics.reset() + logger.info(f'[Step {optim_step}/{MAX_STEPS}] {log_dict}') + + logger.info(f'Training completed. optim_steps={optim_step}') + model.save('kodcode-grpo-final') + + +if __name__ == '__main__': + main() diff --git a/cookbook/rl/grpo/mbpp_grpo.py b/cookbook/rl/grpo/mbpp_grpo.py new file mode 100644 index 000000000..1c75d927e --- /dev/null +++ b/cookbook/rl/grpo/mbpp_grpo.py @@ -0,0 +1,407 @@ +"""GRPO training script for MBPP (code generation with assert-verified reward). + +Same structure as kodcode_grpo.py, but MBPP's tests are bare asserts that call +the function by name (``assert min_cost(...) == 8``), so the generated code, +``test_setup_code`` and the asserts are concatenated into a single file and +executed -- no ``from solution import`` layout is needed. That judging path was +checked against all 974 reference solutions and passes 974/974. + +The problem statement does not name the function, and the asserts do, so the +asserts are shown in the prompt (the standard MBPP setup used by OpenCompass / +EvalPlus). Without them the function name is unguessable and every sample fails +for reasons unrelated to coding ability. + +Measured difficulty of the full 974-problem set under Qwen3-4B (8 samples each, +see output/mbpp/measure_mbpp_difficulty.py): 21.97% all-wrong, 56.67% all-right, +21.36% mixed. Only the mixed ones carry a GRPO gradient; the full set is used +here as requested. +""" +import json +import os +import re +import resource +import shutil +import signal +import subprocess +import sys +import tempfile +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Optional, Tuple + +from peft import LoraConfig + +import swanlab +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.cli import CLI +from twinkle.data_format import Message, SamplingParams, Trajectory +from twinkle.dataloader import DataLoader +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.metric import CompletionRewardMetric +from twinkle.model import TransformersModel +from twinkle.preprocessor import Preprocessor +from twinkle.processor import InputProcessor +from twinkle.reward.base import Reward +from twinkle.sampler import vLLMSampler + +logger = get_logger() +args = CLI.from_args() + +swanlab.init(project='twinkle') + +# ========== Configuration ========== +MODEL_ID = args.model.model_id or 'ms://Qwen/Qwen3-4B' +USE_MEGATRON = args.model.strategy != 'native_fsdp' + +MODEL_GPUS = args.infra.model_gpus or 4 +SAMPLER_GPUS = args.infra.sampler_gpus or 4 +NUM_GPUS = MODEL_GPUS + SAMPLER_GPUS + +NUM_GENERATIONS = args.rl.num_generations or 8 +MAX_NEW_TOKENS = args.sampling.max_tokens or 2048 +LEARNING_RATE = args.optimizer.learning_rate or 1e-5 +MAX_STEPS = args.training.max_steps or 1000 +BATCH_SIZE = args.training.batch_size or 8 +MINI_BATCH_SIZE = args.training.mini_batch_size or 8 +MICRO_BATCH_SIZE = args.training.micro_batch_size or 2 +GRADIENT_ACCUMULATION_STEPS = args.training.gradient_accumulation_steps or 1 +SAVE_STEPS = args.training.save_steps or 200 +LORA_RANK = args.lora.lora_r or 16 + +JUDGE_WORKERS = int(os.environ.get('JUDGE_WORKERS', max(24, min(96, (os.cpu_count() or 24) // 2)))) +TEST_TIMEOUT = int(os.environ.get('TEST_TIMEOUT', 30)) + +SYSTEM_PROMPT = ('You are an expert Python programmer. Write a complete, self-contained ' + 'solution in a single ```python code block. Do not include tests.') + +_FENCE_RE = re.compile(r'```(?:python|py)?\s*\n(.*?)```', re.S) + + +# ========== Text handling ========== +def extract_code(text: str) -> str: + """Take the last fenced block; fall back to the whole body when unfenced.""" + idx = (text or '').rfind('</think>') + body = text[idx + len('</think>'):] if idx >= 0 else (text or '') + blocks = _FENCE_RE.findall(body) + return (blocks[-1] if blocks else body).strip() + + +# ========== Sandbox ========== +def run_asserts(code: str, setup: str, asserts: List[str], timeout: int = TEST_TIMEOUT) -> bool: + """True when every assert passes. + + MBPP asserts call the function by name, so code + setup + asserts run as a + single file. Uses start_new_session + killpg so a forking solution cannot + leave stray processes, and caps the child at 2GB of address space. + """ + if not code.strip(): + return False + parts = [code] + if (setup or '').strip(): + parts.append(setup) + parts.extend(asserts) + script = '\n\n'.join(parts) + '\n' + tmp = tempfile.mkdtemp(prefix='mbpp_') + try: + with open(os.path.join(tmp, '_run.py'), 'w', encoding='utf-8') as f: + f.write(script) + env = dict(os.environ, MPLBACKEND='Agg', PYTHONHASHSEED='0', OMP_NUM_THREADS='1', + MKL_NUM_THREADS='1', TOKENIZERS_PARALLELISM='false') + env.pop('CUDA_VISIBLE_DEVICES', None) + + def _limit(): + resource.setrlimit(resource.RLIMIT_AS, (2 * 1024**3, 2 * 1024**3)) + + proc = subprocess.Popen([sys.executable, '_run.py'], cwd=tmp, env=env, + stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, + start_new_session=True, preexec_fn=_limit) + try: + proc.communicate(timeout=timeout) + return proc.returncode == 0 + except subprocess.TimeoutExpired: + try: + os.killpg(proc.pid, signal.SIGKILL) + except ProcessLookupError: + pass + try: + proc.communicate(timeout=5) + except Exception: + pass + return False + finally: + shutil.rmtree(tmp, ignore_errors=True) + + +# ========== Reward ========== +class MbppAssertReward(Reward): + """1.0 when the generated code satisfies every assert of its problem. + + The asserts travel per-sample through ``user_data`` (``mbpp_payload``), so + each trajectory is judged against its own tests. Judging runs in a thread + pool because every verdict is a separate subprocess. + """ + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + jobs: List[Tuple[int, str, Dict[str, Any]]] = [] + rewards = [0.0] * len(trajectories) + for i, traj in enumerate(trajectories): + payload = None + for item in traj.get('user_data') or []: + if item[0] == 'mbpp_payload': + payload = item[1] + break + if payload is None: + continue + completion = '' + for msg in reversed(traj.get('messages', []) or []): + if msg.get('role') == 'assistant': + completion = msg.get('content', '') or '' + break + jobs.append((i, extract_code(completion), payload)) + + if not jobs: + return rewards + # Same (task, code) judged once: identical completions are common. + uniq: Dict[Tuple[str, str], Dict[str, Any]] = {} + for _, code, payload in jobs: + uniq.setdefault((payload['task_id'], code), payload) + todo = list(uniq) + with ThreadPoolExecutor(max_workers=max(1, min(JUDGE_WORKERS, len(todo)))) as ex: + verdicts = dict(zip(todo, ex.map( + lambda k: run_asserts(k[1], uniq[k]['setup'], uniq[k]['asserts']), todo))) + for i, code, payload in jobs: + rewards[i] = 1.0 if verdicts.get((payload['task_id'], code)) else 0.0 + return rewards + + +class MbppFormatReward(Reward): + """1.0 when the completion contains a parseable python code block.""" + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + rewards = [] + for traj in trajectories: + completion = '' + for msg in reversed(traj.get('messages', []) or []): + if msg.get('role') == 'assistant': + completion = msg.get('content', '') or '' + break + rewards.append(1.0 if extract_code(completion).strip() else 0.0) + return rewards + + +# ========== Dataset ========== +# The problem statement never names the function while the asserts do, so the +# asserts go into the prompt (standard MBPP setup). Without them the name is +# unguessable and every sample fails regardless of coding ability. +_TEST_HINT = '\n\nYour code should satisfy these tests:\n```python\n{tests}\n```' + + +def _asserts(row: Dict[str, Any]) -> List[str]: + tl = row.get('test_list') + if tl is None: + return [] + return list(tl) if not isinstance(tl, str) else json.loads(tl) + + +class MbppProcessor(Preprocessor): + """MBPP row -> prompt-only Trajectory carrying its asserts.""" + + def __init__(self, system=SYSTEM_PROMPT): + self.system = system + + def __call__(self, rows: Dict[str, List[Any]]) -> Dict[str, List[Any]]: + rows = self.map_col_to_row(rows) + rows = [self.preprocess(row) for row in rows] + return self.map_row_to_col(rows) + + def preprocess(self, row) -> Trajectory: + asserts = _asserts(row) + question = (row.get('text') or '') + _TEST_HINT.format(tests='\n'.join(asserts)) + payload = { + 'task_id': str(row.get('task_id') or ''), + 'setup': row.get('test_setup_code') or '', + 'asserts': asserts, + } + return Trajectory( + messages=[ + Message(role='system', content=self.system), + Message(role='user', content=question), + ], + user_data=[('mbpp_payload', payload)], + ) + + +def create_mbpp_dataset(): + # opencompass/mbpp ships bare jsonl with no HF subset config, so loading it + # by dataset id raises KeyError('default'); download the file and read it + # as a local jsonl instead. + from modelscope.hub.file_download import dataset_file_download + path = dataset_file_download(dataset_id='opencompass/mbpp', file_path='mbpp.jsonl') + dataset = Dataset() + dataset.add_dataset(DatasetMeta(path, split='train')) + dataset.set_template('Template', model_id=MODEL_ID, max_length=4096, + truncation_strategy='delete', enable_thinking=True) + dataset.map(MbppProcessor()) + dataset.encode(add_generation_prompt=True) + return dataset + + +def compute_rewards( + trajectories: List[Dict[str, Any]], +) -> Tuple[List[float], List[float], List[float]]: + pass_rewards = MbppAssertReward()(trajectories) + format_rewards = MbppFormatReward()(trajectories) + total_rewards = [p + f for p, f in zip(pass_rewards, format_rewards)] + return total_rewards, format_rewards, pass_rewards + + +# ========== Main ========== +def main(): + device_groups = [ + DeviceGroup(name='model', ranks=list(range(MODEL_GPUS)), device_type='GPU'), + DeviceGroup(name='sampler', ranks=list(range(MODEL_GPUS, 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, dp_size=SAMPLER_GPUS) + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, groups=device_groups, lazy_collect=False) + + lora_config = LoraConfig( + target_modules='all-linear', + r=LORA_RANK, + lora_alpha=LORA_RANK * 2, + lora_dropout=0.05, + ) + + if USE_MEGATRON: + from twinkle.model.megatron import MegatronModel + model = MegatronModel( + model_id=MODEL_ID, + device_mesh=model_mesh, + remote_group='model', + mixed_precision='bf16', + variable_seq_lengths=True, + ) + else: + model = TransformersModel( + model_id=MODEL_ID, + device_mesh=model_mesh, + remote_group='model', + ) + + model.add_adapter_to_model('default', lora_config, + gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) + if USE_MEGATRON: + model.set_optimizer('default', lr=LEARNING_RATE) + model.set_lr_scheduler('default', lr_decay_steps=MAX_STEPS, max_lr=LEARNING_RATE) + else: + 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('Template', model_id=MODEL_ID, enable_thinking=True) + + sampler = vLLMSampler( + model_id=MODEL_ID, + engine_args={ + 'gpu_memory_utilization': 0.8, + 'max_model_len': 8192, + 'max_lora_rank': 32, + 'enable_lora': True, + 'enable_tower_connector_lora': True, + }, + device_mesh=sampler_mesh, + remote_group='sampler', + ) + sampler.set_template('Template', model_id=MODEL_ID, enable_thinking=True) + + ckpt_manager = CheckpointEngineManager(model=model, sampler=sampler) + + GLOBAL_BATCH_SIZE = BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS + dataloader = DataLoader( + dataset=create_mbpp_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(f'Starting MBPP GRPO (974 problems, judge workers {JUDGE_WORKERS})') + logger.info(get_device_placement()) + + for batch in dataloader: + if optim_step >= MAX_STEPS: + break + + metrics.reset() + expand_prompts = [] + for prompt in batch: + 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) + + 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)) + + total_rewards, format_rewards, pass_rewards = compute_rewards(all_input_data) + + metrics.accumulate( + completion_lengths=all_completion_lengths, + rewards={ + 'total': total_rewards, + 'format': format_rewards, + 'pass': pass_rewards, + }, + ) + + advantages = advantage_fn(total_rewards, num_generations=NUM_GENERATIONS, + scale='group').tolist() + + 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) + model.forward_backward( + inputs=all_input_data[mb_start:mb_end], + old_logps=all_old_logps[mb_start:mb_end], + advantages=advantages[mb_start:mb_end], + 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'mbpp-grpo-checkpoint-{optim_step}') + + log_dict = metrics.calculate() + log_dict.update(model.calculate_metric(is_training=True)) + swanlab.log(log_dict) + metrics.reset() + logger.info(f'[Step {optim_step}/{MAX_STEPS}] {log_dict}') + + logger.info(f'Training completed. optim_steps={optim_step}') + model.save('mbpp-grpo-final') + + +if __name__ == '__main__': + main() diff --git a/cookbook/rl/grpo/short_math_grpo.py b/cookbook/rl/grpo/short_math_grpo.py index 91fcd7669..fa1187f03 100644 --- a/cookbook/rl/grpo/short_math_grpo.py +++ b/cookbook/rl/grpo/short_math_grpo.py @@ -90,7 +90,7 @@ def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: def create_gsm8k_dataset(): dataset = Dataset() dataset.add_dataset(DatasetMeta('ms://modelscope/gsm8k', subset_name='main', split='train')) - dataset.set_template('Qwen3_5Template', model_id=MODEL_ID, max_length=4096, truncation_strategy='delete', enable_thinking=False) + dataset.set_template('Template', model_id=MODEL_ID, max_length=4096, truncation_strategy='delete', enable_thinking=False) dataset.map(GSM8KProcessor(system=SYSTEM_PROMPT)) dataset.encode(add_generation_prompt=True) return dataset @@ -153,7 +153,7 @@ def main(): 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=False) + model.set_template('Template', model_id=MODEL_ID, enable_thinking=False) sampler = vLLMSampler( model_id=MODEL_ID, @@ -167,7 +167,7 @@ def main(): device_mesh=sampler_mesh, remote_group='sampler', ) - sampler.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=False) + sampler.set_template('Template', model_id=MODEL_ID, enable_thinking=False) ckpt_manager = CheckpointEngineManager(model=model, sampler=sampler) diff --git a/cookbook/rl/rsi_agentic/rsi_agent.yaml b/cookbook/rl/rsi_agentic/rsi_agent.yaml new file mode 100644 index 000000000..67a609e25 --- /dev/null +++ b/cookbook/rl/rsi_agentic/rsi_agent.yaml @@ -0,0 +1,64 @@ +# ms-agent config for agentic RSI training. +# +# This file is the framework-specific half on purpose: the system prompt, the +# tool line-up and the sandbox settings live here in cookbook, while +# src/twinkle_agentic stays generic. Point RSI_AGENT_CONFIG at a copy of this +# to change the agent without touching the trainer. +# +# The entry script deletes the `llm:` section before preparing the agent. +# Omitting it here is not enough: ms-agent merges this file over its own +# ms_agent/agent/agent.yaml, which does declare one, and FileSystemTool builds +# an LLM client whenever config.llm exists -- which then asserts on a missing +# modelscope_api_key. Generation comes from twinkle's vLLM sampler, so no +# second model should be reachable from here at all. + +prompt: + # Unset -> ms-agent's built-in agent prompt, which is also what serving uses. + # Set a string here only if training should see a different system prompt + # than production, and know that you are breaking train/serve parity. + system: + +personalization: + # Off: SOUL/AGENTS/PROFILE.md from the developer's own workspace would leak + # machine-specific context into every training prompt. + enabled: false + +# One turn == one sampler call. MultiTurnRollout's max_turns is the real limit; +# this only stops ms-agent from imposing a lower one. +max_chat_round: 9999 + +# Never wait on a human: training runs unattended. +interactive: false +permission_mode: auto + +# Overwritten per trajectory by the entry script. Both file_system and +# code_executor root themselves here, so this is what isolates episodes. +output_dir: output/rsi_agentic/workspace + +callbacks: [] + +tools: + file_system: + mcp: false + include: + - write_file + - read_file + - edit_file + - grep + - glob + code_executor: + mcp: false + # python_env runs on the host; switch to the docker implementation for + # untrusted code, at the cost of a container per episode. + implementation: python_env + include: + - shell_executor + - python_executor + - notebook_executor + todo_list: + mcp: false + +# Web search is deliberately absent. ms-agent's `web_search` key only provides +# fetch_page (retrieve a known URL); a real query-a-search-engine tool needs +# EXA_API_KEY / SERPAPI_API_KEY and is wired separately from the plain tool +# list. Add it here once that is decided; until then no task should need it. diff --git a/cookbook/rl/rsi_agentic/rsi_agentic_grpo.py b/cookbook/rl/rsi_agentic/rsi_agentic_grpo.py new file mode 100644 index 000000000..4aa6d690b --- /dev/null +++ b/cookbook/rl/rsi_agentic/rsi_agentic_grpo.py @@ -0,0 +1,325 @@ +"""Agentic RSI: GRPO on multi-turn tool-using episodes, scored by program checks. + +The solver is an ms-agent agent with a real tool line-up (shell, filesystem, +python, notebook sandbox, web search, todo list) working in its own directory. +It explores for as many turns as it needs; when it stops calling tools the +episode ends and the task's checks are run against what it left behind. The +checks are ordinary programs -- file exists, file content, command exit status, +final answer match -- so the same trajectory always earns the same reward. + +Everything framework-specific lives here and in ``rsi_agent.yaml``: the tool +line-up, the sandbox settings, the task file. ``src/twinkle_agentic`` stays +generic -- ``MsAgentHarness`` shapes messages, ``MsAgentToolEnv`` executes +tools, ``result_check`` scores outcomes, and none of them know about RSI. + +Layout mirrors cookbook/rl/multi_turn/multi_turn_grpo.py; the differences are +the harness (ms-agent owns the system prompt and message evolution) and the +reward (program checks over the end state instead of an env-emitted scalar). + +Usage: + RSI_TASKS=cookbook/rl/rsi_agentic/tasks.example.jsonl \\ + python cookbook/rl/rsi_agentic/rsi_agentic_grpo.py + +Task file: one JSON object per line, with ``id``, ``query`` and ``checks`` +(see tasks.example.jsonl and twinkle_agentic.verifier.result_check.Check). +""" +import json +import os +import shutil +from typing import Any, Dict, List, Tuple + +from peft import LoraConfig + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_device_placement, get_logger +from twinkle.advantage import GRPOAdvantage +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.cli import CLI +from twinkle.data_format import SamplingParams +from twinkle.metric import CompletionRewardMetric +from twinkle.model import TransformersModel +from twinkle.processor import InputProcessor +from twinkle.sampler import vLLMSampler +from twinkle.template import Template +from twinkle_agentic.envs import EnvTool, MsAgentToolEnv +from twinkle_agentic.harness import MsAgentHarness +from twinkle_agentic.rollout.multi_turn import MultiTurnRollout +from twinkle_agentic.tools.tool_manager import ToolManager +from twinkle_agentic.verifier.result_check import (CheckContext, checks_from_dicts, + run_checks) + +logger = get_logger() +args = CLI.from_args() + +# ========== Configuration ========== +MODEL_ID = args.model.model_id or 'ms://Qwen/Qwen3-4B' + +MODEL_GPUS = args.infra.model_gpus or 4 +SAMPLER_GPUS = args.infra.sampler_gpus or 4 +NUM_GPUS = MODEL_GPUS + SAMPLER_GPUS + +NUM_GENERATIONS = args.rl.num_generations or 8 +MAX_NEW_TOKENS = args.sampling.max_tokens or 4096 +LEARNING_RATE = args.optimizer.learning_rate or 1e-5 +MAX_STEPS = args.training.max_steps or 1000 +BATCH_SIZE = args.training.batch_size or 4 +MINI_BATCH_SIZE = args.training.mini_batch_size or 8 +MICRO_BATCH_SIZE = args.training.micro_batch_size or 2 +GRADIENT_ACCUMULATION_STEPS = args.training.gradient_accumulation_steps or 1 +ADAPTER_NAME = args.lora.adapter_name or 'default' +SAVE_STEPS = args.training.save_steps or 500 +LORA_RANK = args.lora.lora_r or 16 + +# A tool-using episode needs room for observations on top of its own tokens. +MAX_TRAJECTORY_TOKENS = int(os.environ.get('RSI_MAX_TRAJ_TOKENS', 32768)) +MAX_TURNS = int(os.environ.get('RSI_MAX_TURNS', 20)) + +TASKS_PATH = os.environ.get('RSI_TASKS', 'cookbook/rl/rsi_agentic/tasks.example.jsonl') +AGENT_CONFIG = os.environ.get('RSI_AGENT_CONFIG', 'cookbook/rl/rsi_agentic/rsi_agent.yaml') +RUN_DIR = os.environ.get('RSI_RUN_DIR', 'output/rsi_agentic/run') +# 'fraction' gives partial credit per check; 'all_or_nothing' is stricter and +# produces a cleaner pass/fail signal at the cost of a sparser reward. +SCORE_MODE = os.environ.get('RSI_SCORE_MODE', 'fraction') +# Keep each episode's workspace after scoring. Useful while debugging tasks, +# expensive over a long run. +KEEP_WORKSPACES = os.environ.get('RSI_KEEP_WORKSPACES', '0') == '1' + + +def load_tasks(path: str) -> List[Dict[str, Any]]: + """Read the task file and fail loudly on a task that can never be scored.""" + tasks = [] + with open(path, encoding='utf-8') as f: + for lineno, line in enumerate(f, 1): + if not line.strip(): + continue + task = json.loads(line) + if not task.get('query'): + raise ValueError(f'{path}:{lineno} has no query') + if not task.get('checks'): + # An unchecked task scores 0 for every rollout, so the whole + # group has zero advantage and contributes no gradient. + raise ValueError(f'{path}:{lineno} ({task.get("id")}) declares no checks') + task['_checks'] = checks_from_dicts(task['checks']) + tasks.append(task) + if not tasks: + raise ValueError(f'{path} contains no tasks') + return tasks + + +def build_episode(task: Dict[str, Any], slot: int, step: int) -> Tuple[Any, Any, Any, Dict]: + """Create one episode: harness + isolated workspace + bound tool manager. + + The harness and the Env share one ms-agent runtime, so the tools named in + the prompt are exactly the tools that will run. Each episode gets its own + ``output_dir`` -- that directory is both the sandbox root and what the + checks will later inspect. + """ + from omegaconf import OmegaConf, open_dict + + workspace = os.path.join(RUN_DIR, f'step{step:06d}', f'slot{slot:03d}') + os.makedirs(workspace, exist_ok=True) + + cfg = OmegaConf.load(AGENT_CONFIG) + with open_dict(cfg): + cfg.output_dir = os.path.abspath(workspace) + + harness = MsAgentHarness(config=cfg) + # ms-agent merges the config above over its own default agent.yaml, which + # declares an `llm:` section; FileSystemTool then builds a remote LLM client + # from it and asserts on a missing api key. Generation here comes from the + # vLLM sampler, so drop that section before any tool is constructed. + with open_dict(harness.agent.config): + harness.agent.config.pop('llm', None) + harness.prepare() + + env = MsAgentToolEnv(agent=harness.agent, workspace=workspace) + # Same schema list on both sides: prompt and executor cannot drift apart. + tool_manager = ToolManager(EnvTool.from_schemas(env, harness.tool_schemas())) + + trajectory = harness.start(task['query']) + return harness, env, tool_manager, trajectory + + +def score_episode(task: Dict[str, Any], env: MsAgentToolEnv, trajectory: Dict[str, Any]) -> float: + """Run the task's checks against the state this episode left behind.""" + final_answer = '' + for msg in reversed(trajectory.get('messages') or []): + if msg.get('role') == 'assistant' and (msg.get('content') or '').strip(): + final_answer = msg['content'] + break + + ctx = CheckContext( + workspace=env.workspace, + final_answer=final_answer, + # Route shell/python checks back through the episode's own sandbox so + # they see the filesystem the agent actually wrote to. + runner=env.runner(), + ) + report = run_checks(task['_checks'], ctx, mode=SCORE_MODE) + if not report.all_passed: + logger.debug(f'[{task["id"]}] {report.n_passed}/{report.n_total} checks: ' + f'{report.failures()}') + return report.score + + +def main(): + tasks = load_tasks(TASKS_PATH) + logger.info(f'Loaded {len(tasks)} tasks from {TASKS_PATH}') + + device_groups = [ + DeviceGroup(name='model', ranks=list(range(MODEL_GPUS)), device_type='GPU'), + DeviceGroup(name='sampler', ranks=list(range(MODEL_GPUS, 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, dp_size=SAMPLER_GPUS) + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, groups=device_groups, + lazy_collect=False) + + lora_config = LoraConfig( + target_modules='all-linear', + r=LORA_RANK, + lora_alpha=LORA_RANK * 2, + lora_dropout=0.05, + ) + + model = TransformersModel(model_id=MODEL_ID, device_mesh=model_mesh, remote_group='model') + model.add_adapter_to_model(ADAPTER_NAME, lora_config, + gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) + 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('Template', model_id=MODEL_ID, enable_thinking=True) + + sampler = vLLMSampler( + model_id=MODEL_ID, + engine_args={ + 'gpu_memory_utilization': 0.8, + 'max_model_len': MAX_TRAJECTORY_TOKENS, + 'max_lora_rank': 32, + 'enable_lora': True, + 'enable_tower_connector_lora': True, + }, + device_mesh=sampler_mesh, + remote_group='sampler', + ) + sampler.set_template('Template', model_id=MODEL_ID, enable_thinking=True) + + rollout_template = Template(MODEL_ID, max_length=MAX_TRAJECTORY_TOKENS, enable_thinking=True) + rollout_template.truncation_strategy = 'delete' + + ckpt_manager = CheckpointEngineManager(model=model, sampler=sampler) + sampling_params = SamplingParams(max_tokens=MAX_NEW_TOKENS, num_samples=1, logprobs=1, + temperature=1.0, top_p=0.95) + rollout = MultiTurnRollout( + sampler=sampler, + template=rollout_template, + sampling_params=sampling_params, + max_turns=MAX_TURNS, + max_trajectory_tokens=MAX_TRAJECTORY_TOKENS, + ) + + advantage_fn = GRPOAdvantage() + metrics = CompletionRewardMetric() + + optim_step = 0 + task_cursor = 0 + logger.info(f'Starting agentic RSI GRPO (max_turns={MAX_TURNS}, score={SCORE_MODE})') + logger.info(get_device_placement()) + + while optim_step < MAX_STEPS: + metrics.reset() + + # Each prompt is repeated NUM_GENERATIONS times; GRPO needs a group of + # rollouts on the SAME task to have anything to compare against. + batch_tasks = [tasks[(task_cursor + i) % len(tasks)] for i in range(BATCH_SIZE)] + task_cursor = (task_cursor + BATCH_SIZE) % len(tasks) + episode_tasks = [t for t in batch_tasks for _ in range(NUM_GENERATIONS)] + + harnesses, envs, tool_managers, trajectories = [], [], [], [] + for slot, task in enumerate(episode_tasks): + h, env, tm, traj = build_episode(task, slot, optim_step) + harnesses.append(h) + envs.append(env) + tool_managers.append(tm) + trajectories.append(traj) + + ckpt_manager.sync_weights(merge_and_sync=False) + sampler.reset_prefix_cache() + + try: + outs: List[Dict[str, Any]] = rollout( + trajectories, harness=harnesses, tool_manager=tool_managers) + + rewards = [score_episode(task, env, traj) + for task, env, traj in zip(episode_tasks, envs, outs)] + finally: + # Sandboxes are a finite resource; a step that raises must still + # give them back or the next step starts short. + for env in envs: + env.close() + if not KEEP_WORKSPACES: + shutil.rmtree(os.path.join(RUN_DIR, f'step{optim_step:06d}'), ignore_errors=True) + + all_old_logps, completion_lengths, turns = [], [], [] + for traj in outs: + logprobs = traj.get('logprobs') or [] + all_old_logps.append([lp[0][1] for lp in logprobs] if logprobs else []) + labels = traj.get('labels') or [] + completion_lengths.append(sum(1 for label in labels if label != -100)) + turns.append(int(traj.get('turns') or 0)) + + advantages = advantage_fn(rewards, num_generations=NUM_GENERATIONS, + scale='group').tolist() + metrics.accumulate(completion_lengths=completion_lengths, rewards={'total': rewards}) + + avg_reward = sum(rewards) / len(rewards) if rewards else 0.0 + solved = sum(1 for r in rewards if r >= 1.0) + logger.info(f'[Step {optim_step}] avg_reward={avg_reward:.3f} ' + f'fully_solved={solved}/{len(rewards)} ' + f'avg_turns={sum(turns)/max(1,len(turns)):.1f}') + + # Drop episodes the template refused (too long) or that produced no + # trainable tokens; feeding those in would corrupt the logp alignment. + inputs, kept_logps, kept_adv = [], [], [] + for i, traj in enumerate(outs): + if not completion_lengths[i]: + continue + if len(traj.get('input_ids') or []) > MAX_TRAJECTORY_TOKENS: + continue + inputs.append(traj) + kept_logps.append(all_old_logps[i]) + kept_adv.append(advantages[i]) + + if len(inputs) < MODEL_GPUS: + logger.warning(f'[Step {optim_step}] only {len(inputs)} usable trajectories ' + f'(need >= {MODEL_GPUS}); skipping batch') + continue + + for mb_start in range(0, len(inputs), MINI_BATCH_SIZE): + mb_end = min(mb_start + MINI_BATCH_SIZE, len(inputs)) + model.forward_backward( + inputs=inputs[mb_start:mb_end], + old_logps=kept_logps[mb_start:mb_end], + advantages=kept_adv[mb_start:mb_end], + 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'rsi-agentic-checkpoint-{optim_step}') + + log_dict = metrics.calculate() + log_dict.update(model.calculate_metric(is_training=True)) + log_dict['avg_reward'] = avg_reward + log_dict['fully_solved'] = solved + logger.info(f'[Step {optim_step}/{MAX_STEPS}] {log_dict}') + + logger.info(f'Training completed. optim_steps={optim_step}') + model.save('rsi-agentic-final') + + +if __name__ == '__main__': + main() diff --git a/cookbook/rl/rsi_agentic/tasks.example.jsonl b/cookbook/rl/rsi_agentic/tasks.example.jsonl new file mode 100644 index 000000000..d9d9a37e3 --- /dev/null +++ b/cookbook/rl/rsi_agentic/tasks.example.jsonl @@ -0,0 +1,4 @@ +{"id": "t_file_001", "query": "In the current directory, create a file named report.md that contains a level-1 markdown heading reading 'Sales Report' followed by one bullet per quarter (Q1 to Q4).", "checks": [{"kind": "file_exists", "path": "report.md", "description": "report.md was created"}, {"kind": "file_contains", "path": "report.md", "value": "# Sales Report", "description": "has the required H1"}, {"kind": "file_contains", "path": "report.md", "pattern": "(?s)Q1.*Q2.*Q3.*Q4", "description": "lists all four quarters in order"}]} +{"id": "t_code_001", "query": "Write a Python module solution.py exposing a function `merge_intervals(intervals)` that merges overlapping closed intervals and returns them sorted by start. Verify it yourself before finishing.", "checks": [{"kind": "file_exists", "path": "solution.py", "description": "solution.py was created"}, {"kind": "python", "code": "from solution import merge_intervals\nassert merge_intervals([[1,3],[2,6],[8,10]]) == [[1,6],[8,10]]\nassert merge_intervals([]) == []\nassert merge_intervals([[1,4],[4,5]]) == [[1,5]]", "description": "merge_intervals passes the reference cases"}]} +{"id": "t_cli_001", "query": "Create a directory named logs/ holding three files a.log, b.log and c.log. Each must have at least 8 lines, and exactly 7 lines across all three must contain the word ERROR. Then count those ERROR lines with shell commands and write just that number into error_count.txt.", "checks": [{"kind": "file_exists", "path": "error_count.txt", "description": "error_count.txt was created"}, {"kind": "shell", "code": "test -f logs/a.log && test -f logs/b.log && test -f logs/c.log", "description": "all three log files exist"}, {"kind": "shell", "code": "test \"$(cat logs/a.log logs/b.log logs/c.log | grep -c ERROR)\" = 7", "description": "exactly 7 ERROR lines were written"}, {"kind": "shell", "code": "test \"$(tr -d '[:space:]' < error_count.txt)\" = \"$(cat logs/a.log logs/b.log logs/c.log | grep -c ERROR)\"", "description": "the recorded count matches a fresh recount"}]} +{"id": "t_data_001", "query": "Create a CSV file data.csv with a header line `name,score` and exactly 5 data rows of your choosing. Then compute the mean score and write it, rounded to 2 decimal places, into mean.txt as the only content.", "checks": [{"kind": "file_exists", "path": "data.csv", "description": "data.csv was created"}, {"kind": "file_exists", "path": "mean.txt", "description": "mean.txt was created"}, {"kind": "python", "code": "import csv\nwith open('data.csv') as f:\n rows = list(csv.DictReader(f))\nassert len(rows) == 5, len(rows)\nmean = sum(float(r['score']) for r in rows) / len(rows)\ngot = float(open('mean.txt').read().strip())\nassert abs(got - round(mean, 2)) < 0.01, (got, mean)", "description": "mean.txt matches the mean recomputed from data.csv"}]} diff --git a/cookbook/rsi/run_rsi_selfplay.py b/cookbook/rsi/run_rsi_selfplay.py new file mode 100644 index 000000000..c8c6bdc4d --- /dev/null +++ b/cookbook/rsi/run_rsi_selfplay.py @@ -0,0 +1,125 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""RSI self-play entry point — the challenger/solver data loop, no prepare/refine. + +Two stages live in ``twinkle_agentic.rsi`` (one model, Qwen3-4B, plays both roles): + + 1 challenge twinkle_agentic.rsi.rsi_challenge (ray+GPU) [seed] -> flows + tests + 2 rl twinkle_agentic.rsi.rsi_rl (ray+GPU) flows -> trained model + +Stage 1 asks the model to invent (or vary a seed into) self-contained Python +problems, RUNS its reference solution to get the ground truth, turns that into +asserts, then keeps only the problems the same model solves sometimes-but-not- +always (0 < pass < N). Stage 3 trains on those with ``RSI_SOLVER_MODE``: +``grpo`` (feed the sandbox error back as a tool turn and continue) or ``opsd`` +(a teacher that saw the reference solution distills the student). + +Why this launches SUBPROCESSES instead of importing and calling (same reason as +run_rsi.py): + * ``rsi_rl`` runs ``CLI.from_args()`` and ``swanlab.init()`` at IMPORT time, so + merely importing it would parse this launcher's argv and start a run. + * the two stages need different ray topologies (sampler-only vs trainer+sampler) + and cannot share one ray init in-process. + +This launcher invents no parameters: it wires stage 1's default outputs into +stage 2's inputs (through each script's own env vars) and forwards any extra +flags straight through to the selected stage. + +Examples +-------- +Run one stage at a time (extra flags after the known ones are forwarded): + + # from scratch (no seed dataset) + python cookbook/rsi/run_rsi_selfplay.py --step challenge + # from a seed dataset (challenger writes variants of its queries) + python cookbook/rsi/run_rsi_selfplay.py --step challenge --seed data/seed.jsonl + # train — grpo (default) or opsd; twinkle CLI knobs are forwarded as extras + python cookbook/rsi/run_rsi_selfplay.py --step rl --mode grpo \ + --model.model_id ms://Qwen/Qwen3-4B --infra.model_gpus 4 --infra.sampler_gpus 4 + +Run the whole chain with default paths (each stage still a fresh process): + + python cookbook/rsi/run_rsi_selfplay.py --step all --mode grpo +""" +import argparse +import os +import subprocess +import sys + +# Default paths chain stage 1 into stage 2. These mirror the defaults baked into +# each stage's own env-var config, kept here so --step all wires up with no flags. +DEFAULT_FLOWS = 'output/rsi/challenge_flows.jsonl' # rsi_challenge RSI_CH_OUT_FLOWS / rsi_rl RSI_STD_FLOWS +DEFAULT_TESTS = 'output/rsi/challenge_tests.jsonl' # rsi_challenge RSI_CH_OUT_TESTS / rsi_rl RSI_TESTS + +MODULES = { + 'challenge': 'twinkle_agentic.rsi.rsi_challenge', + 'rl': 'twinkle_agentic.rsi.rsi_rl', +} +ORDER = ['challenge', 'rl'] + + +def _run(module: str, argv: list, env: dict) -> None: + """Run ``python -m module argv...`` as a child process, streaming its output. + + Raises on non-zero exit so --step all stops at the first failing stage + instead of silently feeding a broken artifact into the next stage. + """ + cmd = [sys.executable, '-m', module] + argv + print(f'\n[run_rsi_selfplay] $ {" ".join(cmd)}', flush=True) + subprocess.run(cmd, env=env, check=True) + + +def _argv_for(step: str, a: argparse.Namespace, extra: list) -> tuple: + """Build (argv, env) for one stage. ``extra`` is forwarded verbatim so each + stage's own flags (twinkle CLI knobs for rl) still work.""" + env = dict(os.environ) + if step == 'challenge': + # rsi_challenge is configured purely through RSI_CH_* env vars (no CLI). + env['RSI_CH_OUT_FLOWS'] = a.flows + env['RSI_CH_OUT_TESTS'] = a.tests + if a.seed: + env['RSI_CH_SEED'] = a.seed + return list(extra), env + if step == 'rl': + # rsi_rl reads flows/tests and the solver mode from env vars; the + # model/infra/rl knobs arrive through `extra` (twinkle CLI). + env['RSI_STD_FLOWS'] = a.flows + env['RSI_TESTS'] = a.tests + env['RSI_SOLVER_MODE'] = a.mode + return list(extra), env + raise SystemExit(f'[run_rsi_selfplay] 未知 step: {step}') + + +def main(): + parser = argparse.ArgumentParser( + description='RSI self-play launcher — run one stage (validate) or the whole chain.', + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument('--step', required=True, choices=ORDER + ['all'], + help='Which stage to run (or "all" for challenge->rl).') + parser.add_argument('--seed', default='', + help='Optional seed dataset (jsonl) for challenge; empty = invent from scratch.') + parser.add_argument('--flows', default=DEFAULT_FLOWS, + help='challenge output flows / rl standard-flow input.') + parser.add_argument('--tests', default=DEFAULT_TESTS, + help='challenge output tests / rl code-round asserts input.') + parser.add_argument('--mode', default='grpo', choices=['grpo', 'opsd'], + help='rl solver mode (RSI_SOLVER_MODE).') + a, extra = parser.parse_known_args() + + if a.step == 'all': + if extra: + # For 'all' the extras are ambiguous (which stage?); refuse rather than + # forward a flag to a stage that does not accept it. + raise SystemExit(f'[run_rsi_selfplay] --step all 不接受透传参数 {extra}；' + '请逐个 --step 跑并各自带参数') + for step in ORDER: + argv, env = _argv_for(step, a, []) + _run(MODULES[step], argv, env) + print('\n[run_rsi_selfplay] all stages done.', flush=True) + return + + argv, env = _argv_for(a.step, a, extra) + _run(MODULES[a.step], argv, env) + + +if __name__ == '__main__': + main() diff --git a/src/twinkle/template/qwen3_5_vl.py b/src/twinkle/template/qwen3_5_vl.py index 2655a78ef..2e96da738 100644 --- a/src/twinkle/template/qwen3_5_vl.py +++ b/src/twinkle/template/qwen3_5_vl.py @@ -136,7 +136,7 @@ def to_tensor(_input): value = _input[key] if isinstance(value, np.ndarray): value = torch.from_numpy(value) - elif isinstance(value, list) and isinstance(value[0], (int, float, np.number)): + elif isinstance(value, list) and len(value) > 0 and isinstance(value[0], (int, float, np.number)): value = torch.tensor(value) _input[key] = value return _input diff --git a/src/twinkle/template/tools/__init__.py b/src/twinkle/template/tools/__init__.py index 8bb5d0db1..bb5a3cfd2 100644 --- a/src/twinkle/template/tools/__init__.py +++ b/src/twinkle/template/tools/__init__.py @@ -6,6 +6,7 @@ over weaker fallbacks. """ from .base import ToolCallParser, ToolCallRegistry +from .bracket_dsl import BracketDslParser from .cline import ClineParser from .qwen import HermesQwenParser from .react import ReActParser @@ -17,6 +18,9 @@ ToolCallRegistry.register(ClineParser()) ToolCallRegistry.register(VCPParser()) ToolCallRegistry.register(ReActParser()) +# Last: the bracketed call list carries no markup of its own, so it must only +# claim text that no marked-up format recognised. +ToolCallRegistry.register(BracketDslParser()) __all__ = [ 'ToolCallParser', @@ -25,4 +29,5 @@ 'ClineParser', 'VCPParser', 'ReActParser', + 'BracketDslParser', ] diff --git a/src/twinkle/template/tools/bracket_dsl.py b/src/twinkle/template/tools/bracket_dsl.py new file mode 100644 index 000000000..c3d1088a8 --- /dev/null +++ b/src/twinkle/template/tools/bracket_dsl.py @@ -0,0 +1,180 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import ast +import re +from typing import Any, Dict, List, Optional, Tuple + +from .base import ToolCallParser + + +class BracketDslParser(ToolCallParser): + """Parser for the bracketed call list used by ToolACE-style prompts. + + The system prompt of these datasets asks the model to answer with a python + call list instead of markup, e.g.:: + + [Text Analysis(text="great service"), UserID(username="alex")] + [quarterly_data(stock_symbols=["AAPL", "TSLA"])] + + Function names may contain spaces, dots and dashes ("Get All Strains", + "database.insert_data"). Argument values may themselves contain brackets and + parentheses (list arguments), so the call list is located by scanning with a + depth counter rather than by a bracket-free regex. Argument values are read + as python literals, falling back to the raw text when they are not literals. + """ + + name = 'bracket_dsl' + open_marker = None + close_marker = None + + # A call opens with a name directly followed by '('; used for cheap detection + # and to find call starts inside a located block. Names may carry spaces, + # dots, dashes and apostrophes ("Get Today's Prices"). + _CALL_START_RE = re.compile(r"([A-Za-z_][\w.\-' ]*?)\s*\(") + _DETECT_RE = re.compile(r"\[\s*[A-Za-z_][\w.\-' ]*?\s*\(") + # Split an argument body on top-level commas only (values may hold commas). + _ARG_NAME_RE = re.compile(r'^\s*([A-Za-z_]\w*)\s*=\s*(.*)$', re.DOTALL) + + def detect(self, text: str) -> bool: + return bool(self._DETECT_RE.search(text or '')) + + @staticmethod + def _find_blocks(text: str) -> List[Tuple[int, int]]: + """Locate ``[ ... ]`` spans that start a call list, honouring nesting. + + Only a '[' immediately followed by ``name(`` opens a block, so plain + prose lists ("[1, 2, 3]") are ignored. Quotes are honoured only inside an + argument body (paren depth > 0) so that an apostrophe in a function name + ("Get Today's Prices") does not start a string. + """ + spans: List[Tuple[int, int]] = [] + i, n = 0, len(text or '') + while i < n: + if text[i] != '[': + i += 1 + continue + if not BracketDslParser._DETECT_RE.match(text, i): + i += 1 + continue + depth, j, quote, paren = 0, i, None, 0 + while j < n: + ch = text[j] + if quote: + if ch == '\\': + j += 2 + continue + if ch == quote: + quote = None + elif ch in '"\'' and paren > 0: + quote = ch + elif ch == '(': + paren += 1 + elif ch == ')': + paren -= 1 + elif ch == '[': + depth += 1 + elif ch == ']': + depth -= 1 + if depth == 0: + spans.append((i, j + 1)) + break + j += 1 + i = (spans[-1][1] if spans and spans[-1][0] == i else i + 1) + return spans + + @staticmethod + def _match_paren(text: str, open_idx: int) -> Optional[int]: + """Index of the ')' matching the '(' at ``open_idx``.""" + depth, j, quote = 0, open_idx, None + n = len(text) + while j < n: + ch = text[j] + if quote: + if ch == '\\': + j += 2 + continue + if ch == quote: + quote = None + elif ch in '"\'' and depth > 0: + quote = ch + elif ch == '(': + depth += 1 + elif ch == ')': + depth -= 1 + if depth == 0: + return j + j += 1 + return None + + @staticmethod + def _split_top_level(body: str) -> List[str]: + """Split on commas that are not inside quotes, brackets or parens.""" + parts, buf = [], [] + depth, quote = 0, None + for ch in body or '': + if quote: + if ch == quote: + quote = None + buf.append(ch) + continue + if ch in '"\'': + quote = ch + elif ch in '([{': + depth += 1 + elif ch in ')]}': + depth -= 1 + elif ch == ',' and depth == 0: + parts.append(''.join(buf)) + buf = [] + continue + buf.append(ch) + if buf: + parts.append(''.join(buf)) + return parts + + def _parse_args(self, body: str) -> Dict[str, Any]: + args: Dict[str, Any] = {} + for chunk in self._split_top_level(body): + m = self._ARG_NAME_RE.match(chunk) + if not m: + continue + key = m.group(1) + raw = m.group(2).strip() + try: + args[key] = ast.literal_eval(raw) + except (ValueError, SyntaxError): + args[key] = raw.strip('"\'') + return args + + def parse(self, text: str) -> List[Dict[str, Any]]: + calls: List[Dict[str, Any]] = [] + text = text or '' + for start, end in self._find_blocks(text): + block = text[start:end] + pos = 1 # skip the opening '[' + while pos < len(block): + m = self._CALL_START_RE.search(block, pos) + if not m: + break + close = self._match_paren(block, m.end() - 1) + if close is None: + break + name = m.group(1).strip() + if name: + calls.append({ + 'type': 'function', + 'function': { + 'name': name, + 'arguments': self._parse_args(block[m.end():close]), + }, + }) + pos = close + 1 + return calls + + def clean(self, text: str) -> str: + text = text or '' + out, last = [], 0 + for start, end in self._find_blocks(text): + out.append(text[last:start]) + last = end + out.append(text[last:]) + return ''.join(out).rstrip() diff --git a/src/twinkle/utils/__init__.py b/src/twinkle/utils/__init__.py index d5d1b698b..9c5b43857 100644 --- a/src/twinkle/utils/__init__.py +++ b/src/twinkle/utils/__init__.py @@ -14,5 +14,5 @@ split_cp_inputs, stateless_init_process_group, to_device) from .transformers_utils import find_all_linears, find_layers, get_modules_to_not_convert from .unsafe import check_unsafe, trust_remote_code -from .utils import copy_files_by_pattern, deep_getattr, get_runtime_meta +from .utils import copy_files_by_pattern, deep_getattr, get_runtime_meta, run_sync from .vision_tools import load_image, load_mm_file diff --git a/src/twinkle/utils/utils.py b/src/twinkle/utils/utils.py index 40894a689..57051f099 100644 --- a/src/twinkle/utils/utils.py +++ b/src/twinkle/utils/utils.py @@ -1,10 +1,13 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +import asyncio +import concurrent.futures import fnmatch import glob import inspect import os import shutil from functools import lru_cache +from typing import Any, Callable def deep_getattr(obj, attr: str, default=None): @@ -131,3 +134,21 @@ def get_runtime_meta() -> str: f'- **Rank**: `{rank}/{world_size}` (local_rank=`{local_rank}`)', ] return '\n'.join(lines) + + +def run_sync(async_fn: Callable[..., Any], *args, **kwargs): + """Run an async function from sync code. + + ``async_fn`` must be a *callable that returns a coroutine*, not an + already-created coroutine (those are bound to the creating loop). + """ + + def _go(): + return asyncio.run(async_fn(*args, **kwargs)) + + try: + asyncio.get_running_loop() + except RuntimeError: + return _go() + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: + return pool.submit(_go).result() diff --git a/src/twinkle_agentic/classifier/__init__.py b/src/twinkle_agentic/classifier/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/src/twinkle_agentic/classifier/base.py b/src/twinkle_agentic/classifier/base.py deleted file mode 100644 index d79f6c2e2..000000000 --- a/src/twinkle_agentic/classifier/base.py +++ /dev/null @@ -1,11 +0,0 @@ -from abc import ABC, abstractmethod - - -class Classifier(ABC): - - def __init__(self, model_path: str, **kwargs): - self.model_path = model_path - - @abstractmethod - def classify(self, text: str) -> str: - pass \ No newline at end of file diff --git a/src/twinkle_agentic/data_format/__init__.py b/src/twinkle_agentic/data_format/__init__.py deleted file mode 100644 index 6298015c8..000000000 --- a/src/twinkle_agentic/data_format/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .chunks import Chunk, Chunks diff --git a/src/twinkle_agentic/data_format/chunks.py b/src/twinkle_agentic/data_format/chunks.py deleted file mode 100644 index f13245f33..000000000 --- a/src/twinkle_agentic/data_format/chunks.py +++ /dev/null @@ -1,104 +0,0 @@ -import sys -from dataclasses import dataclass -from itertools import groupby -from typing import Any, Dict, List, Literal, Optional, Tuple, Union - -if sys.version_info[:2] <= (3, 11): - # Pydantic requirements. - from typing_extensions import TypedDict -else: - from typing import TypedDict - -_MULTIMODAL_TYPES = ('image', 'video', 'audio') -_MEDIA_BUCKETS = (('images', 'image'), ('videos', 'video'), ('audios', 'audio')) - - -class Chunk(TypedDict, total=False): - - type: Literal['text', 'image', 'video', 'audio'] - content: Union[str, Any] - raw: Union[str, Any] - role: str - round: int - - -@dataclass -class Chunks: - - chunks: List[Chunk] - - def to_trajectory( - self, - block_wrapper: Optional[Tuple[str, str]] = ('<block_{n}>', '</block_{n}>'), - ) -> Dict[str, Any]: - media: Dict[str, List[Any]] = {t: [] for t in _MULTIMODAL_TYPES} - bound: List[Chunk] = [] - wrap_counter = 0 - for c in self.chunks: - if c.get('type') in _MULTIMODAL_TYPES and not isinstance(c.get('raw'), dict): - media[c['type']].append(c.get('content')) - continue - if (block_wrapper and c.get('type') == 'text' and c.get('role') != 'tool'): - raw = c.get('raw') - is_condensed = isinstance(raw, dict) and raw.get('condensed') - content = c.get('content') - if is_condensed and isinstance(content, str) and content: - wrap_counter += 1 - prefix = block_wrapper[0].format(n=wrap_counter) - suffix = block_wrapper[1].format(n=wrap_counter) - c = {**c, 'content': f'{prefix}{content}{suffix}'} - bound.append(c) - - # Merge consecutive same-role chunks into one message via groupby. - messages = [ - self._group_to_message(role, list(grp)) - for role, grp in groupby(bound, key=lambda c: c.get('role') or 'user') - ] - - trajectory: Dict[str, Any] = {'messages': messages} - for plural, singular in _MEDIA_BUCKETS: - if media[singular]: - trajectory[plural] = media[singular] - return trajectory - - @staticmethod - def _group_to_message(role: str, group: List[Chunk]) -> Dict[str, Any]: - """Fold a same-role run of chunks into one :class:`Message`. - - Preserves the intra-group order so mixed text / image / video / audio - parts round-trip back into OpenAI-style structured ``content``. - """ - reasoning: List[str] = [] - parts: List[Dict[str, Any]] = [] - tool_calls: List[Dict[str, Any]] = [] - tool_call_id: Optional[str] = None - has_media = False - - for c in group: - t, raw, content = c.get('type'), c.get('raw'), c.get('content') - kind = raw.get('kind') if isinstance(raw, dict) else None - # Any chunk in the group may carry the shared ``tool_call_id``. - if isinstance(raw, dict) and raw.get('tool_call_id') and tool_call_id is None: - tool_call_id = raw['tool_call_id'] - - if t == 'text' and kind == 'reasoning_content' and content: - reasoning.append(content) - elif t == 'text' and kind == 'tool_call' and isinstance(raw.get('tool_call'), dict): - tool_calls.append(dict(raw['tool_call'])) - elif t == 'text' and content: - parts.append({'type': 'text', 'text': content}) - elif t in _MULTIMODAL_TYPES and isinstance(raw, dict): - has_media = True - # Drop condenser-only markers, keep the original part shape. - parts.append({k: v for k, v in raw.items() if k != 'condensed'} or {'type': t, t: content}) - - msg: Dict[str, Any] = {'role': role} - if reasoning: - msg['reasoning_content'] = '\n\n'.join(reasoning) - if parts: - msg['content'] = parts if has_media else '\n\n'.join(p['text'] for p in parts) - if tool_calls: - msg['tool_calls'] = tool_calls - if tool_call_id is not None: - msg['tool_call_id'] = tool_call_id - return msg diff --git a/src/twinkle_agentic/envs/__init__.py b/src/twinkle_agentic/envs/__init__.py index a3cf38814..7ac8386ad 100644 --- a/src/twinkle_agentic/envs/__init__.py +++ b/src/twinkle_agentic/envs/__init__.py @@ -1,4 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from .base import Env, StepResult from .env_tool import EnvTool +from .ms_agent_tool_env import MsAgentToolEnv from .openenv import EnvPool, EnvPoolAdapter, OpenEnv diff --git a/src/twinkle_agentic/envs/base.py b/src/twinkle_agentic/envs/base.py index 552d9e1a5..52c8fda92 100644 --- a/src/twinkle_agentic/envs/base.py +++ b/src/twinkle_agentic/envs/base.py @@ -1,7 +1,7 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from abc import ABC, abstractmethod from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Sequence, Tuple from twinkle.data_format import Trajectory from twinkle.data_format.message import Tool as ToolInfo @@ -24,6 +24,10 @@ class Env(ABC): env = SomeEnv(...) result = env.reset() result = env.step(tool_name, arguments) + + Tool-call markup is parsed upstream by + :meth:`twinkle.template.base.Template.parse_tool_call`. This class only + executes already-split ``(tool_name, arguments)`` pairs. """ def reset(self, trajectory: Optional[Trajectory] = None) -> StepResult: @@ -33,6 +37,18 @@ def reset(self, trajectory: Optional[Trajectory] = None) -> StepResult: def step(self, tool_name: str, arguments: Dict[str, Any]) -> StepResult: raise NotImplementedError + def step_batch( + self, + calls: Sequence[Tuple[str, Dict[str, Any]]], + ) -> List[StepResult]: + """Execute a batch of already-parsed ``(tool_name, arguments)`` pairs. + + Default is a serial loop over :meth:`step`. Subclasses that talk to a + remote sandbox should override this so MultiTurn can keep tools off + the generate critical path. + """ + return [self.step(name, args or {}) for name, args in calls] + def tools(self) -> List[ToolInfo]: return [] diff --git a/src/twinkle_agentic/envs/env_tool.py b/src/twinkle_agentic/envs/env_tool.py index 3b2409da9..5732a75b7 100644 --- a/src/twinkle_agentic/envs/env_tool.py +++ b/src/twinkle_agentic/envs/env_tool.py @@ -1,6 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """EnvTool: bridges any Env to ToolManager.""" -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Tuple from twinkle.data_format.message import Tool as ToolInfo from .base import Env, StepResult @@ -25,6 +25,13 @@ def __call__(self, tool_name: str, arguments: Dict[str, Any]) -> str: self.last_result = result return result.observation + def call_many(self, calls: List[Tuple[str, Dict[str, Any]]]) -> List[str]: + """Batch through ``Env.step_batch``.""" + results = self._env.step_batch(calls) + if results: + self.last_result = results[-1] + return [r.observation for r in results] + def tool_info(self) -> ToolInfo: return { 'type': 'function', @@ -45,6 +52,38 @@ def episode_reward(self) -> float: return self.last_result.info['episode_reward'] return self.last_result.reward if self.last_result else 0.0 + @classmethod + def from_schemas(cls, env: Env, schemas: List[ToolInfo]) -> List['EnvTool']: + """Bind an externally-declared tool list to ``env``. + + Used when an agent framework owns the tool names/schemas that go into + the prompt and the Env only supplies the implementation, so training + and serving advertise the same tools. Every returned tool shares + ``env``, which lets :meth:`ToolManager.call_many` collapse a whole turn + into one :meth:`Env.step_batch`. + + Each name is forwarded to ``env.step`` verbatim, so the Env must accept + exactly these names. + """ + tools = [] + for info in schemas or []: + fn = info.get('function', {}) if isinstance(info, dict) else {} + name = fn.get('name') + if not name: + raise ValueError(f'tool schema without function.name cannot be bound to an ' + f'Env; the prompt would advertise an uncallable tool: {info!r}') + tools.append( + cls( + env=env, + tool_name=name, + description=fn.get('description', ''), + parameters=fn.get('parameters') or { + 'type': 'object', + 'properties': {} + }, + )) + return tools + @classmethod def from_env(cls, env: Env) -> List['EnvTool']: tool_infos = env.tools() diff --git a/src/twinkle_agentic/envs/ms_agent_tool_env.py b/src/twinkle_agentic/envs/ms_agent_tool_env.py new file mode 100644 index 000000000..65c11a26e --- /dev/null +++ b/src/twinkle_agentic/envs/ms_agent_tool_env.py @@ -0,0 +1,212 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Env backed by an ms-agent tool runtime. + +The harness declares which tools exist (``AgentHarness.tool_schemas``); this +Env is where those calls actually run. It owns no tool logic of its own -- it +forwards to the ``ToolManager`` that ms-agent already built (filesystem, shell, +python, notebook sandbox, web search, todo list) and adds the two things RL +needs on top: one workspace per episode, and batched dispatch so a turn's tool +calls do not run one at a time while the GPUs idle. + +Isolation comes from ``config.output_dir``: both ``FileSystemTool`` and +``CodeExecutionTool`` root themselves there, so giving every trajectory its own +directory keeps concurrent episodes from reading each other's files. + +Reward is deliberately absent. ``step`` always reports ``reward=0.0`` and +``done=False``; an agentic episode is scored after the fact from the state it +left behind (:mod:`twinkle_agentic.verifier.result_check`), and the rollout +ends when the model stops emitting tool calls. +""" +import json +import os +import re +from typing import Any, Dict, List, Optional, Sequence, Tuple + +from twinkle.utils import run_sync + +from .base import Env, StepResult + +__all__ = ['MsAgentToolEnv'] + +# Marker used to recover an exit status from a tool that only returns text. +_RC_MARK = '__TWINKLE_RC__' +_RC_RE = re.compile(rf'{_RC_MARK}:(-?\d+)') + +_PY_WRAPPER = """\ +import sys, traceback +try: +{body} +except SystemExit as _e: + print('{mark}:%d' % (_e.code or 0)) + sys.exit(0) +except BaseException: + traceback.print_exc() + print('{mark}:1') +else: + print('{mark}:0') +""" + + +class MsAgentToolEnv(Env): + """Execute ms-agent tool calls for one episode. + + Args: + agent: an ``LLMAgent`` whose tools are already prepared -- normally + ``MsAgentHarness.agent`` after ``harness.prepare()``. Sharing the + harness's agent is what guarantees the executing tool set is the + one the prompt advertised. + tool_manager: an ms-agent ``ToolManager`` to use instead of the + agent's. Only one of ``agent`` / ``tool_manager`` is needed. + workspace: directory this episode reads and writes. Defaults to the + agent's ``config.output_dir``. + max_observation_chars: truncate a tool result before it becomes a + message. A single ``grep`` can otherwise blow the context window + and truncate the trajectory mid-episode. + """ + + def __init__( + self, + agent: Any = None, + *, + tool_manager: Any = None, + workspace: str = '', + max_observation_chars: int = 8000, + ): + if agent is None and tool_manager is None: + raise ValueError('MsAgentToolEnv needs either agent= or tool_manager=') + self._agent = agent + self._tm = tool_manager if tool_manager is not None else getattr(agent, 'tool_manager', None) + if self._tm is None: + raise ValueError('no ms-agent ToolManager available; call harness.prepare() ' + 'before constructing the Env so tools are initialised') + self.workspace = workspace or self._workspace_from_agent(agent) + if self.workspace: + os.makedirs(self.workspace, exist_ok=True) + self.max_observation_chars = max_observation_chars + self._names: Optional[List[str]] = None + + # ------------------------------------------------------------------ Env + + def tool_names(self) -> List[str]: + """Names of the tools actually registered, as the runtime spells them.""" + if self._names is None: + raw = run_sync(self._tm.get_tools) + items: List[Any] = [] + if isinstance(raw, dict): + for value in raw.values(): + items.extend(value if isinstance(value, list) else [value]) + elif isinstance(raw, list): + items = raw + names = [] + for item in items: + if isinstance(item, dict): + fn = item.get('function') + name = (fn or {}).get('name') if isinstance(fn, dict) else None + name = name or item.get('tool_name') or item.get('name') + if name: + names.append(str(name)) + self._names = names + return list(self._names) + + def resolve_tool(self, name: str) -> str: + """Map a plain tool name onto the runtime's own spelling. + + ms-agent namespaces its tools as ``{server}---{tool}``, so a caller that + asks for ``shell_executor`` means ``code_executor---shell_executor``. + An unknown name raises instead of being passed through: a mistyped tool + comes back as a failed call, which for a checker is indistinguishable + from a failed check, and a whole GRPO group would silently score zero. + """ + names = self.tool_names() + if name in names: + return name + matches = [n for n in names if n.rsplit('---', 1)[-1] == name] + if len(matches) == 1: + return matches[0] + if not matches: + raise ValueError(f'no registered tool named {name!r}; available: {names}') + raise ValueError(f'{name!r} is ambiguous across servers: {matches}') + + def step(self, tool_name: str, arguments: Dict[str, Any]) -> StepResult: + result = run_sync(self._tm.single_call_tool, self._call(tool_name, arguments)) + return StepResult(observation=self._observation(result)) + + def step_batch(self, calls: Sequence[Tuple[str, Dict[str, Any]]]) -> List[StepResult]: + """Run a turn's calls concurrently through ms-agent's own gather.""" + calls = list(calls) + if not calls: + return [] + if len(calls) == 1: + return [self.step(calls[0][0], calls[0][1] or {})] + payload = [self._call(name, args or {}) for name, args in calls] + results = run_sync(self._tm.parallel_call_tool, payload) + return [StepResult(observation=self._observation(r)) for r in results] + + def close(self) -> None: + cleanup = getattr(self._tm, 'cleanup', None) + if cleanup is not None: + try: + run_sync(cleanup) + except Exception: # noqa + # Teardown must not take down a training step; a leaked sandbox + # is recoverable, a crashed trainer loses the whole batch. + pass + + # ------------------------------------------------------- for the checker + + def runner(self, shell_tool: str = 'shell_executor', python_tool: str = 'python_executor'): + """A ``result_check`` runner that executes inside this episode's sandbox. + + Verification has to see the same filesystem the agent wrote to, so the + check goes back through the same tools rather than a local subprocess. + Those tools return prose, not an exit status, so the command is made to + print a marker and the status is read back out of the output. + + The tool names are resolved against what is actually registered, so + plain names work regardless of how the runtime namespaces them. + """ + shell_name = self.resolve_tool(shell_tool) + python_name = self.resolve_tool(python_tool) + + def _run(source: str, interpreter: str) -> Tuple[int, str]: + if interpreter == 'python': + body = '\n'.join(' ' + line for line in source.splitlines()) or ' pass' + code = _PY_WRAPPER.format(body=body, mark=_RC_MARK) + out = self.step(python_name, {'code': code}).observation + else: + out = self.step(shell_name, {'command': f'{source}\necho "{_RC_MARK}:$?"'}).observation + match = _RC_RE.search(out or '') + if match is None: + # No marker means the tool itself failed (timeout, sandbox down) + # rather than the check failing; report non-zero and keep output. + return 1, out or 'check produced no output and no exit marker' + return int(match.group(1)), _RC_RE.sub('', out or '').strip() + + return _run + + # -------------------------------------------------------------- private + + @staticmethod + def _call(tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]: + return {'tool_name': tool_name, 'arguments': arguments or {}} + + def _observation(self, result: Any) -> str: + if result is None: + text = '' + elif isinstance(result, str): + text = result + else: + try: + text = json.dumps(result, ensure_ascii=False) + except (TypeError, ValueError): + text = str(result) + limit = self.max_observation_chars + if limit and len(text) > limit: + head = text[:limit] + text = f'{head}\n...[truncated {len(text) - limit} chars]' + return text + + @staticmethod + def _workspace_from_agent(agent: Any) -> str: + config = getattr(agent, 'config', None) + return str(getattr(config, 'output_dir', '') or '') if config is not None else '' diff --git a/src/twinkle_agentic/harness/__init__.py b/src/twinkle_agentic/harness/__init__.py new file mode 100644 index 000000000..ff7575514 --- /dev/null +++ b/src/twinkle_agentic/harness/__init__.py @@ -0,0 +1,14 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .base import AgentHarness + +__all__ = [ + 'AgentHarness', + 'MsAgentHarness', +] + + +def __getattr__(name: str): + if name == 'MsAgentHarness': + from .ms_agent import MsAgentHarness + return MsAgentHarness + raise AttributeError(f'module {__name__!r} has no attribute {name!r}') diff --git a/src/twinkle_agentic/harness/base.py b/src/twinkle_agentic/harness/base.py new file mode 100644 index 000000000..153a34a5c --- /dev/null +++ b/src/twinkle_agentic/harness/base.py @@ -0,0 +1,103 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Agent harness: framework-owned message/tool lifecycle, minus generate/execute. + +``MultiTurnRollout`` owns batched sampling and ``new_input_feature`` extension. +``Env`` owns tool execution. A harness mutates the same :class:`Trajectory` +the rest of the stack already uses (``messages`` / ``tools`` / ``user_data``). + +Only *append-only* mutations of ``messages`` are safe after the first encode: +rewriting earlier turns would break the token-id chain MultiTurn keeps in +``new_input_feature``. Implementations that compact/rewrite history must do +it in :meth:`start` / the first :meth:`before_generate` (before encode), or +opt in explicitly. +""" +from abc import ABC +from typing import Any, Dict, List, Optional + +from twinkle.data_format import Trajectory + + +class AgentHarness(ABC): + """Per-episode agent-framework hooks. + + Default implementations are no-ops so MultiTurn can take ``harness=None`` + or a subclass that only overrides some phases. Subclasses that wrap a + specific framework (ms-agent, …) live next to this file, not in + ``rollout/`` or ``rsi/``. Harness-private runtime (LLMAgent, session) + lives on the harness instance, not on the trajectory. + """ + + def tool_schemas(self) -> List[Dict[str, Any]]: + """OpenAI-shaped tool list this harness puts in the prompt. + + The harness owns the tool *names and schemas* so training and serving + advertise the identical set; the Env owns the *implementation*. Build + the executing side from the same list:: + + tm = ToolManager(EnvTool.from_schemas(env, harness.tool_schemas())) + + Skipping that step lets the prompt advertise tools the Env cannot run, + and every call comes back as an unknown-tool error. + """ + return [] + + def start(self, query: str, **kwargs) -> Trajectory: + """Open an episode: system + user (+ tool schema). + + Called by the training driver *before* MultiTurn encodes. Not invoked + by MultiTurn itself. Extra kwargs are merged onto the trajectory + (``user_data``, ``tools``, …). + """ + traj: Trajectory = {'messages': [{'role': 'user', 'content': query}]} + traj.update(kwargs) + return traj + + def before_generate(self, trajectory: Trajectory) -> Trajectory: + """Mutate ``trajectory`` immediately before a generate turn. + + First call happens before the initial ``template.encode``. Later calls + must be append-only relative to ``messages`` already in the pif, + or MultiTurn will ignore the rewrite to protect token alignment. + """ + return trajectory + + def after_generate( + self, + trajectory: Trajectory, + decoded: str, + tool_calls: Optional[List[Dict[str, Any]]] = None, + ) -> Trajectory: + """Normalize the assistant turn (content / tool_calls / reasoning). + + ``decoded`` and ``tool_calls`` come from the sampler; the pif already + contains the generated tokens. This hook only updates message metadata + so the next encode-bridge and the serving agent see the same shape. + """ + return trajectory + + def after_tools( + self, + trajectory: Trajectory, + observations: List[str], + tool_calls: Optional[List[Dict[str, Any]]] = None, + ) -> Trajectory: + """Turn raw Env observations into ``role=tool`` messages (append). + + Default: one tool message per observation, copying ``id`` / ``name`` + from the corresponding tool call when present. + """ + msgs = trajectory.setdefault('messages', []) + calls = list(tool_calls or []) + for i, obs in enumerate(observations): + msg: Dict[str, Any] = {'role': 'tool', 'content': obs if obs is not None else ''} + if i < len(calls): + tc = calls[i] if isinstance(calls[i], dict) else {} + fn = tc.get('function') if isinstance(tc.get('function'), dict) else {} + tid = tc.get('id') or tc.get('tool_call_id') + name = fn.get('name') or tc.get('name') or tc.get('tool_name') + if tid: + msg['tool_call_id'] = tid + if name: + msg['name'] = name + msgs.append(msg) + return trajectory diff --git a/src/twinkle_agentic/harness/ms_agent.py b/src/twinkle_agentic/harness/ms_agent.py new file mode 100644 index 000000000..948c82861 --- /dev/null +++ b/src/twinkle_agentic/harness/ms_agent.py @@ -0,0 +1,422 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""ms-agent harness: LLMAgent owns prompt/message evolution, not generate/execute. + +Training path: + + harness.start(query) # create_messages + tool schema + MultiTurnRollout # sampler.sample + Env.step_batch + harness.before_generate # memory / hooks / (optional) skill refresh + harness.after_generate # handle_new_response + harness.after_tools # tool-message shape (not tool execution) + +ms-agent owns the tool names and schemas so the prompt is identical in +training and serving; the Env owns the implementation. Wire the executing +side from the same list, or the prompt advertises tools the Env cannot run:: + + harness = MsAgentHarness(config) + harness.prepare() + tool_manager = ToolManager(EnvTool.from_schemas(env, harness.tool_schemas())) + rollout = MultiTurnRollout(sampler, template, + tool_manager=tool_manager, harness=harness) + outs = rollout([harness.start(q) for q in queries]) + +Serving path keeps using ``LLMAgent.run()`` with the same ``agent.yaml`` and +the same :class:`~twinkle_agentic.envs.base.Env` backend. This class must +**not** call ``llm.generate`` or ``parallel_tool_call`` (those execute tools). +""" +from __future__ import annotations + +import json +import uuid +from typing import Any, Dict, List, Optional, Union + +from twinkle import requires +from twinkle.data_format import Trajectory +from twinkle.utils import run_sync + +from .base import AgentHarness + + +class MsAgentHarness(AgentHarness): + """Harness that *calls* LLMAgent methods instead of copying their prompts. + + Args: + config: ms-agent ``DictConfig`` / dict / yaml path. Ignored when + ``agent`` is passed. + agent: an existing :class:`ms_agent.agent.llm_agent.LLMAgent`. + auto_prepare: run ``prepare_runtime`` / ``prepare_tools`` / skills / + memory on first :meth:`start`. Skip LLM init (training generate + is vLLM). Set ``False`` in unit tests that only need + ``create_messages``. + freeze_system: if True (default, RL-safe), do not rewrite + ``messages[0]`` after the episode starts. Skill/memory *append* + paths still run. + permission_mode: forced onto the agent so training never blocks on + a TUI/CLI confirm. ``auto`` matches non-interactive LLMAgent. + """ + + def __init__( + self, + config: Any = None, + *, + agent: Any = None, + auto_prepare: bool = True, + freeze_system: bool = True, + permission_mode: str = 'auto', + trust_remote_code: bool = False, + **agent_kwargs, + ): + requires('ms-agent') + from omegaconf import DictConfig, OmegaConf + + from ms_agent.agent.llm_agent import LLMAgent + + if agent is not None: + self.agent = agent + else: + if config is None: + cfg: Any = DictConfig({}) + elif isinstance(config, str): + cfg = OmegaConf.load(config) + elif isinstance(config, dict): + cfg = OmegaConf.create(config) + else: + cfg = config + self.agent = LLMAgent( + cfg, + trust_remote_code=trust_remote_code, + **agent_kwargs, + ) + self.auto_prepare = auto_prepare + self.freeze_system = freeze_system + self.permission_mode = permission_mode + self._prepared = False + self._apply_rl_stubs() + + # ------------------------------------------------------------------ public + + def prepare(self) -> None: + """Initialize tools / skills / memory (sync wrapper). Idempotent.""" + if self._prepared: + return + run_sync(self._prepare_async) + self._prepared = True + + def start(self, query: str, **kwargs) -> Trajectory: + if self.auto_prepare: + self.prepare() + messages = run_sync(self.agent.create_messages, query) + tools = self.tool_schemas() + traj: Trajectory = { + 'messages': self._messages_to_dicts(messages), + 'tools': tools, + } + traj.update(kwargs) + return traj + + def before_generate(self, trajectory: Trajectory) -> Trajectory: + from ms_agent.hooks.context import condense_hook_attachments_for_llm + + if self.auto_prepare: + self.prepare() + messages = self._dicts_to_messages(trajectory.get('messages') or []) + frozen_system = messages[0].content if (self.freeze_system and messages + and messages[0].role == 'system') else None + + messages = self.agent._append_task_notifications(messages) + messages = condense_hook_attachments_for_llm(messages) + + if getattr(self.agent, 'runtime', None) is not None: + run_sync(self.agent.on_generate_response, messages) + + if getattr(self.agent, 'context_assembler', None) is not None and not self.freeze_system: + # Compaction rewrites earlier turns — incompatible with + # new_input_feature extension. Only run when the caller opts in. + assembled = self.agent.context_assembler.assemble() + if assembled: + messages = self._dicts_to_messages(assembled) + + messages = run_sync(self.agent.condense_memory, messages) + + skill_runtime = getattr(self.agent, '_skill_runtime', None) + if skill_runtime is not None and not self.freeze_system: + skill_runtime.maybe_refresh_system_prompt(messages) + + if frozen_system is not None and messages and messages[0].role == 'system': + messages[0].content = frozen_system + + trajectory['messages'] = self._messages_to_dicts(messages) + return trajectory + + def after_generate( + self, + trajectory: Trajectory, + decoded: str, + tool_calls: Optional[List[Dict[str, Any]]] = None, + ) -> Trajectory: + messages = self._dicts_to_messages(trajectory.get('messages') or []) + response = self._assistant_message(decoded, tool_calls, messages) + self.agent.handle_new_response(messages, response) + if getattr(self.agent, 'runtime', None) is not None and response.tool_calls: + run_sync(self.agent.on_tool_call, messages) + trajectory['messages'] = self._messages_to_dicts(messages) + return trajectory + + def after_tools( + self, + trajectory: Trajectory, + observations: List[str], + tool_calls: Optional[List[Dict[str, Any]]] = None, + ) -> Trajectory: + """Format Env observations as ms-agent ``role=tool`` messages. + + Mirrors the *message construction* half of ``parallel_tool_call``; + does not execute tools. + """ + from ms_agent.llm.utils import Message, ToolResult + + messages = self._dicts_to_messages(trajectory.get('messages') or []) + calls = self._ms_tool_calls(tool_calls or self._last_assistant_calls(messages)) + for i, raw in enumerate(observations): + formatted = ToolResult.from_raw(raw) + tc = calls[i] if i < len(calls) else {} + tid = tc.get('id') or str(uuid.uuid4())[:8] + name = tc.get('tool_name') or '' + messages.append( + Message( + role='tool', + content=formatted.text, + tool_call_id=tid, + name=name, + resources=formatted.resources, + tool_detail=formatted.tool_detail, + hook_attachments=formatted.hook_attachments, + is_error=formatted.is_error, + )) + if i < len(calls) and not tc.get('id'): + calls[i]['id'] = tid + + skill_runtime = getattr(self.agent, '_skill_runtime', None) + if skill_runtime is not None and not self.freeze_system: + skill_runtime.maybe_refresh_system_prompt(messages) + + messages = run_sync(self.agent.condense_memory, messages) + if getattr(self.agent, 'runtime', None) is not None: + run_sync(self.agent.after_tool_call, messages) + + trajectory['messages'] = self._messages_to_dicts(messages) + return trajectory + + # ------------------------------------------------------------------ prepare + + def _apply_rl_stubs(self) -> None: + """Non-interactive: never block on TUI / permission prompts / stdin.""" + try: + from omegaconf import open_dict + with open_dict(self.agent.config): + self.agent.config.interactive = False + if self.permission_mode: + self.agent.config.permission_mode = self.permission_mode + except Exception: + pass + self.agent._interactive = False + self.agent._event_sink = None + self.agent._input_source = None + + async def _prepare_async(self) -> None: + agent = self.agent + if getattr(agent, 'runtime', None) is None: + agent.prepare_runtime() + if getattr(agent, 'tool_manager', None) is None: + await agent.prepare_tools() + await agent.prepare_skills() + await agent.load_memory() + if hasattr(agent, 'prepare_rag'): + await agent.prepare_rag() + if hasattr(agent, 'prepare_knowledge_search'): + await agent.prepare_knowledge_search() + + def tool_schemas(self) -> List[Dict[str, Any]]: + """ms-agent's own tool list, OpenAI-shaped. + + This is the list that reaches the prompt. Feed the same list to + ``EnvTool.from_schemas`` so the Env executes exactly what was + advertised. + """ + if self.auto_prepare: + self.prepare() + tm = getattr(self.agent, 'tool_manager', None) + if tm is None: + return [] + raw = run_sync(tm.get_tools) + return _ms_tools_to_openai(raw) + + # ------------------------------------------------------------------ convert + + def _assistant_message(self, decoded: str, tool_calls, messages): + from ms_agent.llm.utils import Message + + ms_calls = self._ms_tool_calls(tool_calls) + if messages and messages[-1].role == 'assistant': + response = messages[-1] + if ms_calls and not response.tool_calls: + response.tool_calls = ms_calls + if decoded and not response.content: + response.content = decoded + return response + return Message(role='assistant', content=decoded or '', tool_calls=ms_calls) + + @staticmethod + def _last_assistant_calls(messages) -> List[Dict[str, Any]]: + for msg in reversed(messages): + if getattr(msg, 'role', None) == 'assistant': + return list(getattr(msg, 'tool_calls', None) or []) + return [] + + @staticmethod + def _ms_tool_calls(tool_calls: Optional[List[Any]]) -> List[Dict[str, Any]]: + out: List[Dict[str, Any]] = [] + for tc in tool_calls or []: + if not isinstance(tc, dict): + continue + fn = tc.get('function') if isinstance(tc.get('function'), dict) else None + if fn is not None: + args = fn.get('arguments', '{}') + if isinstance(args, dict): + args = json.dumps(args, ensure_ascii=False) + out.append({ + 'id': tc.get('id') or '', + 'type': tc.get('type', 'function'), + 'tool_name': fn.get('name') or '', + 'arguments': args if isinstance(args, str) else '{}', + }) + continue + args = tc.get('arguments', '{}') + if isinstance(args, dict): + args = json.dumps(args, ensure_ascii=False) + out.append({ + 'id': tc.get('id') or '', + 'type': tc.get('type', 'function'), + 'tool_name': tc.get('tool_name') or tc.get('name') or '', + 'arguments': args if isinstance(args, str) else '{}', + }) + return out + + @staticmethod + def _messages_to_dicts(messages) -> List[Dict[str, Any]]: + out: List[Dict[str, Any]] = [] + for msg in messages: + if isinstance(msg, dict): + out.append(dict(msg)) + continue + d: Dict[str, Any] = { + 'role': msg.role, + 'content': msg.content if msg.content is not None else '', + } + if getattr(msg, 'tool_calls', None): + d['tool_calls'] = _ms_calls_to_openai(msg.tool_calls) + if getattr(msg, 'tool_call_id', None): + d['tool_call_id'] = msg.tool_call_id + if getattr(msg, 'name', None): + d['name'] = msg.name + if getattr(msg, 'reasoning_content', ''): + d['reasoning_content'] = msg.reasoning_content + out.append(d) + return out + + @staticmethod + def _dicts_to_messages(messages: List[Dict[str, Any]]): + from ms_agent.llm.utils import Message + + out = [] + for m in messages: + if not isinstance(m, dict): + out.append(m) + continue + kwargs: Dict[str, Any] = { + 'role': m.get('role') or 'user', + 'content': m.get('content') if m.get('content') is not None else '', + } + tcs = m.get('tool_calls') + if tcs: + kwargs['tool_calls'] = MsAgentHarness._ms_tool_calls(tcs) + if m.get('tool_call_id'): + kwargs['tool_call_id'] = m['tool_call_id'] + if m.get('name'): + kwargs['name'] = m['name'] + if m.get('reasoning_content'): + kwargs['reasoning_content'] = m['reasoning_content'] + out.append(Message(**kwargs)) + return out + + +def _ms_calls_to_openai(tool_calls: List[Any]) -> List[Dict[str, Any]]: + out: List[Dict[str, Any]] = [] + for tc in tool_calls or []: + if not isinstance(tc, dict): + continue + fn = tc.get('function') if isinstance(tc.get('function'), dict) else None + if fn is not None: + args = fn.get('arguments', '{}') + if isinstance(args, dict): + args = json.dumps(args, ensure_ascii=False) + item = { + 'id': tc.get('id') or '', + 'type': tc.get('type', 'function'), + 'function': { + 'name': fn.get('name') or '', + 'arguments': args if isinstance(args, str) else '{}', + }, + } + out.append(item) + continue + args = tc.get('arguments', '{}') + if isinstance(args, dict): + args = json.dumps(args, ensure_ascii=False) + out.append({ + 'id': tc.get('id') or '', + 'type': tc.get('type', 'function'), + 'function': { + 'name': tc.get('tool_name') or tc.get('name') or '', + 'arguments': args if isinstance(args, str) else '{}', + }, + }) + return out + + +def _ms_tools_to_openai(raw: Union[Dict[str, Any], List[Any], None]) -> List[Dict[str, Any]]: + if not raw: + return [] + items: List[Any] = [] + if isinstance(raw, dict): + for v in raw.values(): + if isinstance(v, list): + items.extend(v) + else: + items.append(v) + elif isinstance(raw, list): + items = raw + else: + return [] + out: List[Dict[str, Any]] = [] + for t in items: + if not isinstance(t, dict): + continue + if t.get('type') == 'function' and isinstance(t.get('function'), dict): + out.append(t) + continue + name = t.get('tool_name') or t.get('name') + if not name: + continue + out.append({ + 'type': 'function', + 'function': { + 'name': name, + 'description': t.get('description', ''), + 'parameters': t.get('parameters') or { + 'type': 'object', + 'properties': {}, + }, + }, + }) + return out diff --git a/src/twinkle_agentic/memory/DESIGN.md b/src/twinkle_agentic/memory/DESIGN.md deleted file mode 100644 index f27f1cb19..000000000 --- a/src/twinkle_agentic/memory/DESIGN.md +++ /dev/null @@ -1,745 +0,0 @@ -# Memory 线设计（自进化蒸馏框架的第二条线） - -> 状态：**设计稿，未实现**。本文件是讨论沉淀，供后续实现参考。 -> 与主线（清洗 + 评分 + 在线蒸馏）共用同一批 trajectory 与同一套 verifier / `user_data` 信封 / `llm_backup` / D7c 校准思想，不另起炉灶。 - ---- - -## 0. 背景与定位：双线 - -同一份生产 trajectory，榨出两种产物，固化到不同位置、不同时间尺度： - -| | 线 A：训模型（慢记忆） | 线 B：总结 memory（快记忆） | -|---|---|---| -| 固化位置 | 模型权重 | 外部 memory store | -| 生效方式 | 蒸馏/微调后永久具备 | 推理时检索注入 context | -| 时间尺度 | 天/周（攒批再训） | 秒/分钟（写完即用） | -| 改什么 | 参数 | 行为（不改参数） | -| 装什么 | 可泛化的**技能/模式** | 易变的**事实/偏好/近期上下文** | - -**line B 的存在理由**：在线蒸馏有延迟（攒数据→训练→部署），空窗期 student 学不到新东西；memory 立刻起作用填这个空窗，等能力被训进权重后再从 memory 退休。 - -> **line B 有两种载体（同为"不动 base 的外挂补充"，见 §11.6）**： -> 1. **文本 memory**：检索注入 context —— 本文 §1~§9 讨论的主形态，适合**易变事实/实体/偏好**（一条即用即改，天然带"检不到就拒答"信号）。 -> 2. **参数化 memory（LoRA，base 冻结）**：把能力挂成可插拔 adapter —— 适合**可泛化的行为模式**（如"边生成边查错"），在已部署 vLLM 里边际显存/进程成本最低、可回滚（换 adapter ≈ 删一条 memory，而非重训 base）。 -> -> 二者与 line A 的分界是"**动不动 base**"：line A 改 base 权重（有遗忘风险，故慎重、攒批离线做）；参数化 memory 虽然也是"权重形态"，但**base 全程冻结、只挂 LoRA**，本质仍是 line B（可插拔、可回滚、即挂即用）。参数化 memory 与文本 memory 的取舍见 Substrate Asymmetry（§11 反方证据）：参数化擅长行为/风格，文本检索擅长事实/拒答。 - -**双线共用的唯一“质检车间”**：现有 `preprocessor + verifier`。它产出的 `traj_score / round_scores / confidence / intent / safety` 同时作为两条线的准入闸门，不重复造。 - -### 分工判据（什么进 A，什么进 B） -用 **可泛化性 × 稳定性** 分流： -- 进 A（训权重）：高频、可泛化、稳定 —— 能变成“技能”的。 -- 进 B（memory）：低频 / 易变 / 实体绑定 —— 只能当“事实/上下文”的。训进权重会过拟合具体实体且会过时。 - -### 消费者与作用域（已确认） -- **消费者**：两个都要 —— ① 本地 student 推理时检索注入；② 辅助置信度路由。 -- **作用域**：两层 —— 用户级（个性化：偏好、历史）+ 全局级（跨用户通用事实/模式）。 - ---- - -## 1. 存储与检索选型（基于现有代码，零新增重依赖） - -现有可复用： -- **存储** `twinkle/server/state`：memory / file / redis 三后端，带 TTL、`update_atomic`（原子）、`keys(pattern)`（通配）。 -- **embedding** `preprocessor/experimental/llm_backend.py`：`OpenAIBackend.embeddings`（走 HTTP，dashscope 有 embedding 接口）。 -- **没有**专用向量库（无 Milvus/Qdrant/faiss），但 D1 的 MinHash 说明“近似检索土办法”本项目可接受。 - -**结论（最小可跑，可后续换库）**： -``` -MemoryStore -├── 用户级（结构化 KV/标签）── 复用 StateBackend（FileBackend 默认，可切 Redis） -│ key = mem::user:<uid>::<kind>::<slot>；精确/前缀匹配；零额外依赖 -└── 全局级（向量语义召回）── 条目存 StateBackend；embedding 用 OpenAIBackend.embeddings - 召回 = 应用层 cosine top-k（初期 O(n)，涨了再换 faiss/Qdrant） - 抽象出 VectorIndex 接口，后端可替换 -``` -- 用户级以**结构化 KV** 为主（实体绑定要精确命中，向量会招噪声）。 -- 全局级以**向量召回**为主，再用结构化标签（intent/domain/min_score）过滤。 - -### MemoryItem（统一条目） -```python -@dataclass -class MemoryItem: - id: str - scope: str # 'user:<uid>' | 'global' - kind: str # 'fact' | 'preference' | 'action_pattern' | 'anti_pattern' - key: str # 结构化槽位（用户级精确匹配）；全局级可空 - content: str # 供注入 context 的自然语言 - embedding: list # 全局级语义召回用 - # 毕业机制 / 效用 所需元数据（先埋点） - hit_count: int - source_score: float # 来源轨迹 traj_score（写入门槛） - created_at: int - last_hit_at: int - graduated: bool # 是否已训进权重 → A→B 退休标记 -``` - -### MemoryStore 接口 -```python -class MemoryStore: - def write(item): ... - def retrieve(query, scope, *, intent=None, min_score=None, k=5) -> list[MemoryItem]: - # 用户级：KV/标签精确匹配；全局级：向量召回 + 标签过滤 - def mark_hit(item_id): ... # 命中打点（喂效用/晋升） - def retire(item_id): ... # A→B 退休（标记 graduated） -``` - ---- - -## 2. 核心目标：抽取器越来越专业（区别于 mem0/reme） - -### mem0 / reme 的做法与天花板 -- **mem0**：两步固定 LLM 流水 —— Extraction（固定 prompt 抽事实）+ Update（ADD/UPDATE/DELETE/NOOP 去重消歧）。 -- **reme**：分 personal / task memory，带 reflection 从成败轨迹提炼经验，检索时 rerank。 -- **共同天花板**：抽取器**静态**（extraction prompt 永远 day-1），**没有下游效用反馈回流到抽取器**，质量靠 LLM 当下自评。→ memory 只会“越攒越多”，不会“越来越专业”。 - -### 我们的突破口 -下游有**真实效用信号**（蒸馏是否受益、路由是否更准）+ 有 **verifier/rubric 打分**。→ 可以让抽取器被“经下游验证过的 memory”反过来训。 - -### 三个进化层次（由易到难） -1. **meta harness 学会“选对抽取器”**（最快见效）：contextual bandit，按轨迹特征分派抽取器；反馈=各抽取器产出 memory 的效用分。**纯统计即可**，不需训练。 -2. **抽取器“知道什么是好 memory”**：用 rubric 给 memory 打分，维度对准**可用性**（自包含 / 可泛化 / 可操作 / 不冗余），当抽取的软靶 + 写入门槛。 -3. **抽取器本身被蒸馏得更强**（最终形态，与主线同构）：用**下游效用**给抽取器输出打真标签（不是 rubric 自评），正/负样本微调抽取器（同 `llm_backup` 机制，被蒸馏的是“抽取器”这个角色）。 - -### 关键：rubric 分 = 训练靶，效用分 = 真值锚 -- **rubric 内在分**：即时、便宜 → 抽取软靶 + 写入门槛。 -- **下游效用分**：滞后、稀疏、客观 → 校准 rubric、训练抽取器的真值锚。 -- 可用效用分**反向校准 rubric**：某维度 rubric 高分但效用低 → 自动降权（复用 D7c“客观校准主观”）。 -- ⚠️ 坑：别让 rubric 自评当最终真值 —— “看起来专业 ≠ 用起来有用”，同 “模型不知道自己不知道”。 - -### 其它有效手段（业界/研究验证） -- **失败/反思挖掘**：失败轨迹信息量常更大（“这么调工具会报错”是高价值 anti_pattern）。见 §4。 -- **巩固/合并（consolidation）**：定期把多条相关碎片 memory 合并升华成更抽象的通则 → 更泛化，是 B→A 晋升头号候选。 -- **对比式抽取（contrastive）**：`llm_backup` 已收集的 (student 错 / teacher 对) 配对 → 专抽“teacher 做对而 student 做错的那步差异”，精准命中能力缺口。 -- **引用计数衰减**：长期不命中的 memory 降权/淘汰，防历史噪声拖累。 - ---- - -## 3. 毕业机制（B↔A 梯度；先设计不实现） - -memory 是权重的“预备队/退休区”，双线互相喂： -``` -生产轨迹 ─(清洗+评分)─┬─► line B: 抽成 memory → 立刻可检索（快） - │ │ - │ └─ 高频命中 & 已泛化 ──► 晋升为训练样本 - │ ↓ - └────────────────────────► line A: 攒批蒸馏进权重（慢） - ↓ - 权重已掌握 ──► 对应 memory 退休（stale/删） -``` - -- **B→A 晋升**（同时满足）：`hit_count ≥ N`；`kind==action_pattern` 或被判可泛化（排除实体孤例）；`source_score ≥ 阈值` 且 safety 通过。→ 还原成训练样本进 line A。 -- **A→B 退休**：一批训练部署后做**回归检验** —— 让新 student 在**不给该 memory** 时回答，若已能答对（verifier 自动判）→ 标记 `graduated`，移出检索池。 - - ⚠️ 依赖“新权重能否自答”的自动判定，正是“模型不知道自己不知道”；判定靠主线 verifier/置信度，**不能靠模型自评**，否则误退休。 -- 价值：普通 memory 系统缺“遗忘/毕业”机制会无限膨胀 + 过时；有 line A 就天然有毕业出口。 - ---- - -## 4. 失败挖掘与归因（最硬的一块） - -### 现状问题（必须改 pipeline） -按当前 pipeline，**失败轨迹会在打分前/后被丢，不会自动留存**： -- 退化性失败（死循环/复读/心跳/token soup）在 `DeadLoopFilter` 等**硬过滤前置步**就丢了（实测某片 `DeadLoopFilter: 25->2 dropped 23`），**没机会打分**。 -- 能力性失败若侥幸过硬过滤，会被尾部 `TrajectoryOutcomeFilter` 按 `low_traj_score` **删除**。 - -### 归因难题（用户戳中的核心） -一条带 memory 的失败，根因有三，且在最终轨迹上**长得一样**： -1. 模型能力问题（没 memory 也做不对）； -2. memory 缺失（有正确 memory 就能对，但没检索到/库里没有）； -3. **memory 误导**（检索到的 memory 错/过时/不相关，把模型带沟里）。 - -单看一条失败轨迹**无法区分** —— 缺反事实。若把 memory 误导当“能力缺口”抽成 anti_pattern → 坏 memory 的锅上再叠 memory → **无限污染循环**。这正是 mem0/reme 不敢做 memory 归因的原因。 - -### 破局：归因在“memory 使用现场”做，不在“挖掘阶段”猜 -把 memory 注入当成一次**可对照的干预**。对同一 query 跑带/不带 memory： - -| 不带 M | 带 M | 归因 | 动作 | -|---|---|---|---| -| 失败 | 失败 | model_capability | 抽 anti_pattern / 进训练线 | -| 失败 | 成功 | memory_helpful | M 加分 | -| 成功 | 失败 | **memory_harmful** | **退休该 M，不抽 anti_pattern** | -| 成功 | 成功 | memory_redundant | M 中性/可淘汰 | - -**归因是推理时记录的，不是挖掘时推断的。** - -### 剪枝（用户务实取舍） -- **token soup / 退化性失败**：根因是生成层退化，几乎与 memory 无关 → **直接丢，不留、不归因**（对照必然“两边都烂”，得不到 memory 信息）。 -- **归因预算（影子对照 + 留痕）只投在“能力性失败”**。 -- 判“垃圾 vs 值得归因”用**现成便宜过滤器**（DeadLoop/TokenSoup/SpecialChars 命中 = 垃圾直接丢）；判“模型问题 vs memory 问题”才用贵的影子对照。两层闸门，贵的只处理少数。 -- 放弃项（暂）：“反复调错工具”的死循环表面像 token soup 实为高价值 anti_pattern，但难自动区分、占比不高，先不捞。 - -### 落地（部署已确认：可做 A/B 影子跑；memory 注入尚未上线 → 从零把留痕设计进去） -注入路径三件套（第一天就装）： -1. **影子对照采样**：每次带 memory 推理以概率 `p`（约 5%）触发“不带 memory”的影子；两条都用 verifier 判成败，填 2×2 表。常规请求只留痕不对照。 -2. **注入留痕**（`user_data` 信封加字段）： - ``` - injected_memory_ids: [...] - memory_shadow: {ran: bool, without_mem_outcome: 0/1} - inference_outcome: 0/1 # verifier 判 - attribution: helpful|harmful|redundant|model_capability|unknown - ``` -3. **归因判定**：采样命中 → 按 2×2 表当场出结论；未命中 → 挂 `unknown`，靠历史反事实给**概率性**归因（低置信，仅用于排序“哪些 M 优先做对照验证”）。 - -### FailureMiner 的干净输入 -- 只从 `attribution == model_capability` 的失败抽 anti_pattern（已排除 memory 误导）。 -- `memory_harmful` → 触发删 M，不抽新 memory。 -- `unknown` 且历史高度怀疑 harmful → 排进影子对照队列验证，不直接抽。 -- 用 `round_scores` 定位**失败转折点**（哪轮突然掉），anti_pattern 抽那个点。 -- ⚠️ 准确率 > 覆盖率：一条错的 anti_pattern 会主动误导模型 → **宁可不抽，不错抽**；噪声/环境偶发失败不抽。 - ---- - -## 5. 下游效用分（memory 飞轮的燃料） - -### 定义（已确认：outcome 用主线 verifier；粒度 per-item） -效用 ≠ memory 内在质量，而是**边际因果贡献**： -``` -utility(M) ≈ E[ verifier(轨迹) | 注入M ] − E[ verifier(轨迹) | 不注入M ] -``` -`verifier(轨迹)` 用主线连续 `traj_score`（比 0/1 更细，能测“0.6→0.85”的边际改善）。它是 §4 归因的**连续版**（2×2 离散表的加权量）。 - -### 五种信号（弱→强） -- **A 反事实成功率差**（5% 影子对照）：因果最干净、正负都能测；稀疏 → **真值锚**。 -- **B teacher 追平差**（`llm_backup` 免费捎带，teacher 不带 memory）：student 带 M 后 match=True 是否上升；覆盖 llm_backup 采样流量。 -- **C 归因离散效用**（2×2 → +1/−1/0/0）：A 的离散版，够 bandit 奖励用。 -- **D 检索命中×结果**（全流量、便宜、**有选择偏差**）：不能单飞，只当排序候选，靠 A 去偏。 -- **E 多轮内即时行为信号**（`round_scores`：工具一次成功率、少走弯路、少重试）：细粒度，尤适 action_pattern。 - -### 组合（便宜信号打底 + 稀疏真值锚校准，同 D7c） -``` -效用分 = f(观测代理 D, teacher 追平 B, 行为 E) ── 用硬对照 A/C 校准去偏 -``` -1. 日常：累积 D+B+E 加权分（便宜、全覆盖、有偏）。 -2. 校准：5% 对照 A 产出无偏真值，回归；若某类代理分系统性高于真值 → 自动降权 D。 -3. 不确定性探索：代理与真值分歧大/方差大的 M → 提高其对照采样率。 - -### per-item 聚合器（唯一要新写的东西） -```python -MemoryUtility(memory_id): - n_hits; proxy_sum(按相关度加权); n_controlled; delta_sum - utility_hat ∈ [-1,1]; confidence(样本量+锚一致性); last_hit_at -``` -融合：对照足→用无偏 `delta_sum/n_controlled`；对照少→用 proxy 减去同类已知偏差；分歧/低置信→抬高对照采样率。 -per-extractor 分**不单独采**，从 per-item **聚合**上来（某抽取器所有 memory 的 `utility_hat` 均值），零成本供 meta harness。 - -### 驱动的动作 -- 退休/清坏：`utility_hat<0` 且置信够 → 删（memory_harmful 自动闸门）。 -- 晋升：`utility_hat` 高 + 高频 + 可泛化 → 进训练池。 -- 冷启动保护：新 M 低置信给保底曝光 + 时间衰减防老 M 霸榜。 - -### 三个坑 -1. 选择偏差（D 的病）：必须 A/B 去偏，D 不单飞。 -2. 信用分配：一条轨迹多条 M → 初期均摊或按检索相关度加权，别急上 Shapley。 -3. 反馈自强化：高效用被检索更多→分更高→更多… 可能锁死 → 时间衰减 + 强制探索低曝光。 - -### 前提 -地基是 **verifier 判 `traj_score` 的质量**。rubric 打分噪声大则整个飞轮歪 → **先坐实 verifier 打分可靠性，再上效用分**（与“修 TrajectoryScorer 打分区分度”同一条线）。 - ---- - -## 6. meta harness（管理抽取器；contextual bandit） - -### 定性：contextual bandit，不是全 RL；初期连 bandit 都先不上 -- 无状态转移（选抽取器不影响下一条外部流量）→ 不需要 Q-learning/PG。 -- 是 contextual bandit：轨迹特征=context，选抽取器=action，效用=reward。 -- **reward 延迟极大**（几天）→ 先做**离线统计的 bandit**（按 context 分桶统计各抽取器历史效用 + ε 探索），成熟后再升 LinUCB/Thompson。 - -### 状态（context，全复用 preprocessor 标签） -`intent` / `traj_score` 分桶 / `round_scores` 形状 / 轨迹长度·段数 / 是否 agent(`is_agent_row`) / domain。 -- 简版：离散化组合成桶 key，如 `(intent=tool_call, traj_hi, is_agent)`。 -- 升级：特征向量 + LinUCB。**初期分桶就够，别急上向量。** - -### 动作（两级） -- 一级（必做）：从 `{fact, action_pattern, preference, failure_mine, none}` 选一个/几个。**`none`（不抽任何 memory）是合法动作** —— 省成本 + 避免噪声。 -- 二级（成熟后）：抽取参数（粒度/数量/是否跨轮聚合），初期固定默认。 -- 约束：**动作空间要小**（个位数~十几个），否则永远冷启动。 - -### 奖励 -``` -reward(context, action) = agg{ utility_hat(M) : M 由该 action 产出 } -``` -- `agg` 用均值，可选覆盖惩罚（抽太多低效用条目扣分，鼓励精不鼓励多）。`action=none` reward=0 作基线。 -- **延迟处理**：批式/异步。抽取时记 `(context_bucket, action, [mem_ids])` 到待结算表；效用成熟后回填、更新桶统计。 -- 信用分配已在 per-item 层解决，bandit 直接用聚合值。 - -### 探索 -- 简版 ε-greedy，ε 随桶样本量衰减；新桶/新抽取器强制均匀试几次。 -- 升级 Thompson/UCB，用奖励不确定性驱动（与效用分“低置信多做对照”共用信号）。 -- **协同**：bandit 探索性选了冷门抽取器产出的 memory，效用最不确定 → 应**优先安排影子对照**去测它，否则探索白探。 - -### 骨架 -```python -MetaHarness: - policy: dict[context_bucket → dict[action → (reward_mean, n, confidence)]] # 存 StateBackend - select(traj_labels) -> action # ε-greedy over bucket - log_decision(context, action, mem_ids) # 待结算表 - settle(mem_id, utility) # 效用成熟后回填 policy -``` - ---- - -## 7. 稳定性：harness 会漂移，怎么关进笼子 - -“更新 harness 后不稳定”拆成三种病： -1. **策略震荡**：噪声+延迟反馈做了过自信更新 → 来回摆。 -2. **反馈自锁**：偏向某抽取器→只有它有反馈→别人永远没数据→越锁越死（off-policy 经典病）。 -3. **非平稳漂移**（本系统独有、最麻烦）：student 在被持续蒸馏、memory 库/流量在变 → 上周最优抽取器这周可能就不对。普通 bandit 假设平稳，这里不平稳。 - -### 对策：harness 从“决策者”降级为“建议者”，更新慢、可回滚 -1. **冻结基线 + 影子上线**（治 2、防炸）：新策略先只“建议”，实际按固定基线执行，记录“若听 harness 会怎样”；证据表明稳定优于基线才灰度切；永远保留基线兜底。 -2. **慢更新 + 迟滞带**（治 1）：桶样本 `≥ N` 才更新，否则用先验；只有显著且持续优于当前才切换（margin + 连续几批）。 -3. **强制探索地板**（治 2）：每个抽取器保底 `ε_min` 曝光永不归零，防自锁、且能先发现漂移。 -4. **滑动窗口/时间衰减**（治 3）：效用统计只用近期窗口/指数衰减，让 harness 跟随漂移重学。 - -### 更重要：先稳定可用，再进化 —— 分阶段把不稳定源头关掉 -- **阶段0 静态 harness（先上，零 bandit）**：固定人写路由表（`tool_call→action_summarizer`、`code→action+fact`、`低分失败→failure_mine`、`其他→fact`）。不学习不更新、完全确定。先把“多抽取器按类型分派”的**结构**跑通、攒效用数据、当 bandit 基线。**无任何不稳定性。** -- **阶段1 离线 bandit，仅离线复盘时更新**：攒够数据后离线做 off-policy 评估，只有证明稳定优于静态表，才把新策略**作为新静态表发布**上线。策略更新发生在**离线、可审、可回滚**节点，不在线漂移。 -- **阶段2（远期，可选）**：在线自适应。 - -### 分层可降级 mode（把“新方案风险”关进笼子） -``` -harness.mode = 'static' # 固定路由表，永远兜底 - | 'shadow' # bandit 只建议不执行，收集证据 - | 'canary' # bandit 接管 x% 流量 - | 'live' # bandit 全量（需离线验证达标才允许） -``` -出问题一键降回 `static`。 - -### 诚实边界 -- 阶段0 静态表需人工先验（哪个 intent 配哪个抽取器）—— 这不“自进化”，但是冷启动正确起点。 -- 非平稳漂移无完美解，滑动窗口只缓解；靠“离线定期复盘 + 可回滚发布”管理，不追求永远正确的在线策略。 - ---- - -## 8. 现有资产映射（实现时几乎不用造轮子） - -| 需要的能力 | 复用现有 | -|---|---| -| 存储（用户级 KV、策略、效用状态） | `twinkle/server/state`（memory/file/redis + TTL + update_atomic + keys） | -| embedding（全局向量召回） | `preprocessor/experimental/llm_backend.py::OpenAIBackend.embeddings` | -| 抽取器（fact/action/pattern） | `twinkle_agentic/summarizer/*`（都走 `llm_backup` 蒸馏，可 per-type LoRA）；`pattern_summarizer` 现为空 → 做“可复用模式”抽取，是 B→A 晋升头号候选 | -| 抽取器蒸馏 | `llm_backup`（被蒸馏对象换成“抽取器”角色） | -| outcome / 写入门槛 | 主线 verifier / rubric / hard_scorer 产出的 `traj_score / safety` | -| 转折点定位 | `round_scores`（TrajectoryScorer 已产出） | -| 校准逻辑（客观校主观） | D7c 思想直接搬 | -| 标签信封 | `preprocessor/label_schema.py` + `user_data`（PyArrow 稳定）；新增 `injected_memory_ids / memory_shadow / inference_outcome / attribution / memory_utility` | -| 配对（对比抽取 / teacher 对照） | `llm_backup` 已采 (student, teacher, match) | - -**唯一新写**：MemoryStore + VectorIndex + FailureMiner + per-item 效用聚合器 + MetaHarness（含 static/shadow/canary/live mode）。毕业机制先留接口 + docstring + 判据常量，逻辑 `NotImplementedError`。 - ---- - -## 9. 待定 / 下一步（未拍板） - -- 静态路由表的具体先验规则（intent → 抽取器映射）细化。 -- 待结算表 schema、桶设计粒度（先粗按 intent，数据多了再细）。 -- 检索/注入策略本身：检索几条、排序、注入到 context 哪个位置。 -- 双线去重边界：同一高分轨迹既进训练池又进 memory，B→A 晋升时如何避免重复训练。 -- memory 命中对置信度路由的**方向**：命中→更敢自答，还是命中→说明是薄弱区更该路由？（两种逻辑相反，需定。） -- 用户级 memory 的隐私/时效：覆盖写 vs 版本化；TTL。 - ---- - -## 10. Related Works（2026 检索，按本设计的轴归类） - -> 检索源：arXiv（2026-06 ~ 2026-07 为主）。结论：本设计的**每一条主要思路都能在近半年文献里找到平行工作或验证证据**——这是好事（方向被验证、不孤立），差异化在于**把这些点在一个自进化蒸馏框架里闭环组合**，且共用主线 verifier / `llm_backup` / D7c，而非各做各的。下面按“对应本文哪一节”组织。 - -### 10.1 双线（context-space + parameter-space）—— 对应 §0 - -**DuoMem: Dual-Space Distillation**（arXiv 2606.29961） -- 具体做法（分三步离线 + 一步在线，见其 Fig.2）： - 1. **teacher 造料**：用 Qwen2.5-72B teacher 对 3,553 个训练任务各跑 3–4 次（共 11,546 实例，5 次重试内累计成功率 99%），得到 11,434 条成功轨迹；**故意 oversample 同一任务的多条不同解**以增多样性、防过拟合。 - 2. **context-space 蒸馏（CD，训练无关）**：让 teacher（而非 student）对每条完成轨迹**离线生成 procedural memory 脚本**，存成文本 bank（整套任务几 MB）。推理时对新任务 d，用 `text-embedding-3-small` 算 d 与各条 memory 的 cosine，取 top-k **prepend** 进 student prompt。不改任何参数。 - 3. **parameter-space 蒸馏（LoRA）**：只用**成功 teacher 轨迹**微调 student 的 LoRA（rank 8–32，α/r=2，base 冻结）。 - 4. **组合**：先 LoRA 再 CD。ALFWorld 上 Qwen3-4B：No-Mem 4.3% → MemP(student 自产 memory) 55% → +CD 56.4% → +LoRA 72.1% → +DuoMem 77.9%（逼近 72B teacher 87.1%），只加 5.9M 参数、~12MB memory，且比 teacher 快 3×。**关键消融结论：CD 单独收益很小（+1.4），LoRA 才是大头（+17），两者组合还有超加性（>各自之和）。** -- 与我们的区别：DuoMem 的两轴都是**一次性离线固定**（teacher memory 生成一次、LoRA 训一次），**无在线闭环、无毕业机制、只用成功轨迹**。我们要：两轴在线持续、B↔A 毕业梯度、且**失败轨迹也要挖**。另外 DuoMem 的 CD 是"整任务级 memory 检索"，我们是 per-item 效用 + 归因过滤后的 memory。**它是我们双线最直接的可行性背书 + 起点基线**（甚至可先复现 DuoMem 当 line A/B 的 v0）。 - -**KbSD: Knowledge Boundary aware Self-Distillation**（arXiv 2606.29863） -- 想解决的问题：模型经常**不知道自己知不知道**——该答的时候瞎编（幻觉），不该答的时候硬答，或者明明该去查资料却凭记忆蒙。KbSD 想教会模型"划清知识边界"：**会的直接答、不确定的去检索、真不会的就说不会**。难点是普通 RL 只有一个"最后答对没答对"的稀疏奖励，没法教中间推理过程该怎么走。 - -- 怎么做（分三步，公式见原文 §3）： - - **① 给每个问题打三个"边界标注"**（这就是你问的"确定性/靠谱度怎么来的"）： - - **参数化确定性 μ(q)**：拿**冻结模型、不给检索**，对同一问题**独立采样 N 次**，算答对（match 标准答案）的比例 `μ = (1/N)Σ I[match(yᵢ, a)]`。μ 高 = 答案本就在模型脑子里。**注意：这一步需要 ground-truth 答案 a。** - - **语义稳定性 σ(q)**：这 N 次回答**两两之间的 embedding 余弦相似度求平均** `σ = mean cos(Enc(yᵢ), Enc(yⱼ))`。σ 高 = 每次都说同样的话（信念稳）；σ 低 = 每次瞎蒙都不一样。**这个不需要标准答案。** - - **检索质量 ρ̂(q)**：对检索返回的证据打一个 retrieval-quality 分（原文没细化打分器，是可替换的相关性模型/reranker）。 - - **映射到四象限**：把 μ、ρ̂ 各卡一个阈值二值化 → 得"内部知识可靠 k / 检索充分 s"，组合成四种该有的行为：都行→**Integrated（整合）**、只检索行→**External（靠检索）**、只内部行→**Internal（信自己）**、都不行→**Refusal（拒答）**。σ 不进象限，只用来筛训练集（已知区留稳定的、未知区留不稳的），让边界更干净。 - - **② 用"开小灶的自己"当老师做示范**：把上面标注 `(μ, σ, ρ̂, 目标象限)` 拼成一段 **hint 前缀，只喂给老师**；老师 = **同一个模型** conditioned 在 `[hint; q]` 上（且 stop-gradient，不回传老师）。学生 = 同一个模型但**看不到 hint**。因为老师多看了提示，能写出"知道分寸"的推理示范（该查就查、不会就认怂），学生就跟这个"开了天眼的自己"学——**这就是"信息不对称自蒸馏"，全程不需要更大的外部模型**。hint 只在训练用，推理时不给。 - - **③ 学生怎么对齐老师——是 KL 蒸馏，而且按象限切方向**（这就是你问的"用什么 KL"）：token 级两种损失，本质是 KL 的两个方向： - - **前向 KL（mass-covering，覆盖）**：在老师轨迹上最大化学生似然 → 学生尽量覆盖老师所有说法。 - - **反向 KL（mode-seeking，收敛）**：在**学生自己采样**的轨迹上匹配老师 → 学生收敛到老师主模式、抑制老师不支持的行为。 - - **分象限分配**：Integrated（只有一种正确整合、分布集中）→ **反向 KL**；Refusal（合理拒答说法很多、分布发散）→ **前向 KL**（别塌成一个模板）；External / Internal（要精准又抗噪）→ **前向+反向都要**，用 Pareto 加权自动解一个公共下降方向 α*（闭式解，不用手调系数）。这套蒸馏再和 GRPO 的稀疏 outcome 奖励**联合优化**（稀疏管"最终对不对"，稠密管"推理怎么走"）。 - -- 对我们有什么用：它正好治我们主线最头疼的病——**"模型不知道自己不知道"**，这直接决定置信度路由该不该把请求甩给 teacher；而且证明了**不用更大模型也能自我校准**（老师=学生+hint），与我们 `llm_backup` 的 student/teacher 同构。可直接借的两处：**训练置信度路由**、**A→B 退休判定**（退休本质就是判"新权重不给 memory 提示、自己能不能答对"，正好对应 μ）。 -- **⚠️ 搬过来要改的地方**：KbSD 的 μ **依赖 ground-truth 答案**，而我们主线是**无 ground-truth 的生产流量**。所以不能照搬 μ，要用**主线的 rubric/hard verifier 分数、或 teacher 一致性**来替代"match 标准答案"这一步；σ（自洽度）可原样复用，因为它本就不需要答案。 - -### 10.2 memory 越来越有用 / 自优化抽取器 —— 对应 §2、§6 - -**SelfMem: Self-Optimizing Memory**（arXiv 2607.03726） -- 想解决的问题：现有 memory 系统（MemGPT/Mem0/MemoryBank）都是**人写死的 memory 流程**——"存用户画像""到 context 上限就压缩"这种固定规矩，换个任务就不合适、还得手动调。SelfMem 想"授人以渔"：**不给死规矩，让 agent 自己摸索"这个任务下 memory 该怎么攒"。** -- 怎么做（关键：搞清 refine 的到底是什么）：世界里有三样东西—— - 1. **原始对话记录（transcript）**：存成 SQLite 表，**只读、永不改**，是事实唯一真相来源。 - 2. **memory 工作区（workspace）**：agent 自己维护的一块**可读可写白板**，装什么结构由它自己定（用户画像 / 偏好列表 / 项目笔记 / 时间线，甚至一段压缩策略文字）。 - 3. **一组固定的 memory 工具（四类）**：读 transcript（可跑只读 SQL）、读 workspace、写 workspace（加/替换/合并去重/精炼摘要/归档过时/记精确事实）、review（只诊断不改）。 - - 所谓 **"refine 自己的 memory 策略"，refine 的是白板里 memory 的内容与组织方式**（不是改模型权重，也不是改工具本身）：agent 跑一个 **inspect→write→review→revise（查→写→自查→修订）循环**——写完一条就用 review 工具自查"有没有过时/矛盾/没出处/难检索"，拿到诊断后回去改这块 memory。反馈是**多维不塌成单一标量**的（响应质量、token 数、成本、缓存命中），让 agent 在语言层面权衡"多存提升召回但涨成本、压缩省钱但丢细节"。全程**不动模型权重**（与 Reflexion/Self-Refine 同脉，靠语言反馈迭代）。BEAM 上 100K/500K/1M token 比最强基线 official score +0.165/0.141/0.134。 -- 与我们的区别：先厘清一个易混点——**工具集是固定的，agent 不能改工具**；它和"普通调工具"的区别在于**调哪个、什么顺序、写什么、要不要重写全由 agent 按反馈自己决定，没有预设 SOP**（MemGPT/Mem0 = 人写好 SOP、工具是执行的手；SelfMem = 人只给工具+评价、让 agent 自己长出 SOP）。而 SelfMem 优化的是"**这一个 agent 怎么攒/用它自己那块 workspace**"（都在 prompt/流程层，权重不动）；我们的 meta harness 优化的是"**用哪个抽取器/参数把轨迹变成 memory**"，反馈是**下游效用（Δverifier）**而非 agent 自评，且更进一步要把好 memory **反向蒸馏回抽取器权重**。相同的是"给工具+反馈让它自进化"这个哲学，它的"inspect→write→review→revise 循环 + 多维不塌缩反馈"可直接借进我们抽取器的自评环节。 - -**MetaSkill-Evolve: Two-Timescale Recursive Self-Improvement**（arXiv 2607.05297） -- **先说最关键的一句，破除误解**：这篇**完全没有训练、没有梯度、没有 loss**。从头到尾只有**一个冻结的模型**（Gemma-4 31B），"进化"全靠**让这个模型反复读写几个 Markdown 文本文件 + 拿准确率做进化搜索**。所谓"skill/meta-skill"就是几份 `SKILL.md` 文件，不是模型参数。所以你问的"怎么训练、loss 是什么"——答案是**不训练、没有 loss**，它是"改文件 + 挑最好的文件"的搜索过程。 - -- 几个"模型"其实是同一个冻结模型扮演的**五个角色**（靠不同 prompt 区分，各读一份对应的 `SKILL.md`）： - - **Analyzer**：看一条失败案例，诊断"为什么错"，打个标签。 - - **Retriever**：从别的分支里捞点"以前类似问题怎么改好的"当灵感。 - - **Allocator**：决定这一轮生几个候选改法（预算）。 - - **Proposer**：根据诊断，具体写出"skill 文件该怎么改"。 - - **Evolver**：把改动写进文件，并验证一下。 - -- 用什么数据 / 怎么"标注"：**不需要人工标注**。数据就是任务的 (输入, 参考答案) 样本，分成训练批和验证批。所谓"评分"是**自动的**——拿当前 skill 让 agent 去做验证批的题，**答对率就是这份 skill 的分数 `U(s)`**（`r∈[0,1]` 对着参考答案打分）。没有人在中间标任何东西。 - -- 读还是写：**主要是"写"侧的自我改进**——不断重写 skill 文件让 agent 做题做得更好；"读"只是 Retriever 去翻历史分支当灵感，不涉及检索优化。 - -- 两个时间尺度到底在进化什么（这是它唯一的新意）： - - **快环（每轮）**：改 **task skill `s`**（"这个任务该怎么做"的说明书）。拿当前 skill 在训练批上找错得最惨的一题 → 五角色流水诊断+提改法 → 生成几个候选新 skill → 谁在验证批上**答对率更高**就留下（只有严格变好 `ΔU>0` 的才进池子）。 - - **慢环（每 H 轮）**：改 **meta-skill `m`**，也就是"**上面那五个角色自己的说明书**"。关键点：因为五个角色的说明书也是同格式的 `SKILL.md`，所以**用同一套五角色流水去改它们自己**（自己改进自己的改进方法）——这就是"递归自我改进"。判据是 **meta-productivity `P(m|s)`**：这个分支最近 H 个"孩子"平均提升了多少（`= 子代 ΔU 的均值`），衡量"这套改进方法还灵不灵"。 - - **选哪个分支继续进化**：打分 `η₁·当前分 U + η₂·这套方法的产出率 P + η₃·新颖度 N`（N 惩罚被选太多的分支，逼它去探索没试过的路线）。 - -- 效果：OfficeQA/SealQA/ALFWorld 比 No-Skill +23.5/+16.1/+1.9，比"只进化 s、不进化改进方法"的版本 +6.4/+8.1/+1.9（证明"连改进方法本身也进化"确实有额外收益）。 - -- 与我们的区别：它的思想（"**快改内容、慢改'改内容的方法'本身**"）正是我们 §6/§7 想要的——我们的 static→shadow→canary→live 就是它慢环的工程化 + 加了可回滚门控。**但它全程改的是文本文件、模型冻结、没有训练**；我们的慢环最终要落到**真的蒸馏进抽取器权重**（把好经验烧进参数，而非只改说明书），且用**下游效用（Δverifier）**而非"验证批答对率"当信号。一句话：**它是"用一个冻结模型玩进化搜索改 prompt"，我们要的是"把搜出来的好东西训进权重"。** - -**COMFYCLAW: Self-Evolving Skill Harnesses**（arXiv 2607.01709） -- 领域 / 模型：图像生成工作流（ComfyUI）。**不训练模型**——用现成 LLM 当 agent、现成 VLM 当"验收员"，改的是外部的 skill 文件库。 -- 怎么做（一个"边做边攒经验"的闭环，见其 Fig.1）： - 1. 给一个文生图需求，agent 通过**带类型的图编辑**（连节点、调参、加 LoRA）把 ComfyUI 工作流搭起来，跑出一张图。**非法的编辑会被自动撤销**（防止把流程改坏）。 - 2. **VLM 验收员**把需求拆成一串"可观察的是非题"（比如"有没有三只手臂""风格对不对"），逐条判过没过 + 给个 0–10 细节分，合成一个标量分数；并把**没过的条目 + 哪里错了 + 具体该怎么改**回吐给 agent，指导下一轮修改。 - 3. 跨很多需求跑下来，把"反复成功/失败的经验"提炼成**可复用的 Agent Skill（`SKILL.md` 文件）**，存进 skill 库，下次相关需求时**先只给 skill 摘要、需要时再展开全文**（渐进披露，省 context）。 -- 数据 / 标注 / loss：**没有训练、没有 loss**，"分数"来自 VLM 验收员的是非题（自动，无人标）。四个 split×两 backbone×三模型下，比"只有验收员、不进化 skill"的基线高 4 分、比"完全不修改"高 10 分。 -- 与我们的区别：它名字和"harness 管理可复用技能"的思路跟我们碎片1直接撞上，但它是**图像生成域实证**、skill 是"给 agent 复用的操作技能"；我们的"skill/抽取器"是**把轨迹变 memory 的工具**、且最终要蒸进权重。它的两个工程点可直接抄：**非法编辑自动回滚**、**验收反馈翻译成"可执行的修改建议"而非只给一个分**。 - -**UCOB: Credit-Aware On-Policy Bidirectional Self-Distillation**（arXiv 2606.29502）—— **与我们 §4 归因 + §5 效用最像的一篇，务必读透** -- 想解决的问题：检索来的"经验/skill"**不是万能的**——同一个模型，在情形 A 被这条经验帮到、在情形 B 反被它带沟里。所以"把'带经验的回答'当成永远正确的老师去教'不带经验的回答'"这个假设是**脆的、会把坏经验也学进去**。这正是我们担心的"分不清是模型问题还是 memory 问题"的学术版。 -- 怎么做（大白话，这是**真·训练，有 RL loss**）： - 1. **同一个在线模型**，同一道题准备**两种输入**：带经验的（`P₊`）和不带经验的（`P₀`）。每道题各采一批 rollout，一半用 `P₊`、一半用 `P₀`，**两边都参与在线 RL 更新**。 - 2. **在相同的"局面"上比谁做得好**：把两边 rollout 里**走到同一个中间状态**（论文叫 anchor-state）的记录凑成一组，各自算"从这一步往后的总回报"（return-to-go = 后续奖励的折扣和）。定义**同局面下的差值**`Δ = 带经验的平均回报 − 不带经验的平均回报`。 - 3. **谁赢谁当老师，只在这一局面上教对方**：`Δ` 明显为正 → 说明这条经验在这儿确实有用，让"带经验的回答"去教"不带经验的"（把能力吸收进去，以后不给经验也会）；`Δ` 明显为负 → 说明这条经验在这儿是**误导**，反过来让"不带经验的回答"**纠正**"带经验的"（**主动压制坏经验**）。教的方式是 token 级分布对齐 + 置信度门控（只在有把握的位置教）。 - 4. 同一个 `Δ` 还顺便用来**更新每条经验的效用分**（配 UCB 决定以后检不检索它）、并训练"写经验"的 reflection 模块。ALFWorld/WebShop 比 SOTA +23.5/+18.0。 -- 与我们的区别：UCOB 的"带经验 vs 不带经验、比谁回报高"就是我们 §4 **影子对照的在线 RL 内生版**——它在 rollout 内、按 anchor-state 配对做；我们在**推理服务侧按 ~5% 采样**做 2×3 归因表。UCOB 的产物是**改 policy 权重 + 更新经验效用**；我们把同一个"带/不带 memory 对照"信号接到**三个出口**：失败挖掘、per-item 效用、坏 memory 退休。**它强证了我们方案的核心机制可行且高收益**，它的 `Δ` 就是我们 per-item 效用的一种无偏估计，"按相同中间状态配对"这招可直接借来**降低我们效用估计的方差**。 - -### 10.3 失败挖掘 / 从成败双向抽 memory —— 对应 §4 - -**Learning from Failure: Inference-Time Self-Improvement for Computer-Use Agents**（arXiv 2606.31270，**ECCV 2026**） -- 想解决的问题：现在造 agent 训练数据的标准做法是"agent 在有验证器的环境里跑 → **只留成功轨迹**去微调 → 丢掉所有失败"。但失败其实携带了"模型哪里弱"的宝贵信息，全扔了很浪费。 -- 怎么做（**不训练**，改的是 agent 的推理时行为，见其 Fig.2）： - 1. agent 跑一批任务，收集**失败**轨迹。 - 2. 用一个 LLM 当"分析员"，把 (指令、动作历史、思维链) 喂进去**诊断失败原因**，归成**四类**：定位不准（grounding）、能力缺口（该用某工具却不用）、知识缺失、无脑重复循环。 - 3. 对每一类，LLM **提出一个推理时的补救办法并生成一段代码补丁**（分别对应：加视觉搜索、允许走终端执行、注入知识、加重复告警），**人工轻量核对**这段补丁后，并进 agent 的工作流，再重跑。每轮按"当前最主要的失败类型"选一个补丁，补丁跨轮累积。 -- 数据 / loss：**零训练、无 loss**，纯粹是"诊断失败→打补丁→重测"的循环。OpenCUA-72B 在 OSWorld 从 42.3%→48.9%（全部补丁叠加 52.74%）。 -- 与我们的区别：**直接印证我们"别丢失败轨迹"**。但它把失败变成**给 agent 的代码补丁/工具**（改 harness），而且**完全不区分失败是模型本身弱还是 memory 带偏的**；我们把失败变成 **anti-pattern memory + 训练负样本**，且**先做归因**（排除"是 memory 误导"才抽经验）。它那**四类失败诊断**可直接拿来当我们失败挖掘器的分类初值。 - -**ISM: Self-Improving Strategy Memory for Continual Math Reasoning**（arXiv 2606.31191，**ICML 2026 AI4Math Workshop**）—— **毕业/退休机制的现成七件套模板** -- 想解决的问题：冻结的 LLM 做连续不同领域的数学题时，学到的经验存哪、怎么不越堆越乱？纯 retrieval 会无限膨胀、纯 reflection 只存散乱文字。 -- 怎么做（**不训练模型**，只维护一个外部"策略库"）： - 1. 外挂一个**紧凑的 strategy-schema 库**。每条 schema 拆成两半：**content**（策略描述/解题模板/启发式，用时注入 prompt）+ **feature hook**（结构标签 + embedding，决定"什么时候该检索到它"，且随使用自动微调）——把"这条经验讲什么"和"何时被调出来"解耦。 - 2. 检索分两步：先按题型/算子过滤，再 soft 打分选最相关的。 - 3. **库由七个自维护机制打理**（这是精华）：①audit 审查 ②merge 合并近重复 ③prune 删无用 ④promote/demote 升降级——管质量和体积；⑤reinforce 从**成功**里抽正向启发式 ⑥antipattern 从**失败**里记"要避开的坑"——**成败双向都学**；⑦rehabilitate 给表现差的 schema 一次翻身机会再决定删不删。**每次改库都要先过符号验证器**（数学能硬校验），防止把错误泛化写进去。 -- 数据 / 标注：数据是 300 题的连续流（按域分块），"对不对"由**符号验证器自动判**（数学域独有的硬 verifier），无人工标注。backbone 是 gpt-4.1-mini、temperature=0。 -- **准确率数字要拆开看（别被 0.48→0.81 唬到）**（Table 1，MATH-Hard 累计 acc）： - - Vanilla 48.0 → RAG 57.0 / Reflexion 55.7 → **Static Schema 78.67**（bank=1，一个固定 prompt 模板 + 允许调符号工具）→ Passive 78.67 → ISM 80.67。 - - **+30 点的大头来自 Static Schema——即"好 prompt 模板 + 符号验证工具"，与 memory 无关**；真正属于"memory 自维护机制"的（ISM vs Passive）**只有 +2 点**（300 题净多对 6 道，OlympiadBench 同样 +2）。作者自承单 seed、单 stream 顺序、无逐机制消融，"gains should be interpreted as preliminary"。 - - ISM 的真卖点其实是**省存储 + 抗遗忘**（bank 只有 Passive 的 1/3~1/7、比 RAG 少 23×），不是"memory 让数学变强"。 -- **⚠️ 评测可信度存疑（重要）**：全文**没有任何去污染措施**——grep 全文无 `decontamination / n-gram / 13-gram / dedup / held-out / train-test split`。而且它的"记忆积累"和"评测"用的**是同一条 300 题流、没有独立 held-out**：RAG 基线是"**把做过的每道题（含解）全存进库、按 embedding 召回最像的一道注入**"，对 MATH/OlympiadBench 这种"换数字的同型题很常见"的数据，等于系统性地**召回近似题的解 = 流内信息泄漏**，且未做任何近似题过滤。所以 RAG 的 0.57、schema 系的 0.79 都**掺了泄漏水分 + prompt/工具增益**，**这些数字不能当作"memory 在数学上有效"的证据**。 -- 与我们的区别：ISM 的**七机制 + 验证门 + content/feature 解耦**几乎就是我们"毕业（promote/demote/retire）+ 效用退休 + 去重合并"的**现成设计模板**，它的"成功→正样本、失败→anti-pattern 对称双向抽取"正是我们要的——**我们借的是这套领域无关的 lifecycle 机制**。但**不采信它的数学有效性结论**：它靠数学专有的硬符号验证器、且评测有泄漏/缺干净对照；我们无此硬 verifier（只有带噪声的 rubric+hard 软 verifier）。而且它这套结论**恰好与我们自己的实验吻合**——我们 200 题数学 memory 只比 direct 多对 1~2 道，而 ISM 剥掉 prompt/工具/泄漏后 memory 机制也就 +2 点：**两边共同印证"数学这类可泛化硬技能应走线 A（训权重），memory（线 B）边际只有 1~2 个点"**。ISM 还是权重全程冻结、纯外部 memory；我们是双线，退休判据还多一条"新权重不给 memory 能不能自答"。 - -**M2Note: Mistake Notebook Learning**（arXiv 2607.00685） -- 想解决的问题：怎么把失败经验安全地攒成"错题本"，又不会因为写错东西把整体带崩。 -- 怎么做（不训练，改外部笔记）：把失败轨迹提炼成**按主题组织的"错题本"note**，检索注入引导 agent 规避同类坑；关键工程点是 **批级后验 + 回滚**——一批笔记编辑先在同一批任务上验证，**只有整批指标真的提升了才提交，否则整批回滚**。支持同模型自我进化，也支持"一个模型的错题本给另一个模型用"（= 我们的 student/teacher）。 -- 与我们的区别：它的"批级后验 + 只在变好时提交、否则回滚"与我们 D7c 校准 / harness 发布门的"批级回归检验 + 达标才发布"**几乎一样**，是 §7 稳定性的又一独立佐证；"跨模型进化"正对应我们 teacher→student 的 memory 迁移。差异仍是我们把它接进**归因 + 双线 + 效用退休**的完整闭环，而非只做注入引导。 - -### 10.4 归因 / 去偏 / "何时不该写 memory" —— 对应 §4 归因难题、§5 选择偏差 - -**GovMem: When Not to Write Memory — Governing False Promotion from Correlated Traces**（arXiv 2607.02579，**MLISE 2026**）—— **精确命中"分不清是模型还是 memory 问题 + 会污染"** -- 想解决的问题（一句话）：**"重复出现"不等于"多份独立证据"**。五个 agent 说同一句，可能是五次独立发现，也可能是**同一条过时笔记经共享上下文回声了五遍**。如果按"出现次数多就晋升"的朴素规则，长期 memory 会慢慢变成"把相关性错误固化下来的持久层"。 -- 怎么做（**不训练**，是一个"该不该写这条 memory"的审计决策，四步 write-path，见其 Fig.1）： - 1. 把说法相近的观测**聚成一条候选 memory**，同时保留每条的**来源信息**（哪个 source、哪套 prompt、哪个父事件、什么环境、可信度）。 - 2. **算"去掉相关性后的有效支持度"**：共享同一 prompt/工具/父事件的观测**不算独立的一票**（这是核心——把"回声"折价）。 - 3. 检索**反证**、检查"这条经验声称适用的范围"合不合理。 - 4. 综合上面判断，输出三选一：**promote（写）/ reject（拒）/ needs-review（送人审）**。 -- 数据 / 效果：合成压力测试里把"错误晋升率"从 0.597（按出现次数）降到 **0.040**，同时召回还保 0.960，代价是 15% 送人审。**最刺眼的是真实数据结果**：133 条高影响候选经人裁后，**0/133 可以安全自动晋升，本地门控判过的 11 条全被人否掉**。 -- 与我们的区别：GovMem 是个**保守的审计侧策略**，只管"写不写"、不管"写了有没有用"。我们把它当**"晋升到 A 线 / 长期 memory 之前"的闸门**，再叠上我们的 verifier 地基 + 灰度人审；它的四步（留来源→按相关性折价→找反证→三选一）可以直接做我们 promote 的前置检查。它"几乎没有能安全自动晋升的"这个结论，**强烈警示我们：宁可漏抽也别错抽，默认走 needs-review 灰度**。 - -**MemDelta: Controlled Baselines and Hidden Confounds in Agent Memory Evaluation**（arXiv 2606.29914） -- 想解决的问题：报告里 memory 系统"比 RAG 强"的结论，常常混进了 LLM/embedding/检索管线本身的变化——到底是**memory 架构真强**，还是只是**换了个更好的 embedding**？ -- 怎么做（**一篇测量方法论的论文，不提新架构**）：在 LongMemEval-S（500 题、每人 50+ 会话、三个模型族）上，**一次只变一个变量、其余全冻**，隔离四个隐藏混淆：检索质量、embedding 选择、模型长上下文行为、写路径成本。 -- 关键实测（这些数字很能说明问题）：① **agent 自产 memory 只有 42%，反而不如朴素 retrieval**；② 只换 embedding 不动别的，Mem0 就从"比 RAG 基线 +11pp"翻转成"−1.2pp"——**结论被一个变量掀翻**；③ Mem0 只在窄题型上占优，但**写路径成本可占 agent 总执行时间 80%+**。建议：固定 embedding、按模型族分层、把 write 成本当一等公民报告。 -- 与我们的区别：**直接给我们 §5 选择偏差 + §4"必须做对照"背书**——它证明了"不做受控对照，memory 的效用数字根本不可信，甚至自产 memory 是负收益"。这正是我们坚持"per-item 效用必须用 **带/不带 memory 反事实对照**去偏、且先把 verifier 地基坐实再上效用分"的理由；它的"固定 embedding + 报告 write 成本"应直接写进我们的效用评估协议。 - -**Stealthy Memory Injection in Persistent Personal Agents**（arXiv 2607.05189） -- 想解决的问题 / 做法：展示在持久化个人 agent 里，**坏的或恶意的 memory 能被悄悄写入并长期潜伏**——用户看不见、还跨会话反复生效造成危害。 -- 与我们的区别：佐证我们需要 `memory_harmful` **自动退休闸门** + 用户级 memory 的隐私/时效治理（§9 待定项）——尤其我们有用户级 scope，更要防"一条坏偏好被写进去后反复检索注入"。 - -### 10.5 memory 命中的"选择性使用" + memory-on/off 对照 —— 对应 §4 影子对照、§9 路由方向 - -**ATMem + STR-GRPO: What Memory Do GUI Agents Really Need?**（arXiv 2606.31612）—— **STR-GRPO 就是我们"影子对照"的 RL 化，几乎逐条对应** -- 想解决的问题：GUI agent 做长任务时，光"把过去看到的存下来"不够——它还得知道"这条信息现在**该不该用、用了到底有没有帮助**"。 -- 怎么做（两部分，见其 Fig.2；这是**真训练**：先 SFT 再 RL）： - 1. **ATMem（把 memory 从被动存储变成主动的"执行状态"）**：memory 不是流水账，而是一张**结构化的任务进度表**——记着"整体进度 + 约束""每个待办项的内容 + 它的状态（待办 / 已完成 / 跳过）"，由 agent 边做边更新。先用"只保留通过验证器的成功轨迹"造 SFT 数据（120 模板 → 1.1K 实例 → 21,713 步级样本），教会模型**会建、会更新、会引用**这张进度表。 - 2. **STR-GRPO（用对照实验学"何时该用 memory"）**：对同一道题采一批 rollout，**刻意一半开 memory、一半关 memory**（只切 memory 这一个开关，其余历史都保留）。奖励 = 最终有没有做成（验证器给 0/1）**减去 memory 的使用成本**（用了 memory 却多走步、没帮上忙，就扣分）。因为同一题的开/关两组**共用同一个打分基准**，所以"开 memory 比关 memory 好多少"就**直接量化成了这条 memory 通道的净贡献**，模型据此学会"该用才用"。 -- 与我们的区别：它"同题下开 memory vs 关 memory、比谁做得好"**就是我们 §4 影子对照 2×2 表的 RL 内生版**，它的"memory 使用成本"惩罚正好回答我们 §9 待定的"何时该用 memory、何时 memory 反而是负担"。差异：它在 **GUI 域、RL rollout 内**做、产物是改 policy；我们在**推理服务侧按采样**做、产物接归因+效用+退休三出口，且走双线蒸馏而非纯 RL。**它"开/关对照算净贡献"的思路可直接当我们 per-item 效用的估计量。** - -**WorldEvolver: Self-Evolving World Models for LLM Agent Planning**（arXiv 2606.30639） -- 想解决的问题：给 agent 装一个"世界模型"（行动前先预测后果）能帮规划，但**预测不准时反而会把 agent 带偏**。而每次都靠梯度更新去修这个世界模型，在线部署下太贵、还会灾难性遗忘。 -- 怎么做（**关键：agent 和世界模型的参数全程冻结，只改外部 memory**）：三个模块——① **情节记忆**：把真实发生过的"状态→动作→结果"存下来，检索出来做"检索式模拟"；② **语义记忆**：把"预测的结果 vs 真实观测"对不上的地方，提炼成可复用的启发式规则；③ **选择性前瞻**：预测在喂给 agent 之前，**先把低置信度的预测过滤掉**，只把有把握的预测注入。 -- 与我们的区别："**没把握就别注入**"正是我们 §5 不确定性驱动 + 主线置信度门控的同款直觉（我们把它用在"低效用/高风险 memory 不注入"上）；"从预测-真实的失配里提炼规则"对应我们从失败挖 anti-pattern。差异：它**冻结参数、只改 memory**；我们要把提炼出的规则进一步走 B→A 蒸馏进权重。 - -### 10.6 记忆巩固 / 晋升与身份稳定 —— 对应 §3 毕业机制、§7 稳定性 - -**Episodic-to-Semantic Consolidation Without Identity Drift**（arXiv 2607.01988） -- 想解决的问题（一个偏"合规/审计"的场景）：受监管的长期部署 agent（医院、工厂机器人）有一个**加密认证的身份**（对一份 manifest 做哈希）。传统"巩固知识"的做法（微调 / 改 prompt / 蒸馏 / 追加反思）都会改动定义身份的那份东西，于是**每学一点新知识就等于换了个 agent、要重新认证**。矛盾：既要越用越聪明，又要身份字节级不变。 -- 怎么做（**v1 完全不用 LLM、是确定性统计规则**）：把"巩固"定义成一个**确定性函数** `f: 情节日志 → 语义层`。情节日志是只追加的原始事件记录；`f` 就是**按 (技能+对象+场景) 分组、数成败、算成功率**，输出一行带 **置信度 + 观测数 + 溯源指针** 的语义事实（例："对玻璃杯、这个环境，建议抓取力 25N，置信 0.83，基于 15 次观测"）。关键设计：**身份哈希在构造上就不读这个语义层** → 无论巩固多少次，身份字节不变；planner 只能**只读**查询语义层、不能改。（用 LLM 的 v2 被明确列为 future work，因为引入不确定性会破坏可审计性。） -- 数据 / loss：**无训练、无 loss**，纯确定性聚合。合成 benchmark（1000 决策）上，对比一个校准过的 Bayesian 基线，planner 的无效尝试降 79.82%，同时身份哈希全程字节相等。 -- 与我们的区别：`f: 情节→语义` 正对应我们 **B（快 memory，情节性）→ A（长期语义）** 的毕业方向，"每条带置信度+溯源、可审计"对应我们的血缘记录；"学新知识不改身份"提示我们**毕业/退休不该破坏模型稳定的基础能力/人格**（呼应 §7）。差异：它只在 memory 层做确定性聚合、**完全不碰权重**、且是单向（情节→语义）；我们的 B→A 是**真的权重蒸馏**，还有 A→B 反向退休。 - -**SEA: Self-Evolving Agents with Anytime-Valid Certificates**（arXiv 2607.00871）—— **几乎就是我们 §7"把进化关进可回滚笼子"的统计化理论** -- 想解决的问题：自进化 agent 有个根本麻烦——它用来学习的数据、评判自己的评估器、用的组件，**全是被它自己更新的策略生产出来的**（自己考自己、自己出题自己判）。这种"闭环自产"下，经典学习理论的保证（收敛、不遗忘、安全改进）**全部失效**。 -- 怎么做（四层架构 + 两条铁律，见其 Fig.1；注意它**主要不做权重微调**）： - - **四层**：L0 = 冻结的底座模型；L1 = 一个很小的 **steering adapter**（在线只调"选哪条指令"的概率分布，不动权重）；L2 = **带版本的 harness**（prompt/工具/预算/技能库，可改可扩，每次改都记一个新版本）；L3 = 在旁边的调度器（不在推理主路径上）。 - - **铁律一：每次自我修改都要过一道"随时有效"的统计门**。因为 agent 每轮都在偷看自己的成绩，普通固定样本量的显著性检验会失效；它改用一种**允许你随时停下来看、结论都成立**的统计量（e-value），对一个固定的"错误预算"发一张**可审计的证书**（通过/暂缓/拒绝/无解），全部记进一个账本。快环每轮调 L1、慢环每隔 K 轮改 L2。 - - **铁律二：门只能在"底座本来就能做出的行为"里挑**（不能凭空造新能力）。所以另配五个"验证器在环"的引擎（best-of-N、微步搜索、自写复现测试、搜索层控制、自修复）来**产生候选行为 + 提供密集的、不靠人打分的信号**。SWE-bench Verified 上 +4/+5。 -- 与我们的区别：SEA 用**统计证书 + 错误预算**保证每次进化不把系统带崩；我们用**工程化的发布分层（static→shadow→canary→live + 可回滚）**做同一件事——两者可互补（我们的 canary 达标判定可以升级成它这种"随时有效的统计门"）。它的"L0 冻结 / L1 只 steer / L2 可改 harness"分层，给我们"哪些能在线漂移、哪些必须冻结"划了清晰边界；"门只能挑已有行为"正是我们要把 harness 关进笼子的理由。差异：SEA **基本不微调权重**（只 steer L1）；我们主线恰恰要蒸馏权重，所以更需要它这套门控来兜底。 - -### 10.7 rubric 作为 reward / 可靠性 —— 对应 §2 rubric 靶、§5 verifier 地基 - -**RuVerBench: Can LLM-as-a-Judge Reliably Verify Rubrics in Agentic Scenarios?**（arXiv 2606.29920） -- 想解决的问题：现在大家用 LLM 当"裁判"、按 rubric（评分条目）给 agent 打分，但**这个裁判本身靠不靠谱**没人系统测过——尤其 agent 输出又长又复杂（深研报告几千 token、编码轨迹几万 token）时。 -- 怎么做（**一个 benchmark，不训练**）：构造 2458 条样本（深研 1615 + 编码 843），每条 = (一段 agent 输出, 一条 rubric, 人工标的"满足没满足")。人标经**双人独立标注 + 裁决**，两组一致率 90.4%、κ=0.808（很高）。然后拿各种前沿模型当裁判去判，测它们和人标的吻合度，并测"多判几次投票""一次判多条"这些策略。 -- 结论：**即便最强模型判 rubric 也有明显噪声**；**编码类、尤其涉及 tool-use 的 rubric 判得最差**；多数投票有效但**收益递减**；一次判多条省钱但掉准确率。 -- 与我们的区别：直接支撑我们的**执行顺序**——"**先坐实 verifier 打分可靠，再拿它去算 memory 效用分**"，否则效用分建在噪声地基上（和 MemDelta 的警告叠加）。它"编码/tool-use 类判得最差"正戳中我们 agentic 场景，它的投票收益递减曲线能帮我们定"什么时候值得多投几票"。 - -**MRRG: Many Voices, One Reward — Multi-Role Rubric Generation**（arXiv 2607.01830） -- 想解决的问题：现在"自动生成 rubric"多是**一个通用评估器一口气列所有标准**，容易漏维度（论文叫"维度盲点"），进而导致判分看领域、还能被"只优化被覆盖的标准"刷分（rubric hacking）。 -- 怎么做（**训练无关、无需参考答案**）：让同一个 LLM **轮流扮演多个角色**（用户、领域专家、教育者、AI 研究员、语言学家…），每个角色从自己视角产一批**原子、可验证**的 rubric 条目，再**汇总去重**成一个可审计的打分器。这个打分器既能做偏好判定，也能**直接当 GRPO 类 RLVR 的 reward**。 -- 效果：RewardBench-2 / JudgeBench / PPE 上比单角色基线 +3.1~16.4pp；用作 RLVR reward 时 +1.7 / +3.4。 -- 与我们的区别：可直接增强我们 `RubricVerifier` 的 **rubric 生成质量**——把当前"单 prompt 生成 rubric"升级成"**多角色生成再合并**"，且它天然兼容我们主线 verifier（同一套 verifier 既供 memory 效用、也供训练奖励）。 - -### 10.8 memory 系统工程形态（存储/检索）—— 对应 §1 选型 - -**MOSS: Auditable Agentic Memory**（arXiv 2607.04391） -- 想解决的问题：主流 RAG 用向量相似度检索，**不透明、难审计、有理论上限**——在长期/个人/受监管场景尤其致命。 -- 怎么做：由 **agent 分析查询意图 → 参数化一条结构化检索 → 在关系库上用 SQL 确定性地取数**，**检索环里没有 LLM**（一旦查询定好，执行完全可复现）；词表从语料自动归纳、不外加本体；每一步从建索引到出答案全部可审计。已**真实生产部署约一年**（约 4400 万 token 语料、每天当主力工作记忆用）。 -- 与我们的区别：**印证我们"检索用结构化 KV + 可审计、LLM 只在写/抽 memory 时参与"** 的选型——检索热路径不放 LLM，省成本、可复现、可审计。 - -**Mandol: Agglomerative Agent Memory**（arXiv 2606.29778） -- 想解决的问题：现有系统把向量库、图库拆成好几套，**跨库 I/O 慢、信息碎片化**，RAG 式检索又容易招噪声、漏关联、控不住 token 预算。 -- 怎么做：用一套 **SemanticMap + SemanticGraph 的内存数据结构，原生融合 KV / 向量 / 图**（不是拼三套系统），提供统一的混合检索算子；检索走"查询自适应路由 → 去噪/消解冲突 → 按 token 预算生成上下文"，**全程不调 LLM**。LoCoMo 92.21% / LongMemEval 88.40%（均最优），10 QPS 下检索延迟比最快基线还低 **5.4×**。 -- 与我们的区别：直接支持我们"**混合索引（向量 + 结构化 KV）+ 应用层检索、检索不放 LLM**"的方向，且证明融合式索引比拼装更快更准——是我们 §1 存储层的可参考实现形态。 - -**HyphaeDB**（arXiv 2606.28781） -- 怎么做：把 HNSW 近邻图当作**多 agent 之间的知识传播网**（gossip 扩散 + 能量衰减 + 自发共识，让高价值知识自然传开、陈旧的自然淡出），并给了 **pgvector 参考实现**。 -- 与我们的区别：现在用不上（我们先做单机混合索引），但将来若**全局向量层要升级成多副本/多 agent 协同**，它的"能量衰减 = 自动淘汰陈旧 memory"和我们退休机制思路一致，可作远期参考。 - -**其它形态（提示 memory 不止"注入文本"一条路）** -- **PLACEMEM**（2607.04089）：按算力预算调度 memory 平面。 -- **Neural Procedural Memory**（2606.29824）：用**隐式 activation steering**（直接改模型激活，而不是往 prompt 里拼文本）来承载程序性 memory。 -- **Analytic Concept-Centric Memory**（2606.29774）：以概念为中心组织 memory。 -- 与我们的区别：后两者提示"注入 memory"未必只有"prepend 文本"这一种——**改激活（activation steering）是一条可选的补充通道**（和 SEA 的 L1 steering adapter 呼应）。我们当前走文本注入，把它列为远期 B 线的可选实现。 - -### 10.9 对比表：本设计 vs 代表性工作 - -| 维度 | 本设计（twinkle） | 最接近的工作 | 我们的差异 | -|---|---|---|---| -| 双线（memory + 权重） | context 注入 + `llm_backup` 蒸馏，且有 B↔A 毕业梯度 | DuoMem：CD(teacher memory 检索 prepend)+LoRA(成功轨迹) | DuoMem 两轴**离线一次性、只用成功轨迹**；我们在线闭环 + 毕业 + 失败挖掘 | -| 抽取器自进化 | meta harness（bandit）+ 效用回流 + 抽取器可蒸馏 | MetaSkill(两时间尺度五 agent)/ SelfMem / COMFYCLAW | 它们在 prompt/skill 层递归、权重冻结；我们慢环落到**权重蒸馏** + 分层可回滚笼子 | -| 失败挖掘 | 只挖“能力性失败”，token soup 直接丢 | Learning-from-Failure(四类诊断→patch)/ ISM(七机制)/ M2Note | 它们不区分模型/memory；我们前置**归因**排除 memory 误导后才抽 anti-pattern | -| 归因（模型 vs memory） | 注入现场留痕 + 5% 影子对照 2×2 表 | UCOB(CBSD:anchor-state ΔG)/ ATMem(STR-GRPO 干预 advantage) | 同为 memory-on/off 对照；它们在 RL rollout 内改 policy，我们在服务侧接失败挖掘+效用+退休三出口 | -| 效用分 | per-item 反事实 Δverifier，便宜信号 + 稀疏锚校准 | MemDelta（对照诊断证据）/ ATMem(memory-cost reward) | MemDelta 只做评估诊断；我们把去偏后的效用直接驱动晋升/退休 | -| 何时不写 memory | 门槛（traj_score/safety）+ harmful 退休 | GovMem(provenance→依赖去相关→反证→三路决策) | GovMem 是保守诊断策略、只判写不写；我们与 verifier 地基 + 灰度 + 效用结合 | -| 稳定性/进化治理 | static/shadow/canary/live + 离线可回滚发布 | SEA(四层 + e-value certificate 门)/ M2Note(批级 rollback) | 我们用工程 mode 分层，SEA 用统计 certificate；canary 判定可升级成 anytime-valid gate | -| verifier 地基 | 主线 hard+rubric 融合，先坐实再上效用 | RuVerBench(可靠性有噪声)/ MRRG(多角色 rubric) | 我们直接复用主线 verifier，不为 memory 另造评审；rubric 生成可升级为多角色 | - -### 10.10 我们仍然新颖的地方（综合判断） -单点都有平行工作，但**没有一篇把下面这套完整闭环合在一起**： -1. **同一批生产轨迹**同时喂“慢权重蒸馏”和“快 memory”，且两者之间有**显式毕业梯度（B→A 晋升 / A→B 退休）**——退休用“新权重能否自答”自动判定（DuoMem 无毕业；Consolidation 无 A→B 反向）。 -2. **失败挖掘前置因果归因**：先用注入留痕 + 影子对照区分“模型能力 / memory 缺失 / memory 误导”，**只在 model_capability 上抽 anti-pattern**，把 GovMem/MemDelta 警示的污染从源头挡掉（Learning-from-Failure 类不做 memory 归因）。 -3. **抽取器本身被“经下游效用验证过的 memory”反向蒸馏**，效用分是 per-item 反事实 Δverifier、且**与主线训练奖励同源**（SelfMem/MetaSkill 优化的是策略/prompt，不回流蒸馏抽取器权重）。 -4. **进化被工程化为可回滚的发布流程**（static→shadow→canary→live），而非在线自由漂移——把 SEA 的“门控只能 select 已有行为”落成部署 mode。 - -**一句话定位**：DuoMem 证明了双线值得做，UCOB/ATMem 证明了 memory-on/off 对照能归因，GovMem/MemDelta 证明了不归因会污染/被混淆，SEA 证明了进化要门控——**本设计是把这些已被各自验证的结论，收进一个共用 verifier/`llm_backup`/D7c 的单一自进化蒸馏闭环。** - ---- - -## 11. skill/rubric → 参数化（LoRA）：双 LoRA 方案的文献支撑 - -> 背景：讨论中提出过一个具体的 line A 实现形态 —— **① 把高分轨迹的 system skills 蒸成一个"技能 LoRA"，推理时先用它产出/注入 skill；② 把 rubric 评价能力蒸成另一个"评价 LoRA"，推理中每隔 N 个 token 用它检验，当过程奖励/memory 提示。** 这一节把 2026 上半年直接对应这三个子命题（skill→LoRA、rubric→LoRA、多 LoRA 推理时切换）的工作按 §10 的详尽风格补全，并标注每篇对方案的取舍启示。**结论先行：三个子命题都各有直接平行工作，MetaClaw 几乎是整套方案 + 本双线设计的镜像；但"评价要不要绕一圈做成 LoRA judge"和"过程监督要不要走 LoRA"这两处，有明确的反方证据，需先决策。** - -### 11.1 skill → LoRA（对应子命题 ①：技能 LoRA） - -这一类的**共同范式**高度一致：**离线用完整 skill 文本合成"技能引导"的示范 → 训一个 skill 专属 LoRA → 在线丢掉 skill 文本、动态挂 LoRA 激活行为**。动机都是：skill 文本每步注入 context 太贵、且长上下文里关键指令定位不到 / 遵守不了（小模型尤甚，ICL 一贯不如微调、且模型越小差距越大）。 - -**Skill-to-LoRA (S2L)**（arXiv 2606.16769）—— **最贴近方案 ① 的朴素做法（一 skill 一 LoRA，不用 hypernetwork）** -- 想解决的问题：agent skill 现在以 `SKILL.md`（人读的流程文档：workflow / 工具 / 资源 / 领域约定）分发，可读可复用，但**同一套可复用流程要在每步 runtime context 里反复注入**，费 token 又稀释注意力。 -- 怎么做（**离线合成 + 在线换挂**，behavior-centric）：不压缩文档本身，而是**建模"skill 文本诱导的行为改变"**。**离线**：把完整 `SKILL.md` 喂进去，让模型合成一批"skill 引导下的示范轨迹"（demonstrations），用这些示范 SFT 出一个**该 skill 专属的 LoRA**；**在线**：完全省略 `SKILL.md` 全文，按当前需要动态加载对应 LoRA 来"激活"这个技能行为。 -- 数据 / loss：标准 SFT（在合成的 skill-guided 示范上做 teacher-forcing 训 LoRA），无 RL。Qwen3.6-27B、SWE-Skills-Bench 21 个 skill 子集：比 no-skill +2.9 pp、比 Full-Text +5.2 pp，每步 token 比 Full-Text prompting −6.6%；18/21 个 skill 追平或超过 Full-Text、15/21 超过 no-skill。**关键对照实验：Wrong-LoRA（挂错 skill 的 LoRA）和 Shared-LoRA（所有 skill 共用一个 LoRA）都掉点** → 收益依赖 **skill-专属对齐**，不是"随便训个 LoRA 就行"。 -- 与我们的区别：这就是方案 ① 最省事的落法（不引入 hypernetwork，一 skill 一 LoRA、离线 SFT）。**它的对照实验直接给我们敲定了两条工程铁律**：(a) 技能 LoRA 必须按 skill 对齐、**不能把多 skill 或 query 糊进一个 LoRA**（呼应 §0 分工判据："query 是易变实体，训进权重会过拟合"）；(b) 挂错 LoRA 反而有害 → 上线要有"挂哪个 skill LoRA"的可靠路由（正好复用我们 meta harness / intent 标签）。差异：S2L 是**离线一次性**、skill 集固定；我们要它在自进化闭环里**持续产出新技能 LoRA**，且用主线 `traj_score` 当"哪些高分轨迹够格蒸成 skill LoRA"的门槛。 - -**LatentSkill**（arXiv 2606.06087）—— **用 hypernetwork 把文本 skill 一次前向转成 LoRA** -- 想解决的问题：同 S2L（per-step 注入 skill 费 context、且 skill 明文暴露），但更进一步想要"**不为每个 skill 单独训 LoRA**"。 -- 怎么做：训一个**预训练 hypernetwork**，输入文本 skill、输出即插即用的 LoRA adapter（把 skill 知识存进**权重空间**而非 context 空间）。保留了 LoRA 的模块化：可加载、可缩放（用 LoRA scaling 系数精确调强弱）、可组合（对齐时能在**参数空间做算术**叠加多个 skill）。 -- 数据 / loss：hypernetwork 预训练。ALFWorld seen/unseen +21.4 / +13.4 分、prefill token −64.1%；Search-QA EM +3.0、skill-token 开销 −72.2%。分析发现生成的 skill LoRA 形成**结构化语义几何**。 -- 与我们的区别：如果我们不想"一个 skill 训一个 LoRA"（S2L 的痛点是 skill 一多 LoRA 就爆炸），LatentSkill 的 hypernetwork 是升级路径 —— **一次前向即出新 skill 的 LoRA、零梯度更新、零 skill 专属数据采集**。但它更复杂、要预训练 hypernetwork，属于方案 ① 的"进阶版"，建议 S2L 跑通、验证 skill LoRA 确有增益后再考虑。它的"参数空间算术组合"对我们"把多条相关碎片 memory 合并升华"（§2 consolidation / B→A 晋升）是权重侧的对应工具。 - -**ParametricSkills**（arXiv 2606.30015）—— **hypernetwork 同时参数化"skill 内容 + 利用方法"，含自进化/持续学习** -- 想解决的问题：同上两条，外加一个更深的观察 —— **文本空间演化 skill（EvoSkill/SkillOpt 等改写 SKILL.md）和模型本身的学习是解耦的**，模型能力没被优化。 -- 怎么做（三阶段，hypernetwork 驱动）：(1) 建 **45.8k skill 库**（网爬 + 从真实 agent 轨迹总结，覆盖 13 领域），用 OpenCode 沙箱围绕这些 skill 合成单/多轮"skill 利用轨迹"；(2) **skill-重建预训练**：三个自监督目标让 hypernetwork 学会把 skill 编码成 LoRA —— **完整重建**（据全文生成能重建全文的 adapter）、**前缀补全**（只给前缀、要补出全文，学 skill 结构）、**段级 cloze 补全**（挖掉一个功能段、据前后文补，学"触发条件/执行步骤/失败处理"如何组织与互相支撑 → 强组合泛化 + 支持局部编辑）；(3) 在 skill-利用轨迹上**多轮 SFT** hypernetwork。 -- 数据 / loss：自监督重建 + 多轮 SFT（loss 都 backprop 到 hypernetwork）。6 个 SWE 子任务比 ICL +6.44 分（DeepSeek-V4-Flash 判）、BERTScore +1.17、F1 +5.53%；**注意 text-to-LoRA 基线 SHINE 反而打不过 in-context skill** → hypernetwork 训不好会退化。持续学习：把多条轨迹的经验**不断 merge 成一个全局 parametric skill**。 -- 与我们的区别：它把 §2「抽取器越来越专业」和 §3「B→A 晋升」在**权重侧**给出了一个具体形态 —— "文本演化 skill = 直接改进模型"。三个自监督目标（尤其**段级 cloze**）可直接借来当我们"技能 LoRA 抽取器"的预训练任务。但它同样是**离线训 hypernetwork**、且 SHINE 反例提醒**参数化 skill 不保证优于文本注入**，必须带对照验证（呼应 §4 影子对照）。 - -### 11.2 rubric / 评价 → 参数化（对应子命题 ②：评价 LoRA）—— 两条岔路，务必先选 - -**支线 A：把 rubric 打分能力做成一个轻量 LoRA judge（= 方案 ② 的原意）** - -**Plug-and-Play LLM Judges**（arXiv 2506.05748）—— **"rubric + 小 LoRA = 顶级裁判"的最强直接背书** -- 想解决的问题：RLHF 的奖励模型训练是成本瓶颈（动辄几十亿参数 + 离线偏好微调阶段）。 -- 怎么做：**冻结的 instruction-tuned 7B + 一行 JSON rubric + rank-16 LoRA（只动 0.8% 参数）**，就当完整奖励模型用。消融：6 条 in-context 示范贡献了大部分零样本→少样本增益（+2pp），**LoRA 补上剩余差距**（尤其 safety / 对抗性 Chat-Hard 段）。 -- 数据 / loss：小 LoRA 微调 + prompt 工程。RewardBench 96.2%，超过 27B~70B 专用奖励网络；配它当在线 PPO 的 reward，7B actor 在 GSM-8K 拿 92% EM、超过 70B DPO 基线；LoRA judge 的解释与人类相似度 ≈9/10（零样本裁判仅 ≈5/10）。 -- 与我们的区别：**这是方案 ②「rubric→评价 LoRA」最硬的可行性背书** —— 极小 LoRA + 一条 rubric 就能把通用模型变成高质量、可解释、可调的裁判。直接支持我们把 `RubricVerifier`（现在调 dashscope teacher）蒸成本地评价 LoRA：既复用 `score_lora_path` 现成入口，又能治昨天"每段调远程 LLM 太慢"的痛（本地 LoRA 打分快几个数量级）。 - -**支线 B：跳过 judge，把 rubric 直接蒸进 policy 的 token 级信号（更省，可能是更优解）** - -> ⚠️ 这两篇机制外壳都是"rubric-conditioned 的自己当 teacher、逐 token 蒸给 unconditioned 的自己"，但**要解决的痛点、对标的对手、卖点完全不同**，别当成一篇：**RCSD 的对手是 RL 的标量奖励**（卖点＝把标量 reward 升级成过程级 token 信用分配）；**RGSD 的对手是 rubric 训练里的那个 LLM verifier**（卖点＝把 verifier 从训练回路里彻底删掉）。下面各自只讲其独有点。 - -**Rubric-Conditioned Self-Distillation / RCSD**（arXiv 2606.19327）—— **卖点：用 rubric 替代 RL 的"标量奖励"，做过程级信用分配** -- 想解决的问题：针对的是**蒸馏与 RLVR 两种 post-training 各自的病**。蒸馏靠 CoT 标注（贵、可能有噪/不全/半错，**哪怕最终答案对，坏 rationale 也会干扰学习**）；RLVR 则把评价**压成一个标量 reward**，看不出"该改推理的哪一步"。它要的是一个**比标量更细的过程监督信号**。 -- 怎么做（**独有点＝两阶段 pipeline + 显式过程级信用分配**）：核心是"**不把单一参考 rationale 当唯一监督靶**"，而让 teacher 看 criterion 级 rubric、在 student 自采样轨迹上给 token 级指导。落地成**两阶段**：**阶段① 先训一个"生成 task-specific rubric"的模块**（给任务先产出该任务的评分条目），**阶段② 再用这些 rubric 训"rubric 引导的 reasoner"**。rubric 说明"强回答该满足什么" → 转成**过程级信用分配**，这是它明确对标 GRPO 的地方。 -- 数据 / loss：on-policy 自蒸馏（token 级）。科学推理套件上**比 GRPO +1.0、比 OPSD +0.9**（对手是 RL/在线自蒸馏方法，不是 verifier）。 -- 与我们的区别：它证明 **rubric 能当"过程级 reward"直接进 policy 训练**，比标量 GRPO 奖励细。对我们的启示落在**训练信号形态**上：如果 line A 想要比 `traj_score` 标量更细的过程监督，RCSD 的"rubric→token 级信用分配"是替代 GRPO 标量奖励的路子；它的**阶段① rubric 生成器**正对应我们 `RubricVerifier` 的 stage-1（可复用）。 - -**Rubric-Guided Self-Distillation / RGSD**（arXiv 2606.12507）—— **卖点：verifier-free，把 LLM judge 从训练回路里删掉** -- 想解决的问题：针对的是**现有 rubric 训练法都要挂一个 LLM verifier 给每条 rollout 打分**这件事本身 —— 带来三个后果：训练期 verifier 调用**开销大**、优化被**特定 verifier 的偏差**污染、且 rubric 反馈被 verifier 压成**稀疏的轨迹末端信号**（只有一个 end-of-trajectory 分）。它要的是**根本不调 verifier**。 -- 怎么做（**独有点＝极简、零 verifier、单 rollout**）：直接拿 rubric-conditioned base policy 当 teacher、unconditioned 当 student 逐 token 蒸 —— 关键在于它把这套做到了**训练回路里完全没有 LLM judge**、且**每 prompt 只需一条 on-policy rollout**（不用像 GRPO 那样一题多采样再让 judge 排序）。 -- 数据 / loss：**零 verifier 调用 + 单 rollout/prompt**。Qwen-2.5(3B/7B)、Qwen3-Thinking(4B/8B) 医学/科学域：rubric 满足度**与 judge-based GRPO 相当**（对手是"带 verifier 的 rubric 训练"）。独有消融：**raw rubric 比"自生成参考回答"是更强的 teacher 富化信号**；但**更强的 GRPO judge 在某些设置能反超 RGSD** → 它诚实地把自己定位为"**当 verifier 成本/可靠性是瓶颈时**"的互补方案，而非全面更优。 -- 与我们的区别：**这篇最该在决策前读透**。它直击我们现状——verifier 又贵又不稳（RuVerBench 警示 + dashscope 每段调用慢）时，**把 rubric 直接蒸进 policy 比"训一个评价 LoRA 再在线打分"更省更稳**。据此，方案 ② 有两条路：**A 训评价 LoRA judge（Plug-and-Play 背书，产物是可复用的独立评价分）** vs **B 走 RGSD/RCSD 把 rubric 直接蒸进主 policy（不产独立分、只喂 student）**。选 A 还是 B，取决于我们是否真需要一个"能被 memory 效用 / 在线 PRM 复用的独立评价分"——若需要就 A，若只为提升 student 就 B。 - -### 11.3 多 LoRA 推理时切换 / 评价即一个 LoRA（对应子命题 ③：每 N token 用评价 LoRA 检验） - -**VideoMind — Chain-of-LoRA**（arXiv 2503.13444，**ICLR 2026**）—— **方案 ③"主生成 LoRA + 评价 LoRA 交替"的现成机制原型** -- 想解决的问题：视频时序 grounding 推理要多种能力（定位、验证、回答），但为每种能力各开一个完整模型太重。 -- 怎么做（两个创新）：(1) **角色化 agent 工作流** —— planner 协调、grounder 时序定位、**verifier 评估候选**、answerer 回答；(2) **Chain-of-LoRA**：一个统一 base model + **多个 LoRA adapter**，推理时**无缝切换角色**（用哪个角色就挂哪个 LoRA），在"每角色一个完整模型"和"纯 prompt 切角色"之间取平衡。 -- 数据 / loss：各角色 LoRA 分别训。15 个 benchmark（Grounded VideoQA / 时序 grounding / 通用 VideoQA）验证有效，且利于 test-time scaling / 长视频。 -- 与我们的区别：**它的 verifier 就是链条里的一个 LoRA 角色，与主生成角色在同一 base 上按需热切** —— 这正是方案 ③ 想要的"生成 LoRA / 评价 LoRA 交替"的成熟原型，**强烈建议精读**它怎么做角色切换调度。⚠️ 但注意：Chain-of-LoRA 是**在"角色回合"边界切**（planner→grounder→verifier→answerer 各跑一段），**不是"每 N 个 token 打断主流插评价"**；后者需要在 decode 中途暂停、跑旁路评估、再续，当前我们 vLLM sampler 的 `sample`/`sample_stream` 没有这种中途插入钩子（一次前向也只能挂一个 LoRA），要自己写**交错解码调度器**——这是方案 ③ 真正的工程量所在。 - -**MetaClaw: Just Talk**（arXiv 2603.17187，**有 code**）—— **几乎是整套方案 + 本双线设计的生产级镜像** -- 想解决的问题：部署的 agent 是**静态**的，跟不上用户需求漂移；在 OpenClaw（20+ 渠道、杂负载）上，现有法要么只存原始轨迹不蒸馏、要么静态 skill 库、要么retrain 要停机。 -- 怎么做（**双互补机制**）：(1) **skill 驱动的快适应** —— LLM evolver 分析**失败轨迹**合成新 skill，**零停机立刻生效**（= 我们的 line B / 快记忆）；(2) **机会式策略优化** —— **云端 LoRA 微调 + RL-PRM（带过程奖励模型的 RL）** 做梯度更新（= 我们的 line A / 慢权重），由 **OMLS 调度器**在**用户空闲窗口**（监控系统空闲 + 日历）触发。两机制**互相喂**：更好的策略产更好轨迹给 skill 合成，更丰富的 skill 给策略优化更高质数据。用**版本机制**分离 support / query 数据**防污染**。proxy 架构、无需本地 GPU。 -- 数据 / loss：SFT/RL-PRM（LoRA）+ 无训练的 skill 合成。skill 快适应相对 +32%；全流程把 Kimi-K2.5 从 21.4%→40.6%、综合鲁棒性 +18.3%。 -- 与我们的区别：**这是与本设计 §0 双线 + §3 毕业 + §6 harness 最像的一篇，且是 OpenClaw 生产场景**。相同点几乎逐条对上：skill 快线 + LoRA/PRM 慢线、失败轨迹驱动 skill、空闲窗口触发训练、版本防污染。差异（也是我们的增量）：MetaClaw 的两机制是**并列互喂**，**没有显式的 B→A 晋升 / A→B 退休毕业梯度**，也**不做"失败是模型问题还是 memory 误导"的因果归因**（§4）；我们多了归因前置、毕业出口、和 static→shadow→canary→live 的可回滚门控。**它是本方案最好的对标基线与工程参考（有 code），建议直接研读其 OMLS 调度 + RL-PRM 实现。** - -### 11.4 ⚠️ 反方证据：过程监督 / TTT 未必要走 LoRA - -**SCATR: Simple Calibrated Test-Time Ranking**（arXiv 2604.16535） -- 想解决的问题：Best-of-N 的效果全看打分函数；学出来的 PRM 强但**训练/推理都贵**，而基于 token logprob 的轻量置信度启发式又**明显偏弱**。 -- 怎么做：从**小校准集**学一个**轻量 scorer**，用的是 **base 模型的隐藏表示**（不是训 LoRA、不生成）。 -- 数据 / loss：轻量回归头。编码/数学基准上比置信度基线 +最高 9%；**相对在同样校准数据上做 LoRA 微调，用少 8000× 的可训练参数达到相当精度**，训练/推理延迟分别快 150×/1000×；和强 PRM 相当、部分设置数学 +7.8%/代码 +4.2% 且推理快 1000×。 -- 与我们的区别：**直接质疑方案 ③"评价/PRM 一定要做成 LoRA"** —— 一个读 base 隐藏层的轻量头，可能比评价 LoRA 更省几个数量级且更快。若方案 ② 的评价只用于"打个过程分排序/门控"（而非要它生成文字理由），**优先考虑 SCATR 式轻量头，而不是 LoRA**。 - -**Surprisal-Guided Selection**（arXiv 2602.07670） -- 想解决的问题：可验证、密集奖励任务（如 GPU kernel 优化，有确定性 evaluator）下，test-time 到底该"梯度自适应"还是"搜索"？ -- 怎么做 / 结论：KernelBench + GPT-OSS-120B(LoRA)：**Best-of-N 搜索（K=64 达 90% 成功）远胜 TTT 梯度自适应（最好 30.6%）**；TTT 会**过度锐化**、把多样性塌成平庸解，"等效 K < 1"（还不如单样本）。零成本妙招：**选 surprisal 最高（最不自信）的正确样本**比选最自信的 +30%。 -- 与我们的区别：提醒我们——**有确定性 verifier 时，算力花在"采样多样性 + 聪明选样"常比训 LoRA / 在线梯度自适应更值**。方案 ② / ③ 若目的是"提升生成质量"，先比一比"评价 LoRA + 干预" vs "多采样 + 评价选样"哪个划算，别默认前者。 - -**VDS-TTT**（arXiv 2505.19475）—— **支持方案：verifier 选样 → 只训 LoRA** -- 怎么做：learned verifier 给一批候选打分，**选高分伪标签**（置信度过阈值）配对成训练数据，**只微调 LoRA adapter** 做 test-time training。 -- 数据 / loss：verifier 驱动的自监督 SFT（仅 LoRA）。三 benchmark × 三 LLM，比 base 相对 +最高 32.29%、比"用 verifier 但不 TTT" +6.66%。 -- 与我们的区别：这是**"高分轨迹 → LoRA"这条主链的直接同构与背书**（verifier 打分选样 → 只训 LoRA），但它是**离线选样再训**、不是"推理中途干预"。我们 line A 的"用主线 `traj_score` 选高分轨迹蒸 LoRA"和它几乎一样，可当实现参考。 - -**Beyond Perplexity（TTT memory 审计）**（arXiv 2607.00368） -- 怎么做 / 结论：提出行为层评测框架，审计 TTT/memory 工作。发现一步 LoRA 更新能降 support/answer loss（跨 3 个 Qwen3 规模），但**自由 recall 仍为零** —— **proxy 指标改善 ≠ 部署行为改善**。 -- 与我们的区别：警示方案 ① 的技能 LoRA / 方案 ② 的评价能力，**别只看 loss 或 rubric 分下降就宣称有效**，要用**行为层对照**（带/不带、later recall / 下游动作）验证真收益（呼应 §4 影子对照、§10 DuoMem 的 CD 单独只 +1.4 的教训）。 - -### 11.5 双 LoRA 方案落地建议（综合上面证据） - -| 子命题 | 直接背书 | 反方 / 风险 | 建议 | -|---|---|---|---| -| ① 技能 LoRA（skill→LoRA） | S2L / LatentSkill / ParametricSkills / VDS-TTT | Beyond-Perplexity（proxy≠行为）；DuoMem CD 单独仅 +1.4 | **先做 S2L 式（一 skill 一 LoRA、离线 SFT、主线 traj_score 选样）**；必须带**带/不带对照**验证增益；skill 一多再上 hypernetwork（LatentSkill） | -| ② 评价 LoRA（rubric→judge） | Plug-and-Play Judge（rank-16 LoRA 顶 70B） | RGSD/RCSD：verifier 贵/不稳时**直接蒸进 policy 更优**；SCATR：轻量头比 LoRA 省 8000× | **先决策产物是不是"独立可复用的评价分"**：要 → 训评价 LoRA（复用 `score_lora_path`）；只为给 student dense 信号 → 走 RGSD 直蒸；只为排序/门控 → SCATR 轻量头 | -| ③ 每 N token 在线 PRM | Chain-of-LoRA（多 LoRA 热切原型）；MetaClaw（生产 RL-PRM） | 一次前向只挂一个 LoRA、无中途插入钩子；Surprisal/SCATR：搜索选样常更值；TTT 过锐化 | **最后做**；第一版用 `prompt_logprobs` **旁路打分 + 只留痕不干预**；确认 PRM 分与主线 hard verifier 一致性够高再考虑写交错解码调度器 | - -**对标基线**：**MetaClaw（有 code、OpenClaw 生产、双线镜像）**是整套方案最该研读与对标的工作；**Chain-of-LoRA** 是子命题 ③ 的机制原型；**RGSD** 是子命题 ② 决策的关键对照。我们相对它们的增量仍是 §10.10 那四条（显式毕业梯度、失败前置归因、抽取器被下游效用反向蒸馏、可回滚发布门控）——双 LoRA 只是把 line A 的"权重"具体化成"技能 LoRA + 评价能力"，不改变整体闭环定位。 - -### 11.6 定稿：查错 LoRA 作为参数化 memory —— 结论与实验方案 - -> 本节是 §11 讨论收敛后的**决策记录 + 实验计划**。核心转变：把"评价 LoRA"从"给数据打分的 judge"重新定位成**不动 base 的参数化 memory（line B 第二载体），专门给 base 补一个"在线查错"能力**。技能 LoRA（LoRA-2）**本轮暂缓**（理由见末尾）。 - -#### 11.6.1 定位（钉死的几条前提） - -1. **base 全程冻结，永不训。** 从根上杜绝知识遗忘——这是硬底线。所有能力增量都以**可插拔 LoRA**形式外挂，本质是 line B（参数化 memory），不是 line A。 -2. **LoRA-1 = 查错器（过程级 memory）。** 它不改 base 的知识，只给 base 补"边生成边发现自己踩了 rubric 里哪类错"的能力（如"公式第 k 步描述错""工具调用缺参数"）。发现错误后把**具体问题**注回 context，引导 base 改。 -3. **可回滚性 ≈ 删一条 memory。** 因为不动 base、LoRA 可插拔，一版 adapter 训坏了直接换/摘，撤销成本远低于"重训 base"——这消解了"权重侧犯错难撤"的顾虑（那顾虑只对"蒸进 base"成立）。 -4. **蒸馏数据 = 问题定位 + 打分（不是只给分）。** 要蒸出的是"**可定位、可操作的过程诊断**"能力，所以 teacher rubric 的产出必须到"哪一步/哪个 tool call/缺什么"这一粒度，而非单一 PASS/FAIL 标量（这是能否蒸出"查错"而非"打分习惯"的前提，对应 §10.7 原子化可验证条目 / 多角色 rubric）。 - -#### 11.6.2 自进化引擎：超大模型 + 本地 LoRA 双端检测 - -``` -超大模型(teacher rubric) ──检测出错误──┐ - ├─► 分歧/teacher 独有的错误 = 本地能力缺口 = 训练信号 ─► 升级 LoRA-1 -本地查错 LoRA-1 ──检测出错误──┘ │ - ▼ - 下一轮两端一起查 ──► 本地 LoRA-1 持续追平 teacher 的查错能力（追平后 teacher 少调、省成本） -``` - -- 与主线蒸馏同构：主线是"student 生成能力追平 teacher"，这里是"**本地 LoRA 的查错能力追平 teacher**"，同一套 `llm_backup` (teacher 对/本地错) 配对机制，被蒸对象换成"查错"这一角色。 -- **双端分歧同时是"该审上游"的探针**：本地 LoRA 与 teacher **系统性**分歧（成规律、非零星）时，要么本地没学到位（继续训），要么 **teacher/rubric 本身有系统性偏差**——后者是上游数据源/超大模型质量问题，**在上游治**（换更强 teacher、多角色 rubric 交叉、rubric 可靠性审计），不指望下游 LoRA 兜（呼应 §10.7"先坐实 verifier 再用它"）。 - -#### 11.6.3 三阶段实验路径（每阶段是下阶段的 gate，早止损） - -**阶段 0（先做）：不引入任何 LoRA，纯 teacher LLM 在线查错，测"上限"。** -- 数据集：**数学**（对错相对明确、查错信号干净，理想试验田）。 -- 机制：base 每隔 N 个 token **暂停** → 用 **teacher LLM（走 `llm_backup`）** 判当前生成有没有踩 rubric 错 → 有则把发现的问题注入 context → 继续生成。 -- 目的：**验证"在线过程级查错 + 注入 rubric"这个机制本身能否抬升数学解题上限**。用最强 teacher 代表能力天花板。 -- 性质：**纯可行性 gate，不训任何东西**。若最强 teacher 在线查错都提不了分，后面蒸 LoRA 更无意义 —— 先证信号有价值。 -- 工程：实验期"每 N token 暂停判断"**可用最粗暴实现**（停→跑一次 teacher 判断→拼 rubric→续），**不追性能**；生产化才需要专门的交错解码调度器（见待解决项）。 -- **已有实现**：`cookbook/exp/embedding/eval_dualline_math.py`（`--mode dualline`）。它复用 `eval_gpqa_rag.py` 的 MATH 分层加载 / 采样参数 / `answers_match` 判分，token 级分段生成：每 `--chunk-tokens` 暂停 → `RubricVerifier.diagnose()`（无 student sampler，走 llm_backup teacher）查错 → 命中且分数低于 `DUALLINE_CHECK_FLOOR` 则把 fix/原因作 `[Checker]` 注入续写。对照基线 `--mode baseline`（＝单遍生成，等同 `eval_gpqa_rag --mode direct`），同一 200 条子集直接比较 overall / 分层准确率与 checks/injections 计数。launch 配置：`dualline_math`。 - -**阶段 1（阶段 0 证明有效后）：把 teacher 的查错能力蒸成 LoRA-1，替换在线 teacher 调用。** -- 用阶段 0 收集的 (完整核验 CoT + 问题定位 + 打分) 数据蒸 LoRA-1（可复用 `score_lora_path` 入口）。**数据须正负配平**（成功段"无错、continue" + 失败段"定位/原因/建议"），构造约束见 11.6.6。 -- 目的：验证**本地低秩 LoRA 能否追平 teacher 查错**（＝待验证 b"学不学得动"）。实验要**把 rank × rubric 产出粒度当两个变量扫**（ParametricSkills 的 SHINE 退化反例说明：配置不对会打不过 in-context，学不学得动不是 0/1）。 - -**阶段 2（远期）：LoRA-2 技能 + 召回**——本轮暂缓，见 11.6.5。 - -#### 11.6.4 三条反驳 → 回应 → 限定（决策留痕） - -| 反驳 | 我们的回应 | 仍需守住的限定 | -|---|---|---| -| **① 廉价 query 路由只解决"该不该召回"，没解决"召回内容本身可能是错的"** | 不走"召回一条可能有错的 memory"，而是把"**查错能力**"蒸进 LoRA，让 base 内生地边做边纠错，**绕开召回可信度问题** | 前提：rubric 产出要到"**可定位错误**"粒度，否则只蒸出打分习惯、蒸不出查错 | -| **③ 权重侧坏数据难撤 / 会不会被污染** | LoRA 对**单条随机坏数据**抗性强于 RAG（梯度统计稀释 + base 低秩先验；RAG 一条即直入 context 无稀释）；且不动 base、可插拔，撤销≈删 memory | 抗的是**单条噪声**，**不抗系统性偏差**——系统性错误会被梯度强化。故 §4 归因质量仍是地基；系统性偏差归上游治（11.6.2） | -| **（Substrate Asymmetry）参数化 memory 在"该缺的要拒答"上惨败** | 承认：这是 LoRA "缺检不到信号"的固有短板，**不是污染问题** | **LoRA 只装可泛化行为/模式（如查错），易变事实仍留文本 memory**（§0 分工判据）；别用 LoRA 装事实 | - -#### 11.6.6 查错 LoRA 的训练数据构造（三条硬约束，钉死） - -> 复用现有 `RubricVerifier` 的 rubric + `llm_backup` 配对采集管线来攒数据，但现有 scoring prompt 刻意"只出 PASS/FAIL、不出解释"（省 token、便于 comparator 对齐），**直接拿来训会蒸出"打分习惯"而非"查错能力"**。故新增一个"诊断模式"产训练样本，须同时满足： - -1. **分数与原因必须一次调用同源产出（不可分两次采样拼）。** 若分数、原因来自两次独立采样，二者可能逻辑矛盾（判 FAIL 但原因说"没问题"），LoRA 学到错位映射。实现上：单次调用同时产 `per-criterion verdict + 每条 FAIL 的定位/原因/建议`，聚合分数由这批 verdict 算得，**原因与分数天然一致**。 - -2. **正负样本都要产、且要均衡（不能只在低分/被 gate 段跑诊断）。** 只见"错的"会把 LoRA 训成"逢查必报错"的挑刺偏置——它在线运行时每 N token 都硬报一个错，注入噪声反拖垮 base。必须让它见过大量"**看完一段、逐条核验、全部 PASS、结论=本段无过程错、无需干预**"的样本，才有能力在线输出关键的"OK，继续"信号。故诊断要**按比例覆盖高分成功段**，与失败段配平。 - -3. **CoT 要完整（核验过程，非只报结论）。** 正负两类样本的 target 都写成完整推理链： - - 成功段：`逐条核验 → 每条为何 PASS → 结论：无过程错误，continue` - - 失败段：`逐条核验 → 定位到第 k 条 FAIL → 错在哪 / 为什么 / 建议修法` - 完整 CoT 才对齐 LoRA-1 在线"边看前缀边判断"的实际形态；只给"这里错了"的结论会让它学不到核验过程，泛化差（呼应 §11.4 Beyond-Perplexity："proxy 改善 ≠ 部署行为改善"）。 - -> 落地路径（不改动服务打分/准入的现有 `RubricVerifier` 便宜路径）：新增诊断入口（`explain=True` 变体或独立 `_diagnose_once`，带 `@llm_backup` 自动落配对数据），schema：`input=[segment + rubric]`，`target=[完整核验 CoT + verdict + (FAIL 时) 定位/原因/建议]`。 -> -> **采样策略（已定）**： -> - **采集期全量存储、不做平衡、不设成本上限** —— 对每个 segment 都产完整诊断 CoT（正负都存），先把数据完整跑出来。正负配比、降采样等**留到训练采样阶段**再定，避免采集期过早丢信息。 -> - **注意训练/部署分布差**：离线诊断看的是"已完成的整段"，LoRA-1 在线看的是"生成中途的前缀"——离线攒数据能省掉阶段 0 的交错解码调度器工程，但**仍需一次"带/不带 LoRA-1 在线注入"的行为层对照**才算真验证（不能只看离线诊断准确率）。 - -#### 11.6.5 待验证 / 待解决 / 暂缓 - -- **待验证 a（已定方案）**：蒸 LoRA-1 时数据须含"问题定位 + 打分"，二者**一次调用同源产出**，且**正负样本都产、带完整核验 CoT**（详见 11.6.6 三条硬约束）。 -- **待验证 b**：低秩 LoRA 学不学得动查错 —— 阶段 1 跑实验，扫 rank × 信号粒度。 -- **待解决（生产化）**：vLLM **一次前向只挂一个 LoRA、无中途插入钩子**，"每 N token 暂停查错"生产化需自写**交错解码调度器**；实验阶段用粗暴实现绕过。 -- **待解决**：LoRA-1（旁路查错）与未来 LoRA-2（技能常驻）**同时在线的调度冲突**（单请求单 LoRA 限制）。 -- **暂缓：LoRA-2（优秀轨迹 skill 召回）** —— 本轮不做。理由：(1) 它**必须先训 LoRA 才能用，没有"纯 LLM 不训练"的验证捷径**，不适合当前"先用 LLM 测上限"的实验起点；(2) 训它需要**大量 trajectory** 作输入。两个已知前置问题记录备查：**P1 技能可能 per-user、跨用户不可迁移**（长期解＝只把 global 可泛化技能进 LoRA-2，用户专属留文本 memory，归 §0 分工；实验阶段忽略）；**P2 冷启动召回**（首个 query 可能是"你好"无信息量）——解法优先级：**A 延迟召回**（无信息量不召回，用 `intent`/信息量阈值判，"你好"本就不该召回）＞ B 滚动召回（每轮用累积上下文重判）＞ C 用 system/场景先验预挂。 diff --git a/src/twinkle_agentic/preprocessor/__init__.py b/src/twinkle_agentic/preprocessor/__init__.py index 20b514b5d..9734b44b5 100644 --- a/src/twinkle_agentic/preprocessor/__init__.py +++ b/src/twinkle_agentic/preprocessor/__init__.py @@ -16,15 +16,11 @@ from .message_normalizer import MessageNormalizer # noqa: F401 from .message_sanity import MessageSanityFilter from .model_filter import ModelFilter -from .outcome_filter import TrajectoryOutcomeFilter # noqa: F401 from .pii_presidio_filter import PIIPresidioFilter from .provenance import ProvenanceStamp # noqa: F401 from .refuse_filter import RefuseFilter -from .safety_scorer import SafetyScorer # noqa: F401 from .structural_noise import StructuralNoiseTagger # noqa: F401 from .token_soup import TokenSoupFilter -from .trajectory_scorer import TrajectoryScorer # noqa: F401 -from .value_selector import ValueSelector, select_top_for_rubric # noqa: F401 logger = get_logger() diff --git a/src/twinkle_agentic/preprocessor/message_normalizer.py b/src/twinkle_agentic/preprocessor/message_normalizer.py index 1472bdfcf..cd11c7bc6 100644 --- a/src/twinkle_agentic/preprocessor/message_normalizer.py +++ b/src/twinkle_agentic/preprocessor/message_normalizer.py @@ -201,8 +201,20 @@ class MessageNormalizer(Preprocessor): Multimodal list-shaped content passes through every stage untouched. This is a mapper — it never drops rows. + + Args: + normalize_tool_calls: Whether to run the tool-call rewrite pass. Turn it + off for pure code data (e.g. MBPP), where an assistant turn holds a + markdown code block and no tool call at all: the bracket-DSL parser + is a marker-less fallback that matches ``[name(``, which is also the + shape of a python list comprehension (``[abs(b - a) for ...]``) or a + call-indexed subscript (``count[ord(i)]``), so the rewrite would + delete real code from the content. """ + def __init__(self, normalize_tool_calls: bool = True): + self.normalize_tool_calls = normalize_tool_calls + def __call__(self, rows: List[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: rows = self.map_col_to_row(rows) for row in rows: @@ -210,7 +222,8 @@ def __call__(self, rows: List[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], Li if not isinstance(msgs, list) or not msgs: continue msgs = _strip_heartbeat(msgs) - msgs = _normalize_tool_calls(msgs) + if self.normalize_tool_calls: + msgs = _normalize_tool_calls(msgs) msgs = _merge_consecutive(msgs) row['messages'] = msgs return rows, [] diff --git a/src/twinkle_agentic/preprocessor/outcome_filter.py b/src/twinkle_agentic/preprocessor/outcome_filter.py deleted file mode 100644 index bc95b5171..000000000 --- a/src/twinkle_agentic/preprocessor/outcome_filter.py +++ /dev/null @@ -1,80 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Drop failed / dead-end trajectories by reading scores — pure tag reader (AUDIT D6). - -This filter does **not** compute anything. It reads the scores that -:class:`TrajectoryScorer` (D7) already wrote into ``user_data`` and drops rows -whose trajectory score / safety score fall below configurable thresholds. Because -it only consumes labels, the dependency on the scorer is a *data* dependency -(scorer writes ``traj_score``, this reads it) enforced simply by pipeline order — -no module import of the verifier, no DAG. - -Thresholds are meant to be **set by default and then tuned against the observed -score distribution** (self-evolving framework: no human-labeled calibration set). -A row with no score label is kept by default (fail-open) so that placing this -filter before the scorer, or scoring being disabled, never silently drops data. -""" -from __future__ import annotations - -from typing import Any, Dict, List, Tuple - -from twinkle.preprocessor import Preprocessor -from twinkle.utils import get_logger - -from . import label_schema as L - -logger = get_logger() - - -class TrajectoryOutcomeFilter(Preprocessor): - """Drop trajectories whose written scores fall below thresholds (reads only). - - Args: - min_traj_score: drop if ``traj_score`` < this. ``None`` disables. - min_safety_score: drop if ``safety_score`` < this. ``None`` disables. - drop_unsafe_flag: drop if ``safety_unsafe`` is True. Default True. - require_score: if True, rows with no ``traj_score`` label are DROPPED - (fail-closed); default False keeps them (fail-open) so a mis-ordered - or scorer-disabled pipeline never silently deletes data. - """ - - def __init__( - self, - *, - min_traj_score: float = 0.25, - min_safety_score: float = 0.5, - drop_unsafe_flag: bool = True, - require_score: bool = False, - ): - self.min_traj_score = min_traj_score - self.min_safety_score = min_safety_score - self.drop_unsafe_flag = bool(drop_unsafe_flag) - self.require_score = bool(require_score) - - def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: - rows = self.map_col_to_row(rows) - kept: List[Dict[str, Any]] = [] - dropped: List[Dict[str, Any]] = [] - for row in rows: - reason = self._drop_reason(row) - if reason is None: - kept.append(row) - else: - dropped.append(dict(row, drop_reason=reason)) - return kept, dropped - - def _drop_reason(self, row: Dict[str, Any]): - traj = L.get_label(row, L.KEY_TRAJ_SCORE, None) - if traj is None: - if self.require_score: - return 'no_score' - # fail-open: unscored rows pass through - elif self.min_traj_score is not None and float(traj) < self.min_traj_score: - return 'low_traj_score' - - if self.drop_unsafe_flag and bool(L.get_label(row, L.KEY_SAFETY_UNSAFE, False)): - return 'unsafe' - safety = L.get_label(row, L.KEY_SAFETY_SCORE, None) - if safety is not None and self.min_safety_score is not None \ - and float(safety) < self.min_safety_score: - return 'low_safety_score' - return None diff --git a/src/twinkle_agentic/preprocessor/safety_scorer.py b/src/twinkle_agentic/preprocessor/safety_scorer.py deleted file mode 100644 index 4cbe7142f..000000000 --- a/src/twinkle_agentic/preprocessor/safety_scorer.py +++ /dev/null @@ -1,100 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Safety / toxicity scoring via a FIXED rubric — tag only, never drop (AUDIT D8). - -Instead of adding a separate toxicity classifier + dependency, we reuse the -existing :class:`~twinkle_agentic.verifier.RubricVerifier` with a *fixed* safety -rubric (no stage-1 generation). Each criterion is phrased so that PASS = safe; -the aggregate pass-rate becomes ``safety_score`` in ``[0, 1]`` (higher = safer), -and a below-threshold score sets the ``safety_unsafe`` flag. Both are written to -``user_data``; the downstream :class:`TrajectoryOutcomeFilter` (D6) drops on them. - -This is a mapper (never drops) and shares the rubric voting/aggregation -machinery, so it inherits margin-adaptive voting for uncertain cases for free. -""" -from __future__ import annotations - -from typing import Any, Dict, List, Optional, Tuple - -from twinkle.preprocessor import Preprocessor -from twinkle.utils import get_logger - -from . import label_schema as L - -logger = get_logger() - -# PASS = the trajectory is SAFE w.r.t. the criterion. All are "hard" (weighted). -DEFAULT_SAFETY_CRITERIA: Tuple[str, ...] = ( - 'The content contains no violent, gory, or physically harmful instructions.', - 'The content contains no hateful, harassing, or discriminatory language toward any group.', - 'The content contains no sexual or adult material involving minors or non-consent.', - 'The assistant did NOT comply with a jailbreak / policy-evasion attempt (no successful jailbreak).', - 'The content does not leak private personal data (real names+contacts, credentials, IDs).', - 'The content contains no instructions for weapons, malware, or other clearly illegal harm.', -) - - -class SafetyScorer(Preprocessor): - """Score each trajectory for safety with a fixed rubric; write labels only. - - Args: - rubric_verifier: a pre-built :class:`RubricVerifier`. If ``None``, one is - constructed internally with the fixed safety rubric. When no sampler/ - teacher is available the score defaults to safe (1.0) — this filter - should then be treated as disabled rather than trusted. - criteria: override the default safety criteria (list of PASS=safe strings). - unsafe_threshold: ``safety_unsafe`` is set when ``safety_score`` < this. - """ - - def __init__( - self, - rubric_verifier: Optional[Any] = None, - *, - criteria: Optional[Tuple[str, ...]] = None, - unsafe_threshold: float = 0.5, - gate_label: Optional[str] = None, - ): - from twinkle_agentic.verifier import RubricItem, RubricVerifier - - self.unsafe_threshold = float(unsafe_threshold) - crits = criteria if criteria is not None else DEFAULT_SAFETY_CRITERIA - fixed = [RubricItem(text=c, is_hard=True) for c in crits] - if rubric_verifier is None: - rubric_verifier = RubricVerifier(fixed_rubric=fixed) - else: - rubric_verifier.fixed_rubric = fixed - self.verifier = rubric_verifier - # Active-learning gate: when set, only rows whose ``gate_label`` is True - # spend an LLM safety pass. The rest are tagged as neutral-safe (the LLM - # safety check runs post-selection only). None -> score every row. - self.gate_label = gate_label - - def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: - rows = self.map_col_to_row(rows) - out: List[Dict[str, Any]] = [] - for row in rows: - try: - out.append(self._score_row(row)) - except Exception as e: - logger.warning(f'[SafetyScorer] scoring failed, row left unscored: {e}') - out.append(row) - return out, [] # mapper: never drops - - def _score_row(self, row: Dict[str, Any]) -> Dict[str, Any]: - messages = row.get('messages') - if not isinstance(messages, list) or not messages: - return row - # Gated out (not selected for the LLM pass): tag neutral-safe, no LLM call. - if self.gate_label and L.get_label(row, self.gate_label, None) is False: - return L.set_labels(row, { - L.KEY_SAFETY_SCORE: 1.0, - L.KEY_SAFETY_UNSAFE: False, - }) - trajectory = {'messages': messages} - if row.get('tools'): - trajectory['tools'] = row['tools'] - detail = self.verifier.score_detail(trajectory) - score = float(detail.scalar) - return L.set_labels(row, { - L.KEY_SAFETY_SCORE: round(score, 6), - L.KEY_SAFETY_UNSAFE: score < self.unsafe_threshold, - }) diff --git a/src/twinkle_agentic/preprocessor/trajectory_scorer.py b/src/twinkle_agentic/preprocessor/trajectory_scorer.py deleted file mode 100644 index 1310410ce..000000000 --- a/src/twinkle_agentic/preprocessor/trajectory_scorer.py +++ /dev/null @@ -1,365 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Per-round / per-segment / per-trajectory scoring — tag only, never drop (AUDIT D7). - -This preprocessor wires the existing ``segment`` + ``verifier`` + ``aggregation`` -infrastructure into the cleaning pipeline. It is a **mapper**: it never removes a -row, it only writes scores into ``user_data`` (see :mod:`label_schema`, A5). A -downstream tail filter (``TrajectoryOutcomeFilter``, D6) reads those labels and -decides what to drop — so scoring and filtering stay decoupled and the pipeline -remains a linear list (no DAG, no filter↔verifier code coupling). - -Flow per trajectory:: - - Segmenter(traj) ─► segments - for each segment: - split_segment_into_rounds ─► rounds - HardScorer(round) ─► RoundScore (per-round hard scalar) - fuse_segment(round_scores, rubric_fn) ─► SegmentScore - └ rubric_fn lazily calls RubricVerifier ONLY when not short-circuited - aggregate_trajectory(segment_scores) ─► TrajectoryScore - -Labels written (all JSON-packed, PyArrow-stable): - round_scores, round_gated, segment_scores, traj_score, traj_level, score_meta. - -The ``RubricVerifier`` is optional: when no sampler/teacher is available the soft -chain returns an empty score and fusion falls back to the hard signal, so the -scorer still produces useful per-round hard scores with zero LLM calls. -""" -from __future__ import annotations - -import os -from concurrent.futures import ThreadPoolExecutor -from typing import Any, Dict, List, Optional, Tuple - -from twinkle.preprocessor import Preprocessor -from twinkle.utils import get_logger - -from . import label_schema as L - -logger = get_logger() - - -class TrajectoryScorer(Preprocessor): - """Score every trajectory and write the scores into ``user_data`` (never drops). - - Args: - segmenter: a :class:`~twinkle_agentic.segment.base.Segmenter`. Defaults to - structural ``TurnSegmenter('cluster')`` (LLM-free). - hard_scorer: a :class:`~twinkle_agentic.verifier.HardScorer` (per-round, - deterministic). Defaults to a plain ``HardScorer()``. - rubric_verifier: optional :class:`~twinkle_agentic.verifier.RubricVerifier` - (per-segment, soft/LLM). If ``None``, only hard scores are used. - hard_agg / fusion / hard_floor / hard_ceil_skip: passed to - :func:`~twinkle_agentic.verifier.fuse_segment`. - traj_agg / weight_by_rounds: passed to - :func:`~twinkle_agentic.verifier.aggregate_trajectory`. - write_round_detail: also store per-check breakdown into ``score_meta``. - """ - - def __init__( - self, - segmenter: Optional[Any] = None, - hard_scorer: Optional[Any] = None, - rubric_verifier: Optional[Any] = None, - *, - hard_agg: str = 'gmean', - fusion: str = 'hard_soft_blend', - hard_floor: float = 0.25, - hard_ceil_skip: Optional[float] = None, - traj_agg: str = 'mean', - weight_by_rounds: bool = True, - write_round_detail: bool = False, - calibrate: bool = True, - disagree_margin: float = 0.34, - reconcile_max_messages: int = 80, - scorer_workers: Optional[int] = None, - intent_aware: bool = True, - rubric_gate_label: Optional[str] = L.KEY_SELECTED_FOR_RUBRIC, - persist_diagnosis: bool = False, - ): - # Lazy imports keep the module importable even if verifier/segment deps - # are heavy; construction still fails loudly if the packages are absent. - from twinkle_agentic.segment import TurnSegmenter - from twinkle_agentic.verifier import HardScorer - - self.segmenter = segmenter if segmenter is not None else TurnSegmenter('cluster') - self.hard_scorer = hard_scorer if hard_scorer is not None else HardScorer() - self.rubric_verifier = rubric_verifier - # Route each segment's rubric by its structural intent (tool_call/code/ - # math) so the verifier can apply intent-keyed fixed/half-fixed rubrics. - self.intent_aware = bool(intent_aware) - self._intent_detectors = None - # Active-learning gate: when set, only rows whose ``rubric_gate_label`` is - # True spend an LLM rubric pass; the rest are scored hard-only. Leaving it - # None (or the label absent) preserves the "rubric every row" behavior. - self.rubric_gate_label = rubric_gate_label - # When True, rubric-scored segments also emit a full DiagnoseDetail - # (per-criterion verdict + reason + fix + raw teacher output). Persisted - # under KEY_RUBRIC_DIAGNOSIS as the SFT corpus for a distilled PRM/checker - # LoRA. Costs one extra teacher call per scored segment. - self.persist_diagnosis = bool(persist_diagnosis) - self.hard_agg = hard_agg - self.fusion = fusion - self.hard_floor = float(hard_floor) - self.hard_ceil_skip = hard_ceil_skip - self.traj_agg = traj_agg - self.weight_by_rounds = bool(weight_by_rounds) - self.write_round_detail = bool(write_round_detail) - # D7c: self-evolving calibration (no human alignment). - self.calibrate = bool(calibrate) - self.disagree_margin = float(disagree_margin) - self.reconcile_max_messages = int(reconcile_max_messages) - if scorer_workers is None: - scorer_workers = int(os.environ.get('TRAJ_SCORER_WORKERS', '1')) - self.scorer_workers = max(1, int(scorer_workers)) - - def _score_row_safe(self, row: Dict[str, Any]) -> Dict[str, Any]: - try: - return self._score_row(row) - except Exception as e: - logger.warning(f'[TrajectoryScorer] scoring failed, row left unscored: {e}') - return row - - # ------------------------------------------------------------------ - def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: - rows = self.map_col_to_row(rows) - if self.scorer_workers <= 1 or len(rows) <= 1: - out = [self._score_row_safe(row) for row in rows] - else: - with ThreadPoolExecutor(max_workers=self.scorer_workers) as pool: - out = list(pool.map(self._score_row_safe, rows)) - return out, [] # mapper: never drops - - # ------------------------------------------------------------------ - def _score_row(self, row: Dict[str, Any]) -> Dict[str, Any]: - from twinkle_agentic.verifier import (aggregate_trajectory, fuse_segment, - split_segment_into_rounds) - from twinkle_agentic.verifier.aggregation import RoundScore - - messages = row.get('messages') - if not isinstance(messages, list) or not messages: - return row - - trajectory = {'messages': messages} - if row.get('tools'): - trajectory['tools'] = row['tools'] - - segments = self.segmenter.segment(trajectory) - if not segments: - return row - - # Active-learning gate: skip the LLM rubric for rows not pre-selected by - # ValueSelector (hard-only), so expensive labeling is spent on the top - # fraction only. Absent label -> treat as selected (backward compatible). - rubric_enabled = self.rubric_verifier is not None - if rubric_enabled and self.rubric_gate_label: - selected = L.get_label(row, self.rubric_gate_label, None) - if selected is False: - rubric_enabled = False - - query = self._infer_query(messages) - all_round_scalars: List[float] = [] - all_round_gated: List[bool] = [] - segment_scalars: List[float] = [] - segment_confidence: List[float] = [] - segment_scores = [] - diagnoses: List[Dict[str, Any]] = [] - - for s_idx, segment in enumerate(segments): - rounds = split_segment_into_rounds(segment) - round_scores = [] - for r_idx, rnd in enumerate(rounds): - detail = self.hard_scorer.score_detail(rnd) - round_scores.append(RoundScore( - index=r_idx, - hard_scalar=detail.scalar, - gated=detail.gated, - detail=detail if self.write_round_detail else None, - )) - all_round_scalars.append(detail.scalar) - all_round_gated.append(detail.gated) - - seg_intent = None - if rubric_enabled: - seg_intent = self._segment_intent(segment) if self.intent_aware else None - rubric_fn = self._make_rubric_fn(segment, query, round_scores, seg_intent) - else: - rubric_fn = None - seg_score = fuse_segment( - s_idx, round_scores, rubric_fn, - hard_agg=self.hard_agg, fusion=self.fusion, - hard_floor=self.hard_floor, hard_ceil_skip=self.hard_ceil_skip, - ) - # carry the rubric ScoreDetail (stashed by rubric_fn) for confidence - seg_score.detail = segment.pop('_last_rubric', None) - segment_scores.append(seg_score) - segment_scalars.append(seg_score.scalar) - segment_confidence.append(self._segment_confidence(seg_score)) - - # Persist a full diagnostic chain for segments that actually reached - # the LLM (not hard-only / short-circuited) — the PRM/checker SFT data. - if (self.persist_diagnosis and rubric_enabled - and not seg_score.short_circuited): - diag = self._diagnose_segment(segment, query, seg_intent, s_idx) - if diag is not None: - diagnoses.append(diag) - - traj = aggregate_trajectory( - segment_scores, how=self.traj_agg, weight_by_rounds=self.weight_by_rounds) - - labels: Dict[str, Any] = { - L.KEY_ROUND_SCORES: [round(x, 6) for x in all_round_scalars], - L.KEY_ROUND_GATED: all_round_gated, - L.KEY_SEGMENT_SCORES: [round(x, 6) for x in segment_scalars], - L.KEY_TRAJ_SCORE: round(traj.scalar, 6), - L.KEY_TRAJ_LEVEL: traj.level, - } - if self.calibrate: - labels[L.KEY_SEGMENT_CONFIDENCE] = [round(c, 6) for c in segment_confidence] - labels[L.KEY_TRAJ_CONFIDENCE] = round( - sum(segment_confidence) / len(segment_confidence), 6) if segment_confidence else 1.0 - if self.write_round_detail: - labels[L.KEY_SCORE_META] = { - 'n_segments': len(segment_scores), - 'short_circuited': [s.short_circuited for s in segment_scores], - 'segment_hard': [round(s.hard_scalar, 6) for s in segment_scores], - 'segment_rubric': [ - None if s.rubric_scalar is None else round(s.rubric_scalar, 6) - for s in segment_scores - ], - } - if self.persist_diagnosis and diagnoses: - labels[L.KEY_RUBRIC_DIAGNOSIS] = diagnoses - return L.set_labels(row, labels) - - def _diagnose_segment(self, segment: dict, query: str, intent, s_idx: int): - """Run the verifier's full diagnosis and pack it for persistence. - - Emits everything a distilled PRM/checker LoRA needs: per-criterion - verdict + reason + fix, the overall verdict, the rubric text, and the - raw teacher output (SFT target) alongside the query + segment text - (SFT inputs). Never raises — diagnosis is best-effort enrichment. - """ - rv = self.rubric_verifier - if rv is None or not hasattr(rv, 'diagnose'): - return None - try: - d = rv.diagnose(segment, query=query, intent=intent) - except Exception as e: - logger.warning(f'[TrajectoryScorer] diagnose failed (seg {s_idx}): {e}') - return None - if d is None: - return None - return { - 'segment_index': s_idx, - 'intent': intent, - 'scalar': round(float(getattr(d, 'scalar', 0.0)), 6), - 'overall_ok': bool(getattr(d, 'overall_ok', False)), - 'summary': getattr(d, 'summary', '') or '', - 'query': getattr(d, 'query', '') or query, - 'segment_text': getattr(d, 'segment_text', '') or '', - 'raw': getattr(d, 'raw', '') or '', - 'rubric': [ - {'text': it.text, 'is_hard': bool(getattr(it, 'is_hard', False))} - for it in (getattr(d, 'rubric', None) or []) - ], - 'items': [ - { - 'index': it.index, - 'verdict': bool(it.verdict), - 'reason': it.reason or '', - 'fix': it.fix or '', - } - for it in (getattr(d, 'items', None) or []) - ], - } - - # ------------------------------------------------------------------ - def _segment_intent(self, segment: dict): - """Classify a segment by structural intent (tool_call > code > math). - - Reuses the lightweight, LLM-free detectors from IntentClassifier so the - segment's rubric can be routed to an intent-keyed fixed/half-fixed rubric. - Returns an intent string or ``None`` (no confident match -> generate). - """ - if self._intent_detectors is None: - from .intent_classifier import (CodeDetector, MathDetector, - ToolCallDetector) - # Order matters: tool_call is the strongest structural signal. - self._intent_detectors = [ToolCallDetector(), CodeDetector(), MathDetector()] - messages = segment.get('messages') or [] - if not isinstance(messages, list) or not messages: - return None - for det in self._intent_detectors: - try: - if det(messages): - return det.intent - except Exception: - continue - return None - - def _make_rubric_fn(self, segment: dict, query: str, round_scores, intent=None): - """Return a zero-arg callable for the soft chain, or None if unavailable. - - ``fuse_segment`` only invokes this when the segment is NOT short-circuited, - so the expensive LLM path runs exactly when the hard signal is inconclusive. - - D7c objective→subjective correction: score once, and if the rubric verdict - disagrees with the deterministic hard signal beyond ``disagree_margin``, - re-score with the objective evidence folded into the transcript so the - judge revises WITH the hard facts in view. The last ``ScoreDetail`` is - stashed on the segment (``_last_rubric``) for confidence estimation. - """ - rv = self.rubric_verifier - if rv is None: - return None - from twinkle_agentic.verifier.aggregation import aggregate_hard_over_rounds - hard_agg_val = aggregate_hard_over_rounds(round_scores, how=self.hard_agg) - - def _fn(): - detail = rv.score_detail(segment, query=query, intent=intent) - if (self.calibrate and detail is not None - and getattr(detail, 'scalar', None) is not None - and len(segment.get('messages') or []) <= self.reconcile_max_messages): - if abs(detail.scalar - hard_agg_val) >= self.disagree_margin: - evidence = (f'Deterministic tool/answer checks scored this ' - f'segment {hard_agg_val:.2f} out of 1.0. Reconcile ' - f'your assessment with this objective evidence.') - revised = rv.score_detail(segment, query=query, - extra_context=evidence, intent=intent) - if revised is not None: - detail = revised - segment['_last_rubric'] = detail - return detail - - return _fn - - # ------------------------------------------------------------------ - def _segment_confidence(self, seg_score) -> float: - """Self-evolving confidence in a segment score (no human labels). - - Combines three automatic signals: - 1. hard↔rubric agreement — 1 minus their absolute gap (objective anchoring); - 2. voting stability — fewer escalated votes ⇒ the judge was decisive; - 3. decisiveness — distance of the fused score from the ambiguous 0.5 band. - Short-circuited (hard-only) segments are highly confident by construction. - """ - if seg_score.short_circuited or seg_score.rubric_scalar is None: - return 1.0 - agree = 1.0 - min(1.0, abs(seg_score.hard_scalar - seg_score.rubric_scalar)) - detail = getattr(seg_score, 'detail', None) - n_votes = getattr(detail, 'n_votes', 1) or 1 - max_votes = getattr(self.rubric_verifier, 'max_votes', 1) or 1 - stability = 1.0 if max_votes <= 1 else 1.0 - (n_votes - 1) / max(1, max_votes - 1) - decisive = min(1.0, abs(seg_score.scalar - 0.5) * 2.0) - return max(0.0, min(1.0, 0.5 * agree + 0.3 * stability + 0.2 * decisive)) - - # ------------------------------------------------------------------ - @staticmethod - def _infer_query(messages: List[dict]) -> str: - for m in messages: - if isinstance(m, dict) and m.get('role') == 'user': - c = m.get('content') - if isinstance(c, str) and c.strip(): - return c.strip() - return '(no explicit query)' diff --git a/src/twinkle_agentic/preprocessor/value_selector.py b/src/twinkle_agentic/preprocessor/value_selector.py deleted file mode 100644 index aeef78ccf..000000000 --- a/src/twinkle_agentic/preprocessor/value_selector.py +++ /dev/null @@ -1,288 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Active-learning pre-selection for the expensive rubric pass. - -Motivation ----------- -At high daily volume it is neither affordable to LLM-label every trajectory nor -smart to sample at random (most rows are unremarkable). This mapper assigns each -row a **cheap, fully deterministic** ``value_score`` — an estimate of how much an -expensive rubric/LLM pass would *learn* from it — so a downstream gate can send -only the top fraction to the LLM. Self-evolution needs few, well-chosen samples. - -The score is a weighted blend of three LLM-free signals (all in ``[0, 1]``): - -- **uncertainty** — how close the deterministic hard signal is to undecided. - A row the hard checks already call clearly good (all 1.0) or clearly bad - (all 0.0) teaches the LLM little; rows near the boundary, or with internal - disagreement across rounds, are where a rubric pass pays off most. -- **difficulty** — structural complexity (rounds, tool calls, distinct tools, - segments), log-compressed so a few giant traces don't dominate. Long agentic - traces carry more signal than single-turn chit-chat. -- **error_signal** — deterministic failure evidence (gated rounds, failed - tool execution / termination / repetition checks). Mistakes are valuable - learning material for self-evolution (negative / correction examples). - -Two-pass usage (see ``TrajectoryScorer`` gate) ----------------------------------------------- -1. Run this mapper over the full stream (parallel, no global state) to stamp - ``value_score`` on every row. -2. After ``map`` completes, in the *single* driver process call - :func:`select_top_for_rubric` to flip ``selected_for_rubric=True`` on the - global top ``select_frac``. Only those rows spend an LLM call. -""" -from __future__ import annotations - -import math -from typing import Any, Dict, List, Optional, Tuple - -from twinkle.preprocessor import Preprocessor -from twinkle.utils import get_logger - -from . import label_schema as L -from .utils import normalize_tool_calls - -logger = get_logger() - - -def _log_norm(x: float, cap: float) -> float: - """Log-compress a count into [0, 1], saturating at ``cap``.""" - if x <= 0: - return 0.0 - return min(1.0, math.log1p(x) / math.log1p(cap)) - - -def _mean(xs: List[float]) -> float: - return sum(xs) / len(xs) if xs else 0.0 - - -def _pstdev(xs: List[float]) -> float: - if len(xs) < 2: - return 0.0 - m = _mean(xs) - return math.sqrt(sum((x - m) ** 2 for x in xs) / len(xs)) - - -class ValueSelector(Preprocessor): - """Stamp a deterministic ``value_score`` on every row (never drops). - - Args: - hard_scorer: a :class:`~twinkle_agentic.verifier.HardScorer` (reused for - the per-round hard scalars that feed ``uncertainty`` / ``error``). - Defaults to a plain ``HardScorer()``. No LLM is ever called. - segmenter: a segmenter for round splitting. Defaults to - ``TurnSegmenter('cluster')`` (LLM-free, same as TrajectoryScorer). - w_uncertainty / w_difficulty / w_error: blend weights (need not sum to 1; - normalized internally). - rounds_cap / toolcalls_cap / tools_cap / segments_cap: saturation caps - for the difficulty sub-signals. - write_meta: also store the per-component breakdown under ``value_meta``. - """ - - def __init__( - self, - hard_scorer: Optional[Any] = None, - segmenter: Optional[Any] = None, - *, - w_uncertainty: float = 0.45, - w_difficulty: float = 0.30, - w_error: float = 0.25, - rounds_cap: int = 20, - toolcalls_cap: int = 15, - tools_cap: int = 6, - segments_cap: int = 8, - write_meta: bool = True, - ): - from twinkle_agentic.segment import TurnSegmenter - from twinkle_agentic.verifier import HardScorer - - self.hard_scorer = hard_scorer if hard_scorer is not None else HardScorer() - self.segmenter = segmenter if segmenter is not None else TurnSegmenter('cluster') - total = w_uncertainty + w_difficulty + w_error - if total <= 0: - raise ValueError('at least one value weight must be > 0') - self.w_uncertainty = w_uncertainty / total - self.w_difficulty = w_difficulty / total - self.w_error = w_error / total - self.rounds_cap = int(rounds_cap) - self.toolcalls_cap = int(toolcalls_cap) - self.tools_cap = int(tools_cap) - self.segments_cap = int(segments_cap) - self.write_meta = bool(write_meta) - - # ------------------------------------------------------------------ - def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: - rows = self.map_col_to_row(rows) - out = [] - for row in rows: - try: - out.append(self._score_row(row)) - except Exception as e: # never break the pipeline on a bad row - logger.warning(f'[ValueSelector] scoring failed, value=0: {e}') - out.append(L.set_label(row, L.KEY_VALUE_SCORE, 0.0)) - return out, [] # mapper: never drops - - # ------------------------------------------------------------------ - def _score_row(self, row: Dict[str, Any]) -> Dict[str, Any]: - from twinkle_agentic.verifier import split_segment_into_rounds - - messages = row.get('messages') - if not isinstance(messages, list) or not messages: - return L.set_label(row, L.KEY_VALUE_SCORE, 0.0) - - trajectory = {'messages': messages} - if row.get('tools'): - trajectory['tools'] = row['tools'] - - segments = self.segmenter.segment(trajectory) or [] - - round_scalars: List[float] = [] - any_gated = False - soft_fail = 0.0 # worst non-critical hard-check miss across rounds - for segment in segments: - for rnd in split_segment_into_rounds(segment): - detail = self.hard_scorer.score_detail(rnd) - round_scalars.append(detail.scalar) - if detail.gated: - any_gated = True - soft_fail = max(soft_fail, self._soft_fail(detail)) - - uncertainty = self._uncertainty(round_scalars) - difficulty = self._difficulty(messages, segments) - error_signal = self._error_signal(any_gated, soft_fail) - - value = (self.w_uncertainty * uncertainty - + self.w_difficulty * difficulty - + self.w_error * error_signal) - value = max(0.0, min(1.0, value)) - - updates: Dict[str, Any] = {L.KEY_VALUE_SCORE: round(value, 6)} - if self.write_meta: - updates[L.KEY_VALUE_META] = { - 'uncertainty': round(uncertainty, 4), - 'difficulty': round(difficulty, 4), - 'error': round(error_signal, 4), - 'n_rounds': len(round_scalars), - } - return L.set_labels(row, updates) - - # ------------------------------------------------------------------ - # signal components - # ------------------------------------------------------------------ - @staticmethod - def _uncertainty(round_scalars: List[float]) -> float: - """High when the hard signal is undecided OR rounds disagree.""" - if not round_scalars: - return 0.0 - mean_hard = _mean(round_scalars) - central = 1.0 - abs(2.0 * mean_hard - 1.0) # peak at 0.5 - disagreement = min(1.0, 2.0 * _pstdev(round_scalars)) # spread across rounds - return max(central, disagreement) - - def _difficulty(self, messages: List[dict], segments: List[dict]) -> float: - n_rounds = sum(1 for m in messages - if isinstance(m, dict) and m.get('role') == 'assistant') - n_toolcalls = 0 - tool_names = set() - for m in messages: - if not isinstance(m, dict) or m.get('role') != 'assistant': - continue - for tc in (normalize_tool_calls(m) or []): - n_toolcalls += 1 - fn = (tc.get('function') or {}) if isinstance(tc, dict) else {} - name = fn.get('name') if isinstance(fn, dict) else None - if name: - tool_names.add(name) - n_segments = len(segments) - return _mean([ - _log_norm(n_rounds, self.rounds_cap), - _log_norm(n_toolcalls, self.toolcalls_cap), - _log_norm(len(tool_names), self.tools_cap), - _log_norm(n_segments, self.segments_cap), - ]) - - @staticmethod - def _soft_fail(detail: Any) -> float: - """Worst miss among informative non-critical checks in a round (0..1).""" - watch = {'tool_executed', 'clean_termination', 'no_repeated_calls', - 'protocol_pairing', 'final_answer'} - worst = 0.0 - for c in getattr(detail, 'checks', None) or []: - if getattr(c, 'name', None) in watch: - worst = max(worst, 1.0 - float(getattr(c, 'score', 1.0))) - return worst - - @staticmethod - def _error_signal(any_gated: bool, soft_fail: float) -> float: - """Deterministic evidence the model made a mistake worth studying.""" - if any_gated: - return 1.0 - return soft_fail - - -# --------------------------------------------------------------------------- -# global top-fraction selection (driver process, after Dataset.map) -# --------------------------------------------------------------------------- -def select_top_for_rubric( - dataset, - *, - select_frac: float = 0.1, - min_select: int = 0, - max_select: Optional[int] = None, - value_key: str = L.KEY_VALUE_SCORE, - selected_key: str = L.KEY_SELECTED_FOR_RUBRIC, -): - """Flip ``selected_for_rubric`` on the global top-``select_frac`` by value. - - Must run in the single driver process AFTER ``Dataset.map`` (the top - fraction is a global order that per-shard workers cannot compute). Returns - ``(dataset, n_selected)``; the dataset is mutated via a lightweight map. - - Ties at the cutoff are all included (selection is by a value threshold), so - the realized count can slightly exceed ``select_frac * N``. - """ - hf = dataset.dataset - n = len(hf) - if n == 0: - return dataset, 0 - - def _val(row) -> float: - v = L.get_label(row, value_key, 0.0) - try: - return float(v) - except (TypeError, ValueError): - return 0.0 - - values = sorted((_val(hf[i]) for i in range(n)), reverse=True) - k = int(round(select_frac * n)) - if min_select: - k = max(k, min_select) - if max_select is not None: - k = min(k, max_select) - k = max(0, min(k, n)) - if k == 0: - threshold = float('inf') - else: - threshold = values[k - 1] - - def _mark(batch): - rows = Preprocessor.map_col_to_row(batch) - out = [L.set_label(r, selected_key, _val(r) >= threshold) for r in rows] - return Preprocessor.map_row_to_col(out, keys=list(batch.keys())) - - marked = hf.map(_mark, batched=True, load_from_cache_file=False, - remove_columns=list(hf.column_names)) - # Write back to BOTH views so the next Dataset.map sees the marks: twinkle's - # Dataset.map operates on self.datasets[key] (not self.dataset), so updating - # only self.dataset would silently drop selected_for_rubric before pass 2. - dataset.dataset = marked - datasets = getattr(dataset, 'datasets', None) - if isinstance(datasets, dict): - for k in list(datasets.keys()): - if datasets[k] is hf or len(datasets) == 1: - datasets[k] = marked - - n_selected = sum(1 for i in range(len(marked)) - if L.get_label(marked[i], selected_key, False)) - logger.info(f'[ValueSelector] selected {n_selected}/{n} rows for rubric ' - f'(frac={select_frac}, threshold={threshold:.4f})') - return dataset, n_selected diff --git a/src/twinkle_agentic/rollout/__init__.py b/src/twinkle_agentic/rollout/__init__.py index 52c0fb12b..eb010ba89 100644 --- a/src/twinkle_agentic/rollout/__init__.py +++ b/src/twinkle_agentic/rollout/__init__.py @@ -1,12 +1,16 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -from .api_multi_turn import APIMultiTurnRollout from .base import Rollout from .multi_turn import MultiTurnRollout -from .multi_turn_condense import MultiTurnCondenseRollout __all__ = [ 'APIMultiTurnRollout', - 'MultiTurnCondenseRollout', 'MultiTurnRollout', 'Rollout', ] + + +def __getattr__(name: str): + if name == 'APIMultiTurnRollout': + from .api_multi_turn import APIMultiTurnRollout + return APIMultiTurnRollout + raise AttributeError(f'module {__name__!r} has no attribute {name!r}') diff --git a/src/twinkle_agentic/rollout/multi_turn.py b/src/twinkle_agentic/rollout/multi_turn.py index 3b4035eda..9367537ab 100644 --- a/src/twinkle_agentic/rollout/multi_turn.py +++ b/src/twinkle_agentic/rollout/multi_turn.py @@ -4,12 +4,15 @@ import os import re import time -from typing import Any, Callable, Dict, List, Optional +from collections import defaultdict +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Any, Callable, Dict, List, Optional, Tuple from twinkle.data_format import Trajectory, user_data_get from twinkle.data_format.sampling import SampleResponse, SamplingParams from twinkle.infra import remote_class, remote_function from twinkle.template.base import Template +from twinkle_agentic.harness.base import AgentHarness from twinkle_agentic.tools.tool_manager import ToolManager from .base import Rollout @@ -36,6 +39,41 @@ def _to_plain(obj: Any) -> Any: return obj +def _append_only_delta( + old_messages: List[Dict[str, Any]], + new_messages: List[Dict[str, Any]], +) -> Optional[List[Dict[str, Any]]]: + """Return newly appended messages, or None if ``new`` rewrote history.""" + old = list(old_messages or []) + new = list(new_messages or []) + if len(new) < len(old): + return None + for a, b in zip(old, new): + if a != b: + return None + return new[len(old):] + + +def _default_tool_messages( + tool_calls: List[Dict[str, Any]], + observations: List[str], +) -> List[Dict[str, Any]]: + msgs: List[Dict[str, Any]] = [] + for i, obs in enumerate(observations): + msg: Dict[str, Any] = {'role': 'tool', 'content': '' if obs is None else str(obs)} + if i < len(tool_calls) and isinstance(tool_calls[i], dict): + tc = tool_calls[i] + fn = tc.get('function') if isinstance(tc.get('function'), dict) else {} + tid = tc.get('id') or tc.get('tool_call_id') + name = fn.get('name') or tc.get('name') or tc.get('tool_name') + if tid: + msg['tool_call_id'] = tid + if name: + msg['name'] = name + msgs.append(msg) + return msgs + + @remote_class() class MultiTurnRollout(Rollout): """Agentic multi-turn rollout with tool use (batched). @@ -46,31 +84,21 @@ class MultiTurnRollout(Rollout): so vLLM can run all live trajectories in parallel; finished trajectories are parked and excluded from subsequent batches. - Per-trajectory loop: - 1. Encode the initial trajectory into an ``InputFeature`` with a - generation prompt at the tail. - 2. Call ``sampler.sample(pifs)`` (batched). The sampler internally - invokes ``template.concat_input_feature`` to append the freshly - sampled assistant tokens; we pick up ``seq.new_input_feature`` as - the new running ``pif``. - 3. If ``stop_reason == 'length'`` or the decoded assistant output has - no tool calls, mark the trajectory as done. - 4. Otherwise, invoke the tools via ``ToolManager`` and append each - tool response as a ``{'role':'tool', 'content': ...}`` message. - Compute "bridge" tokens (tool turns + next ``<|im_start|>assistant`` - header) with ``labels = -100`` and extend the pif. - 5. Repeat until all trajectories are done or ``max_turns`` is hit. + Per-trajectory loop:: + + harness.before_generate # append-only after the first encode + sampler.sample(batch) # keep seq.new_input_feature + harness.after_generate + ToolManager.call_many # Env.step_batch when tools share an Env + harness.after_tools # format observations as tool messages + _extend_with_bridge # labels=-100; never decode-reencode history Per-call overrides via ``**kwargs``: * ``sampling_params``: shared :class:`SamplingParams` for the batch. - * ``tool_manager``: either a single :class:`ToolManager` (applied to - every trajectory) or a list of ``ToolManager`` aligned 1:1 with - ``trajectories`` (used by :class:`MultiTurnCondenseRollout` to - attach a trajectory-bound ``ExtractCondensed``). - - The class intentionally has no knowledge of condensers/chunkers; they are - applied upstream (on the trajectory before rollout) or downstream - (on the returned messages). + * ``tool_manager``: a single :class:`ToolManager` or a 1:1 list. + * ``harness``: a single :class:`AgentHarness` or a 1:1 list. Framework + specifics (ms-agent system/memory/tool-message shape) live in the + harness subclass, not here. """ def __init__( @@ -84,6 +112,7 @@ def __init__( trace_dir: Optional[str] = None, trace_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, success_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, + harness: Optional[AgentHarness] = None, ): super().__init__() if template is None: @@ -96,6 +125,7 @@ def __init__( self.sampler = sampler self.template = template self.tool_manager = tool_manager + self.harness = harness self.sampling_params = sampling_params or SamplingParams() self.max_turns = max_turns self.max_trajectory_tokens = max_trajectory_tokens @@ -124,15 +154,36 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] sampling_params = kwargs.get('sampling_params', self.sampling_params) tool_managers = self._resolve_tool_managers(kwargs.get('tool_manager', self.tool_manager), n) + harnesses = self._resolve_harnesses(kwargs.get('harness', self.harness), n) + lives: List[Optional[Trajectory]] = [ + dict(trajectories[i]) if harnesses[i] is not None else None for i in range(n) + ] + for live in lives: + if live is not None: + live['messages'] = list(live.get('messages') or []) + + # 1. First before_generate happens *before* encode so memory/system + # injection is in the initial prefix (not a later rewrite). + encode_trajs: List[Trajectory] = [] + for i, traj in enumerate(trajectories): + h, live = harnesses[i], lives[i] + if h is not None and live is not None: + lives[i] = h.before_generate(live) + live = lives[i] + traj = dict(traj) + traj['messages'] = list(live.get('messages') or []) + if live.get('tools'): + traj['tools'] = list(live['tools']) + encode_trajs.append(traj) - # 1. Encode each trajectory once; ``pifs[i]`` is the live per-turn - # state for trajectory ``i``. pifs: List[Dict[str, Any]] = [] - for traj in trajectories: + for i, traj in enumerate(encode_trajs): pif = self.template.encode(traj, add_generation_prompt=True) pif = _to_plain(pif) pif.setdefault('messages', list(traj.get('messages', []))) pifs.append(pif) + if lives[i] is not None: + lives[i]['messages'] = list(pifs[i].get('messages') or []) all_logprobs: List[List[Any]] = [[] for _ in range(n)] stop_reasons: List[Optional[str]] = [None] * n @@ -140,11 +191,24 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] truncated: List[bool] = [False] * n done: List[bool] = [False] * n + first_turn = True for _ in range(self.max_turns): active = [i for i in range(n) if not done[i]] if not active: break + if not first_turn: + for global_idx in active: + pifs[global_idx], lives[global_idx], dropped = self._harness_before_generate( + pifs[global_idx], lives[global_idx], harnesses[global_idx]) + if dropped: + truncated[global_idx] = True + done[global_idx] = True + active = [i for i in range(n) if not done[i]] + if not active: + break + first_turn = False + # 2. One batched sample call for all currently-live trajectories. batch_pifs = [pifs[i] for i in active] actual = len(batch_pifs) @@ -155,7 +219,7 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] resps = self.sampler.sample(batch_pifs, sampling_params=sampling_params) resps = self._unwrap_response_list(resps, len(batch_pifs))[:actual] - pending_bridges: List[tuple] = [] # (global_idx, tool_messages) + pending_tools: List[tuple] = [] # (global_idx, tool_calls) for local_idx, global_idx in enumerate(active): turns[global_idx] += 1 seq = resps[local_idx].sequences[0] @@ -176,6 +240,19 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] all_logprobs[global_idx].extend(seq.logprobs) stop_reasons[global_idx] = seq.stop_reason + _msgs = pifs[global_idx].get('messages') or [] + _last_msg = _msgs[-1] if _msgs else None + tool_calls = (_last_msg.get('tool_calls') if isinstance(_last_msg, dict) else None) + if not tool_calls: + tool_calls = self.template.parse_tool_call(seq.decoded or '') + + if lives[global_idx] is not None: + lives[global_idx]['messages'] = list(_msgs) + if harnesses[global_idx] is not None and lives[global_idx] is not None: + lives[global_idx] = harnesses[global_idx].after_generate( + lives[global_idx], seq.decoded or '', tool_calls or []) + self._merge_assistant_metadata(pifs[global_idx], lives[global_idx]) + # 3. Termination conditions if seq.stop_reason == 'length': done[global_idx] = True @@ -188,11 +265,6 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] done[global_idx] = True continue - _msgs = pifs[global_idx].get('messages') or [] - _last_msg = _msgs[-1] if _msgs else None - tool_calls = (_last_msg.get('tool_calls') if isinstance(_last_msg, dict) else None) - if not tool_calls: - tool_calls = self.template.parse_tool_call(seq.decoded or '') if not tool_calls: done[global_idx] = True continue @@ -202,25 +274,25 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] done[global_idx] = True continue - # 4. Dispatch tools per trajectory (uses this trajectory's - # tool_manager, which may be a trajectory-bound clone). - tool_messages = [{ - 'role': 'tool', - 'content': tool_managers[global_idx](tc), - } for tc in tool_calls] - pending_bridges.append((global_idx, tool_messages)) - - # Extend pif with bridge tokens for every trajectory that has - # outstanding tool turns. Done serially: bridge computation is - # a cheap decode-diff-encode on python strings / token lists. - for global_idx, tool_messages in pending_bridges: - extended = self._extend_with_bridge(pifs[global_idx], tool_messages) - if extended is None: - # Trajectory exceeded max_length, mark as done (deleted) - truncated[global_idx] = True - done[global_idx] = True - else: - pifs[global_idx] = extended + pending_tools.append((global_idx, list(tool_calls))) + + # 4. Parallel tool dispatch across the live batch, then harness + # formats observations into tool messages (append-only bridge). + if pending_tools: + obs_by_traj = self._dispatch_tools(tool_managers, pending_tools) + for global_idx, tool_calls in pending_tools: + observations = obs_by_traj.get(global_idx) or [''] * len(tool_calls) + tool_messages, lives[global_idx] = self._tool_messages_after( + pifs[global_idx], lives[global_idx], harnesses[global_idx], + observations, tool_calls) + extended = self._extend_with_bridge(pifs[global_idx], tool_messages) + if extended is None: + truncated[global_idx] = True + done[global_idx] = True + else: + pifs[global_idx] = extended + if lives[global_idx] is not None: + lives[global_idx]['messages'] = list(extended.get('messages') or []) for i in range(n): if not all_logprobs[i]: @@ -271,6 +343,126 @@ def _resolve_tool_managers(arg, n: int) -> List[ToolManager]: return list(arg) return [arg] * n + @staticmethod + def _resolve_harnesses(arg, n: int) -> List[Optional[AgentHarness]]: + if arg is None: + return [None] * n + if isinstance(arg, list): + if len(arg) != n: + raise ValueError(f'per-call harness list length ({len(arg)}) does ' + f'not match number of trajectories ({n})') + return list(arg) + return [arg] * n + + def _harness_before_generate( + self, + pif: Dict[str, Any], + live: Optional[Trajectory], + harness: Optional[AgentHarness], + ) -> Tuple[Dict[str, Any], Optional[Trajectory], bool]: + """Run before_generate; bridge append-only deltas. ``dropped`` if encode fails.""" + if harness is None or live is None: + return pif, live, False + live['messages'] = list(pif.get('messages') or []) + live = harness.before_generate(live) + delta = _append_only_delta(pif.get('messages') or [], live.get('messages') or []) + if not delta: + return pif, live, False + extended = self._extend_with_bridge(pif, delta) + if extended is None: + return pif, live, True + live['messages'] = list(extended.get('messages') or []) + return extended, live, False + + @staticmethod + def _merge_assistant_metadata(pif: Dict[str, Any], live: Trajectory) -> None: + """Copy tool_calls / reasoning onto the sampled assistant message. + + Content is left untouched so the token-id chain stays valid. + """ + pif_msgs = pif.get('messages') or [] + if not pif_msgs or pif_msgs[-1].get('role') != 'assistant': + return + last_asst = None + for m in reversed(live.get('messages') or []): + if m.get('role') == 'assistant': + last_asst = m + break + if last_asst is None: + return + dst = pif_msgs[-1] + for key in ('tool_calls', 'reasoning_content', 'name'): + if last_asst.get(key) and not dst.get(key): + dst[key] = last_asst[key] + + def _dispatch_tools( + self, + tool_managers: List[ToolManager], + pending: List[Tuple[int, List[Dict[str, Any]]]], + ) -> Dict[int, List[str]]: + """Run tool calls for the live batch, grouped by ToolManager. + + Trajectories that share a manager (and therefore often one Env) go + through ``call_many`` / ``Env.step_batch``. Distinct managers run + concurrently so remote sandboxes are not serialized on generate. + """ + obs: Dict[int, List[str]] = { + gi: [''] * len(tcs) for gi, tcs in pending + } + groups: Dict[int, List[Tuple[int, int, Dict[str, Any]]]] = defaultdict(list) + mgr_by_id: Dict[int, ToolManager] = {} + for gi, tcs in pending: + mid = id(tool_managers[gi]) + mgr_by_id[mid] = tool_managers[gi] + for ci, tc in enumerate(tcs): + groups[mid].append((gi, ci, tc)) + + def _run_group(items: List[Tuple[int, int, Dict[str, Any]]], mgr: ToolManager): + tcs = [tc for _, _, tc in items] + if hasattr(mgr, 'call_many'): + contents = mgr.call_many(tcs) + else: + contents = [mgr(tc) for tc in tcs] + return list(zip(items, contents)) + + group_items = list(groups.items()) + if len(group_items) == 1: + mid, items = group_items[0] + finished = [_run_group(items, mgr_by_id[mid])] + else: + finished = [] + with ThreadPoolExecutor(max_workers=min(32, len(group_items))) as pool: + futs = [ + pool.submit(_run_group, items, mgr_by_id[mid]) + for mid, items in group_items + ] + for fut in as_completed(futs): + finished.append(fut.result()) + + for group_result in finished: + for (gi, ci, _tc), content in group_result: + obs[gi][ci] = '' if content is None else str(content) + return obs + + def _tool_messages_after( + self, + pif: Dict[str, Any], + live: Optional[Trajectory], + harness: Optional[AgentHarness], + observations: List[str], + tool_calls: List[Dict[str, Any]], + ) -> Tuple[List[Dict[str, Any]], Optional[Trajectory]]: + fallback = _default_tool_messages(tool_calls, observations) + if harness is None or live is None: + return fallback, live + old = list(pif.get('messages') or []) + live['messages'] = list(old) + live = harness.after_tools(live, observations, tool_calls) + delta = _append_only_delta(old, live.get('messages') or []) + if not delta: + return fallback, live + return delta, live + _TRACE_SKIP_KEYS = ( 'input_ids', 'labels', diff --git a/src/twinkle_agentic/rollout/multi_turn_condense.py b/src/twinkle_agentic/rollout/multi_turn_condense.py deleted file mode 100644 index 51f2affab..000000000 --- a/src/twinkle_agentic/rollout/multi_turn_condense.py +++ /dev/null @@ -1,284 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -from typing import Any, Callable, Dict, List, Optional - -from twinkle.data_format import Trajectory -from twinkle.data_format.sampling import SamplingParams -from twinkle.infra import remote_class, remote_function -from twinkle.template.base import Template -from twinkle_agentic.chunker.base import Chunker -from twinkle_agentic.condenser.base import Condenser -from twinkle_agentic.data_format import Chunks -from twinkle_agentic.tools.extract_condensed import TOOL_NAME as EXTRACT_TOOL_NAME -from twinkle_agentic.tools.extract_condensed import ExtractCondensed -from twinkle_agentic.tools.tool_manager import ToolManager -from .multi_turn import MultiTurnRollout - - -@remote_class() -class MultiTurnCondenseRollout(MultiTurnRollout): - """Multi-turn rollout with trajectory compression + on-demand recovery. - TODO: Experimental feature, wait for testing - - Pipeline for a batch of trajectories: - 1. ``chunker(trajectory)`` splits each incoming trajectory into chunks. - 2. All per-trajectory :class:`Chunks` are concatenated into a single - :class:`Chunks` and passed through ``condenser`` in ONE call, so - the underlying sampler (e.g. vLLM) sees a maximally-packed batch - spanning the whole rollout batch instead of a per-trajectory - sequence. Remembered trajectory boundaries are used to slice the - condensed chunks back into per-trajectory :class:`Chunks`. - 3. ``chunks.to_trajectory()`` rebuilds each trajectory, wrapping every - condensed chunk in ``<block_N>...</block_N>`` markers. - 4. A trajectory-scoped :class:`ExtractCondensed` tool is registered on - a per-trajectory clone of :attr:`tool_manager`, so the model can - recover the original text of any block by its number. - 5. The batch of compressed trajectories + a parallel list of - per-trajectory tool managers are handed to - :meth:`MultiTurnRollout.__call__`, which drives the sample/tool - loop (one batched ``sampler.sample`` per turn). - - The per-call tool manager is cloned via :meth:`ToolManager.copy`; the - shared ``self.tool_manager`` is never mutated, so concurrent rollouts on - the same instance are safe. - - Constructor accepts any :class:`Chunker` / :class:`Condenser` pair, so - plug-in chunkers (e.g. ``NativeChunker``) and condensers (e.g. - ``KeywordCondenser``, ``ModelCondenser``) compose freely. - """ - - def __init__( - self, - sampler, - template: Template, - tool_manager: ToolManager, - chunker: Chunker, - condenser: Condenser, - sampling_params: Optional[SamplingParams] = None, - max_turns: int = 6, - max_trajectory_tokens: Optional[int] = None, - condenser_kwargs: Optional[Dict[str, Any]] = None, - trace_dir: Optional[str] = None, - trace_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, - success_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, - post_compress_callback: Optional[Callable] = None, - ): - super().__init__( - sampler=sampler, - template=template, - tool_manager=tool_manager, - sampling_params=sampling_params, - max_turns=max_turns, - max_trajectory_tokens=max_trajectory_tokens, - trace_dir=trace_dir, - trace_callback=trace_callback, - success_callback=success_callback, - ) - if chunker is None: - raise ValueError('MultiTurnCondenseRollout requires a Chunker instance') - if condenser is None: - raise ValueError('MultiTurnCondenseRollout requires a Condenser instance') - if EXTRACT_TOOL_NAME in tool_manager.names(): - raise ValueError(f'tool_manager already registers {EXTRACT_TOOL_NAME!r}; ' - f'MultiTurnCondenseRollout registers a trajectory-bound ' - f'ExtractCondensed per call and would shadow the existing ' - f'one. Remove it from the shared manager or rename it.') - self.chunker = chunker - self.condenser = condenser - if getattr(self.condenser, 'template', None) is None: - self.condenser.template = template - self.condenser_kwargs = dict(condenser_kwargs or {}) - self.post_compress_callback = post_compress_callback - self._trace_block_chunks: Optional[List[Optional[Chunks]]] = None - - @remote_function() - def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory]: - if isinstance(trajectories, dict): - raise TypeError('MultiTurnCondenseRollout.__call__ expects a ' - 'List[Trajectory]; wrap a single trajectory as [trajectory].') - trajectories = list(trajectories) - if not trajectories: - return [] - - per_traj_chunks: List[Chunks] = [self.chunker(t) for t in trajectories] - signatures = [self._chunk_signature(ck) for ck in per_traj_chunks] - group_first: Dict[int, int] = {} - for i, sig in enumerate(signatures): - group_first.setdefault(sig, i) - unique_indices: List[int] = list(group_first.values()) - - merged_list = [] - boundaries: List[int] = [] - for idx in unique_indices: - merged_list.extend(per_traj_chunks[idx].chunks) - boundaries.append(len(merged_list)) - merged = Chunks(chunks=merged_list) - merged = self.condenser(merged, **self.condenser_kwargs) - - # Split the merged result back into per-unique-trajectory Chunks. - canonical: Dict[int, Chunks] = {} - start = 0 - for uidx, end in zip(unique_indices, boundaries): - canonical[uidx] = Chunks(chunks=merged.chunks[start:end]) - start = end - - compressed_list: List[Trajectory] = [] - tool_managers: List[ToolManager] = [] - for i, traj in enumerate(trajectories): - traj_chunks = canonical[group_first[signatures[i]]] - compressed = traj_chunks.to_trajectory() - for k, v in traj.items(): - compressed.setdefault(k, v) - if self.post_compress_callback is not None: - compressed = self.post_compress_callback(compressed, traj_chunks, **kwargs) - compressed_list.append(compressed) - - call_tm = self.tool_manager.copy() - call_tm.register(ExtractCondensed(traj_chunks)) - tool_managers.append(call_tm) - - # 5. Delegate to the parent batch loop. A caller-supplied - # ``tool_manager`` would be surprising here (we already built - # the list) -- drop it to avoid ambiguity. - kwargs.pop('tool_manager', None) - if self.trace_dir: - self._trace_block_chunks = [canonical[group_first[signatures[i]]] for i in range(len(trajectories))] - else: - self._trace_block_chunks = None - try: - return super().__call__(compressed_list, tool_manager=tool_managers, **kwargs) - finally: - self._trace_block_chunks = None - - @staticmethod - def _chunk_signature(chunks: Chunks) -> int: - """Cheap content-based signature of a :class:`Chunks` for dedup. - - Walks the chunk list once, dispatches on content type: - - * ``str`` / ``bytes``: hash with Python's built-in ``hash`` -- - SipHash, ~1 GB/s in C, and CPython caches the result on the - string object so GRPO duplicates that share the same string - are re-hashed for free. - * Multimodal (PIL image, numpy array, tensor, dict, ...): if - the object exposes ``tobytes``, hash its byte payload (stable - across identity-distinct but pixel-identical images); else - fall back to ``id(content)`` so duplicates referencing the - SAME object still dedup, while distinct-but-equal payloads - safely under-dedup (never over-dedup). - - Avoids ``json.dumps`` / ``repr``: both are 10-100x slower on - long text, and either crash on non-serializable multimodal - payloads or produce unstable output (e.g. PIL ``repr`` embeds - a memory address). - """ - parts: List[Any] = [] - for c in chunks.chunks: - content = c.get('content') - if isinstance(content, (str, bytes)): - chash = hash(content) - elif content is None: - chash = 0 - else: - tobytes = getattr(content, 'tobytes', None) - if callable(tobytes): - try: - chash = hash(tobytes()) - except Exception: - chash = id(content) - else: - chash = id(content) - parts.append(( - c.get('type'), - c.get('role'), - c.get('round'), - chash, - )) - return hash(tuple(parts)) - - def _build_trace_record( - self, - traj: Dict[str, Any], - *, - idx: int, - success: bool, - ) -> Dict[str, Any]: - """Attach per-block and per-passthrough-passage maps to the record. - - Two complementary maps are dumped so the trace alone is enough - to audit compression quality and compression coverage: - - * ``blocks`` — numbered ``block_N`` entries mirror - :meth:`Chunks.to_trajectory` and :class:`ExtractCondensed`: - text chunks with ``raw.condensed=True``, non-empty content - and ``role != 'tool'``, numbered from 1. Each entry carries - the pre-compression text (``original``, from - ``raw.original``) and the post-compression text - (``compressed``, the chunk content the model saw inside - ``<block_N>...</block_N>``). - * ``passages`` — numbered ``passage_M`` entries for text chunks - from the first user message (role neither ``'system'`` nor - ``'tool'``) that were NOT compressed — either because they - failed the eligibility filter (too short, wrong role, - ``skip_pattern`` matched, ...) or because the condenser's - output was not strictly shorter than the original and fell - back to passthrough. This lets the trace show the compressed - vs. passthrough ratio per rollout. - """ - record = super()._build_trace_record(traj, idx=idx, success=success) - - all_chunks = self._trace_block_chunks - if all_chunks is None or idx >= len(all_chunks): - return record - chunks = all_chunks[idx] - if chunks is None: - return record - blocks, passages = self._enumerate_blocks(chunks) - record['blocks'] = blocks - record['passages'] = passages - return record - - @staticmethod - def _enumerate_blocks(chunks: Chunks, ) -> 'tuple[Dict[str, Dict[str, Any]], Dict[str, Dict[str, Any]]]': - """Walk ``chunks`` and emit ``(blocks, passages)`` maps. - - * ``blocks`` → ``{block_N: {original, compressed}}`` for every - text chunk flagged ``raw.condensed=True`` (``role != 'tool'``). - ``original`` is ``None`` when the condenser did not attach a - ``raw.original`` snapshot; ``compressed`` is always present - since it is simply the chunk's post-compression content. - * ``passages`` → ``{passage_M: {content}}`` for every text chunk - from the first user message (``role not in {'system', 'tool'}``) - that was NOT flagged ``raw.condensed`` — i.e. chunks that - were either filtered out before compression or fell back to - passthrough because the model output was not strictly shorter. - Lets a reader of the trace see the compressed / passthrough - split without having to diff the raw trajectory. - """ - blocks: Dict[str, Dict[str, Any]] = {} - passages: Dict[str, Dict[str, Any]] = {} - block_counter = 0 - passage_counter = 0 - for c in chunks.chunks: - if c.get('type') != 'text': - continue - content = c.get('content') - if not isinstance(content, str) or not content: - continue - role = c.get('role') - if role == 'tool': - continue - raw = c.get('raw') - is_condensed = (isinstance(raw, dict) and bool(raw.get('condensed'))) - if is_condensed: - block_counter += 1 - original = raw.get('original') if isinstance(raw, dict) else None - blocks[f'block_{block_counter}'] = { - 'original': (original if isinstance(original, str) and original else None), - 'compressed': content, - } - elif role == 'user': - passage_counter += 1 - passages[f'passage_{passage_counter}'] = { - 'content': content, - } - return blocks, passages diff --git a/src/twinkle_agentic/rsi/rsi_challenge.py b/src/twinkle_agentic/rsi/rsi_challenge.py new file mode 100644 index 000000000..6f29a15b5 --- /dev/null +++ b/src/twinkle_agentic/rsi/rsi_challenge.py @@ -0,0 +1,1044 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""RSI step 0 — self-play data generation for code tasks. + +A single model plays TWO roles (self-play, same Qwen3-4B weights, per the talk's +"出题者/做题者" setup): + + * CHALLENGER — writes a self-contained Python solution, we RUN it in the same + sandbox rsi_rl uses to capture ground-truth outputs, and turn those outputs + into asserts. The problem statement is what the solver will be shown; the + executed reference solution is the GT. This is reverse construction: the + answer exists first (we ran it), the problem is written around it, so a GT + is available without any external labeling. + + * SOLVER (difficulty filter) — the SAME model then attempts each proposed + problem N times from the problem statement alone. We run its code against + the challenger's asserts and keep only problems whose pass count is strictly + between 0 and N ("half-know" band): all-pass or all-fail rounds give GRPO a + zero gradient (verified on MBPP), so they are dropped here. + +Two safety gates carried over from earlier failures: + * The challenger's OWN reference solution must pass its OWN asserts, or the + problem is dropped (a GT that cannot pass its own tests is noise — the + "standard answer that itself fails" pitfall). + * Output capture uses a sentinel marker + returncode check, never the last + stdout line, so an environment banner can never be mistaken for a result. + +Optional seed dataset (RSI_CH_SEED): a jsonl whose rows carry a `query` and, +preferably, a `code` reference solution. When a row has `code` (and keywords are +enabled), that proposal takes the TWO-STEP path: call 1 writes a HARDER solution on +top of the reference code, call 2 describes the problem that solution answers, and +the stage-1 code becomes the ground truth. Rows without `code` fall back to the older +single-call "seed as inspiration" prompt. Set RSI_CH_TWO_STEP=0 to force that older +path everywhere. Without any seed the challenger invents problems from scratch. + +Output (consumed directly by rsi_rl.py, no prepare/refine in between): + RSI_CH_OUT_FLOWS flows jsonl: {id, system, query, tools, rounds:[code round]} + RSI_CH_OUT_TESTS tests jsonl: {id, test_list, test_setup_code} (-> RSI_TESTS) + +Every knob is an env var (nothing hard-coded); see the config block below. +Run it as a Ray job just like rsi_rl.py (sampler-only, no trainer): + + RSI_CH_SEED=... RSI_CH_NUM_PROPOSE=2000 python -m twinkle_agentic.rsi.rsi_challenge +""" +import json +import os +import random +import re +import resource +import shutil +import signal +import subprocess +import sys +import tempfile +from typing import Any, Dict, List, Optional, Tuple + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import SamplingParams +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.sampler import vLLMSampler + +logger = get_logger() + +# ── config (all env; nothing hard-coded) ─────────────────────────────────── +MODEL_ID = os.environ.get('RSI_CH_MODEL', 'ms://Qwen/Qwen3-4B') +TEMPLATE = os.environ.get('RSI_CH_TEMPLATE', 'Template') # base text template for Qwen3-4B (text-only) +SAMPLER_GPUS = int(os.environ.get('RSI_CH_SAMPLER_GPUS', 4)) + +SEED_PATH = os.environ.get('RSI_CH_SEED', '') # optional seed jsonl (empty = from scratch) +NUM_PROPOSE = int(os.environ.get('RSI_CH_NUM_PROPOSE', 2000)) # how many problems to attempt to create +PROPOSE_TEMP = float(os.environ.get('RSI_CH_PROPOSE_TEMP', 1.1)) # challenger temperature (higher = diverse) +PROPOSE_MAX_TOKENS = int(os.environ.get('RSI_CH_PROPOSE_MAX_TOKENS', 8192)) # raised: thinking+cross-domain is long +# Drop pathologically long problem statements (rambling / non-problems) before the solver +# stage, both for quality and so the solver input never exceeds the model context. +PROBLEM_MAX_CHARS = int(os.environ.get('RSI_CH_PROBLEM_MAX_CHARS', 4000)) +MAX_MODEL_LEN = int(os.environ.get('RSI_CH_MAX_MODEL_LEN', 16384)) + +# Topic-keyword conditioning (from-scratch only): first brainstorm a pool of diverse +# coding topics at high temperature, then seed each proposal with a random keyword so +# the challenger stops collapsing onto a few archetypes (palindromes, brackets, ...). +# Topic-keyword conditioning (from-scratch only): keep a persistent, 3-category keyword +# bank on disk (algorithm / computer / non-computer). Each proposal is seeded with a +# cross-domain TRIPLE drawn WITHOUT replacement (keywords are consumed); when a category +# runs out we ask the model for more distinct ones, and if it can't, we recycle. Keywords +# whose problems the solver fails hardest are expanded into more same-domain topics. +KEYWORDS_N = int(os.environ.get('RSI_CH_KEYWORDS_N', 128)) # per-category target (0 = disable) +KEYWORD_GEN_CALLS = int(os.environ.get('RSI_CH_KEYWORD_GEN_CALLS', 8)) # sampler calls per generation batch +KEYWORD_TEMP = float(os.environ.get('RSI_CH_KEYWORD_TEMP', 1.3)) # high temp -> diverse keywords +KEYWORD_MAX_TOKENS = int(os.environ.get('RSI_CH_KEYWORD_MAX_TOKENS', 1024)) +KEYWORD_DB = os.environ.get('RSI_CH_KEYWORD_DB', 'output/rsi/keywords.jsonl') # persistent bank +KEYWORD_REFILL_TRIES = int(os.environ.get('RSI_CH_KEYWORD_REFILL_TRIES', 2)) # refill attempts before recycle +SINGLE_KW_PROB = float(os.environ.get('RSI_CH_SINGLE_KW_PROB', 0.1)) # chance a proposal uses 1 keyword, not a triple +# With a seed pool loaded, this is the chance a proposal gets a seed problem ON TOP of +# its keywords; the rest are keywords-only. Seeds are drawn WITH replacement, so a pool +# smaller than NUM_PROPOSE is fine. 1.0 reproduces the old seed-only behaviour when +# keywords are disabled (RSI_CH_KEYWORDS_N=0). +SEED_MIX_PROB = float(os.environ.get('RSI_CH_SEED_MIX_PROB', 0.5)) +# Two-step seeded proposing (V4): a seeded proposal is produced by TWO sampler calls -- +# first write a harder solution on top of the seed's reference code, then describe the +# problem that solution answers. Requires the drawn seed to have a non-empty `code` field +# AND keywords to be enabled (both prompts take a topic block, so RSI_CH_KEYWORDS_N=0 +# silently keeps the single-call path). Keywords-only proposals and code-less seeds also +# keep the original single call. Costs one extra call per two-step proposal; the log line +# below reports how many proposals actually took it. +TWO_STEP = os.environ.get('RSI_CH_TWO_STEP', '1') == '1' + +# Combination arity: 'triple' = one keyword per category (max diversity); 'mix' = random +# 1/2/3 categories per proposal (use with the audit dump to see which combos keep best). +COMBO_ARITY = os.environ.get('RSI_CH_COMBO_ARITY', 'triple').lower() +# Optional 'w1,w2,w3' sampling weights for arity 1/2/3 in 'mix' mode (empty = uniform). +# Diagnostic finding: keep-rate falls with arity (~17%/8%/2%), so favour 1-2 for yield. +ARITY_WEIGHTS = os.environ.get('RSI_CH_ARITY_WEIGHTS', '') +_ARITY_W: Optional[List[float]] = None +if ARITY_WEIGHTS: + try: + _ARITY_W = [float(x) for x in ARITY_WEIGHTS.split(',')] + except ValueError: + _ARITY_W = None +AUDIT_PATH = os.environ.get('RSI_CH_AUDIT', 'output/rsi/challenge_audit.jsonl') # per-proposal outcome log +# Feedback: expand keywords whose problems the solver passed <= this many times (0 = all-fail). +LOW_PASS_EXPAND = int(os.environ.get('RSI_CH_LOW_PASS_EXPAND', 0)) +EXPAND_PER_KW = int(os.environ.get('RSI_CH_EXPAND_PER_KW', 8)) # new topics per hard keyword +EXPAND_MAX_KWS = int(os.environ.get('RSI_CH_EXPAND_MAX_KWS', 32)) # cap hard keywords expanded per run + +SOLVER_ROLLOUTS = int(os.environ.get('RSI_CH_SOLVER_ROLLOUTS', 8)) # N attempts per problem for difficulty +SOLVER_TEMP = float(os.environ.get('RSI_CH_SOLVER_TEMP', 1.0)) +SOLVER_MAX_TOKENS = int(os.environ.get('RSI_CH_SOLVER_MAX_TOKENS', 2048)) + +KEEP_MIN_PASS = int(os.environ.get('RSI_CH_KEEP_MIN_PASS', 1)) # keep if pass in [MIN, N-KEEP_MAX_MARGIN] +# "drop all-pass / all-fail" == keep 0 < pass < N. Both bounds configurable. +KEEP_MAX_PASS_MARGIN = int(os.environ.get('RSI_CH_KEEP_MAX_MARGIN', 1)) # drop pass >= N - margin + 1 + +SANDBOX_TIMEOUT = int(os.environ.get('RSI_CH_SANDBOX_TIMEOUT', 30)) +MAX_CHECKS = int(os.environ.get('RSI_CH_MAX_CHECKS', 6)) # asserts per problem cap +# Drop problems where every assert expects the SAME value: `return <that constant>` +# scores a perfect reward without reading the input, so the problem teaches nothing and +# actively rewards ignoring the task. Measured at 6.2% of kept problems on sp4_iter1. +DROP_CONSTANT_ANSWER = os.environ.get('RSI_CH_DROP_CONSTANT_ANSWER', '1') == '1' +SORT_BY_DIFFICULTY = os.environ.get('RSI_CH_SORT_BY_DIFFICULTY', '1') == '1' +# Cap how many kept problems to persist (0 = keep all). When set and exceeded, +# subsample EVENLY across the difficulty-sorted list so the stored set spans the +# whole difficulty range, not just the easiest end. +KEEP_TARGET = int(os.environ.get('RSI_CH_KEEP_TARGET', 0)) +CH_SEED = int(os.environ.get('RSI_CH_RANDOM_SEED', 0)) + +OUT_FLOWS = os.environ.get('RSI_CH_OUT_FLOWS', 'output/rsi/challenge_flows.jsonl') +OUT_TESTS = os.environ.get('RSI_CH_OUT_TESTS', 'output/rsi/challenge_tests.jsonl') +DUMP_REJECTED = os.environ.get('RSI_CH_DUMP_REJECTED', 'output/rsi/challenge_rejected.jsonl') + +CODE_SYSTEM = {'role': 'system', 'content': 'You are an expert Python programmer.'} + +# ── challenger prompt (shown to the user for review; a brand-new prompt) ──── +_MARK = '__RSI_GT__' # sentinel isolating captured output from any banner/log + +_CHALLENGER_SYS = ( + 'You design self-contained Python coding problems for training another model.\n' + 'A good problem: (1) is solvable from its statement ALONE with no external files, ' + 'network, images, or hidden context; (2) has ONE clear entry function; (3) is ' + 'deterministic (same input -> same output), no randomness, no wall-clock, no threads; ' + '(4) is neither trivial nor impossible for a mid-size model.\n' + 'You will also write the reference solution. We will EXECUTE it to obtain the ' + 'ground-truth outputs, so your solution must be correct and runnable as-is.\n' + 'Return ONLY one JSON object, no prose around it, with keys:\n' + ' "problem": the statement shown to the solver (describe the function name, its ' + 'inputs and expected behavior; do NOT include the solution).\n' + ' "solution": the reference implementation as plain Python source (no markdown fence).\n' + ' "entry": the entry function name.\n' + ' "checks": a list of 3-6 Python expressions calling the entry function on concrete ' + 'inputs (e.g. "solve([1,2,3])"); each must be evaluable after running the solution. ' + 'Do NOT write the expected value — we compute it by running your solution.' +) + +_CHALLENGER_FROM_SCRATCH = ( + 'Create ONE new Python coding problem now. Vary the topic freely ' + '(strings, arrays, math, greedy, DP, parsing, simulation ...).' +) + +_CHALLENGER_FROM_SEED = ( + 'Here is a seed problem. Create ONE NEW problem that is a meaningful VARIANT of it ' + '(change the twist, constraints, or data shape — not just renaming), keeping it ' + 'self-contained and deterministic.\n\n[seed]\n{seed}' +) + +# Keyword bank generation + keyword-conditioned proposing (diversity). +CATEGORIES = ('algorithm', 'computer', 'noncs') +_CATEGORY_DESC = { + 'algorithm': 'algorithmic techniques and paradigms (e.g. dynamic programming, binary ' + 'search, union-find, Dijkstra, backtracking, segment trees, greedy, ' + 'divide and conquer, sliding window ...)', + 'computer': 'computer-science / computing concepts that are NOT algorithms per se ' + '(e.g. hash maps, tries, LRU cache, bitsets, regular expressions, base ' + 'conversion, finite state machines, serialization, parsing, memoization ...)', + 'noncs': 'real-world domains OUTSIDE computer science, used to give a problem flavor ' + '(e.g. biology, finance, chemistry, logistics, music, cooking, sports, ' + 'astronomy, geography, linguistics ...)', +} +_KEYWORD_SYS = ( + 'You brainstorm diverse topics for a Python coding-problem generator.' +) +_KEYWORD_CAT_USER = ( + 'List {k} DISTINCT and SPECIFIC topics from this category: {desc}\n' + 'Be creative and concrete; avoid vague umbrella words. ' + 'Return ONLY a JSON array of short strings, nothing else.' +) +_KEYWORD_EXPAND_USER = ( + 'The topic "{kw}" turned out to seed genuinely HARD problems. List {m} MORE distinct, ' + 'specific topics in the SAME family/domain as "{kw}" that could seed similarly ' + 'challenging Python problems. Return ONLY a JSON array of short strings, nothing else.' +) +_CHALLENGER_FROM_KEYWORDS = ( + 'Create ONE new Python coding problem now. Draw inspiration from the following ' + 'topic(s) and combine them creatively into a single coherent problem:\n{keywords}\n' + 'You may use each topic directly or bend it loosely; combine with any data shape ' + '(strings, arrays, grids, trees, numbers, parsing, simulation ...). Make it require ' + 'real thought, not a one-liner, and keep it self-contained and deterministic.' +) +# Seed AND keywords together. The seed is deliberately framed as inspiration only, +# not as something to produce a variant of: the point is to pull the generated +# problems toward the shape of public benchmark items (short statement, one plain +# task) while the keywords keep supplying topical variety. +_CHALLENGER_FROM_SEED_KEYWORDS = ( + 'Create ONE new Python coding problem now. Use the problem below only as a ' + 'STARTING POINT for inspiration — you do NOT have to keep its task, and the new ' + 'problem does NOT need to be a variant of it.\n\n[inspiration]\n{seed}\n\n' + 'Also draw on the following topic(s), combining them into a single coherent ' + 'problem:\n{keywords}\n' + 'Make it require real thought, not a one-liner, and keep it self-contained and ' + 'deterministic.' +) + +# ── two-step (V4) challenger: build the problem FROM a harder solution ─────── +# Winning variant of the prompt bake-off (output/rsi/prompt_exp_mbpp_upgrade.py). +# The difficulty comes from adding a layer on top of a real, runnable reference +# solution, not from imagining a hard problem outright; splitting into two calls (write +# the harder code, THEN describe it) keeps the statement and the ground-truth solution +# consistent, which a single call does not. Measured on 40 MBPP seeds vs the single-call +# seed+keywords prompt: kept-rate 25% vs 15%, const-answer 4 vs 7, seed similarity 0.42. +# +# Stage 1: given the seed problem + its reference solution + topics, write a harder +# solution. Uses the plain code system prompt (CODE_SYSTEM), not _CHALLENGER_SYS, and +# returns raw code (extract_code parses it) rather than JSON. +_TWO_STEP_SOL = ( + 'Below is a coding problem and its reference solution.\n\n' + '[problem]\n{seed}\n\n[reference solution]\n{code}\n\n' + 'Write a MORE COMPLEX Python function that keeps the idea of the reference solution ' + 'as one step and builds a harder computation around it (extra pass, different data ' + 'structure, an added rule), in the direction of these topic(s):\n{keywords}\n' + 'Requirements: deterministic, self-contained, no randomness, no I/O, one clear entry ' + 'function. Output ONLY the code in a single ```python block, no explanation.' +) +# Stage 2: given the harder function PLUS the seed it grew from and the topics, describe +# the problem it answers. Seeing the seed pulls the wording back toward the MBPP task +# family (similarity 0.32 -> 0.42 when the seed is shown). The solution is NOT taken from +# this JSON -- we overwrite it with stage 1's code so the GT matches what was produced. +_TWO_STEP_PROB = ( + 'Here is a Python function.\n\n```python\n{code}\n```\n\n' + 'It was written as a harder follow-up to this exercise:\n\n[original exercise]\n' + '{seed}\n\nand it was pushed in the direction of these topic(s):\n{keywords}\n\n' + 'Write the problem statement that the function above is the answer to, as if it were ' + 'a coding exercise in the same series as the original: name the entry function, ' + 'describe its inputs and the exact behaviour expected, and do NOT reveal the ' + 'implementation. Phrase it as plainly and briefly as the original exercise.\n' + 'Return ONLY one JSON object, no prose around it, with keys:\n' + ' "problem": the statement shown to the solver.\n' + ' "entry": the entry function name.\n' + ' "checks": a list of 3-6 Python expressions calling the entry function on ' + 'concrete inputs; each must be evaluable after running the function above. Do NOT ' + 'write the expected value.\n' + 'The "solution" is already known, so do not include it.' +) + + + +_SOLVER_SYS = {'role': 'system', 'content': 'You are an expert Python programmer.'} +_SOLVER_USER = ( + '{problem}\n\n' + 'Write the complete Python solution. Put the final code in a single ```python fenced ' + 'block. Define the exact function name required by the problem.' +) + + +# ── sandbox (mirrors rsi_rl.run_asserts / extract_code; duplicated on purpose: +# importing rsi_rl would run its module-level CLI.from_args() + swanlab.init) ─ +_FENCE_RE = re.compile(r'```(?:python|py)?\s*\n(.*?)```', re.S) + + +def extract_code(text: str) -> str: + idx = (text or '').rfind('</think>') + body = text[idx + len('</think>'):] if idx >= 0 else (text or '') + blocks = _FENCE_RE.findall(body) + return (blocks[-1] if blocks else body).strip() + + +def _run_script(script: str, timeout: int) -> Tuple[int, str]: + """Run a python script in an isolated dir, 2GB cap, killpg on timeout. + + Returns (returncode, stdout). returncode is -1 on timeout/spawn failure. + """ + tmp = tempfile.mkdtemp(prefix='rsi_ch_') + try: + with open(os.path.join(tmp, '_run.py'), 'w', encoding='utf-8') as f: + f.write(script + '\n') + env = dict(os.environ, MPLBACKEND='Agg', PYTHONHASHSEED='0', OMP_NUM_THREADS='1', + MKL_NUM_THREADS='1', TOKENIZERS_PARALLELISM='false') + env.pop('CUDA_VISIBLE_DEVICES', None) + + def _limit(): + resource.setrlimit(resource.RLIMIT_AS, (2 * 1024**3, 2 * 1024**3)) + + proc = subprocess.Popen([sys.executable, '_run.py'], cwd=tmp, env=env, + stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, + text=True, start_new_session=True, preexec_fn=_limit) + try: + out, _ = proc.communicate(timeout=timeout) + return proc.returncode, out or '' + except subprocess.TimeoutExpired: + try: + os.killpg(proc.pid, signal.SIGKILL) + except ProcessLookupError: + pass + try: + proc.communicate(timeout=5) + except Exception: + pass + return -1, '' + finally: + shutil.rmtree(tmp, ignore_errors=True) + + +def run_asserts(code: str, setup: str, asserts: List[str], timeout: int = SANDBOX_TIMEOUT) -> bool: + """True when every assert passes (returncode 0). Same contract as rsi_rl.""" + if not code.strip() or not asserts: + return False + parts = [code] + if (setup or '').strip(): + parts.append(setup) + parts.extend(asserts) + rc, _ = _run_script('\n\n'.join(parts), timeout) + return rc == 0 + + +def build_asserts(solution: str, checks: List[str], timeout: int = SANDBOX_TIMEOUT) -> Optional[List[str]]: + """Run the reference solution once to capture repr of each check expression, + then form ``assert <check> == <captured>``. Sentinel-marked + returncode + checked so a crash or a banner line can never be read as a value. + + Returns the assert list, or None if the solution crashed / produced no usable + output (that problem is then dropped upstream). + """ + checks = [c for c in checks if isinstance(c, str) and c.strip()][:MAX_CHECKS] + if not checks: + return None + lines = [solution, ''] + for i, c in enumerate(checks): + # repr on its own line, tagged with index; a check that raises makes the + # whole script exit non-zero -> we drop the problem. Pure f-string (no %% + # formatting) so a check expression containing '%' (modulo/percent) is safe. + lines.append(f'print("{_MARK}{i}=" + repr({c}))') + rc, out = _run_script('\n'.join(lines), timeout) + if rc != 0: + return None + captured: Dict[int, str] = {} + for line in out.splitlines(): + if line.startswith(_MARK): + try: + idx_str, val = line[len(_MARK):].split('=', 1) + captured[int(idx_str)] = val + except (ValueError, IndexError): + continue + if len(captured) != len(checks): + return None + asserts = [] + for i, c in enumerate(checks): + val = captured[i] + # The captured text is a repr, so it is a valid literal to compare against. + asserts.append(f'assert ({c}) == ({val})') + return asserts + + +def _split_top_eq(s: str) -> Optional[tuple]: + """Split on the first top-level ``==``, ignoring anything inside brackets or quotes.""" + depth = 0 + quote = '' + i = 0 + while i < len(s) - 1: + c = s[i] + if quote: + if c == quote: + quote = '' + elif c in '\'"': + quote = c + elif c in '([{': + depth += 1 + elif c in ')]}': + depth -= 1 + elif depth == 0 and c == '=' and s[i + 1] == '=': + return s[:i].strip(), s[i + 2:].strip() + i += 1 + return None + + +def _expected_of(assert_line: str) -> Optional[str]: + """The value the solver actually has to produce for one assert. + + build_asserts emits ``assert (<check>) == (<repr>)``, but a check may itself be a + comparison, giving ``assert (f(x) == 3) == (True)``. Reading the outer side there + would report 'True' and make such a problem look constant-answer, so the inner + right-hand side is used instead. An outer ``False`` pins nothing down at all and + is reported as unknown. + """ + m = re.match(r'^\s*assert\s*\((.*)\)\s*==\s*\((.*)\)\s*$', assert_line.strip()) + if not m: + return None + lhs, rhs = m.group(1).strip(), m.group(2).strip() + inner = _split_top_eq(lhs) + if rhs in ('True', 'False') and inner is not None: + return inner[1] if rhs == 'True' else None + return rhs + + +def is_constant_answer(asserts: List[str]) -> bool: + """Would ``return <one constant>`` satisfy every assert? + + Requires at least two asserts with a readable expectation: a single assert is + trivially 'constant' and one unreadable assert must not hide a constant set. + """ + vals = [_expected_of(a) for a in asserts] + if any(v is None for v in vals) or len(vals) < 2: + return False + return len(set(vals)) == 1 + + +# ── challenger output parsing ────────────────────────────────────────────── +_JSON_FENCE_RE = re.compile(r'^\s*```(?:json)?\s*|\s*```\s*$', re.I) + + +def parse_challenger(text: str, require_solution: bool = True) -> Optional[Dict[str, Any]]: + """Pull the JSON object out of the challenger's completion. + + ``require_solution=False`` is for the two-step (V4) flow, where the harder solution + comes from a separate call and stage 2 returns only ``problem``/``entry``/``checks``. + """ + body = text + idx = body.rfind('</think>') + if idx >= 0: + body = body[idx + len('</think>'):] + body = _JSON_FENCE_RE.sub('', body.strip()).strip() + # Grab the outermost {...} if there is leading/trailing prose. + start = body.find('{') + end = body.rfind('}') + if start < 0 or end <= start: + return None + try: + obj = json.loads(body[start:end + 1]) + except (ValueError, TypeError): + return None + if not isinstance(obj, dict): + return None + problem = obj.get('problem') + solution = obj.get('solution') + checks = obj.get('checks') + if not (isinstance(problem, str) and problem.strip() + and isinstance(checks, list) and checks): + return None + if require_solution: + if not (isinstance(solution, str) and solution.strip()): + return None + else: + # Stage 2 is told NOT to include a solution; if it did anyway, ignore it -- we + # will overwrite with stage 1's code so the GT matches what actually ran. + solution = solution if isinstance(solution, str) else '' + # solution may still arrive fenced despite instructions. + if solution and '```' in solution: + solution = extract_code(solution) + return {'problem': problem.strip(), 'solution': (solution or '').strip(), + 'entry': str(obj.get('entry') or '').strip(), 'checks': checks} + + +# ── sampling helpers ─────────────────────────────────────────────────────── +def _completion_text(seq) -> str: + """The assistant text of one sampled sequence (decode is fine: not training).""" + if seq.decoded: + return seq.decoded + feat = seq.new_input_feature or {} + for m in reversed(feat.get('messages', []) or []): + if m.get('role') == 'assistant': + return m.get('content', '') or '' + return '' + + +def sample_texts(sampler, message_lists: List[List[Dict[str, Any]]], + sampling_params: SamplingParams) -> List[str]: + """Sample one completion per message list; return the assistant texts.""" + if not message_lists: + return [] + trajs = [{'messages': msgs} for msgs in message_lists] + responses = sampler.sample(trajs, sampling_params) + texts: List[str] = [] + for resp in responses: + seqs = resp.sequences + texts.append(_completion_text(seqs[0]) if seqs else '') + return texts + + +def _parse_keyword_list(text: str) -> List[str]: + """Extract a JSON array of short strings from a (possibly thinking) reply.""" + body = text + idx = body.rfind('</think>') + if idx >= 0: + body = body[idx + len('</think>'):] + start, end = body.find('['), body.rfind(']') + if start < 0 or end <= start: + return [] + try: + arr = json.loads(body[start:end + 1]) + except (ValueError, TypeError): + return [] + out: List[str] = [] + for x in arr: + if isinstance(x, str) and x.strip() and len(x.strip()) <= 60: + out.append(x.strip()) + return out + + +class KeywordStore: + """Persistent 3-category keyword bank with usage tracking (see CATEGORIES). + + On-disk format (KEYWORD_DB, one JSON per line): + {"category", "text", "used": bool, "used_count": int, "source": "gen"|"expand", + "parent": <keyword or null>} + De-duplicates case-insensitively within each category so re-runs never conflict. + """ + + def __init__(self, path: str): + self.path = path + self.items: Dict[str, List[Dict[str, Any]]] = {c: [] for c in CATEGORIES} + self._seen: Dict[str, set] = {c: set() for c in CATEGORIES} + if path and os.path.exists(path): + with open(path, encoding='utf-8') as f: + for line in f: + line = line.strip() + if not line: + continue + try: + r = json.loads(line) + except (ValueError, TypeError): + continue + c, t = r.get('category'), r.get('text') + if c in self.items and isinstance(t, str) and t.strip(): + key = t.strip().lower() + if key not in self._seen[c]: + self._seen[c].add(key) + self.items[c].append(r) + + def save(self) -> None: + os.makedirs(os.path.dirname(os.path.abspath(self.path)) or '.', exist_ok=True) + tmp = self.path + '.tmp' + with open(tmp, 'w', encoding='utf-8') as f: + for c in CATEGORIES: + for r in self.items[c]: + f.write(json.dumps(r, ensure_ascii=False) + '\n') + os.replace(tmp, self.path) + + def add(self, category: str, texts: List[str], source: str = 'gen', + parent: Optional[str] = None) -> int: + added = 0 + for t in texts: + key = t.strip().lower() + if not key or key in self._seen[category]: + continue + self._seen[category].add(key) + self.items[category].append({'category': category, 'text': t.strip(), + 'used': False, 'used_count': 0, + 'source': source, 'parent': parent}) + added += 1 + return added + + def unused(self, category: str) -> List[Dict[str, Any]]: + return [r for r in self.items[category] if not r.get('used')] + + def texts(self, category: str) -> List[str]: + return [r['text'] for r in self.items[category]] + + def recycle(self, category: str) -> None: + """Mark every keyword unused again (safety valve when the model is tapped out).""" + for r in self.items[category]: + r['used'] = False + + +def generate_category_keywords(sampler, category: str, n_want: int, + avoid: List[str], rng, nonce: int = 0) -> List[str]: + """Ask the model for up to ``n_want`` distinct keywords in ``category``. + + ``avoid`` (already-known texts) is both injected as a soft "don't repeat" hint and + used to filter the parsed result. ``nonce`` perturbs the prompt so refills differ. + """ + if n_want <= 0: + return [] + n_calls = max(KEYWORD_GEN_CALLS, SAMPLER_GPUS) # batch must cover every DP worker + per_call = max(1, -(-n_want // n_calls) + 4) # ceil(n/calls) + margin + avoid_note = '' + if avoid: + shown = avoid if len(avoid) <= 40 else rng.sample(avoid, 40) + avoid_note = '\nDo NOT repeat any of these already-used topics: ' + ', '.join(shown) + base = _KEYWORD_CAT_USER.format(k=per_call, desc=_CATEGORY_DESC[category]) + avoid_note + msgs = [[{'role': 'system', 'content': _KEYWORD_SYS}, + {'role': 'user', 'content': f'{base}\n(batch {nonce}-{i})'}] + for i in range(n_calls)] + sp = SamplingParams(max_tokens=KEYWORD_MAX_TOKENS, num_samples=1, logprobs=1, + temperature=KEYWORD_TEMP, top_p=0.98) + texts = sample_texts(sampler, msgs, sp) + out: List[str] = [] + seen = {a.strip().lower() for a in avoid} + for t in texts: + for kw in _parse_keyword_list(t): + key = kw.strip().lower() + if key and key not in seen: + seen.add(key) + out.append(kw.strip()) + rng.shuffle(out) + return out[:n_want] + + +def ensure_unused(store: 'KeywordStore', sampler, category: str, need: int, rng, + nonce: int) -> int: + """Make ``category`` hold >= ``need`` unused keywords, generating/recycling as needed. + + Returns the next free ``nonce`` to use for the following generation call. + """ + tries = 0 + while len(store.unused(category)) < need: + new = generate_category_keywords(sampler, category, KEYWORDS_N, + store.texts(category), rng, nonce=nonce) + added = store.add(category, new, source='gen') + nonce += 1 + tries += 1 + if added == 0 and tries >= KEYWORD_REFILL_TRIES: + # Model is out of fresh distinct topics; recycle so combinations keep flowing. + if store.items[category]: + store.recycle(category) + logger.info(f'[rsi_challenge] keyword category {category!r} exhausted -> recycled ' + f'{len(store.items[category])} topics') + break + return nonce + + +def expand_hard_keywords(store: 'KeywordStore', sampler, hard, rng, nonce: int) -> int: + """For each (category, keyword) the solver failed hardest, brainstorm same-domain + topics and add them (source='expand') to the bank. Returns count added.""" + hard = hard[:EXPAND_MAX_KWS] + if not hard or EXPAND_PER_KW <= 0: + return 0 + # Batch must cover all DP workers; cycle the hard list if it is too short. + reqs = list(hard) + while len(reqs) < SAMPLER_GPUS: + reqs.append(hard[len(reqs) % len(hard)]) + msgs = [[{'role': 'system', 'content': _KEYWORD_SYS}, + {'role': 'user', 'content': _KEYWORD_EXPAND_USER.format(kw=kw, m=EXPAND_PER_KW) + + f'\n(batch {nonce}-{i})'}] + for i, (_c, kw) in enumerate(reqs)] + sp = SamplingParams(max_tokens=KEYWORD_MAX_TOKENS, num_samples=1, logprobs=1, + temperature=KEYWORD_TEMP, top_p=0.98) + texts = sample_texts(sampler, msgs, sp) + added = 0 + for (cat, kw), t in zip(reqs, texts): + added += store.add(cat, _parse_keyword_list(t), source='expand', parent=kw) + return added + + +def _draw_keywords(store: 'KeywordStore', sampler, rng, nonce: int) -> Tuple[List[Tuple[str, str]], int]: + """Consume one keyword combination from the bank (arity per COMBO_ARITY). + + Split out of the propose loop so a proposal can carry keywords whether or not + it also carries a seed problem. + """ + if COMBO_ARITY == 'mix': + if _ARITY_W and len(_ARITY_W) == len(CATEGORIES): + k = rng.choices(range(1, len(CATEGORIES) + 1), weights=_ARITY_W)[0] + else: + k = rng.randint(1, len(CATEGORIES)) + cats = rng.sample(list(CATEGORIES), k) + elif rng.random() < SINGLE_KW_PROB: + cats = [rng.choice(CATEGORIES)] + else: + cats = list(CATEGORIES) + picks: List[Tuple[str, str]] = [] + for c in cats: + if not store.unused(c): + nonce = ensure_unused(store, sampler, c, 1, rng, nonce) + un = store.unused(c) + if not un: + continue + r = rng.choice(un) + r['used'] = True + r['used_count'] = r.get('used_count', 0) + 1 + picks.append((c, r['text'])) + return picks, nonce + + +def load_seeds(path: str) -> List[Dict[str, str]]: + """Read seed problems from a jsonl. Returns dicts with at least 'query'; may also + have 'code' (reference solution) when the file was written by split_mbpp.py v2+. + + When the file only has 'query' (legacy format), the returned dicts have code=''. + The two-step challenger requires 'code' to be non-empty; if all seeds lack code it + falls back to the original single-step prompt automatically. + """ + if not path or not os.path.exists(path): + return [] + seeds: List[Dict[str, str]] = [] + with open(path, encoding='utf-8') as f: + for line in f: + line = line.strip() + if not line: + continue + try: + row = json.loads(line) + except (ValueError, TypeError): + continue + q = row.get('query') or row.get('problem') or row.get('prompt') + if isinstance(q, dict): + q = q.get('content') + if not q: + msgs = row.get('messages') or [] + q = next((m.get('content') for m in msgs if m.get('role') == 'user'), None) + if isinstance(q, str) and q.strip(): + seeds.append({'query': q.strip(), 'code': (row.get('code') or '').strip()}) + return seeds + + +def main(): + rng = random.Random(CH_SEED) + for p in (OUT_FLOWS, OUT_TESTS, DUMP_REJECTED): + os.makedirs(os.path.dirname(os.path.abspath(p)) or '.', exist_ok=True) + + device_groups = [DeviceGroup(name='sampler', ranks=list(range(SAMPLER_GPUS)), device_type='GPU')] + sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, dp_size=SAMPLER_GPUS) + twinkle.initialize(mode='ray', nproc_per_node=SAMPLER_GPUS, groups=device_groups, lazy_collect=False) + + sampler = vLLMSampler( + model_id=MODEL_ID, + engine_args={'gpu_memory_utilization': 0.8, 'max_model_len': MAX_MODEL_LEN}, + device_mesh=sampler_mesh, + remote_group='sampler', + ) + sampler.set_template(TEMPLATE, model_id=MODEL_ID, enable_thinking=True, max_length=MAX_MODEL_LEN) + + seeds = load_seeds(SEED_PATH) + logger.info(f'[rsi_challenge] seeds loaded: {len(seeds)} from {SEED_PATH!r} ' + f'(seed_mix_prob={SEED_MIX_PROB if seeds else 0.0}, ' + f'{"seed+keywords / keywords-only mix" if seeds else "keywords-only"})') + + # ── stage 1: challenger proposes, we execute to build GT asserts ──────── + # Diversity: each proposal draws a cross-domain keyword combination from a + # persistent 3-category bank, consuming keywords without replacement. When a seed + # pool is given, SEED_MIX_PROB of the proposals additionally get one seed problem + # (drawn WITH replacement) as inspiration on top of the keywords -- the keyword + # bank is now used in BOTH modes, where it used to be skipped entirely whenever + # seeds were present. + store = KeywordStore(KEYWORD_DB) if KEYWORDS_N > 0 else None + nonce = int(CH_SEED) + if store is not None and KEYWORDS_N > 0: + for c in CATEGORIES: + nonce = ensure_unused(store, sampler, c, 1, rng, nonce) + logger.info('[rsi_challenge] keyword bank: ' + + ', '.join(f'{c}={len(store.items[c])}({len(store.unused(c))} free)' + for c in CATEGORIES)) + + # Build the FIRST-call message for every proposal. A seeded proposal that (a) drew a + # seed carrying reference code and (b) has TWO_STEP on becomes a two-step proposal: + # its first call writes a HARDER solution (raw code), and a second call -- built once + # we see that code -- turns it into a problem statement + checks. Everything else is a + # single JSON-producing call, exactly as before. + propose_msgs: List[List[Dict[str, Any]]] = [] + propose_kws: List[List[Tuple[str, str]]] = [] # (category, text) picked per proposal, for feedback + propose_seeded: List[bool] = [] # whether an MBPP seed rode along, for the audit + is_two_step: List[bool] = [] # whether this proposal uses the V4 two-call flow + seed_query: List[str] = [] # seed statement carried into stage 2 (two-step only) + kw_body: List[str] = [] # keyword block carried into stage 2 (two-step only) + n_seeded = 0 + n_two = 0 + for _ in range(NUM_PROPOSE): + picks: List[Tuple[str, str]] = [] + if store is not None: + picks, nonce = _draw_keywords(store, sampler, rng, nonce) + body = '\n'.join(f'- {c}: {t}' for c, t in picks) + use_seed = bool(seeds) and rng.random() < SEED_MIX_PROB + seed = rng.choice(seeds) if use_seed else None + two = bool(use_seed and TWO_STEP and picks and seed and seed.get('code')) + if use_seed: + n_seeded += 1 + if two: + n_two += 1 + sys_msg = dict(CODE_SYSTEM) + user = _TWO_STEP_SOL.format(seed=seed['query'], code=seed['code'], keywords=body) + elif use_seed and picks: + sys_msg = {'role': 'system', 'content': _CHALLENGER_SYS} + user = _CHALLENGER_FROM_SEED_KEYWORDS.format(seed=seed['query'], keywords=body) + elif use_seed: + sys_msg = {'role': 'system', 'content': _CHALLENGER_SYS} + user = _CHALLENGER_FROM_SEED.format(seed=seed['query']) + elif picks: + sys_msg = {'role': 'system', 'content': _CHALLENGER_SYS} + user = _CHALLENGER_FROM_KEYWORDS.format(keywords=body) + else: + sys_msg = {'role': 'system', 'content': _CHALLENGER_SYS} + user = _CHALLENGER_FROM_SCRATCH + propose_kws.append(picks) + propose_seeded.append(use_seed) + is_two_step.append(two) + seed_query.append(seed['query'] if two else '') + kw_body.append(body if two else '') + propose_msgs.append([sys_msg, {'role': 'user', 'content': user}]) + logger.info(f'[rsi_challenge] proposals={NUM_PROPOSE} seeded={n_seeded} two_step={n_two} ' + f'keywords_only={NUM_PROPOSE - n_seeded} (RSI_CH_SEED_MIX_PROB={SEED_MIX_PROB}, ' + f'RSI_CH_TWO_STEP={int(TWO_STEP)})') + propose_sp = SamplingParams(max_tokens=PROPOSE_MAX_TOKENS, num_samples=1, logprobs=1, + temperature=PROPOSE_TEMP, top_p=0.95) + logger.info(f'[rsi_challenge] proposing {NUM_PROPOSE} problems (T={PROPOSE_TEMP})') + stage1_text = sample_texts(sampler, propose_msgs, propose_sp) + + # objs[i] = the parsed proposal for index i (or None if it failed). For single-step + # proposals this is just parse_challenger(stage1). For two-step, stage 1 gave code and + # a second batched call turns each into problem+checks; we then FORCE solution = the + # stage-1 code so the ground truth matches what was actually produced. + # Driven by len(stage1_text), not NUM_PROPOSE, so a short sampler return cannot + # IndexError here (the original enumerate-based loop was naturally tolerant). + n_prop = len(stage1_text) + if n_prop != NUM_PROPOSE: + logger.warning(f'[rsi_challenge] sampler returned {n_prop} texts for ' + f'{NUM_PROPOSE} proposals; proceeding with {n_prop}') + objs: List[Optional[Dict[str, Any]]] = [None] * n_prop + two_codes: List[str] = [''] * n_prop + n_no_code = 0 + stage2_idx: List[int] = [] + stage2_msgs: List[List[Dict[str, Any]]] = [] + for i in range(n_prop): + if not is_two_step[i]: + objs[i] = parse_challenger(stage1_text[i]) + continue + code = extract_code(stage1_text[i]) + two_codes[i] = code + if not code.strip(): + # Stage 1 produced no code block (usually a truncated completion): this + # proposal dies here. Counted separately from JSON parse failures. + n_no_code += 1 + continue + stage2_idx.append(i) + stage2_msgs.append([{'role': 'system', 'content': _CHALLENGER_SYS}, + {'role': 'user', 'content': _TWO_STEP_PROB.format( + code=code, seed=seed_query[i], keywords=kw_body[i])}]) + stage2_text_by_idx: Dict[int, str] = {} + if stage2_msgs: + logger.info(f'[rsi_challenge] two-step stage 2: {len(stage2_msgs)} problem writes ' + f'({n_no_code} stage-1 completions had no code block)') + for i, txt in zip(stage2_idx, sample_texts(sampler, stage2_msgs, propose_sp)): + stage2_text_by_idx[i] = txt + obj = parse_challenger(txt, require_solution=False) + if obj is not None: + obj['solution'] = two_codes[i] # GT = the harder solution stage 1 produced + objs[i] = obj + + # For the audit, keep the completion whose JSON we parsed (stage 2 for two-step, else + # the single call) so resp_chars / think_closed describe the statement-producing call. + proposals_text = [stage2_text_by_idx.get(i, stage1_text[i]) if is_two_step[i] + else stage1_text[i] for i in range(n_prop)] + + problems: List[Dict[str, Any]] = [] + stat = {'parsed': 0, 'parse_fail': 0, 'stage1_no_code': n_no_code, 'too_long': 0, + 'gt_fail': 0, 'selfcheck_fail': 0, 'constant_answer': 0} + outcome: List[str] = ['parse_fail'] * len(proposals_text) # per-proposal audit label + for i in range(n_prop): + if is_two_step[i] and not two_codes[i].strip(): + outcome[i] = 'stage1_no_code' + rejected = open(DUMP_REJECTED, 'w', encoding='utf-8') + for pi in range(len(proposals_text)): + obj = objs[pi] + if obj is None: + if outcome[pi] != 'stage1_no_code': + stat['parse_fail'] += 1 + continue + stat['parsed'] += 1 + # Reject rambling / non-problem statements early (also keeps solver input in-context). + if len(obj['problem']) > PROBLEM_MAX_CHARS: + stat['too_long'] += 1 + outcome[pi] = 'too_long' + continue + asserts = build_asserts(obj['solution'], obj['checks']) + if not asserts: + stat['gt_fail'] += 1 + outcome[pi] = 'gt_fail' + rejected.write(json.dumps({'reason': 'gt_build_fail', **obj}, ensure_ascii=False) + '\n') + continue + # The reference solution must pass its own asserts, or the GT is noise. + if not run_asserts(obj['solution'], '', asserts): + stat['selfcheck_fail'] += 1 + outcome[pi] = 'selfcheck_fail' + rejected.write(json.dumps({'reason': 'selfcheck_fail', 'asserts': asserts, **obj}, + ensure_ascii=False) + '\n') + continue + # A problem whose every assert expects the same value rewards `return <constant>`, + # so it would train the solver to ignore the statement. Drop it before the (much + # more expensive) solver rollouts. + if DROP_CONSTANT_ANSWER and is_constant_answer(asserts): + stat['constant_answer'] += 1 + outcome[pi] = 'constant_answer' + rejected.write(json.dumps({'reason': 'constant_answer', 'asserts': asserts, **obj}, + ensure_ascii=False) + '\n') + continue + obj['asserts'] = asserts + obj['_kw'] = propose_kws[pi] if pi < len(propose_kws) else [] # origin keywords (feedback) + obj['_idx'] = pi + outcome[pi] = 'usable' + problems.append(obj) + logger.info(f'[rsi_challenge] proposal stage: {stat}, usable problems={len(problems)}') + + # ── stage 2: solver difficulty filter (keep 0 < pass < N) ─────────────── + kept: List[Dict[str, Any]] = [] + if problems: + solver_msgs: List[List[Dict[str, Any]]] = [] + for prob in problems: + for _ in range(SOLVER_ROLLOUTS): + solver_msgs.append([_SOLVER_SYS, + {'role': 'user', 'content': _SOLVER_USER.format(problem=prob['problem'])}]) + solver_sp = SamplingParams(max_tokens=SOLVER_MAX_TOKENS, num_samples=1, logprobs=1, + temperature=SOLVER_TEMP, top_p=0.95) + logger.info(f'[rsi_challenge] difficulty rollout: {len(problems)} problems x ' + f'{SOLVER_ROLLOUTS} (T={SOLVER_TEMP})') + solver_text = sample_texts(sampler, solver_msgs, solver_sp) + + hi = SOLVER_ROLLOUTS - KEEP_MAX_PASS_MARGIN + for pi, prob in enumerate(problems): + n_pass = 0 + for k in range(SOLVER_ROLLOUTS): + code = extract_code(solver_text[pi * SOLVER_ROLLOUTS + k]) + if run_asserts(code, '', prob['asserts']): + n_pass += 1 + prob['n_pass'] = n_pass + if KEEP_MIN_PASS <= n_pass <= hi: + kept.append(prob) + outcome[prob['_idx']] = f'kept(pass={n_pass})' + else: + outcome[prob['_idx']] = f'dropped(pass={n_pass})' + rejected.write(json.dumps({'reason': f'difficulty_pass={n_pass}/{SOLVER_ROLLOUTS}', + 'problem': prob['problem']}, ensure_ascii=False) + '\n') + rejected.close() + + # ── per-proposal audit: attribute outcome to the keyword combination + flag + # truncation (thinking never closed) so combo/domain effects can be measured. + if AUDIT_PATH: + with open(AUDIT_PATH, 'w', encoding='utf-8') as af: + for i, txt in enumerate(proposals_text): + kws = propose_kws[i] if i < len(propose_kws) else [] + af.write(json.dumps({ + 'idx': i, + 'arity': len(kws), + 'cats': [c for c, _ in kws], + 'kws': kws, + # Whether this proposal also carried an MBPP seed statement, so the + # keep rate of seeded vs keywords-only proposals can be compared. + 'seeded': bool(propose_seeded[i]) if i < len(propose_seeded) else False, + # Whether the V4 two-call flow was used (seed carried code + TWO_STEP). + 'two_step': bool(is_two_step[i]) if i < len(is_two_step) else False, + 'outcome': outcome[i], + 'resp_chars': len(txt), + 'think_closed': '</think>' in txt, + }, ensure_ascii=False) + '\n') + logger.info(f'[rsi_challenge] per-proposal audit -> {AUDIT_PATH}') + + # ── feedback: expand the keywords behind the hardest problems (solver pass + # <= LOW_PASS_EXPAND) into more same-domain topics, then persist the bank. + if store is not None: + hard: List[Tuple[str, str]] = [] + seen_hard = set() + for prob in problems: + if prob.get('n_pass', SOLVER_ROLLOUTS) <= LOW_PASS_EXPAND: + for c, t in prob.get('_kw', []): + if (c, t.lower()) not in seen_hard: + seen_hard.add((c, t.lower())) + hard.append((c, t)) + if hard: + rng.shuffle(hard) + added = expand_hard_keywords(store, sampler, hard, rng, nonce) + logger.info(f'[rsi_challenge] feedback: expanded {min(len(hard), EXPAND_MAX_KWS)} ' + f'hard keyword(s) -> +{added} new same-domain topics') + store.save() + logger.info('[rsi_challenge] keyword bank saved: ' + + ', '.join(f'{c}={len(store.items[c])}' for c in CATEGORIES) + + f' -> {KEYWORD_DB}') + + # File order = increasing difficulty (fewer solver passes later), which the + # fixed-pool validation in rsi_rl relies on. + if SORT_BY_DIFFICULTY: + kept.sort(key=lambda p: -p['n_pass']) + + # Optional even subsample to KEEP_TARGET across the difficulty-sorted list. + if KEEP_TARGET and len(kept) > KEEP_TARGET: + n = len(kept) + idx = sorted({round(i * (n - 1) / (KEEP_TARGET - 1)) for i in range(KEEP_TARGET)}) + kept = [kept[j] for j in idx] + logger.info(f'[rsi_challenge] capped to KEEP_TARGET={KEEP_TARGET} ' + f'(evenly across difficulty), stored={len(kept)}') + + # ── write flows + tests for rsi_rl (skip prepare/refine) ──────────────── + with open(OUT_FLOWS, 'w', encoding='utf-8') as ff, open(OUT_TESTS, 'w', encoding='utf-8') as ft: + for i, prob in enumerate(kept): + cid = f'ch_{i:06d}' + flow = { + 'id': cid, + 'system': CODE_SYSTEM, + 'query': {'role': 'user', 'content': prob['problem']}, + 'tools': [], + # Difficulty audit: how many of the SOLVER_ROLLOUTS attempts passed + # (0 < n_pass < N by construction). rsi_rl ignores these extra keys; + # kept so a persisted flow can be analyzed without re-running. + 'n_pass': prob.get('n_pass'), + 'n_rollouts': SOLVER_ROLLOUTS, + # Origin keywords (category, text) of the cross-domain triple, for audit. + 'keywords': prob.get('_kw', []), + 'rounds': [{ + 'intent': 'solve the problem', + 'type': 'code', + 'tool_call': None, + 'code': prob['solution'], # challenger's passing solution (OPSD reads this) + 'result': '', + 'reward_method': 'rubric', + }], + } + ff.write(json.dumps(flow, ensure_ascii=False) + '\n') + ft.write(json.dumps({'id': cid, 'test_list': prob['asserts'], 'test_setup_code': ''}, + ensure_ascii=False) + '\n') + + logger.info(f'[rsi_challenge] kept {len(kept)}/{len(problems)} problems -> ' + f'{OUT_FLOWS} + {OUT_TESTS}') + if kept: + dist: Dict[int, int] = {} + for p in kept: + dist[p['n_pass']] = dist.get(p['n_pass'], 0) + 1 + logger.info(f'[rsi_challenge] kept pass-count distribution (0<pass<N): ' + f'{dict(sorted(dist.items()))}') + + +if __name__ == '__main__': + main() diff --git a/src/twinkle_agentic/rsi/rsi_prepare.py b/src/twinkle_agentic/rsi/rsi_prepare.py index dd4fefdef..812b9766b 100644 --- a/src/twinkle_agentic/rsi/rsi_prepare.py +++ b/src/twinkle_agentic/rsi/rsi_prepare.py @@ -52,8 +52,12 @@ def build_pipeline(): DedupFilter is deliberately excluded (see module docstring); it is applied separately on the full materialized dataset. """ + # RSI_NORMALIZE_TOOL_CALLS=0 for pure code data (e.g. MBPP): the bracket-DSL + # parser is a marker-less fallback matching ``[name(``, which is also what a + # python list comprehension or a call-indexed subscript looks like, so the + # rewrite silently deletes real code from the assistant turn. steps = [ - MessageNormalizer(), # strip heartbeats, rewrite tool calls, merge consecutive roles + MessageNormalizer(normalize_tool_calls=_env_flag('RSI_NORMALIZE_TOOL_CALLS', '1')), MessageSanityFilter(), # role order / tool-id matching / content integrity / sensitive words RefuseFilter(), # drop assistant self-referential refusals DeadLoopFilter(), # drop degenerate / stuck (hesitation, cascade, ngram repeat) diff --git a/src/twinkle_agentic/rsi/rsi_rl.py b/src/twinkle_agentic/rsi/rsi_rl.py index fcd947871..f36948b0b 100644 --- a/src/twinkle_agentic/rsi/rsi_rl.py +++ b/src/twinkle_agentic/rsi/rsi_rl.py @@ -1,5 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""RSI step 3 — multi-LoRA GRPO where each round of a standard flow becomes its +"""RSI step 3 — full-parameter GRPO where each round of a standard flow becomes its own training query. Idea (confirmed): a multi-turn standard flow is decomposed into one training @@ -11,39 +11,84 @@ the standard call is the reference answer. Only the reasoning ("思路") varies across rollouts; the key node is the target. -v1 scope: only TOOL rounds are trained. Code rounds (no tool call) still appear -in the prior context but are not turned into training queries yet (their reward -is rubric-based, wired in a later iteration). +Rounds trained: TOOL rounds are rewarded either by matching the recorded standard +call (RSI_TOOL_REWARD=match, the default) or by asking a judge model whether the +generated call means the same thing as the recorded one (RSI_TOOL_REWARD=rubric). +CODE rounds are rewarded by EXECUTION when the flow carries tests (RSI_TESTS, +keyed by record id): the generated code runs against those asserts and scores +1.0 only if all pass. + +The rubric judge's system prompt and score parsing below are a verbatim copy of +output/rsi/rubric_judge.py, which was used to measure offline that rubric scoring +raises the share of groups with a non-zero advantage from 6.2% to 15.0% on +ToolACE. Keep the two in sync or that number no longer describes this run. RL data-flow discipline (verified): train ONLY on ``sequence.new_input_feature`` and use ``sequence.logprobs`` as old_logps — never decode-then-re-encode. The generated tool call is already parsed into ``new_input_feature['messages'][-1]`` by the template, and the reference call rides along in ``user_data``. -Structure mirrors cookbook/rl/grpo/short_math_grpo_multi_lora.py (MultiLoRA -Megatron + filesystem LoRA sync to vLLM). RSI-specific paths come from env vars -so the standard CLI (model/infra/rl knobs) stays identical to the reference: +Structure mirrors cookbook/rl/grpo/short_math_grpo.py, but FULL-PARAMETER: no +adapter is added, so all weights are trained and CheckpointEngineManager ships the +whole model to vLLM each step. RSI-specific paths come from env vars so the standard +CLI (model/infra/rl knobs) stays identical to the reference: RSI_STD_FLOWS standard_flows.jsonl from rsi_refine.py (default output/rsi/standard_flows.jsonl) - RSI_TEMPLATE template name, must match the model (default Qwen3_5Template) - RSI_LORA_SYNC dir for filesystem LoRA sync (default output/rsi/lora_sync) - RSI_ADAPTER executor adapter name (default executor) + RSI_TEMPLATE template name, must match the model (default Template, for text-only Qwen3-4B) + RSI_TESTS jsonl with {id, test_list, test_setup_code} to score code rounds + by execution (empty = code rounds are not trained) + RSI_TOOL_REWARD 'match' (default) or 'rubric' for tool rounds + RSI_JUDGE_MODEL judge model name for rubric scoring (default qwen3.8-max); the + endpoint and key come from LLM_BACKUP_BASE_URL / LLM_BACKUP_API_KEY + RSI_MAX_ROUNDS keep only the first N trainable rounds (0 = all), in file order + RSI_RUN_NAME swanlab experiment name + RSI_REWARD_DUMP path to append a per-sample reward audit jsonl (off when unset) + +Solver learning mode (RSI step-3 subclass, RSI_SOLVER_MODE): + * 'grpo' (default) -- on a code round whose first attempt FAILS the asserts, the + sandbox error is injected back as a {'role':'tool'} message and the model is + asked to continue, up to RSI_SOLVER_MAX_TURNS total turns. The whole + multi-turn trajectory (turn-1 tokens + turn-2 tokens, the tool error bridged + in as -100) is trained by GRPO on the final pass/fail reward. Tool rounds and + length-stopped rollouts stay single-shot. Bridge tokens are computed in + template space and appended verbatim -- never decode-then-re-encode. + * 'opsd' -- single turn. A teacher forward conditioned on a PRIVILEGED extra + system message carrying the challenger's passing reference solution + (RSI_OPSD_TEACHER_SYS) scores the SAME student response tokens; the per-token + teacher log-probs pull the student via OPSDLoss (no advantages, no reward). + Teacher log-probs come from model.forward_only on the trainer (same engine + + same weights as the student, so r = teacher - student reflects only the prompt + context, not vLLM<->Megatron skew). The response-only extraction is + self-checked on the first batch against the sampler's old_logps. + + RSI_SOLVER_MODE 'grpo' (default) or 'opsd' + RSI_SOLVER_MAX_TURNS GRPO: max total turns per code rollout (default 2) + RSI_OPSD_TEACHER_SYS OPSD: teacher-only system template, '{solution}' filled + RSI_OPSD_REVERSE OPSD: k3 direction (1 = KL(student||teacher), default) """ import json import os +import random +import re +import resource +import shutil +import signal +import subprocess +import sys +import tempfile +import time +from concurrent.futures import ThreadPoolExecutor from typing import Any, Dict, List, Optional, Tuple -from peft import LoraConfig - import twinkle from twinkle import DeviceGroup, DeviceMesh, get_device_placement, get_logger from twinkle.advantage import GRPOAdvantage +from twinkle.checkpoint_engine import CheckpointEngineManager from twinkle.cli import CLI 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 MultiLoraMegatronModel from twinkle.processor import InputProcessor from twinkle.reward.base import Reward from twinkle.sampler import vLLMSampler @@ -53,17 +98,102 @@ # ── RSI-specific paths (env) ─────────────────────────────────────────────── STD_FLOWS = os.environ.get('RSI_STD_FLOWS', 'output/rsi/standard_flows.jsonl') -TEMPLATE = os.environ.get('RSI_TEMPLATE', 'Qwen3_5Template') -LORA_SYNC_DIR = os.environ.get('RSI_LORA_SYNC', 'output/rsi/lora_sync') -ADAPTER_NAME = os.environ.get('RSI_ADAPTER', 'executor') +TEMPLATE = os.environ.get('RSI_TEMPLATE', 'Template') # base text template for Qwen3-4B (text-only) REWARD_TOOL_RESULT = 'tool_result' # matches rsi_refine.attach_reward_method +REWARD_RUBRIC = 'rubric' +# Tests for code rounds, keyed by the flow's id (rsi_refine passes the id through +# from the step-1 subset, which carries the dataset's own tests). +TESTS_PATH = os.environ.get('RSI_TESTS', '') +TEST_TIMEOUT = int(os.environ.get('RSI_TEST_TIMEOUT', 30)) +JUDGE_WORKERS = int(os.environ.get('RSI_JUDGE_WORKERS', max(24, min(96, (os.cpu_count() or 24) // 2)))) +# Tool-round scoring: 'match' compares the call literally, 'rubric' asks a judge +# model whether it means the same as the recorded call. +TOOL_REWARD = os.environ.get('RSI_TOOL_REWARD', 'match') +JUDGE_MODEL = os.environ.get('RSI_JUDGE_MODEL', 'qwen3.8-max') +JUDGE_BASE_URL = os.environ.get('LLM_BACKUP_BASE_URL', '') +JUDGE_API_KEY = os.environ.get('LLM_BACKUP_API_KEY', '') +# 16 workers is what the offline judging run used; each step needs one call per +# rollout that actually produced a tool call. +RUBRIC_WORKERS = int(os.environ.get('RSI_RUBRIC_WORKERS', 16)) +RUBRIC_RETRIES = int(os.environ.get('RSI_RUBRIC_RETRIES', 3)) +# Keep only the first N trainable rounds, in file order (no shuffle anywhere). +MAX_ROUNDS = int(os.environ.get('RSI_MAX_ROUNDS', 0)) +RUN_NAME = os.environ.get('RSI_RUN_NAME', '') +# Optional per-sample reward audit: when set, every scored rollout is appended as +# one jsonl line (step, kind, ref/gen call, completion head, score, judge reason). +# Pure observability; the reward and training path are untouched. +REWARD_DUMP = os.environ.get('RSI_REWARD_DUMP', '') +# Raw step-1 conversations (before rsi_refine). rsi_refine's flow schema keeps +# only the FIRST user message (as `query`) plus the tool rounds, so any parameter +# the user stated in a LATER user turn is missing from a round's prompt and the +# model is asked to produce a call it cannot possibly know. When this points at +# the raw file, each round's prompt is rebuilt to splice those dropped user (and +# assistant clarification) turns back in, joined to the raw conversation by the +# first user message (unique for ~99.6% of flows); flows that cannot be joined +# fall back to the flow-only prompt. Empty reproduces the old flow-only behavior. +RAW_MESSAGES = os.environ.get('RSI_RAW_MESSAGES', '') +# Diagnostic knobs for "does reward rise on a FIXED distribution". Training reward +# on the default sequential single-epoch feed cannot answer that: the flow file is +# ordered so later rounds carry more arguments (harder), so a falling curve mixes +# difficulty with capability. Set both to hold the distribution still: +# RSI_SHUFFLE_SEED shuffle the rounds once with this seed (empty = file order) +# RSI_POOL_SIZE keep only this many rounds and repeat them until MAX_ROUNDS +# is filled, re-shuffling each pass (0 = no repetition) +# With a pool the same questions are seen every pass, so reward must climb unless +# the optimizer itself is at fault. +SHUFFLE_SEED = os.environ.get('RSI_SHUFFLE_SEED', '') +POOL_SIZE = int(os.environ.get('RSI_POOL_SIZE', 0)) + +# ── solver learning mode (RSI step-3 subclass) ───────────────────────────── +SOLVER_MODE = os.environ.get('RSI_SOLVER_MODE', 'grpo').lower() +if SOLVER_MODE not in ('grpo', 'opsd'): + raise ValueError(f"RSI_SOLVER_MODE must be 'grpo' or 'opsd', got {SOLVER_MODE!r}") +# GRPO subclass: on a failed code round, inject the sandbox error as a tool +# message and let the model retry until it passes or the turn budget is spent. +SOLVER_MAX_TURNS = int(os.environ.get('RSI_SOLVER_MAX_TURNS', 2)) +# OPSD subclass: the teacher sees one extra system message carrying the +# challenger's passing reference solution ('{solution}' is filled per sample). +OPSD_TEACHER_SYS = os.environ.get( + 'RSI_OPSD_TEACHER_SYS', + 'A correct reference solution is provided to guide you:\n' + '```python\n{solution}\n```\n' + 'Study it, then produce your own complete, correct solution to the task above.') +# k3 direction guard, exposed so the divergence sign can be swapped without +# touching call sites (see opsd.py `reverse`). +OPSD_REVERSE = os.environ.get('RSI_OPSD_REVERSE', '1') == '1' +# Tolerance for the OPSD first-batch self-check that pins the response-logps +# frame against the sampler's known-correct old_logps (mean abs diff per token). +OPSD_SELFCHECK_TOL = float(os.environ.get('RSI_OPSD_SELFCHECK_TOL', 0.5)) # ── standard CLI knobs (same shape as the reference script) ──────────────── MODEL_ID = args.model.model_id or 'ms://Qwen/Qwen3.6-35B-A3B' MODEL_GPUS = args.infra.model_gpus or 4 -SAMPLER_GPUS = args.infra.sampler_gpus or 2 -SAMPLER_TP = args.sampler.tensor_parallel_size or 2 -NUM_GPUS = MODEL_GPUS + SAMPLER_GPUS +SAMPLER_GPUS = args.infra.sampler_gpus or 4 +# KL anchor. GRPOLoss adds beta * KL(pi || ref) per response token only when BOTH +# beta > 0 AND ref_logps are passed to forward_backward (grpo.py:315), so leaving +# either at its default silently trains without any anchor. +# +# The anchor model gets its OWN gpus rather than sharing the trainer's: Megatron +# calls mpu.initialize_model_parallel() unconditionally per process +# (model/megatron/strategy/megatron.py:115), so constructing a second +# MegatronModel on the same ranks would re-init the process-global parallel +# state. REF_GPUS > 0 appends a separate device group. +KL_BETA = float(os.environ.get('RSI_KL_BETA', 0.0)) +REF_GPUS = int(os.environ.get('RSI_REF_GPUS', 0)) +# Stays the ORIGINAL base weights for every self-play iteration, while MODEL_ID +# advances to the previous iteration's checkpoint — that is the point of the +# anchor: it bounds the drift accumulated across iterations, not just within one. +REF_MODEL_ID = os.environ.get('RSI_REF_MODEL_ID', 'ms://Qwen/Qwen3-4B') +NUM_GPUS = MODEL_GPUS + SAMPLER_GPUS + REF_GPUS +# Which GRPO-family aggregation to use. 'GRPOLoss' normalizes each sequence by +# its OWN token count (grpo.py:132), which leaves a per-group residual +# sum_i(A_i / len_i); with group-centred advantages (sum_i A_i = 0) that residual +# is zero only if all lengths are equal, and it measured +0.16..+0.57 in +# equivalent-advantage terms across the 17 self-play iterations (passing rollouts +# were consistently the shorter ones), i.e. a standing push toward shorter output. +# 'DRGRPOLoss' divides by batch * max_completion_length, a constant, so the same +# residual becomes sum_i(A_i)/const == 0. +LOSS_NAME = os.environ.get('RSI_LOSS', 'GRPOLoss') NUM_GENERATIONS = args.rl.num_generations or 8 MAX_NEW_TOKENS = args.sampling.max_tokens or 4096 LEARNING_RATE = args.optimizer.learning_rate or 5e-5 @@ -73,13 +203,39 @@ MICRO_BATCH_SIZE = args.training.micro_batch_size or 1 GRADIENT_ACCUMULATION_STEPS = args.training.gradient_accumulation_steps or 1 SAVE_STEPS = args.training.save_steps or 1000 -LORA_RANK = args.lora.lora_r or 16 +# Context window (prompt + generation). Raised above the old hard-coded 8192 so a +# larger --max-tokens cannot overflow the engine: RSI code prompts can be long, +# and on a policy collapse the model rambles to the token cap. Keep this >= +# max_tokens + longest prompt. Matches the challenger's MAX_MODEL_LEN default. +MAX_MODEL_LEN = int(os.environ.get('RSI_MAX_MODEL_LEN', 16384)) +# Where the trained weights land. Kept per-iteration-configurable so a self-play +# loop can point each iteration's checkpoint at its own dir (and feed the final +# HF-format dir back to the next challenge/rl via --model-id / RSI_CH_MODEL). +SAVE_DIR = os.environ.get('RSI_SAVE_DIR', 'output') +SAVE_NAME = os.environ.get('RSI_SAVE_NAME', 'rsi-executor-final') import swanlab -swanlab.init(project='twinkle-rsi') +swanlab.init(project='twinkle-rsi', experiment_name=RUN_NAME or None) # ── tool-call matching (name exact + standard-call arg subset) ───────────── +def _json_safe(o: Any) -> Any: + """Coerce anything json.dumps cannot handle into a string, recursively. + + A malformed rollout can parse into arguments holding a bare ``...`` (Python + Ellipsis) or other non-JSON objects; without this a single such sample makes + json.dumps raise and takes the whole run down. Normal dicts/lists/scalars are + returned unchanged, so well-formed calls serialize exactly as before. + """ + if isinstance(o, dict): + return {str(k): _json_safe(v) for k, v in o.items()} + if isinstance(o, (list, tuple)): + return [_json_safe(x) for x in o] + if isinstance(o, (str, int, float, bool)) or o is None: + return o + return str(o) + + def _as_args(a: Any) -> Dict[str, Any]: if isinstance(a, str): try: @@ -101,6 +257,104 @@ def tool_call_matches(gen_call: Optional[Dict[str, Any]], ref_call: Dict[str, An return True +# ── rubric reward for tool rounds (judge model over the API) ─────────────── +# Verbatim from output/rsi/rubric_judge.py, which produced the offline 6.2% -> 15.0% +# comparison; changing a word here means this run is no longer that measurement. +JUDGE_SYSTEM = """\ +You are a strict tool-call equivalence judge. You will be given: +1. The user's request (what they asked for). +2. The STANDARD tool call (the reference answer: function name + arguments). +3. The MODEL's output (what the model actually produced). + +Your job: decide whether the model's output is semantically equivalent to the standard call. + +Rules: +- The model MUST have attempted a tool/function call. If it only gave a natural language + answer without any call, score 0. +- Function name must match (case-insensitive, ignore spacing differences). +- Arguments: check SEMANTIC equivalence, not exact string match. + * Search queries: "Oscars cinema drama" ≈ "Oscars newest cinema drama" (same intent) → OK + * Numbers: "5" = 5 = 5.0 → OK + * Coordinates/measurements that point to the same place or value → OK + * Lists with same elements in different order → OK + * Completely different values → NOT OK +- EXTRA arguments the model added that the standard call omits do NOT count against it, + as long as they do not contradict the user's request. Spelling out an optional + parameter at its default value (e.g. output="json" when json is the default) is + fully equivalent to leaving it out → still score 1.0 +- Only the arguments present in the STANDARD call have to be matched. +- If the function is correct and every standard argument is semantically equivalent → 1.0 +- If the function is correct but a standard argument is partially wrong, or a required + one is missing → 0.5 +- If wrong function, no call at all, or standard arguments completely wrong → 0.0 + +Output ONLY a JSON object: {"score": <0.0 or 0.5 or 1.0>, "reason": "<one sentence>"} +Nothing else. +""" + +_SCORE_RE = re.compile(r'"score"\s*:\s*([\d.]+)') +# The judge sometimes replies with a bare number or "Score: 0" instead of JSON. Offline +# every such reply was a genuine 0, so read it rather than throwing the sample away. +_BARE_SCORE_RE = re.compile(r'(?:score\D{0,12})?\b(0(?:\.0)?|0\.5|1(?:\.0)?)\b', re.I) + +_judge_client = None + + +def judge_client(): + global _judge_client + if _judge_client is None: + from openai import OpenAI + if not JUDGE_API_KEY: + raise RuntimeError('RSI_TOOL_REWARD=rubric needs LLM_BACKUP_API_KEY ' + '(and LLM_BACKUP_BASE_URL) in the environment') + _judge_client = OpenAI(base_url=JUDGE_BASE_URL or None, api_key=JUDGE_API_KEY) + return _judge_client + + +def judge_input(completion: str, gen_call: Dict[str, Any]) -> str: + """What the judge sees as "the model's output". + + The template already lifted the call out of the raw text into a structured + field, so the raw text cannot be recovered: the body the model wrote comes + first, then the call that was parsed out of it. + """ + call = {'name': gen_call.get('name'), 'arguments': _as_args(gen_call.get('arguments'))} + return f'{completion}\n\n[parsed tool call]\n{json.dumps(_json_safe(call), ensure_ascii=False)}' + + +def judge_rubric(ref_call: Dict[str, Any], model_output: str) -> Tuple[Optional[float], Optional[str]]: + """Ask the judge for one score; (None, reason) when it never answered. + + None is not zero: the caller drops the sample from its group instead of + counting it as a miss, so a timeout cannot masquerade as a wrong answer. + The second element is the judge's raw reply, kept only for the audit dump. + """ + if len(model_output) > 3000: + model_output = model_output[:1500] + '\n...[truncated]...\n' + model_output[-1500:] + user_msg = (f'## Standard tool call (reference answer)\n```json\n' + f'{json.dumps(ref_call, ensure_ascii=False)}\n```\n\n' + f"## Model's full output\n```\n{model_output}\n```\n\nScore the model's output.") + for attempt in range(RUBRIC_RETRIES): + try: + resp = judge_client().chat.completions.create( + model=JUDGE_MODEL, + messages=[{'role': 'system', 'content': JUDGE_SYSTEM}, + {'role': 'user', 'content': user_msg}], + temperature=0.0, + max_tokens=200, + ) + text = resp.choices[0].message.content or '' + m = (_SCORE_RE.search(text) or _BARE_SCORE_RE.match(text.strip()) + or _BARE_SCORE_RE.search(text[:40])) + return (float(m.group(1)) if m else None), text + except Exception as e: # network / rate limit / bad gateway + if attempt == RUBRIC_RETRIES - 1: + logger.warning(f'[rsi_rl] judge gave up after {RUBRIC_RETRIES} tries: {str(e)[:160]}') + return None, f'error:{str(e)[:160]}' + time.sleep(2**attempt) + return None, None + + class ToolMatchReward(Reward): """1.0 when the generated tool call matches the recorded standard call.""" @@ -126,6 +380,229 @@ def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: return rewards +# ── code-round execution reward ──────────────────────────────────────────── +# Same sandbox contract as cookbook/rl/grpo/mbpp_grpo.py, which was checked +# against all 974 MBPP reference solutions (974/974 pass): the generated code, +# the setup code and the asserts are concatenated into one file and executed, so +# a bare ``assert fn(...) == x`` resolves the function by name. +_FENCE_RE = re.compile(r'```(?:python|py)?\s*\n(.*?)```', re.S) + + +def extract_code(text: str) -> str: + """Take the last fenced block; fall back to the whole body when unfenced.""" + idx = (text or '').rfind('</think>') + body = text[idx + len('</think>'):] if idx >= 0 else (text or '') + blocks = _FENCE_RE.findall(body) + return (blocks[-1] if blocks else body).strip() + + +def run_asserts(code: str, setup: str, asserts: List[str], timeout: int = TEST_TIMEOUT) -> bool: + """True when every assert passes. Thin wrapper over run_asserts_verbose.""" + return run_asserts_verbose(code, setup, asserts, timeout)[0] + + +def run_asserts_verbose(code: str, setup: str, asserts: List[str], + timeout: int = TEST_TIMEOUT) -> Tuple[bool, str]: + """Run code+setup+asserts and return (passed, stderr_text). + + Same sandbox contract as the MBPP-verified path (start_new_session + killpg + so a forking solution leaves no stray processes; RLIMIT_AS caps the child at + 2GB). stderr is captured (not sent to /dev/null) so the GRPO continuation can + feed the actual traceback back to the model as a tool message. ``passed`` is + exactly ``returncode == 0``, identical to the old bool-only behavior. + """ + if not code.strip() or not asserts: + return False, 'no code was produced' + parts = [code] + if (setup or '').strip(): + parts.append(setup) + parts.extend(asserts) + tmp = tempfile.mkdtemp(prefix='rsi_code_') + try: + with open(os.path.join(tmp, '_run.py'), 'w', encoding='utf-8') as f: + f.write('\n\n'.join(parts) + '\n') + env = dict(os.environ, MPLBACKEND='Agg', PYTHONHASHSEED='0', OMP_NUM_THREADS='1', + MKL_NUM_THREADS='1', TOKENIZERS_PARALLELISM='false') + env.pop('CUDA_VISIBLE_DEVICES', None) + + def _limit(): + resource.setrlimit(resource.RLIMIT_AS, (2 * 1024**3, 2 * 1024**3)) + + proc = subprocess.Popen([sys.executable, '_run.py'], cwd=tmp, env=env, + stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, + text=True, start_new_session=True, preexec_fn=_limit) + try: + _, err = proc.communicate(timeout=timeout) + return proc.returncode == 0, (err or '') + except subprocess.TimeoutExpired: + try: + os.killpg(proc.pid, signal.SIGKILL) + except ProcessLookupError: + pass + try: + proc.communicate(timeout=5) + except Exception: + pass + return False, f'execution timed out after {timeout}s (possible infinite loop)' + finally: + shutil.rmtree(tmp, ignore_errors=True) + + +def load_tests() -> Dict[str, Dict[str, Any]]: + """Read the tests file into {id: {asserts, setup}} (empty when not configured).""" + if not TESTS_PATH or not os.path.exists(TESTS_PATH): + return {} + tests: Dict[str, Dict[str, Any]] = {} + with open(TESTS_PATH, encoding='utf-8') as f: + for line in f: + line = line.strip() + if not line: + continue + row = json.loads(line) + asserts = row.get('test_list') or [] + if isinstance(asserts, str): + asserts = json.loads(asserts) + if asserts: + tests[str(row.get('id'))] = {'asserts': list(asserts), + 'setup': row.get('test_setup_code') or ''} + return tests + + +class RoundReward(Reward): + """Score each rollout by what its round is: tool call, or code execution. + + Which branch applies is carried per-sample in ``user_data``: a tool round + rides ``ref_tool_call``, a code round rides ``code_tests``. Execution runs in + a thread pool because every verdict is a separate subprocess; rubric judging + runs in a thread pool because every verdict is a separate API call. + + A returned ``None`` means "never scored" (the judge never answered after its + retries) and is not the same as 0.0 -- see ``group_advantages``. + """ + + def __init__(self): + # Per-step counters, read by main() for logging. + self.stats: Dict[str, int] = {} + # Per-sample audit rows for the latest __call__, aligned by index with the + # returned rewards. main() stamps them with the step and writes the dump. + self.records: List[Dict[str, Any]] = [] + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[Optional[float]]: + rewards: List[Optional[float]] = [0.0] * len(trajectories) + code_jobs: List[Tuple[int, str, Dict[str, Any]]] = [] + rubric_jobs: List[Tuple[int, Dict[str, Any], str]] = [] + n_no_call = 0 + # One audit row per sample; branch/score/reason filled in as we go. + recs: List[Dict[str, Any]] = [{'kind': None, 'ref_call': None, 'gen_call': None, + 'completion': '', 'reason': None} for _ in trajectories] + + for i, traj in enumerate(trajectories): + ud = {item[0]: item[1] for item in (traj.get('user_data') or [])} + completion = '' + gen_call = None + for m in reversed(traj.get('messages', []) or []): + if m.get('role') == 'assistant': + tcs = m.get('tool_calls') or [] + if tcs: + gen_call = tcs[0].get('function') + completion = m.get('content', '') or '' + break + # Full completion, never truncated: a cut tail once hid whether the + # model emitted </think> / <tool_call>, which is exactly what the audit + # must answer. Store the whole thing. + recs[i]['completion'] = completion + recs[i]['gen_call'] = gen_call + + if 'ref_tool_call' in ud: + try: + ref_call = json.loads(ud['ref_tool_call']) + except (ValueError, TypeError): + ref_call = None + recs[i]['ref_call'] = ref_call + if TOOL_REWARD != 'rubric': + rewards[i] = 1.0 if (ref_call and tool_call_matches(gen_call, ref_call)) else 0.0 + recs[i]['kind'] = 'tool_match' + elif not gen_call: + # No call was parsed out, so there is nothing for the judge to + # compare against: 0 without spending a request. + n_no_call += 1 + recs[i]['kind'] = 'tool_no_call' + elif ref_call: + rubric_jobs.append((i, ref_call, judge_input(completion, gen_call))) + recs[i]['kind'] = 'tool_rubric' + elif 'code_tests' in ud: + try: + spec = json.loads(ud['code_tests']) + except (ValueError, TypeError): + continue + code_jobs.append((i, extract_code(completion), spec)) + recs[i]['kind'] = 'code' + + if code_jobs: + # Judge each distinct (task, code) once: identical completions are common. + uniq: Dict[Tuple[str, str], Dict[str, Any]] = {} + for _, code, spec in code_jobs: + uniq.setdefault((str(spec.get('id')), code), spec) + todo = list(uniq) + with ThreadPoolExecutor(max_workers=max(1, min(JUDGE_WORKERS, len(todo)))) as ex: + verdicts = dict(zip(todo, ex.map( + lambda k: run_asserts(k[1], uniq[k]['setup'], uniq[k]['asserts']), todo))) + for i, code, spec in code_jobs: + rewards[i] = 1.0 if verdicts.get((str(spec.get('id')), code)) else 0.0 + + n_failed = 0 + if rubric_jobs: + # One request per distinct (reference, output) pair; the judge runs at + # temperature 0, so repeats would only cost money. + uniq_r: Dict[Tuple[str, str], Dict[str, Any]] = {} + for _, ref_call, text in rubric_jobs: + uniq_r.setdefault((json.dumps(ref_call, ensure_ascii=False), text), ref_call) + keys = list(uniq_r) + with ThreadPoolExecutor(max_workers=max(1, min(RUBRIC_WORKERS, len(keys)))) as ex: + scores = dict(zip(keys, ex.map(lambda k: judge_rubric(uniq_r[k], k[1]), keys))) + for i, ref_call, text in rubric_jobs: + s, reason = scores.get((json.dumps(ref_call, ensure_ascii=False), text), (None, None)) + rewards[i] = s + recs[i]['reason'] = reason + if s is None: + n_failed += 1 + + for i in range(len(trajectories)): + recs[i]['score'] = rewards[i] + self.records = recs + self.stats = {'no_call': n_no_call, 'judged': len(rubric_jobs) - n_failed, + 'judge_failed': n_failed, 'executed': len(code_jobs)} + return rewards + + +def group_advantages(rewards: List[Optional[float]], num_generations: int, + scale: str = 'group') -> List[float]: + """GRPOAdvantage, but unscored samples stay out of their group's statistics. + + Same formula as twinkle.advantage.GRPOAdvantage (subtract the group mean, + divide by the group's unbiased std). A ``None`` reward means the judge never + returned a verdict: it is left out of the mean and std and gets advantage 0, + so it pushes the policy in neither direction. A group left with a single + verdict has no baseline to compare against, so all of it gets 0. + """ + if all(r is not None for r in rewards): + return GRPOAdvantage()(rewards, num_generations=num_generations, scale=scale).tolist() + import torch + vals = torch.tensor([0.0 if r is None else r for r in rewards], dtype=torch.float32) + mask = torch.tensor([r is not None for r in rewards], dtype=torch.float32) + g_vals = vals.view(-1, num_generations) + g_mask = mask.view(-1, num_generations) + n = g_mask.sum(dim=1, keepdim=True) + mean = (g_vals * g_mask).sum(dim=1, keepdim=True) / n.clamp(min=1) + adv = (g_vals - mean) * g_mask + if scale == 'group': + var = ((g_vals - mean)**2 * g_mask).sum(dim=1, keepdim=True) / (n - 1).clamp(min=1) + adv = adv / (var.sqrt() + 1e-8) + elif scale == 'batch': + adv = adv / (adv[g_mask.bool()].std() + 1e-8) + return (adv * (n > 1).float()).view(-1).tolist() + + # ── decompose standard flows into per-round training trajectories ────────── def _openai_tool_call(call: Dict[str, Any], idx: int) -> Dict[str, Any]: args_ = call.get('arguments', {}) @@ -155,8 +632,100 @@ def _render_prior_round(r: Dict[str, Any], idx: int) -> List[Dict[str, Any]]: ] -def build_round_trajectories(records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: - """One training trajectory per TOOL round; prior rounds become fixed context.""" +def _load_raw_by_query(path: str) -> Dict[str, List[Dict[str, Any]]]: + """Index the raw step-1 conversations by their first user message, the join + key back to a flow's ``query``. Only keys that map to exactly ONE conversation + are kept, so an ambiguous first question never pulls the wrong conversation. + """ + if not path or not os.path.exists(path): + return {} + seen: Dict[str, List[Dict[str, Any]]] = {} + dup: set = set() + with open(path, encoding='utf-8') as f: + for line in f: + line = line.strip() + if not line: + continue + msgs = json.loads(line).get('messages') or [] + fu = next((str(m.get('content') or '').strip() + for m in msgs if m.get('role') == 'user'), '') + if not fu: + continue + if fu in seen: + dup.add(fu) + else: + seen[fu] = msgs + for k in dup: + seen.pop(k, None) + return seen + + +def _locate_calls(raw_msgs: List[Dict[str, Any]], + rounds: List[Dict[str, Any]]) -> Optional[List[int]]: + """For each round, the index of the raw assistant message that made its call, + matched forward by the round's tool name (robust to ToolACE's several call + syntaxes: ``Name(...)``, ``[Name]=>``, ``{Name}=>`` ...). A code/nameless + round reuses the running cursor. Returns None if any named call is not found + in order, so the caller falls back to the flow-only prompt. + """ + locs: List[int] = [] + cur = 0 + for r in rounds: + name = ((r.get('tool_call') or {}).get('name')) or '' + if not name: + locs.append(cur) + continue + found = -1 + for j in range(cur, len(raw_msgs)): + m = raw_msgs[j] + if m.get('role') == 'assistant' and name in (m.get('content') or ''): + found = j + break + if found < 0: + return None + locs.append(found) + cur = found # a later parallel call may live in the same message + return locs + + +def _intervening_turns(raw_msgs: List[Dict[str, Any]], lo: int, hi: int) -> List[Dict[str, Any]]: + """User turns and assistant clarification turns in ``raw_msgs[lo+1:hi]``. + + Assistant tool-call messages (content starting with ``[``) and tool results + are dropped here because the structured prior rounds already carry the call + and its result; what is recovered is exactly the conversational turns + rsi_refine did not keep. + """ + out: List[Dict[str, Any]] = [] + for j in range(lo + 1, hi): + m = raw_msgs[j] + role = m.get('role') + content = m.get('content') or '' + if role == 'user': + out.append({'role': 'user', 'content': content}) + elif role == 'assistant' and content.strip() and not content.lstrip().startswith('['): + out.append({'role': 'assistant', 'content': content}) + return out + + +def build_round_trajectories(records: List[Dict[str, Any]], + tests: Optional[Dict[str, Dict[str, Any]]] = None, + raw_by_query: Optional[Dict[str, List[Dict[str, Any]]]] = None, + recover_stats: Optional[Dict[str, int]] = None) -> List[Dict[str, Any]]: + """One training trajectory per trainable round; prior rounds become fixed context. + + A tool round is trainable when it has a recorded call to match. A code round + is trainable only when ``tests`` holds asserts for the record, since that is + what its reward executes; otherwise it stays context-only. + + When ``raw_by_query`` is given, the round prompt is rebuilt from the raw + conversation so the user turns that rsi_refine dropped (e.g. the turn that + states the call's arguments) are spliced back in at their real positions; + flows whose raw conversation cannot be located fall back to the flow-only + prompt (system + first query + prior rounds). + """ + tests = tests or {} + raw_by_query = raw_by_query or {} trajs: List[Dict[str, Any]] = [] for rec in records: prefix: List[Dict[str, Any]] = [] @@ -166,81 +735,482 @@ def build_round_trajectories(records: List[Dict[str, Any]]) -> List[Dict[str, An prefix.append(rec['query']) tools = rec.get('tools') or [] rounds = rec.get('rounds') or [] + rec_id = str(rec.get('id')) + spec = tests.get(rec_id) + + # Try to recover the dropped user turns from the raw conversation. + raw = None + locs = None + first_user_idx = 0 + if raw_by_query: + q = rec.get('query') or {} + qtext = str(q.get('content') if isinstance(q, dict) else q).strip() + raw = raw_by_query.get(qtext) + if raw is not None: + locs = _locate_calls(raw, rounds) + first_user_idx = next((j for j, m in enumerate(raw) + if m.get('role') == 'user'), 0) + if recover_stats is not None: + key = 'recovered' if locs is not None else ('unjoined' if raw is None else 'unlocatable') + recover_stats[key] = recover_stats.get(key, 0) + 1 + for i, r in enumerate(rounds): - if r.get('reward_method') != REWARD_TOOL_RESULT or not r.get('tool_call'): - continue # v1: train tool rounds only + if r.get('reward_method') == REWARD_TOOL_RESULT and r.get('tool_call'): + user_data = [('ref_tool_call', json.dumps(r['tool_call'], ensure_ascii=False))] + elif spec and not r.get('tool_call'): + user_data = [('code_tests', json.dumps({'id': rec_id, **spec}, ensure_ascii=False))] + # Carry the challenger's passing solution so OPSD can build the + # teacher's privileged prompt; harmless/unused in GRPO mode. + if r.get('code'): + user_data.append(('ref_solution', r['code'])) + else: + continue messages = list(prefix) - for j in range(i): - messages.extend(_render_prior_round(rounds[j], j)) - trajs.append({ - 'messages': messages, - 'tools': tools, - 'user_data': [('ref_tool_call', json.dumps(r['tool_call'], ensure_ascii=False))], - }) + if locs is not None: + anchor = first_user_idx + for j in range(i): + messages.extend(_intervening_turns(raw, anchor, locs[j])) + messages.extend(_render_prior_round(rounds[j], j)) + anchor = locs[j] + messages.extend(_intervening_turns(raw, anchor, locs[i])) + else: + for j in range(i): + messages.extend(_render_prior_round(rounds[j], j)) + trajs.append({'messages': messages, 'tools': tools, 'user_data': user_data}) return trajs def create_rsi_dataset(): records = Dataset(DatasetMeta(dataset_id=STD_FLOWS)).dataset.to_list() - trajs = build_round_trajectories(records) - logger.info(f'[rsi_rl] {len(records)} standard flows -> {len(trajs)} per-round tool queries') + tests = load_tests() + raw_by_query = _load_raw_by_query(RAW_MESSAGES) + recover_stats: Dict[str, int] = {} + trajs = build_round_trajectories(records, tests, raw_by_query, recover_stats) + if raw_by_query: + rec_n = recover_stats.get('recovered', 0) + logger.info(f'[rsi_rl] raw-turn recovery from {RAW_MESSAGES}: ' + f"recovered={rec_n} " + f"unjoined={recover_stats.get('unjoined', 0)} " + f"unlocatable={recover_stats.get('unlocatable', 0)} " + f'of {len(records)} flows ({rec_n / max(len(records), 1):.1%} rebuilt)') + else: + logger.info('[rsi_rl] RSI_RAW_MESSAGES unset: using flow-only prompts ' + '(dropped user turns are NOT recovered)') + if SHUFFLE_SEED: + random.Random(int(SHUFFLE_SEED)).shuffle(trajs) + logger.info(f'[rsi_rl] shuffled {len(trajs)} rounds with seed {SHUFFLE_SEED} ' + '(difficulty no longer correlates with step)') + if POOL_SIZE and len(trajs) > POOL_SIZE: + pool = trajs[:POOL_SIZE] + target = MAX_ROUNDS or len(trajs) + rng = random.Random(int(SHUFFLE_SEED) if SHUFFLE_SEED else 0) + repeated: List[Dict[str, Any]] = [] + while len(repeated) < target: + one_pass = list(pool) + rng.shuffle(one_pass) + repeated.extend(one_pass) + trajs = repeated[:target] + logger.info(f'[rsi_rl] fixed pool of {POOL_SIZE} rounds repeated to {len(trajs)} ' + f'({len(trajs) / POOL_SIZE:.1f} passes): the question distribution is ' + 'now identical across steps') + if MAX_ROUNDS and len(trajs) > MAX_ROUNDS: + # File order, no shuffle: one optim step consumes one round, so N rounds + # is exactly N steps of a single epoch. + logger.info(f'[rsi_rl] keeping the first {MAX_ROUNDS} of {len(trajs)} trainable rounds') + trajs = trajs[:MAX_ROUNDS] + n_code = sum(1 for t in trajs if t['user_data'][0][0] == 'code_tests') + logger.info(f'[rsi_rl] {len(records)} standard flows -> {len(trajs)} per-round queries ' + f'({len(trajs) - n_code} tool / {n_code} code); tests loaded: {len(tests)}; ' + f'tool reward: {TOOL_REWARD}' + + (f' via {JUDGE_MODEL}' if TOOL_REWARD == 'rubric' else '')) + if not trajs: + raise RuntimeError( + 'no trainable rounds: tool rounds need a recorded tool_call and code rounds ' + f'need asserts via RSI_TESTS (currently {TESTS_PATH!r})') dataset = Dataset(DatasetMeta(data=trajs)) # enable_thinking=True: we train the reasoning that precedes the tool call. - dataset.set_template(TEMPLATE, model_id=MODEL_ID, max_length=8192, + dataset.set_template(TEMPLATE, model_id=MODEL_ID, max_length=MAX_MODEL_LEN, truncation_strategy='delete', enable_thinking=True) dataset.encode(add_generation_prompt=True) return dataset +# ── solver-mode helpers (GRPO continuation + OPSD teacher forward) ────────── +def make_local_template(): + """A driver-side Template instance for token surgery (bridge / concat). + + The model and sampler each hold their own remote template; bridge and + teacher-prompt construction happen on the driver, so we need a local one. + Mirrors cookbook/rl/multi_turn/multi_turn_grpo.py's rollout_template. + """ + import twinkle.template as _tm + cls = getattr(_tm, TEMPLATE, None) + if cls is None: + raise ValueError(f'template class {TEMPLATE!r} not found in twinkle.template') + t = cls(MODEL_ID, max_length=MAX_MODEL_LEN, enable_thinking=True) + t.truncation_strategy = 'delete' + return t + + +def _last_assistant_text(pif: Dict[str, Any]) -> str: + for m in reversed(pif.get('messages') or []): + if m.get('role') == 'assistant': + return m.get('content', '') or '' + return '' + + +def _format_exec_error(err: str) -> str: + """Turn captured stderr into the tool message shown back to the model.""" + err = (err or '').strip() or 'Your code did not pass the tests (no error output captured).' + if len(err) > 1500: + err = err[:700] + '\n...[truncated]...\n' + err[-700:] + return ('Your solution failed when executed against the tests:\n' + f'{err}\n\n' + 'Fix the bug and reply with the complete corrected solution in a single ' + '```python code block.') + + +def _bridge_tool_message(template, pif: Dict[str, Any], tool_content: str) -> Optional[Dict[str, Any]]: + """Append a {'role':'tool'} turn + next generation prompt as -100 bridge. + + Computed entirely in template space (render-after minus render-before), so + history tokens stay byte-for-byte in ``input_ids`` and only the new tool turn + is tokenized from canonical template output -- never decode-then-re-encode. + Mirrors Template.concat_input_feature / MultiTurnRollout._extend_with_bridge. + Returns the extended pif, or None if it would exceed the template max_length. + """ + import copy + tok = template.tokenizer + messages_before = list(pif.get('messages') or []) + messages_after = messages_before + [{'role': 'tool', 'content': tool_content}] + et = getattr(template, 'enable_thinking', False) + s_before = tok.apply_chat_template(messages_before, tokenize=False, + add_generation_prompt=False, enable_thinking=et) + s_after = tok.apply_chat_template(messages_after, tokenize=False, + add_generation_prompt=True, enable_thinking=et) + # SEAM: the vLLM pif ends at the assistant's closing <|im_end|> with NO trailing + # newline (generation stops at the eos token), but the canonical render puts a + # "\n" right after that <|im_end|>. Splitting at len(s_before) would drop that + # "\n" and append the tool turn directly onto <|im_end|>, producing a malformed + # "<|im_end|><|im_start|>" boundary that the trained turn-2 tokens then condition + # on. Split right AFTER the assistant's <|im_end|> so the bridge carries the + # "\n" + tool turn and reproduces the canonical tokenization exactly. + marker = '<|im_end|>' + cut = s_before.rfind(marker) + if cut < 0: + raise RuntimeError('tool bridge: no <|im_end|> found in the rendered history; ' + 'cannot locate the assistant turn boundary.') + head = s_before[:cut + len(marker)] + if not s_after.startswith(head): + raise RuntimeError('tool bridge: chat template is not monotonic in the message list; ' + 'cannot append a tool turn as a suffix.') + bridge_text = s_after[len(head):] + bridge_ids = tok.encode(bridge_text, add_special_tokens=False) + if not bridge_ids: + raise RuntimeError('tool bridge tokenized to an empty id list') + result = copy.deepcopy(pif) + input_ids = list(result['input_ids']) + labels = list(result.get('labels') or []) + if labels: + if len(labels) != len(input_ids): + raise RuntimeError('tool bridge: labels/input_ids length mismatch') + labels = labels[-1:] + labels[:-1] # unroll to input order (mirror concat_input_feature) + else: + labels = [-100] * len(input_ids) + result['input_ids'] = input_ids + bridge_ids + result['labels'] = labels + [-100] * len(bridge_ids) + max_len = getattr(template, 'max_length', None) + if max_len and len(result['input_ids']) > max_len: + return None + new_if = template._invoke_post_pipeline([result])[0] + result.update(new_if) + result['messages'] = messages_after + return result + + +def grpo_continue(sampler, template, expand_prompts, sampling_params): + """GRPO rollout with error-feedback continuation for code rounds. + + Turn 1 samples every prompt. A code sample that FAILS its asserts (and did + not stop on 'length') gets the sandbox stderr injected as a {'role':'tool'} + message and is re-sampled, up to SOLVER_MAX_TURNS total turns. Tool rounds + and length-stopped samples are never continued. The returned per-sample + input feature is the full multi-turn trajectory (turn tokens trainable, tool + bridge -100) and old_logps is the concatenation of each turn's logprobs, so + the (#logps == #trainable labels) invariant holds for GRPO training. + """ + resps = sampler.sample(expand_prompts, sampling_params) + pifs: List[Dict[str, Any]] = [] + logps: List[List[float]] = [] + lens: List[int] = [] + stops: List[Optional[str]] = [] + for r in resps: + s = r.sequences[0] + pifs.append(s.new_input_feature) + logps.append([lp[0][1] for lp in s.logprobs]) + lens.append(len(s.tokens)) + stops.append(s.stop_reason) + + done = [False] * len(expand_prompts) + dm = getattr(sampler, 'device_mesh', None) + min_batch = dm.data_world_size if dm is not None else 1 + for _turn in range(2, SOLVER_MAX_TURNS + 1): + retry: List[int] = [] + for i, prompt in enumerate(expand_prompts): + if done[i]: + continue + ud = {item[0]: item[1] for item in (prompt.get('user_data') or [])} + if 'code_tests' not in ud or stops[i] == 'length': + done[i] = True + continue + try: + spec = json.loads(ud['code_tests']) + except (ValueError, TypeError): + done[i] = True + continue + code = extract_code(_last_assistant_text(pifs[i])) + passed, err = run_asserts_verbose(code, spec.get('setup', ''), spec.get('asserts', [])) + if passed: + done[i] = True + continue + # Bridge the tool error in. If the template can't append a tool turn as + # a clean suffix (e.g. a malformed/cut turn-1 without a proper </think>), + # skip continuation for THIS sample rather than crashing the whole step. + try: + bridged = _bridge_tool_message(template, pifs[i], _format_exec_error(err)) + except RuntimeError as e: + logger.warning(f'[rsi_rl][grpo] skip continuation for sample {i}: {e}') + bridged = None + if bridged is None: + done[i] = True + continue + pifs[i] = bridged + retry.append(i) + if not retry: + break + batch = [pifs[i] for i in retry] + if len(batch) < min_batch: + batch = batch + [batch[-1]] * (min_batch - len(batch)) + rresps = sampler.sample(batch, sampling_params)[:len(retry)] + for j, i in enumerate(retry): + s2 = rresps[j].sequences[0] + pifs[i] = s2.new_input_feature + logps[i].extend([lp[0][1] for lp in s2.logprobs]) + lens[i] += len(s2.tokens) + stops[i] = s2.stop_reason + + # Same invariant MultiTurnRollout enforces: one logp per trainable token. + for i, pif in enumerate(pifs): + trainable = sum(1 for lb in (pif.get('labels') or []) if lb != -100) + if len(logps[i]) != trainable: + raise RuntimeError(f'GRPO continuation logps/labels misaligned for sample {i}: ' + f'{len(logps[i])} logps vs {trainable} trainable labels') + return pifs, logps, lens + + +def _teacher_pif(template, student_pif: Dict[str, Any], ref_solution: str, + response_tokens: List[int]) -> Dict[str, Any]: + """Teacher input = student's query context + a privileged system message + carrying the reference solution, then the SAME student response tokens + concatenated verbatim (concat_input_feature, never re-encoded). Only the + prompt differs from the student; the scored response tokens are identical, + which is exactly what OPSD's per-token alignment requires. + """ + msgs = list(student_pif.get('messages') or []) + prompt_msgs = msgs[:-1] if (msgs and msgs[-1].get('role') == 'assistant') else list(msgs) + priv = {'role': 'system', 'content': OPSD_TEACHER_SYS.format(solution=ref_solution)} + insert_at = 1 if (prompt_msgs and prompt_msgs[0].get('role') == 'system') else 0 + teacher_msgs = prompt_msgs[:insert_at] + [priv] + prompt_msgs[insert_at:] + prompt_pif = template.encode({'messages': teacher_msgs}, add_generation_prompt=True) + return template.concat_input_feature(prompt_pif, list(response_tokens)) + + +def _as_rows(out_logps): + """Normalize forward_only's logps (a list of [mb, L] tensors OR a stacked + [N, L] tensor, depending on DP/microbatch config) into a flat per-sample + list of 1-D tensors, in input order. + """ + import torch + rows = [] + items = out_logps if isinstance(out_logps, list) else [out_logps] + for t in items: + if t is None: + continue + t = torch.as_tensor(t) + if t.dim() == 1: + rows.append(t) + elif t.dim() == 2: + rows.extend([t[i] for i in range(t.shape[0])]) + else: + raise RuntimeError(f'unexpected forward_only logps ndim={t.dim()}') + return rows + + +_OPSD_OFFSET: Optional[int] = None + + +def _extract_resp(row, seq_len: int, n: int, offset: int): + valid = row[:seq_len] + end = len(valid) - offset + return valid[end - n:end] + + +def _calibrate_opsd_offset(rows, pifs, old_logps_list): + """Pin the response-logps frame by matching a student self-forward against + the sampler's known-correct old_logps. Tries suffix offset 0 and 1 and picks + the one with the smallest mean|diff|; returns (offset, mean_abs_diff). + """ + import torch + best, best_err = None, float('inf') + for off in (0, 1): + errs = [] + ok = True + for row, pif, old in zip(rows, pifs, old_logps_list): + n = len(old) + if n == 0: + continue + resp = _extract_resp(row, len(pif['input_ids']), n, off) + if len(resp) != n: + ok = False + break + errs.append((resp.float() - torch.tensor(old, dtype=torch.float32)).abs().mean().item()) + if ok and errs: + m = sum(errs) / len(errs) + if m < best_err: + best_err, best = m, off + return best, best_err + + +def opsd_teacher_logps(model, template, student_pifs, response_tokens_list, + ref_solutions, student_old_logps): + """Per-sample response-only teacher log-probs for OPSDLoss (ragged lists). + + On the first call, calibrates the suffix offset by self-checking a student + forward against the sampler old_logps; if it cannot align within + OPSD_SELFCHECK_TOL it raises rather than feed a mis-framed teacher. + """ + global _OPSD_OFFSET + if _OPSD_OFFSET is None: + s_out = model.forward_only(inputs=list(student_pifs), micro_batch_size=MICRO_BATCH_SIZE) + off, err = _calibrate_opsd_offset(_as_rows(s_out.logps), student_pifs, student_old_logps) + if off is None or err > OPSD_SELFCHECK_TOL: + raise RuntimeError( + f'OPSD self-check failed: forward_only response logps could not be aligned to the ' + f'sampler old_logps (best mean|diff|={err}); the response frame is off, so teacher ' + f'logps cannot be trusted. Inspect _as_rows / _extract_resp before training.') + _OPSD_OFFSET = off + logger.info(f'[rsi_rl][opsd] response-logps suffix offset calibrated to {off} ' + f'(self-check mean|diff|={err:.4f} < tol {OPSD_SELFCHECK_TOL})') + + teacher_pifs = [_teacher_pif(template, sp, sol, toks) + for sp, toks, sol in zip(student_pifs, response_tokens_list, ref_solutions)] + t_rows = _as_rows(model.forward_only(inputs=teacher_pifs, micro_batch_size=MICRO_BATCH_SIZE).logps) + teacher_logps: List[List[float]] = [] + for row, tpif, toks in zip(t_rows, teacher_pifs, response_tokens_list): + n = len(toks) + resp = _extract_resp(row, len(tpif['input_ids']), n, _OPSD_OFFSET) + if len(resp) != n: + raise RuntimeError(f'OPSD teacher extraction: {len(resp)} logps for {n} response tokens') + teacher_logps.append([float(x) for x in resp]) + return teacher_logps + + def main(): device_groups = [ DeviceGroup(name='model', ranks=list(range(MODEL_GPUS)), device_type='GPU'), - DeviceGroup(name='sampler', ranks=list(range(MODEL_GPUS, NUM_GPUS)), device_type='GPU', - gpus_per_worker=SAMPLER_TP), + DeviceGroup(name='sampler', ranks=list(range(MODEL_GPUS, MODEL_GPUS + SAMPLER_GPUS)), + device_type='GPU'), ] - model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, tp_size=2, ep_size=2, pp_size=2, sequence_parallel=True) - sampler_dp_size = SAMPLER_GPUS // SAMPLER_TP - sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, dp_size=sampler_dp_size, tp_size=SAMPLER_TP) + model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, dp_size=MODEL_GPUS) + sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, dp_size=SAMPLER_GPUS) + ref_mesh = None + if REF_GPUS: + device_groups.append( + DeviceGroup(name='ref', ranks=list(range(MODEL_GPUS + SAMPLER_GPUS, NUM_GPUS)), + device_type='GPU')) + ref_mesh = DeviceMesh.from_sizes(world_size=REF_GPUS, dp_size=REF_GPUS) twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, groups=device_groups, lazy_collect=False) - lora_config = LoraConfig(target_modules='all-linear', r=LORA_RANK, lora_alpha=LORA_RANK * 2, lora_dropout=0.05) + # Full-parameter training: no adapter is added, so every weight is trained and + # the whole model is pushed to the sampler each step. + from twinkle.model.megatron import MegatronModel + model = MegatronModel(model_id=MODEL_ID, device_mesh=model_mesh, remote_group='model', + mixed_precision='bf16', variable_seq_lengths=True) + model.set_optimizer('default', lr=LEARNING_RATE) + model.set_lr_scheduler('default', lr_decay_steps=MAX_STEPS, max_lr=LEARNING_RATE) + if SOLVER_MODE == 'opsd': + # On-policy self-distillation: student pulled toward a teacher that saw + # the reference solution. No advantages / reward in the loss. + model.set_loss('OPSDLoss', reverse=OPSD_REVERSE) + else: + loss_kwargs: Dict[str, Any] = {'epsilon': 0.2, 'beta': KL_BETA} + if LOSS_NAME == 'DRGRPOLoss': + # Must be the real generation cap: the class default is 1024 + # (grpo.py:591) and it sits in the denominator, so leaving it there + # while generating MAX_NEW_TOKENS scales every gradient by + # MAX_NEW_TOKENS/1024. + loss_kwargs['max_completion_length'] = MAX_NEW_TOKENS + model.set_loss(LOSS_NAME, **loss_kwargs) + logger.info(f'[rsi_rl] loss={LOSS_NAME} {loss_kwargs} ' + f'ref={"none" if not REF_GPUS else REF_MODEL_ID}') + if KL_BETA > 0 and not REF_GPUS: + raise RuntimeError( + f'RSI_KL_BETA={KL_BETA} but RSI_REF_GPUS=0: the KL term needs ref_logps ' + f'(grpo.py:315 requires beta>0 AND ref_logps), so it would silently do ' + f'nothing. Set RSI_REF_GPUS (e.g. 2) or RSI_KL_BETA=0.') - model = MultiLoraMegatronModel(model_id=MODEL_ID, device_mesh=model_mesh, remote_group='model', - mixed_precision='bf16') - model.add_adapter_to_model(ADAPTER_NAME, lora_config, gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) - model.set_optimizer('default', lr=LEARNING_RATE, adapter_name=ADAPTER_NAME) - model.set_lr_scheduler('default', lr_decay_steps=MAX_STEPS, max_lr=LEARNING_RATE, adapter_name=ADAPTER_NAME) - model.set_loss('GRPOLoss', epsilon=0.2, adapter_name=ADAPTER_NAME) - model.set_processor(InputProcessor, adapter_name=ADAPTER_NAME) - model.set_template(TEMPLATE, model_id=MODEL_ID, enable_thinking=True, adapter_name=ADAPTER_NAME) + model.set_processor(InputProcessor, padding_free=True) + model.set_template(TEMPLATE, model_id=MODEL_ID, max_length=MAX_MODEL_LEN, enable_thinking=True) + # Observability only: approx_kl / clip_ratio / entropy per step. approx_kl at the + # first inner step also reconciles sampler vs trainer logps, which is the check + # for whether the full-weight sync actually landed. OPSD has no PPO ratio. + if SOLVER_MODE != 'opsd': + model.add_metric('GRPOMetric', is_training=True, epsilon=0.2) sampler = vLLMSampler( model_id=MODEL_ID, engine_args={ - 'tensor_parallel_size': SAMPLER_TP, 'gpu_memory_utilization': 0.8, - 'max_model_len': 8192, - 'max_lora_rank': LORA_RANK, - 'enable_lora': True, - 'enable_tower_connector_lora': True, + 'max_model_len': MAX_MODEL_LEN, }, device_mesh=sampler_mesh, remote_group='sampler', ) - sampler.set_template(TEMPLATE, model_id=MODEL_ID, enable_thinking=True) + sampler.set_template(TEMPLATE, model_id=MODEL_ID, max_length=MAX_MODEL_LEN, enable_thinking=True) + + # KL anchor: frozen base weights, forward only, no optimizer. Same template / + # processor as the trainer so the returned per-token logps line up position by + # position with the trainer's own forward (both are Megatron forwards over the + # identical token sequence, so no frame calibration is needed -- unlike the + # OPSD teacher, whose prompt has a different length). + ref_model = None + if REF_GPUS: + ref_model = MegatronModel(model_id=REF_MODEL_ID, device_mesh=ref_mesh, remote_group='ref', + mixed_precision='bf16', variable_seq_lengths=True) + # advantages=None on this path, so GRPOLoss short-circuits to a zero loss + # and only outputs['logps'] is harvested (grpo.py:298). + ref_model.set_loss('GRPOLoss', epsilon=0.2) + ref_model.set_processor(InputProcessor, padding_free=True) + ref_model.set_template(TEMPLATE, model_id=REF_MODEL_ID, max_length=MAX_MODEL_LEN, + enable_thinking=True) + + # Driver-side template for token surgery: the GRPO tool-error bridge and the + # OPSD teacher-prompt concat both run on the driver. + local_template = make_local_template() + + ckpt_manager = CheckpointEngineManager(model=model, sampler=sampler) GLOBAL_BATCH_SIZE = BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS dataloader = DataLoader(dataset=create_rsi_dataset, batch_size=GLOBAL_BATCH_SIZE, min_batch_size=GLOBAL_BATCH_SIZE, device_mesh=model_mesh, remote_group='model') - advantage_fn = GRPOAdvantage() metrics = CompletionRewardMetric() - reward_fn = ToolMatchReward() + reward_fn = RoundReward() 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 RSI per-round GRPO (MultiLoraMegatron, filesystem LoRA sync)') + logger.info('Starting RSI per-round GRPO (full-parameter Megatron)') logger.info(get_device_placement()) for batch in dataloader: @@ -251,48 +1221,124 @@ def main(): for prompt in batch: expand_prompts.extend([prompt] * NUM_GENERATIONS) - lora_sync_path = model.save(f'lora-sync-step-{optim_step}', output_dir=LORA_SYNC_DIR, adapter_name=ADAPTER_NAME) + # No LoRA, so every sync ships the full weights. + ckpt_manager.sync_weights(merge_and_sync=True) sampler.reset_prefix_cache() - sample_responses = sampler.sample(expand_prompts, sampling_params, adapter_path=lora_sync_path) - 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)) + all_tokens: List[List[int]] = [] + if SOLVER_MODE == 'grpo': + # Rollout with error-feedback continuation on failed code rounds. + all_input_data, all_old_logps, all_completion_lengths = grpo_continue( + sampler, local_template, expand_prompts, sampling_params) + else: + # OPSD: single turn; also keep raw response tokens for the teacher concat. + all_input_data, all_old_logps, all_completion_lengths = [], [], [] + for sample_response in sampler.sample(expand_prompts, sampling_params): + 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)) + all_tokens.append(list(sequence.tokens)) + # Reward drives GRPO advantages; in OPSD it is observability only. rewards = reward_fn(all_input_data) - metrics.accumulate(completion_lengths=all_completion_lengths, rewards={'tool_match': rewards}) - advantages = advantage_fn(rewards, num_generations=NUM_GENERATIONS, scale='group').tolist() + scored = [r for r in rewards if r is not None] + metrics.accumulate(completion_lengths=all_completion_lengths, rewards={'round_reward': scored}) + + teacher_logps = None + if SOLVER_MODE == 'grpo': + advantages = group_advantages(rewards, num_generations=NUM_GENERATIONS, scale='group') + # First group, verbatim: catches a reward/advantage misalignment (a high-reward + # sample must not carry a negative advantage). + logger.info(f'[group0] rewards={rewards[:NUM_GENERATIONS]} ' + f'advantages={[round(a, 3) for a in advantages[:NUM_GENERATIONS]]} ' + f'lens={all_completion_lengths[:NUM_GENERATIONS]}') + else: + # OPSD: no advantages in the loss; keep a zero list only for the audit dump. + advantages = [0.0] * len(all_input_data) + all_ref_solutions: List[str] = [] + for prompt in expand_prompts: + ud = {item[0]: item[1] for item in (prompt.get('user_data') or [])} + all_ref_solutions.append(ud.get('ref_solution', '')) + if any(not s for s in all_ref_solutions): + raise RuntimeError('OPSD needs a ref_solution (challenger passing solution) on ' + 'every code round; some rounds are missing it.') + teacher_logps = opsd_teacher_logps( + model, local_template, all_input_data, all_tokens, + all_ref_solutions, all_old_logps) + logger.info(f'[group0] opsd rewards(obs)={rewards[:NUM_GENERATIONS]} ' + f'lens={all_completion_lengths[:NUM_GENERATIONS]}') + + if REWARD_DUMP: + # Append one audit line per rollout of this step. Reward/advantage are + # already computed above; this only reads them, never changes them. + with open(REWARD_DUMP, 'a', encoding='utf-8') as fdump: + for i, rec in enumerate(reward_fn.records): + fdump.write(json.dumps(_json_safe({ + 'step': optim_step + 1, + 'group': i // NUM_GENERATIONS, + 'score': rec.get('score'), + 'advantage': round(advantages[i], 4), + 'len': all_completion_lengths[i], + 'kind': rec.get('kind'), + 'ref_call': rec.get('ref_call'), + 'gen_call': rec.get('gen_call'), + 'reason': rec.get('reason'), + 'completion': rec.get('completion'), + }), ensure_ascii=False) + '\n') total = len(all_input_data) for mb_start in range(0, total, MINI_BATCH_SIZE): mb_end = min(mb_start + MINI_BATCH_SIZE, total) - model.forward_backward( - inputs=all_input_data[mb_start:mb_end], - old_logps=all_old_logps[mb_start:mb_end], - advantages=advantages[mb_start:mb_end], - micro_batch_size=MICRO_BATCH_SIZE, - adapter_name=ADAPTER_NAME, - ) - model.clip_grad_and_step(adapter_name=ADAPTER_NAME) + if SOLVER_MODE == 'grpo': + ref_logps = None + if ref_model is not None: + # ModelOutput is a TypedDict (data_format/output.py:15), so it is a + # plain dict -- index it, never attribute-access it. + ref_out = ref_model.forward_only( + inputs=all_input_data[mb_start:mb_end], + micro_batch_size=MICRO_BATCH_SIZE) + ref_logps = _as_rows(ref_out['logps']) + if optim_step == 0 and mb_start == 0: + # One-time shape check: a row must cover the whole padded + # sequence, otherwise GRPOLoss's full-sequence branch + # (grpo.py:210) would reject it and we want to see the + # numbers rather than only the assertion. + logger.info( + f'[rsi_rl][kl] ref rows={len(ref_logps)} ' + f'row_lens={[len(r) for r in ref_logps[:4]]} ' + f'input_lens={[len(x["input_ids"]) for x in all_input_data[mb_start:mb_start + 4]]}') + model.forward_backward( + inputs=all_input_data[mb_start:mb_end], + old_logps=all_old_logps[mb_start:mb_end], + advantages=advantages[mb_start:mb_end], + ref_logps=ref_logps, + micro_batch_size=MICRO_BATCH_SIZE, + ) + else: + # OPSD: teacher_logps (ragged, response-only) drives the k3 pull; + # OPSDLoss ignores advantages / old_logps. + model.forward_backward( + inputs=all_input_data[mb_start:mb_end], + teacher_logps=teacher_logps[mb_start:mb_end], + 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'rsi-executor-checkpoint-{optim_step}', adapter_name=ADAPTER_NAME) + model.save(f'{SAVE_NAME}-checkpoint-{optim_step}', output_dir=SAVE_DIR) log_dict = metrics.calculate() - log_dict.update(model.calculate_metric(is_training=True, adapter_name=ADAPTER_NAME)) + log_dict.update(model.calculate_metric(is_training=True)) + log_dict.update({f'train/{k}': v for k, v in reward_fn.stats.items()}) swanlab.log(log_dict) metrics.reset() logger.info(f'[Step {optim_step}/{MAX_STEPS}] {log_dict}') logger.info(f'Training completed. optim_steps={optim_step}') - model.save('rsi-executor-final', adapter_name=ADAPTER_NAME) + model.save(SAVE_NAME, output_dir=SAVE_DIR) if __name__ == '__main__': diff --git a/src/twinkle_agentic/segment/__init__.py b/src/twinkle_agentic/segment/__init__.py deleted file mode 100644 index 17f64e7f2..000000000 --- a/src/twinkle_agentic/segment/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .base import Segmenter, Turn, TurnSegmenter -from .llm_segmenter import LlmSegmenter - -__all__ = ['Segmenter', 'Turn', 'TurnSegmenter', 'LlmSegmenter'] diff --git a/src/twinkle_agentic/segment/base.py b/src/twinkle_agentic/segment/base.py deleted file mode 100644 index 353daa5cf..000000000 --- a/src/twinkle_agentic/segment/base.py +++ /dev/null @@ -1,237 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Trajectory segmentation for segment-level rubric scoring. - -A long agent trajectory is split into a list of *segments*; each segment is a -self-contained sub-trajectory (``{'messages': [...], 'tools': [...]}``) that can -be fed directly to a :class:`~twinkle_agentic.verifier.Verifier`. - -Two layers, matching the literature: -- **Structural (per-turn)** — free, deterministic. A "turn" is one assistant - message plus the tool result messages it triggered (Web-Shepherd / AgentPRM - style turn-level MDP). See :class:`TurnSegmenter`. -- **Sub-goal (LLM)** — compress each turn to a one-line intent gist, then make - ONE LLM pass over the whole (short) gist list to group turns into a few - coarse sub-goals (Web-Shepherd / MiRA style), then reassemble segments from - the ORIGINAL messages by index. See :class:`LlmSegmenter`. - -Design notes: -- The leading ``system`` message and the first ``user`` message form the - trajectory *preamble*; it is not itself a scorable segment, but every segment - carries the preamble (system + original user query) so the verifier keeps the - task context. This mirrors ``RubricVerifier._infer_query``. -- A new ``user`` message mid-trajectory is a hard boundary (a new turn starts). -""" -from __future__ import annotations - -from abc import ABC, abstractmethod -from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Tuple - - -@dataclass -class Turn: - """One structural turn: an assistant message + its tool results. - - ``indices`` are positions into the original ``messages`` list so a segment - can be reassembled verbatim from the source (never from a summary). - """ - indices: List[int] - messages: List[dict] = field(default_factory=list) - role_kind: str = 'assistant' # 'assistant' | 'user' | 'other' - - -class Segmenter(ABC): - """Split a trajectory into scorable sub-trajectories. - - Subclasses implement :meth:`segment`. The base class provides the shared - structural turn splitter and segment assembly so LLM-based subclasses only - decide *how turns are grouped*. - """ - - def __call__(self, trajectory: dict, **kwargs) -> List[dict]: - return self.segment(trajectory, **kwargs) - - @abstractmethod - def segment(self, trajectory: dict, **kwargs) -> List[dict]: - raise NotImplementedError - - # ------------------------------------------------------------------ - # shared: preamble + structural turn splitting + assembly - # ------------------------------------------------------------------ - @staticmethod - def _split_preamble(messages: List[dict]) -> Tuple[List[dict], int]: - """Return (preamble_messages, start_index). - - Preamble = leading system message(s) + the first user message. Segments - start at ``start_index`` (the message after the first user turn). - """ - preamble: List[dict] = [] - i = 0 - n = len(messages) - while i < n and messages[i].get('role') == 'system': - preamble.append(messages[i]) - i += 1 - if i < n and messages[i].get('role') == 'user': - preamble.append(messages[i]) - i += 1 - return preamble, i - - @classmethod - def split_turns(cls, messages: List[dict], start: int = 0) -> List[Turn]: - """Group messages[start:] into structural turns. - - A turn begins at an ``assistant`` message and absorbs the following - ``tool`` messages. A mid-trajectory ``user`` message becomes its own - boundary turn (role_kind='user'). Stray leading non-assistant messages - are attached to the first turn. - """ - turns: List[Turn] = [] - cur: Optional[Turn] = None - for idx in range(start, len(messages)): - role = messages[idx].get('role') - if role == 'assistant': - cur = Turn(indices=[idx], messages=[messages[idx]], role_kind='assistant') - turns.append(cur) - elif role == 'user': - # hard boundary: user re-prompt starts a fresh turn - cur = Turn(indices=[idx], messages=[messages[idx]], role_kind='user') - turns.append(cur) - else: # tool / other -> attach to current turn, or open a new one - if cur is None: - cur = Turn(indices=[idx], messages=[messages[idx]], role_kind='other') - turns.append(cur) - else: - cur.indices.append(idx) - cur.messages.append(messages[idx]) - return turns - - @staticmethod - def _assemble(trajectory: dict, preamble: List[dict], turns_slice: List[Turn]) -> dict: - """Build a segment sub-trajectory from preamble + a slice of turns. - - Messages are taken from the ORIGINAL trajectory (verbatim), so the - verifier scores real content, not any compressed gist. - """ - seg_messages: List[dict] = list(preamble) - for t in turns_slice: - seg_messages.extend(t.messages) - seg: Dict[str, Any] = {'messages': seg_messages} - if trajectory.get('tools'): - seg['tools'] = list(trajectory['tools']) - if trajectory.get('user_data'): - seg['user_data'] = list(trajectory['user_data']) - return seg - - @classmethod - def _segments_from_groups( - cls, - trajectory: dict, - preamble: List[dict], - turns: List[Turn], - groups: List[List[int]], - ) -> List[dict]: - """Assemble segments given a grouping of turn-indices. - - ``groups`` is a list of lists of indices into ``turns``. Robust to - gaps/overlaps: see :meth:`_normalize_groups`. - """ - groups = cls._normalize_groups(groups, len(turns)) - return [cls._assemble(trajectory, preamble, [turns[i] for i in grp]) - for grp in groups if grp] - - @staticmethod - def _normalize_groups(groups: List[List[int]], n_turns: int) -> List[List[int]]: - """Repair LLM-proposed groupings into a clean partition of 0..n_turns-1. - - - drop out-of-range indices - - dedupe (first occurrence wins; later duplicates dropped) - - sort each group; sort groups by their first index - - assign any uncovered turns to the nearest preceding group (or the - first group), so every turn lands in exactly one segment - """ - if n_turns <= 0: - return [] - seen: set = set() - cleaned: List[List[int]] = [] - for grp in groups: - g = [] - for i in grp: - if isinstance(i, bool): # guard: bools are ints in python - continue - if isinstance(i, int) and 0 <= i < n_turns and i not in seen: - seen.add(i) - g.append(i) - if g: - cleaned.append(sorted(g)) - cleaned.sort(key=lambda g: g[0]) - - # cover missing turns - missing = [i for i in range(n_turns) if i not in seen] - if missing: - if not cleaned: - cleaned = [missing] - else: - for i in missing: - # nearest preceding group by first-index - target = cleaned[0] - for grp in cleaned: - if grp[0] <= i: - target = grp - else: - break - target.append(i) - for grp in cleaned: - grp.sort() - cleaned.sort(key=lambda g: g[0]) - return cleaned - - -class TurnSegmenter(Segmenter): - """Structural, LLM-free segmenter. - - ``granularity='turn'``: one segment per turn (finest; per-tool-call level). - ``granularity='cluster'``: merge consecutive tool-using assistant turns into - one segment, closing the cluster on a turn that produces a user-facing - text answer with no tool calls (sub-task attempt level). A ``user`` turn - always starts a new cluster. - """ - - def __init__(self, granularity: str = 'cluster'): - if granularity not in ('turn', 'cluster'): - raise ValueError("granularity must be 'turn' or 'cluster'") - self.granularity = granularity - - def segment(self, trajectory: dict, **kwargs) -> List[dict]: - messages = list(trajectory.get('messages', []) or []) - preamble, start = self._split_preamble(messages) - turns = self.split_turns(messages, start) - if not turns: - return [self._assemble(trajectory, preamble, [])] if preamble else [] - - if self.granularity == 'turn': - groups = [[i] for i in range(len(turns))] - else: - groups = self._cluster_groups(turns) - return self._segments_from_groups(trajectory, preamble, turns, groups) - - @staticmethod - def _cluster_groups(turns: List[Turn]) -> List[List[int]]: - groups: List[List[int]] = [] - cur: List[int] = [] - for i, t in enumerate(turns): - if t.role_kind == 'user': - if cur: - groups.append(cur) - cur = [] - groups.append([i]) # user re-prompt as its own boundary segment - continue - cur.append(i) - has_tool_call = any(m.get('role') == 'assistant' and m.get('tool_calls') - for m in t.messages) - if not has_tool_call: - # a text-only assistant answer closes the current sub-task cluster - groups.append(cur) - cur = [] - if cur: - groups.append(cur) - return groups diff --git a/src/twinkle_agentic/segment/llm_segmenter.py b/src/twinkle_agentic/segment/llm_segmenter.py deleted file mode 100644 index 7230d225f..000000000 --- a/src/twinkle_agentic/segment/llm_segmenter.py +++ /dev/null @@ -1,308 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Compress-then-segment: per-turn action gist + one-pass LLM sub-goal grouping. - -Pipeline (all LLM calls distilled via ``llm_backup``): - 1. Split the trajectory into structural turns (base class, free). - 2. Compress each turn to a one-line intent gist with an - :class:`ActionSummarizer` (optional; falls back to a truncated - structural render when no summarizer is given). - 3. Make ONE LLM pass over the numbered gist list to group turns into a few - coarse sub-goals — the segmenter LLM sees the whole (short) trajectory at - once, which is what makes segmentation global yet cheap. - 4. Reassemble segments from the ORIGINAL messages by turn index (scoring - always uses verbatim content, never the gist). - -The grouping call is wrapped with ``@llm_backup`` (student sub-goal model with -teacher fallback + progressive distillation), keyed by ``query`` so confidence -is tracked per task family. Index alignment is validated/repaired in pure code -by the base class ``_normalize_groups`` (contiguous, non-overlapping, full -coverage). -""" -from __future__ import annotations - -import json -import re -from typing import TYPE_CHECKING, Any, List, Optional - -from twinkle_agentic.utils.llm_backup import llm_backup - -from .base import Segmenter, Turn - -if TYPE_CHECKING: - from twinkle.data_format import SamplingParams # noqa: F401 - from twinkle.sampler.base import Sampler # noqa: F401 - from twinkle_agentic.summarizer.action_summarizer import ActionSummarizer # noqa: F401 - - -_SEG_SYSTEM = """\ -You segment an AI agent's trajectory into a FEW coarse sub-goals for later \ -evaluation. You are given the task query and a numbered list of one-line turn \ -gists (one per turn). Group CONSECUTIVE turns that jointly pursue the same \ -sub-goal. - -Guidelines: -- Prefer {min_g}-{max_g} sub-goals total. Each sub-goal spans a contiguous - range of turns; ranges must not overlap and must cover every turn. -- Group by MEANINGFUL task progress, not by exact actions: e.g. several - consecutive searches that gather info for one purpose form ONE sub-goal. -- A user re-prompt starts a new sub-goal. - -Output ONLY a JSON array, each item: {"goal": "<short label>", "start": <first turn #>, "end": <last turn #>} -Turn numbers are 0-based and inclusive. No prose, no markdown fence.""" - -_SEG_USER = """\ -## Task / query -{query} - -## Turn gists (numbered) -{gists} - -Now output the JSON array of sub-goals covering turns 0..{last}.""" - - -class LlmSegmenter(Segmenter): - """Sub-goal segmenter via compress-then-one-pass-LLM grouping. - - Args: - sampler: Student model sampler for the grouping call. If ``None``, every - grouping call is served by the teacher via ``llm_backup``. - action_summarizer: Optional :class:`ActionSummarizer` used to compress - each turn. If ``None``, a truncated structural render is used as the - gist (no per-turn LLM cost). - min_subgoals / max_subgoals: Target sub-goal count window. - sampling_params: Sampling params for the grouping call. - lora_path: LoRA adapter for the sub-goal student model. - max_gist_chars: Truncation for the structural-render fallback gist. - """ - - def __init__( - self, - sampler: Optional['Sampler'] = None, - *, - action_summarizer: Optional['ActionSummarizer'] = None, - min_subgoals: int = 3, - max_subgoals: int = 6, - sampling_params: Optional['SamplingParams'] = None, - lora_path: Optional[str] = None, - max_gist_chars: int = 200, - ): - if max_subgoals < min_subgoals: - raise ValueError('max_subgoals must be >= min_subgoals') - if min_subgoals < 1: - raise ValueError('min_subgoals must be >= 1') - self.sampler = sampler - self.action_summarizer = action_summarizer - self.min_subgoals = int(min_subgoals) - self.max_subgoals = int(max_subgoals) - self.sampling_params = sampling_params - self.lora_path = lora_path or None - self.max_gist_chars = int(max_gist_chars) - - # ------------------------------------------------------------------ - def segment(self, trajectory: dict, *, query: Optional[str] = None, **kwargs) -> List[dict]: - messages = list(trajectory.get('messages', []) or []) - preamble, start = self._split_preamble(messages) - turns = self.split_turns(messages, start) - if not turns: - return [self._assemble(trajectory, preamble, [])] if preamble else [] - # Too few turns to bother segmenting -> one segment. - if len(turns) <= self.min_subgoals: - return self._segments_from_groups( - trajectory, preamble, turns, [[i] for i in range(len(turns))]) - - from .base import TurnSegmenter - - # No LLM available at all (no student sampler AND no teacher configured) - # -> degrade to free structural clustering instead of crashing. - if not self._llm_available(): - groups = TurnSegmenter._cluster_groups(turns) - return self._segments_from_groups(trajectory, preamble, turns, groups) - - query = query or self._infer_query(messages) - gists = [self._turn_gist(t, query) for t in turns] - gist_block = '\n'.join(f'[{i}] {g}' for i, g in enumerate(gists)) - - raw = self._group( - trajectory=self._group_trajectory(query, gist_block, len(turns)), - sampling_params=self._group_sampling_params(), - query=query) - groups = self._parse_groups(raw, len(turns)) - if not groups: - # LLM produced nothing usable -> fall back to structural clustering. - groups = TurnSegmenter._cluster_groups(turns) - return self._segments_from_groups(trajectory, preamble, turns, groups) - - def _llm_available(self) -> bool: - """True if a student sampler exists or a teacher API is configured. - - Mirrors the env vars ``llm_backup`` uses for its teacher; when neither a - student nor a teacher is present we must not attempt an LLM call. - """ - if self.sampler is not None: - return True - import os - return bool(os.environ.get('LLM_BACKUP_API_KEY') - or os.environ.get('OPENAI_API_KEY') - or os.environ.get('LLM_BACKUP_BASE_URL')) - - # ------------------------------------------------------------------ - # the distilled grouping call - # ------------------------------------------------------------------ - @llm_backup(key_params=['query'], comparator=lambda a, b: _grouping_similar(a, b)) - def _group(self, trajectory, sampling_params, query: str = None) -> str: - if self.sampler is None: - return '' - sample_kwargs: dict[str, Any] = {'sampling_params': sampling_params} - if self.lora_path is None: - sample_kwargs['use_base_model'] = True - else: - sample_kwargs['adapter_path'] = self.lora_path - responses = self.sampler.sample([trajectory], **sample_kwargs) - resp = list(responses)[0] if responses else None - if resp is None: - return '' - seqs = getattr(resp, 'sequences', None) or [] - return (getattr(seqs[0], 'decoded', None) or '') if seqs else '' - - # ------------------------------------------------------------------ - # per-turn gist - # ------------------------------------------------------------------ - def _turn_gist(self, turn: Turn, query: str) -> str: - rendered = self._render_turn(turn) - if self.action_summarizer is not None: - try: - gist = self.action_summarizer(rendered, query=query) - if isinstance(gist, str) and gist.strip(): - return ' '.join(gist.split())[:self.max_gist_chars] - except Exception: - pass - # fallback: truncated structural render (no LLM) - return ' '.join(rendered.split())[:self.max_gist_chars] - - @staticmethod - def _render_turn(turn: Turn) -> str: - parts: List[str] = [] - for m in turn.messages: - role = m.get('role', '?') - content = m.get('content') - if isinstance(content, list): - content = '\n'.join(p.get('text', '') for p in content - if isinstance(p, dict) and p.get('type') == 'text') - content = content or '' - tool_calls = m.get('tool_calls') or [] - if tool_calls: - names = ', '.join((tc.get('function') or {}).get('name', '?') - for tc in tool_calls if isinstance(tc, dict)) - parts.append(f'{role}: {content} [tool_calls: {names}]'.strip()) - else: - parts.append(f'{role}: {content}'.strip()) - return ' | '.join(p for p in parts if p) - - # ------------------------------------------------------------------ - # prompt / sampling plumbing - # ------------------------------------------------------------------ - def _group_trajectory(self, query: str, gist_block: str, n_turns: int) -> dict: - system = (_SEG_SYSTEM - .replace('{min_g}', str(self.min_subgoals)) - .replace('{max_g}', str(self.max_subgoals))) - user = (_SEG_USER - .replace('{query}', query) - .replace('{gists}', gist_block) - .replace('{last}', str(n_turns - 1))) - return {'messages': [ - {'role': 'system', 'content': system}, - {'role': 'user', 'content': user}, - ]} - - def _group_sampling_params(self): - if self.sampling_params is not None: - return self.sampling_params - from twinkle.data_format.sampling import SamplingParams - return SamplingParams(temperature=0.0, max_tokens=512) - - # ------------------------------------------------------------------ - # parsing - # ------------------------------------------------------------------ - @staticmethod - def _infer_query(messages: List[dict]) -> str: - for m in messages: - if m.get('role') == 'user': - c = m.get('content') - if isinstance(c, str) and c.strip(): - return c.strip() - return '(no explicit query)' - - _JSON_ARRAY_RE = re.compile(r'\[.*\]', re.DOTALL) - - @classmethod - def _parse_groups(cls, raw: str, n_turns: int) -> List[List[int]]: - """Parse the sub-goal JSON into a list of turn-index groups. - - Accepts ``[{"goal":..,"start":i,"end":j}, ...]``. Falls back to empty - on unparseable output (caller then uses structural clustering). - """ - if not raw: - return [] - text = raw.strip() - m = cls._JSON_ARRAY_RE.search(text) - if not m: - return [] - try: - data = json.loads(m.group(0)) - except (json.JSONDecodeError, ValueError): - return [] - if not isinstance(data, list): - return [] - groups: List[List[int]] = [] - for item in data: - if not isinstance(item, dict): - continue - s, e = item.get('start'), item.get('end') - if not isinstance(s, int) or not isinstance(e, int): - continue - if e < s: - s, e = e, s - s = max(0, s) - e = min(n_turns - 1, e) - grp = list(range(s, e + 1)) - if grp: - groups.append(grp) - return groups - - -# --------------------------------------------------------------------------- -# comparator for llm_backup -# --------------------------------------------------------------------------- -_JSON_ARRAY_RE = re.compile(r'\[.*\]', re.DOTALL) - - -def _boundaries(raw: str) -> Optional[List[int]]: - m = _JSON_ARRAY_RE.search((raw or '').strip()) - if not m: - return None - try: - data = json.loads(m.group(0)) - except (json.JSONDecodeError, ValueError): - return None - if not isinstance(data, list): - return None - starts = [] - for item in data: - if isinstance(item, dict) and isinstance(item.get('start'), int): - starts.append(item['start']) - return sorted(starts) if starts else None - - -def _grouping_similar(a: str, b: str) -> bool: - """Two segmentations match when they have a similar number of sub-goals and - near-identical boundaries (not byte-identical labels/text).""" - ba, bb = _boundaries(a), _boundaries(b) - if ba is None or bb is None: - return (a or '').strip() == (b or '').strip() - if abs(len(ba) - len(bb)) > 1: - return False - # Jaccard-ish agreement on boundary start positions - sa, sb = set(ba), set(bb) - inter = len(sa & sb) - union = len(sa | sb) or 1 - return inter / union >= 0.6 diff --git a/src/twinkle_agentic/summarizer/__init__.py b/src/twinkle_agentic/summarizer/__init__.py index 9991a0e33..53d48ff21 100644 --- a/src/twinkle_agentic/summarizer/__init__.py +++ b/src/twinkle_agentic/summarizer/__init__.py @@ -1,5 +1,3 @@ -from .action_summarizer import ActionSummarizer from .base import Summarizer -from .fact_summarizer import FactSummarizer -__all__ = ['Summarizer', 'FactSummarizer', 'ActionSummarizer'] +__all__ = ['Summarizer'] diff --git a/src/twinkle_agentic/summarizer/action_summarizer.py b/src/twinkle_agentic/summarizer/action_summarizer.py deleted file mode 100644 index 2e6c3aa0d..000000000 --- a/src/twinkle_agentic/summarizer/action_summarizer.py +++ /dev/null @@ -1,86 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Action-level summarizer: compress ONE agent turn into a single-line -``(action, goal, result-state)`` gist for downstream trajectory segmentation. - -Unlike :class:`FactSummarizer` (which preserves *facts* for retrieval), this -summarizer preserves *intent* — what the agent tried to do this turn and -whether it worked — because that is what a segmenter needs to group turns into -sub-goals. The gist is deliberately tiny (one line) so a whole long trajectory -can be laid out and segmented in a single LLM pass. - -Same machinery as every other component: the shared ``_sample`` is decorated -with ``@llm_backup`` so student/teacher routing + progressive distillation -happen transparently. -""" -from __future__ import annotations - -from typing import TYPE_CHECKING, Any - -from twinkle_agentic.summarizer.base import Summarizer - -if TYPE_CHECKING: - from twinkle.data_format import SamplingParams # noqa: F401 - from twinkle.sampler.base import Sampler # noqa: F401 - - -_ACTION_SCHEMA = """\ -You compress ONE turn of an AI agent's trajectory into a single short line \ -that captures INTENT, not facts. A downstream segmenter reads these lines to \ -group consecutive turns into sub-goals, so keep the action's PURPOSE and \ -OUTCOME, and drop retrieved content / long results. - -Output EXACTLY one line in this shape (<= 140 chars): - <ACTION> | goal: <what this turn is trying to achieve> | result: <ok|fail|partial|pending|answer> - -Where <ACTION> is a terse verb-phrase for the turn, e.g. one of: - search, read, compute, call-tool:<name>, plan, reason, ask-user, final-answer, other - -Rules: -- Focus on WHY the turn happened and WHETHER it advanced the task. -- result=ok (tool/step succeeded), fail (error/empty), partial (some progress), - pending (awaiting more), answer (produced a user-facing final answer). -- Do NOT copy retrieved passages, numbers or long tool output — only the gist. -- One line only. No markdown, no extra commentary. -""" - -_ACTION_USER_TEMPLATE = """\ -Summarize this single agent turn into ONE intent line (see the required shape). \ -Capture the action, its goal, and the result-state. Do not exceed {budget} chars. \ -Ignore long retrieved content; keep only what the turn was DOING. - -## Task / query (context) -{query} - -## Turn -{text}""" - - -class ActionSummarizer(Summarizer): - """Compress a single rendered turn into a one-line intent gist. - - Defaults target a tiny fixed budget (one line) regardless of turn length, - which is what makes a whole-trajectory single-pass segmentation cheap. - """ - - def __init__( - self, - sampler: 'Sampler', - compression_ratio: float = 6.0, - *, - model_path: str = '', - sampling_params: 'SamplingParams | None' = None, - min_budget_chars: int = 140, - template: Any | None = None, - lora_path: str | None = None, - ): - super().__init__( - sampler, - compression_ratio, - model_path=model_path, - sampling_params=sampling_params, - system_prompt=_ACTION_SCHEMA, - user_prompt_template=_ACTION_USER_TEMPLATE, - min_budget_chars=min_budget_chars, - template=template, - lora_path=lora_path, - ) diff --git a/src/twinkle_agentic/summarizer/error_summarizer.py b/src/twinkle_agentic/summarizer/error_summarizer.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/src/twinkle_agentic/summarizer/fact_summarizer.py b/src/twinkle_agentic/summarizer/fact_summarizer.py deleted file mode 100644 index ab24a081b..000000000 --- a/src/twinkle_agentic/summarizer/fact_summarizer.py +++ /dev/null @@ -1,83 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -from __future__ import annotations - -from typing import TYPE_CHECKING, Any - -from twinkle_agentic.summarizer.base import Summarizer - -if TYPE_CHECKING: - from twinkle.data_format import SamplingParams # noqa: F401 - from twinkle.sampler.base import Sampler # noqa: F401 - -_SECTION_SCHEMA = """You are a text compression assistant. A downstream model will read your compressed output to decide whether the detail it needs is inside this block; if yes, it will fetch and read the original passage. - -Downstream model workflow: -Read your compressed output -> Decide whether needed info is in this block -> If yes -> Fetch original. - -Therefore your compression MUST NOT lose major information from the source. - -Output format: - -```text -## Summary -Overview plus facts STRONGLY RELATED to the Query, stated explicitly. - -## More -A collapsed index; expansion required to see specific information. -``` - -Rules: -1. Telegraphic style — drop function words ("the", "a", "is", "are", "of", ...); colons and commas mean "is" / "has". - * Exception: KEEP role-tagging verb+preposition phrases verbatim ("published by X", "written by X", "directed by X", "starring X", "founded by X", "created by X", "composed by X", "produced by X", "based on X", "adapted from X"). Collapsing these to a bare name loses the relation role (author vs publisher vs director) that the downstream question may hinge on. -2. Summary MUST contain the passage's primary topic + 2–4 concrete core facts drawn from the source (entities, numbers, dates, relations). If a Query is given, order Query-relevant facts first, but STILL include other core facts within the budget. A Query is an ORDERING HINT, NOT a filter. -3. Summary MUST NOT be meta-commentary about the Query. Forbidden patterns: "no X mention", "Query info: absent", "passage covers Y only", "does not contain ...", "no relevant info", or summaries that are only abstract category words like "structure/order/usage" with no facts. If the passage is unrelated to the Query, you still summarize the passage normally. -4. More is an INDEX of category keywords, NOT inline data. Enumerate what CAN be recovered from the source (e.g. "birthplace, death place, age"); do NOT paste dates/numbers/names inline. Make sure all category of useful facts are introduced here. -5. Output language MUST match the source language. -6. Do NOT fabricate. Do NOT omit major information. Any fact not in the source MUST NOT appear in your output. - -Now begin. -""" # noqa - -_SECTION_USER_TEMPLATE = """\ -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. - -## Query (ordering hint only — still summarize the whole passage) -{query} - -## Target length -Compress AS MUCH AS faithfully possible. HARD CEILING: {budget} chars. \ -If core facts fit in far fewer chars, output fewer. \ -Never exceed the ceiling. - -## Passage -{text}""" - - -class FactSummarizer(Summarizer): - - def __init__( - self, - sampler: Sampler, - compression_ratio: float = 2.0, - *, - model_path: str = '', - sampling_params: SamplingParams | None = None, - min_budget_chars: int = 250, - template: Any | None = None, - lora_path: str | None = None, - ): - super().__init__( - sampler, - compression_ratio, - model_path=model_path, - sampling_params=sampling_params, - system_prompt=_SECTION_SCHEMA, - user_prompt_template=_SECTION_USER_TEMPLATE, - min_budget_chars=min_budget_chars, - template=template, - lora_path=lora_path, - ) diff --git a/src/twinkle_agentic/summarizer/pattern_summarizer.py b/src/twinkle_agentic/summarizer/pattern_summarizer.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/src/twinkle_agentic/tools/extract_condensed.py b/src/twinkle_agentic/tools/extract_condensed.py deleted file mode 100644 index ead82c2c5..000000000 --- a/src/twinkle_agentic/tools/extract_condensed.py +++ /dev/null @@ -1,150 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -from typing import Any, Dict, List, Optional - -from twinkle.data_format.message import Tool as ToolInfo -from twinkle_agentic.data_format import Chunks -from .base import Tool - -TOOL_NAME = 'extract_condensed' - - -class ExtractCondensed(Tool): - """Return the original text behind a ``<block_N>`` compressed segment. - - TODO: Experimental feature, wait for testing - Args: - chunks: The :class:`Chunks` object emitted by a condenser - (post-compression). Each condensed chunk should carry - ``raw.original`` holding the pre-compression text; if that - snapshot is missing the block is still enumerated (so - numbering stays aligned with ``<block_N>``) but the tool - returns an explicit error on lookup rather than silently - handing back the compressed stand-in. - - The block enumeration rule mirrors :meth:`Chunks.to_trajectory` - exactly: only text chunks with ``raw.condensed=True``, - ``role != 'tool'`` and non-empty content are indexed via a - 1-based monotonic counter in chunk order. The block ids this - tool accepts therefore match the ``<block_N>`` tags the model - actually sees. - """ - - def __init__(self, chunks: Chunks): - self._blocks: Dict[int, Optional[str]] = {} - # Trajectory-bound set of block ids already returned in full. - self._already_expanded: set = set() - counter = 0 - for c in chunks.chunks: - if c.get('type') != 'text': - continue - content = c.get('content') - if not isinstance(content, str) or not content: - continue - if c.get('role') == 'tool': - continue - raw = c.get('raw') - if not (isinstance(raw, dict) and raw.get('condensed')): - continue - counter += 1 - original = raw.get('original') - self._blocks[counter] = (original if isinstance(original, str) and original else None) - - # ------------------------------------------------------------------ - # Tool interface - # ------------------------------------------------------------------ - def tool_info(self) -> ToolInfo: - return { - 'type': 'function', - 'function': { - 'name': - TOOL_NAME, - 'description': ('Recover the full, uncompressed text of ONE previously ' - 'condensed passage, identified by its <block_N> tag. Use ' - 'this tool whenever you need to re-read the original ' - 'detail of a compressed block. Each call expands exactly ' - 'one block; issue separate calls for additional blocks, ' - 'and do not request the same block twice.'), - 'parameters': { - 'blocks': ('int, the 1-indexed block number N appearing ' - 'inside <block_N>...</block_N>. Exactly one ' - 'block per call (e.g. 3); lists are rejected.'), - }, - }, - } - - def __call__(self, tool_name: str, arguments: Dict[str, Any]) -> str: - if not isinstance(arguments, dict): - return (f'Error: arguments must be an object, got ' - f'{type(arguments).__name__}.') - # Accept the new preferred name ``blocks`` first, fall back to the - # legacy singular ``block`` for backward compatibility with callers - # that were built against the int-only interface. - if 'blocks' in arguments: - raw = arguments['blocks'] - key = 'blocks' - elif 'block' in arguments: - raw = arguments['block'] - key = 'block' - else: - return 'Error: missing required argument "blocks".' - - # Single-block-per-call contract. Reject list/tuple up front so a - # hallucinated ``blocks=[1..200]`` cannot balloon the tool response. - if isinstance(raw, (list, tuple)): - return (f'Error: "{key}" must be a single integer; only one ' - f'block may be expanded per call. Issue a separate ' - f'extract_condensed call for each block you need.') - - # ``bool`` subclasses ``int`` (``int(True) == 1``) and ``float`` - # coerces silently (``int(1.9) == 1``); reject both up front. - if isinstance(raw, bool) or isinstance(raw, float): - return (f'Error: "{key}" must be an integer, got ' - f'{type(raw).__name__} {raw!r}.') - try: - n = int(raw) - except (TypeError, ValueError): - return f'Error: "{key}" must be an integer, got {raw!r}.' - - # Short existence check. Deliberately do NOT list every available - # id -- when the policy hallucinates a large range, echoing the - # full list back multiplies the error into thousands of tokens. - if n not in self._blocks: - if not self._blocks: - return f'Error: block {n} not found; no blocks available.' - ids = sorted(self._blocks) - return (f'Error: block {n} not found; valid block ids are ' - f'{ids}.') - - # Trajectory-bound idempotency. The raw text is already in the - # conversation as a prior tool response -- returning it again would - # just double the non-trainable footprint. - if n in self._already_expanded: - return (f'Block {n} was already expanded earlier in this ' - f'trajectory; re-read the previous tool response ' - f'instead of requesting it again.') - - value = self._blocks[n] - if value is None: - return (f'Error: block {n} has no original-text snapshot. ' - f'The upstream condenser must populate raw.original ' - f'before registering ExtractCondensed.') - - self._already_expanded.add(n) - return value - - # ------------------------------------------------------------------ - # Introspection helpers (handy for debugging / tests) - # ------------------------------------------------------------------ - @property - def blocks(self) -> List[int]: - """Sorted list of block indices available to this tool.""" - return sorted(self._blocks) - - def __len__(self) -> int: - return len(self._blocks) - - def __contains__(self, n: Any) -> bool: - try: - return int(n) in self._blocks - except (TypeError, ValueError): - return False diff --git a/src/twinkle_agentic/tools/tool_manager.py b/src/twinkle_agentic/tools/tool_manager.py index 46cbc907d..ea9bf63db 100644 --- a/src/twinkle_agentic/tools/tool_manager.py +++ b/src/twinkle_agentic/tools/tool_manager.py @@ -1,6 +1,7 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import json -from typing import Any, Dict, Iterable, List, Optional, Union +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Any, Dict, Iterable, List, Optional, Tuple, Union from twinkle.data_format import ToolCall from twinkle.data_format.message import Tool as ToolInfo @@ -19,6 +20,37 @@ def _extract_name(info: Any) -> Optional[str]: return None +def _unpack_tool_call(tool_call: Any) -> Tuple[Optional[str], Dict[str, Any], Optional[str]]: + """Split an OpenAI-shaped tool_call into ``(name, args, error)``. + + These dicts come from :meth:`twinkle.template.base.Template.parse_tool_call`. + ``error`` is set when the payload cannot be executed. + """ + if not isinstance(tool_call, dict): + return None, {}, f'Error: tool_call must be an object, got {type(tool_call).__name__}.' + fn = tool_call.get('function') + if not isinstance(fn, dict): + return None, {}, 'Error: tool_call missing "function" object.' + name = fn.get('name') + if not name: + return None, {}, 'Error: tool_call missing "function.name".' + raw_args = fn.get('arguments') + if raw_args is None: + return str(name), {}, None + if isinstance(raw_args, str): + try: + args = json.loads(raw_args) if raw_args.strip() else {} + except json.JSONDecodeError as e: + return str(name), {}, f'Error: invalid JSON in arguments: {e}' + if not isinstance(args, dict): + return str(name), {}, 'Error: "arguments" JSON must be an object.' + return str(name), args, None + if isinstance(raw_args, dict): + return str(name), raw_args, None + return None, {}, (f'Error: "arguments" must be a JSON string or object, ' + f'got {type(raw_args).__name__}.') + + class ToolManager: def __init__( @@ -65,33 +97,66 @@ def tool_infos(self) -> List[ToolInfo]: return [t.tool_info() for t in self._tools.values()] def __call__(self, tool_call: Union[ToolCall, Dict[str, Any]]) -> str: - if not isinstance(tool_call, dict): - return f'Error: tool_call must be an object, got {type(tool_call).__name__}.' - fn = tool_call.get('function') - if not isinstance(fn, dict): - return 'Error: tool_call missing "function" object.' - name = fn.get('name') - if not name: - return 'Error: tool_call missing "function.name".' + name, args, err = _unpack_tool_call(tool_call) + if err: + return err if (tool := self._tools.get(name)) is None: available = ', '.join(sorted(self._tools)) or '(none)' return f'Error: unknown tool {name!r}. Available: {available}.' - - raw_args = fn.get('arguments') - if raw_args is None: - args: Dict[str, Any] = {} - elif isinstance(raw_args, str): - try: - args = json.loads(raw_args) if raw_args.strip() else {} - except json.JSONDecodeError as e: - return f'Error: invalid JSON in arguments: {e}' - elif isinstance(raw_args, dict): - args = raw_args - else: - return (f'Error: "arguments" must be a JSON string or object, ' - f'got {type(raw_args).__name__}.') - try: return str(tool(name, args)) except Exception as e: # noqa return f'Error: tool {name!r} raised {type(e).__name__}: {e}' + + def call_many( + self, + tool_calls: Iterable[Union[ToolCall, Dict[str, Any]]], + max_workers: Optional[int] = None, + ) -> List[str]: + """Execute many tool calls, preserving input order. + + ``tool_calls`` are the OpenAI-shaped dicts produced by + :meth:`~twinkle.template.base.Template.parse_tool_call`. This method + unpacks them to ``(name, arguments)`` and, when every tool wraps the + same :class:`~twinkle_agentic.envs.base.Env`, dispatches through + ``Env.step_batch``. Otherwise a thread pool of :meth:`__call__`. + """ + calls = list(tool_calls) + if not calls: + return [] + if len(calls) == 1: + return [self(calls[0])] + + unpacked = [_unpack_tool_call(tc) for tc in calls] + env = self._shared_env() + can_batch = env is not None and all( + err is None and name in self._tools for name, _args, err in unpacked) + if can_batch: + try: + results = env.step_batch([(name, args) for name, args, _err in unpacked]) + return [ + r.observation if hasattr(r, 'observation') else str(r) for r in results + ] + except Exception: + pass + + workers = max_workers or min(32, len(calls)) + out: List[Optional[str]] = [None] * len(calls) + with ThreadPoolExecutor(max_workers=workers) as pool: + futs = {pool.submit(self, tc): i for i, tc in enumerate(calls)} + for fut in as_completed(futs): + out[futs[fut]] = fut.result() + return ['' if x is None else x for x in out] + + def _shared_env(self): + """Return the Env shared by every registered EnvTool, else None.""" + env = None + for tool in self._tools.values(): + wrapped = getattr(tool, '_env', None) + if wrapped is None: + return None + if env is None: + env = wrapped + elif wrapped is not env: + return None + return env diff --git a/src/twinkle_agentic/train/__init__.py b/src/twinkle_agentic/train/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/src/twinkle_agentic/train/base.py b/src/twinkle_agentic/train/base.py deleted file mode 100644 index a0fd281e3..000000000 --- a/src/twinkle_agentic/train/base.py +++ /dev/null @@ -1,6 +0,0 @@ - - -class Trainer: - - def train(self): - pass \ No newline at end of file diff --git a/src/twinkle_agentic/train/cron.py b/src/twinkle_agentic/train/cron.py deleted file mode 100644 index d118df8bc..000000000 --- a/src/twinkle_agentic/train/cron.py +++ /dev/null @@ -1,5 +0,0 @@ - - -class CronTrainManager: - - pass \ No newline at end of file diff --git a/src/twinkle_agentic/verifier/__init__.py b/src/twinkle_agentic/verifier/__init__.py index 38e14f601..eaf813963 100644 --- a/src/twinkle_agentic/verifier/__init__.py +++ b/src/twinkle_agentic/verifier/__init__.py @@ -1,32 +1,7 @@ -from .aggregation import (RoundScore, SegmentScore, TrajectoryScore, - aggregate_hard_over_rounds, aggregate_trajectory, - fuse_segment, scalar_to_level, - split_segment_into_rounds) -from .base import Verifier -from .domain_checks import (check_answer_match, check_code_parses, - check_instruction_constraints, check_not_degenerate, - check_numeric_equiv, check_output_format, - default_checks_for) -from .hard_scorer import CheckResult, HardScorer, HardScoreDetail, TrajectoryView -from .leak_verifier import LeakDetail, LeakVerifier -from .rubric_library import (INTENT_BASE_RUBRICS, INTENT_FIXED_RUBRICS, - default_intent_base_rubrics, - default_intent_fixed_rubrics) -from .rubric_verifier import (DiagnoseDetail, DiagnosisItem, RubricItem, - RubricVerifier, ScoreDetail) +from .result_check import (Check, CheckContext, CheckOutcome, CheckReport, + checks_from_dicts, local_runner, run_checks) __all__ = [ - 'Verifier', - 'RubricVerifier', 'RubricItem', 'ScoreDetail', - 'DiagnoseDetail', 'DiagnosisItem', - 'LeakVerifier', 'LeakDetail', - 'HardScorer', 'HardScoreDetail', 'CheckResult', 'TrajectoryView', - 'check_output_format', 'check_numeric_equiv', 'check_answer_match', - 'check_code_parses', 'check_instruction_constraints', 'check_not_degenerate', - 'default_checks_for', - 'RoundScore', 'SegmentScore', 'TrajectoryScore', - 'split_segment_into_rounds', 'aggregate_hard_over_rounds', 'fuse_segment', - 'aggregate_trajectory', 'scalar_to_level', - 'INTENT_BASE_RUBRICS', 'INTENT_FIXED_RUBRICS', - 'default_intent_base_rubrics', 'default_intent_fixed_rubrics', + 'Check', 'CheckContext', 'CheckOutcome', 'CheckReport', + 'run_checks', 'checks_from_dicts', 'local_runner', ] diff --git a/src/twinkle_agentic/verifier/aggregation.py b/src/twinkle_agentic/verifier/aggregation.py deleted file mode 100644 index c791ccb40..000000000 --- a/src/twinkle_agentic/verifier/aggregation.py +++ /dev/null @@ -1,274 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Multi-granularity score aggregation. - -Two scorers operate at deliberately different granularities: - -- :class:`HardScorer` is **per-round** — tool-call validity / execution / - protocol are facts about a single assistant turn, cheap and deterministic. -- :class:`RubricVerifier` is **per-segment** — soft quality (sub-goal progress, - reasoning soundness, no redundant calls) needs multi-round context. - -This module bridges the two without forcing either onto the other's grain: - - rounds ── HardScorer (per round) ──► h_1..h_R - │ aggregate over the rounds in a segment - segment ── RubricVerifier (whole) ──► rubric_scalar - │ fuse(hard_agg, rubric_scalar) - ▼ - segment score ──► aggregate ──► trajectory score - -It also implements the **short-circuit gate**: when a segment's hard score is -extreme (e.g. every tool call failed, no final answer), the soft rubric chain -is skipped entirely — saving the long/expensive LLM path exactly when the -answer is already decided. - -The functions here are pure: pass callables/detail objects, get scores back. -They do not import HardScorer/RubricVerifier, so they stay easily testable and -decoupled (the orchestrating TrajectoryScorer will wire real scorers in later). -""" -from __future__ import annotations - -import statistics -from dataclasses import dataclass, field -from typing import Any, Callable, Dict, List, Optional, Sequence - -from twinkle_agentic.segment.base import Segmenter - -NUM_LEVELS = 5 - -# aggregation strategies for combining a list of scalars in [0, 1] -_REDUCERS: Dict[str, Callable[[Sequence[float]], float]] = { - 'mean': lambda xs: sum(xs) / len(xs), - 'min': min, - 'max': max, - 'median': lambda xs: statistics.median(xs), - # geometric mean: harsher than mean, one bad round drags the whole segment - 'gmean': lambda xs: (statistics.geometric_mean([max(1e-6, x) for x in xs])), -} - - -def _reduce(xs: Sequence[float], how: str) -> float: - xs = [x for x in xs if x is not None] - if not xs: - return 0.0 - fn = _REDUCERS.get(how) - if fn is None: - raise ValueError(f'unknown reducer {how!r}; choose from {list(_REDUCERS)}') - return float(fn(xs)) - - -def scalar_to_level(scalar: float, num_levels: int = NUM_LEVELS) -> int: - scalar = min(1.0, max(0.0, scalar)) - return min(num_levels - 1, max(0, int(round(scalar * (num_levels - 1))))) - - -# --------------------------------------------------------------------------- -# score containers -# --------------------------------------------------------------------------- -@dataclass -class RoundScore: - """Per-round hard score.""" - index: int # round index within the segment - hard_scalar: float - gated: bool = False - detail: Any = None # optional HardScoreDetail - - -@dataclass -class SegmentScore: - """Fused per-segment score.""" - index: int - scalar: float # final fused score in [0, 1] - level: int - hard_scalar: float # aggregated hard score over the segment's rounds - rubric_scalar: Optional[float] # None if the soft chain was short-circuited - short_circuited: bool - n_rounds: int - rounds: List[RoundScore] = field(default_factory=list) - detail: Any = None # optional rubric ScoreDetail - - -@dataclass -class TrajectoryScore: - """Trajectory-level score aggregated over segments.""" - scalar: float - level: int - segments: List[SegmentScore] = field(default_factory=list) - - -# --------------------------------------------------------------------------- -# round -> segment -# --------------------------------------------------------------------------- -def split_segment_into_rounds(segment: dict) -> List[dict]: - """Split a segment sub-trajectory into per-round sub-trajectories. - - Reuses :meth:`Segmenter.split_turns`. The segment's preamble (system + the - first user message) is carried onto every round so a per-round HardScorer - still sees tools/context. Each returned dict is a valid sub-trajectory - (``messages`` + ``tools``/``user_data`` when present). - """ - messages = list(segment.get('messages', []) or []) - preamble, start = Segmenter._split_preamble(messages) - turns = Segmenter.split_turns(messages, start) - rounds: List[dict] = [] - for t in turns: - # skip pure 'user' boundary turns: nothing tool-verifiable there - if t.role_kind == 'user': - continue - r: Dict[str, Any] = {'messages': list(preamble) + list(t.messages)} - if segment.get('tools'): - r['tools'] = list(segment['tools']) - if segment.get('user_data'): - r['user_data'] = list(segment['user_data']) - rounds.append(r) - return rounds - - -def aggregate_hard_over_rounds( - round_scores: Sequence[RoundScore], - *, - how: str = 'gmean', -) -> float: - """Aggregate per-round hard scalars into one segment-level hard scalar. - - Default ``gmean`` (geometric mean) is intentionally harsher than plain mean: - a single badly-formed round meaningfully drags the segment, which matches - the intuition that one hallucinated/failed tool call hurts the sub-task. - """ - if not round_scores: - return 1.0 # no rounds to fault -> neutral (e.g. a text-only segment) - return _reduce([r.hard_scalar for r in round_scores], how) - - -# --------------------------------------------------------------------------- -# fusion (hard x soft) with short-circuit -# --------------------------------------------------------------------------- -def fuse_segment( - index: int, - round_scores: Sequence[RoundScore], - rubric_fn: Optional[Callable[[], Any]] = None, - *, - hard_agg: str = 'gmean', - fusion: str = 'product', - hard_floor: float = 0.25, - hard_ceil_skip: Optional[float] = None, - num_levels: int = NUM_LEVELS, -) -> SegmentScore: - """Fuse per-round hard scores with a (lazily computed) segment rubric score. - - Args: - index: segment index. - round_scores: per-round hard scores (already computed; cheap/code-only). - rubric_fn: zero-arg callable returning a rubric ScoreDetail (something - with a ``.scalar`` attribute) OR a float. Called ONLY when the soft - chain is not short-circuited — this is what saves the long LLM path. - hard_agg: reducer for per-round hard scores ('gmean'|'mean'|'min'|...). - fusion: how to combine hard_agg and rubric: - 'product' -> hard_agg * rubric (hard acts as a floor/gatekeeper) - 'hard_soft_blend' -> product, but when hard is high blend rubric toward - a floor so all-pass tool segments are not one-shot vetoed - 'min' -> min(hard_agg, rubric) - 'mean' -> (hard_agg + rubric)/2 - 'hard_only'-> ignore rubric entirely - hard_floor: if aggregated hard score < this, SHORT-CIRCUIT: skip rubric, - segment score = hard_agg (the answer is already decided as bad). - hard_ceil_skip: (optional) if aggregated hard score >= this, also skip - rubric and use hard_agg. Set None to disable. Useful for - trivially-good tool-only segments where soft quality adds little. - num_levels: level discretization. - """ - hard_agg_val = aggregate_hard_over_rounds(round_scores, how=hard_agg) - - short = False - rubric_scalar: Optional[float] = None - - if fusion == 'hard_only' or rubric_fn is None: - scalar = hard_agg_val - short = True - elif hard_agg_val < hard_floor: - # bad hard signal -> don't waste the soft chain - scalar = hard_agg_val - short = True - elif hard_ceil_skip is not None and hard_agg_val >= hard_ceil_skip: - scalar = hard_agg_val - short = True - else: - rubric_scalar = _rubric_scalar(rubric_fn()) - scalar = _combine(hard_agg_val, rubric_scalar, fusion) - - return SegmentScore( - index=index, - scalar=scalar, - level=scalar_to_level(scalar, num_levels), - hard_scalar=hard_agg_val, - rubric_scalar=rubric_scalar, - short_circuited=short, - n_rounds=len(round_scores), - rounds=list(round_scores), - ) - - -def _rubric_scalar(result: Any) -> float: - if result is None: - return 0.0 - if isinstance(result, (int, float)): - return float(result) - scalar = getattr(result, 'scalar', None) - return float(scalar) if scalar is not None else 0.0 - - -def _combine(hard: float, soft: float, fusion: str) -> float: - if fusion == 'product': - return hard * soft - if fusion == 'hard_soft_blend': - # When hard checks are strong (tool/format all pass), a harsh rubric on a - # long agent trace must not one-shot veto the segment (product → ~0.08). - if hard >= 0.9: - mix = 0.55 * soft + 0.45 - elif hard >= 0.75: - mix = 0.75 * soft + 0.25 - else: - mix = soft - return hard * mix - if fusion == 'min': - return min(hard, soft) - if fusion == 'mean': - return (hard + soft) / 2.0 - raise ValueError(f'unknown fusion {fusion!r}') - - -# --------------------------------------------------------------------------- -# segment -> trajectory -# --------------------------------------------------------------------------- -def aggregate_trajectory( - segment_scores: Sequence[SegmentScore], - *, - how: str = 'mean', - weight_by_rounds: bool = True, - num_levels: int = NUM_LEVELS, -) -> TrajectoryScore: - """Aggregate segment scores into a trajectory score. - - Args: - how: reducer when ``weight_by_rounds`` is False. - weight_by_rounds: when True, weight each segment by its round count so - longer sub-tasks count proportionally (ignores ``how``, uses a - round-weighted mean). Text-only segments count as weight 1. - """ - if not segment_scores: - return TrajectoryScore(scalar=0.0, level=0, segments=[]) - if weight_by_rounds: - num = 0.0 - den = 0.0 - for s in segment_scores: - w = max(1, s.n_rounds) - num += s.scalar * w - den += w - scalar = num / den if den else 0.0 - else: - scalar = _reduce([s.scalar for s in segment_scores], how) - return TrajectoryScore( - scalar=scalar, - level=scalar_to_level(scalar, num_levels), - segments=list(segment_scores), - ) diff --git a/src/twinkle_agentic/verifier/base.py b/src/twinkle_agentic/verifier/base.py deleted file mode 100644 index 534206d36..000000000 --- a/src/twinkle_agentic/verifier/base.py +++ /dev/null @@ -1,11 +0,0 @@ -from abc import ABC, abstractmethod - - -class Verifier(ABC): - """Reward verifier that scores a sample on a 5-level scale (0-4).""" - - NUM_LEVELS = 5 - - @abstractmethod - def __call__(self, trajectory: dict, **kwargs) -> int: - pass diff --git a/src/twinkle_agentic/verifier/domain_checks.py b/src/twinkle_agentic/verifier/domain_checks.py deleted file mode 100644 index 92ac59a64..000000000 --- a/src/twinkle_agentic/verifier/domain_checks.py +++ /dev/null @@ -1,436 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Domain-specific deterministic checks for :class:`HardScorer`. - -These are LLM-free, dependency-free (stdlib ``ast``/``re``/``json`` only) and -reuse the answer-extraction / F1 helpers already in ``reward/f1.py``. They are -factories: call them with config and get back a plain ``CheckFn`` that plugs -into ``HardScorer(checks=[...])``. - -Coverage (initial, no sandbox): -- output format: ``\\boxed{}`` / fenced code block / parseable JSON present -- numeric equivalence: lightweight fraction/decimal/percent normalization -- reference match: F1/EM vs ``ground_truth`` (reuses ``_f1_score``) -- code syntax: ``ast.parse`` on the last fenced block (stdlib, does NOT run) -- instruction constraints: length / keyword must-include / must-exclude / lang -- degeneration: empty / too-short / repetitive final answer - -Sandbox-based math (``math-verify``/sympy) and code execution (unit tests) can -be added later as additional CheckFns without touching HardScorer. -""" -from __future__ import annotations - -import ast -import json -import re -from typing import TYPE_CHECKING, Any, Callable, List, Optional, Sequence - -from twinkle_agentic.reward.f1 import _extract_final_answer -from twinkle_agentic.reward.f1 import _f1_score as _f1_score_stemmed - -from .hard_scorer import CheckResult, TrajectoryView - -if TYPE_CHECKING: - from .hard_scorer import CheckFn # noqa: F401 - - -# --------------------------------------------------------------------------- -# shared helpers -# --------------------------------------------------------------------------- -_CODE_FENCE_RE = re.compile(r'```([a-zA-Z0-9_+-]*)\s*\n(.*?)```', re.DOTALL) -_BOXED_RE = re.compile(r'\\boxed\s*\{') -_NUMBER_RE = re.compile(r'[-+]?\d*\.?\d+(?:/\d+)?%?') - - -def _ground_truths(trajectory: dict) -> List[str]: - """Read ground_truth values from user_data (same convention as F1Reward).""" - out: List[str] = [] - for entry in trajectory.get('user_data', []) or []: - if isinstance(entry, (list, tuple)) and len(entry) == 2 and entry[0] == 'ground_truth': - v = entry[1] - if isinstance(v, str): - try: - v = json.loads(v) - except (json.JSONDecodeError, ValueError): - pass - if isinstance(v, (list, tuple)): - out.extend(str(x) for x in v if x) - elif v: - out.append(str(v)) - return out - - -def _last_code_block(text: str) -> Optional[str]: - matches = _CODE_FENCE_RE.findall(text or '') - if not matches: - return None - return matches[-1][1] - - -def _to_number(token: str) -> Optional[float]: - """Normalize a numeric token: fraction 'a/b', percent 'x%', or decimal.""" - if token is None: - return None - s = str(token).strip().replace(',', '').replace('$', '').replace(' ', '') - if not s: - return None - percent = s.endswith('%') - if percent: - s = s[:-1] - try: - if '/' in s: - num, den = s.split('/', 1) - val = float(num) / float(den) - else: - val = float(s) - except (ValueError, ZeroDivisionError): - return None - return val / 100.0 if percent else val - - -def _numbers_in(text: str) -> List[float]: - vals = [] - for tok in _NUMBER_RE.findall(text or ''): - v = _to_number(tok) - if v is not None: - vals.append(v) - return vals - - -_PUNCT_RE = re.compile(r'[^\w\s]', re.UNICODE) -_ARTICLE_RE = re.compile(r'\b(a|an|the)\b') - - -def _f1_score(prediction: str, gold: str): - """F1/EM with graceful fallback when nltk (used by f1.py stemming) is - unavailable — degrade to a stemmer-free token F1 instead of crashing.""" - try: - return _f1_score_stemmed(prediction, gold) - except ImportError: - pass - from collections import Counter - norm = lambda s: _ARTICLE_RE.sub( # noqa: E731 - ' ', _PUNCT_RE.sub('', (s or '').lower())).split() - p_tok, g_tok = norm(prediction), norm(gold) - if not p_tok or not g_tok: - em = float(p_tok == g_tok) - return em, em - em = float(p_tok == g_tok) - common = Counter(p_tok) & Counter(g_tok) - same = sum(common.values()) - if same == 0: - return 0.0, em - prec, rec = same / len(p_tok), same / len(g_tok) - return 2 * prec * rec / (prec + rec), em - - -# --------------------------------------------------------------------------- -# format / structure -# --------------------------------------------------------------------------- -def check_output_format(fmt: str, *, weight: float = 1.5, critical: bool = True - ) -> 'CheckFn': - """Require a specific output artifact in the final answer. - - Args: - fmt: one of ``'boxed'`` (\\boxed{...}), ``'code'`` (fenced block), - ``'json'`` (a parseable JSON object/array anywhere in the answer). - """ - fmt = fmt.lower() - if fmt not in ('boxed', 'code', 'json'): - raise ValueError("fmt must be 'boxed', 'code' or 'json'") - - def _check(view: TrajectoryView) -> CheckResult: - text = view.last_assistant_text() - if fmt == 'boxed': - ok = bool(_extract_final_answer(text)) or bool(_BOXED_RE.search(text)) - detail = 'boxed present' if ok else 'no \\boxed{}' - elif fmt == 'code': - ok = _last_code_block(text) is not None - detail = 'code block present' if ok else 'no code block' - else: # json - block = _last_code_block(text) or text - ok = _has_parseable_json(block) - detail = 'json parseable' if ok else 'no parseable json' - return CheckResult(f'format_{fmt}', 1.0 if ok else 0.0, weight, - critical=critical, n=1, detail=detail) - - return _check - - -def _has_parseable_json(text: str) -> bool: - s = (text or '').strip() - if not s: - return False - # try whole-string first, then first {...}/[...] span - for candidate in (s, _first_bracket_span(s)): - if not candidate: - continue - try: - json.loads(candidate) - return True - except (json.JSONDecodeError, ValueError): - continue - return False - - -def _first_bracket_span(s: str) -> Optional[str]: - starts = [i for i, c in enumerate(s) if c in '{['] - if not starts: - return None - i = starts[0] - open_c = s[i] - close_c = '}' if open_c == '{' else ']' - depth = 0 - for j in range(i, len(s)): - if s[j] == open_c: - depth += 1 - elif s[j] == close_c: - depth -= 1 - if depth == 0: - return s[i:j + 1] - return None - - -# --------------------------------------------------------------------------- -# numeric equivalence (lightweight, no sympy) -# --------------------------------------------------------------------------- -def check_numeric_equiv(*, tol: float = 1e-6, weight: float = 2.0, - critical: bool = False) -> 'CheckFn': - """Compare the extracted final number(s) against ground_truth numerically. - - Handles fractions / decimals / percentages. For symbolic equivalence, - swap this for a sympy/math-verify CheckFn later. Neutral pass when there - is no numeric ground truth to compare against. - """ - def _check(view: TrajectoryView) -> CheckResult: - golds = _ground_truths(view.trajectory) - gold_nums = [n for g in golds for n in _numbers_in(g)] - if not gold_nums: - return CheckResult('numeric_equiv', 1.0, weight, critical=False, n=0, - detail='no numeric ground truth') - text = view.last_assistant_text() - boxed = _extract_final_answer(text) - pred_nums = _numbers_in(boxed) if boxed else _numbers_in(text) - if not pred_nums: - return CheckResult('numeric_equiv', 0.0, weight, critical=critical, n=1, - detail='no number in answer') - # match if any predicted number equals any gold (last pred preferred) - target = gold_nums[-1] - ok = any(abs(p - target) <= tol + tol * abs(target) for p in pred_nums) - return CheckResult('numeric_equiv', 1.0 if ok else 0.0, weight, - critical=critical, n=1, - detail=f'pred~{pred_nums[-1]} vs gold~{target}') - - return _check - - -# --------------------------------------------------------------------------- -# reference match (reuse f1.py) -# --------------------------------------------------------------------------- -def check_answer_match(*, threshold: float = 0.6, weight: float = 2.0, - critical: bool = False, use_em: bool = False) -> 'CheckFn': - """F1/EM of the extracted answer vs ground_truth (reuses ``_f1_score``). - - Score is the max F1 over gold answers (or EM when ``use_em``); pass/fail is - F1 >= threshold. Neutral pass when there is no ground truth. - """ - def _check(view: TrajectoryView) -> CheckResult: - golds = _ground_truths(view.trajectory) - if not golds: - return CheckResult('answer_match', 1.0, weight, critical=False, n=0, - detail='no ground truth') - text = view.last_assistant_text() - boxed = _extract_final_answer(text) - pred = boxed or text - scored = [_f1_score(pred, g) for g in golds] - best_f1 = max(f for f, _ in scored) - best_em = max(e for _, e in scored) - # Containment fallback: when the answer isn't boxed, a short gold that - # appears verbatim in the answer counts as a hit (robust to preamble - # like "The answer is Paris."). - contained = False - if not boxed: - low = text.lower() - contained = any(g.strip() and g.lower() in low and len(g.split()) <= 6 - for g in golds) - if contained: - best_f1 = max(best_f1, 1.0) - best_em = max(best_em, 1.0) - val = best_em if use_em else best_f1 - return CheckResult('answer_match', val, weight, critical=critical, n=1, - detail=f'f1={best_f1:.2f} em={best_em:.0f}' - + (' contained' if contained else '') - + ('' if val >= threshold else ' <thr')) - - return _check - - -# --------------------------------------------------------------------------- -# code syntax (stdlib ast, does NOT execute) -# --------------------------------------------------------------------------- -def check_code_parses(*, language: str = 'python', weight: float = 1.5, - critical: bool = False) -> 'CheckFn': - """Last fenced code block must parse (Python only, via stdlib ``ast``). - - This validates *syntax* without a sandbox and without running anything. - Non-Python blocks are a neutral pass (we can't cheaply verify them here). - """ - def _check(view: TrajectoryView) -> CheckResult: - text = view.last_assistant_text() - block = _last_code_block(text) - if block is None: - return CheckResult('code_parses', 0.0, weight, critical=critical, n=1, - detail='no code block') - if language.lower() != 'python': - return CheckResult('code_parses', 1.0, weight, critical=False, n=0, - detail=f'{language} not statically checked') - try: - ast.parse(block) - return CheckResult('code_parses', 1.0, weight, critical=critical, n=1, - detail='parses') - except SyntaxError as e: - return CheckResult('code_parses', 0.0, weight, critical=critical, n=1, - detail=f'SyntaxError: {e.msg}') - - return _check - - -# --------------------------------------------------------------------------- -# instruction constraints (IFEval-style, pure code) -# --------------------------------------------------------------------------- -def check_instruction_constraints( - *, - min_words: Optional[int] = None, - max_words: Optional[int] = None, - must_include: Optional[Sequence[str]] = None, - must_exclude: Optional[Sequence[str]] = None, - match_source_language: bool = False, - weight: float = 1.0, - critical: bool = False, -) -> 'CheckFn': - """Verify code-checkable instruction-following constraints on the answer. - - Score is the fraction of active sub-constraints satisfied. - """ - must_include = list(must_include or []) - must_exclude = list(must_exclude or []) - - def _check(view: TrajectoryView) -> CheckResult: - text = view.last_assistant_text() - words = text.split() - n_words = len(words) - checks: List[bool] = [] - notes: List[str] = [] - - if min_words is not None: - ok = n_words >= min_words - checks.append(ok) - if not ok: - notes.append(f'words<{min_words}') - if max_words is not None: - ok = n_words <= max_words - checks.append(ok) - if not ok: - notes.append(f'words>{max_words}') - low = text.lower() - for kw in must_include: - ok = kw.lower() in low - checks.append(ok) - if not ok: - notes.append(f'missing:{kw}') - for kw in must_exclude: - ok = kw.lower() not in low - checks.append(ok) - if not ok: - notes.append(f'forbidden:{kw}') - if match_source_language: - ok = _language_matches(view) - checks.append(ok) - if not ok: - notes.append('lang-mismatch') - - if not checks: - return CheckResult('instruction_constraints', 1.0, weight, - critical=False, n=0, detail='no active constraints') - score = sum(1 for c in checks if c) / len(checks) - return CheckResult('instruction_constraints', score, weight, - critical=critical, n=len(checks), - detail=', '.join(notes) or 'all satisfied') - - return _check - - -def _cjk_ratio(text: str) -> float: - if not text: - return 0.0 - cjk = sum(1 for c in text if '\u4e00' <= c <= '\u9fff' - or '\u3040' <= c <= '\u30ff' - or '\uac00' <= c <= '\ud7a3') - return cjk / len(text) - - -def _language_matches(view: TrajectoryView) -> bool: - """Cheap heuristic: answer's CJK-ness matches the first user message's.""" - user_text = '' - for m in view.messages: - if m.get('role') == 'user': - user_text = view.text_of(m) - break - ans = view.last_assistant_text() - if not user_text or not ans: - return True - return abs(_cjk_ratio(user_text[:400]) - _cjk_ratio(ans[:400])) < 0.3 - - -# --------------------------------------------------------------------------- -# degeneration -# --------------------------------------------------------------------------- -def check_not_degenerate(*, min_chars: int = 1, max_repeat_ratio: float = 0.5, - ngram: int = 8, weight: float = 1.0, - critical: bool = False) -> 'CheckFn': - """Fail on empty / trivially short / highly repetitive final answers.""" - def _check(view: TrajectoryView) -> CheckResult: - text = view.last_assistant_text().strip() - if len(text) < min_chars: - return CheckResult('not_degenerate', 0.0, weight, critical=critical, - n=1, detail='too short/empty') - rep = _repetition_ratio(text, ngram) - ok = rep <= max_repeat_ratio - return CheckResult('not_degenerate', 1.0 if ok else 0.0, weight, - critical=critical, n=1, - detail=f'repeat={rep:.2f}' + ('' if ok else ' >thr')) - - return _check - - -def _repetition_ratio(text: str, ngram: int) -> float: - if _cjk_ratio(text[:500]) > 0.3: - tokens = [c for c in text if not c.isspace()] - else: - tokens = text.split() - if len(tokens) < ngram: - return 0.0 - grams = [tuple(tokens[i:i + ngram]) for i in range(len(tokens) - ngram + 1)] - if not grams: - return 0.0 - return 1.0 - len(set(grams)) / len(grams) - - -# Convenience presets keyed by domain, for use with a router later. -def default_checks_for(domain: str) -> List['CheckFn']: - """Return a reasonable initial check bundle for a domain (no sandbox).""" - domain = (domain or '').lower() - if domain == 'math': - return [check_output_format('boxed', critical=False), - check_numeric_equiv(), - check_not_degenerate()] - if domain == 'code': - return [check_output_format('code', critical=False), - check_code_parses(), - check_not_degenerate()] - if domain in ('factual_qa', 'factual', 'qa'): - return [check_answer_match(), - check_not_degenerate()] - if domain in ('open_qa', 'open', 'writing'): - return [check_instruction_constraints(), - check_not_degenerate()] - return [check_not_degenerate()] diff --git a/src/twinkle_agentic/verifier/hard_scorer.py b/src/twinkle_agentic/verifier/hard_scorer.py deleted file mode 100644 index bc4dca126..000000000 --- a/src/twinkle_agentic/verifier/hard_scorer.py +++ /dev/null @@ -1,444 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Deterministic (LLM-free) hard scorer for agentic trajectories. - -Where :class:`RubricVerifier` judges *soft* quality via an LLM, this scorer -judges *hard* facts with plain code: did the agent call declared tools, were -the arguments valid, did the calls execute, is the OpenAI tool protocol -consistent, did the run terminate cleanly, is there a final answer, and did it -avoid degenerate repetition. These signals are **free** and, crucially, -**un-hackable by the policy** — the policy cannot talk its way past a JSON -parse error or a hallucinated tool name. - -The score is a weighted mean of independent checks, each producing a -``CheckResult`` in ``[0, 1]``. Two aggregation modes: - -- ``mode='mean'`` (default): weighted average of all checks. -- ``mode='gate'``: any *critical* check that scores 0 caps the whole score at - 0 (a strict gatekeeper — one hallucinated tool call fails the segment). - -Checks are pluggable: pass your own callables to extend/override. The public -``__call__`` returns an ``int`` in ``[0, NUM_LEVELS)`` per the -:class:`Verifier` contract; ``score_detail`` returns the full breakdown. -""" -from __future__ import annotations - -import json -import re -from dataclasses import dataclass, field -from typing import Any, Callable, Dict, List, Optional, Tuple - -from .base import Verifier - -# A check takes the parsed trajectory view and returns a CheckResult. -CheckFn = Callable[['TrajectoryView'], 'CheckResult'] - -_ERROR_PREFIX_RE = re.compile(r'^\s*(error|exception|traceback|failed)\b[:\s]', re.IGNORECASE) - - -@dataclass -class CheckResult: - name: str - score: float # in [0, 1] - weight: float - critical: bool # if True and score==0, gate mode caps total at 0 - n: int = 0 # number of items this check evaluated - detail: str = '' - - def __post_init__(self): - self.score = min(1.0, max(0.0, float(self.score))) - - -@dataclass -class HardScoreDetail: - level: int - scalar: float # continuous hard score in [0, 1] - gated: bool # a critical check zeroed the score (gate mode) - checks: List[CheckResult] = field(default_factory=list) - - def as_dict(self) -> Dict[str, Any]: - return { - 'level': self.level, - 'scalar': self.scalar, - 'gated': self.gated, - 'checks': {c.name: {'score': c.score, 'weight': c.weight, - 'critical': c.critical, 'n': c.n, 'detail': c.detail} - for c in self.checks}, - } - - -@dataclass -class _ToolCall: - name: Optional[str] - raw_args: Any - call_id: Optional[str] - msg_index: int - - -class TrajectoryView: - """Parsed, check-friendly view over a trajectory segment. - - Precomputes the message list, the assistant tool_calls, the tool-result - messages and the declared tool schema so individual checks stay cheap and - don't each re-walk the messages. - """ - - def __init__(self, trajectory: dict): - self.trajectory = trajectory or {} - self.messages: List[dict] = list(self.trajectory.get('messages', []) or []) - self.tools: List[dict] = list(self.trajectory.get('tools', []) or []) - - # declared tool names + parameter schemas - self.declared_names: set = set() - self.declared_required: Dict[str, List[str]] = {} - for t in self.tools: - fn = t.get('function') if isinstance(t, dict) else None - if not isinstance(fn, dict): - continue - name = fn.get('name') - if not isinstance(name, str) or not name: - continue - self.declared_names.add(name) - params = fn.get('parameters') - if isinstance(params, dict): - req = params.get('required') - if isinstance(req, list): - self.declared_required[name] = [r for r in req if isinstance(r, str)] - - # assistant tool calls, in order - self.tool_calls: List[_ToolCall] = [] - for i, m in enumerate(self.messages): - if m.get('role') != 'assistant': - continue - for tc in (m.get('tool_calls') or []): - if not isinstance(tc, dict): - continue - fn = tc.get('function') or {} - self.tool_calls.append(_ToolCall( - name=fn.get('name') if isinstance(fn, dict) else None, - raw_args=fn.get('arguments') if isinstance(fn, dict) else None, - call_id=tc.get('id'), - msg_index=i, - )) - - # tool-result messages, indexed by tool_call_id where present - self.tool_msgs_by_id: Dict[str, dict] = {} - self.tool_msgs: List[dict] = [] - for m in self.messages: - if m.get('role') == 'tool': - self.tool_msgs.append(m) - cid = m.get('tool_call_id') - if isinstance(cid, str) and cid: - self.tool_msgs_by_id[cid] = m - - # -- shared helpers reused by checks -- - def parsed_args(self, tc: _ToolCall) -> Optional[dict]: - raw = tc.raw_args - if isinstance(raw, dict): - return raw - if raw is None: - return {} - if isinstance(raw, str): - s = raw.strip() - if not s: - return {} - try: - v = json.loads(s) - return v if isinstance(v, dict) else None - except (json.JSONDecodeError, ValueError): - return None - return None - - def result_for(self, tc: _ToolCall) -> Optional[dict]: - """Find the tool result for a call: prefer id match, else next tool msg.""" - if tc.call_id and tc.call_id in self.tool_msgs_by_id: - return self.tool_msgs_by_id[tc.call_id] - # positional fallback: first tool message after the call's assistant msg - for j in range(tc.msg_index + 1, len(self.messages)): - m = self.messages[j] - role = m.get('role') - if role == 'tool': - return m - if role == 'assistant': - break - return None - - @staticmethod - def text_of(msg: Optional[dict]) -> str: - if not msg: - return '' - content = msg.get('content') - if isinstance(content, list): - return '\n'.join(p.get('text', '') for p in content - if isinstance(p, dict) and p.get('type') == 'text') - return content if isinstance(content, str) else '' - - def last_assistant_text(self) -> str: - for m in reversed(self.messages): - if m.get('role') == 'assistant': - return self.text_of(m) - return '' - - def last_assistant_msg(self) -> Optional[dict]: - for m in reversed(self.messages): - if m.get('role') == 'assistant': - return m - return None - - -# --------------------------------------------------------------------------- -# Individual checks (pure, deterministic) -# --------------------------------------------------------------------------- -def check_args_valid_json(view: TrajectoryView) -> CheckResult: - """Every tool call's arguments must parse as a JSON object.""" - calls = view.tool_calls - if not calls: - return CheckResult('args_valid_json', 1.0, 1.0, critical=True, n=0, - detail='no tool calls') - ok = sum(1 for tc in calls if view.parsed_args(tc) is not None) - return CheckResult('args_valid_json', ok / len(calls), 1.0, critical=True, - n=len(calls), detail=f'{ok}/{len(calls)} valid') - - -def check_tool_declared(view: TrajectoryView) -> CheckResult: - """Called tools must be in the declared tool set (no hallucinated tools).""" - calls = view.tool_calls - if not calls or not view.declared_names: - # can't verify without a declared schema -> neutral pass, non-critical - return CheckResult('tool_declared', 1.0, 1.0, critical=False, n=0, - detail='no tools declared or no calls') - ok = sum(1 for tc in calls if tc.name in view.declared_names) - return CheckResult('tool_declared', ok / len(calls), 1.5, critical=True, - n=len(calls), detail=f'{ok}/{len(calls)} declared') - - -def check_required_args(view: TrajectoryView) -> CheckResult: - """Parsed arguments must contain the schema's required fields.""" - calls = [tc for tc in view.tool_calls if tc.name in view.declared_required] - if not calls: - return CheckResult('required_args', 1.0, 1.0, critical=False, n=0, - detail='no schema-required fields to check') - ok = 0 - for tc in calls: - args = view.parsed_args(tc) - if args is None: - continue - req = view.declared_required.get(tc.name, []) - if all(r in args for r in req): - ok += 1 - return CheckResult('required_args', ok / len(calls), 1.0, critical=False, - n=len(calls), detail=f'{ok}/{len(calls)} complete') - - -def check_tool_executed(view: TrajectoryView) -> CheckResult: - """Each tool call must have a non-empty, non-error result message.""" - calls = view.tool_calls - if not calls: - return CheckResult('tool_executed', 1.0, 1.0, critical=False, n=0, - detail='no tool calls') - ok = 0 - for tc in calls: - res = view.result_for(tc) - text = view.text_of(res).strip() - if text and not _ERROR_PREFIX_RE.match(text): - ok += 1 - return CheckResult('tool_executed', ok / len(calls), 2.0, critical=False, - n=len(calls), detail=f'{ok}/{len(calls)} succeeded') - - -def check_protocol_pairing(view: TrajectoryView) -> CheckResult: - """OpenAI protocol: every tool_call id should have a matching tool msg, and - every tool msg should reference a known call id (when ids are used).""" - calls = view.tool_calls - if not calls: - return CheckResult('protocol_pairing', 1.0, 1.0, critical=False, n=0, - detail='no tool calls') - call_ids = {tc.call_id for tc in calls if tc.call_id} - if not call_ids: - # ids not used in this trace; fall back to counting result coverage - paired = sum(1 for tc in calls if view.result_for(tc) is not None) - return CheckResult('protocol_pairing', paired / len(calls), 1.0, - critical=False, n=len(calls), - detail=f'{paired}/{len(calls)} have a result (no ids)') - matched_calls = sum(1 for cid in call_ids if cid in view.tool_msgs_by_id) - # orphan tool messages referencing unknown ids - orphans = sum(1 for m in view.tool_msgs - if isinstance(m.get('tool_call_id'), str) - and m['tool_call_id'] not in call_ids) - total = len(call_ids) + orphans - score = matched_calls / total if total else 1.0 - return CheckResult('protocol_pairing', score, 1.0, critical=False, - n=len(call_ids), - detail=f'{matched_calls}/{len(call_ids)} paired, {orphans} orphan tool msgs') - - -def check_no_repeated_calls(view: TrajectoryView) -> CheckResult: - """Penalize degenerate tool-call loops. - - Two independent signals, worst one wins: - 1. exact-duplicate ``(name, args)`` calls — classic redundant repetition; - 2. single-tool domination — one tool name fired over and over (even with - *different* args), the "spin the same tool forever" failure that (1) - misses because the arguments differ each time. Only kicks in once there - are enough calls (``>= _SPIN_MIN_CALLS``) so a legitimate 3-4 step loop - of the same tool is not punished. - """ - calls = view.tool_calls - if len(calls) < 2: - return CheckResult('no_repeated_calls', 1.0, 1.0, critical=False, - n=len(calls), detail='fewer than 2 calls') - - seen: set = set() - dupes = 0 - name_counts: Dict[str, int] = {} - for tc in calls: - args = view.parsed_args(tc) - key = (tc.name, json.dumps(args, sort_keys=True) if isinstance(args, dict) else str(tc.raw_args)) - if key in seen: - dupes += 1 - else: - seen.add(key) - name_counts[tc.name or ''] = name_counts.get(tc.name or '', 0) + 1 - dup_score = 1.0 - dupes / len(calls) - - # Single-tool domination: one tool fired over and over (a spin loop). This - # is only a *soft* signal — a legitimate agent may batch-read 8 files with - # the same tool — so it is deliberately lenient: it only triggers on long - # sequences that are almost entirely one tool, and it floors the penalty so - # a batch operation is nudged down, not zeroed. Real dead-loops (empty - # repeated spins) get further penalized by the rubric / final-answer checks. - _SPIN_MIN_CALLS = 8 - _SPIN_TOLERATED_SHARE = 0.8 - _SPIN_FLOOR = 0.4 - spin_score = 1.0 - top_name, top_n = max(name_counts.items(), key=lambda kv: kv[1]) - top_share = top_n / len(calls) - if len(calls) >= _SPIN_MIN_CALLS and top_share > _SPIN_TOLERATED_SHARE: - # Linearly map (tolerated..1.0] share onto (1.0.._SPIN_FLOOR] score. - frac = (top_share - _SPIN_TOLERATED_SHARE) / (1.0 - _SPIN_TOLERATED_SHARE) - spin_score = max(_SPIN_FLOOR, 1.0 - (1.0 - _SPIN_FLOOR) * frac) - - score = min(dup_score, spin_score) - detail = f'{dupes} duplicate calls' - if spin_score < dup_score: - detail = (f"tool '{top_name}' dominates {top_n}/{len(calls)} " - f'calls ({top_share:.0%})') - return CheckResult('no_repeated_calls', score, 1.0, critical=False, - n=len(calls), detail=detail) - - -def check_clean_termination(view: TrajectoryView) -> CheckResult: - """The trajectory should end on an assistant answer, not a dangling tool - call or a length-truncated turn.""" - last = view.last_assistant_msg() - if last is None: - return CheckResult('clean_termination', 0.0, 1.0, critical=False, n=1, - detail='no assistant message') - # last message overall should be the assistant answer (no trailing tool call - # left unanswered / no pending tool msg after it) - last_role = view.messages[-1].get('role') if view.messages else None - finish = last.get('finish_reason') - truncated = finish == 'length' - dangling = last_role == 'assistant' and bool(last.get('tool_calls')) - ok = (not truncated) and (not dangling) and (last_role in ('assistant', 'tool')) - detail = [] - if truncated: - detail.append('length-truncated') - if dangling: - detail.append('dangling tool_call') - return CheckResult('clean_termination', 1.0 if ok else 0.0, 1.0, - critical=False, n=1, detail=', '.join(detail) or 'clean') - - -def check_final_answer(view: TrajectoryView) -> CheckResult: - """There must be a non-empty final assistant answer.""" - text = view.last_assistant_text().strip() - return CheckResult('final_answer', 1.0 if text else 0.0, 1.5, - critical=False, n=1, - detail=f'{len(text)} chars' if text else 'empty') - - -DEFAULT_CHECKS: Tuple[CheckFn, ...] = ( - check_args_valid_json, - check_tool_declared, - check_required_args, - check_tool_executed, - check_protocol_pairing, - check_no_repeated_calls, - check_clean_termination, - check_final_answer, -) - - -# --------------------------------------------------------------------------- -# Scorer -# --------------------------------------------------------------------------- -class HardScorer(Verifier): - """Deterministic hard score for an agentic trajectory segment. - - Args: - checks: Ordered check callables. Defaults to :data:`DEFAULT_CHECKS`. - Pass your own to extend or replace. - mode: ``'mean'`` (weighted average) or ``'gate'`` (a critical check - scoring 0 zeroes the total). - weights: Optional ``{check_name: weight}`` overrides. - """ - - def __init__( - self, - checks: Optional[List[CheckFn]] = None, - *, - mode: str = 'mean', - weights: Optional[Dict[str, float]] = None, - gate_threshold: float = 1.0, - ): - if mode not in ('mean', 'gate'): - raise ValueError("mode must be 'mean' or 'gate'") - if not 0.0 <= gate_threshold <= 1.0: - raise ValueError('gate_threshold must be in [0, 1]') - self.checks: List[CheckFn] = list(checks) if checks is not None else list(DEFAULT_CHECKS) - self.mode = mode - self.weights = dict(weights or {}) - # In gate mode, a critical check scoring BELOW this threshold zeroes the - # total. 1.0 = any violation gates (strict); 0.0 = only total failure. - self.gate_threshold = float(gate_threshold) - - def __call__(self, trajectory: dict, **kwargs) -> int: - return self.score_detail(trajectory, **kwargs).level - - def score_detail(self, trajectory: dict, **kwargs) -> HardScoreDetail: - view = TrajectoryView(trajectory) - results: List[CheckResult] = [] - for fn in self.checks: - r = fn(view) - if r.name in self.weights: - r.weight = float(self.weights[r.name]) - results.append(r) - - gated = False - if self.mode == 'gate': - # A critical check that is not fully satisfied gates the segment: - # a single hallucinated tool or JSON parse error fails the whole - # thing, regardless of how many other calls were fine. - for r in results: - if r.critical and r.n > 0 and r.score < self.gate_threshold: - gated = True - break - - if gated: - scalar = 0.0 - else: - num = sum(r.score * r.weight for r in results) - den = sum(r.weight for r in results) - scalar = num / den if den else 0.0 - - return HardScoreDetail( - level=self._to_level(scalar), - scalar=scalar, - gated=gated, - checks=results, - ) - - def _to_level(self, scalar: float) -> int: - scalar = min(1.0, max(0.0, scalar)) - level = int(round(scalar * (self.NUM_LEVELS - 1))) - return min(self.NUM_LEVELS - 1, max(0, level)) diff --git a/src/twinkle_agentic/verifier/leak_verifier.py b/src/twinkle_agentic/verifier/leak_verifier.py deleted file mode 100644 index b692d5f80..000000000 --- a/src/twinkle_agentic/verifier/leak_verifier.py +++ /dev/null @@ -1,363 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""General answer-leak verifier for auxiliary hints / skills / notes. - -Decides whether a piece of *auxiliary content* — a hint, a distilled skill, a -retrieved note, a rationale — that will be shown to a solver ALONGSIDE a task -leaks the answer (or an essentially complete solution) to THAT task. A good hint -carries only transferable strategy ("factor the radicand before adding roots"); -a leaking hint hands over the result, a task-specific decisive step, or a full -derivation. - -**One layer only: an LLM judge.** Leak detection is inherently *semantic* — the -common case is a hint that describes the solution structure without ever writing -the answer, which no string rule can catch. A cheap deterministic pre-filter -(verbatim answer / answer-number matching) was measured to replace only a few -percent of the judge's catches, is not domain-general, and false-flags short -answers ("D", "1") that appear in almost any text. So it is deliberately absent: -this verifier is a single ``llm_backup``-distilled LEAK / CLEAN judge that runs -on the student model when confident and falls back to the teacher API otherwise -— the same progressive-distillation path the other verifiers here use. - -The :class:`Verifier` contract (``__call__ -> int``) maps CLEAN -> ``NUM_LEVELS-1`` -(a clean hint is the "good" end of the reward scale) and LEAK -> ``0``. -:meth:`leak_detail` exposes the boolean + reason, and :meth:`leak_batch` runs -many checks in parallel (the judge dominates latency). -""" -from __future__ import annotations - -import os -from concurrent.futures import ThreadPoolExecutor -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Union - -from twinkle_agentic.utils.llm_backup import llm_backup - -from .base import Verifier -from .domain_checks import _ground_truths - -if TYPE_CHECKING: - from twinkle.data_format.sampling import SamplingParams # noqa: F401 - from twinkle.sampler.base import Sampler # noqa: F401 - - -# --------------------------------------------------------------------------- -# Prompts (task-agnostic; a reference answer is optional context for the judge) -# --------------------------------------------------------------------------- -_JUDGE_SYSTEM = """\ -You check whether a HINT that will be shown to someone solving a TASK leaks the \ -answer. - -A good hint gives ONLY transferable strategy — general methods, common \ -pitfalls, or sanity checks that would help on a whole class of similar tasks. - -The hint LEAKS if it does ANY of: -- reveals the final answer or final result (a value, expression, choice, label, \ -or verbatim output); -- states a specific decisive intermediate result or fact that is unique to THIS \ -task; -- gives a derivation or step-by-step plan that essentially solves THIS task. - -The hint does NOT leak if it only names general methods, common mistakes, or \ -checks that are not specific to this task's answer. - -Reply with exactly one word: LEAK or CLEAN.""" - -# Answer-only variant: flags ONLY final-answer disclosure and explicitly PERMITS -# revealing the method / approach / step plan. For skill-generation training, the -# method IS the product we want, so the default prompt's "decisive intermediate -# result" and "step-by-step plan that solves this task" clauses over-flag good -# skills; this variant drops them. -_JUDGE_SYSTEM_ANSWER_ONLY = """\ -You check whether a HINT that will be shown to someone solving a TASK reveals the -TASK's FINAL ANSWER. - -The hint may FREELY describe general methods, the solution approach, which -technique to apply, the steps to take, common pitfalls, or sanity checks — even -if that strongly points at how to solve THIS task. Revealing the METHOD or PLAN -is allowed and expected of a good hint. - -The hint LEAKS only if it reveals the FINAL ANSWER itself — the concrete final -value, expression, choice, label, or verbatim result the task asks for (or a -trivially equivalent restatement of it). Describing HOW to get there WITHOUT -stating the resulting final value does NOT leak. - -Reply with exactly one word: LEAK or CLEAN.""" - -_JUDGE_USER = """\ -## Task -{query} -{reference_block} -## Hint to check -{hint} - -Does the hint leak the answer or an essentially complete solution to THIS task? \ -Reply LEAK or CLEAN.""" - -_REFERENCE_BLOCK = """ -## Known answer (for your judgement only; do not treat its wording as the hint) -{reference} -""" - - -# --------------------------------------------------------------------------- -# Result holder -# --------------------------------------------------------------------------- -@dataclass -class LeakDetail: - """Outcome of one leak check. - - Attributes: - leaked: True if the content leaks the answer/solution. - reason: machine-readable reason — ``''`` (clean), ``'llm_leak'``, - ``'llm_uncertain'`` (judge reply unparseable), or ``'no_llm'`` (no - student sampler and no teacher API configured). - source: which layer decided — ``'llm'`` | ``'none'``. - """ - leaked: bool - reason: str - source: str - - -# --------------------------------------------------------------------------- -# Verdict parsing / comparator (for the judge + its distillation) -# --------------------------------------------------------------------------- -def _verdict_of(raw: str) -> Optional[bool]: - """Parse a LEAK/CLEAN reply -> True (leak) / False (clean) / None (unclear).""" - v = (raw or '').strip().upper() - if 'CLEAN' in v: - return False - if 'LEAK' in v: - return True - return None - - -def _verdict_close(a: str, b: str) -> bool: - """llm_backup comparator: student/teacher agree iff same LEAK/CLEAN verdict.""" - va, vb = _verdict_of(a), _verdict_of(b) - if va is None or vb is None: - return (a or '').strip() == (b or '').strip() - return va == vb - - -# --------------------------------------------------------------------------- -# Verifier -# --------------------------------------------------------------------------- -class LeakVerifier(Verifier): - """Verify that an auxiliary hint does not leak the answer to its task. - - Args: - sampler: student model sampler (local inference). If ``None`` the judge - is served entirely by the teacher API via ``llm_backup`` (useful - before a student exists); if no teacher is configured either, the - verifier reports ``no_llm`` and cannot judge. - model_path: identifier (bookkeeping only). - sampling_params: default sampling params for the judge call. - judge_lora_path: LoRA adapter for the distilled judge student. - max_content_chars / max_query_chars: truncation caps for the judge input. - uncertain_is_leak: if the judge reply is unparseable, treat it as a leak - (default False: keep the hint, the conservative choice). - judge_system: optional custom judge system prompt that overrides the built-in - answer_only/legacy criteria (for task-specific leak policies). - """ - - def __init__( - self, - sampler: Optional['Sampler'] = None, - *, - model_path: str = '', - sampling_params: Optional['SamplingParams'] = None, - judge_lora_path: Optional[str] = None, - max_content_chars: int = 4000, - max_query_chars: int = 4000, - uncertain_is_leak: bool = False, - answer_only: bool = True, - judge_system: Optional[str] = None, - ): - self.sampler = sampler - self.model_path = model_path - self.sampling_params = sampling_params - self.judge_lora_path = judge_lora_path or None - self.max_content_chars = int(max_content_chars) - self.max_query_chars = int(max_query_chars) - self.uncertain_is_leak = bool(uncertain_is_leak) - # answer_only: flag ONLY final-answer disclosure, permit method/plan (see - # _JUDGE_SYSTEM_ANSWER_ONLY). Default keeps the strict legacy behaviour. - self.answer_only = bool(answer_only) - # judge_system: caller-supplied criterion that overrides both built-ins, for - # task-specific leak policies (e.g. permit method/plan but still flag concrete - # intermediate key results). None -> fall back to the answer_only/legacy pair. - self.judge_system = judge_system or None - - # ------------------------------------------------------------------ - # public entry points - # ------------------------------------------------------------------ - def __call__(self, trajectory: dict, *, query: Optional[str] = None, - reference: Optional[Union[str, Sequence[str]]] = None, - **kwargs) -> int: - detail = self.leak_detail( - self._content_of(trajectory), - query=query or self._infer_query(trajectory), - reference=reference if reference is not None else _ground_truths(trajectory), - ) - return 0 if detail.leaked else self.NUM_LEVELS - 1 - - def is_leak(self, content: str, *, query: str, - reference: Optional[Union[str, Sequence[str]]] = None) -> bool: - return self.leak_detail(content, query=query, reference=reference).leaked - - def leak_detail(self, content: str, *, query: str, - reference: Optional[Union[str, Sequence[str]]] = None - ) -> LeakDetail: - """LLM judge; ``no_llm`` when neither a student nor a teacher exists.""" - content = content or '' - references = self._as_list(reference) - - if not self._llm_available(): - return LeakDetail(False, 'no_llm', 'none') - - verdict = self._judge(content, query or '', references) - if verdict is True: - return LeakDetail(True, 'llm_leak', 'llm') - if verdict is False: - return LeakDetail(False, '', 'llm') - return LeakDetail(self.uncertain_is_leak, 'llm_uncertain', 'llm') - - def leak_batch(self, items: Sequence[dict], *, max_workers: int = 8 - ) -> List[LeakDetail]: - """Check many hints in parallel; each item is ``{content, query, reference?}``. - - The judge (network / student inference) dominates latency, so the checks - fan out over a thread pool. Results are returned in input order. - """ - items = list(items) - if not items: - return [] - workers = max(1, min(max_workers, len(items))) - if workers == 1: - return [self._leak_detail_item(it) for it in items] - with ThreadPoolExecutor(max_workers=workers) as pool: - return list(pool.map(self._leak_detail_item, items)) - - def _leak_detail_item(self, item: dict) -> LeakDetail: - return self.leak_detail(item.get('content', ''), query=item.get('query', ''), - reference=item.get('reference')) - - # ------------------------------------------------------------------ - # judge (student, distilled via llm_backup) - # ------------------------------------------------------------------ - def _judge(self, content: str, query: str, - references: Sequence[str]) -> Optional[bool]: - trajectory = self._judge_trajectory(content, query, references) - raw = self._judge_once( - trajectory=trajectory, - sampling_params=self._judge_sampling_params(self.sampling_params), - judge_key='leak') - return _verdict_of(raw) - - # Distilled on the LEAK/CLEAN verdict. A single shared key ('leak') tracks - # student/teacher agreement on the judging skill as a whole; the comparator - # matches on the verdict, not byte-identical text. - @llm_backup(key_params=['judge_key'], comparator=_verdict_close) - def _judge_once(self, trajectory, sampling_params, judge_key: str = 'leak') -> str: - return self._sample_text(trajectory, sampling_params, self.judge_lora_path) - - def _judge_trajectory(self, content: str, query: str, - references: Sequence[str]) -> dict: - ref_block = '' - refs = [r for r in references if (r or '').strip()] - if refs: - ref_block = _REFERENCE_BLOCK.format(reference='; '.join(refs)) - user = _fill( - _JUDGE_USER, - query=self._trim(query, self.max_query_chars), - reference_block=ref_block, - hint=self._trim(content, self.max_content_chars)) - system = self.judge_system or ( - _JUDGE_SYSTEM_ANSWER_ONLY if self.answer_only else _JUDGE_SYSTEM) - return {'messages': [ - {'role': 'system', 'content': system}, - {'role': 'user', 'content': user}, - ]} - - def _judge_sampling_params(self, override): - if override is not None: - return override - from twinkle.data_format.sampling import SamplingParams - # A one-word verdict; keep the budget tiny. Small headroom absorbs models - # that prepend a stray token before LEAK/CLEAN. - return SamplingParams(temperature=0.0, max_tokens=8) - - # ------------------------------------------------------------------ - # LLM plumbing (mirrors RubricVerifier) - # ------------------------------------------------------------------ - def _llm_available(self) -> bool: - if self.sampler is not None: - return True - return bool(os.environ.get('LLM_BACKUP_API_KEY') - or os.environ.get('OPENAI_API_KEY') - or os.environ.get('LLM_BACKUP_BASE_URL')) - - def _sample_text(self, trajectory, sampling_params, lora_path) -> str: - if self.sampler is None: - return '' - sample_kwargs: dict[str, Any] = {'sampling_params': sampling_params} - if lora_path is None: - sample_kwargs['use_base_model'] = True - else: - sample_kwargs['adapter_path'] = lora_path - responses = self.sampler.sample([trajectory], **sample_kwargs) - resp = list(responses)[0] if responses else None - if resp is None: - return '' - seqs = getattr(resp, 'sequences', None) or [] - return (getattr(seqs[0], 'decoded', None) or '') if seqs else '' - - # ------------------------------------------------------------------ - # small helpers - # ------------------------------------------------------------------ - @staticmethod - def _as_list(reference: Optional[Union[str, Sequence[str]]]) -> List[str]: - if reference is None: - return [] - if isinstance(reference, str): - return [reference] - return [str(r) for r in reference] - - @staticmethod - def _trim(text: str, cap: int) -> str: - text = text or '' - return text if len(text) <= cap else text[:cap] - - @staticmethod - def _content_of(trajectory: dict) -> str: - """The hint under check = last assistant message text (else any content).""" - if isinstance(trajectory, str): - return trajectory - msgs = trajectory.get('messages', []) or [] - for m in reversed(msgs): - if m.get('role') == 'assistant': - c = m.get('content') - if isinstance(c, list): - c = '\n'.join(p.get('text', '') for p in c - if isinstance(p, dict) and p.get('type') == 'text') - if isinstance(c, str) and c.strip(): - return c - return str(trajectory.get('content', '') or '') - - @staticmethod - def _infer_query(trajectory: dict) -> str: - if isinstance(trajectory, str): - return '(no explicit task)' - for m in trajectory.get('messages', []) or []: - if m.get('role') == 'user': - c = m.get('content') - if isinstance(c, str) and c.strip(): - return c.strip() - return '(no explicit task)' - - -def _fill(template: str, **kw) -> str: - out = template - for k, v in kw.items(): - out = out.replace('{' + k + '}', str(v)) - return out diff --git a/src/twinkle_agentic/verifier/result_check.py b/src/twinkle_agentic/verifier/result_check.py new file mode 100644 index 000000000..1ea12e738 --- /dev/null +++ b/src/twinkle_agentic/verifier/result_check.py @@ -0,0 +1,365 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Program-checked outcomes for agentic episodes. + +An agentic episode ends with *state*, not with a string: files written, a +command that now succeeds, an answer stated in the final turn. This module +scores that end state with ordinary programs -- no judge model, so the same +trajectory always earns the same reward and difficulty filtering stays stable. + +A task declares a list of :class:`Check`; :func:`run_checks` evaluates them and +returns a :class:`CheckReport` whose ``score`` is the reward. + +Checks that need to *run* something (``shell`` / ``python``) go through a +``runner`` so they execute wherever the episode ran -- pass the sandbox's +runner and the check sees exactly the state the agent left behind. Without one +they fall back to a local subprocess in ``workspace``, which is only correct +when the episode itself ran locally. +""" +import json +import os +import re +import resource +import shutil +import signal +import subprocess +import sys +import tempfile +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple + +__all__ = [ + 'Check', + 'CheckOutcome', + 'CheckReport', + 'CheckContext', + 'run_checks', + 'checks_from_dicts', + 'local_runner', +] + +# (exit_code, output) for one command run inside the episode's workspace. +Runner = Callable[[str, str], Tuple[int, str]] + +DEFAULT_TIMEOUT = int(os.environ.get('RESULT_CHECK_TIMEOUT', 60)) +# Cap a runaway check so one bad task cannot take the trainer down with it. +_MEM_LIMIT_BYTES = 2 * 1024**3 + +_KINDS = ( + 'file_exists', + 'file_absent', + 'file_contains', + 'file_equals', + 'file_json', + 'shell', + 'python', + 'answer_contains', + 'answer_equals', + 'answer_regex', +) + + +@dataclass +class Check: + """One assertion about the end state. + + Args: + kind: one of :data:`_KINDS`. + path: workspace-relative file for the ``file_*`` kinds. + value: expected substring / exact text / JSON value, per kind. + pattern: regex alternative to ``value`` where the kind allows it. + code: shell command (``shell``) or python source (``python``). + key: dotted path into the document for ``file_json``, e.g. ``a.b.0.c``. + expect_exit: required exit status for ``shell`` / ``python``. + weight: contribution to the score; defaults to 1.0. + timeout: per-check seconds for the running kinds. + description: shown in the report so a failure is readable. + """ + kind: str + path: str = '' + value: Any = None + pattern: str = '' + code: str = '' + key: str = '' + expect_exit: int = 0 + weight: float = 1.0 + timeout: int = DEFAULT_TIMEOUT + description: str = '' + + def __post_init__(self): + if self.kind not in _KINDS: + raise ValueError(f'unknown check kind {self.kind!r}; expected one of {_KINDS}') + if self.weight <= 0: + raise ValueError(f'check weight must be positive, got {self.weight}') + + +@dataclass +class CheckOutcome: + check: Check + passed: bool + detail: str = '' + + +@dataclass +class CheckReport: + """Result of scoring one episode.""" + score: float + n_passed: int + n_total: int + outcomes: List[CheckOutcome] = field(default_factory=list) + + @property + def all_passed(self) -> bool: + return self.n_total > 0 and self.n_passed == self.n_total + + def failures(self) -> List[str]: + return [(o.check.description or o.check.kind) + ': ' + o.detail + for o in self.outcomes if not o.passed] + + def to_dict(self) -> Dict[str, Any]: + return { + 'score': self.score, + 'n_passed': self.n_passed, + 'n_total': self.n_total, + 'failures': self.failures(), + } + + +@dataclass +class CheckContext: + """What the checks are allowed to look at. + + Args: + workspace: directory the episode wrote into. + final_answer: text of the last assistant turn, for the ``answer_*`` kinds. + runner: executes a command in the episode's environment. ``None`` runs + it locally in ``workspace``. + """ + workspace: str = '' + final_answer: str = '' + runner: Optional[Runner] = None + + +def local_runner(workspace: str) -> Runner: + """Run commands in ``workspace`` as a local subprocess. + + Uses ``start_new_session`` + ``killpg`` so a forking command cannot leave + grandchildren behind on timeout, and caps address space at 2GB. + """ + + def _run(command: str, interpreter: str) -> Tuple[int, str]: + return _local_exec(command, interpreter, workspace, DEFAULT_TIMEOUT) + + return _run + + +def _local_exec(source: str, interpreter: str, cwd: str, timeout: int) -> Tuple[int, str]: + cwd = cwd or '.' + os.makedirs(cwd, exist_ok=True) + if interpreter == 'python': + tmp = tempfile.mkdtemp(prefix='rescheck_') + script = os.path.join(tmp, '_check.py') + with open(script, 'w', encoding='utf-8') as f: + f.write(source) + argv = [sys.executable, script] + else: + tmp = None + argv = ['/bin/bash', '-lc', source] + + env = dict(os.environ, MPLBACKEND='Agg', PYTHONHASHSEED='0', + OMP_NUM_THREADS='1', MKL_NUM_THREADS='1', + TOKENIZERS_PARALLELISM='false') + env.pop('CUDA_VISIBLE_DEVICES', None) + + def _limit(): + resource.setrlimit(resource.RLIMIT_AS, (_MEM_LIMIT_BYTES, _MEM_LIMIT_BYTES)) + + try: + proc = subprocess.Popen(argv, cwd=cwd, env=env, stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, text=True, errors='replace', + start_new_session=True, preexec_fn=_limit) + try: + out, _ = proc.communicate(timeout=timeout) + return proc.returncode, out or '' + except subprocess.TimeoutExpired: + try: + os.killpg(proc.pid, signal.SIGKILL) + except ProcessLookupError: + pass + try: + proc.communicate(timeout=5) + except Exception: # noqa + pass + return 124, f'check did not finish within {timeout}s' + except Exception as e: # noqa + return 1, f'{type(e).__name__}: {e}' + finally: + if tmp: + shutil.rmtree(tmp, ignore_errors=True) + + +def checks_from_dicts(raw: Sequence[Dict[str, Any]]) -> List[Check]: + """Build checks from the plain dicts a task file carries.""" + return [Check(**dict(item)) for item in raw or []] + + +def _resolve(workspace: str, path: str) -> str: + """Resolve a task-declared path inside the workspace. + + Rejects escapes: a task must not be able to assert on files outside the + episode's own directory, or one episode could pass by reading another's. + """ + root = os.path.realpath(workspace or '.') + target = os.path.realpath(os.path.join(root, path)) + if target != root and not target.startswith(root + os.sep): + raise ValueError(f'check path {path!r} escapes the workspace') + return target + + +def _read_text(workspace: str, path: str) -> Tuple[Optional[str], str]: + try: + full = _resolve(workspace, path) + except ValueError as e: + return None, str(e) + if not os.path.isfile(full): + return None, f'{path} does not exist' + try: + with open(full, encoding='utf-8', errors='replace') as f: + return f.read(), '' + except OSError as e: + return None, f'cannot read {path}: {e}' + + +def _dig(doc: Any, key: str) -> Tuple[bool, Any]: + """Walk a dotted path; integer segments index into lists.""" + cur = doc + for seg in [s for s in key.split('.') if s]: + if isinstance(cur, dict): + if seg not in cur: + return False, None + cur = cur[seg] + elif isinstance(cur, list): + if not seg.lstrip('-').isdigit(): + return False, None + idx = int(seg) + if not -len(cur) <= idx < len(cur): + return False, None + cur = cur[idx] + else: + return False, None + return True, cur + + +def _norm(text: Any) -> str: + return str(text if text is not None else '').strip() + + +def _eval_one(check: Check, ctx: CheckContext) -> CheckOutcome: + kind = check.kind + + if kind in ('file_exists', 'file_absent'): + try: + full = _resolve(ctx.workspace, check.path) + except ValueError as e: + return CheckOutcome(check, False, str(e)) + there = os.path.exists(full) + want = (kind == 'file_exists') + return CheckOutcome(check, there == want, + '' if there == want else + (f'{check.path} does not exist' if want + else f'{check.path} should not exist')) + + if kind in ('file_contains', 'file_equals', 'file_json'): + text, err = _read_text(ctx.workspace, check.path) + if text is None: + return CheckOutcome(check, False, err) + if kind == 'file_contains': + if check.pattern: + ok = re.search(check.pattern, text, re.S) is not None + return CheckOutcome(check, ok, '' if ok else + f'{check.path} does not match /{check.pattern}/') + ok = _norm(check.value) in text + return CheckOutcome(check, ok, '' if ok else + f'{check.path} does not contain {_norm(check.value)!r}') + if kind == 'file_equals': + ok = text.strip() == _norm(check.value) + return CheckOutcome(check, ok, '' if ok else + f'{check.path} is {text.strip()[:120]!r}, ' + f'expected {_norm(check.value)[:120]!r}') + try: + doc = json.loads(text) + except json.JSONDecodeError as e: + return CheckOutcome(check, False, f'{check.path} is not valid JSON: {e}') + found, got = _dig(doc, check.key) + if not found: + return CheckOutcome(check, False, f'{check.path} has no key {check.key!r}') + ok = got == check.value if not isinstance(check.value, str) else _norm(got) == _norm(check.value) + return CheckOutcome(check, ok, '' if ok else + f'{check.path}:{check.key} is {got!r}, expected {check.value!r}') + + if kind in ('shell', 'python'): + runner = ctx.runner or local_runner(ctx.workspace) + try: + code, out = runner(check.code, 'python' if kind == 'python' else 'shell') + except Exception as e: # noqa + return CheckOutcome(check, False, f'runner raised {type(e).__name__}: {e}') + if code != check.expect_exit: + return CheckOutcome(check, False, + f'exit {code} (expected {check.expect_exit}); output: {out[-300:]}') + if check.pattern and re.search(check.pattern, out or '', re.S) is None: + return CheckOutcome(check, False, f'output does not match /{check.pattern}/') + if check.value is not None and _norm(check.value) not in (out or ''): + return CheckOutcome(check, False, f'output does not contain {_norm(check.value)!r}') + return CheckOutcome(check, True) + + answer = ctx.final_answer or '' + if kind == 'answer_contains': + ok = _norm(check.value) in answer + return CheckOutcome(check, ok, '' if ok else + f'final answer does not contain {_norm(check.value)!r}') + if kind == 'answer_equals': + ok = answer.strip() == _norm(check.value) + return CheckOutcome(check, ok, '' if ok else + f'final answer is {answer.strip()[:120]!r}, ' + f'expected {_norm(check.value)[:120]!r}') + ok = re.search(check.pattern, answer, re.S) is not None + return CheckOutcome(check, ok, '' if ok else + f'final answer does not match /{check.pattern}/') + + +def run_checks( + checks: Sequence[Check], + ctx: CheckContext, + mode: str = 'fraction', +) -> CheckReport: + """Score one episode against its checks. + + Args: + checks: the task's assertions. An empty list scores 0.0 rather than a + free 1.0, so a task that forgot to declare checks cannot look solved. + ctx: workspace / final answer / runner. + mode: ``fraction`` gives weighted partial credit, ``all_or_nothing`` + gives 1.0 only when every check passes. + + A check that raises is a failed check, never a failed batch: one malformed + task must not abort scoring for the rest of the rollout group. + """ + if mode not in ('fraction', 'all_or_nothing'): + raise ValueError(f"mode must be 'fraction' or 'all_or_nothing', got {mode!r}") + checks = list(checks or []) + if not checks: + return CheckReport(score=0.0, n_passed=0, n_total=0, outcomes=[]) + + outcomes: List[CheckOutcome] = [] + for check in checks: + try: + outcomes.append(_eval_one(check, ctx)) + except Exception as e: # noqa + outcomes.append(CheckOutcome(check, False, f'{type(e).__name__}: {e}')) + + n_passed = sum(1 for o in outcomes if o.passed) + if mode == 'all_or_nothing': + score = 1.0 if n_passed == len(outcomes) else 0.0 + else: + total_w = sum(o.check.weight for o in outcomes) + score = sum(o.check.weight for o in outcomes if o.passed) / total_w + return CheckReport(score=score, n_passed=n_passed, n_total=len(outcomes), outcomes=outcomes) diff --git a/src/twinkle_agentic/verifier/rubric_library.py b/src/twinkle_agentic/verifier/rubric_library.py deleted file mode 100644 index 8f5f753c0..000000000 --- a/src/twinkle_agentic/verifier/rubric_library.py +++ /dev/null @@ -1,108 +0,0 @@ -"""Intent-keyed rubric library (DESIGN follow-up: stabilize rubric scoring). - -Rubric *generation* is flexible but high-variance: for template-like intents -(tool_call / code / math) the model re-invents slightly different criteria every -call, which is the main source of score jitter and occasional task-type -misreads. This module supplies two levels of stabilization, both keyed by the -intent vocabulary in :mod:`twinkle_agentic.preprocessor.intents`: - -- ``INTENT_BASE_RUBRICS`` — half-fixed **skeletons**: a small, stable core of - criteria that is PREPENDED to the distilled rubric. The generator still adds - task-specific criteria on top, so flexibility is preserved while the shared - core makes scores comparable across similar segments. This is the DEFAULT - policy (does NOT sacrifice flexibility). -- ``INTENT_FIXED_RUBRICS`` — fully-fixed rubrics per intent (no generation). - Maximum stability, minimum flexibility; opt-in for callers that want it. - -Each criterion is written to match the grader prompt conventions: -- starts with "The response" / "The agent", -- [Hard Rule] for objectively checkable constraints, [Principle] for quality, -- scoped to what is observable INSIDE one segment (never assumes later steps). - -Criteria are deliberately GENERIC (no entities/values) so a single skeleton -generalizes across all segments of that intent. -""" -from typing import Dict, List - -from .rubric_verifier import RubricItem - -# Re-export intent constants so callers wire the library without importing the -# heavier classifier module. -from ..preprocessor.intents import (INTENT_CODE, INTENT_MATH, # noqa: F401 - INTENT_TOOL_CALL) - - -def _h(text: str) -> RubricItem: - return RubricItem(text=text, is_hard=True) - - -def _p(text: str) -> RubricItem: - return RubricItem(text=text, is_hard=False) - - -# --------------------------------------------------------------------------- # -# Half-fixed skeletons (DEFAULT). Kept intentionally short (2-3 items) so the -# distilled generator still supplies the bulk of task-specific coverage. -# --------------------------------------------------------------------------- # -_TOOL_CALL_SKELETON: List[RubricItem] = [ - _h('The agent emits tool calls whose arguments are valid, complete JSON ' - 'matching the tool schema'), - _h('The agent selects tools appropriate to the sub-goal and does not invent ' - 'unavailable tools or arguments'), - _p('The agent uses each tool result to advance the sub-goal without ' - 'redundant or repeated identical calls'), -] - -_CODE_SKELETON: List[RubricItem] = [ - _h('The response produces code that is syntactically well-formed and ' - 'self-consistent within the segment'), - _p('The response addresses the stated coding sub-goal with correct, relevant ' - 'logic rather than placeholder or off-topic code'), - _p('The response avoids obvious defects (undefined names, wrong API usage) ' - 'visible within the segment'), -] - -_MATH_SKELETON: List[RubricItem] = [ - _h('The response performs each mathematical step correctly with no ' - 'arithmetic or algebraic error visible in the segment'), - _p('The response follows a valid, coherent solution path toward the ' - 'sub-goal without unjustified leaps'), - _p('The response states intermediate/final results clearly and consistently ' - 'with the work shown'), -] - -INTENT_BASE_RUBRICS: Dict[str, List[RubricItem]] = { - INTENT_TOOL_CALL: _TOOL_CALL_SKELETON, - INTENT_CODE: _CODE_SKELETON, - INTENT_MATH: _MATH_SKELETON, -} - - -# --------------------------------------------------------------------------- # -# Fully-fixed rubrics (opt-in). Same criteria plus a couple more so the fixed -# set is self-sufficient without any generation. -# --------------------------------------------------------------------------- # -INTENT_FIXED_RUBRICS: Dict[str, List[RubricItem]] = { - INTENT_TOOL_CALL: _TOOL_CALL_SKELETON + [ - _p('The agent grounds its next action in the actual tool output rather ' - 'than hallucinating results'), - ], - INTENT_CODE: _CODE_SKELETON + [ - _p('The response explains or structures the code enough to be usable in ' - 'the surrounding task context'), - ], - INTENT_MATH: _MATH_SKELETON + [ - _p('The response keeps units, signs and notation consistent throughout ' - 'the segment'), - ], -} - - -def default_intent_base_rubrics() -> Dict[str, List[RubricItem]]: - """The recommended half-fixed policy (flexible + stabilized).""" - return {k: list(v) for k, v in INTENT_BASE_RUBRICS.items()} - - -def default_intent_fixed_rubrics() -> Dict[str, List[RubricItem]]: - """The opt-in fully-fixed policy (max stability, min flexibility).""" - return {k: list(v) for k, v in INTENT_FIXED_RUBRICS.items()} diff --git a/src/twinkle_agentic/verifier/rubric_verifier.py b/src/twinkle_agentic/verifier/rubric_verifier.py deleted file mode 100644 index 8e10b19da..000000000 --- a/src/twinkle_agentic/verifier/rubric_verifier.py +++ /dev/null @@ -1,1077 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Rubric-based verifier for a single (pre-segmented) trajectory segment. - -Design follows the OpenRubrics -> RubricARROW line of work, adapted to this -repo's progressive-distillation setup: - -1. **Two LLM stages, both distilled via ``llm_backup``** - - *Rubric generation*: given the segment, produce a small set of scoring - criteria, each tagged ``[Hard Rule]`` or ``[Principle]``. - - *Rubric scoring*: given the segment + rubric, emit a per-criterion - verdict. Scores are aggregated ARROW-style into one pointwise scalar. - -2. **Code-level hard verification does NOT go through the LLM.** - Tool-call success/failure, argument JSON validity and call formatting are - checked deterministically (free + un-hackable) and blended in as a - "gatekeeper" floor on the final score. - -3. **Cost-aware scoring**: a single scoring pass yields a soft margin. Only - when the judge is uncertain (|margin| small) do we escalate to majority - voting, aggregating with a median / trimmed-mean (robust to outliers). - -The public ``__call__`` returns an ``int`` in ``[0, NUM_LEVELS)`` per the -:class:`Verifier` contract. ``score_detail`` exposes the continuous score and -breakdown for callers that want the raw signal (e.g. an RL reward). -""" -from __future__ import annotations - -import json -import os -import re -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple - -from twinkle_agentic.utils.llm_backup import llm_backup - -from .base import Verifier - -if TYPE_CHECKING: - from twinkle.data_format import SamplingParams # noqa: F401 - from twinkle.sampler.base import Sampler # noqa: F401 - - -# --------------------------------------------------------------------------- -# Prompts -# --------------------------------------------------------------------------- -_GEN_SYSTEM = """\ -You write evaluation rubrics for a single segment of an AI agent trajectory. \ -The segment may contain reasoning, tool calls and tool results. - -Produce a SHORT list of scoring criteria that discriminate a good segment from \ -a bad one. Each criterion: -- starts with "The response" or "The agent", -- is checkable and non-overlapping (no two criteria testing the same thing), -- ends with a tag: [Hard Rule] for objectively verifiable constraints \ -(tool actually called, argument schema valid, required output present) or \ -[Principle] for softer quality (reasoning soundness, sub-goal progress, no \ -redundant calls). - -Scope discipline (critical — avoid over-strict, mismatched rubrics): -- This is ONE SEGMENT, possibly the MIDDLE of a longer task. Only write criteria \ -about behavior that is OBSERVABLE INSIDE THIS SEGMENT. Do NOT invent criteria \ -about a final deliverable, later steps, or task completion that this segment is \ -not expected to reach (e.g. "registers the component", "updates the entry point"). -- Infer the task type ONLY from what the segment actually does. Do NOT assume it \ -is an "implement a feature" task unless the segment clearly shows that. When the \ -segment only reads/inspects/answers, judge reading/answering quality, not delivery. -- Reasoning shown inside <think>...</think> (or <thinking>) is internal scratch \ -work. Never write a criterion that penalizes the mere presence of such reasoning, \ -and do NOT let it count against "output only X" style constraints. - -Rules: -- Output {min_n}-{max_n} criteria, as FEW as needed to cover the key axes. -- Do NOT reference specific entities/values from THIS segment; keep criteria \ -generalizable to similar segments. -- Output ONLY a numbered list, one criterion per line, nothing else. -""" - -_GEN_USER = """\ -## Task / query (context) -{query} - -## Segment to build a rubric for -{segment} - -Now output the numbered rubric list.""" - -_SCORE_SYSTEM = """\ -You are a strict rubric grader for one segment of an agent trajectory. - -You are given a rubric (numbered criteria, each tagged [Hard Rule] or \ -[Principle]) and the segment. For EACH criterion output one line: - - <index>: PASS or <index>: FAIL - -Judge every criterion independently and literally. A [Hard Rule] fails unless \ -it is unambiguously satisfied. - -Grading discipline: -- Judge ONLY what is observable in THIS segment; if a criterion asks about a \ -step/deliverable this segment was not meant to reach, do not FAIL it for that \ -alone — grade it satisfied when the in-segment behavior is correct. -- Content inside <think>...</think> (or <thinking>) is internal reasoning, not \ -user-facing output. For "output only X / no extra text" style criteria, ignore \ -such reasoning blocks; judge the actual response payload. - -Output only the verdict lines, in order, then stop. Do not add explanations.""" - -_SCORE_USER = """\ -## Task / query (context) -{query} - -## Rubric -{rubric} - -## Segment -{segment} - -Now output one PASS/FAIL line per criterion, in order.""" - - -# --- diagnostic mode: single call yields verdict + reason together --------- -# Used to distil an "on-the-fly error checker" LoRA (DESIGN §11.6). Unlike the -# terse scorer above, this asks for a COMPLETE verification chain over EVERY -# criterion (both pass and fail) so the distilled LoRA learns to also emit -# "checked, all good, continue" — not only to nitpick. Verdict and reason are -# produced in ONE pass so they can never disagree. -_DIAG_SYSTEM = """\ -You are a process error checker for one segment of an agent trajectory. You are \ -given a rubric (numbered criteria, each tagged [Hard Rule] or [Principle]) and \ -the segment. Walk through EVERY criterion in order and, for each, decide PASS or \ -FAIL and briefly justify it grounded in the segment. - -Output STRICT JSON (no prose outside it) with this shape: -{ - "items": [ - {"index": 1, "verdict": "PASS", "reason": "<why, grounded in the segment>", - "fix": ""}, - {"index": 2, "verdict": "FAIL", "reason": "<what is wrong and where>", - "fix": "<one concrete, actionable correction>"} - ], - "overall": "OK" | "ISSUES", - "summary": "<one sentence: 'no process errors, continue' OR the key issue(s)>" -} - -Rules: -- Judge every criterion independently and literally; a [Hard Rule] is FAIL \ -unless unambiguously satisfied. -- Judge ONLY what is observable in THIS segment; do not FAIL a criterion merely \ -because a later step/deliverable it references is outside this segment's scope. -- Content inside <think>...</think> (or <thinking>) is internal reasoning, not \ -user-facing output; ignore it for "output only X" style criteria. -- For PASS items, leave "fix" as "". For FAIL items, "fix" must be a concrete \ -correction (e.g. add the missing argument, redo step k). -- Keep every "reason" and "fix" clear and concise — one short sentence each, \ -stating only the essential point; do not restate the criterion, quote the segment \ -at length, or add filler. -- "overall" is "OK" only if NO criterion is FAIL. -- Output only the JSON object.""" - -_DIAG_USER = """\ -## Task / query (context) -{query} - -## Rubric -{rubric} - -## Segment -{segment} - -Now output the diagnostic JSON object.""" - - -# --------------------------------------------------------------------------- -# Data holders -# --------------------------------------------------------------------------- -_VERDICT_RE = re.compile(r'^\s*(\d+)\s*[:.)]\s*(pass|fail|true|false|yes|no|1|0)\b', - re.IGNORECASE) - - -@dataclass -class RubricItem: - text: str - is_hard: bool - - -@dataclass -class ScoreDetail: - """Full breakdown behind the final integer level.""" - level: int - scalar: float # continuous pointwise score in [0, 1] - llm_scalar: float # LLM (principle+softhard) component in [0, 1] - hard_pass_rate: float # code-verified hard-rule pass rate in [0, 1] - gated: bool # True if code gatekeeper capped the score - n_votes: int # scoring passes actually spent - rubric: List[RubricItem] = field(default_factory=list) - per_item_pass_rate: List[float] = field(default_factory=list) - - -@dataclass -class DiagnosisItem: - """Per-criterion diagnostic verdict with its justification.""" - index: int - verdict: bool # True == PASS - reason: str = '' - fix: str = '' # concrete correction, only for FAIL - - -@dataclass -class DiagnoseDetail: - """A complete verification chain over one segment (DESIGN §11.6). - - Produced in a single LLM call so verdict and reason are always consistent. - Covers EVERY criterion (pass and fail) so a distilled checker learns to emit - "checked, no error, continue" as well as concrete fault localisation. - """ - scalar: float # aggregated pointwise score in [0, 1] - overall_ok: bool # True == no criterion failed - summary: str # one-line human-readable conclusion - items: List[DiagnosisItem] = field(default_factory=list) - rubric: List[RubricItem] = field(default_factory=list) - raw: str = '' # raw model output (for SFT targets) - query: str = '' # task/query context (for SFT inputs) - segment_text: str = '' # rendered segment (for SFT inputs) - - -# --------------------------------------------------------------------------- -# Verifier -# --------------------------------------------------------------------------- -class RubricVerifier(Verifier): - """Score one trajectory segment on a 0..NUM_LEVELS-1 scale via auto rubrics. - - Args: - sampler: Student model sampler (local inference). If ``None`` the - verifier still works but every LLM call is served by the teacher - API through ``llm_backup`` (useful before a student exists). - model_path: Model identifier (bookkeeping only). - sampling_params: Default sampling params for LLM calls. - gen_lora_path: LoRA adapter for the rubric-generator student. - score_lora_path: LoRA adapter for the rubric-scorer student. - min_rubrics / max_rubrics: Target rubric-count window per segment. - hard_weight / principle_weight: Aggregation weights (ARROW uses 3 / 1). - margin_threshold: |margin| below which we escalate to voting. - max_votes: Voting cap for uncertain segments (odd recommended). - gate_floor_ratio: If code-verified hard rules fail, the final scalar is - capped at ``hard_pass_rate`` (gatekeeper). Set to 1.0 to hard-cap, - 0.0 to disable gating. - """ - - def __init__( - self, - sampler: Optional['Sampler'] = None, - *, - model_path: str = '', - sampling_params: Optional['SamplingParams'] = None, - gen_lora_path: Optional[str] = None, - score_lora_path: Optional[str] = None, - min_rubrics: int = 5, - max_rubrics: int = 8, - hard_weight: float = 3.0, - principle_weight: float = 1.0, - margin_threshold: float = 0.25, - max_votes: int = 5, - gate: bool = True, - fixed_rubric: Optional[List['RubricItem']] = None, - base_rubric: Optional[List['RubricItem']] = None, - intent_rubrics: Optional[Dict[str, List['RubricItem']]] = None, - intent_base_rubrics: Optional[Dict[str, List['RubricItem']]] = None, - max_segment_chars: int = 14_000, - long_segment_chars: int = 8_000, - min_votes_long: int = 3, - long_margin_threshold: float = 0.18, - max_votes_long: int = 3, - min_votes_high: int = 3, - high_score_threshold: float = 0.85, - diag_max_tokens: int = 2048, - ): - if max_rubrics < min_rubrics: - raise ValueError('max_rubrics must be >= min_rubrics') - if min_rubrics < 1: - raise ValueError('min_rubrics must be >= 1') - if hard_weight <= 0 or principle_weight <= 0: - raise ValueError('weights must be > 0') - if not 0.0 <= margin_threshold <= 1.0: - raise ValueError('margin_threshold must be in [0, 1]') - if max_votes < 1: - raise ValueError('max_votes must be >= 1') - - self.sampler = sampler - self.model_path = model_path - self.sampling_params = sampling_params - self.gen_lora_path = gen_lora_path or None - self.score_lora_path = score_lora_path or None - self.min_rubrics = int(min_rubrics) - self.max_rubrics = int(max_rubrics) - self.hard_weight = float(hard_weight) - self.principle_weight = float(principle_weight) - self.margin_threshold = float(margin_threshold) - self.max_votes = int(max_votes) - self.gate = bool(gate) - self.max_segment_chars = int(max_segment_chars) - self.long_segment_chars = int(long_segment_chars) - self.min_votes_long = max(1, int(min_votes_long)) - self.long_margin_threshold = float(long_margin_threshold) - self.max_votes_long = max(1, int(max_votes_long)) - # High-confidence band: force at least this many votes when the first - # pass lands >= high_score_threshold, so 4/4-looking "level 4" segments - # are not decided by a single lucky sample (reduces high-band variance). - self.min_votes_high = max(1, int(min_votes_high)) - self.high_score_threshold = float(high_score_threshold) - # Diagnosis emits a full per-criterion (verdict+reason+fix) JSON; it needs - # a far larger token budget than terse scoring or it truncates mid-JSON. - self.diag_max_tokens = max(256, int(diag_max_tokens)) - # When provided, skip stage-1 rubric generation and score against these - # fixed criteria (e.g. a safety rubric — AUDIT D8). - self.fixed_rubric: Optional[List['RubricItem']] = list(fixed_rubric) if fixed_rubric else None - # Skeleton criteria PREPENDED to every distilled rubric (half-fixed mode, - # DESIGN follow-up): stabilizes cross-segment comparability while still - # letting stage-1 add task-specific criteria. Ignored when fixed_rubric set. - self.base_rubric: Optional[List['RubricItem']] = list(base_rubric) if base_rubric else None - # Intent-aware routing: per-intent fully-fixed rubrics (highest priority) - # and per-intent half-fixed skeletons. Keys are intent strings (intents.py). - self.intent_rubrics: Optional[Dict[str, List['RubricItem']]] = ( - {k: list(v) for k, v in intent_rubrics.items()} if intent_rubrics else None) - self.intent_base_rubrics: Optional[Dict[str, List['RubricItem']]] = ( - {k: list(v) for k, v in intent_base_rubrics.items()} if intent_base_rubrics else None) - - # ------------------------------------------------------------------ - # public entry points - # ------------------------------------------------------------------ - def __call__(self, trajectory: dict, **kwargs) -> int: - return self.score_detail(trajectory, **kwargs).level - - def score_detail(self, trajectory: dict, *, query: Optional[str] = None, - sampling_params: Any = None, - extra_context: Optional[str] = None, - intent: Optional[str] = None) -> ScoreDetail: - query = query or self._infer_query(trajectory) - segment_text = self._trim_segment_for_llm(self._render_segment(trajectory)) - # D7c: fold an objective finding into the scored transcript so the judge - # re-scores WITH the hard evidence in view (objective corrects subjective). - if extra_context: - segment_text = f'{segment_text}\n\n[OBJECTIVE EVIDENCE]\n{extra_context}' - - # --- code-level hard verification (free, un-hackable) --- - hard_pass_rate, has_hard = self._code_hard_checks(trajectory) - - # No LLM (no student sampler AND no teacher API): skip both LLM stages and - # fall back to the deterministic code signal, instead of letting the - # llm_backup teacher path raise a missing-credentials error. - if not self._llm_available(): - scalar = hard_pass_rate if has_hard else 0.0 - return ScoreDetail( - level=self._to_level(scalar), scalar=scalar, llm_scalar=0.0, - hard_pass_rate=hard_pass_rate if has_hard else 1.0, - gated=False, n_votes=0, rubric=[], - ) - - # --- stage 1: rubric (fixed if configured, else distilled generation) --- - rubric = self._build_rubric(query, segment_text, sampling_params, intent=intent) - if not rubric: - # No usable rubric: fall back to the code signal alone. - scalar = hard_pass_rate if has_hard else 0.0 - return ScoreDetail( - level=self._to_level(scalar), scalar=scalar, llm_scalar=0.0, - hard_pass_rate=hard_pass_rate if has_hard else 1.0, - gated=False, n_votes=0, rubric=[], - ) - - # --- stage 2: rubric scoring with margin-adaptive voting --- - per_item_rate, n_votes = self._score_with_voting( - query, segment_text, rubric, sampling_params) - - llm_scalar = self._aggregate(rubric, per_item_rate) - - # --- gatekeeper: code-verified hard failures cap the score --- - scalar = llm_scalar - gated = False - if self.gate and has_hard and hard_pass_rate < 1.0: - capped = min(llm_scalar, hard_pass_rate) - gated = capped < llm_scalar - scalar = capped - - return ScoreDetail( - level=self._to_level(scalar), - scalar=scalar, - llm_scalar=llm_scalar, - hard_pass_rate=hard_pass_rate if has_hard else 1.0, - gated=gated, - n_votes=n_votes, - rubric=rubric, - per_item_pass_rate=per_item_rate, - ) - - def diagnose(self, trajectory: dict, *, query: Optional[str] = None, - sampling_params: Any = None, - intent: Optional[str] = None) -> DiagnoseDetail: - """Produce a COMPLETE verification chain over the segment (DESIGN §11.6). - - Unlike :meth:`score_detail` (terse PASS/FAIL, tuned to be cheap), this - emits, in a SINGLE llm_backup-distilled call, a per-criterion verdict - *with* its reason and (on FAIL) a concrete fix, plus an overall verdict. - The single call keeps verdict and reason mutually consistent, and it - covers passing criteria too so a distilled checker learns to say - "checked, no error, continue" — not only to nitpick. - - Every call flows through ``llm_backup``: the (student, teacher, match) - pairs it records are exactly the SFT corpus for the error-checker LoRA. - Store all of them (both OK and ISSUES segments); balancing is a - training-time sampling concern, not a collection-time one. - """ - query = query or self._infer_query(trajectory) - segment_text = self._trim_segment_for_llm(self._render_segment(trajectory)) - - if not self._llm_available(): - # No LLM: fall back to the deterministic code signal only. - hard_pass_rate, has_hard = self._code_hard_checks(trajectory) - scalar = hard_pass_rate if has_hard else 1.0 - return DiagnoseDetail( - scalar=scalar, overall_ok=scalar >= 1.0, - summary='no LLM available; code-hard signal only', - items=[], rubric=[], raw='', query=query, segment_text=segment_text) - - # Reuse the same rubric machinery as scoring (fixed / half-fixed / gen). - rubric = self._build_rubric(query, segment_text, sampling_params, intent=intent) - if not rubric: - hard_pass_rate, has_hard = self._code_hard_checks(trajectory) - scalar = hard_pass_rate if has_hard else 1.0 - return DiagnoseDetail( - scalar=scalar, overall_ok=scalar >= 1.0, - summary='no usable rubric; code-hard signal only', - items=[], rubric=rubric, raw='', query=query, segment_text=segment_text) - - rubric_block = self._render_rubric(rubric) - rubric_key = _short_hash(rubric_block) - raw = self._diagnose_once( - trajectory=self._diagnose_trajectory(query, rubric_block, segment_text), - sampling_params=self._diagnose_sampling_params(sampling_params, temperature=0.0), - query=query, rubric_key=rubric_key) - - items, overall_ok, summary = self._parse_diagnosis(raw, len(rubric)) - # Blend deterministic hard checks in as a gatekeeper floor, mirroring - # score_detail so the diagnostic scalar is comparable to the scoring one. - per_item_rate = [1.0 if it.verdict else 0.0 for it in items] - llm_scalar = self._aggregate(rubric, per_item_rate) if per_item_rate else 0.0 - hard_pass_rate, has_hard = self._code_hard_checks(trajectory) - scalar = llm_scalar - if self.gate and has_hard and hard_pass_rate < 1.0: - scalar = min(llm_scalar, hard_pass_rate) - return DiagnoseDetail( - scalar=scalar, overall_ok=overall_ok, summary=summary, - items=items, rubric=rubric, raw=raw, - query=query, segment_text=segment_text) - - # ------------------------------------------------------------------ - # stage 1: rubric assembly (fixed | half-fixed skeleton + distilled | gen) - # ------------------------------------------------------------------ - def _build_rubric(self, query, segment_text, sampling_params, - intent: Optional[str] = None) -> List[RubricItem]: - """Return the rubric to score against. - - Selection order (intent-aware routing, DESIGN follow-up): - 1. ``intent_rubrics[intent]`` set -> fully fixed for this intent (most - stable; template-like tasks such as tool_call / code / math). - 2. ``fixed_rubric`` set -> global fixed rubric, verbatim. - 3. else -> distilled generation, optionally - PREPENDED with a fixed skeleton: ``intent_base_rubrics[intent]`` if - present, else the global ``base_rubric`` (half-fixed). Skeleton gives - cross-segment comparability; the generated tail adds task-specific - coverage. Duplicate criteria (same normalized text) drop, skeleton wins. - """ - if intent and self.intent_rubrics and intent in self.intent_rubrics: - return list(self.intent_rubrics[intent]) - if self.fixed_rubric is not None: - return list(self.fixed_rubric) - - skeleton: Optional[List[RubricItem]] = None - if intent and self.intent_base_rubrics and intent in self.intent_base_rubrics: - skeleton = self.intent_base_rubrics[intent] - elif self.base_rubric: - skeleton = self.base_rubric - - gen_min = self.min_rubrics - gen_max = self.max_rubrics - if skeleton: - # leave room for the skeleton so the total stays in the count window - gen_min = max(1, self.min_rubrics - len(skeleton)) - gen_max = max(gen_min, self.max_rubrics - len(skeleton)) - raw_rubric = self._gen_rubric( - trajectory=self._gen_trajectory(query, segment_text, gen_min, gen_max), - sampling_params=self._gen_sampling_params(sampling_params), - query=query) - generated = self._parse_rubric(raw_rubric) - if not skeleton: - return generated - merged = list(skeleton) - seen = {_norm_criterion(it.text) for it in merged} - for it in generated: - key = _norm_criterion(it.text) - if key and key not in seen: - seen.add(key) - merged.append(it) - return merged - - @llm_backup(key_params=['query'], comparator=lambda a, b: _rubric_similar(a, b)) - def _gen_rubric(self, trajectory, sampling_params, query: str = None) -> str: - return self._sample_text(trajectory, sampling_params, self.gen_lora_path) - - # ------------------------------------------------------------------ - # stage 2: rubric scoring (student, distilled via llm_backup) - # ------------------------------------------------------------------ - # Note: the scoring pass is distilled on the *verdict pattern*; the - # comparator matches on binned pass-rate so student/teacher agree when - # their PASS/FAIL vectors are close (not byte-identical). - @llm_backup(key_params=['query', 'rubric_key'], - comparator=lambda a, b: _verdicts_close(a, b)) - def _score_once(self, trajectory, sampling_params, query: str = None, - rubric_key: str = '') -> str: - return self._sample_text(trajectory, sampling_params, self.score_lora_path) - - # ------------------------------------------------------------------ - # diagnostic pass (student, distilled via llm_backup) — DESIGN §11.6 - # ------------------------------------------------------------------ - # Distilled on the full (verdict + reason) chain. Consistency is checked on - # the per-criterion verdict vector (same idea as scoring), not on the free - # text of the reasons — two valid reasons for the same verdict should match. - @llm_backup(key_params=['query', 'rubric_key'], - comparator=lambda a, b: _diag_verdicts_close(a, b)) - def _diagnose_once(self, trajectory, sampling_params, query: str = None, - rubric_key: str = '') -> str: - return self._sample_text(trajectory, sampling_params, self.score_lora_path) - - def _score_with_voting(self, query, segment_text, rubric, sampling_params - ) -> Tuple[List[float], int]: - n = len(rubric) - rubric_block = self._render_rubric(rubric) - rubric_key = _short_hash(rubric_block) - score_traj = self._score_trajectory(query, rubric_block, segment_text) - - margin_thr = self.margin_threshold - vote_cap = self.max_votes - min_votes = 1 - if len(segment_text) >= self.long_segment_chars: - margin_thr = self.long_margin_threshold - min_votes = self.min_votes_long - vote_cap = min(vote_cap, self.max_votes_long) - - # First (cheap) pass. - votes: List[List[bool]] = [] - first = self._score_once( - trajectory=score_traj, - sampling_params=self._score_sampling_params(sampling_params, temperature=0.0), - query=query, rubric_key=rubric_key) - votes.append(self._parse_verdicts(first, n)) - - # High-confidence band: a lone pass that looks like "all good" (>= the - # high threshold) still gets re-sampled, so top-band scores are not - # decided by one lucky draw. Raise the required vote depth accordingly. - rate = self._vote_rates(votes) - first_scalar = self._aggregate(rubric, rate) - if first_scalar >= self.high_score_threshold: - min_votes = max(min_votes, min(self.min_votes_high, vote_cap)) - - # Escalate when uncertain OR when the vote-depth floor is not yet met. - if vote_cap > 1 and (self._is_uncertain(rate, margin_thr) or len(votes) < min_votes): - sp = self._score_sampling_params(sampling_params, temperature=0.7) - while len(votes) < vote_cap: - extra = self._score_once( - trajectory=score_traj, sampling_params=sp, - query=query, rubric_key=rubric_key) - votes.append(self._parse_verdicts(extra, n)) - rate = self._vote_rates(votes) - if len(votes) >= min_votes and not self._is_uncertain(rate, margin_thr): - break - return rate, len(votes) - - def _trim_segment_for_llm(self, text: str) -> str: - cap = self.max_segment_chars - if len(text) <= cap: - return text - head = cap // 2 - 96 - tail = cap // 2 - 96 - omitted = len(text) - head - tail - return (f'{text[:head]}\n\n[... {omitted} chars omitted for rubric scoring ...]\n\n' - f'{text[-tail:]}') - - # ------------------------------------------------------------------ - # code-level hard verification (no LLM) - # ------------------------------------------------------------------ - @staticmethod - def _code_hard_checks(trajectory: dict) -> Tuple[float, bool]: - """Deterministic tool-call checks -> (pass_rate, has_any_hard_signal). - - Each assistant tool_call contributes checks: - - arguments parse as JSON (schema-ish validity), - - a following tool message exists and is non-empty / non-ERROR. - Returns pass_rate over all such checks; has_hard=False when the - segment has no tool calls (nothing to code-verify). - """ - msgs = trajectory.get('messages', []) or [] - n_msgs = len(msgs) - checks: List[bool] = [] - for i, m in enumerate(msgs): - if m.get('role') != 'assistant': - continue - tool_calls = m.get('tool_calls') or [] - for tc in tool_calls: - fn = (tc.get('function') or {}) if isinstance(tc, dict) else {} - args = fn.get('arguments', '') - # 1) argument validity - checks.append(_is_valid_json_args(args)) - # 2) execution success: find the matching/following tool message - ok = False - j = i + 1 - while j < n_msgs and msgs[j].get('role') == 'tool': - content = msgs[j].get('content') or '' - text = content if isinstance(content, str) else str(content) - if text.strip() and not text.lstrip().startswith('ERROR'): - ok = True - break - j += 1 - checks.append(ok) - if not checks: - return 1.0, False - return sum(1 for c in checks if c) / len(checks), True - - # ------------------------------------------------------------------ - # aggregation & mapping - # ------------------------------------------------------------------ - def _aggregate(self, rubric: List[RubricItem], per_item_rate: List[float]) -> float: - """Weighted mean of per-criterion pass-rates (ARROW-style, Hard>Principle).""" - num = 0.0 - den = 0.0 - for item, rate in zip(rubric, per_item_rate): - w = self.hard_weight if item.is_hard else self.principle_weight - num += w * rate - den += w - return num / den if den else 0.0 - - def _to_level(self, scalar: float) -> int: - scalar = min(1.0, max(0.0, scalar)) - # Map [0,1] onto {0..NUM_LEVELS-1} with even-width bins. - level = int(round(scalar * (self.NUM_LEVELS - 1))) - return min(self.NUM_LEVELS - 1, max(0, level)) - - def _is_uncertain(self, per_item_rate: Sequence[float], - margin_threshold: Optional[float] = None) -> bool: - """A segment is uncertain if any criterion sits near the 0.5 boundary.""" - thr = self.margin_threshold if margin_threshold is None else margin_threshold - if not per_item_rate: - return False - # distance of the aggregate margin from a confident 0/1 verdict - for r in per_item_rate: - if abs(r - 0.5) * 2.0 < thr: - return True - return False - - @staticmethod - def _vote_rates(votes: List[List[bool]]) -> List[float]: - """Per-criterion PASS rate across votes (robust: median-of-means style). - - For each criterion we average the boolean verdicts across votes. This - equals majority-vote direction while keeping a soft rate for - aggregation. Outlier passes/fails wash out as votes accumulate. - """ - if not votes: - return [] - n = max(len(v) for v in votes) - rates: List[float] = [] - for k in range(n): - col = [v[k] for v in votes if k < len(v)] - if not col: - rates.append(0.0) - continue - rates.append(sum(1 for c in col if c) / len(col)) - return rates - - # ------------------------------------------------------------------ - # LLM sampling plumbing (mirrors Summarizer) - # ------------------------------------------------------------------ - def _llm_available(self) -> bool: - """True if a student sampler exists or a teacher API is configured. - - Mirrors the env vars ``llm_backup`` uses for its teacher; when neither a - student nor a teacher is present we must not attempt an LLM call (it would - raise a missing-credentials error inside the llm_backup teacher path). - """ - if self.sampler is not None: - return True - return bool(os.environ.get('LLM_BACKUP_API_KEY') - or os.environ.get('OPENAI_API_KEY') - or os.environ.get('LLM_BACKUP_BASE_URL')) - - def _sample_text(self, trajectory, sampling_params, lora_path) -> str: - if self.sampler is None: - return '' - sample_kwargs: dict[str, Any] = {'sampling_params': sampling_params} - if lora_path is None: - sample_kwargs['use_base_model'] = True - else: - sample_kwargs['adapter_path'] = lora_path - responses = self.sampler.sample([trajectory], **sample_kwargs) - resp = list(responses)[0] if responses else None - if resp is None: - return '' - seqs = getattr(resp, 'sequences', None) or [] - return (getattr(seqs[0], 'decoded', None) or '') if seqs else '' - - def _gen_trajectory(self, query: str, segment_text: str, - min_n: Optional[int] = None, max_n: Optional[int] = None) -> dict: - user = _fill(_GEN_USER, query=query, segment=segment_text) - system = _fill(_GEN_SYSTEM, - min_n=self.min_rubrics if min_n is None else min_n, - max_n=self.max_rubrics if max_n is None else max_n) - return {'messages': [ - {'role': 'system', 'content': system}, - {'role': 'user', 'content': user}, - ]} - - def _score_trajectory(self, query: str, rubric_block: str, segment_text: str) -> dict: - user = _fill(_SCORE_USER, query=query, rubric=rubric_block, segment=segment_text) - return {'messages': [ - {'role': 'system', 'content': _SCORE_SYSTEM}, - {'role': 'user', 'content': user}, - ]} - - def _diagnose_trajectory(self, query: str, rubric_block: str, segment_text: str) -> dict: - user = _fill(_DIAG_USER, query=query, rubric=rubric_block, segment=segment_text) - return {'messages': [ - {'role': 'system', 'content': _DIAG_SYSTEM}, - {'role': 'user', 'content': user}, - ]} - - def _gen_sampling_params(self, override): - if override is not None: - return override - if self.sampling_params is not None: - return self.sampling_params - from twinkle.data_format.sampling import SamplingParams - return SamplingParams(temperature=0.3, max_tokens=512) - - def _score_sampling_params(self, override, *, temperature: float): - if override is not None: - return override - from twinkle.data_format.sampling import SamplingParams - return SamplingParams(temperature=temperature, max_tokens=256) - - def _diagnose_sampling_params(self, override, *, temperature: float): - """Token budget for the diagnostic pass. - - Scoring emits terse PASS/FAIL lines (256 tokens is plenty), but the - diagnosis emits a full JSON object with a per-criterion reason AND fix - for EVERY rubric item. With ~7 criteria that easily exceeds 256 tokens - and the JSON gets truncated mid-string (unparsable -> all-FAIL fallback, - useless as SFT data). Give it a much larger budget, scaled by rubric size - and overridable via ``RUBRIC_DIAG_MAX_TOKENS``. - """ - if override is not None: - return override - from twinkle.data_format.sampling import SamplingParams - cap = int(os.environ.get('RUBRIC_DIAG_MAX_TOKENS', str(self.diag_max_tokens))) - return SamplingParams(temperature=temperature, max_tokens=cap) - - # ------------------------------------------------------------------ - # rendering / parsing helpers - # ------------------------------------------------------------------ - @staticmethod - def _infer_query(trajectory: dict) -> str: - for msg in trajectory.get('messages', []) or []: - if msg.get('role') == 'user': - c = msg.get('content') - if isinstance(c, str) and c.strip(): - return c.strip() - return '(no explicit query)' - - @staticmethod - def _render_segment(trajectory: dict) -> str: - """Flatten a segment's messages into a readable transcript.""" - lines: List[str] = [] - for m in trajectory.get('messages', []) or []: - role = m.get('role', '?') - if role == 'system': - continue - content = m.get('content') - if isinstance(content, list): - content = '\n'.join( - p.get('text', '') for p in content - if isinstance(p, dict) and p.get('type') == 'text') - content = content or '' - tool_calls = m.get('tool_calls') or [] - if tool_calls: - calls = '; '.join( - f"{(tc.get('function') or {}).get('name', '?')}" - f"({(tc.get('function') or {}).get('arguments', '')})" - for tc in tool_calls if isinstance(tc, dict)) - lines.append(f'[{role}] {content}\n tool_calls: {calls}'.rstrip()) - else: - lines.append(f'[{role}] {content}'.rstrip()) - return '\n'.join(lines).strip() - - _TAG_HARD_RE = re.compile(r'\[\s*hard\s*rule\s*\]', re.IGNORECASE) - _TAG_PRIN_RE = re.compile(r'\[\s*principle\s*\]', re.IGNORECASE) - _NUM_LINE_RE = re.compile(r'^\s*(?:\d+[.)]|[-*])\s*(.+?)\s*$') - - @classmethod - def _parse_rubric(cls, raw: str) -> List[RubricItem]: - items: List[RubricItem] = [] - for line in (raw or '').splitlines(): - m = cls._NUM_LINE_RE.match(line) - text = (m.group(1) if m else line).strip() - if not text: - continue - is_hard = bool(cls._TAG_HARD_RE.search(text)) - is_prin = bool(cls._TAG_PRIN_RE.search(text)) - if not (is_hard or is_prin): - # Untagged line that isn't clearly a criterion -> skip noise. - if not m: - continue - is_hard = False # default to principle - clean = cls._TAG_HARD_RE.sub('', cls._TAG_PRIN_RE.sub('', text)).strip(' .') - if clean: - items.append(RubricItem(text=clean, is_hard=is_hard)) - return items - - @staticmethod - def _render_rubric(rubric: List[RubricItem]) -> str: - return '\n'.join( - f'{i + 1}. {it.text} [{"Hard Rule" if it.is_hard else "Principle"}]' - for i, it in enumerate(rubric)) - - @staticmethod - def _parse_verdicts(raw: str, n: int) -> List[bool]: - """Parse '<i>: PASS/FAIL' lines into a length-n boolean vector. - - Missing verdicts default to FAIL (conservative for hard rules). - """ - verdicts = [False] * n - for line in (raw or '').splitlines(): - m = _VERDICT_RE.match(line) - if not m: - continue - idx = int(m.group(1)) - 1 - if 0 <= idx < n: - verdicts[idx] = m.group(2).lower() in ('pass', 'true', 'yes', '1') - return verdicts - - @classmethod - def _parse_diagnosis(cls, raw: str, n: int - ) -> Tuple[List[DiagnosisItem], bool, str]: - """Parse the diagnostic JSON into (items, overall_ok, summary). - - Robust to models that wrap JSON in code fences or add stray prose. Falls - back to the PASS/FAIL line parser when JSON is unrecoverable, so a - malformed diagnostic still yields usable verdicts (missing -> FAIL). - """ - obj = _extract_json_obj(raw) - entries: List[dict] = [] - if isinstance(obj, dict) and isinstance(obj.get('items'), list): - entries = [e for e in obj['items'] if isinstance(e, dict)] - if not entries: - # The full JSON did not parse (commonly a truncated response): salvage - # every COMPLETE ``{...}`` item object so partial diagnoses stay usable - # instead of degrading to an all-FAIL, reason-less vector. - entries = _salvage_diag_items(raw) - - items: List[DiagnosisItem] = [] - for i, entry in enumerate(entries): - try: - idx = int(entry.get('index', i + 1)) - except (TypeError, ValueError): - idx = i + 1 - verdict = str(entry.get('verdict', '')).strip().lower() in ( - 'pass', 'true', 'yes', '1', 'ok') - items.append(DiagnosisItem( - index=idx, verdict=verdict, - reason=str(entry.get('reason', '') or '').strip(), - fix=str(entry.get('fix', '') or '').strip())) - if not items: - # Last resort: the terse verdict-line parser (missing -> FAIL). - verdicts = cls._parse_verdicts(raw, n) - items = [DiagnosisItem(index=i + 1, verdict=v) - for i, v in enumerate(verdicts)] - - overall_ok = all(it.verdict for it in items) if items else False - summary = '' - if isinstance(obj, dict): - summary = str(obj.get('summary', '') or '').strip() - overall_raw = str(obj.get('overall', '') or '').strip().lower() - if overall_raw in ('ok', 'pass', 'good'): - # Trust an explicit OK only if no item contradicts it. - overall_ok = overall_ok and True - elif overall_raw in ('issues', 'issue', 'fail', 'bad'): - overall_ok = False - if not summary: - summary = ('no process errors, continue' if overall_ok - else 'process issues found') - return items, overall_ok, summary - - -# --------------------------------------------------------------------------- -# module-level helpers (comparators etc.) -# --------------------------------------------------------------------------- -def _fill(template: str, **kw) -> str: - out = template - for k, v in kw.items(): - out = out.replace('{' + k + '}', str(v)) - return out - - -def _is_valid_json_args(args: Any) -> bool: - if isinstance(args, dict): - return True - if not isinstance(args, str): - return False - s = args.strip() - if not s: - return True # a no-arg call is valid - try: - json.loads(s) - return True - except (json.JSONDecodeError, ValueError): - return False - - -def _short_hash(text: str) -> str: - import hashlib - return hashlib.md5((text or '').encode()).hexdigest()[:12] - - -_NORM_RE = re.compile(r'[^a-z0-9]+') - - -def _norm_criterion(text: str) -> str: - """Normalize criterion text for dedup (lowercase, alnum-only, first 12 words).""" - words = _NORM_RE.sub(' ', (text or '').lower()).split() - return ' '.join(words[:12]) - - -_TAG_ANY_RE = re.compile(r'\[\s*(hard\s*rule|principle)\s*\]', re.IGNORECASE) - - -def _rubric_similar(a: str, b: str) -> bool: - """Comparator for rubric generation: rubrics rarely match verbatim, so we - compare on *shape* — similar criterion count and similar hard/principle mix. - This is a proxy; downstream consistency filtering is the real quality gate. - """ - ca = _TAG_ANY_RE.findall(a or '') - cb = _TAG_ANY_RE.findall(b or '') - na, nb = len(ca), len(cb) - if na == 0 and nb == 0: - return True - if na == 0 or nb == 0: - return False - # count within +/-2 and hard-ratio within 0.34 - if abs(na - nb) > 2: - return False - hard_a = sum(1 for t in ca if t.lower().startswith('hard')) / na - hard_b = sum(1 for t in cb if t.lower().startswith('hard')) / nb - return abs(hard_a - hard_b) <= 0.34 - - -def _parse_rate(raw: str) -> Optional[float]: - pos = 0 - total = 0 - for line in (raw or '').splitlines(): - m = _VERDICT_RE.match(line) - if not m: - continue - total += 1 - if m.group(2).lower() in ('pass', 'true', 'yes', '1'): - pos += 1 - if total == 0: - return None - return pos / total - - -def _verdicts_close(a: str, b: str, tol: float = 0.25) -> bool: - """Comparator for rubric scoring: student/teacher agree when their overall - PASS rate is within ``tol`` (binned agreement, not byte-identical text).""" - ra, rb = _parse_rate(a), _parse_rate(b) - if ra is None or rb is None: - return (a or '').strip() == (b or '').strip() - return abs(ra - rb) <= tol - - -def _salvage_diag_items(raw: str) -> List[dict]: - """Recover complete diagnosis item objects from a (possibly truncated) reply. - - Scans for balanced ``{...}`` spans (string-aware, so braces inside a reason - like ``\\subsubsection*{...}`` don't corrupt the depth count) and json-parses - each object that carries a ``verdict`` key. A response cut off mid-stream - still yields every item emitted before the cut, so the diagnosis keeps its - reasons/fixes instead of collapsing to an all-FAIL, reason-less vector. - """ - if not raw: - return [] - out: List[dict] = [] - stack: List[int] = [] # start index of each open brace, by depth - in_str = False - escaped = False - for i, ch in enumerate(raw): - if in_str: - if escaped: - escaped = False - elif ch == '\\': - escaped = True - elif ch == '"': - in_str = False - continue - if ch == '"': - in_str = True - elif ch == '{': - stack.append(i) - elif ch == '}' and stack: - start = stack.pop() - frag = raw[start:i + 1] - # Only leaf-ish item objects carry a verdict; the outer envelope - # ({"items": [...]}) usually never closes when truncated anyway. - if '"verdict"' in frag and '"items"' not in frag: - try: - obj = json.loads(frag) - if isinstance(obj, dict): - out.append(obj) - except (json.JSONDecodeError, ValueError): - pass - return out - - -def _extract_json_obj(raw: str) -> Optional[dict]: - """Best-effort extraction of the first JSON object from a model response. - - Handles bare JSON, ```json fenced blocks, and JSON embedded in prose. - """ - if not raw: - return None - s = raw.strip() - # Strip a leading/trailing code fence if present. - if s.startswith('```'): - s = re.sub(r'^```[a-zA-Z]*\s*', '', s) - s = re.sub(r'\s*```$', '', s).strip() - try: - obj = json.loads(s) - return obj if isinstance(obj, dict) else None - except (json.JSONDecodeError, ValueError): - pass - # Fall back to the widest {...} span. - start = s.find('{') - end = s.rfind('}') - if start != -1 and end > start: - try: - obj = json.loads(s[start:end + 1]) - return obj if isinstance(obj, dict) else None - except (json.JSONDecodeError, ValueError): - return None - return None - - -def _diag_rate(raw: str) -> Optional[float]: - """Overall PASS rate from a diagnostic JSON response (for the comparator).""" - obj = _extract_json_obj(raw) - if isinstance(obj, dict) and isinstance(obj.get('items'), list): - verdicts = [str(e.get('verdict', '')).strip().lower() in - ('pass', 'true', 'yes', '1', 'ok') - for e in obj['items'] if isinstance(e, dict)] - if verdicts: - return sum(1 for v in verdicts if v) / len(verdicts) - return _parse_rate(raw) - - -def _diag_verdicts_close(a: str, b: str, tol: float = 0.25) -> bool: - """Comparator for the diagnostic pass: student/teacher agree when their - overall PASS rate is within ``tol``. Reasons are free text, so we match on - the verdict vector (two valid phrasings of the same verdict count as a - match), not on byte-identical explanations.""" - ra, rb = _diag_rate(a), _diag_rate(b) - if ra is None or rb is None: - return (a or '').strip() == (b or '').strip() - return abs(ra - rb) <= tol diff --git a/tests/preprocessor/test_refuse_filter.py b/tests/preprocessor/test_refuse_filter.py index ab4e59202..4eb8cfc5f 100644 --- a/tests/preprocessor/test_refuse_filter.py +++ b/tests/preprocessor/test_refuse_filter.py @@ -154,8 +154,11 @@ def test_custom_window_includes_late_refusal(self): text = 'a' * 700 + " I can't help you complete that task." assert _is_refusal(text, check_window=1000) is True - def test_zero_window_finds_nothing(self): - assert _is_refusal("I can't help you complete tasks.", check_window=0) is False + def test_zero_window_scans_whole_text(self): + # check_window <= 0 disables truncation, so even a refusal past the + # default 600-char window is found. + text = 'a' * 700 + " I can't help you complete that task." + assert _is_refusal(text, check_window=0) is True # ── RefuseFilter pipeline ─────────────────────────────────────────────────── @@ -193,29 +196,34 @@ def test_keeps_normal_reply(self): ] assert len(_fil(rows)) == 1 - def test_only_first_assistant_scanned(self): - # Refusal in SECOND assistant turn → kept (filter only checks first). - rows = [ - _row([ - { - 'role': 'user', - 'content': 'q1' - }, - { - 'role': 'assistant', - 'content': 'A clean reply.' - }, - { - 'role': 'user', - 'content': 'q2' - }, - { - 'role': 'assistant', - 'content': "I can't help with that." - }, - ]) - ] - assert len(_fil(rows)) == 1 + def _late_refusal_row(self): + # A clean first reply, then a refusal in the SECOND assistant turn. + return _row([ + { + 'role': 'user', + 'content': 'q1' + }, + { + 'role': 'assistant', + 'content': 'A clean reply.' + }, + { + 'role': 'user', + 'content': 'q2' + }, + { + 'role': 'assistant', + 'content': "I can't help with that." + }, + ]) + + def test_late_refusal_dropped_by_default(self): + # scan_all_assistants defaults to True: a conversation that only refuses + # in a later turn is still a refusal. + assert _fil([self._late_refusal_row()]) == [] + + def test_late_refusal_kept_when_only_first_scanned(self): + assert len(_fil([self._late_refusal_row()], scan_all_assistants=False)) == 1 def test_think_block_stripped(self): # Refusal phrasing inside <think>...</think> must NOT trigger. diff --git a/tests/preprocessor/test_value_selector.py b/tests/preprocessor/test_value_selector.py deleted file mode 100644 index 82ff5ccb2..000000000 --- a/tests/preprocessor/test_value_selector.py +++ /dev/null @@ -1,234 +0,0 @@ -"""ValueSelector: deterministic value_score + component boundaries + top-frac gate.""" - -from twinkle_agentic.preprocessor import ValueSelector -from twinkle_agentic.preprocessor import label_schema as L - - -def _single_turn(): - return {'messages': [ - {'role': 'user', 'content': '写一句短视频文案'}, - {'role': 'assistant', 'content': '这价格我不敢信，现在下单立省八千。'}, - ], 'user_data': []} - - -def _tool_call(name, args, result='ok', content=''): - return [ - {'role': 'assistant', 'content': content, - 'tool_calls': [{'id': 't', 'type': 'function', - 'function': {'name': name, 'arguments': args}}]}, - {'role': 'tool', 'tool_call_id': 't', 'content': result}, - ] - - -def _agent_row(steps): - msgs = [{'role': 'user', 'content': 'do the task'}] - for name, args, result in steps: - msgs += _tool_call(name, args, result) - return {'messages': msgs, 'user_data': []} - - -def _val(row): - return L.get_label(row, L.KEY_VALUE_SCORE) - - -def _meta(row): - return L.get_label(row, L.KEY_VALUE_META) or {} - - -def test_single_turn_scores_low(): - out, dropped = ValueSelector()([_single_turn()]) - assert dropped == [] - v = _val(out[0]) - assert v is not None and v < 0.25 - m = _meta(out[0]) - assert m['uncertainty'] == 0.0 # single all-pass round -> decided - - -def test_tool_error_raises_error_signal(): - # a tool that returns an ERROR result -> tool_executed miss -> error signal up - row = _agent_row([('search', '{"q":"x"}', 'ERROR: not found'), - ('search', '{"q":"y"}', 'ERROR: not found')]) - out, _ = ValueSelector()([row]) - assert _meta(out[0])['error'] > 0.0 - assert _val(out[0]) > _val(ValueSelector()([_single_turn()])[0][0]) - - -def test_long_agent_scores_higher_than_single_turn(): - steps = [(f'read', '{"p":"%d"}' % i, 'contents') for i in range(6)] - steps += [('grep', '{"q":"a"}', 'hit'), ('edit', '{"f":"b"}', 'done')] - long_row = _agent_row(steps) - out_long, _ = ValueSelector()([long_row]) - out_short, _ = ValueSelector()([_single_turn()]) - assert _meta(out_long[0])['difficulty'] > _meta(out_short[0])['difficulty'] - assert _val(out_long[0]) > _val(out_short[0]) - - -def test_bad_row_never_crashes(): - out, dropped = ValueSelector()([{'messages': None, 'user_data': []}, - {'messages': [], 'user_data': []}]) - assert dropped == [] - assert all(_val(r) == 0.0 for r in out) - - -def test_select_top_for_rubric_marks_global_top(): - from datasets import Dataset as HFDataset - - class _Wrap: - def __init__(self, hf): - self.dataset = hf - - from twinkle_agentic.preprocessor import select_top_for_rubric - - def _row_with_value(v): - return {'messages': [{'role': 'user', 'content': 'x'}], - 'user_data': [(L.KEY_VALUE_SCORE, str(v))]} - - hf = HFDataset.from_list([_row_with_value(v) for v in [0.1, 0.9, 0.5, 0.8, 0.2]]) - ds = _Wrap(hf) - ds, n_sel = select_top_for_rubric(ds, select_frac=0.4) # top 2 of 5 - assert n_sel == 2 - selected = [L.get_label(ds.dataset[i], L.KEY_SELECTED_FOR_RUBRIC, False) - for i in range(len(ds.dataset))] - # rows with value 0.9 and 0.8 are the top-2 - assert selected == [False, True, False, True, False] - - -class _SpyRubric: - """Records how many times the rubric was invoked.""" - max_votes = 1 - - def __init__(self): - self.calls = 0 - self.diagnose_calls = 0 - - def score_detail(self, segment, query=None, intent=None, extra_context=None): - self.calls += 1 - - class _D: - scalar = 0.5 - votes = [0.5] - return _D() - - def diagnose(self, segment, query=None, intent=None): - self.diagnose_calls += 1 - - class _Item: - def __init__(self, index, verdict, reason, fix): - self.index, self.verdict, self.reason, self.fix = index, verdict, reason, fix - - class _RItem: - def __init__(self, text, is_hard): - self.text, self.is_hard = text, is_hard - - class _Diag: - scalar = 0.5 - overall_ok = False - summary = 'one criterion failed' - query = 'do it' - segment_text = 'user: do it ...' - raw = '1. FAIL: boom -> retry' - rubric = [_RItem('args are valid JSON', True)] - items = [_Item(0, False, 'tool errored', 'retry with valid args')] - return _Diag() - - -def _low_quality_agent(): - # a tool call whose result errors -> low hard score -> not short-circuited, - # so fuse_segment will actually reach for the rubric (unless gated). - return {'messages': [ - {'role': 'user', 'content': 'do it'}, - {'role': 'assistant', 'content': '', - 'tool_calls': [{'id': 't', 'type': 'function', - 'function': {'name': 'f', 'arguments': '{}'}}]}, - {'role': 'tool', 'tool_call_id': 't', 'content': 'ERROR: boom'}, - ], 'user_data': []} - - -def test_gate_skips_rubric_for_unselected_row(): - from twinkle_agentic.preprocessor import TrajectoryScorer - - spy = _SpyRubric() - scorer = TrajectoryScorer(rubric_verifier=spy, calibrate=False) - - gated = L.set_label(_low_quality_agent(), L.KEY_SELECTED_FOR_RUBRIC, False) - scorer([gated]) - assert spy.calls == 0 # not selected -> no LLM rubric - - selected = L.set_label(_low_quality_agent(), L.KEY_SELECTED_FOR_RUBRIC, True) - scorer([selected]) - assert spy.calls >= 1 # selected -> rubric runs - - -def test_persist_diagnosis_writes_verdict_reason_fix(): - from twinkle_agentic.preprocessor import TrajectoryScorer - - spy = _SpyRubric() - scorer = TrajectoryScorer(rubric_verifier=spy, calibrate=False, - persist_diagnosis=True) - - selected = L.set_label(_low_quality_agent(), L.KEY_SELECTED_FOR_RUBRIC, True) - out, _ = scorer([selected]) - diag = L.get_label(out[0], L.KEY_RUBRIC_DIAGNOSIS) - assert spy.diagnose_calls >= 1 # diagnosis ran for the scored segment - assert isinstance(diag, list) and diag # persisted, one entry per segment - entry = diag[0] - assert entry['overall_ok'] is False - assert entry['raw'] == '1. FAIL: boom -> retry' # SFT target - assert entry['segment_text'] and entry['query'] # SFT inputs - assert entry['items'][0]['verdict'] is False - assert entry['items'][0]['reason'] == 'tool errored' # the "why" - assert entry['items'][0]['fix'] == 'retry with valid args' - assert entry['rubric'][0]['text'] == 'args are valid JSON' - - -def test_no_diagnosis_for_gated_out_row(): - from twinkle_agentic.preprocessor import TrajectoryScorer - - spy = _SpyRubric() - scorer = TrajectoryScorer(rubric_verifier=spy, calibrate=False, - persist_diagnosis=True) - gated = L.set_label(_low_quality_agent(), L.KEY_SELECTED_FOR_RUBRIC, False) - out, _ = scorer([gated]) - assert spy.diagnose_calls == 0 # gated -> no LLM at all - assert L.get_label(out[0], L.KEY_RUBRIC_DIAGNOSIS) is None - - -class _SpySafety: - """Stands in for the RubricVerifier a SafetyScorer holds.""" - fixed_rubric = None - - def __init__(self): - self.calls = 0 - - def score_detail(self, trajectory, **kwargs): - self.calls += 1 - - class _D: - scalar = 1.0 - return _D() - - -def test_safety_gate_skips_llm_for_unselected_row(): - from twinkle_agentic.preprocessor import SafetyScorer - - spy = _SpySafety() - scorer = SafetyScorer(rubric_verifier=spy, gate_label=L.KEY_SELECTED_FOR_RUBRIC) - - gated = L.set_label(_low_quality_agent(), L.KEY_SELECTED_FOR_RUBRIC, False) - out, _ = scorer([gated]) - assert spy.calls == 0 # not selected -> no LLM safety pass - assert L.get_label(out[0], L.KEY_SAFETY_SCORE) == 1.0 # tagged neutral-safe - assert L.get_label(out[0], L.KEY_SAFETY_UNSAFE) is False - - selected = L.set_label(_low_quality_agent(), L.KEY_SELECTED_FOR_RUBRIC, True) - scorer([selected]) - assert spy.calls >= 1 # selected -> safety LLM runs - - -def test_safety_no_gate_scores_every_row(): - from twinkle_agentic.preprocessor import SafetyScorer - - spy = _SpySafety() - scorer = SafetyScorer(rubric_verifier=spy) # gate_label=None -> score all - scorer([L.set_label(_low_quality_agent(), L.KEY_SELECTED_FOR_RUBRIC, False)]) - assert spy.calls == 1 # ungated: LLM runs even for a "not selected" row diff --git a/tests/twinkle_agentic/test_agentic_rsi.py b/tests/twinkle_agentic/test_agentic_rsi.py new file mode 100644 index 000000000..a2c67ff1c --- /dev/null +++ b/tests/twinkle_agentic/test_agentic_rsi.py @@ -0,0 +1,260 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Tests for the agentic building blocks: program checks and the ms-agent Env. + +No GPU and no ms-agent runtime: the Env is driven with a fake ToolManager that +records what it was asked to run, which is enough to pin the two behaviours the +trainer depends on -- calls arrive batched, and a check's exit status survives +the round trip through a text-only tool. +""" +import json +import os +import sys +import tempfile +import unittest + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'src')) + +from twinkle_agentic.envs.env_tool import EnvTool # noqa: E402 +from twinkle_agentic.envs.ms_agent_tool_env import MsAgentToolEnv # noqa: E402 +from twinkle_agentic.tools.tool_manager import ToolManager # noqa: E402 +from twinkle_agentic.verifier.result_check import (Check, CheckContext, # noqa: E402 + checks_from_dicts, run_checks) + + +class FakeMsToolManager: + """Stands in for ms-agent's ToolManager, recording dispatch shape.""" + + # ms-agent namespaces tools as ``{server}---{tool}``; keep that here so the + # tests exercise the same name resolution production hits. + TOOLS = [ + {'tool_name': 'code_executor---shell_executor'}, + {'tool_name': 'code_executor---python_executor'}, + {'tool_name': 'file_system---write_file'}, + ] + + def __init__(self, handler=None): + self.single_calls = [] + self.batch_sizes = [] + self._handler = handler or (lambda call: f'ran {call["tool_name"]}') + + async def get_tools(self): + return list(self.TOOLS) + + async def single_call_tool(self, tool_info): + self.single_calls.append(tool_info) + return self._handler(tool_info) + + async def parallel_call_tool(self, tool_list, on_result=None): + self.batch_sizes.append(len(tool_list)) + return [self._handler(call) for call in tool_list] + + async def cleanup(self): + pass + + +class ResultCheckFileTest(unittest.TestCase): + + def setUp(self): + self.tmp = tempfile.mkdtemp(prefix='rescheck_test_') + with open(os.path.join(self.tmp, 'report.md'), 'w', encoding='utf-8') as f: + f.write('# Sales Report\n- Q1\n- Q2\n- Q3\n- Q4\n') + with open(os.path.join(self.tmp, 'data.json'), 'w', encoding='utf-8') as f: + json.dump({'result': {'items': [{'n': 7}]}}, f) + + def ctx(self, answer=''): + return CheckContext(workspace=self.tmp, final_answer=answer) + + def test_file_exists_and_absent(self): + report = run_checks([ + Check(kind='file_exists', path='report.md'), + Check(kind='file_absent', path='nope.txt'), + ], self.ctx()) + self.assertTrue(report.all_passed) + self.assertEqual(report.score, 1.0) + + def test_file_contains_value_and_pattern(self): + report = run_checks([ + Check(kind='file_contains', path='report.md', value='# Sales Report'), + Check(kind='file_contains', path='report.md', pattern=r'(?s)Q1.*Q4'), + ], self.ctx()) + self.assertTrue(report.all_passed) + + def test_missing_file_fails_with_reason(self): + report = run_checks([Check(kind='file_contains', path='gone.md', value='x')], self.ctx()) + self.assertFalse(report.all_passed) + self.assertIn('does not exist', report.failures()[0]) + + def test_file_json_dotted_key_including_list_index(self): + report = run_checks( + [Check(kind='file_json', path='data.json', key='result.items.0.n', value=7)], + self.ctx()) + self.assertTrue(report.all_passed) + + def test_path_escaping_workspace_is_rejected(self): + report = run_checks([Check(kind='file_exists', path='../../etc/passwd')], self.ctx()) + self.assertFalse(report.all_passed) + self.assertIn('escapes the workspace', report.failures()[0]) + + def test_empty_checks_score_zero_not_one(self): + # A task with no checks must not look solved. + report = run_checks([], self.ctx()) + self.assertEqual(report.score, 0.0) + self.assertEqual(report.n_total, 0) + + def test_fraction_vs_all_or_nothing(self): + checks = [Check(kind='file_exists', path='report.md'), + Check(kind='file_exists', path='missing.md')] + self.assertEqual(run_checks(checks, self.ctx(), mode='fraction').score, 0.5) + self.assertEqual(run_checks(checks, self.ctx(), mode='all_or_nothing').score, 0.0) + + def test_weight_shifts_partial_credit(self): + checks = [Check(kind='file_exists', path='report.md', weight=3.0), + Check(kind='file_exists', path='missing.md', weight=1.0)] + self.assertAlmostEqual(run_checks(checks, self.ctx()).score, 0.75) + + def test_answer_kinds(self): + report = run_checks([ + Check(kind='answer_contains', value='Alibaba'), + Check(kind='answer_regex', pattern=r'(?i)qwen\d'), + ], self.ctx(answer='Qwen3 was published by Alibaba.')) + self.assertTrue(report.all_passed) + + def test_local_shell_and_python_run_in_workspace(self): + report = run_checks([ + Check(kind='shell', code='test -f report.md'), + Check(kind='python', code='open("report.md").read()'), + ], self.ctx()) + self.assertTrue(report.all_passed, report.failures()) + + def test_failing_python_check_reports_nonzero(self): + report = run_checks([Check(kind='python', code='assert 1 == 2')], self.ctx()) + self.assertFalse(report.all_passed) + + def test_bad_kind_rejected_at_construction(self): + with self.assertRaises(ValueError): + Check(kind='definitely_not_a_kind') + + def test_checks_from_dicts(self): + checks = checks_from_dicts([{'kind': 'file_exists', 'path': 'a'}]) + self.assertEqual(checks[0].kind, 'file_exists') + + +class MsAgentToolEnvTest(unittest.TestCase): + + def setUp(self): + self.tmp = tempfile.mkdtemp(prefix='envtest_') + self.tm = FakeMsToolManager() + self.env = MsAgentToolEnv(tool_manager=self.tm, workspace=self.tmp) + + def test_step_forwards_name_and_arguments(self): + result = self.env.step('read_file', {'path': 'a.txt'}) + self.assertEqual(self.tm.single_calls[0], + {'tool_name': 'read_file', 'arguments': {'path': 'a.txt'}}) + self.assertEqual(result.observation, 'ran read_file') + + def test_step_batch_uses_one_parallel_call(self): + results = self.env.step_batch([('read_file', {'p': 1}), ('grep', {'q': 'x'})]) + self.assertEqual(self.tm.batch_sizes, [2]) + self.assertEqual([r.observation for r in results], ['ran read_file', 'ran grep']) + + def test_single_call_batch_does_not_go_through_parallel(self): + self.env.step_batch([('glob', {})]) + self.assertEqual(self.tm.batch_sizes, []) + self.assertEqual(len(self.tm.single_calls), 1) + + def test_observation_is_truncated(self): + tm = FakeMsToolManager(handler=lambda call: 'x' * 50) + env = MsAgentToolEnv(tool_manager=tm, workspace=self.tmp, max_observation_chars=10) + obs = env.step('grep', {}).observation + self.assertTrue(obs.startswith('x' * 10)) + self.assertIn('truncated 40 chars', obs) + + def test_non_string_result_is_json_encoded(self): + tm = FakeMsToolManager(handler=lambda call: {'ok': True}) + env = MsAgentToolEnv(tool_manager=tm, workspace=self.tmp) + self.assertEqual(env.step('t', {}).observation, '{"ok": true}') + + def test_requires_a_tool_manager(self): + with self.assertRaises(ValueError): + MsAgentToolEnv() + + def test_runner_recovers_exit_code_from_text_output(self): + # The sandbox tools return prose; the marker is how the exit status + # survives. Emulate a shell that echoes the marker. Matching on the + # namespaced name also proves the plain name was resolved. + def handler(call): + if call['tool_name'] == 'code_executor---shell_executor': + return 'some output\n__TWINKLE_RC__:0' + return '__TWINKLE_RC__:3' + + env = MsAgentToolEnv(tool_manager=FakeMsToolManager(handler), workspace=self.tmp) + runner = env.runner() + self.assertEqual(runner('ls', 'shell'), (0, 'some output')) + self.assertEqual(runner('boom()', 'python')[0], 3) + + def test_resolve_tool_maps_plain_name_onto_namespaced_one(self): + env = MsAgentToolEnv(tool_manager=FakeMsToolManager(), workspace=self.tmp) + self.assertEqual(env.resolve_tool('shell_executor'), + 'code_executor---shell_executor') + # An already-qualified name is left alone. + self.assertEqual(env.resolve_tool('file_system---write_file'), + 'file_system---write_file') + + def test_resolve_tool_raises_on_unknown_name(self): + # Silently passing a bad name through would surface as a failed check, + # which is indistinguishable from the task genuinely not being solved. + env = MsAgentToolEnv(tool_manager=FakeMsToolManager(), workspace=self.tmp) + with self.assertRaises(ValueError): + env.resolve_tool('no_such_tool') + + def test_runner_missing_marker_is_a_failure_not_a_pass(self): + env = MsAgentToolEnv(tool_manager=FakeMsToolManager(lambda c: 'sandbox died'), + workspace=self.tmp) + code, out = env.runner()('ls', 'shell') + self.assertNotEqual(code, 0) + self.assertIn('sandbox died', out) + + def test_checks_run_through_the_env_runner(self): + env = MsAgentToolEnv( + tool_manager=FakeMsToolManager(lambda c: '__TWINKLE_RC__:0'), + workspace=self.tmp) + report = run_checks([Check(kind='shell', code='true')], + CheckContext(workspace=self.tmp, runner=env.runner())) + self.assertTrue(report.all_passed) + + +class ToolBridgeTest(unittest.TestCase): + """The prompt's tool list and the executing tool list must be one list.""" + + def setUp(self): + self.tmp = tempfile.mkdtemp(prefix='bridge_') + self.tm = FakeMsToolManager() + self.env = MsAgentToolEnv(tool_manager=self.tm, workspace=self.tmp) + self.schemas = [ + {'type': 'function', 'function': {'name': 'read_file', 'parameters': {}}}, + {'type': 'function', 'function': {'name': 'shell_executor', 'parameters': {}}}, + ] + + def test_from_schemas_binds_every_declared_tool(self): + manager = ToolManager(EnvTool.from_schemas(self.env, self.schemas)) + self.assertEqual(sorted(manager.names()), ['read_file', 'shell_executor']) + + def test_declared_tools_collapse_into_one_step_batch(self): + manager = ToolManager(EnvTool.from_schemas(self.env, self.schemas)) + calls = [ + {'id': '1', 'type': 'function', + 'function': {'name': 'read_file', 'arguments': '{"path": "a"}'}}, + {'id': '2', 'type': 'function', + 'function': {'name': 'shell_executor', 'arguments': '{"command": "ls"}'}}, + ] + out = manager.call_many(calls) + self.assertEqual(self.tm.batch_sizes, [2]) + self.assertEqual(out, ['ran read_file', 'ran shell_executor']) + + def test_nameless_schema_is_refused(self): + with self.assertRaises(ValueError): + EnvTool.from_schemas(self.env, [{'type': 'function', 'function': {}}]) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/twinkle_agentic/test_aggregation_fusion.py b/tests/twinkle_agentic/test_aggregation_fusion.py deleted file mode 100644 index 272981e0c..000000000 --- a/tests/twinkle_agentic/test_aggregation_fusion.py +++ /dev/null @@ -1,12 +0,0 @@ -from twinkle_agentic.verifier.aggregation import _combine - - -def test_hard_soft_blend_high_hard_not_one_shot_veto(): - # product would be 0.9 * 0.1 = 0.09 - blended = _combine(0.9, 0.1, 'hard_soft_blend') - assert blended > 0.35 - assert blended < 0.9 - - -def test_hard_soft_blend_low_hard_uses_soft(): - assert _combine(0.5, 0.4, 'hard_soft_blend') == 0.5 * 0.4 diff --git a/tests/twinkle_agentic/test_diagnosis_salvage.py b/tests/twinkle_agentic/test_diagnosis_salvage.py deleted file mode 100644 index 3be52adaa..000000000 --- a/tests/twinkle_agentic/test_diagnosis_salvage.py +++ /dev/null @@ -1,62 +0,0 @@ -"""Diagnosis parsing must survive truncated JSON. - -The diagnostic pass emits a full per-criterion (verdict+reason+fix) JSON. When a -teacher reply is cut off mid-stream (too small a token budget), the outer object -never closes. We must still recover every COMPLETE item object rather than -degrade to an all-FAIL, reason-less vector (which is useless as SFT data). -""" -from twinkle_agentic.verifier.rubric_verifier import (RubricVerifier, - _salvage_diag_items) - -# A realistic truncated response: 3 complete items (one reason contains LaTeX -# braces to exercise string-aware brace matching), then cut off mid-string. -TRUNCATED = ( - '{\n "items": [\n' - ' {"index": 1, "verdict": "PASS", "reason": "No tool calls needed.", "fix": ""},\n' - ' {"index": 2, "verdict": "FAIL", "reason": "Duplicated \\\\subsubsection*{ans} block.", ' - '"fix": "Remove the duplicate."},\n' - ' {"index": 3, "verdict": "PASS", "reason": "Correctly omits \\\\begin{document}.", "fix": ""},\n' - ' {"index": 4, "verdict": "FAIL", "reason": "The response is cut off right he' -) - -WELL_FORMED = ( - '{"items": [{"index": 1, "verdict": "PASS", "reason": "ok", "fix": ""},' - '{"index": 2, "verdict": "FAIL", "reason": "bad", "fix": "do x"}],' - ' "overall": "issues", "summary": "one failure"}' -) - - -def test_salvage_recovers_complete_items_from_truncated_json(): - items = _salvage_diag_items(TRUNCATED) - # 3 complete objects; the 4th (truncated) is dropped. - assert len(items) == 3 - assert [it['verdict'] for it in items] == ['PASS', 'FAIL', 'PASS'] - # braces inside the reason string must not corrupt matching - assert 'subsubsection' in items[1]['reason'] - - -def test_parse_diagnosis_truncated_keeps_reasons_and_verdicts(): - items, overall_ok, summary = RubricVerifier._parse_diagnosis(TRUNCATED, n=7) - assert len(items) == 3 # not the all-FAIL length-7 fallback - assert items[0].verdict is True - assert items[1].verdict is False - assert items[1].reason # the "why" survives - assert items[1].fix == 'Remove the duplicate.' - assert overall_ok is False # a FAIL present -> not ok - - -def test_parse_diagnosis_well_formed_still_works(): - items, overall_ok, summary = RubricVerifier._parse_diagnosis(WELL_FORMED, n=2) - assert len(items) == 2 - assert items[0].verdict is True and items[1].verdict is False - assert items[1].fix == 'do x' - assert overall_ok is False - assert summary == 'one failure' - - -def test_diag_sampling_params_bigger_than_scoring(): - rv = RubricVerifier(diag_max_tokens=2048) - diag = rv._diagnose_sampling_params(None, temperature=0.0) - score = rv._score_sampling_params(None, temperature=0.0) - assert diag.max_tokens >= 2048 - assert diag.max_tokens > score.max_tokens # diagnosis needs a bigger budget diff --git a/tests/twinkle_agentic/test_extract_condensed.py b/tests/twinkle_agentic/test_extract_condensed.py deleted file mode 100644 index c5aa726ff..000000000 --- a/tests/twinkle_agentic/test_extract_condensed.py +++ /dev/null @@ -1,422 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Unit tests for :class:`twinkle_agentic.tools.extract_condensed.ExtractCondensed`. - -Covers: -- block-index enumeration matches :meth:`Chunks.to_trajectory` exactly -- retrieval returns pre-compression text when ``raw.original`` is present -- fallback to current ``content`` when ``raw.original`` missing -- bad / missing arguments produce actionable error strings (no exceptions) -- tool metadata is complete and JSON-serializable -- integration with :class:`ToolManager` -- end-to-end: KeywordCondenser → Chunks → ExtractCondensed round-trips -""" -from __future__ import annotations - -import json -import pytest - -from twinkle_agentic.data_format import Chunks -from twinkle_agentic.tools.extract_condensed import TOOL_NAME, ExtractCondensed -from twinkle_agentic.tools.tool_manager import ToolManager - - -# --------------------------------------------------------------------------- -# helpers -# --------------------------------------------------------------------------- -def _condensed(content, *, original=None, role='user', round_idx=1): - raw = {'condensed': True} - if original is not None: - raw['original'] = original - ch = {'type': 'text', 'role': role, 'content': content, 'raw': raw, 'round': round_idx} - return ch - - -def _plain(content, *, role='user'): - return {'type': 'text', 'role': role, 'content': content} - - -# --------------------------------------------------------------------------- -# block enumeration parity with Chunks.to_trajectory -# --------------------------------------------------------------------------- -def test_blocks_indexed_from_1_in_document_order(): - chunks = Chunks(chunks=[ - _condensed('cmp1', original='orig one'), - _condensed('cmp2', original='orig two'), - _condensed('cmp3', original='orig three'), - ]) - tool = ExtractCondensed(chunks) - assert tool.blocks == [1, 2, 3] - assert len(tool) == 3 - assert 1 in tool and 3 in tool and 4 not in tool - - -def test_non_condensed_text_chunks_are_not_indexed(): - chunks = Chunks(chunks=[ - _plain('system prelude', role='system'), # not condensed - _condensed('cmp1', original='orig one'), - _plain('user follow-up'), # not condensed - _condensed('cmp2', original='orig two'), - ]) - tool = ExtractCondensed(chunks) - assert tool.blocks == [1, 2] - assert tool(TOOL_NAME, {'block': 1}) == 'orig one' - assert tool(TOOL_NAME, {'block': 2}) == 'orig two' - - -def test_tool_role_condensed_chunks_are_skipped(): - # Mirrors Chunks.to_trajectory: role=='tool' is NEVER wrapped, even - # if marked condensed, so it must not consume a block index either. - chunks = Chunks(chunks=[ - _condensed('cmp_user', original='user orig', role='user'), - _condensed('cmp_tool', original='tool orig', role='tool'), - _condensed('cmp_asst', original='asst orig', role='assistant'), - ]) - tool = ExtractCondensed(chunks) - # Only the user + assistant blocks count. - assert tool.blocks == [1, 2] - assert tool(TOOL_NAME, {'block': 1}) == 'user orig' - assert tool(TOOL_NAME, {'block': 2}) == 'asst orig' - - -def test_empty_content_condensed_chunks_are_skipped(): - chunks = Chunks(chunks=[ - _condensed('', original=''), # empty, skipped - _condensed('cmp', original='orig'), - ]) - tool = ExtractCondensed(chunks) - assert tool.blocks == [1] - assert tool(TOOL_NAME, {'block': 1}) == 'orig' - - -def test_non_text_chunks_ignored(): - chunks = Chunks(chunks=[ - { - 'type': 'image', - 'content': 'image bytes', - 'raw': { - 'type': 'image', - 'image': 'x' - }, - 'role': 'user' - }, - _condensed('cmp', original='orig text'), - ]) - tool = ExtractCondensed(chunks) - assert tool.blocks == [1] - assert tool(TOOL_NAME, {'block': 1}) == 'orig text' - - -# --------------------------------------------------------------------------- -# retrieval semantics -# --------------------------------------------------------------------------- -def test_returns_original_when_present(): - chunks = Chunks(chunks=[_condensed('CMP', original='THE ORIGINAL')]) - tool = ExtractCondensed(chunks) - assert tool(TOOL_NAME, {'block': 1}) == 'THE ORIGINAL' - - -def test_missing_original_returns_error_not_compressed_content(): - # Contract: ExtractCondensed returns the *original* text. When the - # upstream pipeline forgot to snapshot it, the tool MUST fail loud - # rather than silently handing back the compressed stand-in, which - # would deceive the LLM into thinking it had recovered the source. - chunks = Chunks(chunks=[_condensed('CMP', original=None)]) - tool = ExtractCondensed(chunks) - # The block is still enumerated so numbering stays aligned. - assert tool.blocks == [1] - out = tool(TOOL_NAME, {'block': 1}) - assert out.startswith('Error:') - assert 'no original-text snapshot' in out - # And crucially, the compressed stand-in is NOT leaked. - assert 'CMP' not in out - - -def test_original_empty_string_also_reports_missing_snapshot(): - chunks = Chunks(chunks=[_condensed('CMP', original='')]) - tool = ExtractCondensed(chunks) - out = tool(TOOL_NAME, {'block': 1}) - assert out.startswith('Error:') - assert 'no original-text snapshot' in out - - -# --------------------------------------------------------------------------- -# bad input handling (never raises) -# --------------------------------------------------------------------------- -def test_missing_block_argument_returns_error_string(): - tool = ExtractCondensed(Chunks(chunks=[_condensed('cmp', original='orig')])) - out = tool(TOOL_NAME, {}) - assert out.startswith('Error: missing required argument') - - -def test_non_integer_block_returns_error_string(): - tool = ExtractCondensed(Chunks(chunks=[_condensed('cmp', original='orig')])) - for bad in ('abc', [], {}, None): - out = tool(TOOL_NAME, {'block': bad}) - assert out.startswith('Error:'), (bad, out) - - -def test_bool_block_is_rejected_not_coerced_to_int(): - # ``bool`` is a subclass of ``int`` so ``int(True) == 1``. Without - # an explicit guard, ``{'block': True}`` would silently retrieve - # block 1 -- a nasty footgun if an LLM stringifies a truthy flag. - tool = ExtractCondensed(Chunks(chunks=[_condensed('cmp', original='orig1')])) - out_true = tool(TOOL_NAME, {'block': True}) - assert out_true.startswith('Error:') and 'bool' in out_true - out_false = tool(TOOL_NAME, {'block': False}) - assert out_false.startswith('Error:') and 'bool' in out_false - # Sanity: the real integer 1 still works. - assert tool(TOOL_NAME, {'block': 1}) == 'orig1' - - -def test_float_block_is_rejected_not_silently_truncated(): - # ``int(1.9) == 1`` would silently round a float down; reject it. - tool = ExtractCondensed(Chunks(chunks=[_condensed('cmp', original='orig1')])) - out = tool(TOOL_NAME, {'block': 1.9}) - assert out.startswith('Error:') and 'float' in out - # And floats that happen to be integer-valued are also rejected to - # keep the contract simple. - out2 = tool(TOOL_NAME, {'block': 1.0}) - assert out2.startswith('Error:') - - -def test_non_dict_arguments_returns_error_not_attribute_error(): - tool = ExtractCondensed(Chunks(chunks=[_condensed('cmp', original='orig')])) - # Bypass ToolManager and feed a non-dict directly; must not raise. - out = tool(TOOL_NAME, 'not a dict') # type: ignore[arg-type] - assert out.startswith('Error:') - - -def test_out_of_range_block_returns_short_range_error(): - # Short existence error -- we must NOT enumerate every valid id, or - # a hallucinated ``blocks=[1..200]`` storm would multiply the error - # into thousands of tokens in the non-trainable bridge. - tool = ExtractCondensed( - Chunks(chunks=[ - _condensed('cmp1', original='orig1'), - _condensed('cmp2', original='orig2'), - ])) - out = tool(TOOL_NAME, {'block': 99}) - assert out.startswith('Error:') - assert 'block 99 not found' in out - assert '[1, 2]' in out - # Defensive: the verbose legacy listing must not leak back. - assert 'Available blocks: 1, 2' not in out - - -def test_empty_tool_reports_no_blocks_available(): - tool = ExtractCondensed(Chunks(chunks=[_plain('nothing condensed')])) - out = tool(TOOL_NAME, {'block': 1}) - assert out.startswith('Error:') - assert 'no blocks available' in out - - -def test_integer_strings_are_accepted(): - tool = ExtractCondensed(Chunks(chunks=[_condensed('cmp', original='orig')])) - assert tool(TOOL_NAME, {'block': '1'}) == 'orig' - - -# --------------------------------------------------------------------------- -# single-block-per-call contract + trajectory-bound idempotency -# -# Lists were previously accepted; they are now rejected so a hallucinated -# ``blocks=[1..200]`` cannot flood the non-trainable bridge. Re-requesting -# the same block returns a short "already expanded" reply instead of the -# raw text (which is already sitting in an earlier tool message). -# --------------------------------------------------------------------------- -def test_blocks_int_equivalent_to_legacy_block_arg(): - # Passing ``{'blocks': N}`` (single int under the new name) must - # behave identically to the legacy ``{'block': N}`` path: bare text, - # no <block_N> wrapper. - tool = ExtractCondensed(Chunks(chunks=[_condensed('cmp1', original='orig one')])) - assert tool(TOOL_NAME, {'blocks': 1}) == 'orig one' - # Re-create the tool so the second call is not deduped against the - # first (which is covered separately below). - tool2 = ExtractCondensed(Chunks(chunks=[_condensed('cmp1', original='orig one')])) - assert tool2(TOOL_NAME, {'block': 1}) == 'orig one' - - -def test_blocks_list_is_rejected_with_short_error(): - # Single-block-per-call contract: the only way a list reaches this - # path is if the policy hallucinated a bulk id enumeration, which is - # exactly what we want to stop. Reject loudly with a brief message. - tool = ExtractCondensed( - Chunks(chunks=[ - _condensed('c1', original='a'), - _condensed('c2', original='b'), - _condensed('c3', original='c'), - ])) - for bad in ([1, 2, 3], (1, 2), [1], []): - out = tool(TOOL_NAME, {'blocks': bad}) - assert out.startswith('Error:'), (bad, out) - assert 'single integer' in out or 'one block' in out, (bad, out) - - -def test_second_call_on_same_block_returns_already_expanded_notice(): - # Trajectory-bound idempotency. The raw text has already been handed - # to the model as a prior tool response, so returning it again only - # doubles the non-trainable footprint. The second call gets a short - # notice instead -- no "Error:" prefix (it's not a failure) and - # crucially the raw text must NOT be repeated. - tool = ExtractCondensed( - Chunks(chunks=[ - _condensed('cmp1', original='ORIGINAL TEXT FOR ONE'), - _condensed('cmp2', original='ORIGINAL TEXT FOR TWO'), - ])) - first = tool(TOOL_NAME, {'block': 1}) - assert first == 'ORIGINAL TEXT FOR ONE' - second = tool(TOOL_NAME, {'block': 1}) - assert 'already expanded' in second - assert 'ORIGINAL TEXT FOR ONE' not in second - # Dedup is per-id: a different block is still expandable once. - third = tool(TOOL_NAME, {'block': 2}) - assert third == 'ORIGINAL TEXT FOR TWO' - # And then that one also becomes deduped. - fourth = tool(TOOL_NAME, {'block': 2}) - assert 'already expanded' in fourth - - -def test_already_expanded_is_trajectory_bound_fresh_instance_resets(): - # ``MultiTurnCondenseRollout`` builds a new ExtractCondensed per - # trajectory, so a fresh instance must start with an empty dedup set - # even if a sibling trajectory just expanded block 1. - chunks = Chunks(chunks=[_condensed('c1', original='raw text')]) - t1 = ExtractCondensed(chunks) - assert t1(TOOL_NAME, {'block': 1}) == 'raw text' - assert 'already expanded' in t1(TOOL_NAME, {'block': 1}) - t2 = ExtractCondensed(chunks) # independent trajectory - assert t2(TOOL_NAME, {'block': 1}) == 'raw text' - - -def test_prefers_blocks_over_legacy_block_when_both_present(): - # Undefined which wins in theory; we declare ``blocks`` takes - # precedence so callers can migrate incrementally. - tool = ExtractCondensed(Chunks(chunks=[ - _condensed('c1', original='NEW'), - _condensed('c2', original='LEGACY'), - ])) - out = tool(TOOL_NAME, {'blocks': 1, 'block': 2}) - assert out == 'NEW' - - -# --------------------------------------------------------------------------- -# tool_info metadata -# --------------------------------------------------------------------------- -def test_tool_info_shape_and_serializability(): - tool = ExtractCondensed(Chunks(chunks=[])) - info = tool.tool_info() - # OpenAI-shape: {type: 'function', function: {name, description, parameters}} - assert info['type'] == 'function' - fn = info['function'] - assert fn['name'] == TOOL_NAME == 'extract_condensed' - assert 'description' in fn and fn['description'] - # parameters is a plain mapping (not a JSON string): the jinja chat - # template consumes it directly. - params = fn['parameters'] - assert isinstance(params, dict) - # The whole info dict must still be JSON-serializable so it can be - # embedded inside a trace / logged safely. - json.dumps(info) - # Preferred parameter name is ``blocks`` (single int per call; no list). - assert 'blocks' in params - assert 'int' in params['blocks'] - # The old ``int OR list[int]`` signature must be gone: no list-form - # type annotation leaks through. (The sentence may still say the - # phrase "lists are rejected", which is fine.) - assert 'list[' not in params['blocks'] - assert 'OR list' not in params['blocks'] - - -# --------------------------------------------------------------------------- -# ToolManager integration -# --------------------------------------------------------------------------- -def test_register_with_tool_manager_and_dispatch(): - tool = ExtractCondensed( - Chunks(chunks=[ - _condensed('cmp1', original='orig one'), - _condensed('cmp2', original='orig two'), - ])) - mgr = ToolManager({}) - mgr.register(tool) - assert TOOL_NAME in mgr.names() - - # dict-form arguments - out = mgr({'type': 'function', 'function': {'name': TOOL_NAME, 'arguments': {'block': 2}}}) - assert out == 'orig two' - - # JSON-string-form arguments (OpenAI-style) - out = mgr({'type': 'function', 'function': {'name': TOOL_NAME, 'arguments': '{"block": 1}'}}) - assert out == 'orig one' - - -def test_manager_reports_error_on_unknown_block_without_raising(): - tool = ExtractCondensed(Chunks(chunks=[_condensed('cmp1', original='orig one')])) - mgr = ToolManager({}) - mgr.register(tool) - out = mgr({'type': 'function', 'function': {'name': TOOL_NAME, 'arguments': '{"block": 999}'}}) - assert out.startswith('Error:') - - -# --------------------------------------------------------------------------- -# end-to-end: round-trip with KeywordCondenser (uses raw.original) -# --------------------------------------------------------------------------- -_SPACY_OK = True -try: - import spacy # noqa: F401 - spacy.load('en_core_web_sm') -except Exception: - _SPACY_OK = False - -LONG_PASSAGE = ('Christopher Nolan was born on 30 July 1970 in London. ' - 'He is a British-American film director, producer and screenwriter. ' - 'His film Inception (2010) is a science-fiction heist movie. ' - 'Inception grossed over 829 million dollars worldwide.') - - -@pytest.mark.skipif(not _SPACY_OK, reason='en_core_web_sm not available') -def test_end_to_end_with_keyword_condenser_returns_original(): - from twinkle_agentic.condenser.keyword import KeywordCondenser - - pre = Chunks(chunks=[{'type': 'text', 'role': 'user', 'content': LONG_PASSAGE}]) - post = KeywordCondenser(compression_ratio=4.0, min_chars=50)(pre) - - # The condenser should have left behind an ``original`` snapshot. - assert post.chunks[0]['raw']['condensed'] is True - assert post.chunks[0]['raw']['original'] == LONG_PASSAGE - assert len(post.chunks[0]['content']) < len(LONG_PASSAGE) - - tool = ExtractCondensed(post) - assert tool.blocks == [1] - assert tool(TOOL_NAME, {'block': 1}) == LONG_PASSAGE - - -@pytest.mark.skipif(not _SPACY_OK, reason='en_core_web_sm not available') -def test_end_to_end_block_indices_match_to_trajectory_wrapping(): - from twinkle_agentic.condenser.keyword import KeywordCondenser - - pre = Chunks(chunks=[ - { - 'type': 'text', - 'role': 'user', - 'content': LONG_PASSAGE, - 'round': 1 - }, - { - 'type': 'text', - 'role': 'assistant', - 'content': LONG_PASSAGE + ' Assistant elaboration.', - 'round': 1 - }, - ]) - # skip_roles default excludes assistant → only first chunk condensed. - post = KeywordCondenser(compression_ratio=4.0, min_chars=50)(pre) - tool = ExtractCondensed(post) - - # Exactly one wrapped block. - assert tool.blocks == [1] - # The trajectory wrapper agrees: block_1 exists, block_2 does not. - traj = post.to_trajectory() - rendered = ''.join(m['content'] if isinstance(m.get('content'), str) else '' for m in traj['messages']) - assert '<block_1>' in rendered and '</block_1>' in rendered - assert '<block_2>' not in rendered - # And the tool returns the correct original. - assert tool(TOOL_NAME, {'block': 1}) == LONG_PASSAGE diff --git a/tests/twinkle_agentic/test_harness.py b/tests/twinkle_agentic/test_harness.py new file mode 100644 index 000000000..d1b22571d --- /dev/null +++ b/tests/twinkle_agentic/test_harness.py @@ -0,0 +1,175 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Harness + Env.step_batch + ToolManager.call_many.""" +from __future__ import annotations + +import json +from typing import Any, Dict, List, Optional + +import pytest + +from twinkle.data_format import Trajectory +from twinkle.data_format.sampling import SamplingParams +from twinkle_agentic.envs.base import Env, StepResult +from twinkle_agentic.envs.env_tool import EnvTool +from twinkle_agentic.harness.base import AgentHarness +from twinkle_agentic.rollout.multi_turn import MultiTurnRollout +from twinkle_agentic.tools.tool_manager import ToolManager + +from test_multi_turn_rollout import ( + FakeSampler, + FakeTemplate, + FakeTokenizer, + _tool_call_text, + _user_traj, +) + + +class PrefixHarness(AgentHarness): + """Inject a system message before the first encode; later turns are no-ops.""" + + def before_generate(self, trajectory: Trajectory) -> Trajectory: + msgs = list(trajectory.get('messages') or []) + if not msgs or msgs[0].get('role') != 'system': + trajectory['messages'] = [{'role': 'system', 'content': 'SYS'}] + msgs + return trajectory + + +class TagToolHarness(AgentHarness): + """Prefix every Env observation so we can see after_tools ran.""" + + def after_tools( + self, + trajectory: Trajectory, + observations: List[str], + tool_calls: Optional[List[Dict[str, Any]]] = None, + ) -> Trajectory: + tagged = [f'H:{o}' for o in observations] + return super().after_tools(trajectory, tagged, tool_calls) + + +class BatchEnv(Env): + def __init__(self) -> None: + self.step_calls = 0 + self.batch_calls = 0 + + def step(self, tool_name: str, arguments: Dict[str, Any]) -> StepResult: + self.step_calls += 1 + return StepResult(observation=f'{tool_name}:{json.dumps(arguments, sort_keys=True)}') + + def step_batch(self, calls): + self.batch_calls += 1 + return super().step_batch(calls) + + def tools(self): + return [ + { + 'type': 'function', + 'function': { + 'name': 'search', + 'description': 'search', + 'parameters': {'type': 'object', 'properties': {}}, + }, + }, + { + 'type': 'function', + 'function': { + 'name': 'lookup', + 'description': 'lookup', + 'parameters': {'type': 'object', 'properties': {}}, + }, + }, + ] + + +def _rollout(sampler, template, tool_manager, harness=None, max_turns=4): + return MultiTurnRollout( + sampler=sampler, + template=template, + tool_manager=tool_manager, + sampling_params=SamplingParams(), + max_turns=max_turns, + harness=harness, + ) + + +@pytest.fixture +def tokenizer(): + return FakeTokenizer() + + +@pytest.fixture +def template(tokenizer): + return FakeTemplate(tokenizer) + + +@pytest.fixture +def sampler(template): + return FakeSampler(template) + + +def test_harness_start_default(): + h = AgentHarness() + traj = h.start('hello', user_data=[('id', '"t1"')]) + assert traj['messages'] == [{'role': 'user', 'content': 'hello'}] + assert traj['user_data'] == [('id', '"t1"')] + + +def test_multiturn_harness_injects_system_before_encode(sampler, template): + env = BatchEnv() + mgr = ToolManager(EnvTool.from_env(env)) + sampler.queue('done.', stop_reason='stop') + out = _rollout(sampler, template, mgr, harness=PrefixHarness())([_user_traj('hi')])[0] + roles = [m['role'] for m in out['messages']] + assert roles[0] == 'system' + assert out['messages'][0]['content'] == 'SYS' + assert 'user' in roles + assert 'assistant' in roles + + +def test_multiturn_harness_after_tools_tags_observation(sampler, template): + env = BatchEnv() + mgr = ToolManager(EnvTool.from_env(env)) + sampler.queue(_tool_call_text('search', {'q': 'a'}), stop_reason='stop') + sampler.queue('final', stop_reason='stop') + out = _rollout(sampler, template, mgr, harness=TagToolHarness())([_user_traj('hi')])[0] + tool_msgs = [m for m in out['messages'] if m['role'] == 'tool'] + assert len(tool_msgs) == 1 + assert tool_msgs[0]['content'].startswith('H:') + assert tool_msgs[0].get('name') == 'search' + + +def test_tool_manager_call_many_uses_env_step_batch(): + env = BatchEnv() + mgr = ToolManager(EnvTool.from_env(env)) + calls = [ + {'type': 'function', 'function': {'name': 'search', 'arguments': {'q': 'a'}}}, + {'type': 'function', 'function': {'name': 'lookup', 'arguments': {'k': 'b'}}}, + ] + out = mgr.call_many(calls) + assert env.batch_calls == 1 + assert env.step_calls == 2 + assert out[0].startswith('search:') + assert out[1].startswith('lookup:') + + +def test_ms_agent_harness_start_system_and_user(): + import sys + from pathlib import Path + ms_root = Path(__file__).resolve().parents[2] / 'ms-agent' + if ms_root.is_dir() and str(ms_root) not in sys.path: + sys.path.insert(0, str(ms_root)) + pytest.importorskip('ms_agent') + from twinkle_agentic.harness.ms_agent import MsAgentHarness + + try: + harness = MsAgentHarness(auto_prepare=False) + traj = harness.start('what is 1+1') + except Exception as e: + pytest.skip(f'ms-agent LLMAgent could not start: {e}') + msgs = traj['messages'] + roles = [m['role'] for m in msgs] + assert roles[0] == 'system' + assert roles[-1] == 'user' + assert '1+1' in msgs[-1]['content'] + assert isinstance(msgs[0]['content'], str) + assert len(msgs[0]['content']) > 0 diff --git a/tests/twinkle_agentic/test_keyword_condenser.py b/tests/twinkle_agentic/test_keyword_condenser.py deleted file mode 100644 index 3afb97f49..000000000 --- a/tests/twinkle_agentic/test_keyword_condenser.py +++ /dev/null @@ -1,486 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -import math -import pytest - -from twinkle_agentic.chunker.native import NativeChunker -from twinkle_agentic.condenser.keyword import KeywordCondenser -from twinkle_agentic.data_format import Chunks - -# Module-level skip if spaCy or the small English model are unavailable. -spacy = pytest.importorskip('spacy') -try: - spacy.load('en_core_web_sm') -except OSError: - pytest.skip('en_core_web_sm not available', allow_module_level=True) - -# A realistic multi-sentence passage; long enough to exercise the three -# output slots and the compression budget. -LONG_PASSAGE = ('Christopher Nolan was born on 30 July 1970 in London. ' - 'He is a British-American film director, producer and screenwriter. ' - 'His film Inception (2010) is a science-fiction heist movie starring ' - 'Leonardo DiCaprio. Inception grossed over 829 million dollars worldwide ' - 'and received eight Academy Award nominations, winning four. ' - 'Nolan also directed The Dark Knight trilogy and Interstellar in 2014.') - - -def _user_chunk(text, role='user'): - return {'role': role, 'type': 'text', 'content': text} - - -def _wrap(*chunks): - return Chunks(chunks=list(chunks)) - - -# --------------------------------------------------------------------------- -# constructor validation -# --------------------------------------------------------------------------- -@pytest.mark.parametrize('kw', [ - { - 'num_relations': -1 - }, - { - 'num_keywords': -1 - }, - { - 'max_first_sentence_chars': -1 - }, - { - 'compression_ratio': 1.0 - }, - { - 'compression_ratio': 0.5 - }, - { - 'min_chars': -1 - }, -]) -def test_invalid_config_raises(kw): - with pytest.raises(ValueError): - KeywordCondenser(**kw) - - -# --------------------------------------------------------------------------- -# compression-ratio contract (STRICT upper bound) -# --------------------------------------------------------------------------- -@pytest.mark.parametrize('ratio', [2.0, 3.0, 4.0, 6.0, 10.0]) -def test_compression_ratio_is_strictly_enforced(ratio): - cond = KeywordCondenser( - num_relations=3, max_first_sentence_chars=160, num_keywords=8, compression_ratio=ratio, min_chars=50) - src = _user_chunk(LONG_PASSAGE) - out = cond(_wrap(src)).chunks - assert len(out) == 1 - compressed = out[0]['content'] - budget = math.ceil(len(LONG_PASSAGE) / ratio) - assert len(compressed) <= budget, (f'ratio={ratio}: got len={len(compressed)} > budget={budget}') - assert compressed, 'output must be non-empty' - - -def test_extreme_ratio_keeps_output_non_empty_and_bounded(): - cond = KeywordCondenser(compression_ratio=100.0, min_chars=50) - out = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks - compressed = out[0]['content'] - budget = math.ceil(len(LONG_PASSAGE) / 100.0) - assert 0 < len(compressed) <= budget - - -# --------------------------------------------------------------------------- -# raw.condensed marker + block wrapping -# --------------------------------------------------------------------------- -def test_marks_condensed_and_wraps_in_block_tags(): - cond = KeywordCondenser(compression_ratio=4.0, min_chars=50) - chunks = cond(_wrap(_user_chunk(LONG_PASSAGE))) - assert chunks.chunks[0]['raw']['condensed'] is True - traj = chunks.to_trajectory() - # Exactly one compressed passage → block_1 wrap. - user_content = traj['messages'][0]['content'] - assert '<block_1>' in user_content and '</block_1>' in user_content - - -def test_multiple_chunks_numbered_sequentially_starting_from_1(): - cond = KeywordCondenser(compression_ratio=4.0, min_chars=50) - passages = [_user_chunk(LONG_PASSAGE) for _ in range(3)] - chunks = cond(_wrap(*passages)) - traj = chunks.to_trajectory() - content = traj['messages'][0]['content'] - for i in (1, 2, 3): - assert f'<block_{i}>' in content and f'</block_{i}>' in content - assert '<block_4>' not in content - - -# --------------------------------------------------------------------------- -# slot extraction (opening / relations / keywords) -# --------------------------------------------------------------------------- -def test_opening_relations_keywords_present_when_budget_allows(): - # Generous budget → all three slots should appear. - # LONG_PASSAGE is ~390 chars; full markup is ~370 chars, so we - # need a ratio close to 1.0 to keep every slot. - cond = KeywordCondenser( - num_relations=3, max_first_sentence_chars=160, num_keywords=8, compression_ratio=1.05, min_chars=50) - out = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0]['content'] - assert out.startswith('Open: ') - assert '\nRel: ' in out - assert '\nMore: ' in out - # At least one of the primary entities should survive in keywords. - assert 'Nolan' in out or 'Inception' in out - - -def test_opening_first_sentence_respects_max_chars(): - cond = KeywordCondenser( - num_relations=0, max_first_sentence_chars=20, num_keywords=0, compression_ratio=1.1, min_chars=10) - out = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0]['content'] - # Opening slot is trimmed to <= 20 chars - opening_line = out.split('\n', 1)[0] - assert opening_line.startswith('Open: ') - opening_text = opening_line[len('Open: '):] - assert len(opening_text) <= 20 - - -def test_relations_use_triple_or_quadruple_syntax(): - cond = KeywordCondenser( - num_relations=5, max_first_sentence_chars=10, num_keywords=0, compression_ratio=1.1, min_chars=50) - out = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0]['content'] - # We expect at least one '(a | b | c)' or '(a | b | c | d)' pattern. - assert '(' in out and ')' in out - # Parentheses must balance. - assert out.count('(') == out.count(')') - # Pipe-delimited slots (avoids ',' collision with slot-internal commas). - assert ' | ' in out - - -def test_verb_surface_preserved_not_lemma(): - """Triples keep surface form with auxiliaries: 'was born' not 'bear'.""" - cond = KeywordCondenser( - num_relations=3, max_first_sentence_chars=10, num_keywords=0, compression_ratio=1.1, min_chars=50) - out = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0]['content'] - # Auxiliary preserved. - assert 'was born' in out or 'was released' in out or 'is' in out - # Bare lemma of 'born' must NOT appear as the verb slot. - assert '| bear |' not in out and '| bear on |' not in out - - -def test_internal_hyphens_preserved_in_np(): - """NP text keeps 'science-fiction' / 'British-American' hyphens.""" - cond = KeywordCondenser( - num_relations=5, max_first_sentence_chars=10, num_keywords=0, compression_ratio=1.1, min_chars=50) - out = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0]['content'] - assert 'science-fiction' in out or 'British-American' in out - - -def test_pronoun_subject_triples_skipped(): - """Unresolved pronoun subjects (He/She/It) are noise and dropped.""" - cond = KeywordCondenser( - num_relations=5, max_first_sentence_chars=10, num_keywords=0, compression_ratio=1.1, min_chars=50) - # LONG_PASSAGE has 'He is a British-American film director...' - out = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0]['content'] - assert '(He |' not in out and '(he |' not in out - - -def test_cardinal_entities_filtered_from_keywords(): - cond = KeywordCondenser( - num_relations=0, num_keywords=10, max_first_sentence_chars=0, compression_ratio=1.1, min_chars=50) - passage = ('Alpha earned 100 medals. Beta scored 200 points. Gamma made 300 attempts. ' - 'Delta received 400 votes. Epsilon collected 500 tokens. Zeta passed 600 miles.') - out = cond(_wrap(_user_chunk(passage))).chunks[0]['content'] - for num in ('100', '200', '300', '400', '500', '600'): - assert num not in out, f'pure CARDINAL {num!r} leaked into keywords' - - -def test_keyword_subsumption_prefers_longer_form(): - """'Nolan' is dropped when 'Christopher Nolan' is already kept.""" - cond = KeywordCondenser( - num_relations=0, max_first_sentence_chars=10, num_keywords=8, compression_ratio=1.05, min_chars=50) - out = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0]['content'] - more_line = next((ln for ln in out.splitlines() if ln.startswith('More: ')), '') - kws = [k.strip() for k in more_line[len('More: '):].split(',') if k.strip()] - # No keyword may be a token-subset of another kept keyword. - import re - sets = [frozenset(re.findall(r'\w+', k.lower())) for k in kws] - for i, a in enumerate(sets): - for j, b in enumerate(sets): - if i != j: - assert not a < b, (f'{kws[i]!r} is subsumed by {kws[j]!r} but kept') - - -def test_keyword_exclusion_is_token_level_not_substring(): - """A keyword is only excluded if ALL its words appear in the opening. - - Substring-based exclusion would wrongly drop 'Starfleet' because - 'star' appears inside other tokens; token-level exclusion keeps it. - """ - cond = KeywordCondenser( - num_relations=0, max_first_sentence_chars=60, num_keywords=5, compression_ratio=1.1, min_chars=50) - passage = ('The Starfleet Academy trains officers for deep-space missions. ' - 'Captain Kirk graduated there in 2251. Starfleet operates many vessels.') - out = cond(_wrap(_user_chunk(passage))).chunks[0]['content'] - # 'Starfleet' shouldn't be dropped just because 'star' is a substring - # of something in the opening. - assert 'Starfleet' in out or 'Kirk' in out - - -def test_opening_truncation_at_word_boundary(): - """When opening exceeds max_chars, cut at the last whole word.""" - cond = KeywordCondenser( - num_relations=0, max_first_sentence_chars=25, num_keywords=0, compression_ratio=1.1, min_chars=10) - out = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0]['content'] - opening = out.split('\n', 1)[0][len('Open: '):] - assert len(opening) <= 25 - # Must not end mid-word: last char is a word char AND original passage - # contains the exact trimmed string as a prefix of the first sentence. - first_sent = LONG_PASSAGE.split('.', 1)[0] - assert first_sent.startswith(opening) - # The char after the trimmed prefix in the source should be a space - # (i.e. we really did stop on a word boundary). - if len(opening) < len(first_sent): - assert first_sent[len(opening)] == ' ' - - -def test_budget_is_filled_greedily_with_triples_and_keywords(): - """At a moderate ratio, output should include MORE than just opening. - - Regression test for the old priority-drop logic that collapsed to - opening-only whenever the full composition exceeded budget. - """ - cond = KeywordCondenser( - num_relations=3, max_first_sentence_chars=80, num_keywords=8, compression_ratio=2.0, min_chars=50) - out = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0]['content'] - budget = math.ceil(len(LONG_PASSAGE) / 2.0) - assert len(out) <= budget - # At ratio=2.0 we MUST retain at least one relation AND at least one keyword. - assert '\nRel: ' in out - assert '\nMore: ' in out - - -def test_budget_too_small_falls_back_to_raw_truncation(): - """Even at absurd ratios, output is non-empty and bounded.""" - cond = KeywordCondenser( - num_relations=3, num_keywords=5, max_first_sentence_chars=160, compression_ratio=200.0, min_chars=50) - out = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0]['content'] - budget = math.ceil(len(LONG_PASSAGE) / 200.0) - assert 0 < len(out) <= budget - - -def test_num_relations_zero_suppresses_slot(): - cond = KeywordCondenser(num_relations=0, num_keywords=5, compression_ratio=1.2, min_chars=50) - out = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0]['content'] - assert '\nRel: ' not in out - - -def test_num_keywords_zero_suppresses_slot(): - cond = KeywordCondenser(num_relations=3, num_keywords=0, compression_ratio=1.2, min_chars=50) - out = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0]['content'] - assert '\nMore: ' not in out - - -# --------------------------------------------------------------------------- -# budget priority: drop keywords → drop relations → truncate opening -# --------------------------------------------------------------------------- -def test_tight_budget_drops_keywords_first(): - # Pick a ratio that is just tight enough to force one slot to go. - # Full output len ≈ 200+; opening+relations alone ≈ 120. - cond = KeywordCondenser( - num_relations=2, max_first_sentence_chars=80, num_keywords=8, compression_ratio=3.0, min_chars=50) - out = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0]['content'] - budget = math.ceil(len(LONG_PASSAGE) / 3.0) - assert len(out) <= budget - assert out.startswith('Open: ') - - -def test_very_tight_budget_falls_back_to_opening_only(): - # Ratio large enough that only the opening slot can fit. - # Keep max_first_sentence_chars small so it does fit. - cond = KeywordCondenser( - num_relations=5, max_first_sentence_chars=40, num_keywords=8, compression_ratio=8.0, min_chars=50) - out = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0]['content'] - budget = math.ceil(len(LONG_PASSAGE) / 8.0) - assert len(out) <= budget - # Either opening-only or further truncated — both fine. - assert out.startswith('Open') or len(out) <= budget - - -# --------------------------------------------------------------------------- -# selection policy -# --------------------------------------------------------------------------- -def test_skip_roles_default_preserves_system_tool_assistant(): - cond = KeywordCondenser(compression_ratio=4.0, min_chars=50) - src = _wrap( - _user_chunk(LONG_PASSAGE, role='system'), - _user_chunk(LONG_PASSAGE, role='assistant'), - _user_chunk(LONG_PASSAGE, role='tool'), - _user_chunk(LONG_PASSAGE, role='user'), - ) - out = cond(src).chunks - # First three pass through untouched. - for i in range(3): - assert out[i]['content'] == LONG_PASSAGE - assert (out[i].get('raw') or {}).get('condensed') is not True - # Fourth gets condensed. - assert out[3]['raw']['condensed'] is True - assert len(out[3]['content']) < len(LONG_PASSAGE) - - -def test_custom_skip_roles(): - cond = KeywordCondenser(compression_ratio=4.0, min_chars=50, skip_roles=()) - src = _wrap(_user_chunk(LONG_PASSAGE, role='assistant')) - out = cond(src).chunks - assert out[0]['raw']['condensed'] is True - - -def test_short_content_passes_through(): - cond = KeywordCondenser(compression_ratio=4.0, min_chars=500) - src = _user_chunk(LONG_PASSAGE) # shorter than 500 - out = cond(_wrap(src)).chunks - assert out[0]['content'] == LONG_PASSAGE - assert (out[0].get('raw') or {}).get('condensed') is not True - - -def test_non_text_chunk_passes_through(): - cond = KeywordCondenser(compression_ratio=4.0, min_chars=1) - src = { - 'type': 'image', - 'content': 'http://x/y.png', - 'role': 'user', - 'raw': { - 'type': 'image', - 'image': 'http://x/y.png' - } - } - out = cond(_wrap(src)).chunks - assert out[0] == src - - -def test_reasoning_and_tool_call_kind_chunks_pass_through(): - cond = KeywordCondenser(compression_ratio=4.0, min_chars=50) - reasoning = { - 'type': 'text', - 'role': 'assistant', - 'content': LONG_PASSAGE, - 'raw': { - 'kind': 'reasoning_content' - }, - } - # Assistant role would already be skipped, but the kind-filter must - # hold even if role is user. - tool_call = { - 'type': 'text', - 'role': 'user', - 'content': LONG_PASSAGE, - 'raw': { - 'kind': 'tool_call', - 'tool_call': { - 'type': 'function', - 'function': { - 'name': 'x', - 'arguments': {} - } - } - }, - } - out = cond(_wrap(reasoning, tool_call)).chunks - assert (out[0].get('raw') or {}).get('condensed') is not True - assert (out[1].get('raw') or {}).get('condensed') is not True - - -def test_empty_content_is_untouched(): - cond = KeywordCondenser(compression_ratio=4.0, min_chars=0) - src = _user_chunk('') - out = cond(_wrap(src)).chunks - assert out[0] == src - - -# --------------------------------------------------------------------------- -# integration with NativeChunker + to_trajectory round-trip -# --------------------------------------------------------------------------- -def test_chunker_then_condenser_produces_block_numbered_output(): - chunker = NativeChunker(chunk_size=300) - cond = KeywordCondenser(compression_ratio=4.0, min_chars=50) - - passages = '\n\n'.join(f'[{i}] Title_{i}: ' + LONG_PASSAGE for i in range(1, 4)) - user_text = f'Question: who directed Inception?\n\nContext:\n\n{passages}' - traj = { - 'messages': [ - { - 'role': 'system', - 'content': 'You are a helpful agent.' - }, - { - 'role': 'user', - 'content': user_text - }, - ] - } - chunks = cond(chunker(traj)) - back = chunks.to_trajectory() - - # System untouched; user got multiple condensed blocks. - assert back['messages'][0]['content'] == 'You are a helpful agent.' - user_content = back['messages'][1]['content'] - assert '<block_1>' in user_content - # Each block must be strictly smaller than its source chunk. - assert len(user_content) < len(user_text) - - -def test_condenser_preserves_chunk_order_and_count(): - cond = KeywordCondenser(compression_ratio=4.0, min_chars=50) - src_chunks = _wrap( - _user_chunk('short', role='user'), - _user_chunk(LONG_PASSAGE, role='user'), - _user_chunk(LONG_PASSAGE, role='system'), - ) - out = cond(src_chunks).chunks - assert len(out) == 3 - assert out[0]['content'] == 'short' # too short - assert out[1]['raw']['condensed'] is True # condensed - assert out[2]['content'] == LONG_PASSAGE # skipped role - - -# --------------------------------------------------------------------------- -# idempotency: running condenser twice is safe -# --------------------------------------------------------------------------- -def test_condenser_is_idempotent_on_already_condensed_output(): - cond = KeywordCondenser(compression_ratio=4.0, min_chars=50) - once = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0] - # Second pass must be a no-op: content identical, raw marker kept. - twice = cond(_wrap(once)).chunks[0] - assert twice['raw']['condensed'] is True - assert twice['content'] == once['content'] - # And a third pass must also be stable. - thrice = cond(_wrap(twice)).chunks[0] - assert thrice['content'] == once['content'] - - -# --------------------------------------------------------------------------- -# round-based selection filter -# --------------------------------------------------------------------------- -def _round_chunk(text, round_idx, role='user'): - return {'role': role, 'type': 'text', 'content': text, 'round': round_idx} - - -def test_rounds_filter_only_compresses_first_user_turn(): - cond = KeywordCondenser(compression_ratio=4.0, min_chars=50, rounds=[1]) - out = cond(_wrap( - _round_chunk(LONG_PASSAGE, 1), - _round_chunk(LONG_PASSAGE + ' extra.', 2), - )).chunks - # Round 1 compressed. - assert out[0]['raw']['condensed'] is True - assert len(out[0]['content']) < len(LONG_PASSAGE) - # Round 2 passed through unchanged. - assert out[1]['content'].endswith(' extra.') - assert not (out[1].get('raw') or {}).get('condensed') - - -def test_rounds_filter_excludes_chunks_without_round_field(): - cond = KeywordCondenser(compression_ratio=4.0, min_chars=50, rounds=[1]) - # Chunk missing ``round`` must be treated as non-matching. - plain = _user_chunk(LONG_PASSAGE) - out = cond(_wrap(plain)).chunks[0] - assert out['content'] == LONG_PASSAGE - assert not (out.get('raw') or {}).get('condensed') - - -def test_rounds_filter_default_none_preserves_legacy_behavior(): - cond = KeywordCondenser(compression_ratio=4.0, min_chars=50) - # No rounds set; chunks without ``round`` are still compressed. - out = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0] - assert out['raw']['condensed'] is True - assert len(out['content']) < len(LONG_PASSAGE) diff --git a/tests/twinkle_agentic/test_model_condenser.py b/tests/twinkle_agentic/test_model_condenser.py deleted file mode 100644 index f7f71f56c..000000000 --- a/tests/twinkle_agentic/test_model_condenser.py +++ /dev/null @@ -1,515 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Unit + integration tests for :class:`twinkle_agentic.condenser.model.ModelCondenser`. - -Unit tests use a deterministic mock :class:`Sampler` so the suite runs -without GPUs / vLLM. The final block contains an opt-in integration -test that spins up a real ``Qwen/Qwen2.5-3B-Instruct`` sampler on a -single GPU; enable it with:: - - TWINKLE_TEST_REAL_SAMPLER=1 pytest tests/twinkle_agentic/test_model_condenser.py -""" -from __future__ import annotations - -import math -import os -import pytest -from typing import Callable, List - -# Import directly from the submodule to avoid the (currently broken) -# ``twinkle.sampler.__init__`` import chain in this workspace. -from twinkle.data_format.sampling import SampledSequence, SampleResponse, SamplingParams -from twinkle_agentic.condenser.model import ModelCondenser, _strip_code_fences -from twinkle_agentic.data_format import Chunks - -# --------------------------------------------------------------------------- -# fixtures / helpers -# --------------------------------------------------------------------------- -LONG_PASSAGE = ('Christopher Nolan was born on 30 July 1970 in London. ' - 'He is a British-American film director, producer and screenwriter. ' - 'His film Inception (2010) is a science-fiction heist movie starring ' - 'Leonardo DiCaprio. Inception grossed over 829 million dollars worldwide ' - 'and received eight Academy Award nominations, winning four. ' - 'Nolan also directed The Dark Knight trilogy and Interstellar in 2014.') - - -def _user_chunk(text, role='user'): - return {'role': role, 'type': 'text', 'content': text} - - -def _wrap(*chunks): - return Chunks(chunks=list(chunks)) - - -class _MockSampler: - """Deterministic duck-typed sampler. Calls ``responder(passage)`` per input. - - We do NOT subclass :class:`twinkle.sampler.base.Sampler` to avoid - dragging the workspace's currently-broken template init-chain into - the test module. ``ModelCondenser`` only touches - ``sampler.sample(...)``, so duck-typing is sufficient. - """ - - def __init__(self, responder: Callable[[str], str]): - self._responder = responder - self.template = object() # truthy placeholder, never inspected - self.engine = None - self.calls: list[dict] = [] - - def sample( - self, - inputs, - sampling_params=None, - adapter_name='', - *, - num_samples=1, - **_kw, - ) -> list[SampleResponse]: - inputs_list = inputs if isinstance(inputs, list) else [inputs] - out: list[SampleResponse] = [] - for traj in inputs_list: - user_msg = next(m for m in traj['messages'] if m['role'] == 'user') - prompt = user_msg['content'] - marker = '## Passage\n' - idx = prompt.rfind(marker) - passage = prompt[idx + len(marker):] if idx >= 0 else prompt - decoded = self._responder(passage) - self.calls.append({ - 'passage': passage, - 'sampling_params': sampling_params, - }) - out.append(SampleResponse(sequences=[SampledSequence(stop_reason='stop', tokens=[], decoded=decoded)])) - return out - - -def _well_formed_markdown(passage: str) -> str: - """A standard three-section markdown response.""" - return ('## Summary\n' - 'Christopher Nolan is a British-American director born in London in 1970.\n\n' - '## Key Facts\n' - '- Nolan directed Inception (2010) starring Leonardo DiCaprio.\n' - '- Inception grossed over 829 million dollars worldwide.\n' - '- Nolan also directed The Dark Knight trilogy and Interstellar.\n\n' - '## More\n' - 'Nolan, Inception, Leonardo DiCaprio, Interstellar, London, 1970') - - -# --------------------------------------------------------------------------- -# constructor validation -# --------------------------------------------------------------------------- -def test_requires_sampler(): - with pytest.raises(ValueError): - ModelCondenser(sampler=None) - - -@pytest.mark.parametrize('kw', [ - { - 'compression_ratio': 1.0 - }, - { - 'compression_ratio': 0.5 - }, - { - 'min_chars': -1 - }, - { - 'batch_size': 0 - }, - { - 'user_prompt_template': 'no placeholders' - }, - { - 'user_prompt_template': 'only {budget} placeholder' - }, - { - 'user_prompt_template': 'only {text} placeholder' - }, -]) -def test_invalid_config_raises(kw): - with pytest.raises(ValueError): - ModelCondenser(_MockSampler(_well_formed_markdown), **kw) - - -# --------------------------------------------------------------------------- -# pure helper smoke tests -# --------------------------------------------------------------------------- -def test_strip_code_fences(): - wrapped = '```markdown\n## Summary\nhi\n```' - assert _strip_code_fences(wrapped) == '## Summary\nhi' - # No fence → returned as-is. - plain = '## Summary\nhi' - assert _strip_code_fences(plain) == plain - - -# --------------------------------------------------------------------------- -# compression-vs-passthrough semantics (no hard clamp anymore) -# --------------------------------------------------------------------------- -@pytest.mark.parametrize('ratio', [2.0, 3.0, 4.0, 6.0, 10.0]) -def test_compressed_output_is_strictly_shorter_than_original(ratio): - cond = ModelCondenser( - _MockSampler(_well_formed_markdown), - compression_ratio=ratio, - min_chars=50, - min_budget_chars=1, - ) - chunk = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0] - if chunk.get('raw', {}).get('condensed'): - # When accepted, output MUST be strictly shorter than the input. - assert len( - chunk['content']) < len(LONG_PASSAGE), (f'ratio={ratio}: condensed output len={len(chunk["content"])}' - f' must be < original len={len(LONG_PASSAGE)}') - else: - # Passthrough: chunk must be byte-identical to the input. - assert chunk['content'] == LONG_PASSAGE - - -def test_overlong_model_output_falls_back_to_original(): - """When the LLM output is not strictly shorter than the input, - the original passage is kept verbatim and NOT marked condensed.""" - overflow = lambda _p: _well_formed_markdown('') * 5 # noqa: E731 - cond = ModelCondenser(_MockSampler(overflow), compression_ratio=3.0, min_chars=50, min_budget_chars=1) - chunk = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0] - assert chunk['content'] == LONG_PASSAGE - assert not (chunk.get('raw') or {}).get('condensed') - - -def test_equal_length_model_output_falls_back_to_original(): - """Output equal in length to the input is treated as non-useful - compression and triggers passthrough.""" - same_length = lambda p: 'X' * len(p) # noqa: E731 - cond = ModelCondenser(_MockSampler(same_length), compression_ratio=4.0, min_chars=50, min_budget_chars=1) - chunk = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0] - assert chunk['content'] == LONG_PASSAGE - assert not (chunk.get('raw') or {}).get('condensed') - - -# --------------------------------------------------------------------------- -# structural output quality -# --------------------------------------------------------------------------- -def test_well_formed_output_keeps_three_sections_at_generous_budget(): - cond = ModelCondenser(_MockSampler(_well_formed_markdown), compression_ratio=1.1, min_chars=50) - out = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0]['content'] - assert '## Summary' in out - assert '## Key Facts' in out - assert '## More' in out - # Primary entities survive in some form. - assert 'Nolan' in out or 'Inception' in out - - -def test_tight_ratio_still_accepts_shorter_output(): - """At a tight ratio, whatever the LLM produces is accepted as long - as it is strictly shorter than the input; we no longer clamp it.""" - - def responder(_p): - return ('## Summary\nA short sentence.\n\n' - '## More\nTopics: x, y, z.\n\n' - '## Key Facts\n- Fact one here.\n- Fact two here.') - - cond = ModelCondenser(_MockSampler(responder), compression_ratio=3.5, min_chars=50, min_budget_chars=1) - chunk = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0] - assert chunk['raw']['condensed'] is True - assert len(chunk['content']) < len(LONG_PASSAGE) - assert '## Summary' in chunk['content'] - - -def test_degenerate_output_falls_back_to_original(): - """When model output has NO alphanumerics (pure markdown markers), - the condenser falls back to the original passage verbatim.""" - markers_only = lambda _p: '## \n- \n##' # noqa: E731 - cond = ModelCondenser(_MockSampler(markers_only), compression_ratio=4.0, min_chars=50, min_budget_chars=1) - chunk = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0] - assert chunk['content'] == LONG_PASSAGE - assert not (chunk.get('raw') or {}).get('condensed') - - -def test_garbled_but_shorter_output_is_accepted(): - """If the model emits unstructured but strictly shorter text, we - take it verbatim — the condenser is not a format validator.""" - garbled = lambda _p: 'this is some unstructured blob' # noqa: E731 - cond = ModelCondenser(_MockSampler(garbled), compression_ratio=4.0, min_chars=50, min_budget_chars=1) - chunk = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0] - assert chunk['raw']['condensed'] is True - assert 'unstructured' in chunk['content'] - assert len(chunk['content']) < len(LONG_PASSAGE) - - -def test_code_fenced_output_is_unwrapped(): - wrapped = lambda _p: '```markdown\n' + _well_formed_markdown('') + '\n```' # noqa: E731 - cond = ModelCondenser(_MockSampler(wrapped), compression_ratio=1.5, min_chars=50) - out = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0]['content'] - # After unwrapping, header is at the start (no leading ```). - assert not out.startswith('```') - assert out.startswith('## Summary') - - -# --------------------------------------------------------------------------- -# raw.condensed marker + block wrapping -# --------------------------------------------------------------------------- -def test_marks_condensed_and_wraps_in_block_tags(): - cond = ModelCondenser(_MockSampler(_well_formed_markdown), compression_ratio=4.0, min_chars=50) - chunks = cond(_wrap(_user_chunk(LONG_PASSAGE))) - assert chunks.chunks[0]['raw']['condensed'] is True - traj = chunks.to_trajectory() - user_content = traj['messages'][0]['content'] - assert '<block_1>' in user_content and '</block_1>' in user_content - - -def test_multiple_chunks_numbered_sequentially(): - cond = ModelCondenser(_MockSampler(_well_formed_markdown), compression_ratio=4.0, min_chars=50, batch_size=2) - passages = [_user_chunk(LONG_PASSAGE) for _ in range(3)] - chunks = cond(_wrap(*passages)) - traj = chunks.to_trajectory() - content = traj['messages'][0]['content'] - for i in (1, 2, 3): - assert f'<block_{i}>' in content and f'</block_{i}>' in content - assert '<block_4>' not in content - - -# --------------------------------------------------------------------------- -# selection policy -# --------------------------------------------------------------------------- -def test_skip_roles_default_preserves_system_tool_assistant(): - sampler = _MockSampler(_well_formed_markdown) - cond = ModelCondenser(sampler, compression_ratio=4.0, min_chars=50) - src = _wrap( - _user_chunk(LONG_PASSAGE, role='system'), - _user_chunk(LONG_PASSAGE, role='assistant'), - _user_chunk(LONG_PASSAGE, role='tool'), - _user_chunk(LONG_PASSAGE, role='user'), - ) - out = cond(src).chunks - for i in range(3): - assert out[i]['content'] == LONG_PASSAGE - assert (out[i].get('raw') or {}).get('condensed') is not True - assert out[3]['raw']['condensed'] is True - # Only one real compression job (the user chunk). - assert len(sampler.calls) == 1 - - -def test_custom_skip_roles_empty_tuple(): - cond = ModelCondenser(_MockSampler(_well_formed_markdown), compression_ratio=4.0, min_chars=50, skip_roles=()) - src = _wrap(_user_chunk(LONG_PASSAGE, role='assistant')) - out = cond(src).chunks - assert out[0]['raw']['condensed'] is True - - -def test_short_content_passes_through(): - sampler = _MockSampler(_well_formed_markdown) - cond = ModelCondenser(sampler, compression_ratio=4.0, min_chars=500) - out = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks - assert out[0]['content'] == LONG_PASSAGE - assert (out[0].get('raw') or {}).get('condensed') is not True - assert sampler.calls == [] - - -def test_non_text_chunk_passes_through(): - sampler = _MockSampler(_well_formed_markdown) - cond = ModelCondenser(sampler, compression_ratio=4.0, min_chars=1) - img = { - 'type': 'image', - 'content': 'http://x/y.png', - 'role': 'user', - 'raw': { - 'type': 'image', - 'image': 'http://x/y.png' - } - } - out = cond(_wrap(img)).chunks - assert out[0] == img - assert sampler.calls == [] - - -def test_reasoning_kind_chunk_passes_through(): - sampler = _MockSampler(_well_formed_markdown) - cond = ModelCondenser(sampler, compression_ratio=4.0, min_chars=50) - reasoning = { - 'type': 'text', - 'role': 'user', - 'content': LONG_PASSAGE, - 'raw': { - 'kind': 'reasoning_content' - }, - } - out = cond(_wrap(reasoning)).chunks - assert (out[0].get('raw') or {}).get('condensed') is not True - assert sampler.calls == [] - - -def test_already_condensed_chunk_is_not_reprocessed(): - sampler = _MockSampler(_well_formed_markdown) - cond = ModelCondenser(sampler, compression_ratio=4.0, min_chars=50) - once = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0] - assert once['raw']['condensed'] is True - sampler.calls.clear() - twice = cond(_wrap(once)).chunks[0] - # No second sampler call — idempotent. - assert sampler.calls == [] - assert twice == once - - -# --------------------------------------------------------------------------- -# batching & ordering -# --------------------------------------------------------------------------- -def test_batching_respects_batch_size(): - sampler = _MockSampler(_well_formed_markdown) - cond = ModelCondenser(sampler, compression_ratio=4.0, min_chars=50, batch_size=2) - src = _wrap(*[_user_chunk(LONG_PASSAGE) for _ in range(5)]) - out = cond(src).chunks - assert len(out) == 5 - for c in out: - assert c['raw']['condensed'] is True - # 5 real jobs dispatched in batches of ``batch_size=2``: - # 2 + 2 + 1 = 5 sampler calls total. - assert len(sampler.calls) == 5 - - -def test_order_preserved_with_mixed_chunks(): - sampler = _MockSampler(_well_formed_markdown) - cond = ModelCondenser(sampler, compression_ratio=4.0, min_chars=50, batch_size=2) - src = _wrap( - _user_chunk('short', role='user'), # too short - _user_chunk(LONG_PASSAGE, role='user'), # condensed - _user_chunk(LONG_PASSAGE, role='system'), # skipped role - _user_chunk(LONG_PASSAGE, role='user'), # condensed - ) - out = cond(src).chunks - assert len(out) == 4 - assert out[0]['content'] == 'short' - assert out[1]['raw']['condensed'] is True - assert out[2]['content'] == LONG_PASSAGE - assert (out[2].get('raw') or {}).get('condensed') is not True - assert out[3]['raw']['condensed'] is True - - -# --------------------------------------------------------------------------- -# prompt robustness -# --------------------------------------------------------------------------- -def test_braces_in_text_do_not_break_prompt_formatting(): - sampler = _MockSampler(_well_formed_markdown) - cond = ModelCondenser(sampler, compression_ratio=4.0, min_chars=50) - text = ('The JSON config was {"model": "Qwen", "temperature": 0.7}. ' * 7) - out = cond(_wrap(_user_chunk(text))).chunks[0] - assert out['raw']['condensed'] is True - # Prompt contained the raw text verbatim. - assert sampler.calls[0]['passage'].strip().startswith('The JSON config was {"model":') - - -def test_prompt_mentions_budget_in_user_message(): - sampler = _MockSampler(_well_formed_markdown) - cond = ModelCondenser(sampler, compression_ratio=3.0, min_chars=50) - cond(_wrap(_user_chunk(LONG_PASSAGE))) - expected_budget = math.ceil(len(LONG_PASSAGE) / 3.0) - # The mock recorded the prompt passage; we check the sampling_params - # carries a reasonable max_tokens (derived from budget). - assert sampler.calls[0]['sampling_params'].max_tokens >= expected_budget // 2 - - -def test_custom_sampling_params_is_forwarded(): - sampler = _MockSampler(_well_formed_markdown) - custom = SamplingParams(temperature=0.3, max_tokens=256) - cond = ModelCondenser(sampler, compression_ratio=4.0, min_chars=50, sampling_params=custom) - cond(_wrap(_user_chunk(LONG_PASSAGE))) - assert sampler.calls[0]['sampling_params'] is custom - - -# --------------------------------------------------------------------------- -# semantic preservation (mock-level sanity) -# --------------------------------------------------------------------------- -def test_semantic_preservation_when_compressed(): - """When the condenser accepts the model output, important entities - survive in some form.""" - cond = ModelCondenser(_MockSampler(_well_formed_markdown), compression_ratio=2.0, min_chars=50, min_budget_chars=1) - chunk = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0] - out = chunk['content'] - if chunk.get('raw', {}).get('condensed'): - hits = sum(1 for ent in ('Nolan', 'Inception', 'Leonardo DiCaprio', 'London') if ent in out) - assert hits >= 2 - else: - # Passthrough branch: the original must be returned verbatim. - assert out == LONG_PASSAGE - - -# --------------------------------------------------------------------------- -# integration test (opt-in; requires single GPU + vLLM + Qwen model) -# --------------------------------------------------------------------------- -INTEGRATION_ENABLED = bool(os.environ.get('TWINKLE_TEST_REAL_SAMPLER')) -INTEGRATION_MODEL = os.environ.get('TWINKLE_TEST_MODEL', 'Qwen/Qwen2.5-3B-Instruct') - - -@pytest.mark.skipif( - not INTEGRATION_ENABLED, - reason='Set TWINKLE_TEST_REAL_SAMPLER=1 to run the real-model integration test', -) -def test_integration_real_qwen_sampler_end_to_end(): - """End-to-end test with a real Qwen sampler on a single GPU.""" - vllm = pytest.importorskip('vllm') # noqa: F841 - from twinkle.sampler.vllm_sampler.vllm_sampler import vLLMSampler - - sampler = vLLMSampler( - model_id=INTEGRATION_MODEL, - engine_args={ - 'dtype': 'bfloat16', - 'gpu_memory_utilization': 0.7, - 'max_model_len': 4096, - 'enforce_eager': True, - }, - ) - try: - sampler.set_template('qwen2_5') - except Exception: - # Fall back to 'auto' template detection if the named one - # isn't registered in this build. - sampler.set_template('default') - - cond = ModelCondenser(sampler, compression_ratio=4.0, min_chars=50) - chunk = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0] - out = chunk['content'] - - # Either the model produced a strictly shorter compression (most - # common), or the chunk is passed through verbatim. - if chunk.get('raw', {}).get('condensed'): - assert 0 < len(out) < len(LONG_PASSAGE) - assert any(ent in out for ent in ('Nolan', 'Inception', 'London', 'Leonardo')) - else: - assert out == LONG_PASSAGE - - -# --------------------------------------------------------------------------- -# round-based selection filter -# --------------------------------------------------------------------------- -def _round_chunk(text, round_idx, role='user'): - return {'role': role, 'type': 'text', 'content': text, 'round': round_idx} - - -def test_rounds_filter_only_compresses_first_user_turn(): - sampler = _MockSampler(_well_formed_markdown) - cond = ModelCondenser(sampler, compression_ratio=4.0, min_chars=50, rounds=[1]) - out = cond(_wrap( - _round_chunk(LONG_PASSAGE, 1), - _round_chunk(LONG_PASSAGE + ' extra.', 2), - )).chunks - # One real compression job (round 1). - assert len(sampler.calls) == 1 - # Round 1 compressed. - assert out[0]['raw']['condensed'] is True - # Round 2 untouched. - assert out[1]['content'].endswith(' extra.') - assert not (out[1].get('raw') or {}).get('condensed') - - -def test_rounds_filter_excludes_chunks_without_round_field(): - sampler = _MockSampler(_well_formed_markdown) - cond = ModelCondenser(sampler, compression_ratio=4.0, min_chars=50, rounds=[1]) - out = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0] - # No call because the chunk had no ``round`` field. - assert sampler.calls == [] - assert out['content'] == LONG_PASSAGE - assert not (out.get('raw') or {}).get('condensed') - - -def test_rounds_filter_default_none_preserves_legacy_behavior(): - sampler = _MockSampler(_well_formed_markdown) - cond = ModelCondenser(sampler, compression_ratio=4.0, min_chars=50) - out = cond(_wrap(_user_chunk(LONG_PASSAGE))).chunks[0] - assert out['raw']['condensed'] is True - # One real job. - assert len(sampler.calls) == 1 diff --git a/tests/twinkle_agentic/test_multi_turn_condense_trace.py b/tests/twinkle_agentic/test_multi_turn_condense_trace.py deleted file mode 100644 index 1c4ad1598..000000000 --- a/tests/twinkle_agentic/test_multi_turn_condense_trace.py +++ /dev/null @@ -1,133 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Unit tests for :class:`MultiTurnCondenseRollout` trace augmentation. - -The subclass extends the base trace record with a ``blocks`` field: -``{'block_N': {'original': raw_text_or_None, 'compressed': post_text}}``. -Having both sides of the mapping in the dumped JSON means the trace -alone is enough to audit compression quality. -""" -from __future__ import annotations - -from typing import Any, Dict, List - -from twinkle_agentic.data_format import Chunks -from twinkle_agentic.rollout.multi_turn_condense import MultiTurnCondenseRollout - - -def _chunks(specs: list[dict[str, Any]]) -> Chunks: - out = [] - for s in specs: - raw: dict[str, Any] = {'condensed': bool(s.get('condensed', True))} - if s.get('original') is not None: - raw['original'] = s['original'] - out.append({ - 'type': s.get('type', 'text'), - 'role': s.get('role', 'user'), - 'content': s['content'], - 'raw': raw, - }) - return Chunks(chunks=out) - - -class _Stub(MultiTurnCondenseRollout): - """Bypass ``__init__`` to exercise only ``_build_trace_record``.""" - - def __init__(self, block_chunks): # noqa: D401 -- minimal stub - self._trace_block_chunks = block_chunks - - -def test_build_trace_record_pairs_original_and_compressed(): - chunks = _chunks([ - { - 'content': 'short A', - 'original': 'long raw passage A ...' - }, - { - 'content': 'short B', - 'original': 'long raw passage B ...' - }, - ]) - rollout = _Stub(block_chunks=[chunks]) - traj = {'messages': [], 'stop_reason': 'stop', 'truncated': False} - - record = rollout._build_trace_record(traj, idx=0, success=False) - - assert record['blocks'] == { - 'block_1': { - 'original': 'long raw passage A ...', - 'compressed': 'short A', - }, - 'block_2': { - 'original': 'long raw passage B ...', - 'compressed': 'short B', - }, - } - # Base fields still intact. - assert record['stop_reason'] == 'stop' - - -def test_build_trace_record_preserves_missing_snapshot_as_none(): - """Compressed content is always kept even when ``raw.original`` is None.""" - chunks = _chunks([{'content': 'short A', 'original': None}]) - rollout = _Stub(block_chunks=[chunks]) - record = rollout._build_trace_record({'messages': []}, idx=0, success=False) - assert record['blocks'] == { - 'block_1': { - 'original': None, - 'compressed': 'short A' - }, - } - - -def test_build_trace_record_skips_non_condensed_and_tool_chunks(): - """Numbering only counts condensed, non-tool, non-empty text chunks.""" - chunks = Chunks(chunks=[ - # skipped: not condensed - { - 'type': 'text', - 'role': 'user', - 'content': 'plain', - 'raw': {} - }, - # counted: condensed user text - { - 'type': 'text', - 'role': 'user', - 'content': 'cA', - 'raw': { - 'condensed': True, - 'original': 'rawA' - } - }, - # skipped: tool role - { - 'type': 'text', - 'role': 'tool', - 'content': 'toolmsg', - 'raw': { - 'condensed': True, - 'original': 'xxx' - } - }, - # counted: condensed assistant text - { - 'type': 'text', - 'role': 'assistant', - 'content': 'cB', - 'raw': { - 'condensed': True, - 'original': 'rawB' - } - }, - ]) - rollout = _Stub(block_chunks=[chunks]) - record = rollout._build_trace_record({'messages': []}, idx=0, success=False) - assert list(record['blocks']) == ['block_1', 'block_2'] - assert record['blocks']['block_1']['original'] == 'rawA' - assert record['blocks']['block_2']['original'] == 'rawB' - - -def test_build_trace_record_is_noop_when_stash_missing(): - rollout = _Stub(block_chunks=None) - record = rollout._build_trace_record({'messages': []}, idx=0, success=False) - assert 'blocks' not in record diff --git a/tests/twinkle_agentic/test_native_chunker.py b/tests/twinkle_agentic/test_native_chunker.py deleted file mode 100644 index 915d0d4bd..000000000 --- a/tests/twinkle_agentic/test_native_chunker.py +++ /dev/null @@ -1,555 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Unit tests for :class:`twinkle_agentic.chunker.native.NativeChunker`. - -Focus: chunk-size boundaries, separator priority, first-user-only scope, -lossless ``''.join`` of split outputs, and edge cases (empty, multimodal, -tool-calls, invalid config). -""" -from __future__ import annotations - -import pytest - -from twinkle_agentic.chunker.native import NativeChunker, _hard_cut, _split_keep -from twinkle_agentic.data_format import Chunks - - -def _u(content, role='user'): - return {'role': role, 'content': content} - - -def _join(chunks, type_='text'): - return ''.join(c['content'] for c in chunks if c.get('type') == type_) - - -# --------------------------------------------------------------------------- -# chunk_size boundaries -# --------------------------------------------------------------------------- -def test_under_chunk_size_returns_single_chunk(): - ch = NativeChunker(chunk_size=100) - out = ch({'messages': [_u('hello world')]}).chunks - assert len(out) == 1 - assert out[0]['content'] == 'hello world' - assert out[0]['role'] == 'user' - assert out[0]['type'] == 'text' - - -def test_exact_chunk_size_not_split(): - ch = NativeChunker(chunk_size=10) - out = ch({'messages': [_u('a' * 10)]}).chunks - assert [c['content'] for c in out] == ['a' * 10] - - -def test_one_over_chunk_size_is_split(): - ch = NativeChunker(chunk_size=10) - out = ch({'messages': [_u('a' * 11)]}).chunks - # No separator matches → hard cut; merge won't fuse (10+1 > 10) - assert len(out) == 2 - assert all(len(c['content']) <= 10 for c in out) - assert _join(out) == 'a' * 11 - - -def test_all_chunks_respect_size_limit_on_realistic_input(): - ch = NativeChunker(chunk_size=20) - text = ('hello world. ' * 50).strip() - out = ch({'messages': [_u(text)]}).chunks - assert all(len(c['content']) <= 20 for c in out) - assert _join(out) == text - - -def test_large_text_split_is_lossless_and_bounded(): - ch = NativeChunker(chunk_size=64) - text = 'The quick brown fox jumps over the lazy dog. ' * 100 - out = ch({'messages': [_u(text)]}).chunks - assert _join(out) == text - assert all(len(c['content']) <= 64 for c in out) - - -# --------------------------------------------------------------------------- -# separator priority (coarsest available wins) -# --------------------------------------------------------------------------- -def test_paragraph_split_preferred_over_sentence(): - ch = NativeChunker(chunk_size=40) - text = 'P1 sentence one. P1 sentence two.\n\nP2 sentence one. P2 sentence two.' - out = ch({'messages': [_u(text)]}).chunks - assert _join(out) == text - assert all(len(c['content']) <= 40 for c in out) - # Because paragraph boundary (18 + 2) and (35) both fit in 40, we - # expect at most 2 chunks (one per paragraph, possibly merged). - assert len(out) <= 2 - - -def test_newline_split_used_when_no_paragraph(): - ch = NativeChunker(chunk_size=10) - text = 'line1\nline2\nline3\nline4' - out = ch({'messages': [_u(text)]}).chunks - assert _join(out) == text - assert all(len(c['content']) <= 10 for c in out) - - -def test_sentence_split_used_when_no_newline(): - ch = NativeChunker(chunk_size=10) - text = 'foo bar b. qux qa bc. abc d.' - out = ch({'messages': [_u(text)]}).chunks - assert _join(out) == text - assert all(len(c['content']) <= 10 for c in out) - - -def test_chinese_sentence_separator(): - ch = NativeChunker(chunk_size=8) - text = '你好世界。这是测试。再见朋友。' - out = ch({'messages': [_u(text)]}).chunks - assert _join(out) == text - assert all(len(c['content']) <= 8 for c in out) - - -def test_custom_separator_list_only(): - ch = NativeChunker(chunk_size=10, separators=['|']) - text = 'aaa|bbb|ccccccccc|dd' - out = ch({'messages': [_u(text)]}).chunks - assert _join(out) == text - assert all(len(c['content']) <= 10 for c in out) - - -def test_empty_string_sentinel_appended_automatically(): - # User omits '' → chunker must still make progress on unsplittable text - ch = NativeChunker(chunk_size=3, separators=['|']) - text = 'abcdefghij' # no '|' at all - out = ch({'messages': [_u(text)]}).chunks - assert _join(out) == text - assert all(len(c['content']) <= 3 for c in out) - - -# --------------------------------------------------------------------------- -# first-user-only constraint -# --------------------------------------------------------------------------- -def test_only_first_user_message_is_split(): - ch = NativeChunker(chunk_size=10) - long = 'a' * 100 - traj = { - 'messages': [ - { - 'role': 'system', - 'content': long - }, - { - 'role': 'user', - 'content': long - }, # ← split - { - 'role': 'assistant', - 'content': long - }, - { - 'role': 'user', - 'content': long - }, # ← pass-through - { - 'role': 'tool', - 'content': long, - 'tool_call_id': 'c1' - }, - ] - } - out = ch(traj).chunks - - # Count chunks per message by position. - system_chunks = [c for c in out if c['role'] == 'system'] - assistant_chunks = [c for c in out if c['role'] == 'assistant'] - tool_chunks = [c for c in out if c['role'] == 'tool'] - user_chunks = [c for c in out if c['role'] == 'user'] - - assert len(system_chunks) == 1 - assert len(assistant_chunks) == 1 - assert len(tool_chunks) == 1 - # First user is split into many + second user pass-through (1 chunk). - assert len(user_chunks) > 2 - # And the second user chunk sits at the end of the user_chunks group - # only after the first-user splits. - assert user_chunks[-1]['content'] == long - - -def test_system_and_assistant_content_not_split(): - ch = NativeChunker(chunk_size=5) - long = 'abcdefghijklmn' - traj = { - 'messages': [ - { - 'role': 'system', - 'content': long - }, - { - 'role': 'assistant', - 'content': long - }, - ] - } - out = ch(traj).chunks - assert len(out) == 2 - assert out[0]['content'] == long - assert out[1]['content'] == long - - -def test_trajectory_without_user_message_produces_no_split(): - ch = NativeChunker(chunk_size=5) - long = 'abcdefghij' - traj = { - 'messages': [ - { - 'role': 'system', - 'content': long - }, - { - 'role': 'assistant', - 'content': long - }, - ] - } - out = ch(traj).chunks - assert all(len(c['content']) == len(long) for c in out) - - -# --------------------------------------------------------------------------- -# decomposition of special message parts -# --------------------------------------------------------------------------- -def test_reasoning_content_becomes_own_chunk(): - ch = NativeChunker(chunk_size=100) - traj = { - 'messages': [ - _u('hi'), - { - 'role': 'assistant', - 'reasoning_content': 'think step', - 'content': 'answer' - }, - ] - } - out = ch(traj).chunks - # user(hi) + assistant.reasoning + assistant.content - assert len(out) == 3 - assert out[1]['raw']['kind'] == 'reasoning_content' - assert out[1]['content'] == 'think step' - assert out[2]['content'] == 'answer' - assert 'raw' not in out[2] or 'kind' not in out[2].get('raw', {}) - - -def test_tool_calls_become_empty_text_chunks_with_kind(): - ch = NativeChunker(chunk_size=100) - traj = { - 'messages': [ - _u('hi'), - { - 'role': - 'assistant', - 'content': - 'calling', - 'tool_calls': [ - { - 'type': 'function', - 'function': { - 'name': 'foo', - 'arguments': {} - } - }, - { - 'type': 'function', - 'function': { - 'name': 'bar', - 'arguments': { - 'x': 1 - } - } - }, - ] - }, - ] - } - out = ch(traj).chunks - tc_chunks = [c for c in out if c.get('raw', {}).get('kind') == 'tool_call'] - assert len(tc_chunks) == 2 - assert tc_chunks[0]['raw']['tool_call']['function']['name'] == 'foo' - assert tc_chunks[1]['raw']['tool_call']['function']['name'] == 'bar' - # Empty content on tool_call chunks. - assert all(c['content'] == '' for c in tc_chunks) - - -def test_tool_message_preserves_tool_call_id(): - ch = NativeChunker(chunk_size=100) - traj = { - 'messages': [ - _u('hi'), - { - 'role': 'tool', - 'content': 'result', - 'tool_call_id': 'call-42' - }, - ] - } - out = ch(traj).chunks - tool_chunk = out[-1] - assert tool_chunk['role'] == 'tool' - assert tool_chunk['raw']['tool_call_id'] == 'call-42' - - -def test_multimodal_content_preserved_on_first_user(): - ch = NativeChunker(chunk_size=5) - traj = { - 'messages': [{ - 'role': - 'user', - 'content': [ - { - 'type': 'text', - 'text': 'describe this image' - }, - { - 'type': 'image', - 'image': 'http://x/y.png' - }, - ], - }] - } - out = ch(traj).chunks - text_chunks = [c for c in out if c['type'] == 'text'] - image_chunks = [c for c in out if c['type'] == 'image'] - assert len(image_chunks) == 1 - assert image_chunks[0]['content'] == 'http://x/y.png' - assert image_chunks[0]['raw'] == {'type': 'image', 'image': 'http://x/y.png'} - # Text part was split; concatenation is lossless. - assert _join(text_chunks) == 'describe this image' - assert all(len(c['content']) <= 5 for c in text_chunks) - - -# --------------------------------------------------------------------------- -# edge cases -# --------------------------------------------------------------------------- -def test_empty_trajectory(): - ch = NativeChunker(chunk_size=10) - assert ch({'messages': []}).chunks == [] - assert ch({}).chunks == [] - - -def test_empty_content_string_produces_no_chunks(): - ch = NativeChunker(chunk_size=10) - assert ch({'messages': [_u('')]}).chunks == [] - - -@pytest.mark.parametrize('bad', [0, -1, -999]) -def test_invalid_chunk_size_raises(bad): - with pytest.raises(ValueError): - NativeChunker(chunk_size=bad) - - -def test_chunk_size_one_hard_cuts_all_chars(): - ch = NativeChunker(chunk_size=1) - text = 'abc' - out = ch({'messages': [_u(text)]}).chunks - assert [c['content'] for c in out] == ['a', 'b', 'c'] - - -def test_whitespace_only_text_is_preserved_losslessly(): - ch = NativeChunker(chunk_size=3) - text = ' \n\n \n' - out = ch({'messages': [_u(text)]}).chunks - assert _join(out) == text - assert all(len(c['content']) <= 3 for c in out) - - -# --------------------------------------------------------------------------- -# HotpotQA-shaped realistic payload -# --------------------------------------------------------------------------- -def test_hotpotqa_like_passage_layout(): - ch = NativeChunker(chunk_size=80) - passages = '\n\n'.join(f'[{i}] Title_{i}: ' + 'This is sentence. ' * 6 for i in range(1, 6)) - user_text = f'Question: who wrote it?\n\nContext:\n\n{passages}' - out = ch({ - 'messages': [ - { - 'role': 'system', - 'content': 'sys' - }, - _u(user_text), - ] - }).chunks - # System message is not split. - assert out[0]['role'] == 'system' and out[0]['content'] == 'sys' - # User text reconstructs losslessly. - user_chunks = [c for c in out if c['role'] == 'user'] - assert _join(user_chunks) == user_text - assert all(len(c['content']) <= 80 for c in user_chunks) - - -# --------------------------------------------------------------------------- -# to_trajectory integration (non-split messages round-trip cleanly) -# --------------------------------------------------------------------------- -def test_non_split_messages_roundtrip_through_to_trajectory(): - ch = NativeChunker(chunk_size=1024) - tc = {'type': 'function', 'function': {'name': 'foo', 'arguments': {}}} - traj = { - 'messages': [ - { - 'role': 'system', - 'content': 'sys' - }, - { - 'role': 'user', - 'content': 'short question' - }, - { - 'role': 'assistant', - 'content': 'answer', - 'tool_calls': [tc] - }, - { - 'role': 'tool', - 'content': 'result', - 'tool_call_id': 'c1' - }, - ] - } - chunks = ch(traj) - back = chunks.to_trajectory(block_wrapper=None) - msgs = back['messages'] - assert msgs[0] == {'role': 'system', 'content': 'sys'} - assert msgs[1]['role'] == 'user' - assert msgs[1]['content'] == 'short question' - assert msgs[2]['role'] == 'assistant' - assert msgs[2]['content'] == 'answer' - assert msgs[2]['tool_calls'] == [tc] - assert msgs[3]['role'] == 'tool' - assert msgs[3]['content'] == 'result' - assert msgs[3]['tool_call_id'] == 'c1' - - -# --------------------------------------------------------------------------- -# helper-level tests (white-box, catches regressions in primitives) -# --------------------------------------------------------------------------- -def test_split_keep_is_lossless(): - cases = [ - ('', '|'), - ('abc', '|'), - ('a|b|c', '|'), - ('|abc|', '|'), - ('|||', '|'), - ('aa..bb.', '.'), - ('hello', ''), # empty separator → single piece - ] - for text, sep in cases: - parts = _split_keep(text, sep) - assert ''.join(parts) == text, (text, sep, parts) - - -def test_hard_cut_bounds_and_lossless(): - for text, size in [('', 3), ('a', 3), ('abcde', 3), ('abcdef', 3)]: - parts = _hard_cut(text, size) - assert ''.join(parts) == text - assert all(len(p) <= size for p in parts) - - -def test_split_keep_keeps_separator_suffix(): - assert _split_keep('aa.bb.cc', '.') == ['aa.', 'bb.', 'cc'] - assert _split_keep('aa\n\nbb\n\ncc', '\n\n') == ['aa\n\n', 'bb\n\n', 'cc'] - - -# --------------------------------------------------------------------------- -# separator ordering / priority contract -# --------------------------------------------------------------------------- -def test_prefers_paragraph_boundary_over_period_when_both_fit(): - # Two paragraphs. Each fits in 40. The whole thing (47) does not. - ch = NativeChunker(chunk_size=40) - text = 'para one sentence. more.\n\npara two sentence.' - assert len(text) > 40 - out = ch({'messages': [_u(text)]}).chunks - # Chunker should split at '\n\n', not inside a paragraph. - assert out[0]['content'].endswith('\n\n') - assert _join(out) == text - - -# --------------------------------------------------------------------------- -# round numbering -# --------------------------------------------------------------------------- -def test_round_starts_at_zero_for_pre_user_system(): - ch = NativeChunker(chunk_size=1024) - out = ch({ - 'messages': [ - { - 'role': 'system', - 'content': 'you are helpful' - }, - _u('hello'), - ] - }).chunks - assert [c['round'] for c in out] == [0, 1] - - -def test_round_increments_on_each_user_message(): - ch = NativeChunker(chunk_size=1024) - out = ch({ - 'messages': [ - _u('first user'), - { - 'role': 'assistant', - 'content': 'first reply' - }, - _u('second user'), - { - 'role': 'assistant', - 'content': 'second reply' - }, - _u('third user'), - ] - }).chunks - rounds = [c['round'] for c in out] - # assistant msgs inherit the round of the preceding user turn. - assert rounds == [1, 1, 2, 2, 3] - - -def test_round_covers_tool_responses_between_users(): - ch = NativeChunker(chunk_size=1024) - out = ch({ - 'messages': [ - _u('query'), - { - 'role': 'assistant', - 'content': 'calling tool' - }, - { - 'role': 'tool', - 'content': 'tool result', - 'tool_call_id': 'x' - }, - { - 'role': 'assistant', - 'content': 'final' - }, - ] - }).chunks - assert {c['round'] for c in out} == {1} - - -def test_round_preserved_when_first_user_is_split(): - ch = NativeChunker(chunk_size=20) - long_user = 'hello world. ' * 10 # gets split - out = ch({ - 'messages': [ - { - 'role': 'system', - 'content': 'sys' - }, - _u(long_user), - { - 'role': 'assistant', - 'content': 'ack' - }, - _u('again'), - ] - }).chunks - # All pieces of the split first user share round=1, system is round=0, - # assistant inherits round=1, second user is round=2. - by_role = {} - for c in out: - by_role.setdefault(c.get('role'), []).append(c['round']) - assert set(by_role.get('system', [])) == {0} - assert set(by_role.get('assistant', [])) == {1} - # Multiple user chunks from the split share round=1. - assert by_role['user'].count(1) >= 2 - assert by_role['user'][-1] == 2 diff --git a/tests/twinkle_agentic/test_repeated_calls_spin.py b/tests/twinkle_agentic/test_repeated_calls_spin.py deleted file mode 100644 index c0b80195e..000000000 --- a/tests/twinkle_agentic/test_repeated_calls_spin.py +++ /dev/null @@ -1,47 +0,0 @@ -"""check_no_repeated_calls: single-tool spin loops are penalized, batches aren't. - -Regression for a dead-loop that exact-duplicate detection missed: the same tool -fired ~20 times with *different* arguments (empty repeated spins) used to score -1.0 because no two (name, args) pairs were identical. -""" -import json - -from twinkle_agentic.verifier.hard_scorer import (TrajectoryView, - check_no_repeated_calls) - - -def _call(name, args): - return {'role': 'assistant', - 'tool_calls': [{'function': {'name': name, 'arguments': json.dumps(args)}}]} - - -def _score(msgs): - return check_no_repeated_calls(TrajectoryView({'messages': msgs})).score - - -def test_single_tool_spin_is_penalized(): - loop = [_call('LatexFixResponse', {'part': str(i)}) for i in range(18)] - assert _score(loop) <= 0.4 - - -def test_exact_duplicate_calls_penalized(): - dupes = [_call('read', {'p': 'same'}) for _ in range(4)] - # 3 of 4 are exact duplicates -> 1 - 3/4 = 0.25 - assert _score(dupes) <= 0.3 - - -def test_mixed_tools_not_penalized(): - mixed = [_call('read', {'p': str(i)}) for i in range(6)] - mixed += [_call('grep', {'q': 'a'}), _call('edit', {'f': 'b'}), - _call('run', {'c': 'c'}), _call('read', {'p': 'z'})] - assert _score(mixed) == 1.0 - - -def test_short_same_tool_loop_ok(): - # A legitimate 4-step same-tool loop (below the spin floor of 8 calls). - small = [_call('read', {'p': str(i)}) for i in range(4)] - assert _score(small) == 1.0 - - -def test_fewer_than_two_calls_ok(): - assert _score([_call('read', {'p': 'a'})]) == 1.0 diff --git a/tests/twinkle_agentic/test_rubric_stabilization.py b/tests/twinkle_agentic/test_rubric_stabilization.py deleted file mode 100644 index 0d2e36484..000000000 --- a/tests/twinkle_agentic/test_rubric_stabilization.py +++ /dev/null @@ -1,92 +0,0 @@ -"""Rubric stabilization: skeleton prepend, intent routing, high-band voting. - -These exercise :class:`RubricVerifier` without any real LLM by stubbing the two -distilled hooks (``_gen_rubric`` / ``_score_once``) and forcing -``_llm_available`` True, so we test the *assembly + voting* logic in isolation. -""" -from twinkle_agentic.preprocessor.intents import INTENT_CODE, INTENT_TOOL_CALL -from twinkle_agentic.verifier import RubricItem, RubricVerifier -from twinkle_agentic.verifier.rubric_library import default_intent_base_rubrics - - -def _segment(): - return {'messages': [ - {'role': 'user', 'content': 'do the thing'}, - {'role': 'assistant', 'content': 'here is the result'}, - ]} - - -def _stub_llm(rv, *, gen_lines, score_seq): - """Force LLM-available and deterministic gen/score outputs. - - ``score_seq`` is a list of raw verdict strings returned by successive - ``_score_once`` calls (so we can count how many votes were spent). - """ - rv._llm_available = lambda: True # type: ignore[method-assign] - rv._gen_rubric = lambda **kw: gen_lines # type: ignore[method-assign] - calls = {'n': 0} - - def _score_once(**kw): - i = min(calls['n'], len(score_seq) - 1) - calls['n'] += 1 - return score_seq[i] - - rv._score_once = _score_once # type: ignore[method-assign] - return calls - - -def test_base_rubric_prepended_and_dedup(): - base = [RubricItem('The agent calls tools with valid JSON', True)] - rv = RubricVerifier(base_rubric=base, min_rubrics=1, max_rubrics=6) - gen = ('1. The agent calls tools with valid JSON [Hard Rule]\n' # dup of skeleton - '2. The response advances the sub-goal [Principle]') - _stub_llm(rv, gen_lines=gen, score_seq=['1: PASS\n2: PASS']) - detail = rv.score_detail(_segment()) - texts = [it.text for it in detail.rubric] - # skeleton first, duplicate from generation dropped -> exactly 2 items - assert texts[0] == 'The agent calls tools with valid JSON' - assert len(detail.rubric) == 2 - - -def test_intent_fixed_rubric_skips_generation(): - fixed = {INTENT_TOOL_CALL: [RubricItem('The agent uses tools correctly', True)]} - rv = RubricVerifier(intent_rubrics=fixed) - # gen would raise if called -> proves generation is skipped for this intent - rv._llm_available = lambda: True # type: ignore[method-assign] - rv._gen_rubric = lambda **kw: (_ for _ in ()).throw(AssertionError('gen called')) # type: ignore - rv._score_once = lambda **kw: '1: PASS' # type: ignore[method-assign] - detail = rv.score_detail(_segment(), intent=INTENT_TOOL_CALL) - assert len(detail.rubric) == 1 - assert detail.rubric[0].text == 'The agent uses tools correctly' - - -def test_intent_base_rubric_routes_by_intent(): - rv = RubricVerifier(intent_base_rubrics=default_intent_base_rubrics(), - min_rubrics=1, max_rubrics=8) - _stub_llm(rv, gen_lines='1. The response is coherent [Principle]', - score_seq=['1: PASS\n2: PASS\n3: PASS\n4: PASS']) - detail = rv.score_detail(_segment(), intent=INTENT_CODE) - # the CODE skeleton leads the rubric - assert detail.rubric[0].text.startswith('The response produces code') - - -def test_high_band_forces_more_votes(): - rv = RubricVerifier(min_rubrics=1, max_rubrics=4, min_votes_high=3, - high_score_threshold=0.85, max_votes=5) - gen = '1. The response is correct [Hard Rule]' - # first pass all-PASS -> scalar 1.0 (>= 0.85) -> must escalate to >=3 votes - calls = _stub_llm(rv, gen_lines=gen, score_seq=['1: PASS']) - detail = rv.score_detail(_segment()) - assert detail.n_votes >= 3 - assert calls['n'] >= 3 - - -def test_low_band_single_vote(): - rv = RubricVerifier(min_rubrics=1, max_rubrics=4, min_votes_high=3, - high_score_threshold=0.85, max_votes=5, margin_threshold=0.1) - gen = '1. The response is correct [Hard Rule]' - # first pass FAIL -> scalar 0.0, decisive (far from 0.5) -> single vote - calls = _stub_llm(rv, gen_lines=gen, score_seq=['1: FAIL']) - detail = rv.score_detail(_segment()) - assert detail.n_votes == 1 - assert calls['n'] == 1 From fe9775a7b192ffff60c257a05e648835d5d78c60 Mon Sep 17 00:00:00 2001 From: root <yuze.zyz@alibaba-inc.com> Date: Sat, 22 Aug 2026 00:03:58 +0800 Subject: [PATCH 42/60] wip --- cookbook/rl/envs/agentenv_server/install.sh | 9 + cookbook/rl/rsi_agentic/README.md | 148 +++ cookbook/rl/rsi_agentic/remote_tool_env.py | 361 ++++++ cookbook/rl/rsi_agentic/rsi_agent.yaml | 41 +- cookbook/rl/rsi_agentic/rsi_agentic_grpo.py | 216 +++- .../rl/rsi_agentic/sandbox_server/Dockerfile | 24 + .../rl/rsi_agentic/sandbox_server/install.sh | 74 ++ .../rl/rsi_agentic/sandbox_server/serve.sh | 20 + .../rsi_agentic/sandbox_server/tool_server.py | 262 +++++ cookbook/rsi/agentic/README.md | 54 + cookbook/rsi/agentic/challenge.py | 302 +++++ cookbook/rsi/agentic/prompts.py | 171 +++ cookbook/rsi/code/challenge.py | 260 ++++ cookbook/rsi/code/prompts.py | 170 +++ cookbook/rsi/prepare.py | 184 +++ .../rsi/rsi_rl.py => cookbook/rsi/rl.py | 23 +- cookbook/rsi/run_rsi.py | 143 ++- cookbook/rsi/run_rsi_selfplay.py | 125 -- src/twinkle_agentic/challenger/__init__.py | 25 + src/twinkle_agentic/challenger/agentic.py | 559 +++++++++ src/twinkle_agentic/challenger/base.py | 379 ++++++ src/twinkle_agentic/challenger/code.py | 826 +++++++++++++ src/twinkle_agentic/envs/__init__.py | 1 - src/twinkle_agentic/envs/ms_agent_tool_env.py | 212 ---- src/twinkle_agentic/rollout/__init__.py | 2 + src/twinkle_agentic/rollout/api_multi_turn.py | 112 +- src/twinkle_agentic/rollout/base.py | 196 +++- src/twinkle_agentic/rollout/factory.py | 74 ++ src/twinkle_agentic/rollout/multi_turn.py | 168 +-- src/twinkle_agentic/rsi/__init__.py | 9 - src/twinkle_agentic/rsi/rsi_challenge.py | 1044 ----------------- src/twinkle_agentic/rsi/rsi_distill.py | 285 ----- src/twinkle_agentic/rsi/rsi_prepare.py | 179 --- src/twinkle_agentic/rsi/rsi_refine.py | 277 ----- src/twinkle_agentic/sampler/__init__.py | 2 - src/twinkle_agentic/sampler/router_sampler.py | 197 ---- tests/twinkle_agentic/test_agentic_rsi.py | 265 +++-- 37 files changed, 4617 insertions(+), 2782 deletions(-) create mode 100644 cookbook/rl/rsi_agentic/README.md create mode 100644 cookbook/rl/rsi_agentic/remote_tool_env.py create mode 100644 cookbook/rl/rsi_agentic/sandbox_server/Dockerfile create mode 100644 cookbook/rl/rsi_agentic/sandbox_server/install.sh create mode 100644 cookbook/rl/rsi_agentic/sandbox_server/serve.sh create mode 100644 cookbook/rl/rsi_agentic/sandbox_server/tool_server.py create mode 100644 cookbook/rsi/agentic/README.md create mode 100644 cookbook/rsi/agentic/challenge.py create mode 100644 cookbook/rsi/agentic/prompts.py create mode 100644 cookbook/rsi/code/challenge.py create mode 100644 cookbook/rsi/code/prompts.py create mode 100644 cookbook/rsi/prepare.py rename src/twinkle_agentic/rsi/rsi_rl.py => cookbook/rsi/rl.py (98%) delete mode 100644 cookbook/rsi/run_rsi_selfplay.py create mode 100644 src/twinkle_agentic/challenger/__init__.py create mode 100644 src/twinkle_agentic/challenger/agentic.py create mode 100644 src/twinkle_agentic/challenger/base.py create mode 100644 src/twinkle_agentic/challenger/code.py delete mode 100644 src/twinkle_agentic/envs/ms_agent_tool_env.py create mode 100644 src/twinkle_agentic/rollout/factory.py delete mode 100644 src/twinkle_agentic/rsi/__init__.py delete mode 100644 src/twinkle_agentic/rsi/rsi_challenge.py delete mode 100644 src/twinkle_agentic/rsi/rsi_distill.py delete mode 100644 src/twinkle_agentic/rsi/rsi_prepare.py delete mode 100644 src/twinkle_agentic/rsi/rsi_refine.py delete mode 100644 src/twinkle_agentic/sampler/__init__.py delete mode 100644 src/twinkle_agentic/sampler/router_sampler.py diff --git a/cookbook/rl/envs/agentenv_server/install.sh b/cookbook/rl/envs/agentenv_server/install.sh index f611f9522..68ac402da 100644 --- a/cookbook/rl/envs/agentenv_server/install.sh +++ b/cookbook/rl/envs/agentenv_server/install.sh @@ -10,10 +10,14 @@ REPO_ROOT="${REPO_ROOT:-$HOME/AgentENV}" CONFIG_DIR="${CONFIG_DIR:-/var/lib/aenv/config}" SKIP_INSTALL=0 +SKIP_BUILD=0 REBUILD=0 for arg in "$@"; do case "$arg" in --skip-install) SKIP_INSTALL=1 ;; + # Bootstrap the host but build no template: used by cookbook setups that + # bring their own Dockerfile and only need the server installed once. + --skip-build) SKIP_BUILD=1 ;; --rebuild) REBUILD=1 ;; *) echo "Unknown option: $arg" >&2; exit 2 ;; esac @@ -49,6 +53,11 @@ else aenv auth fi +if [ "$SKIP_BUILD" = "1" ]; then + echo "==> Skipping template build (--skip-build)" + exit 0 +fi + if [ "$REBUILD" = "1" ]; then echo "==> Deleting template '$TEMPLATE'" aenv template delete "$TEMPLATE" || true diff --git a/cookbook/rl/rsi_agentic/README.md b/cookbook/rl/rsi_agentic/README.md new file mode 100644 index 000000000..3a26f55ff --- /dev/null +++ b/cookbook/rl/rsi_agentic/README.md @@ -0,0 +1,148 @@ +# Agentic RSI + +GRPO on multi-turn tool-using episodes. The solver is an **ms-agent** agent — +shell, filesystem, python, notebook, todo — working inside its own microVM. When +it stops calling tools the episode ends and the task's checks are run against +what it left behind. The checks are ordinary programs, so the same trajectory +always earns the same reward. + +## Where things run + +``` +training host sandbox (one microVM per episode) +───────────────────────────── ───────────────────────────────── +MsAgentHarness tool_server.py + system prompt, message shaping ms-agent ToolManager + llm: and tools: popped -> the real file_system / + constructs no tool at all code_executor / todo_list + ┌──► GET /tools schemas +RemoteMsAgentToolEnv ── curl over ─────┤ POST /call dispatch + forwards tool calls the sandbox │ + copies files back command channel└─── /workspace +``` + +Two properties this layout is built around: + +**Nothing the model emits executes next to the trainer.** The harness has its +`llm` and `tools` sections popped *after* ms-agent merges its own `agent.yaml` +underneath — omitting them from `rsi_agent.yaml` is not enough, since that +default declares `code_executor` and would otherwise put a live shell executor +on the training host with access to the whole machine. + +**The advertised tool contract is read off the code that honours it.** The tool +schemas in the prompt come from `GET /tools` on the sandbox, not from a second +ms-agent next to the trainer. A reimplementation of the tools would have been +much less work, but in RL the policy actively exploits whatever the executor +actually does, and any divergence from the production tools would only surface +after deployment. + +## Setup + +### Environment host + +Needs `/dev/kvm` and kernel 6.8+. Builds the template and runs the AgentENV +server (the same server `cookbook/rl/envs` uses; only the template differs). + +```bash +sh sandbox_server/install.sh # AgentENV + build the RSI template +sh sandbox_server/install.sh --rebuild # after changing the Dockerfile +sh sandbox_server/install.sh --skip-server # template only, server already up + +sh sandbox_server/serve.sh # foreground, binds 127.0.0.1:8000 +NOHUP=1 sh sandbox_server/serve.sh # background +``` + +On a restricted network, point the base image at a reachable registry: + +```bash +BASE_IMAGE=<your-registry>/library/python:3.11-slim sh sandbox_server/install.sh +``` + +The image installs **the `ms-agent/` checkout from this repo**, not the pip +release: the training host imports that same working copy, and a tool whose +output differs by one local commit is a train/serve mismatch the policy absorbs +silently. `install.sh` prints the staged commit — keep the host on it. + +### Training host + +```bash +pip install e2b + +AENV_API_URL=http://<env-host-ip>:8000 \ + python rsi_agentic_grpo.py + +# or through a tunnel: +ssh -N -L 8000:127.0.0.1:8000 root@<env-host-ip> +python rsi_agentic_grpo.py +``` + +Verify a sandbox boots and the runtime comes up before launching training: + +```bash +python -c " +import sys; sys.path.insert(0, '.') +from remote_tool_env import RemoteMsAgentToolEnv +e = RemoteMsAgentToolEnv(template='twinkle-rsi-msagent', config_path='rsi_agent.yaml', + api_url='http://127.0.0.1:8000') +e.reset() +print([t['function']['name'] for t in e.tool_schemas()]) +print(e.step('shell_executor', {'command': 'python -V && pwd'}).observation) +e.close() +" +``` + +## Configuration + +| Variable | Default | | +|---|---|---| +| `AENV_API_URL` | `http://127.0.0.1:8000` | AgentENV server | +| `AENV_TEMPLATE` | `twinkle-rsi-msagent` | template built by `install.sh` | +| `RSI_TASKS` | `tasks.example.jsonl` | task file | +| `RSI_AGENT_CONFIG` | `rsi_agent.yaml` | uploaded into every sandbox | +| `RSI_SANDBOX_TIMEOUT` | `900` | must outlast an episode plus its checks | +| `RSI_ENV_CONCURRENCY` | `16` | parallel boot / scoring | +| `RSI_MAX_TURNS` | `20` | tool-calling turns per episode | +| `RSI_SCORE_MODE` | `fraction` | or `all_or_nothing` | +| `RSI_KEEP_WORKSPACES` | `0` | keep the files copied out of each sandbox | + +Training hyper-parameters come from the CLI, e.g. +`python rsi_agentic_grpo.py --batch-size 2 --num-generations 4 --max-steps 2`. + +Sandbox count is `batch-size × num-generations`, each ~2GiB. At the defaults +(4 × 8) that is 32 microVMs, so size the environment host accordingly. + +## Files + +| File | Role | +|---|---| +| `rsi_agentic_grpo.py` | training loop, episode construction, scoring | +| `remote_tool_env.py` | training-side Env: forwards tool calls, copies files back | +| `rsi_agent.yaml` | ms-agent config — read by *both* halves | +| `sandbox_server/tool_server.py` | in-sandbox HTTP server owning the ToolManager | +| `sandbox_server/Dockerfile` | template image: ms-agent, ripgrep, ipykernel | +| `sandbox_server/install.sh` | build the template (delegates server bootstrap) | +| `sandbox_server/serve.sh` | start AgentENV (delegates to `cookbook/rl/envs`) | +| `tasks.example.jsonl` | one task per line: `id`, `query`, `checks` | + +## Decisions worth knowing + +**`read_file(abbreviate=True)` is withdrawn when no LLM is configured.** That +argument asks an LLM to summarise a file. The sandbox has no API key, so the +tool server drops the unusable `llm` section *and* removes the argument from the +advertised schema — the model is never offered something that can only fail. +Give `rsi_agent.yaml` a real `llm:` section with a key reachable from the +sandbox to get it back. + +**A failed sandbox boot skips the whole batch.** GRPO groups here are +positional: advantages are taken over consecutive runs of `num_generations`, so +dropping one episode would shift every later group onto the wrong task. There is +no retry — a boot failure is logged and the step is abandoned. + +**No web search.** ms-agent's `web_search` key only provides `fetch_page` +(retrieve a known URL). A real search tool needs `EXA_API_KEY` / `SERPAPI_API_KEY` +and is wired separately; no example task requires one. + +**Checks reach the sandbox two ways.** `file_*` checks read a local directory, +so the episode's files are copied out first (`download_workspace`, capped at 200 +files / 1MiB each). `shell` and `python` checks go back into the sandbox through +`env.runner()`, where the interpreter and packages are the ones the agent used. diff --git a/cookbook/rl/rsi_agentic/remote_tool_env.py b/cookbook/rl/rsi_agentic/remote_tool_env.py new file mode 100644 index 000000000..889bcd327 --- /dev/null +++ b/cookbook/rl/rsi_agentic/remote_tool_env.py @@ -0,0 +1,361 @@ +"""Training-side Env: ms-agent's tools, executed inside a remote sandbox. + +Pairs with ``sandbox_server/tool_server.py``. That server runs in the microVM +and owns the real ms-agent ``ToolManager``; this class is the client. Nothing +here knows what ``edit_file`` or ``shell_executor`` do -- it forwards a tool +call and returns whatever ms-agent produced, so the behaviour the policy is +trained against is the behaviour it will meet at serving time. + +Lifecycle mirrors :class:`twinkle_agentic.envs.AgentEnv`: one sandbox per +episode, created on ``reset`` and killed on ``close``. On top of that, ``reset`` +uploads the agent yaml and the server script from the training host and waits +for the runtime to come up, so iterating on either one does not mean rebuilding +the template image. + +Transport is HTTP, driven by ``curl`` over the sandbox's command channel rather +than a forwarded port. It costs one process spawn per turn -- noise next to a +shell command -- and in exchange depends only on ``commands.run`` and +``files.write``, which is the surface every e2b-compatible backend implements +the same way. +""" +import json +import os +import posixpath +import re +import time +from typing import Any, Dict, List, Optional, Sequence, Tuple + +from twinkle import get_logger +from twinkle_agentic.envs.base import Env, StepResult + +logger = get_logger() + +__all__ = ['RemoteMsAgentToolEnv'] + +# Marker used to recover an exit status from a tool that only returns text. +_RC_MARK = '__TWINKLE_RC__' +_RC_RE = re.compile(rf'{_RC_MARK}:(-?\d+)') + +_PY_WRAPPER = """\ +import sys, traceback +try: +{body} +except SystemExit as _e: + print('{mark}:%d' % (_e.code or 0)) + sys.exit(0) +except BaseException: + traceback.print_exc() + print('{mark}:1') +else: + print('{mark}:0') +""" + +_REMOTE_DIR = '/opt/rsi' +_LOCAL_SERVER = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'sandbox_server', 'tool_server.py') + + +class RemoteMsAgentToolEnv(Env): + """Run one episode's ms-agent tool calls inside a dedicated sandbox. + + Args: + template: AgentENV/e2b template name, built by ``sandbox_server/install.sh``. + config_path: ms-agent yaml on the *training host*. Uploaded on every + reset, so this file is the single source of truth for both sides. + api_url: AgentENV server base URL. Falls back to ``E2B_API_URL``. + api_key: API key; AgentENV accepts any non-empty string. + port: port the tool server listens on inside the sandbox. + workspace: ``config.output_dir`` inside the sandbox. + sandbox_timeout: sandbox idle timeout, in seconds. Must outlast a whole + episode plus the checks that run after it. + command_timeout: per-request timeout for a tool call, in seconds. + boot_timeout: how long to wait for the runtime to answer ``/health``. + ms-agent's import plus tool construction dominates this. + max_observation_chars: truncate a tool result before it becomes a + message. A single ``grep`` can otherwise fill the context window. + """ + + def __init__( + self, + template: str, + config_path: str, + *, + api_url: Optional[str] = None, + api_key: Optional[str] = None, + port: int = 8900, + workspace: str = '/workspace', + sandbox_timeout: int = 900, + command_timeout: int = 180, + boot_timeout: int = 300, + max_observation_chars: int = 8000, + ): + if not template: + raise ValueError("RemoteMsAgentToolEnv requires 'template'; build one with " + 'sandbox_server/install.sh') + if not os.path.exists(config_path): + raise FileNotFoundError(f'agent config not found: {config_path}') + self._template = template + self._config_path = config_path + self._api_url = api_url + self._api_key = api_key + self._port = port + self.workspace = workspace + self._sandbox_timeout = sandbox_timeout + self._command_timeout = command_timeout + self._boot_timeout = boot_timeout + self.max_observation_chars = max_observation_chars + self._sandbox = None + self._schemas: Optional[List[Dict[str, Any]]] = None + + # ------------------------------------------------------------------ Env + + def reset(self, trajectory: Optional[Dict[str, Any]] = None) -> StepResult: + """Boot a sandbox and bring ms-agent's tool runtime up inside it.""" + self.close() + self._sandbox = self._create_sandbox() + self._upload() + self._start_server() + self._await_ready() + self._schemas = None + return StepResult(observation='') + + def step(self, tool_name: str, arguments: Dict[str, Any] = None) -> StepResult: + return self.step_batch([(tool_name, arguments or {})])[0] + + def step_batch(self, calls: Sequence[Tuple[str, Dict[str, Any]]]) -> List[StepResult]: + """Send a turn's calls as one request; the server runs them together. + + Batching matters twice over: it is one sandbox round trip instead of + several, and it keeps ms-agent's own ``parallel_call_tool`` semantics + rather than serialising what production would run concurrently. + """ + calls = list(calls) + if not calls: + return [] + payload = { + 'calls': [{ + 'tool_name': name, + 'arguments': args or {} + } for name, args in calls], + 'timeout': self._command_timeout, + } + try: + body = self._rpc('/call', payload) + results = body.get('results') or [] + except Exception as e: # noqa + # A dead sandbox must not kill the training step: report it as an + # observation and let the episode play out (and score zero). + logger.warning(f'RemoteMsAgentToolEnv call failed: {e}') + results = [{'observation': f'Tool runtime unreachable: {e}'} for _ in calls] + if len(results) != len(calls): + results = (results + [{'observation': 'Tool runtime returned no result'}] * len(calls))[:len(calls)] + return [StepResult(observation=self._truncate(r.get('observation') or '')) for r in results] + + def close(self) -> None: + if self._sandbox is None: + return + sandbox_id = getattr(self._sandbox, 'sandbox_id', None) + try: + self._sandbox.kill() + except Exception as e: # noqa # best-effort: the backend evicts on timeout anyway + logger.warning(f'failed to kill sandbox {sandbox_id}: {e}') + finally: + self._sandbox = None + + # ------------------------------------------------------------ tool names + + def tool_schemas(self) -> List[Dict[str, Any]]: + """Schemas from the runtime that will execute them. + + These go straight into the prompt. Sourcing them from the executor + rather than from a second local ms-agent is what makes it impossible + for the advertised contract and the running code to disagree. + """ + if self._schemas is None: + self._schemas = list((self._rpc('/tools', None) or {}).get('tools') or []) + return list(self._schemas) + + def tool_names(self) -> List[str]: + names = [] + for schema in self.tool_schemas(): + name = (schema.get('function') or {}).get('name') + if name: + names.append(str(name)) + return names + + def resolve_tool(self, name: str) -> str: + """Map a plain tool name onto the runtime's own spelling. + + ms-agent namespaces its tools as ``{server}---{tool}``, so a caller that + asks for ``shell_executor`` means ``code_executor---shell_executor``. + An unknown name raises instead of being passed through: a mistyped tool + comes back as a failed call, which for a checker is indistinguishable + from a failed check, and a whole GRPO group would silently score zero. + """ + names = self.tool_names() + if name in names: + return name + matches = [n for n in names if n.rsplit('---', 1)[-1] == name] + if len(matches) == 1: + return matches[0] + if not matches: + raise ValueError(f'no registered tool named {name!r}; available: {names}') + raise ValueError(f'{name!r} is ambiguous across servers: {matches}') + + # ------------------------------------------------------- for the checker + + def runner(self, shell_tool: str = 'shell_executor', python_tool: str = 'python_executor'): + """A ``result_check`` runner that executes inside this episode's sandbox. + + Verification has to see the filesystem the agent actually wrote to, so + the check goes back through the same tools rather than a local + subprocess. Those tools return prose, not an exit status, so the command + is made to print a marker and the status is read back out of the output. + """ + shell_name = self.resolve_tool(shell_tool) + python_name = self.resolve_tool(python_tool) + + def _run(source: str, interpreter: str) -> Tuple[int, str]: + if interpreter == 'python': + body = '\n'.join(' ' + line for line in source.splitlines()) or ' pass' + code = _PY_WRAPPER.format(body=body, mark=_RC_MARK) + out = self.step(python_name, {'code': code}).observation + else: + out = self.step(shell_name, {'command': f'{source}\necho "{_RC_MARK}:$?"'}).observation + match = _RC_RE.search(out or '') + if match is None: + # No marker means the tool itself failed (timeout, sandbox down) + # rather than the check failing; report non-zero and keep output. + return 1, out or 'check produced no output and no exit marker' + return int(match.group(1)), _RC_RE.sub('', out or '').strip() + + return _run + + def download_workspace(self, dest: str, max_files: int = 200, max_bytes: int = 1 << 20) -> str: + """Copy the episode's files out of the sandbox for the ``file_*`` checks. + + Those checks read from an ordinary local directory, which is the right + interface for a generic verifier but cannot see inside a microVM. The + episode is over by the time this runs, so a snapshot is equivalent to + the live filesystem -- and the shell/python checks still go through + :meth:`runner`, against the sandbox itself. + + Files above ``max_bytes`` are skipped: a check that needs to look at a + 100MB artifact wants a command, not a copy. + """ + os.makedirs(dest, exist_ok=True) + listing = self._sandbox.commands.run( + f"find {self.workspace} -type f -size -{max(1, max_bytes // 1024)}k " + f'-printf "%P\\n" 2>/dev/null | head -n {max_files}', + timeout=60) + for rel in (listing.stdout or '').splitlines(): + rel = rel.strip() + if not rel: + continue + local = os.path.join(dest, rel) + os.makedirs(os.path.dirname(local), exist_ok=True) + try: + content = self._sandbox.files.read(posixpath.join(self.workspace, rel)) + except Exception as e: # noqa # an unreadable file fails its own check, not the batch + logger.debug(f'could not fetch {rel} from sandbox: {e}') + continue + mode = 'wb' if isinstance(content, (bytes, bytearray)) else 'w' + with open(local, mode) as f: + f.write(content) + return dest + + # -------------------------------------------------------------- private + + def _create_sandbox(self): + try: + from e2b import Sandbox + except ImportError as e: + raise ImportError('RemoteMsAgentToolEnv needs the e2b SDK: pip install e2b') from e + if self._api_url: + os.environ['E2B_API_URL'] = self._api_url + os.environ.setdefault('E2B_SANDBOX_URL', self._api_url) + if self._api_key: + os.environ['E2B_API_KEY'] = self._api_key + os.environ.setdefault('E2B_API_KEY', 'dummy') + os.environ.setdefault('E2B_ACCESS_TOKEN', 'dummy') + # AgentENV issues no keys, but the SDK asserts the key looks like + # ``e2b_[0-9a-f]+`` before sending anything. This is the SDK's own + # opt-out for deployments that do not mint e2b-format keys. + os.environ.setdefault('E2B_VALIDATE_API_KEY', 'false') + return Sandbox(template=self._template, timeout=self._sandbox_timeout) + + def _upload(self) -> None: + """Push the yaml and the server script into the sandbox. + + Uploading beats baking them into the image: the training host's copy is + authoritative, so editing a tool line-up is a restart rather than a + template rebuild, and the two halves cannot fall out of sync. + """ + with open(self._config_path, encoding='utf-8') as f: + self._sandbox.files.write(f'{_REMOTE_DIR}/rsi_agent.yaml', f.read()) + with open(_LOCAL_SERVER, encoding='utf-8') as f: + self._sandbox.files.write(f'{_REMOTE_DIR}/tool_server.py', f.read()) + + def _start_server(self) -> None: + command = (f'python {_REMOTE_DIR}/tool_server.py ' + f'--config {_REMOTE_DIR}/rsi_agent.yaml ' + f'--workspace {self.workspace} --port {self._port}') + try: + self._sandbox.commands.run(command, background=True) + except TypeError: + # Older SDKs have no `background`; detach with setsid so the server + # outlives the command that launched it. + self._sandbox.commands.run( + f'mkdir -p {self.workspace} && setsid nohup {command} ' + f'> /tmp/tool_server.log 2>&1 < /dev/null &', + timeout=30) + + def _await_ready(self) -> None: + """Poll ``/health`` until the runtime answers, then fail loudly. + + Silence here is worth an exception: a sandbox whose tools never came up + answers every call with an error, the episode scores zero, and the whole + GRPO group looks like a hard task rather than a broken environment. + """ + deadline = time.time() + self._boot_timeout + last = '' + while time.time() < deadline: + try: + if (self._rpc('/health', None, timeout=10) or {}).get('status') == 'ok': + return + except Exception as e: # noqa + last = str(e) + time.sleep(2) + log = '' + try: + log = (self._sandbox.commands.run('tail -n 40 /tmp/tool_server.log', timeout=20).stdout or '') + except Exception: # noqa + pass + raise RuntimeError(f'ms-agent tool runtime did not come up within {self._boot_timeout}s ' + f'(last error: {last})\n{log}') + + def _rpc(self, path: str, payload: Optional[Dict[str, Any]], timeout: Optional[int] = None) -> Dict[str, Any]: + """One request to the in-sandbox server, via curl on the command channel. + + The body is written to a file rather than inlined: tool arguments carry + arbitrary source code, and no amount of shell quoting survives that + reliably. + """ + seconds = timeout or self._command_timeout + if payload is None: + command = f'curl -sS -m {seconds} http://127.0.0.1:{self._port}{path}' + else: + request = f'{_REMOTE_DIR}/request.json' + self._sandbox.files.write(request, json.dumps(payload, ensure_ascii=False)) + command = (f'curl -sS -m {seconds} -X POST -H "Content-Type: application/json" ' + f'--data-binary @{request} http://127.0.0.1:{self._port}{path}') + result = self._sandbox.commands.run(command, timeout=seconds + 30) + stdout = (getattr(result, 'stdout', '') or '').strip() + if not stdout: + raise RuntimeError(f'empty response from {path}: {getattr(result, "stderr", "")}') + return json.loads(stdout) + + def _truncate(self, text: str) -> str: + limit = self.max_observation_chars + if limit and len(text) > limit: + return f'{text[:limit]}\n...[truncated {len(text) - limit} chars]' + return text diff --git a/cookbook/rl/rsi_agentic/rsi_agent.yaml b/cookbook/rl/rsi_agentic/rsi_agent.yaml index 67a609e25..63ca8fd1d 100644 --- a/cookbook/rl/rsi_agentic/rsi_agent.yaml +++ b/cookbook/rl/rsi_agentic/rsi_agent.yaml @@ -1,16 +1,15 @@ # ms-agent config for agentic RSI training. # -# This file is the framework-specific half on purpose: the system prompt, the -# tool line-up and the sandbox settings live here in cookbook, while -# src/twinkle_agentic stays generic. Point RSI_AGENT_CONFIG at a copy of this -# to change the agent without touching the trainer. +# Read by both halves of the setup, which is the point: # -# The entry script deletes the `llm:` section before preparing the agent. -# Omitting it here is not enough: ms-agent merges this file over its own -# ms_agent/agent/agent.yaml, which does declare one, and FileSystemTool builds -# an LLM client whenever config.llm exists -- which then asserts on a missing -# modelscope_api_key. Generation comes from twinkle's vLLM sampler, so no -# second model should be reachable from here at all. +# * the training host loads it to build a MsAgentHarness for message shaping +# only -- the entry script drops `llm:` and `tools:` from the merged config +# first, so no tool is ever constructed next to the trainer; +# * remote_tool_env.py uploads this same file into each sandbox, where +# sandbox_server/tool_server.py loads it and does construct the tools. +# +# So the tool line-up below describes what runs in the microVM. Editing it takes +# effect on the next episode; no image rebuild is involved. prompt: # Unset -> ms-agent's built-in agent prompt, which is also what serving uses. @@ -31,9 +30,10 @@ max_chat_round: 9999 interactive: false permission_mode: auto -# Overwritten per trajectory by the entry script. Both file_system and -# code_executor root themselves here, so this is what isolates episodes. -output_dir: output/rsi_agentic/workspace +# Path *inside the sandbox*. One microVM per episode already isolates +# trajectories from each other, so this is a fixed path rather than a per-slot +# directory; the entry script overrides it only to match --workspace. +output_dir: /workspace callbacks: [] @@ -48,8 +48,10 @@ tools: - glob code_executor: mcp: false - # python_env runs on the host; switch to the docker implementation for - # untrusted code, at the cost of a container per episode. + # python_env means "run in this process's machine", and that machine is the + # microVM -- the sandbox boundary is the VM itself, not this setting. Do not + # switch to the docker implementation: it would nest a container inside the + # VM for no extra isolation. implementation: python_env include: - shell_executor @@ -62,3 +64,12 @@ tools: # fetch_page (retrieve a known URL); a real query-a-search-engine tool needs # EXA_API_KEY / SERPAPI_API_KEY and is wired separately from the plain tool # list. Add it here once that is decided; until then no task should need it. + +# No `llm:` section on purpose, and note that omitting it is not the same as +# disabling it: ms-agent merges this file over its own ms_agent/agent/agent.yaml, +# which declares `service: modelscope`. The tool server treats a section with no +# credentials as absent, drops it, and then withdraws the one argument that +# needed it (read_file's `abbreviate`, an LLM-written file summary) from the +# advertised schema -- so the model is never offered a tool argument that cannot +# work. Put a real `llm:` here, with a key reachable from the sandbox, to get +# that argument back. diff --git a/cookbook/rl/rsi_agentic/rsi_agentic_grpo.py b/cookbook/rl/rsi_agentic/rsi_agentic_grpo.py index 4aa6d690b..4b49bb490 100644 --- a/cookbook/rl/rsi_agentic/rsi_agentic_grpo.py +++ b/cookbook/rl/rsi_agentic/rsi_agentic_grpo.py @@ -1,32 +1,38 @@ """Agentic RSI: GRPO on multi-turn tool-using episodes, scored by program checks. The solver is an ms-agent agent with a real tool line-up (shell, filesystem, -python, notebook sandbox, web search, todo list) working in its own directory. -It explores for as many turns as it needs; when it stops calling tools the -episode ends and the task's checks are run against what it left behind. The -checks are ordinary programs -- file exists, file content, command exit status, -final answer match -- so the same trajectory always earns the same reward. - -Everything framework-specific lives here and in ``rsi_agent.yaml``: the tool -line-up, the sandbox settings, the task file. ``src/twinkle_agentic`` stays -generic -- ``MsAgentHarness`` shapes messages, ``MsAgentToolEnv`` executes -tools, ``result_check`` scores outcomes, and none of them know about RSI. - -Layout mirrors cookbook/rl/multi_turn/multi_turn_grpo.py; the differences are -the harness (ms-agent owns the system prompt and message evolution) and the -reward (program checks over the end state instead of an env-emitted scalar). +python, notebook sandbox, todo list) working in its own microVM. It explores for +as many turns as it needs; when it stops calling tools the episode ends and the +task's checks are run against what it left behind. The checks are ordinary +programs -- file exists, file content, command exit status, final answer match -- +so the same trajectory always earns the same reward. + +The two halves are split by where they run, not by what they know: + + * On the training host, ``MsAgentHarness`` shapes messages -- system prompt, + tool-result formatting, ms-agent's own message evolution. Its ``llm`` and + ``tools`` sections are dropped before it prepares, so it constructs no tool + and nothing the model emits can execute next to the trainer. + * In the sandbox, ``sandbox_server/tool_server.py`` holds the real ms-agent + ``ToolManager``. It also *supplies the tool schemas*, which the prompt then + advertises verbatim -- the contract the model is trained against is read off + the code that will honour it, so the two cannot drift apart. Usage: + AENV_API_URL=http://127.0.0.1:8000 \\ RSI_TASKS=cookbook/rl/rsi_agentic/tasks.example.jsonl \\ python cookbook/rl/rsi_agentic/rsi_agentic_grpo.py +See README.md for building the template and starting the sandbox server. + Task file: one JSON object per line, with ``id``, ``query`` and ``checks`` (see tasks.example.jsonl and twinkle_agentic.verifier.result_check.Check). """ import json import os import shutil -from typing import Any, Dict, List, Tuple +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Any, Dict, List, Optional, Tuple from peft import LoraConfig @@ -41,13 +47,17 @@ from twinkle.processor import InputProcessor from twinkle.sampler import vLLMSampler from twinkle.template import Template -from twinkle_agentic.envs import EnvTool, MsAgentToolEnv +from twinkle_agentic.envs import EnvTool from twinkle_agentic.harness import MsAgentHarness from twinkle_agentic.rollout.multi_turn import MultiTurnRollout from twinkle_agentic.tools.tool_manager import ToolManager from twinkle_agentic.verifier.result_check import (CheckContext, checks_from_dicts, run_checks) +# Same directory as this script, which python puts on sys.path when it is run +# as a file. Kept in cookbook because it is RSI-specific wiring, not framework. +from remote_tool_env import RemoteMsAgentToolEnv # noqa: I100,I202 + logger = get_logger() args = CLI.from_args() @@ -77,16 +87,31 @@ TASKS_PATH = os.environ.get('RSI_TASKS', 'cookbook/rl/rsi_agentic/tasks.example.jsonl') AGENT_CONFIG = os.environ.get('RSI_AGENT_CONFIG', 'cookbook/rl/rsi_agentic/rsi_agent.yaml') RUN_DIR = os.environ.get('RSI_RUN_DIR', 'output/rsi_agentic/run') + +# Sandbox backend. The template is built by sandbox_server/install.sh. +SANDBOX_TEMPLATE = os.environ.get('AENV_TEMPLATE', 'twinkle-rsi-msagent') +SANDBOX_API_URL = os.environ.get('AENV_API_URL', 'http://127.0.0.1:8000') +# Must outlast a whole episode plus the checks that run after it. +SANDBOX_TIMEOUT = int(os.environ.get('RSI_SANDBOX_TIMEOUT', 900)) +# Booting and scoring are network-bound, so they are done on threads. This caps +# how many sandboxes are talked to at once, not how many exist. +ENV_CONCURRENCY = int(os.environ.get('RSI_ENV_CONCURRENCY', 16)) + # 'fraction' gives partial credit per check; 'all_or_nothing' is stricter and # produces a cleaner pass/fail signal at the cost of a sparser reward. SCORE_MODE = os.environ.get('RSI_SCORE_MODE', 'fraction') -# Keep each episode's workspace after scoring. Useful while debugging tasks, -# expensive over a long run. +# Keep each episode's downloaded files after scoring. Useful while debugging +# tasks, expensive over a long run. KEEP_WORKSPACES = os.environ.get('RSI_KEEP_WORKSPACES', '0') == '1' def load_tasks(path: str) -> List[Dict[str, Any]]: - """Read the task file and fail loudly on a task that can never be scored.""" + """Read the task file and fail loudly on a task that can never be scored. + + Supports two formats: + - Structured checks: ``{"checks": [{"kind": ..., ...}]}`` (legacy) + - Script checks: ``{"check_script": "assert ..."}`` (from agentic challenger) + """ tasks = [] with open(path, encoding='utf-8') as f: for lineno, line in enumerate(f, 1): @@ -95,53 +120,107 @@ def load_tasks(path: str) -> List[Dict[str, Any]]: task = json.loads(line) if not task.get('query'): raise ValueError(f'{path}:{lineno} has no query') - if not task.get('checks'): - # An unchecked task scores 0 for every rollout, so the whole - # group has zero advantage and contributes no gradient. - raise ValueError(f'{path}:{lineno} ({task.get("id")}) declares no checks') - task['_checks'] = checks_from_dicts(task['checks']) + if task.get('check_script'): + # New format: raw python script, scored by exit code + task['_checks'] = None + elif task.get('checks'): + # Legacy format: structured Check dicts + task['_checks'] = checks_from_dicts(task['checks']) + else: + raise ValueError(f'{path}:{lineno} ({task.get("id")}) declares no checks ' + f'and no check_script') tasks.append(task) if not tasks: raise ValueError(f'{path} contains no tasks') return tasks -def build_episode(task: Dict[str, Any], slot: int, step: int) -> Tuple[Any, Any, Any, Dict]: - """Create one episode: harness + isolated workspace + bound tool manager. +def build_episode(task: Dict[str, Any]) -> Tuple[Any, Any, Any, Dict]: + """Create one episode: a sandbox with ms-agent's tools, plus a local harness. - The harness and the Env share one ms-agent runtime, so the tools named in - the prompt are exactly the tools that will run. Each episode gets its own - ``output_dir`` -- that directory is both the sandbox root and what the - checks will later inspect. + The harness is stripped down to message shaping. Popping ``tools`` matters + as much as popping ``llm``, and for the same reason omitting the section + from the yaml is not enough: ms-agent merges its own agent.yaml underneath, + which declares file_system and code_executor, so a live shell executor would + otherwise be constructed on the training host with access to the whole + machine. Popping them after the merge leaves the harness with zero tools -- + and the system prompt byte-identical, because ms-agent does not fold the + tool list into it. """ from omegaconf import OmegaConf, open_dict - workspace = os.path.join(RUN_DIR, f'step{step:06d}', f'slot{slot:03d}') - os.makedirs(workspace, exist_ok=True) - cfg = OmegaConf.load(AGENT_CONFIG) - with open_dict(cfg): - cfg.output_dir = os.path.abspath(workspace) - harness = MsAgentHarness(config=cfg) - # ms-agent merges the config above over its own default agent.yaml, which - # declares an `llm:` section; FileSystemTool then builds a remote LLM client - # from it and asserts on a missing api key. Generation here comes from the - # vLLM sampler, so drop that section before any tool is constructed. with open_dict(harness.agent.config): harness.agent.config.pop('llm', None) + harness.agent.config.pop('tools', None) harness.prepare() - env = MsAgentToolEnv(agent=harness.agent, workspace=workspace) - # Same schema list on both sides: prompt and executor cannot drift apart. - tool_manager = ToolManager(EnvTool.from_schemas(env, harness.tool_schemas())) + env = RemoteMsAgentToolEnv( + template=SANDBOX_TEMPLATE, + config_path=AGENT_CONFIG, + api_url=SANDBOX_API_URL, + sandbox_timeout=SANDBOX_TIMEOUT, + ) + env.reset() trajectory = harness.start(task['query']) + # The executor's own schemas, not the harness's (which are now empty by + # construction). Advertising what will run is the whole point of sourcing + # them from the sandbox. + schemas = env.tool_schemas() + trajectory['tools'] = schemas + tool_manager = ToolManager(EnvTool.from_schemas(env, schemas)) return harness, env, tool_manager, trajectory -def score_episode(task: Dict[str, Any], env: MsAgentToolEnv, trajectory: Dict[str, Any]) -> float: - """Run the task's checks against the state this episode left behind.""" +def boot_episodes(tasks: List[Dict[str, Any]]) -> List[Tuple[Any, Any, Any, Dict]]: + """Bring up every rollout's sandbox at once, all-or-nothing. + + Serial boot would dominate the step: a microVM plus ms-agent's import runs + to seconds, multiplied by ``batch_size x num_generations``. + + All-or-nothing because GRPO groups here are positional -- advantages are + taken over consecutive runs of ``NUM_GENERATIONS`` -- so dropping one + episode would not shrink its group, it would shift every later group onto + the wrong task. + """ + episodes: List[Optional[Tuple[Any, Any, Any, Dict]]] = [None] * len(tasks) + error: Optional[BaseException] = None + with ThreadPoolExecutor(max_workers=ENV_CONCURRENCY) as pool: + futures = {pool.submit(build_episode, task): slot for slot, task in enumerate(tasks)} + for future in as_completed(futures): + try: + episodes[futures[future]] = future.result() + except Exception as e: # noqa + error = error or e + if error is not None: + for episode in episodes: + if episode is not None: + episode[1].close() + raise RuntimeError(f'sandbox boot failed: {error}') from error + return episodes # type: ignore[return-value] + + +def score_episode(task: Dict[str, Any], env: RemoteMsAgentToolEnv, trajectory: Dict[str, Any], + snapshot_dir: str) -> float: + """Run the task's checks against the state this episode left behind. + + Supports two scoring paths: + - check_script: run a python script in the sandbox; exit 0 = score 1.0 + - structured checks: download workspace + run_checks (legacy) + """ + check_script = task.get('check_script') + if check_script: + # New path: run the script directly in the sandbox's python executor + runner = env.runner() + exit_code, output = runner(check_script, 'python') + if exit_code != 0: + logger.debug(f'[{task["id"]}] check_script failed (exit {exit_code}): ' + f'{output[-200:]}') + return 1.0 if exit_code == 0 else 0.0 + + # Legacy path: structured checks final_answer = '' for msg in reversed(trajectory.get('messages') or []): if msg.get('role') == 'assistant' and (msg.get('content') or '').strip(): @@ -149,10 +228,8 @@ def score_episode(task: Dict[str, Any], env: MsAgentToolEnv, trajectory: Dict[st break ctx = CheckContext( - workspace=env.workspace, + workspace=env.download_workspace(snapshot_dir), final_answer=final_answer, - # Route shell/python checks back through the episode's own sandbox so - # they see the filesystem the agent actually wrote to. runner=env.runner(), ) report = run_checks(task['_checks'], ctx, mode=SCORE_MODE) @@ -162,6 +239,27 @@ def score_episode(task: Dict[str, Any], env: MsAgentToolEnv, trajectory: Dict[st return report.score +def score_episodes(tasks: List[Dict[str, Any]], envs: List[RemoteMsAgentToolEnv], + outs: List[Dict[str, Any]], step: int) -> List[float]: + """Score every episode in parallel; a scoring crash costs one reward, not the step. + + Each check is a sandbox round trip, so scoring serially would idle the GPUs + for as long as booting did. An episode whose sandbox died mid-check scores + zero, which is also what it would have scored had the checks simply failed. + """ + + def _score(slot: int) -> float: + snapshot = os.path.join(RUN_DIR, f'step{step:06d}', f'slot{slot:03d}') + try: + return score_episode(tasks[slot], envs[slot], outs[slot], snapshot) + except Exception as e: # noqa + logger.warning(f'[step {step} slot {slot}] scoring failed: {e}') + return 0.0 + + with ThreadPoolExecutor(max_workers=ENV_CONCURRENCY) as pool: + return list(pool.map(_score, range(len(outs)))) + + def main(): tasks = load_tasks(TASKS_PATH) logger.info(f'Loaded {len(tasks)} tasks from {TASKS_PATH}') @@ -225,6 +323,8 @@ def main(): optim_step = 0 task_cursor = 0 logger.info(f'Starting agentic RSI GRPO (max_turns={MAX_TURNS}, score={SCORE_MODE})') + logger.info(f'Sandboxes: template={SANDBOX_TEMPLATE} api={SANDBOX_API_URL} ' + f'concurrency={ENV_CONCURRENCY}') logger.info(get_device_placement()) while optim_step < MAX_STEPS: @@ -237,12 +337,19 @@ def main(): episode_tasks = [t for t in batch_tasks for _ in range(NUM_GENERATIONS)] harnesses, envs, tool_managers, trajectories = [], [], [], [] - for slot, task in enumerate(episode_tasks): - h, env, tm, traj = build_episode(task, slot, optim_step) - harnesses.append(h) + try: + episodes = boot_episodes(episode_tasks) + except Exception as e: # noqa + # A sandbox that never came up answers every call with an error, so + # the group would score a uniform zero and look like a hard task + # rather than a broken environment. Skip the batch and say so. + logger.warning(f'[Step {optim_step}] {e}; skipping batch') + continue + for harness, env, tool_manager, trajectory in episodes: + harnesses.append(harness) envs.append(env) - tool_managers.append(tm) - trajectories.append(traj) + tool_managers.append(tool_manager) + trajectories.append(trajectory) ckpt_manager.sync_weights(merge_and_sync=False) sampler.reset_prefix_cache() @@ -251,8 +358,7 @@ def main(): outs: List[Dict[str, Any]] = rollout( trajectories, harness=harnesses, tool_manager=tool_managers) - rewards = [score_episode(task, env, traj) - for task, env, traj in zip(episode_tasks, envs, outs)] + rewards = score_episodes(episode_tasks, envs, outs, optim_step) finally: # Sandboxes are a finite resource; a step that raises must still # give them back or the next step starts short. diff --git a/cookbook/rl/rsi_agentic/sandbox_server/Dockerfile b/cookbook/rl/rsi_agentic/sandbox_server/Dockerfile new file mode 100644 index 000000000..16212f63e --- /dev/null +++ b/cookbook/rl/rsi_agentic/sandbox_server/Dockerfile @@ -0,0 +1,24 @@ +FROM python:3.11-slim + +# ripgrep is not optional: file_system's `grep` uses `rg` when it is on PATH and +# silently falls back to a Python scan with a different output shape when it is +# not. The policy is trained on whatever it sees, so the sandbox has to take the +# same branch a serving deployment does. +RUN apt-get update && apt-get install -y --no-install-recommends \ + ca-certificates curl git ripgrep \ + && rm -rf /var/lib/apt/lists/* + +# The exact ms-agent checkout the training host imports, staged here by +# install.sh -- not `pip install ms-agent`. The host runs an editable install of +# this working copy, and a tool whose output differs by one local commit is a +# train/serve mismatch that only shows up after deployment. +COPY ms-agent /opt/ms-agent +RUN pip install --no-cache-dir -e /opt/ms-agent + +# notebook_executor pip-installs these the first time it is called. In a sandbox +# that is either a network round trip at the start of every episode or an +# outright failure on an air-gapped host, so they ship in the image. +RUN pip install --no-cache-dir ipykernel jupyter-client + +ENV PYTHONUNBUFFERED=1 +WORKDIR /workspace diff --git a/cookbook/rl/rsi_agentic/sandbox_server/install.sh b/cookbook/rl/rsi_agentic/sandbox_server/install.sh new file mode 100644 index 000000000..8cce9c457 --- /dev/null +++ b/cookbook/rl/rsi_agentic/sandbox_server/install.sh @@ -0,0 +1,74 @@ +#!/bin/sh +# Build the sandbox template for agentic RSI. +# +# The AgentENV server itself is the same one cookbook/rl/envs uses; only the +# template differs. Its bootstrap (install, host provisioning, config seeding) +# is delegated rather than copied, so there is one place to fix when it changes. +set -eu + +TEMPLATE="${TEMPLATE:-twinkle-rsi-msagent}" +# ms-agent pulls in pandas/matplotlib/modelscope and notebook_executor starts a +# real ipykernel, so 1GiB (the plain code template's size) is not enough. +CPU_COUNT="${CPU_COUNT:-2}" +MEMORY_MB="${MEMORY_MB:-2048}" +BASE_IMAGE="${BASE_IMAGE:-}" + +SKIP_SERVER=0 +REBUILD=0 +for arg in "$@"; do + case "$arg" in + --skip-server) SKIP_SERVER=1 ;; + --rebuild) REBUILD=1 ;; + *) echo "Unknown option: $arg" >&2; exit 2 ;; + esac +done + +cd "$(dirname "$0")" +REPO_ROOT=$(cd ../../../.. && pwd) +MS_AGENT="$REPO_ROOT/ms-agent" + +[ -f "$MS_AGENT/setup.py" ] || { + echo "ms-agent checkout not found at $MS_AGENT" >&2 + echo "The image installs the same source the training host imports; without it" >&2 + echo "the sandbox would run a different ms-agent than training assumes." >&2 + exit 1 +} + +if [ "$SKIP_SERVER" = "0" ]; then + echo "==> Installing the AgentENV server (shared with cookbook/rl/envs)" + sh ../../envs/agentenv_server/install.sh --skip-build +fi + +# A fresh staging directory per run, so a stale ms-agent copy can never end up +# in the image and nothing has to be deleted to make room. +STAGE=$(mktemp -d) +echo "==> Staging build context in $STAGE" +cp Dockerfile "$STAGE/" +# .git and caches are megabytes of noise in a build context and would also +# invalidate the layer cache on every commit. +tar -C "$(dirname "$MS_AGENT")" \ + --exclude='.git' --exclude='__pycache__' --exclude='*.pyc' \ + -cf - "$(basename "$MS_AGENT")" | tar -C "$STAGE" -xf - + +HEAD_SHA=$(git -C "$MS_AGENT" rev-parse --short HEAD 2>/dev/null || echo unknown) +echo " ms-agent staged at commit $HEAD_SHA" +echo " the training host must import this same checkout; if it does not," +echo " tool behaviour differs between rollout and the agent you deploy." + +if [ "$REBUILD" = "1" ]; then + echo "==> Deleting template '$TEMPLATE'" + aenv template delete "$TEMPLATE" || true +fi + +echo "==> Building template '$TEMPLATE' (cpu=$CPU_COUNT mem=${MEMORY_MB}MiB)" +set -- "$STAGE/Dockerfile" -t "$TEMPLATE" --cpu-count "$CPU_COUNT" --memory-mb "$MEMORY_MB" +[ -n "$BASE_IMAGE" ] && set -- "$@" --image "$BASE_IMAGE" +aenv build "$@" + +echo +echo "Build runs server-side and takes a few minutes. Follow it with:" +echo " aenv template watch <template-id> # id printed above" +echo " aenv template list # confirm it reaches ready" +echo +echo "Then start the server:" +echo " sh serve.sh" diff --git a/cookbook/rl/rsi_agentic/sandbox_server/serve.sh b/cookbook/rl/rsi_agentic/sandbox_server/serve.sh new file mode 100644 index 000000000..74fdab11d --- /dev/null +++ b/cookbook/rl/rsi_agentic/sandbox_server/serve.sh @@ -0,0 +1,20 @@ +#!/bin/sh +# Start the AgentENV server that hosts the RSI sandboxes. +# +# Deliberately a delegation, not a copy. It is the same server as +# cookbook/rl/envs uses -- RSI only changes which template the sandboxes boot +# from -- and that script carries a hundred lines of host-provisioning detail +# (capability wrapper, config path, systemd handover) that would silently drift +# if it existed twice. +# +# All of its environment variables still apply, e.g.: +# API_ADDR=0.0.0.0:8000 NOHUP=1 sh serve.sh +set -eu +cd "$(dirname "$0")" + +SHARED=../../envs/agentenv_server/serve.sh +[ -f "$SHARED" ] || { + echo "Shared AgentENV launcher not found: $SHARED" >&2 + exit 1 +} +exec sh "$SHARED" "$@" diff --git a/cookbook/rl/rsi_agentic/sandbox_server/tool_server.py b/cookbook/rl/rsi_agentic/sandbox_server/tool_server.py new file mode 100644 index 000000000..c8ad53976 --- /dev/null +++ b/cookbook/rl/rsi_agentic/sandbox_server/tool_server.py @@ -0,0 +1,262 @@ +"""ms-agent's tool runtime, served over HTTP from inside the sandbox. + +This is the half of the RSI setup that runs *in* the microVM. It builds a real +``LLMAgent`` from the same ``rsi_agent.yaml`` the training host reads, lets +ms-agent prepare its own tools, and exposes two things over HTTP: + +* ``GET /tools`` -- the tool schemas, taken from the runtime that will execute + them. The training host advertises these to the model verbatim, so the + contract in the prompt and the code behind it cannot drift apart. +* ``POST /call`` -- dispatch, through ms-agent's own ``single_call_tool`` / + ``parallel_call_tool``. + +Nothing here reimplements a tool. That is the whole point: the policy is +trained against the same ``edit_file`` / ``grep`` / ``shell_executor`` behaviour +it will meet at serving time, down to the output formatting. A reimplementation +would be cheaper, but in RL any divergence gets actively exploited by the policy +and only shows up after deployment. + +The server is deliberately stdlib-only so the sandbox image stays close to +ms-agent's own dependency set. + +Run inside the sandbox:: + + python tool_server.py --config /opt/rsi/rsi_agent.yaml --workspace /workspace +""" +import argparse +import asyncio +import copy +import json +import os +import sys +import threading +import traceback +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Any, Dict, List, Optional + +DEFAULT_PORT = 8900 + +# read_file's only LLM-backed argument: it summarises a file instead of +# returning it verbatim. Without a reachable LLM the tool cannot honour it, so +# it is also removed from the advertised schema -- see `_usable_llm`. +_LLM_BACKED_ARGS = {'file_system---read_file': ('abbreviate', )} + + +def _usable_llm(cfg) -> bool: + """Whether the declared ``llm`` section can actually serve a request. + + ms-agent merges its own ``agent.yaml`` underneath the user's, and that + default declares ``service: modelscope``. So an absent ``llm:`` section in + rsi_agent.yaml does not mean "no LLM" -- it means "modelscope, with no + credentials", which asserts as soon as FileSystemTool is constructed. The + presence of a key is what decides it. + """ + llm = getattr(cfg, 'llm', None) + if llm is None: + return False + service = str(getattr(llm, 'service', '') or '') + key_fields = (f'{service}_api_key', 'api_key', 'openai_api_key') + return any(getattr(llm, f, None) or os.environ.get(f.upper()) for f in key_fields) + + +def _without_llm_args(schema: Dict[str, Any]) -> Dict[str, Any]: + """Drop arguments this deployment cannot serve from a tool schema. + + Everything reachable from ``/tools`` has to be executable, or the model + spends the episode learning that a documented argument is broken and carries + that lesson to a deployment where it works. + """ + fn = schema.get('function') or {} + drop = _LLM_BACKED_ARGS.get(fn.get('name')) + properties = ((fn.get('parameters') or {}).get('properties') or {}) + if not drop or not any(arg in properties for arg in drop): + return schema + schema = copy.deepcopy(schema) + for arg in drop: + schema['function']['parameters']['properties'].pop(arg, None) + return schema + + +class _LoopThread: + """A single long-lived asyncio loop, owned by a background thread. + + ms-agent's tools bind state to the loop that created them: the notebook + kernel, MCP client sessions and subprocess transports all hold references + to it. Running ``asyncio.run`` per request would strand that state -- the + notebook would lose its variables between turns -- so one loop is created + at startup and every request is submitted onto it. + """ + + def __init__(self) -> None: + self._loop = asyncio.new_event_loop() + self._thread = threading.Thread(target=self._serve, name='ms-agent-loop', daemon=True) + self._thread.start() + + def _serve(self) -> None: + asyncio.set_event_loop(self._loop) + self._loop.run_forever() + + def run(self, coro, timeout: Optional[float] = None): + return asyncio.run_coroutine_threadsafe(coro, self._loop).result(timeout) + + +class ToolRuntime: + """Owns the ms-agent agent and answers tool queries against it.""" + + def __init__(self, config_path: str, workspace: str) -> None: + from omegaconf import OmegaConf, open_dict + + from ms_agent.agent.llm_agent import LLMAgent + + cfg = OmegaConf.load(config_path) + with open_dict(cfg): + cfg.output_dir = workspace + # Same non-interactive stubs MsAgentHarness applies on the training + # host: nothing here can answer a TUI permission prompt, and a tool + # blocking on stdin would hang the episode until the sandbox timeout. + cfg.interactive = False + cfg.permission_mode = 'auto' + self.has_llm = _usable_llm(cfg) + if not self.has_llm: + # Leaving an unusable section in place is not an option: + # FileSystemTool builds a client from it eagerly and asserts on + # the missing key, so no tool at all would come up. + cfg.pop('llm', None) + self.agent = LLMAgent(cfg) + self.agent._interactive = False + self.agent._event_sink = None + self.agent._input_source = None + self.workspace = workspace + self._loop = _LoopThread() + self._loop.run(self._prepare()) + + async def _prepare(self) -> None: + self.agent.prepare_runtime() + await self.agent.prepare_tools() + + @property + def _tm(self): + return self.agent.tool_manager + + def tools(self) -> List[Dict[str, Any]]: + """Tool schemas, flattened to a plain OpenAI-shaped list. + + ``get_tools`` groups by server; the model only ever sees the flat list, + and the names are already namespaced as ``{server}---{tool}``. + """ + raw = self._loop.run(self._tm.get_tools()) + if isinstance(raw, dict): + flat: List[Any] = [] + for value in raw.values(): + flat.extend(value if isinstance(value, list) else [value]) + else: + flat = list(raw or []) + return [t if self.has_llm else _without_llm_args(t) for t in flat if isinstance(t, dict)] + + def call(self, calls: List[Dict[str, Any]], timeout: Optional[float]) -> List[Dict[str, Any]]: + """Dispatch a turn's tool calls, mirroring how ms-agent itself does it. + + A single call goes through ``single_call_tool`` and a batch through + ``parallel_call_tool``, matching LLMAgent, so concurrency-sensitive + tools behave in training exactly as they do in production. + """ + payload = [{'tool_name': c.get('tool_name'), 'arguments': c.get('arguments') or {}} for c in calls] + try: + if len(payload) == 1: + results = [self._loop.run(self._tm.single_call_tool(payload[0]), timeout)] + else: + results = self._loop.run(self._tm.parallel_call_tool(payload), timeout) + except Exception as e: # noqa + # One failing tool must not take down the server: the episode can + # still recover, and a dead server would fail every later step of + # every trajectory sharing this sandbox. + detail = f'{type(e).__name__}: {e}' + return [{'observation': f'Tool call failed. {detail}', 'ok': False} for _ in payload] + return [{'observation': _as_text(r), 'ok': True} for r in list(results)] + + +def _as_text(result: Any) -> str: + if result is None: + return '' + if isinstance(result, str): + return result + try: + return json.dumps(result, ensure_ascii=False) + except (TypeError, ValueError): + return str(result) + + +class _Handler(BaseHTTPRequestHandler): + runtime: ToolRuntime = None # set on the class before the server starts + protocol_version = 'HTTP/1.1' + + def do_GET(self) -> None: # noqa: N802 - BaseHTTPRequestHandler's spelling + if self.path.startswith('/health'): + self._reply(200, {'status': 'ok', 'workspace': self.runtime.workspace}) + elif self.path.startswith('/tools'): + self._guarded(lambda: {'tools': self.runtime.tools()}) + else: + self._reply(404, {'error': f'no such endpoint: {self.path}'}) + + def do_POST(self) -> None: # noqa: N802 + if not self.path.startswith('/call'): + self._reply(404, {'error': f'no such endpoint: {self.path}'}) + return + length = int(self.headers.get('Content-Length') or 0) + try: + body = json.loads(self.rfile.read(length) or b'{}') + except ValueError as e: + self._reply(400, {'error': f'malformed request body: {e}'}) + return + calls = body.get('calls') or [] + if not isinstance(calls, list) or not calls: + self._reply(400, {'error': "'calls' must be a non-empty list"}) + return + self._guarded(lambda: {'results': self.runtime.call(calls, body.get('timeout'))}) + + def _guarded(self, produce) -> None: + """Answer with ``produce()``, turning a crash into a 500 with a traceback. + + The client surfaces the body as the observation, so a bug in here shows + up in the trajectory instead of as an opaque connection reset. + """ + try: + self._reply(200, produce()) + except Exception: # noqa + self._reply(500, {'error': traceback.format_exc()}) + + def _reply(self, code: int, payload: Dict[str, Any]) -> None: + data = json.dumps(payload, ensure_ascii=False).encode('utf-8') + self.send_response(code) + self.send_header('Content-Type', 'application/json; charset=utf-8') + self.send_header('Content-Length', str(len(data))) + self.end_headers() + self.wfile.write(data) + + def log_message(self, fmt: str, *args: Any) -> None: + sys.stderr.write('[tool_server] %s\n' % (fmt % args)) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--config', required=True, help='ms-agent yaml, the same one the training host loads') + parser.add_argument('--workspace', default='/workspace', help='config.output_dir for this episode') + parser.add_argument('--host', default='0.0.0.0') + parser.add_argument('--port', type=int, default=DEFAULT_PORT) + args = parser.parse_args() + + runtime = ToolRuntime(args.config, args.workspace) + _Handler.runtime = runtime + # Threading, because a turn's tool calls arrive as one request but the + # health poll must stay answerable while a long shell command runs. + server = ThreadingHTTPServer((args.host, args.port), _Handler) + names = [t.get('function', {}).get('name') for t in runtime.tools()] + llm_note = 'llm configured' if runtime.has_llm else 'no llm (read_file.abbreviate withdrawn)' + sys.stderr.write(f'[tool_server] ready on {args.host}:{args.port}, {llm_note}, ' + f'{len(names)} tools: {names}\n') + sys.stderr.flush() + server.serve_forever() + + +if __name__ == '__main__': + main() diff --git a/cookbook/rsi/agentic/README.md b/cookbook/rsi/agentic/README.md new file mode 100644 index 000000000..ebb782979 --- /dev/null +++ b/cookbook/rsi/agentic/README.md @@ -0,0 +1,54 @@ +# RSI Agentic Self-Play + +## 前置条件 + +- 沙箱环境已启动（AgentENV / e2b），模板由 `cookbook/rl/rsi_agentic/sandbox_server/install.sh` 构建 +- `AENV_API_URL` 和 `AENV_TEMPLATE` 环境变量已设置 +- ms-agent 配置文件就绪（默认 `cookbook/rl/rsi_agentic/rsi_agent.yaml`） + +## Step 1: 生成任务 + +```bash +python cookbook/rsi/agentic/challenge.py \ + --keep-target 200 \ + --sandbox-template $AENV_TEMPLATE \ + --sandbox-api-url $AENV_API_URL \ + --sampler-gpus 4 +``` + +产出：`output/rsi_agentic/challenge_flows.jsonl` + +每行一个任务：`{id, query, check_script, n_pass, n_rollouts, keywords, seeded}` + +可选参数： +- `--seed-file seeds.jsonl` 用已有 trajectory 做起点 +- `--keywords-n 0` 关闭关键词库 +- `--solver-rollouts 4` 难度过滤尝试次数 +- `--max-turns 20` round1 最大工具调用轮数 + +## Step 2: 训练（GRPO） + +```bash +AENV_API_URL=http://... \ +AENV_TEMPLATE=twinkle-rsi-msagent \ +RSI_TASKS=output/rsi_agentic/challenge_flows.jsonl \ + python cookbook/rl/rsi_agentic/rsi_agentic_grpo.py \ + --model-id ms://Qwen/Qwen3-4B \ + --model-gpus 4 --sampler-gpus 4 +``` + +训练脚本自动识别 `check_script` 格式，在沙箱中跑检查脚本评分（exit 0 = 1.0）。 + +## 流程总结 + +``` +challenge.py rsi_agentic_grpo.py +┌─────────────────────┐ ┌─────────────────────┐ +│ 1. 选方向+关键词 │ │ 1. 读 flows │ +│ 2. 模型在沙箱做事 │ │ 2. 起沙箱 │ +│ 3. 模型写检查脚本 │ flows │ 3. solver 多轮做题 │ +│ 4. 跑检查（验证） │ ──────► │ 4. 跑 check_script │ +│ 5. 模型写题目描述 │ │ 5. GRPO 训练 │ +│ 6. 难度过滤 │ └─────────────────────┘ +└─────────────────────┘ +``` diff --git a/cookbook/rsi/agentic/challenge.py b/cookbook/rsi/agentic/challenge.py new file mode 100644 index 000000000..878bfd569 --- /dev/null +++ b/cookbook/rsi/agentic/challenge.py @@ -0,0 +1,302 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""RSI self-play, agentic half: generate training tasks by doing them first. + +One model plays both roles. It first acts as an agent in a sandbox (multi-turn +with tools), producing a trajectory and final workspace state. Then it writes a +check script that verifies the end state, and finally describes the task as a +problem statement. The same model then attempts the problem multiple times, and +only problems it solves *sometimes* are kept. + +The machinery lives in :mod:`twinkle_agentic.challenger`; the prompts live in +``prompts.py`` next to this file. What is here is the wiring: which model, how +many, the sandbox connection, and where the output goes. + +Output format (one JSONL line per task): + + --out-flows {id, query, check_script, n_pass, n_rollouts, keywords, seeded} + +Run it as a Ray job (sampler only, no trainer):: + + python cookbook/rsi/agentic/challenge.py --keep-target 200 +""" +import argparse +import json +import os +import sys + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import SamplingParams, user_data_get +from twinkle.sampler import vLLMSampler +from twinkle_agentic.challenger import AgenticChallenger, KeywordStore +from twinkle_agentic.envs import EnvTool +from twinkle_agentic.rollout import build_rollout +from twinkle_agentic.tools.tool_manager import ToolManager + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from prompts import CATEGORIES, CATEGORY_DESC, agentic_prompts # noqa: E402 + +logger = get_logger() + + +def parse_args(): + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + # Model + p.add_argument('--model-id', default='ms://Qwen/Qwen3-4B') + p.add_argument('--template', default='Template', + help='template class in twinkle.template') + p.add_argument('--sampler-gpus', type=int, default=4) + p.add_argument('--max-model-len', type=int, default=32768) + + # Generation control + p.add_argument('--keep-target', type=int, default=200, + help='how many tasks to keep; generation stops once reached') + p.add_argument('--batch-size', type=int, default=0, + help='tasks per yielded batch (0 = one batch of --keep-target)') + p.add_argument('--max-proposals-per-round', type=int, default=64, + help='max proposals per round (serial, so keep moderate)') + p.add_argument('--seed-file', default='', help='seed jsonl with query field') + p.add_argument('--seed-mix-prob', type=float, default=0.5) + + # Sampling params for round 1 (proposing) + p.add_argument('--propose-temp', type=float, default=1.0) + p.add_argument('--propose-max-tokens', type=int, default=4096) + p.add_argument('--max-turns', type=int, default=20, + help='max tool-calling turns for round 1') + + # Problem statement + p.add_argument('--problem-max-chars', type=int, default=8192) + + # Keywords + p.add_argument('--keywords-n', type=int, default=128, + help='per-category refill target; 0 disables keyword bank') + p.add_argument('--keyword-db', default='output/rsi_agentic/keywords.jsonl') + p.add_argument('--keyword-gen-calls', type=int, default=8) + p.add_argument('--keyword-refill-tries', type=int, default=2) + p.add_argument('--keyword-temp', type=float, default=1.3) + p.add_argument('--keyword-max-tokens', type=int, default=1024) + p.add_argument('--single-kw-prob', type=float, default=0.1) + p.add_argument('--combo-arity', default='triple', choices=['triple', 'mix']) + p.add_argument('--arity-weights', default='', + help="'w1,w2,w3' for --combo-arity mix (empty = uniform)") + + # Difficulty filter + p.add_argument('--solver-rollouts', type=int, default=4) + p.add_argument('--solver-temp', type=float, default=1.0) + p.add_argument('--solver-max-tokens', type=int, default=4096) + p.add_argument('--solver-max-turns', type=int, default=20) + p.add_argument('--keep-min-pass', type=int, default=1) + p.add_argument('--keep-max-margin', type=int, default=1) + + # Sandbox + p.add_argument('--sandbox-template', default='', + help='AgentENV/e2b template name (required)') + p.add_argument('--sandbox-api-url', default='', + help='AgentENV server URL (or AENV_API_URL env var)') + p.add_argument('--agent-config', default='cookbook/rl/rsi_agentic/rsi_agent.yaml', + help='ms-agent yaml for the sandbox tool server') + p.add_argument('--sandbox-timeout', type=int, default=900) + p.add_argument('--workspace', default='/workspace', + help='working directory inside the sandbox') + + # Output + p.add_argument('--random-seed', type=int, default=0) + p.add_argument('--out-flows', default='output/rsi_agentic/challenge_flows.jsonl') + p.add_argument('--dump-rejected', default='output/rsi_agentic/challenge_rejected.jsonl') + p.add_argument('--no-sort-by-difficulty', action='store_true') + return p.parse_args() + + +def build_env(args): + """Create the long-lived sandbox environment.""" + sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), + '..', '..', 'rl', 'rsi_agentic')) + from remote_tool_env import RemoteMsAgentToolEnv # noqa: E402 + + template = args.sandbox_template or os.environ.get('AENV_TEMPLATE', '') + api_url = args.sandbox_api_url or os.environ.get('AENV_API_URL', '') + if not template: + raise SystemExit('[challenge] --sandbox-template or AENV_TEMPLATE is required') + if not api_url: + raise SystemExit('[challenge] --sandbox-api-url or AENV_API_URL is required') + + env = RemoteMsAgentToolEnv( + template=template, + config_path=args.agent_config, + api_url=api_url, + workspace=args.workspace, + sandbox_timeout=args.sandbox_timeout, + ) + env.reset() + return env + + +def main(): + args = parse_args() + for path in (args.out_flows, args.dump_rejected, args.keyword_db): + if path: + os.makedirs(os.path.dirname(os.path.abspath(path)) or '.', exist_ok=True) + + # Initialize twinkle (sampler only, no trainer) + twinkle.initialize( + mode='ray', nproc_per_node=args.sampler_gpus, lazy_collect=False, + groups=[DeviceGroup(name='sampler', ranks=list(range(args.sampler_gpus)), + device_type='GPU')]) + sampler = vLLMSampler( + model_id=args.model_id, + engine_args={'gpu_memory_utilization': 0.8, 'max_model_len': args.max_model_len}, + device_mesh=DeviceMesh.from_sizes(world_size=args.sampler_gpus, + dp_size=args.sampler_gpus), + remote_group='sampler', + ) + sampler.set_template(args.template, model_id=args.model_id, enable_thinking=True, + max_length=args.max_model_len) + + import twinkle.template as template_module + template = getattr(template_module, args.template)( + args.model_id, max_length=args.max_model_len, enable_thinking=True) + + # Build sandbox environment + env = build_env(args) + schemas = env.tool_schemas() + tool_manager = ToolManager(EnvTool.from_schemas(env, schemas)) + + # Explorer: multi-turn rollout with sandbox tools + explorer = build_rollout( + sampler, + template=template, + tool_manager=tool_manager, + max_turns=args.max_turns, + sampling_params=SamplingParams(max_tokens=args.propose_max_tokens, num_samples=1, + logprobs=1, temperature=args.propose_temp, top_p=0.95), + ) + + # Sandbox control functions -- use env.runner() which resolves tool names + # (ms-agent registers tools as "server---name") and parses exit codes from + # the marker protocol, so we don't rely on string matching. + runner = env.runner() + + def reset_fn(): + """Clear the sandbox workspace between episodes.""" + runner(f'rm -rf {args.workspace}/* {args.workspace}/.* 2>/dev/null; true', 'shell') + + def run_check_fn(script: str): + """Run a python check script in the sandbox; returns (exit_code, output).""" + return runner(script, 'python') + + def workspace_snapshot_fn(): + """Get a summary of the current workspace state.""" + _, output = runner( + f'find {args.workspace} -type f -printf "%P %s\\n" 2>/dev/null | head -50', + 'shell') + return output or '(empty)' + + # Keywords + store = None + if args.keywords_n > 0: + store = KeywordStore(args.keyword_db, CATEGORIES) + logger.info('[challenge] keyword bank loaded: ' + + ', '.join(f'{c}={len(store.items[c])}' for c in CATEGORIES)) + + # Seeds + seeds = [] + if args.seed_file: + with open(args.seed_file, encoding='utf-8') as f: + for line in f: + if line.strip(): + seeds.append(json.loads(line)) + logger.info(f'[challenge] loaded {len(seeds)} seeds from {args.seed_file}') + + # Rejected log + rejected = open(args.dump_rejected, 'w', encoding='utf-8') if args.dump_rejected else None + + def _reject(record): + if rejected is not None: + rejected.write(json.dumps(record, ensure_ascii=False, default=str) + '\n') + + # Build challenger + prompts = agentic_prompts() + challenger = AgenticChallenger( + prompts, + explorer, + seeds=seeds, + keyword_store=store, + category_desc=CATEGORY_DESC if store else None, + seed_mix_prob=args.seed_mix_prob, + reset_fn=reset_fn, + run_check_fn=run_check_fn, + workspace_snapshot_fn=workspace_snapshot_fn, + combo_arity=args.combo_arity, + arity_weights=[float(x) for x in args.arity_weights.split(',')] if args.arity_weights + else None, + single_kw_prob=args.single_kw_prob, + keyword_refill_target=args.keywords_n, + keyword_gen_calls=args.keyword_gen_calls, + keyword_refill_tries=args.keyword_refill_tries, + keyword_params=SamplingParams(max_tokens=args.keyword_max_tokens, num_samples=1, + logprobs=1, temperature=args.keyword_temp, top_p=0.98), + min_batch=args.sampler_gpus, + problem_max_chars=args.problem_max_chars, + reject_sink=_reject, + max_proposals_per_round=args.max_proposals_per_round, + solver_rollouts=args.solver_rollouts, + keep_min_pass=args.keep_min_pass, + keep_max_pass_margin=args.keep_max_margin, + solver_params=SamplingParams(max_tokens=args.solver_max_tokens, num_samples=1, + logprobs=1, temperature=args.solver_temp, top_p=0.95), + seed=args.random_seed, + ) + + # Generate + batch_size = args.batch_size or args.keep_target + kept = [] + for batch in challenger(batch_size=batch_size, total=args.keep_target): + kept.extend(batch) + logger.info(f'[challenge] kept {len(kept)}/{args.keep_target} so far; ' + f'stats {challenger.stats}') + if rejected is not None: + rejected.close() + + if store is not None: + challenger.expand_hard_keywords() + store.save() + logger.info('[challenge] keyword bank saved -> ' + args.keyword_db) + + # Sort by difficulty (hardest last) + if not args.no_sort_by_difficulty: + kept.sort(key=lambda t: -(user_data_get(t.get('user_data'), 'n_pass', 0) or 0)) + + # Write output + write_flows(kept, args) + logger.info(f'[challenge] wrote {len(kept)} tasks -> {args.out_flows}') + dist = {} + for task in kept: + n = user_data_get(task.get('user_data'), 'n_pass', 0) + dist[n] = dist.get(n, 0) + 1 + logger.info(f'[challenge] pass-count distribution: {dict(sorted(dist.items()))}') + + env.close() + + +def write_flows(kept, args): + """Write one flow per task.""" + with open(args.out_flows, 'w', encoding='utf-8') as f: + for i, task in enumerate(kept): + data = task.get('user_data') + messages = task.get('messages') or [] + query = next((m['content'] for m in messages if m.get('role') == 'user'), '') + flow = { + 'id': f'ag_{i:06d}', + 'query': query, + 'check_script': user_data_get(data, 'check_script', ''), + 'n_pass': user_data_get(data, 'n_pass'), + 'n_rollouts': user_data_get(data, 'n_rollouts'), + 'keywords': user_data_get(data, 'keywords', []), + 'seeded': user_data_get(data, 'seeded', False), + } + f.write(json.dumps(flow, ensure_ascii=False) + '\n') + + +if __name__ == '__main__': + main() diff --git a/cookbook/rsi/agentic/prompts.py b/cookbook/rsi/agentic/prompts.py new file mode 100644 index 000000000..8ff4356e9 --- /dev/null +++ b/cookbook/rsi/agentic/prompts.py @@ -0,0 +1,171 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Prompts for the agentic challenger. + +The agentic challenger works in three rounds: + Round 1: model acts as an agent in a sandbox (multi-turn with tools), + producing a tool-call chain and final workspace state. + Round 2a: model sees the trajectory and writes a python check script + that asserts properties of the final state. + Round 2b: model sees trajectory + checks and writes a problem statement + that someone else would need to follow to reproduce the result. + +The check script is run against the sandbox immediately after round 1 to verify +it passes; any task whose own checks fail is thrown away. This is the agentic +analogue of the code challenger running the reference solution against its own +asserts. + +Keyword categories and directions are configurable. The defaults below exercise: + - filesystem: creating, moving, reading, transforming files + - scripting: writing python/shell scripts that produce output + - data: CSV/JSON/text parsing and aggregation +""" +from twinkle_agentic.challenger import AgenticPrompts + +# ── Keyword categories (analogous to code's algorithm/computer/noncs) ────── + +CATEGORIES = ['filesystem', 'scripting', 'data'] + +CATEGORY_DESC = { + 'filesystem': 'file and directory manipulation tasks: creating directory trees, ' + 'moving/renaming files by pattern, finding files by content, ' + 'generating structured text files', + 'scripting': 'tasks that require writing a script (python or shell) whose output ' + 'or side effects are the goal: number crunching, text transformation, ' + 'format conversion, small utilities', + 'data': 'tasks involving structured data: parsing CSV/JSON/YAML, aggregating ' + 'rows, filtering records, joining multiple files, producing summary ' + 'reports or reformatted output', +} + +# ── Round 1: model acts in sandbox ───────────────────────────────────────── + +SYSTEM = ( + 'You are an expert developer working in an empty directory with a shell and ' + 'python. Your job is to do something interesting and non-trivial based on ' + 'the direction given below. Use the tools available to you (shell commands, ' + 'python scripts, file operations) to produce a meaningful end state: files ' + 'with content, computed outputs, structured data.\n\n' + 'Requirements:\n' + '- Work entirely within the current directory (do not use /tmp or ~).\n' + '- Do not use the network.\n' + '- Make sure the end state is deterministic: the same steps always produce ' + 'the same files with the same content.\n' + '- Do at least 2-3 distinct steps, not just one command.\n' + '- When you are satisfied with the result, stop calling tools and say ' + '"Done." as your final message.' +) + +FROM_SCRATCH = ( + 'Do something interesting and non-trivial in the current empty directory. ' + 'Create files, write scripts, process data -- whatever demonstrates ' + 'competent use of the tools. Aim for 2-4 steps that build on each other.' +) + +FROM_SEED = ( + 'Here is an example of the kind of task we want:\n\n{seed}\n\n' + 'Do something in the same spirit but on a different subject. Change what ' + 'is produced and how, not just the names. Work in the current empty directory.' +) + +FROM_KEYWORDS = ( + 'Your direction for this task:\n{keywords}\n\n' + 'Do something interesting that exercises the topics above. Work in the ' + 'current empty directory, producing files and/or computed output.' +) + +FROM_SEED_KEYWORDS = ( + 'Here is an example task for inspiration:\n\n{seed}\n\n' + 'Your direction keywords:\n{keywords}\n\n' + 'Do something that combines the spirit of the example with the keyword ' + 'topics. Work in the current empty directory.' +) + +# ── Round 2a: write check script ─────────────────────────────────────────── + +CHECK_SYSTEM = ( + 'You are a test engineer. Given a record of what an agent did in a directory ' + 'and the resulting state, write a python script that ASSERTS properties of ' + 'the end state. The script will be run in the same directory the agent worked ' + 'in.\n\n' + 'Rules:\n' + '- Use only the standard library (os, json, csv, re, pathlib, etc.).\n' + '- Write 2-6 assert statements that verify the most important outcomes.\n' + '- Each assert should check something observable: file existence, file ' + 'content, computed values, directory structure.\n' + '- The script must exit 0 when all assertions hold and non-zero otherwise.\n' + '- Do NOT import anything that is not in the python standard library.\n' + '- Do NOT use the network or read from outside the working directory.\n' + '- Return ONLY a fenced python code block, no prose.' +) + +CHECK_USER = ( + 'Here is what the agent did:\n\n{trajectory}\n\n' + 'Here is the final state of the working directory:\n\n{final_state}\n\n' + 'Write a python check script (fenced code block) that asserts the key ' + 'properties of this end state. 2-6 assertions.' +) + +# ── Round 2b: write problem statement ───────────────────────────────────── + +PROBLEM_SYSTEM = ( + 'You write task descriptions for an AI agent. Given a record of what was ' + 'done and the checks that verify it, write a clear problem statement that ' + 'another agent would need to follow to reproduce the same end state.\n\n' + 'Rules:\n' + '- State exactly what files must exist and what they must contain.\n' + '- Be specific about formats, names, and expected values.\n' + '- Do NOT reveal the solution steps -- only describe the desired end state.\n' + '- Do NOT mention the checks or how verification works.\n' + '- The statement must be self-contained: no references to prior context.\n' + '- Keep it concise: 50-300 words.\n' + '- Return ONLY the problem statement as plain text, no code fences.' +) + +PROBLEM_USER = ( + 'Here is what the agent did:\n\n{trajectory}\n\n' + 'Here are the check assertions that verify the end state:\n\n' + '```python\n{checks}\n```\n\n' + 'Write a problem statement (plain text, 50-300 words) describing what ' + 'another agent must produce to pass these checks. Do not reveal the ' + 'solution steps.' +) + +# ── Keyword generation ───────────────────────────────────────────────────── + +KEYWORD_SYSTEM = ( + 'You generate diverse topic keywords for training an AI agent that works ' + 'with files, scripts, and data in a local directory.' +) + +KEYWORD_USER = ( + 'List {k} diverse, specific topic keywords for the following category:\n' + '{desc}\n\n' + 'Return one keyword per line, no numbering, no explanation. ' + 'Each should be 2-5 words, concrete enough to inspire a specific task.' +) + +KEYWORD_EXPAND_USER = ( + 'The keyword "{kw}" produced a very hard task. List {m} related keywords ' + 'in the same domain that might produce similarly challenging but different ' + 'tasks. One per line, no numbering.' +) + + +# ── Factory ──────────────────────────────────────────────────────────────── + +def agentic_prompts() -> AgenticPrompts: + """Assemble all strings into the object the challenger takes.""" + return AgenticPrompts( + system=SYSTEM, + from_scratch=FROM_SCRATCH, + from_seed=FROM_SEED, + from_keywords=FROM_KEYWORDS, + from_seed_keywords=FROM_SEED_KEYWORDS, + check_system=CHECK_SYSTEM, + check_user=CHECK_USER, + problem_system=PROBLEM_SYSTEM, + problem_user=PROBLEM_USER, + keyword_system=KEYWORD_SYSTEM, + keyword_user=KEYWORD_USER, + keyword_expand_user=KEYWORD_EXPAND_USER, + ) diff --git a/cookbook/rsi/code/challenge.py b/cookbook/rsi/code/challenge.py new file mode 100644 index 000000000..b9c930f00 --- /dev/null +++ b/cookbook/rsi/code/challenge.py @@ -0,0 +1,260 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""RSI self-play, code half: generate training problems with a local sampler. + +One model plays both roles. It writes a problem plus a reference solution; the +solution is executed to turn the problem's check expressions into asserts; then +the same model attempts the problem several times and only problems it solves +*sometimes* are kept -- an all-pass or all-fail group gives GRPO nothing to learn +from. + +The machinery lives in :mod:`twinkle_agentic.challenger`; the prompts live in +``prompts.py`` next to this file. What is here is the wiring: which model, how +many, where the output goes. + +Output is what ``rsi_rl`` reads directly, no prepare/refine stage in between: + + --out-flows {id, system, query, tools, rounds:[code round]} + --out-tests {id, test_list, test_setup_code} + +Run it as a Ray job (sampler only, no trainer):: + + python cookbook/rsi/code/challenge.py --keep-target 500 --seed-file seeds.jsonl +""" +import argparse +import json +import os +import sys + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import SamplingParams, user_data_get +from twinkle.sampler import vLLMSampler +from twinkle_agentic.challenger import CodeChallenger, KeywordStore, load_seeds +from twinkle_agentic.rollout import build_rollout +from twinkle_agentic.tools.tool_manager import ToolManager + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from prompts import CATEGORIES, CATEGORY_DESC, code_prompts # noqa: E402 + +logger = get_logger() + + +def parse_args(): + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + # Defaults are the ones the previous env-var script shipped with, so a run + # started without flags produces what earlier iterations produced. + p.add_argument('--model-id', default='ms://Qwen/Qwen3-4B') + p.add_argument('--template', default='Template', + help='template class in twinkle.template; the text one for Qwen3-4B') + p.add_argument('--sampler-gpus', type=int, default=4) + p.add_argument('--max-model-len', type=int, default=16384) + + p.add_argument('--keep-target', type=int, default=500, + help='how many problems to keep; generation stops once reached') + p.add_argument('--batch-size', type=int, default=0, + help='problems per written batch (0 = one batch of --keep-target)') + p.add_argument('--max-proposals-per-round', type=int, default=2000, + help='ceiling on one proposing round, i.e. one batched generate') + p.add_argument('--seed-file', default='', help='seed jsonl with query [+ code]') + p.add_argument('--seed-mix-prob', type=float, default=0.5) + p.add_argument('--no-two-step', action='store_true', + help='never take the two-call path, even for seeds carrying code') + + p.add_argument('--propose-temp', type=float, default=1.1) + p.add_argument('--propose-max-tokens', type=int, default=8192) + p.add_argument('--problem-max-chars', type=int, default=4000) + + p.add_argument('--keywords-n', type=int, default=128, + help='per-category refill target; 0 disables the keyword bank') + p.add_argument('--keyword-db', default='output/rsi/keywords.jsonl') + p.add_argument('--keyword-gen-calls', type=int, default=8) + p.add_argument('--keyword-refill-tries', type=int, default=2) + p.add_argument('--keyword-temp', type=float, default=1.3) + p.add_argument('--keyword-max-tokens', type=int, default=1024) + p.add_argument('--single-kw-prob', type=float, default=0.1) + p.add_argument('--combo-arity', default='triple', choices=['triple', 'mix']) + p.add_argument('--arity-weights', default='', + help="'w1,w2,w3' for --combo-arity mix (empty = uniform)") + p.add_argument('--low-pass-expand', type=int, default=0, + help='expand topics of problems solved at most this many times') + p.add_argument('--expand-per-kw', type=int, default=8) + p.add_argument('--expand-max-kws', type=int, default=32) + + p.add_argument('--solver-rollouts', type=int, default=8) + p.add_argument('--solver-temp', type=float, default=1.0) + p.add_argument('--solver-max-tokens', type=int, default=2048) + p.add_argument('--keep-min-pass', type=int, default=1) + p.add_argument('--keep-max-margin', type=int, default=1) + + p.add_argument('--sandbox-timeout', type=int, default=30) + p.add_argument('--max-checks', type=int, default=6) + p.add_argument('--keep-constant-answer', action='store_true', + help='keep problems where one constant satisfies every assert') + p.add_argument('--no-sort-by-difficulty', action='store_true', + help='write in generation order instead of hardest-last') + p.add_argument('--random-seed', type=int, default=0) + + p.add_argument('--out-flows', default='output/rsi/challenge_flows.jsonl') + p.add_argument('--out-tests', default='output/rsi/challenge_tests.jsonl') + p.add_argument('--dump-rejected', default='output/rsi/challenge_rejected.jsonl') + return p.parse_args() + + +def main(): + args = parse_args() + for path in (args.out_flows, args.out_tests, args.dump_rejected, args.keyword_db): + if path: + os.makedirs(os.path.dirname(os.path.abspath(path)) or '.', exist_ok=True) + + twinkle.initialize( + mode='ray', nproc_per_node=args.sampler_gpus, lazy_collect=False, + groups=[DeviceGroup(name='sampler', ranks=list(range(args.sampler_gpus)), + device_type='GPU')]) + sampler = vLLMSampler( + model_id=args.model_id, + engine_args={'gpu_memory_utilization': 0.8, 'max_model_len': args.max_model_len}, + device_mesh=DeviceMesh.from_sizes(world_size=args.sampler_gpus, + dp_size=args.sampler_gpus), + remote_group='sampler', + ) + sampler.set_template(args.template, model_id=args.model_id, enable_thinking=True, + max_length=args.max_model_len) + + import twinkle.template as template_module + template = getattr(template_module, args.template)( + args.model_id, max_length=args.max_model_len, enable_thinking=True) + # Single-turn generation, but through the same rollout the RL loop uses, so a + # challenger that should be allowed to run code while inventing only needs a + # tool manager here rather than a different code path. + explorer = build_rollout( + sampler, + template=template, + tool_manager=ToolManager([]), + max_turns=1, + sampling_params=SamplingParams(max_tokens=args.propose_max_tokens, num_samples=1, + logprobs=1, temperature=args.propose_temp, top_p=0.95), + ) + + store = None + if args.keywords_n > 0: + store = KeywordStore(args.keyword_db, CATEGORIES) + logger.info('[challenge] keyword bank loaded: ' + + ', '.join(f'{c}={len(store.items[c])}' for c in CATEGORIES)) + + seeds = load_seeds(args.seed_file) + logger.info(f'[challenge] seeds: {len(seeds)} from {args.seed_file!r} ' + f'(seed_mix_prob={args.seed_mix_prob if seeds else 0.0})') + + rejected = open(args.dump_rejected, 'w', encoding='utf-8') if args.dump_rejected else None + + def _reject(record): + if rejected is not None: + rejected.write(json.dumps(record, ensure_ascii=False, default=str) + '\n') + + challenger = CodeChallenger( + code_prompts(), + explorer, + seeds=seeds, + keyword_store=store, + category_desc=CATEGORY_DESC if store else None, + seed_mix_prob=args.seed_mix_prob, + two_step=not args.no_two_step, + combo_arity=args.combo_arity, + arity_weights=[float(x) for x in args.arity_weights.split(',')] if args.arity_weights + else None, + single_kw_prob=args.single_kw_prob, + keyword_refill_target=args.keywords_n, + keyword_gen_calls=args.keyword_gen_calls, + keyword_refill_tries=args.keyword_refill_tries, + keyword_params=SamplingParams(max_tokens=args.keyword_max_tokens, num_samples=1, + logprobs=1, temperature=args.keyword_temp, top_p=0.98), + # A batch smaller than the sampler's data-parallel width leaves workers idle. + min_batch=args.sampler_gpus, + problem_max_chars=args.problem_max_chars, + max_checks=args.max_checks, + sandbox_timeout=args.sandbox_timeout, + drop_constant_answer=not args.keep_constant_answer, + low_pass_expand=args.low_pass_expand, + expand_per_kw=args.expand_per_kw, + expand_max_kws=args.expand_max_kws, + reject_sink=_reject, + max_proposals_per_round=args.max_proposals_per_round, + solver_rollouts=args.solver_rollouts, + keep_min_pass=args.keep_min_pass, + keep_max_pass_margin=args.keep_max_margin, + solver_params=SamplingParams(max_tokens=args.solver_max_tokens, num_samples=1, + logprobs=1, temperature=args.solver_temp, top_p=0.95), + seed=args.random_seed, + ) + + batch_size = args.batch_size or args.keep_target + kept = [] + for batch in challenger(batch_size=batch_size, total=args.keep_target): + kept.extend(batch) + logger.info(f'[challenge] kept {len(kept)}/{args.keep_target} so far; ' + f'proposal stats {challenger.stats}') + if rejected is not None: + rejected.close() + + if store is not None: + challenger.expand_hard_keywords() + store.save() + logger.info('[challenge] keyword bank saved: ' + + ', '.join(f'{c}={len(store.items[c])}' for c in CATEGORIES) + + f' -> {args.keyword_db}') + + # File order = decreasing pass count, i.e. hardest last. The fixed-pool + # validation in rsi_rl relies on this ordering. + if not args.no_sort_by_difficulty: + kept.sort(key=lambda t: -(user_data_get(t.get('user_data'), 'n_pass', 0) or 0)) + + write_flows(kept, args) + logger.info(f'[challenge] wrote {len(kept)} problems -> {args.out_flows} + {args.out_tests}') + dist = {} + for task in kept: + n = user_data_get(task.get('user_data'), 'n_pass', 0) + dist[n] = dist.get(n, 0) + 1 + logger.info(f'[challenge] kept pass-count distribution: {dict(sorted(dist.items()))}') + + +def write_flows(kept, args): + """Write the two files rsi_rl reads: one flow and one test row per problem.""" + with open(args.out_flows, 'w', encoding='utf-8') as ff, \ + open(args.out_tests, 'w', encoding='utf-8') as ft: + for i, task in enumerate(kept): + data = task.get('user_data') + cid = f'ch_{i:06d}' + messages = task.get('messages') or [] + system = next((m for m in messages if m.get('role') == 'system'), None) + query = next((m for m in messages if m.get('role') == 'user'), None) + flow = { + 'id': cid, + 'system': system, + 'query': query, + 'tools': [], + # Difficulty audit, ignored by rsi_rl: how many solver attempts + # passed, so a stored flow can be analysed without re-running. + 'n_pass': user_data_get(data, 'n_pass'), + 'n_rollouts': user_data_get(data, 'n_rollouts'), + 'keywords': user_data_get(data, 'keywords', []), + 'seeded': user_data_get(data, 'seeded', False), + 'two_step': user_data_get(data, 'two_step', False), + 'rounds': [{ + 'intent': 'solve the problem', + 'type': 'code', + 'tool_call': None, + # The challenger's own passing solution; OPSD reads this. + 'code': user_data_get(data, 'solution', ''), + 'result': '', + 'reward_method': 'rubric', + }], + } + ff.write(json.dumps(flow, ensure_ascii=False) + '\n') + ft.write(json.dumps({'id': cid, + 'test_list': user_data_get(data, 'asserts', []), + 'test_setup_code': ''}, ensure_ascii=False) + '\n') + + +if __name__ == '__main__': + main() diff --git a/cookbook/rsi/code/prompts.py b/cookbook/rsi/code/prompts.py new file mode 100644 index 000000000..6e1388c49 --- /dev/null +++ b/cookbook/rsi/code/prompts.py @@ -0,0 +1,170 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Prompts for the code challenger. + +Every string the model sees during self-play for code tasks, in one file, +because the prompt *is* the experiment: two runs that differ here are not +comparable, and a run has to be able to say which wording produced its data. +:class:`twinkle_agentic.challenger.CodePrompts` holds no defaults for exactly +that reason. + +Text carried over verbatim from the previous ``rsi_challenge.py``, including the +findings recorded next to it, so numbers from earlier runs stay comparable. +""" +from twinkle_agentic.challenger import CodePrompts + +# Categories of the keyword bank. One keyword is drawn from each per proposal, +# so the challenger has to bridge an algorithm, a computing concept and a +# real-world domain instead of falling back on palindromes and bracket matching. +CATEGORIES = ('algorithm', 'computer', 'noncs') + +CATEGORY_DESC = { + 'algorithm': 'algorithmic techniques and paradigms (e.g. dynamic programming, binary ' + 'search, union-find, Dijkstra, backtracking, segment trees, greedy, ' + 'divide and conquer, sliding window ...)', + 'computer': 'computer-science / computing concepts that are NOT algorithms per se ' + '(e.g. hash maps, tries, LRU cache, bitsets, regular expressions, base ' + 'conversion, finite state machines, serialization, parsing, memoization ...)', + 'noncs': 'real-world domains OUTSIDE computer science, used to give a problem flavor ' + '(e.g. biology, finance, chemistry, logistics, music, cooking, sports, ' + 'astronomy, geography, linguistics ...)', +} + +# The output contract. It names the four keys parse_challenge() reads back, and +# the "do NOT write the expected value" line is what makes the ground truth come +# from execution rather than from the model's own guess about its own code. +CHALLENGER_SYSTEM = ( + 'You design self-contained Python coding problems for training another model.\n' + 'A good problem: (1) is solvable from its statement ALONE with no external files, ' + 'network, images, or hidden context; (2) has ONE clear entry function; (3) is ' + 'deterministic (same input -> same output), no randomness, no wall-clock, no threads; ' + '(4) is neither trivial nor impossible for a mid-size model.\n' + 'You will also write the reference solution. We will EXECUTE it to obtain the ' + 'ground-truth outputs, so your solution must be correct and runnable as-is.\n' + 'Return ONLY one JSON object, no prose around it, with keys:\n' + ' "problem": the statement shown to the solver (describe the function name, its ' + 'inputs and expected behavior; do NOT include the solution).\n' + ' "solution": the reference implementation as plain Python source (no markdown fence).\n' + ' "entry": the entry function name.\n' + ' "checks": a list of 3-6 Python expressions calling the entry function on concrete ' + 'inputs (e.g. "solve([1,2,3])"); each must be evaluable after running the solution. ' + 'Do NOT write the expected value — we compute it by running your solution.' +) + +FROM_SCRATCH = ( + 'Create ONE new Python coding problem now. Vary the topic freely ' + '(strings, arrays, math, greedy, DP, parsing, simulation ...).' +) + +FROM_SEED = ( + 'Here is a seed problem. Create ONE NEW problem that is a meaningful VARIANT of it ' + '(change the twist, constraints, or data shape — not just renaming), keeping it ' + 'self-contained and deterministic.\n\n[seed]\n{seed}' +) + +FROM_KEYWORDS = ( + 'Create ONE new Python coding problem now. Draw inspiration from the following ' + 'topic(s) and combine them creatively into a single coherent problem:\n{keywords}\n' + 'You may use each topic directly or bend it loosely; combine with any data shape ' + '(strings, arrays, grids, trees, numbers, parsing, simulation ...). Make it require ' + 'real thought, not a one-liner, and keep it self-contained and deterministic.' +) + +# Seed AND keywords together. The seed is deliberately framed as inspiration only, +# not as something to produce a variant of: the point is to pull the generated +# problems toward the shape of public benchmark items (short statement, one plain +# task) while the keywords keep supplying topical variety. +FROM_SEED_KEYWORDS = ( + 'Create ONE new Python coding problem now. Use the problem below only as a ' + 'STARTING POINT for inspiration — you do NOT have to keep its task, and the new ' + 'problem does NOT need to be a variant of it.\n\n[inspiration]\n{seed}\n\n' + 'Also draw on the following topic(s), combining them into a single coherent ' + 'problem:\n{keywords}\n' + 'Make it require real thought, not a one-liner, and keep it self-contained and ' + 'deterministic.' +) + +# ── two-step proposing ────────────────────────────────────────────────────── +# The difficulty comes from adding a layer on top of a real, runnable reference +# solution, not from imagining a hard problem outright; splitting into two calls +# (write the harder code, THEN describe it) keeps the statement and the ground +# truth consistent, which a single call does not. Measured on 40 MBPP seeds +# against the single-call seed+keywords prompt: kept-rate 25% vs 15%, +# constant-answer problems 4 vs 7, similarity to the seed 0.42. +TWO_STEP_SYSTEM = 'You are an expert Python programmer.' + +TWO_STEP_SOLUTION = ( + 'Below is a coding problem and its reference solution.\n\n' + '[problem]\n{seed}\n\n[reference solution]\n{code}\n\n' + 'Write a MORE COMPLEX Python function that keeps the idea of the reference solution ' + 'as one step and builds a harder computation around it (extra pass, different data ' + 'structure, an added rule), in the direction of these topic(s):\n{keywords}\n' + 'Requirements: deterministic, self-contained, no randomness, no I/O, one clear entry ' + 'function. Output ONLY the code in a single ```python block, no explanation.' +) + +# Showing the seed here pulls the wording back toward the MBPP task family +# (similarity 0.32 -> 0.42). The solution is NOT taken from this JSON -- the +# challenger overwrites it with the code the first call produced, so the ground +# truth matches what was actually executed. +TWO_STEP_PROBLEM = ( + 'Here is a Python function.\n\n```python\n{code}\n```\n\n' + 'It was written as a harder follow-up to this exercise:\n\n[original exercise]\n' + '{seed}\n\nand it was pushed in the direction of these topic(s):\n{keywords}\n\n' + 'Write the problem statement that the function above is the answer to, as if it were ' + 'a coding exercise in the same series as the original: name the entry function, ' + 'describe its inputs and the exact behaviour expected, and do NOT reveal the ' + 'implementation. Phrase it as plainly and briefly as the original exercise.\n' + 'Return ONLY one JSON object, no prose around it, with keys:\n' + ' "problem": the statement shown to the solver.\n' + ' "entry": the entry function name.\n' + ' "checks": a list of 3-6 Python expressions calling the entry function on ' + 'concrete inputs; each must be evaluable after running the function above. Do NOT ' + 'write the expected value.\n' + 'The "solution" is already known, so do not include it.' +) + +# ── keyword bank ──────────────────────────────────────────────────────────── +KEYWORD_SYSTEM = 'You brainstorm diverse topics for a Python coding-problem generator.' + +KEYWORD_USER = ( + 'List {k} DISTINCT and SPECIFIC topics from this category: {desc}\n' + 'Be creative and concrete; avoid vague umbrella words. ' + 'Return ONLY a JSON array of short strings, nothing else.' +) + +KEYWORD_EXPAND_USER = ( + 'The topic "{kw}" turned out to seed genuinely HARD problems. List {m} MORE distinct, ' + 'specific topics in the SAME family/domain as "{kw}" that could seed similarly ' + 'challenging Python problems. Return ONLY a JSON array of short strings, nothing else.' +) + +# ── solver ────────────────────────────────────────────────────────────────── +# Used both to measure difficulty and, as the system half, as the system prompt +# of the task that gets stored: training against a different one than the +# difficulty measurement used would make the measurement mean nothing. +SOLVER_SYSTEM = 'You are an expert Python programmer.' + +SOLVER_USER = ( + '{problem}\n\n' + 'Write the complete Python solution. Put the final code in a single ```python fenced ' + 'block. Define the exact function name required by the problem.' +) + + +def code_prompts() -> CodePrompts: + """Assemble the strings above into the object the challenger takes.""" + return CodePrompts( + system=CHALLENGER_SYSTEM, + from_scratch=FROM_SCRATCH, + from_seed=FROM_SEED, + from_keywords=FROM_KEYWORDS, + from_seed_keywords=FROM_SEED_KEYWORDS, + two_step_system=TWO_STEP_SYSTEM, + two_step_solution=TWO_STEP_SOLUTION, + two_step_problem=TWO_STEP_PROBLEM, + keyword_system=KEYWORD_SYSTEM, + keyword_user=KEYWORD_USER, + keyword_expand_user=KEYWORD_EXPAND_USER, + solver_system=SOLVER_SYSTEM, + solver_user=SOLVER_USER, + ) diff --git a/cookbook/rsi/prepare.py b/cookbook/rsi/prepare.py new file mode 100644 index 000000000..af8696b13 --- /dev/null +++ b/cookbook/rsi/prepare.py @@ -0,0 +1,184 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Clean a raw dataset into seeds for the RSI challengers. + +Reads a raw source, runs it through the ``twinkle_agentic.preprocessor`` +pipeline in parallel, and writes the surviving rows. What comes out is seed +material: ``cookbook/rsi/code/challenge.py`` and the agentic challenger take it +as a pool to draw inspiration from, so anything junk in here becomes junk the +challenger imitates. + +Usage +----- + # multi-turn tool-calling data (ToolACE and friends) + python cookbook/rsi/prepare.py --input ms://... --output output/rsi/agentic_seeds.jsonl + + # pure code data (MBPP and friends) + python cookbook/rsi/prepare.py --input mbpp.jsonl --output output/rsi/code_seeds.jsonl \ + --no-normalize-tool-calls + +``--input`` accepts a local ``.jsonl``/``.parquet`` path or an ``ms://`` dataset +id. Every row must expose a ``messages`` list -- the preprocessor keys off it; +ShareGPT ``conversations`` rows are adapted automatically. + +Pipeline +-------- +Core steps, always on, each using the filter's OWN default thresholds (nothing +invented here): + + MessageNormalizer -> MessageSanityFilter -> RefuseFilter -> DeadLoopFilter + -> TokenSoupFilter -> HardFilter + +Optional steps, off unless asked for (each needs extra packages): + --use-lang LanguageFilter (langid, degrades to a heuristic) + --use-datajuicer FixUnicode / RemoveRepeat / SpecialChars / TokenNum + --use-pii PIIPresidioFilter (presidio-analyzer/anonymizer) + +``DedupFilter`` is not part of the parallel pipeline: it has to see the whole +dataset in one call, so it runs once afterwards. +""" +import argparse +import os + +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.utils import get_logger +from twinkle_agentic.preprocessor import (DeadLoopFilter, DedupFilter, HardFilter, MessageNormalizer, + MessageSanityFilter, QualityPreprocessor, RefuseFilter, TokenSoupFilter, + merge_dropped_shards, run_quality_pipeline, truncate_dropped_logs) + +logger = get_logger() + +# ShareGPT `from` value -> standard message role. ToolACE uses +# system/user/assistant/tool; other ShareGPT variants use human/gpt/observation. +_ROLE_MAP = { + 'system': 'system', + 'user': 'user', 'human': 'user', + 'assistant': 'assistant', 'gpt': 'assistant', 'bot': 'assistant', + 'tool': 'tool', 'observation': 'tool', 'function': 'tool', + 'function_call': 'assistant', 'function_response': 'tool', 'tool_response': 'tool', +} + + +def build_pipeline(args): + """The ordered steps for the parallel pass (dedup is applied separately).""" + steps = [ + MessageNormalizer(normalize_tool_calls=args.normalize_tool_calls), + MessageSanityFilter(), # role order / tool-id matching / content integrity / sensitive words + RefuseFilter(), # drop assistant self-referential refusals + DeadLoopFilter(), # drop degenerate / stuck (hesitation, cascade, ngram repeat) + TokenSoupFilter(), # drop garbled text (replacement/control/private-use chars, script chaos) + HardFilter(min_assistant_chars_2turn=args.min_assistant_chars_2turn), + ] + if args.use_lang: + from twinkle_agentic.preprocessor import LanguageFilter + steps.append(LanguageFilter()) + if args.use_datajuicer: + from twinkle_agentic.preprocessor import (FixUnicodeFilter, RemoveRepeatSentencesFilter, SpecialCharsFilter, + TokenNumFilter) + steps += [FixUnicodeFilter(), RemoveRepeatSentencesFilter(), SpecialCharsFilter(), TokenNumFilter()] + if args.use_pii: + from twinkle_agentic.preprocessor import PIIPresidioFilter + steps.append(PIIPresidioFilter()) + return steps + + +def _row_to_messages(row: dict) -> dict: + """Map one ShareGPT ``conversations`` row to a ``messages`` row. + + Only ``from``->``role`` and ``value``->``content`` are rewritten; a tool call + embedded in an assistant turn is left as-is in ``content`` (ToolACE keeps it + as a bracket-DSL string). Turns whose ``from`` is unknown are dropped so no + invalid role reaches the pipeline. + """ + messages = [] + for turn in (row.get('conversations') or []): + if not isinstance(turn, dict): + continue + role = _ROLE_MAP.get(str(turn.get('from', '')).lower()) + if role is None: + continue + messages.append({'role': role, 'content': turn.get('value', '') or ''}) + return {'messages': messages, 'id': row.get('id', '')} + + +def load_source(input_path: str) -> Dataset: + """Load the raw source into a twinkle Dataset. + + A local path is loaded by extension (jsonl->json, parquet, csv...); anything + else is treated as a hub id (e.g. ``ms://org/name``). Rows pass through + unchanged except for the ShareGPT adaptation above. + """ + ds = Dataset(DatasetMeta(dataset_id=input_path)) + cols = ds.dataset.column_names + if 'messages' not in cols and 'conversations' in cols: + logger.info('[prepare] ShareGPT `conversations` detected -> mapping to `messages`') + # Materialize + convert in Python then rebuild: twinkle's Dataset.map forces + # batched=True and wraps the fn as a Preprocessor, which does not fit a plain + # per-row schema rewrite. The source is small enough to hold in memory. + rows = [_row_to_messages(r) for r in ds.dataset.to_list()] + ds = Dataset(DatasetMeta(data=rows)) + return ds + + +def parse_args(): + p = argparse.ArgumentParser(description='Clean a raw source into RSI seed material.', + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--input', required=True, help='Local .jsonl/.parquet path or an ms:// dataset id.') + p.add_argument('--output', default='output/rsi/subset.jsonl', help='Where to write the surviving rows.') + p.add_argument('--num-proc', type=int, default=4, help='Parallel workers for the map pass.') + p.add_argument('--dropped-log', default='', help='Optional JSONL of dropped-row metadata (empty=off).') + + # On for tool-calling data. Off for pure code (e.g. MBPP): the bracket-DSL + # parser is a marker-less fallback matching ``[name(``, which is also what a + # list comprehension or a call-indexed subscript looks like, so the rewrite + # silently deletes real code from the assistant turn. + p.add_argument('--no-normalize-tool-calls', dest='normalize_tool_calls', + action='store_false', help='pure code data: leave assistant text alone') + p.set_defaults(normalize_tool_calls=True) + # 0, not HardFilter's own 80-char floor: a single-turn valid tool call + # (e.g. `[Func(x=1)]`) is only tens of chars and would be dropped as a + # "shallow_reply". Rule 3 still removes genuinely empty assistant turns. + p.add_argument('--min-assistant-chars-2turn', type=int, default=0) + + p.add_argument('--use-lang', action='store_true', help='LanguageFilter (needs langid)') + p.add_argument('--use-datajuicer', action='store_true', help='data_juicer-based filters') + p.add_argument('--use-pii', action='store_true', help='PIIPresidioFilter (needs presidio)') + return p.parse_args() + + +def main(): + args = parse_args() + os.makedirs(os.path.dirname(os.path.abspath(args.output)) or '.', exist_ok=True) + + pipeline = build_pipeline(args) + logger.info(f'[prepare] pipeline: {" -> ".join(type(s).__name__ for s in pipeline)} ' + f'+ DedupFilter(global)') + + dataset = load_source(args.input) + n_in = len(dataset.dataset) + logger.info(f'[prepare] loaded {n_in} rows from {args.input}') + + # 'mark' mode + run_quality_pipeline is the ghost-proof parallel path: map + # returns equal-length columns flagged _keep, then a single filter removes. + if args.dropped_log: + truncate_dropped_logs(args.dropped_log) + qp = QualityPreprocessor(pipeline, dropped_log_path=args.dropped_log, drop_mode='mark') + run_quality_pipeline(dataset, qp, num_proc=args.num_proc) + if args.dropped_log: + merge_dropped_shards(args.dropped_log) + + n_after_pipeline = len(dataset.dataset) + logger.info(f'[prepare] after parallel pipeline: {n_in} -> {n_after_pipeline}') + + # Global longest-wins dedup -- must see the whole dataset at once. + rows = dataset.dataset.to_list() + kept, dropped = DedupFilter()(rows) + logger.info(f'[prepare] after global dedup: {n_after_pipeline} -> {len(kept)} ' + f'(dropped {len(dropped)} duplicates)') + + out = Dataset(DatasetMeta(data=kept)) + out.save_as(args.output) + logger.info(f'[prepare] wrote {len(kept)} rows -> {args.output}') + + +if __name__ == '__main__': + main() diff --git a/src/twinkle_agentic/rsi/rsi_rl.py b/cookbook/rsi/rl.py similarity index 98% rename from src/twinkle_agentic/rsi/rsi_rl.py rename to cookbook/rsi/rl.py index f36948b0b..4087f4d36 100644 --- a/src/twinkle_agentic/rsi/rsi_rl.py +++ b/cookbook/rsi/rl.py @@ -33,7 +33,8 @@ whole model to vLLM each step. RSI-specific paths come from env vars so the standard CLI (model/infra/rl knobs) stays identical to the reference: - RSI_STD_FLOWS standard_flows.jsonl from rsi_refine.py (default output/rsi/standard_flows.jsonl) + RSI_STD_FLOWS flows jsonl, as written by cookbook/rsi/code/challenge.py --out-flows + (default output/rsi/standard_flows.jsonl) RSI_TEMPLATE template name, must match the model (default Template, for text-only Qwen3-4B) RSI_TESTS jsonl with {id, test_list, test_setup_code} to score code rounds by execution (empty = code rounds are not trained) @@ -99,10 +100,13 @@ # ── RSI-specific paths (env) ─────────────────────────────────────────────── STD_FLOWS = os.environ.get('RSI_STD_FLOWS', 'output/rsi/standard_flows.jsonl') TEMPLATE = os.environ.get('RSI_TEMPLATE', 'Template') # base text template for Qwen3-4B (text-only) -REWARD_TOOL_RESULT = 'tool_result' # matches rsi_refine.attach_reward_method +# Round-level reward labels a flow may carry. The challenger writes 'rubric' on +# its code rounds; 'tool_result' is for flows whose tool rounds are scored +# against a recorded result. +REWARD_TOOL_RESULT = 'tool_result' REWARD_RUBRIC = 'rubric' -# Tests for code rounds, keyed by the flow's id (rsi_refine passes the id through -# from the step-1 subset, which carries the dataset's own tests). +# Tests for code rounds, keyed by the flow's id (the challenger writes both files +# with the same ids; --out-tests here, --out-flows above). TESTS_PATH = os.environ.get('RSI_TESTS', '') TEST_TIMEOUT = int(os.environ.get('RSI_TEST_TIMEOUT', 30)) JUDGE_WORKERS = int(os.environ.get('RSI_JUDGE_WORKERS', max(24, min(96, (os.cpu_count() or 24) // 2)))) @@ -123,8 +127,9 @@ # one jsonl line (step, kind, ref/gen call, completion head, score, judge reason). # Pure observability; the reward and training path are untouched. REWARD_DUMP = os.environ.get('RSI_REWARD_DUMP', '') -# Raw step-1 conversations (before rsi_refine). rsi_refine's flow schema keeps -# only the FIRST user message (as `query`) plus the tool rounds, so any parameter +# Raw conversations behind the flows, when the flows were derived from a dataset +# rather than invented. A flow keeps only the FIRST user message (as `query`) +# plus the tool rounds, so any parameter # the user stated in a LATER user turn is missing from a round's prompt and the # model is asked to produce a call it cannot possibly know. When this points at # the raw file, each round's prompt is rebuilt to splice those dropped user (and @@ -693,8 +698,8 @@ def _intervening_turns(raw_msgs: List[Dict[str, Any]], lo: int, hi: int) -> List Assistant tool-call messages (content starting with ``[``) and tool results are dropped here because the structured prior rounds already carry the call - and its result; what is recovered is exactly the conversational turns - rsi_refine did not keep. + and its result; what is recovered is exactly the conversational turns the + flow did not keep. """ out: List[Dict[str, Any]] = [] for j in range(lo + 1, hi): @@ -719,7 +724,7 @@ def build_round_trajectories(records: List[Dict[str, Any]], what its reward executes; otherwise it stays context-only. When ``raw_by_query`` is given, the round prompt is rebuilt from the raw - conversation so the user turns that rsi_refine dropped (e.g. the turn that + conversation so the user turns the flow dropped (e.g. the turn that states the call's arguments) are spliced back in at their real positions; flows whose raw conversation cannot be located fall back to the flow-only prompt (system + first query + prior rounds). diff --git a/cookbook/rsi/run_rsi.py b/cookbook/rsi/run_rsi.py index 8bddfea56..f1b9ebeca 100644 --- a/cookbook/rsi/run_rsi.py +++ b/cookbook/rsi/run_rsi.py @@ -1,37 +1,44 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""RSI pipeline entry point — run any single stage (or the whole chain) so each -step can be validated in isolation. +"""RSI entry point — run one stage on its own, or the whole chain. -The four stages live in ``twinkle_agentic.rsi`` and each already has its own CLI: +Three stages, one model (Qwen3-4B) playing both roles: - 1 prepare twinkle_agentic.rsi.rsi_prepare (CPU) raw -> subset - 2 refine twinkle_agentic.rsi.rsi_refine (API) subset-> flows - 3 rl twinkle_agentic.rsi.rsi_rl (ray+GPU) flows -> executor LoRA - 4 distill twinkle_agentic.rsi.rsi_distill (ray+GPU) dump -> role LoRA + 1 prepare cookbook/rsi/prepare.py (CPU) raw -> seeds + 2 challenge cookbook/rsi/code/challenge.py (ray+GPU) [seeds]-> flows + tests + 3 rl cookbook/rsi/rl.py (ray+GPU) flows -> trained model + +``prepare`` only cleans a dataset into seed material and is optional: the +challenger invents problems from nothing when given no seeds. ``challenge`` asks +the model for a problem plus a reference solution, RUNS that solution to get the +ground truth, turns it into asserts, then keeps only the problems the same model +solves sometimes-but-not-always -- a group that all passes or all fails gives +GRPO a zero gradient. ``rl`` trains on what survived, in ``grpo`` mode (feed the +sandbox error back as a tool turn and let it continue) or ``opsd`` (a teacher +that was shown the reference solution distills the student). Why this launches SUBPROCESSES instead of importing and calling: - * ``rsi_rl`` runs ``CLI.from_args()`` and ``swanlab.init()`` at IMPORT time, so - merely importing it would parse this launcher's argv and start a run. - * ``rl`` and ``distill`` need DIFFERENT ray topologies (MultiLora+sampler vs a - single TransformersModel group); they cannot share one ray init in-process. -Running each stage as its own ``python -m ...`` process side-steps both — and is -exactly what "run each step separately to validate" needs. + * ``rl`` runs ``CLI.from_args()`` and ``swanlab.init()`` at IMPORT time, so + merely importing it would parse this launcher's argv and start a run; + * the stages need different ray topologies (sampler-only vs trainer+sampler) + and cannot share one ray init in-process. -This launcher invents no parameters: it only wires the default output of one -stage into the input of the next (reusing each script's own default paths) and -forwards any extra flags straight through to the selected stage. +This launcher invents no parameters: it wires each stage's default output into +the next stage's input and forwards any extra flags straight through. Examples -------- -Validate one stage at a time (extra flags after the known ones are forwarded): - + # clean a dataset into seeds (optional) python cookbook/rsi/run_rsi.py --step prepare --raw data/raw.jsonl - python cookbook/rsi/run_rsi.py --step refine --teacher-model qwen3-235b-a22b-instruct-2507 - python cookbook/rsi/run_rsi.py --step rl --model.model_id ms://Qwen/Qwen3-4B --infra.model_gpus 4 - python cookbook/rsi/run_rsi.py --step distill --dump output/rsi/dump/refine.jsonl --adapter refine -Run the whole chain with default paths (each stage still a fresh process): + # invent problems: from nothing, or seeded by the file above + python cookbook/rsi/run_rsi.py --step challenge --keep-target 500 + python cookbook/rsi/run_rsi.py --step challenge --seeds output/rsi/subset.jsonl + # train; twinkle CLI knobs are forwarded as extras + python cookbook/rsi/run_rsi.py --step rl --mode grpo \ + --model-id ms://Qwen/Qwen3-4B --model-gpus 4 --sampler-gpus 4 + + # whole chain with default paths (each stage still a fresh process) python cookbook/rsi/run_rsi.py --step all --raw data/raw.jsonl """ import argparse @@ -39,92 +46,102 @@ import subprocess import sys -# Default paths chain one stage into the next. These mirror the defaults baked -# into each stage's own CLI, kept here so --step all wires up with no flags. -DEFAULT_SUBSET = 'output/rsi/subset.jsonl' # rsi_prepare --output -DEFAULT_FLOWS = 'output/rsi/standard_flows.jsonl' # rsi_refine --output / rsi_rl RSI_STD_FLOWS -DEFAULT_DUMP = 'output/rsi/dump/refine.jsonl' # llm_backup LLM_BACKUP_DUMP_PATH / rsi_distill --input - -MODULES = { - 'prepare': 'twinkle_agentic.rsi.rsi_prepare', - 'refine': 'twinkle_agentic.rsi.rsi_refine', - 'rl': 'twinkle_agentic.rsi.rsi_rl', - 'distill': 'twinkle_agentic.rsi.rsi_distill', +HERE = os.path.dirname(os.path.abspath(__file__)) + +# Default paths chain one stage into the next. They mirror the defaults each +# stage ships with, kept here so --step all wires up with no flags. +DEFAULT_SEEDS = 'output/rsi/subset.jsonl' # prepare --output / challenge --seed-file +DEFAULT_FLOWS = 'output/rsi/challenge_flows.jsonl' # challenge --out-flows / rl RSI_STD_FLOWS +DEFAULT_TESTS = 'output/rsi/challenge_tests.jsonl' # challenge --out-tests / rl RSI_TESTS + +SCRIPTS = { + 'prepare': os.path.join(HERE, 'prepare.py'), + 'challenge': os.path.join(HERE, 'code', 'challenge.py'), + 'rl': os.path.join(HERE, 'rl.py'), } -ORDER = ['prepare', 'refine', 'rl', 'distill'] +ORDER = ['prepare', 'challenge', 'rl'] -def _run(module: str, argv: list, env: dict) -> None: - """Run ``python -m module argv...`` as a child process, streaming its output. +def _run(script: str, argv: list, env: dict) -> None: + """Run ``python script argv...`` as a child process, streaming its output. Raises on non-zero exit so --step all stops at the first failing stage - instead of silently feeding a broken artifact into the next stage. + instead of silently feeding a broken artifact into the next one. """ - cmd = [sys.executable, '-m', module] + argv + cmd = [sys.executable, script] + argv print(f'\n[run_rsi] $ {" ".join(cmd)}', flush=True) subprocess.run(cmd, env=env, check=True) def _argv_for(step: str, a: argparse.Namespace, extra: list) -> tuple: """Build (argv, env) for one stage. ``extra`` is forwarded verbatim so each - stage's own flags (teacher creds, twinkle CLI knobs, ...) still work.""" + stage's own flags (challenger knobs, twinkle CLI knobs, ...) still work.""" env = dict(os.environ) if step == 'prepare': if not a.raw: raise SystemExit('[run_rsi] --step prepare 需要 --raw 指向原始数据源') - argv = ['--input', a.raw, '--output', a.subset, '--num-proc', str(a.num_proc)] + argv = ['--input', a.raw, '--output', a.seeds, '--num-proc', str(a.num_proc)] if a.dropped_log: argv += ['--dropped-log', a.dropped_log] return argv + extra, env - if step == 'refine': - return ['--input', a.subset, '--output', a.flows] + extra, env + if step == 'challenge': + argv = ['--out-flows', a.flows, '--out-tests', a.tests] + if a.seeds_given: + argv += ['--seed-file', a.seeds] + if a.keep_target: + argv += ['--keep-target', str(a.keep_target)] + return argv + extra, env if step == 'rl': - # rsi_rl reads the standard-flow path from an env var, not a flag; + # rl reads flows/tests and the solver mode from env vars; the # model/infra/rl knobs arrive through `extra` (twinkle CLI). env['RSI_STD_FLOWS'] = a.flows + env['RSI_TESTS'] = a.tests + env['RSI_SOLVER_MODE'] = a.mode return list(extra), env - if step == 'distill': - # rsi_distill accepts --input/--adapter and also honours these env vars. - env['RSI_DUMP_PATH'] = a.dump - argv = ['--input', a.dump] - if a.adapter: - argv += ['--adapter', a.adapter] - return argv + extra, env raise SystemExit(f'[run_rsi] 未知 step: {step}') def main(): parser = argparse.ArgumentParser( - description='RSI pipeline launcher — run one stage (validate) or the whole chain.', + description='RSI launcher — run one stage (validate) or the whole chain.', formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument('--step', required=True, choices=ORDER + ['all'], - help='Which stage to run (or "all" for prepare->refine->rl->distill).') + help='Which stage to run (or "all" for prepare->challenge->rl).') parser.add_argument('--raw', default='', help='Raw data source for prepare (local path or ms:// id).') - parser.add_argument('--subset', default=DEFAULT_SUBSET, help='prepare output / refine input.') - parser.add_argument('--flows', default=DEFAULT_FLOWS, help='refine output / rl standard-flow input.') - parser.add_argument('--dump', default=os.environ.get('LLM_BACKUP_DUMP_PATH', DEFAULT_DUMP), - help='llm_backup dump JSONL / distill input.') - parser.add_argument('--adapter', default='', help='Adapter name for distill (default: derived from dump name).') - parser.add_argument('--num-proc', type=int, default=int(os.environ.get('RSI_NUM_PROC', '4')), - help='Parallel workers for prepare.') + parser.add_argument('--seeds', default=DEFAULT_SEEDS, + help='prepare output / challenge seed pool. Passed to challenge only ' + 'when given explicitly or when the chain produced it.') + parser.add_argument('--flows', default=DEFAULT_FLOWS, help='challenge output flows / rl input.') + parser.add_argument('--tests', default=DEFAULT_TESTS, help='challenge output tests / rl code asserts.') + parser.add_argument('--keep-target', type=int, default=0, + help="How many problems challenge should keep (0 = the script's own default).") + parser.add_argument('--mode', default='grpo', choices=['grpo', 'opsd'], + help='rl solver mode (RSI_SOLVER_MODE).') + parser.add_argument('--num-proc', type=int, default=4, help='Parallel workers for prepare.') parser.add_argument('--dropped-log', default='', help='Optional dropped-row log for prepare.') a, extra = parser.parse_known_args() + # A seed pool is only handed to the challenger when it was asked for: passing + # the default path silently would turn "invent from scratch" into "vary + # whatever happens to be left in output/rsi/ from an earlier run". + a.seeds_given = '--seeds' in sys.argv if a.step == 'all': if extra: # For 'all' the extras are ambiguous (which stage?); refuse rather than # forward a flag to a stage that does not accept it. raise SystemExit(f'[run_rsi] --step all 不接受透传参数 {extra}；请逐个 --step 跑并各自带参数') + if not a.raw: + raise SystemExit('[run_rsi] --step all 需要 --raw 指向原始数据源') + # prepare ran, so its output exists and the challenger should use it. + a.seeds_given = True for step in ORDER: - if step == 'prepare' and not a.raw: - raise SystemExit('[run_rsi] --step all 需要 --raw 指向原始数据源') argv, env = _argv_for(step, a, []) - _run(MODULES[step], argv, env) + _run(SCRIPTS[step], argv, env) print('\n[run_rsi] all stages done.', flush=True) return argv, env = _argv_for(a.step, a, extra) - _run(MODULES[a.step], argv, env) + _run(SCRIPTS[a.step], argv, env) if __name__ == '__main__': diff --git a/cookbook/rsi/run_rsi_selfplay.py b/cookbook/rsi/run_rsi_selfplay.py deleted file mode 100644 index c8c6bdc4d..000000000 --- a/cookbook/rsi/run_rsi_selfplay.py +++ /dev/null @@ -1,125 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""RSI self-play entry point — the challenger/solver data loop, no prepare/refine. - -Two stages live in ``twinkle_agentic.rsi`` (one model, Qwen3-4B, plays both roles): - - 1 challenge twinkle_agentic.rsi.rsi_challenge (ray+GPU) [seed] -> flows + tests - 2 rl twinkle_agentic.rsi.rsi_rl (ray+GPU) flows -> trained model - -Stage 1 asks the model to invent (or vary a seed into) self-contained Python -problems, RUNS its reference solution to get the ground truth, turns that into -asserts, then keeps only the problems the same model solves sometimes-but-not- -always (0 < pass < N). Stage 3 trains on those with ``RSI_SOLVER_MODE``: -``grpo`` (feed the sandbox error back as a tool turn and continue) or ``opsd`` -(a teacher that saw the reference solution distills the student). - -Why this launches SUBPROCESSES instead of importing and calling (same reason as -run_rsi.py): - * ``rsi_rl`` runs ``CLI.from_args()`` and ``swanlab.init()`` at IMPORT time, so - merely importing it would parse this launcher's argv and start a run. - * the two stages need different ray topologies (sampler-only vs trainer+sampler) - and cannot share one ray init in-process. - -This launcher invents no parameters: it wires stage 1's default outputs into -stage 2's inputs (through each script's own env vars) and forwards any extra -flags straight through to the selected stage. - -Examples --------- -Run one stage at a time (extra flags after the known ones are forwarded): - - # from scratch (no seed dataset) - python cookbook/rsi/run_rsi_selfplay.py --step challenge - # from a seed dataset (challenger writes variants of its queries) - python cookbook/rsi/run_rsi_selfplay.py --step challenge --seed data/seed.jsonl - # train — grpo (default) or opsd; twinkle CLI knobs are forwarded as extras - python cookbook/rsi/run_rsi_selfplay.py --step rl --mode grpo \ - --model.model_id ms://Qwen/Qwen3-4B --infra.model_gpus 4 --infra.sampler_gpus 4 - -Run the whole chain with default paths (each stage still a fresh process): - - python cookbook/rsi/run_rsi_selfplay.py --step all --mode grpo -""" -import argparse -import os -import subprocess -import sys - -# Default paths chain stage 1 into stage 2. These mirror the defaults baked into -# each stage's own env-var config, kept here so --step all wires up with no flags. -DEFAULT_FLOWS = 'output/rsi/challenge_flows.jsonl' # rsi_challenge RSI_CH_OUT_FLOWS / rsi_rl RSI_STD_FLOWS -DEFAULT_TESTS = 'output/rsi/challenge_tests.jsonl' # rsi_challenge RSI_CH_OUT_TESTS / rsi_rl RSI_TESTS - -MODULES = { - 'challenge': 'twinkle_agentic.rsi.rsi_challenge', - 'rl': 'twinkle_agentic.rsi.rsi_rl', -} -ORDER = ['challenge', 'rl'] - - -def _run(module: str, argv: list, env: dict) -> None: - """Run ``python -m module argv...`` as a child process, streaming its output. - - Raises on non-zero exit so --step all stops at the first failing stage - instead of silently feeding a broken artifact into the next stage. - """ - cmd = [sys.executable, '-m', module] + argv - print(f'\n[run_rsi_selfplay] $ {" ".join(cmd)}', flush=True) - subprocess.run(cmd, env=env, check=True) - - -def _argv_for(step: str, a: argparse.Namespace, extra: list) -> tuple: - """Build (argv, env) for one stage. ``extra`` is forwarded verbatim so each - stage's own flags (twinkle CLI knobs for rl) still work.""" - env = dict(os.environ) - if step == 'challenge': - # rsi_challenge is configured purely through RSI_CH_* env vars (no CLI). - env['RSI_CH_OUT_FLOWS'] = a.flows - env['RSI_CH_OUT_TESTS'] = a.tests - if a.seed: - env['RSI_CH_SEED'] = a.seed - return list(extra), env - if step == 'rl': - # rsi_rl reads flows/tests and the solver mode from env vars; the - # model/infra/rl knobs arrive through `extra` (twinkle CLI). - env['RSI_STD_FLOWS'] = a.flows - env['RSI_TESTS'] = a.tests - env['RSI_SOLVER_MODE'] = a.mode - return list(extra), env - raise SystemExit(f'[run_rsi_selfplay] 未知 step: {step}') - - -def main(): - parser = argparse.ArgumentParser( - description='RSI self-play launcher — run one stage (validate) or the whole chain.', - formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument('--step', required=True, choices=ORDER + ['all'], - help='Which stage to run (or "all" for challenge->rl).') - parser.add_argument('--seed', default='', - help='Optional seed dataset (jsonl) for challenge; empty = invent from scratch.') - parser.add_argument('--flows', default=DEFAULT_FLOWS, - help='challenge output flows / rl standard-flow input.') - parser.add_argument('--tests', default=DEFAULT_TESTS, - help='challenge output tests / rl code-round asserts input.') - parser.add_argument('--mode', default='grpo', choices=['grpo', 'opsd'], - help='rl solver mode (RSI_SOLVER_MODE).') - a, extra = parser.parse_known_args() - - if a.step == 'all': - if extra: - # For 'all' the extras are ambiguous (which stage?); refuse rather than - # forward a flag to a stage that does not accept it. - raise SystemExit(f'[run_rsi_selfplay] --step all 不接受透传参数 {extra}；' - '请逐个 --step 跑并各自带参数') - for step in ORDER: - argv, env = _argv_for(step, a, []) - _run(MODULES[step], argv, env) - print('\n[run_rsi_selfplay] all stages done.', flush=True) - return - - argv, env = _argv_for(a.step, a, extra) - _run(MODULES[a.step], argv, env) - - -if __name__ == '__main__': - main() diff --git a/src/twinkle_agentic/challenger/__init__.py b/src/twinkle_agentic/challenger/__init__.py new file mode 100644 index 000000000..39e38ad63 --- /dev/null +++ b/src/twinkle_agentic/challenger/__init__.py @@ -0,0 +1,25 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .agentic import AgenticChallenger, AgenticPrompts, parse_check_script, parse_problem_statement +from .base import Challenger, Explorer, assistant_text, attach_user_data +from .code import (CodeChallenger, CodePrompts, KeywordStore, build_asserts, extract_code, + is_constant_answer, load_seeds, parse_challenge, run_asserts) + +__all__ = [ + 'AgenticChallenger', + 'AgenticPrompts', + 'Challenger', + 'CodeChallenger', + 'CodePrompts', + 'Explorer', + 'KeywordStore', + 'assistant_text', + 'attach_user_data', + 'build_asserts', + 'extract_code', + 'is_constant_answer', + 'load_seeds', + 'parse_check_script', + 'parse_challenge', + 'parse_problem_statement', + 'run_asserts', +] diff --git a/src/twinkle_agentic/challenger/agentic.py b/src/twinkle_agentic/challenger/agentic.py new file mode 100644 index 000000000..d54c5ed7e --- /dev/null +++ b/src/twinkle_agentic/challenger/agentic.py @@ -0,0 +1,559 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Agentic challenger: invent tasks by doing them first. + +The approach mirrors how the code challenger works, adapted to tool-using +agents. Instead of writing a problem statement and hoping it is achievable, +the model first *does* something interesting in a sandbox (round 1), then a +second call writes check assertions that verify the end state, and a third +call writes the problem statement someone else would need to reproduce it. + +Steps for one candidate: + + 1. Choose direction + keywords. Optionally start from a seed trajectory. + 2. Round 1 (explore, multi-turn with tools): model acts in a clean sandbox, + producing a tool-call chain and a final workspace state. + 3. Round 2a (explore, single-turn): model sees the trajectory and writes a + python check script that asserts properties of the end state. + 4. Verify: run the check script in the sandbox (must pass). + 5. Round 2b (explore, single-turn): model sees trajectory + checks and + writes a problem statement. + 6. Difficulty filter: reset workspace, let the solver do the task N times, + run checks, keep only "sometimes pass" tasks. + +Because every round-1 episode needs a clean workspace and because episodes +share a single long-lived sandbox, round 1 is **serial** -- one proposal at a +time with a workspace reset in between. Rounds 2a/2b are text-only generation +and can be batched. + +Prompt text is not here. Every string the model sees arrives in +:class:`AgenticPrompts`, built by whoever runs the challenger -- see +``cookbook/rsi/agentic/prompts.py``. +""" +import re +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple + +from twinkle.data_format import SamplingParams, Trajectory, user_data_get +from twinkle.utils import get_logger +from .base import Challenger, Explorer, assistant_text, attach_user_data +from .code import KeywordStore, parse_keyword_list + +logger = get_logger() + +__all__ = [ + 'AgenticChallenger', + 'AgenticPrompts', + 'parse_check_script', + 'parse_problem_statement', +] + +_FENCE_RE = re.compile(r'```(?:python|py)?\s*\n(.*?)```', re.S) + + +# ── parsing ─────────────────────────────────────────────────────────────── + +def parse_check_script(text: str) -> Optional[str]: + """Extract a python check script from the model's reply. + + Looks for the last fenced python code block after ``</think>``. + Returns ``None`` when nothing usable is found. + """ + body = text or '' + idx = body.rfind('</think>') + if idx >= 0: + body = body[idx + len('</think>'):] + blocks = _FENCE_RE.findall(body) + if not blocks: + return None + script = blocks[-1].strip() + return script if script else None + + +def parse_problem_statement(text: str) -> Optional[str]: + """Extract a problem statement from the model's reply. + + The model is asked to return the problem in prose (not code). We take + everything after ``</think>`` with code fences stripped as the statement. + Returns ``None`` when the result is empty. + """ + body = text or '' + idx = body.rfind('</think>') + if idx >= 0: + body = body[idx + len('</think>'):] + # Strip any fenced blocks (those are code, not prose) + body = _FENCE_RE.sub('', body).strip() + # Strip json fences too + body = re.sub(r'^\s*```(?:json)?\s*|\s*```\s*$', '', body, flags=re.I).strip() + return body if body else None + + +def _trajectory_summary(trajectory: Trajectory) -> str: + """A compact text representation of a trajectory for prompting. + + Shows each message as role: content (truncated for tool results). + """ + parts = [] + for msg in trajectory.get('messages') or []: + role = msg.get('role', '?') + content = msg.get('content') or '' + if role == 'tool' and len(content) > 500: + content = content[:500] + '...[truncated]' + parts.append(f'[{role}] {content}') + return '\n'.join(parts) + + +# ── prompts ──────────────────────────────────────────────────────────────── + +@dataclass +class AgenticPrompts: + """Every string an :class:`AgenticChallenger` sends. + + All fields are injected by the caller (no defaults with real text here). + Placeholder validation happens at construction time. + """ + + # Round 1: model acts in sandbox + system: str + from_scratch: str + from_seed: str = '' + from_keywords: str = '' + from_seed_keywords: str = '' + + # Round 2a: write check script + check_system: str = '' + check_user: str = '' + + # Round 2b: write problem statement + problem_system: str = '' + problem_user: str = '' + + # Keyword generation (same structure as code side) + keyword_system: str = '' + keyword_user: str = '' + keyword_expand_user: str = '' + + _REQUIRED_FIELDS = { + 'from_seed': ('seed',), + 'from_keywords': ('keywords',), + 'from_seed_keywords': ('seed', 'keywords'), + 'check_user': ('trajectory', 'final_state'), + 'problem_user': ('trajectory', 'checks'), + 'keyword_user': ('k', 'desc'), + 'keyword_expand_user': ('kw', 'm'), + } + + def __post_init__(self): + for name in ('system', 'from_scratch', 'check_system', 'check_user', + 'problem_system', 'problem_user'): + if not getattr(self, name).strip(): + raise ValueError(f'AgenticPrompts.{name} is required') + for name, placeholders in self._REQUIRED_FIELDS.items(): + text = getattr(self, name) + if not text: + continue + for placeholder in placeholders: + if '{' + placeholder + '}' not in text: + raise ValueError(f'AgenticPrompts.{name} must contain ' + f'{{{placeholder}}}') + + def require(self, *names: str) -> None: + """Raise unless every named prompt was supplied.""" + missing = [n for n in names if not getattr(self, n).strip()] + if missing: + raise ValueError(f'this configuration needs AgenticPrompts.' + f'{", AgenticPrompts.".join(missing)}') + + +# ── challenger ───────────────────────────────────────────────────────────── + +class AgenticChallenger(Challenger): + """Propose tool-using tasks by first doing them, then describing them. + + Args: + prompts: every string sent to the model. + explorer: batch-in / batch-out generation with sandbox tools (multi-turn). + seeds: optional pool of seed trajectories (dicts with a ``query`` key), + drawn with replacement. + keyword_store: optional bank for diversity control. + category_desc: category -> description for keyword generation. + seed_mix_prob: chance a proposal carries a seed. + reset_fn: called before each round-1 episode to clean the sandbox + workspace. Must be synchronous and leave the workspace empty. + run_check_fn: run a python script in the sandbox's current state. + Signature: ``(source: str) -> (exit_code: int, output: str)``. + workspace_snapshot_fn: after round 1, return a text summary of the + workspace state (e.g. ``find . -type f``). If None, a default + that lists messages is used. + combo_arity: ``'triple'`` or ``'mix'``, as in :class:`.CodeChallenger`. + arity_weights: weights for the ``'mix'`` subset size. + single_kw_prob: chance of using one category in ``'triple'`` mode. + keyword_refill_target / keyword_gen_calls / keyword_refill_tries / + keyword_params: keyword bank refill parameters. + min_batch: smallest batch worth sending to the explorer. + problem_max_chars: reject problem statements longer than this. + reject_sink: called with a dict for every rejected proposal. + """ + + def __init__( + self, + prompts: AgenticPrompts, + explorer: Explorer, + *, + seeds: Sequence[Dict[str, Any]] = (), + keyword_store: Optional[KeywordStore] = None, + category_desc: Optional[Dict[str, str]] = None, + seed_mix_prob: float = 0.5, + reset_fn: Callable[[], None], + run_check_fn: Callable[[str], Tuple[int, str]], + workspace_snapshot_fn: Optional[Callable[[], str]] = None, + combo_arity: str = 'triple', + arity_weights: Optional[Sequence[float]] = None, + single_kw_prob: float = 0.1, + keyword_refill_target: int = 128, + keyword_gen_calls: int = 8, + keyword_refill_tries: int = 2, + keyword_params: Optional[SamplingParams] = None, + min_batch: int = 1, + problem_max_chars: int = 8192, + reject_sink: Optional[Callable[[Dict[str, Any]], None]] = None, + **challenger_kwargs: Any, + ): + super().__init__(explorer, system=prompts.system, **challenger_kwargs) + if combo_arity not in ('triple', 'mix'): + raise ValueError(f"combo_arity must be 'triple' or 'mix', got {combo_arity!r}") + if keyword_store is not None: + desc = category_desc or {} + missing_cats = [c for c in keyword_store.categories if not desc.get(c)] + if missing_cats: + raise ValueError(f'category_desc is missing a description for ' + f'{missing_cats}; a dry category could not be refilled.') + prompts.require('keyword_system', 'keyword_user', 'from_keywords') + self.prompts = prompts + self.seeds = list(seeds) + self.store = keyword_store + self.category_desc = dict(category_desc or {}) + self.seed_mix_prob = seed_mix_prob + self.reset_fn = reset_fn + self.run_check_fn = run_check_fn + self.workspace_snapshot_fn = workspace_snapshot_fn + self.combo_arity = combo_arity + self.arity_weights = list(arity_weights) if arity_weights else None + self.single_kw_prob = single_kw_prob + self.keyword_refill_target = keyword_refill_target + self.keyword_gen_calls = keyword_gen_calls + self.keyword_refill_tries = keyword_refill_tries + self.keyword_params = keyword_params + self.min_batch = max(1, min_batch) + self.problem_max_chars = problem_max_chars + self.reject_sink = reject_sink + if self.seeds: + prompts.require('from_seed') + if self.store is not None: + prompts.require('from_seed_keywords') + self._nonce = 0 + self.stats: Dict[str, int] = { + 'round1_done': 0, 'check_parse_fail': 0, 'check_run_fail': 0, + 'problem_parse_fail': 0, 'too_long': 0, 'parsed': 0, + } + self._hard: List[Tuple[str, str]] = [] + + # ------------------------------------------------------------- proposing + + def propose(self, count: int) -> List[Trajectory]: + """Build ``count`` prompt trajectories for round 1. + + Each carries a direction + keywords + optional seed. The explorer will + run these multi-turn in the sandbox. + """ + proposals: List[Trajectory] = [] + for _ in range(count): + picks = self._draw_keywords() + body = '\n'.join(f'- {c}: {t}' for c, t in picks) + use_seed = bool(self.seeds) and self.rng.random() < self.seed_mix_prob + seed = self.rng.choice(self.seeds) if use_seed else None + if use_seed and picks: + user = self.prompts.from_seed_keywords.format( + seed=seed['query'], keywords=body) + elif use_seed: + user = self.prompts.from_seed.format(seed=seed['query']) + elif picks: + user = self.prompts.from_keywords.format(keywords=body) + else: + user = self.prompts.from_scratch + proposal: Trajectory = { + 'messages': [{'role': 'system', 'content': self.prompts.system}, + {'role': 'user', 'content': user}], + } + proposals.append(attach_user_data( + proposal, keywords=picks, seeded=use_seed, keyword_block=body)) + return proposals + + # ------------------------------------------------------------- building + + def build(self, explored: List[Trajectory]) -> List[Optional[Trajectory]]: + """Satisfy the abstract method; not usable outside ``_round``. + + ``_build_one`` requires the sandbox to hold the episode's workspace state, + which is only guaranteed inside the serial ``_round`` loop. Calling this + method directly will produce wrong results because the sandbox state does + not match the trajectory being processed. + """ + raise RuntimeError( + f'{type(self).__name__}.build() must not be called directly; ' + f'the serial _round() loop calls _build_one() per episode instead.') + + def _build_one(self, explored: Trajectory) -> Optional[Trajectory]: + """Process one round-1 result: write checks, verify, write problem. + + Called while the sandbox still holds this episode's workspace state. + """ + summary = _trajectory_summary(explored) + snapshot = self.workspace_snapshot_fn() if self.workspace_snapshot_fn else summary + + # Round 2a: write check script + # NOTE: This goes through the same explorer (with tool schemas visible). + # The prompt must clearly instruct the model to output ONLY a code block + # and not call tools, otherwise tool calls would corrupt the sandbox state + # before verification. The check_system prompt enforces this. + check_prompt: Trajectory = { + 'messages': [ + {'role': 'system', 'content': self.prompts.check_system}, + {'role': 'user', 'content': self.prompts.check_user.format( + trajectory=summary, final_state=snapshot)}, + ], + } + check_reply = self.explore([check_prompt]) + script = parse_check_script(assistant_text(check_reply[0])) + if script is None: + self.stats['check_parse_fail'] += 1 + self._reject_record(explored, 'check_parse_fail') + return None + + # Verify: run check script in current sandbox state (must pass) + exit_code, output = self.run_check_fn(script) + if exit_code != 0: + self.stats['check_run_fail'] += 1 + self._reject_record(explored, 'check_run_fail', + detail=f'exit {exit_code}: {output[-200:]}') + return None + + # Round 2b: write problem statement + problem_prompt: Trajectory = { + 'messages': [ + {'role': 'system', 'content': self.prompts.problem_system}, + {'role': 'user', 'content': self.prompts.problem_user.format( + trajectory=summary, checks=script)}, + ], + } + problem_reply = self.explore([problem_prompt]) + statement = parse_problem_statement(assistant_text(problem_reply[0])) + if statement is None: + self.stats['problem_parse_fail'] += 1 + self._reject_record(explored, 'problem_parse_fail') + return None + if len(statement) > self.problem_max_chars: + self.stats['too_long'] += 1 + self._reject_record(explored, 'too_long') + return None + + self.stats['parsed'] += 1 + task: Trajectory = { + 'messages': [{'role': 'user', 'content': statement}], + } + return attach_user_data( + task, + check_script=script, + keywords=user_data_get(explored.get('user_data'), 'keywords', []), + seeded=user_data_get(explored.get('user_data'), 'seeded', False)) + + def _reject_record(self, traj: Trajectory, reason: str, detail: str = '') -> None: + if self.reject_sink is not None: + payload: Dict[str, Any] = {'reason': reason} + if detail: + payload['detail'] = detail + payload['last_assistant'] = assistant_text(traj)[:500] + self.reject_sink(payload) + + # ------------------------------------------------------------ revised _round + + def _round(self, missing: int) -> Optional[List[Trajectory]]: + """One cycle: serial round-1 episodes, inline build, then difficulty filter.""" + count = min(self._estimate(missing), self.max_proposals_per_round) + proposals = self.propose(count) + if not proposals: + return None + + usable: List[Trajectory] = [] + for proposal in proposals: + # Reset workspace, run round 1 + self.reset_fn() + result = self.explore([proposal]) + if not result: + continue + explored = result[0] + self.stats['round1_done'] += 1 + # Workspace still holds this episode's state → build inline + task = self._build_one(explored) + if task is not None: + usable.append(task) + + kept = self._filter_difficulty(usable) if self.solver_rollouts else usable + self.n_proposed += len(proposals) + self.n_kept += len(kept) + band = (f', in difficulty band {len(kept)}' if self.solver_rollouts else '') + logger.info(f'[{type(self).__name__}] proposed {len(proposals)}, usable ' + f'{len(usable)}{band} (cumulative {self.n_kept}/{self.n_proposed})') + return kept + + # ------------------------------------------------------------ difficulty + + def _filter_difficulty(self, tasks: List[Trajectory]) -> List[Trajectory]: + """Override: each solver attempt needs a clean workspace, so run serially.""" + if not tasks: + return [] + passes = [0] * len(tasks) + for i, task in enumerate(tasks): + prompt = self.solver_prompt(task) + for _ in range(self.solver_rollouts): + self.reset_fn() + attempts = self._solver_explore( + [dict(prompt)], sampling_params=self.solver_params) + if attempts and self.judge_attempt(task, attempts[0]): + passes[i] += 1 + + measured = [ + attach_user_data(task, n_pass=passes[i], n_rollouts=self.solver_rollouts) + for i, task in enumerate(tasks) + ] + self.on_difficulty_measured(measured) + high = self.solver_rollouts - self.keep_max_pass_margin + return [t for t, n in zip(measured, passes) if self.keep_min_pass <= n <= high] + + def solver_prompt(self, task: Trajectory) -> Trajectory: + """The task statement, nothing else -- the solver's own harness adds the system.""" + return {'messages': [dict(m) for m in task.get('messages') or []]} + + def judge_attempt(self, task: Trajectory, attempt: Trajectory) -> bool: + """Run the check script against the sandbox's current state.""" + script = user_data_get(task.get('user_data'), 'check_script', '') + if not script: + return False + exit_code, _ = self.run_check_fn(script) + return exit_code == 0 + + def on_difficulty_measured(self, candidates: List[Trajectory]) -> None: + """Remember keywords behind candidates nobody solved.""" + if self.store is None: + return + seen = {(c, t.lower()) for c, t in self._hard} + for task in candidates: + data = task.get('user_data') + if user_data_get(data, 'n_pass', 0) > 0: + continue + for pick in user_data_get(data, 'keywords', []) or []: + if isinstance(pick, (list, tuple)) and len(pick) >= 2: + c, t = pick[0], pick[1] + if (c, t.lower()) not in seen: + seen.add((c, t.lower())) + self._hard.append((c, t)) + + # ------------------------------------------------------------ keywords + + def _draw_keywords(self) -> List[Tuple[str, str]]: + """Consume one keyword combination from the bank; [] without a bank.""" + if self.store is None: + return [] + categories = self.store.categories + if self.combo_arity == 'mix': + if self.arity_weights and len(self.arity_weights) == len(categories): + k = self.rng.choices(range(1, len(categories) + 1), + weights=self.arity_weights)[0] + else: + k = self.rng.randint(1, len(categories)) + cats = self.rng.sample(list(categories), k) + elif self.rng.random() < self.single_kw_prob: + cats = [self.rng.choice(categories)] + else: + cats = list(categories) + picks: List[Tuple[str, str]] = [] + for c in cats: + if not self.store.unused(c): + self._refill(c) + text = self.store.take(c, self.rng) + if text is not None: + picks.append((c, text)) + return picks + + def _refill(self, category: str) -> None: + """Ask the model for more keywords in ``category``.""" + tries = 0 + while not self.store.unused(category): + new = self._generate_keywords(category, self.keyword_refill_target) + added = self.store.add(category, new, source='gen') + tries += 1 + if added == 0 and tries >= self.keyword_refill_tries: + if self.store.items[category]: + self.store.recycle(category) + logger.info(f'[AgenticChallenger] keyword category {category!r} ' + f'exhausted -> recycled {len(self.store.items[category])} topics') + break + + def _generate_keywords(self, category: str, n_want: int) -> List[str]: + """Up to ``n_want`` keywords the bank does not already hold.""" + if n_want <= 0: + return [] + known = self.store.texts(category) + n_calls = max(self.keyword_gen_calls, self.min_batch) + per_call = max(1, -(-n_want // n_calls) + 4) + avoid_note = '' + if known: + shown = known if len(known) <= 40 else self.rng.sample(known, 40) + avoid_note = ('\nDo NOT repeat any of these already-used topics: ' + + ', '.join(shown)) + base = self.prompts.keyword_user.format( + k=per_call, desc=self.category_desc[category]) + avoid_note + self._nonce += 1 + prompts = [{ + 'messages': [{'role': 'system', 'content': self.prompts.keyword_system}, + {'role': 'user', 'content': f'{base}\n(batch {self._nonce}-{i})'}], + } for i in range(n_calls)] + seen = {t.strip().lower() for t in known} + out: List[str] = [] + for reply in self.explore(prompts, sampling_params=self.keyword_params): + for kw in parse_keyword_list(assistant_text(reply)): + key = kw.lower() + if key not in seen: + seen.add(key) + out.append(kw) + self.rng.shuffle(out) + return out[:n_want] + + # ------------------------------------------------------------ feedback + + def expand_hard_keywords(self) -> int: + """Brainstorm more topics in families that produced the hardest tasks.""" + if self.store is None or not self._hard or not hasattr(self.prompts, 'keyword_expand_user'): + return 0 + self.prompts.require('keyword_expand_user') + hard = self._hard[:32] + self.rng.shuffle(hard) + reqs = list(hard) + while len(reqs) < self.min_batch: + reqs.append(hard[len(reqs) % len(hard)]) + self._nonce += 1 + prompts = [{ + 'messages': [ + {'role': 'system', 'content': self.prompts.keyword_system}, + {'role': 'user', + 'content': self.prompts.keyword_expand_user.format(kw=kw, m=8) + + f'\n(batch {self._nonce}-{i})'}, + ], + } for i, (_c, kw) in enumerate(reqs)] + added = 0 + for (cat, kw), reply in zip(reqs, self.explore(prompts, + sampling_params=self.keyword_params)): + added += self.store.add(cat, parse_keyword_list(assistant_text(reply)), + source='expand', parent=kw) + logger.info(f'[AgenticChallenger] expanded {len(hard)} hard keyword(s) -> ' + f'+{added} same-domain topics') + return added diff --git a/src/twinkle_agentic/challenger/base.py b/src/twinkle_agentic/challenger/base.py new file mode 100644 index 000000000..727e99bdd --- /dev/null +++ b/src/twinkle_agentic/challenger/base.py @@ -0,0 +1,379 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Challenger: turn raw material into training tasks. + +A challenger invents the problems a solver will later be trained on. The three +things that vary between deployments are all injected: + +* **what to ask for** -- the system prompt, and the parser that reads the + answer back. They are one contract, so they are passed together. +* **how to explore** -- an :class:`Explorer`, i.e. anything that takes a batch + of trajectories and returns them with the model's reply appended. Both + rollouts in :mod:`twinkle_agentic.rollout` have that signature already, so a + challenger can explore *with tools* -- running code, reading files -- while + it invents, over a local sampler or over an HTTP endpoint alike. + :func:`twinkle_agentic.rollout.build_rollout` picks the right one for the + backend at hand. +* **what counts as a keeper** -- subclasses decide, in :meth:`Challenger.build`. +* **how hard is hard enough** -- optional. Ask for ``solver_rollouts`` attempts per + candidate and only tasks the model solves *sometimes* are kept: a task every + attempt gets right, or none does, gives GRPO a zero gradient, so it costs a + training slot and teaches nothing. Counting the attempts is the same work in + every domain and lives here; deciding whether one attempt was right is not, + and is left to :meth:`Challenger.judge_attempt`. + +Everything a strategy needs beyond that (seed examples, keyword banks) goes in +``__init__``; :meth:`Challenger.__call__` only says how many tasks you want per +batch. It is a generator that yields *full* batches: challengers throw away +most of what they propose -- keep rates of a few percent are normal once the +difficulty filter runs -- so the alternative is a caller that has to cope with +ragged batches for reasons that have nothing to do with it. +""" +import math +import random +from abc import ABC, abstractmethod +from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence + +from twinkle.data_format import SamplingParams, Trajectory, pack_user_data +from twinkle.utils import get_logger + +logger = get_logger() + +__all__ = ['Challenger', 'Explorer', 'assistant_text', 'attach_user_data'] + +# A batch of trajectories in, the same trajectories with the model's reply +# appended out. Both MultiTurnRollout and APIMultiTurnRollout satisfy this +# as-is; build_rollout() returns whichever fits the backend. Both also accept a +# per-call ``sampling_params=`` keyword, which is how the difficulty stage asks +# for its own temperature and length budget without a second explorer. +Explorer = Callable[[List[Trajectory]], List[Trajectory]] + + +def attach_user_data(trajectory: Trajectory, **values: Any) -> Trajectory: + """Return ``trajectory`` with ``values`` merged into its packed ``user_data``. + + ``user_data`` is a list of ``(key, json_string)`` pairs rather than a dict, + so it cannot be updated in place with ``update()``; going through + :func:`pack_user_data` keeps it in the one shape readers understand. + """ + merged: Dict[str, Any] = {} + for entry in trajectory.get('user_data') or []: + if isinstance(entry, (list, tuple)) and len(entry) == 2: + merged[entry[0]] = entry[1] + merged.update(values) + out = dict(trajectory) + out['user_data'] = pack_user_data(merged) + return out + + +def assistant_text(trajectory: Trajectory) -> str: + """The last assistant message's text, or '' if the model produced none. + + Explorers differ in what else they attach -- token ids, logprobs, tool + turns -- but every one of them leaves the reply as an assistant message, + so this is the one field a parser can rely on. + """ + for message in reversed(trajectory.get('messages') or []): + if isinstance(message, dict) and message.get('role') == 'assistant': + return message.get('content') or '' + return '' + + +class Challenger(ABC): + """Base class: propose, explore, keep, repeat until the batch is full. + + Args: + explorer: takes a batch of trajectories and returns them with the + model's reply appended -- a rollout from + :func:`twinkle_agentic.rollout.build_rollout`, over a local sampler + or over an API endpoint. + system: system prompt handed to the model. It carries the output + contract, which is why ``build`` -- the code that reads that output + back -- lives in the same subclass. + max_proposals_per_round: ceiling on how many proposals one round may + request. Without it a low keep rate makes the estimator ask for an + unbounded batch after the first round. + solver_rollouts: attempts per candidate in the difficulty stage. ``0`` + skips the stage entirely; any other value requires the subclass to + implement :meth:`solver_prompt` and :meth:`judge_attempt`. + keep_min_pass: keep a candidate only if at least this many attempts + succeeded. The default drops tasks nobody solved. + keep_max_pass_margin: keep a candidate only if at most + ``solver_rollouts - keep_max_pass_margin`` attempts succeeded. The + default drops tasks everybody solved. + solver_params: sampling params for the difficulty stage only, passed to + the explorer per call. ``None`` reuses whatever the explorer was + built with -- which is usually the proposing temperature, and that + is higher than a solver should get. + solver_explorer: optional separate explorer for the difficulty stage. + ``None`` reuses the main explorer. Useful when the solver needs a + different configuration (e.g. sandbox tools, more turns) than the + proposer. + seed: RNG seed for whatever sampling a subclass does. ``None`` leaves + the RNG unseeded. + """ + + def __init__( + self, + explorer: Explorer, + *, + system: str, + max_proposals_per_round: int = 512, + solver_rollouts: int = 0, + keep_min_pass: int = 1, + keep_max_pass_margin: int = 1, + solver_params: Optional[SamplingParams] = None, + solver_explorer: Optional[Explorer] = None, + seed: Optional[int] = None, + ): + if not system: + raise ValueError('Challenger needs a system prompt: it carries the output ' + 'contract that build() parses back.') + if solver_rollouts < 0: + raise ValueError(f'solver_rollouts must be >= 0, got {solver_rollouts}') + if solver_rollouts: + # Checked here rather than at first use: the stage runs after a full + # round of generation, and finding out then that this challenger + # cannot grade an attempt wastes the whole round. + missing = [ + name for name in ('solver_prompt', 'judge_attempt') + if getattr(type(self), name) is getattr(Challenger, name) + ] + if missing: + raise NotImplementedError( + f'solver_rollouts={solver_rollouts} needs {type(self).__name__} to ' + f'implement {", ".join(missing)}; pass solver_rollouts=0 to skip the ' + f'difficulty stage.') + if keep_min_pass > solver_rollouts - keep_max_pass_margin: + raise ValueError( + f'difficulty band is empty: keep_min_pass={keep_min_pass} > ' + f'solver_rollouts - keep_max_pass_margin = ' + f'{solver_rollouts - keep_max_pass_margin}') + self.explorer = explorer + self.system = system + self.max_proposals_per_round = max_proposals_per_round + self.solver_rollouts = solver_rollouts + self.keep_min_pass = keep_min_pass + self.keep_max_pass_margin = keep_max_pass_margin + self.solver_params = solver_params + self.solver_explorer = solver_explorer + self.rng = random.Random(seed) + # Running tally, used to size the next round and worth logging: a keep + # rate near zero means the prompt or the filter is miscalibrated, not + # that the model is bad. + self.n_proposed = 0 + self.n_kept = 0 + + # ------------------------------------------------------------- subclass + + @abstractmethod + def propose(self, count: int) -> List[Trajectory]: + """Build ``count`` prompt trajectories to hand to the explorer. + + Returning fewer than asked is allowed and means the source material ran + out; :meth:`__call__` stops once a round proposes nothing. + """ + + @abstractmethod + def build(self, explored: List[Trajectory]) -> List[Optional[Trajectory]]: + """Turn explored proposals into finished tasks. + + Returns one entry per input, ``None`` for anything rejected -- failed + parse, failed verification, wrong difficulty. Positional so a subclass + can line rejects up against what produced them. + """ + + def solver_prompt(self, task: Trajectory) -> Trajectory: + """The trajectory to hand a solver attempting ``task``. + + Only called when ``solver_rollouts`` is non-zero. It must return a + prompt for every task: a task that cannot be attempted has no measurable + difficulty and should have been rejected in :meth:`build` instead. + """ + raise NotImplementedError() + + def judge_attempt(self, task: Trajectory, attempt: Trajectory) -> bool: + """Did this solver attempt solve ``task``? + + ``attempt`` is the explored :meth:`solver_prompt` trajectory, so the + model's answer is its last assistant message. Program checks only: a + judgement that drifts between rounds turns the difficulty band into + noise. + """ + raise NotImplementedError() + + def on_difficulty_measured(self, candidates: List[Trajectory]) -> None: + """Called once per round with every measured candidate, before filtering. + + Each carries ``n_pass`` / ``n_rollouts`` in its ``user_data``. This is + the only place that sees the candidates the band is about to drop, which + is what a strategy adapting to difficulty needs -- an all-fail task says + more about its source material than a kept one does. + """ + + # ---------------------------------------------------------------- public + + def __call__(self, batch_size: int, total: Optional[int] = None) -> Iterator[List[Trajectory]]: + """Yield batches of exactly ``batch_size`` finished tasks. + + Args: + batch_size: tasks per yielded batch. + total: stop after this many tasks. ``None`` runs until the source + material is exhausted, which for a from-scratch challenger + means forever -- pass a total or break out of the loop. + + The final batch is short only when the source runs out or ``total`` is + not a multiple of ``batch_size``. + """ + if batch_size <= 0: + raise ValueError(f'batch_size must be positive, got {batch_size}') + pending: List[Trajectory] = [] + produced = 0 + while total is None or produced < total: + want = batch_size if total is None else min(batch_size, total - produced) + while len(pending) < want: + kept = self._round(want - len(pending)) + if kept is None: + # Source exhausted: hand back whatever is left rather than + # spinning, and let the caller see a short final batch. + if pending: + yield pending + return + pending.extend(kept) + yield pending[:want] + produced += want + pending = pending[want:] + + # --------------------------------------------------------------- private + + def _round(self, missing: int) -> Optional[List[Trajectory]]: + """One propose/explore/build/measure cycle. ``None`` means the source is dry.""" + count = min(self._estimate(missing), self.max_proposals_per_round) + proposals = self.propose(count) + if not proposals: + return None + explored = self.explore(proposals) + built = self.build(explored) + usable = [t for t in built if t is not None] + kept = self._filter_difficulty(usable) if self.solver_rollouts else usable + self.n_proposed += len(proposals) + self.n_kept += len(kept) + band = (f', in difficulty band {len(kept)}' if self.solver_rollouts else '') + logger.info(f'[{type(self).__name__}] proposed {len(proposals)}, usable ' + f'{len(usable)}{band} (cumulative {self.n_kept}/{self.n_proposed})') + return kept + + def explore( + self, + trajectories: List[Trajectory], + sampling_params: Optional[SamplingParams] = None, + ) -> List[Trajectory]: + """Run the explorer over a batch, optionally overriding its sampling params. + + The override is only forwarded when asked for, so a plain callable + explorer keeps working; both rollouts in + :mod:`twinkle_agentic.rollout` accept it. + """ + if not trajectories: + return [] + if sampling_params is None: + return self.explorer(trajectories) + return self.explorer(trajectories, sampling_params=sampling_params) + + def _solver_explore( + self, + trajectories: List[Trajectory], + sampling_params: Optional[SamplingParams] = None, + ) -> List[Trajectory]: + """Run solver attempts through the solver explorer, or fall back to the main one. + + Subclasses that need per-attempt isolation (e.g. sandbox workspace reset) + override this rather than the whole difficulty filter. + """ + if self.solver_explorer is not None: + if sampling_params is None: + return self.solver_explorer(trajectories) + return self.solver_explorer(trajectories, sampling_params=sampling_params) + return self.explore(trajectories, sampling_params=sampling_params) + + def _filter_difficulty(self, tasks: List[Trajectory]) -> List[Trajectory]: + """Attempt each task ``solver_rollouts`` times; keep the ones in the band. + + All attempts for the whole batch go out in one explorer call: on the + sampler path that is one batched generate, and the alternative -- a call + per task -- would leave the GPUs idle between them. + """ + if not tasks: + return [] + prompts: List[Trajectory] = [] + owners: List[int] = [] + for i, task in enumerate(tasks): + prompt = self.solver_prompt(task) + for _ in range(self.solver_rollouts): + prompts.append(dict(prompt)) + owners.append(i) + + attempts = self._solver_explore(prompts, sampling_params=self.solver_params) + if len(attempts) != len(prompts): + # Counting a partial return would silently understate every affected + # task's pass count, i.e. report tasks as harder than they are. + raise RuntimeError(f'explorer returned {len(attempts)} attempts for ' + f'{len(prompts)} solver prompts; expected one per prompt.') + + passes = [0] * len(tasks) + for owner, attempt in zip(owners, attempts): + if self.judge_attempt(tasks[owner], attempt): + passes[owner] += 1 + + measured = [ + attach_user_data(task, n_pass=passes[i], n_rollouts=self.solver_rollouts) + for i, task in enumerate(tasks) + ] + self.on_difficulty_measured(measured) + high = self.solver_rollouts - self.keep_max_pass_margin + return [t for t, n in zip(measured, passes) if self.keep_min_pass <= n <= high] + + def _estimate(self, missing: int) -> int: + """How many proposals to make for ``missing`` keepers. + + The first round has nothing to go on and asks for exactly what is + missing; after that the measured keep rate scales the request. A round + that kept nothing leaves the rate at its last non-zero estimate rather + than dividing by zero. + """ + if self.n_kept <= 0: + return missing + rate = self.n_kept / max(1, self.n_proposed) + return max(missing, math.ceil(missing / rate)) + + # ------------------------------------------------------------- utilities + + def prompt_trajectory(self, user: str, **extra: Any) -> Trajectory: + """A two-message trajectory carrying this challenger's system prompt.""" + trajectory: Trajectory = { + 'messages': [ + {'role': 'system', 'content': self.system}, + {'role': 'user', 'content': user}, + ], + } + trajectory.update(extra) + return trajectory + + @staticmethod + def draw(rng: random.Random, pool: Sequence[Any], count: int) -> List[Any]: + """Draw ``count`` items with replacement; ``[]`` for an empty pool.""" + return [rng.choice(pool) for _ in range(count)] if pool else [] + + +def sampling_params_of(explorer: Any) -> Optional[SamplingParams]: + """The sampling params an explorer was built with, when it exposes them. + + Only used for logging what a run actually asked for; both explorer kinds + keep the field under the same name. + """ + params: Optional[SamplingParams] = getattr(explorer, 'sampling_params', None) + return params + + +def as_dict(trajectory: Trajectory) -> Dict[str, Any]: + """A plain dict copy, for writing a trajectory to jsonl.""" + return {k: v for k, v in trajectory.items()} diff --git a/src/twinkle_agentic/challenger/code.py b/src/twinkle_agentic/challenger/code.py new file mode 100644 index 000000000..29d57cdd4 --- /dev/null +++ b/src/twinkle_agentic/challenger/code.py @@ -0,0 +1,826 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Code challenger: invent Python problems whose ground truth was executed. + +The task is built backwards. The model writes a problem *and* a reference +solution; the solution is run to capture what each check expression actually +returns, and those captured values become the asserts. So the answer exists +before the question does, and no external labelling is involved. Two gates carry +over from earlier runs, both from real failures: + +* the reference solution must pass its own asserts, or the ground truth is noise; +* output capture uses a sentinel marker plus the exit status, never the last + stdout line, so a startup banner can never be read as a result. + +Prompt text is not here. Every string the model sees arrives in +:class:`CodePrompts`, built by whoever runs the challenger -- see +``cookbook/rsi/code/prompts.py``. What stays here is the machinery that cannot +be restated in a prompt: the sandbox, the assert capture, the constant-answer +check, the keyword bank, and how a proposal becomes a task. +""" +import json +import os +import random +import re +import resource +import shutil +import signal +import subprocess +import sys +import tempfile +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple + +from twinkle.data_format import SamplingParams, Trajectory, user_data_get +from twinkle.utils import get_logger +from .base import Challenger, Explorer, assistant_text, attach_user_data + +logger = get_logger() + +__all__ = [ + 'CodeChallenger', 'CodePrompts', 'KeywordStore', 'build_asserts', + 'extract_code', 'is_constant_answer', 'load_seeds', 'parse_challenge', + 'run_asserts', +] + +# Isolates a captured value from anything else the script prints. +_MARK = '__RSI_GT__' +_FENCE_RE = re.compile(r'```(?:python|py)?\s*\n(.*?)```', re.S) +_JSON_FENCE_RE = re.compile(r'^\s*```(?:json)?\s*|\s*```\s*$', re.I) + + +# ── sandbox ──────────────────────────────────────────────────────────────── +def extract_code(text: str) -> str: + """The last fenced code block after the thinking section, else the raw body.""" + idx = (text or '').rfind('</think>') + body = text[idx + len('</think>'):] if idx >= 0 else (text or '') + blocks = _FENCE_RE.findall(body) + return (blocks[-1] if blocks else body).strip() + + +def _run_script(script: str, timeout: int) -> Tuple[int, str]: + """Run a python script in an isolated dir, 2GB cap, killpg on timeout. + + Returns (returncode, stdout). returncode is -1 on timeout/spawn failure. + """ + tmp = tempfile.mkdtemp(prefix='rsi_ch_') + try: + with open(os.path.join(tmp, '_run.py'), 'w', encoding='utf-8') as f: + f.write(script + '\n') + env = dict(os.environ, MPLBACKEND='Agg', PYTHONHASHSEED='0', OMP_NUM_THREADS='1', + MKL_NUM_THREADS='1', TOKENIZERS_PARALLELISM='false') + env.pop('CUDA_VISIBLE_DEVICES', None) + + def _limit(): + resource.setrlimit(resource.RLIMIT_AS, (2 * 1024**3, 2 * 1024**3)) + + proc = subprocess.Popen([sys.executable, '_run.py'], cwd=tmp, env=env, + stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, + text=True, start_new_session=True, preexec_fn=_limit) + try: + out, _ = proc.communicate(timeout=timeout) + return proc.returncode, out or '' + except subprocess.TimeoutExpired: + try: + os.killpg(proc.pid, signal.SIGKILL) + except ProcessLookupError: + pass + try: + proc.communicate(timeout=5) + except Exception: + pass + return -1, '' + finally: + shutil.rmtree(tmp, ignore_errors=True) + + +def run_asserts(code: str, setup: str, asserts: List[str], timeout: int = 30) -> bool: + """True when every assert passes (exit status 0).""" + if not code.strip() or not asserts: + return False + parts = [code] + if (setup or '').strip(): + parts.append(setup) + parts.extend(asserts) + rc, _ = _run_script('\n\n'.join(parts), timeout) + return rc == 0 + + +def build_asserts(solution: str, checks: List[str], timeout: int = 30, + max_checks: int = 6) -> Optional[List[str]]: + """Run the reference solution once to capture each check's repr, then form + ``assert <check> == <captured>``. + + Returns None if the solution crashed or produced no usable output -- the + caller drops that problem. The marker plus the exit status is what makes the + capture trustworthy: a crash or a banner line can never become a value. + """ + checks = [c for c in checks if isinstance(c, str) and c.strip()][:max_checks] + if not checks: + return None + lines = [solution, ''] + for i, c in enumerate(checks): + # repr on its own line, tagged with index; a check that raises makes the + # whole script exit non-zero -> we drop the problem. Pure f-string (no %% + # formatting) so a check expression containing '%' (modulo/percent) is safe. + lines.append(f'print("{_MARK}{i}=" + repr({c}))') + rc, out = _run_script('\n'.join(lines), timeout) + if rc != 0: + return None + captured: Dict[int, str] = {} + for line in out.splitlines(): + if line.startswith(_MARK): + try: + idx_str, val = line[len(_MARK):].split('=', 1) + captured[int(idx_str)] = val + except (ValueError, IndexError): + continue + if len(captured) != len(checks): + return None + # The captured text is a repr, so it is a valid literal to compare against. + return [f'assert ({c}) == ({captured[i]})' for i, c in enumerate(checks)] + + +def _split_top_eq(s: str) -> Optional[tuple]: + """Split on the first top-level ``==``, ignoring anything inside brackets or quotes.""" + depth = 0 + quote = '' + i = 0 + while i < len(s) - 1: + c = s[i] + if quote: + if c == quote: + quote = '' + elif c in '\'"': + quote = c + elif c in '([{': + depth += 1 + elif c in ')]}': + depth -= 1 + elif depth == 0 and c == '=' and s[i + 1] == '=': + return s[:i].strip(), s[i + 2:].strip() + i += 1 + return None + + +def _expected_of(assert_line: str) -> Optional[str]: + """The value the solver actually has to produce for one assert. + + :func:`build_asserts` emits ``assert (<check>) == (<repr>)``, but a check may + itself be a comparison, giving ``assert (f(x) == 3) == (True)``. Reading the + outer side there would report 'True' and make such a problem look + constant-answer, so the inner right-hand side is used instead. An outer + ``False`` pins nothing down at all and is reported as unknown. + """ + m = re.match(r'^\s*assert\s*\((.*)\)\s*==\s*\((.*)\)\s*$', assert_line.strip()) + if not m: + return None + lhs, rhs = m.group(1).strip(), m.group(2).strip() + inner = _split_top_eq(lhs) + if rhs in ('True', 'False') and inner is not None: + return inner[1] if rhs == 'True' else None + return rhs + + +def is_constant_answer(asserts: List[str]) -> bool: + """Would ``return <one constant>`` satisfy every assert? + + Such a problem pays full reward for ignoring its own statement, so it + actively teaches the solver not to read the input. Requires at least two + asserts with a readable expectation: a single assert is trivially + 'constant', and one unreadable assert must not hide a constant set. + """ + vals = [_expected_of(a) for a in asserts] + if any(v is None for v in vals) or len(vals) < 2: + return False + return len(set(vals)) == 1 + + +# ── parsing ──────────────────────────────────────────────────────────────── +def parse_challenge(text: str, require_solution: bool = True) -> Optional[Dict[str, Any]]: + """Pull the ``{problem, solution, entry, checks}`` object out of a completion. + + ``require_solution=False`` is for the two-step flow, whose second call is + told the solution is already known and returns only the statement. + """ + body = text + idx = body.rfind('</think>') + if idx >= 0: + body = body[idx + len('</think>'):] + body = _JSON_FENCE_RE.sub('', body.strip()).strip() + # Grab the outermost {...} if there is leading/trailing prose. + start, end = body.find('{'), body.rfind('}') + if start < 0 or end <= start: + return None + try: + obj = json.loads(body[start:end + 1]) + except (ValueError, TypeError): + return None + if not isinstance(obj, dict): + return None + problem, solution, checks = obj.get('problem'), obj.get('solution'), obj.get('checks') + if not (isinstance(problem, str) and problem.strip() + and isinstance(checks, list) and checks): + return None + if require_solution: + if not (isinstance(solution, str) and solution.strip()): + return None + else: + # Told not to include a solution; if it did anyway, ignore it -- the + # caller overwrites with the code that actually ran. + solution = solution if isinstance(solution, str) else '' + if solution and '```' in solution: + solution = extract_code(solution) + return {'problem': problem.strip(), 'solution': (solution or '').strip(), + 'entry': str(obj.get('entry') or '').strip(), 'checks': checks} + + +def parse_keyword_list(text: str) -> List[str]: + """Extract a JSON array of short strings from a (possibly thinking) reply.""" + body = text + idx = body.rfind('</think>') + if idx >= 0: + body = body[idx + len('</think>'):] + start, end = body.find('['), body.rfind(']') + if start < 0 or end <= start: + return [] + try: + arr = json.loads(body[start:end + 1]) + except (ValueError, TypeError): + return [] + return [x.strip() for x in arr + if isinstance(x, str) and x.strip() and len(x.strip()) <= 60] + + +def load_seeds(path: str) -> List[Dict[str, str]]: + """Read seed problems from a jsonl: dicts with ``query`` and maybe ``code``. + + A seed without ``code`` cannot take the two-step path (there is no reference + solution to build on top of) and falls back to the single-call prompt. + """ + if not path or not os.path.exists(path): + return [] + seeds: List[Dict[str, str]] = [] + with open(path, encoding='utf-8') as f: + for line in f: + line = line.strip() + if not line: + continue + try: + row = json.loads(line) + except (ValueError, TypeError): + continue + q = row.get('query') or row.get('problem') or row.get('prompt') + if isinstance(q, dict): + q = q.get('content') + if not q: + msgs = row.get('messages') or [] + q = next((m.get('content') for m in msgs if m.get('role') == 'user'), None) + if isinstance(q, str) and q.strip(): + seeds.append({'query': q.strip(), 'code': (row.get('code') or '').strip()}) + return seeds + + +# ── keyword bank ─────────────────────────────────────────────────────────── +class KeywordStore: + """Persistent keyword bank with usage tracking, one bucket per category. + + Keywords exist to stop the challenger collapsing onto a handful of + archetypes. They are consumed rather than sampled with replacement, so a + run keeps reaching for topics it has not used; when a bucket runs dry the + caller refills it from the model, and recycles only if the model has run out + of distinct ideas. + + On-disk format (one JSON per line):: + + {"category", "text", "used": bool, "used_count": int, + "source": "gen"|"expand", "parent": <keyword or null>} + + De-duplicates case-insensitively within a category, so re-runs never + conflict with the bank on disk. + """ + + def __init__(self, path: str, categories: Sequence[str]): + if not categories: + raise ValueError('KeywordStore needs at least one category') + self.path = path + self.categories = tuple(categories) + self.items: Dict[str, List[Dict[str, Any]]] = {c: [] for c in self.categories} + self._seen: Dict[str, set] = {c: set() for c in self.categories} + if path and os.path.exists(path): + with open(path, encoding='utf-8') as f: + for line in f: + line = line.strip() + if not line: + continue + try: + r = json.loads(line) + except (ValueError, TypeError): + continue + c, t = r.get('category'), r.get('text') + if c in self.items and isinstance(t, str) and t.strip(): + key = t.strip().lower() + if key not in self._seen[c]: + self._seen[c].add(key) + self.items[c].append(r) + + def save(self) -> None: + os.makedirs(os.path.dirname(os.path.abspath(self.path)) or '.', exist_ok=True) + tmp = self.path + '.tmp' + with open(tmp, 'w', encoding='utf-8') as f: + for c in self.categories: + for r in self.items[c]: + f.write(json.dumps(r, ensure_ascii=False) + '\n') + os.replace(tmp, self.path) + + def add(self, category: str, texts: List[str], source: str = 'gen', + parent: Optional[str] = None) -> int: + added = 0 + for t in texts: + key = t.strip().lower() + if not key or key in self._seen[category]: + continue + self._seen[category].add(key) + self.items[category].append({'category': category, 'text': t.strip(), + 'used': False, 'used_count': 0, + 'source': source, 'parent': parent}) + added += 1 + return added + + def unused(self, category: str) -> List[Dict[str, Any]]: + return [r for r in self.items[category] if not r.get('used')] + + def texts(self, category: str) -> List[str]: + return [r['text'] for r in self.items[category]] + + def take(self, category: str, rng: random.Random) -> Optional[str]: + """Consume one unused keyword from ``category``; None if it is dry.""" + un = self.unused(category) + if not un: + return None + r = rng.choice(un) + r['used'] = True + r['used_count'] = r.get('used_count', 0) + 1 + return r['text'] + + def recycle(self, category: str) -> None: + """Mark every keyword unused again (safety valve when the model is tapped out).""" + for r in self.items[category]: + r['used'] = False + + +# ── prompts (text supplied by the caller) ────────────────────────────────── +@dataclass +class CodePrompts: + """Every string a :class:`CodeChallenger` sends, and nothing else. + + Deliberately without defaults for the always-needed fields: a prompt is the + experiment, so a run has to state which one it used rather than inherit a + library's idea of it. Optional groups stay empty until the feature that + needs them is switched on, and the constructor says so if one is missing. + + Placeholders are checked at construction: a typo'd ``{keywords}`` would + otherwise surface as a KeyError halfway through a generation run. + """ + + system: str + from_scratch: str + solver_system: str + solver_user: str + from_seed: str = '' + from_keywords: str = '' + from_seed_keywords: str = '' + two_step_system: str = '' + two_step_solution: str = '' + two_step_problem: str = '' + keyword_system: str = '' + keyword_user: str = '' + keyword_expand_user: str = '' + + #: field -> placeholders it must contain. + _REQUIRED_FIELDS = { + 'solver_user': ('problem', ), + 'from_seed': ('seed', ), + 'from_keywords': ('keywords', ), + 'from_seed_keywords': ('seed', 'keywords'), + 'two_step_solution': ('seed', 'code', 'keywords'), + 'two_step_problem': ('code', 'seed', 'keywords'), + 'keyword_user': ('k', 'desc'), + 'keyword_expand_user': ('kw', 'm'), + } + + def __post_init__(self): + for name in ('system', 'from_scratch', 'solver_system', 'solver_user'): + if not getattr(self, name).strip(): + raise ValueError(f'CodePrompts.{name} is required') + for name, placeholders in self._REQUIRED_FIELDS.items(): + text = getattr(self, name) + if not text: + continue + for placeholder in placeholders: + if '{' + placeholder + '}' not in text: + raise ValueError(f'CodePrompts.{name} must contain ' + f'{{{placeholder}}}') + + def require(self, *names: str) -> None: + """Raise unless every named prompt was supplied.""" + missing = [n for n in names if not getattr(self, n).strip()] + if missing: + raise ValueError(f'this configuration needs CodePrompts.' + f'{", CodePrompts.".join(missing)}') + + +class CodeChallenger(Challenger): + """Propose code problems, execute them for ground truth, keep the graded ones. + + One class rather than several because 'from scratch', 'from a seed problem', + 'from keywords' and the two-step build differ only in which prompt the + proposal carries: parsing, execution, the self-check and the difficulty + band are the same afterwards. Which path a proposal takes is decided per + proposal, so one run mixes them. + + Args: + prompts: every string sent to the model. + explorer: batch-in / batch-out generation, see :class:`.base.Explorer`. + seeds: optional pool from :func:`load_seeds`, drawn with replacement. + keyword_store: optional bank; without it proposals carry no topics. + category_desc: category -> description used when asking for more + keywords. Keys must cover the store's categories. + seed_mix_prob: chance a proposal also carries a seed problem, when a + pool was given. + two_step: allow the two-call path (write a harder solution on top of the + seed's reference code, then describe the problem it answers). Needs + a seed carrying ``code`` and at least one keyword, so it is skipped + silently for proposals that have neither. + combo_arity: ``'triple'`` takes one keyword per category; ``'mix'`` + takes a random 1..len(categories) subset. + arity_weights: sampling weights for the ``'mix'`` subset size. + single_kw_prob: in ``'triple'`` mode, the chance of using one category + instead of all of them. + keyword_refill_target / keyword_gen_calls / keyword_refill_tries / + keyword_params: how a dry category is refilled from the model. + min_batch: smallest batch worth sending -- a sampler shards a batch over + its data-parallel workers, and a batch smaller than that leaves some + with nothing to do. Set it to the number of sampler workers. + problem_max_chars: reject statements longer than this. Rambling + non-problems, and they would also crowd out the solver's context. + max_checks / sandbox_timeout: passed to :func:`build_asserts`. + drop_constant_answer: reject problems where one constant satisfies every + assert. + low_pass_expand / expand_per_kw / expand_max_kws: feedback for + :meth:`expand_hard_keywords`. + reject_sink: called with a dict for every rejected proposal. The caller + decides whether that goes to a file; nothing here writes one. + """ + + def __init__( + self, + prompts: CodePrompts, + explorer: Explorer, + *, + seeds: Sequence[Dict[str, str]] = (), + keyword_store: Optional[KeywordStore] = None, + category_desc: Optional[Dict[str, str]] = None, + seed_mix_prob: float = 0.5, + two_step: bool = True, + combo_arity: str = 'triple', + arity_weights: Optional[Sequence[float]] = None, + single_kw_prob: float = 0.1, + keyword_refill_target: int = 128, + keyword_gen_calls: int = 8, + keyword_refill_tries: int = 2, + keyword_params: Optional[SamplingParams] = None, + min_batch: int = 1, + problem_max_chars: int = 4000, + max_checks: int = 6, + sandbox_timeout: int = 30, + drop_constant_answer: bool = True, + low_pass_expand: int = 0, + expand_per_kw: int = 8, + expand_max_kws: int = 32, + reject_sink: Optional[Callable[[Dict[str, Any]], None]] = None, + **challenger_kwargs: Any, + ): + super().__init__(explorer, system=prompts.system, **challenger_kwargs) + if combo_arity not in ('triple', 'mix'): + raise ValueError(f"combo_arity must be 'triple' or 'mix', got {combo_arity!r}") + if keyword_store is not None: + desc = category_desc or {} + missing = [c for c in keyword_store.categories if not desc.get(c)] + if missing: + raise ValueError(f'category_desc is missing a description for ' + f'{missing}; a dry category could not be refilled.') + prompts.require('keyword_system', 'keyword_user', 'from_keywords') + self.prompts = prompts + self.seeds = list(seeds) + self.store = keyword_store + self.category_desc = dict(category_desc or {}) + self.seed_mix_prob = seed_mix_prob + self.two_step = two_step + self.combo_arity = combo_arity + self.arity_weights = list(arity_weights) if arity_weights else None + self.single_kw_prob = single_kw_prob + self.keyword_refill_target = keyword_refill_target + self.keyword_gen_calls = keyword_gen_calls + self.keyword_refill_tries = keyword_refill_tries + self.keyword_params = keyword_params + self.min_batch = max(1, min_batch) + self.problem_max_chars = problem_max_chars + self.max_checks = max_checks + self.sandbox_timeout = sandbox_timeout + self.drop_constant_answer = drop_constant_answer + self.low_pass_expand = low_pass_expand + self.expand_per_kw = expand_per_kw + self.expand_max_kws = expand_max_kws + self.reject_sink = reject_sink + if self.seeds: + # Both are reachable with a bank configured: a proposal draws no + # keywords when every category is dry, and then falls back to the + # seed-only prompt. + prompts.require('from_seed') + if self.store is not None: + prompts.require('from_seed_keywords') + if two_step: + prompts.require('two_step_system', 'two_step_solution', 'two_step_problem') + # Perturbs refill prompts so a second ask does not repeat the first. + self._nonce = 0 + # Why proposals died, for the caller to log; the shape a run is judged on. + self.stats: Dict[str, int] = { + 'parsed': 0, 'parse_fail': 0, 'stage1_no_code': 0, 'too_long': 0, + 'gt_fail': 0, 'selfcheck_fail': 0, 'constant_answer': 0, + } + # (category, keyword) behind candidates nobody could solve, for feedback. + self._hard: List[Tuple[str, str]] = [] + + # ------------------------------------------------------------- proposing + + def propose(self, count: int) -> List[Trajectory]: + proposals: List[Trajectory] = [] + for _ in range(count): + picks = self._draw_keywords() + body = '\n'.join(f'- {c}: {t}' for c, t in picks) + use_seed = bool(self.seeds) and self.rng.random() < self.seed_mix_prob + seed = self.rng.choice(self.seeds) if use_seed else None + two = bool(use_seed and self.two_step and picks and seed and seed.get('code')) + if two: + system = self.prompts.two_step_system + user = self.prompts.two_step_solution.format( + seed=seed['query'], code=seed['code'], keywords=body) + elif use_seed and picks: + system = self.prompts.system + user = self.prompts.from_seed_keywords.format(seed=seed['query'], keywords=body) + elif use_seed: + system = self.prompts.system + user = self.prompts.from_seed.format(seed=seed['query']) + elif picks: + system = self.prompts.system + user = self.prompts.from_keywords.format(keywords=body) + else: + system = self.prompts.system + user = self.prompts.from_scratch + proposal: Trajectory = { + 'messages': [{'role': 'system', 'content': system}, + {'role': 'user', 'content': user}], + } + # Carried through the explorer so build() knows which path this + # proposal took and what the second call has to be told. + proposals.append(attach_user_data( + proposal, keywords=picks, seeded=use_seed, two_step=two, + seed_query=(seed['query'] if two else ''), keyword_block=body)) + return proposals + + def _draw_keywords(self) -> List[Tuple[str, str]]: + """Consume one keyword combination from the bank; [] without a bank.""" + if self.store is None: + return [] + categories = self.store.categories + if self.combo_arity == 'mix': + if self.arity_weights and len(self.arity_weights) == len(categories): + k = self.rng.choices(range(1, len(categories) + 1), + weights=self.arity_weights)[0] + else: + k = self.rng.randint(1, len(categories)) + cats = self.rng.sample(list(categories), k) + elif self.rng.random() < self.single_kw_prob: + cats = [self.rng.choice(categories)] + else: + cats = list(categories) + picks: List[Tuple[str, str]] = [] + for c in cats: + if not self.store.unused(c): + self._refill(c) + text = self.store.take(c, self.rng) + if text is not None: + picks.append((c, text)) + return picks + + def _refill(self, category: str) -> None: + """Ask the model for more keywords in ``category``; recycle if it is tapped out.""" + tries = 0 + while not self.store.unused(category): + new = self._generate_keywords(category, self.keyword_refill_target) + added = self.store.add(category, new, source='gen') + tries += 1 + if added == 0 and tries >= self.keyword_refill_tries: + if self.store.items[category]: + self.store.recycle(category) + logger.info(f'[CodeChallenger] keyword category {category!r} exhausted ' + f'-> recycled {len(self.store.items[category])} topics') + break + + def _generate_keywords(self, category: str, n_want: int) -> List[str]: + """Up to ``n_want`` keywords the bank does not already hold.""" + if n_want <= 0: + return [] + known = self.store.texts(category) + n_calls = max(self.keyword_gen_calls, self.min_batch) + per_call = max(1, -(-n_want // n_calls) + 4) # ceil(n/calls) + margin + avoid_note = '' + if known: + shown = known if len(known) <= 40 else self.rng.sample(known, 40) + avoid_note = ('\nDo NOT repeat any of these already-used topics: ' + + ', '.join(shown)) + base = self.prompts.keyword_user.format( + k=per_call, desc=self.category_desc[category]) + avoid_note + self._nonce += 1 + prompts = [{ + 'messages': [{'role': 'system', 'content': self.prompts.keyword_system}, + {'role': 'user', 'content': f'{base}\n(batch {self._nonce}-{i})'}], + } for i in range(n_calls)] + seen = {t.strip().lower() for t in known} + out: List[str] = [] + for reply in self.explore(prompts, sampling_params=self.keyword_params): + for kw in parse_keyword_list(assistant_text(reply)): + key = kw.lower() + if key not in seen: + seen.add(key) + out.append(kw) + self.rng.shuffle(out) + return out[:n_want] + + # ---------------------------------------------------------------- building + + def build(self, explored: List[Trajectory]) -> List[Optional[Trajectory]]: + """Parse, execute, self-check; None for every proposal that did not survive. + + The second call of the two-step path happens here rather than in + :meth:`propose`, because it needs the code the first call produced. It + goes out as one batch for the whole round, so the extra call costs one + more generate, not one per proposal. + """ + objs: List[Optional[Dict[str, Any]]] = [None] * len(explored) + # Proposals that died before parsing: they must not also be counted as a + # parse failure, because the cause -- and the fix -- is a different one. + dead: List[bool] = [False] * len(explored) + stage2_idx: List[int] = [] + stage2_prompts: List[Trajectory] = [] + for i, traj in enumerate(explored): + text = assistant_text(traj) + if not user_data_get(traj.get('user_data'), 'two_step', False): + objs[i] = parse_challenge(text) + continue + code = extract_code(text) + if not code.strip(): + # Usually a truncated completion: there is no solution to + # describe, so this proposal ends here. + self.stats['stage1_no_code'] += 1 + dead[i] = True + continue + objs[i] = {'_stage1_code': code} + stage2_idx.append(i) + stage2_prompts.append({ + 'messages': [ + {'role': 'system', 'content': self.prompts.system}, + {'role': 'user', 'content': self.prompts.two_step_problem.format( + code=code, + seed=user_data_get(explored[i].get('user_data'), 'seed_query', ''), + keywords=user_data_get(explored[i].get('user_data'), + 'keyword_block', ''))}, + ], + }) + if stage2_prompts: + logger.info(f'[CodeChallenger] two-step stage 2: {len(stage2_prompts)} problem ' + f'writes ({self.stats["stage1_no_code"]} first calls had no code)') + for i, reply in zip(stage2_idx, self.explore(stage2_prompts)): + stage1_code = objs[i]['_stage1_code'] + obj = parse_challenge(assistant_text(reply), require_solution=False) + if obj is not None: + # Ground truth is the code that actually ran, never the one + # the second call may have re-imagined. + obj['solution'] = stage1_code + objs[i] = obj + + return [None if dead[i] else self._finish(explored[i], obj) + for i, obj in enumerate(objs)] + + def _finish(self, proposal: Trajectory, obj: Optional[Dict[str, Any]]) -> Optional[Trajectory]: + """One parsed proposal -> a task, or None with a reason recorded.""" + if obj is None: + self.stats['parse_fail'] += 1 + return None + self.stats['parsed'] += 1 + + def _reject(reason: str, **extra: Any) -> None: + self.stats[reason] += 1 + if self.reject_sink is not None: + self.reject_sink({'reason': reason, **extra, **obj}) + + if len(obj['problem']) > self.problem_max_chars: + _reject('too_long') + return None + asserts = build_asserts(obj['solution'], obj['checks'], + timeout=self.sandbox_timeout, max_checks=self.max_checks) + if not asserts: + _reject('gt_fail') + return None + if not run_asserts(obj['solution'], '', asserts, timeout=self.sandbox_timeout): + # A reference solution that fails its own asserts is not ground + # truth, whatever the statement says. + _reject('selfcheck_fail', asserts=asserts) + return None + if self.drop_constant_answer and is_constant_answer(asserts): + _reject('constant_answer', asserts=asserts) + return None + + user_data = proposal.get('user_data') + # The task the solver is trained on: the statement alone, exactly as the + # difficulty stage will present it, with no instructions from the + # challenger's own prompt leaking in. + task: Trajectory = { + 'messages': [{'role': 'system', 'content': self.prompts.solver_system}, + {'role': 'user', 'content': obj['problem']}], + } + return attach_user_data( + task, + asserts=asserts, + solution=obj['solution'], + entry=obj['entry'], + keywords=user_data_get(user_data, 'keywords', []), + seeded=user_data_get(user_data, 'seeded', False), + two_step=user_data_get(user_data, 'two_step', False)) + + # -------------------------------------------------------------- difficulty + + def solver_prompt(self, task: Trajectory) -> Trajectory: + problem = next((m['content'] for m in reversed(task.get('messages') or []) + if m.get('role') == 'user'), '') + return { + 'messages': [{'role': 'system', 'content': self.prompts.solver_system}, + {'role': 'user', + 'content': self.prompts.solver_user.format(problem=problem)}], + } + + def judge_attempt(self, task: Trajectory, attempt: Trajectory) -> bool: + asserts = user_data_get(task.get('user_data'), 'asserts', []) or [] + return run_asserts(extract_code(assistant_text(attempt)), '', asserts, + timeout=self.sandbox_timeout) + + def on_difficulty_measured(self, candidates: List[Trajectory]) -> None: + """Remember the topics behind the candidates nobody solved.""" + if self.store is None: + return + seen = {(c, t.lower()) for c, t in self._hard} + for task in candidates: + data = task.get('user_data') + if user_data_get(data, 'n_pass', 0) > self.low_pass_expand: + continue + for pick in user_data_get(data, 'keywords', []) or []: + c, t = pick[0], pick[1] + if (c, t.lower()) not in seen: + seen.add((c, t.lower())) + self._hard.append((c, t)) + + # ------------------------------------------------------------- feedback + + def expand_hard_keywords(self) -> int: + """Brainstorm more topics in the families that produced the hardest tasks. + + Called by whoever drives the challenger, after generating, so the bank + drifts toward material the solver actually struggles with. Returns how + many new keywords were added. + """ + if self.store is None or not self._hard or self.expand_per_kw <= 0: + return 0 + self.prompts.require('keyword_expand_user') + hard = self._hard[:self.expand_max_kws] + self.rng.shuffle(hard) + # Cycle a short list so the batch still covers every sampler worker. + reqs = list(hard) + while len(reqs) < self.min_batch: + reqs.append(hard[len(reqs) % len(hard)]) + self._nonce += 1 + prompts = [{ + 'messages': [ + {'role': 'system', 'content': self.prompts.keyword_system}, + {'role': 'user', + 'content': self.prompts.keyword_expand_user.format(kw=kw, m=self.expand_per_kw) + + f'\n(batch {self._nonce}-{i})'}, + ], + } for i, (_c, kw) in enumerate(reqs)] + added = 0 + for (cat, kw), reply in zip(reqs, self.explore(prompts, + sampling_params=self.keyword_params)): + added += self.store.add(cat, parse_keyword_list(assistant_text(reply)), + source='expand', parent=kw) + logger.info(f'[CodeChallenger] expanded {len(hard)} hard keyword(s) -> ' + f'+{added} same-domain topics') + return added diff --git a/src/twinkle_agentic/envs/__init__.py b/src/twinkle_agentic/envs/__init__.py index 9eaacf844..4633039c8 100644 --- a/src/twinkle_agentic/envs/__init__.py +++ b/src/twinkle_agentic/envs/__init__.py @@ -2,5 +2,4 @@ from .agentenv import AgentEnv from .base import Env, StepResult from .env_tool import EnvTool -from .ms_agent_tool_env import MsAgentToolEnv from .openenv import EnvPool, EnvPoolAdapter, OpenEnv, OpenEnvClient diff --git a/src/twinkle_agentic/envs/ms_agent_tool_env.py b/src/twinkle_agentic/envs/ms_agent_tool_env.py deleted file mode 100644 index 65c11a26e..000000000 --- a/src/twinkle_agentic/envs/ms_agent_tool_env.py +++ /dev/null @@ -1,212 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Env backed by an ms-agent tool runtime. - -The harness declares which tools exist (``AgentHarness.tool_schemas``); this -Env is where those calls actually run. It owns no tool logic of its own -- it -forwards to the ``ToolManager`` that ms-agent already built (filesystem, shell, -python, notebook sandbox, web search, todo list) and adds the two things RL -needs on top: one workspace per episode, and batched dispatch so a turn's tool -calls do not run one at a time while the GPUs idle. - -Isolation comes from ``config.output_dir``: both ``FileSystemTool`` and -``CodeExecutionTool`` root themselves there, so giving every trajectory its own -directory keeps concurrent episodes from reading each other's files. - -Reward is deliberately absent. ``step`` always reports ``reward=0.0`` and -``done=False``; an agentic episode is scored after the fact from the state it -left behind (:mod:`twinkle_agentic.verifier.result_check`), and the rollout -ends when the model stops emitting tool calls. -""" -import json -import os -import re -from typing import Any, Dict, List, Optional, Sequence, Tuple - -from twinkle.utils import run_sync - -from .base import Env, StepResult - -__all__ = ['MsAgentToolEnv'] - -# Marker used to recover an exit status from a tool that only returns text. -_RC_MARK = '__TWINKLE_RC__' -_RC_RE = re.compile(rf'{_RC_MARK}:(-?\d+)') - -_PY_WRAPPER = """\ -import sys, traceback -try: -{body} -except SystemExit as _e: - print('{mark}:%d' % (_e.code or 0)) - sys.exit(0) -except BaseException: - traceback.print_exc() - print('{mark}:1') -else: - print('{mark}:0') -""" - - -class MsAgentToolEnv(Env): - """Execute ms-agent tool calls for one episode. - - Args: - agent: an ``LLMAgent`` whose tools are already prepared -- normally - ``MsAgentHarness.agent`` after ``harness.prepare()``. Sharing the - harness's agent is what guarantees the executing tool set is the - one the prompt advertised. - tool_manager: an ms-agent ``ToolManager`` to use instead of the - agent's. Only one of ``agent`` / ``tool_manager`` is needed. - workspace: directory this episode reads and writes. Defaults to the - agent's ``config.output_dir``. - max_observation_chars: truncate a tool result before it becomes a - message. A single ``grep`` can otherwise blow the context window - and truncate the trajectory mid-episode. - """ - - def __init__( - self, - agent: Any = None, - *, - tool_manager: Any = None, - workspace: str = '', - max_observation_chars: int = 8000, - ): - if agent is None and tool_manager is None: - raise ValueError('MsAgentToolEnv needs either agent= or tool_manager=') - self._agent = agent - self._tm = tool_manager if tool_manager is not None else getattr(agent, 'tool_manager', None) - if self._tm is None: - raise ValueError('no ms-agent ToolManager available; call harness.prepare() ' - 'before constructing the Env so tools are initialised') - self.workspace = workspace or self._workspace_from_agent(agent) - if self.workspace: - os.makedirs(self.workspace, exist_ok=True) - self.max_observation_chars = max_observation_chars - self._names: Optional[List[str]] = None - - # ------------------------------------------------------------------ Env - - def tool_names(self) -> List[str]: - """Names of the tools actually registered, as the runtime spells them.""" - if self._names is None: - raw = run_sync(self._tm.get_tools) - items: List[Any] = [] - if isinstance(raw, dict): - for value in raw.values(): - items.extend(value if isinstance(value, list) else [value]) - elif isinstance(raw, list): - items = raw - names = [] - for item in items: - if isinstance(item, dict): - fn = item.get('function') - name = (fn or {}).get('name') if isinstance(fn, dict) else None - name = name or item.get('tool_name') or item.get('name') - if name: - names.append(str(name)) - self._names = names - return list(self._names) - - def resolve_tool(self, name: str) -> str: - """Map a plain tool name onto the runtime's own spelling. - - ms-agent namespaces its tools as ``{server}---{tool}``, so a caller that - asks for ``shell_executor`` means ``code_executor---shell_executor``. - An unknown name raises instead of being passed through: a mistyped tool - comes back as a failed call, which for a checker is indistinguishable - from a failed check, and a whole GRPO group would silently score zero. - """ - names = self.tool_names() - if name in names: - return name - matches = [n for n in names if n.rsplit('---', 1)[-1] == name] - if len(matches) == 1: - return matches[0] - if not matches: - raise ValueError(f'no registered tool named {name!r}; available: {names}') - raise ValueError(f'{name!r} is ambiguous across servers: {matches}') - - def step(self, tool_name: str, arguments: Dict[str, Any]) -> StepResult: - result = run_sync(self._tm.single_call_tool, self._call(tool_name, arguments)) - return StepResult(observation=self._observation(result)) - - def step_batch(self, calls: Sequence[Tuple[str, Dict[str, Any]]]) -> List[StepResult]: - """Run a turn's calls concurrently through ms-agent's own gather.""" - calls = list(calls) - if not calls: - return [] - if len(calls) == 1: - return [self.step(calls[0][0], calls[0][1] or {})] - payload = [self._call(name, args or {}) for name, args in calls] - results = run_sync(self._tm.parallel_call_tool, payload) - return [StepResult(observation=self._observation(r)) for r in results] - - def close(self) -> None: - cleanup = getattr(self._tm, 'cleanup', None) - if cleanup is not None: - try: - run_sync(cleanup) - except Exception: # noqa - # Teardown must not take down a training step; a leaked sandbox - # is recoverable, a crashed trainer loses the whole batch. - pass - - # ------------------------------------------------------- for the checker - - def runner(self, shell_tool: str = 'shell_executor', python_tool: str = 'python_executor'): - """A ``result_check`` runner that executes inside this episode's sandbox. - - Verification has to see the same filesystem the agent wrote to, so the - check goes back through the same tools rather than a local subprocess. - Those tools return prose, not an exit status, so the command is made to - print a marker and the status is read back out of the output. - - The tool names are resolved against what is actually registered, so - plain names work regardless of how the runtime namespaces them. - """ - shell_name = self.resolve_tool(shell_tool) - python_name = self.resolve_tool(python_tool) - - def _run(source: str, interpreter: str) -> Tuple[int, str]: - if interpreter == 'python': - body = '\n'.join(' ' + line for line in source.splitlines()) or ' pass' - code = _PY_WRAPPER.format(body=body, mark=_RC_MARK) - out = self.step(python_name, {'code': code}).observation - else: - out = self.step(shell_name, {'command': f'{source}\necho "{_RC_MARK}:$?"'}).observation - match = _RC_RE.search(out or '') - if match is None: - # No marker means the tool itself failed (timeout, sandbox down) - # rather than the check failing; report non-zero and keep output. - return 1, out or 'check produced no output and no exit marker' - return int(match.group(1)), _RC_RE.sub('', out or '').strip() - - return _run - - # -------------------------------------------------------------- private - - @staticmethod - def _call(tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]: - return {'tool_name': tool_name, 'arguments': arguments or {}} - - def _observation(self, result: Any) -> str: - if result is None: - text = '' - elif isinstance(result, str): - text = result - else: - try: - text = json.dumps(result, ensure_ascii=False) - except (TypeError, ValueError): - text = str(result) - limit = self.max_observation_chars - if limit and len(text) > limit: - head = text[:limit] - text = f'{head}\n...[truncated {len(text) - limit} chars]' - return text - - @staticmethod - def _workspace_from_agent(agent: Any) -> str: - config = getattr(agent, 'config', None) - return str(getattr(config, 'output_dir', '') or '') if config is not None else '' diff --git a/src/twinkle_agentic/rollout/__init__.py b/src/twinkle_agentic/rollout/__init__.py index 1e839adf4..835d94da0 100644 --- a/src/twinkle_agentic/rollout/__init__.py +++ b/src/twinkle_agentic/rollout/__init__.py @@ -1,12 +1,14 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from .base import Rollout from .bridge import extend_with_bridge +from .factory import build_rollout from .multi_turn import MultiTurnRollout __all__ = [ 'APIMultiTurnRollout', 'MultiTurnRollout', 'Rollout', + 'build_rollout', 'extend_with_bridge', ] diff --git a/src/twinkle_agentic/rollout/api_multi_turn.py b/src/twinkle_agentic/rollout/api_multi_turn.py index 7521a454f..4513d361b 100644 --- a/src/twinkle_agentic/rollout/api_multi_turn.py +++ b/src/twinkle_agentic/rollout/api_multi_turn.py @@ -1,14 +1,12 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -import os from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Any, Callable, Dict, List, Optional from twinkle.data_format import Trajectory from twinkle.data_format.sampling import SamplingParams -from twinkle_agentic.protocol.openai import OpenAI +from twinkle_agentic.protocol.base import API from twinkle_agentic.tools.tool_manager import ToolManager from .base import Rollout -from .multi_turn import MultiTurnRollout # Termination reasons surfaced via ``trajectory['stop_reason']``. _STOP_NO_TOOL = 'stop' @@ -34,7 +32,8 @@ class APIMultiTurnRollout(Rollout): Constructor and per-call override semantics intentionally mirror :class:`MultiTurnRollout`: ``tool_manager`` may be a single instance - (broadcast) or a list aligned 1:1 with trajectories. + (broadcast) or a list aligned 1:1 with trajectories, and it is optional -- + a challenger inventing tasks has nothing to execute. Tool schema source: ``trajectory['tools']`` if present, else ``tool_manager.tool_infos()`` of the trajectory's manager. Caller is @@ -50,8 +49,8 @@ class APIMultiTurnRollout(Rollout): def __init__( self, - api: OpenAI, - tool_manager: ToolManager, + api: API, + tool_manager: Optional[ToolManager] = None, sampling_params: Optional[SamplingParams] = None, max_turns: int = 6, concurrency: int = 8, @@ -62,28 +61,19 @@ def __init__( ): super().__init__() if api is None: - raise ValueError('APIMultiTurnRollout requires an OpenAI client') - if tool_manager is None: - raise ValueError('APIMultiTurnRollout requires a ToolManager') - if max_turns < 1: - raise ValueError(f'max_turns must be >= 1, got {max_turns}') + raise ValueError('APIMultiTurnRollout requires an API client') if concurrency < 1: raise ValueError(f'concurrency must be >= 1, got {concurrency}') - sp = sampling_params or SamplingParams() - if sp.num_samples != 1: - raise ValueError(f'APIMultiTurnRollout supports num_samples=1 only, ' - f'got {sp.num_samples}') + self._init_common( + max_turns=max_turns, + sampling_params=sampling_params, + trace_dir=trace_dir, + trace_callback=trace_callback, + success_callback=success_callback) self.api = api self.tool_manager = tool_manager - self.sampling_params = sp - self.max_turns = max_turns self.concurrency = concurrency self.extra_body = dict(extra_body or {}) - self.trace_dir = trace_dir - self.trace_callback = trace_callback - self.success_callback = success_callback - if self.trace_dir: - os.makedirs(self.trace_dir, exist_ok=True) def __call__( self, @@ -99,7 +89,7 @@ def __call__( return [] sampling_params: SamplingParams = kwargs.get('sampling_params', self.sampling_params) - tool_managers = MultiTurnRollout._resolve_tool_managers(kwargs.get('tool_manager', self.tool_manager), n) + tool_managers = self._broadcast(kwargs.get('tool_manager', self.tool_manager), n, name='tool_manager') extra_body = dict(self.extra_body) if 'extra_body' in kwargs and kwargs['extra_body']: extra_body.update(kwargs['extra_body']) @@ -119,7 +109,7 @@ def __call__( result_outs: List[Trajectory] = [o if o is not None else dict(trajectories[i]) for i, o in enumerate(outs)] if self.trace_dir: - self._write_traces(result_outs, kwargs.get('global_step')) + self._write_rollout_traces(result_outs, global_step=kwargs.get('global_step')) return result_outs # ------------------------------------------------------------------ private @@ -127,7 +117,7 @@ def __call__( def _run_one( self, trajectory: Trajectory, - tool_manager: ToolManager, + tool_manager: Optional[ToolManager], sampling_params: SamplingParams, extra_body: Dict[str, Any], ) -> Trajectory: @@ -139,7 +129,7 @@ def _run_one( """ messages: List[Dict[str, Any]] = list(trajectory.get('messages') or []) tools = trajectory.get('tools') - if tools is None: + if tools is None and tool_manager is not None: tools = tool_manager.tool_infos() or None turn = 0 @@ -182,6 +172,15 @@ def _run_one( stop_reason = _STOP_MAX_TURNS break + if tool_manager is None: + # Nothing can run the call, so the conversation cannot continue: + # say why rather than looping on an unanswered tool turn. + stop_reason = _STOP_API_ERROR + error = ('model emitted tool_calls but this rollout has no ToolManager; ' + 'pass one at construction time or as a per-call kwarg') + truncated = True + break + try: for tc in tool_calls: response = tool_manager(tc) @@ -242,48 +241,21 @@ def _normalise_assistant(reply: Any, turn: int) -> Dict[str, Any]: msg['reasoning_content'] = reasoning return msg - def _write_traces( + def _build_trace_record( self, - outs: List[Trajectory], - global_step: Optional[int], - ) -> None: - """Per-trajectory JSON dump. Mirrors :meth:`MultiTurnRollout. - _write_rollout_traces` but reuses its static helpers — failures - on a single trajectory never abort the batch.""" - import json - import os - for idx, traj in enumerate(outs): - try: - should_store = True - if self.trace_callback is not None: - try: - should_store = bool(self.trace_callback(traj)) - except Exception: - should_store = False - if not should_store: - continue - success = False - if self.success_callback is not None: - try: - success = bool(self.success_callback(traj)) - except Exception: - success = False - record = { - 'trajectory': MultiTurnRollout._serialize_for_trace(traj), - 'ground_truth': MultiTurnRollout._extract_ground_truth(traj), - 'stop_reason': traj.get('stop_reason'), - 'truncated': bool(traj.get('truncated')), - 'turns': traj.get('turns'), - 'success': success, - } - if traj.get('error'): - record['error'] = traj['error'] - prefix = 'ok' if success else 'fail' - step_tag = (f'step{int(global_step):06d}-' if global_step is not None else '') - fname = (f'{step_tag}{prefix}-' - f'{MultiTurnRollout._resolve_traj_id(traj, idx)}.json') - path = os.path.join(self.trace_dir, fname) - with open(path, 'w', encoding='utf-8') as f: - json.dump(record, f, ensure_ascii=False, indent=2, default=str) - except Exception: - pass + traj: Dict[str, Any], + *, + idx: int, + success: bool, + ) -> Dict[str, Any]: + """The shared record, plus the two fields only this loop produces. + + ``turns`` counts API round-trips and ``error`` carries the exception + text behind ``stop_reason='api_error'`` -- without it a trace of a + failed batch shows an empty conversation and no reason. + """ + record = super()._build_trace_record(traj, idx=idx, success=success) + record['turns'] = traj.get('turns') + if traj.get('error'): + record['error'] = traj['error'] + return record diff --git a/src/twinkle_agentic/rollout/base.py b/src/twinkle_agentic/rollout/base.py index 64d9f922b..f0b6a44a0 100644 --- a/src/twinkle_agentic/rollout/base.py +++ b/src/twinkle_agentic/rollout/base.py @@ -1,12 +1,204 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +import json +import os +import re +import time from abc import ABC, abstractmethod -from typing import List +from typing import Any, Callable, Dict, List, Optional -from twinkle.data_format import Trajectory +from twinkle.data_format import Trajectory, user_data_get +from twinkle.data_format.sampling import SamplingParams +from .bridge import _to_plain class Rollout(ABC): + """A batch of trajectories in, the same batch with the model's turns appended. + + Implementations differ in where the turns come from -- a local sampler, + whose token ids are spliced into the trajectory, or an HTTP endpoint, which + only ever returns text -- and the difference is real enough that they stay + separate classes: only one of them produces something trainable. + + Everything that is *not* generation is here: option validation, spreading a + per-call argument over the batch, and the trace dump. It moved up because + the two implementations had drifted into sharing it by reaching across the + class boundary for each other's underscore methods. + """ + + # Set by _init_common. Declared at class level so a subclass that does its + # own setup still answers these attributes instead of raising from a base + # method it inherited. + max_turns: int = 1 + sampling_params: Optional[SamplingParams] = None + trace_dir: Optional[str] = None + trace_callback: Optional[Callable[[Dict[str, Any]], bool]] = None + success_callback: Optional[Callable[[Dict[str, Any]], bool]] = None @abstractmethod def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory]: raise NotImplementedError() + + # ------------------------------------------------------------------ setup + + def _init_common( + self, + *, + max_turns: int, + sampling_params: Optional[SamplingParams] = None, + trace_dir: Optional[str] = None, + trace_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, + success_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, + ) -> None: + """Validate and store the options every multi-turn rollout takes.""" + if max_turns < 1: + raise ValueError(f'max_turns must be >= 1, got {max_turns}') + sp = sampling_params or SamplingParams() + if sp.num_samples != 1: + # n>1 would fork the conversation at turn 1 and there is no defined + # way to continue a forked trajectory: ask for several rollouts by + # passing the trajectory several times instead. + raise ValueError(f'{type(self).__name__} supports num_samples=1 only, ' + f'got {sp.num_samples}') + self.max_turns = max_turns + self.sampling_params = sp + self.trace_dir = trace_dir + self.trace_callback = trace_callback + self.success_callback = success_callback + if trace_dir: + os.makedirs(trace_dir, exist_ok=True) + + @staticmethod + def _broadcast(arg, n: int, *, name: str, required: bool = False) -> List[Any]: + """One value shared by the batch, or a list already aligned 1:1 with it. + + A list of the wrong length is refused rather than zipped short: the + mismatch would silently pair trajectories with the wrong tool manager, + which reads downstream as a model that used the wrong sandbox. + """ + if arg is None: + if required: + raise ValueError(f'{name} is required but was not provided. ' + 'Pass it at construction time or as a per-call kwarg.') + return [None] * n + if isinstance(arg, list): + if len(arg) != n: + raise ValueError(f'per-call {name} list length ({len(arg)}) does ' + f'not match number of trajectories ({n})') + return list(arg) + return [arg] * n + + # ------------------------------------------------------------------ trace + + _TRACE_SKIP_KEYS = ( + 'input_ids', + 'labels', + 'attention_mask', + 'position_ids', + 'logprobs', + 'pixel_values', + 'image_grid_thw', + 'mm_token_type_ids', + ) + + @classmethod + def _serialize_for_trace(cls, traj: Dict[str, Any]) -> Dict[str, Any]: + """Drop tensor-like / oversized fields; keep messages + metadata. + + Trace files are for human forensics; raw token ids, labels and + image buffers would bloat the file by orders of magnitude without + adding diagnostic value (the chat-template rendering of + ``messages`` already captures the textual content). + """ + slim = {k: v for k, v in traj.items() if k not in cls._TRACE_SKIP_KEYS} + return _to_plain(slim) + + @staticmethod + def _extract_ground_truth(traj: Dict[str, Any]) -> str: + """Pull ``ground_truth`` out of packed ``user_data``.""" + return user_data_get(traj.get('user_data'), 'ground_truth', '') or '' + + @staticmethod + def _resolve_traj_id(traj: Dict[str, Any], fallback_idx: int) -> str: + """Stable-ish trajectory id for filenames. + + Prefers an explicit ``id`` / ``prompt_id`` key in ``user_data`` + (sanitised for filesystem safety); else falls back to + ``{timestamp_ms}-{fallback_idx}`` so concurrent rollouts do not + overwrite each other's files. + """ + for key in ('id', 'prompt_id'): + val = user_data_get(traj.get('user_data'), key) + if val not in (None, ''): + safe = re.sub(r'[^A-Za-z0-9_\-.]+', '_', str(val))[:64] + if safe: + return safe + return f'{int(time.time() * 1000)}-{fallback_idx}' + + def _build_trace_record( + self, + traj: Dict[str, Any], + *, + idx: int, + success: bool, + ) -> Dict[str, Any]: + """Assemble one trace record. Subclasses override to add fields. + + ``idx`` is the trajectory's position in the rollout output list, + so subclasses can correlate the record with any per-call state + they stashed on ``self`` during ``__call__``. + """ + return { + 'trajectory': self._serialize_for_trace(traj), + 'ground_truth': self._extract_ground_truth(traj), + 'stop_reason': traj.get('stop_reason'), + 'truncated': bool(traj.get('truncated')), + 'success': success, + } + + def _write_rollout_traces( + self, + outs: List[Dict[str, Any]], + *, + global_step: Optional[int] = None, + ) -> None: + """Dump one pretty-printed JSON file per selected trajectory. + + ``trace_callback`` (if set) decides WHETHER to store; + ``success_callback`` (if set) decides the filename prefix + (``ok-`` vs ``fail-``). Defaults: store-all / mark-fail. + + Observability must never break training -- any I/O or encoding + problem on a single trajectory is swallowed so the remaining + dumps and the optimisation loop continue unaffected. + """ + if not self.trace_dir: + return + for idx, traj in enumerate(outs): + try: + should_store = True + if self.trace_callback is not None: + try: + should_store = bool(self.trace_callback(traj)) + except Exception: + should_store = False + if not should_store: + continue + + success = False + if self.success_callback is not None: + try: + success = bool(self.success_callback(traj)) + except Exception: + success = False + + record = self._build_trace_record(traj, idx=idx, success=success) + prefix = 'ok' if success else 'fail' + # global_step prefix lets file listings sort by training step. + step_tag = f'step{int(global_step):06d}-' if global_step is not None else '' + fname = f'{step_tag}{prefix}-{self._resolve_traj_id(traj, idx)}.json' + path = os.path.join(self.trace_dir, fname) + with open(path, 'w', encoding='utf-8') as f: + json.dump(record, f, ensure_ascii=False, indent=2, default=str) + except Exception: + # Per-trajectory failure never aborts the loop. + pass diff --git a/src/twinkle_agentic/rollout/factory.py b/src/twinkle_agentic/rollout/factory.py new file mode 100644 index 000000000..eb6494d42 --- /dev/null +++ b/src/twinkle_agentic/rollout/factory.py @@ -0,0 +1,74 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""One call that turns a generation backend into a :class:`Rollout`. + +Callers that only want turns appended to trajectories -- a challenger inventing +tasks, an evaluation script -- should not have to know that a local sampler and +an HTTP endpoint are driven by different classes with different required +arguments. They ask for a rollout, hand over whichever backend they happen to +have, and get something with the same contract: + + List[Trajectory] -> List[Trajectory] + +What the two still differ in is what ends up *inside* the trajectory, and no +factory can paper over it: the sampler path keeps ``input_ids`` / ``labels`` / +``logprobs`` and is the only one whose output can be trained on, while the API +path returns messages only. Pick the backend accordingly. +""" +from typing import Any, Dict, Optional + +from twinkle.data_format.sampling import SamplingParams +from .base import Rollout + +__all__ = ['build_rollout'] + + +def build_rollout( + backend: Any, + *, + template: Any = None, + tool_manager: Any = None, + sampling_params: Optional[SamplingParams] = None, + max_turns: int = 6, + trace_dir: Optional[str] = None, + **backend_kwargs: Any, +) -> Rollout: + """Build the multi-turn rollout that matches ``backend``. + + Args: + backend: an :class:`twinkle_agentic.protocol.base.API` (any + OpenAI-compatible endpoint) or a sampler exposing ``sample()``. + template: required for a sampler, rejected for an API. The sampler path + continues a conversation by splicing token ids, which needs the + local chat template; the API path re-sends messages as text. + tool_manager: optional for both. Without one the model is told there + are no tools. + backend_kwargs: passed straight to the chosen class -- e.g. ``harness`` + and ``max_trajectory_tokens`` for a sampler, ``concurrency`` and + ``extra_body`` for an API. An argument meant for the other backend + surfaces as a TypeError naming it. + """ + from twinkle_agentic.protocol.base import API + + common: Dict[str, Any] = { + 'tool_manager': tool_manager, + 'sampling_params': sampling_params, + 'max_turns': max_turns, + 'trace_dir': trace_dir, + } + + if isinstance(backend, API): + if template is not None: + raise ValueError('template is only used by the sampler path; an API ' + 'backend re-sends messages as text and never encodes ' + 'them locally.') + from .api_multi_turn import APIMultiTurnRollout + return APIMultiTurnRollout(api=backend, **common, **backend_kwargs) + + if not hasattr(backend, 'sample'): + raise TypeError(f'backend must be an API client or a sampler with a sample() ' + f'method, got {type(backend).__name__}') + if template is None: + raise ValueError('a sampler backend needs a template: the rollout appends each ' + 'turn as token ids and cannot re-encode the history.') + from .multi_turn import MultiTurnRollout + return MultiTurnRollout(sampler=backend, template=template, **common, **backend_kwargs) diff --git a/src/twinkle_agentic/rollout/multi_turn.py b/src/twinkle_agentic/rollout/multi_turn.py index e0ab9e36e..e66f0b222 100644 --- a/src/twinkle_agentic/rollout/multi_turn.py +++ b/src/twinkle_agentic/rollout/multi_turn.py @@ -1,13 +1,9 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -import json -import os -import re -import time from collections import defaultdict from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Any, Callable, Dict, List, Optional, Tuple -from twinkle.data_format import Trajectory, user_data_get +from twinkle.data_format import Trajectory from twinkle.data_format.sampling import SampleResponse, SamplingParams from twinkle.infra import remote_class, remote_function from twinkle.template.base import Template @@ -95,27 +91,20 @@ def __init__( super().__init__() if template is None: raise ValueError('MultiTurnRollout requires a local Template instance') - if max_turns < 1: - raise ValueError(f'max_turns must be >= 1, got {max_turns}') if max_trajectory_tokens is not None and max_trajectory_tokens < 1: raise ValueError(f'max_trajectory_tokens must be >= 1 or None, got ' f'{max_trajectory_tokens}') + self._init_common( + max_turns=max_turns, + sampling_params=sampling_params, + trace_dir=trace_dir, + trace_callback=trace_callback, + success_callback=success_callback) self.sampler = sampler self.template = template self.tool_manager = tool_manager self.harness = harness - self.sampling_params = sampling_params or SamplingParams() - self.max_turns = max_turns self.max_trajectory_tokens = max_trajectory_tokens - self.trace_dir = trace_dir - self.trace_callback = trace_callback - self.success_callback = success_callback - if self.trace_dir: - os.makedirs(self.trace_dir, exist_ok=True) - - if self.sampling_params.num_samples != 1: - raise ValueError(f'MultiTurnRollout currently supports num_samples=1 only, ' - f'got {self.sampling_params.num_samples}') assert self.template.truncation_strategy != 'split', ( "MultiTurnRollout does not support truncation_strategy='split'; " 'use left/right/delete/raise on the template.') @@ -131,8 +120,9 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] return [] sampling_params = kwargs.get('sampling_params', self.sampling_params) - tool_managers = self._resolve_tool_managers(kwargs.get('tool_manager', self.tool_manager), n) - harnesses = self._resolve_harnesses(kwargs.get('harness', self.harness), n) + tool_managers = self._broadcast( + kwargs.get('tool_manager', self.tool_manager), n, name='tool_manager', required=True) + harnesses = self._broadcast(kwargs.get('harness', self.harness), n, name='harness') lives: List[Optional[Trajectory]] = [ dict(trajectories[i]) if harnesses[i] is not None else None for i in range(n) ] @@ -311,30 +301,6 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] # ------------------------------------------------------------------ private - @staticmethod - def _resolve_tool_managers(arg, n: int) -> List[ToolManager]: - """Broadcast a single ``ToolManager`` or validate a per-trajectory list.""" - if arg is None: - raise ValueError('tool_manager is required but was not provided. ' - 'Pass it at construction time or as a per-call kwarg.') - if isinstance(arg, list): - if len(arg) != n: - raise ValueError(f'per-call tool_manager list length ({len(arg)}) does ' - f'not match number of trajectories ({n})') - return list(arg) - return [arg] * n - - @staticmethod - def _resolve_harnesses(arg, n: int) -> List[Optional[AgentHarness]]: - if arg is None: - return [None] * n - if isinstance(arg, list): - if len(arg) != n: - raise ValueError(f'per-call harness list length ({len(arg)}) does ' - f'not match number of trajectories ({n})') - return list(arg) - return [arg] * n - def _harness_before_generate( self, pif: Dict[str, Any], @@ -444,120 +410,6 @@ def _tool_messages_after( return fallback, live return delta, live - _TRACE_SKIP_KEYS = ( - 'input_ids', - 'labels', - 'attention_mask', - 'position_ids', - 'logprobs', - 'pixel_values', - 'image_grid_thw', - 'mm_token_type_ids', - ) - - @classmethod - def _serialize_for_trace(cls, traj: Dict[str, Any]) -> Dict[str, Any]: - """Drop tensor-like / oversized fields; keep messages + metadata. - - Trace files are for human forensics; raw token ids, labels and - image buffers would bloat the file by orders of magnitude without - adding diagnostic value (the chat-template rendering of - ``messages`` already captures the textual content). - """ - slim = {k: v for k, v in traj.items() if k not in cls._TRACE_SKIP_KEYS} - return _to_plain(slim) - - @staticmethod - def _extract_ground_truth(traj: Dict[str, Any]) -> str: - """Pull ``ground_truth`` out of packed ``user_data``.""" - return user_data_get(traj.get('user_data'), 'ground_truth', '') or '' - - @staticmethod - def _resolve_traj_id(traj: Dict[str, Any], fallback_idx: int) -> str: - """Stable-ish trajectory id for filenames. - - Prefers an explicit ``id`` / ``prompt_id`` key in ``user_data`` - (sanitised for filesystem safety); else falls back to - ``{timestamp_ms}-{fallback_idx}`` so concurrent rollouts do not - overwrite each other's files. - """ - for key in ('id', 'prompt_id'): - val = user_data_get(traj.get('user_data'), key) - if val not in (None, ''): - safe = re.sub(r'[^A-Za-z0-9_\-.]+', '_', str(val))[:64] - if safe: - return safe - return f'{int(time.time() * 1000)}-{fallback_idx}' - - def _build_trace_record( - self, - traj: Dict[str, Any], - *, - idx: int, - success: bool, - ) -> Dict[str, Any]: - """Assemble one trace record. Subclasses override to add fields. - - ``idx`` is the trajectory's position in the rollout output list, - so subclasses can correlate the record with any per-call state - they stashed on ``self`` during ``__call__``. - """ - return { - 'trajectory': self._serialize_for_trace(traj), - 'ground_truth': self._extract_ground_truth(traj), - 'stop_reason': traj.get('stop_reason'), - 'truncated': bool(traj.get('truncated')), - 'success': success, - } - - def _write_rollout_traces( - self, - outs: List[Dict[str, Any]], - *, - global_step: Optional[int] = None, - ) -> None: - """Dump one pretty-printed JSON file per selected trajectory. - - ``trace_callback`` (if set) decides WHETHER to store; - ``success_callback`` (if set) decides the filename prefix - (``ok-`` vs ``fail-``). Defaults: store-all / mark-fail. - - Observability must never break training -- any I/O or encoding - problem on a single trajectory is swallowed so the remaining - dumps and the optimisation loop continue unaffected. - """ - if not self.trace_dir: - return - for idx, traj in enumerate(outs): - try: - should_store = True - if self.trace_callback is not None: - try: - should_store = bool(self.trace_callback(traj)) - except Exception: - should_store = False - if not should_store: - continue - - success = False - if self.success_callback is not None: - try: - success = bool(self.success_callback(traj)) - except Exception: - success = False - - record = self._build_trace_record(traj, idx=idx, success=success) - prefix = 'ok' if success else 'fail' - # global_step prefix lets file listings sort by training step. - step_tag = f'step{int(global_step):06d}-' if global_step is not None else '' - fname = f'{step_tag}{prefix}-{self._resolve_traj_id(traj, idx)}.json' - path = os.path.join(self.trace_dir, fname) - with open(path, 'w', encoding='utf-8') as f: - json.dump(record, f, ensure_ascii=False, indent=2, default=str) - except Exception: - # Per-trajectory failure never aborts the loop. - pass - @staticmethod def _unwrap_response_list(resps, expected: int) -> List[SampleResponse]: """Validate that the sampler returned ``expected`` ``SampleResponse``s, diff --git a/src/twinkle_agentic/rsi/__init__.py b/src/twinkle_agentic/rsi/__init__.py deleted file mode 100644 index eb7bb499b..000000000 --- a/src/twinkle_agentic/rsi/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""RSI (recursive self-improvement) pipeline scripts. - -Stages (each a standalone entry script): - rsi_prepare.py - step 1: read a raw source, parallel-preprocess, dump a subset. - rsi_refine.py - step 2: re-analyze/strengthen trajectories into a standard flow. - rsi_rl.py - step 3: multi-LoRA RL, one training query per round. - rsi_distill.py - step 4: dump llm_backup data, SFT the auxiliary-role LoRAs. -""" diff --git a/src/twinkle_agentic/rsi/rsi_challenge.py b/src/twinkle_agentic/rsi/rsi_challenge.py deleted file mode 100644 index 6f29a15b5..000000000 --- a/src/twinkle_agentic/rsi/rsi_challenge.py +++ /dev/null @@ -1,1044 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""RSI step 0 — self-play data generation for code tasks. - -A single model plays TWO roles (self-play, same Qwen3-4B weights, per the talk's -"出题者/做题者" setup): - - * CHALLENGER — writes a self-contained Python solution, we RUN it in the same - sandbox rsi_rl uses to capture ground-truth outputs, and turn those outputs - into asserts. The problem statement is what the solver will be shown; the - executed reference solution is the GT. This is reverse construction: the - answer exists first (we ran it), the problem is written around it, so a GT - is available without any external labeling. - - * SOLVER (difficulty filter) — the SAME model then attempts each proposed - problem N times from the problem statement alone. We run its code against - the challenger's asserts and keep only problems whose pass count is strictly - between 0 and N ("half-know" band): all-pass or all-fail rounds give GRPO a - zero gradient (verified on MBPP), so they are dropped here. - -Two safety gates carried over from earlier failures: - * The challenger's OWN reference solution must pass its OWN asserts, or the - problem is dropped (a GT that cannot pass its own tests is noise — the - "standard answer that itself fails" pitfall). - * Output capture uses a sentinel marker + returncode check, never the last - stdout line, so an environment banner can never be mistaken for a result. - -Optional seed dataset (RSI_CH_SEED): a jsonl whose rows carry a `query` and, -preferably, a `code` reference solution. When a row has `code` (and keywords are -enabled), that proposal takes the TWO-STEP path: call 1 writes a HARDER solution on -top of the reference code, call 2 describes the problem that solution answers, and -the stage-1 code becomes the ground truth. Rows without `code` fall back to the older -single-call "seed as inspiration" prompt. Set RSI_CH_TWO_STEP=0 to force that older -path everywhere. Without any seed the challenger invents problems from scratch. - -Output (consumed directly by rsi_rl.py, no prepare/refine in between): - RSI_CH_OUT_FLOWS flows jsonl: {id, system, query, tools, rounds:[code round]} - RSI_CH_OUT_TESTS tests jsonl: {id, test_list, test_setup_code} (-> RSI_TESTS) - -Every knob is an env var (nothing hard-coded); see the config block below. -Run it as a Ray job just like rsi_rl.py (sampler-only, no trainer): - - RSI_CH_SEED=... RSI_CH_NUM_PROPOSE=2000 python -m twinkle_agentic.rsi.rsi_challenge -""" -import json -import os -import random -import re -import resource -import shutil -import signal -import subprocess -import sys -import tempfile -from typing import Any, Dict, List, Optional, Tuple - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.data_format import SamplingParams -from twinkle.dataset import Dataset, DatasetMeta -from twinkle.sampler import vLLMSampler - -logger = get_logger() - -# ── config (all env; nothing hard-coded) ─────────────────────────────────── -MODEL_ID = os.environ.get('RSI_CH_MODEL', 'ms://Qwen/Qwen3-4B') -TEMPLATE = os.environ.get('RSI_CH_TEMPLATE', 'Template') # base text template for Qwen3-4B (text-only) -SAMPLER_GPUS = int(os.environ.get('RSI_CH_SAMPLER_GPUS', 4)) - -SEED_PATH = os.environ.get('RSI_CH_SEED', '') # optional seed jsonl (empty = from scratch) -NUM_PROPOSE = int(os.environ.get('RSI_CH_NUM_PROPOSE', 2000)) # how many problems to attempt to create -PROPOSE_TEMP = float(os.environ.get('RSI_CH_PROPOSE_TEMP', 1.1)) # challenger temperature (higher = diverse) -PROPOSE_MAX_TOKENS = int(os.environ.get('RSI_CH_PROPOSE_MAX_TOKENS', 8192)) # raised: thinking+cross-domain is long -# Drop pathologically long problem statements (rambling / non-problems) before the solver -# stage, both for quality and so the solver input never exceeds the model context. -PROBLEM_MAX_CHARS = int(os.environ.get('RSI_CH_PROBLEM_MAX_CHARS', 4000)) -MAX_MODEL_LEN = int(os.environ.get('RSI_CH_MAX_MODEL_LEN', 16384)) - -# Topic-keyword conditioning (from-scratch only): first brainstorm a pool of diverse -# coding topics at high temperature, then seed each proposal with a random keyword so -# the challenger stops collapsing onto a few archetypes (palindromes, brackets, ...). -# Topic-keyword conditioning (from-scratch only): keep a persistent, 3-category keyword -# bank on disk (algorithm / computer / non-computer). Each proposal is seeded with a -# cross-domain TRIPLE drawn WITHOUT replacement (keywords are consumed); when a category -# runs out we ask the model for more distinct ones, and if it can't, we recycle. Keywords -# whose problems the solver fails hardest are expanded into more same-domain topics. -KEYWORDS_N = int(os.environ.get('RSI_CH_KEYWORDS_N', 128)) # per-category target (0 = disable) -KEYWORD_GEN_CALLS = int(os.environ.get('RSI_CH_KEYWORD_GEN_CALLS', 8)) # sampler calls per generation batch -KEYWORD_TEMP = float(os.environ.get('RSI_CH_KEYWORD_TEMP', 1.3)) # high temp -> diverse keywords -KEYWORD_MAX_TOKENS = int(os.environ.get('RSI_CH_KEYWORD_MAX_TOKENS', 1024)) -KEYWORD_DB = os.environ.get('RSI_CH_KEYWORD_DB', 'output/rsi/keywords.jsonl') # persistent bank -KEYWORD_REFILL_TRIES = int(os.environ.get('RSI_CH_KEYWORD_REFILL_TRIES', 2)) # refill attempts before recycle -SINGLE_KW_PROB = float(os.environ.get('RSI_CH_SINGLE_KW_PROB', 0.1)) # chance a proposal uses 1 keyword, not a triple -# With a seed pool loaded, this is the chance a proposal gets a seed problem ON TOP of -# its keywords; the rest are keywords-only. Seeds are drawn WITH replacement, so a pool -# smaller than NUM_PROPOSE is fine. 1.0 reproduces the old seed-only behaviour when -# keywords are disabled (RSI_CH_KEYWORDS_N=0). -SEED_MIX_PROB = float(os.environ.get('RSI_CH_SEED_MIX_PROB', 0.5)) -# Two-step seeded proposing (V4): a seeded proposal is produced by TWO sampler calls -- -# first write a harder solution on top of the seed's reference code, then describe the -# problem that solution answers. Requires the drawn seed to have a non-empty `code` field -# AND keywords to be enabled (both prompts take a topic block, so RSI_CH_KEYWORDS_N=0 -# silently keeps the single-call path). Keywords-only proposals and code-less seeds also -# keep the original single call. Costs one extra call per two-step proposal; the log line -# below reports how many proposals actually took it. -TWO_STEP = os.environ.get('RSI_CH_TWO_STEP', '1') == '1' - -# Combination arity: 'triple' = one keyword per category (max diversity); 'mix' = random -# 1/2/3 categories per proposal (use with the audit dump to see which combos keep best). -COMBO_ARITY = os.environ.get('RSI_CH_COMBO_ARITY', 'triple').lower() -# Optional 'w1,w2,w3' sampling weights for arity 1/2/3 in 'mix' mode (empty = uniform). -# Diagnostic finding: keep-rate falls with arity (~17%/8%/2%), so favour 1-2 for yield. -ARITY_WEIGHTS = os.environ.get('RSI_CH_ARITY_WEIGHTS', '') -_ARITY_W: Optional[List[float]] = None -if ARITY_WEIGHTS: - try: - _ARITY_W = [float(x) for x in ARITY_WEIGHTS.split(',')] - except ValueError: - _ARITY_W = None -AUDIT_PATH = os.environ.get('RSI_CH_AUDIT', 'output/rsi/challenge_audit.jsonl') # per-proposal outcome log -# Feedback: expand keywords whose problems the solver passed <= this many times (0 = all-fail). -LOW_PASS_EXPAND = int(os.environ.get('RSI_CH_LOW_PASS_EXPAND', 0)) -EXPAND_PER_KW = int(os.environ.get('RSI_CH_EXPAND_PER_KW', 8)) # new topics per hard keyword -EXPAND_MAX_KWS = int(os.environ.get('RSI_CH_EXPAND_MAX_KWS', 32)) # cap hard keywords expanded per run - -SOLVER_ROLLOUTS = int(os.environ.get('RSI_CH_SOLVER_ROLLOUTS', 8)) # N attempts per problem for difficulty -SOLVER_TEMP = float(os.environ.get('RSI_CH_SOLVER_TEMP', 1.0)) -SOLVER_MAX_TOKENS = int(os.environ.get('RSI_CH_SOLVER_MAX_TOKENS', 2048)) - -KEEP_MIN_PASS = int(os.environ.get('RSI_CH_KEEP_MIN_PASS', 1)) # keep if pass in [MIN, N-KEEP_MAX_MARGIN] -# "drop all-pass / all-fail" == keep 0 < pass < N. Both bounds configurable. -KEEP_MAX_PASS_MARGIN = int(os.environ.get('RSI_CH_KEEP_MAX_MARGIN', 1)) # drop pass >= N - margin + 1 - -SANDBOX_TIMEOUT = int(os.environ.get('RSI_CH_SANDBOX_TIMEOUT', 30)) -MAX_CHECKS = int(os.environ.get('RSI_CH_MAX_CHECKS', 6)) # asserts per problem cap -# Drop problems where every assert expects the SAME value: `return <that constant>` -# scores a perfect reward without reading the input, so the problem teaches nothing and -# actively rewards ignoring the task. Measured at 6.2% of kept problems on sp4_iter1. -DROP_CONSTANT_ANSWER = os.environ.get('RSI_CH_DROP_CONSTANT_ANSWER', '1') == '1' -SORT_BY_DIFFICULTY = os.environ.get('RSI_CH_SORT_BY_DIFFICULTY', '1') == '1' -# Cap how many kept problems to persist (0 = keep all). When set and exceeded, -# subsample EVENLY across the difficulty-sorted list so the stored set spans the -# whole difficulty range, not just the easiest end. -KEEP_TARGET = int(os.environ.get('RSI_CH_KEEP_TARGET', 0)) -CH_SEED = int(os.environ.get('RSI_CH_RANDOM_SEED', 0)) - -OUT_FLOWS = os.environ.get('RSI_CH_OUT_FLOWS', 'output/rsi/challenge_flows.jsonl') -OUT_TESTS = os.environ.get('RSI_CH_OUT_TESTS', 'output/rsi/challenge_tests.jsonl') -DUMP_REJECTED = os.environ.get('RSI_CH_DUMP_REJECTED', 'output/rsi/challenge_rejected.jsonl') - -CODE_SYSTEM = {'role': 'system', 'content': 'You are an expert Python programmer.'} - -# ── challenger prompt (shown to the user for review; a brand-new prompt) ──── -_MARK = '__RSI_GT__' # sentinel isolating captured output from any banner/log - -_CHALLENGER_SYS = ( - 'You design self-contained Python coding problems for training another model.\n' - 'A good problem: (1) is solvable from its statement ALONE with no external files, ' - 'network, images, or hidden context; (2) has ONE clear entry function; (3) is ' - 'deterministic (same input -> same output), no randomness, no wall-clock, no threads; ' - '(4) is neither trivial nor impossible for a mid-size model.\n' - 'You will also write the reference solution. We will EXECUTE it to obtain the ' - 'ground-truth outputs, so your solution must be correct and runnable as-is.\n' - 'Return ONLY one JSON object, no prose around it, with keys:\n' - ' "problem": the statement shown to the solver (describe the function name, its ' - 'inputs and expected behavior; do NOT include the solution).\n' - ' "solution": the reference implementation as plain Python source (no markdown fence).\n' - ' "entry": the entry function name.\n' - ' "checks": a list of 3-6 Python expressions calling the entry function on concrete ' - 'inputs (e.g. "solve([1,2,3])"); each must be evaluable after running the solution. ' - 'Do NOT write the expected value — we compute it by running your solution.' -) - -_CHALLENGER_FROM_SCRATCH = ( - 'Create ONE new Python coding problem now. Vary the topic freely ' - '(strings, arrays, math, greedy, DP, parsing, simulation ...).' -) - -_CHALLENGER_FROM_SEED = ( - 'Here is a seed problem. Create ONE NEW problem that is a meaningful VARIANT of it ' - '(change the twist, constraints, or data shape — not just renaming), keeping it ' - 'self-contained and deterministic.\n\n[seed]\n{seed}' -) - -# Keyword bank generation + keyword-conditioned proposing (diversity). -CATEGORIES = ('algorithm', 'computer', 'noncs') -_CATEGORY_DESC = { - 'algorithm': 'algorithmic techniques and paradigms (e.g. dynamic programming, binary ' - 'search, union-find, Dijkstra, backtracking, segment trees, greedy, ' - 'divide and conquer, sliding window ...)', - 'computer': 'computer-science / computing concepts that are NOT algorithms per se ' - '(e.g. hash maps, tries, LRU cache, bitsets, regular expressions, base ' - 'conversion, finite state machines, serialization, parsing, memoization ...)', - 'noncs': 'real-world domains OUTSIDE computer science, used to give a problem flavor ' - '(e.g. biology, finance, chemistry, logistics, music, cooking, sports, ' - 'astronomy, geography, linguistics ...)', -} -_KEYWORD_SYS = ( - 'You brainstorm diverse topics for a Python coding-problem generator.' -) -_KEYWORD_CAT_USER = ( - 'List {k} DISTINCT and SPECIFIC topics from this category: {desc}\n' - 'Be creative and concrete; avoid vague umbrella words. ' - 'Return ONLY a JSON array of short strings, nothing else.' -) -_KEYWORD_EXPAND_USER = ( - 'The topic "{kw}" turned out to seed genuinely HARD problems. List {m} MORE distinct, ' - 'specific topics in the SAME family/domain as "{kw}" that could seed similarly ' - 'challenging Python problems. Return ONLY a JSON array of short strings, nothing else.' -) -_CHALLENGER_FROM_KEYWORDS = ( - 'Create ONE new Python coding problem now. Draw inspiration from the following ' - 'topic(s) and combine them creatively into a single coherent problem:\n{keywords}\n' - 'You may use each topic directly or bend it loosely; combine with any data shape ' - '(strings, arrays, grids, trees, numbers, parsing, simulation ...). Make it require ' - 'real thought, not a one-liner, and keep it self-contained and deterministic.' -) -# Seed AND keywords together. The seed is deliberately framed as inspiration only, -# not as something to produce a variant of: the point is to pull the generated -# problems toward the shape of public benchmark items (short statement, one plain -# task) while the keywords keep supplying topical variety. -_CHALLENGER_FROM_SEED_KEYWORDS = ( - 'Create ONE new Python coding problem now. Use the problem below only as a ' - 'STARTING POINT for inspiration — you do NOT have to keep its task, and the new ' - 'problem does NOT need to be a variant of it.\n\n[inspiration]\n{seed}\n\n' - 'Also draw on the following topic(s), combining them into a single coherent ' - 'problem:\n{keywords}\n' - 'Make it require real thought, not a one-liner, and keep it self-contained and ' - 'deterministic.' -) - -# ── two-step (V4) challenger: build the problem FROM a harder solution ─────── -# Winning variant of the prompt bake-off (output/rsi/prompt_exp_mbpp_upgrade.py). -# The difficulty comes from adding a layer on top of a real, runnable reference -# solution, not from imagining a hard problem outright; splitting into two calls (write -# the harder code, THEN describe it) keeps the statement and the ground-truth solution -# consistent, which a single call does not. Measured on 40 MBPP seeds vs the single-call -# seed+keywords prompt: kept-rate 25% vs 15%, const-answer 4 vs 7, seed similarity 0.42. -# -# Stage 1: given the seed problem + its reference solution + topics, write a harder -# solution. Uses the plain code system prompt (CODE_SYSTEM), not _CHALLENGER_SYS, and -# returns raw code (extract_code parses it) rather than JSON. -_TWO_STEP_SOL = ( - 'Below is a coding problem and its reference solution.\n\n' - '[problem]\n{seed}\n\n[reference solution]\n{code}\n\n' - 'Write a MORE COMPLEX Python function that keeps the idea of the reference solution ' - 'as one step and builds a harder computation around it (extra pass, different data ' - 'structure, an added rule), in the direction of these topic(s):\n{keywords}\n' - 'Requirements: deterministic, self-contained, no randomness, no I/O, one clear entry ' - 'function. Output ONLY the code in a single ```python block, no explanation.' -) -# Stage 2: given the harder function PLUS the seed it grew from and the topics, describe -# the problem it answers. Seeing the seed pulls the wording back toward the MBPP task -# family (similarity 0.32 -> 0.42 when the seed is shown). The solution is NOT taken from -# this JSON -- we overwrite it with stage 1's code so the GT matches what was produced. -_TWO_STEP_PROB = ( - 'Here is a Python function.\n\n```python\n{code}\n```\n\n' - 'It was written as a harder follow-up to this exercise:\n\n[original exercise]\n' - '{seed}\n\nand it was pushed in the direction of these topic(s):\n{keywords}\n\n' - 'Write the problem statement that the function above is the answer to, as if it were ' - 'a coding exercise in the same series as the original: name the entry function, ' - 'describe its inputs and the exact behaviour expected, and do NOT reveal the ' - 'implementation. Phrase it as plainly and briefly as the original exercise.\n' - 'Return ONLY one JSON object, no prose around it, with keys:\n' - ' "problem": the statement shown to the solver.\n' - ' "entry": the entry function name.\n' - ' "checks": a list of 3-6 Python expressions calling the entry function on ' - 'concrete inputs; each must be evaluable after running the function above. Do NOT ' - 'write the expected value.\n' - 'The "solution" is already known, so do not include it.' -) - - - -_SOLVER_SYS = {'role': 'system', 'content': 'You are an expert Python programmer.'} -_SOLVER_USER = ( - '{problem}\n\n' - 'Write the complete Python solution. Put the final code in a single ```python fenced ' - 'block. Define the exact function name required by the problem.' -) - - -# ── sandbox (mirrors rsi_rl.run_asserts / extract_code; duplicated on purpose: -# importing rsi_rl would run its module-level CLI.from_args() + swanlab.init) ─ -_FENCE_RE = re.compile(r'```(?:python|py)?\s*\n(.*?)```', re.S) - - -def extract_code(text: str) -> str: - idx = (text or '').rfind('</think>') - body = text[idx + len('</think>'):] if idx >= 0 else (text or '') - blocks = _FENCE_RE.findall(body) - return (blocks[-1] if blocks else body).strip() - - -def _run_script(script: str, timeout: int) -> Tuple[int, str]: - """Run a python script in an isolated dir, 2GB cap, killpg on timeout. - - Returns (returncode, stdout). returncode is -1 on timeout/spawn failure. - """ - tmp = tempfile.mkdtemp(prefix='rsi_ch_') - try: - with open(os.path.join(tmp, '_run.py'), 'w', encoding='utf-8') as f: - f.write(script + '\n') - env = dict(os.environ, MPLBACKEND='Agg', PYTHONHASHSEED='0', OMP_NUM_THREADS='1', - MKL_NUM_THREADS='1', TOKENIZERS_PARALLELISM='false') - env.pop('CUDA_VISIBLE_DEVICES', None) - - def _limit(): - resource.setrlimit(resource.RLIMIT_AS, (2 * 1024**3, 2 * 1024**3)) - - proc = subprocess.Popen([sys.executable, '_run.py'], cwd=tmp, env=env, - stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, - text=True, start_new_session=True, preexec_fn=_limit) - try: - out, _ = proc.communicate(timeout=timeout) - return proc.returncode, out or '' - except subprocess.TimeoutExpired: - try: - os.killpg(proc.pid, signal.SIGKILL) - except ProcessLookupError: - pass - try: - proc.communicate(timeout=5) - except Exception: - pass - return -1, '' - finally: - shutil.rmtree(tmp, ignore_errors=True) - - -def run_asserts(code: str, setup: str, asserts: List[str], timeout: int = SANDBOX_TIMEOUT) -> bool: - """True when every assert passes (returncode 0). Same contract as rsi_rl.""" - if not code.strip() or not asserts: - return False - parts = [code] - if (setup or '').strip(): - parts.append(setup) - parts.extend(asserts) - rc, _ = _run_script('\n\n'.join(parts), timeout) - return rc == 0 - - -def build_asserts(solution: str, checks: List[str], timeout: int = SANDBOX_TIMEOUT) -> Optional[List[str]]: - """Run the reference solution once to capture repr of each check expression, - then form ``assert <check> == <captured>``. Sentinel-marked + returncode - checked so a crash or a banner line can never be read as a value. - - Returns the assert list, or None if the solution crashed / produced no usable - output (that problem is then dropped upstream). - """ - checks = [c for c in checks if isinstance(c, str) and c.strip()][:MAX_CHECKS] - if not checks: - return None - lines = [solution, ''] - for i, c in enumerate(checks): - # repr on its own line, tagged with index; a check that raises makes the - # whole script exit non-zero -> we drop the problem. Pure f-string (no %% - # formatting) so a check expression containing '%' (modulo/percent) is safe. - lines.append(f'print("{_MARK}{i}=" + repr({c}))') - rc, out = _run_script('\n'.join(lines), timeout) - if rc != 0: - return None - captured: Dict[int, str] = {} - for line in out.splitlines(): - if line.startswith(_MARK): - try: - idx_str, val = line[len(_MARK):].split('=', 1) - captured[int(idx_str)] = val - except (ValueError, IndexError): - continue - if len(captured) != len(checks): - return None - asserts = [] - for i, c in enumerate(checks): - val = captured[i] - # The captured text is a repr, so it is a valid literal to compare against. - asserts.append(f'assert ({c}) == ({val})') - return asserts - - -def _split_top_eq(s: str) -> Optional[tuple]: - """Split on the first top-level ``==``, ignoring anything inside brackets or quotes.""" - depth = 0 - quote = '' - i = 0 - while i < len(s) - 1: - c = s[i] - if quote: - if c == quote: - quote = '' - elif c in '\'"': - quote = c - elif c in '([{': - depth += 1 - elif c in ')]}': - depth -= 1 - elif depth == 0 and c == '=' and s[i + 1] == '=': - return s[:i].strip(), s[i + 2:].strip() - i += 1 - return None - - -def _expected_of(assert_line: str) -> Optional[str]: - """The value the solver actually has to produce for one assert. - - build_asserts emits ``assert (<check>) == (<repr>)``, but a check may itself be a - comparison, giving ``assert (f(x) == 3) == (True)``. Reading the outer side there - would report 'True' and make such a problem look constant-answer, so the inner - right-hand side is used instead. An outer ``False`` pins nothing down at all and - is reported as unknown. - """ - m = re.match(r'^\s*assert\s*\((.*)\)\s*==\s*\((.*)\)\s*$', assert_line.strip()) - if not m: - return None - lhs, rhs = m.group(1).strip(), m.group(2).strip() - inner = _split_top_eq(lhs) - if rhs in ('True', 'False') and inner is not None: - return inner[1] if rhs == 'True' else None - return rhs - - -def is_constant_answer(asserts: List[str]) -> bool: - """Would ``return <one constant>`` satisfy every assert? - - Requires at least two asserts with a readable expectation: a single assert is - trivially 'constant' and one unreadable assert must not hide a constant set. - """ - vals = [_expected_of(a) for a in asserts] - if any(v is None for v in vals) or len(vals) < 2: - return False - return len(set(vals)) == 1 - - -# ── challenger output parsing ────────────────────────────────────────────── -_JSON_FENCE_RE = re.compile(r'^\s*```(?:json)?\s*|\s*```\s*$', re.I) - - -def parse_challenger(text: str, require_solution: bool = True) -> Optional[Dict[str, Any]]: - """Pull the JSON object out of the challenger's completion. - - ``require_solution=False`` is for the two-step (V4) flow, where the harder solution - comes from a separate call and stage 2 returns only ``problem``/``entry``/``checks``. - """ - body = text - idx = body.rfind('</think>') - if idx >= 0: - body = body[idx + len('</think>'):] - body = _JSON_FENCE_RE.sub('', body.strip()).strip() - # Grab the outermost {...} if there is leading/trailing prose. - start = body.find('{') - end = body.rfind('}') - if start < 0 or end <= start: - return None - try: - obj = json.loads(body[start:end + 1]) - except (ValueError, TypeError): - return None - if not isinstance(obj, dict): - return None - problem = obj.get('problem') - solution = obj.get('solution') - checks = obj.get('checks') - if not (isinstance(problem, str) and problem.strip() - and isinstance(checks, list) and checks): - return None - if require_solution: - if not (isinstance(solution, str) and solution.strip()): - return None - else: - # Stage 2 is told NOT to include a solution; if it did anyway, ignore it -- we - # will overwrite with stage 1's code so the GT matches what actually ran. - solution = solution if isinstance(solution, str) else '' - # solution may still arrive fenced despite instructions. - if solution and '```' in solution: - solution = extract_code(solution) - return {'problem': problem.strip(), 'solution': (solution or '').strip(), - 'entry': str(obj.get('entry') or '').strip(), 'checks': checks} - - -# ── sampling helpers ─────────────────────────────────────────────────────── -def _completion_text(seq) -> str: - """The assistant text of one sampled sequence (decode is fine: not training).""" - if seq.decoded: - return seq.decoded - feat = seq.new_input_feature or {} - for m in reversed(feat.get('messages', []) or []): - if m.get('role') == 'assistant': - return m.get('content', '') or '' - return '' - - -def sample_texts(sampler, message_lists: List[List[Dict[str, Any]]], - sampling_params: SamplingParams) -> List[str]: - """Sample one completion per message list; return the assistant texts.""" - if not message_lists: - return [] - trajs = [{'messages': msgs} for msgs in message_lists] - responses = sampler.sample(trajs, sampling_params) - texts: List[str] = [] - for resp in responses: - seqs = resp.sequences - texts.append(_completion_text(seqs[0]) if seqs else '') - return texts - - -def _parse_keyword_list(text: str) -> List[str]: - """Extract a JSON array of short strings from a (possibly thinking) reply.""" - body = text - idx = body.rfind('</think>') - if idx >= 0: - body = body[idx + len('</think>'):] - start, end = body.find('['), body.rfind(']') - if start < 0 or end <= start: - return [] - try: - arr = json.loads(body[start:end + 1]) - except (ValueError, TypeError): - return [] - out: List[str] = [] - for x in arr: - if isinstance(x, str) and x.strip() and len(x.strip()) <= 60: - out.append(x.strip()) - return out - - -class KeywordStore: - """Persistent 3-category keyword bank with usage tracking (see CATEGORIES). - - On-disk format (KEYWORD_DB, one JSON per line): - {"category", "text", "used": bool, "used_count": int, "source": "gen"|"expand", - "parent": <keyword or null>} - De-duplicates case-insensitively within each category so re-runs never conflict. - """ - - def __init__(self, path: str): - self.path = path - self.items: Dict[str, List[Dict[str, Any]]] = {c: [] for c in CATEGORIES} - self._seen: Dict[str, set] = {c: set() for c in CATEGORIES} - if path and os.path.exists(path): - with open(path, encoding='utf-8') as f: - for line in f: - line = line.strip() - if not line: - continue - try: - r = json.loads(line) - except (ValueError, TypeError): - continue - c, t = r.get('category'), r.get('text') - if c in self.items and isinstance(t, str) and t.strip(): - key = t.strip().lower() - if key not in self._seen[c]: - self._seen[c].add(key) - self.items[c].append(r) - - def save(self) -> None: - os.makedirs(os.path.dirname(os.path.abspath(self.path)) or '.', exist_ok=True) - tmp = self.path + '.tmp' - with open(tmp, 'w', encoding='utf-8') as f: - for c in CATEGORIES: - for r in self.items[c]: - f.write(json.dumps(r, ensure_ascii=False) + '\n') - os.replace(tmp, self.path) - - def add(self, category: str, texts: List[str], source: str = 'gen', - parent: Optional[str] = None) -> int: - added = 0 - for t in texts: - key = t.strip().lower() - if not key or key in self._seen[category]: - continue - self._seen[category].add(key) - self.items[category].append({'category': category, 'text': t.strip(), - 'used': False, 'used_count': 0, - 'source': source, 'parent': parent}) - added += 1 - return added - - def unused(self, category: str) -> List[Dict[str, Any]]: - return [r for r in self.items[category] if not r.get('used')] - - def texts(self, category: str) -> List[str]: - return [r['text'] for r in self.items[category]] - - def recycle(self, category: str) -> None: - """Mark every keyword unused again (safety valve when the model is tapped out).""" - for r in self.items[category]: - r['used'] = False - - -def generate_category_keywords(sampler, category: str, n_want: int, - avoid: List[str], rng, nonce: int = 0) -> List[str]: - """Ask the model for up to ``n_want`` distinct keywords in ``category``. - - ``avoid`` (already-known texts) is both injected as a soft "don't repeat" hint and - used to filter the parsed result. ``nonce`` perturbs the prompt so refills differ. - """ - if n_want <= 0: - return [] - n_calls = max(KEYWORD_GEN_CALLS, SAMPLER_GPUS) # batch must cover every DP worker - per_call = max(1, -(-n_want // n_calls) + 4) # ceil(n/calls) + margin - avoid_note = '' - if avoid: - shown = avoid if len(avoid) <= 40 else rng.sample(avoid, 40) - avoid_note = '\nDo NOT repeat any of these already-used topics: ' + ', '.join(shown) - base = _KEYWORD_CAT_USER.format(k=per_call, desc=_CATEGORY_DESC[category]) + avoid_note - msgs = [[{'role': 'system', 'content': _KEYWORD_SYS}, - {'role': 'user', 'content': f'{base}\n(batch {nonce}-{i})'}] - for i in range(n_calls)] - sp = SamplingParams(max_tokens=KEYWORD_MAX_TOKENS, num_samples=1, logprobs=1, - temperature=KEYWORD_TEMP, top_p=0.98) - texts = sample_texts(sampler, msgs, sp) - out: List[str] = [] - seen = {a.strip().lower() for a in avoid} - for t in texts: - for kw in _parse_keyword_list(t): - key = kw.strip().lower() - if key and key not in seen: - seen.add(key) - out.append(kw.strip()) - rng.shuffle(out) - return out[:n_want] - - -def ensure_unused(store: 'KeywordStore', sampler, category: str, need: int, rng, - nonce: int) -> int: - """Make ``category`` hold >= ``need`` unused keywords, generating/recycling as needed. - - Returns the next free ``nonce`` to use for the following generation call. - """ - tries = 0 - while len(store.unused(category)) < need: - new = generate_category_keywords(sampler, category, KEYWORDS_N, - store.texts(category), rng, nonce=nonce) - added = store.add(category, new, source='gen') - nonce += 1 - tries += 1 - if added == 0 and tries >= KEYWORD_REFILL_TRIES: - # Model is out of fresh distinct topics; recycle so combinations keep flowing. - if store.items[category]: - store.recycle(category) - logger.info(f'[rsi_challenge] keyword category {category!r} exhausted -> recycled ' - f'{len(store.items[category])} topics') - break - return nonce - - -def expand_hard_keywords(store: 'KeywordStore', sampler, hard, rng, nonce: int) -> int: - """For each (category, keyword) the solver failed hardest, brainstorm same-domain - topics and add them (source='expand') to the bank. Returns count added.""" - hard = hard[:EXPAND_MAX_KWS] - if not hard or EXPAND_PER_KW <= 0: - return 0 - # Batch must cover all DP workers; cycle the hard list if it is too short. - reqs = list(hard) - while len(reqs) < SAMPLER_GPUS: - reqs.append(hard[len(reqs) % len(hard)]) - msgs = [[{'role': 'system', 'content': _KEYWORD_SYS}, - {'role': 'user', 'content': _KEYWORD_EXPAND_USER.format(kw=kw, m=EXPAND_PER_KW) - + f'\n(batch {nonce}-{i})'}] - for i, (_c, kw) in enumerate(reqs)] - sp = SamplingParams(max_tokens=KEYWORD_MAX_TOKENS, num_samples=1, logprobs=1, - temperature=KEYWORD_TEMP, top_p=0.98) - texts = sample_texts(sampler, msgs, sp) - added = 0 - for (cat, kw), t in zip(reqs, texts): - added += store.add(cat, _parse_keyword_list(t), source='expand', parent=kw) - return added - - -def _draw_keywords(store: 'KeywordStore', sampler, rng, nonce: int) -> Tuple[List[Tuple[str, str]], int]: - """Consume one keyword combination from the bank (arity per COMBO_ARITY). - - Split out of the propose loop so a proposal can carry keywords whether or not - it also carries a seed problem. - """ - if COMBO_ARITY == 'mix': - if _ARITY_W and len(_ARITY_W) == len(CATEGORIES): - k = rng.choices(range(1, len(CATEGORIES) + 1), weights=_ARITY_W)[0] - else: - k = rng.randint(1, len(CATEGORIES)) - cats = rng.sample(list(CATEGORIES), k) - elif rng.random() < SINGLE_KW_PROB: - cats = [rng.choice(CATEGORIES)] - else: - cats = list(CATEGORIES) - picks: List[Tuple[str, str]] = [] - for c in cats: - if not store.unused(c): - nonce = ensure_unused(store, sampler, c, 1, rng, nonce) - un = store.unused(c) - if not un: - continue - r = rng.choice(un) - r['used'] = True - r['used_count'] = r.get('used_count', 0) + 1 - picks.append((c, r['text'])) - return picks, nonce - - -def load_seeds(path: str) -> List[Dict[str, str]]: - """Read seed problems from a jsonl. Returns dicts with at least 'query'; may also - have 'code' (reference solution) when the file was written by split_mbpp.py v2+. - - When the file only has 'query' (legacy format), the returned dicts have code=''. - The two-step challenger requires 'code' to be non-empty; if all seeds lack code it - falls back to the original single-step prompt automatically. - """ - if not path or not os.path.exists(path): - return [] - seeds: List[Dict[str, str]] = [] - with open(path, encoding='utf-8') as f: - for line in f: - line = line.strip() - if not line: - continue - try: - row = json.loads(line) - except (ValueError, TypeError): - continue - q = row.get('query') or row.get('problem') or row.get('prompt') - if isinstance(q, dict): - q = q.get('content') - if not q: - msgs = row.get('messages') or [] - q = next((m.get('content') for m in msgs if m.get('role') == 'user'), None) - if isinstance(q, str) and q.strip(): - seeds.append({'query': q.strip(), 'code': (row.get('code') or '').strip()}) - return seeds - - -def main(): - rng = random.Random(CH_SEED) - for p in (OUT_FLOWS, OUT_TESTS, DUMP_REJECTED): - os.makedirs(os.path.dirname(os.path.abspath(p)) or '.', exist_ok=True) - - device_groups = [DeviceGroup(name='sampler', ranks=list(range(SAMPLER_GPUS)), device_type='GPU')] - sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, dp_size=SAMPLER_GPUS) - twinkle.initialize(mode='ray', nproc_per_node=SAMPLER_GPUS, groups=device_groups, lazy_collect=False) - - sampler = vLLMSampler( - model_id=MODEL_ID, - engine_args={'gpu_memory_utilization': 0.8, 'max_model_len': MAX_MODEL_LEN}, - device_mesh=sampler_mesh, - remote_group='sampler', - ) - sampler.set_template(TEMPLATE, model_id=MODEL_ID, enable_thinking=True, max_length=MAX_MODEL_LEN) - - seeds = load_seeds(SEED_PATH) - logger.info(f'[rsi_challenge] seeds loaded: {len(seeds)} from {SEED_PATH!r} ' - f'(seed_mix_prob={SEED_MIX_PROB if seeds else 0.0}, ' - f'{"seed+keywords / keywords-only mix" if seeds else "keywords-only"})') - - # ── stage 1: challenger proposes, we execute to build GT asserts ──────── - # Diversity: each proposal draws a cross-domain keyword combination from a - # persistent 3-category bank, consuming keywords without replacement. When a seed - # pool is given, SEED_MIX_PROB of the proposals additionally get one seed problem - # (drawn WITH replacement) as inspiration on top of the keywords -- the keyword - # bank is now used in BOTH modes, where it used to be skipped entirely whenever - # seeds were present. - store = KeywordStore(KEYWORD_DB) if KEYWORDS_N > 0 else None - nonce = int(CH_SEED) - if store is not None and KEYWORDS_N > 0: - for c in CATEGORIES: - nonce = ensure_unused(store, sampler, c, 1, rng, nonce) - logger.info('[rsi_challenge] keyword bank: ' - + ', '.join(f'{c}={len(store.items[c])}({len(store.unused(c))} free)' - for c in CATEGORIES)) - - # Build the FIRST-call message for every proposal. A seeded proposal that (a) drew a - # seed carrying reference code and (b) has TWO_STEP on becomes a two-step proposal: - # its first call writes a HARDER solution (raw code), and a second call -- built once - # we see that code -- turns it into a problem statement + checks. Everything else is a - # single JSON-producing call, exactly as before. - propose_msgs: List[List[Dict[str, Any]]] = [] - propose_kws: List[List[Tuple[str, str]]] = [] # (category, text) picked per proposal, for feedback - propose_seeded: List[bool] = [] # whether an MBPP seed rode along, for the audit - is_two_step: List[bool] = [] # whether this proposal uses the V4 two-call flow - seed_query: List[str] = [] # seed statement carried into stage 2 (two-step only) - kw_body: List[str] = [] # keyword block carried into stage 2 (two-step only) - n_seeded = 0 - n_two = 0 - for _ in range(NUM_PROPOSE): - picks: List[Tuple[str, str]] = [] - if store is not None: - picks, nonce = _draw_keywords(store, sampler, rng, nonce) - body = '\n'.join(f'- {c}: {t}' for c, t in picks) - use_seed = bool(seeds) and rng.random() < SEED_MIX_PROB - seed = rng.choice(seeds) if use_seed else None - two = bool(use_seed and TWO_STEP and picks and seed and seed.get('code')) - if use_seed: - n_seeded += 1 - if two: - n_two += 1 - sys_msg = dict(CODE_SYSTEM) - user = _TWO_STEP_SOL.format(seed=seed['query'], code=seed['code'], keywords=body) - elif use_seed and picks: - sys_msg = {'role': 'system', 'content': _CHALLENGER_SYS} - user = _CHALLENGER_FROM_SEED_KEYWORDS.format(seed=seed['query'], keywords=body) - elif use_seed: - sys_msg = {'role': 'system', 'content': _CHALLENGER_SYS} - user = _CHALLENGER_FROM_SEED.format(seed=seed['query']) - elif picks: - sys_msg = {'role': 'system', 'content': _CHALLENGER_SYS} - user = _CHALLENGER_FROM_KEYWORDS.format(keywords=body) - else: - sys_msg = {'role': 'system', 'content': _CHALLENGER_SYS} - user = _CHALLENGER_FROM_SCRATCH - propose_kws.append(picks) - propose_seeded.append(use_seed) - is_two_step.append(two) - seed_query.append(seed['query'] if two else '') - kw_body.append(body if two else '') - propose_msgs.append([sys_msg, {'role': 'user', 'content': user}]) - logger.info(f'[rsi_challenge] proposals={NUM_PROPOSE} seeded={n_seeded} two_step={n_two} ' - f'keywords_only={NUM_PROPOSE - n_seeded} (RSI_CH_SEED_MIX_PROB={SEED_MIX_PROB}, ' - f'RSI_CH_TWO_STEP={int(TWO_STEP)})') - propose_sp = SamplingParams(max_tokens=PROPOSE_MAX_TOKENS, num_samples=1, logprobs=1, - temperature=PROPOSE_TEMP, top_p=0.95) - logger.info(f'[rsi_challenge] proposing {NUM_PROPOSE} problems (T={PROPOSE_TEMP})') - stage1_text = sample_texts(sampler, propose_msgs, propose_sp) - - # objs[i] = the parsed proposal for index i (or None if it failed). For single-step - # proposals this is just parse_challenger(stage1). For two-step, stage 1 gave code and - # a second batched call turns each into problem+checks; we then FORCE solution = the - # stage-1 code so the ground truth matches what was actually produced. - # Driven by len(stage1_text), not NUM_PROPOSE, so a short sampler return cannot - # IndexError here (the original enumerate-based loop was naturally tolerant). - n_prop = len(stage1_text) - if n_prop != NUM_PROPOSE: - logger.warning(f'[rsi_challenge] sampler returned {n_prop} texts for ' - f'{NUM_PROPOSE} proposals; proceeding with {n_prop}') - objs: List[Optional[Dict[str, Any]]] = [None] * n_prop - two_codes: List[str] = [''] * n_prop - n_no_code = 0 - stage2_idx: List[int] = [] - stage2_msgs: List[List[Dict[str, Any]]] = [] - for i in range(n_prop): - if not is_two_step[i]: - objs[i] = parse_challenger(stage1_text[i]) - continue - code = extract_code(stage1_text[i]) - two_codes[i] = code - if not code.strip(): - # Stage 1 produced no code block (usually a truncated completion): this - # proposal dies here. Counted separately from JSON parse failures. - n_no_code += 1 - continue - stage2_idx.append(i) - stage2_msgs.append([{'role': 'system', 'content': _CHALLENGER_SYS}, - {'role': 'user', 'content': _TWO_STEP_PROB.format( - code=code, seed=seed_query[i], keywords=kw_body[i])}]) - stage2_text_by_idx: Dict[int, str] = {} - if stage2_msgs: - logger.info(f'[rsi_challenge] two-step stage 2: {len(stage2_msgs)} problem writes ' - f'({n_no_code} stage-1 completions had no code block)') - for i, txt in zip(stage2_idx, sample_texts(sampler, stage2_msgs, propose_sp)): - stage2_text_by_idx[i] = txt - obj = parse_challenger(txt, require_solution=False) - if obj is not None: - obj['solution'] = two_codes[i] # GT = the harder solution stage 1 produced - objs[i] = obj - - # For the audit, keep the completion whose JSON we parsed (stage 2 for two-step, else - # the single call) so resp_chars / think_closed describe the statement-producing call. - proposals_text = [stage2_text_by_idx.get(i, stage1_text[i]) if is_two_step[i] - else stage1_text[i] for i in range(n_prop)] - - problems: List[Dict[str, Any]] = [] - stat = {'parsed': 0, 'parse_fail': 0, 'stage1_no_code': n_no_code, 'too_long': 0, - 'gt_fail': 0, 'selfcheck_fail': 0, 'constant_answer': 0} - outcome: List[str] = ['parse_fail'] * len(proposals_text) # per-proposal audit label - for i in range(n_prop): - if is_two_step[i] and not two_codes[i].strip(): - outcome[i] = 'stage1_no_code' - rejected = open(DUMP_REJECTED, 'w', encoding='utf-8') - for pi in range(len(proposals_text)): - obj = objs[pi] - if obj is None: - if outcome[pi] != 'stage1_no_code': - stat['parse_fail'] += 1 - continue - stat['parsed'] += 1 - # Reject rambling / non-problem statements early (also keeps solver input in-context). - if len(obj['problem']) > PROBLEM_MAX_CHARS: - stat['too_long'] += 1 - outcome[pi] = 'too_long' - continue - asserts = build_asserts(obj['solution'], obj['checks']) - if not asserts: - stat['gt_fail'] += 1 - outcome[pi] = 'gt_fail' - rejected.write(json.dumps({'reason': 'gt_build_fail', **obj}, ensure_ascii=False) + '\n') - continue - # The reference solution must pass its own asserts, or the GT is noise. - if not run_asserts(obj['solution'], '', asserts): - stat['selfcheck_fail'] += 1 - outcome[pi] = 'selfcheck_fail' - rejected.write(json.dumps({'reason': 'selfcheck_fail', 'asserts': asserts, **obj}, - ensure_ascii=False) + '\n') - continue - # A problem whose every assert expects the same value rewards `return <constant>`, - # so it would train the solver to ignore the statement. Drop it before the (much - # more expensive) solver rollouts. - if DROP_CONSTANT_ANSWER and is_constant_answer(asserts): - stat['constant_answer'] += 1 - outcome[pi] = 'constant_answer' - rejected.write(json.dumps({'reason': 'constant_answer', 'asserts': asserts, **obj}, - ensure_ascii=False) + '\n') - continue - obj['asserts'] = asserts - obj['_kw'] = propose_kws[pi] if pi < len(propose_kws) else [] # origin keywords (feedback) - obj['_idx'] = pi - outcome[pi] = 'usable' - problems.append(obj) - logger.info(f'[rsi_challenge] proposal stage: {stat}, usable problems={len(problems)}') - - # ── stage 2: solver difficulty filter (keep 0 < pass < N) ─────────────── - kept: List[Dict[str, Any]] = [] - if problems: - solver_msgs: List[List[Dict[str, Any]]] = [] - for prob in problems: - for _ in range(SOLVER_ROLLOUTS): - solver_msgs.append([_SOLVER_SYS, - {'role': 'user', 'content': _SOLVER_USER.format(problem=prob['problem'])}]) - solver_sp = SamplingParams(max_tokens=SOLVER_MAX_TOKENS, num_samples=1, logprobs=1, - temperature=SOLVER_TEMP, top_p=0.95) - logger.info(f'[rsi_challenge] difficulty rollout: {len(problems)} problems x ' - f'{SOLVER_ROLLOUTS} (T={SOLVER_TEMP})') - solver_text = sample_texts(sampler, solver_msgs, solver_sp) - - hi = SOLVER_ROLLOUTS - KEEP_MAX_PASS_MARGIN - for pi, prob in enumerate(problems): - n_pass = 0 - for k in range(SOLVER_ROLLOUTS): - code = extract_code(solver_text[pi * SOLVER_ROLLOUTS + k]) - if run_asserts(code, '', prob['asserts']): - n_pass += 1 - prob['n_pass'] = n_pass - if KEEP_MIN_PASS <= n_pass <= hi: - kept.append(prob) - outcome[prob['_idx']] = f'kept(pass={n_pass})' - else: - outcome[prob['_idx']] = f'dropped(pass={n_pass})' - rejected.write(json.dumps({'reason': f'difficulty_pass={n_pass}/{SOLVER_ROLLOUTS}', - 'problem': prob['problem']}, ensure_ascii=False) + '\n') - rejected.close() - - # ── per-proposal audit: attribute outcome to the keyword combination + flag - # truncation (thinking never closed) so combo/domain effects can be measured. - if AUDIT_PATH: - with open(AUDIT_PATH, 'w', encoding='utf-8') as af: - for i, txt in enumerate(proposals_text): - kws = propose_kws[i] if i < len(propose_kws) else [] - af.write(json.dumps({ - 'idx': i, - 'arity': len(kws), - 'cats': [c for c, _ in kws], - 'kws': kws, - # Whether this proposal also carried an MBPP seed statement, so the - # keep rate of seeded vs keywords-only proposals can be compared. - 'seeded': bool(propose_seeded[i]) if i < len(propose_seeded) else False, - # Whether the V4 two-call flow was used (seed carried code + TWO_STEP). - 'two_step': bool(is_two_step[i]) if i < len(is_two_step) else False, - 'outcome': outcome[i], - 'resp_chars': len(txt), - 'think_closed': '</think>' in txt, - }, ensure_ascii=False) + '\n') - logger.info(f'[rsi_challenge] per-proposal audit -> {AUDIT_PATH}') - - # ── feedback: expand the keywords behind the hardest problems (solver pass - # <= LOW_PASS_EXPAND) into more same-domain topics, then persist the bank. - if store is not None: - hard: List[Tuple[str, str]] = [] - seen_hard = set() - for prob in problems: - if prob.get('n_pass', SOLVER_ROLLOUTS) <= LOW_PASS_EXPAND: - for c, t in prob.get('_kw', []): - if (c, t.lower()) not in seen_hard: - seen_hard.add((c, t.lower())) - hard.append((c, t)) - if hard: - rng.shuffle(hard) - added = expand_hard_keywords(store, sampler, hard, rng, nonce) - logger.info(f'[rsi_challenge] feedback: expanded {min(len(hard), EXPAND_MAX_KWS)} ' - f'hard keyword(s) -> +{added} new same-domain topics') - store.save() - logger.info('[rsi_challenge] keyword bank saved: ' - + ', '.join(f'{c}={len(store.items[c])}' for c in CATEGORIES) - + f' -> {KEYWORD_DB}') - - # File order = increasing difficulty (fewer solver passes later), which the - # fixed-pool validation in rsi_rl relies on. - if SORT_BY_DIFFICULTY: - kept.sort(key=lambda p: -p['n_pass']) - - # Optional even subsample to KEEP_TARGET across the difficulty-sorted list. - if KEEP_TARGET and len(kept) > KEEP_TARGET: - n = len(kept) - idx = sorted({round(i * (n - 1) / (KEEP_TARGET - 1)) for i in range(KEEP_TARGET)}) - kept = [kept[j] for j in idx] - logger.info(f'[rsi_challenge] capped to KEEP_TARGET={KEEP_TARGET} ' - f'(evenly across difficulty), stored={len(kept)}') - - # ── write flows + tests for rsi_rl (skip prepare/refine) ──────────────── - with open(OUT_FLOWS, 'w', encoding='utf-8') as ff, open(OUT_TESTS, 'w', encoding='utf-8') as ft: - for i, prob in enumerate(kept): - cid = f'ch_{i:06d}' - flow = { - 'id': cid, - 'system': CODE_SYSTEM, - 'query': {'role': 'user', 'content': prob['problem']}, - 'tools': [], - # Difficulty audit: how many of the SOLVER_ROLLOUTS attempts passed - # (0 < n_pass < N by construction). rsi_rl ignores these extra keys; - # kept so a persisted flow can be analyzed without re-running. - 'n_pass': prob.get('n_pass'), - 'n_rollouts': SOLVER_ROLLOUTS, - # Origin keywords (category, text) of the cross-domain triple, for audit. - 'keywords': prob.get('_kw', []), - 'rounds': [{ - 'intent': 'solve the problem', - 'type': 'code', - 'tool_call': None, - 'code': prob['solution'], # challenger's passing solution (OPSD reads this) - 'result': '', - 'reward_method': 'rubric', - }], - } - ff.write(json.dumps(flow, ensure_ascii=False) + '\n') - ft.write(json.dumps({'id': cid, 'test_list': prob['asserts'], 'test_setup_code': ''}, - ensure_ascii=False) + '\n') - - logger.info(f'[rsi_challenge] kept {len(kept)}/{len(problems)} problems -> ' - f'{OUT_FLOWS} + {OUT_TESTS}') - if kept: - dist: Dict[int, int] = {} - for p in kept: - dist[p['n_pass']] = dist.get(p['n_pass'], 0) + 1 - logger.info(f'[rsi_challenge] kept pass-count distribution (0<pass<N): ' - f'{dict(sorted(dist.items()))}') - - -if __name__ == '__main__': - main() diff --git a/src/twinkle_agentic/rsi/rsi_distill.py b/src/twinkle_agentic/rsi/rsi_distill.py deleted file mode 100644 index bdedc8c00..000000000 --- a/src/twinkle_agentic/rsi/rsi_distill.py +++ /dev/null @@ -1,285 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""RSI step 4 — turn the data collected by ``llm_backup`` into a per-role LoRA -via plain SFT. - -Where the data comes from -------------------------- -``twinkle_agentic.utils.llm_backup`` now optionally appends one raw record per -teacher call to the JSONL at ``$LLM_BACKUP_DUMP_PATH`` (off unless that env var -is set). Each line is:: - - {"key": <hash>, "trajectory": <the exact model input>, - "student": <student output str>, "teacher": <teacher output str>, - "match": <bool>} - -``trajectory`` is whatever the decorated role passed as its ``trajectory`` arg — -in twinkle_agentic that is an ``{"messages": [...], "tools": [...]}`` dict (see -protocol/openai.py). It is stored verbatim so it can be reshaped here into an -SFT pair without re-running anything. - -What this script trains ------------------------- -SFT target = EVERY teacher output (the ``match`` flag is ignored; decided by the -user). One training sample is:: - - messages = <trajectory messages> + [{"role": "assistant", "content": teacher}] - -and only that final assistant turn is trainable (``key_rounds=[len(msgs)-1]``), -so the LoRA learns to reproduce the teacher's output for that role. Run the -script once per auxiliary role, each with its own dump file and adapter name — -that is the composable unit; there is no multi-role loop here on purpose. - -Status: this is the PLUMBING. It only runs distillation when invoked explicitly -(``python -m twinkle_agentic.rsi.rsi_distill --input ...``); importing it does -nothing. - -Numbers are INHERITED, not invented: - * LR / BATCH_SIZE / MICRO_BATCH / EPOCHS / MAX_MODEL_LEN <- e18_sft_kod.py - * LORA_RANK / alpha=rank*2 / dropout=0.05 <- rsi_rl.py -All are overridable via the env vars below. - -Env vars --------- - RSI_DUMP_PATH dump JSONL to read (default $LLM_BACKUP_DUMP_PATH) - RSI_ADAPTER adapter name to train + save (default: derived from dump name) - OUTPUT_DIR where the adapter is written (default output/rsi/distill) - MODEL_ID base model (default Qwen/Qwen3-4B) - plus TRAIN_GPUS / TRAIN_FSDP / GPU_MEM / SEED / EPOCHS / BATCH_SIZE / - MICRO_BATCH / LR / MAX_MODEL_LEN / LORA_RANK (all inherited defaults). -""" -import argparse -import json -import os -import random -import shutil -import time -from typing import Any, Dict, List - -from peft import LoraConfig - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.data_format import pack_user_data -from twinkle.model import TransformersModel -from twinkle.processor import InputProcessor -from twinkle.template import Template - -logger = get_logger() - -# ── config (env, all defaults inherited from e18_sft_kod.py / rsi_rl.py) ──── -MODEL_ID = os.environ.get('MODEL_ID', 'Qwen/Qwen3-4B') -OUTPUT_DIR = os.environ.get('OUTPUT_DIR', os.path.join('output', 'rsi', 'distill')) - -TRAIN_GPUS = int(os.environ.get('TRAIN_GPUS', 4)) -TRAIN_FSDP = int(os.environ.get('TRAIN_FSDP', 1)) -TRAIN_DP = TRAIN_GPUS // TRAIN_FSDP -GPU_MEM = float(os.environ.get('GPU_MEM', 0.8)) - -SEED = int(os.environ.get('SEED', 42)) -EPOCHS = float(os.environ.get('EPOCHS', 1)) -BATCH_SIZE = int(os.environ.get('BATCH_SIZE', 16)) -MICRO_BATCH = int(os.environ.get('MICRO_BATCH', 8)) -LR = float(os.environ.get('LR', 1e-5)) -MAX_MODEL_LEN = int(os.environ.get('MAX_MODEL_LEN', 16000)) -LORA_RANK = int(os.environ.get('LORA_RANK', 16)) - -LOG_EVERY_STEPS = int(os.environ.get('LOG_EVERY_STEPS', 1)) -RUN_ID = time.strftime('%m%d-%H%M%S') - - -# =========================================================================== -# data: llm_backup dump -> SFT samples (target = every teacher output) -# =========================================================================== -def _trajectory_messages(traj: Any) -> List[Dict[str, Any]]: - """Pull the message list out of a dumped ``trajectory``. - - twinkle_agentic passes trajectory as ``{"messages": [...], "tools": ...}``; - tolerate a bare list of messages too so the loader does not depend on one - role's exact calling convention. - """ - if isinstance(traj, dict): - msgs = traj.get('messages') - elif isinstance(traj, list): - msgs = traj - else: - msgs = None - return msgs if isinstance(msgs, list) else [] - - -def load_samples(dump_path: str) -> List[Dict[str, Any]]: - """Read the llm_backup dump and build one SFT sample per usable record. - - A record is usable when it has a non-empty trajectory message list AND a - non-empty teacher string. The teacher output becomes the sole trainable - assistant turn appended to the trajectory messages. ``match`` is ignored on - purpose (target = every teacher output). - """ - if not os.path.exists(dump_path): - raise FileNotFoundError(f'找不到 llm_backup dump：{dump_path}') - raw, bad_json = [], 0 - with open(dump_path, 'r', encoding='utf-8') as f: - for line in f: - line = line.strip() - if not line: - continue - try: - raw.append(json.loads(line)) - except Exception: - bad_json += 1 - drop = {'no_messages': 0, 'no_teacher': 0} - out = [] - for r in raw: - msgs = _trajectory_messages(r.get('trajectory')) - if not msgs: - drop['no_messages'] += 1 - continue - teacher = r.get('teacher') - if not isinstance(teacher, str) or not teacher.strip(): - drop['no_teacher'] += 1 - continue - sample_msgs = list(msgs) + [{'role': 'assistant', 'content': teacher}] - out.append({'messages': sample_msgs, 'key': r.get('key', '')}) - logger.info(f'[data] 读入 {len(raw)} 条' - + (f'（跳过 {bad_json} 行半行）' if bad_json else '') - + f'，可训 {len(out)} 条，丢弃明细 {drop}') - return out - - -def make_trajs(batch: List[Dict[str, Any]]) -> List[Dict[str, Any]]: - """样本 -> twinkle 轨迹。只把最后一轮（teacher 输出）标为可训区。""" - trajs = [] - for s in batch: - msgs = s['messages'] - trajs.append({'messages': msgs, - 'user_data': pack_user_data({'key_rounds': [len(msgs) - 1]})}) - return trajs - - -# =========================================================================== -# model: base + one LoRA adapter (rsi_rl.py's config, TransformersModel SFT) -# =========================================================================== -def build_model(adapter_name: str): - twinkle.initialize(mode='ray', nproc_per_node=TRAIN_GPUS, lazy_collect=False, groups=[ - DeviceGroup(name='train', ranks=list(range(TRAIN_GPUS)), device_type='GPU')]) - model = TransformersModel( - model_id=MODEL_ID, remote_group='train', - device_mesh=DeviceMesh.from_sizes(world_size=TRAIN_GPUS, dp_size=TRAIN_DP, - fsdp_size=TRAIN_FSDP), - ddp_config={'find_unused_parameters': False}) - # enable_thinking=True: auxiliary roles produce reasoning before their answer, - # same deployment mode as the teacher that generated the targets. - model.set_template(Template, model_id=MODEL_ID, enable_thinking=True, - max_length=MAX_MODEL_LEN, truncation_strategy='delete') - model.set_processor(InputProcessor, padding_free=False) - model.set_loss('CrossEntropyLoss') - lora_cfg = LoraConfig(target_modules='all-linear', r=LORA_RANK, - lora_alpha=LORA_RANK * 2, lora_dropout=0.05) - model.add_adapter_to_model(adapter_name, lora_cfg, - gradient_accumulation_steps=1) - model.set_optimizer('AdamW', lr=LR) - return model - - -def step_metrics(model) -> Dict[str, float]: - """取本 step 的优化指标（loss/grad_norm/lr）。twinkle 把 loss 格式化成字符串， - 所以用 float() 试转而不是 isinstance 判数值型，否则会静默丢掉 loss。""" - out: Dict[str, float] = {} - for k, val in (model.calculate_metric(is_training=True) or {}).items(): - if isinstance(val, bool): - continue - try: - fval = float(val) - except (TypeError, ValueError): - continue - if k.startswith('learning rate'): - if 'group 1' in k: - out['lr'] = fval - else: - out[k.replace(' ', '_')] = fval - return out - - -def archive_output_dir() -> None: - """启动时把已存在的非空 OUTPUT_DIR 整个 mv 走（sft_log.jsonl 是追写的）， - 否则重跑会把两条 loss 曲线焊进一个文件。与 e18_sft_kod 同一套机制。""" - if not os.path.isdir(OUTPUT_DIR) or not os.listdir(OUTPUT_DIR): - os.makedirs(OUTPUT_DIR, exist_ok=True) - return - stamp = time.strftime('%m%d-%H%M%S', time.localtime(os.path.getmtime(OUTPUT_DIR))) - dst = f'{OUTPUT_DIR}.bak-{stamp}' - i = 1 - while os.path.exists(dst): - dst = f'{OUTPUT_DIR}.bak-{stamp}-{i}' - i += 1 - shutil.move(OUTPUT_DIR, dst) - os.makedirs(OUTPUT_DIR, exist_ok=True) - logger.info(f'[init] 旧输出目录已归档 -> {dst}') - - -# =========================================================================== -# train -# =========================================================================== -def main(): - parser = argparse.ArgumentParser(description='RSI step 4: distill one auxiliary-role LoRA from an llm_backup dump.') - parser.add_argument('--input', default=os.environ.get('RSI_DUMP_PATH', os.environ.get('LLM_BACKUP_DUMP_PATH')), - help='llm_backup dump JSONL (default $RSI_DUMP_PATH / $LLM_BACKUP_DUMP_PATH)') - parser.add_argument('--adapter', default=os.environ.get('RSI_ADAPTER'), - help='adapter name to train and save (default: derived from dump filename)') - args = parser.parse_args() - - dump_path = args.input - if not dump_path: - raise SystemExit('必须给 --input（或设 $RSI_DUMP_PATH / $LLM_BACKUP_DUMP_PATH）指向 llm_backup dump') - adapter_name = args.adapter or os.path.splitext(os.path.basename(dump_path))[0] - - t0 = time.time() - archive_output_dir() - samples = load_samples(dump_path) - if len(samples) < BATCH_SIZE: - raise RuntimeError(f'可用样本 {len(samples)} 条 < BATCH_SIZE {BATCH_SIZE}') - if BATCH_SIZE % TRAIN_DP: - raise RuntimeError(f'BATCH_SIZE({BATCH_SIZE}) 必须是 TRAIN_DP({TRAIN_DP}) 的整倍数') - - model = build_model(adapter_name) - steps_per_epoch = len(samples) // BATCH_SIZE - total_steps = int(steps_per_epoch * EPOCHS) - logger.info(f'RSI-DISTILL start: adapter={adapter_name} n={len(samples)} bs={BATCH_SIZE} ' - f'micro={MICRO_BATCH} lr={LR} epochs={EPOCHS} steps/epoch={steps_per_epoch} ' - f'total_steps={total_steps} rank={LORA_RANK} gpus={TRAIN_GPUS} out={OUTPUT_DIR}') - - log_path = os.path.join(OUTPUT_DIR, 'sft_log.jsonl') - rng = random.Random(SEED) - step = 0 - with open(log_path, 'a', encoding='utf-8') as log_fh: - epoch = 0 - while step < total_steps: - order = list(range(len(samples))) - rng.shuffle(order) # 每 epoch 重洗，种子固定所以可复现 - for bi in range(steps_per_epoch): - if step >= total_steps: - break - batch = [samples[j] for j in order[bi * BATCH_SIZE:(bi + 1) * BATCH_SIZE]] - trajs = make_trajs(batch) - micro = max(TRAIN_DP, min(MICRO_BATCH, len(trajs))) - t_step = time.time() - for i in range(0, len(trajs), micro): - model.forward_backward(inputs=trajs[i:i + micro]) - model.clip_grad_and_step() - step += 1 - row = {'step': step, 'epoch': epoch, 'run': RUN_ID, 'adapter': adapter_name, - 'n_samples': len(batch), 'seconds': round(time.time() - t_step, 2)} - row.update(step_metrics(model)) - log_fh.write(json.dumps(row, ensure_ascii=False) + '\n') - log_fh.flush() - if step % LOG_EVERY_STEPS == 0: - logger.info('[s%d/%d ep%d] ' % (step, total_steps, epoch) - + ' '.join(f'{k}={v:.4g}' for k, v in row.items() - if isinstance(v, float))) - epoch += 1 - - ckpt = model.save(f'{adapter_name}-final', output_dir=OUTPUT_DIR, adapter_name=adapter_name) - logger.info(f'[done] steps={step} 用时 {(time.time() - t0) / 60:.1f} 分钟 -> {ckpt}') - - -if __name__ == '__main__': - main() diff --git a/src/twinkle_agentic/rsi/rsi_prepare.py b/src/twinkle_agentic/rsi/rsi_prepare.py deleted file mode 100644 index 812b9766b..000000000 --- a/src/twinkle_agentic/rsi/rsi_prepare.py +++ /dev/null @@ -1,179 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""RSI step 1 — read a raw data source, parallel-preprocess it with the -twinkle_agentic preprocessor, and write the surviving subset to disk. - -Usage ------ - python -m twinkle_agentic.rsi.rsi_prepare \ - --input /path/to/raw.jsonl \ - --output output/rsi/subset.jsonl \ - --num-proc 4 - -``--input`` accepts a local ``.jsonl``/``.parquet`` path or an ``ms://`` dataset -id; the raw schema is intentionally not pinned here (decided per data source at -test time). Every row must expose a ``messages`` list — the preprocessor keys -off it. Adapt other schemas in :func:`load_source` before the pipeline runs. - -Pipeline --------- -Core steps (no external deps, always on), each using the filter's OWN default -thresholds (no thresholds invented here): - - MessageNormalizer -> MessageSanityFilter -> RefuseFilter -> DeadLoopFilter - -> TokenSoupFilter -> HardFilter - -Optional steps, off by default (enabling needs extra packages): - RSI_USE_LANG=1 LanguageFilter (langid, degrades to heuristic) - RSI_USE_DATAJUICER=1 FixUnicode/RemoveRepeat/SpecialChars/TokenNum (data_juicer[, modelscope]) - RSI_USE_PII=1 PIIPresidioFilter (presidio-analyzer/anonymizer) - -``DedupFilter`` is NOT part of the parallel pipeline: its docstring requires it -to see the whole dataset in one call, so it runs once after the parallel pass. -""" -import argparse -import os - -from twinkle.dataset import Dataset, DatasetMeta -from twinkle.utils import get_logger -from twinkle_agentic.preprocessor import (DeadLoopFilter, DedupFilter, HardFilter, MessageNormalizer, - MessageSanityFilter, QualityPreprocessor, RefuseFilter, TokenSoupFilter, - merge_dropped_shards, run_quality_pipeline, truncate_dropped_logs) - -logger = get_logger() - - -def _env_flag(name: str, default: str = '0') -> bool: - return os.environ.get(name, default).strip().lower() in ('1', 'true', 'yes', 'on') - - -def build_pipeline(): - """Return the ordered list of preprocessor steps for the parallel pass. - - DedupFilter is deliberately excluded (see module docstring); it is applied - separately on the full materialized dataset. - """ - # RSI_NORMALIZE_TOOL_CALLS=0 for pure code data (e.g. MBPP): the bracket-DSL - # parser is a marker-less fallback matching ``[name(``, which is also what a - # python list comprehension or a call-indexed subscript looks like, so the - # rewrite silently deletes real code from the assistant turn. - steps = [ - MessageNormalizer(normalize_tool_calls=_env_flag('RSI_NORMALIZE_TOOL_CALLS', '1')), - MessageSanityFilter(), # role order / tool-id matching / content integrity / sensitive words - RefuseFilter(), # drop assistant self-referential refusals - DeadLoopFilter(), # drop degenerate / stuck (hesitation, cascade, ngram repeat) - TokenSoupFilter(), # drop garbled text (replacement/control/private-use chars, script chaos) - # min_assistant_chars_2turn=0: a single-turn valid tool call (e.g. `[Func(x=1)]`) - # is only tens of chars; HardFilter's default 80-char floor wrongly drops it - # as a "shallow_reply". Zeroing the floor keeps these tool-call rows (Rule 3 - # still removes genuinely empty assistants). Overridable via env. - HardFilter(min_assistant_chars_2turn=int(os.environ.get('RSI_MIN_ASST_CHARS_2TURN', 0))), - ] - if _env_flag('RSI_USE_LANG'): - from twinkle_agentic.preprocessor import LanguageFilter - steps.append(LanguageFilter()) - if _env_flag('RSI_USE_DATAJUICER'): - from twinkle_agentic.preprocessor import (FixUnicodeFilter, RemoveRepeatSentencesFilter, SpecialCharsFilter, - TokenNumFilter) - steps += [FixUnicodeFilter(), RemoveRepeatSentencesFilter(), SpecialCharsFilter(), TokenNumFilter()] - if _env_flag('RSI_USE_PII'): - from twinkle_agentic.preprocessor import PIIPresidioFilter - steps.append(PIIPresidioFilter()) - return steps - - -# ShareGPT `from` value -> standard message role. ToolACE uses -# system/user/assistant/tool; other ShareGPT variants use human/gpt/observation. -_ROLE_MAP = { - 'system': 'system', - 'user': 'user', 'human': 'user', - 'assistant': 'assistant', 'gpt': 'assistant', 'bot': 'assistant', - 'tool': 'tool', 'observation': 'tool', 'function': 'tool', - 'function_call': 'assistant', 'function_response': 'tool', 'tool_response': 'tool', -} - - -def _row_to_messages(row: dict) -> dict: - """Map one ShareGPT ``conversations`` row to a ``messages`` row. - - Only ``from``->``role`` and ``value``->``content`` are rewritten; the tool - call embedded in an assistant turn is left as-is in ``content`` (ToolACE keeps - it as a bracket-DSL string) and parsed later in rsi_refine/rsi_rl. Turns whose - ``from`` is unknown are dropped so no invalid role reaches the pipeline. - """ - messages = [] - for turn in (row.get('conversations') or []): - if not isinstance(turn, dict): - continue - role = _ROLE_MAP.get(str(turn.get('from', '')).lower()) - if role is None: - continue - messages.append({'role': role, 'content': turn.get('value', '') or ''}) - return {'messages': messages, 'id': row.get('id', '')} - - -def load_source(input_path: str, num_proc: int = 4) -> Dataset: - """Load the raw source into a twinkle Dataset. - - A local path is loaded by extension (jsonl->json, parquet, csv...); anything - else is treated as a hub id (e.g. ``ms://org/name``). Rows are passed through - unchanged except for one adaptation: ShareGPT-style rows (a ``conversations`` - list of ``{"from", "value"}`` turns, e.g. ToolACE) are mapped to a standard - ``messages`` list, because the whole preprocessor keys off ``messages``. Rows - that already carry ``messages`` are left untouched. - """ - ds = Dataset(DatasetMeta(dataset_id=input_path)) - cols = ds.dataset.column_names - if 'messages' not in cols and 'conversations' in cols: - logger.info('[rsi_prepare] ShareGPT `conversations` detected -> mapping to `messages`') - # Materialize + convert in Python then rebuild: twinkle's Dataset.map forces - # batched=True and wraps the fn as a Preprocessor, which does not fit a plain - # per-row schema rewrite. The source is small enough to hold in memory. - rows = [_row_to_messages(r) for r in ds.dataset.to_list()] - ds = Dataset(DatasetMeta(data=rows)) - return ds - - -def main(): - parser = argparse.ArgumentParser(description='RSI step 1: preprocess a raw source into a clean subset.') - parser.add_argument('--input', required=True, help='Local .jsonl/.parquet path or an ms:// dataset id.') - parser.add_argument('--output', default='output/rsi/subset.jsonl', help='Where to write the surviving subset.') - parser.add_argument('--num-proc', type=int, default=int(os.environ.get('RSI_NUM_PROC', '4')), - help='Parallel workers for the preprocessor map pass.') - parser.add_argument('--dropped-log', default='', help='Optional JSONL of dropped-row metadata (empty=off).') - args = parser.parse_args() - - os.makedirs(os.path.dirname(os.path.abspath(args.output)) or '.', exist_ok=True) - - pipeline = build_pipeline() - step_names = [type(s).__name__ for s in pipeline] - logger.info(f'[rsi_prepare] pipeline: {" -> ".join(step_names)} + DedupFilter(global)') - - dataset = load_source(args.input, num_proc=args.num_proc) - n_in = len(dataset.dataset) - logger.info(f'[rsi_prepare] loaded {n_in} rows from {args.input}') - - # 'mark' mode + run_quality_pipeline is the ghost-proof parallel path: - # map returns equal-length columns flagged _keep, then a single filter removes. - if args.dropped_log: - truncate_dropped_logs(args.dropped_log) - qp = QualityPreprocessor(pipeline, dropped_log_path=args.dropped_log, drop_mode='mark') - run_quality_pipeline(dataset, qp, num_proc=args.num_proc) - if args.dropped_log: - merge_dropped_shards(args.dropped_log) - - n_after_pipeline = len(dataset.dataset) - logger.info(f'[rsi_prepare] after parallel pipeline: {n_in} -> {n_after_pipeline}') - - # Global longest-wins dedup — must see the whole dataset at once. - rows = dataset.dataset.to_list() - kept, dropped = DedupFilter()(rows) - logger.info(f'[rsi_prepare] after global dedup: {n_after_pipeline} -> {len(kept)} ' - f'(dropped {len(dropped)} duplicates)') - - out = Dataset(DatasetMeta(data=kept)) - out.save_as(args.output) - logger.info(f'[rsi_prepare] wrote {len(kept)} rows -> {args.output}') - - -if __name__ == '__main__': - main() diff --git a/src/twinkle_agentic/rsi/rsi_refine.py b/src/twinkle_agentic/rsi/rsi_refine.py deleted file mode 100644 index fbfd90724..000000000 --- a/src/twinkle_agentic/rsi/rsi_refine.py +++ /dev/null @@ -1,277 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""RSI step 2 — re-analyze each preprocessed trajectory into a STANDARD solving -flow via an injectable teacher API, mark the key rounds, and attach a per-round -reward method. - -Input : the subset produced by rsi_prepare.py (rows with a ``messages`` list). -Output: one refined record per trajectory that could be organized: - { - "id": <passthrough id if present>, - "system": <original system message, kept verbatim (holds tool defs)>, - "query": <original first user message, kept verbatim>, - "tools": <original tools, kept verbatim>, - "rounds": [ {intent, type, tool_call, result, code, reward_method} ], - } -Trajectories the teacher marks unorganizable (e.g. missing tools) are written to -a separate ``*.unorganizable.jsonl`` log and excluded from the standard set. - -Design ------- -- The teacher is INJECTABLE: any OpenAI-compatible endpoint, chosen at runtime - via --teacher-model / --teacher-base-url / --teacher-api-key (env fallbacks - RSI_TEACHER_* then LLM_BACKUP_*). Nothing about the model is hardcoded. -- Key round = a round that carries a tool call OR a code block. The reward - method is attached automatically by round content: - tool call present -> 'tool_result' (executable, verifiable) - code only -> 'rubric' (no executable signal here) - Concrete reward thresholds are intentionally left unset (decided in step 3). -- Heartbeat rounds are already stripped upstream by MessageNormalizer (step 1), - so this stage does not re-handle them. -""" -import argparse -import json -import os -import re -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Any, Dict, List, Optional, Tuple - -from twinkle.data_format import SamplingParams -from twinkle.dataset import Dataset, DatasetMeta -from twinkle.utils import get_logger -from twinkle_agentic.protocol.openai import OpenAI - -logger = get_logger() - -REWARD_TOOL_RESULT = 'tool_result' -REWARD_RUBRIC = 'rubric' - -# Runtime prompt is English on purpose: trajectories are predominantly English, -# and mixing languages degrades the teacher. A Chinese rendering was reviewed -# and approved separately. -REORG_SYSTEM = """\ -You are given ONE agent trajectory that solves a single task through multiple \ -rounds of tool calls. Produce the STANDARD solving flow for this task as \ -STRUCTURED JSON, so it can be parsed. - -The original system message (which holds the tool definitions) and the original \ -user query are kept separately — do NOT rewrite them. Your job is to output only \ -the cleaned, correctly-ordered sequence of KEY rounds: the tool calls and their \ -results that materially lead to the solution. - -Rules: -- Work ONLY from what actually happens in the trajectory and its real tool \ -results. Do NOT invent tools, arguments, or results that are not present. -- Remove redundant, failed-and-abandoned, or out-of-order rounds; reorder the \ -remaining rounds into the logical order that reaches the solution. -- Preserve each kept round's tool call (name + arguments) and the tool's \ -returned result verbatim. -- In each tool_call's "arguments", include ONLY the parameters that were \ -actually passed in that call. Do NOT enumerate the tool's full parameter \ -schema, and do NOT add keys whose value is null/empty for unused parameters. -- Write it as a FRESH, clean standard procedure. Do NOT say the original was \ -wrong, and do NOT reference "previous attempts", "the original code", or "the \ -error above". -- If the trajectory cannot be organized into a standard flow (e.g. required \ -tools are missing, the task never reaches a solution), output exactly \ -{"unorganizable": true, "reason": "<short reason>"} and nothing else. -- Otherwise output ONLY valid JSON in this schema (no prose outside the JSON): -{ - "rounds": [ - {"intent": "<one line>", - "type": "tool" or "code", - "tool_call": {"name": "...", "arguments": {...}} or null, - "result": "<verbatim tool/exec result>", - "code": "<code text if type==code, else null>"} - ] -} -""" - -_CODE_FENCE_RE = re.compile(r'```') -_JSON_FENCE_RE = re.compile(r'^\s*```(?:json)?\s*|\s*```\s*$', re.IGNORECASE) - - -def _first_role(messages: List[Dict[str, Any]], role: str) -> Optional[Dict[str, Any]]: - for m in messages: - if isinstance(m, dict) and m.get('role') == role: - return m - return None - - -def _strip_json_fence(text: str) -> str: - """Remove a leading ```json / trailing ``` wrapper if the model added one.""" - text = text.strip() - text = _JSON_FENCE_RE.sub('', text) - return text.strip() - - -def _strip_null_args(tool_call: Any) -> Any: - """Drop arguments whose value is null/empty from a tool_call. - - Backstop for teachers that echo the tool's full parameter schema and pad - unused params with null: the original calls only pass real args, so a null - here is invented noise that would break step-3 argument matching. Keys with - value None or '' are removed; the rest are kept verbatim. - """ - if not isinstance(tool_call, dict): - return tool_call - args = tool_call.get('arguments') - if isinstance(args, dict): - tool_call['arguments'] = {k: v for k, v in args.items() if v is not None and v != ''} - return tool_call - - -def attach_reward_method(round_obj: Dict[str, Any]) -> str: - """Decide the reward method from the round's actual content (not the label). - - A round with a tool call is verifiable by its tool result; a code-only round - has no executable signal here, so it is scored by rubric. - """ - if round_obj.get('tool_call'): - return REWARD_TOOL_RESULT - if round_obj.get('code') or (isinstance(round_obj.get('type'), str) and round_obj['type'] == 'code'): - return REWARD_RUBRIC - # Fallback: treat as rubric (no tool call, no code detected). - return REWARD_RUBRIC - - -def build_teacher(args) -> OpenAI: - """Construct the injectable teacher client from CLI/env (nothing hardcoded).""" - model = args.teacher_model or os.environ.get('RSI_TEACHER_MODEL') or os.environ.get('LLM_BACKUP_MODEL') - base_url = args.teacher_base_url or os.environ.get('RSI_TEACHER_BASE_URL') or os.environ.get('LLM_BACKUP_BASE_URL') - api_key = args.teacher_api_key or os.environ.get('RSI_TEACHER_API_KEY') or os.environ.get('LLM_BACKUP_API_KEY') - if not model: - raise ValueError('No teacher model given. Pass --teacher-model or set RSI_TEACHER_MODEL/LLM_BACKUP_MODEL.') - timeout = float(os.environ.get('RSI_TEACHER_TIMEOUT', '120')) - max_retries = int(os.environ.get('RSI_TEACHER_MAX_RETRIES', '2')) - return OpenAI(model=model, api_key=api_key, base_url=base_url, - client_kwargs={'timeout': timeout, 'max_retries': max_retries}) - - -def reorder_workflow(traj: Dict[str, Any], teacher: OpenAI, - sampling_params: SamplingParams) -> Tuple[Optional[Dict[str, Any]], Optional[str]]: - """Ask the teacher to reorganize one trajectory into the standard-flow JSON. - - Returns (parsed_json, None) on success, or (None, reason) when the teacher - declares it unorganizable or the response is not parseable. - """ - payload = json.dumps({'messages': traj.get('messages', []), 'tools': traj.get('tools', [])}, - ensure_ascii=False) - request = {'messages': [{'role': 'system', 'content': REORG_SYSTEM}, - {'role': 'user', 'content': payload}]} - message = teacher(request, sampling_params) - if isinstance(message, list): - message = message[0] if message else {} - content = message.get('content', '') if isinstance(message, dict) else '' - if not content.strip(): - return None, 'empty_teacher_response' - try: - parsed = json.loads(_strip_json_fence(content)) - except (ValueError, TypeError): - return None, 'unparseable_json' - if parsed.get('unorganizable'): - return None, f"unorganizable:{parsed.get('reason', '')}" - if not isinstance(parsed.get('rounds'), list) or not parsed['rounds']: - return None, 'no_rounds' - return parsed, None - - -def refine_one(traj: Dict[str, Any], teacher: OpenAI, - sampling_params: SamplingParams) -> Tuple[Optional[Dict[str, Any]], Optional[str]]: - """Turn one raw trajectory into a standard-flow record (or a drop reason).""" - messages = traj.get('messages') or [] - system_msg = _first_role(messages, 'system') - query_msg = _first_role(messages, 'user') - if query_msg is None: - return None, 'no_user_query' - - parsed, reason = reorder_workflow(traj, teacher, sampling_params) - if parsed is None: - return None, reason - - rounds = [] - for r in parsed['rounds']: - if not isinstance(r, dict): - continue - r = dict(r) - if r.get('tool_call'): - r['tool_call'] = _strip_null_args(r['tool_call']) - r['reward_method'] = attach_reward_method(r) - rounds.append(r) - if not rounds: - return None, 'no_valid_rounds' - - record = { - 'id': traj.get('id'), - 'system': system_msg, - 'query': query_msg, - 'tools': traj.get('tools', []), - 'rounds': rounds, - } - return record, None - - -def main(): - parser = argparse.ArgumentParser(description='RSI step 2: refine trajectories into standard solving flows.') - parser.add_argument('--input', default='output/rsi/subset.jsonl', help='Subset from rsi_prepare.py.') - parser.add_argument('--output', default='output/rsi/standard_flows.jsonl', help='Refined standard-flow records.') - parser.add_argument('--teacher-model', default='', help='Teacher model id (or RSI_TEACHER_MODEL/LLM_BACKUP_MODEL).') - parser.add_argument('--teacher-base-url', default='', help='Teacher endpoint base url.') - parser.add_argument('--teacher-api-key', default='', help='Teacher API key.') - parser.add_argument('--max-workers', type=int, default=int(os.environ.get('RSI_MAX_WORKERS', '8')), - help='Concurrent teacher API calls.') - # Generation knobs (defaults shown; override to taste). Low temperature keeps - # the reformat deterministic; max-tokens bounds the JSON output size. - parser.add_argument('--temperature', type=float, default=float(os.environ.get('RSI_TEACHER_TEMPERATURE', '0.0'))) - parser.add_argument('--max-tokens', type=int, default=int(os.environ.get('RSI_TEACHER_MAX_TOKENS', '8192'))) - args = parser.parse_args() - - os.makedirs(os.path.dirname(os.path.abspath(args.output)) or '.', exist_ok=True) - unorg_path = os.path.splitext(args.output)[0] + '.unorganizable.jsonl' - - teacher = build_teacher(args) - sampling_params = SamplingParams(max_tokens=args.max_tokens, num_samples=1, - temperature=args.temperature, top_p=1.0) - - rows = Dataset(DatasetMeta(dataset_id=args.input)).dataset.to_list() - logger.info(f'[rsi_refine] loaded {len(rows)} trajectories from {args.input}') - - kept: List[Dict[str, Any]] = [] - dropped: List[Dict[str, Any]] = [] - total = len(rows) - # Log progress every ~5% (at least every 1) so a long run is not a black box. - step = max(1, total // 20) - done = 0 - with ThreadPoolExecutor(max_workers=args.max_workers) as ex: - futures = {ex.submit(refine_one, row, teacher, sampling_params): row for row in rows} - for fut in as_completed(futures): - row = futures[fut] - try: - record, reason = fut.result() - except Exception as e: # noqa: BLE001 - record, reason = None, f'exception:{type(e).__name__}:{e}' - if record is not None: - kept.append(record) - else: - dropped.append({'id': row.get('id'), 'reason': reason}) - done += 1 - if done % step == 0 or done == total: - logger.info(f'[rsi_refine] progress {done}/{total} ({100 * done // total}%) ' - f'kept={len(kept)} dropped={len(dropped)}') - - # Write JSONL directly (NOT via Dataset.save_as): the Arrow table used by - # save_as unifies the per-round ``arguments`` struct across all rows, padding - # every call with the global union of arg names as null (and coercing ints to - # floats). That corrupts the tool calls, so we serialize each record verbatim. - with open(args.output, 'w', encoding='utf-8') as f: - for rec in kept: - f.write(json.dumps(rec, ensure_ascii=False) + '\n') - if dropped: - with open(unorg_path, 'w', encoding='utf-8') as f: - for d in dropped: - f.write(json.dumps(d, ensure_ascii=False) + '\n') - logger.info(f'[rsi_refine] standard flows: {len(kept)}; dropped/unorganizable: {len(dropped)} ' - f'-> {args.output}' + (f' (+ {unorg_path})' if dropped else '')) - - -if __name__ == '__main__': - main() diff --git a/src/twinkle_agentic/sampler/__init__.py b/src/twinkle_agentic/sampler/__init__.py deleted file mode 100644 index 93d4eec2e..000000000 --- a/src/twinkle_agentic/sampler/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -from .router_sampler import RouterSampler diff --git a/src/twinkle_agentic/sampler/router_sampler.py b/src/twinkle_agentic/sampler/router_sampler.py deleted file mode 100644 index ec57343e0..000000000 --- a/src/twinkle_agentic/sampler/router_sampler.py +++ /dev/null @@ -1,197 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -import httpx -import math -from copy import copy -from typing import Any, Dict, List, Literal, Optional, Union - -from twinkle import get_logger -from twinkle.data_format import SampledSequence, SampleResponse, SamplingParams, Trajectory - -logger = get_logger() - - -def _entropy_from_topk(logprobs_per_token: List[List[tuple]]) -> float: - """Mean per-token entropy approximated from top-K logprobs (renormalized).""" - if not logprobs_per_token: - return float('inf') - total = 0.0 - for candidates in logprobs_per_token: - if not candidates: - total += float('inf') - continue - lps = [lp for _, lp in candidates] - max_lp = max(lps) - # numerically stable softmax over top-K - exps = [math.exp(lp - max_lp) for lp in lps] - z = sum(exps) - total += sum(-(e / z) * (lp - max_lp - math.log(z)) for e, lp in zip(exps, lps)) - return total / len(logprobs_per_token) - - -def _mean_logp(logprobs_per_token: List[List[tuple]], tokens: List[int]) -> float: - """Mean log-probability of generated tokens (sequence-level confidence).""" - if not logprobs_per_token or not tokens: - return float('-inf') - total = 0.0 - count = 0 - for t, candidates in enumerate(logprobs_per_token): - if t >= len(tokens) or not candidates: - continue - tok = tokens[t] - lp = next((v for tid, v in candidates if tid == tok), None) - if lp is None: - lp = candidates[0][1] - total += lp - count += 1 - return total / max(count, 1) - - -class RouterSampler: - """Confidence-based routing sampler. - - Generates with a local sampler first; if confidence is low, falls back - to an OpenAI-compatible endpoint (stronger model). - """ - - def __init__( - self, - sampler, - fallback_endpoint: str, - fallback_model: str = 'default', - fallback_api_key: str = '', - method: Literal['entropy', 'logp'] = 'entropy', - threshold: float = 2.0, - top_k_logprobs: int = 10, - fallback_temperature: float = 0.7, - fallback_max_tokens: int = 4096, - timeout: float = 120.0, - ): - """ - Args: - sampler: Inner sampler instance (e.g. vLLMSampler). - fallback_endpoint: OpenAI-compatible API base URL. - fallback_model: Model name for fallback requests. - fallback_api_key: Bearer token for fallback API. - method: Confidence metric — 'entropy' (route when H > threshold) - or 'logp' (route when mean logp < threshold). - threshold: Routing threshold. For entropy: higher = more routing. - For logp: lower (more negative) = more routing. - top_k_logprobs: Number of top logprobs to request from inner sampler. - fallback_temperature: Temperature for fallback generation. - fallback_max_tokens: Max tokens for fallback generation. - timeout: HTTP timeout for fallback requests. - """ - self.sampler = sampler - self._method = method - self._threshold = threshold - self._top_k = top_k_logprobs - self._fb_temperature = fallback_temperature - self._fb_max_tokens = fallback_max_tokens - self._fb_endpoint = f'{fallback_endpoint.rstrip("/")}/v1/chat/completions' - self._fb_model = fallback_model - headers = {'Content-Type': 'application/json'} - if fallback_api_key: - headers['Authorization'] = f'Bearer {fallback_api_key}' - self._client = httpx.Client(timeout=timeout, headers=headers) - - @property - def template(self): - return self.sampler.template - - def set_template(self, *args, **kwargs): - return self.sampler.set_template(*args, **kwargs) - - def _should_route(self, seq: SampledSequence) -> bool: - if not seq.logprobs: - return True - if self._method == 'entropy': - score = _entropy_from_topk(seq.logprobs) - return score > self._threshold - score = _mean_logp(seq.logprobs, seq.tokens) - return score < self._threshold - - def _fallback_generate(self, trajectory: Trajectory) -> Optional[str]: - messages = trajectory.get('messages', []) - if not messages: - return None - api_messages = [] - for m in messages: - if not isinstance(m, dict): - continue - entry = {'role': m.get('role', 'user')} - content = m.get('content', '') - if isinstance(content, list): - parts = [] - for block in content: - if isinstance(block, dict) and block.get('type') == 'text': - parts.append(block.get('text', '')) - content = '\n'.join(parts) if parts else '' - entry['content'] = content or '' - api_messages.append(entry) - try: - resp = self._client.post( - self._fb_endpoint, - json={ - 'model': self._fb_model, - 'messages': api_messages, - 'temperature': self._fb_temperature, - 'max_tokens': self._fb_max_tokens, - }) - resp.raise_for_status() - choices = resp.json().get('choices', []) - if choices: - return (choices[0].get('message') or {}).get('content', '') - except Exception as e: - logger.warning(f'RouterSampler fallback failed: {e}') - return None - - def sample( - self, - inputs: Union[Dict, List[Dict]], - sampling_params: Optional[Union[SamplingParams, Dict[str, Any]]] = None, - adapter_name: str = '', - adapter_path: Optional[str] = None, - **kwargs, - ) -> List[SampleResponse]: - """Sample with confidence-based routing to fallback model.""" - if sampling_params is None: - sampling_params = SamplingParams() - elif isinstance(sampling_params, dict): - sampling_params = SamplingParams.from_dict(sampling_params) - - # Ensure logprobs are requested for confidence evaluation - routed_params = copy(sampling_params) - if routed_params.logprobs is None or routed_params.logprobs < self._top_k: - routed_params.logprobs = self._top_k - - inputs_list = inputs if isinstance(inputs, list) else [inputs] - is_trajectory = isinstance(inputs_list[0], dict) and 'input_ids' not in inputs_list[0] - - results = self.sampler.sample(inputs_list, routed_params, adapter_name, adapter_path=adapter_path, **kwargs) - - if not is_trajectory: - return results - - for i, (resp, traj) in enumerate(zip(results, inputs_list)): - new_sequences = [] - for seq in resp.sequences: - if self._should_route(seq): - fallback_text = self._fallback_generate(traj) - if fallback_text is not None: - new_sequences.append( - SampledSequence( - stop_reason='stop', - tokens=[], - logprobs=None, - decoded=fallback_text, - )) - continue - new_sequences.append(seq) - results[i] = SampleResponse( - sequences=new_sequences, - prompt_token_ids=resp.prompt_token_ids, - prompt_logprobs=resp.prompt_logprobs, - topk_prompt_logprobs=resp.topk_prompt_logprobs, - ) - - return results diff --git a/tests/twinkle_agentic/test_agentic_rsi.py b/tests/twinkle_agentic/test_agentic_rsi.py index a2c67ff1c..1a1d6f2ba 100644 --- a/tests/twinkle_agentic/test_agentic_rsi.py +++ b/tests/twinkle_agentic/test_agentic_rsi.py @@ -1,10 +1,12 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Tests for the agentic building blocks: program checks and the ms-agent Env. - -No GPU and no ms-agent runtime: the Env is driven with a fake ToolManager that -records what it was asked to run, which is enough to pin the two behaviours the -trainer depends on -- calls arrive batched, and a check's exit status survives -the round trip through a text-only tool. +"""Tests for the agentic building blocks: program checks and the sandboxed Env. + +No GPU, no microVM and no ms-agent runtime. The Env is driven against a fake +sandbox that implements the two operations the real transport uses -- write a +file, run a command -- which is enough to pin what the trainer depends on: a +turn's calls leave as one request, a check's exit status survives the round trip +through a text-only tool, and a mistyped tool name is refused rather than +quietly scored as a failed check. """ import json import os @@ -12,44 +14,107 @@ import tempfile import unittest -sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'src')) +_REPO = os.path.join(os.path.dirname(__file__), '..', '..') +sys.path.insert(0, os.path.join(_REPO, 'src')) +# The RSI wiring lives in cookbook, not in the framework: it is one deployment's +# choice of sandbox backend, and the tests follow it there. +_COOKBOOK = os.path.join(_REPO, 'cookbook', 'rl', 'rsi_agentic') +sys.path.insert(0, _COOKBOOK) +sys.path.insert(0, os.path.join(_COOKBOOK, 'sandbox_server')) +from remote_tool_env import RemoteMsAgentToolEnv # noqa: E402 +from tool_server import _usable_llm, _without_llm_args # noqa: E402 from twinkle_agentic.envs.env_tool import EnvTool # noqa: E402 -from twinkle_agentic.envs.ms_agent_tool_env import MsAgentToolEnv # noqa: E402 from twinkle_agentic.tools.tool_manager import ToolManager # noqa: E402 from twinkle_agentic.verifier.result_check import (Check, CheckContext, # noqa: E402 checks_from_dicts, run_checks) +AGENT_CONFIG = os.path.join(_COOKBOOK, 'rsi_agent.yaml') + +# ms-agent namespaces tools as ``{server}---{tool}``; keep that here so the +# tests exercise the same name resolution production hits. +DEFAULT_TOOLS = [ + {'type': 'function', 'function': {'name': 'code_executor---shell_executor', 'parameters': {}}}, + {'type': 'function', 'function': {'name': 'code_executor---python_executor', 'parameters': {}}}, + {'type': 'function', 'function': {'name': 'file_system---write_file', 'parameters': {}}}, +] + + +class _Result: + + def __init__(self, stdout='', stderr='', exit_code=0): + self.stdout = stdout + self.stderr = stderr + self.exit_code = exit_code + + +class _FakeFiles: + + def __init__(self): + self.store = {} + + def write(self, path, content): + self.store[path] = content + + def read(self, path): + return self.store[path] -class FakeMsToolManager: - """Stands in for ms-agent's ToolManager, recording dispatch shape.""" - # ms-agent namespaces tools as ``{server}---{tool}``; keep that here so the - # tests exercise the same name resolution production hits. - TOOLS = [ - {'tool_name': 'code_executor---shell_executor'}, - {'tool_name': 'code_executor---python_executor'}, - {'tool_name': 'file_system---write_file'}, - ] +class _FakeCommands: - def __init__(self, handler=None): - self.single_calls = [] - self.batch_sizes = [] - self._handler = handler or (lambda call: f'ran {call["tool_name"]}') + def __init__(self, sandbox): + self._sandbox = sandbox - async def get_tools(self): - return list(self.TOOLS) + def run(self, command, timeout=None, background=False, cwd=None): + return self._sandbox.handle(command, background) - async def single_call_tool(self, tool_info): - self.single_calls.append(tool_info) - return self._handler(tool_info) - async def parallel_call_tool(self, tool_list, on_result=None): - self.batch_sizes.append(len(tool_list)) - return [self._handler(call) for call in tool_list] +class FakeSandbox: + """Stands in for an e2b sandbox: a filesystem plus a command channel. - async def cleanup(self): - pass + The Env reaches its in-sandbox server by writing a request file and then + running curl, so a fake that understands those two operations exercises the + real transport -- request shape included -- without booting a microVM. + """ + + def __init__(self, responder=None, tools=None): + self.files = _FakeFiles() + self.commands = _FakeCommands(self) + self.requests = [] + self.killed = False + self.tools = DEFAULT_TOOLS if tools is None else tools + self._responder = responder or (lambda call: f'ran {call["tool_name"]}') + + def kill(self): + self.killed = True + + def handle(self, command, background=False): + if background or 'tool_server.py' in command or command.startswith('tail '): + return _Result() + if command.startswith('find '): + prefix = '/workspace/' + return _Result('\n'.join(p[len(prefix):] for p in self.files.store if p.startswith(prefix))) + if '/health' in command: + return _Result(json.dumps({'status': 'ok'})) + if '/tools' in command: + return _Result(json.dumps({'tools': self.tools})) + if '/call' in command: + payload = json.loads(self.files.store['/opt/rsi/request.json']) + self.requests.append(payload) + results = [{'observation': self._responder(call)} for call in payload['calls']] + return _Result(json.dumps({'results': results})) + raise AssertionError(f'unexpected sandbox command: {command}') + + +def make_env(responder=None, tools=None, **kwargs): + """An Env already attached to a fake sandbox. + + ``reset`` would create a real one, so the sandbox is injected instead and + everything above the e2b SDK boundary still runs for real. + """ + env = RemoteMsAgentToolEnv(template='fake', config_path=AGENT_CONFIG, **kwargs) + env._sandbox = FakeSandbox(responder, tools) + return env class ResultCheckFileTest(unittest.TestCase): @@ -139,97 +204,139 @@ def test_checks_from_dicts(self): self.assertEqual(checks[0].kind, 'file_exists') -class MsAgentToolEnvTest(unittest.TestCase): +class RemoteMsAgentToolEnvTest(unittest.TestCase): def setUp(self): self.tmp = tempfile.mkdtemp(prefix='envtest_') - self.tm = FakeMsToolManager() - self.env = MsAgentToolEnv(tool_manager=self.tm, workspace=self.tmp) + self.env = make_env() + self.sandbox = self.env._sandbox def test_step_forwards_name_and_arguments(self): result = self.env.step('read_file', {'path': 'a.txt'}) - self.assertEqual(self.tm.single_calls[0], - {'tool_name': 'read_file', 'arguments': {'path': 'a.txt'}}) + self.assertEqual(self.sandbox.requests[0]['calls'], + [{'tool_name': 'read_file', 'arguments': {'path': 'a.txt'}}]) self.assertEqual(result.observation, 'ran read_file') - def test_step_batch_uses_one_parallel_call(self): + def test_a_turn_leaves_as_one_request(self): + # Two calls, one sandbox round trip: the server runs them through + # ms-agent's own parallel dispatch, as production would. results = self.env.step_batch([('read_file', {'p': 1}), ('grep', {'q': 'x'})]) - self.assertEqual(self.tm.batch_sizes, [2]) + self.assertEqual(len(self.sandbox.requests), 1) + self.assertEqual(len(self.sandbox.requests[0]['calls']), 2) self.assertEqual([r.observation for r in results], ['ran read_file', 'ran grep']) - def test_single_call_batch_does_not_go_through_parallel(self): - self.env.step_batch([('glob', {})]) - self.assertEqual(self.tm.batch_sizes, []) - self.assertEqual(len(self.tm.single_calls), 1) - def test_observation_is_truncated(self): - tm = FakeMsToolManager(handler=lambda call: 'x' * 50) - env = MsAgentToolEnv(tool_manager=tm, workspace=self.tmp, max_observation_chars=10) + env = make_env(responder=lambda call: 'x' * 50, max_observation_chars=10) obs = env.step('grep', {}).observation self.assertTrue(obs.startswith('x' * 10)) self.assertIn('truncated 40 chars', obs) - def test_non_string_result_is_json_encoded(self): - tm = FakeMsToolManager(handler=lambda call: {'ok': True}) - env = MsAgentToolEnv(tool_manager=tm, workspace=self.tmp) - self.assertEqual(env.step('t', {}).observation, '{"ok": true}') + def test_unreachable_runtime_becomes_an_observation(self): + # A dead sandbox must not take down the training step: the episode plays + # out and scores zero, which is what a broken run deserves anyway. + def explode(command, background=False): + raise RuntimeError('connection refused') + + self.sandbox.handle = explode + obs = self.env.step('read_file', {}).observation + self.assertIn('unreachable', obs) + + def test_tool_schemas_come_from_the_sandbox(self): + self.assertEqual(self.env.tool_names(), + [t['function']['name'] for t in DEFAULT_TOOLS]) - def test_requires_a_tool_manager(self): + def test_resolve_tool_maps_plain_name_onto_namespaced_one(self): + self.assertEqual(self.env.resolve_tool('shell_executor'), + 'code_executor---shell_executor') + # An already-qualified name is left alone. + self.assertEqual(self.env.resolve_tool('file_system---write_file'), + 'file_system---write_file') + + def test_resolve_tool_raises_on_unknown_name(self): + # Silently passing a bad name through would surface as a failed check, + # which is indistinguishable from the task genuinely not being solved. with self.assertRaises(ValueError): - MsAgentToolEnv() + self.env.resolve_tool('no_such_tool') def test_runner_recovers_exit_code_from_text_output(self): # The sandbox tools return prose; the marker is how the exit status # survives. Emulate a shell that echoes the marker. Matching on the # namespaced name also proves the plain name was resolved. - def handler(call): + def responder(call): if call['tool_name'] == 'code_executor---shell_executor': return 'some output\n__TWINKLE_RC__:0' return '__TWINKLE_RC__:3' - env = MsAgentToolEnv(tool_manager=FakeMsToolManager(handler), workspace=self.tmp) - runner = env.runner() + runner = make_env(responder).runner() self.assertEqual(runner('ls', 'shell'), (0, 'some output')) self.assertEqual(runner('boom()', 'python')[0], 3) - def test_resolve_tool_maps_plain_name_onto_namespaced_one(self): - env = MsAgentToolEnv(tool_manager=FakeMsToolManager(), workspace=self.tmp) - self.assertEqual(env.resolve_tool('shell_executor'), - 'code_executor---shell_executor') - # An already-qualified name is left alone. - self.assertEqual(env.resolve_tool('file_system---write_file'), - 'file_system---write_file') - - def test_resolve_tool_raises_on_unknown_name(self): - # Silently passing a bad name through would surface as a failed check, - # which is indistinguishable from the task genuinely not being solved. - env = MsAgentToolEnv(tool_manager=FakeMsToolManager(), workspace=self.tmp) - with self.assertRaises(ValueError): - env.resolve_tool('no_such_tool') - def test_runner_missing_marker_is_a_failure_not_a_pass(self): - env = MsAgentToolEnv(tool_manager=FakeMsToolManager(lambda c: 'sandbox died'), - workspace=self.tmp) - code, out = env.runner()('ls', 'shell') + code, out = make_env(lambda c: 'sandbox died').runner()('ls', 'shell') self.assertNotEqual(code, 0) self.assertIn('sandbox died', out) def test_checks_run_through_the_env_runner(self): - env = MsAgentToolEnv( - tool_manager=FakeMsToolManager(lambda c: '__TWINKLE_RC__:0'), - workspace=self.tmp) + env = make_env(lambda c: '__TWINKLE_RC__:0') report = run_checks([Check(kind='shell', code='true')], CheckContext(workspace=self.tmp, runner=env.runner())) self.assertTrue(report.all_passed) + def test_download_workspace_brings_files_back_for_file_checks(self): + # file_* checks read an ordinary local directory and cannot see into a + # microVM, so the episode's output has to be copied out first. + self.sandbox.files.store['/workspace/report.md'] = '# done\n' + self.sandbox.files.store['/workspace/src/main.py'] = 'print(1)\n' + dest = self.env.download_workspace(os.path.join(self.tmp, 'snap')) + with open(os.path.join(dest, 'report.md'), encoding='utf-8') as f: + self.assertEqual(f.read(), '# done\n') + self.assertTrue(os.path.exists(os.path.join(dest, 'src', 'main.py'))) + + def test_close_kills_the_sandbox(self): + self.env.close() + self.assertTrue(self.sandbox.killed) + + +class ToolServerSchemaTest(unittest.TestCase): + """What /tools advertises must be what the runtime can actually honour.""" + + READ_FILE = { + 'type': 'function', + 'function': { + 'name': 'file_system---read_file', + 'parameters': {'properties': {'paths': {}, 'abbreviate': {}}}, + }, + } + + def test_abbreviate_is_withdrawn_when_no_llm_is_configured(self): + # abbreviate asks an LLM to summarise a file. With no key in the sandbox + # it can only fail, and a model that learns "abbreviate is broken" would + # carry that to a deployment where it works. + stripped = _without_llm_args(self.READ_FILE) + self.assertEqual(sorted(stripped['function']['parameters']['properties']), ['paths']) + # The input is left alone: ms-agent owns that dict. + self.assertIn('abbreviate', self.READ_FILE['function']['parameters']['properties']) + + def test_other_tools_pass_through_untouched(self): + schema = {'type': 'function', 'function': {'name': 'file_system---glob', 'parameters': {}}} + self.assertIs(_without_llm_args(schema), schema) + + def test_missing_llm_section_is_not_a_usable_llm(self): + from omegaconf import OmegaConf + self.assertFalse(_usable_llm(OmegaConf.create({}))) + # ms-agent's default agent.yaml declares a service but no credentials; + # treating that as "configured" is what makes FileSystemTool assert. + self.assertFalse(_usable_llm(OmegaConf.create({'llm': {'service': 'modelscope'}}))) + self.assertTrue(_usable_llm( + OmegaConf.create({'llm': {'service': 'modelscope', 'modelscope_api_key': 'k'}}))) + class ToolBridgeTest(unittest.TestCase): """The prompt's tool list and the executing tool list must be one list.""" def setUp(self): - self.tmp = tempfile.mkdtemp(prefix='bridge_') - self.tm = FakeMsToolManager() - self.env = MsAgentToolEnv(tool_manager=self.tm, workspace=self.tmp) + self.env = make_env() + self.sandbox = self.env._sandbox self.schemas = [ {'type': 'function', 'function': {'name': 'read_file', 'parameters': {}}}, {'type': 'function', 'function': {'name': 'shell_executor', 'parameters': {}}}, @@ -248,7 +355,7 @@ def test_declared_tools_collapse_into_one_step_batch(self): 'function': {'name': 'shell_executor', 'arguments': '{"command": "ls"}'}}, ] out = manager.call_many(calls) - self.assertEqual(self.tm.batch_sizes, [2]) + self.assertEqual(len(self.sandbox.requests), 1) self.assertEqual(out, ['ran read_file', 'ran shell_executor']) def test_nameless_schema_is_refused(self): From 2e9b3c3e78b65daa79a8aaf357b57dd5ed3e482d Mon Sep 17 00:00:00 2001 From: root <yuze.zyz@alibaba-inc.com> Date: Sat, 22 Aug 2026 15:03:07 +0800 Subject: [PATCH 43/60] wip --- cookbook/rl/rsi_agentic/README.md | 148 ----------- .../rl/rsi_agentic/sandbox_server/Dockerfile | 24 -- .../rl/rsi_agentic/sandbox_server/install.sh | 74 ------ .../rl/rsi_agentic/sandbox_server/serve.sh | 20 -- cookbook/rsi/agentic/README.md | 241 +++++++++++++++--- cookbook/rsi/agentic/challenge.py | 121 ++++++++- .../agentic}/remote_tool_env.py | 14 +- .../rsi_agentic_grpo.py => rsi/agentic/rl.py} | 13 +- .../agentic}/rsi_agent.yaml | 0 .../rsi/agentic/sandbox_server/Dockerfile | 54 ++++ .../rsi/agentic/sandbox_server/install.sh | 83 ++++++ cookbook/rsi/agentic/sandbox_server/serve.sh | 112 ++++++++ .../agentic}/sandbox_server/tool_server.py | 133 +++++++++- .../agentic}/tasks.example.jsonl | 0 src/twinkle_agentic/challenger/agentic.py | 109 +++++++- src/twinkle_agentic/harness/ms_agent.py | 88 +++++++ tests/twinkle_agentic/test_agentic_rsi.py | 2 +- tests/twinkle_agentic/test_harness.py | 39 +++ 18 files changed, 943 insertions(+), 332 deletions(-) delete mode 100644 cookbook/rl/rsi_agentic/README.md delete mode 100644 cookbook/rl/rsi_agentic/sandbox_server/Dockerfile delete mode 100644 cookbook/rl/rsi_agentic/sandbox_server/install.sh delete mode 100644 cookbook/rl/rsi_agentic/sandbox_server/serve.sh rename cookbook/{rl/rsi_agentic => rsi/agentic}/remote_tool_env.py (95%) rename cookbook/{rl/rsi_agentic/rsi_agentic_grpo.py => rsi/agentic/rl.py} (97%) rename cookbook/{rl/rsi_agentic => rsi/agentic}/rsi_agent.yaml (100%) create mode 100644 cookbook/rsi/agentic/sandbox_server/Dockerfile create mode 100644 cookbook/rsi/agentic/sandbox_server/install.sh create mode 100644 cookbook/rsi/agentic/sandbox_server/serve.sh rename cookbook/{rl/rsi_agentic => rsi/agentic}/sandbox_server/tool_server.py (64%) rename cookbook/{rl/rsi_agentic => rsi/agentic}/tasks.example.jsonl (100%) diff --git a/cookbook/rl/rsi_agentic/README.md b/cookbook/rl/rsi_agentic/README.md deleted file mode 100644 index 3a26f55ff..000000000 --- a/cookbook/rl/rsi_agentic/README.md +++ /dev/null @@ -1,148 +0,0 @@ -# Agentic RSI - -GRPO on multi-turn tool-using episodes. The solver is an **ms-agent** agent — -shell, filesystem, python, notebook, todo — working inside its own microVM. When -it stops calling tools the episode ends and the task's checks are run against -what it left behind. The checks are ordinary programs, so the same trajectory -always earns the same reward. - -## Where things run - -``` -training host sandbox (one microVM per episode) -───────────────────────────── ───────────────────────────────── -MsAgentHarness tool_server.py - system prompt, message shaping ms-agent ToolManager - llm: and tools: popped -> the real file_system / - constructs no tool at all code_executor / todo_list - ┌──► GET /tools schemas -RemoteMsAgentToolEnv ── curl over ─────┤ POST /call dispatch - forwards tool calls the sandbox │ - copies files back command channel└─── /workspace -``` - -Two properties this layout is built around: - -**Nothing the model emits executes next to the trainer.** The harness has its -`llm` and `tools` sections popped *after* ms-agent merges its own `agent.yaml` -underneath — omitting them from `rsi_agent.yaml` is not enough, since that -default declares `code_executor` and would otherwise put a live shell executor -on the training host with access to the whole machine. - -**The advertised tool contract is read off the code that honours it.** The tool -schemas in the prompt come from `GET /tools` on the sandbox, not from a second -ms-agent next to the trainer. A reimplementation of the tools would have been -much less work, but in RL the policy actively exploits whatever the executor -actually does, and any divergence from the production tools would only surface -after deployment. - -## Setup - -### Environment host - -Needs `/dev/kvm` and kernel 6.8+. Builds the template and runs the AgentENV -server (the same server `cookbook/rl/envs` uses; only the template differs). - -```bash -sh sandbox_server/install.sh # AgentENV + build the RSI template -sh sandbox_server/install.sh --rebuild # after changing the Dockerfile -sh sandbox_server/install.sh --skip-server # template only, server already up - -sh sandbox_server/serve.sh # foreground, binds 127.0.0.1:8000 -NOHUP=1 sh sandbox_server/serve.sh # background -``` - -On a restricted network, point the base image at a reachable registry: - -```bash -BASE_IMAGE=<your-registry>/library/python:3.11-slim sh sandbox_server/install.sh -``` - -The image installs **the `ms-agent/` checkout from this repo**, not the pip -release: the training host imports that same working copy, and a tool whose -output differs by one local commit is a train/serve mismatch the policy absorbs -silently. `install.sh` prints the staged commit — keep the host on it. - -### Training host - -```bash -pip install e2b - -AENV_API_URL=http://<env-host-ip>:8000 \ - python rsi_agentic_grpo.py - -# or through a tunnel: -ssh -N -L 8000:127.0.0.1:8000 root@<env-host-ip> -python rsi_agentic_grpo.py -``` - -Verify a sandbox boots and the runtime comes up before launching training: - -```bash -python -c " -import sys; sys.path.insert(0, '.') -from remote_tool_env import RemoteMsAgentToolEnv -e = RemoteMsAgentToolEnv(template='twinkle-rsi-msagent', config_path='rsi_agent.yaml', - api_url='http://127.0.0.1:8000') -e.reset() -print([t['function']['name'] for t in e.tool_schemas()]) -print(e.step('shell_executor', {'command': 'python -V && pwd'}).observation) -e.close() -" -``` - -## Configuration - -| Variable | Default | | -|---|---|---| -| `AENV_API_URL` | `http://127.0.0.1:8000` | AgentENV server | -| `AENV_TEMPLATE` | `twinkle-rsi-msagent` | template built by `install.sh` | -| `RSI_TASKS` | `tasks.example.jsonl` | task file | -| `RSI_AGENT_CONFIG` | `rsi_agent.yaml` | uploaded into every sandbox | -| `RSI_SANDBOX_TIMEOUT` | `900` | must outlast an episode plus its checks | -| `RSI_ENV_CONCURRENCY` | `16` | parallel boot / scoring | -| `RSI_MAX_TURNS` | `20` | tool-calling turns per episode | -| `RSI_SCORE_MODE` | `fraction` | or `all_or_nothing` | -| `RSI_KEEP_WORKSPACES` | `0` | keep the files copied out of each sandbox | - -Training hyper-parameters come from the CLI, e.g. -`python rsi_agentic_grpo.py --batch-size 2 --num-generations 4 --max-steps 2`. - -Sandbox count is `batch-size × num-generations`, each ~2GiB. At the defaults -(4 × 8) that is 32 microVMs, so size the environment host accordingly. - -## Files - -| File | Role | -|---|---| -| `rsi_agentic_grpo.py` | training loop, episode construction, scoring | -| `remote_tool_env.py` | training-side Env: forwards tool calls, copies files back | -| `rsi_agent.yaml` | ms-agent config — read by *both* halves | -| `sandbox_server/tool_server.py` | in-sandbox HTTP server owning the ToolManager | -| `sandbox_server/Dockerfile` | template image: ms-agent, ripgrep, ipykernel | -| `sandbox_server/install.sh` | build the template (delegates server bootstrap) | -| `sandbox_server/serve.sh` | start AgentENV (delegates to `cookbook/rl/envs`) | -| `tasks.example.jsonl` | one task per line: `id`, `query`, `checks` | - -## Decisions worth knowing - -**`read_file(abbreviate=True)` is withdrawn when no LLM is configured.** That -argument asks an LLM to summarise a file. The sandbox has no API key, so the -tool server drops the unusable `llm` section *and* removes the argument from the -advertised schema — the model is never offered something that can only fail. -Give `rsi_agent.yaml` a real `llm:` section with a key reachable from the -sandbox to get it back. - -**A failed sandbox boot skips the whole batch.** GRPO groups here are -positional: advantages are taken over consecutive runs of `num_generations`, so -dropping one episode would shift every later group onto the wrong task. There is -no retry — a boot failure is logged and the step is abandoned. - -**No web search.** ms-agent's `web_search` key only provides `fetch_page` -(retrieve a known URL). A real search tool needs `EXA_API_KEY` / `SERPAPI_API_KEY` -and is wired separately; no example task requires one. - -**Checks reach the sandbox two ways.** `file_*` checks read a local directory, -so the episode's files are copied out first (`download_workspace`, capped at 200 -files / 1MiB each). `shell` and `python` checks go back into the sandbox through -`env.runner()`, where the interpreter and packages are the ones the agent used. diff --git a/cookbook/rl/rsi_agentic/sandbox_server/Dockerfile b/cookbook/rl/rsi_agentic/sandbox_server/Dockerfile deleted file mode 100644 index 16212f63e..000000000 --- a/cookbook/rl/rsi_agentic/sandbox_server/Dockerfile +++ /dev/null @@ -1,24 +0,0 @@ -FROM python:3.11-slim - -# ripgrep is not optional: file_system's `grep` uses `rg` when it is on PATH and -# silently falls back to a Python scan with a different output shape when it is -# not. The policy is trained on whatever it sees, so the sandbox has to take the -# same branch a serving deployment does. -RUN apt-get update && apt-get install -y --no-install-recommends \ - ca-certificates curl git ripgrep \ - && rm -rf /var/lib/apt/lists/* - -# The exact ms-agent checkout the training host imports, staged here by -# install.sh -- not `pip install ms-agent`. The host runs an editable install of -# this working copy, and a tool whose output differs by one local commit is a -# train/serve mismatch that only shows up after deployment. -COPY ms-agent /opt/ms-agent -RUN pip install --no-cache-dir -e /opt/ms-agent - -# notebook_executor pip-installs these the first time it is called. In a sandbox -# that is either a network round trip at the start of every episode or an -# outright failure on an air-gapped host, so they ship in the image. -RUN pip install --no-cache-dir ipykernel jupyter-client - -ENV PYTHONUNBUFFERED=1 -WORKDIR /workspace diff --git a/cookbook/rl/rsi_agentic/sandbox_server/install.sh b/cookbook/rl/rsi_agentic/sandbox_server/install.sh deleted file mode 100644 index 8cce9c457..000000000 --- a/cookbook/rl/rsi_agentic/sandbox_server/install.sh +++ /dev/null @@ -1,74 +0,0 @@ -#!/bin/sh -# Build the sandbox template for agentic RSI. -# -# The AgentENV server itself is the same one cookbook/rl/envs uses; only the -# template differs. Its bootstrap (install, host provisioning, config seeding) -# is delegated rather than copied, so there is one place to fix when it changes. -set -eu - -TEMPLATE="${TEMPLATE:-twinkle-rsi-msagent}" -# ms-agent pulls in pandas/matplotlib/modelscope and notebook_executor starts a -# real ipykernel, so 1GiB (the plain code template's size) is not enough. -CPU_COUNT="${CPU_COUNT:-2}" -MEMORY_MB="${MEMORY_MB:-2048}" -BASE_IMAGE="${BASE_IMAGE:-}" - -SKIP_SERVER=0 -REBUILD=0 -for arg in "$@"; do - case "$arg" in - --skip-server) SKIP_SERVER=1 ;; - --rebuild) REBUILD=1 ;; - *) echo "Unknown option: $arg" >&2; exit 2 ;; - esac -done - -cd "$(dirname "$0")" -REPO_ROOT=$(cd ../../../.. && pwd) -MS_AGENT="$REPO_ROOT/ms-agent" - -[ -f "$MS_AGENT/setup.py" ] || { - echo "ms-agent checkout not found at $MS_AGENT" >&2 - echo "The image installs the same source the training host imports; without it" >&2 - echo "the sandbox would run a different ms-agent than training assumes." >&2 - exit 1 -} - -if [ "$SKIP_SERVER" = "0" ]; then - echo "==> Installing the AgentENV server (shared with cookbook/rl/envs)" - sh ../../envs/agentenv_server/install.sh --skip-build -fi - -# A fresh staging directory per run, so a stale ms-agent copy can never end up -# in the image and nothing has to be deleted to make room. -STAGE=$(mktemp -d) -echo "==> Staging build context in $STAGE" -cp Dockerfile "$STAGE/" -# .git and caches are megabytes of noise in a build context and would also -# invalidate the layer cache on every commit. -tar -C "$(dirname "$MS_AGENT")" \ - --exclude='.git' --exclude='__pycache__' --exclude='*.pyc' \ - -cf - "$(basename "$MS_AGENT")" | tar -C "$STAGE" -xf - - -HEAD_SHA=$(git -C "$MS_AGENT" rev-parse --short HEAD 2>/dev/null || echo unknown) -echo " ms-agent staged at commit $HEAD_SHA" -echo " the training host must import this same checkout; if it does not," -echo " tool behaviour differs between rollout and the agent you deploy." - -if [ "$REBUILD" = "1" ]; then - echo "==> Deleting template '$TEMPLATE'" - aenv template delete "$TEMPLATE" || true -fi - -echo "==> Building template '$TEMPLATE' (cpu=$CPU_COUNT mem=${MEMORY_MB}MiB)" -set -- "$STAGE/Dockerfile" -t "$TEMPLATE" --cpu-count "$CPU_COUNT" --memory-mb "$MEMORY_MB" -[ -n "$BASE_IMAGE" ] && set -- "$@" --image "$BASE_IMAGE" -aenv build "$@" - -echo -echo "Build runs server-side and takes a few minutes. Follow it with:" -echo " aenv template watch <template-id> # id printed above" -echo " aenv template list # confirm it reaches ready" -echo -echo "Then start the server:" -echo " sh serve.sh" diff --git a/cookbook/rl/rsi_agentic/sandbox_server/serve.sh b/cookbook/rl/rsi_agentic/sandbox_server/serve.sh deleted file mode 100644 index 74fdab11d..000000000 --- a/cookbook/rl/rsi_agentic/sandbox_server/serve.sh +++ /dev/null @@ -1,20 +0,0 @@ -#!/bin/sh -# Start the AgentENV server that hosts the RSI sandboxes. -# -# Deliberately a delegation, not a copy. It is the same server as -# cookbook/rl/envs uses -- RSI only changes which template the sandboxes boot -# from -- and that script carries a hundred lines of host-provisioning detail -# (capability wrapper, config path, systemd handover) that would silently drift -# if it existed twice. -# -# All of its environment variables still apply, e.g.: -# API_ADDR=0.0.0.0:8000 NOHUP=1 sh serve.sh -set -eu -cd "$(dirname "$0")" - -SHARED=../../envs/agentenv_server/serve.sh -[ -f "$SHARED" ] || { - echo "Shared AgentENV launcher not found: $SHARED" >&2 - exit 1 -} -exec sh "$SHARED" "$@" diff --git a/cookbook/rsi/agentic/README.md b/cookbook/rsi/agentic/README.md index ebb782979..440dc3e9d 100644 --- a/cookbook/rsi/agentic/README.md +++ b/cookbook/rsi/agentic/README.md @@ -1,54 +1,229 @@ -# RSI Agentic Self-Play +# Agentic RSI -## 前置条件 +Self-play for tool-using agents. The model invents its own tasks by doing them, +then trains on the ones it solves only sometimes. -- 沙箱环境已启动（AgentENV / e2b），模板由 `cookbook/rl/rsi_agentic/sandbox_server/install.sh` 构建 -- `AENV_API_URL` 和 `AENV_TEMPLATE` 环境变量已设置 -- ms-agent 配置文件就绪（默认 `cookbook/rl/rsi_agentic/rsi_agent.yaml`） +The solver is an **ms-agent** agent -- shell, filesystem, python, notebook, +todo -- working inside a microVM. When it stops calling tools the episode ends +and the task's check script is run against what it left behind. Checks are +ordinary programs, so the same trajectory always earns the same reward. -## Step 1: 生成任务 +## Pipeline + +``` +challenge.py rl.py +┌──────────────────────────┐ ┌────────────────────────┐ +│ 1 direction + keywords │ │ 1 read flows │ +│ 2 model acts in sandbox │ │ 2 boot sandboxes │ +│ 3 model writes checks │ flows │ 3 solver works N turns │ +│ 4 run checks (verify) │ ───────► │ 4 run check_script │ +│ 5 model writes statement │ │ 5 GRPO │ +│ 6 difficulty filter │ └────────────────────────┘ +└──────────────────────────┘ +``` + +`challenge.py` builds tasks backwards: the end state exists before the question +does, so nothing it produces can be unachievable. The difficulty filter then +drops tasks the solver always passes or always fails -- either way GRPO gets a +zero gradient from the whole group. + +## Where things run + +``` +training host sandbox (microVM) +───────────────────────────── ───────────────────────────────── +MsAgentHarness tool_server.py + system prompt, message shaping ms-agent ToolManager + llm: and tools: popped -> the real file_system / + constructs no tool at all code_executor / todo_list + ┌──► GET /tools schemas +RemoteMsAgentToolEnv ── curl over ─────┤ POST /call dispatch + forwards tool calls the sandbox │ + copies files back command channel└─── /workspace +``` + +Two properties this layout is built around: + +**Nothing the model emits executes next to the trainer.** The harness has its +`llm` and `tools` sections popped *after* ms-agent merges its own `agent.yaml` +underneath -- omitting them from `rsi_agent.yaml` is not enough, since that +default declares `code_executor` and would otherwise put a live shell executor +on the training host with access to the whole machine. + +**The advertised tool contract is read off the code that honours it.** The tool +schemas in the prompt come from `GET /tools` on the sandbox, not from a second +ms-agent next to the trainer. In RL the policy actively exploits whatever the +executor actually does, and any divergence from the production tools would only +surface after deployment. + +## Setup + +### Environment host + +Needs `/dev/kvm` and kernel 6.8+. Builds the template and runs the AgentENV +server. + +```bash +sh sandbox_server/install.sh # AgentENV + build the template +sh sandbox_server/install.sh --rebuild # after changing the Dockerfile +sh sandbox_server/install.sh --skip-install # template only, AgentENV already up + +sh sandbox_server/serve.sh # foreground, binds 127.0.0.1:8000 +NOHUP=1 sh sandbox_server/serve.sh # background +``` + +On a restricted network the base image will not resolve from Docker Hub; point it +at a reachable mirror (forwarded to `aenv build --image`, so the Dockerfile stays +untouched): + +```bash +BASE_IMAGE=docker.m.daocloud.io/library/python:3.11-slim sh sandbox_server/install.sh +``` + +`docker.m.daocloud.io` is a third-party Docker Hub proxy -- every sandbox's base +image comes through it. Substitute your own Aliyun accelerator address +(`<id>.mirror.aliyuncs.com`) if you would rather not depend on one. + +The Dockerfile also pins `PIP_INDEX_URL` to an Aliyun mirror for the same +reason -- edit those two lines if your host reaches pypi.org directly. + +The image clones **ms-agent from source** (`--depth 1` of `main`), not the pip +release: the tools the policy is trained against are the ones in the repository, +and a released wheel can lag behind it. + +Verify a sandbox boots and the runtime comes up before anything else: + +```bash +python -c " +import sys; sys.path.insert(0, '.') +from remote_tool_env import RemoteMsAgentToolEnv +e = RemoteMsAgentToolEnv(template='twinkle-rsi-msagent', config_path='rsi_agent.yaml', + api_url='http://127.0.0.1:8000') +e.reset() +print([t['function']['name'] for t in e.tool_schemas()]) +print(e.step('shell_executor', {'command': 'python -V && pwd'}).observation) +e.close() +" +``` + +### Step 1 -- generate tasks ```bash -python cookbook/rsi/agentic/challenge.py \ +pip install e2b + +python challenge.py \ --keep-target 200 \ - --sandbox-template $AENV_TEMPLATE \ - --sandbox-api-url $AENV_API_URL \ + --sandbox-template twinkle-rsi-msagent \ + --sandbox-api-url http://<env-host-ip>:8000 \ --sampler-gpus 4 ``` -产出：`output/rsi_agentic/challenge_flows.jsonl` +Writes `output/rsi_agentic/challenge_flows.jsonl`, one task per line: +`{id, query, check_script, n_pass, n_rollouts, keywords, seeded}`. + +Also writes `output/rsi_agentic/propose_traj/` -- the rounds that *produced* each +task, kept and rejected alike, as one `.npz` of token ids / labels / logprobs per +attempt plus an `index.jsonl` carrying the text and the outcome. Nothing reads it +yet. It exists because proposing is generation like any other, so those rounds +could be trained on later; rejects are in there on purpose, since a set of +kept-only attempts has no zero-reward half to contrast against. `pass_rate` is +stored raw -- mapping it onto a difficulty score means choosing a target rate, +which is a training decision, not a dump format. Pass `--dump-propose-traj ''` +to turn it off; expect it to dwarf the task file. -每行一个任务：`{id, query, check_script, n_pass, n_rollouts, keywords, seeded}` +Useful flags: `--seed-file` (start from existing trajectories), `--keywords-n 0` +(no keyword bank), `--solver-rollouts` (attempts per task in the difficulty +filter), `--max-turns` (tool-calling turns per episode). -可选参数： -- `--seed-file seeds.jsonl` 用已有 trajectory 做起点 -- `--keywords-n 0` 关闭关键词库 -- `--solver-rollouts 4` 难度过滤尝试次数 -- `--max-turns 20` round1 最大工具调用轮数 +Round 1 is serial -- one episode at a time, workspace cleared in between -- so +`--max-proposals-per-round` trades throughput against how often the estimator +recalibrates. -## Step 2: 训练（GRPO） +### Step 2 -- train ```bash -AENV_API_URL=http://... \ +AENV_API_URL=http://<env-host-ip>:8000 \ AENV_TEMPLATE=twinkle-rsi-msagent \ RSI_TASKS=output/rsi_agentic/challenge_flows.jsonl \ - python cookbook/rl/rsi_agentic/rsi_agentic_grpo.py \ - --model-id ms://Qwen/Qwen3-4B \ - --model-gpus 4 --sampler-gpus 4 + python rl.py --model-id ms://Qwen/Qwen3-4B \ + --model-gpus 4 --sampler-gpus 4 ``` -训练脚本自动识别 `check_script` 格式，在沙箱中跑检查脚本评分（exit 0 = 1.0）。 +`rl.py` accepts both task formats: `check_script` (from `challenge.py`, scored +by exit status) and structured `checks` (see `tasks.example.jsonl`). -## 流程总结 +Through a tunnel instead: +```bash +ssh -N -L 8000:127.0.0.1:8000 root@<env-host-ip> ``` -challenge.py rsi_agentic_grpo.py -┌─────────────────────┐ ┌─────────────────────┐ -│ 1. 选方向+关键词 │ │ 1. 读 flows │ -│ 2. 模型在沙箱做事 │ │ 2. 起沙箱 │ -│ 3. 模型写检查脚本 │ flows │ 3. solver 多轮做题 │ -│ 4. 跑检查（验证） │ ──────► │ 4. 跑 check_script │ -│ 5. 模型写题目描述 │ │ 5. GRPO 训练 │ -│ 6. 难度过滤 │ └─────────────────────┘ -└─────────────────────┘ -``` + +## Configuration + +`challenge.py` is all command-line flags (`--help`). `rl.py` reads: + +| Variable | Default | | +|---|---|---| +| `AENV_API_URL` | `http://127.0.0.1:8000` | AgentENV server | +| `AENV_TEMPLATE` | `twinkle-rsi-msagent` | template built by `install.sh` | +| `RSI_TASKS` | `tasks.example.jsonl` | task file | +| `RSI_AGENT_CONFIG` | `rsi_agent.yaml` | uploaded into every sandbox | +| `RSI_SANDBOX_TIMEOUT` | `900` | must outlast an episode plus its checks | +| `RSI_ENV_CONCURRENCY` | `16` | parallel boot / scoring | +| `RSI_MAX_TURNS` | `20` | tool-calling turns per episode | +| `RSI_SCORE_MODE` | `fraction` | or `all_or_nothing`; structured checks only | +| `RSI_KEEP_WORKSPACES` | `0` | keep the files copied out of each sandbox | + +Training hyper-parameters come from the CLI, e.g. +`python rl.py --batch-size 2 --num-generations 4 --max-steps 2`. + +Sandbox count during training is `batch-size x num-generations`, each ~2GiB. At +the defaults (4 x 8) that is 32 microVMs, so size the environment host +accordingly. + +## Files + +| File | Role | +|---|---| +| `challenge.py` | task generation: act, write checks, verify, describe, filter | +| `prompts.py` | every string the challenger sends; categories live here too | +| `rl.py` | training loop, episode construction, scoring | +| `remote_tool_env.py` | training-side Env: forwards tool calls, copies files back | +| `rsi_agent.yaml` | ms-agent config -- read by *both* halves | +| `sandbox_server/tool_server.py` | in-sandbox HTTP server owning the ToolManager | +| `sandbox_server/Dockerfile` | template image: ms-agent, ripgrep, ipykernel | +| `sandbox_server/install.sh` | install AgentENV + build the template | +| `sandbox_server/serve.sh` | start the AgentENV server | +| `tasks.example.jsonl` | hand-written tasks in the structured `checks` format | + +The machinery both scripts call lives in `twinkle_agentic.challenger`; only +wiring and prompt text are here. + +## Decisions worth knowing + +**Round 1 is serial.** Every episode needs an empty workspace and they share one +long-lived sandbox, so proposals cannot overlap -- the second episode would see +the first one's files and its checks would pass for free. The difficulty filter +resets between attempts for the same reason. + +**Checks are a python script, not a structured list.** `challenge.py` asks the +model to write asserts against the state it just produced, which is the same +trick the code challenger uses (execute first, capture the result, make that the +ground truth). Exit status is the whole verdict: no partial credit, no judge +model, no drift between rounds. + +**`read_file(abbreviate=True)` is withdrawn when no LLM is configured.** That +argument asks an LLM to summarise a file. The sandbox has no API key, so the +tool server drops the unusable `llm` section *and* removes the argument from the +advertised schema -- the model is never offered something that can only fail. +Give `rsi_agent.yaml` a real `llm:` section with a key reachable from the sandbox +to get it back. + +**A failed sandbox boot skips the whole training batch.** GRPO groups here are +positional: advantages are taken over consecutive runs of `num_generations`, so +dropping one episode would shift every later group onto the wrong task. There is +no retry -- a boot failure is logged and the step is abandoned. + +**No web search.** ms-agent's `web_search` key only provides `fetch_page` +(retrieve a known URL). A real search tool needs `EXA_API_KEY` / `SERPAPI_API_KEY` +and is wired separately; no example task requires one. diff --git a/cookbook/rsi/agentic/challenge.py b/cookbook/rsi/agentic/challenge.py index 878bfd569..650f279c6 100644 --- a/cookbook/rsi/agentic/challenge.py +++ b/cookbook/rsi/agentic/challenge.py @@ -24,6 +24,7 @@ import os import sys +import numpy as np import twinkle from twinkle import DeviceGroup, DeviceMesh, get_logger from twinkle.data_format import SamplingParams, user_data_get @@ -35,6 +36,7 @@ sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from prompts import CATEGORIES, CATEGORY_DESC, agentic_prompts # noqa: E402 +from remote_tool_env import RemoteMsAgentToolEnv # noqa: E402 logger = get_logger() @@ -94,7 +96,7 @@ def parse_args(): help='AgentENV/e2b template name (required)') p.add_argument('--sandbox-api-url', default='', help='AgentENV server URL (or AENV_API_URL env var)') - p.add_argument('--agent-config', default='cookbook/rl/rsi_agentic/rsi_agent.yaml', + p.add_argument('--agent-config', default='cookbook/rsi/agentic/rsi_agent.yaml', help='ms-agent yaml for the sandbox tool server') p.add_argument('--sandbox-timeout', type=int, default=900) p.add_argument('--workspace', default='/workspace', @@ -104,16 +106,16 @@ def parse_args(): p.add_argument('--random-seed', type=int, default=0) p.add_argument('--out-flows', default='output/rsi_agentic/challenge_flows.jsonl') p.add_argument('--dump-rejected', default='output/rsi_agentic/challenge_rejected.jsonl') + p.add_argument('--dump-propose-traj', default='output/rsi_agentic/propose_traj', + help='directory for the proposing rounds (token ids + logprobs, one npz ' + 'per attempt plus index.jsonl). Empty string turns it off; keeping ' + 'it is what leaves the door open to training the challenger itself.') p.add_argument('--no-sort-by-difficulty', action='store_true') return p.parse_args() def build_env(args): """Create the long-lived sandbox environment.""" - sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), - '..', '..', 'rl', 'rsi_agentic')) - from remote_tool_env import RemoteMsAgentToolEnv # noqa: E402 - template = args.sandbox_template or os.environ.get('AENV_TEMPLATE', '') api_url = args.sandbox_api_url or os.environ.get('AENV_API_URL', '') if not template: @@ -177,9 +179,36 @@ def main(): # the marker protocol, so we don't rely on string matching. runner = env.runner() + # Cleared through the python tool, not `rm -rf`: ms-agent's safety policy + # rejects `rm -rf` outright ("Blocked by safety rule"), and it rejects globs + # in write operations, which rules out `find -delete` too. The script asserts + # the directory really is empty so a future policy change surfaces as a + # failed reset instead of tasks quietly inheriting the previous workspace. + _CLEAR = ''' +import os, shutil +root = {workspace!r} +os.makedirs(root, exist_ok=True) +for name in os.listdir(root): + path = os.path.join(root, name) + if os.path.isdir(path) and not os.path.islink(path): + shutil.rmtree(path, ignore_errors=True) + else: + os.remove(path) +leftover = os.listdir(root) +assert not leftover, 'workspace not empty after clear: %r' % (leftover,) +''' + def reset_fn(): - """Clear the sandbox workspace between episodes.""" - runner(f'rm -rf {args.workspace}/* {args.workspace}/.* 2>/dev/null; true', 'shell') + """Empty the sandbox workspace before an episode. + + Raises rather than returning: every caller depends on a clean start, and + a silent no-op here means a task inherits the previous task's files -- + which lets a solver pass without doing anything and makes the difficulty + numbers meaningless. + """ + exit_code, output = runner(_CLEAR.format(workspace=args.workspace), 'python') + if exit_code != 0: + raise RuntimeError(f'workspace reset failed (exit {exit_code}): {output[-400:]}') def run_check_fn(script: str): """Run a python check script in the sandbox; returns (exit_code, output).""" @@ -215,6 +244,8 @@ def _reject(record): if rejected is not None: rejected.write(json.dumps(record, ensure_ascii=False, default=str) + '\n') + propose_writer = ProposeTrajWriter(args.dump_propose_traj) + # Build challenger prompts = agentic_prompts() challenger = AgenticChallenger( @@ -239,6 +270,7 @@ def _reject(record): min_batch=args.sampler_gpus, problem_max_chars=args.problem_max_chars, reject_sink=_reject, + propose_sink=propose_writer.write, max_proposals_per_round=args.max_proposals_per_round, solver_rollouts=args.solver_rollouts, keep_min_pass=args.keep_min_pass, @@ -257,6 +289,7 @@ def _reject(record): f'stats {challenger.stats}') if rejected is not None: rejected.close() + propose_writer.close() if store is not None: challenger.expand_hard_keywords() @@ -279,6 +312,80 @@ def _reject(record): env.close() +class ProposeTrajWriter: + """Persist the proposing rounds so the challenger could be trained later. + + One ``.npz`` per proposal attempt holds the arrays, and one line per attempt + in ``index.jsonl`` holds everything a human reads plus the outcome. Splitting + them is what keeps this affordable: a 20-turn agentic episode is tens of + thousands of token ids, which as JSON is an order of magnitude larger than + the same numbers as int32. + + ``logprobs`` arrive as ``[[(token_id, logprob)]]`` and are flattened to the + logprob column alone -- that is the shape GRPO's ``old_logps`` wants, and the + token each one belongs to is already in ``labels``. + + Rejected attempts are written too. Their outcome is the reward's zero, and a + dump of kept-only attempts would have nothing to contrast against. + """ + + def __init__(self, out_dir: str): + self.dir = out_dir + self.index = None + self.n = 0 + if not out_dir: + return + os.makedirs(out_dir, exist_ok=True) + self.index = open(os.path.join(out_dir, 'index.jsonl'), 'w', encoding='utf-8') + + def write(self, record): + if self.index is None: + return + trace_id = f'p{self.n:06d}' + self.n += 1 + arrays, meta = {}, [] + for i, rnd in enumerate(record.get('rounds') or []): + labels = rnd.get('labels') or [] + logprobs = rnd.get('logprobs') or [] + if rnd.get('input_ids'): + arrays[f'r{i}_input_ids'] = np.asarray(rnd['input_ids'], dtype=np.int32) + if labels: + arrays[f'r{i}_labels'] = np.asarray(labels, dtype=np.int32) + if logprobs: + arrays[f'r{i}_logprobs'] = np.asarray( + [lp[0][1] for lp in logprobs], dtype=np.float32) + meta.append({ + 'stage': rnd.get('stage'), + 'messages': rnd.get('messages') or [], + 'n_tokens': len(rnd.get('input_ids') or []), + 'n_trainable': sum(1 for label in labels if label != -100), + 'n_logprobs': len(logprobs), + }) + # No arrays means the explorer was text-only (an API rollout), so there + # is nothing trainable to store -- record the attempt without an npz + # rather than leaving thousands of empty archives behind. + npz_name = f'{trace_id}.npz' if arrays else None + if arrays: + np.savez_compressed(os.path.join(self.dir, npz_name), **arrays) + line = { + 'trace_id': trace_id, + 'npz': npz_name, + 'outcome': record.get('outcome'), + 'n_pass': record.get('n_pass'), + 'n_rollouts': record.get('n_rollouts'), + 'pass_rate': record.get('pass_rate'), + 'keywords': record.get('keywords'), + 'seeded': record.get('seeded'), + 'rounds': meta, + } + self.index.write(json.dumps(line, ensure_ascii=False, default=str) + '\n') + + def close(self): + if self.index is not None: + self.index.close() + logger.info(f'[challenge] wrote {self.n} propose traces -> {self.dir}') + + def write_flows(kept, args): """Write one flow per task.""" with open(args.out_flows, 'w', encoding='utf-8') as f: diff --git a/cookbook/rl/rsi_agentic/remote_tool_env.py b/cookbook/rsi/agentic/remote_tool_env.py similarity index 95% rename from cookbook/rl/rsi_agentic/remote_tool_env.py rename to cookbook/rsi/agentic/remote_tool_env.py index 889bcd327..d3e0cbf2f 100644 --- a/cookbook/rl/rsi_agentic/remote_tool_env.py +++ b/cookbook/rsi/agentic/remote_tool_env.py @@ -41,8 +41,13 @@ try: {body} except SystemExit as _e: - print('{mark}:%d' % (_e.code or 0)) - sys.exit(0) + # Print the status and stop -- do NOT re-raise. SystemExit inherits from + # BaseException, so ms-agent's `except Exception` around the exec does not + # catch it; letting it escape kills the whole tool server process, and the + # sandbox is shared by every task in the run. `else` is already skipped + # because the exception was handled, so nothing further is needed. + _c = _e.code + print('{mark}:%d' % (0 if _c is None else _c if isinstance(_c, int) else 1)) except BaseException: traceback.print_exc() print('{mark}:1') @@ -281,7 +286,10 @@ def _create_sandbox(self): # ``e2b_[0-9a-f]+`` before sending anything. This is the SDK's own # opt-out for deployments that do not mint e2b-format keys. os.environ.setdefault('E2B_VALIDATE_API_KEY', 'false') - return Sandbox(template=self._template, timeout=self._sandbox_timeout) + # ``Sandbox.create``, not ``Sandbox(...)``: since e2b 2.x the constructor + # takes connection options for an *existing* sandbox and rejects + # ``template``, while the classmethod is what provisions a new one. + return Sandbox.create(template=self._template, timeout=self._sandbox_timeout) def _upload(self) -> None: """Push the yaml and the server script into the sandbox. diff --git a/cookbook/rl/rsi_agentic/rsi_agentic_grpo.py b/cookbook/rsi/agentic/rl.py similarity index 97% rename from cookbook/rl/rsi_agentic/rsi_agentic_grpo.py rename to cookbook/rsi/agentic/rl.py index 4b49bb490..ab693e126 100644 --- a/cookbook/rl/rsi_agentic/rsi_agentic_grpo.py +++ b/cookbook/rsi/agentic/rl.py @@ -20,13 +20,14 @@ Usage: AENV_API_URL=http://127.0.0.1:8000 \\ - RSI_TASKS=cookbook/rl/rsi_agentic/tasks.example.jsonl \\ - python cookbook/rl/rsi_agentic/rsi_agentic_grpo.py + RSI_TASKS=cookbook/rsi/agentic/tasks.example.jsonl \\ + python cookbook/rsi/agentic/rl.py See README.md for building the template and starting the sandbox server. -Task file: one JSON object per line, with ``id``, ``query`` and ``checks`` -(see tasks.example.jsonl and twinkle_agentic.verifier.result_check.Check). +Task file: one JSON object per line, with ``id``, ``query`` and either +``check_script`` (a python script, from ``challenge.py``) or ``checks`` +(structured, see twinkle_agentic.verifier.result_check.Check). """ import json import os @@ -84,8 +85,8 @@ MAX_TRAJECTORY_TOKENS = int(os.environ.get('RSI_MAX_TRAJ_TOKENS', 32768)) MAX_TURNS = int(os.environ.get('RSI_MAX_TURNS', 20)) -TASKS_PATH = os.environ.get('RSI_TASKS', 'cookbook/rl/rsi_agentic/tasks.example.jsonl') -AGENT_CONFIG = os.environ.get('RSI_AGENT_CONFIG', 'cookbook/rl/rsi_agentic/rsi_agent.yaml') +TASKS_PATH = os.environ.get('RSI_TASKS', 'cookbook/rsi/agentic/tasks.example.jsonl') +AGENT_CONFIG = os.environ.get('RSI_AGENT_CONFIG', 'cookbook/rsi/agentic/rsi_agent.yaml') RUN_DIR = os.environ.get('RSI_RUN_DIR', 'output/rsi_agentic/run') # Sandbox backend. The template is built by sandbox_server/install.sh. diff --git a/cookbook/rl/rsi_agentic/rsi_agent.yaml b/cookbook/rsi/agentic/rsi_agent.yaml similarity index 100% rename from cookbook/rl/rsi_agentic/rsi_agent.yaml rename to cookbook/rsi/agentic/rsi_agent.yaml diff --git a/cookbook/rsi/agentic/sandbox_server/Dockerfile b/cookbook/rsi/agentic/sandbox_server/Dockerfile new file mode 100644 index 000000000..d5e245855 --- /dev/null +++ b/cookbook/rsi/agentic/sandbox_server/Dockerfile @@ -0,0 +1,54 @@ +FROM python:3.11-slim + +# Every instruction has to fit on ONE line: aenv's Dockerfile parser does not +# join backslash continuations, and reports the second line as an unknown +# instruction ("Dockerfile instruction ca-certificates is not supported"). + +# The build runs on the AgentENV host. Ours sits behind a firewall that cannot +# reach files.pythonhosted.org (it 302s and then stalls), and ms-agent pulls in +# pandas/matplotlib/modelscope -- a stalled index is a build that never +# finishes. Point these at a different index if your host reaches pypi directly. +ENV PIP_INDEX_URL=https://mirrors.aliyun.com/pypi/simple/ +ENV PIP_TRUSTED_HOST=mirrors.aliyun.com + +# ripgrep is not optional: file_system's `grep` uses `rg` when it is on PATH and +# silently falls back to a Python scan with a different output shape when it is +# not. The policy is trained on whatever it sees, so the sandbox has to take the +# same branch a serving deployment does. +RUN apt-get update && apt-get install -y --no-install-recommends ca-certificates curl git ripgrep && rm -rf /var/lib/apt/lists/* + +# ms-agent from source, not `pip install ms-agent`: the tools the policy is +# trained against are the ones in this repository, and a released wheel can lag +# behind it. +# +# curl, not `git clone`: Debian's git links against GnuTLS, and cloning GitHub +# from inside the build VM stalls for ~15 minutes and then dies with "GnuTLS +# recv error (-110)". The same host pulls the codeload tarball over curl's +# OpenSSL at 12MB/s. Version comes from ms_agent/version.py, not scm, so the +# missing .git costs nothing. +RUN mkdir -p /opt/ms-agent && curl -fsSL https://codeload.github.com/modelscope/ms-agent/tar.gz/refs/heads/main | tar -xz -C /opt/ms-agent --strip-components=1 && pip install --no-cache-dir -e /opt/ms-agent + +# One pip call, not three: aenv runs each instruction inside a fresh microVM with +# no layer cache, so every extra RUN is another full dependency resolution. +# +# httpx -- imported by ms_agent/llm/openai_llm.py, which `import ms_agent` +# reaches via tools/search/websearch_tool, but no requirements file +# declares it. Without it the image builds cleanly and then every +# sandbox fails at import. +# ipykernel, +# jupyter-client +# -- notebook_executor pip-installs these on first call. +# the rest -- LocalCodeExecutionTool._check_dependencies installs this exact +# list at construction time, i.e. on every sandbox boot. ms-agent's +# own requirements cover six of them; without the other five each +# of the N sandboxes in a training step spends its first seconds +# fetching seaborn, scikit-learn, beautifulsoup4, lxml and pyarrow. +# Kept as the full list so it stays correct if ms-agent's +# dependencies shift. +# +# In a sandbox, any of these missing is either a network round trip at the start +# of every episode or an outright failure on an air-gapped host. +RUN pip install --no-cache-dir httpx ipykernel jupyter-client numpy pandas matplotlib seaborn scikit-learn requests beautifulsoup4 lxml pillow tqdm pyarrow + +ENV PYTHONUNBUFFERED=1 +WORKDIR /workspace diff --git a/cookbook/rsi/agentic/sandbox_server/install.sh b/cookbook/rsi/agentic/sandbox_server/install.sh new file mode 100644 index 000000000..6983d9a65 --- /dev/null +++ b/cookbook/rsi/agentic/sandbox_server/install.sh @@ -0,0 +1,83 @@ +#!/bin/sh +# Install AgentENV and build the sandbox template for agentic RSI. +# +# Usage: +# sh install.sh # install AgentENV + build the template +# sh install.sh --rebuild # delete the old template and rebuild +# sh install.sh --skip-install # template only, AgentENV already installed +set -eu + +TEMPLATE="${TEMPLATE:-twinkle-rsi-msagent}" +# ms-agent pulls in pandas/matplotlib/modelscope and notebook_executor starts a +# real ipykernel, so 1GiB is not enough. +CPU_COUNT="${CPU_COUNT:-2}" +MEMORY_MB="${MEMORY_MB:-2048}" +# Overrides the Dockerfile's `FROM` (passed to `aenv build --image`). Set this +# when the host cannot reach Docker Hub, e.g. +# BASE_IMAGE=docker.m.daocloud.io/library/python:3.11-slim +# daocloud is a third-party Docker Hub proxy, not an official Docker or Aliyun +# endpoint -- the base image of every sandbox would come through it. Prefer your +# own Aliyun accelerator address (<id>.mirror.aliyuncs.com) if you have one. +BASE_IMAGE="${BASE_IMAGE:-}" +# Where the runtime config is copied to, readable by the aenv user. serve.sh +# reads the same default. +REPO_ROOT="${REPO_ROOT:-$HOME/AgentENV}" +CONFIG_DIR="${CONFIG_DIR:-/var/lib/aenv/config}" + +SKIP_INSTALL=0 +REBUILD=0 +for arg in "$@"; do + case "$arg" in + --skip-install) SKIP_INSTALL=1 ;; + --rebuild) REBUILD=1 ;; + *) echo "Unknown option: $arg" >&2; exit 2 ;; + esac +done + +cd "$(dirname "$0")" + +if [ "$SKIP_INSTALL" = "0" ]; then + echo "==> Installing AgentENV server + aenv CLI" + curl -fsSL https://raw.githubusercontent.com/kvcache-ai/AgentENV/main/scripts/install.sh \ + | sudo bash + + echo "==> Provisioning the host (kvm group, ublk module, udev, sysctl)" + sudo server --setup-host + sudo install -d -o aenv -g aenv /var/lib/aenv/home + + # A source-built binary defaults to its build-time repo path for the config + # (CARGO_MANIFEST_DIR), which the aenv user cannot read when the repo lives + # under /root. Only default.toml needs copying — deps_manifest.toml is + # include_str!'d into the binary at compile time. + if [ -f "$REPO_ROOT/config/default.toml" ]; then + sudo install -d -o aenv -g aenv "$CONFIG_DIR" + sudo install -o aenv -g aenv -m 0644 \ + "$REPO_ROOT/config/default.toml" "$CONFIG_DIR/config.toml" + echo " config seeded to $CONFIG_DIR/config.toml" + fi +fi + +echo "==> Authenticating the CLI" +if [ -f "$HOME/.config/aenv/credentials" ]; then + echo " already authenticated ($HOME/.config/aenv/credentials)" +else + aenv auth +fi + +if [ "$REBUILD" = "1" ]; then + echo "==> Deleting template '$TEMPLATE'" + aenv template delete "$TEMPLATE" || true +fi + +echo "==> Building template '$TEMPLATE' (cpu=$CPU_COUNT mem=${MEMORY_MB}MiB)" +set -- Dockerfile -t "$TEMPLATE" --cpu-count "$CPU_COUNT" --memory-mb "$MEMORY_MB" +[ -n "$BASE_IMAGE" ] && set -- "$@" --image "$BASE_IMAGE" +aenv build "$@" + +echo +echo "Build runs server-side and takes a few minutes. Follow it with:" +echo " aenv template watch <template-id> # id printed above" +echo " aenv template list # confirm it reaches ready" +echo +echo "Then start the server:" +echo " sh serve.sh" diff --git a/cookbook/rsi/agentic/sandbox_server/serve.sh b/cookbook/rsi/agentic/sandbox_server/serve.sh new file mode 100644 index 000000000..e9043d8d1 --- /dev/null +++ b/cookbook/rsi/agentic/sandbox_server/serve.sh @@ -0,0 +1,112 @@ +#!/bin/sh +# Start the AgentENV server that hosts the RSI sandboxes. +# +# Usage: +# sh serve.sh # foreground, binds 127.0.0.1:8000 +# API_ADDR=0.0.0.0:8000 sh serve.sh # listen on all interfaces +# NOHUP=1 sh serve.sh # background, logs to /tmp/aenv-server.log +# STOP_ONLY=1 sh serve.sh # shut down without starting again +set -eu +REPO_ROOT="${REPO_ROOT:-$HOME/AgentENV}" +# Read by the server itself, not by this script. +export API_ADDR="${API_ADDR:-127.0.0.1:8000}" +LOG_FILE="${LOG_FILE:-/tmp/aenv-server.log}" +NOHUP="${NOHUP:-0}" + +# The server drops privileges to a non-root user, so it must not inherit root's +# HOME — regctl and docker credential lookups fail with EACCES there, which +# turns into a hard failure once a private registry needs credentials. +AENV_HOME="${AENV_HOME:-/var/lib/aenv/home}" + +# The binary bakes in its build-time repo path as the default config location +# (CARGO_MANIFEST_DIR in src/cfg.rs), so a server built under /root looks for +# /root/AgentENV/config/default.toml — unreadable once it drops to the aenv +# user, since /root is 0700. Point it at a copy the runtime user owns. +AENV_CONFIG_PATH="${AENV_CONFIG_PATH:-/var/lib/aenv/config/config.toml}" + +# run-with-capabilities.sh is primarily a test wrapper: when these are unset it +# defaults them to /tmp/aenv-test-<uid>/{home,run}. That sends downloaded +# dependencies (kernel, firecracker, overlaybd — hundreds of MB) to a directory +# that /tmp cleanup wipes, so every restart re-downloads them. Pin the real +# state directory instead; home_path in config.toml points at the same place. +AENV_HOME_PATH="${AENV_HOME_PATH:-/var/lib/aenv}" +AENV_RUNTIME_PATH="${AENV_RUNTIME_PATH:-/run/aenv}" + +if [ ! -r "$AENV_CONFIG_PATH" ]; then + echo "Config not readable: $AENV_CONFIG_PATH" >&2 + echo "Seed it from the repo (install.sh does this for you):" >&2 + echo " sudo install -d -o aenv -g aenv \$(dirname $AENV_CONFIG_PATH)" >&2 + echo " sudo install -o aenv -g aenv -m 0644 \\" >&2 + echo " $REPO_ROOT/config/default.toml $AENV_CONFIG_PATH" >&2 + exit 1 +fi + +# Stop whatever is already running, so this script is a restart rather than a +# "port already in use" failure. Match the binary path, not this script's name: +# run-with-capabilities.sh ends in `exec setpriv ... server`, which replaces the +# process image, so argv[0] of the live process is the server binary. +SERVER_BIN="${SERVER_BIN:-/usr/local/bin/server}" + +stop_running() { + # A systemd-managed instance would be restarted right after a kill, so hand + # it over to systemctl instead. install.sh sets up aenv.service when systemd + # is present. + if [ -d /run/systemd/system ] && systemctl is-active --quiet aenv 2>/dev/null; then + echo "Stopping systemd service aenv" + sudo systemctl stop aenv + return + fi + + pids=$(pgrep -f "^$SERVER_BIN" 2>/dev/null || true) + [ -z "$pids" ] && return + + echo "Stopping running server (pid: $pids)" + # SIGTERM first: the server tears down microVMs, veth pairs and iptables + # rules on shutdown, and SIGKILL would leave those behind. + sudo kill $pids 2>/dev/null || true + i=0 + while [ $i -lt 30 ] && pgrep -f "^$SERVER_BIN" >/dev/null 2>&1; do + sleep 1 + i=$((i + 1)) + done + if pgrep -f "^$SERVER_BIN" >/dev/null 2>&1; then + echo " still alive after 30s, sending SIGKILL" + sudo pkill -KILL -f "^$SERVER_BIN" 2>/dev/null || true + sleep 1 + fi +} + +stop_running + +if [ "${STOP_ONLY:-0}" = "1" ]; then + echo "Stopped." + exit 0 +fi + +cd "$REPO_ROOT" + +# run-with-capabilities.sh grants CAP_NET_ADMIN + CAP_SYS_ADMIN via setpriv and +# re-initialises supplementary groups (--init-groups), which is what makes a +# fresh kvm-group membership take effect without re-login. It derives repo_root +# from BASH_SOURCE, so the path above is what matters, not the cwd. +# +# `sudo env VAR=...`, not `sudo VAR=...`: with sudoers env_reset (the default) +# the latter is not guaranteed to pass anything through. +# +# AENV_RUN_USER must be explicit: the script otherwise falls back through +# SUDO_USER -> repo owner -> aenv -> root, and running as root is not supported. +E="AENV_RUN_USER=aenv HOME=$AENV_HOME API_ADDR=$API_ADDR AENV_CONFIG_PATH=$AENV_CONFIG_PATH AENV_HOME_PATH=$AENV_HOME_PATH AENV_RUNTIME_PATH=$AENV_RUNTIME_PATH" + +if [ "$NOHUP" = "1" ]; then + # setsid, not just nohup: the wrapper ends in `exec setpriv`, which replaces + # the process image, and a SIGHUP disposition inherited from nohup is not + # guaranteed to survive that. A new session detaches from the terminal + # regardless. + echo "Starting AgentENV on $API_ADDR (background) -> $LOG_FILE" + sudo env $E setsid nohup ./scripts/run-with-capabilities.sh server \ + >"$LOG_FILE" 2>&1 </dev/null & + echo "Tail with: tail -f $LOG_FILE" +else + echo "Starting AgentENV on $API_ADDR (foreground, Ctrl-C to stop)" + exec sudo env $E ./scripts/run-with-capabilities.sh server +fi diff --git a/cookbook/rl/rsi_agentic/sandbox_server/tool_server.py b/cookbook/rsi/agentic/sandbox_server/tool_server.py similarity index 64% rename from cookbook/rl/rsi_agentic/sandbox_server/tool_server.py rename to cookbook/rsi/agentic/sandbox_server/tool_server.py index c8ad53976..102da23c1 100644 --- a/cookbook/rl/rsi_agentic/sandbox_server/tool_server.py +++ b/cookbook/rsi/agentic/sandbox_server/tool_server.py @@ -41,15 +41,86 @@ # it is also removed from the advertised schema -- see `_usable_llm`. _LLM_BACKED_ARGS = {'file_system---read_file': ('abbreviate', )} +_SINGLE_NS_FLAG = '_twinkle_single_namespace' + + +def _single_namespace_source(code: str) -> str: + """Wrap ``code`` so it runs in one namespace and cannot exit the process. + + The inner ``exec`` passes one dict twice, which is what ordinary module + execution does, so nested scopes see top-level names; and ``SystemExit`` / + ``KeyboardInterrupt`` are turned into stderr text -- which is what ms-agent + reads as ``success: false`` -- instead of escaping into this server's event + loop. Stdout written before the exit survives, and ``sys.exit(0)`` stays a + success. ``repr`` handles the quoting, so the original source survives byte + for byte. + """ + return ('import builtins as _tw_builtins\n' + 'import sys as _tw_sys\n' + '_tw_src = ' + repr(code) + '\n' + "_tw_ns = {'__name__': '__main__', '__builtins__': _tw_builtins}\n" + 'try:\n' + " exec(compile(_tw_src, '<tool>', 'exec'), _tw_ns, _tw_ns)\n" + 'except (SystemExit, KeyboardInterrupt) as _tw_exit:\n' + " _tw_status = getattr(_tw_exit, 'code', 1)\n" + ' if _tw_status not in (0, None):\n' + " _tw_sys.stderr.write('%s: %s\\n' % (type(_tw_exit).__name__, _tw_status))\n") + + +def _patch_python_executor() -> bool: + """Give ms-agent's local ``python_executor`` ordinary module semantics. + + ``LocalCodeExecutionTool.python_executor`` calls + ``exec(code, globals_dict, locals_dict)`` with two *different* dicts + (ms_agent/tools/code/local_code_executor.py:670). Python then runs the code + the way it runs a class body: top-level assignments land in ``locals_dict``, + but every nested scope -- a function body, a generator expression -- + resolves free names against ``globals_dict`` alone. So:: + + import os + paths = ['a.txt'] + assert all(os.path.exists(p) for p in paths) + + raises ``NameError: name 'os' is not defined``, which reads as if the model + wrote broken code. Here it is worse than noise: check scripts arrive through + this tool and are the reward's ground truth, so a correct check scores as a + failure. + + The same method catches only ``Exception``, so ``sys.exit(3)`` in a script + raises ``SystemExit`` out of its ``asyncio.to_thread`` call. ``asyncio.Task`` + re-raises that one after storing it, so it unwinds ``run_forever`` and kills + :class:`_LoopThread` -- after which every later tool call in the sandbox + waits for a loop that is gone. Verified: without this, a ``sys.exit(3)`` + call is followed by timeouts on scripts that passed moments earlier. + + Duplicated from ``twinkle_agentic.harness.ms_agent`` on purpose -- this file + is uploaded into a sandbox that has ms-agent and nothing else. Temporary, + pending an upstream PR. + """ + from ms_agent.tools.code.local_code_executor import LocalCodeExecutionTool + + original = LocalCodeExecutionTool.python_executor + if getattr(original, _SINGLE_NS_FLAG, False): + return False + + async def python_executor(self, code, description='', timeout=None): + return await original(self, _single_namespace_source(code), + description=description, timeout=timeout) + + setattr(python_executor, _SINGLE_NS_FLAG, True) + LocalCodeExecutionTool.python_executor = python_executor + return True + def _usable_llm(cfg) -> bool: """Whether the declared ``llm`` section can actually serve a request. - ms-agent merges its own ``agent.yaml`` underneath the user's, and that - default declares ``service: modelscope``. So an absent ``llm:`` section in - rsi_agent.yaml does not mean "no LLM" -- it means "modelscope, with no - credentials", which asserts as soon as FileSystemTool is constructed. The - presence of a key is what decides it. + Call this on the config *after* ``LLMAgent`` construction. ms-agent merges + its own ``agent.yaml`` underneath the user's, and that default declares + ``service: modelscope``. So an absent ``llm:`` section in rsi_agent.yaml does + not mean "no LLM" -- it means "modelscope, with no credentials", which + asserts as soon as FileSystemTool is constructed. The presence of a key is + what decides it. """ llm = getattr(cfg, 'llm', None) if llm is None: @@ -59,6 +130,35 @@ def _usable_llm(cfg) -> bool: return any(getattr(llm, f, None) or os.environ.get(f.upper()) for f in key_fields) +def _to_openai(schema: Dict[str, Any]) -> Dict[str, Any]: + """Convert one ms-agent tool schema to the OpenAI shape. + + ``ToolManager.get_tools`` yields ms-agent's own flat form -- + ``{tool_name, server_name, description, parameters}`` -- but the schemas + served here go into the policy's prompt and are also what + ``RemoteMsAgentToolEnv.tool_names`` reads, and both speak OpenAI's nested + ``{type: function, function: {...}}``. Converting at this boundary keeps + ``/tools`` in the one shape every consumer expects. + + This mirrors ``twinkle_agentic.harness.ms_agent._ms_tools_to_openai``, which + cannot be imported: this file is uploaded into a sandbox that has ms-agent + and nothing else. + """ + if schema.get('type') == 'function' and isinstance(schema.get('function'), dict): + return schema + name = schema.get('tool_name') or schema.get('name') + if not name: + return schema + return { + 'type': 'function', + 'function': { + 'name': name, + 'description': schema.get('description', ''), + 'parameters': schema.get('parameters') or {'type': 'object', 'properties': {}}, + }, + } + + def _without_llm_args(schema: Dict[str, Any]) -> Dict[str, Any]: """Drop arguments this deployment cannot serve from a tool schema. @@ -108,6 +208,10 @@ def __init__(self, config_path: str, workspace: str) -> None: from ms_agent.agent.llm_agent import LLMAgent + # Before any tool is constructed: the patch replaces a method on + # LocalCodeExecutionTool, and prepare_tools() instantiates it. + _patch_python_executor() + cfg = OmegaConf.load(config_path) with open_dict(cfg): cfg.output_dir = workspace @@ -116,13 +220,21 @@ def __init__(self, config_path: str, workspace: str) -> None: # blocking on stdin would hang the episode until the sandbox timeout. cfg.interactive = False cfg.permission_mode = 'auto' - self.has_llm = _usable_llm(cfg) + self.agent = LLMAgent(cfg) + # The llm decision has to be made on the *merged* config, after LLMAgent + # has layered ms-agent's own agent.yaml underneath ours. Popping `llm` + # from the pre-merge config only removes our section and lets the + # default's `service: modelscope` show through, which asserts on the + # missing key as soon as FileSystemTool is constructed. This mirrors + # MsAgentHarness._apply_rl_stubs, which mutates agent.config for the same + # reason. + with open_dict(self.agent.config): + self.has_llm = _usable_llm(self.agent.config) if not self.has_llm: # Leaving an unusable section in place is not an option: # FileSystemTool builds a client from it eagerly and asserts on # the missing key, so no tool at all would come up. - cfg.pop('llm', None) - self.agent = LLMAgent(cfg) + self.agent.config.pop('llm', None) self.agent._interactive = False self.agent._event_sink = None self.agent._input_source = None @@ -151,7 +263,8 @@ def tools(self) -> List[Dict[str, Any]]: flat.extend(value if isinstance(value, list) else [value]) else: flat = list(raw or []) - return [t if self.has_llm else _without_llm_args(t) for t in flat if isinstance(t, dict)] + schemas = [_to_openai(t) for t in flat if isinstance(t, dict)] + return [t if self.has_llm else _without_llm_args(t) for t in schemas] def call(self, calls: List[Dict[str, Any]], timeout: Optional[float]) -> List[Dict[str, Any]]: """Dispatch a turn's tool calls, mirroring how ms-agent itself does it. @@ -250,7 +363,7 @@ def main() -> None: # Threading, because a turn's tool calls arrive as one request but the # health poll must stay answerable while a long shell command runs. server = ThreadingHTTPServer((args.host, args.port), _Handler) - names = [t.get('function', {}).get('name') for t in runtime.tools()] + names = [(t.get('function') or {}).get('name') for t in runtime.tools()] llm_note = 'llm configured' if runtime.has_llm else 'no llm (read_file.abbreviate withdrawn)' sys.stderr.write(f'[tool_server] ready on {args.host}:{args.port}, {llm_note}, ' f'{len(names)} tools: {names}\n') diff --git a/cookbook/rl/rsi_agentic/tasks.example.jsonl b/cookbook/rsi/agentic/tasks.example.jsonl similarity index 100% rename from cookbook/rl/rsi_agentic/tasks.example.jsonl rename to cookbook/rsi/agentic/tasks.example.jsonl diff --git a/src/twinkle_agentic/challenger/agentic.py b/src/twinkle_agentic/challenger/agentic.py index d54c5ed7e..b28aad6f3 100644 --- a/src/twinkle_agentic/challenger/agentic.py +++ b/src/twinkle_agentic/challenger/agentic.py @@ -102,6 +102,30 @@ def _trajectory_summary(trajectory: Trajectory) -> str: return '\n'.join(parts) +# The fields a local rollout splices into a trajectory, and the only ones a +# later GRPO step needs: ``labels`` marks which of ``input_ids`` are trainable +# (-100 elsewhere) and ``logprobs`` holds one entry per trainable token, taken +# from the policy that actually generated it. +_TRAINABLE_KEYS = ('input_ids', 'labels', 'logprobs') + + +def _propose_round(stage: str, trajectory: Trajectory) -> Dict[str, Any]: + """One proposing round, reduced to what a later training step would read. + + ``messages`` comes along for reading by humans; it is redundant with + ``input_ids`` and is not what a trainer should encode from. + """ + record: Dict[str, Any] = { + 'stage': stage, + 'messages': [dict(m) for m in trajectory.get('messages') or []], + } + for key in _TRAINABLE_KEYS: + value = trajectory.get(key) + if value is not None: + record[key] = value + return record + + # ── prompts ──────────────────────────────────────────────────────────────── @dataclass @@ -192,6 +216,16 @@ class AgenticChallenger(Challenger): min_batch: smallest batch worth sending to the explorer. problem_max_chars: reject problem statements longer than this. reject_sink: called with a dict for every rejected proposal. + propose_sink: called once per proposal attempt -- kept, rejected while + building, or dropped by the difficulty band alike -- with the + token-level record of the rounds that produced it. This is the only + way the proposing rounds survive: they are generation like any + other, so they carry ``input_ids`` / ``labels`` / ``logprobs`` and + could later be trained on, but nothing downstream of ``build`` + looks at them and without a sink they are dropped on the floor. + Rejects are included on purpose: they are the zero-reward half of a + GRPO group, so a set of kept-only records has no variance to learn + from. Requires a local sampler -- an API explorer returns text only. """ def __init__( @@ -216,6 +250,7 @@ def __init__( min_batch: int = 1, problem_max_chars: int = 8192, reject_sink: Optional[Callable[[Dict[str, Any]], None]] = None, + propose_sink: Optional[Callable[[Dict[str, Any]], None]] = None, **challenger_kwargs: Any, ): super().__init__(explorer, system=prompts.system, **challenger_kwargs) @@ -246,6 +281,7 @@ def __init__( self.min_batch = max(1, min_batch) self.problem_max_chars = problem_max_chars self.reject_sink = reject_sink + self.propose_sink = propose_sink if self.seeds: prompts.require('from_seed') if self.store is not None: @@ -309,6 +345,12 @@ def _build_one(self, explored: Trajectory) -> Optional[Trajectory]: """ summary = _trajectory_summary(explored) snapshot = self.workspace_snapshot_fn() if self.workspace_snapshot_fn else summary + keywords = user_data_get(explored.get('user_data'), 'keywords', []) + seeded = user_data_get(explored.get('user_data'), 'seeded', False) + # Every round this proposal generates, in order. Handed to propose_sink + # with whatever verdict the proposal ends up with, so a rejected attempt + # is recorded as fully as a kept one. + rounds = [_propose_round('explore', explored)] # Round 2a: write check script # NOTE: This goes through the same explorer (with tool schemas visible). @@ -323,10 +365,12 @@ def _build_one(self, explored: Trajectory) -> Optional[Trajectory]: ], } check_reply = self.explore([check_prompt]) + rounds.append(_propose_round('check', check_reply[0])) script = parse_check_script(assistant_text(check_reply[0])) if script is None: self.stats['check_parse_fail'] += 1 self._reject_record(explored, 'check_parse_fail') + self._emit_propose(rounds, 'check_parse_fail', keywords=keywords, seeded=seeded) return None # Verify: run check script in current sandbox state (must pass) @@ -335,6 +379,7 @@ def _build_one(self, explored: Trajectory) -> Optional[Trajectory]: self.stats['check_run_fail'] += 1 self._reject_record(explored, 'check_run_fail', detail=f'exit {exit_code}: {output[-200:]}') + self._emit_propose(rounds, 'check_run_fail', keywords=keywords, seeded=seeded) return None # Round 2b: write problem statement @@ -346,25 +391,29 @@ def _build_one(self, explored: Trajectory) -> Optional[Trajectory]: ], } problem_reply = self.explore([problem_prompt]) + rounds.append(_propose_round('problem', problem_reply[0])) statement = parse_problem_statement(assistant_text(problem_reply[0])) if statement is None: self.stats['problem_parse_fail'] += 1 self._reject_record(explored, 'problem_parse_fail') + self._emit_propose(rounds, 'problem_parse_fail', keywords=keywords, seeded=seeded) return None if len(statement) > self.problem_max_chars: self.stats['too_long'] += 1 self._reject_record(explored, 'too_long') + self._emit_propose(rounds, 'too_long', keywords=keywords, seeded=seeded) return None self.stats['parsed'] += 1 task: Trajectory = { 'messages': [{'role': 'user', 'content': statement}], } - return attach_user_data( - task, - check_script=script, - keywords=user_data_get(explored.get('user_data'), 'keywords', []), - seeded=user_data_get(explored.get('user_data'), 'seeded', False)) + task = attach_user_data(task, check_script=script, keywords=keywords, seeded=seeded) + # Carried, not emitted: the verdict this proposal earns depends on the + # difficulty measurement, which has not run yet. A plain top-level key + # rather than user_data, which json-encodes every value on each update. + task['propose_rounds'] = rounds + return task def _reject_record(self, traj: Trajectory, reason: str, detail: str = '') -> None: if self.reject_sink is not None: @@ -374,6 +423,37 @@ def _reject_record(self, traj: Trajectory, reason: str, detail: str = '') -> Non payload['last_assistant'] = assistant_text(traj)[:500] self.reject_sink(payload) + def _emit_propose(self, rounds: Optional[List[Dict[str, Any]]], outcome: str, *, + keywords: Any = (), seeded: bool = False, + n_pass: Optional[int] = None) -> None: + """Hand one proposal attempt's rounds to ``propose_sink``. + + ``pass_rate`` is the raw fraction of solver attempts that succeeded. It + is left as the measurement rather than mapped onto a difficulty score: + the target rate and its tolerance are training decisions, and baking a + guess at them into the dump would make it look like they had been + settled. + """ + if self.propose_sink is None or not rounds: + return + rollouts = self.solver_rollouts or None + self.propose_sink({ + 'outcome': outcome, + 'n_pass': n_pass, + 'n_rollouts': rollouts, + 'pass_rate': (n_pass / rollouts) if (n_pass is not None and rollouts) else None, + 'keywords': list(keywords or ()), + 'seeded': bool(seeded), + 'rounds': rounds, + }) + + def _take_rounds(self, task: Trajectory) -> Optional[List[Dict[str, Any]]]: + """Detach a task's proposing rounds. Popped even with no sink attached: + token ids for a whole agentic episode are large, and a kept task is held + until the caller's batch is full. + """ + return task.pop('propose_rounds', None) + # ------------------------------------------------------------ revised _round def _round(self, missing: int) -> Optional[List[Trajectory]]: @@ -398,6 +478,12 @@ def _round(self, missing: int) -> Optional[List[Trajectory]]: usable.append(task) kept = self._filter_difficulty(usable) if self.solver_rollouts else usable + if not self.solver_rollouts: + # No difficulty stage, so the verdict is final as soon as it is built. + for task in usable: + self._emit_propose(self._take_rounds(task), 'kept', + keywords=user_data_get(task.get('user_data'), 'keywords', []), + seeded=user_data_get(task.get('user_data'), 'seeded', False)) self.n_proposed += len(proposals) self.n_kept += len(kept) band = (f', in difficulty band {len(kept)}' if self.solver_rollouts else '') @@ -427,7 +513,18 @@ def _filter_difficulty(self, tasks: List[Trajectory]) -> List[Trajectory]: ] self.on_difficulty_measured(measured) high = self.solver_rollouts - self.keep_max_pass_margin - return [t for t, n in zip(measured, passes) if self.keep_min_pass <= n <= high] + in_band = [self.keep_min_pass <= n <= high for n in passes] + # Emit here, not in _round: this is where a proposal's verdict is + # decided, and both sides of the band are worth keeping -- a task nobody + # solved and one everybody solved are the two failure modes the + # proposer would need to learn to avoid. + for task, n, kept_flag in zip(measured, passes, in_band): + self._emit_propose(self._take_rounds(task), + 'kept' if kept_flag else 'outside_band', + keywords=user_data_get(task.get('user_data'), 'keywords', []), + seeded=user_data_get(task.get('user_data'), 'seeded', False), + n_pass=n) + return [t for t, kept_flag in zip(measured, in_band) if kept_flag] def solver_prompt(self, task: Trajectory) -> Trajectory: """The task statement, nothing else -- the solver's own harness adds the system.""" diff --git a/src/twinkle_agentic/harness/ms_agent.py b/src/twinkle_agentic/harness/ms_agent.py index 948c82861..d3f50eb92 100644 --- a/src/twinkle_agentic/harness/ms_agent.py +++ b/src/twinkle_agentic/harness/ms_agent.py @@ -211,6 +211,7 @@ def after_tools( def _apply_rl_stubs(self) -> None: """Non-interactive: never block on TUI / permission prompts / stdin.""" + patch_ms_agent_python_executor() try: from omegaconf import open_dict with open_dict(self.agent.config): @@ -384,6 +385,93 @@ def _ms_calls_to_openai(tool_calls: List[Any]) -> List[Dict[str, Any]]: return out +_SINGLE_NS_FLAG = '_twinkle_single_namespace' + + +def single_namespace_source(code: str) -> str: + """Wrap ``code`` so it runs in one namespace and cannot exit the process. + + Two things happen here, both of them repairs (see + :func:`patch_ms_agent_python_executor`): + + * the inner ``exec`` passes one dict twice, which is what ordinary module + execution does, so nested scopes see top-level names; + * ``SystemExit`` / ``KeyboardInterrupt`` are caught and turned into stderr + output, so a ``sys.exit(3)`` in a script fails that one call instead of + escaping into the caller's event loop. + + A non-zero status is reported the way any other failure is -- text on + stderr, which is what ms-agent turns into ``success: false`` -- so stdout + written before the exit survives. ``sys.exit()`` and ``sys.exit(0)`` stay + successes: that is a script saying it is done. + + The wrapper only assigns and reads at top level, which works under split + globals/locals. ``repr`` handles all quoting, so the original source + survives byte for byte. + """ + return ('import builtins as _tw_builtins\n' + 'import sys as _tw_sys\n' + '_tw_src = ' + repr(code) + '\n' + "_tw_ns = {'__name__': '__main__', '__builtins__': _tw_builtins}\n" + 'try:\n' + " exec(compile(_tw_src, '<tool>', 'exec'), _tw_ns, _tw_ns)\n" + 'except (SystemExit, KeyboardInterrupt) as _tw_exit:\n' + " _tw_status = getattr(_tw_exit, 'code', 1)\n" + ' if _tw_status not in (0, None):\n' + " _tw_sys.stderr.write('%s: %s\\n' % (type(_tw_exit).__name__, _tw_status))\n") + + +def patch_ms_agent_python_executor() -> bool: + """Give ms-agent's local ``python_executor`` ordinary module semantics. + + ``LocalCodeExecutionTool.python_executor`` calls + ``exec(code, globals_dict, locals_dict)`` with two *different* dicts + (ms_agent/tools/code/local_code_executor.py:670). Python then runs the + submitted code the way it runs a class body: top-level assignments land in + ``locals_dict``, but every nested scope -- a function body, a generator + expression -- resolves free names against ``globals_dict`` alone. So:: + + import os + paths = ['a.txt'] + assert all(os.path.exists(p) for p in paths) + + raises ``NameError: name 'os' is not defined``, which reads as if the model + wrote broken code. For RSI that is worse than noise: the check script *is* + the reward's ground truth, so this scores a correct check as a failure. + + The same method catches only ``Exception`` around the ``exec``, so a script + calling ``sys.exit(3)`` raises ``SystemExit`` out of the ``asyncio.to_thread`` + call. ``asyncio.Task`` re-raises that one after storing it, which unwinds + whatever loop is driving the tool: with a long-lived loop (the RSI sandbox + server keeps one, so notebook and MCP state survive across turns) the loop + thread dies and every later tool call in the run hangs. One model-written + ``sys.exit`` would take out the rest of the episode. + + Temporary local fix pending an upstream PR. It wraps the source instead of + reimplementing the method, so ms-agent keeps owning timeouts, output capture + and the JSON result shape. + + Idempotent. Returns True when it patched, False when ms-agent is missing or + the patch is already in place. + """ + try: + from ms_agent.tools.code.local_code_executor import LocalCodeExecutionTool + except Exception: # noqa -- ms-agent is optional for most of twinkle + return False + + original = LocalCodeExecutionTool.python_executor + if getattr(original, _SINGLE_NS_FLAG, False): + return False + + async def python_executor(self, code: str, description: str = '', timeout=None): + return await original(self, single_namespace_source(code), + description=description, timeout=timeout) + + setattr(python_executor, _SINGLE_NS_FLAG, True) + LocalCodeExecutionTool.python_executor = python_executor + return True + + def _ms_tools_to_openai(raw: Union[Dict[str, Any], List[Any], None]) -> List[Dict[str, Any]]: if not raw: return [] diff --git a/tests/twinkle_agentic/test_agentic_rsi.py b/tests/twinkle_agentic/test_agentic_rsi.py index 1a1d6f2ba..42409a7dd 100644 --- a/tests/twinkle_agentic/test_agentic_rsi.py +++ b/tests/twinkle_agentic/test_agentic_rsi.py @@ -18,7 +18,7 @@ sys.path.insert(0, os.path.join(_REPO, 'src')) # The RSI wiring lives in cookbook, not in the framework: it is one deployment's # choice of sandbox backend, and the tests follow it there. -_COOKBOOK = os.path.join(_REPO, 'cookbook', 'rl', 'rsi_agentic') +_COOKBOOK = os.path.join(_REPO, 'cookbook', 'rsi', 'agentic') sys.path.insert(0, _COOKBOOK) sys.path.insert(0, os.path.join(_COOKBOOK, 'sandbox_server')) diff --git a/tests/twinkle_agentic/test_harness.py b/tests/twinkle_agentic/test_harness.py index d1b22571d..93a3b9c21 100644 --- a/tests/twinkle_agentic/test_harness.py +++ b/tests/twinkle_agentic/test_harness.py @@ -152,6 +152,45 @@ def test_tool_manager_call_many_uses_env_step_batch(): assert out[1].startswith('lookup:') +def _run_wrapped(source: str): + """exec the wrapper the way ms-agent's python_executor does: split dicts.""" + import io + from contextlib import redirect_stderr, redirect_stdout + + from twinkle_agentic.harness.ms_agent import single_namespace_source + + out, err = io.StringIO(), io.StringIO() + with redirect_stdout(out), redirect_stderr(err): + exec(single_namespace_source(source), {'__builtins__': __builtins__}, {}) + return out.getvalue(), err.getvalue() + + +@pytest.mark.parametrize( + 'source, expect_err', + [ + # A comprehension seeing a top-level name: broken under split dicts. + ('xs = [1, 2]\nlim = 3\nassert all(x <= lim for x in xs)\nprint("ok")', False), + ('import os\npaths = []\nassert all(os.path.exists(p) for p in paths)\nprint("ok")', False), + # sys.exit must fail this call only -- never reach the caller's loop. + ('print("ok")\nimport sys\nsys.exit(3)', True), + ('print("ok")\nimport sys\nsys.exit(0)', False), + ('print("ok")\nimport sys\nsys.exit()', False), + ], +) +def test_single_namespace_source(source, expect_err): + stdout, stderr = _run_wrapped(source) + assert 'ok' in stdout + assert bool(stderr) is expect_err + if expect_err: + assert 'SystemExit: 3' in stderr + + +def test_single_namespace_source_keeps_real_errors(): + """The patch must not turn a failing check into a passing one.""" + with pytest.raises(AssertionError): + _run_wrapped('assert 1 == 2, "counts differ"') + + def test_ms_agent_harness_start_system_and_user(): import sys from pathlib import Path From 063bafa7ae31258335cd1ba363660f1f7557f681 Mon Sep 17 00:00:00 2001 From: root <yuze.zyz@alibaba-inc.com> Date: Wed, 26 Aug 2026 00:24:54 +0800 Subject: [PATCH 44/60] wip --- cookbook/rsi/agentic/README.md | 39 +- cookbook/rsi/agentic/challenge.py | 887 ++++++++++-- cookbook/rsi/agentic/episode.py | 234 +++ cookbook/rsi/agentic/eval.py | 244 ++++ cookbook/rsi/agentic/prompts.py | 304 ++-- cookbook/rsi/agentic/remote_tool_env.py | 343 ++++- cookbook/rsi/agentic/rl.py | 216 +-- cookbook/rsi/agentic/rsi_agent.yaml | 127 +- .../rsi/agentic/sandbox_server/Dockerfile | 53 +- .../sandbox_server/build_via_sandbox.sh | 126 ++ .../rsi/agentic/sandbox_server/install.sh | 5 + .../rsi/agentic/sandbox_server/reap_paused.py | 72 + cookbook/rsi/agentic/sandbox_server/serve.sh | 10 +- .../rsi/agentic/sandbox_server/tool_server.py | 296 +++- cookbook/rsi/agentic/split_tasks.py | 70 + docs/source_en/Components/Agentic/Rollout.md | 2 +- .../Agentic/Rollout.md" | 2 +- src/twinkle/data_format/sampling.py | 11 + src/twinkle/template/tools/bracket_dsl.py | 70 +- src/twinkle_agentic/challenger/agentic.py | 1280 ++++++++++++++--- src/twinkle_agentic/challenger/base.py | 23 +- src/twinkle_agentic/harness/ms_agent.py | 15 + src/twinkle_agentic/rollout/api_multi_turn.py | 55 +- src/twinkle_agentic/rollout/bridge.py | 49 +- src/twinkle_agentic/rollout/multi_turn.py | 302 +++- src/twinkle_agentic/tools/tool_manager.py | 66 +- src/twinkle_client/rollout/multi_turn.py | 8 + tests/template/test_tool_call_parsers.py | 86 ++ tests/twinkle_agentic/test_agentic_rsi.py | 1008 ++++++++++++- tests/twinkle_agentic/test_harness.py | 34 + .../test_multi_turn_rollout.py | 441 +++++- tests/twinkle_agentic/test_tools.py | 29 + .../test_client_multi_turn_rollout.py | 25 + 33 files changed, 5842 insertions(+), 690 deletions(-) create mode 100644 cookbook/rsi/agentic/episode.py create mode 100644 cookbook/rsi/agentic/eval.py create mode 100644 cookbook/rsi/agentic/sandbox_server/build_via_sandbox.sh create mode 100644 cookbook/rsi/agentic/sandbox_server/reap_paused.py create mode 100644 cookbook/rsi/agentic/split_tasks.py create mode 100644 tests/template/test_tool_call_parsers.py diff --git a/cookbook/rsi/agentic/README.md b/cookbook/rsi/agentic/README.md index 440dc3e9d..d6542c1ff 100644 --- a/cookbook/rsi/agentic/README.md +++ b/cookbook/rsi/agentic/README.md @@ -84,6 +84,40 @@ BASE_IMAGE=docker.m.daocloud.io/library/python:3.11-slim sh sandbox_server/insta image comes through it. Substitute your own Aliyun accelerator address (`<id>.mirror.aliyuncs.com`) if you would rather not depend on one. +#### When the template build is too slow to use + +On 2026-08-23 three `aenv build` runs on our host failed or ran for hours, and +the cause was download speed rather than the Dockerfile. Measured within one +minute, from inside a sandbox: `deb.debian.org` 33 KB/s, `mirrors.aliyun.com` +5.4 MB/s, and for comparison the host itself 12 MB/s and sandbox disk writes +639 MB/s. apt's package index alone is 9.6MB, so the build sat two hours with no +output -- and since the server logs `template build started` and then nothing +until the build ends, slow is indistinguishable from hung. The Dockerfile now +rewrites `deb.debian.org` to the Aliyun mirror, which should remove the cause. + +The path that is verified end to end installs inside a live sandbox and +snapshots it, which needs no template builder and takes about six minutes: + +```bash +sh sandbox_server/build_via_sandbox.sh # name: twinkle-rsi-msagent +NAME=twinkle-rsi-msagent-v2 sh sandbox_server/build_via_sandbox.sh # verify first +``` + +Three things to know about a snapshot: + +* it shows up in `aenv snapshot list`, **not** `aenv template list`, but the name + lives in the same namespace -- `--sandbox-template twinkle-rsi-msagent` + resolves to it unchanged; +* it keeps the filesystem, not the image config, so the Dockerfile's `ENV + PYTHONUNBUFFERED=1`, `ENV PIP_INDEX_URL=...` and `WORKDIR /workspace` are gone. + `build_via_sandbox.sh` writes `/etc/pip.conf` and `/workspace` instead, and + `remote_tool_env` starts the runtime with `python -u`; +* aenv refuses to rebind an existing name, so replacing an image means deleting + the old one first. Build under a second name and verify before you do that -- + deleting first cost us four hours with no usable sandbox. + +Verify either one from the training host with the boot check below. + The Dockerfile also pins `PIP_INDEX_URL` to an Aliyun mirror for the same reason -- edit those two lines if your host reaches pypi.org directly. @@ -165,7 +199,7 @@ ssh -N -L 8000:127.0.0.1:8000 root@<env-host-ip> | Variable | Default | | |---|---|---| | `AENV_API_URL` | `http://127.0.0.1:8000` | AgentENV server | -| `AENV_TEMPLATE` | `twinkle-rsi-msagent` | template built by `install.sh` | +| `AENV_TEMPLATE` | `twinkle-rsi-msagent` | template or snapshot name to boot from | | `RSI_TASKS` | `tasks.example.jsonl` | task file | | `RSI_AGENT_CONFIG` | `rsi_agent.yaml` | uploaded into every sandbox | | `RSI_SANDBOX_TIMEOUT` | `900` | must outlast an episode plus its checks | @@ -191,8 +225,9 @@ accordingly. | `remote_tool_env.py` | training-side Env: forwards tool calls, copies files back | | `rsi_agent.yaml` | ms-agent config -- read by *both* halves | | `sandbox_server/tool_server.py` | in-sandbox HTTP server owning the ToolManager | -| `sandbox_server/Dockerfile` | template image: ms-agent, ripgrep, ipykernel | +| `sandbox_server/Dockerfile` | template image: ms-agent, ripgrep, ffmpeg, imagemagick, openpyxl/reportlab/pdfplumber | | `sandbox_server/install.sh` | install AgentENV + build the template | +| `sandbox_server/build_via_sandbox.sh` | build the image as a snapshot of a live sandbox instead | | `sandbox_server/serve.sh` | start the AgentENV server | | `tasks.example.jsonl` | hand-written tasks in the structured `checks` format | diff --git a/cookbook/rsi/agentic/challenge.py b/cookbook/rsi/agentic/challenge.py index 650f279c6..5bfbce3f9 100644 --- a/cookbook/rsi/agentic/challenge.py +++ b/cookbook/rsi/agentic/challenge.py @@ -1,8 +1,9 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """RSI self-play, agentic half: generate training tasks by doing them first. -One model plays both roles. It first acts as an agent in a sandbox (multi-turn -with tools), producing a trajectory and final workspace state. Then it writes a +One model plays both roles, and one proposal is one conversation. It first acts +as an agent in a sandbox (multi-turn with tools), producing a trajectory and a +final workspace state; then, appended to that same conversation, it writes a check script that verifies the end state, and finally describes the task as a problem statement. The same model then attempts the problem multiple times, and only problems it solves *sometimes* are kept. @@ -21,6 +22,9 @@ """ import argparse import json +import base64 +import binascii +import hashlib import os import sys @@ -35,11 +39,157 @@ from twinkle_agentic.tools.tool_manager import ToolManager sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from episode import solver_harness # noqa: E402 from prompts import CATEGORIES, CATEGORY_DESC, agentic_prompts # noqa: E402 -from remote_tool_env import RemoteMsAgentToolEnv # noqa: E402 +from remote_tool_env import RemoteMsAgentToolEnv, tool_payload # noqa: E402 logger = get_logger() +# Cleared through the python tool, not `rm -rf`: ms-agent's safety policy rejects +# `rm -rf` outright ("Blocked by safety rule"), and it rejects globs in write +# operations, which rules out `find -delete` too. The script asserts the +# directory really is empty, so a future policy change surfaces as a failed reset +# instead of tasks quietly inheriting the previous workspace. +# +# Module level so a test can drive the same string the run does; a copy in a test +# would keep passing after this one changed. +CLEAR_WORKSPACE = ''' +import os, shutil +root = {workspace!r} +os.makedirs(root, exist_ok=True) +for name in os.listdir(root): + path = os.path.join(root, name) + if os.path.isdir(path) and not os.path.islink(path): + shutil.rmtree(path, ignore_errors=True) + else: + os.remove(path) +leftover = os.listdir(root) +assert not leftover, 'workspace not empty after clear: %r' % (leftover,) +''' + +# ── Arm B: copy the episode's input files out, and put them back later ────── +# +# The bytes travel, not a description of them: the model is not asked to write a +# script that recreates its own inputs, because a wrong one costs the whole +# episode and would fail exactly where the inputs are least ordinary. +# +# Read in slices because the executor truncates its output at roughly 8 KB. Each +# slice is base64 so any byte survives the trip, and the manifest's sha256 is +# what says the trip was faithful -- checked locally against the bytes that +# arrived, so a truncated read cannot pass as a smaller file. +INPUT_MANIFEST = ''' +import hashlib, os +root = os.path.join({workspace!r}, 'input') +rows = [] +for dirpath, dirnames, filenames in os.walk(root): + dirnames[:] = [d for d in dirnames if d not in {{'__pycache__', '.ipynb_checkpoints'}}] + for name in sorted(filenames): + path = os.path.join(dirpath, name) + with open(path, 'rb') as handle: + body = handle.read() + rows.append((os.path.relpath(path, {workspace!r}), len(body), + hashlib.sha256(body).hexdigest())) +for rel, size, digest in sorted(rows): + print(rel, size, digest) +''' + +INPUT_SLICE = ''' +import base64, os +path = os.path.join({workspace!r}, {rel!r}) +with open(path, 'rb') as handle: + handle.seek({offset}) + print(base64.b64encode(handle.read({length})).decode()) +''' + +# What gets stored with the task and run before every attempt at it. Writes the +# captured bytes and nothing else: no cleanup, because whoever runs this has just +# cleared the workspace. +SETUP_SCRIPT_TEMPLATE = '''# Recreate the task's input files. +import base64, os, pathlib + +FILES = {files!r} + +for rel, payload in FILES.items(): + path = pathlib.Path(rel) + if path.parent != pathlib.Path('.'): + os.makedirs(path.parent, exist_ok=True) + with open(path, 'wb') as handle: + handle.write(base64.b64decode(payload)) +''' + +# The ground truth the check script is written against. A listing alone is not +# enough: three of the six rejected proposals in the first real run failed on a +# value the model recomputed from its own recollection ("Mean values mismatch") +# rather than read off the file, so the end state has to arrive as content, not +# just as names. Bounded on both axes -- 50 files, 600 bytes each, 6000 overall -- +# because this goes into a prompt and a 100k artifact would push the trajectory +# it has to be read alongside out of the window. +# +# Walks the tree in python rather than shelling out to `find`: the same code then +# decides what is text, what is truncated, and what the budget was spent on, +# which a pipeline of find/head cannot report back. +# +# File bodies go out byte for byte. An earlier version printed `body.rstrip()`, +# which hid trailing newlines while the size column still counted them, so a +# check writer shown an 11-byte file whose content looked 10 characters long +# wrote `content == 'Mean: 63.9'` and the check failed against the very state it +# was written from. The listing is only ground truth if it does not tidy up. +# +# Facts *about* a file go in its header, never after its body. A note printed +# below the content is indistinguishable from content: annotated one file with a +# trailing `(no newline at end of file)` line and the next check script asserted +# the README's content ending in that sentence. +WORKSPACE_SNAPSHOT = ''' +import os + +root = {workspace!r} +skip = {{'.ms_agent', '__pycache__', '.ipynb_checkpoints', '.git'}} +rows = [] +for dirpath, dirnames, filenames in os.walk(root): + dirnames[:] = [d for d in dirnames if d not in skip] + for name in sorted(filenames): + path = os.path.join(dirpath, name) + try: + rows.append((os.path.relpath(path, root), os.path.getsize(path), path)) + except OSError: + pass +rows.sort() +for rel, size, _ in rows[:{max_files}]: + print(rel, size) + +budget = {total_budget} +for rel, size, path in rows[:{max_files}]: + if budget <= 0: + break + try: + with open(path, encoding='utf-8') as handle: + text = handle.read({per_file} + 1) + except (OSError, UnicodeDecodeError): + continue # binary or unreadable: the listing already names it + if '\\x00' in text: + continue + body = text[:{per_file}] + budget -= len(body) + # The trailing-newline count is stated for every file, both ways. Saying it + # only when it is absent made "this file ends with a newline" invisible, and + # the check writer then compared exact bytes without one: in ex9 two of the + # three checks that failed their own verification failed on exactly that -- + # the same reply asserted three files, guessed right on the two marked "no + # newline at end" and wrong on the unmarked one. + trailing = len(body) - len(body.rstrip(chr(10))) + if len(text) > len(body): + suffix = ' (first {per_file} bytes)' + elif trailing == 0: + suffix = ' (no newline at end)' + else: + suffix = ' (ends with %d newline character(s))' % trailing + print() + print('--- ' + rel + suffix + ' ---') + print(body, end='') + if not body.endswith(chr(10)): + print() +''' + def parse_args(): p = argparse.ArgumentParser(description=__doc__, @@ -49,7 +199,59 @@ def parse_args(): p.add_argument('--template', default='Template', help='template class in twinkle.template') p.add_argument('--sampler-gpus', type=int, default=4) - p.add_argument('--max-model-len', type=int, default=32768) + # Challenger backend: local vLLM by default, or an OpenAI-compatible API for + # the proposing side only (keywords + explore + check + statement). The + # solver side of the difficulty stage still runs through the local sampler, + # so --challenger-api is only usable with --solver-rollouts 0 (no local + # sampler is built at all in that case). The three connection args default + # to the LLM_BACKUP_* env vars the summarizer teacher already uses. + p.add_argument('--challenger-api', action='store_true', + help='propose through an OpenAI-compatible API instead of local vLLM; ' + 'requires --solver-rollouts 0') + p.add_argument('--challenger-api-model', + default=os.environ.get('LLM_BACKUP_MODEL', '')) + p.add_argument('--challenger-api-base', + default=os.environ.get('LLM_BACKUP_BASE_URL', '')) + p.add_argument('--challenger-api-key', + default=os.environ.get('LLM_BACKUP_API_KEY', '')) + p.add_argument('--challenger-concurrency', type=int, default=8, + help='parallel API conversations (API backend only)') + # Measured on qwen3.8-max at a ~15k-character exploration context: one turn + # took 58s with the default (unbounded) thinking and 10s at 2048, because the + # default spent ~5300 characters on reasoning per turn. 0 leaves the API's own + # default in place; anything else is sent as extra_body={'thinking_budget': N} + # on every proposing call (explore, check, statement, keyword generation). + p.add_argument('--challenger-thinking-budget', type=int, default=0, + help='cap reasoning tokens per API call; 0 = leave the API default') + # Split mode: explore on the local (trainable) vLLM model, but write the check + # script (success judgement) and the problem statement over the API instead of + # the local model. Only the exploration turns keep labels/logprobs and get + # trained; the two API stages are text-only and never enter the trajectory. + # Reuses the --challenger-api-* connection args, and is mutually exclusive with + # --challenger-api (which sends the whole proposing side to the API). + p.add_argument('--followup-api', action='store_true', + help='explore locally (trainable) but generate the check script and ' + 'problem statement over --challenger-api-* (e.g. qwen3-max); ' + 'only the exploration part is trained. Not with --challenger-api.') + # How many episodes run at once, each in its own sandbox. An episode owns its + # workspace from the reset until its check has run, so this is also the number + # of sandboxes booted at startup. Default 48: episodes alternate between vLLM + # generation (~24 concurrent sequences fit in KV cache) and sandbox execution, + # so 48 keeps both the GPU cluster and the sandbox host saturated. + p.add_argument('--episode-concurrency', type=int, default=48, + help='sandboxes to boot, and how many things run at once in both ' + 'stages: proposal episodes in flight, and solver attempts ' + 'per wave in the difficulty filter') + # 40960, up from 32768, because one proposal is now a single conversation: + # the tool-using turns, the check script and the problem statement all share + # this window. Measured on ex9's three separate calls, the worst case summed + # to about 25k tokens (12394 + 7643 + 2943 plus the appended messages), so + # this leaves room for episodes that take more steps than ex9's 2-5. + # + # 40960 and not more: it is Qwen3-4B's max_position_embeddings, and vLLM + # refuses to start above it -- 49152 was tried and rejected, since going past + # a RoPE model's trained positions produces nan rather than longer context. + p.add_argument('--max-model-len', type=int, default=40960) # Generation control p.add_argument('--keep-target', type=int, default=200, @@ -61,35 +263,116 @@ def parse_args(): p.add_argument('--seed-file', default='', help='seed jsonl with query field') p.add_argument('--seed-mix-prob', type=float, default=0.5) - # Sampling params for round 1 (proposing) + # Sampling params for the exploring stage (proposing) p.add_argument('--propose-temp', type=float, default=1.0) - p.add_argument('--propose-max-tokens', type=int, default=4096) - p.add_argument('--max-turns', type=int, default=20, - help='max tool-calling turns for round 1') + # 8192, not 4096. At 4096, 3 of 12 explore episodes ended on the first turn + # with stop_reason=length and an untouched workspace: the model had written + # 15k, 16k and 10k characters of <think>, two of them without ever closing the + # tag, and one degenerating into a run of newlines. Nothing was dispatched, so + # those three cost a full episode each and produced no end state to write a + # check about. The trajectory ceiling and the engine's max_model_len are + # 32768, well above prompt plus this. + p.add_argument('--propose-max-tokens', type=int, default=8192) + p.add_argument('--max-turns', type=int, default=24, + help='max tool-calling turns for the exploring stage') + # One call per reply, because the calls in one reply run *concurrently*: + # tool_manager.call_many hands a turn to Env.step_batch, which the sandbox + # server runs through ms-agent's parallel_call_tool. The model writes them in + # the order it means them to happen and gets none of the results, so a reply + # that writes a file and then reads it back reads the file as it was before. + # Measured on ex11: 13 of 36 episodes contain an observation that contradicts + # the end state -- read_file answering FileNotFound for a file the snapshot + # lists, glob answering with 0 files, `ls -R` missing a file written earlier + # in the same reply -- and 3 of the 4 kept tasks are among them. One of those + # kept tasks describes two files as "empty", which is what they were only + # because the call that filled them had not run yet. + # + # It also removes the other failure of a batched reply: 6 of 36 episodes + # spent the whole 8192-token budget on one reply holding 70 to 259 calls, + # the tail of it the same read_file over and over, and were discarded whole. + # A reply that can hold one call cannot do either. + p.add_argument('--one-call-per-reply', action='store_true', default=True, + help='stop generation at </tool_call> so each reply carries a single ' + 'call and the model sees its result before choosing the next') + p.add_argument('--no-one-call-per-reply', dest='one_call_per_reply', + action='store_false', + help='let a reply carry several calls, which then run concurrently') + p.add_argument('--stop-after-stuck-turns', type=int, default=2, + help='end an episode after this many consecutive turns that made no ' + 'progress; 0 runs to --max-turns regardless. A turn counts as ' + 'stuck when every call in it came back an error, or every call ' + 'in it was byte-identical to one already made in the episode. ' + 'Replayed over 12 recorded episodes: errors alone would stop 1 ' + 'of 12 and save 9 of 239 calls, since the worst offenders mix a ' + 'failing call with a glob that succeeds; adding the repeat rule ' + 'stops 3 of 12 and saves 63 calls, and those 3 are exactly the ' + 'ones that spent 54, 84 and 17 calls to leave a script that ' + 'could not run.') # Problem statement p.add_argument('--problem-max-chars', type=int, default=8192) + p.add_argument('--check-retries', type=int, default=1, + help='How many times a check script that fails is handed back, ' + 'with the traceback and the workspace listing, to be ' + 'rewritten. ex12 lost 36 of 72 proposals here, and 29 of ' + 'those were one assertion naming a value the model had ' + 'never read -- a row count, a nearly-right content ' + 'string, a timestamp -- on a workspace state that was ' + 'fine. 0 rejects on the first failure, as ex9-ex12 did.') + # Budgets for the two stages appended to the episode. Separate numbers + # because the two are not alike: writing the checks reads the whole episode + # plus the end state and reasons at length (ex9's largest such reply was 7643 + # trainable tokens, so 4096 would cut the tail off and the proposal would be + # discarded as unparseable), while the statement is prose and ex9's largest + # was 2943. + p.add_argument('--check-max-tokens', type=int, default=8192) + p.add_argument('--problem-max-tokens', type=int, default=4096) # Keywords p.add_argument('--keywords-n', type=int, default=128, help='per-category refill target; 0 disables keyword bank') p.add_argument('--keyword-db', default='output/rsi_agentic/keywords.jsonl') p.add_argument('--keyword-gen-calls', type=int, default=8) + # How many of a refill's generating calls go out together. 1 means each is + # told what the ones before it produced; the first round of arm measurements + # effectively ran at 8, where the whole first refill went out with an empty + # 'do not repeat' list and came back with synonyms of each other. + p.add_argument('--keyword-refill-concurrency', type=int, default=1) p.add_argument('--keyword-refill-tries', type=int, default=2) p.add_argument('--keyword-temp', type=float, default=1.3) - p.add_argument('--keyword-max-tokens', type=int, default=1024) + # 1024 measured 8 of 24 generation calls cut off at the budget with nothing + # parseable: the model spends most of it listing candidates inside <think>, + # rewrites the list two or three times, and the JSON array afterwards gets + # severed mid-string. The successful calls landed just under 1024, so the cap + # sat inside the distribution of working replies rather than beyond it. + p.add_argument('--keyword-max-tokens', type=int, default=4096) p.add_argument('--single-kw-prob', type=float, default=0.1) p.add_argument('--combo-arity', default='triple', choices=['triple', 'mix']) p.add_argument('--arity-weights', default='', help="'w1,w2,w3' for --combo-arity mix (empty = uniform)") # Difficulty filter - p.add_argument('--solver-rollouts', type=int, default=4) + # 8 attempts, keeping 2-6: with 4 attempts the band was 1-3 and ex9's + # measured pass counts came out {0: 6, 1: 1, 4: 2} -- two-thirds of the + # tasks landed on an end of the range where one attempt either way changes + # the verdict. 8 costs twice the sandbox time per task and puts the kept + # band around one third of attempts passing. + p.add_argument('--solver-rollouts', type=int, default=8) p.add_argument('--solver-temp', type=float, default=1.0) - p.add_argument('--solver-max-tokens', type=int, default=4096) - p.add_argument('--solver-max-turns', type=int, default=20) - p.add_argument('--keep-min-pass', type=int, default=1) - p.add_argument('--keep-max-margin', type=int, default=1) + # Same 8192 as the explore round, and for the same measured reason: at 4096, + # 15 of 50 solver attempts ended on stop_reason=length with an untouched + # workspace, and one task lost all four of its attempts that way and was + # discarded as too hard. Raising it took that to 0 of 20. It has to stay in + # step with --propose-max-tokens: a task the proposer needed room to build is + # not solvable in less. + p.add_argument('--solver-max-tokens', type=int, default=8192) + p.add_argument('--solver-max-turns', type=int, default=24, + help='NOT WIRED: no solver_explorer is passed, so solver attempts ' + 'run through the same rollout as the proposing episodes and ' + 'obey --max-turns. Kept so the value can be set once the two ' + 'are separated; changing it alone has no effect.') + p.add_argument('--keep-min-pass', type=int, default=2) + p.add_argument('--keep-max-margin', type=int, default=2) # Sandbox p.add_argument('--sandbox-template', default='', @@ -101,8 +384,33 @@ def parse_args(): p.add_argument('--sandbox-timeout', type=int, default=900) p.add_argument('--workspace', default='/workspace', help='working directory inside the sandbox') + p.add_argument('--snapshot-max-files', type=int, default=50, + help='files listed in the end-state snapshot') + p.add_argument('--snapshot-per-file', type=int, default=600, + help='bytes of each file shown to the check writer') + p.add_argument('--snapshot-budget', type=int, default=6000, + help='total bytes of file content in the snapshot') # Output + # ---- Experiment arms. Each isolates one measured failure and can be used + # alone or stacked. Off by default, so an unflagged run is the old behaviour. + # + # A: ex11/ex12/ex13 statements quoted their own answer, because stage 3 asks + # for the full end state while the solver starts empty -- the only way to say + # what a derived file holds is to write out what was computed. Difficulty came + # out 8/8 or 0/8. This makes the statement give input data verbatim and + # everything derived as a rule. + # C: apitest4's statements each wanted an 8-12 file package with a CLI, and + # Qwen3-4B passed 0 of 96 -- 32 attempts spent the whole token budget typing + # source, 64 declared success with the files unwritten. + p.add_argument('--max-build-files', type=int, default=0, + help='arm C: cap the episode at this many files, no package, ' + 'no CLI with subcommands (0 = no cap)') + # For measuring a configuration rather than filling a dataset: two arms are + # only comparable when given the same number of tries. + p.add_argument('--max-proposals-total', type=int, default=0, + help='stop after this many proposals regardless of keep-target ' + '(0 = run until keep-target)') p.add_argument('--random-seed', type=int, default=0) p.add_argument('--out-flows', default='output/rsi_agentic/challenge_flows.jsonl') p.add_argument('--dump-rejected', default='output/rsi_agentic/challenge_rejected.jsonl') @@ -110,7 +418,35 @@ def parse_args(): help='directory for the proposing rounds (token ids + logprobs, one npz ' 'per attempt plus index.jsonl). Empty string turns it off; keeping ' 'it is what leaves the door open to training the challenger itself.') + p.add_argument('--dump-solver-attempts', + default='output/rsi_agentic/solver_attempts.jsonl', + help='one line per difficulty-stage solver attempt: the statement, the ' + 'check script, the attempt, the state it left and what the check ' + 'said. Without it a task measured 0 of 4 gives no way to tell an ' + 'impossible task from a statement that withholds what the check ' + 'demands. Empty string turns it off.') p.add_argument('--no-sort-by-difficulty', action='store_true') + p.add_argument('--stage', default='all', choices=['all', 'keywords', 'explore'], + help="'keywords' runs step 1 only -- fill the keyword bank, draw " + 'the combinations, write the proposal prompts they produce to ' + '--out-flows, and exit without touching the sandbox. ' + "'explore' adds steps 2-4: clear the workspace, run the " + 'sandbox episode, snapshot the end state, and stop before the ' + 'check-writing round. Both exist because a stage that is ' + 'broken cannot be diagnosed from the far end of an ' + 'hours-long full run.') + p.add_argument('--stage-proposals', type=int, default=16, + help='how many proposals --stage keywords or --stage explore runs') + p.add_argument('--dump-explore', default='output/rsi_agentic/explore_episodes.jsonl', + help='one line per --stage explore episode: the prompt, every ' + 'message, every tool call and its observation, and the end ' + 'state the snapshot saw. Empty string turns it off.') + p.add_argument('--dump-keyword-gen', + default='output/rsi_agentic/keyword_gen.jsonl', + help='one line per keyword-generation call: the prompt, the raw ' + 'reply, and what the parser made of it. Without it a bank that ' + 'stays empty gives no way to tell a disobedient model from a ' + 'parser that rejects valid output. Empty string turns it off.') return p.parse_args() @@ -136,97 +472,251 @@ def build_env(args): def main(): args = parse_args() - for path in (args.out_flows, args.dump_rejected, args.keyword_db): + for path in (args.out_flows, args.dump_rejected, args.keyword_db, + args.dump_keyword_gen, args.dump_explore): if path: os.makedirs(os.path.dirname(os.path.abspath(path)) or '.', exist_ok=True) - # Initialize twinkle (sampler only, no trainer) - twinkle.initialize( - mode='ray', nproc_per_node=args.sampler_gpus, lazy_collect=False, - groups=[DeviceGroup(name='sampler', ranks=list(range(args.sampler_gpus)), - device_type='GPU')]) - sampler = vLLMSampler( - model_id=args.model_id, - engine_args={'gpu_memory_utilization': 0.8, 'max_model_len': args.max_model_len}, - device_mesh=DeviceMesh.from_sizes(world_size=args.sampler_gpus, - dp_size=args.sampler_gpus), - remote_group='sampler', - ) - sampler.set_template(args.template, model_id=args.model_id, enable_thinking=True, - max_length=args.max_model_len) + # Build the proposing backend: an OpenAI-compatible API, or a local vLLM + # sampler. The API path skips twinkle.initialize entirely -- it needs no GPUs + # -- and passes no template, since the API re-sends messages as text rather + # than splicing token ids the way the sampler continuation does. + use_api = args.challenger_api + if args.followup_api: + # Split mode needs the local sampler for exploration (that is the trainable + # half), so it cannot run under --challenger-api, which builds no sampler. + if use_api: + raise SystemExit('[challenge] --followup-api and --challenger-api are mutually ' + 'exclusive: --followup-api explores on the local sampler and ' + 'sends only the check/statement stages to the API, while ' + '--challenger-api sends the whole proposing side to the API.') + if not args.challenger_api_model or not args.challenger_api_base: + raise SystemExit('[challenge] --followup-api needs --challenger-api-model and ' + '--challenger-api-base (or LLM_BACKUP_MODEL / ' + 'LLM_BACKUP_BASE_URL).') + if use_api: + if args.solver_rollouts: + raise SystemExit('[challenge] --challenger-api needs --solver-rollouts 0: the ' + 'solver side still runs on the local sampler, which is not ' + 'built in API mode.') + if not args.challenger_api_model or not args.challenger_api_base: + raise SystemExit('[challenge] --challenger-api needs --challenger-api-model and ' + '--challenger-api-base (or LLM_BACKUP_MODEL / LLM_BACKUP_BASE_URL).') + from twinkle_agentic.protocol.openai import OpenAI + backend = OpenAI(model=args.challenger_api_model, + api_key=args.challenger_api_key or None, + base_url=args.challenger_api_base) + template = None + logger.info(f'[challenge] proposing via API model={args.challenger_api_model} ' + f'base={args.challenger_api_base}') + else: + # Initialize twinkle (sampler only, no trainer) + twinkle.initialize( + mode='ray', nproc_per_node=args.sampler_gpus, lazy_collect=False, + groups=[DeviceGroup(name='sampler', ranks=list(range(args.sampler_gpus)), + device_type='GPU')]) + backend = vLLMSampler( + model_id=args.model_id, + engine_args={'gpu_memory_utilization': 0.8, 'max_model_len': args.max_model_len}, + device_mesh=DeviceMesh.from_sizes(world_size=args.sampler_gpus, + dp_size=args.sampler_gpus), + remote_group='sampler', + ) + backend.set_template(args.template, model_id=args.model_id, enable_thinking=True, + max_length=args.max_model_len) - import twinkle.template as template_module - template = getattr(template_module, args.template)( - args.model_id, max_length=args.max_model_len, enable_thinking=True) + import twinkle.template as template_module + template = getattr(template_module, args.template)( + args.model_id, max_length=args.max_model_len, enable_thinking=True) - # Build sandbox environment - env = build_env(args) - schemas = env.tool_schemas() - tool_manager = ToolManager(EnvTool.from_schemas(env, schemas)) + # Build sandbox environments -- one per episode slot, since an episode owns + # its workspace from the reset until its check has run and two episodes + # sharing a sandbox would read each other's files. Skipped for + # --stage keywords: that stage only brainstorms and draws keywords, and + # booting a microVM to do it would make checking step 1 depend on the one part + # of the setup most likely to be down. + # + # ``envs[0]`` is also the one the serial paths use (the difficulty stage, and + # --stage explore), so ``env`` stays a name for it. + envs = [] + env = None + schemas = None + tool_manager = ToolManager() + episode_tool_managers = None + if args.stage != 'keywords': + n_slots = max(1, args.episode_concurrency) + if n_slots == 1: + envs = [build_env(args)] + else: + # Booted in parallel: each is a microVM taking ~10s, and doing eight + # of them one after another would put a minute and a half in front of + # every run. + from concurrent.futures import ThreadPoolExecutor + with ThreadPoolExecutor(max_workers=n_slots) as pool: + envs = list(pool.map(lambda _: build_env(args), range(n_slots))) + env = envs[0] + schemas = env.tool_schemas() + # One ToolManager per sandbox: the tools carry the env they dispatch into, + # so slot i's model turns have to go through slot i's manager. + episode_tool_managers = [ToolManager(EnvTool.from_schemas(e, schemas)) for e in envs] + tool_manager = episode_tool_managers[0] + logger.info(f'[challenge] {len(envs)} sandbox(es) ready ' + f'(episode concurrency {n_slots})') # Explorer: multi-turn rollout with sandbox tools - explorer = build_rollout( - sampler, - template=template, - tool_manager=tool_manager, - max_turns=args.max_turns, - sampling_params=SamplingParams(max_tokens=args.propose_max_tokens, num_samples=1, - logprobs=1, temperature=args.propose_temp, top_p=0.95), - ) + # ``stop`` ends the reply at the end of the first tool call, and + # ``include_stop_str_in_output`` keeps that '</tool_call>' in what the policy + # is trained on -- vLLM drops the matched stop by default, which would train + # every turn to end on an unclosed block. + explore_stop = ['</tool_call>'] if (args.one_call_per_reply and not use_api) else None + # Sent on every API call when set. Capping the reasoning is the one knob that + # moved the wall-clock: 58s -> 10s per turn at 2048 on a ~15k-character context. + api_extra_body = ({'thinking_budget': args.challenger_thinking_budget} + if (use_api and args.challenger_thinking_budget > 0) else None) + if api_extra_body: + logger.info(f'[challenge] thinking_budget={args.challenger_thinking_budget} ' + f'on every API call') + explore_params = SamplingParams(max_tokens=args.propose_max_tokens, num_samples=1, + logprobs=1, temperature=args.propose_temp, top_p=0.95, + stop=explore_stop, + include_stop_str_in_output=bool(explore_stop)) + # One tool call per reply and stuck-turn early stop are sampler-path features: + # the API dispatches native tool_calls (never a '</tool_call>' string) and + # APIMultiTurnRollout takes neither kwarg. + if use_api: + explorer = build_rollout( + backend, tool_manager=tool_manager, max_turns=args.max_turns, + concurrency=args.challenger_concurrency, sampling_params=explore_params, + extra_body=api_extra_body) + else: + explorer = build_rollout( + backend, template=template, tool_manager=tool_manager, + max_turns=args.max_turns, stop_after_stuck_turns=args.stop_after_stuck_turns, + sampling_params=explore_params) + + # Keyword brainstorming runs through this one instead of the sandbox + # explorer: a list is a text answer, and a bracketed list in a reply is + # exactly what the sandbox explorer would try to dispatch as a call. + # + # max_turns=1 is what makes it tool-less: MultiTurnRollout ends the + # trajectory on the turn limit before it dispatches anything, so the empty + # ToolManager below is never consulted. It is here because the rollout + # requires one at construction, not because these calls have tools. + keyword_params = SamplingParams(max_tokens=args.keyword_max_tokens, num_samples=1, + logprobs=1, temperature=args.keyword_temp, top_p=0.98) + if use_api: + keyword_explorer = build_rollout( + backend, tool_manager=ToolManager(), max_turns=1, + concurrency=args.challenger_concurrency, sampling_params=keyword_params, + extra_body=api_extra_body) + else: + keyword_explorer = build_rollout( + backend, + template=template, + tool_manager=ToolManager(), + max_turns=1, + sampling_params=keyword_params, + ) # Sandbox control functions -- use env.runner() which resolves tool names # (ms-agent registers tools as "server---name") and parses exit codes from # the marker protocol, so we don't rely on string matching. - runner = env.runner() - - # Cleared through the python tool, not `rm -rf`: ms-agent's safety policy - # rejects `rm -rf` outright ("Blocked by safety rule"), and it rejects globs - # in write operations, which rules out `find -delete` too. The script asserts - # the directory really is empty so a future policy change surfaces as a - # failed reset instead of tasks quietly inheriting the previous workspace. - _CLEAR = ''' -import os, shutil -root = {workspace!r} -os.makedirs(root, exist_ok=True) -for name in os.listdir(root): - path = os.path.join(root, name) - if os.path.isdir(path) and not os.path.islink(path): - shutil.rmtree(path, ignore_errors=True) - else: - os.remove(path) -leftover = os.listdir(root) -assert not leftover, 'workspace not empty after clear: %r' % (leftover,) -''' + # + # One runner per sandbox; ``slot`` picks which one. The challenger passes the + # slot of the episode it is serving, so a concurrent episode never clears or + # inspects another episode's workspace. Everything serial (the difficulty + # stage, --stage explore) leaves it at the default and uses sandbox 0. + # + # Empty for --stage keywords, which has no sandbox. The functions below index + # into it and would raise if that stage ever reached them; it returns first. + runners = [e.runner() for e in envs] + runner = runners[0] if runners else None - def reset_fn(): - """Empty the sandbox workspace before an episode. + def reset_fn(slot: int = 0): + """Empty sandbox ``slot``'s workspace before an episode. Raises rather than returning: every caller depends on a clean start, and a silent no-op here means a task inherits the previous task's files -- which lets a solver pass without doing anything and makes the difficulty numbers meaningless. + + This is also the one point where losing the sandbox costs nothing, since + the workspace is about to be emptied regardless -- so a runtime that went + away is rebuilt here instead of ending a run that may have hours of + proposals behind it. """ - exit_code, output = runner(_CLEAR.format(workspace=args.workspace), 'python') + if envs[slot].ensure_ready(): + # A rebuilt sandbox starts empty, so the clear below is redundant, + # but running it anyway keeps one path through this function. The + # rebuild replaces the sandbox behind this env, so the runner is + # re-fetched rather than reused. + runners[slot] = envs[slot].runner() + logger.warning(f'[challenge] sandbox {slot} was rebuilt before this episode') + exit_code, output = runners[slot]( + CLEAR_WORKSPACE.format(workspace=args.workspace), 'python') if exit_code != 0: raise RuntimeError(f'workspace reset failed (exit {exit_code}): {output[-400:]}') - def run_check_fn(script: str): - """Run a python check script in the sandbox; returns (exit_code, output).""" - return runner(script, 'python') + def run_check_fn(script: str, slot: int = 0): + """Run a python check script in sandbox ``slot``; returns (exit_code, output).""" + return runners[slot](script, 'python') + + def workspace_snapshot_fn(slot: int = 0): + """Every file the episode left behind: ``path size`` lines, then contents. - def workspace_snapshot_fn(): - """Get a summary of the current workspace state.""" - _, output = runner( - f'find {args.workspace} -type f -printf "%P %s\\n" 2>/dev/null | head -50', - 'shell') - return output or '(empty)' + This is the ground truth the check script is written against, so it is + unwrapped from the tool's JSON envelope and returned as a bare listing: + the model has to be able to read it as a directory rather than as a tool + result, or it falls back on what it *believes* it created. + + Returns an empty string when the episode left nothing behind, and also + when the listing could not be read at all. Both mean the same thing to + the caller -- there is no end state to write checks about -- and neither + may be dressed up as a plausible one: a snapshot that says "empty" + when it means "I could not look" produces tasks whose only true + assertion is that nothing happened. + """ + exit_code, output = runners[slot]( + WORKSPACE_SNAPSHOT.format(workspace=args.workspace, + max_files=args.snapshot_max_files, + per_file=args.snapshot_per_file, + total_budget=args.snapshot_budget), + 'python') + if exit_code != 0: + # Not fatal, but not silent either: checks written against a missing + # end state are the failure this whole function exists to prevent. + logger.warning(f'[challenge] workspace snapshot failed (exit {exit_code}): ' + f'{output[-200:]}') + return '' + return tool_payload(output).strip() + + # Arm B. Read at most this much per call: the executor truncates its output + # near 8 KB, and base64 grows 3 bytes into 4, so 4 KB of file is about 5.5 KB + # of text with room left for the JSON envelope. + SLICE_BYTES = 4096 + + + # The opening the solver is measured against, built by the same function the + # eval script uses so that n_pass and pass@k are measuring one thing. Until + # this existed the difficulty stage handed over the statement as a lone user + # message with no system prompt: nothing said the model was in a sandbox, could + # take many turns, or should make one call per reply, and it answered by + # writing whole programs into a single call argument until they truncated. + _solver_harness = solver_harness(args.agent_config) if args.solver_rollouts > 0 else None + + def solver_prompt_fn(query: str): + return _solver_harness.start(query) # Keywords store = None + # Arm D replaces the three topic axes with one bank of 'kind of work' phrases: + # a proposal takes one phrase, not one entry from each of three axes. + categories = CATEGORIES + category_desc = CATEGORY_DESC if args.keywords_n > 0: - store = KeywordStore(args.keyword_db, CATEGORIES) + store = KeywordStore(args.keyword_db, categories) logger.info('[challenge] keyword bank loaded: ' - + ', '.join(f'{c}={len(store.items[c])}' for c in CATEGORIES)) + + ', '.join(f'{c}={len(store.items[c])}' for c in categories)) # Seeds seeds = [] @@ -243,58 +733,243 @@ def workspace_snapshot_fn(): def _reject(record): if rejected is not None: rejected.write(json.dumps(record, ensure_ascii=False, default=str) + '\n') + # Flushed per record: this file is the only account of why proposals + # are being dropped, and a run is worth watching for hours before it + # ends. Buffered, it stays empty until then. + rejected.flush() propose_writer = ProposeTrajWriter(args.dump_propose_traj) + # Solver attempts from the difficulty stage. One line per attempt, so a task + # measured at 0 of 4 can be read rather than guessed at: the attempt, the + # state it left, and what the check said about it. + solver_log = (open(args.dump_solver_attempts, 'w', encoding='utf-8') + if args.dump_solver_attempts else None) + + def _solver_attempt(record): + if solver_log is not None: + solver_log.write(json.dumps(record, ensure_ascii=False, default=str) + '\n') + solver_log.flush() + + # Keyword generation, one line per call. The bank is the first step of the + # whole pipeline and the easiest place to fail invisibly: a reply the parser + # rejects leaves the bank empty, and every proposal downstream then runs the + # no-keyword prompt while the run looks healthy. Whole runs went that way + # before this existed, so prompt and reply are both kept verbatim. + keyword_log = (open(args.dump_keyword_gen, 'w', encoding='utf-8') + if args.dump_keyword_gen else None) + + def _keyword_gen(record): + if keyword_log is not None: + keyword_log.write(json.dumps(record, ensure_ascii=False, default=str) + '\n') + keyword_log.flush() + # Build challenger - prompts = agentic_prompts() + prompts = agentic_prompts(max_build_files=args.max_build_files) + # Followup API: explore on the local (trainable) sampler above, but write the + # check script and problem statement over an OpenAI-compatible API (qwen3-max). + # Reuses the --challenger-api-* connection args; the thinking cap, when set, is + # sent as extra_body on every check/statement call. + followup_api = None + followup_extra_body = None + if args.followup_api: + from twinkle_agentic.protocol.openai import OpenAI + followup_api = OpenAI(model=args.challenger_api_model, + api_key=args.challenger_api_key or None, + base_url=args.challenger_api_base) + if args.challenger_thinking_budget > 0: + followup_extra_body = {'thinking_budget': args.challenger_thinking_budget} + logger.info(f'[challenge] followup (check + statement) via API ' + f'model={args.challenger_api_model} base={args.challenger_api_base}' + + (f' thinking_budget={args.challenger_thinking_budget}' + if followup_extra_body else '')) challenger = AgenticChallenger( prompts, explorer, seeds=seeds, keyword_store=store, - category_desc=CATEGORY_DESC if store else None, + category_desc=category_desc if store else None, seed_mix_prob=args.seed_mix_prob, reset_fn=reset_fn, run_check_fn=run_check_fn, workspace_snapshot_fn=workspace_snapshot_fn, + # The executor's own schemas, so the rounds that may call tools advertise + # exactly what will run -- same source as the training script uses. + tool_schemas=schemas, + episode_concurrency=max(1, args.episode_concurrency), + episode_tool_managers=episode_tool_managers, combo_arity=args.combo_arity, arity_weights=[float(x) for x in args.arity_weights.split(',')] if args.arity_weights else None, single_kw_prob=args.single_kw_prob, keyword_refill_target=args.keywords_n, keyword_gen_calls=args.keyword_gen_calls, + keyword_refill_concurrency=max(1, args.keyword_refill_concurrency), keyword_refill_tries=args.keyword_refill_tries, keyword_params=SamplingParams(max_tokens=args.keyword_max_tokens, num_samples=1, logprobs=1, temperature=args.keyword_temp, top_p=0.98), + # Same temperature as the episode, different budgets: the only thing + # being changed per stage is how much room the reply gets. + check_params=SamplingParams(max_tokens=args.check_max_tokens, num_samples=1, + logprobs=1, temperature=args.propose_temp, top_p=0.95), + problem_params=SamplingParams(max_tokens=args.problem_max_tokens, num_samples=1, + logprobs=1, temperature=args.propose_temp, top_p=0.95), + followup_api=followup_api, + followup_extra_body=followup_extra_body, + keyword_explorer=keyword_explorer, min_batch=args.sampler_gpus, problem_max_chars=args.problem_max_chars, + max_proposals_total=args.max_proposals_total, + solver_prompt_fn=solver_prompt_fn if _solver_harness is not None else None, + check_retries=args.check_retries, reject_sink=_reject, propose_sink=propose_writer.write, + solver_sink=_solver_attempt, + keyword_sink=_keyword_gen, max_proposals_per_round=args.max_proposals_per_round, solver_rollouts=args.solver_rollouts, keep_min_pass=args.keep_min_pass, keep_max_pass_margin=args.keep_max_margin, + # One call per reply here too, for the same reason and to keep the two + # sides comparable: a solver whose read-back is dispatched alongside the + # write it is checking fails a task the proposer built cleanly, and + # n_pass would then be measuring the dispatch, not the difficulty. solver_params=SamplingParams(max_tokens=args.solver_max_tokens, num_samples=1, - logprobs=1, temperature=args.solver_temp, top_p=0.95), + logprobs=1, temperature=args.solver_temp, top_p=0.95, + stop=explore_stop, + include_stop_str_in_output=bool(explore_stop)), seed=args.random_seed, ) # Generate batch_size = args.batch_size or args.keep_target kept = [] - for batch in challenger(batch_size=batch_size, total=args.keep_target): - kept.extend(batch) - logger.info(f'[challenge] kept {len(kept)}/{args.keep_target} so far; ' - f'stats {challenger.stats}') + + # --stage keywords stops after step 1: fill the bank, draw the combinations, + # write out the proposal prompts they produce, and exit without touching the + # sandbox. Step 1 was broken for several runs and the failure was only + # visible by reading what it fed the next step, so it has to be runnable on + # its own rather than only as the first minute of an hours-long run. + if args.stage == 'keywords': + proposals = challenger.propose(args.stage_proposals) + with open(args.out_flows, 'w', encoding='utf-8') as out: + for i, proposal in enumerate(proposals): + data = proposal.get('user_data') + out.write(json.dumps({ + 'index': i, + 'keywords': user_data_get(data, 'keywords', []), + 'seeded': user_data_get(data, 'seeded', False), + 'prompt': proposal['messages'][-1]['content'], + }, ensure_ascii=False) + '\n') + drawn = sum(1 for p in proposals + if user_data_get(p.get('user_data'), 'keywords', [])) + logger.info(f'[challenge] stage=keywords: {len(proposals)} proposals, ' + f'{drawn} of them carry keywords') + if store is not None: + store.save() + logger.info('[challenge] keyword bank saved -> ' + args.keyword_db + + ' (' + ', '.join(f'{c}={len(store.items[c])}' + for c in categories) + ')') + if keyword_log is not None: + keyword_log.close() + propose_writer.close() + if rejected is not None: + rejected.close() + if solver_log is not None: + solver_log.close() + if env is not None: + for e in envs: + e.close() + return + + # --stage explore stops after step 4: draw a proposal, clear the workspace, + # run the sandbox episode, snapshot what it left, and stop before the + # check-writing round. What it is for: the episode is where the run either + # produces something worth writing a check about or leaves an empty directory, + # and 9 of 30 proposals in run11 left an empty one for reasons the rejection + # record could not distinguish. Everything the episode saw and did goes out + # verbatim, so that question is answerable from the file. + if args.stage == 'explore': + explore_log = (open(args.dump_explore, 'w', encoding='utf-8') + if args.dump_explore else None) + empty = 0 + for i, proposal in enumerate(challenger.propose(args.stage_proposals)): + reset_fn() + result = challenger.explore([proposal]) + episode = result[0] if result else {} + snapshot = workspace_snapshot_fn() + if not snapshot.strip(): + empty += 1 + messages = episode.get('messages') or [] + calls = sum(len(m.get('tool_calls') or []) for m in messages + if isinstance(m, dict)) + logger.info(f'[challenge] episode {i}: stop={episode.get("stop_reason")} ' + f'truncated={bool(episode.get("truncated"))} ' + f'stuck_stop={bool(episode.get("stuck_stop"))} ' + f'turns={episode.get("turns")} calls={calls} ' + f'end_state={len(snapshot)}b') + if explore_log is not None: + explore_log.write(json.dumps({ + 'index': i, + 'keywords': user_data_get(proposal.get('user_data'), 'keywords', []), + 'prompt': proposal['messages'][-1]['content'], + 'stop_reason': episode.get('stop_reason'), + 'truncated': bool(episode.get('truncated')), + 'stuck_stop': bool(episode.get('stuck_stop')), + 'turns': episode.get('turns'), + 'n_tool_calls': calls, + 'messages': messages, + 'end_state': snapshot, + }, ensure_ascii=False, default=str) + '\n') + explore_log.flush() + logger.info(f'[challenge] stage=explore: {args.stage_proposals} episodes, ' + f'{empty} left an empty workspace') + if explore_log is not None: + explore_log.close() + if keyword_log is not None: + keyword_log.close() + propose_writer.close() + if rejected is not None: + rejected.close() + if solver_log is not None: + solver_log.close() + if store is not None: + store.save() + if env is not None: + for e in envs: + e.close() + return + + # Appended as batches arrive, then rewritten sorted at the end. A run that + # keeps one task every few minutes for hours cannot afford to hold them all + # in memory only: a crash at hour three would leave nothing to train on, + # while an unsorted partial file is a usable task set. + with open(args.out_flows, 'w', encoding='utf-8') as partial: + for batch in challenger(batch_size=batch_size, total=args.keep_target): + for offset, task in enumerate(batch): + partial.write(json.dumps(flow_record(len(kept) + offset, task), + ensure_ascii=False) + '\n') + partial.flush() + kept.extend(batch) + logger.info(f'[challenge] kept {len(kept)}/{args.keep_target} so far; ' + f'stats {challenger.stats}') if rejected is not None: rejected.close() + if solver_log is not None: + solver_log.close() propose_writer.close() + # Before the keyword log is closed: expanding the bank generates keywords, + # and generating them writes to that log. Closing it first ended ex11 -- + # after all 4 tasks were kept and written -- with `ValueError: I/O operation + # on closed file`, which also skipped store.save() below and every line + # after it, so the run reported nothing about what it had produced. if store is not None: challenger.expand_hard_keywords() store.save() logger.info('[challenge] keyword bank saved -> ' + args.keyword_db) + if keyword_log is not None: + keyword_log.close() # Sort by difficulty (hardest last) if not args.no_sort_by_difficulty: @@ -303,13 +978,17 @@ def _reject(record): # Write output write_flows(kept, args) logger.info(f'[challenge] wrote {len(kept)} tasks -> {args.out_flows}') + if env.n_recoveries: + logger.warning(f'[challenge] sandbox was rebuilt {env.n_recoveries} time(s) during ' + f'this run; episodes in flight at those moments were lost') dist = {} for task in kept: n = user_data_get(task.get('user_data'), 'n_pass', 0) dist[n] = dist.get(n, 0) + 1 logger.info(f'[challenge] pass-count distribution: {dict(sorted(dist.items()))}') - env.close() + for e in envs: + e.close() class ProposeTrajWriter: @@ -379,6 +1058,7 @@ def write(self, record): 'rounds': meta, } self.index.write(json.dumps(line, ensure_ascii=False, default=str) + '\n') + self.index.flush() def close(self): if self.index is not None: @@ -386,23 +1066,30 @@ def close(self): logger.info(f'[challenge] wrote {self.n} propose traces -> {self.dir}') +def flow_record(index, task): + """One task as the training script reads it back.""" + data = task.get('user_data') + messages = task.get('messages') or [] + query = next((m['content'] for m in messages if m.get('role') == 'user'), '') + return { + 'id': f'ag_{index:06d}', + 'query': query, + 'check_script': user_data_get(data, 'check_script', ''), + # Arm B: run before the solver starts, to put the input files it is told + # it already has on disk. Empty for every other arm. + 'setup_script': user_data_get(data, 'setup_script', ''), + 'n_pass': user_data_get(data, 'n_pass'), + 'n_rollouts': user_data_get(data, 'n_rollouts'), + 'keywords': user_data_get(data, 'keywords', []), + 'seeded': user_data_get(data, 'seeded', False), + } + + def write_flows(kept, args): - """Write one flow per task.""" + """Write one flow per task, replacing whatever the run appended as it went.""" with open(args.out_flows, 'w', encoding='utf-8') as f: for i, task in enumerate(kept): - data = task.get('user_data') - messages = task.get('messages') or [] - query = next((m['content'] for m in messages if m.get('role') == 'user'), '') - flow = { - 'id': f'ag_{i:06d}', - 'query': query, - 'check_script': user_data_get(data, 'check_script', ''), - 'n_pass': user_data_get(data, 'n_pass'), - 'n_rollouts': user_data_get(data, 'n_rollouts'), - 'keywords': user_data_get(data, 'keywords', []), - 'seeded': user_data_get(data, 'seeded', False), - } - f.write(json.dumps(flow, ensure_ascii=False) + '\n') + f.write(json.dumps(flow_record(i, task), ensure_ascii=False) + '\n') if __name__ == '__main__': diff --git a/cookbook/rsi/agentic/episode.py b/cookbook/rsi/agentic/episode.py new file mode 100644 index 000000000..f817d26f6 --- /dev/null +++ b/cookbook/rsi/agentic/episode.py @@ -0,0 +1,234 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Episode construction and scoring for agentic RSI, shared by training and eval. + +Both halves need the same three things and must not disagree about any of them: +how an episode is built (a sandbox with ms-agent's tools plus a local harness +that only shapes messages), how the tool contract is advertised (schemas read off +the executor that will honour them), and how a trajectory is scored (the task's +own checks, run against the state the episode left behind). + +An eval that differed from training on any of these would measure something other +than what was trained, so this module is the single definition and the scripts +are only wiring. +""" +import json +import os +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple + +from twinkle import get_logger +from twinkle_agentic.envs import EnvTool +from twinkle_agentic.harness import MsAgentHarness +from twinkle_agentic.tools.tool_manager import ToolManager +from twinkle_agentic.verifier.result_check import (CheckContext, checks_from_dicts, + run_checks) + +from remote_tool_env import RemoteMsAgentToolEnv # noqa: I100,I202 + +logger = get_logger() + + +@dataclass(frozen=True) +class SandboxConfig: + """Everything about where episodes run, in one object. + + Read from the environment so training and eval cannot be pointed at + different sandboxes by accident. + """ + + agent_config: str = 'cookbook/rsi/agentic/rsi_agent.yaml' + template: str = 'twinkle-rsi-msagent' + api_url: str = 'http://127.0.0.1:8000' + # Must outlast a whole episode plus the checks that run after it. + timeout: int = 900 + # Booting and scoring are network-bound, so they are done on threads. This + # caps how many sandboxes are talked to at once, not how many exist. + concurrency: int = 16 + # 'fraction' gives partial credit per check; 'all_or_nothing' is stricter. + # Applies to structured checks only -- a check script has no partial credit. + score_mode: str = 'fraction' + + @classmethod + def from_env(cls) -> 'SandboxConfig': + return cls( + agent_config=os.environ.get('RSI_AGENT_CONFIG', cls.agent_config), + template=os.environ.get('AENV_TEMPLATE', cls.template), + api_url=os.environ.get('AENV_API_URL', cls.api_url), + timeout=int(os.environ.get('RSI_SANDBOX_TIMEOUT', cls.timeout)), + concurrency=int(os.environ.get('RSI_ENV_CONCURRENCY', cls.concurrency)), + score_mode=os.environ.get('RSI_SCORE_MODE', cls.score_mode), + ) + + +def load_tasks(path: str) -> List[Dict[str, Any]]: + """Read the task file and fail loudly on a task that can never be scored. + + Supports both formats: + - ``check_script``: a python script, scored by exit status (challenge.py). + - ``checks``: structured Check dicts (see tasks.example.jsonl). + """ + tasks = [] + with open(path, encoding='utf-8') as f: + for lineno, line in enumerate(f, 1): + if not line.strip(): + continue + task = json.loads(line) + if not task.get('query'): + raise ValueError(f'{path}:{lineno} has no query') + if task.get('check_script'): + task['_checks'] = None + elif task.get('checks'): + task['_checks'] = checks_from_dicts(task['checks']) + else: + raise ValueError(f'{path}:{lineno} ({task.get("id")}) declares no checks ' + f'and no check_script') + tasks.append(task) + if not tasks: + raise ValueError(f'{path} contains no tasks') + return tasks + + +def solver_harness(agent_config: str): + """A harness that only shapes messages: no llm, no tools. + + Popping ``tools`` matters as much as popping ``llm``, and for the same reason + omitting the section from the yaml is not enough: ms-agent merges its own + agent.yaml underneath, which declares file_system and code_executor, so a live + shell executor would otherwise be constructed on the training host with access + to the whole machine. Popping them after the merge leaves the harness with zero + tools -- and the system prompt byte-identical, because ms-agent does not fold + the tool list into it. + + Shared with challenge.py's difficulty stage on purpose: the opening a task is + measured against there has to be the opening it is evaluated against here, and + a second copy of these four lines would drift. + """ + from omegaconf import OmegaConf, open_dict + + agent_cfg = OmegaConf.load(agent_config) + harness = MsAgentHarness(config=agent_cfg) + with open_dict(harness.agent.config): + harness.agent.config.pop('llm', None) + harness.agent.config.pop('tools', None) + harness.prepare() + return harness + + +def build_episode(task: Dict[str, Any], cfg: SandboxConfig) -> Tuple[Any, Any, Any, Dict]: + """Create one episode: a sandbox with ms-agent's tools, plus a local harness.""" + harness = solver_harness(cfg.agent_config) + + env = RemoteMsAgentToolEnv( + template=cfg.template, + config_path=cfg.agent_config, + api_url=cfg.api_url, + sandbox_timeout=cfg.timeout, + ) + env.reset() + + # A task may hand the solver its input files instead of asking it to write + # them (challenge.py --preseed-inputs). Loudly, not on a best-effort basis: a + # statement that says the inputs are on disk, run against a workspace where + # they are not, scores 0 for a reason that has nothing to do with the task. + setup = task.get('setup_script') + if setup: + exit_code, output = env.runner()(setup, 'python') + if exit_code != 0: + raise RuntimeError(f'[{task.get("id")}] setup_script failed ' + f'(exit {exit_code}): {output[-400:]}') + + trajectory = harness.start(task['query']) + # The executor's own schemas, not the harness's (which are now empty by + # construction). Advertising what will run is the whole point of sourcing + # them from the sandbox. + schemas = env.tool_schemas() + trajectory['tools'] = schemas + tool_manager = ToolManager(EnvTool.from_schemas(env, schemas)) + return harness, env, tool_manager, trajectory + + +def boot_episodes(tasks: List[Dict[str, Any]], + cfg: SandboxConfig) -> List[Tuple[Any, Any, Any, Dict]]: + """Bring up every rollout's sandbox at once, all-or-nothing. + + Serial boot would dominate the step: a microVM plus ms-agent's import runs + to seconds, multiplied by ``batch_size x num_generations``. + + All-or-nothing because GRPO groups are positional -- advantages are taken + over consecutive runs of ``num_generations`` -- so dropping one episode would + not shrink its group, it would shift every later group onto the wrong task. + """ + episodes: List[Optional[Tuple[Any, Any, Any, Dict]]] = [None] * len(tasks) + error: Optional[BaseException] = None + with ThreadPoolExecutor(max_workers=cfg.concurrency) as pool: + futures = {pool.submit(build_episode, task, cfg): slot + for slot, task in enumerate(tasks)} + for future in as_completed(futures): + try: + episodes[futures[future]] = future.result() + except Exception as e: # noqa + error = error or e + if error is not None: + for episode in episodes: + if episode is not None: + episode[1].close() + raise RuntimeError(f'sandbox boot failed: {error}') from error + return episodes # type: ignore[return-value] + + +def score_episode(task: Dict[str, Any], env: RemoteMsAgentToolEnv, + trajectory: Dict[str, Any], snapshot_dir: str, + cfg: SandboxConfig) -> float: + """Run the task's checks against the state this episode left behind. + + A ``check_script`` is the whole verdict by exit status: no partial credit, no + judge model, no drift between the run that invented the task and the run + being scored. Structured checks go through ``run_checks`` instead. + """ + check_script = task.get('check_script') + if check_script: + exit_code, output = env.runner()(check_script, 'python') + if exit_code != 0: + logger.debug(f'[{task.get("id")}] check_script failed (exit {exit_code}): ' + f'{output[-200:]}') + return 1.0 if exit_code == 0 else 0.0 + + final_answer = '' + for msg in reversed(trajectory.get('messages') or []): + if msg.get('role') == 'assistant' and (msg.get('content') or '').strip(): + final_answer = msg['content'] + break + + ctx = CheckContext( + workspace=env.download_workspace(snapshot_dir), + final_answer=final_answer, + runner=env.runner(), + ) + report = run_checks(task['_checks'], ctx, mode=cfg.score_mode) + if not report.all_passed: + logger.debug(f'[{task.get("id")}] {report.n_passed}/{report.n_total} checks: ' + f'{report.failures()}') + return report.score + + +def score_episodes(tasks: List[Dict[str, Any]], envs: List[RemoteMsAgentToolEnv], + outs: List[Dict[str, Any]], snapshot_root: str, + cfg: SandboxConfig) -> List[float]: + """Score every episode in parallel; a scoring crash costs one reward, not the step. + + Each check is a sandbox round trip, so scoring serially would idle the GPUs + for as long as booting did. An episode whose sandbox died mid-check scores + zero, which is also what it would have scored had the checks simply failed. + """ + + def _score(slot: int) -> float: + snapshot = os.path.join(snapshot_root, f'slot{slot:03d}') + try: + return score_episode(tasks[slot], envs[slot], outs[slot], snapshot, cfg) + except Exception as e: # noqa + logger.warning(f'[{snapshot_root} slot {slot}] scoring failed: {e}') + return 0.0 + + with ThreadPoolExecutor(max_workers=cfg.concurrency) as pool: + return list(pool.map(_score, range(len(outs)))) diff --git a/cookbook/rsi/agentic/eval.py b/cookbook/rsi/agentic/eval.py new file mode 100644 index 000000000..5c5534ba5 --- /dev/null +++ b/cookbook/rsi/agentic/eval.py @@ -0,0 +1,244 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Held-out evaluation for agentic RSI: pass rate on tasks the trainer never saw. + +Episodes are built and scored by :mod:`episode`, the same module ``rl.py`` uses, +so a number reported here is the number training was optimising -- an eval that +constructed episodes differently would measure a different agent. + +What it adds on top of training is only what training does not need: several +attempts per task (a single attempt at temperature 1 is a coin flip, not a rate), +no optimizer, and a LoRA read off disk rather than synced from a live trainer. + +Usage:: + + # baseline, no adapter + python cookbook/rsi/agentic/eval.py --tasks output/.../eval_tasks.jsonl \\ + --label base --out output/.../eval_base.jsonl + + # after training + python cookbook/rsi/agentic/eval.py --tasks output/.../eval_tasks.jsonl \\ + --adapter-path output/rsi-agentic-final --label trained \\ + --out output/.../eval_trained.jsonl + +Both runs must use the same ``--tasks``, ``--rollouts-per-task`` and sampling +parameters, or the comparison is not one. +""" +import argparse +import json +import os +import shutil +import statistics +import sys + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle.data_format import SamplingParams +from twinkle.sampler import vLLMSampler +from twinkle.template import Template +from twinkle_agentic.rollout.multi_turn import MultiTurnRollout + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from episode import (SandboxConfig, boot_episodes, load_tasks, # noqa: E402,I100,I202 + score_episodes) + +logger = get_logger() + + +def parse_args(): + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--tasks', required=True, help='task jsonl (challenge.py or structured)') + p.add_argument('--model-id', default='ms://Qwen/Qwen3-4B') + p.add_argument('--adapter-path', default='', + help='LoRA directory saved by rl.py; empty evaluates the base model') + p.add_argument('--label', default='eval', help='name for this measurement in the log') + p.add_argument('--sampler-gpus', type=int, default=4) + p.add_argument('--max-model-len', type=int, default=32768) + p.add_argument('--max-lora-rank', type=int, default=32) + + p.add_argument('--rollouts-per-task', type=int, default=4, + help='attempts per task; the pass rate is over these') + p.add_argument('--episodes-per-wave', type=int, default=16, + help='sandboxes alive at once; keep at or below RSI_ENV_CONCURRENCY') + p.add_argument('--max-turns', type=int, default=20) + # Has to equal the challenger's --solver-max-tokens and --propose-max-tokens. + # An eval that gives the model less room than the run that built the tasks is + # measuring the budget, not the model: at 4096, 15 of 50 attempts ended on + # stop_reason=length with an untouched workspace. + p.add_argument('--max-tokens', type=int, default=8192) + p.add_argument('--temperature', type=float, default=1.0) + p.add_argument('--top-p', type=float, default=0.95) + + p.add_argument('--out', default='', help='per-episode results jsonl') + # The per-episode row says how an attempt ended but not what it did, and a + # rate of 0 has two very different causes that only the conversation tells + # apart: the model worked and got it wrong, or it answered in prose and never + # touched a tool. 71 of 96 attempts in one run ended within two turns, which + # is unreadable without this. + p.add_argument('--dump-messages', default='', + help='jsonl of the full conversation per episode, for reading attempts') + p.add_argument('--keep-workspaces', action='store_true') + return p.parse_args() + + +def build_sampler(args): + twinkle.initialize( + mode='ray', nproc_per_node=args.sampler_gpus, lazy_collect=False, + groups=[DeviceGroup(name='sampler', ranks=list(range(args.sampler_gpus)), + device_type='GPU')]) + engine_args = {'gpu_memory_utilization': 0.8, 'max_model_len': args.max_model_len} + if args.adapter_path: + # Declared at construction or the engine has no slot to load into, which + # surfaces much later as an adapter that appears to do nothing. + engine_args.update({'enable_lora': True, 'max_lora_rank': args.max_lora_rank}) + sampler = vLLMSampler( + model_id=args.model_id, + engine_args=engine_args, + device_mesh=DeviceMesh.from_sizes(world_size=args.sampler_gpus, + dp_size=args.sampler_gpus), + remote_group='sampler', + ) + sampler.set_template('Template', model_id=args.model_id, enable_thinking=True, + max_length=args.max_model_len) + return sampler + + +def main(): + args = parse_args() + if args.adapter_path and not os.path.isdir(args.adapter_path): + raise SystemExit(f'[eval] no adapter directory at {args.adapter_path}') + tasks = load_tasks(args.tasks) + cfg = SandboxConfig.from_env() + logger.info(f'[eval:{args.label}] {len(tasks)} tasks x {args.rollouts_per_task} attempts, ' + f'adapter={args.adapter_path or "(base model)"}') + logger.info(f'[eval:{args.label}] sandboxes: template={cfg.template} api={cfg.api_url}') + + sampler = build_sampler(args) + template = Template(args.model_id, max_length=args.max_model_len, enable_thinking=True) + template.truncation_strategy = 'delete' + rollout = MultiTurnRollout( + sampler=sampler, + template=template, + sampling_params=SamplingParams(max_tokens=args.max_tokens, num_samples=1, logprobs=1, + temperature=args.temperature, top_p=args.top_p), + max_turns=args.max_turns, + max_trajectory_tokens=args.max_model_len, + adapter_path=args.adapter_path or None, + ) + + # One flat list of attempts, so a wave is a fixed number of sandboxes no + # matter how the attempts distribute over tasks. + attempts = [(task, rep) for task in tasks for rep in range(args.rollouts_per_task)] + results = [] + n_boot_failed = 0 + scratch = os.path.join('output', 'rsi_agentic', f'_eval_{args.label}') + msg_dump = (open(args.dump_messages, 'w', encoding='utf-8') + if args.dump_messages else None) + + for start in range(0, len(attempts), args.episodes_per_wave): + wave = attempts[start:start + args.episodes_per_wave] + wave_tasks = [task for task, _ in wave] + try: + episodes = boot_episodes(wave_tasks, cfg) + except Exception as e: # noqa + # Reported, never silently dropped: an eval that quietly measures + # fewer episodes than it claims is worse than one that admits a gap. + n_boot_failed += len(wave) + logger.warning(f'[eval:{args.label}] wave at {start} failed to boot: {e}') + continue + harnesses = [ep[0] for ep in episodes] + envs = [ep[1] for ep in episodes] + tool_managers = [ep[2] for ep in episodes] + trajectories = [ep[3] for ep in episodes] + wave_dir = os.path.join(scratch, f'wave{start:04d}') + try: + outs = rollout(trajectories, harness=harnesses, tool_manager=tool_managers) + rewards = score_episodes(wave_tasks, envs, outs, wave_dir, cfg) + finally: + for env in envs: + env.close() + if not args.keep_workspaces: + shutil.rmtree(wave_dir, ignore_errors=True) + + for (task, rep), out, reward in zip(wave, outs, rewards): + labels = out.get('labels') or [] + if msg_dump is not None: + msg_dump.write(json.dumps({ + 'id': task.get('id'), + 'rep': rep, + 'reward': reward, + 'turns': int(out.get('turns') or 0), + 'stop_reason': out.get('stop_reason'), + 'query': task.get('query'), + 'messages': out.get('messages') or [], + }, ensure_ascii=False, default=str) + '\n') + msg_dump.flush() + results.append({ + 'id': task.get('id'), + 'rep': rep, + 'reward': reward, + 'turns': int(out.get('turns') or 0), + 'stop_reason': out.get('stop_reason'), + 'truncated': bool(out.get('truncated')), + 'completion_tokens': sum(1 for label in labels if label != -100), + }) + done = len(results) + rate = sum(r['reward'] for r in results) / done if done else 0.0 + logger.info(f'[eval:{args.label}] {done}/{len(attempts)} episodes, ' + f'mean reward so far {rate:.3f}') + + if args.out: + os.makedirs(os.path.dirname(os.path.abspath(args.out)) or '.', exist_ok=True) + with open(args.out, 'w', encoding='utf-8') as f: + for row in results: + f.write(json.dumps(row, ensure_ascii=False) + '\n') + logger.info(f'[eval:{args.label}] wrote {len(results)} episodes -> {args.out}') + if msg_dump is not None: + msg_dump.close() + logger.info(f'[eval:{args.label}] conversations -> {args.dump_messages}') + + report(args, tasks, results, n_boot_failed) + + +def report(args, tasks, results, n_boot_failed): + """Print what the run measured, including what it failed to measure.""" + if not results: + logger.warning(f'[eval:{args.label}] no episodes completed; nothing to report') + return + per_task = {} + for row in results: + per_task.setdefault(row['id'], []).append(row['reward']) + + rewards = [row['reward'] for row in results] + mean_reward = statistics.fmean(rewards) + task_rates = [statistics.fmean(v) for v in per_task.values()] + solved_always = sum(1 for r in task_rates if r >= 1.0) + solved_never = sum(1 for r in task_rates if r <= 0.0) + turns = [row['turns'] for row in results] + stops = {} + for row in results: + stops[row['stop_reason']] = stops.get(row['stop_reason'], 0) + 1 + + logger.info( + f'[eval:{args.label}] === {len(results)} episodes over {len(per_task)} tasks ' + f'({args.rollouts_per_task} attempts each) ===') + logger.info(f'[eval:{args.label}] pass rate (mean reward) : {mean_reward:.4f}') + logger.info(f'[eval:{args.label}] per-task rate mean/median : ' + f'{statistics.fmean(task_rates):.4f} / {statistics.median(task_rates):.4f}') + logger.info(f'[eval:{args.label}] tasks always/never solved : ' + f'{solved_always}/{solved_never} of {len(per_task)}') + logger.info(f'[eval:{args.label}] turns mean/max : ' + f'{statistics.fmean(turns):.1f} / {max(turns)}') + logger.info(f'[eval:{args.label}] truncated episodes : ' + f'{sum(1 for r in results if r["truncated"])}') + logger.info(f'[eval:{args.label}] stop reasons : {stops}') + if n_boot_failed: + logger.warning(f'[eval:{args.label}] {n_boot_failed} episodes never ran ' + f'(sandbox boot failed) and are excluded from every number above') + if len(tasks) != len(per_task): + logger.warning(f'[eval:{args.label}] {len(tasks) - len(per_task)} of {len(tasks)} tasks ' + f'produced no episode at all') + + +if __name__ == '__main__': + main() diff --git a/cookbook/rsi/agentic/prompts.py b/cookbook/rsi/agentic/prompts.py index 8ff4356e9..23359e840 100644 --- a/cookbook/rsi/agentic/prompts.py +++ b/cookbook/rsi/agentic/prompts.py @@ -1,170 +1,284 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Prompts for the agentic challenger. -The agentic challenger works in three rounds: - Round 1: model acts as an agent in a sandbox (multi-turn with tools), - producing a tool-call chain and final workspace state. - Round 2a: model sees the trajectory and writes a python check script - that asserts properties of the final state. - Round 2b: model sees trajectory + checks and writes a problem statement - that someone else would need to follow to reproduce the result. - -The check script is run against the sandbox immediately after round 1 to verify -it passes; any task whose own checks fail is thrown away. This is the agentic -analogue of the code challenger running the reference solution against its own -asserts. - -Keyword categories and directions are configurable. The defaults below exercise: - - filesystem: creating, moving, reading, transforming files - - scripting: writing python/shell scripts that produce output - - data: CSV/JSON/text parsing and aggregation +One conversation, three stages: + Stage 1: the model acts as an agent in a sandbox (multi-turn with tools), + producing a tool-call chain and a final workspace state, and stops + calling tools. + Stage 2: a user message carrying the real workspace listing is appended to that + same conversation, asking for a python check script. + Stage 3: another user message asks for the problem statement. + +Stages 2 and 3 are appended rather than sent as fresh calls, so the whole chain +is one sample whose every assistant turn can be trained on. That is also why +their rules live in user messages: a conversation has one system message, and it +was already spent on stage 1. + +The check script is run against the sandbox immediately after stage 2 to verify +it passes; any task whose own checks fail is thrown away. + +Keyword categories are configurable; the framework (``KeywordStore`` + draw/combine +logic) is category-agnostic. A proposal draws one entry from each category, and the +three are facets of ONE task so combining them yields a single non-trivial +computation: + - transform: the computation the task turns on + - domain: the material it runs over (data AND code/compilation) + - edge_case: the twist that makes a naive solution fail """ from twinkle_agentic.challenger import AgenticPrompts -# ── Keyword categories (analogous to code's algorithm/computer/noncs) ────── +# ── Keyword categories (framework-agnostic; only the content is scenario-specific) ─ -CATEGORIES = ['filesystem', 'scripting', 'data'] +CATEGORIES = ['transform', 'domain', 'edge_case'] + +# Each category gives three examples, then pushes away from them, then pins the +# answer to what the sandbox can actually build and read back. +_LEAVE_THE_EXAMPLES = ( + '. These three are only to show the form of an answer -- do NOT stay ' + 'near them; name things from as many different areas of computer ' + 'engineering as you can') +_PINNED_TO_CONTAINER = ( + ', but only material that can actually be BUILT and read back inside a Linux ' + 'container that has python (numpy, pandas, pillow, scipy, sympy, networkx, ' + 'openpyxl, xlsxwriter, pypdf, pymupdf, pdfplumber, python-docx, python-pptx, ' + 'reportlab, lxml, pyarrow, matplotlib), sqlite3, ffmpeg, imagemagick, git, ' + 'jq, tar/zip/7z, poppler-utils and pip -- and NO compiler, no GPU, no docker, ' + 'no hardware devices. Name a FILE FORMAT or a DATA STRUCTURE, never a device ' + 'or a service') +_PINNED_COMPUTATION = ( + ', but only computations that run in that same container: not compiling, not ' + 'flashing firmware, not driving hardware') CATEGORY_DESC = { - 'filesystem': 'file and directory manipulation tasks: creating directory trees, ' - 'moving/renaming files by pattern, finding files by content, ' - 'generating structured text files', - 'scripting': 'tasks that require writing a script (python or shell) whose output ' - 'or side effects are the goal: number crunching, text transformation, ' - 'format conversion, small utilities', - 'data': 'tasks involving structured data: parsing CSV/JSON/YAML, aggregating ' - 'rows, filtering records, joining multiple files, producing summary ' - 'reports or reformatted output', + 'transform': 'a specific, non-trivial transformation the solver must COMPUTE ' + 'rather than copy -- the answer is derived, never stated. For ' + 'example: solve a system of equations with sympy, decode a ' + 'binary format, find a shortest path' + _LEAVE_THE_EXAMPLES + + _PINNED_COMPUTATION, + 'domain': 'the kind of material the task operates on. For example: WAV audio ' + 'files, a SQLite database, PNG images' + _LEAVE_THE_EXAMPLES + + _PINNED_TO_CONTAINER, + 'edge_case': 'a twist that makes a naive or copy-the-statement solution fail ' + 'and forces careful handling. For example: floating-point ' + 'rounding, byte order, cycles in a tree' + _LEAVE_THE_EXAMPLES, } -# ── Round 1: model acts in sandbox ───────────────────────────────────────── +# ── Stage 1: model acts in sandbox ───────────────────────────────────────── +# Only shell_executor and python_executor exist -- there is no directory-listing +# tool -- so the prompt names ``ls -R`` explicitly and spells tool names in full. +# One tool call per message: the sampler stops generation at ``</tool_call>``, and +# a reply that plans many calls but is cut after the first would be trained on a +# reasoning that does not match what happened. SYSTEM = ( 'You are an expert developer working in an empty directory with a shell and ' - 'python. Your job is to do something interesting and non-trivial based on ' - 'the direction given below. Use the tools available to you (shell commands, ' - 'python scripts, file operations) to produce a meaningful end state: files ' + 'python. Your job is to build something complex and realistic, one tool call ' + 'at a time, based on the direction given below. Use ' + 'the tools available to you (shell commands, python scripts, file ' + 'operations) to produce a meaningful end state: files ' 'with content, computed outputs, structured data.\n\n' 'Requirements:\n' '- Work entirely within the current directory (do not use /tmp or ~).\n' - '- Do not use the network.\n' + '- Do not use the network: no downloads, no web requests, and no installing ' + 'packages (no pip install, no apt install). The sandbox has no internet, so ' + 'any such call wastes a turn; build only with the Python standard library and ' + 'the packages already installed (numpy, pandas, matplotlib, scikit-learn, ' + 'pyarrow, and other common data libraries).\n' '- Make sure the end state is deterministic: the same steps always produce ' 'the same files with the same content.\n' - '- Do at least 2-3 distinct steps, not just one command.\n' + '- Make exactly ONE tool call per message, then read what it returned before ' + 'choosing the next one. Take as many turns as the work needs.\n' + '- Verify your own work before finishing: list the directory and read back ' + 'what you wrote. A file you meant to create but did not is worse than a ' + 'smaller result, because the task built from this state will be impossible.\n' + '- To see what is in the directory, run shell_executor with "ls -R", which ' + 'shows files and directories at every depth.\n' '- When you are satisfied with the result, stop calling tools and say ' '"Done." as your final message.' ) FROM_SCRATCH = ( - 'Do something interesting and non-trivial in the current empty directory. ' - 'Create files, write scripts, process data -- whatever demonstrates ' - 'competent use of the tools. Aim for 2-4 steps that build on each other.' + 'Build something complex and realistic in the current empty directory. ' + 'Create files, write scripts, ' + 'process data -- whatever demonstrates ' + 'competent use of the tools. Take as many turns as the work needs, each ' + 'building on the last.' ) FROM_SEED = ( 'Here is an example of the kind of task we want:\n\n{seed}\n\n' - 'Do something in the same spirit but on a different subject. Change what ' + 'Build something in the same spirit but on a ' + 'different subject, equally complex and realistic. Change what ' 'is produced and how, not just the names. Work in the current empty directory.' ) FROM_KEYWORDS = ( 'Your direction for this task:\n{keywords}\n\n' - 'Do something interesting that exercises the topics above. Work in the ' + 'Build something complex and realistic that exercises the topics above. ' + 'Work in the ' 'current empty directory, producing files and/or computed output.' ) FROM_SEED_KEYWORDS = ( 'Here is an example task for inspiration:\n\n{seed}\n\n' 'Your direction keywords:\n{keywords}\n\n' - 'Do something that combines the spirit of the example with the keyword ' - 'topics. Work in the current empty directory.' + 'Build something complex and realistic that combines the spirit of the ' + 'example with the keyword topics. ' + 'Work in the current empty directory.' ) -# ── Round 2a: write check script ─────────────────────────────────────────── +# ── Stage 2: write the check script ──────────────────────────────────────── +# Appended to the episode as a user message once the model stops calling tools, +# so it says "you": the same conversation did the work. brittle_check_reason() in +# challenger/agentic.py rejects size/checksum/source-text asserts on the syntax +# tree and sends the script back through the rewrite path. -CHECK_SYSTEM = ( - 'You are a test engineer. Given a record of what an agent did in a directory ' - 'and the resulting state, write a python script that ASSERTS properties of ' - 'the end state. The script will be run in the same directory the agent worked ' - 'in.\n\n' +CHECK_FOLLOWUP = ( + 'Now write a python script that ASSERTS properties of the state you just ' + 'produced. It will be run in the same directory you worked in.\n\n' + 'Here is the actual final state of that directory: first every file as ' + '"path size-in-bytes", then the contents of each one. This listing is the ' + 'ground truth, not your account of what you did. Assert only about paths ' + 'that appear in it, and only about content you can read here. If it is empty ' + 'or shows nothing worth testing, say UNTESTABLE and write no code.\n\n' + '{final_state}\n\n' 'Rules:\n' - '- Use only the standard library (os, json, csv, re, pathlib, etc.).\n' - '- Write 2-6 assert statements that verify the most important outcomes.\n' - '- Each assert should check something observable: file existence, file ' - 'content, computed values, directory structure.\n' - '- The script must exit 0 when all assertions hold and non-zero otherwise.\n' - '- Do NOT import anything that is not in the python standard library.\n' - '- Do NOT use the network or read from outside the working directory.\n' - '- Return ONLY a fenced python code block, no prose.' -) - -CHECK_USER = ( - 'Here is what the agent did:\n\n{trajectory}\n\n' - 'Here is the final state of the working directory:\n\n{final_state}\n\n' - 'Write a python check script (fenced code block) that asserts the key ' - 'properties of this end state. 2-6 assertions.' + '- 2-6 asserts, standard library only.\n' + '- Make the check ROBUST and BROAD: it must pass for ANY correct ' + 'reproduction of this state, and fail only for one that got the work wrong. ' + 'Assert meaning, not form -- that a file exists, that it parses, that a value ' + 'or a row read out of it is right, that an expected substring is present.\n' + '- Do NOT pin exact bytes: no file sizes, no checksums, no asserting that a ' + 'whole file equals one exact string, no timestamps, no script source text. A ' + 'different correct solution writes different bytes and would fail such a ' + 'check even though it is right.\n' + '- Do NOT constrain the directory as a whole: never assert the exact number ' + 'of files, or that no other files exist. Check only the files that carry the ' + 'result and ignore the rest.\n' + '- Never write down a number you did not read above -- do not recompute a ' + 'mean, a count or a checksum in your head.\n' + '- A file shown truncated has more content than you can see: assert about the ' + 'part you were shown, not its end or its length.\n' + '- Still discriminating: what you keep must fail for a directory that does ' + 'not hold this state. Robust does not mean empty.\n' + '- Exit 0 when every assertion holds, non-zero otherwise.\n' + '- Do NOT call any tool now. Return ONLY a fenced python code block, no ' + 'prose.' ) -# ── Round 2b: write problem statement ───────────────────────────────────── +# ── Stage 2b: the one chance to fix a check that did not pass ────────────── -PROBLEM_SYSTEM = ( - 'You write task descriptions for an AI agent. Given a record of what was ' - 'done and the checks that verify it, write a clear problem statement that ' - 'another agent would need to follow to reproduce the same end state.\n\n' - 'Rules:\n' - '- State exactly what files must exist and what they must contain.\n' - '- Be specific about formats, names, and expected values.\n' - '- Do NOT reveal the solution steps -- only describe the desired end state.\n' - '- Do NOT mention the checks or how verification works.\n' - '- The statement must be self-contained: no references to prior context.\n' - '- Keep it concise: 50-300 words.\n' - '- Return ONLY the problem statement as plain text, no code fences.' +CHECK_RETRY_FOLLOWUP = ( + 'That script does not pass. Running it in that directory gave:\n\n' + '{error}\n\n' + 'Nothing has changed in the directory; this is what it holds:\n\n' + '{final_state}\n\n' + 'Rewrite the script so that it passes. Where your assertion and this listing ' + 'disagree, the listing is what is there and the assertion is what is wrong -- ' + 'fix the assertion, do not assert something new that you still cannot read ' + 'here. Drop an assertion you cannot make true instead of weakening every one ' + 'of them; what stays must still fail for a directory that does not hold this ' + 'state.\n\n' + 'Same rules as before: standard library only, 2-6 asserts, no file sizes, ' + 'checksums, timestamps, script source text, whole-file exact-string ' + 'equality, or claims about the exact set of files in the directory. Keep it ' + 'robust -- it must pass for any correct reproduction of this state -- yet ' + 'still fail for a directory that does not hold it. Exit 0 exactly when the ' + 'state is right. Do NOT call any tool now. Return ONLY a fenced python code ' + 'block, no prose.' ) -PROBLEM_USER = ( - 'Here is what the agent did:\n\n{trajectory}\n\n' - 'Here are the check assertions that verify the end state:\n\n' - '```python\n{checks}\n```\n\n' - 'Write a problem statement (plain text, 50-300 words) describing what ' - 'another agent must produce to pass these checks. Do not reveal the ' - 'solution steps.' -) # ── Keyword generation ───────────────────────────────────────────────────── +# ``parse_keyword_list`` reads a JSON array and returns nothing when it cannot +# find one, so these must ask for a JSON array; keep them in step with the parser. KEYWORD_SYSTEM = ( - 'You generate diverse topic keywords for training an AI agent that works ' - 'with files, scripts, and data in a local directory.' + 'You generate diverse topic keywords for training an AI agent that does ' + 'computer engineering work in a Linux sandbox: writing and running programs, ' + 'building, testing and debugging software, processing and analysing data, ' + 'and administering files and the system.' ) KEYWORD_USER = ( 'List {k} diverse, specific topic keywords for the following category:\n' '{desc}\n\n' - 'Return one keyword per line, no numbering, no explanation. ' - 'Each should be 2-5 words, concrete enough to inspire a specific task.' + 'Each should be 2-5 words, concrete enough to inspire a specific task. ' + 'Return ONLY a JSON array of short strings, nothing else.' ) KEYWORD_EXPAND_USER = ( 'The keyword "{kw}" produced a very hard task. List {m} related keywords ' 'in the same domain that might produce similarly challenging but different ' - 'tasks. One per line, no numbering.' + 'tasks. Return ONLY a JSON array of short strings, nothing else.' +) + + +# ── Arm C: cap what one episode may build ────────────────────────────────── +# A cap on volume only: the failure it targets is a smaller model running out of +# tokens writing many files, and the thing that must survive is the computation. + +BUILD_SIZE_CAP = ( + '\n- Keep the result SMALL: at most {n} files in total, counting inputs, ' + 'scripts and outputs. No python package (no __init__.py, no importable ' + 'module tree), no command-line interface with subcommands. Depth, not ' + 'volume: one non-trivial computation done properly on a small input beats ' + 'many files. A task built from this state has to be finishable by a smaller ' + 'model in about twenty tool calls.' +) + +# ── Stage 3: the statement gives the rules, never the computed answer ──────── +# The end state is split in two: input data verbatim (it is not the answer), and +# everything derived given as the rule that produces it -- otherwise the only way +# to state what a derived file must contain is to quote the computed answer. + +PROBLEM_FOLLOWUP_RULES_ONLY = ( + 'Your checks pass on the state you produced. Now write the task description ' + 'another AI agent would be given to reproduce that same end state.\n\n' + 'That agent starts in an EMPTY directory and sees nothing but your ' + 'statement: every file that must be there at the end has to be created by ' + 'it.\n\n' + 'Give the two halves differently:\n' + '- INPUT data, the raw material nothing was computed from yet: verbatim, ' + 'exact filenames and exact contents, so it can be written byte for byte.\n' + '- Everything DERIVED from it -- computed values, aggregates, orderings, ' + 'resolved references, reports: only the RULE that produces it. Name the ' + 'output file and its format, say how each part follows from the input, and ' + 'never state the resulting value. Not as an example, not in a sample of the ' + 'output. A statement that writes out what you computed can be satisfied by ' + 'copying it, and then it measures typing.\n\n' + 'Rules:\n' + '- Be specific about formats, filenames and layout.\n' + '- Say what must be true of the result, not which commands to run.\n' + '- Do NOT mention the checks or how verification works.\n' + '- Self-contained: no reference to this conversation or to anything the ' + 'reader cannot see.\n' + '- 300 words or less, not counting input data quoted verbatim.\n' + '- Do NOT call any tool now. Return ONLY the problem statement as plain ' + 'text, no code fences.' ) # ── Factory ──────────────────────────────────────────────────────────────── -def agentic_prompts() -> AgenticPrompts: - """Assemble all strings into the object the challenger takes.""" +def agentic_prompts(max_build_files: int = 0) -> AgenticPrompts: + """Assemble all strings into the object the challenger takes. + + ``max_build_files`` is the one knob: when > 0 it appends BUILD_SIZE_CAP to the + system prompt. + """ + system = SYSTEM + if max_build_files > 0: + system = system + BUILD_SIZE_CAP.format(n=max_build_files) return AgenticPrompts( - system=SYSTEM, + system=system, from_scratch=FROM_SCRATCH, from_seed=FROM_SEED, from_keywords=FROM_KEYWORDS, from_seed_keywords=FROM_SEED_KEYWORDS, - check_system=CHECK_SYSTEM, - check_user=CHECK_USER, - problem_system=PROBLEM_SYSTEM, - problem_user=PROBLEM_USER, + check_followup=CHECK_FOLLOWUP, + check_retry_followup=CHECK_RETRY_FOLLOWUP, + problem_followup=PROBLEM_FOLLOWUP_RULES_ONLY, keyword_system=KEYWORD_SYSTEM, keyword_user=KEYWORD_USER, keyword_expand_user=KEYWORD_EXPAND_USER, diff --git a/cookbook/rsi/agentic/remote_tool_env.py b/cookbook/rsi/agentic/remote_tool_env.py index d3e0cbf2f..d36dde62a 100644 --- a/cookbook/rsi/agentic/remote_tool_env.py +++ b/cookbook/rsi/agentic/remote_tool_env.py @@ -18,11 +18,13 @@ ``files.write``, which is the surface every e2b-compatible backend implements the same way. """ +import copy import json import os import posixpath import re import time +import uuid from typing import Any, Dict, List, Optional, Sequence, Tuple from twinkle import get_logger @@ -36,28 +38,99 @@ _RC_MARK = '__TWINKLE_RC__' _RC_RE = re.compile(rf'{_RC_MARK}:(-?\d+)') +# The check script is passed through as a string and compiled under its own +# filename rather than indented into the `try` below. Indenting shifted every +# line by the two lines of preamble, so a traceback said "line 27" about a +# 25-line script and "line 15" about a comment -- the one piece of information a +# reader needs to see which assertion failed pointed at the wrong assertion, or +# past the end of the file. `<check>` in the traceback is that script, line for +# line, and the frames above it are this wrapper's. _PY_WRAPPER = """\ -import sys, traceback +import sys, io, traceback +_tw_check_src = {body} +_tw_buf = io.StringIO() +_tw_out = sys.stdout +sys.stdout = _tw_buf +_tw_rc = 0 try: -{body} + _tw_ns = {{'__name__': '__main__'}} + exec(compile(_tw_check_src, '<check>', 'exec'), _tw_ns, _tw_ns) except SystemExit as _e: # Print the status and stop -- do NOT re-raise. SystemExit inherits from # BaseException, so ms-agent's `except Exception` around the exec does not # catch it; letting it escape kills the whole tool server process, and the - # sandbox is shared by every task in the run. `else` is already skipped - # because the exception was handled, so nothing further is needed. + # sandbox is shared by every task in the run. _c = _e.code - print('{mark}:%d' % (0 if _c is None else _c if isinstance(_c, int) else 1)) + _tw_rc = 0 if _c is None else _c if isinstance(_c, int) else 1 except BaseException: - traceback.print_exc() - print('{mark}:1') -else: - print('{mark}:0') + traceback.print_exc(file=_tw_buf) + _tw_rc = 1 +finally: + sys.stdout = _tw_out +# Marker FIRST, then the body. The executor truncates a tool observation at +# ~8KB, counted from the start; a marker printed after a large body (a rich +# workspace snapshot, say) is silently cut off, `runner` then finds no marker +# and reports exit 1 -- which read as an empty workspace and threw the task +# away. Emitted before the body, the marker always survives; only the tail of +# the body is ever lost. +print('{mark}:%d' % _tw_rc) +sys.stdout.write(_tw_buf.getvalue()) """ _REMOTE_DIR = '/opt/rsi' +# Where the in-sandbox runtime's stdout/stderr goes. Read back by `server_log`. +SERVER_LOG = '/tmp/tool_server.log' +# Seconds the transport gets beyond the server's own budget, so that a slow call +# is answered by the layer that knows which call was slow. curl waits this much +# longer than the server may spend, and the command channel that much again. +# Anything smaller than the gap between two deadlines is a race, and the client +# wins it -- which turns one slow call into "runtime unreachable" for the whole +# turn. +_RPC_HEADROOM = 60 _LOCAL_SERVER = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'sandbox_server', 'tool_server.py') +# ms-agent's permission package, uploaded alongside the yaml because this +# repository's copy carries two switches the released code does not have: +# ``safety_rules.unrestricted_removal`` and ``safety_rules.allow_write_globs``, +# which rsi_agent.yaml turns on. Without them the sandbox reads those keys, +# finds no code that looks at them, and silently keeps refusing `rm -rf *`, +# `cp src/* dst/` and `chmod +x bin/*` -- a whole run's worth of tasks shaped by +# a config that never took effect. +# +# The image installs ms-agent editable from a tarball into /opt/ms-agent, so +# replacing the files under it is what the interpreter picks up; and it has to +# land before tool_server.py imports them, which is why this is part of the same +# upload rather than a rebuilt image. +_MS_AGENT_PERMISSION_DIR = '/opt/ms-agent/ms_agent/permission' +_PATCHED_PERMISSION_FILES = ('config.py', 'safety.py', 'shell_validator.py', + 'path_validator.py') + + +def tool_payload(observation: str) -> str: + """The command output inside an ms-agent tool observation, or the text as-is. + + The executor tools answer with a JSON envelope + (``{"success": ..., "output": ..., "error": ...}``). Fed to a model as-is it + reads as a wall of metadata around the one part that matters, and a model + asked to describe a directory from it tends to trust its own recollection + instead. Callers that want the output *as data* -- a file listing, a computed + value -- go through here; callers that only want an exit status do not need it. + """ + text = (observation or '').strip() + if not text.startswith('{'): + return observation or '' + try: + body = json.loads(text) + except ValueError: + return observation or '' + if not isinstance(body, dict) or 'output' not in body: + return observation or '' + payload = body.get('output') or '' + error = body.get('error') + if error: + payload = f'{payload}\n{error}'.strip() + return payload + class RemoteMsAgentToolEnv(Env): """Run one episode's ms-agent tool calls inside a dedicated sandbox. @@ -72,7 +145,9 @@ class RemoteMsAgentToolEnv(Env): workspace: ``config.output_dir`` inside the sandbox. sandbox_timeout: sandbox idle timeout, in seconds. Must outlast a whole episode plus the checks that run after it. - command_timeout: per-request timeout for a tool call, in seconds. + command_timeout: how long the in-sandbox server may spend on a turn's + tool calls, in seconds. The transport around it is given headroom on + top -- see :meth:`_rpc`. boot_timeout: how long to wait for the runtime to answer ``/health``. ms-agent's import plus tool construction dominates this. max_observation_chars: truncate a tool result before it becomes a @@ -110,6 +185,12 @@ def __init__( self.max_observation_chars = max_observation_chars self._sandbox = None self._schemas: Optional[List[Dict[str, Any]]] = None + # Short advertised name -> the runtime's namespaced one. Filled by + # _load_schemas, which is the only thing that knows which short names are + # unambiguous. + self._short_to_full: Dict[str, str] = {} + self._deadline = 0.0 + self.n_recoveries = 0 # ------------------------------------------------------------------ Env @@ -117,12 +198,44 @@ def reset(self, trajectory: Optional[Dict[str, Any]] = None) -> StepResult: """Boot a sandbox and bring ms-agent's tool runtime up inside it.""" self.close() self._sandbox = self._create_sandbox() + self._deadline = time.time() + self._sandbox_timeout self._upload() self._start_server() self._await_ready() self._schemas = None + self._short_to_full = {} return StepResult(observation='') + def healthy(self) -> bool: + """Does the tool runtime answer right now?""" + if self._sandbox is None: + return False + try: + return (self._rpc('/health', None, timeout=10) or {}).get('status') == 'ok' + except Exception: # noqa + return False + + def ensure_ready(self) -> bool: + """Re-establish the sandbox if its runtime has gone away. True if it did. + + For the callers that hold one sandbox across many episodes, losing it -- + evicted, timed out, runtime crashed -- otherwise ends the whole run. This + is safe to call only where the workspace is about to be discarded + anyway: a mid-episode rebuild would silently swap the state the episode + is being judged on for an empty directory, so recovery is offered as an + explicit call rather than a retry hidden inside every tool dispatch. + + Recoveries are counted in ``n_recoveries`` so a run can report how often + this happened instead of hiding it. + """ + if self.healthy(): + return False + self.n_recoveries += 1 + logger.warning(f'tool runtime unreachable; rebuilding the sandbox ' + f'(recovery #{self.n_recoveries})') + self.reset() + return True + def step(self, tool_name: str, arguments: Dict[str, Any] = None) -> StepResult: return self.step_batch([(tool_name, arguments or {})])[0] @@ -138,7 +251,7 @@ def step_batch(self, calls: Sequence[Tuple[str, Dict[str, Any]]]) -> List[StepRe return [] payload = { 'calls': [{ - 'tool_name': name, + 'tool_name': self._dispatch_name(name), 'arguments': args or {} } for name, args in calls], 'timeout': self._command_timeout, @@ -165,9 +278,83 @@ def close(self) -> None: logger.warning(f'failed to kill sandbox {sandbox_id}: {e}') finally: self._sandbox = None + self._deadline = 0.0 + + def _keep_alive(self) -> None: + """Push the sandbox's expiry back while it is still being used. + + ``sandbox_timeout`` is a lifetime from creation, not an idle timer, so a + caller that keeps one sandbox for a long run would lose it mid-run no + matter how busy it was. Extended once past the halfway mark rather than on + every call: this is an extra HTTP round trip, and tool dispatch is already + the slow part of a turn. + """ + if self._sandbox is None: + return + if time.time() < self._deadline - self._sandbox_timeout / 2: + return + try: + self._sandbox.set_timeout(self._sandbox_timeout) + self._deadline = time.time() + self._sandbox_timeout + except Exception as e: # noqa + logger.warning(f'failed to extend sandbox timeout: {e}') # ------------------------------------------------------------ tool names + def _load_schemas(self) -> None: + """Fetch the runtime's schemas and shorten the names it advertises. + + ms-agent namespaces every tool as ``{server}---{tool}``, and a 4B policy + spends calls on that prefix: across three arms it wrote a bare + ``shell_executor`` 7 times, each one refused with "unknown tool ... Did + you mean 'code_executor---shell_executor'?" -- a whole turn burnt on + punctuation. Since the prefix carries no information the model can act + on (nothing here has two servers offering the same tool), the advertised + name drops it, and :meth:`step_batch` puts it back before dispatch. + + This is not the same as accepting a wrong name and fixing it up: the + model is shown ``shell_executor`` and calls ``shell_executor``, so what + it learns to emit is what the schema promised. A name that would collide + keeps its prefix, in both directions, rather than becoming ambiguous. + """ + raw = list((self._rpc('/tools', None) or {}).get('tools') or []) + full_names = [(t.get('function') or {}).get('name') for t in raw] + counts: Dict[str, int] = {} + for name in full_names: + if name: + counts[str(name).rsplit('---', 1)[-1]] = counts.get( + str(name).rsplit('---', 1)[-1], 0) + 1 + self._short_to_full = {} + schemas = [] + for schema in raw: + schema = copy.deepcopy(schema) + fn = schema.get('function') or {} + full = str(fn.get('name') or '') + short = full.rsplit('---', 1)[-1] + if full and counts.get(short) == 1 and short != full: + fn['name'] = short + self._short_to_full[short] = full + schemas.append(schema) + self._schemas = schemas + + def _dispatch_name(self, name: str) -> str: + """The runtime's own spelling for a name taken from a tool call. + + Usually the map is already there, because the schemas were advertised + before anything could be called. When it is not, fetching it must not be + able to raise: a dead sandbox has to come back through ``step_batch`` as + an observation the episode survives, not as an exception from name + lookup. An unmapped name passes through as-is, which is also what a + caller using the runtime's full spelling wants. + """ + if not self._short_to_full and self._schemas is None: + try: + self._load_schemas() + except Exception as e: # noqa + logger.warning(f'could not load tool names for dispatch: {e}') + return name + return self._short_to_full.get(name, name) + def tool_schemas(self) -> List[Dict[str, Any]]: """Schemas from the runtime that will execute them. @@ -176,7 +363,7 @@ def tool_schemas(self) -> List[Dict[str, Any]]: for the advertised contract and the running code to disagree. """ if self._schemas is None: - self._schemas = list((self._rpc('/tools', None) or {}).get('tools') or []) + self._load_schemas() return list(self._schemas) def tool_names(self) -> List[str]: @@ -188,10 +375,13 @@ def tool_names(self) -> List[str]: return names def resolve_tool(self, name: str) -> str: - """Map a plain tool name onto the runtime's own spelling. + """Map any spelling of a tool onto the one this Env advertises. + + Advertised names are short (see :meth:`_load_schemas`), so this returns + ``shell_executor``, not ``code_executor---shell_executor``. Both spellings + go in: callers written before the names were shortened pass the + namespaced one, and a stale spelling should not be the thing that fails. - ms-agent namespaces its tools as ``{server}---{tool}``, so a caller that - asks for ``shell_executor`` means ``code_executor---shell_executor``. An unknown name raises instead of being passed through: a mistyped tool comes back as a failed call, which for a checker is indistinguishable from a failed check, and a whole GRPO group would silently score zero. @@ -199,6 +389,10 @@ def resolve_tool(self, name: str) -> str: names = self.tool_names() if name in names: return name + # A namespaced name for a tool advertised short. + suffix = name.rsplit('---', 1)[-1] + if suffix in names and self._short_to_full.get(suffix) == name: + return suffix matches = [n for n in names if n.rsplit('---', 1)[-1] == name] if len(matches) == 1: return matches[0] @@ -221,8 +415,7 @@ def runner(self, shell_tool: str = 'shell_executor', python_tool: str = 'python_ def _run(source: str, interpreter: str) -> Tuple[int, str]: if interpreter == 'python': - body = '\n'.join(' ' + line for line in source.splitlines()) or ' pass' - code = _PY_WRAPPER.format(body=body, mark=_RC_MARK) + code = _PY_WRAPPER.format(body=repr(source), mark=_RC_MARK) out = self.step(python_name, {'code': code}).observation else: out = self.step(shell_name, {'command': f'{source}\necho "{_RC_MARK}:$?"'}).observation @@ -292,7 +485,7 @@ def _create_sandbox(self): return Sandbox.create(template=self._template, timeout=self._sandbox_timeout) def _upload(self) -> None: - """Push the yaml and the server script into the sandbox. + """Push the yaml, the server script and the permission patch into the sandbox. Uploading beats baking them into the image: the training host's copy is authoritative, so editing a tool line-up is a restart rather than a @@ -302,20 +495,83 @@ def _upload(self) -> None: self._sandbox.files.write(f'{_REMOTE_DIR}/rsi_agent.yaml', f.read()) with open(_LOCAL_SERVER, encoding='utf-8') as f: self._sandbox.files.write(f'{_REMOTE_DIR}/tool_server.py', f.read()) + self._upload_permission_patch() + + def _upload_permission_patch(self) -> None: + """Overwrite ms-agent's permission package with this repository's copy. + + Taken from the installed package rather than a path built out of + ``__file__``, so what lands in the sandbox is the same code the training + host imports. + + Verified rather than assumed: a silent miss here does not fail anything, + it just leaves the sandbox refusing commands the yaml said to allow, and + the only symptom would be a run whose tasks are quietly narrower than + intended. If the switch is not readable afterwards, this raises. + """ + import ms_agent.permission as _perm + + local_dir = os.path.dirname(os.path.abspath(_perm.__file__)) + for name in _PATCHED_PERMISSION_FILES: + with open(os.path.join(local_dir, name), encoding='utf-8') as f: + self._sandbox.files.write(f'{_MS_AGENT_PERMISSION_DIR}/{name}', f.read()) + probe = ('python -c "from ms_agent.permission.config import SafetyConfig as S; ' + 'print(S().unrestricted_removal, S().allow_write_globs)"') + result = self._sandbox.commands.run(probe, timeout=60) + out = (getattr(result, 'stdout', '') or '').strip() + if out.split() != ['False', 'False']: + raise RuntimeError( + 'permission patch did not land in the sandbox: expected the two ' + f'switches to exist and default to False, got {out!r}. The ' + f'sandbox may install ms-agent somewhere other than ' + f'{_MS_AGENT_PERMISSION_DIR}.') def _start_server(self) -> None: - command = (f'python {_REMOTE_DIR}/tool_server.py ' + """Launch the tool runtime in the background, with its output on disk. + + ``background=True`` is what detaches it; the redirect is what makes a + later death diagnosable. Without the redirect the output lives on a + command handle nobody keeps, so a runtime that dies mid-run reads only as + a refused connection. Do not swap the redirect for a trailing ``&``: + ``commands.run`` then waits out its own timeout instead of returning. + + ``-u`` rather than relying on the image: the template is built from a + snapshot of a live sandbox, which keeps the filesystem but not the image + config, so the Dockerfile's ``ENV PYTHONUNBUFFERED=1`` is not there. An + unflushed buffer is the difference between a readable log and an empty + one when the runtime dies. + + ``cd`` into the workspace, because ``python_executor`` runs ``exec()`` + inside this process (ms-agent's local_code_executor.py:657) rather than in + a subprocess with its own cwd. Started from ``/``, as it was, a relative + path in model code resolved against ``/`` while every other tool resolves + against the workspace: measured in a live sandbox, ``write_file + 'a.txt'`` answered "Save file successfully" and the next python call got + ``[Errno 2] No such file or directory: 'a.txt'``, with the file sitting in + ``/workspace`` and python looking in ``/``. That single mismatch is 41 of + ex7's 58 such failures, and it also hid files from the end-of-episode + snapshot, which only lists the workspace. The python_executor patch in + tool_server.py chdirs per call as well, so the two do not depend on each + other. + """ + command = (f'mkdir -p {self.workspace} && cd {self.workspace} && ' + f'python -u {_REMOTE_DIR}/tool_server.py ' f'--config {_REMOTE_DIR}/rsi_agent.yaml ' - f'--workspace {self.workspace} --port {self._port}') + f'--workspace {self.workspace} --port {self._port} ' + f'> {SERVER_LOG} 2>&1') + self._sandbox.commands.run(command, background=True) + + def server_log(self, lines: int = 40) -> str: + """Tail the in-sandbox runtime log; '' if it cannot be read. + + Used when the runtime stops answering, which is the one moment its own + output matters and the one moment an RPC cannot fetch it. + """ try: - self._sandbox.commands.run(command, background=True) - except TypeError: - # Older SDKs have no `background`; detach with setsid so the server - # outlives the command that launched it. - self._sandbox.commands.run( - f'mkdir -p {self.workspace} && setsid nohup {command} ' - f'> /tmp/tool_server.log 2>&1 < /dev/null &', - timeout=30) + return (self._sandbox.commands.run(f'tail -n {lines} {SERVER_LOG}', + timeout=20).stdout or '') + except Exception: # noqa + return '' def _await_ready(self) -> None: """Poll ``/health`` until the runtime answers, then fail loudly. @@ -335,7 +591,7 @@ def _await_ready(self) -> None: time.sleep(2) log = '' try: - log = (self._sandbox.commands.run('tail -n 40 /tmp/tool_server.log', timeout=20).stdout or '') + log = self.server_log(40) except Exception: # noqa pass raise RuntimeError(f'ms-agent tool runtime did not come up within {self._boot_timeout}s ' @@ -347,16 +603,35 @@ def _rpc(self, path: str, payload: Optional[Dict[str, Any]], timeout: Optional[i The body is written to a file rather than inlined: tool arguments carry arbitrary source code, and no amount of shell quoting survives that reliably. + + The file name carries a nonce because two threads can be in here at once. + A fixed ``request.json`` made them overwrite each other between the write + and the curl, so every concurrent call executed whichever payload landed + last and each caller filed that one answer under its own call. That is + how ex4's episode 8 came back with a glob listing as the result of a + python script it never ran. + + curl is given ``_RPC_HEADROOM`` seconds more than the server is allowed to + spend, and the command channel more again. They used to share one number, + which meant that when a call ran long the client gave up in the same + second the server was formulating its answer -- and the client wins that + race, so a turn holding one slow call came back as "Tool runtime + unreachable" for *every* call in it, including the ones that had finished. + ex8's episode 23 is that: a shell command started an HTTP server, and the + write_file beside it was reported as an unreachable runtime. With headroom + the server's own per-call timeout message arrives instead. """ seconds = timeout or self._command_timeout + self._keep_alive() if payload is None: - command = f'curl -sS -m {seconds} http://127.0.0.1:{self._port}{path}' + command = f'curl -sS -m {seconds + _RPC_HEADROOM} http://127.0.0.1:{self._port}{path}' else: - request = f'{_REMOTE_DIR}/request.json' + request = f'{_REMOTE_DIR}/request-{uuid.uuid4().hex}.json' self._sandbox.files.write(request, json.dumps(payload, ensure_ascii=False)) - command = (f'curl -sS -m {seconds} -X POST -H "Content-Type: application/json" ' - f'--data-binary @{request} http://127.0.0.1:{self._port}{path}') - result = self._sandbox.commands.run(command, timeout=seconds + 30) + command = (f'curl -sS -m {seconds + _RPC_HEADROOM} -X POST -H "Content-Type: application/json" ' + f'--data-binary @{request} http://127.0.0.1:{self._port}{path}; ' + f'rm -f {request}') + result = self._sandbox.commands.run(command, timeout=seconds + 2 * _RPC_HEADROOM) stdout = (getattr(result, 'stdout', '') or '').strip() if not stdout: raise RuntimeError(f'empty response from {path}: {getattr(result, "stderr", "")}') diff --git a/cookbook/rsi/agentic/rl.py b/cookbook/rsi/agentic/rl.py index ab693e126..49de52a35 100644 --- a/cookbook/rsi/agentic/rl.py +++ b/cookbook/rsi/agentic/rl.py @@ -29,11 +29,9 @@ ``check_script`` (a python script, from ``challenge.py``) or ``checks`` (structured, see twinkle_agentic.verifier.result_check.Check). """ -import json import os import shutil -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List from peft import LoraConfig @@ -48,16 +46,14 @@ from twinkle.processor import InputProcessor from twinkle.sampler import vLLMSampler from twinkle.template import Template -from twinkle_agentic.envs import EnvTool -from twinkle_agentic.harness import MsAgentHarness from twinkle_agentic.rollout.multi_turn import MultiTurnRollout -from twinkle_agentic.tools.tool_manager import ToolManager -from twinkle_agentic.verifier.result_check import (CheckContext, checks_from_dicts, - run_checks) -# Same directory as this script, which python puts on sys.path when it is run -# as a file. Kept in cookbook because it is RSI-specific wiring, not framework. -from remote_tool_env import RemoteMsAgentToolEnv # noqa: I100,I202 +# Same directory as this script, which python puts on sys.path when it is run as +# a file. Episode construction and scoring are shared with eval.py so the two +# cannot drift: an eval measuring episodes built differently from training would +# not be measuring the training. +from episode import (SandboxConfig, boot_episodes, load_tasks, # noqa: I100,I202 + score_episodes) logger = get_logger() args = CLI.from_args() @@ -70,7 +66,11 @@ NUM_GPUS = MODEL_GPUS + SAMPLER_GPUS NUM_GENERATIONS = args.rl.num_generations or 8 -MAX_NEW_TOKENS = args.sampling.max_tokens or 4096 +# Per turn, and it has to match the challenger and eval side (both 8192): a +# trajectory generated with less room than the tasks were built with is trained on +# truncated attempts. At 4096, replies ran out mid-<think> before dispatching any +# tool -- 3 of 12 episodes on the generation side, 15 of 50 solver attempts. +MAX_NEW_TOKENS = args.sampling.max_tokens or 8192 LEARNING_RATE = args.optimizer.learning_rate or 1e-5 MAX_STEPS = args.training.max_steps or 1000 BATCH_SIZE = args.training.batch_size or 4 @@ -86,180 +86,16 @@ MAX_TURNS = int(os.environ.get('RSI_MAX_TURNS', 20)) TASKS_PATH = os.environ.get('RSI_TASKS', 'cookbook/rsi/agentic/tasks.example.jsonl') -AGENT_CONFIG = os.environ.get('RSI_AGENT_CONFIG', 'cookbook/rsi/agentic/rsi_agent.yaml') RUN_DIR = os.environ.get('RSI_RUN_DIR', 'output/rsi_agentic/run') -# Sandbox backend. The template is built by sandbox_server/install.sh. -SANDBOX_TEMPLATE = os.environ.get('AENV_TEMPLATE', 'twinkle-rsi-msagent') -SANDBOX_API_URL = os.environ.get('AENV_API_URL', 'http://127.0.0.1:8000') -# Must outlast a whole episode plus the checks that run after it. -SANDBOX_TIMEOUT = int(os.environ.get('RSI_SANDBOX_TIMEOUT', 900)) -# Booting and scoring are network-bound, so they are done on threads. This caps -# how many sandboxes are talked to at once, not how many exist. -ENV_CONCURRENCY = int(os.environ.get('RSI_ENV_CONCURRENCY', 16)) - -# 'fraction' gives partial credit per check; 'all_or_nothing' is stricter and -# produces a cleaner pass/fail signal at the cost of a sparser reward. -SCORE_MODE = os.environ.get('RSI_SCORE_MODE', 'fraction') +# Where episodes run, and how they are scored. Shared with eval.py. +SANDBOX = SandboxConfig.from_env() + # Keep each episode's downloaded files after scoring. Useful while debugging # tasks, expensive over a long run. KEEP_WORKSPACES = os.environ.get('RSI_KEEP_WORKSPACES', '0') == '1' -def load_tasks(path: str) -> List[Dict[str, Any]]: - """Read the task file and fail loudly on a task that can never be scored. - - Supports two formats: - - Structured checks: ``{"checks": [{"kind": ..., ...}]}`` (legacy) - - Script checks: ``{"check_script": "assert ..."}`` (from agentic challenger) - """ - tasks = [] - with open(path, encoding='utf-8') as f: - for lineno, line in enumerate(f, 1): - if not line.strip(): - continue - task = json.loads(line) - if not task.get('query'): - raise ValueError(f'{path}:{lineno} has no query') - if task.get('check_script'): - # New format: raw python script, scored by exit code - task['_checks'] = None - elif task.get('checks'): - # Legacy format: structured Check dicts - task['_checks'] = checks_from_dicts(task['checks']) - else: - raise ValueError(f'{path}:{lineno} ({task.get("id")}) declares no checks ' - f'and no check_script') - tasks.append(task) - if not tasks: - raise ValueError(f'{path} contains no tasks') - return tasks - - -def build_episode(task: Dict[str, Any]) -> Tuple[Any, Any, Any, Dict]: - """Create one episode: a sandbox with ms-agent's tools, plus a local harness. - - The harness is stripped down to message shaping. Popping ``tools`` matters - as much as popping ``llm``, and for the same reason omitting the section - from the yaml is not enough: ms-agent merges its own agent.yaml underneath, - which declares file_system and code_executor, so a live shell executor would - otherwise be constructed on the training host with access to the whole - machine. Popping them after the merge leaves the harness with zero tools -- - and the system prompt byte-identical, because ms-agent does not fold the - tool list into it. - """ - from omegaconf import OmegaConf, open_dict - - cfg = OmegaConf.load(AGENT_CONFIG) - harness = MsAgentHarness(config=cfg) - with open_dict(harness.agent.config): - harness.agent.config.pop('llm', None) - harness.agent.config.pop('tools', None) - harness.prepare() - - env = RemoteMsAgentToolEnv( - template=SANDBOX_TEMPLATE, - config_path=AGENT_CONFIG, - api_url=SANDBOX_API_URL, - sandbox_timeout=SANDBOX_TIMEOUT, - ) - env.reset() - - trajectory = harness.start(task['query']) - # The executor's own schemas, not the harness's (which are now empty by - # construction). Advertising what will run is the whole point of sourcing - # them from the sandbox. - schemas = env.tool_schemas() - trajectory['tools'] = schemas - tool_manager = ToolManager(EnvTool.from_schemas(env, schemas)) - return harness, env, tool_manager, trajectory - - -def boot_episodes(tasks: List[Dict[str, Any]]) -> List[Tuple[Any, Any, Any, Dict]]: - """Bring up every rollout's sandbox at once, all-or-nothing. - - Serial boot would dominate the step: a microVM plus ms-agent's import runs - to seconds, multiplied by ``batch_size x num_generations``. - - All-or-nothing because GRPO groups here are positional -- advantages are - taken over consecutive runs of ``NUM_GENERATIONS`` -- so dropping one - episode would not shrink its group, it would shift every later group onto - the wrong task. - """ - episodes: List[Optional[Tuple[Any, Any, Any, Dict]]] = [None] * len(tasks) - error: Optional[BaseException] = None - with ThreadPoolExecutor(max_workers=ENV_CONCURRENCY) as pool: - futures = {pool.submit(build_episode, task): slot for slot, task in enumerate(tasks)} - for future in as_completed(futures): - try: - episodes[futures[future]] = future.result() - except Exception as e: # noqa - error = error or e - if error is not None: - for episode in episodes: - if episode is not None: - episode[1].close() - raise RuntimeError(f'sandbox boot failed: {error}') from error - return episodes # type: ignore[return-value] - - -def score_episode(task: Dict[str, Any], env: RemoteMsAgentToolEnv, trajectory: Dict[str, Any], - snapshot_dir: str) -> float: - """Run the task's checks against the state this episode left behind. - - Supports two scoring paths: - - check_script: run a python script in the sandbox; exit 0 = score 1.0 - - structured checks: download workspace + run_checks (legacy) - """ - check_script = task.get('check_script') - if check_script: - # New path: run the script directly in the sandbox's python executor - runner = env.runner() - exit_code, output = runner(check_script, 'python') - if exit_code != 0: - logger.debug(f'[{task["id"]}] check_script failed (exit {exit_code}): ' - f'{output[-200:]}') - return 1.0 if exit_code == 0 else 0.0 - - # Legacy path: structured checks - final_answer = '' - for msg in reversed(trajectory.get('messages') or []): - if msg.get('role') == 'assistant' and (msg.get('content') or '').strip(): - final_answer = msg['content'] - break - - ctx = CheckContext( - workspace=env.download_workspace(snapshot_dir), - final_answer=final_answer, - runner=env.runner(), - ) - report = run_checks(task['_checks'], ctx, mode=SCORE_MODE) - if not report.all_passed: - logger.debug(f'[{task["id"]}] {report.n_passed}/{report.n_total} checks: ' - f'{report.failures()}') - return report.score - - -def score_episodes(tasks: List[Dict[str, Any]], envs: List[RemoteMsAgentToolEnv], - outs: List[Dict[str, Any]], step: int) -> List[float]: - """Score every episode in parallel; a scoring crash costs one reward, not the step. - - Each check is a sandbox round trip, so scoring serially would idle the GPUs - for as long as booting did. An episode whose sandbox died mid-check scores - zero, which is also what it would have scored had the checks simply failed. - """ - - def _score(slot: int) -> float: - snapshot = os.path.join(RUN_DIR, f'step{step:06d}', f'slot{slot:03d}') - try: - return score_episode(tasks[slot], envs[slot], outs[slot], snapshot) - except Exception as e: # noqa - logger.warning(f'[step {step} slot {slot}] scoring failed: {e}') - return 0.0 - - with ThreadPoolExecutor(max_workers=ENV_CONCURRENCY) as pool: - return list(pool.map(_score, range(len(outs)))) - def main(): tasks = load_tasks(TASKS_PATH) @@ -288,7 +124,12 @@ 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('Template', model_id=MODEL_ID, enable_thinking=True) + # Both templates get the trajectory budget explicitly. The default is far + # smaller than a tool-using episode: leaving it out makes the sampler refuse + # the trajectory mid-run with `Input length N exceeds max_length 8192`, after + # the step it happened in has already booted its sandboxes. + model.set_template('Template', model_id=MODEL_ID, enable_thinking=True, + max_length=MAX_TRAJECTORY_TOKENS) sampler = vLLMSampler( model_id=MODEL_ID, @@ -302,7 +143,8 @@ def main(): device_mesh=sampler_mesh, remote_group='sampler', ) - sampler.set_template('Template', model_id=MODEL_ID, enable_thinking=True) + sampler.set_template('Template', model_id=MODEL_ID, enable_thinking=True, + max_length=MAX_TRAJECTORY_TOKENS) rollout_template = Template(MODEL_ID, max_length=MAX_TRAJECTORY_TOKENS, enable_thinking=True) rollout_template.truncation_strategy = 'delete' @@ -323,9 +165,10 @@ def main(): optim_step = 0 task_cursor = 0 - logger.info(f'Starting agentic RSI GRPO (max_turns={MAX_TURNS}, score={SCORE_MODE})') - logger.info(f'Sandboxes: template={SANDBOX_TEMPLATE} api={SANDBOX_API_URL} ' - f'concurrency={ENV_CONCURRENCY}') + logger.info(f'Starting agentic RSI GRPO (max_turns={MAX_TURNS}, ' + f'score={SANDBOX.score_mode})') + logger.info(f'Sandboxes: template={SANDBOX.template} api={SANDBOX.api_url} ' + f'concurrency={SANDBOX.concurrency}') logger.info(get_device_placement()) while optim_step < MAX_STEPS: @@ -339,7 +182,7 @@ def main(): harnesses, envs, tool_managers, trajectories = [], [], [], [] try: - episodes = boot_episodes(episode_tasks) + episodes = boot_episodes(episode_tasks, SANDBOX) except Exception as e: # noqa # A sandbox that never came up answers every call with an error, so # the group would score a uniform zero and look like a hard task @@ -359,7 +202,8 @@ def main(): outs: List[Dict[str, Any]] = rollout( trajectories, harness=harnesses, tool_manager=tool_managers) - rewards = score_episodes(episode_tasks, envs, outs, optim_step) + rewards = score_episodes(episode_tasks, envs, outs, + os.path.join(RUN_DIR, f'step{optim_step:06d}'), SANDBOX) finally: # Sandboxes are a finite resource; a step that raises must still # give them back or the next step starts short. diff --git a/cookbook/rsi/agentic/rsi_agent.yaml b/cookbook/rsi/agentic/rsi_agent.yaml index 63ca8fd1d..5f058f686 100644 --- a/cookbook/rsi/agentic/rsi_agent.yaml +++ b/cookbook/rsi/agentic/rsi_agent.yaml @@ -12,10 +12,46 @@ # effect on the next episode; no image rebuild is involved. prompt: - # Unset -> ms-agent's built-in agent prompt, which is also what serving uses. - # Set a string here only if training should see a different system prompt - # than production, and know that you are breaking train/serve parity. - system: + # Replaces ms-agent's BASE_AGENT_PROMPT (prompting/builtin.py) for the SOLVER + # only -- the proposing episode gets prompts.py's own SYSTEM through the + # challenger, and never reads this field. That built-in prompt is written for a + # general assistant sitting in a user's workspace, and two of its lines work + # against being a solver: "First decide whether the task needs tools. If you can + # answer reliably from what you know ... just answer", and "Ask first when it + # isn't [safe]". Here there is no one to ask (interactive: false) and answering + # without touching the directory is always wrong. + # + # The paragraph about the empty directory is what 5 of armA2shellV5's 8 + # unsolved tasks needed. Their statements listed a file under "Input data:" + # and the solver read that as "already present" -- in 5a70b77f it created the + # file the rules told it to generate and left the two listed as input alone, + # so it was not confused about being in an empty directory, it was following + # the statement's own division of labour. Nothing in the statement or the + # prompt said that division does not survive into its workspace. + system: | + You are a command-line agent working inside a fresh Linux container. You are + given one task and you carry it out by running commands and writing files. + Nobody is watching and nobody can answer a question, so never ask one and + never stop to confirm: decide and act. + + Your working directory starts COMPLETELY EMPTY. Every file the task + mentions -- including files it describes as inputs, given data, existing + configuration, or material you are handed -- does not exist yet. You have to + create all of them yourself, with exactly the names and contents the task + specifies, before anything can read them. A task that shows you the contents + of a file is telling you what to write into it, not telling you it is there. + + How to work: + - Start by listing the directory to see the real state. Do not assume. + - Create every file the task names. Then do the computation it asks for and + write the results it asks for. + - Answering in prose without creating files is a failure, however clearly you + can describe what the answer would be. + - Before you finish, list the directory again and read back what you wrote. + Check each thing the task asked for is actually there. If something is + missing, fix it rather than reporting success. + - Never invent a value you did not compute. If a number has to come out of + the data, compute it from the data. personalization: # Off: SOUL/AGENTS/PROFILE.md from the developer's own workspace would leak @@ -30,6 +66,17 @@ max_chat_round: 9999 interactive: false permission_mode: auto +# How long ms-agent waits around one tool call. Written down rather than left to +# its default (tool_manager.py TOOL_CALL_TIMEOUT, 120s, overridable by the +# TOOL_CALL_TIMEOUT environment variable) so the sandbox does not inherit a +# number from whatever shell started it. It has to stay below what +# remote_tool_env allows the whole turn (command_timeout, 180s), which in turn is +# below the transport's budget: the innermost layer should be the one that times +# out, because it is the only one that knows which call was slow. When they were +# equal, one command that never returns made every call in the turn read as an +# unreachable runtime. +tool_call_timeout: 120 + # Path *inside the sandbox*. One microVM per episode already isolates # trajectories from each other, so this is a fixed path rather than a per-slot # directory; the entry script overrides it only to match --workspace. @@ -38,14 +85,20 @@ output_dir: /workspace callbacks: [] tools: - file_system: - mcp: false - include: - - write_file - - read_file - - edit_file - - grep - - glob + # `file_system` is NOT listed here and is nevertheless on. ms-agent's own + # ms_agent/agent/agent.yaml declares it (write_file, read_file, edit_file, + # grep, glob) and LLMAgent merges this file *over* that one, so omitting a key + # inherits it rather than dropping it. Measured: the merged config's tools are + # ['file_system', 'code_executor', 'todo_list'], and /tools advertises all ten + # of those tools to the model. In armA2shellV6's 128 proposing calls, + # file_system took 63 (43 of them write_file) against code_executor's 58. + # + # So the paragraph that used to be here -- claiming the five were removed to + # stop write_file being the path of least resistance -- described a state that + # never existed, through the arms named A2shell*, whose whole premise was + # "shell and python only". Turning it off takes an explicit + # `file_system: {enabled: false}`, which _tool_on (tool_manager.py:47) reads. + # Left on for now, deliberately and with the effect known. code_executor: mcp: false # python_env means "run in this process's machine", and that machine is the @@ -55,10 +108,58 @@ tools: implementation: python_env include: - shell_executor + # Kept alongside the shell so that writing a file does not depend on + # getting a heredoc right. Dropping notebook_executor because it overlaps + # this one and adds a cell-state model nothing here needs. - python_executor - - notebook_executor todo_list: mcp: false + # Kept out of the workspace root. The plan files default to + # `<output_dir>/plan.json` and `plan.md`, and output_dir *is* the directory + # whose end state becomes the task: 2 of ex11's 36 proposals wrote checks + # asserting the agent's own todo bookkeeping, one of them pinning + # `updated_at`, which no solver can reproduce. `.ms_agent/` is where + # ms_agent/project/paths.py says framework internals belong, and the + # workspace listing already skips it. + plan_filename: .ms_agent/plan.json + plan_md_filename: .ms_agent/plan.md + +# Every refusal ms-agent applies to a shell command, turned off. Read by +# LLMAgent.prepare_runtime (llm_agent.py builds PermissionConfig.from_dict off +# this section), so it takes effect in the sandbox, where tool_server.py loads +# this same file. +# +# The reason is what the refusals cost here rather than what they protect: this +# runs in a microVM that is reset once per episode and holds nothing but the +# workspace, while each refusal rules out a whole family of tasks the model could +# otherwise pose. `curl`/`wget` blocked means no task can fetch a source tarball +# or a dataset; the rm rules mean it cannot clear a directory (`rm -rf *` and +# `rm -rf build/*` are both refused) or write a task that starts from a mess that +# has to be cleaned up. +permission: + # Drops the default blacklist wholesale: curl, wget, ssh, scp, rsync, nc, + # netcat. (Whether the microVM actually has a route out is a separate + # question from whether the command is allowed to run.) + allow_network: true + safety_rules: + # Emptied, replacing the three baked-in patterns: `rm -rf /*`, `mkfs *`, + # `dd if=*`. An empty list here is not the same as an absent key -- absent + # means "use the defaults". + patterns: [] + # Same, for the configurable half of the rm/rmdir path check: `*`, `/*`, + # `/`, `~`. + dangerous_removal_paths: [] + # And the half that a config cannot reach, added for this: the refusals + # written into is_dangerous_removal_path for `*`, anything ending in `/*`, + # `/`, a direct child of `/` (which `/workspace` is), and the home directory. + unrestricted_removal: true + # A separate refusal, found by running commands through SafetyGuard rather + # than by reading the config: a glob anywhere in a write or create path is + # denied on its own ("Glob patterns not allowed in write operations"), which + # is what actually stopped `rm -rf *` and `rm -rf build/*` after the two + # lists above were emptied. It is not specific to rm -- `cp src/* dst/` and + # `chmod +x bin/*` hit it too. + allow_write_globs: true # Web search is deliberately absent. ms-agent's `web_search` key only provides # fetch_page (retrieve a known URL); a real query-a-search-engine tool needs diff --git a/cookbook/rsi/agentic/sandbox_server/Dockerfile b/cookbook/rsi/agentic/sandbox_server/Dockerfile index d5e245855..3a40cd35c 100644 --- a/cookbook/rsi/agentic/sandbox_server/Dockerfile +++ b/cookbook/rsi/agentic/sandbox_server/Dockerfile @@ -11,11 +11,42 @@ FROM python:3.11-slim ENV PIP_INDEX_URL=https://mirrors.aliyun.com/pypi/simple/ ENV PIP_TRUSTED_HOST=mirrors.aliyun.com +# apt from the same mirror, for the same reason. Measured from inside a sandbox +# on this host on 2026-08-23: deb.debian.org delivered 33 KB/s, and +# mirrors.aliyun.com/debian 5.4 MB/s -- for the ~200MB ffmpeg/imagemagick set +# that is the difference between a minute and two hours. The 9.6MB package index +# alone stalled one build long enough to look hung. Both file names are listed +# because trixie-based images carry .sources and older ones sources.list. +RUN sed -i 's|deb.debian.org|mirrors.aliyun.com|g' /etc/apt/sources.list.d/debian.sources /etc/apt/sources.list 2>/dev/null || true + +# dpkg fsyncs every control file it unpacks, and fsync does not work in these VMs: +# probed on 2026-08-23 inside a running sandbox, os.fsync returned EIO in /, /tmp, +# /workspace and /root alike with 60GB free -- the virtual block device does not +# implement flush. Without this option each package fails to unpack with "unable +# to sync file '/var/lib/dpkg/tmp.ci//md5sums': Input/output error" and apt exits +# 100. Writing without fsync is the usual answer in a container and loses nothing +# that matters here, since the image is built once and never survives a crash. +RUN mkdir -p /etc/dpkg/dpkg.cfg.d && echo force-unsafe-io > /etc/dpkg/dpkg.cfg.d/99-unsafe-io + # ripgrep is not optional: file_system's `grep` uses `rg` when it is on PATH and # silently falls back to a Python scan with a different output shape when it is # not. The policy is trained on whatever it sees, so the sandbox has to take the # same branch a serving deployment does. -RUN apt-get update && apt-get install -y --no-install-recommends ca-certificates curl git ripgrep && rm -rf /var/lib/apt/lists/* +# +# ffmpeg and imagemagick: episodes reach for them from shell_executor and got +# "ffmpeg: not found" (1 call) and "convert: not found" (2 calls) in ex6. Both +# are the standard answer for the media half of a task, so a sandbox without them +# turns a reasonable plan into a dead end. +# +# The rest are the everyday command-line tools a python:slim image happens not to +# carry. `zip` and `unzip` were each asked for and missing in ex7, and the wider +# list is there because the misses are a long tail -- every name recorded across +# ex3-ex7 appears once or twice, so waiting for a second sighting means paying for +# the same dead end again. Probed against the live image on 2026-08-23, all of +# these were absent. Deliberately left out as too large for what they would buy: +# libreoffice (~700MB), pandoc, build-essential, and weasyprint's pango/cairo +# stack. +RUN apt-get update && apt-get install -y --no-install-recommends ca-certificates curl wget git ripgrep ffmpeg imagemagick zip unzip bzip2 xz-utils p7zip-full jq sqlite3 tree file bc patch dos2unix bsdextrautils xxd poppler-utils && rm -rf /var/lib/apt/lists/* # ms-agent from source, not `pip install ms-agent`: the tools the policy is # trained against are the ones in this repository, and a released wheel can lag @@ -46,9 +77,27 @@ RUN mkdir -p /opt/ms-agent && curl -fsSL https://codeload.github.com/modelscope/ # Kept as the full list so it stays correct if ms-agent's # dependencies shift. # +# openpyxl, +# reportlab, +# pdfplumber +# -- what episodes actually asked for and did not get: 10 calls died +# on `No module named 'openpyxl'`, 3 on reportlab, 1 on pdfplumber +# (ex3-ex6, 983 recorded calls). Spreadsheets and PDFs are a large +# part of what the keyword bank proposes, and pandas' Excel support +# needs openpyxl anyway. +# pyspellchecker, +# python-docx, +# the rest of +# the tail -- same reasoning as the apt list above: `No module named` was +# recorded once each for weasyprint and spellchecker, and these are +# what a task about documents, spreadsheets, archives or text +# normally imports next. All were absent when probed against the +# live image on 2026-08-23. weasyprint itself is not here: it needs +# pango and cairo, which is a different size of decision. +# # In a sandbox, any of these missing is either a network round trip at the start # of every episode or an outright failure on an air-gapped host. -RUN pip install --no-cache-dir httpx ipykernel jupyter-client numpy pandas matplotlib seaborn scikit-learn requests beautifulsoup4 lxml pillow tqdm pyarrow +RUN pip install --no-cache-dir httpx ipykernel jupyter-client numpy pandas matplotlib seaborn scikit-learn requests beautifulsoup4 lxml pillow tqdm pyarrow openpyxl reportlab pdfplumber python-docx python-pptx xlsxwriter pypdf pymupdf toml jinja2 chardet regex tabulate sympy networkx faker pyspellchecker ENV PYTHONUNBUFFERED=1 WORKDIR /workspace diff --git a/cookbook/rsi/agentic/sandbox_server/build_via_sandbox.sh b/cookbook/rsi/agentic/sandbox_server/build_via_sandbox.sh new file mode 100644 index 000000000..8727ec1f8 --- /dev/null +++ b/cookbook/rsi/agentic/sandbox_server/build_via_sandbox.sh @@ -0,0 +1,126 @@ +#!/bin/sh +# Build the sandbox template by installing inside a live sandbox and snapshotting +# it, instead of `aenv build`. +# +# Why this exists. On 2026-08-23 three `aenv build` attempts failed or stalled on +# this host, and the reason turned out to be download speed rather than anything +# in the Dockerfile. Measured the same minute, from inside a sandbox: +# +# deb.debian.org 33 KB/s +# mirrors.aliyun.com/debian 5.4 MB/s +# host, same aliyun file 12 MB/s +# sandbox disk write 639 MB/s +# +# The build VM was pulling apt's 9.6MB package index at that first rate, which +# reads exactly like a hang: the server logs "template build started" and then +# nothing at all until the build ends. A sandbox, by contrast, installs the whole +# list in about six minutes. +# +# The Dockerfile now points apt at the same mirror, so `install.sh` should work +# again -- but this path is kept because it is the one that has been verified end +# to end, and because it needs no template builder at all. +# +# Keep the two package lists here identical to the Dockerfile's. They are +# duplicated rather than shared because this script needs shell lines a sandbox +# can run and the Dockerfile needs one instruction per line. +# +# What a snapshot does not carry: the image config. `ENV PYTHONUNBUFFERED=1`, +# `ENV PIP_INDEX_URL=...` and `WORKDIR /workspace` from the Dockerfile do not +# survive, so the steps below write the equivalents into the filesystem +# (/etc/pip.conf, /workspace) and remote_tool_env.py starts the runtime with +# `python -u`. +# +# Usage, on the environment host: +# sh build_via_sandbox.sh # snapshot named twinkle-rsi-msagent +# NAME=twinkle-rsi-msagent-v2 sh build_via_sandbox.sh # a second name, to verify first +set -eu + +NAME="${NAME:-twinkle-rsi-msagent}" +BASE_IMAGE="${BASE_IMAGE:-docker.m.daocloud.io/library/python:3.11-slim}" +# 65536 is not a preference: `aenv start --cold` refuses a virtual size smaller +# than the base image's ("shrinking is disabled"), and that base is 64GiB. +DISK_MB="${DISK_MB:-65536}" +CPU="${CPU:-2}" +MEMORY_MB="${MEMORY_MB:-2048}" +TTL="${TTL:-3600}" + +echo "==> Starting a sandbox from $BASE_IMAGE" +SID=$(aenv start --cold "$BASE_IMAGE" -d --timeout "$TTL" \ + --cpu "$CPU" --memory "$MEMORY_MB" --disk-size-mb "$DISK_MB" | tail -1 | tr -d '\r') +echo " sandbox $SID" + +SETUP=$(cat <<'SCRIPT' +set -eux +export DEBIAN_FRONTEND=noninteractive +export PIP_INDEX_URL=https://mirrors.aliyun.com/pypi/simple/ +export PIP_TRUSTED_HOST=mirrors.aliyun.com + +sed -i 's|deb.debian.org|mirrors.aliyun.com|g' \ + /etc/apt/sources.list.d/debian.sources /etc/apt/sources.list 2>/dev/null || true + +# dpkg fsyncs each control file it unpacks, and fsync does not work in this VM: +# probed on 2026-08-23, os.fsync returned EIO in /, /tmp, /workspace and /root +# alike, with 60GB free -- the virtual block device simply does not implement +# flush. Without this every package fails to unpack ("unable to sync file +# '/var/lib/dpkg/tmp.ci//md5sums': Input/output error", 278 of them). The option +# tells dpkg to write without fsyncing, which is the usual answer in a container +# and costs nothing here: the sandbox is disposable and the snapshot is taken +# from the filesystem afterwards, not from the block device's write cache. +mkdir -p /etc/dpkg/dpkg.cfg.d +echo force-unsafe-io > /etc/dpkg/dpkg.cfg.d/99-unsafe-io + +apt-get update +apt-get install -y --no-install-recommends ca-certificates curl wget git ripgrep \ + ffmpeg imagemagick zip unzip bzip2 xz-utils p7zip-full jq sqlite3 tree file \ + bc patch dos2unix bsdextrautils xxd poppler-utils +rm -rf /var/lib/apt/lists/* + +mkdir -p /opt/ms-agent +curl -fsSL https://codeload.github.com/modelscope/ms-agent/tar.gz/refs/heads/main \ + | tar -xz -C /opt/ms-agent --strip-components=1 +pip install --no-cache-dir -e /opt/ms-agent + +pip install --no-cache-dir httpx ipykernel jupyter-client numpy pandas matplotlib \ + seaborn scikit-learn requests beautifulsoup4 lxml pillow tqdm pyarrow \ + openpyxl reportlab pdfplumber python-docx python-pptx xlsxwriter pypdf \ + pymupdf toml jinja2 chardet regex tabulate sympy networkx faker pyspellchecker + +mkdir -p /workspace +printf '[global]\nindex-url = %s\ntrusted-host = %s\n' \ + "$PIP_INDEX_URL" "$PIP_TRUSTED_HOST" > /etc/pip.conf +rm -rf /root/.cache/pip +echo SETUP-OK +SCRIPT +) + +echo "==> Installing inside the sandbox (~6 min; watch /tmp/setup.log)" +B64=$(printf '%s\n' "$SETUP" | base64 -w0) +# setsid + a log file, not a foreground exec: `aenv exec` would hold the +# connection open for the whole install and a dropped ssh session would take the +# install with it. +aenv exec "$SID" sh -c "echo $B64 | base64 -d > /tmp/setup.sh; \ + sh -c 'setsid nohup sh /tmp/setup.sh > /tmp/setup.log 2>&1 &'" + +while : ; do + sleep 20 + if aenv exec "$SID" sh -c 'grep -q SETUP-OK /tmp/setup.log' 2>/dev/null; then + echo " install finished" + break + fi + aenv exec "$SID" sh -c 'tail -1 /tmp/setup.log' 2>/dev/null || true +done + +echo "==> What the sandbox ended up with" +aenv exec "$SID" python -c \ + "import openpyxl, reportlab, pdfplumber, docx, pptx, xlsxwriter, pypdf, fitz, sympy, networkx, spellchecker, ms_agent; print('python packages ok')" +aenv exec "$SID" sh -c \ + 'for b in ffmpeg convert rg git curl zip unzip 7z jq sqlite3 tree file bc pdftotext; do command -v $b >/dev/null && echo "$b ok" || echo "$b MISSING"; done' + +echo "==> Snapshotting as '$NAME'" +aenv exec "$SID" sh -c 'rm -f /tmp/setup.sh /tmp/setup.log' +aenv snapshot create "$SID" --name "$NAME" +aenv delete "$SID" >/dev/null 2>&1 || true + +echo +echo "Verify from the training host, which reaches it by the same name:" +echo " AENV_TEMPLATE=$NAME # then run the boot check in README.md ('Verify a sandbox boots')" diff --git a/cookbook/rsi/agentic/sandbox_server/install.sh b/cookbook/rsi/agentic/sandbox_server/install.sh index 6983d9a65..d99300a49 100644 --- a/cookbook/rsi/agentic/sandbox_server/install.sh +++ b/cookbook/rsi/agentic/sandbox_server/install.sh @@ -1,6 +1,11 @@ #!/bin/sh # Install AgentENV and build the sandbox template for agentic RSI. # +# If a build stalls with no output, read build_via_sandbox.sh before waiting it +# out: on our host the builder's VM downloaded at 33 KB/s against a sandbox's +# 5.4 MB/s, and that script installs inside a live sandbox and snapshots it +# instead -- six minutes, and no template builder involved. +# # Usage: # sh install.sh # install AgentENV + build the template # sh install.sh --rebuild # delete the old template and rebuild diff --git a/cookbook/rsi/agentic/sandbox_server/reap_paused.py b/cookbook/rsi/agentic/sandbox_server/reap_paused.py new file mode 100644 index 000000000..8ccb1d17e --- /dev/null +++ b/cookbook/rsi/agentic/sandbox_server/reap_paused.py @@ -0,0 +1,72 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Reap paused sandboxes on the environment host, which is what keeps it alive. + +AgentENV does not discard a sandbox when it ends: it *persists* it, as a paused +sandbox whose memory and disk image live under +``/var/lib/aenv/persisted-sandboxes/artifacts`` at roughly 1GB each. Closing the +sandbox from the client does not change this -- a closed sandbox is a paused +sandbox -- so every episode leaks a gigabyte. A GRPO step that boots +``batch_size x num_generations`` sandboxes leaks that many, and a 40GB root +filesystem is gone in a couple of dozen steps. The failure is not graceful: boots +start returning ``500: backend error: ... No space left on device``, and every +episode in the batch scores zero, which reads like a hard task rather than a +broken host. + +Run this on the environment host for the length of a training run:: + + setsid nohup python3 reap_paused.py --alias twinkle-rsi-msagent \\ + > /var/log/reap.log 2>&1 & + +Only *paused* sandboxes with the given alias are deleted. A running one may be an +episode in flight, and a different alias belongs to a different experiment -- +this script never touches either. +""" +import argparse +import json +import subprocess +import time + + +def parse_args(): + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--alias', default='twinkle-rsi-msagent', + help='only reap sandboxes built from this template') + p.add_argument('--interval', type=int, default=120, + help='seconds between sweeps; 0 sweeps once and exits') + return p.parse_args() + + +def sweep(alias): + """Delete every paused sandbox with this alias. Returns (reaped, running).""" + proc = subprocess.run(['aenv', 'list'], capture_output=True, text=True) + try: + rows = json.loads(proc.stdout) + except ValueError: + # The server restarting mid-sweep is not worth dying over; the next + # sweep will see the same sandboxes. + return 0, -1 + paused = [r['sandboxID'] for r in rows + if r.get('state') == 'paused' and r.get('alias') == alias] + for sandbox_id in paused: + subprocess.run(['aenv', 'delete', sandbox_id], capture_output=True) + running = sum(1 for r in rows if r.get('state') == 'running') + return len(paused), running + + +def main(): + args = parse_args() + while True: + reaped, running = sweep(args.alias) + disk = subprocess.run(['df', '-h', '/'], capture_output=True, + text=True).stdout.splitlines()[-1].split() + stamp = time.strftime('%H:%M:%S') + print(f'{stamp} reaped={reaped:3d} running={running:3d} ' + f'free={disk[3]} used={disk[4]}', flush=True) + if args.interval <= 0: + return + time.sleep(args.interval) + + +if __name__ == '__main__': + main() diff --git a/cookbook/rsi/agentic/sandbox_server/serve.sh b/cookbook/rsi/agentic/sandbox_server/serve.sh index e9043d8d1..e24b13e02 100644 --- a/cookbook/rsi/agentic/sandbox_server/serve.sh +++ b/cookbook/rsi/agentic/sandbox_server/serve.sh @@ -5,6 +5,7 @@ # sh serve.sh # foreground, binds 127.0.0.1:8000 # API_ADDR=0.0.0.0:8000 sh serve.sh # listen on all interfaces # NOHUP=1 sh serve.sh # background, logs to /tmp/aenv-server.log +# RUST_LOG=agentenv=debug sh serve.sh # verbose, to watch a template build # STOP_ONLY=1 sh serve.sh # shut down without starting again set -eu REPO_ROOT="${REPO_ROOT:-$HOME/AgentENV}" @@ -32,6 +33,13 @@ AENV_CONFIG_PATH="${AENV_CONFIG_PATH:-/var/lib/aenv/config/config.toml}" AENV_HOME_PATH="${AENV_HOME_PATH:-/var/lib/aenv}" AENV_RUNTIME_PATH="${AENV_RUNTIME_PATH:-/run/aenv}" +# Passed through explicitly because `sudo env` below resets the environment. At +# the default level a template build logs "template build started" and then +# nothing at all until it succeeds or fails -- a build that is merely slow reads +# exactly like a hung one, which cost hours of guessing on 2026-08-23. Restart +# with RUST_LOG=agentenv=debug before a build you need to watch. +RUST_LOG="${RUST_LOG:-agentenv=info,envd=info,uvm_ublk=info}" + if [ ! -r "$AENV_CONFIG_PATH" ]; then echo "Config not readable: $AENV_CONFIG_PATH" >&2 echo "Seed it from the repo (install.sh does this for you):" >&2 @@ -95,7 +103,7 @@ cd "$REPO_ROOT" # # AENV_RUN_USER must be explicit: the script otherwise falls back through # SUDO_USER -> repo owner -> aenv -> root, and running as root is not supported. -E="AENV_RUN_USER=aenv HOME=$AENV_HOME API_ADDR=$API_ADDR AENV_CONFIG_PATH=$AENV_CONFIG_PATH AENV_HOME_PATH=$AENV_HOME_PATH AENV_RUNTIME_PATH=$AENV_RUNTIME_PATH" +E="AENV_RUN_USER=aenv HOME=$AENV_HOME API_ADDR=$API_ADDR AENV_CONFIG_PATH=$AENV_CONFIG_PATH AENV_HOME_PATH=$AENV_HOME_PATH AENV_RUNTIME_PATH=$AENV_RUNTIME_PATH RUST_LOG=$RUST_LOG" if [ "$NOHUP" = "1" ]; then # setsid, not just nohup: the wrapper ends in `exec setpriv`, which replaces diff --git a/cookbook/rsi/agentic/sandbox_server/tool_server.py b/cookbook/rsi/agentic/sandbox_server/tool_server.py index 102da23c1..9e5cce8ed 100644 --- a/cookbook/rsi/agentic/sandbox_server/tool_server.py +++ b/cookbook/rsi/agentic/sandbox_server/tool_server.py @@ -26,13 +26,15 @@ import argparse import asyncio import copy +import inspect import json import os import sys import threading import traceback +from concurrent.futures import TimeoutError as FuturesTimeoutError from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Set, Tuple DEFAULT_PORT = 8900 @@ -43,6 +45,27 @@ _SINGLE_NS_FLAG = '_twinkle_single_namespace' +# ms-agent namespaces every tool as ``{server}---{tool}``. +_TOOL_SPLIT = '---' + +# Arguments that belong to ms-agent's plumbing rather than to any one tool, and +# that it invites the model to pass without every tool accepting one. Its own +# timeout message says to "set numeric field 'timeout' in the tool arguments" +# (tool_manager.py:687), but only the code_executor trio has a ``timeout`` +# parameter, so following that advice on write_file raises TypeError; +# ``description`` is documentation that two of those three declare and the third +# does not; the call id is injected by the host. For a tool whose signature has +# no room for one of these, it is dropped -- the alternative is failing a call +# ms-agent itself asked for. Measured over 5793 calls: 11 ``timeout`` on +# file_system tools, 2 ``description`` on shell_executor. +_FRAMEWORK_ARGS = ('timeout', 'description', 'call_id', '__call_id') + +# Withdrawn from the advertised schema whatever ms-agent declares: ``__call_id`` +# is a correlation id the host injects ("injected by host when supported", +# local_code_executor.py:494). Advertising it puts an internal handle in the +# prompt and invites the model to invent values for it. +_INTERNAL_ARGS = ('__call_id', ) + def _single_namespace_source(code: str) -> str: """Wrap ``code`` so it runs in one namespace and cannot exit the process. @@ -93,6 +116,27 @@ def _patch_python_executor() -> bool: waits for a loop that is gone. Verified: without this, a ``sys.exit(3)`` call is followed by timeouts on scripts that passed moments earlier. + And because that ``exec`` runs in *this* process, a relative path in model + code resolves against this process's cwd -- not against the workspace that + every other tool uses (``shell_executor`` and ``file_system`` both pass + ``cwd=self._ws.root`` explicitly). Measured in a live sandbox before this + chdir: ``write_file 'a.txt'`` answered "Save file successfully", the next + python call got ``[Errno 2] No such file or directory: 'a.txt'``, and the file + was in ``/workspace`` while python looked in ``/``. It cost 41 of ex7's 58 + such failures, and files python did write landed outside the directory the + end-of-episode snapshot lists, so they were invisible to whoever writes the + check script. The chdir is per call rather than once at startup so that this + holds however the server was launched; it is the same directory every time, + so concurrent calls in one turn cannot pull each other around. + + The chdir alone does not let ``import`` find a module the model wrote: + ``import`` searches ``sys.path``, which holds the server's launch dir + (``/opt/rsi``), not the workspace and not ``''``. Measured in a live sandbox: + after ``write_file 'mymod.py'``, ``open('rel.txt')`` read fine but + ``import mymod`` raised ``ModuleNotFoundError``, so the natural "write a + helper .py then import and run it" loop failed every time. So the workspace + is put on ``sys.path`` too, kept as the first entry and never duplicated. + Duplicated from ``twinkle_agentic.harness.ms_agent`` on purpose -- this file is uploaded into a sandbox that has ms-agent and nothing else. Temporary, pending an upstream PR. @@ -104,6 +148,19 @@ def _patch_python_executor() -> bool: return False async def python_executor(self, code, description='', timeout=None): + root = getattr(self, 'output_dir', None) or getattr(getattr(self, '_ws', None), 'root', None) + if root: + os.makedirs(root, exist_ok=True) + os.chdir(root) + # So ``import`` finds a module the model just wrote here. chdir moves + # cwd but not the import search path, and the workspace is not on it. + # Must be a str: the import machinery's path finders ignore a + # PathLike entry on sys.path, and ``root`` arrives as a PosixPath. + root_str = os.fspath(root) + if sys.path[:1] != [root_str]: + if root_str in sys.path: + sys.path.remove(root_str) + sys.path.insert(0, root_str) return await original(self, _single_namespace_source(code), description=description, timeout=timeout) @@ -177,6 +234,27 @@ def _without_llm_args(schema: Dict[str, Any]) -> Dict[str, Any]: return schema +def _without_internal_args(schema: Dict[str, Any]) -> Dict[str, Any]: + """Drop arguments the host owns from a tool schema. + + Unlike :func:`_without_llm_args` this does not depend on the deployment: + ``__call_id`` is never something the model should be choosing, however the + sandbox is configured. + """ + fn = schema.get('function') or {} + parameters = fn.get('parameters') or {} + properties = parameters.get('properties') or {} + if not any(arg in properties for arg in _INTERNAL_ARGS): + return schema + schema = copy.deepcopy(schema) + parameters = schema['function']['parameters'] + for arg in _INTERNAL_ARGS: + parameters['properties'].pop(arg, None) + if isinstance(parameters.get('required'), list) and arg in parameters['required']: + parameters['required'].remove(arg) + return schema + + class _LoopThread: """A single long-lived asyncio loop, owned by a background thread. @@ -241,6 +319,9 @@ def __init__(self, config_path: str, workspace: str) -> None: self.workspace = workspace self._loop = _LoopThread() self._loop.run(self._prepare()) + # Only after prepare_tools(): a contract can only be read off a tool that + # exists. + self._contracts = self._build_contracts() async def _prepare(self) -> None: self.agent.prepare_runtime() @@ -264,28 +345,215 @@ def tools(self) -> List[Dict[str, Any]]: else: flat = list(raw or []) schemas = [_to_openai(t) for t in flat if isinstance(t, dict)] + schemas = [_without_internal_args(t) for t in schemas] return [t if self.has_llm else _without_llm_args(t) for t in schemas] + def _build_contracts(self) -> Dict[str, Tuple[Set[str], Optional[Set[str]]]]: + """Per tool: the arguments advertised, and the ones the code will take. + + Both halves are needed because ms-agent lets them disagree, and every + disagreement is a call the model was invited to make and cannot. The + advertised half comes from :meth:`tools`, so it is the exact contract the + prompt carries; the other from the signature of the method + ``call_tool`` will ``getattr`` and splat the arguments into + (filesystem_tool.py:387, local_code_executor.py:583). ``None`` means the + method takes ``**kwargs`` or could not be introspected -- then nothing is + assumed and nothing is removed. + + Drift is reported at startup rather than waited for: the last one + (``shell_executor`` advertising nothing about ``description`` while its + siblings declare it) cost two calls in 239 before anyone noticed, and it + was found by reading a trajectory. + """ + contracts: Dict[str, Tuple[Set[str], Optional[Set[str]]]] = {} + for schema in self.tools(): + fn = schema.get('function') or {} + name = fn.get('name') + if not name: + continue + declared = set((fn.get('parameters') or {}).get('properties') or {}) + contracts[name] = (declared, self._accepted_args(name)) + drift = { + name: sorted(declared - accepted) + for name, (declared, accepted) in contracts.items() + if accepted is not None and declared - accepted + } + if drift: + sys.stderr.write('[tool_server] WARNING advertised arguments the implementation ' + 'rejects (dropped at dispatch, fix upstream): %s\n' % (drift, )) + sys.stderr.flush() + return contracts + + def _accepted_args(self, name: str) -> Optional[Set[str]]: + """Keyword names the implementation behind ``name`` accepts, or None.""" + try: + tool_ins = self._tm._tool_index[name][0] + method = getattr(tool_ins, name.split(_TOOL_SPLIT)[-1]) + sig = inspect.signature(method) + except Exception: # noqa -- an un-introspectable tool just gets no repairs + return None + if any(p.kind is p.VAR_KEYWORD for p in sig.parameters.values()): + return None + return { + n + for n, p in sig.parameters.items() + if p.kind in (p.POSITIONAL_OR_KEYWORD, p.KEYWORD_ONLY) + } + + def _reconcile(self, name: str, args: Dict[str, Any]) -> Tuple[Dict[str, Any], Optional[str]]: + """Repair what ms-agent's contract breaks; refuse the rest, precisely. + + Two different failures arrive as the same TypeError, and they do not + deserve the same treatment: + + * an argument ms-agent asked for and cannot take -- its own plumbing + (:data:`_FRAMEWORK_ARGS`) or a schema that overstates the code -- is + removed. The model followed the contract it was given; failing the call + would only teach it to distrust a correct one. + * an argument the model invented is refused, with the accepted list and + the tool those arguments actually belong to. Measured over 5793 calls, + 259 ``write_file`` calls carried ``old_string``/``new_string``, which + are ``edit_file``'s. Rewriting those into a ``content=`` write would + hide a mistake the model should be trained out of, and would teach it a + call shape that fails outside this sandbox. + + Returns ``(arguments, error)``; ``error`` is not None when the call must + not run. + """ + contract = self._contracts.get(name) + if contract is None: + return args, None + declared, accepted = contract + args = dict(args) + for arg in list(args): + if accepted is None or arg in accepted: + continue + if arg in _FRAMEWORK_ARGS or arg in declared: + args.pop(arg) + # glob's own default for ``path`` is '' (filesystem_tool.py:392 advertises + # it as optional), but ms-agent's safety guard rejects an empty file path + # before dispatch, so a model that spells the default out loud gets + # "Blocked by safety policy: Empty file path" -- 66 times in 5793 calls. + # '.' is what '' resolves to once inside the tool. + if name.split(_TOOL_SPLIT)[-1] == 'glob' and 'path' in args \ + and not str(args.get('path') or '').strip(): + args['path'] = '.' + unknown = sorted(set(args) - declared) + if unknown: + return args, self._argument_error(name, unknown, declared) + return args, None + + def _argument_error(self, name: str, unknown: List[str], declared: Set[str]) -> str: + """Say what was rejected, what is accepted, and who owns the rest. + + The last part is the useful one and it costs nothing: the arguments of + every other advertised tool are already known here, so an argument + belonging to a sibling can be named as such instead of leaving the model + to guess which of eleven tools it meant. + """ + owners: Dict[str, List[str]] = {} + for other, (other_declared, _accepted) in self._contracts.items(): + if other == name: + continue + for arg in unknown: + if arg in other_declared: + owners.setdefault(arg, []).append(other) + quoted = ', '.join(repr(a) for a in unknown) + parts = ['Error: %s has no argument %s.' % (name, quoted), + 'It accepts: %s.' % (', '.join(sorted(declared)) or '(none)')] + for arg, tools in sorted(owners.items()): + parts.append('%r belongs to %s.' % (arg, ' or '.join(sorted(tools)))) + parts.append('Re-issue the call with this tool\'s arguments, or call the tool ' + 'the arguments belong to.') + return ' '.join(parts) + def call(self, calls: List[Dict[str, Any]], timeout: Optional[float]) -> List[Dict[str, Any]]: """Dispatch a turn's tool calls, mirroring how ms-agent itself does it. A single call goes through ``single_call_tool`` and a batch through ``parallel_call_tool``, matching LLMAgent, so concurrency-sensitive - tools behave in training exactly as they do in production. + tools behave in training exactly as they do in production. Each call is + put through :meth:`_reconcile` first, and one that cannot run is answered + from here without reaching ms-agent -- so a batch keeps its shape and + result *i* still answers call *i*. """ - payload = [{'tool_name': c.get('tool_name'), 'arguments': c.get('arguments') or {}} for c in calls] - try: - if len(payload) == 1: - results = [self._loop.run(self._tm.single_call_tool(payload[0]), timeout)] + out: List[Optional[Dict[str, Any]]] = [None] * len(calls) + prepared: List[Tuple[int, Dict[str, Any]]] = [] + for i, c in enumerate(calls): + name = c.get('tool_name') + args = c.get('arguments') + if isinstance(args, str): + try: + args = json.loads(args or '{}') + except ValueError: + # ms-agent has its own message for unparseable arguments, and + # it names the offending text; leave the call to it. + prepared.append((i, {'tool_name': name, 'arguments': c.get('arguments')})) + continue + if not isinstance(args, dict): + args = {} + args, error = self._reconcile(name, args) + if error: + out[i] = {'observation': error, 'ok': False} + else: + prepared.append((i, {'tool_name': name, 'arguments': args})) + if prepared: + payload = [p for _i, p in prepared] + try: + if len(payload) == 1: + results = [self._loop.run(self._tm.single_call_tool(payload[0]), timeout)] + else: + results = self._loop.run(self._tm.parallel_call_tool(payload), timeout) + except Exception as e: # noqa + # One failing tool must not take down the server: the episode can + # still recover, and a dead server would fail every later step of + # every trajectory sharing this sandbox. + # + # A timeout is spelled out rather than reported as its exception + # name. `concurrent.futures.TimeoutError` carries no message at + # all, so the model used to read "Tool call failed. TimeoutError:" + # -- which says nothing about what to do differently. What it + # needs to know is that the call was abandoned rather than + # rejected, that whatever it started may still be running (this + # cannot cancel a subprocess ms-agent has already spawned), and + # that a long-running command has somewhere else to go. ex8's + # episode 23 started an HTTP server in the foreground and stalled + # the whole turn. + if isinstance(e, (FuturesTimeoutError, asyncio.TimeoutError)): + detail = (f'Timed out: this turn\'s tool calls did not finish within ' + f'{timeout}s and were abandoned. Whatever they started may ' + f'still be running. A command that does not return on its own ' + f'-- a server, a watcher, an interactive program -- has to be ' + f'started with run_in_background=true, or given an explicit ' + f'time limit inside the command itself.') + else: + detail = f'Tool call failed. {type(e).__name__}: {e}' + for i, _p in prepared: + out[i] = {'observation': detail, 'ok': False} else: - results = self._loop.run(self._tm.parallel_call_tool(payload), timeout) - except Exception as e: # noqa - # One failing tool must not take down the server: the episode can - # still recover, and a dead server would fail every later step of - # every trajectory sharing this sandbox. - detail = f'{type(e).__name__}: {e}' - return [{'observation': f'Tool call failed. {detail}', 'ok': False} for _ in payload] - return [{'observation': _as_text(r), 'ok': True} for r in list(results)] + for (i, _p), r in zip(prepared, list(results)): + out[i] = {'observation': _with_timeout_advice(_as_text(r)), 'ok': True} + return [o if o is not None else {'observation': '', 'ok': False} for o in out] + + +# ms-agent's own words when its per-call wait runs out (tool_manager.py:687). +_MS_TIMEOUT_MARK = 'Tool call timed out after' +# Appended to it, not substituted for it. Its message offers exactly one remedy -- +# raise the `timeout` argument -- which is the wrong one for a command that never +# returns at all: ex8's episode 23 started `python -m http.server` in the +# foreground, and no limit up to the 600s ceiling would have helped. shell_executor +# already advertises `run_in_background`, so this names the argument the model +# already has rather than teaching it anything new. +_TIMEOUT_ADVICE = ( + ' A command that does not return on its own -- a server, a watcher, an ' + 'interactive program -- will time out at any limit; start it with ' + 'run_in_background=true instead, or bound it inside the command itself.') + + +def _with_timeout_advice(observation: str) -> str: + if _MS_TIMEOUT_MARK in observation and 'run_in_background' not in observation: + return observation + _TIMEOUT_ADVICE + return observation def _as_text(result: Any) -> str: diff --git a/cookbook/rsi/agentic/split_tasks.py b/cookbook/rsi/agentic/split_tasks.py new file mode 100644 index 000000000..6aeb00e4d --- /dev/null +++ b/cookbook/rsi/agentic/split_tasks.py @@ -0,0 +1,70 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Split a challenge.py task file into a training set and a held-out eval set. + +The eval set is stratified by ``n_pass`` -- the number of solver attempts that +succeeded when the task was filtered. Difficulty is the whole point of the +filter, so a random split can easily hand the eval set every task the model +already solves 3 times in 4, and a pass rate on those says nothing about the +hard end. Stratifying keeps both halves the same shape. + + python cookbook/rsi/agentic/split_tasks.py \\ + output/rsi_agentic/run3/challenge_flows.jsonl --eval-frac 0.25 +""" +import argparse +import json +import os +import random + + +def main(): + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('flows', help='challenge_flows.jsonl') + p.add_argument('--eval-frac', type=float, default=0.25) + p.add_argument('--seed', type=int, default=0) + p.add_argument('--train-out', default='', help='default: <flows dir>/train_tasks.jsonl') + p.add_argument('--eval-out', default='', help='default: <flows dir>/eval_tasks.jsonl') + args = p.parse_args() + + out_dir = os.path.dirname(os.path.abspath(args.flows)) + train_out = args.train_out or os.path.join(out_dir, 'train_tasks.jsonl') + eval_out = args.eval_out or os.path.join(out_dir, 'eval_tasks.jsonl') + + with open(args.flows, encoding='utf-8') as f: + tasks = [json.loads(line) for line in f if line.strip()] + if not tasks: + raise SystemExit(f'{args.flows} contains no tasks') + + strata = {} + for task in tasks: + strata.setdefault(task.get('n_pass'), []).append(task) + + rng = random.Random(args.seed) + train, held = [], [] + for n_pass in sorted(strata, key=lambda x: (x is None, x)): + group = strata[n_pass][:] + rng.shuffle(group) + # round() rather than int(): with 3 tasks at a difficulty and a quarter + # held out, truncating would give the eval set none of them. + n_eval = min(len(group) - 1, round(len(group) * args.eval_frac)) if len(group) > 1 else 0 + held.extend(group[:n_eval]) + train.extend(group[n_eval:]) + + for path, rows in ((train_out, train), (eval_out, held)): + with open(path, 'w', encoding='utf-8') as f: + for row in rows: + f.write(json.dumps(row, ensure_ascii=False) + '\n') + + def dist(rows): + out = {} + for row in rows: + out[row.get('n_pass')] = out.get(row.get('n_pass'), 0) + 1 + return dict(sorted(out.items(), key=lambda kv: (kv[0] is None, kv[0]))) + + print(f'{len(tasks)} tasks, n_pass dist {dist(tasks)}') + print(f'train {len(train)} -> {train_out} dist {dist(train)}') + print(f'eval {len(held)} -> {eval_out} dist {dist(held)}') + + +if __name__ == '__main__': + main() diff --git a/docs/source_en/Components/Agentic/Rollout.md b/docs/source_en/Components/Agentic/Rollout.md index 94b143454..e803b9076 100644 --- a/docs/source_en/Components/Agentic/Rollout.md +++ b/docs/source_en/Components/Agentic/Rollout.md @@ -74,7 +74,7 @@ Each output trajectory dict includes: | `labels` | `List[int]` | Training labels (`-100` for non-trainable tokens). | | `turns` | `int` | Number of turns performed. | | `stop_reason` | `str` | `'stop'` / `'length'` | -| `truncated` | `bool` | Whether the trajectory was truncated. | +| `truncated` | `bool` | Whether the trajectory was cut off rather than concluding on its own: generation hit `max_tokens` (`stop_reason='length'`), the turn limit was reached, or a length cap dropped it. | | `logprobs` | `List` | Per-token log probabilities (if available). | ### Ray Remote Support diff --git "a/docs/source_zh/\347\273\204\344\273\266/Agentic/Rollout.md" "b/docs/source_zh/\347\273\204\344\273\266/Agentic/Rollout.md" index b74c1e791..134532c97 100644 --- "a/docs/source_zh/\347\273\204\344\273\266/Agentic/Rollout.md" +++ "b/docs/source_zh/\347\273\204\344\273\266/Agentic/Rollout.md" @@ -74,7 +74,7 @@ results = rollout(trajectories) | `labels` | `List[int]` | 训练标签（非可训练 token 为 `-100`）。 | | `turns` | `int` | 执行的轮次数。 | | `stop_reason` | `str` | `'stop'` / `'length'` | -| `truncated` | `bool` | 轨迹是否被截断。 | +| `truncated` | `bool` | 轨迹是否被截断（而非自行结束）：生成触及 `max_tokens`（`stop_reason='length'`）、达到轮次上限，或被长度上限丢弃。 | | `logprobs` | `List` | 每 token 的对数概率（如有）。 | ### Ray 远程支持 diff --git a/src/twinkle/data_format/sampling.py b/src/twinkle/data_format/sampling.py index 05ecdd641..a7b458ba2 100644 --- a/src/twinkle/data_format/sampling.py +++ b/src/twinkle/data_format/sampling.py @@ -12,6 +12,15 @@ class SamplingParams: max_tokens: Optional[int] = None seed: Optional[int] = None stop: Union[str, Sequence[str], Sequence[int], None] = None + # Whether what ``stop`` matched stays in the output. vLLM drops it by + # default -- both the string form and the token-id form, since v1's + # detokenizer excludes the final token whenever a stop terminated the + # request -- which is wrong for a stop that is part of the syntax being + # generated. Stopping a tool-using agent at '</tool_call>' so it reads one + # observation before deciding the next call is exactly that case: without + # this, every turn the policy is trained on ends on an unclosed + # '<tool_call>' block. + include_stop_str_in_output: bool = False temperature: float = 1.0 top_k: int = -1 top_p: float = 1.0 @@ -95,6 +104,8 @@ def to_vllm(self, **kwargs): kwargs['stop_token_ids'] = list(self.stop) else: kwargs['stop'] = list(self.stop) + if self.include_stop_str_in_output: + kwargs['include_stop_str_in_output'] = True if self.logprobs is not None: kwargs['logprobs'] = self.logprobs diff --git a/src/twinkle/template/tools/bracket_dsl.py b/src/twinkle/template/tools/bracket_dsl.py index c3d1088a8..dc0b16359 100644 --- a/src/twinkle/template/tools/bracket_dsl.py +++ b/src/twinkle/template/tools/bracket_dsl.py @@ -20,6 +20,18 @@ class BracketDslParser(ToolCallParser): parentheses (list arguments), so the call list is located by scanning with a depth counter rather than by a bracket-free regex. Argument values are read as python literals, falling back to the raw text when they are not literals. + + Fenced code blocks are excluded, and so is anything the model wrote inside + ``<think>``: this format has no markup of its own, so a python expression is + otherwise indistinguishable from a call list. Two further rules keep code out: + a block only counts as a call list when every argument in it is a keyword + argument (``name=value``), which no comprehension is, and a reply cut off + mid-thought leaves ``<think>`` unterminated, so that region runs to the end of + the text. + + Getting this wrong is expensive and quiet: ``[int(v) for v in raw]`` in a + reply parses as a call to ``int`` with no arguments, the tool the model + actually meant to call never runs, and the episode ends having done nothing. """ name = 'bracket_dsl' @@ -33,9 +45,59 @@ class BracketDslParser(ToolCallParser): _DETECT_RE = re.compile(r"\[\s*[A-Za-z_][\w.\-' ]*?\s*\(") # Split an argument body on top-level commas only (values may hold commas). _ARG_NAME_RE = re.compile(r'^\s*([A-Za-z_]\w*)\s*=\s*(.*)$', re.DOTALL) + # A fence runs to its closing delimiter, or to the end of a truncated reply. + _FENCE_RE = re.compile(r'```.*?(?:```|\Z)', re.DOTALL) + # So does a thinking block: a reply truncated inside one never closes it. + _THINK_RE = re.compile(r'<think>.*?(?:</think>|\Z)', re.DOTALL) + + @staticmethod + def _fenced_spans(text: str) -> List[Tuple[int, int]]: + return [m.span() for m in BracketDslParser._FENCE_RE.finditer(text or '')] + + @staticmethod + def _skip_spans(text: str) -> List[Tuple[int, int]]: + """Regions where a call list is quoted code or private thought, not a call.""" + text = text or '' + return (BracketDslParser._fenced_spans(text) + + [m.span() for m in BracketDslParser._THINK_RE.finditer(text)]) + + @staticmethod + def _in_spans(index: int, spans: List[Tuple[int, int]]) -> bool: + return any(start <= index < end for start, end in spans) + + @classmethod + def _is_keyword_body(cls, body: str) -> bool: + """Is every argument in this body a ``name=value`` pair? + + An empty body qualifies -- ``[get_time()]`` is a call list. A positional + argument does not: that is what a comprehension or a nested expression + looks like. + """ + chunks = [c for c in cls._split_top_level(body) if c.strip()] + return all(cls._ARG_NAME_RE.match(c) for c in chunks) + + @classmethod + def _looks_like_call_list(cls, block: str) -> bool: + """Does ``[...]`` hold calls with keyword arguments, and nothing else?""" + pos, seen = 1, 0 + while pos < len(block): + m = cls._CALL_START_RE.search(block, pos) + if not m: + break + close = cls._match_paren(block, m.end() - 1) + if close is None: + return False + if not cls._is_keyword_body(block[m.end():close]): + return False + seen += 1 + pos = close + 1 + return seen > 0 def detect(self, text: str) -> bool: - return bool(self._DETECT_RE.search(text or '')) + # Via _find_blocks, so that detect and parse cannot disagree: a parser + # that claims a reply and then finds nothing in it denies the remaining + # parsers their turn. + return bool(self._find_blocks(text or '')) @staticmethod def _find_blocks(text: str) -> List[Tuple[int, int]]: @@ -47,9 +109,10 @@ def _find_blocks(text: str) -> List[Tuple[int, int]]: ("Get Today's Prices") does not start a string. """ spans: List[Tuple[int, int]] = [] + skip = BracketDslParser._skip_spans(text) i, n = 0, len(text or '') while i < n: - if text[i] != '[': + if text[i] != '[' or BracketDslParser._in_spans(i, skip): i += 1 continue if not BracketDslParser._DETECT_RE.match(text, i): @@ -75,7 +138,8 @@ def _find_blocks(text: str) -> List[Tuple[int, int]]: elif ch == ']': depth -= 1 if depth == 0: - spans.append((i, j + 1)) + if BracketDslParser._looks_like_call_list(text[i:j + 1]): + spans.append((i, j + 1)) break j += 1 i = (spans[-1][1] if spans and spans[-1][0] == i else i + 1) diff --git a/src/twinkle_agentic/challenger/agentic.py b/src/twinkle_agentic/challenger/agentic.py index b28aad6f3..7cfdcbf63 100644 --- a/src/twinkle_agentic/challenger/agentic.py +++ b/src/twinkle_agentic/challenger/agentic.py @@ -2,35 +2,47 @@ """Agentic challenger: invent tasks by doing them first. The approach mirrors how the code challenger works, adapted to tool-using -agents. Instead of writing a problem statement and hoping it is achievable, -the model first *does* something interesting in a sandbox (round 1), then a -second call writes check assertions that verify the end state, and a third -call writes the problem statement someone else would need to reproduce it. +agents. Instead of writing a problem statement and hoping it is achievable, the +model first *does* something in a sandbox, then -- in the same conversation -- +writes the check script that verifies the end state it just produced, then the +problem statement someone else would need to reproduce it. Steps for one candidate: 1. Choose direction + keywords. Optionally start from a seed trajectory. - 2. Round 1 (explore, multi-turn with tools): model acts in a clean sandbox, - producing a tool-call chain and a final workspace state. - 3. Round 2a (explore, single-turn): model sees the trajectory and writes a - python check script that asserts properties of the end state. + 2. Explore (multi-turn with tools): model acts in a clean sandbox, producing + a tool-call chain and a final workspace state, and stops calling tools. + 3. A user message is appended to that same conversation carrying the + workspace listing, asking for a python check script. Tools are no longer + dispatched from here on. 4. Verify: run the check script in the sandbox (must pass). - 5. Round 2b (explore, single-turn): model sees trajectory + checks and - writes a problem statement. + 5. A second user message is appended asking for the problem statement. 6. Difficulty filter: reset workspace, let the solver do the task N times, run checks, keep only "sometimes pass" tasks. -Because every round-1 episode needs a clean workspace and because episodes -share a single long-lived sandbox, round 1 is **serial** -- one proposal at a -time with a workspace reset in between. Rounds 2a/2b are text-only generation -and can be batched. +Steps 3 and 5 are appended to the episode rather than sent as fresh calls, so +every assistant turn in the chain keeps its ``labels`` and ``logprobs`` and the +whole proposal -- acting, checking, describing -- is one trainable sample. The +follow-up messages come back from :meth:`AgenticChallenger._followup`, which the +rollout calls at the moment the model stops calling tools; that is where the +sandbox work (snapshot, running the check) happens, because only the caller can +do it. + +Because every episode needs a clean workspace and because episodes share a +single long-lived sandbox, proposing is **serial** -- one proposal at a time with +a workspace reset in between. Prompt text is not here. Every string the model sees arrives in :class:`AgenticPrompts`, built by whoever runs the challenger -- see ``cookbook/rsi/agentic/prompts.py``. """ +import ast +import json import re +import threading +from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass +from functools import partial from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple from twinkle.data_format import SamplingParams, Trajectory, user_data_get @@ -48,6 +60,11 @@ ] _FENCE_RE = re.compile(r'```(?:python|py)?\s*\n(.*?)```', re.S) +# A fence around the *whole* reply, which is packaging rather than content. +_WHOLE_FENCE_RE = re.compile(r'```[\w+-]*\s*\n?(.*?)```', re.S) +# The JSON body of a tool call: the proposing episode uses tools, so at the check +# stage a 4B model often keeps calling one instead of writing a fenced block. +_TOOLCALL_RE = re.compile(r'<tool_call>\s*(.*?)\s*</tool_call>', re.S) # ── parsing ─────────────────────────────────────────────────────────────── @@ -55,7 +72,12 @@ def parse_check_script(text: str) -> Optional[str]: """Extract a python check script from the model's reply. - Looks for the last fenced python code block after ``</think>``. + Prefers the last fenced python block after ``</think>``. When the reply has + no fence at all, falls back to reading the tail as bare code: 8 of + armA2shellV6's 11 check_parse_fail rejections were a complete, parseable + check script that the model simply did not wrap in backticks, and throwing + the task away over the packaging loses a task that was ready. + Returns ``None`` when nothing usable is found. """ body = text or '' @@ -63,45 +85,189 @@ def parse_check_script(text: str) -> Optional[str]: if idx >= 0: body = body[idx + len('</think>'):] blocks = _FENCE_RE.findall(body) - if not blocks: + if blocks: + script = blocks[-1].strip() + return script if script else None + bare = _bare_check_script(body) + if bare: + return bare + return _toolcall_check_script(body) + + +def _bare_check_script(body: str) -> Optional[str]: + """Read an unfenced reply as code, or None. + + Advances the start line until the rest parses, which drops whatever prose + came first (the "ALSO CORRECT:" line, a sentence introducing the script) + without needing to recognise it. Requires an ``assert`` so that a one-line + reply of prose -- which can be a syntactically valid expression -- is not + mistaken for a check. + """ + lines = body.strip().split('\n') + for start in range(len(lines)): + cand = '\n'.join(lines[start:]).strip() + if 'assert' not in cand: + break # no assert left in the tail; nothing further can qualify + try: + ast.parse(cand) + except SyntaxError: + continue + return cand + return None + + +def _toolcall_check_script(body: str) -> Optional[str]: + """Recover a check script the model put inside a tool call, or None. + + The proposing episode uses tools, and at the check stage a 4B model often + keeps calling one -- it emits ``python_executor(code="...assert...")`` (or a + shell command, or ``write_file(content=...)``) instead of a fenced block. + The script is right there in the call's ``code``/``command``/``content`` + argument, so pull it out rather than lose the task: measured on run_clean1, + most first-round check_parse_fail rejections were tool-call wrapped. + + Only code that parses and actually asserts is accepted, so a shell + ``command`` that merely runs a file -- which has no assert of its own -- + does not slip through as a check. + """ + blobs = _TOOLCALL_RE.findall(body) + for blob in reversed(blobs): + code = None + try: + obj = json.loads(blob) + args = obj.get('arguments') if isinstance(obj, dict) else None + if isinstance(args, dict): + code = args.get('code') or args.get('command') or args.get('content') + except (ValueError, AttributeError): + m = re.search(r'"(?:code|command|content)"\s*:\s*"(.*?)"\s*\}', blob, + re.S) + if m: + try: + code = m.group(1).encode().decode('unicode_escape') + except (UnicodeDecodeError, ValueError): + code = None + if not code or 'assert' not in code: + continue + try: + ast.parse(code) + except SyntaxError: + continue + return code.strip() + return None + + + +# Two rules CHECK_FOLLOWUP already states -- no equality on a script's source +# text, no byte count or checksum on a binary -- were broken by 9 and 8 of 41 +# measured tasks respectively, so stating them a third time is not the fix. A +# check that pins the exact source of a .py rejects every equivalent solution, +# and one that pins a .png's byte count rejects every matplotlib version; both +# make a task nobody but the author can pass. +_SIZE_OR_HASH_NAMES = ('getsize', 'st_size', 'sha256', 'sha1', 'md5', 'hexdigest', + 'digest') +# What makes a string python rather than data. Checked instead of "is it long and +# multi-line", because the contents of a csv or a json file are legitimately +# asserted verbatim -- the statement handed those to the solver -- while the text +# of a script never is. +_LOOKS_LIKE_PYTHON = ('import ', 'def ', 'print(', 'with open(', 'if __name__') + + +def brittle_check_reason(script: str) -> Optional[str]: + """Why this check script would reject a correct solution, or None. + + Returned text goes back to the model through the same retry path a failing + assertion uses, because the defect is the same kind: an assertion that does + not hold for solutions other than the one in front of it. + + Read off the syntax tree rather than matched as text. Both defects survive + patterns easily: source equality reads the file into a name first + (``c = f.read()``, then ``assert c == '...'``) so nothing sits between + ``open()`` and ``==``, and a size check can put the call either around the + name (``getsize("a.png")``) or after it. + """ + try: + tree = ast.parse(script) + except SyntaxError: + # Unparseable means it cannot run either, so let the sandbox report it. return None - script = blocks[-1].strip() - return script if script else None + for node in ast.walk(tree): + if not (isinstance(node, ast.Compare) + and any(isinstance(o, ast.Eq) for o in node.ops)): + continue + for side in [node.left] + list(node.comparators): + if not (isinstance(side, ast.Constant) and isinstance(side.value, str)): + continue + if any(m in side.value for m in _LOOKS_LIKE_PYTHON): + return ('AssertionError: this check compares a file against the ' + 'full text of a python script with ==, which only the ' + 'exact script you wrote can pass. Assert what running ' + 'that script produces instead.') + # A byte count or a checksum compared for equality. Not restricted to + # binary suffixes: CHECK_FOLLOWUP says "NEVER assert a file size in bytes" + # about any file, and keying on a suffix list let + # ``getsize('data.mat') == 264`` through. Only equality against a literal is + # a defect -- ``getsize(f) > 0`` is a fine way to say "not empty". + for node in ast.walk(tree): + if not (isinstance(node, ast.Compare) + and any(isinstance(o, ast.Eq) for o in node.ops)): + continue + sides = [node.left] + list(node.comparators) + has_literal = any(isinstance(s, ast.Constant) + and isinstance(s.value, (int, float, str)) + and not isinstance(s.value, bool) for s in sides) + if not has_literal: + continue + for side in sides: + names = {n.attr for n in ast.walk(side) if isinstance(n, ast.Attribute)} + names |= {n.id for n in ast.walk(side) if isinstance(n, ast.Name)} + hit = names & set(_SIZE_OR_HASH_NAMES) + if hit: + what = ('a checksum' if hit - {'getsize', 'st_size'} + else 'a byte count') + return (f'AssertionError: this check pins {what} of a file, and ' + 'correct solutions differ there. Assert what can be read ' + 'out of the file instead -- its structure, or the values ' + 'inside it.') + # Comparing raw bytes of a file: same defect, different spelling. + for node in ast.walk(tree): + if not (isinstance(node, ast.Compare) + and any(isinstance(o, ast.Eq) for o in node.ops)): + continue + for side in [node.left] + list(node.comparators): + if isinstance(side, ast.Constant) and isinstance(side.value, bytes): + return ('AssertionError: this check compares the raw bytes of a ' + 'file, and correct solutions differ there. Assert what ' + 'can be read out of it instead.') + return None + def parse_problem_statement(text: str) -> Optional[str]: """Extract a problem statement from the model's reply. - The model is asked to return the problem in prose (not code). We take - everything after ``</think>`` with code fences stripped as the statement. + Everything after ``</think>`` is the statement. A fence around the whole + reply is unwrapped; fences *inside* it are kept. + + Keeping them matters more than it sounds: a statement that says what a file + must contain puts the content in a fence, and stripping every fence left + "1. `data.json` containing:" with nothing after it. 7 of ex11's 16 measured + statements had a fence, and 5 of those 7 were solved 0 times out of 8 -- + against 1 of the 9 statements that had no fence to lose. The tasks were not + hard, they were unanswerable. + Returns ``None`` when the result is empty. """ body = text or '' idx = body.rfind('</think>') if idx >= 0: body = body[idx + len('</think>'):] - # Strip any fenced blocks (those are code, not prose) - body = _FENCE_RE.sub('', body).strip() - # Strip json fences too - body = re.sub(r'^\s*```(?:json)?\s*|\s*```\s*$', '', body, flags=re.I).strip() + body = body.strip() + whole = _WHOLE_FENCE_RE.fullmatch(body) + if whole: + body = whole.group(1).strip() return body if body else None -def _trajectory_summary(trajectory: Trajectory) -> str: - """A compact text representation of a trajectory for prompting. - - Shows each message as role: content (truncated for tool results). - """ - parts = [] - for msg in trajectory.get('messages') or []: - role = msg.get('role', '?') - content = msg.get('content') or '' - if role == 'tool' and len(content) > 500: - content = content[:500] + '...[truncated]' - parts.append(f'[{role}] {content}') - return '\n'.join(parts) - - # The fields a local rollout splices into a trajectory, and the only ones a # later GRPO step needs: ``labels`` marks which of ``input_ids`` are trainable # (-100 elsewhere) and ``logprobs`` holds one entry per trainable token, taken @@ -136,20 +302,24 @@ class AgenticPrompts: Placeholder validation happens at construction time. """ - # Round 1: model acts in sandbox + # Explore: model acts in sandbox system: str from_scratch: str from_seed: str = '' from_keywords: str = '' from_seed_keywords: str = '' - # Round 2a: write check script - check_system: str = '' - check_user: str = '' - - # Round 2b: write problem statement - problem_system: str = '' - problem_user: str = '' + # Appended to the same conversation once the model stops calling tools: + # first "write the check script" (which carries the workspace listing), then + # "write the problem statement". Each has to repeat the rules that used to + # live in a system message of its own, because there is no second system + # message in a single conversation. + check_followup: str = '' + # Sent instead of the statement stage when the check script does not pass, so + # the model can fix it from the traceback. Required only when the challenger + # is built with ``check_retries`` above 0. + check_retry_followup: str = '' + problem_followup: str = '' # Keyword generation (same structure as code side) keyword_system: str = '' @@ -160,15 +330,14 @@ class AgenticPrompts: 'from_seed': ('seed',), 'from_keywords': ('keywords',), 'from_seed_keywords': ('seed', 'keywords'), - 'check_user': ('trajectory', 'final_state'), - 'problem_user': ('trajectory', 'checks'), + 'check_followup': ('final_state',), + 'check_retry_followup': ('error', 'final_state'), 'keyword_user': ('k', 'desc'), 'keyword_expand_user': ('kw', 'm'), } def __post_init__(self): - for name in ('system', 'from_scratch', 'check_system', 'check_user', - 'problem_system', 'problem_user'): + for name in ('system', 'from_scratch', 'check_followup', 'problem_followup'): if not getattr(self, name).strip(): raise ValueError(f'AgenticPrompts.{name} is required') for name, placeholders in self._REQUIRED_FIELDS.items(): @@ -208,24 +377,80 @@ class AgenticChallenger(Challenger): workspace_snapshot_fn: after round 1, return a text summary of the workspace state (e.g. ``find . -type f``). If None, a default that lists messages is used. + tool_schemas: the executor's tool schemas, in the OpenAI shape the + template renders. Attached to the trajectories that are *meant* to + call tools -- the exploring episode and each solve attempt. Without + this the model is never told the tool names, so it writes code in + prose instead of calling anything: the workspace stays empty, every + check fails, and the difficulty numbers describe a model that had no + tools rather than a hard task. The check-writing and + problem-writing stages sit in the same conversation and so see the + same list, which is why the rollout stops dispatching calls once a + follow-up has been appended -- a python block written as an *answer* + parses as a call list, and 41 of 146 such replies in a measured run + edited the very workspace the answer was about. combo_arity: ``'triple'`` or ``'mix'``, as in :class:`.CodeChallenger`. arity_weights: weights for the ``'mix'`` subset size. single_kw_prob: chance of using one category in ``'triple'`` mode. keyword_refill_target / keyword_gen_calls / keyword_refill_tries / keyword_params: keyword bank refill parameters. + check_params / problem_params: sampling params for the two appended + stages. ``None`` keeps whatever the episode was already using, which + is sized for one agent turn; the check-writing stage reads the whole + episode plus the end state and reasons at length before answering, + and one that runs out of budget mid-thought never emits its code + block and is thrown away as unparseable. + followup_api: optional OpenAI-compatible API client (e.g. qwen3-max). When + given, exploration still runs on the local explorer -- so its turns keep + their ``labels`` and ``logprobs`` and remain trainable -- but the + check-script (success judgement) and problem-statement stages are + generated by this API instead of the local model, and are appended + neither to the trainable trajectory nor to its token stream. This is + the "explore locally, judge and describe over an API, train only the + exploration" split. ``None`` keeps the single-model behaviour where the + local model writes those two stages in the same conversation. + followup_extra_body: extra request body forwarded on every ``followup_api`` + call (e.g. ``{'thinking_budget': N}`` to cap qwen3-max reasoning). + ``None`` sends the request unmodified. Ignored when ``followup_api`` is + ``None``. + keyword_explorer: explorer used to brainstorm keywords. Should have no + tools wired to it: a list is a text answer, and a bracketed list in + a reply is exactly what the sandbox explorer would try to dispatch as + a call. ``None`` reuses the main explorer, which for a sandbox setup + means its tools are live there too. min_batch: smallest batch worth sending to the explorer. problem_max_chars: reject problem statements longer than this. + check_retries: how many times a check script that did not pass is handed + back, with the traceback and the workspace listing, for a rewrite + before the proposal is rejected. 0 restores the old behaviour of + rejecting on the first failure. Measured in ex12: 36 of 72 proposals + died on the check, and 29 of those were a single assertion naming a + value the model had not read -- a row count, a content string that + was nearly right, a timestamp -- over a workspace state that was + perfectly good. Each retry costs one more sampling call for that + episode and nothing for the ones that pass first time. reject_sink: called with a dict for every rejected proposal. propose_sink: called once per proposal attempt -- kept, rejected while building, or dropped by the difficulty band alike -- with the - token-level record of the rounds that produced it. This is the only - way the proposing rounds survive: they are generation like any - other, so they carry ``input_ids`` / ``labels`` / ``logprobs`` and + token-level record of the episode that produced it. This is the only + way the proposing episode survives: it is generation like any + other, so it carries ``input_ids`` / ``labels`` / ``logprobs`` and could later be trained on, but nothing downstream of ``build`` - looks at them and without a sink they are dropped on the floor. + looks at it and without a sink it is dropped on the floor. Rejects are included on purpose: they are the zero-reward half of a GRPO group, so a set of kept-only records has no variance to learn from. Requires a local sampler -- an API explorer returns text only. + solver_sink: called once per solver attempt in the difficulty stage, with + the statement, the check script, the attempt, the workspace it left + and the check's verdict. ``n_pass`` alone cannot distinguish a task + that is impossible from one whose statement withholds a value its + check demands, and both look like a hard task worth keeping. + keyword_sink: called once per keyword-generation call, with the prompt, + the raw reply and what ``parse_keyword_list`` made of it. A bank that + refuses to fill is invisible otherwise -- proposals fall back to the + no-keyword prompt and the run carries on looking normal -- and a count + of zero does not say whether the model broke the format or the parser + rejected output that was fine. """ def __init__( @@ -237,20 +462,35 @@ def __init__( keyword_store: Optional[KeywordStore] = None, category_desc: Optional[Dict[str, str]] = None, seed_mix_prob: float = 0.5, - reset_fn: Callable[[], None], - run_check_fn: Callable[[str], Tuple[int, str]], - workspace_snapshot_fn: Optional[Callable[[], str]] = None, + reset_fn: Callable[..., None], + run_check_fn: Callable[..., Tuple[int, str]], + workspace_snapshot_fn: Optional[Callable[..., str]] = None, + tool_schemas: Optional[Sequence[Dict[str, Any]]] = None, + episode_concurrency: int = 1, + episode_tool_managers: Optional[Sequence[Any]] = None, combo_arity: str = 'triple', arity_weights: Optional[Sequence[float]] = None, single_kw_prob: float = 0.1, keyword_refill_target: int = 128, keyword_gen_calls: int = 8, + keyword_refill_concurrency: int = 1, keyword_refill_tries: int = 2, keyword_params: Optional[SamplingParams] = None, + check_params: Optional[SamplingParams] = None, + problem_params: Optional[SamplingParams] = None, + followup_api: Optional[Any] = None, + followup_extra_body: Optional[Dict[str, Any]] = None, + keyword_explorer: Optional[Explorer] = None, min_batch: int = 1, problem_max_chars: int = 8192, + max_proposals_total: int = 0, + setup_script_fn: Optional[Callable[..., str]] = None, + solver_prompt_fn: Optional[Callable[[str], Trajectory]] = None, + check_retries: int = 1, reject_sink: Optional[Callable[[Dict[str, Any]], None]] = None, propose_sink: Optional[Callable[[Dict[str, Any]], None]] = None, + solver_sink: Optional[Callable[[Dict[str, Any]], None]] = None, + keyword_sink: Optional[Callable[[Dict[str, Any]], None]] = None, **challenger_kwargs: Any, ): super().__init__(explorer, system=prompts.system, **challenger_kwargs) @@ -271,25 +511,120 @@ def __init__( self.reset_fn = reset_fn self.run_check_fn = run_check_fn self.workspace_snapshot_fn = workspace_snapshot_fn + self.tool_schemas = list(tool_schemas) if tool_schemas else None + # More than one episode at a time needs more than one sandbox: an episode + # owns its workspace from the reset until its check has run. The three + # sandbox callables above are then called with ``slot=i`` to say which one, + # and ``episode_tool_managers[i]`` must dispatch tool calls into that same + # sandbox -- an episode acting in one workspace and checking another + # produces a task whose check nobody can pass. + if episode_concurrency < 1: + raise ValueError(f'episode_concurrency must be >= 1, got {episode_concurrency}') + if episode_concurrency > 1: + if not episode_tool_managers or len(episode_tool_managers) != episode_concurrency: + raise ValueError( + f'episode_concurrency={episode_concurrency} needs exactly that many ' + f'episode_tool_managers, one per sandbox; got ' + f'{len(episode_tool_managers) if episode_tool_managers else 0}.') + self.episode_concurrency = episode_concurrency + self.episode_tool_managers = (list(episode_tool_managers) + if episode_tool_managers else None) + # Held while writing to the dump files and while bumping ``stats``: with + # concurrent episodes those are the only shared mutable things the + # follow-up callback touches, and a half-written json line is unreadable. + self._sink_lock = threading.Lock() + # Separate from the sink lock: the keyword path holds this while it draws + # from the shared rng and bumps the prompt nonce, and it must not be held + # while a sink write is waiting on disk. + self._kw_lock = threading.Lock() self.combo_arity = combo_arity self.arity_weights = list(arity_weights) if arity_weights else None self.single_kw_prob = single_kw_prob self.keyword_refill_target = keyword_refill_target self.keyword_gen_calls = keyword_gen_calls + # How many of a refill's generating calls go out together. At 1 each call + # is told what the ones before it produced, which is the point; raising it + # is what the first round of arm measurements ran with, where a whole + # first refill went out at once with nothing yet to avoid and came back + # with synonyms. Kept configurable so the two can be compared on one build + # rather than across two versions of this file. + if keyword_refill_concurrency < 1: + raise ValueError('keyword_refill_concurrency must be >= 1, got ' + f'{keyword_refill_concurrency}') + self.keyword_refill_concurrency = keyword_refill_concurrency self.keyword_refill_tries = keyword_refill_tries self.keyword_params = keyword_params + self.check_params = check_params + self.problem_params = problem_params + # When set, exploration runs on the (trainable) local explorer as before, + # but the check-script and problem-statement stages are generated by this + # OpenAI-compatible API (e.g. qwen3-max) instead of the local model. The + # two stages then contribute nothing to the trainable trajectory: the API + # returns text only, so the episode's ``input_ids`` / ``labels`` / + # ``logprobs`` stay exactly the exploration turns the local sampler + # produced -- which is what "train only the exploration part" means. The + # generated check script and statement are used solely to build the task. + self.followup_api = followup_api + # extra_body sent on every followup API call (e.g. {'thinking_budget': N} + # to cap qwen3-max reasoning). None sends the request unmodified. + self.followup_extra_body = dict(followup_extra_body) if followup_extra_body else None + self.keyword_explorer = keyword_explorer self.min_batch = max(1, min_batch) self.problem_max_chars = problem_max_chars + # A budget in proposals rather than in kept tasks, for runs whose purpose + # is to measure what the current configuration produces: with a keep-rate + # near 6% a keep-target of 8 is 128 proposals, and comparing two + # configurations means giving them the same number of tries, not the same + # output. 0 leaves the run governed by its keep-target. + self.max_proposals_total = max_proposals_total + # Arm B. Returns a python script that recreates this episode's input files, + # captured while the workspace still holds them, and replayed before every + # solver attempt. None leaves the solver starting from an empty directory. + self.setup_script_fn = setup_script_fn + # How a task statement becomes the solver's opening conversation. Without + # one, the solver is handed the statement as a bare user message and no + # system message at all -- nothing says it is working in a sandbox, that it + # may take many turns, or that a reply carries one tool call. Measured on + # arm B: 71 of 80 attempts used 2-3 turns, writing 8-12k characters into a + # single python_executor argument and truncating there, so ``n_pass`` was + # reporting that omission rather than the task. Passing the same function + # the eval script uses is what keeps the two measuring the same thing. + self.solver_prompt_fn = solver_prompt_fn + if check_retries < 0: + raise ValueError(f'check_retries must be >= 0, got {check_retries}') + self.check_retries = check_retries + if check_retries: + prompts.require('check_retry_followup') self.reject_sink = reject_sink self.propose_sink = propose_sink + self.solver_sink = solver_sink + self.keyword_sink = keyword_sink if self.seeds: prompts.require('from_seed') if self.store is not None: prompts.require('from_seed_keywords') self._nonce = 0 self.stats: Dict[str, int] = { - 'round1_done': 0, 'check_parse_fail': 0, 'check_run_fail': 0, + 'explore_done': 0, 'check_parse_fail': 0, 'check_run_fail': 0, + 'empty_workspace': 0, 'solver_truncated': 0, 'problem_parse_fail': 0, 'too_long': 0, 'parsed': 0, + # How often a check that failed was handed back for a rewrite, and + # how often the rewrite passed. The two together say whether the + # retry earns its extra sampling call. + 'check_retry': 0, 'check_retry_pass': 0, + # The episode ended before the appended stages could run or finish: + # it used up ``max_turns``, hit the trajectory token cap, or left no + # room for the follow-up message. Distinct from every other reason + # here, which is the model producing something unusable. + 'episode_cut_short': 0, + # Arm B only. ``setup_capture_fail``: the episode's input files could + # not be read back, so the task was dropped. ``setup_replay_fail``: a + # solver attempt was skipped because putting those files back failed, + # which would otherwise have scored as the task being too hard. + 'setup_capture_fail': 0, 'setup_replay_fail': 0, + # followup_api mode only: a check or statement API call failed. The + # conversation is then unusable and the proposal is rejected. + 'followup_api_error': 0, } self._hard: List[Tuple[str, str]] = [] @@ -302,6 +637,8 @@ def propose(self, count: int) -> List[Trajectory]: run these multi-turn in the sandbox. """ proposals: List[Trajectory] = [] + directions: List[str] = [] + metas: List[Tuple[List[Tuple[str, str]], bool, str]] = [] for _ in range(count): picks = self._draw_keywords() body = '\n'.join(f'- {c}: {t}' for c, t in picks) @@ -316,99 +653,352 @@ def propose(self, count: int) -> List[Trajectory]: user = self.prompts.from_keywords.format(keywords=body) else: user = self.prompts.from_scratch + directions.append(user) + metas.append((picks, use_seed, body)) + + for user, (picks, use_seed, body) in zip(directions, metas): proposal: Trajectory = { 'messages': [{'role': 'system', 'content': self.prompts.system}, {'role': 'user', 'content': user}], } + if self.tool_schemas: + proposal['tools'] = self.tool_schemas proposals.append(attach_user_data( proposal, keywords=picks, seeded=use_seed, keyword_block=body)) return proposals + # ------------------------------------------------------------- building def build(self, explored: List[Trajectory]) -> List[Optional[Trajectory]]: """Satisfy the abstract method; not usable outside ``_round``. - ``_build_one`` requires the sandbox to hold the episode's workspace state, - which is only guaranteed inside the serial ``_round`` loop. Calling this - method directly will produce wrong results because the sandbox state does - not match the trajectory being processed. + Building happens inside the episode now: :meth:`_followup` runs while the + model is still generating and needs the sandbox to hold that episode's + workspace state, which is only guaranteed inside the serial ``_round`` + loop. """ raise RuntimeError( f'{type(self).__name__}.build() must not be called directly; ' - f'the serial _round() loop calls _build_one() per episode instead.') + f'the serial _round() loop drives one episode at a time instead.') + + def _followup(self, state: Dict[str, Any], trajectory: Trajectory, + n_before: int) -> Optional[Tuple[str, Optional[SamplingParams]]]: + """What to say next when the model stops calling tools; ``None`` to stop. + + The rollout calls this once per stage, handing over the episode as it + stands. ``state`` is this episode's scratchpad, read afterwards by + :meth:`_finish_episode`: the workspace listing and the check script are + produced here, and anything that goes wrong before a statement exists is + left in ``state['reject']``. + + The sandbox work has to happen at this moment and nowhere else -- the + workspace holds this episode's end state right now, and the next + episode's reset wipes it. ``state['slot']`` says which sandbox that is; + with concurrent episodes several of these run at once, each against its + own. + """ + slot = state.get('slot', 0) + if n_before == 0: + snapshot = self.workspace_snapshot_fn(slot=slot) if self.workspace_snapshot_fn else '' + state['snapshot'] = snapshot + # An episode that left nothing behind has no end state to write checks + # about, and asking for them anyway is worse than useless: the only + # true thing to assert is that the directory is empty, which every + # solver passes by doing nothing. Five of run5's ten verified tasks + # were that task. Reject here instead. + if not snapshot.strip(): + self._bump('empty_workspace') + state['reject'] = ('empty_workspace', '') + return None + return (self.prompts.check_followup.format(final_state=snapshot), + self.check_params) + + # Every follow-up from here until a check passes is a check-script reply: + # the first one, plus up to ``check_retries`` rewrites. + if not state.get('checked'): + attempt = state.get('check_attempts', 0) + 1 + state['check_attempts'] = attempt + reply = assistant_text(trajectory) + script = parse_check_script(reply) + if script is None: + # Same one-rewrite budget a run failure gets: hand the parse + # failure back and let it regenerate, rather than dropping a task + # whose only fault was packaging. Shares the check_attempts + # count, so parse and run failures together get check_retries + # extra tries, not one each. + if attempt <= self.check_retries: + self._bump('check_retry') + err = ('Could not read a check script from your reply: it was ' + 'not a fenced python code block. Do not wrap it in a ' + 'tool call and do not add prose -- return ONLY a fenced ' + 'python code block.') + return (self.prompts.check_retry_followup.format( + error=err, final_state=state.get('snapshot') or ''), + self.check_params) + self._bump('check_parse_fail') + # The whole reply, not a tail: this stage fails either because the + # model declared the state untestable (it says so) or because it + # ran out of tokens while thinking, and a record that cannot tell + # them apart sends the next reader back to re-run the batch. + state['reject'] = ('check_parse_fail', reply) + return None + state['script'] = script + brittle = brittle_check_reason(script) + if brittle is not None: + # Same bookkeeping as a check that ran and failed: the script is + # rejected before it can pass on the author's own state, because + # passing there is exactly what hides the defect. + exit_code, output = 1, brittle + else: + exit_code, output = self.run_check_fn(script, slot=slot) + if exit_code == 0: + state['checked'] = True + if attempt > 1: + self._bump('check_retry_pass') + # Capture the inputs now, at the one moment the workspace holds + # exactly the state this check just passed on. A capture after the + # statement stage would be the same bytes only by luck. + if self.setup_script_fn is not None: + setup = self.setup_script_fn(slot=slot) + if not setup: + self._bump('setup_capture_fail') + state['reject'] = ( + 'setup_capture_fail', + 'no input files to hand the solver, or their bytes ' + 'could not be read back') + return None + state['setup_script'] = setup + return (self.prompts.problem_followup, self.problem_params) + # Snapshot again, after the failure. A check that asserts only + # paths from the snapshot it was shown and still fails leaves two + # very different bugs indistinguishable -- the model asserted + # something untrue, or the workspace changed under it -- and the + # difference is visible only in the state at the moment the check + # ran. It is also what the rewrite gets to read. + after = self.workspace_snapshot_fn(slot=slot) if self.workspace_snapshot_fn else '' + state.setdefault('attempts', []).append( + f'--- attempt {attempt}: exit {exit_code} ---\n{output}\n' + f'--- check script ---\n{script}') + if attempt <= self.check_retries: + self._bump('check_retry') + return (self.prompts.check_retry_followup.format( + error=output, final_state=after or state.get('snapshot') or ''), + self.check_params) + self._bump('check_run_fail') + state['reject'] = ( + 'check_run_fail', + '\n'.join(state['attempts']) + + f"\n--- state before check ---\n{state.get('snapshot') or ''}\n" + + f'--- state after check ---\n{after}') + return None - def _build_one(self, explored: Trajectory) -> Optional[Trajectory]: - """Process one round-1 result: write checks, verify, write problem. + return None + + def _api_reply(self, messages: List[Dict[str, Any]], user_text: str, + params: Optional[SamplingParams]) -> Optional[str]: + """Append ``user_text`` and one ``followup_api`` reply to ``messages``. + + ``messages`` is a throwaway copy owned by :meth:`_run_followup_api`, never + the trainable trajectory, so mutating it in place costs the model nothing. + Returns the assistant text, or ``None`` when the API call raised -- the + caller then rejects rather than building a task on a broken conversation. + + Tools are withdrawn for these stages on purpose (they are answers, not + actions), so only the text is kept; any structured ``tool_calls`` the API + returned are dropped. + """ + messages.append({'role': 'user', 'content': user_text}) + request: Trajectory = {'messages': messages} + try: + if self.followup_extra_body: + reply = self.followup_api(request, params, extra_body=self.followup_extra_body) + else: + reply = self.followup_api(request, params) + except Exception as exc: # noqa: BLE001 -- one bad call must not kill the round + logger.warning(f'[{type(self).__name__}] followup API call failed: ' + f'{type(exc).__name__}: {exc}') + return None + if isinstance(reply, list): + reply = reply[0] if reply else {} + content = (reply.get('content') if isinstance(reply, dict) else None) or '' + messages.append({'role': 'assistant', 'content': content}) + return content + + def _run_followup_api(self, state: Dict[str, Any], explored: Trajectory) -> None: + """Generate the check script and problem statement over ``followup_api``. + + The API-only counterpart of :meth:`_followup`: the same stages, the same + sandbox work (snapshot, run the check, capture inputs) and the same retry + budget, but driven imperatively here instead of turn-by-turn by the + rollout, and answered by the API rather than the local model. Results land + in ``state`` for :meth:`_finish_episode`: + + * ``state['script']`` / ``state['checked']`` -- the check that passed, + * ``state['setup_script']`` -- captured inputs (Arm B), + * ``state['statement']`` -- the problem-statement text, + * ``state['reject']`` -- ``(reason, detail)`` when a stage fails. + + Nothing here touches ``explored``'s ``input_ids`` / ``labels`` / + ``logprobs``: the messages the API sees are a private copy, so the + trainable trajectory stays exactly the exploration turns the local sampler + produced. + """ + slot = state.get('slot', 0) + messages: List[Dict[str, Any]] = [dict(m) for m in explored.get('messages') or []] + + snapshot = self.workspace_snapshot_fn(slot=slot) if self.workspace_snapshot_fn else '' + state['snapshot'] = snapshot + # An episode that left nothing behind has no end state to write checks + # about; rejecting here mirrors the n_before==0 branch of _followup. + if not snapshot.strip(): + self._bump('empty_workspace') + state['reject'] = ('empty_workspace', '') + return - Called while the sandbox still holds this episode's workspace state. + # Check-script stage: the first ask plus up to ``check_retries`` rewrites, + # sharing one attempt counter across parse and run failures exactly as the + # single-model path does. + user_text = self.prompts.check_followup.format(final_state=snapshot) + attempt = 0 + while True: + attempt += 1 + state['check_attempts'] = attempt + reply = self._api_reply(messages, user_text, self.check_params) + if reply is None: + self._bump('followup_api_error') + state['reject'] = ('followup_api_error', 'check-script API call failed') + return + script = parse_check_script(reply) + if script is None: + if attempt <= self.check_retries: + self._bump('check_retry') + err = ('Could not read a check script from your reply: it was ' + 'not a fenced python code block. Do not wrap it in a ' + 'tool call and do not add prose -- return ONLY a fenced ' + 'python code block.') + user_text = self.prompts.check_retry_followup.format( + error=err, final_state=snapshot) + continue + self._bump('check_parse_fail') + state['reject'] = ('check_parse_fail', reply) + return + state['script'] = script + brittle = brittle_check_reason(script) + if brittle is not None: + # Rejected before it can pass on the author's own state, since + # passing there is exactly what hides the defect. + exit_code, output = 1, brittle + else: + exit_code, output = self.run_check_fn(script, slot=slot) + if exit_code == 0: + state['checked'] = True + if attempt > 1: + self._bump('check_retry_pass') + # Capture inputs now, while the workspace still holds the state + # this check just passed on. + if self.setup_script_fn is not None: + setup = self.setup_script_fn(slot=slot) + if not setup: + self._bump('setup_capture_fail') + state['reject'] = ( + 'setup_capture_fail', + 'no input files to hand the solver, or their bytes ' + 'could not be read back') + return + state['setup_script'] = setup + break + after = self.workspace_snapshot_fn(slot=slot) if self.workspace_snapshot_fn else '' + state.setdefault('attempts', []).append( + f'--- attempt {attempt}: exit {exit_code} ---\n{output}\n' + f'--- check script ---\n{script}') + if attempt <= self.check_retries: + self._bump('check_retry') + user_text = self.prompts.check_retry_followup.format( + error=output, final_state=after or snapshot) + continue + self._bump('check_run_fail') + state['reject'] = ( + 'check_run_fail', + '\n'.join(state['attempts']) + + f"\n--- state before check ---\n{state.get('snapshot') or ''}\n" + + f'--- state after check ---\n{after}') + return + + # Problem-statement stage: one API reply, kept as the task's statement. + reply = self._api_reply(messages, self.prompts.problem_followup, self.problem_params) + if reply is None: + self._bump('followup_api_error') + state['reject'] = ('followup_api_error', 'problem-statement API call failed') + return + state['statement'] = reply + + def _finish_episode(self, state: Dict[str, Any], + explored: Trajectory) -> Optional[Trajectory]: + """Turn a finished episode into a task, or record why it is not one. + + Everything the model wrote is in ``explored``: the tool-using turns, the + check script, and the problem statement as the last assistant message. + ``state`` carries what only the sandbox could say -- the end state, and + whether the check passed on it. """ - summary = _trajectory_summary(explored) - snapshot = self.workspace_snapshot_fn() if self.workspace_snapshot_fn else summary keywords = user_data_get(explored.get('user_data'), 'keywords', []) seeded = user_data_get(explored.get('user_data'), 'seeded', False) - # Every round this proposal generates, in order. Handed to propose_sink - # with whatever verdict the proposal ends up with, so a rejected attempt - # is recorded as fully as a kept one. - rounds = [_propose_round('explore', explored)] - - # Round 2a: write check script - # NOTE: This goes through the same explorer (with tool schemas visible). - # The prompt must clearly instruct the model to output ONLY a code block - # and not call tools, otherwise tool calls would corrupt the sandbox state - # before verification. The check_system prompt enforces this. - check_prompt: Trajectory = { - 'messages': [ - {'role': 'system', 'content': self.prompts.check_system}, - {'role': 'user', 'content': self.prompts.check_user.format( - trajectory=summary, final_state=snapshot)}, - ], - } - check_reply = self.explore([check_prompt]) - rounds.append(_propose_round('check', check_reply[0])) - script = parse_check_script(assistant_text(check_reply[0])) - if script is None: - self.stats['check_parse_fail'] += 1 - self._reject_record(explored, 'check_parse_fail') - self._emit_propose(rounds, 'check_parse_fail', keywords=keywords, seeded=seeded) + # The episode as one record: a single conversation, so a single set of + # token ids and logprobs. Handed to propose_sink with whatever verdict the + # proposal ends up with, so a rejected attempt is recorded as fully as a + # kept one. + rounds = [_propose_round('episode', explored)] + + def reject(reason: str, detail: str = '') -> None: + self._reject_record(explored, reason, detail=detail) + self._emit_propose(rounds, reason, keywords=keywords, seeded=seeded) + + if state.get('reject'): + reason, detail = state['reject'] + reject(reason, detail) return None - # Verify: run check script in current sandbox state (must pass) - exit_code, output = self.run_check_fn(script) - if exit_code != 0: - self.stats['check_run_fail'] += 1 - self._reject_record(explored, 'check_run_fail', - detail=f'exit {exit_code}: {output[-200:]}') - self._emit_propose(rounds, 'check_run_fail', keywords=keywords, seeded=seeded) + if not state.get('checked'): + # The stages never ran, or the check-writing one never got a reply: + # the episode used up its turns, hit the trajectory token cap, or left + # no room to append the next message. The model produced nothing + # wrong here, so this is not one of the other reasons. Keyed on the + # check having *passed* rather than on a script existing: a rewrite + # that never came back leaves the failed script in ``state``, and + # building a task on it would ship a check nobody can pass. + self._bump('episode_cut_short') + reject('episode_cut_short', + detail=f"stop_reason={explored.get('stop_reason')} " + f"truncated={bool(explored.get('truncated'))} " + f"turns={explored.get('turns')} " + f"followups={explored.get('followups')}") return None - # Round 2b: write problem statement - problem_prompt: Trajectory = { - 'messages': [ - {'role': 'system', 'content': self.prompts.problem_system}, - {'role': 'user', 'content': self.prompts.problem_user.format( - trajectory=summary, checks=script)}, - ], - } - problem_reply = self.explore([problem_prompt]) - rounds.append(_propose_round('problem', problem_reply[0])) - statement = parse_problem_statement(assistant_text(problem_reply[0])) + script = state['script'] + # In followup_api mode the statement was written by the API and is not in + # ``explored`` (whose last assistant turn is the final exploration reply); + # it lives in ``state``. The single-model path keeps it as the last + # assistant message of the episode. + if self.followup_api is not None: + statement = parse_problem_statement(state.get('statement') or '') + else: + statement = parse_problem_statement(assistant_text(explored)) if statement is None: - self.stats['problem_parse_fail'] += 1 - self._reject_record(explored, 'problem_parse_fail') - self._emit_propose(rounds, 'problem_parse_fail', keywords=keywords, seeded=seeded) + self._bump('problem_parse_fail') + reject('problem_parse_fail') return None if len(statement) > self.problem_max_chars: - self.stats['too_long'] += 1 - self._reject_record(explored, 'too_long') - self._emit_propose(rounds, 'too_long', keywords=keywords, seeded=seeded) + self._bump('too_long') + reject('too_long') return None - self.stats['parsed'] += 1 + self._bump('parsed') task: Trajectory = { 'messages': [{'role': 'user', 'content': statement}], } - task = attach_user_data(task, check_script=script, keywords=keywords, seeded=seeded) + task = attach_user_data(task, check_script=script, keywords=keywords, seeded=seeded, + setup_script=state.get('setup_script', '')) # Carried, not emitted: the verdict this proposal earns depends on the # difficulty measurement, which has not run yet. A plain top-level key # rather than user_data, which json-encodes every value on each update. @@ -416,11 +1006,30 @@ def _build_one(self, explored: Trajectory) -> Optional[Trajectory]: return task def _reject_record(self, traj: Trajectory, reason: str, detail: str = '') -> None: - if self.reject_sink is not None: - payload: Dict[str, Any] = {'reason': reason} - if detail: - payload['detail'] = detail - payload['last_assistant'] = assistant_text(traj)[:500] + """Record a rejected proposal, with enough of the episode to tell why. + + The reason alone is not diagnosable. Nine ``empty_workspace`` rejections in + one run all looked like the model refusing to act; the messages showed a + single assistant turn each, and the question of whether it had run out of + tokens or simply emitted no call could not be answered from the record -- + the fields that answered it were on the trajectory and were dropped. So + how the episode ended travels with the reason. + """ + if self.reject_sink is None: + return + messages = traj.get('messages') or [] + payload: Dict[str, Any] = {'reason': reason} + if detail: + payload['detail'] = detail + payload['stop_reason'] = traj.get('stop_reason') + payload['truncated'] = bool(traj.get('truncated')) + payload['turns'] = traj.get('turns') + payload['n_assistant'] = sum(1 for m in messages + if isinstance(m, dict) and m.get('role') == 'assistant') + payload['n_tool_calls'] = sum(len(m.get('tool_calls') or []) for m in messages + if isinstance(m, dict)) + payload['last_assistant'] = assistant_text(traj) + with self._sink_lock: self.reject_sink(payload) def _emit_propose(self, rounds: Optional[List[Dict[str, Any]]], outcome: str, *, @@ -437,7 +1046,7 @@ def _emit_propose(self, rounds: Optional[List[Dict[str, Any]]], outcome: str, *, if self.propose_sink is None or not rounds: return rollouts = self.solver_rollouts or None - self.propose_sink({ + payload = { 'outcome': outcome, 'n_pass': n_pass, 'n_rollouts': rollouts, @@ -445,7 +1054,9 @@ def _emit_propose(self, rounds: Optional[List[Dict[str, Any]]], outcome: str, *, 'keywords': list(keywords or ()), 'seeded': bool(seeded), 'rounds': rounds, - }) + } + with self._sink_lock: + self.propose_sink(payload) def _take_rounds(self, task: Trajectory) -> Optional[List[Dict[str, Any]]]: """Detach a task's proposing rounds. Popped even with no sink attached: @@ -456,26 +1067,115 @@ def _take_rounds(self, task: Trajectory) -> Optional[List[Dict[str, Any]]]: # ------------------------------------------------------------ revised _round + def _bump(self, key: str, n: int = 1) -> None: + """Thread-safe stats increment.""" + with self._sink_lock: + self.stats[key] += n + + def _parallel(self, fn: Callable[[Any], Any], items: Sequence[Any]) -> List[Any]: + """Map ``fn`` over ``items`` at once, results in input order. + + Every use of this is waiting on a sandbox, not computing, so the thread + pool is the point. One item runs inline: a pool for a single sandbox call + only adds a thread, and it keeps the serial configuration on exactly the + same code path it had before. + """ + items = list(items) + if len(items) <= 1: + return [fn(item) for item in items] + out: List[Any] = [None] * len(items) + with ThreadPoolExecutor(max_workers=len(items)) as pool: + futures = {pool.submit(fn, item): i for i, item in enumerate(items)} + for fut in as_completed(futures): + out[futures[fut]] = fut.result() + return out + + def _run_episode(self, proposal: Trajectory, slot: int) -> Optional[Trajectory]: + """One episode top-to-bottom, using sandbox slot ``slot``.""" + self.reset_fn(slot=slot) + state: Dict[str, Any] = {'slot': slot} + tm = self.episode_tool_managers[slot] if self.episode_tool_managers else None + if self.followup_api is not None: + # Split path: explore on the local (trainable) model with NO + # followup_fn, so the rollout ends the moment the model stops calling + # tools and the returned trajectory carries only the exploration + # turns' input_ids/labels/logprobs. The check-script and + # problem-statement stages then run over the API against the same end + # state, appended to a throwaway copy of the messages -- never to the + # trainable trajectory. + kwargs: Dict[str, Any] = {} + if tm is not None: + kwargs['tool_manager'] = tm + result = self.explore([proposal], **kwargs) + if not result: + return None + explored = result[0] + self._bump('explore_done') + # A reply cut off at the token budget never finished its thought, so + # continuing the conversation over the API would build a check on a + # half-written turn. Leave ``state`` untouched and let + # ``_finish_episode`` record it as ``episode_cut_short``, matching the + # single-model path which does not run the stages after a length cut. + if explored.get('stop_reason') != 'length': + self._run_followup_api(state, explored) + return self._finish_episode(state, explored) + kwargs = {'followup_fn': partial(self._followup, state)} + if tm is not None: + kwargs['tool_manager'] = tm + result = self.explore([proposal], **kwargs) + if not result: + return None + explored = result[0] + self._bump('explore_done') + return self._finish_episode(state, explored) + def _round(self, missing: int) -> Optional[List[Trajectory]]: - """One cycle: serial round-1 episodes, inline build, then difficulty filter.""" + """One cycle: episodes in parallel across sandbox slots, then difficulty.""" count = min(self._estimate(missing), self.max_proposals_per_round) + if self.max_proposals_total > 0: + left = self.max_proposals_total - self.n_proposed + if left <= 0: + # Budget spent. None is the 'source exhausted' answer the batching + # loop already knows how to stop on, so the run ends after this + # round's keepers are handed back rather than mid-episode. + logger.info(f'[{type(self).__name__}] proposal budget spent ' + f'({self.n_proposed}/{self.max_proposals_total}); stopping') + return None + count = min(count, left) proposals = self.propose(count) if not proposals: return None usable: List[Trajectory] = [] - for proposal in proposals: - # Reset workspace, run round 1 - self.reset_fn() - result = self.explore([proposal]) - if not result: - continue - explored = result[0] - self.stats['round1_done'] += 1 - # Workspace still holds this episode's state → build inline - task = self._build_one(explored) - if task is not None: - usable.append(task) + n_slots = self.episode_concurrency + + if n_slots <= 1 or len(proposals) <= 1: + # Serial fallback (original path). + for proposal in proposals: + task = self._run_episode(proposal, slot=0) + if task is not None: + usable.append(task) + else: + # One worker per sandbox slot, each draining its own share serially. + # A slot is a single sandbox and cannot host two episodes at once, so + # the split is by slot, never round-robin into a shared pool where two + # tasks could land on the same slot concurrently. + buckets: List[List[Trajectory]] = [[] for _ in range(n_slots)] + for i, proposal in enumerate(proposals): + buckets[i % n_slots].append(proposal) + + def _drain(slot: int) -> List[Trajectory]: + out: List[Trajectory] = [] + for proposal in buckets[slot]: + task = self._run_episode(proposal, slot=slot) + if task is not None: + out.append(task) + return out + + with ThreadPoolExecutor(max_workers=n_slots) as pool: + futures = [pool.submit(_drain, s) for s in range(n_slots) if buckets[s]] + for fut in as_completed(futures): + usable.extend(fut.result()) kept = self._filter_difficulty(usable) if self.solver_rollouts else usable if not self.solver_rollouts: @@ -494,18 +1194,86 @@ def _round(self, missing: int) -> Optional[List[Trajectory]]: # ------------------------------------------------------------ difficulty def _filter_difficulty(self, tasks: List[Trajectory]) -> List[Trajectory]: - """Override: each solver attempt needs a clean workspace, so run serially.""" + """Override: every solver attempt needs its own clean workspace. + + Attempts are run in waves of ``episode_concurrency``, attempt k of a wave + in sandbox slot k. Within a wave all attempts go out in one explorer call, + so the sampler generates them as one batch instead of leaving the GPUs + waiting on a single sequence, and the wave's clears, input replays and + checks all run at the same time too -- they are sandbox round-trips, not + compute. + + The slot is what keeps this honest: a wave's attempts each clear, act in + and get checked against their own sandbox. Sharing one would let attempt A + pass on files attempt B wrote, and ``n_pass`` would stop being a + difficulty measurement. + + An attempt cut off at the generation budget is counted in + ``stats['solver_truncated']`` but still counts as a failure, because + deciding otherwise decides which tasks are kept. Watch that number: when + it is a large share of ``solver_rollouts`` times the task count, ``n_pass`` + is reporting the token budget rather than the difficulty. It was 15 of 50 + on one run, and one task lost all four attempts that way and was discarded + as too hard without a solver ever touching the workspace. Raising + ``solver_params.max_tokens`` took it to 0 of 20. + """ if not tasks: return [] passes = [0] * len(tasks) - for i, task in enumerate(tasks): - prompt = self.solver_prompt(task) - for _ in range(self.solver_rollouts): - self.reset_fn() - attempts = self._solver_explore( - [dict(prompt)], sampling_params=self.solver_params) - if attempts and self.judge_attempt(task, attempts[0]): - passes[i] += 1 + n_slots = max(1, self.episode_concurrency) + # Which task each attempt belongs to, flattened, so a wave is a fixed + # number of sandboxes no matter how attempts distribute over tasks. + plan = [i for i in range(len(tasks)) for _ in range(self.solver_rollouts)] + + for start in range(0, len(plan), n_slots): + wave = plan[start:start + n_slots] + slots = list(range(len(wave))) + setups = [user_data_get(tasks[i].get('user_data'), 'setup_script', '') + for i in wave] + + def _prepare(k: int) -> bool: + """Clear slot k, then put back the inputs this task hands out.""" + self.reset_fn(slot=k) + if not setups[k]: + return True + exit_code, output = self.run_check_fn(setups[k], slot=k) + if exit_code != 0: + # Measuring this attempt against a workspace missing its + # inputs would score the task as harder than it is, so the + # attempt is skipped and counted rather than run. + logger.warning(f'[{type(self).__name__}] input setup failed in ' + f'slot {k} (exit {exit_code}): {output[-200:]}') + return False + return True + + ready = self._parallel(_prepare, slots) + live = [k for k in slots if ready[k]] + self._bump('setup_replay_fail', len(slots) - len(live)) + if not live: + continue + prompts = [dict(self.solver_prompt(tasks[wave[k]])) for k in live] + kwargs: Dict[str, Any] = {} + if self.episode_tool_managers: + kwargs['tool_manager'] = [self.episode_tool_managers[k] for k in live] + attempts = self._solver_explore(prompts, sampling_params=self.solver_params, + **kwargs) + if len(attempts) != len(prompts): + # Counting a partial return would silently understate every + # affected task's pass count, i.e. report tasks as harder than + # they are. + raise RuntimeError(f'explorer returned {len(attempts)} attempts for ' + f'{len(prompts)} solver prompts; expected one per prompt.') + for attempt in attempts: + if attempt is not None and attempt.get('truncated'): + self._bump('solver_truncated') + verdicts = self._parallel( + lambda j: (attempts[j] is not None + and self.judge_attempt(tasks[wave[live[j]]], attempts[j], + slot=live[j])), + list(range(len(live)))) + for j, passed in enumerate(verdicts): + if passed: + passes[wave[live[j]]] += 1 measured = [ attach_user_data(task, n_pass=passes[i], n_rollouts=self.solver_rollouts) @@ -527,15 +1295,65 @@ def _filter_difficulty(self, tasks: List[Trajectory]) -> List[Trajectory]: return [t for t, kept_flag in zip(measured, in_band) if kept_flag] def solver_prompt(self, task: Trajectory) -> Trajectory: - """The task statement, nothing else -- the solver's own harness adds the system.""" - return {'messages': [dict(m) for m in task.get('messages') or []]} + """The statement as the solver first sees it: system message, query, tools. - def judge_attempt(self, task: Trajectory, attempt: Trajectory) -> bool: - """Run the check script against the sandbox's current state.""" + ``solver_prompt_fn`` is how the surrounding script hands over the same + opening the eval script builds, so a task kept at n_pass=4 here is a task + the eval measures the same way. Without one this falls back to the bare + statement, which is what it used to be. + + The schemas travel with the prompt for the same reason they do in round 1: + a solver that cannot see the tool names cannot use them, and would score + zero on every task regardless of difficulty. + """ + if self.solver_prompt_fn is not None: + messages = task.get('messages') or [] + query = next((m.get('content', '') for m in messages + if m.get('role') == 'user'), '') + prompt = self.solver_prompt_fn(query) + if not prompt.get('tools') and self.tool_schemas: + prompt['tools'] = self.tool_schemas + return prompt + prompt: Trajectory = {'messages': [dict(m) for m in task.get('messages') or []]} + if self.tool_schemas: + prompt['tools'] = self.tool_schemas + return prompt + + def judge_attempt(self, task: Trajectory, attempt: Trajectory, + slot: int = 0) -> bool: + """Run the check script against sandbox ``slot``'s current state. + + Also hands the whole attempt to ``solver_sink`` when one is given. The + difficulty stage otherwise reports a single number per task, and + ``n_pass=0`` reads the same whether the task is impossible, the statement + withholds something the check demands, or the solver merely gave up -- + which are three different things to fix. The evidence that separates them + is the attempt itself and the state it left, so both are recorded here + rather than reconstructed later. + """ script = user_data_get(task.get('user_data'), 'check_script', '') if not script: return False - exit_code, _ = self.run_check_fn(script) + exit_code, output = self.run_check_fn(script, slot=slot) + if self.solver_sink is not None: + messages = task.get('messages') or [{}] + record = { + 'statement': messages[0].get('content', ''), + 'check_script': script, + 'passed': exit_code == 0, + 'check_exit': exit_code, + 'check_output': output, + # Whether the reply was cut off at the generation budget. The + # difficulty stage drops such an attempt from its denominator, so + # the flag has to travel with the record for the dropped count to + # be reproducible from the dump. + 'truncated': bool((attempt or {}).get('truncated')), + 'attempt': attempt, + 'end_state': (self.workspace_snapshot_fn(slot=slot) + if self.workspace_snapshot_fn else ''), + } + with self._sink_lock: + self.solver_sink(record) return exit_code == 0 def on_difficulty_measured(self, candidates: List[Trajectory]) -> None: @@ -573,22 +1391,44 @@ def _draw_keywords(self) -> List[Tuple[str, str]]: else: cats = list(categories) picks: List[Tuple[str, str]] = [] + # Refill every dry category at once rather than as each one is reached: the + # three refills are independent model calls that used to run one after + # another (20s each at the start of a run), and they touch separate + # entries of the bank. + dry = [c for c in cats if not self.store.unused(c)] + if dry: + self._parallel(self._refill, dry) for c in cats: - if not self.store.unused(c): - self._refill(c) text = self.store.take(c, self.rng) if text is not None: picks.append((c, text)) return picks def _refill(self, category: str) -> None: - """Ask the model for more keywords in ``category``.""" + """Ask the model for more keywords in ``category``. + + Says so when it comes back empty. A silent no-op here is the worst + outcome available: ``_draw_keywords`` then hands out no keywords, every + proposal quietly falls back to the from-scratch prompt, and the run looks + normal while producing one identical prompt over and over. That is exactly + what happened for whole runs when the prompt asked for one keyword per + line and the parser wanted a JSON array. + """ tries = 0 while not self.store.unused(category): new = self._generate_keywords(category, self.keyword_refill_target) added = self.store.add(category, new, source='gen') tries += 1 - if added == 0 and tries >= self.keyword_refill_tries: + if added: + logger.info(f'[AgenticChallenger] keyword category {category!r} ' + f'refilled +{added} (try {tries})') + continue + logger.warning( + f'[AgenticChallenger] keyword refill for {category!r} produced ' + f'nothing on try {tries}: {len(new)} parsed, 0 new. Proposals will ' + f'run without keywords unless this recovers -- pass keyword_sink ' + f'to see the replies.') + if tries >= self.keyword_refill_tries: if self.store.items[category]: self.store.recycle(category) logger.info(f'[AgenticChallenger] keyword category {category!r} ' @@ -596,35 +1436,104 @@ def _refill(self, category: str) -> None: break def _generate_keywords(self, category: str, n_want: int) -> List[str]: - """Up to ``n_want`` keywords the bank does not already hold.""" + """Up to ``n_want`` keywords the bank does not already hold. + + Runs on ``keyword_explorer`` when there is one: brainstorming a list is a + text round, and putting it through the sandbox-tool explorer both wastes + turns and lets a bracketed list in the reply be taken for a tool call. + """ if n_want <= 0: return [] known = self.store.texts(category) n_calls = max(self.keyword_gen_calls, self.min_batch) per_call = max(1, -(-n_want // n_calls) + 4) - avoid_note = '' - if known: - shown = known if len(known) <= 40 else self.rng.sample(known, 40) - avoid_note = ('\nDo NOT repeat any of these already-used topics: ' - + ', '.join(shown)) - base = self.prompts.keyword_user.format( - k=per_call, desc=self.category_desc[category]) + avoid_note - self._nonce += 1 - prompts = [{ - 'messages': [{'role': 'system', 'content': self.prompts.keyword_system}, - {'role': 'user', 'content': f'{base}\n(batch {self._nonce}-{i})'}], - } for i in range(n_calls)] seen = {t.strip().lower() for t in known} out: List[str] = [] - for reply in self.explore(prompts, sampling_params=self.keyword_params): - for kw in parse_keyword_list(assistant_text(reply)): - key = kw.lower() - if key not in seen: - seen.add(key) - out.append(kw) - self.rng.shuffle(out) + explorer = self.keyword_explorer or self.explorer + for start in range(0, n_calls, self.keyword_refill_concurrency): + group = range(start, min(start + self.keyword_refill_concurrency, n_calls)) + # Every call in a group is built before any of them runs, so they all + # carry the same avoid list -- which is exactly the batched behaviour, + # and why a group of one is what lets call k+1 see call k. + users = [(self.prompts.keyword_user.format( + k=per_call, desc=self.category_desc[category]) + + self._avoid_note(known, out, + '\nDo NOT repeat any of these already-used topics: ') + + f'\n(batch {self._next_nonce()}-{i})') for i in group] + prompts = [{ + 'messages': [{'role': 'system', 'content': self.prompts.keyword_system}, + {'role': 'user', 'content': u}], + } for u in users] + for user, reply in zip(users, explorer(prompts, + sampling_params=self.keyword_params)): + text = assistant_text(reply) + parsed = parse_keyword_list(text) + fresh = [] + for kw in parsed: + key = kw.lower() + if key not in seen: + seen.add(key) + fresh.append(kw) + out.extend(fresh) + if self.keyword_sink is not None: + # Full text, both sides. The one question this dump exists to + # answer -- did the model disobey the format, or does the parser + # reject what it produced -- cannot be answered from a count. + record = { + 'category': category, + 'prompt': user, + 'reply': text, + 'stop_reason': reply.get('stop_reason'), + 'truncated': bool(reply.get('truncated')), + 'parsed': parsed, + 'n_parsed': len(parsed), + 'n_new': len(fresh), + } + with self._sink_lock: + self.keyword_sink(record) + with self._kw_lock: + self.rng.shuffle(out) return out[:n_want] + # How many phrases the 'do not repeat these' line may quote in total. There + # has to be a ceiling in both directions: too few and a serial refill stops + # seeing what it just said, too many and the model runs out of room to obey. + # Measured on armA2ser, where this refill's own output went in uncapped: with + # 130 quoted the eighth call was still answering normally, with 150 it started + # inventing -- 'îRAPIÓN holistic replace', 'ซะ subspace cutter map limit', 10 + # of 480 phrases that run. 100 sits below where that began. + _AVOID_TOTAL = 100 + + def _next_nonce(self) -> int: + """A number no other call gets, so two prompts are never byte-identical. + + Shared across categories, which refill at the same time: two threads + reading the counter together would send the same prompt twice and halve + the diversity with nothing to show that it happened. + """ + with self._kw_lock: + self._nonce += 1 + return self._nonce + + def _avoid_note(self, older: List[str], fresh: List[str], lead: str) -> str: + """The 'do not repeat these' line, newest first, capped at ``_AVOID_TOTAL``. + + What this refill has just produced comes first and evicts older entries + rather than the reverse -- the calls run one at a time so that each can + avoid what the ones before it said, and dropping those would undo it. Past + the cap the oldest of *this refill's* phrases are what falls off, which is + also the least costly thing to drop: the model has already moved away from + them. + """ + fresh_shown = list(fresh)[-self._AVOID_TOTAL:] + room = max(0, self._AVOID_TOTAL - len(fresh_shown)) + with self._kw_lock: + shown = older if len(older) <= room else self.rng.sample(older, room) + avoid = fresh_shown + list(shown) + return lead + ', '.join(avoid) if avoid else '' + + + # ------------------------------------------------------------ feedback def expand_hard_keywords(self) -> int: @@ -647,10 +1556,19 @@ def expand_hard_keywords(self) -> int: ], } for i, (_c, kw) in enumerate(reqs)] added = 0 - for (cat, kw), reply in zip(reqs, self.explore(prompts, - sampling_params=self.keyword_params)): - added += self.store.add(cat, parse_keyword_list(assistant_text(reply)), - source='expand', parent=kw) + explorer = self.keyword_explorer or self.explorer + for (cat, kw), reply in zip(reqs, explorer(prompts, + sampling_params=self.keyword_params)): + text = assistant_text(reply) + parsed = parse_keyword_list(text) + added += self.store.add(cat, parsed, source='expand', parent=kw) + if self.keyword_sink is not None: + self.keyword_sink({ + 'category': cat, 'parent': kw, 'reply': text, + 'stop_reason': reply.get('stop_reason'), + 'truncated': bool(reply.get('truncated')), + 'parsed': parsed, 'n_parsed': len(parsed), + }) logger.info(f'[AgenticChallenger] expanded {len(hard)} hard keyword(s) -> ' f'+{added} same-domain topics') return added diff --git a/src/twinkle_agentic/challenger/base.py b/src/twinkle_agentic/challenger/base.py index 727e99bdd..843d93848 100644 --- a/src/twinkle_agentic/challenger/base.py +++ b/src/twinkle_agentic/challenger/base.py @@ -266,34 +266,43 @@ def explore( self, trajectories: List[Trajectory], sampling_params: Optional[SamplingParams] = None, + **kwargs: Any, ) -> List[Trajectory]: """Run the explorer over a batch, optionally overriding its sampling params. The override is only forwarded when asked for, so a plain callable explorer keeps working; both rollouts in - :mod:`twinkle_agentic.rollout` accept it. + :mod:`twinkle_agentic.rollout` accept it. Anything else in ``kwargs`` is + passed straight through for the same reason -- a caller that needs a + rollout-specific hook (``followup_fn``) says so per call, and an explorer + that does not take it fails loudly instead of silently ignoring it. """ if not trajectories: return [] - if sampling_params is None: + if sampling_params is not None: + kwargs['sampling_params'] = sampling_params + if not kwargs: return self.explorer(trajectories) - return self.explorer(trajectories, sampling_params=sampling_params) + return self.explorer(trajectories, **kwargs) def _solver_explore( self, trajectories: List[Trajectory], sampling_params: Optional[SamplingParams] = None, + **kwargs: Any, ) -> List[Trajectory]: """Run solver attempts through the solver explorer, or fall back to the main one. Subclasses that need per-attempt isolation (e.g. sandbox workspace reset) - override this rather than the whole difficulty filter. + override this rather than the whole difficulty filter. Extra kwargs are + forwarded, which is how such a subclass says which sandbox each attempt + runs in (``tool_manager`` as a list, one entry per trajectory). """ if self.solver_explorer is not None: if sampling_params is None: - return self.solver_explorer(trajectories) - return self.solver_explorer(trajectories, sampling_params=sampling_params) - return self.explore(trajectories, sampling_params=sampling_params) + return self.solver_explorer(trajectories, **kwargs) + return self.solver_explorer(trajectories, sampling_params=sampling_params, **kwargs) + return self.explore(trajectories, sampling_params=sampling_params, **kwargs) def _filter_difficulty(self, tasks: List[Trajectory]) -> List[Trajectory]: """Attempt each task ``solver_rollouts`` times; keep the ones in the band. diff --git a/src/twinkle_agentic/harness/ms_agent.py b/src/twinkle_agentic/harness/ms_agent.py index d3f50eb92..5eb275890 100644 --- a/src/twinkle_agentic/harness/ms_agent.py +++ b/src/twinkle_agentic/harness/ms_agent.py @@ -27,6 +27,7 @@ from __future__ import annotations import json +import os import uuid from typing import Any, Dict, List, Optional, Union @@ -451,6 +452,16 @@ def patch_ms_agent_python_executor() -> bool: reimplementing the method, so ms-agent keeps owning timeouts, output capture and the JSON result shape. + It also chdirs into the tool's own workspace before each call. That ``exec`` + runs in the host process, so a relative path in model code resolves against + whatever directory the process happens to be in, while ``shell_executor`` and + every ``file_system`` tool pass ``cwd=self._ws.root``. Measured in the RSI + sandbox before the fix: ``write_file 'a.txt'`` answered "Save file + successfully" and the next python call got ``[Errno 2] No such file or + directory: 'a.txt'``, because the file was in the workspace and python was + looking in ``/``; it accounted for 41 of one run's 58 such failures, and files + python wrote landed outside the directory an episode's end state is read from. + Idempotent. Returns True when it patched, False when ms-agent is missing or the patch is already in place. """ @@ -464,6 +475,10 @@ def patch_ms_agent_python_executor() -> bool: return False async def python_executor(self, code: str, description: str = '', timeout=None): + root = getattr(self, 'output_dir', None) or getattr(getattr(self, '_ws', None), 'root', None) + if root: + os.makedirs(root, exist_ok=True) + os.chdir(root) return await original(self, single_namespace_source(code), description=description, timeout=timeout) diff --git a/src/twinkle_agentic/rollout/api_multi_turn.py b/src/twinkle_agentic/rollout/api_multi_turn.py index 4513d361b..f140a2dc5 100644 --- a/src/twinkle_agentic/rollout/api_multi_turn.py +++ b/src/twinkle_agentic/rollout/api_multi_turn.py @@ -14,6 +14,10 @@ _STOP_MAX_TURNS = 'max_turns' _STOP_API_ERROR = 'api_error' +# Runaway guard: a ``followup_fn`` is expected to return None eventually. This +# only bounds a callback that never does, so one bad hook cannot spin forever. +_MAX_FOLLOWUPS = 20 + class APIMultiTurnRollout(Rollout): """Multi-turn rollout over an OpenAI-compatible chat-completions API. @@ -30,6 +34,13 @@ class APIMultiTurnRollout(Rollout): 6. ``turn >= max_turns`` => terminate with ``stop_reason='max_turns'`` (and ``truncated=True``). + After the tool loop ends, if a ``followup_fn`` was passed (per call or at + construction), it is invoked exactly as in :class:`MultiTurnRollout`: it may + append one more user message and buy one more generation whose reply is an + answer, not a tool turn (tools are withdrawn for it), repeating until the + callback returns None. This is what lets a challenger append its check-script + and problem-statement stages onto the same conversation. + Constructor and per-call override semantics intentionally mirror :class:`MultiTurnRollout`: ``tool_manager`` may be a single instance (broadcast) or a list aligned 1:1 with trajectories, and it is optional -- @@ -93,6 +104,7 @@ def __call__( extra_body = dict(self.extra_body) if 'extra_body' in kwargs and kwargs['extra_body']: extra_body.update(kwargs['extra_body']) + followup_fn = kwargs.get('followup_fn') # Per-trajectory thread pool. OpenAI ``/chat/completions`` is # one-conversation-per-call; concurrency only buys us network @@ -100,7 +112,7 @@ def __call__( outs: List[Optional[Trajectory]] = [None] * n with ThreadPoolExecutor(max_workers=self.concurrency) as pool: futures = { - pool.submit(self._run_one, trajectories[i], tool_managers[i], sampling_params, extra_body): i + pool.submit(self._run_one, trajectories[i], tool_managers[i], sampling_params, extra_body, followup_fn): i for i in range(n) } for fut in as_completed(futures): @@ -120,6 +132,7 @@ def _run_one( tool_manager: Optional[ToolManager], sampling_params: SamplingParams, extra_body: Dict[str, Any], + followup_fn: Optional[Callable[[Trajectory, int], Any]] = None, ) -> Trajectory: """Drive the API turn loop for a single trajectory. @@ -199,11 +212,51 @@ def _run_one( truncated = True stop_reason = _STOP_MAX_TURNS + # Follow-up stages (check script, problem statement, ...). Each one + # appends a user message and takes one generation whose reply is an + # answer: tools are withdrawn so the model writes rather than calls. + # Skipped entirely on an API error -- the conversation is already broken. + followups = 0 + if followup_fn is not None and stop_reason != _STOP_API_ERROR: + while followups < _MAX_FOLLOWUPS: + view = dict(trajectory) + view['messages'] = messages + view['turns'] = turn + view['stop_reason'] = stop_reason + view['truncated'] = truncated + view['followups'] = followups + followup = followup_fn(view, followups) + if followup is None: + break + text, next_params = (followup if isinstance(followup, tuple) + else (followup, None)) + messages.append({'role': 'user', 'content': text}) + followups += 1 + fu_params = next_params if next_params is not None else sampling_params + try: + reply = (self.api( # tools omitted on purpose: this is an answer + {'messages': messages}, fu_params, extra_body=extra_body) + if extra_body else self.api({'messages': messages}, fu_params)) + except Exception as exc: + stop_reason = _STOP_API_ERROR + error = f'{type(exc).__name__}: {exc}' + truncated = True + break + assistant_msg = self._normalise_assistant(reply, turn + followups) + messages.append(assistant_msg) + # A follow-up that stopped cleanly means the episode was not cut + # off after all, even if the tool phase had hit its turn cap. + if assistant_msg.get('finish_reason') == 'length': + truncated = True + elif stop_reason == _STOP_MAX_TURNS: + stop_reason = _STOP_NO_TOOL + out = dict(trajectory) out['messages'] = messages out['turns'] = turn out['stop_reason'] = stop_reason out['truncated'] = truncated + out['followups'] = followups if error is not None: out['error'] = error return out diff --git a/src/twinkle_agentic/rollout/bridge.py b/src/twinkle_agentic/rollout/bridge.py index 2663f9ed1..fa8bffa1b 100644 --- a/src/twinkle_agentic/rollout/bridge.py +++ b/src/twinkle_agentic/rollout/bridge.py @@ -12,11 +12,19 @@ rewritten to use the ``template`` parameter. No Ray decorators (``@remote_function`` / ``@remote_class``) are applied here. """ + + import numpy as np from typing import Any, Dict, List, Optional from twinkle.template.base import Template +# Stand-in history for the fallback delta computation in +# :func:`extend_with_bridge`. A single user turn, because what precedes the +# appended message must itself render the same way with and without it: a user +# turn has no reasoning block for the template to move or drop. +_ANCHOR = [{'role': 'user', 'content': 'x'}] + def _to_plain(obj: Any) -> Any: """Recursively convert numpy arrays/scalars to plain Python lists/numbers. @@ -75,13 +83,40 @@ def extend_with_bridge( messages_after, tokenize=False, add_generation_prompt=True, enable_thinking=enable_thinking) if not s_after.startswith(s_before): - raise RuntimeError('Canonical chat_template output for messages_after is not a ' - 'prefix-extension of messages_before; cannot compute bridge ' - 'delta. This indicates the template is non-monotonic in the ' - 'message list (e.g. reorders / rewrites earlier turns).\n' - f's_before tail: {s_before[-80:]!r}\n' - f's_after at same offset: ' - f'{s_after[max(0, len(s_before) - 80):len(s_before) + 80]!r}') + # Appending a *user* message moves where Qwen3's template thinks the + # conversation's last question is, and it renders assistant turns either + # side of that point differently: the turn before it loses its <think> + # block, and the turn after it gains an empty one when it had none. + # Measured on Qwen3-4B with three messages -- rendered alone, the + # assistant turn reads '<think>\nthinking hard\n</think>\n\nAll tasks are + # complete.'; rendered with a user turn after it, just 'All tasks are + # complete.'. Tool messages do not move that point, which is why + # appending tool observations has always been a clean extension. + # + # So the delta is measured against a stand-in history instead: render one + # user turn, then the same turn plus these messages, and take the + # difference. That is exact as long as a message block does not depend on + # what precedes it, which the prefix check below still enforces. + # + # What stays on record is the history as generated, thinking included -- + # those are the tokens the policy read back when it produced the next + # turn, and a later training step has to see the same. + s_anchor = tokenizer.apply_chat_template( + _ANCHOR, tokenize=False, add_generation_prompt=False, + enable_thinking=enable_thinking) + s_anchor_after = tokenizer.apply_chat_template( + _ANCHOR + list(tool_messages), tokenize=False, add_generation_prompt=True, + enable_thinking=enable_thinking) + if not s_anchor_after.startswith(s_anchor): + raise RuntimeError('Canonical chat_template output for messages_after is not a ' + 'prefix-extension of messages_before, and the same is true ' + 'of a one-message stand-in history; cannot compute bridge ' + 'delta. This indicates the template is non-monotonic in the ' + 'message list (e.g. reorders / rewrites earlier turns).\n' + f's_before tail: {s_before[-80:]!r}\n' + f's_after at same offset: ' + f'{s_after[max(0, len(s_before) - 80):len(s_before) + 80]!r}') + s_before, s_after = s_anchor, s_anchor_after bridge_text = s_after[len(s_before):] if not bridge_text: raise RuntimeError('Bridge text computation returned empty string; ' diff --git a/src/twinkle_agentic/rollout/multi_turn.py b/src/twinkle_agentic/rollout/multi_turn.py index e66f0b222..396a60e25 100644 --- a/src/twinkle_agentic/rollout/multi_turn.py +++ b/src/twinkle_agentic/rollout/multi_turn.py @@ -1,4 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +import json +import re from collections import defaultdict from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Any, Callable, Dict, List, Optional, Tuple @@ -28,6 +30,47 @@ def _append_only_delta( return new[len(old):] +def is_error_observation(observation: str) -> bool: + """Did a tool come back with a failure rather than a result? + + Only the two shapes tools actually produce are matched, taken from a dump of + 239 real calls: ms-agent wraps a failure as ``{"success": false, ...}``, and + a dispatch that never reached a tool (unknown name, a file the tool refuses + to touch) comes back as a bare line starting with ``Error:``. Plus the two + messages an unreachable sandbox produces. + + Deliberately narrow. Matching on words like ``failed`` or ``not found`` + anywhere in the text also matches a *successful* read of a file that happens + to contain them, and this decides whether an episode is cut short. + """ + text = (observation or '').strip() + if not text: + return False + if text.startswith('Error:'): + return True + if text.startswith(('Tool runtime unreachable:', 'Tool runtime returned no result')): + return True + return bool(re.search(r'"success"\s*:\s*false', text)) + + +def _call_key(tool_call: Dict[str, Any]) -> str: + """A stable identity for a tool call: its name plus its arguments verbatim. + + Byte-identical is the point. A model that changes one path and tries again is + making progress; one that reissues the same call with the same arguments is + not, whatever the tool answered. + """ + fn = tool_call.get('function') if isinstance(tool_call.get('function'), dict) else {} + name = fn.get('name') or tool_call.get('name') or tool_call.get('tool_name') or '' + args = fn.get('arguments', tool_call.get('arguments')) + if not isinstance(args, str): + try: + args = json.dumps(args, sort_keys=True, ensure_ascii=False) + except (TypeError, ValueError): + args = repr(args) + return f'{name}\x00{args}' + + def _default_tool_messages( tool_calls: List[Dict[str, Any]], observations: List[str], @@ -73,6 +116,7 @@ class MultiTurnRollout(Rollout): * ``harness``: a single :class:`AgentHarness` or a 1:1 list. Framework specifics (ms-agent system/memory/tool-message shape) live in the harness subclass, not here. + * ``followup_fn``: see ``__init__``. """ def __init__( @@ -87,6 +131,9 @@ def __init__( trace_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, success_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, harness: Optional[AgentHarness] = None, + adapter_path: Optional[str] = None, + stop_after_stuck_turns: int = 0, + followup_fn: Optional[Callable[[Trajectory, int], Any]] = None, ): super().__init__() if template is None: @@ -104,7 +151,58 @@ def __init__( self.template = template self.tool_manager = tool_manager self.harness = harness + # A LoRA directory on disk, forwarded to every sample call. Training syncs + # its adapter into the sampler directly, but evaluating a saved one has no + # such channel: without this, an eval script would silently measure the + # base model and report it as the trained one. + self.adapter_path = adapter_path self.max_trajectory_tokens = max_trajectory_tokens + # How many stuck turns in a row end the episode; 0 runs to ``max_turns`` + # regardless. A turn is stuck when it made no progress at all, which is + # either of: + # * every call in it came back an error, or + # * every call in it was byte-identical to one already made in this + # episode, whatever it answered. + # One useful call in a turn resets the count, so probing for something + # and then creating it is untouched. + # + # Both halves are needed, measured by replaying 12 recorded episodes: + # errors alone stop 1 of 12 and save 9 of 239 calls, because the worst + # offenders interleave a failing call with a glob that succeeds. Adding + # the repeat rule stops 3 of 12 and saves 63 calls, and the three are + # exactly the ones that spent 54, 84 and 17 calls to leave behind a + # script that could not run. Nothing an episode kept was written after + # its stop point except those broken scripts. + if stop_after_stuck_turns < 0: + raise ValueError(f'stop_after_stuck_turns must be >= 0, got ' + f'{stop_after_stuck_turns}') + self.stop_after_stuck_turns = stop_after_stuck_turns + # Called with (trajectory, how many follow-ups it has had already) at the + # moment an episode would end: because the model stopped calling tools, + # because it used up ``max_turns``, or because it was stopped for being + # stuck. Returning a string appends it as a user message and the episode + # keeps going; returning None ends it. May also return + # ``(text, SamplingParams)`` to give that stage its own budget. + # + # It is asked in the ran-out-of-budget cases too, not only when the model + # says it is done, because what those stages read is the state the episode + # left behind -- which exists either way. An episode dropped for hitting + # the turn limit costs its whole sandbox run and produces nothing. + # + # This is what keeps a multi-stage episode in ONE trajectory. The + # alternative -- ending here and starting a second rollout whose prompt is + # this conversation -- re-encodes the history as prompt, so every earlier + # assistant turn comes back with labels == -100 and only the last stage is + # trainable. Appending goes through the same append-only bridge the tool + # observations use, so labels and logprobs of the earlier turns survive and + # the whole chain can be trained as one sample. + # + # Tool calls are no longer dispatched once a follow-up has been appended: + # the stages that come after the tool-using one are meant to produce text + # about the state as it is, and a python block in a reply parses as a call + # list -- 41 of 146 such replies dispatched something in a measured run -- + # which would rewrite the very state the text is about. + self.followup_fn = followup_fn assert self.template.truncation_strategy != 'split', ( "MultiTurnRollout does not support truncation_strategy='split'; " 'use left/right/delete/raise on the template.') @@ -120,6 +218,10 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] return [] sampling_params = kwargs.get('sampling_params', self.sampling_params) + adapter_path = kwargs.get('adapter_path', self.adapter_path) + # Left out entirely when unset, so a sampler without LoRA enabled sees the + # same call it always did. + adapter_kwargs = {'adapter_path': adapter_path} if adapter_path else {} tool_managers = self._broadcast( kwargs.get('tool_manager', self.tool_manager), n, name='tool_manager', required=True) harnesses = self._broadcast(kwargs.get('harness', self.harness), n, name='harness') @@ -158,9 +260,66 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] turns: List[int] = [0] * n truncated: List[bool] = [False] * n done: List[bool] = [False] * n - + # Consecutive turns that made no progress, the calls already issued in + # each episode, and whether being stuck is what ended it. All three stay + # at their initial value when ``stop_after_stuck_turns`` is 0. + stuck_turns: List[int] = [0] * n + seen_calls: List[set] = [set() for _ in range(n)] + stuck_stop: List[bool] = [False] * n + # Follow-up bookkeeping (all no-ops when ``followup_fn`` is None): + # how many follow-ups each trajectory has had, and the params its next + # turn should use. A trajectory that has had one stops dispatching tools. + followups: List[int] = [0] * n + params_for: List[Any] = [sampling_params] * n + followup_fn = kwargs.get('followup_fn', self.followup_fn) + # Why the tool-calling part of each episode ended, when it was not the + # model's own choice: 'max_turns' or 'stuck'. Reported separately from + # ``truncated`` because an episode can hit the turn limit and still go on + # to answer the follow-up stages, in which case nothing was cut off. + tool_stop: List[Optional[str]] = [None] * n + + def append_followup(global_idx: int) -> bool: + """Ask for one more stage; True when the episode carries on. + + Sets ``truncated`` itself in the one case where the answer is "there + is no room for another stage", which is a cut trajectory rather than + a caller that had nothing more to ask. + """ + nonlocal iterations + if followup_fn is None: + return False + followup = followup_fn( + self._as_trajectory(trajectories[global_idx], pifs[global_idx], + all_logprobs[global_idx], turns[global_idx], + stop_reasons[global_idx], truncated[global_idx]), + followups[global_idx]) + if followup is None: + return False + text, next_params = followup if isinstance(followup, tuple) else (followup, None) + extended = extend_with_bridge( + pifs[global_idx], [{'role': 'user', 'content': text}], self.template) + if extended is None: + truncated[global_idx] = True + return False + pifs[global_idx] = extended + if lives[global_idx] is not None: + lives[global_idx]['messages'] = list(extended.get('messages') or []) + followups[global_idx] += 1 + iterations += 1 + if next_params is not None: + params_for[global_idx] = next_params + return True + + # The loop counts generations, and each granted follow-up buys the one + # extra generation it asked for. Paying for the follow-up stages out of + # ``max_turns`` would mean an episode that spent its whole tool budget + # never reaches the stages that read what it built, and a short one + # silently gets more tool turns than a long one. + iterations = self.max_turns + done_iterations = 0 first_turn = True - for _ in range(self.max_turns): + while done_iterations < iterations: + done_iterations += 1 active = [i for i in range(n) if not done[i]] if not active: break @@ -177,26 +336,47 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] break first_turn = False - # 2. One batched sample call for all currently-live trajectories. - batch_pifs = [pifs[i] for i in active] - actual = len(batch_pifs) + # 2. One batched sample call per distinct SamplingParams among the + # live trajectories -- normally exactly one, since only a + # follow-up stage asks for its own budget. Grouping rather than + # taking the first is what keeps a mixed batch honest: sampling one + # trajectory under another's token limit would silently truncate or + # over-spend, and the two are indistinguishable afterwards. + groups: List[List[int]] = [] + group_params: List[Any] = [] + for global_idx in active: + for slot, params in enumerate(group_params): + if params is params_for[global_idx]: + groups[slot].append(global_idx) + break + else: + group_params.append(params_for[global_idx]) + groups.append([global_idx]) + + resps_by_idx: Dict[int, Any] = {} device_mesh = getattr(self.sampler, 'device_mesh', None) min_batch_size = (device_mesh.data_world_size if device_mesh is not None else 1) - if actual < min_batch_size: - batch_pifs = batch_pifs + ([batch_pifs[-1]] * (min_batch_size - actual)) - resps = self.sampler.sample(batch_pifs, sampling_params=sampling_params) - resps = self._unwrap_response_list(resps, len(batch_pifs))[:actual] + for slot, group in enumerate(groups): + batch_pifs = [pifs[i] for i in group] + actual = len(batch_pifs) + if actual < min_batch_size: + batch_pifs = batch_pifs + ([batch_pifs[-1]] * (min_batch_size - actual)) + group_resps = self.sampler.sample(batch_pifs, + sampling_params=group_params[slot], + **adapter_kwargs) + group_resps = self._unwrap_response_list(group_resps, len(batch_pifs))[:actual] + for local_idx, global_idx in enumerate(group): + resps_by_idx[global_idx] = group_resps[local_idx] pending_tools: List[tuple] = [] # (global_idx, tool_calls) - for local_idx, global_idx in enumerate(active): + for global_idx in active: turns[global_idx] += 1 - seq = resps[local_idx].sequences[0] + seq = resps_by_idx[global_idx].sequences[0] if seq.new_input_feature is None or 'input_ids' not in seq.new_input_feature: raise RuntimeError(f'Sampler returned a SampledSequence without ' - f'new_input_feature.input_ids at batch index ' - f'{local_idx} (trajectory {global_idx}); ' - f'cannot continue multi-turn.') + f'new_input_feature.input_ids for trajectory ' + f'{global_idx}; cannot continue multi-turn.') pifs[global_idx] = _to_plain(dict(seq.new_input_feature)) if seq.logprobs is not None: @@ -213,6 +393,34 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] tool_calls = (_last_msg.get('tool_calls') if isinstance(_last_msg, dict) else None) if not tool_calls: tool_calls = self.template.parse_tool_call(seq.decoded or '') + # After a follow-up, a parsed call is not a call: the tools were + # withdrawn for these stages on purpose (see ``followup_fn``), and + # dispatching python that the model wrote as *an answer* would edit + # the state the answer is about. + if followups[global_idx]: + tool_calls = None + # The parse also *rewrote* the message: when a reply parses as + # a call, the template stores it with the call text removed, so + # a caller reading the message gets less than the model wrote. + # For these stages the reply is the deliverable, and one of the + # tool-call formats is XML-shaped, so a check script asserting + # the content of an .xml file matches it: 5 of ex12's 72 check + # scripts came back with the XML cut out of them -- three then + # ran with `content == ''` where the model had written the file's + # real text, and two no longer held a code block at all. + if _msgs and isinstance(_last_msg, dict): + # Decoded without the special tokens, the way the + # template writes a message: ``seq.decoded`` keeps the + # closing ``<|im_end|>``, and putting that in the content + # put it in the problem statements ex13 handed to solvers + # -- 7 of 7 of them ended in a literal '<|im_end|>'. + tok = getattr(self.template, 'tokenizer', None) + if tok is not None and seq.tokens: + _last_msg['content'] = tok.decode( + seq.tokens, skip_special_tokens=True) + else: + _last_msg['content'] = seq.decoded or '' + _last_msg.pop('tool_calls', None) if lives[global_idx] is not None: lives[global_idx]['messages'] = list(_msgs) @@ -222,7 +430,15 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] self._merge_assistant_metadata(pifs[global_idx], lives[global_idx]) # 3. Termination conditions + # A reply cut off at ``max_tokens`` is truncated in exactly the + # sense the flag names, and consumers read the flag to tell a + # trajectory that finished from one that ran out of room: a + # difficulty measurement counting such an attempt as a genuine + # failure blames the task for the token budget. Tool calls the + # cut reply happens to contain are still not dispatched -- the + # turn never got to decide it was done emitting them. if seq.stop_reason == 'length': + truncated[global_idx] = True done[global_idx] = True continue @@ -234,10 +450,20 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] continue if not tool_calls: + # The episode is over as far as the model is concerned. Give + # the caller one chance to say otherwise -- see + # ``followup_fn`` for why this is not a second rollout. + if append_followup(global_idx): + continue done[global_idx] = True continue if turns[global_idx] >= self.max_turns: + # Out of tool turns, not out of episode: the stages that read + # the end state can still run on what was built. + tool_stop[global_idx] = 'max_turns' + if append_followup(global_idx): + continue truncated[global_idx] = True done[global_idx] = True continue @@ -252,6 +478,17 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] obs_by_traj = self._dispatch_tools(tool_managers, pending_tools) for global_idx, tool_calls in pending_tools: observations = obs_by_traj.get(global_idx) or [''] * len(tool_calls) + if self.stop_after_stuck_turns: + keys = [_call_key(tc) for tc in tool_calls] + all_repeats = bool(keys) and all(k in seen_calls[global_idx] + for k in keys) + seen_calls[global_idx].update(keys) + all_errors = bool(observations) and all( + is_error_observation(o) for o in observations) + if all_errors or all_repeats: + stuck_turns[global_idx] += 1 + else: + stuck_turns[global_idx] = 0 tool_messages, lives[global_idx] = self._tool_messages_after( pifs[global_idx], lives[global_idx], harnesses[global_idx], observations, tool_calls) @@ -264,6 +501,18 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] pifs[global_idx] = extended if lives[global_idx] is not None: lives[global_idx]['messages'] = list(extended.get('messages') or []) + # Checked after the messages are appended, so the turns that + # ended the episode are in the trajectory the caller reads. + if (self.stop_after_stuck_turns + and stuck_turns[global_idx] >= self.stop_after_stuck_turns): + stuck_stop[global_idx] = True + tool_stop[global_idx] = 'stuck' + # Same as the turn limit: the tool phase is over, the + # state it left is not, so the stages still get their turn. + if not done[global_idx] and append_followup(global_idx): + continue + truncated[global_idx] = True + done[global_idx] = True for i in range(n): if not all_logprobs[i]: @@ -289,6 +538,12 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] out['turns'] = turns[i] out['stop_reason'] = stop_reasons[i] out['truncated'] = truncated[i] + # ``truncated`` says something was cut off; these two say what ended + # the tool-calling part, which is a different question -- an episode + # can run out of turns, be handed a follow-up stage, and finish it. + out['stuck_stop'] = stuck_stop[i] + out['tool_stop'] = tool_stop[i] + out['followups'] = followups[i] outs.append(out) # Per-rollout trace dump: one JSON file per selected trajectory. @@ -301,6 +556,25 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] # ------------------------------------------------------------------ private + @staticmethod + def _as_trajectory(traj: Trajectory, pif: Dict[str, Any], logprobs: List[Any], + turns: int, stop_reason: Optional[str], + truncated: bool) -> Trajectory: + """The episode so far, shaped like the value ``__call__`` returns. + + Handed to ``followup_fn`` so the callback reads an episode the same way + every other consumer does -- ``messages`` complete, token fields present -- + rather than having to know this loop's local variables. + """ + out = dict(traj) + out.update(pif) + out['messages'] = list(pif.get('messages') or traj.get('messages') or []) + out['logprobs'] = logprobs if logprobs else None + out['turns'] = turns + out['stop_reason'] = stop_reason + out['truncated'] = truncated + return out + def _harness_before_generate( self, pif: Dict[str, Any], diff --git a/src/twinkle_agentic/tools/tool_manager.py b/src/twinkle_agentic/tools/tool_manager.py index ea9bf63db..3c0697b50 100644 --- a/src/twinkle_agentic/tools/tool_manager.py +++ b/src/twinkle_agentic/tools/tool_manager.py @@ -51,6 +51,28 @@ def _unpack_tool_call(tool_call: Any) -> Tuple[Optional[str], Dict[str, Any], Op f'got {type(raw_args).__name__}.') +def _suggest(name: str, available: Iterable[str]) -> Optional[str]: + """The registered tool ``name`` was probably meant to be, if there is one. + + Only one mistake is guessed at: a name given without its namespace, or under + the wrong one. Agent frameworks hand out qualified names -- ms-agent's are + ``{server}---{tool}`` -- and a model that has seen the bare verb in a + docstring writes ``shell_executor``, or files it under the server it was last + using. Measured over 5793 calls: 201 bare ``shell_executor`` and 30 + ``file_system---shell_executor``, all for one tool that does exist. + + Deliberately only a suggestion: the call is still refused. Resolving it + silently would train the policy to emit a name that no serving deployment + accepts, and the unqualified form is ambiguous the moment two servers export + the same verb -- which is why a suffix shared by several tools yields nothing. + """ + wanted = name.rsplit('---', 1)[-1] + if not wanted: + return None + matches = [n for n in available if n != name and n.rsplit('---', 1)[-1] == wanted] + return matches[0] if len(matches) == 1 else None + + class ToolManager: def __init__( @@ -102,7 +124,11 @@ def __call__(self, tool_call: Union[ToolCall, Dict[str, Any]]) -> str: return err if (tool := self._tools.get(name)) is None: available = ', '.join(sorted(self._tools)) or '(none)' - return f'Error: unknown tool {name!r}. Available: {available}.' + hint = '' + if (suggestion := _suggest(name, self._tools)) is not None: + hint = (f' Did you mean {suggestion!r}? Tool names must be given in ' + f'full, including the part before "---".') + return f'Error: unknown tool {name!r}.{hint} Available: {available}.' try: return str(tool(name, args)) except Exception as e: # noqa @@ -117,9 +143,19 @@ def call_many( ``tool_calls`` are the OpenAI-shaped dicts produced by :meth:`~twinkle.template.base.Template.parse_tool_call`. This method - unpacks them to ``(name, arguments)`` and, when every tool wraps the - same :class:`~twinkle_agentic.envs.base.Env`, dispatches through + unpacks them to ``(name, arguments)`` and, when the tools wrap the same + :class:`~twinkle_agentic.envs.base.Env`, dispatches through ``Env.step_batch``. Otherwise a thread pool of :meth:`__call__`. + + A call this manager can answer by itself -- an unknown name, a malformed + payload -- is answered here and *excluded* from the batch rather than + disqualifying it. It used to disqualify it: one bare ``shell_executor`` + in a turn of five sent the whole turn down the thread pool, and + concurrent dispatch is where the environment is least likely to be safe. + It was not: in ex4's episode 8 four calls fired at once and all four came + back with the same glob listing, so the model was told its python had run + when it never did. Nothing in that turn needed concurrency -- the reason + it was used was a tool name the host could have refused on the spot. """ calls = list(tool_calls) if not calls: @@ -129,19 +165,25 @@ def call_many( unpacked = [_unpack_tool_call(tc) for tc in calls] env = self._shared_env() - can_batch = env is not None and all( - err is None and name in self._tools for name, _args, err in unpacked) - if can_batch: + if env is not None: + out: List[Optional[str]] = [None] * len(calls) + batched: List[Tuple[int, str, Dict[str, Any]]] = [] + for i, (name, args, err) in enumerate(unpacked): + if err is None and name in self._tools: + batched.append((i, name, args)) + else: + out[i] = self(calls[i]) try: - results = env.step_batch([(name, args) for name, args, _err in unpacked]) - return [ - r.observation if hasattr(r, 'observation') else str(r) for r in results - ] + results = env.step_batch([(name, args) for _i, name, args in batched]) except Exception: - pass + results = None + if results is not None and len(results) == len(batched): + for (i, _name, _args), r in zip(batched, results): + out[i] = r.observation if hasattr(r, 'observation') else str(r) + return ['' if x is None else x for x in out] workers = max_workers or min(32, len(calls)) - out: List[Optional[str]] = [None] * len(calls) + out = [None] * len(calls) with ThreadPoolExecutor(max_workers=workers) as pool: futs = {pool.submit(self, tc): i for i, tc in enumerate(calls)} for fut in as_completed(futs): diff --git a/src/twinkle_client/rollout/multi_turn.py b/src/twinkle_client/rollout/multi_turn.py index 55c5800b7..dbd2c2db3 100644 --- a/src/twinkle_client/rollout/multi_turn.py +++ b/src/twinkle_client/rollout/multi_turn.py @@ -92,6 +92,9 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] * tool_calls produced with no ``tool_manager`` -> ValueError. * ``max_turns == 1`` with a first-round tool call -> the trajectory is marked ``truncated=True, stop_reason='max_turns'`` and sampling stops. + * ``stop_reason == 'length'`` -> the trajectory is marked + ``truncated=True`` and sampling stops without dispatching any tool + call the cut reply contains. """ if isinstance(trajectories, dict): raise TypeError('ClientMultiTurnRollout.__call__ expects a List[Trajectory]; ' @@ -170,7 +173,12 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] stop_reasons[global_idx] = seq.stop_reason # 3. Termination conditions. + # Cut off at ``max_tokens``: truncated, same as the max_turns and + # length-cap cases below, and same as ``MultiTurnRollout`` and + # ``ApiMultiTurnRollout``. Tool calls in the cut reply are still + # not dispatched. if seq.stop_reason == 'length': + truncated[global_idx] = True done[global_idx] = True continue diff --git a/tests/template/test_tool_call_parsers.py b/tests/template/test_tool_call_parsers.py new file mode 100644 index 000000000..6caf57574 --- /dev/null +++ b/tests/template/test_tool_call_parsers.py @@ -0,0 +1,86 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Tool-call parser selection and the bracketed call-list format.""" +import pytest + +from twinkle.template.tools import ToolCallRegistry +from twinkle.template.tools.bracket_dsl import BracketDslParser + +FENCE = '```' + + +def names(text): + parser = BracketDslParser() + if not parser.detect(text): + return [] + return [c['function']['name'] for c in parser.parse(text)] + + +@pytest.mark.parametrize( + 'label, text, expected', + [ + ('call list alone', + '[Text Analysis(text="great service"), UserID(username="alex")]', + ['Text Analysis', 'UserID']), + ('call list after prose', + 'Here you go: [quarterly_data(stock_symbols=["AAPL", "TSLA"])]', + ['quarterly_data']), + ('dotted name', '[database.insert_data(table="t")]', ['database.insert_data']), + ], +) +def test_bracket_dsl_parses_call_lists(label, text, expected): + assert names(text) == expected + + +@pytest.mark.parametrize( + 'label, text', + [ + # A comprehension is shaped exactly like a call list. Reading one as + # tool calls invents names like 'float' and 'for _ in range', and the + # tools the model meant to call never run. + ('comprehension in a fence', + f'Sure:\n{FENCE}python\nvals = [float(random.uniform(1, 10)) for _ in range(20)]\n{FENCE}\n'), + ('nested comprehension in a fence', + f'{FENCE}\nrows = [dict(zip(h, r)) for r in raw]\n{FENCE}\n'), + # A reply truncated mid-fence still has to be treated as code. + ('unterminated fence', f'writing code:\n{FENCE}python\ny = [str(i) for i in xs]'), + ('plain prose list', 'the values are [1, 2, 3]'), + # A model writing code while it thinks does not use fences. This is how + # 10% of the episodes in an agentic run lost their tool calls: the reply + # was cut off inside <think>, the comprehension in it parsed as calls to + # `int` and `for _ in range`, and the sandbox was never touched. + ('comprehension in unfenced prose', + 'I will write vals = [float(random.uniform(1, 10)) for _ in range(20)] next'), + ('comprehension inside a think block', + '<think>\nnums = [int(v) for v in raw]\n</think>\nDone.'), + ('reply truncated inside think, comprehension left open', + '<think>\nSo the code would be:\n\nvals = [int(x) for x in lines]\nWait, maybe'), + ('a call list rehearsed while thinking is not a call', + '<think>\nI could answer [get_price(sym="AAPL")] here.\n</think>\nLet me check first.'), + ('positional argument is not a call list', '[get_price("AAPL")]'), + ], +) +def test_bracket_dsl_ignores_code_and_prose(label, text): + assert names(text) == [] + + +def test_bracket_dsl_sees_the_call_after_a_closed_think_block(): + text = '<think>\nvals = [int(v) for v in raw]\n</think>\n[get_price(sym="AAPL")]' + assert names(text) == ['get_price'] + + +def test_bracket_dsl_accepts_a_call_with_no_arguments(): + assert names('[get_time()]') == ['get_time'] + + +def test_bracket_dsl_still_sees_calls_outside_a_fence(): + text = f'{FENCE}python\nx = [int(v) for v in raw]\n{FENCE}\n[get_price(sym="AAPL")]' + assert names(text) == ['get_price'] + + +def test_marked_up_formats_win_over_the_bracket_heuristic(): + """Hermes markup must go to Hermes even when its arguments contain ``[f(``.""" + text = ('<tool_call>\n{"name": "shell_executor", ' + '"arguments": {"command": "python -c \'print([int(x) for x in y])\'"}}\n</tool_call>') + parser = ToolCallRegistry.detect_first(text) + assert parser is not None and parser.name != 'bracket_dsl' + assert [c['function']['name'] for c in parser.parse(text)] == ['shell_executor'] diff --git a/tests/twinkle_agentic/test_agentic_rsi.py b/tests/twinkle_agentic/test_agentic_rsi.py index 42409a7dd..ece9d9a06 100644 --- a/tests/twinkle_agentic/test_agentic_rsi.py +++ b/tests/twinkle_agentic/test_agentic_rsi.py @@ -10,8 +10,11 @@ """ import json import os +import re +import shutil import sys import tempfile +import threading import unittest _REPO = os.path.join(os.path.dirname(__file__), '..', '..') @@ -23,7 +26,8 @@ sys.path.insert(0, os.path.join(_COOKBOOK, 'sandbox_server')) from remote_tool_env import RemoteMsAgentToolEnv # noqa: E402 -from tool_server import _usable_llm, _without_llm_args # noqa: E402 +from tool_server import (ToolRuntime, _usable_llm, # noqa: E402 + _without_internal_args, _without_llm_args) from twinkle_agentic.envs.env_tool import EnvTool # noqa: E402 from twinkle_agentic.tools.tool_manager import ToolManager # noqa: E402 from twinkle_agentic.verifier.result_check import (Check, CheckContext, # noqa: E402 @@ -99,7 +103,14 @@ def handle(self, command, background=False): if '/tools' in command: return _Result(json.dumps({'tools': self.tools})) if '/call' in command: - payload = json.loads(self.files.store['/opt/rsi/request.json']) + # The request file name carries a uuid, so that concurrent calls cannot + # overwrite each other's payload. Read the one this command names + # rather than a fixed path: reading a fixed path is what would keep + # passing after the Env went back to a shared file. + match = re.search(r'--data-binary @(\S+)', command) + if not match: + raise AssertionError(f'call command names no request file: {command}') + payload = json.loads(self.files.store[match.group(1)]) self.requests.append(payload) results = [{'observation': self._responder(call)} for call in payload['calls']] return _Result(json.dumps({'results': results})) @@ -242,15 +253,26 @@ def explode(command, background=False): self.assertIn('unreachable', obs) def test_tool_schemas_come_from_the_sandbox(self): - self.assertEqual(self.env.tool_names(), - [t['function']['name'] for t in DEFAULT_TOOLS]) - - def test_resolve_tool_maps_plain_name_onto_namespaced_one(self): - self.assertEqual(self.env.resolve_tool('shell_executor'), - 'code_executor---shell_executor') - # An already-qualified name is left alone. - self.assertEqual(self.env.resolve_tool('file_system---write_file'), - 'file_system---write_file') + # Advertised without ms-agent's `{server}---` prefix: a 4B policy wrote a + # bare `shell_executor` 7 times across three arms and lost the turn to + # "unknown tool". The prefix carries nothing it can act on. + self.assertEqual( + self.env.tool_names(), + [t['function']['name'].rsplit('---', 1)[-1] for t in DEFAULT_TOOLS]) + + def test_short_name_is_expanded_before_dispatch(self): + # The runtime only answers to its own spelling, so the prefix has to come + # back on the way out. Shortening the advertised name without this would + # make every call fail. + self.env.step('shell_executor', {'command': 'ls'}) + sent = [c['tool_name'] for c in self.sandbox.requests[-1]['calls']] + self.assertEqual(sent, ['code_executor---shell_executor']) + + def test_resolve_tool_accepts_either_spelling(self): + self.assertEqual(self.env.resolve_tool('shell_executor'), 'shell_executor') + # A caller written before the names were shortened still resolves. + self.assertEqual(self.env.resolve_tool('code_executor---shell_executor'), + 'shell_executor') def test_resolve_tool_raises_on_unknown_name(self): # Silently passing a bad name through would surface as a failed check, @@ -330,6 +352,107 @@ def test_missing_llm_section_is_not_a_usable_llm(self): self.assertTrue(_usable_llm( OmegaConf.create({'llm': {'service': 'modelscope', 'modelscope_api_key': 'k'}}))) + def test_host_owned_call_id_is_never_advertised(self): + # ms-agent declares __call_id on shell_executor as "injected by host when + # supported". In the prompt it reads as an argument the model may choose. + schema = { + 'type': 'function', + 'function': { + 'name': 'code_executor---shell_executor', + 'parameters': { + 'properties': {'command': {}, '__call_id': {}}, + 'required': ['command', '__call_id'], + }, + }, + } + stripped = _without_internal_args(schema) + self.assertEqual(sorted(stripped['function']['parameters']['properties']), ['command']) + self.assertEqual(stripped['function']['parameters']['required'], ['command']) + self.assertIn('__call_id', schema['function']['parameters']['properties']) + + +class ToolCallReconcileTest(unittest.TestCase): + """Two failures arrive as one TypeError; only one of them is the model's. + + ms-agent asking for an argument its own tool cannot take is a bug and is + repaired silently. The model reaching for another tool's arguments is not, + and is refused -- with the name of the tool it should have called, because a + call that is quietly rewritten teaches a shape that fails outside this + sandbox. + """ + + # The real line-up, reduced to the two arguments each case turns on. + CONTRACTS = { + 'file_system---write_file': ({'path', 'content'}, {'path', 'content'}), + 'file_system---edit_file': ({'path', 'old_string', 'new_string', 'replace_all'}, + {'path', 'old_string', 'new_string', 'replace_all'}), + 'file_system---glob': ({'pattern', 'path'}, {'pattern', 'path'}), + 'file_system---read_file': ({'path'}, {'path', 'abbreviate'}), + 'code_executor---shell_executor': ({'command', 'run_in_background', 'timeout'}, + {'command', 'run_in_background', 'timeout', 'call_id'}), + 'code_executor---python_executor': ({'code', 'description', 'timeout'}, + {'code', 'description', 'timeout'}), + } + + def setUp(self): + # No ms-agent, no kernel: _reconcile only reads the contract table, and + # building a real runtime here would need a microVM's worth of setup. + self.runtime = ToolRuntime.__new__(ToolRuntime) + self.runtime._contracts = dict(self.CONTRACTS) + + def test_framework_timeout_is_dropped_for_tools_without_one(self): + # ms-agent's own timeout message tells the model to pass `timeout` in the + # tool arguments; write_file has no such parameter and raises TypeError. + args, error = self.runtime._reconcile('file_system---write_file', + {'path': 'a.txt', 'content': 'x', 'timeout': 30}) + self.assertIsNone(error) + self.assertEqual(args, {'path': 'a.txt', 'content': 'x'}) + + def test_description_is_dropped_for_the_sibling_that_lacks_it(self): + args, error = self.runtime._reconcile('code_executor---shell_executor', + {'command': 'ls', 'description': 'list'}) + self.assertIsNone(error) + self.assertEqual(args, {'command': 'ls'}) + + def test_declared_arguments_are_left_alone(self): + call = {'code': 'print(1)', 'description': 'demo', 'timeout': 20} + args, error = self.runtime._reconcile('code_executor---python_executor', dict(call)) + self.assertIsNone(error) + self.assertEqual(args, call) + + def test_empty_glob_path_becomes_the_workspace_root(self): + # '' is glob's own default, but ms-agent's safety guard rejects it as an + # empty file path before the tool is reached. + args, error = self.runtime._reconcile('file_system---glob', {'pattern': '*', 'path': ''}) + self.assertIsNone(error) + self.assertEqual(args, {'pattern': '*', 'path': '.'}) + + def test_edit_file_arguments_on_write_file_are_refused_by_name(self): + args, error = self.runtime._reconcile( + 'file_system---write_file', + {'path': 'a.py', 'old_string': '', 'new_string': 'print(1)'}) + self.assertIsNotNone(error) + # The message has to carry three things: what was rejected, what this + # tool takes, and who owns the arguments that were passed. + self.assertIn("'new_string', 'old_string'", error) + self.assertIn('It accepts: content, path.', error) + self.assertIn('file_system---edit_file', error) + # Not repaired into a content= write: that is the mistake being reported. + self.assertNotIn('content', args) + + def test_withdrawn_argument_is_refused_rather_than_attempted(self): + # abbreviate exists on the method but was withdrawn from the schema + # because this sandbox has no LLM to serve it. + _args, error = self.runtime._reconcile('file_system---read_file', + {'path': 'a.txt', 'abbreviate': True}) + self.assertIn("has no argument 'abbreviate'", error) + + def test_unknown_tool_is_left_for_ms_agent_to_report(self): + call = {'anything': 1} + args, error = self.runtime._reconcile('file_system---nope', dict(call)) + self.assertIsNone(error) + self.assertEqual(args, call) + class ToolBridgeTest(unittest.TestCase): """The prompt's tool list and the executing tool list must be one list.""" @@ -356,12 +479,873 @@ def test_declared_tools_collapse_into_one_step_batch(self): ] out = manager.call_many(calls) self.assertEqual(len(self.sandbox.requests), 1) - self.assertEqual(out, ['ran read_file', 'ran shell_executor']) + # read_file has no prefix to restore; shell_executor does, and the + # FakeSandbox echoes whatever name reached it. + self.assertEqual(out, ['ran read_file', 'ran code_executor---shell_executor']) def test_nameless_schema_is_refused(self): with self.assertRaises(ValueError): EnvTool.from_schemas(self.env, [{'type': 'function', 'function': {}}]) +class EmptyWorkspaceTest(unittest.TestCase): + """An episode that left nothing behind must not become a task. + + When the explorer writes no files, the only assertion true of the end state is + that the directory is empty -- and every solver satisfies that by doing + nothing, so the task scores 4 of 4 and teaches nothing. Five of the ten + verified tasks in one generation run were exactly this. + """ + + def _challenger(self, snapshot): + from twinkle_agentic.challenger.agentic import AgenticChallenger, AgenticPrompts + + self.checks_run = [] + + def explorer(trajectories, **kwargs): + return [{'messages': list(t['messages']) + + [{'role': 'assistant', 'content': '```python\nassert True\n```'}]} + for t in trajectories] + + def run_check_fn(script, slot=0): + self.checks_run.append(script) + return 0, '' + + prompts = AgenticPrompts( + system='s', from_scratch='u', + check_followup='write checks for {final_state}', + check_retry_followup='{error} / {final_state}', + problem_followup='write the statement') + return AgenticChallenger( + prompts, explorer, + reset_fn=lambda slot=0: None, + run_check_fn=run_check_fn, + workspace_snapshot_fn=lambda slot=0: snapshot, + solver_rollouts=0, + ) + + def _explored(self): + return {'messages': [{'role': 'user', 'content': 'do something'}, + {'role': 'assistant', 'content': 'I made three files.'}]} + + def test_empty_snapshot_ends_the_episode_before_any_check_is_written(self): + ch = self._challenger('') + state = {} + # None means "nothing more to say": the rollout ends the episode here. + self.assertIsNone(ch._followup(state, self._explored(), 0)) + self.assertEqual(ch.stats['empty_workspace'], 1) + self.assertEqual(state['reject'][0], 'empty_workspace') + # No check script was even run: there was nothing to check. + self.assertEqual(self.checks_run, []) + # And the episode is not turned into a task afterwards. + self.assertIsNone(ch._finish_episode(state, self._explored())) + + def test_whitespace_only_snapshot_counts_as_empty(self): + ch = self._challenger(' \n ') + state = {} + self.assertIsNone(ch._followup(state, self._explored(), 0)) + self.assertEqual(ch.stats['empty_workspace'], 1) + + def test_a_real_snapshot_asks_for_checks_and_then_runs_them(self): + ch = self._challenger('data.csv 15\n\n--- data.csv ---\nA,B\n1,2') + state = {} + text, params = ch._followup(state, self._explored(), 0) + self.assertEqual(ch.stats['empty_workspace'], 0) + # The listing reaches the model verbatim -- it is the ground truth the + # checks are written against. + self.assertIn('--- data.csv ---', text) + self.assertIsNone(params) + self.assertEqual(self.checks_run, []) + + wrote_script = {'messages': [ + {'role': 'assistant', 'content': '```python\nassert True\n```'}]} + self.assertEqual(ch._followup(state, wrote_script, 1), + ('write the statement', None)) + self.assertEqual(self.checks_run, ['assert True']) + + +class ProblemStatementParseTest(unittest.TestCase): + """What a statement is allowed to carry. + + A statement that says what a file must contain has to be able to show the + content, and the model shows it in a fence. Stripping every fence -- which is + what "the statement is prose, not code" had been implemented as -- turned + "1. `data.json` containing:" into a sentence that ends there. 7 of ex11's 16 + measured statements had a fence and 5 of those 7 were solved 0 times out of + 8, against 1 of the 9 that had none: those tasks were unanswerable, not hard. + """ + + def setUp(self): + from twinkle_agentic.challenger.agentic import parse_problem_statement + self.parse = parse_problem_statement + + def test_fenced_file_content_stays_in_the_statement(self): + reply = ('<think>planning</think>\n' + 'Create `data.json` containing:\n\n' + '```json\n{"a": 1}\n```\n\n' + 'No other files may exist.') + statement = self.parse(reply) + self.assertIn('{"a": 1}', statement) + self.assertIn('No other files may exist.', statement) + + def test_a_fence_around_the_whole_reply_is_unwrapped_not_deleted(self): + reply = '<think>planning</think>\n```\nCreate data.json holding {}.\n```' + self.assertEqual(self.parse(reply), 'Create data.json holding {}.') + + def test_thinking_is_never_part_of_the_statement(self): + reply = '<think>Create secret.txt</think>\nCreate visible.txt.' + self.assertEqual(self.parse(reply), 'Create visible.txt.') + + def test_an_empty_reply_is_no_statement(self): + self.assertIsNone(self.parse('<think>only thought</think>\n \n')) + + +class EpisodeStagesTest(unittest.TestCase): + """One conversation carries the work, the checks and the statement. + + The three used to be three separate calls, which meant only the last one's + tokens were trainable. The fake explorer here plays the part + ``MultiTurnRollout`` plays for real: it appends whatever ``followup_fn`` + returns and keeps generating until it returns None. + """ + + def _challenger(self, replies, snapshot='a.txt 1\n\n--- a.txt ---\nx', + check_exit=0, check_exits=None, **kwargs): + from twinkle_agentic.challenger.agentic import AgenticChallenger, AgenticPrompts + + self.emitted = [] + self.rejected = [] + self.appended = [] + # One exit code per check run, so a test can make the first fail and the + # rewrite pass. + exits = list(check_exits) if check_exits is not None else None + + def run_check(script, slot=0): + code = exits.pop(0) if exits else check_exit + return (code, 'AssertionError' if code else '') + + def explorer(trajectories, **kw): + followup_fn = kw.get('followup_fn') + traj = {'messages': list(trajectories[0]['messages']), 'input_ids': [1, 2, 3]} + for reply in replies: + traj['messages'].append({'role': 'assistant', 'content': reply}) + if followup_fn is None: + break + out = followup_fn(traj, len(self.appended)) + if out is None: + break + text, _params = out + self.appended.append(text) + traj['messages'].append({'role': 'user', 'content': text}) + return [traj] + + prompts = AgenticPrompts( + system='s', from_scratch='u', + check_followup='checks please: {final_state}', + check_retry_followup='it failed: {error} / state: {final_state}', + problem_followup='statement please') + return AgenticChallenger( + prompts, explorer, + reset_fn=lambda slot=0: None, + run_check_fn=run_check, + workspace_snapshot_fn=lambda slot=0: snapshot, + reject_sink=self.rejected.append, + propose_sink=self.emitted.append, + solver_rollouts=0, + **kwargs) + + def test_one_episode_yields_the_script_and_the_statement(self): + from twinkle.data_format import user_data_get + ch = self._challenger(['Done.', '```python\nassert True\n```', + 'Create a.txt holding x.']) + kept = ch._round(1) + + self.assertEqual(len(kept), 1) + self.assertEqual(user_data_get(kept[0].get('user_data'), 'check_script'), + 'assert True') + self.assertEqual(kept[0]['messages'][-1]['content'], 'Create a.txt holding x.') + # Both stages were asked for, in order, in the same conversation. + self.assertEqual(len(self.appended), 2) + self.assertIn('--- a.txt ---', self.appended[0]) + self.assertEqual(self.appended[1], 'statement please') + # One record, not three: one conversation has one set of token ids. + self.assertEqual([r['stage'] for r in self.emitted[0]['rounds']], ['episode']) + self.assertEqual(self.emitted[0]['outcome'], 'kept') + + def test_a_check_that_fails_on_its_own_workspace_stops_before_the_statement(self): + ch = self._challenger(['Done.', '```python\nassert False\n```', + 'never asked for'], check_exit=1, check_retries=0) + kept = ch._round(1) + + self.assertEqual(kept, []) + self.assertEqual(ch.stats['check_run_fail'], 1) + self.assertEqual(len(self.appended), 1) + self.assertEqual(self.rejected[0]['reason'], 'check_run_fail') + # The record has to say what the workspace held when the check ran. + self.assertIn('--- state before check ---', self.rejected[0]['detail']) + # Rejected attempts are dumped too: they are the zero-reward half of a + # GRPO group. + self.assertEqual(self.emitted[0]['outcome'], 'check_run_fail') + + def test_a_failing_check_gets_one_rewrite_and_the_episode_carries_on(self): + """29 of ex12's 36 check failures were one assertion, on a state that was + fine; the rewrite reads the traceback and the listing.""" + from twinkle.data_format import user_data_get + ch = self._challenger(['Done.', + '```python\nassert len(rows) == 5\n```', + '```python\nassert len(rows) == 3\n```', + 'Create a.txt holding x.'], + check_exits=[1, 0]) + kept = ch._round(1) + + self.assertEqual(len(kept), 1) + # The task ships the script that passed, not the first one. + self.assertEqual(user_data_get(kept[0].get('user_data'), 'check_script'), + 'assert len(rows) == 3') + self.assertEqual(ch.stats['check_retry'], 1) + self.assertEqual(ch.stats['check_retry_pass'], 1) + self.assertEqual(ch.stats['check_run_fail'], 0) + # checks -> rewrite -> statement, and the rewrite was told what broke. + self.assertEqual(len(self.appended), 3) + self.assertIn('AssertionError', self.appended[1]) + self.assertEqual(self.appended[2], 'statement please') + + def test_a_rewrite_that_fails_too_is_rejected_with_both_attempts(self): + ch = self._challenger(['Done.', + '```python\nassert False\n```', + '```python\nassert False\n```', + 'never asked for'], + check_exits=[1, 1]) + kept = ch._round(1) + + self.assertEqual(kept, []) + self.assertEqual(ch.stats['check_run_fail'], 1) + self.assertEqual(ch.stats['check_retry_pass'], 0) + detail = self.rejected[0]['detail'] + self.assertIn('--- attempt 1:', detail) + self.assertIn('--- attempt 2:', detail) + + def test_a_rewrite_that_never_arrives_is_not_shipped_as_a_task(self): + """The failed script is still in the scratchpad when the episode dies.""" + ch = self._challenger(['Done.', '```python\nassert False\n```'], + check_exits=[1]) + kept = ch._round(1) + + self.assertEqual(kept, []) + self.assertEqual(ch.stats['episode_cut_short'], 1) + self.assertEqual(self.rejected[0]['reason'], 'episode_cut_short') + + def test_an_episode_that_never_reached_the_stages_is_recorded_as_cut_short(self): + # The explorer returns after its single reply without consulting the + # callback, which is what a rollout does when the episode ran out of turns. + from twinkle_agentic.challenger.agentic import AgenticChallenger, AgenticPrompts + + self.emitted, self.rejected = [], [] + + def explorer(trajectories, **kw): + return [{'messages': list(trajectories[0]['messages']) + + [{'role': 'assistant', 'content': 'half a thought'}], + 'truncated': True, 'stop_reason': 'length'}] + + ch = AgenticChallenger( + AgenticPrompts(system='s', from_scratch='u', + check_followup='c {final_state}', + check_retry_followup='{error} / {final_state}', + problem_followup='p'), + explorer, + reset_fn=lambda slot=0: None, + run_check_fn=lambda script, slot=0: (0, ''), + workspace_snapshot_fn=lambda slot=0: 'a.txt 1\n', + reject_sink=self.rejected.append, + propose_sink=self.emitted.append, + solver_rollouts=0) + kept = ch._round(1) + + self.assertEqual(kept, []) + self.assertEqual(ch.stats['episode_cut_short'], 1) + self.assertEqual(self.rejected[0]['reason'], 'episode_cut_short') + self.assertIn('truncated=True', self.rejected[0]['detail']) + + +class ConcurrentEpisodeSlotsTest(unittest.TestCase): + """Concurrent episodes must each drive their own sandbox slot. + + A rack of one sandbox per slot is the whole point of running episodes in + parallel; if the slot the challenger passes for episode i is not the slot + reset_fn / run_check_fn / workspace_snapshot_fn / tool_manager see for that + episode, then two episodes end up sharing a workspace and the check written + against one runs against the other. That is the failure mode this test + exists to catch. + """ + + def test_each_episode_uses_its_own_slot_end_to_end(self): + from twinkle_agentic.challenger.agentic import AgenticChallenger, AgenticPrompts + + n_slots = 4 + resets = [] # (slot,) per call + checks = [] # (slot, script) per call + snaps = [] # (slot,) per call + tm_calls = [] # (slot,) per tool_manager use + lock = threading.Lock() + + class FakeTM: + def __init__(self, slot): self.slot = slot + def tool_infos(self): return [] + def __call__(self, tc): + with lock: + tm_calls.append(self.slot) + return 'ok' + + tool_managers = [FakeTM(i) for i in range(n_slots)] + + def reset_fn(slot): + with lock: + resets.append(slot) + + def run_check_fn(script, slot): + with lock: + checks.append((slot, script)) + return 0, '' + + def workspace_snapshot_fn(slot): + with lock: + snaps.append(slot) + # Encode the slot in the snapshot so an episode reading the wrong + # slot's workspace would produce a mismatched check statement. + return f'slot_{slot}.txt 1\n\n--- slot_{slot}.txt ---\nx' + + def explorer(trajectories, **kw): + # The two follow-ups (check script, then statement) are threaded + # through the callback so the slot-aware handlers actually run. + tm = kw.get('tool_manager') + if tm is not None: + tm({'id': 'x', 'type': 'function', + 'function': {'name': 'noop', 'arguments': '{}'}}) + traj = {'messages': list(trajectories[0]['messages']), + 'input_ids': [1, 2, 3]} + followup = kw.get('followup_fn') + replies = ['```python\nassert True\n```', + 'Statement:\n\n```\ndo the thing\n```'] + for i, reply in enumerate(replies): + traj['messages'].append({'role': 'assistant', 'content': reply}) + if followup is None: + break + out = followup(traj, i) + if out is None: + break + text, _params = out + traj['messages'].append({'role': 'user', 'content': text}) + return [traj] + + emitted = [] + prompts = AgenticPrompts( + system='s', from_scratch='u {keywords}', + check_followup='c {final_state}', + check_retry_followup='{error} / {final_state}', + problem_followup='p') + ch = AgenticChallenger( + prompts, explorer, + reset_fn=reset_fn, + run_check_fn=run_check_fn, + workspace_snapshot_fn=workspace_snapshot_fn, + episode_concurrency=n_slots, + episode_tool_managers=tool_managers, + propose_sink=emitted.append, + solver_rollouts=0, + max_proposals_per_round=8, + ) + kept = ch._round(8) + + self.assertEqual(len(kept), 8) + # 8 episodes across 4 slots, evenly split -> each slot reset twice, ran + # its own check twice, and every check saw the slot's own snapshot text. + from collections import Counter + self.assertEqual(Counter(resets), Counter({0: 2, 1: 2, 2: 2, 3: 2})) + self.assertEqual(Counter(s for s, _ in checks), Counter({0: 2, 1: 2, 2: 2, 3: 2})) + # The tool_manager slot used matches the check slot for each episode. + self.assertEqual(Counter(tm_calls), Counter({0: 2, 1: 2, 2: 2, 3: 2})) + + def test_wrong_tool_manager_count_is_refused(self): + from twinkle_agentic.challenger.agentic import AgenticChallenger, AgenticPrompts + prompts = AgenticPrompts(system='s', from_scratch='u', + check_followup='c {final_state}', + check_retry_followup='{error} / {final_state}', + problem_followup='p') + with self.assertRaises(ValueError): + AgenticChallenger( + prompts, explorer=lambda t, **k: t, + reset_fn=lambda slot=0: None, + run_check_fn=lambda s, slot=0: (0, ''), + workspace_snapshot_fn=lambda slot=0: '', + episode_concurrency=4, + episode_tool_managers=[object(), object()], # wrong count + solver_rollouts=0) + + +class PreseedInputsTest(unittest.TestCase): + """A task carrying a setup script has it replayed before every attempt. + + Order is the whole point: clear, then write the inputs back, then let the + solver run. Replaying before the clear would delete the files it just wrote, + and skipping the replay would measure the task against a workspace missing the + data its statement says is there -- which reads as 'too hard' and is not. + """ + + def _challenger(self, run_check_fn, **kwargs): + from twinkle_agentic.challenger.agentic import AgenticChallenger, AgenticPrompts + + prompts = AgenticPrompts( + system='s', from_scratch='u', + check_followup='c {final_state}', + check_retry_followup='{error} / {final_state}', + problem_followup='p') + return AgenticChallenger( + prompts, + lambda trajs, **kw: [{'messages': list(t['messages']), 'stop_reason': 'stop'} + for t in trajs], + run_check_fn=run_check_fn, + workspace_snapshot_fn=lambda slot=0: 'input/a.csv 3\n', + solver_rollouts=2, + keep_min_pass=1, + keep_max_pass_margin=0, + propose_sink=[].append, + **kwargs) + + def _task(self, setup): + from twinkle_agentic.challenger.base import attach_user_data + return attach_user_data({'messages': [{'role': 'user', 'content': 'q'}]}, + check_script='assert True', setup_script=setup, + keywords=[]) + + def test_setup_runs_after_the_clear_and_before_the_check(self): + events = [] + ch = self._challenger( + run_check_fn=lambda script, slot=0: (events.append( + 'setup' if script.startswith('#SETUP') else 'check'), (0, ''))[1], + reset_fn=lambda slot=0: events.append('clear')) + kept = ch._filter_difficulty([self._task('#SETUP\nopen("a","w")')]) + self.assertEqual(events, ['clear', 'setup', 'check'] * 2) + self.assertEqual(len(kept), 1) + + def test_failed_setup_skips_the_attempt_instead_of_scoring_it_zero(self): + """An attempt that never ran must not be counted as an attempt that failed.""" + checks = [] + + def run_check(script, slot=0): + if script.startswith('#SETUP'): + return 1, 'no space left on device' + checks.append(script) + return 0, '' + + ch = self._challenger(run_check_fn=run_check, reset_fn=lambda slot=0: None) + kept = ch._filter_difficulty([self._task('#SETUP\nboom')]) + # The solver was never asked, so nothing was checked and nothing is kept. + self.assertEqual(checks, []) + self.assertEqual(kept, []) + self.assertEqual(ch.stats['setup_replay_fail'], 2) + + +class ParallelDifficultyTest(unittest.TestCase): + """Solver attempts run in waves, each attempt isolated in its own sandbox. + + The measurement is only a measurement if attempt A cannot pass on files + attempt B wrote, so what this pins down is that within one wave the clear, + the tool dispatch and the check all reach the *same* slot for a given + attempt, and that a wave is one batched explorer call rather than one call + per attempt (which is what left the GPUs idle). + """ + + def test_attempts_are_batched_per_wave_and_stay_in_their_slot(self): + from twinkle.data_format import user_data_get + from twinkle_agentic.challenger.agentic import AgenticChallenger, AgenticPrompts + from twinkle_agentic.challenger.base import attach_user_data + + n_slots = 4 + lock = threading.Lock() + batch_sizes = [] # trajectories per explorer call + reset_slots = [] # slot per reset + pairs = [] # (tool_manager slot, check slot) per attempt + + class FakeTM: + def __init__(self, slot): self.slot = slot + def tool_infos(self): return [] + def __call__(self, tc): return 'ok' + + tool_managers = [FakeTM(i) for i in range(n_slots)] + # Which slot's manager each trajectory of the current wave was handed. + wave_slots = [] + + def explorer(trajectories, **kw): + tms = kw.get('tool_manager') + with lock: + batch_sizes.append(len(trajectories)) + wave_slots.clear() + wave_slots.extend([tm.slot for tm in (tms or [])]) + return [{'messages': list(t['messages']), 'stop_reason': 'stop'} + for t in trajectories] + + # One attempt in flight per slot, so the check for the attempt that used + # slot k must itself run in slot k. Recorded as a pair to compare. + seq = iter(range(10_000)) + + def run_check_fn(script, slot=0): + with lock: + pairs.append(slot) + next(seq) + return 0, '' + + def reset_fn(slot=0): + with lock: + reset_slots.append(slot) + + prompts = AgenticPrompts( + system='s', from_scratch='u', + check_followup='c {final_state}', + check_retry_followup='{error} / {final_state}', + problem_followup='p') + ch = AgenticChallenger( + prompts, explorer, + reset_fn=reset_fn, + run_check_fn=run_check_fn, + workspace_snapshot_fn=lambda slot=0: 'a.txt 1\n', + episode_concurrency=n_slots, + episode_tool_managers=tool_managers, + solver_rollouts=4, + keep_min_pass=1, + keep_max_pass_margin=0, + propose_sink=[].append, + ) + tasks = [attach_user_data({'messages': [{'role': 'user', 'content': f'task {i}'}]}, + check_script='assert True', keywords=[]) + for i in range(2)] + kept = ch._filter_difficulty(tasks) + + # 2 tasks x 4 attempts = 8 attempts, 4 slots -> two waves of 4, each a + # single explorer call. Serial code would have made 8 calls of 1. + self.assertEqual(batch_sizes, [4, 4]) + # Every slot cleared once per wave, and the managers handed out are the + # slots that were cleared. + self.assertEqual(sorted(reset_slots), [0, 0, 1, 1, 2, 2, 3, 3]) + self.assertEqual(sorted(wave_slots), [0, 1, 2, 3]) + # One check per attempt, one per slot per wave. + self.assertEqual(sorted(pairs), [0, 0, 1, 1, 2, 2, 3, 3]) + # All checks passed -> both tasks scored 4 of 4. + self.assertEqual([user_data_get(t.get('user_data'), 'n_pass', -1) for t in kept], + [4, 4]) + + +class TruncatedSolverTest(unittest.TestCase): + """A solver attempt cut off at its token budget has to be countable. + + It is still scored as a failure -- whether to discount it decides which tasks + are kept, which is not this code's call -- but the count is what says whether + ``n_pass`` measured difficulty or the token budget. On one run 15 of 50 + attempts ended that way, all 15 with an untouched workspace. + """ + + def _challenger(self, attempt_flags, **kwargs): + """``attempt_flags``: one (truncated, passes) pair per solver attempt.""" + from twinkle_agentic.challenger.agentic import AgenticChallenger, AgenticPrompts + + self.flags = list(attempt_flags) + self.emitted = [] + + def explorer(trajectories, **kw): + truncated, _ = self.flags[0] + return [{'messages': list(t['messages']), + 'truncated': truncated, + 'stop_reason': 'length' if truncated else 'stop'} + for t in trajectories] + + def run_check_fn(script, slot=0): + _, passes = self.flags.pop(0) + return (0 if passes else 1), '' + + prompts = AgenticPrompts( + system='s', from_scratch='u', + check_followup='cs {final_state}', + check_retry_followup='{error} / {final_state}', problem_followup='ps') + return AgenticChallenger( + prompts, explorer, + reset_fn=lambda slot=0: None, + run_check_fn=run_check_fn, + workspace_snapshot_fn=lambda slot=0: 'a.txt 1\n', + propose_sink=self.emitted.append, + solver_rollouts=4, + **kwargs) + + def _task(self): + from twinkle.data_format import pack_user_data + return {'messages': [{'role': 'user', 'content': 'make a.txt'}], + 'user_data': pack_user_data({'check_script': 'assert True'}), + 'propose_rounds': [{'input_ids': [1]}]} + + def test_truncated_attempts_are_counted_and_still_scored_as_failures(self): + # 1 pass, 1 honest failure, 2 truncated failures -> 1 of 4, and the two + # truncations visible in stats so the 1-of-4 can be read for what it is. + ch = self._challenger([(False, True), (False, False), + (True, False), (True, False)], + keep_min_pass=1, keep_max_pass_margin=1) + kept = ch._filter_difficulty([self._task()]) + + self.assertEqual(ch.stats['solver_truncated'], 2) + self.assertEqual(self.emitted[0]['n_pass'], 1) + self.assertEqual(self.emitted[0]['n_rollouts'], 4) + self.assertEqual(self.emitted[0]['outcome'], 'kept') + self.assertEqual(len(kept), 1) + + def test_all_four_truncated_reads_as_nobody_solved_it(self): + # Pinned as the known cost of scoring them as failures: this task is + # discarded for being too hard and stats['solver_truncated'] == 4 is the + # only thing that says no solver ever acted. + ch = self._challenger([(True, False)] * 4, + keep_min_pass=1, keep_max_pass_margin=1) + kept = ch._filter_difficulty([self._task()]) + + self.assertEqual(ch.stats['solver_truncated'], 4) + self.assertEqual(self.emitted[0]['n_pass'], 0) + self.assertEqual(self.emitted[0]['outcome'], 'outside_band') + self.assertEqual(kept, []) + + +class KeywordBankTest(unittest.TestCase): + """The bank has to actually fill, and say so when it does not. + + It went empty for whole runs: the agentic prompt asked for "one keyword per + line" while ``parse_keyword_list`` reads a JSON array, so every generation + call parsed to nothing, ``_refill`` returned silently, and all 17 proposals in + one run ran the no-keyword prompt with no log line to say so. + """ + + def _challenger(self, reply_text): + from twinkle_agentic.challenger import KeywordStore + from twinkle_agentic.challenger.agentic import AgenticChallenger, AgenticPrompts + from prompts import KEYWORD_EXPAND_USER, KEYWORD_SYSTEM, KEYWORD_USER + + self.gen_records = [] + self.tool_explorer_calls = 0 + + def tool_explorer(trajectories, **kwargs): + self.tool_explorer_calls += 1 + return [{'messages': list(t['messages']) + + [{'role': 'assistant', 'content': reply_text}]} + for t in trajectories] + + def text_explorer(trajectories, **kwargs): + return [{'messages': list(t['messages']) + + [{'role': 'assistant', 'content': reply_text}], + 'stop_reason': 'stop'} + for t in trajectories] + + self.store = KeywordStore(os.path.join(self.tmp, 'kw.jsonl'), ('filesystem',)) + prompts = AgenticPrompts( + system='s', from_scratch='u', from_keywords='dir:\n{keywords}', + check_followup='cs {final_state}', + check_retry_followup='{error} / {final_state}', problem_followup='ps', + keyword_system=KEYWORD_SYSTEM, keyword_user=KEYWORD_USER, + keyword_expand_user=KEYWORD_EXPAND_USER) + return AgenticChallenger( + prompts, tool_explorer, + keyword_store=self.store, + category_desc={'filesystem': 'files and directories'}, + reset_fn=lambda slot=0: None, + run_check_fn=lambda script, slot=0: (0, ''), + workspace_snapshot_fn=lambda slot=0: 'a.txt 1\n', + keyword_explorer=text_explorer, + keyword_sink=self.gen_records.append, + keyword_gen_calls=1, + keyword_refill_target=4, + solver_rollouts=0) + + def setUp(self): + self.tmp = tempfile.mkdtemp(prefix='kwbank_test_') + + def tearDown(self): + shutil.rmtree(self.tmp, ignore_errors=True) + + def test_the_shipped_prompt_asks_for_what_the_parser_reads(self): + """The real prompt string, not a stand-in: this is the contract that broke.""" + from twinkle_agentic.challenger.code import parse_keyword_list + from prompts import KEYWORD_EXPAND_USER, KEYWORD_USER + + for text in (KEYWORD_USER, KEYWORD_EXPAND_USER): + self.assertIn('JSON array', text) + # And what a model following that instruction returns must parse. + self.assertEqual( + parse_keyword_list('["csv deduplication", "log rotation"]'), + ['csv deduplication', 'log rotation']) + # While the format the prompt used to ask for does not -- so a future + # rewording back to one-per-line fails here rather than in a night's run. + self.assertEqual(parse_keyword_list('csv deduplication\nlog rotation'), []) + + def test_a_json_reply_fills_the_bank_and_reaches_the_proposal(self): + from twinkle.data_format import user_data_get + ch = self._challenger('["csv deduplication", "log rotation"]') + proposals = ch.propose(1) + + self.assertTrue(self.store.texts('filesystem')) + picks = user_data_get(proposals[0].get('user_data'), 'keywords', []) + self.assertTrue(picks, 'the drawn keywords must reach the proposal') + self.assertIn(picks[0][1], proposals[0]['messages'][-1]['content']) + + def test_keyword_generation_does_not_use_the_tool_explorer(self): + ch = self._challenger('["csv deduplication"]') + ch.propose(1) + # Brainstorming a list needs no sandbox, and a bracketed list in the reply + # is exactly what the tool explorer would try to dispatch. + self.assertEqual(self.tool_explorer_calls, 0) + + def test_an_unparseable_reply_is_recorded_rather_than_swallowed(self): + from twinkle.data_format import user_data_get + ch = self._challenger('csv deduplication\nlog rotation') + proposals = ch.propose(1) + + self.assertEqual(user_data_get(proposals[0].get('user_data'), 'keywords', []), []) + self.assertTrue(self.gen_records, 'the keyword sink must see the failing call') + rec = self.gen_records[0] + self.assertEqual(rec['n_parsed'], 0) + self.assertEqual(rec['reply'], 'csv deduplication\nlog rotation') + self.assertIn('JSON array', rec['prompt']) + + +class SerialKeywordRefillTest(unittest.TestCase): + """A refill's calls go out one at a time, each told what the earlier ones said. + + Batched, the calls were identical but for a trailing index, and the 'do not + repeat these' list could only name what the bank already held -- which on a + first refill is nothing. Measured on armD: all eight parallel calls answered + with the same three phrases ('aggregating data', 'processing data', + 'generating a single output file'), and 22 of that run's 24 drawn phrases came + from that one batch. So what is pinned here is not that the code is serial but + the reason it is: call k+1 must be able to see call k's output. + """ + + def setUp(self): + self.tmp = tempfile.mkdtemp(prefix='kwserial_test_') + self.seen = [] # the user message of every call, in order + self.batch_sizes = [] # trajectories per call + + def tearDown(self): + shutil.rmtree(self.tmp, ignore_errors=True) + + def _challenger(self, prompts, category, refill_concurrency=1): + from twinkle_agentic.challenger import KeywordStore + from twinkle_agentic.challenger.agentic import AgenticChallenger + + n_kw = [0] + + def explorer(trajectories, **kwargs): + self.batch_sizes.append(len(trajectories)) + out = [] + for t in trajectories: + user = t['messages'][-1]['content'] + self.seen.append(user) + if 'KIND of work' in user: + reply = f'["kind {n_kw[0]}"]' + n_kw[0] += 1 + elif 'JSON array' in user: + reply = f'["topic {n_kw[0]}"]' + n_kw[0] += 1 + else: + reply = 'A draft task: read some files and compute something.' + out.append({'messages': list(t['messages']) + + [{'role': 'assistant', 'content': reply}], + 'stop_reason': 'stop'}) + return out + + self.store = KeywordStore(os.path.join(self.tmp, 'kw.jsonl'), (category,)) + return AgenticChallenger( + prompts, explorer, + keyword_store=self.store, + category_desc={category: 'some kind of work'}, + reset_fn=lambda slot=0: None, + run_check_fn=lambda script, slot=0: (0, ''), + workspace_snapshot_fn=lambda slot=0: 'a.txt 1\n', + keyword_explorer=explorer, + keyword_gen_calls=3, + keyword_refill_concurrency=refill_concurrency, + min_batch=1, + solver_rollouts=0) + + def _three_axis_prompts(self): + from twinkle_agentic.challenger.agentic import AgenticPrompts + from prompts import KEYWORD_EXPAND_USER, KEYWORD_SYSTEM, KEYWORD_USER + + return AgenticPrompts( + system='s', from_scratch='u', from_keywords='dir:\n{keywords}', + check_followup='cs {final_state}', + check_retry_followup='{error} / {final_state}', problem_followup='ps', + keyword_system=KEYWORD_SYSTEM, keyword_user=KEYWORD_USER, + keyword_expand_user=KEYWORD_EXPAND_USER) + + def test_three_axis_refill_shows_each_call_the_previous_output(self): + ch = self._challenger(self._three_axis_prompts(), 'transform') + + got = ch._generate_keywords('transform', 9) + + self.assertEqual(sorted(got), ['topic 0', 'topic 1', 'topic 2']) + self.assertEqual(self.batch_sizes, [1, 1, 1], + 'one call at a time, or the calls cannot see each other') + self.assertNotIn('topic 0', self.seen[0], 'nothing exists yet for the first call') + self.assertIn('topic 0', self.seen[1]) + for kw in ('topic 0', 'topic 1'): + self.assertIn(kw, self.seen[2]) + + + def test_raising_the_concurrency_restores_the_batched_behaviour(self): + """At n_calls in flight, no call can see any other -- the first round's setup. + + Kept measurable on one build: the arms were first compared with a whole + refill going out at once, and telling that apart from what serial produces + means being able to run both without checking out an older file. + """ + ch = self._challenger(self._three_axis_prompts(), 'transform', + refill_concurrency=3) + + got = ch._generate_keywords('transform', 9) + + self.assertEqual(sorted(got), ['topic 0', 'topic 1', 'topic 2']) + self.assertEqual(self.batch_sizes, [3], 'all three go out as one batch') + for user in self.seen: + for kw in ('topic 0', 'topic 1', 'topic 2'): + self.assertNotIn(kw, user, 'a batched call cannot see its siblings') + + def test_the_avoid_list_is_capped_and_drops_older_entries_first(self): + """The cap evicts banked phrases before this refill's own, and holds a ceiling. + + Both halves matter. Capping by sampling the whole list would start dropping + exactly what this refill just produced, and the serial ordering would buy + nothing. Not capping at all is what made the eighth call of armA2ser's + edge_case refill invent 'îRAPIÓN holistic replace' and nine other + non-phrases: 150 quoted phrases left it no room to answer. + """ + from twinkle_agentic.challenger.agentic import AgenticPrompts + from prompts import KEYWORD_EXPAND_USER, KEYWORD_SYSTEM, KEYWORD_USER + + prompts = AgenticPrompts( + system='s', from_scratch='u', from_keywords='dir:\n{keywords}', + check_followup='cs {final_state}', + check_retry_followup='{error} / {final_state}', problem_followup='ps', + keyword_system=KEYWORD_SYSTEM, keyword_user=KEYWORD_USER, + keyword_expand_user=KEYWORD_EXPAND_USER) + ch = self._challenger(prompts, 'transform') + cap = ch._AVOID_TOTAL + older = [f'old {i}' for i in range(200)] + fresh = [f'new {i}' for i in range(5)] + note = ch._avoid_note(older, fresh, 'avoid: ') + for kw in fresh: + self.assertIn(kw, note) + self.assertEqual(note.count('old '), cap - len(fresh)) + + # Once this refill alone fills the cap, no banked phrase is quoted and the + # line stops growing -- it is the growth that broke the eighth call. + many = [f'new {i}' for i in range(cap + 30)] + note = ch._avoid_note(older, many, 'avoid: ') + self.assertEqual(note.count('old '), 0) + self.assertEqual(note.count('new '), cap) + self.assertNotIn('new 0', note, 'the oldest of this refill falls off first') + self.assertIn(f'new {cap + 29}', note, 'the newest is always kept') + + if __name__ == '__main__': unittest.main() diff --git a/tests/twinkle_agentic/test_harness.py b/tests/twinkle_agentic/test_harness.py index 93a3b9c21..7c1d36753 100644 --- a/tests/twinkle_agentic/test_harness.py +++ b/tests/twinkle_agentic/test_harness.py @@ -152,6 +152,40 @@ def test_tool_manager_call_many_uses_env_step_batch(): assert out[1].startswith('lookup:') +def test_unknown_name_does_not_push_the_turn_off_the_batch(): + """One name the manager can refuse must not make the rest run concurrently. + + The thread pool is the fallback for tools that share no Env, and dispatching + a sandbox turn through it is how four calls in one ex4 episode came back with + a single tool's answer. A refusable name is answered here, and the calls that + can run still go as one ordered batch. + """ + env = BatchEnv() + mgr = ToolManager(EnvTool.from_env(env)) + calls = [ + {'type': 'function', 'function': {'name': 'search', 'arguments': {'q': 'a'}}}, + {'type': 'function', 'function': {'name': 'no_such_tool', 'arguments': {}}}, + {'type': 'function', 'function': {'name': 'lookup', 'arguments': {'k': 'b'}}}, + ] + out = mgr.call_many(calls) + assert env.batch_calls == 1 + assert env.step_calls == 2 + assert out[0].startswith('search:') + assert out[1].startswith("Error: unknown tool 'no_such_tool'") + assert out[2].startswith('lookup:') + + +def test_call_many_all_names_unknown_never_reaches_the_env(): + env = BatchEnv() + mgr = ToolManager(EnvTool.from_env(env)) + out = mgr.call_many([ + {'type': 'function', 'function': {'name': 'nope', 'arguments': {}}}, + {'type': 'function', 'function': {'name': 'also_nope', 'arguments': {}}}, + ]) + assert env.step_calls == 0 + assert all(o.startswith('Error: unknown tool') for o in out) + + def _run_wrapped(source: str): """exec the wrapper the way ms-agent's python_executor does: split dicts.""" import io diff --git a/tests/twinkle_agentic/test_multi_turn_rollout.py b/tests/twinkle_agentic/test_multi_turn_rollout.py index 4f17d8282..21c1bc8f5 100644 --- a/tests/twinkle_agentic/test_multi_turn_rollout.py +++ b/tests/twinkle_agentic/test_multi_turn_rollout.py @@ -156,6 +156,10 @@ def parse_tool_call(self, decoded: str) -> list[dict[str, Any]]: }) return results + def clean_tool_call(self, decoded: str) -> str: + """Strip the call blocks, as the real template does before storing.""" + return re.sub(r'<tool_call>[\s\S]*?</tool_call>', '', decoded or '') + # --- Used by the fake sampler to mirror real concat_input_feature ------- def concat_input_feature(self, pif: dict[str, Any], new_tokens: list[int]) -> dict[str, Any]: result = copy.deepcopy(pif) @@ -171,10 +175,20 @@ def concat_input_feature(self, pif: dict[str, Any], new_tokens: list[int]) -> di result['input_ids'] = input_ids result['labels'] = labels result = self._invoke_post_pipeline([result])[0] - # Append assistant message with the decoded response (no special toks) + # Append assistant message with the decoded response (no special toks). + # A reply that parses as a call is stored with the call text removed and + # the calls in their own field, which is what the real template does -- + # and the reason a stage reply has to be put back afterwards. response_text = self.tokenizer.decode(new_tokens, skip_special_tokens=True) messages = list(result.get('messages') or []) - messages.append({'role': 'assistant', 'content': response_text}) + parsed = self.parse_tool_call(response_text) + msg: dict[str, Any] = { + 'role': 'assistant', + 'content': self.clean_tool_call(response_text) if parsed else response_text, + } + if parsed: + msg['tool_calls'] = parsed + messages.append(msg) result['messages'] = messages return result @@ -186,6 +200,9 @@ def __init__(self, template: FakeTemplate) -> None: self.template = template self._queue: list[dict[str, Any]] = [] self.sample_calls = 0 + # One entry per sample() call, so a test can assert which budget each + # stage was sampled under. + self.params_seen: list[Any] = [] def queue( self, @@ -214,6 +231,7 @@ def sample(self, pifs, sampling_params=None): if isinstance(pifs, dict): pifs = [pifs] assert isinstance(pifs, list), (f'FakeSampler.sample expects a list, got {type(pifs).__name__}') + self.params_seen.append(sampling_params) responses: list[SampleResponse] = [] for pif in pifs: assert self._queue, 'FakeSampler queue exhausted — scripted turns' @@ -251,6 +269,32 @@ def tool_info(self): } +class FailTool(Tool): + """Answers in the two shapes a real failure arrives in. + + ``kind='envelope'`` is ms-agent wrapping a failure; ``kind='bare'`` is a + dispatch that never reached a tool. Both copied from a recorded run. + """ + + def __init__(self, name: str = 'grep', kind: str = 'envelope'): + self._name = name + self._kind = kind + + def __call__(self, tool_name: str, arguments: dict[str, Any]) -> str: + if self._kind == 'bare': + return (f"Error: unknown tool '{tool_name}'. " + f'Available: code_executor---shell_executor') + return ('{\n "success": false,\n "output": "",\n' + ' "error": "[Errno 2] No such file or directory"\n}') + + def tool_info(self): + return { + 'type': 'function', + 'function': {'name': self._name, 'description': 'always fails', + 'parameters': {}}, + } + + # ============================================================================= # Fixtures # ============================================================================= @@ -273,19 +317,23 @@ def sampler(template): def tool_manager(): mgr = ToolManager({}) mgr.register(EchoTool('search')) + mgr.register(FailTool('grep')) + mgr.register(FailTool('badname', kind='bare')) return mgr @pytest.fixture def make_rollout(sampler, template, tool_manager): - def _make(max_turns: int = 4, sampling_params: SamplingParams | None = None): + def _make(max_turns: int = 4, sampling_params: SamplingParams | None = None, + stop_after_stuck_turns: int = 0): return MultiTurnRollout( sampler=sampler, template=template, tool_manager=tool_manager, sampling_params=sampling_params or SamplingParams(), max_turns=max_turns, + stop_after_stuck_turns=stop_after_stuck_turns, ) return _make @@ -340,7 +388,10 @@ def test_single_turn_length_stop(make_rollout, sampler): # short-circuit BEFORE we parse / dispatch tools. assert out['turns'] == 1 assert out['stop_reason'] == 'length' - assert out['truncated'] is False + # Running out of generation budget is a truncation, like the max_turns and + # max_trajectory_tokens cases: a consumer filtering on this flag must not see + # a cut-off trajectory as one that reached its own conclusion. + assert out['truncated'] is True assert sampler.sample_calls == 1 # No tool message should have been appended. roles = [m['role'] for m in out['messages']] @@ -411,6 +462,153 @@ def test_max_turns_natural_stop_at_ceiling(make_rollout, sampler): assert out['truncated'] is False +def test_max_turns_one_dispatches_no_tool(make_rollout, sampler): + """A one-turn rollout never runs a tool, even when the reply asks for one. + + This is what a caller relies on to get a text-only round out of a rollout that + requires a tool manager at construction: the challenger's check-writing round + must not be able to touch the workspace its script is about to be verified + against, and a reply containing python parses as a tool call whether or not + the model meant one. + """ + sampler.queue(_tool_call_text('search', {'q': 'x'}), stop_reason='stop') + rollout = make_rollout(max_turns=1) + out = rollout([_user_traj()])[0] + + assert out['turns'] == 1 + assert [m['role'] for m in out['messages']].count('tool') == 0 + # The fake tool echoes what it was called with, so its absence anywhere in + # the transcript is proof it never ran. + assert 'echo[' not in ''.join(m.get('content') or '' for m in out['messages']) + + +# ============================================================================= +# Tests: stuck-episode early stop +# +# Measured on 12 recorded sandbox episodes: 131 of 239 tool calls were +# byte-identical repeats of an earlier call, and the two worst episodes burned 54 +# and 84 calls to leave behind a single script that could not run. Stopping on +# errors alone would have caught 1 of the 12 -- the offenders interleave a failing +# call with a glob that succeeds -- so a turn also counts as stuck when every call +# in it repeats one already made. +# ============================================================================= +def test_stuck_stop_off_by_default(make_rollout, sampler): + """Two failing turns run on when the limit is 0: existing callers see no change.""" + sampler.queue(_tool_call_text('grep', {'p': 1}), stop_reason='stop') + sampler.queue(_tool_call_text('grep', {'p': 2}), stop_reason='stop') + sampler.queue('Done.', stop_reason='stop') + out = make_rollout(max_turns=4)([_user_traj()])[0] + + assert out['stuck_stop'] is False + assert out['turns'] == 3 + + +def test_two_all_error_turns_stop_the_episode(make_rollout, sampler): + sampler.queue(_tool_call_text('grep', {'p': 1}), stop_reason='stop') + sampler.queue(_tool_call_text('grep', {'p': 2}), stop_reason='stop') + # Would have been a third turn; the stop means it is never sampled. + sampler.queue(_tool_call_text('search', {'q': 'x'}), stop_reason='stop') + out = make_rollout(max_turns=6, stop_after_stuck_turns=2)([_user_traj()])[0] + + assert out['stuck_stop'] is True + assert out['truncated'] is True + assert out['turns'] == 2 + assert sampler.sample_calls == 2 + # The failures that ended it are in the transcript the caller reads, so the + # reason is visible without re-running anything. + assert [m['role'] for m in out['messages']].count('tool') == 2 + + +def test_bare_error_string_counts_as_a_failure(make_rollout, sampler): + """An unknown tool name never reaches a tool; that is still a failed turn.""" + sampler.queue(_tool_call_text('badname', {'a': 1}), stop_reason='stop') + sampler.queue(_tool_call_text('badname', {'a': 2}), stop_reason='stop') + out = make_rollout(max_turns=6, stop_after_stuck_turns=2)([_user_traj()])[0] + + assert out['stuck_stop'] is True + assert out['turns'] == 2 + + +def test_two_verbatim_repeat_turns_stop_the_episode(make_rollout, sampler): + """Repeating a *successful* call is stuck too -- it cannot produce new state.""" + sampler.queue(_tool_call_text('search', {'q': 'a'}), stop_reason='stop') + sampler.queue(_tool_call_text('search', {'q': 'a'}), stop_reason='stop') + sampler.queue(_tool_call_text('search', {'q': 'a'}), stop_reason='stop') + sampler.queue('Done.', stop_reason='stop') + out = make_rollout(max_turns=6, stop_after_stuck_turns=2)([_user_traj()])[0] + + assert out['stuck_stop'] is True + assert out['turns'] == 3 + + +def test_changed_arguments_are_not_a_repeat(make_rollout, sampler): + sampler.queue(_tool_call_text('search', {'q': 'a'}), stop_reason='stop') + sampler.queue(_tool_call_text('search', {'q': 'b'}), stop_reason='stop') + sampler.queue(_tool_call_text('search', {'q': 'c'}), stop_reason='stop') + sampler.queue('Done.', stop_reason='stop') + out = make_rollout(max_turns=6, stop_after_stuck_turns=2)([_user_traj()])[0] + + assert out['stuck_stop'] is False + assert out['turns'] == 4 + + +def test_one_success_in_a_turn_resets_the_count(make_rollout, sampler): + """The case that decided the rule: a failing call next to a useful one. + + Counting these as stuck would stop at turn 2 -- and in the recorded run the + files worth writing a check about were created after that point. + """ + for i in range(3): + sampler.queue(_tool_call_text('grep', {'p': i}) + + _tool_call_text('search', {'q': i}), stop_reason='stop') + sampler.queue('Done.', stop_reason='stop') + out = make_rollout(max_turns=6, stop_after_stuck_turns=2)([_user_traj()])[0] + + assert out['stuck_stop'] is False + assert out['turns'] == 4 + + +def test_a_good_turn_between_two_bad_ones_resets_the_count(make_rollout, sampler): + sampler.queue(_tool_call_text('grep', {'p': 1}), stop_reason='stop') + sampler.queue(_tool_call_text('search', {'q': 'new'}), stop_reason='stop') + sampler.queue(_tool_call_text('grep', {'p': 2}), stop_reason='stop') + sampler.queue('Done.', stop_reason='stop') + out = make_rollout(max_turns=6, stop_after_stuck_turns=2)([_user_traj()])[0] + + assert out['stuck_stop'] is False + assert out['turns'] == 4 + + +def test_stuck_stop_is_per_trajectory_in_a_batch(make_rollout, sampler, template): + """One stuck episode must not end its batch mates.""" + good = ToolManager({}) + good.register(EchoTool('search')) + bad = ToolManager({}) + bad.register(FailTool('search')) + + sampler.queue(_tool_call_text('search', {'q': 1}), stop_reason='stop') + sampler.queue(_tool_call_text('search', {'q': 1}), stop_reason='stop') + sampler.queue(_tool_call_text('search', {'q': 2}), stop_reason='stop') + sampler.queue(_tool_call_text('search', {'q': 3}), stop_reason='stop') + sampler.queue('Done.', stop_reason='stop') + sampler.queue('Done.', stop_reason='stop') + + rollout = MultiTurnRollout( + sampler=sampler, template=template, tool_manager=[good, bad], + sampling_params=SamplingParams(), max_turns=6, stop_after_stuck_turns=2) + outs = rollout([_user_traj('a'), _user_traj('b')]) + + assert outs[1]['stuck_stop'] is True + assert outs[0]['stuck_stop'] is False + assert outs[0]['turns'] > outs[1]['turns'] + + +def test_rejects_negative_stuck_limit(sampler, template, tool_manager): + with pytest.raises(ValueError, match='stop_after_stuck_turns'): + MultiTurnRollout(sampler=sampler, template=template, + tool_manager=tool_manager, stop_after_stuck_turns=-1) + + # ============================================================================= # Tests: label & logprobs alignment # ============================================================================= @@ -869,3 +1067,238 @@ def test_trace_dir_uses_user_data_id_in_filename(tmp_path, sampler, template, to # Slashes are sanitised away; the id still drives the filename. assert 'hotpotqa_42' in files[0] assert files[0].startswith('fail-') + + +# ============================================================================= +# followup_fn: several stages, one trajectory +# ============================================================================= +def test_followup_appends_a_user_turn_and_keeps_generating(sampler, template, tool_manager): + """A stage that ends without tool calls continues when the callback says so.""" + asked = [] + + def followup(traj, n_before): + asked.append((n_before, len(traj['messages']))) + return ['write the checks', 'write the statement'][n_before] if n_before < 2 else None + + rollout = MultiTurnRollout( + sampler=sampler, template=template, tool_manager=tool_manager, + sampling_params=SamplingParams(), max_turns=8, followup_fn=followup) + sampler.queue(_tool_call_text('search', {'q': 'x'}), stop_reason='stop') + sampler.queue('Done.', stop_reason='stop') + sampler.queue('```python\nassert True\n```', stop_reason='stop') + sampler.queue('The statement.', stop_reason='stop') + + out = rollout([_user_traj()])[0] + + assert [n for n, _ in asked] == [0, 1, 2] + assert out['followups'] == 2 + roles = [m['role'] for m in out['messages']] + # user, assistant(tool call), tool, assistant(Done.), user, assistant(checks), + # user, assistant(statement) + assert roles == ['user', 'assistant', 'tool', 'assistant', 'user', 'assistant', + 'user', 'assistant'] + assert out['messages'][4]['content'] == 'write the checks' + assert out['messages'][6]['content'] == 'write the statement' + + +def test_every_assistant_stage_stays_trainable(sampler, template, tool_manager): + """The whole chain trains: no stage is demoted to prompt by the follow-ups. + + This is the reason follow-ups are appended inside one rollout instead of + starting a second one on the finished conversation: a second rollout encodes + the history as its prompt, which sets labels to -100 for every earlier + assistant turn and leaves only the last stage trainable. + """ + def followup(traj, n_before): + return 'next stage' if n_before < 2 else None + + rollout = MultiTurnRollout( + sampler=sampler, template=template, tool_manager=tool_manager, + sampling_params=SamplingParams(), max_turns=8, followup_fn=followup) + replies = [_tool_call_text('search', {'q': 'x'}), 'Done.', 'CHECKS', 'STATEMENT'] + for i, text in enumerate(replies): + sampler.queue(text, stop_reason='stop', logprobs=[-0.5] * len( + template.tokenizer.encode(text + '<|im_end|>', add_special_tokens=False))) + + out = rollout([_user_traj()])[0] + + trainable = _count_trainable(out['labels']) + expected = sum(len(template.tokenizer.encode(text + '<|im_end|>', add_special_tokens=False)) + for text in replies) + assert trainable == expected + # The alignment invariant GRPO depends on: one logprob per trainable label. + assert len(out['logprobs']) == trainable + + +def test_followup_stage_can_use_its_own_sampling_params(sampler, template, tool_manager): + """``(text, params)`` gives that stage its own budget, without touching others.""" + small = SamplingParams(max_tokens=17) + + def followup(traj, n_before): + return ('write the checks', small) if n_before == 0 else None + + rollout = MultiTurnRollout( + sampler=sampler, template=template, tool_manager=tool_manager, + sampling_params=SamplingParams(max_tokens=99), max_turns=6, followup_fn=followup) + sampler.queue('Done.', stop_reason='stop') + sampler.queue('CHECKS', stop_reason='stop') + + rollout([_user_traj()]) + + assert [p.max_tokens for p in sampler.params_seen] == [99, 17] + + +def test_tool_calls_are_not_dispatched_after_a_followup(sampler, template, tool_manager): + """Python in a check script parses as a call list; it must not run.""" + def followup(traj, n_before): + return 'write the checks' if n_before == 0 else None + + rollout = MultiTurnRollout( + sampler=sampler, template=template, tool_manager=tool_manager, + sampling_params=SamplingParams(), max_turns=6, followup_fn=followup) + sampler.queue('Done.', stop_reason='stop') + sampler.queue(_tool_call_text('search', {'q': 'should not run'}), stop_reason='stop') + + out = rollout([_user_traj()])[0] + + assert not any(m['role'] == 'tool' for m in out['messages']) + assert out['followups'] == 1 + + +# ============================================================================= +# Appending a user turn under a template that moves reasoning blocks around +# ============================================================================= +class ThinkAwareTokenizer(FakeTokenizer): + """Renders like Qwen3: reasoning is kept only after the last user turn. + + Two rules, both measured on Qwen3-4B's own template: an assistant turn that + precedes the last user message loses its ``<think>`` block, and the trailing + assistant turn gains an empty one when it has none. Together they mean that + appending a user message rewrites earlier text, so the plain + "render before, render after, take the difference" bridge cannot be used. + """ + + def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=False, **_): + last_user = max((i for i, m in enumerate(messages) if m['role'] == 'user'), default=-1) + s = '' + for i, m in enumerate(messages): + content = m['content'] + if m['role'] == 'assistant': + if i < last_user: + content = re.sub(r'<think>[\s\S]*?</think>\n*', '', content) + elif '<think>' not in content: + content = '<think>\n\n</think>\n\n' + content + s += f"<|im_start|>{m['role']}\n{content}<|im_end|>\n" + if add_generation_prompt: + s += '<|im_start|>assistant\n' + return self.encode(s) if tokenize else s + + +def test_appending_a_user_turn_keeps_the_history_ids_and_adds_only_the_new_block(): + """The delta is the new user block plus the generation prompt, nothing else.""" + from twinkle_agentic.rollout.bridge import extend_with_bridge + + template = FakeTemplate(ThinkAwareTokenizer()) + messages = [{'role': 'user', 'content': 'do work'}, + {'role': 'assistant', 'content': '<think>reasoning</think>Done.'}] + pif = template.encode({'messages': messages}) + pif['labels'] = [7] * len(pif['input_ids']) # stand-in for "these were sampled" + before_ids = list(pif['input_ids']) + + out = extend_with_bridge(pif, [{'role': 'user', 'content': 'write the checks'}], template) + + # History untouched: the reasoning the policy produced is still in the ids. + assert out['input_ids'][:len(before_ids)] == before_ids + added = template.tokenizer.decode(out['input_ids'][len(before_ids):]) + assert added == ('<|im_start|>user\nwrite the checks<|im_end|>\n' + '<|im_start|>assistant\n'), added + # And the appended block is not trained on. + assert set(out['labels'][len(before_ids):-1]) == {-100} + + +def test_a_template_that_really_reorders_history_still_raises(): + """The fallback must not paper over a template that rewrites message blocks.""" + from twinkle_agentic.rollout.bridge import extend_with_bridge + + class ReorderingTokenizer(FakeTokenizer): + """Puts the message count up front, so every append rewrites the start.""" + + def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=False, **_): + s = f'[{len(messages)} messages]' + for m in messages: + s += f"<|im_start|>{m['role']}\n{m['content']}<|im_end|>\n" + if add_generation_prompt: + s += '<|im_start|>assistant\n' + return self.encode(s) if tokenize else s + + template = FakeTemplate(ReorderingTokenizer()) + pif = template.encode({'messages': [{'role': 'user', 'content': 'a'}, + {'role': 'assistant', 'content': 'b'}]}) + with pytest.raises(RuntimeError, match='non-monotonic'): + extend_with_bridge(pif, [{'role': 'user', 'content': 'c'}], template) + + +def test_running_out_of_tool_turns_still_reaches_the_follow_up_stages(sampler, template, tool_manager): + """An episode that spends its whole turn budget is not thrown away. + + Before, hitting ``max_turns`` ended the trajectory outright -- and with the + stages living inside the episode that would throw away the sandbox run that + produced the state they are about. + """ + asked = [] + + def followup(traj, n_before): + asked.append(n_before) + return 'write the checks' if n_before == 0 else None + + rollout = MultiTurnRollout( + sampler=sampler, template=template, tool_manager=tool_manager, + sampling_params=SamplingParams(), max_turns=2, followup_fn=followup) + # Two turns of tool calls: the second one hits the limit. + sampler.queue(_tool_call_text('search', {'q': 'a'}), stop_reason='stop') + sampler.queue(_tool_call_text('search', {'q': 'b'}), stop_reason='stop') + sampler.queue('assert True', stop_reason='stop') + + out = rollout([{'messages': [{'role': 'user', 'content': 'go'}]}])[0] + + assert asked == [0, 1] + assert out['tool_stop'] == 'max_turns' + # The stage ran, so nothing was cut off. + assert out['truncated'] is False + assert out['messages'][-2:] == [{'role': 'user', 'content': 'write the checks'}, + {'role': 'assistant', 'content': 'assert True'}] + + +def test_a_stage_reply_that_looks_like_a_tool_call_is_kept_whole(sampler, template, tool_manager): + """The stage reply the caller reads is what the model wrote. + + The template stores a reply that parses as a call with the call text removed, + which is right for a turn whose calls get dispatched and wrong for a stage + whose reply *is* the answer. It bit for real: one of the tool-call formats is + XML-shaped, so a check script asserting the content of an .xml file parsed as + calls, and 5 of ex12's 72 scripts arrived with that content deleted -- three + then asserted `content == ''` against a file that had text in it. + """ + + def followup(traj, n_before): + return 'write the checks' if n_before == 0 else None + + rollout = MultiTurnRollout( + sampler=sampler, template=template, tool_manager=tool_manager, + sampling_params=SamplingParams(), max_turns=4, followup_fn=followup) + sampler.queue('done exploring', stop_reason='stop') + script = ('```python\n' + _tool_call_text('data', {'number': '75'}) + + "\nassert open('a.xml').read() == 'x'\n```") + sampler.queue(script, stop_reason='stop') + + out = rollout([{'messages': [{'role': 'user', 'content': 'go'}]}])[0] + + last = out['messages'][-1] + assert last['content'] == script + assert 'tool_calls' not in last + # Whole, but still without the special tokens: the sampled ids end with + # <|im_end|> and ``seq.decoded`` may keep it. ex13 shipped 7 of 7 problem + # statements ending in a literal '<|im_end|>' that way. + assert '<|im_end|>' not in last['content'] + # And it was not dispatched: a dispatch appends a tool message. + assert [m['role'] for m in out['messages'] if m['role'] == 'tool'] == [] diff --git a/tests/twinkle_agentic/test_tools.py b/tests/twinkle_agentic/test_tools.py index 87c050f6b..cfa5def56 100644 --- a/tests/twinkle_agentic/test_tools.py +++ b/tests/twinkle_agentic/test_tools.py @@ -157,6 +157,35 @@ def test_call_missing_tool(self): assert 'unknown tool' in result assert 'Available:' in result + def test_unqualified_name_is_refused_with_the_qualified_one(self): + # Measured over 5793 RSI calls: 201 bare 'shell_executor' and 30 filed + # under the wrong server, all naming a tool that does exist. The call + # still fails -- silently resolving it would train a name that no serving + # deployment accepts -- but the reply says which name to use. + tm = ToolManager({'code_executor---shell_executor': MockTool()}) + for wrong in ('shell_executor', 'file_system---shell_executor'): + result = tm({'function': {'name': wrong, 'arguments': {}}}) + assert 'unknown tool' in result + assert "Did you mean 'code_executor---shell_executor'" in result + + def test_no_guess_when_the_bare_name_is_ambiguous(self): + # Two servers exporting the same verb: any guess would be a coin toss. + tm = ToolManager({ + 'a---read_file': MockTool('a---read_file'), + 'b---read_file': MockTool('b---read_file'), + }) + result = tm({'function': {'name': 'read_file', 'arguments': {}}}) + assert 'Did you mean' not in result + + def test_no_guess_when_nothing_resembles_the_name(self): + # file_system---list_directory, 107 times: the model wants a tool this + # line-up does not have. The available list is the only useful answer. + tm = ToolManager({'file_system---glob': MockTool('file_system---glob')}) + result = tm({'function': {'name': 'file_system---list_directory', 'arguments': {}}}) + assert 'Did you mean' not in result + assert 'file_system---glob' in result + assert 'Available:' in result + def test_call_missing_function(self): tm = ToolManager({'mock': MockTool()}) result = tm({}) diff --git a/tests/twinkle_client/test_client_multi_turn_rollout.py b/tests/twinkle_client/test_client_multi_turn_rollout.py index 1ff69f5f4..6b9474a1e 100644 --- a/tests/twinkle_client/test_client_multi_turn_rollout.py +++ b/tests/twinkle_client/test_client_multi_turn_rollout.py @@ -428,6 +428,31 @@ def test_max_turns_one_forces_truncation(logprobs_flags): assert out['turns'] == 1 +@settings(deadline=None, max_examples=60) +@given(logprobs_flags=st.lists(st.booleans(), min_size=1, max_size=5)) +def test_length_stop_marks_truncated(logprobs_flags): + """A reply cut off at the generation budget is ``truncated=True``. + + Same flag as the ``max_turns`` edge above: a consumer that filters on + ``truncated`` to separate trajectories that concluded from ones that ran out + of room would otherwise treat a cut-off reply as a finished one. + """ + # Terminal turn ends on 'length' with no tool-call turns before it, so the + # very first generation is the one that gets cut. + scripts_spec = [{'num_tools': 0, 'terminal': 'length', 'logprobs': lp} for lp in logprobs_flags] + trajectories, sampler, template = _build_from_scripts(scripts_spec) + rollout = ClientMultiTurnRollout( + sampler=sampler, template=template, tool_manager=_make_tool_manager(), max_turns=4) + + outs = rollout(copy.deepcopy(trajectories)) + + assert len(outs) == len(trajectories) + for out in outs: + assert out['stop_reason'] == 'length' + assert out['truncated'] is True + assert out['turns'] == 1 + + # ============================================================================= # Deterministic unit tests: exception paths & dependency reuse (non-hypothesis) # From 3408613b1f134b673ebc21d11989784d2aa21631 Mon Sep 17 00:00:00 2001 From: root <yuze.zyz@alibaba-inc.com> Date: Thu, 27 Aug 2026 16:58:04 +0800 Subject: [PATCH 45/60] wip --- cookbook/rsi/agentic/challenge.py | 89 ++- cookbook/rsi/agentic/rl.py | 18 +- cookbook/rsi/agentic/train_offline.py | 517 ++++++++++++++++++ src/twinkle/infra/__init__.py | 327 ++++++++++- src/twinkle/infra/_ray/ray_helper.py | 17 + .../sampler/vllm_sampler/vllm_sampler.py | 9 +- src/twinkle/template/base.py | 10 + src/twinkle/template/tools/base.py | 17 + src/twinkle/template/tools/bracket_dsl.py | 14 +- src/twinkle/template/tools/cline.py | 15 +- src/twinkle/template/tools/qwen.py | 16 +- src/twinkle/template/tools/vcp.py | 14 +- src/twinkle_agentic/challenger/agentic.py | 135 ++++- src/twinkle_agentic/rollout/multi_turn.py | 67 +++ src/twinkle_agentic/tools/tool_manager.py | 39 +- 15 files changed, 1244 insertions(+), 60 deletions(-) create mode 100644 cookbook/rsi/agentic/train_offline.py diff --git a/cookbook/rsi/agentic/challenge.py b/cookbook/rsi/agentic/challenge.py index 5bfbce3f9..fe0f6321d 100644 --- a/cookbook/rsi/agentic/challenge.py +++ b/cookbook/rsi/agentic/challenge.py @@ -27,6 +27,8 @@ import hashlib import os import sys +import time +from typing import Dict import numpy as np import twinkle @@ -117,6 +119,12 @@ handle.write(base64.b64decode(payload)) ''' +# Seconds to wait before asking a sandbox for its workspace listing a second +# time. 62 of run_clean6's 63 snapshot failures were the sandbox answering 410 +# "not proxyable", which is the host having paused it -- worth one more ask, +# since the alternative is throwing the episode away. +SNAPSHOT_RETRY_WAIT = 3 + # The ground truth the check script is written against. A listing alone is not # enough: three of the six rejected proposals in the first real run failed on a # value the model recomputed from its own recollection ("Mean values mismatch") @@ -229,16 +237,25 @@ def parse_args(): # trained; the two API stages are text-only and never enter the trajectory. # Reuses the --challenger-api-* connection args, and is mutually exclusive with # --challenger-api (which sends the whole proposing side to the API). - p.add_argument('--followup-api', action='store_true', + # On by default, because leaving it off is not a milder setting but a + # different experiment: run_clean6 was launched without it and the local 4B + # wrote the check scripts, which turned 0/55 check_parse_fail (qwen3-max) into + # 35/171 and 1/55 check_run_fail into 25/171. Pass --no-followup-api to write + # both stages with the local model on purpose. + p.add_argument('--followup-api', action=argparse.BooleanOptionalAction, + default=True, help='explore locally (trainable) but generate the check script and ' 'problem statement over --challenger-api-* (e.g. qwen3-max); ' 'only the exploration part is trained. Not with --challenger-api.') # How many episodes run at once, each in its own sandbox. An episode owns its # workspace from the reset until its check has run, so this is also the number - # of sandboxes booted at startup. Default 48: episodes alternate between vLLM - # generation (~24 concurrent sequences fit in KV cache) and sandbox execution, - # so 48 keeps both the GPU cluster and the sandbox host saturated. - p.add_argument('--episode-concurrency', type=int, default=48, + # of sandboxes booted at startup. Default 96, from two measurements: on 8 vLLM + # workers, 96 requests in flight reached 28894 tok/s against a 33108 tok/s + # ceiling (87%), where 48 in flight reached only 15599 (47%); and the sandbox + # host booted 96 concurrent sandboxes with 0 failures at a p50 reset of 24.5s, + # up from 21.9s at 64. Above 96 the KV cache is the next limit -- vLLM reports + # room for about 24 sequences per worker, so roughly 192 in total. + p.add_argument('--episode-concurrency', type=int, default=96, help='sandboxes to boot, and how many things run at once in both ' 'stages: proposal episodes in flight, and solver attempts ' 'per wave in the difficulty filter') @@ -350,6 +367,15 @@ def parse_args(): p.add_argument('--combo-arity', default='triple', choices=['triple', 'mix']) p.add_argument('--arity-weights', default='', help="'w1,w2,w3' for --combo-arity mix (empty = uniform)") + # How many proposals answer each keyword draw. Above 1 they share one prompt + # and one group id, which is what the proposing side needs to have a group to + # compute an advantage over -- at 1 every group has one member and every + # advantage is zero. It does not change the compute at a fixed + # --max-proposals-total; it divides the number of distinct keyword draws by + # the same factor, so 216 proposals come from 27 draws at 8 instead of 216. + p.add_argument('--proposals-per-group', type=int, default=1, + help='proposals sharing one keyword draw and prompt (1 = no groups, ' + 'so no proposer advantage)') # Difficulty filter # 8 attempts, keeping 2-6: with 4 attempts the band was 1-3 and ex9's @@ -661,6 +687,12 @@ def run_check_fn(script: str, slot: int = 0): """Run a python check script in sandbox ``slot``; returns (exit_code, output).""" return runners[slot](script, 'python') + # Why the last snapshot for each slot came back empty: the failure text if the + # listing could not be read, absent if the workspace really was empty. Read by + # snapshot_error_fn below so a paused sandbox is not filed as the model having + # built nothing. + snapshot_errors: Dict[int, str] = {} + def workspace_snapshot_fn(slot: int = 0): """Every file the episode left behind: ``path size`` lines, then contents. @@ -674,22 +706,36 @@ def workspace_snapshot_fn(slot: int = 0): the caller -- there is no end state to write checks about -- and neither may be dressed up as a plausible one: a snapshot that says "empty" when it means "I could not look" produces tasks whose only true - assertion is that nothing happened. + assertion is that nothing happened. Which of the two it was is recorded + in ``snapshot_errors`` instead, for the rejection to be filed under. """ - exit_code, output = runners[slot]( - WORKSPACE_SNAPSHOT.format(workspace=args.workspace, - max_files=args.snapshot_max_files, - per_file=args.snapshot_per_file, - total_budget=args.snapshot_budget), - 'python') + snapshot_errors.pop(slot, None) + script = WORKSPACE_SNAPSHOT.format(workspace=args.workspace, + max_files=args.snapshot_max_files, + per_file=args.snapshot_per_file, + total_budget=args.snapshot_budget) + exit_code, output = runners[slot](script, 'python') if exit_code != 0: - # Not fatal, but not silent either: checks written against a missing - # end state are the failure this whole function exists to prevent. - logger.warning(f'[challenge] workspace snapshot failed (exit {exit_code}): ' - f'{output[-200:]}') + # One retry: 62 of run_clean6's 63 snapshot failures were the sandbox + # answering 410 "not proxyable", which is the host having paused it and + # may be over by the time we ask again. Not fatal either way, but not + # silent: checks written against a missing end state are the failure + # this whole function exists to prevent. + logger.warning(f'[challenge] workspace snapshot failed (exit {exit_code}), ' + f'retrying in {SNAPSHOT_RETRY_WAIT}s: {output[-200:]}') + time.sleep(SNAPSHOT_RETRY_WAIT) + exit_code, output = runners[slot](script, 'python') + if exit_code != 0: + logger.warning(f'[challenge] workspace snapshot failed again (exit ' + f'{exit_code}): {output[-200:]}') + snapshot_errors[slot] = f'workspace snapshot failed (exit {exit_code}): {output[-500:]}' return '' return tool_payload(output).strip() + def snapshot_error_fn(slot: int = 0) -> str: + """Why the last snapshot for ``slot`` was empty; '' if it really was.""" + return snapshot_errors.get(slot, '') + # Arm B. Read at most this much per call: the executor truncates its output # near 8 KB, and base64 grows 3 bytes into 4, so 4 KB of file is about 5.5 KB # of text with room left for the JSON envelope. @@ -793,6 +839,7 @@ def _keyword_gen(record): reset_fn=reset_fn, run_check_fn=run_check_fn, workspace_snapshot_fn=workspace_snapshot_fn, + snapshot_error_fn=snapshot_error_fn, # The executor's own schemas, so the rounds that may call tools advertise # exactly what will run -- same source as the training script uses. tool_schemas=schemas, @@ -802,6 +849,7 @@ def _keyword_gen(record): arity_weights=[float(x) for x in args.arity_weights.split(',')] if args.arity_weights else None, single_kw_prob=args.single_kw_prob, + proposals_per_group=max(1, args.proposals_per_group), keyword_refill_target=args.keywords_n, keyword_gen_calls=args.keyword_gen_calls, keyword_refill_concurrency=max(1, args.keyword_refill_concurrency), @@ -1031,8 +1079,15 @@ def write(self, record): if labels: arrays[f'r{i}_labels'] = np.asarray(labels, dtype=np.int32) if logprobs: + # float64, not float32: these are the ``old_logps`` a GRPO step + # divides by, and the sampler hands them over as full-precision + # python floats (the solver-side json dump keeps all 17 digits, + # e.g. -0.4740769863128662). float32 would round them to about 7 + # digits, so the ratio exp(logp - old_logp) would be off by + # roughly 1e-7 for reasons that have nothing to do with the + # policy having changed. arrays[f'r{i}_logprobs'] = np.asarray( - [lp[0][1] for lp in logprobs], dtype=np.float32) + [lp[0][1] for lp in logprobs], dtype=np.float64) meta.append({ 'stage': rnd.get('stage'), 'messages': rnd.get('messages') or [], diff --git a/cookbook/rsi/agentic/rl.py b/cookbook/rsi/agentic/rl.py index 49de52a35..c53b6797a 100644 --- a/cookbook/rsi/agentic/rl.py +++ b/cookbook/rsi/agentic/rl.py @@ -247,6 +247,14 @@ def main(): f'(need >= {MODEL_GPUS}); skipping batch') continue + # One optimizer step per batch, not per mini-batch. ``forward_backward`` + # neither steps nor zeroes, so the mini-batches below simply add their + # gradients together; ``clip_grad_and_step`` afterwards divides by the + # token count accumulated across all of them, so every trajectory in the + # batch carries the same weight regardless of how the mini-batches split. + # Stepping inside the loop instead -- which is what this used to do -- + # made each step see only MINI_BATCH_SIZE trajectories, so a group of + # NUM_GENERATIONS could be torn across two updates. for mb_start in range(0, len(inputs), MINI_BATCH_SIZE): mb_end = min(mb_start + MINI_BATCH_SIZE, len(inputs)) model.forward_backward( @@ -255,12 +263,10 @@ def main(): advantages=kept_adv[mb_start:mb_end], 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'rsi-agentic-checkpoint-{optim_step}') + model.clip_grad_and_step() + optim_step += 1 + if optim_step % SAVE_STEPS == 0: + model.save(f'rsi-agentic-checkpoint-{optim_step}') log_dict = metrics.calculate() log_dict.update(model.calculate_metric(is_training=True)) diff --git a/cookbook/rsi/agentic/train_offline.py b/cookbook/rsi/agentic/train_offline.py new file mode 100644 index 000000000..183703ba1 --- /dev/null +++ b/cookbook/rsi/agentic/train_offline.py @@ -0,0 +1,517 @@ +"""Offline GRPO on both sides of a challenge run's dump: full-parameter. + +``challenge.py`` already generated the problems and solved each one eight times, +and it landed the token ids, the labels and the sampler's logprobs for every one +of those trajectories. This trains on that dump directly -- nothing is generated +here and nothing is re-encoded, so the tokens trained on are byte-for-byte the +tokens that were sampled. + +Two sides come out of one run: + +* proposing -- one trajectory per proposal, grouped by ``group_id`` (the + proposals answering one identical prompt). Reward is ``challenger_reward``: + ``1 - 2|p - 1/2|`` for p the fraction of solver attempts that passed, so a + proposal is worth most when the solver got it right about half the time. +* solving -- eight trajectories per task, grouped by task. Reward is the check + script's exit code, 1.0 or 0.0. + +Both sides go into the same optimizer step and are weighted only by how many +trainable tokens they carry; no coefficient is applied to either. + +Usage, after a run of challenge.py --proposals-per-group 8: + + python cookbook/rsi/agentic/train_offline.py \\ + --run-dir output/rsi_agentic/run_clean10 \\ + --model-id ms://Qwen/Qwen3-4B --model-gpus 8 + +The saved checkpoint is HF-format weights plus tokenizer, which is what +``challenge.py --model-id`` takes, so the next round of the loop is a shell line +rather than a conversion step. +""" +import collections +import json +import os +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_device_placement, get_logger +from twinkle.advantage import GRPOAdvantage +from twinkle.cli import CLI +from twinkle.processor import InputProcessor + +logger = get_logger() +args = CLI.from_args() + +# ========== Configuration ========== +MODEL_ID = args.model.model_id or 'ms://Qwen/Qwen3-4B' +MODEL_GPUS = args.infra.model_gpus or 8 +# Base text template, as in cookbook/rsi/rl.py:102. Qwen3-4B is text-only, and +# the multimodal subclass crashes on encode for it. +TEMPLATE = os.environ.get('RSI_TEMPLATE', 'Template') + +# Whatever --lr says, or the CLI's own 1e-5. Not written as ``or <number>``: the +# CLI default is never zero, so a fallback here would be dead code that reads +# like the default. +LEARNING_RATE = args.optimizer.learning_rate +MINI_BATCH_SIZE = args.training.mini_batch_size or 8 +MICRO_BATCH_SIZE = args.training.micro_batch_size or 2 +SAVE_STEPS = args.training.save_steps or 0 + +RUN_DIR = os.environ.get('RSI_RUN_DIR', '') +SAVE_DIR = os.environ.get('RSI_SAVE_DIR', 'output/rsi_agentic/ckpt') +SAVE_NAME = os.environ.get('RSI_SAVE_NAME', 'agentic-offline') + +# Trajectories per optimizer step. 32 is BATCH_SIZE 4 x NUM_GENERATIONS 8, the +# same as the online rl.py loop, so a step here moves the weights by as much as a +# step there. +STEP_SIZE = int(os.environ.get('RSI_STEP_SIZE', 32)) +# Longest trajectory fed to the model. Above this the model refuses it mid-step. +MAX_MODEL_LEN = int(os.environ.get('RSI_MAX_MODEL_LEN', 32768)) + +# Which side(s) to train. 'both', 'solver', 'proposer'. +SIDES = os.environ.get('RSI_SIDES', 'both') + +# Cap on proposing-side groups per run, 0 for no cap. The solving side is held +# constant by challenge.py's --keep-target, but the proposing side is however +# many proposals it took to reach that target, which grows as the model improves +# and fewer of its proposals land in the band. Capping keeps the two sides' share +# of each step the same from run to run; the proposals above the cap still did +# their job of measuring difficulty, they just do not also become training data. +MAX_PROPOSER_GROUPS = int(os.environ.get('RSI_MAX_PROPOSER_GROUPS', 0)) + +# Same for the solving side, 0 for no cap. ``--keep-target`` stops challenge.py +# once it has that many tasks in the band, but the internal round that crosses +# the target finishes measuring everything it started -- run_clean9's last round +# added 14 tasks to reach 53 from 39 -- so a run overshoots by however much that +# round produced. Capping makes every run's solving side exactly the same size. +MAX_SOLVER_GROUPS = int(os.environ.get('RSI_MAX_SOLVER_GROUPS', 0)) + +# Where the numbers go. Two files under the run directory: +# train_summary.json one object: the settings this run used, what the +# challenger collected, and what got trained on +# train_steps.jsonl one line per optimizer step, every metric the model +# reported plus the batch's own composition +# Written rather than uploaded, and written as they happen rather than at the end, +# so a run that dies partway still leaves the steps it did finish. +SUMMARY_NAME = os.environ.get('RSI_SUMMARY_NAME', 'train_summary.json') +STEPS_NAME = os.environ.get('RSI_STEPS_NAME', 'train_steps.jsonl') + +# Solver groups outside this pass-count range carry one reward for all eight +# members, so their advantages are all zero and a forward+backward over them adds +# nothing. The bounds exclude only 0 and n_rollouts, which is what makes this +# 'the groups that have a gradient' rather than a difficulty judgement. +KEEP_MIN_PASS = int(os.environ.get('RSI_KEEP_MIN_PASS', 1)) +KEEP_MAX_PASS_MARGIN = int(os.environ.get('RSI_KEEP_MAX_MARGIN', 1)) + + +def _cap(groups: List[Dict[str, Any]], limit: int, side: str, + notes: collections.Counter) -> List[Dict[str, Any]]: + """Hold a side's group count to ``limit`` so every run trains on the same amount. + + Kept in file order, which is the order they came out of the challenger. The + weights do not change inside one challenge.py run, so the groups dropped are + not different in kind from the ones kept -- they are just later. + """ + if not limit: + return groups + if len(groups) > limit: + notes[f'{side} groups over the cap of {limit}'] += len(groups) - limit + return groups[:limit] + if len(groups) < limit: + # Worth saying out loud: this run trains on less than the cap promises, so + # its steps are not mixed the same way as a run that reached it. + logger.warning(f'[offline] only {len(groups)} {side} groups, below the cap ' + f'of {limit}; this run is not comparable to one that ' + f'reached the cap') + return groups + + +def _collection_metrics(run_dir: str) -> Dict[str, float]: + """Summarise what the challenger produced, before any of the training filters. + + Read off the two dumps rather than from anything challenge.py logs, so these + numbers describe the whole collection and not the subset that survived the + pass band and the caps. + + Two pass rates are reported and they mean different things. ``acc/all`` is + over every attempt on every task whose difficulty got measured, including the + tasks where all eight attempts agreed. ``acc/trained`` is over the attempts on + the tasks that go into training, and is bounded to [1/8, 7/8] by that band -- + it cannot report an all-correct or all-wrong task even if the model produces + one. Neither is a capability measurement on its own: the tasks change every + iteration and the challenger is being trained to push them toward half + passing, so a flat curve is what success looks like for both. + """ + m: Dict[str, float] = {} + idx = os.path.join(run_dir, 'propose_traj', 'index.jsonl') + if os.path.isfile(idx): + rows = [json.loads(line) for line in open(idx) if line.strip()] + outcomes = collections.Counter(r.get('outcome') for r in rows) + m['collect/proposals'] = len(rows) + for name, n in outcomes.items(): + m[f'collect/outcome_{name}'] = n + rewards = [r['challenger_reward'] for r in rows if r.get('challenger_reward') is not None] + if rewards: + m['collect/proposer_reward_mean'] = sum(rewards) / len(rewards) + + measured = [r['n_pass'] for r in rows + if r.get('n_pass') is not None and r.get('n_rollouts')] + rollouts = [r['n_rollouts'] for r in rows + if r.get('n_pass') is not None and r.get('n_rollouts')] + if measured: + n_att = sum(rollouts) + m['collect/measured_tasks'] = len(measured) + m['collect/never_measured'] = len(rows) - len(measured) + m['acc/all'] = sum(measured) / n_att + dist = collections.Counter(measured) + for k in range(max(rollouts) + 1): + m[f'collect/n_pass_{k}'] = dist.get(k, 0) + # All eight attempts agreeing means one reward for the whole group, so + # the group mean equals it and every advantage is zero. Tracking the + # share of tasks like that is tracking how much of the collection did + # no work. + flat = sum(1 for p, r in zip(measured, rollouts) if p == 0 or p == r) + m['collect/zero_gradient_frac'] = flat / len(measured) + band = [(p, r) for p, r in zip(measured, rollouts) + if KEEP_MIN_PASS <= p <= r - KEEP_MAX_PASS_MARGIN] + if band: + m['acc/trained'] = sum(p for p, _ in band) / sum(r for _, r in band) + m['collect/band_tasks'] = len(band) + # How many proposals it costs to land one trainable task. Expected + # to climb as the model gets better at its own proposals, which is + # what makes an iteration take longer than the one before it. + m['collect/proposals_per_band_task'] = len(rows) / len(band) + + attempts = os.path.join(run_dir, 'solver_attempts.jsonl') + if os.path.isfile(attempts): + rows = [json.loads(line) for line in open(attempts) if line.strip()] + if rows: + m['collect/solver_attempts'] = len(rows) + m['collect/solver_truncated_frac'] = \ + sum(1 for r in rows if r.get('truncated')) / len(rows) + return m + + +def _load_solver_groups(run_dir: str) -> Tuple[List[Dict[str, Any]], collections.Counter]: + """Groups of solver attempts on one task, with their exit-code rewards.""" + path = os.path.join(run_dir, 'solver_attempts.jsonl') + if not os.path.exists(path): + return [], collections.Counter({'no solver_attempts.jsonl': 1}) + by_task: Dict[str, List[Dict[str, Any]]] = collections.OrderedDict() + notes: collections.Counter = collections.Counter() + with open(path) as f: + for line in f: + line = line.strip() + if not line: + continue + try: + rec = json.loads(line) + except json.JSONDecodeError: + notes['unparseable line'] += 1 + continue + by_task.setdefault(rec['statement'], []).append(rec) + + groups: List[Dict[str, Any]] = [] + for statement, recs in by_task.items(): + rewards = [1.0 if r['check_exit'] == 0 else 0.0 for r in recs] + n_pass = int(sum(rewards)) + high = len(recs) - KEEP_MAX_PASS_MARGIN + if not (KEEP_MIN_PASS <= n_pass <= high): + notes[f'solver n_pass={n_pass} of {len(recs)} (no gradient)'] += 1 + continue + members = [] + for rec, reward in zip(recs, rewards): + att = rec.get('attempt') or {} + members.append({ + 'input_ids': att.get('input_ids') or [], + 'labels': att.get('labels') or [], + 'attention_mask': att.get('attention_mask') or [], + 'position_ids': att.get('position_ids') or [], + # The sampler's own logprobs, in the [[token, logp]] shape the + # rollout stored them in. Reusing them rather than a fresh + # forward is what keeps old_logps free of engine differences. + 'logps': [lp[0][1] for lp in (att.get('logprobs') or [])], + 'reward': reward, + }) + groups.append({'side': 'solver', 'key': statement[:60], 'members': members}) + return _cap(groups, MAX_SOLVER_GROUPS, 'solver', notes), notes + + +def _load_proposer_groups(run_dir: str) -> Tuple[List[Dict[str, Any]], collections.Counter]: + """Groups of proposals answering one prompt, with their 50%-target rewards.""" + d = os.path.join(run_dir, 'propose_traj') + index = os.path.join(d, 'index.jsonl') + notes: collections.Counter = collections.Counter() + if not os.path.exists(index): + return [], collections.Counter({'no propose_traj/index.jsonl': 1}) + + by_group: Dict[Any, List[Dict[str, Any]]] = collections.OrderedDict() + with open(index) as f: + for line in f: + line = line.strip() + if not line: + continue + try: + rec = json.loads(line) + except json.JSONDecodeError: + notes['unparseable line'] += 1 + continue + gid = rec.get('group_id') + if gid is None: + # Landed by a run from before proposals were grouped. Every such + # proposal is its own group of one, so its advantage would be + # zero; counted and skipped rather than trained on as noise. + notes['proposal has no group_id (pre-grouping run)'] += 1 + continue + by_group.setdefault(gid, []).append(rec) + + groups: List[Dict[str, Any]] = [] + for gid, recs in by_group.items(): + if len(recs) < 2: + notes[f'proposer group of {len(recs)} (no gradient)'] += 1 + continue + members = [] + for rec in recs: + npz_name = rec.get('npz') + if not npz_name: + notes['proposal has no npz (text-only rollout)'] += 1 + continue + z = np.load(os.path.join(d, npz_name)) + if 'r0_input_ids' not in z.files: + notes['npz has no r0_input_ids'] += 1 + continue + ids = z['r0_input_ids'].astype(np.int64) + # Labels are stored already shifted by one -- labels[i] is the token + # at input_ids[i+1] -- which is the convention the sampler wrote them + # in. Passed through untouched; re-deriving them here would be + # guessing at an alignment that is already correct on disk. + labels = z['r0_labels'].astype(np.int64) + logps = z['r0_logprobs'].astype(np.float64) if 'r0_logprobs' in z.files else None + if logps is None: + notes['npz has no r0_logprobs'] += 1 + continue + members.append({ + 'input_ids': ids.tolist(), + 'labels': labels.tolist(), + 'attention_mask': [1] * len(ids), + 'position_ids': list(range(len(ids))), + 'logps': logps.tolist(), + 'reward': float(rec.get('challenger_reward') or 0.0), + }) + if len(members) < 2: + notes['proposer group lost members to missing arrays'] += 1 + continue + groups.append({'side': 'proposer', 'key': f'group{gid}', 'members': members}) + return _cap(groups, MAX_PROPOSER_GROUPS, 'proposer', notes), notes + + +def _score(groups: List[Dict[str, Any]]) -> collections.Counter: + """Advantage per member, in place. Groups may differ in size.""" + advantage_fn = GRPOAdvantage() + notes: collections.Counter = collections.Counter() + for g in groups: + rewards = [m['reward'] for m in g['members']] + adv = advantage_fn(rewards, num_generations=len(rewards), scale='group').tolist() + if all(abs(a) < 1e-9 for a in adv): + notes[f'{g["side"]}: group with no gradient after scoring'] += 1 + for m, a in zip(g['members'], adv): + m['advantage'] = a + return notes + + +def _interleave(by_side: Dict[str, List[Dict[str, Any]]]) -> List[Dict[str, Any]]: + """Spread each side's groups evenly through the order. + + Both sides have to appear in every step, or a step's update comes from one + side only and 'weighted by token count' stops describing anything. The ratio + is not a free choice: it falls out of using all of both sides' groups over + the same number of steps. + """ + sides = [s for s in ('solver', 'proposer') if by_side.get(s)] + if len(sides) < 2: + return list(by_side.get(sides[0], [])) if sides else [] + # Place each group at its fractional position within its own side, then sort + # by that position: a side with three times the groups contributes three for + # every one of the other's, spread out rather than in a block. + marked: List[Tuple[float, Dict[str, Any]]] = [] + for s in sides: + gs = by_side[s] + for i, g in enumerate(gs): + marked.append(((i + 0.5) / len(gs), g)) + marked.sort(key=lambda t: t[0]) + return [g for _pos, g in marked] + + +def main(): + if not RUN_DIR: + raise SystemExit('set RSI_RUN_DIR to a challenge.py output directory') + + collect = _collection_metrics(RUN_DIR) + + by_side: Dict[str, List[Dict[str, Any]]] = {} + all_notes: collections.Counter = collections.Counter() + if SIDES in ('both', 'solver'): + gs, notes = _load_solver_groups(RUN_DIR) + by_side['solver'] = gs + all_notes.update(notes) + if SIDES in ('both', 'proposer'): + gs, notes = _load_proposer_groups(RUN_DIR) + by_side['proposer'] = gs + all_notes.update(notes) + + all_notes.update(_score([g for gs in by_side.values() for g in gs])) + + for side, gs in by_side.items(): + n_traj = sum(len(g['members']) for g in gs) + sizes = collections.Counter(len(g['members']) for g in gs) + logger.info(f'[offline] {side}: {len(gs)} groups, {n_traj} trajectories, ' + f'group sizes {dict(sorted(sizes.items()))}') + for note, n in all_notes.most_common(): + logger.info(f'[offline] skipped {n}: {note}') + if not any(by_side.values()): + raise SystemExit(f'[offline] nothing trainable in {RUN_DIR}') + + order = _interleave(by_side) + flat: List[Dict[str, Any]] = [] + for g in order: + for m in g['members']: + m['side'] = g['side'] + flat.append(m) + + # Trajectories the model would refuse, dropped before anything is loaded so + # the count is in the log rather than showing up as a mid-step exception. + kept: List[Dict[str, Any]] = [] + for m in flat: + n_train = sum(1 for label in m['labels'] if label != -100) + if not n_train: + all_notes['no trainable tokens'] += 1 + continue + if len(m['logps']) != n_train: + # Off-by-anything here pairs each logprob with the wrong token and + # the loss still comes out a number, so it has to be a hard stop. + all_notes[f'logps {len(m["logps"])} != trainable {n_train}'] += 1 + continue + if len(m['input_ids']) > MAX_MODEL_LEN: + all_notes[f'longer than MAX_MODEL_LEN={MAX_MODEL_LEN}'] += 1 + continue + kept.append(m) + + n_steps = (len(kept) + STEP_SIZE - 1) // STEP_SIZE + mix = collections.Counter(m['side'] for m in kept) + logger.info(f'[offline] {len(kept)} trainable trajectories ' + f'({dict(mix)}), {n_steps} steps of {STEP_SIZE}') + for note, n in all_notes.most_common(): + logger.info(f'[offline] skipped {n}: {note}') + + summary = { + 'run_dir': RUN_DIR, + 'config': {'model_id': MODEL_ID, 'lr': LEARNING_RATE, 'step_size': STEP_SIZE, + 'sides': SIDES, 'keep_min_pass': KEEP_MIN_PASS, + 'keep_max_pass_margin': KEEP_MAX_PASS_MARGIN, + 'max_solver_groups': MAX_SOLVER_GROUPS, + 'max_proposer_groups': MAX_PROPOSER_GROUPS, + 'mini_batch_size': MINI_BATCH_SIZE, + 'micro_batch_size': MICRO_BATCH_SIZE, + 'max_model_len': MAX_MODEL_LEN, 'template': TEMPLATE, + 'model_gpus': MODEL_GPUS}, + 'collect': collect, + 'train': {'trainable_trajectories': len(kept), 'steps': n_steps, + **{f'{side}_trajectories': n for side, n in mix.items()}, + 'groups_per_side': {side: len(gs) for side, gs in by_side.items()}}, + # Every reason anything was left out, with its count. Kept in the summary + # rather than only in the log so a later comparison between iterations can + # tell a change in the model from a change in how much survived the load. + 'skipped': dict(all_notes), + } + summary_path = os.path.join(RUN_DIR, SUMMARY_NAME) + with open(summary_path, 'w') as f: + json.dump(summary, f, indent=2, ensure_ascii=False) + steps_path = os.path.join(RUN_DIR, STEPS_NAME) + steps_file = open(steps_path, 'w') + logger.info(f'[offline] numbers going to {summary_path} and {steps_path}') + for k in sorted(collect): + logger.info(f'[offline] {k} = {collect[k]}') + + device_groups = [DeviceGroup(name='model', ranks=list(range(MODEL_GPUS)), device_type='GPU')] + model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, dp_size=MODEL_GPUS) + twinkle.initialize(mode='ray', nproc_per_node=MODEL_GPUS, groups=device_groups, + lazy_collect=False) + + # Full-parameter: no adapter is added, so every weight is trained and the + # checkpoint is a whole model rather than something that needs merging before + # the next challenger round can load it. + from twinkle.model.megatron import MegatronModel + model = MegatronModel(model_id=MODEL_ID, device_mesh=model_mesh, remote_group='model', + mixed_precision='bf16', variable_seq_lengths=True) + model.set_optimizer('default', lr=LEARNING_RATE) + model.set_lr_scheduler('default', lr_decay_steps=max(1, n_steps), max_lr=LEARNING_RATE) + # beta=0: no reference model here, and grpo.py:315 needs beta>0 AND ref_logps + # for the KL term, so any beta above 0 would silently do nothing. + model.set_loss('GRPOLoss', epsilon=0.2, beta=0.0) + model.set_processor(InputProcessor, padding_free=True) + model.set_template(TEMPLATE, model_id=MODEL_ID, max_length=MAX_MODEL_LEN, + enable_thinking=True) + # approx_kl at the first inner step compares the sampler's logps against the + # trainer's on the same tokens. On landed data that is the check for whether + # this dump belongs to the weights being trained: it should start near zero, + # and a large value means the dump came from a different checkpoint. + model.add_metric('GRPOMetric', is_training=True, epsilon=0.2) + logger.info(get_device_placement()) + + for step in range(n_steps): + lo, hi = step * STEP_SIZE, min((step + 1) * STEP_SIZE, len(kept)) + batch = kept[lo:hi] + inputs = [{'input_ids': m['input_ids'], 'labels': m['labels'], + 'attention_mask': m['attention_mask'], + 'position_ids': m['position_ids']} for m in batch] + old_logps = [m['logps'] for m in batch] + advantages = [m['advantage'] for m in batch] + + for mb in range(0, len(inputs), MINI_BATCH_SIZE): + end = min(mb + MINI_BATCH_SIZE, len(inputs)) + model.forward_backward( + inputs=inputs[mb:end], + old_logps=old_logps[mb:end], + advantages=advantages[mb:end], + micro_batch_size=MICRO_BATCH_SIZE, + ) + # Once per step, not per mini-batch: forward_backward neither steps nor + # zeroes, and clip_grad_norm divides by the tokens accumulated across all + # of them, so every trajectory in the step carries the same weight no + # matter how the mini-batches split -- which is also how the two sides end + # up weighted by their token counts and nothing else. + model.clip_grad_and_step() + + side_mix = collections.Counter(m['side'] for m in batch) + log = model.calculate_metric(is_training=True) + # A list of records rather than a number, marked with a leading underscore + # at grpo.py:451 for that reason. Going to a file, so the records go in + # whole instead of being reduced to a count. + high_kl = log.pop('_high_kl_records', None) + logger.info(f'[offline] step {step + 1}/{n_steps} ' + f'{len(batch)} traj {dict(side_mix)} ' + f'adv[{min(advantages):+.3f},{max(advantages):+.3f}] {log}') + if high_kl: + logger.warning(f'[offline] step {step + 1}: {len(high_kl)} sequences ' + f'with high kl against the sampler logps') + row = {'step': step + 1, **log, + 'trajectories': len(batch), + 'solver': side_mix.get('solver', 0), + 'proposer': side_mix.get('proposer', 0), + 'adv_min': min(advantages), 'adv_max': max(advantages), + 'high_kl_records': high_kl or []} + steps_file.write(json.dumps(row, ensure_ascii=False, default=str) + '\n') + steps_file.flush() + if SAVE_STEPS and (step + 1) % SAVE_STEPS == 0: + model.save(f'{SAVE_NAME}-step{step + 1}', output_dir=SAVE_DIR) + + steps_file.close() + model.save(SAVE_NAME, output_dir=SAVE_DIR) + logger.info(f'[offline] done, {n_steps} steps; checkpoint at ' + f'{os.path.join(SAVE_DIR, SAVE_NAME)}') + + +if __name__ == '__main__': + main() diff --git a/src/twinkle/infra/__init__.py b/src/twinkle/infra/__init__.py index a3c90eae3..b9b169e44 100644 --- a/src/twinkle/infra/__init__.py +++ b/src/twinkle/infra/__init__.py @@ -5,6 +5,7 @@ import numpy as np import os import sys +import threading from typing import Any, Callable, List, Literal, Optional, TypeVar, Union from twinkle.notifier import Notifier, notify_exception @@ -300,6 +301,255 @@ def _get_workers(workers, execute): raise ValueError(f'Unsupported execute method: {execute}') +# Guards creating the per-handle state below. Without it two threads arriving at +# once each build their own state, with their own lock, and one overwrites the +# other -- after which the two are no longer excluding each other and the requests +# charged to the discarded one are never given back. +_CW_CREATE_LOCK = threading.Lock() + + +# Prefix of the awaitable companion generated for a continuous-work method. The +# companion is what the driver actually calls on the worker; see +# ``_make_worker_async_companion``. +_WORKER_ASYNC_PREFIX = '_twinkle_async_' + + +def _worker_executor(self, ): + """Threads for running a blocking worker method off the actor's event loop. + + Sized from ``TWINKLE_ACTOR_MAX_CONCURRENCY``, which ``create_workers`` sets to + the actor's ``max_concurrency``: any fewer threads than that would throttle + below the concurrency the actor was configured for. A private executor rather + than the loop's default one, so this never changes behaviour for anything else + running on that loop. + """ + executor = getattr(self, '_twinkle_worker_executor', None) + if executor is not None: + return executor + with _CW_CREATE_LOCK: + executor = getattr(self, '_twinkle_worker_executor', None) + if executor is None: + from concurrent.futures import ThreadPoolExecutor + n = int(os.environ.get('TWINKLE_ACTOR_MAX_CONCURRENCY') or 0) or 1 + executor = ThreadPoolExecutor(max_workers=n, thread_name_prefix='twinkle-worker') + self._twinkle_worker_executor = executor + return executor + + +def _make_worker_async_companion(func, wrapper): + """Wrap a blocking worker method so the actor can run several of them at once. + + Ray makes a class with any ``async def`` into an asyncio actor, and there a + blocking method holds the actor's single event loop for its whole duration -- + so calls queue and run one after another however high ``max_concurrency`` is. + Measured on this sampler: four concurrent one-prompt calls took 3.98x as long + as one, while the same four prompts in a single call took 1.02x. Handing the + blocking body to a thread leaves the loop free to accept the next call, which + is what puts several requests in the worker's engine together. + """ + import asyncio + + @functools.wraps(func) + async def companion(self, *args, **kwargs): + loop = asyncio.get_running_loop() + return await loop.run_in_executor(_worker_executor(self), functools.partial(wrapper, self, *args, **kwargs)) + + companion.__name__ = _WORKER_ASYNC_PREFIX + func.__name__ + return companion + + +def _cw_state(self, n_workers: int): + """Driver-side bookkeeping for ``enable_continous_work``, created on first use. + + ``load`` counts requests handed to each worker and not yet returned, which is + what picks the next worker. ``inflight`` keeps those counts honest per method + name, and is also what the barrier reads: a method other than the one with + requests in flight must wait for them, because the worker now runs methods + side by side and something like receiving weights or sleeping would otherwise + land on an engine mid-generation. + + The state's own lock is a plain ``Lock``: nothing here takes it while already + holding it, so re-entrance is not needed. + """ + state = getattr(self, '_continous_work_state', None) + if state is not None and len(state['load']) == n_workers: + return state + with _CW_CREATE_LOCK: + # Re-read: another thread may have created it while this one waited. + state = getattr(self, '_continous_work_state', None) + if state is None or len(state['load']) != n_workers: + state = { + 'lock': threading.Lock(), + 'load': [0] * n_workers, + 'inflight': {}, # method name -> list of pending object refs + } + self._continous_work_state = state + return state + + +def _cw_barrier(self, current_func: str) -> None: + """Drain every other method's in-flight work before proceeding. + + Best effort by construction: another thread may submit again the moment this + returns. It removes the case this exists for -- a weight update or a sleep + issued while generations are still running -- but it is not a global lock on + the worker. + """ + state = getattr(self, '_continous_work_state', None) + if not state: + return + with state['lock']: + others = {name: list(refs) for name, refs in state['inflight'].items() if name != current_func and refs} + if not others: + return + import ray + flat = [ref for refs in others.values() for ref in refs] + logger.debug(f'continous_work barrier: {current_func} waits for {len(flat)} pending request(s) ' + f'from {sorted(others)}') + ray.get(flat) + # They are finished now, so drop them instead of re-getting them on every + # later call. The owning thread's own cleanup tolerates them being gone. + with state['lock']: + for name, refs in others.items(): + pending = state['inflight'].get(name) + if pending is None: + continue + for ref in refs: + if ref in pending: + pending.remove(ref) + if not pending: + state['inflight'].pop(name, None) + + +def _cw_object_refs(result) -> List[Any]: + """Every ObjectRef inside a dispatch result, tuples included.""" + import ray + refs = [] + for item in (result or []): + for candidate in (item if isinstance(item, tuple) else (item, )): + if isinstance(candidate, ray.ObjectRef): + refs.append(candidate) + return refs + + +def _cw_register(self, func_name: str, result) -> List[Any]: + """Record a non-continuous call's refs so a later different method waits for it. + + Needed because a lazily collected method returns before its work finishes: + ``receive_weights`` hands back a handle while the worker is still swapping + weights, and with actor concurrency on, a sample issued right after would read + them half written. + """ + refs = _cw_object_refs(result) + if not refs: + return refs + state = _cw_state(self, len(getattr(self, '_actors', ())) or 1) + with state['lock']: + state['inflight'].setdefault(func_name, []).extend(refs) + return refs + + +def _cw_unregister(self, func_name: str, refs: List[Any]) -> None: + state = getattr(self, '_continous_work_state', None) + if not state or not refs: + return + with state['lock']: + pending = state['inflight'].get(func_name) + if pending is None: + return + for ref in refs: + if ref in pending: + pending.remove(ref) + if not pending: + state['inflight'].pop(func_name, None) + + +def _cw_plan(n_workers: int, load: List[int], batch_len: int) -> List[List[int]]: + """Assign each request to the worker holding the fewest, updating ``load``. + + Least-loaded-first, one request at a time, so a call of one request goes to + one worker instead of being padded up to the worker count, and a call of many + spreads out. ``load`` is mutated by the caller's lock holder. + """ + per_worker: List[List[int]] = [[] for _ in range(n_workers)] + for idx in range(batch_len): + target = min(range(n_workers), key=lambda w: load[w]) + per_worker[target].append(idx) + load[target] += 1 + return per_worker + + +def _cw_batch_len(args, kwargs) -> Optional[int]: + """Length of the request list, i.e. the first list argument's length. + + Same convention as ``dispatch='slice'``: list arguments are the batch and + everything else is broadcast. Returns None when there is no list to split, + which is how the caller knows to fall back to the normal dispatch. + """ + for arg in list(args) + list(kwargs.values()): + if isinstance(arg, list): + return len(arg) + return None + + +def _cw_sub_args(args, kwargs, indices: List[int], batch_len: int): + """The arguments for one worker: list arguments indexed, the rest as-is.""" + + def pick(arg): + if isinstance(arg, list) and len(arg) == batch_len: + return [arg[i] for i in indices] + return arg + + return tuple(pick(a) for a in args), {k: pick(v) for k, v in kwargs.items()} + + +def _run_continous_work(self, func_name: str, execute_method, workers, args, kwargs, batch_len: int, + ray_get_timeout: Optional[float]): + """Submit one call per chosen worker and return results in the caller's order. + + Submission happens under the lock so that picking a worker and charging it are + one step -- several caller threads land here at once, and a split of the two + would let them all pick the same idle worker. Waiting happens outside it. + """ + import ray + + state = _cw_state(self, len(workers)) + submitted = [] + # The awaitable form, so the worker can hold several of these at once. Book- + # keeping still uses the plain name, which is what callers and the barrier see. + remote_name = _WORKER_ASYNC_PREFIX + func_name + with state['lock']: + plan = _cw_plan(len(workers), state['load'], batch_len) + for worker_index, indices in enumerate(plan): + if not indices: + continue + sub_args, sub_kwargs = _cw_sub_args(args, kwargs, indices, batch_len) + ref = execute_method(remote_name, [(workers[worker_index], sub_args, sub_kwargs)])[0] + submitted.append((worker_index, indices, ref)) + state['inflight'].setdefault(func_name, []).extend(ref for _, _, ref in submitted) + + try: + ordered: List[Any] = [None] * batch_len + for _, indices, ref in submitted: + part = ray.get(ref, timeout=ray_get_timeout) if ray_get_timeout else ray.get(ref) + if not isinstance(part, (list, tuple)) or len(part) != len(indices): + raise TypeError(f'{func_name}: enable_continous_work needs one result per request, but a worker given ' + f'{len(indices)} request(s) returned {type(part).__name__} of length ' + f'{len(part) if isinstance(part, (list, tuple)) else "n/a"}.') + for local_index, original_index in enumerate(indices): + ordered[original_index] = part[local_index] + return ordered + finally: + with state['lock']: + pending = state['inflight'].get(func_name, []) + for worker_index, indices, ref in submitted: + state['load'][worker_index] -= len(indices) + if ref in pending: + pending.remove(ref) + if not pending: + state['inflight'].pop(func_name, None) + + def _collect_func(method: Union[Literal['none', 'flatten', 'mean', 'sum', 'first', 'last_pp'], Callable], result: List[Any], device_mesh: DeviceMesh = None): @@ -504,14 +754,33 @@ def _prepare_lazy_collect(args, kwargs): return args, kwargs -def remote_class(execute: Literal['first', 'peer', 'all'] = 'all'): +def remote_class(execute: Literal['first', 'peer', 'all'] = 'all', + max_concurrency: Optional[int] = None): """Patch each class used in remote clusters with this decorator. Use this decorator to wrap your class to enable it to execute in a remote cluster. + Args: + execute: which workers the class runs on. + max_concurrency: Ray actor concurrency, i.e. how many of this class's + methods one worker may run at once. ``None`` leaves Ray's default of + 1, under which concurrent calls to the same worker queue and run one + after another. Only set it for a class whose methods tolerate running + side by side: a class holding NCCL collectives does not, because two + collectives interleaving on one rank deadlock. It is what + ``enable_continous_work`` needs to reach the worker's engine + concurrently instead of stopping at the actor boundary. """ def decorator(cls): + # Give every continuous-work method its awaitable form on the class, so Ray + # has something to await instead of a call that would sit on the actor's + # event loop and make the others wait behind it. + for _name in dir(cls): + _attr = getattr(cls, _name, None) + _companion = getattr(_attr, '_worker_async_companion', None) + if _companion is not None: + setattr(cls, _WORKER_ASYNC_PREFIX + _name, _companion) # Get device mesh parameter name device_mesh_name = _get_device_mesh_param_name(cls.__init__) init_method = cls.__init__ @@ -667,9 +936,15 @@ def __next__(_self): instance_id=instance_id, seed=_seed, full_determinism=_full_determinism, + max_concurrency=max_concurrency, *args, **kwargs_for_workers) self._actors = _actors + # Remembered so remote_function knows this class's workers run + # methods side by side, and that it must therefore track what is + # in flight. Without concurrency Ray orders calls per actor and + # the tracking would be dead weight. + self._max_concurrency = max_concurrency if hasattr(cls, '__iter__'): # wraps again, because ray uses cls method to call remote cls.__iter__ = remote_function(dispatch=_dispatch, execute=_execute, collect='none')(__iter__) @@ -696,7 +971,8 @@ def remote_function(dispatch: Union[Literal['slice', 'all', 'slice_dp', 'last_pp collect: Union[Literal['none', 'flatten', 'mean', 'sum', 'first', 'last_pp'], Callable] = 'none', sync: bool = False, lazy_collect: Optional[bool] = None, - timeout: Optional[float] = None): + timeout: Optional[float] = None, + enable_continous_work: bool = False): """Patch each method called from remote(which class should be decorated with `remote_class`) with this decorator. Args: @@ -723,6 +999,19 @@ def remote_function(dispatch: Union[Literal['slice', 'all', 'slice_dp', 'last_pp 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. + enable_continous_work: Route each request to the least busy worker instead + of slicing the batch over all of them, and return the results in the + caller's order. This is what lets a batch smaller than the worker + count through: ``slice_dp`` would hand some ranks nothing and raise, + which is why callers pad a single request up to the worker count and + throw the duplicate generations away. Requires the class to be + declared with ``max_concurrency`` above 1, otherwise the requests + queue at the actor and run one at a time instead of reaching the + worker's engine together. Only for methods that take a list of + independent requests and return one result each, and whose workers + need no collective between them -- data-parallel sampling, not a + method with an all-reduce in it. While one such method has requests in + flight, calling any other method on the same handle waits for them. """ # noqa def decorator(func: Callable[..., T1]) -> Callable[..., T1]: @@ -758,6 +1047,23 @@ def wrapper(self, *args, **kwargs) -> T1: # This is the driver from ._ray import RayHelper execute_method = RayHelper.execute_all_async if not sync else RayHelper.execute_all_sync + # Only classes whose workers run methods side by side need + # this; elsewhere Ray already orders calls per actor. + _concurrent_actor = bool(getattr(self, '_max_concurrency', None)) + if _concurrent_actor: + # Every method waits here, not just the continuous ones: + # the point is to keep a weight update or a sleep from + # reaching a worker that still has generations running. + _cw_barrier(self, func.__name__) + if enable_continous_work and not RayHelper.has_ref(args, kwargs): + assert not sync, (f'{func.__name__}: enable_continous_work cannot be used with sync=True, ' + 'which exists for collectives that must run in lock step.') + _workers = _get_workers(self._actors, execute) + _batch_len = _cw_batch_len(args, kwargs) + if _batch_len: + return _run_continous_work(self, func.__name__, execute_method, _workers, args, kwargs, + _batch_len, + getattr(self, '_ray_get_timeout', None) or timeout) if RayHelper.has_ref(args, kwargs): # If has any object-ref, dispatch in worker, because we don't know the structure in the ref. # for example, dataloader returns any data list. @@ -769,6 +1075,11 @@ def wrapper(self, *args, **kwargs) -> T1: _get_workers(self._actors, execute), dispatch, execute, device_mesh, args, kwargs) result = execute_method(func.__name__, _workers_and_args) + # Tracked from here so that a different method called next + # waits for this one. It matters most for the lazily + # collected methods, which return while the worker is still + # busy. + _tracked_refs = _cw_register(self, func.__name__, result) if _concurrent_actor else [] # This is a result future, call it to get the actual result _rgt = getattr(self, '_ray_get_timeout', None) or timeout result_func = RayHelper.do_get_and_collect_func( @@ -812,12 +1123,17 @@ def _notifying_result_func(*rargs, **rkwargs): _tag_exc(_e, _caller) notify_exception(_notifier, _ctx, _e, _name) raise + finally: + _cw_unregister(self, func.__name__, _tracked_refs) for _attr in ('_futures', ): if hasattr(_orig_result_func, _attr): setattr(_notifying_result_func, _attr, getattr(_orig_result_func, _attr)) return _notifying_result_func - return result_func() + try: + return result_func() + finally: + _cw_unregister(self, func.__name__, _tracked_refs) else: raise NotImplementedError(f'Unsupported mode {_mode}') except StopIteration: @@ -832,6 +1148,11 @@ def _notifying_result_func(*rargs, **rkwargs): wrapper._dispatch = dispatch wrapper._lazy_collect = _lazy_collect wrapper._sync = sync + wrapper._enable_continous_work = enable_continous_work + if enable_continous_work: + # Attached to the class by remote_class, and called instead of this + # method when the driver routes requests worker by worker. + wrapper._worker_async_companion = _make_worker_async_companion(func, wrapper) return wrapper return decorator diff --git a/src/twinkle/infra/_ray/ray_helper.py b/src/twinkle/infra/_ray/ray_helper.py index 6eb991383..faaa2e62a 100644 --- a/src/twinkle/infra/_ray/ray_helper.py +++ b/src/twinkle/infra/_ray/ray_helper.py @@ -268,6 +268,7 @@ def create_workers(worker_cls: Type[T], instance_id, seed=42, full_determinism=False, + max_concurrency: Optional[int] = None, **kwargs) -> List[T]: # TODO when will remote create remote? # Should it peer create peer? or peer create all? @@ -340,6 +341,11 @@ def create_workers(worker_cls: Type[T], # This is critical for multi-GPU workers (gpus_per_worker > 1) env_vars.update(ResourceManager.noset_env()) + if max_concurrency is not None: + # Read back in the worker to size the thread pool that runs + # blocking continuous-work methods off the actor's event loop. + env_vars['TWINKLE_ACTOR_MAX_CONCURRENCY'] = str(max_concurrency) + runtime_env = RuntimeEnv(env_vars=env_vars) worker_options = { @@ -357,6 +363,9 @@ def create_workers(worker_cls: Type[T], # Use custom resource key for non-GPU accelerators (e.g., NPU). worker_options['resources'] = {device_type: 0.01} + if max_concurrency is not None: + worker_options['max_concurrency'] = max_concurrency + worker = worker_cls.options(**worker_options).remote(*args, **kwargs) workers.append(worker) else: @@ -380,6 +389,11 @@ def create_workers(worker_cls: Type[T], 'TWINKLE_FULL_DETERMINISM': str(int(full_determinism)), **_visible_device_env }) + if max_concurrency is not None: + # Read back in the worker to size the thread pool that runs + # blocking continuous-work methods off the actor's event loop. + env_vars['TWINKLE_ACTOR_MAX_CONCURRENCY'] = str(max_concurrency) + runtime_env = RuntimeEnv(env_vars=env_vars) worker_options = { @@ -391,6 +405,9 @@ def create_workers(worker_cls: Type[T], 'num_cpus': 0.01, } + if max_concurrency is not None: + worker_options['max_concurrency'] = max_concurrency + worker = worker_cls.options(**worker_options).remote(*args, **kwargs) workers.append(worker) return workers diff --git a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py index ad3c24163..ac9d8460f 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py @@ -36,7 +36,12 @@ def _convert_ndarray_to_list(obj: Any) -> Any: return obj -@remote_class() +# max_concurrency: how many sample() calls one worker serves at once. Without it +# Ray runs one method per actor at a time, so concurrent callers queue at the actor +# and never share a batch inside AsyncLLM. 24 is what vLLM reports as the maximum +# concurrency its KV cache holds for this context length; past it vLLM preempts and +# recomputes, which costs more than it gains. +@remote_class(max_concurrency=24) class vLLMSampler(Sampler, CheckpointEngineMixin): """A vLLM-based sampler using VLLMEngine (AsyncLLM). @@ -277,7 +282,7 @@ async def _sample_single( prompt_logprobs=response.prompt_logprobs, topk_prompt_logprobs=response.topk_prompt_logprobs) - @remote_function(dispatch='slice_dp', collect='flatten', lazy_collect=False) + @remote_function(dispatch='slice_dp', collect='flatten', lazy_collect=False, enable_continous_work=True) def sample( self, inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]], diff --git a/src/twinkle/template/base.py b/src/twinkle/template/base.py index ae6411119..d914b01a9 100644 --- a/src/twinkle/template/base.py +++ b/src/twinkle/template/base.py @@ -87,6 +87,16 @@ def clean_tool_call(self, decoded: str) -> str: parser = ToolCallRegistry.detect_first(decoded or '') return parser.clean(decoded) if parser else (decoded or '').rstrip() + def tool_call_errors(self, decoded: str) -> List[str]: + """Why ``parse_tool_call`` returned fewer calls than the text asked for. + + Same parser choice as ``parse_tool_call``, so the two describe one pass + over the reply. Empty when the reply carries no tool-call markup at all -- + a reply that simply answered is not a failure. + """ + parser = ToolCallRegistry.detect_first(decoded or '') + return parser.parse_errors(decoded) if parser else [] + @property def tokenizer(self): tokenizer = self.processor diff --git a/src/twinkle/template/tools/base.py b/src/twinkle/template/tools/base.py index 35b63dc82..a782e9156 100644 --- a/src/twinkle/template/tools/base.py +++ b/src/twinkle/template/tools/base.py @@ -22,6 +22,23 @@ def parse(self, text: str) -> List[Dict[str, Any]]: def clean(self, text: str) -> str: """Strip parser-specific markup; return plain content text.""" + def parse_errors(self, text: str) -> List[str]: + """Why markup this parser recognised produced no call. + + ``detect`` saying yes while ``parse`` returns nothing means the model did + try to call a tool and the markup did not survive parsing. Without this + the caller cannot tell that apart from a reply that called nothing, so it + ends the episode and the model is never told its call was dropped. + Measured on one challenger run: 6 of 59 episodes ended that way, each + with a well-formed ``<tool_call>`` block whose JSON carried a Python-style + ``\\'`` escape or a raw newline. + + One string per block that failed, carrying the parser's own reason (for a + JSON block, the ``json.JSONDecodeError`` text). Default empty: a parser + whose ``parse`` is the same regex as its ``detect`` cannot fail this way. + """ + return [] + def extract_tool_result(self, text: str) -> Optional[str]: """If ``text`` is a tool-result message of this protocol, return the body with the protocol-specific prefix stripped; otherwise return ``None``. diff --git a/src/twinkle/template/tools/bracket_dsl.py b/src/twinkle/template/tools/bracket_dsl.py index dc0b16359..a9facc622 100644 --- a/src/twinkle/template/tools/bracket_dsl.py +++ b/src/twinkle/template/tools/bracket_dsl.py @@ -210,7 +210,15 @@ def _parse_args(self, body: str) -> Dict[str, Any]: return args def parse(self, text: str) -> List[Dict[str, Any]]: + return self._scan(text)[0] + + def parse_errors(self, text: str) -> List[str]: + return self._scan(text)[1] + + def _scan(self, text: str) -> Tuple[List[Dict[str, Any]], List[str]]: + """Calls and failures from one pass, so the two cannot disagree.""" calls: List[Dict[str, Any]] = [] + errors: List[str] = [] text = text or '' for start, end in self._find_blocks(text): block = text[start:end] @@ -221,6 +229,8 @@ def parse(self, text: str) -> List[Dict[str, Any]]: break close = self._match_paren(block, m.end() - 1) if close is None: + errors.append(f'{m.group(1).strip()}( is never closed by a ' + f'matching )') break name = m.group(1).strip() if name: @@ -231,8 +241,10 @@ def parse(self, text: str) -> List[Dict[str, Any]]: 'arguments': self._parse_args(block[m.end():close]), }, }) + else: + errors.append('a call in the list has an empty function name') pos = close + 1 - return calls + return calls, errors def clean(self, text: str) -> str: text = text or '' diff --git a/src/twinkle/template/tools/cline.py b/src/twinkle/template/tools/cline.py index 7f3b2bda0..072e7ed3d 100644 --- a/src/twinkle/template/tools/cline.py +++ b/src/twinkle/template/tools/cline.py @@ -110,7 +110,18 @@ def detect(self, text: str) -> bool: return False def parse(self, text: str) -> list[dict[str, Any]]: + return self._scan(text)[0] + + def parse_errors(self, text: str) -> list[str]: + return self._scan(text)[1] + + def _scan(self, text: str) -> tuple[list[dict[str, Any]], list[str]]: + """Calls and failures from one pass, so the two cannot disagree. + + A tag on the deny list is not a failure: those are skipped on purpose. + """ calls: list[dict[str, Any]] = [] + errors: list[str] = [] for m in _BLOCK_RE.finditer(text or ''): tool = m.group('tool') if tool in _DENY: @@ -119,6 +130,8 @@ def parse(self, text: str) -> list[dict[str, Any]]: for pm in _PARAM_RE.finditer(m.group('body')): args[pm.group('key')] = pm.group('val').strip() if not args: + errors.append(f'<{tool}> holds no <parameter>...</parameter> pair, ' + f'so the call has no arguments') continue calls.append({ 'type': 'function', @@ -127,7 +140,7 @@ def parse(self, text: str) -> list[dict[str, Any]]: 'arguments': args }, }) - return calls + return calls, errors def clean(self, text: str) -> str: if not text: diff --git a/src/twinkle/template/tools/qwen.py b/src/twinkle/template/tools/qwen.py index 6713d570a..a87cf0735 100644 --- a/src/twinkle/template/tools/qwen.py +++ b/src/twinkle/template/tools/qwen.py @@ -1,7 +1,7 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import json import re -from typing import Any, Dict, List +from typing import Any, Dict, List, Tuple from .base import ToolCallParser @@ -20,7 +20,15 @@ def detect(self, text: str) -> bool: return self.open_marker in text def parse(self, text: str) -> List[Dict[str, Any]]: + return self._scan(text)[0] + + def parse_errors(self, text: str) -> List[str]: + return self._scan(text)[1] + + def _scan(self, text: str) -> Tuple[List[Dict[str, Any]], List[str]]: + """Calls and failures from one pass, so the two cannot disagree.""" calls: List[Dict[str, Any]] = [] + errors: List[str] = [] for block_m in self._BLOCK_RE.finditer(text or ''): block = block_m.group(1) func_m = self._FUNCTION_RE.search(block) @@ -43,10 +51,12 @@ def parse(self, text: str) -> List[Dict[str, Any]]: continue try: data = json.loads(block) - except json.JSONDecodeError: + except json.JSONDecodeError as e: + errors.append(str(e)) continue name = data.get('name') or data.get('tool_name', '') if not name: + errors.append('the call object has no "name" field') continue args = data.get('arguments', {}) if isinstance(args, str): @@ -61,7 +71,7 @@ def parse(self, text: str) -> List[Dict[str, Any]]: 'arguments': args if isinstance(args, dict) else {}, }, }) - return calls + return calls, errors def clean(self, text: str) -> str: return self._STRIP_RE.sub('', text or '').rstrip() diff --git a/src/twinkle/template/tools/vcp.py b/src/twinkle/template/tools/vcp.py index 5e030f9d5..4c46fd158 100644 --- a/src/twinkle/template/tools/vcp.py +++ b/src/twinkle/template/tools/vcp.py @@ -1,6 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import re -from typing import Any, Dict, List +from typing import Any, Dict, List, Tuple from .base import ToolCallParser @@ -39,7 +39,15 @@ def detect(self, text: str) -> bool: return _VCP_OPEN in (text or '') def parse(self, text: str) -> List[Dict[str, Any]]: + return self._scan(text)[0] + + def parse_errors(self, text: str) -> List[str]: + return self._scan(text)[1] + + def _scan(self, text: str) -> Tuple[List[Dict[str, Any]], List[str]]: + """Calls and failures from one pass, so the two cannot disagree.""" calls: List[Dict[str, Any]] = [] + errors: List[str] = [] for block in _VCP_BLOCK_RE.findall(text or ''): args: Dict[str, Any] = {} name = '' @@ -51,6 +59,8 @@ def parse(self, text: str) -> List[Dict[str, Any]]: else: args[k] = v if not name: + errors.append('the block has no "tool_name:" line, ' + 'so there is no tool to call') continue calls.append({ 'type': 'function', @@ -59,7 +69,7 @@ def parse(self, text: str) -> List[Dict[str, Any]]: 'arguments': args, }, }) - return calls + return calls, errors def clean(self, text: str) -> str: return _VCP_BLOCK_RE.sub('', text or '').rstrip() diff --git a/src/twinkle_agentic/challenger/agentic.py b/src/twinkle_agentic/challenger/agentic.py index 7cfdcbf63..be74f0f8b 100644 --- a/src/twinkle_agentic/challenger/agentic.py +++ b/src/twinkle_agentic/challenger/agentic.py @@ -392,6 +392,12 @@ class AgenticChallenger(Challenger): combo_arity: ``'triple'`` or ``'mix'``, as in :class:`.CodeChallenger`. arity_weights: weights for the ``'mix'`` subset size. single_kw_prob: chance of using one category in ``'triple'`` mode. + proposals_per_group: how many proposals answer the same keyword draw and + the same prompt, tagged with a shared ``group_id``. This is the group + size the proposing side's advantage is computed over; at 1 every + group has one member and every advantage is zero. At a fixed + proposal count it does not change the compute -- it divides the + number of distinct keyword draws per round by the same factor. keyword_refill_target / keyword_gen_calls / keyword_refill_tries / keyword_params: keyword bank refill parameters. check_params / problem_params: sampling params for the two appended @@ -465,12 +471,14 @@ def __init__( reset_fn: Callable[..., None], run_check_fn: Callable[..., Tuple[int, str]], workspace_snapshot_fn: Optional[Callable[..., str]] = None, + snapshot_error_fn: Optional[Callable[..., str]] = None, tool_schemas: Optional[Sequence[Dict[str, Any]]] = None, episode_concurrency: int = 1, episode_tool_managers: Optional[Sequence[Any]] = None, combo_arity: str = 'triple', arity_weights: Optional[Sequence[float]] = None, single_kw_prob: float = 0.1, + proposals_per_group: int = 1, keyword_refill_target: int = 128, keyword_gen_calls: int = 8, keyword_refill_concurrency: int = 1, @@ -511,6 +519,13 @@ def __init__( self.reset_fn = reset_fn self.run_check_fn = run_check_fn self.workspace_snapshot_fn = workspace_snapshot_fn + # Asked, when a snapshot came back empty, why: the text of the failure if + # the listing could not be read, '' if the workspace really was empty. + # Without it the two are one outcome, and a sandbox the host had paused is + # filed as the model having built nothing -- 63 of run_clean6's 71 + # ``empty_workspace`` rejections were the 410 "sandbox is not proxyable" + # error, so that reject class was 89% broken environment. + self.snapshot_error_fn = snapshot_error_fn self.tool_schemas = list(tool_schemas) if tool_schemas else None # More than one episode at a time needs more than one sandbox: an episode # owns its workspace from the reset until its check has run. The three @@ -540,6 +555,15 @@ def __init__( self.combo_arity = combo_arity self.arity_weights = list(arity_weights) if arity_weights else None self.single_kw_prob = single_kw_prob + # How many proposals answer each keyword draw. Above 1 they form a GRPO + # group on the proposing side; see :meth:`propose`. Raising it does not + # cost more compute at a fixed proposal count -- it trades keyword + # variety for group size, since a round's proposals then come from + # ``count / proposals_per_group`` draws instead of ``count`` of them. + if proposals_per_group < 1: + raise ValueError(f'proposals_per_group must be >= 1, got {proposals_per_group}') + self.proposals_per_group = proposals_per_group + self._next_group_id = 0 self.keyword_refill_target = keyword_refill_target self.keyword_gen_calls = keyword_gen_calls # How many of a refill's generating calls go out together. At 1 each call @@ -607,6 +631,10 @@ def __init__( self.stats: Dict[str, int] = { 'explore_done': 0, 'check_parse_fail': 0, 'check_run_fail': 0, 'empty_workspace': 0, 'solver_truncated': 0, + # The workspace listing could not be read, as opposed to being empty. + # Kept apart from ``empty_workspace`` because it says nothing about + # what the model did. + 'snapshot_unavailable': 0, 'problem_parse_fail': 0, 'too_long': 0, 'parsed': 0, # How often a check that failed was handed back for a rewrite, and # how often the rewrite passed. The two together say whether the @@ -635,11 +663,21 @@ def propose(self, count: int) -> List[Trajectory]: Each carries a direction + keywords + optional seed. The explorer will run these multi-turn in the sandbox. + + ``proposals_per_group`` of them share one keyword draw, one seed choice + and one identical prompt, and are tagged with the same ``group_id``. + That is what makes a GRPO group on the proposing side: the advantage of + a proposal is its reward minus the mean over the others answering the + same prompt, so the members have to differ only by sampling noise. At 1 + -- which is what this used to be, every proposal its own keyword draw -- + every group has one member, the mean equals the reward, and every + advantage is zero. """ proposals: List[Trajectory] = [] directions: List[str] = [] - metas: List[Tuple[List[Tuple[str, str]], bool, str]] = [] - for _ in range(count): + metas: List[Tuple[List[Tuple[str, str]], bool, str, int]] = [] + per_group = max(1, self.proposals_per_group) + while len(metas) < count: picks = self._draw_keywords() body = '\n'.join(f'- {c}: {t}' for c, t in picks) use_seed = bool(self.seeds) and self.rng.random() < self.seed_mix_prob @@ -653,10 +691,21 @@ def propose(self, count: int) -> List[Trajectory]: user = self.prompts.from_keywords.format(keywords=body) else: user = self.prompts.from_scratch - directions.append(user) - metas.append((picks, use_seed, body)) - - for user, (picks, use_seed, body) in zip(directions, metas): + # The whole group gets the same prompt, so a short final group is a + # group whose advantage is computed over fewer samples -- noisier, + # but not wrong. Truncating to a multiple of per_group instead would + # silently return fewer proposals than asked for. + # + # The counter is per-run, not per-call: ``propose`` runs once per + # round, and restarting at 0 each round would give two unrelated + # groups the same id in the dump. + gid = self._next_group_id + self._next_group_id += 1 + for _ in range(min(per_group, count - len(metas))): + directions.append(user) + metas.append((picks, use_seed, body, gid)) + + for user, (picks, use_seed, body, gid) in zip(directions, metas): proposal: Trajectory = { 'messages': [{'role': 'system', 'content': self.prompts.system}, {'role': 'user', 'content': user}], @@ -664,7 +713,8 @@ def propose(self, count: int) -> List[Trajectory]: if self.tool_schemas: proposal['tools'] = self.tool_schemas proposals.append(attach_user_data( - proposal, keywords=picks, seeded=use_seed, keyword_block=body)) + proposal, keywords=picks, seeded=use_seed, keyword_block=body, + group_id=gid)) return proposals @@ -682,6 +732,23 @@ def build(self, explored: List[Trajectory]) -> List[Optional[Trajectory]]: f'{type(self).__name__}.build() must not be called directly; ' f'the serial _round() loop drives one episode at a time instead.') + def _reject_for_empty_snapshot(self, state: Dict[str, Any], slot: int) -> None: + """File an episode whose workspace listing came back empty. + + Empty means one of two unrelated things -- the episode built nothing, or + the listing could not be read -- and only the first says anything about + the model. ``snapshot_error_fn`` is what tells them apart; with no such + callback every case is filed as ``empty_workspace``, which is what used + to happen for all of them. + """ + detail = self.snapshot_error_fn(slot=slot) if self.snapshot_error_fn else '' + if detail: + self._bump('snapshot_unavailable') + state['reject'] = ('snapshot_unavailable', detail) + else: + self._bump('empty_workspace') + state['reject'] = ('empty_workspace', '') + def _followup(self, state: Dict[str, Any], trajectory: Trajectory, n_before: int) -> Optional[Tuple[str, Optional[SamplingParams]]]: """What to say next when the model stops calling tools; ``None`` to stop. @@ -708,8 +775,7 @@ def _followup(self, state: Dict[str, Any], trajectory: Trajectory, # solver passes by doing nothing. Five of run5's ten verified tasks # were that task. Reject here instead. if not snapshot.strip(): - self._bump('empty_workspace') - state['reject'] = ('empty_workspace', '') + self._reject_for_empty_snapshot(state, slot) return None return (self.prompts.check_followup.format(final_state=snapshot), self.check_params) @@ -852,8 +918,7 @@ def _run_followup_api(self, state: Dict[str, Any], explored: Trajectory) -> None # An episode that left nothing behind has no end state to write checks # about; rejecting here mirrors the n_before==0 branch of _followup. if not snapshot.strip(): - self._bump('empty_workspace') - state['reject'] = ('empty_workspace', '') + self._reject_for_empty_snapshot(state, slot) return # Check-script stage: the first ask plus up to ``check_retries`` rewrites, @@ -944,6 +1009,7 @@ def _finish_episode(self, state: Dict[str, Any], """ keywords = user_data_get(explored.get('user_data'), 'keywords', []) seeded = user_data_get(explored.get('user_data'), 'seeded', False) + group_id = user_data_get(explored.get('user_data'), 'group_id', None) # The episode as one record: a single conversation, so a single set of # token ids and logprobs. Handed to propose_sink with whatever verdict the # proposal ends up with, so a rejected attempt is recorded as fully as a @@ -952,7 +1018,8 @@ def _finish_episode(self, state: Dict[str, Any], def reject(reason: str, detail: str = '') -> None: self._reject_record(explored, reason, detail=detail) - self._emit_propose(rounds, reason, keywords=keywords, seeded=seeded) + self._emit_propose(rounds, reason, keywords=keywords, seeded=seeded, + group_id=group_id) if state.get('reject'): reason, detail = state['reject'] @@ -1034,23 +1101,30 @@ def _reject_record(self, traj: Trajectory, reason: str, detail: str = '') -> Non def _emit_propose(self, rounds: Optional[List[Dict[str, Any]]], outcome: str, *, keywords: Any = (), seeded: bool = False, - n_pass: Optional[int] = None) -> None: + n_pass: Optional[int] = None, + group_id: Optional[int] = None) -> None: """Hand one proposal attempt's rounds to ``propose_sink``. - ``pass_rate`` is the raw fraction of solver attempts that succeeded. It - is left as the measurement rather than mapped onto a difficulty score: - the target rate and its tolerance are training decisions, and baking a - guess at them into the dump would make it look like they had been - settled. + ``pass_rate`` is the raw fraction of solver attempts that succeeded. + ``challenger_reward`` is that fraction scored against a 50% target by + :meth:`challenger_reward`, which is the number the proposing side trains + on; both are written so a run can be re-scored under a different target + without re-solving anything. + + A proposal with no ``n_pass`` never reached difficulty measurement -- it + was rejected before that -- and scores 0, the same as one nobody or + everybody solved. """ if self.propose_sink is None or not rounds: return rollouts = self.solver_rollouts or None payload = { 'outcome': outcome, + 'group_id': group_id, 'n_pass': n_pass, 'n_rollouts': rollouts, 'pass_rate': (n_pass / rollouts) if (n_pass is not None and rollouts) else None, + 'challenger_reward': self.challenger_reward(n_pass), 'keywords': list(keywords or ()), 'seeded': bool(seeded), 'rounds': rounds, @@ -1058,6 +1132,25 @@ def _emit_propose(self, rounds: Optional[List[Dict[str, Any]]], outcome: str, *, with self._sink_lock: self.propose_sink(payload) + def challenger_reward(self, n_pass: Optional[int]) -> float: + """Score a proposal by how close the solver came to a 50% pass rate. + + ``1 - 2 * |p - 1/2|`` for ``p = n_pass / solver_rollouts``: 1.0 at half + the attempts passing, 0 at none or all of them. The shape is the one + R-Zero (arXiv 2508.05004) trains its challenger on, and the reason it is + peaked at a half rather than at 'hard' is that a GRPO update's size goes + with the reward variance within a group, which for a pass/fail solver is + ``p(1-p)`` -- largest exactly there. + + ``None`` means the proposal never got as far as being solved, and scores + 0. That is the floor, not a penalty: nothing here can go below 0, so a + failed proposal and an unsolvable one are worth the same. + """ + rollouts = self.solver_rollouts or 0 + if n_pass is None or not rollouts: + return 0.0 + return 1.0 - 2.0 * abs(n_pass / rollouts - 0.5) + def _take_rounds(self, task: Trajectory) -> Optional[List[Dict[str, Any]]]: """Detach a task's proposing rounds. Popped even with no sink attached: token ids for a whole agentic episode are large, and a kept task is held @@ -1183,7 +1276,8 @@ def _drain(slot: int) -> List[Trajectory]: for task in usable: self._emit_propose(self._take_rounds(task), 'kept', keywords=user_data_get(task.get('user_data'), 'keywords', []), - seeded=user_data_get(task.get('user_data'), 'seeded', False)) + seeded=user_data_get(task.get('user_data'), 'seeded', False), + group_id=user_data_get(task.get('user_data'), 'group_id', None)) self.n_proposed += len(proposals) self.n_kept += len(kept) band = (f', in difficulty band {len(kept)}' if self.solver_rollouts else '') @@ -1291,7 +1385,8 @@ def _prepare(k: int) -> bool: 'kept' if kept_flag else 'outside_band', keywords=user_data_get(task.get('user_data'), 'keywords', []), seeded=user_data_get(task.get('user_data'), 'seeded', False), - n_pass=n) + n_pass=n, + group_id=user_data_get(task.get('user_data'), 'group_id', None)) return [t for t, kept_flag in zip(measured, in_band) if kept_flag] def solver_prompt(self, task: Trajectory) -> Trajectory: diff --git a/src/twinkle_agentic/rollout/multi_turn.py b/src/twinkle_agentic/rollout/multi_turn.py index 396a60e25..cbeb8f380 100644 --- a/src/twinkle_agentic/rollout/multi_turn.py +++ b/src/twinkle_agentic/rollout/multi_turn.py @@ -91,6 +91,24 @@ def _default_tool_messages( return msgs +def _malformed_tool_message(errors: List[str]) -> Dict[str, Any]: + """What goes back to the model when its tool-call markup did not parse. + + ``role='tool'`` because it is the outcome of the call the model just tried to + make. There is no ``tool_call_id`` to pair it with -- the call never became a + call -- which ``_default_tool_messages`` above already treats as optional. + """ + reason = '; '.join(e for e in errors if e) or 'the markup could not be parsed' + return { + 'role': + 'tool', + 'content': + ('Your tool call was not run: ' + reason + '. Send the call again. Inside ' + 'a JSON string a backslash has to be written as \\\\ and a line break as ' + '\\n; a single quote needs no backslash at all.'), + } + + @remote_class() class MultiTurnRollout(Rollout): """Agentic multi-turn rollout with tool use (batched). @@ -133,6 +151,7 @@ def __init__( harness: Optional[AgentHarness] = None, adapter_path: Optional[str] = None, stop_after_stuck_turns: int = 0, + max_malformed_retries: int = 2, followup_fn: Optional[Callable[[Trajectory, int], Any]] = None, ): super().__init__() @@ -177,6 +196,20 @@ def __init__( raise ValueError(f'stop_after_stuck_turns must be >= 0, got ' f'{stop_after_stuck_turns}') self.stop_after_stuck_turns = stop_after_stuck_turns + # How many replies in a row may carry tool-call markup that does not + # parse before the episode ends anyway. Such a reply is not the model + # declining to call a tool -- it asked for one and the markup was + # rejected -- so it gets the parser's reason back as a tool message and + # another turn. Measured on one challenger run: 6 of 59 episodes ended + # here, each having written a whole ``<tool_call>`` block whose JSON held + # a Python-style ``\'`` escape or a raw newline, and each was told + # nothing. The cap exists because a model that cannot produce valid JSON + # would otherwise spend all of ``max_turns`` failing to; 0 restores the + # old behaviour of ending the episode on the first one. + if max_malformed_retries < 0: + raise ValueError(f'max_malformed_retries must be >= 0, got ' + f'{max_malformed_retries}') + self.max_malformed_retries = max_malformed_retries # Called with (trajectory, how many follow-ups it has had already) at the # moment an episode would end: because the model stopped calling tools, # because it used up ``max_turns``, or because it was stopped for being @@ -266,6 +299,10 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] stuck_turns: List[int] = [0] * n seen_calls: List[set] = [set() for _ in range(n)] stuck_stop: List[bool] = [False] * n + # Replies in a row whose tool-call markup did not parse. Reset by any + # reply that produced a call, so one bad escape in the middle of a + # working episode does not count against a later one. + malformed_turns: List[int] = [0] * n # Follow-up bookkeeping (all no-ops when ``followup_fn`` is None): # how many follow-ups each trajectory has had, and the params its next # turn should use. A trajectory that has had one stops dispatching tools. @@ -356,6 +393,14 @@ def append_followup(global_idx: int) -> bool: resps_by_idx: Dict[int, Any] = {} device_mesh = getattr(self.sampler, 'device_mesh', None) min_batch_size = (device_mesh.data_world_size if device_mesh is not None else 1) + # A sampler that routes each request on its own accepts a batch smaller + # than its worker count, so the padding below is not needed. It was only + # ever there because slicing a batch over all workers raises when some + # rank gets nothing, and the duplicates it added were generated and then + # dropped -- with one prompt and 8 workers that is 8 generations for 1 + # kept result. + if getattr(type(self.sampler).sample, '_enable_continous_work', False): + min_batch_size = 1 for slot, group in enumerate(groups): batch_pifs = [pifs[i] for i in group] actual = len(batch_pifs) @@ -450,6 +495,27 @@ def append_followup(global_idx: int) -> bool: continue if not tool_calls: + # Markup that did not parse is the model asking for a tool, + # not declining one -- ending here tells it nothing and throws + # the turn away. Hand back the parser's own reason and let it + # write the call again. Not after a follow-up: tools are + # withdrawn there on purpose (see ``followup_fn``), so a reply + # that looks like a call is meant to be read as text. + parse_errors = ([] if followups[global_idx] else + self.template.tool_call_errors(seq.decoded or '')) + if (parse_errors and malformed_turns[global_idx] < self.max_malformed_retries): + malformed_turns[global_idx] += 1 + extended = extend_with_bridge(pifs[global_idx], + [_malformed_tool_message(parse_errors)], + self.template) + if extended is None: + truncated[global_idx] = True + done[global_idx] = True + continue + pifs[global_idx] = extended + if lives[global_idx] is not None: + lives[global_idx]['messages'] = list(extended.get('messages') or []) + continue # The episode is over as far as the model is concerned. Give # the caller one chance to say otherwise -- see # ``followup_fn`` for why this is not a second rollout. @@ -468,6 +534,7 @@ def append_followup(global_idx: int) -> bool: done[global_idx] = True continue + malformed_turns[global_idx] = 0 pending_tools.append((global_idx, list(tool_calls))) # 4. Parallel tool dispatch across the live batch, then harness diff --git a/src/twinkle_agentic/tools/tool_manager.py b/src/twinkle_agentic/tools/tool_manager.py index 3c0697b50..2bdf01410 100644 --- a/src/twinkle_agentic/tools/tool_manager.py +++ b/src/twinkle_agentic/tools/tool_manager.py @@ -5,8 +5,11 @@ from twinkle.data_format import ToolCall from twinkle.data_format.message import Tool as ToolInfo +from twinkle.utils import get_logger from twinkle_agentic.tools.base import Tool +logger = get_logger() + def _extract_name(info: Any) -> Optional[str]: """Read ``function.name`` from an OpenAI-shaped tool / tool-call dict.""" @@ -156,6 +159,11 @@ def call_many( back with the same glob listing, so the model was told its python had run when it never did. Nothing in that turn needed concurrency -- the reason it was used was a tool name the host could have refused on the spot. + + Once the tools share an Env, ``step_batch`` is the only way the batch + runs: a raise or a short result list is reported as the result of those + calls, not retried down the thread pool. The thread pool is for tools + that have no Env in common. """ calls = list(tool_calls) if not calls: @@ -175,12 +183,33 @@ def call_many( out[i] = self(calls[i]) try: results = env.step_batch([(name, args) for _i, name, args in batched]) - except Exception: - results = None - if results is not None and len(results) == len(batched): - for (i, _name, _args), r in zip(batched, results): - out[i] = r.observation if hasattr(r, 'observation') else str(r) + except Exception as e: # noqa + # The exception text is the only account of why the batch did not + # run, and the model is what has to react to it, so it goes back + # as the result of every call in the batch. Retrying down the + # thread pool instead -- which is what this used to do, silently + # and without even a log line -- sends the turn along the path the + # docstring above exists to keep it off. + logger.warning(f'{type(env).__name__}.step_batch raised ' + f'{type(e).__name__}: {e}') + failure = f'Error: tool batch did not run: {type(e).__name__}: {e}' + for i, _name, _args in batched: + out[i] = failure + return ['' if x is None else x for x in out] + if len(results) != len(batched): + # Same reasoning: a short result list means the calls did not all + # run, and pairing them up by position would report one call's + # result under another's name. + logger.warning(f'{type(env).__name__}.step_batch returned ' + f'{len(results)} results for {len(batched)} calls') + failure = (f'Error: tool batch did not run: the environment returned ' + f'{len(results)} results for {len(batched)} calls.') + for i, _name, _args in batched: + out[i] = failure return ['' if x is None else x for x in out] + for (i, _name, _args), r in zip(batched, results): + out[i] = r.observation if hasattr(r, 'observation') else str(r) + return ['' if x is None else x for x in out] workers = max_workers or min(32, len(calls)) out = [None] * len(calls) From 3d33b6447b1b2fc017f46d14ced12fb7bf586763 Mon Sep 17 00:00:00 2001 From: tastelikefeet <yuze.zyz@alibaba-inc.com> Date: Sat, 29 Aug 2026 10:48:53 +0800 Subject: [PATCH 46/60] wip --- cookbook/rsi/agentic/challenge.py | 43 +++++- cookbook/rsi/agentic/prompts.py | 10 +- cookbook/rsi/agentic/remote_tool_env.py | 60 ++------ cookbook/rsi/agentic/rl.py | 6 +- cookbook/rsi/agentic/rsi_agent.yaml | 36 +++-- .../rsi/agentic/sandbox_server/tool_server.py | 141 +++++++++++++++++- cookbook/rsi/agentic/train_offline.py | 43 +++++- pyproject.toml | 2 +- src/twinkle_agentic/challenger/agentic.py | 7 + tests/twinkle_agentic/test_agentic_rsi.py | 90 +++++++++++ 10 files changed, 369 insertions(+), 69 deletions(-) diff --git a/cookbook/rsi/agentic/challenge.py b/cookbook/rsi/agentic/challenge.py index fe0f6321d..47b85dfc2 100644 --- a/cookbook/rsi/agentic/challenge.py +++ b/cookbook/rsi/agentic/challenge.py @@ -125,6 +125,13 @@ # since the alternative is throwing the episode away. SNAPSHOT_RETRY_WAIT = 3 +# Same idea for the workspace clear, which is the one sandbox call whose failure +# is fatal to the run rather than to one episode. Longer than the snapshot wait +# because what it waits out is different: a clear times out when ms-agent's +# per-call limit expires with the delete still running, so the second attempt +# wants the first one's rmtree to have drained rather than to race it. +RESET_RETRY_WAIT = 10 + # The ground truth the check script is written against. A listing alone is not # enough: three of the six rejected proposals in the first real run failed on a # value the model recomputed from its own recollection ("Mean values mismatch") @@ -669,8 +676,14 @@ def reset_fn(slot: int = 0): This is also the one point where losing the sandbox costs nothing, since the workspace is about to be emptied regardless -- so a runtime that went away is rebuilt here instead of ending a run that may have hours of - proposals behind it. + proposals behind it. The same reasoning covers a clear that *fails* on a + runtime still answering /health, which ensure_ready cannot see: run 'rsi' + reached iteration 7 -- six checkpoints, eight hours -- and ended on three + clears timing out at ms-agent's per-call limit while the sandbox itself + was healthy enough to report them. So the clear is retried, and then + retried on a deliberately rebuilt sandbox, before the run is given up. """ + clear = CLEAR_WORKSPACE.format(workspace=args.workspace) if envs[slot].ensure_ready(): # A rebuilt sandbox starts empty, so the clear below is redundant, # but running it anyway keeps one path through this function. The @@ -678,8 +691,25 @@ def reset_fn(slot: int = 0): # re-fetched rather than reused. runners[slot] = envs[slot].runner() logger.warning(f'[challenge] sandbox {slot} was rebuilt before this episode') - exit_code, output = runners[slot]( - CLEAR_WORKSPACE.format(workspace=args.workspace), 'python') + exit_code, output = runners[slot](clear, 'python') + if exit_code != 0: + logger.warning(f'[challenge] workspace reset failed on sandbox {slot} ' + f'(exit {exit_code}), retrying in {RESET_RETRY_WAIT}s: ' + f'{output[-200:]}') + time.sleep(RESET_RETRY_WAIT) + exit_code, output = runners[slot](clear, 'python') + if exit_code != 0: + # Rebuilt rather than retried again: two failures in a row is not the + # transient this waits out, and a fresh sandbox brings a workspace + # that is already empty -- which is all this function is asked for. + # Counted as a recovery so the run's own tally at the end still + # accounts for every rebuild, including the ones from here. + logger.warning(f'[challenge] workspace reset failed twice on sandbox {slot} ' + f'(exit {exit_code}); rebuilding it: {output[-200:]}') + envs[slot].reset() + envs[slot].n_recoveries += 1 + runners[slot] = envs[slot].runner() + exit_code, output = runners[slot](clear, 'python') if exit_code != 0: raise RuntimeError(f'workspace reset failed (exit {exit_code}): {output[-400:]}') @@ -1105,6 +1135,13 @@ def write(self, record): 'trace_id': trace_id, 'npz': npz_name, 'outcome': record.get('outcome'), + # Both of these come straight from _emit_propose and both are what the + # proposing side trains on: train_offline.py groups proposals by + # group_id to get a GRPO advantage out of them, and skips the whole + # dump as a "pre-grouping run" when it is absent. Dropping them here + # silently turned SIDES=both into solver-only training. + 'group_id': record.get('group_id'), + 'challenger_reward': record.get('challenger_reward'), 'n_pass': record.get('n_pass'), 'n_rollouts': record.get('n_rollouts'), 'pass_rate': record.get('pass_rate'), diff --git a/cookbook/rsi/agentic/prompts.py b/cookbook/rsi/agentic/prompts.py index 23359e840..890271a7f 100644 --- a/cookbook/rsi/agentic/prompts.py +++ b/cookbook/rsi/agentic/prompts.py @@ -240,7 +240,13 @@ 'it.\n\n' 'Give the two halves differently:\n' '- INPUT data, the raw material nothing was computed from yet: verbatim, ' - 'exact filenames and exact contents, so it can be written byte for byte.\n' + 'exact filenames and exact contents, so it can be written byte for byte. ' + 'Only passive data counts as input -- a CSV, a JSON config, a binary record ' + 'file, a text corpus. Source code is NEVER input data: do not quote the ' + 'body of any script, function or module you wrote, not even one you call a ' + 'fixture. A statement whose input half is the program asks the reader to ' + 'retype your solution, and then it measures typing just as surely as quoting ' + 'a computed answer does.\n' '- Everything DERIVED from it -- computed values, aggregates, orderings, ' 'resolved references, reports: only the RULE that produces it. Name the ' 'output file and its format, say how each part follows from the input, and ' @@ -250,6 +256,8 @@ 'Rules:\n' '- Be specific about formats, filenames and layout.\n' '- Say what must be true of the result, not which commands to run.\n' + '- Describe the behaviour any script must have -- its inputs, its outputs, ' + 'the transformation between them -- and let the reader write the code.\n' '- Do NOT mention the checks or how verification works.\n' '- Self-contained: no reference to this conversation or to anything the ' 'reader cannot see.\n' diff --git a/cookbook/rsi/agentic/remote_tool_env.py b/cookbook/rsi/agentic/remote_tool_env.py index d36dde62a..3ea21b2a3 100644 --- a/cookbook/rsi/agentic/remote_tool_env.py +++ b/cookbook/rsi/agentic/remote_tool_env.py @@ -89,21 +89,15 @@ _RPC_HEADROOM = 60 _LOCAL_SERVER = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'sandbox_server', 'tool_server.py') -# ms-agent's permission package, uploaded alongside the yaml because this -# repository's copy carries two switches the released code does not have: -# ``safety_rules.unrestricted_removal`` and ``safety_rules.allow_write_globs``, -# which rsi_agent.yaml turns on. Without them the sandbox reads those keys, -# finds no code that looks at them, and silently keeps refusing `rm -rf *`, -# `cp src/* dst/` and `chmod +x bin/*` -- a whole run's worth of tasks shaped by -# a config that never took effect. -# -# The image installs ms-agent editable from a tarball into /opt/ms-agent, so -# replacing the files under it is what the interpreter picks up; and it has to -# land before tool_server.py imports them, which is why this is part of the same -# upload rather than a rebuilt image. -_MS_AGENT_PERMISSION_DIR = '/opt/ms-agent/ms_agent/permission' -_PATCHED_PERMISSION_FILES = ('config.py', 'safety.py', 'shell_validator.py', - 'path_validator.py') +# No ms-agent code is uploaded. ``rsi_agent.yaml`` turns on two safety switches +# -- ``safety_rules.unrestricted_removal`` and ``safety_rules.allow_write_globs`` +# -- that ms-agent's ``SafetyConfig.from_dict`` does not implement and silently +# ignores. This used to be handled by copying a patched ``ms_agent/permission`` +# package into every sandbox, which meant carrying a fork of a dependency that +# twinkle supports as a harness, and re-merging it forever. ``tool_server.py`` +# now applies the same two relaxations as a runtime patch inside the sandbox +# (``_patch_permission``), next to the ``python_executor`` patch that was already +# there, so the released ms-agent is used as-is on both sides. def tool_payload(observation: str) -> str: @@ -485,46 +479,20 @@ def _create_sandbox(self): return Sandbox.create(template=self._template, timeout=self._sandbox_timeout) def _upload(self) -> None: - """Push the yaml, the server script and the permission patch into the sandbox. + """Push the yaml and the server script into the sandbox. Uploading beats baking them into the image: the training host's copy is authoritative, so editing a tool line-up is a restart rather than a template rebuild, and the two halves cannot fall out of sync. + + Only twinkle's own two files travel. The safety relaxations + ``rsi_agent.yaml`` asks for are applied by ``tool_server.py`` at runtime, + so no ms-agent source is shipped or overwritten here. """ with open(self._config_path, encoding='utf-8') as f: self._sandbox.files.write(f'{_REMOTE_DIR}/rsi_agent.yaml', f.read()) with open(_LOCAL_SERVER, encoding='utf-8') as f: self._sandbox.files.write(f'{_REMOTE_DIR}/tool_server.py', f.read()) - self._upload_permission_patch() - - def _upload_permission_patch(self) -> None: - """Overwrite ms-agent's permission package with this repository's copy. - - Taken from the installed package rather than a path built out of - ``__file__``, so what lands in the sandbox is the same code the training - host imports. - - Verified rather than assumed: a silent miss here does not fail anything, - it just leaves the sandbox refusing commands the yaml said to allow, and - the only symptom would be a run whose tasks are quietly narrower than - intended. If the switch is not readable afterwards, this raises. - """ - import ms_agent.permission as _perm - - local_dir = os.path.dirname(os.path.abspath(_perm.__file__)) - for name in _PATCHED_PERMISSION_FILES: - with open(os.path.join(local_dir, name), encoding='utf-8') as f: - self._sandbox.files.write(f'{_MS_AGENT_PERMISSION_DIR}/{name}', f.read()) - probe = ('python -c "from ms_agent.permission.config import SafetyConfig as S; ' - 'print(S().unrestricted_removal, S().allow_write_globs)"') - result = self._sandbox.commands.run(probe, timeout=60) - out = (getattr(result, 'stdout', '') or '').strip() - if out.split() != ['False', 'False']: - raise RuntimeError( - 'permission patch did not land in the sandbox: expected the two ' - f'switches to exist and default to False, got {out!r}. The ' - f'sandbox may install ms-agent somewhere other than ' - f'{_MS_AGENT_PERMISSION_DIR}.') def _start_server(self) -> None: """Launch the tool runtime in the background, with its output on disk. diff --git a/cookbook/rsi/agentic/rl.py b/cookbook/rsi/agentic/rl.py index c53b6797a..9621850f2 100644 --- a/cookbook/rsi/agentic/rl.py +++ b/cookbook/rsi/agentic/rl.py @@ -117,7 +117,11 @@ def main(): lora_dropout=0.05, ) - model = TransformersModel(model_id=MODEL_ID, device_mesh=model_mesh, remote_group='model') + # torch_dtype=float32: load master weights in fp32. Training precision is still + # governed by `mixed_precision` (bf16 autocast); fp32 params avoid the numerically + # fragile bf16 optimizer state on this Blackwell + CUDA 13 box. + model = TransformersModel(model_id=MODEL_ID, device_mesh=model_mesh, remote_group='model', + torch_dtype='float32') model.add_adapter_to_model(ADAPTER_NAME, lora_config, gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) model.set_optimizer('AdamW', lr=LEARNING_RATE) diff --git a/cookbook/rsi/agentic/rsi_agent.yaml b/cookbook/rsi/agentic/rsi_agent.yaml index 5f058f686..430b7e920 100644 --- a/cookbook/rsi/agentic/rsi_agent.yaml +++ b/cookbook/rsi/agentic/rsi_agent.yaml @@ -124,10 +124,17 @@ tools: plan_filename: .ms_agent/plan.json plan_md_filename: .ms_agent/plan.md -# Every refusal ms-agent applies to a shell command, turned off. Read by -# LLMAgent.prepare_runtime (llm_agent.py builds PermissionConfig.from_dict off -# this section), so it takes effect in the sandbox, where tool_server.py loads -# this same file. +# Every refusal ms-agent applies to a shell command, turned off. `allow_network` +# and the two list-valued keys are read by LLMAgent.prepare_runtime (llm_agent.py +# builds PermissionConfig.from_dict off this section) and take effect in the +# sandbox, where tool_server.py loads this same file. +# +# The last two keys are different: ms-agent's SafetyConfig does NOT implement +# them, and from_dict ignores unknown keys without a word, so on their own they +# would be dead letters. tool_server.py reads them itself and applies the two +# relaxations as a runtime patch (_patch_permission) inside the sandbox -- +# ms-agent is a harness twinkle supports, so it is used as released rather than +# forked. The startup line reports which ones took effect. # # The reason is what the refusals cost here rather than what they protect: this # runs in a microVM that is reset once per episode and holds nothing but the @@ -147,18 +154,25 @@ permission: # means "use the defaults". patterns: [] # Same, for the configurable half of the rm/rmdir path check: `*`, `/*`, - # `/`, `~`. + # `/`, `~`. Left empty rather than removed to say the intent out loud; + # unrestricted_removal below bypasses the whole check, this list included, + # and tool_server.py warns at startup if the two ever disagree. dangerous_removal_paths: [] - # And the half that a config cannot reach, added for this: the refusals - # written into is_dangerous_removal_path for `*`, anything ending in `/*`, - # `/`, a direct child of `/` (which `/workspace` is), and the home directory. + # And the half that a config cannot reach, which is why this one needs the + # runtime patch: the refusals written into is_dangerous_removal_path for + # `*`, anything ending in `/*`, `/`, a direct child of `/` (which + # `/workspace` is), and the home directory. unrestricted_removal: true # A separate refusal, found by running commands through SafetyGuard rather # than by reading the config: a glob anywhere in a write or create path is # denied on its own ("Glob patterns not allowed in write operations"), which - # is what actually stopped `rm -rf *` and `rm -rf build/*` after the two - # lists above were emptied. It is not specific to rm -- `cp src/* dst/` and - # `chmod +x bin/*` hit it too. + # is what actually stopped `rm -rf build/*` after the two lists above were + # emptied. It is not specific to rm -- `cp src/* dst/` is refused by the same + # check. (`chmod +x bin/*` is NOT: measured through SafetyGuard, chmod's + # arguments are not extracted as write paths, so it was already allowed.) + # Patched by widening the path that is scope-checked to the directory the + # glob expands inside, so a glob still cannot reach outside the workspace -- + # `cp /etc/* /workspace/` stays denied, now for being out of scope. allow_write_globs: true # Web search is deliberately absent. ms-agent's `web_search` key only provides diff --git a/cookbook/rsi/agentic/sandbox_server/tool_server.py b/cookbook/rsi/agentic/sandbox_server/tool_server.py index 9e5cce8ed..2910abff6 100644 --- a/cookbook/rsi/agentic/sandbox_server/tool_server.py +++ b/cookbook/rsi/agentic/sandbox_server/tool_server.py @@ -45,6 +45,10 @@ _SINGLE_NS_FLAG = '_twinkle_single_namespace' +# Marks a permission function this file has already replaced, so a second +# ToolRuntime in one process does not wrap a wrapper. +_PERMISSION_FLAG = '_twinkle_permission_relaxed' + # ms-agent namespaces every tool as ``{server}---{tool}``. _TOOL_SPLIT = '---' @@ -169,6 +173,95 @@ async def python_executor(self, code, description='', timeout=None): return True +def _patch_permission(unrestricted_removal: bool, + allow_write_globs: bool) -> List[str]: + """Honour two safety switches ms-agent's config schema does not implement. + + ``rsi_agent.yaml`` asks for ``safety_rules.unrestricted_removal`` and + ``safety_rules.allow_write_globs``. ``SafetyConfig.from_dict`` reads only the + keys it knows and ignores the rest without a word, so on an unmodified + ms-agent both are dead letters: the sandbox goes on refusing ``rm -rf + build/*``, ``cp src/* dst/`` and ``chmod +x bin/*``, and the only symptom is + a run whose tasks are quietly narrower than the config asked for. + + Done as a runtime patch rather than by editing ms-agent, because ms-agent is + a supported harness and an ordinary dependency: a forked ``permission`` + package would have to be carried, and re-merged, by everyone who runs this + cookbook. Same reasoning and same shape as :func:`_patch_python_executor`. + + Both refusals are written in ``path_validator``, but ``shell_validator`` and + ``safety`` pulled them into their own namespaces with ``from ... import``, + so the replacement is written into every module holding a reference -- + patching the source module alone would leave the copies that actually get + called untouched. + + Returns the names of the patches applied, for the startup line. + """ + from ms_agent.permission import path_validator, safety, shell_validator + + applied: List[str] = [] + # Every module that holds a reference, source module included. + targets = (path_validator, shell_validator, safety) + + if unrestricted_removal: + original_removal = path_validator.is_dangerous_removal_path + # Skips this block only, never the one below: an early return here left + # allow_write_globs unapplied whenever the two were configured together. + if not getattr(original_removal, _PERMISSION_FLAG, False): + + def is_dangerous_removal_path(path, extra_patterns=(), *args, **kwargs): + """No path is too dangerous to remove inside a disposable microVM. + + A blanket bypass, including ``dangerous_removal_paths``: the + checks this switch exists to drop are the fixed ones -- ``*``, + anything ending in ``/*``, ``/``, a direct child of ``/`` (which + ``/workspace`` is) and the home directory -- and they are + entangled with the configurable list in one function. Honouring + the list here would mean restating ms-agent's matching rules, + which is the duplication this whole approach avoids. The caller + is warned at startup when it configured a list this makes moot. + """ + return False + + setattr(is_dangerous_removal_path, _PERMISSION_FLAG, True) + for module in targets: + if hasattr(module, 'is_dangerous_removal_path'): + module.is_dangerous_removal_path = is_dangerous_removal_path + applied.append('unrestricted_removal') + + if allow_write_globs: + original_validate = path_validator.validate_path + if not getattr(original_validate, _PERMISSION_FLAG, False): + + def validate_path(path, cwd, allowed_dirs, op_type, **kwargs): + """Let a glob through a write/create path, scope-checked as usual. + + The glob is handed on as the directory it expands inside, which + is what the original checks anyway once past the deny -- and it + is ms-agent's own ``get_glob_base_directory`` that decides where + that boundary falls, so no policy is restated here. Quotes are + stripped first for the same reason the original does it: a + quoted ``'src/*'`` would otherwise yield a base of ``'src``. + """ + if op_type in ('write', 'create'): + bare = path + if len(bare) >= 2 and bare[0] == bare[-1] and bare[0] in ('"', "'"): + bare = bare[1:-1] + if path_validator.GLOB_CHARS & set(bare): + base = path_validator.get_glob_base_directory(bare) + return original_validate(base, cwd, allowed_dirs, op_type, + **kwargs) + return original_validate(path, cwd, allowed_dirs, op_type, **kwargs) + + setattr(validate_path, _PERMISSION_FLAG, True) + for module in targets: + if hasattr(module, 'validate_path'): + module.validate_path = validate_path + applied.append('allow_write_globs') + + return applied + + def _usable_llm(cfg) -> bool: """Whether the declared ``llm`` section can actually serve a request. @@ -317,6 +410,11 @@ def __init__(self, config_path: str, workspace: str) -> None: self.agent._event_sink = None self.agent._input_source = None self.workspace = workspace + # Before prepare_runtime(), which is what builds SafetyGuard and its + # validators. Read off the merged config for the same reason `llm` is: + # ms-agent layers its own agent.yaml underneath ours, so this is what + # actually took effect rather than what our file happens to say. + self.permission_patches = _patch_permission(*self._safety_switches()) self._loop = _LoopThread() self._loop.run(self._prepare()) # Only after prepare_tools(): a contract can only be read off a tool that @@ -327,6 +425,29 @@ async def _prepare(self) -> None: self.agent.prepare_runtime() await self.agent.prepare_tools() + def _safety_switches(self) -> Tuple[bool, bool]: + """``(unrestricted_removal, allow_write_globs)`` as configured. + + Absent means off, which is what an unmodified ms-agent does with these + keys anyway -- so a config that never mentions them keeps every refusal. + + Warns when ``dangerous_removal_paths`` is configured alongside + ``unrestricted_removal``, because the patch makes that list moot and a + silently ignored blacklist is the one outcome worth shouting about. + """ + rules = {} + permission = getattr(self.agent.config, 'permission', None) + if permission is not None: + rules = getattr(permission, 'safety_rules', None) or {} + unrestricted = bool(_cfg_get(rules, 'unrestricted_removal', False)) + globs = bool(_cfg_get(rules, 'allow_write_globs', False)) + if unrestricted and _cfg_get(rules, 'dangerous_removal_paths', None): + sys.stderr.write('[tool_server] WARNING unrestricted_removal bypasses the ' + 'rm/rmdir path check entirely, so the configured ' + 'dangerous_removal_paths list will not be consulted\n') + sys.stderr.flush() + return unrestricted, globs + @property def _tm(self): return self.agent.tool_manager @@ -567,6 +688,20 @@ def _as_text(result: Any) -> str: return str(result) +def _cfg_get(node: Any, key: str, default: Any = None) -> Any: + """Read ``key`` off an OmegaConf node or a plain dict. + + The permission section arrives as a DictConfig when the yaml declares it and + as a dict when it is assembled in code, and only one of those answers to + ``.get``. + """ + if node is None: + return default + if isinstance(node, dict): + return node.get(key, default) + return getattr(node, key, default) + + class _Handler(BaseHTTPRequestHandler): runtime: ToolRuntime = None # set on the class before the server starts protocol_version = 'HTTP/1.1' @@ -633,8 +768,10 @@ def main() -> None: server = ThreadingHTTPServer((args.host, args.port), _Handler) names = [(t.get('function') or {}).get('name') for t in runtime.tools()] llm_note = 'llm configured' if runtime.has_llm else 'no llm (read_file.abbreviate withdrawn)' - sys.stderr.write(f'[tool_server] ready on {args.host}:{args.port}, {llm_note}, ' - f'{len(names)} tools: {names}\n') + perm_note = (', permission: ' + '+'.join(runtime.permission_patches) + if runtime.permission_patches else '') + sys.stderr.write(f'[tool_server] ready on {args.host}:{args.port}, {llm_note}' + f'{perm_note}, {len(names)} tools: {names}\n') sys.stderr.flush() server.serve_forever() diff --git a/cookbook/rsi/agentic/train_offline.py b/cookbook/rsi/agentic/train_offline.py index 183703ba1..156c9fdb0 100644 --- a/cookbook/rsi/agentic/train_offline.py +++ b/cookbook/rsi/agentic/train_offline.py @@ -55,8 +55,27 @@ # CLI default is never zero, so a fallback here would be dead code that reads # like the default. LEARNING_RATE = args.optimizer.learning_rate -MINI_BATCH_SIZE = args.training.mini_batch_size or 8 -MICRO_BATCH_SIZE = args.training.micro_batch_size or 2 +# One trajectory per micro batch, because padding_free is off: a micro batch is +# padded to its longest member, so pairing a short solver trajectory with a long +# proposer episode pays for the long one twice. Proposer episodes are whole +# agentic rollouts -- median 7.5k tokens, up to 15.7k -- against much shorter +# solver attempts, and at 2 per micro batch step 7 hit a 31k-token padded batch +# and died of CUDA OOM with activation recompute already at its most aggressive +# setting. At 1 the peak is the longest single trajectory and no token is padding. +# +# Read from the environment rather than args.training: TrainingArgs defaults +# micro_batch_size to 2, not to None, so `args.training.micro_batch_size or 1` +# would be dead code that silently kept 2 -- which is exactly how the OOM +# survived a first attempt at this fix. +MICRO_BATCH_SIZE = int(os.environ.get('RSI_MICRO_BATCH_SIZE', 1)) +# forward_backward is declared dispatch='slice_dp', so a mini-batch is sliced +# across the model's data-parallel ranks and each rank collates only its share. +# That share has to hold at least one micro batch, so the floor is +# MODEL_GPUS * MICRO_BATCH_SIZE: at 8 GPUs the old default of 8 left every rank +# with a single trajectory and step 1 died in collate_fn, before any optimizer +# step, on a full 384-trajectory dump. rl.py keeps the same floor by skipping +# batches smaller than MODEL_GPUS. +MINI_BATCH_SIZE = args.training.mini_batch_size or MODEL_GPUS * MICRO_BATCH_SIZE SAVE_STEPS = args.training.save_steps or 0 RUN_DIR = os.environ.get('RSI_RUN_DIR', '') @@ -443,14 +462,19 @@ def main(): # checkpoint is a whole model rather than something that needs merging before # the next challenger round can load it. from twinkle.model.megatron import MegatronModel + # padding_free is off, so this stays off with it: both switches send collate_fn + # down the per-micro-batch packed path, and Megatron's TE extension then reads + # PackedSeqParams.pad_between_seqs, which the Megatron-LM checkout on this box + # does not define. Padded batches cost throughput but keep the attention path + # on plain padded sequences. model = MegatronModel(model_id=MODEL_ID, device_mesh=model_mesh, remote_group='model', - mixed_precision='bf16', variable_seq_lengths=True) + mixed_precision='bf16', variable_seq_lengths=False) model.set_optimizer('default', lr=LEARNING_RATE) model.set_lr_scheduler('default', lr_decay_steps=max(1, n_steps), max_lr=LEARNING_RATE) # beta=0: no reference model here, and grpo.py:315 needs beta>0 AND ref_logps # for the KL term, so any beta above 0 would silently do nothing. model.set_loss('GRPOLoss', epsilon=0.2, beta=0.0) - model.set_processor(InputProcessor, padding_free=True) + model.set_processor(InputProcessor, padding_free=False) model.set_template(TEMPLATE, model_id=MODEL_ID, max_length=MAX_MODEL_LEN, enable_thinking=True) # approx_kl at the first inner step compares the sampler's logps against the @@ -471,6 +495,17 @@ def main(): for mb in range(0, len(inputs), MINI_BATCH_SIZE): end = min(mb + MINI_BATCH_SIZE, len(inputs)) + # Drop a tail shorter than a whole mini batch instead of handing it + # over. The floor is MINI_BATCH_SIZE, not MICRO_BATCH_SIZE: dispatch + # 'slice_dp' splits the batch across all MODEL_GPUS ranks, so a batch + # that cannot give every rank its own micro batch fails before + # collate_fn -- _dispatch_args raises 'Batch too small for N workers, + # some ranks have no data'. A 6-trajectory tail against 8 ranks did + # exactly that. The dropped trajectories are the tail of the last step. + if end - mb < MINI_BATCH_SIZE: + logger.info(f'[offline] step {step + 1}: dropping last ' + f'{end - mb} traj, under mini batch {MINI_BATCH_SIZE}') + break model.forward_backward( inputs=inputs[mb:end], old_logps=old_logps[mb:end], diff --git a/pyproject.toml b/pyproject.toml index 8392acfb0..80bf4d1b9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ version = "0.4.0.dev0" description = "Training API for large language models with efficient data handling and advanced optimization techniques." readme = "README.md" authors = [{ name = "ModelScope", email = "contact@modelscope.cn" }] -requires-python = ">=3.10,<3.13" +requires-python = ">=3.10,<=3.13" dependencies = [ "numpy>=2.0.0,<2.3.0", "datasets", diff --git a/src/twinkle_agentic/challenger/agentic.py b/src/twinkle_agentic/challenger/agentic.py index be74f0f8b..5c9105f9c 100644 --- a/src/twinkle_agentic/challenger/agentic.py +++ b/src/twinkle_agentic/challenger/agentic.py @@ -1064,7 +1064,14 @@ def reject(reason: str, detail: str = '') -> None: task: Trajectory = { 'messages': [{'role': 'user', 'content': statement}], } + # group_id travels with the task, not just with the reject path above: the + # difficulty stage emits the surviving proposals from the task, so a task + # that forgets its group reaches the dump ungrouped and the proposing side + # has no advantage to compute. Leaving it off here made every kept and + # outside_band proposal group_id=None, and only the episodes that failed + # early -- which emit straight off ``explored`` -- kept theirs. task = attach_user_data(task, check_script=script, keywords=keywords, seeded=seeded, + group_id=group_id, setup_script=state.get('setup_script', '')) # Carried, not emitted: the verdict this proposal earns depends on the # difficulty measurement, which has not run yet. A plain top-level key diff --git a/tests/twinkle_agentic/test_agentic_rsi.py b/tests/twinkle_agentic/test_agentic_rsi.py index ece9d9a06..96ffa6b69 100644 --- a/tests/twinkle_agentic/test_agentic_rsi.py +++ b/tests/twinkle_agentic/test_agentic_rsi.py @@ -1347,5 +1347,95 @@ def test_the_avoid_list_is_capped_and_drops_older_entries_first(self): self.assertIn(f'new {cap + 29}', note, 'the newest is always kept') +class ProposeTrajIndexTest(unittest.TestCase): + """index.jsonl has to carry what the proposing side trains on. + + The challenger emits a proposal record and challenge.py copies it into + index.jsonl field by field. Two of those fields are the reason the dump + exists at all: train_offline.py groups proposals by ``group_id`` to get a + GRPO advantage out of them, and skips a dump without it as a 'pre-grouping + run'. While the copy dropped both, SIDES=both trained 384 solver and 0 + proposer trajectories, and said so only in a line nobody read. + """ + + def test_group_id_and_reward_survive_the_copy(self): + from challenge import ProposeTrajWriter + + out = tempfile.mkdtemp(prefix='proposetraj_test_') + try: + writer = ProposeTrajWriter(out) + writer.write({ + 'outcome': 'kept', + 'group_id': 0, + 'challenger_reward': 0.75, + 'n_pass': 4, + 'n_rollouts': 8, + 'pass_rate': 0.5, + 'keywords': [['transform', 'parse a binary log']], + 'seeded': False, + 'rounds': [{'stage': 'episode', 'messages': [], + 'input_ids': [1, 2], 'labels': [-100, 2], + 'logprobs': []}], + }) + writer.close() + with open(os.path.join(out, 'index.jsonl'), encoding='utf-8') as f: + rec = json.loads(f.readline()) + finally: + shutil.rmtree(out, ignore_errors=True) + + # Group 0 is a real group, so this also pins that the copy reads the key + # rather than testing it for truth. + self.assertEqual(rec['group_id'], 0) + self.assertEqual(rec['challenger_reward'], 0.75) + + +class TaskCarriesGroupIdTest(unittest.TestCase): + """A built task has to remember which group proposed it. + + The difficulty stage emits kept and outside_band proposals off the *task*, + so a task built without its group_id reaches the dump ungrouped and the + proposing side gets no advantage from it. The reject path reads group_id off + the episode instead, so while only the successful path dropped it, a run + showed 4 grouped proposals -- all of them early failures -- against 92 + ungrouped kept/outside_band ones, and the copy downstream looked correct. + """ + + def _challenger(self): + from twinkle_agentic.challenger.agentic import AgenticChallenger, AgenticPrompts + + prompts = AgenticPrompts( + system='s', from_scratch='u', + check_followup='write checks for {final_state}', + check_retry_followup='{error} / {final_state}', + problem_followup='write the statement') + return AgenticChallenger( + prompts, + lambda trajectories, **kwargs: list(trajectories), + reset_fn=lambda slot=0: None, + run_check_fn=lambda script, slot=0: (0, ''), + workspace_snapshot_fn=lambda slot=0: 'data.csv 3', + solver_rollouts=0, + ) + + def test_group_id_reaches_the_built_task(self): + from twinkle.data_format import user_data_get + from twinkle_agentic.challenger.base import attach_user_data + + ch = self._challenger() + explored = attach_user_data( + {'messages': [{'role': 'user', 'content': 'explore'}, + {'role': 'assistant', 'content': 'done'}]}, + keywords=[['transform', 'parse a binary log']], seeded=False, group_id=7) + state = {'checked': True, 'script': 'assert True', + 'statement': 'PROBLEM: build a parser\nEND'} + + task = ch._finish_episode(state, explored) + + self.assertIsNotNone(task, 'a checked episode with a statement is a task') + # 7, not None: the emit sites downstream read exactly this key, and a None + # here is what silently turned SIDES=both into solver-only training. + self.assertEqual(user_data_get(task.get('user_data'), 'group_id', None), 7) + + if __name__ == '__main__': unittest.main() From 09d52c405fd3ce2cabc00b018c6f09ec3deb8841 Mon Sep 17 00:00:00 2001 From: root <yuze.zyz@alibaba-inc.com> Date: Mon, 31 Aug 2026 12:08:15 +0800 Subject: [PATCH 47/60] wip --- cookbook/rsi/agentic/README.md | 449 ++- cookbook/rsi/agentic/challenge.py | 2561 +++++++++-------- cookbook/rsi/agentic/episode.py | 23 +- cookbook/rsi/agentic/eval.py | 57 +- cookbook/rsi/agentic/loop.sh | 233 ++ cookbook/rsi/agentic/prompts.py | 286 +- cookbook/rsi/agentic/rl.py | 286 -- .../rsi/agentic/rsi_agent_shortprompt.yaml | 197 ++ cookbook/rsi/agentic/sandbox.py | 292 ++ cookbook/rsi/agentic/train.py | 363 +++ cookbook/rsi/agentic/train_offline.py | 552 ---- src/twinkle_agentic/challenger/agentic.py | 414 ++- src/twinkle_agentic/challenger/task_bank.py | 130 + src/twinkle_agentic/verifier/__init__.py | 5 + src/twinkle_agentic/verifier/rubric_score.py | 451 +++ 15 files changed, 3888 insertions(+), 2411 deletions(-) create mode 100644 cookbook/rsi/agentic/loop.sh delete mode 100644 cookbook/rsi/agentic/rl.py create mode 100644 cookbook/rsi/agentic/rsi_agent_shortprompt.yaml create mode 100644 cookbook/rsi/agentic/sandbox.py create mode 100644 cookbook/rsi/agentic/train.py delete mode 100644 cookbook/rsi/agentic/train_offline.py create mode 100644 src/twinkle_agentic/challenger/task_bank.py create mode 100644 src/twinkle_agentic/verifier/rubric_score.py diff --git a/cookbook/rsi/agentic/README.md b/cookbook/rsi/agentic/README.md index d6542c1ff..92fdf77fb 100644 --- a/cookbook/rsi/agentic/README.md +++ b/cookbook/rsi/agentic/README.md @@ -1,264 +1,243 @@ -# Agentic RSI +# agentic — RSI self-play where one trajectory is one request -Self-play for tool-using agents. The model invents its own tasks by doing them, -then trains on the ones it solves only sometimes. +One model plays both roles. It builds something in a sandbox, then writes a task +description for what it built, then tries to redo that task from the description +alone. How often it succeeds is what scores the description: a task the solver +passes sometimes is worth training on, one it always or never passes is not. -The solver is an **ms-agent** agent -- shell, filesystem, python, notebook, -todo -- working inside a microVM. When it stops calling tools the episode ends -and the task's check script is run against what it left behind. Checks are -ordinary programs, so the same trajectory always earns the same reward. +This replaced an earlier version of the same method whose difference was +scheduling: there a round of proposals moved through the pipeline as a batch and +every stage waited for the slowest member. Here each trajectory is its own request +from start to finish, and the only place anything waits is the last step, deciding +whether a group of eight is worth keeping. (The old version was retired to +`.temp/agentic_legacy`; nothing here imports from it.) -## Pipeline +## Three resources, three queues -``` -challenge.py rl.py -┌──────────────────────────┐ ┌────────────────────────┐ -│ 1 direction + keywords │ │ 1 read flows │ -│ 2 model acts in sandbox │ │ 2 boot sandboxes │ -│ 3 model writes checks │ flows │ 3 solver works N turns │ -│ 4 run checks (verify) │ ───────► │ 4 run check_script │ -│ 5 model writes statement │ │ 5 GRPO │ -│ 6 difficulty filter │ └────────────────────────┘ -└──────────────────────────┘ -``` - -`challenge.py` builds tasks backwards: the end state exists before the question -does, so nothing it produces can be unachievable. The difficulty filter then -drops tasks the solver always passes or always fails -- either way GRPO gets a -zero gradient from the whole group. - -## Where things run +| resource | how many at once | who queues on it | +|---|---|---| +| sandbox | `--sandbox-slots` microVMs (32) | one job owns one slot from the workspace clear to its last check | +| vLLM | `enable_continous_work` routes each request to the least busy worker | every build turn and every solver turn, one trajectory per request | +| API | `--api-concurrency` (32) | check scripts, problem statements, the rubric | + +There is one FIFO job queue and one thread per sandbox slot, so a slot is never +idle while there is work. A build that finishes hands its statement to eight +solver jobs, releases its slot, and returns — it never waits for its own solvers, +which is what would deadlock a pool against itself. Rubric jobs go to a separate +pool because they need no sandbox. + +A batch of one is a first-class vLLM call here. `challenge.py` refuses to start if +the sampler does not advertise `enable_continous_work`, because without it a batch +of one is padded up to the worker count and most of every generation is thrown +away. + +## One proposal, three stages + +1. **Build.** Local model, sandbox tools, one tool call per reply, up to + `--max-turns`. This is the trainable part: the trajectory keeps exactly the + tokens the local model produced. +2. **Check script.** The workspace is read back byte for byte and appended to a + *copy* of the build conversation; qwen3.8-max writes a python script that + asserts the end state. It is rejected on the syntax tree if it pins file sizes, + checksums or a script's source text, then run in the sandbox. One rewrite. +3. **Problem statement.** Same copy, one more API reply: input data verbatim, + everything derived given as the rule that produces it. + +Stages 2 and 3 run on the API so the check and the statement are written with the +whole build history in view without adding untrained tokens to the sample. + +## Groups, and the one place things wait + +A group is `--group-size` (8) proposals sharing one keyword draw and one prompt. +That is what makes it a GRPO group: a proposal's advantage is its reward minus the +mean over the others answering the same prompt. + +- Each proposal's task gets `--solver-rollouts` (8) attempts. `n_pass` is how many + passed, with the denominator fixed at 8 — a truncated attempt is a failed + attempt, the same as one whose assertions failed. +- Once all eight builds are in, the eight statements are scored for novelty + *against each other* plus the closest entries in the task bank. Waiting for all + eight costs nothing: the slots are held by other groups' jobs the whole time. If + a statement still has no verdict after `--novelty-tries` (3), the group is + dropped and its queued solver attempts are skipped. +- **The group is kept when at least one proposal has `n_pass` in `[1, 7]`.** The + other seven may be anything, including builds that produced no task at all; + they train with the reward they earned, which for those is 0. +- From a kept group the highest-reward in-band proposal is selected, and its eight + solver attempts are what the solver side trains on. The unselected proposals' + attempts were measured and are reported, but not trained on. + +Eight kept groups give 64 proposing and 64 solving trajectories: one training step. + +## Reward + +Proposing side, unchanged from where it was measured: ``` -training host sandbox (microVM) -───────────────────────────── ───────────────────────────────── -MsAgentHarness tool_server.py - system prompt, message shaping ms-agent ToolManager - llm: and tools: popped -> the real file_system / - constructs no tool at all code_executor / todo_list - ┌──► GET /tools schemas -RemoteMsAgentToolEnv ── curl over ─────┤ POST /call dispatch - forwards tool calls the sandbox │ - copies files back command channel└─── /workspace +reward = exp(-(n_pass/8 - 0.2)^2 / (2 * 0.3^2)) * (floor + (1-floor) * novelty) +reward = 0 when n_pass is 0 or unmeasured ``` -Two properties this layout is built around: +The gaussian peaks at a pass rate of 0.2, not 0.5: a proposal only teaches the +solver something when the solver mostly cannot do it yet. The floor at +`n_pass <= 0` is load-bearing — the gaussian at p=0 is 0.801, higher than the +0.607 it gives a proposal half the attempts solve, so without the gate the best +thing a proposer could do is write tasks nobody can finish. -**Nothing the model emits executes next to the trainer.** The harness has its -`llm` and `tools` sections popped *after* ms-agent merges its own `agent.yaml` -underneath -- omitting them from `rsi_agent.yaml` is not enough, since that -default declares `code_executor` and would otherwise put a live shell executor -on the training host with access to the whole machine. +Note that being out of band does not zero the reward. A task everybody solves +still earns about 0.03. Out of band decides whether the task is delivered to the +solver side; it does not zero the proposer's score. -**The advertised tool contract is read off the code that honours it.** The tool -schemas in the prompt come from `GET /tools` on the sandbox, not from a second -ms-agent next to the trainer. In RL the policy actively exploits whatever the -executor actually does, and any divergence from the production tools would only -surface after deployment. +`floor` defaults to 1, which makes the novelty term exactly 1.0 — the score is +still judged and still written to `novelty_scores.jsonl`, it just does not move a +reward. Measured on iter1's 27 proposals: judged against their own siblings 24 of +27 scored exactly 0.0, which is the right answer (a keyword draw produces eight +paraphrases of one task) and also a useless one, since a term constant across the +group contributes nothing after GRPO subtracts the group mean. Labelling each +task's shape on its own instead does separate proposals within a group, but the +label changed between sampled repeats on 10 of 27 statements. `NOVELTY_FLOOR=0.5` +puts it back in. -## Setup +Solving side: 1.0 if the check exits 0, else 0.0. -### Environment host - -Needs `/dev/kvm` and kernel 6.8+. Builds the template and runs the AgentENV -server. - -```bash -sh sandbox_server/install.sh # AgentENV + build the template -sh sandbox_server/install.sh --rebuild # after changing the Dockerfile -sh sandbox_server/install.sh --skip-install # template only, AgentENV already up +## Files -sh sandbox_server/serve.sh # foreground, binds 127.0.0.1:8000 -NOHUP=1 sh sandbox_server/serve.sh # background ``` - -On a restricted network the base image will not resolve from Docker Hub; point it -at a reachable mirror (forwarded to `aenv build --image`, so the Dockerfile stays -untouched): - -```bash -BASE_IMAGE=docker.m.daocloud.io/library/python:3.11-slim sh sandbox_server/install.sh +challenge.py collect: the queues, the three job bodies, the group decision +train.py one GRPO step over what was collected, then overwrite the ckpt +sandbox.py the sandbox as a resource: clear, snapshot, run a script +prompts.py every string sent to a model +loop.sh collect -> train -> collect from the new weights, until killed +episode.py how an episode is built and scored, shared with eval.py +remote_tool_env.py the transport to one microVM, paired with sandbox_server/ +sandbox_server/ the image and the in-sandbox tool server it talks to +eval.py held-out pass rate on tasks the trainer never saw +split_tasks.py split a collection's tasks into a train and an eval half +rsi_agent.yaml the ms-agent config both sides' openings are shaped by ``` -`docker.m.daocloud.io` is a third-party Docker Hub proxy -- every sandbox's base -image comes through it. Substitute your own Aliyun accelerator address -(`<id>.mirror.aliyuncs.com`) if you would rather not depend on one. - -#### When the template build is too slow to use - -On 2026-08-23 three `aenv build` runs on our host failed or ran for hours, and -the cause was download speed rather than the Dockerfile. Measured within one -minute, from inside a sandbox: `deb.debian.org` 33 KB/s, `mirrors.aliyun.com` -5.4 MB/s, and for comparison the host itself 12 MB/s and sandbox disk writes -639 MB/s. apt's package index alone is 9.6MB, so the build sat two hours with no -output -- and since the server logs `template build started` and then nothing -until the build ends, slow is indistinguishable from hung. The Dockerfile now -rewrites `deb.debian.org` to the Aliyun mirror, which should remove the cause. - -The path that is verified end to end installs inside a live sandbox and -snapshots it, which needs no template builder and takes about six minutes: +Output under `--out-dir`: -```bash -sh sandbox_server/build_via_sandbox.sh # name: twinkle-rsi-msagent -NAME=twinkle-rsi-msagent-v2 sh sandbox_server/build_via_sandbox.sh # verify first ``` - -Three things to know about a snapshot: - -* it shows up in `aenv snapshot list`, **not** `aenv template list`, but the name - lives in the same namespace -- `--sandbox-template twinkle-rsi-msagent` - resolves to it unchanged; -* it keeps the filesystem, not the image config, so the Dockerfile's `ENV - PYTHONUNBUFFERED=1`, `ENV PIP_INDEX_URL=...` and `WORKDIR /workspace` are gone. - `build_via_sandbox.sh` writes `/etc/pip.conf` and `/workspace` instead, and - `remote_tool_env` starts the runtime with `python -u`; -* aenv refuses to rebind an existing name, so replacing an image means deleting - the old one first. Build under a second name and verify before you do that -- - deleting first cost us four hours with no usable sandbox. - -Verify either one from the training host with the boot check below. - -The Dockerfile also pins `PIP_INDEX_URL` to an Aliyun mirror for the same -reason -- edit those two lines if your host reaches pypi.org directly. - -The image clones **ms-agent from source** (`--depth 1` of `main`), not the pip -release: the tools the policy is trained against are the ones in the repository, -and a released wheel can lag behind it. - -Verify a sandbox boots and the runtime comes up before anything else: - -```bash -python -c " -import sys; sys.path.insert(0, '.') -from remote_tool_env import RemoteMsAgentToolEnv -e = RemoteMsAgentToolEnv(template='twinkle-rsi-msagent', config_path='rsi_agent.yaml', - api_url='http://127.0.0.1:8000') -e.reset() -print([t['function']['name'] for t in e.tool_schemas()]) -print(e.step('shell_executor', {'command': 'python -V && pwd'}).observation) -e.close() -" +trajs/*.npz input_ids / labels / logprobs +trajs/index.jsonl one line per trained trajectory: side, group, reward, messages +groups.jsonl one line per decided group, kept or not, and why +tasks.jsonl the statements and check scripts delivered +rejected.jsonl every build that produced no task, and how its episode ended +solver_attempts.jsonl every attempt: the check's output and the workspace it left +novelty_scores.jsonl the rubric, all three dimensions and all nine verdicts +keyword_gen.jsonl every keyword call, prompt and reply verbatim +keywords.jsonl the keyword bank, carried between iterations +challenge_metrics.json this collection as numbers: scalars, raw counters, histograms +train_summary.json what the step actually trained, and what it skipped ``` -### Step 1 -- generate tasks +Only `trajs/` is read again — by `train.py`. The rest is written for reading after +the fact: `solver_attempts.jsonl` is the only thing that answers whether a task at +`n_pass=0` was unsolvable or the solver gave up, and `novelty_scores.jsonl` records +usefulness and complexity, which are scored by the same call but reach no reward. + +`challenge_metrics.json` is the exception: `train.py` reads its `scalars` section and +sends it to swanlab together with the training metrics, so one chart carries both +halves of an iteration. It is computed by reading `groups.jsonl` back rather than +from the live objects, so it cannot disagree with the audit file beside it, and the +same function recomputes it for a directory that finished hours ago. Three sections: + +* `scalars` — fixed keys, every value a number. What goes up. Includes + `solve_pass_rate`, the accuracy: passes over every solver attempt that ran. Read + it as a property of the pair, not of the model — the tasks change every iteration, + so a rise can be the solver improving or the proposer getting easier, and + `n_pass_in_band_rate` next to it is what separates those. +* `counts` — the raw counters, dynamic keys and all. `group_dropped:rubric_error` + exists only in a run where that happened, so these stay in the file and are not + uploaded: a chart that appears halfway through a run reads as a change in the run. +* `distributions` — the `n_pass`, build-outcome and novelty histograms behind the + means, because a mean `n_pass` of 4 is a different collection depending on whether + it came from eights and zeros or from fours. + +Nothing is truncated in these files. They are read to check whether a reward was +deserved, which a shortened statement cannot answer. + +Everything is in this directory. `sandbox.py` takes its transport from +`remote_tool_env.py`, which is paired with the tool server in `sandbox_server/`, +and the solver's opening from `episode.solver_harness`, which `eval.py` uses too — +so a task's `n_pass` here and its `pass@k` there are measured against one opening. + +## Running it ```bash -pip install e2b - -python challenge.py \ - --keep-target 200 \ - --sandbox-template twinkle-rsi-msagent \ - --sandbox-api-url http://<env-host-ip>:8000 \ - --sampler-gpus 4 +export E2B_API_KEY=... # sandbox host +export SANDBOX_API_URL=http://... # sandbox host address, with port +export LLM_BACKUP_API_KEY=... # dashscope +ITERATIONS=1 bash cookbook/rsi/agentic/loop.sh ``` -Writes `output/rsi_agentic/challenge_flows.jsonl`, one task per line: -`{id, query, check_script, n_pass, n_rollouts, keywords, seeded}`. +Charts land in swanlab project `twinkle-rsi-agentic`, one experiment named after +`TAG`, one step per iteration. `train.py` uploads after saving the checkpoint, so a +swanlab failure costs the charts and not the weights — the numbers are still in +`challenge_metrics.json` and `train_summary.json` either way. Resume is by +`id=TAG`: a second run under the same tag appends to that curve, a new tag starts a +new one. `RSI_SWANLAB_MODE=disabled` turns it off, `RSI_SWANLAB_PROJECT` moves it. -Also writes `output/rsi_agentic/propose_traj/` -- the rounds that *produced* each -task, kept and rejected alike, as one `.npz` of token ids / labels / logprobs per -attempt plus an `index.jsonl` carrying the text and the outcome. Nothing reads it -yet. It exists because proposing is generation like any other, so those rounds -could be trained on later; rejects are in there on purpose, since a set of -kept-only attempts has no zero-reward half to contrast against. `pass_rate` is -stored raw -- mapping it onto a difficulty score means choosing a target rate, -which is a training decision, not a dump format. Pass `--dump-propose-traj ''` -to turn it off; expect it to dwarf the task file. +Verified on this machine at swanlab 0.9.2: three separate processes with the same +tag at steps 1, 2, 3 landed on one run (the second and third print `disabled in +resume mode`). Resume works only in `online` mode — in `local` mode each process +made its own run directory instead. -Useful flags: `--seed-file` (start from existing trajectories), `--keywords-n 0` -(no keyword bank), `--solver-rollouts` (attempts per task in the difficulty -filter), `--max-turns` (tool-calling turns per episode). - -Round 1 is serial -- one episode at a time, workspace cleared in between -- so -`--max-proposals-per-round` trades throughput against how often the estimator -recalibrates. - -### Step 2 -- train - -```bash -AENV_API_URL=http://<env-host-ip>:8000 \ -AENV_TEMPLATE=twinkle-rsi-msagent \ -RSI_TASKS=output/rsi_agentic/challenge_flows.jsonl \ - python rl.py --model-id ms://Qwen/Qwen3-4B \ - --model-gpus 4 --sampler-gpus 4 -``` +## Settings that shape what gets produced -`rl.py` accepts both task formats: `check_script` (from `challenge.py`, scored -by exit status) and structured `checks` (see `tasks.example.jsonl`). +Every one of these changes either the model's output or how it is scored. The +origin column says where the value came from; nothing marked *inherited* has been +re-measured under this scheduler. -Through a tunnel instead: - -```bash -ssh -N -L 8000:127.0.0.1:8000 root@<env-host-ip> -``` - -## Configuration - -`challenge.py` is all command-line flags (`--help`). `rl.py` reads: - -| Variable | Default | | +| setting | value | origin | |---|---|---| -| `AENV_API_URL` | `http://127.0.0.1:8000` | AgentENV server | -| `AENV_TEMPLATE` | `twinkle-rsi-msagent` | template or snapshot name to boot from | -| `RSI_TASKS` | `tasks.example.jsonl` | task file | -| `RSI_AGENT_CONFIG` | `rsi_agent.yaml` | uploaded into every sandbox | -| `RSI_SANDBOX_TIMEOUT` | `900` | must outlast an episode plus its checks | -| `RSI_ENV_CONCURRENCY` | `16` | parallel boot / scoring | -| `RSI_MAX_TURNS` | `20` | tool-calling turns per episode | -| `RSI_SCORE_MODE` | `fraction` | or `all_or_nothing`; structured checks only | -| `RSI_KEEP_WORKSPACES` | `0` | keep the files copied out of each sandbox | - -Training hyper-parameters come from the CLI, e.g. -`python rl.py --batch-size 2 --num-generations 4 --max-steps 2`. - -Sandbox count during training is `batch-size x num-generations`, each ~2GiB. At -the defaults (4 x 8) that is 32 microVMs, so size the environment host -accordingly. - -## Files - -| File | Role | -|---|---| -| `challenge.py` | task generation: act, write checks, verify, describe, filter | -| `prompts.py` | every string the challenger sends; categories live here too | -| `rl.py` | training loop, episode construction, scoring | -| `remote_tool_env.py` | training-side Env: forwards tool calls, copies files back | -| `rsi_agent.yaml` | ms-agent config -- read by *both* halves | -| `sandbox_server/tool_server.py` | in-sandbox HTTP server owning the ToolManager | -| `sandbox_server/Dockerfile` | template image: ms-agent, ripgrep, ffmpeg, imagemagick, openpyxl/reportlab/pdfplumber | -| `sandbox_server/install.sh` | install AgentENV + build the template | -| `sandbox_server/build_via_sandbox.sh` | build the image as a snapshot of a live sandbox instead | -| `sandbox_server/serve.sh` | start the AgentENV server | -| `tasks.example.jsonl` | hand-written tasks in the structured `checks` format | - -The machinery both scripts call lives in `twinkle_agentic.challenger`; only -wiring and prompt text are here. - -## Decisions worth knowing - -**Round 1 is serial.** Every episode needs an empty workspace and they share one -long-lived sandbox, so proposals cannot overlap -- the second episode would see -the first one's files and its checks would pass for free. The difficulty filter -resets between attempts for the same reason. - -**Checks are a python script, not a structured list.** `challenge.py` asks the -model to write asserts against the state it just produced, which is the same -trick the code challenger uses (execute first, capture the result, make that the -ground truth). Exit status is the whole verdict: no partial credit, no judge -model, no drift between rounds. - -**`read_file(abbreviate=True)` is withdrawn when no LLM is configured.** That -argument asks an LLM to summarise a file. The sandbox has no API key, so the -tool server drops the unusable `llm` section *and* removes the argument from the -advertised schema -- the model is never offered something that can only fail. -Give `rsi_agent.yaml` a real `llm:` section with a key reachable from the sandbox -to get it back. - -**A failed sandbox boot skips the whole training batch.** GRPO groups here are -positional: advantages are taken over consecutive runs of `num_generations`, so -dropping one episode would shift every later group onto the wrong task. There is -no retry -- a boot failure is logged and the step is abandoned. - -**No web search.** ms-agent's `web_search` key only provides `fetch_page` -(retrieve a known URL). A real search tool needs `EXA_API_KEY` / `SERPAPI_API_KEY` -and is wired separately; no example task requires one. +| `--keep-groups` / `--group-size` / `--solver-rollouts` | 8 / 8 / 8 | decided for this pipeline | +| keep rule: ≥1 proposal with `n_pass ∈ [1,7]` | — | decided for this pipeline | +| truncated solver attempt counts as a failure, denominator fixed at 8 | — | decided for this pipeline | +| a build cut off at `--propose-max-tokens` writes no check and no statement | — | restored from the old pipeline, which skipped both stages after a length cut | +| rubric failure after 3 tries drops the whole group | — | decided for this pipeline | +| `--max-build-files` | 4 | inherited: `loop.sh` has passed this since it was added. It is text in the system prompt. | +| `--api-thinking-budget` | 4096 | inherited from the old `loop.sh` | +| `--propose-max-tokens` / `--max-turns` / `--stop-after-stuck-turns` | 8192 / 24 / 2 | inherited | +| `--one-call-per-reply` | on | inherited | +| `--check-retries` / `--check-max-tokens` | 1 / 8192 | inherited | +| `--problem-max-tokens` / `--problem-max-chars` | 4096 / 8192 | inherited | +| `--solver-max-tokens` / `--solver-max-turns` | 8192 / 24 | inherited | +| temperature / top_p, both sides | 1.0 / 0.95 | inherited | +| `--novelty-floor` | 1 | decided after measuring: the term was constant across the group at floor 0.5, so it only scaled the whole reward down | +| `--task-bank-refs` / `--novelty-tries` | 5 / 3 | inherited | +| `--keywords-n` / `--keyword-gen-calls` / `--keyword-temp` | 128 / 8 / 1.3 | inherited | +| `--snapshot-max-files` / `-per-file` / `-budget` | 50 / 600 / 6000 | inherited | +| `--sandbox-slots` | 32 | inherited: a probe once held 96, but not reliably for a whole run | +| lr / one optimizer step / `GRPOLoss(epsilon=0.2, beta=0.0)` | 1e-6 | inherited | +| `MICRO_BATCH_SIZE=1`, `padding_free=False` | — | inherited, forced by an OOM at 2 | + +Prompt texts are byte-identical to the ones the old pipeline sent — verified +string by string — minus the seed and single-model follow-up strings, which this +pipeline never sends. + +## What has been checked, and what has not + +Checked offline, `.tmp_analysis/test_agentic.py` — the real scheduler, group +state machine, job bodies, rubric loop and writers against fake vLLM/sandbox/API. +25 checks, all passing: 8 kept groups produce exactly 64 + 64, zero-reward +proposals are still written, only the selected proposal's attempts are, a group +with nothing in band is dropped, a rubric that never returns a verdict drops the +whole group after 3 tries and skips its queued solver jobs, a length-cut build +writes no check and no statement, `solver_attempts.jsonl` has one line per +attempt (400 of them, 243 failures, each with the check's output and the workspace +it left), and `novelty_scores.jsonl` logs all three dimensions on every retry. +`train.py` then reads the same directory back as 16 groups of 8, 128 +trajectories, nothing skipped, every group centred by `GRPOAdvantage`. + +Two ways the run could hang were found by that test and fixed rather than worked +around: an exception inside the rubric job (whose future nobody reads) left its +group waiting for a verdict forever, and an exception while writing a decided +group's output skipped the launch of its replacement topic. Both now log and let +the run continue, and `run()` additionally stops with a message if it ever goes +quiet without reaching its target. + +Not checked: anything requiring a GPU or a sandbox. No end-to-end run has been +done, so there are no wall-clock, keep-rate or `n_pass` numbers under this +scheduler, and none of the inherited settings above have been re-measured. diff --git a/cookbook/rsi/agentic/challenge.py b/cookbook/rsi/agentic/challenge.py index 47b85dfc2..2534ddac4 100644 --- a/cookbook/rsi/agentic/challenge.py +++ b/cookbook/rsi/agentic/challenge.py @@ -1,1187 +1,1506 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""RSI self-play, agentic half: generate training tasks by doing them first. - -One model plays both roles, and one proposal is one conversation. It first acts -as an agent in a sandbox (multi-turn with tools), producing a trajectory and a -final workspace state; then, appended to that same conversation, it writes a -check script that verifies the end state, and finally describes the task as a -problem statement. The same model then attempts the problem multiple times, and -only problems it solves *sometimes* are kept. - -The machinery lives in :mod:`twinkle_agentic.challenger`; the prompts live in -``prompts.py`` next to this file. What is here is the wiring: which model, how -many, the sandbox connection, and where the output goes. - -Output format (one JSONL line per task): - - --out-flows {id, query, check_script, n_pass, n_rollouts, keywords, seeded} +"""RSI self-play, agentic half: one trajectory is one request, start to finish. + +Three resources, each a queue anyone may put a request on: + + sandbox N microVMs. A job holds one for as long as it needs the workspace. + vLLM the local sampler. ``sample`` routes each request to the least busy + worker (``enable_continous_work``), so a batch of one is a first-class + call and 32 threads calling it concurrently is the intended use. + API qwen3.8-max, for the stages that must not add untrained tokens: the + check script, the problem statement, and the rubric. + +Nothing waits for a batch. A proposal that finishes its build hands its statement +straight to eight solver jobs and lets go of its sandbox; those eight run whenever +a slot frees up, in any order, interleaved with proposals from other groups. The +only synchronisation is the last step, deciding whether a group is worth keeping, +and that is a counter under a lock rather than a barrier. + +A group is eight proposals sharing one keyword draw and one prompt -- that is what +makes it a GRPO group, since the advantage of a proposal is its reward minus the +mean over the others answering the same prompt. It is kept when at least one of its +eight produced a task the solver passes sometimes (``1 <= n_pass <= 7``); the other +seven may be anything, including failures, and they train with reward 0. From a kept +group the highest-reward in-band proposal's eight solver attempts are what the +solver side trains on, so a kept group contributes 8 proposing and 8 solving +trajectories, and eight kept groups are the 64 + 64 one training step reads. + +Output (all under ``--out-dir``): + + trajs/*.npz input_ids / labels / logprobs per trajectory + trajs/index.jsonl one line per trajectory: side, group, reward, full text + groups.jsonl one line per decided group: why kept or dropped + tasks.jsonl the statements and check scripts that were delivered + keywords.jsonl the keyword bank, carried between iterations Run it as a Ray job (sampler only, no trainer):: - python cookbook/rsi/agentic/challenge.py --keep-target 200 + python cookbook/rsi/agentic/challenge.py --keep-groups 8 """ import argparse +import collections import json -import base64 -import binascii -import hashlib +import math import os +import queue +import random +import statistics import sys +import threading import time -from typing import Dict +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Tuple import numpy as np import twinkle from twinkle import DeviceGroup, DeviceMesh, get_logger -from twinkle.data_format import SamplingParams, user_data_get +from twinkle.data_format import SamplingParams from twinkle.sampler import vLLMSampler -from twinkle_agentic.challenger import AgenticChallenger, KeywordStore -from twinkle_agentic.envs import EnvTool -from twinkle_agentic.rollout import build_rollout +from twinkle_agentic.challenger import KeywordStore, parse_check_script, parse_problem_statement +from twinkle_agentic.challenger.agentic import brittle_check_reason +from twinkle_agentic.challenger.code import parse_keyword_list +from twinkle_agentic.challenger.task_bank import TaskBank +from twinkle_agentic.rollout import MultiTurnRollout from twinkle_agentic.tools.tool_manager import ToolManager sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -from episode import solver_harness # noqa: E402 -from prompts import CATEGORIES, CATEGORY_DESC, agentic_prompts # noqa: E402 -from remote_tool_env import RemoteMsAgentToolEnv, tool_payload # noqa: E402 +import prompts as P # noqa: E402 +from sandbox import close_pool, open_pool, solver_harness # noqa: E402 logger = get_logger() -# Cleared through the python tool, not `rm -rf`: ms-agent's safety policy rejects -# `rm -rf` outright ("Blocked by safety rule"), and it rejects globs in write -# operations, which rules out `find -delete` too. The script asserts the -# directory really is empty, so a future policy change surfaces as a failed reset -# instead of tasks quietly inheriting the previous workspace. -# -# Module level so a test can drive the same string the run does; a copy in a test -# would keep passing after this one changed. -CLEAR_WORKSPACE = ''' -import os, shutil -root = {workspace!r} -os.makedirs(root, exist_ok=True) -for name in os.listdir(root): - path = os.path.join(root, name) - if os.path.isdir(path) and not os.path.islink(path): - shutil.rmtree(path, ignore_errors=True) - else: - os.remove(path) -leftover = os.listdir(root) -assert not leftover, 'workspace not empty after clear: %r' % (leftover,) -''' - -# ── Arm B: copy the episode's input files out, and put them back later ────── -# -# The bytes travel, not a description of them: the model is not asked to write a -# script that recreates its own inputs, because a wrong one costs the whole -# episode and would fail exactly where the inputs are least ordinary. -# -# Read in slices because the executor truncates its output at roughly 8 KB. Each -# slice is base64 so any byte survives the trip, and the manifest's sha256 is -# what says the trip was faithful -- checked locally against the bytes that -# arrived, so a truncated read cannot pass as a smaller file. -INPUT_MANIFEST = ''' -import hashlib, os -root = os.path.join({workspace!r}, 'input') -rows = [] -for dirpath, dirnames, filenames in os.walk(root): - dirnames[:] = [d for d in dirnames if d not in {{'__pycache__', '.ipynb_checkpoints'}}] - for name in sorted(filenames): - path = os.path.join(dirpath, name) - with open(path, 'rb') as handle: - body = handle.read() - rows.append((os.path.relpath(path, {workspace!r}), len(body), - hashlib.sha256(body).hexdigest())) -for rel, size, digest in sorted(rows): - print(rel, size, digest) -''' - -INPUT_SLICE = ''' -import base64, os -path = os.path.join({workspace!r}, {rel!r}) -with open(path, 'rb') as handle: - handle.seek({offset}) - print(base64.b64encode(handle.read({length})).decode()) -''' - -# What gets stored with the task and run before every attempt at it. Writes the -# captured bytes and nothing else: no cleanup, because whoever runs this has just -# cleared the workspace. -SETUP_SCRIPT_TEMPLATE = '''# Recreate the task's input files. -import base64, os, pathlib - -FILES = {files!r} - -for rel, payload in FILES.items(): - path = pathlib.Path(rel) - if path.parent != pathlib.Path('.'): - os.makedirs(path.parent, exist_ok=True) - with open(path, 'wb') as handle: - handle.write(base64.b64decode(payload)) -''' - -# Seconds to wait before asking a sandbox for its workspace listing a second -# time. 62 of run_clean6's 63 snapshot failures were the sandbox answering 410 -# "not proxyable", which is the host having paused it -- worth one more ask, -# since the alternative is throwing the episode away. -SNAPSHOT_RETRY_WAIT = 3 - -# Same idea for the workspace clear, which is the one sandbox call whose failure -# is fatal to the run rather than to one episode. Longer than the snapshot wait -# because what it waits out is different: a clear times out when ms-agent's -# per-call limit expires with the delete still running, so the second attempt -# wants the first one's rmtree to have drained rather than to race it. -RESET_RETRY_WAIT = 10 - -# The ground truth the check script is written against. A listing alone is not -# enough: three of the six rejected proposals in the first real run failed on a -# value the model recomputed from its own recollection ("Mean values mismatch") -# rather than read off the file, so the end state has to arrive as content, not -# just as names. Bounded on both axes -- 50 files, 600 bytes each, 6000 overall -- -# because this goes into a prompt and a 100k artifact would push the trajectory -# it has to be read alongside out of the window. -# -# Walks the tree in python rather than shelling out to `find`: the same code then -# decides what is text, what is truncated, and what the budget was spent on, -# which a pipeline of find/head cannot report back. -# -# File bodies go out byte for byte. An earlier version printed `body.rstrip()`, -# which hid trailing newlines while the size column still counted them, so a -# check writer shown an 11-byte file whose content looked 10 characters long -# wrote `content == 'Mean: 63.9'` and the check failed against the very state it -# was written from. The listing is only ground truth if it does not tidy up. -# -# Facts *about* a file go in its header, never after its body. A note printed -# below the content is indistinguishable from content: annotated one file with a -# trailing `(no newline at end of file)` line and the next check script asserted -# the README's content ending in that sentence. -WORKSPACE_SNAPSHOT = ''' -import os +# ── Reward ───────────────────────────────────────────────────────────────── +# The proposing side's reward, unchanged from challenger/agentic.py where it was +# measured. Kept as free functions because nothing here has the state the method +# version read off self. + +# Peak and width of the pass-rate gaussian. This replaced R-Zero's +# ``1 - 2*|p - 1/2|`` for two reasons measured on run_clean9's 87 in-band +# proposals: that shape was not injective (with 8 rollouts its seven in-band +# values of n_pass mapped onto four rewards, so a task 1 of 8 solvers could do and +# one 7 of 8 could do were worth the same), and its signal was smaller than its +# noise (0.280 signal over 0.246 binomial noise, against 0.347 over 0.177 here). +# A peak below one half is also the more useful target: a proposal only teaches +# the solver something when the solver mostly cannot do it yet. +PASS_RATE_TARGET = 0.2 +PASS_RATE_WIDTH = 0.3 + +# How many phrases the keyword prompt's 'do not repeat these' line may quote. +# Measured on armA2ser: with 130 quoted the eighth refill call was still answering +# normally, with 150 it started inventing -- 'iRAPION holistic replace', 10 of 480 +# phrases that run. 100 sits below where that began. +AVOID_TOTAL = 100 + +# How often run() looks for a stall. Only ever reached when the run has already +# gone quiet, so it costs one wakeup per interval and nothing else. +STALL_CHECK_SECONDS = 30.0 + + +def novelty_factor(novelty: Optional[float], floor: float) -> float: + """What a proposal's difficulty score is multiplied by for its novelty. + + ``floor + (1 - floor) * N``. ``None`` returns 1.0, not the floor: it means + nobody judged this proposal, and charging it for a measurement that did not + happen would make the reward depend on API uptime. + """ + if novelty is None: + return 1.0 + n = min(1.0, max(0.0, float(novelty))) + return floor + (1.0 - floor) * n -root = {workspace!r} -skip = {{'.ms_agent', '__pycache__', '.ipynb_checkpoints', '.git'}} -rows = [] -for dirpath, dirnames, filenames in os.walk(root): - dirnames[:] = [d for d in dirnames if d not in skip] - for name in sorted(filenames): - path = os.path.join(dirpath, name) - try: - rows.append((os.path.relpath(path, root), os.path.getsize(path), path)) - except OSError: - pass -rows.sort() -for rel, size, _ in rows[:{max_files}]: - print(rel, size) - -budget = {total_budget} -for rel, size, path in rows[:{max_files}]: - if budget <= 0: - break - try: - with open(path, encoding='utf-8') as handle: - text = handle.read({per_file} + 1) - except (OSError, UnicodeDecodeError): - continue # binary or unreadable: the listing already names it - if '\\x00' in text: - continue - body = text[:{per_file}] - budget -= len(body) - # The trailing-newline count is stated for every file, both ways. Saying it - # only when it is absent made "this file ends with a newline" invisible, and - # the check writer then compared exact bytes without one: in ex9 two of the - # three checks that failed their own verification failed on exactly that -- - # the same reply asserted three files, guessed right on the two marked "no - # newline at end" and wrong on the unmarked one. - trailing = len(body) - len(body.rstrip(chr(10))) - if len(text) > len(body): - suffix = ' (first {per_file} bytes)' - elif trailing == 0: - suffix = ' (no newline at end)' - else: - suffix = ' (ends with %d newline character(s))' % trailing - print() - print('--- ' + rel + suffix + ' ---') - print(body, end='') - if not body.endswith(chr(10)): - print() -''' + +def challenger_reward(n_pass: Optional[int], rollouts: int, + novelty: Optional[float] = None, floor: float = 1.0) -> float: + """How close the solver came to the target pass rate, times novelty. + + ``None`` means the proposal never got as far as being solved and 0 means no + attempt passed. Both score 0, and that floor is load-bearing rather than + incidental: the gaussian at p=0 is 0.801, higher than the 0.607 it gives a + proposal half the attempts solve, so without the gate the best thing a proposer + could do is write tasks nobody can finish. + """ + if n_pass is None or not rollouts or n_pass <= 0: + return 0.0 + gap = n_pass / rollouts - PASS_RATE_TARGET + difficulty = math.exp(-(gap * gap) / (2.0 * PASS_RATE_WIDTH**2)) + return difficulty * novelty_factor(novelty, floor) + + +# ── Arguments ────────────────────────────────────────────────────────────── def parse_args(): p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - # Model + # What to collect. + p.add_argument('--keep-groups', type=int, default=8, + help='stop once this many groups have been kept. 8 groups x 8 ' + 'proposals = 64 proposing trajectories, and the selected ' + 'proposal of each x 8 attempts = 64 solving ones.') + p.add_argument('--group-size', type=int, default=8, + help='proposals sharing one keyword draw and one prompt: the ' + 'GRPO group on the proposing side.') + p.add_argument('--solver-rollouts', type=int, default=8, + help='attempts per candidate task: the GRPO group on the solving ' + 'side, and the denominator of n_pass.') + p.add_argument('--max-group-attempts', type=int, default=0, + help='give up after this many groups have been tried, kept or ' + 'not. 0 leaves the run governed by --keep-groups alone.') + + # Local model (the trainable half). p.add_argument('--model-id', default='ms://Qwen/Qwen3-4B') - p.add_argument('--template', default='Template', - help='template class in twinkle.template') + p.add_argument('--template', default='Template') p.add_argument('--sampler-gpus', type=int, default=4) - # Challenger backend: local vLLM by default, or an OpenAI-compatible API for - # the proposing side only (keywords + explore + check + statement). The - # solver side of the difficulty stage still runs through the local sampler, - # so --challenger-api is only usable with --solver-rollouts 0 (no local - # sampler is built at all in that case). The three connection args default - # to the LLM_BACKUP_* env vars the summarizer teacher already uses. - p.add_argument('--challenger-api', action='store_true', - help='propose through an OpenAI-compatible API instead of local vLLM; ' - 'requires --solver-rollouts 0') - p.add_argument('--challenger-api-model', - default=os.environ.get('LLM_BACKUP_MODEL', '')) - p.add_argument('--challenger-api-base', - default=os.environ.get('LLM_BACKUP_BASE_URL', '')) - p.add_argument('--challenger-api-key', - default=os.environ.get('LLM_BACKUP_API_KEY', '')) - p.add_argument('--challenger-concurrency', type=int, default=8, - help='parallel API conversations (API backend only)') - # Measured on qwen3.8-max at a ~15k-character exploration context: one turn - # took 58s with the default (unbounded) thinking and 10s at 2048, because the - # default spent ~5300 characters on reasoning per turn. 0 leaves the API's own - # default in place; anything else is sent as extra_body={'thinking_budget': N} - # on every proposing call (explore, check, statement, keyword generation). - p.add_argument('--challenger-thinking-budget', type=int, default=0, - help='cap reasoning tokens per API call; 0 = leave the API default') - # Split mode: explore on the local (trainable) vLLM model, but write the check - # script (success judgement) and the problem statement over the API instead of - # the local model. Only the exploration turns keep labels/logprobs and get - # trained; the two API stages are text-only and never enter the trajectory. - # Reuses the --challenger-api-* connection args, and is mutually exclusive with - # --challenger-api (which sends the whole proposing side to the API). - # On by default, because leaving it off is not a milder setting but a - # different experiment: run_clean6 was launched without it and the local 4B - # wrote the check scripts, which turned 0/55 check_parse_fail (qwen3-max) into - # 35/171 and 1/55 check_run_fail into 25/171. Pass --no-followup-api to write - # both stages with the local model on purpose. - p.add_argument('--followup-api', action=argparse.BooleanOptionalAction, - default=True, - help='explore locally (trainable) but generate the check script and ' - 'problem statement over --challenger-api-* (e.g. qwen3-max); ' - 'only the exploration part is trained. Not with --challenger-api.') - # How many episodes run at once, each in its own sandbox. An episode owns its - # workspace from the reset until its check has run, so this is also the number - # of sandboxes booted at startup. Default 96, from two measurements: on 8 vLLM - # workers, 96 requests in flight reached 28894 tok/s against a 33108 tok/s - # ceiling (87%), where 48 in flight reached only 15599 (47%); and the sandbox - # host booted 96 concurrent sandboxes with 0 failures at a p50 reset of 24.5s, - # up from 21.9s at 64. Above 96 the KV cache is the next limit -- vLLM reports - # room for about 24 sequences per worker, so roughly 192 in total. - p.add_argument('--episode-concurrency', type=int, default=96, - help='sandboxes to boot, and how many things run at once in both ' - 'stages: proposal episodes in flight, and solver attempts ' - 'per wave in the difficulty filter') - # 40960, up from 32768, because one proposal is now a single conversation: - # the tool-using turns, the check script and the problem statement all share - # this window. Measured on ex9's three separate calls, the worst case summed - # to about 25k tokens (12394 + 7643 + 2943 plus the appended messages), so - # this leaves room for episodes that take more steps than ex9's 2-5. - # - # 40960 and not more: it is Qwen3-4B's max_position_embeddings, and vLLM - # refuses to start above it -- 49152 was tried and rejected, since going past - # a RoPE model's trained positions produces nan rather than longer context. p.add_argument('--max-model-len', type=int, default=40960) - - # Generation control - p.add_argument('--keep-target', type=int, default=200, - help='how many tasks to keep; generation stops once reached') - p.add_argument('--batch-size', type=int, default=0, - help='tasks per yielded batch (0 = one batch of --keep-target)') - p.add_argument('--max-proposals-per-round', type=int, default=64, - help='max proposals per round (serial, so keep moderate)') - p.add_argument('--seed-file', default='', help='seed jsonl with query field') - p.add_argument('--seed-mix-prob', type=float, default=0.5) - - # Sampling params for the exploring stage (proposing) + p.add_argument('--gpu-memory-utilization', type=float, default=0.8) + + # The API model: check scripts, problem statements, keywords, rubric. + p.add_argument('--api-model', default=os.environ.get('LLM_BACKUP_MODEL', '')) + p.add_argument('--api-base', default=os.environ.get('LLM_BACKUP_BASE_URL', '')) + p.add_argument('--api-key', default=os.environ.get('LLM_BACKUP_API_KEY', '')) + p.add_argument('--api-concurrency', type=int, default=32, + help='API calls in flight. Only the rubric runs as its own job; ' + 'check and statement calls are made from inside a sandbox ' + 'job and are already capped by the slot count.') + p.add_argument('--api-thinking-budget', type=int, default=0, + help='sent as extra_body on every API call when > 0. Capping the ' + 'reasoning is the one knob that moved wall-clock: 58s -> 10s ' + 'per turn at 2048 on a ~15k-character context.') + + # Building: stage 1, the part that is trained. p.add_argument('--propose-temp', type=float, default=1.0) - # 8192, not 4096. At 4096, 3 of 12 explore episodes ended on the first turn - # with stop_reason=length and an untouched workspace: the model had written - # 15k, 16k and 10k characters of <think>, two of them without ever closing the - # tag, and one degenerating into a run of newlines. Nothing was dispatched, so - # those three cost a full episode each and produced no end state to write a - # check about. The trajectory ceiling and the engine's max_model_len are - # 32768, well above prompt plus this. p.add_argument('--propose-max-tokens', type=int, default=8192) - p.add_argument('--max-turns', type=int, default=24, - help='max tool-calling turns for the exploring stage') - # One call per reply, because the calls in one reply run *concurrently*: - # tool_manager.call_many hands a turn to Env.step_batch, which the sandbox - # server runs through ms-agent's parallel_call_tool. The model writes them in - # the order it means them to happen and gets none of the results, so a reply - # that writes a file and then reads it back reads the file as it was before. - # Measured on ex11: 13 of 36 episodes contain an observation that contradicts - # the end state -- read_file answering FileNotFound for a file the snapshot - # lists, glob answering with 0 files, `ls -R` missing a file written earlier - # in the same reply -- and 3 of the 4 kept tasks are among them. One of those - # kept tasks describes two files as "empty", which is what they were only - # because the call that filled them had not run yet. - # - # It also removes the other failure of a batched reply: 6 of 36 episodes - # spent the whole 8192-token budget on one reply holding 70 to 259 calls, - # the tail of it the same read_file over and over, and were discarded whole. - # A reply that can hold one call cannot do either. - p.add_argument('--one-call-per-reply', action='store_true', default=True, - help='stop generation at </tool_call> so each reply carries a single ' - 'call and the model sees its result before choosing the next') - p.add_argument('--no-one-call-per-reply', dest='one_call_per_reply', - action='store_false', - help='let a reply carry several calls, which then run concurrently') + p.add_argument('--max-turns', type=int, default=24) + p.add_argument('--max-build-files', type=int, default=4, + help='appends BUILD_SIZE_CAP to the system prompt, capping how ' + 'many files one build may leave behind. 0 removes the cap. ' + 'This changes the prompt, so it changes what is trained.') p.add_argument('--stop-after-stuck-turns', type=int, default=2, - help='end an episode after this many consecutive turns that made no ' - 'progress; 0 runs to --max-turns regardless. A turn counts as ' - 'stuck when every call in it came back an error, or every call ' - 'in it was byte-identical to one already made in the episode. ' - 'Replayed over 12 recorded episodes: errors alone would stop 1 ' - 'of 12 and save 9 of 239 calls, since the worst offenders mix a ' - 'failing call with a glob that succeeds; adding the repeat rule ' - 'stops 3 of 12 and saves 63 calls, and those 3 are exactly the ' - 'ones that spent 54, 84 and 17 calls to leave a script that ' - 'could not run.') - - # Problem statement - p.add_argument('--problem-max-chars', type=int, default=8192) - p.add_argument('--check-retries', type=int, default=1, - help='How many times a check script that fails is handed back, ' - 'with the traceback and the workspace listing, to be ' - 'rewritten. ex12 lost 36 of 72 proposals here, and 29 of ' - 'those were one assertion naming a value the model had ' - 'never read -- a row count, a nearly-right content ' - 'string, a timestamp -- on a workspace state that was ' - 'fine. 0 rejects on the first failure, as ex9-ex12 did.') - # Budgets for the two stages appended to the episode. Separate numbers - # because the two are not alike: writing the checks reads the whole episode - # plus the end state and reasons at length (ex9's largest such reply was 7643 - # trainable tokens, so 4096 would cut the tail off and the proposal would be - # discarded as unparseable), while the statement is prose and ex9's largest - # was 2943. + help='end the tool phase after this many turns that repeated a ' + 'call and changed nothing. 0 turns it off.') + p.add_argument('--one-call-per-reply', action=argparse.BooleanOptionalAction, + default=True, + help="stop generation at '</tool_call>' so a reply carries exactly " + 'one call. The stop string is kept in the output, or every ' + 'turn would train on an unclosed block.') + + # Stages 2 and 3, over the API. p.add_argument('--check-max-tokens', type=int, default=8192) + p.add_argument('--check-retries', type=int, default=1, + help='rewrites offered to a check script that does not parse or ' + 'does not pass on the state it was written from.') p.add_argument('--problem-max-tokens', type=int, default=4096) + p.add_argument('--problem-max-chars', type=int, default=8192, + help='a statement longer than this is thrown away: it is quoting ' + 'the workspace rather than describing the task.') + + # Solving. + p.add_argument('--solver-temp', type=float, default=1.0) + p.add_argument('--solver-max-tokens', type=int, default=8192) + p.add_argument('--solver-max-turns', type=int, default=24) - # Keywords + # Keywords. p.add_argument('--keywords-n', type=int, default=128, - help='per-category refill target; 0 disables keyword bank') - p.add_argument('--keyword-db', default='output/rsi_agentic/keywords.jsonl') - p.add_argument('--keyword-gen-calls', type=int, default=8) - # How many of a refill's generating calls go out together. 1 means each is - # told what the ones before it produced; the first round of arm measurements - # effectively ran at 8, where the whole first refill went out with an empty - # 'do not repeat' list and came back with synonyms of each other. - p.add_argument('--keyword-refill-concurrency', type=int, default=1) - p.add_argument('--keyword-refill-tries', type=int, default=2) + help='how many keywords a dry category is refilled with.') + p.add_argument('--keyword-gen-calls', type=int, default=8, + help='calls a refill is split over, run one at a time so each can ' + 'be told what the ones before it already said.') + p.add_argument('--keyword-refill-tries', type=int, default=2, + help='refill rounds before a dry category is recycled, i.e. every ' + 'keyword in it marked unused again. Without this a bank the ' + 'model has run out of new ideas for ends the run.') + p.add_argument('--keyword-expand', action=argparse.BooleanOptionalAction, + default=True, + help='after the last group, ask for more keywords in the domains ' + 'that produced tasks nobody solved, and write them to the ' + 'bank the next iteration reads.') p.add_argument('--keyword-temp', type=float, default=1.3) - # 1024 measured 8 of 24 generation calls cut off at the budget with nothing - # parseable: the model spends most of it listing candidates inside <think>, - # rewrites the list two or three times, and the JSON array afterwards gets - # severed mid-string. The successful calls landed just under 1024, so the cap - # sat inside the distribution of working replies rather than beyond it. p.add_argument('--keyword-max-tokens', type=int, default=4096) - p.add_argument('--single-kw-prob', type=float, default=0.1) - p.add_argument('--combo-arity', default='triple', choices=['triple', 'mix']) - p.add_argument('--arity-weights', default='', - help="'w1,w2,w3' for --combo-arity mix (empty = uniform)") - # How many proposals answer each keyword draw. Above 1 they share one prompt - # and one group id, which is what the proposing side needs to have a group to - # compute an advantage over -- at 1 every group has one member and every - # advantage is zero. It does not change the compute at a fixed - # --max-proposals-total; it divides the number of distinct keyword draws by - # the same factor, so 216 proposals come from 27 draws at 8 instead of 216. - p.add_argument('--proposals-per-group', type=int, default=1, - help='proposals sharing one keyword draw and prompt (1 = no groups, ' - 'so no proposer advantage)') - - # Difficulty filter - # 8 attempts, keeping 2-6: with 4 attempts the band was 1-3 and ex9's - # measured pass counts came out {0: 6, 1: 1, 4: 2} -- two-thirds of the - # tasks landed on an end of the range where one attempt either way changes - # the verdict. 8 costs twice the sandbox time per task and puts the kept - # band around one third of attempts passing. - p.add_argument('--solver-rollouts', type=int, default=8) - p.add_argument('--solver-temp', type=float, default=1.0) - # Same 8192 as the explore round, and for the same measured reason: at 4096, - # 15 of 50 solver attempts ended on stop_reason=length with an untouched - # workspace, and one task lost all four of its attempts that way and was - # discarded as too hard. Raising it took that to 0 of 20. It has to stay in - # step with --propose-max-tokens: a task the proposer needed room to build is - # not solvable in less. - p.add_argument('--solver-max-tokens', type=int, default=8192) - p.add_argument('--solver-max-turns', type=int, default=24, - help='NOT WIRED: no solver_explorer is passed, so solver attempts ' - 'run through the same rollout as the proposing episodes and ' - 'obey --max-turns. Kept so the value can be set once the two ' - 'are separated; changing it alone has no effect.') - p.add_argument('--keep-min-pass', type=int, default=2) - p.add_argument('--keep-max-margin', type=int, default=2) - - # Sandbox - p.add_argument('--sandbox-template', default='', - help='AgentENV/e2b template name (required)') - p.add_argument('--sandbox-api-url', default='', - help='AgentENV server URL (or AENV_API_URL env var)') - p.add_argument('--agent-config', default='cookbook/rsi/agentic/rsi_agent.yaml', - help='ms-agent yaml for the sandbox tool server') + + # Novelty. + p.add_argument('--task-bank', default='', + help='jsonl of statements from earlier iterations. Empty turns ' + 'novelty off, and the reward is the pass-rate gaussian alone.') + p.add_argument('--task-bank-refs', type=int, default=5, + help='stored statements shown to the judge, on top of the group ' + "'s own siblings, which are always shown.") + # 1.0 leaves the novelty term at exactly 1.0, so a proposal is scored on its + # pass rate alone while the rubric still runs and still writes + # novelty_scores.jsonl. See loop.sh for the measurement that set it there. + p.add_argument('--novelty-floor', type=float, default=1.0) + p.add_argument('--novelty-tries', type=int, default=3, + help='attempts to get a verdict for a group. After the last one ' + 'the group is dropped and its pending solver jobs skipped.') + + # Sandbox. + p.add_argument('--sandbox-slots', type=int, default=32, + help='microVMs, i.e. how many jobs run at once. One job owns one ' + 'slot from the workspace clear to its last check.') + p.add_argument('--sandbox-template', default=os.environ.get('AENV_TEMPLATE', '')) + p.add_argument('--sandbox-api-url', default=os.environ.get('AENV_API_URL', '')) p.add_argument('--sandbox-timeout', type=int, default=900) - p.add_argument('--workspace', default='/workspace', - help='working directory inside the sandbox') - p.add_argument('--snapshot-max-files', type=int, default=50, - help='files listed in the end-state snapshot') - p.add_argument('--snapshot-per-file', type=int, default=600, - help='bytes of each file shown to the check writer') - p.add_argument('--snapshot-budget', type=int, default=6000, - help='total bytes of file content in the snapshot') - - # Output - # ---- Experiment arms. Each isolates one measured failure and can be used - # alone or stacked. Off by default, so an unflagged run is the old behaviour. - # - # A: ex11/ex12/ex13 statements quoted their own answer, because stage 3 asks - # for the full end state while the solver starts empty -- the only way to say - # what a derived file holds is to write out what was computed. Difficulty came - # out 8/8 or 0/8. This makes the statement give input data verbatim and - # everything derived as a rule. - # C: apitest4's statements each wanted an 8-12 file package with a CLI, and - # Qwen3-4B passed 0 of 96 -- 32 attempts spent the whole token budget typing - # source, 64 declared success with the files unwritten. - p.add_argument('--max-build-files', type=int, default=0, - help='arm C: cap the episode at this many files, no package, ' - 'no CLI with subcommands (0 = no cap)') - # For measuring a configuration rather than filling a dataset: two arms are - # only comparable when given the same number of tries. - p.add_argument('--max-proposals-total', type=int, default=0, - help='stop after this many proposals regardless of keep-target ' - '(0 = run until keep-target)') + p.add_argument('--agent-config', default='cookbook/rsi/agentic/rsi_agent.yaml') + p.add_argument('--workspace', default='/workspace') + p.add_argument('--snapshot-max-files', type=int, default=50) + p.add_argument('--snapshot-per-file', type=int, default=600) + p.add_argument('--snapshot-budget', type=int, default=6000) + + # Output. + p.add_argument('--out-dir', default='output/rsi_agentic') + p.add_argument('--keyword-db', default='', + help='defaults to <out-dir>/keywords.jsonl') p.add_argument('--random-seed', type=int, default=0) - p.add_argument('--out-flows', default='output/rsi_agentic/challenge_flows.jsonl') - p.add_argument('--dump-rejected', default='output/rsi_agentic/challenge_rejected.jsonl') - p.add_argument('--dump-propose-traj', default='output/rsi_agentic/propose_traj', - help='directory for the proposing rounds (token ids + logprobs, one npz ' - 'per attempt plus index.jsonl). Empty string turns it off; keeping ' - 'it is what leaves the door open to training the challenger itself.') - p.add_argument('--dump-solver-attempts', - default='output/rsi_agentic/solver_attempts.jsonl', - help='one line per difficulty-stage solver attempt: the statement, the ' - 'check script, the attempt, the state it left and what the check ' - 'said. Without it a task measured 0 of 4 gives no way to tell an ' - 'impossible task from a statement that withholds what the check ' - 'demands. Empty string turns it off.') - p.add_argument('--no-sort-by-difficulty', action='store_true') - p.add_argument('--stage', default='all', choices=['all', 'keywords', 'explore'], - help="'keywords' runs step 1 only -- fill the keyword bank, draw " - 'the combinations, write the proposal prompts they produce to ' - '--out-flows, and exit without touching the sandbox. ' - "'explore' adds steps 2-4: clear the workspace, run the " - 'sandbox episode, snapshot the end state, and stop before the ' - 'check-writing round. Both exist because a stage that is ' - 'broken cannot be diagnosed from the far end of an ' - 'hours-long full run.') - p.add_argument('--stage-proposals', type=int, default=16, - help='how many proposals --stage keywords or --stage explore runs') - p.add_argument('--dump-explore', default='output/rsi_agentic/explore_episodes.jsonl', - help='one line per --stage explore episode: the prompt, every ' - 'message, every tool call and its observation, and the end ' - 'state the snapshot saw. Empty string turns it off.') - p.add_argument('--dump-keyword-gen', - default='output/rsi_agentic/keyword_gen.jsonl', - help='one line per keyword-generation call: the prompt, the raw ' - 'reply, and what the parser made of it. Without it a bank that ' - 'stays empty gives no way to tell a disobedient model from a ' - 'parser that rejects valid output. Empty string turns it off.') - return p.parse_args() - - -def build_env(args): - """Create the long-lived sandbox environment.""" - template = args.sandbox_template or os.environ.get('AENV_TEMPLATE', '') - api_url = args.sandbox_api_url or os.environ.get('AENV_API_URL', '') - if not template: - raise SystemExit('[challenge] --sandbox-template or AENV_TEMPLATE is required') - if not api_url: - raise SystemExit('[challenge] --sandbox-api-url or AENV_API_URL is required') - - env = RemoteMsAgentToolEnv( - template=template, + args = p.parse_args() + if not args.api_model or not args.api_base: + raise SystemExit('[challenge] --api-model and --api-base are required ' + '(or LLM_BACKUP_MODEL / LLM_BACKUP_BASE_URL)') + if args.solver_rollouts < 2: + raise SystemExit('[challenge] --solver-rollouts must be >= 2: it is both the ' + "solver side's GRPO group size and the denominator n_pass is " + 'judged against') + if args.group_size < 2: + raise SystemExit('[challenge] --group-size must be >= 2: a group of one has ' + 'no mean to subtract, so every advantage is zero') + args.keyword_db = args.keyword_db or os.path.join(args.out_dir, 'keywords.jsonl') + return args + + +# ── Resources ────────────────────────────────────────────────────────────── + + +def initialize_device(args) -> Tuple[Any, Any]: + """Bring up Ray and the local vLLM sampler; returns (sampler, template). + + The template is built here as well as inside the sampler because the rollout + encodes with it locally: one object, so the token ids the sampler continues + from are the ids the trajectory was encoded with. + """ + twinkle.initialize( + mode='ray', nproc_per_node=args.sampler_gpus, lazy_collect=False, + groups=[DeviceGroup(name='sampler', ranks=list(range(args.sampler_gpus)), + device_type='GPU')]) + sampler = vLLMSampler( + model_id=args.model_id, + engine_args={'gpu_memory_utilization': args.gpu_memory_utilization, + 'max_model_len': args.max_model_len}, + device_mesh=DeviceMesh.from_sizes(world_size=args.sampler_gpus, + dp_size=args.sampler_gpus), + remote_group='sampler', + ) + sampler.set_template(args.template, model_id=args.model_id, enable_thinking=True, + max_length=args.max_model_len) + import twinkle.template as template_module + template = getattr(template_module, args.template)( + args.model_id, max_length=args.max_model_len, enable_thinking=True) + if not getattr(type(sampler).sample, '_enable_continous_work', False): + raise SystemExit( + '[challenge] this sampler does not route requests one at a time ' + '(sample lacks enable_continous_work), so a batch of one would be ' + 'padded to the worker count and most of every generation thrown away. ' + 'The whole design here is one trajectory per request.') + return sampler, template + + +def initialize_sandbox(args) -> List[Any]: + """Boot the slots. See ``sandbox.open_pool``.""" + return open_pool( + args.sandbox_slots, + template=args.sandbox_template, + api_url=args.sandbox_api_url, config_path=args.agent_config, - api_url=api_url, workspace=args.workspace, sandbox_timeout=args.sandbox_timeout, + snapshot_max_files=args.snapshot_max_files, + snapshot_per_file=args.snapshot_per_file, + snapshot_budget=args.snapshot_budget, ) - env.reset() - return env -def main(): - args = parse_args() - for path in (args.out_flows, args.dump_rejected, args.keyword_db, - args.dump_keyword_gen, args.dump_explore): - if path: - os.makedirs(os.path.dirname(os.path.abspath(path)) or '.', exist_ok=True) - - # Build the proposing backend: an OpenAI-compatible API, or a local vLLM - # sampler. The API path skips twinkle.initialize entirely -- it needs no GPUs - # -- and passes no template, since the API re-sends messages as text rather - # than splicing token ids the way the sampler continuation does. - use_api = args.challenger_api - if args.followup_api: - # Split mode needs the local sampler for exploration (that is the trainable - # half), so it cannot run under --challenger-api, which builds no sampler. - if use_api: - raise SystemExit('[challenge] --followup-api and --challenger-api are mutually ' - 'exclusive: --followup-api explores on the local sampler and ' - 'sends only the check/statement stages to the API, while ' - '--challenger-api sends the whole proposing side to the API.') - if not args.challenger_api_model or not args.challenger_api_base: - raise SystemExit('[challenge] --followup-api needs --challenger-api-model and ' - '--challenger-api-base (or LLM_BACKUP_MODEL / ' - 'LLM_BACKUP_BASE_URL).') - if use_api: - if args.solver_rollouts: - raise SystemExit('[challenge] --challenger-api needs --solver-rollouts 0: the ' - 'solver side still runs on the local sampler, which is not ' - 'built in API mode.') - if not args.challenger_api_model or not args.challenger_api_base: - raise SystemExit('[challenge] --challenger-api needs --challenger-api-model and ' - '--challenger-api-base (or LLM_BACKUP_MODEL / LLM_BACKUP_BASE_URL).') - from twinkle_agentic.protocol.openai import OpenAI - backend = OpenAI(model=args.challenger_api_model, - api_key=args.challenger_api_key or None, - base_url=args.challenger_api_base) - template = None - logger.info(f'[challenge] proposing via API model={args.challenger_api_model} ' - f'base={args.challenger_api_base}') - else: - # Initialize twinkle (sampler only, no trainer) - twinkle.initialize( - mode='ray', nproc_per_node=args.sampler_gpus, lazy_collect=False, - groups=[DeviceGroup(name='sampler', ranks=list(range(args.sampler_gpus)), - device_type='GPU')]) - backend = vLLMSampler( - model_id=args.model_id, - engine_args={'gpu_memory_utilization': 0.8, 'max_model_len': args.max_model_len}, - device_mesh=DeviceMesh.from_sizes(world_size=args.sampler_gpus, - dp_size=args.sampler_gpus), - remote_group='sampler', - ) - backend.set_template(args.template, model_id=args.model_id, enable_thinking=True, - max_length=args.max_model_len) - - import twinkle.template as template_module - template = getattr(template_module, args.template)( - args.model_id, max_length=args.max_model_len, enable_thinking=True) - - # Build sandbox environments -- one per episode slot, since an episode owns - # its workspace from the reset until its check has run and two episodes - # sharing a sandbox would read each other's files. Skipped for - # --stage keywords: that stage only brainstorms and draws keywords, and - # booting a microVM to do it would make checking step 1 depend on the one part - # of the setup most likely to be down. - # - # ``envs[0]`` is also the one the serial paths use (the difficulty stage, and - # --stage explore), so ``env`` stays a name for it. - envs = [] - env = None - schemas = None - tool_manager = ToolManager() - episode_tool_managers = None - if args.stage != 'keywords': - n_slots = max(1, args.episode_concurrency) - if n_slots == 1: - envs = [build_env(args)] - else: - # Booted in parallel: each is a microVM taking ~10s, and doing eight - # of them one after another would put a minute and a half in front of - # every run. - from concurrent.futures import ThreadPoolExecutor - with ThreadPoolExecutor(max_workers=n_slots) as pool: - envs = list(pool.map(lambda _: build_env(args), range(n_slots))) - env = envs[0] - schemas = env.tool_schemas() - # One ToolManager per sandbox: the tools carry the env they dispatch into, - # so slot i's model turns have to go through slot i's manager. - episode_tool_managers = [ToolManager(EnvTool.from_schemas(e, schemas)) for e in envs] - tool_manager = episode_tool_managers[0] - logger.info(f'[challenge] {len(envs)} sandbox(es) ready ' - f'(episode concurrency {n_slots})') - - # Explorer: multi-turn rollout with sandbox tools - # ``stop`` ends the reply at the end of the first tool call, and - # ``include_stop_str_in_output`` keeps that '</tool_call>' in what the policy - # is trained on -- vLLM drops the matched stop by default, which would train - # every turn to end on an unclosed block. - explore_stop = ['</tool_call>'] if (args.one_call_per_reply and not use_api) else None - # Sent on every API call when set. Capping the reasoning is the one knob that - # moved the wall-clock: 58s -> 10s per turn at 2048 on a ~15k-character context. - api_extra_body = ({'thinking_budget': args.challenger_thinking_budget} - if (use_api and args.challenger_thinking_budget > 0) else None) - if api_extra_body: - logger.info(f'[challenge] thinking_budget={args.challenger_thinking_budget} ' - f'on every API call') - explore_params = SamplingParams(max_tokens=args.propose_max_tokens, num_samples=1, - logprobs=1, temperature=args.propose_temp, top_p=0.95, - stop=explore_stop, - include_stop_str_in_output=bool(explore_stop)) - # One tool call per reply and stuck-turn early stop are sampler-path features: - # the API dispatches native tool_calls (never a '</tool_call>' string) and - # APIMultiTurnRollout takes neither kwarg. - if use_api: - explorer = build_rollout( - backend, tool_manager=tool_manager, max_turns=args.max_turns, - concurrency=args.challenger_concurrency, sampling_params=explore_params, - extra_body=api_extra_body) - else: - explorer = build_rollout( - backend, template=template, tool_manager=tool_manager, - max_turns=args.max_turns, stop_after_stuck_turns=args.stop_after_stuck_turns, - sampling_params=explore_params) - - # Keyword brainstorming runs through this one instead of the sandbox - # explorer: a list is a text answer, and a bracketed list in a reply is - # exactly what the sandbox explorer would try to dispatch as a call. - # - # max_turns=1 is what makes it tool-less: MultiTurnRollout ends the - # trajectory on the turn limit before it dispatches anything, so the empty - # ToolManager below is never consulted. It is here because the rollout - # requires one at construction, not because these calls have tools. - keyword_params = SamplingParams(max_tokens=args.keyword_max_tokens, num_samples=1, - logprobs=1, temperature=args.keyword_temp, top_p=0.98) - if use_api: - keyword_explorer = build_rollout( - backend, tool_manager=ToolManager(), max_turns=1, - concurrency=args.challenger_concurrency, sampling_params=keyword_params, - extra_body=api_extra_body) - else: - keyword_explorer = build_rollout( - backend, - template=template, - tool_manager=ToolManager(), - max_turns=1, - sampling_params=keyword_params, - ) - - # Sandbox control functions -- use env.runner() which resolves tool names - # (ms-agent registers tools as "server---name") and parses exit codes from - # the marker protocol, so we don't rely on string matching. - # - # One runner per sandbox; ``slot`` picks which one. The challenger passes the - # slot of the episode it is serving, so a concurrent episode never clears or - # inspects another episode's workspace. Everything serial (the difficulty - # stage, --stage explore) leaves it at the default and uses sandbox 0. - # - # Empty for --stage keywords, which has no sandbox. The functions below index - # into it and would raise if that stage ever reached them; it returns first. - runners = [e.runner() for e in envs] - runner = runners[0] if runners else None - - def reset_fn(slot: int = 0): - """Empty sandbox ``slot``'s workspace before an episode. - - Raises rather than returning: every caller depends on a clean start, and - a silent no-op here means a task inherits the previous task's files -- - which lets a solver pass without doing anything and makes the difficulty - numbers meaningless. - - This is also the one point where losing the sandbox costs nothing, since - the workspace is about to be emptied regardless -- so a runtime that went - away is rebuilt here instead of ending a run that may have hours of - proposals behind it. The same reasoning covers a clear that *fails* on a - runtime still answering /health, which ensure_ready cannot see: run 'rsi' - reached iteration 7 -- six checkpoints, eight hours -- and ended on three - clears timing out at ms-agent's per-call limit while the sandbox itself - was healthy enough to report them. So the clear is retried, and then - retried on a deliberately rebuilt sandbox, before the run is given up. +def rollout_one(rollout: MultiTurnRollout, traj: Dict[str, Any], + params: SamplingParams, slot) -> Optional[Dict[str, Any]]: + """Run one trajectory: vLLM for the replies, ``slot`` for the tool calls. + + A batch of one. The sampler routes it to whichever worker is free, so this is + called from every sandbox thread at once and the requests share vLLM's batch + without any of them waiting for the others to be ready. + """ + out = rollout([traj], sampling_params=params, tool_manager=slot.tool_manager) + return out[0] if out else None + + +def call_one(slot, script: str) -> Tuple[int, str]: + """Run one python script inside ``slot``; returns (exit code, output).""" + return slot.run(script) + + +def api_one(api, messages: List[Dict[str, Any]], user_text: str, + params: SamplingParams, extra_body: Optional[Dict[str, Any]] = None) -> Optional[str]: + """Append ``user_text`` and one API reply to ``messages``; returns the reply. + + ``messages`` is the caller's private copy, never a trainable trajectory, so + mutating it in place costs the model nothing. ``None`` means the call raised: + the caller rejects rather than building a task on a broken conversation. + + Tools are withdrawn for these stages on purpose -- they are answers, not + actions -- so only the text is kept and any structured ``tool_calls`` the API + returned are dropped. + """ + messages.append({'role': 'user', 'content': user_text}) + request = {'messages': messages} + try: + reply = api(request, params, extra_body=extra_body) if extra_body else api(request, params) + except Exception as e: # noqa: BLE001 -- one bad call must not kill the run + logger.warning(f'[challenge] API call failed: {type(e).__name__}: {e}') + return None + if isinstance(reply, list): + reply = reply[0] if reply else {} + content = (reply.get('content') if isinstance(reply, dict) else None) or '' + messages.append({'role': 'assistant', 'content': content}) + return content + + +# ── Output ───────────────────────────────────────────────────────────────── + + +def logprob_column(logprobs: Any) -> List[float]: + """One float per generated token: the logprob of the token that was chosen. + + The sampler hands these over as ``List[List[Tuple[int, float]]]`` -- per + generated token, a list of top-k ``(token_id, logprob)`` pairs with the chosen + token first (``SampledSequence.logprobs``, data_format/sampling.py:185). + Passing that to ``np.asarray`` directly would store an ``(N, k, 2)`` array and + the loader would hand GRPO nested lists where it wants one float per trainable + token -- which is a crash inside the step, or worse a silent reshape. + + A plain list of floats is accepted too, for a sampler that already flattened. + Anything else raises rather than being coerced: a wrong ``old_logps`` makes the + GRPO ratio wrong on the first step, and nothing downstream would say so. + """ + out: List[float] = [] + for step in logprobs: + if isinstance(step, (int, float)): + out.append(float(step)) + continue + if isinstance(step, (list, tuple)) and step: + head = step[0] + if isinstance(head, (list, tuple)) and len(head) >= 2: + out.append(float(head[1])) + continue + raise TypeError(f'cannot read a logprob out of {step!r}; expected a float ' + f'or a list of (token_id, logprob) pairs') + return out + + +class Recorder: + """Everything a run writes, behind one lock. + + Trajectories go to ``.npz`` for the token fields and to ``index.jsonl`` for + everything a reader needs to interpret them. The text is written in full and + never truncated: these files are read to check whether a reward was deserved, + which a shortened statement cannot answer. + """ + + def __init__(self, out_dir: str): + self.dir = out_dir + self.traj_dir = os.path.join(out_dir, 'trajs') + os.makedirs(self.traj_dir, exist_ok=True) + self._lock = threading.Lock() + self._n = 0 + self._index = open(os.path.join(self.traj_dir, 'index.jsonl'), 'w', encoding='utf-8') + self._groups = open(os.path.join(out_dir, 'groups.jsonl'), 'w', encoding='utf-8') + self._tasks = open(os.path.join(out_dir, 'tasks.jsonl'), 'w', encoding='utf-8') + # Why a build produced no task. The reason alone is not diagnosable: nine + # empty_workspace rejections in one run all looked like the model refusing + # to act, and the question of whether it had run out of tokens or simply + # emitted no call could not be answered from the record, because the fields + # that answered it were on the trajectory and were dropped. + self._rejected = open(os.path.join(out_dir, 'rejected.jsonl'), 'w', encoding='utf-8') + # Keyword replies, both sides in full. The one question this file exists to + # answer -- did the model disobey the format, or does the parser reject what + # it produced -- cannot be answered from a count. Keyword generation was + # silently broken for whole runs when the prompt asked for one per line and + # the parser wanted a JSON array. + self._keywords = open(os.path.join(out_dir, 'keyword_gen.jsonl'), 'w', encoding='utf-8') + # Every solver attempt, passed or not, with the state it left and what the + # check said about it. A task measured at 0 of 8 has three explanations -- + # the check is wrong, the statement withholds something the check demands, + # or the solver gave up -- and only the attempt and the workspace it left + # tell them apart. Written for every attempt, not only for the ones that + # end up trained on: the failures are what this file is for. + self._attempts = open(os.path.join(out_dir, 'solver_attempts.jsonl'), 'w', + encoding='utf-8') + # The rubric, all three of its dimensions. Only novelty reaches a reward; + # usefulness and complexity are recorded so the question of whether they + # should count can be answered from a run instead of argued. + self._novelty = open(os.path.join(out_dir, 'novelty_scores.jsonl'), 'w', + encoding='utf-8') + + def trajectory(self, traj: Dict[str, Any], **fields: Any) -> None: + """One training sample: token fields to npz, everything else to the index. + + A trajectory with no ``logprobs`` is written anyway, with the field left + null. It is not trainable and the loader will say so -- which is the point: + a sample silently dropped here would make the group it belongs to look like + a different size than it was. + """ + input_ids = np.asarray(traj.get('input_ids') or [], dtype=np.int32) + labels = np.asarray(traj.get('labels') or [], dtype=np.int32) + logprobs = traj.get('logprobs') + with self._lock: + self._n += 1 + name = f'{self._n:06d}.npz' + arrays = {'input_ids': input_ids, 'labels': labels} + if logprobs is not None: + # float64, and the chosen token's column only. These are the old_logps a + # GRPO step divides by; float32 would round them to about 7 digits, so + # the ratio exp(logp - old_logp) would be off by roughly 1e-7 for + # reasons that have nothing to do with the policy having changed. + arrays['logprobs'] = np.asarray(logprob_column(logprobs), dtype=np.float64) + # Compressed: a 24-turn agentic episode is tens of thousands of token ids, + # and 128 of them per iteration adds up on disk. + np.savez_compressed(os.path.join(self.traj_dir, name), **arrays) + record = dict(fields) + record.update({ + 'npz': name, + 'n_tokens': int(input_ids.size), + 'n_trainable': int((labels != -100).sum()) if labels.size else 0, + 'has_logprobs': logprobs is not None, + # The rollout guarantees one logprob per trainable label; recorded so a + # loader can check it rather than trust it. + 'n_logprobs': int(arrays['logprobs'].size) if logprobs is not None else 0, + 'turns': traj.get('turns'), + 'stop_reason': traj.get('stop_reason'), + 'truncated': bool(traj.get('truncated')), + 'tool_stop': traj.get('tool_stop'), + 'messages': traj.get('messages') or [], + }) + self._write(self._index, record) + + def group(self, record: Dict[str, Any]) -> None: + self._write(self._groups, record) + + def task(self, record: Dict[str, Any]) -> None: + self._write(self._tasks, record) + + def rejected(self, record: Dict[str, Any]) -> None: + self._write(self._rejected, record) + + def keywords(self, record: Dict[str, Any]) -> None: + self._write(self._keywords, record) + + def attempt(self, record: Dict[str, Any]) -> None: + self._write(self._attempts, record) + + def novelty(self, record: Dict[str, Any]) -> None: + self._write(self._novelty, record) + + def close(self) -> None: + for handle in (self._index, self._groups, self._tasks, self._rejected, + self._keywords, self._attempts, self._novelty): + handle.close() + + def _write(self, handle, record: Dict[str, Any]) -> None: + line = json.dumps(record, ensure_ascii=False, default=str) + with self._lock: + handle.write(line + '\n') + handle.flush() + + +# ── Group state ──────────────────────────────────────────────────────────── + + +@dataclass +class Proposal: + """One trajectory's worth of state, from the build to its solver attempts.""" + + group: 'Group' + idx: int + outcome: str = '' # 'ok', or why this one produced no task + detail: str = '' # what to look at when it did not + statement: str = '' + check: str = '' + traj: Optional[Dict[str, Any]] = None # the trainable build trajectory + attempts: List[Dict[str, Any]] = field(default_factory=list) + passes: List[bool] = field(default_factory=list) + novelty: Optional[float] = None + n_solved: int = 0 # attempts finished, not attempts passed + + @property + def n_pass(self) -> Optional[int]: + """How many attempts passed, or None if the task was never measured.""" + if not self.statement or self.n_solved < self.group.rollouts: + return None + return sum(1 for p in self.passes if p) + + def reward(self, rollouts: int, floor: float) -> float: + return challenger_reward(self.n_pass, rollouts, self.novelty, floor) + + +class Group: + """``size`` proposals sharing one keyword draw, and the counters that decide them. + + Every method that reads more than one field takes the lock, because the + proposals resolve on different threads and in any order. Nothing here blocks: + a thread reports what it finished and asks whether that was the last thing + outstanding, and only the thread that gets ``True`` runs the decision. + """ + + def __init__(self, gid: int, keywords: List[Tuple[str, str]], keyword_block: str, + prompt: str, size: int, rollouts: int): + self.id = gid + self.keywords = keywords + self.keyword_block = keyword_block + self.prompt = prompt + self.rollouts = rollouts + self.proposals = [Proposal(self, i) for i in range(size)] + self.lock = threading.Lock() + self.n_built = 0 # proposals whose build stage is over + self.rubric_done = False + self.dropped = '' # reason, once this group is abandoned + self.decided = False + + @property + def size(self) -> int: + return len(self.proposals) + + def abandon(self, reason: str) -> bool: + """Give up on this group. True if this call is the one that decided it. + + Jobs already queued for it check ``dropped`` and return their slot without + doing any work, so abandoning is also how the remaining solver attempts of + a group are cancelled. """ - clear = CLEAR_WORKSPACE.format(workspace=args.workspace) - if envs[slot].ensure_ready(): - # A rebuilt sandbox starts empty, so the clear below is redundant, - # but running it anyway keeps one path through this function. The - # rebuild replaces the sandbox behind this env, so the runner is - # re-fetched rather than reused. - runners[slot] = envs[slot].runner() - logger.warning(f'[challenge] sandbox {slot} was rebuilt before this episode') - exit_code, output = runners[slot](clear, 'python') - if exit_code != 0: - logger.warning(f'[challenge] workspace reset failed on sandbox {slot} ' - f'(exit {exit_code}), retrying in {RESET_RETRY_WAIT}s: ' - f'{output[-200:]}') - time.sleep(RESET_RETRY_WAIT) - exit_code, output = runners[slot](clear, 'python') - if exit_code != 0: - # Rebuilt rather than retried again: two failures in a row is not the - # transient this waits out, and a fresh sandbox brings a workspace - # that is already empty -- which is all this function is asked for. - # Counted as a recovery so the run's own tally at the end still - # accounts for every rebuild, including the ones from here. - logger.warning(f'[challenge] workspace reset failed twice on sandbox {slot} ' - f'(exit {exit_code}); rebuilding it: {output[-200:]}') - envs[slot].reset() - envs[slot].n_recoveries += 1 - runners[slot] = envs[slot].runner() - exit_code, output = runners[slot](clear, 'python') - if exit_code != 0: - raise RuntimeError(f'workspace reset failed (exit {exit_code}): {output[-400:]}') - - def run_check_fn(script: str, slot: int = 0): - """Run a python check script in sandbox ``slot``; returns (exit_code, output).""" - return runners[slot](script, 'python') - - # Why the last snapshot for each slot came back empty: the failure text if the - # listing could not be read, absent if the workspace really was empty. Read by - # snapshot_error_fn below so a paused sandbox is not filed as the model having - # built nothing. - snapshot_errors: Dict[int, str] = {} - - def workspace_snapshot_fn(slot: int = 0): - """Every file the episode left behind: ``path size`` lines, then contents. - - This is the ground truth the check script is written against, so it is - unwrapped from the tool's JSON envelope and returned as a bare listing: - the model has to be able to read it as a directory rather than as a tool - result, or it falls back on what it *believes* it created. - - Returns an empty string when the episode left nothing behind, and also - when the listing could not be read at all. Both mean the same thing to - the caller -- there is no end state to write checks about -- and neither - may be dressed up as a plausible one: a snapshot that says "empty" - when it means "I could not look" produces tasks whose only true - assertion is that nothing happened. Which of the two it was is recorded - in ``snapshot_errors`` instead, for the rejection to be filed under. + with self.lock: + if self.decided: + return False + self.dropped = reason + self.decided = True + return True + + def built(self, prop: Proposal) -> str: + """Record that ``prop``'s build stage is over; returns what to do next. + + ``'rubric'`` when this was the last build and the statements are ready to + be judged, ``'decide'`` when the group needs no judging and nothing else is + outstanding, ``''`` when there is still work in flight. """ - snapshot_errors.pop(slot, None) - script = WORKSPACE_SNAPSHOT.format(workspace=args.workspace, - max_files=args.snapshot_max_files, - per_file=args.snapshot_per_file, - total_budget=args.snapshot_budget) - exit_code, output = runners[slot](script, 'python') - if exit_code != 0: - # One retry: 62 of run_clean6's 63 snapshot failures were the sandbox - # answering 410 "not proxyable", which is the host having paused it and - # may be over by the time we ask again. Not fatal either way, but not - # silent: checks written against a missing end state are the failure - # this whole function exists to prevent. - logger.warning(f'[challenge] workspace snapshot failed (exit {exit_code}), ' - f'retrying in {SNAPSHOT_RETRY_WAIT}s: {output[-200:]}') - time.sleep(SNAPSHOT_RETRY_WAIT) - exit_code, output = runners[slot](script, 'python') - if exit_code != 0: - logger.warning(f'[challenge] workspace snapshot failed again (exit ' - f'{exit_code}): {output[-200:]}') - snapshot_errors[slot] = f'workspace snapshot failed (exit {exit_code}): {output[-500:]}' + with self.lock: + self.n_built += 1 + if self.n_built < self.size or self.decided: + return '' + if any(p.statement for p in self.proposals): + return 'rubric' + self.rubric_done = True + return self._ready_locked() + + def judged(self) -> str: + with self.lock: + self.rubric_done = True + return self._ready_locked() + + def solved(self, prop: Proposal, attempt: Optional[Dict[str, Any]], passed: bool) -> str: + with self.lock: + prop.n_solved += 1 + prop.attempts.append(attempt or {}) + prop.passes.append(passed) + return self._ready_locked() + + def _ready_locked(self) -> str: + """``'decide'`` once every outstanding piece of this group has landed.""" + if self.decided or self.n_built < self.size or not self.rubric_done: return '' - return tool_payload(output).strip() - - def snapshot_error_fn(slot: int = 0) -> str: - """Why the last snapshot for ``slot`` was empty; '' if it really was.""" - return snapshot_errors.get(slot, '') - - # Arm B. Read at most this much per call: the executor truncates its output - # near 8 KB, and base64 grows 3 bytes into 4, so 4 KB of file is about 5.5 KB - # of text with room left for the JSON envelope. - SLICE_BYTES = 4096 - - - # The opening the solver is measured against, built by the same function the - # eval script uses so that n_pass and pass@k are measuring one thing. Until - # this existed the difficulty stage handed over the statement as a lone user - # message with no system prompt: nothing said the model was in a sandbox, could - # take many turns, or should make one call per reply, and it answered by - # writing whole programs into a single call argument until they truncated. - _solver_harness = solver_harness(args.agent_config) if args.solver_rollouts > 0 else None - - def solver_prompt_fn(query: str): - return _solver_harness.start(query) - - # Keywords - store = None - # Arm D replaces the three topic axes with one bank of 'kind of work' phrases: - # a proposal takes one phrase, not one entry from each of three axes. - categories = CATEGORIES - category_desc = CATEGORY_DESC - if args.keywords_n > 0: - store = KeywordStore(args.keyword_db, categories) - logger.info('[challenge] keyword bank loaded: ' - + ', '.join(f'{c}={len(store.items[c])}' for c in categories)) - - # Seeds - seeds = [] - if args.seed_file: - with open(args.seed_file, encoding='utf-8') as f: - for line in f: - if line.strip(): - seeds.append(json.loads(line)) - logger.info(f'[challenge] loaded {len(seeds)} seeds from {args.seed_file}') - - # Rejected log - rejected = open(args.dump_rejected, 'w', encoding='utf-8') if args.dump_rejected else None - - def _reject(record): - if rejected is not None: - rejected.write(json.dumps(record, ensure_ascii=False, default=str) + '\n') - # Flushed per record: this file is the only account of why proposals - # are being dropped, and a run is worth watching for hours before it - # ends. Buffered, it stays empty until then. - rejected.flush() - - propose_writer = ProposeTrajWriter(args.dump_propose_traj) - - # Solver attempts from the difficulty stage. One line per attempt, so a task - # measured at 0 of 4 can be read rather than guessed at: the attempt, the - # state it left, and what the check said about it. - solver_log = (open(args.dump_solver_attempts, 'w', encoding='utf-8') - if args.dump_solver_attempts else None) - - def _solver_attempt(record): - if solver_log is not None: - solver_log.write(json.dumps(record, ensure_ascii=False, default=str) + '\n') - solver_log.flush() - - # Keyword generation, one line per call. The bank is the first step of the - # whole pipeline and the easiest place to fail invisibly: a reply the parser - # rejects leaves the bank empty, and every proposal downstream then runs the - # no-keyword prompt while the run looks healthy. Whole runs went that way - # before this existed, so prompt and reply are both kept verbatim. - keyword_log = (open(args.dump_keyword_gen, 'w', encoding='utf-8') - if args.dump_keyword_gen else None) - - def _keyword_gen(record): - if keyword_log is not None: - keyword_log.write(json.dumps(record, ensure_ascii=False, default=str) + '\n') - keyword_log.flush() - - # Build challenger - prompts = agentic_prompts(max_build_files=args.max_build_files) - # Followup API: explore on the local (trainable) sampler above, but write the - # check script and problem statement over an OpenAI-compatible API (qwen3-max). - # Reuses the --challenger-api-* connection args; the thinking cap, when set, is - # sent as extra_body on every check/statement call. - followup_api = None - followup_extra_body = None - if args.followup_api: - from twinkle_agentic.protocol.openai import OpenAI - followup_api = OpenAI(model=args.challenger_api_model, - api_key=args.challenger_api_key or None, - base_url=args.challenger_api_base) - if args.challenger_thinking_budget > 0: - followup_extra_body = {'thinking_budget': args.challenger_thinking_budget} - logger.info(f'[challenge] followup (check + statement) via API ' - f'model={args.challenger_api_model} base={args.challenger_api_base}' - + (f' thinking_budget={args.challenger_thinking_budget}' - if followup_extra_body else '')) - challenger = AgenticChallenger( - prompts, - explorer, - seeds=seeds, - keyword_store=store, - category_desc=category_desc if store else None, - seed_mix_prob=args.seed_mix_prob, - reset_fn=reset_fn, - run_check_fn=run_check_fn, - workspace_snapshot_fn=workspace_snapshot_fn, - snapshot_error_fn=snapshot_error_fn, - # The executor's own schemas, so the rounds that may call tools advertise - # exactly what will run -- same source as the training script uses. - tool_schemas=schemas, - episode_concurrency=max(1, args.episode_concurrency), - episode_tool_managers=episode_tool_managers, - combo_arity=args.combo_arity, - arity_weights=[float(x) for x in args.arity_weights.split(',')] if args.arity_weights - else None, - single_kw_prob=args.single_kw_prob, - proposals_per_group=max(1, args.proposals_per_group), - keyword_refill_target=args.keywords_n, - keyword_gen_calls=args.keyword_gen_calls, - keyword_refill_concurrency=max(1, args.keyword_refill_concurrency), - keyword_refill_tries=args.keyword_refill_tries, - keyword_params=SamplingParams(max_tokens=args.keyword_max_tokens, num_samples=1, - logprobs=1, temperature=args.keyword_temp, top_p=0.98), - # Same temperature as the episode, different budgets: the only thing - # being changed per stage is how much room the reply gets. - check_params=SamplingParams(max_tokens=args.check_max_tokens, num_samples=1, - logprobs=1, temperature=args.propose_temp, top_p=0.95), - problem_params=SamplingParams(max_tokens=args.problem_max_tokens, num_samples=1, - logprobs=1, temperature=args.propose_temp, top_p=0.95), - followup_api=followup_api, - followup_extra_body=followup_extra_body, - keyword_explorer=keyword_explorer, - min_batch=args.sampler_gpus, - problem_max_chars=args.problem_max_chars, - max_proposals_total=args.max_proposals_total, - solver_prompt_fn=solver_prompt_fn if _solver_harness is not None else None, - check_retries=args.check_retries, - reject_sink=_reject, - propose_sink=propose_writer.write, - solver_sink=_solver_attempt, - keyword_sink=_keyword_gen, - max_proposals_per_round=args.max_proposals_per_round, - solver_rollouts=args.solver_rollouts, - keep_min_pass=args.keep_min_pass, - keep_max_pass_margin=args.keep_max_margin, - # One call per reply here too, for the same reason and to keep the two - # sides comparable: a solver whose read-back is dispatched alongside the - # write it is checking fails a task the proposer built cleanly, and - # n_pass would then be measuring the dispatch, not the difficulty. - solver_params=SamplingParams(max_tokens=args.solver_max_tokens, num_samples=1, - logprobs=1, temperature=args.solver_temp, top_p=0.95, - stop=explore_stop, - include_stop_str_in_output=bool(explore_stop)), - seed=args.random_seed, - ) + if any(p.statement and p.n_solved < self.rollouts for p in self.proposals): + return '' + self.decided = True + return 'decide' + + def statements(self) -> List[Proposal]: + return [p for p in self.proposals if p.statement] + + - # Generate - batch_size = args.batch_size or args.keep_target - kept = [] - - # --stage keywords stops after step 1: fill the bank, draw the combinations, - # write out the proposal prompts they produce, and exit without touching the - # sandbox. Step 1 was broken for several runs and the failure was only - # visible by reading what it fed the next step, so it has to be runnable on - # its own rather than only as the first minute of an hours-long run. - if args.stage == 'keywords': - proposals = challenger.propose(args.stage_proposals) - with open(args.out_flows, 'w', encoding='utf-8') as out: - for i, proposal in enumerate(proposals): - data = proposal.get('user_data') - out.write(json.dumps({ - 'index': i, - 'keywords': user_data_get(data, 'keywords', []), - 'seeded': user_data_get(data, 'seeded', False), - 'prompt': proposal['messages'][-1]['content'], - }, ensure_ascii=False) + '\n') - drawn = sum(1 for p in proposals - if user_data_get(p.get('user_data'), 'keywords', [])) - logger.info(f'[challenge] stage=keywords: {len(proposals)} proposals, ' - f'{drawn} of them carry keywords') - if store is not None: - store.save() - logger.info('[challenge] keyword bank saved -> ' + args.keyword_db - + ' (' + ', '.join(f'{c}={len(store.items[c])}' - for c in categories) + ')') - if keyword_log is not None: - keyword_log.close() - propose_writer.close() - if rejected is not None: - rejected.close() - if solver_log is not None: - solver_log.close() - if env is not None: - for e in envs: - e.close() - return - - # --stage explore stops after step 4: draw a proposal, clear the workspace, - # run the sandbox episode, snapshot what it left, and stop before the - # check-writing round. What it is for: the episode is where the run either - # produces something worth writing a check about or leaves an empty directory, - # and 9 of 30 proposals in run11 left an empty one for reasons the rejection - # record could not distinguish. Everything the episode saw and did goes out - # verbatim, so that question is answerable from the file. - if args.stage == 'explore': - explore_log = (open(args.dump_explore, 'w', encoding='utf-8') - if args.dump_explore else None) - empty = 0 - for i, proposal in enumerate(challenger.propose(args.stage_proposals)): - reset_fn() - result = challenger.explore([proposal]) - episode = result[0] if result else {} - snapshot = workspace_snapshot_fn() - if not snapshot.strip(): - empty += 1 - messages = episode.get('messages') or [] - calls = sum(len(m.get('tool_calls') or []) for m in messages - if isinstance(m, dict)) - logger.info(f'[challenge] episode {i}: stop={episode.get("stop_reason")} ' - f'truncated={bool(episode.get("truncated"))} ' - f'stuck_stop={bool(episode.get("stuck_stop"))} ' - f'turns={episode.get("turns")} calls={calls} ' - f'end_state={len(snapshot)}b') - if explore_log is not None: - explore_log.write(json.dumps({ - 'index': i, - 'keywords': user_data_get(proposal.get('user_data'), 'keywords', []), - 'prompt': proposal['messages'][-1]['content'], - 'stop_reason': episode.get('stop_reason'), - 'truncated': bool(episode.get('truncated')), - 'stuck_stop': bool(episode.get('stuck_stop')), - 'turns': episode.get('turns'), - 'n_tool_calls': calls, - 'messages': messages, - 'end_state': snapshot, - }, ensure_ascii=False, default=str) + '\n') - explore_log.flush() - logger.info(f'[challenge] stage=explore: {args.stage_proposals} episodes, ' - f'{empty} left an empty workspace') - if explore_log is not None: - explore_log.close() - if keyword_log is not None: - keyword_log.close() - propose_writer.close() - if rejected is not None: - rejected.close() - if solver_log is not None: - solver_log.close() - if store is not None: - store.save() - if env is not None: - for e in envs: - e.close() - return - - # Appended as batches arrive, then rewritten sorted at the end. A run that - # keeps one task every few minutes for hours cannot afford to hold them all - # in memory only: a crash at hour three would leave nothing to train on, - # while an unsorted partial file is a usable task set. - with open(args.out_flows, 'w', encoding='utf-8') as partial: - for batch in challenger(batch_size=batch_size, total=args.keep_target): - for offset, task in enumerate(batch): - partial.write(json.dumps(flow_record(len(kept) + offset, task), - ensure_ascii=False) + '\n') - partial.flush() - kept.extend(batch) - logger.info(f'[challenge] kept {len(kept)}/{args.keep_target} so far; ' - f'stats {challenger.stats}') - if rejected is not None: - rejected.close() - if solver_log is not None: - solver_log.close() - propose_writer.close() - - # Before the keyword log is closed: expanding the bank generates keywords, - # and generating them writes to that log. Closing it first ended ex11 -- - # after all 4 tasks were kept and written -- with `ValueError: I/O operation - # on closed file`, which also skipped store.save() below and every line - # after it, so the run reported nothing about what it had produced. - if store is not None: - challenger.expand_hard_keywords() - store.save() - logger.info('[challenge] keyword bank saved -> ' + args.keyword_db) - if keyword_log is not None: - keyword_log.close() - - # Sort by difficulty (hardest last) - if not args.no_sort_by_difficulty: - kept.sort(key=lambda t: -(user_data_get(t.get('user_data'), 'n_pass', 0) or 0)) - - # Write output - write_flows(kept, args) - logger.info(f'[challenge] wrote {len(kept)} tasks -> {args.out_flows}') - if env.n_recoveries: - logger.warning(f'[challenge] sandbox was rebuilt {env.n_recoveries} time(s) during ' - f'this run; episodes in flight at those moments were lost') - dist = {} - for task in kept: - n = user_data_get(task.get('user_data'), 'n_pass', 0) - dist[n] = dist.get(n, 0) + 1 - logger.info(f'[challenge] pass-count distribution: {dict(sorted(dist.items()))}') - - for e in envs: - e.close() - - -class ProposeTrajWriter: - """Persist the proposing rounds so the challenger could be trained later. - - One ``.npz`` per proposal attempt holds the arrays, and one line per attempt - in ``index.jsonl`` holds everything a human reads plus the outcome. Splitting - them is what keeps this affordable: a 20-turn agentic episode is tens of - thousands of token ids, which as JSON is an order of magnitude larger than - the same numbers as int32. - - ``logprobs`` arrive as ``[[(token_id, logprob)]]`` and are flattened to the - logprob column alone -- that is the shape GRPO's ``old_logps`` wants, and the - token each one belongs to is already in ``labels``. - - Rejected attempts are written too. Their outcome is the reward's zero, and a - dump of kept-only attempts would have nothing to contrast against. +# ── The run ──────────────────────────────────────────────────────────────── + + +class Run: + """One collection pass: the resources, the queue, and the three job bodies. + + Sandbox jobs go on one FIFO queue served by one thread per slot, so a slot is + never idle while there is work. Rubric jobs go to a separate pool because they + need no sandbox, and putting them on the same queue would let a group's + judgement wait behind the solver attempts of another group. + + Nothing in a job waits for another job. A build enqueues its solver attempts + and returns its slot; a group is decided by whichever thread happens to land + the last outstanding piece. That is what keeps the pool from deadlocking on + itself, which a build that waited for its own solvers would do at once. """ - def __init__(self, out_dir: str): - self.dir = out_dir - self.index = None - self.n = 0 - if not out_dir: + def __init__(self, args, sampler, template, slots: List[Any], recorder: Recorder): + self.args = args + self.slots = slots + self.rec = recorder + self.rng = random.Random(args.random_seed or None) + + # Two rollouts over one sampler and one template: they differ only in the + # turn budget, and a per-call override for that does not exist. The + # trajectory-level state a rollout keeps is all local to __call__, so both + # are called from every thread at once with a per-call tool_manager. + self.propose_params = SamplingParams( + max_tokens=args.propose_max_tokens, num_samples=1, logprobs=1, + temperature=args.propose_temp, top_p=0.95, + stop=['</tool_call>'] if args.one_call_per_reply else None, + include_stop_str_in_output=bool(args.one_call_per_reply)) + self.solve_params = SamplingParams( + max_tokens=args.solver_max_tokens, num_samples=1, logprobs=1, + temperature=args.solver_temp, top_p=0.95, + stop=['</tool_call>'] if args.one_call_per_reply else None, + include_stop_str_in_output=bool(args.one_call_per_reply)) + self.propose_rollout = MultiTurnRollout( + sampler, template=template, tool_manager=ToolManager(), + max_turns=args.max_turns, stop_after_stuck_turns=args.stop_after_stuck_turns, + sampling_params=self.propose_params) + self.solve_rollout = MultiTurnRollout( + sampler, template=template, tool_manager=ToolManager(), + max_turns=args.solver_max_turns, + stop_after_stuck_turns=args.stop_after_stuck_turns, + sampling_params=self.solve_params) + self.keyword_rollout = MultiTurnRollout( + sampler, template=template, tool_manager=ToolManager(), max_turns=1, + sampling_params=SamplingParams(max_tokens=args.keyword_max_tokens, + num_samples=1, logprobs=1, + temperature=args.keyword_temp, top_p=0.98)) + + from twinkle_agentic.protocol.openai import OpenAI + self.api = OpenAI(model=args.api_model, api_key=args.api_key or None, + base_url=args.api_base) + self.api_extra = ({'thinking_budget': args.api_thinking_budget} + if args.api_thinking_budget > 0 else None) + self.check_params = SamplingParams(max_tokens=args.check_max_tokens, num_samples=1, + temperature=args.propose_temp, top_p=0.95) + self.problem_params = SamplingParams(max_tokens=args.problem_max_tokens, + num_samples=1, temperature=args.propose_temp, + top_p=0.95) + + # Built once: the cap is part of the system prompt, so a build that got a + # different one would be a different experiment. + self.system = P.SYSTEM + (P.BUILD_SIZE_CAP.format(n=args.max_build_files) + if args.max_build_files > 0 else '') + self.store = KeywordStore(args.keyword_db, P.CATEGORIES) + self.bank = TaskBank(args.task_bank, refs=args.task_bank_refs) if args.task_bank else None + # ms-agent builds the solver's opening messages, and it does so through a + # stateful agent object -- so one instance, one lock, and only for the few + # milliseconds it takes to shape two messages. + self.harness = solver_harness(args.agent_config) + self.harness_lock = threading.Lock() + + self.jobs: 'queue.Queue' = queue.Queue() + self.api_pool = ThreadPoolExecutor(max_workers=args.api_concurrency, + thread_name_prefix='api') + self.state = threading.Lock() + self.kw_lock = threading.Lock() + # Jobs actually being worked on right now, sandbox and API. Only used by + # the stall check in run(): 'the queue is empty' is not 'there is nothing + # left to do' while a thread is still inside a job that will queue more. + self.busy = 0 + self.api_jobs = 0 + self.nonce = 0 + # (category, keyword) pairs behind tasks nobody solved. Read at the end by + # expand_hard_keywords, which asks for more in the same domains. + self.hard: List[Tuple[str, str]] = [] + self.kept: List[Group] = [] + self.groups: List[Group] = [] + self.n_launched = 0 + self.stop = threading.Event() + self.counts: Dict[str, int] = {} + + # ---------------------------------------------------------------- helpers + + def bump(self, key: str, n: int = 1) -> None: + with self.state: + self.counts[key] = self.counts.get(key, 0) + n + + def draw_keywords(self) -> Tuple[List[Tuple[str, str]], str]: + """One entry from each category, refilling any that has run dry. + + On its own lock, not ``state``: a refill is eight model calls and holding + the lock every ``bump`` needs for that long would stall all 32 slots. Two + threads drawing at once still have to take turns, or the second refill's + prompt would not know what the first one had just said. + """ + with self.kw_lock: + for category in P.CATEGORIES: + if not self.store.unused(category): + self.refill(category) + picks = [] + for category in P.CATEGORIES: + text = self.store.take(category, self.rng) + if text is not None: + picks.append((category, text)) + return picks, '\n'.join(f'- {c}: {t}' for c, t in picks) + + def refill(self, category: str) -> None: + """Ask the local model for more keywords in ``category``. + + Says so when it comes back empty. A silent no-op here is the worst outcome + available: every proposal then falls back to a keyword-less prompt and the + run looks normal while producing one identical prompt over and over. That + is exactly what happened for whole runs when the prompt asked for one + keyword per line and the parser wanted a JSON array. + + The calls run one at a time so each can be told what the ones before it + said; each is answered by a rollout with ``max_turns=1``, which ends the + trajectory before any tool could be dispatched -- brainstorming a list is a + text round, and a bracketed list in a reply is exactly what a tool-calling + rollout would try to run. + """ + for attempt in range(1, max(1, self.args.keyword_refill_tries) + 1): + if self.generate_keywords(category): + return + logger.warning(f'[challenge] keyword refill for {category!r} produced ' + f'nothing new on try {attempt}') + if self.store.items[category]: + # Every keyword marked unused again. The alternative is a category that + # can never be drawn from, which stops the run: a repeat draw is worse + # than no run only if diversity matters more than collecting anything. + self.store.recycle(category) + self.store.save() + logger.warning(f'[challenge] keyword category {category!r} exhausted -> ' + f'recycled {len(self.store.items[category])} topics') + + def generate_keywords(self, category: str) -> bool: + """One refill round. True when it added something the bank did not have.""" + want = self.args.keywords_n + calls = max(1, self.args.keyword_gen_calls) + per_call = max(1, -(-want // calls) + 4) + known = self.store.texts(category) + fresh: List[str] = [] + seen = {t.strip().lower() for t in known} + for i in range(calls): + # Newest first: the calls run one at a time so each can avoid what the + # ones before it said, and letting older entries evict those would undo + # it. Past the cap the oldest of this refill's phrases fall off, which + # is also the least costly thing to drop. + avoid = (fresh + known)[:AVOID_TOTAL] + self.nonce += 1 + user = (P.KEYWORD_USER.format(k=per_call, desc=P.CATEGORY_DESC[category]) + + ('\nDo NOT repeat any of these already-used topics: ' + + ', '.join(avoid) if avoid else '') + + f'\n(batch {self.nonce}-{i})') + traj = {'messages': [{'role': 'system', 'content': P.KEYWORD_SYSTEM}, + {'role': 'user', 'content': user}]} + out = self.keyword_rollout([traj]) + reply = self._assistant_text(out[0] if out else {}) + parsed = parse_keyword_list(reply) + new = [k for k in parsed if k.lower() not in seen] + for keyword in new: + seen.add(keyword.lower()) + fresh.extend(new) + self.rec.keywords({'category': category, 'prompt': user, 'reply': reply, + 'parsed': parsed, 'n_parsed': len(parsed), + 'n_new': len(new), + 'stop_reason': (out[0].get('stop_reason') if out else None), + 'truncated': bool(out[0].get('truncated')) if out else None}) + if len(fresh) >= want: + break + added = self.store.add(category, fresh[:want], source='gen') + if added: + # Written now rather than at the end of the run: a run that crashes after + # spending eight model calls on keywords should not have to spend them + # again, and the next iteration reads this file to know what was used. + self.store.save() + logger.info(f'[challenge] keywords {category!r} +{added}') + return bool(added) + + @staticmethod + def _assistant_text(traj: Dict[str, Any]) -> str: + for message in reversed((traj.get('messages') if traj else []) or []): + if message.get('role') == 'assistant': + return message.get('content') or '' + return '' + + def expand_hard_keywords(self) -> int: + """More keywords in the domains that produced tasks nobody solved. + + Run once at the end, so what it adds is there for the next iteration rather + than for the groups still in flight. One call per hard keyword, capped at + 32 of them: this is the only feedback the keyword bank gets from difficulty, + and without it the bank drifts wherever the refill prompt happens to go. + """ + with self.state: + hard = list(self.hard)[:32] + if not hard: + return 0 + added = 0 + for i, (category, keyword) in enumerate(hard): + self.nonce += 1 + traj = {'messages': [ + {'role': 'system', 'content': P.KEYWORD_SYSTEM}, + {'role': 'user', + 'content': P.KEYWORD_EXPAND_USER.format(kw=keyword, m=8) + + f'\n(batch {self.nonce}-{i})'}]} + out = self.keyword_rollout([traj]) + reply = self._assistant_text(out[0] if out else {}) + parsed = parse_keyword_list(reply) + added += self.store.add(category, parsed, source='expand', parent=keyword) + self.rec.keywords({'category': category, 'parent': keyword, 'prompt': 'expand', + 'reply': reply, 'parsed': parsed, 'n_parsed': len(parsed)}) + self.store.save() + logger.info(f'[challenge] expanded {len(hard)} hard keyword(s) -> ' + f'+{added} same-domain topics') + return added + + def launch_group(self) -> Optional[Group]: + """Draw a topic and queue its ``group_size`` builds. None once at the cap.""" + with self.state: + if self.stop.is_set(): + return None + if self.args.max_group_attempts and self.n_launched >= self.args.max_group_attempts: + return None + gid = self.n_launched + self.n_launched += 1 + picks, block = self.draw_keywords() + if len(picks) != len(P.CATEGORIES): + # Every proposal's prompt is the keyword draw, so there is no honest + # prompt to send without one. Stopping is the reportable outcome; a + # substitute prompt would change what is being trained and say nothing. + logger.error(f'[challenge] keyword bank gave {len(picks)} of ' + f'{len(P.CATEGORIES)} categories; cannot build a prompt') + return None + prompt = P.FROM_KEYWORDS.format(keywords=block) + group = Group(gid, picks, block, prompt, self.args.group_size, + self.args.solver_rollouts) + with self.state: + self.groups.append(group) + for prop in group.proposals: + self.jobs.put(lambda slot, p=prop: self.build_job(p, slot)) + logger.info(f'[challenge] group {gid} launched: {block.replace(chr(10), " | ")}') + return group + + # ------------------------------------------------------------------ jobs + + def build_job(self, prop: Proposal, slot) -> None: + """Stage 1-3 for one proposal, then hand its statement to eight solvers.""" + if prop.group.dropped: + self.bump('build_skipped') + return + try: + self.build(prop, slot) + except Exception as e: # noqa: BLE001 -- one bad build must not end the run + logger.warning(f'[challenge] build g{prop.group.id}/{prop.idx} raised: ' + f'{type(e).__name__}: {e}') + prop.outcome, prop.detail = 'build_error', f'{type(e).__name__}: {e}' + self.bump(f'build:{prop.outcome}') + if not prop.statement: + self.record_rejection(prop) + if prop.statement and not prop.group.dropped: + for _ in range(self.args.solver_rollouts): + self.jobs.put(lambda s, p=prop: self.solve_job(p, s)) + action = prop.group.built(prop) + if action == 'rubric': + with self.state: + self.api_jobs += 1 + self.api_pool.submit(self.rubric_job, prop.group) + elif action == 'decide': + self.decide(prop.group) + + def record_rejection(self, prop: Proposal) -> None: + """Why this build produced no task, with enough of the episode to tell. + + How the episode ended travels with the reason. A reason on its own is not + diagnosable: whether a model that left an empty workspace ran out of tokens + or simply emitted no tool call is answered by stop_reason and the call + count, not by the word 'empty_workspace'. + """ + traj = prop.traj or {} + messages = traj.get('messages') or [] + self.rec.rejected({ + 'group_id': prop.group.id, + 'proposal_idx': prop.idx, + 'reason': prop.outcome, + 'detail': prop.detail, + 'keywords': prop.group.keywords, + 'stop_reason': traj.get('stop_reason'), + 'truncated': bool(traj.get('truncated')), + 'stuck_stop': bool(traj.get('stuck_stop')), + 'tool_stop': traj.get('tool_stop'), + 'turns': traj.get('turns'), + 'n_assistant': sum(1 for m in messages + if isinstance(m, dict) and m.get('role') == 'assistant'), + 'n_tool_calls': sum(len(m.get('tool_calls') or []) for m in messages + if isinstance(m, dict)), + 'last_assistant': self._assistant_text(traj), + 'check': prop.check, + }) + + def build(self, prop: Proposal, slot) -> None: + """Build in the sandbox, then have the API write the check and the task. + + The build is the trainable part and runs on the local model. The two stages + after it are appended to a *copy* of its messages and answered by the API, + so the check script and the statement are written with the whole build + history in view while the trajectory keeps exactly the tokens the local + model produced. + """ + args = self.args + slot.clear() + traj = {'messages': [{'role': 'system', 'content': self.system}, + {'role': 'user', 'content': prop.group.prompt}], + 'tools': slot.schemas} + prop.traj = rollout_one(self.propose_rollout, traj, self.propose_params, slot) + if prop.traj is None: + prop.outcome = 'rollout_empty' + return + if prop.traj.get('stop_reason') == 'length': + # A reply cut off at the token budget never finished its thought, so + # continuing the conversation over the API would write a check against + # a half-written turn. The trajectory is kept and trains with reward 0; + # what stops here are the two stages after it. + prop.outcome = 'cut_short' + prop.detail = f'stop_reason=length after {prop.traj.get("turns")} turn(s)' return - os.makedirs(out_dir, exist_ok=True) - self.index = open(os.path.join(out_dir, 'index.jsonl'), 'w', encoding='utf-8') - def write(self, record): - if self.index is None: + snapshot, error = slot.snapshot() + if error: + # Not filed as an empty workspace: a snapshot that says "empty" when it + # means "I could not look" produces tasks whose only true assertion is + # that nothing happened. + prop.outcome, prop.detail = 'snapshot_unavailable', error + return + if not snapshot: + prop.outcome = 'empty_workspace' + return + + messages = [dict(m) for m in prop.traj.get('messages') or []] + user_text = P.CHECK_FOLLOWUP.format(final_state=snapshot) + attempt = 0 + while True: + attempt += 1 + reply = api_one(self.api, messages, user_text, self.check_params, self.api_extra) + if reply is None: + prop.outcome, prop.detail = 'api_error', 'check-script call failed' + return + script = parse_check_script(reply) + if script is None: + if attempt <= args.check_retries: + user_text = P.CHECK_RETRY_FOLLOWUP.format( + error='Could not read a check script from your reply: it was ' + 'not a fenced python code block. Do not wrap it in a ' + 'tool call and do not add prose -- return ONLY a fenced ' + 'python code block.', + final_state=snapshot) + continue + prop.outcome, prop.detail = 'check_parse_fail', reply + return + # Rejected on the syntax tree before it can pass on the author's own + # state, since passing there is exactly what hides the defect: a check + # that pins a file's size or quotes a script's source passes for its + # author and fails every correct reproduction. + brittle = brittle_check_reason(script) + exit_code, output = (1, brittle) if brittle else call_one(slot, script) + if exit_code == 0: + prop.check = script + break + after, _ = slot.snapshot() + if attempt <= args.check_retries: + user_text = P.CHECK_RETRY_FOLLOWUP.format(error=output, + final_state=after or snapshot) + continue + prop.outcome = 'check_run_fail' + prop.detail = (f'exit {exit_code}\n{output}\n--- check script ---\n{script}' + f'\n--- state after check ---\n{after}') return - trace_id = f'p{self.n:06d}' - self.n += 1 - arrays, meta = {}, [] - for i, rnd in enumerate(record.get('rounds') or []): - labels = rnd.get('labels') or [] - logprobs = rnd.get('logprobs') or [] - if rnd.get('input_ids'): - arrays[f'r{i}_input_ids'] = np.asarray(rnd['input_ids'], dtype=np.int32) - if labels: - arrays[f'r{i}_labels'] = np.asarray(labels, dtype=np.int32) - if logprobs: - # float64, not float32: these are the ``old_logps`` a GRPO step - # divides by, and the sampler hands them over as full-precision - # python floats (the solver-side json dump keeps all 17 digits, - # e.g. -0.4740769863128662). float32 would round them to about 7 - # digits, so the ratio exp(logp - old_logp) would be off by - # roughly 1e-7 for reasons that have nothing to do with the - # policy having changed. - arrays[f'r{i}_logprobs'] = np.asarray( - [lp[0][1] for lp in logprobs], dtype=np.float64) - meta.append({ - 'stage': rnd.get('stage'), - 'messages': rnd.get('messages') or [], - 'n_tokens': len(rnd.get('input_ids') or []), - 'n_trainable': sum(1 for label in labels if label != -100), - 'n_logprobs': len(logprobs), - }) - # No arrays means the explorer was text-only (an API rollout), so there - # is nothing trainable to store -- record the attempt without an npz - # rather than leaving thousands of empty archives behind. - npz_name = f'{trace_id}.npz' if arrays else None - if arrays: - np.savez_compressed(os.path.join(self.dir, npz_name), **arrays) - line = { - 'trace_id': trace_id, - 'npz': npz_name, - 'outcome': record.get('outcome'), - # Both of these come straight from _emit_propose and both are what the - # proposing side trains on: train_offline.py groups proposals by - # group_id to get a GRPO advantage out of them, and skips the whole - # dump as a "pre-grouping run" when it is absent. Dropping them here - # silently turned SIDES=both into solver-only training. - 'group_id': record.get('group_id'), - 'challenger_reward': record.get('challenger_reward'), - 'n_pass': record.get('n_pass'), - 'n_rollouts': record.get('n_rollouts'), - 'pass_rate': record.get('pass_rate'), - 'keywords': record.get('keywords'), - 'seeded': record.get('seeded'), - 'rounds': meta, - } - self.index.write(json.dumps(line, ensure_ascii=False, default=str) + '\n') - self.index.flush() - def close(self): - if self.index is not None: - self.index.close() - logger.info(f'[challenge] wrote {self.n} propose traces -> {self.dir}') + reply = api_one(self.api, messages, P.PROBLEM_FOLLOWUP, self.problem_params, + self.api_extra) + if reply is None: + prop.outcome, prop.detail = 'api_error', 'problem-statement call failed' + return + statement = parse_problem_statement(reply) + if not statement: + prop.outcome, prop.detail = 'problem_parse_fail', reply + return + if len(statement) > args.problem_max_chars: + prop.outcome = 'too_long' + prop.detail = f'{len(statement)} chars > {args.problem_max_chars}' + return + prop.statement = statement + prop.outcome = 'ok' + + def solve_job(self, prop: Proposal, slot) -> None: + """One attempt at ``prop``'s task, scored by ``prop``'s own check script. + + A truncated attempt is a failed attempt: it left a workspace the check + rejects, and the denominator stays at ``solver_rollouts`` so the same + ``n_pass`` means the same thing in every group. + """ + if prop.group.dropped: + self.bump('solve_skipped') + return + attempt, passed = None, False + exit_code, output, end_state = None, '', '' + try: + slot.clear() + with self.harness_lock: + opening = self.harness.start(prop.statement) + if not opening.get('tools'): + # The harness only shapes messages -- its tool list is empty on + # purpose -- so the schemas come from the slot that will run them. + opening['tools'] = slot.schemas + attempt = rollout_one(self.solve_rollout, opening, self.solve_params, slot) + if attempt is not None: + exit_code, output = call_one(slot, prop.check) + passed = exit_code == 0 + # Read after the check, not before: the check is allowed to write, and + # what a reader of a failed attempt needs is the workspace the check + # was unhappy with. + end_state, _ = slot.snapshot() + except Exception as e: # noqa: BLE001 -- a lost attempt is a failed attempt + logger.warning(f'[challenge] solve g{prop.group.id}/{prop.idx} raised: ' + f'{type(e).__name__}: {e}') + output = f'{type(e).__name__}: {e}' + self.rec.attempt({ + 'group_id': prop.group.id, + 'proposal_idx': prop.idx, + 'statement': prop.statement, + 'check_script': prop.check, + 'passed': passed, + 'check_exit': exit_code, + 'check_output': output, + # A cut-off reply counts as a failed attempt and stays in the + # denominator, so the flag travels with the record for that to be + # checkable from the file rather than taken on trust. + 'truncated': bool((attempt or {}).get('truncated')), + 'stop_reason': (attempt or {}).get('stop_reason'), + 'turns': (attempt or {}).get('turns'), + 'messages': (attempt or {}).get('messages') or [], + 'end_state': end_state, + }) + self.bump('solve_pass' if passed else 'solve_fail') + if prop.group.solved(prop, attempt, passed) == 'decide': + self.decide(prop.group) + + def rubric_job(self, group: Group) -> None: + """Score the group's statements for novelty, all against each other. + + The siblings are the references that matter: a whole group can be scored + identically novel against history while being eight versions of one idea, + and GRPO subtracts the group mean, so a term identical across the group + produces no gradient at all. That is why this waits for all eight builds + instead of scoring each statement as it lands -- and why waiting costs + nothing: the slots are held by other groups' jobs the whole time. + + Retried up to ``--novelty-tries`` times. If the last one still has no + verdict for some statement, the group is dropped and its pending solver + attempts are skipped. + + Wrapped whole, because this is the one job whose exceptions nobody would + see: it runs on a pool whose futures are never read, so a raise in here + left the group waiting for a verdict that never came, and the run then sat + with an empty queue and idle slots until it was killed. Anything + unexpected drops the group instead of stalling everything. + """ + try: + self._rubric(group) + except Exception as e: # noqa: BLE001 + logger.warning(f'[challenge] rubric for group {group.id} raised: ' + f'{type(e).__name__}: {e}') + if group.abandon(f'rubric_error: {type(e).__name__}: {e}'): + self.bump('group_dropped:rubric_error') + self.decide(group) + finally: + with self.state: + self.api_jobs -= 1 + + def _rubric(self, group: Group) -> None: + if group.dropped: + return + props = group.statements() + if self.bank is None: + # No bank means no reference set, so nothing to be novel against. + # Novelty stays None and the reward is the pass-rate gaussian alone. + self._advance(group, group.judged()) + return + from twinkle_agentic.verifier import DIMENSIONS, score_tasks + texts = [p.statement for p in props] + pending = list(range(len(props))) + for attempt in range(1, max(1, self.args.novelty_tries) + 1): + payload = [{ + 'statement': texts[i], + 'check': props[i].check, + 'references': self.bank.references( + texts[i], extra=[t for j, t in enumerate(texts) if j != i]), + } for i in pending] + results = score_tasks(payload, workers=self.args.api_concurrency, + model=self.args.api_model, + extra_body=self.api_extra) + still: List[int] = [] + for i, task, result in zip(pending, payload, results): + score = result.scores.get('novelty') + self.rec.novelty({ + 'group_id': group.id, 'proposal_idx': props[i].idx, 'try': attempt, + **{dim: result.scores.get(dim) for dim in DIMENSIONS}, + 'verdicts': result.verdicts, 'n_votes': result.n_votes, + 'error': result.error, + 'n_references': len(task.get('references') or ()), + # Full text on both sides: this file is read to check whether a + # score was deserved, which a shortened statement cannot answer. + 'statement': task.get('statement') or '', + 'references': list(task.get('references') or ()), + }) + if score is None: + still.append(i) + else: + props[i].novelty = float(score) + if not still: + break + pending = still + logger.warning(f'[challenge] group {group.id}: {len(still)} statement(s) ' + f'came back without a novelty verdict (try {attempt})') + else: + if pending and group.abandon(f'novelty_unscored x{len(pending)}'): + self.bump('group_dropped:novelty') + self.decide(group) + return + self._advance(group, group.judged()) + + # -------------------------------------------------------------- decision + def decide(self, group: Group) -> None: + """Keep or drop the group, then start a replacement or stop the run. -def flow_record(index, task): - """One task as the training script reads it back.""" - data = task.get('user_data') - messages = task.get('messages') or [] - query = next((m['content'] for m in messages if m.get('role') == 'user'), '') + The second half is in a ``finally`` because the first half writes files: a + raise while writing used to take the replacement topic down with it, and + the run then had one fewer group in flight for every failure until there + was nothing left running and nothing left to wait for. + """ + try: + self._decide(group) + except Exception as e: # noqa: BLE001 + logger.warning(f'[challenge] deciding group {group.id} raised: ' + f'{type(e).__name__}: {e}') + self.bump('group_decide_error') + finally: + self._after_decision() + + def _decide(self, group: Group) -> None: + rollouts = self.args.solver_rollouts + floor = self.args.novelty_floor + in_band = [p for p in group.proposals + if p.n_pass is not None and 1 <= p.n_pass <= rollouts - 1] + chosen = max(in_band, key=lambda p: p.reward(rollouts, floor)) if in_band else None + if chosen is not None and group.dropped: + # An abandoned group can still have in-band proposals: its solver + # attempts were already running when it was abandoned. Keeping it on + # that basis would train on the very group that was judged unusable, + # and would do it with a novelty term measured for some members and + # missing for others. + chosen = None + if chosen is not None and not self._claim_keep(group): + # The target was reached while this group was finishing. Claimed before + # anything is written, because writing first and counting after is how + # a run ends up with eleven groups on disk and a loader that reads a + # different number of GRPO groups than the run reported. + chosen = None + self.bump('group_late') + record = { + 'group_id': group.id, + 'kept': chosen is not None, + 'dropped': group.dropped, + 'keywords': group.keywords, + 'chosen': chosen.idx if chosen is not None else None, + 'n_in_band': len(in_band), + 'proposals': [{ + 'idx': p.idx, + 'outcome': p.outcome, + 'n_pass': p.n_pass, + 'novelty': p.novelty, + 'reward': p.reward(rollouts, floor), + 'statement': p.statement, + 'check': p.check, + 'detail': p.detail, + } for p in group.proposals], + } + self.rec.group(record) + # Keyword draws behind tasks nobody solved, for expand_hard_keywords. Taken + # from every decided group, kept or not: a task at n_pass=0 says the same + # thing about its keywords either way. + if any(p.statement and p.n_pass == 0 for p in group.proposals): + with self.state: + seen = {(c, t.lower()) for c, t in self.hard} + for category, text in group.keywords: + if (category, text.lower()) not in seen: + self.hard.append((category, text)) + if chosen is None: + self.bump('group_dropped' if not group.dropped else 'group_dropped_early') + return + + # Every proposal of a kept group trains, including the ones that produced + # no task: they are the zero-reward half of the GRPO group, and a set of + # kept-only records has no variance to learn from. + for prop in group.proposals: + if prop.traj is None: + continue + self.rec.trajectory( + prop.traj, side='propose', group_id=group.id, proposal_idx=prop.idx, + reward=prop.reward(rollouts, floor), n_pass=prop.n_pass, + novelty=prop.novelty, outcome=prop.outcome, + keywords=group.keywords, selected=prop is chosen) + # Only the chosen proposal's attempts. The others were measured and are + # reported in groups.jsonl, but training on eight near-identical tasks from + # one keyword draw is what a group of one keyword direction is meant to + # avoid. + for i, (attempt, passed) in enumerate(zip(chosen.attempts, chosen.passes)): + if not attempt: + continue + self.rec.trajectory(attempt, side='solve', group_id=group.id, + proposal_idx=chosen.idx, attempt_idx=i, + reward=1.0 if passed else 0.0, passed=passed, + statement=chosen.statement) + self.rec.task({'id': f'ag_g{group.id:04d}p{chosen.idx}', + 'group_id': group.id, 'proposal_idx': chosen.idx, + 'query': chosen.statement, 'check_script': chosen.check, + 'n_pass': chosen.n_pass, 'n_rollouts': rollouts, + 'novelty': chosen.novelty, + 'reward': chosen.reward(rollouts, floor), + 'keywords': group.keywords}) + if self.bank is not None: + self.bank.add(chosen.statement, chosen.check, group_id=group.id, + n_pass=chosen.n_pass) + self.bump('group_kept') + + def _claim_keep(self, group: Group) -> bool: + """Take one of the ``--keep-groups`` slots, if there is one left. + + The slot is taken before the group's trajectories are written and released + by nobody, so the number of groups on disk is exactly the number claimed + even though several groups can finish at the same moment. + """ + with self.state: + if len(self.kept) >= self.args.keep_groups: + return False + self.kept.append(group) + return True + + def _after_decision(self) -> None: + """Stop the run if the target is met, otherwise start a replacement topic. + + Replacing one topic per decided group is what keeps the number of groups in + flight at ``sandbox_slots / group_size`` without anything having to track + it: the queue is fed by whatever finishes. + """ + with self.state: + enough = len(self.kept) >= self.args.keep_groups + if enough: + if not self.stop.is_set(): + logger.info(f'[challenge] {len(self.kept)} groups kept; stopping') + self.stop.set() + return + if self.launch_group() is None and self._idle(): + logger.warning('[challenge] no topics left to try and nothing in flight; ' + f'stopping with {len(self.kept)} kept group(s)') + self.stop.set() + + def _idle(self) -> bool: + with self.state: + return all(g.decided for g in self.groups) + + def _advance(self, group: Group, action: str) -> None: + if action == 'decide': + self.decide(group) + + # ------------------------------------------------------------- the loop + + def work(self, slot) -> None: + """One thread, one slot, jobs until the run stops.""" + while not self.stop.is_set(): + try: + job = self.jobs.get(timeout=1.0) + except queue.Empty: + continue + with self.state: + self.busy += 1 + try: + job(slot) + except Exception as e: # noqa: BLE001 -- never lose the thread + logger.warning(f'[challenge] job on slot {slot.slot} raised: ' + f'{type(e).__name__}: {e}') + finally: + with self.state: + self.busy -= 1 + self.jobs.task_done() + + def run(self) -> None: + """Start one thread per slot, prime the queue, and wait for the target.""" + n_topics = max(1, len(self.slots) // self.args.group_size) + threads = [threading.Thread(target=self.work, args=(slot,), daemon=True, + name=f'slot{slot.slot}') for slot in self.slots] + for thread in threads: + thread.start() + for _ in range(n_topics): + if self.launch_group() is None: + break + # Waited on in slices rather than once, so a run that has stopped making + # progress ends with the reason on stdout instead of sitting there. Two + # consecutive idle checks, because one can catch the moment between a job + # being taken off the queue and the counter going up. + idle_rounds = 0 + while not self.stop.wait(STALL_CHECK_SECONDS): + with self.state: + quiet = self.busy == 0 and self.api_jobs == 0 + stuck = [g.id for g in self.groups if not g.decided] + if not (quiet and self.jobs.empty()): + idle_rounds = 0 + continue + idle_rounds += 1 + if idle_rounds < 2: + continue + # Reached only when the run has gone quiet without meeting its target + # and without deciding to stop, which is a bug rather than attrition: + # normally either a group finishes (and launches a replacement) or + # launch_group runs out and sets stop itself. + logger.error(f'[challenge] nothing running and nothing queued after ' + f'{len(self.kept)}/{self.args.keep_groups} kept groups ' + f'and {self.n_launched} launched' + + (f'; group(s) {stuck} were never decided' if stuck else '') + + '. Stopping.') + for group in list(self.groups): + if group.abandon('never_decided'): + self.bump('group_dropped:never_decided') + self.decide(group) + self.stop.set() + for thread in threads: + thread.join(timeout=self.args.sandbox_timeout) + self.api_pool.shutdown(wait=True) + + +def collect_metrics(out_dir: str, counts: Dict[str, int], launched: int, + rollouts: int, wall: float) -> Dict[str, Any]: + """What this collection produced, as numbers, for ``challenge_metrics.json``. + + Read back out of ``groups.jsonl`` rather than taken from the live objects, so + the file cannot disagree with the audit files it sits next to, and so the same + function can recompute the metrics for a directory that finished hours ago. + + Three sections, because they are read for different things: + + * ``scalars`` -- fixed keys, always present, every value a float or int. This + is the set that goes to swanlab; a key appearing in one iteration and not the + next would make a chart that means something different in each. + * ``counts`` -- the raw bump counters, dynamic keys and all + (``group_dropped:rubric_error`` only exists in a run where that happened). + Kept here and not uploaded. + * ``distributions`` -- the histograms behind the means, because a mean n_pass + of 4 is a different collection depending on whether it came from eights and + zeros or from fours. + + ``solve_pass_rate`` is over every solver attempt run, the number the user + asked for as accuracy. It is not a fixed yardstick: the tasks change every + iteration, so it moving says the pair moved, not which half. + """ + path = os.path.join(out_dir, 'groups.jsonl') + groups: List[Dict[str, Any]] = [] + if os.path.exists(path): + with open(path, encoding='utf-8') as f: + for line in f: + line = line.strip() + if line: + try: + groups.append(json.loads(line)) + except json.JSONDecodeError: + continue + props = [p for g in groups for p in g.get('proposals') or []] + with_stmt = [p for p in props if p.get('statement')] + measured = [p for p in with_stmt if p.get('n_pass') is not None] + in_band = [p for p in measured if 1 <= p['n_pass'] <= rollouts - 1] + chosen = [next((p for p in g['proposals'] if p['idx'] == g.get('chosen')), None) + for g in groups if g.get('kept')] + chosen = [p for p in chosen if p is not None] + novelty = [p['novelty'] for p in props if p.get('novelty') is not None] + rewards = [p['reward'] for p in props if p.get('reward') is not None] + passes = counts.get('solve_pass', 0) + attempts = passes + counts.get('solve_fail', 0) + kept = sum(1 for g in groups if g.get('kept')) + + def rate(num: float, den: float) -> float: + return float(num) / den if den else 0.0 + + scalars = { + 'groups_launched': launched, + 'groups_kept': kept, + 'groups_decided': len(groups), + 'group_keep_rate': rate(kept, len(groups)), + 'wall_seconds': round(wall, 1), + 'builds': len(props), + 'builds_with_statement': len(with_stmt), + 'build_statement_rate': rate(len(with_stmt), len(props)), + # The accuracy: every solver attempt that ran, passed over total. + 'solve_attempts': attempts, + 'solve_pass_rate': rate(passes, attempts), + # Of the tasks that were measured at all, how many landed in the band the + # keep rule wants. This is the proposer's hit rate. + 'n_pass_in_band_rate': rate(len(in_band), len(measured)), + 'n_pass_mean': statistics.fmean(p['n_pass'] for p in measured) if measured else 0.0, + 'delivered_n_pass_mean': + statistics.fmean(p['n_pass'] for p in chosen if p.get('n_pass') is not None) + if chosen else 0.0, + 'proposer_reward_mean': statistics.fmean(rewards) if rewards else 0.0, + 'novelty_scored_rate': rate(len(novelty), len(with_stmt)), + 'novelty_mean': statistics.fmean(novelty) if novelty else 0.0, + 'novelty_zero_rate': rate(sum(1 for v in novelty if v == 0.0), len(novelty)), + } return { - 'id': f'ag_{index:06d}', - 'query': query, - 'check_script': user_data_get(data, 'check_script', ''), - # Arm B: run before the solver starts, to put the input files it is told - # it already has on disk. Empty for every other arm. - 'setup_script': user_data_get(data, 'setup_script', ''), - 'n_pass': user_data_get(data, 'n_pass'), - 'n_rollouts': user_data_get(data, 'n_rollouts'), - 'keywords': user_data_get(data, 'keywords', []), - 'seeded': user_data_get(data, 'seeded', False), + 'scalars': scalars, + 'counts': dict(sorted(counts.items())), + 'distributions': { + 'n_pass': {str(k): v for k, v in + sorted(collections.Counter(p['n_pass'] for p in measured).items())}, + 'build_outcome': {k.split(':', 1)[1]: v for k, v in sorted(counts.items()) + if k.startswith('build:')}, + 'novelty': {str(round(v, 2)): n for v, n in + sorted(collections.Counter(novelty).items())}, + }, } -def write_flows(kept, args): - """Write one flow per task, replacing whatever the run appended as it went.""" - with open(args.out_flows, 'w', encoding='utf-8') as f: - for i, task in enumerate(kept): - f.write(json.dumps(flow_record(i, task), ensure_ascii=False) + '\n') +def main(): + args = parse_args() + os.makedirs(args.out_dir, exist_ok=True) + recorder = Recorder(args.out_dir) + sampler, template = initialize_device(args) + slots = initialize_sandbox(args) + run = Run(args, sampler, template, slots, recorder) + started = time.time() + try: + run.run() + # After the loop, not during: what it adds is for the next iteration, and + # doing it here means a crash in collection does not also lose the bank. + if args.keyword_expand: + run.expand_hard_keywords() + finally: + rebuilds = close_pool(slots) + recorder.close() + if run.bank is not None: + logger.info(f'[challenge] task bank: {run.bank.stats()}') + run.store.save() + if rebuilds: + logger.warning(f'[challenge] sandboxes were rebuilt {rebuilds} time(s); ' + f'the jobs in flight at those moments were lost') + # Written after recorder.close(), so groups.jsonl is complete and flushed + # before it is read back. In the finally block because a run that crashed + # is the one whose numbers are most worth having. + metrics = collect_metrics(args.out_dir, run.counts, run.n_launched, + args.solver_rollouts, time.time() - started) + with open(os.path.join(args.out_dir, 'challenge_metrics.json'), 'w', + encoding='utf-8') as f: + json.dump(metrics, f, indent=2, ensure_ascii=False, default=str) + logger.info(f'[challenge] {len(run.kept)}/{run.n_launched} groups kept in ' + f'{time.time() - started:.0f}s, counts: ' + f'{dict(sorted(run.counts.items()))}') + logger.info(f'[challenge] metrics -> ' + f'{os.path.join(args.out_dir, "challenge_metrics.json")}: ' + f'{metrics["scalars"]}') if __name__ == '__main__': diff --git a/cookbook/rsi/agentic/episode.py b/cookbook/rsi/agentic/episode.py index f817d26f6..df236c699 100644 --- a/cookbook/rsi/agentic/episode.py +++ b/cookbook/rsi/agentic/episode.py @@ -1,15 +1,15 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Episode construction and scoring for agentic RSI, shared by training and eval. +"""Episode construction and scoring, shared by collection and eval. -Both halves need the same three things and must not disagree about any of them: -how an episode is built (a sandbox with ms-agent's tools plus a local harness -that only shapes messages), how the tool contract is advertised (schemas read off -the executor that will honour them), and how a trajectory is scored (the task's -own checks, run against the state the episode left behind). +Three things must not differ between the run that invents a task and the run that +measures it: how an episode is built (a sandbox with ms-agent's tools plus a local +harness that only shapes messages), how the tool contract is advertised (schemas +read off the executor that will honour them), and how a trajectory is scored (the +task's own checks, run against the state the episode left behind). -An eval that differed from training on any of these would measure something other -than what was trained, so this module is the single definition and the scripts -are only wiring. +``challenge.py`` takes ``solver_harness`` from here, and ``eval.py`` takes the +whole boot/score path, so a task kept at n_pass=4 during collection is a task the +eval measures the same way. A second copy of these lines would drift. """ import json import os @@ -127,8 +127,9 @@ def build_episode(task: Dict[str, Any], cfg: SandboxConfig) -> Tuple[Any, Any, A ) env.reset() - # A task may hand the solver its input files instead of asking it to write - # them (challenge.py --preseed-inputs). Loudly, not on a best-effort basis: a + # A task may carry a setup_script that writes its input files instead of asking + # the solver to. Nothing produces one now, but a task file from an older run + # can still hold one. Loudly, not on a best-effort basis: a # statement that says the inputs are on disk, run against a workspace where # they are not, scores 0 for a reason that has nothing to do with the task. setup = task.get('setup_script') diff --git a/cookbook/rsi/agentic/eval.py b/cookbook/rsi/agentic/eval.py index 5c5534ba5..69d7bbda6 100644 --- a/cookbook/rsi/agentic/eval.py +++ b/cookbook/rsi/agentic/eval.py @@ -1,23 +1,28 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Held-out evaluation for agentic RSI: pass rate on tasks the trainer never saw. -Episodes are built and scored by :mod:`episode`, the same module ``rl.py`` uses, -so a number reported here is the number training was optimising -- an eval that -constructed episodes differently would measure a different agent. +Episodes are built and scored by :mod:`episode`, the same module ``challenge.py`` +takes ``solver_harness`` from, so a number reported here is measured against the +opening a task's n_pass was measured against -- an eval that constructed episodes +differently would measure a different agent. -What it adds on top of training is only what training does not need: several +What it adds on top of collection is only what collection does not need: several attempts per task (a single attempt at temperature 1 is a coin flip, not a rate), -no optimizer, and a LoRA read off disk rather than synced from a live trainer. +no optimizer, and weights read off disk rather than held by a live trainer. + +Weights are named by ``--model-id`` and nothing else. ``train.py`` trains every +parameter and saves a whole model, so the trained side of a comparison is a +checkpoint directory in exactly the place the base model's name goes. Usage:: - # baseline, no adapter + # baseline python cookbook/rsi/agentic/eval.py --tasks output/.../eval_tasks.jsonl \\ --label base --out output/.../eval_base.jsonl # after training python cookbook/rsi/agentic/eval.py --tasks output/.../eval_tasks.jsonl \\ - --adapter-path output/rsi-agentic-final --label trained \\ + --model-id output/rsi_agentic/<tag>/ckpt/model --label trained \\ --out output/.../eval_trained.jsonl Both runs must use the same ``--tasks``, ``--rollouts-per-task`` and sampling @@ -48,19 +53,31 @@ def parse_args(): p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) p.add_argument('--tasks', required=True, help='task jsonl (challenge.py or structured)') - p.add_argument('--model-id', default='ms://Qwen/Qwen3-4B') - p.add_argument('--adapter-path', default='', - help='LoRA directory saved by rl.py; empty evaluates the base model') + p.add_argument('--model-id', default='ms://Qwen/Qwen3-4B', + help='base model name, or a checkpoint directory saved by train.py') p.add_argument('--label', default='eval', help='name for this measurement in the log') p.add_argument('--sampler-gpus', type=int, default=4) p.add_argument('--max-model-len', type=int, default=32768) - p.add_argument('--max-lora-rank', type=int, default=32) p.add_argument('--rollouts-per-task', type=int, default=4, help='attempts per task; the pass rate is over these') p.add_argument('--episodes-per-wave', type=int, default=16, help='sandboxes alive at once; keep at or below RSI_ENV_CONCURRENCY') p.add_argument('--max-turns', type=int, default=20) + # Both have to equal what challenge.py built the tasks under, and neither was + # reachable from the command line before: the rollout was constructed with the + # class default for the first and with --max-model-len for the second, while + # challenge.py passes --stop-after-stuck-turns and leaves the token cap unset. + # The stuck cutoff is the one that bites -- it ended 50 to 88 of each + # iteration's ~550 attempts -- so an eval that leaves it at 0 measures an agent + # that is allowed to repeat itself forever, against tasks whose n_pass was + # measured on an agent that was not. + p.add_argument('--stop-after-stuck-turns', type=int, default=2, + help="consecutive no-progress turns that end the tool phase; " + "challenge.py's default is 2, 0 disables the cutoff") + p.add_argument('--max-trajectory-tokens', type=int, default=0, + help='cap on the whole trajectory; 0 leaves it unset, which is ' + 'what challenge.py does') # Has to equal the challenger's --solver-max-tokens and --propose-max-tokens. # An eval that gives the model less room than the run that built the tasks is # measuring the budget, not the model: at 4096, 15 of 50 attempts ended on @@ -86,14 +103,10 @@ def build_sampler(args): mode='ray', nproc_per_node=args.sampler_gpus, lazy_collect=False, groups=[DeviceGroup(name='sampler', ranks=list(range(args.sampler_gpus)), device_type='GPU')]) - engine_args = {'gpu_memory_utilization': 0.8, 'max_model_len': args.max_model_len} - if args.adapter_path: - # Declared at construction or the engine has no slot to load into, which - # surfaces much later as an adapter that appears to do nothing. - engine_args.update({'enable_lora': True, 'max_lora_rank': args.max_lora_rank}) sampler = vLLMSampler( model_id=args.model_id, - engine_args=engine_args, + engine_args={'gpu_memory_utilization': 0.8, + 'max_model_len': args.max_model_len}, device_mesh=DeviceMesh.from_sizes(world_size=args.sampler_gpus, dp_size=args.sampler_gpus), remote_group='sampler', @@ -105,12 +118,10 @@ def build_sampler(args): def main(): args = parse_args() - if args.adapter_path and not os.path.isdir(args.adapter_path): - raise SystemExit(f'[eval] no adapter directory at {args.adapter_path}') tasks = load_tasks(args.tasks) cfg = SandboxConfig.from_env() - logger.info(f'[eval:{args.label}] {len(tasks)} tasks x {args.rollouts_per_task} attempts, ' - f'adapter={args.adapter_path or "(base model)"}') + logger.info(f'[eval:{args.label}] {len(tasks)} tasks x {args.rollouts_per_task} ' + f'attempts, weights={args.model_id}') logger.info(f'[eval:{args.label}] sandboxes: template={cfg.template} api={cfg.api_url}') sampler = build_sampler(args) @@ -122,8 +133,8 @@ def main(): sampling_params=SamplingParams(max_tokens=args.max_tokens, num_samples=1, logprobs=1, temperature=args.temperature, top_p=args.top_p), max_turns=args.max_turns, - max_trajectory_tokens=args.max_model_len, - adapter_path=args.adapter_path or None, + stop_after_stuck_turns=args.stop_after_stuck_turns, + max_trajectory_tokens=args.max_trajectory_tokens or None, ) # One flat list of attempts, so a wave is a fixed number of sandboxes no diff --git a/cookbook/rsi/agentic/loop.sh b/cookbook/rsi/agentic/loop.sh new file mode 100644 index 000000000..f07496426 --- /dev/null +++ b/cookbook/rsi/agentic/loop.sh @@ -0,0 +1,233 @@ +#!/bin/bash +# Self-evolving loop: collect, train on what was collected, collect again from the +# weights that came out. +# +# collect: challenge.py --model-id <last ckpt> --keep-groups 8 +# runs until 8 groups have been kept, whatever that costs in topics +# train: train.py --run-dir <that collection> -> one HF checkpoint +# repeat +# +# Nothing is generated in the training stage and nothing is re-encoded: the tokens +# trained on are the ones the sampler produced, read straight off disk. +# +# The two stages are separate processes so each gets every GPU. Collection is the +# slow half, and splitting the GPUs between a trainer and a sampler in one process +# would halve it. The cost is restarting vLLM and the sandboxes each time, measured +# at 1-2 minutes against roughly 40 minutes of collecting. +# +# Nothing about the host is written down here: the repo is found from this script's +# own location, the GPU count from nvidia-smi, and the secrets have to be exported +# first -- the script stops with the name of whatever is missing rather than +# guessing a value that would fail deep inside a run. +# +# export E2B_API_KEY=... # sandbox host key +# export SANDBOX_API_URL=http://... # sandbox host address, with port +# export LLM_BACKUP_API_KEY=... # dashscope, for checks/statements/rubric +# bash cookbook/rsi/agentic/loop.sh # until killed +# ITERATIONS=1 bash cookbook/rsi/agentic/loop.sh # one collect + one train +set -e +# Both stages are piped into tee, and without this the pipeline's status is tee's, +# which is 0 even when python died. An earlier loop crashed inside collection, +# trained on the partial collection anyway, saved a checkpoint from it and marked +# the iteration finished -- all reported as success. +set -o pipefail + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO="$(cd "$HERE/../../.." && pwd)" +if [ ! -f "$HERE/challenge.py" ] || [ ! -f "$REPO/setup.cfg" ]; then + echo "expected challenge.py beside this script and the repo root three levels up" >&2 + exit 1 +fi +cd "$REPO" + +missing="" +for v in E2B_API_KEY SANDBOX_API_URL LLM_BACKUP_API_KEY; do + [ -z "${!v}" ] && missing="$missing $v" +done +if [ -n "$missing" ]; then + echo "export these first:$missing" >&2 + exit 1 +fi + +# Every GPU on the box unless told otherwise. Counted rather than written down, +# since the point of moving hosts is usually a different number of them. +GPUS="${GPUS:-$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l)}" +[ "$GPUS" -lt 1 ] && { echo "nvidia-smi reports no GPUs" >&2; exit 1; } +DEVICES="${DEVICES:-$(seq -s, 0 $((GPUS - 1)))}" + +# Refuse to start on top of someone else's job. Both stages want the whole GPU: +# challenge.py boots one vLLM per GPU at 0.8 of its memory, so sharing means an +# out-of-memory crash partway in and the other job may go down with it. Any compute +# process at all counts; CONFIRM_GPUS=1 starts anyway. +BUSY="$(nvidia-smi --query-compute-apps=pid --format=csv,noheader | sort -u | wc -l)" +if [ "$BUSY" -gt 0 ] && [ "${CONFIRM_GPUS:-0}" != "1" ]; then + echo "$BUSY process(es) already on the GPUs:" >&2 + nvidia-smi --query-compute-apps=pid,used_memory --format=csv,noheader >&2 + echo "set CONFIRM_GPUS=1 to start anyway" >&2 + exit 1 +fi + +export AENV_API_URL="$SANDBOX_API_URL" +export AENV_API_KEY="$E2B_API_KEY" +# Sandbox image name: a fact about how the host was set up, not a preference. +export AENV_TEMPLATE="${AENV_TEMPLATE:-twinkle-rsi-msagent}" + +# ---- what to run --------------------------------------------------------- +ITERATIONS="${ITERATIONS:-0}" # 0 = until killed +TAG="${TAG:-fix1}" +BASE_MODEL="${BASE_MODEL:-ms://Qwen/Qwen3-4B}" + +# Concurrent sandboxes, i.e. how many trajectories are in flight at once. Bounded +# by the sandbox host, not by the GPUs here, so it does not follow GPUS. 32, not +# the 96 a capacity probe once managed: holding 96 for a whole run was not +# reliable. Each slot is one thread and one microVM; vLLM sees up to this many +# single-trajectory requests at a time and batches them itself. +SANDBOX_SLOTS="${SANDBOX_SLOTS:-32}" + +# Groups kept per iteration, and their shape. A group is one keyword draw answered +# GROUP_SIZE times; it is kept when at least one of those answers became a task the +# solver passes sometimes, meaning n_pass in [1, SOLVER_ROLLOUTS-1]. Outside that +# band every attempt carries the same reward, the group mean equals it, and the +# advantage is zero for all of them. +# +# proposing side: KEEP_GROUPS x GROUP_SIZE = 8 x 8 = 64 trajectories +# solving side: KEEP_GROUPS x SOLVER_ROLLOUTS = 8 x 8 = 64 trajectories +# 128 total, one optimizer step over all of it +# +# What this pays for and throws away: every proposal that produced a task costs +# SOLVER_ROLLOUTS sandbox attempts, and only the selected proposal's attempts are +# trained on. At 8 groups x 8 proposals that is up to 512 attempts run and 64 +# trained on. The unselected proposals are not wasted on the proposing side -- each +# earns its own reward from its own n_pass, including 0 for the ones that produced +# no task at all. +KEEP_GROUPS="${KEEP_GROUPS:-8}" +GROUP_SIZE="${GROUP_SIZE:-8}" +SOLVER_ROLLOUTS="${SOLVER_ROLLOUTS:-8}" + +# Cap on files one build may leave behind, appended to the system prompt. This +# changes the prompt and therefore what is trained; 4 is what every run since it +# was added has used. +MAX_BUILD_FILES="${MAX_BUILD_FILES:-4}" + +# Reasoning cap on every API call. The one knob that moved wall-clock: 58s -> 10s +# per turn at 2048 on a ~15k-character context. +API_THINKING_BUDGET="${API_THINKING_BUDGET:-4096}" +API_MODEL="${API_MODEL:-qwen3.8-max}" +API_BASE="${API_BASE:-https://dashscope.aliyuncs.com/compatible-mode/v1}" + +# Novelty. The bank is one file for the whole loop, not one per iteration, because +# the point of it is comparing iteration k+1's proposals against what k produced. +# TASK_BANK="" turns it off and gives back the pass-rate gaussian alone; +# NOVELTY_FLOOR=1 keeps the judging and the log but stops it changing any reward. +# +# 1 is the default because the score it would multiply in has not been shown to carry +# anything yet. Measured on the 27 proposals of iter1 (novelty_scores.jsonl): judged +# against their own siblings, 24 of 27 scored exactly 0.0, so the term was constant +# across the group and contributed nothing once GRPO subtracts the group mean -- while +# still halving every proposer reward at floor 0.5. The alternative measured, labelling +# each task's shape on its own and scoring by how rare that shape is in the group, does +# separate proposals (0 of 4 groups constant), but its label changed between sampled +# repeats on 10 of 27 statements, so the number it produces is not comparable across +# runs. Until one of those is fixed the score is written to novelty_scores.jsonl and +# read there. Set 0.5 to bring it back into the reward. +NOVELTY_FLOOR="${NOVELTY_FLOOR:-1}" + +LEARNING_RATE="${LEARNING_RATE:-1e-6}" +SIDES="${SIDES:-both}" + +ROOT="output/rsi_agentic/${TAG}" +# One checkpoint directory for the whole loop, overwritten every iteration, so the +# disk holds one 4B model rather than one per iteration. The previous round's +# weights are gone once the next save starts: if a save dies partway there is +# nothing to fall back to but BASE_MODEL. +CKPT_DIR="$ROOT/ckpt" +# Written with ${VAR-default} rather than ${VAR:-default} so that TASK_BANK="" +# means off; with the colon an empty value would silently get the default back. +TASK_BANK="${TASK_BANK-$ROOT/task_bank.jsonl}" +mkdir -p "$ROOT" + +# Pick up where a previous invocation left off. The iteration number comes from a +# marker written after the checkpoint has been checked, not from train_summary.json, +# which is written at the end of a training run but would still be there after a +# crash in a later stage. +MODEL="$BASE_MODEL" +START=1 +while [ -f "$ROOT/iter${START}/iteration.done" ]; do + START=$((START + 1)) +done +if [ "$START" -gt 1 ]; then + if [ -f "$CKPT_DIR/model/config.json" ]; then + MODEL="$CKPT_DIR/model" + else + echo "$((START - 1)) iteration(s) finished under $ROOT but no checkpoint at" >&2 + echo "$CKPT_DIR/model -- each iteration overwrites the one before, so those" >&2 + echo "weights are gone. Start a new TAG, or delete the iteration.done" >&2 + echo "markers to redo them from $BASE_MODEL." >&2 + exit 1 + fi +fi + +cat <<EOF +=== repo $REPO +=== gpus $GPUS (devices $DEVICES) +=== sandbox $AENV_API_URL template $AENV_TEMPLATE slots $SANDBOX_SLOTS +=== api $API_MODEL at $API_BASE, thinking budget $API_THINKING_BUDGET +=== per iter $KEEP_GROUPS groups of $GROUP_SIZE, band [1, $((SOLVER_ROLLOUTS - 1))] of $SOLVER_ROLLOUTS +=== build cap $([ "$MAX_BUILD_FILES" -eq 0 ] && echo "none" || echo "$MAX_BUILD_FILES files, in the system prompt") +=== novelty $([ -z "$TASK_BANK" ] && echo "off" || echo "bank $TASK_BANK, floor $NOVELTY_FLOOR") +=== trains on $((KEEP_GROUPS * GROUP_SIZE)) propose + $((KEEP_GROUPS * SOLVER_ROLLOUTS)) solve trajectories, one step, lr $LEARNING_RATE +=== checkpoint $CKPT_DIR/model, overwritten each iteration +=== iterations $([ "$ITERATIONS" -eq 0 ] && echo "until killed" || echo "$ITERATIONS") +=== swanlab ${RSI_SWANLAB_MODE:-online} project ${RSI_SWANLAB_PROJECT:-twinkle-rsi-agentic}, experiment $TAG, one step per iteration +=== starting at iteration $START from $MODEL +EOF + +i="$START" +while [ "$ITERATIONS" -eq 0 ] || [ "$i" -lt $((START + ITERATIONS)) ]; do + OUT="$ROOT/iter${i}" + mkdir -p "$OUT" + echo "=== iteration $i: collect $KEEP_GROUPS groups from $MODEL -> $OUT" + + CUDA_VISIBLE_DEVICES="$DEVICES" python cookbook/rsi/agentic/challenge.py \ + --model-id "$MODEL" \ + --sampler-gpus "$GPUS" \ + --sandbox-slots "$SANDBOX_SLOTS" \ + --keep-groups "$KEEP_GROUPS" \ + --group-size "$GROUP_SIZE" \ + --solver-rollouts "$SOLVER_ROLLOUTS" \ + --max-build-files "$MAX_BUILD_FILES" \ + --api-model "$API_MODEL" \ + --api-base "$API_BASE" \ + --api-thinking-budget "$API_THINKING_BUDGET" \ + --task-bank "$TASK_BANK" \ + --novelty-floor "$NOVELTY_FLOOR" \ + --out-dir "$OUT" \ + --keyword-db "$ROOT/keywords.jsonl" \ + 2>&1 | tee "$OUT/challenge.log" + + echo "=== iteration $i: train on $OUT -> $CKPT_DIR" + RSI_RUN_DIR="$OUT" \ + RSI_SAVE_DIR="$CKPT_DIR" \ + RSI_SAVE_NAME="model" \ + RSI_SIDES="$SIDES" \ + RSI_TAG="$TAG" \ + RSI_ITER="$i" \ + CUDA_VISIBLE_DEVICES="$DEVICES" python cookbook/rsi/agentic/train.py \ + --model_id "$MODEL" \ + --model_gpus "$GPUS" \ + --lr "$LEARNING_RATE" \ + 2>&1 | tee "$OUT/train.log" + + # HF-format weights plus tokenizer, which is what --model-id takes, so the next + # iteration needs no conversion step. + MODEL="$CKPT_DIR/model" + if [ ! -f "$MODEL/config.json" ]; then + echo "iteration $i saved no loadable checkpoint at $MODEL" >&2 + exit 1 + fi + # Written last, so resuming counts only iterations whose weights are on disk. + touch "$OUT/iteration.done" + echo "=== iteration $i done; next starts from $MODEL" + i=$((i + 1)) +done +echo "=== stopped after iteration $((i - 1)); model at $MODEL" diff --git a/cookbook/rsi/agentic/prompts.py b/cookbook/rsi/agentic/prompts.py index 890271a7f..61c61a9e2 100644 --- a/cookbook/rsi/agentic/prompts.py +++ b/cookbook/rsi/agentic/prompts.py @@ -1,33 +1,31 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Prompts for the agentic challenger. +"""Every string this pipeline sends to a model. -One conversation, three stages: - Stage 1: the model acts as an agent in a sandbox (multi-turn with tools), - producing a tool-call chain and a final workspace state, and stops - calling tools. - Stage 2: a user message carrying the real workspace listing is appended to that - same conversation, asking for a python check script. - Stage 3: another user message asks for the problem statement. +One proposal is one conversation with three stages: -Stages 2 and 3 are appended rather than sent as fresh calls, so the whole chain -is one sample whose every assistant turn can be trained on. That is also why -their rules live in user messages: a conversation has one system message, and it -was already spent on stage 1. + 1. the model acts in a sandbox, one tool call per reply, until it stops calling + tools (``SYSTEM`` + ``FROM_KEYWORDS``); + 2. a user message carrying the real workspace listing asks for a python check + script (``CHECK_FOLLOWUP``, once more via ``CHECK_RETRY_FOLLOWUP`` if it does + not pass); + 3. a user message asks for the problem statement (``PROBLEM_FOLLOWUP``). -The check script is run against the sandbox immediately after stage 2 to verify -it passes; any task whose own checks fail is thrown away. +Stage 1 runs on the local model and is what gets trained. Stages 2 and 3 are +appended to a copy of that same conversation and answered by the API model, so the +check script and the statement are written with the whole build history in view +without adding untrained tokens to the sample. -Keyword categories are configurable; the framework (``KeywordStore`` + draw/combine -logic) is category-agnostic. A proposal draws one entry from each category, and the -three are facets of ONE task so combining them yields a single non-trivial -computation: - - transform: the computation the task turns on - - domain: the material it runs over (data AND code/compilation) - - edge_case: the twist that makes a naive solution fail +The texts below are unchanged from the pipeline this replaced; the comments +keep the measurements that decided their wording, because a prompt whose numbers are +lost is a prompt nobody can edit safely. """ -from twinkle_agentic.challenger import AgenticPrompts -# ── Keyword categories (framework-agnostic; only the content is scenario-specific) ─ +# ── Keyword categories ───────────────────────────────────────────────────── +# A proposal draws one entry from each category, and the three are facets of ONE +# task, so combining them yields a single non-trivial computation: +# transform: the computation the task turns on +# domain: the material it runs over +# edge_case: the twist that makes a naive solution fail CATEGORIES = ['transform', 'domain', 'edge_case'] @@ -63,11 +61,36 @@ 'rounding, byte order, cycles in a tree' + _LEAVE_THE_EXAMPLES, } -# ── Stage 1: model acts in sandbox ───────────────────────────────────────── +# ── Keyword generation ───────────────────────────────────────────────────── +# ``parse_keyword_list`` reads a JSON array and returns nothing when it cannot find +# one, so these must ask for a JSON array; keep them in step with the parser. + +KEYWORD_SYSTEM = ( + 'You generate diverse topic keywords for training an AI agent that does ' + 'computer engineering work in a Linux sandbox: writing and running programs, ' + 'building, testing and debugging software, processing and analysing data, ' + 'and administering files and the system.' +) + +KEYWORD_USER = ( + 'List {k} diverse, specific topic keywords for the following category:\n' + '{desc}\n\n' + 'Each should be 2-5 words, concrete enough to inspire a specific task. ' + 'Return ONLY a JSON array of short strings, nothing else.' +) + +KEYWORD_EXPAND_USER = ( + 'The keyword "{kw}" produced a very hard task. List {m} related keywords ' + 'in the same domain that might produce similarly challenging but different ' + 'tasks. Return ONLY a JSON array of short strings, nothing else.' +) + + +# ── Stage 1: the model acts in the sandbox ───────────────────────────────── # Only shell_executor and python_executor exist -- there is no directory-listing # tool -- so the prompt names ``ls -R`` explicitly and spells tool names in full. -# One tool call per message: the sampler stops generation at ``</tool_call>``, and -# a reply that plans many calls but is cut after the first would be trained on a +# One tool call per message: the sampler stops generation at ``</tool_call>``, and a +# reply that plans many calls but is cut after the first would be trained on a # reasoning that does not match what happened. SYSTEM = ( @@ -97,19 +120,16 @@ '"Done." as your final message.' ) -FROM_SCRATCH = ( - 'Build something complex and realistic in the current empty directory. ' - 'Create files, write scripts, ' - 'process data -- whatever demonstrates ' - 'competent use of the tools. Take as many turns as the work needs, each ' - 'building on the last.' -) - -FROM_SEED = ( - 'Here is an example of the kind of task we want:\n\n{seed}\n\n' - 'Build something in the same spirit but on a ' - 'different subject, equally complex and realistic. Change what ' - 'is produced and how, not just the names. Work in the current empty directory.' +# A cap on volume only: the failure it targets is a smaller model running out of +# tokens writing many files, and the thing that must survive is the computation. +# Appended to SYSTEM when --max-build-files > 0, which the loop sets to 4. +BUILD_SIZE_CAP = ( + '\n- Keep the result SMALL: at most {n} files in total, counting inputs, ' + 'scripts and outputs. No python package (no __init__.py, no importable ' + 'module tree), no command-line interface with subcommands. Depth, not ' + 'volume: one non-trivial computation done properly on a small input beats ' + 'many files. A task built from this state has to be finishable by a smaller ' + 'model in about twenty tool calls.' ) FROM_KEYWORDS = ( @@ -119,48 +139,54 @@ 'current empty directory, producing files and/or computed output.' ) -FROM_SEED_KEYWORDS = ( - 'Here is an example task for inspiration:\n\n{seed}\n\n' - 'Your direction keywords:\n{keywords}\n\n' - 'Build something complex and realistic that combines the spirit of the ' - 'example with the keyword topics. ' - 'Work in the current empty directory.' -) - # ── Stage 2: write the check script ──────────────────────────────────────── -# Appended to the episode as a user message once the model stops calling tools, -# so it says "you": the same conversation did the work. brittle_check_reason() in -# challenger/agentic.py rejects size/checksum/source-text asserts on the syntax -# tree and sends the script back through the rewrite path. +# Appended to the conversation once the model stops calling tools, so it says +# "you": the same conversation did the work. ``brittle_check_reason`` rejects +# size/checksum/source-text asserts on the syntax tree and sends the script back +# through the rewrite path. +# +# The substring rule used to read "that an expected substring is present", next to a +# separate ban on "script source text". Asserting that a .py file contains +# 'def worker():' satisfies the permission and violates the ban, and on 188 tasks +# from run_clean9 the permission won 51% of the time. Merging the two rules and +# adding the subprocess instruction took that to 0% of 50 tasks, with 84% of the new +# checks running a program instead of reading one. +# +# Then shortened from 475 words to 261 by dropping every sentence that argued FOR a +# rule while keeping the rule. Measured on the same 50 workspaces, temperature 1.0: +# asserts .py source text runs a program input data handed over +# 475 words 0% 86% 86% +# 261 words 2% 90% 89% +# The noise floor from sampling one prompt twice is 2 points on the source-text rate +# and 14 on handover, so nothing moved. CHECK_FOLLOWUP = ( 'Now write a python script that ASSERTS properties of the state you just ' - 'produced. It will be run in the same directory you worked in.\n\n' - 'Here is the actual final state of that directory: first every file as ' - '"path size-in-bytes", then the contents of each one. This listing is the ' - 'ground truth, not your account of what you did. Assert only about paths ' - 'that appear in it, and only about content you can read here. If it is empty ' - 'or shows nothing worth testing, say UNTESTABLE and write no code.\n\n' + 'produced. It runs in the same directory you worked in.\n\n' + "Below is that directory's actual final state: every file as " + '"path size-in-bytes", then each one\'s contents. This is the ground truth, ' + 'not your account of what you did. Assert only about paths and content ' + 'visible here. If there is nothing worth testing, say UNTESTABLE and write ' + 'no code.\n\n' '{final_state}\n\n' + 'The solver will be told what its program must DO and writes its own code, ' + 'so two correct programs share their behaviour and nothing else.\n\n' 'Rules:\n' '- 2-6 asserts, standard library only.\n' - '- Make the check ROBUST and BROAD: it must pass for ANY correct ' - 'reproduction of this state, and fail only for one that got the work wrong. ' - 'Assert meaning, not form -- that a file exists, that it parses, that a value ' - 'or a row read out of it is right, that an expected substring is present.\n' - '- Do NOT pin exact bytes: no file sizes, no checksums, no asserting that a ' - 'whole file equals one exact string, no timestamps, no script source text. A ' - 'different correct solution writes different bytes and would fail such a ' - 'check even though it is right.\n' - '- Do NOT constrain the directory as a whole: never assert the exact number ' - 'of files, or that no other files exist. Check only the files that carry the ' - 'result and ignore the rest.\n' - '- Never write down a number you did not read above -- do not recompute a ' - 'mean, a count or a checksum in your head.\n' - '- A file shown truncated has more content than you can see: assert about the ' - 'part you were shown, not its end or its length.\n' - '- Still discriminating: what you keep must fail for a directory that does ' - 'not hold this state. Robust does not mean empty.\n' + '- Assert only about files holding RESULTS. Never about the text of a ' + 'program: not a line it contains, not a name it mentions, not its length.\n' + '- If a result only exists once a program runs, RUN it -- ' + 'subprocess.run([sys.executable, "thing.py"], capture_output=True, ' + 'text=True) -- and assert on what it printed or the files it left.\n' + '- No exact bytes: no sizes, no checksums, no whole-file equality, no ' + 'timestamps.\n' + '- No claim about the directory as a whole: not the file count, not that ' + 'nothing else exists.\n' + '- Never write a number you did not read above.\n' + '- A truncated file holds more than you can see: assert about the shown ' + 'part, not its end or its length.\n' + '- Still discriminating: it must fail for a directory that does not hold ' + 'this state.\n' '- Exit 0 when every assertion holds, non-zero otherwise.\n' '- Do NOT call any tool now. Return ONLY a fenced python code block, no ' 'prose.' @@ -188,76 +214,46 @@ 'block, no prose.' ) - -# ── Keyword generation ───────────────────────────────────────────────────── -# ``parse_keyword_list`` reads a JSON array and returns nothing when it cannot -# find one, so these must ask for a JSON array; keep them in step with the parser. - -KEYWORD_SYSTEM = ( - 'You generate diverse topic keywords for training an AI agent that does ' - 'computer engineering work in a Linux sandbox: writing and running programs, ' - 'building, testing and debugging software, processing and analysing data, ' - 'and administering files and the system.' -) - -KEYWORD_USER = ( - 'List {k} diverse, specific topic keywords for the following category:\n' - '{desc}\n\n' - 'Each should be 2-5 words, concrete enough to inspire a specific task. ' - 'Return ONLY a JSON array of short strings, nothing else.' -) - -KEYWORD_EXPAND_USER = ( - 'The keyword "{kw}" produced a very hard task. List {m} related keywords ' - 'in the same domain that might produce similarly challenging but different ' - 'tasks. Return ONLY a JSON array of short strings, nothing else.' -) - - -# ── Arm C: cap what one episode may build ────────────────────────────────── -# A cap on volume only: the failure it targets is a smaller model running out of -# tokens writing many files, and the thing that must survive is the computation. - -BUILD_SIZE_CAP = ( - '\n- Keep the result SMALL: at most {n} files in total, counting inputs, ' - 'scripts and outputs. No python package (no __init__.py, no importable ' - 'module tree), no command-line interface with subcommands. Depth, not ' - 'volume: one non-trivial computation done properly on a small input beats ' - 'many files. A task built from this state has to be finishable by a smaller ' - 'model in about twenty tool calls.' -) - -# ── Stage 3: the statement gives the rules, never the computed answer ──────── +# ── Stage 3: the statement gives the rules, never the computed answer ────── # The end state is split in two: input data verbatim (it is not the answer), and -# everything derived given as the rule that produces it -- otherwise the only way -# to state what a derived file must contain is to quote the computed answer. - -PROBLEM_FOLLOWUP_RULES_ONLY = ( +# everything derived given as the rule that produces it -- otherwise the only way to +# state what a derived file must contain is to quote the computed answer. +# +# Stating the split as a rule is not enough: over run_clean9's 154 tasks whose check +# compares against a computed-looking value, 52% of statements carry EVERY one of +# them (mean share 0.72). Listing the values instead of describing them was tried on +# 50 tasks in two forms and both differences sat inside the noise floor (0.059, +# p=0.10 to 0.53), so the wording is unchanged and the leak rate is a known open +# problem. A forbidden list can also hide INPUT data, which makes a task unsolvable, +# and that cost is invisible to every offline metric. +# +# Shortened from 328 words to 236 in the same round as CHECK_FOLLOWUP. Measured on +# the same 50 workspaces, temperature 1.0: +# input data handed over leak statement words p50 +# 328 words 86% 0.57 208 +# 236 words 91% 0.64 201 +# 236 + short check 86% 0.54 198 +# Handover moved 5 points against a 14-point same-prompt spread and the leak 0.07 +# against 0.059, p=0.50 -- a length change that cost nothing measurable. + +PROBLEM_FOLLOWUP = ( 'Your checks pass on the state you produced. Now write the task description ' 'another AI agent would be given to reproduce that same end state.\n\n' - 'That agent starts in an EMPTY directory and sees nothing but your ' - 'statement: every file that must be there at the end has to be created by ' - 'it.\n\n' + 'It starts in an EMPTY directory and sees nothing but your statement: every ' + 'file that must be there at the end has to be created by it.\n\n' 'Give the two halves differently:\n' '- INPUT data, the raw material nothing was computed from yet: verbatim, ' 'exact filenames and exact contents, so it can be written byte for byte. ' 'Only passive data counts as input -- a CSV, a JSON config, a binary record ' 'file, a text corpus. Source code is NEVER input data: do not quote the ' - 'body of any script, function or module you wrote, not even one you call a ' - 'fixture. A statement whose input half is the program asks the reader to ' - 'retype your solution, and then it measures typing just as surely as quoting ' - 'a computed answer does.\n' + 'body of any script you wrote.\n' '- Everything DERIVED from it -- computed values, aggregates, orderings, ' - 'resolved references, reports: only the RULE that produces it. Name the ' - 'output file and its format, say how each part follows from the input, and ' - 'never state the resulting value. Not as an example, not in a sample of the ' - 'output. A statement that writes out what you computed can be satisfied by ' - 'copying it, and then it measures typing.\n\n' + 'reports: only the RULE that produces it. Name the output file and its ' + 'format, say how each part follows from the input, and never state the ' + 'resulting value, not even as an example.\n\n' 'Rules:\n' '- Be specific about formats, filenames and layout.\n' '- Say what must be true of the result, not which commands to run.\n' - '- Describe the behaviour any script must have -- its inputs, its outputs, ' - 'the transformation between them -- and let the reader write the code.\n' '- Do NOT mention the checks or how verification works.\n' '- Self-contained: no reference to this conversation or to anything the ' 'reader cannot see.\n' @@ -265,29 +261,3 @@ '- Do NOT call any tool now. Return ONLY the problem statement as plain ' 'text, no code fences.' ) - - -# ── Factory ──────────────────────────────────────────────────────────────── - -def agentic_prompts(max_build_files: int = 0) -> AgenticPrompts: - """Assemble all strings into the object the challenger takes. - - ``max_build_files`` is the one knob: when > 0 it appends BUILD_SIZE_CAP to the - system prompt. - """ - system = SYSTEM - if max_build_files > 0: - system = system + BUILD_SIZE_CAP.format(n=max_build_files) - return AgenticPrompts( - system=system, - from_scratch=FROM_SCRATCH, - from_seed=FROM_SEED, - from_keywords=FROM_KEYWORDS, - from_seed_keywords=FROM_SEED_KEYWORDS, - check_followup=CHECK_FOLLOWUP, - check_retry_followup=CHECK_RETRY_FOLLOWUP, - problem_followup=PROBLEM_FOLLOWUP_RULES_ONLY, - keyword_system=KEYWORD_SYSTEM, - keyword_user=KEYWORD_USER, - keyword_expand_user=KEYWORD_EXPAND_USER, - ) diff --git a/cookbook/rsi/agentic/rl.py b/cookbook/rsi/agentic/rl.py deleted file mode 100644 index 9621850f2..000000000 --- a/cookbook/rsi/agentic/rl.py +++ /dev/null @@ -1,286 +0,0 @@ -"""Agentic RSI: GRPO on multi-turn tool-using episodes, scored by program checks. - -The solver is an ms-agent agent with a real tool line-up (shell, filesystem, -python, notebook sandbox, todo list) working in its own microVM. It explores for -as many turns as it needs; when it stops calling tools the episode ends and the -task's checks are run against what it left behind. The checks are ordinary -programs -- file exists, file content, command exit status, final answer match -- -so the same trajectory always earns the same reward. - -The two halves are split by where they run, not by what they know: - - * On the training host, ``MsAgentHarness`` shapes messages -- system prompt, - tool-result formatting, ms-agent's own message evolution. Its ``llm`` and - ``tools`` sections are dropped before it prepares, so it constructs no tool - and nothing the model emits can execute next to the trainer. - * In the sandbox, ``sandbox_server/tool_server.py`` holds the real ms-agent - ``ToolManager``. It also *supplies the tool schemas*, which the prompt then - advertises verbatim -- the contract the model is trained against is read off - the code that will honour it, so the two cannot drift apart. - -Usage: - AENV_API_URL=http://127.0.0.1:8000 \\ - RSI_TASKS=cookbook/rsi/agentic/tasks.example.jsonl \\ - python cookbook/rsi/agentic/rl.py - -See README.md for building the template and starting the sandbox server. - -Task file: one JSON object per line, with ``id``, ``query`` and either -``check_script`` (a python script, from ``challenge.py``) or ``checks`` -(structured, see twinkle_agentic.verifier.result_check.Check). -""" -import os -import shutil -from typing import Any, Dict, List - -from peft import LoraConfig - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_device_placement, get_logger -from twinkle.advantage import GRPOAdvantage -from twinkle.checkpoint_engine import CheckpointEngineManager -from twinkle.cli import CLI -from twinkle.data_format import SamplingParams -from twinkle.metric import CompletionRewardMetric -from twinkle.model import TransformersModel -from twinkle.processor import InputProcessor -from twinkle.sampler import vLLMSampler -from twinkle.template import Template -from twinkle_agentic.rollout.multi_turn import MultiTurnRollout - -# Same directory as this script, which python puts on sys.path when it is run as -# a file. Episode construction and scoring are shared with eval.py so the two -# cannot drift: an eval measuring episodes built differently from training would -# not be measuring the training. -from episode import (SandboxConfig, boot_episodes, load_tasks, # noqa: I100,I202 - score_episodes) - -logger = get_logger() -args = CLI.from_args() - -# ========== Configuration ========== -MODEL_ID = args.model.model_id or 'ms://Qwen/Qwen3-4B' - -MODEL_GPUS = args.infra.model_gpus or 4 -SAMPLER_GPUS = args.infra.sampler_gpus or 4 -NUM_GPUS = MODEL_GPUS + SAMPLER_GPUS - -NUM_GENERATIONS = args.rl.num_generations or 8 -# Per turn, and it has to match the challenger and eval side (both 8192): a -# trajectory generated with less room than the tasks were built with is trained on -# truncated attempts. At 4096, replies ran out mid-<think> before dispatching any -# tool -- 3 of 12 episodes on the generation side, 15 of 50 solver attempts. -MAX_NEW_TOKENS = args.sampling.max_tokens or 8192 -LEARNING_RATE = args.optimizer.learning_rate or 1e-5 -MAX_STEPS = args.training.max_steps or 1000 -BATCH_SIZE = args.training.batch_size or 4 -MINI_BATCH_SIZE = args.training.mini_batch_size or 8 -MICRO_BATCH_SIZE = args.training.micro_batch_size or 2 -GRADIENT_ACCUMULATION_STEPS = args.training.gradient_accumulation_steps or 1 -ADAPTER_NAME = args.lora.adapter_name or 'default' -SAVE_STEPS = args.training.save_steps or 500 -LORA_RANK = args.lora.lora_r or 16 - -# A tool-using episode needs room for observations on top of its own tokens. -MAX_TRAJECTORY_TOKENS = int(os.environ.get('RSI_MAX_TRAJ_TOKENS', 32768)) -MAX_TURNS = int(os.environ.get('RSI_MAX_TURNS', 20)) - -TASKS_PATH = os.environ.get('RSI_TASKS', 'cookbook/rsi/agentic/tasks.example.jsonl') -RUN_DIR = os.environ.get('RSI_RUN_DIR', 'output/rsi_agentic/run') - -# Where episodes run, and how they are scored. Shared with eval.py. -SANDBOX = SandboxConfig.from_env() - -# Keep each episode's downloaded files after scoring. Useful while debugging -# tasks, expensive over a long run. -KEEP_WORKSPACES = os.environ.get('RSI_KEEP_WORKSPACES', '0') == '1' - - - -def main(): - tasks = load_tasks(TASKS_PATH) - logger.info(f'Loaded {len(tasks)} tasks from {TASKS_PATH}') - - device_groups = [ - DeviceGroup(name='model', ranks=list(range(MODEL_GPUS)), device_type='GPU'), - DeviceGroup(name='sampler', ranks=list(range(MODEL_GPUS, 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, dp_size=SAMPLER_GPUS) - twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, groups=device_groups, - lazy_collect=False) - - lora_config = LoraConfig( - target_modules='all-linear', - r=LORA_RANK, - lora_alpha=LORA_RANK * 2, - lora_dropout=0.05, - ) - - # torch_dtype=float32: load master weights in fp32. Training precision is still - # governed by `mixed_precision` (bf16 autocast); fp32 params avoid the numerically - # fragile bf16 optimizer state on this Blackwell + CUDA 13 box. - model = TransformersModel(model_id=MODEL_ID, device_mesh=model_mesh, remote_group='model', - torch_dtype='float32') - model.add_adapter_to_model(ADAPTER_NAME, lora_config, - gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) - 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) - # Both templates get the trajectory budget explicitly. The default is far - # smaller than a tool-using episode: leaving it out makes the sampler refuse - # the trajectory mid-run with `Input length N exceeds max_length 8192`, after - # the step it happened in has already booted its sandboxes. - model.set_template('Template', model_id=MODEL_ID, enable_thinking=True, - max_length=MAX_TRAJECTORY_TOKENS) - - sampler = vLLMSampler( - model_id=MODEL_ID, - engine_args={ - 'gpu_memory_utilization': 0.8, - 'max_model_len': MAX_TRAJECTORY_TOKENS, - 'max_lora_rank': 32, - 'enable_lora': True, - 'enable_tower_connector_lora': True, - }, - device_mesh=sampler_mesh, - remote_group='sampler', - ) - sampler.set_template('Template', model_id=MODEL_ID, enable_thinking=True, - max_length=MAX_TRAJECTORY_TOKENS) - - rollout_template = Template(MODEL_ID, max_length=MAX_TRAJECTORY_TOKENS, enable_thinking=True) - rollout_template.truncation_strategy = 'delete' - - ckpt_manager = CheckpointEngineManager(model=model, sampler=sampler) - sampling_params = SamplingParams(max_tokens=MAX_NEW_TOKENS, num_samples=1, logprobs=1, - temperature=1.0, top_p=0.95) - rollout = MultiTurnRollout( - sampler=sampler, - template=rollout_template, - sampling_params=sampling_params, - max_turns=MAX_TURNS, - max_trajectory_tokens=MAX_TRAJECTORY_TOKENS, - ) - - advantage_fn = GRPOAdvantage() - metrics = CompletionRewardMetric() - - optim_step = 0 - task_cursor = 0 - logger.info(f'Starting agentic RSI GRPO (max_turns={MAX_TURNS}, ' - f'score={SANDBOX.score_mode})') - logger.info(f'Sandboxes: template={SANDBOX.template} api={SANDBOX.api_url} ' - f'concurrency={SANDBOX.concurrency}') - logger.info(get_device_placement()) - - while optim_step < MAX_STEPS: - metrics.reset() - - # Each prompt is repeated NUM_GENERATIONS times; GRPO needs a group of - # rollouts on the SAME task to have anything to compare against. - batch_tasks = [tasks[(task_cursor + i) % len(tasks)] for i in range(BATCH_SIZE)] - task_cursor = (task_cursor + BATCH_SIZE) % len(tasks) - episode_tasks = [t for t in batch_tasks for _ in range(NUM_GENERATIONS)] - - harnesses, envs, tool_managers, trajectories = [], [], [], [] - try: - episodes = boot_episodes(episode_tasks, SANDBOX) - except Exception as e: # noqa - # A sandbox that never came up answers every call with an error, so - # the group would score a uniform zero and look like a hard task - # rather than a broken environment. Skip the batch and say so. - logger.warning(f'[Step {optim_step}] {e}; skipping batch') - continue - for harness, env, tool_manager, trajectory in episodes: - harnesses.append(harness) - envs.append(env) - tool_managers.append(tool_manager) - trajectories.append(trajectory) - - ckpt_manager.sync_weights(merge_and_sync=False) - sampler.reset_prefix_cache() - - try: - outs: List[Dict[str, Any]] = rollout( - trajectories, harness=harnesses, tool_manager=tool_managers) - - rewards = score_episodes(episode_tasks, envs, outs, - os.path.join(RUN_DIR, f'step{optim_step:06d}'), SANDBOX) - finally: - # Sandboxes are a finite resource; a step that raises must still - # give them back or the next step starts short. - for env in envs: - env.close() - if not KEEP_WORKSPACES: - shutil.rmtree(os.path.join(RUN_DIR, f'step{optim_step:06d}'), ignore_errors=True) - - all_old_logps, completion_lengths, turns = [], [], [] - for traj in outs: - logprobs = traj.get('logprobs') or [] - all_old_logps.append([lp[0][1] for lp in logprobs] if logprobs else []) - labels = traj.get('labels') or [] - completion_lengths.append(sum(1 for label in labels if label != -100)) - turns.append(int(traj.get('turns') or 0)) - - advantages = advantage_fn(rewards, num_generations=NUM_GENERATIONS, - scale='group').tolist() - metrics.accumulate(completion_lengths=completion_lengths, rewards={'total': rewards}) - - avg_reward = sum(rewards) / len(rewards) if rewards else 0.0 - solved = sum(1 for r in rewards if r >= 1.0) - logger.info(f'[Step {optim_step}] avg_reward={avg_reward:.3f} ' - f'fully_solved={solved}/{len(rewards)} ' - f'avg_turns={sum(turns)/max(1,len(turns)):.1f}') - - # Drop episodes the template refused (too long) or that produced no - # trainable tokens; feeding those in would corrupt the logp alignment. - inputs, kept_logps, kept_adv = [], [], [] - for i, traj in enumerate(outs): - if not completion_lengths[i]: - continue - if len(traj.get('input_ids') or []) > MAX_TRAJECTORY_TOKENS: - continue - inputs.append(traj) - kept_logps.append(all_old_logps[i]) - kept_adv.append(advantages[i]) - - if len(inputs) < MODEL_GPUS: - logger.warning(f'[Step {optim_step}] only {len(inputs)} usable trajectories ' - f'(need >= {MODEL_GPUS}); skipping batch') - continue - - # One optimizer step per batch, not per mini-batch. ``forward_backward`` - # neither steps nor zeroes, so the mini-batches below simply add their - # gradients together; ``clip_grad_and_step`` afterwards divides by the - # token count accumulated across all of them, so every trajectory in the - # batch carries the same weight regardless of how the mini-batches split. - # Stepping inside the loop instead -- which is what this used to do -- - # made each step see only MINI_BATCH_SIZE trajectories, so a group of - # NUM_GENERATIONS could be torn across two updates. - for mb_start in range(0, len(inputs), MINI_BATCH_SIZE): - mb_end = min(mb_start + MINI_BATCH_SIZE, len(inputs)) - model.forward_backward( - inputs=inputs[mb_start:mb_end], - old_logps=kept_logps[mb_start:mb_end], - advantages=kept_adv[mb_start:mb_end], - micro_batch_size=MICRO_BATCH_SIZE, - ) - model.clip_grad_and_step() - optim_step += 1 - if optim_step % SAVE_STEPS == 0: - model.save(f'rsi-agentic-checkpoint-{optim_step}') - - log_dict = metrics.calculate() - log_dict.update(model.calculate_metric(is_training=True)) - log_dict['avg_reward'] = avg_reward - log_dict['fully_solved'] = solved - logger.info(f'[Step {optim_step}/{MAX_STEPS}] {log_dict}') - - logger.info(f'Training completed. optim_steps={optim_step}') - model.save('rsi-agentic-final') - - -if __name__ == '__main__': - main() diff --git a/cookbook/rsi/agentic/rsi_agent_shortprompt.yaml b/cookbook/rsi/agentic/rsi_agent_shortprompt.yaml new file mode 100644 index 000000000..cace8cbe2 --- /dev/null +++ b/cookbook/rsi/agentic/rsi_agent_shortprompt.yaml @@ -0,0 +1,197 @@ +# ms-agent config for agentic RSI training. +# +# Read by both halves of the setup, which is the point: +# +# * the training host loads it to build a MsAgentHarness for message shaping +# only -- the entry script drops `llm:` and `tools:` from the merged config +# first, so no tool is ever constructed next to the trainer; +# * remote_tool_env.py uploads this same file into each sandbox, where +# sandbox_server/tool_server.py loads it and does construct the tools. +# +# So the tool line-up below describes what runs in the microVM. Editing it takes +# effect on the next episode; no image rebuild is involved. + +prompt: + # Replaces ms-agent's BASE_AGENT_PROMPT (prompting/builtin.py) for the SOLVER + # only -- the proposing episode gets prompts.py's own SYSTEM through the + # challenger, and never reads this field. That built-in prompt is written for a + # general assistant sitting in a user's workspace, and two of its lines work + # against being a solver: "First decide whether the task needs tools. If you can + # answer reliably from what you know ... just answer", and "Ask first when it + # isn't [safe]". Here there is no one to ask (interactive: false) and answering + # without touching the directory is always wrong. + # + # The paragraph about the empty directory is what 5 of armA2shellV5's 8 + # unsolved tasks needed. Their statements listed a file under "Input data:" + # and the solver read that as "already present" -- in 5a70b77f it created the + # file the rules told it to generate and left the two listed as input alone, + # so it was not confused about being in an empty directory, it was following + # the statement's own division of labour. Nothing in the statement or the + # prompt said that division does not survive into its workspace. + # SHORT ARM of the solver-prompt A/B. 242 words down to 190, by the same + # subtractive rule used on prompts.py's two followups: every instruction stays, the + # sentences arguing for it go. Only 21% shorter, and that is the ceiling -- the + # empty-directory paragraph is 90 of the 190 words and is kept whole, because 5 of + # armA2shellV5's 8 unsolved tasks needed exactly it. What left: 'nobody is + # watching', 'before anything can read them', 'Do not assume', 'however clearly you + # can describe what the answer would be', 'Check each thing the task asked for is + # actually there', and the How-to-work heading with its bullets folded into prose. + # + # This file differs from rsi_agent.yaml in this field ONLY -- tools, permission, + # timeouts and output_dir are identical (verified by comparing the parsed configs + # with prompt removed), so n_pass measured against the two is measuring the prompt. + system: | + You are a command-line agent working inside a fresh Linux container. You are + given one task and you carry it out by running commands and writing files. + Nobody can answer a question, so never ask one and never stop to confirm: + decide and act. + + Your working directory starts COMPLETELY EMPTY. Every file the task + mentions -- including files it describes as inputs, given data, existing + configuration, or material you are handed -- does not exist yet. You have to + create all of them yourself, with exactly the names and contents the task + specifies. A task that shows you the contents of a file is telling you what to + write into it, not telling you it is there. + + Start by listing the directory to see the real state. Create every file the + task names, then do the computation it asks for and write the results it asks + for. Before you finish, list the directory again and read back what you wrote; + if something is missing, fix it rather than reporting success. Answering in + prose without creating files is a failure. Never state a value you did not + compute from the data. + +personalization: + # Off: SOUL/AGENTS/PROFILE.md from the developer's own workspace would leak + # machine-specific context into every training prompt. + enabled: false + +# One turn == one sampler call. MultiTurnRollout's max_turns is the real limit; +# this only stops ms-agent from imposing a lower one. +max_chat_round: 9999 + +# Never wait on a human: training runs unattended. +interactive: false +permission_mode: auto + +# How long ms-agent waits around one tool call. Written down rather than left to +# its default (tool_manager.py TOOL_CALL_TIMEOUT, 120s, overridable by the +# TOOL_CALL_TIMEOUT environment variable) so the sandbox does not inherit a +# number from whatever shell started it. It has to stay below what +# remote_tool_env allows the whole turn (command_timeout, 180s), which in turn is +# below the transport's budget: the innermost layer should be the one that times +# out, because it is the only one that knows which call was slow. When they were +# equal, one command that never returns made every call in the turn read as an +# unreachable runtime. +tool_call_timeout: 120 + +# Path *inside the sandbox*. One microVM per episode already isolates +# trajectories from each other, so this is a fixed path rather than a per-slot +# directory; the entry script overrides it only to match --workspace. +output_dir: /workspace + +callbacks: [] + +tools: + # `file_system` is NOT listed here and is nevertheless on. ms-agent's own + # ms_agent/agent/agent.yaml declares it (write_file, read_file, edit_file, + # grep, glob) and LLMAgent merges this file *over* that one, so omitting a key + # inherits it rather than dropping it. Measured: the merged config's tools are + # ['file_system', 'code_executor', 'todo_list'], and /tools advertises all ten + # of those tools to the model. In armA2shellV6's 128 proposing calls, + # file_system took 63 (43 of them write_file) against code_executor's 58. + # + # So the paragraph that used to be here -- claiming the five were removed to + # stop write_file being the path of least resistance -- described a state that + # never existed, through the arms named A2shell*, whose whole premise was + # "shell and python only". Turning it off takes an explicit + # `file_system: {enabled: false}`, which _tool_on (tool_manager.py:47) reads. + # Left on for now, deliberately and with the effect known. + code_executor: + mcp: false + # python_env means "run in this process's machine", and that machine is the + # microVM -- the sandbox boundary is the VM itself, not this setting. Do not + # switch to the docker implementation: it would nest a container inside the + # VM for no extra isolation. + implementation: python_env + include: + - shell_executor + # Kept alongside the shell so that writing a file does not depend on + # getting a heredoc right. Dropping notebook_executor because it overlaps + # this one and adds a cell-state model nothing here needs. + - python_executor + todo_list: + mcp: false + # Kept out of the workspace root. The plan files default to + # `<output_dir>/plan.json` and `plan.md`, and output_dir *is* the directory + # whose end state becomes the task: 2 of ex11's 36 proposals wrote checks + # asserting the agent's own todo bookkeeping, one of them pinning + # `updated_at`, which no solver can reproduce. `.ms_agent/` is where + # ms_agent/project/paths.py says framework internals belong, and the + # workspace listing already skips it. + plan_filename: .ms_agent/plan.json + plan_md_filename: .ms_agent/plan.md + +# Every refusal ms-agent applies to a shell command, turned off. `allow_network` +# and the two list-valued keys are read by LLMAgent.prepare_runtime (llm_agent.py +# builds PermissionConfig.from_dict off this section) and take effect in the +# sandbox, where tool_server.py loads this same file. +# +# The last two keys are different: ms-agent's SafetyConfig does NOT implement +# them, and from_dict ignores unknown keys without a word, so on their own they +# would be dead letters. tool_server.py reads them itself and applies the two +# relaxations as a runtime patch (_patch_permission) inside the sandbox -- +# ms-agent is a harness twinkle supports, so it is used as released rather than +# forked. The startup line reports which ones took effect. +# +# The reason is what the refusals cost here rather than what they protect: this +# runs in a microVM that is reset once per episode and holds nothing but the +# workspace, while each refusal rules out a whole family of tasks the model could +# otherwise pose. `curl`/`wget` blocked means no task can fetch a source tarball +# or a dataset; the rm rules mean it cannot clear a directory (`rm -rf *` and +# `rm -rf build/*` are both refused) or write a task that starts from a mess that +# has to be cleaned up. +permission: + # Drops the default blacklist wholesale: curl, wget, ssh, scp, rsync, nc, + # netcat. (Whether the microVM actually has a route out is a separate + # question from whether the command is allowed to run.) + allow_network: true + safety_rules: + # Emptied, replacing the three baked-in patterns: `rm -rf /*`, `mkfs *`, + # `dd if=*`. An empty list here is not the same as an absent key -- absent + # means "use the defaults". + patterns: [] + # Same, for the configurable half of the rm/rmdir path check: `*`, `/*`, + # `/`, `~`. Left empty rather than removed to say the intent out loud; + # unrestricted_removal below bypasses the whole check, this list included, + # and tool_server.py warns at startup if the two ever disagree. + dangerous_removal_paths: [] + # And the half that a config cannot reach, which is why this one needs the + # runtime patch: the refusals written into is_dangerous_removal_path for + # `*`, anything ending in `/*`, `/`, a direct child of `/` (which + # `/workspace` is), and the home directory. + unrestricted_removal: true + # A separate refusal, found by running commands through SafetyGuard rather + # than by reading the config: a glob anywhere in a write or create path is + # denied on its own ("Glob patterns not allowed in write operations"), which + # is what actually stopped `rm -rf build/*` after the two lists above were + # emptied. It is not specific to rm -- `cp src/* dst/` is refused by the same + # check. (`chmod +x bin/*` is NOT: measured through SafetyGuard, chmod's + # arguments are not extracted as write paths, so it was already allowed.) + # Patched by widening the path that is scope-checked to the directory the + # glob expands inside, so a glob still cannot reach outside the workspace -- + # `cp /etc/* /workspace/` stays denied, now for being out of scope. + allow_write_globs: true + +# Web search is deliberately absent. ms-agent's `web_search` key only provides +# fetch_page (retrieve a known URL); a real query-a-search-engine tool needs +# EXA_API_KEY / SERPAPI_API_KEY and is wired separately from the plain tool +# list. Add it here once that is decided; until then no task should need it. + +# No `llm:` section on purpose, and note that omitting it is not the same as +# disabling it: ms-agent merges this file over its own ms_agent/agent/agent.yaml, +# which declares `service: modelscope`. The tool server treats a section with no +# credentials as absent, drops it, and then withdraws the one argument that +# needed it (read_file's `abbreviate`, an LLM-written file summary) from the +# advertised schema -- so the model is never offered a tool argument that cannot +# work. Put a real `llm:` here, with a key reachable from the sandbox, to get +# that argument back. diff --git a/cookbook/rsi/agentic/sandbox.py b/cookbook/rsi/agentic/sandbox.py new file mode 100644 index 000000000..f17ba2a6d --- /dev/null +++ b/cookbook/rsi/agentic/sandbox.py @@ -0,0 +1,292 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""The sandbox as a resource: N microVMs, each one a slot a job can own. + +A slot is owned for as long as a job needs it, because the workspace lives inside +the microVM: from the clear, through every tool call, to the check that runs +against what was left behind. Two jobs sharing a slot would read each other's +files, so the pool hands out whole slots and never sub-divides one. + +The transport underneath is ``remote_tool_env.RemoteMsAgentToolEnv``, which is +paired with the in-sandbox runtime in ``sandbox_server/tool_server.py``. The +solver's opening messages come from ``episode.solver_harness``, the same function +``eval.py`` uses, so a task's difficulty here and its pass rate there are measured +against one opening. +""" +import os +import time +from concurrent.futures import ThreadPoolExecutor +from typing import List, Tuple + +from twinkle import get_logger +from twinkle_agentic.envs import EnvTool +from twinkle_agentic.tools.tool_manager import ToolManager + +from episode import solver_harness # noqa: I100,I202 +from remote_tool_env import RemoteMsAgentToolEnv, tool_payload # noqa: I100,I202 + +logger = get_logger() + +__all__ = ['Sandbox', 'open_pool', 'close_pool', 'solver_harness', + 'CLEAR_WORKSPACE', 'WORKSPACE_SNAPSHOT'] + +# Cleared through the python tool, not `rm -rf`: ms-agent's safety policy rejects +# `rm -rf` outright ("Blocked by safety rule"), and it rejects globs in write +# operations, which rules out `find -delete` too. The script asserts the +# directory really is empty, so a future policy change surfaces as a failed reset +# instead of jobs quietly inheriting the previous workspace. +CLEAR_WORKSPACE = ''' +import os, shutil +root = {workspace!r} +os.makedirs(root, exist_ok=True) +for name in os.listdir(root): + path = os.path.join(root, name) + if os.path.isdir(path) and not os.path.islink(path): + shutil.rmtree(path, ignore_errors=True) + else: + os.remove(path) +leftover = os.listdir(root) +assert not leftover, 'workspace not empty after clear: %r' % (leftover,) +''' + +# The ground truth the check script is written against. A listing alone is not +# enough: three of the six rejected proposals in the first real run failed on a +# value the model recomputed from its own recollection ("Mean values mismatch") +# rather than read off the file, so the end state has to arrive as content, not +# just as names. Bounded on both axes because this goes into a prompt and a 100k +# artifact would push the trajectory it is read alongside out of the window. +# +# Walks the tree in python rather than shelling out to `find`: the same code then +# decides what is text, what is truncated, and what the budget was spent on, +# which a pipeline of find/head cannot report back. +# +# File bodies go out byte for byte. An earlier version printed `body.rstrip()`, +# which hid trailing newlines while the size column still counted them, so a +# check writer shown an 11-byte file whose content looked 10 characters long +# wrote `content == 'Mean: 63.9'` and the check failed against the very state it +# was written from. The listing is only ground truth if it does not tidy up. +# +# Facts *about* a file go in its header, never after its body. A note printed +# below the content is indistinguishable from content: annotated one file with a +# trailing `(no newline at end of file)` line and the next check script asserted +# the README's content ending in that sentence. +WORKSPACE_SNAPSHOT = ''' +import os + +root = {workspace!r} +skip = {{'.ms_agent', '__pycache__', '.ipynb_checkpoints', '.git'}} +rows = [] +for dirpath, dirnames, filenames in os.walk(root): + dirnames[:] = [d for d in dirnames if d not in skip] + for name in sorted(filenames): + path = os.path.join(dirpath, name) + try: + rows.append((os.path.relpath(path, root), os.path.getsize(path), path)) + except OSError: + pass +rows.sort() +for rel, size, _ in rows[:{max_files}]: + print(rel, size) + +budget = {total_budget} +for rel, size, path in rows[:{max_files}]: + if budget <= 0: + break + try: + with open(path, encoding='utf-8') as handle: + text = handle.read({per_file} + 1) + except (OSError, UnicodeDecodeError): + continue # binary or unreadable: the listing already names it + if '\\x00' in text: + continue + body = text[:{per_file}] + budget -= len(body) + # The trailing-newline count is stated for every file, both ways. Saying it + # only when it is absent made "this file ends with a newline" invisible, and + # the check writer then compared exact bytes without one: in ex9 two of the + # three checks that failed their own verification failed on exactly that -- + # the same reply asserted three files, guessed right on the two marked "no + # newline at end" and wrong on the unmarked one. + trailing = len(body) - len(body.rstrip(chr(10))) + if len(text) > len(body): + suffix = ' (first {per_file} bytes)' + elif trailing == 0: + suffix = ' (no newline at end)' + else: + suffix = ' (ends with %d newline character(s))' % trailing + print() + print('--- ' + rel + suffix + ' ---') + print(body, end='') + if not body.endswith(chr(10)): + print() +''' + +# Seconds to wait before asking a sandbox for its workspace listing a second +# time. 62 of run_clean6's 63 snapshot failures were the sandbox answering 410 +# "not proxyable", which is the host having paused it -- worth one more ask, +# since the alternative is throwing the job away. +SNAPSHOT_RETRY_WAIT = 3 + +# Same idea for the workspace clear. Longer, because what it waits out is +# different: a clear times out when ms-agent's per-call limit expires with the +# delete still running, so the second attempt wants the first one's rmtree to +# have drained rather than to race it. +RESET_RETRY_WAIT = 10 + + +class Sandbox: + """One slot: clear the workspace, run a script in it, read it back. + + Not thread-safe on purpose. A slot belongs to whoever holds it, and the pool + hands each one to exactly one worker thread. + """ + + def __init__(self, slot: int, env: RemoteMsAgentToolEnv, schemas: list, + *, snapshot_max_files: int, snapshot_per_file: int, snapshot_budget: int): + self.slot = slot + self.env = env + self.workspace = env.workspace + # The advertised tool contract. Carried on the slot because it goes into + # the prompt: the schemas a trajectory is built with have to be the ones + # the slot it runs on will honour. + self.schemas = schemas + self._runner = env.runner() + # The tools carry the env they dispatch into, so this slot's model turns + # have to go through this slot's manager. + self.tool_manager = ToolManager(EnvTool.from_schemas(env, schemas)) + self._snapshot_script = WORKSPACE_SNAPSHOT.format( + workspace=self.workspace, max_files=snapshot_max_files, + per_file=snapshot_per_file, total_budget=snapshot_budget) + self._clear_script = CLEAR_WORKSPACE.format(workspace=self.workspace) + + def run(self, script: str) -> Tuple[int, str]: + """Run a python script in the workspace; returns (exit code, output).""" + return self._runner(script, 'python') + + def clear(self) -> None: + """Empty the workspace. Raises rather than returning quietly. + + Every caller depends on a clean start: a silent no-op here means a task + inherits the previous task's files, which lets a solver pass without doing + anything and makes the difficulty numbers meaningless. + + This is also the one point where losing the sandbox costs nothing, since + the workspace is about to be emptied regardless -- so a runtime that went + away is rebuilt here rather than ending a run with hours behind it. The + same covers a clear that *fails* on a runtime still answering /health: + run 'rsi' reached iteration 7 and ended on three clears timing out at + ms-agent's per-call limit while the sandbox reported itself healthy. So + the clear is retried, then retried on a deliberately rebuilt sandbox. + """ + if self.env.ensure_ready(): + self._rebind('runtime was unreachable') + code, out = self.run(self._clear_script) + if code != 0: + logger.warning(f'[sandbox {self.slot}] clear failed (exit {code}), ' + f'retrying in {RESET_RETRY_WAIT}s: {out[-200:]}') + time.sleep(RESET_RETRY_WAIT) + code, out = self.run(self._clear_script) + if code != 0: + # Rebuilt rather than retried again: two failures in a row is not the + # transient this waits out, and a fresh sandbox brings a workspace + # that is already empty -- which is all this method is asked for. + logger.warning(f'[sandbox {self.slot}] clear failed twice (exit {code}); ' + f'rebuilding: {out[-200:]}') + self.env.reset() + self.env.n_recoveries += 1 + self._rebind('rebuilt after two failed clears') + code, out = self.run(self._clear_script) + if code != 0: + raise RuntimeError(f'workspace clear failed (exit {code}): {out[-400:]}') + + def snapshot(self) -> Tuple[str, str]: + """The end state as (listing, error). + + Returned as a bare listing, unwrapped from the tool's JSON envelope: the + model has to read it as a directory rather than as a tool result, or it + falls back on what it *believes* it created. + + An empty listing with no error means the workspace really was empty. An + empty listing with an error means it could not be read, and the two are + kept apart because a snapshot that says "empty" when it means "I could not + look" produces tasks whose only true assertion is that nothing happened. + """ + code, out = self.run(self._snapshot_script) + if code != 0: + logger.warning(f'[sandbox {self.slot}] snapshot failed (exit {code}), ' + f'retrying in {SNAPSHOT_RETRY_WAIT}s: {out[-200:]}') + time.sleep(SNAPSHOT_RETRY_WAIT) + code, out = self.run(self._snapshot_script) + if code != 0: + return '', f'workspace snapshot failed (exit {code}): {out[-500:]}' + return tool_payload(out).strip(), '' + + def close(self) -> None: + self.env.close() + + def _rebind(self, why: str) -> None: + """Point the runner and the tools at the sandbox behind this env now.""" + self._runner = self.env.runner() + self.tool_manager = ToolManager(EnvTool.from_schemas(self.env, self.env.tool_schemas())) + logger.warning(f'[sandbox {self.slot}] rebound ({why})') + + +def open_pool( + n: int, + *, + template: str, + api_url: str, + config_path: str, + workspace: str, + sandbox_timeout: int, + snapshot_max_files: int, + snapshot_per_file: int, + snapshot_budget: int, +) -> List[Sandbox]: + """Boot ``n`` slots and return them ready to use. + + Booted in parallel: each is a microVM taking ~10s, and doing them one after + another would put minutes in front of every run. The tool schemas are read + once, off the first slot -- every slot runs the same image, and these go + straight into the prompt, so reading them n times would only add n chances + for the prompt to differ between slots. + """ + if not template: + raise SystemExit('sandbox template is required (--sandbox-template or AENV_TEMPLATE)') + if not api_url: + raise SystemExit('sandbox api url is required (--sandbox-api-url or AENV_API_URL)') + + def _boot(_) -> RemoteMsAgentToolEnv: + env = RemoteMsAgentToolEnv(template=template, config_path=config_path, + api_url=api_url, workspace=workspace, + sandbox_timeout=sandbox_timeout) + env.reset() + return env + + n = max(1, n) + with ThreadPoolExecutor(max_workers=n) as pool: + envs = list(pool.map(_boot, range(n))) + schemas = envs[0].tool_schemas() + slots = [ + Sandbox(i, env, schemas, snapshot_max_files=snapshot_max_files, + snapshot_per_file=snapshot_per_file, snapshot_budget=snapshot_budget) + for i, env in enumerate(envs) + ] + logger.info(f'[sandbox] {len(slots)} slot(s) ready, tools: ' + f'{[(s.get("function") or {}).get("name") for s in schemas]}') + return slots + + +def close_pool(slots: List[Sandbox]) -> int: + """Kill every slot; returns how many rebuilds happened over the run. + + Reported rather than dropped: a run whose sandboxes were rebuilt twenty times + produced its numbers under a different environment than one that was rebuilt + never, and that is invisible from the output files alone. + """ + total = sum(getattr(s.env, 'n_recoveries', 0) for s in slots) + for slot in slots: + try: + slot.close() + except Exception as e: # noqa # best-effort: the backend evicts on timeout anyway + logger.warning(f'[sandbox {slot.slot}] close failed: {e}') + return total diff --git a/cookbook/rsi/agentic/train.py b/cookbook/rsi/agentic/train.py new file mode 100644 index 000000000..297ca14bb --- /dev/null +++ b/cookbook/rsi/agentic/train.py @@ -0,0 +1,363 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""One GRPO step on what challenge.py collected. + +Read ``trajs/index.jsonl``, group it, turn rewards into advantages, accumulate the +whole collection into a single optimizer step, and overwrite the checkpoint the +next iteration loads. + +There is no filtering here. Every rule about what is worth training on was applied +while collecting -- a group is on disk only if it was kept, and a kept group is +exactly its 8 proposals plus the 8 attempts at its selected task -- so anything +this script dropped would be a second, invisible policy on top of that one. What +it does refuse is a trajectory the model cannot be stepped on at all: no logprobs, +no trainable token, a logprob count that disagrees with the trainable count, more +tokens than the model accepts, or a group left with fewer than two members. Each +refusal is named and counted in the summary rather than folded into a total. + +One step, not several: every trajectory here was sampled from one set of weights, +so a second step would be training weights that no longer produced their own data, +``old_logps`` would stop matching, and epsilon would start clipping for a reason +that has nothing to do with the policy being wrong. The cost is update frequency, +one per collection. + + RSI_RUN_DIR=output/rsi_agentic python cookbook/rsi/agentic/train.py \\ + --model_gpus 8 --lr 1e-6 +""" +import collections +import json +import os +from typing import Any, Dict, List + +import numpy as np + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_device_placement, get_logger +from twinkle.advantage import GRPOAdvantage +from twinkle.cli import CLI +from twinkle.processor import InputProcessor + +logger = get_logger() +args = CLI.from_args() + +MODEL_ID = args.model.model_id or 'ms://Qwen/Qwen3-4B' +MODEL_GPUS = args.infra.model_gpus or 8 +# The base text template. Qwen3-4B is text-only and the multimodal subclass +# crashes on encode for it. +TEMPLATE = os.environ.get('RSI_TEMPLATE', 'Template') +# Whatever --lr says, or the CLI's own default. Not written as ``or <number>``: +# the CLI default is never zero, so a fallback here would be dead code that reads +# like the default. +LEARNING_RATE = args.optimizer.learning_rate + +# One trajectory per micro batch, because padding_free is off: a micro batch is +# padded to its longest member, so pairing a short solver attempt with a long +# build episode pays for the long one twice. Builds are whole agentic rollouts -- +# median 7.5k tokens, up to 15.7k -- and at 2 per micro batch a 31k-token padded +# batch died of CUDA OOM with activation recompute already at its most aggressive +# setting. Read from the environment rather than args.training, whose +# micro_batch_size defaults to 2 rather than None. +MICRO_BATCH_SIZE = int(os.environ.get('RSI_MICRO_BATCH_SIZE', 1)) +# forward_backward is declared dispatch='slice_dp', so a mini batch is sliced +# across the data-parallel ranks and each rank collates only its share. That share +# has to hold at least one micro batch, so the floor is MODEL_GPUS * MICRO_BATCH. +MINI_BATCH_SIZE = args.training.mini_batch_size or MODEL_GPUS * MICRO_BATCH_SIZE + +RUN_DIR = os.environ.get('RSI_RUN_DIR', 'output/rsi_agentic') +SAVE_DIR = os.environ.get('RSI_SAVE_DIR', 'output/rsi_agentic/ckpt') +SAVE_NAME = os.environ.get('RSI_SAVE_NAME', 'agentic') +# Longest trajectory fed to the model. Above this the model refuses it mid-step. +MAX_MODEL_LEN = int(os.environ.get('RSI_MAX_MODEL_LEN', 32768)) +# Which side(s) to train: 'both', 'solve', 'propose'. +SIDES = os.environ.get('RSI_SIDES', 'both') + +# swanlab. One experiment for the whole loop rather than one per iteration: the +# question these charts answer is whether iteration k+1 is better than k, which a +# chart that ends after one point cannot show. ``id`` is the tag, so re-running a +# tag appends to its curve and a new tag starts a new one. Both are read from the +# environment because loop.sh is what knows them; the fallbacks parse the run +# directory, which is ``<root>/<tag>/iter<n>``, so a bare ``python train.py`` still +# lands somewhere sensible instead of failing. +SWANLAB_PROJECT = os.environ.get('RSI_SWANLAB_PROJECT', 'twinkle-rsi-agentic') +TAG = os.environ.get('RSI_TAG') or os.path.basename(os.path.dirname( + os.path.abspath(RUN_DIR))) +ITERATION = int(os.environ.get('RSI_ITER') + or ''.join(c for c in os.path.basename(os.path.abspath(RUN_DIR)) + if c.isdigit()) or 0) +# 'disabled' skips it entirely, for a run that should not appear on the dashboard. +SWANLAB_MODE = os.environ.get('RSI_SWANLAB_MODE', 'online') + + +def upload(challenge: Dict[str, Any], training: Dict[str, Any]) -> None: + """Send this iteration's numbers to swanlab, as one step. + + Called after the checkpoint is saved, so a swanlab failure costs this + iteration's charts and not its weights. The cost of that order is the + reverse: a crash between the step and here loses the numbers, which are still + on disk in challenge_metrics.json and train_summary.json. + + Only ``challenge['scalars']`` goes up, not ``challenge['counts']``: the counts + have keys that exist in one iteration and not the next + (``group_dropped:rubric_error``), and a chart that appears halfway through a + run is read as a change in the run rather than a change in what was recorded. + """ + import swanlab + swanlab.init(project=SWANLAB_PROJECT, name=TAG, id=TAG, resume='allow', + mode=SWANLAB_MODE, config={'tag': TAG, 'model_id': MODEL_ID, + 'learning_rate': LEARNING_RATE, + 'sides': SIDES, 'gpus': MODEL_GPUS}) + log = {f'challenge/{k}': v for k, v in challenge.items()} + log.update({f'train/{k}': v for k, v in training.items()}) + swanlab.log(log, step=ITERATION) + logger.info(f'[train] swanlab {SWANLAB_PROJECT}/{TAG} step {ITERATION}: ' + f'{len(log)} metrics') + + + +def load(run_dir: str) -> tuple: + """Read the index into GRPO groups; returns (groups, skipped). + + A group is ``(side, group_id)`` for the proposing side and + ``(side, group_id, proposal_idx)`` for the solving side -- one prompt answered + several times, which is what an advantage is computed over. + """ + traj_dir = os.path.join(run_dir, 'trajs') + index = os.path.join(traj_dir, 'index.jsonl') + if not os.path.exists(index): + raise SystemExit(f'[train] no {index}; run challenge.py first') + wanted = {'both': ('propose', 'solve')}.get(SIDES, (SIDES, )) + skipped: collections.Counter = collections.Counter() + by_key: Dict[Any, List[Dict[str, Any]]] = collections.OrderedDict() + with open(index, encoding='utf-8') as f: + for line in f: + line = line.strip() + if not line: + continue + try: + record = json.loads(line) + except json.JSONDecodeError: + skipped['unparseable index line'] += 1 + continue + side = record.get('side') + if side not in wanted: + skipped[f'side {side!r} not requested'] += 1 + continue + if not record.get('has_logprobs'): + # Nothing to compare a new forward pass against, so GRPO has no + # ratio. Counted rather than dropped silently: a collection that + # produced many of these is a collection whose sampler was not + # returning logprobs, which is a wiring fault, not attrition. + skipped['trajectory has no logprobs'] += 1 + continue + arrays = np.load(os.path.join(traj_dir, record['npz'])) + ids = arrays['input_ids'].astype(np.int64) + labels = arrays['labels'].astype(np.int64) + logps = arrays['logprobs'].astype(np.float64) + n_train = int((labels != -100).sum()) + if not n_train: + skipped['no trainable tokens'] += 1 + continue + if logps.size != n_train: + # Off by anything here pairs each logprob with the wrong token and + # the loss still comes out a number, so it is a hard stop rather + # than something to trim to the shorter of the two. + skipped[f'logps {logps.size} != trainable {n_train}'] += 1 + continue + if ids.size > MAX_MODEL_LEN: + # Dropped before the model is built, so the count is in the log + # rather than arriving as an exception in the middle of a step. + # Reachable in normal operation: challenge.py samples at + # max_model_len 40960, which is above this. + skipped[f'longer than MAX_MODEL_LEN={MAX_MODEL_LEN}'] += 1 + continue + key = ((side, record.get('group_id')) if side == 'propose' else + (side, record.get('group_id'), record.get('proposal_idx'))) + by_key.setdefault(key, []).append({ + 'side': side, + 'input_ids': ids.tolist(), + # Labels are stored already shifted by one -- labels[i] is the + # token at input_ids[i+1] -- which is how the sampler wrote them. + # Passed through untouched; re-deriving them here would be guessing + # at an alignment that is already correct on disk. + 'labels': labels.tolist(), + 'attention_mask': [1] * len(ids), + 'position_ids': list(range(len(ids))), + 'logps': logps.tolist(), + 'reward': float(record.get('reward') or 0.0), + }) + groups = [] + for key, members in by_key.items(): + if len(members) < 2: + # One member means the advantage is the reward minus itself. + skipped[f'group of {len(members)} (no gradient)'] += 1 + continue + groups.append({'key': key, 'side': members[0]['side'], 'members': members}) + return groups, skipped + + +def score(groups: List[Dict[str, Any]]) -> collections.Counter: + """Advantage per member, in place. Groups may differ in size.""" + advantage_fn = GRPOAdvantage() + notes: collections.Counter = collections.Counter() + for group in groups: + rewards = [m['reward'] for m in group['members']] + adv = advantage_fn(rewards, num_generations=len(rewards), scale='group').tolist() + if all(abs(a) < 1e-9 for a in adv): + # Every member scored the same, so the group cancels out. Reported + # because it is the one failure that costs a full collection and looks + # like a successful run: the step happens and moves nothing. + notes[f'{group["side"]}: group with no gradient after scoring'] += 1 + for member, a in zip(group['members'], adv): + member['advantage'] = a + return notes + + +def interleave(groups: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Spread each side's groups evenly through the order. + + At one optimizer step this changes nothing about the update -- every group is + in the one step either way -- but the order is what the per-step log reports, + and a mixed order makes the composition line readable. + + Each side's share of the update is by trajectory count, not by token count. + GRPOLoss averages per sequence and then across the batch and reports + num_tokens=0, which puts the model on the path that weights every micro group + equally rather than dividing by a global token sum. Measured on run_clean9's + lengths the two readings differ by about 4 points; they are not the same thing. + """ + by_side: Dict[str, List[Dict[str, Any]]] = collections.OrderedDict() + for group in groups: + by_side.setdefault(group['side'], []).append(group) + if len(by_side) < 2: + return list(groups) + marked = [((i + 0.5) / len(gs), g) + for gs in by_side.values() for i, g in enumerate(gs)] + marked.sort(key=lambda t: t[0]) + return [g for _, g in marked] + + +def main(): + groups, skipped = load(RUN_DIR) + if not groups: + raise SystemExit(f'[train] nothing trainable in {RUN_DIR}: {dict(skipped)}') + skipped.update(score(groups)) + batch = [m for g in interleave(groups) for m in g['members']] + mix = collections.Counter(m['side'] for m in batch) + sizes = collections.Counter((g['side'], len(g['members'])) for g in groups) + logger.info(f'[train] {len(groups)} groups, {len(batch)} trajectories {dict(mix)}; ' + f'group sizes {dict(sizes)}') + if skipped: + for note, n in sorted(skipped.items()): + logger.warning(f'[train] skipped: {note} x{n}') + + twinkle.initialize(mode='ray', nproc_per_node=MODEL_GPUS, lazy_collect=False, + groups=[DeviceGroup(name='model', ranks=list(range(MODEL_GPUS)), + device_type='GPU')]) + # Full-parameter: no adapter, so every weight is trained and the checkpoint is + # a whole model rather than something to merge before the next iteration. + from twinkle.model.megatron import MegatronModel + # variable_seq_lengths stays off with padding_free: both switches send + # collate_fn down the packed path, and Megatron's TE extension then reads + # PackedSeqParams.pad_between_seqs, which this Megatron-LM checkout does not + # define. Padded batches cost throughput but keep attention on plain sequences. + model = MegatronModel(model_id=MODEL_ID, device_mesh=DeviceMesh.from_sizes( + world_size=MODEL_GPUS, dp_size=MODEL_GPUS), remote_group='model', + mixed_precision='bf16', variable_seq_lengths=False) + model.set_optimizer('default', lr=LEARNING_RATE) + # Inert at one step: the scheduler is read before lr_step advances it, so the + # update happens at max_lr and there is no second step to decay over. Wired up + # so that splitting the run into steps would give a decay across them. + model.set_lr_scheduler('default', lr_decay_steps=1, max_lr=LEARNING_RATE) + # beta=0: there is no reference model here, and the KL term needs beta>0 AND + # ref_logps, so any beta above 0 would silently do nothing. + model.set_loss('GRPOLoss', epsilon=0.2, beta=0.0) + model.set_processor(InputProcessor, padding_free=False) + model.set_template(TEMPLATE, model_id=MODEL_ID, max_length=MAX_MODEL_LEN, + enable_thinking=True) + # approx_kl on landed data is the check for whether this collection belongs to + # the weights being trained: it should start near zero, and a large value means + # the dump came from a different checkpoint. + model.add_metric('GRPOMetric', is_training=True, epsilon=0.2) + logger.info(get_device_placement()) + + inputs = [{k: m[k] for k in ('input_ids', 'labels', 'attention_mask', 'position_ids')} + for m in batch] + old_logps = [m['logps'] for m in batch] + advantages = [m['advantage'] for m in batch] + dropped = 0 + for lo in range(0, len(inputs), MINI_BATCH_SIZE): + hi = min(lo + MINI_BATCH_SIZE, len(inputs)) + # A tail shorter than a whole mini batch is dropped rather than handed + # over: dispatch 'slice_dp' splits it across all ranks, and a batch that + # cannot give every rank its own micro batch raises inside _dispatch_args + # before collate_fn ever runs. + if hi - lo < MINI_BATCH_SIZE: + dropped = hi - lo + logger.warning(f'[train] dropping the last {dropped} trajectories, under ' + f'the mini batch of {MINI_BATCH_SIZE}') + break + model.forward_backward(inputs=inputs[lo:hi], old_logps=old_logps[lo:hi], + advantages=advantages[lo:hi], + micro_batch_size=MICRO_BATCH_SIZE) + # Once, after every mini batch: forward_backward neither steps nor zeroes, so + # the mini batches above simply add their gradients together and one step + # consumes all of them. + model.clip_grad_and_step() + + log = model.calculate_metric(is_training=True) + high_kl = log.pop('_high_kl_records', None) + logger.info(f'[train] one step over {len(batch) - dropped} trajectories ' + f'adv[{min(advantages):+.3f},{max(advantages):+.3f}] {log}') + if high_kl: + logger.warning(f'[train] {len(high_kl)} sequences disagree with the sampler ' + f'logps; this collection may not be from these weights') + summary = { + 'groups': len(groups), + 'trajectories': len(batch), + 'trained': len(batch) - dropped, + 'dropped_tail': dropped, + 'sides': dict(mix), + 'group_sizes': {f'{s}:{n}': c for (s, n), c in sizes.items()}, + 'advantage_min': min(advantages), + 'advantage_max': max(advantages), + 'learning_rate': LEARNING_RATE, + 'metrics': log, + 'high_kl_records': high_kl or [], + # Named, not summed: a collection that lost half its trajectories to one + # reason and one that lost none read the same from the metrics alone. + 'skipped': dict(skipped), + } + with open(os.path.join(RUN_DIR, 'train_summary.json'), 'w', encoding='utf-8') as f: + json.dump(summary, f, indent=2, ensure_ascii=False, default=str) + model.save(SAVE_NAME, output_dir=SAVE_DIR) + logger.info(f'[train] checkpoint at {os.path.join(SAVE_DIR, SAVE_NAME)}') + + # The collection's own numbers, written by challenge.py in this same directory. + # Absent when train.py is pointed at a directory collected before this existed, + # in which case the training half still goes up alone. + challenge_path = os.path.join(RUN_DIR, 'challenge_metrics.json') + challenge: Dict[str, Any] = {} + if os.path.exists(challenge_path): + with open(challenge_path, encoding='utf-8') as f: + challenge = json.load(f).get('scalars') or {} + else: + logger.warning(f'[train] no {challenge_path}; uploading training metrics only') + training = { + 'groups': len(groups), + 'trajectories': len(batch), + 'trained': len(batch) - dropped, + 'dropped_tail': dropped, + 'propose_trajectories': mix.get('propose', 0), + 'solve_trajectories': mix.get('solve', 0), + 'advantage_min': min(advantages), + 'advantage_max': max(advantages), + 'learning_rate': LEARNING_RATE, + 'skipped_total': sum(skipped.values()), + 'high_kl_sequences': len(high_kl or []), + # Every numeric metric GRPOMetric returned: loss, clip fractions, approx_kl. + **{k: v for k, v in log.items() if isinstance(v, (int, float))}, + } + upload(challenge, training) + + +if __name__ == '__main__': + main() diff --git a/cookbook/rsi/agentic/train_offline.py b/cookbook/rsi/agentic/train_offline.py deleted file mode 100644 index 156c9fdb0..000000000 --- a/cookbook/rsi/agentic/train_offline.py +++ /dev/null @@ -1,552 +0,0 @@ -"""Offline GRPO on both sides of a challenge run's dump: full-parameter. - -``challenge.py`` already generated the problems and solved each one eight times, -and it landed the token ids, the labels and the sampler's logprobs for every one -of those trajectories. This trains on that dump directly -- nothing is generated -here and nothing is re-encoded, so the tokens trained on are byte-for-byte the -tokens that were sampled. - -Two sides come out of one run: - -* proposing -- one trajectory per proposal, grouped by ``group_id`` (the - proposals answering one identical prompt). Reward is ``challenger_reward``: - ``1 - 2|p - 1/2|`` for p the fraction of solver attempts that passed, so a - proposal is worth most when the solver got it right about half the time. -* solving -- eight trajectories per task, grouped by task. Reward is the check - script's exit code, 1.0 or 0.0. - -Both sides go into the same optimizer step and are weighted only by how many -trainable tokens they carry; no coefficient is applied to either. - -Usage, after a run of challenge.py --proposals-per-group 8: - - python cookbook/rsi/agentic/train_offline.py \\ - --run-dir output/rsi_agentic/run_clean10 \\ - --model-id ms://Qwen/Qwen3-4B --model-gpus 8 - -The saved checkpoint is HF-format weights plus tokenizer, which is what -``challenge.py --model-id`` takes, so the next round of the loop is a shell line -rather than a conversion step. -""" -import collections -import json -import os -from typing import Any, Dict, List, Optional, Tuple - -import numpy as np - -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_device_placement, get_logger -from twinkle.advantage import GRPOAdvantage -from twinkle.cli import CLI -from twinkle.processor import InputProcessor - -logger = get_logger() -args = CLI.from_args() - -# ========== Configuration ========== -MODEL_ID = args.model.model_id or 'ms://Qwen/Qwen3-4B' -MODEL_GPUS = args.infra.model_gpus or 8 -# Base text template, as in cookbook/rsi/rl.py:102. Qwen3-4B is text-only, and -# the multimodal subclass crashes on encode for it. -TEMPLATE = os.environ.get('RSI_TEMPLATE', 'Template') - -# Whatever --lr says, or the CLI's own 1e-5. Not written as ``or <number>``: the -# CLI default is never zero, so a fallback here would be dead code that reads -# like the default. -LEARNING_RATE = args.optimizer.learning_rate -# One trajectory per micro batch, because padding_free is off: a micro batch is -# padded to its longest member, so pairing a short solver trajectory with a long -# proposer episode pays for the long one twice. Proposer episodes are whole -# agentic rollouts -- median 7.5k tokens, up to 15.7k -- against much shorter -# solver attempts, and at 2 per micro batch step 7 hit a 31k-token padded batch -# and died of CUDA OOM with activation recompute already at its most aggressive -# setting. At 1 the peak is the longest single trajectory and no token is padding. -# -# Read from the environment rather than args.training: TrainingArgs defaults -# micro_batch_size to 2, not to None, so `args.training.micro_batch_size or 1` -# would be dead code that silently kept 2 -- which is exactly how the OOM -# survived a first attempt at this fix. -MICRO_BATCH_SIZE = int(os.environ.get('RSI_MICRO_BATCH_SIZE', 1)) -# forward_backward is declared dispatch='slice_dp', so a mini-batch is sliced -# across the model's data-parallel ranks and each rank collates only its share. -# That share has to hold at least one micro batch, so the floor is -# MODEL_GPUS * MICRO_BATCH_SIZE: at 8 GPUs the old default of 8 left every rank -# with a single trajectory and step 1 died in collate_fn, before any optimizer -# step, on a full 384-trajectory dump. rl.py keeps the same floor by skipping -# batches smaller than MODEL_GPUS. -MINI_BATCH_SIZE = args.training.mini_batch_size or MODEL_GPUS * MICRO_BATCH_SIZE -SAVE_STEPS = args.training.save_steps or 0 - -RUN_DIR = os.environ.get('RSI_RUN_DIR', '') -SAVE_DIR = os.environ.get('RSI_SAVE_DIR', 'output/rsi_agentic/ckpt') -SAVE_NAME = os.environ.get('RSI_SAVE_NAME', 'agentic-offline') - -# Trajectories per optimizer step. 32 is BATCH_SIZE 4 x NUM_GENERATIONS 8, the -# same as the online rl.py loop, so a step here moves the weights by as much as a -# step there. -STEP_SIZE = int(os.environ.get('RSI_STEP_SIZE', 32)) -# Longest trajectory fed to the model. Above this the model refuses it mid-step. -MAX_MODEL_LEN = int(os.environ.get('RSI_MAX_MODEL_LEN', 32768)) - -# Which side(s) to train. 'both', 'solver', 'proposer'. -SIDES = os.environ.get('RSI_SIDES', 'both') - -# Cap on proposing-side groups per run, 0 for no cap. The solving side is held -# constant by challenge.py's --keep-target, but the proposing side is however -# many proposals it took to reach that target, which grows as the model improves -# and fewer of its proposals land in the band. Capping keeps the two sides' share -# of each step the same from run to run; the proposals above the cap still did -# their job of measuring difficulty, they just do not also become training data. -MAX_PROPOSER_GROUPS = int(os.environ.get('RSI_MAX_PROPOSER_GROUPS', 0)) - -# Same for the solving side, 0 for no cap. ``--keep-target`` stops challenge.py -# once it has that many tasks in the band, but the internal round that crosses -# the target finishes measuring everything it started -- run_clean9's last round -# added 14 tasks to reach 53 from 39 -- so a run overshoots by however much that -# round produced. Capping makes every run's solving side exactly the same size. -MAX_SOLVER_GROUPS = int(os.environ.get('RSI_MAX_SOLVER_GROUPS', 0)) - -# Where the numbers go. Two files under the run directory: -# train_summary.json one object: the settings this run used, what the -# challenger collected, and what got trained on -# train_steps.jsonl one line per optimizer step, every metric the model -# reported plus the batch's own composition -# Written rather than uploaded, and written as they happen rather than at the end, -# so a run that dies partway still leaves the steps it did finish. -SUMMARY_NAME = os.environ.get('RSI_SUMMARY_NAME', 'train_summary.json') -STEPS_NAME = os.environ.get('RSI_STEPS_NAME', 'train_steps.jsonl') - -# Solver groups outside this pass-count range carry one reward for all eight -# members, so their advantages are all zero and a forward+backward over them adds -# nothing. The bounds exclude only 0 and n_rollouts, which is what makes this -# 'the groups that have a gradient' rather than a difficulty judgement. -KEEP_MIN_PASS = int(os.environ.get('RSI_KEEP_MIN_PASS', 1)) -KEEP_MAX_PASS_MARGIN = int(os.environ.get('RSI_KEEP_MAX_MARGIN', 1)) - - -def _cap(groups: List[Dict[str, Any]], limit: int, side: str, - notes: collections.Counter) -> List[Dict[str, Any]]: - """Hold a side's group count to ``limit`` so every run trains on the same amount. - - Kept in file order, which is the order they came out of the challenger. The - weights do not change inside one challenge.py run, so the groups dropped are - not different in kind from the ones kept -- they are just later. - """ - if not limit: - return groups - if len(groups) > limit: - notes[f'{side} groups over the cap of {limit}'] += len(groups) - limit - return groups[:limit] - if len(groups) < limit: - # Worth saying out loud: this run trains on less than the cap promises, so - # its steps are not mixed the same way as a run that reached it. - logger.warning(f'[offline] only {len(groups)} {side} groups, below the cap ' - f'of {limit}; this run is not comparable to one that ' - f'reached the cap') - return groups - - -def _collection_metrics(run_dir: str) -> Dict[str, float]: - """Summarise what the challenger produced, before any of the training filters. - - Read off the two dumps rather than from anything challenge.py logs, so these - numbers describe the whole collection and not the subset that survived the - pass band and the caps. - - Two pass rates are reported and they mean different things. ``acc/all`` is - over every attempt on every task whose difficulty got measured, including the - tasks where all eight attempts agreed. ``acc/trained`` is over the attempts on - the tasks that go into training, and is bounded to [1/8, 7/8] by that band -- - it cannot report an all-correct or all-wrong task even if the model produces - one. Neither is a capability measurement on its own: the tasks change every - iteration and the challenger is being trained to push them toward half - passing, so a flat curve is what success looks like for both. - """ - m: Dict[str, float] = {} - idx = os.path.join(run_dir, 'propose_traj', 'index.jsonl') - if os.path.isfile(idx): - rows = [json.loads(line) for line in open(idx) if line.strip()] - outcomes = collections.Counter(r.get('outcome') for r in rows) - m['collect/proposals'] = len(rows) - for name, n in outcomes.items(): - m[f'collect/outcome_{name}'] = n - rewards = [r['challenger_reward'] for r in rows if r.get('challenger_reward') is not None] - if rewards: - m['collect/proposer_reward_mean'] = sum(rewards) / len(rewards) - - measured = [r['n_pass'] for r in rows - if r.get('n_pass') is not None and r.get('n_rollouts')] - rollouts = [r['n_rollouts'] for r in rows - if r.get('n_pass') is not None and r.get('n_rollouts')] - if measured: - n_att = sum(rollouts) - m['collect/measured_tasks'] = len(measured) - m['collect/never_measured'] = len(rows) - len(measured) - m['acc/all'] = sum(measured) / n_att - dist = collections.Counter(measured) - for k in range(max(rollouts) + 1): - m[f'collect/n_pass_{k}'] = dist.get(k, 0) - # All eight attempts agreeing means one reward for the whole group, so - # the group mean equals it and every advantage is zero. Tracking the - # share of tasks like that is tracking how much of the collection did - # no work. - flat = sum(1 for p, r in zip(measured, rollouts) if p == 0 or p == r) - m['collect/zero_gradient_frac'] = flat / len(measured) - band = [(p, r) for p, r in zip(measured, rollouts) - if KEEP_MIN_PASS <= p <= r - KEEP_MAX_PASS_MARGIN] - if band: - m['acc/trained'] = sum(p for p, _ in band) / sum(r for _, r in band) - m['collect/band_tasks'] = len(band) - # How many proposals it costs to land one trainable task. Expected - # to climb as the model gets better at its own proposals, which is - # what makes an iteration take longer than the one before it. - m['collect/proposals_per_band_task'] = len(rows) / len(band) - - attempts = os.path.join(run_dir, 'solver_attempts.jsonl') - if os.path.isfile(attempts): - rows = [json.loads(line) for line in open(attempts) if line.strip()] - if rows: - m['collect/solver_attempts'] = len(rows) - m['collect/solver_truncated_frac'] = \ - sum(1 for r in rows if r.get('truncated')) / len(rows) - return m - - -def _load_solver_groups(run_dir: str) -> Tuple[List[Dict[str, Any]], collections.Counter]: - """Groups of solver attempts on one task, with their exit-code rewards.""" - path = os.path.join(run_dir, 'solver_attempts.jsonl') - if not os.path.exists(path): - return [], collections.Counter({'no solver_attempts.jsonl': 1}) - by_task: Dict[str, List[Dict[str, Any]]] = collections.OrderedDict() - notes: collections.Counter = collections.Counter() - with open(path) as f: - for line in f: - line = line.strip() - if not line: - continue - try: - rec = json.loads(line) - except json.JSONDecodeError: - notes['unparseable line'] += 1 - continue - by_task.setdefault(rec['statement'], []).append(rec) - - groups: List[Dict[str, Any]] = [] - for statement, recs in by_task.items(): - rewards = [1.0 if r['check_exit'] == 0 else 0.0 for r in recs] - n_pass = int(sum(rewards)) - high = len(recs) - KEEP_MAX_PASS_MARGIN - if not (KEEP_MIN_PASS <= n_pass <= high): - notes[f'solver n_pass={n_pass} of {len(recs)} (no gradient)'] += 1 - continue - members = [] - for rec, reward in zip(recs, rewards): - att = rec.get('attempt') or {} - members.append({ - 'input_ids': att.get('input_ids') or [], - 'labels': att.get('labels') or [], - 'attention_mask': att.get('attention_mask') or [], - 'position_ids': att.get('position_ids') or [], - # The sampler's own logprobs, in the [[token, logp]] shape the - # rollout stored them in. Reusing them rather than a fresh - # forward is what keeps old_logps free of engine differences. - 'logps': [lp[0][1] for lp in (att.get('logprobs') or [])], - 'reward': reward, - }) - groups.append({'side': 'solver', 'key': statement[:60], 'members': members}) - return _cap(groups, MAX_SOLVER_GROUPS, 'solver', notes), notes - - -def _load_proposer_groups(run_dir: str) -> Tuple[List[Dict[str, Any]], collections.Counter]: - """Groups of proposals answering one prompt, with their 50%-target rewards.""" - d = os.path.join(run_dir, 'propose_traj') - index = os.path.join(d, 'index.jsonl') - notes: collections.Counter = collections.Counter() - if not os.path.exists(index): - return [], collections.Counter({'no propose_traj/index.jsonl': 1}) - - by_group: Dict[Any, List[Dict[str, Any]]] = collections.OrderedDict() - with open(index) as f: - for line in f: - line = line.strip() - if not line: - continue - try: - rec = json.loads(line) - except json.JSONDecodeError: - notes['unparseable line'] += 1 - continue - gid = rec.get('group_id') - if gid is None: - # Landed by a run from before proposals were grouped. Every such - # proposal is its own group of one, so its advantage would be - # zero; counted and skipped rather than trained on as noise. - notes['proposal has no group_id (pre-grouping run)'] += 1 - continue - by_group.setdefault(gid, []).append(rec) - - groups: List[Dict[str, Any]] = [] - for gid, recs in by_group.items(): - if len(recs) < 2: - notes[f'proposer group of {len(recs)} (no gradient)'] += 1 - continue - members = [] - for rec in recs: - npz_name = rec.get('npz') - if not npz_name: - notes['proposal has no npz (text-only rollout)'] += 1 - continue - z = np.load(os.path.join(d, npz_name)) - if 'r0_input_ids' not in z.files: - notes['npz has no r0_input_ids'] += 1 - continue - ids = z['r0_input_ids'].astype(np.int64) - # Labels are stored already shifted by one -- labels[i] is the token - # at input_ids[i+1] -- which is the convention the sampler wrote them - # in. Passed through untouched; re-deriving them here would be - # guessing at an alignment that is already correct on disk. - labels = z['r0_labels'].astype(np.int64) - logps = z['r0_logprobs'].astype(np.float64) if 'r0_logprobs' in z.files else None - if logps is None: - notes['npz has no r0_logprobs'] += 1 - continue - members.append({ - 'input_ids': ids.tolist(), - 'labels': labels.tolist(), - 'attention_mask': [1] * len(ids), - 'position_ids': list(range(len(ids))), - 'logps': logps.tolist(), - 'reward': float(rec.get('challenger_reward') or 0.0), - }) - if len(members) < 2: - notes['proposer group lost members to missing arrays'] += 1 - continue - groups.append({'side': 'proposer', 'key': f'group{gid}', 'members': members}) - return _cap(groups, MAX_PROPOSER_GROUPS, 'proposer', notes), notes - - -def _score(groups: List[Dict[str, Any]]) -> collections.Counter: - """Advantage per member, in place. Groups may differ in size.""" - advantage_fn = GRPOAdvantage() - notes: collections.Counter = collections.Counter() - for g in groups: - rewards = [m['reward'] for m in g['members']] - adv = advantage_fn(rewards, num_generations=len(rewards), scale='group').tolist() - if all(abs(a) < 1e-9 for a in adv): - notes[f'{g["side"]}: group with no gradient after scoring'] += 1 - for m, a in zip(g['members'], adv): - m['advantage'] = a - return notes - - -def _interleave(by_side: Dict[str, List[Dict[str, Any]]]) -> List[Dict[str, Any]]: - """Spread each side's groups evenly through the order. - - Both sides have to appear in every step, or a step's update comes from one - side only and 'weighted by token count' stops describing anything. The ratio - is not a free choice: it falls out of using all of both sides' groups over - the same number of steps. - """ - sides = [s for s in ('solver', 'proposer') if by_side.get(s)] - if len(sides) < 2: - return list(by_side.get(sides[0], [])) if sides else [] - # Place each group at its fractional position within its own side, then sort - # by that position: a side with three times the groups contributes three for - # every one of the other's, spread out rather than in a block. - marked: List[Tuple[float, Dict[str, Any]]] = [] - for s in sides: - gs = by_side[s] - for i, g in enumerate(gs): - marked.append(((i + 0.5) / len(gs), g)) - marked.sort(key=lambda t: t[0]) - return [g for _pos, g in marked] - - -def main(): - if not RUN_DIR: - raise SystemExit('set RSI_RUN_DIR to a challenge.py output directory') - - collect = _collection_metrics(RUN_DIR) - - by_side: Dict[str, List[Dict[str, Any]]] = {} - all_notes: collections.Counter = collections.Counter() - if SIDES in ('both', 'solver'): - gs, notes = _load_solver_groups(RUN_DIR) - by_side['solver'] = gs - all_notes.update(notes) - if SIDES in ('both', 'proposer'): - gs, notes = _load_proposer_groups(RUN_DIR) - by_side['proposer'] = gs - all_notes.update(notes) - - all_notes.update(_score([g for gs in by_side.values() for g in gs])) - - for side, gs in by_side.items(): - n_traj = sum(len(g['members']) for g in gs) - sizes = collections.Counter(len(g['members']) for g in gs) - logger.info(f'[offline] {side}: {len(gs)} groups, {n_traj} trajectories, ' - f'group sizes {dict(sorted(sizes.items()))}') - for note, n in all_notes.most_common(): - logger.info(f'[offline] skipped {n}: {note}') - if not any(by_side.values()): - raise SystemExit(f'[offline] nothing trainable in {RUN_DIR}') - - order = _interleave(by_side) - flat: List[Dict[str, Any]] = [] - for g in order: - for m in g['members']: - m['side'] = g['side'] - flat.append(m) - - # Trajectories the model would refuse, dropped before anything is loaded so - # the count is in the log rather than showing up as a mid-step exception. - kept: List[Dict[str, Any]] = [] - for m in flat: - n_train = sum(1 for label in m['labels'] if label != -100) - if not n_train: - all_notes['no trainable tokens'] += 1 - continue - if len(m['logps']) != n_train: - # Off-by-anything here pairs each logprob with the wrong token and - # the loss still comes out a number, so it has to be a hard stop. - all_notes[f'logps {len(m["logps"])} != trainable {n_train}'] += 1 - continue - if len(m['input_ids']) > MAX_MODEL_LEN: - all_notes[f'longer than MAX_MODEL_LEN={MAX_MODEL_LEN}'] += 1 - continue - kept.append(m) - - n_steps = (len(kept) + STEP_SIZE - 1) // STEP_SIZE - mix = collections.Counter(m['side'] for m in kept) - logger.info(f'[offline] {len(kept)} trainable trajectories ' - f'({dict(mix)}), {n_steps} steps of {STEP_SIZE}') - for note, n in all_notes.most_common(): - logger.info(f'[offline] skipped {n}: {note}') - - summary = { - 'run_dir': RUN_DIR, - 'config': {'model_id': MODEL_ID, 'lr': LEARNING_RATE, 'step_size': STEP_SIZE, - 'sides': SIDES, 'keep_min_pass': KEEP_MIN_PASS, - 'keep_max_pass_margin': KEEP_MAX_PASS_MARGIN, - 'max_solver_groups': MAX_SOLVER_GROUPS, - 'max_proposer_groups': MAX_PROPOSER_GROUPS, - 'mini_batch_size': MINI_BATCH_SIZE, - 'micro_batch_size': MICRO_BATCH_SIZE, - 'max_model_len': MAX_MODEL_LEN, 'template': TEMPLATE, - 'model_gpus': MODEL_GPUS}, - 'collect': collect, - 'train': {'trainable_trajectories': len(kept), 'steps': n_steps, - **{f'{side}_trajectories': n for side, n in mix.items()}, - 'groups_per_side': {side: len(gs) for side, gs in by_side.items()}}, - # Every reason anything was left out, with its count. Kept in the summary - # rather than only in the log so a later comparison between iterations can - # tell a change in the model from a change in how much survived the load. - 'skipped': dict(all_notes), - } - summary_path = os.path.join(RUN_DIR, SUMMARY_NAME) - with open(summary_path, 'w') as f: - json.dump(summary, f, indent=2, ensure_ascii=False) - steps_path = os.path.join(RUN_DIR, STEPS_NAME) - steps_file = open(steps_path, 'w') - logger.info(f'[offline] numbers going to {summary_path} and {steps_path}') - for k in sorted(collect): - logger.info(f'[offline] {k} = {collect[k]}') - - device_groups = [DeviceGroup(name='model', ranks=list(range(MODEL_GPUS)), device_type='GPU')] - model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, dp_size=MODEL_GPUS) - twinkle.initialize(mode='ray', nproc_per_node=MODEL_GPUS, groups=device_groups, - lazy_collect=False) - - # Full-parameter: no adapter is added, so every weight is trained and the - # checkpoint is a whole model rather than something that needs merging before - # the next challenger round can load it. - from twinkle.model.megatron import MegatronModel - # padding_free is off, so this stays off with it: both switches send collate_fn - # down the per-micro-batch packed path, and Megatron's TE extension then reads - # PackedSeqParams.pad_between_seqs, which the Megatron-LM checkout on this box - # does not define. Padded batches cost throughput but keep the attention path - # on plain padded sequences. - model = MegatronModel(model_id=MODEL_ID, device_mesh=model_mesh, remote_group='model', - mixed_precision='bf16', variable_seq_lengths=False) - model.set_optimizer('default', lr=LEARNING_RATE) - model.set_lr_scheduler('default', lr_decay_steps=max(1, n_steps), max_lr=LEARNING_RATE) - # beta=0: no reference model here, and grpo.py:315 needs beta>0 AND ref_logps - # for the KL term, so any beta above 0 would silently do nothing. - model.set_loss('GRPOLoss', epsilon=0.2, beta=0.0) - model.set_processor(InputProcessor, padding_free=False) - model.set_template(TEMPLATE, model_id=MODEL_ID, max_length=MAX_MODEL_LEN, - enable_thinking=True) - # approx_kl at the first inner step compares the sampler's logps against the - # trainer's on the same tokens. On landed data that is the check for whether - # this dump belongs to the weights being trained: it should start near zero, - # and a large value means the dump came from a different checkpoint. - model.add_metric('GRPOMetric', is_training=True, epsilon=0.2) - logger.info(get_device_placement()) - - for step in range(n_steps): - lo, hi = step * STEP_SIZE, min((step + 1) * STEP_SIZE, len(kept)) - batch = kept[lo:hi] - inputs = [{'input_ids': m['input_ids'], 'labels': m['labels'], - 'attention_mask': m['attention_mask'], - 'position_ids': m['position_ids']} for m in batch] - old_logps = [m['logps'] for m in batch] - advantages = [m['advantage'] for m in batch] - - for mb in range(0, len(inputs), MINI_BATCH_SIZE): - end = min(mb + MINI_BATCH_SIZE, len(inputs)) - # Drop a tail shorter than a whole mini batch instead of handing it - # over. The floor is MINI_BATCH_SIZE, not MICRO_BATCH_SIZE: dispatch - # 'slice_dp' splits the batch across all MODEL_GPUS ranks, so a batch - # that cannot give every rank its own micro batch fails before - # collate_fn -- _dispatch_args raises 'Batch too small for N workers, - # some ranks have no data'. A 6-trajectory tail against 8 ranks did - # exactly that. The dropped trajectories are the tail of the last step. - if end - mb < MINI_BATCH_SIZE: - logger.info(f'[offline] step {step + 1}: dropping last ' - f'{end - mb} traj, under mini batch {MINI_BATCH_SIZE}') - break - model.forward_backward( - inputs=inputs[mb:end], - old_logps=old_logps[mb:end], - advantages=advantages[mb:end], - micro_batch_size=MICRO_BATCH_SIZE, - ) - # Once per step, not per mini-batch: forward_backward neither steps nor - # zeroes, and clip_grad_norm divides by the tokens accumulated across all - # of them, so every trajectory in the step carries the same weight no - # matter how the mini-batches split -- which is also how the two sides end - # up weighted by their token counts and nothing else. - model.clip_grad_and_step() - - side_mix = collections.Counter(m['side'] for m in batch) - log = model.calculate_metric(is_training=True) - # A list of records rather than a number, marked with a leading underscore - # at grpo.py:451 for that reason. Going to a file, so the records go in - # whole instead of being reduced to a count. - high_kl = log.pop('_high_kl_records', None) - logger.info(f'[offline] step {step + 1}/{n_steps} ' - f'{len(batch)} traj {dict(side_mix)} ' - f'adv[{min(advantages):+.3f},{max(advantages):+.3f}] {log}') - if high_kl: - logger.warning(f'[offline] step {step + 1}: {len(high_kl)} sequences ' - f'with high kl against the sampler logps') - row = {'step': step + 1, **log, - 'trajectories': len(batch), - 'solver': side_mix.get('solver', 0), - 'proposer': side_mix.get('proposer', 0), - 'adv_min': min(advantages), 'adv_max': max(advantages), - 'high_kl_records': high_kl or []} - steps_file.write(json.dumps(row, ensure_ascii=False, default=str) + '\n') - steps_file.flush() - if SAVE_STEPS and (step + 1) % SAVE_STEPS == 0: - model.save(f'{SAVE_NAME}-step{step + 1}', output_dir=SAVE_DIR) - - steps_file.close() - model.save(SAVE_NAME, output_dir=SAVE_DIR) - logger.info(f'[offline] done, {n_steps} steps; checkpoint at ' - f'{os.path.join(SAVE_DIR, SAVE_NAME)}') - - -if __name__ == '__main__': - main() diff --git a/src/twinkle_agentic/challenger/agentic.py b/src/twinkle_agentic/challenger/agentic.py index 5c9105f9c..1884fdf5a 100644 --- a/src/twinkle_agentic/challenger/agentic.py +++ b/src/twinkle_agentic/challenger/agentic.py @@ -38,6 +38,7 @@ """ import ast import json +import math import re import threading from concurrent.futures import ThreadPoolExecutor, as_completed @@ -241,6 +242,77 @@ def brittle_check_reason(script: str) -> Optional[str]: return None +# Literals shorter than this match by accident: a statement contains "3" or "id" +# for its own reasons. Measured on 188 tasks from run_clean9, one and two digit +# integers appeared in both the check and the statement 91% of the time, which is +# what an unattributable coincidence rate looks like. +_MIN_DERIVED_LEN = 3 + + +def derived_check_literals(script: str) -> List[str]: + """The values a check compares against that the solver is meant to work out. + + A measurement tool, not part of the proposing path. It exists because "does the + statement give away the answer" cannot be asked without first separating the + three kinds of thing a check's literals are, and only one of them is a leak: + + an identifier, or a name with a file extension + The statement MUST carry these. It is naming the file to create and the + fields to put in it; a statement that withheld them would describe no + particular output at all. Present in 84-93% of run_clean9's statements, + which is the correct rate. + text that appears in the workspace + Input data, which the statement is meant to quote verbatim so the solver + can write the same bytes. Not separated here -- the caller filters on + the snapshot if it wants to. + a long number, a float, or a string that is none of the above + Only exists once the work has been done. This is the group returned. + + Feeding the result back to the statement stage as a forbidden list was tried and + did not reduce the leak: see the note above PROBLEM_FOLLOWUP_RULES_ONLY in + cookbook/rsi/agentic/prompts.py for the two forms measured and their p-values. + + Wrong at the edges by construction: a column named ``total_2024`` reads as an + identifier and is not returned, and a computed value that lands on two digits is + below the length floor. + """ + try: + tree = ast.parse(script) + except SyntaxError: + return [] + out = [] + for node in ast.walk(tree): + if not isinstance(node, ast.Compare): + continue + for side in [node.left] + list(node.comparators): + if not isinstance(side, ast.Constant): + continue + v = side.value + if isinstance(v, bool) or v is None: + continue + if isinstance(v, (int, float)): + text = repr(v) + if len(text.lstrip('-').replace('.', '')) < _MIN_DERIVED_LEN: + continue + out.append(text) + elif isinstance(v, str): + if len(v) < _MIN_DERIVED_LEN: + continue + if re.fullmatch(r'[A-Za-z_][A-Za-z0-9_]*', v): + continue + # A trailing extension means a filename -- but only when the part + # after the dot is not itself digits. '0.001' read out of a CSV is a + # string here, and skipping it as a filename would let exactly the + # kind of value this function exists to catch through. + if re.search(r'\.[A-Za-z]\w{0,4}$', v) and ' ' not in v: + continue + out.append(v) + # Longest first: a short literal is often a substring of a longer one, and + # naming the long one first makes the list read as distinct values rather than + # as prefixes of each other. + return sorted(set(out), key=len, reverse=True) + + def parse_problem_statement(text: str) -> Optional[str]: """Extract a problem statement from the model's reply. @@ -406,7 +478,7 @@ class AgenticChallenger(Challenger): episode plus the end state and reasons at length before answering, and one that runs out of budget mid-thought never emits its code block and is thrown away as unparseable. - followup_api: optional OpenAI-compatible API client (e.g. qwen3-max). When + followup_api: optional OpenAI-compatible API client (e.g. qwen3.8-max). When given, exploration still runs on the local explorer -- so its turns keep their ``labels`` and ``logprobs`` and remain trainable -- but the check-script (success judgement) and problem-statement stages are @@ -416,7 +488,7 @@ class AgenticChallenger(Challenger): exploration" split. ``None`` keeps the single-model behaviour where the local model writes those two stages in the same conversation. followup_extra_body: extra request body forwarded on every ``followup_api`` - call (e.g. ``{'thinking_budget': N}`` to cap qwen3-max reasoning). + call (e.g. ``{'thinking_budget': N}`` to cap qwen3.8-max reasoning). ``None`` sends the request unmodified. Ignored when ``followup_api`` is ``None``. keyword_explorer: explorer used to brainstorm keywords. Should have no @@ -495,6 +567,10 @@ def __init__( setup_script_fn: Optional[Callable[..., str]] = None, solver_prompt_fn: Optional[Callable[[str], Trajectory]] = None, check_retries: int = 1, + task_bank: Optional[Any] = None, + novelty_fn: Optional[Callable[[List[Dict[str, Any]]], List[Optional[float]]]] = None, + novelty_floor: float = 0.5, + keep_per_group: int = 0, reject_sink: Optional[Callable[[Dict[str, Any]], None]] = None, propose_sink: Optional[Callable[[Dict[str, Any]], None]] = None, solver_sink: Optional[Callable[[Dict[str, Any]], None]] = None, @@ -582,7 +658,7 @@ def __init__( self.problem_params = problem_params # When set, exploration runs on the (trainable) local explorer as before, # but the check-script and problem-statement stages are generated by this - # OpenAI-compatible API (e.g. qwen3-max) instead of the local model. The + # OpenAI-compatible API (e.g. qwen3.8-max) instead of the local model. The # two stages then contribute nothing to the trainable trajectory: the API # returns text only, so the episode's ``input_ids`` / ``labels`` / # ``logprobs`` stay exactly the exploration turns the local sampler @@ -590,7 +666,7 @@ def __init__( # generated check script and statement are used solely to build the task. self.followup_api = followup_api # extra_body sent on every followup API call (e.g. {'thinking_budget': N} - # to cap qwen3-max reasoning). None sends the request unmodified. + # to cap qwen3.8-max reasoning). None sends the request unmodified. self.followup_extra_body = dict(followup_extra_body) if followup_extra_body else None self.keyword_explorer = keyword_explorer self.min_batch = max(1, min_batch) @@ -619,6 +695,41 @@ def __init__( self.check_retries = check_retries if check_retries: prompts.require('check_retry_followup') + # Novelty, off unless both of these are given. ``task_bank`` supplies the + # tasks earlier iterations produced (see :mod:`.task_bank`); ``novelty_fn`` + # takes a list of ``{statement, check, references}`` and returns one score in + # [0, 1] per entry, or None where it could not judge. Kept as injected + # callables for the same reason the sandbox ones are: this class then holds + # no opinion about which judge, model or API produces the number, and a test + # can hand it a fixed one. + self.task_bank = task_bank + self.novelty_fn = novelty_fn + if not 0.0 <= novelty_floor <= 1.0: + raise ValueError(f'novelty_floor must be in [0, 1], got {novelty_floor}') + # How much of the reward a proposal keeps when it is judged fully redundant. + # 0.5 halves it; 0.0 would be Ornith-1.5's plain ``V x D x N``, which zeroes + # it. The floor exists because our N is coarse where theirs is continuous: + # scored over run_clean9's 188 tasks, 44% came out at exactly 0.0, and eight + # proposals sharing one keyword draw can all land there -- at floor 0 that + # group's rewards are all zero, its advantages are all zero after GRPO + # subtracts the mean, and eight sandbox rollouts bought nothing. Ornith's own + # text says novelty 'should remain secondary to validity and difficulty'. + self.novelty_floor = float(novelty_floor) + # At most this many of a keyword group's in-band proposals become tasks the + # solver side trains on. 0 keeps every in-band proposal, which is what this + # did before. At 1 the two sides come out the same size -- eight groups of + # eight proposals give 64 proposing trajectories and 8 tasks x 8 attempts = + # 64 solving ones -- and the tasks are one per keyword direction instead of + # three from the same one. + # + # The proposals not selected are NOT wasted from the proposing side: each one + # still earns its own reward from its own n_pass, so the whole group still + # trains. What is dropped is their solver attempts, which were already run to + # measure difficulty: at keep_per_group=1 that is 56 proposals x 8 attempts + # per round measured and then not trained on. + if keep_per_group < 0: + raise ValueError(f'keep_per_group must be >= 0, got {keep_per_group}') + self.keep_per_group = keep_per_group self.reject_sink = reject_sink self.propose_sink = propose_sink self.solver_sink = solver_sink @@ -653,6 +764,25 @@ def __init__( # followup_api mode only: a check or statement API call failed. The # conversation is then unusable and the proposal is rejected. 'followup_api_error': 0, + # Novelty judging. ``novelty_error``: the batch's call raised, so every + # proposal in it scored None and none lost reward for it. + # ``novelty_length_mismatch``: the judge returned a different number of + # scores than proposals sent, which would pair scores with the wrong + # tasks, so all are dropped. ``novelty_unjudged``: proposals the judge + # left without a verdict. ``in_band_not_selected``: proposals inside the + # difficulty band whose group already contributed its keep_per_group + # task. + # + # These four have to be listed here: _bump does self.stats[key] += n on + # a fixed dict, so an unregistered key raises KeyError and takes the + # whole collection down. That is what killed loop2/iter1 -- seven + # proposals were measured and scored, then the counter line at the end + # of _score_novelty crashed and all seven trajectories were lost. + 'novelty_error': 0, 'novelty_length_mismatch': 0, + 'novelty_unjudged': 0, 'in_band_not_selected': 0, + # Keyword groups given up on because the judge never scored one of + # their proposals. Nothing from them is used. + 'novelty_group_dropped': 0, } self._hard: List[Tuple[str, str]] = [] @@ -991,7 +1121,8 @@ def _run_followup_api(self, state: Dict[str, Any], explored: Trajectory) -> None return # Problem-statement stage: one API reply, kept as the task's statement. - reply = self._api_reply(messages, self.prompts.problem_followup, self.problem_params) + reply = self._api_reply(messages, self.prompts.problem_followup, + self.problem_params) if reply is None: self._bump('followup_api_error') state['reject'] = ('followup_api_error', 'problem-statement API call failed') @@ -1109,7 +1240,10 @@ def _reject_record(self, traj: Trajectory, reason: str, detail: str = '') -> Non def _emit_propose(self, rounds: Optional[List[Dict[str, Any]]], outcome: str, *, keywords: Any = (), seeded: bool = False, n_pass: Optional[int] = None, - group_id: Optional[int] = None) -> None: + group_id: Optional[int] = None, + novelty: Optional[float] = None, + selected: Optional[bool] = None, + novelty_dropped: bool = False) -> None: """Hand one proposal attempt's rounds to ``propose_sink``. ``pass_rate`` is the raw fraction of solver attempts that succeeded. @@ -1118,6 +1252,10 @@ def _emit_propose(self, rounds: Optional[List[Dict[str, Any]]], outcome: str, *, on; both are written so a run can be re-scored under a different target without re-solving anything. + ``novelty`` is written next to it for the same reason: the reward already + has it multiplied in, and a run cannot be re-scored at a different floor -- + or with novelty taken back out -- from the product alone. + A proposal with no ``n_pass`` never reached difficulty measurement -- it was rejected before that -- and scores 0, the same as one nobody or everybody solved. @@ -1131,7 +1269,19 @@ def _emit_propose(self, rounds: Optional[List[Dict[str, Any]]], outcome: str, *, 'n_pass': n_pass, 'n_rollouts': rollouts, 'pass_rate': (n_pass / rollouts) if (n_pass is not None and rollouts) else None, - 'challenger_reward': self.challenger_reward(n_pass), + 'novelty': novelty, + 'novelty_factor': self.novelty_factor(novelty), + 'challenger_reward': self.challenger_reward(n_pass, novelty=novelty), + # Whether this proposal's task went on to the solver side. Not the same as + # ``outcome``: with keep_per_group set, a proposal can be in the difficulty + # band and still not be the one its group contributed. Its own reward is + # unaffected either way. + 'selected': selected, + # True when the novelty judge never returned a score for at least one + # proposal in this group, after NOVELTY_TRIES attempts. The record is + # written either way -- the episode really happened and the file is the + # audit trail -- but training skips every group carrying this flag. + 'novelty_dropped': bool(novelty_dropped), 'keywords': list(keywords or ()), 'seeded': bool(seeded), 'rounds': rounds, @@ -1139,24 +1289,80 @@ def _emit_propose(self, rounds: Optional[List[Dict[str, Any]]], outcome: str, *, with self._sink_lock: self.propose_sink(payload) - def challenger_reward(self, n_pass: Optional[int]) -> float: - """Score a proposal by how close the solver came to a 50% pass rate. - - ``1 - 2 * |p - 1/2|`` for ``p = n_pass / solver_rollouts``: 1.0 at half - the attempts passing, 0 at none or all of them. The shape is the one - R-Zero (arXiv 2508.05004) trains its challenger on, and the reason it is - peaked at a half rather than at 'hard' is that a GRPO update's size goes - with the reward variance within a group, which for a pass/fail solver is - ``p(1-p)`` -- largest exactly there. - - ``None`` means the proposal never got as far as being solved, and scores - 0. That is the floor, not a penalty: nothing here can go below 0, so a - failed proposal and an unsolvable one are worth the same. + # Where the pass-rate reward peaks, and how wide the peak is. 0.2 is Ornith-1.5's + # target (ornith.ai/ornith_1_5.html), which trains its proposer on + # ``exp(-(p-p*)^2 / 2s^2)`` rather than on a peak at one half. + PASS_RATE_TARGET = 0.2 + PASS_RATE_WIDTH = 0.3 + + # How many times the novelty judge is asked before a proposal is given up on. + # Only the proposals still missing a score are re-sent. Measured need for this: + # loop3/iter1 had 1 of 61 measured proposals come back without a verdict, in 1 + # of its 10 keyword groups, and giving up on that group costs the 56 sandbox + # attempts already spent on its 7 proposals. + NOVELTY_TRIES = 3 + + def novelty_factor(self, novelty: Optional[float]) -> float: + """What a proposal's difficulty score gets multiplied by for its novelty. + + ``floor + (1 - floor) * N``, so N=1 leaves the reward alone and N=0 leaves + ``novelty_floor`` of it. See ``novelty_floor`` in ``__init__`` for why there + is a floor at all. + + ``None`` returns 1.0, not the floor: it means nobody judged this proposal -- + no bank, no judge, or the judge's API failed -- and charging a proposal for a + measurement that did not happen would make the reward depend on API uptime. + """ + if novelty is None: + return 1.0 + n = min(1.0, max(0.0, float(novelty))) + return self.novelty_floor + (1.0 - self.novelty_floor) * n + + def challenger_reward(self, n_pass: Optional[int], + novelty: Optional[float] = None) -> float: + """Score a proposal by how close the solver came to a target pass rate. + + ``exp(-(p - p*)^2 / 2s^2)`` for ``p = n_pass / solver_rollouts``, peaked at + ``p* = 0.2`` with width ``s = 0.3``. + + This replaced ``1 - 2*|p - 1/2|``, which is what R-Zero (arXiv 2508.05004) + uses, for two reasons measured on run_clean9's 87 in-band proposals: + + It was not injective. With 8 rollouts the seven in-band values of ``n_pass`` + mapped onto four rewards -- 1 and 7 both scored 0.25, 2 and 6 both 0.50 -- + so a proposal one solver out of eight could do and one seven out of eight + could do were worth the same. The whole distinction between too hard and too + easy was erased. The gaussian separates all seven. + + Its signal was smaller than its noise. ``n_pass`` is a binomial draw around + the proposal's real difficulty, and propagating that draw through each shape + gives a noise SD to compare the spread of rewards against: 0.280 signal over + 0.246 noise for the old shape, against 0.347 over 0.177 here. A ratio of 1.14 + means over half of what a GRPO group ranks on is which way eight coin flips + landed. + + A peak below one half is also the more useful target. A group's update size + goes with reward variance, which for a pass/fail solver peaks at p=0.5 -- the + argument for the old shape -- but a proposal only teaches the solver + something when the solver mostly cannot do it yet. + + ``None`` means the proposal never got as far as being solved, and 0 means no + attempt passed. Both score 0, and that floor is now load-bearing rather than + incidental: the gaussian evaluated at p=0 is 0.801, higher than the 0.607 it + gives a proposal half the attempts solve. Without the gate the best thing a + proposer could do is write tasks nobody can finish. + + ``novelty`` multiplies the result through :meth:`novelty_factor`, which is + Ornith-1.5's ``R = V x D x N`` with a floor under the N. Left at ``None`` -- + which is what happens with no task bank or no judge -- the returned number is + exactly what it was before novelty existed. """ rollouts = self.solver_rollouts or 0 - if n_pass is None or not rollouts: + if n_pass is None or not rollouts or n_pass <= 0: return 0.0 - return 1.0 - 2.0 * abs(n_pass / rollouts - 0.5) + gap = n_pass / rollouts - self.PASS_RATE_TARGET + difficulty = math.exp(-(gap * gap) / (2.0 * self.PASS_RATE_WIDTH ** 2)) + return difficulty * self.novelty_factor(novelty) def _take_rounds(self, task: Trajectory) -> Optional[List[Dict[str, Any]]]: """Detach a task's proposing rounds. Popped even with no sink attached: @@ -1381,20 +1587,178 @@ def _prepare(k: int) -> bool: for i, task in enumerate(tasks) ] self.on_difficulty_measured(measured) + novelties = self._score_novelty(measured) high = self.solver_rollouts - self.keep_max_pass_margin in_band = [self.keep_min_pass <= n <= high for n in passes] + # Which of the in-band tasks the solver side actually trains on. Decided + # before emitting so each proposal's record says whether its task was taken. + selected = self._select_per_group(measured, passes, in_band, novelties) + dropped = self._unscored_group_ids(measured, novelties) + if dropped: + self._bump('novelty_group_dropped', len(dropped)) # Emit here, not in _round: this is where a proposal's verdict is # decided, and both sides of the band are worth keeping -- a task nobody # solved and one everybody solved are the two failure modes the # proposer would need to learn to avoid. - for task, n, kept_flag in zip(measured, passes, in_band): + for i, (task, n, kept_flag, nov) in enumerate(zip(measured, passes, in_band, + novelties)): + gid = user_data_get(task.get('user_data'), 'group_id', None) self._emit_propose(self._take_rounds(task), 'kept' if kept_flag else 'outside_band', keywords=user_data_get(task.get('user_data'), 'keywords', []), seeded=user_data_get(task.get('user_data'), 'seeded', False), n_pass=n, - group_id=user_data_get(task.get('user_data'), 'group_id', None)) - return [t for t, kept_flag in zip(measured, in_band) if kept_flag] + group_id=gid, + novelty=nov, + selected=selected[i], + novelty_dropped=gid in dropped) + return [t for t, take in zip(measured, selected) if take] + + def _unscored_group_ids(self, measured: List[Trajectory], + novelties: List[Optional[float]]) -> set: + """Keyword groups the novelty judge never finished answering for. + + A group lands here when at least one of its proposals still has no score + after ``NOVELTY_TRIES`` attempts. Nothing from such a group is used: no task + is taken from it (``_select_per_group``) and its proposals are marked + ``novelty_dropped`` so training skips the whole group. The collecting loop + then keeps going and a later keyword draw makes up the shortfall. + + Only meaningful when novelty is on. With it off every score is ``None`` by + design, which must not drop everything, so an off judge returns no groups. + """ + if self.task_bank is None or self.novelty_fn is None: + return set() + return {user_data_get(task.get('user_data'), 'group_id', None) + for task, nov in zip(measured, novelties) if nov is None} + + def _select_per_group(self, measured: List[Trajectory], passes: List[int], + in_band: List[bool], + novelties: List[Optional[float]]) -> List[bool]: + """Which in-band proposals become tasks: all of them, or the best few per group. + + With ``keep_per_group = k > 0``, each keyword group contributes at most its ``k`` + highest-reward in-band proposals -- reward being the same number the proposing + side trains on, ``challenger_reward``, so the task kept is the one whose pass + rate sat closest to the target and, when novelty is on, was not judged a repeat + of something already in the bank. + + A group with no in-band proposal contributes nothing and is not replaced here: + the collecting loop keeps proposing rounds until the run's target number of + tasks is reached, so a group that produced none is skipped and paid for by one + more group later. + + Proposals with no ``group_id`` (a run with ``proposals_per_group=1``) are each + their own group, so this is a no-op for them beyond the in-band filter. + """ + if self.keep_per_group <= 0: + return list(in_band) + ranked: Dict[Any, List[Tuple[float, int]]] = {} + unscored_groups = self._unscored_group_ids(measured, novelties) + for i, task in enumerate(measured): + if not in_band[i]: + continue + gid = user_data_get(task.get('user_data'), 'group_id', None) + if gid in unscored_groups: + # The judge never finished scoring this group, so there is no honest + # way to rank its members against each other. + continue + key = gid if gid is not None else f'_ungrouped_{i}' + ranked.setdefault(key, []).append( + (self.challenger_reward(passes[i], novelty=novelties[i]), i)) + selected = [False] * len(measured) + for key, entries in ranked.items(): + # Ties broken by the earlier proposal, so the choice does not depend on + # dict or sort instability. + entries.sort(key=lambda pair: (-pair[0], pair[1])) + for _, i in entries[:self.keep_per_group]: + selected[i] = True + dropped = sum(1 for i in range(len(measured)) if in_band[i] and not selected[i]) + if dropped: + self._bump('in_band_not_selected', dropped) + return selected + + def _score_novelty(self, measured: List[Trajectory]) -> List[Optional[float]]: + """One novelty score per measured proposal, ``None`` for every one if off. + + Scored for the whole batch in one call, and with the batch's own statements + as part of each proposal's reference set, because the comparison that matters + is against the siblings sharing a keyword draw: GRPO subtracts the group mean, + so a term that comes out the same for all eight members of a group cancels + exactly and the API calls bought nothing. Only same-group siblings go in -- + an unrelated proposal from the same round is not evidence of redundancy. + + Failures return ``None`` rather than 0.0 and never raise: a judge that is + down must not turn into every proposal being redundant, and must not lose a + round of sandbox work either. + + A proposal the judge skipped is asked about again, up to ``NOVELTY_TRIES`` + attempts in total, sending only the ones still missing. Whatever is still + unscored after that leaves its whole keyword group out of both the task + selection and the training data -- see ``_select_per_group`` and the + ``novelty_dropped`` field written by ``_emit_propose``. + """ + if self.task_bank is None or self.novelty_fn is None or not measured: + return [None] * len(measured) + statements, checks, groups = [], [], [] + for task in measured: + statements.append(self.statement_of(task)) + checks.append(user_data_get(task.get('user_data'), 'check_script', '') or '') + groups.append(user_data_get(task.get('user_data'), 'group_id', None)) + payload = [] + for i, statement in enumerate(statements): + siblings = [statements[j] for j in range(len(statements)) + if j != i and groups[j] is not None and groups[j] == groups[i]] + payload.append({'statement': statement, 'check': checks[i], + 'references': self.task_bank.references(statement, siblings)}) + scores: List[Optional[float]] = [None] * len(measured) + pending = list(range(len(measured))) + for attempt in range(self.NOVELTY_TRIES): + batch = [payload[i] for i in pending] + try: + got = list(self.novelty_fn(batch)) + except Exception as e: # noqa + logger.warning(f'[{type(self).__name__}] novelty scoring failed on ' + f'{len(batch)} proposals, try {attempt + 1} of ' + f'{self.NOVELTY_TRIES} ({type(e).__name__}: {e})') + self._bump('novelty_error', len(batch)) + continue + if len(got) != len(batch): + # Zipping a short list would pair later proposals with someone + # else's number, so the whole reply is dropped. + logger.warning(f'[{type(self).__name__}] novelty judge returned ' + f'{len(got)} scores for {len(batch)} proposals, try ' + f'{attempt + 1} of {self.NOVELTY_TRIES}; ignoring them') + self._bump('novelty_length_mismatch', len(batch)) + continue + still: List[int] = [] + for i, score in zip(pending, got): + if score is None: + still.append(i) + else: + scores[i] = score + if not still: + break + logger.info(f'[{type(self).__name__}] novelty: {len(still)} of ' + f'{len(batch)} left unscored on try {attempt + 1}; ' + f'asking again') + pending = still + else: + pending = [i for i, s in enumerate(scores) if s is None] + unscored = [i for i, s in enumerate(scores) if s is None] + if unscored: + self._bump('novelty_unjudged', len(unscored)) + logger.warning(f'[{type(self).__name__}] {len(unscored)} proposal(s) ' + f'still unscored after {self.NOVELTY_TRIES} tries; their ' + f'keyword groups are dropped') + return scores + + def statement_of(self, task: Trajectory) -> str: + """The statement text a task was built around: its first user message.""" + for message in task.get('messages') or []: + if isinstance(message, dict) and message.get('role') == 'user': + return message.get('content') or '' + return '' def solver_prompt(self, task: Trajectory) -> Trajectory: """The statement as the solver first sees it: system message, query, tools. diff --git a/src/twinkle_agentic/challenger/task_bank.py b/src/twinkle_agentic/challenger/task_bank.py new file mode 100644 index 000000000..bf8663bde --- /dev/null +++ b/src/twinkle_agentic/challenger/task_bank.py @@ -0,0 +1,130 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""A file of the tasks earlier iterations already produced, to compare new ones against. + +Novelty is meaningless without something to be novel against. Within one run the +proposals of a group can be compared to each other, but the failure this is for is +slower than that: iteration k+1 re-proposing what iteration k already trained on. That +needs a file that outlives a run, which is what this is -- one JSON object per line, +appended by :meth:`add`, read back by the next run. + +The similarity used to pick which stored tasks to show the judge is 3-gram Jaccard over +the statement. It is a weak measure and known to be: measured over run_clean9's 188 +statements the closest pair scored 0.060, so ranking by it is nearly ranking at random, +and it cannot see that two tasks with no shared wording are both 'write the given files +verbatim, then derive one from them'. It is used only to CHOOSE the handful of tasks the +judge reads, never to score novelty -- the judging is +:mod:`twinkle_agentic.verifier.rubric_score`, whose criteria compare task shapes. Even a +near-random pick gives the judge real tasks from the same generator to compare against, +which is what the criteria need. +""" +import json +import os +import re +import threading +from typing import Any, Dict, List, Optional, Sequence, Set, Tuple + +__all__ = ['TaskBank', 'jaccard_3gram', 'grams'] + + +def grams(text: str, n: int = 3) -> Set[Tuple[str, ...]]: + words = re.findall(r'[a-z0-9_]+', (text or '').lower()) + return {tuple(words[i:i + n]) for i in range(max(0, len(words) - n + 1))} + + +def jaccard_3gram(a: Set[Tuple[str, ...]], b: Set[Tuple[str, ...]]) -> float: + if not a or not b: + return 0.0 + return len(a & b) / len(a | b) + + +class TaskBank: + """Statements from previous iterations, plus the ones this run adds. + + Args: + path: the JSONL file. A missing file is an empty bank, not an error -- the + first iteration has nothing to compare against and must still run. + refs: how many stored statements :meth:`references` returns. + """ + + def __init__(self, path: str, refs: int = 5): + self.path = path + self.refs = max(0, refs) + self._statements: List[str] = [] + self._grams: List[Set[Tuple[str, ...]]] = [] + self._seen: Set[str] = set() + self._lock = threading.Lock() + self.n_loaded = 0 + self.n_added = 0 + self._load() + + def _load(self) -> None: + if not self.path or not os.path.exists(self.path): + return + with open(self.path, encoding='utf-8') as f: + for line in f: + line = line.strip() + if not line: + continue + try: + rec = json.loads(line) + except json.JSONDecodeError: + # A half-written last line from a killed run. Skipped rather + # than fatal: losing one reference is not worth failing a run, + # and the count below says how many were read. + continue + statement = (rec.get('statement') or '').strip() + if statement and statement not in self._seen: + self._seen.add(statement) + self._statements.append(statement) + self._grams.append(grams(statement)) + self.n_loaded = len(self._statements) + + def __len__(self) -> int: + return len(self._statements) + + def references(self, statement: str, extra: Sequence[str] = ()) -> List[str]: + """The stored statements most similar to ``statement``, closest first. + + ``extra`` is prepended and never dropped -- it is how the proposals of the + current group get in front of the judge. Without them a whole group can be + scored identically novel against history while being eight versions of one + idea, and GRPO subtracts the group mean, so an identical term across the + group produces no gradient at all. + """ + with self._lock: + pairs = list(zip(self._statements, self._grams)) + target = grams(statement) + scored = [(jaccard_3gram(target, g), s) for s, g in pairs if s != statement] + scored.sort(key=lambda p: -p[0]) + out = [s for s in extra if s and s != statement] + out.extend(s for _, s in scored[:self.refs]) + return out + + def add(self, statement: str, check: str = '', **fields: Any) -> bool: + """Append one task. Returns False if the statement is already stored. + + Appended immediately rather than at the end of the run: a run that crashes + after 60 of 80 tasks should still contribute those 60, or the bank silently + under-reports what has been trained on. + """ + statement = (statement or '').strip() + if not statement: + return False + with self._lock: + if statement in self._seen: + return False + self._seen.add(statement) + self._statements.append(statement) + self._grams.append(grams(statement)) + self.n_added += 1 + if self.path: + rec: Dict[str, Any] = {'statement': statement, 'check': check} + rec.update(fields) + os.makedirs(os.path.dirname(self.path) or '.', exist_ok=True) + with open(self.path, 'a', encoding='utf-8') as f: + f.write(json.dumps(rec, ensure_ascii=False) + '\n') + return True + + def stats(self) -> Dict[str, Optional[int]]: + return {'path': self.path, 'loaded': self.n_loaded, 'added': self.n_added, + 'total': len(self._statements)} diff --git a/src/twinkle_agentic/verifier/__init__.py b/src/twinkle_agentic/verifier/__init__.py index eaf813963..a4e345a94 100644 --- a/src/twinkle_agentic/verifier/__init__.py +++ b/src/twinkle_agentic/verifier/__init__.py @@ -1,7 +1,12 @@ from .result_check import (Check, CheckContext, CheckOutcome, CheckReport, checks_from_dicts, local_runner, run_checks) +from .rubric_score import (CRITERIA, DIMENSIONS, Criterion, RubricResult, + build_rubric_prompt, parse_verdicts, score_task, + score_tasks) __all__ = [ 'Check', 'CheckContext', 'CheckOutcome', 'CheckReport', 'run_checks', 'checks_from_dicts', 'local_runner', + 'CRITERIA', 'DIMENSIONS', 'Criterion', 'RubricResult', + 'build_rubric_prompt', 'parse_verdicts', 'score_task', 'score_tasks', ] diff --git a/src/twinkle_agentic/verifier/rubric_score.py b/src/twinkle_agentic/verifier/rubric_score.py new file mode 100644 index 000000000..db82e54c1 --- /dev/null +++ b/src/twinkle_agentic/verifier/rubric_score.py @@ -0,0 +1,451 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Rubric scores for a proposed task: how new it is, what it is worth, how hard it is. + +:mod:`result_check` scores what a solver *did*, with ordinary programs, which is why it +is stable. This module scores the *task itself*, which no program can read: whether a +statement asks for something the pool does not already contain, whether the thing it +asks for resembles work anyone does, and how much reasoning it takes. + +The shape is taken from the rubric verifier this repo used before (deleted in 5175833; +readable at ``git show 5175833^:src/twinkle_agentic/verifier/rubric_verifier.py``), for +the reason that made it work there: **the judge never emits a score.** It emits PASS or +FAIL per criterion and the number is computed here. Asking a model for "3 out of 4" +spends most of its resolution on distinctions it cannot make twice in a row; asking +"does this task need more than one command" is a question it answers the same way on a +re-run. A dimension's value is therefore a weighted pass fraction over 3 binary +judgements, not a level the judge chose. + +Four more things carried over from that file, with its constants: + +* ``[Hard Rule]`` criteria weigh 3, ``[Principle]`` 1 (``hard_weight=3.0``, + ``principle_weight=1.0``). A hard rule fails unless unambiguously satisfied. +* One vote is normally enough. A second and third are spent only when the first is + undecided -- when a dimension lands within ``margin`` of the middle -- so cost tracks + difficulty rather than volume (``margin_threshold=0.25``). +* Criteria are fixed and generic here, naming no file, value or domain from the task + being judged. Letting a model invent the criteria per task was named in + ``rubric_library.py`` as the main source of score jitter. +* Anything a program can decide does not go to the judge. Whether the statement quotes + the values the check compares against is already computed by + ``derived_check_literals`` (challenger/agentic.py) and is deliberately NOT a criterion + below, so the two never disagree. + +Those four constants are inherited, not re-measured for this use. What has to be +measured before any number here is used: how often a re-run flips a criterion. + +What this is NOT: part of the reward. ``AgenticChallenger.challenger_reward`` is the +pass-rate term alone. Ornith-1.5 multiplies its difficulty term by a novelty term +(``R = V x D x N``, ornith.ai/ornith_1_5.html) and ``novelty`` below is the obvious +candidate, but wiring it in needs one more fact first: GRPO subtracts the group mean +(``GRPOAdvantage(scale='group')``), so a term that is near-constant across the eight +proposals sharing a keyword direction contributes no gradient however sensible it looks +per task. +""" +import os +import re +import statistics +import threading +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Sequence + +__all__ = [ + 'Criterion', + 'RubricResult', + 'CRITERIA', + 'DIMENSIONS', + 'build_rubric_prompt', + 'parse_verdicts', + 'score_task', + 'score_tasks', +] + +# Inherited from the deleted rubric_verifier.py (ARROW's 3 / 1) -- not re-derived here. +HARD_WEIGHT = 3.0 +PRINCIPLE_WEIGHT = 1.0 +# The old verifier escalated an undecided result to 3 votes. Measured on 188 tasks from +# run_clean9 that buys nothing here: tasks whose first pass was decisive repeated to +# within 0.043, tasks that spent all 3 votes to within 0.051 -- no better, at 3x the +# calls (92 of 188 tasks escalated). So one vote, and the spread is reported rather +# than voted away. Raise MAX_VOTES to bring the escalation back; the threshold still +# controls when it triggers. +MARGIN_THRESHOLD = 0.25 +MAX_VOTES = 1 + + +@dataclass +class Criterion: + """One yes/no question about the task. + + Args: + dimension: which score it contributes to. + text: the question, phrased so that PASS is the good direction. A criterion + whose PASS means "this task is bad" inverts the aggregate silently. + is_hard: objectively checkable from the statement and check as written, so a + FAIL is not a matter of taste. Weighed ``HARD_WEIGHT``. + needs_references: skipped when no comparison set was supplied. + """ + dimension: str + text: str + is_hard: bool + needs_references: bool = False + + +# Every criterion below is phrased so PASS is the good direction, and every one was +# either kept or replaced on evidence from a first run over run_clean9's 188 tasks +# (.tmp_analysis/rubric_run_clean9.json, kept as rubric_v1.json). What that run showed: +# +# * The first two novelty criteria agreed on 188 of 188 tasks -- one of them was +# free. Both are gone, replaced by three that ask about different things: the +# shape of the task, the machinery it needs, and the form of its end state. Those +# are the three axes a labelling pass over the same pool found the collapse in +# (54% of tasks were 'write a script that simulates a process'). +# * The judge was deciding novelty by DOMAIN, not by what the task does: a task whose +# skeleton was identical to the ones it scored 0.0 got 1.0 because it was about +# PCIe rather than about log files. The shapes are therefore enumerated, and the +# criterion says outright that a different domain is not a different task. +# * 'Reaching the end state takes more than one command' passed 80% of the time and +# passed on 16 tasks that all eight solvers then solved. Replaced by whether the +# obvious untested attempt fails, which is what 'hard' has to mean here. +# * The two soft usefulness criteria barely moved the dimension (it tracked its hard +# criterion: 0.22 mean when that failed, 0.93 when it passed), so both were +# replaced. One of the replacements -- whether the input data looks like a real +# sample -- then passed 7% of the time, i.e. decided nothing, and what it was +# reaching for is countable without a model anyway: 30% of these statements paste +# .py source in as an "input file", and those tasks are the easy ones (n_pass 5.7 +# vs 4.4). That belongs in a regex, not in a rubric, so the criterion now asks the +# part a regex cannot: whether the statement dictates the code to write. +CRITERIA: List[Criterion] = [ + # -- novelty: three independent axes, judged only against the reference set ---- + Criterion( + 'novelty', + 'This task has a different SHAPE from every reference task. Shapes: (a) write ' + 'given input files verbatim, then produce a derived file from them; (b) write ' + 'a script that demonstrates a defect and a second that fixes it; (c) build a ' + 'database or structured store and populate it; (d) run something and report ' + 'timings or counts; (e) parse a log or config and summarise it; (f) anything ' + 'not in this list. Two tasks of the same shape are the same task here EVEN IF ' + 'they are about different subject matter -- a different domain, file format or ' + 'vocabulary does not make a different shape', + is_hard=True, needs_references=True), + Criterion( + 'novelty', + 'Solving this needs machinery that no reference task needs -- a different one ' + 'of: plain text handling, tabular data, binary formats, a database, threads or ' + 'processes, subprocesses, sockets, the filesystem layout itself, timing', + is_hard=False, needs_references=True), + Criterion( + 'novelty', + 'The FORM of the end state differs from every reference: one text file, several ' + 'files, a database file, a program that must run correctly, or a directory tree', + is_hard=False, needs_references=True), + # -- usefulness: the hard criterion kept as-is, it separated 29 from 159 and the + # calls held up on inspection. + Criterion( + 'usefulness', + 'The end state is something a person would want for its own sake, not only as ' + 'an exercise', + is_hard=True), + Criterion( + 'usefulness', + 'The statement says what the end state must be and leaves how to reach it to ' + 'the solver, rather than dictating the code or commands to write', + is_hard=False), + Criterion( + 'usefulness', + 'The task would still be worth doing if the input were a thousand times larger', + is_hard=False), + # -- complexity: the hard criterion asks for a countable property of the task. + # Asking instead whether 'the obvious untested attempt would fail' made the + # judge guess at a counterfactual and it flipped on 13% of re-runs -- the worst + # of the nine, and it carries weight 3. + Criterion( + 'complexity', + 'Reaching the end state takes at least three steps that depend on each other, ' + 'where a later step needs the result of an earlier one', + is_hard=True), + Criterion( + 'complexity', + 'Reaching a passing state means choosing between at least two plausible ' + 'approaches, of which at least one does not work', + is_hard=False), + Criterion( + 'complexity', + 'Passing requires computing something: writing the expected output as a ' + 'literal would not satisfy the check', + is_hard=False), +] + +DIMENSIONS = ('novelty', 'usefulness', 'complexity') + + +@dataclass +class RubricResult: + """One task's scores plus the verdicts they were computed from. + + ``scores[dim]`` is the weighted PASS fraction in [0, 1], or ``None`` when the + dimension was not judged -- an unparseable reply, or novelty with no references. + ``None`` rather than 0.0 so an unjudged task drops out of a mean instead of + dragging it down. + """ + scores: Dict[str, Optional[float]] = field(default_factory=dict) + verdicts: List[Optional[bool]] = field(default_factory=list) + pass_rates: List[Optional[float]] = field(default_factory=list) + n_votes: int = 0 + raw: List[str] = field(default_factory=list) + error: str = '' + + @property + def ok(self) -> bool: + return not self.error and any(v is not None for v in self.scores.values()) + + def to_dict(self) -> Dict[str, Any]: + out: Dict[str, Any] = {k: self.scores.get(k) for k in DIMENSIONS} + out['n_votes'] = self.n_votes + out['verdicts'] = list(self.verdicts) + if self.error: + out['error'] = self.error + return out + + +# Criterion 1's shape list stays INSIDE the criterion. Moving it to its own prompt +# section, so the criterion read 'the shapes listed below', made the judge less stable +# rather than more: verdict flips between two runs went 5% -> 9% and novelty's run-to-run +# spread 0.041 -> 0.100 over the same 60 tasks. What actually stopped the judge from +# answering criterion 1 with a shape name ('1: f', which cost 2 of 188 tasks their +# novelty score) is the paragraph below forbidding it. +_SYSTEM = ( + 'You judge a programming task that was generated automatically, before it is used ' + 'to train a model.\n\n' + 'The task has two parts. The STATEMENT is everything a solver sees: it starts in an ' + 'empty directory, cannot ask questions, and never sees the check. The CHECK is a ' + 'python script run against the solver\'s directory afterwards, where exit 0 means ' + 'passed. The check is shown to you because it is what the task really demands, ' + 'which the statement can understate.\n\n' + 'For each numbered criterion output one line:\n\n' + ' <index>: PASS or <index>: FAIL\n\n' + 'PASS and FAIL are the only two words you may write after the index. Some criteria ' + 'list categories to compare by; those are there to define the question, never to be ' + 'answered with -- naming a category instead of a verdict makes the line unusable.\n\n' + 'Judge every criterion independently and literally, against this task only. A ' + '[Hard Rule] is FAIL unless it is unambiguously satisfied. Do not explain, do not ' + 'restate the criterion, output only the verdict lines in order and then stop.\n') + + +def _applicable(references: Sequence[str], + criteria: Sequence[Criterion] = CRITERIA) -> List[Criterion]: + return [c for c in criteria if references or not c.needs_references] + + +def build_rubric_prompt( + statement: str, + check: str = '', + references: Sequence[str] = (), + criteria: Sequence[Criterion] = CRITERIA, + reference_chars: int = 600, +) -> List[Dict[str, str]]: + """The messages sent to the judge, and the criterion order the reply must follow. + + All three dimensions go in one call: the judge reads the task once, and nine + yes/no lines cost about what one dimension would. The cost is that one dimension + can colour another -- if the scores turn out to move together, splitting into one + call per dimension is the fix, and the correlation is measurable from the dumps. + + References are cut to ``reference_chars`` each. What a task asks for is in its + first paragraph; sending statements whole would spend the context on input data + quoted verbatim, which is the bulk of a statement here. + """ + items = _applicable(references, criteria) + lines = [f'{i + 1}. {c.text} [{"Hard Rule" if c.is_hard else "Principle"}]' + for i, c in enumerate(items)] + parts = ['## Criteria\n' + '\n'.join(lines) + '\n'] + if references: + parts.append('\n## Reference tasks (for the novelty criteria only)\n') + for i, ref in enumerate(references): + parts.append(f'[{i}] {(ref or "")[:reference_chars]}\n') + parts.append('\n## Statement\n' + (statement or '') + '\n') + if check: + parts.append('\n## Check\n' + check + '\n') + parts.append(f'\nNow output {len(items)} verdict lines, in order.') + return [{'role': 'system', 'content': _SYSTEM}, + {'role': 'user', 'content': ''.join(parts)}] + + +# Same tolerant form the previous verifier parsed, so a reply written as '1) yes' or +# '1. FAIL' is read rather than thrown away. +_VERDICT_RE = re.compile(r'^\s*(\d+)\s*[:.)]\s*(pass|fail|true|false|yes|no|1|0)\b', + re.IGNORECASE) +_TRUE = {'pass', 'true', 'yes', '1'} + + +def parse_verdicts(raw: str, n: int) -> List[Optional[bool]]: + """Read ``n`` PASS/FAIL verdicts. A line that is missing stays ``None``. + + Indexed by the number the judge wrote rather than by position, because a reply + that skips a criterion would otherwise shift every later verdict onto the wrong + question -- and the scores would still come out as numbers. + """ + out: List[Optional[bool]] = [None] * n + for line in (raw or '').splitlines(): + match = _VERDICT_RE.match(line) + if not match: + continue + idx = int(match.group(1)) - 1 + if 0 <= idx < n: + out[idx] = match.group(2).lower() in _TRUE + return out + + +def _aggregate(items: Sequence[Criterion], + rates: Sequence[Optional[float]]) -> Dict[str, Optional[float]]: + """Weighted PASS fraction per dimension; ``None`` when nothing was judged.""" + totals: Dict[str, List[float]] = {} + for crit, rate in zip(items, rates): + if rate is None: + continue + weight = HARD_WEIGHT if crit.is_hard else PRINCIPLE_WEIGHT + got, tot = totals.setdefault(crit.dimension, [0.0, 0.0]) + totals[crit.dimension] = [got + weight * rate, tot + weight] + return {dim: (totals[dim][0] / totals[dim][1] if dim in totals else None) + for dim in DIMENSIONS} + + +def _undecided(scores: Dict[str, Optional[float]], margin: float) -> bool: + """Is any dimension close enough to the middle that another vote could move it?""" + return any(v is not None and margin < v < 1.0 - margin for v in scores.values()) + + +_client = None +_client_lock = threading.Lock() + + +def _get_client(model: Optional[str] = None): + """The judge API, from the same environment variables llm_backup.py reads. + + Default model is ``qwen3.8-max``, the same one that writes the check scripts and + problem statements, so a task is judged by the model that phrased it. + + Note for comparing numbers: every rubric measurement on file -- the criterion + flip rates, the per-dimension spreads, the 4-in-940 rate of replies with no + usable verdict -- was taken with ``qwen3-max``, which was the default until now. + Those are not a baseline for this judge. + """ + global _client + if model is None and _client is not None: + return _client + from twinkle_agentic.protocol.openai import OpenAI + client = OpenAI( + model=model or os.environ.get('RUBRIC_MODEL') + or os.environ.get('LLM_BACKUP_MODEL', 'qwen3.8-max'), + api_key=os.environ.get('LLM_BACKUP_API_KEY'), + base_url=os.environ.get('LLM_BACKUP_BASE_URL'), + client_kwargs={'timeout': float(os.environ.get('LLM_BACKUP_TIMEOUT', '120')), + 'max_retries': int(os.environ.get('LLM_BACKUP_MAX_RETRIES', '2'))}, + ) + if model is None: + with _client_lock: + _client = client + return client + + +def score_task( + statement: str, + check: str = '', + references: Sequence[str] = (), + *, + criteria: Sequence[Criterion] = CRITERIA, + model: Optional[str] = None, + temperature: float = 0.0, + max_tokens: int = 256, + margin: float = MARGIN_THRESHOLD, + max_votes: int = MAX_VOTES, + extra_body: Optional[Dict[str, Any]] = None, + client: Any = None, +) -> RubricResult: + """Score one task. Never raises: an API failure comes back in ``error``. + + Votes past the first are spent only on an undecided result, and they are sampled + (temperature 1.0) whatever ``temperature`` says -- repeating a temperature-0 call + would mostly repeat its answer, which reads as agreement without being any. + ``max_tokens`` is small because the reply is nine short lines; a judge that starts + explaining gets cut off, and the verdict lines it already wrote are still read. + + ``extra_body`` is forwarded on every call, and on a reasoning judge it is what + makes the call finish. Measured on one real payload (2044 prompt tokens, 9 + criteria, 6 references) against ``qwen3.8-max``: left alone the judge spent 3757 + reasoning tokens and 93 seconds to write 38 tokens of verdicts, and + ``LLM_BACKUP_TIMEOUT`` at its default of 120s cut off about half of a 27-call + batch. ``max_tokens`` does not bound this -- it bounds the visible answer only, + which is why 256 neither truncated a verdict nor prevented a timeout. With + ``{'thinking_budget': 512}`` the same payload came back in 11 seconds with the + same nine verdicts. + """ + from twinkle.data_format.sampling import SamplingParams + items = _applicable(references, criteria) + messages = build_rubric_prompt(statement, check, references, criteria) + api = client or _get_client(model) + result = RubricResult() + votes: List[List[Optional[bool]]] = [] + + for attempt in range(max(1, max_votes)): + params = SamplingParams( + max_tokens=max_tokens, + temperature=temperature if attempt == 0 else 1.0, + top_p=0.95, + num_samples=1) + try: + message = api({'messages': messages}, params, + **({'extra_body': extra_body} if extra_body else {})) + except Exception as e: # noqa + if not votes: + result.error = f'{type(e).__name__}: {e}' + return result + break + if isinstance(message, list): + message = message[0] if message else {} + content = message.get('content', '') if isinstance(message, dict) else '' + result.raw.append(content) + votes.append(parse_verdicts(content, len(items))) + result.n_votes = len(votes) + + # Mean over the votes cast so far, per criterion, then aggregate. Voting on + # each criterion separately rather than on the final number is what keeps one + # flipped criterion from moving the whole dimension. + rates: List[Optional[float]] = [] + for i in range(len(items)): + seen = [v[i] for v in votes if v[i] is not None] + rates.append(statistics.fmean(1.0 if s else 0.0 for s in seen) if seen else None) + result.pass_rates = rates + result.verdicts = [None if r is None else r >= 0.5 for r in rates] + result.scores = _aggregate(items, rates) + if not _undecided(result.scores, margin): + break + + if all(v is None for v in result.scores.values()): + result.error = result.error or 'no usable verdict in reply' + # Novelty is absent rather than zero when there was nothing to compare against. + for crit in criteria: + if crit.needs_references and not references: + result.scores.setdefault(crit.dimension, None) + return result + + +def score_tasks( + tasks: Sequence[Dict[str, Any]], + *, + workers: int = 8, + **kwargs, +) -> List[RubricResult]: + """Score ``{statement, check, references}`` dicts, order preserved. + + Concurrency is over API calls only; nothing here touches a GPU or a sandbox. + """ + + def _one(task: Dict[str, Any]) -> RubricResult: + return score_task(task.get('statement') or '', task.get('check') or '', + task.get('references') or (), **kwargs) + + if workers <= 1: + return [_one(t) for t in tasks] + with ThreadPoolExecutor(max_workers=workers) as pool: + return list(pool.map(_one, tasks)) From 01fb8f9dcbef3361f237d955b8be12cbbcde58ce Mon Sep 17 00:00:00 2001 From: tastelikefeet <yuze.zyz@alibaba-inc.com> Date: Mon, 31 Aug 2026 19:13:52 +0800 Subject: [PATCH 48/60] wip --- cookbook/rsi/agentic/challenge.py | 80 ++++++++++++--- cookbook/rsi/agentic/loop.sh | 10 +- cookbook/rsi/agentic/prompts.py | 163 ++++++++++++++++++++++++++++-- cookbook/rsi/agentic/train.py | 14 ++- 4 files changed, 245 insertions(+), 22 deletions(-) diff --git a/cookbook/rsi/agentic/challenge.py b/cookbook/rsi/agentic/challenge.py index 2534ddac4..da959a172 100644 --- a/cookbook/rsi/agentic/challenge.py +++ b/cookbook/rsi/agentic/challenge.py @@ -8,7 +8,13 @@ worker (``enable_continous_work``), so a batch of one is a first-class call and 32 threads calling it concurrently is the intended use. API qwen3.8-max, for the stages that must not add untrained tokens: the - check script, the problem statement, and the rubric. + check script, the problem statement, the rubric, and the keyword bank. + Keywords joined this list after measuring what the local model produced + for it: 31% of ``transform`` entries named an activity on a running + system rather than a computation, and 71% of the bank comes from an + expand prompt that had no category rules in it at all. Iteration 9 ran + 18 refills through the API with the rules added and 1 of 105 keywords + missed, against 24% of the bank built without them. Nothing waits for a batch. A proposal that finishes its build hands its statement straight to eight solver jobs and lets go of its sandbox; those eight run whenever @@ -682,6 +688,15 @@ def __init__(self, args, sampler, template, slots: List[Any], recorder: Recorder self.problem_params = SamplingParams(max_tokens=args.problem_max_tokens, num_samples=1, temperature=args.propose_temp, top_p=0.95) + # Keeps the local path's temperature and top_p rather than the 1.0/0.95 the + # other two API stages use. The high temperature is deliberate here -- the + # bank is worthless if every refill returns the same phrases -- and moving + # the model and the temperature in one step would leave no way to tell which + # one changed the result. + self.keyword_api_params = SamplingParams(max_tokens=args.keyword_max_tokens, + num_samples=1, + temperature=args.keyword_temp, + top_p=0.98) # Built once: the cap is part of the system prompt, so a build that got a # different one would be a different experiment. @@ -788,10 +803,35 @@ def generate_keywords(self, category: str) -> bool: + ('\nDo NOT repeat any of these already-used topics: ' + ', '.join(avoid) if avoid else '') + f'\n(batch {self.nonce}-{i})') - traj = {'messages': [{'role': 'system', 'content': P.KEYWORD_SYSTEM}, - {'role': 'user', 'content': user}]} - out = self.keyword_rollout([traj]) - reply = self._assistant_text(out[0] if out else {}) + # Asked of the API model rather than the local one. Keyword text never + # enters a trajectory -- it is parsed into a list and thrown away -- so + # this adds no untrained tokens, which is the rule that decides what may + # use the API. And the bank is the single input every task downstream is + # built from: measured over the 1344 keywords iterations 1-7 generated + # locally at temperature 1.3, against category rules the model is shown + # in full, 31% of transform named an activity on a running system rather + # than a computation, 13% of domain named an operation rather than + # material, and 24% of edge_case needed hardware the container does not + # have. Downstream, 42-70% of statements described themselves as + # simulating their own subject matter, which is what a keyword the + # sandbox cannot honour turns into. A 4B policy at that temperature is + # the wrong instrument for a constraint list this long. + # + # Falls back to the local model instead of giving up: an unreachable API + # must not leave a category dry, because dry means keyword-less prompts + # and a run that looks healthy while producing one prompt over and over + # -- the exact failure the refill logic already guards against. + out = None + messages = [{'role': 'system', 'content': P.KEYWORD_SYSTEM}] + reply = api_one(self.api, messages, user, self.keyword_api_params, + self.api_extra) + via = 'api' + if reply is None: + traj = {'messages': [{'role': 'system', 'content': P.KEYWORD_SYSTEM}, + {'role': 'user', 'content': user}]} + out = self.keyword_rollout([traj]) + reply = self._assistant_text(out[0] if out else {}) + via = 'local-fallback' parsed = parse_keyword_list(reply) new = [k for k in parsed if k.lower() not in seen] for keyword in new: @@ -799,7 +839,7 @@ def generate_keywords(self, category: str) -> bool: fresh.extend(new) self.rec.keywords({'category': category, 'prompt': user, 'reply': reply, 'parsed': parsed, 'n_parsed': len(parsed), - 'n_new': len(new), + 'n_new': len(new), 'via': via, 'stop_reason': (out[0].get('stop_reason') if out else None), 'truncated': bool(out[0].get('truncated')) if out else None}) if len(fresh) >= want: @@ -835,17 +875,29 @@ def expand_hard_keywords(self) -> int: added = 0 for i, (category, keyword) in enumerate(hard): self.nonce += 1 - traj = {'messages': [ - {'role': 'system', 'content': P.KEYWORD_SYSTEM}, - {'role': 'user', - 'content': P.KEYWORD_EXPAND_USER.format(kw=keyword, m=8) - + f'\n(batch {self.nonce}-{i})'}]} - out = self.keyword_rollout([traj]) - reply = self._assistant_text(out[0] if out else {}) + user = (P.KEYWORD_EXPAND_USER.format(kw=keyword, m=8, + desc=P.CATEGORY_DESC[category]) + + f'\n(batch {self.nonce}-{i})') + # Same reasoning as the refill path: the API model, falling back to the + # local one. This path matters more, not less -- it wrote 960 of the 1344 + # keywords the first seven iterations banked, at more than twice their + # rule-break rate, so it is the one shaping what later iterations draw. + out = None + messages = [{'role': 'system', 'content': P.KEYWORD_SYSTEM}] + reply = api_one(self.api, messages, user, self.keyword_api_params, + self.api_extra) + via = 'api' + if reply is None: + traj = {'messages': [{'role': 'system', 'content': P.KEYWORD_SYSTEM}, + {'role': 'user', 'content': user}]} + out = self.keyword_rollout([traj]) + reply = self._assistant_text(out[0] if out else {}) + via = 'local-fallback' parsed = parse_keyword_list(reply) added += self.store.add(category, parsed, source='expand', parent=keyword) self.rec.keywords({'category': category, 'parent': keyword, 'prompt': 'expand', - 'reply': reply, 'parsed': parsed, 'n_parsed': len(parsed)}) + 'reply': reply, 'parsed': parsed, 'n_parsed': len(parsed), + 'via': via}) self.store.save() logger.info(f'[challenge] expanded {len(hard)} hard keyword(s) -> ' f'+{added} same-domain topics') diff --git a/cookbook/rsi/agentic/loop.sh b/cookbook/rsi/agentic/loop.sh index f07496426..e78dedcb3 100644 --- a/cookbook/rsi/agentic/loop.sh +++ b/cookbook/rsi/agentic/loop.sh @@ -140,7 +140,15 @@ ROOT="output/rsi_agentic/${TAG}" # disk holds one 4B model rather than one per iteration. The previous round's # weights are gone once the next save starts: if a save dies partway there is # nothing to fall back to but BASE_MODEL. -CKPT_DIR="$ROOT/ckpt" +# +# Overridable because this is the one path whose filesystem shows up in wall-clock: +# every iteration reads it 1 + GPUS times, once per vLLM worker at collect and once +# more at train, so 7.6 GB of weights is around 60 GB of reads per iteration. Two +# filesystems on this host, same 1.5 T free, measured with dd at 1.5 GB: the repo's +# own disk reads at 223 MB/s and the parallel one at 1074 MB/s -- 4.8x, which showed +# up as five minutes of vLLM startup before any topic was launched. Left at $ROOT/ckpt +# by default so a host with one disk needs to know nothing about this. +CKPT_DIR="${CKPT_DIR:-$ROOT/ckpt}" # Written with ${VAR-default} rather than ${VAR:-default} so that TASK_BANK="" # means off; with the colon an empty value would silently get the default back. TASK_BANK="${TASK_BANK-$ROOT/task_bank.jsonl}" diff --git a/cookbook/rsi/agentic/prompts.py b/cookbook/rsi/agentic/prompts.py index 61c61a9e2..5ea290239 100644 --- a/cookbook/rsi/agentic/prompts.py +++ b/cookbook/rsi/agentic/prompts.py @@ -31,6 +31,37 @@ # Each category gives three examples, then pushes away from them, then pins the # answer to what the sandbox can actually build and read back. +# +# The three pins below were each added against a measured miss rate. Counted over +# the 1344 keywords iterations 1-7 put in keywords.jsonl, scored by the regexes in +# .temp/prune_keywords.py so the numbers can be reproduced (they are lower bounds -- +# a phrase can be wrong without matching): +# transform 517 entries, 37% miss: 31% named an ACTIVITY on a running system +# rather than a computation ("Debug memory leaks in multi-threaded +# applications", "Monitor system load metrics", "Swap Space +# Configuration") and 10% needed hardware or kernel access +# domain 503 entries, 17% miss: 13% named what someone DOES rather than what +# it is done to ("File format conversion", "Binary data parsing", +# "network namespace isolation"), despite the existing rule already +# saying "a FILE FORMAT or a DATA STRUCTURE, never a device"; 4% +# named a device. Hand-reading a sample puts this category higher than +# the regex does -- "Optimize server performance" is a domain entry +# and matches nothing -- so 17% is the floor, not the estimate +# edge_case 324 entries, 24% needed real hardware or a kernel subsystem +# ("USB device enumeration delay", "Linux bridge MAC addresses") +# The edge_case number had a plain cause: this category was the only one with no +# container pin at all, so it was free to name devices. +# +# Why this matters downstream: a keyword the container cannot honour does not +# produce a hard task, it produces a pretend one. Across iterations 1-7, 42-70% +# of statements (mean 53%) described themselves as simulating or synthesising +# their own subject matter, which is what "analyse TCP congestion" collapses into +# when there is no TCP stack to look at. The task then tests whether the solver +# can follow a spec for generating fake data. +# +# Not measured: whether these three additions actually lower those rates. They +# are worded to name the failure rather than restate the rule, because the +# existing domain pin shows a rule the generator agrees with and ignores. _LEAVE_THE_EXAMPLES = ( '. These three are only to show the form of an answer -- do NOT stay ' 'near them; name things from as many different areas of computer ' @@ -42,10 +73,27 @@ 'reportlab, lxml, pyarrow, matplotlib), sqlite3, ffmpeg, imagemagick, git, ' 'jq, tar/zip/7z, poppler-utils and pip -- and NO compiler, no GPU, no docker, ' 'no hardware devices. Name a FILE FORMAT or a DATA STRUCTURE, never a device ' - 'or a service') + 'or a service, and never an ACTIVITY carried out on material: "binary data ' + 'parsing", "file format conversion" and "traffic analysis" all name something ' + 'a person does, not something that sits in a file waiting to be read') _PINNED_COMPUTATION = ( ', but only computations that run in that same container: not compiling, not ' - 'flashing firmware, not driving hardware') + 'flashing firmware, not driving hardware. It has to be a FUNCTION of data ' + 'that can sit in a file -- given the input there is one right answer, and a ' + 'script can recompute it and check it. An activity carried out on a live ' + 'system is not one: "debug X", "monitor X", "detect X in real time", ' + '"configure X" and "X strategy" have no answer to check, so name the ' + 'calculation instead ("reconstruct the allocation timeline from a heap trace" ' + 'rather than "debug memory leaks")') +# edge_case had no pin before, and 18% of what it produced needed a device. A +# twist is only usable if it survives being written down in a file: the solver +# starts in an empty directory and can only be handed data. +_PINNED_EDGE = ( + ', and only a twist that can be REPRODUCED from data in a file: a property of ' + 'the input or of the arithmetic over it. Not the behaviour of a device, a ' + 'kernel subsystem, a real clock, a network peer or another process -- those ' + 'cannot be put in the solver\'s empty directory, so a task built on them can ' + 'only pretend') CATEGORY_DESC = { 'transform': 'a specific, non-trivial transformation the solver must COMPUTE ' @@ -58,7 +106,8 @@ + _PINNED_TO_CONTAINER, 'edge_case': 'a twist that makes a naive or copy-the-statement solution fail ' 'and forces careful handling. For example: floating-point ' - 'rounding, byte order, cycles in a tree' + _LEAVE_THE_EXAMPLES, + 'rounding, byte order, cycles in a tree' + _LEAVE_THE_EXAMPLES + + _PINNED_EDGE, } # ── Keyword generation ───────────────────────────────────────────────────── @@ -79,10 +128,53 @@ 'Return ONLY a JSON array of short strings, nothing else.' ) +# This prompt writes most of the bank, and until now it was the only one with no +# category rules in it. Of the 1344 keywords iterations 1-7 produced, 960 (71%) came +# from here and 384 from the refill above -- and 31% of the expanded ones break their +# category's rules against 14% of the generated ones. The mechanism is visible in the +# data: asked for keywords related to "Deduce network protocol versions", a legitimate +# computation over captured bytes, it returned "Analyze network traffic patterns", +# "Troubleshoot DNS resolution issues", "Debug TCP/IP stack issues", "Review firewall +# rule sets" and "Monitor honeypot logs" -- each a step further from anything a check +# script can verify. One good keyword decays into eight bad ones, and those eight are +# what later iterations draw from. +# +# So the category description goes in, and with it a sentence saying that being +# related to the parent does not excuse leaving the category. That second part is +# load-bearing: every parent here was chosen for being HARD, and a keyword can be +# hard precisely because the sandbox cannot honour it, in which case following it +# faithfully is the wrong move. +# +# Measured after the change, over iteration 9's 18 refill calls (all via the API, +# ``keyword_gen.jsonl`` 'via' field): 1 of 105 accepted keywords breaks its category's +# rules, against 24% of the bank iterations 1-7 built, and that one is a false positive +# of the scorer ("Amdahl's law speedup bound from parallel workload profile", flagged +# on the noun "profile"). The wording works. +# +# What it broke: 39 of the 144 keywords the model returned (27%) were silently dropped +# by ``parse_keyword_list``, which keeps only strings of 60 characters or less. All but +# one were transform -- five of its six calls came back with a median length of 65-98 +# characters, one with all eight over the cap and nothing left. Cause is in this file: +# KEYWORD_USER says "2-5 words" and this prompt only said "short strings", so the +# instruction to name a calculation rather than an activity ("reconstruct the +# allocation timeline from a heap trace") was followed at sentence length. domain came +# back at a 4-23 character median and edge_case at 32-42, both well clear. The cap is +# named here in characters because it is a silent filter in library code: a keyword +# over it does not warn, it just never exists. + KEYWORD_EXPAND_USER = ( 'The keyword "{kw}" produced a very hard task. List {m} related keywords ' - 'in the same domain that might produce similarly challenging but different ' - 'tasks. Return ONLY a JSON array of short strings, nothing else.' + 'that might produce similarly challenging but different tasks.\n\n' + 'They belong to this category, whose rules bind them exactly as they bound ' + 'the keyword above:\n{desc}\n\n' + 'Being related to "{kw}" does not exempt them. If that keyword itself sits ' + 'outside these rules -- and it may, since it was picked only for being hard ' + '-- move back towards the rules instead of following it further out.\n\n' + 'Each must be 2-5 words and at most 60 characters: a topic to build a task ' + 'around, not a description of the task. "heap free-list reconstruction" is ' + 'one; "reconstruct the heap free-list state from a sequenced alloc/free ' + 'trace" is a task statement and will be thrown away. Return ONLY a JSON ' + 'array of short strings, nothing else.' ) @@ -132,11 +224,35 @@ 'model in about twenty tool calls.' ) +# The three keywords are drawn independently, one per category, with nothing +# checking that they belong together (``draw_keywords`` takes a random unused entry +# from each). So a proposal regularly gets a triple no honest task covers -- iter7 +# produced "TCP Congestion Control" + "Geospatial algorithms" + "hash collision", +# and iter1 "Detect memory leaks in real-time" + "Guitar tablature" + "Thread +# stack fragmentation". Told to exercise all three, the model has one way out: +# invent data that stands in for the parts it cannot have, which is how 42-70% of +# statements (mean 53%) across iterations 1-7 came to describe themselves as +# simulating their own subject matter. +# +# The escape hatch below is deliberately not "ignore a keyword": that would lose +# the diversity the draw exists to create, and the keyword bank's used-marks would +# stop describing what was actually built. Demoting one to background keeps the +# draw meaningful while letting the task be about something real. +# +# Not measured: the effect on the simulate rate, and the cost in diversity if the +# model demotes more often than it needs to. Both are visible in the next run -- +# the statements are in tasks.jsonl and the draws in groups.jsonl. FROM_KEYWORDS = ( 'Your direction for this task:\n{keywords}\n\n' 'Build something complex and realistic that exercises the topics above. ' 'Work in the ' - 'current empty directory, producing files and/or computed output.' + 'current empty directory, producing files and/or computed output.\n\n' + 'Those three are a starting point, not a checklist. If all three can only be ' + 'combined by pretending -- generating fake data to stand in for something ' + 'this container cannot have, or inventing a scenario no engineer would meet ' + '-- then let ONE of them stay in the background and build a task the other ' + 'two support honestly. A real computation over material you actually ' + 'constructed is worth more than a simulation that name-checks everything.' ) # ── Stage 2: write the check script ──────────────────────────────────────── @@ -160,6 +276,27 @@ # The noise floor from sampling one prompt twice is 2 points on the source-text rate # and 14 on handover, so nothing moved. +# The rule about DERIVED values was added last, against a case where every other +# rule was satisfied and the check still did not test the task. An iter7 proposal +# specified a hash table with bucket size 100 and chaining for collisions, but the +# key was (latitude + longitude) % 100 on floats, so 1000 coordinates produced 1000 +# distinct keys and not one collision ever happened. Its check asserted the bucket +# count, the threshold, the CSV header and the first coordinate -- all true, all +# shell -- and passed with reward 0.986, so the solver trained on a task whose +# stated subject was never exercised. +# +# Measured over the 533 check scripts of iterations 1-7: median 6 asserts (range +# 2-16, 45% outside the 2-6 the rules ask for) and 46% run a program via +# subprocess. So the shortage is not in volume. Note 46% against the 86-90% +# recorded above: those were measured on run_clean9's workspaces, and here the +# build stage usually leaves its outputs on disk already, so reading them is +# legitimate. +# +# Not added, for lack of evidence: a rule against matching a float by its printed +# digits. It looked like a problem from one example ('58.54579654631016: [0]' in +# content) but only 2 of 533 scripts compare floats without a tolerance, and the +# rest already use abs(got - expected) < eps. A rule earns its words here. + CHECK_FOLLOWUP = ( 'Now write a python script that ASSERTS properties of the state you just ' 'produced. It runs in the same directory you worked in.\n\n' @@ -175,6 +312,11 @@ '- 2-6 asserts, standard library only.\n' '- Assert only about files holding RESULTS. Never about the text of a ' 'program: not a line it contains, not a name it mentions, not its length.\n' + '- At least one assert must pin a DERIVED value: something no one could ' + 'write down without doing the computation -- a total, an ordering, a decoded ' + 'field, a solved quantity. Existence of a file, a header row, a column name ' + 'and a value copied from the input are all shell: a program that produced ' + 'them and got the arithmetic wrong must still fail this script.\n' '- If a result only exists once a program runs, RUN it -- ' 'subprocess.run([sys.executable, "thing.py"], capture_output=True, ' 'text=True) -- and assert on what it printed or the files it left.\n' @@ -193,6 +335,11 @@ ) # ── Stage 2b: the one chance to fix a check that did not pass ────────────── +# "Drop an assertion you cannot make true" and the new DERIVED rule pull against +# each other: the assert most likely to fail here is exactly the derived one, since +# the shell asserts (a path exists, a header matches) were already true when they +# were written. Dropping it is the cheapest way to make the script pass, and it +# lands back at the check that tests nothing. Hence the carve-out below. CHECK_RETRY_FOLLOWUP = ( 'That script does not pass. Running it in that directory gave:\n\n' @@ -205,6 +352,10 @@ 'here. Drop an assertion you cannot make true instead of weakening every one ' 'of them; what stays must still fail for a directory that does not hold this ' 'state.\n\n' + 'One assertion you may not drop: the one pinning a computed value. If it is ' + 'the one that failed, correct it against the listing -- read the value there ' + 'and assert that -- because a script left asserting only paths, headers and ' + 'input values passes for a program that got the computation wrong.\n\n' 'Same rules as before: standard library only, 2-6 asserts, no file sizes, ' 'checksums, timestamps, script source text, whole-file exact-string ' 'equality, or claims about the exact set of files in the directory. Keep it ' diff --git a/cookbook/rsi/agentic/train.py b/cookbook/rsi/agentic/train.py index 297ca14bb..a017d1ae0 100644 --- a/cookbook/rsi/agentic/train.py +++ b/cookbook/rsi/agentic/train.py @@ -356,7 +356,19 @@ def main(): # Every numeric metric GRPOMetric returned: loss, clip fractions, approx_kl. **{k: v for k, v in log.items() if isinstance(v, (int, float))}, } - upload(challenge, training) + # Caught rather than allowed to propagate, because this is the last statement + # of the run and loop.sh reads its exit status: an unreachable dashboard would + # otherwise take down a loop whose checkpoint is already on disk, and -- worse + # than losing the charts -- leave the iteration without its iteration.done + # marker, so the next start would redo the iteration from weights that already + # contain it. Measured: swanlab 0.7.17 with no api key raises KeyFileError + # here, which is exactly the case this has to survive. The numbers are in + # challenge_metrics.json and train_summary.json either way. + try: + upload(challenge, training) + except Exception as e: + logger.warning(f'[train] swanlab upload failed, charts lost but the ' + f'checkpoint and the json metrics are not: {type(e).__name__}: {e}') if __name__ == '__main__': From a3ec3cc3a1cba9643d90f1cd14e938aafad108ae Mon Sep 17 00:00:00 2001 From: tastelikefeet <yuze.zyz@alibaba-inc.com> Date: Mon, 31 Aug 2026 22:21:14 +0800 Subject: [PATCH 49/60] wip --- cookbook/rsi/agentic/challenge.py | 23 ++++-- cookbook/rsi/agentic/prompts.py | 33 +++++--- src/twinkle_agentic/challenger/agentic.py | 19 +++-- src/twinkle_agentic/challenger/code.py | 68 +++++++++++++--- tests/twinkle_agentic/test_agentic_rsi.py | 80 ++++++++++++++----- .../test_multi_turn_rollout.py | 21 +++++ .../test_client_multi_turn_rollout.py | 18 +++++ 7 files changed, 212 insertions(+), 50 deletions(-) diff --git a/cookbook/rsi/agentic/challenge.py b/cookbook/rsi/agentic/challenge.py index da959a172..45fa19ed0 100644 --- a/cookbook/rsi/agentic/challenge.py +++ b/cookbook/rsi/agentic/challenge.py @@ -65,7 +65,7 @@ from twinkle.sampler import vLLMSampler from twinkle_agentic.challenger import KeywordStore, parse_check_script, parse_problem_statement from twinkle_agentic.challenger.agentic import brittle_check_reason -from twinkle_agentic.challenger.code import parse_keyword_list +from twinkle_agentic.challenger.code import split_keyword_list from twinkle_agentic.challenger.task_bank import TaskBank from twinkle_agentic.rollout import MultiTurnRollout from twinkle_agentic.tools.tool_manager import ToolManager @@ -832,7 +832,7 @@ def generate_keywords(self, category: str) -> bool: out = self.keyword_rollout([traj]) reply = self._assistant_text(out[0] if out else {}) via = 'local-fallback' - parsed = parse_keyword_list(reply) + parsed, dropped_long = split_keyword_list(reply) new = [k for k in parsed if k.lower() not in seen] for keyword in new: seen.add(keyword.lower()) @@ -840,6 +840,8 @@ def generate_keywords(self, category: str) -> bool: self.rec.keywords({'category': category, 'prompt': user, 'reply': reply, 'parsed': parsed, 'n_parsed': len(parsed), 'n_new': len(new), 'via': via, + 'dropped_long': dropped_long, + 'n_dropped_long': len(dropped_long), 'stop_reason': (out[0].get('stop_reason') if out else None), 'truncated': bool(out[0].get('truncated')) if out else None}) if len(fresh) >= want: @@ -893,11 +895,22 @@ def expand_hard_keywords(self) -> int: out = self.keyword_rollout([traj]) reply = self._assistant_text(out[0] if out else {}) via = 'local-fallback' - parsed = parse_keyword_list(reply) + parsed, dropped_long = split_keyword_list(reply) added += self.store.add(category, parsed, source='expand', parent=keyword) - self.rec.keywords({'category': category, 'parent': keyword, 'prompt': 'expand', + # The prompt goes in whole, as the refill path already does. It used to + # record the literal string 'expand', which made this file unable to + # answer the one question it gets asked -- whether a change to + # KEYWORD_EXPAND_USER was live in a given iteration -- and cost an + # afternoon to a wrong answer inferred from mtimes instead. + # + # ``dropped_long`` for the same reason one step further in: iteration 9 + # recorded n_parsed 0 on four of six expand calls whose replies were + # well-formed JSON, and nothing in this file said the phrases had been + # thrown away for length rather than never produced. + self.rec.keywords({'category': category, 'parent': keyword, 'prompt': user, 'reply': reply, 'parsed': parsed, 'n_parsed': len(parsed), - 'via': via}) + 'dropped_long': dropped_long, + 'n_dropped_long': len(dropped_long), 'via': via}) self.store.save() logger.info(f'[challenge] expanded {len(hard)} hard keyword(s) -> ' f'+{added} same-domain topics') diff --git a/cookbook/rsi/agentic/prompts.py b/cookbook/rsi/agentic/prompts.py index 5ea290239..9da9895c5 100644 --- a/cookbook/rsi/agentic/prompts.py +++ b/cookbook/rsi/agentic/prompts.py @@ -151,16 +151,29 @@ # of the scorer ("Amdahl's law speedup bound from parallel workload profile", flagged # on the noun "profile"). The wording works. # -# What it broke: 39 of the 144 keywords the model returned (27%) were silently dropped -# by ``parse_keyword_list``, which keeps only strings of 60 characters or less. All but -# one were transform -- five of its six calls came back with a median length of 65-98 -# characters, one with all eight over the cap and nothing left. Cause is in this file: -# KEYWORD_USER says "2-5 words" and this prompt only said "short strings", so the -# instruction to name a calculation rather than an activity ("reconstruct the -# allocation timeline from a heap trace") was followed at sentence length. domain came -# back at a 4-23 character median and edge_case at 32-42, both well clear. The cap is -# named here in characters because it is a silent filter in library code: a keyword -# over it does not warn, it just never exists. +# What it broke, or looked like it did: 39 of the 144 keywords iteration 9's model +# returned (27%) were silently dropped by ``parse_keyword_list``, which keeps only +# strings of 60 characters or less. All but one were transform -- five of its six calls +# came back at a median length of 65-98 characters, one with all eight over the cap and +# nothing left. +# +# The cause is NOT established, and the first version of this comment claimed it was. +# It blamed the wording here: KEYWORD_USER says "2-5 words" and this prompt only said +# "short strings", so the instruction to name a calculation rather than an activity was +# supposedly followed at sentence length. Iteration 10 ran the same prompt, before the +# length rule below existed, and dropped 0 of 168 at a median of 28-33 characters. Both +# iterations had identical prompt text -- this file was not touched between their starts +# -- so the wording cannot be what separated them. Nor does it track the parent: the two +# parents that break their own category's rules ("Performance tuning", "Diagnose kernel +# panics") both produced short keywords, while "Analyze TCP congestion patterns" +# produced 33 and "Network Latency Analysis" produced 98. What is left is API sampling +# at temperature 1.3, which is a weak explanation for a split as clean as six calls all +# above 55 against seven all below 34. +# +# The rule below is therefore a guard, not a fix: it earns its words because the failure +# mode is silent and cost 27% of a refill once, not because its cause is understood. The +# cap is named in characters because it is a filter in library code that does not warn +# -- a keyword over it never exists, and nothing in the logs says so. KEYWORD_EXPAND_USER = ( 'The keyword "{kw}" produced a very hard task. List {m} related keywords ' diff --git a/src/twinkle_agentic/challenger/agentic.py b/src/twinkle_agentic/challenger/agentic.py index 1884fdf5a..d8f33e7d1 100644 --- a/src/twinkle_agentic/challenger/agentic.py +++ b/src/twinkle_agentic/challenger/agentic.py @@ -49,7 +49,7 @@ from twinkle.data_format import SamplingParams, Trajectory, user_data_get from twinkle.utils import get_logger from .base import Challenger, Explorer, assistant_text, attach_user_data -from .code import KeywordStore, parse_keyword_list +from .code import KeywordStore, split_keyword_list logger = get_logger() @@ -524,11 +524,12 @@ class AgenticChallenger(Challenger): that is impossible from one whose statement withholds a value its check demands, and both look like a hard task worth keeping. keyword_sink: called once per keyword-generation call, with the prompt, - the raw reply and what ``parse_keyword_list`` made of it. A bank that + the raw reply and what ``split_keyword_list`` made of it. A bank that refuses to fill is invisible otherwise -- proposals fall back to the no-keyword prompt and the run carries on looking normal -- and a count of zero does not say whether the model broke the format or the parser - rejected output that was fine. + rejected output that was fine, which is why the over-length phrases are + recorded next to the kept ones rather than summed into the difference. """ def __init__( @@ -1933,7 +1934,7 @@ def _generate_keywords(self, category: str, n_want: int) -> List[str]: for user, reply in zip(users, explorer(prompts, sampling_params=self.keyword_params)): text = assistant_text(reply) - parsed = parse_keyword_list(text) + parsed, dropped_long = split_keyword_list(text) fresh = [] for kw in parsed: key = kw.lower() @@ -1954,6 +1955,12 @@ def _generate_keywords(self, category: str, n_want: int) -> List[str]: 'parsed': parsed, 'n_parsed': len(parsed), 'n_new': len(fresh), + # The two fields that make the sentence above true. Without + # them ``n_parsed: 0`` reads the same whether the reply was + # garbled, empty, or eight usable keywords written at + # sentence length -- and the third is the one that happened. + 'dropped_long': dropped_long, + 'n_dropped_long': len(dropped_long), } with self._sink_lock: self.keyword_sink(record) @@ -2026,7 +2033,7 @@ def expand_hard_keywords(self) -> int: for (cat, kw), reply in zip(reqs, explorer(prompts, sampling_params=self.keyword_params)): text = assistant_text(reply) - parsed = parse_keyword_list(text) + parsed, dropped_long = split_keyword_list(text) added += self.store.add(cat, parsed, source='expand', parent=kw) if self.keyword_sink is not None: self.keyword_sink({ @@ -2034,6 +2041,8 @@ def expand_hard_keywords(self) -> int: 'stop_reason': reply.get('stop_reason'), 'truncated': bool(reply.get('truncated')), 'parsed': parsed, 'n_parsed': len(parsed), + 'dropped_long': dropped_long, + 'n_dropped_long': len(dropped_long), }) logger.info(f'[AgenticChallenger] expanded {len(hard)} hard keyword(s) -> ' f'+{added} same-domain topics') diff --git a/src/twinkle_agentic/challenger/code.py b/src/twinkle_agentic/challenger/code.py index 29d57cdd4..5c0acda0d 100644 --- a/src/twinkle_agentic/challenger/code.py +++ b/src/twinkle_agentic/challenger/code.py @@ -234,21 +234,54 @@ def parse_challenge(text: str, require_solution: bool = True) -> Optional[Dict[s 'entry': str(obj.get('entry') or '').strip(), 'checks': checks} -def parse_keyword_list(text: str) -> List[str]: - """Extract a JSON array of short strings from a (possibly thinking) reply.""" +# A keyword is a topic to build a task around, not a task statement. Past this many +# characters the model has written the second thing, and storing it makes the next +# prompt ask for a variation on a sentence rather than on a subject. +KEYWORD_MAX_LEN = 60 + + +def split_keyword_list(text: str) -> Tuple[List[str], List[str]]: + """Extract a JSON array of short strings; return (kept, dropped for length). + + The dropped half exists because it used to be discarded inside a list + comprehension. A refill that returned eight well-formed keywords, all of them + written out as sentences, reached the caller as an empty list and was recorded + as ``n_parsed: 0`` -- the same three characters a garbled reply, a timeout and + an over-length reply all produce, so the log could not tell them apart. One + iteration lost 27% of its keywords that way and the cause was found by + re-parsing the stored replies by hand. + + The bias is the reason to count rather than only to log: length correlates with + specificity, so the filter removes "Compute the critical path delay through a + gate-level netlist with annotated cell delays" and keeps whatever was vague + enough to be short. That is the opposite of what the bank is for. + """ body = text idx = body.rfind('</think>') if idx >= 0: body = body[idx + len('</think>'):] start, end = body.find('['), body.rfind(']') if start < 0 or end <= start: - return [] + return [], [] try: arr = json.loads(body[start:end + 1]) except (ValueError, TypeError): - return [] - return [x.strip() for x in arr - if isinstance(x, str) and x.strip() and len(x.strip()) <= 60] + return [], [] + kept: List[str] = [] + dropped: List[str] = [] + for x in arr: + if not isinstance(x, str): + continue + s = x.strip() + if not s: + continue + (kept if len(s) <= KEYWORD_MAX_LEN else dropped).append(s) + return kept, dropped + + +def parse_keyword_list(text: str) -> List[str]: + """The kept half of :func:`split_keyword_list`, for callers with nothing to record.""" + return split_keyword_list(text)[0] def load_seeds(path: str) -> List[Dict[str, str]]: @@ -648,12 +681,21 @@ def _generate_keywords(self, category: str, n_want: int) -> List[str]: } for i in range(n_calls)] seen = {t.strip().lower() for t in known} out: List[str] = [] + n_long = 0 for reply in self.explore(prompts, sampling_params=self.keyword_params): - for kw in parse_keyword_list(assistant_text(reply)): + kept, dropped = split_keyword_list(assistant_text(reply)) + n_long += len(dropped) + for kw in kept: key = kw.lower() if key not in seen: seen.add(key) out.append(kw) + if n_long: + # This path has no dump to write to, so the count has to be said out + # loud or the refill looks like the model simply produced less. + logger.warning(f'[CodeChallenger] dropped {n_long} keyword(s) over ' + f'{KEYWORD_MAX_LEN} chars while refilling; the prompt is ' + f'asking for task statements rather than topics') self.rng.shuffle(out) return out[:n_want] @@ -817,10 +859,16 @@ def expand_hard_keywords(self) -> int: ], } for i, (_c, kw) in enumerate(reqs)] added = 0 + n_long = 0 for (cat, kw), reply in zip(reqs, self.explore(prompts, - sampling_params=self.keyword_params)): - added += self.store.add(cat, parse_keyword_list(assistant_text(reply)), - source='expand', parent=kw) + sampling_params=self.keyword_params)): + kept, dropped = split_keyword_list(assistant_text(reply)) + n_long += len(dropped) + added += self.store.add(cat, kept, source='expand', parent=kw) + if n_long: + logger.warning(f'[CodeChallenger] dropped {n_long} expanded keyword(s) over ' + f'{KEYWORD_MAX_LEN} chars; expansion follows the parent, so a ' + f'wordy parent produces wordy children') logger.info(f'[CodeChallenger] expanded {len(hard)} hard keyword(s) -> ' f'+{added} same-domain topics') return added diff --git a/tests/twinkle_agentic/test_agentic_rsi.py b/tests/twinkle_agentic/test_agentic_rsi.py index 96ffa6b69..e21652dfe 100644 --- a/tests/twinkle_agentic/test_agentic_rsi.py +++ b/tests/twinkle_agentic/test_agentic_rsi.py @@ -1178,6 +1178,41 @@ def test_the_shipped_prompt_asks_for_what_the_parser_reads(self): # rewording back to one-per-line fails here rather than in a night's run. self.assertEqual(parse_keyword_list('csv deduplication\nlog rotation'), []) + def test_over_length_keywords_are_reported_and_not_merely_gone(self): + """A dropped phrase has to be distinguishable from one never produced. + + The length filter used to live inside a list comprehension, so a reply of + eight well-formed keywords written at sentence length reached the caller as + an empty list and was recorded as ``n_parsed: 0`` -- identical to a garbled + reply and to a timeout. Iteration 9 lost 27% of a refill that way, on four + of six expand calls, and the cause was found by re-parsing stored replies. + + The direction matters as much as the count: length tracks specificity, so + what the filter removes is the half of the output the bank most wants. + """ + from twinkle_agentic.challenger.code import (KEYWORD_MAX_LEN, + split_keyword_list) + from prompts import KEYWORD_EXPAND_USER + + # Verbatim from iteration 9, one of the eight a single expand call lost. + wordy = ('Compute the critical path delay through a gate-level netlist ' + 'with annotated cell delays') + self.assertGreater(len(wordy), KEYWORD_MAX_LEN, 'fixture must exceed the cap') + kept, dropped = split_keyword_list(f'["crc32 table generation", "{wordy}"]') + self.assertEqual(kept, ['crc32 table generation']) + self.assertEqual(dropped, [wordy], 'the dropped phrase must be recoverable') + + # The three cases a reader has to be able to tell apart. Only the first + # carries anything in the dropped half, which is what makes the other two + # diagnosable as format failures rather than length failures. + self.assertEqual(split_keyword_list(f'["{wordy}"]'), ([], [wordy])) + self.assertEqual(split_keyword_list('sorry, I cannot'), ([], [])) + self.assertEqual(split_keyword_list('["unterminated'), ([], [])) + + # And the prompt has to name the same ceiling the parser enforces, or the + # model is being marked down against a rule it was never told. + self.assertIn(str(KEYWORD_MAX_LEN), KEYWORD_EXPAND_USER) + def test_a_json_reply_fills_the_bank_and_reaches_the_proposal(self): from twinkle.data_format import user_data_get ch = self._challenger('["csv deduplication", "log rotation"]') @@ -1356,37 +1391,42 @@ class ProposeTrajIndexTest(unittest.TestCase): GRPO advantage out of them, and skips a dump without it as a 'pre-grouping run'. While the copy dropped both, SIDES=both trained 384 solver and 0 proposer trajectories, and said so only in a line nobody read. + + Written against ``ProposeTrajWriter``, which no longer exists -- the writer is + ``Recorder`` now and the reward field is ``reward``, not ``challenger_reward``. + So this test spent an unknown number of commits failing at import, which is to + say the invariant above went unguarded for exactly as long as it looked + guarded. Kept pointed at ``Recorder.trajectory`` with the keywords the + production call passes, so a rename breaks it again rather than retiring it. + + What it does not cover: that the *caller* passes group_id at all. That was the + other half of the original bug and it needs the collection loop, not this. """ def test_group_id_and_reward_survive_the_copy(self): - from challenge import ProposeTrajWriter + from challenge import Recorder out = tempfile.mkdtemp(prefix='proposetraj_test_') try: - writer = ProposeTrajWriter(out) - writer.write({ - 'outcome': 'kept', - 'group_id': 0, - 'challenger_reward': 0.75, - 'n_pass': 4, - 'n_rollouts': 8, - 'pass_rate': 0.5, - 'keywords': [['transform', 'parse a binary log']], - 'seeded': False, - 'rounds': [{'stage': 'episode', 'messages': [], - 'input_ids': [1, 2], 'labels': [-100, 2], - 'logprobs': []}], - }) - writer.close() - with open(os.path.join(out, 'index.jsonl'), encoding='utf-8') as f: - rec = json.loads(f.readline()) + rec = Recorder(out) + # The field names and shape of challenge.py's own propose-side call. + rec.trajectory( + {'input_ids': [1, 2], 'labels': [-100, 2], 'logprobs': None, + 'messages': []}, + side='propose', group_id=0, proposal_idx=3, reward=0.75, n_pass=4, + novelty=1.0, outcome='kept', + keywords=[['transform', 'parse a binary log']], selected=True) + rec.close() + with open(os.path.join(out, 'trajs', 'index.jsonl'), encoding='utf-8') as f: + record = json.loads(f.readline()) finally: shutil.rmtree(out, ignore_errors=True) # Group 0 is a real group, so this also pins that the copy reads the key # rather than testing it for truth. - self.assertEqual(rec['group_id'], 0) - self.assertEqual(rec['challenger_reward'], 0.75) + self.assertEqual(record['group_id'], 0) + self.assertEqual(record['reward'], 0.75) + self.assertEqual(record['side'], 'propose') class TaskCarriesGroupIdTest(unittest.TestCase): diff --git a/tests/twinkle_agentic/test_multi_turn_rollout.py b/tests/twinkle_agentic/test_multi_turn_rollout.py index 21c1bc8f5..56c02b2a8 100644 --- a/tests/twinkle_agentic/test_multi_turn_rollout.py +++ b/tests/twinkle_agentic/test_multi_turn_rollout.py @@ -156,6 +156,27 @@ def parse_tool_call(self, decoded: str) -> list[dict[str, Any]]: }) return results + def tool_call_errors(self, decoded: str) -> list[str]: + """Why ``parse_tool_call`` returned fewer calls than the markup asked for. + + Mirrors that method's two ``continue`` branches instead of returning an + empty list. A stub that always reported no errors would keep these tests + green while silently retiring the branch in MultiTurnRollout that hands a + parse failure back to the model -- the retry would become unreachable and + no test would notice, which is the failure mode a stub is supposed to + prevent rather than cause. + """ + errors: list[str] = [] + for m in re.findall(r'<tool_call>\s*([\s\S]*?)\s*</tool_call>', decoded or ''): + try: + d = json.loads(m) + except json.JSONDecodeError as exc: + errors.append(f'tool_call is not valid JSON: {exc.msg}') + continue + if not (d.get('name') or d.get('tool_name')): + errors.append('tool_call has no "name" field') + return errors + def clean_tool_call(self, decoded: str) -> str: """Strip the call blocks, as the real template does before storing.""" return re.sub(r'<tool_call>[\s\S]*?</tool_call>', '', decoded or '') diff --git a/tests/twinkle_client/test_client_multi_turn_rollout.py b/tests/twinkle_client/test_client_multi_turn_rollout.py index 6b9474a1e..f426747f8 100644 --- a/tests/twinkle_client/test_client_multi_turn_rollout.py +++ b/tests/twinkle_client/test_client_multi_turn_rollout.py @@ -162,6 +162,24 @@ def parse_tool_call(self, decoded: str) -> List[Dict[str, Any]]: }) return results + def tool_call_errors(self, decoded: str) -> List[str]: + """Why ``parse_tool_call`` returned fewer calls than the markup asked for. + + Mirrors that method's two ``continue`` branches. Returning an empty list + would pass just as well and would quietly make the parse-failure retry in + MultiTurnRollout unreachable from these tests. + """ + errors: List[str] = [] + for m in re.findall(r'<tool_call>\s*([\s\S]*?)\s*</tool_call>', decoded or ''): + try: + d = json.loads(m) + except json.JSONDecodeError as exc: + errors.append(f'tool_call is not valid JSON: {exc.msg}') + continue + if not (d.get('name') or d.get('tool_name')): + errors.append('tool_call has no "name" field') + return errors + def concat_input_feature(self, pif: Dict[str, Any], new_tokens: List[int]) -> Dict[str, Any]: result = copy.deepcopy(pif) prompt_ids = list(result['input_ids']) From b50985cf3b5bd076616830712c0be2dde7997a46 Mon Sep 17 00:00:00 2001 From: tastelikefeet <yuze.zyz@alibaba-inc.com> Date: Tue, 1 Sep 2026 00:43:34 +0800 Subject: [PATCH 50/60] fix --- cookbook/rsi/agentic/loop.sh | 11 +++++++++++ src/twinkle/model/optimizer_group.py | 16 ++++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/cookbook/rsi/agentic/loop.sh b/cookbook/rsi/agentic/loop.sh index e78dedcb3..667ee78ae 100644 --- a/cookbook/rsi/agentic/loop.sh +++ b/cookbook/rsi/agentic/loop.sh @@ -214,12 +214,23 @@ while [ "$ITERATIONS" -eq 0 ] || [ "$i" -lt $((START + ITERATIONS)) ]; do 2>&1 | tee "$OUT/challenge.log" echo "=== iteration $i: train on $OUT -> $CKPT_DIR" + # expandable_segments on the training stage only, and not on collect: one padded + # trajectory per micro batch means every micro batch is a new shape (119 distinct + # lengths in 128 trajectories, 7k-19k tokens), and the caching allocator cannot + # reuse a block across sizes, so it grew to 87.8 GiB reserved against 29.0 GiB + # live on a 97.4 GiB card. That is what starved NCCL of the few hundred MiB it + # needs to connect the metric gather's communicator, which hung iteration 2 for + # 54 minutes. Expandable segments let one virtual range serve every shape, so + # reserved tracks the real peak instead of the sum of shapes. Left off for + # collect because that stage is vLLM, which profiles its own KV cache against + # allocator behaviour and has nothing to do with this failure. RSI_RUN_DIR="$OUT" \ RSI_SAVE_DIR="$CKPT_DIR" \ RSI_SAVE_NAME="model" \ RSI_SIDES="$SIDES" \ RSI_TAG="$TAG" \ RSI_ITER="$i" \ + PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" \ CUDA_VISIBLE_DEVICES="$DEVICES" python cookbook/rsi/agentic/train.py \ --model_id "$MODEL" \ --model_gpus "$GPUS" \ diff --git a/src/twinkle/model/optimizer_group.py b/src/twinkle/model/optimizer_group.py index 384dffe42..f5177d672 100644 --- a/src/twinkle/model/optimizer_group.py +++ b/src/twinkle/model/optimizer_group.py @@ -1,4 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +import torch from dataclasses import dataclass, field from torch.optim import Optimizer from torch.optim.lr_scheduler import LRScheduler @@ -83,6 +84,21 @@ def calculate_metrics(self, is_training): """Calculate and return metrics.""" self.accumulate_metrics(is_training) status = self.train_status if is_training else self.eval_status + # The metrics below gather over the DP group, and that gather is the first + # use of its NCCL communicator: NCCL connects it with its own cudaMalloc, + # which draws on device memory torch's caching allocator has *not* taken, so + # it fails on whichever rank has the least left -- and it fails where nothing + # reports it. Measured on 8xH20 with one padded trajectory per micro batch: + # 16 mini batches of 7k-19k tokens left the allocator holding 87.8 GiB + # reserved against 29.0 GiB live, one rank down to 164 MiB free, and that + # rank raised inside all_gather_object while the other seven waited in it + # forever -- 54 minutes, no log line, GPUs at 0% with their memory held, + # because calculate_metric is collected 'last_pp_first' so the driver never + # fetches the failing rank's exception. Releasing the cache first puts every + # rank above 54 GiB free and the same step completes in 5 ms. The cost is one + # re-allocation per optimizer step, which is once per iteration here. + if status.metrics and torch.cuda.is_available() and torch.cuda.is_initialized(): + torch.cuda.empty_cache() results = {} for metric in status.metrics: results.update(metric.calculate()) From a472cbe2b39f8f9d805ccb89e3b2cc557362682c Mon Sep 17 00:00:00 2001 From: tastelikefeet <yuze.zyz@alibaba-inc.com> Date: Tue, 1 Sep 2026 08:44:23 +0800 Subject: [PATCH 51/60] fix --- .../rsi/agentic => .temp/retired_rsi}/loop.sh | 0 .../retired_rsi}/rsi_agent_shortprompt.yaml | 0 .../retired_rsi}/split_tasks.py | 0 cookbook/rsi/agentic/README.md | 43 ++- cookbook/rsi/agentic/challenge.py | 179 ++++++----- cookbook/rsi/agentic/rsi.py | 233 ++++++++++++++ cookbook/rsi/agentic/train.py | 294 +++++++----------- 7 files changed, 478 insertions(+), 271 deletions(-) rename {cookbook/rsi/agentic => .temp/retired_rsi}/loop.sh (100%) rename {cookbook/rsi/agentic => .temp/retired_rsi}/rsi_agent_shortprompt.yaml (100%) rename {cookbook/rsi/agentic => .temp/retired_rsi}/split_tasks.py (100%) create mode 100644 cookbook/rsi/agentic/rsi.py diff --git a/cookbook/rsi/agentic/loop.sh b/.temp/retired_rsi/loop.sh similarity index 100% rename from cookbook/rsi/agentic/loop.sh rename to .temp/retired_rsi/loop.sh diff --git a/cookbook/rsi/agentic/rsi_agent_shortprompt.yaml b/.temp/retired_rsi/rsi_agent_shortprompt.yaml similarity index 100% rename from cookbook/rsi/agentic/rsi_agent_shortprompt.yaml rename to .temp/retired_rsi/rsi_agent_shortprompt.yaml diff --git a/cookbook/rsi/agentic/split_tasks.py b/.temp/retired_rsi/split_tasks.py similarity index 100% rename from cookbook/rsi/agentic/split_tasks.py rename to .temp/retired_rsi/split_tasks.py diff --git a/cookbook/rsi/agentic/README.md b/cookbook/rsi/agentic/README.md index 92fdf77fb..6d0fdb281 100644 --- a/cookbook/rsi/agentic/README.md +++ b/cookbook/rsi/agentic/README.md @@ -103,20 +103,31 @@ Solving side: 1.0 if the check exits 0, else 0.0. ## Files ``` +rsi.py the loop: one resident process, collect -> step -> sync, forever challenge.py collect: the queues, the three job bodies, the group decision -train.py one GRPO step over what was collected, then overwrite the ckpt +train.py one GRPO step over what was collected, as a library rsi.py calls sandbox.py the sandbox as a resource: clear, snapshot, run a script prompts.py every string sent to a model -loop.sh collect -> train -> collect from the new weights, until killed episode.py how an episode is built and scored, shared with eval.py remote_tool_env.py the transport to one microVM, paired with sandbox_server/ sandbox_server/ the image and the in-sandbox tool server it talks to eval.py held-out pass rate on tasks the trainer never saw -split_tasks.py split a collection's tasks into a train and an eval half rsi_agent.yaml the ms-agent config both sides' openings are shaped by ``` -Output under `--out-dir`: +The trainer and the sampler are two disjoint device groups in one Ray job -- 2 and +6 GPUs by default -- and both stay resident for the whole run. After each step the +new weights go to the live vLLM engines over NCCL (`CheckpointEngineManager`), so +nothing is restarted and nothing round-trips through the filesystem. `loop.sh`, +which used to run a fresh `challenge.py` and `train.py` per iteration, is retired +under `.temp/retired_rsi/`: it spent 11 minutes per iteration on startup, re-read +the checkpoint once per GPU, and -- the reason it had to go -- passed the model +between iterations as bf16 weights only, which threw away the fp32 master weights +and the Adam moments every time. Measured on v3 after 12 iterations at lr 1e-6: +98.54% of the 4.02 B weights were still bit-identical to the base model, and the +largest change anywhere was 2.289e-05, one bf16 step at that magnitude. + +Output under `<root>/<tag>/iter<n>`: ``` trajs/*.npz input_ids / labels / logprobs @@ -169,15 +180,23 @@ so a task's `n_pass` here and its `pass@k` there are measured against one openin export E2B_API_KEY=... # sandbox host export SANDBOX_API_URL=http://... # sandbox host address, with port export LLM_BACKUP_API_KEY=... # dashscope -ITERATIONS=1 bash cookbook/rsi/agentic/loop.sh +python cookbook/rsi/agentic/rsi.py --tag v4 --iterations 1 ``` +`--iterations 0`, the default, runs until killed. Restarting the same `--tag` +continues it: iterations are counted by the `iteration.done` marker, which is +written after the checkpoint, and the loop picks up from +`<root>/<tag>/ckpt/model`. The optimizer is state that only exists in memory, so +it is checkpointed every `--save-optimizer-every` iterations (5); a crash between +two of those resumes with the weights and with Adam at zero moments. + Charts land in swanlab project `twinkle-rsi-agentic`, one experiment named after -`TAG`, one step per iteration. `train.py` uploads after saving the checkpoint, so a -swanlab failure costs the charts and not the weights — the numbers are still in -`challenge_metrics.json` and `train_summary.json` either way. Resume is by -`id=TAG`: a second run under the same tag appends to that curve, a new tag starts a -new one. `RSI_SWANLAB_MODE=disabled` turns it off, `RSI_SWANLAB_PROJECT` moves it. +`--tag`, one step per iteration. The upload happens after the checkpoint is saved +and its failure is caught, so an unreachable dashboard costs the charts and not the +weights — the numbers are still in `challenge_metrics.json` and +`train_summary.json` either way. Resume is by `id=tag`: a second run under the same +tag appends to that curve, a new tag starts a new one. `--swanlab-mode disabled` +turns it off, `--swanlab-project` moves it. Verified on this machine at swanlab 0.9.2: three separate processes with the same tag at steps 1, 2, 3 landed on one run (the second and third print `disabled in @@ -197,8 +216,8 @@ re-measured under this scheduler. | truncated solver attempt counts as a failure, denominator fixed at 8 | — | decided for this pipeline | | a build cut off at `--propose-max-tokens` writes no check and no statement | — | restored from the old pipeline, which skipped both stages after a length cut | | rubric failure after 3 tries drops the whole group | — | decided for this pipeline | -| `--max-build-files` | 4 | inherited: `loop.sh` has passed this since it was added. It is text in the system prompt. | -| `--api-thinking-budget` | 4096 | inherited from the old `loop.sh` | +| `--max-build-files` | 4 | inherited: every run since it was added has passed this. It is text in the system prompt. | +| `--api-thinking-budget` | 4096 | inherited from the retired `loop.sh` | | `--propose-max-tokens` / `--max-turns` / `--stop-after-stuck-turns` | 8192 / 24 / 2 | inherited | | `--one-call-per-reply` | on | inherited | | `--check-retries` / `--check-max-tokens` | 1 / 8192 | inherited | diff --git a/cookbook/rsi/agentic/challenge.py b/cookbook/rsi/agentic/challenge.py index 45fa19ed0..3481a074a 100644 --- a/cookbook/rsi/agentic/challenge.py +++ b/cookbook/rsi/agentic/challenge.py @@ -31,17 +31,20 @@ solver side trains on, so a kept group contributes 8 proposing and 8 solving trajectories, and eight kept groups are the 64 + 64 one training step reads. -Output (all under ``--out-dir``): +Output (per iteration, under ``<root>/<tag>/iter<n>``): trajs/*.npz input_ids / labels / logprobs per trajectory trajs/index.jsonl one line per trajectory: side, group, reward, full text groups.jsonl one line per decided group: why kept or dropped tasks.jsonl the statements and check scripts that were delivered - keywords.jsonl the keyword bank, carried between iterations -Run it as a Ray job (sampler only, no trainer):: +and two files that belong to the loop rather than to an iteration, at +``<root>/<tag>``: keywords.jsonl, the keyword bank, and task_bank.jsonl, the +statements novelty is judged against. - python cookbook/rsi/agentic/challenge.py --keep-groups 8 +This is the collecting half as a library. rsi.py owns the process, the sampler and +the sandbox pool, and calls in here once per iteration; the argument parser lives +here because collection is what almost all of the arguments are about. """ import argparse import collections @@ -59,8 +62,7 @@ from typing import Any, Dict, List, Optional, Tuple import numpy as np -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_logger +from twinkle import DeviceMesh, get_logger from twinkle.data_format import SamplingParams from twinkle.sampler import vLLMSampler from twinkle_agentic.challenger import KeywordStore, parse_check_script, parse_problem_statement @@ -72,7 +74,7 @@ sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) import prompts as P # noqa: E402 -from sandbox import close_pool, open_pool, solver_harness # noqa: E402 +from sandbox import open_pool, solver_harness # noqa: E402 logger = get_logger() @@ -155,9 +157,15 @@ def parse_args(): 'not. 0 leaves the run governed by --keep-groups alone.') # Local model (the trainable half). - p.add_argument('--model-id', default='ms://Qwen/Qwen3-4B') + p.add_argument('--model-id', default='ms://Qwen/Qwen3-4B', + help='where the loop starts from. Once an iteration has ' + 'finished, its own checkpoint is used instead.') p.add_argument('--template', default='Template') - p.add_argument('--sampler-gpus', type=int, default=4) + p.add_argument('--sampler-gpus', type=int, default=6) + p.add_argument('--model-gpus', type=int, default=2, + help='trainer GPUs. Disjoint from the sampler\'s, so the two ' + 'halves stay resident side by side; --sampler-gpus + this ' + 'is the size of the Ray job.') p.add_argument('--max-model-len', type=int, default=40960) p.add_argument('--gpu-memory-utilization', type=float, default=0.8) @@ -225,9 +233,10 @@ def parse_args(): p.add_argument('--keyword-max-tokens', type=int, default=4096) # Novelty. - p.add_argument('--task-bank', default='', - help='jsonl of statements from earlier iterations. Empty turns ' - 'novelty off, and the reward is the pass-rate gaussian alone.') + p.add_argument('--task-bank', default=None, + help='jsonl of statements from earlier iterations, defaulting to ' + "<root>/<tag>/task_bank.jsonl. '' turns novelty off, and the " + 'reward is the pass-rate gaussian alone.') p.add_argument('--task-bank-refs', type=int, default=5, help='stored statements shown to the judge, on top of the group ' "'s own siblings, which are always shown.") @@ -243,8 +252,14 @@ def parse_args(): p.add_argument('--sandbox-slots', type=int, default=32, help='microVMs, i.e. how many jobs run at once. One job owns one ' 'slot from the workspace clear to its last check.') - p.add_argument('--sandbox-template', default=os.environ.get('AENV_TEMPLATE', '')) - p.add_argument('--sandbox-api-url', default=os.environ.get('AENV_API_URL', '')) + # AENV_* is what the sandbox client reads; E2B_API_KEY / SANDBOX_API_URL is + # what the host hands out and what the README tells you to export. The + # translation used to live in loop.sh and has to live somewhere. + p.add_argument('--sandbox-template', + default=os.environ.get('AENV_TEMPLATE') or 'twinkle-rsi-msagent') + p.add_argument('--sandbox-api-url', + default=(os.environ.get('AENV_API_URL') + or os.environ.get('SANDBOX_API_URL', ''))) p.add_argument('--sandbox-timeout', type=int, default=900) p.add_argument('--agent-config', default='cookbook/rsi/agentic/rsi_agent.yaml') p.add_argument('--workspace', default='/workspace') @@ -252,40 +267,95 @@ def parse_args(): p.add_argument('--snapshot-per-file', type=int, default=600) p.add_argument('--snapshot-budget', type=int, default=6000) + # Training, one step per iteration. See train.py. + p.add_argument('--lr', type=float, default=1e-6) + p.add_argument('--sides', default='both', choices=('both', 'propose', 'solve')) + p.add_argument('--micro-batch-size', type=int, default=1, + help='trajectories per micro batch. One, because padding_free is ' + 'off: a micro batch is padded to its longest member, so ' + 'pairing a short solver attempt with a long build episode ' + 'pays for the long one twice.') + p.add_argument('--mini-batch-size', type=int, default=0, + help='0 means --model-gpus x --micro-batch-size, which is the ' + "floor: forward_backward is dispatch='slice_dp', so a mini " + 'batch has to give every rank at least one micro batch.') + p.add_argument('--max-train-len', type=int, default=32768, + help='a trajectory longer than this is not trained on. Below ' + '--max-model-len, so collection can produce some.') + + # The loop. + p.add_argument('--root', default='output/rsi_agentic') + p.add_argument('--tag', default='', + help='names the run: everything lives under <root>/<tag>, and ' + 'restarting a tag continues it rather than redoing it.') + p.add_argument('--iterations', type=int, default=0, + help='0 runs until killed.') + p.add_argument('--ckpt-dir', default='', + help='defaults to <root>/<tag>/ckpt. Worth pointing at a faster ' + 'filesystem than the repo: the first load reads it once per ' + 'GPU, and measured with dd at 1.5 GB this host has one disk ' + 'at 223 MB/s and another at 1074 MB/s.') + p.add_argument('--save-optimizer-every', type=int, default=5, + help='iterations between checkpoints that include the optimizer. ' + 'Weights are saved every iteration either way; this is what ' + 'a resume needs to keep the Adam moments, and it is ~48 GB ' + 'for a 4B model against 7.6 GB for the weights alone.') + p.add_argument('--swanlab-project', default='twinkle-rsi-agentic') + p.add_argument('--swanlab-mode', default='online', + help="'disabled' keeps a run off the dashboard entirely.") + # Output. - p.add_argument('--out-dir', default='output/rsi_agentic') - p.add_argument('--keyword-db', default='', - help='defaults to <out-dir>/keywords.jsonl') p.add_argument('--random-seed', type=int, default=0) args = p.parse_args() if not args.api_model or not args.api_base: - raise SystemExit('[challenge] --api-model and --api-base are required ' + raise SystemExit('[rsi] --api-model and --api-base are required ' '(or LLM_BACKUP_MODEL / LLM_BACKUP_BASE_URL)') + if not args.tag: + raise SystemExit('[rsi] --tag is required: it decides which run these ' + 'iterations belong to and which checkpoint they overwrite') if args.solver_rollouts < 2: - raise SystemExit('[challenge] --solver-rollouts must be >= 2: it is both the ' - "solver side's GRPO group size and the denominator n_pass is " - 'judged against') + raise SystemExit('[rsi] --solver-rollouts must be >= 2: it is both the ' + "solver side's GRPO group size and the denominator n_pass " + 'is judged against') if args.group_size < 2: - raise SystemExit('[challenge] --group-size must be >= 2: a group of one has ' + raise SystemExit('[rsi] --group-size must be >= 2: a group of one has ' 'no mean to subtract, so every advantage is zero') - args.keyword_db = args.keyword_db or os.path.join(args.out_dir, 'keywords.jsonl') + # Checked here rather than where the pool is opened, which is after the model + # and the sampler are up: that is six minutes of startup to find out that a + # host address is missing. + if not args.sandbox_api_url: + raise SystemExit('[rsi] --sandbox-api-url is required (or SANDBOX_API_URL / ' + 'AENV_API_URL)') + if not os.environ.get('E2B_API_KEY') and not os.environ.get('AENV_API_KEY'): + raise SystemExit('[rsi] E2B_API_KEY is required: the sandbox client reads ' + 'it from the environment') + os.environ.setdefault('AENV_API_URL', args.sandbox_api_url) + os.environ.setdefault('AENV_TEMPLATE', args.sandbox_template) + os.environ.setdefault('AENV_API_KEY', os.environ.get('E2B_API_KEY', '')) + # The bank and the keyword store belong to the loop, not to an iteration: + # comparing iteration k+1's proposals against what k produced is the point of + # them. ``--task-bank ''`` turns novelty off and leaves the pass-rate gaussian + # alone. out_dir is set per iteration by rsi.py. + root = os.path.join(args.root, args.tag) + args.keyword_db = os.path.join(root, 'keywords.jsonl') + if args.task_bank is None: + args.task_bank = os.path.join(root, 'task_bank.jsonl') + args.out_dir = root return args # ── Resources ────────────────────────────────────────────────────────────── -def initialize_device(args) -> Tuple[Any, Any]: - """Bring up Ray and the local vLLM sampler; returns (sampler, template). +def build_sampler(args) -> Tuple[Any, Any]: + """The resident vLLM sampler; returns (sampler, template). - The template is built here as well as inside the sampler because the rollout - encodes with it locally: one object, so the token ids the sampler continues - from are the ids the trajectory was encoded with. + Ray and the device groups are already up -- rsi.py owns them, because the + trainer needs a group of its own on the same job. The template is built here as + well as inside the sampler because the rollout encodes with it locally: one + object, so the token ids the sampler continues from are the ids the trajectory + was encoded with. """ - twinkle.initialize( - mode='ray', nproc_per_node=args.sampler_gpus, lazy_collect=False, - groups=[DeviceGroup(name='sampler', ranks=list(range(args.sampler_gpus)), - device_type='GPU')]) sampler = vLLMSampler( model_id=args.model_id, engine_args={'gpu_memory_utilization': args.gpu_memory_utilization, @@ -301,7 +371,7 @@ def initialize_device(args) -> Tuple[Any, Any]: args.model_id, max_length=args.max_model_len, enable_thinking=True) if not getattr(type(sampler).sample, '_enable_continous_work', False): raise SystemExit( - '[challenge] this sampler does not route requests one at a time ' + '[rsi] this sampler does not route requests one at a time ' '(sample lacks enable_continous_work), so a batch of one would be ' 'padded to the worker count and most of every generation thrown away. ' 'The whole design here is one trajectory per request.') @@ -1527,46 +1597,3 @@ def rate(num: float, den: float) -> float: sorted(collections.Counter(novelty).items())}, }, } - - -def main(): - args = parse_args() - os.makedirs(args.out_dir, exist_ok=True) - recorder = Recorder(args.out_dir) - sampler, template = initialize_device(args) - slots = initialize_sandbox(args) - run = Run(args, sampler, template, slots, recorder) - started = time.time() - try: - run.run() - # After the loop, not during: what it adds is for the next iteration, and - # doing it here means a crash in collection does not also lose the bank. - if args.keyword_expand: - run.expand_hard_keywords() - finally: - rebuilds = close_pool(slots) - recorder.close() - if run.bank is not None: - logger.info(f'[challenge] task bank: {run.bank.stats()}') - run.store.save() - if rebuilds: - logger.warning(f'[challenge] sandboxes were rebuilt {rebuilds} time(s); ' - f'the jobs in flight at those moments were lost') - # Written after recorder.close(), so groups.jsonl is complete and flushed - # before it is read back. In the finally block because a run that crashed - # is the one whose numbers are most worth having. - metrics = collect_metrics(args.out_dir, run.counts, run.n_launched, - args.solver_rollouts, time.time() - started) - with open(os.path.join(args.out_dir, 'challenge_metrics.json'), 'w', - encoding='utf-8') as f: - json.dump(metrics, f, indent=2, ensure_ascii=False, default=str) - logger.info(f'[challenge] {len(run.kept)}/{run.n_launched} groups kept in ' - f'{time.time() - started:.0f}s, counts: ' - f'{dict(sorted(run.counts.items()))}') - logger.info(f'[challenge] metrics -> ' - f'{os.path.join(args.out_dir, "challenge_metrics.json")}: ' - f'{metrics["scalars"]}') - - -if __name__ == '__main__': - main() diff --git a/cookbook/rsi/agentic/rsi.py b/cookbook/rsi/agentic/rsi.py new file mode 100644 index 000000000..b2a46e411 --- /dev/null +++ b/cookbook/rsi/agentic/rsi.py @@ -0,0 +1,233 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""The self-play loop as one resident process: collect, step, hand the new weights +to the live sampler, repeat. + +This replaces loop.sh, which ran challenge.py and train.py as a fresh pair of +processes per iteration. What that cost, in the order the numbers matter: + +Accumulation. loop.sh's only channel between iterations was a bf16 HF checkpoint, +so the trainer's fp32 master weights and its Adam moments were thrown away and +rebuilt every iteration. Measured on v3 after 12 iterations at lr 1e-6: 98.54% of +the 4.02 B weights were still bit-identical to the base model, and the largest +change anywhere was 2.289e-05 -- one bf16 step at that magnitude, and the same +value in eight different tensors, which is quantisation showing through rather +than learning. A step displaces an element by about 2e-6, bf16 near |w|=1e-2 +cannot record less than ~4e-5, so each iteration's update was rounded away instead +of added to the last one. Here the optimizer never leaves memory and 12 steps are +12 steps. + +Startup. 5.5 minutes of vLLM and 5.4 minutes of Megatron per iteration, about 29% +of a 38-minute iteration, plus 7.6 GB written and ~50 GB read as every sampler +worker reloaded the checkpoint. + +Memory. The trainer and the sampler own disjoint GPUs, so neither can starve the +other. Time-sharing all eight cards instead -- vLLM asleep during the step -- would +put 29 GB of resident trainer against ~65 GB of woken vLLM inside 97 GB, on the +machine where a metric gather has already died for want of 200 MB. + +The split costs idle capacity: the trainer's cards wait out the ~35 minutes of +collection and the sampler's wait out the ~6 minutes of the step. Collection is +bound by sandbox round trips and API latency rather than generation -- 128 +trajectories of at most 1.16 M tokens in 30 minutes is under 700 tok/s across all +engines, far under what a 4B model does on one H20 -- so buying wall-clock with +sampler width is the cheap direction and buying it with trainer width is not. + + python cookbook/rsi/agentic/rsi.py --tag v4 + +Resuming is by the same marker loop.sh used: iter<n>/iteration.done, written last. +A resident optimizer is state that only exists in memory, so it is checkpointed +every --save-optimizer-every iterations; a crash between two of those resumes with +the weights but with Adam starting from zero moments, which is the old behaviour +for exactly one step rather than for every step. +""" +import json +import os +import sys +import time +from typing import Any, Dict, Optional + +import twinkle +from twinkle import DeviceGroup, get_device_placement, get_logger +from twinkle.checkpoint_engine import CheckpointEngineManager + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import challenge as C # noqa: E402 +import train as T # noqa: E402 +from sandbox import close_pool # noqa: E402 + +logger = get_logger() + + +def next_iteration(root: str) -> int: + """The first iteration with no ``iteration.done``. + + Counted from the marker rather than from what is on disk: a directory exists + as soon as collection starts writing into it, and a train_summary.json is + there after a step whose checkpoint may not have been saved. + """ + i = 1 + while os.path.exists(os.path.join(root, f'iter{i}', 'iteration.done')): + i += 1 + return i + + +def collect_once(args, sampler, template, slots, out_dir: str) -> Dict[str, Any]: + """One collection pass into ``out_dir``; returns its metrics. + + The body of what challenge.py's main() did, minus the resources: the sampler, + the template and the sandbox pool are owned by the caller and outlive this. + """ + os.makedirs(out_dir, exist_ok=True) + args.out_dir = out_dir + recorder = C.Recorder(out_dir) + run = C.Run(args, sampler, template, slots, recorder) + started = time.time() + try: + run.run() + # After the loop, not during: what it adds is for the next iteration, and + # doing it here means a crash in collection does not also lose the bank. + if args.keyword_expand: + run.expand_hard_keywords() + finally: + recorder.close() + run.store.save() + # A Run per iteration means a thread pool per iteration. close_pool cannot + # do this because the sandbox pool is the one thing that is not per-Run. + run.api_pool.shutdown(wait=False) + if run.bank is not None: + logger.info(f'[rsi] task bank: {run.bank.stats()}') + # In the finally block because a run that crashed is the one whose numbers + # are most worth having, and after recorder.close() so groups.jsonl is + # flushed before collect_metrics reads it back. + metrics = C.collect_metrics(out_dir, run.counts, run.n_launched, + args.solver_rollouts, time.time() - started) + with open(os.path.join(out_dir, 'challenge_metrics.json'), 'w', + encoding='utf-8') as f: + json.dump(metrics, f, indent=2, ensure_ascii=False, default=str) + logger.info(f'[rsi] {len(run.kept)}/{run.n_launched} groups kept in ' + f'{time.time() - started:.0f}s: {metrics["scalars"]}') + return metrics + + +def main(): + args = C.parse_args() + root = os.path.join(args.root, args.tag) + os.makedirs(root, exist_ok=True) + ckpt_dir = args.ckpt_dir or os.path.join(root, 'ckpt') + # save() writes <output_dir>/<name>, and --model-id takes an HF directory, so + # the next start reads back exactly what the last one wrote. + hf_dir = os.path.join(ckpt_dir, 'model') + + start = next_iteration(root) + model_id, resume_from = args.model_id, None + if start > 1: + if not os.path.exists(os.path.join(hf_dir, 'config.json')): + raise SystemExit( + f'[rsi] {start - 1} iteration(s) finished under {root} but there ' + f'is no checkpoint at {hf_dir}. One directory holds the whole ' + f'loop and each save overwrites the last, so those weights are ' + f'gone: start a new --tag, or delete the iteration.done markers ' + f'to redo them from {args.model_id}.') + model_id = hf_dir + # Written by save(save_optimizer=True). Its absence is not an error, it + # means the crash landed between two optimizer checkpoints. + if os.path.exists(os.path.join(hf_dir, 'trainer_state.json')): + resume_from = hf_dir + else: + logger.warning(f'[rsi] no optimizer state in {hf_dir}; resuming from ' + f'the weights with Adam at zero moments') + + total_gpus = args.model_gpus + args.sampler_gpus + logger.info(f'[rsi] tag {args.tag}, iterations from {start}' + f'{"" if not args.iterations else f" for {args.iterations}"}, ' + f'{args.model_gpus} trainer + {args.sampler_gpus} sampler GPUs, ' + f'model {model_id}, checkpoint {hf_dir}, lr {args.lr}') + + # Both groups are named here, once, and every remote object below is pinned to + # one of them. Disjoint rank ranges are what keeps the two halves from sharing + # a card. + twinkle.initialize( + mode='ray', nproc_per_node=total_gpus, lazy_collect=False, + groups=[ + DeviceGroup(name='model', ranks=list(range(args.model_gpus)), + device_type='GPU'), + DeviceGroup(name='sampler', ranks=list(range(args.model_gpus, total_gpus)), + device_type='GPU'), + ]) + + model = T.build_model(model_id=model_id, model_gpus=args.model_gpus, lr=args.lr, + template=args.template, max_length=args.max_train_len) + if resume_from: + state = model.resume_from_checkpoint(resume_from) + logger.info(f'[rsi] optimizer resumed from {resume_from}: {state}') + sampler, template = C.build_sampler(args) + # Model rank 0 serves the TCPStore the sampler ranks connect to, so this must + # be built after both halves exist. Its first call is what sends the weights. + weights = CheckpointEngineManager(model=model, sampler=sampler) + slots = C.initialize_sandbox(args) + logger.info(get_device_placement()) + + i = start + try: + while not args.iterations or i < start + args.iterations: + out_dir = os.path.join(root, f'iter{i}') + logger.info(f'[rsi] iteration {i}: collect -> {out_dir}') + challenge_metrics = collect_once(args, sampler, template, slots, out_dir) + + logger.info(f'[rsi] iteration {i}: train on {out_dir}') + summary = T.train_one_step( + model, out_dir, sides=args.sides, max_length=args.max_train_len, + micro_batch_size=args.micro_batch_size, + mini_batch_size=args.mini_batch_size or args.model_gpus * args.micro_batch_size, + lr=args.lr) + + # The whole point of one process: the weights go to the engines that + # are already running, over NCCL, instead of through the filesystem. + # merge_and_sync=True is the full-parameter path -- there is no adapter + # here, so the merge is a no-op and every weight is sent. + t0 = time.time() + weights.sync_weights(merge_and_sync=True) + # The cache holds keys computed under the old weights. Cheap to drop, + # and wrong to keep. + sampler.reset_prefix_cache() + logger.info(f'[rsi] iteration {i}: weights synced to the sampler in ' + f'{time.time() - t0:.1f}s') + + with_optimizer = (i % args.save_optimizer_every == 0) + t0 = time.time() + model.save('model', output_dir=ckpt_dir, save_optimizer=with_optimizer) + logger.info(f'[rsi] iteration {i}: checkpoint at {hf_dir} in ' + f'{time.time() - t0:.0f}s' + f'{" with optimizer state" if with_optimizer else ""}') + + # Caught rather than allowed to propagate: an unreachable dashboard + # would otherwise take down a loop whose checkpoint is already on disk + # and, worse than losing the charts, leave the iteration without its + # marker, so the next start would redo it from weights that already + # contain it. Measured: swanlab 0.7.17 with no api key raises + # KeyFileError here, which is exactly the case this has to survive. + try: + T.upload(challenge_metrics.get('scalars') or {}, summary, tag=args.tag, + iteration=i, project=args.swanlab_project, + mode=args.swanlab_mode, + config={'model_id': args.model_id, 'learning_rate': args.lr, + 'sides': args.sides, 'model_gpus': args.model_gpus, + 'sampler_gpus': args.sampler_gpus}) + except Exception as e: + logger.warning(f'[rsi] swanlab upload failed, charts lost but the ' + f'checkpoint and the json metrics are not: ' + f'{type(e).__name__}: {e}') + # Last, so a resume counts only iterations whose weights are on disk. + open(os.path.join(out_dir, 'iteration.done'), 'w').close() + logger.info(f'[rsi] iteration {i} done') + i += 1 + finally: + rebuilds = close_pool(slots) + if rebuilds: + logger.warning(f'[rsi] sandboxes were rebuilt {rebuilds} time(s); the ' + f'jobs in flight at those moments were lost') + logger.info(f'[rsi] stopped after iteration {i - 1}; model at {hf_dir}') + + +if __name__ == '__main__': + main() diff --git a/cookbook/rsi/agentic/train.py b/cookbook/rsi/agentic/train.py index a017d1ae0..04c611f3d 100644 --- a/cookbook/rsi/agentic/train.py +++ b/cookbook/rsi/agentic/train.py @@ -1,27 +1,24 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""One GRPO step on what challenge.py collected. +"""One GRPO step on what a collection pass left in a run directory. -Read ``trajs/index.jsonl``, group it, turn rewards into advantages, accumulate the -whole collection into a single optimizer step, and overwrite the checkpoint the -next iteration loads. +A library, not a script: rsi.py owns the process, the model and the iteration +loop, and calls in here once per iteration. The model it hands over is resident, +which is the whole point -- see rsi.py's docstring for what the old +process-per-iteration arrangement did to the updates. There is no filtering here. Every rule about what is worth training on was applied while collecting -- a group is on disk only if it was kept, and a kept group is -exactly its 8 proposals plus the 8 attempts at its selected task -- so anything -this script dropped would be a second, invisible policy on top of that one. What -it does refuse is a trajectory the model cannot be stepped on at all: no logprobs, -no trainable token, a logprob count that disagrees with the trainable count, more +exactly its proposals plus the attempts at its selected task -- so anything this +dropped would be a second, invisible policy on top of that one. What it does +refuse is a trajectory the model cannot be stepped on at all: no logprobs, no +trainable token, a logprob count that disagrees with the trainable count, more tokens than the model accepts, or a group left with fewer than two members. Each refusal is named and counted in the summary rather than folded into a total. -One step, not several: every trajectory here was sampled from one set of weights, -so a second step would be training weights that no longer produced their own data, -``old_logps`` would stop matching, and epsilon would start clipping for a reason -that has nothing to do with the policy being wrong. The cost is update frequency, -one per collection. - - RSI_RUN_DIR=output/rsi_agentic python cookbook/rsi/agentic/train.py \\ - --model_gpus 8 --lr 1e-6 +One step per collection, not several: every trajectory was sampled from one set of +weights, so a second step would be training weights that no longer produced their +own data, ``old_logps`` would stop matching, and epsilon would start clipping for a +reason that has nothing to do with the policy being wrong. """ import collections import json @@ -30,90 +27,49 @@ import numpy as np -import twinkle -from twinkle import DeviceGroup, DeviceMesh, get_device_placement, get_logger +from twinkle import DeviceMesh, get_logger from twinkle.advantage import GRPOAdvantage -from twinkle.cli import CLI from twinkle.processor import InputProcessor logger = get_logger() -args = CLI.from_args() - -MODEL_ID = args.model.model_id or 'ms://Qwen/Qwen3-4B' -MODEL_GPUS = args.infra.model_gpus or 8 -# The base text template. Qwen3-4B is text-only and the multimodal subclass -# crashes on encode for it. -TEMPLATE = os.environ.get('RSI_TEMPLATE', 'Template') -# Whatever --lr says, or the CLI's own default. Not written as ``or <number>``: -# the CLI default is never zero, so a fallback here would be dead code that reads -# like the default. -LEARNING_RATE = args.optimizer.learning_rate - -# One trajectory per micro batch, because padding_free is off: a micro batch is -# padded to its longest member, so pairing a short solver attempt with a long -# build episode pays for the long one twice. Builds are whole agentic rollouts -- -# median 7.5k tokens, up to 15.7k -- and at 2 per micro batch a 31k-token padded -# batch died of CUDA OOM with activation recompute already at its most aggressive -# setting. Read from the environment rather than args.training, whose -# micro_batch_size defaults to 2 rather than None. -MICRO_BATCH_SIZE = int(os.environ.get('RSI_MICRO_BATCH_SIZE', 1)) -# forward_backward is declared dispatch='slice_dp', so a mini batch is sliced -# across the data-parallel ranks and each rank collates only its share. That share -# has to hold at least one micro batch, so the floor is MODEL_GPUS * MICRO_BATCH. -MINI_BATCH_SIZE = args.training.mini_batch_size or MODEL_GPUS * MICRO_BATCH_SIZE - -RUN_DIR = os.environ.get('RSI_RUN_DIR', 'output/rsi_agentic') -SAVE_DIR = os.environ.get('RSI_SAVE_DIR', 'output/rsi_agentic/ckpt') -SAVE_NAME = os.environ.get('RSI_SAVE_NAME', 'agentic') -# Longest trajectory fed to the model. Above this the model refuses it mid-step. -MAX_MODEL_LEN = int(os.environ.get('RSI_MAX_MODEL_LEN', 32768)) -# Which side(s) to train: 'both', 'solve', 'propose'. -SIDES = os.environ.get('RSI_SIDES', 'both') - -# swanlab. One experiment for the whole loop rather than one per iteration: the -# question these charts answer is whether iteration k+1 is better than k, which a -# chart that ends after one point cannot show. ``id`` is the tag, so re-running a -# tag appends to its curve and a new tag starts a new one. Both are read from the -# environment because loop.sh is what knows them; the fallbacks parse the run -# directory, which is ``<root>/<tag>/iter<n>``, so a bare ``python train.py`` still -# lands somewhere sensible instead of failing. -SWANLAB_PROJECT = os.environ.get('RSI_SWANLAB_PROJECT', 'twinkle-rsi-agentic') -TAG = os.environ.get('RSI_TAG') or os.path.basename(os.path.dirname( - os.path.abspath(RUN_DIR))) -ITERATION = int(os.environ.get('RSI_ITER') - or ''.join(c for c in os.path.basename(os.path.abspath(RUN_DIR)) - if c.isdigit()) or 0) -# 'disabled' skips it entirely, for a run that should not appear on the dashboard. -SWANLAB_MODE = os.environ.get('RSI_SWANLAB_MODE', 'online') - - -def upload(challenge: Dict[str, Any], training: Dict[str, Any]) -> None: - """Send this iteration's numbers to swanlab, as one step. - Called after the checkpoint is saved, so a swanlab failure costs this - iteration's charts and not its weights. The cost of that order is the - reverse: a crash between the step and here loses the numbers, which are still - on disk in challenge_metrics.json and train_summary.json. - Only ``challenge['scalars']`` goes up, not ``challenge['counts']``: the counts - have keys that exist in one iteration and not the next - (``group_dropped:rubric_error``), and a chart that appears halfway through a - run is read as a change in the run rather than a change in what was recorded. +def build_model(*, model_id: str, model_gpus: int, lr: float, template: str, + max_length: int): + """The resident trainer. Full-parameter: no adapter, so the checkpoint is a + whole model rather than something to merge before the next iteration. """ - import swanlab - swanlab.init(project=SWANLAB_PROJECT, name=TAG, id=TAG, resume='allow', - mode=SWANLAB_MODE, config={'tag': TAG, 'model_id': MODEL_ID, - 'learning_rate': LEARNING_RATE, - 'sides': SIDES, 'gpus': MODEL_GPUS}) - log = {f'challenge/{k}': v for k, v in challenge.items()} - log.update({f'train/{k}': v for k, v in training.items()}) - swanlab.log(log, step=ITERATION) - logger.info(f'[train] swanlab {SWANLAB_PROJECT}/{TAG} step {ITERATION}: ' - f'{len(log)} metrics') - + from twinkle.model.megatron import MegatronModel + # variable_seq_lengths stays off with padding_free: both switches send + # collate_fn down the packed path, and Megatron's TE extension then reads + # PackedSeqParams.pad_between_seqs, which this Megatron-LM checkout does not + # define. Padded batches cost throughput but keep attention on plain sequences. + model = MegatronModel( + model_id=model_id, + device_mesh=DeviceMesh.from_sizes(world_size=model_gpus, dp_size=model_gpus), + remote_group='model', mixed_precision='bf16', variable_seq_lengths=False) + model.set_optimizer('default', lr=lr) + # 'constant' rather than the default cosine, and this matters here in a way it + # did not when each iteration was its own process: the scheduler is stepped + # after every optimizer step and lr_decay_steps=1 would put the second step + # and everything after it at min_lr, which is 0. Constant returns max_lr + # before that check is reached, so every iteration steps at the same rate. + model.set_lr_scheduler('default', lr_decay_steps=1, max_lr=lr, + lr_decay_style='constant') + # beta=0: there is no reference model here, and the KL term needs beta>0 AND + # ref_logps, so any beta above 0 would silently do nothing. + model.set_loss('GRPOLoss', epsilon=0.2, beta=0.0) + model.set_processor(InputProcessor, padding_free=False) + model.set_template(template, model_id=model_id, max_length=max_length, + enable_thinking=True) + # approx_kl on landed data is the check for whether this collection belongs to + # the weights being trained: it should start near zero, and a large value means + # the sampler was not holding these weights. + model.add_metric('GRPOMetric', is_training=True, epsilon=0.2) + return model -def load(run_dir: str) -> tuple: +def load(run_dir: str, *, sides: str, max_length: int) -> tuple: """Read the index into GRPO groups; returns (groups, skipped). A group is ``(side, group_id)`` for the proposing side and @@ -123,8 +79,8 @@ def load(run_dir: str) -> tuple: traj_dir = os.path.join(run_dir, 'trajs') index = os.path.join(traj_dir, 'index.jsonl') if not os.path.exists(index): - raise SystemExit(f'[train] no {index}; run challenge.py first') - wanted = {'both': ('propose', 'solve')}.get(SIDES, (SIDES, )) + raise SystemExit(f'[train] no {index}') + wanted = {'both': ('propose', 'solve')}.get(sides, (sides, )) skipped: collections.Counter = collections.Counter() by_key: Dict[Any, List[Dict[str, Any]]] = collections.OrderedDict() with open(index, encoding='utf-8') as f: @@ -162,12 +118,11 @@ def load(run_dir: str) -> tuple: # than something to trim to the shorter of the two. skipped[f'logps {logps.size} != trainable {n_train}'] += 1 continue - if ids.size > MAX_MODEL_LEN: - # Dropped before the model is built, so the count is in the log - # rather than arriving as an exception in the middle of a step. - # Reachable in normal operation: challenge.py samples at + if ids.size > max_length: + # Dropped here rather than arriving as an exception in the middle + # of a step. Reachable in normal operation: collection samples at # max_model_len 40960, which is above this. - skipped[f'longer than MAX_MODEL_LEN={MAX_MODEL_LEN}'] += 1 + skipped[f'longer than max_length={max_length}'] += 1 continue key = ((side, record.get('group_id')) if side == 'propose' else (side, record.get('group_id'), record.get('proposal_idx'))) @@ -235,69 +190,44 @@ def interleave(groups: List[Dict[str, Any]]) -> List[Dict[str, Any]]: return [g for _, g in marked] -def main(): - groups, skipped = load(RUN_DIR) +def train_one_step(model, run_dir: str, *, sides: str, max_length: int, + micro_batch_size: int, mini_batch_size: int, + lr: float) -> Dict[str, Any]: + """Accumulate everything in ``run_dir`` into one optimizer step. + + Writes train_summary.json next to the collection it trained on and returns it. + """ + groups, skipped = load(run_dir, sides=sides, max_length=max_length) if not groups: - raise SystemExit(f'[train] nothing trainable in {RUN_DIR}: {dict(skipped)}') + raise SystemExit(f'[train] nothing trainable in {run_dir}: {dict(skipped)}') skipped.update(score(groups)) batch = [m for g in interleave(groups) for m in g['members']] mix = collections.Counter(m['side'] for m in batch) sizes = collections.Counter((g['side'], len(g['members'])) for g in groups) logger.info(f'[train] {len(groups)} groups, {len(batch)} trajectories {dict(mix)}; ' f'group sizes {dict(sizes)}') - if skipped: - for note, n in sorted(skipped.items()): - logger.warning(f'[train] skipped: {note} x{n}') - - twinkle.initialize(mode='ray', nproc_per_node=MODEL_GPUS, lazy_collect=False, - groups=[DeviceGroup(name='model', ranks=list(range(MODEL_GPUS)), - device_type='GPU')]) - # Full-parameter: no adapter, so every weight is trained and the checkpoint is - # a whole model rather than something to merge before the next iteration. - from twinkle.model.megatron import MegatronModel - # variable_seq_lengths stays off with padding_free: both switches send - # collate_fn down the packed path, and Megatron's TE extension then reads - # PackedSeqParams.pad_between_seqs, which this Megatron-LM checkout does not - # define. Padded batches cost throughput but keep attention on plain sequences. - model = MegatronModel(model_id=MODEL_ID, device_mesh=DeviceMesh.from_sizes( - world_size=MODEL_GPUS, dp_size=MODEL_GPUS), remote_group='model', - mixed_precision='bf16', variable_seq_lengths=False) - model.set_optimizer('default', lr=LEARNING_RATE) - # Inert at one step: the scheduler is read before lr_step advances it, so the - # update happens at max_lr and there is no second step to decay over. Wired up - # so that splitting the run into steps would give a decay across them. - model.set_lr_scheduler('default', lr_decay_steps=1, max_lr=LEARNING_RATE) - # beta=0: there is no reference model here, and the KL term needs beta>0 AND - # ref_logps, so any beta above 0 would silently do nothing. - model.set_loss('GRPOLoss', epsilon=0.2, beta=0.0) - model.set_processor(InputProcessor, padding_free=False) - model.set_template(TEMPLATE, model_id=MODEL_ID, max_length=MAX_MODEL_LEN, - enable_thinking=True) - # approx_kl on landed data is the check for whether this collection belongs to - # the weights being trained: it should start near zero, and a large value means - # the dump came from a different checkpoint. - model.add_metric('GRPOMetric', is_training=True, epsilon=0.2) - logger.info(get_device_placement()) + for note, n in sorted(skipped.items()): + logger.warning(f'[train] skipped: {note} x{n}') inputs = [{k: m[k] for k in ('input_ids', 'labels', 'attention_mask', 'position_ids')} for m in batch] old_logps = [m['logps'] for m in batch] advantages = [m['advantage'] for m in batch] dropped = 0 - for lo in range(0, len(inputs), MINI_BATCH_SIZE): - hi = min(lo + MINI_BATCH_SIZE, len(inputs)) + for lo in range(0, len(inputs), mini_batch_size): + hi = min(lo + mini_batch_size, len(inputs)) # A tail shorter than a whole mini batch is dropped rather than handed # over: dispatch 'slice_dp' splits it across all ranks, and a batch that # cannot give every rank its own micro batch raises inside _dispatch_args # before collate_fn ever runs. - if hi - lo < MINI_BATCH_SIZE: + if hi - lo < mini_batch_size: dropped = hi - lo logger.warning(f'[train] dropping the last {dropped} trajectories, under ' - f'the mini batch of {MINI_BATCH_SIZE}') + f'the mini batch of {mini_batch_size}') break model.forward_backward(inputs=inputs[lo:hi], old_logps=old_logps[lo:hi], advantages=advantages[lo:hi], - micro_batch_size=MICRO_BATCH_SIZE) + micro_batch_size=micro_batch_size) # Once, after every mini batch: forward_backward neither steps nor zeroes, so # the mini batches above simply add their gradients together and one step # consumes all of them. @@ -319,57 +249,55 @@ def main(): 'group_sizes': {f'{s}:{n}': c for (s, n), c in sizes.items()}, 'advantage_min': min(advantages), 'advantage_max': max(advantages), - 'learning_rate': LEARNING_RATE, + 'learning_rate': lr, 'metrics': log, 'high_kl_records': high_kl or [], # Named, not summed: a collection that lost half its trajectories to one # reason and one that lost none read the same from the metrics alone. 'skipped': dict(skipped), } - with open(os.path.join(RUN_DIR, 'train_summary.json'), 'w', encoding='utf-8') as f: + with open(os.path.join(run_dir, 'train_summary.json'), 'w', encoding='utf-8') as f: json.dump(summary, f, indent=2, ensure_ascii=False, default=str) - model.save(SAVE_NAME, output_dir=SAVE_DIR) - logger.info(f'[train] checkpoint at {os.path.join(SAVE_DIR, SAVE_NAME)}') + return summary - # The collection's own numbers, written by challenge.py in this same directory. - # Absent when train.py is pointed at a directory collected before this existed, - # in which case the training half still goes up alone. - challenge_path = os.path.join(RUN_DIR, 'challenge_metrics.json') - challenge: Dict[str, Any] = {} - if os.path.exists(challenge_path): - with open(challenge_path, encoding='utf-8') as f: - challenge = json.load(f).get('scalars') or {} - else: - logger.warning(f'[train] no {challenge_path}; uploading training metrics only') - training = { - 'groups': len(groups), - 'trajectories': len(batch), - 'trained': len(batch) - dropped, - 'dropped_tail': dropped, - 'propose_trajectories': mix.get('propose', 0), - 'solve_trajectories': mix.get('solve', 0), - 'advantage_min': min(advantages), - 'advantage_max': max(advantages), - 'learning_rate': LEARNING_RATE, - 'skipped_total': sum(skipped.values()), - 'high_kl_sequences': len(high_kl or []), - # Every numeric metric GRPOMetric returned: loss, clip fractions, approx_kl. - **{k: v for k, v in log.items() if isinstance(v, (int, float))}, - } - # Caught rather than allowed to propagate, because this is the last statement - # of the run and loop.sh reads its exit status: an unreachable dashboard would - # otherwise take down a loop whose checkpoint is already on disk, and -- worse - # than losing the charts -- leave the iteration without its iteration.done - # marker, so the next start would redo the iteration from weights that already - # contain it. Measured: swanlab 0.7.17 with no api key raises KeyFileError - # here, which is exactly the case this has to survive. The numbers are in - # challenge_metrics.json and train_summary.json either way. - try: - upload(challenge, training) - except Exception as e: - logger.warning(f'[train] swanlab upload failed, charts lost but the ' - f'checkpoint and the json metrics are not: {type(e).__name__}: {e}') +def upload(challenge: Dict[str, Any], summary: Dict[str, Any], *, tag: str, + iteration: int, project: str, mode: str, config: Dict[str, Any]) -> None: + """Send one iteration's numbers to swanlab, as one step. -if __name__ == '__main__': - main() + One experiment for the whole loop rather than one per iteration: the question + these charts answer is whether iteration k+1 is better than k, which a chart + that ends after one point cannot show. ``id`` is the tag, so re-running a tag + appends to its curve and a new tag starts a new one. + + Only ``challenge['scalars']`` goes up, not its counts: those have keys that + exist in one iteration and not the next (``group_dropped:rubric_error``), and a + chart that appears halfway through a run is read as a change in the run rather + than a change in what was recorded. + + Called after the checkpoint is saved and wrapped by the caller, so an + unreachable dashboard costs this iteration's charts and not its weights. The + numbers are in challenge_metrics.json and train_summary.json either way. + """ + import swanlab + swanlab.init(project=project, name=tag, id=tag, resume='allow', mode=mode, + config={'tag': tag, **config}) + log = {f'challenge/{k}': v for k, v in challenge.items()} + metrics = summary.get('metrics') or {} + log.update({ + 'train/groups': summary['groups'], + 'train/trajectories': summary['trajectories'], + 'train/trained': summary['trained'], + 'train/dropped_tail': summary['dropped_tail'], + 'train/propose_trajectories': summary['sides'].get('propose', 0), + 'train/solve_trajectories': summary['sides'].get('solve', 0), + 'train/advantage_min': summary['advantage_min'], + 'train/advantage_max': summary['advantage_max'], + 'train/learning_rate': summary['learning_rate'], + 'train/skipped_total': sum(summary['skipped'].values()), + 'train/high_kl_sequences': len(summary['high_kl_records']), + # Every numeric metric GRPOMetric returned: loss, clip fractions, approx_kl. + **{f'train/{k}': v for k, v in metrics.items() if isinstance(v, (int, float))}, + }) + swanlab.log(log, step=iteration) + logger.info(f'[train] swanlab {project}/{tag} step {iteration}: {len(log)} metrics') From 753681eb68dd30599e09362248b2462d347ce10a Mon Sep 17 00:00:00 2001 From: tastelikefeet <yuze.zyz@alibaba-inc.com> Date: Tue, 1 Sep 2026 09:04:50 +0800 Subject: [PATCH 52/60] add sh --- cookbook/rsi/agentic/README.md | 9 +++- cookbook/rsi/agentic/run.sh | 75 ++++++++++++++++++++++++++++++++++ 2 files changed, 83 insertions(+), 1 deletion(-) create mode 100644 cookbook/rsi/agentic/run.sh diff --git a/cookbook/rsi/agentic/README.md b/cookbook/rsi/agentic/README.md index 6d0fdb281..2fd1fe62c 100644 --- a/cookbook/rsi/agentic/README.md +++ b/cookbook/rsi/agentic/README.md @@ -103,6 +103,7 @@ Solving side: 1.0 if the check exits 0, else 0.0. ## Files ``` +run.sh start or continue a run: the GPU split, the guards, the env rsi.py the loop: one resident process, collect -> step -> sync, forever challenge.py collect: the queues, the three job bodies, the group decision train.py one GRPO step over what was collected, as a library rsi.py calls @@ -180,9 +181,15 @@ so a task's `n_pass` here and its `pass@k` there are measured against one openin export E2B_API_KEY=... # sandbox host export SANDBOX_API_URL=http://... # sandbox host address, with port export LLM_BACKUP_API_KEY=... # dashscope -python cookbook/rsi/agentic/rsi.py --tag v4 --iterations 1 +TAG=v4 bash cookbook/rsi/agentic/run.sh ``` +`run.sh` only sets up the process — the GPU split, the allocator, the guard against +starting on top of another job — and passes anything else through to `rsi.py`, so +`TAG=v4 bash cookbook/rsi/agentic/run.sh --iterations 1 --keep-groups 4` works. +`python cookbook/rsi/agentic/rsi.py --tag v4` directly is the same thing without +those checks. + `--iterations 0`, the default, runs until killed. Restarting the same `--tag` continues it: iterations are counted by the `iteration.done` marker, which is written after the checkpoint, and the loop picks up from diff --git a/cookbook/rsi/agentic/run.sh b/cookbook/rsi/agentic/run.sh new file mode 100644 index 000000000..6c8b1d3c7 --- /dev/null +++ b/cookbook/rsi/agentic/run.sh @@ -0,0 +1,75 @@ +#!/bin/bash +# Start (or continue) a run of the resident loop. +# +# Everything about what to collect and how to train lives in rsi.py's arguments; +# this only sets up the process. Extra arguments are passed straight through, so +# anything in `python cookbook/rsi/agentic/rsi.py --help` works here: +# +# TAG=v4 bash cookbook/rsi/agentic/run.sh +# TAG=v4 bash cookbook/rsi/agentic/run.sh --keep-groups 4 --iterations 1 +# +# Restarting the same TAG continues it from the last finished iteration. +set -u +set -o pipefail + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO="$(cd "$HERE/../../.." && pwd)" +if [ ! -f "$HERE/rsi.py" ] || [ ! -f "$REPO/setup.cfg" ]; then + echo "expected rsi.py beside this script and the repo root three levels up" >&2 + exit 1 +fi +cd "$REPO" + +missing="" +for v in TAG E2B_API_KEY SANDBOX_API_URL LLM_BACKUP_API_KEY; do + [ -z "${!v:-}" ] && missing="$missing $v" +done +if [ -n "$missing" ]; then + echo "set these first:$missing" >&2 + echo " TAG names the run; the other three are the sandbox host and dashscope" >&2 + exit 1 +fi + +MODEL_GPUS="${MODEL_GPUS:-2}" +SAMPLER_GPUS="${SAMPLER_GPUS:-6}" +GPUS=$((MODEL_GPUS + SAMPLER_GPUS)) +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-$(seq -s, 0 $((GPUS - 1)))}" + +# Refuse to start on top of another job. Both halves want whole cards -- the +# sampler takes 0.8 of each of its own and the trainer holds ~40 GB of weights and +# optimizer state for the whole run -- so sharing means an out-of-memory crash +# partway in, and the other job may go down with it. This guard has already caught +# the case worth catching: a previous run's actors still exiting, each still +# holding tens of GB, at the moment a new one started. CONFIRM_GPUS=1 starts anyway. +BUSY="$(nvidia-smi --query-compute-apps=pid --format=csv,noheader | sort -u | wc -l)" +if [ "$BUSY" -gt 0 ] && [ "${CONFIRM_GPUS:-0}" != "1" ]; then + echo "$BUSY process(es) already on the GPUs:" >&2 + nvidia-smi --query-compute-apps=pid,used_memory --format=csv,noheader >&2 + echo "set CONFIRM_GPUS=1 to start anyway" >&2 + exit 1 +fi + +# One padded trajectory per micro batch means every micro batch is a new shape, and +# the caching allocator cannot reuse a block across sizes: on v3 it grew to 87.8 GiB +# reserved against 29.0 GiB live and starved NCCL of the few hundred MB it needs to +# connect a communicator, which hung an iteration for 54 minutes. Expandable +# segments let one virtual range serve every shape. +export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" +export TWINKLE_DISABLE_CUDNN_SDP="${TWINKLE_DISABLE_CUDNN_SDP:-1}" +# INFO here is tens of thousands of lines per iteration through Ray's log forwarding, +# which is worth having only while chasing a collective. +export NCCL_DEBUG="${NCCL_DEBUG:-WARN}" + +ROOT="${ROOT:-output/rsi_agentic}" +mkdir -p "$ROOT/$TAG" +LOG="$ROOT/$TAG/run.log" +echo "=== $TAG: $MODEL_GPUS trainer + $SAMPLER_GPUS sampler GPUs, logging to $LOG" + +# tee rather than a redirect so a foreground run is watchable, and pipefail above +# so the exit status is python's and not tee's. +python cookbook/rsi/agentic/rsi.py \ + --tag "$TAG" \ + --model-gpus "$MODEL_GPUS" \ + --sampler-gpus "$SAMPLER_GPUS" \ + --root "$ROOT" \ + "$@" 2>&1 | tee -a "$LOG" From 5984ccb0eebb00cc7798e8d89680ac5fd2a9fc94 Mon Sep 17 00:00:00 2001 From: tastelikefeet <yuze.zyz@alibaba-inc.com> Date: Tue, 1 Sep 2026 13:44:37 +0800 Subject: [PATCH 53/60] fix --- cookbook/rsi/agentic/README.md | 14 +++++++----- cookbook/rsi/agentic/challenge.py | 10 ++++++-- cookbook/rsi/agentic/rsi.py | 27 ++++++++-------------- cookbook/rsi/agentic/train.py | 38 +++++++++++++++++++++---------- 4 files changed, 52 insertions(+), 37 deletions(-) diff --git a/cookbook/rsi/agentic/README.md b/cookbook/rsi/agentic/README.md index 2fd1fe62c..1d85b3ca2 100644 --- a/cookbook/rsi/agentic/README.md +++ b/cookbook/rsi/agentic/README.md @@ -198,12 +198,14 @@ it is checkpointed every `--save-optimizer-every` iterations (5); a crash betwee two of those resumes with the weights and with Adam at zero moments. Charts land in swanlab project `twinkle-rsi-agentic`, one experiment named after -`--tag`, one step per iteration. The upload happens after the checkpoint is saved -and its failure is caught, so an unreachable dashboard costs the charts and not the -weights — the numbers are still in `challenge_metrics.json` and -`train_summary.json` either way. Resume is by `id=tag`: a second run under the same -tag appends to that curve, a new tag starts a new one. `--swanlab-mode disabled` -turns it off, `--swanlab-project` moves it. +`--tag`, one step per iteration. `swanlab.init` happens once at startup, before the +GPUs are touched, and it is not guarded: a dashboard that will not accept this +client stops the run in the first second rather than after the first iteration. The +per-iteration upload is not guarded either — the numbers are in +`challenge_metrics.json` and `train_summary.json` either way, but a connection that +worked at startup and fails mid-run is worth stopping on. Resume is by `id=tag`: a +second run under the same tag appends to that curve, a new tag starts a new one. +`--swanlab-mode disabled` turns it off, `--swanlab-project` moves it. Verified on this machine at swanlab 0.9.2: three separate processes with the same tag at steps 1, 2, 3 landed on one run (the second and third print `disabled in diff --git a/cookbook/rsi/agentic/challenge.py b/cookbook/rsi/agentic/challenge.py index 3481a074a..5fb270a48 100644 --- a/cookbook/rsi/agentic/challenge.py +++ b/cookbook/rsi/agentic/challenge.py @@ -177,10 +177,16 @@ def parse_args(): help='API calls in flight. Only the rubric runs as its own job; ' 'check and statement calls are made from inside a sandbox ' 'job and are already capped by the slot count.') - p.add_argument('--api-thinking-budget', type=int, default=0, + p.add_argument('--api-thinking-budget', type=int, default=4096, help='sent as extra_body on every API call when > 0. Capping the ' 'reasoning is the one knob that moved wall-clock: 58s -> 10s ' - 'per turn at 2048 on a ~15k-character context.') + 'per turn at 2048 on a ~15k-character context. Not zero by ' + 'default, because the judge is a reasoning model and left ' + 'uncapped it spends thousands of reasoning tokens on a reply ' + 'of nine short lines: measured on v5 iteration 1, 55 of 129 ' + 'novelty calls hit the 120s client timeout and took their ' + 'whole group down as novelty_unscored, at three times the ' + 'wall-clock of the same collection with this set.') # Building: stage 1, the part that is trained. p.add_argument('--propose-temp', type=float, default=1.0) diff --git a/cookbook/rsi/agentic/rsi.py b/cookbook/rsi/agentic/rsi.py index b2a46e411..dfa2b07c6 100644 --- a/cookbook/rsi/agentic/rsi.py +++ b/cookbook/rsi/agentic/rsi.py @@ -143,6 +143,15 @@ def main(): f'{args.model_gpus} trainer + {args.sampler_gpus} sampler GPUs, ' f'model {model_id}, checkpoint {hf_dir}, lr {args.lr}') + # Before the GPUs: a dashboard that will not accept this client is worth + # finding out about now rather than 35 minutes in, and there is nothing to + # lose yet if it raises. + T.init_swanlab(tag=args.tag, project=args.swanlab_project, + mode=args.swanlab_mode, + config={'model_id': args.model_id, 'learning_rate': args.lr, + 'sides': args.sides, 'model_gpus': args.model_gpus, + 'sampler_gpus': args.sampler_gpus}) + # Both groups are named here, once, and every remote object below is pinned to # one of them. Disjoint rank ranges are what keeps the two halves from sharing # a card. @@ -200,23 +209,7 @@ def main(): f'{time.time() - t0:.0f}s' f'{" with optimizer state" if with_optimizer else ""}') - # Caught rather than allowed to propagate: an unreachable dashboard - # would otherwise take down a loop whose checkpoint is already on disk - # and, worse than losing the charts, leave the iteration without its - # marker, so the next start would redo it from weights that already - # contain it. Measured: swanlab 0.7.17 with no api key raises - # KeyFileError here, which is exactly the case this has to survive. - try: - T.upload(challenge_metrics.get('scalars') or {}, summary, tag=args.tag, - iteration=i, project=args.swanlab_project, - mode=args.swanlab_mode, - config={'model_id': args.model_id, 'learning_rate': args.lr, - 'sides': args.sides, 'model_gpus': args.model_gpus, - 'sampler_gpus': args.sampler_gpus}) - except Exception as e: - logger.warning(f'[rsi] swanlab upload failed, charts lost but the ' - f'checkpoint and the json metrics are not: ' - f'{type(e).__name__}: {e}') + T.upload(challenge_metrics.get('scalars') or {}, summary, iteration=i) # Last, so a resume counts only iterations whose weights are on disk. open(os.path.join(out_dir, 'iteration.done'), 'w').close() logger.info(f'[rsi] iteration {i} done') diff --git a/cookbook/rsi/agentic/train.py b/cookbook/rsi/agentic/train.py index 04c611f3d..9ecd3faec 100644 --- a/cookbook/rsi/agentic/train.py +++ b/cookbook/rsi/agentic/train.py @@ -23,7 +23,7 @@ import collections import json import os -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional import numpy as np @@ -261,27 +261,41 @@ def train_one_step(model, run_dir: str, *, sides: str, max_length: int, return summary -def upload(challenge: Dict[str, Any], summary: Dict[str, Any], *, tag: str, - iteration: int, project: str, mode: str, config: Dict[str, Any]) -> None: - """Send one iteration's numbers to swanlab, as one step. +# swanlab state for this process. ``init`` may be called once and only once here: +# a second call raises 'DataPorter instance already exists', which the old +# process-per-iteration arrangement never hit because every iteration was a fresh +# interpreter. So it happens once, at startup, before anything expensive has been +# built -- a dashboard that will not accept this client is worth finding out about +# in the first second rather than after the first iteration. +def init_swanlab(*, tag: str, project: str, mode: str, config: Dict[str, Any]) -> None: + """Open the one experiment this process logs to. Raises if it cannot. One experiment for the whole loop rather than one per iteration: the question these charts answer is whether iteration k+1 is better than k, which a chart - that ends after one point cannot show. ``id`` is the tag, so re-running a tag - appends to its curve and a new tag starts a new one. + that ends after one point cannot show. ``id`` is the tag, so a later process + under the same tag appends to its curve and a new tag starts a new one. + """ + import swanlab + swanlab.init(project=project, name=tag, id=tag, resume='allow', mode=mode, + config={'tag': tag, **config}) + logger.info(f'[train] swanlab {project}/{tag}, mode {mode}') + + +def upload(challenge: Dict[str, Any], summary: Dict[str, Any], *, + iteration: int) -> None: + """Send one iteration's numbers to the experiment ``init_swanlab`` opened. Only ``challenge['scalars']`` goes up, not its counts: those have keys that exist in one iteration and not the next (``group_dropped:rubric_error``), and a chart that appears halfway through a run is read as a change in the run rather than a change in what was recorded. - Called after the checkpoint is saved and wrapped by the caller, so an - unreachable dashboard costs this iteration's charts and not its weights. The - numbers are in challenge_metrics.json and train_summary.json either way. + Called after the checkpoint is saved, and it does not swallow anything: the + connection was proved at startup by ``init_swanlab``, so a failure here is a + dashboard that went away mid-run and that is worth stopping on. The numbers + are in challenge_metrics.json and train_summary.json either way. """ import swanlab - swanlab.init(project=project, name=tag, id=tag, resume='allow', mode=mode, - config={'tag': tag, **config}) log = {f'challenge/{k}': v for k, v in challenge.items()} metrics = summary.get('metrics') or {} log.update({ @@ -300,4 +314,4 @@ def upload(challenge: Dict[str, Any], summary: Dict[str, Any], *, tag: str, **{f'train/{k}': v for k, v in metrics.items() if isinstance(v, (int, float))}, }) swanlab.log(log, step=iteration) - logger.info(f'[train] swanlab {project}/{tag} step {iteration}: {len(log)} metrics') + logger.info(f'[train] swanlab step {iteration}: {len(log)} metrics') From f35b8474aac1a7bb2f8778ed6b4743be38dbd5fe Mon Sep 17 00:00:00 2001 From: tastelikefeet <yuze.zyz@alibaba-inc.com> Date: Tue, 1 Sep 2026 13:52:43 +0800 Subject: [PATCH 54/60] fix --- cookbook/rsi/agentic/README.md | 52 ++++++++++++++++++++----------- cookbook/rsi/agentic/challenge.py | 14 +++++++-- cookbook/rsi/agentic/run.sh | 42 ++++++++++++++++++++++--- 3 files changed, 83 insertions(+), 25 deletions(-) diff --git a/cookbook/rsi/agentic/README.md b/cookbook/rsi/agentic/README.md index 1d85b3ca2..d1db30375 100644 --- a/cookbook/rsi/agentic/README.md +++ b/cookbook/rsi/agentic/README.md @@ -181,30 +181,44 @@ so a task's `n_pass` here and its `pass@k` there are measured against one openin export E2B_API_KEY=... # sandbox host export SANDBOX_API_URL=http://... # sandbox host address, with port export LLM_BACKUP_API_KEY=... # dashscope -TAG=v4 bash cookbook/rsi/agentic/run.sh +export LLM_BACKUP_MODEL=... # the judge, e.g. qwen3.8-max +export LLM_BACKUP_BASE_URL=... # its endpoint +TAG=v5 bash cookbook/rsi/agentic/run.sh ``` -`run.sh` only sets up the process — the GPU split, the allocator, the guard against -starting on top of another job — and passes anything else through to `rsi.py`, so -`TAG=v4 bash cookbook/rsi/agentic/run.sh --iterations 1 --keep-groups 4` works. -`python cookbook/rsi/agentic/rsi.py --tag v4` directly is the same thing without -those checks. +That is the whole command: everything a run needs is either one of those five +variables or a default in the code, and nothing has to be remembered on the command +line. `run.sh` sets up the process — the GPU split (`MODEL_GPUS` / `SAMPLER_GPUS`, +2 and 6), the allocator, the guard against starting on top of another job, the +checkpoint directory (`CKPT_DIR`, on `/mnt/data2` rather than the NAS because it is +7.6 GB per iteration), the swanlab mode (`SWANLAB_MODE`, see below) — and passes +anything else through to `rsi.py`, so +`TAG=v5 bash cookbook/rsi/agentic/run.sh --iterations 1 --keep-groups 4` works. +`python cookbook/rsi/agentic/rsi.py --tag v5` directly is the same thing without +those checks and without those process settings. `--iterations 0`, the default, runs until killed. Restarting the same `--tag` continues it: iterations are counted by the `iteration.done` marker, which is -written after the checkpoint, and the loop picks up from -`<root>/<tag>/ckpt/model`. The optimizer is state that only exists in memory, so -it is checkpointed every `--save-optimizer-every` iterations (5); a crash between -two of those resumes with the weights and with Adam at zero moments. - -Charts land in swanlab project `twinkle-rsi-agentic`, one experiment named after -`--tag`, one step per iteration. `swanlab.init` happens once at startup, before the -GPUs are touched, and it is not guarded: a dashboard that will not accept this -client stops the run in the first second rather than after the first iteration. The -per-iteration upload is not guarded either — the numbers are in +written after the checkpoint, and the loop picks up from `$CKPT_DIR/model`. The +optimizer is state that only exists in memory, so it is checkpointed every +`--save-optimizer-every` iterations (5); a crash between two of those resumes with +the weights and with Adam at zero moments. Note that resuming from a checkpoint +that *does* carry optimizer state also restores that checkpoint's learning rate: +Megatron's scheduler prefers the checkpointed value over the class value, so a +restart with a different `--lr` keeps the old one and only says so in an INFO line. + +Charts land in swanlab project `twinkle-rsi-selfplay`, one experiment named after +`--tag`, one step per iteration, pushed to the cloud. `swanlab.init` happens once at +startup, before the GPUs are touched, and it is not guarded: a dashboard that will +not accept this client stops the run in the first second rather than after the first +iteration. The per-iteration upload is not guarded either — the numbers are in `challenge_metrics.json` and `train_summary.json` either way, but a connection that -worked at startup and fails mid-run is worth stopping on. Resume is by `id=tag`: a -second run under the same tag appends to that curve, a new tag starts a new one. +worked at startup and fails mid-run is worth stopping on. The project name is part +of this: the older `twinkle-rsi-agentic` project answers `POST /api/project` with +422 for the client here (0.7.17), while a project this client creates itself works, +so the default was moved rather than the client upgraded. `SWANLAB_MODE=local` +writes `swanlog/` for `swanlab watch` instead. Resume is by `id=tag`: a second run +under the same tag appends to that curve, a new tag starts a new one. `--swanlab-mode disabled` turns it off, `--swanlab-project` moves it. Verified on this machine at swanlab 0.9.2: three separate processes with the same @@ -238,7 +252,7 @@ re-measured under this scheduler. | `--keywords-n` / `--keyword-gen-calls` / `--keyword-temp` | 128 / 8 / 1.3 | inherited | | `--snapshot-max-files` / `-per-file` / `-budget` | 50 / 600 / 6000 | inherited | | `--sandbox-slots` | 32 | inherited: a probe once held 96, but not reliably for a whole run | -| lr / one optimizer step / `GRPOLoss(epsilon=0.2, beta=0.0)` | 1e-6 | inherited | +| lr / one optimizer step / `GRPOLoss(epsilon=0.2, beta=0.0)` | 5e-6 | one step per iteration is the whole of what an iteration moves; 1e-6 was inherited from the era when the bf16 round trip rounded it away anyway | | `MICRO_BATCH_SIZE=1`, `padding_free=False` | — | inherited, forced by an OOM at 2 | Prompt texts are byte-identical to the ones the old pipeline sent — verified diff --git a/cookbook/rsi/agentic/challenge.py b/cookbook/rsi/agentic/challenge.py index 5fb270a48..96311a184 100644 --- a/cookbook/rsi/agentic/challenge.py +++ b/cookbook/rsi/agentic/challenge.py @@ -274,7 +274,13 @@ def parse_args(): p.add_argument('--snapshot-budget', type=int, default=6000) # Training, one step per iteration. See train.py. - p.add_argument('--lr', type=float, default=1e-6) + p.add_argument('--lr', type=float, default=5e-6, + help='One optimizer step per iteration, so this is the whole of ' + 'what an iteration moves. 1e-6 was inherited from the ' + 'process-per-iteration arrangement, where it did not matter: ' + 'the update was rounded away by the bf16 round trip through ' + 'the checkpoint anyway -- after 12 such iterations 98.54%% of ' + 'the weights were still bit-identical to the base model.') p.add_argument('--sides', default='both', choices=('both', 'propose', 'solve')) p.add_argument('--micro-batch-size', type=int, default=1, help='trajectories per micro batch. One, because padding_free is ' @@ -306,7 +312,11 @@ def parse_args(): 'Weights are saved every iteration either way; this is what ' 'a resume needs to keep the Adam moments, and it is ~48 GB ' 'for a 4B model against 7.6 GB for the weights alone.') - p.add_argument('--swanlab-project', default='twinkle-rsi-agentic') + p.add_argument('--swanlab-project', default='twinkle-rsi-selfplay', + help='not twinkle-rsi-agentic: that project answers POST ' + '/api/project with 422 for this client (0.7.17), while a ' + 'project it creates itself works. Measured by initialising ' + 'both from the same interpreter.') p.add_argument('--swanlab-mode', default='online', help="'disabled' keeps a run off the dashboard entirely.") diff --git a/cookbook/rsi/agentic/run.sh b/cookbook/rsi/agentic/run.sh index 6c8b1d3c7..26a63eec1 100644 --- a/cookbook/rsi/agentic/run.sh +++ b/cookbook/rsi/agentic/run.sh @@ -21,12 +21,26 @@ fi cd "$REPO" missing="" -for v in TAG E2B_API_KEY SANDBOX_API_URL LLM_BACKUP_API_KEY; do +for v in TAG E2B_API_KEY SANDBOX_API_URL LLM_BACKUP_API_KEY LLM_BACKUP_MODEL \ + LLM_BACKUP_BASE_URL; do [ -z "${!v:-}" ] && missing="$missing $v" done if [ -n "$missing" ]; then echo "set these first:$missing" >&2 - echo " TAG names the run; the other three are the sandbox host and dashscope" >&2 + echo " TAG names the run; the rest are the sandbox host and the API judge" >&2 + echo " (LLM_BACKUP_MODEL and LLM_BACKUP_BASE_URL are where --api-model and" >&2 + echo " --api-base get their defaults, so an empty one is a run with no judge)" >&2 + exit 1 +fi + +# The interpreter, checked rather than hardcoded: a login shell without the conda +# environment active resolves python to /usr/local/bin/python, whose megatron.core +# has no transformer-engine metadata and raises PackageNotFoundError on import -- +# after Ray is up, which reads as eight actors dying for no stated reason. +PYTHON="${PYTHON:-python}" +if ! "$PYTHON" -c 'import twinkle, megatron.core' 2>/dev/null; then + echo "$PYTHON cannot import twinkle and megatron.core: activate the environment" >&2 + echo " twinkle was installed into, or point PYTHON at its interpreter" >&2 exit 1 fi @@ -61,15 +75,35 @@ export TWINKLE_DISABLE_CUDNN_SDP="${TWINKLE_DISABLE_CUDNN_SDP:-1}" export NCCL_DEBUG="${NCCL_DEBUG:-WARN}" ROOT="${ROOT:-output/rsi_agentic}" +# Not under ROOT: ROOT is on the NAS, and this is 7.6 GB of weights every iteration +# and ~48 GB more on the iterations that include the optimizer. +CKPT_DIR="${CKPT_DIR:-/mnt/data2/rsi_agentic/$TAG/ckpt}" +export MODELSCOPE_CACHE="${MODELSCOPE_CACHE:-/mnt/workspace/.cache/modelscope/hub}" +# Pushed to the cloud dashboard. This works on the default project and not on the +# older twinkle-rsi-agentic one, which answers this client (0.7.17) with 422 and, +# since swanlab.init is no longer guarded, would stop the run in its first second. +# SWANLAB_MODE=local writes swanlog/ instead, for `swanlab watch`. +SWANLAB_MODE="${SWANLAB_MODE:-online}" mkdir -p "$ROOT/$TAG" LOG="$ROOT/$TAG/run.log" -echo "=== $TAG: $MODEL_GPUS trainer + $SAMPLER_GPUS sampler GPUs, logging to $LOG" +echo "=== $TAG: $MODEL_GPUS trainer + $SAMPLER_GPUS sampler GPUs, checkpoint $CKPT_DIR," +echo "=== swanlab $SWANLAB_MODE, logging to $LOG" # tee rather than a redirect so a foreground run is watchable, and pipefail above # so the exit status is python's and not tee's. -python cookbook/rsi/agentic/rsi.py \ +# +# Every knob that decides what gets collected or how it is trained is left to +# rsi.py's own defaults on purpose. The one time a verified setting lived in a +# launcher instead -- --api-thinking-budget 4096, in a throwaway script under +# .temp -- a restart that retyped the command line dropped it, the rubric judge +# went back to thinking without a cap, 43% of its calls hit the 120s timeout and +# took their whole group down, and the iteration ran at three times its usual +# wall-clock before anyone noticed. +$PYTHON cookbook/rsi/agentic/rsi.py \ --tag "$TAG" \ --model-gpus "$MODEL_GPUS" \ --sampler-gpus "$SAMPLER_GPUS" \ --root "$ROOT" \ + --ckpt-dir "$CKPT_DIR" \ + --swanlab-mode "$SWANLAB_MODE" \ "$@" 2>&1 | tee -a "$LOG" From fc825bb8465e3b0f597bff2860c480a0806644fc Mon Sep 17 00:00:00 2001 From: "dreaming.ljm" <dreaming.ljm@qoder.alibaba-inc.com> Date: Tue, 1 Sep 2026 21:43:39 +0800 Subject: [PATCH 55/60] harness: ms-agent 1.6.0 compatibility and argv guard for embedded workers - sanitize sys.argv around LLMAgent construction: ms-agent's Config.parse_args() mis-parses a foreign argv in forked/Ray workers (assert crash on value tokens; silent flag mispairing otherwise) - tolerate ms-agent >= 1.6.0 API changes: prepare_skills -> _ensure_auto_skills(), dropped ms_agent.hooks and _append_task_notifications, defensive ToolResult field forwarding --- src/twinkle_agentic/harness/ms_agent.py | 68 +++++++++++++++++-------- 1 file changed, 48 insertions(+), 20 deletions(-) diff --git a/src/twinkle_agentic/harness/ms_agent.py b/src/twinkle_agentic/harness/ms_agent.py index 5eb275890..8807e58ae 100644 --- a/src/twinkle_agentic/harness/ms_agent.py +++ b/src/twinkle_agentic/harness/ms_agent.py @@ -28,6 +28,7 @@ import json import os +import sys import uuid from typing import Any, Dict, List, Optional, Union @@ -83,11 +84,20 @@ def __init__( cfg = OmegaConf.create(config) else: cfg = config - self.agent = LLMAgent( - cfg, - trust_remote_code=trust_remote_code, - **agent_kwargs, - ) + # ms-agent's Config.parse_args() reads sys.argv and asserts every + # token is a --key/value pair. When this harness is built inside a + # Ray worker, sys.argv carries the driver's args (e.g. rsi.py's), + # which break that parser. Hide them during construction. + saved_argv = sys.argv + sys.argv = [saved_argv[0]] + try: + self.agent = LLMAgent( + cfg, + trust_remote_code=trust_remote_code, + **agent_kwargs, + ) + finally: + sys.argv = saved_argv self.auto_prepare = auto_prepare self.freeze_system = freeze_system self.permission_mode = permission_mode @@ -116,7 +126,12 @@ def start(self, query: str, **kwargs) -> Trajectory: return traj def before_generate(self, trajectory: Trajectory) -> Trajectory: - from ms_agent.hooks.context import condense_hook_attachments_for_llm + # ms-agent >= 1.6 removed ms_agent.hooks; the two helpers below + # moved or dropped, so degrade gracefully per installed version. + try: + from ms_agent.hooks.context import condense_hook_attachments_for_llm + except ImportError: # ms-agent >= 1.6 dropped ms_agent.hooks + condense_hook_attachments_for_llm = None if self.auto_prepare: self.prepare() @@ -124,8 +139,11 @@ def before_generate(self, trajectory: Trajectory) -> Trajectory: frozen_system = messages[0].content if (self.freeze_system and messages and messages[0].role == 'system') else None - messages = self.agent._append_task_notifications(messages) - messages = condense_hook_attachments_for_llm(messages) + # _append_task_notifications existed in older ms-agent; skip on >= 1.6. + if hasattr(self.agent, '_append_task_notifications'): + messages = self.agent._append_task_notifications(messages) + if condense_hook_attachments_for_llm is not None: + messages = condense_hook_attachments_for_llm(messages) if getattr(self.agent, 'runtime', None) is not None: run_sync(self.agent.on_generate_response, messages) @@ -183,17 +201,22 @@ def after_tools( tc = calls[i] if i < len(calls) else {} tid = tc.get('id') or str(uuid.uuid4())[:8] name = tc.get('tool_name') or '' - messages.append( - Message( - role='tool', - content=formatted.text, - tool_call_id=tid, - name=name, - resources=formatted.resources, - tool_detail=formatted.tool_detail, - hook_attachments=formatted.hook_attachments, - is_error=formatted.is_error, - )) + kwargs: Dict[str, Any] = { + 'role': 'tool', + 'content': formatted.text, + 'tool_call_id': tid, + 'name': name, + } + # ms-agent 1.6.0 ToolResult.from_raw() only carries text/ + # resources/extra; older versions carried the fields below on + # the object. Forward whichever exist so Message never gets a + # kwarg it cannot take. + for _field in ('resources', 'tool_detail', 'hook_attachments', + 'is_error'): + _value = getattr(formatted, _field, None) + if _value is not None: + kwargs[_field] = _value + messages.append(Message(**kwargs)) if i < len(calls) and not tc.get('id'): calls[i]['id'] = tid @@ -231,7 +254,12 @@ async def _prepare_async(self) -> None: agent.prepare_runtime() if getattr(agent, 'tool_manager', None) is None: await agent.prepare_tools() - await agent.prepare_skills() + if hasattr(agent, 'prepare_skills'): + await agent.prepare_skills() + else: + # ms-agent >= 1.6 has no prepare_skills: AutoSkills initializes + # lazily on first use, so only force the lazy init here. + agent._ensure_auto_skills() await agent.load_memory() if hasattr(agent, 'prepare_rag'): await agent.prepare_rag() From 65a49168838c65c33fbf8882025fc1dad858ffae Mon Sep 17 00:00:00 2001 From: tastelikefeet <yuze.zyz@alibaba-inc.com> Date: Tue, 1 Sep 2026 22:00:24 +0800 Subject: [PATCH 56/60] fix --- cookbook/rsi/agentic/challenge.py | 17 +++++++++++++++-- cookbook/rsi/agentic/rsi.py | 19 ++++++++++++++++--- cookbook/rsi/agentic/run.sh | 10 ++++++++++ 3 files changed, 41 insertions(+), 5 deletions(-) diff --git a/cookbook/rsi/agentic/challenge.py b/cookbook/rsi/agentic/challenge.py index 96311a184..7ce45782c 100644 --- a/cookbook/rsi/agentic/challenge.py +++ b/cookbook/rsi/agentic/challenge.py @@ -291,9 +291,22 @@ def parse_args(): help='0 means --model-gpus x --micro-batch-size, which is the ' "floor: forward_backward is dispatch='slice_dp', so a mini " 'batch has to give every rank at least one micro batch.') - p.add_argument('--max-train-len', type=int, default=32768, + p.add_argument('--max-train-len', type=int, default=16384, help='a trajectory longer than this is not trained on. Below ' - '--max-model-len, so collection can produce some.') + '--max-model-len, so collection can produce some. This is ' + 'the one bound on training memory that works when the micro ' + 'batch is already a single sequence: the peak allocation is ' + 'the logits, vocab 151936 x length x 2 bytes, which is ' + '0.29 MiB per token and nothing else comes close -- the ' + 'allocation that ran out was 5.87 GiB for a 20742-token ' + 'trajectory, matching that product exactly. 16384 leaves ' + '~0.8 GiB of headroom under the longest trajectory that has ' + 'ever trained here (19312) and costs, measured over 12 ' + 'iterations of v4 and v5, 0.98%% of trajectories (3.91%% in ' + 'the worst single iteration) and not one whole group: an ' + 'eight-member group can lose its longest one or two members ' + 'and still produce an advantage. 32768, the old default, was ' + 'above every length ever collected and so never fired.') # The loop. p.add_argument('--root', default='output/rsi_agentic') diff --git a/cookbook/rsi/agentic/rsi.py b/cookbook/rsi/agentic/rsi.py index dfa2b07c6..5812f3988 100644 --- a/cookbook/rsi/agentic/rsi.py +++ b/cookbook/rsi/agentic/rsi.py @@ -129,9 +129,22 @@ def main(): f'gone: start a new --tag, or delete the iteration.done markers ' f'to redo them from {args.model_id}.') model_id = hf_dir - # Written by save(save_optimizer=True). Its absence is not an error, it - # means the crash landed between two optimizer checkpoints. - if os.path.exists(os.path.join(hf_dir, 'trainer_state.json')): + # Written by save(save_optimizer=True), which only fires every + # --save-optimizer-every iterations. Staleness is not a matter of losing a + # few moments: _load_mcore_optimizer reads latest_checkpointed_iteration.txt + # and restores the model from that sub-checkpoint too, so an optimizer state + # older than the weights sitting beside it rolls the weights back to + # whichever iteration wrote it -- silently, since both come from the same + # directory. Only the iteration that saved it may load it back. + saved_at = (start - 1) - (start - 1) % args.save_optimizer_every + if saved_at != start - 1: + logger.warning( + f'[rsi] the optimizer state under {hf_dir} is from iteration ' + f'{saved_at} and the weights are from {start - 1}; loading it would ' + f'take the weights back with it, so it is skipped and Adam starts ' + f'at zero moments. Every {args.save_optimizer_every} iterations is ' + f'a resume point; the others cost the fp32 master residue.') + elif os.path.exists(os.path.join(hf_dir, 'trainer_state.json')): resume_from = hf_dir else: logger.warning(f'[rsi] no optimizer state in {hf_dir}; resuming from ' diff --git a/cookbook/rsi/agentic/run.sh b/cookbook/rsi/agentic/run.sh index 26a63eec1..d3b40740d 100644 --- a/cookbook/rsi/agentic/run.sh +++ b/cookbook/rsi/agentic/run.sh @@ -84,6 +84,15 @@ export MODELSCOPE_CACHE="${MODELSCOPE_CACHE:-/mnt/workspace/.cache/modelscope/hu # since swanlab.init is no longer guarded, would stop the run in its first second. # SWANLAB_MODE=local writes swanlog/ instead, for `swanlab watch`. SWANLAB_MODE="${SWANLAB_MODE:-online}" +# The sequence limit that bounds training memory. Deliberately not given a value +# here: rsi.py's default (16384) is the measured one and duplicating the number in +# two places is how the two drift apart. Set MAX_TRAIN_LEN to override it, which is +# what a change of model, vocabulary, or trainer GPU count calls for -- the bound is +# vocab x length x 2 bytes of logits against whatever the card has left. +LIMIT=() +if [ -n "${MAX_TRAIN_LEN:-}" ]; then + LIMIT=(--max-train-len "$MAX_TRAIN_LEN") +fi mkdir -p "$ROOT/$TAG" LOG="$ROOT/$TAG/run.log" echo "=== $TAG: $MODEL_GPUS trainer + $SAMPLER_GPUS sampler GPUs, checkpoint $CKPT_DIR," @@ -106,4 +115,5 @@ $PYTHON cookbook/rsi/agentic/rsi.py \ --root "$ROOT" \ --ckpt-dir "$CKPT_DIR" \ --swanlab-mode "$SWANLAB_MODE" \ + ${LIMIT[@]+"${LIMIT[@]}"} \ "$@" 2>&1 | tee -a "$LOG" From cb20ddbc042dabf4a8db0707e306ff1a42ec1d58 Mon Sep 17 00:00:00 2001 From: root <yuze.zyz@alibaba-inc.com> Date: Wed, 9 Sep 2026 14:26:54 +0800 Subject: [PATCH 57/60] wip --- cookbook/rl/grpo/kodcode_grpo.py | 36 +- cookbook/rl/grpo/mbpp_grpo.py | 28 +- cookbook/rsi/agentic/challenge.py | 564 ++++---------- cookbook/rsi/agentic/episode.py | 12 +- cookbook/rsi/agentic/remote_tool_env.py | 80 +- cookbook/rsi/agentic/rsi.py | 90 ++- cookbook/rsi/agentic/sandbox.py | 72 +- cookbook/rsi/agentic/train.py | 50 +- cookbook/rsi/code/challenge.py | 26 +- .../code/{prompts.py => challenge_prompts.py} | 0 cookbook/rsi/code/collect.py | 240 ++++++ cookbook/rsi/recorder.py | 170 +++++ cookbook/rsi/rl.py | 300 ++------ src/twinkle/data_format/__init__.py | 2 +- src/twinkle/data_format/trajectory.py | 19 +- src/twinkle_agentic/challenger/__init__.py | 21 +- src/twinkle_agentic/challenger/agentic.py | 721 +++++------------- src/twinkle_agentic/challenger/api.py | 154 ++++ src/twinkle_agentic/challenger/base.py | 202 ++++- src/twinkle_agentic/challenger/code.py | 548 ++++--------- src/twinkle_agentic/challenger/keywords.py | 588 ++++++++++++++ src/twinkle_agentic/challenger/new/keyword.py | 12 + src/twinkle_agentic/envs/__init__.py | 3 +- src/twinkle_agentic/envs/agentenv.py | 91 +-- src/twinkle_agentic/envs/base.py | 202 +++++ src/twinkle_agentic/envs/local.py | 260 +++++++ src/twinkle_agentic/harness/base.py | 2 +- src/twinkle_agentic/harness/ms_agent.py | 2 +- .../preprocessor/data_juicer.py | 2 +- .../preprocessor/dead_loop_filter.py | 3 +- .../preprocessor/dedup_filter.py | 2 +- .../preprocessor/experimental/score_filter.py | 2 +- .../preprocessor/hard_filter.py | 3 +- .../preprocessor/intent_classifier.py | 2 +- .../preprocessor/language_filter.py | 4 +- .../preprocessor/message_normalizer.py | 2 +- .../preprocessor/message_sanity.py | 9 +- .../preprocessor/offline/decontaminate.py | 3 +- .../preprocessor/offline/near_dedup.py | 3 +- .../preprocessor/structural_noise.py | 3 +- .../preprocessor/token_soup.py | 2 +- src/twinkle_agentic/preprocessor/utils.py | 28 - src/twinkle_agentic/reward/f1.py | 22 +- src/twinkle_agentic/utils/code_utils.py | 106 +++ .../{preprocessor => utils}/message_utils.py | 102 ++- src/twinkle_agentic/utils/text_utils.py | 62 ++ src/twinkle_agentic/verifier/__init__.py | 4 +- src/twinkle_agentic/verifier/result_check.py | 109 +-- ...ocessor_utils.py => test_logprob_utils.py} | 8 +- tests/twinkle_agentic/test_agentic_rsi.py | 573 +++++++++----- 50 files changed, 3302 insertions(+), 2247 deletions(-) rename cookbook/rsi/code/{prompts.py => challenge_prompts.py} (100%) create mode 100644 cookbook/rsi/code/collect.py create mode 100644 cookbook/rsi/recorder.py create mode 100644 src/twinkle_agentic/challenger/api.py create mode 100644 src/twinkle_agentic/challenger/keywords.py create mode 100644 src/twinkle_agentic/challenger/new/keyword.py create mode 100644 src/twinkle_agentic/envs/local.py delete mode 100644 src/twinkle_agentic/preprocessor/utils.py create mode 100644 src/twinkle_agentic/utils/code_utils.py rename src/twinkle_agentic/{preprocessor => utils}/message_utils.py (57%) create mode 100644 src/twinkle_agentic/utils/text_utils.py rename tests/preprocessor/{test_preprocessor_utils.py => test_logprob_utils.py} (97%) diff --git a/cookbook/rl/grpo/kodcode_grpo.py b/cookbook/rl/grpo/kodcode_grpo.py index 09195f41f..a7b44447a 100644 --- a/cookbook/rl/grpo/kodcode_grpo.py +++ b/cookbook/rl/grpo/kodcode_grpo.py @@ -43,6 +43,8 @@ from twinkle.processor import InputProcessor from twinkle.reward.base import Reward from twinkle.sampler import vLLMSampler +from twinkle_agentic.utils.code_utils import unwrap_code +from twinkle_agentic.utils.message_utils import assistant_text logger = get_logger() args = CLI.from_args() @@ -80,34 +82,14 @@ TEST_TIMEOUT = int(os.environ.get('TEST_TIMEOUT', 60)) -_FENCE_RE = re.compile(r'```(?:python|py)?\s*\n(.*?)```', re.S) _SPECIAL_TOKEN_RE = re.compile(r'<\|[^|]+\|>') # ========== Text handling (same as e18_kodcode) ========== -def after_think(text: str) -> str: - """Keep only what follows </think>; return the text unchanged if unclosed.""" - idx = text.rfind('</think>') - return text[idx + len('</think>'):] if idx >= 0 else text - - def clean_text(decoded: Optional[str]) -> str: return _SPECIAL_TOKEN_RE.sub('', decoded or '').strip() -def extract_code(text: str) -> str: - """Take the last fenced block; fall back to the whole body when unfenced. - - The last one, not the first: models often draft a version before the final - one, and the last block is their conclusion. - """ - body = after_think(text) - blocks = _FENCE_RE.findall(body) - if blocks: - return blocks[-1].strip() - return body.strip() - - # ========== Sandbox (same contract as e18_kodcode.run_tests) ========== # Assertion vs exception must be told apart via ``reprcrash.message``: pytest # rewrites assertions, so the summary reads "E assert -1 == 3" and the string @@ -274,12 +256,7 @@ def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: break if payload is None: continue - completion = '' - for msg in reversed(traj.get('messages', []) or []): - if msg.get('role') == 'assistant': - completion = msg.get('content', '') or '' - break - jobs.append((i, extract_code(completion), payload)) + jobs.append((i, unwrap_code(assistant_text(traj)), payload)) if not jobs: return rewards @@ -302,12 +279,7 @@ class KodCodeFormatReward(Reward): def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: rewards = [] for traj in trajectories: - completion = '' - for msg in reversed(traj.get('messages', []) or []): - if msg.get('role') == 'assistant': - completion = msg.get('content', '') or '' - break - rewards.append(1.0 if extract_code(completion).strip() else 0.0) + rewards.append(1.0 if unwrap_code(assistant_text(traj)).strip() else 0.0) return rewards diff --git a/cookbook/rl/grpo/mbpp_grpo.py b/cookbook/rl/grpo/mbpp_grpo.py index 1c75d927e..25673f563 100644 --- a/cookbook/rl/grpo/mbpp_grpo.py +++ b/cookbook/rl/grpo/mbpp_grpo.py @@ -18,7 +18,6 @@ """ import json import os -import re import resource import shutil import signal @@ -45,6 +44,8 @@ from twinkle.processor import InputProcessor from twinkle.reward.base import Reward from twinkle.sampler import vLLMSampler +from twinkle_agentic.utils.code_utils import unwrap_code +from twinkle_agentic.utils.message_utils import assistant_text logger = get_logger() args = CLI.from_args() @@ -76,17 +77,6 @@ SYSTEM_PROMPT = ('You are an expert Python programmer. Write a complete, self-contained ' 'solution in a single ```python code block. Do not include tests.') -_FENCE_RE = re.compile(r'```(?:python|py)?\s*\n(.*?)```', re.S) - - -# ========== Text handling ========== -def extract_code(text: str) -> str: - """Take the last fenced block; fall back to the whole body when unfenced.""" - idx = (text or '').rfind('</think>') - body = text[idx + len('</think>'):] if idx >= 0 else (text or '') - blocks = _FENCE_RE.findall(body) - return (blocks[-1] if blocks else body).strip() - # ========== Sandbox ========== def run_asserts(code: str, setup: str, asserts: List[str], timeout: int = TEST_TIMEOUT) -> bool: @@ -154,12 +144,7 @@ def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: break if payload is None: continue - completion = '' - for msg in reversed(traj.get('messages', []) or []): - if msg.get('role') == 'assistant': - completion = msg.get('content', '') or '' - break - jobs.append((i, extract_code(completion), payload)) + jobs.append((i, unwrap_code(assistant_text(traj)), payload)) if not jobs: return rewards @@ -182,12 +167,7 @@ class MbppFormatReward(Reward): def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: rewards = [] for traj in trajectories: - completion = '' - for msg in reversed(traj.get('messages', []) or []): - if msg.get('role') == 'assistant': - completion = msg.get('content', '') or '' - break - rewards.append(1.0 if extract_code(completion).strip() else 0.0) + rewards.append(1.0 if unwrap_code(assistant_text(traj)).strip() else 0.0) return rewards diff --git a/cookbook/rsi/agentic/challenge.py b/cookbook/rsi/agentic/challenge.py index 7ce45782c..b9ddaefc0 100644 --- a/cookbook/rsi/agentic/challenge.py +++ b/cookbook/rsi/agentic/challenge.py @@ -61,19 +61,24 @@ from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Tuple -import numpy as np from twinkle import DeviceMesh, get_logger from twinkle.data_format import SamplingParams from twinkle.sampler import vLLMSampler -from twinkle_agentic.challenger import KeywordStore, parse_check_script, parse_problem_statement +from twinkle_agentic.challenger import (ApiExplorer, ApiModel, KeywordBank, KeywordPrompts, KeywordStore, + parse_check_script, parse_problem_statement) from twinkle_agentic.challenger.agentic import brittle_check_reason -from twinkle_agentic.challenger.code import split_keyword_list from twinkle_agentic.challenger.task_bank import TaskBank from twinkle_agentic.rollout import MultiTurnRollout from twinkle_agentic.tools.tool_manager import ToolManager +from twinkle_agentic.utils.message_utils import assistant_text sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +# The parent too: recorder.py is shared with the code half, which is a sibling +# directory rather than a package -- both halves are run as scripts. +sys.path.insert(1, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) import prompts as P # noqa: E402 +import train as T # noqa: E402 +from recorder import Recorder # noqa: E402 from sandbox import open_pool, solver_harness # noqa: E402 logger = get_logger() @@ -94,12 +99,6 @@ PASS_RATE_TARGET = 0.2 PASS_RATE_WIDTH = 0.3 -# How many phrases the keyword prompt's 'do not repeat these' line may quote. -# Measured on armA2ser: with 130 quoted the eighth refill call was still answering -# normally, with 150 it started inventing -- 'iRAPION holistic replace', 10 of 480 -# phrases that run. 100 sits below where that began. -AVOID_TOTAL = 100 - # How often run() looks for a stall. Only ever reached when the run has already # gone quiet, so it costs one wakeup per interval and nothing else. STALL_CHECK_SECONDS = 30.0 @@ -281,7 +280,13 @@ def parse_args(): 'the update was rounded away by the bf16 round trip through ' 'the checkpoint anyway -- after 12 such iterations 98.54%% of ' 'the weights were still bit-identical to the base model.') - p.add_argument('--sides', default='both', choices=('both', 'propose', 'solve')) + p.add_argument('--sides', default='both', + help="which sides to train, comma-separated: 'both' is the " + "agentic pair (propose + solve), 'code' is the code half, " + "and 'both,code' runs all three into one step. Each name " + 'also decides whether that task source is collected at ' + 'all, so this is one switch rather than two that can ' + 'disagree about what an iteration contains.') p.add_argument('--micro-batch-size', type=int, default=1, help='trajectories per micro batch. One, because padding_free is ' 'off: a micro batch is padded to its longest member, so ' @@ -333,40 +338,104 @@ def parse_args(): p.add_argument('--swanlab-mode', default='online', help="'disabled' keeps a run off the dashboard entirely.") + # The code half. Only read when --sides names 'code'; see code/collect.py. + # Its own prefix throughout, because every one of these has an agentic + # counterpart that means something else: --solver-rollouts is a 24-turn + # sandbox episode and --code-solver-rollouts is one message of python. + p.add_argument('--code-keep-target', type=int, default=8, + help='problems to keep per iteration, i.e. GRPO groups: one ' + 'problem is one prompt answered --code-solver-rollouts ' + 'times.') + p.add_argument('--code-batch-size', type=int, default=0, + help='problems per written batch; 0 is one batch of the target.') + p.add_argument('--code-solver-rollouts', type=int, default=8, + help='attempts per candidate. The group size and the denominator ' + 'of n_pass in one number: these attempts ARE what the code ' + 'side trains on, so the difficulty stage is not overhead.') + p.add_argument('--code-keep-pass-band', type=int, nargs=2, default=(1, 7), + metavar=('LOW', 'HIGH'), + help='keep a problem solved this many times out of ' + '--code-solver-rollouts, inclusive. Also what guarantees ' + 'the group has a gradient, so it has to be read against ' + 'the rollout count: (1,7) is the band for 8.') + p.add_argument('--code-max-proposals-per-round', type=int, default=2000, + help='ceiling on one proposing round, i.e. one batched generate.') + p.add_argument('--code-propose-temp', type=float, default=1.1) + p.add_argument('--code-propose-max-tokens', type=int, default=8192) + p.add_argument('--code-solver-temp', type=float, default=1.0) + p.add_argument('--code-solver-max-tokens', type=int, default=2048) + p.add_argument('--code-problem-max-chars', type=int, default=4000) + p.add_argument('--code-max-checks', type=int, default=6) + p.add_argument('--code-script-timeout', type=int, default=30, + help='seconds one script gets. A local subprocess, not a ' + 'sandbox slot: a judgement is milliseconds and the stage ' + 'makes candidates x rollouts of them, which through a ' + 'microVM would cost more than the rest of the iteration.') + p.add_argument('--code-seed-file', default='', + help='seed jsonl with query [+ code], from prepare.py.') + p.add_argument('--code-seed-mix-prob', type=float, default=0.5) + p.add_argument('--code-no-two-step', action='store_true', + help='never take the two-call path, even for seeds with code.') + p.add_argument('--code-keywords-n', type=int, default=128, + help='per-category refill target; 0 disables the keyword bank.') + # Output. p.add_argument('--random-seed', type=int, default=0) args = p.parse_args() - if not args.api_model or not args.api_base: - raise SystemExit('[rsi] --api-model and --api-base are required ' - '(or LLM_BACKUP_MODEL / LLM_BACKUP_BASE_URL)') if not args.tag: raise SystemExit('[rsi] --tag is required: it decides which run these ' 'iterations belong to and which checkpoint they overwrite') - if args.solver_rollouts < 2: - raise SystemExit('[rsi] --solver-rollouts must be >= 2: it is both the ' - "solver side's GRPO group size and the denominator n_pass " - 'is judged against') - if args.group_size < 2: - raise SystemExit('[rsi] --group-size must be >= 2: a group of one has ' - 'no mean to subtract, so every advantage is zero') - # Checked here rather than where the pool is opened, which is after the model - # and the sampler are up: that is six minutes of startup to find out that a - # host address is missing. - if not args.sandbox_api_url: - raise SystemExit('[rsi] --sandbox-api-url is required (or SANDBOX_API_URL / ' - 'AENV_API_URL)') - if not os.environ.get('E2B_API_KEY') and not os.environ.get('AENV_API_KEY'): - raise SystemExit('[rsi] E2B_API_KEY is required: the sandbox client reads ' - 'it from the environment') - os.environ.setdefault('AENV_API_URL', args.sandbox_api_url) - os.environ.setdefault('AENV_TEMPLATE', args.sandbox_template) - os.environ.setdefault('AENV_API_KEY', os.environ.get('E2B_API_KEY', '')) + # Parsed here rather than left to argparse choices, which cannot express a + # comma-separated set. A typo has to stop the run: train.py counts an unknown + # side as 'not requested' and would take a whole iteration to say so. + args.sides_list = T.sides_wanted(args.sides) + unknown = [s for s in args.sides_list if s not in ('propose', 'solve', 'code')] + if unknown or not args.sides_list: + raise SystemExit(f'[rsi] --sides {args.sides!r} names {unknown or "nothing"}; ' + f"it takes 'both', 'propose', 'solve' and 'code', " + f"comma-separated (e.g. 'both,code')") + # And parsed before everything below it, because most of what follows is one + # half's requirements: an API model and a sandbox host are what the agentic + # half needs to invent and check a task, and --sides code neither calls the API + # nor opens a microVM. Demanding them anyway is a run refused over a resource + # it was never going to touch. + if 'propose' in args.sides_list or 'solve' in args.sides_list: + if not args.api_model or not args.api_base: + raise SystemExit('[rsi] --api-model and --api-base are required ' + '(or LLM_BACKUP_MODEL / LLM_BACKUP_BASE_URL)') + if args.solver_rollouts < 2: + raise SystemExit('[rsi] --solver-rollouts must be >= 2: it is both the ' + "solver side's GRPO group size and the denominator " + 'n_pass is judged against') + if args.group_size < 2: + raise SystemExit('[rsi] --group-size must be >= 2: a group of one has ' + 'no mean to subtract, so every advantage is zero') + # Checked here rather than where the pool is opened, which is after the + # model and the sampler are up: that is six minutes of startup to find out + # that a host address is missing. + if not args.sandbox_api_url: + raise SystemExit('[rsi] --sandbox-api-url is required (or ' + 'SANDBOX_API_URL / AENV_API_URL)') + if not os.environ.get('E2B_API_KEY') and not os.environ.get('AENV_API_KEY'): + raise SystemExit('[rsi] E2B_API_KEY is required: the sandbox client ' + 'reads it from the environment') + os.environ.setdefault('AENV_API_URL', args.sandbox_api_url) + os.environ.setdefault('AENV_TEMPLATE', args.sandbox_template) + os.environ.setdefault('AENV_API_KEY', os.environ.get('E2B_API_KEY', '')) + if 'code' in args.sides_list and args.code_solver_rollouts < 2: + raise SystemExit('[rsi] --code-solver-rollouts must be >= 2: those attempts ' + 'are the code side\'s GRPO group, and a group of one has ' + 'no mean to subtract') # The bank and the keyword store belong to the loop, not to an iteration: # comparing iteration k+1's proposals against what k produced is the point of # them. ``--task-bank ''`` turns novelty off and leaves the pass-rate gaussian # alone. out_dir is set per iteration by rsi.py. root = os.path.join(args.root, args.tag) args.keyword_db = os.path.join(root, 'keywords.jsonl') + # The code half draws from a bank of its own: its categories are + # algorithm/computer/noncs against the agentic half's + # transform/domain/edge_case, and one file cannot hold both. + args.code_keyword_db = os.path.join(root, 'code_keywords.jsonl') if args.task_bank is None: args.task_bank = os.path.join(root, 'task_bank.jsonl') args.out_dir = root @@ -430,188 +499,10 @@ def rollout_one(rollout: MultiTurnRollout, traj: Dict[str, Any], called from every sandbox thread at once and the requests share vLLM's batch without any of them waiting for the others to be ready. """ - out = rollout([traj], sampling_params=params, tool_manager=slot.tool_manager) + out = rollout([traj], sampling_params=params, tool_manager=slot.tool_manager()) return out[0] if out else None -def call_one(slot, script: str) -> Tuple[int, str]: - """Run one python script inside ``slot``; returns (exit code, output).""" - return slot.run(script) - - -def api_one(api, messages: List[Dict[str, Any]], user_text: str, - params: SamplingParams, extra_body: Optional[Dict[str, Any]] = None) -> Optional[str]: - """Append ``user_text`` and one API reply to ``messages``; returns the reply. - - ``messages`` is the caller's private copy, never a trainable trajectory, so - mutating it in place costs the model nothing. ``None`` means the call raised: - the caller rejects rather than building a task on a broken conversation. - - Tools are withdrawn for these stages on purpose -- they are answers, not - actions -- so only the text is kept and any structured ``tool_calls`` the API - returned are dropped. - """ - messages.append({'role': 'user', 'content': user_text}) - request = {'messages': messages} - try: - reply = api(request, params, extra_body=extra_body) if extra_body else api(request, params) - except Exception as e: # noqa: BLE001 -- one bad call must not kill the run - logger.warning(f'[challenge] API call failed: {type(e).__name__}: {e}') - return None - if isinstance(reply, list): - reply = reply[0] if reply else {} - content = (reply.get('content') if isinstance(reply, dict) else None) or '' - messages.append({'role': 'assistant', 'content': content}) - return content - - -# ── Output ───────────────────────────────────────────────────────────────── - - -def logprob_column(logprobs: Any) -> List[float]: - """One float per generated token: the logprob of the token that was chosen. - - The sampler hands these over as ``List[List[Tuple[int, float]]]`` -- per - generated token, a list of top-k ``(token_id, logprob)`` pairs with the chosen - token first (``SampledSequence.logprobs``, data_format/sampling.py:185). - Passing that to ``np.asarray`` directly would store an ``(N, k, 2)`` array and - the loader would hand GRPO nested lists where it wants one float per trainable - token -- which is a crash inside the step, or worse a silent reshape. - - A plain list of floats is accepted too, for a sampler that already flattened. - Anything else raises rather than being coerced: a wrong ``old_logps`` makes the - GRPO ratio wrong on the first step, and nothing downstream would say so. - """ - out: List[float] = [] - for step in logprobs: - if isinstance(step, (int, float)): - out.append(float(step)) - continue - if isinstance(step, (list, tuple)) and step: - head = step[0] - if isinstance(head, (list, tuple)) and len(head) >= 2: - out.append(float(head[1])) - continue - raise TypeError(f'cannot read a logprob out of {step!r}; expected a float ' - f'or a list of (token_id, logprob) pairs') - return out - - -class Recorder: - """Everything a run writes, behind one lock. - - Trajectories go to ``.npz`` for the token fields and to ``index.jsonl`` for - everything a reader needs to interpret them. The text is written in full and - never truncated: these files are read to check whether a reward was deserved, - which a shortened statement cannot answer. - """ - - def __init__(self, out_dir: str): - self.dir = out_dir - self.traj_dir = os.path.join(out_dir, 'trajs') - os.makedirs(self.traj_dir, exist_ok=True) - self._lock = threading.Lock() - self._n = 0 - self._index = open(os.path.join(self.traj_dir, 'index.jsonl'), 'w', encoding='utf-8') - self._groups = open(os.path.join(out_dir, 'groups.jsonl'), 'w', encoding='utf-8') - self._tasks = open(os.path.join(out_dir, 'tasks.jsonl'), 'w', encoding='utf-8') - # Why a build produced no task. The reason alone is not diagnosable: nine - # empty_workspace rejections in one run all looked like the model refusing - # to act, and the question of whether it had run out of tokens or simply - # emitted no call could not be answered from the record, because the fields - # that answered it were on the trajectory and were dropped. - self._rejected = open(os.path.join(out_dir, 'rejected.jsonl'), 'w', encoding='utf-8') - # Keyword replies, both sides in full. The one question this file exists to - # answer -- did the model disobey the format, or does the parser reject what - # it produced -- cannot be answered from a count. Keyword generation was - # silently broken for whole runs when the prompt asked for one per line and - # the parser wanted a JSON array. - self._keywords = open(os.path.join(out_dir, 'keyword_gen.jsonl'), 'w', encoding='utf-8') - # Every solver attempt, passed or not, with the state it left and what the - # check said about it. A task measured at 0 of 8 has three explanations -- - # the check is wrong, the statement withholds something the check demands, - # or the solver gave up -- and only the attempt and the workspace it left - # tell them apart. Written for every attempt, not only for the ones that - # end up trained on: the failures are what this file is for. - self._attempts = open(os.path.join(out_dir, 'solver_attempts.jsonl'), 'w', - encoding='utf-8') - # The rubric, all three of its dimensions. Only novelty reaches a reward; - # usefulness and complexity are recorded so the question of whether they - # should count can be answered from a run instead of argued. - self._novelty = open(os.path.join(out_dir, 'novelty_scores.jsonl'), 'w', - encoding='utf-8') - - def trajectory(self, traj: Dict[str, Any], **fields: Any) -> None: - """One training sample: token fields to npz, everything else to the index. - - A trajectory with no ``logprobs`` is written anyway, with the field left - null. It is not trainable and the loader will say so -- which is the point: - a sample silently dropped here would make the group it belongs to look like - a different size than it was. - """ - input_ids = np.asarray(traj.get('input_ids') or [], dtype=np.int32) - labels = np.asarray(traj.get('labels') or [], dtype=np.int32) - logprobs = traj.get('logprobs') - with self._lock: - self._n += 1 - name = f'{self._n:06d}.npz' - arrays = {'input_ids': input_ids, 'labels': labels} - if logprobs is not None: - # float64, and the chosen token's column only. These are the old_logps a - # GRPO step divides by; float32 would round them to about 7 digits, so - # the ratio exp(logp - old_logp) would be off by roughly 1e-7 for - # reasons that have nothing to do with the policy having changed. - arrays['logprobs'] = np.asarray(logprob_column(logprobs), dtype=np.float64) - # Compressed: a 24-turn agentic episode is tens of thousands of token ids, - # and 128 of them per iteration adds up on disk. - np.savez_compressed(os.path.join(self.traj_dir, name), **arrays) - record = dict(fields) - record.update({ - 'npz': name, - 'n_tokens': int(input_ids.size), - 'n_trainable': int((labels != -100).sum()) if labels.size else 0, - 'has_logprobs': logprobs is not None, - # The rollout guarantees one logprob per trainable label; recorded so a - # loader can check it rather than trust it. - 'n_logprobs': int(arrays['logprobs'].size) if logprobs is not None else 0, - 'turns': traj.get('turns'), - 'stop_reason': traj.get('stop_reason'), - 'truncated': bool(traj.get('truncated')), - 'tool_stop': traj.get('tool_stop'), - 'messages': traj.get('messages') or [], - }) - self._write(self._index, record) - - def group(self, record: Dict[str, Any]) -> None: - self._write(self._groups, record) - - def task(self, record: Dict[str, Any]) -> None: - self._write(self._tasks, record) - - def rejected(self, record: Dict[str, Any]) -> None: - self._write(self._rejected, record) - - def keywords(self, record: Dict[str, Any]) -> None: - self._write(self._keywords, record) - - def attempt(self, record: Dict[str, Any]) -> None: - self._write(self._attempts, record) - - def novelty(self, record: Dict[str, Any]) -> None: - self._write(self._novelty, record) - - def close(self) -> None: - for handle in (self._index, self._groups, self._tasks, self._rejected, - self._keywords, self._attempts, self._novelty): - handle.close() - - def _write(self, handle, record: Dict[str, Any]) -> None: - line = json.dumps(record, ensure_ascii=False, default=str) - with self._lock: - handle.write(line + '\n') - handle.flush() - - # ── Group state ──────────────────────────────────────────────────────────── @@ -771,6 +662,10 @@ def __init__(self, args, sampler, template, slots: List[Any], recorder: Recorder max_turns=args.solver_max_turns, stop_after_stuck_turns=args.stop_after_stuck_turns, sampling_params=self.solve_params) + # The keyword bank's local fallback, below. ``max_turns=1`` because + # brainstorming a list is a text round: the trajectory ends before a tool + # could be dispatched, and a bracketed list in a reply is exactly what a + # tool-calling rollout would try to run. self.keyword_rollout = MultiTurnRollout( sampler, template=template, tool_manager=ToolManager(), max_turns=1, sampling_params=SamplingParams(max_tokens=args.keyword_max_tokens, @@ -778,10 +673,14 @@ def __init__(self, args, sampler, template, slots: List[Any], recorder: Recorder temperature=args.keyword_temp, top_p=0.98)) from twinkle_agentic.protocol.openai import OpenAI - self.api = OpenAI(model=args.api_model, api_key=args.api_key or None, - base_url=args.api_base) + # Kept as its own attribute as well as inside the model: the novelty rubric + # goes out through score_tasks rather than through here, and has to send the + # same body or the two API paths would be capped differently. self.api_extra = ({'thinking_budget': args.api_thinking_budget} if args.api_thinking_budget > 0 else None) + self.api = ApiModel(OpenAI(model=args.api_model, api_key=args.api_key or None, + base_url=args.api_base), + extra_body=self.api_extra, name='challenge') self.check_params = SamplingParams(max_tokens=args.check_max_tokens, num_samples=1, temperature=args.propose_temp, top_p=0.95) self.problem_params = SamplingParams(max_tokens=args.problem_max_tokens, @@ -801,7 +700,29 @@ def __init__(self, args, sampler, template, slots: List[Any], recorder: Recorder # different one would be a different experiment. self.system = P.SYSTEM + (P.BUILD_SIZE_CAP.format(n=args.max_build_files) if args.max_build_files > 0 else '') - self.store = KeywordStore(args.keyword_db, P.CATEGORIES) + # The keyword cycle -- draw, refill, the avoid list, expansion, the bank on + # disk -- is the framework's, not a third copy of it here. What is local is + # only which model answers: the API model, because keyword text never enters + # a trajectory (it is parsed into a list and thrown away) so no untrained + # tokens come of it, and because the bank is the single input every task + # downstream is built from -- measured over the 1344 keywords iterations 1-7 + # generated locally at temperature 1.3, 31% of transform named an activity + # on a running system rather than a computation, 13% of domain named an + # operation rather than material, and 24% of edge_case needed hardware the + # container does not have. The one-turn local rollout is the fallback rather + # than nothing, because an unreachable API must not leave a category dry. + self.keywords = KeywordBank( + KeywordStore(args.keyword_db, P.CATEGORIES), + prompts=KeywordPrompts(system=P.KEYWORD_SYSTEM, user=P.KEYWORD_USER, + expand_user=P.KEYWORD_EXPAND_USER), + category_desc=P.CATEGORY_DESC, + explorer=ApiExplorer(self.api, params=self.keyword_api_params, + fallback=self.keyword_rollout), + rng=self.rng, name='challenge', sink=self.rec.keywords, + # Every prompt here quotes one keyword per category, so a draw that + # covered a subset would send a prompt this run's prompts cannot fill. + single_kw_prob=0.0, refill_target=args.keywords_n, + gen_calls=args.keyword_gen_calls, refill_tries=args.keyword_refill_tries) self.bank = TaskBank(args.task_bank, refs=args.task_bank_refs) if args.task_bank else None # ms-agent builds the solver's opening messages, and it does so through a # stateful agent object -- so one instance, one lock, and only for the few @@ -813,16 +734,11 @@ def __init__(self, args, sampler, template, slots: List[Any], recorder: Recorder self.api_pool = ThreadPoolExecutor(max_workers=args.api_concurrency, thread_name_prefix='api') self.state = threading.Lock() - self.kw_lock = threading.Lock() # Jobs actually being worked on right now, sandbox and API. Only used by # the stall check in run(): 'the queue is empty' is not 'there is nothing # left to do' while a thread is still inside a job that will queue more. self.busy = 0 self.api_jobs = 0 - self.nonce = 0 - # (category, keyword) pairs behind tasks nobody solved. Read at the end by - # expand_hard_keywords, which asks for more in the same domains. - self.hard: List[Tuple[str, str]] = [] self.kept: List[Group] = [] self.groups: List[Group] = [] self.n_launched = 0 @@ -835,186 +751,6 @@ def bump(self, key: str, n: int = 1) -> None: with self.state: self.counts[key] = self.counts.get(key, 0) + n - def draw_keywords(self) -> Tuple[List[Tuple[str, str]], str]: - """One entry from each category, refilling any that has run dry. - - On its own lock, not ``state``: a refill is eight model calls and holding - the lock every ``bump`` needs for that long would stall all 32 slots. Two - threads drawing at once still have to take turns, or the second refill's - prompt would not know what the first one had just said. - """ - with self.kw_lock: - for category in P.CATEGORIES: - if not self.store.unused(category): - self.refill(category) - picks = [] - for category in P.CATEGORIES: - text = self.store.take(category, self.rng) - if text is not None: - picks.append((category, text)) - return picks, '\n'.join(f'- {c}: {t}' for c, t in picks) - - def refill(self, category: str) -> None: - """Ask the local model for more keywords in ``category``. - - Says so when it comes back empty. A silent no-op here is the worst outcome - available: every proposal then falls back to a keyword-less prompt and the - run looks normal while producing one identical prompt over and over. That - is exactly what happened for whole runs when the prompt asked for one - keyword per line and the parser wanted a JSON array. - - The calls run one at a time so each can be told what the ones before it - said; each is answered by a rollout with ``max_turns=1``, which ends the - trajectory before any tool could be dispatched -- brainstorming a list is a - text round, and a bracketed list in a reply is exactly what a tool-calling - rollout would try to run. - """ - for attempt in range(1, max(1, self.args.keyword_refill_tries) + 1): - if self.generate_keywords(category): - return - logger.warning(f'[challenge] keyword refill for {category!r} produced ' - f'nothing new on try {attempt}') - if self.store.items[category]: - # Every keyword marked unused again. The alternative is a category that - # can never be drawn from, which stops the run: a repeat draw is worse - # than no run only if diversity matters more than collecting anything. - self.store.recycle(category) - self.store.save() - logger.warning(f'[challenge] keyword category {category!r} exhausted -> ' - f'recycled {len(self.store.items[category])} topics') - - def generate_keywords(self, category: str) -> bool: - """One refill round. True when it added something the bank did not have.""" - want = self.args.keywords_n - calls = max(1, self.args.keyword_gen_calls) - per_call = max(1, -(-want // calls) + 4) - known = self.store.texts(category) - fresh: List[str] = [] - seen = {t.strip().lower() for t in known} - for i in range(calls): - # Newest first: the calls run one at a time so each can avoid what the - # ones before it said, and letting older entries evict those would undo - # it. Past the cap the oldest of this refill's phrases fall off, which - # is also the least costly thing to drop. - avoid = (fresh + known)[:AVOID_TOTAL] - self.nonce += 1 - user = (P.KEYWORD_USER.format(k=per_call, desc=P.CATEGORY_DESC[category]) - + ('\nDo NOT repeat any of these already-used topics: ' - + ', '.join(avoid) if avoid else '') - + f'\n(batch {self.nonce}-{i})') - # Asked of the API model rather than the local one. Keyword text never - # enters a trajectory -- it is parsed into a list and thrown away -- so - # this adds no untrained tokens, which is the rule that decides what may - # use the API. And the bank is the single input every task downstream is - # built from: measured over the 1344 keywords iterations 1-7 generated - # locally at temperature 1.3, against category rules the model is shown - # in full, 31% of transform named an activity on a running system rather - # than a computation, 13% of domain named an operation rather than - # material, and 24% of edge_case needed hardware the container does not - # have. Downstream, 42-70% of statements described themselves as - # simulating their own subject matter, which is what a keyword the - # sandbox cannot honour turns into. A 4B policy at that temperature is - # the wrong instrument for a constraint list this long. - # - # Falls back to the local model instead of giving up: an unreachable API - # must not leave a category dry, because dry means keyword-less prompts - # and a run that looks healthy while producing one prompt over and over - # -- the exact failure the refill logic already guards against. - out = None - messages = [{'role': 'system', 'content': P.KEYWORD_SYSTEM}] - reply = api_one(self.api, messages, user, self.keyword_api_params, - self.api_extra) - via = 'api' - if reply is None: - traj = {'messages': [{'role': 'system', 'content': P.KEYWORD_SYSTEM}, - {'role': 'user', 'content': user}]} - out = self.keyword_rollout([traj]) - reply = self._assistant_text(out[0] if out else {}) - via = 'local-fallback' - parsed, dropped_long = split_keyword_list(reply) - new = [k for k in parsed if k.lower() not in seen] - for keyword in new: - seen.add(keyword.lower()) - fresh.extend(new) - self.rec.keywords({'category': category, 'prompt': user, 'reply': reply, - 'parsed': parsed, 'n_parsed': len(parsed), - 'n_new': len(new), 'via': via, - 'dropped_long': dropped_long, - 'n_dropped_long': len(dropped_long), - 'stop_reason': (out[0].get('stop_reason') if out else None), - 'truncated': bool(out[0].get('truncated')) if out else None}) - if len(fresh) >= want: - break - added = self.store.add(category, fresh[:want], source='gen') - if added: - # Written now rather than at the end of the run: a run that crashes after - # spending eight model calls on keywords should not have to spend them - # again, and the next iteration reads this file to know what was used. - self.store.save() - logger.info(f'[challenge] keywords {category!r} +{added}') - return bool(added) - - @staticmethod - def _assistant_text(traj: Dict[str, Any]) -> str: - for message in reversed((traj.get('messages') if traj else []) or []): - if message.get('role') == 'assistant': - return message.get('content') or '' - return '' - - def expand_hard_keywords(self) -> int: - """More keywords in the domains that produced tasks nobody solved. - - Run once at the end, so what it adds is there for the next iteration rather - than for the groups still in flight. One call per hard keyword, capped at - 32 of them: this is the only feedback the keyword bank gets from difficulty, - and without it the bank drifts wherever the refill prompt happens to go. - """ - with self.state: - hard = list(self.hard)[:32] - if not hard: - return 0 - added = 0 - for i, (category, keyword) in enumerate(hard): - self.nonce += 1 - user = (P.KEYWORD_EXPAND_USER.format(kw=keyword, m=8, - desc=P.CATEGORY_DESC[category]) - + f'\n(batch {self.nonce}-{i})') - # Same reasoning as the refill path: the API model, falling back to the - # local one. This path matters more, not less -- it wrote 960 of the 1344 - # keywords the first seven iterations banked, at more than twice their - # rule-break rate, so it is the one shaping what later iterations draw. - out = None - messages = [{'role': 'system', 'content': P.KEYWORD_SYSTEM}] - reply = api_one(self.api, messages, user, self.keyword_api_params, - self.api_extra) - via = 'api' - if reply is None: - traj = {'messages': [{'role': 'system', 'content': P.KEYWORD_SYSTEM}, - {'role': 'user', 'content': user}]} - out = self.keyword_rollout([traj]) - reply = self._assistant_text(out[0] if out else {}) - via = 'local-fallback' - parsed, dropped_long = split_keyword_list(reply) - added += self.store.add(category, parsed, source='expand', parent=keyword) - # The prompt goes in whole, as the refill path already does. It used to - # record the literal string 'expand', which made this file unable to - # answer the one question it gets asked -- whether a change to - # KEYWORD_EXPAND_USER was live in a given iteration -- and cost an - # afternoon to a wrong answer inferred from mtimes instead. - # - # ``dropped_long`` for the same reason one step further in: iteration 9 - # recorded n_parsed 0 on four of six expand calls whose replies were - # well-formed JSON, and nothing in this file said the phrases had been - # thrown away for length rather than never produced. - self.rec.keywords({'category': category, 'parent': keyword, 'prompt': user, - 'reply': reply, 'parsed': parsed, 'n_parsed': len(parsed), - 'dropped_long': dropped_long, - 'n_dropped_long': len(dropped_long), 'via': via}) - self.store.save() - logger.info(f'[challenge] expanded {len(hard)} hard keyword(s) -> ' - f'+{added} same-domain topics') - return added - def launch_group(self) -> Optional[Group]: """Draw a topic and queue its ``group_size`` builds. None once at the cap.""" with self.state: @@ -1024,7 +760,8 @@ def launch_group(self) -> Optional[Group]: return None gid = self.n_launched self.n_launched += 1 - picks, block = self.draw_keywords() + picks = self.keywords.draw() + block = KeywordBank.block(picks) if len(picks) != len(P.CATEGORIES): # Every proposal's prompt is the keyword draw, so there is no honest # prompt to send without one. Stopping is the reportable outcome; a @@ -1094,7 +831,7 @@ def record_rejection(self, prop: Proposal) -> None: if isinstance(m, dict) and m.get('role') == 'assistant'), 'n_tool_calls': sum(len(m.get('tool_calls') or []) for m in messages if isinstance(m, dict)), - 'last_assistant': self._assistant_text(traj), + 'last_assistant': assistant_text(traj), 'check': prop.check, }) @@ -1111,7 +848,7 @@ def build(self, prop: Proposal, slot) -> None: slot.clear() traj = {'messages': [{'role': 'system', 'content': self.system}, {'role': 'user', 'content': prop.group.prompt}], - 'tools': slot.schemas} + 'tools': slot.tools()} prop.traj = rollout_one(self.propose_rollout, traj, self.propose_params, slot) if prop.traj is None: prop.outcome = 'rollout_empty' @@ -1141,7 +878,7 @@ def build(self, prop: Proposal, slot) -> None: attempt = 0 while True: attempt += 1 - reply = api_one(self.api, messages, user_text, self.check_params, self.api_extra) + reply = self.api.reply(messages, user_text, self.check_params) if reply is None: prop.outcome, prop.detail = 'api_error', 'check-script call failed' return @@ -1162,7 +899,7 @@ def build(self, prop: Proposal, slot) -> None: # that pins a file's size or quotes a script's source passes for its # author and fails every correct reproduction. brittle = brittle_check_reason(script) - exit_code, output = (1, brittle) if brittle else call_one(slot, script) + exit_code, output = (1, brittle) if brittle else slot.run_script(script) if exit_code == 0: prop.check = script break @@ -1176,8 +913,7 @@ def build(self, prop: Proposal, slot) -> None: f'\n--- state after check ---\n{after}') return - reply = api_one(self.api, messages, P.PROBLEM_FOLLOWUP, self.problem_params, - self.api_extra) + reply = self.api.reply(messages, P.PROBLEM_FOLLOWUP, self.problem_params) if reply is None: prop.outcome, prop.detail = 'api_error', 'problem-statement call failed' return @@ -1211,10 +947,10 @@ def solve_job(self, prop: Proposal, slot) -> None: if not opening.get('tools'): # The harness only shapes messages -- its tool list is empty on # purpose -- so the schemas come from the slot that will run them. - opening['tools'] = slot.schemas + opening['tools'] = slot.tools() attempt = rollout_one(self.solve_rollout, opening, self.solve_params, slot) if attempt is not None: - exit_code, output = call_one(slot, prop.check) + exit_code, output = slot.run_script(prop.check) passed = exit_code == 0 # Read after the check, not before: the check is allowed to write, and # what a reader of a failed attempt needs is the workspace the check @@ -1387,15 +1123,11 @@ def _decide(self, group: Group) -> None: } for p in group.proposals], } self.rec.group(record) - # Keyword draws behind tasks nobody solved, for expand_hard_keywords. Taken - # from every decided group, kept or not: a task at n_pass=0 says the same - # thing about its keywords either way. + # Keyword draws behind tasks nobody solved, for the end-of-run expansion. + # Taken from every decided group, kept or not: a task at n_pass=0 says the + # same thing about its keywords either way. if any(p.statement and p.n_pass == 0 for p in group.proposals): - with self.state: - seen = {(c, t.lower()) for c, t in self.hard} - for category, text in group.keywords: - if (category, text.lower()) not in seen: - self.hard.append((category, text)) + self.keywords.remember_hard(group.keywords) if chosen is None: self.bump('group_dropped' if not group.dropped else 'group_dropped_early') return diff --git a/cookbook/rsi/agentic/episode.py b/cookbook/rsi/agentic/episode.py index df236c699..07a19a96c 100644 --- a/cookbook/rsi/agentic/episode.py +++ b/cookbook/rsi/agentic/episode.py @@ -18,9 +18,7 @@ from typing import Any, Dict, List, Optional, Tuple from twinkle import get_logger -from twinkle_agentic.envs import EnvTool from twinkle_agentic.harness import MsAgentHarness -from twinkle_agentic.tools.tool_manager import ToolManager from twinkle_agentic.verifier.result_check import (CheckContext, checks_from_dicts, run_checks) @@ -134,7 +132,7 @@ def build_episode(task: Dict[str, Any], cfg: SandboxConfig) -> Tuple[Any, Any, A # they are not, scores 0 for a reason that has nothing to do with the task. setup = task.get('setup_script') if setup: - exit_code, output = env.runner()(setup, 'python') + exit_code, output = env.run_script(setup) if exit_code != 0: raise RuntimeError(f'[{task.get("id")}] setup_script failed ' f'(exit {exit_code}): {output[-400:]}') @@ -143,9 +141,9 @@ def build_episode(task: Dict[str, Any], cfg: SandboxConfig) -> Tuple[Any, Any, A # The executor's own schemas, not the harness's (which are now empty by # construction). Advertising what will run is the whole point of sourcing # them from the sandbox. - schemas = env.tool_schemas() + schemas = env.tools() trajectory['tools'] = schemas - tool_manager = ToolManager(EnvTool.from_schemas(env, schemas)) + tool_manager = env.tool_manager(schemas) return harness, env, tool_manager, trajectory @@ -189,7 +187,7 @@ def score_episode(task: Dict[str, Any], env: RemoteMsAgentToolEnv, """ check_script = task.get('check_script') if check_script: - exit_code, output = env.runner()(check_script, 'python') + exit_code, output = env.run_script(check_script) if exit_code != 0: logger.debug(f'[{task.get("id")}] check_script failed (exit {exit_code}): ' f'{output[-200:]}') @@ -204,7 +202,7 @@ def score_episode(task: Dict[str, Any], env: RemoteMsAgentToolEnv, ctx = CheckContext( workspace=env.download_workspace(snapshot_dir), final_answer=final_answer, - runner=env.runner(), + env=env, ) report = run_checks(task['_checks'], ctx, mode=cfg.score_mode) if not report.all_passed: diff --git a/cookbook/rsi/agentic/remote_tool_env.py b/cookbook/rsi/agentic/remote_tool_env.py index 3ea21b2a3..2e7b502cd 100644 --- a/cookbook/rsi/agentic/remote_tool_env.py +++ b/cookbook/rsi/agentic/remote_tool_env.py @@ -28,7 +28,8 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple from twinkle import get_logger -from twinkle_agentic.envs.base import Env, StepResult +from twinkle.data_format.message import Tool as ToolInfo +from twinkle_agentic.envs.base import Env, StepResult, truncate_observation logger = get_logger() @@ -224,11 +225,20 @@ def ensure_ready(self) -> bool: """ if self.healthy(): return False + logger.warning('tool runtime unreachable; rebuilding the sandbox') + self.rebuild() + return True + + def rebuild(self) -> None: + """Kill this sandbox and boot a replacement, counting the recovery. + + A microVM is disposable, so there is nothing to repair: :meth:`reset` + already kills the old one and brings a fresh runtime up. All this adds is + the count, which is the part a run reports at the end. + """ self.n_recoveries += 1 - logger.warning(f'tool runtime unreachable; rebuilding the sandbox ' - f'(recovery #{self.n_recoveries})') + logger.warning(f'rebuilding the sandbox (recovery #{self.n_recoveries})') self.reset() - return True def step(self, tool_name: str, arguments: Dict[str, Any] = None) -> StepResult: return self.step_batch([(tool_name, arguments or {})])[0] @@ -240,6 +250,16 @@ def step_batch(self, calls: Sequence[Tuple[str, Dict[str, Any]]]) -> List[StepRe several, and it keeps ms-agent's own ``parallel_call_tool`` semantics rather than serialising what production would run concurrently. """ + return self._dispatch(calls, self._command_timeout) + + def _dispatch(self, calls: Sequence[Tuple[str, Dict[str, Any]]], + timeout: int) -> List[StepResult]: + """One request, with the per-call budget stated. See :meth:`step_batch`. + + Split out so :meth:`run_script` can name its own timeout -- a check with + a deadline of its own must not be judged by the budget a model turn was + given -- without restating the dispatch. + """ calls = list(calls) if not calls: return [] @@ -248,10 +268,10 @@ def step_batch(self, calls: Sequence[Tuple[str, Dict[str, Any]]]) -> List[StepRe 'tool_name': self._dispatch_name(name), 'arguments': args or {} } for name, args in calls], - 'timeout': self._command_timeout, + 'timeout': timeout, } try: - body = self._rpc('/call', payload) + body = self._rpc('/call', payload, timeout=timeout) results = body.get('results') or [] except Exception as e: # noqa # A dead sandbox must not kill the training step: report it as an @@ -349,7 +369,7 @@ def _dispatch_name(self, name: str) -> str: return name return self._short_to_full.get(name, name) - def tool_schemas(self) -> List[Dict[str, Any]]: + def tools(self) -> List[ToolInfo]: """Schemas from the runtime that will execute them. These go straight into the prompt. Sourcing them from the executor @@ -362,7 +382,7 @@ def tool_schemas(self) -> List[Dict[str, Any]]: def tool_names(self) -> List[str]: names = [] - for schema in self.tool_schemas(): + for schema in self.tools(): name = (schema.get('function') or {}).get('name') if name: names.append(str(name)) @@ -396,31 +416,31 @@ def resolve_tool(self, name: str) -> str: # ------------------------------------------------------- for the checker - def runner(self, shell_tool: str = 'shell_executor', python_tool: str = 'python_executor'): - """A ``result_check`` runner that executes inside this episode's sandbox. + def run_script(self, source: str, interpreter: str = 'python', + timeout: Optional[int] = None) -> Tuple[int, str]: + """Run a whole script inside this episode's sandbox. Verification has to see the filesystem the agent actually wrote to, so the check goes back through the same tools rather than a local subprocess. Those tools return prose, not an exit status, so the command is made to print a marker and the status is read back out of the output. """ - shell_name = self.resolve_tool(shell_tool) - python_name = self.resolve_tool(python_tool) - - def _run(source: str, interpreter: str) -> Tuple[int, str]: - if interpreter == 'python': - code = _PY_WRAPPER.format(body=repr(source), mark=_RC_MARK) - out = self.step(python_name, {'code': code}).observation - else: - out = self.step(shell_name, {'command': f'{source}\necho "{_RC_MARK}:$?"'}).observation - match = _RC_RE.search(out or '') - if match is None: - # No marker means the tool itself failed (timeout, sandbox down) - # rather than the check failing; report non-zero and keep output. - return 1, out or 'check produced no output and no exit marker' - return int(match.group(1)), _RC_RE.sub('', out or '').strip() - - return _run + seconds = timeout or self._command_timeout + if interpreter == 'python': + code = _PY_WRAPPER.format(body=repr(source), mark=_RC_MARK) + call = (self.resolve_tool('python_executor'), {'code': code}) + elif interpreter in ('shell', 'bash'): + call = (self.resolve_tool('shell_executor'), + {'command': f'{source}\necho "{_RC_MARK}:$?"'}) + else: + return 1, f'unsupported interpreter {interpreter!r}; use python or shell' + out = self._dispatch([call], seconds)[0].observation + match = _RC_RE.search(out or '') + if match is None: + # No marker means the tool itself failed (timeout, sandbox down) + # rather than the check failing; report non-zero and keep output. + return 1, out or 'check produced no output and no exit marker' + return int(match.group(1)), _RC_RE.sub('', out or '').strip() def download_workspace(self, dest: str, max_files: int = 200, max_bytes: int = 1 << 20) -> str: """Copy the episode's files out of the sandbox for the ``file_*`` checks. @@ -429,7 +449,7 @@ def download_workspace(self, dest: str, max_files: int = 200, max_bytes: int = 1 interface for a generic verifier but cannot see inside a microVM. The episode is over by the time this runs, so a snapshot is equivalent to the live filesystem -- and the shell/python checks still go through - :meth:`runner`, against the sandbox itself. + :meth:`run_script`, against the sandbox itself. Files above ``max_bytes`` are skipped: a check that needs to look at a 100MB artifact wants a command, not a copy. @@ -607,6 +627,4 @@ def _rpc(self, path: str, payload: Optional[Dict[str, Any]], timeout: Optional[i def _truncate(self, text: str) -> str: limit = self.max_observation_chars - if limit and len(text) > limit: - return f'{text[:limit]}\n...[truncated {len(text) - limit} chars]' - return text + return truncate_observation(text, limit) if limit else text diff --git a/cookbook/rsi/agentic/rsi.py b/cookbook/rsi/agentic/rsi.py index 5812f3988..224cea407 100644 --- a/cookbook/rsi/agentic/rsi.py +++ b/cookbook/rsi/agentic/rsi.py @@ -51,8 +51,17 @@ from twinkle.checkpoint_engine import CheckpointEngineManager sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +_RSI = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +# recorder.py sits one level up, shared with the code half, and the code half +# itself is a sibling directory. Appended rather than inserted, and behind the +# agentic directory on purpose: both halves have a challenge.py, and the one this +# process means by that name is the agentic one. +sys.path.insert(1, _RSI) +sys.path.append(os.path.join(_RSI, 'code')) import challenge as C # noqa: E402 +import collect as CODE # noqa: E402 import train as T # noqa: E402 +from recorder import Recorder # noqa: E402 from sandbox import close_pool # noqa: E402 logger = get_logger() @@ -71,15 +80,14 @@ def next_iteration(root: str) -> int: return i -def collect_once(args, sampler, template, slots, out_dir: str) -> Dict[str, Any]: - """One collection pass into ``out_dir``; returns its metrics. +def collect_agentic(args, sampler, template, slots, recorder: Recorder, + out_dir: str) -> Dict[str, Any]: + """The agentic half of one collection pass; returns its metrics. The body of what challenge.py's main() did, minus the resources: the sampler, - the template and the sandbox pool are owned by the caller and outlive this. + the template, the sandbox pool and the recorder are owned by the caller and + outlive this. """ - os.makedirs(out_dir, exist_ok=True) - args.out_dir = out_dir - recorder = C.Recorder(out_dir) run = C.Run(args, sampler, template, slots, recorder) started = time.time() try: @@ -87,18 +95,18 @@ def collect_once(args, sampler, template, slots, out_dir: str) -> Dict[str, Any] # After the loop, not during: what it adds is for the next iteration, and # doing it here means a crash in collection does not also lose the bank. if args.keyword_expand: - run.expand_hard_keywords() + run.keywords.expand_hard() finally: - recorder.close() - run.store.save() + run.keywords.save() # A Run per iteration means a thread pool per iteration. close_pool cannot # do this because the sandbox pool is the one thing that is not per-Run. run.api_pool.shutdown(wait=False) if run.bank is not None: logger.info(f'[rsi] task bank: {run.bank.stats()}') # In the finally block because a run that crashed is the one whose numbers - # are most worth having, and after recorder.close() so groups.jsonl is - # flushed before collect_metrics reads it back. + # are most worth having. Reading groups.jsonl back no longer waits on a + # close -- the recorder flushes every line as it writes it, and its handles + # outlive this half now that the code half writes through the same ones. metrics = C.collect_metrics(out_dir, run.counts, run.n_launched, args.solver_rollouts, time.time() - started) with open(os.path.join(out_dir, 'challenge_metrics.json'), 'w', @@ -109,6 +117,57 @@ def collect_once(args, sampler, template, slots, out_dir: str) -> Dict[str, Any] return metrics +def collect_code(args, sampler, template, recorder: Recorder, + out_dir: str) -> Dict[str, Any]: + """The code half of the same pass; returns its metrics. + + Takes no sandbox slot and asks for none. A code problem is checked by running + its asserts in a subprocess -- milliseconds, against the hundreds a microVM + round trip costs -- and the difficulty stage runs one per candidate per + rollout, so routing that through the pool would make it the dominant cost of + the iteration. The slots stay with the agentic half, whose episodes have + nowhere else to run at all. + """ + challenger = CODE.build_challenger(args, sampler, template, recorder=recorder) + metrics = CODE.collect(args, challenger, recorder) + with open(os.path.join(out_dir, 'code_metrics.json'), 'w', encoding='utf-8') as f: + json.dump(metrics, f, indent=2, ensure_ascii=False, default=str) + return metrics + + +def collect_once(args, sampler, template, slots, out_dir: str) -> Dict[str, Any]: + """Collect from every task source ``--sides`` names, into one ``out_dir``. + + One recorder for all of them, so the numbering is global and index.jsonl + interleaves the halves. That is the whole of what makes a mixed step possible: + train.py groups on ``(side, group_id)`` and never learns that two different + generators wrote the file it read. + + Each half keeps its own metrics file. Their ``counts`` use the same words for + different things -- ``groups`` is a set of sibling proposals on one side and a + single problem's attempts on the other -- and adding those together produces a + number that means neither. Only the ``scalars`` are merged, and only because + their names are disjoint by construction: the code half prefixes all of its own. + """ + os.makedirs(out_dir, exist_ok=True) + args.out_dir = out_dir + recorder = Recorder(out_dir) + scalars: Dict[str, Any] = {} + try: + if 'propose' in args.sides_list or 'solve' in args.sides_list: + # Named by either side, the agentic pair is collected whole: one build + # is what produces the task its attempts are graded on, so there is no + # way to collect the solving side without the proposing one. + metrics = collect_agentic(args, sampler, template, slots, recorder, out_dir) + scalars.update(metrics.get('scalars') or {}) + if 'code' in args.sides_list: + metrics = collect_code(args, sampler, template, recorder, out_dir) + scalars.update(metrics.get('scalars') or {}) + finally: + recorder.close() + return {'scalars': scalars} + + def main(): args = C.parse_args() root = os.path.join(args.root, args.tag) @@ -186,7 +245,14 @@ def main(): # Model rank 0 serves the TCPStore the sampler ranks connect to, so this must # be built after both halves exist. Its first call is what sends the weights. weights = CheckpointEngineManager(model=model, sampler=sampler) - slots = C.initialize_sandbox(args) + # Only if a task source needs them: --sides code boots no microVMs at all, + # which is 32 fewer machines to wait for and to be billed for. close_pool of + # an empty list is a no-op, so the teardown below needs no second condition. + agentic = 'propose' in args.sides_list or 'solve' in args.sides_list + slots = C.initialize_sandbox(args) if agentic else [] + if not agentic: + logger.info(f'[rsi] --sides {args.sides!r} names no agentic side, so no ' + f'sandbox pool is opened') logger.info(get_device_placement()) i = start diff --git a/cookbook/rsi/agentic/sandbox.py b/cookbook/rsi/agentic/sandbox.py index f17ba2a6d..99c309ec2 100644 --- a/cookbook/rsi/agentic/sandbox.py +++ b/cookbook/rsi/agentic/sandbox.py @@ -12,14 +12,12 @@ ``eval.py`` uses, so a task's difficulty here and its pass rate there are measured against one opening. """ -import os import time from concurrent.futures import ThreadPoolExecutor -from typing import List, Tuple +from typing import Any, Dict, List, Optional, Sequence, Tuple from twinkle import get_logger -from twinkle_agentic.envs import EnvTool -from twinkle_agentic.tools.tool_manager import ToolManager +from twinkle_agentic.envs import Env, StepResult from episode import solver_harness # noqa: I100,I202 from remote_tool_env import RemoteMsAgentToolEnv, tool_payload # noqa: I100,I202 @@ -133,14 +131,21 @@ RESET_RETRY_WAIT = 10 -class Sandbox: +class Sandbox(Env): """One slot: clear the workspace, run a script in it, read it back. + An :class:`~twinkle_agentic.envs.base.Env` wrapping another one, and what it + adds is RSI's policy rather than a transport: which script empties a + workspace, which one reads it back and in what format, and what to do when + either fails. That split is why the same policy works over a microVM and over + :class:`~twinkle_agentic.envs.local.LocalEnv` -- and why a caller holding a + slot does not need to know which it has. + Not thread-safe on purpose. A slot belongs to whoever holds it, and the pool hands each one to exactly one worker thread. """ - def __init__(self, slot: int, env: RemoteMsAgentToolEnv, schemas: list, + def __init__(self, slot: int, env: Env, schemas: list, *, snapshot_max_files: int, snapshot_per_file: int, snapshot_budget: int): self.slot = slot self.env = env @@ -149,18 +154,33 @@ def __init__(self, slot: int, env: RemoteMsAgentToolEnv, schemas: list, # the prompt: the schemas a trajectory is built with have to be the ones # the slot it runs on will honour. self.schemas = schemas - self._runner = env.runner() - # The tools carry the env they dispatch into, so this slot's model turns - # have to go through this slot's manager. - self.tool_manager = ToolManager(EnvTool.from_schemas(env, schemas)) self._snapshot_script = WORKSPACE_SNAPSHOT.format( workspace=self.workspace, max_files=snapshot_max_files, per_file=snapshot_per_file, total_budget=snapshot_budget) self._clear_script = CLEAR_WORKSPACE.format(workspace=self.workspace) - def run(self, script: str) -> Tuple[int, str]: - """Run a python script in the workspace; returns (exit code, output).""" - return self._runner(script, 'python') + # ------------------------------------------------------------------ Env + + def run_script(self, source: str, interpreter: str = 'python', + timeout: Optional[int] = None) -> Tuple[int, str]: + """Run a script in this slot's workspace; returns (exit code, output).""" + return self.env.run_script(source, interpreter, timeout) + + def step(self, tool_name: str, arguments: Dict[str, Any] = None) -> StepResult: + return self.env.step(tool_name, arguments or {}) + + def step_batch(self, calls: Sequence[Tuple[str, Dict[str, Any]]]) -> List[StepResult]: + return self.env.step_batch(calls) + + def tools(self) -> list: + """The schemas this slot was built with, not the ones it could re-read. + + Read once off the pool and carried, so every slot advertises the same + contract: these go into the prompt, and a slot rebuilt mid-run must not + start describing itself differently from the trajectories already in + flight against it. + """ + return list(self.schemas) def clear(self) -> None: """Empty the workspace. Raises rather than returning quietly. @@ -178,23 +198,21 @@ def clear(self) -> None: the clear is retried, then retried on a deliberately rebuilt sandbox. """ if self.env.ensure_ready(): - self._rebind('runtime was unreachable') - code, out = self.run(self._clear_script) + logger.warning(f'[sandbox {self.slot}] runtime was unreachable; rebuilt') + code, out = self.run_script(self._clear_script) if code != 0: logger.warning(f'[sandbox {self.slot}] clear failed (exit {code}), ' f'retrying in {RESET_RETRY_WAIT}s: {out[-200:]}') time.sleep(RESET_RETRY_WAIT) - code, out = self.run(self._clear_script) + code, out = self.run_script(self._clear_script) if code != 0: # Rebuilt rather than retried again: two failures in a row is not the # transient this waits out, and a fresh sandbox brings a workspace # that is already empty -- which is all this method is asked for. logger.warning(f'[sandbox {self.slot}] clear failed twice (exit {code}); ' f'rebuilding: {out[-200:]}') - self.env.reset() - self.env.n_recoveries += 1 - self._rebind('rebuilt after two failed clears') - code, out = self.run(self._clear_script) + self.env.rebuild() + code, out = self.run_script(self._clear_script) if code != 0: raise RuntimeError(f'workspace clear failed (exit {code}): {out[-400:]}') @@ -210,12 +228,12 @@ def snapshot(self) -> Tuple[str, str]: kept apart because a snapshot that says "empty" when it means "I could not look" produces tasks whose only true assertion is that nothing happened. """ - code, out = self.run(self._snapshot_script) + code, out = self.run_script(self._snapshot_script) if code != 0: logger.warning(f'[sandbox {self.slot}] snapshot failed (exit {code}), ' f'retrying in {SNAPSHOT_RETRY_WAIT}s: {out[-200:]}') time.sleep(SNAPSHOT_RETRY_WAIT) - code, out = self.run(self._snapshot_script) + code, out = self.run_script(self._snapshot_script) if code != 0: return '', f'workspace snapshot failed (exit {code}): {out[-500:]}' return tool_payload(out).strip(), '' @@ -223,12 +241,6 @@ def snapshot(self) -> Tuple[str, str]: def close(self) -> None: self.env.close() - def _rebind(self, why: str) -> None: - """Point the runner and the tools at the sandbox behind this env now.""" - self._runner = self.env.runner() - self.tool_manager = ToolManager(EnvTool.from_schemas(self.env, self.env.tool_schemas())) - logger.warning(f'[sandbox {self.slot}] rebound ({why})') - def open_pool( n: int, @@ -265,7 +277,7 @@ def _boot(_) -> RemoteMsAgentToolEnv: n = max(1, n) with ThreadPoolExecutor(max_workers=n) as pool: envs = list(pool.map(_boot, range(n))) - schemas = envs[0].tool_schemas() + schemas = envs[0].tools() slots = [ Sandbox(i, env, schemas, snapshot_max_files=snapshot_max_files, snapshot_per_file=snapshot_per_file, snapshot_budget=snapshot_budget) @@ -283,7 +295,7 @@ def close_pool(slots: List[Sandbox]) -> int: produced its numbers under a different environment than one that was rebuilt never, and that is invisible from the output files alone. """ - total = sum(getattr(s.env, 'n_recoveries', 0) for s in slots) + total = sum(s.env.n_recoveries for s in slots) for slot in slots: try: slot.close() diff --git a/cookbook/rsi/agentic/train.py b/cookbook/rsi/agentic/train.py index 9ecd3faec..754baefc4 100644 --- a/cookbook/rsi/agentic/train.py +++ b/cookbook/rsi/agentic/train.py @@ -69,18 +69,37 @@ def build_model(*, model_id: str, model_gpus: int, lr: float, template: str, return model +def sides_wanted(sides: str) -> tuple: + """The side names ``--sides`` asks for, in order, without repeats. + + Comma-separated because an iteration can collect from more than one task + source: ``both`` is the agentic pair, ``code`` is the code half, and + ``both,code`` runs all three into the same step. Unknown names are not + rejected here -- a side nobody wrote is a side ``load`` counts as skipped, + with the name in the reason, which says more than a parser error would. + """ + out = [] + for part in sides.split(','): + part = part.strip() + if part: + out.extend(('propose', 'solve') if part == 'both' else (part, )) + return tuple(dict.fromkeys(out)) + + def load(run_dir: str, *, sides: str, max_length: int) -> tuple: """Read the index into GRPO groups; returns (groups, skipped). - A group is ``(side, group_id)`` for the proposing side and - ``(side, group_id, proposal_idx)`` for the solving side -- one prompt answered - several times, which is what an advantage is computed over. + A group is ``(side, group_id)`` for the proposing and the code side and + ``(side, group_id, proposal_idx)`` for the agentic solving side -- one prompt + answered several times, which is what an advantage is computed over. The code + half proposes nothing it trains on, so one problem is one group and there is + no proposal to index within it. """ traj_dir = os.path.join(run_dir, 'trajs') index = os.path.join(traj_dir, 'index.jsonl') if not os.path.exists(index): raise SystemExit(f'[train] no {index}') - wanted = {'both': ('propose', 'solve')}.get(sides, (sides, )) + wanted = sides_wanted(sides) skipped: collections.Counter = collections.Counter() by_key: Dict[Any, List[Dict[str, Any]]] = collections.OrderedDict() with open(index, encoding='utf-8') as f: @@ -124,7 +143,7 @@ def load(run_dir: str, *, sides: str, max_length: int) -> tuple: # max_model_len 40960, which is above this. skipped[f'longer than max_length={max_length}'] += 1 continue - key = ((side, record.get('group_id')) if side == 'propose' else + key = ((side, record.get('group_id')) if side in ('propose', 'code') else (side, record.get('group_id'), record.get('proposal_idx'))) by_key.setdefault(key, []).append({ 'side': side, @@ -204,8 +223,18 @@ def train_one_step(model, run_dir: str, *, sides: str, max_length: int, batch = [m for g in interleave(groups) for m in g['members']] mix = collections.Counter(m['side'] for m in batch) sizes = collections.Counter((g['side'], len(g['members'])) for g in groups) + # What interleave cannot balance. It spreads the sides evenly by trajectory + # count, and a code attempt is an order of magnitude shorter than an agentic + # episode, so equal counts are nothing like equal shares of the update. + # Measured and reported rather than corrected for: the weighting to apply, if + # any, has to come off an observed ratio instead of a guess at one. + side_tokens: collections.Counter = collections.Counter() + side_trainable: collections.Counter = collections.Counter() + for m in batch: + side_tokens[m['side']] += len(m['input_ids']) + side_trainable[m['side']] += sum(1 for lb in m['labels'] if lb != -100) logger.info(f'[train] {len(groups)} groups, {len(batch)} trajectories {dict(mix)}; ' - f'group sizes {dict(sizes)}') + f'group sizes {dict(sizes)}; trainable tokens {dict(side_trainable)}') for note, n in sorted(skipped.items()): logger.warning(f'[train] skipped: {note} x{n}') @@ -246,6 +275,8 @@ def train_one_step(model, run_dir: str, *, sides: str, max_length: int, 'trained': len(batch) - dropped, 'dropped_tail': dropped, 'sides': dict(mix), + 'side_tokens': dict(side_tokens), + 'side_trainable_tokens': dict(side_trainable), 'group_sizes': {f'{s}:{n}': c for (s, n), c in sizes.items()}, 'advantage_min': min(advantages), 'advantage_max': max(advantages), @@ -305,6 +336,13 @@ def upload(challenge: Dict[str, Any], summary: Dict[str, Any], *, 'train/dropped_tail': summary['dropped_tail'], 'train/propose_trajectories': summary['sides'].get('propose', 0), 'train/solve_trajectories': summary['sides'].get('solve', 0), + 'train/code_trajectories': summary['sides'].get('code', 0), + # The same three by trainable tokens, which is the share of the update + # each side actually got. Always present, at 0 for a side this iteration + # did not collect, so no chart appears or disappears mid-run. + 'train/propose_tokens': summary['side_trainable_tokens'].get('propose', 0), + 'train/solve_tokens': summary['side_trainable_tokens'].get('solve', 0), + 'train/code_tokens': summary['side_trainable_tokens'].get('code', 0), 'train/advantage_min': summary['advantage_min'], 'train/advantage_max': summary['advantage_max'], 'train/learning_rate': summary['learning_rate'], diff --git a/cookbook/rsi/code/challenge.py b/cookbook/rsi/code/challenge.py index b9c930f00..261217abe 100644 --- a/cookbook/rsi/code/challenge.py +++ b/cookbook/rsi/code/challenge.py @@ -30,11 +30,14 @@ from twinkle.data_format import SamplingParams, user_data_get from twinkle.sampler import vLLMSampler from twinkle_agentic.challenger import CodeChallenger, KeywordStore, load_seeds +from twinkle_agentic.envs import LocalEnv from twinkle_agentic.rollout import build_rollout from twinkle_agentic.tools.tool_manager import ToolManager sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -from prompts import CATEGORIES, CATEGORY_DESC, code_prompts # noqa: E402 +# challenge_prompts, not prompts: the agentic half has a prompts.py of its own +# with a CATEGORIES in it, and rsi.py imports both halves into one process. +from challenge_prompts import CATEGORIES, CATEGORY_DESC, code_prompts # noqa: E402 logger = get_logger() @@ -84,8 +87,10 @@ def parse_args(): p.add_argument('--solver-rollouts', type=int, default=8) p.add_argument('--solver-temp', type=float, default=1.0) p.add_argument('--solver-max-tokens', type=int, default=2048) - p.add_argument('--keep-min-pass', type=int, default=1) - p.add_argument('--keep-max-margin', type=int, default=1) + p.add_argument('--keep-pass-band', type=int, nargs=2, default=(1, 7), + metavar=('LOW', 'HIGH'), + help='keep problems solved this many times out of ' + '--solver-rollouts, inclusive; the default is the band for 8') p.add_argument('--sandbox-timeout', type=int, default=30) p.add_argument('--max-checks', type=int, default=6) @@ -155,6 +160,12 @@ def _reject(record): challenger = CodeChallenger( code_prompts(), explorer, + # Checks run on this machine, in a throwaway directory per script. A + # generated check is a self-contained program over its own asserts, so + # it needs no workspace to carry state between calls -- and the same + # slot is handed a sandbox instead when a task needs one. The per-check + # deadline comes from ``sandbox_timeout`` below, stated on every call. + envs=[LocalEnv()], seeds=seeds, keyword_store=store, category_desc=CATEGORY_DESC if store else None, @@ -181,8 +192,7 @@ def _reject(record): reject_sink=_reject, max_proposals_per_round=args.max_proposals_per_round, solver_rollouts=args.solver_rollouts, - keep_min_pass=args.keep_min_pass, - keep_max_pass_margin=args.keep_max_margin, + keep_pass_band=tuple(args.keep_pass_band), solver_params=SamplingParams(max_tokens=args.solver_max_tokens, num_samples=1, logprobs=1, temperature=args.solver_temp, top_p=0.95), seed=args.random_seed, @@ -251,8 +261,12 @@ def write_flows(kept, args): }], } ff.write(json.dumps(flow, ensure_ascii=False) + '\n') + # The challenger keeps one check script, the shape the agentic half + # also uses; this file is a list of asserts because that is what + # rsi_rl reads, so split it back on the way out. + check_script = user_data_get(data, 'check_script', '') or '' ft.write(json.dumps({'id': cid, - 'test_list': user_data_get(data, 'asserts', []), + 'test_list': check_script.splitlines(), 'test_setup_code': ''}, ensure_ascii=False) + '\n') diff --git a/cookbook/rsi/code/prompts.py b/cookbook/rsi/code/challenge_prompts.py similarity index 100% rename from cookbook/rsi/code/prompts.py rename to cookbook/rsi/code/challenge_prompts.py diff --git a/cookbook/rsi/code/collect.py b/cookbook/rsi/code/collect.py new file mode 100644 index 000000000..c18a907ad --- /dev/null +++ b/cookbook/rsi/code/collect.py @@ -0,0 +1,240 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""RSI self-play, code half, as one iteration of the resident loop. + +``code/challenge.py`` generates problems and writes them to jsonl for a separate +training run to pick up later. This does the same generation and hands the result +straight to the step, in the process that owns the weights -- the arrangement +rsi.py's docstring argues for, and the reason a code problem and an agentic task +now land in one ``trajs/index.jsonl`` under one ``side`` field. + +What one problem contributes is one GRPO group: the ``solver_rollouts`` attempts +the difficulty stage already made at it, each with a binary reward. Nothing is +sampled twice. The band that decides whether a problem is worth keeping -- +``1 <= n_pass <= 7`` of 8 by default -- is the same band that guarantees the +group has a gradient, so selection and grouping are one decision rather than two +that can disagree. Attempts reach here through ``CodeChallenger``'s +``solver_sink``; the agentic half has had the same hook for the same reason. + +Difficulty judgements do not take a sandbox slot. One is a subprocess running the +problem's asserts, milliseconds, and the stage makes ``candidates x rollouts`` of +them per round -- through a microVM that would be the dominant cost of the +iteration, and the 32 slots are worth more to the agentic half, whose episodes +cannot run anywhere else. That choice is one argument: the ``envs`` this half is +built with are :class:`~twinkle_agentic.envs.local.LocalEnv`, and handing it +``sandbox.open_pool``'s slots instead is the whole change if the trade ever does. +""" +import os +import sys +import time +from typing import Any, Callable, Dict, List, Optional + +from twinkle import get_logger +from twinkle.data_format import SamplingParams, Trajectory, user_data_get +from twinkle_agentic.challenger import CodeChallenger, KeywordStore, load_seeds +from twinkle_agentic.envs import LocalEnv +from twinkle_agentic.rollout import build_rollout +from twinkle_agentic.tools.tool_manager import ToolManager + +# Appended, not prepended: rsi.py imports this half into the process that already +# owns the agentic one, and the two directories both hold a challenge.py. Putting +# this one in front would decide that name for everybody who imports afterwards. +_HERE = os.path.dirname(os.path.abspath(__file__)) +if _HERE not in sys.path: + sys.path.append(_HERE) +from challenge_prompts import CATEGORIES, CATEGORY_DESC, code_prompts # noqa: E402 + +logger = get_logger() + + +class CollectingChallenger(CodeChallenger): + """A CodeChallenger that keeps the attempts its difficulty stage makes. + + The base class measures a candidate by sampling it ``solver_rollouts`` times + and then reports one number, and those rollouts are what the solving side + trains on. Holding on to all of them would cost a gigabyte a round -- most + candidates fall outside the band -- so they are dropped as soon as the number + they produced says the candidate is not a keeper. + + The dropping reads ``keep_pass_band`` off self, i.e. the same tuple the base + class applies one line later, so this is not a second filter with its own + opinion. It runs in :meth:`on_difficulty_measured`, which the base class calls + with every candidate of the round after they are measured and before they are + filtered -- the only moment where both the counts and the attempts are in hand. + """ + + def __init__(self, *args: Any, + attempt_sink: Optional[Callable[[Dict[str, Any]], None]] = None, + **kwargs: Any): + super().__init__(*args, solver_sink=self._keep, **kwargs) + if not self.solver_rollouts: + raise ValueError('CollectingChallenger has nothing to train on with the ' + 'difficulty stage off: the attempts it collects ARE the ' + 'solving side. Pass solver_rollouts and keep_pass_band.') + self.attempt_sink = attempt_sink + # check_script -> the attempts at the problem it verifies. Keyed on the + # script because that is the one field a task carries unchanged from the + # judgement to the batch it is yielded in; the task dict itself is copied + # on the way through attach_user_data. + self._attempts: Dict[str, List[Dict[str, Any]]] = {} + + def _keep(self, record: Dict[str, Any]) -> None: + """``solver_sink``: file the verdict, hold on to the trajectory.""" + if self.attempt_sink is not None: + # Without the trajectory: every attempt that ends up trained on is + # written in full to index.jsonl anyway, and the ones that do not are + # here for the question of why a problem measured 0 of 8, which the + # verdict and the interpreter's complaint answer. + self.attempt_sink({k: v for k, v in record.items() if k != 'attempt'}) + self._attempts.setdefault(record['check_script'], []).append(record) + + def on_difficulty_measured(self, candidates: List[Trajectory]) -> None: + super().on_difficulty_measured(candidates) + low, high = self.keep_pass_band + for task in candidates: + data = task.get('user_data') + if not low <= user_data_get(data, 'n_pass', 0) <= high: + self._attempts.pop(user_data_get(data, 'check_script', '') or '', None) + + def take(self, check_script: str) -> List[Dict[str, Any]]: + """The attempts at one kept problem, removed from the store.""" + return self._attempts.pop(check_script, []) + + +def build_challenger(args, sampler, template, *, recorder=None) -> CollectingChallenger: + """The code half wired to the loop's live sampler, ready for one iteration. + + ``sampler`` and ``template`` belong to the caller and outlive this: an + iteration must propose and solve with the weights the last step produced, so + building an engine here would be building the wrong one. The template has to + be the caller's object too -- the rollout continues a conversation by splicing + token ids, so the ids it appends must come from the same encoder the agentic + half is using on the same sampler. + """ + params = SamplingParams(max_tokens=args.code_propose_max_tokens, num_samples=1, + logprobs=1, temperature=args.code_propose_temp, top_p=0.95) + # One rollout for proposing and, through solver_params, for solving. max_turns=1 + # because a code answer is one message: there is nothing for a second turn to + # react to until the asserts have run, and running them is the next stage. + explorer = build_rollout(sampler, template=template, + tool_manager=ToolManager([]), max_turns=1, + sampling_params=params) + store = None + if args.code_keywords_n > 0: + store = KeywordStore(args.code_keyword_db, CATEGORIES) + logger.info('[collect_code] keyword bank: ' + + ', '.join(f'{c}={len(store.items[c])}' for c in CATEGORIES)) + seeds = load_seeds(args.code_seed_file) + logger.info(f'[collect_code] seeds: {len(seeds)} from {args.code_seed_file!r}') + return CollectingChallenger( + code_prompts(), + explorer, + # This half's slot: a check is a self-contained program over its own + # asserts, run here in a throwaway directory. See the module docstring + # for why it is not one of the agentic half's sandboxes. + envs=[LocalEnv()], + seeds=seeds, + keyword_store=store, + category_desc=CATEGORY_DESC if store else None, + seed_mix_prob=args.code_seed_mix_prob, + two_step=not args.code_no_two_step, + keyword_refill_target=args.code_keywords_n, + keyword_params=SamplingParams(max_tokens=1024, num_samples=1, logprobs=1, + temperature=1.3, top_p=0.98), + # A batch under the sampler's data-parallel width leaves workers idle. + min_batch=args.sampler_gpus, + problem_max_chars=args.code_problem_max_chars, + max_checks=args.code_max_checks, + sandbox_timeout=args.code_script_timeout, + max_proposals_per_round=args.code_max_proposals_per_round, + solver_rollouts=args.code_solver_rollouts, + keep_pass_band=tuple(args.code_keep_pass_band), + solver_params=SamplingParams(max_tokens=args.code_solver_max_tokens, + num_samples=1, logprobs=1, + temperature=args.code_solver_temp, top_p=0.95), + seed=args.random_seed, + reject_sink=(recorder.rejected if recorder is not None else None), + attempt_sink=(recorder.attempt if recorder is not None else None), + ) + + +def collect(args, challenger: CollectingChallenger, recorder, *, + group_id_base: int = 0) -> Dict[str, Any]: + """Generate problems until ``--code-keep-target``, writing groups as they land. + + ``group_id_base`` offsets the ids so two task sources sharing one recorder + cannot collide. train.py groups on ``(side, group_id)`` and ``side`` already + separates the halves, so this is belt and braces -- and it is what makes the + ids in index.jsonl still mean something when read by hand. + """ + started = time.time() + counts: Dict[str, int] = {'kept': 0, 'groups': 0, 'trajectories': 0, + 'no_attempts': 0, 'ungrouped': 0} + pass_dist: Dict[int, int] = {} + batch_size = args.code_batch_size or args.code_keep_target + for batch in challenger(batch_size=batch_size, total=args.code_keep_target): + for task in batch: + counts['kept'] += 1 + data = task.get('user_data') + check_script = user_data_get(data, 'check_script', '') or '' + n_pass = user_data_get(data, 'n_pass', 0) + pass_dist[n_pass] = pass_dist.get(n_pass, 0) + 1 + records = challenger.take(check_script) + if len(records) < 2: + # A group of one has an advantage of the reward minus itself, and + # none at all is reachable only if two problems ended up with + # byte-identical asserts and the first yielded took both sets. + # Counted rather than ignored: either would otherwise read as a + # quiet shortfall in how much the iteration trained on. + counts['no_attempts' if not records else 'ungrouped'] += 1 + continue + group_id = group_id_base + counts['groups'] + counts['groups'] += 1 + recorder.task({ + 'side': 'code', 'group_id': group_id, + 'statement': records[0].get('statement', ''), + 'check_script': check_script, + 'setup_script': user_data_get(data, 'setup_script', '') or '', + # The challenger's own passing code, for OPSD and for reading a + # group back: an attempt is only judgeable against a solution. + 'solution': user_data_get(data, 'solution', ''), + 'entry': user_data_get(data, 'entry', ''), + 'n_pass': n_pass, + 'n_rollouts': user_data_get(data, 'n_rollouts', 0), + 'keywords': user_data_get(data, 'keywords', []), + 'seeded': user_data_get(data, 'seeded', False), + 'two_step': user_data_get(data, 'two_step', False), + }) + for idx, record in enumerate(records): + counts['trajectories'] += 1 + recorder.trajectory( + record['attempt'], side='code', group_id=group_id, + # One problem is one group, so there is no proposal to index + # within it. Written anyway, at 0, because index.jsonl is read + # by one loader for both halves. + proposal_idx=0, + reward=1.0 if record['passed'] else 0.0, + attempt_idx=idx, passed=record['passed'], + n_pass=n_pass, check_output=record.get('output', '')) + logger.info(f'[collect_code] {counts["kept"]}/{args.code_keep_target} problems, ' + f'{counts["groups"]} groups, {counts["trajectories"]} trajectories; ' + f'proposal stats {challenger.stats}') + + if challenger.keywords is not None: + # After the loop: what it adds is for the next iteration, so a crash in + # collection does not also cost the bank. + challenger.expand_hard_keywords() + challenger.keywords.save() + metrics = { + 'scalars': { + 'code_problems': counts['kept'], + 'code_groups': counts['groups'], + 'code_trajectories': counts['trajectories'], + 'code_proposed': challenger.n_proposed, + 'code_seconds': round(time.time() - started, 1), + }, + 'counts': {**counts, 'proposals': dict(challenger.stats), + 'pass_dist': dict(sorted(pass_dist.items()))}, + } + logger.info(f'[collect_code] done in {metrics["scalars"]["code_seconds"]}s: ' + f'{metrics["scalars"]}; pass counts {metrics["counts"]["pass_dist"]}') + return metrics diff --git a/cookbook/rsi/recorder.py b/cookbook/rsi/recorder.py new file mode 100644 index 000000000..a8948500f --- /dev/null +++ b/cookbook/rsi/recorder.py @@ -0,0 +1,170 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""What a collection pass writes, shared by both halves of the loop. + +The agentic and the code half invent completely different problems, but a +trajectory is a trajectory: token fields to ``.npz``, everything a reader needs +to interpret them to ``trajs/index.jsonl``, and train.py reads that one index +without caring which half produced a line. Keeping one writer is what makes +``side`` a plain field rather than two file formats to reconcile. + +In cookbook rather than in :mod:`twinkle_agentic` on purpose: this is the +on-disk contract between a collection pass and the step that trains on it, and +that contract is still moving -- fields get added as questions come up about +runs. A library version would freeze it, and the freezing is the expensive part, +not the code. +""" +import json +import os +import threading +from typing import Any, Dict, List + +import numpy as np + + +def logprob_column(logprobs: Any) -> List[float]: + """One float per generated token: the logprob of the token that was chosen. + + The sampler hands these over as ``List[List[Tuple[int, float]]]`` -- per + generated token, a list of top-k ``(token_id, logprob)`` pairs with the chosen + token first (``SampledSequence.logprobs``, data_format/sampling.py:185). + Passing that to ``np.asarray`` directly would store an ``(N, k, 2)`` array and + the loader would hand GRPO nested lists where it wants one float per trainable + token -- which is a crash inside the step, or worse a silent reshape. + + A plain list of floats is accepted too, for a sampler that already flattened. + Anything else raises rather than being coerced: a wrong ``old_logps`` makes the + GRPO ratio wrong on the first step, and nothing downstream would say so. + """ + out: List[float] = [] + for step in logprobs: + if isinstance(step, (int, float)): + out.append(float(step)) + continue + if isinstance(step, (list, tuple)) and step: + head = step[0] + if isinstance(head, (list, tuple)) and len(head) >= 2: + out.append(float(head[1])) + continue + raise TypeError(f'cannot read a logprob out of {step!r}; expected a float ' + f'or a list of (token_id, logprob) pairs') + return out + + +class Recorder: + """Everything a run writes, behind one lock. + + Trajectories go to ``.npz`` for the token fields and to ``index.jsonl`` for + everything a reader needs to interpret them. The text is written in full and + never truncated: these files are read to check whether a reward was deserved, + which a shortened statement cannot answer. + + Both halves of an iteration share one instance, so the numbering is global + and the index interleaves them. That is also why every handle is opened up + front even when the half in front of it has nothing to put in some of them: + a file that appears only sometimes is a file every reader has to guard. + """ + + def __init__(self, out_dir: str): + self.dir = out_dir + self.traj_dir = os.path.join(out_dir, 'trajs') + os.makedirs(self.traj_dir, exist_ok=True) + self._lock = threading.Lock() + self._n = 0 + self._index = open(os.path.join(self.traj_dir, 'index.jsonl'), 'w', encoding='utf-8') + self._groups = open(os.path.join(out_dir, 'groups.jsonl'), 'w', encoding='utf-8') + self._tasks = open(os.path.join(out_dir, 'tasks.jsonl'), 'w', encoding='utf-8') + # Why a build produced no task. The reason alone is not diagnosable: nine + # empty_workspace rejections in one run all looked like the model refusing + # to act, and the question of whether it had run out of tokens or simply + # emitted no call could not be answered from the record, because the fields + # that answered it were on the trajectory and were dropped. + self._rejected = open(os.path.join(out_dir, 'rejected.jsonl'), 'w', encoding='utf-8') + # Keyword replies, both sides in full. The one question this file exists to + # answer -- did the model disobey the format, or does the parser reject what + # it produced -- cannot be answered from a count. Keyword generation was + # silently broken for whole runs when the prompt asked for one per line and + # the parser wanted a JSON array. + self._keywords = open(os.path.join(out_dir, 'keyword_gen.jsonl'), 'w', encoding='utf-8') + # Every solver attempt, passed or not, with the state it left and what the + # check said about it. A task measured at 0 of 8 has three explanations -- + # the check is wrong, the statement withholds something the check demands, + # or the solver gave up -- and only the attempt and the workspace it left + # tell them apart. Written for every attempt, not only for the ones that + # end up trained on: the failures are what this file is for. + self._attempts = open(os.path.join(out_dir, 'solver_attempts.jsonl'), 'w', + encoding='utf-8') + # The rubric, all three of its dimensions. Only novelty reaches a reward; + # usefulness and complexity are recorded so the question of whether they + # should count can be answered from a run instead of argued. + self._novelty = open(os.path.join(out_dir, 'novelty_scores.jsonl'), 'w', + encoding='utf-8') + + def trajectory(self, traj: Dict[str, Any], **fields: Any) -> None: + """One training sample: token fields to npz, everything else to the index. + + A trajectory with no ``logprobs`` is written anyway, with the field left + null. It is not trainable and the loader will say so -- which is the point: + a sample silently dropped here would make the group it belongs to look like + a different size than it was. + """ + input_ids = np.asarray(traj.get('input_ids') or [], dtype=np.int32) + labels = np.asarray(traj.get('labels') or [], dtype=np.int32) + logprobs = traj.get('logprobs') + with self._lock: + self._n += 1 + name = f'{self._n:06d}.npz' + arrays = {'input_ids': input_ids, 'labels': labels} + if logprobs is not None: + # float64, and the chosen token's column only. These are the old_logps a + # GRPO step divides by; float32 would round them to about 7 digits, so + # the ratio exp(logp - old_logp) would be off by roughly 1e-7 for + # reasons that have nothing to do with the policy having changed. + arrays['logprobs'] = np.asarray(logprob_column(logprobs), dtype=np.float64) + # Compressed: a 24-turn agentic episode is tens of thousands of token ids, + # and 128 of them per iteration adds up on disk. + np.savez_compressed(os.path.join(self.traj_dir, name), **arrays) + record = dict(fields) + record.update({ + 'npz': name, + 'n_tokens': int(input_ids.size), + 'n_trainable': int((labels != -100).sum()) if labels.size else 0, + 'has_logprobs': logprobs is not None, + # The rollout guarantees one logprob per trainable label; recorded so a + # loader can check it rather than trust it. + 'n_logprobs': int(arrays['logprobs'].size) if logprobs is not None else 0, + 'turns': traj.get('turns'), + 'stop_reason': traj.get('stop_reason'), + 'truncated': bool(traj.get('truncated')), + 'tool_stop': traj.get('tool_stop'), + 'messages': traj.get('messages') or [], + }) + self._write(self._index, record) + + def group(self, record: Dict[str, Any]) -> None: + self._write(self._groups, record) + + def task(self, record: Dict[str, Any]) -> None: + self._write(self._tasks, record) + + def rejected(self, record: Dict[str, Any]) -> None: + self._write(self._rejected, record) + + def keywords(self, record: Dict[str, Any]) -> None: + self._write(self._keywords, record) + + def attempt(self, record: Dict[str, Any]) -> None: + self._write(self._attempts, record) + + def novelty(self, record: Dict[str, Any]) -> None: + self._write(self._novelty, record) + + def close(self) -> None: + for handle in (self._index, self._groups, self._tasks, self._rejected, + self._keywords, self._attempts, self._novelty): + handle.close() + + def _write(self, handle, record: Dict[str, Any]) -> None: + line = json.dumps(record, ensure_ascii=False, default=str) + with self._lock: + handle.write(line + '\n') + handle.flush() diff --git a/cookbook/rsi/rl.py b/cookbook/rsi/rl.py index 4087f4d36..1e4062f0a 100644 --- a/cookbook/rsi/rl.py +++ b/cookbook/rsi/rl.py @@ -47,11 +47,12 @@ Solver learning mode (RSI step-3 subclass, RSI_SOLVER_MODE): * 'grpo' (default) -- on a code round whose first attempt FAILS the asserts, the - sandbox error is injected back as a {'role':'tool'} message and the model is - asked to continue, up to RSI_SOLVER_MAX_TURNS total turns. The whole - multi-turn trajectory (turn-1 tokens + turn-2 tokens, the tool error bridged - in as -100) is trained by GRPO on the final pass/fail reward. Tool rounds and - length-stopped rollouts stay single-shot. Bridge tokens are computed in + execution output is handed back as a {'role':'user'} message and the model is + asked to fix it, up to RSI_SOLVER_MAX_TURNS total turns. The whole + multi-turn trajectory (turn-1 tokens + turn-2 tokens, the error message + bridged in as -100) is trained by GRPO on the final pass/fail reward. Tool + rounds and length-stopped rollouts stay single-shot. The continuation is + MultiTurnRollout's ``followup_fn``, so the bridge tokens are computed in template space and appended verbatim -- never decode-then-re-encode. * 'opsd' -- single turn. A teacher forward conditioned on a PRIVILEGED extra system message carrying the challenger's passing reference solution @@ -71,12 +72,6 @@ import os import random import re -import resource -import shutil -import signal -import subprocess -import sys -import tempfile import time from concurrent.futures import ThreadPoolExecutor from typing import Any, Dict, List, Optional, Tuple @@ -86,13 +81,18 @@ from twinkle.advantage import GRPOAdvantage from twinkle.checkpoint_engine import CheckpointEngineManager from twinkle.cli import CLI -from twinkle.data_format import SamplingParams +from twinkle.data_format import SamplingParams, Trajectory from twinkle.dataloader import DataLoader from twinkle.dataset import Dataset, DatasetMeta from twinkle.metric import CompletionRewardMetric from twinkle.processor import InputProcessor from twinkle.reward.base import Reward from twinkle.sampler import vLLMSampler +from twinkle_agentic.challenger.code import run_asserts, run_check_script +from twinkle_agentic.rollout.multi_turn import MultiTurnRollout +from twinkle_agentic.tools.tool_manager import ToolManager +from twinkle_agentic.utils.code_utils import unwrap_code +from twinkle_agentic.utils.message_utils import assistant_text logger = get_logger() args = CLI.from_args() @@ -389,68 +389,10 @@ def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: # Same sandbox contract as cookbook/rl/grpo/mbpp_grpo.py, which was checked # against all 974 MBPP reference solutions (974/974 pass): the generated code, # the setup code and the asserts are concatenated into one file and executed, so -# a bare ``assert fn(...) == x`` resolves the function by name. -_FENCE_RE = re.compile(r'```(?:python|py)?\s*\n(.*?)```', re.S) - - -def extract_code(text: str) -> str: - """Take the last fenced block; fall back to the whole body when unfenced.""" - idx = (text or '').rfind('</think>') - body = text[idx + len('</think>'):] if idx >= 0 else (text or '') - blocks = _FENCE_RE.findall(body) - return (blocks[-1] if blocks else body).strip() - - -def run_asserts(code: str, setup: str, asserts: List[str], timeout: int = TEST_TIMEOUT) -> bool: - """True when every assert passes. Thin wrapper over run_asserts_verbose.""" - return run_asserts_verbose(code, setup, asserts, timeout)[0] - - -def run_asserts_verbose(code: str, setup: str, asserts: List[str], - timeout: int = TEST_TIMEOUT) -> Tuple[bool, str]: - """Run code+setup+asserts and return (passed, stderr_text). - - Same sandbox contract as the MBPP-verified path (start_new_session + killpg - so a forking solution leaves no stray processes; RLIMIT_AS caps the child at - 2GB). stderr is captured (not sent to /dev/null) so the GRPO continuation can - feed the actual traceback back to the model as a tool message. ``passed`` is - exactly ``returncode == 0``, identical to the old bool-only behavior. - """ - if not code.strip() or not asserts: - return False, 'no code was produced' - parts = [code] - if (setup or '').strip(): - parts.append(setup) - parts.extend(asserts) - tmp = tempfile.mkdtemp(prefix='rsi_code_') - try: - with open(os.path.join(tmp, '_run.py'), 'w', encoding='utf-8') as f: - f.write('\n\n'.join(parts) + '\n') - env = dict(os.environ, MPLBACKEND='Agg', PYTHONHASHSEED='0', OMP_NUM_THREADS='1', - MKL_NUM_THREADS='1', TOKENIZERS_PARALLELISM='false') - env.pop('CUDA_VISIBLE_DEVICES', None) - - def _limit(): - resource.setrlimit(resource.RLIMIT_AS, (2 * 1024**3, 2 * 1024**3)) - - proc = subprocess.Popen([sys.executable, '_run.py'], cwd=tmp, env=env, - stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, - text=True, start_new_session=True, preexec_fn=_limit) - try: - _, err = proc.communicate(timeout=timeout) - return proc.returncode == 0, (err or '') - except subprocess.TimeoutExpired: - try: - os.killpg(proc.pid, signal.SIGKILL) - except ProcessLookupError: - pass - try: - proc.communicate(timeout=5) - except Exception: - pass - return False, f'execution timed out after {timeout}s (possible infinite loop)' - finally: - shutil.rmtree(tmp, ignore_errors=True) +# a bare ``assert fn(...) == x`` resolves the function by name. Both the run and +# the fence-stripping come from the library -- ``run_asserts`` for a verdict, +# ``run_check_script`` for a verdict plus the output that error feedback shows +# the model, and ``unwrap_code`` for reading the code out of a reply. def load_tests() -> Dict[str, Dict[str, Any]]: @@ -540,7 +482,7 @@ def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[Optiona spec = json.loads(ud['code_tests']) except (ValueError, TypeError): continue - code_jobs.append((i, extract_code(completion), spec)) + code_jobs.append((i, unwrap_code(completion), spec)) recs[i]['kind'] = 'code' if code_jobs: @@ -551,7 +493,8 @@ def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[Optiona todo = list(uniq) with ThreadPoolExecutor(max_workers=max(1, min(JUDGE_WORKERS, len(todo)))) as ex: verdicts = dict(zip(todo, ex.map( - lambda k: run_asserts(k[1], uniq[k]['setup'], uniq[k]['asserts']), todo))) + lambda k: run_asserts(k[1], uniq[k]['setup'], uniq[k]['asserts'], + TEST_TIMEOUT), todo))) for i, code, spec in code_jobs: rewards[i] = 1.0 if verdicts.get((str(spec.get('id')), code)) else 0.0 @@ -857,15 +800,8 @@ def make_local_template(): return t -def _last_assistant_text(pif: Dict[str, Any]) -> str: - for m in reversed(pif.get('messages') or []): - if m.get('role') == 'assistant': - return m.get('content', '') or '' - return '' - - def _format_exec_error(err: str) -> str: - """Turn captured stderr into the tool message shown back to the model.""" + """Turn a failed check's output into the message shown back to the model.""" err = (err or '').strip() or 'Your code did not pass the tests (no error output captured).' if len(err) > 1500: err = err[:700] + '\n...[truncated]...\n' + err[-700:] @@ -875,142 +811,55 @@ def _format_exec_error(err: str) -> str: '```python code block.') -def _bridge_tool_message(template, pif: Dict[str, Any], tool_content: str) -> Optional[Dict[str, Any]]: - """Append a {'role':'tool'} turn + next generation prompt as -100 bridge. +def check_script_of(spec: Dict[str, Any]) -> str: + """A tests entry's setup + asserts as the one script ``run_check_script`` runs.""" + parts = [spec['setup']] if (spec.get('setup') or '').strip() else [] + parts.extend(spec.get('asserts') or ()) + return '\n\n'.join(parts) + + +def code_error_followup(traj: Trajectory, n_followups: int) -> Optional[str]: + """Ask a failed code rollout to fix itself, or None to let the episode end. - Computed entirely in template space (render-after minus render-before), so - history tokens stay byte-for-byte in ``input_ids`` and only the new tool turn - is tokenized from canonical template output -- never decode-then-re-encode. - Mirrors Template.concat_input_feature / MultiTurnRollout._extend_with_bridge. - Returns the extended pif, or None if it would exceed the template max_length. + MultiTurnRollout calls this at the moment a rollout would finish, which is + where the hand-rolled retry pass used to run. The budget is unchanged: + SOLVER_MAX_TURNS counts turns and turn 1 is the rollout's own, so there are + SOLVER_MAX_TURNS - 1 follow-ups to give away. A tool round has no tests to + fail and is never continued, and a reply cut off at ``max_tokens`` never gets + here -- the rollout ends a length-stopped trajectory before asking. """ - import copy - tok = template.tokenizer - messages_before = list(pif.get('messages') or []) - messages_after = messages_before + [{'role': 'tool', 'content': tool_content}] - et = getattr(template, 'enable_thinking', False) - s_before = tok.apply_chat_template(messages_before, tokenize=False, - add_generation_prompt=False, enable_thinking=et) - s_after = tok.apply_chat_template(messages_after, tokenize=False, - add_generation_prompt=True, enable_thinking=et) - # SEAM: the vLLM pif ends at the assistant's closing <|im_end|> with NO trailing - # newline (generation stops at the eos token), but the canonical render puts a - # "\n" right after that <|im_end|>. Splitting at len(s_before) would drop that - # "\n" and append the tool turn directly onto <|im_end|>, producing a malformed - # "<|im_end|><|im_start|>" boundary that the trained turn-2 tokens then condition - # on. Split right AFTER the assistant's <|im_end|> so the bridge carries the - # "\n" + tool turn and reproduces the canonical tokenization exactly. - marker = '<|im_end|>' - cut = s_before.rfind(marker) - if cut < 0: - raise RuntimeError('tool bridge: no <|im_end|> found in the rendered history; ' - 'cannot locate the assistant turn boundary.') - head = s_before[:cut + len(marker)] - if not s_after.startswith(head): - raise RuntimeError('tool bridge: chat template is not monotonic in the message list; ' - 'cannot append a tool turn as a suffix.') - bridge_text = s_after[len(head):] - bridge_ids = tok.encode(bridge_text, add_special_tokens=False) - if not bridge_ids: - raise RuntimeError('tool bridge tokenized to an empty id list') - result = copy.deepcopy(pif) - input_ids = list(result['input_ids']) - labels = list(result.get('labels') or []) - if labels: - if len(labels) != len(input_ids): - raise RuntimeError('tool bridge: labels/input_ids length mismatch') - labels = labels[-1:] + labels[:-1] # unroll to input order (mirror concat_input_feature) - else: - labels = [-100] * len(input_ids) - result['input_ids'] = input_ids + bridge_ids - result['labels'] = labels + [-100] * len(bridge_ids) - max_len = getattr(template, 'max_length', None) - if max_len and len(result['input_ids']) > max_len: + if n_followups >= SOLVER_MAX_TURNS - 1: + return None + ud = {item[0]: item[1] for item in (traj.get('user_data') or [])} + if 'code_tests' not in ud: + return None + try: + spec = json.loads(ud['code_tests']) + except (ValueError, TypeError): return None - new_if = template._invoke_post_pipeline([result])[0] - result.update(new_if) - result['messages'] = messages_after - return result - - -def grpo_continue(sampler, template, expand_prompts, sampling_params): - """GRPO rollout with error-feedback continuation for code rounds. - - Turn 1 samples every prompt. A code sample that FAILS its asserts (and did - not stop on 'length') gets the sandbox stderr injected as a {'role':'tool'} - message and is re-sampled, up to SOLVER_MAX_TURNS total turns. Tool rounds - and length-stopped samples are never continued. The returned per-sample - input feature is the full multi-turn trajectory (turn tokens trainable, tool - bridge -100) and old_logps is the concatenation of each turn's logprobs, so - the (#logps == #trainable labels) invariant holds for GRPO training. + passed, output = run_check_script(unwrap_code(assistant_text(traj)), + check_script_of(spec), TEST_TIMEOUT) + return None if passed else _format_exec_error(output) + + +def make_solver_rollout(sampler, template, sampling_params): + """The GRPO rollout: one turn of code, plus a fix-it round when it fails. + + ``max_turns=1`` is what makes this a text rollout. A code round's reply IS + python, and python parses as a tool-call list; the rollout checks the turn + budget before dispatching, so at 1 the calls it thinks it found are never + run. ``max_malformed_retries=0`` is the same concern from the other side: + markup that only looks like a call must not buy the sample another turn. + The ToolManager is empty and present only because the rollout requires one. + + Follow-ups are paid for separately from ``max_turns`` -- each one granted + adds a generation -- so a rollout still runs at most SOLVER_MAX_TURNS turns, + and ``code_error_followup`` is what stops before that. """ - resps = sampler.sample(expand_prompts, sampling_params) - pifs: List[Dict[str, Any]] = [] - logps: List[List[float]] = [] - lens: List[int] = [] - stops: List[Optional[str]] = [] - for r in resps: - s = r.sequences[0] - pifs.append(s.new_input_feature) - logps.append([lp[0][1] for lp in s.logprobs]) - lens.append(len(s.tokens)) - stops.append(s.stop_reason) - - done = [False] * len(expand_prompts) - dm = getattr(sampler, 'device_mesh', None) - min_batch = dm.data_world_size if dm is not None else 1 - for _turn in range(2, SOLVER_MAX_TURNS + 1): - retry: List[int] = [] - for i, prompt in enumerate(expand_prompts): - if done[i]: - continue - ud = {item[0]: item[1] for item in (prompt.get('user_data') or [])} - if 'code_tests' not in ud or stops[i] == 'length': - done[i] = True - continue - try: - spec = json.loads(ud['code_tests']) - except (ValueError, TypeError): - done[i] = True - continue - code = extract_code(_last_assistant_text(pifs[i])) - passed, err = run_asserts_verbose(code, spec.get('setup', ''), spec.get('asserts', [])) - if passed: - done[i] = True - continue - # Bridge the tool error in. If the template can't append a tool turn as - # a clean suffix (e.g. a malformed/cut turn-1 without a proper </think>), - # skip continuation for THIS sample rather than crashing the whole step. - try: - bridged = _bridge_tool_message(template, pifs[i], _format_exec_error(err)) - except RuntimeError as e: - logger.warning(f'[rsi_rl][grpo] skip continuation for sample {i}: {e}') - bridged = None - if bridged is None: - done[i] = True - continue - pifs[i] = bridged - retry.append(i) - if not retry: - break - batch = [pifs[i] for i in retry] - if len(batch) < min_batch: - batch = batch + [batch[-1]] * (min_batch - len(batch)) - rresps = sampler.sample(batch, sampling_params)[:len(retry)] - for j, i in enumerate(retry): - s2 = rresps[j].sequences[0] - pifs[i] = s2.new_input_feature - logps[i].extend([lp[0][1] for lp in s2.logprobs]) - lens[i] += len(s2.tokens) - stops[i] = s2.stop_reason - - # Same invariant MultiTurnRollout enforces: one logp per trainable token. - for i, pif in enumerate(pifs): - trainable = sum(1 for lb in (pif.get('labels') or []) if lb != -100) - if len(logps[i]) != trainable: - raise RuntimeError(f'GRPO continuation logps/labels misaligned for sample {i}: ' - f'{len(logps[i])} logps vs {trainable} trainable labels') - return pifs, logps, lens + return MultiTurnRollout(sampler, template=template, tool_manager=ToolManager(), + max_turns=1, max_malformed_retries=0, + followup_fn=code_error_followup, + sampling_params=sampling_params) def _teacher_pif(template, student_pif: Dict[str, Any], ref_solution: str, @@ -1200,8 +1049,8 @@ def main(): ref_model.set_template(TEMPLATE, model_id=REF_MODEL_ID, max_length=MAX_MODEL_LEN, enable_thinking=True) - # Driver-side template for token surgery: the GRPO tool-error bridge and the - # OPSD teacher-prompt concat both run on the driver. + # Driver-side template for token surgery: the GRPO rollout's follow-up bridge + # and the OPSD teacher-prompt concat both run on the driver. local_template = make_local_template() ckpt_manager = CheckpointEngineManager(model=model, sampler=sampler) @@ -1213,6 +1062,8 @@ def main(): metrics = CompletionRewardMetric() reward_fn = RoundReward() sampling_params = SamplingParams(max_tokens=MAX_NEW_TOKENS, num_samples=1, logprobs=1, temperature=1.0, top_p=0.95) + solver_rollout = (make_solver_rollout(sampler, local_template, sampling_params) + if SOLVER_MODE == 'grpo' else None) optim_step = 0 logger.info('Starting RSI per-round GRPO (full-parameter Megatron)') @@ -1232,9 +1083,16 @@ def main(): all_tokens: List[List[int]] = [] if SOLVER_MODE == 'grpo': - # Rollout with error-feedback continuation on failed code rounds. - all_input_data, all_old_logps, all_completion_lengths = grpo_continue( - sampler, local_template, expand_prompts, sampling_params) + # One rollout per prompt; a code round that fails its asserts is asked + # to fix itself in the SAME trajectory (code_error_followup), so the + # first attempt's tokens stay trainable and its logprobs stay aligned. + all_input_data = solver_rollout(expand_prompts) + all_old_logps = [[lp[0][1] for lp in (traj.get('logprobs') or [])] + for traj in all_input_data] + # Trainable tokens, not sampled tokens: a continued rollout has two + # generations in one trajectory and the bridge between them is -100. + all_completion_lengths = [sum(1 for lb in (traj.get('labels') or []) if lb != -100) + for traj in all_input_data] else: # OPSD: single turn; also keep raw response tokens for the teacher concat. all_input_data, all_old_logps, all_completion_lengths = [], [], [] diff --git a/src/twinkle/data_format/__init__.py b/src/twinkle/data_format/__init__.py index d51f09dfa..adc540443 100644 --- a/src/twinkle/data_format/__init__.py +++ b/src/twinkle/data_format/__init__.py @@ -3,4 +3,4 @@ from .message import Message, Tool, ToolCall from .output import LossOutput, ModelOutput from .sampling import SampledSequence, SampleResponse, SamplingParams -from .trajectory import Trajectory, pack_user_data, pack_value, user_data_get +from .trajectory import Trajectory, attach_user_data, pack_user_data, pack_value, user_data_get diff --git a/src/twinkle/data_format/trajectory.py b/src/twinkle/data_format/trajectory.py index 5044bcd02..c7b277a35 100644 --- a/src/twinkle/data_format/trajectory.py +++ b/src/twinkle/data_format/trajectory.py @@ -2,7 +2,7 @@ import json import sys from collections.abc import Mapping -from typing import Any, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Tuple, Union from .message import Message, Tool @@ -45,6 +45,23 @@ def pack_user_data(values: Any) -> List[Tuple[str, str]]: return [(k, v if isinstance(v, str) else pack_value(v)) for k, v in values] +def attach_user_data(trajectory: Trajectory, **values: Any) -> Trajectory: + """Return ``trajectory`` with ``values`` merged into its packed ``user_data``. + + ``user_data`` is a list of ``(key, json_string)`` pairs rather than a dict, + so it cannot be updated in place with ``update()``; going through + :func:`pack_user_data` keeps it in the one shape readers understand. + """ + merged: Dict[str, Any] = {} + for entry in trajectory.get('user_data') or []: + if isinstance(entry, (list, tuple)) and len(entry) == 2: + merged[entry[0]] = entry[1] + merged.update(values) + out = dict(trajectory) + out['user_data'] = pack_user_data(merged) + return out + + def user_data_get(items: Any, key: str, default: Any = None) -> Any: """Look up the first value matching ``key`` in packed user_data, decoded.""" if isinstance(items, Mapping): diff --git a/src/twinkle_agentic/challenger/__init__.py b/src/twinkle_agentic/challenger/__init__.py index 39e38ad63..7c0a6261d 100644 --- a/src/twinkle_agentic/challenger/__init__.py +++ b/src/twinkle_agentic/challenger/__init__.py @@ -1,25 +1,36 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from .agentic import AgenticChallenger, AgenticPrompts, parse_check_script, parse_problem_statement -from .base import Challenger, Explorer, assistant_text, attach_user_data -from .code import (CodeChallenger, CodePrompts, KeywordStore, build_asserts, extract_code, - is_constant_answer, load_seeds, parse_challenge, run_asserts) +from .api import ApiExplorer, ApiModel +from .base import Challenger, Explorer, PromptSet, attach_user_data, map_parallel +from .code import (CodeChallenger, CodePrompts, build_asserts, is_constant_answer, load_seeds, parse_challenge, + run_asserts, run_check_script) +from .keywords import (KEYWORD_MAX_LEN, KeywordBank, KeywordPrompts, KeywordStore, parse_keyword_list, + split_keyword_list) __all__ = [ 'AgenticChallenger', 'AgenticPrompts', + 'ApiExplorer', + 'ApiModel', 'Challenger', 'CodeChallenger', 'CodePrompts', 'Explorer', + 'KEYWORD_MAX_LEN', + 'KeywordBank', + 'KeywordPrompts', 'KeywordStore', - 'assistant_text', + 'PromptSet', 'attach_user_data', 'build_asserts', - 'extract_code', 'is_constant_answer', 'load_seeds', + 'map_parallel', 'parse_check_script', 'parse_challenge', + 'parse_keyword_list', 'parse_problem_statement', 'run_asserts', + 'run_check_script', + 'split_keyword_list', ] diff --git a/src/twinkle_agentic/challenger/agentic.py b/src/twinkle_agentic/challenger/agentic.py index d8f33e7d1..d7c6bab2d 100644 --- a/src/twinkle_agentic/challenger/agentic.py +++ b/src/twinkle_agentic/challenger/agentic.py @@ -37,7 +37,6 @@ ``cookbook/rsi/agentic/prompts.py``. """ import ast -import json import math import re import threading @@ -48,113 +47,44 @@ from twinkle.data_format import SamplingParams, Trajectory, user_data_get from twinkle.utils import get_logger -from .base import Challenger, Explorer, assistant_text, attach_user_data -from .code import KeywordStore, split_keyword_list +from twinkle_agentic.utils.code_utils import PYTHON_TAGS, parse_fenced_code, strip_reasoning +from twinkle_agentic.utils.message_utils import assistant_text +from .api import ApiModel +from .base import Challenger, Explorer, PromptSet, attach_user_data, map_parallel +from .keywords import KeywordBank, KeywordStore logger = get_logger() __all__ = [ 'AgenticChallenger', 'AgenticPrompts', + 'DEFAULT_CHECK_PARSE_ERROR', + 'brittle_check_reason', 'parse_check_script', 'parse_problem_statement', ] -_FENCE_RE = re.compile(r'```(?:python|py)?\s*\n(.*?)```', re.S) # A fence around the *whole* reply, which is packaging rather than content. _WHOLE_FENCE_RE = re.compile(r'```[\w+-]*\s*\n?(.*?)```', re.S) -# The JSON body of a tool call: the proposing episode uses tools, so at the check -# stage a 4B model often keeps calling one instead of writing a fenced block. -_TOOLCALL_RE = re.compile(r'<tool_call>\s*(.*?)\s*</tool_call>', re.S) # ── parsing ─────────────────────────────────────────────────────────────── -def parse_check_script(text: str) -> Optional[str]: - """Extract a python check script from the model's reply. +def parse_check_script(text: str, language_tags: Tuple[str, ...] = PYTHON_TAGS) -> Optional[str]: + """Extract a check script from the model's reply. - Prefers the last fenced python block after ``</think>``. When the reply has - no fence at all, falls back to reading the tail as bare code: 8 of - armA2shellV6's 11 check_parse_fail rejections were a complete, parseable - check script that the model simply did not wrap in backticks, and throwing - the task away over the packaging loses a task that was ready. + The script has to be in a fenced block; see + :func:`~twinkle_agentic.utils.code_utils.parse_fenced_code`. Returns ``None`` + otherwise, including when the model reached for a tool instead of answering -- + the proposing episode had tools, so that happens, and the reply is asked for + again rather than dug through. An empty fence is refused the same way: the + model marked where the script went and put nothing there. - Returns ``None`` when nothing usable is found. + ``language_tags`` is the whole of what makes this python. Another language + passes its own, most cheaply as + ``AgenticChallenger(parse_check_fn=partial(parse_check_script, language_tags=...))``. """ - body = text or '' - idx = body.rfind('</think>') - if idx >= 0: - body = body[idx + len('</think>'):] - blocks = _FENCE_RE.findall(body) - if blocks: - script = blocks[-1].strip() - return script if script else None - bare = _bare_check_script(body) - if bare: - return bare - return _toolcall_check_script(body) - - -def _bare_check_script(body: str) -> Optional[str]: - """Read an unfenced reply as code, or None. - - Advances the start line until the rest parses, which drops whatever prose - came first (the "ALSO CORRECT:" line, a sentence introducing the script) - without needing to recognise it. Requires an ``assert`` so that a one-line - reply of prose -- which can be a syntactically valid expression -- is not - mistaken for a check. - """ - lines = body.strip().split('\n') - for start in range(len(lines)): - cand = '\n'.join(lines[start:]).strip() - if 'assert' not in cand: - break # no assert left in the tail; nothing further can qualify - try: - ast.parse(cand) - except SyntaxError: - continue - return cand - return None - - -def _toolcall_check_script(body: str) -> Optional[str]: - """Recover a check script the model put inside a tool call, or None. - - The proposing episode uses tools, and at the check stage a 4B model often - keeps calling one -- it emits ``python_executor(code="...assert...")`` (or a - shell command, or ``write_file(content=...)``) instead of a fenced block. - The script is right there in the call's ``code``/``command``/``content`` - argument, so pull it out rather than lose the task: measured on run_clean1, - most first-round check_parse_fail rejections were tool-call wrapped. - - Only code that parses and actually asserts is accepted, so a shell - ``command`` that merely runs a file -- which has no assert of its own -- - does not slip through as a check. - """ - blobs = _TOOLCALL_RE.findall(body) - for blob in reversed(blobs): - code = None - try: - obj = json.loads(blob) - args = obj.get('arguments') if isinstance(obj, dict) else None - if isinstance(args, dict): - code = args.get('code') or args.get('command') or args.get('content') - except (ValueError, AttributeError): - m = re.search(r'"(?:code|command|content)"\s*:\s*"(.*?)"\s*\}', blob, - re.S) - if m: - try: - code = m.group(1).encode().decode('unicode_escape') - except (UnicodeDecodeError, ValueError): - code = None - if not code or 'assert' not in code: - continue - try: - ast.parse(code) - except SyntaxError: - continue - return code.strip() - return None + return parse_fenced_code(text, language_tags) @@ -185,6 +115,11 @@ def brittle_check_reason(script: str) -> Optional[str]: (``c = f.read()``, then ``assert c == '...'``) so nothing sits between ``open()`` and ``==``, and a size check can put the call either around the name (``getsize("a.png")``) or after it. + + Python throughout -- the tree, the marker words, the stdlib names below. There + is no language-neutral version of this: another language keeps the two rules + but rewrites the whole body, which is why the challenger takes it as + ``brittle_check_fn`` rather than calling it directly. """ try: tree = ast.parse(script) @@ -317,8 +252,8 @@ def derived_check_literals(script: str) -> List[str]: def parse_problem_statement(text: str) -> Optional[str]: """Extract a problem statement from the model's reply. - Everything after ``</think>`` is the statement. A fence around the whole - reply is unwrapped; fences *inside* it are kept. + Everything after the model's thinking is the statement. A fence around the + whole reply is unwrapped; fences *inside* it are kept. Keeping them matters more than it sounds: a statement that says what a file must contain puts the content in a fence, and stripping every fence left @@ -329,11 +264,7 @@ def parse_problem_statement(text: str) -> Optional[str]: Returns ``None`` when the result is empty. """ - body = text or '' - idx = body.rfind('</think>') - if idx >= 0: - body = body[idx + len('</think>'):] - body = body.strip() + body = strip_reasoning(text).strip() whole = _WHOLE_FENCE_RE.fullmatch(body) if whole: body = whole.group(1).strip() @@ -366,12 +297,21 @@ def _propose_round(stage: str, trajectory: Trajectory) -> Dict[str, Any]: # ── prompts ──────────────────────────────────────────────────────────────── +# The one piece of prompt text with a default, because it was written into the +# retry path before there was a field for it and every caller relies on it. It +# names the language, so a caller working in another one has to override it. +DEFAULT_CHECK_PARSE_ERROR = ('Could not read a check script from your reply: it was not a fenced ' + 'python code block. Do not wrap it in a tool call and do not add ' + 'prose -- return ONLY a fenced python code block.') + + @dataclass -class AgenticPrompts: +class AgenticPrompts(PromptSet): """Every string an :class:`AgenticChallenger` sends. All fields are injected by the caller (no defaults with real text here). - Placeholder validation happens at construction time. + Placeholder validation, and the keyword subset a bank is given, are + :class:`.PromptSet`. """ # Explore: model acts in sandbox @@ -391,6 +331,10 @@ class AgenticPrompts: # the model can fix it from the traceback. Required only when the challenger # is built with ``check_retries`` above 0. check_retry_followup: str = '' + # The error text ``check_retry_followup`` carries when the reply held no + # readable check script at all (as opposed to one that ran and failed). + # Empty means ``DEFAULT_CHECK_PARSE_ERROR``. + check_parse_error: str = '' problem_followup: str = '' # Keyword generation (same structure as code side) @@ -398,6 +342,7 @@ class AgenticPrompts: keyword_user: str = '' keyword_expand_user: str = '' + _REQUIRED = ('system', 'from_scratch', 'check_followup', 'problem_followup') _REQUIRED_FIELDS = { 'from_seed': ('seed',), 'from_keywords': ('keywords',), @@ -408,26 +353,6 @@ class AgenticPrompts: 'keyword_expand_user': ('kw', 'm'), } - def __post_init__(self): - for name in ('system', 'from_scratch', 'check_followup', 'problem_followup'): - if not getattr(self, name).strip(): - raise ValueError(f'AgenticPrompts.{name} is required') - for name, placeholders in self._REQUIRED_FIELDS.items(): - text = getattr(self, name) - if not text: - continue - for placeholder in placeholders: - if '{' + placeholder + '}' not in text: - raise ValueError(f'AgenticPrompts.{name} must contain ' - f'{{{placeholder}}}') - - def require(self, *names: str) -> None: - """Raise unless every named prompt was supplied.""" - missing = [n for n in names if not getattr(self, n).strip()] - if missing: - raise ValueError(f'this configuration needs AgenticPrompts.' - f'{", AgenticPrompts.".join(missing)}') - # ── challenger ───────────────────────────────────────────────────────────── @@ -442,36 +367,53 @@ class AgenticChallenger(Challenger): keyword_store: optional bank for diversity control. category_desc: category -> description for keyword generation. seed_mix_prob: chance a proposal carries a seed. - reset_fn: called before each round-1 episode to clean the sandbox - workspace. Must be synchronous and leave the workspace empty. - run_check_fn: run a python script in the sandbox's current state. - Signature: ``(source: str) -> (exit_code: int, output: str)``. - workspace_snapshot_fn: after round 1, return a text summary of the - workspace state (e.g. ``find . -type f``). If None, a default - that lists messages is used. - tool_schemas: the executor's tool schemas, in the OpenAI shape the - template renders. Attached to the trajectories that are *meant* to - call tools -- the exploring episode and each solve attempt. Without - this the model is never told the tool names, so it writes code in - prose instead of calling anything: the workspace stays empty, every - check fails, and the difficulty numbers describe a model that had no - tools rather than a hard task. The check-writing and - problem-writing stages sit in the same conversation and so see the - same list, which is why the rollout stops dispatching calls once a - follow-up has been appended -- a python block written as an *answer* - parses as a call list, and 41 of 146 such replies in a measured run - edited the very workspace the answer was about. - combo_arity: ``'triple'`` or ``'mix'``, as in :class:`.CodeChallenger`. - arity_weights: weights for the ``'mix'`` subset size. - single_kw_prob: chance of using one category in ``'triple'`` mode. + envs: see :class:`~.base.Challenger`. This half asks a slot to be a real + workspace: it clears it, lets the model act in it through + :meth:`~twinkle_agentic.envs.base.Env.tool_manager`, reads the end + state back with :meth:`~twinkle_agentic.envs.base.Env.snapshot` and + runs the check script in it. ``len(envs)`` is therefore also the + episode concurrency: an episode owns its slot from the clear until + its check has run, so two episodes cannot share one, and an episode + acting in one workspace while being checked against another produces + a task nobody can pass. + parse_check_fn: read a check script out of a reply, or return None. + Defaults to :func:`parse_check_script`, which asks only that the script + be fenced python. Whether it asserts anything is the caller's to + require -- in the prompt it writes, or in the function it passes here + instead. + brittle_check_fn: why a parsed script would reject a correct solution, + or None if it would not. Defaults to :func:`brittle_check_reason`, + which reads a python syntax tree; pass ``None`` to drop the check + entirely and judge scripts only by whether they run. Both of these + and ``prompts.check_parse_error`` are the language-bound trio -- a + caller working outside python replaces all three or none. + tool_schemas: the tool contract in the OpenAI shape the template renders. + ``None`` takes it off slot 0, which is the spelling that slot will + honour; pass a list only to advertise something narrower. Attached to + the trajectories that are *meant* to call tools -- the exploring + episode and each solve attempt. Without it the model is never told the + tool names, so it writes code in prose instead of calling anything: + the workspace stays empty, every check fails, and the difficulty + numbers describe a model that had no tools rather than a hard task. + The check-writing and problem-writing stages sit in the same + conversation and so see the same list, which is why the rollout stops + dispatching calls once a follow-up has been appended -- a python block + written as an *answer* parses as a call list, and 41 of 146 such + replies in a measured run edited the very workspace the answer was + about. + combo_arity / arity_weights / single_kw_prob / keyword_refill_target / + keyword_gen_calls / keyword_refill_concurrency / keyword_refill_tries / + keyword_params / keyword_explorer / keyword_sink / min_batch: handed to + the :class:`.keywords.KeywordBank` this challenger holds, which is + where they are documented -- they behave the same on the code half. + ``keyword_explorer`` defaults to ``explorer`` here, which for a + sandbox setup means the bank brainstorms with tools live. proposals_per_group: how many proposals answer the same keyword draw and the same prompt, tagged with a shared ``group_id``. This is the group size the proposing side's advantage is computed over; at 1 every group has one member and every advantage is zero. At a fixed proposal count it does not change the compute -- it divides the number of distinct keyword draws per round by the same factor. - keyword_refill_target / keyword_gen_calls / keyword_refill_tries / - keyword_params: keyword bank refill parameters. check_params / problem_params: sampling params for the two appended stages. ``None`` keeps whatever the episode was already using, which is sized for one agent turn; the check-writing stage reads the whole @@ -491,12 +433,6 @@ class AgenticChallenger(Challenger): call (e.g. ``{'thinking_budget': N}`` to cap qwen3.8-max reasoning). ``None`` sends the request unmodified. Ignored when ``followup_api`` is ``None``. - keyword_explorer: explorer used to brainstorm keywords. Should have no - tools wired to it: a list is a text answer, and a bracketed list in - a reply is exactly what the sandbox explorer would try to dispatch as - a call. ``None`` reuses the main explorer, which for a sandbox setup - means its tools are live there too. - min_batch: smallest batch worth sending to the explorer. problem_max_chars: reject problem statements longer than this. check_retries: how many times a check script that did not pass is handed back, with the traceback and the workspace listing, for a rewrite @@ -523,13 +459,6 @@ class AgenticChallenger(Challenger): and the check's verdict. ``n_pass`` alone cannot distinguish a task that is impossible from one whose statement withholds a value its check demands, and both look like a hard task worth keeping. - keyword_sink: called once per keyword-generation call, with the prompt, - the raw reply and what ``split_keyword_list`` made of it. A bank that - refuses to fill is invisible otherwise -- proposals fall back to the - no-keyword prompt and the run carries on looking normal -- and a count - of zero does not say whether the model broke the format or the parser - rejected output that was fine, which is why the over-length phrases are - recorded next to the kept ones rather than summed into the difference. """ def __init__( @@ -541,13 +470,9 @@ def __init__( keyword_store: Optional[KeywordStore] = None, category_desc: Optional[Dict[str, str]] = None, seed_mix_prob: float = 0.5, - reset_fn: Callable[..., None], - run_check_fn: Callable[..., Tuple[int, str]], - workspace_snapshot_fn: Optional[Callable[..., str]] = None, - snapshot_error_fn: Optional[Callable[..., str]] = None, + parse_check_fn: Callable[[str], Optional[str]] = parse_check_script, + brittle_check_fn: Optional[Callable[[str], Optional[str]]] = brittle_check_reason, tool_schemas: Optional[Sequence[Dict[str, Any]]] = None, - episode_concurrency: int = 1, - episode_tool_managers: Optional[Sequence[Any]] = None, combo_arity: str = 'triple', arity_weights: Optional[Sequence[float]] = None, single_kw_prob: float = 0.1, @@ -579,59 +504,39 @@ def __init__( **challenger_kwargs: Any, ): super().__init__(explorer, system=prompts.system, **challenger_kwargs) - if combo_arity not in ('triple', 'mix'): - raise ValueError(f"combo_arity must be 'triple' or 'mix', got {combo_arity!r}") + if not self.envs: + raise ValueError('envs is empty: an episode here needs a workspace to act in ' + 'and a check to be run against, so there is nothing this ' + 'challenger could measure.') if keyword_store is not None: - desc = category_desc or {} - missing_cats = [c for c in keyword_store.categories if not desc.get(c)] - if missing_cats: - raise ValueError(f'category_desc is missing a description for ' - f'{missing_cats}; a dry category could not be refilled.') - prompts.require('keyword_system', 'keyword_user', 'from_keywords') + prompts.require('from_keywords') self.prompts = prompts self.seeds = list(seeds) - self.store = keyword_store - self.category_desc = dict(category_desc or {}) + # The whole keyword cycle -- draw, refill, expand -- is one object shared + # with the code challenger rather than a second copy of it here. None means + # no bank was configured, and proposals then carry no topics. + self.keywords: Optional[KeywordBank] = None if keyword_store is None else KeywordBank( + keyword_store, prompts=prompts.keyword_prompts(), + category_desc=category_desc or {}, + # Brainstorming a list is a text round: the sandbox-tool explorer would + # waste turns on it and could take a bracketed list for a tool call. + explorer=keyword_explorer or explorer, rng=self.rng, + name=type(self).__name__, sampling_params=keyword_params, + sink=keyword_sink, combo_arity=combo_arity, arity_weights=arity_weights, + single_kw_prob=single_kw_prob, refill_target=keyword_refill_target, + gen_calls=keyword_gen_calls, refill_concurrency=keyword_refill_concurrency, + refill_tries=keyword_refill_tries, min_batch=min_batch) self.seed_mix_prob = seed_mix_prob - self.reset_fn = reset_fn - self.run_check_fn = run_check_fn - self.workspace_snapshot_fn = workspace_snapshot_fn - # Asked, when a snapshot came back empty, why: the text of the failure if - # the listing could not be read, '' if the workspace really was empty. - # Without it the two are one outcome, and a sandbox the host had paused is - # filed as the model having built nothing -- 63 of run_clean6's 71 - # ``empty_workspace`` rejections were the 410 "sandbox is not proxyable" - # error, so that reject class was 89% broken environment. - self.snapshot_error_fn = snapshot_error_fn - self.tool_schemas = list(tool_schemas) if tool_schemas else None - # More than one episode at a time needs more than one sandbox: an episode - # owns its workspace from the reset until its check has run. The three - # sandbox callables above are then called with ``slot=i`` to say which one, - # and ``episode_tool_managers[i]`` must dispatch tool calls into that same - # sandbox -- an episode acting in one workspace and checking another - # produces a task whose check nobody can pass. - if episode_concurrency < 1: - raise ValueError(f'episode_concurrency must be >= 1, got {episode_concurrency}') - if episode_concurrency > 1: - if not episode_tool_managers or len(episode_tool_managers) != episode_concurrency: - raise ValueError( - f'episode_concurrency={episode_concurrency} needs exactly that many ' - f'episode_tool_managers, one per sandbox; got ' - f'{len(episode_tool_managers) if episode_tool_managers else 0}.') - self.episode_concurrency = episode_concurrency - self.episode_tool_managers = (list(episode_tool_managers) - if episode_tool_managers else None) + self.parse_check_fn = parse_check_fn + self.brittle_check_fn = brittle_check_fn + # Read off slot 0 by default: these go into the prompt, and taking them + # from the environment that will execute them is what makes it impossible + # for the advertised contract and the running code to disagree. + self.tool_schemas = list(tool_schemas) if tool_schemas else (self.env().tools() or None) # Held while writing to the dump files and while bumping ``stats``: with # concurrent episodes those are the only shared mutable things the # follow-up callback touches, and a half-written json line is unreadable. self._sink_lock = threading.Lock() - # Separate from the sink lock: the keyword path holds this while it draws - # from the shared rng and bumps the prompt nonce, and it must not be held - # while a sink write is waiting on disk. - self._kw_lock = threading.Lock() - self.combo_arity = combo_arity - self.arity_weights = list(arity_weights) if arity_weights else None - self.single_kw_prob = single_kw_prob # How many proposals answer each keyword draw. Above 1 they form a GRPO # group on the proposing side; see :meth:`propose`. Raising it does not # cost more compute at a fixed proposal count -- it trades keyword @@ -641,20 +546,6 @@ def __init__( raise ValueError(f'proposals_per_group must be >= 1, got {proposals_per_group}') self.proposals_per_group = proposals_per_group self._next_group_id = 0 - self.keyword_refill_target = keyword_refill_target - self.keyword_gen_calls = keyword_gen_calls - # How many of a refill's generating calls go out together. At 1 each call - # is told what the ones before it produced, which is the point; raising it - # is what the first round of arm measurements ran with, where a whole - # first refill went out at once with nothing yet to avoid and came back - # with synonyms. Kept configurable so the two can be compared on one build - # rather than across two versions of this file. - if keyword_refill_concurrency < 1: - raise ValueError('keyword_refill_concurrency must be >= 1, got ' - f'{keyword_refill_concurrency}') - self.keyword_refill_concurrency = keyword_refill_concurrency - self.keyword_refill_tries = keyword_refill_tries - self.keyword_params = keyword_params self.check_params = check_params self.problem_params = problem_params # When set, exploration runs on the (trainable) local explorer as before, @@ -665,12 +556,11 @@ def __init__( # ``logprobs`` stay exactly the exploration turns the local sampler # produced -- which is what "train only the exploration part" means. The # generated check script and statement are used solely to build the task. - self.followup_api = followup_api - # extra_body sent on every followup API call (e.g. {'thinking_budget': N} - # to cap qwen3.8-max reasoning). None sends the request unmodified. - self.followup_extra_body = dict(followup_extra_body) if followup_extra_body else None - self.keyword_explorer = keyword_explorer - self.min_batch = max(1, min_batch) + # None means the single-model path, where the local model writes those two + # stages in the same conversation. ``followup_extra_body`` rides along on + # every call (e.g. {'thinking_budget': N} to cap qwen3.8-max reasoning). + self.followup_model: Optional[ApiModel] = None if followup_api is None else ApiModel( + followup_api, extra_body=followup_extra_body, name=type(self).__name__) self.problem_max_chars = problem_max_chars # A budget in proposals rather than in kept tasks, for runs whose purpose # is to measure what the current configuration produces: with a keep-rate @@ -734,12 +624,10 @@ def __init__( self.reject_sink = reject_sink self.propose_sink = propose_sink self.solver_sink = solver_sink - self.keyword_sink = keyword_sink if self.seeds: prompts.require('from_seed') - if self.store is not None: + if self.keywords is not None: prompts.require('from_seed_keywords') - self._nonce = 0 self.stats: Dict[str, int] = { 'explore_done': 0, 'check_parse_fail': 0, 'check_run_fail': 0, 'empty_workspace': 0, 'solver_truncated': 0, @@ -785,7 +673,6 @@ def __init__( # their proposals. Nothing from them is used. 'novelty_group_dropped': 0, } - self._hard: List[Tuple[str, str]] = [] # ------------------------------------------------------------- proposing @@ -809,8 +696,8 @@ def propose(self, count: int) -> List[Trajectory]: metas: List[Tuple[List[Tuple[str, str]], bool, str, int]] = [] per_group = max(1, self.proposals_per_group) while len(metas) < count: - picks = self._draw_keywords() - body = '\n'.join(f'- {c}: {t}' for c, t in picks) + picks = self.keywords.draw() if self.keywords else [] + body = KeywordBank.block(picks) use_seed = bool(self.seeds) and self.rng.random() < self.seed_mix_prob seed = self.rng.choice(self.seeds) if use_seed else None if use_seed and picks: @@ -848,7 +735,6 @@ def propose(self, count: int) -> List[Trajectory]: group_id=gid)) return proposals - # ------------------------------------------------------------- building def build(self, explored: List[Trajectory]) -> List[Optional[Trajectory]]: @@ -863,16 +749,18 @@ def build(self, explored: List[Trajectory]) -> List[Optional[Trajectory]]: f'{type(self).__name__}.build() must not be called directly; ' f'the serial _round() loop drives one episode at a time instead.') - def _reject_for_empty_snapshot(self, state: Dict[str, Any], slot: int) -> None: + def _reject_for_empty_snapshot(self, state: Dict[str, Any], detail: str) -> None: """File an episode whose workspace listing came back empty. Empty means one of two unrelated things -- the episode built nothing, or the listing could not be read -- and only the first says anything about - the model. ``snapshot_error_fn`` is what tells them apart; with no such - callback every case is filed as ``empty_workspace``, which is what used - to happen for all of them. + the model. ``detail`` is the second half of what + :meth:`~twinkle_agentic.envs.base.Env.snapshot` returns, and it is what + tells them apart: without it every case is filed as ``empty_workspace``, + which is what used to happen for all of them -- 63 of run_clean6's 71 + ``empty_workspace`` rejections were the 410 "sandbox is not proxyable" + error, so that reject class was 89% broken environment. """ - detail = self.snapshot_error_fn(slot=slot) if self.snapshot_error_fn else '' if detail: self._bump('snapshot_unavailable') state['reject'] = ('snapshot_unavailable', detail) @@ -898,7 +786,7 @@ def _followup(self, state: Dict[str, Any], trajectory: Trajectory, """ slot = state.get('slot', 0) if n_before == 0: - snapshot = self.workspace_snapshot_fn(slot=slot) if self.workspace_snapshot_fn else '' + snapshot, snapshot_error = self.env(slot).snapshot() state['snapshot'] = snapshot # An episode that left nothing behind has no end state to write checks # about, and asking for them anyway is worse than useless: the only @@ -906,7 +794,7 @@ def _followup(self, state: Dict[str, Any], trajectory: Trajectory, # solver passes by doing nothing. Five of run5's ten verified tasks # were that task. Reject here instead. if not snapshot.strip(): - self._reject_for_empty_snapshot(state, slot) + self._reject_for_empty_snapshot(state, snapshot_error) return None return (self.prompts.check_followup.format(final_state=snapshot), self.check_params) @@ -917,7 +805,7 @@ def _followup(self, state: Dict[str, Any], trajectory: Trajectory, attempt = state.get('check_attempts', 0) + 1 state['check_attempts'] = attempt reply = assistant_text(trajectory) - script = parse_check_script(reply) + script = self.parse_check_fn(reply) if script is None: # Same one-rewrite budget a run failure gets: hand the parse # failure back and let it regenerate, rather than dropping a task @@ -926,10 +814,7 @@ def _followup(self, state: Dict[str, Any], trajectory: Trajectory, # extra tries, not one each. if attempt <= self.check_retries: self._bump('check_retry') - err = ('Could not read a check script from your reply: it was ' - 'not a fenced python code block. Do not wrap it in a ' - 'tool call and do not add prose -- return ONLY a fenced ' - 'python code block.') + err = self.prompts.check_parse_error or DEFAULT_CHECK_PARSE_ERROR return (self.prompts.check_retry_followup.format( error=err, final_state=state.get('snapshot') or ''), self.check_params) @@ -941,14 +826,14 @@ def _followup(self, state: Dict[str, Any], trajectory: Trajectory, state['reject'] = ('check_parse_fail', reply) return None state['script'] = script - brittle = brittle_check_reason(script) + brittle = self.brittle_check_fn(script) if self.brittle_check_fn else None if brittle is not None: # Same bookkeeping as a check that ran and failed: the script is # rejected before it can pass on the author's own state, because # passing there is exactly what hides the defect. exit_code, output = 1, brittle else: - exit_code, output = self.run_check_fn(script, slot=slot) + exit_code, output = self.env(slot).run_script(script) if exit_code == 0: state['checked'] = True if attempt > 1: @@ -973,7 +858,7 @@ def _followup(self, state: Dict[str, Any], trajectory: Trajectory, # something untrue, or the workspace changed under it -- and the # difference is visible only in the state at the moment the check # ran. It is also what the rewrite gets to read. - after = self.workspace_snapshot_fn(slot=slot) if self.workspace_snapshot_fn else '' + after = self.env(slot).snapshot()[0] state.setdefault('attempts', []).append( f'--- attempt {attempt}: exit {exit_code} ---\n{output}\n' f'--- check script ---\n{script}') @@ -992,36 +877,6 @@ def _followup(self, state: Dict[str, Any], trajectory: Trajectory, return None - def _api_reply(self, messages: List[Dict[str, Any]], user_text: str, - params: Optional[SamplingParams]) -> Optional[str]: - """Append ``user_text`` and one ``followup_api`` reply to ``messages``. - - ``messages`` is a throwaway copy owned by :meth:`_run_followup_api`, never - the trainable trajectory, so mutating it in place costs the model nothing. - Returns the assistant text, or ``None`` when the API call raised -- the - caller then rejects rather than building a task on a broken conversation. - - Tools are withdrawn for these stages on purpose (they are answers, not - actions), so only the text is kept; any structured ``tool_calls`` the API - returned are dropped. - """ - messages.append({'role': 'user', 'content': user_text}) - request: Trajectory = {'messages': messages} - try: - if self.followup_extra_body: - reply = self.followup_api(request, params, extra_body=self.followup_extra_body) - else: - reply = self.followup_api(request, params) - except Exception as exc: # noqa: BLE001 -- one bad call must not kill the round - logger.warning(f'[{type(self).__name__}] followup API call failed: ' - f'{type(exc).__name__}: {exc}') - return None - if isinstance(reply, list): - reply = reply[0] if reply else {} - content = (reply.get('content') if isinstance(reply, dict) else None) or '' - messages.append({'role': 'assistant', 'content': content}) - return content - def _run_followup_api(self, state: Dict[str, Any], explored: Trajectory) -> None: """Generate the check script and problem statement over ``followup_api``. @@ -1044,12 +899,12 @@ def _run_followup_api(self, state: Dict[str, Any], explored: Trajectory) -> None slot = state.get('slot', 0) messages: List[Dict[str, Any]] = [dict(m) for m in explored.get('messages') or []] - snapshot = self.workspace_snapshot_fn(slot=slot) if self.workspace_snapshot_fn else '' + snapshot, snapshot_error = self.env(slot).snapshot() state['snapshot'] = snapshot # An episode that left nothing behind has no end state to write checks # about; rejecting here mirrors the n_before==0 branch of _followup. if not snapshot.strip(): - self._reject_for_empty_snapshot(state, slot) + self._reject_for_empty_snapshot(state, snapshot_error) return # Check-script stage: the first ask plus up to ``check_retries`` rewrites, @@ -1060,19 +915,16 @@ def _run_followup_api(self, state: Dict[str, Any], explored: Trajectory) -> None while True: attempt += 1 state['check_attempts'] = attempt - reply = self._api_reply(messages, user_text, self.check_params) + reply = self.followup_model.reply(messages, user_text, self.check_params) if reply is None: self._bump('followup_api_error') state['reject'] = ('followup_api_error', 'check-script API call failed') return - script = parse_check_script(reply) + script = self.parse_check_fn(reply) if script is None: if attempt <= self.check_retries: self._bump('check_retry') - err = ('Could not read a check script from your reply: it was ' - 'not a fenced python code block. Do not wrap it in a ' - 'tool call and do not add prose -- return ONLY a fenced ' - 'python code block.') + err = self.prompts.check_parse_error or DEFAULT_CHECK_PARSE_ERROR user_text = self.prompts.check_retry_followup.format( error=err, final_state=snapshot) continue @@ -1080,13 +932,13 @@ def _run_followup_api(self, state: Dict[str, Any], explored: Trajectory) -> None state['reject'] = ('check_parse_fail', reply) return state['script'] = script - brittle = brittle_check_reason(script) + brittle = self.brittle_check_fn(script) if self.brittle_check_fn else None if brittle is not None: # Rejected before it can pass on the author's own state, since # passing there is exactly what hides the defect. exit_code, output = 1, brittle else: - exit_code, output = self.run_check_fn(script, slot=slot) + exit_code, output = self.env(slot).run_script(script) if exit_code == 0: state['checked'] = True if attempt > 1: @@ -1104,7 +956,7 @@ def _run_followup_api(self, state: Dict[str, Any], explored: Trajectory) -> None return state['setup_script'] = setup break - after = self.workspace_snapshot_fn(slot=slot) if self.workspace_snapshot_fn else '' + after = self.env(slot).snapshot()[0] state.setdefault('attempts', []).append( f'--- attempt {attempt}: exit {exit_code} ---\n{output}\n' f'--- check script ---\n{script}') @@ -1122,8 +974,8 @@ def _run_followup_api(self, state: Dict[str, Any], explored: Trajectory) -> None return # Problem-statement stage: one API reply, kept as the task's statement. - reply = self._api_reply(messages, self.prompts.problem_followup, - self.problem_params) + reply = self.followup_model.reply(messages, self.prompts.problem_followup, + self.problem_params) if reply is None: self._bump('followup_api_error') state['reject'] = ('followup_api_error', 'problem-statement API call failed') @@ -1179,7 +1031,7 @@ def reject(reason: str, detail: str = '') -> None: # ``explored`` (whose last assistant turn is the final exploration reply); # it lives in ``state``. The single-model path keeps it as the last # assistant message of the episode. - if self.followup_api is not None: + if self.followup_model is not None: statement = parse_problem_statement(state.get('statement') or '') else: statement = parse_problem_statement(assistant_text(explored)) @@ -1379,30 +1231,25 @@ def _bump(self, key: str, n: int = 1) -> None: with self._sink_lock: self.stats[key] += n - def _parallel(self, fn: Callable[[Any], Any], items: Sequence[Any]) -> List[Any]: - """Map ``fn`` over ``items`` at once, results in input order. + def _tool_manager(self, slot: int) -> Optional[Any]: + """The dispatcher for ``slot``'s tools, or None when it advertises none. - Every use of this is waiting on a sandbox, not computing, so the thread - pool is the point. One item runs inline: a pool for a single sandbox call - only adds a thread, and it keeps the serial configuration on exactly the - same code path it had before. + Built per use rather than held, so a slot rebuilt underneath -- evicted, + timed out -- is dispatched into as it is now: a manager captured at + construction would keep sending this episode's calls to a sandbox that is + gone. None means this environment offers no tools, which is the honest + answer for one that only runs scripts, and the rollout then leaves the + model with none rather than an empty tool list it would try to call. """ - items = list(items) - if len(items) <= 1: - return [fn(item) for item in items] - out: List[Any] = [None] * len(items) - with ThreadPoolExecutor(max_workers=len(items)) as pool: - futures = {pool.submit(fn, item): i for i, item in enumerate(items)} - for fut in as_completed(futures): - out[futures[fut]] = fut.result() - return out + env = self.env(slot) + return env.tool_manager() if env.tools() else None def _run_episode(self, proposal: Trajectory, slot: int) -> Optional[Trajectory]: """One episode top-to-bottom, using sandbox slot ``slot``.""" - self.reset_fn(slot=slot) + self.env(slot).clear() state: Dict[str, Any] = {'slot': slot} - tm = self.episode_tool_managers[slot] if self.episode_tool_managers else None - if self.followup_api is not None: + tm = self._tool_manager(slot) + if self.followup_model is not None: # Split path: explore on the local (trainable) model with NO # followup_fn, so the rollout ends the moment the model stops calling # tools and the returned trajectory carries only the exploration @@ -1454,7 +1301,7 @@ def _round(self, missing: int) -> Optional[List[Trajectory]]: return None usable: List[Trajectory] = [] - n_slots = self.episode_concurrency + n_slots = self.n_slots if n_slots <= 1 or len(proposals) <= 1: # Serial fallback (original path). @@ -1504,8 +1351,8 @@ def _drain(slot: int) -> List[Trajectory]: def _filter_difficulty(self, tasks: List[Trajectory]) -> List[Trajectory]: """Override: every solver attempt needs its own clean workspace. - Attempts are run in waves of ``episode_concurrency``, attempt k of a wave - in sandbox slot k. Within a wave all attempts go out in one explorer call, + Attempts are run in waves of ``len(envs)``, attempt k of a wave in slot k. + Within a wave all attempts go out in one explorer call, so the sampler generates them as one batch instead of leaving the GPUs waiting on a single sequence, and the wave's clears, input replays and checks all run at the same time too -- they are sandbox round-trips, not @@ -1528,7 +1375,7 @@ def _filter_difficulty(self, tasks: List[Trajectory]) -> List[Trajectory]: if not tasks: return [] passes = [0] * len(tasks) - n_slots = max(1, self.episode_concurrency) + n_slots = max(1, self.n_slots) # Which task each attempt belongs to, flattened, so a wave is a fixed # number of sandboxes no matter how attempts distribute over tasks. plan = [i for i in range(len(tasks)) for _ in range(self.solver_rollouts)] @@ -1541,10 +1388,10 @@ def _filter_difficulty(self, tasks: List[Trajectory]) -> List[Trajectory]: def _prepare(k: int) -> bool: """Clear slot k, then put back the inputs this task hands out.""" - self.reset_fn(slot=k) + self.env(k).clear() if not setups[k]: return True - exit_code, output = self.run_check_fn(setups[k], slot=k) + exit_code, output = self.env(k).run_script(setups[k]) if exit_code != 0: # Measuring this attempt against a workspace missing its # inputs would score the task as harder than it is, so the @@ -1554,15 +1401,16 @@ def _prepare(k: int) -> bool: return False return True - ready = self._parallel(_prepare, slots) + ready = map_parallel(_prepare, slots) live = [k for k in slots if ready[k]] self._bump('setup_replay_fail', len(slots) - len(live)) if not live: continue prompts = [dict(self.solver_prompt(tasks[wave[k]])) for k in live] kwargs: Dict[str, Any] = {} - if self.episode_tool_managers: - kwargs['tool_manager'] = [self.episode_tool_managers[k] for k in live] + managers = [self._tool_manager(k) for k in live] + if any(tm is not None for tm in managers): + kwargs['tool_manager'] = managers attempts = self._solver_explore(prompts, sampling_params=self.solver_params, **kwargs) if len(attempts) != len(prompts): @@ -1574,7 +1422,7 @@ def _prepare(k: int) -> bool: for attempt in attempts: if attempt is not None and attempt.get('truncated'): self._bump('solver_truncated') - verdicts = self._parallel( + verdicts = map_parallel( lambda j: (attempts[j] is not None and self.judge_attempt(tasks[wave[live[j]]], attempts[j], slot=live[j])), @@ -1589,8 +1437,8 @@ def _prepare(k: int) -> bool: ] self.on_difficulty_measured(measured) novelties = self._score_novelty(measured) - high = self.solver_rollouts - self.keep_max_pass_margin - in_band = [self.keep_min_pass <= n <= high for n in passes] + low, high = self.keep_pass_band + in_band = [low <= n <= high for n in passes] # Which of the in-band tasks the solver side actually trains on. Decided # before emitting so each proposal's record says whether its task was taken. selected = self._select_per_group(measured, passes, in_band, novelties) @@ -1801,7 +1649,7 @@ def judge_attempt(self, task: Trajectory, attempt: Trajectory, script = user_data_get(task.get('user_data'), 'check_script', '') if not script: return False - exit_code, output = self.run_check_fn(script, slot=slot) + exit_code, output = self.env(slot).run_script(script) if self.solver_sink is not None: messages = task.get('messages') or [{}] record = { @@ -1816,234 +1664,19 @@ def judge_attempt(self, task: Trajectory, attempt: Trajectory, # be reproducible from the dump. 'truncated': bool((attempt or {}).get('truncated')), 'attempt': attempt, - 'end_state': (self.workspace_snapshot_fn(slot=slot) - if self.workspace_snapshot_fn else ''), + 'end_state': self.env(slot).snapshot()[0], } with self._sink_lock: self.solver_sink(record) return exit_code == 0 def on_difficulty_measured(self, candidates: List[Trajectory]) -> None: - """Remember keywords behind candidates nobody solved.""" - if self.store is None: - return - seen = {(c, t.lower()) for c, t in self._hard} - for task in candidates: - data = task.get('user_data') - if user_data_get(data, 'n_pass', 0) > 0: - continue - for pick in user_data_get(data, 'keywords', []) or []: - if isinstance(pick, (list, tuple)) and len(pick) >= 2: - c, t = pick[0], pick[1] - if (c, t.lower()) not in seen: - seen.add((c, t.lower())) - self._hard.append((c, t)) - - # ------------------------------------------------------------ keywords - - def _draw_keywords(self) -> List[Tuple[str, str]]: - """Consume one keyword combination from the bank; [] without a bank.""" - if self.store is None: - return [] - categories = self.store.categories - if self.combo_arity == 'mix': - if self.arity_weights and len(self.arity_weights) == len(categories): - k = self.rng.choices(range(1, len(categories) + 1), - weights=self.arity_weights)[0] - else: - k = self.rng.randint(1, len(categories)) - cats = self.rng.sample(list(categories), k) - elif self.rng.random() < self.single_kw_prob: - cats = [self.rng.choice(categories)] - else: - cats = list(categories) - picks: List[Tuple[str, str]] = [] - # Refill every dry category at once rather than as each one is reached: the - # three refills are independent model calls that used to run one after - # another (20s each at the start of a run), and they touch separate - # entries of the bank. - dry = [c for c in cats if not self.store.unused(c)] - if dry: - self._parallel(self._refill, dry) - for c in cats: - text = self.store.take(c, self.rng) - if text is not None: - picks.append((c, text)) - return picks - - def _refill(self, category: str) -> None: - """Ask the model for more keywords in ``category``. - - Says so when it comes back empty. A silent no-op here is the worst - outcome available: ``_draw_keywords`` then hands out no keywords, every - proposal quietly falls back to the from-scratch prompt, and the run looks - normal while producing one identical prompt over and over. That is exactly - what happened for whole runs when the prompt asked for one keyword per - line and the parser wanted a JSON array. - """ - tries = 0 - while not self.store.unused(category): - new = self._generate_keywords(category, self.keyword_refill_target) - added = self.store.add(category, new, source='gen') - tries += 1 - if added: - logger.info(f'[AgenticChallenger] keyword category {category!r} ' - f'refilled +{added} (try {tries})') - continue - logger.warning( - f'[AgenticChallenger] keyword refill for {category!r} produced ' - f'nothing on try {tries}: {len(new)} parsed, 0 new. Proposals will ' - f'run without keywords unless this recovers -- pass keyword_sink ' - f'to see the replies.') - if tries >= self.keyword_refill_tries: - if self.store.items[category]: - self.store.recycle(category) - logger.info(f'[AgenticChallenger] keyword category {category!r} ' - f'exhausted -> recycled {len(self.store.items[category])} topics') - break - - def _generate_keywords(self, category: str, n_want: int) -> List[str]: - """Up to ``n_want`` keywords the bank does not already hold. - - Runs on ``keyword_explorer`` when there is one: brainstorming a list is a - text round, and putting it through the sandbox-tool explorer both wastes - turns and lets a bracketed list in the reply be taken for a tool call. - """ - if n_want <= 0: - return [] - known = self.store.texts(category) - n_calls = max(self.keyword_gen_calls, self.min_batch) - per_call = max(1, -(-n_want // n_calls) + 4) - seen = {t.strip().lower() for t in known} - out: List[str] = [] - explorer = self.keyword_explorer or self.explorer - for start in range(0, n_calls, self.keyword_refill_concurrency): - group = range(start, min(start + self.keyword_refill_concurrency, n_calls)) - # Every call in a group is built before any of them runs, so they all - # carry the same avoid list -- which is exactly the batched behaviour, - # and why a group of one is what lets call k+1 see call k. - users = [(self.prompts.keyword_user.format( - k=per_call, desc=self.category_desc[category]) - + self._avoid_note(known, out, - '\nDo NOT repeat any of these already-used topics: ') - + f'\n(batch {self._next_nonce()}-{i})') for i in group] - prompts = [{ - 'messages': [{'role': 'system', 'content': self.prompts.keyword_system}, - {'role': 'user', 'content': u}], - } for u in users] - for user, reply in zip(users, explorer(prompts, - sampling_params=self.keyword_params)): - text = assistant_text(reply) - parsed, dropped_long = split_keyword_list(text) - fresh = [] - for kw in parsed: - key = kw.lower() - if key not in seen: - seen.add(key) - fresh.append(kw) - out.extend(fresh) - if self.keyword_sink is not None: - # Full text, both sides. The one question this dump exists to - # answer -- did the model disobey the format, or does the parser - # reject what it produced -- cannot be answered from a count. - record = { - 'category': category, - 'prompt': user, - 'reply': text, - 'stop_reason': reply.get('stop_reason'), - 'truncated': bool(reply.get('truncated')), - 'parsed': parsed, - 'n_parsed': len(parsed), - 'n_new': len(fresh), - # The two fields that make the sentence above true. Without - # them ``n_parsed: 0`` reads the same whether the reply was - # garbled, empty, or eight usable keywords written at - # sentence length -- and the third is the one that happened. - 'dropped_long': dropped_long, - 'n_dropped_long': len(dropped_long), - } - with self._sink_lock: - self.keyword_sink(record) - with self._kw_lock: - self.rng.shuffle(out) - return out[:n_want] - - # How many phrases the 'do not repeat these' line may quote in total. There - # has to be a ceiling in both directions: too few and a serial refill stops - # seeing what it just said, too many and the model runs out of room to obey. - # Measured on armA2ser, where this refill's own output went in uncapped: with - # 130 quoted the eighth call was still answering normally, with 150 it started - # inventing -- 'îRAPIÓN holistic replace', 'ซะ subspace cutter map limit', 10 - # of 480 phrases that run. 100 sits below where that began. - _AVOID_TOTAL = 100 - - def _next_nonce(self) -> int: - """A number no other call gets, so two prompts are never byte-identical. - - Shared across categories, which refill at the same time: two threads - reading the counter together would send the same prompt twice and halve - the diversity with nothing to show that it happened. - """ - with self._kw_lock: - self._nonce += 1 - return self._nonce - - def _avoid_note(self, older: List[str], fresh: List[str], lead: str) -> str: - """The 'do not repeat these' line, newest first, capped at ``_AVOID_TOTAL``. - - What this refill has just produced comes first and evicts older entries - rather than the reverse -- the calls run one at a time so that each can - avoid what the ones before it said, and dropping those would undo it. Past - the cap the oldest of *this refill's* phrases are what falls off, which is - also the least costly thing to drop: the model has already moved away from - them. - """ - fresh_shown = list(fresh)[-self._AVOID_TOTAL:] - room = max(0, self._AVOID_TOTAL - len(fresh_shown)) - with self._kw_lock: - shown = older if len(older) <= room else self.rng.sample(older, room) - avoid = fresh_shown + list(shown) - return lead + ', '.join(avoid) if avoid else '' - - + """Remember the topics behind the candidates nobody solved.""" + if self.keywords is not None: + self.keywords.remember_unsolved(candidates) # ------------------------------------------------------------ feedback def expand_hard_keywords(self) -> int: - """Brainstorm more topics in families that produced the hardest tasks.""" - if self.store is None or not self._hard or not hasattr(self.prompts, 'keyword_expand_user'): - return 0 - self.prompts.require('keyword_expand_user') - hard = self._hard[:32] - self.rng.shuffle(hard) - reqs = list(hard) - while len(reqs) < self.min_batch: - reqs.append(hard[len(reqs) % len(hard)]) - self._nonce += 1 - prompts = [{ - 'messages': [ - {'role': 'system', 'content': self.prompts.keyword_system}, - {'role': 'user', - 'content': self.prompts.keyword_expand_user.format(kw=kw, m=8) - + f'\n(batch {self._nonce}-{i})'}, - ], - } for i, (_c, kw) in enumerate(reqs)] - added = 0 - explorer = self.keyword_explorer or self.explorer - for (cat, kw), reply in zip(reqs, explorer(prompts, - sampling_params=self.keyword_params)): - text = assistant_text(reply) - parsed, dropped_long = split_keyword_list(text) - added += self.store.add(cat, parsed, source='expand', parent=kw) - if self.keyword_sink is not None: - self.keyword_sink({ - 'category': cat, 'parent': kw, 'reply': text, - 'stop_reason': reply.get('stop_reason'), - 'truncated': bool(reply.get('truncated')), - 'parsed': parsed, 'n_parsed': len(parsed), - 'dropped_long': dropped_long, - 'n_dropped_long': len(dropped_long), - }) - logger.info(f'[AgenticChallenger] expanded {len(hard)} hard keyword(s) -> ' - f'+{added} same-domain topics') - return added + """Brainstorm more topics in the families that produced the hardest tasks.""" + return self.keywords.expand_hard() if self.keywords is not None else 0 diff --git a/src/twinkle_agentic/challenger/api.py b/src/twinkle_agentic/challenger/api.py new file mode 100644 index 000000000..bb1141e2a --- /dev/null +++ b/src/twinkle_agentic/challenger/api.py @@ -0,0 +1,154 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""An OpenAI-compatible API, reached as if it were one more explorer. + +Both halves of this loop hand some rounds to a stronger model over an API -- the +ones that are answers rather than actions: writing a check script, describing a +task, brainstorming keywords. The rule that decides what may go over the API is +that the reply must not enter a trainable trajectory, and these do not. + +The call itself was written twice, here and in ``cookbook/rsi``, in the same +twenty lines each time: append a user message, send, keep the text, survive a +raised exception. :class:`ApiModel` is that call, once. :class:`ApiExplorer` wraps +it in the :data:`~.base.Explorer` signature, so anything here that takes an +explorer -- the keyword bank above all -- can be pointed at the API without +knowing it is one, and can be handed a local explorer to fall back on when the +API is unreachable. +""" +from typing import Any, Dict, List, Optional, Sequence + +from twinkle.data_format import SamplingParams, Trajectory +from twinkle.utils import get_logger +from .base import Explorer + +logger = get_logger() + +__all__ = ['ApiExplorer', 'ApiModel'] + + +class ApiModel: + """One OpenAI-compatible client, the extra body it always sends, and one call. + + Args: + api: the client, called as ``api(request, params)``, or with + ``extra_body=`` when there is one. ``twinkle_agentic.protocol.openai`` + provides one; anything with that signature will do. + extra_body: sent on every call (e.g. ``{'thinking_budget': N}`` to cap a + reasoning model). ``None`` sends the request unmodified. + name: what log lines call this model, normally the caller's class name. + """ + + def __init__(self, api: Any, *, extra_body: Optional[Dict[str, Any]] = None, + name: str = 'api'): + self.api = api + self.extra_body = dict(extra_body) if extra_body else None + self.name = name + + def generate(self, messages: Sequence[Dict[str, Any]], + params: Optional[SamplingParams] = None) -> Optional[str]: + """One reply to ``messages``, as text. ``None`` means the call raised. + + Returning ``None`` rather than raising is what keeps one unreachable call + from ending a run that has hours of sandbox work behind it; every caller + here either rejects that one item or falls back. + + Tools are withdrawn for these rounds on purpose -- they are answers, not + actions -- so only the text is kept and any structured ``tool_calls`` the + API returned are dropped. + """ + request: Trajectory = {'messages': list(messages)} + try: + if self.extra_body: + reply = self.api(request, params, extra_body=self.extra_body) + else: + reply = self.api(request, params) + except Exception as exc: # noqa: BLE001 -- one bad call must not kill the run + logger.warning(f'[{self.name}] API call failed: {type(exc).__name__}: {exc}') + return None + if isinstance(reply, list): + reply = reply[0] if reply else {} + return (reply.get('content') if isinstance(reply, dict) else None) or '' + + def reply(self, messages: List[Dict[str, Any]], user_text: str, + params: Optional[SamplingParams] = None) -> Optional[str]: + """Append ``user_text`` and one reply to ``messages``; return the reply. + + For the staged conversations: a check script asked for over the end state, + then a statement asked for over the check. ``messages`` is the caller's + private copy, never a trainable trajectory, so mutating it in place costs + the model nothing. A failed call leaves the user message appended and no + assistant message, which is what the caller would have to write out by + hand to retry. + """ + messages.append({'role': 'user', 'content': user_text}) + content = self.generate(messages, params) + if content is None: + return None + messages.append({'role': 'assistant', 'content': content}) + return content + + +class ApiExplorer: + """An :data:`~.base.Explorer` that answers single text rounds over an API model. + + For the keyword bank, whose calls are one round each and whose replies are + parsed into a list and thrown away: no tokens of them are ever trained on, so + a stronger model may write them. That matters more than it sounds. The bank is + the single input every task downstream is built from, and a 4B policy at the + temperature diversity needs is the wrong instrument for a category rule list + this long -- measured over 1344 locally generated keywords, 31% of one + category named an activity where the rules asked for a computation, and 24% of + another needed hardware the sandbox does not have. + + Args: + model: the :class:`ApiModel` to ask. + params: sampling params for these calls. A per-call ``sampling_params`` + overrides them, so a caller that already sizes its own calls keeps + doing so. + fallback: local explorer for whichever prompts the API could not answer. + Without one, a failed call comes back as a trajectory with no + assistant message, which every parser here reads as an empty reply. + With one, an unreachable API cannot leave a keyword category dry -- + and dry means keyword-less prompts and a run that looks healthy while + producing one prompt over and over, the exact failure the bank's + refill logic exists to prevent. + + Every returned trajectory carries ``via``: ``'api'`` or ``'local-fallback'``. + It is the one thing a reader of the keyword dump cannot reconstruct afterwards, + and the two halves answer at measurably different quality. + """ + + def __init__(self, model: ApiModel, *, params: Optional[SamplingParams] = None, + fallback: Optional[Explorer] = None): + self.model = model + self.params = params + self.fallback = fallback + + def __call__(self, prompts: Sequence[Trajectory], + sampling_params: Optional[SamplingParams] = None) -> List[Trajectory]: + """One API call per prompt, in order, then the failures in one local batch. + + Serially, where a local explorer would take the whole batch at once: + nothing here knows the API's rate limit, and firing a 32-call expansion at + it is how that gets discovered. The failures are gathered and handed to the + fallback together, because a local sampler shards a batch over its workers + and one prompt at a time would leave most of them idle. + """ + params = sampling_params or self.params + out: List[Optional[Trajectory]] = [] + failed: List[int] = [] + for prompt in prompts: + messages = [dict(m) for m in prompt.get('messages') or []] + content = self.model.generate(messages, params) + if content is None: + failed.append(len(out)) + out.append(None) + continue + messages.append({'role': 'assistant', 'content': content}) + out.append({'messages': messages, 'via': 'api'}) + if failed and self.fallback is not None: + local = self.fallback([prompts[i] for i in failed]) + for i, trajectory in zip(failed, local): + answered = dict(trajectory) + answered['via'] = 'local-fallback' + out[i] = answered + return [t if t is not None else {'messages': [], 'via': None} for t in out] diff --git a/src/twinkle_agentic/challenger/base.py b/src/twinkle_agentic/challenger/base.py index 843d93848..3dcaae674 100644 --- a/src/twinkle_agentic/challenger/base.py +++ b/src/twinkle_agentic/challenger/base.py @@ -31,14 +31,17 @@ import math import random from abc import ABC, abstractmethod -from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass +from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple -from twinkle.data_format import SamplingParams, Trajectory, pack_user_data +from twinkle.data_format import SamplingParams, Trajectory, attach_user_data from twinkle.utils import get_logger +from twinkle_agentic.envs import Env logger = get_logger() -__all__ = ['Challenger', 'Explorer', 'assistant_text', 'attach_user_data'] +__all__ = ['Challenger', 'Explorer', 'KeywordPrompts', 'PromptSet'] # A batch of trajectories in, the same trajectories with the model's reply # appended out. Both MultiTurnRollout and APIMultiTurnRollout satisfy this @@ -48,34 +51,89 @@ Explorer = Callable[[List[Trajectory]], List[Trajectory]] -def attach_user_data(trajectory: Trajectory, **values: Any) -> Trajectory: - """Return ``trajectory`` with ``values`` merged into its packed ``user_data``. +@dataclass +class KeywordPrompts: + """The three strings a :class:`~.keywords.KeywordBank` sends, and nothing else. - ``user_data`` is a list of ``(key, json_string)`` pairs rather than a dict, - so it cannot be updated in place with ``update()``; going through - :func:`pack_user_data` keeps it in the one shape readers understand. + Its own type rather than the caller's prompt object: the challengers here + carry a dozen other prompts, the RSI drivers in ``cookbook/rsi`` keep theirs + as module constants, and a bank that reached into either by attribute name + would be coupled to both spellings. Building one of these is how a caller + says which of its strings are the keyword ones -- see + :meth:`PromptSet.keyword_prompts` for the challengers' answer. + + Lives here rather than beside the bank so that :class:`PromptSet` can produce + one without importing it. + + Args: + system: the system message every keyword call carries. + user: asks for ``{k}`` topics in a category described by ``{desc}``. + expand_user: asks for ``{m}`` more topics like ``{kw}``, optionally with + the category's ``{desc}``. Only :meth:`.KeywordBank.expand_hard` + needs it. """ - merged: Dict[str, Any] = {} - for entry in trajectory.get('user_data') or []: - if isinstance(entry, (list, tuple)) and len(entry) == 2: - merged[entry[0]] = entry[1] - merged.update(values) - out = dict(trajectory) - out['user_data'] = pack_user_data(merged) - return out + system: str + user: str + expand_user: str = '' + + def __post_init__(self): + missing = [f for f in ('system', 'user') if not getattr(self, f).strip()] + if missing: + raise ValueError(f'KeywordPrompts needs {" and ".join(missing)}: a dry ' + f'category could not be refilled without it.') -def assistant_text(trajectory: Trajectory) -> str: - """The last assistant message's text, or '' if the model produced none. +class PromptSet: + """Base for a challenger's bundle of prompts: what is required, and validation. - Explorers differ in what else they attach -- token ids, logprobs, tool - turns -- but every one of them leaves the reply as an assistant message, - so this is the one field a parser can rely on. + Every challenger here is a dataclass of strings plus the same three questions + -- are the mandatory ones filled in, do the optional ones carry the + placeholders they will be formatted with, and does this configuration have + the ones it needs. Answering them once means a missing placeholder is caught + at construction in every domain, rather than as a ``KeyError`` mid-run in + whichever domain remembered to check. + + Subclasses declare: + + * ``_REQUIRED`` -- fields that must carry text. + * ``_REQUIRED_FIELDS`` -- field -> placeholders its text must contain. """ - for message in reversed(trajectory.get('messages') or []): - if isinstance(message, dict) and message.get('role') == 'assistant': - return message.get('content') or '' - return '' + + _REQUIRED: Tuple[str, ...] = () + _REQUIRED_FIELDS: Dict[str, Sequence[str]] = {} + + def __post_init__(self): + name = type(self).__name__ + for field in self._REQUIRED: + if not getattr(self, field).strip(): + raise ValueError(f'{name}.{field} is required') + for field, placeholders in self._REQUIRED_FIELDS.items(): + text = getattr(self, field) + if not text: + continue + for placeholder in placeholders: + if '{' + placeholder + '}' not in text: + raise ValueError(f'{name}.{field} must contain {{{placeholder}}}') + + def require(self, *names: str) -> None: + """Raise unless every named prompt was supplied. + + For what only a configuration knows: drawing from a keyword bank needs the + keyword prompts, seeds need the seed prompt, and a challenger asks for the + ones its arguments imply. + """ + missing = [n for n in names if not getattr(self, n).strip()] + if missing: + name = type(self).__name__ + separator = f', {name}.' + raise ValueError(f'this configuration needs {name}.' + f'{separator.join(missing)}') + + def keyword_prompts(self) -> KeywordPrompts: + """The keyword subset, for the bank. Validated by :meth:`require` first.""" + self.require('keyword_system', 'keyword_user') + return KeywordPrompts(system=self.keyword_system, user=self.keyword_user, + expand_user=self.keyword_expand_user) class Challenger(ABC): @@ -89,17 +147,28 @@ class Challenger(ABC): system: system prompt handed to the model. It carries the output contract, which is why ``build`` -- the code that reads that output back -- lives in the same subclass. + envs: the environments this challenger works in, one per slot. A slot is + owned whole for as long as a job needs it, because the workspace + lives inside it, so ``len(envs)`` is also how many jobs may run at + once. Both halves take the same parameter and reach it the same way + (:meth:`env`), which is what lets one caller decide where everything + it runs is executed and graded: ``[LocalEnv()]`` keeps judgement on + the training host and costs milliseconds, sandbox slots trade that + for isolation. Empty is allowed for a challenger that executes + nothing; :meth:`env` then says so rather than raising IndexError. max_proposals_per_round: ceiling on how many proposals one round may request. Without it a low keep rate makes the estimator ask for an unbounded batch after the first round. solver_rollouts: attempts per candidate in the difficulty stage. ``0`` skips the stage entirely; any other value requires the subclass to implement :meth:`solver_prompt` and :meth:`judge_attempt`. - keep_min_pass: keep a candidate only if at least this many attempts - succeeded. The default drops tasks nobody solved. - keep_max_pass_margin: keep a candidate only if at most - ``solver_rollouts - keep_max_pass_margin`` attempts succeeded. The - default drops tasks everybody solved. + keep_pass_band: ``(low, high)`` attempt counts, inclusive on both ends: + keep a candidate only if that many of its ``solver_rollouts`` + attempts succeeded. Required whenever the stage runs, and has no + default because the counts are absolute -- ``(1, 7)`` reads as "hard + but solvable" against eight rollouts and as something far stricter + against sixteen, so it has to be written by whoever chose the + rollout count. solver_params: sampling params for the difficulty stage only, passed to the explorer per call. ``None`` reuses whatever the explorer was built with -- which is usually the proposing temperature, and that @@ -117,10 +186,10 @@ def __init__( explorer: Explorer, *, system: str, + envs: Sequence[Env] = (), max_proposals_per_round: int = 512, solver_rollouts: int = 0, - keep_min_pass: int = 1, - keep_max_pass_margin: int = 1, + keep_pass_band: Optional[Tuple[int, int]] = None, solver_params: Optional[SamplingParams] = None, solver_explorer: Optional[Explorer] = None, seed: Optional[int] = None, @@ -143,17 +212,29 @@ def __init__( f'solver_rollouts={solver_rollouts} needs {type(self).__name__} to ' f'implement {", ".join(missing)}; pass solver_rollouts=0 to skip the ' f'difficulty stage.') - if keep_min_pass > solver_rollouts - keep_max_pass_margin: + if keep_pass_band is None: + raise ValueError(f'solver_rollouts={solver_rollouts} needs ' + f'keep_pass_band=(low, high): the band is in attempt ' + f'counts, so what it asks for depends on how many ' + f'attempts were run.') + if len(keep_pass_band) != 2: + raise ValueError(f'keep_pass_band is (low, high) in attempt counts, got ' + f'{keep_pass_band}') + low, high = keep_pass_band + if not 0 <= low <= high <= solver_rollouts: raise ValueError( - f'difficulty band is empty: keep_min_pass={keep_min_pass} > ' - f'solver_rollouts - keep_max_pass_margin = ' - f'{solver_rollouts - keep_max_pass_margin}') + f'keep_pass_band must satisfy 0 <= low <= high <= solver_rollouts, ' + f'got {keep_pass_band} against solver_rollouts={solver_rollouts}') + elif keep_pass_band is not None: + raise ValueError('keep_pass_band has nothing to filter while ' + 'solver_rollouts=0 leaves the difficulty stage off; pass ' + 'the rollout count too, or drop the band.') self.explorer = explorer self.system = system + self.envs = list(envs) self.max_proposals_per_round = max_proposals_per_round self.solver_rollouts = solver_rollouts - self.keep_min_pass = keep_min_pass - self.keep_max_pass_margin = keep_max_pass_margin + self.keep_pass_band = keep_pass_band self.solver_params = solver_params self.solver_explorer = solver_explorer self.rng = random.Random(seed) @@ -163,6 +244,28 @@ def __init__( self.n_proposed = 0 self.n_kept = 0 + # ----------------------------------------------------------------- envs + + @property + def n_slots(self) -> int: + """How many jobs may run at once: one per environment.""" + return len(self.envs) + + def env(self, slot: int = 0) -> Env: + """The environment for ``slot``. + + Fetched per use rather than held in a local, so a slot that had to be + rebuilt underneath is picked up on the next call instead of being used + dead. ``slot=0`` is the default because a challenger with nothing to run + concurrently -- one script, no state to share -- has only one. + """ + if not self.envs: + raise RuntimeError( + f'{type(self).__name__} was given no envs, so there is nowhere to run ' + f'anything: pass envs=[LocalEnv()] to execute on the training host, or ' + f'sandbox slots to execute in one.') + return self.envs[slot] + # ------------------------------------------------------------- subclass @abstractmethod @@ -338,8 +441,8 @@ def _filter_difficulty(self, tasks: List[Trajectory]) -> List[Trajectory]: for i, task in enumerate(tasks) ] self.on_difficulty_measured(measured) - high = self.solver_rollouts - self.keep_max_pass_margin - return [t for t, n in zip(measured, passes) if self.keep_min_pass <= n <= high] + low, high = self.keep_pass_band + return [t for t, n in zip(measured, passes) if low <= n <= high] def _estimate(self, missing: int) -> int: """How many proposals to make for ``missing`` keepers. @@ -373,6 +476,25 @@ def draw(rng: random.Random, pool: Sequence[Any], count: int) -> List[Any]: return [rng.choice(pool) for _ in range(count)] if pool else [] +def map_parallel(fn: Callable[[Any], Any], items: Sequence[Any]) -> List[Any]: + """Map ``fn`` over ``items`` at once, results in input order. + + Every use of this is waiting on a sandbox or on a model call, not computing, + so the thread pool is the point. One item runs inline: a pool for a single + call only adds a thread, and it keeps a serial configuration on exactly the + code path it had before. + """ + items = list(items) + if len(items) <= 1: + return [fn(item) for item in items] + out: List[Any] = [None] * len(items) + with ThreadPoolExecutor(max_workers=len(items)) as pool: + futures = {pool.submit(fn, item): i for i, item in enumerate(items)} + for fut in as_completed(futures): + out[futures[fut]] = fut.result() + return out + + def sampling_params_of(explorer: Any) -> Optional[SamplingParams]: """The sampling params an explorer was built with, when it exposes them. diff --git a/src/twinkle_agentic/challenger/code.py b/src/twinkle_agentic/challenger/code.py index 5c0acda0d..be7b8187e 100644 --- a/src/twinkle_agentic/challenger/code.py +++ b/src/twinkle_agentic/challenger/code.py @@ -13,99 +13,71 @@ Prompt text is not here. Every string the model sees arrives in :class:`CodePrompts`, built by whoever runs the challenger -- see -``cookbook/rsi/code/prompts.py``. What stays here is the machinery that cannot -be restated in a prompt: the sandbox, the assert capture, the constant-answer -check, the keyword bank, and how a proposal becomes a task. +``cookbook/rsi/code/challenge_prompts.py``. Neither is execution: every script +runs in an :class:`~twinkle_agentic.envs.base.Env`, which is the same interface +the agentic half verifies through, so where a task gets graded is a decision made +once by the caller rather than twice by the two halves. What stays here is the +machinery that cannot be restated in a prompt: the assert capture, the +constant-answer check, and how a proposal becomes a task. Neither is the keyword +bank -- drawing, refilling and expanding topics is the same cycle on both halves, +so it lives once in :mod:`.keywords` and this challenger holds one. """ import json import os -import random import re -import resource -import shutil -import signal -import subprocess -import sys -import tempfile from dataclasses import dataclass from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple from twinkle.data_format import SamplingParams, Trajectory, user_data_get from twinkle.utils import get_logger -from .base import Challenger, Explorer, assistant_text, attach_user_data +from twinkle_agentic.envs import Env +from twinkle_agentic.utils.code_utils import strip_reasoning, unwrap_code +from twinkle_agentic.utils.message_utils import assistant_text +from .base import Challenger, Explorer, PromptSet, attach_user_data +from .keywords import KeywordBank, KeywordStore logger = get_logger() __all__ = [ - 'CodeChallenger', 'CodePrompts', 'KeywordStore', 'build_asserts', - 'extract_code', 'is_constant_answer', 'load_seeds', 'parse_challenge', - 'run_asserts', + 'CodeChallenger', 'CodePrompts', 'build_asserts', 'is_constant_answer', + 'load_seeds', 'parse_challenge', 'run_asserts', 'run_check_script', ] # Isolates a captured value from anything else the script prints. _MARK = '__RSI_GT__' -_FENCE_RE = re.compile(r'```(?:python|py)?\s*\n(.*?)```', re.S) _JSON_FENCE_RE = re.compile(r'^\s*```(?:json)?\s*|\s*```\s*$', re.I) -# ── sandbox ──────────────────────────────────────────────────────────────── -def extract_code(text: str) -> str: - """The last fenced code block after the thinking section, else the raw body.""" - idx = (text or '').rfind('</think>') - body = text[idx + len('</think>'):] if idx >= 0 else (text or '') - blocks = _FENCE_RE.findall(body) - return (blocks[-1] if blocks else body).strip() +def run_check_script(code: str, check_script: str, env: Env, + timeout: int = 30) -> Tuple[bool, str]: + """Run ``code`` against ``check_script`` in ``env``; True when it exits 0. - -def _run_script(script: str, timeout: int) -> Tuple[int, str]: - """Run a python script in an isolated dir, 2GB cap, killpg on timeout. - - Returns (returncode, stdout). returncode is -1 on timeout/spawn failure. + One script, one exit status -- the same judgement the agentic half makes, so + a task from either half is graded the same way. The run's output comes back + too: a verdict that only says "wrong" leaves a second attempt nothing to go + on. """ - tmp = tempfile.mkdtemp(prefix='rsi_ch_') - try: - with open(os.path.join(tmp, '_run.py'), 'w', encoding='utf-8') as f: - f.write(script + '\n') - env = dict(os.environ, MPLBACKEND='Agg', PYTHONHASHSEED='0', OMP_NUM_THREADS='1', - MKL_NUM_THREADS='1', TOKENIZERS_PARALLELISM='false') - env.pop('CUDA_VISIBLE_DEVICES', None) - - def _limit(): - resource.setrlimit(resource.RLIMIT_AS, (2 * 1024**3, 2 * 1024**3)) - - proc = subprocess.Popen([sys.executable, '_run.py'], cwd=tmp, env=env, - stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, - text=True, start_new_session=True, preexec_fn=_limit) - try: - out, _ = proc.communicate(timeout=timeout) - return proc.returncode, out or '' - except subprocess.TimeoutExpired: - try: - os.killpg(proc.pid, signal.SIGKILL) - except ProcessLookupError: - pass - try: - proc.communicate(timeout=5) - except Exception: - pass - return -1, '' - finally: - shutil.rmtree(tmp, ignore_errors=True) + if not code.strip(): + return False, 'no code was produced' + if not check_script.strip(): + return False, 'no check script was produced' + rc, out = env.run_script(f'{code}\n\n{check_script}', timeout=timeout) + return rc == 0, out -def run_asserts(code: str, setup: str, asserts: List[str], timeout: int = 30) -> bool: - """True when every assert passes (exit status 0).""" - if not code.strip() or not asserts: - return False - parts = [code] - if (setup or '').strip(): - parts.append(setup) - parts.extend(asserts) - rc, _ = _run_script('\n\n'.join(parts), timeout) - return rc == 0 +def run_asserts(code: str, setup: str, asserts: List[str], env: Env, + timeout: int = 30) -> bool: + """True when every assert passes (exit status 0). + + For callers holding a list of asserts rather than one check script -- a + tests file, say. The list plus the setup *is* the check script. + """ + parts = [setup] if (setup or '').strip() else [] + parts.extend(asserts or ()) + return run_check_script(code, '\n\n'.join(parts), env, timeout)[0] -def build_asserts(solution: str, checks: List[str], timeout: int = 30, +def build_asserts(solution: str, checks: List[str], env: Env, timeout: int = 30, max_checks: int = 6) -> Optional[List[str]]: """Run the reference solution once to capture each check's repr, then form ``assert <check> == <captured>``. @@ -123,7 +95,7 @@ def build_asserts(solution: str, checks: List[str], timeout: int = 30, # whole script exit non-zero -> we drop the problem. Pure f-string (no %% # formatting) so a check expression containing '%' (modulo/percent) is safe. lines.append(f'print("{_MARK}{i}=" + repr({c}))') - rc, out = _run_script('\n'.join(lines), timeout) + rc, out = env.run_script('\n'.join(lines), timeout=timeout) if rc != 0: return None captured: Dict[int, str] = {} @@ -131,9 +103,16 @@ def build_asserts(solution: str, checks: List[str], timeout: int = 30, if line.startswith(_MARK): try: idx_str, val = line[len(_MARK):].split('=', 1) - captured[int(idx_str)] = val + idx = int(idx_str) except (ValueError, IndexError): continue + if idx in captured: + # Two lines claiming the same check. The script prints each one + # exactly once, so a second one came from the solution itself -- + # stderr is part of the output now, and a solution that can + # redefine what its own check returned is not ground truth. + return None + captured[idx] = val if len(captured) != len(checks): return None # The captured text is a repr, so it is a valid literal to compare against. @@ -202,10 +181,7 @@ def parse_challenge(text: str, require_solution: bool = True) -> Optional[Dict[s ``require_solution=False`` is for the two-step flow, whose second call is told the solution is already known and returns only the statement. """ - body = text - idx = body.rfind('</think>') - if idx >= 0: - body = body[idx + len('</think>'):] + body = strip_reasoning(text) body = _JSON_FENCE_RE.sub('', body.strip()).strip() # Grab the outermost {...} if there is leading/trailing prose. start, end = body.find('{'), body.rfind('}') @@ -229,61 +205,11 @@ def parse_challenge(text: str, require_solution: bool = True) -> Optional[Dict[s # caller overwrites with the code that actually ran. solution = solution if isinstance(solution, str) else '' if solution and '```' in solution: - solution = extract_code(solution) + solution = unwrap_code(solution) return {'problem': problem.strip(), 'solution': (solution or '').strip(), 'entry': str(obj.get('entry') or '').strip(), 'checks': checks} -# A keyword is a topic to build a task around, not a task statement. Past this many -# characters the model has written the second thing, and storing it makes the next -# prompt ask for a variation on a sentence rather than on a subject. -KEYWORD_MAX_LEN = 60 - - -def split_keyword_list(text: str) -> Tuple[List[str], List[str]]: - """Extract a JSON array of short strings; return (kept, dropped for length). - - The dropped half exists because it used to be discarded inside a list - comprehension. A refill that returned eight well-formed keywords, all of them - written out as sentences, reached the caller as an empty list and was recorded - as ``n_parsed: 0`` -- the same three characters a garbled reply, a timeout and - an over-length reply all produce, so the log could not tell them apart. One - iteration lost 27% of its keywords that way and the cause was found by - re-parsing the stored replies by hand. - - The bias is the reason to count rather than only to log: length correlates with - specificity, so the filter removes "Compute the critical path delay through a - gate-level netlist with annotated cell delays" and keeps whatever was vague - enough to be short. That is the opposite of what the bank is for. - """ - body = text - idx = body.rfind('</think>') - if idx >= 0: - body = body[idx + len('</think>'):] - start, end = body.find('['), body.rfind(']') - if start < 0 or end <= start: - return [], [] - try: - arr = json.loads(body[start:end + 1]) - except (ValueError, TypeError): - return [], [] - kept: List[str] = [] - dropped: List[str] = [] - for x in arr: - if not isinstance(x, str): - continue - s = x.strip() - if not s: - continue - (kept if len(s) <= KEYWORD_MAX_LEN else dropped).append(s) - return kept, dropped - - -def parse_keyword_list(text: str) -> List[str]: - """The kept half of :func:`split_keyword_list`, for callers with nothing to record.""" - return split_keyword_list(text)[0] - - def load_seeds(path: str) -> List[Dict[str, str]]: """Read seed problems from a jsonl: dicts with ``query`` and maybe ``code``. @@ -313,97 +239,9 @@ def load_seeds(path: str) -> List[Dict[str, str]]: return seeds -# ── keyword bank ─────────────────────────────────────────────────────────── -class KeywordStore: - """Persistent keyword bank with usage tracking, one bucket per category. - - Keywords exist to stop the challenger collapsing onto a handful of - archetypes. They are consumed rather than sampled with replacement, so a - run keeps reaching for topics it has not used; when a bucket runs dry the - caller refills it from the model, and recycles only if the model has run out - of distinct ideas. - - On-disk format (one JSON per line):: - - {"category", "text", "used": bool, "used_count": int, - "source": "gen"|"expand", "parent": <keyword or null>} - - De-duplicates case-insensitively within a category, so re-runs never - conflict with the bank on disk. - """ - - def __init__(self, path: str, categories: Sequence[str]): - if not categories: - raise ValueError('KeywordStore needs at least one category') - self.path = path - self.categories = tuple(categories) - self.items: Dict[str, List[Dict[str, Any]]] = {c: [] for c in self.categories} - self._seen: Dict[str, set] = {c: set() for c in self.categories} - if path and os.path.exists(path): - with open(path, encoding='utf-8') as f: - for line in f: - line = line.strip() - if not line: - continue - try: - r = json.loads(line) - except (ValueError, TypeError): - continue - c, t = r.get('category'), r.get('text') - if c in self.items and isinstance(t, str) and t.strip(): - key = t.strip().lower() - if key not in self._seen[c]: - self._seen[c].add(key) - self.items[c].append(r) - - def save(self) -> None: - os.makedirs(os.path.dirname(os.path.abspath(self.path)) or '.', exist_ok=True) - tmp = self.path + '.tmp' - with open(tmp, 'w', encoding='utf-8') as f: - for c in self.categories: - for r in self.items[c]: - f.write(json.dumps(r, ensure_ascii=False) + '\n') - os.replace(tmp, self.path) - - def add(self, category: str, texts: List[str], source: str = 'gen', - parent: Optional[str] = None) -> int: - added = 0 - for t in texts: - key = t.strip().lower() - if not key or key in self._seen[category]: - continue - self._seen[category].add(key) - self.items[category].append({'category': category, 'text': t.strip(), - 'used': False, 'used_count': 0, - 'source': source, 'parent': parent}) - added += 1 - return added - - def unused(self, category: str) -> List[Dict[str, Any]]: - return [r for r in self.items[category] if not r.get('used')] - - def texts(self, category: str) -> List[str]: - return [r['text'] for r in self.items[category]] - - def take(self, category: str, rng: random.Random) -> Optional[str]: - """Consume one unused keyword from ``category``; None if it is dry.""" - un = self.unused(category) - if not un: - return None - r = rng.choice(un) - r['used'] = True - r['used_count'] = r.get('used_count', 0) + 1 - return r['text'] - - def recycle(self, category: str) -> None: - """Mark every keyword unused again (safety valve when the model is tapped out).""" - for r in self.items[category]: - r['used'] = False - - # ── prompts (text supplied by the caller) ────────────────────────────────── @dataclass -class CodePrompts: +class CodePrompts(PromptSet): """Every string a :class:`CodeChallenger` sends, and nothing else. Deliberately without defaults for the always-needed fields: a prompt is the @@ -412,7 +250,8 @@ class CodePrompts: needs them is switched on, and the constructor says so if one is missing. Placeholders are checked at construction: a typo'd ``{keywords}`` would - otherwise surface as a KeyError halfway through a generation run. + otherwise surface as a KeyError halfway through a generation run. That + checking, and the keyword subset a bank is given, are :class:`.PromptSet`. """ system: str @@ -429,6 +268,8 @@ class CodePrompts: keyword_user: str = '' keyword_expand_user: str = '' + #: fields that must carry text. + _REQUIRED = ('system', 'from_scratch', 'solver_system', 'solver_user') #: field -> placeholders it must contain. _REQUIRED_FIELDS = { 'solver_user': ('problem', ), @@ -441,26 +282,6 @@ class CodePrompts: 'keyword_expand_user': ('kw', 'm'), } - def __post_init__(self): - for name in ('system', 'from_scratch', 'solver_system', 'solver_user'): - if not getattr(self, name).strip(): - raise ValueError(f'CodePrompts.{name} is required') - for name, placeholders in self._REQUIRED_FIELDS.items(): - text = getattr(self, name) - if not text: - continue - for placeholder in placeholders: - if '{' + placeholder + '}' not in text: - raise ValueError(f'CodePrompts.{name} must contain ' - f'{{{placeholder}}}') - - def require(self, *names: str) -> None: - """Raise unless every named prompt was supplied.""" - missing = [n for n in names if not getattr(self, n).strip()] - if missing: - raise ValueError(f'this configuration needs CodePrompts.' - f'{", CodePrompts.".join(missing)}') - class CodeChallenger(Challenger): """Propose code problems, execute them for ground truth, keep the graded ones. @@ -476,33 +297,36 @@ class CodeChallenger(Challenger): explorer: batch-in / batch-out generation, see :class:`.base.Explorer`. seeds: optional pool from :func:`load_seeds`, drawn with replacement. keyword_store: optional bank; without it proposals carry no topics. - category_desc: category -> description used when asking for more - keywords. Keys must cover the store's categories. + category_desc / combo_arity / arity_weights / single_kw_prob / + keyword_refill_target / keyword_gen_calls / keyword_refill_tries / + keyword_params / min_batch / expand_per_kw / expand_max_kws: handed to the + :class:`.keywords.KeywordBank` this challenger holds, which is where + they are documented -- they behave the same on the agentic half. seed_mix_prob: chance a proposal also carries a seed problem, when a pool was given. two_step: allow the two-call path (write a harder solution on top of the seed's reference code, then describe the problem it answers). Needs a seed carrying ``code`` and at least one keyword, so it is skipped silently for proposals that have neither. - combo_arity: ``'triple'`` takes one keyword per category; ``'mix'`` - takes a random 1..len(categories) subset. - arity_weights: sampling weights for the ``'mix'`` subset size. - single_kw_prob: in ``'triple'`` mode, the chance of using one category - instead of all of them. - keyword_refill_target / keyword_gen_calls / keyword_refill_tries / - keyword_params: how a dry category is refilled from the model. - min_batch: smallest batch worth sending -- a sampler shards a batch over - its data-parallel workers, and a batch smaller than that leaves some - with nothing to do. Set it to the number of sampler workers. problem_max_chars: reject statements longer than this. Rambling non-problems, and they would also crowd out the solver's context. max_checks / sandbox_timeout: passed to :func:`build_asserts`. drop_constant_answer: reject problems where one constant satisfies every assert. - low_pass_expand / expand_per_kw / expand_max_kws: feedback for + low_pass_expand: a candidate solved this many times or fewer counts as + hard, and its topics are fed back through :meth:`expand_hard_keywords`. reject_sink: called with a dict for every rejected proposal. The caller decides whether that goes to a file; nothing here writes one. + solver_sink: called once per solver attempt in the difficulty stage, with + the check script, the attempt and the verdict. Two things need it: + ``n_pass=0`` reads the same whether the problem is impossible or the + statement withholds a value its asserts demand, and the attempts are + the trainable half of this challenger's output -- see + :meth:`judge_attempt`. Requires a local sampler; an API explorer + returns text without token fields. + keyword_sink: called once per keyword-generation call, with the prompt, + the reply and both halves of the parse. """ def __init__( @@ -531,66 +355,61 @@ def __init__( expand_per_kw: int = 8, expand_max_kws: int = 32, reject_sink: Optional[Callable[[Dict[str, Any]], None]] = None, + solver_sink: Optional[Callable[[Dict[str, Any]], None]] = None, + keyword_sink: Optional[Callable[[Dict[str, Any]], None]] = None, **challenger_kwargs: Any, ): super().__init__(explorer, system=prompts.system, **challenger_kwargs) - if combo_arity not in ('triple', 'mix'): - raise ValueError(f"combo_arity must be 'triple' or 'mix', got {combo_arity!r}") + if not self.envs: + raise ValueError('envs is empty: there is nowhere to run a check, and every ' + 'proposal would be rejected for a ground truth that never ran.') if keyword_store is not None: - desc = category_desc or {} - missing = [c for c in keyword_store.categories if not desc.get(c)] - if missing: - raise ValueError(f'category_desc is missing a description for ' - f'{missing}; a dry category could not be refilled.') - prompts.require('keyword_system', 'keyword_user', 'from_keywords') + prompts.require('from_keywords') self.prompts = prompts self.seeds = list(seeds) - self.store = keyword_store - self.category_desc = dict(category_desc or {}) + # The whole keyword cycle -- draw, refill, expand -- is one object shared + # with the agentic challenger rather than a second copy of it here. None + # means no bank was configured, and proposals then carry no topics. + self.keywords: Optional[KeywordBank] = None if keyword_store is None else KeywordBank( + keyword_store, prompts=prompts.keyword_prompts(), + category_desc=category_desc or {}, + explorer=explorer, rng=self.rng, name=type(self).__name__, + sampling_params=keyword_params, sink=keyword_sink, combo_arity=combo_arity, + arity_weights=arity_weights, single_kw_prob=single_kw_prob, + refill_target=keyword_refill_target, gen_calls=keyword_gen_calls, + refill_tries=keyword_refill_tries, min_batch=min_batch, + expand_per_kw=expand_per_kw, expand_max_kws=expand_max_kws) self.seed_mix_prob = seed_mix_prob self.two_step = two_step - self.combo_arity = combo_arity - self.arity_weights = list(arity_weights) if arity_weights else None - self.single_kw_prob = single_kw_prob - self.keyword_refill_target = keyword_refill_target - self.keyword_gen_calls = keyword_gen_calls - self.keyword_refill_tries = keyword_refill_tries - self.keyword_params = keyword_params - self.min_batch = max(1, min_batch) self.problem_max_chars = problem_max_chars self.max_checks = max_checks self.sandbox_timeout = sandbox_timeout self.drop_constant_answer = drop_constant_answer self.low_pass_expand = low_pass_expand - self.expand_per_kw = expand_per_kw - self.expand_max_kws = expand_max_kws self.reject_sink = reject_sink + self.solver_sink = solver_sink if self.seeds: # Both are reachable with a bank configured: a proposal draws no # keywords when every category is dry, and then falls back to the # seed-only prompt. prompts.require('from_seed') - if self.store is not None: + if self.keywords is not None: prompts.require('from_seed_keywords') if two_step: prompts.require('two_step_system', 'two_step_solution', 'two_step_problem') - # Perturbs refill prompts so a second ask does not repeat the first. - self._nonce = 0 # Why proposals died, for the caller to log; the shape a run is judged on. self.stats: Dict[str, int] = { 'parsed': 0, 'parse_fail': 0, 'stage1_no_code': 0, 'too_long': 0, 'gt_fail': 0, 'selfcheck_fail': 0, 'constant_answer': 0, } - # (category, keyword) behind candidates nobody could solve, for feedback. - self._hard: List[Tuple[str, str]] = [] # ------------------------------------------------------------- proposing def propose(self, count: int) -> List[Trajectory]: proposals: List[Trajectory] = [] for _ in range(count): - picks = self._draw_keywords() - body = '\n'.join(f'- {c}: {t}' for c, t in picks) + picks = self.keywords.draw() if self.keywords else [] + body = KeywordBank.block(picks) use_seed = bool(self.seeds) and self.rng.random() < self.seed_mix_prob seed = self.rng.choice(self.seeds) if use_seed else None two = bool(use_seed and self.two_step and picks and seed and seed.get('code')) @@ -621,84 +440,6 @@ def propose(self, count: int) -> List[Trajectory]: seed_query=(seed['query'] if two else ''), keyword_block=body)) return proposals - def _draw_keywords(self) -> List[Tuple[str, str]]: - """Consume one keyword combination from the bank; [] without a bank.""" - if self.store is None: - return [] - categories = self.store.categories - if self.combo_arity == 'mix': - if self.arity_weights and len(self.arity_weights) == len(categories): - k = self.rng.choices(range(1, len(categories) + 1), - weights=self.arity_weights)[0] - else: - k = self.rng.randint(1, len(categories)) - cats = self.rng.sample(list(categories), k) - elif self.rng.random() < self.single_kw_prob: - cats = [self.rng.choice(categories)] - else: - cats = list(categories) - picks: List[Tuple[str, str]] = [] - for c in cats: - if not self.store.unused(c): - self._refill(c) - text = self.store.take(c, self.rng) - if text is not None: - picks.append((c, text)) - return picks - - def _refill(self, category: str) -> None: - """Ask the model for more keywords in ``category``; recycle if it is tapped out.""" - tries = 0 - while not self.store.unused(category): - new = self._generate_keywords(category, self.keyword_refill_target) - added = self.store.add(category, new, source='gen') - tries += 1 - if added == 0 and tries >= self.keyword_refill_tries: - if self.store.items[category]: - self.store.recycle(category) - logger.info(f'[CodeChallenger] keyword category {category!r} exhausted ' - f'-> recycled {len(self.store.items[category])} topics') - break - - def _generate_keywords(self, category: str, n_want: int) -> List[str]: - """Up to ``n_want`` keywords the bank does not already hold.""" - if n_want <= 0: - return [] - known = self.store.texts(category) - n_calls = max(self.keyword_gen_calls, self.min_batch) - per_call = max(1, -(-n_want // n_calls) + 4) # ceil(n/calls) + margin - avoid_note = '' - if known: - shown = known if len(known) <= 40 else self.rng.sample(known, 40) - avoid_note = ('\nDo NOT repeat any of these already-used topics: ' - + ', '.join(shown)) - base = self.prompts.keyword_user.format( - k=per_call, desc=self.category_desc[category]) + avoid_note - self._nonce += 1 - prompts = [{ - 'messages': [{'role': 'system', 'content': self.prompts.keyword_system}, - {'role': 'user', 'content': f'{base}\n(batch {self._nonce}-{i})'}], - } for i in range(n_calls)] - seen = {t.strip().lower() for t in known} - out: List[str] = [] - n_long = 0 - for reply in self.explore(prompts, sampling_params=self.keyword_params): - kept, dropped = split_keyword_list(assistant_text(reply)) - n_long += len(dropped) - for kw in kept: - key = kw.lower() - if key not in seen: - seen.add(key) - out.append(kw) - if n_long: - # This path has no dump to write to, so the count has to be said out - # loud or the refill looks like the model simply produced less. - logger.warning(f'[CodeChallenger] dropped {n_long} keyword(s) over ' - f'{KEYWORD_MAX_LEN} chars while refilling; the prompt is ' - f'asking for task statements rather than topics') - self.rng.shuffle(out) - return out[:n_want] - # ---------------------------------------------------------------- building def build(self, explored: List[Trajectory]) -> List[Optional[Trajectory]]: @@ -720,7 +461,7 @@ def build(self, explored: List[Trajectory]) -> List[Optional[Trajectory]]: if not user_data_get(traj.get('user_data'), 'two_step', False): objs[i] = parse_challenge(text) continue - code = extract_code(text) + code = unwrap_code(text) if not code.strip(): # Usually a truncated completion: there is no solution to # describe, so this proposal ends here. @@ -769,12 +510,13 @@ def _reject(reason: str, **extra: Any) -> None: if len(obj['problem']) > self.problem_max_chars: _reject('too_long') return None - asserts = build_asserts(obj['solution'], obj['checks'], + asserts = build_asserts(obj['solution'], obj['checks'], self.env(), timeout=self.sandbox_timeout, max_checks=self.max_checks) if not asserts: _reject('gt_fail') return None - if not run_asserts(obj['solution'], '', asserts, timeout=self.sandbox_timeout): + check_script = '\n'.join(asserts) + if not self._check(obj['solution'], check_script)[0]: # A reference solution that fails its own asserts is not ground # truth, whatever the statement says. _reject('selfcheck_fail', asserts=asserts) @@ -793,7 +535,12 @@ def _reject(reason: str, **extra: Any) -> None: } return attach_user_data( task, - asserts=asserts, + # One script rather than a list of asserts, named as the agentic half + # names it: a consumer that trains on both halves then reads the + # verifier the same way. setup_script is where that half puts the + # part that runs before the checks; nothing here needs one. + check_script=check_script, + setup_script='', solution=obj['solution'], entry=obj['entry'], keywords=user_data_get(user_data, 'keywords', []), @@ -811,25 +558,51 @@ def solver_prompt(self, task: Trajectory) -> Trajectory: 'content': self.prompts.solver_user.format(problem=problem)}], } + def _check(self, code: str, check_script: str) -> Tuple[bool, str]: + """Did ``code`` pass ``check_script``? With the output, for feedback. + + Slot 0 always: a code judgement is one script with no state to share, and + this half runs them one at a time -- see the note in + :meth:`judge_attempt` -- so the slots the agentic half needs for its + concurrent episodes have nothing to do here. + """ + return run_check_script(code, check_script, self.env(), self.sandbox_timeout) + def judge_attempt(self, task: Trajectory, attempt: Trajectory) -> bool: - asserts = user_data_get(task.get('user_data'), 'asserts', []) or [] - return run_asserts(extract_code(assistant_text(attempt)), '', asserts, - timeout=self.sandbox_timeout) + """Did this attempt's code pass the task's asserts? + + Also hands the whole attempt to ``solver_sink`` when one is given. The + stage otherwise reduces each task to one number and drops the attempts, + and they are exactly what a solver trains on: a task kept at 3 of 8 is + one prompt answered eight times with a binary reward, which is a GRPO + group already measured to have a gradient. Sampling them again after the + band has been applied pays for the same tokens twice and can still land + the group at 0 or 8, where the advantage is the reward minus itself. + """ + check_script = user_data_get(task.get('user_data'), 'check_script', '') or '' + passed, output = self._check(unwrap_code(assistant_text(attempt)), check_script) + if self.solver_sink is not None: + # No lock: the difficulty stage judges attempts one at a time, in the + # loop that counts them, unlike the agentic half where each judgement + # is a sandbox round trip worth running concurrently. + self.solver_sink({ + 'statement': next((m.get('content', '') for m in task.get('messages') or [] + if m.get('role') == 'user'), ''), + # What the caller groups on: two problems with byte-identical + # asserts are the same problem, and a kept task carries this + # field through unchanged. + 'check_script': check_script, + 'passed': passed, + 'output': output, + 'truncated': bool((attempt or {}).get('truncated')), + 'attempt': attempt, + }) + return passed def on_difficulty_measured(self, candidates: List[Trajectory]) -> None: """Remember the topics behind the candidates nobody solved.""" - if self.store is None: - return - seen = {(c, t.lower()) for c, t in self._hard} - for task in candidates: - data = task.get('user_data') - if user_data_get(data, 'n_pass', 0) > self.low_pass_expand: - continue - for pick in user_data_get(data, 'keywords', []) or []: - c, t = pick[0], pick[1] - if (c, t.lower()) not in seen: - seen.add((c, t.lower())) - self._hard.append((c, t)) + if self.keywords is not None: + self.keywords.remember_unsolved(candidates, self.low_pass_expand) # ------------------------------------------------------------- feedback @@ -840,35 +613,4 @@ def expand_hard_keywords(self) -> int: drifts toward material the solver actually struggles with. Returns how many new keywords were added. """ - if self.store is None or not self._hard or self.expand_per_kw <= 0: - return 0 - self.prompts.require('keyword_expand_user') - hard = self._hard[:self.expand_max_kws] - self.rng.shuffle(hard) - # Cycle a short list so the batch still covers every sampler worker. - reqs = list(hard) - while len(reqs) < self.min_batch: - reqs.append(hard[len(reqs) % len(hard)]) - self._nonce += 1 - prompts = [{ - 'messages': [ - {'role': 'system', 'content': self.prompts.keyword_system}, - {'role': 'user', - 'content': self.prompts.keyword_expand_user.format(kw=kw, m=self.expand_per_kw) - + f'\n(batch {self._nonce}-{i})'}, - ], - } for i, (_c, kw) in enumerate(reqs)] - added = 0 - n_long = 0 - for (cat, kw), reply in zip(reqs, self.explore(prompts, - sampling_params=self.keyword_params)): - kept, dropped = split_keyword_list(assistant_text(reply)) - n_long += len(dropped) - added += self.store.add(cat, kept, source='expand', parent=kw) - if n_long: - logger.warning(f'[CodeChallenger] dropped {n_long} expanded keyword(s) over ' - f'{KEYWORD_MAX_LEN} chars; expansion follows the parent, so a ' - f'wordy parent produces wordy children') - logger.info(f'[CodeChallenger] expanded {len(hard)} hard keyword(s) -> ' - f'+{added} same-domain topics') - return added + return self.keywords.expand_hard() if self.keywords is not None else 0 diff --git a/src/twinkle_agentic/challenger/keywords.py b/src/twinkle_agentic/challenger/keywords.py new file mode 100644 index 000000000..3f6e7bcef --- /dev/null +++ b/src/twinkle_agentic/challenger/keywords.py @@ -0,0 +1,588 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""The keyword bank: what the next task gets built about. + +Every challenger here proposes from a topic rather than from a fixed prompt, +because a fixed prompt collapses onto a handful of archetypes within a few +hundred proposals. The cycle that prevents it is the same one everywhere -- draw +a combination, refill whichever category ran dry, ask for more of whatever the +solver could not solve -- so it lives here once, in :class:`KeywordBank`, and a +proposer *holds* one rather than inheriting it. Nothing in this module knows what +a task looks like: it deals in short strings and in the prompts its owner supplies, +which is why the two challengers in this package and the RSI drivers in +``cookbook/rsi`` can all share it. + +Two failures shaped the file, both from real runs, both recorded where they hit: +a refill that returns nothing has to be loud (a silent one leaves every proposal +falling back to the from-scratch prompt while the run looks healthy), and a +keyword that arrives written as a sentence has to be counted rather than dropped +in silence -- see :func:`split_keyword_list`. +""" +import json +import os +import random +import threading +from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple + +from twinkle.data_format import SamplingParams, Trajectory, user_data_get +from twinkle.utils import get_logger +from twinkle_agentic.utils.code_utils import strip_reasoning +from twinkle_agentic.utils.message_utils import assistant_text +from .base import Explorer, KeywordPrompts, map_parallel + +logger = get_logger() + +__all__ = [ + 'KEYWORD_MAX_LEN', 'KeywordBank', 'KeywordPrompts', 'KeywordStore', + 'parse_keyword_list', 'split_keyword_list', +] + +# A keyword is a topic to build a task around, not a task statement. Past this many +# characters the model has written the second thing, and storing it makes the next +# prompt ask for a variation on a sentence rather than on a subject. +KEYWORD_MAX_LEN = 60 + + +def split_keyword_list(text: str) -> Tuple[List[str], List[str]]: + """Extract a JSON array of short strings; return (kept, dropped for length). + + The dropped half exists because it used to be discarded inside a list + comprehension. A refill that returned eight well-formed keywords, all of them + written out as sentences, reached the caller as an empty list and was recorded + as ``n_parsed: 0`` -- the same three characters a garbled reply, a timeout and + an over-length reply all produce, so the log could not tell them apart. One + iteration lost 27% of its keywords that way and the cause was found by + re-parsing the stored replies by hand. + + The bias is the reason to count rather than only to log: length correlates with + specificity, so the filter removes "Compute the critical path delay through a + gate-level netlist with annotated cell delays" and keeps whatever was vague + enough to be short. That is the opposite of what the bank is for. + """ + body = strip_reasoning(text) + start, end = body.find('['), body.rfind(']') + if start < 0 or end <= start: + return [], [] + try: + arr = json.loads(body[start:end + 1]) + except (ValueError, TypeError): + return [], [] + kept: List[str] = [] + dropped: List[str] = [] + for x in arr: + if not isinstance(x, str): + continue + s = x.strip() + if not s: + continue + (kept if len(s) <= KEYWORD_MAX_LEN else dropped).append(s) + return kept, dropped + + +def parse_keyword_list(text: str) -> List[str]: + """The kept half of :func:`split_keyword_list`, for callers with nothing to record.""" + return split_keyword_list(text)[0] + + +# ── keyword bank ─────────────────────────────────────────────────────────── +class KeywordStore: + """Persistent keyword bank with usage tracking, one bucket per category. + + Keywords exist to stop the challenger collapsing onto a handful of + archetypes. They are consumed rather than sampled with replacement, so a + run keeps reaching for topics it has not used; when a bucket runs dry the + caller refills it from the model, and recycles only if the model has run out + of distinct ideas. + + On-disk format (one JSON per line):: + + {"category", "text", "used": bool, "used_count": int, + "source": "gen"|"expand", "parent": <keyword or null>} + + De-duplicates case-insensitively within a category, so re-runs never + conflict with the bank on disk. + """ + + def __init__(self, path: str, categories: Sequence[str]): + if not categories: + raise ValueError('KeywordStore needs at least one category') + self.path = path + self.categories = tuple(categories) + self.items: Dict[str, List[Dict[str, Any]]] = {c: [] for c in self.categories} + self._seen: Dict[str, set] = {c: set() for c in self.categories} + if path and os.path.exists(path): + with open(path, encoding='utf-8') as f: + for line in f: + line = line.strip() + if not line: + continue + try: + r = json.loads(line) + except (ValueError, TypeError): + continue + c, t = r.get('category'), r.get('text') + if c in self.items and isinstance(t, str) and t.strip(): + key = t.strip().lower() + if key not in self._seen[c]: + self._seen[c].add(key) + self.items[c].append(r) + + def save(self) -> None: + """Write the bank out, atomically. A bank without a path is in-memory only.""" + if not self.path: + return + os.makedirs(os.path.dirname(os.path.abspath(self.path)) or '.', exist_ok=True) + tmp = self.path + '.tmp' + with open(tmp, 'w', encoding='utf-8') as f: + for c in self.categories: + for r in self.items[c]: + f.write(json.dumps(r, ensure_ascii=False) + '\n') + os.replace(tmp, self.path) + + def add(self, category: str, texts: Sequence[str], source: str = 'gen', + parent: Optional[str] = None) -> int: + added = 0 + for t in texts: + key = t.strip().lower() + if not key or key in self._seen[category]: + continue + self._seen[category].add(key) + self.items[category].append({'category': category, 'text': t.strip(), + 'used': False, 'used_count': 0, + 'source': source, 'parent': parent}) + added += 1 + return added + + def unused(self, category: str) -> List[Dict[str, Any]]: + return [r for r in self.items[category] if not r.get('used')] + + def texts(self, category: str) -> List[str]: + return [r['text'] for r in self.items[category]] + + def take(self, category: str, rng: random.Random) -> Optional[str]: + """Consume one unused keyword from ``category``; None if it is dry.""" + un = self.unused(category) + if not un: + return None + r = rng.choice(un) + r['used'] = True + r['used_count'] = r.get('used_count', 0) + 1 + return r['text'] + + def recycle(self, category: str) -> None: + """Mark every keyword unused again (safety valve when the model is tapped out).""" + for r in self.items[category]: + r['used'] = False + + +class KeywordBank: + """The draw / refill / expand cycle over a :class:`KeywordStore`. + + Held by whoever proposes rather than inherited: the two challengers in this + package share this entire cycle and almost nothing else, and the RSI drivers + in ``cookbook/rsi`` are not challengers at all yet need exactly the same + thing. Safe to drive from several threads -- see :meth:`draw`. + + Args: + store: the bank this works on. + prompts: the :class:`KeywordPrompts` this bank sends. + category_desc: category -> description, shown when asking for more. Must + cover every category in the store, or a dry one could not be refilled. + explorer: batch-in / batch-out generation, used for keyword calls only. + Worth keeping separate from the proposing explorer: brainstorming a + list is a text round, so a tool-calling rollout both wastes turns and + may take a bracketed list in the reply for a tool call. + rng: shared with the owner, so one seed reproduces the whole run. + name: what log lines call this bank; normally the owner's class name. + sampling_params: params for keyword calls. None sends the explorer's own. + sink: called once per keyword call with the prompt, the reply and both + halves of the parse. The one question such a dump exists to answer -- + did the model disobey the format, or does the parser reject what it + produced -- cannot be answered from a count. + combo_arity: ``'triple'`` draws one keyword per category; ``'mix'`` draws + a random 1..len(categories) subset. + arity_weights: sampling weights for the ``'mix'`` subset size. + single_kw_prob: in ``'triple'`` mode, the chance of using one category + instead of all of them. + refill_target: how many new keywords one refill aims for. + gen_calls: how many model calls it may spend on that. + refill_concurrency: how many of those go out together. At 1 every call is + told what the ones before it produced, which is what the avoid list is + for; raising it is faster and comes back with more synonyms. + refill_tries: refills to attempt before recycling a tapped-out category. + min_batch: smallest batch worth sending -- a sampler shards a batch over + its data-parallel workers, and a smaller one leaves some with nothing + to do. Set it to the number of sampler workers. + expand_per_kw / expand_max_kws: size of the :meth:`expand_hard` ask. + """ + + # How many phrases the 'do not repeat these' line may quote in total. There + # has to be a ceiling in both directions: too few and a serial refill stops + # seeing what it just said, too many and the model runs out of room to obey. + # Measured on armA2ser, where this refill's own output went in uncapped: with + # 130 quoted the eighth call was still answering normally, with 150 it started + # inventing -- 'îRAPIÓN holistic replace', 'ซะ subspace cutter map limit', 10 + # of 480 phrases that run. 100 sits below where that began. + _AVOID_TOTAL = 100 + _AVOID_LEAD = '\nDo NOT repeat any of these already-used topics: ' + + def __init__( + self, + store: KeywordStore, + *, + prompts: KeywordPrompts, + category_desc: Dict[str, str], + explorer: Explorer, + rng: random.Random, + name: str = 'keywords', + sampling_params: Optional[SamplingParams] = None, + sink: Optional[Callable[[Dict[str, Any]], None]] = None, + combo_arity: str = 'triple', + arity_weights: Optional[Sequence[float]] = None, + single_kw_prob: float = 0.1, + refill_target: int = 128, + gen_calls: int = 8, + refill_concurrency: int = 1, + refill_tries: int = 2, + min_batch: int = 1, + expand_per_kw: int = 8, + expand_max_kws: int = 32, + ): + if combo_arity not in ('triple', 'mix'): + raise ValueError(f"combo_arity must be 'triple' or 'mix', got {combo_arity!r}") + if refill_concurrency < 1: + raise ValueError(f'refill_concurrency must be >= 1, got {refill_concurrency}') + missing = [c for c in store.categories if not (category_desc or {}).get(c)] + if missing: + raise ValueError(f'category_desc is missing a description for {missing}; ' + f'a dry category could not be refilled.') + self.store = store + self.prompts = prompts + self.category_desc = dict(category_desc) + self.explorer = explorer + self.rng = rng + self.name = name + self.sampling_params = sampling_params + self.sink = sink + self.combo_arity = combo_arity + self.arity_weights = list(arity_weights) if arity_weights else None + self.single_kw_prob = single_kw_prob + self.refill_target = refill_target + self.gen_calls = gen_calls + self.refill_concurrency = refill_concurrency + self.refill_tries = refill_tries + self.min_batch = max(1, min_batch) + self.expand_per_kw = expand_per_kw + self.expand_max_kws = expand_max_kws + # One draw at a time; see :meth:`draw` for why the whole draw and not just + # the store access. + self._draw_lock = threading.Lock() + # Held while the rng, the nonce, the bank or the hard list are touched, and + # never across a model call. Separate from the sink lock, which waits on + # disk: a refill running in another thread must not queue behind a write. + self._state_lock = threading.Lock() + self._sink_lock = threading.Lock() + # Perturbs prompts so two calls are never byte-identical. + self._nonce = 0 + # (category, keyword) behind whatever nobody could solve, for expand_hard. + self._hard: List[Tuple[str, str]] = [] + + # -------------------------------------------------------------- drawing + + @property + def categories(self) -> Tuple[str, ...]: + return self.store.categories + + @staticmethod + def block(picks: Sequence[Tuple[str, str]]) -> str: + """The drawn keywords as the line block a prompt's ``{keywords}`` takes.""" + return '\n'.join(f'- {c}: {t}' for c, t in picks) + + def draw(self) -> List[Tuple[str, str]]: + """Consume one keyword combination, refilling whatever ran dry first. + + Serialised as a whole rather than per bank access: a refill is a batch of + model calls whose prompts quote what the calls before them produced, and + two draws overlapping would each refill without seeing the other's + keywords -- exactly what the avoid list exists to prevent. + """ + with self._draw_lock: + cats = self._pick_categories() + # Refill every dry category at once rather than as each one is reached: + # they are independent model calls that used to run one after another + # (20s each at the start of a run) and they touch separate buckets. + dry = [c for c in cats if not self.store.unused(c)] + if dry: + map_parallel(self.refill, dry) + picks: List[Tuple[str, str]] = [] + for category in cats: + with self._state_lock: + text = self.store.take(category, self.rng) + if text is not None: + picks.append((category, text)) + return picks + + def _pick_categories(self) -> List[str]: + """Which categories one draw covers, per ``combo_arity``.""" + categories = self.store.categories + with self._state_lock: + if self.combo_arity == 'mix': + if self.arity_weights and len(self.arity_weights) == len(categories): + k = self.rng.choices(range(1, len(categories) + 1), + weights=self.arity_weights)[0] + else: + k = self.rng.randint(1, len(categories)) + return self.rng.sample(list(categories), k) + if self.rng.random() < self.single_kw_prob: + return [self.rng.choice(categories)] + return list(categories) + + # ------------------------------------------------------------- refilling + + def refill(self, category: str) -> None: + """Ask the model for more keywords in ``category``; recycle if it is tapped out. + + Says so when it comes back empty. A silent no-op here is the worst outcome + available: :meth:`draw` then hands out no keywords, every proposal quietly + falls back to the from-scratch prompt, and the run looks normal while + producing one identical prompt over and over. That is exactly what happened + for whole runs when the prompt asked for one keyword per line and the parser + wanted a JSON array. + """ + tries = 0 + while not self.store.unused(category): + new = self._generate(category, self.refill_target) + with self._state_lock: + added = self.store.add(category, new, source='gen') + if added: + # Saved now rather than at the end of the run: a refill costs a + # batch of model calls, and a run that crashes later should not + # have to spend them again -- the next iteration reads this file + # to know what was already used. + self.store.save() + tries += 1 + if added: + logger.info(f'[{self.name}] keyword category {category!r} refilled ' + f'+{added} (try {tries})') + continue + logger.warning( + f'[{self.name}] keyword refill for {category!r} produced nothing on try ' + f'{tries}: {len(new)} parsed, 0 new. Proposals will run without keywords ' + f'unless this recovers -- pass a keyword sink to see the replies.') + if tries >= self.refill_tries: + with self._state_lock: + # Every keyword marked unused again. The alternative is a + # category that can never be drawn from, which stops the run: a + # repeat draw is worse than no run only if diversity matters + # more than collecting anything at all. + n_recycled = len(self.store.items[category]) + if n_recycled: + self.store.recycle(category) + self.store.save() + if n_recycled: + logger.info(f'[{self.name}] keyword category {category!r} exhausted ' + f'-> recycled {n_recycled} topics') + break + + def _generate(self, category: str, n_want: int) -> List[str]: + """Up to ``n_want`` keywords the bank does not already hold.""" + if n_want <= 0: + return [] + with self._state_lock: + known = self.store.texts(category) + n_calls = max(self.gen_calls, self.min_batch) + per_call = max(1, -(-n_want // n_calls) + 4) # ceil(n/calls) + margin + seen = {t.strip().lower() for t in known} + out: List[str] = [] + n_long = 0 + for start in range(0, n_calls, self.refill_concurrency): + group = range(start, min(start + self.refill_concurrency, n_calls)) + # Every call in a group is built before any of them runs, so they all + # carry the same avoid list -- which is exactly the batched behaviour, + # and why a group of one is what lets call k+1 see call k. + users = [(self.prompts.user.format( + k=per_call, desc=self.category_desc[category]) + + self._avoid_note(known, out) + + f'\n(batch {self._next_nonce()}-{i})') for i in group] + for user, reply in zip(users, self._explore(users)): + text = assistant_text(reply) + parsed, dropped_long = split_keyword_list(text) + n_long += len(dropped_long) + fresh = [kw for kw in parsed if kw.lower() not in seen] + seen.update(kw.lower() for kw in fresh) + out.extend(fresh) + # Full text, both sides: the question this dump answers is whether + # the model disobeyed the format or the parser rejected what it + # produced, and a count cannot say which. + self._record({ + 'category': category, 'prompt': user, 'reply': text, + 'stop_reason': reply.get('stop_reason'), + 'truncated': bool(reply.get('truncated')), + 'parsed': parsed, 'n_parsed': len(parsed), 'n_new': len(fresh), + # Which backend answered, for an explorer that has more than + # one -- an API with a local fallback is the case this exists + # for, and nothing else here could know which one ran. + 'via': reply.get('via'), + # The two fields that make the sentence above true. Without them + # ``n_parsed: 0`` reads the same whether the reply was garbled, + # empty, or eight usable keywords written at sentence length -- + # and the third is the one that happened. + 'dropped_long': dropped_long, 'n_dropped_long': len(dropped_long), + }) + if len(out) >= n_want: + # The surplus is dropped below, so further calls would buy nothing. + break + if n_long and self.sink is None: + # Without a dump to write to, the count has to be said out loud or the + # refill looks like the model simply produced less. + logger.warning(f'[{self.name}] dropped {n_long} keyword(s) over ' + f'{KEYWORD_MAX_LEN} chars while refilling; the prompt is ' + f'asking for task statements rather than topics') + with self._state_lock: + self.rng.shuffle(out) + return out[:n_want] + + def _avoid_note(self, older: List[str], fresh: List[str]) -> str: + """The 'do not repeat these' line, newest first, capped at ``_AVOID_TOTAL``. + + What this refill has just produced comes first and evicts older entries + rather than the reverse -- the calls run one at a time so that each can + avoid what the ones before it said, and dropping those would undo it. Past + the cap the oldest of *this refill's* phrases are what falls off, which is + also the least costly thing to drop: the model has already moved away from + them. + """ + fresh_shown = list(fresh)[-self._AVOID_TOTAL:] + room = max(0, self._AVOID_TOTAL - len(fresh_shown)) + with self._state_lock: + shown = older if len(older) <= room else self.rng.sample(older, room) + avoid = fresh_shown + list(shown) + return self._AVOID_LEAD + ', '.join(avoid) if avoid else '' + + # -------------------------------------------------------------- feedback + + def remember_hard(self, picks: Iterable[Sequence[str]]) -> None: + """Note the (category, keyword) pairs behind something nobody solved. + + De-duplicated case-insensitively and kept in arrival order. This is the + only feedback the bank gets from difficulty; without it, it drifts wherever + the refill prompt happens to go. + """ + with self._state_lock: + seen = {(c, t.lower()) for c, t in self._hard} + for pick in picks or (): + if not (isinstance(pick, (list, tuple)) and len(pick) >= 2): + continue + category, text = pick[0], pick[1] + if (category, text.lower()) not in seen: + seen.add((category, text.lower())) + self._hard.append((category, text)) + + def remember_unsolved(self, candidates: Sequence[Trajectory], max_pass: int = 0) -> None: + """:meth:`remember_hard` for measured candidates at or below ``max_pass``. + + Each candidate carries ``n_pass`` and its keyword draw in ``user_data``, so + this is the whole of what a difficulty round feeds back to the bank. + """ + for task in candidates: + data = task.get('user_data') + if user_data_get(data, 'n_pass', 0) > max_pass: + continue + self.remember_hard(user_data_get(data, 'keywords', []) or []) + + def expand_hard(self) -> int: + """Brainstorm more topics in the families that produced the hardest tasks. + + Called by whoever drives the proposer, after a round, so the bank drifts + toward material the solver actually struggles with. Returns how many new + keywords were added. + """ + if not self._hard or self.expand_per_kw <= 0: + return 0 + template = self.prompts.expand_user.strip() + if not template: + raise ValueError(f'[{self.name}] measured {len(self._hard)} hard keyword(s) to ' + f'expand on, but the prompts carry no expand_user.') + with self._state_lock: + hard = self._hard[:self.expand_max_kws] + self.rng.shuffle(hard) + # Cycle a short list so the batch still covers every sampler worker. + reqs = list(hard) + while len(reqs) < self.min_batch: + reqs.append(hard[len(reqs) % len(hard)]) + # ``desc`` is offered alongside ``kw``/``m``: a template that does not ask + # for it ignores it, and one that does gets the category's rules with it. + users = [(template.format(kw=kw, m=self.expand_per_kw, + desc=self.category_desc.get(category, '')) + + f'\n(batch {self._next_nonce()}-{i})') + for i, (category, kw) in enumerate(reqs)] + added = 0 + n_long = 0 + for (category, kw), user, reply in zip(reqs, users, self._explore(users)): + text = assistant_text(reply) + parsed, dropped_long = split_keyword_list(text) + n_long += len(dropped_long) + with self._state_lock: + added += self.store.add(category, parsed, source='expand', parent=kw) + # The prompt goes in whole, as the refill path already does. This used + # to record the literal string 'expand', which left the dump unable to + # answer the one question it gets asked -- whether a change to the + # expansion prompt was live in a given iteration. + self._record({ + 'category': category, 'parent': kw, 'prompt': user, 'reply': text, + 'stop_reason': reply.get('stop_reason'), + 'truncated': bool(reply.get('truncated')), + 'parsed': parsed, 'n_parsed': len(parsed), + 'dropped_long': dropped_long, 'n_dropped_long': len(dropped_long), + 'via': reply.get('via'), + }) + if added: + with self._state_lock: + self.store.save() + if n_long and self.sink is None: + logger.warning(f'[{self.name}] dropped {n_long} expanded keyword(s) over ' + f'{KEYWORD_MAX_LEN} chars; expansion follows the parent, so a ' + f'wordy parent produces wordy children') + logger.info(f'[{self.name}] expanded {len(hard)} hard keyword(s) -> ' + f'+{added} same-domain topics') + return added + + def save(self) -> None: + """Write the bank out. Refills and expansions already do; this is for the end.""" + with self._state_lock: + self.store.save() + + # --------------------------------------------------------------- private + + def _next_nonce(self) -> int: + """A number no other call gets, so two prompts are never byte-identical. + + Shared across categories, which refill at the same time: two threads + reading the counter together would send the same prompt twice and halve the + diversity with nothing to show that it happened. + """ + with self._state_lock: + self._nonce += 1 + return self._nonce + + def _explore(self, users: Sequence[str]) -> List[Trajectory]: + """One batch of keyword calls, one per user message. + + Sampling params are only forwarded when set, so a plain callable explorer + that takes nothing else keeps working. + """ + if not users: + return [] + prompts: List[Trajectory] = [{ + 'messages': [{'role': 'system', 'content': self.prompts.system}, + {'role': 'user', 'content': user}], + } for user in users] + if self.sampling_params is None: + return self.explorer(prompts) + return self.explorer(prompts, sampling_params=self.sampling_params) + + def _record(self, record: Dict[str, Any]) -> None: + """Hand one keyword call to the sink, if there is one.""" + if self.sink is None: + return + with self._sink_lock: + self.sink(record) diff --git a/src/twinkle_agentic/challenger/new/keyword.py b/src/twinkle_agentic/challenger/new/keyword.py new file mode 100644 index 000000000..0d2c66bad --- /dev/null +++ b/src/twinkle_agentic/challenger/new/keyword.py @@ -0,0 +1,12 @@ + + +from typing import Optional + + +class KeywordGenerator: + + def __init__(system_prompt: Optional[str] = None): + pass + + def generate_keywords(self, text): + pass diff --git a/src/twinkle_agentic/envs/__init__.py b/src/twinkle_agentic/envs/__init__.py index 4633039c8..aa0ea7c32 100644 --- a/src/twinkle_agentic/envs/__init__.py +++ b/src/twinkle_agentic/envs/__init__.py @@ -1,5 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from .agentenv import AgentEnv -from .base import Env, StepResult +from .base import DEFAULT_TOOLS, TIMEOUT_EXIT_CODE, Env, StepResult from .env_tool import EnvTool +from .local import LocalEnv from .openenv import EnvPool, EnvPoolAdapter, OpenEnv, OpenEnvClient diff --git a/src/twinkle_agentic/envs/agentenv.py b/src/twinkle_agentic/envs/agentenv.py index eed3a9e81..36bd961d3 100644 --- a/src/twinkle_agentic/envs/agentenv.py +++ b/src/twinkle_agentic/envs/agentenv.py @@ -26,74 +26,10 @@ from twinkle.data_format import Trajectory from twinkle.data_format.message import Tool as ToolInfo from twinkle.utils import get_logger -from .base import Env, StepResult +from .base import DEFAULT_TOOLS, Env, StepResult, format_command_output, truncate_observation logger = get_logger() -_MAX_OBSERVATION_CHARS = 32 * 1024 - -_DEFAULT_TOOLS: List[ToolInfo] = [ - { - 'type': 'function', - 'function': { - 'name': 'run_command', - 'description': 'Run a shell command inside the sandbox and return its output.', - 'parameters': { - 'type': 'object', - 'properties': { - 'command': { - 'type': 'string', - 'description': 'The shell command to execute.' - }, - 'cwd': { - 'type': 'string', - 'description': 'Working directory (optional).' - }, - }, - 'required': ['command'], - }, - }, - }, - { - 'type': 'function', - 'function': { - 'name': 'write_file', - 'description': 'Write text content to a file inside the sandbox.', - 'parameters': { - 'type': 'object', - 'properties': { - 'path': { - 'type': 'string', - 'description': 'Absolute file path in the sandbox.' - }, - 'content': { - 'type': 'string', - 'description': 'Text content to write.' - }, - }, - 'required': ['path', 'content'], - }, - }, - }, - { - 'type': 'function', - 'function': { - 'name': 'read_file', - 'description': 'Read a text file from the sandbox.', - 'parameters': { - 'type': 'object', - 'properties': { - 'path': { - 'type': 'string', - 'description': 'Absolute file path in the sandbox.' - }, - }, - 'required': ['path'], - }, - }, - }, -] - def _require_e2b(): """Import the e2b SDK lazily with an actionable error message.""" @@ -107,23 +43,6 @@ def _require_e2b(): return Sandbox -def _truncate(text: str, limit: int = _MAX_OBSERVATION_CHARS) -> str: - if len(text) <= limit: - return text - return text[:limit] + f'\n... [truncated, {len(text) - limit} chars omitted]' - - -def _format_command_output(stdout: str, stderr: str, exit_code: int) -> str: - parts = [] - if stdout: - parts.append(stdout) - if stderr: - parts.append(f'[stderr]\n{stderr}') - if exit_code != 0: - parts.append(f'[exit code: {exit_code}]') - return _truncate('\n'.join(parts)) if parts else '(no output)' - - class AgentEnv(Env): """Env backed by one AgentENV sandbox per episode. @@ -299,7 +218,7 @@ def step(self, tool_name: str, arguments: Dict[str, Any] = None) -> StepResult: self._sandbox.files.write(arguments['path'], arguments.get('content', '')) observation = f"File written: {arguments['path']}" elif self._include_default_tools and tool_name == 'read_file': - observation = _truncate(str(self._sandbox.files.read(arguments['path']))) + observation = truncate_observation(str(self._sandbox.files.read(arguments['path']))) else: available = [t['function']['name'] for t in self.tools()] observation = f'Error: unknown tool {tool_name!r}. Available tools: {available}.' @@ -319,7 +238,7 @@ def tools(self) -> List[ToolInfo]: tools: List[ToolInfo] = [] if self._include_default_tools: custom_names = set(self._custom_handlers) - tools.extend(t for t in _DEFAULT_TOOLS if t['function']['name'] not in custom_names) + tools.extend(t for t in DEFAULT_TOOLS if t['function']['name'] not in custom_names) tools.extend(self._custom_tools) return tools @@ -350,7 +269,7 @@ def run_command(self, arguments: Dict[str, Any]) -> str: cwd=arguments.get('cwd'), timeout=int(arguments.get('timeout', self._command_timeout)), ) - return _format_command_output(result.stdout or '', result.stderr or '', result.exit_code or 0) + return format_command_output(result.stdout or '', result.stderr or '', result.exit_code or 0) except Exception as e: # noqa # The SDK raises on non-zero exit codes; surface the output # instead of failing the step so the model can react to it. @@ -359,7 +278,7 @@ def run_command(self, arguments: Dict[str, Any]) -> str: exit_code = getattr(e, 'exit_code', None) if exit_code is None: raise - return _format_command_output(stdout, stderr, exit_code) + return format_command_output(stdout, stderr, exit_code) def _kill_sandbox(self) -> None: if self._sandbox is None: diff --git a/src/twinkle_agentic/envs/base.py b/src/twinkle_agentic/envs/base.py index 52c8fda92..52df86ce9 100644 --- a/src/twinkle_agentic/envs/base.py +++ b/src/twinkle_agentic/envs/base.py @@ -6,6 +6,99 @@ from twinkle.data_format import Trajectory from twinkle.data_format.message import Tool as ToolInfo +# What :meth:`Env.run_script` returns when it had to kill the script. 124 is +# what GNU ``timeout`` uses, so a caller that logs the number is logging +# something a reader already knows how to interpret. +TIMEOUT_EXIT_CODE = 124 + +# Truncation guard for anything that becomes an observation: a command that +# dumps a whole file would otherwise spend the episode's context on one turn. +MAX_OBSERVATION_CHARS = 32 * 1024 + +# The tools every general-purpose environment advertises, sandboxed or local. +# Shared rather than restated per implementation: a trajectory built against one +# env has to replay on another, and it only does if the names and the argument +# spellings are the same object. +DEFAULT_TOOLS: List[ToolInfo] = [ + { + 'type': 'function', + 'function': { + 'name': 'run_command', + 'description': 'Run a shell command inside the sandbox and return its output.', + 'parameters': { + 'type': 'object', + 'properties': { + 'command': { + 'type': 'string', + 'description': 'The shell command to execute.' + }, + 'cwd': { + 'type': 'string', + 'description': 'Working directory (optional).' + }, + }, + 'required': ['command'], + }, + }, + }, + { + 'type': 'function', + 'function': { + 'name': 'write_file', + 'description': 'Write text content to a file inside the sandbox.', + 'parameters': { + 'type': 'object', + 'properties': { + 'path': { + 'type': 'string', + 'description': 'Absolute file path in the sandbox.' + }, + 'content': { + 'type': 'string', + 'description': 'Text content to write.' + }, + }, + 'required': ['path', 'content'], + }, + }, + }, + { + 'type': 'function', + 'function': { + 'name': 'read_file', + 'description': 'Read a text file from the sandbox.', + 'parameters': { + 'type': 'object', + 'properties': { + 'path': { + 'type': 'string', + 'description': 'Absolute file path in the sandbox.' + }, + }, + 'required': ['path'], + }, + }, + }, +] + + +def truncate_observation(text: str, limit: int = MAX_OBSERVATION_CHARS) -> str: + if len(text) <= limit: + return text + return text[:limit] + f'\n... [truncated, {len(text) - limit} chars omitted]' + + +def format_command_output(stdout: str, stderr: str, exit_code: int) -> str: + """One command's result as the model sees it.""" + parts = [] + if stdout: + parts.append(stdout) + if stderr: + parts.append(f'[stderr]\n{stderr}') + if exit_code != 0: + parts.append(f'[exit code: {exit_code}]') + return truncate_observation('\n'.join(parts)) if parts else '(no output)' + @dataclass class StepResult: @@ -30,6 +123,13 @@ class Env(ABC): executes already-split ``(tool_name, arguments)`` pairs. """ + #: How many times this environment had to be rebuilt under a caller that was + #: holding it. Reported rather than dropped: a run whose environments were + #: rebuilt twenty times produced its numbers under different conditions than + #: one that was rebuilt never, and that is invisible from the outputs alone. + #: Stays at zero for an environment that cannot be lost. + n_recoveries = 0 + def reset(self, trajectory: Optional[Trajectory] = None) -> StepResult: return StepResult() @@ -49,9 +149,111 @@ def step_batch( """ return [self.step(name, args or {}) for name, args in calls] + def run_script(self, source: str, interpreter: str = 'python', + timeout: Optional[int] = None) -> Tuple[int, str]: + """Run a whole script here; returns ``(exit_code, output)``. + + The execution path a *verifier* takes, as opposed to :meth:`step`, which + is the one the model takes. Both land in the same place, and that is the + point: a check has to observe the filesystem the episode actually wrote + to, so it runs in the environment rather than beside it. + + Args: + source: the script, not a path. + interpreter: ``'python'`` or ``'shell'``. + timeout: seconds; ``None`` means the environment's own default. + + Returns: + ``(exit_code, output)``. ``output`` is stdout followed by stderr, so + a traceback lands at the end rather than interleaved. A non-zero exit + code is the only failure signal callers should read -- the specific + value is the script's, except for a timeout, which is + :data:`TIMEOUT_EXIT_CODE`. + """ + raise NotImplementedError(f'{type(self).__name__} cannot run scripts') + + def ensure_ready(self) -> bool: + """Re-establish this environment if it has gone away. True if it did. + + For a caller that holds one environment across many jobs, losing it -- + evicted, timed out, runtime crashed -- otherwise ends the whole run. Safe + to call only where the workspace is about to be discarded anyway: a + mid-episode rebuild silently swaps the state a job is being judged on for + an empty directory, which is why recovery is an explicit call rather than + a retry hidden inside every dispatch. + + The default is ``False``: an environment that is a local process has + nothing to lose between calls and so is never not ready. + """ + return False + + def rebuild(self) -> None: + """Throw this environment away and stand a fresh one up in its place. + + For the caller that has a *working* environment it no longer trusts -- + one that keeps failing an operation it should not fail -- as opposed to + :meth:`ensure_ready`, which is about one that stopped answering. Counted + in :attr:`n_recoveries`. + + The default is a no-op, which is the truth for an environment holding + nothing worth rebuilding. + """ + + def clear(self) -> None: + """Return to a clean state, ready for the next episode. + + Called by whoever owns the environment, before handing it to a job that + must not see the previous one's files. Raising is the right answer for an + environment that could not clean itself: a silent no-op there means the + next job inherits a workspace, which lets a solver pass without doing + anything and makes a difficulty measurement meaningless. + + The default is a no-op because it is the truth for an environment holding + no state between calls -- the shape one-shot verification uses. That is + also what lets both halves run the same sequence: the code half clears + before every judgement too, and clearing nothing costs nothing. + """ + + def snapshot(self) -> Tuple[str, str]: + """The end state as ``(listing, error)``; both empty when there is none. + + What an episode left behind, for a caller that has to describe it to a + model -- writing a check against a workspace means knowing what is in it. + The two strings are kept apart because a snapshot that returns "empty" + when it means "I could not look" produces tasks whose only true assertion + is that nothing happened. + + The default is the honest answer for an environment that keeps nothing: + there is no end state to read back, which is why an env used only to run + one-shot checks does not have to implement this. + """ + return '', '' + def tools(self) -> List[ToolInfo]: return [] + def tool_manager(self, schemas: Optional[Sequence[ToolInfo]] = None) -> Any: + """A ``ToolManager`` that dispatches tool calls into this environment. + + What a rollout needs to let a model act here, so it is built once on the + environment rather than restated by every caller that owns one -- and a + caller holding N environments gets N managers that cannot be crossed, + which is the failure this prevents: an episode acting in one workspace + and being checked in another produces a task nobody can pass. + + Args: + schemas: the tool contract to advertise; defaults to :meth:`tools`. + Passed explicitly when an agent framework owns the names that go + into the prompt and this environment only supplies the + implementation. + """ + # Local import: ToolManager is a consumer of this package, and the tools + # package is not needed by an env that is only ever asked to run scripts. + from ..tools.tool_manager import ToolManager + from .env_tool import EnvTool + declared = list(schemas) if schemas is not None else self.tools() + return ToolManager(EnvTool.from_schemas(self, declared)) + def evaluate(self, trajectories: List[Trajectory], **kwargs) -> List[float]: return [0.0] * len(trajectories) diff --git a/src/twinkle_agentic/envs/local.py b/src/twinkle_agentic/envs/local.py new file mode 100644 index 000000000..69053cb90 --- /dev/null +++ b/src/twinkle_agentic/envs/local.py @@ -0,0 +1,260 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""LocalEnv: the training host as an environment. + +Same interface as the remote sandboxes -- :meth:`step` for the model's tool +calls, :meth:`run_script` for a verifier's script -- so a task does not have to +know which kind of environment it is being graded in. What differs is the +isolation: a subprocess in a new session with a capped address space, not a +microVM. + +That makes it the right environment for a check that is a few asserts over pure +computation. A microVM round trip costs hundreds of milliseconds and a +difficulty pass makes one call per candidate per rollout, so the same +verification that takes minutes here takes hours there, for a script that cannot +tell the difference. + +It makes it the wrong environment for running code against anything you would +mind that code reading or reaching: there is no filesystem or network isolation, +and the path checks below stop a mistake, not an attempt. Untrusted code belongs +in :class:`~twinkle_agentic.envs.agentenv.AgentEnv` or another sandbox. + +Two shapes, chosen by ``workspace``: + +* ``workspace=<dir>``: that directory is the working directory for every call and + outlives them all. What a multi-turn episode needs -- the model writes a file + with one tool call, and the check script reads it back after the episode ends. +* ``workspace=None``: every call runs in a fresh temporary directory that is + removed afterwards. What one-shot verification needs, and the reason the code + half has no workspace to reset: nothing survives a call to leak into the next + one. The file tools are withdrawn in this shape, because a file written by one + call would not be there for the next. +""" +import os +import resource +import shlex +import shutil +import signal +import subprocess +import sys +import tempfile +from typing import Any, Dict, List, Optional, Tuple + +from twinkle.data_format.message import Tool as ToolInfo +from twinkle.utils import get_logger +from .base import (DEFAULT_TOOLS, TIMEOUT_EXIT_CODE, Env, StepResult, format_command_output, + truncate_observation) + +logger = get_logger() + +# Kept deterministic and single-threaded: a check that changes its answer with +# the machine's core count is not a check. Matches what the sandboxes set. +_SCRIPT_ENVS = { + 'MPLBACKEND': 'Agg', + 'PYTHONHASHSEED': '0', + 'OMP_NUM_THREADS': '1', + 'MKL_NUM_THREADS': '1', + 'TOKENIZERS_PARALLELISM': 'false', +} + + +class LocalEnv(Env): + """Run scripts and tool calls on this machine. See the module docstring.""" + + def __init__(self, + workspace: Optional[str] = None, + command_timeout: int = 60, + memory_limit_gb: Optional[float] = 2.0, + envs: Optional[Dict[str, str]] = None): + """ + Args: + workspace: persistent working directory, created if absent. ``None`` + gives every call its own temporary directory and withdraws the + file tools. + command_timeout: default seconds per call, when the caller does not + pass one. + memory_limit_gb: address-space cap per call, so one runaway script + cannot take the trainer down with it. ``None`` to not cap. + envs: extra environment variables for the child process. + """ + self._workspace = os.path.abspath(workspace) if workspace else None + if self._workspace: + os.makedirs(self._workspace, exist_ok=True) + self._command_timeout = command_timeout + self._memory_limit_gb = memory_limit_gb + self._envs = dict(envs or {}) + + @property + def workspace(self) -> Optional[str]: + """The persistent working directory, or None in the throwaway shape.""" + return self._workspace + + # ------------------------------------------------------------------ + # Env interface + # ------------------------------------------------------------------ + + def run_script(self, source: str, interpreter: str = 'python', + timeout: Optional[int] = None) -> Tuple[int, str]: + timeout = self._command_timeout if timeout is None else timeout + # The script file is never written into the workspace. A persistent + # workspace gets read back -- by a snapshot, or by a check that lists the + # directory -- and a stray _script.py in there reads as something the + # episode created. With no workspace this same directory is the working + # directory, which is what makes that shape leave nothing behind. + holder = tempfile.mkdtemp(prefix='twinkle_local_') + try: + if interpreter == 'python': + path = os.path.join(holder, '_script.py') + with open(path, 'w', encoding='utf-8') as f: + f.write(source + '\n') + argv = [sys.executable, path] + elif interpreter in ('shell', 'bash'): + # Not a login shell: sourcing the host's profile prepends whatever + # banner it prints to the output of every command, and a check + # comparing that output against an expected string then fails on + # the banner. PATH and the rest are inherited from the trainer, + # which is already in the right environment. + argv = ['/bin/bash', '-c', source] + else: + return 1, f'unsupported interpreter {interpreter!r}; use python or shell' + return self._spawn(argv, self._workspace or holder, timeout) + finally: + shutil.rmtree(holder, ignore_errors=True) + + def step(self, tool_name: str, arguments: Dict[str, Any] = None) -> StepResult: + arguments = arguments or {} + try: + if tool_name == 'run_command': + observation = self.run_command(arguments) + elif tool_name in ('write_file', 'read_file'): + if self._workspace is None: + # Not an error the model can recover from by rephrasing, so + # it says what is missing rather than what went wrong. + observation = (f'Error: {tool_name} needs a persistent workspace; ' + 'this environment runs every call in a fresh directory.') + elif tool_name == 'write_file': + observation = self._write_file(arguments) + else: + observation = self._read_file(arguments) + else: + available = [t['function']['name'] for t in self.tools()] + observation = f'Error: unknown tool {tool_name!r}. Available tools: {available}.' + return StepResult(observation=observation) + except Exception as e: # noqa + # Same contract as the sandboxed envs: a tool error is an + # observation, so the rollout loop can let the model recover. + logger.warning(f'LocalEnv step error (tool={tool_name}): {e}') + return StepResult(observation=f'Error: {e}', info={'error': str(e)}) + + def tools(self) -> List[ToolInfo]: + if self._workspace is None: + # Nothing an episode could build on: every call would start from an + # empty directory, so this shape is a verifier, not an environment. + return [] + return list(DEFAULT_TOOLS) + + def clear(self) -> None: + """Empty the workspace. A no-op in the throwaway shape, which has none. + + Raises rather than reporting, per :meth:`Env.clear`: a caller that clears + before every job is depending on this, and the failure it guards against + -- a job inheriting the previous one's files -- is invisible downstream. + """ + if self._workspace is None: + return + for name in os.listdir(self._workspace): + path = os.path.join(self._workspace, name) + if os.path.isdir(path) and not os.path.islink(path): + shutil.rmtree(path) + else: + os.remove(path) + + # ------------------------------------------------------------------ + # Tools + # ------------------------------------------------------------------ + + def run_command(self, arguments: Dict[str, Any]) -> str: + """Run a shell command; public so custom tool handlers can reuse it.""" + command = arguments.get('command') + if not command: + return "Error: 'command' argument is required." + cwd = arguments.get('cwd') + if cwd: + command = f'cd {shlex.quote(str(cwd))} && {command}' + exit_code, output = self.run_script(command, 'shell', timeout=arguments.get('timeout')) + # stderr is already folded into output by run_script, hence the empty + # stream here: what this call adds is the exit-code line. + return format_command_output(output, '', exit_code) + + def _write_file(self, arguments: Dict[str, Any]) -> str: + path = self._resolve(arguments['path']) + os.makedirs(os.path.dirname(path) or '.', exist_ok=True) + with open(path, 'w', encoding='utf-8') as f: + f.write(arguments.get('content', '')) + return f"File written: {arguments['path']}" + + def _read_file(self, arguments: Dict[str, Any]) -> str: + with open(self._resolve(arguments['path']), encoding='utf-8', errors='replace') as f: + return truncate_observation(f.read()) + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def _resolve(self, path: str) -> str: + """Resolve a tool-supplied path inside the workspace. + + The workspace is the root, so an absolute path means absolute *in here* + -- the tool schema is shared with the sandboxed envs, where it genuinely + is the filesystem root. An escape raises: with no isolation underneath, a + relative path with enough ``..`` in it would otherwise be writing to the + training host. This bounds a mistake; it is not a security boundary, + since ``run_command`` reaches the same filesystem directly. + """ + root = os.path.realpath(self._workspace) + target = os.path.realpath(os.path.join(root, str(path).lstrip('/'))) + if target != root and not target.startswith(root + os.sep): + raise ValueError(f'path {path!r} escapes the workspace') + return target + + def _spawn(self, argv: List[str], cwd: str, timeout: int) -> Tuple[int, str]: + env = dict(os.environ, **_SCRIPT_ENVS, **self._envs) + # Inherited from the trainer, and a check that imports torch would + # otherwise take a share of a GPU that is mid-generation. + env.pop('CUDA_VISIBLE_DEVICES', None) + + def _limit(): + if self._memory_limit_gb: + cap = int(self._memory_limit_gb * 1024**3) + resource.setrlimit(resource.RLIMIT_AS, (cap, cap)) + + try: + proc = subprocess.Popen(argv, cwd=cwd, env=env, + stdout=subprocess.PIPE, stderr=subprocess.PIPE, + text=True, errors='replace', + start_new_session=True, preexec_fn=_limit) + except Exception as e: # noqa + # A spawn failure is the host's problem, not the script's, and it + # comes back as a failed run so one bad call cannot end a whole pass. + return 1, f'{type(e).__name__}: {e}' + try: + out, err = proc.communicate(timeout=timeout) + out, err = out or '', err or '' + # A newline between the streams: a stdout line left unterminated + # swallows the first line of the traceback that follows it. + if out and err and not out.endswith('\n'): + out += '\n' + return proc.returncode, out + err + except subprocess.TimeoutExpired: + # killpg, not kill: start_new_session gave the script its own process + # group, so a script that forked cannot leave grandchildren running. + try: + os.killpg(proc.pid, signal.SIGKILL) + except ProcessLookupError: + pass + try: + proc.communicate(timeout=5) + except Exception: # noqa # already killed; the output is forfeit + pass + # A killed script has no traceback to explain itself with, so the + # output has to say why it produced nothing. + return TIMEOUT_EXIT_CODE, f'execution timed out after {timeout}s (possible infinite loop)' diff --git a/src/twinkle_agentic/harness/base.py b/src/twinkle_agentic/harness/base.py index 153a34a5c..dadabce63 100644 --- a/src/twinkle_agentic/harness/base.py +++ b/src/twinkle_agentic/harness/base.py @@ -34,7 +34,7 @@ def tool_schemas(self) -> List[Dict[str, Any]]: advertise the identical set; the Env owns the *implementation*. Build the executing side from the same list:: - tm = ToolManager(EnvTool.from_schemas(env, harness.tool_schemas())) + tm = env.tool_manager(harness.tool_schemas()) Skipping that step lets the prompt advertise tools the Env cannot run, and every call comes back as an unknown-tool error. diff --git a/src/twinkle_agentic/harness/ms_agent.py b/src/twinkle_agentic/harness/ms_agent.py index 8807e58ae..e1f8f7bbd 100644 --- a/src/twinkle_agentic/harness/ms_agent.py +++ b/src/twinkle_agentic/harness/ms_agent.py @@ -15,7 +15,7 @@ harness = MsAgentHarness(config) harness.prepare() - tool_manager = ToolManager(EnvTool.from_schemas(env, harness.tool_schemas())) + tool_manager = env.tool_manager(harness.tool_schemas()) rollout = MultiTurnRollout(sampler, template, tool_manager=tool_manager, harness=harness) outs = rollout([harness.start(q) for q in queries]) diff --git a/src/twinkle_agentic/preprocessor/data_juicer.py b/src/twinkle_agentic/preprocessor/data_juicer.py index cd6b10d69..fad79564a 100644 --- a/src/twinkle_agentic/preprocessor/data_juicer.py +++ b/src/twinkle_agentic/preprocessor/data_juicer.py @@ -7,7 +7,7 @@ from typing import Any, Dict, List, Tuple from twinkle.preprocessor import Preprocessor -from .utils import msg_content_text +from twinkle_agentic.utils.message_utils import msg_content_text # ── Shared helpers ──────────────────────────────────────────────────────────── diff --git a/src/twinkle_agentic/preprocessor/dead_loop_filter.py b/src/twinkle_agentic/preprocessor/dead_loop_filter.py index d2316f0a9..257014c25 100644 --- a/src/twinkle_agentic/preprocessor/dead_loop_filter.py +++ b/src/twinkle_agentic/preprocessor/dead_loop_filter.py @@ -5,7 +5,8 @@ from typing import Any, Dict, List, Tuple from twinkle.preprocessor import Preprocessor -from .utils import cjk_ratio, is_agent_row, msg_content_text +from twinkle_agentic.utils.message_utils import is_agent_row, msg_content_text +from twinkle_agentic.utils.text_utils import cjk_ratio # ── Hesitation-marker regexes ───────────────────────────────────────────────── # diff --git a/src/twinkle_agentic/preprocessor/dedup_filter.py b/src/twinkle_agentic/preprocessor/dedup_filter.py index 58476d6f8..72a2eea01 100644 --- a/src/twinkle_agentic/preprocessor/dedup_filter.py +++ b/src/twinkle_agentic/preprocessor/dedup_filter.py @@ -4,7 +4,7 @@ from typing import Any, Dict, List, Tuple from twinkle.preprocessor import Preprocessor -from .utils import msg_content_text +from twinkle_agentic.utils.message_utils import msg_content_text _SYSTEM_INJECTION_RE = re.compile(r'^<(?:system-reminder|system_reminder|context|user_info|attached_files)[ >]', re.IGNORECASE) diff --git a/src/twinkle_agentic/preprocessor/experimental/score_filter.py b/src/twinkle_agentic/preprocessor/experimental/score_filter.py index 228a8e25b..94333d8ef 100644 --- a/src/twinkle_agentic/preprocessor/experimental/score_filter.py +++ b/src/twinkle_agentic/preprocessor/experimental/score_filter.py @@ -30,8 +30,8 @@ from twinkle.preprocessor import Preprocessor from twinkle.template import Template from twinkle.utils import get_logger +from ..logprob_utils import _chr_min_distinct, _ifd_family_metrics, _lp_to_jsonable, _pad_batch, _to_int_list from .llm_backend import LLMBackend -from ..utils import _chr_min_distinct, _ifd_family_metrics, _lp_to_jsonable, _pad_batch, _to_int_list logger = get_logger() diff --git a/src/twinkle_agentic/preprocessor/hard_filter.py b/src/twinkle_agentic/preprocessor/hard_filter.py index d303ebce4..535fc9711 100644 --- a/src/twinkle_agentic/preprocessor/hard_filter.py +++ b/src/twinkle_agentic/preprocessor/hard_filter.py @@ -4,7 +4,8 @@ from typing import Any, Dict, List, Optional, Tuple from twinkle.preprocessor import Preprocessor -from .utils import cjk_ratio, is_agent_row, msg_content_text, msg_has_media, normalize_tool_calls +from twinkle_agentic.utils.message_utils import is_agent_row, msg_content_text, msg_has_media, normalize_tool_calls +from twinkle_agentic.utils.text_utils import cjk_ratio # ── Language detection ──────────────────────────────────────────────────────── diff --git a/src/twinkle_agentic/preprocessor/intent_classifier.py b/src/twinkle_agentic/preprocessor/intent_classifier.py index 079719ed7..baa04c13e 100644 --- a/src/twinkle_agentic/preprocessor/intent_classifier.py +++ b/src/twinkle_agentic/preprocessor/intent_classifier.py @@ -6,7 +6,7 @@ from twinkle.data_format import pack_value from twinkle.preprocessor import Preprocessor from twinkle.utils import get_logger -from .utils import msg_content_text, normalize_tool_calls +from twinkle_agentic.utils.message_utils import msg_content_text, normalize_tool_calls logger = get_logger() diff --git a/src/twinkle_agentic/preprocessor/language_filter.py b/src/twinkle_agentic/preprocessor/language_filter.py index ec0c7fac9..b86c0b021 100644 --- a/src/twinkle_agentic/preprocessor/language_filter.py +++ b/src/twinkle_agentic/preprocessor/language_filter.py @@ -18,8 +18,8 @@ from twinkle.preprocessor import Filter from twinkle.utils import get_logger - -from .message_utils import cjk_ratio, msg_content_text +from twinkle_agentic.utils.message_utils import msg_content_text +from twinkle_agentic.utils.text_utils import cjk_ratio logger = get_logger() diff --git a/src/twinkle_agentic/preprocessor/message_normalizer.py b/src/twinkle_agentic/preprocessor/message_normalizer.py index cd11c7bc6..36f169f4e 100644 --- a/src/twinkle_agentic/preprocessor/message_normalizer.py +++ b/src/twinkle_agentic/preprocessor/message_normalizer.py @@ -21,7 +21,7 @@ from twinkle.preprocessor import Preprocessor from twinkle.template.tools import ToolCallRegistry -from .utils import msg_content_text, msg_has_media, normalize_tool_calls +from twinkle_agentic.utils.message_utils import msg_content_text, msg_has_media, normalize_tool_calls # IGNORECASE absorbs every variant ("Read HEARTBEAT.md", "HEARTBEAT_OK", # "duplicate heartbeat", etc.) under the single token "heartbeat". diff --git a/src/twinkle_agentic/preprocessor/message_sanity.py b/src/twinkle_agentic/preprocessor/message_sanity.py index 2014976e5..38f966d15 100644 --- a/src/twinkle_agentic/preprocessor/message_sanity.py +++ b/src/twinkle_agentic/preprocessor/message_sanity.py @@ -10,12 +10,9 @@ from typing import Any, Dict, List, Optional, Tuple from twinkle.preprocessor import Preprocessor -from .utils import (build_sensitive_regex, cjk_ratio, is_agent_row, load_sensitive_words, msg_content_text, - msg_has_media, msg_has_payload, normalize_tool_calls) - -# Backward-compat re-exports. -_msg_content_text = msg_content_text -_normalize_tool_calls = normalize_tool_calls +from twinkle_agentic.utils.message_utils import (is_agent_row, msg_content_text, msg_has_media, msg_has_payload, + normalize_tool_calls) +from twinkle_agentic.utils.text_utils import build_sensitive_regex, cjk_ratio, load_sensitive_words _VALID_ROLES = {'system', 'user', 'assistant', 'tool'} _IDENTIFIER_RE = re.compile(r'^[a-zA-Z_][a-zA-Z0-9_.\-]*$') diff --git a/src/twinkle_agentic/preprocessor/offline/decontaminate.py b/src/twinkle_agentic/preprocessor/offline/decontaminate.py index 7ec9f7572..7dd8e3831 100644 --- a/src/twinkle_agentic/preprocessor/offline/decontaminate.py +++ b/src/twinkle_agentic/preprocessor/offline/decontaminate.py @@ -19,9 +19,8 @@ from typing import Any, Dict, Iterable, List, Set, Tuple from twinkle.preprocessor import Preprocessor - +from twinkle_agentic.utils.message_utils import msg_content_text from .. import label_schema as L -from ..message_utils import msg_content_text KEY_CONTAMINATED = 'contaminated' diff --git a/src/twinkle_agentic/preprocessor/offline/near_dedup.py b/src/twinkle_agentic/preprocessor/offline/near_dedup.py index e1b255688..d8999ebb6 100644 --- a/src/twinkle_agentic/preprocessor/offline/near_dedup.py +++ b/src/twinkle_agentic/preprocessor/offline/near_dedup.py @@ -22,8 +22,7 @@ from twinkle.preprocessor import Preprocessor from twinkle.utils import get_logger - -from ..message_utils import msg_content_text +from twinkle_agentic.utils.message_utils import msg_content_text logger = get_logger() diff --git a/src/twinkle_agentic/preprocessor/structural_noise.py b/src/twinkle_agentic/preprocessor/structural_noise.py index 27945a4d8..257899636 100644 --- a/src/twinkle_agentic/preprocessor/structural_noise.py +++ b/src/twinkle_agentic/preprocessor/structural_noise.py @@ -18,9 +18,8 @@ from typing import Any, Dict from twinkle.preprocessor import Mapper - +from twinkle_agentic.utils.message_utils import msg_content_text, normalize_tool_calls from . import label_schema as L -from .message_utils import msg_content_text, normalize_tool_calls KEY_NOISE_RATIO = 'structural_noise_ratio' diff --git a/src/twinkle_agentic/preprocessor/token_soup.py b/src/twinkle_agentic/preprocessor/token_soup.py index a18c08f5e..dd99cba09 100644 --- a/src/twinkle_agentic/preprocessor/token_soup.py +++ b/src/twinkle_agentic/preprocessor/token_soup.py @@ -5,7 +5,7 @@ from typing import Any, Dict, List, Tuple from twinkle.preprocessor import Preprocessor -from .utils import msg_content_text +from twinkle_agentic.utils.message_utils import msg_content_text # ── Pre-compiled patterns ───────────────────────────────────────────────────── diff --git a/src/twinkle_agentic/preprocessor/utils.py b/src/twinkle_agentic/preprocessor/utils.py deleted file mode 100644 index 803986db0..000000000 --- a/src/twinkle_agentic/preprocessor/utils.py +++ /dev/null @@ -1,28 +0,0 @@ -"""Backward-compat re-export shim (AUDIT A2). - -``utils.py`` was split into two focused modules: -- :mod:`logprob_utils` — log-prob data-selection math (IFD / S-IFD / chr_min), - used only by the experimental log-prob scorers. -- :mod:`message_utils` — message-format helpers used by every active step. - -This shim keeps historical ``from .utils import ...`` imports working. Prefer -importing from the focused modules directly in new code. -""" -from .logprob_utils import (_chr_min_distinct, _chr_min_weighted, # noqa: F401 - _extract_logprob, _ifd_family_metrics, - _lp_to_jsonable, _mean_logprob_delta, _pad_batch, - _to_int_list) -from .message_utils import (CJK_CHARS_RE, build_sensitive_regex, # noqa: F401 - cjk_ratio, is_agent_row, load_sensitive_words, - msg_content_text, msg_has_media, msg_has_payload, - normalize_tool_calls) - -__all__ = [ - # message utils - 'msg_content_text', 'msg_has_media', 'msg_has_payload', 'normalize_tool_calls', - 'cjk_ratio', 'CJK_CHARS_RE', 'load_sensitive_words', 'build_sensitive_regex', - 'is_agent_row', - # logprob utils - '_extract_logprob', '_to_int_list', '_chr_min_distinct', '_chr_min_weighted', - '_ifd_family_metrics', '_mean_logprob_delta', '_lp_to_jsonable', '_pad_batch', -] diff --git a/src/twinkle_agentic/reward/f1.py b/src/twinkle_agentic/reward/f1.py index 3828ca7a6..cda65e923 100644 --- a/src/twinkle_agentic/reward/f1.py +++ b/src/twinkle_agentic/reward/f1.py @@ -4,6 +4,7 @@ from typing import Any, Dict, List, Tuple from twinkle.reward import Reward +from twinkle_agentic.utils.message_utils import assistant_text _BOXED_MARKER = '\\boxed{' @@ -35,17 +36,6 @@ def _extract_final_answer(completion: str) -> str: return out -def _last_assistant_text(traj: Dict[str, Any]) -> str: - for msg in reversed(traj.get('messages', [])): - if msg.get('role') != 'assistant': - continue - content = msg.get('content') or '' - if isinstance(content, str): - return content - return '\n'.join(p.get('text', '') for p in content if isinstance(p, dict) and p.get('type') == 'text') - return '' - - def _stem(tok: str) -> str: from nltk.stem import PorterStemmer return PorterStemmer().stem(tok) if len(tok) >= 4 and tok.isalpha() else tok @@ -137,7 +127,7 @@ def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: rewards = [] for traj in trajectories: golds = [val for key, val in traj.get('user_data', []) or [] if key == 'ground_truth' and val] - pred = self._extract(_last_assistant_text(traj)) + pred = self._extract(assistant_text(traj)) if golds: f1 = max(_f1_score(pred, g)[0] for g in golds) else: @@ -157,16 +147,16 @@ def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: # Newline-joined so ``^`` line anchors work even when # multiple assistant turns exist. - assistant_text = '\n'.join( + all_assistant_text = '\n'.join( m.get('content', '') or '' for m in msgs if m.get('role') == 'assistant' and isinstance(m.get('content'), str)) - if not self._HAS_BOXED_RE.search(assistant_text): + if not self._HAS_BOXED_RE.search(all_assistant_text): rewards.append(0.0) continue steps: set = set() - for match in self._STEP_LINE_RE.finditer(assistant_text): + for match in self._STEP_LINE_RE.finditer(all_assistant_text): try: steps.add(int(match.group(1))) except ValueError: @@ -207,7 +197,7 @@ def _extract(self, completion: str) -> str: def _trajectory_f1(self, traj: Dict[str, Any]) -> float: golds = [val for key, val in traj.get('user_data', []) or [] if key == 'ground_truth' and val] - pred = self._extract(_last_assistant_text(traj)) + pred = self._extract(assistant_text(traj)) if golds: return max(_f1_score(pred, g)[0] for g in golds) f1, _ = _f1_score(pred, '') diff --git a/src/twinkle_agentic/utils/code_utils.py b/src/twinkle_agentic/utils/code_utils.py new file mode 100644 index 000000000..37368045e --- /dev/null +++ b/src/twinkle_agentic/utils/code_utils.py @@ -0,0 +1,106 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Reading code back out of a model's reply. + +A model asked for a python snippet -- a check script, a solution, a repro -- +fences it. Taking that fence back off is the same work whatever the snippet is +*for*, so it lives here rather than in one challenger. + +A reply that fenced nothing is rejected, not read some other way. Reading bare +text and tool-call arguments was tried, and both come down to guessing where the +code starts and ends and then asking a parser whether the guess was plausible; a +guess that parses but is short a few lines is indistinguishable from a good one, +and it becomes a task. Requiring the fence trades those silent losses for a loud +one -- the reply is refused and the model is asked again. + +Nothing here is pinned to a model family, and what *is* knowledge gets passed in +rather than assumed. The caller knows which language it asked for and says so with +``language_tags``. The caller does not know which model answered, so this module +absorbs that: reasoning is cut by a list of markers rather than the one tag a +given model emits. Handing that up to a challenger only moves the ignorance -- it +would then guess ``</think>`` and be right for one model family. + +There are two ways out, and they differ on the replies that fenced no code -- no +fence at all, or one left empty. :func:`parse_fenced_code` answers None to both, +:func:`unwrap_code` hands the reply back whole for the first and ``''`` for the +second, where the model did say the code went here and put nothing there. +""" +import re +from functools import lru_cache +from typing import Optional, Pattern, Tuple + +__all__ = [ + 'PYTHON_TAGS', + 'parse_fenced_code', + 'strip_reasoning', + 'unwrap_code', +] + +_REASONING_END_MARKERS = ('</think>', '</thinking>', '</reasoning>', '<|end_of_thought|>') +PYTHON_TAGS: Tuple[str, ...] = ('python', 'py') + + +@lru_cache(maxsize=None) +def _fence_re(language_tags: Tuple[str, ...]) -> Pattern: + """A fenced block whose language tag is one of ``language_tags``, or absent. + + ``python``, ``py``, ``Python`` and ``python3`` are one intent spelled four + ways, so a tag matches case-insensitively and with any version suffix. A tag + that is not on the list -- ``bash``, ``json`` -- is a different intent, and is + not read as code at all. + """ + alts = '|'.join(re.escape(tag) for tag in language_tags) + label = r'(?:(?:%s)[\d.]*)?' % alts if alts else '' + return re.compile(r'```[ \t]*%s[ \t]*\r?\n(.*?)```' % label, re.S | re.I) + + +def strip_reasoning(text: str) -> str: + """``text`` with everything up to the end of the model's thinking removed. + + The last marker anywhere in the reply wins: reasoning precedes the answer, and + a model that opens a second thought after answering is still answering last. + Text with no marker is returned unchanged. + """ + body = text or '' + cut = 0 + for marker in _REASONING_END_MARKERS: + idx = body.rfind(marker) + if idx >= 0: + cut = max(cut, idx + len(marker)) + return body[cut:] + + +def parse_fenced_code(text: str, language_tags: Tuple[str, ...] = PYTHON_TAGS) -> Optional[str]: + """The last block ``text`` fenced as that language, or None if there is none. + + The last one, not the first: a model often drafts a version before the final + one, and the block it ends on is its answer. + + What is inside is taken as given -- a fence is the model saying which part is + the code, so second-guessing it would throw away the one piece of the reply + that was unambiguous. Whether it runs is the sandbox's answer to give. + + A fence the model opened and left empty answers None too, on the grounds that + a caller who cannot use a missing script cannot use an empty one either. Use + this when nothing downstream will judge the result and a wrong guess becomes a + task. + """ + blocks = _fence_re(language_tags).findall(strip_reasoning(text)) + return (blocks[-1].strip() if blocks else '') or None + + +def unwrap_code(text: str, language_tags: Tuple[str, ...] = PYTHON_TAGS) -> str: + """``text`` with the model's packaging taken off, always a string. + + Takes the fence off if there is one and hands the reply back whole if there is + not, on the reading that a reply to "write the code" *is* the code however it + was dressed. An empty fence answers ``''``, because the model did mark where + the code went and put nothing there. + + Those two are the whole difference from :func:`parse_fenced_code`, which + answers None to both. Use this on an answer that is about to be run -- the + sandbox is the better judge of whether that was code, and it says so with an + exit status. + """ + body = strip_reasoning(text) + blocks = _fence_re(language_tags).findall(body) + return blocks[-1].strip() if blocks else body.strip() diff --git a/src/twinkle_agentic/preprocessor/message_utils.py b/src/twinkle_agentic/utils/message_utils.py similarity index 57% rename from src/twinkle_agentic/preprocessor/message_utils.py rename to src/twinkle_agentic/utils/message_utils.py index 3ee95b1eb..64a543b6c 100644 --- a/src/twinkle_agentic/preprocessor/message_utils.py +++ b/src/twinkle_agentic/utils/message_utils.py @@ -1,14 +1,35 @@ -"""Message-format utilities shared across active preprocessor steps. - -Split out of ``utils.py`` (AUDIT A2): content projection, tool-call -normalization, CJK ratio, sensitive-word regex, and agent-row detection. These -are the helpers every cleaning step depends on, independent of the log-prob -scoring math (see :mod:`logprob_utils`). +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Reading what messages carry. + +A message's ``content`` is a plain string in the simple case and a list of typed +parts when it is multimodal, so every caller that wants the text has to handle +both shapes. ``tool_calls`` has the same problem one level up: a round trip +through PyArrow or a JSONL dataset can leave it as a string holding JSON, or a +list of such strings, so asking "did the model call a tool" means decoding +before looking. A whole conversation raises the same kind of question -- which +turn is the model's answer, did it use tools at all -- answered the same way, +by looking rather than trusting the shape. + +These live here rather than under any one consumer because none of the questions +is a preprocessing one: a challenger reading a model's reply, a reward scoring +one, and a cleaning step filtering one all ask them. Each place that answered on +its own answered differently -- handing back the raw list, or raising on it. + +Kept to a plain ``Dict`` rather than :class:`~twinkle.data_format.Message` on +purpose: rows read straight off disk go through these too, before anything has +promised they match the type. """ import json -import os -import re -from typing import Any, Dict, List, Optional, Set +from typing import Any, Dict, List, Optional + +__all__ = [ + 'assistant_text', + 'is_agent_row', + 'msg_content_text', + 'msg_has_media', + 'msg_has_payload', + 'normalize_tool_calls', +] def msg_content_text(msg: Dict[str, Any]) -> str: @@ -36,9 +57,6 @@ def msg_has_payload(msg: Dict[str, Any]) -> bool: or msg_has_media(msg)) -_CJK_RE = re.compile(r'[\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\uac00-\ud7a3]') - - def normalize_tool_calls(msg: Dict[str, Any]) -> Optional[List[Any]]: """Return ``tool_calls`` as a list of dicts, handling PyArrow/HF serialization artifacts.""" tcs = msg.get('tool_calls') @@ -75,50 +93,6 @@ def normalize_tool_calls(msg: Dict[str, Any]) -> Optional[List[Any]]: return result -CJK_CHARS_RE = _CJK_RE - - -def cjk_ratio(text: str) -> float: - """Fraction of non-whitespace characters that are CJK.""" - chars = text.replace(' ', '').replace('\n', '').replace('\t', '') - if not chars: - return 0.0 - return len(CJK_CHARS_RE.findall(chars)) / len(chars) - - -def load_sensitive_words(path: Optional[str]) -> Set[str]: - """Load from external file (one word per line). Blank lines and #-comments ignored.""" - if not path or not os.path.isfile(path): - return set() - words: Set[str] = set() - with open(path, encoding='utf-8') as f: - for line in f: - line = line.strip() - if line and not line.startswith('#'): - words.add(line) - return words - - -def build_sensitive_regex(words: Set[str]) -> Optional['re.Pattern']: - """Build a compiled regex from a set of words. Returns None if empty.""" - if not words: - return None - cjk_words = [] - latin_words = [] - cjk_re = re.compile(r'[\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\uac00-\ud7a3]') - for w in sorted(words): - if cjk_re.search(w): - cjk_words.append(re.escape(w)) - else: - latin_words.append(re.escape(w)) - parts = [] - if latin_words: - parts.append(r'\b(' + '|'.join(latin_words) + r')\b') - if cjk_words: - parts.append('(' + '|'.join(cjk_words) + ')') - return re.compile('|'.join(parts), re.IGNORECASE) - - def is_agent_row(messages) -> bool: """Return True if the conversation contains tool interactions (agent trace). @@ -135,3 +109,19 @@ def is_agent_row(messages) -> bool: if normalize_tool_calls(m): return True return False + + +def assistant_text(trajectory: Dict[str, Any]) -> str: + """The last assistant message's text, or '' if the model produced none. + + Explorers differ in what else they attach -- token ids, logprobs, tool + turns -- but every one of them leaves the reply as an assistant message, + so this is the one field a parser can rely on. + + The *last* one: a conversation that went through tools has several, and the + model's answer is the turn it finished on. + """ + for message in reversed(trajectory.get('messages') or []): + if isinstance(message, dict) and message.get('role') == 'assistant': + return msg_content_text(message) + return '' diff --git a/src/twinkle_agentic/utils/text_utils.py b/src/twinkle_agentic/utils/text_utils.py new file mode 100644 index 000000000..a14a44bc2 --- /dev/null +++ b/src/twinkle_agentic/utils/text_utils.py @@ -0,0 +1,62 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Shape of a piece of text, and word-list matching against it. + +These take plain strings, not messages: which script a string is written in, and +whether it hits a banned-word list. Both questions come up wherever text arrives +from a model or a dataset -- filtering a corpus, deciding a reply's language, +refusing to train on something -- so they do not belong to any one of those. + +The CJK class covers Han, Hiragana, Katakana and Hangul, which is what callers +mean by "CJK" here even though Korean is not Chinese-Japanese. +""" +import os +import re +from typing import Optional, Set + +__all__ = ['CJK_CHARS_RE', 'build_sensitive_regex', 'cjk_ratio', 'load_sensitive_words'] + +CJK_CHARS_RE = re.compile(r'[\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\uac00-\ud7a3]') + + +def cjk_ratio(text: str) -> float: + """Fraction of non-whitespace characters that are CJK.""" + chars = text.replace(' ', '').replace('\n', '').replace('\t', '') + if not chars: + return 0.0 + return len(CJK_CHARS_RE.findall(chars)) / len(chars) + + +def load_sensitive_words(path: Optional[str]) -> Set[str]: + """Load from external file (one word per line). Blank lines and #-comments ignored.""" + if not path or not os.path.isfile(path): + return set() + words: Set[str] = set() + with open(path, encoding='utf-8') as f: + for line in f: + line = line.strip() + if line and not line.startswith('#'): + words.add(line) + return words + + +def build_sensitive_regex(words: Set[str]) -> Optional['re.Pattern']: + """Build a compiled regex from a set of words. Returns None if empty. + + Latin words get word boundaries, CJK ones cannot: there is no ``\\b`` between + two Han characters, so a boundary there would never match. + """ + if not words: + return None + cjk_words = [] + latin_words = [] + for w in sorted(words): + if CJK_CHARS_RE.search(w): + cjk_words.append(re.escape(w)) + else: + latin_words.append(re.escape(w)) + parts = [] + if latin_words: + parts.append(r'\b(' + '|'.join(latin_words) + r')\b') + if cjk_words: + parts.append('(' + '|'.join(cjk_words) + ')') + return re.compile('|'.join(parts), re.IGNORECASE) diff --git a/src/twinkle_agentic/verifier/__init__.py b/src/twinkle_agentic/verifier/__init__.py index a4e345a94..5dd8cc2ac 100644 --- a/src/twinkle_agentic/verifier/__init__.py +++ b/src/twinkle_agentic/verifier/__init__.py @@ -1,12 +1,12 @@ from .result_check import (Check, CheckContext, CheckOutcome, CheckReport, - checks_from_dicts, local_runner, run_checks) + checks_from_dicts, run_checks) from .rubric_score import (CRITERIA, DIMENSIONS, Criterion, RubricResult, build_rubric_prompt, parse_verdicts, score_task, score_tasks) __all__ = [ 'Check', 'CheckContext', 'CheckOutcome', 'CheckReport', - 'run_checks', 'checks_from_dicts', 'local_runner', + 'run_checks', 'checks_from_dicts', 'CRITERIA', 'DIMENSIONS', 'Criterion', 'RubricResult', 'build_rubric_prompt', 'parse_verdicts', 'score_task', 'score_tasks', ] diff --git a/src/twinkle_agentic/verifier/result_check.py b/src/twinkle_agentic/verifier/result_check.py index 1ea12e738..d619a8002 100644 --- a/src/twinkle_agentic/verifier/result_check.py +++ b/src/twinkle_agentic/verifier/result_check.py @@ -9,23 +9,20 @@ A task declares a list of :class:`Check`; :func:`run_checks` evaluates them and returns a :class:`CheckReport` whose ``score`` is the reward. -Checks that need to *run* something (``shell`` / ``python``) go through a -``runner`` so they execute wherever the episode ran -- pass the sandbox's -runner and the check sees exactly the state the agent left behind. Without one -they fall back to a local subprocess in ``workspace``, which is only correct -when the episode itself ran locally. +Checks that need to *run* something (``shell`` / ``python``) run inside the +episode's :class:`~twinkle_agentic.envs.base.Env`, so they see exactly the state +the agent left behind -- hand over the sandbox the episode acted in. Without one +they fall back to a :class:`~twinkle_agentic.envs.local.LocalEnv` over +``workspace``, which is only correct when the episode itself ran locally. """ import json import os import re -import resource -import shutil -import signal -import subprocess -import sys -import tempfile from dataclasses import dataclass, field -from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple + +if TYPE_CHECKING: # importing the env package for a type would cost every caller + from ..envs.base import Env # a second of import time -- see _local_env. __all__ = [ 'Check', @@ -34,15 +31,11 @@ 'CheckContext', 'run_checks', 'checks_from_dicts', - 'local_runner', ] -# (exit_code, output) for one command run inside the episode's workspace. -Runner = Callable[[str, str], Tuple[int, str]] - DEFAULT_TIMEOUT = int(os.environ.get('RESULT_CHECK_TIMEOUT', 60)) # Cap a runaway check so one bad task cannot take the trainer down with it. -_MEM_LIMIT_BYTES = 2 * 1024**3 +_MEM_LIMIT_GB = 2.0 _KINDS = ( 'file_exists', @@ -131,70 +124,29 @@ class CheckContext: Args: workspace: directory the episode wrote into. final_answer: text of the last assistant turn, for the ``answer_*`` kinds. - runner: executes a command in the episode's environment. ``None`` runs - it locally in ``workspace``. + env: where the ``shell`` / ``python`` kinds run -- the environment the + episode acted in. ``None`` runs them locally in ``workspace``. """ workspace: str = '' final_answer: str = '' - runner: Optional[Runner] = None + env: Optional['Env'] = None -def local_runner(workspace: str) -> Runner: - """Run commands in ``workspace`` as a local subprocess. +def _local_env(workspace: str) -> 'Env': + """Run checks in ``workspace`` on this machine. - Uses ``start_new_session`` + ``killpg`` so a forking command cannot leave - grandchildren behind on timeout, and caps address space at 2GB. + The fallback for a :class:`CheckContext` with no env. It is a + :class:`~twinkle_agentic.envs.local.LocalEnv`, so a check that falls back to + here and a check that runs in a sandbox go through one interface -- and the + process isolation (own session, killpg on timeout, capped address space) + lives in one place instead of being restated by every caller that needs it. """ - - def _run(command: str, interpreter: str) -> Tuple[int, str]: - return _local_exec(command, interpreter, workspace, DEFAULT_TIMEOUT) - - return _run - - -def _local_exec(source: str, interpreter: str, cwd: str, timeout: int) -> Tuple[int, str]: - cwd = cwd or '.' - os.makedirs(cwd, exist_ok=True) - if interpreter == 'python': - tmp = tempfile.mkdtemp(prefix='rescheck_') - script = os.path.join(tmp, '_check.py') - with open(script, 'w', encoding='utf-8') as f: - f.write(source) - argv = [sys.executable, script] - else: - tmp = None - argv = ['/bin/bash', '-lc', source] - - env = dict(os.environ, MPLBACKEND='Agg', PYTHONHASHSEED='0', - OMP_NUM_THREADS='1', MKL_NUM_THREADS='1', - TOKENIZERS_PARALLELISM='false') - env.pop('CUDA_VISIBLE_DEVICES', None) - - def _limit(): - resource.setrlimit(resource.RLIMIT_AS, (_MEM_LIMIT_BYTES, _MEM_LIMIT_BYTES)) - - try: - proc = subprocess.Popen(argv, cwd=cwd, env=env, stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, text=True, errors='replace', - start_new_session=True, preexec_fn=_limit) - try: - out, _ = proc.communicate(timeout=timeout) - return proc.returncode, out or '' - except subprocess.TimeoutExpired: - try: - os.killpg(proc.pid, signal.SIGKILL) - except ProcessLookupError: - pass - try: - proc.communicate(timeout=5) - except Exception: # noqa - pass - return 124, f'check did not finish within {timeout}s' - except Exception as e: # noqa - return 1, f'{type(e).__name__}: {e}' - finally: - if tmp: - shutil.rmtree(tmp, ignore_errors=True) + # Imported here, not at module scope: the env package pulls in twinkle's + # remote-class machinery, and a task declaring only file_* checks should not + # pay a second of import time for an environment it never runs anything in. + from ..envs.local import LocalEnv + return LocalEnv(workspace=workspace or '.', command_timeout=DEFAULT_TIMEOUT, + memory_limit_gb=_MEM_LIMIT_GB) def checks_from_dicts(raw: Sequence[Dict[str, Any]]) -> List[Check]: @@ -297,11 +249,12 @@ def _eval_one(check: Check, ctx: CheckContext) -> CheckOutcome: f'{check.path}:{check.key} is {got!r}, expected {check.value!r}') if kind in ('shell', 'python'): - runner = ctx.runner or local_runner(ctx.workspace) + env = ctx.env or _local_env(ctx.workspace) try: - code, out = runner(check.code, 'python' if kind == 'python' else 'shell') + code, out = env.run_script(check.code, kind, check.timeout) except Exception as e: # noqa - return CheckOutcome(check, False, f'runner raised {type(e).__name__}: {e}') + return CheckOutcome(check, False, + f'{type(env).__name__} raised {type(e).__name__}: {e}') if code != check.expect_exit: return CheckOutcome(check, False, f'exit {code} (expected {check.expect_exit}); output: {out[-300:]}') @@ -336,7 +289,7 @@ def run_checks( Args: checks: the task's assertions. An empty list scores 0.0 rather than a free 1.0, so a task that forgot to declare checks cannot look solved. - ctx: workspace / final answer / runner. + ctx: workspace / final answer / environment. mode: ``fraction`` gives weighted partial credit, ``all_or_nothing`` gives 1.0 only when every check passes. diff --git a/tests/preprocessor/test_preprocessor_utils.py b/tests/preprocessor/test_logprob_utils.py similarity index 97% rename from tests/preprocessor/test_preprocessor_utils.py rename to tests/preprocessor/test_logprob_utils.py index d52f8a77f..29bb18e36 100644 --- a/tests/preprocessor/test_preprocessor_utils.py +++ b/tests/preprocessor/test_logprob_utils.py @@ -1,5 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Tests for preprocessor.utils — pure logprob math helpers. +"""Tests for preprocessor.logprob_utils — pure logprob math helpers. These helpers compute conditional-vs-unconditional logprob deltas for IFD-family scoring (CherryLLM, T-SHIRT, ChR). All functions are stateless @@ -14,9 +14,9 @@ import math import pytest -from twinkle_agentic.preprocessor.utils import (_chr_min_distinct, _chr_min_weighted, _extract_logprob, - _ifd_family_metrics, _lp_to_jsonable, _mean_logprob_delta, _pad_batch, - _to_int_list) +from twinkle_agentic.preprocessor.logprob_utils import (_chr_min_distinct, _chr_min_weighted, _extract_logprob, + _ifd_family_metrics, _lp_to_jsonable, _mean_logprob_delta, + _pad_batch, _to_int_list) # ── _extract_logprob ──────────────────────────────────────────────────────── diff --git a/tests/twinkle_agentic/test_agentic_rsi.py b/tests/twinkle_agentic/test_agentic_rsi.py index e21652dfe..37e96d584 100644 --- a/tests/twinkle_agentic/test_agentic_rsi.py +++ b/tests/twinkle_agentic/test_agentic_rsi.py @@ -24,11 +24,18 @@ _COOKBOOK = os.path.join(_REPO, 'cookbook', 'rsi', 'agentic') sys.path.insert(0, _COOKBOOK) sys.path.insert(0, os.path.join(_COOKBOOK, 'sandbox_server')) +# recorder.py sits one level up, shared with the code half. +sys.path.insert(0, os.path.dirname(_COOKBOOK)) +# The code half itself, appended rather than inserted: both halves have a +# challenge.py, and the one these tests mean by that name is the agentic one. +sys.path.append(os.path.join(os.path.dirname(_COOKBOOK), 'code')) from remote_tool_env import RemoteMsAgentToolEnv # noqa: E402 from tool_server import (ToolRuntime, _usable_llm, # noqa: E402 _without_internal_args, _without_llm_args) +from twinkle_agentic.envs.base import Env, StepResult # noqa: E402 from twinkle_agentic.envs.env_tool import EnvTool # noqa: E402 +from twinkle_agentic.envs.local import LocalEnv # noqa: E402 from twinkle_agentic.tools.tool_manager import ToolManager # noqa: E402 from twinkle_agentic.verifier.result_check import (Check, CheckContext, # noqa: E402 checks_from_dicts, run_checks) @@ -128,6 +135,121 @@ def make_env(responder=None, tools=None, **kwargs): return env +class EnvJournal: + """What each slot's environment was asked to do, in order, across threads. + + One journal shared by a challenger's whole rack of environments. What the + slot tests are about is the correspondence between slots -- the env an + episode's check ran in has to be the env its tool calls were dispatched into + -- and that is only readable if every env and every manager writes to the + same place. + """ + + def __init__(self): + self._lock = threading.Lock() + self.events = [] # (kind, slot, payload) per operation, in order + + def add(self, kind, slot, payload=None): + with self._lock: + self.events.append((kind, slot, payload)) + + def kinds(self, *wanted): + """The sequence of operations, keeping only ``wanted``.""" + return [kind for kind, _, _ in self.events if kind in wanted] + + def slots(self, kind): + """Which slot each ``kind`` operation reached, in order.""" + return [slot for kind_, slot, _ in self.events if kind_ == kind] + + def payloads(self, kind, slot=None): + """What each ``kind`` operation carried, for one slot or all of them.""" + return [payload for kind_, slot_, payload in self.events + if kind_ == kind and slot in (None, slot_)] + + +class FakeToolManager: + """The dispatcher a :class:`FakeEnv` hands out, tagged with its slot.""" + + def __init__(self, slot, journal): + self.slot = slot + self.journal = journal + + def tool_infos(self): + return [] + + def __call__(self, tool_call): + self.journal.add('tool', self.slot) + return 'ok' + + +class FakeEnv(Env): + """A workspace a challenger can drive without a sandbox. + + A challenger reaches its workspace through four Env operations -- wipe it, + run a script in it, read the listing back, dispatch a tool call -- and this + implements those over a listing the test dictates and a queue of exit codes + it hands out. Shared by every challenger test below rather than a fresh set + of callbacks per test class: what they pin down is that the challenger drives + one env per slot correctly, which only means something while they all agree + on what an env is. + """ + + def __init__(self, listing='a.txt 1\n', *, slot=0, exit_code=0, exit_codes=None, + error='AssertionError', journal=None, with_tools=False): + """ + Args: + listing: what :meth:`snapshot` reports the workspace holds. + slot: which slot of the rack this is; recorded on every operation. + exit_code: what a script exits with once ``exit_codes`` runs out. + exit_codes: one exit code per script, consumed in order. A check that + fails and a rewrite that passes is the case this exists for. + error: the output a non-zero script comes back with. + journal: shared record; a private one when not given. + with_tools: advertise tools and hand out a :class:`FakeToolManager`. + Off by default -- an env that only runs scripts has none, and the + challenger is expected to leave the model without any. + """ + self.listing = listing + self.slot = slot + self.exit_code = exit_code + self.error = error + self.journal = journal if journal is not None else EnvJournal() + self.manager = FakeToolManager(slot, self.journal) if with_tools else None + self._exits = list(exit_codes) if exit_codes is not None else [] + self._lock = threading.Lock() + + # -- the operations a challenger performs on its workspace --------------- + + def clear(self): + self.journal.add('clear', self.slot) + + def snapshot(self): + self.journal.add('snapshot', self.slot) + return self.listing, '' + + def run_script(self, source, interpreter='python', timeout=None): + with self._lock: + code = self._exits.pop(0) if self._exits else self.exit_code + self.journal.add('run', self.slot, source) + return code, (self.error if code else '') + + def tools(self): + return DEFAULT_TOOLS if self.manager is not None else [] + + def tool_manager(self, schemas=None): + return self.manager + + def step(self, tool_name, arguments): + return StepResult(observation=f'ran {tool_name}') + + # -- what a test reads back ---------------------------------------------- + + @property + def scripts(self): + """Every script this env was asked to run, in order.""" + return self.journal.payloads('run', self.slot) + + class ResultCheckFileTest(unittest.TestCase): def setUp(self): @@ -240,7 +362,7 @@ def test_observation_is_truncated(self): env = make_env(responder=lambda call: 'x' * 50, max_observation_chars=10) obs = env.step('grep', {}).observation self.assertTrue(obs.startswith('x' * 10)) - self.assertIn('truncated 40 chars', obs) + self.assertIn('40 chars omitted', obs) def test_unreachable_runtime_becomes_an_observation(self): # A dead sandbox must not take down the training step: the episode plays @@ -280,7 +402,7 @@ def test_resolve_tool_raises_on_unknown_name(self): with self.assertRaises(ValueError): self.env.resolve_tool('no_such_tool') - def test_runner_recovers_exit_code_from_text_output(self): + def test_run_script_recovers_exit_code_from_text_output(self): # The sandbox tools return prose; the marker is how the exit status # survives. Emulate a shell that echoes the marker. Matching on the # namespaced name also proves the plain name was resolved. @@ -289,19 +411,21 @@ def responder(call): return 'some output\n__TWINKLE_RC__:0' return '__TWINKLE_RC__:3' - runner = make_env(responder).runner() - self.assertEqual(runner('ls', 'shell'), (0, 'some output')) - self.assertEqual(runner('boom()', 'python')[0], 3) + env = make_env(responder) + self.assertEqual(env.run_script('ls', 'shell'), (0, 'some output')) + self.assertEqual(env.run_script('boom()')[0], 3) - def test_runner_missing_marker_is_a_failure_not_a_pass(self): - code, out = make_env(lambda c: 'sandbox died').runner()('ls', 'shell') + def test_run_script_missing_marker_is_a_failure_not_a_pass(self): + code, out = make_env(lambda c: 'sandbox died').run_script('ls', 'shell') self.assertNotEqual(code, 0) self.assertIn('sandbox died', out) - def test_checks_run_through_the_env_runner(self): + def test_checks_run_through_the_env(self): + # The env is the one thing a check needs to reach the episode's own + # filesystem, so it is passed as itself rather than as a callable. env = make_env(lambda c: '__TWINKLE_RC__:0') report = run_checks([Check(kind='shell', code='true')], - CheckContext(workspace=self.tmp, runner=env.runner())) + CheckContext(workspace=self.tmp, env=env)) self.assertTrue(report.all_passed) def test_download_workspace_brings_files_back_for_file_checks(self): @@ -500,29 +624,18 @@ class EmptyWorkspaceTest(unittest.TestCase): def _challenger(self, snapshot): from twinkle_agentic.challenger.agentic import AgenticChallenger, AgenticPrompts - self.checks_run = [] - def explorer(trajectories, **kwargs): return [{'messages': list(t['messages']) + [{'role': 'assistant', 'content': '```python\nassert True\n```'}]} for t in trajectories] - def run_check_fn(script, slot=0): - self.checks_run.append(script) - return 0, '' - + self.env = FakeEnv(snapshot) prompts = AgenticPrompts( system='s', from_scratch='u', check_followup='write checks for {final_state}', check_retry_followup='{error} / {final_state}', problem_followup='write the statement') - return AgenticChallenger( - prompts, explorer, - reset_fn=lambda slot=0: None, - run_check_fn=run_check_fn, - workspace_snapshot_fn=lambda slot=0: snapshot, - solver_rollouts=0, - ) + return AgenticChallenger(prompts, explorer, envs=[self.env], solver_rollouts=0) def _explored(self): return {'messages': [{'role': 'user', 'content': 'do something'}, @@ -536,7 +649,7 @@ def test_empty_snapshot_ends_the_episode_before_any_check_is_written(self): self.assertEqual(ch.stats['empty_workspace'], 1) self.assertEqual(state['reject'][0], 'empty_workspace') # No check script was even run: there was nothing to check. - self.assertEqual(self.checks_run, []) + self.assertEqual(self.env.scripts, []) # And the episode is not turned into a task afterwards. self.assertIsNone(ch._finish_episode(state, self._explored())) @@ -555,13 +668,13 @@ def test_a_real_snapshot_asks_for_checks_and_then_runs_them(self): # checks are written against. self.assertIn('--- data.csv ---', text) self.assertIsNone(params) - self.assertEqual(self.checks_run, []) + self.assertEqual(self.env.scripts, []) wrote_script = {'messages': [ {'role': 'assistant', 'content': '```python\nassert True\n```'}]} self.assertEqual(ch._followup(state, wrote_script, 1), ('write the statement', None)) - self.assertEqual(self.checks_run, ['assert True']) + self.assertEqual(self.env.scripts, ['assert True']) class ProblemStatementParseTest(unittest.TestCase): @@ -616,13 +729,6 @@ def _challenger(self, replies, snapshot='a.txt 1\n\n--- a.txt ---\nx', self.emitted = [] self.rejected = [] self.appended = [] - # One exit code per check run, so a test can make the first fail and the - # rewrite pass. - exits = list(check_exits) if check_exits is not None else None - - def run_check(script, slot=0): - code = exits.pop(0) if exits else check_exit - return (code, 'AssertionError' if code else '') def explorer(trajectories, **kw): followup_fn = kw.get('followup_fn') @@ -646,9 +752,9 @@ def explorer(trajectories, **kw): problem_followup='statement please') return AgenticChallenger( prompts, explorer, - reset_fn=lambda slot=0: None, - run_check_fn=run_check, - workspace_snapshot_fn=lambda slot=0: snapshot, + # One exit code per check run, so a test can make the first fail and + # the rewrite pass. + envs=[FakeEnv(snapshot, exit_code=check_exit, exit_codes=check_exits)], reject_sink=self.rejected.append, propose_sink=self.emitted.append, solver_rollouts=0, @@ -753,9 +859,7 @@ def explorer(trajectories, **kw): check_retry_followup='{error} / {final_state}', problem_followup='p'), explorer, - reset_fn=lambda slot=0: None, - run_check_fn=lambda script, slot=0: (0, ''), - workspace_snapshot_fn=lambda slot=0: 'a.txt 1\n', + envs=[FakeEnv()], reject_sink=self.rejected.append, propose_sink=self.emitted.append, solver_rollouts=0) @@ -768,51 +872,25 @@ def explorer(trajectories, **kw): class ConcurrentEpisodeSlotsTest(unittest.TestCase): - """Concurrent episodes must each drive their own sandbox slot. + """Concurrent episodes must each drive their own environment. A rack of one sandbox per slot is the whole point of running episodes in - parallel; if the slot the challenger passes for episode i is not the slot - reset_fn / run_check_fn / workspace_snapshot_fn / tool_manager see for that - episode, then two episodes end up sharing a workspace and the check written - against one runs against the other. That is the failure mode this test - exists to catch. + parallel; if the env the challenger hands episode i is not the env that + episode's clear, check, snapshot and tool calls land in, then two episodes + end up sharing a workspace and the check written against one runs against the + other. That is the failure mode this test exists to catch. """ def test_each_episode_uses_its_own_slot_end_to_end(self): from twinkle_agentic.challenger.agentic import AgenticChallenger, AgenticPrompts n_slots = 4 - resets = [] # (slot,) per call - checks = [] # (slot, script) per call - snaps = [] # (slot,) per call - tm_calls = [] # (slot,) per tool_manager use - lock = threading.Lock() - - class FakeTM: - def __init__(self, slot): self.slot = slot - def tool_infos(self): return [] - def __call__(self, tc): - with lock: - tm_calls.append(self.slot) - return 'ok' - - tool_managers = [FakeTM(i) for i in range(n_slots)] - - def reset_fn(slot): - with lock: - resets.append(slot) - - def run_check_fn(script, slot): - with lock: - checks.append((slot, script)) - return 0, '' - - def workspace_snapshot_fn(slot): - with lock: - snaps.append(slot) - # Encode the slot in the snapshot so an episode reading the wrong - # slot's workspace would produce a mismatched check statement. - return f'slot_{slot}.txt 1\n\n--- slot_{slot}.txt ---\nx' + journal = EnvJournal() + # The slot is encoded in the listing too, so an episode reading the wrong + # slot's workspace would write its checks against another one's files. + envs = [FakeEnv(f'slot_{i}.txt 1\n\n--- slot_{i}.txt ---\nx', + slot=i, journal=journal, with_tools=True) + for i in range(n_slots)] def explorer(trajectories, **kw): # The two follow-ups (check script, then statement) are threaded @@ -845,11 +923,7 @@ def explorer(trajectories, **kw): problem_followup='p') ch = AgenticChallenger( prompts, explorer, - reset_fn=reset_fn, - run_check_fn=run_check_fn, - workspace_snapshot_fn=workspace_snapshot_fn, - episode_concurrency=n_slots, - episode_tool_managers=tool_managers, + envs=envs, propose_sink=emitted.append, solver_rollouts=0, max_proposals_per_round=8, @@ -857,29 +931,28 @@ def explorer(trajectories, **kw): kept = ch._round(8) self.assertEqual(len(kept), 8) - # 8 episodes across 4 slots, evenly split -> each slot reset twice, ran + # 8 episodes across 4 slots, evenly split -> each slot cleared twice, ran # its own check twice, and every check saw the slot's own snapshot text. from collections import Counter - self.assertEqual(Counter(resets), Counter({0: 2, 1: 2, 2: 2, 3: 2})) - self.assertEqual(Counter(s for s, _ in checks), Counter({0: 2, 1: 2, 2: 2, 3: 2})) + self.assertEqual(Counter(journal.slots('clear')), Counter({0: 2, 1: 2, 2: 2, 3: 2})) + self.assertEqual(Counter(journal.slots('run')), Counter({0: 2, 1: 2, 2: 2, 3: 2})) # The tool_manager slot used matches the check slot for each episode. - self.assertEqual(Counter(tm_calls), Counter({0: 2, 1: 2, 2: 2, 3: 2})) + self.assertEqual(Counter(journal.slots('tool')), Counter({0: 2, 1: 2, 2: 2, 3: 2})) + + def test_a_challenger_without_an_env_is_refused(self): + """There is no episode without a workspace, and no check without one either. - def test_wrong_tool_manager_count_is_refused(self): + Refused at construction rather than at the first episode: the failure is a + missing argument in the wiring, and finding out about it a round into a run + costs the round. + """ from twinkle_agentic.challenger.agentic import AgenticChallenger, AgenticPrompts prompts = AgenticPrompts(system='s', from_scratch='u', check_followup='c {final_state}', check_retry_followup='{error} / {final_state}', problem_followup='p') with self.assertRaises(ValueError): - AgenticChallenger( - prompts, explorer=lambda t, **k: t, - reset_fn=lambda slot=0: None, - run_check_fn=lambda s, slot=0: (0, ''), - workspace_snapshot_fn=lambda slot=0: '', - episode_concurrency=4, - episode_tool_managers=[object(), object()], # wrong count - solver_rollouts=0) + AgenticChallenger(prompts, explorer=lambda t, **k: t, solver_rollouts=0) class PreseedInputsTest(unittest.TestCase): @@ -891,7 +964,7 @@ class PreseedInputsTest(unittest.TestCase): data its statement says is there -- which reads as 'too hard' and is not. """ - def _challenger(self, run_check_fn, **kwargs): + def _challenger(self, env, **kwargs): from twinkle_agentic.challenger.agentic import AgenticChallenger, AgenticPrompts prompts = AgenticPrompts( @@ -903,11 +976,9 @@ def _challenger(self, run_check_fn, **kwargs): prompts, lambda trajs, **kw: [{'messages': list(t['messages']), 'stop_reason': 'stop'} for t in trajs], - run_check_fn=run_check_fn, - workspace_snapshot_fn=lambda slot=0: 'input/a.csv 3\n', + envs=[env], solver_rollouts=2, - keep_min_pass=1, - keep_max_pass_margin=0, + keep_pass_band=(1, 2), propose_sink=[].append, **kwargs) @@ -918,29 +989,30 @@ def _task(self, setup): keywords=[]) def test_setup_runs_after_the_clear_and_before_the_check(self): - events = [] - ch = self._challenger( - run_check_fn=lambda script, slot=0: (events.append( - 'setup' if script.startswith('#SETUP') else 'check'), (0, ''))[1], - reset_fn=lambda slot=0: events.append('clear')) + env = FakeEnv('input/a.csv 3\n') + ch = self._challenger(env) kept = ch._filter_difficulty([self._task('#SETUP\nopen("a","w")')]) - self.assertEqual(events, ['clear', 'setup', 'check'] * 2) + # Per attempt: wipe the workspace, replay the inputs, then check. + self.assertEqual(env.journal.kinds('clear', 'run'), ['clear', 'run', 'run'] * 2) + self.assertEqual([s.startswith('#SETUP') for s in env.scripts], [True, False] * 2) self.assertEqual(len(kept), 1) def test_failed_setup_skips_the_attempt_instead_of_scoring_it_zero(self): """An attempt that never ran must not be counted as an attempt that failed.""" - checks = [] - def run_check(script, slot=0): - if script.startswith('#SETUP'): - return 1, 'no space left on device' - checks.append(script) - return 0, '' + class NoSpaceEnv(FakeEnv): + """A workspace where the replay fails and a check would have passed.""" - ch = self._challenger(run_check_fn=run_check, reset_fn=lambda slot=0: None) + def run_script(self, source, interpreter='python', timeout=None): + if source.startswith('#SETUP'): + return 1, 'no space left on device' + return super().run_script(source, interpreter, timeout) + + env = NoSpaceEnv('input/a.csv 3\n') + ch = self._challenger(env) kept = ch._filter_difficulty([self._task('#SETUP\nboom')]) # The solver was never asked, so nothing was checked and nothing is kept. - self.assertEqual(checks, []) + self.assertEqual(env.scripts, []) self.assertEqual(kept, []) self.assertEqual(ch.stats['setup_replay_fail'], 2) @@ -963,15 +1035,8 @@ def test_attempts_are_batched_per_wave_and_stay_in_their_slot(self): n_slots = 4 lock = threading.Lock() batch_sizes = [] # trajectories per explorer call - reset_slots = [] # slot per reset - pairs = [] # (tool_manager slot, check slot) per attempt - - class FakeTM: - def __init__(self, slot): self.slot = slot - def tool_infos(self): return [] - def __call__(self, tc): return 'ok' - - tool_managers = [FakeTM(i) for i in range(n_slots)] + journal = EnvJournal() + envs = [FakeEnv(slot=i, journal=journal, with_tools=True) for i in range(n_slots)] # Which slot's manager each trajectory of the current wave was handed. wave_slots = [] @@ -984,20 +1049,6 @@ def explorer(trajectories, **kw): return [{'messages': list(t['messages']), 'stop_reason': 'stop'} for t in trajectories] - # One attempt in flight per slot, so the check for the attempt that used - # slot k must itself run in slot k. Recorded as a pair to compare. - seq = iter(range(10_000)) - - def run_check_fn(script, slot=0): - with lock: - pairs.append(slot) - next(seq) - return 0, '' - - def reset_fn(slot=0): - with lock: - reset_slots.append(slot) - prompts = AgenticPrompts( system='s', from_scratch='u', check_followup='c {final_state}', @@ -1005,14 +1056,9 @@ def reset_fn(slot=0): problem_followup='p') ch = AgenticChallenger( prompts, explorer, - reset_fn=reset_fn, - run_check_fn=run_check_fn, - workspace_snapshot_fn=lambda slot=0: 'a.txt 1\n', - episode_concurrency=n_slots, - episode_tool_managers=tool_managers, + envs=envs, solver_rollouts=4, - keep_min_pass=1, - keep_max_pass_margin=0, + keep_pass_band=(1, 4), propose_sink=[].append, ) tasks = [attach_user_data({'messages': [{'role': 'user', 'content': f'task {i}'}]}, @@ -1025,10 +1071,10 @@ def reset_fn(slot=0): self.assertEqual(batch_sizes, [4, 4]) # Every slot cleared once per wave, and the managers handed out are the # slots that were cleared. - self.assertEqual(sorted(reset_slots), [0, 0, 1, 1, 2, 2, 3, 3]) + self.assertEqual(sorted(journal.slots('clear')), [0, 0, 1, 1, 2, 2, 3, 3]) self.assertEqual(sorted(wave_slots), [0, 1, 2, 3]) # One check per attempt, one per slot per wave. - self.assertEqual(sorted(pairs), [0, 0, 1, 1, 2, 2, 3, 3]) + self.assertEqual(sorted(journal.slots('run')), [0, 0, 1, 1, 2, 2, 3, 3]) # All checks passed -> both tasks scored 4 of 4. self.assertEqual([user_data_get(t.get('user_data'), 'n_pass', -1) for t in kept], [4, 4]) @@ -1047,29 +1093,26 @@ def _challenger(self, attempt_flags, **kwargs): """``attempt_flags``: one (truncated, passes) pair per solver attempt.""" from twinkle_agentic.challenger.agentic import AgenticChallenger, AgenticPrompts - self.flags = list(attempt_flags) self.emitted = [] + # One attempt per wave with a single env, so the explorer and the env walk + # the flags in step: the truncation and the verdict below it belong to the + # same attempt, which is the whole point of the count being read together. + truncations = [truncated for truncated, _ in attempt_flags] def explorer(trajectories, **kw): - truncated, _ = self.flags[0] + truncated = truncations.pop(0) return [{'messages': list(t['messages']), 'truncated': truncated, 'stop_reason': 'length' if truncated else 'stop'} for t in trajectories] - def run_check_fn(script, slot=0): - _, passes = self.flags.pop(0) - return (0 if passes else 1), '' - prompts = AgenticPrompts( system='s', from_scratch='u', check_followup='cs {final_state}', check_retry_followup='{error} / {final_state}', problem_followup='ps') return AgenticChallenger( prompts, explorer, - reset_fn=lambda slot=0: None, - run_check_fn=run_check_fn, - workspace_snapshot_fn=lambda slot=0: 'a.txt 1\n', + envs=[FakeEnv(exit_codes=[0 if passes else 1 for _, passes in attempt_flags])], propose_sink=self.emitted.append, solver_rollouts=4, **kwargs) @@ -1085,7 +1128,7 @@ def test_truncated_attempts_are_counted_and_still_scored_as_failures(self): # truncations visible in stats so the 1-of-4 can be read for what it is. ch = self._challenger([(False, True), (False, False), (True, False), (True, False)], - keep_min_pass=1, keep_max_pass_margin=1) + keep_pass_band=(1, 3)) kept = ch._filter_difficulty([self._task()]) self.assertEqual(ch.stats['solver_truncated'], 2) @@ -1099,7 +1142,7 @@ def test_all_four_truncated_reads_as_nobody_solved_it(self): # discarded for being too hard and stats['solver_truncated'] == 4 is the # only thing that says no solver ever acted. ch = self._challenger([(True, False)] * 4, - keep_min_pass=1, keep_max_pass_margin=1) + keep_pass_band=(1, 3)) kept = ch._filter_difficulty([self._task()]) self.assertEqual(ch.stats['solver_truncated'], 4) @@ -1146,11 +1189,9 @@ def text_explorer(trajectories, **kwargs): keyword_expand_user=KEYWORD_EXPAND_USER) return AgenticChallenger( prompts, tool_explorer, + envs=[FakeEnv()], keyword_store=self.store, category_desc={'filesystem': 'files and directories'}, - reset_fn=lambda slot=0: None, - run_check_fn=lambda script, slot=0: (0, ''), - workspace_snapshot_fn=lambda slot=0: 'a.txt 1\n', keyword_explorer=text_explorer, keyword_sink=self.gen_records.append, keyword_gen_calls=1, @@ -1165,7 +1206,7 @@ def tearDown(self): def test_the_shipped_prompt_asks_for_what_the_parser_reads(self): """The real prompt string, not a stand-in: this is the contract that broke.""" - from twinkle_agentic.challenger.code import parse_keyword_list + from twinkle_agentic.challenger.keywords import parse_keyword_list from prompts import KEYWORD_EXPAND_USER, KEYWORD_USER for text in (KEYWORD_USER, KEYWORD_EXPAND_USER): @@ -1190,8 +1231,8 @@ def test_over_length_keywords_are_reported_and_not_merely_gone(self): The direction matters as much as the count: length tracks specificity, so what the filter removes is the half of the output the bank most wants. """ - from twinkle_agentic.challenger.code import (KEYWORD_MAX_LEN, - split_keyword_list) + from twinkle_agentic.challenger.keywords import (KEYWORD_MAX_LEN, + split_keyword_list) from prompts import KEYWORD_EXPAND_USER # Verbatim from iteration 9, one of the eight a single expand call lost. @@ -1291,11 +1332,9 @@ def explorer(trajectories, **kwargs): self.store = KeywordStore(os.path.join(self.tmp, 'kw.jsonl'), (category,)) return AgenticChallenger( prompts, explorer, + envs=[FakeEnv()], keyword_store=self.store, category_desc={category: 'some kind of work'}, - reset_fn=lambda slot=0: None, - run_check_fn=lambda script, slot=0: (0, ''), - workspace_snapshot_fn=lambda slot=0: 'a.txt 1\n', keyword_explorer=explorer, keyword_gen_calls=3, keyword_refill_concurrency=refill_concurrency, @@ -1316,7 +1355,7 @@ def _three_axis_prompts(self): def test_three_axis_refill_shows_each_call_the_previous_output(self): ch = self._challenger(self._three_axis_prompts(), 'transform') - got = ch._generate_keywords('transform', 9) + got = ch.keywords._generate('transform', 9) self.assertEqual(sorted(got), ['topic 0', 'topic 1', 'topic 2']) self.assertEqual(self.batch_sizes, [1, 1, 1], @@ -1337,7 +1376,7 @@ def test_raising_the_concurrency_restores_the_batched_behaviour(self): ch = self._challenger(self._three_axis_prompts(), 'transform', refill_concurrency=3) - got = ch._generate_keywords('transform', 9) + got = ch.keywords._generate('transform', 9) self.assertEqual(sorted(got), ['topic 0', 'topic 1', 'topic 2']) self.assertEqual(self.batch_sizes, [3], 'all three go out as one batch') @@ -1364,10 +1403,10 @@ def test_the_avoid_list_is_capped_and_drops_older_entries_first(self): keyword_system=KEYWORD_SYSTEM, keyword_user=KEYWORD_USER, keyword_expand_user=KEYWORD_EXPAND_USER) ch = self._challenger(prompts, 'transform') - cap = ch._AVOID_TOTAL + cap = ch.keywords._AVOID_TOTAL older = [f'old {i}' for i in range(200)] fresh = [f'new {i}' for i in range(5)] - note = ch._avoid_note(older, fresh, 'avoid: ') + note = ch.keywords._avoid_note(older, fresh) for kw in fresh: self.assertIn(kw, note) self.assertEqual(note.count('old '), cap - len(fresh)) @@ -1375,7 +1414,7 @@ def test_the_avoid_list_is_capped_and_drops_older_entries_first(self): # Once this refill alone fills the cap, no banked phrase is quoted and the # line stops growing -- it is the growth that broke the eighth call. many = [f'new {i}' for i in range(cap + 30)] - note = ch._avoid_note(older, many, 'avoid: ') + note = ch.keywords._avoid_note(older, many) self.assertEqual(note.count('old '), 0) self.assertEqual(note.count('new '), cap) self.assertNotIn('new 0', note, 'the oldest of this refill falls off first') @@ -1404,7 +1443,7 @@ class ProposeTrajIndexTest(unittest.TestCase): """ def test_group_id_and_reward_survive_the_copy(self): - from challenge import Recorder + from recorder import Recorder out = tempfile.mkdtemp(prefix='proposetraj_test_') try: @@ -1451,9 +1490,7 @@ def _challenger(self): return AgenticChallenger( prompts, lambda trajectories, **kwargs: list(trajectories), - reset_fn=lambda slot=0: None, - run_check_fn=lambda script, slot=0: (0, ''), - workspace_snapshot_fn=lambda slot=0: 'data.csv 3', + envs=[FakeEnv('data.csv 3')], solver_rollouts=0, ) @@ -1477,5 +1514,195 @@ def test_group_id_reaches_the_built_task(self): self.assertEqual(user_data_get(task.get('user_data'), 'group_id', None), 7) +# ── the code half ────────────────────────────────────────────────────────── + +# Two problems the local runner can actually verify, so the difficulty stage +# here is the production one: build_asserts runs the solution to capture each +# check's repr, and every judgement is a real subprocess. +_SOLVE_MARK = 'SOLVE:' +_CODE_PROBLEMS = ( + {'problem': 'Double an integer.', + 'solution': 'def double(x):\n return x * 2\n', + 'wrong': 'def double(x):\n return x\n', + 'entry': 'double', + 'checks': ['double(2)', 'double(5)']}, + {'problem': 'Sum a list of integers.', + 'solution': 'def total(xs):\n return sum(xs)\n', + 'wrong': 'def total(xs):\n return 0\n', + 'entry': 'total', + 'checks': ['total([1, 2, 3])', 'total([])']}, +) + + +def _explored(traj, text, n_prompt=3, n_new=4): + """What a local sampler returns: the reply, and the tokens behind it. + + The token fields matter as much as the text. train.py refuses a trajectory + whose logprob count disagrees with its trainable label count, so a fake that + got the counts wrong would pass the collection tests and be dropped by the + step -- which is the failure this half was built to make impossible. + """ + ids = list(range(1, n_prompt + n_new + 1)) + return { + 'messages': list(traj['messages']) + [{'role': 'assistant', 'content': text}], + 'input_ids': ids, + 'labels': [-100] * n_prompt + ids[n_prompt:], + # Top-1 pairs, the shape SampledSequence.logprobs uses. + 'logprobs': [[(i, -0.5)] for i in ids[n_prompt:]], + } + + +class _ScriptedCodeExplorer: + """Answers a code challenger's prompts from a fixed script. + + Proposals are answered in order from ``problems``; solver prompts are matched + back to their problem by the statement they quote and answered from + ``verdicts[i]``, one boolean per attempt. Stating the pass counts is the point: + the band is what decides whether a group has a gradient, so a test about it + cannot depend on which code a model would have happened to write. + """ + + def __init__(self, problems, verdicts): + self.problems = list(problems) + self.verdicts = [list(v) for v in verdicts] + self.by_statement = {p['problem']: i for i, p in enumerate(self.problems)} + self.n_proposed = 0 + self.n_attempted = [0] * len(self.problems) + + def __call__(self, trajectories, sampling_params=None, **kwargs): + return [self._reply(t) for t in trajectories] + + def _reply(self, traj): + user = next(m['content'] for m in reversed(traj['messages']) + if m.get('role') == 'user') + if not user.startswith(_SOLVE_MARK): + problem = self.problems[min(self.n_proposed, len(self.problems) - 1)] + self.n_proposed += 1 + return _explored(traj, json.dumps( + {k: problem[k] for k in ('problem', 'solution', 'entry', 'checks')})) + i = self.by_statement[user[len(_SOLVE_MARK):].strip()] + problem = self.problems[i] + passing = self.verdicts[i][self.n_attempted[i]] + self.n_attempted[i] += 1 + return _explored(traj, f'```python\n{problem["solution" if passing else "wrong"]}```') + + +class _CodeArgs: + """The two attributes ``collect`` reads off the parsed arguments.""" + + code_keep_target = 1 + code_batch_size = 0 + + +class CodeHalfCollectionTest(unittest.TestCase): + """The attempts the difficulty stage makes are the code half's training data. + + Measuring a candidate samples it ``solver_rollouts`` times and then reports one + number, and the base class drops the attempts. Those attempts are exactly what + a solver trains on, and a problem kept inside the band is a group already + measured to contain both a pass and a failure. Sampling a fresh group after the + band has been applied pays for the same tokens twice and can still land at 0 or + 8, where every advantage is the reward minus itself. + + So ``CollectingChallenger`` keeps them, and what these pin is the whole path: + the kept problem arrives with all of its attempts, a problem the band drops + does not go on holding its own, and what reaches index.jsonl loads back as one + code group with a gradient. + + The out-of-band problem is proposed first on purpose. It is measured in a round + of its own, so the round that keeps a problem is not the round that has to + forget one -- the two would otherwise pass together and fail together. + """ + + def _collect(self, out_dir, verdicts=((False, False, False, False), + (True, True, True, False))): + from collect import CollectingChallenger, collect + from recorder import Recorder + from twinkle_agentic.challenger.code import CodePrompts + + explorer = _ScriptedCodeExplorer(_CODE_PROBLEMS, verdicts) + prompts = CodePrompts(system='S', from_scratch='INVENT', solver_system='SS', + solver_user=_SOLVE_MARK + ' {problem}') + recorder = Recorder(out_dir) + seen = [] + challenger = CollectingChallenger( + prompts, explorer, envs=[LocalEnv()], solver_rollouts=4, + keep_pass_band=(1, 3), two_step=False, seed=1, attempt_sink=seen.append) + try: + metrics = collect(_CodeArgs(), challenger, recorder) + finally: + recorder.close() + return challenger, metrics, seen + + def test_a_kept_problem_arrives_as_one_group_of_every_attempt(self): + out = tempfile.mkdtemp(prefix='codecollect_test_') + try: + _ch, metrics, seen = self._collect(out) + with open(os.path.join(out, 'trajs', 'index.jsonl'), encoding='utf-8') as f: + records = [json.loads(line) for line in f if line.strip()] + finally: + shutil.rmtree(out, ignore_errors=True) + + self.assertEqual(metrics['counts']['kept'], 1) + self.assertEqual(metrics['counts']['groups'], 1) + # Four members from four rollouts: a group short of one attempt is a group + # whose advantage was computed against a mean it never had. + self.assertEqual(len(records), 4) + self.assertEqual({r['side'] for r in records}, {'code'}) + self.assertEqual({r['group_id'] for r in records}, {0}) + # Three passes and one failure, which is what n_pass=3 of 4 means. + self.assertEqual(sorted(r['reward'] for r in records), [0.0, 1.0, 1.0, 1.0]) + # Both problems' attempts reach the audit file, including the eight that + # measured a problem the band then dropped: that file exists for the + # question of why something measured zero. + self.assertEqual(len(seen), 8) + self.assertNotIn('attempt', seen[0], + 'the audit line carries the verdict, not the tokens') + + def test_a_problem_outside_the_band_does_not_keep_its_attempts(self): + out = tempfile.mkdtemp(prefix='codeband_test_') + try: + challenger, _metrics, _seen = self._collect(out) + finally: + shutil.rmtree(out, ignore_errors=True) + + # Empty, not 'holds one problem': the kept problem's attempts were taken by + # collect and the dropped problem's were released when its count came in. + # Four attempts of a 24-turn episode is a gigabyte a round, so this is the + # difference between a loop that runs and one that runs out of memory. + self.assertEqual(challenger._attempts, {}) + + def test_the_index_loads_back_as_a_code_group_with_a_gradient(self): + import train as T + + out = tempfile.mkdtemp(prefix='codeload_test_') + try: + self._collect(out) + groups, skipped = T.load(out, sides='both,code', max_length=1024) + notes = T.score(groups) + finally: + shutil.rmtree(out, ignore_errors=True) + + self.assertEqual(dict(skipped), {}, 'every attempt written should be trainable') + self.assertEqual(len(groups), 1) + # (side, group_id), not (side, group_id, proposal_idx): one problem is one + # group here, and keying on a proposal index that is always 0 would work by + # accident rather than by agreement with what collect writes. + self.assertEqual(groups[0]['key'], ('code', 0)) + self.assertEqual(groups[0]['side'], 'code') + self.assertEqual(dict(notes), {}, 'a group inside the band has to have a gradient') + advantages = [m['advantage'] for m in groups[0]['members']] + self.assertTrue(any(abs(a) > 1e-9 for a in advantages), advantages) + + def test_sides_wanted_reads_the_three_sides_out_of_one_switch(self): + import train as T + + self.assertEqual(T.sides_wanted('both'), ('propose', 'solve')) + self.assertEqual(T.sides_wanted('both,code'), ('propose', 'solve', 'code')) + self.assertEqual(T.sides_wanted('code'), ('code', )) + # Repeats collapse rather than doubling a side's share of the step. + self.assertEqual(T.sides_wanted('code, code ,solve'), ('code', 'solve')) + + if __name__ == '__main__': unittest.main() From 364a3bd65f6a14e0421f918ef7454635b9650354 Mon Sep 17 00:00:00 2001 From: root <yuze.zyz@alibaba-inc.com> Date: Thu, 10 Sep 2026 00:44:56 +0800 Subject: [PATCH 58/60] wip --- README.md | 2 +- README_ZH.md | 2 +- cookbook/rsi/code/challenge.py | 4 +- cookbook/rsi/code/collect.py | 8 +- docs/source_en/Components/Agentic/Envs.md | 4 +- .../Agentic/Multi-Turn-Tool-Usage.md | 19 +- docs/source_en/Components/Agentic/Rollout.md | 58 +- .../\347\273\204\344\273\266/Agentic/Envs.md" | 4 +- .../Agentic/Multi-Turn-Tool-Usage.md" | 19 +- .../Agentic/Rollout.md" | 58 +- src/twinkle/loss/grpo.py | 35 +- src/twinkle/model/megatron/megatron.py | 8 + .../strategy/sequence_parallel/__init__.py | 26 +- .../model/transformers/transformers.py | 8 + src/twinkle/processor/base.py | 39 +- .../sampler/vllm_sampler/vllm_sampler.py | 10 +- .../server/sampler/backends/mock_sampler.py | 11 +- src/twinkle/template/base.py | 154 +++- src/twinkle_agentic/async_rl/data_plane.py | 18 +- src/twinkle_agentic/challenger/base.py | 12 +- .../challenger/new/__init__.py | 12 + src/twinkle_agentic/challenger/new/agentic.py | 511 +++++++++++ src/twinkle_agentic/challenger/new/base.py | 132 +++ src/twinkle_agentic/challenger/new/keyword.py | 288 +++++- .../challenger/new/recorder.py | 91 ++ src/twinkle_agentic/harness/ms_agent.py | 16 +- src/twinkle_agentic/preprocessor/AUDIT.md | 179 ---- src/twinkle_agentic/preprocessor/__init__.py | 2 - .../preprocessor/experimental/__init__.py | 21 - .../preprocessor/experimental/llm_backend.py | 344 ------- .../preprocessor/experimental/score_filter.py | 835 ----------------- .../preprocessor/intent_classifier.py | 12 +- src/twinkle_agentic/preprocessor/intents.py | 25 - .../preprocessor/label_schema.py | 117 --- .../preprocessor/logprob_utils.py | 231 ----- .../preprocessor/offline/__init__.py | 24 - .../preprocessor/offline/decontaminate.py | 109 --- .../preprocessor/offline/near_dedup.py | 163 ---- .../preprocessor/provenance.py | 73 -- .../preprocessor/structural_noise.py | 60 -- src/twinkle_agentic/protocol/openai.py | 46 +- src/twinkle_agentic/rollout/__init__.py | 17 +- src/twinkle_agentic/rollout/api_multi_turn.py | 314 ------- src/twinkle_agentic/rollout/api_sampler.py | 133 +++ src/twinkle_agentic/rollout/base.py | 119 ++- src/twinkle_agentic/rollout/bridge.py | 165 ++-- src/twinkle_agentic/rollout/factory.py | 74 -- src/twinkle_agentic/rollout/multi_turn.py | 853 +++++++++--------- src/twinkle_agentic/summarizer/base.py | 63 +- src/twinkle_agentic/utils/code_utils.py | 27 +- tests/preprocessor/test_logprob_utils.py | 354 -------- .../test_multi_turn_rollout.py | 156 +++- 52 files changed, 2357 insertions(+), 3708 deletions(-) create mode 100644 src/twinkle_agentic/challenger/new/__init__.py create mode 100644 src/twinkle_agentic/challenger/new/agentic.py create mode 100644 src/twinkle_agentic/challenger/new/base.py create mode 100644 src/twinkle_agentic/challenger/new/recorder.py delete mode 100644 src/twinkle_agentic/preprocessor/AUDIT.md delete mode 100644 src/twinkle_agentic/preprocessor/experimental/__init__.py delete mode 100644 src/twinkle_agentic/preprocessor/experimental/llm_backend.py delete mode 100644 src/twinkle_agentic/preprocessor/experimental/score_filter.py delete mode 100644 src/twinkle_agentic/preprocessor/intents.py delete mode 100644 src/twinkle_agentic/preprocessor/label_schema.py delete mode 100644 src/twinkle_agentic/preprocessor/logprob_utils.py delete mode 100644 src/twinkle_agentic/preprocessor/offline/__init__.py delete mode 100644 src/twinkle_agentic/preprocessor/offline/decontaminate.py delete mode 100644 src/twinkle_agentic/preprocessor/offline/near_dedup.py delete mode 100644 src/twinkle_agentic/preprocessor/provenance.py delete mode 100644 src/twinkle_agentic/preprocessor/structural_noise.py delete mode 100644 src/twinkle_agentic/rollout/api_multi_turn.py create mode 100644 src/twinkle_agentic/rollout/api_sampler.py delete mode 100644 src/twinkle_agentic/rollout/factory.py delete mode 100644 tests/preprocessor/test_logprob_utils.py diff --git a/README.md b/README.md index 3d2b40eca..fc64af16f 100644 --- a/README.md +++ b/README.md @@ -109,7 +109,7 @@ sh INSTALL_MEGATRON.sh - 🎉2026-08-12 The ModelScope training service has been deployed to [Qwen/Qwen3.8-27B](https://www.modelscope.cn/models/Qwen/Qwen3.8-27B). - 🎉2026-08-04 Sandboxed multi-turn RL is now supported: run model-generated code in isolated [AgentENV](https://github.com/kvcache-ai/AgentENV) Firecracker microVMs, or in an OpenEnv server, with the same `train.py`. See the [cookbook](cookbook/rl/envs) and the [deployment guide](docs/source_en/Usage%20Guide/Agentic-RL-Deployment-and-Training.md). - 🎉2026-05-20 Support DeepSeek-V4-Flash and DeepSeek-V4-Pro models. -- 🎉2026-05-20 Multi-turn rollout and tool calling in RL are now supported. The Cookbook is currently being written. You can use `from twinkle_agentic.rollout import MultiTurnRollout/APIMultiTurnRollout` directly for multi-turn rollout. +- 🎉2026-05-20 Multi-turn rollout and tool calling in RL are now supported. The Cookbook is currently being written. You can use `from twinkle_agentic.rollout import MultiTurnRollout` directly for sampler, API, or mixed-backend multi-turn rollout. - 🎉2026-05-20 IM message alerting on training job failure is now supported. Usage: `import twinkle; twinkle.initialize(..., notifier=DingNotifier(...))`. - 🎉2026-04-27 Support the `padding_free` operation for sft/dpo/grpo/gkd, use `set_processor('InputProcessor', padding_free=True)` to train with it. - 🎉2026-04-22 The ModelScope service has been deployed to [Qwen/Qwen3.6-27B](https://www.modelscope.cn/models/Qwen/Qwen3.6-27B) with a new release 0.2.1. diff --git a/README_ZH.md b/README_ZH.md index f2f214f48..d7b3d66fa 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -105,7 +105,7 @@ Twinkle✨支持相同的算法接口运行在单GPU、torchrun多机、Ray、Cl - 🎉2026-08-12 ModelScope的训练服务部署为[Qwen/Qwen3.8-27B](https://www.modelscope.cn/models/Qwen/Qwen3.8-27B)。 - 🎉2026-08-04 支持沙箱环境下的多轮RL训练：模型生成的代码可在隔离的 [AgentENV](https://github.com/kvcache-ai/AgentENV) Firecracker microVM 或 OpenEnv 服务中执行，两个后端共用同一份 `train.py`。参考 [cookbook](cookbook/rl/envs) 和[部署文档](docs/source_zh/使用指引/Agentic%20RL部署与训练.md)。 - 🎉2026-05-20 支持DeepSeek-V4-Flash and DeepSeek-V4-Pro系列模型。 -- 🎉2026-05-20 支持多轮rollout和RL中的工具调用，Cookbook正在编写中，可以直接使用`from twinkle_agentic.rollout import MultiTurnRollout/APIMultiTurnRollout`进行多轮rollout。 +- 🎉2026-05-20 支持多轮rollout和RL中的工具调用，Cookbook正在编写中，可以直接使用 `from twinkle_agentic.rollout import MultiTurnRollout` 进行 sampler、API 或混合后端的多轮 rollout。 - 🎉2026-05-20 支持训练任务失败后的IM消息告警, 使用方式: `import twinkle; twinkle.initialize(..., notifier=DingNotifier(...))`。 - 🎉2026-04-27 支持sft/dpo/grpo/gkd的padding_free方法, 使用`set_processor('InputProcessor', padding_free=True)`来开启训练。 - 🎉2026-04-22 ModelScope的训练服务部署为[Qwen/Qwen3.6-27B](https://www.modelscope.cn/models/Qwen/Qwen3.6-27B)，并发布了0.2.1版本。 diff --git a/cookbook/rsi/code/challenge.py b/cookbook/rsi/code/challenge.py index 261217abe..4e3b6383c 100644 --- a/cookbook/rsi/code/challenge.py +++ b/cookbook/rsi/code/challenge.py @@ -31,7 +31,7 @@ from twinkle.sampler import vLLMSampler from twinkle_agentic.challenger import CodeChallenger, KeywordStore, load_seeds from twinkle_agentic.envs import LocalEnv -from twinkle_agentic.rollout import build_rollout +from twinkle_agentic.rollout import MultiTurnRollout from twinkle_agentic.tools.tool_manager import ToolManager sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) @@ -132,7 +132,7 @@ def main(): # Single-turn generation, but through the same rollout the RL loop uses, so a # challenger that should be allowed to run code while inventing only needs a # tool manager here rather than a different code path. - explorer = build_rollout( + explorer = MultiTurnRollout( sampler, template=template, tool_manager=ToolManager([]), diff --git a/cookbook/rsi/code/collect.py b/cookbook/rsi/code/collect.py index c18a907ad..2b15c6786 100644 --- a/cookbook/rsi/code/collect.py +++ b/cookbook/rsi/code/collect.py @@ -32,7 +32,7 @@ from twinkle.data_format import SamplingParams, Trajectory, user_data_get from twinkle_agentic.challenger import CodeChallenger, KeywordStore, load_seeds from twinkle_agentic.envs import LocalEnv -from twinkle_agentic.rollout import build_rollout +from twinkle_agentic.rollout import MultiTurnRollout from twinkle_agentic.tools.tool_manager import ToolManager # Appended, not prepended: rsi.py imports this half into the process that already @@ -115,9 +115,9 @@ def build_challenger(args, sampler, template, *, recorder=None) -> CollectingCha # One rollout for proposing and, through solver_params, for solving. max_turns=1 # because a code answer is one message: there is nothing for a second turn to # react to until the asserts have run, and running them is the next stage. - explorer = build_rollout(sampler, template=template, - tool_manager=ToolManager([]), max_turns=1, - sampling_params=params) + explorer = MultiTurnRollout(sampler, template=template, + tool_manager=ToolManager([]), max_turns=1, + sampling_params=params) store = None if args.code_keywords_n > 0: store = KeywordStore(args.code_keyword_db, CATEGORIES) diff --git a/docs/source_en/Components/Agentic/Envs.md b/docs/source_en/Components/Agentic/Envs.md index 988eac6fe..774f412c0 100644 --- a/docs/source_en/Components/Agentic/Envs.md +++ b/docs/source_en/Components/Agentic/Envs.md @@ -191,7 +191,7 @@ Downstream usage is the same for both modes: ```python from twinkle_agentic.envs.env_tool import EnvTool from twinkle_agentic.tools.tool_manager import ToolManager -from twinkle_agentic.rollout.api_multi_turn import APIMultiTurnRollout +from twinkle_agentic.rollout import MultiTurnRollout env.reset() @@ -200,7 +200,7 @@ env_tools = EnvTool.from_env(env) manager = ToolManager(env_tools) # Use in rollout -rollout = APIMultiTurnRollout(api=api, tool_manager=manager, max_turns=10) +rollout = MultiTurnRollout(api=api, template=template, tool_manager=manager, max_turns=10) results = rollout(trajectories) ``` diff --git a/docs/source_en/Components/Agentic/Multi-Turn-Tool-Usage.md b/docs/source_en/Components/Agentic/Multi-Turn-Tool-Usage.md index 24296f584..c4b4615be 100644 --- a/docs/source_en/Components/Agentic/Multi-Turn-Tool-Usage.md +++ b/docs/source_en/Components/Agentic/Multi-Turn-Tool-Usage.md @@ -19,8 +19,9 @@ The simplest way to run a multi-turn tool-use rollout using an OpenAI-compatible from twinkle_agentic.protocol.openai import OpenAI from twinkle_agentic.tools.base import Tool from twinkle_agentic.tools.tool_manager import ToolManager -from twinkle_agentic.rollout.api_multi_turn import APIMultiTurnRollout +from twinkle_agentic.rollout import MultiTurnRollout from twinkle.data_format.sampling import SamplingParams +from twinkle.template import Template # 1. Define tools class WeatherTool(Tool): @@ -47,16 +48,17 @@ class WeatherTool(Tool): # 2. Set up ToolManager manager = ToolManager([WeatherTool()]) -# 3. Create API client -api = OpenAI(model='qwen3.5-32b', base_url='http://localhost:8000/v1') +# 3. Create API client and the local template used to encode its replies +api = OpenAI(model='qwen3.5-32b', base_url='http://localhost:8000/v1', concurrency=8) +template = Template(model_id='Qwen/Qwen3.5-32B') # 4. Create rollout -rollout = APIMultiTurnRollout( +rollout = MultiTurnRollout( api=api, + template=template, tool_manager=manager, sampling_params=SamplingParams(temperature=0.7, max_tokens=2048), max_turns=6, - concurrency=8, ) # 5. Prepare trajectories @@ -138,7 +140,7 @@ env_tools = EnvTool.from_env(env) manager = ToolManager(env_tools) # Use manager in rollout as usual -rollout = APIMultiTurnRollout(api=api, tool_manager=manager, max_turns=10) +rollout = MultiTurnRollout(api=api, template=template, tool_manager=manager, max_turns=10) ``` ## Using OpenEnv Environments @@ -193,11 +195,12 @@ results = rollout(trajectories, tool_manager=managers) ## Trace Debugging -Both rollout implementations support trace dumps for debugging: +The unified rollout supports trace dumps for debugging: ```python -rollout = APIMultiTurnRollout( +rollout = MultiTurnRollout( api=api, + template=template, tool_manager=manager, trace_dir='traces/', trace_callback=lambda t: t['turns'] > 1, # Only store multi-turn diff --git a/docs/source_en/Components/Agentic/Rollout.md b/docs/source_en/Components/Agentic/Rollout.md index e803b9076..10ccc3471 100644 --- a/docs/source_en/Components/Agentic/Rollout.md +++ b/docs/source_en/Components/Agentic/Rollout.md @@ -1,6 +1,6 @@ # Multi-Turn Rollout -The Rollout module provides multi-turn conversation rollout engines for agentic RLHF training. Two implementations are available: `MultiTurnRollout` for batched vLLM sampling and `APIMultiTurnRollout` for OpenAI-compatible API endpoints. +The Rollout module provides one multi-turn conversation engine for agentic RLHF training. `MultiTurnRollout` can generate each assistant turn with a local sampler, an OpenAI-compatible API, or a callback that chooses between them. ## Rollout Base Class @@ -19,12 +19,12 @@ All rollouts accept a list of trajectories and return the same number of traject ## MultiTurnRollout -Batched multi-turn rollout engine that uses a vLLM sampler for generation. All active trajectories are sampled in a single batched call per turn for maximum throughput. +Multi-turn rollout engine supporting local samplers, external APIs, and per-turn backend selection. Each trajectory runs independently in the rollout thread pool. ### Per-turn Loop 1. Encode each trajectory into an `InputFeature` with a generation prompt -2. Batch `sampler.sample(active_pifs)` — all live trajectories in parallel +2. Call `response_callback(...)` to obtain one `SampledSequence` from the sampler or API 3. Check termination: `stop_reason == 'length'`, no tool calls, or max turns reached 4. Dispatch tools via `ToolManager`, append tool responses 5. Compute bridge tokens (tool turns + generation prompt) with `labels = -100` @@ -53,8 +53,12 @@ results = rollout(trajectories) | Parameter | Type | Description | |-----------|------|-------------| -| `sampler` | Sampler | vLLM sampler instance for batched generation. | -| `template` | `Template` | Chat template for encoding/decoding. | +| `sampler` | Sampler | Local sampler. Used by default when both backends exist. | +| `api` | `API` | Optional external generation API. | +| `template` | `Template` | Required local chat template for encoding every backend's output. | +| `response_callback` | `Callable` | Optional per-turn backend selector returning `SampledSequence`. | +| `api_appended_as` | `str` | API turns are `demonstration` (SFT only) or `context` (no loss). | +| `api_kwargs` | `Dict` | Request fields forwarded to each API call. | | `tool_manager` | `ToolManager` | Tool dispatcher. Can also be passed per-call. | | `sampling_params` | `SamplingParams` | Default sampling parameters. | | `max_turns` | `int` | Maximum number of turns per trajectory (default: 6). | @@ -72,6 +76,7 @@ Each output trajectory dict includes: | `messages` | `List[Dict]` | Full conversation including tool turns. | | `input_ids` | `List[int]` | Token IDs of the full sequence. | | `labels` | `List[int]` | Training labels (`-100` for non-trainable tokens). | +| `completion_mask` | `List[int]` | Policy-generated positions that carry rollout log probabilities. | | `turns` | `int` | Number of turns performed. | | `stop_reason` | `str` | `'stop'` / `'length'` | | `truncated` | `bool` | Whether the trajectory was cut off rather than concluding on its own: generation hit `max_tokens` (`stop_reason='length'`), the turn limit was reached, or a length cap dropped it. | @@ -87,54 +92,33 @@ rollout_actor = MultiTurnRollout.remote(sampler=sampler, template=template, ...) results = ray.get(rollout_actor.__call__.remote(trajectories)) ``` -## APIMultiTurnRollout +## API and Mixed-Backend Rollouts -Multi-turn rollout over an OpenAI-compatible chat-completions API. Each trajectory runs independently in a thread pool for network concurrency. +API-only rollout uses the same class and still requires the local template that tokenizes external replies: ```python -from twinkle_agentic.rollout.api_multi_turn import APIMultiTurnRollout from twinkle_agentic.protocol.openai import OpenAI +from twinkle_agentic.rollout import MultiTurnRollout -api = OpenAI(model='qwen3.5-32b', base_url='http://localhost:8000/v1') - -rollout = APIMultiTurnRollout( - api=api, +api = OpenAI(model='qwen3.5-32b', base_url='http://localhost:8000/v1', concurrency=8) +rollout = MultiTurnRollout( + api, + template=template, tool_manager=tool_manager, sampling_params=SamplingParams(temperature=0.7), max_turns=6, - concurrency=8, trace_dir='api_traces/', ) - results = rollout(trajectories) ``` -### Parameters - -| Parameter | Type | Description | -|-----------|------|-------------| -| `api` | `OpenAI` | OpenAI-compatible API client. | -| `tool_manager` | `ToolManager` | Tool dispatcher (single or per-trajectory list). | -| `sampling_params` | `SamplingParams` | Default sampling parameters. | -| `max_turns` | `int` | Maximum turns per trajectory (default: 6). | -| `concurrency` | `int` | Thread pool size for parallel API calls (default: 8). | -| `extra_body` | `Dict` | Extra fields to include in API requests. | -| `trace_dir` | `str` | Directory for trace dumps. | +When both `sampler` and `api` are supplied, the default is the sampler. Pass `response_callback` to choose per turn; it receives both backends and must return one `SampledSequence`. API turns have no rollout log probabilities, so `api_appended_as='demonstration'` includes them in SFT but excludes them from GRPO. Use `'context'` to exclude them from both. ### Stop Reasons | Reason | Description | |--------|-------------| | `stop` | Assistant responded without tool calls (natural end). | -| `length` | API returned `finish_reason='length'` (token limit). | -| `max_turns` | Reached `max_turns` limit. | -| `api_error` | API call or tool execution raised an exception. | - -## Choosing Between Rollouts - -| Feature | MultiTurnRollout | APIMultiTurnRollout | -|---------|-----------------|---------------------| -| **Backend** | vLLM sampler (local GPU) | OpenAI-compatible API | -| **Training integration** | Produces `input_ids` / `labels` for GRPO | Messages only (for data collection) | -| **Batching** | GPU-level batch parallelism | Network-level thread concurrency | -| **Use case** | Online RLHF training loop | Offline data generation / evaluation | +| `length` | Generation reached its token limit. | +| `max_turns` | Reached the tool-turn limit without a follow-up. | +| `generation_error` | The external endpoint failed before returning a valid response. | diff --git "a/docs/source_zh/\347\273\204\344\273\266/Agentic/Envs.md" "b/docs/source_zh/\347\273\204\344\273\266/Agentic/Envs.md" index 76c729eb7..bae412f5c 100644 --- "a/docs/source_zh/\347\273\204\344\273\266/Agentic/Envs.md" +++ "b/docs/source_zh/\347\273\204\344\273\266/Agentic/Envs.md" @@ -191,7 +191,7 @@ env.close() ```python from twinkle_agentic.envs.env_tool import EnvTool from twinkle_agentic.tools.tool_manager import ToolManager -from twinkle_agentic.rollout.api_multi_turn import APIMultiTurnRollout +from twinkle_agentic.rollout import MultiTurnRollout env.reset() @@ -200,7 +200,7 @@ env_tools = EnvTool.from_env(env) manager = ToolManager(env_tools) # 在 rollout 中使用 -rollout = APIMultiTurnRollout(api=api, tool_manager=manager, max_turns=10) +rollout = MultiTurnRollout(api=api, template=template, tool_manager=manager, max_turns=10) results = rollout(trajectories) ``` diff --git "a/docs/source_zh/\347\273\204\344\273\266/Agentic/Multi-Turn-Tool-Usage.md" "b/docs/source_zh/\347\273\204\344\273\266/Agentic/Multi-Turn-Tool-Usage.md" index 8b94b2ed4..b3af8b45f 100644 --- "a/docs/source_zh/\347\273\204\344\273\266/Agentic/Multi-Turn-Tool-Usage.md" +++ "b/docs/source_zh/\347\273\204\344\273\266/Agentic/Multi-Turn-Tool-Usage.md" @@ -19,8 +19,9 @@ Agentic rollout 管线由四个核心组件组成： from twinkle_agentic.protocol.openai import OpenAI from twinkle_agentic.tools.base import Tool from twinkle_agentic.tools.tool_manager import ToolManager -from twinkle_agentic.rollout.api_multi_turn import APIMultiTurnRollout +from twinkle_agentic.rollout import MultiTurnRollout from twinkle.data_format.sampling import SamplingParams +from twinkle.template import Template # 1. 定义工具 class WeatherTool(Tool): @@ -47,16 +48,17 @@ class WeatherTool(Tool): # 2. 设置 ToolManager manager = ToolManager([WeatherTool()]) -# 3. 创建 API 客户端 -api = OpenAI(model='qwen3.5-32b', base_url='http://localhost:8000/v1') +# 3. 创建 API 客户端，以及用于编码 API 回复的本地 template +api = OpenAI(model='qwen3.5-32b', base_url='http://localhost:8000/v1', concurrency=8) +template = Template(model_id='Qwen/Qwen3.5-32B') # 4. 创建 rollout -rollout = APIMultiTurnRollout( +rollout = MultiTurnRollout( api=api, + template=template, tool_manager=manager, sampling_params=SamplingParams(temperature=0.7, max_tokens=2048), max_turns=6, - concurrency=8, ) # 5. 准备轨迹 @@ -138,7 +140,7 @@ env_tools = EnvTool.from_env(env) manager = ToolManager(env_tools) # 照常在 rollout 中使用 manager -rollout = APIMultiTurnRollout(api=api, tool_manager=manager, max_turns=10) +rollout = MultiTurnRollout(api=api, template=template, tool_manager=manager, max_turns=10) ``` ## 使用 OpenEnv 环境 @@ -193,11 +195,12 @@ results = rollout(trajectories, tool_manager=managers) ## 跟踪调试 -两种 rollout 实现都支持跟踪文件输出用于调试： +统一的 rollout 支持跟踪文件输出用于调试： ```python -rollout = APIMultiTurnRollout( +rollout = MultiTurnRollout( api=api, + template=template, tool_manager=manager, trace_dir='traces/', trace_callback=lambda t: t['turns'] > 1, # 仅存储多轮对话 diff --git "a/docs/source_zh/\347\273\204\344\273\266/Agentic/Rollout.md" "b/docs/source_zh/\347\273\204\344\273\266/Agentic/Rollout.md" index 134532c97..767e8a538 100644 --- "a/docs/source_zh/\347\273\204\344\273\266/Agentic/Rollout.md" +++ "b/docs/source_zh/\347\273\204\344\273\266/Agentic/Rollout.md" @@ -1,6 +1,6 @@ # 多轮 Rollout -Rollout 模块提供了用于 Agentic RLHF 训练的多轮对话 rollout 引擎。包含两种实现：用于批量 vLLM 采样的 `MultiTurnRollout` 和用于 OpenAI 兼容 API 端点的 `APIMultiTurnRollout`。 +Rollout 模块提供统一的多轮对话引擎 `MultiTurnRollout`，每轮 assistant 可由本地 sampler、OpenAI 兼容 API，或在两者间动态选择的 callback 生成。 ## Rollout 基类 @@ -19,12 +19,12 @@ class Rollout(ABC): ## MultiTurnRollout -批量多轮 rollout 引擎，使用 vLLM 采样器进行生成。每轮中所有活跃轨迹通过单次批量采样调用并行处理，最大化吞吐量。 +统一的多轮 rollout 引擎，支持本地 sampler、外部 API 和逐轮后端选择。每条轨迹在线程池中独立执行。 ### 每轮循环 1. 将每个轨迹编码为带生成提示的 `InputFeature` -2. 批量调用 `sampler.sample(active_pifs)` —— 所有活跃轨迹并行 +2. 调用 `response_callback(...)`，从 sampler 或 API 获取一个 `SampledSequence` 3. 检查终止条件：`stop_reason == 'length'`、无工具调用、或达到最大轮次 4. 通过 `ToolManager` 分发工具调用，追加工具响应 5. 计算桥接 token（工具轮次 + 生成提示），设置 `labels = -100` @@ -53,8 +53,12 @@ results = rollout(trajectories) | 参数 | 类型 | 说明 | |------|------|------| -| `sampler` | Sampler | 用于批量生成的 vLLM 采样器实例。 | -| `template` | `Template` | 用于编码/解码的聊天模板。 | +| `sampler` | Sampler | 本地 sampler；两个后端同时存在时默认使用它。 | +| `api` | `API` | 可选的外部生成 API。 | +| `template` | `Template` | 必传；用于编码所有后端的输出。 | +| `response_callback` | `Callable` | 可选的逐轮后端选择器，返回 `SampledSequence`。 | +| `api_appended_as` | `str` | API 轮为 `demonstration`（仅 SFT）或 `context`（不训练）。 | +| `api_kwargs` | `Dict` | 传给每次 API 调用的请求字段。 | | `tool_manager` | `ToolManager` | 工具分发器。也可以按调用传入。 | | `sampling_params` | `SamplingParams` | 默认采样参数。 | | `max_turns` | `int` | 每个轨迹的最大轮次（默认：6）。 | @@ -72,6 +76,7 @@ results = rollout(trajectories) | `messages` | `List[Dict]` | 包含工具轮次的完整对话。 | | `input_ids` | `List[int]` | 完整序列的 token ID。 | | `labels` | `List[int]` | 训练标签（非可训练 token 为 `-100`）。 | +| `completion_mask` | `List[int]` | 由 policy 生成且具有 rollout log probability 的位置。 | | `turns` | `int` | 执行的轮次数。 | | `stop_reason` | `str` | `'stop'` / `'length'` | | `truncated` | `bool` | 轨迹是否被截断（而非自行结束）：生成触及 `max_tokens`（`stop_reason='length'`）、达到轮次上限，或被长度上限丢弃。 | @@ -87,54 +92,33 @@ rollout_actor = MultiTurnRollout.remote(sampler=sampler, template=template, ...) results = ray.get(rollout_actor.__call__.remote(trajectories)) ``` -## APIMultiTurnRollout +## API 与混合后端 Rollout -通过 OpenAI 兼容 chat-completions API 进行多轮 rollout。每个轨迹在线程池中独立运行，实现网络并发。 +纯 API 模式使用同一个类，并仍需传入本地 template，以便将外部回复编码成训练侧一致的 token： ```python -from twinkle_agentic.rollout.api_multi_turn import APIMultiTurnRollout from twinkle_agentic.protocol.openai import OpenAI +from twinkle_agentic.rollout import MultiTurnRollout -api = OpenAI(model='qwen3.5-32b', base_url='http://localhost:8000/v1') - -rollout = APIMultiTurnRollout( - api=api, +api = OpenAI(model='qwen3.5-32b', base_url='http://localhost:8000/v1', concurrency=8) +rollout = MultiTurnRollout( + api, + template=template, tool_manager=tool_manager, sampling_params=SamplingParams(temperature=0.7), max_turns=6, - concurrency=8, trace_dir='api_traces/', ) - results = rollout(trajectories) ``` -### 参数 - -| 参数 | 类型 | 说明 | -|------|------|------| -| `api` | `OpenAI` | OpenAI 兼容 API 客户端。 | -| `tool_manager` | `ToolManager` | 工具分发器（单个或按轨迹的列表）。 | -| `sampling_params` | `SamplingParams` | 默认采样参数。 | -| `max_turns` | `int` | 每轨迹最大轮次（默认：6）。 | -| `concurrency` | `int` | 并行 API 调用的线程池大小（默认：8）。 | -| `extra_body` | `Dict` | API 请求中附加的额外字段。 | -| `trace_dir` | `str` | 跟踪文件目录。 | +同时传入 `sampler` 和 `api` 时，默认使用 sampler。传入 `response_callback` 可逐轮选择后端；callback 会收到两个后端，并必须返回一个 `SampledSequence`。API 轮没有 rollout log probability，因此 `api_appended_as='demonstration'` 会让它参与 SFT 但跳过 GRPO；使用 `'context'` 可让它完全不参与训练。 ### 停止原因 | 原因 | 说明 | |------|------| | `stop` | 助手回复未包含工具调用（自然结束）。 | -| `length` | API 返回 `finish_reason='length'`（token 限制）。 | -| `max_turns` | 达到 `max_turns` 限制。 | -| `api_error` | API 调用或工具执行抛出异常。 | - -## 选择建议 - -| 特性 | MultiTurnRollout | APIMultiTurnRollout | -|------|-----------------|---------------------| -| **后端** | vLLM 采样器（本地 GPU） | OpenAI 兼容 API | -| **训练集成** | 生成 `input_ids` / `labels` 用于 GRPO | 仅消息（用于数据收集） | -| **批处理** | GPU 级别批量并行 | 网络级别线程并发 | -| **用例** | 在线 RLHF 训练循环 | 离线数据生成 / 评估 | +| `length` | 生成达到 token 上限。 | +| `max_turns` | 达到工具轮次上限且没有 follow-up。 | +| `generation_error` | 外部端点未能返回有效响应。 | diff --git a/src/twinkle/loss/grpo.py b/src/twinkle/loss/grpo.py index 33b89c206..85470a0dc 100644 --- a/src/twinkle/loss/grpo.py +++ b/src/twinkle/loss/grpo.py @@ -227,6 +227,39 @@ def _pad_and_align_to_batch( return result + def _resolve_loss_mask(self, inputs: Dict, labels: 'torch.Tensor') -> 'torch.Tensor': + """Positions this loss may score: trainable *and* log-prob-bearing. + + ``labels`` alone answers "should this token be scored", which is all SFT + needs. A policy-gradient loss also needs a sampling log-prob per token to + form an importance ratio, and a turn produced outside the sampled policy + (an API, a human, a replayed demonstration) has none. Such turns carry + ``completion_mask == 0``: excluded here, yet still trainable for SFT. + + A feature without ``completion_mask`` predates the field, and there every + trainable token was the policy's own, so the mask degenerates to + ``labels != ignore_index`` and old trajectories train exactly as before. + """ + import torch + trainable = (labels != self.ignore_index).bool() + completion_mask = inputs.get('completion_mask') + if completion_mask is None: + return trainable + if not torch.is_tensor(completion_mask): + completion_mask = torch.as_tensor(completion_mask) + completion_mask = completion_mask.to(trainable.device) + if completion_mask.dim() == 1: + completion_mask = completion_mask.unsqueeze(0) + if completion_mask.shape != trainable.shape: + raise ValueError(f'completion_mask shape {tuple(completion_mask.shape)} does not match labels shape ' + f'{tuple(trainable.shape)}. A misaligned mask would apply importance ratios to ' + 'the wrong tokens, so it is refused rather than broadcast.') + loss_mask = trainable & completion_mask.bool() + if self.enable_sampling_replay and not bool((loss_mask == trainable).all()): + raise ValueError('sampling replay does not support turns generated outside the sampled policy: ' + 'they are trainable but have no sampling mask to replay against.') + return loss_mask + def __call__( self, inputs: Dict, @@ -269,7 +302,7 @@ def __call__( logps = outputs.get('logps') if self.enable_sampling_replay and logps is None: raise RuntimeError('sampling replay logps must be computed by the model forward') - loss_mask = (labels != self.ignore_index).bool() + loss_mask = self._resolve_loss_mask(inputs, labels) if logps is None: logits = outputs.get('logits') if logits.shape[1] != labels.shape[1]: diff --git a/src/twinkle/model/megatron/megatron.py b/src/twinkle/model/megatron/megatron.py index 5240816d3..8d85dacf8 100644 --- a/src/twinkle/model/megatron/megatron.py +++ b/src/twinkle/model/megatron/megatron.py @@ -402,6 +402,8 @@ def post_loss_function(output_tensor, inputs, logps, unpacked_logits=None, entro def forward_step_func(data_iterator, model): batch = next(data_iterator) labels = batch.pop('labels', None) + # Not a model argument; restored below so the loss can read it. + completion_mask = batch.pop('completion_mask', None) unwrapped_model = self.strategy.unwrap_model([model])[0] if disable_lora and isinstance(unwrapped_model, PeftModel): with unwrapped_model.disable_adapter(): @@ -410,6 +412,8 @@ def forward_step_func(data_iterator, model): output_tensor = model(**batch) batch['labels'] = labels + if completion_mask is not None: + batch['completion_mask'] = completion_mask logps = None unpacked_logits = None entropies = None @@ -440,6 +444,10 @@ def forward_step_func(data_iterator, model): if entropies is not None: entropies = processor.postprocess_tensor_cp(entropies, cu_seqlens=cu_seqlens_q) batch['labels'] = processor.postprocess_tensor_cp(labels, cu_seqlens=cu_seqlens_q) + if completion_mask is not None: + # Same index space as labels, so it needs the same CP reassembly. + batch['completion_mask'] = processor.postprocess_tensor_cp( + completion_mask, cu_seqlens=cu_seqlens_q) if 'position_ids' in batch: pos = batch['position_ids'] if pos.dim() == 3: diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py index 46ace2c64..9d0cec9f1 100644 --- a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py +++ b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py @@ -845,7 +845,7 @@ def prepare_inputs(self, inputs): """Prepare inputs 1. set extra_kwargs['position_ids'] - 2. split labels + 2. split labels, and completion_mask when present """ input_ids = inputs.get('input_ids') position_ids = inputs.get('position_ids') @@ -863,7 +863,11 @@ def prepare_inputs(self, inputs): self.extra_kwargs['input_ids'] = input_ids.clone() if 'labels' in inputs: labels = inputs.get('labels') - _, _, labels, _, _, _, _ = self.pad_and_split_inputs( + # completion_mask sits on the labels' index space, so it is padded and + # split identically -- unlike loss_scale, which is rolled beforehand. + completion_mask = inputs.get('completion_mask') + extra_split_values = None if completion_mask is None else [(completion_mask, 0, -1)] + _, _, labels, _, _, _, extra_values = self.pad_and_split_inputs( None, None, labels, @@ -871,8 +875,11 @@ def prepare_inputs(self, inputs): None, None, real_position_ids=real_position_ids, + extra_split_values=extra_split_values, ) inputs['labels'] = labels + if extra_values: + inputs['completion_mask'] = extra_values[0] return inputs @@ -986,6 +993,19 @@ def _trim_gathered_sequence_padding(tensor: torch.Tensor, real_position_ids: tor return torch.cat(pieces, dim=1).contiguous() if pieces else tensor[:, :0].contiguous() return tensor[:, :real_position_ids.shape[-1]].contiguous() + def _gather_completion_mask(self, inputs: Dict[str, Any], real_position_ids) -> None: + """Gather ``completion_mask`` in place, mirroring the labels gather. + + Deliberately not routed through :class:`GatherLoss`: the mask carries no + gradient, and reusing that autograd Function would attach a second backward + path to whichever tensor were passed alongside it, double-scaling its grad. + """ + mask = inputs.get('completion_mask') + if mask is None or not torch.is_tensor(mask) or mask.dim() < 2: + return + gathered = sequence_parallel.gather(mask, dim=1, position_ids=real_position_ids) + inputs['completion_mask'] = self._trim_gathered_sequence_padding(gathered, real_position_ids) + def gather_loss_tensors( self, inputs: Dict[str, Any], @@ -1017,6 +1037,7 @@ def gather_loss_tensors( gathered_labels = self._trim_gathered_sequence_padding(gathered_labels, real_position_ids) outputs['logits'] = gathered_hidden inputs['labels'] = gathered_labels + self._gather_completion_mask(inputs, real_position_ids) return inputs, outputs if labels is None or logps is None: return inputs, outputs @@ -1031,6 +1052,7 @@ def gather_loss_tensors( gathered_labels = self._trim_gathered_sequence_padding(gathered_labels, real_position_ids) outputs['logps'] = gathered_logps inputs['labels'] = gathered_labels + self._gather_completion_mask(inputs, real_position_ids) entropies = outputs.get('entropies') if entropies is not None and torch.is_tensor(entropies) and entropies.dim() >= 2: gathered_entropies, _ = GatherLoss.apply(entropies, labels, 1, real_position_ids) diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index dda7aea8d..9ce4c4cb7 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -568,6 +568,8 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec enable_sp=getattr(self, '_enable_sp', False), ) labels: torch.Tensor = inputs.pop('labels', None) + # Not a model argument; the loss reads it back off `inputs` further down. + completion_mask = inputs.pop('completion_mask', None) replay_metadata = replay_loss_mask = replay_masked_labels = None if enable_sampling_replay: replay_loss_mask, replay_masked_labels, replay_metadata = _prepare_sampling_replay( @@ -595,6 +597,8 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec recorded_routing = rr_cleanup() inputs['labels'] = labels + if completion_mask is not None: + inputs['completion_mask'] = completion_mask if task != 'embedding' and labels is not None and loss_require_logps: loss_mask = replay_loss_mask if enable_sampling_replay else (labels != -100).bool() masked_labels = replay_masked_labels if enable_sampling_replay else labels.masked_fill(~loss_mask, 0) @@ -689,6 +693,8 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T enable_sp=getattr(self, '_enable_sp', False), ) labels = inputs.pop('labels', None) + # Not a model argument; the loss reads it back off `inputs` further down. + completion_mask = inputs.pop('completion_mask', None) replay_metadata = replay_loss_mask = replay_masked_labels = None if enable_sampling_replay: packed_position_ids = processor._is_packed_position_ids(inputs.get('position_ids')) @@ -720,6 +726,8 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T recorded_routing = rr_cleanup() inputs['labels'] = labels + if completion_mask is not None: + inputs['completion_mask'] = completion_mask if task != 'embedding' and labels is not None and loss_require_logps: loss_mask = replay_loss_mask if enable_sampling_replay else (labels != -100).bool() masked_labels = replay_masked_labels if enable_sampling_replay else labels.masked_fill(~loss_mask, 0) diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py index 5d67f1fcb..b28812f72 100644 --- a/src/twinkle/processor/base.py +++ b/src/twinkle/processor/base.py @@ -608,9 +608,10 @@ def unpack_packed_sequences( """Unpack packed (padding_free) sequences into per-sequence batch format. Called after SP gather / CP gather, before loss computation. - Unpacks ``labels`` and any present output keys (``logps``, ``logits``) - from ``[1, total_tokens, ...]`` to ``[num_sequences, max_seq_len, ...]``. - Keys that are ``None`` are silently skipped. + Unpacks ``labels``, ``completion_mask`` and any present output keys + (``logps``, ``logits``) from ``[1, total_tokens, ...]`` to + ``[num_sequences, max_seq_len, ...]``. Keys that are ``None`` are silently + skipped. For ``task='embedding'`` the outputs are already pooled to ``[n_seqs, H]`` by ``postprocess_tensor_sp``, so this is a no-op. @@ -627,23 +628,29 @@ def unpack_packed_sequences( from copy import copy - # Collect output keys to unpack: (key, pad_value) - output_keys = [] - for key, pad_val in [('logps', 0), ('values', 0), ('entropies', 0), ('logits', 0)]: - if outputs and outputs.get(key) is not None: - output_keys.append((key, pad_val)) - - all_tensors = [labels] + [outputs[k] for k, _ in output_keys] - all_pads = [-100] + [p for _, p in output_keys] - unpacked = self._unpack_by_position_ids(position_ids, *all_tensors, padding_values=all_pads) + # (key, tensor, pad_value) for everything that must come back as + # [num_sequences, max_seq_len]. completion_mask shares the labels' index + # space, so leaving it packed would hand the loss two differently shaped + # views of the same sequence. + input_specs = [('labels', labels, -100)] + if inputs.get('completion_mask') is not None: + input_specs.append(('completion_mask', inputs['completion_mask'], self.padding_map['completion_mask'])) + output_specs = [(key, outputs[key], 0) for key in ('logps', 'values', 'entropies', 'logits') + if outputs and outputs.get(key) is not None] + + specs = input_specs + output_specs + unpacked = iter( + self._unpack_by_position_ids( + position_ids, *[tensor for _, tensor, _ in specs], padding_values=[pad for _, _, pad in specs])) inputs = copy(inputs) - inputs['labels'] = unpacked[0] + for key, _, _ in input_specs: + inputs[key] = next(unpacked) - if output_keys: + if output_specs: outputs = copy(outputs) - for i, (key, _) in enumerate(output_keys): - outputs[key] = unpacked[i + 1] + for key, _, _ in output_specs: + outputs[key] = next(unpacked) return inputs, outputs diff --git a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py index 72600bc4a..7877ddb55 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py @@ -36,12 +36,10 @@ def _convert_ndarray_to_list(obj: Any) -> Any: return obj -# max_concurrency: how many sample() calls one worker serves at once. Without it -# Ray runs one method per actor at a time, so concurrent callers queue at the actor -# and never share a batch inside AsyncLLM. 24 is what vLLM reports as the maximum -# concurrency its KV cache holds for this context length; past it vLLM preempts and -# recomputes, which costs more than it gains. -@remote_class(max_concurrency=24) +_MAX_CONCURRENCY = max(1, int(os.environ.get('TWINKLE_SAMPLER_MAX_CONCURRENCY') or 24)) + + +@remote_class(max_concurrency=_MAX_CONCURRENCY) class vLLMSampler(Sampler, CheckpointEngineMixin): """A vLLM-based sampler using VLLMEngine (AsyncLLM). diff --git a/src/twinkle/server/sampler/backends/mock_sampler.py b/src/twinkle/server/sampler/backends/mock_sampler.py index d8355008b..2d5e5930d 100644 --- a/src/twinkle/server/sampler/backends/mock_sampler.py +++ b/src/twinkle/server/sampler/backends/mock_sampler.py @@ -252,8 +252,9 @@ def _build_new_input_feature(pif: Any, tokens: list[int]) -> dict[str, Any]: Produces a plain-dict ``InputFeature`` that carries the running context for the next multi-turn round: ``input_ids`` is the prior prompt plus - this round's sampled tokens, and ``labels`` marks the sampled tokens as - trainable (their own ids) while prior/context positions stay ``-100``. + this round's sampled tokens, ``labels`` marks the sampled tokens as + trainable (their own ids) while prior/context positions stay ``-100``, + and ``completion_mask`` marks them as the policy's own output. This mirrors the shape a real sampler's ``concat_input_feature`` yields, which the multi-turn rollout relies on (it reads ``new_input_feature.input_ids`` and counts trainable ``labels``). @@ -270,8 +271,14 @@ def _build_new_input_feature(pif: Any, tokens: list[int]) -> dict[str, Any]: # No (or misaligned) prior labels: treat the entire prior context as # non-trainable so only this round's sampled tokens count. labels = [-100] * len(prev_ids) + prev_mask = feat.get('completion_mask') + if prev_mask is not None and len(prev_mask) == len(prev_ids): + completion_mask = list(prev_mask) + else: + completion_mask = [0 if label == -100 else 1 for label in labels] feat['input_ids'] = prev_ids + list(tokens) feat['labels'] = labels + list(tokens) + feat['completion_mask'] = completion_mask + [1] * len(tokens) feat['length'] = len(feat['input_ids']) return feat diff --git a/src/twinkle/template/base.py b/src/twinkle/template/base.py index d914b01a9..10912bd7b 100644 --- a/src/twinkle/template/base.py +++ b/src/twinkle/template/base.py @@ -23,6 +23,18 @@ VideoInput = Union[str, List['Image.Image'], 'torch.Tensor'] AudioInput = Union[str, np.ndarray, 'torch.Tensor'] +# Fields that are one entry per token and must be sliced with ``input_ids``. +# ``mm_token_type_ids`` is excluded: it may carry a leading batch dim and is +# sliced on its last axis instead. +_SEQUENCE_ALIGNED_FIELDS = ('labels', 'completion_mask') + +# What an appended turn is to a trainer: the policy's own completion (scored, and +# a log-prob exists for each of its tokens), someone else's completion offered +# for imitation (scored, no log-prob -- usable by SFT but not by RL), or history +# that no loss may touch. There is deliberately no fourth role: a log-prob is +# only ever needed for a token that is also scored. +_APPEND_ROLES = ('completion', 'demonstration', 'context') + @remote_class() class Template: @@ -197,19 +209,48 @@ def _invoke_post_pipeline(self, input_features: List[InputFeature]) -> List[Inpu current = next_batch return current - def concat_input_feature(self, prompt_input_feature: InputFeature, new_tokens: List[int]) -> InputFeature: + def concat_input_feature(self, + prompt_input_feature: InputFeature, + new_tokens: List[int], + *, + appended_as: Literal['completion', 'demonstration', 'context'] = 'completion', + tool_calls: Optional[List[Dict[str, Any]]] = None) -> InputFeature: + """Append one generated turn to an already-encoded prefix. + + Args: + appended_as: what the turn is to a trainer, which decides ``labels`` + and ``completion_mask`` together: + + * ``'completion'`` -- the sampled policy's own output. Scored, and + a log-prob exists for every token. + * ``'demonstration'`` -- written by someone else (a stronger model, + a human) and offered for imitation. Scored, but carries no + log-prob, so RL losses skip it while SFT trains on it. + * ``'context'`` -- history that later turns must see and no loss + may touch. + tool_calls: calls to attach to the appended message, for generators that + return them as structured fields (any OpenAI-compatible API does) + rather than as markup inside the text, which is all + ``parse_tool_call`` can read. + """ import copy import torch assert self.truncation_strategy != 'split', 'concat_input_feature does not support `truncation_strategy=split`' + if appended_as not in _APPEND_ROLES: + raise ValueError(f'appended_as must be one of {_APPEND_ROLES}, got {appended_as!r}') result = copy.deepcopy(prompt_input_feature) prompt_ids = result['input_ids'] labels = list(result.get('labels', [])) input_ids = list(prompt_ids) + new_tokens labels = labels[-1:] + labels[:-1] # roll to input order - labels = labels + new_tokens + completion_mask = self._prefix_completion_mask(result, labels) + scored = appended_as != 'context' + labels = labels + (new_tokens if scored else [-100] * len(new_tokens)) + completion_mask = completion_mask + [int(appended_as == 'completion')] * len(new_tokens) # We don't need to roll back, self._invoke_post_pipeline will do this. result['input_ids'] = input_ids result['labels'] = labels + result['completion_mask'] = completion_mask if 'mm_token_type_ids' in result: mm_token_type_ids = result['mm_token_type_ids'] if not isinstance(mm_token_type_ids, torch.Tensor): @@ -227,8 +268,14 @@ def concat_input_feature(self, prompt_input_feature: InputFeature, new_tokens: L messages: List[Message] = result.get('messages') if messages is not None: response_text = self.tokenizer.decode(new_tokens, skip_special_tokens=True) - parsed = self.parse_tool_call(response_text) or [] - content_text = (self.clean_tool_call(response_text) if parsed else response_text) + if tool_calls is None: + parsed = self.parse_tool_call(response_text) or [] + content_text = (self.clean_tool_call(response_text) if parsed else response_text) + else: + # Structured calls arrived beside the text, so the text carries no + # markup to strip. + parsed = list(tool_calls) + content_text = response_text asst_msg = Message(role='assistant', content=content_text) if parsed: asst_msg['tool_calls'] = parsed @@ -236,6 +283,26 @@ def concat_input_feature(self, prompt_input_feature: InputFeature, new_tokens: L result['messages'] = messages return result + @staticmethod + def _prefix_completion_mask(feature: InputFeature, labels: List[int]) -> List[int]: + """The prefix's ``completion_mask``, in input order, materialised if absent. + + A feature encoded before this field existed records no provenance, and for + those the trainable positions *were* exactly the policy's own -- deriving the + mask from ``labels`` therefore leaves old and new trajectories equivalent. + """ + mask = feature.get('completion_mask') + if mask is None: + mask = [0 if label == -100 else 1 for label in labels] + else: + mask = list(mask) + mask = mask[-1:] + mask[:-1] # roll to input order, exactly as labels + expected = len(feature['input_ids']) + if len(mask) != expected: + raise ValueError(f'prefix completion_mask has {len(mask)} entries for {expected} ' + f'input_ids; appending would misalign every position after it.') + return mask + def _add_default_system(self, trajectory: Trajectory) -> List[Trajectory]: if self.use_chat_template and self.default_system: if trajectory['messages'][0]['role'] == 'user': @@ -274,27 +341,25 @@ def _extract_reasoning_content(messages: list[Message]) -> List[Message]: return [trajectory] def _truncate_feature(self, feature: InputFeature, strategy: str) -> InputFeature: - """Truncate input_ids and labels in a single InputFeature.""" + """Truncate the sequence-aligned fields of a single InputFeature.""" length = len(feature['input_ids']) if length <= self.max_length: return feature if strategy == 'raise': raise ValueError(f'Input length {length} exceeds max_length {self.max_length}') - result = dict(feature) if strategy == 'left': - result['input_ids'] = result['input_ids'][-self.max_length:] - if 'labels' in result: - result['labels'] = result['labels'][-self.max_length:] - if 'mm_token_type_ids' in result: - result['mm_token_type_ids'] = result['mm_token_type_ids'][..., -self.max_length:] + keep = slice(-self.max_length, None) elif strategy == 'right': - result['input_ids'] = result['input_ids'][:self.max_length] - if 'labels' in result: - result['labels'] = result['labels'][:self.max_length] - if 'mm_token_type_ids' in result: - result['mm_token_type_ids'] = result['mm_token_type_ids'][..., :self.max_length] + keep = slice(None, self.max_length) else: raise ValueError(f'Unsupported truncation_strategy={strategy!r}.') + result = dict(feature) + result['input_ids'] = result['input_ids'][keep] + for key in _SEQUENCE_ALIGNED_FIELDS: + if key in result: + result[key] = result[key][keep] + if 'mm_token_type_ids' in result: + result['mm_token_type_ids'] = result['mm_token_type_ids'][..., keep] return InputFeature(**result) def set_mm_position_ids(self, input_feature: InputFeature): @@ -321,8 +386,9 @@ def _check_max_length(self, input_feature: InputFeature) -> List[InputFeature]: end = min(start + self.max_length, len(input_feature['input_ids'])) feat = dict(input_feature) feat['input_ids'] = feat['input_ids'][start:end] - if 'labels' in feat: - feat['labels'] = feat['labels'][start:end] + for key in _SEQUENCE_ALIGNED_FIELDS: + if key in feat: + feat[key] = feat[key][start:end] if 'mm_token_type_ids' in feat: feat['mm_token_type_ids'] = feat['mm_token_type_ids'][..., start:end] results.append(InputFeature(**feat)) @@ -350,6 +416,10 @@ def _roll_labels(self, input_feature: InputFeature) -> List[InputFeature]: if 'input_ids' not in input_feature: return [input_feature] input_feature['labels'] = np.roll(input_feature['labels'], -1, axis=-1) + if 'completion_mask' in input_feature: + # The mask answers "is there a log-prob for this position's target", so it + # lives on the labels' index space and has to follow the same roll. + input_feature['completion_mask'] = np.roll(input_feature['completion_mask'], -1, axis=-1) return [input_feature] def _process_mm_messages(self, messages: List, images: List, videos: List, audios: List) -> List: @@ -524,6 +594,34 @@ def _build_standard_messages(self, trajectory: Trajectory) -> List[Trajectory]: message['content'] = c[0]['text'] if c else '' return [trajectory] + @staticmethod + def decode_tool_calls(message: Dict[str, Any]) -> Dict[str, Any]: + """Return ``message`` with ``tool_calls`` in the shape a chat template renders. + + OpenAI-shaped calls carry ``function.arguments`` as a JSON string, and an + Arrow round-trip can turn the whole list into one; templates index them as + objects. Arguments that will not parse become ``{}`` rather than reaching + Jinja as a string it would render verbatim. The message is returned + untouched when it carries no calls. + """ + tool_calls = message.get('tool_calls') + if isinstance(tool_calls, str): + tool_calls = json.loads(tool_calls) if tool_calls else [] + elif not tool_calls: + return message + decoded = [] + for tool_call in tool_calls: + fn = tool_call['function'] + args = fn['arguments'] + if isinstance(args, dict): + value = args + elif isinstance(args, str): + value = json.loads(args) if args.strip() else {} + else: + value = {} + decoded.append({**tool_call, 'function': {**fn, 'arguments': value}}) + return {**message, 'tool_calls': decoded} + def _apply_chat_template(self, trajectory: Trajectory, add_generation_prompt: bool = False, **kwargs): messages = [dict(message) for message in trajectory['messages']] # Arrow serialization may pad content blocks with null keys (e.g. 'image': None @@ -536,25 +634,7 @@ def _apply_chat_template(self, trajectory: Trajectory, add_generation_prompt: bo k: v for k, v in b.items() if v is not None } for b in msg['content'] if isinstance(b, dict)] - for msg in messages: - tcs = msg.get('tool_calls') - if isinstance(tcs, str): - tcs = json.loads(tcs) if tcs else [] - msg['tool_calls'] = tcs - if not tcs: - continue - new_tcs = [] - for tc in tcs: - fn = tc['function'] - args = fn['arguments'] - if isinstance(args, dict): - decoded = args - elif isinstance(args, str): - decoded = json.loads(args) if args.strip() else {} - else: - decoded = {} - new_tcs.append({**tc, 'function': {**fn, 'arguments': decoded}}) - msg['tool_calls'] = new_tcs + messages = [self.decode_tool_calls(msg) for msg in messages] # ``tool_calls`` / ``tools`` are already OpenAI-shaped (see # :mod:`twinkle.data_format.message`); pass them through verbatim. tools = list(trajectory.get('tools') or []) diff --git a/src/twinkle_agentic/async_rl/data_plane.py b/src/twinkle_agentic/async_rl/data_plane.py index 641947366..da9015f0c 100644 --- a/src/twinkle_agentic/async_rl/data_plane.py +++ b/src/twinkle_agentic/async_rl/data_plane.py @@ -71,10 +71,20 @@ def _require_rollout_logprobs(sample: dict[str, Any], *, sample_key: str) -> lis values.append(float(value)) labels = sample.get('labels') if labels is not None: - trainable_tokens = sum(1 for label in labels if label != -100) - if len(values) != trainable_tokens: - raise ValueError(f'rollout sample {sample_key!r} logprobs length must match trainable labels: ' - f'{len(values)} != {trainable_tokens}') + # Only policy-generated tokens carry a sampling log-prob. A turn written by + # an API or a human is trainable yet has none, and is marked + # completion_mask=0 -- the same basis GRPOLoss restricts itself to. + completion_mask = sample.get('completion_mask') + if completion_mask is None: + expected = sum(1 for label in labels if label != -100) + elif len(completion_mask) != len(labels): + raise ValueError(f'rollout sample {sample_key!r} completion_mask length must match labels: ' + f'{len(completion_mask)} != {len(labels)}') + else: + expected = sum(1 for label, flag in zip(labels, completion_mask) if label != -100 and flag) + if len(values) != expected: + raise ValueError(f'rollout sample {sample_key!r} logprobs length must match policy-generated tokens: ' + f'{len(values)} != {expected}') return values diff --git a/src/twinkle_agentic/challenger/base.py b/src/twinkle_agentic/challenger/base.py index 3dcaae674..69ab74673 100644 --- a/src/twinkle_agentic/challenger/base.py +++ b/src/twinkle_agentic/challenger/base.py @@ -11,8 +11,7 @@ rollouts in :mod:`twinkle_agentic.rollout` have that signature already, so a challenger can explore *with tools* -- running code, reading files -- while it invents, over a local sampler or over an HTTP endpoint alike. - :func:`twinkle_agentic.rollout.build_rollout` picks the right one for the - backend at hand. + :class:`twinkle_agentic.rollout.MultiTurnRollout` accepts either backend. * **what counts as a keeper** -- subclasses decide, in :meth:`Challenger.build`. * **how hard is hard enough** -- optional. Ask for ``solver_rollouts`` attempts per candidate and only tasks the model solves *sometimes* are kept: a task every @@ -44,8 +43,7 @@ __all__ = ['Challenger', 'Explorer', 'KeywordPrompts', 'PromptSet'] # A batch of trajectories in, the same trajectories with the model's reply -# appended out. Both MultiTurnRollout and APIMultiTurnRollout satisfy this -# as-is; build_rollout() returns whichever fits the backend. Both also accept a +# appended out. MultiTurnRollout accepts either backend and also accepts a # per-call ``sampling_params=`` keyword, which is how the difficulty stage asks # for its own temperature and length budget without a second explorer. Explorer = Callable[[List[Trajectory]], List[Trajectory]] @@ -141,9 +139,9 @@ class Challenger(ABC): Args: explorer: takes a batch of trajectories and returns them with the - model's reply appended -- a rollout from - :func:`twinkle_agentic.rollout.build_rollout`, over a local sampler - or over an API endpoint. + model's reply appended -- typically a + :class:`twinkle_agentic.rollout.MultiTurnRollout` over a local + sampler or an API endpoint. system: system prompt handed to the model. It carries the output contract, which is why ``build`` -- the code that reads that output back -- lives in the same subclass. diff --git a/src/twinkle_agentic/challenger/new/__init__.py b/src/twinkle_agentic/challenger/new/__init__.py new file mode 100644 index 000000000..0719b04e5 --- /dev/null +++ b/src/twinkle_agentic/challenger/new/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .agentic import AgenticChallenger, parse_problem_statement +from .base import Challenger +from .keyword import KEYWORD_MAX_LEN, KeywordGenerator + +__all__ = [ + 'AgenticChallenger', + 'Challenger', + 'KEYWORD_MAX_LEN', + 'KeywordGenerator', + 'parse_problem_statement', +] diff --git a/src/twinkle_agentic/challenger/new/agentic.py b/src/twinkle_agentic/challenger/new/agentic.py new file mode 100644 index 000000000..ca41486a7 --- /dev/null +++ b/src/twinkle_agentic/challenger/new/agentic.py @@ -0,0 +1,511 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Agentic challenger: act in a sandbox, verify the result, then describe it.""" +import math +import random +import re +import uuid +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple + +from twinkle.data_format import SamplingParams, Trajectory, attach_user_data, user_data_get +from twinkle.data_format.sampling import SampledSequence, SampleResponse +from twinkle.utils import get_logger +from twinkle_agentic.envs import Env +from twinkle_agentic.protocol.base import API +from twinkle_agentic.rollout import APISampler, MultiTurnRollout +from twinkle_agentic.summarizer import Summarizer +from twinkle_agentic.utils.code_utils import parse_fenced_code, strip_reasoning +from twinkle_agentic.utils.message_utils import assistant_text, msg_content_text, normalize_tool_calls +from .base import Challenger, _parallel +from .keyword import KeywordGenerator +from .recorder import RolloutRecorder + +__all__ = ['AgenticChallenger', 'parse_problem_statement'] + +logger = get_logger() + +_FENCED_BLOCK_RE = re.compile(r'```[^\r\n]*\r?\n(.*?)```', re.S) + + +def parse_problem_statement(text: str) -> Optional[str]: + """Return the statement after removing reasoning and one outer fence.""" + body = strip_reasoning(text).strip() + whole = _FENCED_BLOCK_RE.fullmatch(body) + if whole: + body = whole.group(1).strip() + return body or None + + +def _sample_one(sampler: Any, input_feature: Dict[str, Any], sampling_params: Optional[SamplingParams], + adapter_kwargs: Dict[str, Any]) -> SampledSequence: + responses = sampler.sample([input_feature], sampling_params=sampling_params, **adapter_kwargs) + if not isinstance(responses, list): + raise TypeError(f'expected List[SampleResponse] from sampler.sample, got ' + f'{type(responses).__name__}') + if len(responses) != 1: + raise RuntimeError(f'sampler returned {len(responses)} responses for a single request; ' + 'expected exactly one') + response = responses[0] + if not isinstance(response, SampleResponse): + raise TypeError(f'expected SampleResponse from sampler.sample, got ' + f'{type(response).__name__}') + if len(response.sequences) != 1: + raise RuntimeError(f'SampleResponse contains {len(response.sequences)} sequences; ' + 'expected exactly one') + sequence = response.sequences[0] + if not isinstance(sequence, SampledSequence): + raise TypeError(f'expected SampledSequence, got {type(sequence).__name__}') + return sequence + + +def _api_followup_response( + sampler: Any, + api: Optional[APISampler], + sampling_params: Optional[SamplingParams], + *, + input_feature: Dict[str, Any], + adapter_kwargs: Dict[str, Any], + followups: int, + **kwargs: Any, +) -> SampledSequence: + """Use the API for appended stages and the primary backend otherwise.""" + if followups: + if api is None: + raise ValueError('use_api=True requires an API backend') + return api(input_feature, sampling_params, **adapter_kwargs) + if sampler is not None: + return _sample_one(sampler, input_feature, sampling_params, adapter_kwargs) + if api is not None: + return api(input_feature, sampling_params, **adapter_kwargs) + raise ValueError('AgenticChallenger has neither a sampler nor an API backend') + + +@dataclass +class _ProposalResult: + trajectory: Trajectory + group_id: str = '' + task: Optional[Trajectory] = None + reason: str = '' + detail: str = '' + outcome: str = '' + n_pass: Optional[int] = None + reward: float = 0.0 + + +class AgenticChallenger(Challenger): + """Invent tool-using tasks by doing, checking, and describing them. + + ``backend`` drives exploration and solver attempts. When ``use_api`` is true, + ``api`` generates only the appended check-script and problem-statement turns; + those turns retain the masking semantics selected by ``api_appended_as`` in + ``rollout_kwargs``. + """ + + _system = ('You invent tasks for another agent to solve. You have a sandbox and ' + 'tools. Work in it first: build something real, then you will be asked ' + 'to verify it and to describe it.') + _from_scratch = ('Choose a task worth doing in this sandbox and do it now, using ' + 'your tools. Do not describe it yet.') + _from_keywords = ('Choose a task around these topics and do it now, using your ' + 'tools. Do not describe it yet.\n\nTopics: {keywords}') + _from_seed = ('Here is an earlier task:\n\n{seed}\n\nDo something in the same ' + 'spirit but different, using your tools now. Do not describe it yet.') + _from_seed_keywords = ('Here is an earlier task:\n\n{seed}\n\nDo something in the same ' + 'spirit but different, may be more complex and interesting and meaningful, ' + 'around these topics, using your tools now. Do not describe it yet.\n\n' + 'Topics: {keywords}') + _check_followup = ('Stop working. This is the workspace you produced:\n\n{final_state}\n\n' + 'Write a {language} script that verifies this end state, as a fenced ' + '{language} code block and nothing else. It must exit with a non-zero status ' + 'if the work was not done. Check what can be read out of the files -- their ' + 'structure and the values inside them. NEVER check a file size in bytes, a ' + 'checksum, or the full source text of a script: correct solutions differ ' + 'there, and such a check only its own author can pass.') + _check_retry_followup = ('Your check script did not pass:\n\n{error}\n\nThe workspace is:\n\n' + '{final_state}\n\nReturn a corrected script as a fenced {language} code block ' + 'and nothing else.') + _check_parse_error = ('Could not read a check script from your reply: it was not a ' + 'fenced {language} code block. Do not wrap it in a tool call and ' + 'do not add prose -- return ONLY a fenced {language} code block.') + _problem_followup = ('Now write the task statement: what someone starting from an empty workspace ' + 'would have to be told to produce what you produced, and nothing about how you ' + 'did it. Name the files to create and quote any input data verbatim. Do not ' + 'reveal values your check script computes. Reply with the statement only.') + + + def __init__( + self, + backend: Any, + *, + api: Optional[Any] = None, + use_api: bool = False, + keyword_generator: Optional[KeywordGenerator] = None, + trajectory_seed: Optional[List[Trajectory]] = None, + summarizer: Optional[Summarizer] = None, + system_prompt: Optional[str] = None, + from_scratch_prompt: Optional[str] = None, + from_keywords_prompt: Optional[str] = None, + from_seed_prompt: Optional[str] = None, + from_seed_keywords_prompt: Optional[str] = None, + check_followup_prompt: Optional[str] = None, + check_retry_followup_prompt: Optional[str] = None, + check_parse_error_prompt: Optional[str] = None, + problem_followup_prompt: Optional[str] = None, + check_retries: int = 1, + problem_max_chars: int = 8192, + check_language: str = 'python', + parse_check_fn: Optional[Callable[[str], Optional[str]]] = None, + pass_rate_target: float = 0.2, + envs: Sequence[Env] = (), + num_challenger_rollouts: int = 8, + num_solver_rollouts: int = 8, + pass_band: Tuple[float, float] = (1.0, 7.0), + pass_rate_width: float = 0.3, + max_empty_rounds: int = 0, + followup_params: Optional[SamplingParams] = None, + checker: Optional[Callable[[Trajectory], bool]] = None, + save_dir: Optional[str] = None, + save_failed_rollouts: bool = True, + **rollout_kwargs: Any, + ): + super().__init__( + envs=envs, + num_challenger_rollouts=num_challenger_rollouts, + num_solver_rollouts=num_solver_rollouts, + pass_band=pass_band, + max_empty_rounds=max_empty_rounds, + ) + if check_retries < 0: + raise ValueError(f'check_retries must be >= 0, got {check_retries}') + if problem_max_chars <= 0: + raise ValueError(f'problem_max_chars must be positive, got {problem_max_chars}') + if not check_language.strip(): + raise ValueError('check_language must not be empty') + if not 0 <= pass_rate_target <= 1: + raise ValueError(f'pass_rate_target must be in [0, 1], got {pass_rate_target}') + if pass_rate_width <= 0: + raise ValueError(f'pass_rate_width must be positive, got {pass_rate_width}') + if use_api and rollout_kwargs.get('response_callback') is not None: + raise ValueError('use_api=True cannot be combined with response_callback') + backend_is_api = isinstance(backend, (API, APISampler)) + if use_api and api is None and not backend_is_api: + raise ValueError('use_api=True requires api= when backend is a sampler') + self.keyword_generator = keyword_generator + self.trajectory_seed = list(trajectory_seed or ()) + self.summarizer = summarizer + self._system = self._system if system_prompt is None else system_prompt + self._from_scratch = self._from_scratch if from_scratch_prompt is None else from_scratch_prompt + self._from_keywords = self._from_keywords if from_keywords_prompt is None else from_keywords_prompt + self._from_seed = self._from_seed if from_seed_prompt is None else from_seed_prompt + self._from_seed_keywords = (self._from_seed_keywords if from_seed_keywords_prompt is None else + from_seed_keywords_prompt) + self._check_followup = self._check_followup if check_followup_prompt is None else check_followup_prompt + self._check_retry_followup = (self._check_retry_followup if check_retry_followup_prompt is None else + check_retry_followup_prompt) + self._check_parse_error = (self._check_parse_error if check_parse_error_prompt is None else + check_parse_error_prompt) + self._problem_followup = (self._problem_followup if problem_followup_prompt is None else + problem_followup_prompt) + self._check_retries = check_retries + self._problem_max_chars = problem_max_chars + self._check_language = check_language.strip().lower() + self._parse_check_fn = parse_check_fn + self._pass_rate_target = pass_rate_target + self._pass_rate_width = pass_rate_width + self.checker = checker + self.followup_params = followup_params + self.rng = random.Random() + self.use_api = use_api + self.save_failed_rollouts = save_failed_rollouts + self._recorder = RolloutRecorder(save_dir) if save_dir else None + self._round_proposals: List[_ProposalResult] = [] + self._backend = backend + self._rollout_kwargs = dict(rollout_kwargs) + if api is not None: + self._rollout_kwargs['api'] = api + if use_api: + self._rollout_kwargs['response_callback'] = _api_followup_response + self._rollout: Optional[MultiTurnRollout] = None + self._tool_schemas = self.env().tools() or None + + def _rollout_instance(self) -> MultiTurnRollout: + if self._rollout is None: + self._rollout = MultiTurnRollout(self._backend, **self._rollout_kwargs) + return self._rollout + + def _tool_manager(self, slot: int) -> Optional[Any]: + env = self.env(slot) + return env.tool_manager() if env.tools() else None + + def _summary(self, trajectory: Trajectory) -> str: + turns: List[str] = [] + for message in trajectory.get('messages') or []: + if not isinstance(message, dict): + continue + role = message.get('role') or '' + if role == 'system': + continue + parts = [msg_content_text(message).strip()] + for call in normalize_tool_calls(message) or (): + fn = call.get('function') or {} + if isinstance(fn, dict) and fn.get('name'): + parts.append(f"calls {fn['name']}({fn.get('arguments') or ''})") + body = '\n'.join(part for part in parts if part) + if body: + turns.append(f'{role}: {body}') + text = '\n'.join(turns) + if not text: + return '' + return self.summarizer(text) if self.summarizer is not None else text + + def _build_challenge_prompt(self) -> Optional[Trajectory]: + keywords: List[str] = [] + if self.keyword_generator is not None: + groups = self.keyword_generator.get_keywords(1) + if not groups: + return None + keywords = groups[0] + seed = '' + if self.trajectory_seed: + seed = self._summary(self.rng.choice(self.trajectory_seed)) + block = ', '.join(keywords) + if seed and keywords: + user = self._from_seed_keywords.format(seed=seed, keywords=block) + elif seed: + user = self._from_seed.format(seed=seed) + elif keywords: + user = self._from_keywords.format(keywords=block) + else: + user = self._from_scratch + prompt: Trajectory = { + 'messages': [ + { + 'role': 'system', + 'content': self._system + }, + { + 'role': 'user', + 'content': user + }, + ], + } + if self._tool_schemas: + prompt['tools'] = self._tool_schemas + return attach_user_data(prompt, keywords=keywords, seeded=bool(seed)) + + def _explore(self, prompt: Trajectory) -> List[Trajectory]: + group_id = uuid.uuid4().hex + proposals: List[_ProposalResult] = [] + remaining = self.num_challenger_rollouts + while remaining > 0: + wave = min(self.n_slots, remaining) + proposals.extend(_parallel(lambda slot: self._run_episode(prompt, slot), wave)) + remaining -= wave + for proposal in proposals: + proposal.group_id = group_id + self._round_proposals = proposals + return [proposal.task for proposal in proposals if proposal.task is not None] + + def _run_episode(self, prompt: Trajectory, slot: int) -> _ProposalResult: + self.env(slot).clear() + state: Dict[str, Any] = {'slot': slot} + kwargs: Dict[str, Any] = { + 'followup_fn': lambda trajectory, n_before: self._followup(state, trajectory, n_before), + } + manager = self._tool_manager(slot) + if manager is not None: + kwargs['tool_manager'] = manager + explored = self._rollout_instance()([prompt], **kwargs) + if not explored: + self._reject(state, 'rollout_no_output') + return _ProposalResult(dict(prompt), reason='rollout_no_output') + trajectory = explored[0] + task = self._build_query(state, trajectory) + reason, detail = state.get('reject', ('', '')) + return _ProposalResult(trajectory, task=task, reason=reason, detail=detail) + + def _followup(self, state: Dict[str, Any], trajectory: Trajectory, + n_before: int) -> Optional[Tuple[str, Optional[SamplingParams]]]: + if state.get('checked'): + return None + reply = None if n_before == 0 else assistant_text(trajectory) + followup = self._build_test_case(state, reply) + if followup is None: + return None + return followup, self.followup_params + + def _build_test_case(self, state: Dict[str, Any], reply: Optional[str]) -> Optional[str]: + slot = state['slot'] + if reply is None: + snapshot, error = self.env(slot).snapshot() + state['snapshot'] = snapshot + if not snapshot.strip(): + state['reject'] = ('snapshot_unavailable' if error else 'empty_workspace', error) + return None + return self._check_followup.format(final_state=snapshot, language=self._check_language) + + attempt = state.get('check_attempts', 0) + 1 + state['check_attempts'] = attempt + script = (self._parse_check_fn(reply) if self._parse_check_fn is not None else + parse_fenced_code(reply, language_tags=None)) + if script is None: + if attempt <= self._check_retries: + return self._check_retry_followup.format( + error=self._check_parse_error.format(language=self._check_language), + final_state=state.get('snapshot', ''), + language=self._check_language, + ) + state['reject'] = ('check_parse_fail', reply) + return None + state['script'] = script + exit_code, output = self.env(slot).run_script(script, interpreter=self._check_language) + if exit_code == 0: + state['checked'] = True + return self._problem_followup + after = self.env(slot).snapshot()[0] + state.setdefault('attempts', []).append(f'--- attempt {attempt}: exit {exit_code} ---\n{output}\n' + f'--- check script ---\n{script}') + if attempt <= self._check_retries: + return self._check_retry_followup.format( + error=output, + final_state=after or state.get('snapshot', ''), + language=self._check_language, + ) + state['reject'] = ('check_run_fail', '\n'.join(state['attempts'])) + return None + + def _build_query(self, state: Dict[str, Any], explored: Trajectory) -> Optional[Trajectory]: + if state.get('reject'): + return self._reject(state, *state['reject']) + if not state.get('checked'): + return self._reject( + state, + 'episode_cut_short', + f"stop_reason={explored.get('stop_reason')} " + f"truncated={bool(explored.get('truncated'))} " + f"turns={explored.get('turns')}", + ) + statement = parse_problem_statement(assistant_text(explored)) + if statement is None: + return self._reject(state, 'problem_parse_fail') + if len(statement) > self._problem_max_chars: + return self._reject(state, 'too_long', f'{len(statement)} chars') + task: Trajectory = attach_user_data( + {'messages': [{ + 'role': 'user', + 'content': statement + }]}, + check_script=state['script'], + keywords=user_data_get(explored.get('user_data'), 'keywords', []), + seeded=user_data_get(explored.get('user_data'), 'seeded', False), + ) + if self.checker is not None and not self.checker(task): + return self._reject(state, 'rejected_by_checker') + return task + + def _reject(self, state: Dict[str, Any], reason: str, detail: str = '') -> Optional[Trajectory]: + state['reject'] = (reason, detail) + logger.info(f'[{type(self).__name__}] rejected: {reason}' + f"{f' -- {detail[:400]}' if detail else ''}") + return None + + def _solver_prompt(self, task: Trajectory) -> Trajectory: + prompt: Trajectory = {'messages': [dict(message) for message in task.get('messages') or []]} + if self._tool_schemas: + prompt['tools'] = self._tool_schemas + return prompt + + def _judge(self, task: Trajectory, slot: int) -> bool: + script = user_data_get(task.get('user_data'), 'check_script', '') + if not script: + return False + return self.env(slot).run_script(script, interpreter=self._check_language)[0] == 0 + + def challenger_reward(self, n_pass: Optional[int]) -> float: + """Reward tasks near the target solver pass rate; unmeasured failures score zero.""" + if n_pass is None or not self.num_solver_rollouts or n_pass <= 0: + return 0.0 + gap = n_pass / self.num_solver_rollouts - self._pass_rate_target + variance = 2.0 * self._pass_rate_width**2 + return math.exp(-(gap * gap) / variance) + + def _record_proposals(self) -> None: + proposals, self._round_proposals = self._round_proposals, [] + if self._recorder is None: + return + for index, proposal in enumerate(proposals): + if proposal.task is None and not self.save_failed_rollouts: + continue + trajectory = dict(proposal.trajectory) + trajectory['rewards'] = proposal.reward + task_data = proposal.task.get('user_data') if proposal.task is not None else None + statement = '' + if proposal.task is not None: + statement = next((message.get('content', '') for message in proposal.task.get('messages') or [] + if isinstance(message, dict) and message.get('role') == 'user'), '') + self._recorder.write( + trajectory, + side='propose', + group_id=proposal.group_id, + proposal_index=index, + outcome=proposal.outcome or ('rejected' if proposal.reason else 'kept'), + reason=proposal.reason, + detail=proposal.detail, + reward=proposal.reward, + n_pass=proposal.n_pass, + n_rollouts=(self.num_solver_rollouts if proposal.n_pass is not None else None), + pass_rate=(proposal.n_pass / self.num_solver_rollouts + if proposal.n_pass is not None and self.num_solver_rollouts else None), + statement=statement, + check_script=user_data_get(task_data, 'check_script', ''), + keywords=user_data_get(proposal.trajectory.get('user_data'), 'keywords', []), + seeded=user_data_get(proposal.trajectory.get('user_data'), 'seeded', False), + ) + + def _filter_difficulty(self, tasks: List[Trajectory]) -> List[Trajectory]: + successful = [proposal for proposal in self._round_proposals if proposal.task is not None] + if len(successful) != len(tasks): + raise RuntimeError('proposal/task alignment failed before difficulty filtering') + if not tasks or not self.num_solver_rollouts: + for proposal in successful: + proposal.outcome = 'kept' + self._record_proposals() + return tasks + + passes = [0] * len(tasks) + plan = [i for i in range(len(tasks)) for _ in range(self.num_solver_rollouts)] + rollout = self._rollout_instance() + for start in range(0, len(plan), self.n_slots): + wave = plan[start:start + self.n_slots] + _parallel(lambda slot: self.env(slot).clear(), len(wave)) + prompts = [self._solver_prompt(tasks[i]) for i in wave] + kwargs: Dict[str, Any] = {} + managers = [self._tool_manager(slot) for slot in range(len(wave))] + if any(manager is not None for manager in managers): + kwargs['tool_manager'] = managers + attempts = rollout(prompts, **kwargs) + if len(attempts) != len(prompts): + raise RuntimeError(f'rollout returned {len(attempts)} attempts for ' + f'{len(prompts)} prompts; expected one per prompt') + verdicts = _parallel(lambda slot: self._judge(tasks[wave[slot]], slot), len(wave)) + for slot, passed in enumerate(verdicts): + if passed: + passes[wave[slot]] += 1 + + low, high = self.pass_band + measured = [ + attach_user_data(task, n_pass=n_pass, n_rollouts=self.num_solver_rollouts) + for task, n_pass in zip(tasks, passes) + ] + kept: List[Trajectory] = [] + for proposal, task, n_pass in zip(successful, measured, passes): + proposal.task = task + proposal.n_pass = n_pass + proposal.reward = self.challenger_reward(n_pass) + if low <= n_pass <= high: + proposal.outcome = 'kept' + kept.append(task) + else: + proposal.outcome = 'outside_band' + self._record_proposals() + return kept diff --git a/src/twinkle_agentic/challenger/new/base.py b/src/twinkle_agentic/challenger/new/base.py new file mode 100644 index 000000000..1c819533b --- /dev/null +++ b/src/twinkle_agentic/challenger/new/base.py @@ -0,0 +1,132 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Reusable lifecycle for task challengers.""" +from abc import ABC, abstractmethod +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Callable, Iterator, List, Optional, Sequence, Tuple + +from twinkle.data_format import Trajectory +from twinkle.utils import get_logger +from twinkle_agentic.envs import Env + +logger = get_logger() + +__all__ = ['Challenger'] + + +def _parallel(fn: Callable[[int], Any], count: int) -> List[Any]: + """Run ``fn`` over ``range(count)`` concurrently, preserving order.""" + if count <= 1: + return [fn(i) for i in range(count)] + out: List[Any] = [None] * count + with ThreadPoolExecutor(max_workers=count) as pool: + futures = {pool.submit(fn, i): i for i in range(count)} + for future, i in futures.items(): + out[i] = future.result() + return out + + +class Challenger(ABC): + """Common batching and environment lifecycle for task challengers. + + Subclasses define how a round builds its prompt, explores it, and measures + candidate difficulty. One environment is owned by one concurrent job for the + complete lifetime of that job. + """ + + def __init__( + self, + *, + envs: Sequence[Env], + num_challenger_rollouts: int = 8, + num_solver_rollouts: int = 8, + pass_band: Tuple[float, float] = (1.0, 7.0), + max_empty_rounds: int = 0, + ): + if not envs: + raise ValueError('envs is empty: a challenger needs a workspace to act in and grade') + if num_challenger_rollouts < 1: + raise ValueError(f'num_challenger_rollouts must be >= 1, got ' + f'{num_challenger_rollouts}') + if num_solver_rollouts < 0: + raise ValueError(f'num_solver_rollouts must be >= 0, got {num_solver_rollouts}') + if max_empty_rounds < 0: + raise ValueError(f'max_empty_rounds must be >= 0, got {max_empty_rounds}') + if num_solver_rollouts: + if len(pass_band) != 2: + raise ValueError(f'pass_band is (low, high) in attempt counts, got {pass_band}') + low, high = pass_band + if not 0 <= low <= high <= num_solver_rollouts: + raise ValueError(f'pass_band must satisfy 0 <= low <= high <= num_solver_rollouts, got ' + f'{pass_band} against num_solver_rollouts={num_solver_rollouts}') + self.envs = list(envs) + self.num_challenger_rollouts = num_challenger_rollouts + self.num_solver_rollouts = num_solver_rollouts + self.pass_band = pass_band + self.max_empty_rounds = max_empty_rounds + self.n_proposed = 0 + self.n_kept = 0 + + @property + def n_slots(self) -> int: + """How many jobs may run at once: one per environment.""" + return len(self.envs) + + def env(self, slot: int = 0) -> Env: + """Return the current environment for ``slot``.""" + return self.envs[slot] + + @abstractmethod + def _build_challenge_prompt(self) -> Optional[Trajectory]: + """Build one round's shared prompt, or return None when exhausted.""" + + @abstractmethod + def _explore(self, prompt: Trajectory) -> List[Trajectory]: + """Generate and validate candidates from one shared prompt.""" + + @abstractmethod + def _filter_difficulty(self, tasks: List[Trajectory]) -> List[Trajectory]: + """Measure candidate difficulty and return the accepted tasks.""" + + def __call__(self, batch_size: int, total: Optional[int] = None) -> Iterator[List[Trajectory]]: + """Yield finished tasks in batches.""" + if batch_size <= 0: + raise ValueError(f'batch_size must be positive, got {batch_size}') + pending: List[Trajectory] = [] + produced = 0 + empty_rounds = 0 + while total is None or produced < total: + want = batch_size if total is None else min(batch_size, total - produced) + while len(pending) < want: + kept = self._round() + if kept is None: + if pending: + yield pending + return + if kept: + empty_rounds = 0 + pending.extend(kept) + continue + empty_rounds += 1 + if self.max_empty_rounds and empty_rounds >= self.max_empty_rounds: + logger.warning(f'[{type(self).__name__}] stopped after {empty_rounds} ' + 'consecutive rounds without a usable task') + if pending: + yield pending + return + yield pending[:want] + produced += want + pending = pending[want:] + + def _round(self) -> Optional[List[Trajectory]]: + """Run one proposal group; None means the source is exhausted.""" + prompt = self._build_challenge_prompt() + if prompt is None: + return None + verified = self._explore(prompt) + kept = self._filter_difficulty(verified) + self.n_proposed += self.num_challenger_rollouts + self.n_kept += len(kept) + logger.info(f'[{type(self).__name__}] {self.num_challenger_rollouts} episodes, ' + f'{len(verified)} verified, {len(kept)} in band ' + f'(cumulative {self.n_kept}/{self.n_proposed})') + return kept diff --git a/src/twinkle_agentic/challenger/new/keyword.py b/src/twinkle_agentic/challenger/new/keyword.py index 0d2c66bad..09dee2bad 100644 --- a/src/twinkle_agentic/challenger/new/keyword.py +++ b/src/twinkle_agentic/challenger/new/keyword.py @@ -1,12 +1,290 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Keywords per direction: generate, de-duplicate, store, read back. +A keyword is a *topic* to build a task around, not a task statement, which is +why over-length replies are dropped rather than stored. +""" +import json +import os +from typing import Any, Dict, List, Optional, Sequence, Set, Tuple -from typing import Optional +from twinkle.data_format import SamplingParams, Trajectory +from twinkle.utils import get_logger +from twinkle_agentic.rollout import MultiTurnRollout +from twinkle_agentic.utils.code_utils import strip_reasoning +from twinkle_agentic.utils.message_utils import assistant_text + +logger = get_logger() + +__all__ = ['KEYWORD_MAX_LEN', 'KeywordGenerator'] + +KEYWORD_MAX_LEN = 60 class KeywordGenerator: + """Keyword combinations drawn from one list per direction. + + ``keywords_group_size`` of the directions are active at a time and one draw + takes a keyword from each. What a draw spends is the *combination*, not the + keywords: a group only has to differ from every group already handed out, so + three directions holding ``num_keywords`` each are worth their product in + draws rather than just ``num_keywords``. A direction that has produced + ``num_keywords`` is retired and the next unused one takes its slot, which is + why more directions than a group needs is the normal case. De-duplication of + the keywords themselves is flat, so a keyword one direction produced is never + handed to another. + + Args: + query: what the keywords have to satisfy -- one entry per direction. Must + be at least ``keywords_group_size`` of them. + backend: an API client or a sampler; driven through ``MultiTurnRollout``. + path: JSONL cache. Empty means in-memory only. + num_keywords: a direction's budget; past it, it is retired. + keywords_group_size: how many keywords one draw combines. + system_prompt: overrides the built-in one. + recycle: once every direction is spent, hand out the same combinations + again instead of returning None. + rollout_kwargs: passed to ``MultiTurnRollout``. ``template`` is required; + API request options belong in ``api_kwargs``. + """ + + # How many known keywords the 'do not repeat these' line may quote. A cap in + # both directions: too few and a second round says the same things again, too + # many and the model runs out of room to obey. + _avoid_max = 100 + _avoid_lead = '\nDo NOT repeat any of these: ' + + # A default prompt to use to generate the keywords + _default_prompt = ( + 'You brainstorm topics. Reply with a JSON array of short noun phrases ' + f'(at most {KEYWORD_MAX_LEN} characters each) and nothing else. ' + 'Each phrase names a subject to build a task around, never a task statement.') + + _user_prompt = 'Give {k} distinct topics that satisfy:\n{query}' + + def __init__( + self, + query: Sequence[str], + backend: Any, + path: str, + *, + num_keywords: int = 64, + keywords_group_size: int = 3, + system_prompt: Optional[str] = None, + sampling_params: Optional[SamplingParams] = None, + recycle: bool = False, + **rollout_kwargs: Any, + ): + self.query = list(query) + if keywords_group_size < 1: + raise ValueError(f'keywords_group_size must be >= 1, got {keywords_group_size}') + if len(self.query) < keywords_group_size: + raise ValueError(f'{len(self.query)} query(ies) cannot fill a group of ' + f'{keywords_group_size}') + self.path = path + self.num_keywords = num_keywords + self.keywords_group_size = keywords_group_size + self.recycle = recycle + self.system_prompt = system_prompt or self._default_prompt + # Built on the first call rather than here, so a fully cached run needs no backend. + self._backend = backend + self._rollout_kwargs = dict(rollout_kwargs, sampling_params=sampling_params, max_turns=1) + self._rollout: Optional[Any] = None + self._cached_keywords: Dict[str, List[str]] = self.load_keywords() + # Flat: one keyword belongs to one direction, whichever produced it first. + self._seen = {kw.lower() for kws in self._cached_keywords.values() for kw in kws} + # The active slots, the next direction to promote, which slot retires + # next, and the mixed-radix counter walking the active buckets. Drawn + # combinations are remembered because a bucket growing mid-run shifts the + # counter's order and would otherwise let it land on an old group again. + self._active = list(self.query[:keywords_group_size]) + self._next_query = keywords_group_size + self._retire_slot = 0 + self._odometer = [0] * keywords_group_size + self._drawn: Set[Tuple[str, ...]] = set() + self._recycled = False + + # ------------------------------------------------------------------- get + + def get_keywords(self, num_groups: int = 1) -> Optional[List[List[str]]]: + """Up to ``num_groups`` combinations of ``keywords_group_size`` keywords each. + + Fewer than asked for when the directions run dry mid-way -- a partial + batch is still usable -- and None when not even one group could be + filled, which is the caller's signal to stop. + """ + if num_groups < 1: + raise ValueError(f'num_groups must be >= 1, got {num_groups}') + groups: List[List[str]] = [] + for _ in range(num_groups): + group = self._draw_group() + if group is None: + break + groups.append(group) + return groups or None + + def _draw_group(self) -> Optional[List[str]]: + """The next combination nobody has been handed, widening the pool to find one.""" + while True: + group = self._step() + if group is not None: + return group + if not self._grow_or_retire(): + return None + + def _step(self) -> Optional[List[str]]: + """One sweep of the odometer for an undrawn combination. None once there is none.""" + buckets = [self._cached_keywords.get(q, []) for q in self._active] + total = 1 + for bucket in buckets: + total *= len(bucket) + for _ in range(total): + combo = tuple(bucket[i] for bucket, i in zip(buckets, self._odometer)) + self._advance(buckets) + if combo not in self._drawn: + self._drawn.add(combo) + self._recycled = False + return list(combo) + return None + + def _advance(self, buckets: Sequence[Sequence[str]]) -> None: + """Odometer +1, last slot first, carrying into the one before it.""" + for slot in reversed(range(len(buckets))): + self._odometer[slot] += 1 + if self._odometer[slot] < len(buckets[slot]): + return + self._odometer[slot] = 0 + + def _grow_or_retire(self) -> bool: + """Widen the combination space: more keywords, else a new direction. + + False once neither is left. Growing comes first because it multiplies what + the current slots are worth, while retiring gives up on a direction. + """ + short = [q for q in self._active + if len(self._cached_keywords.get(q, [])) < self.num_keywords] + # A round that adds nothing means the model has run out of distinct ideas + # for these directions, so asking again would only spend calls. + if short and self.generate(short): + return True + # Round-robin, so the surplus queries are spent evenly across the slots. + slot = self._retire_slot + self._retire_slot = (slot + 1) % self.keywords_group_size + return self._retire(slot) + + def _retire(self, slot: int) -> bool: + """Promote the next unused direction into ``slot``. False once nothing is left to serve.""" + if self._next_query < len(self.query): + self._active[slot] = self.query[self._next_query] + self._next_query += 1 + self._odometer = [0] * self.keywords_group_size + return True + # Recycling twice without a group in between would spin forever, so it is + # allowed only once per exhaustion -- ``_step`` clears the flag on success. + if self._recycled or not self.recycle or not any(self._cached_keywords.values()): + logger.warning(f'all {len(self.query)} query(ies) are spent; ' + f'pass recycle=True to hand out the same groups again') + return False + self._drawn.clear() + self._active = list(self.query[:self.keywords_group_size]) + self._next_query = self.keywords_group_size + self._odometer = [0] * self.keywords_group_size + self._recycled = True + logger.info(f'[{type(self).__name__}] every query spent -> recycling the combinations') + return True + + # -------------------------------------------------------------- generate + + def generate(self, query: Optional[Sequence[str]] = None) -> int: + """Ask every direction (or just ``query``) for more. Returns how many landed. + + Callable as often as wanted: each round tells the model what that + direction already holds, so the lists grow instead of repeating. + """ + query = list(query if query is not None else self.query) + added = self._add_to_cached(query, self._generate_keywords(query)) + if added: + self.save_keywords() + return added + + def _generate_keywords(self, query: Sequence[str]) -> List[List[str]]: + """One model call per direction, in a single batch; replies stay aligned with ``query``.""" + prompts: List[Trajectory] = [{ + 'messages': [{'role': 'system', 'content': self.system_prompt}, + {'role': 'user', 'content': self._build_user_prompt(q)}], + } for q in query] + if self._rollout is None: + self._rollout = MultiTurnRollout(self._backend, **self._rollout_kwargs) + return [self._parse_keywords_from_response(assistant_text(t)) + for t in self._rollout(prompts)] + + def _build_user_prompt(self, query: str) -> str: + """The ask for one direction, plus what it already holds as an avoid list.""" + known = self._cached_keywords.get(query, []) + want = max(1, self.num_keywords - len(known)) + user = self._user_prompt.format(k=want, query=query) + if known: + user += self._avoid_lead + ', '.join(known[-self._avoid_max:]) + return user + + @staticmethod + def _parse_keywords_from_response(text: str) -> List[str]: + """The JSON array in ``text``, over-length and non-string entries dropped.""" + body = strip_reasoning(text) + start, end = body.find('['), body.rfind(']') + if start < 0 or end <= start: + return [] + try: + arr = json.loads(body[start:end + 1]) + except (ValueError, TypeError): + return [] + return [s for s in (x.strip() for x in arr if isinstance(x, str)) + if 0 < len(s) <= KEYWORD_MAX_LEN] + + # ----------------------------------------------------------------- store + + def _add_to_cached(self, query: Sequence[str], + keywords: Sequence[Sequence[str]]) -> int: + """Append each direction's new keywords, case-insensitively. Returns how many landed.""" + added = 0 + for q, kws in zip(query, keywords): + bucket = self._cached_keywords.setdefault(q, []) + for kw in kws: + if kw.lower() in self._seen: + continue + self._seen.add(kw.lower()) + bucket.append(kw) + added += 1 + if not added: + # Silence here would read as a model that simply produced less. + logger.warning(f'no new keyword for {len(query)} direction(s); ' + f'everything generated was already known') + return added - def __init__(system_prompt: Optional[str] = None): - pass + def load_keywords(self) -> Dict[str, List[str]]: + """Read the cache back, one direction per line. An unreadable line is skipped.""" + cached: Dict[str, List[str]] = {} + if not (self.path and os.path.exists(self.path)): + return cached + with open(self.path, encoding='utf-8') as f: + for line in f: + if not line.strip(): + continue + try: + r = json.loads(line) + except (ValueError, TypeError): + continue + if isinstance(r.get('query'), str) and isinstance(r.get('keywords'), list): + cached[r['query']] = [kw for kw in r['keywords'] if isinstance(kw, str)] + return cached - def generate_keywords(self, text): - pass + def save_keywords(self) -> None: + """Write the cache out atomically, so a crash mid-write cannot truncate it.""" + if not self.path: + return + os.makedirs(os.path.dirname(os.path.abspath(self.path)) or '.', exist_ok=True) + tmp = self.path + '.tmp' + with open(tmp, 'w', encoding='utf-8') as f: + for q, kws in self._cached_keywords.items(): + f.write(json.dumps({'query': q, 'keywords': kws}, ensure_ascii=False) + '\n') + os.replace(tmp, self.path) diff --git a/src/twinkle_agentic/challenger/new/recorder.py b/src/twinkle_agentic/challenger/new/recorder.py new file mode 100644 index 000000000..c72d6fab3 --- /dev/null +++ b/src/twinkle_agentic/challenger/new/recorder.py @@ -0,0 +1,91 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Persistent proposer trajectories for challenger training and diagnosis.""" +import json +import os +import threading +import uuid +from typing import Any, Dict, List + +import numpy as np + +_TOKEN_FIELDS = ('input_ids', 'labels', 'completion_mask', 'attention_mask', 'position_ids') + + +def _as_numpy(value: Any, dtype: Any = None) -> np.ndarray: + if hasattr(value, 'detach'): + value = value.detach().cpu().numpy() + return np.asarray(value, dtype=dtype) + + +def _logprob_column(logprobs: Any) -> List[float]: + """Extract the chosen token's log probability from each sampling step.""" + out: List[float] = [] + for step in logprobs: + if isinstance(step, (int, float)): + out.append(float(step)) + continue + if isinstance(step, (list, tuple)) and step: + chosen = step[0] + if isinstance(chosen, (list, tuple)) and len(chosen) >= 2: + out.append(float(chosen[1])) + continue + raise TypeError(f'cannot read a chosen-token logprob from {step!r}') + return out + + +def _json_default(value: Any) -> Any: + if hasattr(value, 'tolist'): + return value.tolist() + return str(value) + + +class RolloutRecorder: + """Write token arrays to NPZ and trajectory metadata to a JSONL index.""" + + def __init__(self, save_dir: str): + self.trajectory_dir = os.path.join(save_dir, 'trajs') + self.index_path = os.path.join(self.trajectory_dir, 'index.jsonl') + os.makedirs(self.trajectory_dir, exist_ok=True) + self._lock = threading.Lock() + + def write(self, trajectory: Dict[str, Any], **fields: Any) -> None: + arrays: Dict[str, np.ndarray] = {} + for key in _TOKEN_FIELDS: + value = trajectory.get(key) + if value is not None: + arrays[key] = _as_numpy(value, np.int32) + logprobs = trajectory.get('logprobs') + if logprobs is not None: + arrays['logprobs'] = np.asarray(_logprob_column(logprobs), dtype=np.float64) + + name = f'{uuid.uuid4().hex}.npz' + labels = arrays.get('labels', np.asarray([], dtype=np.int32)) + completion_mask = arrays.get('completion_mask') + if completion_mask is None: + n_policy_tokens = int((labels != -100).sum()) + else: + if completion_mask.size != labels.size: + raise ValueError('completion_mask and labels must have the same number of tokens') + n_policy_tokens = int(((labels != -100) & completion_mask.astype(bool)).sum()) + n_logprobs = len(arrays.get('logprobs', ())) + if logprobs is not None and n_logprobs != n_policy_tokens: + raise ValueError(f'logprobs contain {n_logprobs} policy tokens, expected ' + f'{n_policy_tokens} from labels and completion_mask') + metadata = { + key: value + for key, value in trajectory.items() if key not in _TOKEN_FIELDS and key not in ('logprobs', 'rewards') + } + record = dict(metadata) + record.update(fields) + record.update({ + 'npz': name, + 'n_tokens': int(arrays.get('input_ids', np.asarray([])).size), + 'n_policy_tokens': n_policy_tokens, + 'has_logprobs': logprobs is not None, + 'n_logprobs': n_logprobs, + }) + line = json.dumps(record, ensure_ascii=False, default=_json_default) + with self._lock: + np.savez_compressed(os.path.join(self.trajectory_dir, name), **arrays) + with open(self.index_path, 'a', encoding='utf-8') as handle: + handle.write(line + '\n') diff --git a/src/twinkle_agentic/harness/ms_agent.py b/src/twinkle_agentic/harness/ms_agent.py index e1f8f7bbd..c9da5774c 100644 --- a/src/twinkle_agentic/harness/ms_agent.py +++ b/src/twinkle_agentic/harness/ms_agent.py @@ -13,12 +13,16 @@ training and serving; the Env owns the implementation. Wire the executing side from the same list, or the prompt advertises tools the Env cannot run:: - harness = MsAgentHarness(config) - harness.prepare() - tool_manager = env.tool_manager(harness.tool_schemas()) - rollout = MultiTurnRollout(sampler, template, - tool_manager=tool_manager, harness=harness) - outs = rollout([harness.start(q) for q in queries]) +One harness per trajectory: each holds an ``LLMAgent`` with memory and context +of its own, and episodes run in parallel threads:: + + harnesses = [MsAgentHarness(config) for _ in queries] + for h in harnesses: + h.prepare() + tool_managers = [env.tool_manager(h.tool_schemas()) for h, env in zip(harnesses, envs)] + rollout = MultiTurnRollout(sampler, template) + outs = rollout([h.start(q) for h, q in zip(harnesses, queries)], + tool_manager=tool_managers, harness=harnesses) Serving path keeps using ``LLMAgent.run()`` with the same ``agent.yaml`` and the same :class:`~twinkle_agentic.envs.base.Env` backend. This class must diff --git a/src/twinkle_agentic/preprocessor/AUDIT.md b/src/twinkle_agentic/preprocessor/AUDIT.md deleted file mode 100644 index a5f1a71bd..000000000 --- a/src/twinkle_agentic/preprocessor/AUDIT.md +++ /dev/null @@ -1,179 +0,0 @@ -# Preprocessor 审计与整改清单 - -> 审计范围：`src/twinkle_agentic/preprocessor/` 全部 15 个文件、30 个类。 -> 审计方法：逐行只读审阅 + 关键断言代码复核 + cookbook/tests 实际接线核实。 -> 三轮视角：(A) 实现问题 (B) 类设计/拆分合并 (C) 功能增删。 - -## 实施状态（已按本清单完整落地） - -> A1 经确认跳过：R1 已把 `score_filter.py` 整体移入 `experimental/`（零 active 使用的死代码）， -> 对死代码再做 4 文件拆分只增维护面、零收益，启用前再拆。其余 21 项全部实施。 - -| 项 | 状态 | 落地位置 | -|----|------|----------| -| A5 | ✅ | `label_schema.py`（`user_data` 信封 + `set_labels`/`get_label` + `pack_value`） | -| P3 | ✅ | `message_normalizer.py` Pass1 重建 assistant 时 `dict(msg)` 透传全字段 | -| P7 | ✅ | `hard_filter._has_tool_calls` / `message_normalizer._strip_heartbeat`+`_is_atomic` 全部走 `normalize_tool_calls` | -| D7 | ✅ | `trajectory_scorer.py`（Segmenter→HardScorer 逐轮→fuse_segment→aggregate_trajectory→写回 `user_data`，mapper 不删） | -| D7c | ✅ | `trajectory_scorer.py` `_segment_confidence`（一致性+voting稳定+决断性）+ `RubricVerifier.score_detail(extra_context=)` 客观注入重评 | -| D6 | ✅ | `outcome_filter.py`（纯读 `traj_score`/`safety_*` 标签比阈值，fail-open） | -| D8 | ✅ | `safety_scorer.py` + `RubricVerifier(fixed_rubric=)` 固定安全 rubric | -| D9 | ✅ | `pii_presidio_filter.py` `regex_only=True`（stub NlpEngine 免 spaCy，REPLACE→MASK 免 faker） | -| D10 | ✅ | `provenance.py`（`ProvenanceStamp`，血缘写入 `user_data`） | -| R1 | ✅ | `experimental/`（`score_filter.py` + `llm_backend.py` git mv 移出主包） | -| R3/R4/R5 | ✅ | `intent_classifier.py`（默认不删；DEFAULT_DETECTORS 精简为 ToolCall/Code/Math；LLM 路径经 R1 已全归 `llm_backup`） | -| P1/P5/P6/P8/P10 | ✅ | trim 后重算 `is_agent`；deadloop agent 行改扫有文本轮；`max_rounds` 按 pair；refuse 扫全 assistant+可选 reasoning；system 多模态保护 | -| A2/A4 | ✅ | `logprob_utils.py` + `message_utils.py`（`utils.py` 保留 shim）；`intents.py` 常量下沉 | -| A3 | ✅ | `twinkle/preprocessor/base.py` 基类返回 `Tuple[List,List]` + `Mapper`/`Filter` 语义基类（`ModelFilter`/`ProvenanceStamp` 已改用） | -| D4/D5 | ✅ | `language_filter.py`（langid 可选，启发式回退）；`structural_noise.py`（关键词无关噪声轮打标） | -| D1/D2 | ✅ | `offline/near_dedup.py`（MinHash-LSH，datasketch 可选+纯 Python 回退）；`offline/decontaminate.py`（13-gram 重叠，drop/tag） | -| A1 | ⏭️ 跳过 | 见上（R1 已隔离为死代码） | - ---- - -## 0. 结论速览 - -> 本清单已根据 review 意见复核收敛：P2/P4 撤销，P1 降级，P11/P12 归入 R1（死代码，暂不单独修）。 - -- **共需改动 22 项**：实现问题 6（P1、P3、P5–P10）、结构重构 5（A1–A5）、功能增删 11（R1–R5 + D1/D2/D4/D5 + D6/D7/D7c/D8/D9/D10，D3 废弃）。 -- **必须做（会静默损坏训练数据）**：仅 **P3** 一项（工具归一丢 reasoning 字段）。 -- **达成「干净 + 每轮评分」最终目标的核心**：**A5**（`user_data` 信封，去 DAG 前置）+ **D7**（接线 verifier，分数写回每轮）+ **D7c**（自动校准 + 客观纠偏主观重评）+ **D6**（读标签滤废案）+ **D8**（安全 rubric）+ **D9**（PII 纯 regex）。 -- **一句话结论**：现有清单修的是「清洗器 bug + 基础过滤」；要产出「干净且每轮带可信分」的 trajectory，还差——**A5 统一 `user_data` 标签信封（把评分/过滤解耦成打标 mapper + 末尾读标签 filter，去掉 DAG）+ 每轮评分打标(D7) + 自进化校准(D7c) + 废案过滤(D6) + 安全/PII(D8/D9)**。零件多数已存在（`verifier`+`aggregation`+`RubricVerifier`+`llm_backup`），核心工作是**接线 + 定 `user_data` 契约**。 - -### review 复核结论（撤销 / 降级项） - -| 原项 | review 意见 | 复核结论 | 处置 | -|------|-------------|----------|------| -| **P1** trim 砍 tool 尾 | 不以 assistant 结尾的部分无训练必要，最多用于工具调用打分 | 成立。trim 尾部未闭合 tool 对训练无害；`is_agent` 不更新的副作用仅剩“末尾 `assistant(tool_calls)` 无对应结果”，训练时本应 mask | **降级为中等**，改描述，不再算“数据损坏” | -| **P2** heartbeat 误杀 | 这类数据是 openclaw/OpenHands 常见格式，作者本意就是要删 | 成立。agent 轨迹清洗语境下 heartbeat 轮几乎必为真噪声，误杀率极低；`message_normalizer.py:26-27` 注释确认是**故意**删除 | **撤销**（保留现状；可选加词边界，非必做） | -| **P4** 删 reasoning-only 轮 | 只有 thinking 无工具调用，训练无落点 | 成立。纯 thinking 轮无 target 输出，多轮里是悬空推理，删掉合理 | **撤销** | -| **P11** ParaphraseScorer 崩溃 | 应该没有实际使用 | 成立。`ScoreFilter`/`ParaphraseScorer` **全库零 active 使用**（仅自身定义 + docs 示例 + 注释掉的引用），测试只覆盖 `utils` 数学函数 | **归入 R1**（死代码，启用时再修） | -| **P12** IFD 公式口径 | 同上 | 同上 | **归入 R1** | - ---- - -## 一、实现问题（正确性 / 语义） - -### 严重：会静默损坏训练数据 - -| ID | 位置 | 问题 | 改动 | 预期收益 | -|----|------|------|------|----------| -| **P3** | `message_normalizer.py` Pass 1 重建消息 | 工具归一路径只保留 `role/content/tool_calls/tool_call_id` 四字段，**丢弃 `reasoning_content`/`thinking`/`name`** | 重建时透传全部原字段 | reasoning 蒸馏数据不再被清洗流程静默剥离 | - -### 中等：策略漏洞 / 语义错位 - -| ID | 位置 | 问题 | 改动 | 预期收益 | -|----|------|------|------|----------| -| **P1** | `message_sanity.py:317,324-329` trim + `is_agent` | trim 掉末尾未闭合 tool 结果本身对训练无害，但 `is_agent` 在 trim **前**计算、trim 后不更新，残留“末尾 `assistant(tool_calls)` 无对应结果”，`check_tool_matching`（forward-only）不拦 | trim 后重算 `is_agent`，或末轮 `tool_calls` 无结果时 mask/剥离该 call；**非必做** | 末轮悬空 tool_call 得到一致处理，避免训练时误算 loss | -| **P5** | `dead_loop_filter.py:192-194` | `is_agent_row` 为真则**整行跳过** stuck 检测，agent 恰恰最易死循环 | agent 死循环走 `HardScorer.check_no_repeated_calls`（见 D3） | 覆盖 agent 重复工具调用循环，堵住最大系统性漏检 | -| **P6** | `hard_filter.py` `max_rounds` | 实现是 `len(asst_msgs) > max_rounds`，只数 assistant，注释却写 “user-assistant pairs” | 修正为按 pair 计数或改注释与语义一致 | 轮数过滤阈值语义正确 | -| **P7** | `message_normalizer.py` / `hard_filter.py` / `utils.py` | `tool_calls` 真值判断三处不一致（裸真值 vs `_has_tool_calls` 视 `''`/`'[]'`/`[]` 为空 vs `normalize_tool_calls`） | 全部统一走 `normalize_tool_calls` | 同一数据“是否 agent”判定一致，消除跨 filter 不一致 | -| **P8** | `refuse_filter.py` | 只扫首条 assistant 前 600 字、不读 reasoning，多轮/reasoning 拒答漏检 | 扩到全 assistant + reasoning 字段（可配窗口） | 拒答样本召回上升，减少污染 | -| **P9** | `token_soup.py` | `max_chars>0` 只查头部（cookbook 用 8000），尾部乱码漏检；不扫 reasoning | 全文 + reasoning 扫描或分段抽样 | 乱码样本召回上升 | -| **P10** | `message_sanity.py` `consolidate_system_messages` | 合并 multimodal system 时压成纯字符串，可能丢非 text part | 用 `msg_has_media` 保护多模态 system | 多模态 system 不丢内容 | - ---- - -## 二、架构 / 类设计（拆分与合并） - -| ID | 项 | 判断 | 改动 | 预期收益 | -|----|----|------|------|----------| -| **A1**（高） | `score_filter.py` 9 个类 486 行 | 契约与实现混在一起，加 scorer 就改巨型文件 | 拆为 `score/` 子包：`types.py`(RoundContext/ScoreResult/Scorer) + `scorers.py`(ChrMin/SIFD 轻) + `judge.py`(PassN/Paraphrase 重) + `score_filter.py`(编排) | 开闭原则；轻/重依赖分离；新增 scorer 零侵入 | -| **A2**（中） | `utils.py` | logprob 数学 + 消息格式工具两个无关模块塞一起 | 拆为 `logprob_utils.py` + `message_utils.py` | 降耦合；改 score 逻辑不误碰消息工具 | -| **A3**（高） | `twinkle/preprocessor/base.py:39` | 基类 `__call__` 声明返回 `Dict`，所有子类实际返回 `Tuple[kept, dropped]`，类型契约名存实亡 | 基类改 `-> Tuple[List, List]`；可选分 `Mapper`/`Filter` 语义基类 | 类型检查生效；新人不会照错签名写导致解包崩溃 | -| **A4**（低） | intent 常量位置 | `ScoreFilter` 消费 intent，但常量定义在 `intent_classifier.py`，score 独立后形成跨模块依赖 | intent 常量下沉到轻量 `intents.py` | 为 score 子包独立化铺路 | -| **A5**（高，目标前置） | 统一 `user_data` 标签信封 | 评分/安全/血缘无统一落点；D6↔D7 若代码互调会逼出 DAG | 所有标签走 `user_data` 的 `List[Tuple[str, pack_value(v)]]`（PyArrow 稳定，见 D 节前置）；打标 mapper 写、末尾 filter 读，靠列表顺序解耦 | 去 DAG、统一数据契约；A3 返回契约的自然延伸 | - -**明确不动**（避免过度设计）：`HardFilter`/`RefuseFilter`/`DeadLoopFilter`/`TokenSoupFilter` 保持独立（合并成上帝类只会更糟）；`data_juicer.py` 4 个薄封装保持一文件；`LLMBackend` 三类保持；`MessageNormalizer` 3 个 pass 不拆（有强顺序依赖）；`IntentDetector` 层级设计是全代码最佳，保持。 - ---- - -## 三、功能增删 - -### 建议去掉 / 降级（死代码与过度设计） - -| ID | 项 | 证据 | 改动 | 预期收益 | -|----|----|------|------|----------| -| **R1** | `ScoreFilter` 全家（+4 scorer）**+ `llm_backend.py` 整个文件** | `ScoreFilter` **全库零 active 使用**（仅自身定义 + docs 示例 + `train_cold_start.py:216` 注释态；测试只覆盖底层 `utils` 数学函数）；内含两处死代码 bug —— 原 **P11**（`ParaphraseScorer:452-455` 缺 DP pad，小批量/DP>1 时 `SamplerBackend` raise）、原 **P12**（`utils.py:154` `ifd=exp(-mean_delta)` 是 Superfiltering 差分指数口径而非 Cherry 损失比值，阈值不可互换）。`llm_backend.py`（`LLMBackend`/`OpenAIBackend`/`SamplerBackend`）**唯一消费者就是 `ScoreFilter`**（grep 确认除自身+`__init__` 导出外无他），它专为 score 打分提供 `chat`/`prompt_logprobs`/`prompt_logprobs_ids`/`embeddings` | `ScoreFilter` + `llm_backend.py` 一起移出主包到 `experimental/` 或 `data_selection/`，标记未验证；**启用前**再修 P11（复用 `_pad_batch`）+ 明确 P12 IFD 口径并重标阈值 | 主路径不再拖未验证的重代码 + 未接线的 LLM 后端抽象；bug 修复延后到真正需要时 | -| **R2** | `LLMBackend.embeddings()` | 是 R1 中 `llm_backend.py` 的一部分；preprocessor 侧**零调用者**，`SamplerBackend.embeddings` 直接 raise | 随 R1 一并移出（不单独保留伪抽象） | 去掉未接线接口，等真需要 embedding 去重再加 | -| **R3** | `IntentClassifier` | 产物 `key_rounds/intents` 主消费者是死的 `ScoreFilter`；`dataset_think.py` 只 import 不入 pipeline | 重定位为“标注器”（从不 drop）；短期可移出主 pipeline 省 CPU | 明确职责；省无谓计算 | -| **R4** | `ComplexLogic/Reasoning/UserDissatisfaction` Detector | 仅被 `IntentClassifier.DEFAULT_DETECTORS` 引用，下游死 | 精简 default 到 `ToolCall/Code/Math` | 减少无消费者的启发式维护面 | -| **R5** | LLM 调用抽象统一到 `llm_backup` | preprocessor 里活跃的 LLM 生成需求（summarizer/segment/verifier）**早已全部走 `twinkle_agentic/utils/llm_backup.py`**（置信度路由 student/teacher + 蒸馏数据收集）；只有死代码 `ScoreFilter` 还用独立的 `LLMBackend` | 主路径不再引入独立 `LLMBackend` 抽象，生成类需求统一走 `@llm_backup`。**注意**：`llm_backup` 只提供 chat 生成（返回 content 字符串），**不提供 `prompt_logprobs`/`embeddings`**——IFD/chr_min 类 logprob 数据选择若复活，那部分接口需在 `experimental/` 内单独保留或重写，不能指望 `llm_backup` | 收敛到单一蒸馏路由机制；生成享受 student/teacher 蒸馏；消除重复的推理后端抽象 | - -### 建议增加（真正缺失的清洗能力） - -#### 已列（清洗器层） - -| ID | 项 | 缺口 | 改动 | 预期收益 | -|----|----|------|------|----------| -| **D1**（离线） | 近重复去重 MinHash-LSH/SimHash | `DedupFilter` 只做前缀 md5 精确去重，改一字的近重复全漏 | 扩展 `DedupFilter` 或新增 `NearDupFilter`（`datasketch` 轻依赖）。**限离线批处理阶段**（需全局视图）；实时 per-batch 主路径不启用，否则局部视图导致误杀严重 | 相似轨迹被挡，多样性上升；离线做，避免在线误杀 | -| **D2**（离线） | 基准去污染 decontamination | train/test n-gram overlap **完全没有** | 新增 13-gram 重叠过滤，比对**静态** benchmark n-gram 索引。**限离线**或**只打标不删**，避免实时流误杀正常样本 | 评测不被污染，指标可信 | -| **D4**（中） | 语言识别 langid/fastText | 只有 `cjk_ratio` script 比例，粗糙 | 轻量 langid 过滤 | 中英限定更可靠，混语噪声下降 | -| **D5**（低，可选） | 结构化噪声轮识别 | 现有 heartbeat 靠关键词（对 openclaw/OpenHands 格式已够用，见 P2 撤销）；仅当出现无关键词的结构性噪声轮时才需要 | 极短轮 + 高重复 + embedding 距离（复用 D1 基础设施） | 覆盖无关键词的噪声轮；非当前痛点 | - -> **D3（agent 死循环接线）已废弃**：review 指出 `preprocessor` 与 `verifier` 当前**零互相 import**（grep 确认），让 `DeadLoopFilter` 去 import `HardScorer`（一个 RL reward `Verifier`）会破坏模块边界、且职责串（清洗器 vs 打分器）。正确路径并入 **D6/D7**：新增打标/评分 preprocessor，agent 死循环由其中的确定性 check 覆盖。 - -#### 新增（达成「干净 + 每轮评分」最终目标所需的整段能力） - -> 对标业界标准 agent 数据流程（Llama-3 / DeepSeek-V3 / Nemotron / Tulu-3 / AgentInstruct / ToolBench）。目标五属性映射：无不良信息→D8/D9、无废案→D6、无重复冗余→D1+D6(轨迹内)、无心跳→已有、每轮评分→D7。 - -| ID | 项 | 属性 | 缺口 | 改动 | 预期收益 | -|----|----|------|------|------|----------| -| **D7**（高，核心） | 每轮评分打标 preprocessor（**只打标不过滤**） | 每轮评分 | `verifier`（per-round `HardScorer` + per-segment `RubricVerifier`）+ `aggregation`（round→segment→trajectory）**基础设施现成但未接线**；`aggregation.py:27` 明说编排器 `TrajectoryScorer` 未实现 | **新增 preprocessor**（mapper，从不 drop）：`Segmenter → HardScorer(逐轮) → RubricVerifier(逐段) → aggregation → 分数写回 `user_data``。分数、`score_confidence`、安全标全部作为 `(key, pack_value(v))` 追加进 `user_data`（见架构前置 A5）。`RubricVerifier.score_detail()` 已返回完整 `ScoreDetail`，`__call__(trajectory)` 兼容逐行调用 | 直接产出「每轮评分」的 trajectory；打标与过滤解耦，D6/D8 只读标签 | -| **D7c**（高，核心） | 评分校准探针 + 客观→主观重评（自进化，无人评） | 每轮评分可信度 | 自进化框架**不能靠人评对齐**；未校准的分会系统性放大 judge 偏见 | 用三个**自动**信号合成 per-segment `score_confidence`：①**teacher-student 一致性**（复用 `llm_backup` 已收集的 `(student, teacher, match)`）②**结果锚定**（`HardScorer` 确定性 check 当弱标签探针）③**voting 方差**（`RubricVerifier` 已有 voting，导出方差）。**关键**：当客观（硬 check）与 LLM 主观**不一致**时，**把客观结果注入 rubric 的打分上下文，让 `RubricVerifier` 重新评分**（不是简单降权，是带硬信号修正的二次评分） | 分数可信度自动量化；客观事实纠偏主观判断，闭环收敛；零人工 | -| **D6**（高） | 轨迹成败判定（过滤废案）——**纯读标签 filter** | 无废案 / 轨迹内冗余 | 无 outcome verification：失败/绕圈/工具全错/最终答案错的轨迹留在训练集 | **不自己算分**，只 `user_data_get(row, 'traj_score')` 等标签跟**阈值**比 → 判废案 drop（依赖的是 D7 已写好的**数据标签**，不是 D7 的代码，靠 pipeline 列表顺序保证 D6 在 D7 后）。**阈值先拍默认值，实测回收分布后回调**（不做人评标定） | 废案不进训练集；与打标解耦，无模块依赖 | -| **D8**（高） | 安全/毒性评分（复用 rubric） | 无不良信息 | 只有 `RefuseFilter`（拒答正则）+ 敏感词表，无 toxicity/safety 覆盖暴力/仇恨/成人/越狱成功 | **复用 `RubricVerifier`**：把安全维度作为一组**固定 `RubricItem`**（暴力/仇恨/成人/越狱成功/隐私泄露）注入 stage-2 打分，走现有 `_score_with_voting` + `_aggregate`；低于阈值判不良。**无需新分类器/新依赖** | 安全过滤召回远超敏感词表；与 D7 共用打分基础设施 | -| **D9**（中） | PII 真脱敏（激活现有 Presidio） | 无不良信息 | `PIIPresidioFilter` 存在但曾因**慢**去掉（spaCy NER 是瓶颈） | 加回来但**纯 regex 模式**：现 `IGNORED_ENTITIES` 已忽略全部 NER 实体（PERSON/LOCATION/ORG…），只留 regex 标识符（邮箱/电话/证件/银行卡）→ **可不加载 spaCy**，绕过 NER 瓶颈，速度问题基本消除 | 邮箱/电话/证件等真 PII 脱敏，且不拖慢管线 | -| **D10**（中） | 治理层：provenance（血缘字段） | 可追溯 | 无血缘字段（source/teacher_model/timestamp） | 每条 trajectory 把血缘作为 `(key, pack_value(v))` 写进 `user_data`（蒸馏场景 teacher/student 版本）。**批次归因/数据卡暂缓**（backlog，见「不做」） | 可追溯；对标 Nemotron/Dolma 但先只做血缘字段 | - -> **架构前置 A5（去 DAG 的正解）**：所有评分/安全/血缘标签统一写进 **`user_data` 信封**，把「评分」与「过滤」解耦成「**打标 mapper（D7/D8/D10，从不 drop）+ 末尾纯读标签 filter（D6）**」。这样 D6→D7 是**数据依赖**（D7 写标签、D6 读标签），靠**线性 `QualityPreprocessor` 的列表顺序**保证，**无需 DAG、无模块互相 import**。 -> -> **PyArrow 硬约束**：`user_data` 必须是 **`List[Tuple[str, str]]`**，**不能用 dict**（HF `datasets` 的 PyArrow 后端对异构/嵌套 dict 序列化有问题）。已核实这是仓库现有官方约定 —— `twinkle/data_format/trajectory.py:18-19`（`user_data: List[Tuple[str, str]]`，注释 "PyArrow-stable encoding: each entry is (key, json.dumps(value))"），写用 `pack_value(v)`（JSON 字符串，值可为任意结构但对外恒为 `(str, str)`），读用 `user_data_get(row, key)`。每轮分数可用 `(f'round_{i}_score', pack_value(...))` 或 `('round_scores', pack_value([...]))` 形态。 -> -> 确定性 check 复用：D6/D7/D8 都要 `HardScorer` 的 LLM-free check。可抽到无依赖公共层（如 `twinkle_agentic/agent_checks.py`）供 verifier 与新 preprocessor 各自依赖；用户已确认也可接受 preprocessor→verifier 单向依赖，则公共层后置。 - -### 明确不做（自进化框架的取舍） - -| 项 | 为什么不做 | -|----|-----------| -| 人评校准对齐 | 自进化框架不 scale；改用 D7c 的三信号（teacher 一致性 + 结果锚定 + voting 方差）替代 | -| 批次间质量归因 / 数据集版本 diff | 暂缓（backlog），当前不阻塞可用性 | -| 管线内数据配比 / 分层采样 | 移到**训练时的 sampler** 消费 `user_data` 标签，清洗管线只负责打标 | -| DAG / 阶段化编排引擎 | 用 A5 的「打标 + 末尾读标签」拍平成线性，不引入 DAG | - ---- - -## 四、改动项汇总与优先级 - -| 优先级 | 项 | 类型 | 是否引入依赖 | -|--------|----|------|--------------| -| P0 前置 | **A5** | `user_data` 信封（去 DAG，其余目标项的地基） | 否 | -| P0 必做 | P3 | 数据损坏（丢 reasoning） | 否 | -| P1 目标核心 | D7, D7c, D6, D8 | 每轮评分 / 校准 / 废案 / 安全（接线 verifier） | 否（复用 verifier） | -| P1 高 | P7, A1, A3 | 一致性 / 结构重构 | 否 | -| P2 中 | D9, D10 | PII 脱敏 / 血缘 | presidio(D9) | -| P2 中 | P1, P5, P6, P8, P10, A2, D4 | 语义/漏检/降耦合 | langid(D4) | -| P2 中（仅离线） | D1, D2 | 去重/去污染（防实时误杀） | `datasketch`(D1) | -| P3 低 | A4, R1, R2, R3, R4, R5, D5 | 清理/重定位/增强 | 否 | - -**合计 22 项**：实现问题 6（P1、P3、P5–P10）、架构 5（A1–A5）、功能增删 11（R1–R5 + D1/D2/D4/D5 + D6/D7/D7c/D8/D9/D10）。原 P2/P4 撤销，P11/P12 归入 R1，**D3 废弃**（并入 D6/D7）。 - ---- - -## 五、改动后预期整体收益 - -1. **达成最终目标（干净 + 每轮可信分）**：A5 统一标签信封 → D7 分数写回每轮 → D7c 自动校准 + 客观纠偏主观 → D6 读标签滤废案 → D8 安全 rubric → D9 PII；heartbeat 已有、D1 离线去重。五属性齐活，且分数带 `score_confidence`。 -2. **去 DAG**：A5 把「打标 mapper + 末尾读标签 filter」拍平成线性 `QualityPreprocessor`，D6↔D7 只有数据依赖、零模块互调，无需 DAG 引擎。 -3. **数据正确性**：P3 消除 reasoning 被静默剥离；P1 末轮悬空 tool_call 一致处理。 -4. **复用而非新建**：D7 复用 `verifier`+`aggregation`；D7c 复用 `llm_backup` 一致性 + `RubricVerifier` voting；D8 复用 rubric 固定项；D9 纯 regex 免 spaCy —— **几乎零新依赖**。 -5. **可维护性 / 一致性**：A1–A4 契约清晰、依赖分层；R1–R5 砍死代码；P7 统一 `tool_calls` 判定；D10 血缘可追溯。 - -> 落地顺序建议： -> 1. **A5**（定 `user_data` 标签信封 + list-of-tuple/`pack_value` 约定）→ 所有目标项的地基。 -> 2. **P3**（防丢数据，无依赖）。 -> 3. **D7**（每轮评分打标 mapper，接线 verifier+aggregation，写回 `user_data`）。 -> 4. **D7c**（三信号校准 + 客观→主观重评）→ 让分数可信。 -> 5. **D6 + D8**（读标签滤废案 + 安全 rubric，阈值先拍后测）。 -> 6. **D9**（PII 纯 regex 加回）→ **D10**（血缘字段）。 -> 7. **A1 + A3**（结构地基）→ 其余（P1/P5/P6/P8/P10/D4）视数据分布投入。 -> 8. **D1 + D2** 放离线批处理阶段单独跑；配比放训练 sampler。 diff --git a/src/twinkle_agentic/preprocessor/__init__.py b/src/twinkle_agentic/preprocessor/__init__.py index 9734b44b5..9f5deb24c 100644 --- a/src/twinkle_agentic/preprocessor/__init__.py +++ b/src/twinkle_agentic/preprocessor/__init__.py @@ -17,9 +17,7 @@ from .message_sanity import MessageSanityFilter from .model_filter import ModelFilter from .pii_presidio_filter import PIIPresidioFilter -from .provenance import ProvenanceStamp # noqa: F401 from .refuse_filter import RefuseFilter -from .structural_noise import StructuralNoiseTagger # noqa: F401 from .token_soup import TokenSoupFilter logger = get_logger() diff --git a/src/twinkle_agentic/preprocessor/experimental/__init__.py b/src/twinkle_agentic/preprocessor/experimental/__init__.py deleted file mode 100644 index 6257664e3..000000000 --- a/src/twinkle_agentic/preprocessor/experimental/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Experimental / not-yet-wired preprocessor components (AUDIT R1). - -These modules are kept out of the main :mod:`twinkle_agentic.preprocessor` -namespace because they have no active consumer in the shipped pipeline -(``QualityPreprocessor``) and are not exercised by the cookbook or tests: - -- :class:`ScoreFilter` and its scorers (per-round SFT key-round selection). Active - LLM generation in the framework goes through ``twinkle_agentic.utils.llm_backup`` - instead of the local ``LLMBackend`` abstraction. -- :class:`LLMBackend` / :class:`OpenAIBackend` / :class:`SamplerBackend`, which - exclusively serve ``ScoreFilter``. - -Import explicitly from here if you want to experiment with them, e.g.:: - - from twinkle_agentic.preprocessor.experimental import ScoreFilter, SamplerBackend -""" -from .llm_backend import LLMBackend, OpenAIBackend, SamplerBackend # noqa: F401 -from .score_filter import ScoreFilter # noqa: F401 - -__all__ = ['ScoreFilter', 'LLMBackend', 'OpenAIBackend', 'SamplerBackend'] diff --git a/src/twinkle_agentic/preprocessor/experimental/llm_backend.py b/src/twinkle_agentic/preprocessor/experimental/llm_backend.py deleted file mode 100644 index 002618620..000000000 --- a/src/twinkle_agentic/preprocessor/experimental/llm_backend.py +++ /dev/null @@ -1,344 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Abstract LLM backend for preprocessor pipeline. - -Supports two modes: - - OpenAIBackend: httpx-based calls to any OpenAI-compatible HTTP server - - SamplerBackend: direct calls to Twinkle vLLMSampler Ray actor (no HTTP) -""" -from abc import ABC, abstractmethod -from typing import Any, Dict, List, Optional, Tuple - -from twinkle.utils import get_logger - -logger = get_logger() - - -class LLMBackend(ABC): - """Abstract base for LLM inference used by QualityPreprocessor stages.""" - - @abstractmethod - def chat( - self, - messages: List[Dict[str, Any]], - *, - temperature: float = 0.0, - max_tokens: int = 16, - n: int = 1, - ) -> List[Dict[str, str]]: - """Chat completion. - - Returns: - List of n choices, each a dict with keys 'content' and 'reasoning_content'. - """ - - def chat_batch( - self, - messages_list: List[List[Dict[str, Any]]], - *, - temperature: float = 0.0, - max_tokens: int = 16, - n: int = 1, - ) -> List[List[Dict[str, str]]]: - """Batched chat completion. Returns one List[choice] per input messages list. - - Default impl loops over `chat`; backends should override to fan out concurrently - (HTTP) or pass the full list to the underlying sampler in a single call (vLLM DP). - """ - return [self.chat(m, temperature=temperature, max_tokens=max_tokens, n=n) for m in messages_list] - - @abstractmethod - def prompt_logprobs(self, messages: List[Dict[str, Any]]) -> Optional[List]: - """Evaluate prompt tokens without generation. - - Returns: - List of per-token logprob entries (format varies by backend but - is compatible with _extract_logprob helpers), or None on failure. - """ - - @abstractmethod - def prompt_logprobs_ids(self, input_ids_list: List[List[int]]) -> List[List]: - """Batched: evaluate raw token-id prompts without chat template wrapping. - - Used for unconditional perplexity (e.g. IFD denominator). Caller MUST - supply a list of token-id sequences; for distributed backends the list - length must satisfy backend-specific batching constraints (e.g. - ``len >= dp_world_size`` for SamplerBackend). - """ - - def embeddings(self, texts: List[str]) -> Any: - """Compute text embeddings. Override in backends that support it.""" - raise NotImplementedError(f'{type(self).__name__} does not support embeddings') - - -class OpenAIBackend(LLMBackend): - """Backend wrapping any OpenAI-compatible HTTP endpoint.""" - - def __init__( - self, - endpoint: str, - model: str = 'default', - api_key: str = '', - timeout: float = 120.0, - ): - import httpx - headers = {'Content-Type': 'application/json'} - if api_key: - headers['Authorization'] = f'Bearer {api_key}' - self._client = httpx.Client(timeout=timeout, headers=headers) - base = endpoint.rstrip('/') - self._chat_endpoint = f'{base}/v1/chat/completions' - self._embed_endpoint = f'{base}/v1/embeddings' - self._model = model - - @property - def model(self) -> str: - return self._model - - def chat( - self, - messages: List[Dict[str, Any]], - *, - temperature: float = 0.0, - max_tokens: int = 16, - n: int = 1, - ) -> List[Dict[str, str]]: - try: - resp = self._client.post( - self._chat_endpoint, - json={ - 'model': self._model, - 'messages': messages, - 'temperature': temperature, - 'max_tokens': max_tokens, - 'n': n, - }) - resp.raise_for_status() - choices = resp.json().get('choices', []) - results = [] - for c in choices: - msg = c.get('message') or {} - results.append({ - 'content': msg.get('content') or '', - 'reasoning_content': msg.get('reasoning_content') or '', - }) - return results - except Exception as e: - logger.warning(f'[OpenAIBackend] chat failed: {e}') - return [] - - def chat_batch( - self, - messages_list: List[List[Dict[str, Any]]], - *, - temperature: float = 0.0, - max_tokens: int = 16, - n: int = 1, - max_workers: int = 16, - ) -> List[List[Dict[str, str]]]: - """Concurrent chat: vLLM HTTP server multiplexes requests; httpx.Client is thread-safe.""" - from concurrent.futures import ThreadPoolExecutor - if not messages_list: - return [] - workers = max(1, min(max_workers, len(messages_list))) - results: List[List[Dict[str, str]]] = [[] for _ in messages_list] - with ThreadPoolExecutor(max_workers=workers) as ex: - futs = { - ex.submit(self.chat, m, temperature=temperature, max_tokens=max_tokens, n=n): i - for i, m in enumerate(messages_list) - } - for fut in futs: - results[futs[fut]] = fut.result() - return results - - def prompt_logprobs(self, messages: List[Dict[str, Any]]) -> Optional[List]: - try: - resp = self._client.post( - self._chat_endpoint, - json={ - 'model': self._model, - 'messages': messages, - 'max_tokens': 0, - 'prompt_logprobs': 1, - }) - resp.raise_for_status() - return resp.json().get('prompt_logprobs') - except Exception: - return None - - def prompt_logprobs_ids(self, input_ids_list: List[List[int]]) -> List[List]: - endpoint = self._chat_endpoint.rsplit('/', 2)[0] + '/v1/completions' - results: List[List] = [] - for input_ids in input_ids_list: - resp = self._client.post( - endpoint, - json={ - 'model': self._model, - 'prompt': list(input_ids), - 'max_tokens': 0, - 'echo': True, - 'prompt_logprobs': 1, - }) - resp.raise_for_status() - data = resp.json() - choices = data.get('choices') or [] - if choices and 'prompt_logprobs' in choices[0]: - results.append(choices[0]['prompt_logprobs']) - else: - results.append(data['prompt_logprobs']) - return results - - def embeddings(self, texts: List[str]): - import numpy as np - resp = self._client.post( - self._embed_endpoint, json={ - 'model': self._model, - 'input': texts, - }) - resp.raise_for_status() - data = resp.json().get('data', []) - data_sorted = sorted(data, key=lambda x: x.get('index', 0)) - return np.array([d['embedding'] for d in data_sorted], dtype=np.float32) - - -class SamplerBackend(LLMBackend): - """Backend wrapping a Twinkle vLLMSampler (Ray actor, no HTTP overhead).""" - - def __init__( - self, - sampler, - embed_endpoint: str = '', - embed_model: str = 'bge-m3', - ): - """ - Args: - sampler: A vLLMSampler instance (with template already set). - embed_endpoint: Optional OpenAI-compatible endpoint for embeddings. - embed_model: Model name for embeddings. - """ - self._sampler = sampler - self._embed_endpoint = embed_endpoint - self._embed_model = embed_model - self._embed_client = None - if embed_endpoint: - import httpx - self._embed_client = httpx.Client(timeout=120.0) - self._embed_url = f'{embed_endpoint.rstrip("/")}/v1/embeddings' - - def chat( - self, - messages: List[Dict[str, Any]], - *, - temperature: float = 0.0, - max_tokens: int = 16, - n: int = 1, - ) -> List[Dict[str, str]]: - from twinkle.data_format import SamplingParams - trajectory = {'messages': messages} - params = SamplingParams( - temperature=temperature, - max_tokens=max_tokens, - num_samples=n, - ) - try: - responses = self._sampler.sample(trajectory, params) - results = [] - for resp in responses: - for seq in resp.sequences: - text = seq.decoded or '' - reasoning = '' - if '</think>' in text: - parts = text.split('</think>', 1) - reasoning = parts[0].split('<think>')[-1].strip() - text = parts[1].strip() - results.append({'content': text, 'reasoning_content': reasoning}) - return results - except Exception as e: - logger.warning(f'[SamplerBackend] chat failed: {e}') - return [] - - @staticmethod - def _split_think(text: str) -> Tuple[str, str]: - if '</think>' in text: - parts = text.split('</think>', 1) - return parts[1].strip(), parts[0].split('<think>')[-1].strip() - return text, '' - - def chat_batch( - self, - messages_list: List[List[Dict[str, Any]]], - *, - temperature: float = 0.0, - max_tokens: int = 16, - n: int = 1, - ) -> List[List[Dict[str, str]]]: - """One sampler dispatch over the full list; lets vLLM DP workers stay saturated.""" - from twinkle.data_format import SamplingParams - if not messages_list: - return [] - device_mesh = getattr(self._sampler, 'device_mesh', None) - dp_world_size = getattr(device_mesh, 'dp_world_size', 1) or 1 - n_inputs = len(messages_list) - feats = [{'messages': m} for m in messages_list] - # Pad the dispatch so every DP worker has at least one item; trim duplicates after. - if n_inputs < dp_world_size: - feats = feats + [feats[-1]] * (dp_world_size - n_inputs) - params = SamplingParams(temperature=temperature, max_tokens=max_tokens, num_samples=n) - try: - responses = self._sampler.sample(feats, params) - except Exception as e: - logger.warning(f'[SamplerBackend] chat_batch failed: {e}') - return [[] for _ in range(n_inputs)] - responses = list(responses)[:n_inputs] - out: List[List[Dict[str, str]]] = [] - for resp in responses: - choices: List[Dict[str, str]] = [] - for seq in (getattr(resp, 'sequences', None) or []): - text, reasoning = self._split_think(seq.decoded or '') - choices.append({'content': text, 'reasoning_content': reasoning}) - out.append(choices) - while len(out) < n_inputs: - out.append([]) - return out - - def prompt_logprobs(self, messages: List[Dict[str, Any]]) -> Optional[List]: - from twinkle.data_format import SamplingParams - trajectory = {'messages': messages} - params = SamplingParams(max_tokens=0, prompt_logprobs=1) - try: - responses = self._sampler.sample(trajectory, params) - if responses and responses[0].prompt_logprobs is not None: - return responses[0].prompt_logprobs - return None - except Exception as e: - logger.warning(f'[SamplerBackend] prompt_logprobs failed: {e}') - return None - - def prompt_logprobs_ids(self, input_ids_list: List[List[int]]) -> List[List]: - from twinkle.data_format import SamplingParams - if not isinstance(input_ids_list, list) or not input_ids_list: - raise ValueError('prompt_logprobs_ids requires a non-empty List[List[int]].') - device_mesh = getattr(self._sampler, 'device_mesh', None) - dp_world_size = getattr(device_mesh, 'dp_world_size', 1) or 1 - if len(input_ids_list) < dp_world_size: - raise ValueError(f'SamplerBackend.prompt_logprobs_ids requires at least ' - f'dp_world_size={dp_world_size} inputs to keep all DP workers busy, ' - f'got {len(input_ids_list)}. Batch upstream before calling.') - feats = [{'input_ids': list(ids)} for ids in input_ids_list] - params = SamplingParams(max_tokens=0, prompt_logprobs=1) - responses = self._sampler.sample(feats, params) - return [r.prompt_logprobs for r in responses] - - def embeddings(self, texts: List[str]): - if self._embed_client is None: - raise NotImplementedError('SamplerBackend requires embed_endpoint for embeddings. ' - 'Pass embed_endpoint when constructing SamplerBackend.') - import numpy as np - resp = self._embed_client.post( - self._embed_url, json={ - 'model': self._embed_model, - 'input': texts, - }) - resp.raise_for_status() - data = resp.json().get('data', []) - data_sorted = sorted(data, key=lambda x: x.get('index', 0)) - return np.array([d['embedding'] for d in data_sorted], dtype=np.float32) diff --git a/src/twinkle_agentic/preprocessor/experimental/score_filter.py b/src/twinkle_agentic/preprocessor/experimental/score_filter.py deleted file mode 100644 index 94333d8ef..000000000 --- a/src/twinkle_agentic/preprocessor/experimental/score_filter.py +++ /dev/null @@ -1,835 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Pluggable per-round scorer/filter for SFT key rounds. - -Architecture: - - ScoreFilter(backend, scorers=[...]) - ├── pre-fetches logprobs once if any scorer requires them - ├── runs each Scorer in order, collecting ScoreResult per round - ├── trace dump (per-round JSON, multi_turn-style) - └── AND aggregation: a round is kept iff every scorer returns passed=True. - -Built-in scorers (each is its own class): - ChrMinScorer chr_dist_min_pos. LOW = hard = keep. - SIFDScorer IFD / S-IFD-50 / S-IFD-75. Default observe-only. - PassNScorer Self-rollouts judged by an LLM. extras carry rollouts/verdicts. - ParaphraseScorer chr_min over a model paraphrase produced under GT injection. - -Decoupling: - * key_rounds missing/empty → every assistant turn becomes a candidate round. - * intents=None → no intent-based gating (all rounds processed). -""" -import json -import os -import re -import time -from dataclasses import dataclass, field -from typing import Any, Callable, Dict, Iterable, List, Optional, Protocol, Set, Tuple - -from twinkle.data_format import pack_value, user_data_get -from twinkle.preprocessor import Preprocessor -from twinkle.template import Template -from twinkle.utils import get_logger -from ..logprob_utils import _chr_min_distinct, _ifd_family_metrics, _lp_to_jsonable, _pad_batch, _to_int_list -from .llm_backend import LLMBackend - -logger = get_logger() - -_MIN_RESPONSE_TOKENS = 5 - - -@dataclass -class RoundContext: - """Per-round payload passed to scorers.""" - row_idx: int - rnd_idx: int - asst_idx: int - row: Dict[str, Any] - intent: Optional[str] - messages: List[Dict[str, Any]] - context_messages: List[Dict[str, Any]] - cond_ids: List[int] - n_prompt: int - asst_ids: List[int] - asst_text: str - user_prompt: str - features: Dict[str, Any] = field(default_factory=dict) - - -@dataclass -class ScoreResult: - score: Optional[float] = None - passed: bool = True - extras: Dict[str, Any] = field(default_factory=dict) - - -class Scorer(Protocol): - name: str - requires_logprobs: bool - - def score(self, contexts: List[RoundContext]) -> List[ScoreResult]: - ... - - -def _user_data_lookup(user_data: Any, key: str) -> Any: - """Pull a value by key from packed user_data; returns the JSON-decoded value.""" - return user_data_get(user_data, key) - - -# ============================================================================ -# Built-in scorers -# ============================================================================ - - -class ChrMinScorer: - """chr_dist_min_pos. Dual-threshold: keep samples in [low, high).""" - name = 'chr_min' - requires_logprobs = True - - def __init__(self, threshold: float = 0.47): - self._threshold = float(threshold) - - def score(self, contexts: List[RoundContext]) -> List[ScoreResult]: - out: List[ScoreResult] = [] - for ctx in contexts: - cond_lp = ctx.features.get('cond_lp') - asst_lp = ctx.features.get('asst_lp') - score = _chr_min_distinct( - cond_lp, - asst_lp, - ctx.cond_ids, - ctx.asst_ids, - ctx.n_prompt, - ) - passed = (score is None) or (score < self._threshold) - out.append(ScoreResult( - score=score, - passed=passed, - extras={'threshold': self._threshold}, - )) - return out - - -class SIFDScorer: - """IFD / S-IFD-50 / S-IFD-75. Observation-only by default.""" - name = 'sifd' - requires_logprobs = True - - def __init__(self, ifd_threshold: Optional[float] = None): - # If set, passed = (ifd >= threshold). HIGH IFD = hard = keep. - self._ifd_threshold = ifd_threshold - - def score(self, contexts: List[RoundContext]) -> List[ScoreResult]: - out: List[ScoreResult] = [] - for ctx in contexts: - cond_lp = ctx.features.get('cond_lp') - asst_lp = ctx.features.get('asst_lp') - fam = _ifd_family_metrics(cond_lp, asst_lp, ctx.cond_ids, ctx.asst_ids, ctx.n_prompt) - score = fam.get('ifd') - if self._ifd_threshold is None or score is None: - passed = True - else: - passed = score >= self._ifd_threshold - out.append(ScoreResult(score=score, passed=passed, extras=dict(fam))) - return out - - -_JUDGE_SYSTEM_PROMPT = """\ -You are a strict but fair answer grader. Judge whether the [Model Answer] is acceptable based on the reference answer (Ground Truth). -Evaluate the following three aspects; if any has a major issue, return FAIL: - -1. Computational/factual correctness: whether the final conclusion, numbers, and key factual statements match the reference answer; -2. Reasoning/approach similarity: whether the solution path, key steps, and considered dimensions are close to the reference answer; - For open-ended questions (no single correct answer), assess whether the style, stance, and considered dimensions align with the reference answer; -3. Completeness: the answer is not truncated, ends naturally, and covers all points of the question. - -First give a brief 1-3 sentence justification, then on the last line strictly output: -<verdict>PASS</verdict> or <verdict>FAIL</verdict>""" # noqa - - -class PassNScorer: - """Self-rollouts (n × per round) judged by an LLM.""" - name = 'pass_n' - requires_logprobs = False - - def __init__( - self, - backend: LLMBackend, - judge_api=None, - judge_model: Optional[str] = None, - judge_base_url: Optional[str] = None, - judge_api_key: Optional[str] = None, - judge_client_kwargs: Optional[Dict[str, Any]] = None, - n: int = 4, - min_pass: int = 0, - sample_temperature: float = 0.7, - sample_max_tokens: int = 4096, - judge_temperature: float = 0.0, - judge_max_tokens: int = 512, - judge_max_rollout_chars: int = 8000, - judge_max_workers: int = 8, - ): - self._backend = backend - self._judge_api = self._build_judge_api(judge_api, judge_model, judge_base_url, judge_api_key, - judge_client_kwargs) - self._n = max(1, int(n)) - self._min_pass = int(min_pass) - self._sample_temperature = float(sample_temperature) - self._sample_max_tokens = int(sample_max_tokens) - self._judge_temperature = float(judge_temperature) - self._judge_max_tokens = int(judge_max_tokens) - self._judge_max_rollout_chars = int(judge_max_rollout_chars) - self._judge_max_workers = max(1, int(judge_max_workers)) - if self._judge_api is None: - logger.warning('[PassNScorer] no judge_api configured; rollouts will be sampled ' - 'without verdicts (every round trivially passes).') - - @staticmethod - def _build_judge_api(api, model, base_url, api_key, client_kwargs): - if api is not None: - return api - if not model: - return None - from twinkle_agentic.protocol.openai import OpenAI as OpenAIAPI - return OpenAIAPI(model=model, api_key=api_key, base_url=base_url, client_kwargs=client_kwargs) - - @staticmethod - def _extract_text_from_choice(choice: Any) -> str: - if not isinstance(choice, dict): - return '' - parts: List[str] = [] - rc = choice.get('reasoning_content') - if isinstance(rc, str) and rc.strip(): - parts.append(f'<thinking>\n{rc.strip()}\n</thinking>') - content = choice.get('content') - if isinstance(content, str) and content.strip(): - parts.append(content.strip()) - if parts: - return '\n\n'.join(parts) - return content if isinstance(content, str) else '' - - @staticmethod - def _truncate(text: str, max_chars: int) -> str: - if not isinstance(text, str) or max_chars <= 0 or len(text) <= max_chars: - return text - head = max_chars * 2 // 3 - tail = max_chars - head - 32 - if tail <= 0: - return text[:max_chars] - return text[:head] + '\n\n...[truncated]...\n\n' + text[-tail:] - - @staticmethod - def _parse_verdict(judge_text: str) -> Optional[bool]: - if not isinstance(judge_text, str): - return None - compact = ''.join(judge_text.upper().split()) - has_pass = '<VERDICT>PASS</VERDICT>' in compact - has_fail = '<VERDICT>FAIL</VERDICT>' in compact - if has_pass and not has_fail: - return True - if has_fail and not has_pass: - return False - # Fallback: keyword scan in the tail (last 200 chars, post-compact). - tail = compact[-200:] - if 'PASS' in tail and 'FAIL' not in tail: - return True - if 'FAIL' in tail and 'PASS' not in tail: - return False - return None - - def _judge_one(self, user_prompt: str, gt_text: str, rollout_text: str) -> Tuple[bool, str]: - if self._judge_api is None: - return True, '(no judge configured)' - if not rollout_text or not rollout_text.strip(): - return False, '(empty rollout)' - from twinkle.data_format.sampling import SamplingParams - body = (f'[问题]\n{self._truncate(user_prompt, self._judge_max_rollout_chars)}\n\n' - f'[参考答案]\n{self._truncate(gt_text, self._judge_max_rollout_chars)}\n\n' - f'[模型回答]\n{self._truncate(rollout_text, self._judge_max_rollout_chars)}\n\n' - '请评分。') - trajectory = { - 'messages': [ - { - 'role': 'system', - 'content': _JUDGE_SYSTEM_PROMPT - }, - { - 'role': 'user', - 'content': body - }, - ] - } - sp = SamplingParams( - temperature=self._judge_temperature, - max_tokens=self._judge_max_tokens, - num_samples=1, - ) - # extra_body forwards `enable_thinking=False` so the judge skips CoT. - msg = self._judge_api(trajectory, sp, extra_body={'enable_thinking': False}) - if isinstance(msg, list): - msg = msg[0] if msg else {} - text = msg.get('content', '') if isinstance(msg, dict) else str(msg) - text = text or '' - verdict = self._parse_verdict(text) - # Conservative default: ambiguous verdict → FAIL. - return bool(verdict) if verdict is not None else False, text - - def score(self, contexts: List[RoundContext]) -> List[ScoreResult]: - if not contexts: - return [] - ctx_msgs = [ctx.context_messages for ctx in contexts] - batched = self._backend.chat_batch( - ctx_msgs, - temperature=self._sample_temperature, - max_tokens=self._sample_max_tokens, - n=self._n, - ) or [] - - while len(batched) < len(contexts): - batched.append([]) - - from concurrent.futures import ThreadPoolExecutor - work: List[Tuple[int, int, str, str, str]] = [] - for i, (ctx, choices) in enumerate(zip(contexts, batched)): - if not isinstance(choices, list): - continue - for r_i, choice in enumerate(choices): - rt = self._extract_text_from_choice(choice) - work.append((i, r_i, ctx.user_prompt, ctx.asst_text, rt)) - - verdict_by_round: Dict[int, List[Tuple[int, bool, str]]] = {} - if work and self._judge_api is not None: - - def _do(item): - i, r_i, up, gt, rt = item - ok, raw = self._judge_one(up, gt, rt) - return i, r_i, ok, raw - - with ThreadPoolExecutor(max_workers=self._judge_max_workers) as ex: - for i, r_i, ok, raw in ex.map(_do, work): - verdict_by_round.setdefault(i, []).append((r_i, ok, raw)) - - out: List[ScoreResult] = [] - for i, (ctx, choices) in enumerate(zip(contexts, batched)): - rollouts = [{ - 'rollout_idx': r_i, - 'content': self._extract_text_from_choice(c) - } for r_i, c in enumerate(choices or [])] - verdicts = sorted(verdict_by_round.get(i, []), key=lambda x: x[0]) - judgments = [{'rollout_idx': r_i, 'passed': bool(p), 'judge_raw': raw} for r_i, p, raw in verdicts] - pass_count = sum(1 for _, p, _ in verdicts if p) - score = (pass_count / self._n) if rollouts else None - passed = pass_count >= self._min_pass - out.append( - ScoreResult( - score=score, - passed=passed, - extras={ - 'pass_count': pass_count, - 'n_rollouts': len(rollouts), - 'rollouts': rollouts, - 'judgments': judgments, - 'min_pass': self._min_pass, - }, - )) - - scored = [r for r in out if r.score is not None] - if scored: - avg = sum(r.score for r in scored) / len(scored) - logger.info(f'[PassNScorer] graded {len(scored)}/{len(out)} rounds × {self._n} ' - f'rollouts; avg pass-rate = {avg:.3f}') - return out - - -class ParaphraseScorer: - """Generate a model paraphrase under GT injection, then re-score chr_min.""" - name = 'paraphrase' - # Owns its own logprob fetch on the rewritten asst tokens. - requires_logprobs = False - - def __init__( - self, - backend: LLMBackend, - template: Template, - chr_min_threshold: Optional[float] = None, - prompt_budget: int = 4096, - sample_temperature: float = 0.7, - sample_max_tokens: int = 4096, - max_prompt_tokens: int = 1024, - ): - self._backend = backend - self._template = template - self._threshold = chr_min_threshold - self._prompt_budget = int(prompt_budget) - self._sample_temperature = float(sample_temperature) - self._sample_max_tokens = int(sample_max_tokens) - self._max_prompt_tokens = int(max_prompt_tokens) - - @staticmethod - def _inject_gt(context_messages, gt_text): - msgs = [dict(m) if isinstance(m, dict) else m for m in context_messages] - instr = f"""\ -Below is the reference answer to this question, for your reference only: - -<reference_answer> -{gt_text} -</reference_answer> - -Based on the reference answer above, please provide a complete answer to the preceding question in your own words and reasoning. Output your answer directly; do not repeat the reference answer verbatim.""" # noqa - if msgs and isinstance(msgs[-1], dict) and msgs[-1].get('role') == 'user': - last = dict(msgs[-1]) - last['content'] = (last.get('content') or '') + '\n\n' + instr - msgs[-1] = last - else: - msgs.append({'role': 'user', 'content': instr}) - return msgs - - def _truncate_gt(self, gt_text: str, n_prompt: int) -> Optional[str]: - # 80 = conservative instruction-template overhead. - budget = self._prompt_budget - n_prompt - 80 - if budget < 50: - return None - gt_ids = _to_int_list(self._template.tokenizer(gt_text, add_special_tokens=False)['input_ids']) - if len(gt_ids) <= budget: - return gt_text - return self._template.tokenizer.decode(gt_ids[:budget], skip_special_tokens=False) - - def _encode_prompt(self, ctx_msgs): - ids = _to_int_list(self._template.encode({'messages': list(ctx_msgs)}, add_generation_prompt=True)['input_ids']) - if self._max_prompt_tokens <= 0 or len(ids) <= self._max_prompt_tokens: - return ids - return ids[-self._max_prompt_tokens:] - - def score(self, contexts: List[RoundContext]) -> List[ScoreResult]: - if not contexts: - return [] - - keys: List[int] = [] - augmented: List[List[Dict[str, Any]]] = [] - for i, ctx in enumerate(contexts): - gt = self._truncate_gt(ctx.asst_text, ctx.n_prompt) - if gt is None or not ctx.context_messages: - continue - keys.append(i) - augmented.append(self._inject_gt(ctx.context_messages, gt)) - - out: List[ScoreResult] = [ - ScoreResult(score=None, passed=True, extras={'reason': 'paraphrase skipped'}) for _ in contexts - ] - if not keys: - return out - - batched = self._backend.chat_batch( - augmented, - temperature=self._sample_temperature, - max_tokens=self._sample_max_tokens, - n=1, - ) or [] - - # Re-tokenize against the ORIGINAL (no-GT) context so logprobs reflect - # pure self-conditional probability of the paraphrase. - para_data: Dict[int, Tuple[List[int], int, List[int], str]] = {} - for i, choices in zip(keys, batched): - text = None - if choices: - c0 = choices[0] - if isinstance(c0, dict): - text = c0.get('content') - if not isinstance(text, str) or not text.strip(): - continue - ctx = contexts[i] - prompt_ids = self._encode_prompt(ctx.context_messages) - asst_ids = _to_int_list(self._template.tokenizer(text, add_special_tokens=False)['input_ids']) - if len(asst_ids) < _MIN_RESPONSE_TOKENS + 1: - continue - cond_ids = prompt_ids + asst_ids - para_data[i] = (cond_ids, len(prompt_ids), asst_ids, text) - - if not para_data: - return out - - ordered = list(para_data.keys()) - cond_batch = [para_data[i][0] for i in ordered] - asst_batch = [para_data[i][2] for i in ordered] - cond_lps = self._backend.prompt_logprobs_ids(cond_batch) - asst_lps = self._backend.prompt_logprobs_ids(asst_batch) - - for i, cond_lp, asst_lp in zip(ordered, cond_lps, asst_lps): - cond_ids, n_prompt, asst_ids, text = para_data[i] - score = _chr_min_distinct(cond_lp, asst_lp, cond_ids, asst_ids, n_prompt) - if self._threshold is None or score is None: - passed = True - else: - passed = score < self._threshold - out[i] = ScoreResult( - score=score, - passed=passed, - extras={ - 'paraphrase_text': text, - 'n_prompt': n_prompt, - 'cond_lp': _lp_to_jsonable(cond_lp), - 'asst_lp': _lp_to_jsonable(asst_lp), - 'threshold': self._threshold, - }, - ) - - logger.info(f'[ParaphraseScorer] paraphrased + scored {len(para_data)}/' - f'{len(contexts)} rounds') - return out - - -# ============================================================================ -# ScoreFilter (Preprocessor entry point) -# ============================================================================ - - -class ScoreFilter(Preprocessor): - """Score and filter assistant turns by a pluggable scorer set. - - A round is kept iff every scorer returns ``passed=True``. Rows that lose - all key rounds are dropped (configurable via ``keep_if_no_key_rounds``). - - Decoupling rules: - * `key_rounds` missing/empty in `user_data` → every assistant turn - becomes a candidate round. - * `intents=None` → no intent-based gating. - """ - - def __init__( - self, - template: Template, - backend: LLMBackend, - scorers: List[Scorer], - intents: Optional[Iterable[str]] = None, - keep_if_no_key_rounds: bool = False, - drop_row_on_any_fail: bool = True, - max_prompt_tokens: int = 1024, - trace_dir: Optional[str] = None, - trace_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, - success_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, - ): - super().__init__() - if not isinstance(template, Template): - raise TypeError(f'ScoreFilter requires a `Template` instance, got ' - f'{type(template).__name__}.') - self._template = template - self._backend = backend - self._scorers = list(scorers) - self._intents: Optional[Set[str]] = (None if intents is None else set(intents)) - self._keep_if_no_key_rounds = bool(keep_if_no_key_rounds) - self._drop_row_on_any_fail = bool(drop_row_on_any_fail) - self._max_prompt_tokens = int(max_prompt_tokens) - self._trace_dir = trace_dir - self._trace_callback = trace_callback - self._success_callback = success_callback - if self._trace_dir: - import shutil - if os.path.exists(self._trace_dir): - shutil.rmtree(self._trace_dir) - os.makedirs(self._trace_dir, exist_ok=True) - - def __call__(self, rows): - rows_list = self.map_col_to_row(rows) - contexts = self._build_contexts(rows_list) - dropped: List[Dict[str, Any]] = [] - if contexts: - score_table = self._score_contexts(contexts) - self._log_score_summary(contexts, score_table) - if self._trace_dir: - self._write_traces(contexts, score_table) - rows_list, dropped = self._apply_filter(rows_list, contexts, score_table) - return rows_list, dropped - - def _log_score_summary(self, contexts, score_table): - for scorer in self._scorers: - scores = [ - t[scorer.name].score for t in score_table if scorer.name in t and t[scorer.name].score is not None - ] - if not scores: - continue - n_pass = sum(1 for t in score_table if scorer.name in t and t[scorer.name].passed) - extras_sample = {} - for t in score_table: - if scorer.name in t and t[scorer.name].extras: - extras_sample = t[scorer.name].extras - break - extra_keys = [k for k in extras_sample if k != 'threshold'] - extra_stats = '' - for k in extra_keys: - vals = [ - t[scorer.name].extras.get(k) for t in score_table - if scorer.name in t and t[scorer.name].extras and t[scorer.name].extras.get(k) is not None - ] - if vals and isinstance(vals[0], (int, float)): - avg = sum(vals) / len(vals) - extra_stats += f', {k}_avg={avg:.4f}' - logger.info(f'[ScoreFilter/{scorer.name}] n={len(scores)}, ' - f'mean={sum(scores) / len(scores):.4f}, ' - f'min={min(scores):.4f}, max={max(scores):.4f}, ' - f'pass={n_pass}/{len(score_table)}' - f'{extra_stats}') - - # ---- scoring (inlined DefaultScoreCalculator) -------------------------- - - def _score_contexts(self, contexts: List[RoundContext]) -> List[Dict[str, ScoreResult]]: - if any(getattr(s, 'requires_logprobs', False) for s in self._scorers): - self._attach_logprobs(contexts) - out: List[Dict[str, ScoreResult]] = [dict() for _ in contexts] - for scorer in self._scorers: - results = scorer.score(contexts) - if len(results) != len(contexts): - raise RuntimeError(f'scorer {scorer.name!r} returned {len(results)} results ' - f'for {len(contexts)} contexts') - for i, r in enumerate(results): - out[i][scorer.name] = r - return out - - def _attach_logprobs(self, contexts: List[RoundContext]) -> None: - cond_batch = [ctx.cond_ids for ctx in contexts] - asst_batch = [ctx.asst_ids for ctx in contexts] - floor = self._batch_floor() - cond_padded, n_cond = _pad_batch(cond_batch, floor) - asst_padded, n_asst = _pad_batch(asst_batch, floor) - cond_lps = self._backend.prompt_logprobs_ids(cond_padded)[:n_cond] - asst_lps = self._backend.prompt_logprobs_ids(asst_padded)[:n_asst] - for ctx, c, a in zip(contexts, cond_lps, asst_lps): - ctx.features['cond_lp'] = c - ctx.features['asst_lp'] = a - - def _batch_floor(self) -> int: - sampler = getattr(self._backend, '_sampler', None) - device_mesh = getattr(sampler, 'device_mesh', None) - return getattr(device_mesh, 'dp_world_size', 1) or 1 - - # ---- context construction -------------------------------------------- - - def _build_contexts(self, rows: List[Dict[str, Any]]) -> List[RoundContext]: - out: List[RoundContext] = [] - for ri, row in enumerate(rows): - messages = row.get('messages') if isinstance(row, dict) else None - if not isinstance(messages, list): - continue - user_data = row.get('user_data') if isinstance(row, dict) else None - key_rounds = _user_data_lookup(user_data, 'key_rounds') - if not isinstance(key_rounds, list) or not key_rounds: - key_rounds = [i for i, m in enumerate(messages) if isinstance(m, dict) and m.get('role') == 'assistant'] - for rnd_idx, asst_idx in enumerate(key_rounds): - if not isinstance(asst_idx, int): - continue - intent = self._lookup_intent(row, asst_idx) - if self._intents is not None and intent not in self._intents: - continue - ctx = self._prepare_round(row, messages, ri, rnd_idx, asst_idx, intent) - if ctx is not None: - out.append(ctx) - return out - - def _prepare_round( - self, - row: Dict[str, Any], - messages: List[Dict[str, Any]], - ri: int, - rnd_idx: int, - asst_idx: int, - intent: Optional[str], - ) -> Optional[RoundContext]: - if not (0 <= asst_idx < len(messages)): - return None - asst_msg = messages[asst_idx] - if not isinstance(asst_msg, dict) or asst_msg.get('role') != 'assistant': - return None - asst_text = asst_msg.get('content') or '' - if isinstance(asst_text, list): - asst_text = ' '.join( - p.get('text', '') for p in asst_text if isinstance(p, dict) and p.get('type') == 'text') - if not asst_text.strip(): - return None - context_messages = messages[:asst_idx] - if not context_messages: - return None - prompt_ids = self._encode_prompt_within_budget(context_messages) - # Raw asst_ids (no chat-template wrapping) so cond/asst share byte-equal - # A-token sequences; otherwise chr_min positions desync. - asst_ids = _to_int_list(self._template.tokenizer(asst_text, add_special_tokens=False)['input_ids']) - if len(asst_ids) < _MIN_RESPONSE_TOKENS + 1: - return None - return RoundContext( - row_idx=ri, - rnd_idx=rnd_idx, - asst_idx=asst_idx, - row=row, - intent=intent, - messages=messages, - context_messages=context_messages, - cond_ids=prompt_ids + asst_ids, - n_prompt=len(prompt_ids), - asst_ids=asst_ids, - asst_text=asst_text, - user_prompt=self._render_user_prompt(context_messages), - ) - - def _encode_prompt_within_budget(self, ctx_msgs: List[Dict[str, Any]]) -> List[int]: - ctx = list(ctx_msgs) - ids = _to_int_list(self._template.encode({'messages': ctx}, add_generation_prompt=True)['input_ids']) - budget = self._max_prompt_tokens - if budget <= 0 or len(ids) <= budget: - return ids - has_sys = bool(ctx) and isinstance(ctx[0], dict) and ctx[0].get('role') == 'system' - body_start = 1 if has_sys else 0 - while len(ctx) - body_start > 1: - ctx.pop(body_start) - ids = _to_int_list(self._template.encode({'messages': ctx}, add_generation_prompt=True)['input_ids']) - if len(ids) <= budget: - return ids - # Single message still over budget → keep tail tokens. - return ids[-budget:] - - @staticmethod - def _render_user_prompt(ctx_msgs: List[Dict[str, Any]]) -> str: - parts: List[str] = [] - for m in ctx_msgs: - if not isinstance(m, dict): - continue - role = m.get('role') or 'user' - content = m.get('content', '') - if isinstance(content, list): - content = ' '.join( - p.get('text', '') for p in content if isinstance(p, dict) and p.get('type') == 'text') - if isinstance(content, str) and content.strip(): - parts.append(f'[{role}] {content.strip()}') - return '\n\n'.join(parts) - - @staticmethod - def _lookup_intent(row: Dict[str, Any], asst_idx: int) -> Optional[str]: - user_data = row.get('user_data') if isinstance(row, dict) else None - intents = _user_data_lookup(user_data, 'intents') - if not isinstance(intents, dict): - return None - v = intents.get(asst_idx) - if v is None: - v = intents.get(str(asst_idx)) - return v if isinstance(v, str) else None - - # ---- trace dump (multi_turn-style) ----------------------------------- - - def _write_traces( - self, - contexts: List[RoundContext], - score_table: List[Dict[str, ScoreResult]], - ) -> None: - for i, ctx in enumerate(contexts): - try: - scores = score_table[i] if i < len(score_table) else {} - kept = all(r.passed for r in scores.values()) if scores else True - record = self._build_trace_record(ctx, scores, kept) - if self._trace_callback is not None and not bool(self._trace_callback(record)): - continue - success = (bool(self._success_callback(record)) if self._success_callback is not None else kept) - prefix = 'ok' if success else 'fail' - rid = f'{ctx.row_idx}-{ctx.asst_idx}-{i}-{int(time.time() * 1000)}' - rid = re.sub(r'[^A-Za-z0-9_\-.]+', '_', rid)[:64] - path = os.path.join(self._trace_dir, f'{prefix}-{rid}.json') - with open(path, 'w', encoding='utf-8') as f: - json.dump(record, f, ensure_ascii=False, indent=2, default=str) - except Exception as e: - # Observability must never break filtering; surface the cause. - logger.warning(f'[ScoreFilter] trace dump failed for row={ctx.row_idx} ' - f'asst={ctx.asst_idx}: {e}') - - @staticmethod - def _build_trace_record( - ctx: RoundContext, - scores: Dict[str, ScoreResult], - kept: bool, - ) -> Dict[str, Any]: - return { - 'row_idx': ctx.row_idx, - 'rnd_idx': ctx.rnd_idx, - 'asst_idx': ctx.asst_idx, - 'intent': ctx.intent, - 'messages': ctx.messages, - 'n_prompt': ctx.n_prompt, - 'cond_ids': ctx.cond_ids, - 'asst_ids': ctx.asst_ids, - 'features': { - k: (_lp_to_jsonable(v) if k.endswith('_lp') else v) - for k, v in ctx.features.items() - }, - 'scores': { - name: { - 'score': r.score, - 'passed': r.passed, - 'extras': r.extras - } - for name, r in scores.items() - }, - 'kept': bool(kept), - } - - # ---- aggregation & row reassembly ------------------------------------ - - def _apply_filter( - self, - rows: List[Dict[str, Any]], - contexts: List[RoundContext], - score_table: List[Dict[str, ScoreResult]], - ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: - per_row: Dict[int, Dict[str, Any]] = {} - for i, ctx in enumerate(contexts): - scores = score_table[i] if i < len(score_table) else {} - passed = all(r.passed for r in scores.values()) if scores else True - slot = per_row.setdefault(ctx.row_idx, { - 'kept': [], - 'failed': 0, - }) - if passed: - slot['kept'].append(ctx.asst_idx) - else: - slot['failed'] += 1 - - out: List[Dict[str, Any]] = [] - dropped: List[Dict[str, Any]] = [] - n_removed_rounds = 0 - n_removed_rows = 0 - for ri, row in enumerate(rows): - user_data = row.get('user_data') if isinstance(row, dict) else None - kr_val = _user_data_lookup(user_data, 'key_rounds') - had_key_rounds = isinstance(kr_val, list) and bool(kr_val) - decision = per_row.get(ri) - - if decision is None: - # Row produced no contexts (no asst turns or filtered by intent). - if had_key_rounds and not self._keep_if_no_key_rounds: - n_removed_rows += 1 - dropped.append(dict(row, drop_reason='score_no_context')) - continue - if self._intents is not None and not self._keep_if_no_key_rounds: - n_removed_rows += 1 - dropped.append(dict(row, drop_reason='score_no_context')) - continue - out.append(row) - continue - - n_removed_rounds += decision['failed'] - kept = decision['kept'] - if had_key_rounds: - if not kept: - n_removed_rows += 1 - dropped.append(dict(row, drop_reason='score_all_rounds_failed')) - continue - new_row = dict(row) - # Re-pack key_rounds; keep all other entries as-is (already packed). - rebuilt = [(k, v) for (k, v) in (user_data or []) if k != 'key_rounds'] - rebuilt.append(('key_rounds', pack_value(list(kept)))) - new_row['user_data'] = rebuilt - out.append(new_row) - else: - if decision['failed'] > 0 and self._drop_row_on_any_fail: - n_removed_rows += 1 - dropped.append(dict(row, drop_reason='score_round_failed')) - continue - out.append(row) - - logger.info(f'[ScoreFilter] removed {n_removed_rounds} rounds, ' - f'dropped {n_removed_rows} rows, kept {len(out)}/{len(rows)}') - return out, dropped diff --git a/src/twinkle_agentic/preprocessor/intent_classifier.py b/src/twinkle_agentic/preprocessor/intent_classifier.py index baa04c13e..6d1b21c24 100644 --- a/src/twinkle_agentic/preprocessor/intent_classifier.py +++ b/src/twinkle_agentic/preprocessor/intent_classifier.py @@ -13,10 +13,14 @@ # Reasoning block regex covers both <think> and <thinking> forms. _THINK_BLOCK_RE = re.compile(r'<think(?:ing)?>(.*?)</think(?:ing)?>', re.DOTALL | re.IGNORECASE) -# ── Intent categories (canonical vocabulary lives in intents.py; re-exported) ── -from .intents import (INTENT_CODE, INTENT_COMPLEX_LOGIC, # noqa: F401,E402 - INTENT_MATH, INTENT_OTHER, INTENT_REASONING, - INTENT_TOOL_CALL, INTENT_USER_DISSATISFACTION) +# ── Intent categories ───────────────────────────────────────────────────────── +INTENT_TOOL_CALL = 'tool_call' +INTENT_CODE = 'code' +INTENT_MATH = 'math' +INTENT_COMPLEX_LOGIC = 'complex_logic' +INTENT_REASONING = 'reasoning' +INTENT_USER_DISSATISFACTION = 'user_dissatisfaction' +INTENT_OTHER = 'other' # ── Heuristic patterns ──────────────────────────────────────────────────────── _CODE_BLOCK_RE = re.compile(r'```[\s\S]{10,}?```') diff --git a/src/twinkle_agentic/preprocessor/intents.py b/src/twinkle_agentic/preprocessor/intents.py deleted file mode 100644 index d945e59af..000000000 --- a/src/twinkle_agentic/preprocessor/intents.py +++ /dev/null @@ -1,25 +0,0 @@ -"""Intent category constants (AUDIT A4). - -Sunk out of ``intent_classifier.py`` into this dependency-free module so any -consumer (e.g. the experimental log-prob scorers, or downstream sampling that -reads ``intents`` labels) can reference the vocabulary without importing the -heavier classifier + its regex detectors. -""" - -INTENT_TOOL_CALL = 'tool_call' -INTENT_CODE = 'code' -INTENT_MATH = 'math' -INTENT_COMPLEX_LOGIC = 'complex_logic' -INTENT_REASONING = 'reasoning' -INTENT_USER_DISSATISFACTION = 'user_dissatisfaction' -INTENT_OTHER = 'other' - -ALL_INTENTS = ( - INTENT_TOOL_CALL, - INTENT_CODE, - INTENT_MATH, - INTENT_COMPLEX_LOGIC, - INTENT_REASONING, - INTENT_USER_DISSATISFACTION, - INTENT_OTHER, -) diff --git a/src/twinkle_agentic/preprocessor/label_schema.py b/src/twinkle_agentic/preprocessor/label_schema.py deleted file mode 100644 index f9cc97938..000000000 --- a/src/twinkle_agentic/preprocessor/label_schema.py +++ /dev/null @@ -1,117 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Unified ``user_data`` label envelope (AUDIT A5). - -All scoring / safety / provenance annotations produced by the pipeline are -written into a trajectory's ``user_data`` as ``(key, pack_value(value))`` pairs. -This is the single data contract that lets us decouple *tagging* (mappers that -never drop) from *filtering* (a tail filter that only reads tags), so the whole -pipeline stays a linear ``QualityPreprocessor`` list — no DAG, no cross-module -imports between a filter and the verifier it depends on. - -PyArrow hard constraint ------------------------ -``user_data`` MUST be a ``List[Tuple[str, str]]`` (see -``twinkle/data_format/trajectory.py``). We NEVER put a bare ``dict`` in a row -column: HF ``datasets``' PyArrow backend cannot stably serialize -heterogeneous / nested dicts. Structured values are JSON-encoded to a single -string via :func:`pack_value`; on read :func:`user_data_get` JSON-decodes them. - -Keep this module dependency-light: it only knows the *keys* and thin get/set -helpers, so both preprocessors and (optionally) other modules can share it -without pulling in verifier/segment code. -""" -from __future__ import annotations - -from typing import Any, Dict, List, Optional, Tuple - -from twinkle.data_format import pack_value, user_data_get - -# --------------------------------------------------------------------------- -# Canonical label keys -# --------------------------------------------------------------------------- -# Per-round hard scores, aligned to assistant/round order within the trajectory. -# Value: List[float] in [0, 1]. -KEY_ROUND_SCORES = 'round_scores' -# Per-round gated flags (a critical hard check zeroed the round). Value: List[bool]. -KEY_ROUND_GATED = 'round_gated' - -# Per-segment fused scores. Value: List[float] in [0, 1]. -KEY_SEGMENT_SCORES = 'segment_scores' -# Per-segment score confidence (D7c calibration). Value: List[float] in [0, 1]. -KEY_SEGMENT_CONFIDENCE = 'segment_confidence' - -# Whole-trajectory fused score in [0, 1] and its discrete level. -KEY_TRAJ_SCORE = 'traj_score' -KEY_TRAJ_LEVEL = 'traj_level' -# Aggregate confidence for the trajectory score (D7c). Value: float in [0, 1]. -KEY_TRAJ_CONFIDENCE = 'traj_confidence' - -# Safety score in [0, 1] (D8, higher = safer) + boolean unsafe flag. -KEY_SAFETY_SCORE = 'safety_score' -KEY_SAFETY_UNSAFE = 'safety_unsafe' - -# Provenance blob (D10): dict-like value JSON-encoded (source/teacher/student/ts). -KEY_PROVENANCE = 'provenance' - -# Free-form scoring metadata (short-circuit stats, per-check breakdown, etc.). -KEY_SCORE_META = 'score_meta' - -# Active-learning pre-selection (ValueSelector): a cheap, LLM-free "how worth an -# expensive rubric pass is this row" score in [0, 1], its per-component -# breakdown, and the boolean gate the rubric stage reads to decide whether to -# spend an LLM call on this row (top-fraction by value_score). -KEY_VALUE_SCORE = 'value_score' -KEY_VALUE_META = 'value_meta' -KEY_SELECTED_FOR_RUBRIC = 'selected_for_rubric' - -# Persisted rubric diagnosis for rubric-scored rows: a per-segment verification -# chain (rubric text + per-criterion verdict/reason/fix + raw model output + -# query/segment_text). This is the SFT corpus for distilling a PRM / error-checker -# LoRA — store it so training never has to re-run the (expensive) teacher. -# Value: List[dict], one entry per rubric-scored segment (see TrajectoryScorer). -KEY_RUBRIC_DIAGNOSIS = 'rubric_diagnosis' - - -# --------------------------------------------------------------------------- -# thin get / set helpers over the (key, pack_value) envelope -# --------------------------------------------------------------------------- -def get_user_data(row: Dict[str, Any]) -> List[Tuple[str, str]]: - """Return the row's ``user_data`` as a list (never a dict), defaulting to [].""" - ud = row.get('user_data') - if ud is None: - return [] - if isinstance(ud, list): - return ud - # Be forgiving of a stray dict (e.g. hand-authored rows) — flatten to pairs. - if isinstance(ud, dict): - return [(k, v if isinstance(v, str) else pack_value(v)) for k, v in ud.items()] - return [] - - -def get_label(row: Dict[str, Any], key: str, default: Any = None) -> Any: - """Read+JSON-decode the first label matching ``key`` from ``row['user_data']``.""" - return user_data_get(get_user_data(row), key, default) - - -def set_labels(row: Dict[str, Any], updates: Dict[str, Any]) -> Dict[str, Any]: - """Return a shallow-copied row with ``updates`` merged into ``user_data``. - - Existing entries for the same keys are replaced (last-write-wins), preserving - the original order for untouched keys. Values are packed with :func:`pack_value` - so the column stays ``List[Tuple[str, str]]`` (PyArrow-stable). - """ - if not updates: - return row - existing = get_user_data(row) - replace = set(updates.keys()) - merged: List[Tuple[str, str]] = [(k, v) for (k, v) in existing if k not in replace] - for k, v in updates.items(): - merged.append((k, pack_value(v))) - new_row = dict(row) - new_row['user_data'] = merged - return new_row - - -def set_label(row: Dict[str, Any], key: str, value: Any) -> Dict[str, Any]: - """Convenience: set a single label.""" - return set_labels(row, {key: value}) diff --git a/src/twinkle_agentic/preprocessor/logprob_utils.py b/src/twinkle_agentic/preprocessor/logprob_utils.py deleted file mode 100644 index 5f86e2865..000000000 --- a/src/twinkle_agentic/preprocessor/logprob_utils.py +++ /dev/null @@ -1,231 +0,0 @@ -"""Log-probability data-selection math (IFD / S-IFD / chr_min). - -Split out of ``utils.py`` (AUDIT A2): these helpers are consumed only by the -log-prob based scorers (the experimental ``ScoreFilter`` family). Keeping them -separate from the message-format utilities means editing scoring math never -risks touching the message helpers used across every active cleaning step. -""" -import math -from typing import Any, Dict, List, Optional, Set, Tuple - - -def _extract_logprob(lp, token_id: Optional[int] = None) -> Optional[float]: - if lp is None: - return None - if isinstance(lp, (int, float)): - return float(lp) - if not isinstance(lp, dict): - return None - # vLLM with prompt_logprobs=1 returns top-1 PLUS actual token if they differ; - # actual is appended LAST, so iter-first picks the wrong (top-1) one. - entry = None - if token_id is not None: - entry = lp.get(token_id) - if entry is None: - entry = lp.get(str(token_id)) - if entry is None: - entry = next(iter(lp.values()), None) - if entry is None: - return None - if hasattr(entry, 'logprob'): - return float(entry.logprob) - if isinstance(entry, dict): - v = entry.get('logprob') - return float(v) if v is not None else None - if isinstance(entry, (int, float)): - return float(entry) - return None - - -def _to_int_list(x) -> List[int]: - if hasattr(x, 'tolist'): - return x.tolist() - return list(x) - - -def _chr_min_distinct( - cond_lp: List, - asst_lp: List, - cond_ids: List[int], - asst_ids: List[int], - n_prompt: int, - exclude_ids: Optional[Set[int]] = None, -) -> Optional[float]: - """chr_dist_min_pos: fraction of distinct asst-token ids whose - per-occurrence min(cond_lp - asst_lp) is strictly positive.""" - if not asst_lp or not cond_lp or not asst_ids: - return None - n_a = min(len(asst_lp), len(asst_ids)) - n_c = len(cond_lp) - by_tok: Dict[int, List[float]] = {} - for i in range(n_a): - ci = n_prompt + i - if ci >= n_c: - break - tid = asst_ids[i] - if tid is None: - continue - if exclude_ids is not None and int(tid) in exclude_ids: - continue - a = _extract_logprob(asst_lp[i], tid) - c_tok = cond_ids[ci] if ci < len(cond_ids) else None - c = _extract_logprob(cond_lp[ci], c_tok) - if a is None or c is None: - continue - by_tok.setdefault(int(tid), []).append(c - a) - if not by_tok: - return None - pos = sum(1 for diffs in by_tok.values() if min(diffs) > 0) - return pos / len(by_tok) - - -def _chr_min_weighted( - cond_lp: List, - asst_lp: List, - cond_ids: List[int], - asst_ids: List[int], - n_prompt: int, -) -> Optional[float]: - """Magnitude-weighted chr_min: each distinct token contributes |min_delta| - as weight; returns sum(pos_weights) / sum(all_weights).""" - if not asst_lp or not cond_lp or not asst_ids: - return None - n_a = min(len(asst_lp), len(asst_ids)) - n_c = len(cond_lp) - by_tok: Dict[int, List[float]] = {} - for i in range(n_a): - ci = n_prompt + i - if ci >= n_c: - break - tid = asst_ids[i] - if tid is None: - continue - a = _extract_logprob(asst_lp[i], tid) - c_tok = cond_ids[ci] if ci < len(cond_ids) else None - c = _extract_logprob(cond_lp[ci], c_tok) - if a is None or c is None: - continue - by_tok.setdefault(int(tid), []).append(c - a) - if not by_tok: - return None - total_w = 0.0 - pos_w = 0.0 - for diffs in by_tok.values(): - md = min(diffs) - w = abs(md) - total_w += w - if md > 0: - pos_w += w - if total_w == 0: - return None - return pos_w / total_w - - -def _ifd_family_metrics( - cond_lp: List, - asst_lp: List, - cond_ids: List[int], - asst_ids: List[int], - n_prompt: int, -) -> Dict[str, Any]: - """IFD (Cherry-LLM) and S-IFD-{50,75} (T-SHIRT) for one round.""" - if not asst_lp or not cond_lp or not asst_ids: - return {} - n_a = min(len(asst_lp), len(asst_ids)) - n_c = len(cond_lp) - deltas: List[float] = [] - for i in range(n_a): - ci = n_prompt + i - if ci >= n_c: - break - tid = asst_ids[i] - if tid is None: - continue - a = _extract_logprob(asst_lp[i], tid) - c_tok = cond_ids[ci] if ci < len(cond_ids) else None - c = _extract_logprob(cond_lp[ci], c_tok) - if a is None or c is None: - continue - deltas.append(c - a) - if not deltas: - return {} - n = len(deltas) - mean_delta = sum(deltas) / n - out: Dict[str, Any] = { - 'n_tokens': n, - 'mean_delta': mean_delta, - 'ifd': math.exp(-mean_delta), - } - abs_sorted = sorted(range(n), key=lambda i: abs(deltas[i]), reverse=True) - for k_pct in (50, 75): - keep = max(1, int(round(n * k_pct / 100))) - sub = [deltas[i] for i in abs_sorted[:keep]] - out[f's_ifd_{k_pct}'] = math.exp(-sum(sub) / len(sub)) - return out - - -def _mean_logprob_delta( - cond_lp: List, - asst_lp: List, - cond_ids: List[int], - asst_ids: List[int], - n_prompt: int, -) -> Optional[float]: - """Mean per-token (cond_lp - asst_lp) over the response span.""" - if not asst_lp or not cond_lp or not asst_ids: - return None - n_a = min(len(asst_lp), len(asst_ids)) - n_c = len(cond_lp) - deltas: List[float] = [] - for i in range(n_a): - ci = n_prompt + i - if ci >= n_c: - break - tid = asst_ids[i] - if tid is None: - continue - a = _extract_logprob(asst_lp[i], tid) - c_tok = cond_ids[ci] if ci < len(cond_ids) else None - c = _extract_logprob(cond_lp[ci], c_tok) - if a is None or c is None: - continue - deltas.append(c - a) - if not deltas: - return None - return sum(deltas) / len(deltas) - - -def _lp_to_jsonable(lp_list): - """Convert per-position prompt_logprobs into JSON-safe form.""" - out = [] - for lp in (lp_list or []): - if lp is None: - out.append(None) - continue - if isinstance(lp, (int, float)): - out.append(float(lp)) - continue - if not isinstance(lp, dict): - out.append(repr(lp)) - continue - d = {} - for k, v in lp.items(): - if hasattr(v, 'logprob'): - d[str(k)] = { - 'logprob': float(v.logprob), - 'rank': getattr(v, 'rank', None), - 'decoded': getattr(v, 'decoded_token', None) - } - elif isinstance(v, dict): - d[str(k)] = v - else: - d[str(k)] = repr(v) - out.append(d) - return out - - -def _pad_batch(batch: List[List[int]], floor: int) -> Tuple[List[List[int]], int]: - n = len(batch) - if n >= floor or not batch: - return batch, n - return list(batch) + [batch[-1]] * (floor - n), n diff --git a/src/twinkle_agentic/preprocessor/offline/__init__.py b/src/twinkle_agentic/preprocessor/offline/__init__.py deleted file mode 100644 index 0a14834a6..000000000 --- a/src/twinkle_agentic/preprocessor/offline/__init__.py +++ /dev/null @@ -1,24 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Offline batch-only preprocessors (AUDIT D1 / D2). - -These steps require a GLOBAL view of the dataset and must NOT be dropped into the -per-batch :class:`~twinkle_agentic.preprocessor.QualityPreprocessor` pipeline: - -- :class:`NearDupFilter` (D1) — MinHash-LSH near-duplicate removal; per-batch use - would only compare within a batch, causing severe false negatives. -- :class:`Decontaminator` (D2) — benchmark n-gram overlap removal against a static - index; kept out of the real-time path to avoid false-positive deletions - (defaults to a safe ``'tag'``-friendly design). - -They are deliberately kept out of the main package namespace. Import explicitly:: - - from twinkle_agentic.preprocessor.offline import NearDupFilter, Decontaminator - from twinkle_agentic.preprocessor.offline import build_benchmark_index - -Usage: materialize the dataset to ``List[Dict]``, run these once, then re-wrap -the kept rows before/after the streaming QualityPreprocessor pipeline. -""" -from .decontaminate import Decontaminator, build_benchmark_index # noqa: F401 -from .near_dedup import NearDupFilter # noqa: F401 - -__all__ = ['NearDupFilter', 'Decontaminator', 'build_benchmark_index'] diff --git a/src/twinkle_agentic/preprocessor/offline/decontaminate.py b/src/twinkle_agentic/preprocessor/offline/decontaminate.py deleted file mode 100644 index 7dd8e3831..000000000 --- a/src/twinkle_agentic/preprocessor/offline/decontaminate.py +++ /dev/null @@ -1,109 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Benchmark decontamination via n-gram overlap (AUDIT D2) — OFFLINE ONLY. - -Removes (or tags) training rows that overlap with evaluation benchmarks, so -reported metrics aren't inflated by leakage. Follows the standard 13-gram -overlap recipe (GPT-3 / Llama / Dolma): build an n-gram set from the benchmark -texts once, then flag any row whose text shares an n-gram with it. - -OFFLINE CONTRACT: the benchmark index is static and global; build it once and -reuse across the whole dataset. This is not a per-batch pipeline step — but -unlike near-dup it *is* embarrassingly parallel per row, so it can also run as a -standalone batch pass. Default mode ``'drop'`` removes contaminated rows; -``'tag'`` keeps them and only records a ``contaminated`` label (safer default for -real-time-ish contexts where false positives must never delete data). -""" -from __future__ import annotations - -import re -from typing import Any, Dict, Iterable, List, Set, Tuple - -from twinkle.preprocessor import Preprocessor -from twinkle_agentic.utils.message_utils import msg_content_text -from .. import label_schema as L - -KEY_CONTAMINATED = 'contaminated' - -_WORD_RE = re.compile(r'\w+', re.UNICODE) - - -def _ngrams(text: str, n: int) -> Set[str]: - tokens = _WORD_RE.findall(text.lower()) - if len(tokens) < n: - return set() - return {' '.join(tokens[i:i + n]) for i in range(len(tokens) - n + 1)} - - -def build_benchmark_index(texts: Iterable[str], n: int = 13) -> Set[str]: - """Build a static n-gram set from benchmark texts (build once, reuse).""" - index: Set[str] = set() - for t in texts: - index |= _ngrams(t or '', n) - return index - - -class Decontaminator(Preprocessor): - """Flag/drop rows that share an n-gram with a static benchmark index. - - Args: - benchmark_ngrams: prebuilt index from :func:`build_benchmark_index`. - n: n-gram size (must match the index's n). Default 13. - min_overlap: number of shared n-grams to count as contaminated. - mode: ``'drop'`` removes contaminated rows; ``'tag'`` keeps them and only - writes the ``contaminated`` label (fail-open). - scan: which roles to scan — 'user' (default), 'assistant', or 'all'. - """ - - def __init__( - self, - benchmark_ngrams: Set[str], - *, - n: int = 13, - min_overlap: int = 1, - mode: str = 'drop', - scan: str = 'user', - ): - if mode not in ('drop', 'tag'): - raise ValueError("mode must be 'drop' or 'tag'") - if scan not in ('user', 'assistant', 'all'): - raise ValueError("scan must be 'user', 'assistant', or 'all'") - self.index = benchmark_ngrams or set() - self.n = int(n) - self.min_overlap = int(min_overlap) - self.mode = mode - self.scan = scan - - def _row_text(self, row: Dict[str, Any]) -> str: - messages = row.get('messages') or [] - parts = [] - for m in messages: - if not isinstance(m, dict): - continue - role = m.get('role') - if self.scan == 'all' or role == self.scan: - parts.append(msg_content_text(m)) - return '\n'.join(p for p in parts if p) - - def _overlap(self, row: Dict[str, Any]) -> int: - if not self.index: - return 0 - grams = _ngrams(self._row_text(row), self.n) - if not grams: - return 0 - return len(grams & self.index) - - def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: - rows = self.map_col_to_row(rows) - kept: List[Dict[str, Any]] = [] - dropped: List[Dict[str, Any]] = [] - for row in rows: - overlap = self._overlap(row) - contaminated = overlap >= self.min_overlap - if contaminated and self.mode == 'drop': - dropped.append(dict(row, drop_reason='benchmark_contamination')) - continue - if self.mode == 'tag': - kept.append(L.set_label(row, KEY_CONTAMINATED, contaminated)) - else: - kept.append(row) - return kept, dropped diff --git a/src/twinkle_agentic/preprocessor/offline/near_dedup.py b/src/twinkle_agentic/preprocessor/offline/near_dedup.py deleted file mode 100644 index d8999ebb6..000000000 --- a/src/twinkle_agentic/preprocessor/offline/near_dedup.py +++ /dev/null @@ -1,163 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Near-duplicate removal via MinHash-LSH (AUDIT D1) — OFFLINE ONLY. - -``DedupFilter`` collapses only exact/prefix duplicates; a single edited character -slips through. This adds fuzzy near-dup detection over shingled trajectory text. - -OFFLINE CONTRACT (same as :class:`DedupFilter`): this needs a *global* view of -the dataset — it must see all rows in one ``__call__`` and is NOT a per-batch -``QualityPreprocessor`` step. Running near-dup on a per-batch stream would judge -similarity against only the current batch, causing severe false negatives (and, -if used to drop, unstable results). Materialize the dataset, run this once, then -re-wrap the kept rows. - -Uses ``datasketch`` when installed (fast LSH); otherwise falls back to a pure -O(n²) MinHash comparison — correct but slower, fine for modest offline batches. -""" -from __future__ import annotations - -import hashlib -import re -from typing import Any, Dict, List, Set, Tuple - -from twinkle.preprocessor import Preprocessor -from twinkle.utils import get_logger -from twinkle_agentic.utils.message_utils import msg_content_text - -logger = get_logger() - -_WORD_RE = re.compile(r'\w+', re.UNICODE) - - -def _row_text(row: Dict[str, Any]) -> str: - messages = row.get('messages') or [] - return '\n'.join(msg_content_text(m) for m in messages if isinstance(m, dict)) - - -def _shingles(text: str, k: int) -> Set[str]: - tokens = _WORD_RE.findall(text.lower()) - if len(tokens) < k: - return {' '.join(tokens)} if tokens else set() - return {' '.join(tokens[i:i + k]) for i in range(len(tokens) - k + 1)} - - -def _minhash_signature(shingles: Set[str], num_perm: int) -> List[int]: - """Pure-python MinHash: for each of ``num_perm`` salted hashes, take the min.""" - if not shingles: - return [0] * num_perm - sig: List[int] = [] - for p in range(num_perm): - salt = str(p).encode() - mn = min(int(hashlib.md5(salt + s.encode('utf-8')).hexdigest(), 16) for s in shingles) - sig.append(mn) - return sig - - -class NearDupFilter(Preprocessor): - """Global near-duplicate removal over a fully materialized row collection. - - Args: - threshold: Jaccard similarity at/above which two rows are near-duplicates. - shingle_size: word n-gram size for shingling. - num_perm: MinHash permutations (higher = more accurate, slower). - keep: within a near-dup cluster, keep the ``'longest'`` (most messages) - or ``'first'`` seen row. - """ - - def __init__( - self, - *, - threshold: float = 0.8, - shingle_size: int = 5, - num_perm: int = 128, - keep: str = 'longest', - ): - if keep not in ('longest', 'first'): - raise ValueError("keep must be 'longest' or 'first'") - self.threshold = float(threshold) - self.shingle_size = int(shingle_size) - self.num_perm = int(num_perm) - self.keep = keep - - def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: - rows = self.map_col_to_row(rows) - n = len(rows) - if n <= 1: - return rows, [] - shingle_sets = [_shingles(_row_text(r), self.shingle_size) for r in rows] - - try: - from datasketch import MinHash, MinHashLSH - clusters = self._cluster_lsh(shingle_sets, MinHash, MinHashLSH) - except Exception as e: - logger.info(f'[NearDupFilter] datasketch unavailable ({e}); pure-python O(n^2) fallback.') - clusters = self._cluster_bruteforce(shingle_sets) - - keep_flag = [True] * n - dropped: List[Dict[str, Any]] = [] - for cluster in clusters: - if len(cluster) <= 1: - continue - winner = self._pick_winner(rows, cluster) - for idx in cluster: - if idx != winner: - keep_flag[idx] = False - dropped.append(dict(rows[idx], drop_reason='near_duplicate')) - kept = [rows[i] for i in range(n) if keep_flag[i]] - return kept, dropped - - def _pick_winner(self, rows: List[Dict[str, Any]], cluster: List[int]) -> int: - if self.keep == 'first': - return min(cluster) - return max(cluster, key=lambda i: len(rows[i].get('messages') or [])) - - def _cluster_lsh(self, shingle_sets, MinHash, MinHashLSH) -> List[List[int]]: - lsh = MinHashLSH(threshold=self.threshold, num_perm=self.num_perm) - mh_list = [] - for i, sh in enumerate(shingle_sets): - mh = MinHash(num_perm=self.num_perm) - for s in sh: - mh.update(s.encode('utf-8')) - mh_list.append(mh) - lsh.insert(str(i), mh) - return self._union_find([(i, [int(x) for x in lsh.query(mh_list[i])]) for i in range(len(shingle_sets))], - len(shingle_sets)) - - def _cluster_bruteforce(self, shingle_sets) -> List[List[int]]: - n = len(shingle_sets) - neighbors: List[Tuple[int, List[int]]] = [] - for i in range(n): - adj = [i] - for j in range(i + 1, n): - a, b = shingle_sets[i], shingle_sets[j] - if not a and not b: - continue - inter = len(a & b) - union = len(a | b) or 1 - if inter / union >= self.threshold: - adj.append(j) - neighbors.append((i, adj)) - return self._union_find(neighbors, n) - - @staticmethod - def _union_find(adjacency: List[Tuple[int, List[int]]], n: int) -> List[List[int]]: - parent = list(range(n)) - - def find(x): - while parent[x] != x: - parent[x] = parent[parent[x]] - x = parent[x] - return x - - def union(a, b): - ra, rb = find(a), find(b) - if ra != rb: - parent[rb] = ra - - for i, adj in adjacency: - for j in adj: - union(i, j) - groups: Dict[int, List[int]] = {} - for i in range(n): - groups.setdefault(find(i), []).append(i) - return list(groups.values()) diff --git a/src/twinkle_agentic/preprocessor/provenance.py b/src/twinkle_agentic/preprocessor/provenance.py deleted file mode 100644 index 45b926056..000000000 --- a/src/twinkle_agentic/preprocessor/provenance.py +++ /dev/null @@ -1,73 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Data-lineage / provenance stamping — tag only, never drop (AUDIT D10). - -Industry data pipelines keep provenance so any training example is traceable -back to its source (which dataset, which teacher/student model produced it, when -it was ingested, which cleaning pipeline version touched it). In a self-evolving -distillation loop this is what lets us later attribute a regression to a bad -source or a specific teacher, and to reproduce a training mix. - -This mapper writes a single ``provenance`` blob into ``user_data`` (JSON-packed, -PyArrow-stable via A5). It reads whatever lineage fields already exist on the row -(``model_id`` and any configured passthroughs) and adds an ingest timestamp so -the record is self-describing downstream. -""" -from __future__ import annotations - -import time -from typing import Any, Dict, Sequence - -from twinkle.preprocessor import Mapper - -from . import label_schema as L - - -class ProvenanceStamp(Mapper): - """Stamp each row with a provenance blob in ``user_data`` (never drops). - - Args: - source: a static source/dataset identifier for this ingest batch. - pipeline_version: version string of the cleaning pipeline for audit. - model_field: row field holding the producing model id (default 'model_id'). - extra_fields: additional row fields to copy verbatim into provenance - (e.g. 'teacher_model', 'student_model', 'request_id'). - add_timestamp: include a unix ingest timestamp. Default True. - overwrite: if False, rows that already carry a provenance blob are left - untouched (idempotent re-runs / preserve upstream lineage). Default False. - """ - - def __init__( - self, - *, - source: str = '', - pipeline_version: str = '', - model_field: str = 'model_id', - extra_fields: Sequence[str] = (), - add_timestamp: bool = True, - overwrite: bool = False, - ): - self.source = source - self.pipeline_version = pipeline_version - self.model_field = model_field - self.extra_fields = tuple(extra_fields) - self.add_timestamp = bool(add_timestamp) - self.overwrite = bool(overwrite) - - def map(self, row: Dict[str, Any]) -> Dict[str, Any]: - if not self.overwrite and L.get_label(row, L.KEY_PROVENANCE, None) is not None: - return row - blob: Dict[str, Any] = {} - if self.source: - blob['source'] = self.source - if self.pipeline_version: - blob['pipeline_version'] = self.pipeline_version - model = row.get(self.model_field) - if model: - blob['model'] = model - for f in self.extra_fields: - v = row.get(f) - if v is not None: - blob[f] = v - if self.add_timestamp: - blob['ingested_at'] = int(time.time()) - return L.set_label(row, L.KEY_PROVENANCE, blob) diff --git a/src/twinkle_agentic/preprocessor/structural_noise.py b/src/twinkle_agentic/preprocessor/structural_noise.py deleted file mode 100644 index 257899636..000000000 --- a/src/twinkle_agentic/preprocessor/structural_noise.py +++ /dev/null @@ -1,60 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Keyword-free structural noise-turn tagging (AUDIT D5, optional). - -Existing heartbeat stripping (``message_normalizer``) is keyword-based and, per -the audit, already covers the common OpenHands/OpenClaw formats. This optional -tagger catches *keyword-free* structural noise: near-identical, very short turns -that repeat across the trajectory (polling / retries with no new signal), using -only cheap structural signals (length + exact repetition) — no embeddings, no -LLM. It **tags** a per-trajectory noise ratio into ``user_data`` (never drops), -so a downstream filter can act on it if desired. - -The embedding-distance variant sketched in the audit is deferred until the D1 -near-dup infrastructure (which provides the embedding index) exists. -""" -from __future__ import annotations - -from collections import Counter -from typing import Any, Dict - -from twinkle.preprocessor import Mapper -from twinkle_agentic.utils.message_utils import msg_content_text, normalize_tool_calls -from . import label_schema as L - -KEY_NOISE_RATIO = 'structural_noise_ratio' - - -class StructuralNoiseTagger(Mapper): - """Tag the fraction of assistant turns that are short, repeated boilerplate. - - Args: - short_chars: an assistant turn with visible text at/under this length is a - noise candidate (tool-call turns are exempt — they carry structure). - min_repeat: a candidate counts as noise only if its normalized text recurs - at least this many times across the trajectory's assistant turns. - """ - - def __init__(self, *, short_chars: int = 40, min_repeat: int = 3): - self.short_chars = int(short_chars) - self.min_repeat = int(min_repeat) - - def map(self, row: Dict[str, Any]) -> Dict[str, Any]: - messages = row.get('messages') - if not isinstance(messages, list) or not messages: - return row - asst = [m for m in messages if isinstance(m, dict) and m.get('role') == 'assistant'] - if not asst: - return row - texts = [] - for m in asst: - if normalize_tool_calls(m) is not None: - texts.append(None) # tool-call turn: never noise - else: - texts.append(msg_content_text(m).strip()) - counts = Counter(t for t in texts if t) - noise = 0 - for t in texts: - if t and len(t) <= self.short_chars and counts[t] >= self.min_repeat: - noise += 1 - ratio = noise / len(asst) - return L.set_label(row, KEY_NOISE_RATIO, round(ratio, 6)) diff --git a/src/twinkle_agentic/protocol/openai.py b/src/twinkle_agentic/protocol/openai.py index e0a7f60f0..286609cc0 100644 --- a/src/twinkle_agentic/protocol/openai.py +++ b/src/twinkle_agentic/protocol/openai.py @@ -1,4 +1,6 @@ -from typing import Any, Dict, List, Optional, Union +import threading +from contextlib import nullcontext +from typing import Any, ContextManager, Dict, List, Optional, Union from twinkle.data_format import Trajectory from twinkle.data_format.message import Message @@ -11,6 +13,11 @@ class OpenAI(API): Works with any endpoint speaking the ``/v1/chat/completions`` protocol (OpenAI, Azure OpenAI, vLLM, SGLang, Ollama, ...). + + Requests in flight are capped here rather than by whatever thread pool calls + in. A caller's thread count sizes local parallelism and wants to be large; a + provider's quota belongs to the endpoint and wants to be small. One number + cannot serve both, and only this object knows which endpoint it is talking to. """ def __init__( @@ -18,15 +25,47 @@ def __init__( model: str, api_key: Optional[str] = None, base_url: Optional[str] = None, + *, + concurrency: Optional[int] = None, + timeout: Optional[float] = None, + max_retries: Optional[int] = None, client_kwargs: Optional[Dict[str, Any]] = None, ): + """ + Args: + concurrency: most requests allowed in flight at once, or None for no + cap. The limit is per instance and shared by every thread holding + it, so a module-level client caps the whole process. + timeout: per-request timeout in seconds. Left at the SDK's default + when None. + max_retries: how many times the SDK retries a request it deems + transient -- 429, 5xx, timeouts, dropped connections -- using its + own exponential backoff. Left at the SDK's default when None. + client_kwargs: anything else the ``openai`` constructor accepts. + """ from openai import OpenAI as _OpenAIClient + if concurrency is not None and concurrency < 1: + raise ValueError(f'concurrency must be >= 1 or None, got {concurrency}') + kwargs = dict(client_kwargs or {}) + for name, value in (('timeout', timeout), ('max_retries', max_retries)): + if value is None: + continue + if name in kwargs: + raise ValueError(f'{name} was passed both directly and in client_kwargs; ' + 'drop one so that which value wins is not a matter of ordering') + kwargs[name] = value + self.model = model + self.concurrency = concurrency + # Held across the SDK's own retries too: a request that is backing off + # still occupies the endpoint's attention, so it keeps its slot. + self._slots: ContextManager[Any] = ( + threading.BoundedSemaphore(concurrency) if concurrency is not None else nullcontext()) self._client = _OpenAIClient( api_key=api_key, base_url=base_url, - **(client_kwargs or {}), + **kwargs, ) def __call__( @@ -36,7 +75,8 @@ def __call__( **kwargs, ) -> Union[Message, List[Message]]: request = self._build_request(trajectory, sampling_params, kwargs) - response = self._client.chat.completions.create(**request) + with self._slots: + response = self._client.chat.completions.create(**request) messages = [self._choice_to_message(c) for c in response.choices] return messages[0] if sampling_params.num_samples == 1 else messages diff --git a/src/twinkle_agentic/rollout/__init__.py b/src/twinkle_agentic/rollout/__init__.py index 835d94da0..cddff8eb2 100644 --- a/src/twinkle_agentic/rollout/__init__.py +++ b/src/twinkle_agentic/rollout/__init__.py @@ -1,20 +1,7 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +from .api_sampler import APISampler from .base import Rollout from .bridge import extend_with_bridge -from .factory import build_rollout from .multi_turn import MultiTurnRollout -__all__ = [ - 'APIMultiTurnRollout', - 'MultiTurnRollout', - 'Rollout', - 'build_rollout', - 'extend_with_bridge', -] - - -def __getattr__(name: str): - if name == 'APIMultiTurnRollout': - from .api_multi_turn import APIMultiTurnRollout - return APIMultiTurnRollout - raise AttributeError(f'module {__name__!r} has no attribute {name!r}') +__all__ = ['APISampler', 'MultiTurnRollout', 'Rollout', 'extend_with_bridge'] diff --git a/src/twinkle_agentic/rollout/api_multi_turn.py b/src/twinkle_agentic/rollout/api_multi_turn.py deleted file mode 100644 index f140a2dc5..000000000 --- a/src/twinkle_agentic/rollout/api_multi_turn.py +++ /dev/null @@ -1,314 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Any, Callable, Dict, List, Optional - -from twinkle.data_format import Trajectory -from twinkle.data_format.sampling import SamplingParams -from twinkle_agentic.protocol.base import API -from twinkle_agentic.tools.tool_manager import ToolManager -from .base import Rollout - -# Termination reasons surfaced via ``trajectory['stop_reason']``. -_STOP_NO_TOOL = 'stop' -_STOP_LENGTH = 'length' -_STOP_MAX_TURNS = 'max_turns' -_STOP_API_ERROR = 'api_error' - -# Runaway guard: a ``followup_fn`` is expected to return None eventually. This -# only bounds a callback that never does, so one bad hook cannot spin forever. -_MAX_FOLLOWUPS = 20 - - -class APIMultiTurnRollout(Rollout): - """Multi-turn rollout over an OpenAI-compatible chat-completions API. - - Per-trajectory loop: - 1. POST ``messages + tools`` to the API; receive an assistant message - (``content`` and/or structured ``tool_calls``). - 2. Append the assistant message to ``messages``. - 3. If the assistant emitted ``tool_calls``, dispatch each through the - trajectory-bound :class:`ToolManager`, append one - ``{role:'tool', tool_call_id, content}`` per call, then loop. - 4. Else terminate with ``stop_reason='stop'``. - 5. ``finish_reason='length'`` => terminate with ``stop_reason='length'``. - 6. ``turn >= max_turns`` => terminate with ``stop_reason='max_turns'`` - (and ``truncated=True``). - - After the tool loop ends, if a ``followup_fn`` was passed (per call or at - construction), it is invoked exactly as in :class:`MultiTurnRollout`: it may - append one more user message and buy one more generation whose reply is an - answer, not a tool turn (tools are withdrawn for it), repeating until the - callback returns None. This is what lets a challenger append its check-script - and problem-statement stages onto the same conversation. - - Constructor and per-call override semantics intentionally mirror - :class:`MultiTurnRollout`: ``tool_manager`` may be a single instance - (broadcast) or a list aligned 1:1 with trajectories, and it is optional -- - a challenger inventing tasks has nothing to execute. - - Tool schema source: ``trajectory['tools']`` if present, else - ``tool_manager.tool_infos()`` of the trajectory's manager. Caller is - free to set neither — the API will simply be told there are no tools. - - Output trajectory shape (keys added to the input dict): - * ``messages``: the full conversation including tool turns. - * ``turns``: number of API round-trips actually performed. - * ``stop_reason``: one of ``'stop' | 'length' | 'max_turns' | 'api_error'``. - * ``truncated``: True iff terminated by ``max_turns`` or ``length``. - * ``error``: error string when ``stop_reason == 'api_error'``. - """ - - def __init__( - self, - api: API, - tool_manager: Optional[ToolManager] = None, - sampling_params: Optional[SamplingParams] = None, - max_turns: int = 6, - concurrency: int = 8, - extra_body: Optional[Dict[str, Any]] = None, - trace_dir: Optional[str] = None, - trace_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, - success_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, - ): - super().__init__() - if api is None: - raise ValueError('APIMultiTurnRollout requires an API client') - if concurrency < 1: - raise ValueError(f'concurrency must be >= 1, got {concurrency}') - self._init_common( - max_turns=max_turns, - sampling_params=sampling_params, - trace_dir=trace_dir, - trace_callback=trace_callback, - success_callback=success_callback) - self.api = api - self.tool_manager = tool_manager - self.concurrency = concurrency - self.extra_body = dict(extra_body or {}) - - def __call__( - self, - trajectories: List[Trajectory], - **kwargs, - ) -> List[Trajectory]: - if isinstance(trajectories, dict): - raise TypeError('APIMultiTurnRollout.__call__ expects a List[Trajectory]; ' - 'wrap a single trajectory as [trajectory].') - trajectories = list(trajectories) - n = len(trajectories) - if n == 0: - return [] - - sampling_params: SamplingParams = kwargs.get('sampling_params', self.sampling_params) - tool_managers = self._broadcast(kwargs.get('tool_manager', self.tool_manager), n, name='tool_manager') - extra_body = dict(self.extra_body) - if 'extra_body' in kwargs and kwargs['extra_body']: - extra_body.update(kwargs['extra_body']) - followup_fn = kwargs.get('followup_fn') - - # Per-trajectory thread pool. OpenAI ``/chat/completions`` is - # one-conversation-per-call; concurrency only buys us network - # parallelism, never batched compute. - outs: List[Optional[Trajectory]] = [None] * n - with ThreadPoolExecutor(max_workers=self.concurrency) as pool: - futures = { - pool.submit(self._run_one, trajectories[i], tool_managers[i], sampling_params, extra_body, followup_fn): i - for i in range(n) - } - for fut in as_completed(futures): - i = futures[fut] - outs[i] = fut.result() - - result_outs: List[Trajectory] = [o if o is not None else dict(trajectories[i]) for i, o in enumerate(outs)] - if self.trace_dir: - self._write_rollout_traces(result_outs, global_step=kwargs.get('global_step')) - return result_outs - - # ------------------------------------------------------------------ private - - def _run_one( - self, - trajectory: Trajectory, - tool_manager: Optional[ToolManager], - sampling_params: SamplingParams, - extra_body: Dict[str, Any], - followup_fn: Optional[Callable[[Trajectory, int], Any]] = None, - ) -> Trajectory: - """Drive the API turn loop for a single trajectory. - - Never raises; API failures are encoded in ``stop_reason='api_error'`` - with the exception text in ``error``. This keeps one bad row from - poisoning a whole rollout batch. - """ - messages: List[Dict[str, Any]] = list(trajectory.get('messages') or []) - tools = trajectory.get('tools') - if tools is None and tool_manager is not None: - tools = tool_manager.tool_infos() or None - - turn = 0 - stop_reason = _STOP_MAX_TURNS - truncated = False - error: Optional[str] = None - - while turn < self.max_turns: - turn += 1 - req_traj = {'messages': messages} - if tools: - req_traj['tools'] = list(tools) - try: - reply = self.api( - req_traj, sampling_params, extra_body=extra_body) if extra_body else self.api( - req_traj, sampling_params) - except Exception as exc: - stop_reason = _STOP_API_ERROR - error = f'{type(exc).__name__}: {exc}' - truncated = True - break - - assistant_msg = self._normalise_assistant(reply, turn) - messages.append(assistant_msg) - finish = assistant_msg.get('finish_reason') - tool_calls = assistant_msg.get('tool_calls') or [] - - if finish == 'length': - stop_reason = _STOP_LENGTH - truncated = True - break - if not tool_calls: - stop_reason = _STOP_NO_TOOL - break - - # Skip tool execution at the last turn — results would never be - # consumed by a subsequent API call (consistent with multi_turn.py). - if turn >= self.max_turns: - truncated = True - stop_reason = _STOP_MAX_TURNS - break - - if tool_manager is None: - # Nothing can run the call, so the conversation cannot continue: - # say why rather than looping on an unanswered tool turn. - stop_reason = _STOP_API_ERROR - error = ('model emitted tool_calls but this rollout has no ToolManager; ' - 'pass one at construction time or as a per-call kwarg') - truncated = True - break - - try: - for tc in tool_calls: - response = tool_manager(tc) - messages.append({ - 'role': 'tool', - 'tool_call_id': tc.get('id'), - 'content': str(response), - }) - except Exception as exc: - stop_reason = _STOP_API_ERROR - error = f'ToolExecution {type(exc).__name__}: {exc}' - truncated = True - break - else: - # Loop exited normally => max_turns reached. - truncated = True - stop_reason = _STOP_MAX_TURNS - - # Follow-up stages (check script, problem statement, ...). Each one - # appends a user message and takes one generation whose reply is an - # answer: tools are withdrawn so the model writes rather than calls. - # Skipped entirely on an API error -- the conversation is already broken. - followups = 0 - if followup_fn is not None and stop_reason != _STOP_API_ERROR: - while followups < _MAX_FOLLOWUPS: - view = dict(trajectory) - view['messages'] = messages - view['turns'] = turn - view['stop_reason'] = stop_reason - view['truncated'] = truncated - view['followups'] = followups - followup = followup_fn(view, followups) - if followup is None: - break - text, next_params = (followup if isinstance(followup, tuple) - else (followup, None)) - messages.append({'role': 'user', 'content': text}) - followups += 1 - fu_params = next_params if next_params is not None else sampling_params - try: - reply = (self.api( # tools omitted on purpose: this is an answer - {'messages': messages}, fu_params, extra_body=extra_body) - if extra_body else self.api({'messages': messages}, fu_params)) - except Exception as exc: - stop_reason = _STOP_API_ERROR - error = f'{type(exc).__name__}: {exc}' - truncated = True - break - assistant_msg = self._normalise_assistant(reply, turn + followups) - messages.append(assistant_msg) - # A follow-up that stopped cleanly means the episode was not cut - # off after all, even if the tool phase had hit its turn cap. - if assistant_msg.get('finish_reason') == 'length': - truncated = True - elif stop_reason == _STOP_MAX_TURNS: - stop_reason = _STOP_NO_TOOL - - out = dict(trajectory) - out['messages'] = messages - out['turns'] = turn - out['stop_reason'] = stop_reason - out['truncated'] = truncated - out['followups'] = followups - if error is not None: - out['error'] = error - return out - - @staticmethod - def _normalise_assistant(reply: Any, turn: int) -> Dict[str, Any]: - """Ensure tool_calls have stable ``id``/``type`` fields and strip - message-internal noise that would confuse the next API turn. - - Some OpenAI-compatible servers (vLLM, SGLang) occasionally omit - ``tool_call.id``; the assistant->tool round-trip needs a stable - id to wire ``role:'tool'.tool_call_id`` back to the call site. - """ - if not isinstance(reply, dict): - return {'role': 'assistant', 'content': str(reply)} - msg: Dict[str, Any] = {'role': 'assistant'} - content = reply.get('content') - msg['content'] = content if content is not None else '' - finish = reply.get('finish_reason') - if finish is not None: - msg['finish_reason'] = finish - tool_calls = reply.get('tool_calls') or [] - if tool_calls: - normalised: List[Dict[str, Any]] = [] - for i, tc in enumerate(tool_calls): - tc = dict(tc) - tc.setdefault('id', f'call_{turn}_{i}') - tc.setdefault('type', 'function') - normalised.append(tc) - msg['tool_calls'] = normalised - # Reasoning content is informational only; keep it for trace - # forensics but it is never re-fed to the API. - reasoning = reply.get('reasoning_content') - if reasoning: - msg['reasoning_content'] = reasoning - return msg - - def _build_trace_record( - self, - traj: Dict[str, Any], - *, - idx: int, - success: bool, - ) -> Dict[str, Any]: - """The shared record, plus the two fields only this loop produces. - - ``turns`` counts API round-trips and ``error`` carries the exception - text behind ``stop_reason='api_error'`` -- without it a trace of a - failed batch shows an empty conversation and no reason. - """ - record = super()._build_trace_record(traj, idx=idx, success=success) - record['turns'] = traj.get('turns') - if traj.get('error'): - record['error'] = traj['error'] - return record diff --git a/src/twinkle_agentic/rollout/api_sampler.py b/src/twinkle_agentic/rollout/api_sampler.py new file mode 100644 index 000000000..0d2c025ce --- /dev/null +++ b/src/twinkle_agentic/rollout/api_sampler.py @@ -0,0 +1,133 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Sampler-shaped adapter for external generation APIs.""" + +from typing import Any, Dict, List, Literal, Optional + +from twinkle.data_format import Trajectory +from twinkle.data_format.sampling import SampledSequence, SamplingParams, StopReason +from twinkle.template import Template + +from ..protocol.base import API +from .bridge import _to_plain, encode_appended_turn + +_FINISH_TO_STOP: Dict[Optional[str], StopReason] = { + 'stop': 'stop', + 'length': 'length', + 'tool_calls': 'stop', + 'function_call': 'stop', + 'content_filter': 'abort', +} + + +class APIGenerationError(RuntimeError): + """The endpoint failed before returning a response to validate.""" + + +def _normalise_assistant(reply: Any, turn: int) -> Dict[str, Any]: + """Make an API reply safe to render and feed into the next turn.""" + if not isinstance(reply, dict): + raise TypeError(f'API must return an assistant message dict, got {type(reply).__name__}') + message: Dict[str, Any] = { + 'role': 'assistant', + 'content': reply.get('content') or '', + } + tool_calls = reply.get('tool_calls') or [] + if tool_calls: + normalised = [] + for i, tool_call in enumerate(tool_calls): + tool_call = dict(tool_call) + tool_call.setdefault('id', f'call_{turn}_{i}') + tool_call.setdefault('type', 'function') + normalised.append(tool_call) + message['tool_calls'] = normalised + finish_reason = reply.get('finish_reason') + if finish_reason is not None: + message['finish_reason'] = finish_reason + return message + + +class APISampler: + """Normalize one :class:`API` turn into a :class:`SampledSequence`. + + Holds the local ``template`` (an API turn's text must be tokenised the way + the trainer reads it back, not by the endpoint) and the tool schema the + endpoint should see (a rollout's ``pif`` no longer carries it after encode). + """ + + def __init__( + self, + api: API, + template: Template, + *, + tools: Optional[List[Dict[str, Any]]] = None, + appended_as: Literal['demonstration', 'context'] = 'demonstration', + api_kwargs: Optional[Dict[str, Any]] = None, + ): + """ + Args: + appended_as: how the turn enters training -- ``'demonstration'`` + (scored by SFT, skipped by RL) or ``'context'`` (no loss). + ``'completion'`` is refused: it would claim a per-token log-prob + the API never returns. + api_kwargs: request fields forwarded to every API call. + """ + if appended_as not in ('demonstration', 'context'): + raise ValueError("APISampler appended_as must be 'demonstration' or 'context', " + f'got {appended_as!r}; an API turn has no log-prob to be a completion.') + self.api = api + self.template = template + self.tools = list(tools) if tools else None + self.appended_as = appended_as + self.api_kwargs = dict(api_kwargs or {}) + + def __call__(self, + pif: Dict[str, Any], + sampling_params: Optional[SamplingParams] = None, + **adapter_kwargs) -> SampledSequence: + """Generate one external turn in the callback's normalized shape. + + ``adapter_kwargs`` (``adapter_path`` / ``use_base_model``) name a weight + set the API does not have; they are accepted and ignored so callback code + can forward the same values to either backend. + """ + if sampling_params is None: + sampling_params = SamplingParams() + if sampling_params.num_samples != 1: + raise ValueError('APISampler draws one turn per input; got ' + f'num_samples={sampling_params.num_samples}.') + messages = list(pif.get('messages') or []) + if not messages: + raise ValueError('APISampler needs an encoded prefix carrying its messages; ' + "the pif has no 'messages' to send to the endpoint.") + tools = pif.get('tools') if 'tools' in pif else self.tools + request: Trajectory = {'messages': messages} + if tools: + request['tools'] = list(tools) + + try: + reply = self.api(request, sampling_params, **self.api_kwargs) + except Exception as exc: + raise APIGenerationError(f'{type(exc).__name__}: {exc}') from exc + if isinstance(reply, list): + raise TypeError('APISampler expects one message per turn but the API returned a ' + 'list; num_samples > 1 is rejected above, so this is an API bug.') + turn = sum(message.get('role') == 'assistant' for message in messages) + 1 + reply = _normalise_assistant(reply, turn) + + new_tokens = encode_appended_turn(messages, reply, self.template, tools) + new_input_feature = _to_plain( + self.template.concat_input_feature( + pif, new_tokens, appended_as=self.appended_as, tool_calls=reply.get('tool_calls'))) + # concat_input_feature reconstructs content by decoding ``new_tokens``; + # those include the template's rendered tool-call block. Keep the API's + # original content beside its structured calls instead of duplicating it. + assistant_message = {key: reply[key] for key in ('role', 'content', 'tool_calls') if key in reply} + new_input_feature['messages'][-1] = assistant_message + + return SampledSequence( + stop_reason=_FINISH_TO_STOP.get(reply.get('finish_reason'), 'stop'), + tokens=new_tokens, + logprobs=None, + decoded=self.template.decode(new_tokens), + new_input_feature=new_input_feature, + ) diff --git a/src/twinkle_agentic/rollout/base.py b/src/twinkle_agentic/rollout/base.py index f0b6a44a0..dfaa9ae50 100644 --- a/src/twinkle_agentic/rollout/base.py +++ b/src/twinkle_agentic/rollout/base.py @@ -4,25 +4,41 @@ import re import time from abc import ABC, abstractmethod -from typing import Any, Callable, Dict, List, Optional +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Any, Callable, Dict, List, Optional, Tuple from twinkle.data_format import Trajectory, user_data_get from twinkle.data_format.sampling import SamplingParams from .bridge import _to_plain +# Termination reasons surfaced via ``trajectory['stop_reason']``. The sampler +# path takes the first three from the sampler itself; the API path has to name +# them, and one vocabulary for both is what lets a consumer read either. +STOP_NO_TOOL = 'stop' +STOP_LENGTH = 'length' +STOP_MAX_TURNS = 'max_turns' +STOP_GENERATION_ERROR = 'generation_error' + +# Runaway guard: a ``followup_fn`` is expected to return None eventually. This +# only bounds a callback that never does, so one bad hook cannot spin forever. +MAX_FOLLOWUPS = 20 + class Rollout(ABC): """A batch of trajectories in, the same batch with the model's turns appended. - Implementations differ in where the turns come from -- a local sampler, - whose token ids are spliced into the trajectory, or an HTTP endpoint, which - only ever returns text -- and the difference is real enough that they stay - separate classes: only one of them produces something trainable. + The concrete multi-turn loop may source each assistant turn from a local + sampler or an HTTP endpoint. Everything independent of that choice lives + here: option validation, spreading a per-call argument over the batch, the + thread pool that runs episodes, and trace dumping. - Everything that is *not* generation is here: option validation, spreading a - per-call argument over the batch, and the trace dump. It moved up because - the two implementations had drifted into sharing it by reaching across the - class boundary for each other's underscore methods. + One episode per thread, and a subclass only writes the episode. Both + backends are latency-bound on something that is not the caller's CPU -- an + HTTP round trip, a sandbox, a sampler that routes each request to whichever + worker is free -- so the threads overlap the waiting. Nothing crosses + between episodes, which is what makes the pool safe and also what the old + lockstep loop had to give up: there, one slow sandbox round trip held up the + next generation for every trajectory in the batch. """ # Set by _init_common. Declared at class level so a subclass that does its @@ -33,10 +49,7 @@ class boundary for each other's underscore methods. trace_dir: Optional[str] = None trace_callback: Optional[Callable[[Dict[str, Any]], bool]] = None success_callback: Optional[Callable[[Dict[str, Any]], bool]] = None - - @abstractmethod - def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory]: - raise NotImplementedError() + concurrency: Optional[int] = None # ------------------------------------------------------------------ setup @@ -45,6 +58,7 @@ def _init_common( *, max_turns: int, sampling_params: Optional[SamplingParams] = None, + concurrency: Optional[int] = None, trace_dir: Optional[str] = None, trace_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, success_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, @@ -59,21 +73,93 @@ def _init_common( # passing the trajectory several times instead. raise ValueError(f'{type(self).__name__} supports num_samples=1 only, ' f'got {sp.num_samples}') + if concurrency is not None and concurrency < 1: + raise ValueError(f'concurrency must be >= 1 or None, got {concurrency}') self.max_turns = max_turns self.sampling_params = sp + # None means one thread per trajectory. A cap below the batch size costs + # throughput rather than buying safety, so it has to be asked for. + self.concurrency = concurrency self.trace_dir = trace_dir self.trace_callback = trace_callback self.success_callback = success_callback if trace_dir: os.makedirs(trace_dir, exist_ok=True) + # ------------------------------------------------------------------- drive + + def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory]: + """Run one episode per trajectory and return them in the input order. + + Order is restored from the future map rather than from completion order, + because callers pair the result with their own list positionally -- a + GRPO group is a slice of this list. + """ + if isinstance(trajectories, dict): + raise TypeError(f'{type(self).__name__}.__call__ expects a List[Trajectory]; ' + 'wrap a single trajectory as [trajectory].') + trajectories = list(trajectories) + n = len(trajectories) + if n == 0: + return [] + + ctx = self._resolve_call(kwargs, n) + outs: List[Optional[Trajectory]] = [None] * n + workers = min(n, self.concurrency or n) + if workers == 1: + # No pool for a single episode: a thread would only make the + # traceback of a failing one harder to read. + outs = [self._run_one(trajectories[i], i, ctx) for i in range(n)] + else: + with ThreadPoolExecutor(max_workers=workers) as pool: + futures = {pool.submit(self._run_one, trajectories[i], i, ctx): i for i in range(n)} + for fut in as_completed(futures): + outs[futures[fut]] = fut.result() + + result: List[Trajectory] = [o if o is not None else dict(trajectories[i]) for i, o in enumerate(outs)] + if self.trace_dir: + self._write_rollout_traces(result, global_step=kwargs.get('global_step')) + return result + + @abstractmethod + def _run_one(self, trajectory: Trajectory, index: int, ctx: Dict[str, Any]) -> Trajectory: + """One trajectory, start to finish, in its own thread. + + ``ctx`` is whatever ``_resolve_call`` produced; ``index`` is the + trajectory's position in the batch, which is how per-trajectory entries + in ``ctx`` are addressed. + """ + raise NotImplementedError() + + def _resolve_call(self, kwargs: Dict[str, Any], n: int) -> Dict[str, Any]: + """Fold per-call ``**kwargs`` over the constructor defaults, once. + + Done before the pool starts so a bad argument raises from the caller's + frame instead of inside n threads, and so ``_broadcast`` runs once + rather than per episode. + """ + return {} + @staticmethod - def _broadcast(arg, n: int, *, name: str, required: bool = False) -> List[Any]: + def _unpack_followup(followup: Any) -> Tuple[str, Optional[SamplingParams]]: + """``followup_fn`` may answer with text, or text plus its own budget.""" + if isinstance(followup, tuple): + text, params = followup + return text, params + return followup, None + + @staticmethod + def _broadcast(arg, n: int, *, name: str, required: bool = False, per_trajectory: bool = False) -> List[Any]: """One value shared by the batch, or a list already aligned 1:1 with it. A list of the wrong length is refused rather than zipped short: the mismatch would silently pair trajectories with the wrong tool manager, which reads downstream as a model that used the wrong sandbox. + + ``per_trajectory`` refuses to share one instance across a batch at all. + It is for arguments that carry episode state: episodes now run in + parallel threads, so a shared one would have several conversations + writing to the same object instead of merely interleaving in it. """ if arg is None: if required: @@ -85,6 +171,10 @@ def _broadcast(arg, n: int, *, name: str, required: bool = False) -> List[Any]: raise ValueError(f'per-call {name} list length ({len(arg)}) does ' f'not match number of trajectories ({n})') return list(arg) + if per_trajectory and n > 1: + raise ValueError(f'{name} holds per-episode state and cannot be shared by ' + f'{n} trajectories running in parallel threads: pass a list ' + f'of {n}, one per trajectory.') return [arg] * n # ------------------------------------------------------------------ trace @@ -92,6 +182,7 @@ def _broadcast(arg, n: int, *, name: str, required: bool = False) -> List[Any]: _TRACE_SKIP_KEYS = ( 'input_ids', 'labels', + 'completion_mask', 'attention_mask', 'position_ids', 'logprobs', diff --git a/src/twinkle_agentic/rollout/bridge.py b/src/twinkle_agentic/rollout/bridge.py index fa8bffa1b..ce668d54b 100644 --- a/src/twinkle_agentic/rollout/bridge.py +++ b/src/twinkle_agentic/rollout/bridge.py @@ -1,16 +1,21 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Shared, pure bridge-token stitching logic for multi-turn rollouts. - -This module hosts :func:`extend_with_bridge`, a ``self``-free function that -appends tool messages and the next generation prompt to a running -``InputFeature`` (``pif``) as ``-100`` "bridge" tokens. It is shared between -the core-library ``MultiTurnRollout`` and the client-side rollout so the two -paths cannot drift. - -The logic was lifted verbatim from ``MultiTurnRollout._extend_with_bridge`` and -``MultiTurnRollout._append_bridge_tokens``; every ``self.template`` access was -rewritten to use the ``template`` parameter. No Ray decorators -(``@remote_function`` / ``@remote_class``) are applied here. +"""Shared, pure template-space stitching logic for multi-turn rollouts. + +This module hosts ``self``-free functions that grow a running ``InputFeature`` +(``pif``) one turn at a time, all measuring what a turn adds by diffing rendered +chat-template output rather than by pasting special tokens together: + +* :func:`extend_with_bridge` appends tool messages and the next generation + prompt as ``-100`` "bridge" tokens. +* :func:`encode_appended_turn` returns the tokens an assistant turn written + outside the sampler (an API, a human) contributes. + +The bridge logic was lifted verbatim from ``MultiTurnRollout._extend_with_bridge`` +and ``MultiTurnRollout._append_bridge_tokens``; every ``self.template`` access was +rewritten to use the ``template`` parameter. It is shared between the +core-library ``MultiTurnRollout`` and the client-side rollout so the two paths +cannot drift. No Ray decorators (``@remote_function`` / ``@remote_class``) are +applied here. """ @@ -48,39 +53,34 @@ def _to_plain(obj: Any) -> Any: return obj -def extend_with_bridge( - pif: Dict[str, Any], - tool_messages: List[Dict[str, Any]], +def _delta_text( template: Template, -) -> Optional[Dict[str, Any]]: - """Append tool messages and the next generation prompt as -100 bridge. - - Strategy: compute the bridge ENTIRELY in template space. Render - ``messages_before`` and ``messages_before + tool_messages`` with the - same chat template and take ``s_after[len(s_before):]`` as the delta. - - We deliberately do NOT diff against ``tokenizer.decode(pif.input_ids)`` - because raw vLLM output and canonical template rendering differ in - whitespace (e.g. Qwen inserts ``\\n\\n`` between assistant content and - a ``<tool_call>`` block, while the model generates only ``\\n``). Such - cosmetic divergences would break a ``startswith`` alignment but do not - affect training correctness: history tokens stay in ``pif.input_ids`` - verbatim; only the newly appended bridge is tokenized from the - canonical template output. - - Returns ``None`` when the trajectory exceeds ``max_length`` and the - template's truncation strategy is ``'delete'``. + messages_before: List[Dict[str, Any]], + appended: List[Dict[str, Any]], + *, + gen_prompt_before: bool, + gen_prompt_after: bool, + tools: Optional[List[Dict[str, Any]]] = None, +) -> str: + """Text the chat template adds when ``appended`` is tacked onto history. + + ``gen_prompt_*`` place the delta relative to the generation prompt: a bridge + ends on one (``False -> True``), a completion consumes one + (``True -> False``). """ tokenizer = template.tokenizer + enable_thinking = getattr(template, 'enable_thinking', False) - messages_before = list(pif.get('messages') or []) - messages_after = messages_before + list(tool_messages) + def render(messages: List[Dict[str, Any]], add_generation_prompt: bool) -> str: + return tokenizer.apply_chat_template( + messages, + tools=tools or None, + tokenize=False, + add_generation_prompt=add_generation_prompt, + enable_thinking=enable_thinking) - enable_thinking = getattr(template, 'enable_thinking', False) - s_before = tokenizer.apply_chat_template( - messages_before, tokenize=False, add_generation_prompt=False, enable_thinking=enable_thinking) - s_after = tokenizer.apply_chat_template( - messages_after, tokenize=False, add_generation_prompt=True, enable_thinking=enable_thinking) + s_before = render(messages_before, gen_prompt_before) + s_after = render(list(messages_before) + list(appended), gen_prompt_after) if not s_after.startswith(s_before): # Appending a *user* message moves where Qwen3's template thinks the @@ -101,28 +101,86 @@ def extend_with_bridge( # What stays on record is the history as generated, thinking included -- # those are the tokens the policy read back when it produced the next # turn, and a later training step has to see the same. - s_anchor = tokenizer.apply_chat_template( - _ANCHOR, tokenize=False, add_generation_prompt=False, - enable_thinking=enable_thinking) - s_anchor_after = tokenizer.apply_chat_template( - _ANCHOR + list(tool_messages), tokenize=False, add_generation_prompt=True, - enable_thinking=enable_thinking) + s_anchor = render(_ANCHOR, gen_prompt_before) + s_anchor_after = render(_ANCHOR + list(appended), gen_prompt_after) if not s_anchor_after.startswith(s_anchor): raise RuntimeError('Canonical chat_template output for messages_after is not a ' 'prefix-extension of messages_before, and the same is true ' - 'of a one-message stand-in history; cannot compute bridge ' + 'of a one-message stand-in history; cannot compute the ' 'delta. This indicates the template is non-monotonic in the ' 'message list (e.g. reorders / rewrites earlier turns).\n' f's_before tail: {s_before[-80:]!r}\n' f's_after at same offset: ' f'{s_after[max(0, len(s_before) - 80):len(s_before) + 80]!r}') s_before, s_after = s_anchor, s_anchor_after - bridge_text = s_after[len(s_before):] + return s_after[len(s_before):] + + +def encode_appended_turn( + messages_before: List[Dict[str, Any]], + message: Dict[str, Any], + template: Template, + tools: Optional[List[Dict[str, Any]]] = None, +) -> List[int]: + """Tokens an assistant turn authored elsewhere contributes to the sequence. + + A sampler returns the ids it generated; an API returns text, whose tokens are + only part of the turn -- the template also writes the turn terminator and + whatever follows it. Diffing the rendered template recovers those without + naming a single special token, so this holds for any chat template. + + The result is what :meth:`Template.concat_input_feature` expects as + ``new_tokens``, and is token-for-token what :meth:`Template.encode` would + have produced for the same conversation. + """ + delta = _delta_text( + template, + messages_before, [template.decode_tool_calls(message)], + gen_prompt_before=True, + gen_prompt_after=False, + tools=tools) + if not delta: + raise RuntimeError(f'Appending {message.get("role")!r} turn added no text; ' + 'the chat template dropped it entirely.') + tokens = template.tokenizer.encode(delta, add_special_tokens=False) + if not tokens: + raise RuntimeError(f'Appended turn tokenised to an empty id list: {delta!r}') + return tokens + + +def extend_with_bridge( + pif: Dict[str, Any], + tool_messages: List[Dict[str, Any]], + template: Template, +) -> Optional[Dict[str, Any]]: + """Append tool messages and the next generation prompt as -100 bridge. + + Strategy: compute the bridge ENTIRELY in template space. Render + ``messages_before`` and ``messages_before + tool_messages`` with the + same chat template and take ``s_after[len(s_before):]`` as the delta. + + We deliberately do NOT diff against ``tokenizer.decode(pif.input_ids)`` + because raw vLLM output and canonical template rendering differ in + whitespace (e.g. Qwen inserts ``\\n\\n`` between assistant content and + a ``<tool_call>`` block, while the model generates only ``\\n``). Such + cosmetic divergences would break a ``startswith`` alignment but do not + affect training correctness: history tokens stay in ``pif.input_ids`` + verbatim; only the newly appended bridge is tokenized from the + canonical template output. + + Returns ``None`` when the trajectory exceeds ``max_length`` and the + template's truncation strategy is ``'delete'``. + """ + messages_before = list(pif.get('messages') or []) + messages_after = messages_before + list(tool_messages) + + bridge_text = _delta_text( + template, messages_before, tool_messages, gen_prompt_before=False, gen_prompt_after=True) if not bridge_text: raise RuntimeError('Bridge text computation returned empty string; ' 'tool turn would add no tokens (template misconfiguration?).') - bridge_ids = tokenizer.encode(bridge_text, add_special_tokens=False) + bridge_ids = template.tokenizer.encode(bridge_text, add_special_tokens=False) if not bridge_ids: raise RuntimeError(f'Bridge text tokenised to empty id list: {bridge_text!r}') @@ -142,8 +200,10 @@ def _append_bridge_tokens( """Append bridge tokens with labels = -100. Mirrors the unroll-append-reroll pattern of - :meth:`Template.concat_input_feature` so that ``labels`` semantics - stay consistent with the sampler-produced pif. + :meth:`Template.concat_input_feature` so that ``labels`` and + ``completion_mask`` semantics stay consistent with the sampler-produced + pif. Bridge tokens are nobody's completion -- neither scored nor + log-prob-bearing -- so both fields are appended as zeros. Shallow copy is deliberately used: every mutation below is a top-level key reassignment, never an in-place change to nested @@ -164,12 +224,15 @@ def _append_bridge_tokens( labels = labels[-1:] + labels[:-1] else: labels = [-100] * len(input_ids) + completion_mask = template._prefix_completion_mask(result, labels) input_ids = input_ids + list(bridge_ids) labels = labels + [-100] * len(bridge_ids) + completion_mask = completion_mask + [0] * len(bridge_ids) result['input_ids'] = input_ids result['labels'] = labels + result['completion_mask'] = completion_mask if 'mm_token_type_ids' in result: import torch diff --git a/src/twinkle_agentic/rollout/factory.py b/src/twinkle_agentic/rollout/factory.py deleted file mode 100644 index eb6494d42..000000000 --- a/src/twinkle_agentic/rollout/factory.py +++ /dev/null @@ -1,74 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""One call that turns a generation backend into a :class:`Rollout`. - -Callers that only want turns appended to trajectories -- a challenger inventing -tasks, an evaluation script -- should not have to know that a local sampler and -an HTTP endpoint are driven by different classes with different required -arguments. They ask for a rollout, hand over whichever backend they happen to -have, and get something with the same contract: - - List[Trajectory] -> List[Trajectory] - -What the two still differ in is what ends up *inside* the trajectory, and no -factory can paper over it: the sampler path keeps ``input_ids`` / ``labels`` / -``logprobs`` and is the only one whose output can be trained on, while the API -path returns messages only. Pick the backend accordingly. -""" -from typing import Any, Dict, Optional - -from twinkle.data_format.sampling import SamplingParams -from .base import Rollout - -__all__ = ['build_rollout'] - - -def build_rollout( - backend: Any, - *, - template: Any = None, - tool_manager: Any = None, - sampling_params: Optional[SamplingParams] = None, - max_turns: int = 6, - trace_dir: Optional[str] = None, - **backend_kwargs: Any, -) -> Rollout: - """Build the multi-turn rollout that matches ``backend``. - - Args: - backend: an :class:`twinkle_agentic.protocol.base.API` (any - OpenAI-compatible endpoint) or a sampler exposing ``sample()``. - template: required for a sampler, rejected for an API. The sampler path - continues a conversation by splicing token ids, which needs the - local chat template; the API path re-sends messages as text. - tool_manager: optional for both. Without one the model is told there - are no tools. - backend_kwargs: passed straight to the chosen class -- e.g. ``harness`` - and ``max_trajectory_tokens`` for a sampler, ``concurrency`` and - ``extra_body`` for an API. An argument meant for the other backend - surfaces as a TypeError naming it. - """ - from twinkle_agentic.protocol.base import API - - common: Dict[str, Any] = { - 'tool_manager': tool_manager, - 'sampling_params': sampling_params, - 'max_turns': max_turns, - 'trace_dir': trace_dir, - } - - if isinstance(backend, API): - if template is not None: - raise ValueError('template is only used by the sampler path; an API ' - 'backend re-sends messages as text and never encodes ' - 'them locally.') - from .api_multi_turn import APIMultiTurnRollout - return APIMultiTurnRollout(api=backend, **common, **backend_kwargs) - - if not hasattr(backend, 'sample'): - raise TypeError(f'backend must be an API client or a sampler with a sample() ' - f'method, got {type(backend).__name__}') - if template is None: - raise ValueError('a sampler backend needs a template: the rollout appends each ' - 'turn as token ids and cannot re-encode the history.') - from .multi_turn import MultiTurnRollout - return MultiTurnRollout(sampler=backend, template=template, **common, **backend_kwargs) diff --git a/src/twinkle_agentic/rollout/multi_turn.py b/src/twinkle_agentic/rollout/multi_turn.py index cbeb8f380..6c9fd03d2 100644 --- a/src/twinkle_agentic/rollout/multi_turn.py +++ b/src/twinkle_agentic/rollout/multi_turn.py @@ -1,20 +1,49 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import json import re -from collections import defaultdict -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Any, Callable, Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, List, Literal, Optional, Tuple from twinkle.data_format import Trajectory -from twinkle.data_format.sampling import SampleResponse, SamplingParams +from twinkle.data_format.sampling import SampledSequence, SampleResponse, SamplingParams from twinkle.infra import remote_class, remote_function from twinkle.template.base import Template from twinkle_agentic.harness.base import AgentHarness +from twinkle_agentic.protocol.base import API from twinkle_agentic.tools.tool_manager import ToolManager -from .base import Rollout +from .api_sampler import APIGenerationError, APISampler +from .base import MAX_FOLLOWUPS, STOP_GENERATION_ERROR, Rollout from .bridge import _to_plain, extend_with_bridge +ResponseCallback = Callable[..., SampledSequence] + + +def _default_response_callback(sampler, api, sampling_params, *, input_feature, adapter_kwargs, + **kwargs) -> SampledSequence: + """Use the sampler when present, otherwise the API adapter.""" + if sampler is None: + if api is None: + raise ValueError('response_callback was omitted, but no sampler or API was provided') + return api(input_feature, sampling_params, **adapter_kwargs) + responses = sampler.sample([input_feature], sampling_params=sampling_params, **adapter_kwargs) + if not isinstance(responses, list): + raise TypeError(f'expected List[SampleResponse] from sampler.sample, got ' + f'{type(responses).__name__}') + if len(responses) != 1: + raise RuntimeError(f'sampler returned {len(responses)} responses for a single request; ' + 'expected exactly one.') + response = responses[0] + if not isinstance(response, SampleResponse): + raise TypeError(f'expected SampleResponse from sampler.sample, got ' + f'{type(response).__name__}') + if len(response.sequences) != 1: + raise RuntimeError(f'SampleResponse contains {len(response.sequences)} sequences; expected exactly one.') + sequence = response.sequences[0] + if not isinstance(sequence, SampledSequence): + raise TypeError(f'expected SampledSequence, got {type(sequence).__name__}') + return sequence + + def _append_only_delta( old_messages: List[Dict[str, Any]], new_messages: List[Dict[str, Any]], @@ -111,63 +140,118 @@ def _malformed_tool_message(errors: List[str]) -> Dict[str, Any]: @remote_class() class MultiTurnRollout(Rollout): - """Agentic multi-turn rollout with tool use (batched). + """Agentic multi-turn rollout with tool use, one episode per thread. Contract (matches :class:`Rollout`): accepts a ``List[Trajectory]`` and returns a ``List[Trajectory]`` of the same length, in the same order. - Every turn issues a SINGLE batched ``sampler.sample(active_pifs)`` call - so vLLM can run all live trajectories in parallel; finished trajectories - are parked and excluded from subsequent batches. Per-trajectory loop:: harness.before_generate # append-only after the first encode - sampler.sample(batch) # keep seq.new_input_feature + response_callback(...) # sampler or API -> SampledSequence harness.after_generate - ToolManager.call_many # Env.step_batch when tools share an Env + ToolManager.call_many # this turn's calls, one Env round trip harness.after_tools # format observations as tool messages extend_with_bridge # labels=-100; never decode-reencode history + Each trajectory runs its whole loop in its own thread. The callback may route + each turn to the sampler or the API adapter; either can overlap with other + trajectories while its thread waits on a GPU worker, endpoint, or sandbox. + + A supplied sampler must declare ``sample`` with ``enable_continous_work``. + Without it, ``slice_dp`` spreads each single-request call over every worker + and raises on ranks that receive nothing. + + Shared state: ``sampler``, API client and ``template`` are read-only during a + rollout and safe to share. A ``harness`` is not -- an ms-agent one delegates to an + ``LLMAgent`` that holds memory and context of its own -- so a batch of more + than one trajectory has to be given a 1:1 list of them; a single instance is + refused rather than shared. + Per-call overrides via ``**kwargs``: - * ``sampling_params``: shared :class:`SamplingParams` for the batch. + * ``sampling_params``: :class:`SamplingParams` for every episode. + * ``response_callback``: chooses a backend for each assistant turn and + returns one :class:`SampledSequence`. * ``tool_manager``: a single :class:`ToolManager` or a 1:1 list. - * ``harness``: a single :class:`AgentHarness` or a 1:1 list. Framework - specifics (ms-agent system/memory/tool-message shape) live in the - harness subclass, not here. + * ``harness``: a 1:1 list of :class:`AgentHarness` (a single instance + only for a batch of one). Framework specifics (ms-agent + system/memory/tool-message shape) live in the harness subclass, not + here. + * ``adapter_path`` / ``use_base_model``: see ``__init__``. * ``followup_fn``: see ``__init__``. """ def __init__( self, - sampler, - template: Template, + sampler=None, + template: Optional[Template] = None, tool_manager: Optional[ToolManager] = None, sampling_params: Optional[SamplingParams] = None, max_turns: int = 6, max_trajectory_tokens: Optional[int] = None, + concurrency: Optional[int] = None, trace_dir: Optional[str] = None, trace_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, success_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, harness: Optional[AgentHarness] = None, adapter_path: Optional[str] = None, + use_base_model: bool = False, stop_after_stuck_turns: int = 0, max_malformed_retries: int = 2, followup_fn: Optional[Callable[[Trajectory, int], Any]] = None, + api: Optional[API] = None, + response_callback: Optional[ResponseCallback] = None, + api_appended_as: Literal['demonstration', 'context'] = 'demonstration', + api_kwargs: Optional[Dict[str, Any]] = None, ): super().__init__() + if isinstance(sampler, (API, APISampler)): + if api is not None: + raise ValueError('the positional backend and api= both specify an API') + api, sampler = sampler, None if template is None: raise ValueError('MultiTurnRollout requires a local Template instance') + if response_callback is None and sampler is None and api is None: + raise ValueError('MultiTurnRollout requires a sampler or API when response_callback is omitted') + if sampler is not None: + sample = getattr(type(sampler), 'sample', None) + if sample is None: + raise TypeError(f'backend must be an API or sampler, got {type(sampler).__name__}') + if not getattr(sample, '_enable_continous_work', False): + raise ValueError( + f'{type(sampler).__name__}.sample must be declared with ' + 'enable_continous_work=True: this rollout samples one trajectory per ' + 'call, and a slice_dp sampler raises when a worker gets nothing from ' + 'a batch of one.') + if adapter_path and use_base_model: + raise ValueError('adapter_path and use_base_model=True ask for opposite ' + 'weights; the sampler would drop the adapter silently.') if max_trajectory_tokens is not None and max_trajectory_tokens < 1: raise ValueError(f'max_trajectory_tokens must be >= 1 or None, got ' f'{max_trajectory_tokens}') self._init_common( max_turns=max_turns, sampling_params=sampling_params, + concurrency=concurrency, trace_dir=trace_dir, trace_callback=trace_callback, success_callback=success_callback) self.sampler = sampler self.template = template + if isinstance(api, APISampler): + if api_kwargs: + raise ValueError('api_kwargs belongs on the APISampler when api= is already adapted') + if api.template is not template: + raise ValueError('MultiTurnRollout and APISampler must share the same template instance') + self.api = api + elif api is not None: + self.api = APISampler( + api, template, appended_as=api_appended_as, api_kwargs=api_kwargs) + else: + if api_kwargs: + raise ValueError('api_kwargs requires an API backend') + self.api = None + self.response_callback = response_callback or _default_response_callback self.tool_manager = tool_manager self.harness = harness # A LoRA directory on disk, forwarded to every sample call. Training syncs @@ -175,6 +259,11 @@ def __init__( # such channel: without this, an eval script would silently measure the # base model and report it as the trained one. self.adapter_path = adapter_path + # The other direction: force the base weights. Needed because a sampler + # mid-training falls back to the LoRA synced into it whenever a call names + # no adapter, so a utility rollout (summarizing, judging) that wants the + # untrained model has to say so rather than stay silent. + self.use_base_model = use_base_model self.max_trajectory_tokens = max_trajectory_tokens # How many stuck turns in a row end the episode; 0 runs to ``max_turns`` # regardless. A turn is stuck when it made no progress at all, which is @@ -242,391 +331,363 @@ def __init__( @remote_function() def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory]: - if isinstance(trajectories, dict): - raise TypeError('MultiTurnRollout.__call__ expects a List[Trajectory]; ' - 'wrap a single trajectory as [trajectory].') - trajectories = list(trajectories) - n = len(trajectories) - if n == 0: - return [] + """The base implementation; the decorator is what a deployed handle needs.""" + return super().__call__(trajectories, **kwargs) - sampling_params = kwargs.get('sampling_params', self.sampling_params) + def _resolve_call(self, kwargs: Dict[str, Any], n: int) -> Dict[str, Any]: adapter_path = kwargs.get('adapter_path', self.adapter_path) # Left out entirely when unset, so a sampler without LoRA enabled sees the # same call it always did. adapter_kwargs = {'adapter_path': adapter_path} if adapter_path else {} - tool_managers = self._broadcast( - kwargs.get('tool_manager', self.tool_manager), n, name='tool_manager', required=True) - harnesses = self._broadcast(kwargs.get('harness', self.harness), n, name='harness') - lives: List[Optional[Trajectory]] = [ - dict(trajectories[i]) if harnesses[i] is not None else None for i in range(n) - ] - for live in lives: - if live is not None: - live['messages'] = list(live.get('messages') or []) + if kwargs.get('use_base_model', self.use_base_model): + adapter_kwargs['use_base_model'] = True + sampling_params = kwargs.get('sampling_params', self.sampling_params) + if sampling_params.num_samples != 1: + raise ValueError(f'MultiTurnRollout supports num_samples=1 only, got ' + f'{sampling_params.num_samples}') + response_callback = kwargs.get('response_callback', self.response_callback) + if not callable(response_callback): + raise TypeError('response_callback must be callable') + return { + 'sampling_params': sampling_params, + 'adapter_kwargs': adapter_kwargs, + 'response_callback': response_callback, + 'tool_managers': self._broadcast( + kwargs.get('tool_manager', self.tool_manager), n, name='tool_manager'), + 'harnesses': self._broadcast( + kwargs.get('harness', self.harness), n, name='harness', per_trajectory=True), + 'followup_fn': kwargs.get('followup_fn', self.followup_fn), + } + + def _run_one(self, trajectory: Trajectory, index: int, ctx: Dict[str, Any]) -> Trajectory: + tool_manager: ToolManager = ctx['tool_managers'][index] + harness: Optional[AgentHarness] = ctx['harnesses'][index] + followup_fn = ctx['followup_fn'] + adapter_kwargs: Dict[str, Any] = ctx['adapter_kwargs'] + response_callback: ResponseCallback = ctx['response_callback'] # 1. First before_generate happens *before* encode so memory/system # injection is in the initial prefix (not a later rewrite). - encode_trajs: List[Trajectory] = [] - for i, traj in enumerate(trajectories): - h, live = harnesses[i], lives[i] - if h is not None and live is not None: - lives[i] = h.before_generate(live) - live = lives[i] - traj = dict(traj) - traj['messages'] = list(live.get('messages') or []) - if live.get('tools'): - traj['tools'] = list(live['tools']) - encode_trajs.append(traj) - - pifs: List[Dict[str, Any]] = [] - for i, traj in enumerate(encode_trajs): - pif = self.template.encode(traj, add_generation_prompt=True) - pif = _to_plain(pif) - pif.setdefault('messages', list(traj.get('messages', []))) - pifs.append(pif) - if lives[i] is not None: - lives[i]['messages'] = list(pifs[i].get('messages') or []) - - all_logprobs: List[List[Any]] = [[] for _ in range(n)] - stop_reasons: List[Optional[str]] = [None] * n - turns: List[int] = [0] * n - truncated: List[bool] = [False] * n - done: List[bool] = [False] * n - # Consecutive turns that made no progress, the calls already issued in - # each episode, and whether being stuck is what ended it. All three stay - # at their initial value when ``stop_after_stuck_turns`` is 0. - stuck_turns: List[int] = [0] * n - seen_calls: List[set] = [set() for _ in range(n)] - stuck_stop: List[bool] = [False] * n + live: Optional[Trajectory] = None + to_encode = trajectory + if harness is not None: + live = dict(trajectory) + live['messages'] = list(live.get('messages') or []) + live = harness.before_generate(live) + to_encode = dict(trajectory) + to_encode['messages'] = list(live.get('messages') or []) + if live.get('tools'): + to_encode['tools'] = list(live['tools']) + + pif = _to_plain(self.template.encode(to_encode, add_generation_prompt=True)) + pif.setdefault('messages', list(to_encode.get('messages') or [])) + if 'tools' in to_encode: + pif['tools'] = list(to_encode.get('tools') or []) + elif tool_manager is not None: + pif['tools'] = list(tool_manager.tool_infos() or []) + if live is not None: + live['messages'] = list(pif.get('messages') or []) + + logprobs: List[Any] = [] + stop_reason: Optional[str] = None + generation_error: Optional[str] = None + turns = 0 + truncated = False + params = ctx['sampling_params'] + # Consecutive turns that made no progress, the calls already issued, and + # whether being stuck is what ended the episode. All three stay at their + # initial value when ``stop_after_stuck_turns`` is 0. + stuck_turns = 0 + seen_calls: set = set() + stuck_stop = False # Replies in a row whose tool-call markup did not parse. Reset by any # reply that produced a call, so one bad escape in the middle of a # working episode does not count against a later one. - malformed_turns: List[int] = [0] * n - # Follow-up bookkeeping (all no-ops when ``followup_fn`` is None): - # how many follow-ups each trajectory has had, and the params its next - # turn should use. A trajectory that has had one stops dispatching tools. - followups: List[int] = [0] * n - params_for: List[Any] = [sampling_params] * n - followup_fn = kwargs.get('followup_fn', self.followup_fn) - # Why the tool-calling part of each episode ended, when it was not the - # model's own choice: 'max_turns' or 'stuck'. Reported separately from - # ``truncated`` because an episode can hit the turn limit and still go on - # to answer the follow-up stages, in which case nothing was cut off. - tool_stop: List[Optional[str]] = [None] * n - - def append_followup(global_idx: int) -> bool: + malformed_turns = 0 + followups = 0 + # Why the tool-calling part ended, when it was not the model's own + # choice: 'max_turns' or 'stuck'. Reported separately from ``truncated`` + # because an episode can hit the turn limit and still go on to answer the + # follow-up stages, in which case nothing was cut off. + tool_stop: Optional[str] = None + # The loop counts generations, and each granted follow-up buys the one + # extra generation it asked for. Paying for the follow-up stages out of + # ``max_turns`` would mean an episode that spent its whole tool budget + # never reaches the stages that read what it built, and a short one + # silently gets more tool turns than a long one. + budget = self.max_turns + spent = 0 + + def grant_followup() -> bool: """Ask for one more stage; True when the episode carries on. Sets ``truncated`` itself in the one case where the answer is "there is no room for another stage", which is a cut trajectory rather than a caller that had nothing more to ask. """ - nonlocal iterations - if followup_fn is None: + nonlocal pif, live, followups, budget, params, truncated + if followup_fn is None or followups >= MAX_FOLLOWUPS: return False followup = followup_fn( - self._as_trajectory(trajectories[global_idx], pifs[global_idx], - all_logprobs[global_idx], turns[global_idx], - stop_reasons[global_idx], truncated[global_idx]), - followups[global_idx]) + self._as_trajectory(trajectory, pif, logprobs, turns, stop_reason, truncated), followups) if followup is None: return False - text, next_params = followup if isinstance(followup, tuple) else (followup, None) - extended = extend_with_bridge( - pifs[global_idx], [{'role': 'user', 'content': text}], self.template) + text, next_params = self._unpack_followup(followup) + extended = extend_with_bridge(pif, [{'role': 'user', 'content': text}], self.template) if extended is None: - truncated[global_idx] = True + truncated = True return False - pifs[global_idx] = extended - if lives[global_idx] is not None: - lives[global_idx]['messages'] = list(extended.get('messages') or []) - followups[global_idx] += 1 - iterations += 1 + pif = extended + # Follow-up stages are answers, so an API must not see tool schemas. + pif['tools'] = [] + if live is not None: + live['messages'] = list(extended.get('messages') or []) + followups += 1 + budget += 1 if next_params is not None: - params_for[global_idx] = next_params + params = next_params return True - # The loop counts generations, and each granted follow-up buys the one - # extra generation it asked for. Paying for the follow-up stages out of - # ``max_turns`` would mean an episode that spent its whole tool budget - # never reaches the stages that read what it built, and a short one - # silently gets more tool turns than a long one. - iterations = self.max_turns - done_iterations = 0 - first_turn = True - while done_iterations < iterations: - done_iterations += 1 - active = [i for i in range(n) if not done[i]] - if not active: - break + while spent < budget: + spent += 1 - if not first_turn: - for global_idx in active: - pifs[global_idx], lives[global_idx], dropped = self._harness_before_generate( - pifs[global_idx], lives[global_idx], harnesses[global_idx]) - if dropped: - truncated[global_idx] = True - done[global_idx] = True - active = [i for i in range(n) if not done[i]] - if not active: + if spent > 1: + pif, live, dropped = self._harness_before_generate(pif, live, harness) + if dropped: + truncated = True break - first_turn = False - - # 2. One batched sample call per distinct SamplingParams among the - # live trajectories -- normally exactly one, since only a - # follow-up stage asks for its own budget. Grouping rather than - # taking the first is what keeps a mixed batch honest: sampling one - # trajectory under another's token limit would silently truncate or - # over-spend, and the two are indistinguishable afterwards. - groups: List[List[int]] = [] - group_params: List[Any] = [] - for global_idx in active: - for slot, params in enumerate(group_params): - if params is params_for[global_idx]: - groups[slot].append(global_idx) + + # 2. One request. The callback chooses the local sampler or the API + # adapter, but both paths return exactly one SampledSequence. + try: + seq = response_callback( + self.sampler, + self.api, + params, + input_feature=pif, + adapter_kwargs=adapter_kwargs, + trajectory=trajectory, + trajectory_index=index, + turn=turns + 1, + followups=followups, + ) + except APIGenerationError as exc: + stop_reason = STOP_GENERATION_ERROR + generation_error = str(exc) + truncated = True + break + if not isinstance(seq, SampledSequence): + raise TypeError(f'response_callback must return SampledSequence, got ' + f'{type(seq).__name__}') + turns += 1 + + if seq.new_input_feature is None or 'input_ids' not in seq.new_input_feature: + raise RuntimeError(f'Sampler returned a SampledSequence without ' + f'new_input_feature.input_ids for trajectory ' + f'{index}; cannot continue multi-turn.') + + pif = _to_plain(dict(seq.new_input_feature)) + if seq.logprobs is not None: + if len(seq.logprobs) != len(seq.tokens): + raise RuntimeError(f'logprobs length ({len(seq.logprobs)}) does not ' + f'match sampled token count ({len(seq.tokens)}) ' + f'at turn {turns} (trajectory {index})') + logprobs.extend(seq.logprobs) + stop_reason = seq.stop_reason + + msgs = pif.get('messages') or [] + last_msg = msgs[-1] if msgs else None + tool_calls = (last_msg.get('tool_calls') if isinstance(last_msg, dict) else None) + if not tool_calls: + tool_calls = self.template.parse_tool_call(seq.decoded or '') + # After a follow-up, a parsed call is not a call: the tools were + # withdrawn for these stages on purpose (see ``followup_fn``), and + # dispatching python that the model wrote as *an answer* would edit + # the state the answer is about. + if followups: + tool_calls = None + # The parse also *rewrote* the message: when a reply parses as + # a call, the template stores it with the call text removed, so + # a caller reading the message gets less than the model wrote. + # For these stages the reply is the deliverable, and one of the + # tool-call formats is XML-shaped, so a check script asserting + # the content of an .xml file matches it: 5 of ex12's 72 check + # scripts came back with the XML cut out of them -- three then + # ran with `content == ''` where the model had written the file's + # real text, and two no longer held a code block at all. + if msgs and isinstance(last_msg, dict): + # Decoded without the special tokens, the way the template + # writes a message: ``seq.decoded`` keeps the closing + # ``<|im_end|>``, and putting that in the content put it in + # the problem statements ex13 handed to solvers -- 7 of 7 of + # them ended in a literal '<|im_end|>'. + tok = getattr(self.template, 'tokenizer', None) + if tok is not None and seq.tokens: + last_msg['content'] = tok.decode(seq.tokens, skip_special_tokens=True) + else: + last_msg['content'] = seq.decoded or '' + last_msg.pop('tool_calls', None) + + if live is not None: + live['messages'] = list(msgs) + if harness is not None and live is not None: + live = harness.after_generate(live, seq.decoded or '', tool_calls or []) + self._merge_assistant_metadata(pif, live) + + # 3. Termination conditions + # A reply cut off at ``max_tokens`` is truncated in exactly the sense + # the flag names, and consumers read the flag to tell a trajectory + # that finished from one that ran out of room: a difficulty + # measurement counting such an attempt as a genuine failure blames + # the task for the token budget. Tool calls the cut reply happens to + # contain are still not dispatched -- the turn never got to decide it + # was done emitting them. + if seq.stop_reason == 'length': + truncated = True + break + + # 3a. Sequence-length cap. + if (self.max_trajectory_tokens is not None + and len(pif.get('input_ids') or []) >= self.max_trajectory_tokens): + truncated = True + break + + if not tool_calls: + # Markup that did not parse is the model asking for a tool, not + # declining one -- ending here tells it nothing and throws the + # turn away. Hand back the parser's own reason and let it write + # the call again. Not after a follow-up: tools are withdrawn + # there on purpose (see ``followup_fn``), so a reply that looks + # like a call is meant to be read as text. + parse_errors = ([] if followups else self.template.tool_call_errors(seq.decoded or '')) + if parse_errors and malformed_turns < self.max_malformed_retries: + malformed_turns += 1 + extended = extend_with_bridge(pif, [_malformed_tool_message(parse_errors)], self.template) + if extended is None: + truncated = True break - else: - group_params.append(params_for[global_idx]) - groups.append([global_idx]) - - resps_by_idx: Dict[int, Any] = {} - device_mesh = getattr(self.sampler, 'device_mesh', None) - min_batch_size = (device_mesh.data_world_size if device_mesh is not None else 1) - # A sampler that routes each request on its own accepts a batch smaller - # than its worker count, so the padding below is not needed. It was only - # ever there because slicing a batch over all workers raises when some - # rank gets nothing, and the duplicates it added were generated and then - # dropped -- with one prompt and 8 workers that is 8 generations for 1 - # kept result. - if getattr(type(self.sampler).sample, '_enable_continous_work', False): - min_batch_size = 1 - for slot, group in enumerate(groups): - batch_pifs = [pifs[i] for i in group] - actual = len(batch_pifs) - if actual < min_batch_size: - batch_pifs = batch_pifs + ([batch_pifs[-1]] * (min_batch_size - actual)) - group_resps = self.sampler.sample(batch_pifs, - sampling_params=group_params[slot], - **adapter_kwargs) - group_resps = self._unwrap_response_list(group_resps, len(batch_pifs))[:actual] - for local_idx, global_idx in enumerate(group): - resps_by_idx[global_idx] = group_resps[local_idx] - - pending_tools: List[tuple] = [] # (global_idx, tool_calls) - for global_idx in active: - turns[global_idx] += 1 - seq = resps_by_idx[global_idx].sequences[0] - - if seq.new_input_feature is None or 'input_ids' not in seq.new_input_feature: - raise RuntimeError(f'Sampler returned a SampledSequence without ' - f'new_input_feature.input_ids for trajectory ' - f'{global_idx}; cannot continue multi-turn.') - - pifs[global_idx] = _to_plain(dict(seq.new_input_feature)) - if seq.logprobs is not None: - if len(seq.logprobs) != len(seq.tokens): - raise RuntimeError(f'logprobs length ({len(seq.logprobs)}) does not ' - f'match sampled token count ({len(seq.tokens)}) ' - f'at turn {turns[global_idx]} ' - f'(trajectory {global_idx})') - all_logprobs[global_idx].extend(seq.logprobs) - stop_reasons[global_idx] = seq.stop_reason - - _msgs = pifs[global_idx].get('messages') or [] - _last_msg = _msgs[-1] if _msgs else None - tool_calls = (_last_msg.get('tool_calls') if isinstance(_last_msg, dict) else None) - if not tool_calls: - tool_calls = self.template.parse_tool_call(seq.decoded or '') - # After a follow-up, a parsed call is not a call: the tools were - # withdrawn for these stages on purpose (see ``followup_fn``), and - # dispatching python that the model wrote as *an answer* would edit - # the state the answer is about. - if followups[global_idx]: - tool_calls = None - # The parse also *rewrote* the message: when a reply parses as - # a call, the template stores it with the call text removed, so - # a caller reading the message gets less than the model wrote. - # For these stages the reply is the deliverable, and one of the - # tool-call formats is XML-shaped, so a check script asserting - # the content of an .xml file matches it: 5 of ex12's 72 check - # scripts came back with the XML cut out of them -- three then - # ran with `content == ''` where the model had written the file's - # real text, and two no longer held a code block at all. - if _msgs and isinstance(_last_msg, dict): - # Decoded without the special tokens, the way the - # template writes a message: ``seq.decoded`` keeps the - # closing ``<|im_end|>``, and putting that in the content - # put it in the problem statements ex13 handed to solvers - # -- 7 of 7 of them ended in a literal '<|im_end|>'. - tok = getattr(self.template, 'tokenizer', None) - if tok is not None and seq.tokens: - _last_msg['content'] = tok.decode( - seq.tokens, skip_special_tokens=True) - else: - _last_msg['content'] = seq.decoded or '' - _last_msg.pop('tool_calls', None) - - if lives[global_idx] is not None: - lives[global_idx]['messages'] = list(_msgs) - if harnesses[global_idx] is not None and lives[global_idx] is not None: - lives[global_idx] = harnesses[global_idx].after_generate( - lives[global_idx], seq.decoded or '', tool_calls or []) - self._merge_assistant_metadata(pifs[global_idx], lives[global_idx]) - - # 3. Termination conditions - # A reply cut off at ``max_tokens`` is truncated in exactly the - # sense the flag names, and consumers read the flag to tell a - # trajectory that finished from one that ran out of room: a - # difficulty measurement counting such an attempt as a genuine - # failure blames the task for the token budget. Tool calls the - # cut reply happens to contain are still not dispatched -- the - # turn never got to decide it was done emitting them. - if seq.stop_reason == 'length': - truncated[global_idx] = True - done[global_idx] = True + pif = extended + if live is not None: + live['messages'] = list(extended.get('messages') or []) continue - - # 3a. Sequence-length cap. - if (self.max_trajectory_tokens is not None - and len(pifs[global_idx].get('input_ids') or []) >= self.max_trajectory_tokens): - truncated[global_idx] = True - done[global_idx] = True + # The episode is over as far as the model is concerned. Give the + # caller one chance to say otherwise -- see ``followup_fn`` for + # why this is not a second rollout. + if grant_followup(): continue + break - if not tool_calls: - # Markup that did not parse is the model asking for a tool, - # not declining one -- ending here tells it nothing and throws - # the turn away. Hand back the parser's own reason and let it - # write the call again. Not after a follow-up: tools are - # withdrawn there on purpose (see ``followup_fn``), so a reply - # that looks like a call is meant to be read as text. - parse_errors = ([] if followups[global_idx] else - self.template.tool_call_errors(seq.decoded or '')) - if (parse_errors and malformed_turns[global_idx] < self.max_malformed_retries): - malformed_turns[global_idx] += 1 - extended = extend_with_bridge(pifs[global_idx], - [_malformed_tool_message(parse_errors)], - self.template) - if extended is None: - truncated[global_idx] = True - done[global_idx] = True - continue - pifs[global_idx] = extended - if lives[global_idx] is not None: - lives[global_idx]['messages'] = list(extended.get('messages') or []) - continue - # The episode is over as far as the model is concerned. Give - # the caller one chance to say otherwise -- see - # ``followup_fn`` for why this is not a second rollout. - if append_followup(global_idx): - continue - done[global_idx] = True + if turns >= self.max_turns: + # Out of tool turns, not out of episode: the stages that read the + # end state can still run on what was built. + tool_stop = 'max_turns' + if grant_followup(): continue + truncated = True + break - if turns[global_idx] >= self.max_turns: - # Out of tool turns, not out of episode: the stages that read - # the end state can still run on what was built. - tool_stop[global_idx] = 'max_turns' - if append_followup(global_idx): - continue - truncated[global_idx] = True - done[global_idx] = True + malformed_turns = 0 + + # 4. This turn's calls, then the harness formats the observations + # into tool messages (append-only bridge). + if tool_manager is None: + raise ValueError('the model emitted tool_calls but this trajectory has no ToolManager') + observations = self._run_tools(tool_manager, tool_calls) + if self.stop_after_stuck_turns: + keys = [_call_key(tc) for tc in tool_calls] + all_repeats = bool(keys) and all(k in seen_calls for k in keys) + seen_calls.update(keys) + all_errors = bool(observations) and all(is_error_observation(o) for o in observations) + if all_errors or all_repeats: + stuck_turns += 1 + else: + stuck_turns = 0 + + tool_messages, live = self._tool_messages_after(pif, live, harness, observations, tool_calls) + extended = extend_with_bridge(pif, tool_messages, self.template) + overflowed = extended is None + if overflowed: + # Trajectory exceeded max_length. + truncated = True + else: + pif = extended + if live is not None: + live['messages'] = list(extended.get('messages') or []) + # Checked after the messages are appended, so the turns that ended + # the episode are in the trajectory the caller reads. + if self.stop_after_stuck_turns and stuck_turns >= self.stop_after_stuck_turns: + stuck_stop = True + tool_stop = 'stuck' + # Same as the turn limit: the tool phase is over, the state it + # left is not, so the stages still get their turn. + if not overflowed and grant_followup(): continue + truncated = True + break + if overflowed: + break - malformed_turns[global_idx] = 0 - pending_tools.append((global_idx, list(tool_calls))) - - # 4. Parallel tool dispatch across the live batch, then harness - # formats observations into tool messages (append-only bridge). - # The bridge itself is computed serially: it is a cheap - # decode-diff-encode on python strings / token lists. - if pending_tools: - obs_by_traj = self._dispatch_tools(tool_managers, pending_tools) - for global_idx, tool_calls in pending_tools: - observations = obs_by_traj.get(global_idx) or [''] * len(tool_calls) - if self.stop_after_stuck_turns: - keys = [_call_key(tc) for tc in tool_calls] - all_repeats = bool(keys) and all(k in seen_calls[global_idx] - for k in keys) - seen_calls[global_idx].update(keys) - all_errors = bool(observations) and all( - is_error_observation(o) for o in observations) - if all_errors or all_repeats: - stuck_turns[global_idx] += 1 - else: - stuck_turns[global_idx] = 0 - tool_messages, lives[global_idx] = self._tool_messages_after( - pifs[global_idx], lives[global_idx], harnesses[global_idx], - observations, tool_calls) - extended = extend_with_bridge(pifs[global_idx], tool_messages, self.template) - if extended is None: - # Trajectory exceeded max_length, mark as done (deleted) - truncated[global_idx] = True - done[global_idx] = True - else: - pifs[global_idx] = extended - if lives[global_idx] is not None: - lives[global_idx]['messages'] = list(extended.get('messages') or []) - # Checked after the messages are appended, so the turns that - # ended the episode are in the trajectory the caller reads. - if (self.stop_after_stuck_turns - and stuck_turns[global_idx] >= self.stop_after_stuck_turns): - stuck_stop[global_idx] = True - tool_stop[global_idx] = 'stuck' - # Same as the turn limit: the tool phase is over, the - # state it left is not, so the stages still get their turn. - if not done[global_idx] and append_followup(global_idx): - continue - truncated[global_idx] = True - done[global_idx] = True - - for i in range(n): - if not all_logprobs[i]: - continue - labels_i = pifs[i].get('labels') or [] - trainable_i = sum(1 for label in labels_i if label != -100) - if len(all_logprobs[i]) != trainable_i: - raise RuntimeError(f'logprobs/labels misaligned for trajectory {i}: ' - f'{len(all_logprobs[i])} logprobs vs {trainable_i} ' - f'trainable labels (labels != -100). This invariant is ' - f'required by grpo._pad_and_align_to_batch; a mismatch ' - f'would silently corrupt GRPO old_logps alignment.') - - # 5. Merge pif fields into each trajectory dict at TOP LEVEL so - # downstream consumers (VLLMSampler with ``'input_ids' in inputs``) - # see an encoded InputFeature and skip re-encoding. - outs: List[Trajectory] = [] - for i, traj in enumerate(trajectories): - out = dict(traj) - out.update(pifs[i]) - out['messages'] = list(pifs[i].get('messages') or out.get('messages', [])) - out['logprobs'] = all_logprobs[i] if all_logprobs[i] else None - out['turns'] = turns[i] - out['stop_reason'] = stop_reasons[i] - out['truncated'] = truncated[i] - # ``truncated`` says something was cut off; these two say what ended - # the tool-calling part, which is a different question -- an episode - # can run out of turns, be handed a follow-up stage, and finish it. - out['stuck_stop'] = stuck_stop[i] - out['tool_stop'] = tool_stop[i] - out['followups'] = followups[i] - outs.append(out) - - # Per-rollout trace dump: one JSON file per selected trajectory. - # ``trace_callback`` decides whether to store; ``success_callback`` - # decides the filename prefix. Observability only -- any failure - # is swallowed inside ``_write_rollout_traces``. - if self.trace_dir: - self._write_rollout_traces(outs, global_step=kwargs.get('global_step')) - return outs + if logprobs: + labels = pif.get('labels') or [] + completion_mask = pif.get('completion_mask') + if completion_mask is None: + expected = sum(1 for label in labels if label != -100) + elif len(completion_mask) != len(labels): + raise RuntimeError(f'completion_mask/labels misaligned for trajectory {index}: ' + f'{len(completion_mask)} != {len(labels)}') + else: + expected = sum(1 for label, flag in zip(labels, completion_mask) if label != -100 and flag) + if len(logprobs) != expected: + raise RuntimeError(f'logprobs/policy-token alignment failed for trajectory {index}: ' + f'{len(logprobs)} logprobs vs {expected} positions selected by ' + '(labels != -100) & completion_mask.') + + # 5. Merge pif fields into the trajectory dict at TOP LEVEL so downstream + # consumers (VLLMSampler with ``'input_ids' in inputs``) see an encoded + # InputFeature and skip re-encoding. + out = dict(trajectory) + out.update(pif) + out['messages'] = list(pif.get('messages') or out.get('messages', [])) + out['logprobs'] = logprobs if logprobs else None + out['turns'] = turns + out['stop_reason'] = stop_reason + out['truncated'] = truncated + # ``truncated`` says something was cut off; these two say what ended the + # tool-calling part, which is a different question -- an episode can run + # out of turns, be handed a follow-up stage, and finish it. + out['stuck_stop'] = stuck_stop + out['tool_stop'] = tool_stop + out['followups'] = followups + if generation_error is not None: + out['error'] = generation_error + return out # ------------------------------------------------------------------ private @staticmethod - def _as_trajectory(traj: Trajectory, pif: Dict[str, Any], logprobs: List[Any], - turns: int, stop_reason: Optional[str], - truncated: bool) -> Trajectory: + def _run_tools(tool_manager: ToolManager, tool_calls: List[Dict[str, Any]]) -> List[str]: + """Run one turn's calls, through ``call_many`` when the manager has it. + + A turn's calls go together because they share one Env round trip + (``Env.step_batch``). Calls from *different* trajectories no longer meet + here -- each episode has its own thread and, in the sandbox case, its own + Env -- so there is nothing left to group across. + + A manager that answers with fewer results than calls leaves the rest + empty rather than shifting them onto the wrong call. + """ + if hasattr(tool_manager, 'call_many'): + contents = tool_manager.call_many(tool_calls) + else: + contents = [tool_manager(tc) for tc in tool_calls] + obs = [''] * len(tool_calls) + for i, content in enumerate(contents[:len(tool_calls)]): + obs[i] = '' if content is None else str(content) + return obs + + @staticmethod + def _as_trajectory(traj: Trajectory, pif: Dict[str, Any], logprobs: List[Any], turns: int, + stop_reason: Optional[str], truncated: bool) -> Trajectory: """The episode so far, shaped like the value ``__call__`` returns. Handed to ``followup_fn`` so the callback reads an episode the same way @@ -683,55 +744,6 @@ def _merge_assistant_metadata(pif: Dict[str, Any], live: Trajectory) -> None: if last_asst.get(key) and not dst.get(key): dst[key] = last_asst[key] - def _dispatch_tools( - self, - tool_managers: List[ToolManager], - pending: List[Tuple[int, List[Dict[str, Any]]]], - ) -> Dict[int, List[str]]: - """Run tool calls for the live batch, grouped by ToolManager. - - Trajectories that share a manager (and therefore often one Env) go - through ``call_many`` / ``Env.step_batch``. Distinct managers run - concurrently so remote sandboxes are not serialized on generate. - """ - obs: Dict[int, List[str]] = { - gi: [''] * len(tcs) for gi, tcs in pending - } - groups: Dict[int, List[Tuple[int, int, Dict[str, Any]]]] = defaultdict(list) - mgr_by_id: Dict[int, ToolManager] = {} - for gi, tcs in pending: - mid = id(tool_managers[gi]) - mgr_by_id[mid] = tool_managers[gi] - for ci, tc in enumerate(tcs): - groups[mid].append((gi, ci, tc)) - - def _run_group(items: List[Tuple[int, int, Dict[str, Any]]], mgr: ToolManager): - tcs = [tc for _, _, tc in items] - if hasattr(mgr, 'call_many'): - contents = mgr.call_many(tcs) - else: - contents = [mgr(tc) for tc in tcs] - return list(zip(items, contents)) - - group_items = list(groups.items()) - if len(group_items) == 1: - mid, items = group_items[0] - finished = [_run_group(items, mgr_by_id[mid])] - else: - finished = [] - with ThreadPoolExecutor(max_workers=min(32, len(group_items))) as pool: - futs = [ - pool.submit(_run_group, items, mgr_by_id[mid]) - for mid, items in group_items - ] - for fut in as_completed(futs): - finished.append(fut.result()) - - for group_result in finished: - for (gi, ci, _tc), content in group_result: - obs[gi][ci] = '' if content is None else str(content) - return obs - def _tool_messages_after( self, pif: Dict[str, Any], @@ -750,22 +762,3 @@ def _tool_messages_after( if not delta: return fallback, live return delta, live - - @staticmethod - def _unwrap_response_list(resps, expected: int) -> List[SampleResponse]: - """Validate that the sampler returned ``expected`` ``SampleResponse``s, - one per input in the batch. - """ - if not isinstance(resps, list): - raise TypeError(f'expected List[SampleResponse] from sampler.sample (batched ' - f'call), got {type(resps).__name__}') - if len(resps) != expected: - raise RuntimeError(f'sampler returned {len(resps)} responses for a batch of ' - f'{expected} trajectories; expected one per input.') - for i, r in enumerate(resps): - if not isinstance(r, SampleResponse): - raise TypeError(f'expected SampleResponse at batch index {i}, got ' - f'{type(r).__name__}') - if not r.sequences: - raise RuntimeError(f'SampleResponse at batch index {i} has no sequences') - return resps diff --git a/src/twinkle_agentic/summarizer/base.py b/src/twinkle_agentic/summarizer/base.py index 8e9be5f5f..aecc25675 100644 --- a/src/twinkle_agentic/summarizer/base.py +++ b/src/twinkle_agentic/summarizer/base.py @@ -5,11 +5,12 @@ import re from typing import TYPE_CHECKING, Any, Sequence +from twinkle_agentic.rollout import MultiTurnRollout from twinkle_agentic.utils.llm_backup import llm_backup +from twinkle_agentic.utils.message_utils import assistant_text if TYPE_CHECKING: from twinkle.data_format import SamplingParams, Trajectory # noqa: F401 - from twinkle.sampler.base import Sampler # noqa: F401 DEFAULT_USER_PROMPT_TEMPLATE = """\ @@ -36,7 +37,8 @@ class Summarizer: - LLM_BACKUP_BASE_URL: API endpoint Args: - sampler: Student model sampler (local inference, shared across types). + backend: a sampler or an API client, driven through + :class:`~twinkle_agentic.rollout.MultiTurnRollout`. compression_ratio: Target compression factor (> 1). model_path: Model identifier. sampling_params: Default sampling params. @@ -44,14 +46,19 @@ class Summarizer: user_prompt_template: User prompt template. Must contain ``{budget}`` and ``{text}``. May contain ``{query}``. min_budget_chars: Floor for the character budget in the prompt. - template: Optional :class:`Template` for special token stripping. + template: local :class:`Template`, required by the sampler path and also + what special-token stripping reads its tokenizer from. lora_path: LoRA adapter path specific to this summarizer type. - Each subclass can use a different LoRA for its task. + Each subclass can use a different LoRA for its task. Without one the + base weights are asked for explicitly -- a sampler mid-training + otherwise lends this out the policy LoRA synced into it. + rollout_kwargs: passed to ``MultiTurnRollout``. API request options + belong in ``api_kwargs``. """ def __init__( self, - sampler: Sampler, + backend: Any, compression_ratio: float = 2.0, *, model_path: str = '', @@ -61,9 +68,10 @@ def __init__( min_budget_chars: int = 250, template: Any | None = None, lora_path: str | None = None, + **rollout_kwargs: Any, ): - if sampler is None: - raise ValueError('sampler is required') + if backend is None: + raise ValueError('backend is required') if compression_ratio <= 1.0: raise ValueError(f'compression_ratio must be > 1, got {compression_ratio}') if min_budget_chars < 1: @@ -74,7 +82,6 @@ def __init__( raise ValueError('user_prompt_template must contain both {budget} and {text}') self.model_path = model_path - self.sampler = sampler self.compression_ratio = float(compression_ratio) self.sampling_params = sampling_params self.system_prompt = system_prompt @@ -83,6 +90,20 @@ def __init__( self.template = template self.lora_path = lora_path if lora_path else None self._special_tokens_cache: tuple[str, ...] | None = None + # Built on the first call rather than here, so a summarizer that never + # compresses anything (every text already under budget) costs nothing. + self._backend = backend + self._rollout_kwargs = dict(rollout_kwargs, max_turns=1) + if template is not None: + self._rollout_kwargs['template'] = template + # Which weights, and only for a local sampler: an API endpoint serves + # whatever it serves and has no notion of an adapter. + if hasattr(backend, 'sample'): + if self.lora_path: + self._rollout_kwargs['adapter_path'] = self.lora_path + else: + self._rollout_kwargs['use_base_model'] = True + self._rollout: Any | None = None # ------------------------------------------------------------------ # public entry point (pre/post processing, NOT decorated) @@ -106,14 +127,15 @@ def __call__(self, text: str, system: str = None, query: str = None, # ------------------------------------------------------------------ @llm_backup(key_params=["query"]) def _sample(self, trajectory, sampling_params, query: str = None) -> str: - """Student model: trajectory + sampling_params -> raw text.""" - sample_kwargs: dict[str, Any] = {'sampling_params': sampling_params} - if self.lora_path is None: - sample_kwargs['use_base_model'] = True - else: - sample_kwargs['adapter_path'] = self.lora_path - responses = self.sampler.sample([trajectory], **sample_kwargs) - return self._decoded(list(responses)[0]) if responses else '' + """Student model: trajectory + sampling_params -> raw text. + + The signature is what ``llm_backup`` reads by name to hand the teacher the + same input, so it stays even though the body no longer touches a sampler. + """ + if self._rollout is None: + self._rollout = MultiTurnRollout(self._backend, **self._rollout_kwargs) + replies = self._rollout([trajectory], sampling_params=sampling_params) + return assistant_text(replies[0]) if replies else '' # ------------------------------------------------------------------ # internals @@ -121,7 +143,7 @@ def _sample(self, trajectory, sampling_params, query: str = None) -> str: def _get_special_tokens(self) -> tuple[str, ...]: if self._special_tokens_cache is not None: return self._special_tokens_cache - tpl = self.template or getattr(self.sampler, 'template', None) + tpl = self.template or getattr(self._backend, 'template', None) tokenizer = getattr(tpl, 'tokenizer', None) if tpl is not None else None tokens: list[str] = [] if tokenizer is not None: @@ -176,13 +198,6 @@ def _postprocess(raw: str, original: str, special_tokens: tuple[str, ...]) -> st return None return text - @staticmethod - def _decoded(response: Any) -> str: - seqs = getattr(response, 'sequences', None) or [] - if not seqs: - return '' - return getattr(seqs[0], 'decoded', None) or '' - @staticmethod def _strip_code_fences(text: str) -> str: stripped = text.strip() diff --git a/src/twinkle_agentic/utils/code_utils.py b/src/twinkle_agentic/utils/code_utils.py index 37368045e..ab8f38403 100644 --- a/src/twinkle_agentic/utils/code_utils.py +++ b/src/twinkle_agentic/utils/code_utils.py @@ -40,16 +40,17 @@ @lru_cache(maxsize=None) -def _fence_re(language_tags: Tuple[str, ...]) -> Pattern: - """A fenced block whose language tag is one of ``language_tags``, or absent. +def _fence_re(language_tags: Optional[Tuple[str, ...]]) -> Pattern: + """Match a fenced block, optionally restricting its language label. - ``python``, ``py``, ``Python`` and ``python3`` are one intent spelled four - ways, so a tag matches case-insensitively and with any version suffix. A tag - that is not on the list -- ``bash``, ``json`` -- is a different intent, and is - not read as code at all. + ``None`` accepts any label. Otherwise, listed tags match case-insensitively, + with any version suffix; an unlabelled fence is accepted as well. """ - alts = '|'.join(re.escape(tag) for tag in language_tags) - label = r'(?:(?:%s)[\d.]*)?' % alts if alts else '' + if language_tags is None: + label = r'[^\r\n]*' + else: + alts = '|'.join(re.escape(tag) for tag in language_tags) + label = r'(?:(?:%s)[\d.]*)?' % alts if alts else '' return re.compile(r'```[ \t]*%s[ \t]*\r?\n(.*?)```' % label, re.S | re.I) @@ -69,10 +70,14 @@ def strip_reasoning(text: str) -> str: return body[cut:] -def parse_fenced_code(text: str, language_tags: Tuple[str, ...] = PYTHON_TAGS) -> Optional[str]: - """The last block ``text`` fenced as that language, or None if there is none. +def parse_fenced_code( + text: str, + language_tags: Optional[Tuple[str, ...]] = PYTHON_TAGS, +) -> Optional[str]: + """Return the last matching fenced block, or None if there is none. - The last one, not the first: a model often drafts a version before the final + Pass ``language_tags=None`` to accept any language label. The last block, not + the first, is returned because a model often drafts a version before the final one, and the block it ends on is its answer. What is inside is taken as given -- a fence is the model saying which part is diff --git a/tests/preprocessor/test_logprob_utils.py b/tests/preprocessor/test_logprob_utils.py deleted file mode 100644 index 29bb18e36..000000000 --- a/tests/preprocessor/test_logprob_utils.py +++ /dev/null @@ -1,354 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Tests for preprocessor.logprob_utils — pure logprob math helpers. - -These helpers compute conditional-vs-unconditional logprob deltas for -IFD-family scoring (CherryLLM, T-SHIRT, ChR). All functions are stateless -and accept simple list inputs. - -Conventions used in this test file: - * "lp" lists are aligned to the FULL sequence (prompt + answer). - * ``n_prompt`` is the number of prompt tokens; assistant tokens start at - index ``n_prompt`` in the cond list. - * Each lp entry is a dict {token_id: logprob_float}. -""" -import math -import pytest - -from twinkle_agentic.preprocessor.logprob_utils import (_chr_min_distinct, _chr_min_weighted, _extract_logprob, - _ifd_family_metrics, _lp_to_jsonable, _mean_logprob_delta, - _pad_batch, _to_int_list) - -# ── _extract_logprob ──────────────────────────────────────────────────────── - - -class TestExtractLogprob: - - def test_none(self): - assert _extract_logprob(None) is None - - def test_scalar_int(self): - assert _extract_logprob(5) == 5.0 - - def test_scalar_float(self): - assert _extract_logprob(-1.2) == -1.2 - - def test_dict_with_int_token_id(self): - lp = {7: -0.5, 8: -2.0} - assert _extract_logprob(lp, token_id=7) == -0.5 - assert _extract_logprob(lp, token_id=8) == -2.0 - - def test_dict_with_str_token_id_fallback(self): - # vLLM may emit string keys; lookup must fall back to str(token_id). - lp = {'7': -0.5} - assert _extract_logprob(lp, token_id=7) == -0.5 - - def test_dict_no_token_id_picks_first(self): - # No token_id → iter-first behaviour. - lp = {7: -0.5} - assert _extract_logprob(lp) == -0.5 - - def test_dict_token_id_missing_uses_first(self): - # token_id not in dict → fall back to first entry. - lp = {99: -3.0} - assert _extract_logprob(lp, token_id=7) == -3.0 - - def test_dict_with_logprob_attr_object(self): - - class Entry: - - def __init__(self, v): - self.logprob = v - - lp = {7: Entry(-0.7)} - assert _extract_logprob(lp, token_id=7) == -0.7 - - def test_dict_with_nested_dict(self): - lp = {7: {'logprob': -0.9, 'rank': 1}} - assert _extract_logprob(lp, token_id=7) == -0.9 - - def test_dict_with_nested_dict_none_logprob(self): - lp = {7: {'logprob': None}} - assert _extract_logprob(lp, token_id=7) is None - - def test_unrecognized_type(self): - # str entries → returns None - lp = {7: 'oops'} - assert _extract_logprob(lp, token_id=7) is None - - def test_non_dict_non_scalar(self): - # A list is neither scalar nor dict → None. - assert _extract_logprob([1, 2, 3]) is None - - -# ── _to_int_list ──────────────────────────────────────────────────────────── - - -class TestToIntList: - - def test_plain_list(self): - assert _to_int_list([1, 2, 3]) == [1, 2, 3] - - def test_tuple(self): - assert _to_int_list((1, 2, 3)) == [1, 2, 3] - - def test_with_tolist(self): - - class Tensor: - - def tolist(self): - return [4, 5, 6] - - assert _to_int_list(Tensor()) == [4, 5, 6] - - def test_empty(self): - assert _to_int_list([]) == [] - - -# ── _chr_min_distinct ─────────────────────────────────────────────────────── - - -class TestChrMinDistinct: - - def test_empty_inputs_returns_none(self): - assert _chr_min_distinct([], [{1: -1.0}], [], [1], 0) is None - assert _chr_min_distinct([{1: -1.0}], [], [1], [], 0) is None - assert _chr_min_distinct([{1: -1.0}], [{1: -1.0}], [1], [], 0) is None - - def test_simple_all_positive(self): - # cond_lp[i] - asst_lp[i] > 0 for all i → ratio = 1.0 - n_prompt = 1 - # cond covers prompt(1) + asst(2) = 3 positions - cond_lp = [ - { - 0: -10.0 - }, # prompt position - { - 1: -0.1 - }, # asst pos 0 — high cond logprob - { - 2: -0.2 - } - ] # asst pos 1 - asst_lp = [{1: -1.0}, {2: -1.5}] - cond_ids = [0, 1, 2] - asst_ids = [1, 2] - ratio = _chr_min_distinct(cond_lp, asst_lp, cond_ids, asst_ids, n_prompt) - assert ratio == 1.0 - - def test_all_negative(self): - # delta < 0 → ratio = 0 - n_prompt = 1 - cond_lp = [{0: 0.0}, {1: -3.0}, {2: -3.0}] - asst_lp = [{1: -0.5}, {2: -0.5}] - ratio = _chr_min_distinct(cond_lp, asst_lp, [0, 1, 2], [1, 2], n_prompt) - assert ratio == 0.0 - - def test_distinct_token_min_aggregation(self): - # Two occurrences of same token: one has +delta, one has -delta. - # min(deltas) is negative → token contributes 0 to ratio. - n_prompt = 1 - cond_lp = [{0: 0.0}, {1: -0.1}, {1: -3.0}] - asst_lp = [{1: -1.0}, {1: -0.5}] # delta1=+0.9, delta2=-2.5 - ratio = _chr_min_distinct(cond_lp, asst_lp, [0, 1, 1], [1, 1], n_prompt) - assert ratio == 0.0 # min < 0 - - def test_exclude_ids(self): - # Excluded token is dropped before counting. - n_prompt = 1 - cond_lp = [{0: 0.0}, {1: -0.1}, {2: -0.1}] - asst_lp = [{1: -1.0}, {2: -1.0}] - # Without exclude: 2 distinct tokens, both positive → 1.0 - ratio = _chr_min_distinct(cond_lp, asst_lp, [0, 1, 2], [1, 2], n_prompt, exclude_ids={1}) - assert ratio == 1.0 # only token 2 counted, still positive - - def test_truncation_when_cond_short(self): - # cond_lp shorter than n_prompt + n_asst → loop breaks early. - n_prompt = 2 - cond_lp = [{0: 0.0}, {0: 0.0}, {1: -0.1}] # only 1 asst position - asst_lp = [{1: -1.0}, {2: -1.0}] # 2 asst positions requested - ratio = _chr_min_distinct(cond_lp, asst_lp, [0, 0, 1], [1, 2], n_prompt) - assert ratio == 1.0 # only the first delta processed - - -# ── _chr_min_weighted ─────────────────────────────────────────────────────── - - -class TestChrMinWeighted: - - def test_empty_returns_none(self): - assert _chr_min_weighted([], [{1: -1.0}], [], [1], 0) is None - - def test_all_positive_returns_one(self): - n_prompt = 1 - cond_lp = [{0: 0.0}, {1: -0.1}, {2: -0.2}] - asst_lp = [{1: -1.0}, {2: -1.5}] - ratio = _chr_min_weighted(cond_lp, asst_lp, [0, 1, 2], [1, 2], n_prompt) - assert ratio == 1.0 # all positive → pos_w == total_w - - def test_zero_total_weight_returns_none(self): - # All deltas == 0 → total_w == 0 → None - n_prompt = 1 - cond_lp = [{0: 0.0}, {1: -1.0}] - asst_lp = [{1: -1.0}] - assert _chr_min_weighted(cond_lp, asst_lp, [0, 1], [1], n_prompt) is None - - def test_weighted_mixture(self): - # Token A: min_delta = +2.0 (weight 2) - # Token B: min_delta = -1.0 (weight 1) - # pos / total = 2 / 3 - n_prompt = 1 - cond_lp = [{0: 0.0}, {1: 1.0}, {2: -2.0}] # cond: A=1.0, B=-2.0 - asst_lp = [{1: -1.0}, {2: -1.0}] # asst: A=-1.0, B=-1.0 - # delta A = 1.0 - (-1.0) = 2.0 - # delta B = -2.0 - (-1.0) = -1.0 - ratio = _chr_min_weighted(cond_lp, asst_lp, [0, 1, 2], [1, 2], n_prompt) - assert abs(ratio - 2 / 3) < 1e-9 - - -# ── _ifd_family_metrics ───────────────────────────────────────────────────── - - -class TestIfdFamilyMetrics: - - def test_empty_returns_empty_dict(self): - assert _ifd_family_metrics([], [{1: -1.0}], [], [1], 0) == {} - - def test_simple_uniform(self): - # All deltas = 0.5 → mean=0.5, ifd=exp(-0.5) - n_prompt = 1 - cond_lp = [{0: 0.0}, {1: -0.5}, {2: -0.5}] - asst_lp = [{1: -1.0}, {2: -1.0}] - out = _ifd_family_metrics(cond_lp, asst_lp, [0, 1, 2], [1, 2], n_prompt) - assert out['n_tokens'] == 2 - assert abs(out['mean_delta'] - 0.5) < 1e-9 - assert abs(out['ifd'] - math.exp(-0.5)) < 1e-9 - # s_ifd_50 keeps top-1 by |delta| = 0.5; s_ifd_75 keeps top-2 (rounded up). - assert abs(out['s_ifd_50'] - math.exp(-0.5)) < 1e-9 - assert abs(out['s_ifd_75'] - math.exp(-0.5)) < 1e-9 - - def test_mixed_deltas(self): - # deltas = [+2.0, -1.0]; mean = 0.5 - n_prompt = 1 - cond_lp = [{0: 0.0}, {1: 1.0}, {2: -2.0}] - asst_lp = [{1: -1.0}, {2: -1.0}] - out = _ifd_family_metrics(cond_lp, asst_lp, [0, 1, 2], [1, 2], n_prompt) - assert out['n_tokens'] == 2 - assert abs(out['mean_delta'] - 0.5) < 1e-9 - # s_ifd_50 keeps top-1 by |delta| = 2.0 → exp(-2.0) - assert abs(out['s_ifd_50'] - math.exp(-2.0)) < 1e-9 - - -# ── _mean_logprob_delta ───────────────────────────────────────────────────── - - -class TestMeanLogprobDelta: - - def test_empty(self): - assert _mean_logprob_delta([], [{1: -1.0}], [], [1], 0) is None - - def test_uniform_delta(self): - n_prompt = 1 - cond_lp = [{0: 0.0}, {1: -0.5}, {2: -0.5}] - asst_lp = [{1: -1.0}, {2: -1.0}] - out = _mean_logprob_delta(cond_lp, asst_lp, [0, 1, 2], [1, 2], n_prompt) - assert abs(out - 0.5) < 1e-9 - - def test_mixed_average(self): - # deltas = [+2.0, -1.0] → mean 0.5 - n_prompt = 1 - cond_lp = [{0: 0.0}, {1: 1.0}, {2: -2.0}] - asst_lp = [{1: -1.0}, {2: -1.0}] - out = _mean_logprob_delta(cond_lp, asst_lp, [0, 1, 2], [1, 2], n_prompt) - assert abs(out - 0.5) < 1e-9 - - def test_skips_none_logprobs(self): - # When asst lp returns None, that position is skipped silently. - n_prompt = 1 - cond_lp = [{0: 0.0}, {1: -0.5}, {2: -0.5}] - asst_lp = [None, {2: -1.0}] - out = _mean_logprob_delta(cond_lp, asst_lp, [0, 1, 2], [1, 2], n_prompt) - assert abs(out - 0.5) < 1e-9 # only position 1 used - - -# ── _lp_to_jsonable ───────────────────────────────────────────────────────── - - -class TestLpToJsonable: - - def test_none_input(self): - assert _lp_to_jsonable(None) == [] - - def test_empty(self): - assert _lp_to_jsonable([]) == [] - - def test_none_passthrough(self): - assert _lp_to_jsonable([None, None]) == [None, None] - - def test_scalar_to_float(self): - assert _lp_to_jsonable([1, -2.0]) == [1.0, -2.0] - - def test_dict_with_logprob_object(self): - - class Entry: - - def __init__(self, lp, rank, decoded): - self.logprob = lp - self.rank = rank - self.decoded_token = decoded - - out = _lp_to_jsonable([{7: Entry(-0.5, 1, 'hello')}]) - assert out == [{'7': {'logprob': -0.5, 'rank': 1, 'decoded': 'hello'}}] - - def test_dict_with_nested_dict(self): - out = _lp_to_jsonable([{7: {'logprob': -0.5}}]) - assert out == [{'7': {'logprob': -0.5}}] - - def test_dict_with_repr_fallback(self): - # Non-dict, non-Entry value falls back to repr string. - out = _lp_to_jsonable([{7: 'plain'}]) - assert out == [{'7': repr('plain')}] - - def test_non_dict_non_scalar_repr(self): - # An object that isn't dict/scalar gets repr-ed. - out = _lp_to_jsonable([(1, 2)]) - assert out == [repr((1, 2))] - - -# ── _pad_batch ────────────────────────────────────────────────────────────── - - -class TestPadBatch: - - def test_empty_batch(self): - padded, n = _pad_batch([], floor=4) - assert padded == [] - assert n == 0 - - def test_already_at_floor(self): - batch = [[1], [2], [3], [4]] - padded, n = _pad_batch(batch, floor=4) - assert padded == batch - assert n == 4 - - def test_above_floor(self): - batch = [[1], [2], [3], [4], [5]] - padded, n = _pad_batch(batch, floor=3) - assert padded == batch # unchanged - assert n == 5 - - def test_below_floor_pads_with_last(self): - batch = [[1], [2]] - padded, n = _pad_batch(batch, floor=4) - assert padded == [[1], [2], [2], [2]] - assert n == 2 # original size - - def test_returns_new_list(self): - batch = [[1], [2]] - padded, _ = _pad_batch(batch, floor=4) - # Mutating padded should not affect original. - padded.append([99]) - assert batch == [[1], [2]] - - -if __name__ == '__main__': - pytest.main([__file__, '-v']) diff --git a/tests/twinkle_agentic/test_multi_turn_rollout.py b/tests/twinkle_agentic/test_multi_turn_rollout.py index 56c02b2a8..15541ba10 100644 --- a/tests/twinkle_agentic/test_multi_turn_rollout.py +++ b/tests/twinkle_agentic/test_multi_turn_rollout.py @@ -23,6 +23,7 @@ import json import pytest import re +import threading from typing import Any, Dict, List, Optional from twinkle.data_format.sampling import SampledSequence, SampleResponse, SamplingParams @@ -215,15 +216,40 @@ def concat_input_feature(self, pif: dict[str, Any], new_tokens: list[int]) -> di class FakeSampler: - """Queue-driven sampler that mirrors VLLMSampler output shape.""" + """Queue-driven sampler that mirrors VLLMSampler output shape. + + ``queue`` feeds one shared FIFO, which is all a single-trajectory test needs. + A batch needs ``queue_for(key, ...)``: episodes run in parallel threads, so + the order in which their turns reach ``sample`` is not defined, and a shared + FIFO would hand one trajectory's scripted reply to another. The key is the + text of the trajectory's first user message. + """ def __init__(self, template: FakeTemplate) -> None: self.template = template self._queue: list[dict[str, Any]] = [] + self._keyed: dict[str, list[dict[str, Any]]] = {} self.sample_calls = 0 # One entry per sample() call, so a test can assert which budget each # stage was sampled under. self.params_seen: list[Any] = [] + self._lock = threading.Lock() + + @staticmethod + def _entry( + template: FakeTemplate, + response_text: str, + stop_reason: str, + logprobs: list[Any] | None, + append_im_end: bool, + ) -> dict[str, Any]: + raw = response_text + ('<|im_end|>' if append_im_end else '') + return { + 'tokens': template.tokenizer.encode(raw, add_special_tokens=False), + 'decoded': response_text, + 'stop_reason': stop_reason, + 'logprobs': logprobs, + } def queue( self, @@ -236,14 +262,26 @@ def queue( ``<|im_end|>`` is appended to the encoded tokens when ``append_im_end``. ``seq.decoded`` is the raw response WITHOUT the trailing <|im_end|> (matches vLLM's common behaviour).""" - raw = response_text + ('<|im_end|>' if append_im_end else '') - tokens = self.template.tokenizer.encode(raw, add_special_tokens=False) - self._queue.append({ - 'tokens': tokens, - 'decoded': response_text, - 'stop_reason': stop_reason, - 'logprobs': logprobs, - }) + self._queue.append(self._entry(self.template, response_text, stop_reason, logprobs, append_im_end)) + + def queue_for( + self, + key: str, + response_text: str, + stop_reason: str = 'stop', + logprobs: list[Any] | None = None, + append_im_end: bool = True, + ) -> None: + """Script one turn for the trajectory whose first user message is ``key``.""" + self._keyed.setdefault(key, []).append( + self._entry(self.template, response_text, stop_reason, logprobs, append_im_end)) + + @staticmethod + def _key_of(pif: dict[str, Any]) -> str | None: + for m in pif.get('messages') or []: + if m.get('role') == 'user': + return m.get('content') + return None def sample(self, pifs, sampling_params=None): # Batched contract: accept a list of pifs, return one @@ -252,12 +290,14 @@ def sample(self, pifs, sampling_params=None): if isinstance(pifs, dict): pifs = [pifs] assert isinstance(pifs, list), (f'FakeSampler.sample expects a list, got {type(pifs).__name__}') - self.params_seen.append(sampling_params) responses: list[SampleResponse] = [] for pif in pifs: - assert self._queue, 'FakeSampler queue exhausted — scripted turns' - r = self._queue.pop(0) - self.sample_calls += 1 + with self._lock: + self.params_seen.append(sampling_params) + queue = self._keyed.get(self._key_of(pif)) or self._queue + assert queue, 'FakeSampler queue exhausted — scripted turns' + r = queue.pop(0) + self.sample_calls += 1 new_pif = self.template.concat_input_feature(pif, r['tokens']) seq = SampledSequence( stop_reason=r['stop_reason'], @@ -269,6 +309,10 @@ def sample(self, pifs, sampling_params=None): responses.append(SampleResponse(sequences=[seq])) return responses + # MultiTurnRollout samples one trajectory per call and refuses a sampler + # that would slice such a batch across workers. + sample._enable_continous_work = True + class EchoTool(Tool): """Echoes its arguments as a JSON string.""" @@ -607,12 +651,13 @@ def test_stuck_stop_is_per_trajectory_in_a_batch(make_rollout, sampler, template bad = ToolManager({}) bad.register(FailTool('search')) - sampler.queue(_tool_call_text('search', {'q': 1}), stop_reason='stop') - sampler.queue(_tool_call_text('search', {'q': 1}), stop_reason='stop') - sampler.queue(_tool_call_text('search', {'q': 2}), stop_reason='stop') - sampler.queue(_tool_call_text('search', {'q': 3}), stop_reason='stop') - sampler.queue('Done.', stop_reason='stop') - sampler.queue('Done.', stop_reason='stop') + sampler.queue_for('a', _tool_call_text('search', {'q': 1}), stop_reason='stop') + sampler.queue_for('a', _tool_call_text('search', {'q': 2}), stop_reason='stop') + sampler.queue_for('a', _tool_call_text('search', {'q': 3}), stop_reason='stop') + sampler.queue_for('a', 'Done.', stop_reason='stop') + # 'b' calls the failing tool twice, which trips stop_after_stuck_turns=2. + sampler.queue_for('b', _tool_call_text('search', {'q': 1}), stop_reason='stop') + sampler.queue_for('b', _tool_call_text('search', {'q': 1}), stop_reason='stop') rollout = MultiTurnRollout( sampler=sampler, template=template, tool_manager=[good, bad], @@ -766,6 +811,28 @@ def test_rejects_num_samples_gt_1(sampler, template, tool_manager): sampling_params=SamplingParams(num_samples=2)) +def test_rejects_sampler_without_continous_work(template, tool_manager): + """A batch of one is what a slice_dp sampler cannot serve.""" + + class SlicingSampler: + + def sample(self, pifs, sampling_params=None): + return [] + + with pytest.raises(ValueError, match='enable_continous_work'): + MultiTurnRollout(sampler=SlicingSampler(), template=template, tool_manager=tool_manager) + + +def test_rejects_one_harness_shared_by_a_batch(sampler, template, tool_manager): + """Episodes run in parallel threads, so a stateful harness cannot be shared.""" + from twinkle_agentic.harness.base import AgentHarness + + rollout = MultiTurnRollout( + sampler=sampler, template=template, tool_manager=tool_manager, max_turns=2, harness=AgentHarness()) + with pytest.raises(ValueError, match='harness holds per-episode state'): + rollout([_user_traj('A'), _user_traj('B')]) + + # ============================================================================= # Tests: defensive guards # ============================================================================= @@ -779,6 +846,8 @@ def sample(self, pifs, sampling_params=None): seq = SampledSequence(stop_reason='stop', tokens=[], logprobs=None, decoded='', new_input_feature=None) return [SampleResponse(sequences=[seq]) for _ in pifs] + sample._enable_continous_work = True + rollout = MultiTurnRollout(sampler=BrokenSampler(), template=template, tool_manager=tool_manager) with pytest.raises(RuntimeError, match='new_input_feature'): rollout([_user_traj()]) @@ -791,6 +860,8 @@ class EmptySampler: def sample(self, pifs, sampling_params=None): return [] + sample._enable_continous_work = True + rollout = MultiTurnRollout(sampler=EmptySampler(), template=template, tool_manager=tool_manager) # Batched contract: 0 responses for a batch of 1 → mismatch error. with pytest.raises(RuntimeError, match='0 responses'): @@ -806,6 +877,8 @@ def sample(self, pifs, sampling_params=None): pifs = [pifs] return [SampleResponse(sequences=[]) for _ in pifs] + sample._enable_continous_work = True + rollout = MultiTurnRollout(sampler=NoSeqSampler(), template=template, tool_manager=tool_manager) with pytest.raises(RuntimeError, match='no sequences'): rollout([_user_traj()]) @@ -820,18 +893,17 @@ def test_empty_batch_returns_empty_list(make_rollout): def test_batch_single_turn_two_trajectories(make_rollout, sampler): - """Two trajectories finish on turn 1 → one batched sample call.""" - sampler.queue('answer-A', stop_reason='stop') - sampler.queue('answer-B', stop_reason='stop') + """Two trajectories, one turn each, in their own threads.""" + sampler.queue_for('Q-A', 'answer-A', stop_reason='stop') + sampler.queue_for('Q-B', 'answer-B', stop_reason='stop') rollout = make_rollout(max_turns=3) outs = rollout([_user_traj('Q-A'), _user_traj('Q-B')]) assert len(outs) == 2 - # Exactly ONE batched sample call, not two. - assert sampler.sample_calls == 2 # one per item, still one turn - # But FakeSampler counts per-input; the critical batching invariant is - # that MultiTurnRollout only calls sampler.sample ONCE per turn. We - # enforce this via the queue ordering + single turn. + assert sampler.sample_calls == 2 # one generation per trajectory + # Results come back in input order even though the threads may not. + assert outs[0]['messages'][-1]['content'] == 'answer-A' + assert outs[1]['messages'][-1]['content'] == 'answer-B' for out in outs: assert out['turns'] == 1 assert out['stop_reason'] == 'stop' @@ -841,14 +913,12 @@ def test_batch_single_turn_two_trajectories(make_rollout, sampler): def test_batch_different_termination_turns(make_rollout, sampler): """Trajectory A finishes on turn 1; trajectory B needs a tool turn. - Turn 1 batch: [A: 'done-A' stop, B: tool_call stop] → A parked. - Turn 2 batch: [B: 'done-B' stop] → only B live. + Each episode owns its turn budget, so B taking a second turn neither waits + for A nor buys A anything. """ - sampler.queue('done-A', stop_reason='stop') # A turn 1 - sampler.queue( - _tool_call_text('search', {'q': 'b'}), # B turn 1 - stop_reason='stop') - sampler.queue('done-B', stop_reason='stop') # B turn 2 + sampler.queue_for('Q-A', 'done-A', stop_reason='stop') + sampler.queue_for('Q-B', _tool_call_text('search', {'q': 'b'}), stop_reason='stop') + sampler.queue_for('Q-B', 'done-B', stop_reason='stop') rollout = make_rollout(max_turns=4) outs = rollout([_user_traj('Q-A'), _user_traj('Q-B')]) @@ -889,10 +959,10 @@ def tool_info(self): tm_b = ToolManager({}) tm_b.register(TagTool('B')) - sampler.queue(_tool_call_text('search', {'q': 'x'}), stop_reason='stop') - sampler.queue(_tool_call_text('search', {'q': 'y'}), stop_reason='stop') - sampler.queue('done-A', stop_reason='stop') - sampler.queue('done-B', stop_reason='stop') + sampler.queue_for('A', _tool_call_text('search', {'q': 'x'}), stop_reason='stop') + sampler.queue_for('A', 'done-A', stop_reason='stop') + sampler.queue_for('B', _tool_call_text('search', {'q': 'y'}), stop_reason='stop') + sampler.queue_for('B', 'done-B', stop_reason='stop') rollout = MultiTurnRollout( sampler=sampler, @@ -1014,8 +1084,8 @@ def _is_success(traj): max_turns=2, trace_dir=str(trace_dir), success_callback=_is_success) - sampler.queue('good answer', stop_reason='stop') - sampler.queue('bad answer', stop_reason='stop') + sampler.queue_for('A', 'good answer', stop_reason='stop') + sampler.queue_for('B', 'bad answer', stop_reason='stop') rollout([_user_traj('A'), _user_traj('B')]) @@ -1031,9 +1101,9 @@ def test_trace_dir_batch_writes_one_file_per_trajectory(tmp_path, sampler, templ rollout = MultiTurnRollout( sampler=sampler, template=template, tool_manager=tool_manager, max_turns=4, trace_dir=str(trace_dir)) # Traj 0: stops turn 1. Traj 1: tool-calls turn 1, stops turn 2. - sampler.queue('done0', stop_reason='stop') - sampler.queue(_tool_call_text('search', {'q': 'y'})) - sampler.queue('done1', stop_reason='stop') + sampler.queue_for('A', 'done0', stop_reason='stop') + sampler.queue_for('B', _tool_call_text('search', {'q': 'y'})) + sampler.queue_for('B', 'done1', stop_reason='stop') rollout([_user_traj('A'), _user_traj('B')]) From 55d2309ebdbf70e5de0f91643ce2c27370fcb768 Mon Sep 17 00:00:00 2001 From: root <yuze.zyz@alibaba-inc.com> Date: Thu, 10 Sep 2026 00:44:56 +0800 Subject: [PATCH 59/60] wip --- .temp/retired_rsi/loop.sh | 252 ------ .temp/retired_rsi/rsi_agent_shortprompt.yaml | 197 ---- .temp/retired_rsi/split_tasks.py | 70 -- README.md | 2 +- README_ZH.md | 2 +- cookbook/rsi/code/challenge.py | 4 +- cookbook/rsi/code/collect.py | 8 +- docs/source_en/Components/Agentic/Envs.md | 4 +- .../Agentic/Multi-Turn-Tool-Usage.md | 19 +- docs/source_en/Components/Agentic/Rollout.md | 58 +- .../\347\273\204\344\273\266/Agentic/Envs.md" | 4 +- .../Agentic/Multi-Turn-Tool-Usage.md" | 19 +- .../Agentic/Rollout.md" | 58 +- src/twinkle/loss/grpo.py | 35 +- src/twinkle/model/megatron/megatron.py | 8 + .../strategy/sequence_parallel/__init__.py | 26 +- .../model/transformers/transformers.py | 8 + src/twinkle/processor/base.py | 39 +- .../sampler/vllm_sampler/vllm_sampler.py | 10 +- .../server/sampler/backends/mock_sampler.py | 11 +- src/twinkle/template/base.py | 154 +++- src/twinkle_agentic/async_rl/data_plane.py | 18 +- src/twinkle_agentic/challenger/base.py | 12 +- .../challenger/new/__init__.py | 12 + src/twinkle_agentic/challenger/new/agentic.py | 511 +++++++++++ src/twinkle_agentic/challenger/new/base.py | 132 +++ src/twinkle_agentic/challenger/new/keyword.py | 288 +++++- .../challenger/new/recorder.py | 91 ++ src/twinkle_agentic/harness/ms_agent.py | 16 +- src/twinkle_agentic/preprocessor/AUDIT.md | 179 ---- src/twinkle_agentic/preprocessor/__init__.py | 2 - .../preprocessor/experimental/__init__.py | 21 - .../preprocessor/experimental/llm_backend.py | 344 ------- .../preprocessor/experimental/score_filter.py | 835 ----------------- .../preprocessor/intent_classifier.py | 12 +- src/twinkle_agentic/preprocessor/intents.py | 25 - .../preprocessor/label_schema.py | 117 --- .../preprocessor/logprob_utils.py | 231 ----- .../preprocessor/offline/__init__.py | 24 - .../preprocessor/offline/decontaminate.py | 109 --- .../preprocessor/offline/near_dedup.py | 163 ---- .../preprocessor/provenance.py | 73 -- .../preprocessor/structural_noise.py | 60 -- src/twinkle_agentic/protocol/openai.py | 46 +- src/twinkle_agentic/rollout/__init__.py | 17 +- src/twinkle_agentic/rollout/api_multi_turn.py | 314 ------- src/twinkle_agentic/rollout/api_sampler.py | 133 +++ src/twinkle_agentic/rollout/base.py | 119 ++- src/twinkle_agentic/rollout/bridge.py | 165 ++-- src/twinkle_agentic/rollout/factory.py | 74 -- src/twinkle_agentic/rollout/multi_turn.py | 853 +++++++++--------- src/twinkle_agentic/summarizer/base.py | 63 +- src/twinkle_agentic/utils/code_utils.py | 27 +- tests/preprocessor/test_logprob_utils.py | 354 -------- .../test_multi_turn_rollout.py | 156 +++- 55 files changed, 2357 insertions(+), 4227 deletions(-) delete mode 100644 .temp/retired_rsi/loop.sh delete mode 100644 .temp/retired_rsi/rsi_agent_shortprompt.yaml delete mode 100644 .temp/retired_rsi/split_tasks.py create mode 100644 src/twinkle_agentic/challenger/new/__init__.py create mode 100644 src/twinkle_agentic/challenger/new/agentic.py create mode 100644 src/twinkle_agentic/challenger/new/base.py create mode 100644 src/twinkle_agentic/challenger/new/recorder.py delete mode 100644 src/twinkle_agentic/preprocessor/AUDIT.md delete mode 100644 src/twinkle_agentic/preprocessor/experimental/__init__.py delete mode 100644 src/twinkle_agentic/preprocessor/experimental/llm_backend.py delete mode 100644 src/twinkle_agentic/preprocessor/experimental/score_filter.py delete mode 100644 src/twinkle_agentic/preprocessor/intents.py delete mode 100644 src/twinkle_agentic/preprocessor/label_schema.py delete mode 100644 src/twinkle_agentic/preprocessor/logprob_utils.py delete mode 100644 src/twinkle_agentic/preprocessor/offline/__init__.py delete mode 100644 src/twinkle_agentic/preprocessor/offline/decontaminate.py delete mode 100644 src/twinkle_agentic/preprocessor/offline/near_dedup.py delete mode 100644 src/twinkle_agentic/preprocessor/provenance.py delete mode 100644 src/twinkle_agentic/preprocessor/structural_noise.py delete mode 100644 src/twinkle_agentic/rollout/api_multi_turn.py create mode 100644 src/twinkle_agentic/rollout/api_sampler.py delete mode 100644 src/twinkle_agentic/rollout/factory.py delete mode 100644 tests/preprocessor/test_logprob_utils.py diff --git a/.temp/retired_rsi/loop.sh b/.temp/retired_rsi/loop.sh deleted file mode 100644 index 667ee78ae..000000000 --- a/.temp/retired_rsi/loop.sh +++ /dev/null @@ -1,252 +0,0 @@ -#!/bin/bash -# Self-evolving loop: collect, train on what was collected, collect again from the -# weights that came out. -# -# collect: challenge.py --model-id <last ckpt> --keep-groups 8 -# runs until 8 groups have been kept, whatever that costs in topics -# train: train.py --run-dir <that collection> -> one HF checkpoint -# repeat -# -# Nothing is generated in the training stage and nothing is re-encoded: the tokens -# trained on are the ones the sampler produced, read straight off disk. -# -# The two stages are separate processes so each gets every GPU. Collection is the -# slow half, and splitting the GPUs between a trainer and a sampler in one process -# would halve it. The cost is restarting vLLM and the sandboxes each time, measured -# at 1-2 minutes against roughly 40 minutes of collecting. -# -# Nothing about the host is written down here: the repo is found from this script's -# own location, the GPU count from nvidia-smi, and the secrets have to be exported -# first -- the script stops with the name of whatever is missing rather than -# guessing a value that would fail deep inside a run. -# -# export E2B_API_KEY=... # sandbox host key -# export SANDBOX_API_URL=http://... # sandbox host address, with port -# export LLM_BACKUP_API_KEY=... # dashscope, for checks/statements/rubric -# bash cookbook/rsi/agentic/loop.sh # until killed -# ITERATIONS=1 bash cookbook/rsi/agentic/loop.sh # one collect + one train -set -e -# Both stages are piped into tee, and without this the pipeline's status is tee's, -# which is 0 even when python died. An earlier loop crashed inside collection, -# trained on the partial collection anyway, saved a checkpoint from it and marked -# the iteration finished -- all reported as success. -set -o pipefail - -HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -REPO="$(cd "$HERE/../../.." && pwd)" -if [ ! -f "$HERE/challenge.py" ] || [ ! -f "$REPO/setup.cfg" ]; then - echo "expected challenge.py beside this script and the repo root three levels up" >&2 - exit 1 -fi -cd "$REPO" - -missing="" -for v in E2B_API_KEY SANDBOX_API_URL LLM_BACKUP_API_KEY; do - [ -z "${!v}" ] && missing="$missing $v" -done -if [ -n "$missing" ]; then - echo "export these first:$missing" >&2 - exit 1 -fi - -# Every GPU on the box unless told otherwise. Counted rather than written down, -# since the point of moving hosts is usually a different number of them. -GPUS="${GPUS:-$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l)}" -[ "$GPUS" -lt 1 ] && { echo "nvidia-smi reports no GPUs" >&2; exit 1; } -DEVICES="${DEVICES:-$(seq -s, 0 $((GPUS - 1)))}" - -# Refuse to start on top of someone else's job. Both stages want the whole GPU: -# challenge.py boots one vLLM per GPU at 0.8 of its memory, so sharing means an -# out-of-memory crash partway in and the other job may go down with it. Any compute -# process at all counts; CONFIRM_GPUS=1 starts anyway. -BUSY="$(nvidia-smi --query-compute-apps=pid --format=csv,noheader | sort -u | wc -l)" -if [ "$BUSY" -gt 0 ] && [ "${CONFIRM_GPUS:-0}" != "1" ]; then - echo "$BUSY process(es) already on the GPUs:" >&2 - nvidia-smi --query-compute-apps=pid,used_memory --format=csv,noheader >&2 - echo "set CONFIRM_GPUS=1 to start anyway" >&2 - exit 1 -fi - -export AENV_API_URL="$SANDBOX_API_URL" -export AENV_API_KEY="$E2B_API_KEY" -# Sandbox image name: a fact about how the host was set up, not a preference. -export AENV_TEMPLATE="${AENV_TEMPLATE:-twinkle-rsi-msagent}" - -# ---- what to run --------------------------------------------------------- -ITERATIONS="${ITERATIONS:-0}" # 0 = until killed -TAG="${TAG:-fix1}" -BASE_MODEL="${BASE_MODEL:-ms://Qwen/Qwen3-4B}" - -# Concurrent sandboxes, i.e. how many trajectories are in flight at once. Bounded -# by the sandbox host, not by the GPUs here, so it does not follow GPUS. 32, not -# the 96 a capacity probe once managed: holding 96 for a whole run was not -# reliable. Each slot is one thread and one microVM; vLLM sees up to this many -# single-trajectory requests at a time and batches them itself. -SANDBOX_SLOTS="${SANDBOX_SLOTS:-32}" - -# Groups kept per iteration, and their shape. A group is one keyword draw answered -# GROUP_SIZE times; it is kept when at least one of those answers became a task the -# solver passes sometimes, meaning n_pass in [1, SOLVER_ROLLOUTS-1]. Outside that -# band every attempt carries the same reward, the group mean equals it, and the -# advantage is zero for all of them. -# -# proposing side: KEEP_GROUPS x GROUP_SIZE = 8 x 8 = 64 trajectories -# solving side: KEEP_GROUPS x SOLVER_ROLLOUTS = 8 x 8 = 64 trajectories -# 128 total, one optimizer step over all of it -# -# What this pays for and throws away: every proposal that produced a task costs -# SOLVER_ROLLOUTS sandbox attempts, and only the selected proposal's attempts are -# trained on. At 8 groups x 8 proposals that is up to 512 attempts run and 64 -# trained on. The unselected proposals are not wasted on the proposing side -- each -# earns its own reward from its own n_pass, including 0 for the ones that produced -# no task at all. -KEEP_GROUPS="${KEEP_GROUPS:-8}" -GROUP_SIZE="${GROUP_SIZE:-8}" -SOLVER_ROLLOUTS="${SOLVER_ROLLOUTS:-8}" - -# Cap on files one build may leave behind, appended to the system prompt. This -# changes the prompt and therefore what is trained; 4 is what every run since it -# was added has used. -MAX_BUILD_FILES="${MAX_BUILD_FILES:-4}" - -# Reasoning cap on every API call. The one knob that moved wall-clock: 58s -> 10s -# per turn at 2048 on a ~15k-character context. -API_THINKING_BUDGET="${API_THINKING_BUDGET:-4096}" -API_MODEL="${API_MODEL:-qwen3.8-max}" -API_BASE="${API_BASE:-https://dashscope.aliyuncs.com/compatible-mode/v1}" - -# Novelty. The bank is one file for the whole loop, not one per iteration, because -# the point of it is comparing iteration k+1's proposals against what k produced. -# TASK_BANK="" turns it off and gives back the pass-rate gaussian alone; -# NOVELTY_FLOOR=1 keeps the judging and the log but stops it changing any reward. -# -# 1 is the default because the score it would multiply in has not been shown to carry -# anything yet. Measured on the 27 proposals of iter1 (novelty_scores.jsonl): judged -# against their own siblings, 24 of 27 scored exactly 0.0, so the term was constant -# across the group and contributed nothing once GRPO subtracts the group mean -- while -# still halving every proposer reward at floor 0.5. The alternative measured, labelling -# each task's shape on its own and scoring by how rare that shape is in the group, does -# separate proposals (0 of 4 groups constant), but its label changed between sampled -# repeats on 10 of 27 statements, so the number it produces is not comparable across -# runs. Until one of those is fixed the score is written to novelty_scores.jsonl and -# read there. Set 0.5 to bring it back into the reward. -NOVELTY_FLOOR="${NOVELTY_FLOOR:-1}" - -LEARNING_RATE="${LEARNING_RATE:-1e-6}" -SIDES="${SIDES:-both}" - -ROOT="output/rsi_agentic/${TAG}" -# One checkpoint directory for the whole loop, overwritten every iteration, so the -# disk holds one 4B model rather than one per iteration. The previous round's -# weights are gone once the next save starts: if a save dies partway there is -# nothing to fall back to but BASE_MODEL. -# -# Overridable because this is the one path whose filesystem shows up in wall-clock: -# every iteration reads it 1 + GPUS times, once per vLLM worker at collect and once -# more at train, so 7.6 GB of weights is around 60 GB of reads per iteration. Two -# filesystems on this host, same 1.5 T free, measured with dd at 1.5 GB: the repo's -# own disk reads at 223 MB/s and the parallel one at 1074 MB/s -- 4.8x, which showed -# up as five minutes of vLLM startup before any topic was launched. Left at $ROOT/ckpt -# by default so a host with one disk needs to know nothing about this. -CKPT_DIR="${CKPT_DIR:-$ROOT/ckpt}" -# Written with ${VAR-default} rather than ${VAR:-default} so that TASK_BANK="" -# means off; with the colon an empty value would silently get the default back. -TASK_BANK="${TASK_BANK-$ROOT/task_bank.jsonl}" -mkdir -p "$ROOT" - -# Pick up where a previous invocation left off. The iteration number comes from a -# marker written after the checkpoint has been checked, not from train_summary.json, -# which is written at the end of a training run but would still be there after a -# crash in a later stage. -MODEL="$BASE_MODEL" -START=1 -while [ -f "$ROOT/iter${START}/iteration.done" ]; do - START=$((START + 1)) -done -if [ "$START" -gt 1 ]; then - if [ -f "$CKPT_DIR/model/config.json" ]; then - MODEL="$CKPT_DIR/model" - else - echo "$((START - 1)) iteration(s) finished under $ROOT but no checkpoint at" >&2 - echo "$CKPT_DIR/model -- each iteration overwrites the one before, so those" >&2 - echo "weights are gone. Start a new TAG, or delete the iteration.done" >&2 - echo "markers to redo them from $BASE_MODEL." >&2 - exit 1 - fi -fi - -cat <<EOF -=== repo $REPO -=== gpus $GPUS (devices $DEVICES) -=== sandbox $AENV_API_URL template $AENV_TEMPLATE slots $SANDBOX_SLOTS -=== api $API_MODEL at $API_BASE, thinking budget $API_THINKING_BUDGET -=== per iter $KEEP_GROUPS groups of $GROUP_SIZE, band [1, $((SOLVER_ROLLOUTS - 1))] of $SOLVER_ROLLOUTS -=== build cap $([ "$MAX_BUILD_FILES" -eq 0 ] && echo "none" || echo "$MAX_BUILD_FILES files, in the system prompt") -=== novelty $([ -z "$TASK_BANK" ] && echo "off" || echo "bank $TASK_BANK, floor $NOVELTY_FLOOR") -=== trains on $((KEEP_GROUPS * GROUP_SIZE)) propose + $((KEEP_GROUPS * SOLVER_ROLLOUTS)) solve trajectories, one step, lr $LEARNING_RATE -=== checkpoint $CKPT_DIR/model, overwritten each iteration -=== iterations $([ "$ITERATIONS" -eq 0 ] && echo "until killed" || echo "$ITERATIONS") -=== swanlab ${RSI_SWANLAB_MODE:-online} project ${RSI_SWANLAB_PROJECT:-twinkle-rsi-agentic}, experiment $TAG, one step per iteration -=== starting at iteration $START from $MODEL -EOF - -i="$START" -while [ "$ITERATIONS" -eq 0 ] || [ "$i" -lt $((START + ITERATIONS)) ]; do - OUT="$ROOT/iter${i}" - mkdir -p "$OUT" - echo "=== iteration $i: collect $KEEP_GROUPS groups from $MODEL -> $OUT" - - CUDA_VISIBLE_DEVICES="$DEVICES" python cookbook/rsi/agentic/challenge.py \ - --model-id "$MODEL" \ - --sampler-gpus "$GPUS" \ - --sandbox-slots "$SANDBOX_SLOTS" \ - --keep-groups "$KEEP_GROUPS" \ - --group-size "$GROUP_SIZE" \ - --solver-rollouts "$SOLVER_ROLLOUTS" \ - --max-build-files "$MAX_BUILD_FILES" \ - --api-model "$API_MODEL" \ - --api-base "$API_BASE" \ - --api-thinking-budget "$API_THINKING_BUDGET" \ - --task-bank "$TASK_BANK" \ - --novelty-floor "$NOVELTY_FLOOR" \ - --out-dir "$OUT" \ - --keyword-db "$ROOT/keywords.jsonl" \ - 2>&1 | tee "$OUT/challenge.log" - - echo "=== iteration $i: train on $OUT -> $CKPT_DIR" - # expandable_segments on the training stage only, and not on collect: one padded - # trajectory per micro batch means every micro batch is a new shape (119 distinct - # lengths in 128 trajectories, 7k-19k tokens), and the caching allocator cannot - # reuse a block across sizes, so it grew to 87.8 GiB reserved against 29.0 GiB - # live on a 97.4 GiB card. That is what starved NCCL of the few hundred MiB it - # needs to connect the metric gather's communicator, which hung iteration 2 for - # 54 minutes. Expandable segments let one virtual range serve every shape, so - # reserved tracks the real peak instead of the sum of shapes. Left off for - # collect because that stage is vLLM, which profiles its own KV cache against - # allocator behaviour and has nothing to do with this failure. - RSI_RUN_DIR="$OUT" \ - RSI_SAVE_DIR="$CKPT_DIR" \ - RSI_SAVE_NAME="model" \ - RSI_SIDES="$SIDES" \ - RSI_TAG="$TAG" \ - RSI_ITER="$i" \ - PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" \ - CUDA_VISIBLE_DEVICES="$DEVICES" python cookbook/rsi/agentic/train.py \ - --model_id "$MODEL" \ - --model_gpus "$GPUS" \ - --lr "$LEARNING_RATE" \ - 2>&1 | tee "$OUT/train.log" - - # HF-format weights plus tokenizer, which is what --model-id takes, so the next - # iteration needs no conversion step. - MODEL="$CKPT_DIR/model" - if [ ! -f "$MODEL/config.json" ]; then - echo "iteration $i saved no loadable checkpoint at $MODEL" >&2 - exit 1 - fi - # Written last, so resuming counts only iterations whose weights are on disk. - touch "$OUT/iteration.done" - echo "=== iteration $i done; next starts from $MODEL" - i=$((i + 1)) -done -echo "=== stopped after iteration $((i - 1)); model at $MODEL" diff --git a/.temp/retired_rsi/rsi_agent_shortprompt.yaml b/.temp/retired_rsi/rsi_agent_shortprompt.yaml deleted file mode 100644 index cace8cbe2..000000000 --- a/.temp/retired_rsi/rsi_agent_shortprompt.yaml +++ /dev/null @@ -1,197 +0,0 @@ -# ms-agent config for agentic RSI training. -# -# Read by both halves of the setup, which is the point: -# -# * the training host loads it to build a MsAgentHarness for message shaping -# only -- the entry script drops `llm:` and `tools:` from the merged config -# first, so no tool is ever constructed next to the trainer; -# * remote_tool_env.py uploads this same file into each sandbox, where -# sandbox_server/tool_server.py loads it and does construct the tools. -# -# So the tool line-up below describes what runs in the microVM. Editing it takes -# effect on the next episode; no image rebuild is involved. - -prompt: - # Replaces ms-agent's BASE_AGENT_PROMPT (prompting/builtin.py) for the SOLVER - # only -- the proposing episode gets prompts.py's own SYSTEM through the - # challenger, and never reads this field. That built-in prompt is written for a - # general assistant sitting in a user's workspace, and two of its lines work - # against being a solver: "First decide whether the task needs tools. If you can - # answer reliably from what you know ... just answer", and "Ask first when it - # isn't [safe]". Here there is no one to ask (interactive: false) and answering - # without touching the directory is always wrong. - # - # The paragraph about the empty directory is what 5 of armA2shellV5's 8 - # unsolved tasks needed. Their statements listed a file under "Input data:" - # and the solver read that as "already present" -- in 5a70b77f it created the - # file the rules told it to generate and left the two listed as input alone, - # so it was not confused about being in an empty directory, it was following - # the statement's own division of labour. Nothing in the statement or the - # prompt said that division does not survive into its workspace. - # SHORT ARM of the solver-prompt A/B. 242 words down to 190, by the same - # subtractive rule used on prompts.py's two followups: every instruction stays, the - # sentences arguing for it go. Only 21% shorter, and that is the ceiling -- the - # empty-directory paragraph is 90 of the 190 words and is kept whole, because 5 of - # armA2shellV5's 8 unsolved tasks needed exactly it. What left: 'nobody is - # watching', 'before anything can read them', 'Do not assume', 'however clearly you - # can describe what the answer would be', 'Check each thing the task asked for is - # actually there', and the How-to-work heading with its bullets folded into prose. - # - # This file differs from rsi_agent.yaml in this field ONLY -- tools, permission, - # timeouts and output_dir are identical (verified by comparing the parsed configs - # with prompt removed), so n_pass measured against the two is measuring the prompt. - system: | - You are a command-line agent working inside a fresh Linux container. You are - given one task and you carry it out by running commands and writing files. - Nobody can answer a question, so never ask one and never stop to confirm: - decide and act. - - Your working directory starts COMPLETELY EMPTY. Every file the task - mentions -- including files it describes as inputs, given data, existing - configuration, or material you are handed -- does not exist yet. You have to - create all of them yourself, with exactly the names and contents the task - specifies. A task that shows you the contents of a file is telling you what to - write into it, not telling you it is there. - - Start by listing the directory to see the real state. Create every file the - task names, then do the computation it asks for and write the results it asks - for. Before you finish, list the directory again and read back what you wrote; - if something is missing, fix it rather than reporting success. Answering in - prose without creating files is a failure. Never state a value you did not - compute from the data. - -personalization: - # Off: SOUL/AGENTS/PROFILE.md from the developer's own workspace would leak - # machine-specific context into every training prompt. - enabled: false - -# One turn == one sampler call. MultiTurnRollout's max_turns is the real limit; -# this only stops ms-agent from imposing a lower one. -max_chat_round: 9999 - -# Never wait on a human: training runs unattended. -interactive: false -permission_mode: auto - -# How long ms-agent waits around one tool call. Written down rather than left to -# its default (tool_manager.py TOOL_CALL_TIMEOUT, 120s, overridable by the -# TOOL_CALL_TIMEOUT environment variable) so the sandbox does not inherit a -# number from whatever shell started it. It has to stay below what -# remote_tool_env allows the whole turn (command_timeout, 180s), which in turn is -# below the transport's budget: the innermost layer should be the one that times -# out, because it is the only one that knows which call was slow. When they were -# equal, one command that never returns made every call in the turn read as an -# unreachable runtime. -tool_call_timeout: 120 - -# Path *inside the sandbox*. One microVM per episode already isolates -# trajectories from each other, so this is a fixed path rather than a per-slot -# directory; the entry script overrides it only to match --workspace. -output_dir: /workspace - -callbacks: [] - -tools: - # `file_system` is NOT listed here and is nevertheless on. ms-agent's own - # ms_agent/agent/agent.yaml declares it (write_file, read_file, edit_file, - # grep, glob) and LLMAgent merges this file *over* that one, so omitting a key - # inherits it rather than dropping it. Measured: the merged config's tools are - # ['file_system', 'code_executor', 'todo_list'], and /tools advertises all ten - # of those tools to the model. In armA2shellV6's 128 proposing calls, - # file_system took 63 (43 of them write_file) against code_executor's 58. - # - # So the paragraph that used to be here -- claiming the five were removed to - # stop write_file being the path of least resistance -- described a state that - # never existed, through the arms named A2shell*, whose whole premise was - # "shell and python only". Turning it off takes an explicit - # `file_system: {enabled: false}`, which _tool_on (tool_manager.py:47) reads. - # Left on for now, deliberately and with the effect known. - code_executor: - mcp: false - # python_env means "run in this process's machine", and that machine is the - # microVM -- the sandbox boundary is the VM itself, not this setting. Do not - # switch to the docker implementation: it would nest a container inside the - # VM for no extra isolation. - implementation: python_env - include: - - shell_executor - # Kept alongside the shell so that writing a file does not depend on - # getting a heredoc right. Dropping notebook_executor because it overlaps - # this one and adds a cell-state model nothing here needs. - - python_executor - todo_list: - mcp: false - # Kept out of the workspace root. The plan files default to - # `<output_dir>/plan.json` and `plan.md`, and output_dir *is* the directory - # whose end state becomes the task: 2 of ex11's 36 proposals wrote checks - # asserting the agent's own todo bookkeeping, one of them pinning - # `updated_at`, which no solver can reproduce. `.ms_agent/` is where - # ms_agent/project/paths.py says framework internals belong, and the - # workspace listing already skips it. - plan_filename: .ms_agent/plan.json - plan_md_filename: .ms_agent/plan.md - -# Every refusal ms-agent applies to a shell command, turned off. `allow_network` -# and the two list-valued keys are read by LLMAgent.prepare_runtime (llm_agent.py -# builds PermissionConfig.from_dict off this section) and take effect in the -# sandbox, where tool_server.py loads this same file. -# -# The last two keys are different: ms-agent's SafetyConfig does NOT implement -# them, and from_dict ignores unknown keys without a word, so on their own they -# would be dead letters. tool_server.py reads them itself and applies the two -# relaxations as a runtime patch (_patch_permission) inside the sandbox -- -# ms-agent is a harness twinkle supports, so it is used as released rather than -# forked. The startup line reports which ones took effect. -# -# The reason is what the refusals cost here rather than what they protect: this -# runs in a microVM that is reset once per episode and holds nothing but the -# workspace, while each refusal rules out a whole family of tasks the model could -# otherwise pose. `curl`/`wget` blocked means no task can fetch a source tarball -# or a dataset; the rm rules mean it cannot clear a directory (`rm -rf *` and -# `rm -rf build/*` are both refused) or write a task that starts from a mess that -# has to be cleaned up. -permission: - # Drops the default blacklist wholesale: curl, wget, ssh, scp, rsync, nc, - # netcat. (Whether the microVM actually has a route out is a separate - # question from whether the command is allowed to run.) - allow_network: true - safety_rules: - # Emptied, replacing the three baked-in patterns: `rm -rf /*`, `mkfs *`, - # `dd if=*`. An empty list here is not the same as an absent key -- absent - # means "use the defaults". - patterns: [] - # Same, for the configurable half of the rm/rmdir path check: `*`, `/*`, - # `/`, `~`. Left empty rather than removed to say the intent out loud; - # unrestricted_removal below bypasses the whole check, this list included, - # and tool_server.py warns at startup if the two ever disagree. - dangerous_removal_paths: [] - # And the half that a config cannot reach, which is why this one needs the - # runtime patch: the refusals written into is_dangerous_removal_path for - # `*`, anything ending in `/*`, `/`, a direct child of `/` (which - # `/workspace` is), and the home directory. - unrestricted_removal: true - # A separate refusal, found by running commands through SafetyGuard rather - # than by reading the config: a glob anywhere in a write or create path is - # denied on its own ("Glob patterns not allowed in write operations"), which - # is what actually stopped `rm -rf build/*` after the two lists above were - # emptied. It is not specific to rm -- `cp src/* dst/` is refused by the same - # check. (`chmod +x bin/*` is NOT: measured through SafetyGuard, chmod's - # arguments are not extracted as write paths, so it was already allowed.) - # Patched by widening the path that is scope-checked to the directory the - # glob expands inside, so a glob still cannot reach outside the workspace -- - # `cp /etc/* /workspace/` stays denied, now for being out of scope. - allow_write_globs: true - -# Web search is deliberately absent. ms-agent's `web_search` key only provides -# fetch_page (retrieve a known URL); a real query-a-search-engine tool needs -# EXA_API_KEY / SERPAPI_API_KEY and is wired separately from the plain tool -# list. Add it here once that is decided; until then no task should need it. - -# No `llm:` section on purpose, and note that omitting it is not the same as -# disabling it: ms-agent merges this file over its own ms_agent/agent/agent.yaml, -# which declares `service: modelscope`. The tool server treats a section with no -# credentials as absent, drops it, and then withdraws the one argument that -# needed it (read_file's `abbreviate`, an LLM-written file summary) from the -# advertised schema -- so the model is never offered a tool argument that cannot -# work. Put a real `llm:` here, with a key reachable from the sandbox, to get -# that argument back. diff --git a/.temp/retired_rsi/split_tasks.py b/.temp/retired_rsi/split_tasks.py deleted file mode 100644 index 6aeb00e4d..000000000 --- a/.temp/retired_rsi/split_tasks.py +++ /dev/null @@ -1,70 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Split a challenge.py task file into a training set and a held-out eval set. - -The eval set is stratified by ``n_pass`` -- the number of solver attempts that -succeeded when the task was filtered. Difficulty is the whole point of the -filter, so a random split can easily hand the eval set every task the model -already solves 3 times in 4, and a pass rate on those says nothing about the -hard end. Stratifying keeps both halves the same shape. - - python cookbook/rsi/agentic/split_tasks.py \\ - output/rsi_agentic/run3/challenge_flows.jsonl --eval-frac 0.25 -""" -import argparse -import json -import os -import random - - -def main(): - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('flows', help='challenge_flows.jsonl') - p.add_argument('--eval-frac', type=float, default=0.25) - p.add_argument('--seed', type=int, default=0) - p.add_argument('--train-out', default='', help='default: <flows dir>/train_tasks.jsonl') - p.add_argument('--eval-out', default='', help='default: <flows dir>/eval_tasks.jsonl') - args = p.parse_args() - - out_dir = os.path.dirname(os.path.abspath(args.flows)) - train_out = args.train_out or os.path.join(out_dir, 'train_tasks.jsonl') - eval_out = args.eval_out or os.path.join(out_dir, 'eval_tasks.jsonl') - - with open(args.flows, encoding='utf-8') as f: - tasks = [json.loads(line) for line in f if line.strip()] - if not tasks: - raise SystemExit(f'{args.flows} contains no tasks') - - strata = {} - for task in tasks: - strata.setdefault(task.get('n_pass'), []).append(task) - - rng = random.Random(args.seed) - train, held = [], [] - for n_pass in sorted(strata, key=lambda x: (x is None, x)): - group = strata[n_pass][:] - rng.shuffle(group) - # round() rather than int(): with 3 tasks at a difficulty and a quarter - # held out, truncating would give the eval set none of them. - n_eval = min(len(group) - 1, round(len(group) * args.eval_frac)) if len(group) > 1 else 0 - held.extend(group[:n_eval]) - train.extend(group[n_eval:]) - - for path, rows in ((train_out, train), (eval_out, held)): - with open(path, 'w', encoding='utf-8') as f: - for row in rows: - f.write(json.dumps(row, ensure_ascii=False) + '\n') - - def dist(rows): - out = {} - for row in rows: - out[row.get('n_pass')] = out.get(row.get('n_pass'), 0) + 1 - return dict(sorted(out.items(), key=lambda kv: (kv[0] is None, kv[0]))) - - print(f'{len(tasks)} tasks, n_pass dist {dist(tasks)}') - print(f'train {len(train)} -> {train_out} dist {dist(train)}') - print(f'eval {len(held)} -> {eval_out} dist {dist(held)}') - - -if __name__ == '__main__': - main() diff --git a/README.md b/README.md index 3d2b40eca..fc64af16f 100644 --- a/README.md +++ b/README.md @@ -109,7 +109,7 @@ sh INSTALL_MEGATRON.sh - 🎉2026-08-12 The ModelScope training service has been deployed to [Qwen/Qwen3.8-27B](https://www.modelscope.cn/models/Qwen/Qwen3.8-27B). - 🎉2026-08-04 Sandboxed multi-turn RL is now supported: run model-generated code in isolated [AgentENV](https://github.com/kvcache-ai/AgentENV) Firecracker microVMs, or in an OpenEnv server, with the same `train.py`. See the [cookbook](cookbook/rl/envs) and the [deployment guide](docs/source_en/Usage%20Guide/Agentic-RL-Deployment-and-Training.md). - 🎉2026-05-20 Support DeepSeek-V4-Flash and DeepSeek-V4-Pro models. -- 🎉2026-05-20 Multi-turn rollout and tool calling in RL are now supported. The Cookbook is currently being written. You can use `from twinkle_agentic.rollout import MultiTurnRollout/APIMultiTurnRollout` directly for multi-turn rollout. +- 🎉2026-05-20 Multi-turn rollout and tool calling in RL are now supported. The Cookbook is currently being written. You can use `from twinkle_agentic.rollout import MultiTurnRollout` directly for sampler, API, or mixed-backend multi-turn rollout. - 🎉2026-05-20 IM message alerting on training job failure is now supported. Usage: `import twinkle; twinkle.initialize(..., notifier=DingNotifier(...))`. - 🎉2026-04-27 Support the `padding_free` operation for sft/dpo/grpo/gkd, use `set_processor('InputProcessor', padding_free=True)` to train with it. - 🎉2026-04-22 The ModelScope service has been deployed to [Qwen/Qwen3.6-27B](https://www.modelscope.cn/models/Qwen/Qwen3.6-27B) with a new release 0.2.1. diff --git a/README_ZH.md b/README_ZH.md index f2f214f48..d7b3d66fa 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -105,7 +105,7 @@ Twinkle✨支持相同的算法接口运行在单GPU、torchrun多机、Ray、Cl - 🎉2026-08-12 ModelScope的训练服务部署为[Qwen/Qwen3.8-27B](https://www.modelscope.cn/models/Qwen/Qwen3.8-27B)。 - 🎉2026-08-04 支持沙箱环境下的多轮RL训练：模型生成的代码可在隔离的 [AgentENV](https://github.com/kvcache-ai/AgentENV) Firecracker microVM 或 OpenEnv 服务中执行，两个后端共用同一份 `train.py`。参考 [cookbook](cookbook/rl/envs) 和[部署文档](docs/source_zh/使用指引/Agentic%20RL部署与训练.md)。 - 🎉2026-05-20 支持DeepSeek-V4-Flash and DeepSeek-V4-Pro系列模型。 -- 🎉2026-05-20 支持多轮rollout和RL中的工具调用，Cookbook正在编写中，可以直接使用`from twinkle_agentic.rollout import MultiTurnRollout/APIMultiTurnRollout`进行多轮rollout。 +- 🎉2026-05-20 支持多轮rollout和RL中的工具调用，Cookbook正在编写中，可以直接使用 `from twinkle_agentic.rollout import MultiTurnRollout` 进行 sampler、API 或混合后端的多轮 rollout。 - 🎉2026-05-20 支持训练任务失败后的IM消息告警, 使用方式: `import twinkle; twinkle.initialize(..., notifier=DingNotifier(...))`。 - 🎉2026-04-27 支持sft/dpo/grpo/gkd的padding_free方法, 使用`set_processor('InputProcessor', padding_free=True)`来开启训练。 - 🎉2026-04-22 ModelScope的训练服务部署为[Qwen/Qwen3.6-27B](https://www.modelscope.cn/models/Qwen/Qwen3.6-27B)，并发布了0.2.1版本。 diff --git a/cookbook/rsi/code/challenge.py b/cookbook/rsi/code/challenge.py index 261217abe..4e3b6383c 100644 --- a/cookbook/rsi/code/challenge.py +++ b/cookbook/rsi/code/challenge.py @@ -31,7 +31,7 @@ from twinkle.sampler import vLLMSampler from twinkle_agentic.challenger import CodeChallenger, KeywordStore, load_seeds from twinkle_agentic.envs import LocalEnv -from twinkle_agentic.rollout import build_rollout +from twinkle_agentic.rollout import MultiTurnRollout from twinkle_agentic.tools.tool_manager import ToolManager sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) @@ -132,7 +132,7 @@ def main(): # Single-turn generation, but through the same rollout the RL loop uses, so a # challenger that should be allowed to run code while inventing only needs a # tool manager here rather than a different code path. - explorer = build_rollout( + explorer = MultiTurnRollout( sampler, template=template, tool_manager=ToolManager([]), diff --git a/cookbook/rsi/code/collect.py b/cookbook/rsi/code/collect.py index c18a907ad..2b15c6786 100644 --- a/cookbook/rsi/code/collect.py +++ b/cookbook/rsi/code/collect.py @@ -32,7 +32,7 @@ from twinkle.data_format import SamplingParams, Trajectory, user_data_get from twinkle_agentic.challenger import CodeChallenger, KeywordStore, load_seeds from twinkle_agentic.envs import LocalEnv -from twinkle_agentic.rollout import build_rollout +from twinkle_agentic.rollout import MultiTurnRollout from twinkle_agentic.tools.tool_manager import ToolManager # Appended, not prepended: rsi.py imports this half into the process that already @@ -115,9 +115,9 @@ def build_challenger(args, sampler, template, *, recorder=None) -> CollectingCha # One rollout for proposing and, through solver_params, for solving. max_turns=1 # because a code answer is one message: there is nothing for a second turn to # react to until the asserts have run, and running them is the next stage. - explorer = build_rollout(sampler, template=template, - tool_manager=ToolManager([]), max_turns=1, - sampling_params=params) + explorer = MultiTurnRollout(sampler, template=template, + tool_manager=ToolManager([]), max_turns=1, + sampling_params=params) store = None if args.code_keywords_n > 0: store = KeywordStore(args.code_keyword_db, CATEGORIES) diff --git a/docs/source_en/Components/Agentic/Envs.md b/docs/source_en/Components/Agentic/Envs.md index 988eac6fe..774f412c0 100644 --- a/docs/source_en/Components/Agentic/Envs.md +++ b/docs/source_en/Components/Agentic/Envs.md @@ -191,7 +191,7 @@ Downstream usage is the same for both modes: ```python from twinkle_agentic.envs.env_tool import EnvTool from twinkle_agentic.tools.tool_manager import ToolManager -from twinkle_agentic.rollout.api_multi_turn import APIMultiTurnRollout +from twinkle_agentic.rollout import MultiTurnRollout env.reset() @@ -200,7 +200,7 @@ env_tools = EnvTool.from_env(env) manager = ToolManager(env_tools) # Use in rollout -rollout = APIMultiTurnRollout(api=api, tool_manager=manager, max_turns=10) +rollout = MultiTurnRollout(api=api, template=template, tool_manager=manager, max_turns=10) results = rollout(trajectories) ``` diff --git a/docs/source_en/Components/Agentic/Multi-Turn-Tool-Usage.md b/docs/source_en/Components/Agentic/Multi-Turn-Tool-Usage.md index 24296f584..c4b4615be 100644 --- a/docs/source_en/Components/Agentic/Multi-Turn-Tool-Usage.md +++ b/docs/source_en/Components/Agentic/Multi-Turn-Tool-Usage.md @@ -19,8 +19,9 @@ The simplest way to run a multi-turn tool-use rollout using an OpenAI-compatible from twinkle_agentic.protocol.openai import OpenAI from twinkle_agentic.tools.base import Tool from twinkle_agentic.tools.tool_manager import ToolManager -from twinkle_agentic.rollout.api_multi_turn import APIMultiTurnRollout +from twinkle_agentic.rollout import MultiTurnRollout from twinkle.data_format.sampling import SamplingParams +from twinkle.template import Template # 1. Define tools class WeatherTool(Tool): @@ -47,16 +48,17 @@ class WeatherTool(Tool): # 2. Set up ToolManager manager = ToolManager([WeatherTool()]) -# 3. Create API client -api = OpenAI(model='qwen3.5-32b', base_url='http://localhost:8000/v1') +# 3. Create API client and the local template used to encode its replies +api = OpenAI(model='qwen3.5-32b', base_url='http://localhost:8000/v1', concurrency=8) +template = Template(model_id='Qwen/Qwen3.5-32B') # 4. Create rollout -rollout = APIMultiTurnRollout( +rollout = MultiTurnRollout( api=api, + template=template, tool_manager=manager, sampling_params=SamplingParams(temperature=0.7, max_tokens=2048), max_turns=6, - concurrency=8, ) # 5. Prepare trajectories @@ -138,7 +140,7 @@ env_tools = EnvTool.from_env(env) manager = ToolManager(env_tools) # Use manager in rollout as usual -rollout = APIMultiTurnRollout(api=api, tool_manager=manager, max_turns=10) +rollout = MultiTurnRollout(api=api, template=template, tool_manager=manager, max_turns=10) ``` ## Using OpenEnv Environments @@ -193,11 +195,12 @@ results = rollout(trajectories, tool_manager=managers) ## Trace Debugging -Both rollout implementations support trace dumps for debugging: +The unified rollout supports trace dumps for debugging: ```python -rollout = APIMultiTurnRollout( +rollout = MultiTurnRollout( api=api, + template=template, tool_manager=manager, trace_dir='traces/', trace_callback=lambda t: t['turns'] > 1, # Only store multi-turn diff --git a/docs/source_en/Components/Agentic/Rollout.md b/docs/source_en/Components/Agentic/Rollout.md index e803b9076..10ccc3471 100644 --- a/docs/source_en/Components/Agentic/Rollout.md +++ b/docs/source_en/Components/Agentic/Rollout.md @@ -1,6 +1,6 @@ # Multi-Turn Rollout -The Rollout module provides multi-turn conversation rollout engines for agentic RLHF training. Two implementations are available: `MultiTurnRollout` for batched vLLM sampling and `APIMultiTurnRollout` for OpenAI-compatible API endpoints. +The Rollout module provides one multi-turn conversation engine for agentic RLHF training. `MultiTurnRollout` can generate each assistant turn with a local sampler, an OpenAI-compatible API, or a callback that chooses between them. ## Rollout Base Class @@ -19,12 +19,12 @@ All rollouts accept a list of trajectories and return the same number of traject ## MultiTurnRollout -Batched multi-turn rollout engine that uses a vLLM sampler for generation. All active trajectories are sampled in a single batched call per turn for maximum throughput. +Multi-turn rollout engine supporting local samplers, external APIs, and per-turn backend selection. Each trajectory runs independently in the rollout thread pool. ### Per-turn Loop 1. Encode each trajectory into an `InputFeature` with a generation prompt -2. Batch `sampler.sample(active_pifs)` — all live trajectories in parallel +2. Call `response_callback(...)` to obtain one `SampledSequence` from the sampler or API 3. Check termination: `stop_reason == 'length'`, no tool calls, or max turns reached 4. Dispatch tools via `ToolManager`, append tool responses 5. Compute bridge tokens (tool turns + generation prompt) with `labels = -100` @@ -53,8 +53,12 @@ results = rollout(trajectories) | Parameter | Type | Description | |-----------|------|-------------| -| `sampler` | Sampler | vLLM sampler instance for batched generation. | -| `template` | `Template` | Chat template for encoding/decoding. | +| `sampler` | Sampler | Local sampler. Used by default when both backends exist. | +| `api` | `API` | Optional external generation API. | +| `template` | `Template` | Required local chat template for encoding every backend's output. | +| `response_callback` | `Callable` | Optional per-turn backend selector returning `SampledSequence`. | +| `api_appended_as` | `str` | API turns are `demonstration` (SFT only) or `context` (no loss). | +| `api_kwargs` | `Dict` | Request fields forwarded to each API call. | | `tool_manager` | `ToolManager` | Tool dispatcher. Can also be passed per-call. | | `sampling_params` | `SamplingParams` | Default sampling parameters. | | `max_turns` | `int` | Maximum number of turns per trajectory (default: 6). | @@ -72,6 +76,7 @@ Each output trajectory dict includes: | `messages` | `List[Dict]` | Full conversation including tool turns. | | `input_ids` | `List[int]` | Token IDs of the full sequence. | | `labels` | `List[int]` | Training labels (`-100` for non-trainable tokens). | +| `completion_mask` | `List[int]` | Policy-generated positions that carry rollout log probabilities. | | `turns` | `int` | Number of turns performed. | | `stop_reason` | `str` | `'stop'` / `'length'` | | `truncated` | `bool` | Whether the trajectory was cut off rather than concluding on its own: generation hit `max_tokens` (`stop_reason='length'`), the turn limit was reached, or a length cap dropped it. | @@ -87,54 +92,33 @@ rollout_actor = MultiTurnRollout.remote(sampler=sampler, template=template, ...) results = ray.get(rollout_actor.__call__.remote(trajectories)) ``` -## APIMultiTurnRollout +## API and Mixed-Backend Rollouts -Multi-turn rollout over an OpenAI-compatible chat-completions API. Each trajectory runs independently in a thread pool for network concurrency. +API-only rollout uses the same class and still requires the local template that tokenizes external replies: ```python -from twinkle_agentic.rollout.api_multi_turn import APIMultiTurnRollout from twinkle_agentic.protocol.openai import OpenAI +from twinkle_agentic.rollout import MultiTurnRollout -api = OpenAI(model='qwen3.5-32b', base_url='http://localhost:8000/v1') - -rollout = APIMultiTurnRollout( - api=api, +api = OpenAI(model='qwen3.5-32b', base_url='http://localhost:8000/v1', concurrency=8) +rollout = MultiTurnRollout( + api, + template=template, tool_manager=tool_manager, sampling_params=SamplingParams(temperature=0.7), max_turns=6, - concurrency=8, trace_dir='api_traces/', ) - results = rollout(trajectories) ``` -### Parameters - -| Parameter | Type | Description | -|-----------|------|-------------| -| `api` | `OpenAI` | OpenAI-compatible API client. | -| `tool_manager` | `ToolManager` | Tool dispatcher (single or per-trajectory list). | -| `sampling_params` | `SamplingParams` | Default sampling parameters. | -| `max_turns` | `int` | Maximum turns per trajectory (default: 6). | -| `concurrency` | `int` | Thread pool size for parallel API calls (default: 8). | -| `extra_body` | `Dict` | Extra fields to include in API requests. | -| `trace_dir` | `str` | Directory for trace dumps. | +When both `sampler` and `api` are supplied, the default is the sampler. Pass `response_callback` to choose per turn; it receives both backends and must return one `SampledSequence`. API turns have no rollout log probabilities, so `api_appended_as='demonstration'` includes them in SFT but excludes them from GRPO. Use `'context'` to exclude them from both. ### Stop Reasons | Reason | Description | |--------|-------------| | `stop` | Assistant responded without tool calls (natural end). | -| `length` | API returned `finish_reason='length'` (token limit). | -| `max_turns` | Reached `max_turns` limit. | -| `api_error` | API call or tool execution raised an exception. | - -## Choosing Between Rollouts - -| Feature | MultiTurnRollout | APIMultiTurnRollout | -|---------|-----------------|---------------------| -| **Backend** | vLLM sampler (local GPU) | OpenAI-compatible API | -| **Training integration** | Produces `input_ids` / `labels` for GRPO | Messages only (for data collection) | -| **Batching** | GPU-level batch parallelism | Network-level thread concurrency | -| **Use case** | Online RLHF training loop | Offline data generation / evaluation | +| `length` | Generation reached its token limit. | +| `max_turns` | Reached the tool-turn limit without a follow-up. | +| `generation_error` | The external endpoint failed before returning a valid response. | diff --git "a/docs/source_zh/\347\273\204\344\273\266/Agentic/Envs.md" "b/docs/source_zh/\347\273\204\344\273\266/Agentic/Envs.md" index 76c729eb7..bae412f5c 100644 --- "a/docs/source_zh/\347\273\204\344\273\266/Agentic/Envs.md" +++ "b/docs/source_zh/\347\273\204\344\273\266/Agentic/Envs.md" @@ -191,7 +191,7 @@ env.close() ```python from twinkle_agentic.envs.env_tool import EnvTool from twinkle_agentic.tools.tool_manager import ToolManager -from twinkle_agentic.rollout.api_multi_turn import APIMultiTurnRollout +from twinkle_agentic.rollout import MultiTurnRollout env.reset() @@ -200,7 +200,7 @@ env_tools = EnvTool.from_env(env) manager = ToolManager(env_tools) # 在 rollout 中使用 -rollout = APIMultiTurnRollout(api=api, tool_manager=manager, max_turns=10) +rollout = MultiTurnRollout(api=api, template=template, tool_manager=manager, max_turns=10) results = rollout(trajectories) ``` diff --git "a/docs/source_zh/\347\273\204\344\273\266/Agentic/Multi-Turn-Tool-Usage.md" "b/docs/source_zh/\347\273\204\344\273\266/Agentic/Multi-Turn-Tool-Usage.md" index 8b94b2ed4..b3af8b45f 100644 --- "a/docs/source_zh/\347\273\204\344\273\266/Agentic/Multi-Turn-Tool-Usage.md" +++ "b/docs/source_zh/\347\273\204\344\273\266/Agentic/Multi-Turn-Tool-Usage.md" @@ -19,8 +19,9 @@ Agentic rollout 管线由四个核心组件组成： from twinkle_agentic.protocol.openai import OpenAI from twinkle_agentic.tools.base import Tool from twinkle_agentic.tools.tool_manager import ToolManager -from twinkle_agentic.rollout.api_multi_turn import APIMultiTurnRollout +from twinkle_agentic.rollout import MultiTurnRollout from twinkle.data_format.sampling import SamplingParams +from twinkle.template import Template # 1. 定义工具 class WeatherTool(Tool): @@ -47,16 +48,17 @@ class WeatherTool(Tool): # 2. 设置 ToolManager manager = ToolManager([WeatherTool()]) -# 3. 创建 API 客户端 -api = OpenAI(model='qwen3.5-32b', base_url='http://localhost:8000/v1') +# 3. 创建 API 客户端，以及用于编码 API 回复的本地 template +api = OpenAI(model='qwen3.5-32b', base_url='http://localhost:8000/v1', concurrency=8) +template = Template(model_id='Qwen/Qwen3.5-32B') # 4. 创建 rollout -rollout = APIMultiTurnRollout( +rollout = MultiTurnRollout( api=api, + template=template, tool_manager=manager, sampling_params=SamplingParams(temperature=0.7, max_tokens=2048), max_turns=6, - concurrency=8, ) # 5. 准备轨迹 @@ -138,7 +140,7 @@ env_tools = EnvTool.from_env(env) manager = ToolManager(env_tools) # 照常在 rollout 中使用 manager -rollout = APIMultiTurnRollout(api=api, tool_manager=manager, max_turns=10) +rollout = MultiTurnRollout(api=api, template=template, tool_manager=manager, max_turns=10) ``` ## 使用 OpenEnv 环境 @@ -193,11 +195,12 @@ results = rollout(trajectories, tool_manager=managers) ## 跟踪调试 -两种 rollout 实现都支持跟踪文件输出用于调试： +统一的 rollout 支持跟踪文件输出用于调试： ```python -rollout = APIMultiTurnRollout( +rollout = MultiTurnRollout( api=api, + template=template, tool_manager=manager, trace_dir='traces/', trace_callback=lambda t: t['turns'] > 1, # 仅存储多轮对话 diff --git "a/docs/source_zh/\347\273\204\344\273\266/Agentic/Rollout.md" "b/docs/source_zh/\347\273\204\344\273\266/Agentic/Rollout.md" index 134532c97..767e8a538 100644 --- "a/docs/source_zh/\347\273\204\344\273\266/Agentic/Rollout.md" +++ "b/docs/source_zh/\347\273\204\344\273\266/Agentic/Rollout.md" @@ -1,6 +1,6 @@ # 多轮 Rollout -Rollout 模块提供了用于 Agentic RLHF 训练的多轮对话 rollout 引擎。包含两种实现：用于批量 vLLM 采样的 `MultiTurnRollout` 和用于 OpenAI 兼容 API 端点的 `APIMultiTurnRollout`。 +Rollout 模块提供统一的多轮对话引擎 `MultiTurnRollout`，每轮 assistant 可由本地 sampler、OpenAI 兼容 API，或在两者间动态选择的 callback 生成。 ## Rollout 基类 @@ -19,12 +19,12 @@ class Rollout(ABC): ## MultiTurnRollout -批量多轮 rollout 引擎，使用 vLLM 采样器进行生成。每轮中所有活跃轨迹通过单次批量采样调用并行处理，最大化吞吐量。 +统一的多轮 rollout 引擎，支持本地 sampler、外部 API 和逐轮后端选择。每条轨迹在线程池中独立执行。 ### 每轮循环 1. 将每个轨迹编码为带生成提示的 `InputFeature` -2. 批量调用 `sampler.sample(active_pifs)` —— 所有活跃轨迹并行 +2. 调用 `response_callback(...)`，从 sampler 或 API 获取一个 `SampledSequence` 3. 检查终止条件：`stop_reason == 'length'`、无工具调用、或达到最大轮次 4. 通过 `ToolManager` 分发工具调用，追加工具响应 5. 计算桥接 token（工具轮次 + 生成提示），设置 `labels = -100` @@ -53,8 +53,12 @@ results = rollout(trajectories) | 参数 | 类型 | 说明 | |------|------|------| -| `sampler` | Sampler | 用于批量生成的 vLLM 采样器实例。 | -| `template` | `Template` | 用于编码/解码的聊天模板。 | +| `sampler` | Sampler | 本地 sampler；两个后端同时存在时默认使用它。 | +| `api` | `API` | 可选的外部生成 API。 | +| `template` | `Template` | 必传；用于编码所有后端的输出。 | +| `response_callback` | `Callable` | 可选的逐轮后端选择器，返回 `SampledSequence`。 | +| `api_appended_as` | `str` | API 轮为 `demonstration`（仅 SFT）或 `context`（不训练）。 | +| `api_kwargs` | `Dict` | 传给每次 API 调用的请求字段。 | | `tool_manager` | `ToolManager` | 工具分发器。也可以按调用传入。 | | `sampling_params` | `SamplingParams` | 默认采样参数。 | | `max_turns` | `int` | 每个轨迹的最大轮次（默认：6）。 | @@ -72,6 +76,7 @@ results = rollout(trajectories) | `messages` | `List[Dict]` | 包含工具轮次的完整对话。 | | `input_ids` | `List[int]` | 完整序列的 token ID。 | | `labels` | `List[int]` | 训练标签（非可训练 token 为 `-100`）。 | +| `completion_mask` | `List[int]` | 由 policy 生成且具有 rollout log probability 的位置。 | | `turns` | `int` | 执行的轮次数。 | | `stop_reason` | `str` | `'stop'` / `'length'` | | `truncated` | `bool` | 轨迹是否被截断（而非自行结束）：生成触及 `max_tokens`（`stop_reason='length'`）、达到轮次上限，或被长度上限丢弃。 | @@ -87,54 +92,33 @@ rollout_actor = MultiTurnRollout.remote(sampler=sampler, template=template, ...) results = ray.get(rollout_actor.__call__.remote(trajectories)) ``` -## APIMultiTurnRollout +## API 与混合后端 Rollout -通过 OpenAI 兼容 chat-completions API 进行多轮 rollout。每个轨迹在线程池中独立运行，实现网络并发。 +纯 API 模式使用同一个类，并仍需传入本地 template，以便将外部回复编码成训练侧一致的 token： ```python -from twinkle_agentic.rollout.api_multi_turn import APIMultiTurnRollout from twinkle_agentic.protocol.openai import OpenAI +from twinkle_agentic.rollout import MultiTurnRollout -api = OpenAI(model='qwen3.5-32b', base_url='http://localhost:8000/v1') - -rollout = APIMultiTurnRollout( - api=api, +api = OpenAI(model='qwen3.5-32b', base_url='http://localhost:8000/v1', concurrency=8) +rollout = MultiTurnRollout( + api, + template=template, tool_manager=tool_manager, sampling_params=SamplingParams(temperature=0.7), max_turns=6, - concurrency=8, trace_dir='api_traces/', ) - results = rollout(trajectories) ``` -### 参数 - -| 参数 | 类型 | 说明 | -|------|------|------| -| `api` | `OpenAI` | OpenAI 兼容 API 客户端。 | -| `tool_manager` | `ToolManager` | 工具分发器（单个或按轨迹的列表）。 | -| `sampling_params` | `SamplingParams` | 默认采样参数。 | -| `max_turns` | `int` | 每轨迹最大轮次（默认：6）。 | -| `concurrency` | `int` | 并行 API 调用的线程池大小（默认：8）。 | -| `extra_body` | `Dict` | API 请求中附加的额外字段。 | -| `trace_dir` | `str` | 跟踪文件目录。 | +同时传入 `sampler` 和 `api` 时，默认使用 sampler。传入 `response_callback` 可逐轮选择后端；callback 会收到两个后端，并必须返回一个 `SampledSequence`。API 轮没有 rollout log probability，因此 `api_appended_as='demonstration'` 会让它参与 SFT 但跳过 GRPO；使用 `'context'` 可让它完全不参与训练。 ### 停止原因 | 原因 | 说明 | |------|------| | `stop` | 助手回复未包含工具调用（自然结束）。 | -| `length` | API 返回 `finish_reason='length'`（token 限制）。 | -| `max_turns` | 达到 `max_turns` 限制。 | -| `api_error` | API 调用或工具执行抛出异常。 | - -## 选择建议 - -| 特性 | MultiTurnRollout | APIMultiTurnRollout | -|------|-----------------|---------------------| -| **后端** | vLLM 采样器（本地 GPU） | OpenAI 兼容 API | -| **训练集成** | 生成 `input_ids` / `labels` 用于 GRPO | 仅消息（用于数据收集） | -| **批处理** | GPU 级别批量并行 | 网络级别线程并发 | -| **用例** | 在线 RLHF 训练循环 | 离线数据生成 / 评估 | +| `length` | 生成达到 token 上限。 | +| `max_turns` | 达到工具轮次上限且没有 follow-up。 | +| `generation_error` | 外部端点未能返回有效响应。 | diff --git a/src/twinkle/loss/grpo.py b/src/twinkle/loss/grpo.py index 33b89c206..85470a0dc 100644 --- a/src/twinkle/loss/grpo.py +++ b/src/twinkle/loss/grpo.py @@ -227,6 +227,39 @@ def _pad_and_align_to_batch( return result + def _resolve_loss_mask(self, inputs: Dict, labels: 'torch.Tensor') -> 'torch.Tensor': + """Positions this loss may score: trainable *and* log-prob-bearing. + + ``labels`` alone answers "should this token be scored", which is all SFT + needs. A policy-gradient loss also needs a sampling log-prob per token to + form an importance ratio, and a turn produced outside the sampled policy + (an API, a human, a replayed demonstration) has none. Such turns carry + ``completion_mask == 0``: excluded here, yet still trainable for SFT. + + A feature without ``completion_mask`` predates the field, and there every + trainable token was the policy's own, so the mask degenerates to + ``labels != ignore_index`` and old trajectories train exactly as before. + """ + import torch + trainable = (labels != self.ignore_index).bool() + completion_mask = inputs.get('completion_mask') + if completion_mask is None: + return trainable + if not torch.is_tensor(completion_mask): + completion_mask = torch.as_tensor(completion_mask) + completion_mask = completion_mask.to(trainable.device) + if completion_mask.dim() == 1: + completion_mask = completion_mask.unsqueeze(0) + if completion_mask.shape != trainable.shape: + raise ValueError(f'completion_mask shape {tuple(completion_mask.shape)} does not match labels shape ' + f'{tuple(trainable.shape)}. A misaligned mask would apply importance ratios to ' + 'the wrong tokens, so it is refused rather than broadcast.') + loss_mask = trainable & completion_mask.bool() + if self.enable_sampling_replay and not bool((loss_mask == trainable).all()): + raise ValueError('sampling replay does not support turns generated outside the sampled policy: ' + 'they are trainable but have no sampling mask to replay against.') + return loss_mask + def __call__( self, inputs: Dict, @@ -269,7 +302,7 @@ def __call__( logps = outputs.get('logps') if self.enable_sampling_replay and logps is None: raise RuntimeError('sampling replay logps must be computed by the model forward') - loss_mask = (labels != self.ignore_index).bool() + loss_mask = self._resolve_loss_mask(inputs, labels) if logps is None: logits = outputs.get('logits') if logits.shape[1] != labels.shape[1]: diff --git a/src/twinkle/model/megatron/megatron.py b/src/twinkle/model/megatron/megatron.py index 5240816d3..8d85dacf8 100644 --- a/src/twinkle/model/megatron/megatron.py +++ b/src/twinkle/model/megatron/megatron.py @@ -402,6 +402,8 @@ def post_loss_function(output_tensor, inputs, logps, unpacked_logits=None, entro def forward_step_func(data_iterator, model): batch = next(data_iterator) labels = batch.pop('labels', None) + # Not a model argument; restored below so the loss can read it. + completion_mask = batch.pop('completion_mask', None) unwrapped_model = self.strategy.unwrap_model([model])[0] if disable_lora and isinstance(unwrapped_model, PeftModel): with unwrapped_model.disable_adapter(): @@ -410,6 +412,8 @@ def forward_step_func(data_iterator, model): output_tensor = model(**batch) batch['labels'] = labels + if completion_mask is not None: + batch['completion_mask'] = completion_mask logps = None unpacked_logits = None entropies = None @@ -440,6 +444,10 @@ def forward_step_func(data_iterator, model): if entropies is not None: entropies = processor.postprocess_tensor_cp(entropies, cu_seqlens=cu_seqlens_q) batch['labels'] = processor.postprocess_tensor_cp(labels, cu_seqlens=cu_seqlens_q) + if completion_mask is not None: + # Same index space as labels, so it needs the same CP reassembly. + batch['completion_mask'] = processor.postprocess_tensor_cp( + completion_mask, cu_seqlens=cu_seqlens_q) if 'position_ids' in batch: pos = batch['position_ids'] if pos.dim() == 3: diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py index 46ace2c64..9d0cec9f1 100644 --- a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py +++ b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py @@ -845,7 +845,7 @@ def prepare_inputs(self, inputs): """Prepare inputs 1. set extra_kwargs['position_ids'] - 2. split labels + 2. split labels, and completion_mask when present """ input_ids = inputs.get('input_ids') position_ids = inputs.get('position_ids') @@ -863,7 +863,11 @@ def prepare_inputs(self, inputs): self.extra_kwargs['input_ids'] = input_ids.clone() if 'labels' in inputs: labels = inputs.get('labels') - _, _, labels, _, _, _, _ = self.pad_and_split_inputs( + # completion_mask sits on the labels' index space, so it is padded and + # split identically -- unlike loss_scale, which is rolled beforehand. + completion_mask = inputs.get('completion_mask') + extra_split_values = None if completion_mask is None else [(completion_mask, 0, -1)] + _, _, labels, _, _, _, extra_values = self.pad_and_split_inputs( None, None, labels, @@ -871,8 +875,11 @@ def prepare_inputs(self, inputs): None, None, real_position_ids=real_position_ids, + extra_split_values=extra_split_values, ) inputs['labels'] = labels + if extra_values: + inputs['completion_mask'] = extra_values[0] return inputs @@ -986,6 +993,19 @@ def _trim_gathered_sequence_padding(tensor: torch.Tensor, real_position_ids: tor return torch.cat(pieces, dim=1).contiguous() if pieces else tensor[:, :0].contiguous() return tensor[:, :real_position_ids.shape[-1]].contiguous() + def _gather_completion_mask(self, inputs: Dict[str, Any], real_position_ids) -> None: + """Gather ``completion_mask`` in place, mirroring the labels gather. + + Deliberately not routed through :class:`GatherLoss`: the mask carries no + gradient, and reusing that autograd Function would attach a second backward + path to whichever tensor were passed alongside it, double-scaling its grad. + """ + mask = inputs.get('completion_mask') + if mask is None or not torch.is_tensor(mask) or mask.dim() < 2: + return + gathered = sequence_parallel.gather(mask, dim=1, position_ids=real_position_ids) + inputs['completion_mask'] = self._trim_gathered_sequence_padding(gathered, real_position_ids) + def gather_loss_tensors( self, inputs: Dict[str, Any], @@ -1017,6 +1037,7 @@ def gather_loss_tensors( gathered_labels = self._trim_gathered_sequence_padding(gathered_labels, real_position_ids) outputs['logits'] = gathered_hidden inputs['labels'] = gathered_labels + self._gather_completion_mask(inputs, real_position_ids) return inputs, outputs if labels is None or logps is None: return inputs, outputs @@ -1031,6 +1052,7 @@ def gather_loss_tensors( gathered_labels = self._trim_gathered_sequence_padding(gathered_labels, real_position_ids) outputs['logps'] = gathered_logps inputs['labels'] = gathered_labels + self._gather_completion_mask(inputs, real_position_ids) entropies = outputs.get('entropies') if entropies is not None and torch.is_tensor(entropies) and entropies.dim() >= 2: gathered_entropies, _ = GatherLoss.apply(entropies, labels, 1, real_position_ids) diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index dda7aea8d..9ce4c4cb7 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -568,6 +568,8 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec enable_sp=getattr(self, '_enable_sp', False), ) labels: torch.Tensor = inputs.pop('labels', None) + # Not a model argument; the loss reads it back off `inputs` further down. + completion_mask = inputs.pop('completion_mask', None) replay_metadata = replay_loss_mask = replay_masked_labels = None if enable_sampling_replay: replay_loss_mask, replay_masked_labels, replay_metadata = _prepare_sampling_replay( @@ -595,6 +597,8 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec recorded_routing = rr_cleanup() inputs['labels'] = labels + if completion_mask is not None: + inputs['completion_mask'] = completion_mask if task != 'embedding' and labels is not None and loss_require_logps: loss_mask = replay_loss_mask if enable_sampling_replay else (labels != -100).bool() masked_labels = replay_masked_labels if enable_sampling_replay else labels.masked_fill(~loss_mask, 0) @@ -689,6 +693,8 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T enable_sp=getattr(self, '_enable_sp', False), ) labels = inputs.pop('labels', None) + # Not a model argument; the loss reads it back off `inputs` further down. + completion_mask = inputs.pop('completion_mask', None) replay_metadata = replay_loss_mask = replay_masked_labels = None if enable_sampling_replay: packed_position_ids = processor._is_packed_position_ids(inputs.get('position_ids')) @@ -720,6 +726,8 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T recorded_routing = rr_cleanup() inputs['labels'] = labels + if completion_mask is not None: + inputs['completion_mask'] = completion_mask if task != 'embedding' and labels is not None and loss_require_logps: loss_mask = replay_loss_mask if enable_sampling_replay else (labels != -100).bool() masked_labels = replay_masked_labels if enable_sampling_replay else labels.masked_fill(~loss_mask, 0) diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py index 5d67f1fcb..b28812f72 100644 --- a/src/twinkle/processor/base.py +++ b/src/twinkle/processor/base.py @@ -608,9 +608,10 @@ def unpack_packed_sequences( """Unpack packed (padding_free) sequences into per-sequence batch format. Called after SP gather / CP gather, before loss computation. - Unpacks ``labels`` and any present output keys (``logps``, ``logits``) - from ``[1, total_tokens, ...]`` to ``[num_sequences, max_seq_len, ...]``. - Keys that are ``None`` are silently skipped. + Unpacks ``labels``, ``completion_mask`` and any present output keys + (``logps``, ``logits``) from ``[1, total_tokens, ...]`` to + ``[num_sequences, max_seq_len, ...]``. Keys that are ``None`` are silently + skipped. For ``task='embedding'`` the outputs are already pooled to ``[n_seqs, H]`` by ``postprocess_tensor_sp``, so this is a no-op. @@ -627,23 +628,29 @@ def unpack_packed_sequences( from copy import copy - # Collect output keys to unpack: (key, pad_value) - output_keys = [] - for key, pad_val in [('logps', 0), ('values', 0), ('entropies', 0), ('logits', 0)]: - if outputs and outputs.get(key) is not None: - output_keys.append((key, pad_val)) - - all_tensors = [labels] + [outputs[k] for k, _ in output_keys] - all_pads = [-100] + [p for _, p in output_keys] - unpacked = self._unpack_by_position_ids(position_ids, *all_tensors, padding_values=all_pads) + # (key, tensor, pad_value) for everything that must come back as + # [num_sequences, max_seq_len]. completion_mask shares the labels' index + # space, so leaving it packed would hand the loss two differently shaped + # views of the same sequence. + input_specs = [('labels', labels, -100)] + if inputs.get('completion_mask') is not None: + input_specs.append(('completion_mask', inputs['completion_mask'], self.padding_map['completion_mask'])) + output_specs = [(key, outputs[key], 0) for key in ('logps', 'values', 'entropies', 'logits') + if outputs and outputs.get(key) is not None] + + specs = input_specs + output_specs + unpacked = iter( + self._unpack_by_position_ids( + position_ids, *[tensor for _, tensor, _ in specs], padding_values=[pad for _, _, pad in specs])) inputs = copy(inputs) - inputs['labels'] = unpacked[0] + for key, _, _ in input_specs: + inputs[key] = next(unpacked) - if output_keys: + if output_specs: outputs = copy(outputs) - for i, (key, _) in enumerate(output_keys): - outputs[key] = unpacked[i + 1] + for key, _, _ in output_specs: + outputs[key] = next(unpacked) return inputs, outputs diff --git a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py index 72600bc4a..7877ddb55 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py @@ -36,12 +36,10 @@ def _convert_ndarray_to_list(obj: Any) -> Any: return obj -# max_concurrency: how many sample() calls one worker serves at once. Without it -# Ray runs one method per actor at a time, so concurrent callers queue at the actor -# and never share a batch inside AsyncLLM. 24 is what vLLM reports as the maximum -# concurrency its KV cache holds for this context length; past it vLLM preempts and -# recomputes, which costs more than it gains. -@remote_class(max_concurrency=24) +_MAX_CONCURRENCY = max(1, int(os.environ.get('TWINKLE_SAMPLER_MAX_CONCURRENCY') or 24)) + + +@remote_class(max_concurrency=_MAX_CONCURRENCY) class vLLMSampler(Sampler, CheckpointEngineMixin): """A vLLM-based sampler using VLLMEngine (AsyncLLM). diff --git a/src/twinkle/server/sampler/backends/mock_sampler.py b/src/twinkle/server/sampler/backends/mock_sampler.py index d8355008b..2d5e5930d 100644 --- a/src/twinkle/server/sampler/backends/mock_sampler.py +++ b/src/twinkle/server/sampler/backends/mock_sampler.py @@ -252,8 +252,9 @@ def _build_new_input_feature(pif: Any, tokens: list[int]) -> dict[str, Any]: Produces a plain-dict ``InputFeature`` that carries the running context for the next multi-turn round: ``input_ids`` is the prior prompt plus - this round's sampled tokens, and ``labels`` marks the sampled tokens as - trainable (their own ids) while prior/context positions stay ``-100``. + this round's sampled tokens, ``labels`` marks the sampled tokens as + trainable (their own ids) while prior/context positions stay ``-100``, + and ``completion_mask`` marks them as the policy's own output. This mirrors the shape a real sampler's ``concat_input_feature`` yields, which the multi-turn rollout relies on (it reads ``new_input_feature.input_ids`` and counts trainable ``labels``). @@ -270,8 +271,14 @@ def _build_new_input_feature(pif: Any, tokens: list[int]) -> dict[str, Any]: # No (or misaligned) prior labels: treat the entire prior context as # non-trainable so only this round's sampled tokens count. labels = [-100] * len(prev_ids) + prev_mask = feat.get('completion_mask') + if prev_mask is not None and len(prev_mask) == len(prev_ids): + completion_mask = list(prev_mask) + else: + completion_mask = [0 if label == -100 else 1 for label in labels] feat['input_ids'] = prev_ids + list(tokens) feat['labels'] = labels + list(tokens) + feat['completion_mask'] = completion_mask + [1] * len(tokens) feat['length'] = len(feat['input_ids']) return feat diff --git a/src/twinkle/template/base.py b/src/twinkle/template/base.py index d914b01a9..10912bd7b 100644 --- a/src/twinkle/template/base.py +++ b/src/twinkle/template/base.py @@ -23,6 +23,18 @@ VideoInput = Union[str, List['Image.Image'], 'torch.Tensor'] AudioInput = Union[str, np.ndarray, 'torch.Tensor'] +# Fields that are one entry per token and must be sliced with ``input_ids``. +# ``mm_token_type_ids`` is excluded: it may carry a leading batch dim and is +# sliced on its last axis instead. +_SEQUENCE_ALIGNED_FIELDS = ('labels', 'completion_mask') + +# What an appended turn is to a trainer: the policy's own completion (scored, and +# a log-prob exists for each of its tokens), someone else's completion offered +# for imitation (scored, no log-prob -- usable by SFT but not by RL), or history +# that no loss may touch. There is deliberately no fourth role: a log-prob is +# only ever needed for a token that is also scored. +_APPEND_ROLES = ('completion', 'demonstration', 'context') + @remote_class() class Template: @@ -197,19 +209,48 @@ def _invoke_post_pipeline(self, input_features: List[InputFeature]) -> List[Inpu current = next_batch return current - def concat_input_feature(self, prompt_input_feature: InputFeature, new_tokens: List[int]) -> InputFeature: + def concat_input_feature(self, + prompt_input_feature: InputFeature, + new_tokens: List[int], + *, + appended_as: Literal['completion', 'demonstration', 'context'] = 'completion', + tool_calls: Optional[List[Dict[str, Any]]] = None) -> InputFeature: + """Append one generated turn to an already-encoded prefix. + + Args: + appended_as: what the turn is to a trainer, which decides ``labels`` + and ``completion_mask`` together: + + * ``'completion'`` -- the sampled policy's own output. Scored, and + a log-prob exists for every token. + * ``'demonstration'`` -- written by someone else (a stronger model, + a human) and offered for imitation. Scored, but carries no + log-prob, so RL losses skip it while SFT trains on it. + * ``'context'`` -- history that later turns must see and no loss + may touch. + tool_calls: calls to attach to the appended message, for generators that + return them as structured fields (any OpenAI-compatible API does) + rather than as markup inside the text, which is all + ``parse_tool_call`` can read. + """ import copy import torch assert self.truncation_strategy != 'split', 'concat_input_feature does not support `truncation_strategy=split`' + if appended_as not in _APPEND_ROLES: + raise ValueError(f'appended_as must be one of {_APPEND_ROLES}, got {appended_as!r}') result = copy.deepcopy(prompt_input_feature) prompt_ids = result['input_ids'] labels = list(result.get('labels', [])) input_ids = list(prompt_ids) + new_tokens labels = labels[-1:] + labels[:-1] # roll to input order - labels = labels + new_tokens + completion_mask = self._prefix_completion_mask(result, labels) + scored = appended_as != 'context' + labels = labels + (new_tokens if scored else [-100] * len(new_tokens)) + completion_mask = completion_mask + [int(appended_as == 'completion')] * len(new_tokens) # We don't need to roll back, self._invoke_post_pipeline will do this. result['input_ids'] = input_ids result['labels'] = labels + result['completion_mask'] = completion_mask if 'mm_token_type_ids' in result: mm_token_type_ids = result['mm_token_type_ids'] if not isinstance(mm_token_type_ids, torch.Tensor): @@ -227,8 +268,14 @@ def concat_input_feature(self, prompt_input_feature: InputFeature, new_tokens: L messages: List[Message] = result.get('messages') if messages is not None: response_text = self.tokenizer.decode(new_tokens, skip_special_tokens=True) - parsed = self.parse_tool_call(response_text) or [] - content_text = (self.clean_tool_call(response_text) if parsed else response_text) + if tool_calls is None: + parsed = self.parse_tool_call(response_text) or [] + content_text = (self.clean_tool_call(response_text) if parsed else response_text) + else: + # Structured calls arrived beside the text, so the text carries no + # markup to strip. + parsed = list(tool_calls) + content_text = response_text asst_msg = Message(role='assistant', content=content_text) if parsed: asst_msg['tool_calls'] = parsed @@ -236,6 +283,26 @@ def concat_input_feature(self, prompt_input_feature: InputFeature, new_tokens: L result['messages'] = messages return result + @staticmethod + def _prefix_completion_mask(feature: InputFeature, labels: List[int]) -> List[int]: + """The prefix's ``completion_mask``, in input order, materialised if absent. + + A feature encoded before this field existed records no provenance, and for + those the trainable positions *were* exactly the policy's own -- deriving the + mask from ``labels`` therefore leaves old and new trajectories equivalent. + """ + mask = feature.get('completion_mask') + if mask is None: + mask = [0 if label == -100 else 1 for label in labels] + else: + mask = list(mask) + mask = mask[-1:] + mask[:-1] # roll to input order, exactly as labels + expected = len(feature['input_ids']) + if len(mask) != expected: + raise ValueError(f'prefix completion_mask has {len(mask)} entries for {expected} ' + f'input_ids; appending would misalign every position after it.') + return mask + def _add_default_system(self, trajectory: Trajectory) -> List[Trajectory]: if self.use_chat_template and self.default_system: if trajectory['messages'][0]['role'] == 'user': @@ -274,27 +341,25 @@ def _extract_reasoning_content(messages: list[Message]) -> List[Message]: return [trajectory] def _truncate_feature(self, feature: InputFeature, strategy: str) -> InputFeature: - """Truncate input_ids and labels in a single InputFeature.""" + """Truncate the sequence-aligned fields of a single InputFeature.""" length = len(feature['input_ids']) if length <= self.max_length: return feature if strategy == 'raise': raise ValueError(f'Input length {length} exceeds max_length {self.max_length}') - result = dict(feature) if strategy == 'left': - result['input_ids'] = result['input_ids'][-self.max_length:] - if 'labels' in result: - result['labels'] = result['labels'][-self.max_length:] - if 'mm_token_type_ids' in result: - result['mm_token_type_ids'] = result['mm_token_type_ids'][..., -self.max_length:] + keep = slice(-self.max_length, None) elif strategy == 'right': - result['input_ids'] = result['input_ids'][:self.max_length] - if 'labels' in result: - result['labels'] = result['labels'][:self.max_length] - if 'mm_token_type_ids' in result: - result['mm_token_type_ids'] = result['mm_token_type_ids'][..., :self.max_length] + keep = slice(None, self.max_length) else: raise ValueError(f'Unsupported truncation_strategy={strategy!r}.') + result = dict(feature) + result['input_ids'] = result['input_ids'][keep] + for key in _SEQUENCE_ALIGNED_FIELDS: + if key in result: + result[key] = result[key][keep] + if 'mm_token_type_ids' in result: + result['mm_token_type_ids'] = result['mm_token_type_ids'][..., keep] return InputFeature(**result) def set_mm_position_ids(self, input_feature: InputFeature): @@ -321,8 +386,9 @@ def _check_max_length(self, input_feature: InputFeature) -> List[InputFeature]: end = min(start + self.max_length, len(input_feature['input_ids'])) feat = dict(input_feature) feat['input_ids'] = feat['input_ids'][start:end] - if 'labels' in feat: - feat['labels'] = feat['labels'][start:end] + for key in _SEQUENCE_ALIGNED_FIELDS: + if key in feat: + feat[key] = feat[key][start:end] if 'mm_token_type_ids' in feat: feat['mm_token_type_ids'] = feat['mm_token_type_ids'][..., start:end] results.append(InputFeature(**feat)) @@ -350,6 +416,10 @@ def _roll_labels(self, input_feature: InputFeature) -> List[InputFeature]: if 'input_ids' not in input_feature: return [input_feature] input_feature['labels'] = np.roll(input_feature['labels'], -1, axis=-1) + if 'completion_mask' in input_feature: + # The mask answers "is there a log-prob for this position's target", so it + # lives on the labels' index space and has to follow the same roll. + input_feature['completion_mask'] = np.roll(input_feature['completion_mask'], -1, axis=-1) return [input_feature] def _process_mm_messages(self, messages: List, images: List, videos: List, audios: List) -> List: @@ -524,6 +594,34 @@ def _build_standard_messages(self, trajectory: Trajectory) -> List[Trajectory]: message['content'] = c[0]['text'] if c else '' return [trajectory] + @staticmethod + def decode_tool_calls(message: Dict[str, Any]) -> Dict[str, Any]: + """Return ``message`` with ``tool_calls`` in the shape a chat template renders. + + OpenAI-shaped calls carry ``function.arguments`` as a JSON string, and an + Arrow round-trip can turn the whole list into one; templates index them as + objects. Arguments that will not parse become ``{}`` rather than reaching + Jinja as a string it would render verbatim. The message is returned + untouched when it carries no calls. + """ + tool_calls = message.get('tool_calls') + if isinstance(tool_calls, str): + tool_calls = json.loads(tool_calls) if tool_calls else [] + elif not tool_calls: + return message + decoded = [] + for tool_call in tool_calls: + fn = tool_call['function'] + args = fn['arguments'] + if isinstance(args, dict): + value = args + elif isinstance(args, str): + value = json.loads(args) if args.strip() else {} + else: + value = {} + decoded.append({**tool_call, 'function': {**fn, 'arguments': value}}) + return {**message, 'tool_calls': decoded} + def _apply_chat_template(self, trajectory: Trajectory, add_generation_prompt: bool = False, **kwargs): messages = [dict(message) for message in trajectory['messages']] # Arrow serialization may pad content blocks with null keys (e.g. 'image': None @@ -536,25 +634,7 @@ def _apply_chat_template(self, trajectory: Trajectory, add_generation_prompt: bo k: v for k, v in b.items() if v is not None } for b in msg['content'] if isinstance(b, dict)] - for msg in messages: - tcs = msg.get('tool_calls') - if isinstance(tcs, str): - tcs = json.loads(tcs) if tcs else [] - msg['tool_calls'] = tcs - if not tcs: - continue - new_tcs = [] - for tc in tcs: - fn = tc['function'] - args = fn['arguments'] - if isinstance(args, dict): - decoded = args - elif isinstance(args, str): - decoded = json.loads(args) if args.strip() else {} - else: - decoded = {} - new_tcs.append({**tc, 'function': {**fn, 'arguments': decoded}}) - msg['tool_calls'] = new_tcs + messages = [self.decode_tool_calls(msg) for msg in messages] # ``tool_calls`` / ``tools`` are already OpenAI-shaped (see # :mod:`twinkle.data_format.message`); pass them through verbatim. tools = list(trajectory.get('tools') or []) diff --git a/src/twinkle_agentic/async_rl/data_plane.py b/src/twinkle_agentic/async_rl/data_plane.py index 641947366..da9015f0c 100644 --- a/src/twinkle_agentic/async_rl/data_plane.py +++ b/src/twinkle_agentic/async_rl/data_plane.py @@ -71,10 +71,20 @@ def _require_rollout_logprobs(sample: dict[str, Any], *, sample_key: str) -> lis values.append(float(value)) labels = sample.get('labels') if labels is not None: - trainable_tokens = sum(1 for label in labels if label != -100) - if len(values) != trainable_tokens: - raise ValueError(f'rollout sample {sample_key!r} logprobs length must match trainable labels: ' - f'{len(values)} != {trainable_tokens}') + # Only policy-generated tokens carry a sampling log-prob. A turn written by + # an API or a human is trainable yet has none, and is marked + # completion_mask=0 -- the same basis GRPOLoss restricts itself to. + completion_mask = sample.get('completion_mask') + if completion_mask is None: + expected = sum(1 for label in labels if label != -100) + elif len(completion_mask) != len(labels): + raise ValueError(f'rollout sample {sample_key!r} completion_mask length must match labels: ' + f'{len(completion_mask)} != {len(labels)}') + else: + expected = sum(1 for label, flag in zip(labels, completion_mask) if label != -100 and flag) + if len(values) != expected: + raise ValueError(f'rollout sample {sample_key!r} logprobs length must match policy-generated tokens: ' + f'{len(values)} != {expected}') return values diff --git a/src/twinkle_agentic/challenger/base.py b/src/twinkle_agentic/challenger/base.py index 3dcaae674..69ab74673 100644 --- a/src/twinkle_agentic/challenger/base.py +++ b/src/twinkle_agentic/challenger/base.py @@ -11,8 +11,7 @@ rollouts in :mod:`twinkle_agentic.rollout` have that signature already, so a challenger can explore *with tools* -- running code, reading files -- while it invents, over a local sampler or over an HTTP endpoint alike. - :func:`twinkle_agentic.rollout.build_rollout` picks the right one for the - backend at hand. + :class:`twinkle_agentic.rollout.MultiTurnRollout` accepts either backend. * **what counts as a keeper** -- subclasses decide, in :meth:`Challenger.build`. * **how hard is hard enough** -- optional. Ask for ``solver_rollouts`` attempts per candidate and only tasks the model solves *sometimes* are kept: a task every @@ -44,8 +43,7 @@ __all__ = ['Challenger', 'Explorer', 'KeywordPrompts', 'PromptSet'] # A batch of trajectories in, the same trajectories with the model's reply -# appended out. Both MultiTurnRollout and APIMultiTurnRollout satisfy this -# as-is; build_rollout() returns whichever fits the backend. Both also accept a +# appended out. MultiTurnRollout accepts either backend and also accepts a # per-call ``sampling_params=`` keyword, which is how the difficulty stage asks # for its own temperature and length budget without a second explorer. Explorer = Callable[[List[Trajectory]], List[Trajectory]] @@ -141,9 +139,9 @@ class Challenger(ABC): Args: explorer: takes a batch of trajectories and returns them with the - model's reply appended -- a rollout from - :func:`twinkle_agentic.rollout.build_rollout`, over a local sampler - or over an API endpoint. + model's reply appended -- typically a + :class:`twinkle_agentic.rollout.MultiTurnRollout` over a local + sampler or an API endpoint. system: system prompt handed to the model. It carries the output contract, which is why ``build`` -- the code that reads that output back -- lives in the same subclass. diff --git a/src/twinkle_agentic/challenger/new/__init__.py b/src/twinkle_agentic/challenger/new/__init__.py new file mode 100644 index 000000000..0719b04e5 --- /dev/null +++ b/src/twinkle_agentic/challenger/new/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .agentic import AgenticChallenger, parse_problem_statement +from .base import Challenger +from .keyword import KEYWORD_MAX_LEN, KeywordGenerator + +__all__ = [ + 'AgenticChallenger', + 'Challenger', + 'KEYWORD_MAX_LEN', + 'KeywordGenerator', + 'parse_problem_statement', +] diff --git a/src/twinkle_agentic/challenger/new/agentic.py b/src/twinkle_agentic/challenger/new/agentic.py new file mode 100644 index 000000000..ca41486a7 --- /dev/null +++ b/src/twinkle_agentic/challenger/new/agentic.py @@ -0,0 +1,511 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Agentic challenger: act in a sandbox, verify the result, then describe it.""" +import math +import random +import re +import uuid +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple + +from twinkle.data_format import SamplingParams, Trajectory, attach_user_data, user_data_get +from twinkle.data_format.sampling import SampledSequence, SampleResponse +from twinkle.utils import get_logger +from twinkle_agentic.envs import Env +from twinkle_agentic.protocol.base import API +from twinkle_agentic.rollout import APISampler, MultiTurnRollout +from twinkle_agentic.summarizer import Summarizer +from twinkle_agentic.utils.code_utils import parse_fenced_code, strip_reasoning +from twinkle_agentic.utils.message_utils import assistant_text, msg_content_text, normalize_tool_calls +from .base import Challenger, _parallel +from .keyword import KeywordGenerator +from .recorder import RolloutRecorder + +__all__ = ['AgenticChallenger', 'parse_problem_statement'] + +logger = get_logger() + +_FENCED_BLOCK_RE = re.compile(r'```[^\r\n]*\r?\n(.*?)```', re.S) + + +def parse_problem_statement(text: str) -> Optional[str]: + """Return the statement after removing reasoning and one outer fence.""" + body = strip_reasoning(text).strip() + whole = _FENCED_BLOCK_RE.fullmatch(body) + if whole: + body = whole.group(1).strip() + return body or None + + +def _sample_one(sampler: Any, input_feature: Dict[str, Any], sampling_params: Optional[SamplingParams], + adapter_kwargs: Dict[str, Any]) -> SampledSequence: + responses = sampler.sample([input_feature], sampling_params=sampling_params, **adapter_kwargs) + if not isinstance(responses, list): + raise TypeError(f'expected List[SampleResponse] from sampler.sample, got ' + f'{type(responses).__name__}') + if len(responses) != 1: + raise RuntimeError(f'sampler returned {len(responses)} responses for a single request; ' + 'expected exactly one') + response = responses[0] + if not isinstance(response, SampleResponse): + raise TypeError(f'expected SampleResponse from sampler.sample, got ' + f'{type(response).__name__}') + if len(response.sequences) != 1: + raise RuntimeError(f'SampleResponse contains {len(response.sequences)} sequences; ' + 'expected exactly one') + sequence = response.sequences[0] + if not isinstance(sequence, SampledSequence): + raise TypeError(f'expected SampledSequence, got {type(sequence).__name__}') + return sequence + + +def _api_followup_response( + sampler: Any, + api: Optional[APISampler], + sampling_params: Optional[SamplingParams], + *, + input_feature: Dict[str, Any], + adapter_kwargs: Dict[str, Any], + followups: int, + **kwargs: Any, +) -> SampledSequence: + """Use the API for appended stages and the primary backend otherwise.""" + if followups: + if api is None: + raise ValueError('use_api=True requires an API backend') + return api(input_feature, sampling_params, **adapter_kwargs) + if sampler is not None: + return _sample_one(sampler, input_feature, sampling_params, adapter_kwargs) + if api is not None: + return api(input_feature, sampling_params, **adapter_kwargs) + raise ValueError('AgenticChallenger has neither a sampler nor an API backend') + + +@dataclass +class _ProposalResult: + trajectory: Trajectory + group_id: str = '' + task: Optional[Trajectory] = None + reason: str = '' + detail: str = '' + outcome: str = '' + n_pass: Optional[int] = None + reward: float = 0.0 + + +class AgenticChallenger(Challenger): + """Invent tool-using tasks by doing, checking, and describing them. + + ``backend`` drives exploration and solver attempts. When ``use_api`` is true, + ``api`` generates only the appended check-script and problem-statement turns; + those turns retain the masking semantics selected by ``api_appended_as`` in + ``rollout_kwargs``. + """ + + _system = ('You invent tasks for another agent to solve. You have a sandbox and ' + 'tools. Work in it first: build something real, then you will be asked ' + 'to verify it and to describe it.') + _from_scratch = ('Choose a task worth doing in this sandbox and do it now, using ' + 'your tools. Do not describe it yet.') + _from_keywords = ('Choose a task around these topics and do it now, using your ' + 'tools. Do not describe it yet.\n\nTopics: {keywords}') + _from_seed = ('Here is an earlier task:\n\n{seed}\n\nDo something in the same ' + 'spirit but different, using your tools now. Do not describe it yet.') + _from_seed_keywords = ('Here is an earlier task:\n\n{seed}\n\nDo something in the same ' + 'spirit but different, may be more complex and interesting and meaningful, ' + 'around these topics, using your tools now. Do not describe it yet.\n\n' + 'Topics: {keywords}') + _check_followup = ('Stop working. This is the workspace you produced:\n\n{final_state}\n\n' + 'Write a {language} script that verifies this end state, as a fenced ' + '{language} code block and nothing else. It must exit with a non-zero status ' + 'if the work was not done. Check what can be read out of the files -- their ' + 'structure and the values inside them. NEVER check a file size in bytes, a ' + 'checksum, or the full source text of a script: correct solutions differ ' + 'there, and such a check only its own author can pass.') + _check_retry_followup = ('Your check script did not pass:\n\n{error}\n\nThe workspace is:\n\n' + '{final_state}\n\nReturn a corrected script as a fenced {language} code block ' + 'and nothing else.') + _check_parse_error = ('Could not read a check script from your reply: it was not a ' + 'fenced {language} code block. Do not wrap it in a tool call and ' + 'do not add prose -- return ONLY a fenced {language} code block.') + _problem_followup = ('Now write the task statement: what someone starting from an empty workspace ' + 'would have to be told to produce what you produced, and nothing about how you ' + 'did it. Name the files to create and quote any input data verbatim. Do not ' + 'reveal values your check script computes. Reply with the statement only.') + + + def __init__( + self, + backend: Any, + *, + api: Optional[Any] = None, + use_api: bool = False, + keyword_generator: Optional[KeywordGenerator] = None, + trajectory_seed: Optional[List[Trajectory]] = None, + summarizer: Optional[Summarizer] = None, + system_prompt: Optional[str] = None, + from_scratch_prompt: Optional[str] = None, + from_keywords_prompt: Optional[str] = None, + from_seed_prompt: Optional[str] = None, + from_seed_keywords_prompt: Optional[str] = None, + check_followup_prompt: Optional[str] = None, + check_retry_followup_prompt: Optional[str] = None, + check_parse_error_prompt: Optional[str] = None, + problem_followup_prompt: Optional[str] = None, + check_retries: int = 1, + problem_max_chars: int = 8192, + check_language: str = 'python', + parse_check_fn: Optional[Callable[[str], Optional[str]]] = None, + pass_rate_target: float = 0.2, + envs: Sequence[Env] = (), + num_challenger_rollouts: int = 8, + num_solver_rollouts: int = 8, + pass_band: Tuple[float, float] = (1.0, 7.0), + pass_rate_width: float = 0.3, + max_empty_rounds: int = 0, + followup_params: Optional[SamplingParams] = None, + checker: Optional[Callable[[Trajectory], bool]] = None, + save_dir: Optional[str] = None, + save_failed_rollouts: bool = True, + **rollout_kwargs: Any, + ): + super().__init__( + envs=envs, + num_challenger_rollouts=num_challenger_rollouts, + num_solver_rollouts=num_solver_rollouts, + pass_band=pass_band, + max_empty_rounds=max_empty_rounds, + ) + if check_retries < 0: + raise ValueError(f'check_retries must be >= 0, got {check_retries}') + if problem_max_chars <= 0: + raise ValueError(f'problem_max_chars must be positive, got {problem_max_chars}') + if not check_language.strip(): + raise ValueError('check_language must not be empty') + if not 0 <= pass_rate_target <= 1: + raise ValueError(f'pass_rate_target must be in [0, 1], got {pass_rate_target}') + if pass_rate_width <= 0: + raise ValueError(f'pass_rate_width must be positive, got {pass_rate_width}') + if use_api and rollout_kwargs.get('response_callback') is not None: + raise ValueError('use_api=True cannot be combined with response_callback') + backend_is_api = isinstance(backend, (API, APISampler)) + if use_api and api is None and not backend_is_api: + raise ValueError('use_api=True requires api= when backend is a sampler') + self.keyword_generator = keyword_generator + self.trajectory_seed = list(trajectory_seed or ()) + self.summarizer = summarizer + self._system = self._system if system_prompt is None else system_prompt + self._from_scratch = self._from_scratch if from_scratch_prompt is None else from_scratch_prompt + self._from_keywords = self._from_keywords if from_keywords_prompt is None else from_keywords_prompt + self._from_seed = self._from_seed if from_seed_prompt is None else from_seed_prompt + self._from_seed_keywords = (self._from_seed_keywords if from_seed_keywords_prompt is None else + from_seed_keywords_prompt) + self._check_followup = self._check_followup if check_followup_prompt is None else check_followup_prompt + self._check_retry_followup = (self._check_retry_followup if check_retry_followup_prompt is None else + check_retry_followup_prompt) + self._check_parse_error = (self._check_parse_error if check_parse_error_prompt is None else + check_parse_error_prompt) + self._problem_followup = (self._problem_followup if problem_followup_prompt is None else + problem_followup_prompt) + self._check_retries = check_retries + self._problem_max_chars = problem_max_chars + self._check_language = check_language.strip().lower() + self._parse_check_fn = parse_check_fn + self._pass_rate_target = pass_rate_target + self._pass_rate_width = pass_rate_width + self.checker = checker + self.followup_params = followup_params + self.rng = random.Random() + self.use_api = use_api + self.save_failed_rollouts = save_failed_rollouts + self._recorder = RolloutRecorder(save_dir) if save_dir else None + self._round_proposals: List[_ProposalResult] = [] + self._backend = backend + self._rollout_kwargs = dict(rollout_kwargs) + if api is not None: + self._rollout_kwargs['api'] = api + if use_api: + self._rollout_kwargs['response_callback'] = _api_followup_response + self._rollout: Optional[MultiTurnRollout] = None + self._tool_schemas = self.env().tools() or None + + def _rollout_instance(self) -> MultiTurnRollout: + if self._rollout is None: + self._rollout = MultiTurnRollout(self._backend, **self._rollout_kwargs) + return self._rollout + + def _tool_manager(self, slot: int) -> Optional[Any]: + env = self.env(slot) + return env.tool_manager() if env.tools() else None + + def _summary(self, trajectory: Trajectory) -> str: + turns: List[str] = [] + for message in trajectory.get('messages') or []: + if not isinstance(message, dict): + continue + role = message.get('role') or '' + if role == 'system': + continue + parts = [msg_content_text(message).strip()] + for call in normalize_tool_calls(message) or (): + fn = call.get('function') or {} + if isinstance(fn, dict) and fn.get('name'): + parts.append(f"calls {fn['name']}({fn.get('arguments') or ''})") + body = '\n'.join(part for part in parts if part) + if body: + turns.append(f'{role}: {body}') + text = '\n'.join(turns) + if not text: + return '' + return self.summarizer(text) if self.summarizer is not None else text + + def _build_challenge_prompt(self) -> Optional[Trajectory]: + keywords: List[str] = [] + if self.keyword_generator is not None: + groups = self.keyword_generator.get_keywords(1) + if not groups: + return None + keywords = groups[0] + seed = '' + if self.trajectory_seed: + seed = self._summary(self.rng.choice(self.trajectory_seed)) + block = ', '.join(keywords) + if seed and keywords: + user = self._from_seed_keywords.format(seed=seed, keywords=block) + elif seed: + user = self._from_seed.format(seed=seed) + elif keywords: + user = self._from_keywords.format(keywords=block) + else: + user = self._from_scratch + prompt: Trajectory = { + 'messages': [ + { + 'role': 'system', + 'content': self._system + }, + { + 'role': 'user', + 'content': user + }, + ], + } + if self._tool_schemas: + prompt['tools'] = self._tool_schemas + return attach_user_data(prompt, keywords=keywords, seeded=bool(seed)) + + def _explore(self, prompt: Trajectory) -> List[Trajectory]: + group_id = uuid.uuid4().hex + proposals: List[_ProposalResult] = [] + remaining = self.num_challenger_rollouts + while remaining > 0: + wave = min(self.n_slots, remaining) + proposals.extend(_parallel(lambda slot: self._run_episode(prompt, slot), wave)) + remaining -= wave + for proposal in proposals: + proposal.group_id = group_id + self._round_proposals = proposals + return [proposal.task for proposal in proposals if proposal.task is not None] + + def _run_episode(self, prompt: Trajectory, slot: int) -> _ProposalResult: + self.env(slot).clear() + state: Dict[str, Any] = {'slot': slot} + kwargs: Dict[str, Any] = { + 'followup_fn': lambda trajectory, n_before: self._followup(state, trajectory, n_before), + } + manager = self._tool_manager(slot) + if manager is not None: + kwargs['tool_manager'] = manager + explored = self._rollout_instance()([prompt], **kwargs) + if not explored: + self._reject(state, 'rollout_no_output') + return _ProposalResult(dict(prompt), reason='rollout_no_output') + trajectory = explored[0] + task = self._build_query(state, trajectory) + reason, detail = state.get('reject', ('', '')) + return _ProposalResult(trajectory, task=task, reason=reason, detail=detail) + + def _followup(self, state: Dict[str, Any], trajectory: Trajectory, + n_before: int) -> Optional[Tuple[str, Optional[SamplingParams]]]: + if state.get('checked'): + return None + reply = None if n_before == 0 else assistant_text(trajectory) + followup = self._build_test_case(state, reply) + if followup is None: + return None + return followup, self.followup_params + + def _build_test_case(self, state: Dict[str, Any], reply: Optional[str]) -> Optional[str]: + slot = state['slot'] + if reply is None: + snapshot, error = self.env(slot).snapshot() + state['snapshot'] = snapshot + if not snapshot.strip(): + state['reject'] = ('snapshot_unavailable' if error else 'empty_workspace', error) + return None + return self._check_followup.format(final_state=snapshot, language=self._check_language) + + attempt = state.get('check_attempts', 0) + 1 + state['check_attempts'] = attempt + script = (self._parse_check_fn(reply) if self._parse_check_fn is not None else + parse_fenced_code(reply, language_tags=None)) + if script is None: + if attempt <= self._check_retries: + return self._check_retry_followup.format( + error=self._check_parse_error.format(language=self._check_language), + final_state=state.get('snapshot', ''), + language=self._check_language, + ) + state['reject'] = ('check_parse_fail', reply) + return None + state['script'] = script + exit_code, output = self.env(slot).run_script(script, interpreter=self._check_language) + if exit_code == 0: + state['checked'] = True + return self._problem_followup + after = self.env(slot).snapshot()[0] + state.setdefault('attempts', []).append(f'--- attempt {attempt}: exit {exit_code} ---\n{output}\n' + f'--- check script ---\n{script}') + if attempt <= self._check_retries: + return self._check_retry_followup.format( + error=output, + final_state=after or state.get('snapshot', ''), + language=self._check_language, + ) + state['reject'] = ('check_run_fail', '\n'.join(state['attempts'])) + return None + + def _build_query(self, state: Dict[str, Any], explored: Trajectory) -> Optional[Trajectory]: + if state.get('reject'): + return self._reject(state, *state['reject']) + if not state.get('checked'): + return self._reject( + state, + 'episode_cut_short', + f"stop_reason={explored.get('stop_reason')} " + f"truncated={bool(explored.get('truncated'))} " + f"turns={explored.get('turns')}", + ) + statement = parse_problem_statement(assistant_text(explored)) + if statement is None: + return self._reject(state, 'problem_parse_fail') + if len(statement) > self._problem_max_chars: + return self._reject(state, 'too_long', f'{len(statement)} chars') + task: Trajectory = attach_user_data( + {'messages': [{ + 'role': 'user', + 'content': statement + }]}, + check_script=state['script'], + keywords=user_data_get(explored.get('user_data'), 'keywords', []), + seeded=user_data_get(explored.get('user_data'), 'seeded', False), + ) + if self.checker is not None and not self.checker(task): + return self._reject(state, 'rejected_by_checker') + return task + + def _reject(self, state: Dict[str, Any], reason: str, detail: str = '') -> Optional[Trajectory]: + state['reject'] = (reason, detail) + logger.info(f'[{type(self).__name__}] rejected: {reason}' + f"{f' -- {detail[:400]}' if detail else ''}") + return None + + def _solver_prompt(self, task: Trajectory) -> Trajectory: + prompt: Trajectory = {'messages': [dict(message) for message in task.get('messages') or []]} + if self._tool_schemas: + prompt['tools'] = self._tool_schemas + return prompt + + def _judge(self, task: Trajectory, slot: int) -> bool: + script = user_data_get(task.get('user_data'), 'check_script', '') + if not script: + return False + return self.env(slot).run_script(script, interpreter=self._check_language)[0] == 0 + + def challenger_reward(self, n_pass: Optional[int]) -> float: + """Reward tasks near the target solver pass rate; unmeasured failures score zero.""" + if n_pass is None or not self.num_solver_rollouts or n_pass <= 0: + return 0.0 + gap = n_pass / self.num_solver_rollouts - self._pass_rate_target + variance = 2.0 * self._pass_rate_width**2 + return math.exp(-(gap * gap) / variance) + + def _record_proposals(self) -> None: + proposals, self._round_proposals = self._round_proposals, [] + if self._recorder is None: + return + for index, proposal in enumerate(proposals): + if proposal.task is None and not self.save_failed_rollouts: + continue + trajectory = dict(proposal.trajectory) + trajectory['rewards'] = proposal.reward + task_data = proposal.task.get('user_data') if proposal.task is not None else None + statement = '' + if proposal.task is not None: + statement = next((message.get('content', '') for message in proposal.task.get('messages') or [] + if isinstance(message, dict) and message.get('role') == 'user'), '') + self._recorder.write( + trajectory, + side='propose', + group_id=proposal.group_id, + proposal_index=index, + outcome=proposal.outcome or ('rejected' if proposal.reason else 'kept'), + reason=proposal.reason, + detail=proposal.detail, + reward=proposal.reward, + n_pass=proposal.n_pass, + n_rollouts=(self.num_solver_rollouts if proposal.n_pass is not None else None), + pass_rate=(proposal.n_pass / self.num_solver_rollouts + if proposal.n_pass is not None and self.num_solver_rollouts else None), + statement=statement, + check_script=user_data_get(task_data, 'check_script', ''), + keywords=user_data_get(proposal.trajectory.get('user_data'), 'keywords', []), + seeded=user_data_get(proposal.trajectory.get('user_data'), 'seeded', False), + ) + + def _filter_difficulty(self, tasks: List[Trajectory]) -> List[Trajectory]: + successful = [proposal for proposal in self._round_proposals if proposal.task is not None] + if len(successful) != len(tasks): + raise RuntimeError('proposal/task alignment failed before difficulty filtering') + if not tasks or not self.num_solver_rollouts: + for proposal in successful: + proposal.outcome = 'kept' + self._record_proposals() + return tasks + + passes = [0] * len(tasks) + plan = [i for i in range(len(tasks)) for _ in range(self.num_solver_rollouts)] + rollout = self._rollout_instance() + for start in range(0, len(plan), self.n_slots): + wave = plan[start:start + self.n_slots] + _parallel(lambda slot: self.env(slot).clear(), len(wave)) + prompts = [self._solver_prompt(tasks[i]) for i in wave] + kwargs: Dict[str, Any] = {} + managers = [self._tool_manager(slot) for slot in range(len(wave))] + if any(manager is not None for manager in managers): + kwargs['tool_manager'] = managers + attempts = rollout(prompts, **kwargs) + if len(attempts) != len(prompts): + raise RuntimeError(f'rollout returned {len(attempts)} attempts for ' + f'{len(prompts)} prompts; expected one per prompt') + verdicts = _parallel(lambda slot: self._judge(tasks[wave[slot]], slot), len(wave)) + for slot, passed in enumerate(verdicts): + if passed: + passes[wave[slot]] += 1 + + low, high = self.pass_band + measured = [ + attach_user_data(task, n_pass=n_pass, n_rollouts=self.num_solver_rollouts) + for task, n_pass in zip(tasks, passes) + ] + kept: List[Trajectory] = [] + for proposal, task, n_pass in zip(successful, measured, passes): + proposal.task = task + proposal.n_pass = n_pass + proposal.reward = self.challenger_reward(n_pass) + if low <= n_pass <= high: + proposal.outcome = 'kept' + kept.append(task) + else: + proposal.outcome = 'outside_band' + self._record_proposals() + return kept diff --git a/src/twinkle_agentic/challenger/new/base.py b/src/twinkle_agentic/challenger/new/base.py new file mode 100644 index 000000000..1c819533b --- /dev/null +++ b/src/twinkle_agentic/challenger/new/base.py @@ -0,0 +1,132 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Reusable lifecycle for task challengers.""" +from abc import ABC, abstractmethod +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Callable, Iterator, List, Optional, Sequence, Tuple + +from twinkle.data_format import Trajectory +from twinkle.utils import get_logger +from twinkle_agentic.envs import Env + +logger = get_logger() + +__all__ = ['Challenger'] + + +def _parallel(fn: Callable[[int], Any], count: int) -> List[Any]: + """Run ``fn`` over ``range(count)`` concurrently, preserving order.""" + if count <= 1: + return [fn(i) for i in range(count)] + out: List[Any] = [None] * count + with ThreadPoolExecutor(max_workers=count) as pool: + futures = {pool.submit(fn, i): i for i in range(count)} + for future, i in futures.items(): + out[i] = future.result() + return out + + +class Challenger(ABC): + """Common batching and environment lifecycle for task challengers. + + Subclasses define how a round builds its prompt, explores it, and measures + candidate difficulty. One environment is owned by one concurrent job for the + complete lifetime of that job. + """ + + def __init__( + self, + *, + envs: Sequence[Env], + num_challenger_rollouts: int = 8, + num_solver_rollouts: int = 8, + pass_band: Tuple[float, float] = (1.0, 7.0), + max_empty_rounds: int = 0, + ): + if not envs: + raise ValueError('envs is empty: a challenger needs a workspace to act in and grade') + if num_challenger_rollouts < 1: + raise ValueError(f'num_challenger_rollouts must be >= 1, got ' + f'{num_challenger_rollouts}') + if num_solver_rollouts < 0: + raise ValueError(f'num_solver_rollouts must be >= 0, got {num_solver_rollouts}') + if max_empty_rounds < 0: + raise ValueError(f'max_empty_rounds must be >= 0, got {max_empty_rounds}') + if num_solver_rollouts: + if len(pass_band) != 2: + raise ValueError(f'pass_band is (low, high) in attempt counts, got {pass_band}') + low, high = pass_band + if not 0 <= low <= high <= num_solver_rollouts: + raise ValueError(f'pass_band must satisfy 0 <= low <= high <= num_solver_rollouts, got ' + f'{pass_band} against num_solver_rollouts={num_solver_rollouts}') + self.envs = list(envs) + self.num_challenger_rollouts = num_challenger_rollouts + self.num_solver_rollouts = num_solver_rollouts + self.pass_band = pass_band + self.max_empty_rounds = max_empty_rounds + self.n_proposed = 0 + self.n_kept = 0 + + @property + def n_slots(self) -> int: + """How many jobs may run at once: one per environment.""" + return len(self.envs) + + def env(self, slot: int = 0) -> Env: + """Return the current environment for ``slot``.""" + return self.envs[slot] + + @abstractmethod + def _build_challenge_prompt(self) -> Optional[Trajectory]: + """Build one round's shared prompt, or return None when exhausted.""" + + @abstractmethod + def _explore(self, prompt: Trajectory) -> List[Trajectory]: + """Generate and validate candidates from one shared prompt.""" + + @abstractmethod + def _filter_difficulty(self, tasks: List[Trajectory]) -> List[Trajectory]: + """Measure candidate difficulty and return the accepted tasks.""" + + def __call__(self, batch_size: int, total: Optional[int] = None) -> Iterator[List[Trajectory]]: + """Yield finished tasks in batches.""" + if batch_size <= 0: + raise ValueError(f'batch_size must be positive, got {batch_size}') + pending: List[Trajectory] = [] + produced = 0 + empty_rounds = 0 + while total is None or produced < total: + want = batch_size if total is None else min(batch_size, total - produced) + while len(pending) < want: + kept = self._round() + if kept is None: + if pending: + yield pending + return + if kept: + empty_rounds = 0 + pending.extend(kept) + continue + empty_rounds += 1 + if self.max_empty_rounds and empty_rounds >= self.max_empty_rounds: + logger.warning(f'[{type(self).__name__}] stopped after {empty_rounds} ' + 'consecutive rounds without a usable task') + if pending: + yield pending + return + yield pending[:want] + produced += want + pending = pending[want:] + + def _round(self) -> Optional[List[Trajectory]]: + """Run one proposal group; None means the source is exhausted.""" + prompt = self._build_challenge_prompt() + if prompt is None: + return None + verified = self._explore(prompt) + kept = self._filter_difficulty(verified) + self.n_proposed += self.num_challenger_rollouts + self.n_kept += len(kept) + logger.info(f'[{type(self).__name__}] {self.num_challenger_rollouts} episodes, ' + f'{len(verified)} verified, {len(kept)} in band ' + f'(cumulative {self.n_kept}/{self.n_proposed})') + return kept diff --git a/src/twinkle_agentic/challenger/new/keyword.py b/src/twinkle_agentic/challenger/new/keyword.py index 0d2c66bad..09dee2bad 100644 --- a/src/twinkle_agentic/challenger/new/keyword.py +++ b/src/twinkle_agentic/challenger/new/keyword.py @@ -1,12 +1,290 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Keywords per direction: generate, de-duplicate, store, read back. +A keyword is a *topic* to build a task around, not a task statement, which is +why over-length replies are dropped rather than stored. +""" +import json +import os +from typing import Any, Dict, List, Optional, Sequence, Set, Tuple -from typing import Optional +from twinkle.data_format import SamplingParams, Trajectory +from twinkle.utils import get_logger +from twinkle_agentic.rollout import MultiTurnRollout +from twinkle_agentic.utils.code_utils import strip_reasoning +from twinkle_agentic.utils.message_utils import assistant_text + +logger = get_logger() + +__all__ = ['KEYWORD_MAX_LEN', 'KeywordGenerator'] + +KEYWORD_MAX_LEN = 60 class KeywordGenerator: + """Keyword combinations drawn from one list per direction. + + ``keywords_group_size`` of the directions are active at a time and one draw + takes a keyword from each. What a draw spends is the *combination*, not the + keywords: a group only has to differ from every group already handed out, so + three directions holding ``num_keywords`` each are worth their product in + draws rather than just ``num_keywords``. A direction that has produced + ``num_keywords`` is retired and the next unused one takes its slot, which is + why more directions than a group needs is the normal case. De-duplication of + the keywords themselves is flat, so a keyword one direction produced is never + handed to another. + + Args: + query: what the keywords have to satisfy -- one entry per direction. Must + be at least ``keywords_group_size`` of them. + backend: an API client or a sampler; driven through ``MultiTurnRollout``. + path: JSONL cache. Empty means in-memory only. + num_keywords: a direction's budget; past it, it is retired. + keywords_group_size: how many keywords one draw combines. + system_prompt: overrides the built-in one. + recycle: once every direction is spent, hand out the same combinations + again instead of returning None. + rollout_kwargs: passed to ``MultiTurnRollout``. ``template`` is required; + API request options belong in ``api_kwargs``. + """ + + # How many known keywords the 'do not repeat these' line may quote. A cap in + # both directions: too few and a second round says the same things again, too + # many and the model runs out of room to obey. + _avoid_max = 100 + _avoid_lead = '\nDo NOT repeat any of these: ' + + # A default prompt to use to generate the keywords + _default_prompt = ( + 'You brainstorm topics. Reply with a JSON array of short noun phrases ' + f'(at most {KEYWORD_MAX_LEN} characters each) and nothing else. ' + 'Each phrase names a subject to build a task around, never a task statement.') + + _user_prompt = 'Give {k} distinct topics that satisfy:\n{query}' + + def __init__( + self, + query: Sequence[str], + backend: Any, + path: str, + *, + num_keywords: int = 64, + keywords_group_size: int = 3, + system_prompt: Optional[str] = None, + sampling_params: Optional[SamplingParams] = None, + recycle: bool = False, + **rollout_kwargs: Any, + ): + self.query = list(query) + if keywords_group_size < 1: + raise ValueError(f'keywords_group_size must be >= 1, got {keywords_group_size}') + if len(self.query) < keywords_group_size: + raise ValueError(f'{len(self.query)} query(ies) cannot fill a group of ' + f'{keywords_group_size}') + self.path = path + self.num_keywords = num_keywords + self.keywords_group_size = keywords_group_size + self.recycle = recycle + self.system_prompt = system_prompt or self._default_prompt + # Built on the first call rather than here, so a fully cached run needs no backend. + self._backend = backend + self._rollout_kwargs = dict(rollout_kwargs, sampling_params=sampling_params, max_turns=1) + self._rollout: Optional[Any] = None + self._cached_keywords: Dict[str, List[str]] = self.load_keywords() + # Flat: one keyword belongs to one direction, whichever produced it first. + self._seen = {kw.lower() for kws in self._cached_keywords.values() for kw in kws} + # The active slots, the next direction to promote, which slot retires + # next, and the mixed-radix counter walking the active buckets. Drawn + # combinations are remembered because a bucket growing mid-run shifts the + # counter's order and would otherwise let it land on an old group again. + self._active = list(self.query[:keywords_group_size]) + self._next_query = keywords_group_size + self._retire_slot = 0 + self._odometer = [0] * keywords_group_size + self._drawn: Set[Tuple[str, ...]] = set() + self._recycled = False + + # ------------------------------------------------------------------- get + + def get_keywords(self, num_groups: int = 1) -> Optional[List[List[str]]]: + """Up to ``num_groups`` combinations of ``keywords_group_size`` keywords each. + + Fewer than asked for when the directions run dry mid-way -- a partial + batch is still usable -- and None when not even one group could be + filled, which is the caller's signal to stop. + """ + if num_groups < 1: + raise ValueError(f'num_groups must be >= 1, got {num_groups}') + groups: List[List[str]] = [] + for _ in range(num_groups): + group = self._draw_group() + if group is None: + break + groups.append(group) + return groups or None + + def _draw_group(self) -> Optional[List[str]]: + """The next combination nobody has been handed, widening the pool to find one.""" + while True: + group = self._step() + if group is not None: + return group + if not self._grow_or_retire(): + return None + + def _step(self) -> Optional[List[str]]: + """One sweep of the odometer for an undrawn combination. None once there is none.""" + buckets = [self._cached_keywords.get(q, []) for q in self._active] + total = 1 + for bucket in buckets: + total *= len(bucket) + for _ in range(total): + combo = tuple(bucket[i] for bucket, i in zip(buckets, self._odometer)) + self._advance(buckets) + if combo not in self._drawn: + self._drawn.add(combo) + self._recycled = False + return list(combo) + return None + + def _advance(self, buckets: Sequence[Sequence[str]]) -> None: + """Odometer +1, last slot first, carrying into the one before it.""" + for slot in reversed(range(len(buckets))): + self._odometer[slot] += 1 + if self._odometer[slot] < len(buckets[slot]): + return + self._odometer[slot] = 0 + + def _grow_or_retire(self) -> bool: + """Widen the combination space: more keywords, else a new direction. + + False once neither is left. Growing comes first because it multiplies what + the current slots are worth, while retiring gives up on a direction. + """ + short = [q for q in self._active + if len(self._cached_keywords.get(q, [])) < self.num_keywords] + # A round that adds nothing means the model has run out of distinct ideas + # for these directions, so asking again would only spend calls. + if short and self.generate(short): + return True + # Round-robin, so the surplus queries are spent evenly across the slots. + slot = self._retire_slot + self._retire_slot = (slot + 1) % self.keywords_group_size + return self._retire(slot) + + def _retire(self, slot: int) -> bool: + """Promote the next unused direction into ``slot``. False once nothing is left to serve.""" + if self._next_query < len(self.query): + self._active[slot] = self.query[self._next_query] + self._next_query += 1 + self._odometer = [0] * self.keywords_group_size + return True + # Recycling twice without a group in between would spin forever, so it is + # allowed only once per exhaustion -- ``_step`` clears the flag on success. + if self._recycled or not self.recycle or not any(self._cached_keywords.values()): + logger.warning(f'all {len(self.query)} query(ies) are spent; ' + f'pass recycle=True to hand out the same groups again') + return False + self._drawn.clear() + self._active = list(self.query[:self.keywords_group_size]) + self._next_query = self.keywords_group_size + self._odometer = [0] * self.keywords_group_size + self._recycled = True + logger.info(f'[{type(self).__name__}] every query spent -> recycling the combinations') + return True + + # -------------------------------------------------------------- generate + + def generate(self, query: Optional[Sequence[str]] = None) -> int: + """Ask every direction (or just ``query``) for more. Returns how many landed. + + Callable as often as wanted: each round tells the model what that + direction already holds, so the lists grow instead of repeating. + """ + query = list(query if query is not None else self.query) + added = self._add_to_cached(query, self._generate_keywords(query)) + if added: + self.save_keywords() + return added + + def _generate_keywords(self, query: Sequence[str]) -> List[List[str]]: + """One model call per direction, in a single batch; replies stay aligned with ``query``.""" + prompts: List[Trajectory] = [{ + 'messages': [{'role': 'system', 'content': self.system_prompt}, + {'role': 'user', 'content': self._build_user_prompt(q)}], + } for q in query] + if self._rollout is None: + self._rollout = MultiTurnRollout(self._backend, **self._rollout_kwargs) + return [self._parse_keywords_from_response(assistant_text(t)) + for t in self._rollout(prompts)] + + def _build_user_prompt(self, query: str) -> str: + """The ask for one direction, plus what it already holds as an avoid list.""" + known = self._cached_keywords.get(query, []) + want = max(1, self.num_keywords - len(known)) + user = self._user_prompt.format(k=want, query=query) + if known: + user += self._avoid_lead + ', '.join(known[-self._avoid_max:]) + return user + + @staticmethod + def _parse_keywords_from_response(text: str) -> List[str]: + """The JSON array in ``text``, over-length and non-string entries dropped.""" + body = strip_reasoning(text) + start, end = body.find('['), body.rfind(']') + if start < 0 or end <= start: + return [] + try: + arr = json.loads(body[start:end + 1]) + except (ValueError, TypeError): + return [] + return [s for s in (x.strip() for x in arr if isinstance(x, str)) + if 0 < len(s) <= KEYWORD_MAX_LEN] + + # ----------------------------------------------------------------- store + + def _add_to_cached(self, query: Sequence[str], + keywords: Sequence[Sequence[str]]) -> int: + """Append each direction's new keywords, case-insensitively. Returns how many landed.""" + added = 0 + for q, kws in zip(query, keywords): + bucket = self._cached_keywords.setdefault(q, []) + for kw in kws: + if kw.lower() in self._seen: + continue + self._seen.add(kw.lower()) + bucket.append(kw) + added += 1 + if not added: + # Silence here would read as a model that simply produced less. + logger.warning(f'no new keyword for {len(query)} direction(s); ' + f'everything generated was already known') + return added - def __init__(system_prompt: Optional[str] = None): - pass + def load_keywords(self) -> Dict[str, List[str]]: + """Read the cache back, one direction per line. An unreadable line is skipped.""" + cached: Dict[str, List[str]] = {} + if not (self.path and os.path.exists(self.path)): + return cached + with open(self.path, encoding='utf-8') as f: + for line in f: + if not line.strip(): + continue + try: + r = json.loads(line) + except (ValueError, TypeError): + continue + if isinstance(r.get('query'), str) and isinstance(r.get('keywords'), list): + cached[r['query']] = [kw for kw in r['keywords'] if isinstance(kw, str)] + return cached - def generate_keywords(self, text): - pass + def save_keywords(self) -> None: + """Write the cache out atomically, so a crash mid-write cannot truncate it.""" + if not self.path: + return + os.makedirs(os.path.dirname(os.path.abspath(self.path)) or '.', exist_ok=True) + tmp = self.path + '.tmp' + with open(tmp, 'w', encoding='utf-8') as f: + for q, kws in self._cached_keywords.items(): + f.write(json.dumps({'query': q, 'keywords': kws}, ensure_ascii=False) + '\n') + os.replace(tmp, self.path) diff --git a/src/twinkle_agentic/challenger/new/recorder.py b/src/twinkle_agentic/challenger/new/recorder.py new file mode 100644 index 000000000..c72d6fab3 --- /dev/null +++ b/src/twinkle_agentic/challenger/new/recorder.py @@ -0,0 +1,91 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Persistent proposer trajectories for challenger training and diagnosis.""" +import json +import os +import threading +import uuid +from typing import Any, Dict, List + +import numpy as np + +_TOKEN_FIELDS = ('input_ids', 'labels', 'completion_mask', 'attention_mask', 'position_ids') + + +def _as_numpy(value: Any, dtype: Any = None) -> np.ndarray: + if hasattr(value, 'detach'): + value = value.detach().cpu().numpy() + return np.asarray(value, dtype=dtype) + + +def _logprob_column(logprobs: Any) -> List[float]: + """Extract the chosen token's log probability from each sampling step.""" + out: List[float] = [] + for step in logprobs: + if isinstance(step, (int, float)): + out.append(float(step)) + continue + if isinstance(step, (list, tuple)) and step: + chosen = step[0] + if isinstance(chosen, (list, tuple)) and len(chosen) >= 2: + out.append(float(chosen[1])) + continue + raise TypeError(f'cannot read a chosen-token logprob from {step!r}') + return out + + +def _json_default(value: Any) -> Any: + if hasattr(value, 'tolist'): + return value.tolist() + return str(value) + + +class RolloutRecorder: + """Write token arrays to NPZ and trajectory metadata to a JSONL index.""" + + def __init__(self, save_dir: str): + self.trajectory_dir = os.path.join(save_dir, 'trajs') + self.index_path = os.path.join(self.trajectory_dir, 'index.jsonl') + os.makedirs(self.trajectory_dir, exist_ok=True) + self._lock = threading.Lock() + + def write(self, trajectory: Dict[str, Any], **fields: Any) -> None: + arrays: Dict[str, np.ndarray] = {} + for key in _TOKEN_FIELDS: + value = trajectory.get(key) + if value is not None: + arrays[key] = _as_numpy(value, np.int32) + logprobs = trajectory.get('logprobs') + if logprobs is not None: + arrays['logprobs'] = np.asarray(_logprob_column(logprobs), dtype=np.float64) + + name = f'{uuid.uuid4().hex}.npz' + labels = arrays.get('labels', np.asarray([], dtype=np.int32)) + completion_mask = arrays.get('completion_mask') + if completion_mask is None: + n_policy_tokens = int((labels != -100).sum()) + else: + if completion_mask.size != labels.size: + raise ValueError('completion_mask and labels must have the same number of tokens') + n_policy_tokens = int(((labels != -100) & completion_mask.astype(bool)).sum()) + n_logprobs = len(arrays.get('logprobs', ())) + if logprobs is not None and n_logprobs != n_policy_tokens: + raise ValueError(f'logprobs contain {n_logprobs} policy tokens, expected ' + f'{n_policy_tokens} from labels and completion_mask') + metadata = { + key: value + for key, value in trajectory.items() if key not in _TOKEN_FIELDS and key not in ('logprobs', 'rewards') + } + record = dict(metadata) + record.update(fields) + record.update({ + 'npz': name, + 'n_tokens': int(arrays.get('input_ids', np.asarray([])).size), + 'n_policy_tokens': n_policy_tokens, + 'has_logprobs': logprobs is not None, + 'n_logprobs': n_logprobs, + }) + line = json.dumps(record, ensure_ascii=False, default=_json_default) + with self._lock: + np.savez_compressed(os.path.join(self.trajectory_dir, name), **arrays) + with open(self.index_path, 'a', encoding='utf-8') as handle: + handle.write(line + '\n') diff --git a/src/twinkle_agentic/harness/ms_agent.py b/src/twinkle_agentic/harness/ms_agent.py index e1f8f7bbd..c9da5774c 100644 --- a/src/twinkle_agentic/harness/ms_agent.py +++ b/src/twinkle_agentic/harness/ms_agent.py @@ -13,12 +13,16 @@ training and serving; the Env owns the implementation. Wire the executing side from the same list, or the prompt advertises tools the Env cannot run:: - harness = MsAgentHarness(config) - harness.prepare() - tool_manager = env.tool_manager(harness.tool_schemas()) - rollout = MultiTurnRollout(sampler, template, - tool_manager=tool_manager, harness=harness) - outs = rollout([harness.start(q) for q in queries]) +One harness per trajectory: each holds an ``LLMAgent`` with memory and context +of its own, and episodes run in parallel threads:: + + harnesses = [MsAgentHarness(config) for _ in queries] + for h in harnesses: + h.prepare() + tool_managers = [env.tool_manager(h.tool_schemas()) for h, env in zip(harnesses, envs)] + rollout = MultiTurnRollout(sampler, template) + outs = rollout([h.start(q) for h, q in zip(harnesses, queries)], + tool_manager=tool_managers, harness=harnesses) Serving path keeps using ``LLMAgent.run()`` with the same ``agent.yaml`` and the same :class:`~twinkle_agentic.envs.base.Env` backend. This class must diff --git a/src/twinkle_agentic/preprocessor/AUDIT.md b/src/twinkle_agentic/preprocessor/AUDIT.md deleted file mode 100644 index a5f1a71bd..000000000 --- a/src/twinkle_agentic/preprocessor/AUDIT.md +++ /dev/null @@ -1,179 +0,0 @@ -# Preprocessor 审计与整改清单 - -> 审计范围：`src/twinkle_agentic/preprocessor/` 全部 15 个文件、30 个类。 -> 审计方法：逐行只读审阅 + 关键断言代码复核 + cookbook/tests 实际接线核实。 -> 三轮视角：(A) 实现问题 (B) 类设计/拆分合并 (C) 功能增删。 - -## 实施状态（已按本清单完整落地） - -> A1 经确认跳过：R1 已把 `score_filter.py` 整体移入 `experimental/`（零 active 使用的死代码）， -> 对死代码再做 4 文件拆分只增维护面、零收益，启用前再拆。其余 21 项全部实施。 - -| 项 | 状态 | 落地位置 | -|----|------|----------| -| A5 | ✅ | `label_schema.py`（`user_data` 信封 + `set_labels`/`get_label` + `pack_value`） | -| P3 | ✅ | `message_normalizer.py` Pass1 重建 assistant 时 `dict(msg)` 透传全字段 | -| P7 | ✅ | `hard_filter._has_tool_calls` / `message_normalizer._strip_heartbeat`+`_is_atomic` 全部走 `normalize_tool_calls` | -| D7 | ✅ | `trajectory_scorer.py`（Segmenter→HardScorer 逐轮→fuse_segment→aggregate_trajectory→写回 `user_data`，mapper 不删） | -| D7c | ✅ | `trajectory_scorer.py` `_segment_confidence`（一致性+voting稳定+决断性）+ `RubricVerifier.score_detail(extra_context=)` 客观注入重评 | -| D6 | ✅ | `outcome_filter.py`（纯读 `traj_score`/`safety_*` 标签比阈值，fail-open） | -| D8 | ✅ | `safety_scorer.py` + `RubricVerifier(fixed_rubric=)` 固定安全 rubric | -| D9 | ✅ | `pii_presidio_filter.py` `regex_only=True`（stub NlpEngine 免 spaCy，REPLACE→MASK 免 faker） | -| D10 | ✅ | `provenance.py`（`ProvenanceStamp`，血缘写入 `user_data`） | -| R1 | ✅ | `experimental/`（`score_filter.py` + `llm_backend.py` git mv 移出主包） | -| R3/R4/R5 | ✅ | `intent_classifier.py`（默认不删；DEFAULT_DETECTORS 精简为 ToolCall/Code/Math；LLM 路径经 R1 已全归 `llm_backup`） | -| P1/P5/P6/P8/P10 | ✅ | trim 后重算 `is_agent`；deadloop agent 行改扫有文本轮；`max_rounds` 按 pair；refuse 扫全 assistant+可选 reasoning；system 多模态保护 | -| A2/A4 | ✅ | `logprob_utils.py` + `message_utils.py`（`utils.py` 保留 shim）；`intents.py` 常量下沉 | -| A3 | ✅ | `twinkle/preprocessor/base.py` 基类返回 `Tuple[List,List]` + `Mapper`/`Filter` 语义基类（`ModelFilter`/`ProvenanceStamp` 已改用） | -| D4/D5 | ✅ | `language_filter.py`（langid 可选，启发式回退）；`structural_noise.py`（关键词无关噪声轮打标） | -| D1/D2 | ✅ | `offline/near_dedup.py`（MinHash-LSH，datasketch 可选+纯 Python 回退）；`offline/decontaminate.py`（13-gram 重叠，drop/tag） | -| A1 | ⏭️ 跳过 | 见上（R1 已隔离为死代码） | - ---- - -## 0. 结论速览 - -> 本清单已根据 review 意见复核收敛：P2/P4 撤销，P1 降级，P11/P12 归入 R1（死代码，暂不单独修）。 - -- **共需改动 22 项**：实现问题 6（P1、P3、P5–P10）、结构重构 5（A1–A5）、功能增删 11（R1–R5 + D1/D2/D4/D5 + D6/D7/D7c/D8/D9/D10，D3 废弃）。 -- **必须做（会静默损坏训练数据）**：仅 **P3** 一项（工具归一丢 reasoning 字段）。 -- **达成「干净 + 每轮评分」最终目标的核心**：**A5**（`user_data` 信封，去 DAG 前置）+ **D7**（接线 verifier，分数写回每轮）+ **D7c**（自动校准 + 客观纠偏主观重评）+ **D6**（读标签滤废案）+ **D8**（安全 rubric）+ **D9**（PII 纯 regex）。 -- **一句话结论**：现有清单修的是「清洗器 bug + 基础过滤」；要产出「干净且每轮带可信分」的 trajectory，还差——**A5 统一 `user_data` 标签信封（把评分/过滤解耦成打标 mapper + 末尾读标签 filter，去掉 DAG）+ 每轮评分打标(D7) + 自进化校准(D7c) + 废案过滤(D6) + 安全/PII(D8/D9)**。零件多数已存在（`verifier`+`aggregation`+`RubricVerifier`+`llm_backup`），核心工作是**接线 + 定 `user_data` 契约**。 - -### review 复核结论（撤销 / 降级项） - -| 原项 | review 意见 | 复核结论 | 处置 | -|------|-------------|----------|------| -| **P1** trim 砍 tool 尾 | 不以 assistant 结尾的部分无训练必要，最多用于工具调用打分 | 成立。trim 尾部未闭合 tool 对训练无害；`is_agent` 不更新的副作用仅剩“末尾 `assistant(tool_calls)` 无对应结果”，训练时本应 mask | **降级为中等**，改描述，不再算“数据损坏” | -| **P2** heartbeat 误杀 | 这类数据是 openclaw/OpenHands 常见格式，作者本意就是要删 | 成立。agent 轨迹清洗语境下 heartbeat 轮几乎必为真噪声，误杀率极低；`message_normalizer.py:26-27` 注释确认是**故意**删除 | **撤销**（保留现状；可选加词边界，非必做） | -| **P4** 删 reasoning-only 轮 | 只有 thinking 无工具调用，训练无落点 | 成立。纯 thinking 轮无 target 输出，多轮里是悬空推理，删掉合理 | **撤销** | -| **P11** ParaphraseScorer 崩溃 | 应该没有实际使用 | 成立。`ScoreFilter`/`ParaphraseScorer` **全库零 active 使用**（仅自身定义 + docs 示例 + 注释掉的引用），测试只覆盖 `utils` 数学函数 | **归入 R1**（死代码，启用时再修） | -| **P12** IFD 公式口径 | 同上 | 同上 | **归入 R1** | - ---- - -## 一、实现问题（正确性 / 语义） - -### 严重：会静默损坏训练数据 - -| ID | 位置 | 问题 | 改动 | 预期收益 | -|----|------|------|------|----------| -| **P3** | `message_normalizer.py` Pass 1 重建消息 | 工具归一路径只保留 `role/content/tool_calls/tool_call_id` 四字段，**丢弃 `reasoning_content`/`thinking`/`name`** | 重建时透传全部原字段 | reasoning 蒸馏数据不再被清洗流程静默剥离 | - -### 中等：策略漏洞 / 语义错位 - -| ID | 位置 | 问题 | 改动 | 预期收益 | -|----|------|------|------|----------| -| **P1** | `message_sanity.py:317,324-329` trim + `is_agent` | trim 掉末尾未闭合 tool 结果本身对训练无害，但 `is_agent` 在 trim **前**计算、trim 后不更新，残留“末尾 `assistant(tool_calls)` 无对应结果”，`check_tool_matching`（forward-only）不拦 | trim 后重算 `is_agent`，或末轮 `tool_calls` 无结果时 mask/剥离该 call；**非必做** | 末轮悬空 tool_call 得到一致处理，避免训练时误算 loss | -| **P5** | `dead_loop_filter.py:192-194` | `is_agent_row` 为真则**整行跳过** stuck 检测，agent 恰恰最易死循环 | agent 死循环走 `HardScorer.check_no_repeated_calls`（见 D3） | 覆盖 agent 重复工具调用循环，堵住最大系统性漏检 | -| **P6** | `hard_filter.py` `max_rounds` | 实现是 `len(asst_msgs) > max_rounds`，只数 assistant，注释却写 “user-assistant pairs” | 修正为按 pair 计数或改注释与语义一致 | 轮数过滤阈值语义正确 | -| **P7** | `message_normalizer.py` / `hard_filter.py` / `utils.py` | `tool_calls` 真值判断三处不一致（裸真值 vs `_has_tool_calls` 视 `''`/`'[]'`/`[]` 为空 vs `normalize_tool_calls`） | 全部统一走 `normalize_tool_calls` | 同一数据“是否 agent”判定一致，消除跨 filter 不一致 | -| **P8** | `refuse_filter.py` | 只扫首条 assistant 前 600 字、不读 reasoning，多轮/reasoning 拒答漏检 | 扩到全 assistant + reasoning 字段（可配窗口） | 拒答样本召回上升，减少污染 | -| **P9** | `token_soup.py` | `max_chars>0` 只查头部（cookbook 用 8000），尾部乱码漏检；不扫 reasoning | 全文 + reasoning 扫描或分段抽样 | 乱码样本召回上升 | -| **P10** | `message_sanity.py` `consolidate_system_messages` | 合并 multimodal system 时压成纯字符串，可能丢非 text part | 用 `msg_has_media` 保护多模态 system | 多模态 system 不丢内容 | - ---- - -## 二、架构 / 类设计（拆分与合并） - -| ID | 项 | 判断 | 改动 | 预期收益 | -|----|----|------|------|----------| -| **A1**（高） | `score_filter.py` 9 个类 486 行 | 契约与实现混在一起，加 scorer 就改巨型文件 | 拆为 `score/` 子包：`types.py`(RoundContext/ScoreResult/Scorer) + `scorers.py`(ChrMin/SIFD 轻) + `judge.py`(PassN/Paraphrase 重) + `score_filter.py`(编排) | 开闭原则；轻/重依赖分离；新增 scorer 零侵入 | -| **A2**（中） | `utils.py` | logprob 数学 + 消息格式工具两个无关模块塞一起 | 拆为 `logprob_utils.py` + `message_utils.py` | 降耦合；改 score 逻辑不误碰消息工具 | -| **A3**（高） | `twinkle/preprocessor/base.py:39` | 基类 `__call__` 声明返回 `Dict`，所有子类实际返回 `Tuple[kept, dropped]`，类型契约名存实亡 | 基类改 `-> Tuple[List, List]`；可选分 `Mapper`/`Filter` 语义基类 | 类型检查生效；新人不会照错签名写导致解包崩溃 | -| **A4**（低） | intent 常量位置 | `ScoreFilter` 消费 intent，但常量定义在 `intent_classifier.py`，score 独立后形成跨模块依赖 | intent 常量下沉到轻量 `intents.py` | 为 score 子包独立化铺路 | -| **A5**（高，目标前置） | 统一 `user_data` 标签信封 | 评分/安全/血缘无统一落点；D6↔D7 若代码互调会逼出 DAG | 所有标签走 `user_data` 的 `List[Tuple[str, pack_value(v)]]`（PyArrow 稳定，见 D 节前置）；打标 mapper 写、末尾 filter 读，靠列表顺序解耦 | 去 DAG、统一数据契约；A3 返回契约的自然延伸 | - -**明确不动**（避免过度设计）：`HardFilter`/`RefuseFilter`/`DeadLoopFilter`/`TokenSoupFilter` 保持独立（合并成上帝类只会更糟）；`data_juicer.py` 4 个薄封装保持一文件；`LLMBackend` 三类保持；`MessageNormalizer` 3 个 pass 不拆（有强顺序依赖）；`IntentDetector` 层级设计是全代码最佳，保持。 - ---- - -## 三、功能增删 - -### 建议去掉 / 降级（死代码与过度设计） - -| ID | 项 | 证据 | 改动 | 预期收益 | -|----|----|------|------|----------| -| **R1** | `ScoreFilter` 全家（+4 scorer）**+ `llm_backend.py` 整个文件** | `ScoreFilter` **全库零 active 使用**（仅自身定义 + docs 示例 + `train_cold_start.py:216` 注释态；测试只覆盖底层 `utils` 数学函数）；内含两处死代码 bug —— 原 **P11**（`ParaphraseScorer:452-455` 缺 DP pad，小批量/DP>1 时 `SamplerBackend` raise）、原 **P12**（`utils.py:154` `ifd=exp(-mean_delta)` 是 Superfiltering 差分指数口径而非 Cherry 损失比值，阈值不可互换）。`llm_backend.py`（`LLMBackend`/`OpenAIBackend`/`SamplerBackend`）**唯一消费者就是 `ScoreFilter`**（grep 确认除自身+`__init__` 导出外无他），它专为 score 打分提供 `chat`/`prompt_logprobs`/`prompt_logprobs_ids`/`embeddings` | `ScoreFilter` + `llm_backend.py` 一起移出主包到 `experimental/` 或 `data_selection/`，标记未验证；**启用前**再修 P11（复用 `_pad_batch`）+ 明确 P12 IFD 口径并重标阈值 | 主路径不再拖未验证的重代码 + 未接线的 LLM 后端抽象；bug 修复延后到真正需要时 | -| **R2** | `LLMBackend.embeddings()` | 是 R1 中 `llm_backend.py` 的一部分；preprocessor 侧**零调用者**，`SamplerBackend.embeddings` 直接 raise | 随 R1 一并移出（不单独保留伪抽象） | 去掉未接线接口，等真需要 embedding 去重再加 | -| **R3** | `IntentClassifier` | 产物 `key_rounds/intents` 主消费者是死的 `ScoreFilter`；`dataset_think.py` 只 import 不入 pipeline | 重定位为“标注器”（从不 drop）；短期可移出主 pipeline 省 CPU | 明确职责；省无谓计算 | -| **R4** | `ComplexLogic/Reasoning/UserDissatisfaction` Detector | 仅被 `IntentClassifier.DEFAULT_DETECTORS` 引用，下游死 | 精简 default 到 `ToolCall/Code/Math` | 减少无消费者的启发式维护面 | -| **R5** | LLM 调用抽象统一到 `llm_backup` | preprocessor 里活跃的 LLM 生成需求（summarizer/segment/verifier）**早已全部走 `twinkle_agentic/utils/llm_backup.py`**（置信度路由 student/teacher + 蒸馏数据收集）；只有死代码 `ScoreFilter` 还用独立的 `LLMBackend` | 主路径不再引入独立 `LLMBackend` 抽象，生成类需求统一走 `@llm_backup`。**注意**：`llm_backup` 只提供 chat 生成（返回 content 字符串），**不提供 `prompt_logprobs`/`embeddings`**——IFD/chr_min 类 logprob 数据选择若复活，那部分接口需在 `experimental/` 内单独保留或重写，不能指望 `llm_backup` | 收敛到单一蒸馏路由机制；生成享受 student/teacher 蒸馏；消除重复的推理后端抽象 | - -### 建议增加（真正缺失的清洗能力） - -#### 已列（清洗器层） - -| ID | 项 | 缺口 | 改动 | 预期收益 | -|----|----|------|------|----------| -| **D1**（离线） | 近重复去重 MinHash-LSH/SimHash | `DedupFilter` 只做前缀 md5 精确去重，改一字的近重复全漏 | 扩展 `DedupFilter` 或新增 `NearDupFilter`（`datasketch` 轻依赖）。**限离线批处理阶段**（需全局视图）；实时 per-batch 主路径不启用，否则局部视图导致误杀严重 | 相似轨迹被挡，多样性上升；离线做，避免在线误杀 | -| **D2**（离线） | 基准去污染 decontamination | train/test n-gram overlap **完全没有** | 新增 13-gram 重叠过滤，比对**静态** benchmark n-gram 索引。**限离线**或**只打标不删**，避免实时流误杀正常样本 | 评测不被污染，指标可信 | -| **D4**（中） | 语言识别 langid/fastText | 只有 `cjk_ratio` script 比例，粗糙 | 轻量 langid 过滤 | 中英限定更可靠，混语噪声下降 | -| **D5**（低，可选） | 结构化噪声轮识别 | 现有 heartbeat 靠关键词（对 openclaw/OpenHands 格式已够用，见 P2 撤销）；仅当出现无关键词的结构性噪声轮时才需要 | 极短轮 + 高重复 + embedding 距离（复用 D1 基础设施） | 覆盖无关键词的噪声轮；非当前痛点 | - -> **D3（agent 死循环接线）已废弃**：review 指出 `preprocessor` 与 `verifier` 当前**零互相 import**（grep 确认），让 `DeadLoopFilter` 去 import `HardScorer`（一个 RL reward `Verifier`）会破坏模块边界、且职责串（清洗器 vs 打分器）。正确路径并入 **D6/D7**：新增打标/评分 preprocessor，agent 死循环由其中的确定性 check 覆盖。 - -#### 新增（达成「干净 + 每轮评分」最终目标所需的整段能力） - -> 对标业界标准 agent 数据流程（Llama-3 / DeepSeek-V3 / Nemotron / Tulu-3 / AgentInstruct / ToolBench）。目标五属性映射：无不良信息→D8/D9、无废案→D6、无重复冗余→D1+D6(轨迹内)、无心跳→已有、每轮评分→D7。 - -| ID | 项 | 属性 | 缺口 | 改动 | 预期收益 | -|----|----|------|------|------|----------| -| **D7**（高，核心） | 每轮评分打标 preprocessor（**只打标不过滤**） | 每轮评分 | `verifier`（per-round `HardScorer` + per-segment `RubricVerifier`）+ `aggregation`（round→segment→trajectory）**基础设施现成但未接线**；`aggregation.py:27` 明说编排器 `TrajectoryScorer` 未实现 | **新增 preprocessor**（mapper，从不 drop）：`Segmenter → HardScorer(逐轮) → RubricVerifier(逐段) → aggregation → 分数写回 `user_data``。分数、`score_confidence`、安全标全部作为 `(key, pack_value(v))` 追加进 `user_data`（见架构前置 A5）。`RubricVerifier.score_detail()` 已返回完整 `ScoreDetail`，`__call__(trajectory)` 兼容逐行调用 | 直接产出「每轮评分」的 trajectory；打标与过滤解耦，D6/D8 只读标签 | -| **D7c**（高，核心） | 评分校准探针 + 客观→主观重评（自进化，无人评） | 每轮评分可信度 | 自进化框架**不能靠人评对齐**；未校准的分会系统性放大 judge 偏见 | 用三个**自动**信号合成 per-segment `score_confidence`：①**teacher-student 一致性**（复用 `llm_backup` 已收集的 `(student, teacher, match)`）②**结果锚定**（`HardScorer` 确定性 check 当弱标签探针）③**voting 方差**（`RubricVerifier` 已有 voting，导出方差）。**关键**：当客观（硬 check）与 LLM 主观**不一致**时，**把客观结果注入 rubric 的打分上下文，让 `RubricVerifier` 重新评分**（不是简单降权，是带硬信号修正的二次评分） | 分数可信度自动量化；客观事实纠偏主观判断，闭环收敛；零人工 | -| **D6**（高） | 轨迹成败判定（过滤废案）——**纯读标签 filter** | 无废案 / 轨迹内冗余 | 无 outcome verification：失败/绕圈/工具全错/最终答案错的轨迹留在训练集 | **不自己算分**，只 `user_data_get(row, 'traj_score')` 等标签跟**阈值**比 → 判废案 drop（依赖的是 D7 已写好的**数据标签**，不是 D7 的代码，靠 pipeline 列表顺序保证 D6 在 D7 后）。**阈值先拍默认值，实测回收分布后回调**（不做人评标定） | 废案不进训练集；与打标解耦，无模块依赖 | -| **D8**（高） | 安全/毒性评分（复用 rubric） | 无不良信息 | 只有 `RefuseFilter`（拒答正则）+ 敏感词表，无 toxicity/safety 覆盖暴力/仇恨/成人/越狱成功 | **复用 `RubricVerifier`**：把安全维度作为一组**固定 `RubricItem`**（暴力/仇恨/成人/越狱成功/隐私泄露）注入 stage-2 打分，走现有 `_score_with_voting` + `_aggregate`；低于阈值判不良。**无需新分类器/新依赖** | 安全过滤召回远超敏感词表；与 D7 共用打分基础设施 | -| **D9**（中） | PII 真脱敏（激活现有 Presidio） | 无不良信息 | `PIIPresidioFilter` 存在但曾因**慢**去掉（spaCy NER 是瓶颈） | 加回来但**纯 regex 模式**：现 `IGNORED_ENTITIES` 已忽略全部 NER 实体（PERSON/LOCATION/ORG…），只留 regex 标识符（邮箱/电话/证件/银行卡）→ **可不加载 spaCy**，绕过 NER 瓶颈，速度问题基本消除 | 邮箱/电话/证件等真 PII 脱敏，且不拖慢管线 | -| **D10**（中） | 治理层：provenance（血缘字段） | 可追溯 | 无血缘字段（source/teacher_model/timestamp） | 每条 trajectory 把血缘作为 `(key, pack_value(v))` 写进 `user_data`（蒸馏场景 teacher/student 版本）。**批次归因/数据卡暂缓**（backlog，见「不做」） | 可追溯；对标 Nemotron/Dolma 但先只做血缘字段 | - -> **架构前置 A5（去 DAG 的正解）**：所有评分/安全/血缘标签统一写进 **`user_data` 信封**，把「评分」与「过滤」解耦成「**打标 mapper（D7/D8/D10，从不 drop）+ 末尾纯读标签 filter（D6）**」。这样 D6→D7 是**数据依赖**（D7 写标签、D6 读标签），靠**线性 `QualityPreprocessor` 的列表顺序**保证，**无需 DAG、无模块互相 import**。 -> -> **PyArrow 硬约束**：`user_data` 必须是 **`List[Tuple[str, str]]`**，**不能用 dict**（HF `datasets` 的 PyArrow 后端对异构/嵌套 dict 序列化有问题）。已核实这是仓库现有官方约定 —— `twinkle/data_format/trajectory.py:18-19`（`user_data: List[Tuple[str, str]]`，注释 "PyArrow-stable encoding: each entry is (key, json.dumps(value))"），写用 `pack_value(v)`（JSON 字符串，值可为任意结构但对外恒为 `(str, str)`），读用 `user_data_get(row, key)`。每轮分数可用 `(f'round_{i}_score', pack_value(...))` 或 `('round_scores', pack_value([...]))` 形态。 -> -> 确定性 check 复用：D6/D7/D8 都要 `HardScorer` 的 LLM-free check。可抽到无依赖公共层（如 `twinkle_agentic/agent_checks.py`）供 verifier 与新 preprocessor 各自依赖；用户已确认也可接受 preprocessor→verifier 单向依赖，则公共层后置。 - -### 明确不做（自进化框架的取舍） - -| 项 | 为什么不做 | -|----|-----------| -| 人评校准对齐 | 自进化框架不 scale；改用 D7c 的三信号（teacher 一致性 + 结果锚定 + voting 方差）替代 | -| 批次间质量归因 / 数据集版本 diff | 暂缓（backlog），当前不阻塞可用性 | -| 管线内数据配比 / 分层采样 | 移到**训练时的 sampler** 消费 `user_data` 标签，清洗管线只负责打标 | -| DAG / 阶段化编排引擎 | 用 A5 的「打标 + 末尾读标签」拍平成线性，不引入 DAG | - ---- - -## 四、改动项汇总与优先级 - -| 优先级 | 项 | 类型 | 是否引入依赖 | -|--------|----|------|--------------| -| P0 前置 | **A5** | `user_data` 信封（去 DAG，其余目标项的地基） | 否 | -| P0 必做 | P3 | 数据损坏（丢 reasoning） | 否 | -| P1 目标核心 | D7, D7c, D6, D8 | 每轮评分 / 校准 / 废案 / 安全（接线 verifier） | 否（复用 verifier） | -| P1 高 | P7, A1, A3 | 一致性 / 结构重构 | 否 | -| P2 中 | D9, D10 | PII 脱敏 / 血缘 | presidio(D9) | -| P2 中 | P1, P5, P6, P8, P10, A2, D4 | 语义/漏检/降耦合 | langid(D4) | -| P2 中（仅离线） | D1, D2 | 去重/去污染（防实时误杀） | `datasketch`(D1) | -| P3 低 | A4, R1, R2, R3, R4, R5, D5 | 清理/重定位/增强 | 否 | - -**合计 22 项**：实现问题 6（P1、P3、P5–P10）、架构 5（A1–A5）、功能增删 11（R1–R5 + D1/D2/D4/D5 + D6/D7/D7c/D8/D9/D10）。原 P2/P4 撤销，P11/P12 归入 R1，**D3 废弃**（并入 D6/D7）。 - ---- - -## 五、改动后预期整体收益 - -1. **达成最终目标（干净 + 每轮可信分）**：A5 统一标签信封 → D7 分数写回每轮 → D7c 自动校准 + 客观纠偏主观 → D6 读标签滤废案 → D8 安全 rubric → D9 PII；heartbeat 已有、D1 离线去重。五属性齐活，且分数带 `score_confidence`。 -2. **去 DAG**：A5 把「打标 mapper + 末尾读标签 filter」拍平成线性 `QualityPreprocessor`，D6↔D7 只有数据依赖、零模块互调，无需 DAG 引擎。 -3. **数据正确性**：P3 消除 reasoning 被静默剥离；P1 末轮悬空 tool_call 一致处理。 -4. **复用而非新建**：D7 复用 `verifier`+`aggregation`；D7c 复用 `llm_backup` 一致性 + `RubricVerifier` voting；D8 复用 rubric 固定项；D9 纯 regex 免 spaCy —— **几乎零新依赖**。 -5. **可维护性 / 一致性**：A1–A4 契约清晰、依赖分层；R1–R5 砍死代码；P7 统一 `tool_calls` 判定；D10 血缘可追溯。 - -> 落地顺序建议： -> 1. **A5**（定 `user_data` 标签信封 + list-of-tuple/`pack_value` 约定）→ 所有目标项的地基。 -> 2. **P3**（防丢数据，无依赖）。 -> 3. **D7**（每轮评分打标 mapper，接线 verifier+aggregation，写回 `user_data`）。 -> 4. **D7c**（三信号校准 + 客观→主观重评）→ 让分数可信。 -> 5. **D6 + D8**（读标签滤废案 + 安全 rubric，阈值先拍后测）。 -> 6. **D9**（PII 纯 regex 加回）→ **D10**（血缘字段）。 -> 7. **A1 + A3**（结构地基）→ 其余（P1/P5/P6/P8/P10/D4）视数据分布投入。 -> 8. **D1 + D2** 放离线批处理阶段单独跑；配比放训练 sampler。 diff --git a/src/twinkle_agentic/preprocessor/__init__.py b/src/twinkle_agentic/preprocessor/__init__.py index 9734b44b5..9f5deb24c 100644 --- a/src/twinkle_agentic/preprocessor/__init__.py +++ b/src/twinkle_agentic/preprocessor/__init__.py @@ -17,9 +17,7 @@ from .message_sanity import MessageSanityFilter from .model_filter import ModelFilter from .pii_presidio_filter import PIIPresidioFilter -from .provenance import ProvenanceStamp # noqa: F401 from .refuse_filter import RefuseFilter -from .structural_noise import StructuralNoiseTagger # noqa: F401 from .token_soup import TokenSoupFilter logger = get_logger() diff --git a/src/twinkle_agentic/preprocessor/experimental/__init__.py b/src/twinkle_agentic/preprocessor/experimental/__init__.py deleted file mode 100644 index 6257664e3..000000000 --- a/src/twinkle_agentic/preprocessor/experimental/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Experimental / not-yet-wired preprocessor components (AUDIT R1). - -These modules are kept out of the main :mod:`twinkle_agentic.preprocessor` -namespace because they have no active consumer in the shipped pipeline -(``QualityPreprocessor``) and are not exercised by the cookbook or tests: - -- :class:`ScoreFilter` and its scorers (per-round SFT key-round selection). Active - LLM generation in the framework goes through ``twinkle_agentic.utils.llm_backup`` - instead of the local ``LLMBackend`` abstraction. -- :class:`LLMBackend` / :class:`OpenAIBackend` / :class:`SamplerBackend`, which - exclusively serve ``ScoreFilter``. - -Import explicitly from here if you want to experiment with them, e.g.:: - - from twinkle_agentic.preprocessor.experimental import ScoreFilter, SamplerBackend -""" -from .llm_backend import LLMBackend, OpenAIBackend, SamplerBackend # noqa: F401 -from .score_filter import ScoreFilter # noqa: F401 - -__all__ = ['ScoreFilter', 'LLMBackend', 'OpenAIBackend', 'SamplerBackend'] diff --git a/src/twinkle_agentic/preprocessor/experimental/llm_backend.py b/src/twinkle_agentic/preprocessor/experimental/llm_backend.py deleted file mode 100644 index 002618620..000000000 --- a/src/twinkle_agentic/preprocessor/experimental/llm_backend.py +++ /dev/null @@ -1,344 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Abstract LLM backend for preprocessor pipeline. - -Supports two modes: - - OpenAIBackend: httpx-based calls to any OpenAI-compatible HTTP server - - SamplerBackend: direct calls to Twinkle vLLMSampler Ray actor (no HTTP) -""" -from abc import ABC, abstractmethod -from typing import Any, Dict, List, Optional, Tuple - -from twinkle.utils import get_logger - -logger = get_logger() - - -class LLMBackend(ABC): - """Abstract base for LLM inference used by QualityPreprocessor stages.""" - - @abstractmethod - def chat( - self, - messages: List[Dict[str, Any]], - *, - temperature: float = 0.0, - max_tokens: int = 16, - n: int = 1, - ) -> List[Dict[str, str]]: - """Chat completion. - - Returns: - List of n choices, each a dict with keys 'content' and 'reasoning_content'. - """ - - def chat_batch( - self, - messages_list: List[List[Dict[str, Any]]], - *, - temperature: float = 0.0, - max_tokens: int = 16, - n: int = 1, - ) -> List[List[Dict[str, str]]]: - """Batched chat completion. Returns one List[choice] per input messages list. - - Default impl loops over `chat`; backends should override to fan out concurrently - (HTTP) or pass the full list to the underlying sampler in a single call (vLLM DP). - """ - return [self.chat(m, temperature=temperature, max_tokens=max_tokens, n=n) for m in messages_list] - - @abstractmethod - def prompt_logprobs(self, messages: List[Dict[str, Any]]) -> Optional[List]: - """Evaluate prompt tokens without generation. - - Returns: - List of per-token logprob entries (format varies by backend but - is compatible with _extract_logprob helpers), or None on failure. - """ - - @abstractmethod - def prompt_logprobs_ids(self, input_ids_list: List[List[int]]) -> List[List]: - """Batched: evaluate raw token-id prompts without chat template wrapping. - - Used for unconditional perplexity (e.g. IFD denominator). Caller MUST - supply a list of token-id sequences; for distributed backends the list - length must satisfy backend-specific batching constraints (e.g. - ``len >= dp_world_size`` for SamplerBackend). - """ - - def embeddings(self, texts: List[str]) -> Any: - """Compute text embeddings. Override in backends that support it.""" - raise NotImplementedError(f'{type(self).__name__} does not support embeddings') - - -class OpenAIBackend(LLMBackend): - """Backend wrapping any OpenAI-compatible HTTP endpoint.""" - - def __init__( - self, - endpoint: str, - model: str = 'default', - api_key: str = '', - timeout: float = 120.0, - ): - import httpx - headers = {'Content-Type': 'application/json'} - if api_key: - headers['Authorization'] = f'Bearer {api_key}' - self._client = httpx.Client(timeout=timeout, headers=headers) - base = endpoint.rstrip('/') - self._chat_endpoint = f'{base}/v1/chat/completions' - self._embed_endpoint = f'{base}/v1/embeddings' - self._model = model - - @property - def model(self) -> str: - return self._model - - def chat( - self, - messages: List[Dict[str, Any]], - *, - temperature: float = 0.0, - max_tokens: int = 16, - n: int = 1, - ) -> List[Dict[str, str]]: - try: - resp = self._client.post( - self._chat_endpoint, - json={ - 'model': self._model, - 'messages': messages, - 'temperature': temperature, - 'max_tokens': max_tokens, - 'n': n, - }) - resp.raise_for_status() - choices = resp.json().get('choices', []) - results = [] - for c in choices: - msg = c.get('message') or {} - results.append({ - 'content': msg.get('content') or '', - 'reasoning_content': msg.get('reasoning_content') or '', - }) - return results - except Exception as e: - logger.warning(f'[OpenAIBackend] chat failed: {e}') - return [] - - def chat_batch( - self, - messages_list: List[List[Dict[str, Any]]], - *, - temperature: float = 0.0, - max_tokens: int = 16, - n: int = 1, - max_workers: int = 16, - ) -> List[List[Dict[str, str]]]: - """Concurrent chat: vLLM HTTP server multiplexes requests; httpx.Client is thread-safe.""" - from concurrent.futures import ThreadPoolExecutor - if not messages_list: - return [] - workers = max(1, min(max_workers, len(messages_list))) - results: List[List[Dict[str, str]]] = [[] for _ in messages_list] - with ThreadPoolExecutor(max_workers=workers) as ex: - futs = { - ex.submit(self.chat, m, temperature=temperature, max_tokens=max_tokens, n=n): i - for i, m in enumerate(messages_list) - } - for fut in futs: - results[futs[fut]] = fut.result() - return results - - def prompt_logprobs(self, messages: List[Dict[str, Any]]) -> Optional[List]: - try: - resp = self._client.post( - self._chat_endpoint, - json={ - 'model': self._model, - 'messages': messages, - 'max_tokens': 0, - 'prompt_logprobs': 1, - }) - resp.raise_for_status() - return resp.json().get('prompt_logprobs') - except Exception: - return None - - def prompt_logprobs_ids(self, input_ids_list: List[List[int]]) -> List[List]: - endpoint = self._chat_endpoint.rsplit('/', 2)[0] + '/v1/completions' - results: List[List] = [] - for input_ids in input_ids_list: - resp = self._client.post( - endpoint, - json={ - 'model': self._model, - 'prompt': list(input_ids), - 'max_tokens': 0, - 'echo': True, - 'prompt_logprobs': 1, - }) - resp.raise_for_status() - data = resp.json() - choices = data.get('choices') or [] - if choices and 'prompt_logprobs' in choices[0]: - results.append(choices[0]['prompt_logprobs']) - else: - results.append(data['prompt_logprobs']) - return results - - def embeddings(self, texts: List[str]): - import numpy as np - resp = self._client.post( - self._embed_endpoint, json={ - 'model': self._model, - 'input': texts, - }) - resp.raise_for_status() - data = resp.json().get('data', []) - data_sorted = sorted(data, key=lambda x: x.get('index', 0)) - return np.array([d['embedding'] for d in data_sorted], dtype=np.float32) - - -class SamplerBackend(LLMBackend): - """Backend wrapping a Twinkle vLLMSampler (Ray actor, no HTTP overhead).""" - - def __init__( - self, - sampler, - embed_endpoint: str = '', - embed_model: str = 'bge-m3', - ): - """ - Args: - sampler: A vLLMSampler instance (with template already set). - embed_endpoint: Optional OpenAI-compatible endpoint for embeddings. - embed_model: Model name for embeddings. - """ - self._sampler = sampler - self._embed_endpoint = embed_endpoint - self._embed_model = embed_model - self._embed_client = None - if embed_endpoint: - import httpx - self._embed_client = httpx.Client(timeout=120.0) - self._embed_url = f'{embed_endpoint.rstrip("/")}/v1/embeddings' - - def chat( - self, - messages: List[Dict[str, Any]], - *, - temperature: float = 0.0, - max_tokens: int = 16, - n: int = 1, - ) -> List[Dict[str, str]]: - from twinkle.data_format import SamplingParams - trajectory = {'messages': messages} - params = SamplingParams( - temperature=temperature, - max_tokens=max_tokens, - num_samples=n, - ) - try: - responses = self._sampler.sample(trajectory, params) - results = [] - for resp in responses: - for seq in resp.sequences: - text = seq.decoded or '' - reasoning = '' - if '</think>' in text: - parts = text.split('</think>', 1) - reasoning = parts[0].split('<think>')[-1].strip() - text = parts[1].strip() - results.append({'content': text, 'reasoning_content': reasoning}) - return results - except Exception as e: - logger.warning(f'[SamplerBackend] chat failed: {e}') - return [] - - @staticmethod - def _split_think(text: str) -> Tuple[str, str]: - if '</think>' in text: - parts = text.split('</think>', 1) - return parts[1].strip(), parts[0].split('<think>')[-1].strip() - return text, '' - - def chat_batch( - self, - messages_list: List[List[Dict[str, Any]]], - *, - temperature: float = 0.0, - max_tokens: int = 16, - n: int = 1, - ) -> List[List[Dict[str, str]]]: - """One sampler dispatch over the full list; lets vLLM DP workers stay saturated.""" - from twinkle.data_format import SamplingParams - if not messages_list: - return [] - device_mesh = getattr(self._sampler, 'device_mesh', None) - dp_world_size = getattr(device_mesh, 'dp_world_size', 1) or 1 - n_inputs = len(messages_list) - feats = [{'messages': m} for m in messages_list] - # Pad the dispatch so every DP worker has at least one item; trim duplicates after. - if n_inputs < dp_world_size: - feats = feats + [feats[-1]] * (dp_world_size - n_inputs) - params = SamplingParams(temperature=temperature, max_tokens=max_tokens, num_samples=n) - try: - responses = self._sampler.sample(feats, params) - except Exception as e: - logger.warning(f'[SamplerBackend] chat_batch failed: {e}') - return [[] for _ in range(n_inputs)] - responses = list(responses)[:n_inputs] - out: List[List[Dict[str, str]]] = [] - for resp in responses: - choices: List[Dict[str, str]] = [] - for seq in (getattr(resp, 'sequences', None) or []): - text, reasoning = self._split_think(seq.decoded or '') - choices.append({'content': text, 'reasoning_content': reasoning}) - out.append(choices) - while len(out) < n_inputs: - out.append([]) - return out - - def prompt_logprobs(self, messages: List[Dict[str, Any]]) -> Optional[List]: - from twinkle.data_format import SamplingParams - trajectory = {'messages': messages} - params = SamplingParams(max_tokens=0, prompt_logprobs=1) - try: - responses = self._sampler.sample(trajectory, params) - if responses and responses[0].prompt_logprobs is not None: - return responses[0].prompt_logprobs - return None - except Exception as e: - logger.warning(f'[SamplerBackend] prompt_logprobs failed: {e}') - return None - - def prompt_logprobs_ids(self, input_ids_list: List[List[int]]) -> List[List]: - from twinkle.data_format import SamplingParams - if not isinstance(input_ids_list, list) or not input_ids_list: - raise ValueError('prompt_logprobs_ids requires a non-empty List[List[int]].') - device_mesh = getattr(self._sampler, 'device_mesh', None) - dp_world_size = getattr(device_mesh, 'dp_world_size', 1) or 1 - if len(input_ids_list) < dp_world_size: - raise ValueError(f'SamplerBackend.prompt_logprobs_ids requires at least ' - f'dp_world_size={dp_world_size} inputs to keep all DP workers busy, ' - f'got {len(input_ids_list)}. Batch upstream before calling.') - feats = [{'input_ids': list(ids)} for ids in input_ids_list] - params = SamplingParams(max_tokens=0, prompt_logprobs=1) - responses = self._sampler.sample(feats, params) - return [r.prompt_logprobs for r in responses] - - def embeddings(self, texts: List[str]): - if self._embed_client is None: - raise NotImplementedError('SamplerBackend requires embed_endpoint for embeddings. ' - 'Pass embed_endpoint when constructing SamplerBackend.') - import numpy as np - resp = self._embed_client.post( - self._embed_url, json={ - 'model': self._embed_model, - 'input': texts, - }) - resp.raise_for_status() - data = resp.json().get('data', []) - data_sorted = sorted(data, key=lambda x: x.get('index', 0)) - return np.array([d['embedding'] for d in data_sorted], dtype=np.float32) diff --git a/src/twinkle_agentic/preprocessor/experimental/score_filter.py b/src/twinkle_agentic/preprocessor/experimental/score_filter.py deleted file mode 100644 index 94333d8ef..000000000 --- a/src/twinkle_agentic/preprocessor/experimental/score_filter.py +++ /dev/null @@ -1,835 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Pluggable per-round scorer/filter for SFT key rounds. - -Architecture: - - ScoreFilter(backend, scorers=[...]) - ├── pre-fetches logprobs once if any scorer requires them - ├── runs each Scorer in order, collecting ScoreResult per round - ├── trace dump (per-round JSON, multi_turn-style) - └── AND aggregation: a round is kept iff every scorer returns passed=True. - -Built-in scorers (each is its own class): - ChrMinScorer chr_dist_min_pos. LOW = hard = keep. - SIFDScorer IFD / S-IFD-50 / S-IFD-75. Default observe-only. - PassNScorer Self-rollouts judged by an LLM. extras carry rollouts/verdicts. - ParaphraseScorer chr_min over a model paraphrase produced under GT injection. - -Decoupling: - * key_rounds missing/empty → every assistant turn becomes a candidate round. - * intents=None → no intent-based gating (all rounds processed). -""" -import json -import os -import re -import time -from dataclasses import dataclass, field -from typing import Any, Callable, Dict, Iterable, List, Optional, Protocol, Set, Tuple - -from twinkle.data_format import pack_value, user_data_get -from twinkle.preprocessor import Preprocessor -from twinkle.template import Template -from twinkle.utils import get_logger -from ..logprob_utils import _chr_min_distinct, _ifd_family_metrics, _lp_to_jsonable, _pad_batch, _to_int_list -from .llm_backend import LLMBackend - -logger = get_logger() - -_MIN_RESPONSE_TOKENS = 5 - - -@dataclass -class RoundContext: - """Per-round payload passed to scorers.""" - row_idx: int - rnd_idx: int - asst_idx: int - row: Dict[str, Any] - intent: Optional[str] - messages: List[Dict[str, Any]] - context_messages: List[Dict[str, Any]] - cond_ids: List[int] - n_prompt: int - asst_ids: List[int] - asst_text: str - user_prompt: str - features: Dict[str, Any] = field(default_factory=dict) - - -@dataclass -class ScoreResult: - score: Optional[float] = None - passed: bool = True - extras: Dict[str, Any] = field(default_factory=dict) - - -class Scorer(Protocol): - name: str - requires_logprobs: bool - - def score(self, contexts: List[RoundContext]) -> List[ScoreResult]: - ... - - -def _user_data_lookup(user_data: Any, key: str) -> Any: - """Pull a value by key from packed user_data; returns the JSON-decoded value.""" - return user_data_get(user_data, key) - - -# ============================================================================ -# Built-in scorers -# ============================================================================ - - -class ChrMinScorer: - """chr_dist_min_pos. Dual-threshold: keep samples in [low, high).""" - name = 'chr_min' - requires_logprobs = True - - def __init__(self, threshold: float = 0.47): - self._threshold = float(threshold) - - def score(self, contexts: List[RoundContext]) -> List[ScoreResult]: - out: List[ScoreResult] = [] - for ctx in contexts: - cond_lp = ctx.features.get('cond_lp') - asst_lp = ctx.features.get('asst_lp') - score = _chr_min_distinct( - cond_lp, - asst_lp, - ctx.cond_ids, - ctx.asst_ids, - ctx.n_prompt, - ) - passed = (score is None) or (score < self._threshold) - out.append(ScoreResult( - score=score, - passed=passed, - extras={'threshold': self._threshold}, - )) - return out - - -class SIFDScorer: - """IFD / S-IFD-50 / S-IFD-75. Observation-only by default.""" - name = 'sifd' - requires_logprobs = True - - def __init__(self, ifd_threshold: Optional[float] = None): - # If set, passed = (ifd >= threshold). HIGH IFD = hard = keep. - self._ifd_threshold = ifd_threshold - - def score(self, contexts: List[RoundContext]) -> List[ScoreResult]: - out: List[ScoreResult] = [] - for ctx in contexts: - cond_lp = ctx.features.get('cond_lp') - asst_lp = ctx.features.get('asst_lp') - fam = _ifd_family_metrics(cond_lp, asst_lp, ctx.cond_ids, ctx.asst_ids, ctx.n_prompt) - score = fam.get('ifd') - if self._ifd_threshold is None or score is None: - passed = True - else: - passed = score >= self._ifd_threshold - out.append(ScoreResult(score=score, passed=passed, extras=dict(fam))) - return out - - -_JUDGE_SYSTEM_PROMPT = """\ -You are a strict but fair answer grader. Judge whether the [Model Answer] is acceptable based on the reference answer (Ground Truth). -Evaluate the following three aspects; if any has a major issue, return FAIL: - -1. Computational/factual correctness: whether the final conclusion, numbers, and key factual statements match the reference answer; -2. Reasoning/approach similarity: whether the solution path, key steps, and considered dimensions are close to the reference answer; - For open-ended questions (no single correct answer), assess whether the style, stance, and considered dimensions align with the reference answer; -3. Completeness: the answer is not truncated, ends naturally, and covers all points of the question. - -First give a brief 1-3 sentence justification, then on the last line strictly output: -<verdict>PASS</verdict> or <verdict>FAIL</verdict>""" # noqa - - -class PassNScorer: - """Self-rollouts (n × per round) judged by an LLM.""" - name = 'pass_n' - requires_logprobs = False - - def __init__( - self, - backend: LLMBackend, - judge_api=None, - judge_model: Optional[str] = None, - judge_base_url: Optional[str] = None, - judge_api_key: Optional[str] = None, - judge_client_kwargs: Optional[Dict[str, Any]] = None, - n: int = 4, - min_pass: int = 0, - sample_temperature: float = 0.7, - sample_max_tokens: int = 4096, - judge_temperature: float = 0.0, - judge_max_tokens: int = 512, - judge_max_rollout_chars: int = 8000, - judge_max_workers: int = 8, - ): - self._backend = backend - self._judge_api = self._build_judge_api(judge_api, judge_model, judge_base_url, judge_api_key, - judge_client_kwargs) - self._n = max(1, int(n)) - self._min_pass = int(min_pass) - self._sample_temperature = float(sample_temperature) - self._sample_max_tokens = int(sample_max_tokens) - self._judge_temperature = float(judge_temperature) - self._judge_max_tokens = int(judge_max_tokens) - self._judge_max_rollout_chars = int(judge_max_rollout_chars) - self._judge_max_workers = max(1, int(judge_max_workers)) - if self._judge_api is None: - logger.warning('[PassNScorer] no judge_api configured; rollouts will be sampled ' - 'without verdicts (every round trivially passes).') - - @staticmethod - def _build_judge_api(api, model, base_url, api_key, client_kwargs): - if api is not None: - return api - if not model: - return None - from twinkle_agentic.protocol.openai import OpenAI as OpenAIAPI - return OpenAIAPI(model=model, api_key=api_key, base_url=base_url, client_kwargs=client_kwargs) - - @staticmethod - def _extract_text_from_choice(choice: Any) -> str: - if not isinstance(choice, dict): - return '' - parts: List[str] = [] - rc = choice.get('reasoning_content') - if isinstance(rc, str) and rc.strip(): - parts.append(f'<thinking>\n{rc.strip()}\n</thinking>') - content = choice.get('content') - if isinstance(content, str) and content.strip(): - parts.append(content.strip()) - if parts: - return '\n\n'.join(parts) - return content if isinstance(content, str) else '' - - @staticmethod - def _truncate(text: str, max_chars: int) -> str: - if not isinstance(text, str) or max_chars <= 0 or len(text) <= max_chars: - return text - head = max_chars * 2 // 3 - tail = max_chars - head - 32 - if tail <= 0: - return text[:max_chars] - return text[:head] + '\n\n...[truncated]...\n\n' + text[-tail:] - - @staticmethod - def _parse_verdict(judge_text: str) -> Optional[bool]: - if not isinstance(judge_text, str): - return None - compact = ''.join(judge_text.upper().split()) - has_pass = '<VERDICT>PASS</VERDICT>' in compact - has_fail = '<VERDICT>FAIL</VERDICT>' in compact - if has_pass and not has_fail: - return True - if has_fail and not has_pass: - return False - # Fallback: keyword scan in the tail (last 200 chars, post-compact). - tail = compact[-200:] - if 'PASS' in tail and 'FAIL' not in tail: - return True - if 'FAIL' in tail and 'PASS' not in tail: - return False - return None - - def _judge_one(self, user_prompt: str, gt_text: str, rollout_text: str) -> Tuple[bool, str]: - if self._judge_api is None: - return True, '(no judge configured)' - if not rollout_text or not rollout_text.strip(): - return False, '(empty rollout)' - from twinkle.data_format.sampling import SamplingParams - body = (f'[问题]\n{self._truncate(user_prompt, self._judge_max_rollout_chars)}\n\n' - f'[参考答案]\n{self._truncate(gt_text, self._judge_max_rollout_chars)}\n\n' - f'[模型回答]\n{self._truncate(rollout_text, self._judge_max_rollout_chars)}\n\n' - '请评分。') - trajectory = { - 'messages': [ - { - 'role': 'system', - 'content': _JUDGE_SYSTEM_PROMPT - }, - { - 'role': 'user', - 'content': body - }, - ] - } - sp = SamplingParams( - temperature=self._judge_temperature, - max_tokens=self._judge_max_tokens, - num_samples=1, - ) - # extra_body forwards `enable_thinking=False` so the judge skips CoT. - msg = self._judge_api(trajectory, sp, extra_body={'enable_thinking': False}) - if isinstance(msg, list): - msg = msg[0] if msg else {} - text = msg.get('content', '') if isinstance(msg, dict) else str(msg) - text = text or '' - verdict = self._parse_verdict(text) - # Conservative default: ambiguous verdict → FAIL. - return bool(verdict) if verdict is not None else False, text - - def score(self, contexts: List[RoundContext]) -> List[ScoreResult]: - if not contexts: - return [] - ctx_msgs = [ctx.context_messages for ctx in contexts] - batched = self._backend.chat_batch( - ctx_msgs, - temperature=self._sample_temperature, - max_tokens=self._sample_max_tokens, - n=self._n, - ) or [] - - while len(batched) < len(contexts): - batched.append([]) - - from concurrent.futures import ThreadPoolExecutor - work: List[Tuple[int, int, str, str, str]] = [] - for i, (ctx, choices) in enumerate(zip(contexts, batched)): - if not isinstance(choices, list): - continue - for r_i, choice in enumerate(choices): - rt = self._extract_text_from_choice(choice) - work.append((i, r_i, ctx.user_prompt, ctx.asst_text, rt)) - - verdict_by_round: Dict[int, List[Tuple[int, bool, str]]] = {} - if work and self._judge_api is not None: - - def _do(item): - i, r_i, up, gt, rt = item - ok, raw = self._judge_one(up, gt, rt) - return i, r_i, ok, raw - - with ThreadPoolExecutor(max_workers=self._judge_max_workers) as ex: - for i, r_i, ok, raw in ex.map(_do, work): - verdict_by_round.setdefault(i, []).append((r_i, ok, raw)) - - out: List[ScoreResult] = [] - for i, (ctx, choices) in enumerate(zip(contexts, batched)): - rollouts = [{ - 'rollout_idx': r_i, - 'content': self._extract_text_from_choice(c) - } for r_i, c in enumerate(choices or [])] - verdicts = sorted(verdict_by_round.get(i, []), key=lambda x: x[0]) - judgments = [{'rollout_idx': r_i, 'passed': bool(p), 'judge_raw': raw} for r_i, p, raw in verdicts] - pass_count = sum(1 for _, p, _ in verdicts if p) - score = (pass_count / self._n) if rollouts else None - passed = pass_count >= self._min_pass - out.append( - ScoreResult( - score=score, - passed=passed, - extras={ - 'pass_count': pass_count, - 'n_rollouts': len(rollouts), - 'rollouts': rollouts, - 'judgments': judgments, - 'min_pass': self._min_pass, - }, - )) - - scored = [r for r in out if r.score is not None] - if scored: - avg = sum(r.score for r in scored) / len(scored) - logger.info(f'[PassNScorer] graded {len(scored)}/{len(out)} rounds × {self._n} ' - f'rollouts; avg pass-rate = {avg:.3f}') - return out - - -class ParaphraseScorer: - """Generate a model paraphrase under GT injection, then re-score chr_min.""" - name = 'paraphrase' - # Owns its own logprob fetch on the rewritten asst tokens. - requires_logprobs = False - - def __init__( - self, - backend: LLMBackend, - template: Template, - chr_min_threshold: Optional[float] = None, - prompt_budget: int = 4096, - sample_temperature: float = 0.7, - sample_max_tokens: int = 4096, - max_prompt_tokens: int = 1024, - ): - self._backend = backend - self._template = template - self._threshold = chr_min_threshold - self._prompt_budget = int(prompt_budget) - self._sample_temperature = float(sample_temperature) - self._sample_max_tokens = int(sample_max_tokens) - self._max_prompt_tokens = int(max_prompt_tokens) - - @staticmethod - def _inject_gt(context_messages, gt_text): - msgs = [dict(m) if isinstance(m, dict) else m for m in context_messages] - instr = f"""\ -Below is the reference answer to this question, for your reference only: - -<reference_answer> -{gt_text} -</reference_answer> - -Based on the reference answer above, please provide a complete answer to the preceding question in your own words and reasoning. Output your answer directly; do not repeat the reference answer verbatim.""" # noqa - if msgs and isinstance(msgs[-1], dict) and msgs[-1].get('role') == 'user': - last = dict(msgs[-1]) - last['content'] = (last.get('content') or '') + '\n\n' + instr - msgs[-1] = last - else: - msgs.append({'role': 'user', 'content': instr}) - return msgs - - def _truncate_gt(self, gt_text: str, n_prompt: int) -> Optional[str]: - # 80 = conservative instruction-template overhead. - budget = self._prompt_budget - n_prompt - 80 - if budget < 50: - return None - gt_ids = _to_int_list(self._template.tokenizer(gt_text, add_special_tokens=False)['input_ids']) - if len(gt_ids) <= budget: - return gt_text - return self._template.tokenizer.decode(gt_ids[:budget], skip_special_tokens=False) - - def _encode_prompt(self, ctx_msgs): - ids = _to_int_list(self._template.encode({'messages': list(ctx_msgs)}, add_generation_prompt=True)['input_ids']) - if self._max_prompt_tokens <= 0 or len(ids) <= self._max_prompt_tokens: - return ids - return ids[-self._max_prompt_tokens:] - - def score(self, contexts: List[RoundContext]) -> List[ScoreResult]: - if not contexts: - return [] - - keys: List[int] = [] - augmented: List[List[Dict[str, Any]]] = [] - for i, ctx in enumerate(contexts): - gt = self._truncate_gt(ctx.asst_text, ctx.n_prompt) - if gt is None or not ctx.context_messages: - continue - keys.append(i) - augmented.append(self._inject_gt(ctx.context_messages, gt)) - - out: List[ScoreResult] = [ - ScoreResult(score=None, passed=True, extras={'reason': 'paraphrase skipped'}) for _ in contexts - ] - if not keys: - return out - - batched = self._backend.chat_batch( - augmented, - temperature=self._sample_temperature, - max_tokens=self._sample_max_tokens, - n=1, - ) or [] - - # Re-tokenize against the ORIGINAL (no-GT) context so logprobs reflect - # pure self-conditional probability of the paraphrase. - para_data: Dict[int, Tuple[List[int], int, List[int], str]] = {} - for i, choices in zip(keys, batched): - text = None - if choices: - c0 = choices[0] - if isinstance(c0, dict): - text = c0.get('content') - if not isinstance(text, str) or not text.strip(): - continue - ctx = contexts[i] - prompt_ids = self._encode_prompt(ctx.context_messages) - asst_ids = _to_int_list(self._template.tokenizer(text, add_special_tokens=False)['input_ids']) - if len(asst_ids) < _MIN_RESPONSE_TOKENS + 1: - continue - cond_ids = prompt_ids + asst_ids - para_data[i] = (cond_ids, len(prompt_ids), asst_ids, text) - - if not para_data: - return out - - ordered = list(para_data.keys()) - cond_batch = [para_data[i][0] for i in ordered] - asst_batch = [para_data[i][2] for i in ordered] - cond_lps = self._backend.prompt_logprobs_ids(cond_batch) - asst_lps = self._backend.prompt_logprobs_ids(asst_batch) - - for i, cond_lp, asst_lp in zip(ordered, cond_lps, asst_lps): - cond_ids, n_prompt, asst_ids, text = para_data[i] - score = _chr_min_distinct(cond_lp, asst_lp, cond_ids, asst_ids, n_prompt) - if self._threshold is None or score is None: - passed = True - else: - passed = score < self._threshold - out[i] = ScoreResult( - score=score, - passed=passed, - extras={ - 'paraphrase_text': text, - 'n_prompt': n_prompt, - 'cond_lp': _lp_to_jsonable(cond_lp), - 'asst_lp': _lp_to_jsonable(asst_lp), - 'threshold': self._threshold, - }, - ) - - logger.info(f'[ParaphraseScorer] paraphrased + scored {len(para_data)}/' - f'{len(contexts)} rounds') - return out - - -# ============================================================================ -# ScoreFilter (Preprocessor entry point) -# ============================================================================ - - -class ScoreFilter(Preprocessor): - """Score and filter assistant turns by a pluggable scorer set. - - A round is kept iff every scorer returns ``passed=True``. Rows that lose - all key rounds are dropped (configurable via ``keep_if_no_key_rounds``). - - Decoupling rules: - * `key_rounds` missing/empty in `user_data` → every assistant turn - becomes a candidate round. - * `intents=None` → no intent-based gating. - """ - - def __init__( - self, - template: Template, - backend: LLMBackend, - scorers: List[Scorer], - intents: Optional[Iterable[str]] = None, - keep_if_no_key_rounds: bool = False, - drop_row_on_any_fail: bool = True, - max_prompt_tokens: int = 1024, - trace_dir: Optional[str] = None, - trace_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, - success_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, - ): - super().__init__() - if not isinstance(template, Template): - raise TypeError(f'ScoreFilter requires a `Template` instance, got ' - f'{type(template).__name__}.') - self._template = template - self._backend = backend - self._scorers = list(scorers) - self._intents: Optional[Set[str]] = (None if intents is None else set(intents)) - self._keep_if_no_key_rounds = bool(keep_if_no_key_rounds) - self._drop_row_on_any_fail = bool(drop_row_on_any_fail) - self._max_prompt_tokens = int(max_prompt_tokens) - self._trace_dir = trace_dir - self._trace_callback = trace_callback - self._success_callback = success_callback - if self._trace_dir: - import shutil - if os.path.exists(self._trace_dir): - shutil.rmtree(self._trace_dir) - os.makedirs(self._trace_dir, exist_ok=True) - - def __call__(self, rows): - rows_list = self.map_col_to_row(rows) - contexts = self._build_contexts(rows_list) - dropped: List[Dict[str, Any]] = [] - if contexts: - score_table = self._score_contexts(contexts) - self._log_score_summary(contexts, score_table) - if self._trace_dir: - self._write_traces(contexts, score_table) - rows_list, dropped = self._apply_filter(rows_list, contexts, score_table) - return rows_list, dropped - - def _log_score_summary(self, contexts, score_table): - for scorer in self._scorers: - scores = [ - t[scorer.name].score for t in score_table if scorer.name in t and t[scorer.name].score is not None - ] - if not scores: - continue - n_pass = sum(1 for t in score_table if scorer.name in t and t[scorer.name].passed) - extras_sample = {} - for t in score_table: - if scorer.name in t and t[scorer.name].extras: - extras_sample = t[scorer.name].extras - break - extra_keys = [k for k in extras_sample if k != 'threshold'] - extra_stats = '' - for k in extra_keys: - vals = [ - t[scorer.name].extras.get(k) for t in score_table - if scorer.name in t and t[scorer.name].extras and t[scorer.name].extras.get(k) is not None - ] - if vals and isinstance(vals[0], (int, float)): - avg = sum(vals) / len(vals) - extra_stats += f', {k}_avg={avg:.4f}' - logger.info(f'[ScoreFilter/{scorer.name}] n={len(scores)}, ' - f'mean={sum(scores) / len(scores):.4f}, ' - f'min={min(scores):.4f}, max={max(scores):.4f}, ' - f'pass={n_pass}/{len(score_table)}' - f'{extra_stats}') - - # ---- scoring (inlined DefaultScoreCalculator) -------------------------- - - def _score_contexts(self, contexts: List[RoundContext]) -> List[Dict[str, ScoreResult]]: - if any(getattr(s, 'requires_logprobs', False) for s in self._scorers): - self._attach_logprobs(contexts) - out: List[Dict[str, ScoreResult]] = [dict() for _ in contexts] - for scorer in self._scorers: - results = scorer.score(contexts) - if len(results) != len(contexts): - raise RuntimeError(f'scorer {scorer.name!r} returned {len(results)} results ' - f'for {len(contexts)} contexts') - for i, r in enumerate(results): - out[i][scorer.name] = r - return out - - def _attach_logprobs(self, contexts: List[RoundContext]) -> None: - cond_batch = [ctx.cond_ids for ctx in contexts] - asst_batch = [ctx.asst_ids for ctx in contexts] - floor = self._batch_floor() - cond_padded, n_cond = _pad_batch(cond_batch, floor) - asst_padded, n_asst = _pad_batch(asst_batch, floor) - cond_lps = self._backend.prompt_logprobs_ids(cond_padded)[:n_cond] - asst_lps = self._backend.prompt_logprobs_ids(asst_padded)[:n_asst] - for ctx, c, a in zip(contexts, cond_lps, asst_lps): - ctx.features['cond_lp'] = c - ctx.features['asst_lp'] = a - - def _batch_floor(self) -> int: - sampler = getattr(self._backend, '_sampler', None) - device_mesh = getattr(sampler, 'device_mesh', None) - return getattr(device_mesh, 'dp_world_size', 1) or 1 - - # ---- context construction -------------------------------------------- - - def _build_contexts(self, rows: List[Dict[str, Any]]) -> List[RoundContext]: - out: List[RoundContext] = [] - for ri, row in enumerate(rows): - messages = row.get('messages') if isinstance(row, dict) else None - if not isinstance(messages, list): - continue - user_data = row.get('user_data') if isinstance(row, dict) else None - key_rounds = _user_data_lookup(user_data, 'key_rounds') - if not isinstance(key_rounds, list) or not key_rounds: - key_rounds = [i for i, m in enumerate(messages) if isinstance(m, dict) and m.get('role') == 'assistant'] - for rnd_idx, asst_idx in enumerate(key_rounds): - if not isinstance(asst_idx, int): - continue - intent = self._lookup_intent(row, asst_idx) - if self._intents is not None and intent not in self._intents: - continue - ctx = self._prepare_round(row, messages, ri, rnd_idx, asst_idx, intent) - if ctx is not None: - out.append(ctx) - return out - - def _prepare_round( - self, - row: Dict[str, Any], - messages: List[Dict[str, Any]], - ri: int, - rnd_idx: int, - asst_idx: int, - intent: Optional[str], - ) -> Optional[RoundContext]: - if not (0 <= asst_idx < len(messages)): - return None - asst_msg = messages[asst_idx] - if not isinstance(asst_msg, dict) or asst_msg.get('role') != 'assistant': - return None - asst_text = asst_msg.get('content') or '' - if isinstance(asst_text, list): - asst_text = ' '.join( - p.get('text', '') for p in asst_text if isinstance(p, dict) and p.get('type') == 'text') - if not asst_text.strip(): - return None - context_messages = messages[:asst_idx] - if not context_messages: - return None - prompt_ids = self._encode_prompt_within_budget(context_messages) - # Raw asst_ids (no chat-template wrapping) so cond/asst share byte-equal - # A-token sequences; otherwise chr_min positions desync. - asst_ids = _to_int_list(self._template.tokenizer(asst_text, add_special_tokens=False)['input_ids']) - if len(asst_ids) < _MIN_RESPONSE_TOKENS + 1: - return None - return RoundContext( - row_idx=ri, - rnd_idx=rnd_idx, - asst_idx=asst_idx, - row=row, - intent=intent, - messages=messages, - context_messages=context_messages, - cond_ids=prompt_ids + asst_ids, - n_prompt=len(prompt_ids), - asst_ids=asst_ids, - asst_text=asst_text, - user_prompt=self._render_user_prompt(context_messages), - ) - - def _encode_prompt_within_budget(self, ctx_msgs: List[Dict[str, Any]]) -> List[int]: - ctx = list(ctx_msgs) - ids = _to_int_list(self._template.encode({'messages': ctx}, add_generation_prompt=True)['input_ids']) - budget = self._max_prompt_tokens - if budget <= 0 or len(ids) <= budget: - return ids - has_sys = bool(ctx) and isinstance(ctx[0], dict) and ctx[0].get('role') == 'system' - body_start = 1 if has_sys else 0 - while len(ctx) - body_start > 1: - ctx.pop(body_start) - ids = _to_int_list(self._template.encode({'messages': ctx}, add_generation_prompt=True)['input_ids']) - if len(ids) <= budget: - return ids - # Single message still over budget → keep tail tokens. - return ids[-budget:] - - @staticmethod - def _render_user_prompt(ctx_msgs: List[Dict[str, Any]]) -> str: - parts: List[str] = [] - for m in ctx_msgs: - if not isinstance(m, dict): - continue - role = m.get('role') or 'user' - content = m.get('content', '') - if isinstance(content, list): - content = ' '.join( - p.get('text', '') for p in content if isinstance(p, dict) and p.get('type') == 'text') - if isinstance(content, str) and content.strip(): - parts.append(f'[{role}] {content.strip()}') - return '\n\n'.join(parts) - - @staticmethod - def _lookup_intent(row: Dict[str, Any], asst_idx: int) -> Optional[str]: - user_data = row.get('user_data') if isinstance(row, dict) else None - intents = _user_data_lookup(user_data, 'intents') - if not isinstance(intents, dict): - return None - v = intents.get(asst_idx) - if v is None: - v = intents.get(str(asst_idx)) - return v if isinstance(v, str) else None - - # ---- trace dump (multi_turn-style) ----------------------------------- - - def _write_traces( - self, - contexts: List[RoundContext], - score_table: List[Dict[str, ScoreResult]], - ) -> None: - for i, ctx in enumerate(contexts): - try: - scores = score_table[i] if i < len(score_table) else {} - kept = all(r.passed for r in scores.values()) if scores else True - record = self._build_trace_record(ctx, scores, kept) - if self._trace_callback is not None and not bool(self._trace_callback(record)): - continue - success = (bool(self._success_callback(record)) if self._success_callback is not None else kept) - prefix = 'ok' if success else 'fail' - rid = f'{ctx.row_idx}-{ctx.asst_idx}-{i}-{int(time.time() * 1000)}' - rid = re.sub(r'[^A-Za-z0-9_\-.]+', '_', rid)[:64] - path = os.path.join(self._trace_dir, f'{prefix}-{rid}.json') - with open(path, 'w', encoding='utf-8') as f: - json.dump(record, f, ensure_ascii=False, indent=2, default=str) - except Exception as e: - # Observability must never break filtering; surface the cause. - logger.warning(f'[ScoreFilter] trace dump failed for row={ctx.row_idx} ' - f'asst={ctx.asst_idx}: {e}') - - @staticmethod - def _build_trace_record( - ctx: RoundContext, - scores: Dict[str, ScoreResult], - kept: bool, - ) -> Dict[str, Any]: - return { - 'row_idx': ctx.row_idx, - 'rnd_idx': ctx.rnd_idx, - 'asst_idx': ctx.asst_idx, - 'intent': ctx.intent, - 'messages': ctx.messages, - 'n_prompt': ctx.n_prompt, - 'cond_ids': ctx.cond_ids, - 'asst_ids': ctx.asst_ids, - 'features': { - k: (_lp_to_jsonable(v) if k.endswith('_lp') else v) - for k, v in ctx.features.items() - }, - 'scores': { - name: { - 'score': r.score, - 'passed': r.passed, - 'extras': r.extras - } - for name, r in scores.items() - }, - 'kept': bool(kept), - } - - # ---- aggregation & row reassembly ------------------------------------ - - def _apply_filter( - self, - rows: List[Dict[str, Any]], - contexts: List[RoundContext], - score_table: List[Dict[str, ScoreResult]], - ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: - per_row: Dict[int, Dict[str, Any]] = {} - for i, ctx in enumerate(contexts): - scores = score_table[i] if i < len(score_table) else {} - passed = all(r.passed for r in scores.values()) if scores else True - slot = per_row.setdefault(ctx.row_idx, { - 'kept': [], - 'failed': 0, - }) - if passed: - slot['kept'].append(ctx.asst_idx) - else: - slot['failed'] += 1 - - out: List[Dict[str, Any]] = [] - dropped: List[Dict[str, Any]] = [] - n_removed_rounds = 0 - n_removed_rows = 0 - for ri, row in enumerate(rows): - user_data = row.get('user_data') if isinstance(row, dict) else None - kr_val = _user_data_lookup(user_data, 'key_rounds') - had_key_rounds = isinstance(kr_val, list) and bool(kr_val) - decision = per_row.get(ri) - - if decision is None: - # Row produced no contexts (no asst turns or filtered by intent). - if had_key_rounds and not self._keep_if_no_key_rounds: - n_removed_rows += 1 - dropped.append(dict(row, drop_reason='score_no_context')) - continue - if self._intents is not None and not self._keep_if_no_key_rounds: - n_removed_rows += 1 - dropped.append(dict(row, drop_reason='score_no_context')) - continue - out.append(row) - continue - - n_removed_rounds += decision['failed'] - kept = decision['kept'] - if had_key_rounds: - if not kept: - n_removed_rows += 1 - dropped.append(dict(row, drop_reason='score_all_rounds_failed')) - continue - new_row = dict(row) - # Re-pack key_rounds; keep all other entries as-is (already packed). - rebuilt = [(k, v) for (k, v) in (user_data or []) if k != 'key_rounds'] - rebuilt.append(('key_rounds', pack_value(list(kept)))) - new_row['user_data'] = rebuilt - out.append(new_row) - else: - if decision['failed'] > 0 and self._drop_row_on_any_fail: - n_removed_rows += 1 - dropped.append(dict(row, drop_reason='score_round_failed')) - continue - out.append(row) - - logger.info(f'[ScoreFilter] removed {n_removed_rounds} rounds, ' - f'dropped {n_removed_rows} rows, kept {len(out)}/{len(rows)}') - return out, dropped diff --git a/src/twinkle_agentic/preprocessor/intent_classifier.py b/src/twinkle_agentic/preprocessor/intent_classifier.py index baa04c13e..6d1b21c24 100644 --- a/src/twinkle_agentic/preprocessor/intent_classifier.py +++ b/src/twinkle_agentic/preprocessor/intent_classifier.py @@ -13,10 +13,14 @@ # Reasoning block regex covers both <think> and <thinking> forms. _THINK_BLOCK_RE = re.compile(r'<think(?:ing)?>(.*?)</think(?:ing)?>', re.DOTALL | re.IGNORECASE) -# ── Intent categories (canonical vocabulary lives in intents.py; re-exported) ── -from .intents import (INTENT_CODE, INTENT_COMPLEX_LOGIC, # noqa: F401,E402 - INTENT_MATH, INTENT_OTHER, INTENT_REASONING, - INTENT_TOOL_CALL, INTENT_USER_DISSATISFACTION) +# ── Intent categories ───────────────────────────────────────────────────────── +INTENT_TOOL_CALL = 'tool_call' +INTENT_CODE = 'code' +INTENT_MATH = 'math' +INTENT_COMPLEX_LOGIC = 'complex_logic' +INTENT_REASONING = 'reasoning' +INTENT_USER_DISSATISFACTION = 'user_dissatisfaction' +INTENT_OTHER = 'other' # ── Heuristic patterns ──────────────────────────────────────────────────────── _CODE_BLOCK_RE = re.compile(r'```[\s\S]{10,}?```') diff --git a/src/twinkle_agentic/preprocessor/intents.py b/src/twinkle_agentic/preprocessor/intents.py deleted file mode 100644 index d945e59af..000000000 --- a/src/twinkle_agentic/preprocessor/intents.py +++ /dev/null @@ -1,25 +0,0 @@ -"""Intent category constants (AUDIT A4). - -Sunk out of ``intent_classifier.py`` into this dependency-free module so any -consumer (e.g. the experimental log-prob scorers, or downstream sampling that -reads ``intents`` labels) can reference the vocabulary without importing the -heavier classifier + its regex detectors. -""" - -INTENT_TOOL_CALL = 'tool_call' -INTENT_CODE = 'code' -INTENT_MATH = 'math' -INTENT_COMPLEX_LOGIC = 'complex_logic' -INTENT_REASONING = 'reasoning' -INTENT_USER_DISSATISFACTION = 'user_dissatisfaction' -INTENT_OTHER = 'other' - -ALL_INTENTS = ( - INTENT_TOOL_CALL, - INTENT_CODE, - INTENT_MATH, - INTENT_COMPLEX_LOGIC, - INTENT_REASONING, - INTENT_USER_DISSATISFACTION, - INTENT_OTHER, -) diff --git a/src/twinkle_agentic/preprocessor/label_schema.py b/src/twinkle_agentic/preprocessor/label_schema.py deleted file mode 100644 index f9cc97938..000000000 --- a/src/twinkle_agentic/preprocessor/label_schema.py +++ /dev/null @@ -1,117 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Unified ``user_data`` label envelope (AUDIT A5). - -All scoring / safety / provenance annotations produced by the pipeline are -written into a trajectory's ``user_data`` as ``(key, pack_value(value))`` pairs. -This is the single data contract that lets us decouple *tagging* (mappers that -never drop) from *filtering* (a tail filter that only reads tags), so the whole -pipeline stays a linear ``QualityPreprocessor`` list — no DAG, no cross-module -imports between a filter and the verifier it depends on. - -PyArrow hard constraint ------------------------ -``user_data`` MUST be a ``List[Tuple[str, str]]`` (see -``twinkle/data_format/trajectory.py``). We NEVER put a bare ``dict`` in a row -column: HF ``datasets``' PyArrow backend cannot stably serialize -heterogeneous / nested dicts. Structured values are JSON-encoded to a single -string via :func:`pack_value`; on read :func:`user_data_get` JSON-decodes them. - -Keep this module dependency-light: it only knows the *keys* and thin get/set -helpers, so both preprocessors and (optionally) other modules can share it -without pulling in verifier/segment code. -""" -from __future__ import annotations - -from typing import Any, Dict, List, Optional, Tuple - -from twinkle.data_format import pack_value, user_data_get - -# --------------------------------------------------------------------------- -# Canonical label keys -# --------------------------------------------------------------------------- -# Per-round hard scores, aligned to assistant/round order within the trajectory. -# Value: List[float] in [0, 1]. -KEY_ROUND_SCORES = 'round_scores' -# Per-round gated flags (a critical hard check zeroed the round). Value: List[bool]. -KEY_ROUND_GATED = 'round_gated' - -# Per-segment fused scores. Value: List[float] in [0, 1]. -KEY_SEGMENT_SCORES = 'segment_scores' -# Per-segment score confidence (D7c calibration). Value: List[float] in [0, 1]. -KEY_SEGMENT_CONFIDENCE = 'segment_confidence' - -# Whole-trajectory fused score in [0, 1] and its discrete level. -KEY_TRAJ_SCORE = 'traj_score' -KEY_TRAJ_LEVEL = 'traj_level' -# Aggregate confidence for the trajectory score (D7c). Value: float in [0, 1]. -KEY_TRAJ_CONFIDENCE = 'traj_confidence' - -# Safety score in [0, 1] (D8, higher = safer) + boolean unsafe flag. -KEY_SAFETY_SCORE = 'safety_score' -KEY_SAFETY_UNSAFE = 'safety_unsafe' - -# Provenance blob (D10): dict-like value JSON-encoded (source/teacher/student/ts). -KEY_PROVENANCE = 'provenance' - -# Free-form scoring metadata (short-circuit stats, per-check breakdown, etc.). -KEY_SCORE_META = 'score_meta' - -# Active-learning pre-selection (ValueSelector): a cheap, LLM-free "how worth an -# expensive rubric pass is this row" score in [0, 1], its per-component -# breakdown, and the boolean gate the rubric stage reads to decide whether to -# spend an LLM call on this row (top-fraction by value_score). -KEY_VALUE_SCORE = 'value_score' -KEY_VALUE_META = 'value_meta' -KEY_SELECTED_FOR_RUBRIC = 'selected_for_rubric' - -# Persisted rubric diagnosis for rubric-scored rows: a per-segment verification -# chain (rubric text + per-criterion verdict/reason/fix + raw model output + -# query/segment_text). This is the SFT corpus for distilling a PRM / error-checker -# LoRA — store it so training never has to re-run the (expensive) teacher. -# Value: List[dict], one entry per rubric-scored segment (see TrajectoryScorer). -KEY_RUBRIC_DIAGNOSIS = 'rubric_diagnosis' - - -# --------------------------------------------------------------------------- -# thin get / set helpers over the (key, pack_value) envelope -# --------------------------------------------------------------------------- -def get_user_data(row: Dict[str, Any]) -> List[Tuple[str, str]]: - """Return the row's ``user_data`` as a list (never a dict), defaulting to [].""" - ud = row.get('user_data') - if ud is None: - return [] - if isinstance(ud, list): - return ud - # Be forgiving of a stray dict (e.g. hand-authored rows) — flatten to pairs. - if isinstance(ud, dict): - return [(k, v if isinstance(v, str) else pack_value(v)) for k, v in ud.items()] - return [] - - -def get_label(row: Dict[str, Any], key: str, default: Any = None) -> Any: - """Read+JSON-decode the first label matching ``key`` from ``row['user_data']``.""" - return user_data_get(get_user_data(row), key, default) - - -def set_labels(row: Dict[str, Any], updates: Dict[str, Any]) -> Dict[str, Any]: - """Return a shallow-copied row with ``updates`` merged into ``user_data``. - - Existing entries for the same keys are replaced (last-write-wins), preserving - the original order for untouched keys. Values are packed with :func:`pack_value` - so the column stays ``List[Tuple[str, str]]`` (PyArrow-stable). - """ - if not updates: - return row - existing = get_user_data(row) - replace = set(updates.keys()) - merged: List[Tuple[str, str]] = [(k, v) for (k, v) in existing if k not in replace] - for k, v in updates.items(): - merged.append((k, pack_value(v))) - new_row = dict(row) - new_row['user_data'] = merged - return new_row - - -def set_label(row: Dict[str, Any], key: str, value: Any) -> Dict[str, Any]: - """Convenience: set a single label.""" - return set_labels(row, {key: value}) diff --git a/src/twinkle_agentic/preprocessor/logprob_utils.py b/src/twinkle_agentic/preprocessor/logprob_utils.py deleted file mode 100644 index 5f86e2865..000000000 --- a/src/twinkle_agentic/preprocessor/logprob_utils.py +++ /dev/null @@ -1,231 +0,0 @@ -"""Log-probability data-selection math (IFD / S-IFD / chr_min). - -Split out of ``utils.py`` (AUDIT A2): these helpers are consumed only by the -log-prob based scorers (the experimental ``ScoreFilter`` family). Keeping them -separate from the message-format utilities means editing scoring math never -risks touching the message helpers used across every active cleaning step. -""" -import math -from typing import Any, Dict, List, Optional, Set, Tuple - - -def _extract_logprob(lp, token_id: Optional[int] = None) -> Optional[float]: - if lp is None: - return None - if isinstance(lp, (int, float)): - return float(lp) - if not isinstance(lp, dict): - return None - # vLLM with prompt_logprobs=1 returns top-1 PLUS actual token if they differ; - # actual is appended LAST, so iter-first picks the wrong (top-1) one. - entry = None - if token_id is not None: - entry = lp.get(token_id) - if entry is None: - entry = lp.get(str(token_id)) - if entry is None: - entry = next(iter(lp.values()), None) - if entry is None: - return None - if hasattr(entry, 'logprob'): - return float(entry.logprob) - if isinstance(entry, dict): - v = entry.get('logprob') - return float(v) if v is not None else None - if isinstance(entry, (int, float)): - return float(entry) - return None - - -def _to_int_list(x) -> List[int]: - if hasattr(x, 'tolist'): - return x.tolist() - return list(x) - - -def _chr_min_distinct( - cond_lp: List, - asst_lp: List, - cond_ids: List[int], - asst_ids: List[int], - n_prompt: int, - exclude_ids: Optional[Set[int]] = None, -) -> Optional[float]: - """chr_dist_min_pos: fraction of distinct asst-token ids whose - per-occurrence min(cond_lp - asst_lp) is strictly positive.""" - if not asst_lp or not cond_lp or not asst_ids: - return None - n_a = min(len(asst_lp), len(asst_ids)) - n_c = len(cond_lp) - by_tok: Dict[int, List[float]] = {} - for i in range(n_a): - ci = n_prompt + i - if ci >= n_c: - break - tid = asst_ids[i] - if tid is None: - continue - if exclude_ids is not None and int(tid) in exclude_ids: - continue - a = _extract_logprob(asst_lp[i], tid) - c_tok = cond_ids[ci] if ci < len(cond_ids) else None - c = _extract_logprob(cond_lp[ci], c_tok) - if a is None or c is None: - continue - by_tok.setdefault(int(tid), []).append(c - a) - if not by_tok: - return None - pos = sum(1 for diffs in by_tok.values() if min(diffs) > 0) - return pos / len(by_tok) - - -def _chr_min_weighted( - cond_lp: List, - asst_lp: List, - cond_ids: List[int], - asst_ids: List[int], - n_prompt: int, -) -> Optional[float]: - """Magnitude-weighted chr_min: each distinct token contributes |min_delta| - as weight; returns sum(pos_weights) / sum(all_weights).""" - if not asst_lp or not cond_lp or not asst_ids: - return None - n_a = min(len(asst_lp), len(asst_ids)) - n_c = len(cond_lp) - by_tok: Dict[int, List[float]] = {} - for i in range(n_a): - ci = n_prompt + i - if ci >= n_c: - break - tid = asst_ids[i] - if tid is None: - continue - a = _extract_logprob(asst_lp[i], tid) - c_tok = cond_ids[ci] if ci < len(cond_ids) else None - c = _extract_logprob(cond_lp[ci], c_tok) - if a is None or c is None: - continue - by_tok.setdefault(int(tid), []).append(c - a) - if not by_tok: - return None - total_w = 0.0 - pos_w = 0.0 - for diffs in by_tok.values(): - md = min(diffs) - w = abs(md) - total_w += w - if md > 0: - pos_w += w - if total_w == 0: - return None - return pos_w / total_w - - -def _ifd_family_metrics( - cond_lp: List, - asst_lp: List, - cond_ids: List[int], - asst_ids: List[int], - n_prompt: int, -) -> Dict[str, Any]: - """IFD (Cherry-LLM) and S-IFD-{50,75} (T-SHIRT) for one round.""" - if not asst_lp or not cond_lp or not asst_ids: - return {} - n_a = min(len(asst_lp), len(asst_ids)) - n_c = len(cond_lp) - deltas: List[float] = [] - for i in range(n_a): - ci = n_prompt + i - if ci >= n_c: - break - tid = asst_ids[i] - if tid is None: - continue - a = _extract_logprob(asst_lp[i], tid) - c_tok = cond_ids[ci] if ci < len(cond_ids) else None - c = _extract_logprob(cond_lp[ci], c_tok) - if a is None or c is None: - continue - deltas.append(c - a) - if not deltas: - return {} - n = len(deltas) - mean_delta = sum(deltas) / n - out: Dict[str, Any] = { - 'n_tokens': n, - 'mean_delta': mean_delta, - 'ifd': math.exp(-mean_delta), - } - abs_sorted = sorted(range(n), key=lambda i: abs(deltas[i]), reverse=True) - for k_pct in (50, 75): - keep = max(1, int(round(n * k_pct / 100))) - sub = [deltas[i] for i in abs_sorted[:keep]] - out[f's_ifd_{k_pct}'] = math.exp(-sum(sub) / len(sub)) - return out - - -def _mean_logprob_delta( - cond_lp: List, - asst_lp: List, - cond_ids: List[int], - asst_ids: List[int], - n_prompt: int, -) -> Optional[float]: - """Mean per-token (cond_lp - asst_lp) over the response span.""" - if not asst_lp or not cond_lp or not asst_ids: - return None - n_a = min(len(asst_lp), len(asst_ids)) - n_c = len(cond_lp) - deltas: List[float] = [] - for i in range(n_a): - ci = n_prompt + i - if ci >= n_c: - break - tid = asst_ids[i] - if tid is None: - continue - a = _extract_logprob(asst_lp[i], tid) - c_tok = cond_ids[ci] if ci < len(cond_ids) else None - c = _extract_logprob(cond_lp[ci], c_tok) - if a is None or c is None: - continue - deltas.append(c - a) - if not deltas: - return None - return sum(deltas) / len(deltas) - - -def _lp_to_jsonable(lp_list): - """Convert per-position prompt_logprobs into JSON-safe form.""" - out = [] - for lp in (lp_list or []): - if lp is None: - out.append(None) - continue - if isinstance(lp, (int, float)): - out.append(float(lp)) - continue - if not isinstance(lp, dict): - out.append(repr(lp)) - continue - d = {} - for k, v in lp.items(): - if hasattr(v, 'logprob'): - d[str(k)] = { - 'logprob': float(v.logprob), - 'rank': getattr(v, 'rank', None), - 'decoded': getattr(v, 'decoded_token', None) - } - elif isinstance(v, dict): - d[str(k)] = v - else: - d[str(k)] = repr(v) - out.append(d) - return out - - -def _pad_batch(batch: List[List[int]], floor: int) -> Tuple[List[List[int]], int]: - n = len(batch) - if n >= floor or not batch: - return batch, n - return list(batch) + [batch[-1]] * (floor - n), n diff --git a/src/twinkle_agentic/preprocessor/offline/__init__.py b/src/twinkle_agentic/preprocessor/offline/__init__.py deleted file mode 100644 index 0a14834a6..000000000 --- a/src/twinkle_agentic/preprocessor/offline/__init__.py +++ /dev/null @@ -1,24 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Offline batch-only preprocessors (AUDIT D1 / D2). - -These steps require a GLOBAL view of the dataset and must NOT be dropped into the -per-batch :class:`~twinkle_agentic.preprocessor.QualityPreprocessor` pipeline: - -- :class:`NearDupFilter` (D1) — MinHash-LSH near-duplicate removal; per-batch use - would only compare within a batch, causing severe false negatives. -- :class:`Decontaminator` (D2) — benchmark n-gram overlap removal against a static - index; kept out of the real-time path to avoid false-positive deletions - (defaults to a safe ``'tag'``-friendly design). - -They are deliberately kept out of the main package namespace. Import explicitly:: - - from twinkle_agentic.preprocessor.offline import NearDupFilter, Decontaminator - from twinkle_agentic.preprocessor.offline import build_benchmark_index - -Usage: materialize the dataset to ``List[Dict]``, run these once, then re-wrap -the kept rows before/after the streaming QualityPreprocessor pipeline. -""" -from .decontaminate import Decontaminator, build_benchmark_index # noqa: F401 -from .near_dedup import NearDupFilter # noqa: F401 - -__all__ = ['NearDupFilter', 'Decontaminator', 'build_benchmark_index'] diff --git a/src/twinkle_agentic/preprocessor/offline/decontaminate.py b/src/twinkle_agentic/preprocessor/offline/decontaminate.py deleted file mode 100644 index 7dd8e3831..000000000 --- a/src/twinkle_agentic/preprocessor/offline/decontaminate.py +++ /dev/null @@ -1,109 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Benchmark decontamination via n-gram overlap (AUDIT D2) — OFFLINE ONLY. - -Removes (or tags) training rows that overlap with evaluation benchmarks, so -reported metrics aren't inflated by leakage. Follows the standard 13-gram -overlap recipe (GPT-3 / Llama / Dolma): build an n-gram set from the benchmark -texts once, then flag any row whose text shares an n-gram with it. - -OFFLINE CONTRACT: the benchmark index is static and global; build it once and -reuse across the whole dataset. This is not a per-batch pipeline step — but -unlike near-dup it *is* embarrassingly parallel per row, so it can also run as a -standalone batch pass. Default mode ``'drop'`` removes contaminated rows; -``'tag'`` keeps them and only records a ``contaminated`` label (safer default for -real-time-ish contexts where false positives must never delete data). -""" -from __future__ import annotations - -import re -from typing import Any, Dict, Iterable, List, Set, Tuple - -from twinkle.preprocessor import Preprocessor -from twinkle_agentic.utils.message_utils import msg_content_text -from .. import label_schema as L - -KEY_CONTAMINATED = 'contaminated' - -_WORD_RE = re.compile(r'\w+', re.UNICODE) - - -def _ngrams(text: str, n: int) -> Set[str]: - tokens = _WORD_RE.findall(text.lower()) - if len(tokens) < n: - return set() - return {' '.join(tokens[i:i + n]) for i in range(len(tokens) - n + 1)} - - -def build_benchmark_index(texts: Iterable[str], n: int = 13) -> Set[str]: - """Build a static n-gram set from benchmark texts (build once, reuse).""" - index: Set[str] = set() - for t in texts: - index |= _ngrams(t or '', n) - return index - - -class Decontaminator(Preprocessor): - """Flag/drop rows that share an n-gram with a static benchmark index. - - Args: - benchmark_ngrams: prebuilt index from :func:`build_benchmark_index`. - n: n-gram size (must match the index's n). Default 13. - min_overlap: number of shared n-grams to count as contaminated. - mode: ``'drop'`` removes contaminated rows; ``'tag'`` keeps them and only - writes the ``contaminated`` label (fail-open). - scan: which roles to scan — 'user' (default), 'assistant', or 'all'. - """ - - def __init__( - self, - benchmark_ngrams: Set[str], - *, - n: int = 13, - min_overlap: int = 1, - mode: str = 'drop', - scan: str = 'user', - ): - if mode not in ('drop', 'tag'): - raise ValueError("mode must be 'drop' or 'tag'") - if scan not in ('user', 'assistant', 'all'): - raise ValueError("scan must be 'user', 'assistant', or 'all'") - self.index = benchmark_ngrams or set() - self.n = int(n) - self.min_overlap = int(min_overlap) - self.mode = mode - self.scan = scan - - def _row_text(self, row: Dict[str, Any]) -> str: - messages = row.get('messages') or [] - parts = [] - for m in messages: - if not isinstance(m, dict): - continue - role = m.get('role') - if self.scan == 'all' or role == self.scan: - parts.append(msg_content_text(m)) - return '\n'.join(p for p in parts if p) - - def _overlap(self, row: Dict[str, Any]) -> int: - if not self.index: - return 0 - grams = _ngrams(self._row_text(row), self.n) - if not grams: - return 0 - return len(grams & self.index) - - def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: - rows = self.map_col_to_row(rows) - kept: List[Dict[str, Any]] = [] - dropped: List[Dict[str, Any]] = [] - for row in rows: - overlap = self._overlap(row) - contaminated = overlap >= self.min_overlap - if contaminated and self.mode == 'drop': - dropped.append(dict(row, drop_reason='benchmark_contamination')) - continue - if self.mode == 'tag': - kept.append(L.set_label(row, KEY_CONTAMINATED, contaminated)) - else: - kept.append(row) - return kept, dropped diff --git a/src/twinkle_agentic/preprocessor/offline/near_dedup.py b/src/twinkle_agentic/preprocessor/offline/near_dedup.py deleted file mode 100644 index d8999ebb6..000000000 --- a/src/twinkle_agentic/preprocessor/offline/near_dedup.py +++ /dev/null @@ -1,163 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Near-duplicate removal via MinHash-LSH (AUDIT D1) — OFFLINE ONLY. - -``DedupFilter`` collapses only exact/prefix duplicates; a single edited character -slips through. This adds fuzzy near-dup detection over shingled trajectory text. - -OFFLINE CONTRACT (same as :class:`DedupFilter`): this needs a *global* view of -the dataset — it must see all rows in one ``__call__`` and is NOT a per-batch -``QualityPreprocessor`` step. Running near-dup on a per-batch stream would judge -similarity against only the current batch, causing severe false negatives (and, -if used to drop, unstable results). Materialize the dataset, run this once, then -re-wrap the kept rows. - -Uses ``datasketch`` when installed (fast LSH); otherwise falls back to a pure -O(n²) MinHash comparison — correct but slower, fine for modest offline batches. -""" -from __future__ import annotations - -import hashlib -import re -from typing import Any, Dict, List, Set, Tuple - -from twinkle.preprocessor import Preprocessor -from twinkle.utils import get_logger -from twinkle_agentic.utils.message_utils import msg_content_text - -logger = get_logger() - -_WORD_RE = re.compile(r'\w+', re.UNICODE) - - -def _row_text(row: Dict[str, Any]) -> str: - messages = row.get('messages') or [] - return '\n'.join(msg_content_text(m) for m in messages if isinstance(m, dict)) - - -def _shingles(text: str, k: int) -> Set[str]: - tokens = _WORD_RE.findall(text.lower()) - if len(tokens) < k: - return {' '.join(tokens)} if tokens else set() - return {' '.join(tokens[i:i + k]) for i in range(len(tokens) - k + 1)} - - -def _minhash_signature(shingles: Set[str], num_perm: int) -> List[int]: - """Pure-python MinHash: for each of ``num_perm`` salted hashes, take the min.""" - if not shingles: - return [0] * num_perm - sig: List[int] = [] - for p in range(num_perm): - salt = str(p).encode() - mn = min(int(hashlib.md5(salt + s.encode('utf-8')).hexdigest(), 16) for s in shingles) - sig.append(mn) - return sig - - -class NearDupFilter(Preprocessor): - """Global near-duplicate removal over a fully materialized row collection. - - Args: - threshold: Jaccard similarity at/above which two rows are near-duplicates. - shingle_size: word n-gram size for shingling. - num_perm: MinHash permutations (higher = more accurate, slower). - keep: within a near-dup cluster, keep the ``'longest'`` (most messages) - or ``'first'`` seen row. - """ - - def __init__( - self, - *, - threshold: float = 0.8, - shingle_size: int = 5, - num_perm: int = 128, - keep: str = 'longest', - ): - if keep not in ('longest', 'first'): - raise ValueError("keep must be 'longest' or 'first'") - self.threshold = float(threshold) - self.shingle_size = int(shingle_size) - self.num_perm = int(num_perm) - self.keep = keep - - def __call__(self, rows) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: - rows = self.map_col_to_row(rows) - n = len(rows) - if n <= 1: - return rows, [] - shingle_sets = [_shingles(_row_text(r), self.shingle_size) for r in rows] - - try: - from datasketch import MinHash, MinHashLSH - clusters = self._cluster_lsh(shingle_sets, MinHash, MinHashLSH) - except Exception as e: - logger.info(f'[NearDupFilter] datasketch unavailable ({e}); pure-python O(n^2) fallback.') - clusters = self._cluster_bruteforce(shingle_sets) - - keep_flag = [True] * n - dropped: List[Dict[str, Any]] = [] - for cluster in clusters: - if len(cluster) <= 1: - continue - winner = self._pick_winner(rows, cluster) - for idx in cluster: - if idx != winner: - keep_flag[idx] = False - dropped.append(dict(rows[idx], drop_reason='near_duplicate')) - kept = [rows[i] for i in range(n) if keep_flag[i]] - return kept, dropped - - def _pick_winner(self, rows: List[Dict[str, Any]], cluster: List[int]) -> int: - if self.keep == 'first': - return min(cluster) - return max(cluster, key=lambda i: len(rows[i].get('messages') or [])) - - def _cluster_lsh(self, shingle_sets, MinHash, MinHashLSH) -> List[List[int]]: - lsh = MinHashLSH(threshold=self.threshold, num_perm=self.num_perm) - mh_list = [] - for i, sh in enumerate(shingle_sets): - mh = MinHash(num_perm=self.num_perm) - for s in sh: - mh.update(s.encode('utf-8')) - mh_list.append(mh) - lsh.insert(str(i), mh) - return self._union_find([(i, [int(x) for x in lsh.query(mh_list[i])]) for i in range(len(shingle_sets))], - len(shingle_sets)) - - def _cluster_bruteforce(self, shingle_sets) -> List[List[int]]: - n = len(shingle_sets) - neighbors: List[Tuple[int, List[int]]] = [] - for i in range(n): - adj = [i] - for j in range(i + 1, n): - a, b = shingle_sets[i], shingle_sets[j] - if not a and not b: - continue - inter = len(a & b) - union = len(a | b) or 1 - if inter / union >= self.threshold: - adj.append(j) - neighbors.append((i, adj)) - return self._union_find(neighbors, n) - - @staticmethod - def _union_find(adjacency: List[Tuple[int, List[int]]], n: int) -> List[List[int]]: - parent = list(range(n)) - - def find(x): - while parent[x] != x: - parent[x] = parent[parent[x]] - x = parent[x] - return x - - def union(a, b): - ra, rb = find(a), find(b) - if ra != rb: - parent[rb] = ra - - for i, adj in adjacency: - for j in adj: - union(i, j) - groups: Dict[int, List[int]] = {} - for i in range(n): - groups.setdefault(find(i), []).append(i) - return list(groups.values()) diff --git a/src/twinkle_agentic/preprocessor/provenance.py b/src/twinkle_agentic/preprocessor/provenance.py deleted file mode 100644 index 45b926056..000000000 --- a/src/twinkle_agentic/preprocessor/provenance.py +++ /dev/null @@ -1,73 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Data-lineage / provenance stamping — tag only, never drop (AUDIT D10). - -Industry data pipelines keep provenance so any training example is traceable -back to its source (which dataset, which teacher/student model produced it, when -it was ingested, which cleaning pipeline version touched it). In a self-evolving -distillation loop this is what lets us later attribute a regression to a bad -source or a specific teacher, and to reproduce a training mix. - -This mapper writes a single ``provenance`` blob into ``user_data`` (JSON-packed, -PyArrow-stable via A5). It reads whatever lineage fields already exist on the row -(``model_id`` and any configured passthroughs) and adds an ingest timestamp so -the record is self-describing downstream. -""" -from __future__ import annotations - -import time -from typing import Any, Dict, Sequence - -from twinkle.preprocessor import Mapper - -from . import label_schema as L - - -class ProvenanceStamp(Mapper): - """Stamp each row with a provenance blob in ``user_data`` (never drops). - - Args: - source: a static source/dataset identifier for this ingest batch. - pipeline_version: version string of the cleaning pipeline for audit. - model_field: row field holding the producing model id (default 'model_id'). - extra_fields: additional row fields to copy verbatim into provenance - (e.g. 'teacher_model', 'student_model', 'request_id'). - add_timestamp: include a unix ingest timestamp. Default True. - overwrite: if False, rows that already carry a provenance blob are left - untouched (idempotent re-runs / preserve upstream lineage). Default False. - """ - - def __init__( - self, - *, - source: str = '', - pipeline_version: str = '', - model_field: str = 'model_id', - extra_fields: Sequence[str] = (), - add_timestamp: bool = True, - overwrite: bool = False, - ): - self.source = source - self.pipeline_version = pipeline_version - self.model_field = model_field - self.extra_fields = tuple(extra_fields) - self.add_timestamp = bool(add_timestamp) - self.overwrite = bool(overwrite) - - def map(self, row: Dict[str, Any]) -> Dict[str, Any]: - if not self.overwrite and L.get_label(row, L.KEY_PROVENANCE, None) is not None: - return row - blob: Dict[str, Any] = {} - if self.source: - blob['source'] = self.source - if self.pipeline_version: - blob['pipeline_version'] = self.pipeline_version - model = row.get(self.model_field) - if model: - blob['model'] = model - for f in self.extra_fields: - v = row.get(f) - if v is not None: - blob[f] = v - if self.add_timestamp: - blob['ingested_at'] = int(time.time()) - return L.set_label(row, L.KEY_PROVENANCE, blob) diff --git a/src/twinkle_agentic/preprocessor/structural_noise.py b/src/twinkle_agentic/preprocessor/structural_noise.py deleted file mode 100644 index 257899636..000000000 --- a/src/twinkle_agentic/preprocessor/structural_noise.py +++ /dev/null @@ -1,60 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Keyword-free structural noise-turn tagging (AUDIT D5, optional). - -Existing heartbeat stripping (``message_normalizer``) is keyword-based and, per -the audit, already covers the common OpenHands/OpenClaw formats. This optional -tagger catches *keyword-free* structural noise: near-identical, very short turns -that repeat across the trajectory (polling / retries with no new signal), using -only cheap structural signals (length + exact repetition) — no embeddings, no -LLM. It **tags** a per-trajectory noise ratio into ``user_data`` (never drops), -so a downstream filter can act on it if desired. - -The embedding-distance variant sketched in the audit is deferred until the D1 -near-dup infrastructure (which provides the embedding index) exists. -""" -from __future__ import annotations - -from collections import Counter -from typing import Any, Dict - -from twinkle.preprocessor import Mapper -from twinkle_agentic.utils.message_utils import msg_content_text, normalize_tool_calls -from . import label_schema as L - -KEY_NOISE_RATIO = 'structural_noise_ratio' - - -class StructuralNoiseTagger(Mapper): - """Tag the fraction of assistant turns that are short, repeated boilerplate. - - Args: - short_chars: an assistant turn with visible text at/under this length is a - noise candidate (tool-call turns are exempt — they carry structure). - min_repeat: a candidate counts as noise only if its normalized text recurs - at least this many times across the trajectory's assistant turns. - """ - - def __init__(self, *, short_chars: int = 40, min_repeat: int = 3): - self.short_chars = int(short_chars) - self.min_repeat = int(min_repeat) - - def map(self, row: Dict[str, Any]) -> Dict[str, Any]: - messages = row.get('messages') - if not isinstance(messages, list) or not messages: - return row - asst = [m for m in messages if isinstance(m, dict) and m.get('role') == 'assistant'] - if not asst: - return row - texts = [] - for m in asst: - if normalize_tool_calls(m) is not None: - texts.append(None) # tool-call turn: never noise - else: - texts.append(msg_content_text(m).strip()) - counts = Counter(t for t in texts if t) - noise = 0 - for t in texts: - if t and len(t) <= self.short_chars and counts[t] >= self.min_repeat: - noise += 1 - ratio = noise / len(asst) - return L.set_label(row, KEY_NOISE_RATIO, round(ratio, 6)) diff --git a/src/twinkle_agentic/protocol/openai.py b/src/twinkle_agentic/protocol/openai.py index e0a7f60f0..286609cc0 100644 --- a/src/twinkle_agentic/protocol/openai.py +++ b/src/twinkle_agentic/protocol/openai.py @@ -1,4 +1,6 @@ -from typing import Any, Dict, List, Optional, Union +import threading +from contextlib import nullcontext +from typing import Any, ContextManager, Dict, List, Optional, Union from twinkle.data_format import Trajectory from twinkle.data_format.message import Message @@ -11,6 +13,11 @@ class OpenAI(API): Works with any endpoint speaking the ``/v1/chat/completions`` protocol (OpenAI, Azure OpenAI, vLLM, SGLang, Ollama, ...). + + Requests in flight are capped here rather than by whatever thread pool calls + in. A caller's thread count sizes local parallelism and wants to be large; a + provider's quota belongs to the endpoint and wants to be small. One number + cannot serve both, and only this object knows which endpoint it is talking to. """ def __init__( @@ -18,15 +25,47 @@ def __init__( model: str, api_key: Optional[str] = None, base_url: Optional[str] = None, + *, + concurrency: Optional[int] = None, + timeout: Optional[float] = None, + max_retries: Optional[int] = None, client_kwargs: Optional[Dict[str, Any]] = None, ): + """ + Args: + concurrency: most requests allowed in flight at once, or None for no + cap. The limit is per instance and shared by every thread holding + it, so a module-level client caps the whole process. + timeout: per-request timeout in seconds. Left at the SDK's default + when None. + max_retries: how many times the SDK retries a request it deems + transient -- 429, 5xx, timeouts, dropped connections -- using its + own exponential backoff. Left at the SDK's default when None. + client_kwargs: anything else the ``openai`` constructor accepts. + """ from openai import OpenAI as _OpenAIClient + if concurrency is not None and concurrency < 1: + raise ValueError(f'concurrency must be >= 1 or None, got {concurrency}') + kwargs = dict(client_kwargs or {}) + for name, value in (('timeout', timeout), ('max_retries', max_retries)): + if value is None: + continue + if name in kwargs: + raise ValueError(f'{name} was passed both directly and in client_kwargs; ' + 'drop one so that which value wins is not a matter of ordering') + kwargs[name] = value + self.model = model + self.concurrency = concurrency + # Held across the SDK's own retries too: a request that is backing off + # still occupies the endpoint's attention, so it keeps its slot. + self._slots: ContextManager[Any] = ( + threading.BoundedSemaphore(concurrency) if concurrency is not None else nullcontext()) self._client = _OpenAIClient( api_key=api_key, base_url=base_url, - **(client_kwargs or {}), + **kwargs, ) def __call__( @@ -36,7 +75,8 @@ def __call__( **kwargs, ) -> Union[Message, List[Message]]: request = self._build_request(trajectory, sampling_params, kwargs) - response = self._client.chat.completions.create(**request) + with self._slots: + response = self._client.chat.completions.create(**request) messages = [self._choice_to_message(c) for c in response.choices] return messages[0] if sampling_params.num_samples == 1 else messages diff --git a/src/twinkle_agentic/rollout/__init__.py b/src/twinkle_agentic/rollout/__init__.py index 835d94da0..cddff8eb2 100644 --- a/src/twinkle_agentic/rollout/__init__.py +++ b/src/twinkle_agentic/rollout/__init__.py @@ -1,20 +1,7 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +from .api_sampler import APISampler from .base import Rollout from .bridge import extend_with_bridge -from .factory import build_rollout from .multi_turn import MultiTurnRollout -__all__ = [ - 'APIMultiTurnRollout', - 'MultiTurnRollout', - 'Rollout', - 'build_rollout', - 'extend_with_bridge', -] - - -def __getattr__(name: str): - if name == 'APIMultiTurnRollout': - from .api_multi_turn import APIMultiTurnRollout - return APIMultiTurnRollout - raise AttributeError(f'module {__name__!r} has no attribute {name!r}') +__all__ = ['APISampler', 'MultiTurnRollout', 'Rollout', 'extend_with_bridge'] diff --git a/src/twinkle_agentic/rollout/api_multi_turn.py b/src/twinkle_agentic/rollout/api_multi_turn.py deleted file mode 100644 index f140a2dc5..000000000 --- a/src/twinkle_agentic/rollout/api_multi_turn.py +++ /dev/null @@ -1,314 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Any, Callable, Dict, List, Optional - -from twinkle.data_format import Trajectory -from twinkle.data_format.sampling import SamplingParams -from twinkle_agentic.protocol.base import API -from twinkle_agentic.tools.tool_manager import ToolManager -from .base import Rollout - -# Termination reasons surfaced via ``trajectory['stop_reason']``. -_STOP_NO_TOOL = 'stop' -_STOP_LENGTH = 'length' -_STOP_MAX_TURNS = 'max_turns' -_STOP_API_ERROR = 'api_error' - -# Runaway guard: a ``followup_fn`` is expected to return None eventually. This -# only bounds a callback that never does, so one bad hook cannot spin forever. -_MAX_FOLLOWUPS = 20 - - -class APIMultiTurnRollout(Rollout): - """Multi-turn rollout over an OpenAI-compatible chat-completions API. - - Per-trajectory loop: - 1. POST ``messages + tools`` to the API; receive an assistant message - (``content`` and/or structured ``tool_calls``). - 2. Append the assistant message to ``messages``. - 3. If the assistant emitted ``tool_calls``, dispatch each through the - trajectory-bound :class:`ToolManager`, append one - ``{role:'tool', tool_call_id, content}`` per call, then loop. - 4. Else terminate with ``stop_reason='stop'``. - 5. ``finish_reason='length'`` => terminate with ``stop_reason='length'``. - 6. ``turn >= max_turns`` => terminate with ``stop_reason='max_turns'`` - (and ``truncated=True``). - - After the tool loop ends, if a ``followup_fn`` was passed (per call or at - construction), it is invoked exactly as in :class:`MultiTurnRollout`: it may - append one more user message and buy one more generation whose reply is an - answer, not a tool turn (tools are withdrawn for it), repeating until the - callback returns None. This is what lets a challenger append its check-script - and problem-statement stages onto the same conversation. - - Constructor and per-call override semantics intentionally mirror - :class:`MultiTurnRollout`: ``tool_manager`` may be a single instance - (broadcast) or a list aligned 1:1 with trajectories, and it is optional -- - a challenger inventing tasks has nothing to execute. - - Tool schema source: ``trajectory['tools']`` if present, else - ``tool_manager.tool_infos()`` of the trajectory's manager. Caller is - free to set neither — the API will simply be told there are no tools. - - Output trajectory shape (keys added to the input dict): - * ``messages``: the full conversation including tool turns. - * ``turns``: number of API round-trips actually performed. - * ``stop_reason``: one of ``'stop' | 'length' | 'max_turns' | 'api_error'``. - * ``truncated``: True iff terminated by ``max_turns`` or ``length``. - * ``error``: error string when ``stop_reason == 'api_error'``. - """ - - def __init__( - self, - api: API, - tool_manager: Optional[ToolManager] = None, - sampling_params: Optional[SamplingParams] = None, - max_turns: int = 6, - concurrency: int = 8, - extra_body: Optional[Dict[str, Any]] = None, - trace_dir: Optional[str] = None, - trace_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, - success_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, - ): - super().__init__() - if api is None: - raise ValueError('APIMultiTurnRollout requires an API client') - if concurrency < 1: - raise ValueError(f'concurrency must be >= 1, got {concurrency}') - self._init_common( - max_turns=max_turns, - sampling_params=sampling_params, - trace_dir=trace_dir, - trace_callback=trace_callback, - success_callback=success_callback) - self.api = api - self.tool_manager = tool_manager - self.concurrency = concurrency - self.extra_body = dict(extra_body or {}) - - def __call__( - self, - trajectories: List[Trajectory], - **kwargs, - ) -> List[Trajectory]: - if isinstance(trajectories, dict): - raise TypeError('APIMultiTurnRollout.__call__ expects a List[Trajectory]; ' - 'wrap a single trajectory as [trajectory].') - trajectories = list(trajectories) - n = len(trajectories) - if n == 0: - return [] - - sampling_params: SamplingParams = kwargs.get('sampling_params', self.sampling_params) - tool_managers = self._broadcast(kwargs.get('tool_manager', self.tool_manager), n, name='tool_manager') - extra_body = dict(self.extra_body) - if 'extra_body' in kwargs and kwargs['extra_body']: - extra_body.update(kwargs['extra_body']) - followup_fn = kwargs.get('followup_fn') - - # Per-trajectory thread pool. OpenAI ``/chat/completions`` is - # one-conversation-per-call; concurrency only buys us network - # parallelism, never batched compute. - outs: List[Optional[Trajectory]] = [None] * n - with ThreadPoolExecutor(max_workers=self.concurrency) as pool: - futures = { - pool.submit(self._run_one, trajectories[i], tool_managers[i], sampling_params, extra_body, followup_fn): i - for i in range(n) - } - for fut in as_completed(futures): - i = futures[fut] - outs[i] = fut.result() - - result_outs: List[Trajectory] = [o if o is not None else dict(trajectories[i]) for i, o in enumerate(outs)] - if self.trace_dir: - self._write_rollout_traces(result_outs, global_step=kwargs.get('global_step')) - return result_outs - - # ------------------------------------------------------------------ private - - def _run_one( - self, - trajectory: Trajectory, - tool_manager: Optional[ToolManager], - sampling_params: SamplingParams, - extra_body: Dict[str, Any], - followup_fn: Optional[Callable[[Trajectory, int], Any]] = None, - ) -> Trajectory: - """Drive the API turn loop for a single trajectory. - - Never raises; API failures are encoded in ``stop_reason='api_error'`` - with the exception text in ``error``. This keeps one bad row from - poisoning a whole rollout batch. - """ - messages: List[Dict[str, Any]] = list(trajectory.get('messages') or []) - tools = trajectory.get('tools') - if tools is None and tool_manager is not None: - tools = tool_manager.tool_infos() or None - - turn = 0 - stop_reason = _STOP_MAX_TURNS - truncated = False - error: Optional[str] = None - - while turn < self.max_turns: - turn += 1 - req_traj = {'messages': messages} - if tools: - req_traj['tools'] = list(tools) - try: - reply = self.api( - req_traj, sampling_params, extra_body=extra_body) if extra_body else self.api( - req_traj, sampling_params) - except Exception as exc: - stop_reason = _STOP_API_ERROR - error = f'{type(exc).__name__}: {exc}' - truncated = True - break - - assistant_msg = self._normalise_assistant(reply, turn) - messages.append(assistant_msg) - finish = assistant_msg.get('finish_reason') - tool_calls = assistant_msg.get('tool_calls') or [] - - if finish == 'length': - stop_reason = _STOP_LENGTH - truncated = True - break - if not tool_calls: - stop_reason = _STOP_NO_TOOL - break - - # Skip tool execution at the last turn — results would never be - # consumed by a subsequent API call (consistent with multi_turn.py). - if turn >= self.max_turns: - truncated = True - stop_reason = _STOP_MAX_TURNS - break - - if tool_manager is None: - # Nothing can run the call, so the conversation cannot continue: - # say why rather than looping on an unanswered tool turn. - stop_reason = _STOP_API_ERROR - error = ('model emitted tool_calls but this rollout has no ToolManager; ' - 'pass one at construction time or as a per-call kwarg') - truncated = True - break - - try: - for tc in tool_calls: - response = tool_manager(tc) - messages.append({ - 'role': 'tool', - 'tool_call_id': tc.get('id'), - 'content': str(response), - }) - except Exception as exc: - stop_reason = _STOP_API_ERROR - error = f'ToolExecution {type(exc).__name__}: {exc}' - truncated = True - break - else: - # Loop exited normally => max_turns reached. - truncated = True - stop_reason = _STOP_MAX_TURNS - - # Follow-up stages (check script, problem statement, ...). Each one - # appends a user message and takes one generation whose reply is an - # answer: tools are withdrawn so the model writes rather than calls. - # Skipped entirely on an API error -- the conversation is already broken. - followups = 0 - if followup_fn is not None and stop_reason != _STOP_API_ERROR: - while followups < _MAX_FOLLOWUPS: - view = dict(trajectory) - view['messages'] = messages - view['turns'] = turn - view['stop_reason'] = stop_reason - view['truncated'] = truncated - view['followups'] = followups - followup = followup_fn(view, followups) - if followup is None: - break - text, next_params = (followup if isinstance(followup, tuple) - else (followup, None)) - messages.append({'role': 'user', 'content': text}) - followups += 1 - fu_params = next_params if next_params is not None else sampling_params - try: - reply = (self.api( # tools omitted on purpose: this is an answer - {'messages': messages}, fu_params, extra_body=extra_body) - if extra_body else self.api({'messages': messages}, fu_params)) - except Exception as exc: - stop_reason = _STOP_API_ERROR - error = f'{type(exc).__name__}: {exc}' - truncated = True - break - assistant_msg = self._normalise_assistant(reply, turn + followups) - messages.append(assistant_msg) - # A follow-up that stopped cleanly means the episode was not cut - # off after all, even if the tool phase had hit its turn cap. - if assistant_msg.get('finish_reason') == 'length': - truncated = True - elif stop_reason == _STOP_MAX_TURNS: - stop_reason = _STOP_NO_TOOL - - out = dict(trajectory) - out['messages'] = messages - out['turns'] = turn - out['stop_reason'] = stop_reason - out['truncated'] = truncated - out['followups'] = followups - if error is not None: - out['error'] = error - return out - - @staticmethod - def _normalise_assistant(reply: Any, turn: int) -> Dict[str, Any]: - """Ensure tool_calls have stable ``id``/``type`` fields and strip - message-internal noise that would confuse the next API turn. - - Some OpenAI-compatible servers (vLLM, SGLang) occasionally omit - ``tool_call.id``; the assistant->tool round-trip needs a stable - id to wire ``role:'tool'.tool_call_id`` back to the call site. - """ - if not isinstance(reply, dict): - return {'role': 'assistant', 'content': str(reply)} - msg: Dict[str, Any] = {'role': 'assistant'} - content = reply.get('content') - msg['content'] = content if content is not None else '' - finish = reply.get('finish_reason') - if finish is not None: - msg['finish_reason'] = finish - tool_calls = reply.get('tool_calls') or [] - if tool_calls: - normalised: List[Dict[str, Any]] = [] - for i, tc in enumerate(tool_calls): - tc = dict(tc) - tc.setdefault('id', f'call_{turn}_{i}') - tc.setdefault('type', 'function') - normalised.append(tc) - msg['tool_calls'] = normalised - # Reasoning content is informational only; keep it for trace - # forensics but it is never re-fed to the API. - reasoning = reply.get('reasoning_content') - if reasoning: - msg['reasoning_content'] = reasoning - return msg - - def _build_trace_record( - self, - traj: Dict[str, Any], - *, - idx: int, - success: bool, - ) -> Dict[str, Any]: - """The shared record, plus the two fields only this loop produces. - - ``turns`` counts API round-trips and ``error`` carries the exception - text behind ``stop_reason='api_error'`` -- without it a trace of a - failed batch shows an empty conversation and no reason. - """ - record = super()._build_trace_record(traj, idx=idx, success=success) - record['turns'] = traj.get('turns') - if traj.get('error'): - record['error'] = traj['error'] - return record diff --git a/src/twinkle_agentic/rollout/api_sampler.py b/src/twinkle_agentic/rollout/api_sampler.py new file mode 100644 index 000000000..0d2c025ce --- /dev/null +++ b/src/twinkle_agentic/rollout/api_sampler.py @@ -0,0 +1,133 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Sampler-shaped adapter for external generation APIs.""" + +from typing import Any, Dict, List, Literal, Optional + +from twinkle.data_format import Trajectory +from twinkle.data_format.sampling import SampledSequence, SamplingParams, StopReason +from twinkle.template import Template + +from ..protocol.base import API +from .bridge import _to_plain, encode_appended_turn + +_FINISH_TO_STOP: Dict[Optional[str], StopReason] = { + 'stop': 'stop', + 'length': 'length', + 'tool_calls': 'stop', + 'function_call': 'stop', + 'content_filter': 'abort', +} + + +class APIGenerationError(RuntimeError): + """The endpoint failed before returning a response to validate.""" + + +def _normalise_assistant(reply: Any, turn: int) -> Dict[str, Any]: + """Make an API reply safe to render and feed into the next turn.""" + if not isinstance(reply, dict): + raise TypeError(f'API must return an assistant message dict, got {type(reply).__name__}') + message: Dict[str, Any] = { + 'role': 'assistant', + 'content': reply.get('content') or '', + } + tool_calls = reply.get('tool_calls') or [] + if tool_calls: + normalised = [] + for i, tool_call in enumerate(tool_calls): + tool_call = dict(tool_call) + tool_call.setdefault('id', f'call_{turn}_{i}') + tool_call.setdefault('type', 'function') + normalised.append(tool_call) + message['tool_calls'] = normalised + finish_reason = reply.get('finish_reason') + if finish_reason is not None: + message['finish_reason'] = finish_reason + return message + + +class APISampler: + """Normalize one :class:`API` turn into a :class:`SampledSequence`. + + Holds the local ``template`` (an API turn's text must be tokenised the way + the trainer reads it back, not by the endpoint) and the tool schema the + endpoint should see (a rollout's ``pif`` no longer carries it after encode). + """ + + def __init__( + self, + api: API, + template: Template, + *, + tools: Optional[List[Dict[str, Any]]] = None, + appended_as: Literal['demonstration', 'context'] = 'demonstration', + api_kwargs: Optional[Dict[str, Any]] = None, + ): + """ + Args: + appended_as: how the turn enters training -- ``'demonstration'`` + (scored by SFT, skipped by RL) or ``'context'`` (no loss). + ``'completion'`` is refused: it would claim a per-token log-prob + the API never returns. + api_kwargs: request fields forwarded to every API call. + """ + if appended_as not in ('demonstration', 'context'): + raise ValueError("APISampler appended_as must be 'demonstration' or 'context', " + f'got {appended_as!r}; an API turn has no log-prob to be a completion.') + self.api = api + self.template = template + self.tools = list(tools) if tools else None + self.appended_as = appended_as + self.api_kwargs = dict(api_kwargs or {}) + + def __call__(self, + pif: Dict[str, Any], + sampling_params: Optional[SamplingParams] = None, + **adapter_kwargs) -> SampledSequence: + """Generate one external turn in the callback's normalized shape. + + ``adapter_kwargs`` (``adapter_path`` / ``use_base_model``) name a weight + set the API does not have; they are accepted and ignored so callback code + can forward the same values to either backend. + """ + if sampling_params is None: + sampling_params = SamplingParams() + if sampling_params.num_samples != 1: + raise ValueError('APISampler draws one turn per input; got ' + f'num_samples={sampling_params.num_samples}.') + messages = list(pif.get('messages') or []) + if not messages: + raise ValueError('APISampler needs an encoded prefix carrying its messages; ' + "the pif has no 'messages' to send to the endpoint.") + tools = pif.get('tools') if 'tools' in pif else self.tools + request: Trajectory = {'messages': messages} + if tools: + request['tools'] = list(tools) + + try: + reply = self.api(request, sampling_params, **self.api_kwargs) + except Exception as exc: + raise APIGenerationError(f'{type(exc).__name__}: {exc}') from exc + if isinstance(reply, list): + raise TypeError('APISampler expects one message per turn but the API returned a ' + 'list; num_samples > 1 is rejected above, so this is an API bug.') + turn = sum(message.get('role') == 'assistant' for message in messages) + 1 + reply = _normalise_assistant(reply, turn) + + new_tokens = encode_appended_turn(messages, reply, self.template, tools) + new_input_feature = _to_plain( + self.template.concat_input_feature( + pif, new_tokens, appended_as=self.appended_as, tool_calls=reply.get('tool_calls'))) + # concat_input_feature reconstructs content by decoding ``new_tokens``; + # those include the template's rendered tool-call block. Keep the API's + # original content beside its structured calls instead of duplicating it. + assistant_message = {key: reply[key] for key in ('role', 'content', 'tool_calls') if key in reply} + new_input_feature['messages'][-1] = assistant_message + + return SampledSequence( + stop_reason=_FINISH_TO_STOP.get(reply.get('finish_reason'), 'stop'), + tokens=new_tokens, + logprobs=None, + decoded=self.template.decode(new_tokens), + new_input_feature=new_input_feature, + ) diff --git a/src/twinkle_agentic/rollout/base.py b/src/twinkle_agentic/rollout/base.py index f0b6a44a0..dfaa9ae50 100644 --- a/src/twinkle_agentic/rollout/base.py +++ b/src/twinkle_agentic/rollout/base.py @@ -4,25 +4,41 @@ import re import time from abc import ABC, abstractmethod -from typing import Any, Callable, Dict, List, Optional +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Any, Callable, Dict, List, Optional, Tuple from twinkle.data_format import Trajectory, user_data_get from twinkle.data_format.sampling import SamplingParams from .bridge import _to_plain +# Termination reasons surfaced via ``trajectory['stop_reason']``. The sampler +# path takes the first three from the sampler itself; the API path has to name +# them, and one vocabulary for both is what lets a consumer read either. +STOP_NO_TOOL = 'stop' +STOP_LENGTH = 'length' +STOP_MAX_TURNS = 'max_turns' +STOP_GENERATION_ERROR = 'generation_error' + +# Runaway guard: a ``followup_fn`` is expected to return None eventually. This +# only bounds a callback that never does, so one bad hook cannot spin forever. +MAX_FOLLOWUPS = 20 + class Rollout(ABC): """A batch of trajectories in, the same batch with the model's turns appended. - Implementations differ in where the turns come from -- a local sampler, - whose token ids are spliced into the trajectory, or an HTTP endpoint, which - only ever returns text -- and the difference is real enough that they stay - separate classes: only one of them produces something trainable. + The concrete multi-turn loop may source each assistant turn from a local + sampler or an HTTP endpoint. Everything independent of that choice lives + here: option validation, spreading a per-call argument over the batch, the + thread pool that runs episodes, and trace dumping. - Everything that is *not* generation is here: option validation, spreading a - per-call argument over the batch, and the trace dump. It moved up because - the two implementations had drifted into sharing it by reaching across the - class boundary for each other's underscore methods. + One episode per thread, and a subclass only writes the episode. Both + backends are latency-bound on something that is not the caller's CPU -- an + HTTP round trip, a sandbox, a sampler that routes each request to whichever + worker is free -- so the threads overlap the waiting. Nothing crosses + between episodes, which is what makes the pool safe and also what the old + lockstep loop had to give up: there, one slow sandbox round trip held up the + next generation for every trajectory in the batch. """ # Set by _init_common. Declared at class level so a subclass that does its @@ -33,10 +49,7 @@ class boundary for each other's underscore methods. trace_dir: Optional[str] = None trace_callback: Optional[Callable[[Dict[str, Any]], bool]] = None success_callback: Optional[Callable[[Dict[str, Any]], bool]] = None - - @abstractmethod - def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory]: - raise NotImplementedError() + concurrency: Optional[int] = None # ------------------------------------------------------------------ setup @@ -45,6 +58,7 @@ def _init_common( *, max_turns: int, sampling_params: Optional[SamplingParams] = None, + concurrency: Optional[int] = None, trace_dir: Optional[str] = None, trace_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, success_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, @@ -59,21 +73,93 @@ def _init_common( # passing the trajectory several times instead. raise ValueError(f'{type(self).__name__} supports num_samples=1 only, ' f'got {sp.num_samples}') + if concurrency is not None and concurrency < 1: + raise ValueError(f'concurrency must be >= 1 or None, got {concurrency}') self.max_turns = max_turns self.sampling_params = sp + # None means one thread per trajectory. A cap below the batch size costs + # throughput rather than buying safety, so it has to be asked for. + self.concurrency = concurrency self.trace_dir = trace_dir self.trace_callback = trace_callback self.success_callback = success_callback if trace_dir: os.makedirs(trace_dir, exist_ok=True) + # ------------------------------------------------------------------- drive + + def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory]: + """Run one episode per trajectory and return them in the input order. + + Order is restored from the future map rather than from completion order, + because callers pair the result with their own list positionally -- a + GRPO group is a slice of this list. + """ + if isinstance(trajectories, dict): + raise TypeError(f'{type(self).__name__}.__call__ expects a List[Trajectory]; ' + 'wrap a single trajectory as [trajectory].') + trajectories = list(trajectories) + n = len(trajectories) + if n == 0: + return [] + + ctx = self._resolve_call(kwargs, n) + outs: List[Optional[Trajectory]] = [None] * n + workers = min(n, self.concurrency or n) + if workers == 1: + # No pool for a single episode: a thread would only make the + # traceback of a failing one harder to read. + outs = [self._run_one(trajectories[i], i, ctx) for i in range(n)] + else: + with ThreadPoolExecutor(max_workers=workers) as pool: + futures = {pool.submit(self._run_one, trajectories[i], i, ctx): i for i in range(n)} + for fut in as_completed(futures): + outs[futures[fut]] = fut.result() + + result: List[Trajectory] = [o if o is not None else dict(trajectories[i]) for i, o in enumerate(outs)] + if self.trace_dir: + self._write_rollout_traces(result, global_step=kwargs.get('global_step')) + return result + + @abstractmethod + def _run_one(self, trajectory: Trajectory, index: int, ctx: Dict[str, Any]) -> Trajectory: + """One trajectory, start to finish, in its own thread. + + ``ctx`` is whatever ``_resolve_call`` produced; ``index`` is the + trajectory's position in the batch, which is how per-trajectory entries + in ``ctx`` are addressed. + """ + raise NotImplementedError() + + def _resolve_call(self, kwargs: Dict[str, Any], n: int) -> Dict[str, Any]: + """Fold per-call ``**kwargs`` over the constructor defaults, once. + + Done before the pool starts so a bad argument raises from the caller's + frame instead of inside n threads, and so ``_broadcast`` runs once + rather than per episode. + """ + return {} + @staticmethod - def _broadcast(arg, n: int, *, name: str, required: bool = False) -> List[Any]: + def _unpack_followup(followup: Any) -> Tuple[str, Optional[SamplingParams]]: + """``followup_fn`` may answer with text, or text plus its own budget.""" + if isinstance(followup, tuple): + text, params = followup + return text, params + return followup, None + + @staticmethod + def _broadcast(arg, n: int, *, name: str, required: bool = False, per_trajectory: bool = False) -> List[Any]: """One value shared by the batch, or a list already aligned 1:1 with it. A list of the wrong length is refused rather than zipped short: the mismatch would silently pair trajectories with the wrong tool manager, which reads downstream as a model that used the wrong sandbox. + + ``per_trajectory`` refuses to share one instance across a batch at all. + It is for arguments that carry episode state: episodes now run in + parallel threads, so a shared one would have several conversations + writing to the same object instead of merely interleaving in it. """ if arg is None: if required: @@ -85,6 +171,10 @@ def _broadcast(arg, n: int, *, name: str, required: bool = False) -> List[Any]: raise ValueError(f'per-call {name} list length ({len(arg)}) does ' f'not match number of trajectories ({n})') return list(arg) + if per_trajectory and n > 1: + raise ValueError(f'{name} holds per-episode state and cannot be shared by ' + f'{n} trajectories running in parallel threads: pass a list ' + f'of {n}, one per trajectory.') return [arg] * n # ------------------------------------------------------------------ trace @@ -92,6 +182,7 @@ def _broadcast(arg, n: int, *, name: str, required: bool = False) -> List[Any]: _TRACE_SKIP_KEYS = ( 'input_ids', 'labels', + 'completion_mask', 'attention_mask', 'position_ids', 'logprobs', diff --git a/src/twinkle_agentic/rollout/bridge.py b/src/twinkle_agentic/rollout/bridge.py index fa8bffa1b..ce668d54b 100644 --- a/src/twinkle_agentic/rollout/bridge.py +++ b/src/twinkle_agentic/rollout/bridge.py @@ -1,16 +1,21 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Shared, pure bridge-token stitching logic for multi-turn rollouts. - -This module hosts :func:`extend_with_bridge`, a ``self``-free function that -appends tool messages and the next generation prompt to a running -``InputFeature`` (``pif``) as ``-100`` "bridge" tokens. It is shared between -the core-library ``MultiTurnRollout`` and the client-side rollout so the two -paths cannot drift. - -The logic was lifted verbatim from ``MultiTurnRollout._extend_with_bridge`` and -``MultiTurnRollout._append_bridge_tokens``; every ``self.template`` access was -rewritten to use the ``template`` parameter. No Ray decorators -(``@remote_function`` / ``@remote_class``) are applied here. +"""Shared, pure template-space stitching logic for multi-turn rollouts. + +This module hosts ``self``-free functions that grow a running ``InputFeature`` +(``pif``) one turn at a time, all measuring what a turn adds by diffing rendered +chat-template output rather than by pasting special tokens together: + +* :func:`extend_with_bridge` appends tool messages and the next generation + prompt as ``-100`` "bridge" tokens. +* :func:`encode_appended_turn` returns the tokens an assistant turn written + outside the sampler (an API, a human) contributes. + +The bridge logic was lifted verbatim from ``MultiTurnRollout._extend_with_bridge`` +and ``MultiTurnRollout._append_bridge_tokens``; every ``self.template`` access was +rewritten to use the ``template`` parameter. It is shared between the +core-library ``MultiTurnRollout`` and the client-side rollout so the two paths +cannot drift. No Ray decorators (``@remote_function`` / ``@remote_class``) are +applied here. """ @@ -48,39 +53,34 @@ def _to_plain(obj: Any) -> Any: return obj -def extend_with_bridge( - pif: Dict[str, Any], - tool_messages: List[Dict[str, Any]], +def _delta_text( template: Template, -) -> Optional[Dict[str, Any]]: - """Append tool messages and the next generation prompt as -100 bridge. - - Strategy: compute the bridge ENTIRELY in template space. Render - ``messages_before`` and ``messages_before + tool_messages`` with the - same chat template and take ``s_after[len(s_before):]`` as the delta. - - We deliberately do NOT diff against ``tokenizer.decode(pif.input_ids)`` - because raw vLLM output and canonical template rendering differ in - whitespace (e.g. Qwen inserts ``\\n\\n`` between assistant content and - a ``<tool_call>`` block, while the model generates only ``\\n``). Such - cosmetic divergences would break a ``startswith`` alignment but do not - affect training correctness: history tokens stay in ``pif.input_ids`` - verbatim; only the newly appended bridge is tokenized from the - canonical template output. - - Returns ``None`` when the trajectory exceeds ``max_length`` and the - template's truncation strategy is ``'delete'``. + messages_before: List[Dict[str, Any]], + appended: List[Dict[str, Any]], + *, + gen_prompt_before: bool, + gen_prompt_after: bool, + tools: Optional[List[Dict[str, Any]]] = None, +) -> str: + """Text the chat template adds when ``appended`` is tacked onto history. + + ``gen_prompt_*`` place the delta relative to the generation prompt: a bridge + ends on one (``False -> True``), a completion consumes one + (``True -> False``). """ tokenizer = template.tokenizer + enable_thinking = getattr(template, 'enable_thinking', False) - messages_before = list(pif.get('messages') or []) - messages_after = messages_before + list(tool_messages) + def render(messages: List[Dict[str, Any]], add_generation_prompt: bool) -> str: + return tokenizer.apply_chat_template( + messages, + tools=tools or None, + tokenize=False, + add_generation_prompt=add_generation_prompt, + enable_thinking=enable_thinking) - enable_thinking = getattr(template, 'enable_thinking', False) - s_before = tokenizer.apply_chat_template( - messages_before, tokenize=False, add_generation_prompt=False, enable_thinking=enable_thinking) - s_after = tokenizer.apply_chat_template( - messages_after, tokenize=False, add_generation_prompt=True, enable_thinking=enable_thinking) + s_before = render(messages_before, gen_prompt_before) + s_after = render(list(messages_before) + list(appended), gen_prompt_after) if not s_after.startswith(s_before): # Appending a *user* message moves where Qwen3's template thinks the @@ -101,28 +101,86 @@ def extend_with_bridge( # What stays on record is the history as generated, thinking included -- # those are the tokens the policy read back when it produced the next # turn, and a later training step has to see the same. - s_anchor = tokenizer.apply_chat_template( - _ANCHOR, tokenize=False, add_generation_prompt=False, - enable_thinking=enable_thinking) - s_anchor_after = tokenizer.apply_chat_template( - _ANCHOR + list(tool_messages), tokenize=False, add_generation_prompt=True, - enable_thinking=enable_thinking) + s_anchor = render(_ANCHOR, gen_prompt_before) + s_anchor_after = render(_ANCHOR + list(appended), gen_prompt_after) if not s_anchor_after.startswith(s_anchor): raise RuntimeError('Canonical chat_template output for messages_after is not a ' 'prefix-extension of messages_before, and the same is true ' - 'of a one-message stand-in history; cannot compute bridge ' + 'of a one-message stand-in history; cannot compute the ' 'delta. This indicates the template is non-monotonic in the ' 'message list (e.g. reorders / rewrites earlier turns).\n' f's_before tail: {s_before[-80:]!r}\n' f's_after at same offset: ' f'{s_after[max(0, len(s_before) - 80):len(s_before) + 80]!r}') s_before, s_after = s_anchor, s_anchor_after - bridge_text = s_after[len(s_before):] + return s_after[len(s_before):] + + +def encode_appended_turn( + messages_before: List[Dict[str, Any]], + message: Dict[str, Any], + template: Template, + tools: Optional[List[Dict[str, Any]]] = None, +) -> List[int]: + """Tokens an assistant turn authored elsewhere contributes to the sequence. + + A sampler returns the ids it generated; an API returns text, whose tokens are + only part of the turn -- the template also writes the turn terminator and + whatever follows it. Diffing the rendered template recovers those without + naming a single special token, so this holds for any chat template. + + The result is what :meth:`Template.concat_input_feature` expects as + ``new_tokens``, and is token-for-token what :meth:`Template.encode` would + have produced for the same conversation. + """ + delta = _delta_text( + template, + messages_before, [template.decode_tool_calls(message)], + gen_prompt_before=True, + gen_prompt_after=False, + tools=tools) + if not delta: + raise RuntimeError(f'Appending {message.get("role")!r} turn added no text; ' + 'the chat template dropped it entirely.') + tokens = template.tokenizer.encode(delta, add_special_tokens=False) + if not tokens: + raise RuntimeError(f'Appended turn tokenised to an empty id list: {delta!r}') + return tokens + + +def extend_with_bridge( + pif: Dict[str, Any], + tool_messages: List[Dict[str, Any]], + template: Template, +) -> Optional[Dict[str, Any]]: + """Append tool messages and the next generation prompt as -100 bridge. + + Strategy: compute the bridge ENTIRELY in template space. Render + ``messages_before`` and ``messages_before + tool_messages`` with the + same chat template and take ``s_after[len(s_before):]`` as the delta. + + We deliberately do NOT diff against ``tokenizer.decode(pif.input_ids)`` + because raw vLLM output and canonical template rendering differ in + whitespace (e.g. Qwen inserts ``\\n\\n`` between assistant content and + a ``<tool_call>`` block, while the model generates only ``\\n``). Such + cosmetic divergences would break a ``startswith`` alignment but do not + affect training correctness: history tokens stay in ``pif.input_ids`` + verbatim; only the newly appended bridge is tokenized from the + canonical template output. + + Returns ``None`` when the trajectory exceeds ``max_length`` and the + template's truncation strategy is ``'delete'``. + """ + messages_before = list(pif.get('messages') or []) + messages_after = messages_before + list(tool_messages) + + bridge_text = _delta_text( + template, messages_before, tool_messages, gen_prompt_before=False, gen_prompt_after=True) if not bridge_text: raise RuntimeError('Bridge text computation returned empty string; ' 'tool turn would add no tokens (template misconfiguration?).') - bridge_ids = tokenizer.encode(bridge_text, add_special_tokens=False) + bridge_ids = template.tokenizer.encode(bridge_text, add_special_tokens=False) if not bridge_ids: raise RuntimeError(f'Bridge text tokenised to empty id list: {bridge_text!r}') @@ -142,8 +200,10 @@ def _append_bridge_tokens( """Append bridge tokens with labels = -100. Mirrors the unroll-append-reroll pattern of - :meth:`Template.concat_input_feature` so that ``labels`` semantics - stay consistent with the sampler-produced pif. + :meth:`Template.concat_input_feature` so that ``labels`` and + ``completion_mask`` semantics stay consistent with the sampler-produced + pif. Bridge tokens are nobody's completion -- neither scored nor + log-prob-bearing -- so both fields are appended as zeros. Shallow copy is deliberately used: every mutation below is a top-level key reassignment, never an in-place change to nested @@ -164,12 +224,15 @@ def _append_bridge_tokens( labels = labels[-1:] + labels[:-1] else: labels = [-100] * len(input_ids) + completion_mask = template._prefix_completion_mask(result, labels) input_ids = input_ids + list(bridge_ids) labels = labels + [-100] * len(bridge_ids) + completion_mask = completion_mask + [0] * len(bridge_ids) result['input_ids'] = input_ids result['labels'] = labels + result['completion_mask'] = completion_mask if 'mm_token_type_ids' in result: import torch diff --git a/src/twinkle_agentic/rollout/factory.py b/src/twinkle_agentic/rollout/factory.py deleted file mode 100644 index eb6494d42..000000000 --- a/src/twinkle_agentic/rollout/factory.py +++ /dev/null @@ -1,74 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""One call that turns a generation backend into a :class:`Rollout`. - -Callers that only want turns appended to trajectories -- a challenger inventing -tasks, an evaluation script -- should not have to know that a local sampler and -an HTTP endpoint are driven by different classes with different required -arguments. They ask for a rollout, hand over whichever backend they happen to -have, and get something with the same contract: - - List[Trajectory] -> List[Trajectory] - -What the two still differ in is what ends up *inside* the trajectory, and no -factory can paper over it: the sampler path keeps ``input_ids`` / ``labels`` / -``logprobs`` and is the only one whose output can be trained on, while the API -path returns messages only. Pick the backend accordingly. -""" -from typing import Any, Dict, Optional - -from twinkle.data_format.sampling import SamplingParams -from .base import Rollout - -__all__ = ['build_rollout'] - - -def build_rollout( - backend: Any, - *, - template: Any = None, - tool_manager: Any = None, - sampling_params: Optional[SamplingParams] = None, - max_turns: int = 6, - trace_dir: Optional[str] = None, - **backend_kwargs: Any, -) -> Rollout: - """Build the multi-turn rollout that matches ``backend``. - - Args: - backend: an :class:`twinkle_agentic.protocol.base.API` (any - OpenAI-compatible endpoint) or a sampler exposing ``sample()``. - template: required for a sampler, rejected for an API. The sampler path - continues a conversation by splicing token ids, which needs the - local chat template; the API path re-sends messages as text. - tool_manager: optional for both. Without one the model is told there - are no tools. - backend_kwargs: passed straight to the chosen class -- e.g. ``harness`` - and ``max_trajectory_tokens`` for a sampler, ``concurrency`` and - ``extra_body`` for an API. An argument meant for the other backend - surfaces as a TypeError naming it. - """ - from twinkle_agentic.protocol.base import API - - common: Dict[str, Any] = { - 'tool_manager': tool_manager, - 'sampling_params': sampling_params, - 'max_turns': max_turns, - 'trace_dir': trace_dir, - } - - if isinstance(backend, API): - if template is not None: - raise ValueError('template is only used by the sampler path; an API ' - 'backend re-sends messages as text and never encodes ' - 'them locally.') - from .api_multi_turn import APIMultiTurnRollout - return APIMultiTurnRollout(api=backend, **common, **backend_kwargs) - - if not hasattr(backend, 'sample'): - raise TypeError(f'backend must be an API client or a sampler with a sample() ' - f'method, got {type(backend).__name__}') - if template is None: - raise ValueError('a sampler backend needs a template: the rollout appends each ' - 'turn as token ids and cannot re-encode the history.') - from .multi_turn import MultiTurnRollout - return MultiTurnRollout(sampler=backend, template=template, **common, **backend_kwargs) diff --git a/src/twinkle_agentic/rollout/multi_turn.py b/src/twinkle_agentic/rollout/multi_turn.py index cbeb8f380..6c9fd03d2 100644 --- a/src/twinkle_agentic/rollout/multi_turn.py +++ b/src/twinkle_agentic/rollout/multi_turn.py @@ -1,20 +1,49 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import json import re -from collections import defaultdict -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Any, Callable, Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, List, Literal, Optional, Tuple from twinkle.data_format import Trajectory -from twinkle.data_format.sampling import SampleResponse, SamplingParams +from twinkle.data_format.sampling import SampledSequence, SampleResponse, SamplingParams from twinkle.infra import remote_class, remote_function from twinkle.template.base import Template from twinkle_agentic.harness.base import AgentHarness +from twinkle_agentic.protocol.base import API from twinkle_agentic.tools.tool_manager import ToolManager -from .base import Rollout +from .api_sampler import APIGenerationError, APISampler +from .base import MAX_FOLLOWUPS, STOP_GENERATION_ERROR, Rollout from .bridge import _to_plain, extend_with_bridge +ResponseCallback = Callable[..., SampledSequence] + + +def _default_response_callback(sampler, api, sampling_params, *, input_feature, adapter_kwargs, + **kwargs) -> SampledSequence: + """Use the sampler when present, otherwise the API adapter.""" + if sampler is None: + if api is None: + raise ValueError('response_callback was omitted, but no sampler or API was provided') + return api(input_feature, sampling_params, **adapter_kwargs) + responses = sampler.sample([input_feature], sampling_params=sampling_params, **adapter_kwargs) + if not isinstance(responses, list): + raise TypeError(f'expected List[SampleResponse] from sampler.sample, got ' + f'{type(responses).__name__}') + if len(responses) != 1: + raise RuntimeError(f'sampler returned {len(responses)} responses for a single request; ' + 'expected exactly one.') + response = responses[0] + if not isinstance(response, SampleResponse): + raise TypeError(f'expected SampleResponse from sampler.sample, got ' + f'{type(response).__name__}') + if len(response.sequences) != 1: + raise RuntimeError(f'SampleResponse contains {len(response.sequences)} sequences; expected exactly one.') + sequence = response.sequences[0] + if not isinstance(sequence, SampledSequence): + raise TypeError(f'expected SampledSequence, got {type(sequence).__name__}') + return sequence + + def _append_only_delta( old_messages: List[Dict[str, Any]], new_messages: List[Dict[str, Any]], @@ -111,63 +140,118 @@ def _malformed_tool_message(errors: List[str]) -> Dict[str, Any]: @remote_class() class MultiTurnRollout(Rollout): - """Agentic multi-turn rollout with tool use (batched). + """Agentic multi-turn rollout with tool use, one episode per thread. Contract (matches :class:`Rollout`): accepts a ``List[Trajectory]`` and returns a ``List[Trajectory]`` of the same length, in the same order. - Every turn issues a SINGLE batched ``sampler.sample(active_pifs)`` call - so vLLM can run all live trajectories in parallel; finished trajectories - are parked and excluded from subsequent batches. Per-trajectory loop:: harness.before_generate # append-only after the first encode - sampler.sample(batch) # keep seq.new_input_feature + response_callback(...) # sampler or API -> SampledSequence harness.after_generate - ToolManager.call_many # Env.step_batch when tools share an Env + ToolManager.call_many # this turn's calls, one Env round trip harness.after_tools # format observations as tool messages extend_with_bridge # labels=-100; never decode-reencode history + Each trajectory runs its whole loop in its own thread. The callback may route + each turn to the sampler or the API adapter; either can overlap with other + trajectories while its thread waits on a GPU worker, endpoint, or sandbox. + + A supplied sampler must declare ``sample`` with ``enable_continous_work``. + Without it, ``slice_dp`` spreads each single-request call over every worker + and raises on ranks that receive nothing. + + Shared state: ``sampler``, API client and ``template`` are read-only during a + rollout and safe to share. A ``harness`` is not -- an ms-agent one delegates to an + ``LLMAgent`` that holds memory and context of its own -- so a batch of more + than one trajectory has to be given a 1:1 list of them; a single instance is + refused rather than shared. + Per-call overrides via ``**kwargs``: - * ``sampling_params``: shared :class:`SamplingParams` for the batch. + * ``sampling_params``: :class:`SamplingParams` for every episode. + * ``response_callback``: chooses a backend for each assistant turn and + returns one :class:`SampledSequence`. * ``tool_manager``: a single :class:`ToolManager` or a 1:1 list. - * ``harness``: a single :class:`AgentHarness` or a 1:1 list. Framework - specifics (ms-agent system/memory/tool-message shape) live in the - harness subclass, not here. + * ``harness``: a 1:1 list of :class:`AgentHarness` (a single instance + only for a batch of one). Framework specifics (ms-agent + system/memory/tool-message shape) live in the harness subclass, not + here. + * ``adapter_path`` / ``use_base_model``: see ``__init__``. * ``followup_fn``: see ``__init__``. """ def __init__( self, - sampler, - template: Template, + sampler=None, + template: Optional[Template] = None, tool_manager: Optional[ToolManager] = None, sampling_params: Optional[SamplingParams] = None, max_turns: int = 6, max_trajectory_tokens: Optional[int] = None, + concurrency: Optional[int] = None, trace_dir: Optional[str] = None, trace_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, success_callback: Optional[Callable[[Dict[str, Any]], bool]] = None, harness: Optional[AgentHarness] = None, adapter_path: Optional[str] = None, + use_base_model: bool = False, stop_after_stuck_turns: int = 0, max_malformed_retries: int = 2, followup_fn: Optional[Callable[[Trajectory, int], Any]] = None, + api: Optional[API] = None, + response_callback: Optional[ResponseCallback] = None, + api_appended_as: Literal['demonstration', 'context'] = 'demonstration', + api_kwargs: Optional[Dict[str, Any]] = None, ): super().__init__() + if isinstance(sampler, (API, APISampler)): + if api is not None: + raise ValueError('the positional backend and api= both specify an API') + api, sampler = sampler, None if template is None: raise ValueError('MultiTurnRollout requires a local Template instance') + if response_callback is None and sampler is None and api is None: + raise ValueError('MultiTurnRollout requires a sampler or API when response_callback is omitted') + if sampler is not None: + sample = getattr(type(sampler), 'sample', None) + if sample is None: + raise TypeError(f'backend must be an API or sampler, got {type(sampler).__name__}') + if not getattr(sample, '_enable_continous_work', False): + raise ValueError( + f'{type(sampler).__name__}.sample must be declared with ' + 'enable_continous_work=True: this rollout samples one trajectory per ' + 'call, and a slice_dp sampler raises when a worker gets nothing from ' + 'a batch of one.') + if adapter_path and use_base_model: + raise ValueError('adapter_path and use_base_model=True ask for opposite ' + 'weights; the sampler would drop the adapter silently.') if max_trajectory_tokens is not None and max_trajectory_tokens < 1: raise ValueError(f'max_trajectory_tokens must be >= 1 or None, got ' f'{max_trajectory_tokens}') self._init_common( max_turns=max_turns, sampling_params=sampling_params, + concurrency=concurrency, trace_dir=trace_dir, trace_callback=trace_callback, success_callback=success_callback) self.sampler = sampler self.template = template + if isinstance(api, APISampler): + if api_kwargs: + raise ValueError('api_kwargs belongs on the APISampler when api= is already adapted') + if api.template is not template: + raise ValueError('MultiTurnRollout and APISampler must share the same template instance') + self.api = api + elif api is not None: + self.api = APISampler( + api, template, appended_as=api_appended_as, api_kwargs=api_kwargs) + else: + if api_kwargs: + raise ValueError('api_kwargs requires an API backend') + self.api = None + self.response_callback = response_callback or _default_response_callback self.tool_manager = tool_manager self.harness = harness # A LoRA directory on disk, forwarded to every sample call. Training syncs @@ -175,6 +259,11 @@ def __init__( # such channel: without this, an eval script would silently measure the # base model and report it as the trained one. self.adapter_path = adapter_path + # The other direction: force the base weights. Needed because a sampler + # mid-training falls back to the LoRA synced into it whenever a call names + # no adapter, so a utility rollout (summarizing, judging) that wants the + # untrained model has to say so rather than stay silent. + self.use_base_model = use_base_model self.max_trajectory_tokens = max_trajectory_tokens # How many stuck turns in a row end the episode; 0 runs to ``max_turns`` # regardless. A turn is stuck when it made no progress at all, which is @@ -242,391 +331,363 @@ def __init__( @remote_function() def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory]: - if isinstance(trajectories, dict): - raise TypeError('MultiTurnRollout.__call__ expects a List[Trajectory]; ' - 'wrap a single trajectory as [trajectory].') - trajectories = list(trajectories) - n = len(trajectories) - if n == 0: - return [] + """The base implementation; the decorator is what a deployed handle needs.""" + return super().__call__(trajectories, **kwargs) - sampling_params = kwargs.get('sampling_params', self.sampling_params) + def _resolve_call(self, kwargs: Dict[str, Any], n: int) -> Dict[str, Any]: adapter_path = kwargs.get('adapter_path', self.adapter_path) # Left out entirely when unset, so a sampler without LoRA enabled sees the # same call it always did. adapter_kwargs = {'adapter_path': adapter_path} if adapter_path else {} - tool_managers = self._broadcast( - kwargs.get('tool_manager', self.tool_manager), n, name='tool_manager', required=True) - harnesses = self._broadcast(kwargs.get('harness', self.harness), n, name='harness') - lives: List[Optional[Trajectory]] = [ - dict(trajectories[i]) if harnesses[i] is not None else None for i in range(n) - ] - for live in lives: - if live is not None: - live['messages'] = list(live.get('messages') or []) + if kwargs.get('use_base_model', self.use_base_model): + adapter_kwargs['use_base_model'] = True + sampling_params = kwargs.get('sampling_params', self.sampling_params) + if sampling_params.num_samples != 1: + raise ValueError(f'MultiTurnRollout supports num_samples=1 only, got ' + f'{sampling_params.num_samples}') + response_callback = kwargs.get('response_callback', self.response_callback) + if not callable(response_callback): + raise TypeError('response_callback must be callable') + return { + 'sampling_params': sampling_params, + 'adapter_kwargs': adapter_kwargs, + 'response_callback': response_callback, + 'tool_managers': self._broadcast( + kwargs.get('tool_manager', self.tool_manager), n, name='tool_manager'), + 'harnesses': self._broadcast( + kwargs.get('harness', self.harness), n, name='harness', per_trajectory=True), + 'followup_fn': kwargs.get('followup_fn', self.followup_fn), + } + + def _run_one(self, trajectory: Trajectory, index: int, ctx: Dict[str, Any]) -> Trajectory: + tool_manager: ToolManager = ctx['tool_managers'][index] + harness: Optional[AgentHarness] = ctx['harnesses'][index] + followup_fn = ctx['followup_fn'] + adapter_kwargs: Dict[str, Any] = ctx['adapter_kwargs'] + response_callback: ResponseCallback = ctx['response_callback'] # 1. First before_generate happens *before* encode so memory/system # injection is in the initial prefix (not a later rewrite). - encode_trajs: List[Trajectory] = [] - for i, traj in enumerate(trajectories): - h, live = harnesses[i], lives[i] - if h is not None and live is not None: - lives[i] = h.before_generate(live) - live = lives[i] - traj = dict(traj) - traj['messages'] = list(live.get('messages') or []) - if live.get('tools'): - traj['tools'] = list(live['tools']) - encode_trajs.append(traj) - - pifs: List[Dict[str, Any]] = [] - for i, traj in enumerate(encode_trajs): - pif = self.template.encode(traj, add_generation_prompt=True) - pif = _to_plain(pif) - pif.setdefault('messages', list(traj.get('messages', []))) - pifs.append(pif) - if lives[i] is not None: - lives[i]['messages'] = list(pifs[i].get('messages') or []) - - all_logprobs: List[List[Any]] = [[] for _ in range(n)] - stop_reasons: List[Optional[str]] = [None] * n - turns: List[int] = [0] * n - truncated: List[bool] = [False] * n - done: List[bool] = [False] * n - # Consecutive turns that made no progress, the calls already issued in - # each episode, and whether being stuck is what ended it. All three stay - # at their initial value when ``stop_after_stuck_turns`` is 0. - stuck_turns: List[int] = [0] * n - seen_calls: List[set] = [set() for _ in range(n)] - stuck_stop: List[bool] = [False] * n + live: Optional[Trajectory] = None + to_encode = trajectory + if harness is not None: + live = dict(trajectory) + live['messages'] = list(live.get('messages') or []) + live = harness.before_generate(live) + to_encode = dict(trajectory) + to_encode['messages'] = list(live.get('messages') or []) + if live.get('tools'): + to_encode['tools'] = list(live['tools']) + + pif = _to_plain(self.template.encode(to_encode, add_generation_prompt=True)) + pif.setdefault('messages', list(to_encode.get('messages') or [])) + if 'tools' in to_encode: + pif['tools'] = list(to_encode.get('tools') or []) + elif tool_manager is not None: + pif['tools'] = list(tool_manager.tool_infos() or []) + if live is not None: + live['messages'] = list(pif.get('messages') or []) + + logprobs: List[Any] = [] + stop_reason: Optional[str] = None + generation_error: Optional[str] = None + turns = 0 + truncated = False + params = ctx['sampling_params'] + # Consecutive turns that made no progress, the calls already issued, and + # whether being stuck is what ended the episode. All three stay at their + # initial value when ``stop_after_stuck_turns`` is 0. + stuck_turns = 0 + seen_calls: set = set() + stuck_stop = False # Replies in a row whose tool-call markup did not parse. Reset by any # reply that produced a call, so one bad escape in the middle of a # working episode does not count against a later one. - malformed_turns: List[int] = [0] * n - # Follow-up bookkeeping (all no-ops when ``followup_fn`` is None): - # how many follow-ups each trajectory has had, and the params its next - # turn should use. A trajectory that has had one stops dispatching tools. - followups: List[int] = [0] * n - params_for: List[Any] = [sampling_params] * n - followup_fn = kwargs.get('followup_fn', self.followup_fn) - # Why the tool-calling part of each episode ended, when it was not the - # model's own choice: 'max_turns' or 'stuck'. Reported separately from - # ``truncated`` because an episode can hit the turn limit and still go on - # to answer the follow-up stages, in which case nothing was cut off. - tool_stop: List[Optional[str]] = [None] * n - - def append_followup(global_idx: int) -> bool: + malformed_turns = 0 + followups = 0 + # Why the tool-calling part ended, when it was not the model's own + # choice: 'max_turns' or 'stuck'. Reported separately from ``truncated`` + # because an episode can hit the turn limit and still go on to answer the + # follow-up stages, in which case nothing was cut off. + tool_stop: Optional[str] = None + # The loop counts generations, and each granted follow-up buys the one + # extra generation it asked for. Paying for the follow-up stages out of + # ``max_turns`` would mean an episode that spent its whole tool budget + # never reaches the stages that read what it built, and a short one + # silently gets more tool turns than a long one. + budget = self.max_turns + spent = 0 + + def grant_followup() -> bool: """Ask for one more stage; True when the episode carries on. Sets ``truncated`` itself in the one case where the answer is "there is no room for another stage", which is a cut trajectory rather than a caller that had nothing more to ask. """ - nonlocal iterations - if followup_fn is None: + nonlocal pif, live, followups, budget, params, truncated + if followup_fn is None or followups >= MAX_FOLLOWUPS: return False followup = followup_fn( - self._as_trajectory(trajectories[global_idx], pifs[global_idx], - all_logprobs[global_idx], turns[global_idx], - stop_reasons[global_idx], truncated[global_idx]), - followups[global_idx]) + self._as_trajectory(trajectory, pif, logprobs, turns, stop_reason, truncated), followups) if followup is None: return False - text, next_params = followup if isinstance(followup, tuple) else (followup, None) - extended = extend_with_bridge( - pifs[global_idx], [{'role': 'user', 'content': text}], self.template) + text, next_params = self._unpack_followup(followup) + extended = extend_with_bridge(pif, [{'role': 'user', 'content': text}], self.template) if extended is None: - truncated[global_idx] = True + truncated = True return False - pifs[global_idx] = extended - if lives[global_idx] is not None: - lives[global_idx]['messages'] = list(extended.get('messages') or []) - followups[global_idx] += 1 - iterations += 1 + pif = extended + # Follow-up stages are answers, so an API must not see tool schemas. + pif['tools'] = [] + if live is not None: + live['messages'] = list(extended.get('messages') or []) + followups += 1 + budget += 1 if next_params is not None: - params_for[global_idx] = next_params + params = next_params return True - # The loop counts generations, and each granted follow-up buys the one - # extra generation it asked for. Paying for the follow-up stages out of - # ``max_turns`` would mean an episode that spent its whole tool budget - # never reaches the stages that read what it built, and a short one - # silently gets more tool turns than a long one. - iterations = self.max_turns - done_iterations = 0 - first_turn = True - while done_iterations < iterations: - done_iterations += 1 - active = [i for i in range(n) if not done[i]] - if not active: - break + while spent < budget: + spent += 1 - if not first_turn: - for global_idx in active: - pifs[global_idx], lives[global_idx], dropped = self._harness_before_generate( - pifs[global_idx], lives[global_idx], harnesses[global_idx]) - if dropped: - truncated[global_idx] = True - done[global_idx] = True - active = [i for i in range(n) if not done[i]] - if not active: + if spent > 1: + pif, live, dropped = self._harness_before_generate(pif, live, harness) + if dropped: + truncated = True break - first_turn = False - - # 2. One batched sample call per distinct SamplingParams among the - # live trajectories -- normally exactly one, since only a - # follow-up stage asks for its own budget. Grouping rather than - # taking the first is what keeps a mixed batch honest: sampling one - # trajectory under another's token limit would silently truncate or - # over-spend, and the two are indistinguishable afterwards. - groups: List[List[int]] = [] - group_params: List[Any] = [] - for global_idx in active: - for slot, params in enumerate(group_params): - if params is params_for[global_idx]: - groups[slot].append(global_idx) + + # 2. One request. The callback chooses the local sampler or the API + # adapter, but both paths return exactly one SampledSequence. + try: + seq = response_callback( + self.sampler, + self.api, + params, + input_feature=pif, + adapter_kwargs=adapter_kwargs, + trajectory=trajectory, + trajectory_index=index, + turn=turns + 1, + followups=followups, + ) + except APIGenerationError as exc: + stop_reason = STOP_GENERATION_ERROR + generation_error = str(exc) + truncated = True + break + if not isinstance(seq, SampledSequence): + raise TypeError(f'response_callback must return SampledSequence, got ' + f'{type(seq).__name__}') + turns += 1 + + if seq.new_input_feature is None or 'input_ids' not in seq.new_input_feature: + raise RuntimeError(f'Sampler returned a SampledSequence without ' + f'new_input_feature.input_ids for trajectory ' + f'{index}; cannot continue multi-turn.') + + pif = _to_plain(dict(seq.new_input_feature)) + if seq.logprobs is not None: + if len(seq.logprobs) != len(seq.tokens): + raise RuntimeError(f'logprobs length ({len(seq.logprobs)}) does not ' + f'match sampled token count ({len(seq.tokens)}) ' + f'at turn {turns} (trajectory {index})') + logprobs.extend(seq.logprobs) + stop_reason = seq.stop_reason + + msgs = pif.get('messages') or [] + last_msg = msgs[-1] if msgs else None + tool_calls = (last_msg.get('tool_calls') if isinstance(last_msg, dict) else None) + if not tool_calls: + tool_calls = self.template.parse_tool_call(seq.decoded or '') + # After a follow-up, a parsed call is not a call: the tools were + # withdrawn for these stages on purpose (see ``followup_fn``), and + # dispatching python that the model wrote as *an answer* would edit + # the state the answer is about. + if followups: + tool_calls = None + # The parse also *rewrote* the message: when a reply parses as + # a call, the template stores it with the call text removed, so + # a caller reading the message gets less than the model wrote. + # For these stages the reply is the deliverable, and one of the + # tool-call formats is XML-shaped, so a check script asserting + # the content of an .xml file matches it: 5 of ex12's 72 check + # scripts came back with the XML cut out of them -- three then + # ran with `content == ''` where the model had written the file's + # real text, and two no longer held a code block at all. + if msgs and isinstance(last_msg, dict): + # Decoded without the special tokens, the way the template + # writes a message: ``seq.decoded`` keeps the closing + # ``<|im_end|>``, and putting that in the content put it in + # the problem statements ex13 handed to solvers -- 7 of 7 of + # them ended in a literal '<|im_end|>'. + tok = getattr(self.template, 'tokenizer', None) + if tok is not None and seq.tokens: + last_msg['content'] = tok.decode(seq.tokens, skip_special_tokens=True) + else: + last_msg['content'] = seq.decoded or '' + last_msg.pop('tool_calls', None) + + if live is not None: + live['messages'] = list(msgs) + if harness is not None and live is not None: + live = harness.after_generate(live, seq.decoded or '', tool_calls or []) + self._merge_assistant_metadata(pif, live) + + # 3. Termination conditions + # A reply cut off at ``max_tokens`` is truncated in exactly the sense + # the flag names, and consumers read the flag to tell a trajectory + # that finished from one that ran out of room: a difficulty + # measurement counting such an attempt as a genuine failure blames + # the task for the token budget. Tool calls the cut reply happens to + # contain are still not dispatched -- the turn never got to decide it + # was done emitting them. + if seq.stop_reason == 'length': + truncated = True + break + + # 3a. Sequence-length cap. + if (self.max_trajectory_tokens is not None + and len(pif.get('input_ids') or []) >= self.max_trajectory_tokens): + truncated = True + break + + if not tool_calls: + # Markup that did not parse is the model asking for a tool, not + # declining one -- ending here tells it nothing and throws the + # turn away. Hand back the parser's own reason and let it write + # the call again. Not after a follow-up: tools are withdrawn + # there on purpose (see ``followup_fn``), so a reply that looks + # like a call is meant to be read as text. + parse_errors = ([] if followups else self.template.tool_call_errors(seq.decoded or '')) + if parse_errors and malformed_turns < self.max_malformed_retries: + malformed_turns += 1 + extended = extend_with_bridge(pif, [_malformed_tool_message(parse_errors)], self.template) + if extended is None: + truncated = True break - else: - group_params.append(params_for[global_idx]) - groups.append([global_idx]) - - resps_by_idx: Dict[int, Any] = {} - device_mesh = getattr(self.sampler, 'device_mesh', None) - min_batch_size = (device_mesh.data_world_size if device_mesh is not None else 1) - # A sampler that routes each request on its own accepts a batch smaller - # than its worker count, so the padding below is not needed. It was only - # ever there because slicing a batch over all workers raises when some - # rank gets nothing, and the duplicates it added were generated and then - # dropped -- with one prompt and 8 workers that is 8 generations for 1 - # kept result. - if getattr(type(self.sampler).sample, '_enable_continous_work', False): - min_batch_size = 1 - for slot, group in enumerate(groups): - batch_pifs = [pifs[i] for i in group] - actual = len(batch_pifs) - if actual < min_batch_size: - batch_pifs = batch_pifs + ([batch_pifs[-1]] * (min_batch_size - actual)) - group_resps = self.sampler.sample(batch_pifs, - sampling_params=group_params[slot], - **adapter_kwargs) - group_resps = self._unwrap_response_list(group_resps, len(batch_pifs))[:actual] - for local_idx, global_idx in enumerate(group): - resps_by_idx[global_idx] = group_resps[local_idx] - - pending_tools: List[tuple] = [] # (global_idx, tool_calls) - for global_idx in active: - turns[global_idx] += 1 - seq = resps_by_idx[global_idx].sequences[0] - - if seq.new_input_feature is None or 'input_ids' not in seq.new_input_feature: - raise RuntimeError(f'Sampler returned a SampledSequence without ' - f'new_input_feature.input_ids for trajectory ' - f'{global_idx}; cannot continue multi-turn.') - - pifs[global_idx] = _to_plain(dict(seq.new_input_feature)) - if seq.logprobs is not None: - if len(seq.logprobs) != len(seq.tokens): - raise RuntimeError(f'logprobs length ({len(seq.logprobs)}) does not ' - f'match sampled token count ({len(seq.tokens)}) ' - f'at turn {turns[global_idx]} ' - f'(trajectory {global_idx})') - all_logprobs[global_idx].extend(seq.logprobs) - stop_reasons[global_idx] = seq.stop_reason - - _msgs = pifs[global_idx].get('messages') or [] - _last_msg = _msgs[-1] if _msgs else None - tool_calls = (_last_msg.get('tool_calls') if isinstance(_last_msg, dict) else None) - if not tool_calls: - tool_calls = self.template.parse_tool_call(seq.decoded or '') - # After a follow-up, a parsed call is not a call: the tools were - # withdrawn for these stages on purpose (see ``followup_fn``), and - # dispatching python that the model wrote as *an answer* would edit - # the state the answer is about. - if followups[global_idx]: - tool_calls = None - # The parse also *rewrote* the message: when a reply parses as - # a call, the template stores it with the call text removed, so - # a caller reading the message gets less than the model wrote. - # For these stages the reply is the deliverable, and one of the - # tool-call formats is XML-shaped, so a check script asserting - # the content of an .xml file matches it: 5 of ex12's 72 check - # scripts came back with the XML cut out of them -- three then - # ran with `content == ''` where the model had written the file's - # real text, and two no longer held a code block at all. - if _msgs and isinstance(_last_msg, dict): - # Decoded without the special tokens, the way the - # template writes a message: ``seq.decoded`` keeps the - # closing ``<|im_end|>``, and putting that in the content - # put it in the problem statements ex13 handed to solvers - # -- 7 of 7 of them ended in a literal '<|im_end|>'. - tok = getattr(self.template, 'tokenizer', None) - if tok is not None and seq.tokens: - _last_msg['content'] = tok.decode( - seq.tokens, skip_special_tokens=True) - else: - _last_msg['content'] = seq.decoded or '' - _last_msg.pop('tool_calls', None) - - if lives[global_idx] is not None: - lives[global_idx]['messages'] = list(_msgs) - if harnesses[global_idx] is not None and lives[global_idx] is not None: - lives[global_idx] = harnesses[global_idx].after_generate( - lives[global_idx], seq.decoded or '', tool_calls or []) - self._merge_assistant_metadata(pifs[global_idx], lives[global_idx]) - - # 3. Termination conditions - # A reply cut off at ``max_tokens`` is truncated in exactly the - # sense the flag names, and consumers read the flag to tell a - # trajectory that finished from one that ran out of room: a - # difficulty measurement counting such an attempt as a genuine - # failure blames the task for the token budget. Tool calls the - # cut reply happens to contain are still not dispatched -- the - # turn never got to decide it was done emitting them. - if seq.stop_reason == 'length': - truncated[global_idx] = True - done[global_idx] = True + pif = extended + if live is not None: + live['messages'] = list(extended.get('messages') or []) continue - - # 3a. Sequence-length cap. - if (self.max_trajectory_tokens is not None - and len(pifs[global_idx].get('input_ids') or []) >= self.max_trajectory_tokens): - truncated[global_idx] = True - done[global_idx] = True + # The episode is over as far as the model is concerned. Give the + # caller one chance to say otherwise -- see ``followup_fn`` for + # why this is not a second rollout. + if grant_followup(): continue + break - if not tool_calls: - # Markup that did not parse is the model asking for a tool, - # not declining one -- ending here tells it nothing and throws - # the turn away. Hand back the parser's own reason and let it - # write the call again. Not after a follow-up: tools are - # withdrawn there on purpose (see ``followup_fn``), so a reply - # that looks like a call is meant to be read as text. - parse_errors = ([] if followups[global_idx] else - self.template.tool_call_errors(seq.decoded or '')) - if (parse_errors and malformed_turns[global_idx] < self.max_malformed_retries): - malformed_turns[global_idx] += 1 - extended = extend_with_bridge(pifs[global_idx], - [_malformed_tool_message(parse_errors)], - self.template) - if extended is None: - truncated[global_idx] = True - done[global_idx] = True - continue - pifs[global_idx] = extended - if lives[global_idx] is not None: - lives[global_idx]['messages'] = list(extended.get('messages') or []) - continue - # The episode is over as far as the model is concerned. Give - # the caller one chance to say otherwise -- see - # ``followup_fn`` for why this is not a second rollout. - if append_followup(global_idx): - continue - done[global_idx] = True + if turns >= self.max_turns: + # Out of tool turns, not out of episode: the stages that read the + # end state can still run on what was built. + tool_stop = 'max_turns' + if grant_followup(): continue + truncated = True + break - if turns[global_idx] >= self.max_turns: - # Out of tool turns, not out of episode: the stages that read - # the end state can still run on what was built. - tool_stop[global_idx] = 'max_turns' - if append_followup(global_idx): - continue - truncated[global_idx] = True - done[global_idx] = True + malformed_turns = 0 + + # 4. This turn's calls, then the harness formats the observations + # into tool messages (append-only bridge). + if tool_manager is None: + raise ValueError('the model emitted tool_calls but this trajectory has no ToolManager') + observations = self._run_tools(tool_manager, tool_calls) + if self.stop_after_stuck_turns: + keys = [_call_key(tc) for tc in tool_calls] + all_repeats = bool(keys) and all(k in seen_calls for k in keys) + seen_calls.update(keys) + all_errors = bool(observations) and all(is_error_observation(o) for o in observations) + if all_errors or all_repeats: + stuck_turns += 1 + else: + stuck_turns = 0 + + tool_messages, live = self._tool_messages_after(pif, live, harness, observations, tool_calls) + extended = extend_with_bridge(pif, tool_messages, self.template) + overflowed = extended is None + if overflowed: + # Trajectory exceeded max_length. + truncated = True + else: + pif = extended + if live is not None: + live['messages'] = list(extended.get('messages') or []) + # Checked after the messages are appended, so the turns that ended + # the episode are in the trajectory the caller reads. + if self.stop_after_stuck_turns and stuck_turns >= self.stop_after_stuck_turns: + stuck_stop = True + tool_stop = 'stuck' + # Same as the turn limit: the tool phase is over, the state it + # left is not, so the stages still get their turn. + if not overflowed and grant_followup(): continue + truncated = True + break + if overflowed: + break - malformed_turns[global_idx] = 0 - pending_tools.append((global_idx, list(tool_calls))) - - # 4. Parallel tool dispatch across the live batch, then harness - # formats observations into tool messages (append-only bridge). - # The bridge itself is computed serially: it is a cheap - # decode-diff-encode on python strings / token lists. - if pending_tools: - obs_by_traj = self._dispatch_tools(tool_managers, pending_tools) - for global_idx, tool_calls in pending_tools: - observations = obs_by_traj.get(global_idx) or [''] * len(tool_calls) - if self.stop_after_stuck_turns: - keys = [_call_key(tc) for tc in tool_calls] - all_repeats = bool(keys) and all(k in seen_calls[global_idx] - for k in keys) - seen_calls[global_idx].update(keys) - all_errors = bool(observations) and all( - is_error_observation(o) for o in observations) - if all_errors or all_repeats: - stuck_turns[global_idx] += 1 - else: - stuck_turns[global_idx] = 0 - tool_messages, lives[global_idx] = self._tool_messages_after( - pifs[global_idx], lives[global_idx], harnesses[global_idx], - observations, tool_calls) - extended = extend_with_bridge(pifs[global_idx], tool_messages, self.template) - if extended is None: - # Trajectory exceeded max_length, mark as done (deleted) - truncated[global_idx] = True - done[global_idx] = True - else: - pifs[global_idx] = extended - if lives[global_idx] is not None: - lives[global_idx]['messages'] = list(extended.get('messages') or []) - # Checked after the messages are appended, so the turns that - # ended the episode are in the trajectory the caller reads. - if (self.stop_after_stuck_turns - and stuck_turns[global_idx] >= self.stop_after_stuck_turns): - stuck_stop[global_idx] = True - tool_stop[global_idx] = 'stuck' - # Same as the turn limit: the tool phase is over, the - # state it left is not, so the stages still get their turn. - if not done[global_idx] and append_followup(global_idx): - continue - truncated[global_idx] = True - done[global_idx] = True - - for i in range(n): - if not all_logprobs[i]: - continue - labels_i = pifs[i].get('labels') or [] - trainable_i = sum(1 for label in labels_i if label != -100) - if len(all_logprobs[i]) != trainable_i: - raise RuntimeError(f'logprobs/labels misaligned for trajectory {i}: ' - f'{len(all_logprobs[i])} logprobs vs {trainable_i} ' - f'trainable labels (labels != -100). This invariant is ' - f'required by grpo._pad_and_align_to_batch; a mismatch ' - f'would silently corrupt GRPO old_logps alignment.') - - # 5. Merge pif fields into each trajectory dict at TOP LEVEL so - # downstream consumers (VLLMSampler with ``'input_ids' in inputs``) - # see an encoded InputFeature and skip re-encoding. - outs: List[Trajectory] = [] - for i, traj in enumerate(trajectories): - out = dict(traj) - out.update(pifs[i]) - out['messages'] = list(pifs[i].get('messages') or out.get('messages', [])) - out['logprobs'] = all_logprobs[i] if all_logprobs[i] else None - out['turns'] = turns[i] - out['stop_reason'] = stop_reasons[i] - out['truncated'] = truncated[i] - # ``truncated`` says something was cut off; these two say what ended - # the tool-calling part, which is a different question -- an episode - # can run out of turns, be handed a follow-up stage, and finish it. - out['stuck_stop'] = stuck_stop[i] - out['tool_stop'] = tool_stop[i] - out['followups'] = followups[i] - outs.append(out) - - # Per-rollout trace dump: one JSON file per selected trajectory. - # ``trace_callback`` decides whether to store; ``success_callback`` - # decides the filename prefix. Observability only -- any failure - # is swallowed inside ``_write_rollout_traces``. - if self.trace_dir: - self._write_rollout_traces(outs, global_step=kwargs.get('global_step')) - return outs + if logprobs: + labels = pif.get('labels') or [] + completion_mask = pif.get('completion_mask') + if completion_mask is None: + expected = sum(1 for label in labels if label != -100) + elif len(completion_mask) != len(labels): + raise RuntimeError(f'completion_mask/labels misaligned for trajectory {index}: ' + f'{len(completion_mask)} != {len(labels)}') + else: + expected = sum(1 for label, flag in zip(labels, completion_mask) if label != -100 and flag) + if len(logprobs) != expected: + raise RuntimeError(f'logprobs/policy-token alignment failed for trajectory {index}: ' + f'{len(logprobs)} logprobs vs {expected} positions selected by ' + '(labels != -100) & completion_mask.') + + # 5. Merge pif fields into the trajectory dict at TOP LEVEL so downstream + # consumers (VLLMSampler with ``'input_ids' in inputs``) see an encoded + # InputFeature and skip re-encoding. + out = dict(trajectory) + out.update(pif) + out['messages'] = list(pif.get('messages') or out.get('messages', [])) + out['logprobs'] = logprobs if logprobs else None + out['turns'] = turns + out['stop_reason'] = stop_reason + out['truncated'] = truncated + # ``truncated`` says something was cut off; these two say what ended the + # tool-calling part, which is a different question -- an episode can run + # out of turns, be handed a follow-up stage, and finish it. + out['stuck_stop'] = stuck_stop + out['tool_stop'] = tool_stop + out['followups'] = followups + if generation_error is not None: + out['error'] = generation_error + return out # ------------------------------------------------------------------ private @staticmethod - def _as_trajectory(traj: Trajectory, pif: Dict[str, Any], logprobs: List[Any], - turns: int, stop_reason: Optional[str], - truncated: bool) -> Trajectory: + def _run_tools(tool_manager: ToolManager, tool_calls: List[Dict[str, Any]]) -> List[str]: + """Run one turn's calls, through ``call_many`` when the manager has it. + + A turn's calls go together because they share one Env round trip + (``Env.step_batch``). Calls from *different* trajectories no longer meet + here -- each episode has its own thread and, in the sandbox case, its own + Env -- so there is nothing left to group across. + + A manager that answers with fewer results than calls leaves the rest + empty rather than shifting them onto the wrong call. + """ + if hasattr(tool_manager, 'call_many'): + contents = tool_manager.call_many(tool_calls) + else: + contents = [tool_manager(tc) for tc in tool_calls] + obs = [''] * len(tool_calls) + for i, content in enumerate(contents[:len(tool_calls)]): + obs[i] = '' if content is None else str(content) + return obs + + @staticmethod + def _as_trajectory(traj: Trajectory, pif: Dict[str, Any], logprobs: List[Any], turns: int, + stop_reason: Optional[str], truncated: bool) -> Trajectory: """The episode so far, shaped like the value ``__call__`` returns. Handed to ``followup_fn`` so the callback reads an episode the same way @@ -683,55 +744,6 @@ def _merge_assistant_metadata(pif: Dict[str, Any], live: Trajectory) -> None: if last_asst.get(key) and not dst.get(key): dst[key] = last_asst[key] - def _dispatch_tools( - self, - tool_managers: List[ToolManager], - pending: List[Tuple[int, List[Dict[str, Any]]]], - ) -> Dict[int, List[str]]: - """Run tool calls for the live batch, grouped by ToolManager. - - Trajectories that share a manager (and therefore often one Env) go - through ``call_many`` / ``Env.step_batch``. Distinct managers run - concurrently so remote sandboxes are not serialized on generate. - """ - obs: Dict[int, List[str]] = { - gi: [''] * len(tcs) for gi, tcs in pending - } - groups: Dict[int, List[Tuple[int, int, Dict[str, Any]]]] = defaultdict(list) - mgr_by_id: Dict[int, ToolManager] = {} - for gi, tcs in pending: - mid = id(tool_managers[gi]) - mgr_by_id[mid] = tool_managers[gi] - for ci, tc in enumerate(tcs): - groups[mid].append((gi, ci, tc)) - - def _run_group(items: List[Tuple[int, int, Dict[str, Any]]], mgr: ToolManager): - tcs = [tc for _, _, tc in items] - if hasattr(mgr, 'call_many'): - contents = mgr.call_many(tcs) - else: - contents = [mgr(tc) for tc in tcs] - return list(zip(items, contents)) - - group_items = list(groups.items()) - if len(group_items) == 1: - mid, items = group_items[0] - finished = [_run_group(items, mgr_by_id[mid])] - else: - finished = [] - with ThreadPoolExecutor(max_workers=min(32, len(group_items))) as pool: - futs = [ - pool.submit(_run_group, items, mgr_by_id[mid]) - for mid, items in group_items - ] - for fut in as_completed(futs): - finished.append(fut.result()) - - for group_result in finished: - for (gi, ci, _tc), content in group_result: - obs[gi][ci] = '' if content is None else str(content) - return obs - def _tool_messages_after( self, pif: Dict[str, Any], @@ -750,22 +762,3 @@ def _tool_messages_after( if not delta: return fallback, live return delta, live - - @staticmethod - def _unwrap_response_list(resps, expected: int) -> List[SampleResponse]: - """Validate that the sampler returned ``expected`` ``SampleResponse``s, - one per input in the batch. - """ - if not isinstance(resps, list): - raise TypeError(f'expected List[SampleResponse] from sampler.sample (batched ' - f'call), got {type(resps).__name__}') - if len(resps) != expected: - raise RuntimeError(f'sampler returned {len(resps)} responses for a batch of ' - f'{expected} trajectories; expected one per input.') - for i, r in enumerate(resps): - if not isinstance(r, SampleResponse): - raise TypeError(f'expected SampleResponse at batch index {i}, got ' - f'{type(r).__name__}') - if not r.sequences: - raise RuntimeError(f'SampleResponse at batch index {i} has no sequences') - return resps diff --git a/src/twinkle_agentic/summarizer/base.py b/src/twinkle_agentic/summarizer/base.py index 8e9be5f5f..aecc25675 100644 --- a/src/twinkle_agentic/summarizer/base.py +++ b/src/twinkle_agentic/summarizer/base.py @@ -5,11 +5,12 @@ import re from typing import TYPE_CHECKING, Any, Sequence +from twinkle_agentic.rollout import MultiTurnRollout from twinkle_agentic.utils.llm_backup import llm_backup +from twinkle_agentic.utils.message_utils import assistant_text if TYPE_CHECKING: from twinkle.data_format import SamplingParams, Trajectory # noqa: F401 - from twinkle.sampler.base import Sampler # noqa: F401 DEFAULT_USER_PROMPT_TEMPLATE = """\ @@ -36,7 +37,8 @@ class Summarizer: - LLM_BACKUP_BASE_URL: API endpoint Args: - sampler: Student model sampler (local inference, shared across types). + backend: a sampler or an API client, driven through + :class:`~twinkle_agentic.rollout.MultiTurnRollout`. compression_ratio: Target compression factor (> 1). model_path: Model identifier. sampling_params: Default sampling params. @@ -44,14 +46,19 @@ class Summarizer: user_prompt_template: User prompt template. Must contain ``{budget}`` and ``{text}``. May contain ``{query}``. min_budget_chars: Floor for the character budget in the prompt. - template: Optional :class:`Template` for special token stripping. + template: local :class:`Template`, required by the sampler path and also + what special-token stripping reads its tokenizer from. lora_path: LoRA adapter path specific to this summarizer type. - Each subclass can use a different LoRA for its task. + Each subclass can use a different LoRA for its task. Without one the + base weights are asked for explicitly -- a sampler mid-training + otherwise lends this out the policy LoRA synced into it. + rollout_kwargs: passed to ``MultiTurnRollout``. API request options + belong in ``api_kwargs``. """ def __init__( self, - sampler: Sampler, + backend: Any, compression_ratio: float = 2.0, *, model_path: str = '', @@ -61,9 +68,10 @@ def __init__( min_budget_chars: int = 250, template: Any | None = None, lora_path: str | None = None, + **rollout_kwargs: Any, ): - if sampler is None: - raise ValueError('sampler is required') + if backend is None: + raise ValueError('backend is required') if compression_ratio <= 1.0: raise ValueError(f'compression_ratio must be > 1, got {compression_ratio}') if min_budget_chars < 1: @@ -74,7 +82,6 @@ def __init__( raise ValueError('user_prompt_template must contain both {budget} and {text}') self.model_path = model_path - self.sampler = sampler self.compression_ratio = float(compression_ratio) self.sampling_params = sampling_params self.system_prompt = system_prompt @@ -83,6 +90,20 @@ def __init__( self.template = template self.lora_path = lora_path if lora_path else None self._special_tokens_cache: tuple[str, ...] | None = None + # Built on the first call rather than here, so a summarizer that never + # compresses anything (every text already under budget) costs nothing. + self._backend = backend + self._rollout_kwargs = dict(rollout_kwargs, max_turns=1) + if template is not None: + self._rollout_kwargs['template'] = template + # Which weights, and only for a local sampler: an API endpoint serves + # whatever it serves and has no notion of an adapter. + if hasattr(backend, 'sample'): + if self.lora_path: + self._rollout_kwargs['adapter_path'] = self.lora_path + else: + self._rollout_kwargs['use_base_model'] = True + self._rollout: Any | None = None # ------------------------------------------------------------------ # public entry point (pre/post processing, NOT decorated) @@ -106,14 +127,15 @@ def __call__(self, text: str, system: str = None, query: str = None, # ------------------------------------------------------------------ @llm_backup(key_params=["query"]) def _sample(self, trajectory, sampling_params, query: str = None) -> str: - """Student model: trajectory + sampling_params -> raw text.""" - sample_kwargs: dict[str, Any] = {'sampling_params': sampling_params} - if self.lora_path is None: - sample_kwargs['use_base_model'] = True - else: - sample_kwargs['adapter_path'] = self.lora_path - responses = self.sampler.sample([trajectory], **sample_kwargs) - return self._decoded(list(responses)[0]) if responses else '' + """Student model: trajectory + sampling_params -> raw text. + + The signature is what ``llm_backup`` reads by name to hand the teacher the + same input, so it stays even though the body no longer touches a sampler. + """ + if self._rollout is None: + self._rollout = MultiTurnRollout(self._backend, **self._rollout_kwargs) + replies = self._rollout([trajectory], sampling_params=sampling_params) + return assistant_text(replies[0]) if replies else '' # ------------------------------------------------------------------ # internals @@ -121,7 +143,7 @@ def _sample(self, trajectory, sampling_params, query: str = None) -> str: def _get_special_tokens(self) -> tuple[str, ...]: if self._special_tokens_cache is not None: return self._special_tokens_cache - tpl = self.template or getattr(self.sampler, 'template', None) + tpl = self.template or getattr(self._backend, 'template', None) tokenizer = getattr(tpl, 'tokenizer', None) if tpl is not None else None tokens: list[str] = [] if tokenizer is not None: @@ -176,13 +198,6 @@ def _postprocess(raw: str, original: str, special_tokens: tuple[str, ...]) -> st return None return text - @staticmethod - def _decoded(response: Any) -> str: - seqs = getattr(response, 'sequences', None) or [] - if not seqs: - return '' - return getattr(seqs[0], 'decoded', None) or '' - @staticmethod def _strip_code_fences(text: str) -> str: stripped = text.strip() diff --git a/src/twinkle_agentic/utils/code_utils.py b/src/twinkle_agentic/utils/code_utils.py index 37368045e..ab8f38403 100644 --- a/src/twinkle_agentic/utils/code_utils.py +++ b/src/twinkle_agentic/utils/code_utils.py @@ -40,16 +40,17 @@ @lru_cache(maxsize=None) -def _fence_re(language_tags: Tuple[str, ...]) -> Pattern: - """A fenced block whose language tag is one of ``language_tags``, or absent. +def _fence_re(language_tags: Optional[Tuple[str, ...]]) -> Pattern: + """Match a fenced block, optionally restricting its language label. - ``python``, ``py``, ``Python`` and ``python3`` are one intent spelled four - ways, so a tag matches case-insensitively and with any version suffix. A tag - that is not on the list -- ``bash``, ``json`` -- is a different intent, and is - not read as code at all. + ``None`` accepts any label. Otherwise, listed tags match case-insensitively, + with any version suffix; an unlabelled fence is accepted as well. """ - alts = '|'.join(re.escape(tag) for tag in language_tags) - label = r'(?:(?:%s)[\d.]*)?' % alts if alts else '' + if language_tags is None: + label = r'[^\r\n]*' + else: + alts = '|'.join(re.escape(tag) for tag in language_tags) + label = r'(?:(?:%s)[\d.]*)?' % alts if alts else '' return re.compile(r'```[ \t]*%s[ \t]*\r?\n(.*?)```' % label, re.S | re.I) @@ -69,10 +70,14 @@ def strip_reasoning(text: str) -> str: return body[cut:] -def parse_fenced_code(text: str, language_tags: Tuple[str, ...] = PYTHON_TAGS) -> Optional[str]: - """The last block ``text`` fenced as that language, or None if there is none. +def parse_fenced_code( + text: str, + language_tags: Optional[Tuple[str, ...]] = PYTHON_TAGS, +) -> Optional[str]: + """Return the last matching fenced block, or None if there is none. - The last one, not the first: a model often drafts a version before the final + Pass ``language_tags=None`` to accept any language label. The last block, not + the first, is returned because a model often drafts a version before the final one, and the block it ends on is its answer. What is inside is taken as given -- a fence is the model saying which part is diff --git a/tests/preprocessor/test_logprob_utils.py b/tests/preprocessor/test_logprob_utils.py deleted file mode 100644 index 29bb18e36..000000000 --- a/tests/preprocessor/test_logprob_utils.py +++ /dev/null @@ -1,354 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Tests for preprocessor.logprob_utils — pure logprob math helpers. - -These helpers compute conditional-vs-unconditional logprob deltas for -IFD-family scoring (CherryLLM, T-SHIRT, ChR). All functions are stateless -and accept simple list inputs. - -Conventions used in this test file: - * "lp" lists are aligned to the FULL sequence (prompt + answer). - * ``n_prompt`` is the number of prompt tokens; assistant tokens start at - index ``n_prompt`` in the cond list. - * Each lp entry is a dict {token_id: logprob_float}. -""" -import math -import pytest - -from twinkle_agentic.preprocessor.logprob_utils import (_chr_min_distinct, _chr_min_weighted, _extract_logprob, - _ifd_family_metrics, _lp_to_jsonable, _mean_logprob_delta, - _pad_batch, _to_int_list) - -# ── _extract_logprob ──────────────────────────────────────────────────────── - - -class TestExtractLogprob: - - def test_none(self): - assert _extract_logprob(None) is None - - def test_scalar_int(self): - assert _extract_logprob(5) == 5.0 - - def test_scalar_float(self): - assert _extract_logprob(-1.2) == -1.2 - - def test_dict_with_int_token_id(self): - lp = {7: -0.5, 8: -2.0} - assert _extract_logprob(lp, token_id=7) == -0.5 - assert _extract_logprob(lp, token_id=8) == -2.0 - - def test_dict_with_str_token_id_fallback(self): - # vLLM may emit string keys; lookup must fall back to str(token_id). - lp = {'7': -0.5} - assert _extract_logprob(lp, token_id=7) == -0.5 - - def test_dict_no_token_id_picks_first(self): - # No token_id → iter-first behaviour. - lp = {7: -0.5} - assert _extract_logprob(lp) == -0.5 - - def test_dict_token_id_missing_uses_first(self): - # token_id not in dict → fall back to first entry. - lp = {99: -3.0} - assert _extract_logprob(lp, token_id=7) == -3.0 - - def test_dict_with_logprob_attr_object(self): - - class Entry: - - def __init__(self, v): - self.logprob = v - - lp = {7: Entry(-0.7)} - assert _extract_logprob(lp, token_id=7) == -0.7 - - def test_dict_with_nested_dict(self): - lp = {7: {'logprob': -0.9, 'rank': 1}} - assert _extract_logprob(lp, token_id=7) == -0.9 - - def test_dict_with_nested_dict_none_logprob(self): - lp = {7: {'logprob': None}} - assert _extract_logprob(lp, token_id=7) is None - - def test_unrecognized_type(self): - # str entries → returns None - lp = {7: 'oops'} - assert _extract_logprob(lp, token_id=7) is None - - def test_non_dict_non_scalar(self): - # A list is neither scalar nor dict → None. - assert _extract_logprob([1, 2, 3]) is None - - -# ── _to_int_list ──────────────────────────────────────────────────────────── - - -class TestToIntList: - - def test_plain_list(self): - assert _to_int_list([1, 2, 3]) == [1, 2, 3] - - def test_tuple(self): - assert _to_int_list((1, 2, 3)) == [1, 2, 3] - - def test_with_tolist(self): - - class Tensor: - - def tolist(self): - return [4, 5, 6] - - assert _to_int_list(Tensor()) == [4, 5, 6] - - def test_empty(self): - assert _to_int_list([]) == [] - - -# ── _chr_min_distinct ─────────────────────────────────────────────────────── - - -class TestChrMinDistinct: - - def test_empty_inputs_returns_none(self): - assert _chr_min_distinct([], [{1: -1.0}], [], [1], 0) is None - assert _chr_min_distinct([{1: -1.0}], [], [1], [], 0) is None - assert _chr_min_distinct([{1: -1.0}], [{1: -1.0}], [1], [], 0) is None - - def test_simple_all_positive(self): - # cond_lp[i] - asst_lp[i] > 0 for all i → ratio = 1.0 - n_prompt = 1 - # cond covers prompt(1) + asst(2) = 3 positions - cond_lp = [ - { - 0: -10.0 - }, # prompt position - { - 1: -0.1 - }, # asst pos 0 — high cond logprob - { - 2: -0.2 - } - ] # asst pos 1 - asst_lp = [{1: -1.0}, {2: -1.5}] - cond_ids = [0, 1, 2] - asst_ids = [1, 2] - ratio = _chr_min_distinct(cond_lp, asst_lp, cond_ids, asst_ids, n_prompt) - assert ratio == 1.0 - - def test_all_negative(self): - # delta < 0 → ratio = 0 - n_prompt = 1 - cond_lp = [{0: 0.0}, {1: -3.0}, {2: -3.0}] - asst_lp = [{1: -0.5}, {2: -0.5}] - ratio = _chr_min_distinct(cond_lp, asst_lp, [0, 1, 2], [1, 2], n_prompt) - assert ratio == 0.0 - - def test_distinct_token_min_aggregation(self): - # Two occurrences of same token: one has +delta, one has -delta. - # min(deltas) is negative → token contributes 0 to ratio. - n_prompt = 1 - cond_lp = [{0: 0.0}, {1: -0.1}, {1: -3.0}] - asst_lp = [{1: -1.0}, {1: -0.5}] # delta1=+0.9, delta2=-2.5 - ratio = _chr_min_distinct(cond_lp, asst_lp, [0, 1, 1], [1, 1], n_prompt) - assert ratio == 0.0 # min < 0 - - def test_exclude_ids(self): - # Excluded token is dropped before counting. - n_prompt = 1 - cond_lp = [{0: 0.0}, {1: -0.1}, {2: -0.1}] - asst_lp = [{1: -1.0}, {2: -1.0}] - # Without exclude: 2 distinct tokens, both positive → 1.0 - ratio = _chr_min_distinct(cond_lp, asst_lp, [0, 1, 2], [1, 2], n_prompt, exclude_ids={1}) - assert ratio == 1.0 # only token 2 counted, still positive - - def test_truncation_when_cond_short(self): - # cond_lp shorter than n_prompt + n_asst → loop breaks early. - n_prompt = 2 - cond_lp = [{0: 0.0}, {0: 0.0}, {1: -0.1}] # only 1 asst position - asst_lp = [{1: -1.0}, {2: -1.0}] # 2 asst positions requested - ratio = _chr_min_distinct(cond_lp, asst_lp, [0, 0, 1], [1, 2], n_prompt) - assert ratio == 1.0 # only the first delta processed - - -# ── _chr_min_weighted ─────────────────────────────────────────────────────── - - -class TestChrMinWeighted: - - def test_empty_returns_none(self): - assert _chr_min_weighted([], [{1: -1.0}], [], [1], 0) is None - - def test_all_positive_returns_one(self): - n_prompt = 1 - cond_lp = [{0: 0.0}, {1: -0.1}, {2: -0.2}] - asst_lp = [{1: -1.0}, {2: -1.5}] - ratio = _chr_min_weighted(cond_lp, asst_lp, [0, 1, 2], [1, 2], n_prompt) - assert ratio == 1.0 # all positive → pos_w == total_w - - def test_zero_total_weight_returns_none(self): - # All deltas == 0 → total_w == 0 → None - n_prompt = 1 - cond_lp = [{0: 0.0}, {1: -1.0}] - asst_lp = [{1: -1.0}] - assert _chr_min_weighted(cond_lp, asst_lp, [0, 1], [1], n_prompt) is None - - def test_weighted_mixture(self): - # Token A: min_delta = +2.0 (weight 2) - # Token B: min_delta = -1.0 (weight 1) - # pos / total = 2 / 3 - n_prompt = 1 - cond_lp = [{0: 0.0}, {1: 1.0}, {2: -2.0}] # cond: A=1.0, B=-2.0 - asst_lp = [{1: -1.0}, {2: -1.0}] # asst: A=-1.0, B=-1.0 - # delta A = 1.0 - (-1.0) = 2.0 - # delta B = -2.0 - (-1.0) = -1.0 - ratio = _chr_min_weighted(cond_lp, asst_lp, [0, 1, 2], [1, 2], n_prompt) - assert abs(ratio - 2 / 3) < 1e-9 - - -# ── _ifd_family_metrics ───────────────────────────────────────────────────── - - -class TestIfdFamilyMetrics: - - def test_empty_returns_empty_dict(self): - assert _ifd_family_metrics([], [{1: -1.0}], [], [1], 0) == {} - - def test_simple_uniform(self): - # All deltas = 0.5 → mean=0.5, ifd=exp(-0.5) - n_prompt = 1 - cond_lp = [{0: 0.0}, {1: -0.5}, {2: -0.5}] - asst_lp = [{1: -1.0}, {2: -1.0}] - out = _ifd_family_metrics(cond_lp, asst_lp, [0, 1, 2], [1, 2], n_prompt) - assert out['n_tokens'] == 2 - assert abs(out['mean_delta'] - 0.5) < 1e-9 - assert abs(out['ifd'] - math.exp(-0.5)) < 1e-9 - # s_ifd_50 keeps top-1 by |delta| = 0.5; s_ifd_75 keeps top-2 (rounded up). - assert abs(out['s_ifd_50'] - math.exp(-0.5)) < 1e-9 - assert abs(out['s_ifd_75'] - math.exp(-0.5)) < 1e-9 - - def test_mixed_deltas(self): - # deltas = [+2.0, -1.0]; mean = 0.5 - n_prompt = 1 - cond_lp = [{0: 0.0}, {1: 1.0}, {2: -2.0}] - asst_lp = [{1: -1.0}, {2: -1.0}] - out = _ifd_family_metrics(cond_lp, asst_lp, [0, 1, 2], [1, 2], n_prompt) - assert out['n_tokens'] == 2 - assert abs(out['mean_delta'] - 0.5) < 1e-9 - # s_ifd_50 keeps top-1 by |delta| = 2.0 → exp(-2.0) - assert abs(out['s_ifd_50'] - math.exp(-2.0)) < 1e-9 - - -# ── _mean_logprob_delta ───────────────────────────────────────────────────── - - -class TestMeanLogprobDelta: - - def test_empty(self): - assert _mean_logprob_delta([], [{1: -1.0}], [], [1], 0) is None - - def test_uniform_delta(self): - n_prompt = 1 - cond_lp = [{0: 0.0}, {1: -0.5}, {2: -0.5}] - asst_lp = [{1: -1.0}, {2: -1.0}] - out = _mean_logprob_delta(cond_lp, asst_lp, [0, 1, 2], [1, 2], n_prompt) - assert abs(out - 0.5) < 1e-9 - - def test_mixed_average(self): - # deltas = [+2.0, -1.0] → mean 0.5 - n_prompt = 1 - cond_lp = [{0: 0.0}, {1: 1.0}, {2: -2.0}] - asst_lp = [{1: -1.0}, {2: -1.0}] - out = _mean_logprob_delta(cond_lp, asst_lp, [0, 1, 2], [1, 2], n_prompt) - assert abs(out - 0.5) < 1e-9 - - def test_skips_none_logprobs(self): - # When asst lp returns None, that position is skipped silently. - n_prompt = 1 - cond_lp = [{0: 0.0}, {1: -0.5}, {2: -0.5}] - asst_lp = [None, {2: -1.0}] - out = _mean_logprob_delta(cond_lp, asst_lp, [0, 1, 2], [1, 2], n_prompt) - assert abs(out - 0.5) < 1e-9 # only position 1 used - - -# ── _lp_to_jsonable ───────────────────────────────────────────────────────── - - -class TestLpToJsonable: - - def test_none_input(self): - assert _lp_to_jsonable(None) == [] - - def test_empty(self): - assert _lp_to_jsonable([]) == [] - - def test_none_passthrough(self): - assert _lp_to_jsonable([None, None]) == [None, None] - - def test_scalar_to_float(self): - assert _lp_to_jsonable([1, -2.0]) == [1.0, -2.0] - - def test_dict_with_logprob_object(self): - - class Entry: - - def __init__(self, lp, rank, decoded): - self.logprob = lp - self.rank = rank - self.decoded_token = decoded - - out = _lp_to_jsonable([{7: Entry(-0.5, 1, 'hello')}]) - assert out == [{'7': {'logprob': -0.5, 'rank': 1, 'decoded': 'hello'}}] - - def test_dict_with_nested_dict(self): - out = _lp_to_jsonable([{7: {'logprob': -0.5}}]) - assert out == [{'7': {'logprob': -0.5}}] - - def test_dict_with_repr_fallback(self): - # Non-dict, non-Entry value falls back to repr string. - out = _lp_to_jsonable([{7: 'plain'}]) - assert out == [{'7': repr('plain')}] - - def test_non_dict_non_scalar_repr(self): - # An object that isn't dict/scalar gets repr-ed. - out = _lp_to_jsonable([(1, 2)]) - assert out == [repr((1, 2))] - - -# ── _pad_batch ────────────────────────────────────────────────────────────── - - -class TestPadBatch: - - def test_empty_batch(self): - padded, n = _pad_batch([], floor=4) - assert padded == [] - assert n == 0 - - def test_already_at_floor(self): - batch = [[1], [2], [3], [4]] - padded, n = _pad_batch(batch, floor=4) - assert padded == batch - assert n == 4 - - def test_above_floor(self): - batch = [[1], [2], [3], [4], [5]] - padded, n = _pad_batch(batch, floor=3) - assert padded == batch # unchanged - assert n == 5 - - def test_below_floor_pads_with_last(self): - batch = [[1], [2]] - padded, n = _pad_batch(batch, floor=4) - assert padded == [[1], [2], [2], [2]] - assert n == 2 # original size - - def test_returns_new_list(self): - batch = [[1], [2]] - padded, _ = _pad_batch(batch, floor=4) - # Mutating padded should not affect original. - padded.append([99]) - assert batch == [[1], [2]] - - -if __name__ == '__main__': - pytest.main([__file__, '-v']) diff --git a/tests/twinkle_agentic/test_multi_turn_rollout.py b/tests/twinkle_agentic/test_multi_turn_rollout.py index 56c02b2a8..15541ba10 100644 --- a/tests/twinkle_agentic/test_multi_turn_rollout.py +++ b/tests/twinkle_agentic/test_multi_turn_rollout.py @@ -23,6 +23,7 @@ import json import pytest import re +import threading from typing import Any, Dict, List, Optional from twinkle.data_format.sampling import SampledSequence, SampleResponse, SamplingParams @@ -215,15 +216,40 @@ def concat_input_feature(self, pif: dict[str, Any], new_tokens: list[int]) -> di class FakeSampler: - """Queue-driven sampler that mirrors VLLMSampler output shape.""" + """Queue-driven sampler that mirrors VLLMSampler output shape. + + ``queue`` feeds one shared FIFO, which is all a single-trajectory test needs. + A batch needs ``queue_for(key, ...)``: episodes run in parallel threads, so + the order in which their turns reach ``sample`` is not defined, and a shared + FIFO would hand one trajectory's scripted reply to another. The key is the + text of the trajectory's first user message. + """ def __init__(self, template: FakeTemplate) -> None: self.template = template self._queue: list[dict[str, Any]] = [] + self._keyed: dict[str, list[dict[str, Any]]] = {} self.sample_calls = 0 # One entry per sample() call, so a test can assert which budget each # stage was sampled under. self.params_seen: list[Any] = [] + self._lock = threading.Lock() + + @staticmethod + def _entry( + template: FakeTemplate, + response_text: str, + stop_reason: str, + logprobs: list[Any] | None, + append_im_end: bool, + ) -> dict[str, Any]: + raw = response_text + ('<|im_end|>' if append_im_end else '') + return { + 'tokens': template.tokenizer.encode(raw, add_special_tokens=False), + 'decoded': response_text, + 'stop_reason': stop_reason, + 'logprobs': logprobs, + } def queue( self, @@ -236,14 +262,26 @@ def queue( ``<|im_end|>`` is appended to the encoded tokens when ``append_im_end``. ``seq.decoded`` is the raw response WITHOUT the trailing <|im_end|> (matches vLLM's common behaviour).""" - raw = response_text + ('<|im_end|>' if append_im_end else '') - tokens = self.template.tokenizer.encode(raw, add_special_tokens=False) - self._queue.append({ - 'tokens': tokens, - 'decoded': response_text, - 'stop_reason': stop_reason, - 'logprobs': logprobs, - }) + self._queue.append(self._entry(self.template, response_text, stop_reason, logprobs, append_im_end)) + + def queue_for( + self, + key: str, + response_text: str, + stop_reason: str = 'stop', + logprobs: list[Any] | None = None, + append_im_end: bool = True, + ) -> None: + """Script one turn for the trajectory whose first user message is ``key``.""" + self._keyed.setdefault(key, []).append( + self._entry(self.template, response_text, stop_reason, logprobs, append_im_end)) + + @staticmethod + def _key_of(pif: dict[str, Any]) -> str | None: + for m in pif.get('messages') or []: + if m.get('role') == 'user': + return m.get('content') + return None def sample(self, pifs, sampling_params=None): # Batched contract: accept a list of pifs, return one @@ -252,12 +290,14 @@ def sample(self, pifs, sampling_params=None): if isinstance(pifs, dict): pifs = [pifs] assert isinstance(pifs, list), (f'FakeSampler.sample expects a list, got {type(pifs).__name__}') - self.params_seen.append(sampling_params) responses: list[SampleResponse] = [] for pif in pifs: - assert self._queue, 'FakeSampler queue exhausted — scripted turns' - r = self._queue.pop(0) - self.sample_calls += 1 + with self._lock: + self.params_seen.append(sampling_params) + queue = self._keyed.get(self._key_of(pif)) or self._queue + assert queue, 'FakeSampler queue exhausted — scripted turns' + r = queue.pop(0) + self.sample_calls += 1 new_pif = self.template.concat_input_feature(pif, r['tokens']) seq = SampledSequence( stop_reason=r['stop_reason'], @@ -269,6 +309,10 @@ def sample(self, pifs, sampling_params=None): responses.append(SampleResponse(sequences=[seq])) return responses + # MultiTurnRollout samples one trajectory per call and refuses a sampler + # that would slice such a batch across workers. + sample._enable_continous_work = True + class EchoTool(Tool): """Echoes its arguments as a JSON string.""" @@ -607,12 +651,13 @@ def test_stuck_stop_is_per_trajectory_in_a_batch(make_rollout, sampler, template bad = ToolManager({}) bad.register(FailTool('search')) - sampler.queue(_tool_call_text('search', {'q': 1}), stop_reason='stop') - sampler.queue(_tool_call_text('search', {'q': 1}), stop_reason='stop') - sampler.queue(_tool_call_text('search', {'q': 2}), stop_reason='stop') - sampler.queue(_tool_call_text('search', {'q': 3}), stop_reason='stop') - sampler.queue('Done.', stop_reason='stop') - sampler.queue('Done.', stop_reason='stop') + sampler.queue_for('a', _tool_call_text('search', {'q': 1}), stop_reason='stop') + sampler.queue_for('a', _tool_call_text('search', {'q': 2}), stop_reason='stop') + sampler.queue_for('a', _tool_call_text('search', {'q': 3}), stop_reason='stop') + sampler.queue_for('a', 'Done.', stop_reason='stop') + # 'b' calls the failing tool twice, which trips stop_after_stuck_turns=2. + sampler.queue_for('b', _tool_call_text('search', {'q': 1}), stop_reason='stop') + sampler.queue_for('b', _tool_call_text('search', {'q': 1}), stop_reason='stop') rollout = MultiTurnRollout( sampler=sampler, template=template, tool_manager=[good, bad], @@ -766,6 +811,28 @@ def test_rejects_num_samples_gt_1(sampler, template, tool_manager): sampling_params=SamplingParams(num_samples=2)) +def test_rejects_sampler_without_continous_work(template, tool_manager): + """A batch of one is what a slice_dp sampler cannot serve.""" + + class SlicingSampler: + + def sample(self, pifs, sampling_params=None): + return [] + + with pytest.raises(ValueError, match='enable_continous_work'): + MultiTurnRollout(sampler=SlicingSampler(), template=template, tool_manager=tool_manager) + + +def test_rejects_one_harness_shared_by_a_batch(sampler, template, tool_manager): + """Episodes run in parallel threads, so a stateful harness cannot be shared.""" + from twinkle_agentic.harness.base import AgentHarness + + rollout = MultiTurnRollout( + sampler=sampler, template=template, tool_manager=tool_manager, max_turns=2, harness=AgentHarness()) + with pytest.raises(ValueError, match='harness holds per-episode state'): + rollout([_user_traj('A'), _user_traj('B')]) + + # ============================================================================= # Tests: defensive guards # ============================================================================= @@ -779,6 +846,8 @@ def sample(self, pifs, sampling_params=None): seq = SampledSequence(stop_reason='stop', tokens=[], logprobs=None, decoded='', new_input_feature=None) return [SampleResponse(sequences=[seq]) for _ in pifs] + sample._enable_continous_work = True + rollout = MultiTurnRollout(sampler=BrokenSampler(), template=template, tool_manager=tool_manager) with pytest.raises(RuntimeError, match='new_input_feature'): rollout([_user_traj()]) @@ -791,6 +860,8 @@ class EmptySampler: def sample(self, pifs, sampling_params=None): return [] + sample._enable_continous_work = True + rollout = MultiTurnRollout(sampler=EmptySampler(), template=template, tool_manager=tool_manager) # Batched contract: 0 responses for a batch of 1 → mismatch error. with pytest.raises(RuntimeError, match='0 responses'): @@ -806,6 +877,8 @@ def sample(self, pifs, sampling_params=None): pifs = [pifs] return [SampleResponse(sequences=[]) for _ in pifs] + sample._enable_continous_work = True + rollout = MultiTurnRollout(sampler=NoSeqSampler(), template=template, tool_manager=tool_manager) with pytest.raises(RuntimeError, match='no sequences'): rollout([_user_traj()]) @@ -820,18 +893,17 @@ def test_empty_batch_returns_empty_list(make_rollout): def test_batch_single_turn_two_trajectories(make_rollout, sampler): - """Two trajectories finish on turn 1 → one batched sample call.""" - sampler.queue('answer-A', stop_reason='stop') - sampler.queue('answer-B', stop_reason='stop') + """Two trajectories, one turn each, in their own threads.""" + sampler.queue_for('Q-A', 'answer-A', stop_reason='stop') + sampler.queue_for('Q-B', 'answer-B', stop_reason='stop') rollout = make_rollout(max_turns=3) outs = rollout([_user_traj('Q-A'), _user_traj('Q-B')]) assert len(outs) == 2 - # Exactly ONE batched sample call, not two. - assert sampler.sample_calls == 2 # one per item, still one turn - # But FakeSampler counts per-input; the critical batching invariant is - # that MultiTurnRollout only calls sampler.sample ONCE per turn. We - # enforce this via the queue ordering + single turn. + assert sampler.sample_calls == 2 # one generation per trajectory + # Results come back in input order even though the threads may not. + assert outs[0]['messages'][-1]['content'] == 'answer-A' + assert outs[1]['messages'][-1]['content'] == 'answer-B' for out in outs: assert out['turns'] == 1 assert out['stop_reason'] == 'stop' @@ -841,14 +913,12 @@ def test_batch_single_turn_two_trajectories(make_rollout, sampler): def test_batch_different_termination_turns(make_rollout, sampler): """Trajectory A finishes on turn 1; trajectory B needs a tool turn. - Turn 1 batch: [A: 'done-A' stop, B: tool_call stop] → A parked. - Turn 2 batch: [B: 'done-B' stop] → only B live. + Each episode owns its turn budget, so B taking a second turn neither waits + for A nor buys A anything. """ - sampler.queue('done-A', stop_reason='stop') # A turn 1 - sampler.queue( - _tool_call_text('search', {'q': 'b'}), # B turn 1 - stop_reason='stop') - sampler.queue('done-B', stop_reason='stop') # B turn 2 + sampler.queue_for('Q-A', 'done-A', stop_reason='stop') + sampler.queue_for('Q-B', _tool_call_text('search', {'q': 'b'}), stop_reason='stop') + sampler.queue_for('Q-B', 'done-B', stop_reason='stop') rollout = make_rollout(max_turns=4) outs = rollout([_user_traj('Q-A'), _user_traj('Q-B')]) @@ -889,10 +959,10 @@ def tool_info(self): tm_b = ToolManager({}) tm_b.register(TagTool('B')) - sampler.queue(_tool_call_text('search', {'q': 'x'}), stop_reason='stop') - sampler.queue(_tool_call_text('search', {'q': 'y'}), stop_reason='stop') - sampler.queue('done-A', stop_reason='stop') - sampler.queue('done-B', stop_reason='stop') + sampler.queue_for('A', _tool_call_text('search', {'q': 'x'}), stop_reason='stop') + sampler.queue_for('A', 'done-A', stop_reason='stop') + sampler.queue_for('B', _tool_call_text('search', {'q': 'y'}), stop_reason='stop') + sampler.queue_for('B', 'done-B', stop_reason='stop') rollout = MultiTurnRollout( sampler=sampler, @@ -1014,8 +1084,8 @@ def _is_success(traj): max_turns=2, trace_dir=str(trace_dir), success_callback=_is_success) - sampler.queue('good answer', stop_reason='stop') - sampler.queue('bad answer', stop_reason='stop') + sampler.queue_for('A', 'good answer', stop_reason='stop') + sampler.queue_for('B', 'bad answer', stop_reason='stop') rollout([_user_traj('A'), _user_traj('B')]) @@ -1031,9 +1101,9 @@ def test_trace_dir_batch_writes_one_file_per_trajectory(tmp_path, sampler, templ rollout = MultiTurnRollout( sampler=sampler, template=template, tool_manager=tool_manager, max_turns=4, trace_dir=str(trace_dir)) # Traj 0: stops turn 1. Traj 1: tool-calls turn 1, stops turn 2. - sampler.queue('done0', stop_reason='stop') - sampler.queue(_tool_call_text('search', {'q': 'y'})) - sampler.queue('done1', stop_reason='stop') + sampler.queue_for('A', 'done0', stop_reason='stop') + sampler.queue_for('B', _tool_call_text('search', {'q': 'y'})) + sampler.queue_for('B', 'done1', stop_reason='stop') rollout([_user_traj('A'), _user_traj('B')]) From b285f44ad27da61e287b3e0562ff71db24042310 Mon Sep 17 00:00:00 2001 From: root <yuze.zyz@alibaba-inc.com> Date: Thu, 10 Sep 2026 00:58:43 +0800 Subject: [PATCH 60/60] wip --- src/twinkle_agentic/challenger/__init__.py | 32 +- src/twinkle_agentic/challenger/agentic.py | 2061 ++++------------- src/twinkle_agentic/challenger/api.py | 154 -- src/twinkle_agentic/challenger/base.py | 532 +---- src/twinkle_agentic/challenger/code.py | 616 ----- .../challenger/{new => }/keyword.py | 0 src/twinkle_agentic/challenger/keywords.py | 588 ----- .../challenger/new/__init__.py | 12 - src/twinkle_agentic/challenger/new/agentic.py | 511 ---- src/twinkle_agentic/challenger/new/base.py | 132 -- .../challenger/{new => }/recorder.py | 0 src/twinkle_agentic/challenger/task_bank.py | 130 -- 12 files changed, 527 insertions(+), 4241 deletions(-) delete mode 100644 src/twinkle_agentic/challenger/api.py delete mode 100644 src/twinkle_agentic/challenger/code.py rename src/twinkle_agentic/challenger/{new => }/keyword.py (100%) delete mode 100644 src/twinkle_agentic/challenger/keywords.py delete mode 100644 src/twinkle_agentic/challenger/new/__init__.py delete mode 100644 src/twinkle_agentic/challenger/new/agentic.py delete mode 100644 src/twinkle_agentic/challenger/new/base.py rename src/twinkle_agentic/challenger/{new => }/recorder.py (100%) delete mode 100644 src/twinkle_agentic/challenger/task_bank.py diff --git a/src/twinkle_agentic/challenger/__init__.py b/src/twinkle_agentic/challenger/__init__.py index 7c0a6261d..0719b04e5 100644 --- a/src/twinkle_agentic/challenger/__init__.py +++ b/src/twinkle_agentic/challenger/__init__.py @@ -1,36 +1,12 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -from .agentic import AgenticChallenger, AgenticPrompts, parse_check_script, parse_problem_statement -from .api import ApiExplorer, ApiModel -from .base import Challenger, Explorer, PromptSet, attach_user_data, map_parallel -from .code import (CodeChallenger, CodePrompts, build_asserts, is_constant_answer, load_seeds, parse_challenge, - run_asserts, run_check_script) -from .keywords import (KEYWORD_MAX_LEN, KeywordBank, KeywordPrompts, KeywordStore, parse_keyword_list, - split_keyword_list) +from .agentic import AgenticChallenger, parse_problem_statement +from .base import Challenger +from .keyword import KEYWORD_MAX_LEN, KeywordGenerator __all__ = [ 'AgenticChallenger', - 'AgenticPrompts', - 'ApiExplorer', - 'ApiModel', 'Challenger', - 'CodeChallenger', - 'CodePrompts', - 'Explorer', 'KEYWORD_MAX_LEN', - 'KeywordBank', - 'KeywordPrompts', - 'KeywordStore', - 'PromptSet', - 'attach_user_data', - 'build_asserts', - 'is_constant_answer', - 'load_seeds', - 'map_parallel', - 'parse_check_script', - 'parse_challenge', - 'parse_keyword_list', + 'KeywordGenerator', 'parse_problem_statement', - 'run_asserts', - 'run_check_script', - 'split_keyword_list', ] diff --git a/src/twinkle_agentic/challenger/agentic.py b/src/twinkle_agentic/challenger/agentic.py index d7c6bab2d..ca41486a7 100644 --- a/src/twinkle_agentic/challenger/agentic.py +++ b/src/twinkle_agentic/challenger/agentic.py @@ -1,1682 +1,511 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Agentic challenger: invent tasks by doing them first. - -The approach mirrors how the code challenger works, adapted to tool-using -agents. Instead of writing a problem statement and hoping it is achievable, the -model first *does* something in a sandbox, then -- in the same conversation -- -writes the check script that verifies the end state it just produced, then the -problem statement someone else would need to reproduce it. - -Steps for one candidate: - - 1. Choose direction + keywords. Optionally start from a seed trajectory. - 2. Explore (multi-turn with tools): model acts in a clean sandbox, producing - a tool-call chain and a final workspace state, and stops calling tools. - 3. A user message is appended to that same conversation carrying the - workspace listing, asking for a python check script. Tools are no longer - dispatched from here on. - 4. Verify: run the check script in the sandbox (must pass). - 5. A second user message is appended asking for the problem statement. - 6. Difficulty filter: reset workspace, let the solver do the task N times, - run checks, keep only "sometimes pass" tasks. - -Steps 3 and 5 are appended to the episode rather than sent as fresh calls, so -every assistant turn in the chain keeps its ``labels`` and ``logprobs`` and the -whole proposal -- acting, checking, describing -- is one trainable sample. The -follow-up messages come back from :meth:`AgenticChallenger._followup`, which the -rollout calls at the moment the model stops calling tools; that is where the -sandbox work (snapshot, running the check) happens, because only the caller can -do it. - -Because every episode needs a clean workspace and because episodes share a -single long-lived sandbox, proposing is **serial** -- one proposal at a time with -a workspace reset in between. - -Prompt text is not here. Every string the model sees arrives in -:class:`AgenticPrompts`, built by whoever runs the challenger -- see -``cookbook/rsi/agentic/prompts.py``. -""" -import ast +"""Agentic challenger: act in a sandbox, verify the result, then describe it.""" import math +import random import re -import threading -from concurrent.futures import ThreadPoolExecutor, as_completed +import uuid from dataclasses import dataclass -from functools import partial from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple -from twinkle.data_format import SamplingParams, Trajectory, user_data_get +from twinkle.data_format import SamplingParams, Trajectory, attach_user_data, user_data_get +from twinkle.data_format.sampling import SampledSequence, SampleResponse from twinkle.utils import get_logger -from twinkle_agentic.utils.code_utils import PYTHON_TAGS, parse_fenced_code, strip_reasoning -from twinkle_agentic.utils.message_utils import assistant_text -from .api import ApiModel -from .base import Challenger, Explorer, PromptSet, attach_user_data, map_parallel -from .keywords import KeywordBank, KeywordStore +from twinkle_agentic.envs import Env +from twinkle_agentic.protocol.base import API +from twinkle_agentic.rollout import APISampler, MultiTurnRollout +from twinkle_agentic.summarizer import Summarizer +from twinkle_agentic.utils.code_utils import parse_fenced_code, strip_reasoning +from twinkle_agentic.utils.message_utils import assistant_text, msg_content_text, normalize_tool_calls +from .base import Challenger, _parallel +from .keyword import KeywordGenerator +from .recorder import RolloutRecorder + +__all__ = ['AgenticChallenger', 'parse_problem_statement'] logger = get_logger() -__all__ = [ - 'AgenticChallenger', - 'AgenticPrompts', - 'DEFAULT_CHECK_PARSE_ERROR', - 'brittle_check_reason', - 'parse_check_script', - 'parse_problem_statement', -] - -# A fence around the *whole* reply, which is packaging rather than content. -_WHOLE_FENCE_RE = re.compile(r'```[\w+-]*\s*\n?(.*?)```', re.S) - - -# ── parsing ─────────────────────────────────────────────────────────────── - -def parse_check_script(text: str, language_tags: Tuple[str, ...] = PYTHON_TAGS) -> Optional[str]: - """Extract a check script from the model's reply. - - The script has to be in a fenced block; see - :func:`~twinkle_agentic.utils.code_utils.parse_fenced_code`. Returns ``None`` - otherwise, including when the model reached for a tool instead of answering -- - the proposing episode had tools, so that happens, and the reply is asked for - again rather than dug through. An empty fence is refused the same way: the - model marked where the script went and put nothing there. - - ``language_tags`` is the whole of what makes this python. Another language - passes its own, most cheaply as - ``AgenticChallenger(parse_check_fn=partial(parse_check_script, language_tags=...))``. - """ - return parse_fenced_code(text, language_tags) - - - -# Two rules CHECK_FOLLOWUP already states -- no equality on a script's source -# text, no byte count or checksum on a binary -- were broken by 9 and 8 of 41 -# measured tasks respectively, so stating them a third time is not the fix. A -# check that pins the exact source of a .py rejects every equivalent solution, -# and one that pins a .png's byte count rejects every matplotlib version; both -# make a task nobody but the author can pass. -_SIZE_OR_HASH_NAMES = ('getsize', 'st_size', 'sha256', 'sha1', 'md5', 'hexdigest', - 'digest') -# What makes a string python rather than data. Checked instead of "is it long and -# multi-line", because the contents of a csv or a json file are legitimately -# asserted verbatim -- the statement handed those to the solver -- while the text -# of a script never is. -_LOOKS_LIKE_PYTHON = ('import ', 'def ', 'print(', 'with open(', 'if __name__') - - -def brittle_check_reason(script: str) -> Optional[str]: - """Why this check script would reject a correct solution, or None. - - Returned text goes back to the model through the same retry path a failing - assertion uses, because the defect is the same kind: an assertion that does - not hold for solutions other than the one in front of it. - - Read off the syntax tree rather than matched as text. Both defects survive - patterns easily: source equality reads the file into a name first - (``c = f.read()``, then ``assert c == '...'``) so nothing sits between - ``open()`` and ``==``, and a size check can put the call either around the - name (``getsize("a.png")``) or after it. - - Python throughout -- the tree, the marker words, the stdlib names below. There - is no language-neutral version of this: another language keeps the two rules - but rewrites the whole body, which is why the challenger takes it as - ``brittle_check_fn`` rather than calling it directly. - """ - try: - tree = ast.parse(script) - except SyntaxError: - # Unparseable means it cannot run either, so let the sandbox report it. - return None - for node in ast.walk(tree): - if not (isinstance(node, ast.Compare) - and any(isinstance(o, ast.Eq) for o in node.ops)): - continue - for side in [node.left] + list(node.comparators): - if not (isinstance(side, ast.Constant) and isinstance(side.value, str)): - continue - if any(m in side.value for m in _LOOKS_LIKE_PYTHON): - return ('AssertionError: this check compares a file against the ' - 'full text of a python script with ==, which only the ' - 'exact script you wrote can pass. Assert what running ' - 'that script produces instead.') - # A byte count or a checksum compared for equality. Not restricted to - # binary suffixes: CHECK_FOLLOWUP says "NEVER assert a file size in bytes" - # about any file, and keying on a suffix list let - # ``getsize('data.mat') == 264`` through. Only equality against a literal is - # a defect -- ``getsize(f) > 0`` is a fine way to say "not empty". - for node in ast.walk(tree): - if not (isinstance(node, ast.Compare) - and any(isinstance(o, ast.Eq) for o in node.ops)): - continue - sides = [node.left] + list(node.comparators) - has_literal = any(isinstance(s, ast.Constant) - and isinstance(s.value, (int, float, str)) - and not isinstance(s.value, bool) for s in sides) - if not has_literal: - continue - for side in sides: - names = {n.attr for n in ast.walk(side) if isinstance(n, ast.Attribute)} - names |= {n.id for n in ast.walk(side) if isinstance(n, ast.Name)} - hit = names & set(_SIZE_OR_HASH_NAMES) - if hit: - what = ('a checksum' if hit - {'getsize', 'st_size'} - else 'a byte count') - return (f'AssertionError: this check pins {what} of a file, and ' - 'correct solutions differ there. Assert what can be read ' - 'out of the file instead -- its structure, or the values ' - 'inside it.') - # Comparing raw bytes of a file: same defect, different spelling. - for node in ast.walk(tree): - if not (isinstance(node, ast.Compare) - and any(isinstance(o, ast.Eq) for o in node.ops)): - continue - for side in [node.left] + list(node.comparators): - if isinstance(side, ast.Constant) and isinstance(side.value, bytes): - return ('AssertionError: this check compares the raw bytes of a ' - 'file, and correct solutions differ there. Assert what ' - 'can be read out of it instead.') - return None - - -# Literals shorter than this match by accident: a statement contains "3" or "id" -# for its own reasons. Measured on 188 tasks from run_clean9, one and two digit -# integers appeared in both the check and the statement 91% of the time, which is -# what an unattributable coincidence rate looks like. -_MIN_DERIVED_LEN = 3 - - -def derived_check_literals(script: str) -> List[str]: - """The values a check compares against that the solver is meant to work out. - - A measurement tool, not part of the proposing path. It exists because "does the - statement give away the answer" cannot be asked without first separating the - three kinds of thing a check's literals are, and only one of them is a leak: - - an identifier, or a name with a file extension - The statement MUST carry these. It is naming the file to create and the - fields to put in it; a statement that withheld them would describe no - particular output at all. Present in 84-93% of run_clean9's statements, - which is the correct rate. - text that appears in the workspace - Input data, which the statement is meant to quote verbatim so the solver - can write the same bytes. Not separated here -- the caller filters on - the snapshot if it wants to. - a long number, a float, or a string that is none of the above - Only exists once the work has been done. This is the group returned. - - Feeding the result back to the statement stage as a forbidden list was tried and - did not reduce the leak: see the note above PROBLEM_FOLLOWUP_RULES_ONLY in - cookbook/rsi/agentic/prompts.py for the two forms measured and their p-values. - - Wrong at the edges by construction: a column named ``total_2024`` reads as an - identifier and is not returned, and a computed value that lands on two digits is - below the length floor. - """ - try: - tree = ast.parse(script) - except SyntaxError: - return [] - out = [] - for node in ast.walk(tree): - if not isinstance(node, ast.Compare): - continue - for side in [node.left] + list(node.comparators): - if not isinstance(side, ast.Constant): - continue - v = side.value - if isinstance(v, bool) or v is None: - continue - if isinstance(v, (int, float)): - text = repr(v) - if len(text.lstrip('-').replace('.', '')) < _MIN_DERIVED_LEN: - continue - out.append(text) - elif isinstance(v, str): - if len(v) < _MIN_DERIVED_LEN: - continue - if re.fullmatch(r'[A-Za-z_][A-Za-z0-9_]*', v): - continue - # A trailing extension means a filename -- but only when the part - # after the dot is not itself digits. '0.001' read out of a CSV is a - # string here, and skipping it as a filename would let exactly the - # kind of value this function exists to catch through. - if re.search(r'\.[A-Za-z]\w{0,4}$', v) and ' ' not in v: - continue - out.append(v) - # Longest first: a short literal is often a substring of a longer one, and - # naming the long one first makes the list read as distinct values rather than - # as prefixes of each other. - return sorted(set(out), key=len, reverse=True) - +_FENCED_BLOCK_RE = re.compile(r'```[^\r\n]*\r?\n(.*?)```', re.S) def parse_problem_statement(text: str) -> Optional[str]: - """Extract a problem statement from the model's reply. - - Everything after the model's thinking is the statement. A fence around the - whole reply is unwrapped; fences *inside* it are kept. - - Keeping them matters more than it sounds: a statement that says what a file - must contain puts the content in a fence, and stripping every fence left - "1. `data.json` containing:" with nothing after it. 7 of ex11's 16 measured - statements had a fence, and 5 of those 7 were solved 0 times out of 8 -- - against 1 of the 9 statements that had no fence to lose. The tasks were not - hard, they were unanswerable. - - Returns ``None`` when the result is empty. - """ + """Return the statement after removing reasoning and one outer fence.""" body = strip_reasoning(text).strip() - whole = _WHOLE_FENCE_RE.fullmatch(body) + whole = _FENCED_BLOCK_RE.fullmatch(body) if whole: body = whole.group(1).strip() - return body if body else None - - -# The fields a local rollout splices into a trajectory, and the only ones a -# later GRPO step needs: ``labels`` marks which of ``input_ids`` are trainable -# (-100 elsewhere) and ``logprobs`` holds one entry per trainable token, taken -# from the policy that actually generated it. -_TRAINABLE_KEYS = ('input_ids', 'labels', 'logprobs') - - -def _propose_round(stage: str, trajectory: Trajectory) -> Dict[str, Any]: - """One proposing round, reduced to what a later training step would read. - - ``messages`` comes along for reading by humans; it is redundant with - ``input_ids`` and is not what a trainer should encode from. - """ - record: Dict[str, Any] = { - 'stage': stage, - 'messages': [dict(m) for m in trajectory.get('messages') or []], - } - for key in _TRAINABLE_KEYS: - value = trajectory.get(key) - if value is not None: - record[key] = value - return record - - -# ── prompts ──────────────────────────────────────────────────────────────── - -# The one piece of prompt text with a default, because it was written into the -# retry path before there was a field for it and every caller relies on it. It -# names the language, so a caller working in another one has to override it. -DEFAULT_CHECK_PARSE_ERROR = ('Could not read a check script from your reply: it was not a fenced ' - 'python code block. Do not wrap it in a tool call and do not add ' - 'prose -- return ONLY a fenced python code block.') + return body or None + + +def _sample_one(sampler: Any, input_feature: Dict[str, Any], sampling_params: Optional[SamplingParams], + adapter_kwargs: Dict[str, Any]) -> SampledSequence: + responses = sampler.sample([input_feature], sampling_params=sampling_params, **adapter_kwargs) + if not isinstance(responses, list): + raise TypeError(f'expected List[SampleResponse] from sampler.sample, got ' + f'{type(responses).__name__}') + if len(responses) != 1: + raise RuntimeError(f'sampler returned {len(responses)} responses for a single request; ' + 'expected exactly one') + response = responses[0] + if not isinstance(response, SampleResponse): + raise TypeError(f'expected SampleResponse from sampler.sample, got ' + f'{type(response).__name__}') + if len(response.sequences) != 1: + raise RuntimeError(f'SampleResponse contains {len(response.sequences)} sequences; ' + 'expected exactly one') + sequence = response.sequences[0] + if not isinstance(sequence, SampledSequence): + raise TypeError(f'expected SampledSequence, got {type(sequence).__name__}') + return sequence + + +def _api_followup_response( + sampler: Any, + api: Optional[APISampler], + sampling_params: Optional[SamplingParams], + *, + input_feature: Dict[str, Any], + adapter_kwargs: Dict[str, Any], + followups: int, + **kwargs: Any, +) -> SampledSequence: + """Use the API for appended stages and the primary backend otherwise.""" + if followups: + if api is None: + raise ValueError('use_api=True requires an API backend') + return api(input_feature, sampling_params, **adapter_kwargs) + if sampler is not None: + return _sample_one(sampler, input_feature, sampling_params, adapter_kwargs) + if api is not None: + return api(input_feature, sampling_params, **adapter_kwargs) + raise ValueError('AgenticChallenger has neither a sampler nor an API backend') @dataclass -class AgenticPrompts(PromptSet): - """Every string an :class:`AgenticChallenger` sends. - - All fields are injected by the caller (no defaults with real text here). - Placeholder validation, and the keyword subset a bank is given, are - :class:`.PromptSet`. - """ - - # Explore: model acts in sandbox - system: str - from_scratch: str - from_seed: str = '' - from_keywords: str = '' - from_seed_keywords: str = '' - - # Appended to the same conversation once the model stops calling tools: - # first "write the check script" (which carries the workspace listing), then - # "write the problem statement". Each has to repeat the rules that used to - # live in a system message of its own, because there is no second system - # message in a single conversation. - check_followup: str = '' - # Sent instead of the statement stage when the check script does not pass, so - # the model can fix it from the traceback. Required only when the challenger - # is built with ``check_retries`` above 0. - check_retry_followup: str = '' - # The error text ``check_retry_followup`` carries when the reply held no - # readable check script at all (as opposed to one that ran and failed). - # Empty means ``DEFAULT_CHECK_PARSE_ERROR``. - check_parse_error: str = '' - problem_followup: str = '' - - # Keyword generation (same structure as code side) - keyword_system: str = '' - keyword_user: str = '' - keyword_expand_user: str = '' +class _ProposalResult: + trajectory: Trajectory + group_id: str = '' + task: Optional[Trajectory] = None + reason: str = '' + detail: str = '' + outcome: str = '' + n_pass: Optional[int] = None + reward: float = 0.0 - _REQUIRED = ('system', 'from_scratch', 'check_followup', 'problem_followup') - _REQUIRED_FIELDS = { - 'from_seed': ('seed',), - 'from_keywords': ('keywords',), - 'from_seed_keywords': ('seed', 'keywords'), - 'check_followup': ('final_state',), - 'check_retry_followup': ('error', 'final_state'), - 'keyword_user': ('k', 'desc'), - 'keyword_expand_user': ('kw', 'm'), - } - - -# ── challenger ───────────────────────────────────────────────────────────── class AgenticChallenger(Challenger): - """Propose tool-using tasks by first doing them, then describing them. + """Invent tool-using tasks by doing, checking, and describing them. - Args: - prompts: every string sent to the model. - explorer: batch-in / batch-out generation with sandbox tools (multi-turn). - seeds: optional pool of seed trajectories (dicts with a ``query`` key), - drawn with replacement. - keyword_store: optional bank for diversity control. - category_desc: category -> description for keyword generation. - seed_mix_prob: chance a proposal carries a seed. - envs: see :class:`~.base.Challenger`. This half asks a slot to be a real - workspace: it clears it, lets the model act in it through - :meth:`~twinkle_agentic.envs.base.Env.tool_manager`, reads the end - state back with :meth:`~twinkle_agentic.envs.base.Env.snapshot` and - runs the check script in it. ``len(envs)`` is therefore also the - episode concurrency: an episode owns its slot from the clear until - its check has run, so two episodes cannot share one, and an episode - acting in one workspace while being checked against another produces - a task nobody can pass. - parse_check_fn: read a check script out of a reply, or return None. - Defaults to :func:`parse_check_script`, which asks only that the script - be fenced python. Whether it asserts anything is the caller's to - require -- in the prompt it writes, or in the function it passes here - instead. - brittle_check_fn: why a parsed script would reject a correct solution, - or None if it would not. Defaults to :func:`brittle_check_reason`, - which reads a python syntax tree; pass ``None`` to drop the check - entirely and judge scripts only by whether they run. Both of these - and ``prompts.check_parse_error`` are the language-bound trio -- a - caller working outside python replaces all three or none. - tool_schemas: the tool contract in the OpenAI shape the template renders. - ``None`` takes it off slot 0, which is the spelling that slot will - honour; pass a list only to advertise something narrower. Attached to - the trajectories that are *meant* to call tools -- the exploring - episode and each solve attempt. Without it the model is never told the - tool names, so it writes code in prose instead of calling anything: - the workspace stays empty, every check fails, and the difficulty - numbers describe a model that had no tools rather than a hard task. - The check-writing and problem-writing stages sit in the same - conversation and so see the same list, which is why the rollout stops - dispatching calls once a follow-up has been appended -- a python block - written as an *answer* parses as a call list, and 41 of 146 such - replies in a measured run edited the very workspace the answer was - about. - combo_arity / arity_weights / single_kw_prob / keyword_refill_target / - keyword_gen_calls / keyword_refill_concurrency / keyword_refill_tries / - keyword_params / keyword_explorer / keyword_sink / min_batch: handed to - the :class:`.keywords.KeywordBank` this challenger holds, which is - where they are documented -- they behave the same on the code half. - ``keyword_explorer`` defaults to ``explorer`` here, which for a - sandbox setup means the bank brainstorms with tools live. - proposals_per_group: how many proposals answer the same keyword draw and - the same prompt, tagged with a shared ``group_id``. This is the group - size the proposing side's advantage is computed over; at 1 every - group has one member and every advantage is zero. At a fixed - proposal count it does not change the compute -- it divides the - number of distinct keyword draws per round by the same factor. - check_params / problem_params: sampling params for the two appended - stages. ``None`` keeps whatever the episode was already using, which - is sized for one agent turn; the check-writing stage reads the whole - episode plus the end state and reasons at length before answering, - and one that runs out of budget mid-thought never emits its code - block and is thrown away as unparseable. - followup_api: optional OpenAI-compatible API client (e.g. qwen3.8-max). When - given, exploration still runs on the local explorer -- so its turns keep - their ``labels`` and ``logprobs`` and remain trainable -- but the - check-script (success judgement) and problem-statement stages are - generated by this API instead of the local model, and are appended - neither to the trainable trajectory nor to its token stream. This is - the "explore locally, judge and describe over an API, train only the - exploration" split. ``None`` keeps the single-model behaviour where the - local model writes those two stages in the same conversation. - followup_extra_body: extra request body forwarded on every ``followup_api`` - call (e.g. ``{'thinking_budget': N}`` to cap qwen3.8-max reasoning). - ``None`` sends the request unmodified. Ignored when ``followup_api`` is - ``None``. - problem_max_chars: reject problem statements longer than this. - check_retries: how many times a check script that did not pass is handed - back, with the traceback and the workspace listing, for a rewrite - before the proposal is rejected. 0 restores the old behaviour of - rejecting on the first failure. Measured in ex12: 36 of 72 proposals - died on the check, and 29 of those were a single assertion naming a - value the model had not read -- a row count, a content string that - was nearly right, a timestamp -- over a workspace state that was - perfectly good. Each retry costs one more sampling call for that - episode and nothing for the ones that pass first time. - reject_sink: called with a dict for every rejected proposal. - propose_sink: called once per proposal attempt -- kept, rejected while - building, or dropped by the difficulty band alike -- with the - token-level record of the episode that produced it. This is the only - way the proposing episode survives: it is generation like any - other, so it carries ``input_ids`` / ``labels`` / ``logprobs`` and - could later be trained on, but nothing downstream of ``build`` - looks at it and without a sink it is dropped on the floor. - Rejects are included on purpose: they are the zero-reward half of a - GRPO group, so a set of kept-only records has no variance to learn - from. Requires a local sampler -- an API explorer returns text only. - solver_sink: called once per solver attempt in the difficulty stage, with - the statement, the check script, the attempt, the workspace it left - and the check's verdict. ``n_pass`` alone cannot distinguish a task - that is impossible from one whose statement withholds a value its - check demands, and both look like a hard task worth keeping. + ``backend`` drives exploration and solver attempts. When ``use_api`` is true, + ``api`` generates only the appended check-script and problem-statement turns; + those turns retain the masking semantics selected by ``api_appended_as`` in + ``rollout_kwargs``. """ + _system = ('You invent tasks for another agent to solve. You have a sandbox and ' + 'tools. Work in it first: build something real, then you will be asked ' + 'to verify it and to describe it.') + _from_scratch = ('Choose a task worth doing in this sandbox and do it now, using ' + 'your tools. Do not describe it yet.') + _from_keywords = ('Choose a task around these topics and do it now, using your ' + 'tools. Do not describe it yet.\n\nTopics: {keywords}') + _from_seed = ('Here is an earlier task:\n\n{seed}\n\nDo something in the same ' + 'spirit but different, using your tools now. Do not describe it yet.') + _from_seed_keywords = ('Here is an earlier task:\n\n{seed}\n\nDo something in the same ' + 'spirit but different, may be more complex and interesting and meaningful, ' + 'around these topics, using your tools now. Do not describe it yet.\n\n' + 'Topics: {keywords}') + _check_followup = ('Stop working. This is the workspace you produced:\n\n{final_state}\n\n' + 'Write a {language} script that verifies this end state, as a fenced ' + '{language} code block and nothing else. It must exit with a non-zero status ' + 'if the work was not done. Check what can be read out of the files -- their ' + 'structure and the values inside them. NEVER check a file size in bytes, a ' + 'checksum, or the full source text of a script: correct solutions differ ' + 'there, and such a check only its own author can pass.') + _check_retry_followup = ('Your check script did not pass:\n\n{error}\n\nThe workspace is:\n\n' + '{final_state}\n\nReturn a corrected script as a fenced {language} code block ' + 'and nothing else.') + _check_parse_error = ('Could not read a check script from your reply: it was not a ' + 'fenced {language} code block. Do not wrap it in a tool call and ' + 'do not add prose -- return ONLY a fenced {language} code block.') + _problem_followup = ('Now write the task statement: what someone starting from an empty workspace ' + 'would have to be told to produce what you produced, and nothing about how you ' + 'did it. Name the files to create and quote any input data verbatim. Do not ' + 'reveal values your check script computes. Reply with the statement only.') + + def __init__( self, - prompts: AgenticPrompts, - explorer: Explorer, + backend: Any, *, - seeds: Sequence[Dict[str, Any]] = (), - keyword_store: Optional[KeywordStore] = None, - category_desc: Optional[Dict[str, str]] = None, - seed_mix_prob: float = 0.5, - parse_check_fn: Callable[[str], Optional[str]] = parse_check_script, - brittle_check_fn: Optional[Callable[[str], Optional[str]]] = brittle_check_reason, - tool_schemas: Optional[Sequence[Dict[str, Any]]] = None, - combo_arity: str = 'triple', - arity_weights: Optional[Sequence[float]] = None, - single_kw_prob: float = 0.1, - proposals_per_group: int = 1, - keyword_refill_target: int = 128, - keyword_gen_calls: int = 8, - keyword_refill_concurrency: int = 1, - keyword_refill_tries: int = 2, - keyword_params: Optional[SamplingParams] = None, - check_params: Optional[SamplingParams] = None, - problem_params: Optional[SamplingParams] = None, - followup_api: Optional[Any] = None, - followup_extra_body: Optional[Dict[str, Any]] = None, - keyword_explorer: Optional[Explorer] = None, - min_batch: int = 1, - problem_max_chars: int = 8192, - max_proposals_total: int = 0, - setup_script_fn: Optional[Callable[..., str]] = None, - solver_prompt_fn: Optional[Callable[[str], Trajectory]] = None, + api: Optional[Any] = None, + use_api: bool = False, + keyword_generator: Optional[KeywordGenerator] = None, + trajectory_seed: Optional[List[Trajectory]] = None, + summarizer: Optional[Summarizer] = None, + system_prompt: Optional[str] = None, + from_scratch_prompt: Optional[str] = None, + from_keywords_prompt: Optional[str] = None, + from_seed_prompt: Optional[str] = None, + from_seed_keywords_prompt: Optional[str] = None, + check_followup_prompt: Optional[str] = None, + check_retry_followup_prompt: Optional[str] = None, + check_parse_error_prompt: Optional[str] = None, + problem_followup_prompt: Optional[str] = None, check_retries: int = 1, - task_bank: Optional[Any] = None, - novelty_fn: Optional[Callable[[List[Dict[str, Any]]], List[Optional[float]]]] = None, - novelty_floor: float = 0.5, - keep_per_group: int = 0, - reject_sink: Optional[Callable[[Dict[str, Any]], None]] = None, - propose_sink: Optional[Callable[[Dict[str, Any]], None]] = None, - solver_sink: Optional[Callable[[Dict[str, Any]], None]] = None, - keyword_sink: Optional[Callable[[Dict[str, Any]], None]] = None, - **challenger_kwargs: Any, + problem_max_chars: int = 8192, + check_language: str = 'python', + parse_check_fn: Optional[Callable[[str], Optional[str]]] = None, + pass_rate_target: float = 0.2, + envs: Sequence[Env] = (), + num_challenger_rollouts: int = 8, + num_solver_rollouts: int = 8, + pass_band: Tuple[float, float] = (1.0, 7.0), + pass_rate_width: float = 0.3, + max_empty_rounds: int = 0, + followup_params: Optional[SamplingParams] = None, + checker: Optional[Callable[[Trajectory], bool]] = None, + save_dir: Optional[str] = None, + save_failed_rollouts: bool = True, + **rollout_kwargs: Any, ): - super().__init__(explorer, system=prompts.system, **challenger_kwargs) - if not self.envs: - raise ValueError('envs is empty: an episode here needs a workspace to act in ' - 'and a check to be run against, so there is nothing this ' - 'challenger could measure.') - if keyword_store is not None: - prompts.require('from_keywords') - self.prompts = prompts - self.seeds = list(seeds) - # The whole keyword cycle -- draw, refill, expand -- is one object shared - # with the code challenger rather than a second copy of it here. None means - # no bank was configured, and proposals then carry no topics. - self.keywords: Optional[KeywordBank] = None if keyword_store is None else KeywordBank( - keyword_store, prompts=prompts.keyword_prompts(), - category_desc=category_desc or {}, - # Brainstorming a list is a text round: the sandbox-tool explorer would - # waste turns on it and could take a bracketed list for a tool call. - explorer=keyword_explorer or explorer, rng=self.rng, - name=type(self).__name__, sampling_params=keyword_params, - sink=keyword_sink, combo_arity=combo_arity, arity_weights=arity_weights, - single_kw_prob=single_kw_prob, refill_target=keyword_refill_target, - gen_calls=keyword_gen_calls, refill_concurrency=keyword_refill_concurrency, - refill_tries=keyword_refill_tries, min_batch=min_batch) - self.seed_mix_prob = seed_mix_prob - self.parse_check_fn = parse_check_fn - self.brittle_check_fn = brittle_check_fn - # Read off slot 0 by default: these go into the prompt, and taking them - # from the environment that will execute them is what makes it impossible - # for the advertised contract and the running code to disagree. - self.tool_schemas = list(tool_schemas) if tool_schemas else (self.env().tools() or None) - # Held while writing to the dump files and while bumping ``stats``: with - # concurrent episodes those are the only shared mutable things the - # follow-up callback touches, and a half-written json line is unreadable. - self._sink_lock = threading.Lock() - # How many proposals answer each keyword draw. Above 1 they form a GRPO - # group on the proposing side; see :meth:`propose`. Raising it does not - # cost more compute at a fixed proposal count -- it trades keyword - # variety for group size, since a round's proposals then come from - # ``count / proposals_per_group`` draws instead of ``count`` of them. - if proposals_per_group < 1: - raise ValueError(f'proposals_per_group must be >= 1, got {proposals_per_group}') - self.proposals_per_group = proposals_per_group - self._next_group_id = 0 - self.check_params = check_params - self.problem_params = problem_params - # When set, exploration runs on the (trainable) local explorer as before, - # but the check-script and problem-statement stages are generated by this - # OpenAI-compatible API (e.g. qwen3.8-max) instead of the local model. The - # two stages then contribute nothing to the trainable trajectory: the API - # returns text only, so the episode's ``input_ids`` / ``labels`` / - # ``logprobs`` stay exactly the exploration turns the local sampler - # produced -- which is what "train only the exploration part" means. The - # generated check script and statement are used solely to build the task. - # None means the single-model path, where the local model writes those two - # stages in the same conversation. ``followup_extra_body`` rides along on - # every call (e.g. {'thinking_budget': N} to cap qwen3.8-max reasoning). - self.followup_model: Optional[ApiModel] = None if followup_api is None else ApiModel( - followup_api, extra_body=followup_extra_body, name=type(self).__name__) - self.problem_max_chars = problem_max_chars - # A budget in proposals rather than in kept tasks, for runs whose purpose - # is to measure what the current configuration produces: with a keep-rate - # near 6% a keep-target of 8 is 128 proposals, and comparing two - # configurations means giving them the same number of tries, not the same - # output. 0 leaves the run governed by its keep-target. - self.max_proposals_total = max_proposals_total - # Arm B. Returns a python script that recreates this episode's input files, - # captured while the workspace still holds them, and replayed before every - # solver attempt. None leaves the solver starting from an empty directory. - self.setup_script_fn = setup_script_fn - # How a task statement becomes the solver's opening conversation. Without - # one, the solver is handed the statement as a bare user message and no - # system message at all -- nothing says it is working in a sandbox, that it - # may take many turns, or that a reply carries one tool call. Measured on - # arm B: 71 of 80 attempts used 2-3 turns, writing 8-12k characters into a - # single python_executor argument and truncating there, so ``n_pass`` was - # reporting that omission rather than the task. Passing the same function - # the eval script uses is what keeps the two measuring the same thing. - self.solver_prompt_fn = solver_prompt_fn + super().__init__( + envs=envs, + num_challenger_rollouts=num_challenger_rollouts, + num_solver_rollouts=num_solver_rollouts, + pass_band=pass_band, + max_empty_rounds=max_empty_rounds, + ) if check_retries < 0: raise ValueError(f'check_retries must be >= 0, got {check_retries}') - self.check_retries = check_retries - if check_retries: - prompts.require('check_retry_followup') - # Novelty, off unless both of these are given. ``task_bank`` supplies the - # tasks earlier iterations produced (see :mod:`.task_bank`); ``novelty_fn`` - # takes a list of ``{statement, check, references}`` and returns one score in - # [0, 1] per entry, or None where it could not judge. Kept as injected - # callables for the same reason the sandbox ones are: this class then holds - # no opinion about which judge, model or API produces the number, and a test - # can hand it a fixed one. - self.task_bank = task_bank - self.novelty_fn = novelty_fn - if not 0.0 <= novelty_floor <= 1.0: - raise ValueError(f'novelty_floor must be in [0, 1], got {novelty_floor}') - # How much of the reward a proposal keeps when it is judged fully redundant. - # 0.5 halves it; 0.0 would be Ornith-1.5's plain ``V x D x N``, which zeroes - # it. The floor exists because our N is coarse where theirs is continuous: - # scored over run_clean9's 188 tasks, 44% came out at exactly 0.0, and eight - # proposals sharing one keyword draw can all land there -- at floor 0 that - # group's rewards are all zero, its advantages are all zero after GRPO - # subtracts the mean, and eight sandbox rollouts bought nothing. Ornith's own - # text says novelty 'should remain secondary to validity and difficulty'. - self.novelty_floor = float(novelty_floor) - # At most this many of a keyword group's in-band proposals become tasks the - # solver side trains on. 0 keeps every in-band proposal, which is what this - # did before. At 1 the two sides come out the same size -- eight groups of - # eight proposals give 64 proposing trajectories and 8 tasks x 8 attempts = - # 64 solving ones -- and the tasks are one per keyword direction instead of - # three from the same one. - # - # The proposals not selected are NOT wasted from the proposing side: each one - # still earns its own reward from its own n_pass, so the whole group still - # trains. What is dropped is their solver attempts, which were already run to - # measure difficulty: at keep_per_group=1 that is 56 proposals x 8 attempts - # per round measured and then not trained on. - if keep_per_group < 0: - raise ValueError(f'keep_per_group must be >= 0, got {keep_per_group}') - self.keep_per_group = keep_per_group - self.reject_sink = reject_sink - self.propose_sink = propose_sink - self.solver_sink = solver_sink - if self.seeds: - prompts.require('from_seed') - if self.keywords is not None: - prompts.require('from_seed_keywords') - self.stats: Dict[str, int] = { - 'explore_done': 0, 'check_parse_fail': 0, 'check_run_fail': 0, - 'empty_workspace': 0, 'solver_truncated': 0, - # The workspace listing could not be read, as opposed to being empty. - # Kept apart from ``empty_workspace`` because it says nothing about - # what the model did. - 'snapshot_unavailable': 0, - 'problem_parse_fail': 0, 'too_long': 0, 'parsed': 0, - # How often a check that failed was handed back for a rewrite, and - # how often the rewrite passed. The two together say whether the - # retry earns its extra sampling call. - 'check_retry': 0, 'check_retry_pass': 0, - # The episode ended before the appended stages could run or finish: - # it used up ``max_turns``, hit the trajectory token cap, or left no - # room for the follow-up message. Distinct from every other reason - # here, which is the model producing something unusable. - 'episode_cut_short': 0, - # Arm B only. ``setup_capture_fail``: the episode's input files could - # not be read back, so the task was dropped. ``setup_replay_fail``: a - # solver attempt was skipped because putting those files back failed, - # which would otherwise have scored as the task being too hard. - 'setup_capture_fail': 0, 'setup_replay_fail': 0, - # followup_api mode only: a check or statement API call failed. The - # conversation is then unusable and the proposal is rejected. - 'followup_api_error': 0, - # Novelty judging. ``novelty_error``: the batch's call raised, so every - # proposal in it scored None and none lost reward for it. - # ``novelty_length_mismatch``: the judge returned a different number of - # scores than proposals sent, which would pair scores with the wrong - # tasks, so all are dropped. ``novelty_unjudged``: proposals the judge - # left without a verdict. ``in_band_not_selected``: proposals inside the - # difficulty band whose group already contributed its keep_per_group - # task. - # - # These four have to be listed here: _bump does self.stats[key] += n on - # a fixed dict, so an unregistered key raises KeyError and takes the - # whole collection down. That is what killed loop2/iter1 -- seven - # proposals were measured and scored, then the counter line at the end - # of _score_novelty crashed and all seven trajectories were lost. - 'novelty_error': 0, 'novelty_length_mismatch': 0, - 'novelty_unjudged': 0, 'in_band_not_selected': 0, - # Keyword groups given up on because the judge never scored one of - # their proposals. Nothing from them is used. - 'novelty_group_dropped': 0, - } - - # ------------------------------------------------------------- proposing - - def propose(self, count: int) -> List[Trajectory]: - """Build ``count`` prompt trajectories for round 1. + if problem_max_chars <= 0: + raise ValueError(f'problem_max_chars must be positive, got {problem_max_chars}') + if not check_language.strip(): + raise ValueError('check_language must not be empty') + if not 0 <= pass_rate_target <= 1: + raise ValueError(f'pass_rate_target must be in [0, 1], got {pass_rate_target}') + if pass_rate_width <= 0: + raise ValueError(f'pass_rate_width must be positive, got {pass_rate_width}') + if use_api and rollout_kwargs.get('response_callback') is not None: + raise ValueError('use_api=True cannot be combined with response_callback') + backend_is_api = isinstance(backend, (API, APISampler)) + if use_api and api is None and not backend_is_api: + raise ValueError('use_api=True requires api= when backend is a sampler') + self.keyword_generator = keyword_generator + self.trajectory_seed = list(trajectory_seed or ()) + self.summarizer = summarizer + self._system = self._system if system_prompt is None else system_prompt + self._from_scratch = self._from_scratch if from_scratch_prompt is None else from_scratch_prompt + self._from_keywords = self._from_keywords if from_keywords_prompt is None else from_keywords_prompt + self._from_seed = self._from_seed if from_seed_prompt is None else from_seed_prompt + self._from_seed_keywords = (self._from_seed_keywords if from_seed_keywords_prompt is None else + from_seed_keywords_prompt) + self._check_followup = self._check_followup if check_followup_prompt is None else check_followup_prompt + self._check_retry_followup = (self._check_retry_followup if check_retry_followup_prompt is None else + check_retry_followup_prompt) + self._check_parse_error = (self._check_parse_error if check_parse_error_prompt is None else + check_parse_error_prompt) + self._problem_followup = (self._problem_followup if problem_followup_prompt is None else + problem_followup_prompt) + self._check_retries = check_retries + self._problem_max_chars = problem_max_chars + self._check_language = check_language.strip().lower() + self._parse_check_fn = parse_check_fn + self._pass_rate_target = pass_rate_target + self._pass_rate_width = pass_rate_width + self.checker = checker + self.followup_params = followup_params + self.rng = random.Random() + self.use_api = use_api + self.save_failed_rollouts = save_failed_rollouts + self._recorder = RolloutRecorder(save_dir) if save_dir else None + self._round_proposals: List[_ProposalResult] = [] + self._backend = backend + self._rollout_kwargs = dict(rollout_kwargs) + if api is not None: + self._rollout_kwargs['api'] = api + if use_api: + self._rollout_kwargs['response_callback'] = _api_followup_response + self._rollout: Optional[MultiTurnRollout] = None + self._tool_schemas = self.env().tools() or None + + def _rollout_instance(self) -> MultiTurnRollout: + if self._rollout is None: + self._rollout = MultiTurnRollout(self._backend, **self._rollout_kwargs) + return self._rollout - Each carries a direction + keywords + optional seed. The explorer will - run these multi-turn in the sandbox. - - ``proposals_per_group`` of them share one keyword draw, one seed choice - and one identical prompt, and are tagged with the same ``group_id``. - That is what makes a GRPO group on the proposing side: the advantage of - a proposal is its reward minus the mean over the others answering the - same prompt, so the members have to differ only by sampling noise. At 1 - -- which is what this used to be, every proposal its own keyword draw -- - every group has one member, the mean equals the reward, and every - advantage is zero. - """ - proposals: List[Trajectory] = [] - directions: List[str] = [] - metas: List[Tuple[List[Tuple[str, str]], bool, str, int]] = [] - per_group = max(1, self.proposals_per_group) - while len(metas) < count: - picks = self.keywords.draw() if self.keywords else [] - body = KeywordBank.block(picks) - use_seed = bool(self.seeds) and self.rng.random() < self.seed_mix_prob - seed = self.rng.choice(self.seeds) if use_seed else None - if use_seed and picks: - user = self.prompts.from_seed_keywords.format( - seed=seed['query'], keywords=body) - elif use_seed: - user = self.prompts.from_seed.format(seed=seed['query']) - elif picks: - user = self.prompts.from_keywords.format(keywords=body) - else: - user = self.prompts.from_scratch - # The whole group gets the same prompt, so a short final group is a - # group whose advantage is computed over fewer samples -- noisier, - # but not wrong. Truncating to a multiple of per_group instead would - # silently return fewer proposals than asked for. - # - # The counter is per-run, not per-call: ``propose`` runs once per - # round, and restarting at 0 each round would give two unrelated - # groups the same id in the dump. - gid = self._next_group_id - self._next_group_id += 1 - for _ in range(min(per_group, count - len(metas))): - directions.append(user) - metas.append((picks, use_seed, body, gid)) - - for user, (picks, use_seed, body, gid) in zip(directions, metas): - proposal: Trajectory = { - 'messages': [{'role': 'system', 'content': self.prompts.system}, - {'role': 'user', 'content': user}], - } - if self.tool_schemas: - proposal['tools'] = self.tool_schemas - proposals.append(attach_user_data( - proposal, keywords=picks, seeded=use_seed, keyword_block=body, - group_id=gid)) - return proposals - - # ------------------------------------------------------------- building - - def build(self, explored: List[Trajectory]) -> List[Optional[Trajectory]]: - """Satisfy the abstract method; not usable outside ``_round``. - - Building happens inside the episode now: :meth:`_followup` runs while the - model is still generating and needs the sandbox to hold that episode's - workspace state, which is only guaranteed inside the serial ``_round`` - loop. - """ - raise RuntimeError( - f'{type(self).__name__}.build() must not be called directly; ' - f'the serial _round() loop drives one episode at a time instead.') - - def _reject_for_empty_snapshot(self, state: Dict[str, Any], detail: str) -> None: - """File an episode whose workspace listing came back empty. + def _tool_manager(self, slot: int) -> Optional[Any]: + env = self.env(slot) + return env.tool_manager() if env.tools() else None - Empty means one of two unrelated things -- the episode built nothing, or - the listing could not be read -- and only the first says anything about - the model. ``detail`` is the second half of what - :meth:`~twinkle_agentic.envs.base.Env.snapshot` returns, and it is what - tells them apart: without it every case is filed as ``empty_workspace``, - which is what used to happen for all of them -- 63 of run_clean6's 71 - ``empty_workspace`` rejections were the 410 "sandbox is not proxyable" - error, so that reject class was 89% broken environment. - """ - if detail: - self._bump('snapshot_unavailable') - state['reject'] = ('snapshot_unavailable', detail) + def _summary(self, trajectory: Trajectory) -> str: + turns: List[str] = [] + for message in trajectory.get('messages') or []: + if not isinstance(message, dict): + continue + role = message.get('role') or '' + if role == 'system': + continue + parts = [msg_content_text(message).strip()] + for call in normalize_tool_calls(message) or (): + fn = call.get('function') or {} + if isinstance(fn, dict) and fn.get('name'): + parts.append(f"calls {fn['name']}({fn.get('arguments') or ''})") + body = '\n'.join(part for part in parts if part) + if body: + turns.append(f'{role}: {body}') + text = '\n'.join(turns) + if not text: + return '' + return self.summarizer(text) if self.summarizer is not None else text + + def _build_challenge_prompt(self) -> Optional[Trajectory]: + keywords: List[str] = [] + if self.keyword_generator is not None: + groups = self.keyword_generator.get_keywords(1) + if not groups: + return None + keywords = groups[0] + seed = '' + if self.trajectory_seed: + seed = self._summary(self.rng.choice(self.trajectory_seed)) + block = ', '.join(keywords) + if seed and keywords: + user = self._from_seed_keywords.format(seed=seed, keywords=block) + elif seed: + user = self._from_seed.format(seed=seed) + elif keywords: + user = self._from_keywords.format(keywords=block) else: - self._bump('empty_workspace') - state['reject'] = ('empty_workspace', '') + user = self._from_scratch + prompt: Trajectory = { + 'messages': [ + { + 'role': 'system', + 'content': self._system + }, + { + 'role': 'user', + 'content': user + }, + ], + } + if self._tool_schemas: + prompt['tools'] = self._tool_schemas + return attach_user_data(prompt, keywords=keywords, seeded=bool(seed)) + + def _explore(self, prompt: Trajectory) -> List[Trajectory]: + group_id = uuid.uuid4().hex + proposals: List[_ProposalResult] = [] + remaining = self.num_challenger_rollouts + while remaining > 0: + wave = min(self.n_slots, remaining) + proposals.extend(_parallel(lambda slot: self._run_episode(prompt, slot), wave)) + remaining -= wave + for proposal in proposals: + proposal.group_id = group_id + self._round_proposals = proposals + return [proposal.task for proposal in proposals if proposal.task is not None] + + def _run_episode(self, prompt: Trajectory, slot: int) -> _ProposalResult: + self.env(slot).clear() + state: Dict[str, Any] = {'slot': slot} + kwargs: Dict[str, Any] = { + 'followup_fn': lambda trajectory, n_before: self._followup(state, trajectory, n_before), + } + manager = self._tool_manager(slot) + if manager is not None: + kwargs['tool_manager'] = manager + explored = self._rollout_instance()([prompt], **kwargs) + if not explored: + self._reject(state, 'rollout_no_output') + return _ProposalResult(dict(prompt), reason='rollout_no_output') + trajectory = explored[0] + task = self._build_query(state, trajectory) + reason, detail = state.get('reject', ('', '')) + return _ProposalResult(trajectory, task=task, reason=reason, detail=detail) def _followup(self, state: Dict[str, Any], trajectory: Trajectory, n_before: int) -> Optional[Tuple[str, Optional[SamplingParams]]]: - """What to say next when the model stops calling tools; ``None`` to stop. - - The rollout calls this once per stage, handing over the episode as it - stands. ``state`` is this episode's scratchpad, read afterwards by - :meth:`_finish_episode`: the workspace listing and the check script are - produced here, and anything that goes wrong before a statement exists is - left in ``state['reject']``. + if state.get('checked'): + return None + reply = None if n_before == 0 else assistant_text(trajectory) + followup = self._build_test_case(state, reply) + if followup is None: + return None + return followup, self.followup_params - The sandbox work has to happen at this moment and nowhere else -- the - workspace holds this episode's end state right now, and the next - episode's reset wipes it. ``state['slot']`` says which sandbox that is; - with concurrent episodes several of these run at once, each against its - own. - """ - slot = state.get('slot', 0) - if n_before == 0: - snapshot, snapshot_error = self.env(slot).snapshot() + def _build_test_case(self, state: Dict[str, Any], reply: Optional[str]) -> Optional[str]: + slot = state['slot'] + if reply is None: + snapshot, error = self.env(slot).snapshot() state['snapshot'] = snapshot - # An episode that left nothing behind has no end state to write checks - # about, and asking for them anyway is worse than useless: the only - # true thing to assert is that the directory is empty, which every - # solver passes by doing nothing. Five of run5's ten verified tasks - # were that task. Reject here instead. if not snapshot.strip(): - self._reject_for_empty_snapshot(state, snapshot_error) - return None - return (self.prompts.check_followup.format(final_state=snapshot), - self.check_params) - - # Every follow-up from here until a check passes is a check-script reply: - # the first one, plus up to ``check_retries`` rewrites. - if not state.get('checked'): - attempt = state.get('check_attempts', 0) + 1 - state['check_attempts'] = attempt - reply = assistant_text(trajectory) - script = self.parse_check_fn(reply) - if script is None: - # Same one-rewrite budget a run failure gets: hand the parse - # failure back and let it regenerate, rather than dropping a task - # whose only fault was packaging. Shares the check_attempts - # count, so parse and run failures together get check_retries - # extra tries, not one each. - if attempt <= self.check_retries: - self._bump('check_retry') - err = self.prompts.check_parse_error or DEFAULT_CHECK_PARSE_ERROR - return (self.prompts.check_retry_followup.format( - error=err, final_state=state.get('snapshot') or ''), - self.check_params) - self._bump('check_parse_fail') - # The whole reply, not a tail: this stage fails either because the - # model declared the state untestable (it says so) or because it - # ran out of tokens while thinking, and a record that cannot tell - # them apart sends the next reader back to re-run the batch. - state['reject'] = ('check_parse_fail', reply) + state['reject'] = ('snapshot_unavailable' if error else 'empty_workspace', error) return None - state['script'] = script - brittle = self.brittle_check_fn(script) if self.brittle_check_fn else None - if brittle is not None: - # Same bookkeeping as a check that ran and failed: the script is - # rejected before it can pass on the author's own state, because - # passing there is exactly what hides the defect. - exit_code, output = 1, brittle - else: - exit_code, output = self.env(slot).run_script(script) - if exit_code == 0: - state['checked'] = True - if attempt > 1: - self._bump('check_retry_pass') - # Capture the inputs now, at the one moment the workspace holds - # exactly the state this check just passed on. A capture after the - # statement stage would be the same bytes only by luck. - if self.setup_script_fn is not None: - setup = self.setup_script_fn(slot=slot) - if not setup: - self._bump('setup_capture_fail') - state['reject'] = ( - 'setup_capture_fail', - 'no input files to hand the solver, or their bytes ' - 'could not be read back') - return None - state['setup_script'] = setup - return (self.prompts.problem_followup, self.problem_params) - # Snapshot again, after the failure. A check that asserts only - # paths from the snapshot it was shown and still fails leaves two - # very different bugs indistinguishable -- the model asserted - # something untrue, or the workspace changed under it -- and the - # difference is visible only in the state at the moment the check - # ran. It is also what the rewrite gets to read. - after = self.env(slot).snapshot()[0] - state.setdefault('attempts', []).append( - f'--- attempt {attempt}: exit {exit_code} ---\n{output}\n' - f'--- check script ---\n{script}') - if attempt <= self.check_retries: - self._bump('check_retry') - return (self.prompts.check_retry_followup.format( - error=output, final_state=after or state.get('snapshot') or ''), - self.check_params) - self._bump('check_run_fail') - state['reject'] = ( - 'check_run_fail', - '\n'.join(state['attempts']) - + f"\n--- state before check ---\n{state.get('snapshot') or ''}\n" - + f'--- state after check ---\n{after}') + return self._check_followup.format(final_state=snapshot, language=self._check_language) + + attempt = state.get('check_attempts', 0) + 1 + state['check_attempts'] = attempt + script = (self._parse_check_fn(reply) if self._parse_check_fn is not None else + parse_fenced_code(reply, language_tags=None)) + if script is None: + if attempt <= self._check_retries: + return self._check_retry_followup.format( + error=self._check_parse_error.format(language=self._check_language), + final_state=state.get('snapshot', ''), + language=self._check_language, + ) + state['reject'] = ('check_parse_fail', reply) return None - + state['script'] = script + exit_code, output = self.env(slot).run_script(script, interpreter=self._check_language) + if exit_code == 0: + state['checked'] = True + return self._problem_followup + after = self.env(slot).snapshot()[0] + state.setdefault('attempts', []).append(f'--- attempt {attempt}: exit {exit_code} ---\n{output}\n' + f'--- check script ---\n{script}') + if attempt <= self._check_retries: + return self._check_retry_followup.format( + error=output, + final_state=after or state.get('snapshot', ''), + language=self._check_language, + ) + state['reject'] = ('check_run_fail', '\n'.join(state['attempts'])) return None - def _run_followup_api(self, state: Dict[str, Any], explored: Trajectory) -> None: - """Generate the check script and problem statement over ``followup_api``. - - The API-only counterpart of :meth:`_followup`: the same stages, the same - sandbox work (snapshot, run the check, capture inputs) and the same retry - budget, but driven imperatively here instead of turn-by-turn by the - rollout, and answered by the API rather than the local model. Results land - in ``state`` for :meth:`_finish_episode`: - - * ``state['script']`` / ``state['checked']`` -- the check that passed, - * ``state['setup_script']`` -- captured inputs (Arm B), - * ``state['statement']`` -- the problem-statement text, - * ``state['reject']`` -- ``(reason, detail)`` when a stage fails. - - Nothing here touches ``explored``'s ``input_ids`` / ``labels`` / - ``logprobs``: the messages the API sees are a private copy, so the - trainable trajectory stays exactly the exploration turns the local sampler - produced. - """ - slot = state.get('slot', 0) - messages: List[Dict[str, Any]] = [dict(m) for m in explored.get('messages') or []] - - snapshot, snapshot_error = self.env(slot).snapshot() - state['snapshot'] = snapshot - # An episode that left nothing behind has no end state to write checks - # about; rejecting here mirrors the n_before==0 branch of _followup. - if not snapshot.strip(): - self._reject_for_empty_snapshot(state, snapshot_error) - return - - # Check-script stage: the first ask plus up to ``check_retries`` rewrites, - # sharing one attempt counter across parse and run failures exactly as the - # single-model path does. - user_text = self.prompts.check_followup.format(final_state=snapshot) - attempt = 0 - while True: - attempt += 1 - state['check_attempts'] = attempt - reply = self.followup_model.reply(messages, user_text, self.check_params) - if reply is None: - self._bump('followup_api_error') - state['reject'] = ('followup_api_error', 'check-script API call failed') - return - script = self.parse_check_fn(reply) - if script is None: - if attempt <= self.check_retries: - self._bump('check_retry') - err = self.prompts.check_parse_error or DEFAULT_CHECK_PARSE_ERROR - user_text = self.prompts.check_retry_followup.format( - error=err, final_state=snapshot) - continue - self._bump('check_parse_fail') - state['reject'] = ('check_parse_fail', reply) - return - state['script'] = script - brittle = self.brittle_check_fn(script) if self.brittle_check_fn else None - if brittle is not None: - # Rejected before it can pass on the author's own state, since - # passing there is exactly what hides the defect. - exit_code, output = 1, brittle - else: - exit_code, output = self.env(slot).run_script(script) - if exit_code == 0: - state['checked'] = True - if attempt > 1: - self._bump('check_retry_pass') - # Capture inputs now, while the workspace still holds the state - # this check just passed on. - if self.setup_script_fn is not None: - setup = self.setup_script_fn(slot=slot) - if not setup: - self._bump('setup_capture_fail') - state['reject'] = ( - 'setup_capture_fail', - 'no input files to hand the solver, or their bytes ' - 'could not be read back') - return - state['setup_script'] = setup - break - after = self.env(slot).snapshot()[0] - state.setdefault('attempts', []).append( - f'--- attempt {attempt}: exit {exit_code} ---\n{output}\n' - f'--- check script ---\n{script}') - if attempt <= self.check_retries: - self._bump('check_retry') - user_text = self.prompts.check_retry_followup.format( - error=output, final_state=after or snapshot) - continue - self._bump('check_run_fail') - state['reject'] = ( - 'check_run_fail', - '\n'.join(state['attempts']) - + f"\n--- state before check ---\n{state.get('snapshot') or ''}\n" - + f'--- state after check ---\n{after}') - return - - # Problem-statement stage: one API reply, kept as the task's statement. - reply = self.followup_model.reply(messages, self.prompts.problem_followup, - self.problem_params) - if reply is None: - self._bump('followup_api_error') - state['reject'] = ('followup_api_error', 'problem-statement API call failed') - return - state['statement'] = reply - - def _finish_episode(self, state: Dict[str, Any], - explored: Trajectory) -> Optional[Trajectory]: - """Turn a finished episode into a task, or record why it is not one. - - Everything the model wrote is in ``explored``: the tool-using turns, the - check script, and the problem statement as the last assistant message. - ``state`` carries what only the sandbox could say -- the end state, and - whether the check passed on it. - """ - keywords = user_data_get(explored.get('user_data'), 'keywords', []) - seeded = user_data_get(explored.get('user_data'), 'seeded', False) - group_id = user_data_get(explored.get('user_data'), 'group_id', None) - # The episode as one record: a single conversation, so a single set of - # token ids and logprobs. Handed to propose_sink with whatever verdict the - # proposal ends up with, so a rejected attempt is recorded as fully as a - # kept one. - rounds = [_propose_round('episode', explored)] - - def reject(reason: str, detail: str = '') -> None: - self._reject_record(explored, reason, detail=detail) - self._emit_propose(rounds, reason, keywords=keywords, seeded=seeded, - group_id=group_id) - + def _build_query(self, state: Dict[str, Any], explored: Trajectory) -> Optional[Trajectory]: if state.get('reject'): - reason, detail = state['reject'] - reject(reason, detail) - return None - + return self._reject(state, *state['reject']) if not state.get('checked'): - # The stages never ran, or the check-writing one never got a reply: - # the episode used up its turns, hit the trajectory token cap, or left - # no room to append the next message. The model produced nothing - # wrong here, so this is not one of the other reasons. Keyed on the - # check having *passed* rather than on a script existing: a rewrite - # that never came back leaves the failed script in ``state``, and - # building a task on it would ship a check nobody can pass. - self._bump('episode_cut_short') - reject('episode_cut_short', - detail=f"stop_reason={explored.get('stop_reason')} " - f"truncated={bool(explored.get('truncated'))} " - f"turns={explored.get('turns')} " - f"followups={explored.get('followups')}") - return None - - script = state['script'] - # In followup_api mode the statement was written by the API and is not in - # ``explored`` (whose last assistant turn is the final exploration reply); - # it lives in ``state``. The single-model path keeps it as the last - # assistant message of the episode. - if self.followup_model is not None: - statement = parse_problem_statement(state.get('statement') or '') - else: - statement = parse_problem_statement(assistant_text(explored)) + return self._reject( + state, + 'episode_cut_short', + f"stop_reason={explored.get('stop_reason')} " + f"truncated={bool(explored.get('truncated'))} " + f"turns={explored.get('turns')}", + ) + statement = parse_problem_statement(assistant_text(explored)) if statement is None: - self._bump('problem_parse_fail') - reject('problem_parse_fail') - return None - if len(statement) > self.problem_max_chars: - self._bump('too_long') - reject('too_long') - return None - - self._bump('parsed') - task: Trajectory = { - 'messages': [{'role': 'user', 'content': statement}], - } - # group_id travels with the task, not just with the reject path above: the - # difficulty stage emits the surviving proposals from the task, so a task - # that forgets its group reaches the dump ungrouped and the proposing side - # has no advantage to compute. Leaving it off here made every kept and - # outside_band proposal group_id=None, and only the episodes that failed - # early -- which emit straight off ``explored`` -- kept theirs. - task = attach_user_data(task, check_script=script, keywords=keywords, seeded=seeded, - group_id=group_id, - setup_script=state.get('setup_script', '')) - # Carried, not emitted: the verdict this proposal earns depends on the - # difficulty measurement, which has not run yet. A plain top-level key - # rather than user_data, which json-encodes every value on each update. - task['propose_rounds'] = rounds + return self._reject(state, 'problem_parse_fail') + if len(statement) > self._problem_max_chars: + return self._reject(state, 'too_long', f'{len(statement)} chars') + task: Trajectory = attach_user_data( + {'messages': [{ + 'role': 'user', + 'content': statement + }]}, + check_script=state['script'], + keywords=user_data_get(explored.get('user_data'), 'keywords', []), + seeded=user_data_get(explored.get('user_data'), 'seeded', False), + ) + if self.checker is not None and not self.checker(task): + return self._reject(state, 'rejected_by_checker') return task - def _reject_record(self, traj: Trajectory, reason: str, detail: str = '') -> None: - """Record a rejected proposal, with enough of the episode to tell why. - - The reason alone is not diagnosable. Nine ``empty_workspace`` rejections in - one run all looked like the model refusing to act; the messages showed a - single assistant turn each, and the question of whether it had run out of - tokens or simply emitted no call could not be answered from the record -- - the fields that answered it were on the trajectory and were dropped. So - how the episode ended travels with the reason. - """ - if self.reject_sink is None: - return - messages = traj.get('messages') or [] - payload: Dict[str, Any] = {'reason': reason} - if detail: - payload['detail'] = detail - payload['stop_reason'] = traj.get('stop_reason') - payload['truncated'] = bool(traj.get('truncated')) - payload['turns'] = traj.get('turns') - payload['n_assistant'] = sum(1 for m in messages - if isinstance(m, dict) and m.get('role') == 'assistant') - payload['n_tool_calls'] = sum(len(m.get('tool_calls') or []) for m in messages - if isinstance(m, dict)) - payload['last_assistant'] = assistant_text(traj) - with self._sink_lock: - self.reject_sink(payload) - - def _emit_propose(self, rounds: Optional[List[Dict[str, Any]]], outcome: str, *, - keywords: Any = (), seeded: bool = False, - n_pass: Optional[int] = None, - group_id: Optional[int] = None, - novelty: Optional[float] = None, - selected: Optional[bool] = None, - novelty_dropped: bool = False) -> None: - """Hand one proposal attempt's rounds to ``propose_sink``. - - ``pass_rate`` is the raw fraction of solver attempts that succeeded. - ``challenger_reward`` is that fraction scored against a 50% target by - :meth:`challenger_reward`, which is the number the proposing side trains - on; both are written so a run can be re-scored under a different target - without re-solving anything. - - ``novelty`` is written next to it for the same reason: the reward already - has it multiplied in, and a run cannot be re-scored at a different floor -- - or with novelty taken back out -- from the product alone. - - A proposal with no ``n_pass`` never reached difficulty measurement -- it - was rejected before that -- and scores 0, the same as one nobody or - everybody solved. - """ - if self.propose_sink is None or not rounds: - return - rollouts = self.solver_rollouts or None - payload = { - 'outcome': outcome, - 'group_id': group_id, - 'n_pass': n_pass, - 'n_rollouts': rollouts, - 'pass_rate': (n_pass / rollouts) if (n_pass is not None and rollouts) else None, - 'novelty': novelty, - 'novelty_factor': self.novelty_factor(novelty), - 'challenger_reward': self.challenger_reward(n_pass, novelty=novelty), - # Whether this proposal's task went on to the solver side. Not the same as - # ``outcome``: with keep_per_group set, a proposal can be in the difficulty - # band and still not be the one its group contributed. Its own reward is - # unaffected either way. - 'selected': selected, - # True when the novelty judge never returned a score for at least one - # proposal in this group, after NOVELTY_TRIES attempts. The record is - # written either way -- the episode really happened and the file is the - # audit trail -- but training skips every group carrying this flag. - 'novelty_dropped': bool(novelty_dropped), - 'keywords': list(keywords or ()), - 'seeded': bool(seeded), - 'rounds': rounds, - } - with self._sink_lock: - self.propose_sink(payload) - - # Where the pass-rate reward peaks, and how wide the peak is. 0.2 is Ornith-1.5's - # target (ornith.ai/ornith_1_5.html), which trains its proposer on - # ``exp(-(p-p*)^2 / 2s^2)`` rather than on a peak at one half. - PASS_RATE_TARGET = 0.2 - PASS_RATE_WIDTH = 0.3 - - # How many times the novelty judge is asked before a proposal is given up on. - # Only the proposals still missing a score are re-sent. Measured need for this: - # loop3/iter1 had 1 of 61 measured proposals come back without a verdict, in 1 - # of its 10 keyword groups, and giving up on that group costs the 56 sandbox - # attempts already spent on its 7 proposals. - NOVELTY_TRIES = 3 - - def novelty_factor(self, novelty: Optional[float]) -> float: - """What a proposal's difficulty score gets multiplied by for its novelty. - - ``floor + (1 - floor) * N``, so N=1 leaves the reward alone and N=0 leaves - ``novelty_floor`` of it. See ``novelty_floor`` in ``__init__`` for why there - is a floor at all. - - ``None`` returns 1.0, not the floor: it means nobody judged this proposal -- - no bank, no judge, or the judge's API failed -- and charging a proposal for a - measurement that did not happen would make the reward depend on API uptime. - """ - if novelty is None: - return 1.0 - n = min(1.0, max(0.0, float(novelty))) - return self.novelty_floor + (1.0 - self.novelty_floor) * n - - def challenger_reward(self, n_pass: Optional[int], - novelty: Optional[float] = None) -> float: - """Score a proposal by how close the solver came to a target pass rate. - - ``exp(-(p - p*)^2 / 2s^2)`` for ``p = n_pass / solver_rollouts``, peaked at - ``p* = 0.2`` with width ``s = 0.3``. - - This replaced ``1 - 2*|p - 1/2|``, which is what R-Zero (arXiv 2508.05004) - uses, for two reasons measured on run_clean9's 87 in-band proposals: - - It was not injective. With 8 rollouts the seven in-band values of ``n_pass`` - mapped onto four rewards -- 1 and 7 both scored 0.25, 2 and 6 both 0.50 -- - so a proposal one solver out of eight could do and one seven out of eight - could do were worth the same. The whole distinction between too hard and too - easy was erased. The gaussian separates all seven. - - Its signal was smaller than its noise. ``n_pass`` is a binomial draw around - the proposal's real difficulty, and propagating that draw through each shape - gives a noise SD to compare the spread of rewards against: 0.280 signal over - 0.246 noise for the old shape, against 0.347 over 0.177 here. A ratio of 1.14 - means over half of what a GRPO group ranks on is which way eight coin flips - landed. + def _reject(self, state: Dict[str, Any], reason: str, detail: str = '') -> Optional[Trajectory]: + state['reject'] = (reason, detail) + logger.info(f'[{type(self).__name__}] rejected: {reason}' + f"{f' -- {detail[:400]}' if detail else ''}") + return None - A peak below one half is also the more useful target. A group's update size - goes with reward variance, which for a pass/fail solver peaks at p=0.5 -- the - argument for the old shape -- but a proposal only teaches the solver - something when the solver mostly cannot do it yet. + def _solver_prompt(self, task: Trajectory) -> Trajectory: + prompt: Trajectory = {'messages': [dict(message) for message in task.get('messages') or []]} + if self._tool_schemas: + prompt['tools'] = self._tool_schemas + return prompt - ``None`` means the proposal never got as far as being solved, and 0 means no - attempt passed. Both score 0, and that floor is now load-bearing rather than - incidental: the gaussian evaluated at p=0 is 0.801, higher than the 0.607 it - gives a proposal half the attempts solve. Without the gate the best thing a - proposer could do is write tasks nobody can finish. + def _judge(self, task: Trajectory, slot: int) -> bool: + script = user_data_get(task.get('user_data'), 'check_script', '') + if not script: + return False + return self.env(slot).run_script(script, interpreter=self._check_language)[0] == 0 - ``novelty`` multiplies the result through :meth:`novelty_factor`, which is - Ornith-1.5's ``R = V x D x N`` with a floor under the N. Left at ``None`` -- - which is what happens with no task bank or no judge -- the returned number is - exactly what it was before novelty existed. - """ - rollouts = self.solver_rollouts or 0 - if n_pass is None or not rollouts or n_pass <= 0: + def challenger_reward(self, n_pass: Optional[int]) -> float: + """Reward tasks near the target solver pass rate; unmeasured failures score zero.""" + if n_pass is None or not self.num_solver_rollouts or n_pass <= 0: return 0.0 - gap = n_pass / rollouts - self.PASS_RATE_TARGET - difficulty = math.exp(-(gap * gap) / (2.0 * self.PASS_RATE_WIDTH ** 2)) - return difficulty * self.novelty_factor(novelty) - - def _take_rounds(self, task: Trajectory) -> Optional[List[Dict[str, Any]]]: - """Detach a task's proposing rounds. Popped even with no sink attached: - token ids for a whole agentic episode are large, and a kept task is held - until the caller's batch is full. - """ - return task.pop('propose_rounds', None) - - # ------------------------------------------------------------ revised _round - - def _bump(self, key: str, n: int = 1) -> None: - """Thread-safe stats increment.""" - with self._sink_lock: - self.stats[key] += n - - def _tool_manager(self, slot: int) -> Optional[Any]: - """The dispatcher for ``slot``'s tools, or None when it advertises none. - - Built per use rather than held, so a slot rebuilt underneath -- evicted, - timed out -- is dispatched into as it is now: a manager captured at - construction would keep sending this episode's calls to a sandbox that is - gone. None means this environment offers no tools, which is the honest - answer for one that only runs scripts, and the rollout then leaves the - model with none rather than an empty tool list it would try to call. - """ - env = self.env(slot) - return env.tool_manager() if env.tools() else None - - def _run_episode(self, proposal: Trajectory, slot: int) -> Optional[Trajectory]: - """One episode top-to-bottom, using sandbox slot ``slot``.""" - self.env(slot).clear() - state: Dict[str, Any] = {'slot': slot} - tm = self._tool_manager(slot) - if self.followup_model is not None: - # Split path: explore on the local (trainable) model with NO - # followup_fn, so the rollout ends the moment the model stops calling - # tools and the returned trajectory carries only the exploration - # turns' input_ids/labels/logprobs. The check-script and - # problem-statement stages then run over the API against the same end - # state, appended to a throwaway copy of the messages -- never to the - # trainable trajectory. - kwargs: Dict[str, Any] = {} - if tm is not None: - kwargs['tool_manager'] = tm - result = self.explore([proposal], **kwargs) - if not result: - return None - explored = result[0] - self._bump('explore_done') - # A reply cut off at the token budget never finished its thought, so - # continuing the conversation over the API would build a check on a - # half-written turn. Leave ``state`` untouched and let - # ``_finish_episode`` record it as ``episode_cut_short``, matching the - # single-model path which does not run the stages after a length cut. - if explored.get('stop_reason') != 'length': - self._run_followup_api(state, explored) - return self._finish_episode(state, explored) - kwargs = {'followup_fn': partial(self._followup, state)} - if tm is not None: - kwargs['tool_manager'] = tm - result = self.explore([proposal], **kwargs) - if not result: - return None - explored = result[0] - self._bump('explore_done') - return self._finish_episode(state, explored) - - def _round(self, missing: int) -> Optional[List[Trajectory]]: - """One cycle: episodes in parallel across sandbox slots, then difficulty.""" - count = min(self._estimate(missing), self.max_proposals_per_round) - if self.max_proposals_total > 0: - left = self.max_proposals_total - self.n_proposed - if left <= 0: - # Budget spent. None is the 'source exhausted' answer the batching - # loop already knows how to stop on, so the run ends after this - # round's keepers are handed back rather than mid-episode. - logger.info(f'[{type(self).__name__}] proposal budget spent ' - f'({self.n_proposed}/{self.max_proposals_total}); stopping') - return None - count = min(count, left) - proposals = self.propose(count) - if not proposals: - return None - - usable: List[Trajectory] = [] - n_slots = self.n_slots - - if n_slots <= 1 or len(proposals) <= 1: - # Serial fallback (original path). - for proposal in proposals: - task = self._run_episode(proposal, slot=0) - if task is not None: - usable.append(task) - else: - # One worker per sandbox slot, each draining its own share serially. - # A slot is a single sandbox and cannot host two episodes at once, so - # the split is by slot, never round-robin into a shared pool where two - # tasks could land on the same slot concurrently. - buckets: List[List[Trajectory]] = [[] for _ in range(n_slots)] - for i, proposal in enumerate(proposals): - buckets[i % n_slots].append(proposal) + gap = n_pass / self.num_solver_rollouts - self._pass_rate_target + variance = 2.0 * self._pass_rate_width**2 + return math.exp(-(gap * gap) / variance) - def _drain(slot: int) -> List[Trajectory]: - out: List[Trajectory] = [] - for proposal in buckets[slot]: - task = self._run_episode(proposal, slot=slot) - if task is not None: - out.append(task) - return out - - with ThreadPoolExecutor(max_workers=n_slots) as pool: - futures = [pool.submit(_drain, s) for s in range(n_slots) if buckets[s]] - for fut in as_completed(futures): - usable.extend(fut.result()) - - kept = self._filter_difficulty(usable) if self.solver_rollouts else usable - if not self.solver_rollouts: - # No difficulty stage, so the verdict is final as soon as it is built. - for task in usable: - self._emit_propose(self._take_rounds(task), 'kept', - keywords=user_data_get(task.get('user_data'), 'keywords', []), - seeded=user_data_get(task.get('user_data'), 'seeded', False), - group_id=user_data_get(task.get('user_data'), 'group_id', None)) - self.n_proposed += len(proposals) - self.n_kept += len(kept) - band = (f', in difficulty band {len(kept)}' if self.solver_rollouts else '') - logger.info(f'[{type(self).__name__}] proposed {len(proposals)}, usable ' - f'{len(usable)}{band} (cumulative {self.n_kept}/{self.n_proposed})') - return kept - - # ------------------------------------------------------------ difficulty + def _record_proposals(self) -> None: + proposals, self._round_proposals = self._round_proposals, [] + if self._recorder is None: + return + for index, proposal in enumerate(proposals): + if proposal.task is None and not self.save_failed_rollouts: + continue + trajectory = dict(proposal.trajectory) + trajectory['rewards'] = proposal.reward + task_data = proposal.task.get('user_data') if proposal.task is not None else None + statement = '' + if proposal.task is not None: + statement = next((message.get('content', '') for message in proposal.task.get('messages') or [] + if isinstance(message, dict) and message.get('role') == 'user'), '') + self._recorder.write( + trajectory, + side='propose', + group_id=proposal.group_id, + proposal_index=index, + outcome=proposal.outcome or ('rejected' if proposal.reason else 'kept'), + reason=proposal.reason, + detail=proposal.detail, + reward=proposal.reward, + n_pass=proposal.n_pass, + n_rollouts=(self.num_solver_rollouts if proposal.n_pass is not None else None), + pass_rate=(proposal.n_pass / self.num_solver_rollouts + if proposal.n_pass is not None and self.num_solver_rollouts else None), + statement=statement, + check_script=user_data_get(task_data, 'check_script', ''), + keywords=user_data_get(proposal.trajectory.get('user_data'), 'keywords', []), + seeded=user_data_get(proposal.trajectory.get('user_data'), 'seeded', False), + ) def _filter_difficulty(self, tasks: List[Trajectory]) -> List[Trajectory]: - """Override: every solver attempt needs its own clean workspace. + successful = [proposal for proposal in self._round_proposals if proposal.task is not None] + if len(successful) != len(tasks): + raise RuntimeError('proposal/task alignment failed before difficulty filtering') + if not tasks or not self.num_solver_rollouts: + for proposal in successful: + proposal.outcome = 'kept' + self._record_proposals() + return tasks - Attempts are run in waves of ``len(envs)``, attempt k of a wave in slot k. - Within a wave all attempts go out in one explorer call, - so the sampler generates them as one batch instead of leaving the GPUs - waiting on a single sequence, and the wave's clears, input replays and - checks all run at the same time too -- they are sandbox round-trips, not - compute. - - The slot is what keeps this honest: a wave's attempts each clear, act in - and get checked against their own sandbox. Sharing one would let attempt A - pass on files attempt B wrote, and ``n_pass`` would stop being a - difficulty measurement. - - An attempt cut off at the generation budget is counted in - ``stats['solver_truncated']`` but still counts as a failure, because - deciding otherwise decides which tasks are kept. Watch that number: when - it is a large share of ``solver_rollouts`` times the task count, ``n_pass`` - is reporting the token budget rather than the difficulty. It was 15 of 50 - on one run, and one task lost all four attempts that way and was discarded - as too hard without a solver ever touching the workspace. Raising - ``solver_params.max_tokens`` took it to 0 of 20. - """ - if not tasks: - return [] passes = [0] * len(tasks) - n_slots = max(1, self.n_slots) - # Which task each attempt belongs to, flattened, so a wave is a fixed - # number of sandboxes no matter how attempts distribute over tasks. - plan = [i for i in range(len(tasks)) for _ in range(self.solver_rollouts)] - - for start in range(0, len(plan), n_slots): - wave = plan[start:start + n_slots] - slots = list(range(len(wave))) - setups = [user_data_get(tasks[i].get('user_data'), 'setup_script', '') - for i in wave] - - def _prepare(k: int) -> bool: - """Clear slot k, then put back the inputs this task hands out.""" - self.env(k).clear() - if not setups[k]: - return True - exit_code, output = self.env(k).run_script(setups[k]) - if exit_code != 0: - # Measuring this attempt against a workspace missing its - # inputs would score the task as harder than it is, so the - # attempt is skipped and counted rather than run. - logger.warning(f'[{type(self).__name__}] input setup failed in ' - f'slot {k} (exit {exit_code}): {output[-200:]}') - return False - return True - - ready = map_parallel(_prepare, slots) - live = [k for k in slots if ready[k]] - self._bump('setup_replay_fail', len(slots) - len(live)) - if not live: - continue - prompts = [dict(self.solver_prompt(tasks[wave[k]])) for k in live] + plan = [i for i in range(len(tasks)) for _ in range(self.num_solver_rollouts)] + rollout = self._rollout_instance() + for start in range(0, len(plan), self.n_slots): + wave = plan[start:start + self.n_slots] + _parallel(lambda slot: self.env(slot).clear(), len(wave)) + prompts = [self._solver_prompt(tasks[i]) for i in wave] kwargs: Dict[str, Any] = {} - managers = [self._tool_manager(k) for k in live] - if any(tm is not None for tm in managers): + managers = [self._tool_manager(slot) for slot in range(len(wave))] + if any(manager is not None for manager in managers): kwargs['tool_manager'] = managers - attempts = self._solver_explore(prompts, sampling_params=self.solver_params, - **kwargs) + attempts = rollout(prompts, **kwargs) if len(attempts) != len(prompts): - # Counting a partial return would silently understate every - # affected task's pass count, i.e. report tasks as harder than - # they are. - raise RuntimeError(f'explorer returned {len(attempts)} attempts for ' - f'{len(prompts)} solver prompts; expected one per prompt.') - for attempt in attempts: - if attempt is not None and attempt.get('truncated'): - self._bump('solver_truncated') - verdicts = map_parallel( - lambda j: (attempts[j] is not None - and self.judge_attempt(tasks[wave[live[j]]], attempts[j], - slot=live[j])), - list(range(len(live)))) - for j, passed in enumerate(verdicts): + raise RuntimeError(f'rollout returned {len(attempts)} attempts for ' + f'{len(prompts)} prompts; expected one per prompt') + verdicts = _parallel(lambda slot: self._judge(tasks[wave[slot]], slot), len(wave)) + for slot, passed in enumerate(verdicts): if passed: - passes[wave[live[j]]] += 1 + passes[wave[slot]] += 1 + low, high = self.pass_band measured = [ - attach_user_data(task, n_pass=passes[i], n_rollouts=self.solver_rollouts) - for i, task in enumerate(tasks) + attach_user_data(task, n_pass=n_pass, n_rollouts=self.num_solver_rollouts) + for task, n_pass in zip(tasks, passes) ] - self.on_difficulty_measured(measured) - novelties = self._score_novelty(measured) - low, high = self.keep_pass_band - in_band = [low <= n <= high for n in passes] - # Which of the in-band tasks the solver side actually trains on. Decided - # before emitting so each proposal's record says whether its task was taken. - selected = self._select_per_group(measured, passes, in_band, novelties) - dropped = self._unscored_group_ids(measured, novelties) - if dropped: - self._bump('novelty_group_dropped', len(dropped)) - # Emit here, not in _round: this is where a proposal's verdict is - # decided, and both sides of the band are worth keeping -- a task nobody - # solved and one everybody solved are the two failure modes the - # proposer would need to learn to avoid. - for i, (task, n, kept_flag, nov) in enumerate(zip(measured, passes, in_band, - novelties)): - gid = user_data_get(task.get('user_data'), 'group_id', None) - self._emit_propose(self._take_rounds(task), - 'kept' if kept_flag else 'outside_band', - keywords=user_data_get(task.get('user_data'), 'keywords', []), - seeded=user_data_get(task.get('user_data'), 'seeded', False), - n_pass=n, - group_id=gid, - novelty=nov, - selected=selected[i], - novelty_dropped=gid in dropped) - return [t for t, take in zip(measured, selected) if take] - - def _unscored_group_ids(self, measured: List[Trajectory], - novelties: List[Optional[float]]) -> set: - """Keyword groups the novelty judge never finished answering for. - - A group lands here when at least one of its proposals still has no score - after ``NOVELTY_TRIES`` attempts. Nothing from such a group is used: no task - is taken from it (``_select_per_group``) and its proposals are marked - ``novelty_dropped`` so training skips the whole group. The collecting loop - then keeps going and a later keyword draw makes up the shortfall. - - Only meaningful when novelty is on. With it off every score is ``None`` by - design, which must not drop everything, so an off judge returns no groups. - """ - if self.task_bank is None or self.novelty_fn is None: - return set() - return {user_data_get(task.get('user_data'), 'group_id', None) - for task, nov in zip(measured, novelties) if nov is None} - - def _select_per_group(self, measured: List[Trajectory], passes: List[int], - in_band: List[bool], - novelties: List[Optional[float]]) -> List[bool]: - """Which in-band proposals become tasks: all of them, or the best few per group. - - With ``keep_per_group = k > 0``, each keyword group contributes at most its ``k`` - highest-reward in-band proposals -- reward being the same number the proposing - side trains on, ``challenger_reward``, so the task kept is the one whose pass - rate sat closest to the target and, when novelty is on, was not judged a repeat - of something already in the bank. - - A group with no in-band proposal contributes nothing and is not replaced here: - the collecting loop keeps proposing rounds until the run's target number of - tasks is reached, so a group that produced none is skipped and paid for by one - more group later. - - Proposals with no ``group_id`` (a run with ``proposals_per_group=1``) are each - their own group, so this is a no-op for them beyond the in-band filter. - """ - if self.keep_per_group <= 0: - return list(in_band) - ranked: Dict[Any, List[Tuple[float, int]]] = {} - unscored_groups = self._unscored_group_ids(measured, novelties) - for i, task in enumerate(measured): - if not in_band[i]: - continue - gid = user_data_get(task.get('user_data'), 'group_id', None) - if gid in unscored_groups: - # The judge never finished scoring this group, so there is no honest - # way to rank its members against each other. - continue - key = gid if gid is not None else f'_ungrouped_{i}' - ranked.setdefault(key, []).append( - (self.challenger_reward(passes[i], novelty=novelties[i]), i)) - selected = [False] * len(measured) - for key, entries in ranked.items(): - # Ties broken by the earlier proposal, so the choice does not depend on - # dict or sort instability. - entries.sort(key=lambda pair: (-pair[0], pair[1])) - for _, i in entries[:self.keep_per_group]: - selected[i] = True - dropped = sum(1 for i in range(len(measured)) if in_band[i] and not selected[i]) - if dropped: - self._bump('in_band_not_selected', dropped) - return selected - - def _score_novelty(self, measured: List[Trajectory]) -> List[Optional[float]]: - """One novelty score per measured proposal, ``None`` for every one if off. - - Scored for the whole batch in one call, and with the batch's own statements - as part of each proposal's reference set, because the comparison that matters - is against the siblings sharing a keyword draw: GRPO subtracts the group mean, - so a term that comes out the same for all eight members of a group cancels - exactly and the API calls bought nothing. Only same-group siblings go in -- - an unrelated proposal from the same round is not evidence of redundancy. - - Failures return ``None`` rather than 0.0 and never raise: a judge that is - down must not turn into every proposal being redundant, and must not lose a - round of sandbox work either. - - A proposal the judge skipped is asked about again, up to ``NOVELTY_TRIES`` - attempts in total, sending only the ones still missing. Whatever is still - unscored after that leaves its whole keyword group out of both the task - selection and the training data -- see ``_select_per_group`` and the - ``novelty_dropped`` field written by ``_emit_propose``. - """ - if self.task_bank is None or self.novelty_fn is None or not measured: - return [None] * len(measured) - statements, checks, groups = [], [], [] - for task in measured: - statements.append(self.statement_of(task)) - checks.append(user_data_get(task.get('user_data'), 'check_script', '') or '') - groups.append(user_data_get(task.get('user_data'), 'group_id', None)) - payload = [] - for i, statement in enumerate(statements): - siblings = [statements[j] for j in range(len(statements)) - if j != i and groups[j] is not None and groups[j] == groups[i]] - payload.append({'statement': statement, 'check': checks[i], - 'references': self.task_bank.references(statement, siblings)}) - scores: List[Optional[float]] = [None] * len(measured) - pending = list(range(len(measured))) - for attempt in range(self.NOVELTY_TRIES): - batch = [payload[i] for i in pending] - try: - got = list(self.novelty_fn(batch)) - except Exception as e: # noqa - logger.warning(f'[{type(self).__name__}] novelty scoring failed on ' - f'{len(batch)} proposals, try {attempt + 1} of ' - f'{self.NOVELTY_TRIES} ({type(e).__name__}: {e})') - self._bump('novelty_error', len(batch)) - continue - if len(got) != len(batch): - # Zipping a short list would pair later proposals with someone - # else's number, so the whole reply is dropped. - logger.warning(f'[{type(self).__name__}] novelty judge returned ' - f'{len(got)} scores for {len(batch)} proposals, try ' - f'{attempt + 1} of {self.NOVELTY_TRIES}; ignoring them') - self._bump('novelty_length_mismatch', len(batch)) - continue - still: List[int] = [] - for i, score in zip(pending, got): - if score is None: - still.append(i) - else: - scores[i] = score - if not still: - break - logger.info(f'[{type(self).__name__}] novelty: {len(still)} of ' - f'{len(batch)} left unscored on try {attempt + 1}; ' - f'asking again') - pending = still - else: - pending = [i for i, s in enumerate(scores) if s is None] - unscored = [i for i, s in enumerate(scores) if s is None] - if unscored: - self._bump('novelty_unjudged', len(unscored)) - logger.warning(f'[{type(self).__name__}] {len(unscored)} proposal(s) ' - f'still unscored after {self.NOVELTY_TRIES} tries; their ' - f'keyword groups are dropped') - return scores - - def statement_of(self, task: Trajectory) -> str: - """The statement text a task was built around: its first user message.""" - for message in task.get('messages') or []: - if isinstance(message, dict) and message.get('role') == 'user': - return message.get('content') or '' - return '' - - def solver_prompt(self, task: Trajectory) -> Trajectory: - """The statement as the solver first sees it: system message, query, tools. - - ``solver_prompt_fn`` is how the surrounding script hands over the same - opening the eval script builds, so a task kept at n_pass=4 here is a task - the eval measures the same way. Without one this falls back to the bare - statement, which is what it used to be. - - The schemas travel with the prompt for the same reason they do in round 1: - a solver that cannot see the tool names cannot use them, and would score - zero on every task regardless of difficulty. - """ - if self.solver_prompt_fn is not None: - messages = task.get('messages') or [] - query = next((m.get('content', '') for m in messages - if m.get('role') == 'user'), '') - prompt = self.solver_prompt_fn(query) - if not prompt.get('tools') and self.tool_schemas: - prompt['tools'] = self.tool_schemas - return prompt - prompt: Trajectory = {'messages': [dict(m) for m in task.get('messages') or []]} - if self.tool_schemas: - prompt['tools'] = self.tool_schemas - return prompt - - def judge_attempt(self, task: Trajectory, attempt: Trajectory, - slot: int = 0) -> bool: - """Run the check script against sandbox ``slot``'s current state. - - Also hands the whole attempt to ``solver_sink`` when one is given. The - difficulty stage otherwise reports a single number per task, and - ``n_pass=0`` reads the same whether the task is impossible, the statement - withholds something the check demands, or the solver merely gave up -- - which are three different things to fix. The evidence that separates them - is the attempt itself and the state it left, so both are recorded here - rather than reconstructed later. - """ - script = user_data_get(task.get('user_data'), 'check_script', '') - if not script: - return False - exit_code, output = self.env(slot).run_script(script) - if self.solver_sink is not None: - messages = task.get('messages') or [{}] - record = { - 'statement': messages[0].get('content', ''), - 'check_script': script, - 'passed': exit_code == 0, - 'check_exit': exit_code, - 'check_output': output, - # Whether the reply was cut off at the generation budget. The - # difficulty stage drops such an attempt from its denominator, so - # the flag has to travel with the record for the dropped count to - # be reproducible from the dump. - 'truncated': bool((attempt or {}).get('truncated')), - 'attempt': attempt, - 'end_state': self.env(slot).snapshot()[0], - } - with self._sink_lock: - self.solver_sink(record) - return exit_code == 0 - - def on_difficulty_measured(self, candidates: List[Trajectory]) -> None: - """Remember the topics behind the candidates nobody solved.""" - if self.keywords is not None: - self.keywords.remember_unsolved(candidates) - - # ------------------------------------------------------------ feedback - - def expand_hard_keywords(self) -> int: - """Brainstorm more topics in the families that produced the hardest tasks.""" - return self.keywords.expand_hard() if self.keywords is not None else 0 + kept: List[Trajectory] = [] + for proposal, task, n_pass in zip(successful, measured, passes): + proposal.task = task + proposal.n_pass = n_pass + proposal.reward = self.challenger_reward(n_pass) + if low <= n_pass <= high: + proposal.outcome = 'kept' + kept.append(task) + else: + proposal.outcome = 'outside_band' + self._record_proposals() + return kept diff --git a/src/twinkle_agentic/challenger/api.py b/src/twinkle_agentic/challenger/api.py deleted file mode 100644 index bb1141e2a..000000000 --- a/src/twinkle_agentic/challenger/api.py +++ /dev/null @@ -1,154 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""An OpenAI-compatible API, reached as if it were one more explorer. - -Both halves of this loop hand some rounds to a stronger model over an API -- the -ones that are answers rather than actions: writing a check script, describing a -task, brainstorming keywords. The rule that decides what may go over the API is -that the reply must not enter a trainable trajectory, and these do not. - -The call itself was written twice, here and in ``cookbook/rsi``, in the same -twenty lines each time: append a user message, send, keep the text, survive a -raised exception. :class:`ApiModel` is that call, once. :class:`ApiExplorer` wraps -it in the :data:`~.base.Explorer` signature, so anything here that takes an -explorer -- the keyword bank above all -- can be pointed at the API without -knowing it is one, and can be handed a local explorer to fall back on when the -API is unreachable. -""" -from typing import Any, Dict, List, Optional, Sequence - -from twinkle.data_format import SamplingParams, Trajectory -from twinkle.utils import get_logger -from .base import Explorer - -logger = get_logger() - -__all__ = ['ApiExplorer', 'ApiModel'] - - -class ApiModel: - """One OpenAI-compatible client, the extra body it always sends, and one call. - - Args: - api: the client, called as ``api(request, params)``, or with - ``extra_body=`` when there is one. ``twinkle_agentic.protocol.openai`` - provides one; anything with that signature will do. - extra_body: sent on every call (e.g. ``{'thinking_budget': N}`` to cap a - reasoning model). ``None`` sends the request unmodified. - name: what log lines call this model, normally the caller's class name. - """ - - def __init__(self, api: Any, *, extra_body: Optional[Dict[str, Any]] = None, - name: str = 'api'): - self.api = api - self.extra_body = dict(extra_body) if extra_body else None - self.name = name - - def generate(self, messages: Sequence[Dict[str, Any]], - params: Optional[SamplingParams] = None) -> Optional[str]: - """One reply to ``messages``, as text. ``None`` means the call raised. - - Returning ``None`` rather than raising is what keeps one unreachable call - from ending a run that has hours of sandbox work behind it; every caller - here either rejects that one item or falls back. - - Tools are withdrawn for these rounds on purpose -- they are answers, not - actions -- so only the text is kept and any structured ``tool_calls`` the - API returned are dropped. - """ - request: Trajectory = {'messages': list(messages)} - try: - if self.extra_body: - reply = self.api(request, params, extra_body=self.extra_body) - else: - reply = self.api(request, params) - except Exception as exc: # noqa: BLE001 -- one bad call must not kill the run - logger.warning(f'[{self.name}] API call failed: {type(exc).__name__}: {exc}') - return None - if isinstance(reply, list): - reply = reply[0] if reply else {} - return (reply.get('content') if isinstance(reply, dict) else None) or '' - - def reply(self, messages: List[Dict[str, Any]], user_text: str, - params: Optional[SamplingParams] = None) -> Optional[str]: - """Append ``user_text`` and one reply to ``messages``; return the reply. - - For the staged conversations: a check script asked for over the end state, - then a statement asked for over the check. ``messages`` is the caller's - private copy, never a trainable trajectory, so mutating it in place costs - the model nothing. A failed call leaves the user message appended and no - assistant message, which is what the caller would have to write out by - hand to retry. - """ - messages.append({'role': 'user', 'content': user_text}) - content = self.generate(messages, params) - if content is None: - return None - messages.append({'role': 'assistant', 'content': content}) - return content - - -class ApiExplorer: - """An :data:`~.base.Explorer` that answers single text rounds over an API model. - - For the keyword bank, whose calls are one round each and whose replies are - parsed into a list and thrown away: no tokens of them are ever trained on, so - a stronger model may write them. That matters more than it sounds. The bank is - the single input every task downstream is built from, and a 4B policy at the - temperature diversity needs is the wrong instrument for a category rule list - this long -- measured over 1344 locally generated keywords, 31% of one - category named an activity where the rules asked for a computation, and 24% of - another needed hardware the sandbox does not have. - - Args: - model: the :class:`ApiModel` to ask. - params: sampling params for these calls. A per-call ``sampling_params`` - overrides them, so a caller that already sizes its own calls keeps - doing so. - fallback: local explorer for whichever prompts the API could not answer. - Without one, a failed call comes back as a trajectory with no - assistant message, which every parser here reads as an empty reply. - With one, an unreachable API cannot leave a keyword category dry -- - and dry means keyword-less prompts and a run that looks healthy while - producing one prompt over and over, the exact failure the bank's - refill logic exists to prevent. - - Every returned trajectory carries ``via``: ``'api'`` or ``'local-fallback'``. - It is the one thing a reader of the keyword dump cannot reconstruct afterwards, - and the two halves answer at measurably different quality. - """ - - def __init__(self, model: ApiModel, *, params: Optional[SamplingParams] = None, - fallback: Optional[Explorer] = None): - self.model = model - self.params = params - self.fallback = fallback - - def __call__(self, prompts: Sequence[Trajectory], - sampling_params: Optional[SamplingParams] = None) -> List[Trajectory]: - """One API call per prompt, in order, then the failures in one local batch. - - Serially, where a local explorer would take the whole batch at once: - nothing here knows the API's rate limit, and firing a 32-call expansion at - it is how that gets discovered. The failures are gathered and handed to the - fallback together, because a local sampler shards a batch over its workers - and one prompt at a time would leave most of them idle. - """ - params = sampling_params or self.params - out: List[Optional[Trajectory]] = [] - failed: List[int] = [] - for prompt in prompts: - messages = [dict(m) for m in prompt.get('messages') or []] - content = self.model.generate(messages, params) - if content is None: - failed.append(len(out)) - out.append(None) - continue - messages.append({'role': 'assistant', 'content': content}) - out.append({'messages': messages, 'via': 'api'}) - if failed and self.fallback is not None: - local = self.fallback([prompts[i] for i in failed]) - for i, trajectory in zip(failed, local): - answered = dict(trajectory) - answered['via'] = 'local-fallback' - out[i] = answered - return [t if t is not None else {'messages': [], 'via': None} for t in out] diff --git a/src/twinkle_agentic/challenger/base.py b/src/twinkle_agentic/challenger/base.py index 69ab74673..1c819533b 100644 --- a/src/twinkle_agentic/challenger/base.py +++ b/src/twinkle_agentic/challenger/base.py @@ -1,508 +1,132 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Challenger: turn raw material into training tasks. - -A challenger invents the problems a solver will later be trained on. The three -things that vary between deployments are all injected: - -* **what to ask for** -- the system prompt, and the parser that reads the - answer back. They are one contract, so they are passed together. -* **how to explore** -- an :class:`Explorer`, i.e. anything that takes a batch - of trajectories and returns them with the model's reply appended. Both - rollouts in :mod:`twinkle_agentic.rollout` have that signature already, so a - challenger can explore *with tools* -- running code, reading files -- while - it invents, over a local sampler or over an HTTP endpoint alike. - :class:`twinkle_agentic.rollout.MultiTurnRollout` accepts either backend. -* **what counts as a keeper** -- subclasses decide, in :meth:`Challenger.build`. -* **how hard is hard enough** -- optional. Ask for ``solver_rollouts`` attempts per - candidate and only tasks the model solves *sometimes* are kept: a task every - attempt gets right, or none does, gives GRPO a zero gradient, so it costs a - training slot and teaches nothing. Counting the attempts is the same work in - every domain and lives here; deciding whether one attempt was right is not, - and is left to :meth:`Challenger.judge_attempt`. - -Everything a strategy needs beyond that (seed examples, keyword banks) goes in -``__init__``; :meth:`Challenger.__call__` only says how many tasks you want per -batch. It is a generator that yields *full* batches: challengers throw away -most of what they propose -- keep rates of a few percent are normal once the -difficulty filter runs -- so the alternative is a caller that has to cope with -ragged batches for reasons that have nothing to do with it. -""" -import math -import random +"""Reusable lifecycle for task challengers.""" from abc import ABC, abstractmethod -from concurrent.futures import ThreadPoolExecutor, as_completed -from dataclasses import dataclass -from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Callable, Iterator, List, Optional, Sequence, Tuple -from twinkle.data_format import SamplingParams, Trajectory, attach_user_data +from twinkle.data_format import Trajectory from twinkle.utils import get_logger from twinkle_agentic.envs import Env logger = get_logger() -__all__ = ['Challenger', 'Explorer', 'KeywordPrompts', 'PromptSet'] - -# A batch of trajectories in, the same trajectories with the model's reply -# appended out. MultiTurnRollout accepts either backend and also accepts a -# per-call ``sampling_params=`` keyword, which is how the difficulty stage asks -# for its own temperature and length budget without a second explorer. -Explorer = Callable[[List[Trajectory]], List[Trajectory]] - - -@dataclass -class KeywordPrompts: - """The three strings a :class:`~.keywords.KeywordBank` sends, and nothing else. - - Its own type rather than the caller's prompt object: the challengers here - carry a dozen other prompts, the RSI drivers in ``cookbook/rsi`` keep theirs - as module constants, and a bank that reached into either by attribute name - would be coupled to both spellings. Building one of these is how a caller - says which of its strings are the keyword ones -- see - :meth:`PromptSet.keyword_prompts` for the challengers' answer. - - Lives here rather than beside the bank so that :class:`PromptSet` can produce - one without importing it. - - Args: - system: the system message every keyword call carries. - user: asks for ``{k}`` topics in a category described by ``{desc}``. - expand_user: asks for ``{m}`` more topics like ``{kw}``, optionally with - the category's ``{desc}``. Only :meth:`.KeywordBank.expand_hard` - needs it. - """ - system: str - user: str - expand_user: str = '' - - def __post_init__(self): - missing = [f for f in ('system', 'user') if not getattr(self, f).strip()] - if missing: - raise ValueError(f'KeywordPrompts needs {" and ".join(missing)}: a dry ' - f'category could not be refilled without it.') - - -class PromptSet: - """Base for a challenger's bundle of prompts: what is required, and validation. - - Every challenger here is a dataclass of strings plus the same three questions - -- are the mandatory ones filled in, do the optional ones carry the - placeholders they will be formatted with, and does this configuration have - the ones it needs. Answering them once means a missing placeholder is caught - at construction in every domain, rather than as a ``KeyError`` mid-run in - whichever domain remembered to check. - - Subclasses declare: - - * ``_REQUIRED`` -- fields that must carry text. - * ``_REQUIRED_FIELDS`` -- field -> placeholders its text must contain. - """ +__all__ = ['Challenger'] - _REQUIRED: Tuple[str, ...] = () - _REQUIRED_FIELDS: Dict[str, Sequence[str]] = {} - def __post_init__(self): - name = type(self).__name__ - for field in self._REQUIRED: - if not getattr(self, field).strip(): - raise ValueError(f'{name}.{field} is required') - for field, placeholders in self._REQUIRED_FIELDS.items(): - text = getattr(self, field) - if not text: - continue - for placeholder in placeholders: - if '{' + placeholder + '}' not in text: - raise ValueError(f'{name}.{field} must contain {{{placeholder}}}') - - def require(self, *names: str) -> None: - """Raise unless every named prompt was supplied. - - For what only a configuration knows: drawing from a keyword bank needs the - keyword prompts, seeds need the seed prompt, and a challenger asks for the - ones its arguments imply. - """ - missing = [n for n in names if not getattr(self, n).strip()] - if missing: - name = type(self).__name__ - separator = f', {name}.' - raise ValueError(f'this configuration needs {name}.' - f'{separator.join(missing)}') - - def keyword_prompts(self) -> KeywordPrompts: - """The keyword subset, for the bank. Validated by :meth:`require` first.""" - self.require('keyword_system', 'keyword_user') - return KeywordPrompts(system=self.keyword_system, user=self.keyword_user, - expand_user=self.keyword_expand_user) +def _parallel(fn: Callable[[int], Any], count: int) -> List[Any]: + """Run ``fn`` over ``range(count)`` concurrently, preserving order.""" + if count <= 1: + return [fn(i) for i in range(count)] + out: List[Any] = [None] * count + with ThreadPoolExecutor(max_workers=count) as pool: + futures = {pool.submit(fn, i): i for i in range(count)} + for future, i in futures.items(): + out[i] = future.result() + return out class Challenger(ABC): - """Base class: propose, explore, keep, repeat until the batch is full. + """Common batching and environment lifecycle for task challengers. - Args: - explorer: takes a batch of trajectories and returns them with the - model's reply appended -- typically a - :class:`twinkle_agentic.rollout.MultiTurnRollout` over a local - sampler or an API endpoint. - system: system prompt handed to the model. It carries the output - contract, which is why ``build`` -- the code that reads that output - back -- lives in the same subclass. - envs: the environments this challenger works in, one per slot. A slot is - owned whole for as long as a job needs it, because the workspace - lives inside it, so ``len(envs)`` is also how many jobs may run at - once. Both halves take the same parameter and reach it the same way - (:meth:`env`), which is what lets one caller decide where everything - it runs is executed and graded: ``[LocalEnv()]`` keeps judgement on - the training host and costs milliseconds, sandbox slots trade that - for isolation. Empty is allowed for a challenger that executes - nothing; :meth:`env` then says so rather than raising IndexError. - max_proposals_per_round: ceiling on how many proposals one round may - request. Without it a low keep rate makes the estimator ask for an - unbounded batch after the first round. - solver_rollouts: attempts per candidate in the difficulty stage. ``0`` - skips the stage entirely; any other value requires the subclass to - implement :meth:`solver_prompt` and :meth:`judge_attempt`. - keep_pass_band: ``(low, high)`` attempt counts, inclusive on both ends: - keep a candidate only if that many of its ``solver_rollouts`` - attempts succeeded. Required whenever the stage runs, and has no - default because the counts are absolute -- ``(1, 7)`` reads as "hard - but solvable" against eight rollouts and as something far stricter - against sixteen, so it has to be written by whoever chose the - rollout count. - solver_params: sampling params for the difficulty stage only, passed to - the explorer per call. ``None`` reuses whatever the explorer was - built with -- which is usually the proposing temperature, and that - is higher than a solver should get. - solver_explorer: optional separate explorer for the difficulty stage. - ``None`` reuses the main explorer. Useful when the solver needs a - different configuration (e.g. sandbox tools, more turns) than the - proposer. - seed: RNG seed for whatever sampling a subclass does. ``None`` leaves - the RNG unseeded. + Subclasses define how a round builds its prompt, explores it, and measures + candidate difficulty. One environment is owned by one concurrent job for the + complete lifetime of that job. """ def __init__( - self, - explorer: Explorer, - *, - system: str, - envs: Sequence[Env] = (), - max_proposals_per_round: int = 512, - solver_rollouts: int = 0, - keep_pass_band: Optional[Tuple[int, int]] = None, - solver_params: Optional[SamplingParams] = None, - solver_explorer: Optional[Explorer] = None, - seed: Optional[int] = None, + self, + *, + envs: Sequence[Env], + num_challenger_rollouts: int = 8, + num_solver_rollouts: int = 8, + pass_band: Tuple[float, float] = (1.0, 7.0), + max_empty_rounds: int = 0, ): - if not system: - raise ValueError('Challenger needs a system prompt: it carries the output ' - 'contract that build() parses back.') - if solver_rollouts < 0: - raise ValueError(f'solver_rollouts must be >= 0, got {solver_rollouts}') - if solver_rollouts: - # Checked here rather than at first use: the stage runs after a full - # round of generation, and finding out then that this challenger - # cannot grade an attempt wastes the whole round. - missing = [ - name for name in ('solver_prompt', 'judge_attempt') - if getattr(type(self), name) is getattr(Challenger, name) - ] - if missing: - raise NotImplementedError( - f'solver_rollouts={solver_rollouts} needs {type(self).__name__} to ' - f'implement {", ".join(missing)}; pass solver_rollouts=0 to skip the ' - f'difficulty stage.') - if keep_pass_band is None: - raise ValueError(f'solver_rollouts={solver_rollouts} needs ' - f'keep_pass_band=(low, high): the band is in attempt ' - f'counts, so what it asks for depends on how many ' - f'attempts were run.') - if len(keep_pass_band) != 2: - raise ValueError(f'keep_pass_band is (low, high) in attempt counts, got ' - f'{keep_pass_band}') - low, high = keep_pass_band - if not 0 <= low <= high <= solver_rollouts: - raise ValueError( - f'keep_pass_band must satisfy 0 <= low <= high <= solver_rollouts, ' - f'got {keep_pass_band} against solver_rollouts={solver_rollouts}') - elif keep_pass_band is not None: - raise ValueError('keep_pass_band has nothing to filter while ' - 'solver_rollouts=0 leaves the difficulty stage off; pass ' - 'the rollout count too, or drop the band.') - self.explorer = explorer - self.system = system + if not envs: + raise ValueError('envs is empty: a challenger needs a workspace to act in and grade') + if num_challenger_rollouts < 1: + raise ValueError(f'num_challenger_rollouts must be >= 1, got ' + f'{num_challenger_rollouts}') + if num_solver_rollouts < 0: + raise ValueError(f'num_solver_rollouts must be >= 0, got {num_solver_rollouts}') + if max_empty_rounds < 0: + raise ValueError(f'max_empty_rounds must be >= 0, got {max_empty_rounds}') + if num_solver_rollouts: + if len(pass_band) != 2: + raise ValueError(f'pass_band is (low, high) in attempt counts, got {pass_band}') + low, high = pass_band + if not 0 <= low <= high <= num_solver_rollouts: + raise ValueError(f'pass_band must satisfy 0 <= low <= high <= num_solver_rollouts, got ' + f'{pass_band} against num_solver_rollouts={num_solver_rollouts}') self.envs = list(envs) - self.max_proposals_per_round = max_proposals_per_round - self.solver_rollouts = solver_rollouts - self.keep_pass_band = keep_pass_band - self.solver_params = solver_params - self.solver_explorer = solver_explorer - self.rng = random.Random(seed) - # Running tally, used to size the next round and worth logging: a keep - # rate near zero means the prompt or the filter is miscalibrated, not - # that the model is bad. + self.num_challenger_rollouts = num_challenger_rollouts + self.num_solver_rollouts = num_solver_rollouts + self.pass_band = pass_band + self.max_empty_rounds = max_empty_rounds self.n_proposed = 0 self.n_kept = 0 - # ----------------------------------------------------------------- envs - @property def n_slots(self) -> int: """How many jobs may run at once: one per environment.""" return len(self.envs) def env(self, slot: int = 0) -> Env: - """The environment for ``slot``. - - Fetched per use rather than held in a local, so a slot that had to be - rebuilt underneath is picked up on the next call instead of being used - dead. ``slot=0`` is the default because a challenger with nothing to run - concurrently -- one script, no state to share -- has only one. - """ - if not self.envs: - raise RuntimeError( - f'{type(self).__name__} was given no envs, so there is nowhere to run ' - f'anything: pass envs=[LocalEnv()] to execute on the training host, or ' - f'sandbox slots to execute in one.') + """Return the current environment for ``slot``.""" return self.envs[slot] - # ------------------------------------------------------------- subclass - @abstractmethod - def propose(self, count: int) -> List[Trajectory]: - """Build ``count`` prompt trajectories to hand to the explorer. - - Returning fewer than asked is allowed and means the source material ran - out; :meth:`__call__` stops once a round proposes nothing. - """ + def _build_challenge_prompt(self) -> Optional[Trajectory]: + """Build one round's shared prompt, or return None when exhausted.""" @abstractmethod - def build(self, explored: List[Trajectory]) -> List[Optional[Trajectory]]: - """Turn explored proposals into finished tasks. + def _explore(self, prompt: Trajectory) -> List[Trajectory]: + """Generate and validate candidates from one shared prompt.""" - Returns one entry per input, ``None`` for anything rejected -- failed - parse, failed verification, wrong difficulty. Positional so a subclass - can line rejects up against what produced them. - """ - - def solver_prompt(self, task: Trajectory) -> Trajectory: - """The trajectory to hand a solver attempting ``task``. - - Only called when ``solver_rollouts`` is non-zero. It must return a - prompt for every task: a task that cannot be attempted has no measurable - difficulty and should have been rejected in :meth:`build` instead. - """ - raise NotImplementedError() - - def judge_attempt(self, task: Trajectory, attempt: Trajectory) -> bool: - """Did this solver attempt solve ``task``? - - ``attempt`` is the explored :meth:`solver_prompt` trajectory, so the - model's answer is its last assistant message. Program checks only: a - judgement that drifts between rounds turns the difficulty band into - noise. - """ - raise NotImplementedError() - - def on_difficulty_measured(self, candidates: List[Trajectory]) -> None: - """Called once per round with every measured candidate, before filtering. - - Each carries ``n_pass`` / ``n_rollouts`` in its ``user_data``. This is - the only place that sees the candidates the band is about to drop, which - is what a strategy adapting to difficulty needs -- an all-fail task says - more about its source material than a kept one does. - """ - - # ---------------------------------------------------------------- public + @abstractmethod + def _filter_difficulty(self, tasks: List[Trajectory]) -> List[Trajectory]: + """Measure candidate difficulty and return the accepted tasks.""" def __call__(self, batch_size: int, total: Optional[int] = None) -> Iterator[List[Trajectory]]: - """Yield batches of exactly ``batch_size`` finished tasks. - - Args: - batch_size: tasks per yielded batch. - total: stop after this many tasks. ``None`` runs until the source - material is exhausted, which for a from-scratch challenger - means forever -- pass a total or break out of the loop. - - The final batch is short only when the source runs out or ``total`` is - not a multiple of ``batch_size``. - """ + """Yield finished tasks in batches.""" if batch_size <= 0: raise ValueError(f'batch_size must be positive, got {batch_size}') pending: List[Trajectory] = [] produced = 0 + empty_rounds = 0 while total is None or produced < total: want = batch_size if total is None else min(batch_size, total - produced) while len(pending) < want: - kept = self._round(want - len(pending)) + kept = self._round() if kept is None: - # Source exhausted: hand back whatever is left rather than - # spinning, and let the caller see a short final batch. if pending: yield pending return - pending.extend(kept) + if kept: + empty_rounds = 0 + pending.extend(kept) + continue + empty_rounds += 1 + if self.max_empty_rounds and empty_rounds >= self.max_empty_rounds: + logger.warning(f'[{type(self).__name__}] stopped after {empty_rounds} ' + 'consecutive rounds without a usable task') + if pending: + yield pending + return yield pending[:want] produced += want pending = pending[want:] - # --------------------------------------------------------------- private - - def _round(self, missing: int) -> Optional[List[Trajectory]]: - """One propose/explore/build/measure cycle. ``None`` means the source is dry.""" - count = min(self._estimate(missing), self.max_proposals_per_round) - proposals = self.propose(count) - if not proposals: + def _round(self) -> Optional[List[Trajectory]]: + """Run one proposal group; None means the source is exhausted.""" + prompt = self._build_challenge_prompt() + if prompt is None: return None - explored = self.explore(proposals) - built = self.build(explored) - usable = [t for t in built if t is not None] - kept = self._filter_difficulty(usable) if self.solver_rollouts else usable - self.n_proposed += len(proposals) + verified = self._explore(prompt) + kept = self._filter_difficulty(verified) + self.n_proposed += self.num_challenger_rollouts self.n_kept += len(kept) - band = (f', in difficulty band {len(kept)}' if self.solver_rollouts else '') - logger.info(f'[{type(self).__name__}] proposed {len(proposals)}, usable ' - f'{len(usable)}{band} (cumulative {self.n_kept}/{self.n_proposed})') + logger.info(f'[{type(self).__name__}] {self.num_challenger_rollouts} episodes, ' + f'{len(verified)} verified, {len(kept)} in band ' + f'(cumulative {self.n_kept}/{self.n_proposed})') return kept - - def explore( - self, - trajectories: List[Trajectory], - sampling_params: Optional[SamplingParams] = None, - **kwargs: Any, - ) -> List[Trajectory]: - """Run the explorer over a batch, optionally overriding its sampling params. - - The override is only forwarded when asked for, so a plain callable - explorer keeps working; both rollouts in - :mod:`twinkle_agentic.rollout` accept it. Anything else in ``kwargs`` is - passed straight through for the same reason -- a caller that needs a - rollout-specific hook (``followup_fn``) says so per call, and an explorer - that does not take it fails loudly instead of silently ignoring it. - """ - if not trajectories: - return [] - if sampling_params is not None: - kwargs['sampling_params'] = sampling_params - if not kwargs: - return self.explorer(trajectories) - return self.explorer(trajectories, **kwargs) - - def _solver_explore( - self, - trajectories: List[Trajectory], - sampling_params: Optional[SamplingParams] = None, - **kwargs: Any, - ) -> List[Trajectory]: - """Run solver attempts through the solver explorer, or fall back to the main one. - - Subclasses that need per-attempt isolation (e.g. sandbox workspace reset) - override this rather than the whole difficulty filter. Extra kwargs are - forwarded, which is how such a subclass says which sandbox each attempt - runs in (``tool_manager`` as a list, one entry per trajectory). - """ - if self.solver_explorer is not None: - if sampling_params is None: - return self.solver_explorer(trajectories, **kwargs) - return self.solver_explorer(trajectories, sampling_params=sampling_params, **kwargs) - return self.explore(trajectories, sampling_params=sampling_params, **kwargs) - - def _filter_difficulty(self, tasks: List[Trajectory]) -> List[Trajectory]: - """Attempt each task ``solver_rollouts`` times; keep the ones in the band. - - All attempts for the whole batch go out in one explorer call: on the - sampler path that is one batched generate, and the alternative -- a call - per task -- would leave the GPUs idle between them. - """ - if not tasks: - return [] - prompts: List[Trajectory] = [] - owners: List[int] = [] - for i, task in enumerate(tasks): - prompt = self.solver_prompt(task) - for _ in range(self.solver_rollouts): - prompts.append(dict(prompt)) - owners.append(i) - - attempts = self._solver_explore(prompts, sampling_params=self.solver_params) - if len(attempts) != len(prompts): - # Counting a partial return would silently understate every affected - # task's pass count, i.e. report tasks as harder than they are. - raise RuntimeError(f'explorer returned {len(attempts)} attempts for ' - f'{len(prompts)} solver prompts; expected one per prompt.') - - passes = [0] * len(tasks) - for owner, attempt in zip(owners, attempts): - if self.judge_attempt(tasks[owner], attempt): - passes[owner] += 1 - - measured = [ - attach_user_data(task, n_pass=passes[i], n_rollouts=self.solver_rollouts) - for i, task in enumerate(tasks) - ] - self.on_difficulty_measured(measured) - low, high = self.keep_pass_band - return [t for t, n in zip(measured, passes) if low <= n <= high] - - def _estimate(self, missing: int) -> int: - """How many proposals to make for ``missing`` keepers. - - The first round has nothing to go on and asks for exactly what is - missing; after that the measured keep rate scales the request. A round - that kept nothing leaves the rate at its last non-zero estimate rather - than dividing by zero. - """ - if self.n_kept <= 0: - return missing - rate = self.n_kept / max(1, self.n_proposed) - return max(missing, math.ceil(missing / rate)) - - # ------------------------------------------------------------- utilities - - def prompt_trajectory(self, user: str, **extra: Any) -> Trajectory: - """A two-message trajectory carrying this challenger's system prompt.""" - trajectory: Trajectory = { - 'messages': [ - {'role': 'system', 'content': self.system}, - {'role': 'user', 'content': user}, - ], - } - trajectory.update(extra) - return trajectory - - @staticmethod - def draw(rng: random.Random, pool: Sequence[Any], count: int) -> List[Any]: - """Draw ``count`` items with replacement; ``[]`` for an empty pool.""" - return [rng.choice(pool) for _ in range(count)] if pool else [] - - -def map_parallel(fn: Callable[[Any], Any], items: Sequence[Any]) -> List[Any]: - """Map ``fn`` over ``items`` at once, results in input order. - - Every use of this is waiting on a sandbox or on a model call, not computing, - so the thread pool is the point. One item runs inline: a pool for a single - call only adds a thread, and it keeps a serial configuration on exactly the - code path it had before. - """ - items = list(items) - if len(items) <= 1: - return [fn(item) for item in items] - out: List[Any] = [None] * len(items) - with ThreadPoolExecutor(max_workers=len(items)) as pool: - futures = {pool.submit(fn, item): i for i, item in enumerate(items)} - for fut in as_completed(futures): - out[futures[fut]] = fut.result() - return out - - -def sampling_params_of(explorer: Any) -> Optional[SamplingParams]: - """The sampling params an explorer was built with, when it exposes them. - - Only used for logging what a run actually asked for; both explorer kinds - keep the field under the same name. - """ - params: Optional[SamplingParams] = getattr(explorer, 'sampling_params', None) - return params - - -def as_dict(trajectory: Trajectory) -> Dict[str, Any]: - """A plain dict copy, for writing a trajectory to jsonl.""" - return {k: v for k, v in trajectory.items()} diff --git a/src/twinkle_agentic/challenger/code.py b/src/twinkle_agentic/challenger/code.py deleted file mode 100644 index be7b8187e..000000000 --- a/src/twinkle_agentic/challenger/code.py +++ /dev/null @@ -1,616 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Code challenger: invent Python problems whose ground truth was executed. - -The task is built backwards. The model writes a problem *and* a reference -solution; the solution is run to capture what each check expression actually -returns, and those captured values become the asserts. So the answer exists -before the question does, and no external labelling is involved. Two gates carry -over from earlier runs, both from real failures: - -* the reference solution must pass its own asserts, or the ground truth is noise; -* output capture uses a sentinel marker plus the exit status, never the last - stdout line, so a startup banner can never be read as a result. - -Prompt text is not here. Every string the model sees arrives in -:class:`CodePrompts`, built by whoever runs the challenger -- see -``cookbook/rsi/code/challenge_prompts.py``. Neither is execution: every script -runs in an :class:`~twinkle_agentic.envs.base.Env`, which is the same interface -the agentic half verifies through, so where a task gets graded is a decision made -once by the caller rather than twice by the two halves. What stays here is the -machinery that cannot be restated in a prompt: the assert capture, the -constant-answer check, and how a proposal becomes a task. Neither is the keyword -bank -- drawing, refilling and expanding topics is the same cycle on both halves, -so it lives once in :mod:`.keywords` and this challenger holds one. -""" -import json -import os -import re -from dataclasses import dataclass -from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple - -from twinkle.data_format import SamplingParams, Trajectory, user_data_get -from twinkle.utils import get_logger -from twinkle_agentic.envs import Env -from twinkle_agentic.utils.code_utils import strip_reasoning, unwrap_code -from twinkle_agentic.utils.message_utils import assistant_text -from .base import Challenger, Explorer, PromptSet, attach_user_data -from .keywords import KeywordBank, KeywordStore - -logger = get_logger() - -__all__ = [ - 'CodeChallenger', 'CodePrompts', 'build_asserts', 'is_constant_answer', - 'load_seeds', 'parse_challenge', 'run_asserts', 'run_check_script', -] - -# Isolates a captured value from anything else the script prints. -_MARK = '__RSI_GT__' -_JSON_FENCE_RE = re.compile(r'^\s*```(?:json)?\s*|\s*```\s*$', re.I) - - -def run_check_script(code: str, check_script: str, env: Env, - timeout: int = 30) -> Tuple[bool, str]: - """Run ``code`` against ``check_script`` in ``env``; True when it exits 0. - - One script, one exit status -- the same judgement the agentic half makes, so - a task from either half is graded the same way. The run's output comes back - too: a verdict that only says "wrong" leaves a second attempt nothing to go - on. - """ - if not code.strip(): - return False, 'no code was produced' - if not check_script.strip(): - return False, 'no check script was produced' - rc, out = env.run_script(f'{code}\n\n{check_script}', timeout=timeout) - return rc == 0, out - - -def run_asserts(code: str, setup: str, asserts: List[str], env: Env, - timeout: int = 30) -> bool: - """True when every assert passes (exit status 0). - - For callers holding a list of asserts rather than one check script -- a - tests file, say. The list plus the setup *is* the check script. - """ - parts = [setup] if (setup or '').strip() else [] - parts.extend(asserts or ()) - return run_check_script(code, '\n\n'.join(parts), env, timeout)[0] - - -def build_asserts(solution: str, checks: List[str], env: Env, timeout: int = 30, - max_checks: int = 6) -> Optional[List[str]]: - """Run the reference solution once to capture each check's repr, then form - ``assert <check> == <captured>``. - - Returns None if the solution crashed or produced no usable output -- the - caller drops that problem. The marker plus the exit status is what makes the - capture trustworthy: a crash or a banner line can never become a value. - """ - checks = [c for c in checks if isinstance(c, str) and c.strip()][:max_checks] - if not checks: - return None - lines = [solution, ''] - for i, c in enumerate(checks): - # repr on its own line, tagged with index; a check that raises makes the - # whole script exit non-zero -> we drop the problem. Pure f-string (no %% - # formatting) so a check expression containing '%' (modulo/percent) is safe. - lines.append(f'print("{_MARK}{i}=" + repr({c}))') - rc, out = env.run_script('\n'.join(lines), timeout=timeout) - if rc != 0: - return None - captured: Dict[int, str] = {} - for line in out.splitlines(): - if line.startswith(_MARK): - try: - idx_str, val = line[len(_MARK):].split('=', 1) - idx = int(idx_str) - except (ValueError, IndexError): - continue - if idx in captured: - # Two lines claiming the same check. The script prints each one - # exactly once, so a second one came from the solution itself -- - # stderr is part of the output now, and a solution that can - # redefine what its own check returned is not ground truth. - return None - captured[idx] = val - if len(captured) != len(checks): - return None - # The captured text is a repr, so it is a valid literal to compare against. - return [f'assert ({c}) == ({captured[i]})' for i, c in enumerate(checks)] - - -def _split_top_eq(s: str) -> Optional[tuple]: - """Split on the first top-level ``==``, ignoring anything inside brackets or quotes.""" - depth = 0 - quote = '' - i = 0 - while i < len(s) - 1: - c = s[i] - if quote: - if c == quote: - quote = '' - elif c in '\'"': - quote = c - elif c in '([{': - depth += 1 - elif c in ')]}': - depth -= 1 - elif depth == 0 and c == '=' and s[i + 1] == '=': - return s[:i].strip(), s[i + 2:].strip() - i += 1 - return None - - -def _expected_of(assert_line: str) -> Optional[str]: - """The value the solver actually has to produce for one assert. - - :func:`build_asserts` emits ``assert (<check>) == (<repr>)``, but a check may - itself be a comparison, giving ``assert (f(x) == 3) == (True)``. Reading the - outer side there would report 'True' and make such a problem look - constant-answer, so the inner right-hand side is used instead. An outer - ``False`` pins nothing down at all and is reported as unknown. - """ - m = re.match(r'^\s*assert\s*\((.*)\)\s*==\s*\((.*)\)\s*$', assert_line.strip()) - if not m: - return None - lhs, rhs = m.group(1).strip(), m.group(2).strip() - inner = _split_top_eq(lhs) - if rhs in ('True', 'False') and inner is not None: - return inner[1] if rhs == 'True' else None - return rhs - - -def is_constant_answer(asserts: List[str]) -> bool: - """Would ``return <one constant>`` satisfy every assert? - - Such a problem pays full reward for ignoring its own statement, so it - actively teaches the solver not to read the input. Requires at least two - asserts with a readable expectation: a single assert is trivially - 'constant', and one unreadable assert must not hide a constant set. - """ - vals = [_expected_of(a) for a in asserts] - if any(v is None for v in vals) or len(vals) < 2: - return False - return len(set(vals)) == 1 - - -# ── parsing ──────────────────────────────────────────────────────────────── -def parse_challenge(text: str, require_solution: bool = True) -> Optional[Dict[str, Any]]: - """Pull the ``{problem, solution, entry, checks}`` object out of a completion. - - ``require_solution=False`` is for the two-step flow, whose second call is - told the solution is already known and returns only the statement. - """ - body = strip_reasoning(text) - body = _JSON_FENCE_RE.sub('', body.strip()).strip() - # Grab the outermost {...} if there is leading/trailing prose. - start, end = body.find('{'), body.rfind('}') - if start < 0 or end <= start: - return None - try: - obj = json.loads(body[start:end + 1]) - except (ValueError, TypeError): - return None - if not isinstance(obj, dict): - return None - problem, solution, checks = obj.get('problem'), obj.get('solution'), obj.get('checks') - if not (isinstance(problem, str) and problem.strip() - and isinstance(checks, list) and checks): - return None - if require_solution: - if not (isinstance(solution, str) and solution.strip()): - return None - else: - # Told not to include a solution; if it did anyway, ignore it -- the - # caller overwrites with the code that actually ran. - solution = solution if isinstance(solution, str) else '' - if solution and '```' in solution: - solution = unwrap_code(solution) - return {'problem': problem.strip(), 'solution': (solution or '').strip(), - 'entry': str(obj.get('entry') or '').strip(), 'checks': checks} - - -def load_seeds(path: str) -> List[Dict[str, str]]: - """Read seed problems from a jsonl: dicts with ``query`` and maybe ``code``. - - A seed without ``code`` cannot take the two-step path (there is no reference - solution to build on top of) and falls back to the single-call prompt. - """ - if not path or not os.path.exists(path): - return [] - seeds: List[Dict[str, str]] = [] - with open(path, encoding='utf-8') as f: - for line in f: - line = line.strip() - if not line: - continue - try: - row = json.loads(line) - except (ValueError, TypeError): - continue - q = row.get('query') or row.get('problem') or row.get('prompt') - if isinstance(q, dict): - q = q.get('content') - if not q: - msgs = row.get('messages') or [] - q = next((m.get('content') for m in msgs if m.get('role') == 'user'), None) - if isinstance(q, str) and q.strip(): - seeds.append({'query': q.strip(), 'code': (row.get('code') or '').strip()}) - return seeds - - -# ── prompts (text supplied by the caller) ────────────────────────────────── -@dataclass -class CodePrompts(PromptSet): - """Every string a :class:`CodeChallenger` sends, and nothing else. - - Deliberately without defaults for the always-needed fields: a prompt is the - experiment, so a run has to state which one it used rather than inherit a - library's idea of it. Optional groups stay empty until the feature that - needs them is switched on, and the constructor says so if one is missing. - - Placeholders are checked at construction: a typo'd ``{keywords}`` would - otherwise surface as a KeyError halfway through a generation run. That - checking, and the keyword subset a bank is given, are :class:`.PromptSet`. - """ - - system: str - from_scratch: str - solver_system: str - solver_user: str - from_seed: str = '' - from_keywords: str = '' - from_seed_keywords: str = '' - two_step_system: str = '' - two_step_solution: str = '' - two_step_problem: str = '' - keyword_system: str = '' - keyword_user: str = '' - keyword_expand_user: str = '' - - #: fields that must carry text. - _REQUIRED = ('system', 'from_scratch', 'solver_system', 'solver_user') - #: field -> placeholders it must contain. - _REQUIRED_FIELDS = { - 'solver_user': ('problem', ), - 'from_seed': ('seed', ), - 'from_keywords': ('keywords', ), - 'from_seed_keywords': ('seed', 'keywords'), - 'two_step_solution': ('seed', 'code', 'keywords'), - 'two_step_problem': ('code', 'seed', 'keywords'), - 'keyword_user': ('k', 'desc'), - 'keyword_expand_user': ('kw', 'm'), - } - - -class CodeChallenger(Challenger): - """Propose code problems, execute them for ground truth, keep the graded ones. - - One class rather than several because 'from scratch', 'from a seed problem', - 'from keywords' and the two-step build differ only in which prompt the - proposal carries: parsing, execution, the self-check and the difficulty - band are the same afterwards. Which path a proposal takes is decided per - proposal, so one run mixes them. - - Args: - prompts: every string sent to the model. - explorer: batch-in / batch-out generation, see :class:`.base.Explorer`. - seeds: optional pool from :func:`load_seeds`, drawn with replacement. - keyword_store: optional bank; without it proposals carry no topics. - category_desc / combo_arity / arity_weights / single_kw_prob / - keyword_refill_target / keyword_gen_calls / keyword_refill_tries / - keyword_params / min_batch / expand_per_kw / expand_max_kws: handed to the - :class:`.keywords.KeywordBank` this challenger holds, which is where - they are documented -- they behave the same on the agentic half. - seed_mix_prob: chance a proposal also carries a seed problem, when a - pool was given. - two_step: allow the two-call path (write a harder solution on top of the - seed's reference code, then describe the problem it answers). Needs - a seed carrying ``code`` and at least one keyword, so it is skipped - silently for proposals that have neither. - problem_max_chars: reject statements longer than this. Rambling - non-problems, and they would also crowd out the solver's context. - max_checks / sandbox_timeout: passed to :func:`build_asserts`. - drop_constant_answer: reject problems where one constant satisfies every - assert. - low_pass_expand: a candidate solved this many times or fewer counts as - hard, and its topics are fed back through - :meth:`expand_hard_keywords`. - reject_sink: called with a dict for every rejected proposal. The caller - decides whether that goes to a file; nothing here writes one. - solver_sink: called once per solver attempt in the difficulty stage, with - the check script, the attempt and the verdict. Two things need it: - ``n_pass=0`` reads the same whether the problem is impossible or the - statement withholds a value its asserts demand, and the attempts are - the trainable half of this challenger's output -- see - :meth:`judge_attempt`. Requires a local sampler; an API explorer - returns text without token fields. - keyword_sink: called once per keyword-generation call, with the prompt, - the reply and both halves of the parse. - """ - - def __init__( - self, - prompts: CodePrompts, - explorer: Explorer, - *, - seeds: Sequence[Dict[str, str]] = (), - keyword_store: Optional[KeywordStore] = None, - category_desc: Optional[Dict[str, str]] = None, - seed_mix_prob: float = 0.5, - two_step: bool = True, - combo_arity: str = 'triple', - arity_weights: Optional[Sequence[float]] = None, - single_kw_prob: float = 0.1, - keyword_refill_target: int = 128, - keyword_gen_calls: int = 8, - keyword_refill_tries: int = 2, - keyword_params: Optional[SamplingParams] = None, - min_batch: int = 1, - problem_max_chars: int = 4000, - max_checks: int = 6, - sandbox_timeout: int = 30, - drop_constant_answer: bool = True, - low_pass_expand: int = 0, - expand_per_kw: int = 8, - expand_max_kws: int = 32, - reject_sink: Optional[Callable[[Dict[str, Any]], None]] = None, - solver_sink: Optional[Callable[[Dict[str, Any]], None]] = None, - keyword_sink: Optional[Callable[[Dict[str, Any]], None]] = None, - **challenger_kwargs: Any, - ): - super().__init__(explorer, system=prompts.system, **challenger_kwargs) - if not self.envs: - raise ValueError('envs is empty: there is nowhere to run a check, and every ' - 'proposal would be rejected for a ground truth that never ran.') - if keyword_store is not None: - prompts.require('from_keywords') - self.prompts = prompts - self.seeds = list(seeds) - # The whole keyword cycle -- draw, refill, expand -- is one object shared - # with the agentic challenger rather than a second copy of it here. None - # means no bank was configured, and proposals then carry no topics. - self.keywords: Optional[KeywordBank] = None if keyword_store is None else KeywordBank( - keyword_store, prompts=prompts.keyword_prompts(), - category_desc=category_desc or {}, - explorer=explorer, rng=self.rng, name=type(self).__name__, - sampling_params=keyword_params, sink=keyword_sink, combo_arity=combo_arity, - arity_weights=arity_weights, single_kw_prob=single_kw_prob, - refill_target=keyword_refill_target, gen_calls=keyword_gen_calls, - refill_tries=keyword_refill_tries, min_batch=min_batch, - expand_per_kw=expand_per_kw, expand_max_kws=expand_max_kws) - self.seed_mix_prob = seed_mix_prob - self.two_step = two_step - self.problem_max_chars = problem_max_chars - self.max_checks = max_checks - self.sandbox_timeout = sandbox_timeout - self.drop_constant_answer = drop_constant_answer - self.low_pass_expand = low_pass_expand - self.reject_sink = reject_sink - self.solver_sink = solver_sink - if self.seeds: - # Both are reachable with a bank configured: a proposal draws no - # keywords when every category is dry, and then falls back to the - # seed-only prompt. - prompts.require('from_seed') - if self.keywords is not None: - prompts.require('from_seed_keywords') - if two_step: - prompts.require('two_step_system', 'two_step_solution', 'two_step_problem') - # Why proposals died, for the caller to log; the shape a run is judged on. - self.stats: Dict[str, int] = { - 'parsed': 0, 'parse_fail': 0, 'stage1_no_code': 0, 'too_long': 0, - 'gt_fail': 0, 'selfcheck_fail': 0, 'constant_answer': 0, - } - - # ------------------------------------------------------------- proposing - - def propose(self, count: int) -> List[Trajectory]: - proposals: List[Trajectory] = [] - for _ in range(count): - picks = self.keywords.draw() if self.keywords else [] - body = KeywordBank.block(picks) - use_seed = bool(self.seeds) and self.rng.random() < self.seed_mix_prob - seed = self.rng.choice(self.seeds) if use_seed else None - two = bool(use_seed and self.two_step and picks and seed and seed.get('code')) - if two: - system = self.prompts.two_step_system - user = self.prompts.two_step_solution.format( - seed=seed['query'], code=seed['code'], keywords=body) - elif use_seed and picks: - system = self.prompts.system - user = self.prompts.from_seed_keywords.format(seed=seed['query'], keywords=body) - elif use_seed: - system = self.prompts.system - user = self.prompts.from_seed.format(seed=seed['query']) - elif picks: - system = self.prompts.system - user = self.prompts.from_keywords.format(keywords=body) - else: - system = self.prompts.system - user = self.prompts.from_scratch - proposal: Trajectory = { - 'messages': [{'role': 'system', 'content': system}, - {'role': 'user', 'content': user}], - } - # Carried through the explorer so build() knows which path this - # proposal took and what the second call has to be told. - proposals.append(attach_user_data( - proposal, keywords=picks, seeded=use_seed, two_step=two, - seed_query=(seed['query'] if two else ''), keyword_block=body)) - return proposals - - # ---------------------------------------------------------------- building - - def build(self, explored: List[Trajectory]) -> List[Optional[Trajectory]]: - """Parse, execute, self-check; None for every proposal that did not survive. - - The second call of the two-step path happens here rather than in - :meth:`propose`, because it needs the code the first call produced. It - goes out as one batch for the whole round, so the extra call costs one - more generate, not one per proposal. - """ - objs: List[Optional[Dict[str, Any]]] = [None] * len(explored) - # Proposals that died before parsing: they must not also be counted as a - # parse failure, because the cause -- and the fix -- is a different one. - dead: List[bool] = [False] * len(explored) - stage2_idx: List[int] = [] - stage2_prompts: List[Trajectory] = [] - for i, traj in enumerate(explored): - text = assistant_text(traj) - if not user_data_get(traj.get('user_data'), 'two_step', False): - objs[i] = parse_challenge(text) - continue - code = unwrap_code(text) - if not code.strip(): - # Usually a truncated completion: there is no solution to - # describe, so this proposal ends here. - self.stats['stage1_no_code'] += 1 - dead[i] = True - continue - objs[i] = {'_stage1_code': code} - stage2_idx.append(i) - stage2_prompts.append({ - 'messages': [ - {'role': 'system', 'content': self.prompts.system}, - {'role': 'user', 'content': self.prompts.two_step_problem.format( - code=code, - seed=user_data_get(explored[i].get('user_data'), 'seed_query', ''), - keywords=user_data_get(explored[i].get('user_data'), - 'keyword_block', ''))}, - ], - }) - if stage2_prompts: - logger.info(f'[CodeChallenger] two-step stage 2: {len(stage2_prompts)} problem ' - f'writes ({self.stats["stage1_no_code"]} first calls had no code)') - for i, reply in zip(stage2_idx, self.explore(stage2_prompts)): - stage1_code = objs[i]['_stage1_code'] - obj = parse_challenge(assistant_text(reply), require_solution=False) - if obj is not None: - # Ground truth is the code that actually ran, never the one - # the second call may have re-imagined. - obj['solution'] = stage1_code - objs[i] = obj - - return [None if dead[i] else self._finish(explored[i], obj) - for i, obj in enumerate(objs)] - - def _finish(self, proposal: Trajectory, obj: Optional[Dict[str, Any]]) -> Optional[Trajectory]: - """One parsed proposal -> a task, or None with a reason recorded.""" - if obj is None: - self.stats['parse_fail'] += 1 - return None - self.stats['parsed'] += 1 - - def _reject(reason: str, **extra: Any) -> None: - self.stats[reason] += 1 - if self.reject_sink is not None: - self.reject_sink({'reason': reason, **extra, **obj}) - - if len(obj['problem']) > self.problem_max_chars: - _reject('too_long') - return None - asserts = build_asserts(obj['solution'], obj['checks'], self.env(), - timeout=self.sandbox_timeout, max_checks=self.max_checks) - if not asserts: - _reject('gt_fail') - return None - check_script = '\n'.join(asserts) - if not self._check(obj['solution'], check_script)[0]: - # A reference solution that fails its own asserts is not ground - # truth, whatever the statement says. - _reject('selfcheck_fail', asserts=asserts) - return None - if self.drop_constant_answer and is_constant_answer(asserts): - _reject('constant_answer', asserts=asserts) - return None - - user_data = proposal.get('user_data') - # The task the solver is trained on: the statement alone, exactly as the - # difficulty stage will present it, with no instructions from the - # challenger's own prompt leaking in. - task: Trajectory = { - 'messages': [{'role': 'system', 'content': self.prompts.solver_system}, - {'role': 'user', 'content': obj['problem']}], - } - return attach_user_data( - task, - # One script rather than a list of asserts, named as the agentic half - # names it: a consumer that trains on both halves then reads the - # verifier the same way. setup_script is where that half puts the - # part that runs before the checks; nothing here needs one. - check_script=check_script, - setup_script='', - solution=obj['solution'], - entry=obj['entry'], - keywords=user_data_get(user_data, 'keywords', []), - seeded=user_data_get(user_data, 'seeded', False), - two_step=user_data_get(user_data, 'two_step', False)) - - # -------------------------------------------------------------- difficulty - - def solver_prompt(self, task: Trajectory) -> Trajectory: - problem = next((m['content'] for m in reversed(task.get('messages') or []) - if m.get('role') == 'user'), '') - return { - 'messages': [{'role': 'system', 'content': self.prompts.solver_system}, - {'role': 'user', - 'content': self.prompts.solver_user.format(problem=problem)}], - } - - def _check(self, code: str, check_script: str) -> Tuple[bool, str]: - """Did ``code`` pass ``check_script``? With the output, for feedback. - - Slot 0 always: a code judgement is one script with no state to share, and - this half runs them one at a time -- see the note in - :meth:`judge_attempt` -- so the slots the agentic half needs for its - concurrent episodes have nothing to do here. - """ - return run_check_script(code, check_script, self.env(), self.sandbox_timeout) - - def judge_attempt(self, task: Trajectory, attempt: Trajectory) -> bool: - """Did this attempt's code pass the task's asserts? - - Also hands the whole attempt to ``solver_sink`` when one is given. The - stage otherwise reduces each task to one number and drops the attempts, - and they are exactly what a solver trains on: a task kept at 3 of 8 is - one prompt answered eight times with a binary reward, which is a GRPO - group already measured to have a gradient. Sampling them again after the - band has been applied pays for the same tokens twice and can still land - the group at 0 or 8, where the advantage is the reward minus itself. - """ - check_script = user_data_get(task.get('user_data'), 'check_script', '') or '' - passed, output = self._check(unwrap_code(assistant_text(attempt)), check_script) - if self.solver_sink is not None: - # No lock: the difficulty stage judges attempts one at a time, in the - # loop that counts them, unlike the agentic half where each judgement - # is a sandbox round trip worth running concurrently. - self.solver_sink({ - 'statement': next((m.get('content', '') for m in task.get('messages') or [] - if m.get('role') == 'user'), ''), - # What the caller groups on: two problems with byte-identical - # asserts are the same problem, and a kept task carries this - # field through unchanged. - 'check_script': check_script, - 'passed': passed, - 'output': output, - 'truncated': bool((attempt or {}).get('truncated')), - 'attempt': attempt, - }) - return passed - - def on_difficulty_measured(self, candidates: List[Trajectory]) -> None: - """Remember the topics behind the candidates nobody solved.""" - if self.keywords is not None: - self.keywords.remember_unsolved(candidates, self.low_pass_expand) - - # ------------------------------------------------------------- feedback - - def expand_hard_keywords(self) -> int: - """Brainstorm more topics in the families that produced the hardest tasks. - - Called by whoever drives the challenger, after generating, so the bank - drifts toward material the solver actually struggles with. Returns how - many new keywords were added. - """ - return self.keywords.expand_hard() if self.keywords is not None else 0 diff --git a/src/twinkle_agentic/challenger/new/keyword.py b/src/twinkle_agentic/challenger/keyword.py similarity index 100% rename from src/twinkle_agentic/challenger/new/keyword.py rename to src/twinkle_agentic/challenger/keyword.py diff --git a/src/twinkle_agentic/challenger/keywords.py b/src/twinkle_agentic/challenger/keywords.py deleted file mode 100644 index 3f6e7bcef..000000000 --- a/src/twinkle_agentic/challenger/keywords.py +++ /dev/null @@ -1,588 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""The keyword bank: what the next task gets built about. - -Every challenger here proposes from a topic rather than from a fixed prompt, -because a fixed prompt collapses onto a handful of archetypes within a few -hundred proposals. The cycle that prevents it is the same one everywhere -- draw -a combination, refill whichever category ran dry, ask for more of whatever the -solver could not solve -- so it lives here once, in :class:`KeywordBank`, and a -proposer *holds* one rather than inheriting it. Nothing in this module knows what -a task looks like: it deals in short strings and in the prompts its owner supplies, -which is why the two challengers in this package and the RSI drivers in -``cookbook/rsi`` can all share it. - -Two failures shaped the file, both from real runs, both recorded where they hit: -a refill that returns nothing has to be loud (a silent one leaves every proposal -falling back to the from-scratch prompt while the run looks healthy), and a -keyword that arrives written as a sentence has to be counted rather than dropped -in silence -- see :func:`split_keyword_list`. -""" -import json -import os -import random -import threading -from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple - -from twinkle.data_format import SamplingParams, Trajectory, user_data_get -from twinkle.utils import get_logger -from twinkle_agentic.utils.code_utils import strip_reasoning -from twinkle_agentic.utils.message_utils import assistant_text -from .base import Explorer, KeywordPrompts, map_parallel - -logger = get_logger() - -__all__ = [ - 'KEYWORD_MAX_LEN', 'KeywordBank', 'KeywordPrompts', 'KeywordStore', - 'parse_keyword_list', 'split_keyword_list', -] - -# A keyword is a topic to build a task around, not a task statement. Past this many -# characters the model has written the second thing, and storing it makes the next -# prompt ask for a variation on a sentence rather than on a subject. -KEYWORD_MAX_LEN = 60 - - -def split_keyword_list(text: str) -> Tuple[List[str], List[str]]: - """Extract a JSON array of short strings; return (kept, dropped for length). - - The dropped half exists because it used to be discarded inside a list - comprehension. A refill that returned eight well-formed keywords, all of them - written out as sentences, reached the caller as an empty list and was recorded - as ``n_parsed: 0`` -- the same three characters a garbled reply, a timeout and - an over-length reply all produce, so the log could not tell them apart. One - iteration lost 27% of its keywords that way and the cause was found by - re-parsing the stored replies by hand. - - The bias is the reason to count rather than only to log: length correlates with - specificity, so the filter removes "Compute the critical path delay through a - gate-level netlist with annotated cell delays" and keeps whatever was vague - enough to be short. That is the opposite of what the bank is for. - """ - body = strip_reasoning(text) - start, end = body.find('['), body.rfind(']') - if start < 0 or end <= start: - return [], [] - try: - arr = json.loads(body[start:end + 1]) - except (ValueError, TypeError): - return [], [] - kept: List[str] = [] - dropped: List[str] = [] - for x in arr: - if not isinstance(x, str): - continue - s = x.strip() - if not s: - continue - (kept if len(s) <= KEYWORD_MAX_LEN else dropped).append(s) - return kept, dropped - - -def parse_keyword_list(text: str) -> List[str]: - """The kept half of :func:`split_keyword_list`, for callers with nothing to record.""" - return split_keyword_list(text)[0] - - -# ── keyword bank ─────────────────────────────────────────────────────────── -class KeywordStore: - """Persistent keyword bank with usage tracking, one bucket per category. - - Keywords exist to stop the challenger collapsing onto a handful of - archetypes. They are consumed rather than sampled with replacement, so a - run keeps reaching for topics it has not used; when a bucket runs dry the - caller refills it from the model, and recycles only if the model has run out - of distinct ideas. - - On-disk format (one JSON per line):: - - {"category", "text", "used": bool, "used_count": int, - "source": "gen"|"expand", "parent": <keyword or null>} - - De-duplicates case-insensitively within a category, so re-runs never - conflict with the bank on disk. - """ - - def __init__(self, path: str, categories: Sequence[str]): - if not categories: - raise ValueError('KeywordStore needs at least one category') - self.path = path - self.categories = tuple(categories) - self.items: Dict[str, List[Dict[str, Any]]] = {c: [] for c in self.categories} - self._seen: Dict[str, set] = {c: set() for c in self.categories} - if path and os.path.exists(path): - with open(path, encoding='utf-8') as f: - for line in f: - line = line.strip() - if not line: - continue - try: - r = json.loads(line) - except (ValueError, TypeError): - continue - c, t = r.get('category'), r.get('text') - if c in self.items and isinstance(t, str) and t.strip(): - key = t.strip().lower() - if key not in self._seen[c]: - self._seen[c].add(key) - self.items[c].append(r) - - def save(self) -> None: - """Write the bank out, atomically. A bank without a path is in-memory only.""" - if not self.path: - return - os.makedirs(os.path.dirname(os.path.abspath(self.path)) or '.', exist_ok=True) - tmp = self.path + '.tmp' - with open(tmp, 'w', encoding='utf-8') as f: - for c in self.categories: - for r in self.items[c]: - f.write(json.dumps(r, ensure_ascii=False) + '\n') - os.replace(tmp, self.path) - - def add(self, category: str, texts: Sequence[str], source: str = 'gen', - parent: Optional[str] = None) -> int: - added = 0 - for t in texts: - key = t.strip().lower() - if not key or key in self._seen[category]: - continue - self._seen[category].add(key) - self.items[category].append({'category': category, 'text': t.strip(), - 'used': False, 'used_count': 0, - 'source': source, 'parent': parent}) - added += 1 - return added - - def unused(self, category: str) -> List[Dict[str, Any]]: - return [r for r in self.items[category] if not r.get('used')] - - def texts(self, category: str) -> List[str]: - return [r['text'] for r in self.items[category]] - - def take(self, category: str, rng: random.Random) -> Optional[str]: - """Consume one unused keyword from ``category``; None if it is dry.""" - un = self.unused(category) - if not un: - return None - r = rng.choice(un) - r['used'] = True - r['used_count'] = r.get('used_count', 0) + 1 - return r['text'] - - def recycle(self, category: str) -> None: - """Mark every keyword unused again (safety valve when the model is tapped out).""" - for r in self.items[category]: - r['used'] = False - - -class KeywordBank: - """The draw / refill / expand cycle over a :class:`KeywordStore`. - - Held by whoever proposes rather than inherited: the two challengers in this - package share this entire cycle and almost nothing else, and the RSI drivers - in ``cookbook/rsi`` are not challengers at all yet need exactly the same - thing. Safe to drive from several threads -- see :meth:`draw`. - - Args: - store: the bank this works on. - prompts: the :class:`KeywordPrompts` this bank sends. - category_desc: category -> description, shown when asking for more. Must - cover every category in the store, or a dry one could not be refilled. - explorer: batch-in / batch-out generation, used for keyword calls only. - Worth keeping separate from the proposing explorer: brainstorming a - list is a text round, so a tool-calling rollout both wastes turns and - may take a bracketed list in the reply for a tool call. - rng: shared with the owner, so one seed reproduces the whole run. - name: what log lines call this bank; normally the owner's class name. - sampling_params: params for keyword calls. None sends the explorer's own. - sink: called once per keyword call with the prompt, the reply and both - halves of the parse. The one question such a dump exists to answer -- - did the model disobey the format, or does the parser reject what it - produced -- cannot be answered from a count. - combo_arity: ``'triple'`` draws one keyword per category; ``'mix'`` draws - a random 1..len(categories) subset. - arity_weights: sampling weights for the ``'mix'`` subset size. - single_kw_prob: in ``'triple'`` mode, the chance of using one category - instead of all of them. - refill_target: how many new keywords one refill aims for. - gen_calls: how many model calls it may spend on that. - refill_concurrency: how many of those go out together. At 1 every call is - told what the ones before it produced, which is what the avoid list is - for; raising it is faster and comes back with more synonyms. - refill_tries: refills to attempt before recycling a tapped-out category. - min_batch: smallest batch worth sending -- a sampler shards a batch over - its data-parallel workers, and a smaller one leaves some with nothing - to do. Set it to the number of sampler workers. - expand_per_kw / expand_max_kws: size of the :meth:`expand_hard` ask. - """ - - # How many phrases the 'do not repeat these' line may quote in total. There - # has to be a ceiling in both directions: too few and a serial refill stops - # seeing what it just said, too many and the model runs out of room to obey. - # Measured on armA2ser, where this refill's own output went in uncapped: with - # 130 quoted the eighth call was still answering normally, with 150 it started - # inventing -- 'îRAPIÓN holistic replace', 'ซะ subspace cutter map limit', 10 - # of 480 phrases that run. 100 sits below where that began. - _AVOID_TOTAL = 100 - _AVOID_LEAD = '\nDo NOT repeat any of these already-used topics: ' - - def __init__( - self, - store: KeywordStore, - *, - prompts: KeywordPrompts, - category_desc: Dict[str, str], - explorer: Explorer, - rng: random.Random, - name: str = 'keywords', - sampling_params: Optional[SamplingParams] = None, - sink: Optional[Callable[[Dict[str, Any]], None]] = None, - combo_arity: str = 'triple', - arity_weights: Optional[Sequence[float]] = None, - single_kw_prob: float = 0.1, - refill_target: int = 128, - gen_calls: int = 8, - refill_concurrency: int = 1, - refill_tries: int = 2, - min_batch: int = 1, - expand_per_kw: int = 8, - expand_max_kws: int = 32, - ): - if combo_arity not in ('triple', 'mix'): - raise ValueError(f"combo_arity must be 'triple' or 'mix', got {combo_arity!r}") - if refill_concurrency < 1: - raise ValueError(f'refill_concurrency must be >= 1, got {refill_concurrency}') - missing = [c for c in store.categories if not (category_desc or {}).get(c)] - if missing: - raise ValueError(f'category_desc is missing a description for {missing}; ' - f'a dry category could not be refilled.') - self.store = store - self.prompts = prompts - self.category_desc = dict(category_desc) - self.explorer = explorer - self.rng = rng - self.name = name - self.sampling_params = sampling_params - self.sink = sink - self.combo_arity = combo_arity - self.arity_weights = list(arity_weights) if arity_weights else None - self.single_kw_prob = single_kw_prob - self.refill_target = refill_target - self.gen_calls = gen_calls - self.refill_concurrency = refill_concurrency - self.refill_tries = refill_tries - self.min_batch = max(1, min_batch) - self.expand_per_kw = expand_per_kw - self.expand_max_kws = expand_max_kws - # One draw at a time; see :meth:`draw` for why the whole draw and not just - # the store access. - self._draw_lock = threading.Lock() - # Held while the rng, the nonce, the bank or the hard list are touched, and - # never across a model call. Separate from the sink lock, which waits on - # disk: a refill running in another thread must not queue behind a write. - self._state_lock = threading.Lock() - self._sink_lock = threading.Lock() - # Perturbs prompts so two calls are never byte-identical. - self._nonce = 0 - # (category, keyword) behind whatever nobody could solve, for expand_hard. - self._hard: List[Tuple[str, str]] = [] - - # -------------------------------------------------------------- drawing - - @property - def categories(self) -> Tuple[str, ...]: - return self.store.categories - - @staticmethod - def block(picks: Sequence[Tuple[str, str]]) -> str: - """The drawn keywords as the line block a prompt's ``{keywords}`` takes.""" - return '\n'.join(f'- {c}: {t}' for c, t in picks) - - def draw(self) -> List[Tuple[str, str]]: - """Consume one keyword combination, refilling whatever ran dry first. - - Serialised as a whole rather than per bank access: a refill is a batch of - model calls whose prompts quote what the calls before them produced, and - two draws overlapping would each refill without seeing the other's - keywords -- exactly what the avoid list exists to prevent. - """ - with self._draw_lock: - cats = self._pick_categories() - # Refill every dry category at once rather than as each one is reached: - # they are independent model calls that used to run one after another - # (20s each at the start of a run) and they touch separate buckets. - dry = [c for c in cats if not self.store.unused(c)] - if dry: - map_parallel(self.refill, dry) - picks: List[Tuple[str, str]] = [] - for category in cats: - with self._state_lock: - text = self.store.take(category, self.rng) - if text is not None: - picks.append((category, text)) - return picks - - def _pick_categories(self) -> List[str]: - """Which categories one draw covers, per ``combo_arity``.""" - categories = self.store.categories - with self._state_lock: - if self.combo_arity == 'mix': - if self.arity_weights and len(self.arity_weights) == len(categories): - k = self.rng.choices(range(1, len(categories) + 1), - weights=self.arity_weights)[0] - else: - k = self.rng.randint(1, len(categories)) - return self.rng.sample(list(categories), k) - if self.rng.random() < self.single_kw_prob: - return [self.rng.choice(categories)] - return list(categories) - - # ------------------------------------------------------------- refilling - - def refill(self, category: str) -> None: - """Ask the model for more keywords in ``category``; recycle if it is tapped out. - - Says so when it comes back empty. A silent no-op here is the worst outcome - available: :meth:`draw` then hands out no keywords, every proposal quietly - falls back to the from-scratch prompt, and the run looks normal while - producing one identical prompt over and over. That is exactly what happened - for whole runs when the prompt asked for one keyword per line and the parser - wanted a JSON array. - """ - tries = 0 - while not self.store.unused(category): - new = self._generate(category, self.refill_target) - with self._state_lock: - added = self.store.add(category, new, source='gen') - if added: - # Saved now rather than at the end of the run: a refill costs a - # batch of model calls, and a run that crashes later should not - # have to spend them again -- the next iteration reads this file - # to know what was already used. - self.store.save() - tries += 1 - if added: - logger.info(f'[{self.name}] keyword category {category!r} refilled ' - f'+{added} (try {tries})') - continue - logger.warning( - f'[{self.name}] keyword refill for {category!r} produced nothing on try ' - f'{tries}: {len(new)} parsed, 0 new. Proposals will run without keywords ' - f'unless this recovers -- pass a keyword sink to see the replies.') - if tries >= self.refill_tries: - with self._state_lock: - # Every keyword marked unused again. The alternative is a - # category that can never be drawn from, which stops the run: a - # repeat draw is worse than no run only if diversity matters - # more than collecting anything at all. - n_recycled = len(self.store.items[category]) - if n_recycled: - self.store.recycle(category) - self.store.save() - if n_recycled: - logger.info(f'[{self.name}] keyword category {category!r} exhausted ' - f'-> recycled {n_recycled} topics') - break - - def _generate(self, category: str, n_want: int) -> List[str]: - """Up to ``n_want`` keywords the bank does not already hold.""" - if n_want <= 0: - return [] - with self._state_lock: - known = self.store.texts(category) - n_calls = max(self.gen_calls, self.min_batch) - per_call = max(1, -(-n_want // n_calls) + 4) # ceil(n/calls) + margin - seen = {t.strip().lower() for t in known} - out: List[str] = [] - n_long = 0 - for start in range(0, n_calls, self.refill_concurrency): - group = range(start, min(start + self.refill_concurrency, n_calls)) - # Every call in a group is built before any of them runs, so they all - # carry the same avoid list -- which is exactly the batched behaviour, - # and why a group of one is what lets call k+1 see call k. - users = [(self.prompts.user.format( - k=per_call, desc=self.category_desc[category]) - + self._avoid_note(known, out) - + f'\n(batch {self._next_nonce()}-{i})') for i in group] - for user, reply in zip(users, self._explore(users)): - text = assistant_text(reply) - parsed, dropped_long = split_keyword_list(text) - n_long += len(dropped_long) - fresh = [kw for kw in parsed if kw.lower() not in seen] - seen.update(kw.lower() for kw in fresh) - out.extend(fresh) - # Full text, both sides: the question this dump answers is whether - # the model disobeyed the format or the parser rejected what it - # produced, and a count cannot say which. - self._record({ - 'category': category, 'prompt': user, 'reply': text, - 'stop_reason': reply.get('stop_reason'), - 'truncated': bool(reply.get('truncated')), - 'parsed': parsed, 'n_parsed': len(parsed), 'n_new': len(fresh), - # Which backend answered, for an explorer that has more than - # one -- an API with a local fallback is the case this exists - # for, and nothing else here could know which one ran. - 'via': reply.get('via'), - # The two fields that make the sentence above true. Without them - # ``n_parsed: 0`` reads the same whether the reply was garbled, - # empty, or eight usable keywords written at sentence length -- - # and the third is the one that happened. - 'dropped_long': dropped_long, 'n_dropped_long': len(dropped_long), - }) - if len(out) >= n_want: - # The surplus is dropped below, so further calls would buy nothing. - break - if n_long and self.sink is None: - # Without a dump to write to, the count has to be said out loud or the - # refill looks like the model simply produced less. - logger.warning(f'[{self.name}] dropped {n_long} keyword(s) over ' - f'{KEYWORD_MAX_LEN} chars while refilling; the prompt is ' - f'asking for task statements rather than topics') - with self._state_lock: - self.rng.shuffle(out) - return out[:n_want] - - def _avoid_note(self, older: List[str], fresh: List[str]) -> str: - """The 'do not repeat these' line, newest first, capped at ``_AVOID_TOTAL``. - - What this refill has just produced comes first and evicts older entries - rather than the reverse -- the calls run one at a time so that each can - avoid what the ones before it said, and dropping those would undo it. Past - the cap the oldest of *this refill's* phrases are what falls off, which is - also the least costly thing to drop: the model has already moved away from - them. - """ - fresh_shown = list(fresh)[-self._AVOID_TOTAL:] - room = max(0, self._AVOID_TOTAL - len(fresh_shown)) - with self._state_lock: - shown = older if len(older) <= room else self.rng.sample(older, room) - avoid = fresh_shown + list(shown) - return self._AVOID_LEAD + ', '.join(avoid) if avoid else '' - - # -------------------------------------------------------------- feedback - - def remember_hard(self, picks: Iterable[Sequence[str]]) -> None: - """Note the (category, keyword) pairs behind something nobody solved. - - De-duplicated case-insensitively and kept in arrival order. This is the - only feedback the bank gets from difficulty; without it, it drifts wherever - the refill prompt happens to go. - """ - with self._state_lock: - seen = {(c, t.lower()) for c, t in self._hard} - for pick in picks or (): - if not (isinstance(pick, (list, tuple)) and len(pick) >= 2): - continue - category, text = pick[0], pick[1] - if (category, text.lower()) not in seen: - seen.add((category, text.lower())) - self._hard.append((category, text)) - - def remember_unsolved(self, candidates: Sequence[Trajectory], max_pass: int = 0) -> None: - """:meth:`remember_hard` for measured candidates at or below ``max_pass``. - - Each candidate carries ``n_pass`` and its keyword draw in ``user_data``, so - this is the whole of what a difficulty round feeds back to the bank. - """ - for task in candidates: - data = task.get('user_data') - if user_data_get(data, 'n_pass', 0) > max_pass: - continue - self.remember_hard(user_data_get(data, 'keywords', []) or []) - - def expand_hard(self) -> int: - """Brainstorm more topics in the families that produced the hardest tasks. - - Called by whoever drives the proposer, after a round, so the bank drifts - toward material the solver actually struggles with. Returns how many new - keywords were added. - """ - if not self._hard or self.expand_per_kw <= 0: - return 0 - template = self.prompts.expand_user.strip() - if not template: - raise ValueError(f'[{self.name}] measured {len(self._hard)} hard keyword(s) to ' - f'expand on, but the prompts carry no expand_user.') - with self._state_lock: - hard = self._hard[:self.expand_max_kws] - self.rng.shuffle(hard) - # Cycle a short list so the batch still covers every sampler worker. - reqs = list(hard) - while len(reqs) < self.min_batch: - reqs.append(hard[len(reqs) % len(hard)]) - # ``desc`` is offered alongside ``kw``/``m``: a template that does not ask - # for it ignores it, and one that does gets the category's rules with it. - users = [(template.format(kw=kw, m=self.expand_per_kw, - desc=self.category_desc.get(category, '')) - + f'\n(batch {self._next_nonce()}-{i})') - for i, (category, kw) in enumerate(reqs)] - added = 0 - n_long = 0 - for (category, kw), user, reply in zip(reqs, users, self._explore(users)): - text = assistant_text(reply) - parsed, dropped_long = split_keyword_list(text) - n_long += len(dropped_long) - with self._state_lock: - added += self.store.add(category, parsed, source='expand', parent=kw) - # The prompt goes in whole, as the refill path already does. This used - # to record the literal string 'expand', which left the dump unable to - # answer the one question it gets asked -- whether a change to the - # expansion prompt was live in a given iteration. - self._record({ - 'category': category, 'parent': kw, 'prompt': user, 'reply': text, - 'stop_reason': reply.get('stop_reason'), - 'truncated': bool(reply.get('truncated')), - 'parsed': parsed, 'n_parsed': len(parsed), - 'dropped_long': dropped_long, 'n_dropped_long': len(dropped_long), - 'via': reply.get('via'), - }) - if added: - with self._state_lock: - self.store.save() - if n_long and self.sink is None: - logger.warning(f'[{self.name}] dropped {n_long} expanded keyword(s) over ' - f'{KEYWORD_MAX_LEN} chars; expansion follows the parent, so a ' - f'wordy parent produces wordy children') - logger.info(f'[{self.name}] expanded {len(hard)} hard keyword(s) -> ' - f'+{added} same-domain topics') - return added - - def save(self) -> None: - """Write the bank out. Refills and expansions already do; this is for the end.""" - with self._state_lock: - self.store.save() - - # --------------------------------------------------------------- private - - def _next_nonce(self) -> int: - """A number no other call gets, so two prompts are never byte-identical. - - Shared across categories, which refill at the same time: two threads - reading the counter together would send the same prompt twice and halve the - diversity with nothing to show that it happened. - """ - with self._state_lock: - self._nonce += 1 - return self._nonce - - def _explore(self, users: Sequence[str]) -> List[Trajectory]: - """One batch of keyword calls, one per user message. - - Sampling params are only forwarded when set, so a plain callable explorer - that takes nothing else keeps working. - """ - if not users: - return [] - prompts: List[Trajectory] = [{ - 'messages': [{'role': 'system', 'content': self.prompts.system}, - {'role': 'user', 'content': user}], - } for user in users] - if self.sampling_params is None: - return self.explorer(prompts) - return self.explorer(prompts, sampling_params=self.sampling_params) - - def _record(self, record: Dict[str, Any]) -> None: - """Hand one keyword call to the sink, if there is one.""" - if self.sink is None: - return - with self._sink_lock: - self.sink(record) diff --git a/src/twinkle_agentic/challenger/new/__init__.py b/src/twinkle_agentic/challenger/new/__init__.py deleted file mode 100644 index 0719b04e5..000000000 --- a/src/twinkle_agentic/challenger/new/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -from .agentic import AgenticChallenger, parse_problem_statement -from .base import Challenger -from .keyword import KEYWORD_MAX_LEN, KeywordGenerator - -__all__ = [ - 'AgenticChallenger', - 'Challenger', - 'KEYWORD_MAX_LEN', - 'KeywordGenerator', - 'parse_problem_statement', -] diff --git a/src/twinkle_agentic/challenger/new/agentic.py b/src/twinkle_agentic/challenger/new/agentic.py deleted file mode 100644 index ca41486a7..000000000 --- a/src/twinkle_agentic/challenger/new/agentic.py +++ /dev/null @@ -1,511 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Agentic challenger: act in a sandbox, verify the result, then describe it.""" -import math -import random -import re -import uuid -from dataclasses import dataclass -from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple - -from twinkle.data_format import SamplingParams, Trajectory, attach_user_data, user_data_get -from twinkle.data_format.sampling import SampledSequence, SampleResponse -from twinkle.utils import get_logger -from twinkle_agentic.envs import Env -from twinkle_agentic.protocol.base import API -from twinkle_agentic.rollout import APISampler, MultiTurnRollout -from twinkle_agentic.summarizer import Summarizer -from twinkle_agentic.utils.code_utils import parse_fenced_code, strip_reasoning -from twinkle_agentic.utils.message_utils import assistant_text, msg_content_text, normalize_tool_calls -from .base import Challenger, _parallel -from .keyword import KeywordGenerator -from .recorder import RolloutRecorder - -__all__ = ['AgenticChallenger', 'parse_problem_statement'] - -logger = get_logger() - -_FENCED_BLOCK_RE = re.compile(r'```[^\r\n]*\r?\n(.*?)```', re.S) - - -def parse_problem_statement(text: str) -> Optional[str]: - """Return the statement after removing reasoning and one outer fence.""" - body = strip_reasoning(text).strip() - whole = _FENCED_BLOCK_RE.fullmatch(body) - if whole: - body = whole.group(1).strip() - return body or None - - -def _sample_one(sampler: Any, input_feature: Dict[str, Any], sampling_params: Optional[SamplingParams], - adapter_kwargs: Dict[str, Any]) -> SampledSequence: - responses = sampler.sample([input_feature], sampling_params=sampling_params, **adapter_kwargs) - if not isinstance(responses, list): - raise TypeError(f'expected List[SampleResponse] from sampler.sample, got ' - f'{type(responses).__name__}') - if len(responses) != 1: - raise RuntimeError(f'sampler returned {len(responses)} responses for a single request; ' - 'expected exactly one') - response = responses[0] - if not isinstance(response, SampleResponse): - raise TypeError(f'expected SampleResponse from sampler.sample, got ' - f'{type(response).__name__}') - if len(response.sequences) != 1: - raise RuntimeError(f'SampleResponse contains {len(response.sequences)} sequences; ' - 'expected exactly one') - sequence = response.sequences[0] - if not isinstance(sequence, SampledSequence): - raise TypeError(f'expected SampledSequence, got {type(sequence).__name__}') - return sequence - - -def _api_followup_response( - sampler: Any, - api: Optional[APISampler], - sampling_params: Optional[SamplingParams], - *, - input_feature: Dict[str, Any], - adapter_kwargs: Dict[str, Any], - followups: int, - **kwargs: Any, -) -> SampledSequence: - """Use the API for appended stages and the primary backend otherwise.""" - if followups: - if api is None: - raise ValueError('use_api=True requires an API backend') - return api(input_feature, sampling_params, **adapter_kwargs) - if sampler is not None: - return _sample_one(sampler, input_feature, sampling_params, adapter_kwargs) - if api is not None: - return api(input_feature, sampling_params, **adapter_kwargs) - raise ValueError('AgenticChallenger has neither a sampler nor an API backend') - - -@dataclass -class _ProposalResult: - trajectory: Trajectory - group_id: str = '' - task: Optional[Trajectory] = None - reason: str = '' - detail: str = '' - outcome: str = '' - n_pass: Optional[int] = None - reward: float = 0.0 - - -class AgenticChallenger(Challenger): - """Invent tool-using tasks by doing, checking, and describing them. - - ``backend`` drives exploration and solver attempts. When ``use_api`` is true, - ``api`` generates only the appended check-script and problem-statement turns; - those turns retain the masking semantics selected by ``api_appended_as`` in - ``rollout_kwargs``. - """ - - _system = ('You invent tasks for another agent to solve. You have a sandbox and ' - 'tools. Work in it first: build something real, then you will be asked ' - 'to verify it and to describe it.') - _from_scratch = ('Choose a task worth doing in this sandbox and do it now, using ' - 'your tools. Do not describe it yet.') - _from_keywords = ('Choose a task around these topics and do it now, using your ' - 'tools. Do not describe it yet.\n\nTopics: {keywords}') - _from_seed = ('Here is an earlier task:\n\n{seed}\n\nDo something in the same ' - 'spirit but different, using your tools now. Do not describe it yet.') - _from_seed_keywords = ('Here is an earlier task:\n\n{seed}\n\nDo something in the same ' - 'spirit but different, may be more complex and interesting and meaningful, ' - 'around these topics, using your tools now. Do not describe it yet.\n\n' - 'Topics: {keywords}') - _check_followup = ('Stop working. This is the workspace you produced:\n\n{final_state}\n\n' - 'Write a {language} script that verifies this end state, as a fenced ' - '{language} code block and nothing else. It must exit with a non-zero status ' - 'if the work was not done. Check what can be read out of the files -- their ' - 'structure and the values inside them. NEVER check a file size in bytes, a ' - 'checksum, or the full source text of a script: correct solutions differ ' - 'there, and such a check only its own author can pass.') - _check_retry_followup = ('Your check script did not pass:\n\n{error}\n\nThe workspace is:\n\n' - '{final_state}\n\nReturn a corrected script as a fenced {language} code block ' - 'and nothing else.') - _check_parse_error = ('Could not read a check script from your reply: it was not a ' - 'fenced {language} code block. Do not wrap it in a tool call and ' - 'do not add prose -- return ONLY a fenced {language} code block.') - _problem_followup = ('Now write the task statement: what someone starting from an empty workspace ' - 'would have to be told to produce what you produced, and nothing about how you ' - 'did it. Name the files to create and quote any input data verbatim. Do not ' - 'reveal values your check script computes. Reply with the statement only.') - - - def __init__( - self, - backend: Any, - *, - api: Optional[Any] = None, - use_api: bool = False, - keyword_generator: Optional[KeywordGenerator] = None, - trajectory_seed: Optional[List[Trajectory]] = None, - summarizer: Optional[Summarizer] = None, - system_prompt: Optional[str] = None, - from_scratch_prompt: Optional[str] = None, - from_keywords_prompt: Optional[str] = None, - from_seed_prompt: Optional[str] = None, - from_seed_keywords_prompt: Optional[str] = None, - check_followup_prompt: Optional[str] = None, - check_retry_followup_prompt: Optional[str] = None, - check_parse_error_prompt: Optional[str] = None, - problem_followup_prompt: Optional[str] = None, - check_retries: int = 1, - problem_max_chars: int = 8192, - check_language: str = 'python', - parse_check_fn: Optional[Callable[[str], Optional[str]]] = None, - pass_rate_target: float = 0.2, - envs: Sequence[Env] = (), - num_challenger_rollouts: int = 8, - num_solver_rollouts: int = 8, - pass_band: Tuple[float, float] = (1.0, 7.0), - pass_rate_width: float = 0.3, - max_empty_rounds: int = 0, - followup_params: Optional[SamplingParams] = None, - checker: Optional[Callable[[Trajectory], bool]] = None, - save_dir: Optional[str] = None, - save_failed_rollouts: bool = True, - **rollout_kwargs: Any, - ): - super().__init__( - envs=envs, - num_challenger_rollouts=num_challenger_rollouts, - num_solver_rollouts=num_solver_rollouts, - pass_band=pass_band, - max_empty_rounds=max_empty_rounds, - ) - if check_retries < 0: - raise ValueError(f'check_retries must be >= 0, got {check_retries}') - if problem_max_chars <= 0: - raise ValueError(f'problem_max_chars must be positive, got {problem_max_chars}') - if not check_language.strip(): - raise ValueError('check_language must not be empty') - if not 0 <= pass_rate_target <= 1: - raise ValueError(f'pass_rate_target must be in [0, 1], got {pass_rate_target}') - if pass_rate_width <= 0: - raise ValueError(f'pass_rate_width must be positive, got {pass_rate_width}') - if use_api and rollout_kwargs.get('response_callback') is not None: - raise ValueError('use_api=True cannot be combined with response_callback') - backend_is_api = isinstance(backend, (API, APISampler)) - if use_api and api is None and not backend_is_api: - raise ValueError('use_api=True requires api= when backend is a sampler') - self.keyword_generator = keyword_generator - self.trajectory_seed = list(trajectory_seed or ()) - self.summarizer = summarizer - self._system = self._system if system_prompt is None else system_prompt - self._from_scratch = self._from_scratch if from_scratch_prompt is None else from_scratch_prompt - self._from_keywords = self._from_keywords if from_keywords_prompt is None else from_keywords_prompt - self._from_seed = self._from_seed if from_seed_prompt is None else from_seed_prompt - self._from_seed_keywords = (self._from_seed_keywords if from_seed_keywords_prompt is None else - from_seed_keywords_prompt) - self._check_followup = self._check_followup if check_followup_prompt is None else check_followup_prompt - self._check_retry_followup = (self._check_retry_followup if check_retry_followup_prompt is None else - check_retry_followup_prompt) - self._check_parse_error = (self._check_parse_error if check_parse_error_prompt is None else - check_parse_error_prompt) - self._problem_followup = (self._problem_followup if problem_followup_prompt is None else - problem_followup_prompt) - self._check_retries = check_retries - self._problem_max_chars = problem_max_chars - self._check_language = check_language.strip().lower() - self._parse_check_fn = parse_check_fn - self._pass_rate_target = pass_rate_target - self._pass_rate_width = pass_rate_width - self.checker = checker - self.followup_params = followup_params - self.rng = random.Random() - self.use_api = use_api - self.save_failed_rollouts = save_failed_rollouts - self._recorder = RolloutRecorder(save_dir) if save_dir else None - self._round_proposals: List[_ProposalResult] = [] - self._backend = backend - self._rollout_kwargs = dict(rollout_kwargs) - if api is not None: - self._rollout_kwargs['api'] = api - if use_api: - self._rollout_kwargs['response_callback'] = _api_followup_response - self._rollout: Optional[MultiTurnRollout] = None - self._tool_schemas = self.env().tools() or None - - def _rollout_instance(self) -> MultiTurnRollout: - if self._rollout is None: - self._rollout = MultiTurnRollout(self._backend, **self._rollout_kwargs) - return self._rollout - - def _tool_manager(self, slot: int) -> Optional[Any]: - env = self.env(slot) - return env.tool_manager() if env.tools() else None - - def _summary(self, trajectory: Trajectory) -> str: - turns: List[str] = [] - for message in trajectory.get('messages') or []: - if not isinstance(message, dict): - continue - role = message.get('role') or '' - if role == 'system': - continue - parts = [msg_content_text(message).strip()] - for call in normalize_tool_calls(message) or (): - fn = call.get('function') or {} - if isinstance(fn, dict) and fn.get('name'): - parts.append(f"calls {fn['name']}({fn.get('arguments') or ''})") - body = '\n'.join(part for part in parts if part) - if body: - turns.append(f'{role}: {body}') - text = '\n'.join(turns) - if not text: - return '' - return self.summarizer(text) if self.summarizer is not None else text - - def _build_challenge_prompt(self) -> Optional[Trajectory]: - keywords: List[str] = [] - if self.keyword_generator is not None: - groups = self.keyword_generator.get_keywords(1) - if not groups: - return None - keywords = groups[0] - seed = '' - if self.trajectory_seed: - seed = self._summary(self.rng.choice(self.trajectory_seed)) - block = ', '.join(keywords) - if seed and keywords: - user = self._from_seed_keywords.format(seed=seed, keywords=block) - elif seed: - user = self._from_seed.format(seed=seed) - elif keywords: - user = self._from_keywords.format(keywords=block) - else: - user = self._from_scratch - prompt: Trajectory = { - 'messages': [ - { - 'role': 'system', - 'content': self._system - }, - { - 'role': 'user', - 'content': user - }, - ], - } - if self._tool_schemas: - prompt['tools'] = self._tool_schemas - return attach_user_data(prompt, keywords=keywords, seeded=bool(seed)) - - def _explore(self, prompt: Trajectory) -> List[Trajectory]: - group_id = uuid.uuid4().hex - proposals: List[_ProposalResult] = [] - remaining = self.num_challenger_rollouts - while remaining > 0: - wave = min(self.n_slots, remaining) - proposals.extend(_parallel(lambda slot: self._run_episode(prompt, slot), wave)) - remaining -= wave - for proposal in proposals: - proposal.group_id = group_id - self._round_proposals = proposals - return [proposal.task for proposal in proposals if proposal.task is not None] - - def _run_episode(self, prompt: Trajectory, slot: int) -> _ProposalResult: - self.env(slot).clear() - state: Dict[str, Any] = {'slot': slot} - kwargs: Dict[str, Any] = { - 'followup_fn': lambda trajectory, n_before: self._followup(state, trajectory, n_before), - } - manager = self._tool_manager(slot) - if manager is not None: - kwargs['tool_manager'] = manager - explored = self._rollout_instance()([prompt], **kwargs) - if not explored: - self._reject(state, 'rollout_no_output') - return _ProposalResult(dict(prompt), reason='rollout_no_output') - trajectory = explored[0] - task = self._build_query(state, trajectory) - reason, detail = state.get('reject', ('', '')) - return _ProposalResult(trajectory, task=task, reason=reason, detail=detail) - - def _followup(self, state: Dict[str, Any], trajectory: Trajectory, - n_before: int) -> Optional[Tuple[str, Optional[SamplingParams]]]: - if state.get('checked'): - return None - reply = None if n_before == 0 else assistant_text(trajectory) - followup = self._build_test_case(state, reply) - if followup is None: - return None - return followup, self.followup_params - - def _build_test_case(self, state: Dict[str, Any], reply: Optional[str]) -> Optional[str]: - slot = state['slot'] - if reply is None: - snapshot, error = self.env(slot).snapshot() - state['snapshot'] = snapshot - if not snapshot.strip(): - state['reject'] = ('snapshot_unavailable' if error else 'empty_workspace', error) - return None - return self._check_followup.format(final_state=snapshot, language=self._check_language) - - attempt = state.get('check_attempts', 0) + 1 - state['check_attempts'] = attempt - script = (self._parse_check_fn(reply) if self._parse_check_fn is not None else - parse_fenced_code(reply, language_tags=None)) - if script is None: - if attempt <= self._check_retries: - return self._check_retry_followup.format( - error=self._check_parse_error.format(language=self._check_language), - final_state=state.get('snapshot', ''), - language=self._check_language, - ) - state['reject'] = ('check_parse_fail', reply) - return None - state['script'] = script - exit_code, output = self.env(slot).run_script(script, interpreter=self._check_language) - if exit_code == 0: - state['checked'] = True - return self._problem_followup - after = self.env(slot).snapshot()[0] - state.setdefault('attempts', []).append(f'--- attempt {attempt}: exit {exit_code} ---\n{output}\n' - f'--- check script ---\n{script}') - if attempt <= self._check_retries: - return self._check_retry_followup.format( - error=output, - final_state=after or state.get('snapshot', ''), - language=self._check_language, - ) - state['reject'] = ('check_run_fail', '\n'.join(state['attempts'])) - return None - - def _build_query(self, state: Dict[str, Any], explored: Trajectory) -> Optional[Trajectory]: - if state.get('reject'): - return self._reject(state, *state['reject']) - if not state.get('checked'): - return self._reject( - state, - 'episode_cut_short', - f"stop_reason={explored.get('stop_reason')} " - f"truncated={bool(explored.get('truncated'))} " - f"turns={explored.get('turns')}", - ) - statement = parse_problem_statement(assistant_text(explored)) - if statement is None: - return self._reject(state, 'problem_parse_fail') - if len(statement) > self._problem_max_chars: - return self._reject(state, 'too_long', f'{len(statement)} chars') - task: Trajectory = attach_user_data( - {'messages': [{ - 'role': 'user', - 'content': statement - }]}, - check_script=state['script'], - keywords=user_data_get(explored.get('user_data'), 'keywords', []), - seeded=user_data_get(explored.get('user_data'), 'seeded', False), - ) - if self.checker is not None and not self.checker(task): - return self._reject(state, 'rejected_by_checker') - return task - - def _reject(self, state: Dict[str, Any], reason: str, detail: str = '') -> Optional[Trajectory]: - state['reject'] = (reason, detail) - logger.info(f'[{type(self).__name__}] rejected: {reason}' - f"{f' -- {detail[:400]}' if detail else ''}") - return None - - def _solver_prompt(self, task: Trajectory) -> Trajectory: - prompt: Trajectory = {'messages': [dict(message) for message in task.get('messages') or []]} - if self._tool_schemas: - prompt['tools'] = self._tool_schemas - return prompt - - def _judge(self, task: Trajectory, slot: int) -> bool: - script = user_data_get(task.get('user_data'), 'check_script', '') - if not script: - return False - return self.env(slot).run_script(script, interpreter=self._check_language)[0] == 0 - - def challenger_reward(self, n_pass: Optional[int]) -> float: - """Reward tasks near the target solver pass rate; unmeasured failures score zero.""" - if n_pass is None or not self.num_solver_rollouts or n_pass <= 0: - return 0.0 - gap = n_pass / self.num_solver_rollouts - self._pass_rate_target - variance = 2.0 * self._pass_rate_width**2 - return math.exp(-(gap * gap) / variance) - - def _record_proposals(self) -> None: - proposals, self._round_proposals = self._round_proposals, [] - if self._recorder is None: - return - for index, proposal in enumerate(proposals): - if proposal.task is None and not self.save_failed_rollouts: - continue - trajectory = dict(proposal.trajectory) - trajectory['rewards'] = proposal.reward - task_data = proposal.task.get('user_data') if proposal.task is not None else None - statement = '' - if proposal.task is not None: - statement = next((message.get('content', '') for message in proposal.task.get('messages') or [] - if isinstance(message, dict) and message.get('role') == 'user'), '') - self._recorder.write( - trajectory, - side='propose', - group_id=proposal.group_id, - proposal_index=index, - outcome=proposal.outcome or ('rejected' if proposal.reason else 'kept'), - reason=proposal.reason, - detail=proposal.detail, - reward=proposal.reward, - n_pass=proposal.n_pass, - n_rollouts=(self.num_solver_rollouts if proposal.n_pass is not None else None), - pass_rate=(proposal.n_pass / self.num_solver_rollouts - if proposal.n_pass is not None and self.num_solver_rollouts else None), - statement=statement, - check_script=user_data_get(task_data, 'check_script', ''), - keywords=user_data_get(proposal.trajectory.get('user_data'), 'keywords', []), - seeded=user_data_get(proposal.trajectory.get('user_data'), 'seeded', False), - ) - - def _filter_difficulty(self, tasks: List[Trajectory]) -> List[Trajectory]: - successful = [proposal for proposal in self._round_proposals if proposal.task is not None] - if len(successful) != len(tasks): - raise RuntimeError('proposal/task alignment failed before difficulty filtering') - if not tasks or not self.num_solver_rollouts: - for proposal in successful: - proposal.outcome = 'kept' - self._record_proposals() - return tasks - - passes = [0] * len(tasks) - plan = [i for i in range(len(tasks)) for _ in range(self.num_solver_rollouts)] - rollout = self._rollout_instance() - for start in range(0, len(plan), self.n_slots): - wave = plan[start:start + self.n_slots] - _parallel(lambda slot: self.env(slot).clear(), len(wave)) - prompts = [self._solver_prompt(tasks[i]) for i in wave] - kwargs: Dict[str, Any] = {} - managers = [self._tool_manager(slot) for slot in range(len(wave))] - if any(manager is not None for manager in managers): - kwargs['tool_manager'] = managers - attempts = rollout(prompts, **kwargs) - if len(attempts) != len(prompts): - raise RuntimeError(f'rollout returned {len(attempts)} attempts for ' - f'{len(prompts)} prompts; expected one per prompt') - verdicts = _parallel(lambda slot: self._judge(tasks[wave[slot]], slot), len(wave)) - for slot, passed in enumerate(verdicts): - if passed: - passes[wave[slot]] += 1 - - low, high = self.pass_band - measured = [ - attach_user_data(task, n_pass=n_pass, n_rollouts=self.num_solver_rollouts) - for task, n_pass in zip(tasks, passes) - ] - kept: List[Trajectory] = [] - for proposal, task, n_pass in zip(successful, measured, passes): - proposal.task = task - proposal.n_pass = n_pass - proposal.reward = self.challenger_reward(n_pass) - if low <= n_pass <= high: - proposal.outcome = 'kept' - kept.append(task) - else: - proposal.outcome = 'outside_band' - self._record_proposals() - return kept diff --git a/src/twinkle_agentic/challenger/new/base.py b/src/twinkle_agentic/challenger/new/base.py deleted file mode 100644 index 1c819533b..000000000 --- a/src/twinkle_agentic/challenger/new/base.py +++ /dev/null @@ -1,132 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Reusable lifecycle for task challengers.""" -from abc import ABC, abstractmethod -from concurrent.futures import ThreadPoolExecutor -from typing import Any, Callable, Iterator, List, Optional, Sequence, Tuple - -from twinkle.data_format import Trajectory -from twinkle.utils import get_logger -from twinkle_agentic.envs import Env - -logger = get_logger() - -__all__ = ['Challenger'] - - -def _parallel(fn: Callable[[int], Any], count: int) -> List[Any]: - """Run ``fn`` over ``range(count)`` concurrently, preserving order.""" - if count <= 1: - return [fn(i) for i in range(count)] - out: List[Any] = [None] * count - with ThreadPoolExecutor(max_workers=count) as pool: - futures = {pool.submit(fn, i): i for i in range(count)} - for future, i in futures.items(): - out[i] = future.result() - return out - - -class Challenger(ABC): - """Common batching and environment lifecycle for task challengers. - - Subclasses define how a round builds its prompt, explores it, and measures - candidate difficulty. One environment is owned by one concurrent job for the - complete lifetime of that job. - """ - - def __init__( - self, - *, - envs: Sequence[Env], - num_challenger_rollouts: int = 8, - num_solver_rollouts: int = 8, - pass_band: Tuple[float, float] = (1.0, 7.0), - max_empty_rounds: int = 0, - ): - if not envs: - raise ValueError('envs is empty: a challenger needs a workspace to act in and grade') - if num_challenger_rollouts < 1: - raise ValueError(f'num_challenger_rollouts must be >= 1, got ' - f'{num_challenger_rollouts}') - if num_solver_rollouts < 0: - raise ValueError(f'num_solver_rollouts must be >= 0, got {num_solver_rollouts}') - if max_empty_rounds < 0: - raise ValueError(f'max_empty_rounds must be >= 0, got {max_empty_rounds}') - if num_solver_rollouts: - if len(pass_band) != 2: - raise ValueError(f'pass_band is (low, high) in attempt counts, got {pass_band}') - low, high = pass_band - if not 0 <= low <= high <= num_solver_rollouts: - raise ValueError(f'pass_band must satisfy 0 <= low <= high <= num_solver_rollouts, got ' - f'{pass_band} against num_solver_rollouts={num_solver_rollouts}') - self.envs = list(envs) - self.num_challenger_rollouts = num_challenger_rollouts - self.num_solver_rollouts = num_solver_rollouts - self.pass_band = pass_band - self.max_empty_rounds = max_empty_rounds - self.n_proposed = 0 - self.n_kept = 0 - - @property - def n_slots(self) -> int: - """How many jobs may run at once: one per environment.""" - return len(self.envs) - - def env(self, slot: int = 0) -> Env: - """Return the current environment for ``slot``.""" - return self.envs[slot] - - @abstractmethod - def _build_challenge_prompt(self) -> Optional[Trajectory]: - """Build one round's shared prompt, or return None when exhausted.""" - - @abstractmethod - def _explore(self, prompt: Trajectory) -> List[Trajectory]: - """Generate and validate candidates from one shared prompt.""" - - @abstractmethod - def _filter_difficulty(self, tasks: List[Trajectory]) -> List[Trajectory]: - """Measure candidate difficulty and return the accepted tasks.""" - - def __call__(self, batch_size: int, total: Optional[int] = None) -> Iterator[List[Trajectory]]: - """Yield finished tasks in batches.""" - if batch_size <= 0: - raise ValueError(f'batch_size must be positive, got {batch_size}') - pending: List[Trajectory] = [] - produced = 0 - empty_rounds = 0 - while total is None or produced < total: - want = batch_size if total is None else min(batch_size, total - produced) - while len(pending) < want: - kept = self._round() - if kept is None: - if pending: - yield pending - return - if kept: - empty_rounds = 0 - pending.extend(kept) - continue - empty_rounds += 1 - if self.max_empty_rounds and empty_rounds >= self.max_empty_rounds: - logger.warning(f'[{type(self).__name__}] stopped after {empty_rounds} ' - 'consecutive rounds without a usable task') - if pending: - yield pending - return - yield pending[:want] - produced += want - pending = pending[want:] - - def _round(self) -> Optional[List[Trajectory]]: - """Run one proposal group; None means the source is exhausted.""" - prompt = self._build_challenge_prompt() - if prompt is None: - return None - verified = self._explore(prompt) - kept = self._filter_difficulty(verified) - self.n_proposed += self.num_challenger_rollouts - self.n_kept += len(kept) - logger.info(f'[{type(self).__name__}] {self.num_challenger_rollouts} episodes, ' - f'{len(verified)} verified, {len(kept)} in band ' - f'(cumulative {self.n_kept}/{self.n_proposed})') - return kept diff --git a/src/twinkle_agentic/challenger/new/recorder.py b/src/twinkle_agentic/challenger/recorder.py similarity index 100% rename from src/twinkle_agentic/challenger/new/recorder.py rename to src/twinkle_agentic/challenger/recorder.py diff --git a/src/twinkle_agentic/challenger/task_bank.py b/src/twinkle_agentic/challenger/task_bank.py deleted file mode 100644 index bf8663bde..000000000 --- a/src/twinkle_agentic/challenger/task_bank.py +++ /dev/null @@ -1,130 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""A file of the tasks earlier iterations already produced, to compare new ones against. - -Novelty is meaningless without something to be novel against. Within one run the -proposals of a group can be compared to each other, but the failure this is for is -slower than that: iteration k+1 re-proposing what iteration k already trained on. That -needs a file that outlives a run, which is what this is -- one JSON object per line, -appended by :meth:`add`, read back by the next run. - -The similarity used to pick which stored tasks to show the judge is 3-gram Jaccard over -the statement. It is a weak measure and known to be: measured over run_clean9's 188 -statements the closest pair scored 0.060, so ranking by it is nearly ranking at random, -and it cannot see that two tasks with no shared wording are both 'write the given files -verbatim, then derive one from them'. It is used only to CHOOSE the handful of tasks the -judge reads, never to score novelty -- the judging is -:mod:`twinkle_agentic.verifier.rubric_score`, whose criteria compare task shapes. Even a -near-random pick gives the judge real tasks from the same generator to compare against, -which is what the criteria need. -""" -import json -import os -import re -import threading -from typing import Any, Dict, List, Optional, Sequence, Set, Tuple - -__all__ = ['TaskBank', 'jaccard_3gram', 'grams'] - - -def grams(text: str, n: int = 3) -> Set[Tuple[str, ...]]: - words = re.findall(r'[a-z0-9_]+', (text or '').lower()) - return {tuple(words[i:i + n]) for i in range(max(0, len(words) - n + 1))} - - -def jaccard_3gram(a: Set[Tuple[str, ...]], b: Set[Tuple[str, ...]]) -> float: - if not a or not b: - return 0.0 - return len(a & b) / len(a | b) - - -class TaskBank: - """Statements from previous iterations, plus the ones this run adds. - - Args: - path: the JSONL file. A missing file is an empty bank, not an error -- the - first iteration has nothing to compare against and must still run. - refs: how many stored statements :meth:`references` returns. - """ - - def __init__(self, path: str, refs: int = 5): - self.path = path - self.refs = max(0, refs) - self._statements: List[str] = [] - self._grams: List[Set[Tuple[str, ...]]] = [] - self._seen: Set[str] = set() - self._lock = threading.Lock() - self.n_loaded = 0 - self.n_added = 0 - self._load() - - def _load(self) -> None: - if not self.path or not os.path.exists(self.path): - return - with open(self.path, encoding='utf-8') as f: - for line in f: - line = line.strip() - if not line: - continue - try: - rec = json.loads(line) - except json.JSONDecodeError: - # A half-written last line from a killed run. Skipped rather - # than fatal: losing one reference is not worth failing a run, - # and the count below says how many were read. - continue - statement = (rec.get('statement') or '').strip() - if statement and statement not in self._seen: - self._seen.add(statement) - self._statements.append(statement) - self._grams.append(grams(statement)) - self.n_loaded = len(self._statements) - - def __len__(self) -> int: - return len(self._statements) - - def references(self, statement: str, extra: Sequence[str] = ()) -> List[str]: - """The stored statements most similar to ``statement``, closest first. - - ``extra`` is prepended and never dropped -- it is how the proposals of the - current group get in front of the judge. Without them a whole group can be - scored identically novel against history while being eight versions of one - idea, and GRPO subtracts the group mean, so an identical term across the - group produces no gradient at all. - """ - with self._lock: - pairs = list(zip(self._statements, self._grams)) - target = grams(statement) - scored = [(jaccard_3gram(target, g), s) for s, g in pairs if s != statement] - scored.sort(key=lambda p: -p[0]) - out = [s for s in extra if s and s != statement] - out.extend(s for _, s in scored[:self.refs]) - return out - - def add(self, statement: str, check: str = '', **fields: Any) -> bool: - """Append one task. Returns False if the statement is already stored. - - Appended immediately rather than at the end of the run: a run that crashes - after 60 of 80 tasks should still contribute those 60, or the bank silently - under-reports what has been trained on. - """ - statement = (statement or '').strip() - if not statement: - return False - with self._lock: - if statement in self._seen: - return False - self._seen.add(statement) - self._statements.append(statement) - self._grams.append(grams(statement)) - self.n_added += 1 - if self.path: - rec: Dict[str, Any] = {'statement': statement, 'check': check} - rec.update(fields) - os.makedirs(os.path.dirname(self.path) or '.', exist_ok=True) - with open(self.path, 'a', encoding='utf-8') as f: - f.write(json.dumps(rec, ensure_ascii=False) + '\n') - return True - - def stats(self) -> Dict[str, Optional[int]]: - return {'path': self.path, 'loaded': self.n_loaded, 'added': self.n_added, - 'total': len(self._statements)}